fix(dataset): tui back-navigation bugs + kmsKeyArn support

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);
 

src/cli/commands/add/types.ts

@@ -106,6 +106,7 @@ export interface AddDatasetOptions {
   name: string;
   schemaType: DatasetSchemaType;
   description?: string;
+  kmsKeyArn?: string;
   json?: boolean;
 }
 

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 };
 }
 

src/cli/primitives/DatasetPrimitive.ts

@@ -45,6 +45,7 @@ export class DatasetPrimitive extends BasePrimitive<AddDatasetOptions, Removable
         name: options.name,
         schemaType: options.schemaType,
         ...(options.description && { description: options.description }),
+        ...(options.kmsKeyArn && { kmsKeyArn: options.kmsKeyArn }),
         config: {
           managed: { location },
         },
@@ -140,70 +141,81 @@ export class DatasetPrimitive extends BasePrimitive<AddDatasetOptions, Removable
         'Dataset schema type: AGENTCORE_EVALUATION_PREDEFINED_V1 | AGENTCORE_EVALUATION_SIMULATED_V1 [non-interactive]'
       )
       .option('--description <description>', 'Dataset description [non-interactive]')
+      .option('--kms-key-arn <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);
   }

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'));
 

src/cli/tui/screens/dataset-hub/DatasetFlow.tsx

@@ -54,6 +54,7 @@ interface DatasetFlowProps {
 
 export function DatasetFlow({ onExit }: DatasetFlowProps) {
   const [flow, setFlow] = useState<FlowState>({ name: 'loading' });
+  const [loadedDatasets, setLoadedDatasets] = useState<ResolvedDatasetInfo[]>([]);
 
   // 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) {
       <VersionPickerScreen
         versions={flow.versions}
         onSelect={version => 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) {
       <DeleteVersionPickerScreen
         versions={flow.versions}
         onSelect={version => 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 })}
       />
     );
   }

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({