feat: generate Cedar policies from natural language in project add policy

src/core/index.tsx

@@ -7,6 +7,7 @@ import { GatewayClient } from "./gateway";
 import { HarnessClient } from "./harness";
 import { IdentityClient } from "./identity";
 import { MemoryClient } from "./memory";
+import { PolicyClient } from "./policy";
 import { RuntimeClient } from "./runtime";
 import type {
   AwsClients,
@@ -64,6 +65,7 @@ export class CoreClient implements AwsClients {
   readonly runtime: RuntimeClient;
   readonly gateway: GatewayClient;
   readonly eval: EvalClient;
+  readonly policy: PolicyClient;
 
   readonly projectManager: ProjectManager;
 
@@ -85,6 +87,7 @@ export class CoreClient implements AwsClients {
       this.logger.child({ module: "eval" }),
       config.newSessionId,
     );
+    this.policy = new PolicyClient(this, this.logger.child({ module: "policy" }));
 
     this.projectManager = new FsProjectManager({
       logger: this.logger.child({ module: "projectManager" }),

src/core/policy.tsx

@@ -0,0 +1,137 @@
+import { setTimeout as sleep } from "node:timers/promises";
+import {
+  GetGatewayCommand,
+  GetPolicyGenerationCommand,
+  ListGatewaysCommand,
+  ListPolicyEngineSummariesCommand,
+  ListPolicyGenerationAssetsCommand,
+  StartPolicyGenerationCommand,
+  type GatewaySummary,
+  type PolicyEngineSummary,
+} from "@aws-sdk/client-bedrock-agentcore-control";
+import { AgentCoreCLIError, ResourceNotFoundError } from "../errors";
+import type {
+  CorePolicyClient,
+  GeneratedPolicy,
+  GeneratePolicyInput,
+} from "../handlers/project/add/policy/types";
+import type { Logger } from "../logging";
+import type { AwsClients, CoreOptions } from "./types";
+import { toClientConfig } from "./utils";
+
+const GENERATION_POLL_DELAY_MS = 3_000;
+const GENERATION_MAX_POLLS = 40;
+
+export class PolicyClient implements CorePolicyClient {
+  constructor(
+    private readonly clients: AwsClients,
+    private readonly logger: Logger,
+    private readonly pollDelayMs = GENERATION_POLL_DELAY_MS,
+  ) {}
+
+  async *generatePolicy(
+    input: GeneratePolicyInput,
+    options: CoreOptions,
+  ): AsyncGenerator<{ message: string }, GeneratedPolicy> {
+    const control = this.clients.control(toClientConfig(options));
+
+    yield { message: `Resolving deployed policy engine '${input.engineName}'` };
+    let engine: PolicyEngineSummary | undefined;
+    let engineToken: string | undefined;
+    do {
+      const page = await control.send(
+        new ListPolicyEngineSummariesCommand({ nextToken: engineToken }),
+      );
+      engine = page.policyEngines?.find((candidate) => candidate.name === input.engineServiceName);
+      engineToken = page.nextToken;
+    } while (!engine && engineToken);
+    if (!engine?.policyEngineId) {
+      throw new ResourceNotFoundError(
+        `policy engine '${input.engineName}' is not deployed; run 'agentcore project deploy' first`,
+      );
+    }
+
+    yield { message: `Resolving deployed gateway '${input.gatewayName}'` };
+    let deployed: GatewaySummary | undefined;
+    let gatewayToken: string | undefined;
+    do {
+      const page = await control.send(new ListGatewaysCommand({ nextToken: gatewayToken }));
+      deployed = page.items?.find((candidate) => candidate.name === input.gatewayServiceName);
+      gatewayToken = page.nextToken;
+    } while (!deployed && gatewayToken);
+    if (!deployed) {
+      throw new ResourceNotFoundError(
+        `gateway '${input.gatewayName}' is not deployed; run 'agentcore project deploy' first`,
+      );
+    }
+    const gateway = await control.send(
+      new GetGatewayCommand({ gatewayIdentifier: deployed.gatewayId }),
+    );
+    if (!gateway.gatewayArn) {
+      throw new AgentCoreCLIError(`could not resolve the ARN of gateway '${input.gatewayName}'`);
+    }
+
+    yield { message: "Generating a Cedar policy from the description (may take a minute)" };
+    const started = await control.send(
+      new StartPolicyGenerationCommand({
+        policyEngineId: engine.policyEngineId,
+        resource: { arn: gateway.gatewayArn },
+        content: { rawText: input.description },
+        name: `cli_generation_${Date.now()}`,
+      }),
+    );
+    if (!started.policyGenerationId) {
+      throw new AgentCoreCLIError("StartPolicyGeneration returned no generation id");
+    }
+
+    let status: string | undefined = "GENERATING";
+    let statusReasons: string[] | undefined;
+    for (let poll = 0; poll < GENERATION_MAX_POLLS && status === "GENERATING"; poll++) {
+      await sleep(this.pollDelayMs);
+      const current = await control.send(
+        new GetPolicyGenerationCommand({
+          policyGenerationId: started.policyGenerationId,
+          policyEngineId: engine.policyEngineId,
+        }),
+      );
+      status = current.status;
+      statusReasons = current.statusReasons;
+      this.logger.debug(`policy generation ${started.policyGenerationId} status: ${status}`);
+      if (status === "GENERATING") yield { message: "Still generating" };
+    }
+    if (status !== "GENERATED") {
+      throw new AgentCoreCLIError(
+        status === "GENERATING"
+          ? "policy generation did not finish within the CLI's wait window; it may still complete, retry the command in a few minutes"
+          : `policy generation did not complete: ${statusReasons?.join(", ") ?? status}`,
+      );
+    }
+
+    const assets = await control.send(
+      new ListPolicyGenerationAssetsCommand({
+        policyGenerationId: started.policyGenerationId,
+        policyEngineId: engine.policyEngineId,
+      }),
+    );
+    const asset = assets.policyGenerationAssets?.[0];
+    // The service returns either plain Cedar or its Dogwood superset member.
+    const statement = asset?.definition?.cedar?.statement ?? asset?.definition?.policy?.statement;
+    if (!asset || !statement) {
+      const findings = (asset?.findings ?? [])
+        .map((finding) => `[${finding.type}] ${finding.description}`)
+        .join("; ");
+      throw new AgentCoreCLIError(
+        findings
+          ? `the description could not be translated into a Cedar policy: ${findings}`
+          : "generation completed but returned no generated policy statement",
+      );
+    }
+    return {
+      statement,
+      findings: (asset.findings ?? []).map((finding) => ({
+        type: finding.type ?? "UNKNOWN",
+        description: finding.description ?? "",
+      })),
+    };
+  }
+}

src/handlers/index.tsx

@@ -50,7 +50,9 @@ export function createRootHandler(core: Core, config: RootHandlerConfig): Router
   root.handler(createGatewayHandler(core, io));
   root.handler(createEvalHandler(core, io));
   root.handler(createConfigHandler());
-  root.handler(createProjectHandler({ projectManager: core.projectManager, io }));
+  root.handler(
+    createProjectHandler({ projectManager: core.projectManager, policy: core.policy, io }),
+  );
 
   // Invoking with no subcommand launches the interactive TUI.
   root.default(renderTui(core, io));

src/handlers/project/add/gateway-test-support.ts

@@ -2,6 +2,7 @@ import { mkdtemp, rm } from "node:fs/promises";
 import { join } from "node:path";
 import { tmpdir } from "node:os";
 import { createRootHandler } from "../../index";
+import type { Core } from "../../types";
 import {
   createSilentLogger,
   TestCoreClient,
@@ -24,10 +25,10 @@ export function createGatewayProjectTestHarness(directoryPrefix: string) {
   const originalCwd = process.cwd();
   const tempDirectories: string[] = [];
 
-  async function run(args: string[], stdin?: string) {
+  async function run(args: string[], stdin?: string, core: Core = new TestCoreClient()) {
     const io = testIO();
     if (stdin !== undefined) io.io.stdin.end(stdin);
-    const root = createRootHandler(new TestCoreClient(), {
+    const root = createRootHandler(core, {
       io: io.io,
       globalConfigAccessor: new TestGlobalConfigAccessor(),
       logger: createSilentLogger(),

src/handlers/project/add/policy/generate.test.ts

@@ -0,0 +1,153 @@
+import { afterEach, describe, expect, test } from "bun:test";
+import {
+  GetGatewayCommand,
+  GetPolicyGenerationCommand,
+  ListGatewaysCommand,
+  ListPolicyEngineSummariesCommand,
+  ListPolicyGenerationAssetsCommand,
+  StartPolicyGenerationCommand,
+} from "@aws-sdk/client-bedrock-agentcore-control";
+import { PolicyClient } from "../../../../core/policy";
+import type { AwsClients } from "../../../../core/types";
+import { createSilentLogger, TestCoreClient } from "../../../../testing";
+import { createGatewayProjectTestHarness } from "../gateway-test-support";
+
+const { addGateway, cleanup, inProject, projectSpec, run } =
+  createGatewayProjectTestHarness("policy-generate");
+
+afterEach(cleanup);
+
+/**
+ Command-flow tests for `project add policy --generate`, driven through the real
+ root handler with the real PolicyClient over a control client mocked at .send().
+ These cover the client edges the TestPolicyClient-backed tests in index.test.ts
+ cannot reach: deployed-resource resolution, polling, and asset parsing.
+**/
+
+function fakeClients(responses: {
+  engines?: unknown;
+  gateways?: unknown;
+  getGateway?: unknown;
+  start?: unknown;
+  get?: unknown;
+  assets?: unknown;
+}): AwsClients {
+  const control = {
+    send: async (command: unknown) => {
+      if (command instanceof ListPolicyEngineSummariesCommand) return responses.engines;
+      if (command instanceof ListGatewaysCommand) return responses.gateways;
+      if (command instanceof GetGatewayCommand) return responses.getGateway;
+      if (command instanceof StartPolicyGenerationCommand) return responses.start;
+      if (command instanceof GetPolicyGenerationCommand) return responses.get;
+      if (command instanceof ListPolicyGenerationAssetsCommand) return responses.assets;
+      throw new Error(`unexpected command: ${command?.constructor?.name}`);
+    },
+  };
+  return { control: () => control } as unknown as AwsClients;
+}
+
+const CEDAR = "forbid (principal, action, resource is AgentCore::Gateway);";
+
+const HAPPY = {
+  engines: {
+    policyEngines: [{ name: "TestProject_Guardrails", policyEngineId: "pe-abc123" }],
+  },
+  gateways: { items: [{ name: "TestProject-tools", gatewayId: "gw-1" }] },
+  getGateway: { gatewayArn: "arn:aws:bedrock-agentcore:us-west-2:1:gateway/gw-1" },
+  start: { policyGenerationId: "gen-1" },
+  get: { status: "GENERATED" },
+  assets: {
+    policyGenerationAssets: [
+      {
+        definition: { cedar: { statement: CEDAR } },
+        findings: [{ type: "VALID", description: "ok" }],
+      },
+    ],
+  },
+};
+
+function coreWith(responses: Parameters<typeof fakeClients>[0]) {
+  return {
+    ...new TestCoreClient(),
+    policy: new PolicyClient(fakeClients(responses), createSilentLogger(), 0),
+  };
+}
+
+async function generate(responses: Parameters<typeof fakeClients>[0]) {
+  const projectRoot = await inProject();
+  await run(["add", "policy-engine", "--name", "Guardrails"]);
+  await addGateway("tools");
+  const io = await run(
+    [
+      "add",
+      "policy",
+      "--engine",
+      "Guardrails",
+      "--name",
+      "Gen",
+      "--generate",
+      "forbid everything",
+      "--gateway",
+      "tools",
+    ],
+    undefined,
+    coreWith(responses),
+  );
+  return { projectRoot, io };
+}
+
+describe("project add policy --generate against the control plane", () => {
+  test("resolves deployed ids, prints the Cedar and findings, writes the spec", async () => {
+    const { projectRoot, io } = await generate(HAPPY);
+
+    expect(io.stderr()).toContain(`Generated Cedar policy:\n${CEDAR}`);
+    expect(io.stderr()).toContain("finding [VALID]: ok");
+    expect((await projectSpec(projectRoot)).policyEngines[0].policies[0]).toMatchObject({
+      name: "Gen",
+      statement: CEDAR,
+    });
+  });
+
+  test("reads a Dogwood policy definition member", async () => {
+    const { projectRoot } = await generate({
+      ...HAPPY,
+      assets: { policyGenerationAssets: [{ definition: { policy: { statement: CEDAR } } }] },
+    });
+    expect((await projectSpec(projectRoot)).policyEngines[0].policies[0]).toMatchObject({
+      statement: CEDAR,
+    });
+  });
+
+  test.each([
+    ["engine not deployed", { ...HAPPY, engines: { policyEngines: [] } }, "is not deployed"],
+    ["gateway not deployed", { ...HAPPY, gateways: { items: [] } }, "not deployed"],
+    [
+      "generation failed",
+      { ...HAPPY, get: { status: "GENERATE_FAILED", statusReasons: ["bad input"] } },
+      "bad input",
+    ],
+    ["no assets", { ...HAPPY, assets: { policyGenerationAssets: [] } }, "no generated policy"],
+    [
+      "the description is not translatable",
+      {
+        ...HAPPY,
+        assets: {
+          policyGenerationAssets: [
+            {
+              rawTextFragment: "do the thing",
+              findings: [{ type: "INVALID", description: "Non-translatable" }],
+            },
+          ],
+        },
+      },
+      "could not be translated into a Cedar policy: [INVALID] Non-translatable",
+    ],
+    [
+      "polling exhausts while still generating",
+      { ...HAPPY, get: { status: "GENERATING" } },
+      "may still complete",
+    ],
+  ])("fails when %s", async (_label, responses, message) => {
+    await expect(generate(responses)).rejects.toThrow(message);
+  });
+});