fix(ai): reject truncated bedrock frames (#46281)

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<object> => ({
   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"

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<readonly [string, object]>)
 const fixedBytes = (bytes: Uint8Array) =>
   fixedResponse(bytes.slice().buffer, { headers: { "content-type": "application/vnd.amazon.eventstream" } })
 
+const fixedByteChunks = (...chunks: ReadonlyArray<Uint8Array>) =>
+  fixedResponse(
+    new ReadableStream<Uint8Array>({
+      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(