From c0670677714884c9b5ca4c67f768cc5ef1ade226 Mon Sep 17 00:00:00 2001 From: Tejas Kashinath Date: Thu, 27 Aug 2026 14:26:10 -0400 Subject: [PATCH 1/2] Revert "refactor: extract --generate and PolicyClient to a follow-up PR per review" This reverts commit 529b9facf70745f6a27b3c9dcdd89fc1e460b40d. --- src/core/index.tsx | 3 + src/core/policy.test.ts | 133 +++++++++++++++++ src/core/policy.tsx | 137 ++++++++++++++++++ src/handlers/index.tsx | 4 +- .../project/add/gateway-test-support.ts | 4 +- src/handlers/project/add/policy/index.test.ts | 134 ++++++++++++++++- src/handlers/project/add/policy/index.ts | 89 +++++++++++- src/handlers/project/add/policy/types.ts | 23 +++ src/handlers/project/add/types.ts | 2 + src/handlers/project/index.ts | 2 + src/handlers/types.tsx | 2 + src/testing/TestCoreClient.tsx | 24 +++ 12 files changed, 544 insertions(+), 13 deletions(-) create mode 100644 src/core/policy.test.ts create mode 100644 src/core/policy.tsx create mode 100644 src/handlers/project/add/policy/types.ts diff --git a/src/core/index.tsx b/src/core/index.tsx index d7275b912..b6919a866 100644 --- a/src/core/index.tsx +++ b/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" }), diff --git a/src/core/policy.test.ts b/src/core/policy.test.ts new file mode 100644 index 000000000..2e8d69a9c --- /dev/null +++ b/src/core/policy.test.ts @@ -0,0 +1,133 @@ +import { describe, expect, test } from "bun:test"; +import { + GetGatewayCommand, + GetPolicyGenerationCommand, + ListGatewaysCommand, + ListPolicyEngineSummariesCommand, + ListPolicyGenerationAssetsCommand, + StartPolicyGenerationCommand, +} from "@aws-sdk/client-bedrock-agentcore-control"; +import { PolicyClient } from "./policy"; +import { createSilentLogger } from "../testing"; +import type { AwsClients } from "./types"; + +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 HAPPY = { + engines: { + policyEngines: [{ name: "Proj_Guardrails", policyEngineId: "pe-abc123" }], + }, + gateways: { items: [{ name: "Proj-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: "forbid (principal, action, resource);" } }, + findings: [{ type: "VALID", description: "ok" }], + }, + ], + }, +}; + +async function drain(client: PolicyClient, input: Parameters[0]) { + const generator = client.generatePolicy(input, { region: "us-west-2" }); + const messages: string[] = []; + while (true) { + const next = await generator.next(); + if (next.done) return { messages, result: next.value }; + messages.push(next.value.message); + } +} + +describe("PolicyClient.generatePolicy", () => { + const input = { + engineName: "Guardrails", + gatewayName: "tools", + engineServiceName: "Proj_Guardrails", + gatewayServiceName: "Proj-tools", + description: "block hate speech", + }; + + function client(responses: Parameters[0]) { + return new PolicyClient(fakeClients(responses), createSilentLogger(), 0); + } + + test("resolves deployed ids, generates, and returns statement with findings", async () => { + const { messages, result } = await drain(client(HAPPY), input); + + expect(result.statement).toBe("forbid (principal, action, resource);"); + expect(result.findings).toEqual([{ type: "VALID", description: "ok" }]); + expect(messages.some((message) => message.includes("Generating"))).toBe(true); + }); + + test("reads a Dogwood policy definition member", async () => { + const { result } = await drain( + client({ + ...HAPPY, + assets: { + policyGenerationAssets: [ + { definition: { policy: { statement: "forbid (principal, action, resource);" } } }, + ], + }, + }), + input, + ); + expect(result.statement).toBe("forbid (principal, action, resource);"); + expect(result.findings).toEqual([]); + }); + + 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"], + [ + "polling exhausts while still generating", + { ...HAPPY, get: { status: "GENERATING" } }, + "may still complete", + ], + [ + "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", + ], + ])("fails when %s", async (_label, responses, message) => { + await expect(drain(client(responses), input)).rejects.toThrow(message); + }); +}); diff --git a/src/core/policy.tsx b/src/core/policy.tsx new file mode 100644 index 000000000..5f40f9e70 --- /dev/null +++ b/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 ?? "", + })), + }; + } +} diff --git a/src/handlers/index.tsx b/src/handlers/index.tsx index 87def85c6..36712557f 100644 --- a/src/handlers/index.tsx +++ b/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)); diff --git a/src/handlers/project/add/gateway-test-support.ts b/src/handlers/project/add/gateway-test-support.ts index ec7e6c578..4a6fdd10a 100644 --- a/src/handlers/project/add/gateway-test-support.ts +++ b/src/handlers/project/add/gateway-test-support.ts @@ -24,10 +24,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 = 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(), diff --git a/src/handlers/project/add/policy/index.test.ts b/src/handlers/project/add/policy/index.test.ts index 3d13bdb7a..c99d8e467 100644 --- a/src/handlers/project/add/policy/index.test.ts +++ b/src/handlers/project/add/policy/index.test.ts @@ -1,4 +1,5 @@ import { afterEach, describe, expect, test } from "bun:test"; +import { TestCoreClient } from "../../../../testing"; import { createGatewayProjectTestHarness } from "../gateway-test-support"; import { inferAuthorizationPhase } from "./index"; @@ -113,20 +114,149 @@ describe("project add policy", () => { "--name", ], [ - "missing --statement", + "no statement source", ["add", "policy", "--engine", "Guardrails", "--name", "P"], - "required option '--statement", + "one of '--statement' or '--generate'", ], [ "unknown engine", ["add", "policy", "--engine", "Missing", "--name", "P", "--statement", FORBID_ALL], "policy engine 'Missing' does not exist", ], + [ + "--gateway without --generate", + [ + "add", + "policy", + "--engine", + "Guardrails", + "--name", + "P", + "--statement", + FORBID_ALL, + "--gateway", + "tools", + ], + "--gateway is valid only with --generate", + ], ])("rejects %s", async (_label, args, message) => { await withEngine(); await expect(run(args)).rejects.toThrow(message); }); + test("adds a generated policy and surfaces findings", async () => { + const projectRoot = await withEngine(); + await run(["add", "gateway", "--name", "tools"]); + const core = new TestCoreClient(); + core.policy.generateResult = { + statement: SUPPRESS, + findings: [{ type: "VALID", description: "ok" }], + }; + + const io = await run( + [ + "add", + "policy", + "--engine", + "Guardrails", + "--name", + "Gen", + "--generate", + "block hate speech", + "--gateway", + "tools", + ], + undefined, + core, + ); + + expect(core.policy.generateCalls[0]).toMatchObject({ + engineName: "Guardrails", + gatewayName: "tools", + engineServiceName: "TestProject_Guardrails", + gatewayServiceName: "TestProject-tools", + description: "block hate speech", + }); + expect(io.stderr()).toContain("Generated Cedar policy:"); + expect(io.stderr()).toContain("VALID"); + expect((await projectSpec(projectRoot)).policyEngines[0].policies[0]).toMatchObject({ + name: "Gen", + statement: SUPPRESS, + authorizationPhase: "RETURN_OUTPUT", + }); + }); + + test("rejects --generate when multiple gateways exist and none is named", async () => { + await withEngine(); + await run(["add", "gateway", "--name", "tools"]); + await run(["add", "gateway", "--name", "search"]); + await expect( + run(["add", "policy", "--engine", "Guardrails", "--name", "Gen", "--generate", "x"]), + ).rejects.toThrow("multiple gateways: tools, search; pass --gateway"); + }); + + test.each([ + ["--engine", "Missing", [], "does not exist in policyEngines[]"], + ["--gateway", "Guardrails", ["--gateway", "missing"], "does not exist in agentCoreGateways[]"], + ["no gateways in the project via", "Guardrails", [], "add one to this project"], + ])("rejects --generate with an unknown %s", async (_label, engine, gatewayArgs, message) => { + await withEngine(); + await expect( + run([ + "add", + "policy", + "--engine", + engine, + "--name", + "Gen", + "--generate", + "x", + ...gatewayArgs, + ]), + ).rejects.toThrow(message); + }); + + test("rejects a duplicate policy name before generating", async () => { + await withEngine(); + await run(["add", "gateway", "--name", "tools"]); + await run([ + "add", + "policy", + "--engine", + "Guardrails", + "--name", + "DenyAll", + "--statement", + FORBID_ALL, + ]); + const core = new TestCoreClient(); + + await expect( + run( + ["add", "policy", "--engine", "Guardrails", "--name", "DenyAll", "--generate", "x"], + undefined, + core, + ), + ).rejects.toThrow("already exists in policy engine 'Guardrails'"); + expect(core.policy.generateCalls).toEqual([]); + }); + + test("fails without writing when generation fails", async () => { + const projectRoot = await withEngine(); + await run(["add", "gateway", "--name", "tools"]); + const core = new TestCoreClient(); + core.policy.generateError = new Error("policy engine 'Guardrails' is not deployed"); + + await expect( + run( + ["add", "policy", "--engine", "Guardrails", "--name", "Gen", "--generate", "x"], + undefined, + core, + ), + ).rejects.toThrow("is not deployed"); + expect((await projectSpec(projectRoot)).policyEngines[0].policies).toEqual([]); + }); + test("rejects a duplicate policy name across engines", async () => { await withEngine(); await run(["add", "policy-engine", "--name", "Second"]); diff --git a/src/handlers/project/add/policy/index.ts b/src/handlers/project/add/policy/index.ts index 7d245a0a8..3f64e04c0 100644 --- a/src/handlers/project/add/policy/index.ts +++ b/src/handlers/project/add/policy/index.ts @@ -1,8 +1,10 @@ import z from "zod"; import { InputValidationError } from "../../../../errors"; import { SourceResolver } from "../../../../io"; -import type { PolicySchema } from "../../../../projectSchemas/policy"; +import { gatewayResourceName } from "../../../../projectSchemas/gateway"; +import { policyEngineResourceName, type PolicySchema } from "../../../../projectSchemas/policy"; import { createHandler, flag, ProjectKey } from "../../../../router"; +import { coreOptsFromCtx } from "../../../utils"; import type { AddProjectResourceConfig } from "../types"; /** @@ -32,6 +34,12 @@ export const createAddPolicyHandler = (config: AddProjectResourceConfig) => "Cedar policy statement (inline, file://, or - for stdin)", z.string().optional(), ), + flag( + "generate", + "generate the Cedar statement from a natural-language description", + z.string().optional(), + ), + flag("gateway", "deployed Gateway name that scopes --generate", z.string().optional()), flag( "validation-mode", "validation mode: fail-on-any-findings or ignore-all-findings", @@ -55,16 +63,81 @@ export const createAddPolicyHandler = (config: AddProjectResourceConfig) => if (!flags.name) { throw new InputValidationError("required option '--name ' not specified"); } - if (!flags.statement) { - throw new InputValidationError("required option '--statement ' not specified"); + const sources = [flags.statement, flags.generate].filter((value) => value !== undefined); + if (sources.length !== 1) { + throw new InputValidationError("specify exactly one of '--statement' or '--generate'"); + } + if (flags.gateway !== undefined && flags.generate === undefined) { + throw new InputValidationError("--gateway is valid only with --generate"); } const project = ctx.require(ProjectKey); - const source = new SourceResolver({ stdin: config.io.stdin }); - const statement = (await source.resolveText("statement", flags.statement))!; - const sourceFile = flags.statement.startsWith("file://") - ? flags.statement.slice("file://".length) - : undefined; + let statement: string; + let sourceFile: string | undefined; + if (flags.statement !== undefined) { + const source = new SourceResolver({ stdin: config.io.stdin }); + statement = (await source.resolveText("statement", flags.statement))!; + if (flags.statement.startsWith("file://")) { + sourceFile = flags.statement.slice("file://".length); + } + } else { + // Fail before the minute-long generation call; the manager re-checks on write. + if (!project.spec.policyEngines.some((engine) => engine.name === flags.engine)) { + throw new InputValidationError( + `policy engine '${flags.engine}' does not exist in policyEngines[]`, + ); + } + const owner = project.spec.policyEngines.find((engine) => + engine.policies.some((policy) => policy.name === flags.name), + ); + if (owner) { + throw new InputValidationError( + `a policy with name '${flags.name}' already exists in policy engine '${owner.name}'`, + ); + } + const gateways = project.spec.agentCoreGateways; + const gateway = flags.gateway + ? gateways.find((candidate) => candidate.name === flags.gateway) + : gateways.length === 1 + ? gateways[0] + : undefined; + if (flags.gateway && !gateway) { + throw new InputValidationError( + `gateway '${flags.gateway}' does not exist in agentCoreGateways[]`, + ); + } + if (!gateway) { + throw new InputValidationError( + gateways.length === 0 + ? "--generate needs a deployed gateway; add one to this project and deploy it first" + : `this project declares multiple gateways: ${gateways + .map((candidate) => candidate.name) + .join(", ")}; pass --gateway to choose one`, + ); + } + + const generator = config.policy.generatePolicy( + { + engineName: flags.engine, + gatewayName: gateway.name, + engineServiceName: policyEngineResourceName(project.name, flags.engine), + gatewayServiceName: gatewayResourceName(project.name, gateway), + description: flags.generate!, + }, + coreOptsFromCtx(ctx), + ); + let next = await generator.next(); + while (!next.done) { + config.io.stderr.write(`${next.value.message}\n`); + next = await generator.next(); + } + const generated = next.value; + statement = generated.statement; + config.io.stderr.write(`Generated Cedar policy:\n${statement}\n`); + for (const finding of generated.findings) { + config.io.stderr.write(`finding [${finding.type}]: ${finding.description}\n`); + } + } const authorizationPhase = flags["authorization-phase"] ? PHASES[flags["authorization-phase"]] diff --git a/src/handlers/project/add/policy/types.ts b/src/handlers/project/add/policy/types.ts new file mode 100644 index 000000000..7b5620fe4 --- /dev/null +++ b/src/handlers/project/add/policy/types.ts @@ -0,0 +1,23 @@ +import type { CoreOptions } from "../../../../core/types"; + +export type GeneratePolicyInput = { + /** Project-spec names, used in progress and error messages. */ + engineName: string; + gatewayName: string; + /** Exact deployed service names the control-plane lookups match against. */ + engineServiceName: string; + gatewayServiceName: string; + description: string; +}; + +export type GeneratedPolicy = { + statement: string; + findings: { type: string; description: string }[]; +}; + +export interface CorePolicyClient { + generatePolicy( + input: GeneratePolicyInput, + options: CoreOptions, + ): AsyncGenerator<{ message: string }, GeneratedPolicy>; +} diff --git a/src/handlers/project/add/types.ts b/src/handlers/project/add/types.ts index 26943b932..0a78352d5 100644 --- a/src/handlers/project/add/types.ts +++ b/src/handlers/project/add/types.ts @@ -1,7 +1,9 @@ import type { AppIO } from "../../../io"; import type { ProjectManager } from "../types"; +import type { CorePolicyClient } from "./policy/types"; export type AddProjectResourceConfig = { projectManager: ProjectManager; + policy: CorePolicyClient; io: AppIO; }; diff --git a/src/handlers/project/index.ts b/src/handlers/project/index.ts index 2eb42eeaa..3bb94958c 100644 --- a/src/handlers/project/index.ts +++ b/src/handlers/project/index.ts @@ -13,9 +13,11 @@ import { createStatusProjectHandler } from "./status"; import { createBuildProjectHandler } from "./build"; import type { ProjectManager } from "./types"; import { createAddProjectResourceHandler } from "./add"; +import type { CorePolicyClient } from "./add/policy/types"; type ProjectHandlerConfig = { projectManager: ProjectManager; + policy: CorePolicyClient; io: AppIO; }; diff --git a/src/handlers/types.tsx b/src/handlers/types.tsx index f129805a8..c1eec39dc 100644 --- a/src/handlers/types.tsx +++ b/src/handlers/types.tsx @@ -6,6 +6,7 @@ import type { CoreMemoryClient } from "./memory/types.tsx"; import type { CoreRuntimeClient } from "./runtime/types.tsx"; import type { Context } from "../router"; import type { ProjectManager } from "./project/types.ts"; +import type { CorePolicyClient } from "./project/add/policy/types"; export interface Core { harness: CoreHarnessClient; @@ -14,6 +15,7 @@ export interface Core { runtime: CoreRuntimeClient; gateway: CoreGatewayClient; eval: CoreEvalClient; + policy: CorePolicyClient; projectManager: ProjectManager; } diff --git a/src/testing/TestCoreClient.tsx b/src/testing/TestCoreClient.tsx index e58ef5209..d1d16c1fd 100644 --- a/src/testing/TestCoreClient.tsx +++ b/src/testing/TestCoreClient.tsx @@ -161,6 +161,11 @@ import { isTerminalStatus } from "../core/batchEvaluationResults"; import { abortable } from "../core/abortable"; import type { CoreOptions } from "../core/types"; import type { ProjectManager } from "../handlers/project/types"; +import type { + CorePolicyClient, + GeneratedPolicy, + GeneratePolicyInput, +} from "../handlers/project/add/policy/types"; import type { Logger } from "../logging"; import type { ReadWriteJson } from "../io"; import { createSilentLogger } from "./logging"; @@ -2221,6 +2226,24 @@ export class TestEvalClient implements CoreEvalClient { } } +export class TestPolicyClient implements CorePolicyClient { + generateResult: GeneratedPolicy = { + statement: "forbid (principal, action, resource);", + findings: [], + }; + generateError: Error | undefined; + generateCalls: GeneratePolicyInput[] = []; + + async *generatePolicy( + input: GeneratePolicyInput, + ): AsyncGenerator<{ message: string }, GeneratedPolicy> { + this.generateCalls.push(input); + if (this.generateError) throw this.generateError; + yield { message: "generating" }; + return this.generateResult; + } +} + // TestCoreClient implements the Core contract with fully controllable sub-clients. export class TestCoreClient implements Core { readonly harness = new TestHarnessClient(); @@ -2229,6 +2252,7 @@ export class TestCoreClient implements Core { readonly runtime = new TestRuntimeClient(); readonly gateway = new TestGatewayClient(); readonly eval = new TestEvalClient(); + readonly policy = new TestPolicyClient(); readonly projectManager: ProjectManager; // Commands the project manager would have run (npm install, git init, ...), From 20de14092415dc1237dd3e3133bea0f632023df5 Mon Sep 17 00:00:00 2001 From: Tejas Kashinath Date: Thu, 27 Aug 2026 14:44:30 -0400 Subject: [PATCH 2/2] test: cover generate edges through the handler per review --- src/core/policy.test.ts | 133 --------------- .../project/add/gateway-test-support.ts | 3 +- .../project/add/policy/generate.test.ts | 153 ++++++++++++++++++ src/handlers/project/add/policy/index.ts | 3 +- 4 files changed, 157 insertions(+), 135 deletions(-) delete mode 100644 src/core/policy.test.ts create mode 100644 src/handlers/project/add/policy/generate.test.ts diff --git a/src/core/policy.test.ts b/src/core/policy.test.ts deleted file mode 100644 index 2e8d69a9c..000000000 --- a/src/core/policy.test.ts +++ /dev/null @@ -1,133 +0,0 @@ -import { describe, expect, test } from "bun:test"; -import { - GetGatewayCommand, - GetPolicyGenerationCommand, - ListGatewaysCommand, - ListPolicyEngineSummariesCommand, - ListPolicyGenerationAssetsCommand, - StartPolicyGenerationCommand, -} from "@aws-sdk/client-bedrock-agentcore-control"; -import { PolicyClient } from "./policy"; -import { createSilentLogger } from "../testing"; -import type { AwsClients } from "./types"; - -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 HAPPY = { - engines: { - policyEngines: [{ name: "Proj_Guardrails", policyEngineId: "pe-abc123" }], - }, - gateways: { items: [{ name: "Proj-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: "forbid (principal, action, resource);" } }, - findings: [{ type: "VALID", description: "ok" }], - }, - ], - }, -}; - -async function drain(client: PolicyClient, input: Parameters[0]) { - const generator = client.generatePolicy(input, { region: "us-west-2" }); - const messages: string[] = []; - while (true) { - const next = await generator.next(); - if (next.done) return { messages, result: next.value }; - messages.push(next.value.message); - } -} - -describe("PolicyClient.generatePolicy", () => { - const input = { - engineName: "Guardrails", - gatewayName: "tools", - engineServiceName: "Proj_Guardrails", - gatewayServiceName: "Proj-tools", - description: "block hate speech", - }; - - function client(responses: Parameters[0]) { - return new PolicyClient(fakeClients(responses), createSilentLogger(), 0); - } - - test("resolves deployed ids, generates, and returns statement with findings", async () => { - const { messages, result } = await drain(client(HAPPY), input); - - expect(result.statement).toBe("forbid (principal, action, resource);"); - expect(result.findings).toEqual([{ type: "VALID", description: "ok" }]); - expect(messages.some((message) => message.includes("Generating"))).toBe(true); - }); - - test("reads a Dogwood policy definition member", async () => { - const { result } = await drain( - client({ - ...HAPPY, - assets: { - policyGenerationAssets: [ - { definition: { policy: { statement: "forbid (principal, action, resource);" } } }, - ], - }, - }), - input, - ); - expect(result.statement).toBe("forbid (principal, action, resource);"); - expect(result.findings).toEqual([]); - }); - - 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"], - [ - "polling exhausts while still generating", - { ...HAPPY, get: { status: "GENERATING" } }, - "may still complete", - ], - [ - "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", - ], - ])("fails when %s", async (_label, responses, message) => { - await expect(drain(client(responses), input)).rejects.toThrow(message); - }); -}); diff --git a/src/handlers/project/add/gateway-test-support.ts b/src/handlers/project/add/gateway-test-support.ts index 4a6fdd10a..bec15119e 100644 --- a/src/handlers/project/add/gateway-test-support.ts +++ b/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,7 +25,7 @@ export function createGatewayProjectTestHarness(directoryPrefix: string) { const originalCwd = process.cwd(); const tempDirectories: string[] = []; - async function run(args: string[], stdin?: string, core = new TestCoreClient()) { + 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(core, { diff --git a/src/handlers/project/add/policy/generate.test.ts b/src/handlers/project/add/policy/generate.test.ts new file mode 100644 index 000000000..50fd26d75 --- /dev/null +++ b/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[0]) { + return { + ...new TestCoreClient(), + policy: new PolicyClient(fakeClients(responses), createSilentLogger(), 0), + }; +} + +async function generate(responses: Parameters[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); + }); +}); diff --git a/src/handlers/project/add/policy/index.ts b/src/handlers/project/add/policy/index.ts index 3f64e04c0..83ff192af 100644 --- a/src/handlers/project/add/policy/index.ts +++ b/src/handlers/project/add/policy/index.ts @@ -2,7 +2,8 @@ import z from "zod"; import { InputValidationError } from "../../../../errors"; import { SourceResolver } from "../../../../io"; import { gatewayResourceName } from "../../../../projectSchemas/gateway"; -import { policyEngineResourceName, type PolicySchema } from "../../../../projectSchemas/policy"; +import type { PolicySchema } from "../../../../projectSchemas/policy"; +import { policyEngineResourceName } from "../policy-engine"; import { createHandler, flag, ProjectKey } from "../../../../router"; import { coreOptsFromCtx } from "../../../utils"; import type { AddProjectResourceConfig } from "../types";