From 82f35eb1bcfeeafbac362a74fcbcfec10f11a159 Mon Sep 17 00:00:00 2001 From: jariy17 Date: Tue, 3 Nov 2026 21:15:20 +0000 Subject: [PATCH 1/3] fix(runtime): report rejected custom JWT invocations --- src/core/core.test.ts | 172 ++++++++++---- src/core/eval.tsx | 19 ++ src/core/invokeRuntime.ts | 210 +++++++++++++++++- .../runtime/invoke/invoke.screen.test.tsx | 31 +++ src/handlers/runtime/invoke/invoke.test.tsx | 59 ++++- src/handlers/runtime/invoke/response.test.ts | 65 ++++++ src/handlers/runtime/invoke/response.ts | 22 +- src/handlers/runtime/invoke/screen.tsx | 12 + src/handlers/runtime/types.tsx | 6 + 9 files changed, 551 insertions(+), 45 deletions(-) diff --git a/src/core/core.test.ts b/src/core/core.test.ts index b4135e9e5c..df7abfe299 100644 --- a/src/core/core.test.ts +++ b/src/core/core.test.ts @@ -7,6 +7,7 @@ import type { IAMClient } from "@aws-sdk/client-iam"; import type { CloudWatchLogsClient } from "@aws-sdk/client-cloudwatch-logs"; import type { XRayClient } from "@aws-sdk/client-xray"; import type { ApplicationSignalsClient } from "@aws-sdk/client-application-signals"; +import { SilentCLIError } from "../errors"; import { GetEventCommand, GetMemoryRecordCommand, @@ -824,25 +825,23 @@ test("CUSTOM_JWT modeled header failures do not expose their values", async () = expect(fetched).toBe(false); }); -test("CUSTOM_JWT non-2xx cancels without reading and exposes status only", async () => { +test("CUSTOM_JWT non-2xx returns safe service details for the response writer", async () => { const { logger, logs } = captureLogs(); - let cancelled = 0; - let read = false; - const response = { - ok: false, - status: 401, - headers: new Headers(), - body: { - cancel: async () => { - cancelled++; - }, - async *[Symbol.asyncIterator]() { - read = true; - yield new TextEncoder().encode("secret response body"); - }, - }, - } as unknown as Response; - const core = customJwtCore(async () => response, { logger }); + const core = customJwtCore( + async () => + new Response( + JSON.stringify({ + __type: "com.amazonaws.bedrockagentcore#ExpiredTokenException", + message: + "Bearer secret-token expired for secret-header-value in runtime-session-value with secret payload; application/json remains useful", + }), + { + status: 401, + headers: { "X-Amzn-RequestId": "request-secret-token" }, + }, + ), + { logger }, + ); const request: RuntimeInvokeRequest = { runtimeId: "runtime-123", accountId: "123456789012", @@ -850,14 +849,23 @@ test("CUSTOM_JWT non-2xx cancels without reading and exposes status only", async payload: new TextEncoder().encode("secret payload"), contentType: "application/json", bearerToken: "secret-token", + runtimeSessionId: "runtime-session-value", applicationHeaders: [["X-Secret", "secret-header-value"]], }; - await expect(core.runtime.invokeRuntime(request, { region: "us-east-1" })).rejects.toThrow( - /^HTTP 401$/, - ); - expect(cancelled).toBe(1); - expect(read).toBe(false); + const response = await core.runtime.invokeRuntime(request, { region: "us-east-1" }); + + expect(response).toMatchObject({ + statusCode: 401, + error: { + detail: + "ExpiredTokenException: Bearer [REDACTED] expired for [REDACTED] in [REDACTED] with [REDACTED]; application/json remains useful", + requestId: "request-[REDACTED]", + }, + }); + const chunks: Uint8Array[] = []; + for await (const chunk of response.body) chunks.push(chunk); + expect(chunks).toEqual([]); expectSafeDebugLog( logs, "Runtime invocation returned a non-success response", @@ -868,11 +876,97 @@ test("CUSTOM_JWT non-2xx cancels without reading and exposes status only", async region: "us-east-1", httpStatusCode: 401, }, - ["secret response body", "secret payload", "secret-token", "secret-header-value"], + ["secret payload", "secret-token", "secret-header-value", "runtime-session-value"], + ); +}); + +test.each([ + [ + "nested JSON", + JSON.stringify({ error: { code: "invalid_token", message: "The token has expired" } }), + "invalid_token: The token has expired", + ], + ["plain text", "Unauthorized\nrequest", "Unauthorized request"], + ["unrecognized JSON", JSON.stringify({ unexpected: "value" }), undefined], +] as const)("CUSTOM_JWT non-2xx extracts %s error details", async (_name, body, detail) => { + const core = customJwtCore(async () => new Response(body, { status: 403 })); + const request: RuntimeInvokeRequest = { + runtimeId: "runtime-123", + accountId: "123456789012", + qualifier: "DEFAULT", + payload: new Uint8Array(), + contentType: "application/json", + bearerToken: "secret-token", + }; + + const response = await core.runtime.invokeRuntime(request, { region: "us-east-1" }); + expect(response).toMatchObject({ + statusCode: 403, + error: detail === undefined ? {} : { detail }, + }); +}); + +test("CUSTOM_JWT non-2xx cancels an oversized error response and exposes status only", async () => { + let cancelled = 0; + const body = new ReadableStream({ + pull(controller) { + controller.enqueue(new Uint8Array(32 * 1024).fill(65)); + }, + cancel() { + cancelled++; + }, + }); + const core = customJwtCore(async () => new Response(body, { status: 502 })); + + const response = await core.runtime.invokeRuntime( + { + runtimeId: "runtime-123", + accountId: "123456789012", + qualifier: "DEFAULT", + payload: new Uint8Array(), + contentType: "application/json", + bearerToken: "secret-token", + }, + { region: "us-east-1" }, + ); + expect(response).toMatchObject({ + statusCode: 502, + error: {}, + }); + expect(cancelled).toBe(1); +}); + +test("CUSTOM_JWT non-2xx preserves cancellation while reading an error response", async () => { + let cancelled = 0; + const body = new ReadableStream({ + cancel() { + cancelled++; + }, + }); + const core = customJwtCore(async () => new Response(body, { status: 401 })); + const controller = new AbortController(); + const reason = new Error("cancelled by user"); + const pending = core.runtime.invokeRuntime( + { + runtimeId: "runtime-123", + accountId: "123456789012", + qualifier: "DEFAULT", + payload: new Uint8Array(), + contentType: "application/json", + bearerToken: "secret-token", + }, + { region: "us-east-1" }, + controller.signal, ); + + await new Promise((resolve) => setTimeout(resolve, 10)); + controller.abort(reason); + + await expect(pending).rejects.toBe(reason); + expect(cancelled).toBe(1); }); -test("CUSTOM_JWT transport failures do not expose arbitrary causes", async () => { +test("CUSTOM_JWT transport failures are visible without exposing arbitrary causes", async () => { const { logger, logs } = captureLogs(); const core = customJwtCore( async () => { @@ -883,20 +977,20 @@ test("CUSTOM_JWT transport failures do not expose arbitrary causes", async () => { logger }, ); - await expect( - core.runtime.invokeRuntime( - { - runtimeId: "runtime-123", - accountId: "123456789012", - qualifier: "DEFAULT", - payload: new TextEncoder().encode("secret payload"), - contentType: "application/json", - bearerToken: "secret-token", - applicationHeaders: [["X-Secret", "secret-header-value"]], - }, - { region: "us-east-1" }, - ), - ).rejects.toThrow(/^Runtime invocation failed$/); + const pending = core.runtime.invokeRuntime( + { + runtimeId: "runtime-123", + accountId: "123456789012", + qualifier: "DEFAULT", + payload: new TextEncoder().encode("secret payload"), + contentType: "application/json", + bearerToken: "secret-token", + applicationHeaders: [["X-Secret", "secret-header-value"]], + }, + { region: "us-east-1" }, + ); + await expect(pending).rejects.toThrow(/^Runtime invocation failed$/); + await expect(pending).rejects.not.toBeInstanceOf(SilentCLIError); expectSafeDebugLog( logs, "Runtime invocation transport failed", diff --git a/src/core/eval.tsx b/src/core/eval.tsx index a82229096b..6bddce850e 100644 --- a/src/core/eval.tsx +++ b/src/core/eval.tsx @@ -928,6 +928,25 @@ export class EvalClient implements CoreEvalClient { options, signal, ); + if (response.error) { + throw new AgentCoreCLIError( + `Runtime invocation failed (HTTP ${response.statusCode})${ + response.error.detail ? `: ${response.error.detail}` : "" + }`, + { + source: + response.statusCode >= 400 && response.statusCode < 500 + ? ERROR_SOURCE.USER + : ERROR_SOURCE.SERVICE, + meta: { + httpStatusCode: response.statusCode, + ...(response.error.requestId !== undefined && { + requestId: response.error.requestId, + }), + }, + }, + ); + } // Read to completion to free the socket; a scripted example ignores the text. let text = ""; const decoder = new TextDecoder(); diff --git a/src/core/invokeRuntime.ts b/src/core/invokeRuntime.ts index 9555f2516a..8d93a83144 100644 --- a/src/core/invokeRuntime.ts +++ b/src/core/invokeRuntime.ts @@ -1,6 +1,6 @@ import { randomUUID } from "node:crypto"; import { InvokeAgentRuntimeCommand } from "@aws-sdk/client-bedrock-agentcore"; -import { InputValidationError, RuntimeInvokeResponseError } from "../errors"; +import { InputValidationError, NetworkingError } from "../errors"; import type { Logger } from "../logging"; import type { AwsClients, CoreFetch, CoreOptions } from "./types"; import { abortable } from "./abortable"; @@ -40,9 +40,15 @@ export type RuntimeInvokeResponse = { traceParent?: string; traceState?: string; baggage?: string; + error?: RuntimeInvokeErrorResponse; body: AsyncIterable; }; +export type RuntimeInvokeErrorResponse = { + detail?: string; + requestId?: string; +}; + // The CUSTOM_JWT path puts the account id in the invocation URL, so a bad ARN must fail here. export function accountIdFromRuntimeArn(arn: string | undefined): string { const id = arn?.match(/^arn:[^:]+:bedrock-agentcore:[^:]*:(\d{12}):runtime\//)?.[1]; @@ -61,6 +67,178 @@ export type InvokeRuntimeDeps = { async function* emptyBody(): AsyncGenerator {} +class RuntimeErrorResponseReader { + private static readonly maxBytes = 64 * 1024; + private static readonly maxDetailCharacters = 1024; + + static forRequest( + response: Response, + request: RuntimeInvokeRequest, + bearerToken: string, + signal?: AbortSignal, + ): RuntimeErrorResponseReader { + return new RuntimeErrorResponseReader( + response, + [ + bearerToken, + ...RuntimeErrorResponseReader.requestValues(request), + RuntimeErrorResponseReader.textPayload(request.payload), + ].filter((value): value is string => value !== undefined && value.length > 0), + signal, + ); + } + + constructor( + private readonly response: Response, + private readonly redactions: string[], + private readonly signal?: AbortSignal, + ) {} + + async detail(): Promise { + const body = await this.readBody(); + const candidate = body === undefined ? this.response.statusText : this.parse(body); + return candidate === undefined ? undefined : this.sanitize(candidate); + } + + header(name: string): string | undefined { + const value = this.response.headers.get(name); + return value === null ? undefined : this.sanitize(value, 256); + } + + private async readBody(): Promise { + if (!this.response.body) return undefined; + const reader = this.response.body.getReader(); + const chunks: Uint8Array[] = []; + let size = 0; + const read = () => { + if (!this.signal) return reader.read(); + this.signal.throwIfAborted(); + const signal = this.signal; + return new Promise>>((resolve, reject) => { + const onAbort = () => reject(signal.reason); + signal.addEventListener("abort", onAbort, { once: true }); + reader + .read() + .then(resolve, reject) + .finally(() => { + signal.removeEventListener("abort", onAbort); + }); + }); + }; + try { + while (true) { + const { done, value } = await read(); + if (done) break; + size += value.byteLength; + if (size > RuntimeErrorResponseReader.maxBytes) { + await reader.cancel().catch(() => undefined); + return undefined; + } + chunks.push(value); + } + return new TextDecoder("utf-8", { fatal: true }).decode(Buffer.concat(chunks)); + } catch (error) { + await reader.cancel().catch(() => undefined); + if (this.signal?.aborted) throw this.signal.reason ?? error; + return undefined; + } finally { + reader.releaseLock(); + } + } + + private parse(body: string): string | undefined { + const trimmed = body.trim(); + if (!trimmed) return undefined; + let parsed: unknown; + try { + parsed = JSON.parse(trimmed); + } catch { + return trimmed; + } + if (typeof parsed === "string") return parsed; + const root = RuntimeErrorResponseReader.record(parsed); + if (!root) return undefined; + const nested = RuntimeErrorResponseReader.record(root.error); + const type = + RuntimeErrorResponseReader.string(root, "__type") ?? + RuntimeErrorResponseReader.string(root, "errorType") ?? + RuntimeErrorResponseReader.string(root, "code") ?? + RuntimeErrorResponseReader.string(root, "type") ?? + RuntimeErrorResponseReader.string(nested, "type") ?? + RuntimeErrorResponseReader.string(nested, "code"); + const message = + RuntimeErrorResponseReader.string(root, "message") ?? + (typeof root.error === "string" ? root.error : undefined) ?? + RuntimeErrorResponseReader.string(root, "detail") ?? + RuntimeErrorResponseReader.string(root, "error_description") ?? + RuntimeErrorResponseReader.string(nested, "message"); + const shortType = type?.split("#").at(-1); + if (shortType && message && shortType !== message) return `${shortType}: ${message}`; + return message ?? shortType; + } + + private sanitize( + candidate: string, + maxCharacters = RuntimeErrorResponseReader.maxDetailCharacters, + ): string | undefined { + let value = candidate; + for (const redaction of this.redactions) { + if (redaction) value = value.split(redaction).join("[REDACTED]"); + } + value = Array.from(value, (character) => { + const codePoint = character.codePointAt(0)!; + return codePoint <= 0x1f || (codePoint >= 0x7f && codePoint <= 0x9f) ? " " : character; + }) + .join("") + .replace(/\s+/g, " ") + .trim(); + if (!value) return undefined; + return value.slice(0, maxCharacters); + } + + private static record(value: unknown): Record | undefined { + return value !== null && typeof value === "object" + ? (value as Record) + : undefined; + } + + private static string( + record: Record | undefined, + key: string, + ): string | undefined { + const value = record?.[key]; + return typeof value === "string" && value.trim() ? value : undefined; + } + + private static requestValues(request: RuntimeInvokeRequest): (string | undefined)[] { + return [ + ...(request.applicationHeaders?.map(([, value]) => value) ?? []), + request.runtimeSessionId, + request.runtimeUserId === "default" ? undefined : request.runtimeUserId, + request.mcpSessionId, + request.mcpProtocolVersion, + request.mcpMethod, + request.mcpName, + request.traceId, + request.traceParent, + request.traceState, + request.baggage, + ]; + } + + private static textPayload(payload: Uint8Array): string | undefined { + if (payload.byteLength < 4 || payload.byteLength > RuntimeErrorResponseReader.maxBytes) { + return undefined; + } + try { + const value = new TextDecoder("utf-8", { fatal: true }).decode(payload); + return value.trim() ? value : undefined; + } catch { + return undefined; + } + } +} + // invokeRuntime dispatches by auth mode: a bearer token routes to the CUSTOM_JWT // (raw fetch) path, otherwise the SigV4 SDK path. export async function invokeRuntime( @@ -152,14 +330,38 @@ async function invokeRuntimeWithCustomJwt( : typeof error, }) .debug("Runtime invocation transport failed"); - throw new RuntimeInvokeResponseError("Runtime invocation failed", error); + throw new NetworkingError("Runtime invocation failed"); } if (!response.ok) { logger .child({ httpStatusCode: response.status }) .debug("Runtime invocation returned a non-success response"); - await response.body?.cancel().catch(() => undefined); - throw new RuntimeInvokeResponseError(`HTTP ${response.status}`); + const errorResponse = RuntimeErrorResponseReader.forRequest( + response, + request, + bearerToken, + signal, + ); + const detail = await errorResponse.detail(); + const requestId = + errorResponse.header("x-amzn-requestid") ?? errorResponse.header("x-amz-request-id"); + return { + statusCode: response.status, + contentType: response.headers.get("content-type") ?? "", + runtimeSessionId: + response.headers.get("x-amzn-bedrock-agentcore-runtime-session-id") ?? undefined, + mcpSessionId: response.headers.get("mcp-session-id") ?? undefined, + mcpProtocolVersion: response.headers.get("mcp-protocol-version") ?? undefined, + traceId: response.headers.get("x-amzn-trace-id") ?? undefined, + traceParent: response.headers.get("traceparent") ?? undefined, + traceState: response.headers.get("tracestate") ?? undefined, + baggage: response.headers.get("baggage") ?? undefined, + error: { + ...(detail !== undefined && { detail }), + ...(requestId !== undefined && { requestId }), + }, + body: emptyBody(), + }; } const body = (response.body as AsyncIterable | null) ?? emptyBody(); return { diff --git a/src/handlers/runtime/invoke/invoke.screen.test.tsx b/src/handlers/runtime/invoke/invoke.screen.test.tsx index f6d1308ad2..638c9e9aa1 100644 --- a/src/handlers/runtime/invoke/invoke.screen.test.tsx +++ b/src/handlers/runtime/invoke/invoke.screen.test.tsx @@ -612,6 +612,37 @@ describe("Runtime invoke JSON console", () => { expect(screen.lastFrame()).toContain("failed · 0 bytes"); }); + test("shows a sanitized rejected response without treating it as streamed output", async () => { + let iterations = 0; + const core = new TestCoreClient(); + core.runtime + .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) + .setInvokeResponse({ + statusCode: 401, + contentType: "application/json", + error: { + detail: "ExpiredTokenException: The token has expired", + requestId: "request-123", + }, + body: (async function* () { + iterations++; + yield Buffer.from("must-not-be-rendered"); + })(), + }); + const screen = renderScreen(CONSOLE_PATH, { core }); + + await waitForText(screen.lastFrame, "Ready"); + await screen.write("{}"); + await screen.press("return"); + + await waitForText(screen.lastFrame, "Runtime invocation failed · HTTP 401"); + expect(screen.lastFrame()).toContain("ExpiredTokenException: The token has expired"); + expect(screen.lastFrame()).toContain("Request ID: request-123"); + expect(screen.lastFrame()).toContain("failed · 0 bytes"); + expect(screen.lastFrame()).not.toContain("must-not-be-rendered"); + expect(iterations).toBe(0); + }); + test("formats modeled AWS service errors with their diagnostics", async () => { const core = new TestCoreClient(); core.runtime.setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse); diff --git a/src/handlers/runtime/invoke/invoke.test.tsx b/src/handlers/runtime/invoke/invoke.test.tsx index 909e82eb35..2097f88439 100644 --- a/src/handlers/runtime/invoke/invoke.test.tsx +++ b/src/handlers/runtime/invoke/invoke.test.tsx @@ -32,12 +32,18 @@ function body(...chunks: Uint8Array[]): AsyncIterable { })(); } -function captureIO(input?: Uint8Array): { io: AppIO; bytes: () => Buffer } { +function captureIO(input?: Uint8Array): { + io: AppIO; + bytes: () => Buffer; + stderrBytes: () => Buffer; +} { const stdin = new PassThrough() as unknown as NodeJS.ReadStream; const stdout = new PassThrough(); const stderr = new PassThrough(); const chunks: Buffer[] = []; + const stderrChunks: Buffer[] = []; stdout.on("data", (chunk) => chunks.push(Buffer.from(chunk))); + stderr.on("data", (chunk) => stderrChunks.push(Buffer.from(chunk))); if (input) stdin.end(input); return { io: { @@ -46,6 +52,7 @@ function captureIO(input?: Uint8Array): { io: AppIO; bytes: () => Buffer } { stderr: stderr as unknown as NodeJS.WriteStream, }, bytes: () => Buffer.concat(chunks), + stderrBytes: () => Buffer.concat(stderrChunks), }; } @@ -769,4 +776,54 @@ describe("runtime invoke", () => { ExitCode.FAILURE, ]); }); + + test("prints rejected CUSTOM_JWT responses before exiting unsuccessfully", async () => { + const core = new TestCoreClient(); + core.runtime.setGetResponse({ + agentRuntimeArn: RUNTIME_ARN, + authorizerConfiguration: { customJWTAuthorizer: {} }, + } as GetAgentRuntimeResponse); + core.runtime.setInvokeResponse({ + statusCode: 401, + contentType: "application/json", + error: { + detail: "ExpiredTokenException: The token has expired", + requestId: "request-123", + }, + body: body(), + }); + const output = captureIO(); + const errors: string[] = []; + const errorLog = spyOn(console, "error").mockImplementation((message) => { + errors.push(String(message)); + }); + + try { + const code = await runWithExitCode(async () => + runCommand(core, output.io, [ + "runtime", + "invoke", + "--id", + RUNTIME_ID, + "--payload", + "{}", + "--bearer-token", + "secret-token", + ]), + ); + + expect(code).toBe(ExitCode.FAILURE); + expect(errors).toEqual([]); + expect(output.bytes()).toHaveLength(0); + expect(output.stderrBytes().toString()).toBe( + "status=401 content-type=application/json runtime-session-id=- mcp-session-id=- " + + "mcp-protocol-version=- trace-id=- trace-parent=- trace-state=- baggage=- " + + "complete=false bytes=0 error=ExpiredTokenException: The token has expired " + + "request-id=request-123\n", + ); + expect(output.stderrBytes().toString()).not.toContain("secret-token"); + } finally { + errorLog.mockRestore(); + } + }); }); diff --git a/src/handlers/runtime/invoke/response.test.ts b/src/handlers/runtime/invoke/response.test.ts index 8b2850ca52..2aefa55257 100644 --- a/src/handlers/runtime/invoke/response.test.ts +++ b/src/handlers/runtime/invoke/response.test.ts @@ -239,6 +239,71 @@ describe("Runtime invoke response output", () => { expect(stderr.bytes()).toHaveLength(0); }); + test.each([false, true])( + "reports a rejected Runtime response once without writing response output (json=%s)", + async (json) => { + let iterations = 0; + const stdout = capture(); + const stderr = capture(); + const result = response({ + statusCode: 401, + contentType: "application/json", + error: { + detail: "ExpiredTokenException: The token has expired", + requestId: "request-123", + }, + body: (async function* () { + iterations++; + yield Buffer.from("must-not-be-written"); + })(), + }); + + const pending = writeRuntimeInvokeResponse(result, { + stdout: stdout.stream, + stderr: stderr.stream, + json, + }); + + await expect(pending).rejects.toBeInstanceOf(SilentCLIError); + await expect(pending).rejects.toThrow("HTTP 401"); + expect(iterations).toBe(0); + expect(stdout.bytes()).toHaveLength(0); + expect(stderr.bytes().toString()).toBe( + "status=401 content-type=application/json runtime-session-id=- mcp-session-id=- " + + "mcp-protocol-version=- trace-id=- trace-parent=- trace-state=- baggage=- " + + "complete=false bytes=0 error=ExpiredTokenException: The token has expired " + + "request-id=request-123\n", + ); + }, + ); + + test("does not create an output file for a rejected Runtime response", async () => { + const file = join(tmpdir(), `runtime-invoke-rejected-${process.pid}-${files.length}`); + files.push(file); + const stdout = capture(); + const stderr = capture(); + + await expect( + writeRuntimeInvokeResponse( + response({ + statusCode: 403, + error: { detail: "AccessDeniedException: Access denied" }, + }), + { + stdout: stdout.stream, + stderr: stderr.stream, + outputFile: file, + }, + ), + ).rejects.toThrow("HTTP 403"); + + expect(await Bun.file(file).exists()).toBe(false); + expect(stdout.bytes()).toHaveLength(0); + expect(stderr.bytes().toString()).toContain( + "complete=false bytes=0 error=AccessDeniedException: Access denied", + ); + }); + test("preserves partial stdout and reports one static incomplete summary", async () => { const stdout = capture(); const stderr = capture(); diff --git a/src/handlers/runtime/invoke/response.ts b/src/handlers/runtime/invoke/response.ts index 9c51911cd0..875efdebbb 100644 --- a/src/handlers/runtime/invoke/response.ts +++ b/src/handlers/runtime/invoke/response.ts @@ -1,3 +1,4 @@ +import { pipeline } from "node:stream/promises"; import { RuntimeInvokeResponseError, UserCancellationError } from "../../../errors"; import { classifyStreamingResponse, @@ -42,14 +43,33 @@ function summary( `mcp-protocol-version=${value(response.mcpProtocolVersion)} ` + `trace-id=${value(response.traceId)} trace-parent=${value(response.traceParent)} ` + `trace-state=${value(response.traceState)} baggage=${value(response.baggage)} ` + - `complete=${complete} bytes=${byteCount}${error ? ` error=${error}` : ""}\n` + `complete=${complete} bytes=${byteCount}${error ? ` error=${error}` : ""}` + + `${response.error?.requestId ? ` request-id=${response.error.requestId}` : ""}\n` ); } +async function writeRejectedRuntimeResponse( + response: RuntimeInvokeResponse, + output: StreamingResponseOutput, +): Promise { + try { + await output.beforeOutput?.(); + await pipeline( + [summary(response, 0, false, response.error?.detail ?? `HTTP ${response.statusCode}`)], + output.stderr, + { end: false, signal: output.signal }, + ); + } catch (error) { + failure(error, output.signal); + } + throw new RuntimeInvokeResponseError(`HTTP ${response.statusCode}`); +} + export async function writeRuntimeInvokeResponse( response: RuntimeInvokeResponse, output: StreamingResponseOutput, ): Promise { + if (response.error) await writeRejectedRuntimeResponse(response, output); await writeStreamingResponse(response, output, { metadata: ({ body: _body, ...metadata }) => metadata, summary, diff --git a/src/handlers/runtime/invoke/screen.tsx b/src/handlers/runtime/invoke/screen.tsx index 6182544269..4e9d90d7d4 100644 --- a/src/handlers/runtime/invoke/screen.tsx +++ b/src/handlers/runtime/invoke/screen.tsx @@ -242,6 +242,18 @@ export function RuntimeInvokeConsole({ ...(mcp && { mcpSessionId, mcpProtocolVersion }), }); const response = await core.runtime.invokeRuntime(request, opts, controller.signal); + if (response.error) { + updateExchange({ + error: { + name: "Runtime invocation failed", + message: response.error.detail, + statusCode: response.statusCode, + requestId: response.error.requestId, + }, + state: "failed", + }); + return; + } updateExchange({ heading: `Response · ${response.statusCode} · ${response.contentType || "-"}`, metadata: metadata(response), diff --git a/src/handlers/runtime/types.tsx b/src/handlers/runtime/types.tsx index 2a5521f019..9bf079bf5a 100644 --- a/src/handlers/runtime/types.tsx +++ b/src/handlers/runtime/types.tsx @@ -42,9 +42,15 @@ export type RuntimeInvokeResponse = { traceParent?: string; traceState?: string; baggage?: string; + error?: RuntimeInvokeErrorResponse; body: AsyncIterable; }; +export type RuntimeInvokeErrorResponse = { + detail?: string; + requestId?: string; +}; + export type RuntimeShellRequest = { runtimeArn: string; qualifier: string; From 6f0ec85fe50b25fd1aaf74c5eeeb0673e21f06da Mon Sep 17 00:00:00 2001 From: jariy17 Date: Tue, 3 Nov 2026 21:30:10 +0000 Subject: [PATCH 2/3] refactor(runtime): distinguish rejected invoke responses --- src/core/core.test.ts | 8 +++-- src/core/dev/localInvoke.ts | 1 + src/core/eval.tsx | 2 +- src/core/invokeRuntime.ts | 22 ++++++++++---- src/handlers/project/invoke/index.test.tsx | 1 + .../runtime/invoke/invoke.screen.test.tsx | 26 ++++++++++++----- src/handlers/runtime/invoke/invoke.test.tsx | 21 ++++++++++++-- src/handlers/runtime/invoke/response.test.ts | 29 ++++++++++++------- src/handlers/runtime/invoke/response.ts | 18 +++++++----- src/handlers/runtime/invoke/screen.tsx | 2 +- src/handlers/runtime/types.tsx | 18 ++++++++---- src/testing/TestCoreClient.tsx | 2 ++ 12 files changed, 107 insertions(+), 43 deletions(-) diff --git a/src/core/core.test.ts b/src/core/core.test.ts index df7abfe299..d14b0c994a 100644 --- a/src/core/core.test.ts +++ b/src/core/core.test.ts @@ -574,6 +574,7 @@ test("invokeRuntime maps all modeled IAM fields and response metadata", async () expect(sent[0]!.options).toEqual({ abortSignal: controller.signal }); const { response, ...metadata } = sdkResponse; expect(result).toEqual({ + kind: "stream", ...metadata, body: response, }); @@ -631,6 +632,7 @@ test("invokeRuntime aborts an established IAM response stream", async () => { { region: "us-east-1" }, controller.signal, ); + if (response.kind !== "stream") throw new Error("expected a streaming response"); const iterator = response.body[Symbol.asyncIterator](); expect(await iterator.next()).toEqual({ done: false, value: Buffer.from("partial") }); @@ -746,6 +748,7 @@ test("CUSTOM_JWT invoke uses the generated endpoint and exact fetch request", as ); expect(result.runtimeSessionId).toBe("returned-runtime"); expect(result.mcpSessionId).toBe("returned-mcp"); + if (result.kind !== "stream") throw new Error("expected a streaming response"); const resultBytes: Uint8Array[] = []; for await (const chunk of result.body) resultBytes.push(chunk); expect(Buffer.concat(resultBytes)).toEqual(Buffer.from([1, 2])); @@ -856,6 +859,7 @@ test("CUSTOM_JWT non-2xx returns safe service details for the response writer", const response = await core.runtime.invokeRuntime(request, { region: "us-east-1" }); expect(response).toMatchObject({ + kind: "rejected", statusCode: 401, error: { detail: @@ -863,9 +867,6 @@ test("CUSTOM_JWT non-2xx returns safe service details for the response writer", requestId: "request-[REDACTED]", }, }); - const chunks: Uint8Array[] = []; - for await (const chunk of response.body) chunks.push(chunk); - expect(chunks).toEqual([]); expectSafeDebugLog( logs, "Runtime invocation returned a non-success response", @@ -1022,6 +1023,7 @@ test("invokeRuntime exposes an empty async iterable when the SDK omits the body" { region: "us-east-1" }, ); + if (response.kind !== "stream") throw new Error("expected a streaming response"); const chunks: Uint8Array[] = []; for await (const chunk of response.body) chunks.push(chunk); expect(chunks).toEqual([]); diff --git a/src/core/dev/localInvoke.ts b/src/core/dev/localInvoke.ts index 92db71231b..7312d18cb9 100644 --- a/src/core/dev/localInvoke.ts +++ b/src/core/dev/localInvoke.ts @@ -86,6 +86,7 @@ export async function invokeLocalRuntime( const body = (response.body as AsyncIterable | null) ?? emptyBody(); return { + kind: "stream", statusCode: response.status, contentType: response.headers.get("content-type") ?? "", runtimeSessionId: diff --git a/src/core/eval.tsx b/src/core/eval.tsx index 6bddce850e..f18d224543 100644 --- a/src/core/eval.tsx +++ b/src/core/eval.tsx @@ -928,7 +928,7 @@ export class EvalClient implements CoreEvalClient { options, signal, ); - if (response.error) { + if (response.kind === "rejected") { throw new AgentCoreCLIError( `Runtime invocation failed (HTTP ${response.statusCode})${ response.error.detail ? `: ${response.error.detail}` : "" diff --git a/src/core/invokeRuntime.ts b/src/core/invokeRuntime.ts index 8d93a83144..ee95b4e262 100644 --- a/src/core/invokeRuntime.ts +++ b/src/core/invokeRuntime.ts @@ -30,7 +30,7 @@ export type RuntimeInvokeRequest = { baggage?: string; }; -export type RuntimeInvokeResponse = { +export type RuntimeInvokeResponseMetadata = { statusCode: number; contentType: string; runtimeSessionId?: string; @@ -40,15 +40,23 @@ export type RuntimeInvokeResponse = { traceParent?: string; traceState?: string; baggage?: string; - error?: RuntimeInvokeErrorResponse; +}; + +export type RuntimeInvokeStreamResponse = RuntimeInvokeResponseMetadata & { + kind: "stream"; body: AsyncIterable; }; -export type RuntimeInvokeErrorResponse = { - detail?: string; - requestId?: string; +export type RuntimeInvokeRejectedResponse = RuntimeInvokeResponseMetadata & { + kind: "rejected"; + error: { + detail?: string; + requestId?: string; + }; }; +export type RuntimeInvokeResponse = RuntimeInvokeStreamResponse | RuntimeInvokeRejectedResponse; + // The CUSTOM_JWT path puts the account id in the invocation URL, so a bad ARN must fail here. export function accountIdFromRuntimeArn(arn: string | undefined): string { const id = arn?.match(/^arn:[^:]+:bedrock-agentcore:[^:]*:(\d{12}):runtime\//)?.[1]; @@ -346,6 +354,7 @@ async function invokeRuntimeWithCustomJwt( const requestId = errorResponse.header("x-amzn-requestid") ?? errorResponse.header("x-amz-request-id"); return { + kind: "rejected", statusCode: response.status, contentType: response.headers.get("content-type") ?? "", runtimeSessionId: @@ -360,11 +369,11 @@ async function invokeRuntimeWithCustomJwt( ...(detail !== undefined && { detail }), ...(requestId !== undefined && { requestId }), }, - body: emptyBody(), }; } const body = (response.body as AsyncIterable | null) ?? emptyBody(); return { + kind: "stream", statusCode: response.status, contentType: response.headers.get("content-type") ?? "", runtimeSessionId: @@ -409,6 +418,7 @@ async function invokeRuntimeWithIam( const body = (response.response as AsyncIterable | undefined) ?? emptyBody(); return { + kind: "stream", statusCode: response.statusCode ?? 0, contentType: response.contentType ?? "", runtimeSessionId: response.runtimeSessionId, diff --git a/src/handlers/project/invoke/index.test.tsx b/src/handlers/project/invoke/index.test.tsx index 2daa77dcfd..ddc40b5f20 100644 --- a/src/handlers/project/invoke/index.test.tsx +++ b/src/handlers/project/invoke/index.test.tsx @@ -124,6 +124,7 @@ async function routedCommand( core.runtime .setGetResponse({ agentRuntimeArn: arn(`runtime/${RUNTIME_ID}`) } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: (async function* () { diff --git a/src/handlers/runtime/invoke/invoke.screen.test.tsx b/src/handlers/runtime/invoke/invoke.screen.test.tsx index 638c9e9aa1..15d965a998 100644 --- a/src/handlers/runtime/invoke/invoke.screen.test.tsx +++ b/src/handlers/runtime/invoke/invoke.screen.test.tsx @@ -207,6 +207,7 @@ describe("Runtime invoke routing", () => { runtimeEndpoints: [endpoint({ name: nextQualifier, id: nextQualifier })], }) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: responseBody(Buffer.from("ok")), @@ -328,6 +329,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: responseBody(Buffer.from("ok")), @@ -371,6 +373,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: responseBody(Buffer.from("ok")), @@ -398,6 +401,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: responseBody(Buffer.from("ok")), @@ -475,6 +479,7 @@ describe("Runtime invoke JSON console", () => { core.runtime.invokeRuntime = async (request) => { requests.push(request); return { + kind: "stream" as const, statusCode: 200, contentType: "text/plain", body: @@ -509,6 +514,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", runtimeSessionId: "returned-runtime", @@ -545,6 +551,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/event-stream", runtimeSessionId: "returned-runtime", @@ -581,6 +588,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType, body: responseBody(body), @@ -613,21 +621,17 @@ describe("Runtime invoke JSON console", () => { }); test("shows a sanitized rejected response without treating it as streamed output", async () => { - let iterations = 0; const core = new TestCoreClient(); core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "rejected", statusCode: 401, contentType: "application/json", error: { detail: "ExpiredTokenException: The token has expired", requestId: "request-123", }, - body: (async function* () { - iterations++; - yield Buffer.from("must-not-be-rendered"); - })(), }); const screen = renderScreen(CONSOLE_PATH, { core }); @@ -639,8 +643,6 @@ describe("Runtime invoke JSON console", () => { expect(screen.lastFrame()).toContain("ExpiredTokenException: The token has expired"); expect(screen.lastFrame()).toContain("Request ID: request-123"); expect(screen.lastFrame()).toContain("failed · 0 bytes"); - expect(screen.lastFrame()).not.toContain("must-not-be-rendered"); - expect(iterations).toBe(0); }); test("formats modeled AWS service errors with their diagnostics", async () => { @@ -689,6 +691,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: (async function* () { @@ -714,6 +717,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: responseBody(Buffer.from(response)), @@ -742,6 +746,7 @@ describe("Runtime invoke JSON console", () => { core.runtime.invokeRuntime = async (request) => { requests.push(request); return { + kind: "stream" as const, statusCode: 200, contentType: "text/event-stream", runtimeSessionId: "returned-runtime", @@ -792,6 +797,7 @@ describe("Runtime invoke JSON console", () => { core.runtime.invokeRuntime = async (request) => { requests.push(request); return { + kind: "stream" as const, statusCode: 200, contentType: "application/json", body: responseBody(Buffer.from('{"ok":true}')), @@ -843,6 +849,7 @@ describe("Runtime invoke JSON console", () => { runtimeEndpoints: [endpoint({ name: nextQualifier, id: nextQualifier })], }) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", runtimeSessionId: "returned-runtime", @@ -881,6 +888,7 @@ describe("Runtime invoke JSON console", () => { expect(screen.lastFrame()).toContain("MCP session ID: Not set"); core.runtime.setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: responseBody(Buffer.from("new response")), @@ -922,6 +930,7 @@ describe("Runtime invoke JSON console", () => { runtimeEndpoints: [endpoint({ name: nextQualifier, id: nextQualifier })], }) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: responseBody(Buffer.from("ok")), @@ -966,6 +975,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "application/json", body: responseBody(Buffer.from(raw)), @@ -990,6 +1000,7 @@ describe("Runtime invoke JSON console", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "application/octet-stream", body: (async function* () { @@ -1034,6 +1045,7 @@ describe("Runtime invoke JSON console", () => { core.runtime.invokeRuntime = async (_request, _options, nextSignal) => { signal = nextSignal; return { + kind: "stream" as const, statusCode: 200, contentType: "text/event-stream", runtimeSessionId: "interrupted-runtime", diff --git a/src/handlers/runtime/invoke/invoke.test.tsx b/src/handlers/runtime/invoke/invoke.test.tsx index 2097f88439..7d21899e08 100644 --- a/src/handlers/runtime/invoke/invoke.test.tsx +++ b/src/handlers/runtime/invoke/invoke.test.tsx @@ -84,6 +84,7 @@ async function run( const core = new TestCoreClient(); core.runtime.setGetResponse(getResponse); core.runtime.setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "application/json", body: body(Uint8Array.from([0, 255]), Uint8Array.from([10, 1])), @@ -105,7 +106,12 @@ describe("runtime invoke", () => { agentRuntimeArn: RUNTIME_ARN, ...(input === "bearer-token" && { authorizerConfiguration: { customJWTAuthorizer: {} } }), } as GetAgentRuntimeResponse) - .setInvokeResponse({ statusCode: 200, contentType: "text/plain", body: stream }); + .setInvokeResponse({ + kind: "stream", + statusCode: 200, + contentType: "text/plain", + body: stream, + }); const io = testIO({ isTTY: true }); const stdin = io.io.stdin as unknown as PassThrough; const pending = runCommand(core, io.io, [ @@ -140,7 +146,12 @@ describe("runtime invoke", () => { const stream = new StreamController(); core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) - .setInvokeResponse({ statusCode: 200, contentType: "text/plain", body: stream }); + .setInvokeResponse({ + kind: "stream", + statusCode: 200, + contentType: "text/plain", + body: stream, + }); const io = testIO({ isTTY: true }); const pending = runCommand(core, io.io, [ "runtime", @@ -300,6 +311,7 @@ describe("runtime invoke", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/event-stream", body: (async function* () { @@ -400,6 +412,7 @@ describe("runtime invoke", () => { core.runtime.invokeRuntime = async (request, options, signal) => { core.runtime.calls.push({ method: "invokeRuntime", args: [request, options, signal] }); return { + kind: "stream" as const, statusCode: 200, contentType: "text/event-stream", body: (async function* () { @@ -503,6 +516,7 @@ describe("runtime invoke", () => { core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "application/octet-stream", body: (async function* () { @@ -734,6 +748,7 @@ describe("runtime invoke", () => { outputCore.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) .setInvokeResponse({ + kind: "stream", statusCode: 200, contentType: "text/plain", body: body(Buffer.from("response")), @@ -784,13 +799,13 @@ describe("runtime invoke", () => { authorizerConfiguration: { customJWTAuthorizer: {} }, } as GetAgentRuntimeResponse); core.runtime.setInvokeResponse({ + kind: "rejected", statusCode: 401, contentType: "application/json", error: { detail: "ExpiredTokenException: The token has expired", requestId: "request-123", }, - body: body(), }); const output = captureIO(); const errors: string[] = []; diff --git a/src/handlers/runtime/invoke/response.test.ts b/src/handlers/runtime/invoke/response.test.ts index 2aefa55257..6b03a2242c 100644 --- a/src/handlers/runtime/invoke/response.test.ts +++ b/src/handlers/runtime/invoke/response.test.ts @@ -5,7 +5,7 @@ import { join } from "node:path"; import { PassThrough, Writable } from "node:stream"; import { SilentCLIError, UserCancellationError } from "../../../errors"; import { waitFor } from "../../../testing"; -import type { RuntimeInvokeResponse } from "../types"; +import type { RuntimeInvokeRejectedResponse, RuntimeInvokeStreamResponse } from "../types"; import { writeRuntimeInvokeResponse } from "./response"; const files: string[] = []; @@ -27,8 +27,11 @@ function failingBody(error: Error): AsyncIterable { })(); } -function response(overrides: Partial = {}): RuntimeInvokeResponse { +function response( + overrides: Partial = {}, +): RuntimeInvokeStreamResponse { return { + kind: "stream", statusCode: 200, contentType: "text/plain", body: body(Buffer.from("ok")), @@ -36,6 +39,18 @@ function response(overrides: Partial = {}): RuntimeInvoke }; } +function rejectedResponse( + overrides: Partial = {}, +): RuntimeInvokeRejectedResponse { + return { + kind: "rejected", + statusCode: 401, + contentType: "application/json", + error: {}, + ...overrides, + }; +} + function capture() { const stream = new PassThrough(); const chunks: Buffer[] = []; @@ -242,20 +257,15 @@ describe("Runtime invoke response output", () => { test.each([false, true])( "reports a rejected Runtime response once without writing response output (json=%s)", async (json) => { - let iterations = 0; const stdout = capture(); const stderr = capture(); - const result = response({ + const result = rejectedResponse({ statusCode: 401, contentType: "application/json", error: { detail: "ExpiredTokenException: The token has expired", requestId: "request-123", }, - body: (async function* () { - iterations++; - yield Buffer.from("must-not-be-written"); - })(), }); const pending = writeRuntimeInvokeResponse(result, { @@ -266,7 +276,6 @@ describe("Runtime invoke response output", () => { await expect(pending).rejects.toBeInstanceOf(SilentCLIError); await expect(pending).rejects.toThrow("HTTP 401"); - expect(iterations).toBe(0); expect(stdout.bytes()).toHaveLength(0); expect(stderr.bytes().toString()).toBe( "status=401 content-type=application/json runtime-session-id=- mcp-session-id=- " + @@ -285,7 +294,7 @@ describe("Runtime invoke response output", () => { await expect( writeRuntimeInvokeResponse( - response({ + rejectedResponse({ statusCode: 403, error: { detail: "AccessDeniedException: Access denied" }, }), diff --git a/src/handlers/runtime/invoke/response.ts b/src/handlers/runtime/invoke/response.ts index 875efdebbb..52e3545049 100644 --- a/src/handlers/runtime/invoke/response.ts +++ b/src/handlers/runtime/invoke/response.ts @@ -6,7 +6,11 @@ import { writeStreamingResponseFile, type StreamingResponseOutput, } from "../../../io"; -import type { RuntimeInvokeResponse } from "../types"; +import type { + RuntimeInvokeRejectedResponse, + RuntimeInvokeResponse, + RuntimeInvokeStreamResponse, +} from "../types"; const RESPONSE_STREAM_FAILED = "response stream failed"; @@ -15,7 +19,7 @@ export function classifyRuntimeResponse(contentType: string) { } export async function writeRuntimeInvokeFile( - response: RuntimeInvokeResponse, + response: RuntimeInvokeStreamResponse, path: string, signal?: AbortSignal, onBytes?: (size: number) => void, @@ -44,18 +48,18 @@ function summary( `trace-id=${value(response.traceId)} trace-parent=${value(response.traceParent)} ` + `trace-state=${value(response.traceState)} baggage=${value(response.baggage)} ` + `complete=${complete} bytes=${byteCount}${error ? ` error=${error}` : ""}` + - `${response.error?.requestId ? ` request-id=${response.error.requestId}` : ""}\n` + `${response.kind === "rejected" && response.error.requestId ? ` request-id=${response.error.requestId}` : ""}\n` ); } async function writeRejectedRuntimeResponse( - response: RuntimeInvokeResponse, + response: RuntimeInvokeRejectedResponse, output: StreamingResponseOutput, ): Promise { try { await output.beforeOutput?.(); await pipeline( - [summary(response, 0, false, response.error?.detail ?? `HTTP ${response.statusCode}`)], + [summary(response, 0, false, response.error.detail ?? `HTTP ${response.statusCode}`)], output.stderr, { end: false, signal: output.signal }, ); @@ -69,9 +73,9 @@ export async function writeRuntimeInvokeResponse( response: RuntimeInvokeResponse, output: StreamingResponseOutput, ): Promise { - if (response.error) await writeRejectedRuntimeResponse(response, output); + if (response.kind === "rejected") return writeRejectedRuntimeResponse(response, output); await writeStreamingResponse(response, output, { - metadata: ({ body: _body, ...metadata }) => metadata, + metadata: ({ body: _body, kind: _kind, ...metadata }) => metadata, summary, fail: (error) => failure(error, output.signal), binaryTtyError: "Binary or unknown response content requires --output-file or --json", diff --git a/src/handlers/runtime/invoke/screen.tsx b/src/handlers/runtime/invoke/screen.tsx index 4e9d90d7d4..69bbfcf2bb 100644 --- a/src/handlers/runtime/invoke/screen.tsx +++ b/src/handlers/runtime/invoke/screen.tsx @@ -242,7 +242,7 @@ export function RuntimeInvokeConsole({ ...(mcp && { mcpSessionId, mcpProtocolVersion }), }); const response = await core.runtime.invokeRuntime(request, opts, controller.signal); - if (response.error) { + if (response.kind === "rejected") { updateExchange({ error: { name: "Runtime invocation failed", diff --git a/src/handlers/runtime/types.tsx b/src/handlers/runtime/types.tsx index 9bf079bf5a..0f083e8fb3 100644 --- a/src/handlers/runtime/types.tsx +++ b/src/handlers/runtime/types.tsx @@ -32,7 +32,7 @@ export type RuntimeInvokeRequest = { baggage?: string; }; -export type RuntimeInvokeResponse = { +export type RuntimeInvokeResponseMetadata = { statusCode: number; contentType: string; runtimeSessionId?: string; @@ -42,15 +42,23 @@ export type RuntimeInvokeResponse = { traceParent?: string; traceState?: string; baggage?: string; - error?: RuntimeInvokeErrorResponse; +}; + +export type RuntimeInvokeStreamResponse = RuntimeInvokeResponseMetadata & { + kind: "stream"; body: AsyncIterable; }; -export type RuntimeInvokeErrorResponse = { - detail?: string; - requestId?: string; +export type RuntimeInvokeRejectedResponse = RuntimeInvokeResponseMetadata & { + kind: "rejected"; + error: { + detail?: string; + requestId?: string; + }; }; +export type RuntimeInvokeResponse = RuntimeInvokeStreamResponse | RuntimeInvokeRejectedResponse; + export type RuntimeShellRequest = { runtimeArn: string; qualifier: string; diff --git a/src/testing/TestCoreClient.tsx b/src/testing/TestCoreClient.tsx index 3000b8483d..ef03d534ab 100644 --- a/src/testing/TestCoreClient.tsx +++ b/src/testing/TestCoreClient.tsx @@ -377,6 +377,7 @@ async function* events(items: T[]): AsyncGenerator { } const DEFAULT_RUNTIME_INVOKE_RESPONSE: RuntimeInvokeResponse = { + kind: "stream", statusCode: 200, contentType: "application/json", body: events([]), @@ -852,6 +853,7 @@ export class TestRuntimeClient implements CoreRuntimeClient { ): Promise { this.calls.push({ method: "invokeRuntime", args: [request, options, signal] }); if (this.error) throw this.error; + if (this.invokeResponse.kind === "rejected") return this.invokeResponse; return { ...this.invokeResponse, body: this.invokeBodies.shift() ?? this.invokeResponse.body, From e1fecae58d0abe9682f2ae884a5ee7b195020937 Mon Sep 17 00:00:00 2001 From: jariy17 Date: Tue, 3 Nov 2026 22:08:36 +0000 Subject: [PATCH 3/3] refactor(runtime): avoid reading rejected response bodies --- src/core/core.test.ts | 104 +--------- src/core/eval.tsx | 27 ++- src/core/invokeRuntime.ts | 186 +----------------- .../runtime/invoke/invoke.screen.test.tsx | 4 +- src/handlers/runtime/invoke/invoke.test.tsx | 4 +- src/handlers/runtime/invoke/response.test.ts | 10 +- src/handlers/runtime/invoke/response.ts | 9 +- src/handlers/runtime/invoke/screen.tsx | 1 - src/handlers/runtime/types.tsx | 1 - 9 files changed, 33 insertions(+), 313 deletions(-) diff --git a/src/core/core.test.ts b/src/core/core.test.ts index d14b0c994a..ed42479add 100644 --- a/src/core/core.test.ts +++ b/src/core/core.test.ts @@ -828,19 +828,20 @@ test("CUSTOM_JWT modeled header failures do not expose their values", async () = expect(fetched).toBe(false); }); -test("CUSTOM_JWT non-2xx returns safe service details for the response writer", async () => { +test("CUSTOM_JWT non-2xx cancels the response body and returns status metadata", async () => { const { logger, logs } = captureLogs(); + let cancelled = 0; const core = customJwtCore( async () => new Response( - JSON.stringify({ - __type: "com.amazonaws.bedrockagentcore#ExpiredTokenException", - message: - "Bearer secret-token expired for secret-header-value in runtime-session-value with secret payload; application/json remains useful", + new ReadableStream({ + cancel() { + cancelled++; + }, }), { status: 401, - headers: { "X-Amzn-RequestId": "request-secret-token" }, + headers: { "X-Amzn-RequestId": "request-123" }, }, ), { logger }, @@ -862,11 +863,10 @@ test("CUSTOM_JWT non-2xx returns safe service details for the response writer", kind: "rejected", statusCode: 401, error: { - detail: - "ExpiredTokenException: Bearer [REDACTED] expired for [REDACTED] in [REDACTED] with [REDACTED]; application/json remains useful", - requestId: "request-[REDACTED]", + requestId: "request-123", }, }); + expect(cancelled).toBe(1); expectSafeDebugLog( logs, "Runtime invocation returned a non-success response", @@ -881,92 +881,6 @@ test("CUSTOM_JWT non-2xx returns safe service details for the response writer", ); }); -test.each([ - [ - "nested JSON", - JSON.stringify({ error: { code: "invalid_token", message: "The token has expired" } }), - "invalid_token: The token has expired", - ], - ["plain text", "Unauthorized\nrequest", "Unauthorized request"], - ["unrecognized JSON", JSON.stringify({ unexpected: "value" }), undefined], -] as const)("CUSTOM_JWT non-2xx extracts %s error details", async (_name, body, detail) => { - const core = customJwtCore(async () => new Response(body, { status: 403 })); - const request: RuntimeInvokeRequest = { - runtimeId: "runtime-123", - accountId: "123456789012", - qualifier: "DEFAULT", - payload: new Uint8Array(), - contentType: "application/json", - bearerToken: "secret-token", - }; - - const response = await core.runtime.invokeRuntime(request, { region: "us-east-1" }); - expect(response).toMatchObject({ - statusCode: 403, - error: detail === undefined ? {} : { detail }, - }); -}); - -test("CUSTOM_JWT non-2xx cancels an oversized error response and exposes status only", async () => { - let cancelled = 0; - const body = new ReadableStream({ - pull(controller) { - controller.enqueue(new Uint8Array(32 * 1024).fill(65)); - }, - cancel() { - cancelled++; - }, - }); - const core = customJwtCore(async () => new Response(body, { status: 502 })); - - const response = await core.runtime.invokeRuntime( - { - runtimeId: "runtime-123", - accountId: "123456789012", - qualifier: "DEFAULT", - payload: new Uint8Array(), - contentType: "application/json", - bearerToken: "secret-token", - }, - { region: "us-east-1" }, - ); - expect(response).toMatchObject({ - statusCode: 502, - error: {}, - }); - expect(cancelled).toBe(1); -}); - -test("CUSTOM_JWT non-2xx preserves cancellation while reading an error response", async () => { - let cancelled = 0; - const body = new ReadableStream({ - cancel() { - cancelled++; - }, - }); - const core = customJwtCore(async () => new Response(body, { status: 401 })); - const controller = new AbortController(); - const reason = new Error("cancelled by user"); - const pending = core.runtime.invokeRuntime( - { - runtimeId: "runtime-123", - accountId: "123456789012", - qualifier: "DEFAULT", - payload: new Uint8Array(), - contentType: "application/json", - bearerToken: "secret-token", - }, - { region: "us-east-1" }, - controller.signal, - ); - - await new Promise((resolve) => setTimeout(resolve, 10)); - controller.abort(reason); - - await expect(pending).rejects.toBe(reason); - expect(cancelled).toBe(1); -}); - test("CUSTOM_JWT transport failures are visible without exposing arbitrary causes", async () => { const { logger, logs } = captureLogs(); const core = customJwtCore( diff --git a/src/core/eval.tsx b/src/core/eval.tsx index f18d224543..60da2a22db 100644 --- a/src/core/eval.tsx +++ b/src/core/eval.tsx @@ -929,23 +929,18 @@ export class EvalClient implements CoreEvalClient { signal, ); if (response.kind === "rejected") { - throw new AgentCoreCLIError( - `Runtime invocation failed (HTTP ${response.statusCode})${ - response.error.detail ? `: ${response.error.detail}` : "" - }`, - { - source: - response.statusCode >= 400 && response.statusCode < 500 - ? ERROR_SOURCE.USER - : ERROR_SOURCE.SERVICE, - meta: { - httpStatusCode: response.statusCode, - ...(response.error.requestId !== undefined && { - requestId: response.error.requestId, - }), - }, + throw new AgentCoreCLIError(`Runtime invocation failed (HTTP ${response.statusCode})`, { + source: + response.statusCode >= 400 && response.statusCode < 500 + ? ERROR_SOURCE.USER + : ERROR_SOURCE.SERVICE, + meta: { + httpStatusCode: response.statusCode, + ...(response.error.requestId !== undefined && { + requestId: response.error.requestId, + }), }, - ); + }); } // Read to completion to free the socket; a scripted example ignores the text. let text = ""; diff --git a/src/core/invokeRuntime.ts b/src/core/invokeRuntime.ts index ee95b4e262..4c299ee339 100644 --- a/src/core/invokeRuntime.ts +++ b/src/core/invokeRuntime.ts @@ -50,7 +50,6 @@ export type RuntimeInvokeStreamResponse = RuntimeInvokeResponseMetadata & { export type RuntimeInvokeRejectedResponse = RuntimeInvokeResponseMetadata & { kind: "rejected"; error: { - detail?: string; requestId?: string; }; }; @@ -75,178 +74,6 @@ export type InvokeRuntimeDeps = { async function* emptyBody(): AsyncGenerator {} -class RuntimeErrorResponseReader { - private static readonly maxBytes = 64 * 1024; - private static readonly maxDetailCharacters = 1024; - - static forRequest( - response: Response, - request: RuntimeInvokeRequest, - bearerToken: string, - signal?: AbortSignal, - ): RuntimeErrorResponseReader { - return new RuntimeErrorResponseReader( - response, - [ - bearerToken, - ...RuntimeErrorResponseReader.requestValues(request), - RuntimeErrorResponseReader.textPayload(request.payload), - ].filter((value): value is string => value !== undefined && value.length > 0), - signal, - ); - } - - constructor( - private readonly response: Response, - private readonly redactions: string[], - private readonly signal?: AbortSignal, - ) {} - - async detail(): Promise { - const body = await this.readBody(); - const candidate = body === undefined ? this.response.statusText : this.parse(body); - return candidate === undefined ? undefined : this.sanitize(candidate); - } - - header(name: string): string | undefined { - const value = this.response.headers.get(name); - return value === null ? undefined : this.sanitize(value, 256); - } - - private async readBody(): Promise { - if (!this.response.body) return undefined; - const reader = this.response.body.getReader(); - const chunks: Uint8Array[] = []; - let size = 0; - const read = () => { - if (!this.signal) return reader.read(); - this.signal.throwIfAborted(); - const signal = this.signal; - return new Promise>>((resolve, reject) => { - const onAbort = () => reject(signal.reason); - signal.addEventListener("abort", onAbort, { once: true }); - reader - .read() - .then(resolve, reject) - .finally(() => { - signal.removeEventListener("abort", onAbort); - }); - }); - }; - try { - while (true) { - const { done, value } = await read(); - if (done) break; - size += value.byteLength; - if (size > RuntimeErrorResponseReader.maxBytes) { - await reader.cancel().catch(() => undefined); - return undefined; - } - chunks.push(value); - } - return new TextDecoder("utf-8", { fatal: true }).decode(Buffer.concat(chunks)); - } catch (error) { - await reader.cancel().catch(() => undefined); - if (this.signal?.aborted) throw this.signal.reason ?? error; - return undefined; - } finally { - reader.releaseLock(); - } - } - - private parse(body: string): string | undefined { - const trimmed = body.trim(); - if (!trimmed) return undefined; - let parsed: unknown; - try { - parsed = JSON.parse(trimmed); - } catch { - return trimmed; - } - if (typeof parsed === "string") return parsed; - const root = RuntimeErrorResponseReader.record(parsed); - if (!root) return undefined; - const nested = RuntimeErrorResponseReader.record(root.error); - const type = - RuntimeErrorResponseReader.string(root, "__type") ?? - RuntimeErrorResponseReader.string(root, "errorType") ?? - RuntimeErrorResponseReader.string(root, "code") ?? - RuntimeErrorResponseReader.string(root, "type") ?? - RuntimeErrorResponseReader.string(nested, "type") ?? - RuntimeErrorResponseReader.string(nested, "code"); - const message = - RuntimeErrorResponseReader.string(root, "message") ?? - (typeof root.error === "string" ? root.error : undefined) ?? - RuntimeErrorResponseReader.string(root, "detail") ?? - RuntimeErrorResponseReader.string(root, "error_description") ?? - RuntimeErrorResponseReader.string(nested, "message"); - const shortType = type?.split("#").at(-1); - if (shortType && message && shortType !== message) return `${shortType}: ${message}`; - return message ?? shortType; - } - - private sanitize( - candidate: string, - maxCharacters = RuntimeErrorResponseReader.maxDetailCharacters, - ): string | undefined { - let value = candidate; - for (const redaction of this.redactions) { - if (redaction) value = value.split(redaction).join("[REDACTED]"); - } - value = Array.from(value, (character) => { - const codePoint = character.codePointAt(0)!; - return codePoint <= 0x1f || (codePoint >= 0x7f && codePoint <= 0x9f) ? " " : character; - }) - .join("") - .replace(/\s+/g, " ") - .trim(); - if (!value) return undefined; - return value.slice(0, maxCharacters); - } - - private static record(value: unknown): Record | undefined { - return value !== null && typeof value === "object" - ? (value as Record) - : undefined; - } - - private static string( - record: Record | undefined, - key: string, - ): string | undefined { - const value = record?.[key]; - return typeof value === "string" && value.trim() ? value : undefined; - } - - private static requestValues(request: RuntimeInvokeRequest): (string | undefined)[] { - return [ - ...(request.applicationHeaders?.map(([, value]) => value) ?? []), - request.runtimeSessionId, - request.runtimeUserId === "default" ? undefined : request.runtimeUserId, - request.mcpSessionId, - request.mcpProtocolVersion, - request.mcpMethod, - request.mcpName, - request.traceId, - request.traceParent, - request.traceState, - request.baggage, - ]; - } - - private static textPayload(payload: Uint8Array): string | undefined { - if (payload.byteLength < 4 || payload.byteLength > RuntimeErrorResponseReader.maxBytes) { - return undefined; - } - try { - const value = new TextDecoder("utf-8", { fatal: true }).decode(payload); - return value.trim() ? value : undefined; - } catch { - return undefined; - } - } -} - // invokeRuntime dispatches by auth mode: a bearer token routes to the CUSTOM_JWT // (raw fetch) path, otherwise the SigV4 SDK path. export async function invokeRuntime( @@ -344,15 +171,11 @@ async function invokeRuntimeWithCustomJwt( logger .child({ httpStatusCode: response.status }) .debug("Runtime invocation returned a non-success response"); - const errorResponse = RuntimeErrorResponseReader.forRequest( - response, - request, - bearerToken, - signal, - ); - const detail = await errorResponse.detail(); const requestId = - errorResponse.header("x-amzn-requestid") ?? errorResponse.header("x-amz-request-id"); + response.headers.get("x-amzn-requestid") ?? + response.headers.get("x-amz-request-id") ?? + undefined; + await response.body?.cancel().catch(() => undefined); return { kind: "rejected", statusCode: response.status, @@ -366,7 +189,6 @@ async function invokeRuntimeWithCustomJwt( traceState: response.headers.get("tracestate") ?? undefined, baggage: response.headers.get("baggage") ?? undefined, error: { - ...(detail !== undefined && { detail }), ...(requestId !== undefined && { requestId }), }, }; diff --git a/src/handlers/runtime/invoke/invoke.screen.test.tsx b/src/handlers/runtime/invoke/invoke.screen.test.tsx index 15d965a998..e4dfc2f47e 100644 --- a/src/handlers/runtime/invoke/invoke.screen.test.tsx +++ b/src/handlers/runtime/invoke/invoke.screen.test.tsx @@ -620,7 +620,7 @@ describe("Runtime invoke JSON console", () => { expect(screen.lastFrame()).toContain("failed · 0 bytes"); }); - test("shows a sanitized rejected response without treating it as streamed output", async () => { + test("shows a rejected response without treating it as streamed output", async () => { const core = new TestCoreClient(); core.runtime .setGetResponse({ agentRuntimeArn: RUNTIME_ARN } as GetAgentRuntimeResponse) @@ -629,7 +629,6 @@ describe("Runtime invoke JSON console", () => { statusCode: 401, contentType: "application/json", error: { - detail: "ExpiredTokenException: The token has expired", requestId: "request-123", }, }); @@ -640,7 +639,6 @@ describe("Runtime invoke JSON console", () => { await screen.press("return"); await waitForText(screen.lastFrame, "Runtime invocation failed · HTTP 401"); - expect(screen.lastFrame()).toContain("ExpiredTokenException: The token has expired"); expect(screen.lastFrame()).toContain("Request ID: request-123"); expect(screen.lastFrame()).toContain("failed · 0 bytes"); }); diff --git a/src/handlers/runtime/invoke/invoke.test.tsx b/src/handlers/runtime/invoke/invoke.test.tsx index 7d21899e08..0b87e19360 100644 --- a/src/handlers/runtime/invoke/invoke.test.tsx +++ b/src/handlers/runtime/invoke/invoke.test.tsx @@ -803,7 +803,6 @@ describe("runtime invoke", () => { statusCode: 401, contentType: "application/json", error: { - detail: "ExpiredTokenException: The token has expired", requestId: "request-123", }, }); @@ -833,8 +832,7 @@ describe("runtime invoke", () => { expect(output.stderrBytes().toString()).toBe( "status=401 content-type=application/json runtime-session-id=- mcp-session-id=- " + "mcp-protocol-version=- trace-id=- trace-parent=- trace-state=- baggage=- " + - "complete=false bytes=0 error=ExpiredTokenException: The token has expired " + - "request-id=request-123\n", + "complete=false bytes=0 error=HTTP 401 request-id=request-123\n", ); expect(output.stderrBytes().toString()).not.toContain("secret-token"); } finally { diff --git a/src/handlers/runtime/invoke/response.test.ts b/src/handlers/runtime/invoke/response.test.ts index 6b03a2242c..ceb9bf71e6 100644 --- a/src/handlers/runtime/invoke/response.test.ts +++ b/src/handlers/runtime/invoke/response.test.ts @@ -263,7 +263,6 @@ describe("Runtime invoke response output", () => { statusCode: 401, contentType: "application/json", error: { - detail: "ExpiredTokenException: The token has expired", requestId: "request-123", }, }); @@ -280,8 +279,7 @@ describe("Runtime invoke response output", () => { expect(stderr.bytes().toString()).toBe( "status=401 content-type=application/json runtime-session-id=- mcp-session-id=- " + "mcp-protocol-version=- trace-id=- trace-parent=- trace-state=- baggage=- " + - "complete=false bytes=0 error=ExpiredTokenException: The token has expired " + - "request-id=request-123\n", + "complete=false bytes=0 error=HTTP 401 request-id=request-123\n", ); }, ); @@ -296,7 +294,7 @@ describe("Runtime invoke response output", () => { writeRuntimeInvokeResponse( rejectedResponse({ statusCode: 403, - error: { detail: "AccessDeniedException: Access denied" }, + error: {}, }), { stdout: stdout.stream, @@ -308,9 +306,7 @@ describe("Runtime invoke response output", () => { expect(await Bun.file(file).exists()).toBe(false); expect(stdout.bytes()).toHaveLength(0); - expect(stderr.bytes().toString()).toContain( - "complete=false bytes=0 error=AccessDeniedException: Access denied", - ); + expect(stderr.bytes().toString()).toContain("complete=false bytes=0 error=HTTP 403"); }); test("preserves partial stdout and reports one static incomplete summary", async () => { diff --git a/src/handlers/runtime/invoke/response.ts b/src/handlers/runtime/invoke/response.ts index 52e3545049..3eacb58b8b 100644 --- a/src/handlers/runtime/invoke/response.ts +++ b/src/handlers/runtime/invoke/response.ts @@ -58,11 +58,10 @@ async function writeRejectedRuntimeResponse( ): Promise { try { await output.beforeOutput?.(); - await pipeline( - [summary(response, 0, false, response.error.detail ?? `HTTP ${response.statusCode}`)], - output.stderr, - { end: false, signal: output.signal }, - ); + await pipeline([summary(response, 0, false, `HTTP ${response.statusCode}`)], output.stderr, { + end: false, + signal: output.signal, + }); } catch (error) { failure(error, output.signal); } diff --git a/src/handlers/runtime/invoke/screen.tsx b/src/handlers/runtime/invoke/screen.tsx index 69bbfcf2bb..a3fe3dc67d 100644 --- a/src/handlers/runtime/invoke/screen.tsx +++ b/src/handlers/runtime/invoke/screen.tsx @@ -246,7 +246,6 @@ export function RuntimeInvokeConsole({ updateExchange({ error: { name: "Runtime invocation failed", - message: response.error.detail, statusCode: response.statusCode, requestId: response.error.requestId, }, diff --git a/src/handlers/runtime/types.tsx b/src/handlers/runtime/types.tsx index 0f083e8fb3..788777acd8 100644 --- a/src/handlers/runtime/types.tsx +++ b/src/handlers/runtime/types.tsx @@ -52,7 +52,6 @@ export type RuntimeInvokeStreamResponse = RuntimeInvokeResponseMetadata & { export type RuntimeInvokeRejectedResponse = RuntimeInvokeResponseMetadata & { kind: "rejected"; error: { - detail?: string; requestId?: string; }; };