diff --git a/src/cli/operations/agent/generate/schema-mapper.ts b/src/cli/operations/agent/generate/schema-mapper.ts index 06d420de0..322cca523 100644 --- a/src/cli/operations/agent/generate/schema-mapper.ts +++ b/src/cli/operations/agent/generate/schema-mapper.ts @@ -115,6 +115,7 @@ export function mapGenerateConfigToAgent(config: GenerateConfig): AgentEnvSpec { return { name: config.projectName, build: config.buildType ?? 'CodeZip', + ...(config.dockerfile && { dockerfile: config.dockerfile }), entrypoint: DEFAULT_PYTHON_ENTRYPOINT as FilePath, codeLocation: codeLocation as DirectoryPath, runtimeVersion: DEFAULT_PYTHON_VERSION, @@ -276,5 +277,6 @@ export async function mapGenerateConfigToRenderConfig( gatewayProviders, gatewayAuthTypes: [...new Set(gatewayProviders.map(g => g.authType))], protocol: config.protocol, + dockerfile: config.dockerfile, }; } diff --git a/src/cli/operations/deploy/__tests__/preflight-container.test.ts b/src/cli/operations/deploy/__tests__/preflight-container.test.ts index 5aeb0ce86..0bf4a7af7 100644 --- a/src/cli/operations/deploy/__tests__/preflight-container.test.ts +++ b/src/cli/operations/deploy/__tests__/preflight-container.test.ts @@ -9,6 +9,11 @@ vi.mock('node:fs', () => ({ vi.mock('../../../../lib', () => ({ DOCKERFILE_NAME: 'Dockerfile', + getDockerfilePath: (codeLocation: string, dockerfile?: string) => { + // eslint-disable-next-line @typescript-eslint/no-require-imports + const p = require('node:path') as typeof import('node:path'); + return p.join(codeLocation, dockerfile ?? 'Dockerfile'); + }, resolveCodeLocation: vi.fn((codeLocation: string, configBaseDir: string) => { // eslint-disable-next-line @typescript-eslint/no-require-imports const p = require('node:path') as typeof import('node:path'); @@ -96,4 +101,27 @@ describe('validateContainerAgents', () => { expect(() => validateContainerAgents(spec, CONFIG_ROOT)).toThrow(/agent-a.*agent-b/s); }); + + it('checks for custom dockerfile name when specified', () => { + mockedExistsSync.mockReturnValue(true); + + const spec = makeSpec([ + { name: 'gpu-agent', build: 'Container', codeLocation: dir('agents/gpu'), dockerfile: 'Dockerfile.gpu' }, + ]); + + expect(() => validateContainerAgents(spec, CONFIG_ROOT)).not.toThrow(); + // Should check for Dockerfile.gpu, not the default Dockerfile + const calledPath = mockedExistsSync.mock.calls[0]?.[0] as string; + expect(calledPath).toContain('Dockerfile.gpu'); + }); + + it('throws with custom dockerfile name in error message when missing', () => { + mockedExistsSync.mockReturnValue(false); + + const spec = makeSpec([ + { name: 'gpu-agent', build: 'Container', codeLocation: dir('agents/gpu'), dockerfile: 'Dockerfile.gpu' }, + ]); + + expect(() => validateContainerAgents(spec, CONFIG_ROOT)).toThrow(/Dockerfile\.gpu not found/); + }); }); diff --git a/src/cli/operations/deploy/preflight.ts b/src/cli/operations/deploy/preflight.ts index e427fdf53..4aa24e71b 100644 --- a/src/cli/operations/deploy/preflight.ts +++ b/src/cli/operations/deploy/preflight.ts @@ -1,4 +1,4 @@ -import { ConfigIO, DOCKERFILE_NAME, requireConfigRoot, resolveCodeLocation } from '../../../lib'; +import { ConfigIO, DOCKERFILE_NAME, getDockerfilePath, requireConfigRoot, resolveCodeLocation } from '../../../lib'; import type { AgentCoreProjectSpec, AwsDeploymentTarget } from '../../../schema'; import { validateAwsCredentials } from '../../aws/account'; import { LocalCdkProject } from '../../cdk/local-cdk-project'; @@ -147,11 +147,11 @@ export function validateContainerAgents(projectSpec: AgentCoreProjectSpec, confi for (const agent of projectSpec.runtimes || []) { if (agent.build === 'Container') { const codeLocation = resolveCodeLocation(agent.codeLocation, configRoot); - const dockerfilePath = path.join(codeLocation, DOCKERFILE_NAME); + const dockerfilePath = getDockerfilePath(codeLocation, agent.dockerfile); if (!existsSync(dockerfilePath)) { errors.push( - `Agent "${agent.name}": Dockerfile not found at ${dockerfilePath}. Container agents require a Dockerfile.` + `Agent "${agent.name}": ${agent.dockerfile ?? DOCKERFILE_NAME} not found at ${dockerfilePath}. Container agents require a Dockerfile.` ); } } diff --git a/src/cli/operations/dev/__tests__/config.test.ts b/src/cli/operations/dev/__tests__/config.test.ts index 56b97ef64..b6967ac6e 100644 --- a/src/cli/operations/dev/__tests__/config.test.ts +++ b/src/cli/operations/dev/__tests__/config.test.ts @@ -367,6 +367,35 @@ describe('getDevConfig', () => { expect(config).not.toBeNull(); expect(config!.isPython).toBe(true); }); + + it('threads dockerfile from Container agent spec to DevConfig', () => { + const project: AgentCoreProjectSpec = { + name: 'TestProject', + version: 1, + managedBy: 'CDK' as const, + runtimes: [ + { + name: 'ContainerAgent', + build: 'Container', + runtimeVersion: 'PYTHON_3_12', + entrypoint: filePath('main.py'), + codeLocation: dirPath('./agents/container'), + protocol: 'HTTP', + dockerfile: 'Dockerfile.gpu', + }, + ], + memories: [], + credentials: [], + evaluators: [], + onlineEvalConfigs: [], + agentCoreGateways: [], + policyEngines: [], + }; + + const config = getDevConfig(workingDir, project, '/test/project/agentcore'); + expect(config).not.toBeNull(); + expect(config?.dockerfile).toBe('Dockerfile.gpu'); + }); }); describe('getAgentPort', () => { diff --git a/src/cli/operations/dev/config.ts b/src/cli/operations/dev/config.ts index 36b174748..fd13637bb 100644 --- a/src/cli/operations/dev/config.ts +++ b/src/cli/operations/dev/config.ts @@ -10,6 +10,7 @@ export interface DevConfig { isPython: boolean; buildType: BuildType; protocol: ProtocolMode; + dockerfile?: string; } interface DevSupportResult { @@ -140,6 +141,7 @@ export function getDevConfig( isPython: isPythonAgent(targetAgent), buildType: targetAgent.build, protocol: targetAgent.protocol ?? 'HTTP', + dockerfile: targetAgent.dockerfile, }; } diff --git a/src/cli/operations/dev/container-dev-server.ts b/src/cli/operations/dev/container-dev-server.ts index 3690e8ee8..1159f7bbe 100644 --- a/src/cli/operations/dev/container-dev-server.ts +++ b/src/cli/operations/dev/container-dev-server.ts @@ -1,4 +1,4 @@ -import { CONTAINER_INTERNAL_PORT, DOCKERFILE_NAME } from '../../../lib'; +import { CONTAINER_INTERNAL_PORT, DOCKERFILE_NAME, getDockerfilePath } from '../../../lib'; import { getUvBuildArgs } from '../../../lib/packaging/build-args'; import { detectContainerRuntime, getStartHint } from '../../external-requirements/detect'; import { DevServer, type LogLevel, type SpawnConfig } from './dev-server'; @@ -73,9 +73,10 @@ export class ContainerDevServer extends DevServer { this.runtimeBinary = runtime.binary; // 2. Verify Dockerfile exists - const dockerfilePath = join(this.config.directory, DOCKERFILE_NAME); + const dockerfileName = this.config.dockerfile ?? DOCKERFILE_NAME; + const dockerfilePath = getDockerfilePath(this.config.directory, this.config.dockerfile); if (!existsSync(dockerfilePath)) { - onLog('error', `Dockerfile not found at ${dockerfilePath}. Container agents require a Dockerfile.`); + onLog('error', `${dockerfileName} not found at ${dockerfilePath}. Container agents require a Dockerfile.`); return false; } diff --git a/src/cli/templates/BaseRenderer.ts b/src/cli/templates/BaseRenderer.ts index 7de56b01a..325aeb392 100644 --- a/src/cli/templates/BaseRenderer.ts +++ b/src/cli/templates/BaseRenderer.ts @@ -71,7 +71,8 @@ export abstract class BaseRenderer { const containerTemplateDir = path.join(this.baseTemplateDir, 'container', language); if (existsSync(containerTemplateDir)) { - await copyAndRenderDir(containerTemplateDir, projectDir, { ...templateData, entrypoint: 'main' }); + const exclude = this.config.dockerfile ? new Set(['Dockerfile']) : undefined; + await copyAndRenderDir(containerTemplateDir, projectDir, { ...templateData, entrypoint: 'main' }, { exclude }); } } } diff --git a/src/cli/templates/render.ts b/src/cli/templates/render.ts index e56a84907..55c228555 100644 --- a/src/cli/templates/render.ts +++ b/src/cli/templates/render.ts @@ -47,13 +47,19 @@ export async function copyDir(src: string, dest: string): Promise { /** * Recursively copies a directory, rendering Handlebars templates. */ -export async function copyAndRenderDir(src: string, dest: string, data: T): Promise { +export async function copyAndRenderDir( + src: string, + dest: string, + data: T, + options?: { exclude?: Set } +): Promise { await fs.mkdir(dest, { recursive: true }); const entries = await fs.readdir(src, { withFileTypes: true }); for (const entry of entries) { - const srcPath = path.join(src, entry.name); const destName = resolveTemplateName(entry.name); + if (options?.exclude?.has(destName)) continue; + const srcPath = path.join(src, entry.name); const destPath = path.join(dest, destName); if (entry.isDirectory()) { diff --git a/src/cli/templates/types.ts b/src/cli/templates/types.ts index 337cacf4a..094af9782 100644 --- a/src/cli/templates/types.ts +++ b/src/cli/templates/types.ts @@ -66,4 +66,6 @@ export interface AgentRenderConfig { gatewayAuthTypes: string[]; /** Protocol (HTTP, MCP, A2A). Defaults to HTTP. */ protocol?: ProtocolMode; + /** Custom Dockerfile name — when set, the template Dockerfile is not scaffolded */ + dockerfile?: string; } diff --git a/src/cli/tui/components/PathInput.tsx b/src/cli/tui/components/PathInput.tsx index 9db300d3d..ebb16f956 100644 --- a/src/cli/tui/components/PathInput.tsx +++ b/src/cli/tui/components/PathInput.tsx @@ -19,6 +19,10 @@ interface PathInputProps { allowCreate?: boolean; /** Show hidden files (dotfiles) in completions (default: false) */ showHidden?: boolean; + /** Allow empty input (user presses Enter without selecting a file) */ + allowEmpty?: boolean; + /** Message shown when user submits empty input (only if allowEmpty is true) */ + emptyHelpText?: string; } interface CompletionItem { @@ -133,6 +137,8 @@ export function PathInput({ maxVisibleItems = 8, allowCreate = false, showHidden = false, + allowEmpty = false, + emptyHelpText, }: PathInputProps) { const [value, setValue] = useState(initialValue); const [cursor, setCursor] = useState(initialValue.length); @@ -207,6 +213,10 @@ export function PathInput({ if (key.return) { const trimmed = value.trim(); if (!trimmed) { + if (allowEmpty) { + onSubmit(''); + return; + } setError('Please enter a path'); return; } @@ -322,8 +332,9 @@ export function PathInput({ )} {/* Help text */} - + ↑↓ move → open ← back Enter submit Esc cancel + {allowEmpty && emptyHelpText && !value && {emptyHelpText}} ); diff --git a/src/cli/tui/screens/agent/AddAgentScreen.tsx b/src/cli/tui/screens/agent/AddAgentScreen.tsx index 2b40dd28f..0e54613c9 100644 --- a/src/cli/tui/screens/agent/AddAgentScreen.tsx +++ b/src/cli/tui/screens/agent/AddAgentScreen.tsx @@ -17,19 +17,22 @@ import { ConfirmReview, Cursor, Panel, + PathInput, Screen, StepIndicator, TextInput, + WizardMultiSelect, WizardSelect, } from '../../components'; import type { SelectableItem } from '../../components'; import { JwtConfigInput, useJwtConfigFlow } from '../../components/jwt-config'; import { HELP_TEXT } from '../../constants'; -import { useListNavigation, useProject } from '../../hooks'; +import { useListNavigation, useMultiSelectNavigation, useProject } from '../../hooks'; import { generateUniqueName } from '../../utils'; import { BUILD_TYPE_OPTIONS, GenerateWizardUI, getWizardHelpText, useGenerateWizard } from '../generate'; import type { BuildType, MemoryOption } from '../generate'; -import { ADVANCED_OPTIONS, MEMORY_OPTIONS } from '../generate/types'; +import type { AdvancedSettingId } from '../generate/types'; +import { ADVANCED_SETTING_OPTIONS, MEMORY_OPTIONS } from '../generate/types'; import type { AddAgentConfig, AddAgentStep, AgentType } from './types'; import { ADD_AGENT_STEP_LABELS, @@ -42,6 +45,7 @@ import { } from './types'; import { Box, Text, useInput } from 'ink'; import Spinner from 'ink-spinner'; +import { basename, resolve } from 'path'; import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'; // Helper to get provider display name and env var name from ModelProvider @@ -70,6 +74,7 @@ type InitialStep = 'name' | 'agentType'; type ByoStep = | 'codeLocation' | 'buildType' + | 'dockerfile' | 'modelProvider' | 'apiKey' | 'advanced' @@ -84,12 +89,52 @@ type ByoStep = | 'confirm'; const INITIAL_STEPS: InitialStep[] = ['name', 'agentType']; -const ADVANCED_ITEMS: SelectableItem[] = ADVANCED_OPTIONS.map(o => ({ - id: o.id, - title: o.title, - description: o.description, -})); -const BYO_STEPS: ByoStep[] = ['codeLocation', 'buildType', 'modelProvider', 'apiKey', 'advanced', 'confirm']; +const BYO_BASE_STEPS: ByoStep[] = ['codeLocation', 'buildType', 'modelProvider', 'apiKey', 'advanced', 'confirm']; + +export interface ComputeByoStepsInput { + modelProvider: string; + buildType: string; + networkMode: string; + authorizerType: string; + advancedSettings: Set; +} + +/** Pure function to compute BYO wizard steps from config. Exported for testing. */ +export function computeByoSteps(input: ComputeByoStepsInput): ByoStep[] { + let steps = [...BYO_BASE_STEPS]; + if (input.modelProvider === 'Bedrock') { + steps = steps.filter(s => s !== 'apiKey'); + } + if (input.advancedSettings.size > 0) { + const advancedIndex = steps.indexOf('advanced'); + const afterAdvanced = advancedIndex + 1; + const subSteps: ByoStep[] = []; + if (input.advancedSettings.has('dockerfile') && input.buildType === 'Container') { + subSteps.push('dockerfile'); + } + if (input.advancedSettings.has('network')) { + subSteps.push('networkMode'); + if (input.networkMode === 'VPC') { + subSteps.push('subnets', 'securityGroups'); + } + } + if (input.advancedSettings.has('headers')) { + subSteps.push('requestHeaderAllowlist'); + } + if (input.advancedSettings.has('auth')) { + subSteps.push('authorizerType'); + } + if (input.advancedSettings.has('lifecycle')) { + subSteps.push('idleTimeout', 'maxLifetime'); + } + steps = [...steps.slice(0, afterAdvanced), ...subSteps, ...steps.slice(afterAdvanced)]; + } + if (input.authorizerType === 'CUSTOM_JWT' && steps.includes('authorizerType')) { + const authIndex = steps.indexOf('authorizerType'); + steps = [...steps.slice(0, authIndex + 1), 'jwtConfig', ...steps.slice(authIndex + 1)]; + } + return steps; +} type ImportStep = | 'region' @@ -126,6 +171,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg codeLocation: '', entrypoint: DEFAULT_ENTRYPOINT, buildType: 'CodeZip' as BuildType, + dockerfile: '' as string, modelProvider: 'Bedrock' as ModelProvider, apiKey: undefined as string | undefined, networkMode: 'PUBLIC' as NetworkMode, @@ -135,7 +181,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg idleTimeout: '' as string, maxLifetime: '' as string, }); - const [byoAdvancedSelected, setByoAdvancedSelected] = useState(false); + const [byoAdvancedSettings, setByoAdvancedSettings] = useState>(new Set()); const [byoAuthorizerType, setByoAuthorizerType] = useState('AWS_IAM'); const [byoJwtConfig, setByoJwtConfig] = useState(undefined); @@ -237,6 +283,10 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg entrypoint: 'main.py', language: generateWizard.config.language, buildType: generateWizard.config.buildType, + ...(generateWizard.config.buildType === 'Container' && + generateWizard.config.dockerfile && { + dockerfile: generateWizard.config.dockerfile, + }), protocol: generateWizard.config.protocol, framework: generateWizard.config.sdk, modelProvider: generateWizard.config.modelProvider, @@ -275,37 +325,50 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg // BYO Path // ───────────────────────────────────────────────────────────────────────────── - // BYO steps filtering (remove apiKey for Bedrock, subnets/securityGroups for non-VPC, jwtConfig for non-CUSTOM_JWT) - const byoSteps = useMemo(() => { - let steps = [...BYO_STEPS]; - if (byoConfig.modelProvider === 'Bedrock') { - steps = steps.filter(s => s !== 'apiKey'); - } - if (byoAdvancedSelected) { - const advancedIndex = steps.indexOf('advanced'); - const afterAdvanced = advancedIndex + 1; - const networkSteps: ByoStep[] = - byoConfig.networkMode === 'VPC' ? ['networkMode', 'subnets', 'securityGroups'] : ['networkMode']; - steps = [ - ...steps.slice(0, afterAdvanced), - ...networkSteps, - 'requestHeaderAllowlist', - 'authorizerType', - 'idleTimeout', - 'maxLifetime', - ...steps.slice(afterAdvanced), - ]; - } - // Add jwtConfig step after authorizerType when CUSTOM_JWT is selected - if (byoAuthorizerType === 'CUSTOM_JWT') { - const authIndex = steps.indexOf('authorizerType'); - steps = [...steps.slice(0, authIndex + 1), 'jwtConfig', ...steps.slice(authIndex + 1)]; - } - return steps; - }, [byoConfig.modelProvider, byoConfig.networkMode, byoAdvancedSelected, byoAuthorizerType]); + // BYO steps filtering (apiKey for Bedrock, advanced sub-steps based on multi-select, jwtConfig for CUSTOM_JWT) + const byoAdvancedActive = byoAdvancedSettings.size > 0; + const byoSteps = useMemo( + () => + computeByoSteps({ + modelProvider: byoConfig.modelProvider, + buildType: byoConfig.buildType, + networkMode: byoConfig.networkMode, + authorizerType: byoAuthorizerType, + advancedSettings: byoAdvancedSettings, + }), + [ + byoConfig.buildType, + byoConfig.modelProvider, + byoConfig.networkMode, + byoAdvancedActive, + byoAdvancedSettings, + byoAuthorizerType, + ] + ); const byoCurrentIndex = byoSteps.indexOf(byoStep); + /** Navigate to the next step after the given step in the BYO steps array */ + const goToNextByoStep = useCallback( + (afterStep: ByoStep) => { + const idx = byoSteps.indexOf(afterStep); + const next = idx >= 0 ? byoSteps[idx + 1] : undefined; + setByoStep(next ?? 'confirm'); + }, + [byoSteps] + ); + + // Advanced multi-select items — filter out dockerfile when not a Container build + const byoAdvancedItems: SelectableItem[] = useMemo( + () => + ADVANCED_SETTING_OPTIONS.filter(o => o.id !== 'dockerfile' || byoConfig.buildType === 'Container').map(o => ({ + id: o.id, + title: o.title, + description: o.description, + })), + [byoConfig.buildType] + ); + // BYO build type options const buildTypeItems: SelectableItem[] = useMemo( () => @@ -350,6 +413,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg entrypoint: byoConfig.entrypoint, language: 'Python', // Default - not used for BYO agents buildType: byoConfig.buildType, + ...(byoConfig.buildType === 'Container' && byoConfig.dockerfile && { dockerfile: byoConfig.dockerfile }), protocol: 'HTTP', // Default for BYO agents framework: 'Strands', // Default - not used for BYO agents modelProvider: byoConfig.modelProvider, @@ -371,7 +435,8 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg const buildTypeNav = useListNavigation({ items: buildTypeItems, onSelect: item => { - setByoConfig(c => ({ ...c, buildType: item.id as BuildType })); + const build = item.id as BuildType; + setByoConfig(c => ({ ...c, buildType: build, dockerfile: '' })); setByoStep('modelProvider'); }, onExit: handleByoBack, @@ -404,27 +469,49 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg [] ); - const advancedNav = useListNavigation({ - items: ADVANCED_ITEMS, - onSelect: item => { - if (item.id === 'yes') { - setByoAdvancedSelected(true); - setByoStep('networkMode'); - } else { - setByoAdvancedSelected(false); + const advancedNav = useMultiSelectNavigation({ + items: byoAdvancedItems, + getId: item => item.id, + onConfirm: selectedIds => { + const selected = new Set(selectedIds as AdvancedSettingId[]); + setByoAdvancedSettings(selected); + if (selected.size === 0) { + // No advanced settings — reset defaults and go to confirm setByoConfig(c => ({ ...c, + dockerfile: '', networkMode: 'PUBLIC' as NetworkMode, subnets: '', securityGroups: '', + requestHeaderAllowlist: '', idleTimeout: '', maxLifetime: '', })); + setByoAuthorizerType('AWS_IAM'); + setByoJwtConfig(undefined); setByoStep('confirm'); + } else { + // Navigate to first advanced sub-step (steps memo hasn't updated yet) + setTimeout(() => { + if (selected.has('dockerfile') && byoConfig.buildType === 'Container') { + setByoStep('dockerfile'); + } else if (selected.has('network')) { + setByoStep('networkMode'); + } else if (selected.has('headers')) { + setByoStep('requestHeaderAllowlist'); + } else if (selected.has('auth')) { + setByoStep('authorizerType'); + } else if (selected.has('lifecycle')) { + setByoStep('idleTimeout'); + } else { + setByoStep('confirm'); + } + }, 0); } }, onExit: handleByoBack, isActive: isByoPath && byoStep === 'advanced', + requireSelection: false, }); const networkModeNav = useListNavigation({ @@ -435,7 +522,8 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg if (mode === 'VPC') { setByoStep('subnets'); } else { - setByoStep('requestHeaderAllowlist'); + // Skip subnets/securityGroups — go to next step after networkMode + setTimeout(() => goToNextByoStep('networkMode'), 0); } }, onExit: handleByoBack, @@ -457,7 +545,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg setByoStep('jwtConfig'); } else { setByoJwtConfig(undefined); - setByoStep('confirm'); + setTimeout(() => goToNextByoStep('authorizerType'), 0); } }, onExit: handleByoBack, @@ -468,7 +556,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg const byoJwtFlow = useJwtConfigFlow({ onComplete: jwtConfig => { setByoJwtConfig(jwtConfig); - setByoStep('confirm'); + setTimeout(() => goToNextByoStep('jwtConfig'), 0); }, onBack: () => { setByoStep('authorizerType'); @@ -716,6 +804,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg } if ( byoStep === 'codeLocation' || + byoStep === 'dockerfile' || byoStep === 'apiKey' || byoStep === 'subnets' || byoStep === 'securityGroups' || @@ -725,6 +814,9 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg ) { return HELP_TEXT.TEXT_INPUT; } + if (byoStep === 'advanced') { + return 'Space toggle · Enter confirm · Esc back'; + } if (byoStep === 'confirm') { return HELP_TEXT.CONFIRM_CANCEL; } @@ -974,6 +1066,21 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg )} + {byoStep === 'dockerfile' && ( + { + setByoConfig(c => ({ ...c, dockerfile: value ? basename(value) : '' })); + goToNextByoStep('dockerfile'); + }} + onCancel={handleByoBack} + /> + )} + {byoStep === 'modelProvider' && ( )} @@ -1031,7 +1140,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg customValidation={validateSecurityGroupIds} onSubmit={value => { setByoConfig(c => ({ ...c, securityGroups: value })); - setByoStep('requestHeaderAllowlist'); + goToNextByoStep('securityGroups'); }} onCancel={handleByoBack} /> @@ -1042,13 +1151,14 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg { const result = validateHeaderAllowlist(value); return result.success ? true : result.error!; }} onSubmit={value => { setByoConfig(c => ({ ...c, requestHeaderAllowlist: value })); - setByoStep('authorizerType'); + goToNextByoStep('requestHeaderAllowlist'); }} onCancel={handleByoBack} /> @@ -1098,6 +1208,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg { if (!value) return true; const n = Number(value); @@ -1117,6 +1228,7 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg { if (!value) return true; const n = Number(value); @@ -1147,6 +1259,9 @@ export function AddAgentScreen({ existingAgentNames, onComplete, onExit }: AddAg label: 'Build', value: BUILD_TYPE_OPTIONS.find(o => o.id === byoConfig.buildType)?.title ?? byoConfig.buildType, }, + ...(byoConfig.buildType === 'Container' && byoConfig.dockerfile + ? [{ label: 'Dockerfile', value: byoConfig.dockerfile }] + : []), { label: 'Model Provider', value: `${byoConfig.modelProvider} (${DEFAULT_MODEL_IDS[byoConfig.modelProvider]})`, diff --git a/src/cli/tui/screens/agent/__tests__/computeByoSteps.test.ts b/src/cli/tui/screens/agent/__tests__/computeByoSteps.test.ts new file mode 100644 index 000000000..2d706e1d0 --- /dev/null +++ b/src/cli/tui/screens/agent/__tests__/computeByoSteps.test.ts @@ -0,0 +1,62 @@ +import type { AdvancedSettingId } from '../../generate/types'; +import { computeByoSteps } from '../AddAgentScreen'; +import type { ComputeByoStepsInput } from '../AddAgentScreen'; +import { describe, expect, it } from 'vitest'; + +function makeInput(overrides: Partial = {}): ComputeByoStepsInput { + return { + modelProvider: 'Bedrock', + buildType: 'CodeZip', + networkMode: 'PUBLIC', + authorizerType: 'AWS_IAM', + advancedSettings: new Set(), + ...overrides, + }; +} + +describe('computeByoSteps - dockerfile', () => { + it('Container build with dockerfile selected includes dockerfile step', () => { + const steps = computeByoSteps( + makeInput({ + buildType: 'Container', + advancedSettings: new Set(['dockerfile']), + }) + ); + expect(steps).toContain('dockerfile'); + const advIdx = steps.indexOf('advanced'); + expect(steps[advIdx + 1]).toBe('dockerfile'); + }); + + it('CodeZip build with dockerfile selected does NOT include dockerfile step', () => { + const steps = computeByoSteps( + makeInput({ + buildType: 'CodeZip', + advancedSettings: new Set(['dockerfile']), + }) + ); + expect(steps).not.toContain('dockerfile'); + }); + + it('dockerfile-only selection on Container has steps: advanced, dockerfile, confirm', () => { + const steps = computeByoSteps( + makeInput({ + buildType: 'Container', + advancedSettings: new Set(['dockerfile']), + }) + ); + const advIdx = steps.indexOf('advanced'); + expect(steps.slice(advIdx)).toEqual(['advanced', 'dockerfile', 'confirm']); + }); + + it('dockerfile + lifecycle on Container includes both groups', () => { + const steps = computeByoSteps( + makeInput({ + buildType: 'Container', + advancedSettings: new Set(['dockerfile', 'lifecycle']), + }) + ); + const advIdx = steps.indexOf('advanced'); + expect(steps.slice(advIdx)).toEqual(['advanced', 'dockerfile', 'idleTimeout', 'maxLifetime', 'confirm']); + expect(steps).not.toContain('networkMode'); + }); +}); diff --git a/src/cli/tui/screens/agent/types.ts b/src/cli/tui/screens/agent/types.ts index bcaaa9d3d..d8dc6c899 100644 --- a/src/cli/tui/screens/agent/types.ts +++ b/src/cli/tui/screens/agent/types.ts @@ -40,6 +40,7 @@ export type AddAgentStep = | 'agentType' | 'codeLocation' | 'buildType' + | 'dockerfile' | 'language' | 'protocol' | 'framework' @@ -69,6 +70,8 @@ export interface AddAgentConfig { entrypoint: string; language: TargetLanguage; buildType: BuildType; + /** Path to custom Dockerfile (copied into code directory at setup) or filename already in code directory. */ + dockerfile?: string; /** Protocol (HTTP, MCP, A2A). Defaults to HTTP. */ protocol: ProtocolMode; framework: SDKFramework; @@ -108,6 +111,7 @@ export const ADD_AGENT_STEP_LABELS: Record = { agentType: 'Type', codeLocation: 'Code', buildType: 'Build', + dockerfile: 'Dockerfile', language: 'Language', protocol: 'Protocol', framework: 'Framework', diff --git a/src/cli/tui/screens/agent/useAddAgent.ts b/src/cli/tui/screens/agent/useAddAgent.ts index 574953ae4..63407d2f1 100644 --- a/src/cli/tui/screens/agent/useAddAgent.ts +++ b/src/cli/tui/screens/agent/useAddAgent.ts @@ -15,8 +15,8 @@ import { credentialPrimitive } from '../../../primitives/registry'; import { createRenderer } from '../../../templates'; import type { GenerateConfig } from '../generate/types'; import type { AddAgentConfig } from './types'; -import { mkdirSync } from 'fs'; -import { dirname, join } from 'path'; +import { copyFileSync, existsSync, mkdirSync } from 'fs'; +import { basename, dirname, join, resolve } from 'path'; import { useCallback, useState } from 'react'; // ───────────────────────────────────────────────────────────────────────────── @@ -58,6 +58,7 @@ export function mapByoConfigToAgent(config: AddAgentConfig): AgentEnvSpec { return { name: config.name, build: config.buildType, + ...(config.dockerfile && { dockerfile: config.dockerfile }), entrypoint: config.entrypoint as FilePath, codeLocation: config.codeLocation as DirectoryPath, runtimeVersion: config.pythonVersion, @@ -99,6 +100,7 @@ function mapAddAgentConfigToGenerateConfig(config: AddAgentConfig): GenerateConf return { projectName: config.name, // In create context, this is the agent name buildType: config.buildType, + ...(config.dockerfile && { dockerfile: config.dockerfile }), protocol: config.protocol, sdk: config.framework, modelProvider: config.modelProvider, @@ -212,6 +214,17 @@ async function handleCreatePath( const renderer = createRenderer(renderConfig); await renderer.render({ outputDir: projectRoot }); + // If dockerfile is a path (contains /), copy it into the agent directory (overwriting template default) + if (generateConfig.dockerfile?.includes('/')) { + const sourcePath = resolve(projectRoot, generateConfig.dockerfile); + if (!existsSync(sourcePath)) { + return { ok: false, error: `Dockerfile not found at ${sourcePath}` }; + } + const filename = basename(sourcePath); + copyFileSync(sourcePath, join(agentPath, filename)); + generateConfig.dockerfile = filename; + } + // Write agent to project config if (strategy) { await writeAgentToProject(generateConfig, { configBaseDir, credentialStrategy: strategy }); @@ -304,8 +317,19 @@ async function handleByoPath( const codeDir = join(projectRoot, config.codeLocation.replace(/\/$/, '')); mkdirSync(codeDir, { recursive: true }); + // If dockerfile is a path (contains /), copy it into the code directory and use the filename + let dockerfileName = config.dockerfile; + if (dockerfileName?.includes('/')) { + const sourcePath = resolve(projectRoot, dockerfileName); + if (!existsSync(sourcePath)) { + return { ok: false, error: `Dockerfile not found at ${sourcePath}` }; + } + dockerfileName = basename(sourcePath); + copyFileSync(sourcePath, join(codeDir, dockerfileName)); + } + const project = await configIO.readProjectSpec(); - const agent = mapByoConfigToAgent(config); + const agent = mapByoConfigToAgent({ ...config, dockerfile: dockerfileName }); // Append new agent project.runtimes.push(agent); diff --git a/src/cli/tui/screens/create/useCreateFlow.ts b/src/cli/tui/screens/create/useCreateFlow.ts index 8f6b43430..91610b47a 100644 --- a/src/cli/tui/screens/create/useCreateFlow.ts +++ b/src/cli/tui/screens/create/useCreateFlow.ts @@ -260,6 +260,7 @@ export function useCreateFlow(cwd: string): CreateFlowState { const generateConfig: GenerateConfig = { projectName: addAgentConfig.name, buildType: addAgentConfig.buildType, + ...(addAgentConfig.dockerfile && { dockerfile: addAgentConfig.dockerfile }), protocol: addAgentConfig.protocol, sdk: addAgentConfig.framework, modelProvider: addAgentConfig.modelProvider, @@ -272,6 +273,8 @@ export function useCreateFlow(cwd: string): CreateFlowState { requestHeaderAllowlist: addAgentConfig.requestHeaderAllowlist, authorizerType: addAgentConfig.authorizerType, jwtConfig: addAgentConfig.jwtConfig, + idleRuntimeSessionTimeout: addAgentConfig.idleRuntimeSessionTimeout, + maxLifetime: addAgentConfig.maxLifetime, }; logger.logSubStep(`Framework: ${generateConfig.sdk}`); diff --git a/src/cli/tui/screens/generate/GenerateWizardUI.tsx b/src/cli/tui/screens/generate/GenerateWizardUI.tsx index c1f2c0866..b4f25f660 100644 --- a/src/cli/tui/screens/generate/GenerateWizardUI.tsx +++ b/src/cli/tui/screens/generate/GenerateWizardUI.tsx @@ -3,14 +3,22 @@ import { DEFAULT_MODEL_IDS, LIFECYCLE_TIMEOUT_MAX, LIFECYCLE_TIMEOUT_MIN, Projec import { parseAndNormalizeHeaders, validateHeaderAllowlist } from '../../../commands/shared/header-utils'; import { validateSecurityGroupIds, validateSubnetIds } from '../../../commands/shared/vpc-utils'; import { computeDefaultCredentialEnvVarName } from '../../../primitives/credential-utils'; -import { ApiKeySecretInput, Panel, SelectList, StepIndicator, TextInput } from '../../components'; +import { + ApiKeySecretInput, + Panel, + PathInput, + SelectList, + StepIndicator, + TextInput, + WizardMultiSelect, +} from '../../components'; import type { SelectableItem } from '../../components'; import { JwtConfigInput, useJwtConfigFlow } from '../../components/jwt-config'; -import { useListNavigation } from '../../hooks'; +import { useListNavigation, useMultiSelectNavigation } from '../../hooks'; import { RUNTIME_AUTHORIZER_TYPE_OPTIONS } from '../agent/types'; -import type { BuildType, GenerateConfig, GenerateStep, MemoryOption, ProtocolMode } from './types'; +import type { AdvancedSettingId, BuildType, GenerateConfig, GenerateStep, MemoryOption, ProtocolMode } from './types'; import { - ADVANCED_OPTIONS, + ADVANCED_SETTING_OPTIONS, BUILD_TYPE_OPTIONS, LANGUAGE_OPTIONS, MEMORY_OPTIONS, @@ -22,6 +30,7 @@ import { } from './types'; import type { useGenerateWizard } from './useGenerateWizard'; import { Box, Text, useInput } from 'ink'; +import { basename } from 'path'; // Helper to get provider display name and env var name from ModelProvider function getProviderInfo(provider: ModelProvider): { name: string; envVarName: string } { @@ -83,8 +92,6 @@ export function GenerateWizardUI({ })); case 'memory': return MEMORY_OPTIONS.map(o => ({ id: o.id, title: o.title, description: o.description })); - case 'advanced': - return ADVANCED_OPTIONS.map(o => ({ id: o.id, title: o.title, description: o.description })); case 'networkMode': return NETWORK_MODE_OPTIONS.map(o => ({ id: o.id, title: o.title, description: o.description })); case 'authorizerType': @@ -96,7 +103,9 @@ export function GenerateWizardUI({ const items = getItems(); const isSelectStep = items.length > 0; + const isAdvancedStep = wizard.step === 'advanced'; const isTextStep = wizard.step === 'projectName'; + const isDockerfileStep = wizard.step === 'dockerfile'; const isApiKeyStep = wizard.step === 'apiKey'; const isSubnetsStep = wizard.step === 'subnets'; const isSecurityGroupsStep = wizard.step === 'securityGroups'; @@ -106,6 +115,11 @@ export function GenerateWizardUI({ const isMaxLifetimeStep = wizard.step === 'maxLifetime'; const isConfirmStep = wizard.step === 'confirm'; + // Advanced multi-select items — filter out dockerfile when not a Container build + const advancedItems: SelectableItem[] = ADVANCED_SETTING_OPTIONS.filter( + o => o.id !== 'dockerfile' || wizard.config.buildType === 'Container' + ).map(o => ({ id: o.id, title: o.title, description: o.description })); + const handleSelect = (item: SelectableItem) => { switch (wizard.step) { case 'language': @@ -126,9 +140,6 @@ export function GenerateWizardUI({ case 'memory': wizard.setMemory(item.id as MemoryOption); break; - case 'advanced': - wizard.setAdvanced(item.id === 'yes'); - break; case 'networkMode': wizard.setNetworkMode(item.id as NetworkMode); break; @@ -142,11 +153,20 @@ export function GenerateWizardUI({ items, onSelect: handleSelect, onExit: onBack, - isActive: isActive && isSelectStep, + isActive: isActive && isSelectStep && !isAdvancedStep, isDisabled: item => item.disabled ?? false, resetKey: wizard.step, }); + const advancedNav = useMultiSelectNavigation({ + items: advancedItems, + getId: item => item.id, + onConfirm: selectedIds => wizard.setAdvanced(selectedIds as AdvancedSettingId[]), + onExit: onBack, + isActive: isActive && isAdvancedStep, + requireSelection: false, + }); + // JWT config flow for CUSTOM_JWT authorizer const jwtFlow = useJwtConfigFlow({ onComplete: jwtConfig => { @@ -188,7 +208,30 @@ export function GenerateWizardUI({ )} - {isSelectStep && } + {isSelectStep && !isAdvancedStep && } + + {isAdvancedStep && ( + + )} + + {isDockerfileStep && ( + { + wizard.setDockerfile(value ? basename(value) : undefined); + }} + onCancel={onBack} + /> + )} {isApiKeyStep && ( { const result = validateHeaderAllowlist(value); return result.success ? true : result.error!; @@ -291,6 +335,7 @@ export function GenerateWizardUI({ { if (!value) return true; const n = Number(value); @@ -313,6 +358,7 @@ export function GenerateWizardUI({ { if (!value) return true; const n = Number(value); @@ -347,6 +393,7 @@ export function getWizardHelpText(step: GenerateStep): string { if (step === 'confirm') return 'Enter/Y confirm · Esc back'; if ( step === 'projectName' || + step === 'dockerfile' || step === 'subnets' || step === 'securityGroups' || step === 'requestHeaderAllowlist' || @@ -356,6 +403,7 @@ export function getWizardHelpText(step: GenerateStep): string { return 'Enter submit · Esc cancel'; if (step === 'apiKey') return 'Enter submit · Tab show/hide · Esc back'; if (step === 'jwtConfig') return 'Enter submit · Esc back'; + if (step === 'advanced') return 'Space toggle · Enter confirm · Esc back'; return '↑↓ navigate · Enter select · Esc back'; } @@ -405,6 +453,12 @@ function ConfirmView({ config, credentialProjectName }: { config: GenerateConfig Build: {buildTypeLabel} + {config.buildType === 'Container' && config.dockerfile && ( + + Dockerfile: + {config.dockerfile} + + )} Protocol: {protocolLabel} diff --git a/src/cli/tui/screens/generate/__tests__/useGenerateWizard.test.tsx b/src/cli/tui/screens/generate/__tests__/useGenerateWizard.test.tsx index cbaeb008a..d303be1d6 100644 --- a/src/cli/tui/screens/generate/__tests__/useGenerateWizard.test.tsx +++ b/src/cli/tui/screens/generate/__tests__/useGenerateWizard.test.tsx @@ -1,8 +1,9 @@ +import { validateDockerfileInput } from '../types'; import { useGenerateWizard } from '../useGenerateWizard'; import { Text } from 'ink'; import { render } from 'ink-testing-library'; import React, { act, useImperativeHandle } from 'react'; -import { describe, expect, it } from 'vitest'; +import { describe, expect, it, vi } from 'vitest'; // --------------------------------------------------------------------------- // Imperative harness — exposes wizard methods via ref for act()-based tests @@ -20,7 +21,7 @@ const Harness = React.forwardRef((props return ( step:{wizard.step} steps:{wizard.steps.join(',')} networkMode:{wizard.config.networkMode ?? 'undefined'}{' '} - advancedSelected:{String(wizard.advancedSelected)} + advancedSelected:{String(wizard.advancedSelected)} dockerfile:{wizard.config.dockerfile ?? 'undefined'} ); }); @@ -89,12 +90,12 @@ describe('useGenerateWizard — advanced config gate', () => { }); } - it('setAdvanced(false) jumps to confirm with PUBLIC defaults', () => { + it('setAdvanced([]) jumps to confirm with PUBLIC defaults', () => { const { ref, lastFrame } = setup(); walkToAdvanced(ref); expect(lastFrame()).toContain('step:advanced'); - act(() => ref.current!.wizard.setAdvanced(false)); + act(() => ref.current!.wizard.setAdvanced([])); const frame = lastFrame()!; expect(frame).toContain('step:confirm'); @@ -102,22 +103,28 @@ describe('useGenerateWizard — advanced config gate', () => { expect(frame).toContain('advancedSelected:false'); }); - it('setAdvanced(true) navigates to networkMode', () => { + it('setAdvanced with network selected navigates to networkMode', () => { + vi.useFakeTimers(); const { ref, lastFrame } = setup(); walkToAdvanced(ref); - act(() => ref.current!.wizard.setAdvanced(true)); + act(() => ref.current!.wizard.setAdvanced(['network', 'headers', 'auth', 'lifecycle'])); + // Flush setTimeout used for navigating to first sub-step + act(() => { + vi.runAllTimers(); + }); const frame = lastFrame()!; expect(frame).toContain('step:networkMode'); expect(frame).toContain('advancedSelected:true'); + vi.useRealTimers(); }); - it('setAdvanced(true) injects networkMode and requestHeaderAllowlist into steps', () => { + it('setAdvanced with settings injects sub-steps after advanced', () => { const { ref } = setup(); walkToAdvanced(ref); - act(() => ref.current!.wizard.setAdvanced(true)); + act(() => ref.current!.wizard.setAdvanced(['network', 'headers', 'auth', 'lifecycle'])); const steps = ref.current!.wizard.steps; const advIdx = steps.indexOf('advanced'); @@ -132,11 +139,11 @@ describe('useGenerateWizard — advanced config gate', () => { ]); }); - it('setAdvanced(true) then VPC injects subnets and securityGroups', () => { + it('network setting with VPC injects subnets and securityGroups', () => { const { ref } = setup(); walkToAdvanced(ref); - act(() => ref.current!.wizard.setAdvanced(true)); + act(() => ref.current!.wizard.setAdvanced(['network', 'headers', 'auth', 'lifecycle'])); act(() => ref.current!.wizard.setNetworkMode('VPC')); const steps = ref.current!.wizard.steps; @@ -156,7 +163,7 @@ describe('useGenerateWizard — advanced config gate', () => { }); describe('state cleanup on toggle', () => { - function walkToAdvancedAndSelectYes(ref: React.RefObject) { + function walkToAdvancedAndSelectSettings(ref: React.RefObject) { act(() => { ref.current!.wizard.setProjectName('Test'); ref.current!.wizard.setLanguage('Python'); @@ -166,18 +173,18 @@ describe('useGenerateWizard — advanced config gate', () => { ref.current!.wizard.setModelProvider('Bedrock'); ref.current!.wizard.setMemory('none'); }); - act(() => ref.current!.wizard.setAdvanced(true)); + act(() => ref.current!.wizard.setAdvanced(['network', 'headers', 'auth', 'lifecycle'])); act(() => ref.current!.wizard.setNetworkMode('VPC')); act(() => ref.current!.wizard.setSubnets(['subnet-123'])); act(() => ref.current!.wizard.setSecurityGroups(['sg-456'])); } - it('switching from Yes to No clears VPC config', () => { + it('switching to empty selection clears VPC config', () => { const { ref } = setup(); - walkToAdvancedAndSelectYes(ref); + walkToAdvancedAndSelectSettings(ref); - // Now go back and select No - act(() => ref.current!.wizard.setAdvanced(false)); + // Now go back and deselect all + act(() => ref.current!.wizard.setAdvanced([])); const w = ref.current!.wizard; expect(w.step).toBe('confirm'); @@ -191,9 +198,9 @@ describe('useGenerateWizard — advanced config gate', () => { it('config subnets and securityGroups are cleared to undefined', () => { const { ref } = setup(); - walkToAdvancedAndSelectYes(ref); + walkToAdvancedAndSelectSettings(ref); - act(() => ref.current!.wizard.setAdvanced(false)); + act(() => ref.current!.wizard.setAdvanced([])); expect(ref.current!.wizard.config.subnets).toBeUndefined(); expect(ref.current!.wizard.config.securityGroups).toBeUndefined(); @@ -272,6 +279,112 @@ describe('useGenerateWizard — advanced config gate', () => { }); }); + describe('dockerfile advanced setting', () => { + function walkToAdvancedWithContainer(ref: React.RefObject) { + act(() => { + ref.current!.wizard.setProjectName('Test'); + ref.current!.wizard.setLanguage('Python'); + ref.current!.wizard.setBuildType('Container'); + ref.current!.wizard.setProtocol('HTTP'); + ref.current!.wizard.setSdk('Strands'); + ref.current!.wizard.setModelProvider('Bedrock'); + ref.current!.wizard.setMemory('none'); + }); + } + + it('setAdvanced with only dockerfile navigates to dockerfile step', () => { + vi.useFakeTimers(); + const { ref, lastFrame } = setup(); + walkToAdvancedWithContainer(ref); + expect(lastFrame()).toContain('step:advanced'); + + act(() => ref.current!.wizard.setAdvanced(['dockerfile'])); + act(() => { + vi.runAllTimers(); + }); + + expect(lastFrame()).toContain('step:dockerfile'); + vi.useRealTimers(); + }); + + it('setDockerfile navigates to confirm when only dockerfile is selected', () => { + vi.useFakeTimers(); + const { ref, lastFrame } = setup(); + walkToAdvancedWithContainer(ref); + + act(() => ref.current!.wizard.setAdvanced(['dockerfile'])); + act(() => { + vi.runAllTimers(); + }); + expect(lastFrame()).toContain('step:dockerfile'); + + act(() => ref.current!.wizard.setDockerfile('Dockerfile.gpu')); + act(() => { + vi.runAllTimers(); + }); + + expect(lastFrame()).toContain('step:confirm'); + expect(lastFrame()).toContain('dockerfile:Dockerfile.gpu'); + vi.useRealTimers(); + }); + + it('dockerfile + lifecycle injects both sub-steps but not networkMode', () => { + const { ref } = setup(); + walkToAdvancedWithContainer(ref); + + act(() => ref.current!.wizard.setAdvanced(['dockerfile', 'lifecycle'])); + + const steps = ref.current!.wizard.steps; + const advIdx = steps.indexOf('advanced'); + expect(steps.slice(advIdx)).toEqual(['advanced', 'dockerfile', 'idleTimeout', 'maxLifetime', 'confirm']); + expect(steps).not.toContain('networkMode'); + }); + + it('dockerfile is hidden for CodeZip builds even when selected', () => { + const { ref } = setup(); + // Use CodeZip (default from walkToAdvanced) + act(() => { + ref.current!.wizard.setProjectName('Test'); + ref.current!.wizard.setLanguage('Python'); + ref.current!.wizard.setBuildType('CodeZip'); + ref.current!.wizard.setProtocol('HTTP'); + ref.current!.wizard.setSdk('Strands'); + ref.current!.wizard.setModelProvider('Bedrock'); + ref.current!.wizard.setMemory('none'); + }); + + // Even if dockerfile somehow gets into the set, it shouldn't appear for CodeZip + act(() => ref.current!.wizard.setAdvanced(['dockerfile', 'lifecycle'])); + + const steps = ref.current!.wizard.steps; + expect(steps).not.toContain('dockerfile'); + expect(steps).toContain('idleTimeout'); + }); + + it('deselecting all advanced clears dockerfile config', () => { + vi.useFakeTimers(); + const { ref } = setup(); + walkToAdvancedWithContainer(ref); + + act(() => ref.current!.wizard.setAdvanced(['dockerfile'])); + act(() => { + vi.runAllTimers(); + }); + act(() => ref.current!.wizard.setDockerfile('Dockerfile.gpu')); + act(() => { + vi.runAllTimers(); + }); + expect(ref.current!.wizard.config.dockerfile).toBe('Dockerfile.gpu'); + + // Go back and deselect all + act(() => ref.current!.wizard.setAdvanced([])); + + expect(ref.current!.wizard.config.dockerfile).toBeUndefined(); + expect(ref.current!.wizard.step).toBe('confirm'); + vi.useRealTimers(); + }); + }); + describe('reset clears advancedSelected', () => { it('reset returns advancedSelected to false', () => { const { ref, lastFrame } = setup(); @@ -283,7 +396,7 @@ describe('useGenerateWizard — advanced config gate', () => { ref.current!.wizard.setSdk('Strands'); ref.current!.wizard.setModelProvider('Bedrock'); ref.current!.wizard.setMemory('none'); - ref.current!.wizard.setAdvanced(true); + ref.current!.wizard.setAdvanced(['network']); }); expect(lastFrame()).toContain('advancedSelected:true'); @@ -293,3 +406,32 @@ describe('useGenerateWizard — advanced config gate', () => { }); }); }); + +describe('validateDockerfileInput', () => { + it('accepts empty string (use default)', () => { + expect(validateDockerfileInput('')).toBe(true); + expect(validateDockerfileInput(' ')).toBe(true); + }); + + it.each(['Dockerfile', 'Dockerfile.gpu', 'Dockerfile.dev-v2', 'my.Dockerfile'])( + 'accepts valid filename "%s"', + name => { + expect(validateDockerfileInput(name)).toBe(true); + } + ); + + it('accepts path input (delegates existence check to caller)', () => { + expect(validateDockerfileInput('../shared/Dockerfile.gpu')).toBe(true); + expect(validateDockerfileInput('/absolute/path/Dockerfile')).toBe(true); + }); + + it('rejects name exceeding 255 characters', () => { + const longName = 'D' + 'a'.repeat(255); + expect(validateDockerfileInput(longName)).toContain('255 characters'); + }); + + it.each(['.hidden', '-bad', '_under'])('rejects invalid filename "%s"', name => { + const result = validateDockerfileInput(name); + expect(result).not.toBe(true); + }); +}); diff --git a/src/cli/tui/screens/generate/types.ts b/src/cli/tui/screens/generate/types.ts index 1b5d3e7ed..eb1916399 100644 --- a/src/cli/tui/screens/generate/types.ts +++ b/src/cli/tui/screens/generate/types.ts @@ -14,6 +14,7 @@ export type GenerateStep = | 'projectName' | 'language' | 'buildType' + | 'dockerfile' | 'protocol' | 'sdk' | 'modelProvider' @@ -38,6 +39,8 @@ export type { BuildType, ModelProvider, ProtocolMode, SDKFramework, TargetLangua export interface GenerateConfig { projectName: string; buildType: BuildType; + /** Path to custom Dockerfile (copied into code directory at setup) or filename already in code directory. */ + dockerfile?: string; protocol: ProtocolMode; sdk: SDKFramework; modelProvider: ModelProvider; @@ -77,6 +80,7 @@ export const STEP_LABELS: Record = { projectName: 'Name', language: 'Target Language', buildType: 'Build', + dockerfile: 'Dockerfile', protocol: 'Protocol', sdk: 'Framework', modelProvider: 'Model', @@ -149,11 +153,40 @@ export const NETWORK_MODE_OPTIONS = [ { id: 'VPC', title: 'VPC', description: 'Attach to your VPC' }, ] as const; -export const ADVANCED_OPTIONS = [ - { id: 'no', title: 'No, use defaults', description: 'Public network, no VPC' }, - { id: 'yes', title: 'Yes, customize', description: undefined }, +export type AdvancedSettingId = 'dockerfile' | 'network' | 'headers' | 'auth' | 'lifecycle'; + +export const ADVANCED_SETTING_OPTIONS = [ + { id: 'dockerfile', title: 'Custom Dockerfile', description: 'Specify a custom Dockerfile path' }, + { id: 'network', title: 'VPC networking', description: 'Attach to your VPC with subnets & security groups' }, + { id: 'headers', title: 'Request header allowlist', description: 'Allow custom headers through to your agent' }, + { id: 'auth', title: 'Custom auth (JWT)', description: 'OIDC-based token validation for inbound requests' }, + { id: 'lifecycle', title: 'Lifecycle timeouts', description: 'Idle timeout & max instance lifetime' }, ] as const; +/** Dockerfile filename regex — must match the Zod schema in agent-env.ts */ +const DOCKERFILE_NAME_REGEX = /^[a-zA-Z0-9][a-zA-Z0-9._-]*$/; + +/** + * Validate a Dockerfile input value from the TUI. + * Returns `true` if valid, or an error message string if invalid. + * Does NOT check file existence for path inputs — callers handle that. + */ +export function validateDockerfileInput(value: string): true | string { + const trimmed = value.trim(); + if (!trimmed) return true; // empty is valid (means "use default") + if (trimmed.includes('/')) { + // Path input — caller must check existsSync separately + return true; + } + if (trimmed.length > 255) { + return 'Dockerfile name must be 255 characters or fewer'; + } + if (!DOCKERFILE_NAME_REGEX.test(trimmed)) { + return 'Must be a valid filename (starts with alphanumeric)'; + } + return true; +} + export const MEMORY_OPTIONS = [ { id: 'none', title: 'None', description: 'No memory' }, { id: 'shortTerm', title: 'Short-term memory', description: 'Context within a session' }, diff --git a/src/cli/tui/screens/generate/useGenerateWizard.ts b/src/cli/tui/screens/generate/useGenerateWizard.ts index 4808ab22c..982457aa2 100644 --- a/src/cli/tui/screens/generate/useGenerateWizard.ts +++ b/src/cli/tui/screens/generate/useGenerateWizard.ts @@ -1,7 +1,7 @@ import type { NetworkMode, RuntimeAuthorizerType } from '../../../../schema'; import { ProjectNameSchema } from '../../../../schema'; import type { JwtConfigOptions } from '../../../primitives/auth-utils'; -import type { BuildType, GenerateConfig, GenerateStep, MemoryOption, ProtocolMode } from './types'; +import type { AdvancedSettingId, BuildType, GenerateConfig, GenerateStep, MemoryOption, ProtocolMode } from './types'; import { BASE_GENERATE_STEPS, getModelProviderOptionsForSdk } from './types'; import { useCallback, useMemo, useState } from 'react'; @@ -35,11 +35,13 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { // Track if user has selected a framework (moved past sdk step) const [sdkSelected, setSdkSelected] = useState(false); - const [advancedSelected, setAdvancedSelected] = useState(false); + const [advancedSettings, setAdvancedSettings] = useState>(new Set()); + + const advancedSelected = advancedSettings.size > 0; // Steps depend on protocol, SDK, model provider, network mode, and whether we have an initial name // MCP skips sdk, modelProvider, apiKey, memory - // Filter out: projectName if initialName, apiKey for Bedrock, subnets/securityGroups for non-VPC + // Advanced sub-steps only appear for settings the user selected in the multi-select const steps = useMemo(() => { let filtered = BASE_GENERATE_STEPS; if (hasInitialName) { @@ -59,25 +61,40 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { if (advancedSelected) { const advancedIndex = filtered.indexOf('advanced'); const afterAdvanced = advancedIndex + 1; - const networkSteps: GenerateStep[] = - config.networkMode === 'VPC' ? ['networkMode', 'subnets', 'securityGroups'] : ['networkMode']; - filtered = [ - ...filtered.slice(0, afterAdvanced), - ...networkSteps, - 'requestHeaderAllowlist', - 'authorizerType', - 'idleTimeout', - 'maxLifetime', - ...filtered.slice(afterAdvanced), - ]; + const subSteps: GenerateStep[] = []; + // Dockerfile — only for Container builds when user selected it + if (advancedSettings.has('dockerfile') && config.buildType === 'Container') { + subSteps.push('dockerfile'); + } + // Network — always networkMode, plus subnets/securityGroups for VPC + if (advancedSettings.has('network')) { + subSteps.push('networkMode'); + if (config.networkMode === 'VPC') { + subSteps.push('subnets', 'securityGroups'); + } + } + // Headers + if (advancedSettings.has('headers')) { + subSteps.push('requestHeaderAllowlist'); + } + // Auth + if (advancedSettings.has('auth')) { + subSteps.push('authorizerType'); + } + // Lifecycle + if (advancedSettings.has('lifecycle')) { + subSteps.push('idleTimeout', 'maxLifetime'); + } + filtered = [...filtered.slice(0, afterAdvanced), ...subSteps, ...filtered.slice(afterAdvanced)]; } // Add jwtConfig step after authorizerType when CUSTOM_JWT is selected - if (config.authorizerType === 'CUSTOM_JWT') { + if (config.authorizerType === 'CUSTOM_JWT' && filtered.includes('authorizerType')) { const authIndex = filtered.indexOf('authorizerType'); filtered = [...filtered.slice(0, authIndex + 1), 'jwtConfig', ...filtered.slice(authIndex + 1)]; } return filtered; }, [ + config.buildType, config.modelProvider, config.sdk, config.protocol, @@ -86,6 +103,7 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { hasInitialName, sdkSelected, advancedSelected, + advancedSettings, ]); const currentIndex = steps.indexOf(step); @@ -108,7 +126,7 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { }, []); const setBuildType = useCallback((buildType: BuildType) => { - setConfig(c => ({ ...c, buildType })); + setConfig(c => ({ ...c, buildType, dockerfile: undefined })); setStep('protocol'); }, []); @@ -175,67 +193,128 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { setStep('advanced'); }, []); - const setAdvanced = useCallback((wantsAdvanced: boolean) => { - if (wantsAdvanced) { - setAdvancedSelected(true); - setStep('networkMode'); - } else { - setAdvancedSelected(false); - setConfig(c => ({ - ...c, - networkMode: 'PUBLIC', - subnets: undefined, - securityGroups: undefined, - requestHeaderAllowlist: undefined, - idleRuntimeSessionTimeout: undefined, - maxLifetime: undefined, - })); - setStep('confirm'); - } - }, []); + /** Navigate to the next step after the current one in the steps array */ + const goToNextStep = useCallback( + (afterStep: GenerateStep) => { + // Find the step after afterStep in the current steps array, or fall back to confirm + const idx = steps.indexOf(afterStep); + const next = idx >= 0 ? steps[idx + 1] : undefined; + setStep(next ?? 'confirm'); + }, + [steps] + ); - const setNetworkMode = useCallback((networkMode: NetworkMode) => { - setConfig(c => ({ ...c, networkMode })); - if (networkMode === 'VPC') { - setStep('subnets'); - } else { - setStep('requestHeaderAllowlist'); - } - }, []); + const setDockerfile = useCallback( + (dockerfile: string | undefined) => { + setConfig(c => ({ ...c, dockerfile })); + setTimeout(() => goToNextStep('dockerfile'), 0); + }, + [goToNextStep] + ); + + const setAdvanced = useCallback( + (selectedIds: AdvancedSettingId[]) => { + const selected = new Set(selectedIds); + setAdvancedSettings(selected); + if (selected.size === 0) { + // No advanced settings — reset defaults and go to confirm + setConfig(c => ({ + ...c, + dockerfile: undefined, + networkMode: 'PUBLIC', + subnets: undefined, + securityGroups: undefined, + requestHeaderAllowlist: undefined, + authorizerType: undefined, + jwtConfig: undefined, + idleRuntimeSessionTimeout: undefined, + maxLifetime: undefined, + })); + setStep('confirm'); + } else { + // Navigate to first advanced sub-step — determined by the steps memo on next render. + // Use setTimeout so the steps memo recalculates with the new advancedSettings first. + setTimeout(() => { + // The steps array hasn't updated yet, so we compute the first sub-step manually + if (selected.has('dockerfile') && config.buildType === 'Container') { + setStep('dockerfile'); + } else if (selected.has('network')) { + setStep('networkMode'); + } else if (selected.has('headers')) { + setStep('requestHeaderAllowlist'); + } else if (selected.has('auth')) { + setStep('authorizerType'); + } else if (selected.has('lifecycle')) { + setStep('idleTimeout'); + } else { + setStep('confirm'); + } + }, 0); + } + }, + [config.buildType] + ); + + const setNetworkMode = useCallback( + (networkMode: NetworkMode) => { + setConfig(c => ({ ...c, networkMode })); + if (networkMode === 'VPC') { + setStep('subnets'); + } else { + // Skip subnets/securityGroups, go to next step after networkMode + // We need to find next step after where securityGroups would be, or after networkMode + // Since steps array adapts, just go to next after networkMode + setTimeout(() => goToNextStep('networkMode'), 0); + } + }, + [goToNextStep] + ); const setSubnets = useCallback((subnets: string[]) => { setConfig(c => ({ ...c, subnets })); setStep('securityGroups'); }, []); - const setSecurityGroups = useCallback((securityGroups: string[]) => { - setConfig(c => ({ ...c, securityGroups })); - setStep('requestHeaderAllowlist'); - }, []); + const setSecurityGroups = useCallback( + (securityGroups: string[]) => { + setConfig(c => ({ ...c, securityGroups })); + setTimeout(() => goToNextStep('securityGroups'), 0); + }, + [goToNextStep] + ); - const setRequestHeaderAllowlist = useCallback((requestHeaderAllowlist: string[]) => { - setConfig(c => ({ ...c, requestHeaderAllowlist })); - setStep('authorizerType'); - }, []); + const setRequestHeaderAllowlist = useCallback( + (requestHeaderAllowlist: string[]) => { + setConfig(c => ({ ...c, requestHeaderAllowlist })); + setTimeout(() => goToNextStep('requestHeaderAllowlist'), 0); + }, + [goToNextStep] + ); const skipRequestHeaderAllowlist = useCallback(() => { - setStep('authorizerType'); - }, []); - - const setAuthorizerType = useCallback((authorizerType: RuntimeAuthorizerType) => { - setConfig(c => ({ ...c, authorizerType })); - if (authorizerType === 'CUSTOM_JWT') { - setStep('jwtConfig'); - } else { - setConfig(c => ({ ...c, authorizerType, jwtConfig: undefined })); - setStep('idleTimeout'); - } - }, []); + setTimeout(() => goToNextStep('requestHeaderAllowlist'), 0); + }, [goToNextStep]); + + const setAuthorizerType = useCallback( + (authorizerType: RuntimeAuthorizerType) => { + setConfig(c => ({ ...c, authorizerType })); + if (authorizerType === 'CUSTOM_JWT') { + setStep('jwtConfig'); + } else { + setConfig(c => ({ ...c, authorizerType, jwtConfig: undefined })); + setTimeout(() => goToNextStep('authorizerType'), 0); + } + }, + [goToNextStep] + ); - const setJwtConfig = useCallback((jwtConfig: JwtConfigOptions) => { - setConfig(c => ({ ...c, jwtConfig })); - setStep('idleTimeout'); - }, []); + const setJwtConfig = useCallback( + (jwtConfig: JwtConfigOptions) => { + setConfig(c => ({ ...c, jwtConfig })); + setTimeout(() => goToNextStep('jwtConfig'), 0); + }, + [goToNextStep] + ); const setIdleTimeout = useCallback((value: number | undefined) => { setConfig(c => ({ ...c, idleRuntimeSessionTimeout: value })); @@ -266,7 +345,7 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { setConfig(getDefaultConfig()); setError(null); setSdkSelected(false); - setAdvancedSelected(false); + setAdvancedSettings(new Set()); }, []); /** @@ -290,6 +369,7 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { setProjectName, setLanguage, setBuildType, + setDockerfile, setProtocol, setSdk, setModelProvider, @@ -298,6 +378,7 @@ export function useGenerateWizard(options?: UseGenerateWizardOptions) { setMemory, setAdvanced, advancedSelected, + advancedSettings, setNetworkMode, setSubnets, setSecurityGroups, diff --git a/src/lib/__tests__/constants.test.ts b/src/lib/__tests__/constants.test.ts index 72f823be6..21b952b69 100644 --- a/src/lib/__tests__/constants.test.ts +++ b/src/lib/__tests__/constants.test.ts @@ -1,4 +1,5 @@ -import { getArtifactZipName } from '../constants.js'; +import { getArtifactZipName, getDockerfilePath } from '../constants.js'; +import { join } from 'path'; import { describe, expect, it } from 'vitest'; describe('getArtifactZipName', () => { @@ -18,3 +19,29 @@ describe('getArtifactZipName', () => { expect(getArtifactZipName('agent.tar')).toBe('agent.tar.zip'); }); }); + +describe('getDockerfilePath', () => { + it('returns default Dockerfile when no custom name given', () => { + expect(getDockerfilePath('/app/code')).toBe(join('/app/code', 'Dockerfile')); + }); + + it('returns custom dockerfile name joined to code location', () => { + expect(getDockerfilePath('/app/code', 'Dockerfile.gpu')).toBe(join('/app/code', 'Dockerfile.gpu')); + }); + + it('rejects forward slash in dockerfile name', () => { + expect(() => getDockerfilePath('/app/code', '../Dockerfile')).toThrow(/Invalid dockerfile name/); + }); + + it('rejects backslash in dockerfile name', () => { + expect(() => getDockerfilePath('/app/code', 'Dockerfile\\..\\secret')).toThrow(/Invalid dockerfile name/); + }); + + it('rejects dot-dot traversal in dockerfile name', () => { + expect(() => getDockerfilePath('/app/code', '..')).toThrow(/Invalid dockerfile name/); + }); + + it('rejects path/to/Dockerfile', () => { + expect(() => getDockerfilePath('/app/code', 'path/to/Dockerfile')).toThrow(/Invalid dockerfile name/); + }); +}); diff --git a/src/lib/constants.ts b/src/lib/constants.ts index fac7f0a9c..8916c6052 100644 --- a/src/lib/constants.ts +++ b/src/lib/constants.ts @@ -59,7 +59,13 @@ export const START_HINTS: Record = { /** * Get the Dockerfile path for a given code location. + * @param codeLocation - Directory containing the Dockerfile + * @param dockerfile - Custom Dockerfile name (default: 'Dockerfile') */ -export function getDockerfilePath(codeLocation: string): string { - return join(codeLocation, DOCKERFILE_NAME); +export function getDockerfilePath(codeLocation: string, dockerfile?: string): string { + const name = dockerfile ?? DOCKERFILE_NAME; + if (name.includes('/') || name.includes('\\') || name.includes('..')) { + throw new Error(`Invalid dockerfile name: must be a filename without path separators or traversal`); + } + return join(codeLocation, name); } diff --git a/src/lib/packaging/__tests__/container.test.ts b/src/lib/packaging/__tests__/container.test.ts index 0dc6f285a..cb5685dbf 100644 --- a/src/lib/packaging/__tests__/container.test.ts +++ b/src/lib/packaging/__tests__/container.test.ts @@ -178,6 +178,38 @@ describe('ContainerPackager', () => { expect(result.artifactPath).toBe('finch://agentcore-package-agent'); }); + it('uses custom dockerfile name from spec', async () => { + mockResolveCodeLocation.mockReturnValue('/resolved/src'); + mockExistsSync.mockReturnValue(true); + mockSpawnSync.mockImplementation((cmd: string, args: string[]) => { + if (cmd === 'which' && args[0] === 'docker') return { status: 0 }; + if (cmd === 'docker' && args[0] === '--version') return { status: 0 }; + if (cmd === 'docker' && args[0] === 'build') return { status: 0 }; + if (cmd === 'docker' && args[0] === 'image') return { status: 0, stdout: Buffer.from('1000') }; + return { status: 1 }; + }); + + const specWithDockerfile = { ...baseSpec, dockerfile: 'Dockerfile.gpu' }; + await packager.pack(specWithDockerfile as any); + + // Verify build was called with custom dockerfile path + const buildCall = mockSpawnSync.mock.calls.find( + (c: unknown[]) => c[0] === 'docker' && (c[1] as string[])[0] === 'build' + ); + expect(buildCall).toBeDefined(); + const buildArgs = buildCall![1] as string[]; + const fIdx = buildArgs.indexOf('-f'); + expect(buildArgs[fIdx + 1]).toBe('/resolved/src/Dockerfile.gpu'); + }); + + it('rejects when custom dockerfile not found', async () => { + mockResolveCodeLocation.mockReturnValue('/resolved/src'); + mockExistsSync.mockReturnValue(false); + + const specWithDockerfile = { ...baseSpec, dockerfile: 'Dockerfile.custom' }; + await expect(packager.pack(specWithDockerfile as any)).rejects.toThrow('Dockerfile.custom not found'); + }); + it('detects podman runtime last', async () => { mockResolveCodeLocation.mockReturnValue('/resolved/src'); mockExistsSync.mockReturnValue(true); diff --git a/src/lib/packaging/container.ts b/src/lib/packaging/container.ts index 79cf7325e..a166a08f3 100644 --- a/src/lib/packaging/container.ts +++ b/src/lib/packaging/container.ts @@ -1,12 +1,11 @@ import type { AgentEnvSpec } from '../../schema'; -import { CONTAINER_RUNTIMES, DOCKERFILE_NAME, ONE_GB } from '../constants'; +import { CONTAINER_RUNTIMES, DOCKERFILE_NAME, ONE_GB, getDockerfilePath } from '../constants'; import { getUvBuildArgs } from './build-args'; import { PackagingError } from './errors'; import { resolveCodeLocation } from './helpers'; import type { ArtifactResult, PackageOptions, RuntimePackager } from './types/packaging'; import { spawnSync } from 'child_process'; import { existsSync } from 'fs'; -import { join } from 'path'; /** * Detect container runtime synchronously. @@ -36,12 +35,14 @@ export class ContainerPackager implements RuntimePackager { const agentName = options.agentName ?? spec.name; const configBaseDir = options.artifactDir ?? options.projectRoot ?? process.cwd(); const codeLocation = resolveCodeLocation(spec.codeLocation, configBaseDir); - const dockerfilePath = join(codeLocation, DOCKERFILE_NAME); + const dockerfilePath = getDockerfilePath(codeLocation, spec.dockerfile); // Preflight: Dockerfile must exist if (!existsSync(dockerfilePath)) { return Promise.reject( - new PackagingError(`Dockerfile not found at ${dockerfilePath}. Container agents require a Dockerfile.`) + new PackagingError( + `${spec.dockerfile ?? DOCKERFILE_NAME} not found at ${dockerfilePath}. Container agents require a Dockerfile.` + ) ); } diff --git a/src/schema/llm-compacted/agentcore.ts b/src/schema/llm-compacted/agentcore.ts index 334a5ce8a..7fcda0245 100644 --- a/src/schema/llm-compacted/agentcore.ts +++ b/src/schema/llm-compacted/agentcore.ts @@ -46,6 +46,7 @@ interface AgentEnvSpec { build: BuildType; entrypoint: string; // @regex ^[a-zA-Z0-9_][a-zA-Z0-9_/.-]*\.(py|ts|js)(:[a-zA-Z_][a-zA-Z0-9_]*)?$ e.g. "main.py:handler" or "index.ts" codeLocation: string; // Directory path + dockerfile?: string; // Custom Dockerfile name for Container builds (default: 'Dockerfile'). Must be a filename, not a path. runtimeVersion?: RuntimeVersion; envVars?: EnvVar[]; networkMode?: NetworkMode; // default 'PUBLIC' diff --git a/src/schema/schemas/__tests__/agent-env.test.ts b/src/schema/schemas/__tests__/agent-env.test.ts index 7fcbc17c3..477c2489a 100644 --- a/src/schema/schemas/__tests__/agent-env.test.ts +++ b/src/schema/schemas/__tests__/agent-env.test.ts @@ -405,6 +405,88 @@ describe('LifecycleConfigurationSchema', () => { }); }); +describe('AgentEnvSpecSchema - dockerfile', () => { + const validContainerAgent = { + name: 'ContainerAgent', + build: 'Container', + entrypoint: 'main.py', + codeLocation: './agents/container', + }; + + const validCodeZipAgent = { + name: 'CodeZipAgent', + build: 'CodeZip', + entrypoint: 'main.py:handler', + codeLocation: './agents/test', + runtimeVersion: 'PYTHON_3_12', + }; + + it('accepts Container agent with custom dockerfile', () => { + const result = AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: 'Dockerfile.gpu' }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.dockerfile).toBe('Dockerfile.gpu'); + } + }); + + it('accepts Container agent without dockerfile (optional)', () => { + const result = AgentEnvSpecSchema.safeParse(validContainerAgent); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.dockerfile).toBeUndefined(); + } + }); + + it.each(['Dockerfile', 'Dockerfile.dev', 'Dockerfile.gpu-v2', 'my.Dockerfile', 'dockerfile_test'])( + 'accepts valid dockerfile name "%s"', + name => { + expect(AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: name }).success).toBe(true); + } + ); + + it('rejects dockerfile on CodeZip builds', () => { + const result = AgentEnvSpecSchema.safeParse({ ...validCodeZipAgent, dockerfile: 'Dockerfile.custom' }); + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error.issues.some(i => i.message.includes('only allowed for Container'))).toBe(true); + } + }); + + it.each(['../Dockerfile', '/etc/Dockerfile', 'path/to/Dockerfile', '.hidden'])( + 'rejects path traversal or path separator in dockerfile "%s"', + name => { + expect(AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: name }).success).toBe(false); + } + ); + + it('rejects empty string dockerfile', () => { + expect(AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: '' }).success).toBe(false); + }); + + it('rejects shell metacharacters in dockerfile', () => { + expect(AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: 'Dockerfile;rm -rf /' }).success).toBe( + false + ); + }); + + it('rejects dockerfile exceeding 255 characters', () => { + const longName = 'D' + 'a'.repeat(255); + expect(AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: longName }).success).toBe(false); + }); + + it('accepts dockerfile at exactly 255 characters', () => { + const maxName = 'D' + 'a'.repeat(254); + expect(AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: maxName }).success).toBe(true); + }); + + it.each(['\\\\server\\share', 'Dockerfile\\..\\secret', '..\\Dockerfile'])( + 'rejects backslash path traversal in dockerfile "%s"', + name => { + expect(AgentEnvSpecSchema.safeParse({ ...validContainerAgent, dockerfile: name }).success).toBe(false); + } + ); +}); + describe('AgentEnvSpecSchema - lifecycleConfiguration', () => { const validAgent = { name: 'TestAgent', diff --git a/src/schema/schemas/agent-env.ts b/src/schema/schemas/agent-env.ts index 1d67728e8..fc2c54d1e 100644 --- a/src/schema/schemas/agent-env.ts +++ b/src/schema/schemas/agent-env.ts @@ -183,6 +183,13 @@ export const AgentEnvSpecSchema = z build: BuildTypeSchema, entrypoint: EntrypointSchema, codeLocation: DirectoryPathSchema, + /** Custom Dockerfile name for Container builds. Must be a filename, not a path. Default: 'Dockerfile' */ + dockerfile: z + .string() + .min(1) + .max(255) + .regex(/^[a-zA-Z0-9][a-zA-Z0-9._-]*$/, 'Must be a filename (no path separators or traversal)') + .optional(), runtimeVersion: RuntimeVersionSchemaFromConstants.optional(), /** Environment variables to set on the runtime */ envVars: z.array(EnvVarSchema).optional(), @@ -235,6 +242,14 @@ export const AgentEnvSpecSchema = z path: ['authorizerConfiguration'], }); } + // If adding more Container-specific fields, consider consolidating into a containerConfig object (see networkConfig pattern) + if (data.build !== 'Container' && data.dockerfile) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: 'dockerfile is only allowed for Container builds', + path: ['dockerfile'], + }); + } }); export type AgentEnvSpec = z.infer;