From 72b94a264977d10a06ad23b0af093ee2691a433a Mon Sep 17 00:00:00 2001
From: Aidan Daly <aidandal@amazon.com>
Date: Tue, 18 Aug 2026 14:54:37 +0000
Subject: [PATCH] fix(policy): support generated policy definitions

---
 .../aws/__tests__/policy-generation.test.ts   | 103 ++++++++++++++++++
 src/cli/aws/policy-generation.ts              |  18 ++-
 2 files changed, 115 insertions(+), 6 deletions(-)
 create mode 100644 src/cli/aws/__tests__/policy-generation.test.ts

diff --git a/src/cli/aws/__tests__/policy-generation.test.ts b/src/cli/aws/__tests__/policy-generation.test.ts
new file mode 100644
index 000000000..05ac73baf
--- /dev/null
+++ b/src/cli/aws/__tests__/policy-generation.test.ts
@@ -0,0 +1,103 @@
+import { getPolicyGeneration } from '../policy-generation';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+const { mockSend, mockWaitUntilPolicyGenerationCompleted } = vi.hoisted(() => ({
+  mockSend: vi.fn(),
+  mockWaitUntilPolicyGenerationCompleted: vi.fn(),
+}));
+
+vi.mock('@aws-sdk/client-bedrock-agentcore-control', () => ({
+  BedrockAgentCoreControlClient: class {
+    send = mockSend;
+  },
+  GetPolicyGenerationCommand: class {
+    constructor(public input: unknown) {}
+  },
+  ListPolicyGenerationAssetsCommand: class {
+    constructor(public input: unknown) {}
+  },
+  StartPolicyGenerationCommand: class {
+    constructor(public input: unknown) {}
+  },
+  waitUntilPolicyGenerationCompleted: mockWaitUntilPolicyGenerationCompleted,
+}));
+
+vi.mock('../account', () => ({
+  getCredentialProvider: vi.fn().mockReturnValue({}),
+}));
+
+describe('getPolicyGeneration', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockWaitUntilPolicyGenerationCompleted.mockResolvedValue({ state: 'SUCCESS' });
+  });
+
+  it('returns a generated Cedar definition statement', async () => {
+    mockSend.mockResolvedValueOnce({ status: 'GENERATED' }).mockResolvedValueOnce({
+      policyGenerationAssets: [
+        {
+          definition: {
+            cedar: { statement: 'permit(principal, action, resource);' },
+          },
+        },
+      ],
+    });
+
+    await expect(
+      getPolicyGeneration({
+        generationId: 'generation-1',
+        policyEngineId: 'engine-1',
+        region: 'us-east-1',
+      })
+    ).resolves.toEqual({
+      status: 'GENERATED',
+      statement: 'permit(principal, action, resource);',
+    });
+  });
+
+  it('returns a generated Policy definition statement', async () => {
+    mockSend.mockResolvedValueOnce({ status: 'GENERATED' }).mockResolvedValueOnce({
+      policyGenerationAssets: [
+        {
+          definition: {
+            policy: { statement: 'forbid(principal, action, resource);' },
+          },
+        },
+      ],
+    });
+
+    await expect(
+      getPolicyGeneration({
+        generationId: 'generation-1',
+        policyEngineId: 'engine-1',
+        region: 'us-west-2',
+      })
+    ).resolves.toEqual({
+      status: 'GENERATED',
+      statement: 'forbid(principal, action, resource);',
+    });
+  });
+
+  it('surfaces findings when generation produces no statement', async () => {
+    mockSend.mockResolvedValueOnce({ status: 'GENERATED' }).mockResolvedValueOnce({
+      policyGenerationAssets: [
+        {
+          findings: [
+            {
+              type: 'INVALID',
+              description: 'Non-translatable: cannot be expressed in Dogwood',
+            },
+          ],
+        },
+      ],
+    });
+
+    await expect(
+      getPolicyGeneration({
+        generationId: 'generation-1',
+        policyEngineId: 'engine-1',
+        region: 'us-west-2',
+      })
+    ).rejects.toThrow('Non-translatable: cannot be expressed in Dogwood');
+  });
+});
diff --git a/src/cli/aws/policy-generation.ts b/src/cli/aws/policy-generation.ts
index 9ec35fa22..c6e630266 100644
--- a/src/cli/aws/policy-generation.ts
+++ b/src/cli/aws/policy-generation.ts
@@ -102,16 +102,22 @@ export async function getPolicyGeneration(options: GetPolicyGenerationOptions):
     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');
+  const statement = firstAsset.definition?.cedar?.statement ?? firstAsset.definition?.policy?.statement;
+
+  if (!statement) {
+    const findingDetails =
+      firstAsset.findings
+        ?.map(finding => finding.description ?? finding.type)
+        .filter((finding): finding is string => !!finding)
+        .join(', ') ?? '';
+    throw new Error(
+      `Policy generation completed but no policy statement was found in the assets${findingDetails ? `: ${findingDetails}` : ''}`
+    );
   }
 
   return {
     status: statusResponse.status ?? 'GENERATED',
-    statement: cedarStatement,
+    statement,
   };
 }