diff --git a/README.md b/README.md index 6b7f37f70..41a92ba40 100644 --- a/README.md +++ b/README.md @@ -101,7 +101,7 @@ agentcore # interactive TUI │ │ ├── get # get a Rule under a Gateway │ │ └── list # list Rules under a Gateway │ └── policy -│ └── generate # generate Cedar for a Gateway from a natural-language prompt +│ └── generate # generate Cedar for a Gateway from a prompt (TUI when run bare) ├── eval # evaluate and optimize AgentCore agents │ └── evaluator # manage AgentCore evaluators │ ├── llm-as-a-judge # LLM-as-a-Judge evaluators diff --git a/src/components/ErrorPanel.tsx b/src/components/ErrorPanel.tsx new file mode 100644 index 000000000..48add5c14 --- /dev/null +++ b/src/components/ErrorPanel.tsx @@ -0,0 +1,17 @@ +import { Box, Text, useInput } from "ink"; +import { darkTheme } from "./ui/_core.js"; + +const theme = darkTheme; + +export function ErrorPanel({ message, onBack }: { message: string; onBack: () => void }) { + useInput((_input, key) => { + if (key.escape || key.return) onBack(); + }); + + return ( + + ✗ {message} + {" esc returns to the form"} + + ); +} diff --git a/src/components/HarnessWizard.tsx b/src/components/HarnessWizard.tsx index 1e69b460f..a86573ffe 100644 --- a/src/components/HarnessWizard.tsx +++ b/src/components/HarnessWizard.tsx @@ -11,6 +11,7 @@ import type { CreateHarnessInput } from "../handlers/harness/types"; import type { ScreenProps } from "../handlers/types"; import { coreOptsFromCtx } from "../handlers/utils"; import { Layout } from "./Layout"; +import { ErrorPanel } from "./ErrorPanel"; import { FormTextInput } from "./FormTextInput"; import { FormRadioGroup, type FormRadioOption } from "./FormRadioGroup"; import { FormCheckboxMultiSelect, type FormCheckboxOption } from "./FormCheckboxMultiSelect"; @@ -1188,16 +1189,3 @@ function SuccessPanel({ ); } - -function ErrorPanel({ message, onBack }: { message: string; onBack: () => void }) { - useInput((_input, key) => { - if (key.escape || key.return) onBack(); - }); - - return ( - - ✗ {message} - {" esc returns to the form"} - - ); -} diff --git a/src/components/Root.tsx b/src/components/Root.tsx index b05b27bef..6b33779ed 100644 --- a/src/components/Root.tsx +++ b/src/components/Root.tsx @@ -107,6 +107,7 @@ import { GatewayRuleScreen } from "../handlers/gateway/rule/screen.tsx"; import { GatewayRuleListScreen } from "../handlers/gateway/rule/list/screen.tsx"; import { GatewayRuleGetScreen } from "../handlers/gateway/rule/get/screen.tsx"; import { GatewayInvokeScreen } from "../handlers/gateway/invoke/screen.tsx"; +import { GatewayPolicyGenerateScreen } from "../handlers/gateway/policy/screen.tsx"; import { ProjectScreen, ProjectCommandNotImplementedScreen } from "../handlers/project/screen.tsx"; import { ProjectCreateScreen } from "../handlers/project/create/screen.tsx"; import { ProjectInvokePickerScreen } from "../handlers/project/invoke/screen.tsx"; @@ -464,6 +465,18 @@ export function Root({ path, ctx, core, queryClient }: RootProps) { path="agentcore/gateway/rule/get/:gatewayId/:ruleId" element={} /> + } + /> + } + /> + } + /> } /> { const control = this.clients.control(toClientConfig(options)); const gatewayId = resourceIdFromArn(input.gatewayId); yield { type: "step", message: `Resolving gateway ${gatewayId}` }; - const gateway = await control.send(new GetGatewayCommand({ gatewayIdentifier: gatewayId })); + const gateway = await control.send(new GetGatewayCommand({ gatewayIdentifier: gatewayId }), { + abortSignal: signal, + }); const gatewayArn = gateway.gatewayArn!; const engine = input.policyEngineId ?? gateway.policyEngineConfiguration?.arn; if (!engine) { @@ -63,16 +72,20 @@ export class PolicyClient implements CorePolicyClient { content: { rawText: input.prompt }, name: input.name, }), + { abortSignal: signal }, ); const policyGenerationId = started.policyGenerationId!; const meta = { policyGenerationId, policyEngineId }; yield { type: "step", message: "Waiting for generation to complete" }; const waited = await waitForPolicyGenerationCompleted( - { client: control, ...this.wait }, + { client: control, abortSignal: signal, ...this.wait }, { policyEngineId, policyGenerationId }, ); this.logger.debug(`policy generation ${policyGenerationId} waiter state: ${waited.state}`); + if (waited.state === WaiterState.ABORTED) { + throw signal?.reason ?? new UserCancellationError(); + } if (waited.state === WaiterState.TIMEOUT) { throw new NetworkingError( `policy generation '${policyGenerationId}' did not finish within ${this.wait.maxWaitTime}s; ` + @@ -95,6 +108,7 @@ export class PolicyClient implements CorePolicyClient { do { const page = await control.send( new ListPolicyGenerationAssetsCommand({ policyEngineId, policyGenerationId, nextToken }), + { abortSignal: signal }, ); for (const asset of page.policyGenerationAssets ?? []) { policies.push({ diff --git a/src/handlers/gateway/gateway.screen.test.tsx b/src/handlers/gateway/gateway.screen.test.tsx index ece9b52f7..bcf5f75d1 100644 --- a/src/handlers/gateway/gateway.screen.test.tsx +++ b/src/handlers/gateway/gateway.screen.test.tsx @@ -110,7 +110,7 @@ describe("Gateway menu and list", () => { await waitForText(screen.lastFrame, "inspect AgentCore Gateways"); const frame = screen.lastFrame()!; - for (const command of ["get", "list", "invoke", "target", "connector", "rule"]) { + for (const command of ["get", "list", "invoke", "target", "connector", "rule", "policy"]) { expect(frame).toContain(command); } expect(frame).not.toMatch(/\bcreate\b/); diff --git a/src/handlers/gateway/gateway.test.tsx b/src/handlers/gateway/gateway.test.tsx index f61218b05..2c25b1865 100644 --- a/src/handlers/gateway/gateway.test.tsx +++ b/src/handlers/gateway/gateway.test.tsx @@ -6,7 +6,7 @@ import { } from "@aws-sdk/client-bedrock-agentcore-control"; import { PolicyClient } from "../../core/policy"; import type { AwsClients } from "../../core/types"; -import { NetworkingError } from "../../errors"; +import { NetworkingError, UserCancellationError } from "../../errors"; import { createSilentLogger, TestCoreClient, @@ -15,6 +15,9 @@ import { } from "../../testing"; import { compile, isTuiCommandSupported, ValueContext } from "../../router"; import { createRootHandler } from "../index"; +import { PathKey } from "../../router"; +import { JsonKey } from "../keys"; +import { createGeneratePolicyHandler } from "./policy/generate"; import type { Core } from "../types"; const REGION = "us-west-2"; @@ -148,7 +151,7 @@ describe("gateway validation", () => { ["Policy generate gateway", ["gateway", "policy", "generate", "--prompt", "x"], /--gateway-id/], [ "Policy generate prompt", - ["gateway", "policy", "generate", "--gateway-id", GATEWAY_ID], + ["gateway", "policy", "generate", "--gateway-id", GATEWAY_ID, "--json"], /--prompt/, ], ] as const)( @@ -212,9 +215,53 @@ describe("gateway policy generate against a faked control plane", () => { ).rejects.toThrow("policy generation 'gen-1' failed: bad prompt; try again"); }); + test("stops waiting when the signal aborts", async () => { + const generation = coreWith("GENERATING").policy.generatePolicy( + { gatewayId: GATEWAY_ID, prompt: "x", name: "n" }, + { region: REGION }, + AbortSignal.abort(new UserCancellationError()), + ); + await expect( + (async () => { + for await (const _event of generation) { + // drain + } + })(), + ).rejects.toBeInstanceOf(UserCancellationError); + }); + test("times out when the generation keeps running", async () => { const attempt = run(args, coreWith("GENERATING")); await expect(attempt).rejects.toBeInstanceOf(NetworkingError); await expect(attempt).rejects.toThrow("did not finish within 2s"); }, 10_000); }); + +describe("gateway policy generate deep link", () => { + test.each([ + ["opens the form when only --gateway-id is given", { "gateway-id": "gw-1" }, 1], + ["stays headless when --name is also given", { "gateway-id": "gw-1", name: "n" }, 0], + ])("%s", async (_label, flags, expectedRenders) => { + const core = new TestCoreClient(); + let renders = 0; + const handler = createGeneratePolicyHandler( + core, + testIO().io, + async (path, _ctx, renderedCore) => { + renders++; + expect(path).toBe("/agentcore/gateway/policy/generate/gw-1"); + expect(renderedCore).toBe(core); + }, + ); + const ctx = ValueContext.EmptyContext() + .withValue(PathKey, "/agentcore/gateway/policy/generate") + .withValue(JsonKey, false); + + const attempt = handler.handle(ctx, { prompt: undefined, ...flags }, {}); + if (expectedRenders === 0) await expect(attempt).rejects.toThrow(/--prompt/); + else await attempt; + + expect(renders).toBe(expectedRenders); + expect(core.policy.calls).toEqual([]); + }); +}); diff --git a/src/handlers/gateway/index.tsx b/src/handlers/gateway/index.tsx index 62b3374b5..466065be7 100644 --- a/src/handlers/gateway/index.tsx +++ b/src/handlers/gateway/index.tsx @@ -18,7 +18,7 @@ export function createGatewayHandler(core: Core, io: AppIO): Router { return new Router("gateway", "inspect AgentCore Gateways") .use(withTuiOnEmptyFlagsAndArgs(core, io)) .default(renderTui(core, io)) - .supportedTuiCommands("get", "list", "invoke", "target", "connector", "rule") + .supportedTuiCommands("get", "list", "invoke", "target", "connector", "rule", "policy") .handler(createCreateGatewayHandler(core, io)) .handler(createUpdateGatewayHandler(core, io)) .handler(createGetGatewayHandler(core)) diff --git a/src/handlers/gateway/policy/generate.screen.test.tsx b/src/handlers/gateway/policy/generate.screen.test.tsx new file mode 100644 index 000000000..1e3855abc --- /dev/null +++ b/src/handlers/gateway/policy/generate.screen.test.tsx @@ -0,0 +1,125 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import type { GatewaySummary, GetGatewayResponse } from "@aws-sdk/client-bedrock-agentcore-control"; +import { cleanupScreens, renderScreen, TestCoreClient, waitForText } from "../../../testing"; + +afterEach(cleanupScreens); + +const GATEWAY_ID = "gw-1"; +const ENGINE_ARN = "arn:aws:bedrock-agentcore:us-east-1:123456789012:policy-engine/pe-1"; +const FORBID = + "forbid (principal is AgentCore::IamEntity, action, resource is AgentCore::Gateway);"; +const PERMIT = "permit (principal, action, resource is AgentCore::Gateway);"; +const PLACEHOLDER = "Describe what the policy should allow or deny"; + +function coreWith(engineArn: string | undefined): TestCoreClient { + const core = new TestCoreClient(); + const summary: GatewaySummary = { + gatewayId: GATEWAY_ID, + name: "checkout-gateway", + status: "READY", + createdAt: new Date("2026-08-01T01:02:03.000Z"), + updatedAt: new Date("2026-08-02T03:04:05.000Z"), + authorizerType: "AWS_IAM", + }; + core.gateway.setListResponse({ items: [summary] }); + core.gateway.setGetResponse({ + gatewayId: GATEWAY_ID, + name: "checkout-gateway", + gatewayArn: "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/gw-1", + policyEngineConfiguration: engineArn ? { arn: engineArn, mode: "ENFORCE" } : undefined, + } as GetGatewayResponse); + return core; +} + +describe("gateway policy generate screen", () => { + test("picks a gateway, then shows its engine and the prompt", async () => { + const screen = renderScreen("/agentcore/gateway/policy/generate", { + core: coreWith(ENGINE_ARN), + }); + + await waitForText(screen.lastFrame, "checkout-gateway"); + await screen.press("return"); + + await waitForText(screen.lastFrame, ENGINE_ARN); + const frame = screen.lastFrame()!; + expect(frame).toContain(`agentcore → gateway → policy → generate → ${GATEWAY_ID}`); + expect(frame).toContain(PLACEHOLDER); + expect(frame).toContain("[enter] generate"); + }); + + test("explains when the gateway has no engine and offers no prompt", async () => { + const screen = renderScreen(`/agentcore/gateway/policy/generate/${GATEWAY_ID}`, { + core: coreWith(undefined), + }); + + await waitForText(screen.lastFrame, "no Policy Engine attached"); + expect(screen.lastFrame()).not.toContain(PLACEHOLDER); + }); + + test("generates, renders the Cedar and findings, and edits again on e", async () => { + const core = coreWith(ENGINE_ARN); + core.policy.result = { + policyGenerationId: "gen-1", + policyEngineId: "pe-1", + gatewayArn: "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/gw-1", + policies: [ + { + statement: FORBID, + findings: [{ type: "DENY_ALL", description: "denies every request" }], + }, + { statement: PERMIT, findings: [] }, + ], + }; + const screen = renderScreen(`/agentcore/gateway/policy/generate/${GATEWAY_ID}`, { core }); + await waitForText(screen.lastFrame, PLACEHOLDER); + + await screen.write("forbid IAM callers"); + await screen.press("return"); + + await waitForText(screen.lastFrame, PERMIT, 3000); + const frame = screen.lastFrame()!; + expect(frame).toContain(FORBID); + expect(frame).toContain("✓ Resolving gateway"); + expect(frame).toContain("[DENY_ALL] denies every request"); + expect(frame).toContain("[e] edit prompt"); + expect(core.policy.calls[0]).toMatchObject({ + gatewayId: GATEWAY_ID, + prompt: "forbid IAM callers", + }); + expect(core.policy.calls[0]).not.toHaveProperty("policyEngineId"); + expect(core.policy.calls[0]!.name).toMatch(/^cli_generation_\d+$/); + + await screen.write("e"); + await waitForText(screen.lastFrame, "[enter] generate"); + expect(screen.lastFrame()).toContain("forbid IAM callers"); + }); + + test("aborts the run and returns to the picker on esc while generating", async () => { + const core = coreWith(ENGINE_ARN); + core.policy.hang = true; + const screen = renderScreen(`/agentcore/gateway/policy/generate/${GATEWAY_ID}`, { core }); + await waitForText(screen.lastFrame, PLACEHOLDER); + + await screen.write("x"); + await screen.press("return"); + await waitForText(screen.lastFrame, "[esc] cancel"); + + await screen.press("escape"); + await waitForText(screen.lastFrame, "choose a Gateway to generate a policy for"); + expect(core.policy.signals[0]!.aborted).toBe(true); + }); + + test("shows the error and returns to the form on esc", async () => { + const core = coreWith(ENGINE_ARN); + core.policy.error = new Error("policy generation 'gen-1' failed: bad prompt"); + const screen = renderScreen(`/agentcore/gateway/policy/generate/${GATEWAY_ID}`, { core }); + await waitForText(screen.lastFrame, PLACEHOLDER); + + await screen.write("x"); + await screen.press("return"); + + await waitForText(screen.lastFrame, "✗ policy generation 'gen-1' failed: bad prompt", 3000); + await screen.press("escape"); + await waitForText(screen.lastFrame, "[enter] generate"); + }); +}); diff --git a/src/handlers/gateway/policy/generate.tsx b/src/handlers/gateway/policy/generate.tsx index 3557c71fa..11dbdedfa 100644 --- a/src/handlers/gateway/policy/generate.tsx +++ b/src/handlers/gateway/policy/generate.tsx @@ -1,15 +1,19 @@ import z from "zod"; import { InputValidationError } from "../../../errors"; import { type AppIO, SourceResolver } from "../../../io"; -import { createHandler, flag } from "../../../router"; -import { JsonRendererKey } from "../../../tui"; +import { createHandler, flag, PathKey } from "../../../router"; +import { JsonRendererKey, renderTuiAt } from "../../../tui"; import { runWithProgress } from "../../../tui/progress"; import { JsonKey } from "../../keys"; import type { Core } from "../../types"; import { coreOptsFromCtx, renderJsonError } from "../../utils"; import type { PolicyGenerationResult } from "./types"; -export const createGeneratePolicyHandler = (core: Core, io: AppIO) => +export const createGeneratePolicyHandler = ( + core: Core, + io: AppIO, + renderGenerateTui: typeof renderTuiAt = renderTuiAt, +) => createHandler({ name: "generate", description: "generate a Cedar policy for a Gateway from a natural-language prompt", @@ -39,6 +43,21 @@ export const createGeneratePolicyHandler = (core: Core, io: AppIO) => if (!flags["gateway-id"]) { throw new InputValidationError("required option '--gateway-id ' not specified"); } + const jsonOutput = ctx.require(JsonKey); + if ( + flags.prompt === undefined && + !jsonOutput && + flags["policy-engine-id"] === undefined && + flags.name === undefined + ) { + await renderGenerateTui( + `${ctx.require(PathKey)}/${encodeURIComponent(flags["gateway-id"])}`, + ctx, + core, + io, + ); + return; + } if (flags.prompt === undefined) { throw new InputValidationError("required option '--prompt ' not specified"); } @@ -46,7 +65,6 @@ export const createGeneratePolicyHandler = (core: Core, io: AppIO) => "prompt", flags.prompt, ))!; - const jsonOutput = ctx.require(JsonKey); const generation = core.policy.generatePolicy( { diff --git a/src/handlers/gateway/policy/index.tsx b/src/handlers/gateway/policy/index.tsx index 50efe78ec..778685080 100644 --- a/src/handlers/gateway/policy/index.tsx +++ b/src/handlers/gateway/policy/index.tsx @@ -1,10 +1,12 @@ import type { AppIO } from "../../../io"; import { Router } from "../../../router"; +import { renderTui } from "../../../tui"; import type { Core } from "../../types"; import { createGeneratePolicyHandler } from "./generate"; export function createGatewayPolicyHandler(core: Core, io: AppIO): Router { - return new Router("policy", "generate Cedar policies for an AgentCore Gateway").handler( - createGeneratePolicyHandler(core, io), - ); + return new Router("policy", "generate Cedar policies for an AgentCore Gateway") + .default(renderTui(core, io)) + .supportedTuiCommands("generate") + .handler(createGeneratePolicyHandler(core, io)); } diff --git a/src/handlers/gateway/policy/screen.tsx b/src/handlers/gateway/policy/screen.tsx new file mode 100644 index 000000000..e8dc680c6 --- /dev/null +++ b/src/handlers/gateway/policy/screen.tsx @@ -0,0 +1,208 @@ +import { useEffect, useRef, useState } from "react"; +import { Box, Text, useInput } from "ink"; +import { ScrollView, type ScrollViewRef } from "ink-scroll-view"; +import { useQuery } from "@tanstack/react-query"; +import { useNavigate, useParams } from "react-router"; +import { ErrorPanel } from "../../../components/ErrorPanel"; +import { GatewayPicker } from "../../../components/GatewayPicker"; +import { KeyValueTable } from "../../../components/KeyValueTable"; +import { Layout } from "../../../components/Layout"; +import { MultilineInput } from "../../../components/MultilineInput"; +import { Divider } from "../../../components/ui/divider"; +import { Spinner } from "../../../components/ui/spinner"; +import { TaskList, type Task } from "../../../components/ui/task-list"; +import { darkTheme } from "../../../components/ui/_core.js"; +import { UserCancellationError } from "../../../errors"; +import type { ScreenProps } from "../../types"; +import { coreOptsFromCtx } from "../../utils"; +import type { PolicyGenerationResult } from "./types"; + +const theme = darkTheme; +const PICKER_PATH = "/agentcore/gateway/policy/generate"; +const PROMPT_PLACEHOLDER = "Describe what the policy should allow or deny"; + +type Phase = + | { kind: "form" } + | { kind: "running" } + | { kind: "result"; result: PolicyGenerationResult } + | { kind: "error"; message: string }; + +export function GatewayPolicyGenerateScreen(props: ScreenProps) { + const { gatewayId } = useParams(); + const navigate = useNavigate(); + + if (!gatewayId) { + return ( + navigate(`${PICKER_PATH}/${encodeURIComponent(id)}`)} + /> + ); + } + return ; +} + +function finishTasks(tasks: Task[], state: Task["state"]): Task[] { + return tasks.map((task, index) => (index === tasks.length - 1 ? { ...task, state } : task)); +} + +function GeneratePolicyForm({ ctx, core, gatewayId }: ScreenProps & { gatewayId: string }) { + const navigate = useNavigate(); + const opts = coreOptsFromCtx(ctx); + const gateway = useQuery({ + queryKey: ["gateway", opts.region, gatewayId], + queryFn: () => core.gateway.getGateway(gatewayId, opts), + }); + const engineArn = gateway.data?.policyEngineConfiguration?.arn; + + const [phase, setPhase] = useState({ kind: "form" }); + const [prompt, setPrompt] = useState(""); + const [tasks, setTasks] = useState([]); + const scrollRef = useRef(null); + const runRef = useRef<{ + controller: AbortController; + generation: AsyncGenerator; + }>(null); + + const cancel = () => { + const run = runRef.current; + if (!run) return; + runRef.current = null; + run.controller.abort(new UserCancellationError()); + void run.generation.return(undefined as never); + }; + useEffect(() => cancel, []); + + const submit = async () => { + const controller = new AbortController(); + const generation = core.policy.generatePolicy( + { gatewayId, prompt, name: `cli_generation_${Date.now()}` }, + opts, + controller.signal, + ); + runRef.current = { controller, generation }; + setTasks([]); + setPhase({ kind: "running" }); + try { + let next = await generation.next(); + while (!next.done) { + if (controller.signal.aborted) return; + if (next.value.type === "step") { + const title = next.value.message; + setTasks((current) => [ + ...finishTasks(current, "done"), + { title, state: "running", tail: [] }, + ]); + } + next = await generation.next(); + } + if (controller.signal.aborted) return; + setTasks((current) => finishTasks(current, "done")); + setPhase({ kind: "result", result: next.value }); + } catch (error) { + if (controller.signal.aborted) return; + setTasks((current) => finishTasks(current, "failed")); + setPhase({ kind: "error", message: (error as Error).message }); + } finally { + if (runRef.current?.controller === controller) runRef.current = null; + } + }; + + useInput( + (input, key) => { + if (key.escape) { + cancel(); + navigate(PICKER_PATH); + return; + } + if (phase.kind !== "result") return; + if (input === "e") setPhase({ kind: "form" }); + if (key.upArrow || input === "k") scrollRef.current?.scrollBy(-1); + if (key.downArrow || input === "j") scrollRef.current?.scrollBy(1); + }, + { isActive: phase.kind !== "error" }, + ); + + const keyHints = + phase.kind === "form" && engineArn + ? [ + { key: "enter", label: "generate" }, + { key: "⇧↵", label: "newline" }, + { key: "esc", label: "back" }, + { key: "ctl+c", label: "quit" }, + ] + : phase.kind === "result" + ? [ + { key: "↑↓/kj", label: "scroll" }, + { key: "e", label: "edit prompt" }, + { key: "esc", label: "back" }, + { key: "ctl+c", label: "quit" }, + ] + : [ + { key: "esc", label: phase.kind === "running" ? "cancel" : "back" }, + { key: "ctl+c", label: "quit" }, + ]; + + return ( + + {gateway.isPending ? ( + + ) : gateway.isError ? ( + Error: {(gateway.error as Error).message} + ) : ( + + + + {!engineArn ? ( + + This Gateway has no Policy Engine attached. Attach one and deploy, then come back. + + ) : phase.kind === "form" ? ( + void submit()} + placeholder={PROMPT_PLACEHOLDER} + submitDisabled={prompt.trim().length === 0} + /> + ) : phase.kind === "running" ? ( + + ) : phase.kind === "error" ? ( + + + setPhase({ kind: "form" })} /> + + ) : ( + + + + + {phase.result.policies + .flatMap((policy) => (policy.statement ? [policy.statement.trimEnd()] : [])) + .join("\n\n")} + + {phase.result.policies.flatMap((policy, index) => + policy.findings.map((finding, findingIndex) => ( + + [{finding.type}] {finding.description} + + )), + )} + + )} + + )} + + ); +} diff --git a/src/handlers/gateway/policy/types.tsx b/src/handlers/gateway/policy/types.tsx index dabd678fc..13cfecbdc 100644 --- a/src/handlers/gateway/policy/types.tsx +++ b/src/handlers/gateway/policy/types.tsx @@ -27,5 +27,6 @@ export interface CorePolicyClient { generatePolicy( input: GeneratePolicyInput, options: CoreOptions, + signal?: AbortSignal, ): AsyncGenerator; } diff --git a/src/testing/TestCoreClient.tsx b/src/testing/TestCoreClient.tsx index f4e1e013f..43f81ffa8 100644 --- a/src/testing/TestCoreClient.tsx +++ b/src/testing/TestCoreClient.tsx @@ -2404,21 +2404,35 @@ export class TestObservabilityClient implements CoreObservabilityClient { } export class TestPolicyClient implements CorePolicyClient { + result: PolicyGenerationResult = { + policyGenerationId: "gen-1", + policyEngineId: "pe-1", + gatewayArn: "arn:aws:bedrock-agentcore:us-west-2:111122223333:gateway/gw-1", + policies: [ + { statement: "forbid (principal, action, resource is AgentCore::Gateway);", findings: [] }, + ], + }; + error: Error | undefined; + // hang keeps the generator waiting after its first step until the signal aborts. + hang = false; readonly calls: GeneratePolicyInput[] = []; + readonly signals: (AbortSignal | undefined)[] = []; async *generatePolicy( input: GeneratePolicyInput, + _options: CoreOptions, + signal?: AbortSignal, ): AsyncGenerator { this.calls.push(input); - yield { type: "step", message: "Generating policy" }; - return { - policyGenerationId: "gen-1", - policyEngineId: "pe-1", - gatewayArn: "arn:aws:bedrock-agentcore:us-west-2:111122223333:gateway/gw-1", - policies: [ - { statement: "forbid (principal, action, resource is AgentCore::Gateway);", findings: [] }, - ], - }; + this.signals.push(signal); + yield { type: "step", message: "Resolving gateway" }; + if (this.hang) { + await new Promise((_, reject) => + signal?.addEventListener("abort", () => reject(signal.reason), { once: true }), + ); + } + if (this.error) throw this.error; + return this.result; } }