From bf50194f99a64875938d2780986c1fba7aa10aad Mon Sep 17 00:00:00 2001 From: Aiden Cline <63023139+rekram1-node@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:47:33 -0500 Subject: [PATCH] fix(ai): reject truncated bedrock frames (#46281) --- .../ai/src/protocols/bedrock-event-stream.ts | 31 +++++++++-- .../ai/test/provider/bedrock-converse.test.ts | 51 ++++++++++++++++++- 2 files changed, 77 insertions(+), 5 deletions(-) diff --git a/packages/ai/src/protocols/bedrock-event-stream.ts b/packages/ai/src/protocols/bedrock-event-stream.ts index a1b133d3fa8..a29cc1c529d 100644 --- a/packages/ai/src/protocols/bedrock-event-stream.ts +++ b/packages/ai/src/protocols/bedrock-event-stream.ts @@ -1,7 +1,7 @@ import { EventStreamCodec } from "@smithy/eventstream-codec" import { fromUtf8, toUtf8 } from "@smithy/util-utf8" import { Effect, Encoding, Stream } from "effect" -import { AIError, AIErrorReason } from "../schema/index.js" +import { AIError, AIErrorReason, InvalidProviderOutputError } from "../schema/index.js" import { Framing } from "../route/framing.js" import { ProviderShared } from "./shared.js" @@ -22,6 +22,10 @@ interface FrameBufferState { const initialFrameBuffer: FrameBufferState = { buffer: new Uint8Array(0), offset: 0 } +type FrameInput = { readonly _tag: "Chunk"; readonly bytes: Uint8Array } | { readonly _tag: "End" } + +const endOfStream: FrameInput = { _tag: "End" } + const appendChunk = (state: FrameBufferState, chunk: Uint8Array): FrameBufferState => { const remaining = state.buffer.length - state.offset // Compact: drop the consumed prefix and append the new chunk in one alloc. @@ -33,9 +37,23 @@ const appendChunk = (state: FrameBufferState, chunk: Uint8Array): FrameBufferSta return { buffer: next, offset: 0 } } -const consumeFrames = (route: string) => (state: FrameBufferState, chunk: Uint8Array) => +const consumeFrames = (route: string) => (state: FrameBufferState, input: FrameInput) => Effect.gen(function* () { - let cursor = appendChunk(state, chunk) + if (input._tag === "End") { + const remaining = state.buffer.subarray(state.offset) + if (remaining.length > 0) + return yield* new AIError({ + reason: new InvalidProviderOutputError({ + route, + classification: "incomplete-stream", + message: `Incomplete Bedrock Converse event-stream frame: ${remaining.length} buffered bytes remain at end of stream`, + body: Encoding.encodeBase64(remaining), + }), + }) + return [state, []] as const + } + + let cursor = appendChunk(state, input.bytes) const out: object[] = [] while (cursor.buffer.length - cursor.offset >= 4) { const view = cursor.buffer.subarray(cursor.offset) @@ -113,7 +131,12 @@ const consumeFrames = (route: string) => (state: FrameBufferState, chunk: Uint8A export const framing = (route: string): Framing.Definition => ({ id: "aws-event-stream", body: (frame) => ("rawBody" in frame && typeof frame.rawBody === "string" ? frame.rawBody : undefined), - frame: (bytes) => bytes.pipe(Stream.mapAccumEffect(() => initialFrameBuffer, consumeFrames(route))), + frame: (bytes) => + bytes.pipe( + Stream.map((bytes): FrameInput => ({ _tag: "Chunk", bytes })), + Stream.concat(Stream.succeed(endOfStream)), + Stream.mapAccumEffect(() => initialFrameBuffer, consumeFrames(route)), + ), }) export * as BedrockEventStream from "./bedrock-event-stream.js" diff --git a/packages/ai/test/provider/bedrock-converse.test.ts b/packages/ai/test/provider/bedrock-converse.test.ts index 331390cc5e1..828a650a986 100644 --- a/packages/ai/test/provider/bedrock-converse.test.ts +++ b/packages/ai/test/provider/bedrock-converse.test.ts @@ -1,7 +1,7 @@ import { EventStreamCodec } from "@smithy/eventstream-codec" import { fromUtf8, toUtf8 } from "@smithy/util-utf8" import { describe, expect } from "bun:test" -import { Effect, Ref, Stream } from "effect" +import { Effect, Encoding, Ref, Stream } from "effect" import { CacheHint, GenerationOptions, @@ -84,6 +84,17 @@ const eventStreamBody = (...payloads: ReadonlyArray) const fixedBytes = (bytes: Uint8Array) => fixedResponse(bytes.slice().buffer, { headers: { "content-type": "application/vnd.amazon.eventstream" } }) +const fixedByteChunks = (...chunks: ReadonlyArray) => + fixedResponse( + new ReadableStream({ + start(controller) { + chunks.forEach((chunk) => controller.enqueue(chunk)) + controller.close() + }, + }), + { headers: { "content-type": "application/vnd.amazon.eventstream" } }, + ) + const model = AmazonBedrock.configure({ baseURL: "https://bedrock-runtime.test", apiKey: "test-bearer", @@ -386,6 +397,44 @@ describe("Bedrock Converse route", () => { }), ) + it.effect("rejects truncated event-stream frames after message stop", () => + Effect.gen(function* () { + const partialFrames = [ + eventFrame("metadata", { usage: { inputTokens: 5, outputTokens: 2, totalTokens: 7 } }).subarray(0, 3), + exceptionFrame("modelStreamErrorException", { originalMessage: "Upstream model failed" }).subarray(0, -1), + ] + + for (const partial of partialFrames) { + const error = yield* LLMClient.generate(baseRequest).pipe( + Effect.provide(fixedBytes(concat([eventFrame("messageStop", { stopReason: "end_turn" }), partial]))), + Effect.flip, + ) + + expect(error).toMatchObject({ + reason: { _tag: "InvalidProviderOutput", classification: "incomplete-stream" }, + message: `Incomplete Bedrock Converse event-stream frame: ${partial.length} buffered bytes remain at end of stream`, + }) + expect(error.reason.body).toBe(Encoding.encodeBase64(partial)) + } + }), + ) + + it.effect("decodes frames split across transport chunks through exact-boundary EOF", () => + Effect.gen(function* () { + const body = eventStreamBody( + ["messageStart", { role: "assistant" }], + ["contentBlockDelta", { contentBlockIndex: 0, delta: { text: "Hello" } }], + ["messageStop", { stopReason: "end_turn" }], + ) + const response = yield* LLMClient.generate(baseRequest).pipe( + Effect.provide(fixedByteChunks(body.subarray(0, 2), body.subarray(2, 17), body.subarray(17))), + ) + + expect(response.text).toBe("Hello") + expect(response.finishReason).toEqual({ normalized: "stop", raw: "end_turn" }) + }), + ) + it.effect("maps model context window exhaustion to length", () => Effect.gen(function* () { const response = yield* LLMClient.generate(baseRequest).pipe(