feat: add policy engine and policy support

src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap

@@ -374,6 +374,7 @@ test('AgentCoreStack synthesizes with empty spec', () => {
       credentials: [],
       evaluators: [],
       onlineEvalConfigs: [],
+      policyEngines: [],
     },
   });
   const template = Template.fromStack(stack);

src/assets/cdk/test/cdk.test.ts

@@ -13,6 +13,7 @@ test('AgentCoreStack synthesizes with empty spec', () => {
       credentials: [],
       evaluators: [],
       onlineEvalConfigs: [],
+      policyEngines: [],
     },
   });
   const template = Template.fromStack(stack);

src/cli/aws/index.ts

@@ -15,6 +15,14 @@ export {
 } from './agentcore-control';
 export { streamLogs, searchLogs, type LogEvent, type StreamLogsOptions, type SearchLogsOptions } from './cloudwatch';
 export { enableTransactionSearch, type TransactionSearchEnableResult } from './transaction-search';
+export {
+  startPolicyGeneration,
+  getPolicyGeneration,
+  type StartPolicyGenerationOptions,
+  type StartPolicyGenerationResult,
+  type GetPolicyGenerationOptions,
+  type GetPolicyGenerationResult,
+} from './policy-generation';
 export {
   DEFAULT_RUNTIME_USER_ID,
   invokeA2ARuntime,

src/cli/aws/policy-generation.ts

@@ -0,0 +1,116 @@
+import { getCredentialProvider } from './account';
+import {
+  BedrockAgentCoreControlClient,
+  GetPolicyGenerationCommand,
+  ListPolicyGenerationAssetsCommand,
+  StartPolicyGenerationCommand,
+  waitUntilPolicyGenerationCompleted,
+} from '@aws-sdk/client-bedrock-agentcore-control';
+import { WaiterState } from '@smithy/util-waiter';
+
+export interface StartPolicyGenerationOptions {
+  policyEngineId: string;
+  description: string;
+  region: string;
+  resourceArn: string;
+}
+
+export interface StartPolicyGenerationResult {
+  generationId: string;
+}
+
+export interface GetPolicyGenerationOptions {
+  generationId: string;
+  policyEngineId: string;
+  region: string;
+}
+
+export interface GetPolicyGenerationResult {
+  status: string;
+  statement: string;
+}
+
+export async function startPolicyGeneration(
+  options: StartPolicyGenerationOptions
+): Promise<StartPolicyGenerationResult> {
+  const client = new BedrockAgentCoreControlClient({
+    region: options.region,
+    credentials: getCredentialProvider(),
+  });
+
+  const command = new StartPolicyGenerationCommand({
+    policyEngineId: options.policyEngineId,
+    resource: { arn: options.resourceArn },
+    content: {
+      rawText: options.description,
+    },
+    name: `cli_generation_${Date.now()}`,
+  });
+
+  const response = await client.send(command);
+
+  if (!response.policyGenerationId) {
+    throw new Error('No generation ID returned from StartPolicyGeneration');
+  }
+
+  return { generationId: response.policyGenerationId };
+}
+
+export async function getPolicyGeneration(options: GetPolicyGenerationOptions): Promise<GetPolicyGenerationResult> {
+  const client = new BedrockAgentCoreControlClient({
+    region: options.region,
+    credentials: getCredentialProvider(),
+  });
+
+  // Use the SDK waiter to poll until generation completes
+  const waiterResult = await waitUntilPolicyGenerationCompleted(
+    { client, maxWaitTime: 120, minDelay: 2, maxDelay: 5 },
+    { policyGenerationId: options.generationId, policyEngineId: options.policyEngineId }
+  );
+
+  if (waiterResult.state !== WaiterState.SUCCESS) {
+    throw new Error(
+      `Policy generation did not complete within the timeout period (state: ${waiterResult.state}). ` +
+        `Generation ID: ${options.generationId}`
+    );
+  }
+
+  // Check the final status
+  const getCommand = new GetPolicyGenerationCommand({
+    policyGenerationId: options.generationId,
+    policyEngineId: options.policyEngineId,
+  });
+
+  const statusResponse = await client.send(getCommand);
+
+  if (statusResponse.status === 'GENERATE_FAILED') {
+    const reasons = statusResponse.statusReasons?.join(', ') ?? 'Unknown reason';
+    throw new Error(`Policy generation failed: ${reasons}`);
+  }
+
+  // Fetch the generated assets
+  const assetsCommand = new ListPolicyGenerationAssetsCommand({
+    policyGenerationId: options.generationId,
+    policyEngineId: options.policyEngineId,
+  });
+
+  const assetsResponse = await client.send(assetsCommand);
+  const assets = assetsResponse.policyGenerationAssets ?? [];
+
+  if (assets.length === 0) {
+    throw new Error('Policy generation completed but no assets were returned');
+  }
+
+  // Get the Cedar statement from the first asset
+  const firstAsset = assets[0]!;
+  const cedarStatement = firstAsset.definition?.cedar?.statement;
+
+  if (!cedarStatement) {
+    throw new Error('Policy generation completed but no Cedar policy statement was found in the assets');
+  }
+
+  return {
+    status: statusResponse.status ?? 'GENERATED',
+    statement: cedarStatement,
+  };
+}

src/cli/cloudformation/__tests__/outputs.test.ts

@@ -1,4 +1,10 @@
-import { buildDeployedState, parseGatewayOutputs, parseMemoryOutputs } from '../outputs';
+import {
+  buildDeployedState,
+  parseGatewayOutputs,
+  parseMemoryOutputs,
+  parsePolicyEngineOutputs,
+  parsePolicyOutputs,
+} from '../outputs';
 import { describe, expect, it } from 'vitest';
 
 describe('buildDeployedState', () => {
@@ -285,3 +291,181 @@ describe('parseMemoryOutputs', () => {
     expect(result).toEqual({});
   });
 });
+
+describe('parsePolicyEngineOutputs', () => {
+  it('extracts policy engine outputs matching pattern', () => {
+    const outputs = {
+      ApplicationPolicyEngineMyEngineIdOutputABC123: 'pe-123',
+      ApplicationPolicyEngineMyEngineArnOutputDEF456: 'arn:aws:bedrock:us-east-1:123456789012:policy-engine/pe-123',
+      UnrelatedOutput: 'some-value',
+    };
+
+    const result = parsePolicyEngineOutputs(outputs, ['MyEngine']);
+
+    expect(result).toEqual({
+      MyEngine: {
+        policyEngineId: 'pe-123',
+        policyEngineArn: 'arn:aws:bedrock:us-east-1:123456789012:policy-engine/pe-123',
+      },
+    });
+  });
+
+  it('handles multiple policy engines', () => {
+    const outputs = {
+      ApplicationPolicyEngineFirstEngineIdOutput123: 'pe-1',
+      ApplicationPolicyEngineFirstEngineArnOutput123: 'arn:pe-1',
+      ApplicationPolicyEngineSecondEngineIdOutput456: 'pe-2',
+      ApplicationPolicyEngineSecondEngineArnOutput456: 'arn:pe-2',
+    };
+
+    const result = parsePolicyEngineOutputs(outputs, ['FirstEngine', 'SecondEngine']);
+
+    expect(Object.keys(result)).toHaveLength(2);
+    expect(result.FirstEngine?.policyEngineId).toBe('pe-1');
+    expect(result.SecondEngine?.policyEngineId).toBe('pe-2');
+  });
+
+  it('returns empty record when no policy engine outputs found', () => {
+    const outputs = {
+      UnrelatedOutput: 'some-value',
+    };
+
+    const result = parsePolicyEngineOutputs(outputs, ['MyEngine']);
+
+    expect(result).toEqual({});
+  });
+
+  it('skips incomplete policy engine outputs (missing ARN)', () => {
+    const outputs = {
+      ApplicationPolicyEngineMyEngineIdOutputABC123: 'pe-123',
+    };
+
+    const result = parsePolicyEngineOutputs(outputs, ['MyEngine']);
+
+    expect(result).toEqual({});
+  });
+});
+
+describe('parsePolicyOutputs', () => {
+  it('extracts policy outputs matching pattern', () => {
+    const outputs = {
+      ApplicationPolicyMyEngineDenyAllIdOutputABC123: 'pol-123',
+      ApplicationPolicyMyEngineDenyAllArnOutputDEF456: 'arn:aws:bedrock:us-east-1:123456789012:policy/pol-123',
+      UnrelatedOutput: 'some-value',
+    };
+
+    const result = parsePolicyOutputs(outputs, [{ engineName: 'MyEngine', policyName: 'DenyAll' }]);
+
+    expect(result).toEqual({
+      'MyEngine/DenyAll': {
+        policyId: 'pol-123',
+        policyArn: 'arn:aws:bedrock:us-east-1:123456789012:policy/pol-123',
+        engineName: 'MyEngine',
+      },
+    });
+  });
+
+  it('handles multiple policies across engines', () => {
+    const outputs = {
+      ApplicationPolicyEngine1Policy1IdOutput123: 'pol-1',
+      ApplicationPolicyEngine1Policy1ArnOutput123: 'arn:pol-1',
+      ApplicationPolicyEngine1Policy2IdOutput456: 'pol-2',
+      ApplicationPolicyEngine1Policy2ArnOutput456: 'arn:pol-2',
+    };
+
+    const result = parsePolicyOutputs(outputs, [
+      { engineName: 'Engine1', policyName: 'Policy1' },
+      { engineName: 'Engine1', policyName: 'Policy2' },
+    ]);
+
+    expect(Object.keys(result)).toHaveLength(2);
+    expect(result['Engine1/Policy1']?.policyId).toBe('pol-1');
+    expect(result['Engine1/Policy2']?.policyId).toBe('pol-2');
+  });
+
+  it('returns empty record when no policy outputs found', () => {
+    const outputs = {
+      UnrelatedOutput: 'some-value',
+    };
+
+    const result = parsePolicyOutputs(outputs, [{ engineName: 'MyEngine', policyName: 'DenyAll' }]);
+
+    expect(result).toEqual({});
+  });
+});
+
+describe('buildDeployedState with policy data', () => {
+  it('includes policyEngines in deployed state when provided', () => {
+    const policyEngines = {
+      MyEngine: {
+        policyEngineId: 'pe-123',
+        policyEngineArn: 'arn:aws:bedrock:us-east-1:123456789012:policy-engine/pe-123',
+      },
+    };
+
+    const result = buildDeployedState({
+      targetName: 'default',
+      stackName: 'TestStack',
+      agents: {},
+      gateways: {},
+      policyEngines,
+    });
+
+    expect(result.targets.default!.resources?.policyEngines).toEqual(policyEngines);
+  });
+
+  it('includes policies in deployed state when provided', () => {
+    const policies = {
+      'MyEngine/DenyAll': {
+        policyId: 'pol-123',
+        policyArn: 'arn:aws:bedrock:us-east-1:123456789012:policy/pol-123',
+        engineName: 'MyEngine',
+      },
+    };
+
+    const result = buildDeployedState({
+      targetName: 'default',
+      stackName: 'TestStack',
+      agents: {},
+      gateways: {},
+      policies,
+    });
+
+    expect(result.targets.default!.resources?.policies).toEqual(policies);
+  });
+
+  it('omits policyEngines field when policyEngines is empty object', () => {
+    const result = buildDeployedState({
+      targetName: 'default',
+      stackName: 'TestStack',
+      agents: {},
+      gateways: {},
+      policyEngines: {},
+    });
+
+    expect(result.targets.default!.resources?.policyEngines).toBeUndefined();
+  });
+
+  it('omits policies field when policies is empty object', () => {
+    const result = buildDeployedState({
+      targetName: 'default',
+      stackName: 'TestStack',
+      agents: {},
+      gateways: {},
+      policies: {},
+    });
+
+    expect(result.targets.default!.resources?.policies).toBeUndefined();
+  });
+
+  it('omits policyEngines field when not provided', () => {
+    const result = buildDeployedState({
+      targetName: 'default',
+      stackName: 'TestStack',
+      agents: {},
+      gateways: {},
+    });
+
+    expect(result.targets.default!.resources?.policyEngines).toBeUndefined();
+  });
+});