diff --git a/src/cli.test.ts b/src/cli.test.ts index bb58b7a1..79964066 100644 --- a/src/cli.test.ts +++ b/src/cli.test.ts @@ -1,5 +1,5 @@ import { Command } from 'commander'; -import { parseEnvironmentVariables, parseDomains, parseDomainsFile, escapeShellArg, joinShellArgs, parseVolumeMounts, isValidIPv4, isValidIPv6, parseDnsServers, validateAgentImage, isAgentImagePreset, AGENT_IMAGE_PRESETS, processAgentImageOption, processLocalhostKeyword, validateSkipPullWithBuildLocal, validateAllowHostPorts, parseMemoryLimit, validateFormat, validateApiProxyConfig, buildRateLimitConfig, validateRateLimitFlags, validateApiTargetInAllowedDomains, DEFAULT_OPENAI_API_TARGET, DEFAULT_ANTHROPIC_API_TARGET, emitApiProxyTargetWarnings, formatItem, program, parseAgentTimeout, applyAgentTimeout } from './cli'; +import { parseEnvironmentVariables, parseDomains, parseDomainsFile, escapeShellArg, joinShellArgs, parseVolumeMounts, isValidIPv4, isValidIPv6, parseDnsServers, validateAgentImage, isAgentImagePreset, AGENT_IMAGE_PRESETS, processAgentImageOption, processLocalhostKeyword, validateSkipPullWithBuildLocal, validateAllowHostPorts, parseMemoryLimit, validateFormat, validateApiProxyConfig, buildRateLimitConfig, validateRateLimitFlags, validateApiTargetInAllowedDomains, DEFAULT_OPENAI_API_TARGET, DEFAULT_ANTHROPIC_API_TARGET, emitApiProxyTargetWarnings, formatItem, program, parseAgentTimeout, applyAgentTimeout, handlePredownloadAction } from './cli'; import { redactSecrets } from './redact-secrets'; import * as fs from 'fs'; import * as path from 'path'; @@ -1881,4 +1881,28 @@ describe('cli', () => { expect(result).toBe(' --flag'); }); }); + + describe('handlePredownloadAction', () => { + it('should delegate to predownloadCommand with correct options', async () => { + // Mock the predownload module that handlePredownloadAction dynamically imports + const mockPredownloadCommand = jest.fn().mockResolvedValue(undefined); + jest.mock('./commands/predownload', () => ({ + predownloadCommand: mockPredownloadCommand, + })); + + await handlePredownloadAction({ + imageRegistry: 'ghcr.io/test', + imageTag: 'v1.0', + agentImage: 'default', + enableApiProxy: false, + }); + + expect(mockPredownloadCommand).toHaveBeenCalledWith({ + imageRegistry: 'ghcr.io/test', + imageTag: 'v1.0', + agentImage: 'default', + enableApiProxy: false, + }); + }); + }); }); diff --git a/src/cli.ts b/src/cli.ts index ae9b9971..e1a705cd 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -1521,6 +1521,45 @@ export function validateFormat(format: string, validFormats: string[]): void { } } +// Predownload action handler - exported for testing +export async function handlePredownloadAction(options: { + imageRegistry: string; + imageTag: string; + agentImage: string; + enableApiProxy: boolean; +}): Promise { + const { predownloadCommand } = await import('./commands/predownload'); + try { + await predownloadCommand({ + imageRegistry: options.imageRegistry, + imageTag: options.imageTag, + agentImage: options.agentImage, + enableApiProxy: options.enableApiProxy, + }); + } catch (error) { + const exitCode = (error as Error & { exitCode?: number }).exitCode ?? 1; + process.exit(exitCode); + } +} + +// Predownload subcommand - pre-pull container images +program + .command('predownload') + .description('Pre-download Docker images for offline use or faster startup') + .option( + '--image-registry ', + 'Container image registry', + 'ghcr.io/github/gh-aw-firewall' + ) + .option('--image-tag ', 'Container image tag (applies to squid, agent, and api-proxy images)', 'latest') + .option( + '--agent-image ', + 'Agent image preset (default, act) or custom image', + 'default' + ) + .option('--enable-api-proxy', 'Also download the API proxy image', false) + .action(handlePredownloadAction); + // Logs subcommand - view Squid proxy logs const logsCmd = program .command('logs') diff --git a/src/commands/predownload.test.ts b/src/commands/predownload.test.ts new file mode 100644 index 00000000..679092a9 --- /dev/null +++ b/src/commands/predownload.test.ts @@ -0,0 +1,156 @@ +import { resolveImages, predownloadCommand, PredownloadOptions } from './predownload'; + +// Mock execa +jest.mock('execa', () => { + const mockExeca = jest.fn().mockResolvedValue({ stdout: '', stderr: '' }); + return { __esModule: true, default: mockExeca }; +}); + +// eslint-disable-next-line @typescript-eslint/no-require-imports +const execa = require('execa').default as jest.Mock; + +describe('predownload', () => { + describe('resolveImages', () => { + const defaults: PredownloadOptions = { + imageRegistry: 'ghcr.io/github/gh-aw-firewall', + imageTag: 'latest', + agentImage: 'default', + enableApiProxy: false, + }; + + it('should resolve squid and default agent images', () => { + const images = resolveImages(defaults); + expect(images).toEqual([ + 'ghcr.io/github/gh-aw-firewall/squid:latest', + 'ghcr.io/github/gh-aw-firewall/agent:latest', + ]); + }); + + it('should resolve agent-act image for act preset', () => { + const images = resolveImages({ ...defaults, agentImage: 'act' }); + expect(images).toEqual([ + 'ghcr.io/github/gh-aw-firewall/squid:latest', + 'ghcr.io/github/gh-aw-firewall/agent-act:latest', + ]); + }); + + it('should include api-proxy when enabled', () => { + const images = resolveImages({ ...defaults, enableApiProxy: true }); + expect(images).toEqual([ + 'ghcr.io/github/gh-aw-firewall/squid:latest', + 'ghcr.io/github/gh-aw-firewall/agent:latest', + 'ghcr.io/github/gh-aw-firewall/api-proxy:latest', + ]); + }); + + it('should use custom registry and tag', () => { + const images = resolveImages({ + ...defaults, + imageRegistry: 'my-registry.io/awf', + imageTag: 'v1.0.0', + }); + expect(images).toEqual([ + 'my-registry.io/awf/squid:v1.0.0', + 'my-registry.io/awf/agent:v1.0.0', + ]); + }); + + it('should use custom agent image as-is', () => { + const images = resolveImages({ ...defaults, agentImage: 'ubuntu:22.04' }); + expect(images).toEqual([ + 'ghcr.io/github/gh-aw-firewall/squid:latest', + 'ubuntu:22.04', + ]); + }); + + it('should reject custom image starting with dash', () => { + expect(() => resolveImages({ ...defaults, agentImage: '--help' })).toThrow( + 'must not start with "-"', + ); + }); + + it('should reject custom image containing whitespace', () => { + expect(() => resolveImages({ ...defaults, agentImage: 'ubuntu 22.04' })).toThrow( + 'must not contain whitespace', + ); + }); + }); + + describe('predownloadCommand', () => { + const defaults: PredownloadOptions = { + imageRegistry: 'ghcr.io/github/gh-aw-firewall', + imageTag: 'latest', + agentImage: 'default', + enableApiProxy: false, + }; + + beforeEach(() => { + execa.mockReset(); + execa.mockResolvedValue({ stdout: '', stderr: '' }); + }); + + it('should pull all resolved images', async () => { + await predownloadCommand(defaults); + + expect(execa).toHaveBeenCalledTimes(2); + expect(execa).toHaveBeenCalledWith( + 'docker', + ['pull', 'ghcr.io/github/gh-aw-firewall/squid:latest'], + { stdio: 'inherit' }, + ); + expect(execa).toHaveBeenCalledWith( + 'docker', + ['pull', 'ghcr.io/github/gh-aw-firewall/agent:latest'], + { stdio: 'inherit' }, + ); + }); + + it('should pull api-proxy when enabled', async () => { + await predownloadCommand({ ...defaults, enableApiProxy: true }); + + expect(execa).toHaveBeenCalledTimes(3); + expect(execa).toHaveBeenCalledWith( + 'docker', + ['pull', 'ghcr.io/github/gh-aw-firewall/api-proxy:latest'], + { stdio: 'inherit' }, + ); + }); + + it('should throw with exitCode 1 when a pull fails', async () => { + execa + .mockResolvedValueOnce({ stdout: '', stderr: '' }) + .mockRejectedValueOnce(new Error('pull failed')); + + try { + await predownloadCommand(defaults); + fail('Expected predownloadCommand to throw'); + } catch (error) { + expect((error as Error).message).toBe('1 of 2 image(s) failed to pull'); + expect((error as Error & { exitCode?: number }).exitCode).toBe(1); + } + }); + + it('should continue pulling remaining images after a failure', async () => { + execa.mockRejectedValueOnce(new Error('pull failed')).mockResolvedValueOnce({ stdout: '', stderr: '' }); + + await expect(predownloadCommand(defaults)).rejects.toThrow( + '1 of 2 image(s) failed to pull', + ); + + // Both images should have been attempted + expect(execa).toHaveBeenCalledTimes(2); + }); + + it('should handle non-Error rejection', async () => { + execa.mockRejectedValueOnce('string error').mockResolvedValueOnce({ stdout: '', stderr: '' }); + + try { + await predownloadCommand(defaults); + fail('Expected predownloadCommand to throw'); + } catch (error) { + expect((error as Error).message).toBe('1 of 2 image(s) failed to pull'); + expect((error as Error & { exitCode?: number }).exitCode).toBe(1); + } + }); + }); +}); diff --git a/src/commands/predownload.ts b/src/commands/predownload.ts new file mode 100644 index 00000000..d79ffa6c --- /dev/null +++ b/src/commands/predownload.ts @@ -0,0 +1,83 @@ +import execa from 'execa'; +import { logger } from '../logger'; + +export interface PredownloadOptions { + imageRegistry: string; + imageTag: string; + agentImage: string; + enableApiProxy: boolean; +} + +/** + * Validates a custom Docker image reference. + * Rejects values that could be interpreted as Docker CLI flags or contain whitespace. + */ +function validateImageReference(image: string): void { + if (image.startsWith('-')) { + throw new Error(`Invalid image reference "${image}": must not start with "-"`); + } + if (/\s/.test(image)) { + throw new Error(`Invalid image reference "${image}": must not contain whitespace`); + } +} + +/** + * Resolves the list of image references to pull based on the given options. + */ +export function resolveImages(options: PredownloadOptions): string[] { + const { imageRegistry, imageTag, agentImage, enableApiProxy } = options; + const images: string[] = []; + + // Always pull squid + images.push(`${imageRegistry}/squid:${imageTag}`); + + // Pull agent image based on preset + const isPreset = agentImage === 'default' || agentImage === 'act'; + if (isPreset) { + const imageName = agentImage === 'act' ? 'agent-act' : 'agent'; + images.push(`${imageRegistry}/${imageName}:${imageTag}`); + } else { + // Custom image - validate and pull as-is + validateImageReference(agentImage); + images.push(agentImage); + } + + // Optionally pull api-proxy + if (enableApiProxy) { + images.push(`${imageRegistry}/api-proxy:${imageTag}`); + } + + return images; +} + +/** + * Pre-download Docker images for offline use or faster startup. + */ +export async function predownloadCommand(options: PredownloadOptions): Promise { + const images = resolveImages(options); + + logger.info(`Pre-downloading ${images.length} image(s)...`); + + let failed = 0; + for (const image of images) { + logger.info(`Pulling ${image}...`); + try { + await execa('docker', ['pull', image], { stdio: 'inherit' }); + logger.info(`Successfully pulled ${image}`); + } catch (error) { + logger.error(`Failed to pull ${image}: ${error instanceof Error ? error.message : error}`); + failed++; + } + } + + if (failed > 0) { + const message = `${failed} of ${images.length} image(s) failed to pull`; + logger.error(message); + const error: Error & { exitCode?: number } = new Error(message); + error.exitCode = 1; + throw error; + } + + logger.info(`All ${images.length} image(s) pre-downloaded successfully`); + logger.info('You can now use --skip-pull to skip pulling images at runtime'); +}