From f9207db3ca957d77efe293dddbeaeb8cb122ed40 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 31 Jul 2026 03:10:54 -0700 Subject: [PATCH] fix(bedrock): reject truncated streams and preserve audio results (#116743) Co-authored-by: Peter Steinberger --- .../amazon-bedrock/stream.runtime.test.ts | 123 ++++++++++++++++++ extensions/amazon-bedrock/stream.runtime.ts | 9 +- 2 files changed, 129 insertions(+), 3 deletions(-) diff --git a/extensions/amazon-bedrock/stream.runtime.test.ts b/extensions/amazon-bedrock/stream.runtime.test.ts index 114bbd539c39..cf5499377ac8 100644 --- a/extensions/amazon-bedrock/stream.runtime.test.ts +++ b/extensions/amazon-bedrock/stream.runtime.test.ts @@ -123,6 +123,74 @@ describe("Bedrock inbound image base64", () => { }); describe("Bedrock tool-result replay", () => { + it("replays unsupported audio attachments as their canonical text placeholder", () => { + const messages = testing.convertMessages( + { + messages: [ + { + role: "toolResult", + toolCallId: "call_audio", + toolName: "listen", + content: [{ type: "audio", mimeType: "audio/wav", data: "YXVkaW8=" }], + isError: false, + }, + ], + } as never, + bedrockModel({ input: ["text", "image"] }), + "none", + ); + + expect(messages).toHaveLength(1); + expect(messages[0]).toMatchObject({ + role: ConversationRole.USER, + content: [ + { + toolResult: { + toolUseId: "call_audio", + content: [{ text: "(see attached audio)" }], + }, + }, + ], + }); + }); + + it("preserves valid text and image attachments alongside unsupported audio", () => { + const messages = testing.convertMessages( + { + messages: [ + { + role: "toolResult", + toolCallId: "call_media", + toolName: "inspect", + content: [ + { type: "audio", mimeType: "audio/wav", data: "YXVkaW8=" }, + { type: "text", text: "actual tool output" }, + { type: "image", mimeType: "image/png", data: "aW1hZ2U=" }, + ], + isError: false, + }, + ], + } as never, + bedrockModel({ input: ["text", "image"] }), + "none", + ); + + expect(messages[0]).toMatchObject({ + role: ConversationRole.USER, + content: [ + { + toolResult: { + toolUseId: "call_media", + content: [ + { text: "actual tool output" }, + { image: { format: "png", source: { bytes: expect.any(Uint8Array) } } }, + ], + }, + }, + ], + }); + }); + it("drops payload-less image husks from consecutive tool results", () => { const messages = testing.convertMessages( { @@ -335,6 +403,61 @@ describe("Bedrock profile endpoint resolution", () => { }); describe("Bedrock stop reasons", () => { + it.each([ + { + name: "text", + events: [ + { contentBlockDelta: { contentBlockIndex: 0, delta: { text: "truncated response" } } }, + { contentBlockStop: { contentBlockIndex: 0 } }, + ], + contentType: "text", + }, + { + name: "tool call", + events: [ + { + contentBlockStart: { + contentBlockIndex: 0, + start: { toolUse: { toolUseId: "call_lookup", name: "lookup" } }, + }, + }, + { + contentBlockDelta: { + contentBlockIndex: 0, + delta: { toolUse: { input: '{"query":"partial"}' } }, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + ], + contentType: "toolCall", + }, + ])( + "reports truncated $name streams without a terminal messageStop", + async ({ events, contentType }) => { + vi.spyOn(BedrockRuntimeClient.prototype, "send").mockResolvedValue({ + $metadata: { httpStatusCode: 200 }, + stream: streamEvents([{ messageStart: { role: ConversationRole.ASSISTANT } }, ...events]), + } as never); + + const stream = streamBedrockForTest(bedrockModel({}), { + messages: [{ role: "user", content: "Hello", timestamp: 0 }], + } as never); + const eventTypes: string[] = []; + for await (const event of stream) { + eventTypes.push(event.type); + } + const result = await stream.result(); + + expect(eventTypes.at(-1)).toBe("error"); + expect(eventTypes).not.toContain("done"); + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toBe("Bedrock stream ended before messageStop"); + expect(result.content).toEqual([expect.objectContaining({ type: contentType })]); + expect(result.content[0]).not.toHaveProperty("index"); + expect(result.content[0]).not.toHaveProperty("partialJson"); + }, + ); + it.each([ BedrockStopReason.CONTENT_FILTERED, BedrockStopReason.GUARDRAIL_INTERVENED, diff --git a/extensions/amazon-bedrock/stream.runtime.ts b/extensions/amazon-bedrock/stream.runtime.ts index 2137af2981c9..aa12bbc0ed7f 100644 --- a/extensions/amazon-bedrock/stream.runtime.ts +++ b/extensions/amazon-bedrock/stream.runtime.ts @@ -332,7 +332,7 @@ const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOptions> = } } - if (refusalBuffer && !sawMessageStop) { + if (!sawMessageStop) { throw new Error("Bedrock stream ended before messageStop"); } if (options.signal?.aborted) { @@ -812,7 +812,7 @@ function createBedrockToolResult(message: ToolResultMessage): ContentBlock.ToolR content.push({ text: sanitizeSurrogates(block.text) }); continue; } - if (describeToolResultMediaPlaceholder([block])) { + if (block.type === "image" && describeToolResultMediaPlaceholder([block])) { content.push({ image: createImageBlock(block.mimeType, block.data) }); } } @@ -820,7 +820,10 @@ function createBedrockToolResult(message: ToolResultMessage): ContentBlock.ToolR return { toolResult: { toolUseId: message.toolCallId, - content: content.length > 0 ? content : [{ text: "(no output)" }], + content: + content.length > 0 + ? content + : [{ text: describeToolResultMediaPlaceholder(message.content) ?? "(no output)" }], status: message.isError ? ToolResultStatus.ERROR : ToolResultStatus.SUCCESS, }, };