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

This commit is contained in:
Aiden Cline 2026-07-24 11:15:23 -05:00 committed by GitHub
parent 5aa276c117
commit 4f201f87a9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 208 additions and 19 deletions

View file

@ -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],
}),
})

View file

@ -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

View file

@ -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({