From e4d69154e1079e623ce69b22cc92a50b15896cb1 Mon Sep 17 00:00:00 2001 From: jariy17 <107276415+jariy17@users.noreply.github.com> Date: Tue, 26 May 2026 18:26:03 +0000 Subject: [PATCH 1/3] fix(tui): back-navigation bugs in dataset and batch eval flows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three back-navigation issues reported on PR #1347: 1. BatchEvalWizard Escape skips source-picker — onExit passed the top-level exit handler instead of returning to the source-picker state. Now returns to source-picker on first-step back-nav. 2. AddDatasetScreen loses typed name — initialValue always called generateUniqueName() ignoring the stored name state. Now uses the previously submitted name when available. 3. DatasetFlow back-nav resets datasets to [] — all four sub-screen onExit/onCancel handlers hardcoded an empty array. Added loadedDatasets state that persists across flow transitions. --- src/cli/tui/screens/dataset-hub/DatasetFlow.tsx | 10 ++++++---- src/cli/tui/screens/dataset/AddDatasetScreen.tsx | 2 +- src/cli/tui/screens/run-eval/RunBatchEvalFlow.tsx | 2 +- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/cli/tui/screens/dataset-hub/DatasetFlow.tsx b/src/cli/tui/screens/dataset-hub/DatasetFlow.tsx index 9ae3584fd..a080dc25c 100644 --- a/src/cli/tui/screens/dataset-hub/DatasetFlow.tsx +++ b/src/cli/tui/screens/dataset-hub/DatasetFlow.tsx @@ -54,6 +54,7 @@ interface DatasetFlowProps { export function DatasetFlow({ onExit }: DatasetFlowProps) { const [flow, setFlow] = useState({ name: 'loading' }); + const [loadedDatasets, setLoadedDatasets] = useState([]); // Load datasets on mount useEffect(() => { @@ -100,6 +101,7 @@ export function DatasetFlow({ onExit }: DatasetFlowProps) { return; } + setLoadedDatasets(resolved); setFlow({ name: 'hub', datasets: resolved }); } catch (err) { setFlow({ name: 'error', message: err instanceof Error ? err.message : String(err) }); @@ -194,7 +196,7 @@ export function DatasetFlow({ onExit }: DatasetFlowProps) { setFlow({ name: 'confirm-pull', dataset: flow.dataset, version })} - onExit={() => setFlow({ name: 'hub', datasets: [] })} + onExit={() => setFlow({ name: 'hub', datasets: loadedDatasets })} /> ); } @@ -206,7 +208,7 @@ export function DatasetFlow({ onExit }: DatasetFlowProps) { location={flow.dataset.location} versionLabel={versionLabel} onConfirm={() => void executeAction('download', flow.dataset, flow.version)} - onCancel={() => setFlow({ name: 'hub', datasets: [] })} + onCancel={() => setFlow({ name: 'hub', datasets: loadedDatasets })} /> ); } @@ -224,7 +226,7 @@ export function DatasetFlow({ onExit }: DatasetFlowProps) { setFlow({ name: 'confirm-delete', dataset: flow.dataset, version })} - onExit={() => setFlow({ name: 'hub', datasets: [] })} + onExit={() => setFlow({ name: 'hub', datasets: loadedDatasets })} /> ); } @@ -235,7 +237,7 @@ export function DatasetFlow({ onExit }: DatasetFlowProps) { datasetName={flow.dataset.name} version={flow.version} onConfirm={() => void executeAction('confirm-delete', flow.dataset, flow.version)} - onCancel={() => setFlow({ name: 'hub', datasets: [] })} + onCancel={() => setFlow({ name: 'hub', datasets: loadedDatasets })} /> ); } diff --git a/src/cli/tui/screens/dataset/AddDatasetScreen.tsx b/src/cli/tui/screens/dataset/AddDatasetScreen.tsx index fb41b8baa..9bec86a29 100644 --- a/src/cli/tui/screens/dataset/AddDatasetScreen.tsx +++ b/src/cli/tui/screens/dataset/AddDatasetScreen.tsx @@ -101,7 +101,7 @@ export function AddDatasetScreen({ onComplete, onExit, existingDatasetNames }: A { setName(value); setStep('schema-type'); diff --git a/src/cli/tui/screens/run-eval/RunBatchEvalFlow.tsx b/src/cli/tui/screens/run-eval/RunBatchEvalFlow.tsx index b46ad4e36..a9935d28f 100644 --- a/src/cli/tui/screens/run-eval/RunBatchEvalFlow.tsx +++ b/src/cli/tui/screens/run-eval/RunBatchEvalFlow.tsx @@ -384,7 +384,7 @@ export function RunBatchEvalFlow({ onExit }: RunBatchEvalFlowProps) { source={flow.source} dataset={flow.dataset} onComplete={handleWizardComplete} - onExit={onExit} + onExit={() => setFlow({ name: 'source-picker', agents: flow.agents, evaluators: flow.evaluators })} /> ); } From 385c8e25c3d745057ed6c6733f777f5e04eb5bf0 Mon Sep 17 00:00:00 2001 From: jariy17 <107276415+jariy17@users.noreply.github.com> Date: Tue, 26 May 2026 18:35:28 +0000 Subject: [PATCH 2/3] feat(dataset): add --kms-key-arn support to add dataset command and TUI Wire kmsKeyArn through the full add-dataset flow: - CLI: --kms-key-arn flag with ARN validation - TUI: optional step after Description (Enter to skip) - Primitive: writes kmsKeyArn to agentcore.json when provided - CDK construct already passes KmsKeyArn to CloudFormation The field is immutable after dataset creation (createOnly CFN property). --- src/cli/commands/add/types.ts | 1 + src/cli/commands/add/validate.ts | 9 ++ src/cli/primitives/DatasetPrimitive.ts | 128 ++++++++++-------- .../tui/screens/dataset/AddDatasetFlow.tsx | 7 +- .../tui/screens/dataset/AddDatasetScreen.tsx | 36 ++++- 5 files changed, 115 insertions(+), 66 deletions(-) diff --git a/src/cli/commands/add/types.ts b/src/cli/commands/add/types.ts index 111c3f851..cf069bcca 100644 --- a/src/cli/commands/add/types.ts +++ b/src/cli/commands/add/types.ts @@ -106,6 +106,7 @@ export interface AddDatasetOptions { name: string; schemaType: DatasetSchemaType; description?: string; + kmsKeyArn?: string; json?: boolean; } diff --git a/src/cli/commands/add/validate.ts b/src/cli/commands/add/validate.ts index 9a4d75928..d28c29ce6 100644 --- a/src/cli/commands/add/validate.ts +++ b/src/cli/commands/add/validate.ts @@ -17,6 +17,7 @@ import { TargetLanguageSchema, getSupportedFrameworksForProtocol, getSupportedModelProviders, + isValidKmsKeyArn, matchEnumValue, } from '../../../schema'; import { ARN_VALIDATION_MESSAGE, isValidArn } from '../shared/arn-utils'; @@ -836,6 +837,14 @@ export function validateAddDatasetOptions(options: AddDatasetOptions): Validatio return { valid: false, error: `Invalid schema type: ${options.schemaType}. Valid options: ${valid}` }; } + if (options.kmsKeyArn && !isValidKmsKeyArn(options.kmsKeyArn)) { + return { + valid: false, + error: + '--kms-key-arn must be a valid KMS key ARN (e.g. arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012)', + }; + } + return { valid: true }; } diff --git a/src/cli/primitives/DatasetPrimitive.ts b/src/cli/primitives/DatasetPrimitive.ts index d6a20d968..36423c568 100644 --- a/src/cli/primitives/DatasetPrimitive.ts +++ b/src/cli/primitives/DatasetPrimitive.ts @@ -45,6 +45,7 @@ export class DatasetPrimitive extends BasePrimitive', 'Dataset description [non-interactive]') + .option('--kms-key-arn ', 'KMS key ARN for dataset encryption (optional) [non-interactive]') .option('--json', 'Output as JSON [non-interactive]') - .action(async (cliOptions: { name?: string; schemaType?: string; description?: string; json?: boolean }) => { - if (!findConfigRoot()) { - console.error('No agentcore project found. Run `agentcore create` first.'); - process.exit(1); - } - - if (cliOptions.name || cliOptions.json) { - // CLI mode - await runCliCommand('add.dataset', !!cliOptions.json, async () => { - const validation = validateAddDatasetOptions({ - name: cliOptions.name ?? '', - schemaType: (cliOptions.schemaType ?? '') as DatasetSchemaType, - description: cliOptions.description, - }); - - if (!validation.valid) { - throw new Error(validation.error); - } + .action( + async (cliOptions: { + name?: string; + schemaType?: string; + description?: string; + kmsKeyArn?: string; + json?: boolean; + }) => { + if (!findConfigRoot()) { + console.error('No agentcore project found. Run `agentcore create` first.'); + process.exit(1); + } - const result = await this.add({ - name: cliOptions.name!, - schemaType: cliOptions.schemaType! as DatasetSchemaType, - description: cliOptions.description, + if (cliOptions.name || cliOptions.json) { + // CLI mode + await runCliCommand('add.dataset', !!cliOptions.json, async () => { + const validation = validateAddDatasetOptions({ + name: cliOptions.name ?? '', + schemaType: (cliOptions.schemaType ?? '') as DatasetSchemaType, + description: cliOptions.description, + kmsKeyArn: cliOptions.kmsKeyArn, + }); + + if (!validation.valid) { + throw new Error(validation.error); + } + + const result = await this.add({ + name: cliOptions.name!, + schemaType: cliOptions.schemaType! as DatasetSchemaType, + description: cliOptions.description, + kmsKeyArn: cliOptions.kmsKeyArn, + }); + + if (!result.success) { + throw result.error; + } + + if (cliOptions.json) { + console.log(JSON.stringify(result)); + } else { + console.log(`Added dataset '${result.datasetName}'`); + console.log(` File: ${result.location}`); + } + + return {}; }); - - if (!result.success) { - throw result.error; - } - - if (cliOptions.json) { - console.log(JSON.stringify(result)); - } else { - console.log(`Added dataset '${result.datasetName}'`); - console.log(` File: ${result.location}`); + } else { + try { + // TUI fallback — dynamic imports to avoid pulling ink (async) into registry + requireTTY(); + const [{ render }, { default: React }, { AddFlow }] = await Promise.all([ + import('ink'), + import('react'), + import('../tui/screens/add/AddFlow'), + ]); + const { unmount } = render( + React.createElement(AddFlow, { + isInteractive: false, + initialResource: 'dataset', + onExit: () => { + unmount(); + process.exit(0); + }, + }) + ); + } catch (error) { + console.error(getErrorMessage(error)); + process.exit(1); } - - return {}; - }); - } else { - try { - // TUI fallback — dynamic imports to avoid pulling ink (async) into registry - requireTTY(); - const [{ render }, { default: React }, { AddFlow }] = await Promise.all([ - import('ink'), - import('react'), - import('../tui/screens/add/AddFlow'), - ]); - const { unmount } = render( - React.createElement(AddFlow, { - isInteractive: false, - initialResource: 'dataset', - onExit: () => { - unmount(); - process.exit(0); - }, - }) - ); - } catch (error) { - console.error(getErrorMessage(error)); - process.exit(1); } } - }); + ); this.registerRemoveSubcommand(removeCmd); } diff --git a/src/cli/tui/screens/dataset/AddDatasetFlow.tsx b/src/cli/tui/screens/dataset/AddDatasetFlow.tsx index 31b5cde6e..2c94524cb 100644 --- a/src/cli/tui/screens/dataset/AddDatasetFlow.tsx +++ b/src/cli/tui/screens/dataset/AddDatasetFlow.tsx @@ -36,7 +36,12 @@ export function AddDatasetFlow({ isInteractive = true, onExit, onBack, onDev, on const handleCreateComplete = useCallback((config: AddDatasetConfig) => { void datasetPrimitive - .add({ name: config.name, schemaType: config.schemaType, description: config.description }) + .add({ + name: config.name, + schemaType: config.schemaType, + description: config.description, + kmsKeyArn: config.kmsKeyArn, + }) .then(result => { if (result.success) { setFlow({ diff --git a/src/cli/tui/screens/dataset/AddDatasetScreen.tsx b/src/cli/tui/screens/dataset/AddDatasetScreen.tsx index 9bec86a29..5360efdd4 100644 --- a/src/cli/tui/screens/dataset/AddDatasetScreen.tsx +++ b/src/cli/tui/screens/dataset/AddDatasetScreen.tsx @@ -1,5 +1,5 @@ import type { DatasetSchemaType } from '../../../../schema'; -import { DatasetNameSchema } from '../../../../schema'; +import { DatasetNameSchema, isValidKmsKeyArn } from '../../../../schema'; import { ConfirmReview, Panel, Screen, StepIndicator, TextInput, WizardSelect } from '../../components'; import type { SelectableItem } from '../../components'; import { HELP_TEXT } from '../../constants'; @@ -24,18 +24,20 @@ export interface AddDatasetConfig { name: string; schemaType: DatasetSchemaType; description?: string; + kmsKeyArn?: string; } -type Step = 'name' | 'schema-type' | 'description' | 'confirm'; +type Step = 'name' | 'schema-type' | 'description' | 'kms-key' | 'confirm'; const STEP_LABELS: Record = { name: 'Name', 'schema-type': 'Schema Type', description: 'Description', + 'kms-key': 'KMS Key', confirm: 'Confirm', }; -const STEPS: Step[] = ['name', 'schema-type', 'description', 'confirm']; +const STEPS: Step[] = ['name', 'schema-type', 'description', 'kms-key', 'confirm']; interface AddDatasetScreenProps { onComplete: (config: AddDatasetConfig) => void; @@ -48,10 +50,12 @@ export function AddDatasetScreen({ onComplete, onExit, existingDatasetNames }: A const [name, setName] = useState(''); const [schemaType, setSchemaType] = useState('AGENTCORE_EVALUATION_PREDEFINED_V1'); const [description, setDescription] = useState(''); + const [kmsKeyArn, setKmsKeyArn] = useState(''); const isNameStep = step === 'name'; const isSchemaTypeStep = step === 'schema-type'; const isDescriptionStep = step === 'description'; + const isKmsKeyStep = step === 'kms-key'; const isConfirmStep = step === 'confirm'; const schemaTypeNav = useListNavigation({ @@ -66,8 +70,9 @@ export function AddDatasetScreen({ onComplete, onExit, existingDatasetNames }: A useListNavigation({ items: [{ id: 'confirm', title: 'Confirm' }], - onSelect: () => onComplete({ name, schemaType, description: description || undefined }), - onExit: () => setStep('description'), + onSelect: () => + onComplete({ name, schemaType, description: description || undefined, kmsKeyArn: kmsKeyArn || undefined }), + onExit: () => setStep('kms-key'), isActive: isConfirmStep, }); @@ -84,8 +89,9 @@ export function AddDatasetScreen({ onComplete, onExit, existingDatasetNames }: A { label: 'Name', value: name }, { label: 'Schema Type', value: schemaType }, ...(description ? [{ label: 'Description', value: description }] : []), + ...(kmsKeyArn ? [{ label: 'KMS Key', value: kmsKeyArn }] : []), ], - [name, schemaType, description] + [name, schemaType, description, kmsKeyArn] ); return ( @@ -125,15 +131,31 @@ export function AddDatasetScreen({ onComplete, onExit, existingDatasetNames }: A { setDescription(value); - setStep('confirm'); + setStep('kms-key'); }} onCancel={() => setStep('schema-type')} allowEmpty /> )} + {isKmsKeyStep && ( + { + setKmsKeyArn(value); + setStep('confirm'); + }} + onCancel={() => setStep('description')} + allowEmpty + customValidation={value => !value || isValidKmsKeyArn(value) || 'Must be a valid KMS key ARN'} + /> + )} + {isConfirmStep && } From ed30251be4ccb35cfd4d767694710593bef146ed Mon Sep 17 00:00:00 2001 From: jariy17 Date: Tue, 26 May 2026 19:05:22 +0000 Subject: [PATCH 3/3] test(dataset): add kmsKeyArn unit and integration tests Cover the kmsKeyArn flow with: - Unit: DatasetPrimitive.add() persists kmsKeyArn when provided, omits when absent - Integ: CLI --kms-key-arn writes to agentcore.json, invalid ARN rejected, absent ARN omitted --- integ-tests/add-remove-dataset.test.ts | 63 +++++++++++++++++++ .../__tests__/DatasetPrimitive.test.ts | 36 +++++++++++ 2 files changed, 99 insertions(+) diff --git a/integ-tests/add-remove-dataset.test.ts b/integ-tests/add-remove-dataset.test.ts index 82fa94ebb..64e50209d 100644 --- a/integ-tests/add-remove-dataset.test.ts +++ b/integ-tests/add-remove-dataset.test.ts @@ -120,6 +120,69 @@ describe('add/remove dataset', () => { expect(dataset.description).toBe('Test scenarios for billing'); }); + it('adds dataset with --kms-key-arn and persists to agentcore.json', async () => { + const result = await runCLI( + [ + 'add', + 'dataset', + '--name', + 'KmsDataset', + '--schema-type', + 'AGENTCORE_EVALUATION_PREDEFINED_V1', + '--kms-key-arn', + 'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012', + '--json', + ], + projectDir + ); + + expect(result.exitCode, `stdout: ${result.stdout}, stderr: ${result.stderr}`).toBe(0); + const json = parseJsonOutput(result.stdout) as { success: boolean; datasetName: string }; + expect(json.success).toBe(true); + expect(json.datasetName).toBe('KmsDataset'); + + const spec = JSON.parse(await readFile(join(projectDir, 'agentcore/agentcore.json'), 'utf-8')); + const dataset = spec.datasets.find((d: { name: string }) => d.name === 'KmsDataset'); + expect(dataset.kmsKeyArn).toBe('arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012'); + }); + + it('rejects invalid --kms-key-arn', async () => { + const result = await runCLI( + [ + 'add', + 'dataset', + '--name', + 'BadKms', + '--schema-type', + 'AGENTCORE_EVALUATION_PREDEFINED_V1', + '--kms-key-arn', + 'not-a-valid-arn', + '--json', + ], + projectDir + ); + + expect(result.exitCode).toBe(1); + const json = parseJsonOutput(result.stdout) as { success: boolean; error: string }; + expect(json.success).toBe(false); + expect(json.error).toContain('kms-key-arn'); + }); + + it('omits kmsKeyArn from agentcore.json when not provided', async () => { + const result = await runCLI( + ['add', 'dataset', '--name', 'NoKms', '--schema-type', 'AGENTCORE_EVALUATION_PREDEFINED_V1', '--json'], + projectDir + ); + + expect(result.exitCode).toBe(0); + + const spec = JSON.parse(await readFile(join(projectDir, 'agentcore/agentcore.json'), 'utf-8')); + const dataset = spec.datasets.find((d: { name: string }) => d.name === 'NoKms'); + expect(dataset).toBeTruthy(); + expect(dataset.kmsKeyArn).toBeUndefined(); + expect('kmsKeyArn' in dataset).toBe(false); + }); + it('removes a dataset', async () => { const result = await runCLI(['remove', 'dataset', '--name', 'MyPredefined', '--json'], projectDir); diff --git a/src/cli/primitives/__tests__/DatasetPrimitive.test.ts b/src/cli/primitives/__tests__/DatasetPrimitive.test.ts index 453852a3a..2a9deb2b9 100644 --- a/src/cli/primitives/__tests__/DatasetPrimitive.test.ts +++ b/src/cli/primitives/__tests__/DatasetPrimitive.test.ts @@ -85,6 +85,42 @@ describe('DatasetPrimitive', () => { } }); + it('adds dataset with kmsKeyArn and persists to spec', async () => { + mockReadProjectSpec.mockResolvedValue(makeProject()); + mockWriteProjectSpec.mockResolvedValue(undefined); + mockMkdir.mockResolvedValue(undefined); + mockCopyFile.mockResolvedValue(undefined); + + const result = await primitive.add({ + name: 'KmsDataset', + schemaType: 'AGENTCORE_EVALUATION_PREDEFINED_V1', + kmsKeyArn: 'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012', + }); + + expect(result.success).toBe(true); + + const writtenSpec = mockWriteProjectSpec.mock.calls[0]![0]; + expect(writtenSpec.datasets[0].kmsKeyArn).toBe( + 'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012' + ); + }); + + it('omits kmsKeyArn from spec when not provided', async () => { + mockReadProjectSpec.mockResolvedValue(makeProject()); + mockWriteProjectSpec.mockResolvedValue(undefined); + mockMkdir.mockResolvedValue(undefined); + mockCopyFile.mockResolvedValue(undefined); + + await primitive.add({ + name: 'PlainDataset', + schemaType: 'AGENTCORE_EVALUATION_PREDEFINED_V1', + }); + + const writtenSpec = mockWriteProjectSpec.mock.calls[0]![0]; + expect(writtenSpec.datasets[0].kmsKeyArn).toBeUndefined(); + expect('kmsKeyArn' in writtenSpec.datasets[0]).toBe(false); + }); + it('returns error when readProjectSpec rejects', async () => { mockReadProjectSpec.mockRejectedValue(new Error('disk failure'));