Просмотр исходного кода

fix(ai): align Bedrock stream handling (#38712)

Aiden Cline 3 недель назад
Родитель
Сommit
4f201f87a9

+ 27 - 12
packages/ai/src/protocols/bedrock-converse.ts

@@ -155,6 +155,12 @@ const BedrockUsageSchema = Schema.Struct({
 })
 type BedrockUsageSchema = Schema.Schema.Type<typeof BedrockUsageSchema>
 
+const BedrockStreamException = Schema.Struct({
+  message: Schema.optional(Schema.String),
+  originalMessage: Schema.optional(Schema.String),
+  originalStatusCode: Schema.optional(Schema.Number),
+})
+
 // Streaming event shape — the AWS event stream wraps each JSON payload by its
 // `:event-type` header (e.g. `messageStart`, `contentBlockDelta`). We
 // reconstruct that wrapping in `decodeFrames` below so the event schema can
@@ -184,6 +190,9 @@ const BedrockEvent = Schema.Struct({
               signature: Schema.optional(Schema.String),
               // Blob fields in Bedrock's JSON event stream are base64 strings.
               redactedContent: Schema.optional(Schema.String),
+              // Vercel's Bedrock provider exposes the same delta under
+              // Anthropic's shorter `data` spelling.
+              data: Schema.optional(Schema.String),
             }),
           ),
         }),
@@ -203,11 +212,11 @@ const BedrockEvent = Schema.Struct({
       metrics: Schema.optional(Schema.Unknown),
     }),
   ),
-  internalServerException: Schema.optional(Schema.Struct({ message: Schema.String })),
-  modelStreamErrorException: Schema.optional(Schema.Struct({ message: Schema.String })),
-  validationException: Schema.optional(Schema.Struct({ message: Schema.String })),
-  throttlingException: Schema.optional(Schema.Struct({ message: Schema.String })),
-  serviceUnavailableException: Schema.optional(Schema.Struct({ message: Schema.String })),
+  internalServerException: Schema.optional(BedrockStreamException),
+  modelStreamErrorException: Schema.optional(BedrockStreamException),
+  validationException: Schema.optional(BedrockStreamException),
+  throttlingException: Schema.optional(BedrockStreamException),
+  serviceUnavailableException: Schema.optional(BedrockStreamException),
 })
 type BedrockEvent = Schema.Schema.Type<typeof BedrockEvent>
 
@@ -531,14 +540,20 @@ const step = (state: ParserState, event: BedrockEvent) =>
       const index = event.contentBlockDelta.contentBlockIndex
       const reasoning = event.contentBlockDelta.delta.reasoningContent
       const events: LLMEvent[] = []
-      const lifecycle = reasoning.text
-        ? Lifecycle.reasoningDelta(state.lifecycle, events, `reasoning-${index}`, reasoning.text)
-        : reasoning.redactedContent !== undefined
-          ? Lifecycle.reasoningStart(
+      const redactedData = reasoning.redactedContent ?? reasoning.data
+      const providerMetadata = reasoning.signature
+        ? bedrockMetadata({ signature: reasoning.signature })
+        : redactedData !== undefined
+          ? bedrockMetadata({ redactedData })
+          : undefined
+      const lifecycle =
+        reasoning.text !== undefined || providerMetadata !== undefined
+          ? Lifecycle.reasoningDelta(
               state.lifecycle,
               events,
               `reasoning-${index}`,
-              bedrockMetadata({ redactedData: reasoning.redactedContent }),
+              reasoning.text ?? "",
+              providerMetadata,
             )
           : state.lifecycle
       return [
@@ -618,7 +633,7 @@ const step = (state: ParserState, event: BedrockEvent) =>
     }
 
     if (event.metadata) {
-      const usage = mapUsage(event.metadata.usage)
+      const usage = mapUsage(event.metadata.usage) ?? state.pendingFinish?.usage
       return [
         {
           ...state,
@@ -645,7 +660,7 @@ const step = (state: ParserState, event: BedrockEvent) =>
         module: ADAPTER,
         method: "stream",
         reason: classifyProviderFailure({
-          message: exception[1]?.message ?? "Bedrock Converse stream error",
+          message: exception[1]?.message ?? exception[1]?.originalMessage ?? "Bedrock Converse stream error",
           code: exception[0],
         }),
       })

+ 16 - 2
packages/ai/src/protocols/bedrock-event-stream.ts

@@ -53,8 +53,22 @@ const consumeFrames = (route: string) => (state: FrameBufferState, chunk: Uint8A
       })
       cursor = { buffer: cursor.buffer, offset: cursor.offset + totalLength }
 
-      if (decoded.headers[":message-type"]?.value !== "event") continue
-      const eventType = decoded.headers[":event-type"]?.value
+      const messageType = decoded.headers[":message-type"]?.value
+      if (messageType === "error") {
+        const code = decoded.headers[":error-code"]?.value
+        const message = decoded.headers[":error-message"]?.value
+        return yield* ProviderShared.eventError(
+          route,
+          [code, message].filter((value): value is string => typeof value === "string").join(": ") ||
+            "Bedrock Converse event-stream error",
+        )
+      }
+      const eventType =
+        messageType === "event"
+          ? decoded.headers[":event-type"]?.value
+          : messageType === "exception"
+            ? decoded.headers[":exception-type"]?.value
+            : undefined
       if (typeof eventType !== "string") continue
       const payload = utf8.decode(decoded.body)
       if (!payload) continue

+ 165 - 5
packages/ai/test/provider/bedrock-converse.test.ts

@@ -43,6 +43,26 @@ const eventFrame = (type: string, payload: object) =>
     body: utf8Encoder.encode(JSON.stringify(payload)),
   })
 
+const exceptionFrame = (type: string, payload: object) =>
+  codec.encode({
+    headers: {
+      ":message-type": { type: "string", value: "exception" },
+      ":exception-type": { type: "string", value: type },
+      ":content-type": { type: "string", value: "application/json" },
+    },
+    body: utf8Encoder.encode(JSON.stringify(payload)),
+  })
+
+const errorFrame = (code: string, message: string) =>
+  codec.encode({
+    headers: {
+      ":message-type": { type: "string", value: "error" },
+      ":error-code": { type: "string", value: code },
+      ":error-message": { type: "string", value: message },
+    },
+    body: new Uint8Array(),
+  })
+
 const concat = (frames: ReadonlyArray<Uint8Array>) => {
   const total = frames.reduce((sum, frame) => sum + frame.length, 0)
   const out = new Uint8Array(total)
@@ -363,6 +383,19 @@ describe("Bedrock Converse route", () => {
     }),
   )
 
+  it.effect("preserves usage across later metadata events without usage", () =>
+    Effect.gen(function* () {
+      const body = eventStreamBody(
+        ["messageStop", { stopReason: "end_turn" }],
+        ["metadata", { usage: { inputTokens: 5, outputTokens: 2, totalTokens: 7 } }],
+        ["metadata", { metrics: { latencyMs: 100 } }],
+      )
+      const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
+
+      expect(response.usage).toMatchObject({ inputTokens: 5, outputTokens: 2, totalTokens: 7 })
+    }),
+  )
+
   it.effect("assembles streamed tool call input", () =>
     Effect.gen(function* () {
       const body = eventStreamBody(
@@ -478,6 +511,95 @@ describe("Bedrock Converse route", () => {
     }),
   )
 
+  it.effect("preserves reasoning signatures when contentBlockStop is missing", () =>
+    Effect.gen(function* () {
+      const response = yield* LLMClient.generate(baseRequest).pipe(
+        Effect.provide(
+          fixedBytes(
+            eventStreamBody(
+              ["messageStart", { role: "assistant" }],
+              [
+                "contentBlockDelta",
+                { contentBlockIndex: 0, delta: { reasoningContent: { text: "Let me think." } } },
+              ],
+              [
+                "contentBlockDelta",
+                { contentBlockIndex: 0, delta: { reasoningContent: { signature: "sig_1" } } },
+              ],
+              ["messageStop", { stopReason: "end_turn" }],
+            ),
+          ),
+        ),
+      )
+
+      expect(response.events.find((event) => event.type === "reasoning-delta" && event.text === "")).toEqual({
+        type: "reasoning-delta",
+        id: "reasoning-0",
+        text: "",
+        providerMetadata: { bedrock: { signature: "sig_1" } },
+      })
+      expect(response.message.content).toEqual([
+        {
+          type: "reasoning",
+          text: "Let me think.",
+          providerMetadata: { bedrock: { signature: "sig_1" } },
+        },
+      ])
+
+      const prepared = yield* LLMClient.prepare<BedrockConverse.BedrockConverseBody>(
+        LLM.request({ model, messages: [response.message], cache: "none" }),
+      )
+      expect(prepared.body.messages).toEqual([
+        {
+          role: "assistant",
+          content: [{ reasoningContent: { reasoningText: { text: "Let me think.", signature: "sig_1" } } }],
+        },
+      ])
+    }),
+  )
+
+  it.effect("preserves signature-only reasoning blocks", () =>
+    Effect.gen(function* () {
+      const body = eventStreamBody(
+        ["messageStart", { role: "assistant" }],
+        [
+          "contentBlockDelta",
+          { contentBlockIndex: 0, delta: { reasoningContent: { signature: "sig_1" } } },
+        ],
+        ["contentBlockStop", { contentBlockIndex: 0 }],
+        ["messageStop", { stopReason: "end_turn" }],
+      )
+      const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
+
+      expect(response.message.content).toEqual([
+        { type: "reasoning", text: "", providerMetadata: { bedrock: { signature: "sig_1" } } },
+      ])
+    }),
+  )
+
+  it.effect("accepts Vercel-compatible redacted reasoning data deltas", () =>
+    Effect.gen(function* () {
+      const redactedData = "cmVkYWN0ZWQtdGhpbmtpbmc="
+      const body = eventStreamBody(
+        ["messageStart", { role: "assistant" }],
+        ["contentBlockDelta", { contentBlockIndex: 0, delta: { reasoningContent: { data: redactedData } } }],
+        ["contentBlockStop", { contentBlockIndex: 0 }],
+        ["messageStop", { stopReason: "end_turn" }],
+      )
+      const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
+
+      expect(response.events.find((event) => event.type === "reasoning-delta" && event.text === "")).toEqual({
+        type: "reasoning-delta",
+        id: "reasoning-0",
+        text: "",
+        providerMetadata: { bedrock: { redactedData } },
+      })
+      expect(response.message.content).toEqual([
+        { type: "reasoning", text: "", providerMetadata: { bedrock: { redactedData } } },
+      ])
+    }),
+  )
+
   it.effect("round-trips streamed redacted reasoning with tool use into a continuation request", () =>
     Effect.gen(function* () {
       // Bedrock represents redactedContent blobs as base64 strings on its JSON
@@ -511,6 +633,12 @@ describe("Bedrock Converse route", () => {
           ),
         ),
       )
+      expect(response.events.find((event) => event.type === "reasoning-delta" && event.text === "")).toEqual({
+        type: "reasoning-delta",
+        id: "reasoning-0",
+        text: "",
+        providerMetadata: { bedrock: { redactedData } },
+      })
       const prepared = yield* LLMClient.prepare<BedrockConverse.BedrockConverseBody>(
         LLM.request({
           model,
@@ -543,10 +671,10 @@ describe("Bedrock Converse route", () => {
 
   it.effect("classifies throttlingException as a rate limit", () =>
     Effect.gen(function* () {
-      const body = eventStreamBody(
-        ["messageStart", { role: "assistant" }],
-        ["throttlingException", { message: "Slow down" }],
-      )
+      const body = concat([
+        eventFrame("messageStart", { role: "assistant" }),
+        exceptionFrame("throttlingException", { message: "Slow down" }),
+      ])
       const error = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)), Effect.flip)
 
       expect(error.reason).toMatchObject({ _tag: "RateLimit", message: "Slow down" })
@@ -557,7 +685,7 @@ describe("Bedrock Converse route", () => {
     Effect.gen(function* () {
       const error = yield* LLMClient.generate(baseRequest).pipe(
         Effect.provide(
-          fixedBytes(eventStreamBody(["validationException", { message: "Input is too long for requested model" }])),
+          fixedBytes(exceptionFrame("validationException", { message: "Input is too long for requested model" })),
         ),
         Effect.flip,
       )
@@ -570,6 +698,38 @@ describe("Bedrock Converse route", () => {
     }),
   )
 
+  it.effect("uses originalMessage from model stream exception frames", () =>
+    Effect.gen(function* () {
+      const error = yield* LLMClient.generate(baseRequest).pipe(
+        Effect.provide(
+          fixedBytes(
+            exceptionFrame("modelStreamErrorException", {
+              originalMessage: "Upstream model failed",
+              originalStatusCode: 500,
+            }),
+          ),
+        ),
+        Effect.flip,
+      )
+
+      expect(error.reason).toMatchObject({ _tag: "ProviderInternal", message: "Upstream model failed" })
+    }),
+  )
+
+  it.effect("fails unmodeled AWS event-stream errors", () =>
+    Effect.gen(function* () {
+      const error = yield* LLMClient.generate(baseRequest).pipe(
+        Effect.provide(fixedBytes(errorFrame("BadStream", "Stream failed"))),
+        Effect.flip,
+      )
+
+      expect(error.reason).toMatchObject({
+        _tag: "InvalidProviderOutput",
+        message: "BadStream: Stream failed",
+      })
+    }),
+  )
+
   it.effect("rejects requests with no auth path", () =>
     Effect.gen(function* () {
       const unsignedModel = AmazonBedrock.configure({