refactor: extract StreamAnthropicChatEvents/StreamAnthropicResponsesEvents and route Bedrock Claude streams through shared Anthropic loop via bedrockAnthropicEventReader - #4360
Conversation
|
|
📝 WalkthroughWalkthroughThis PR extracts shared Anthropic streaming event-loop logic into reusable handlers and bridges them to Bedrock via an eventstream adapter, eliminating duplicate code and enabling proper error propagation through a new ChangesStreaming Refactoring and Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
Confidence Score: 3/5The Anthropic-native and Bedrock-Claude streaming paths look correct; there is a token-accounting inconsistency in the Bedrock Converse path that affects non-Anthropic models with cached inputs. The Converse-API ChatCompletionStream loop lost the usage.TotalTokens cached-token adjustment while keeping the identical adjustment on usage.PromptTokens. Requests that hit the cache through non-Anthropic Bedrock models will emit a final usage chunk where total_tokens is lower than prompt_tokens + completion_tokens, a data inconsistency shipped to callers on every cached request through that path. core/providers/bedrock/bedrock.go — Converse-API ChatCompletionStream cached-token accounting (lines 1528–1530); core/providers/bedrock/anthropicstream.go — new adapter warrants integration tests for the Anthropic-via-Bedrock streaming path. Important Files Changed
Reviews (1): Last reviewed commit: "feat: adds anthropic stream delegation i..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if usage.PromptTokensDetails != nil { | ||
| usage.PromptTokens = usage.PromptTokens + usage.PromptTokensDetails.CachedReadTokens + usage.PromptTokensDetails.CachedWriteTokens | ||
| } |
There was a problem hiding this comment.
TotalTokens adjustment for cached tokens in Converse path
The Converse-path ChatCompletionStream loop previously adjusted both usage.PromptTokens and usage.TotalTokens by CachedReadTokens + CachedWriteTokens. After the refactor the TotalTokens line was dropped, so for any Bedrock Converse-API request that carries cached tokens (non-Anthropic models) the emitted final chunk will report a total_tokens that is lower than prompt_tokens + completion_tokens, breaking any downstream consumer that relies on the relationship between those two fields.
| if usage.PromptTokensDetails != nil { | |
| usage.PromptTokens = usage.PromptTokens + usage.PromptTokensDetails.CachedReadTokens + usage.PromptTokensDetails.CachedWriteTokens | |
| } | |
| if usage.PromptTokensDetails != nil { | |
| usage.PromptTokens = usage.PromptTokens + usage.PromptTokensDetails.CachedReadTokens + usage.PromptTokensDetails.CachedWriteTokens | |
| usage.TotalTokens = usage.TotalTokens + usage.PromptTokensDetails.CachedReadTokens + usage.PromptTokensDetails.CachedWriteTokens | |
| } |
| if len(message.Payload) == 0 { | ||
| continue | ||
| } | ||
|
|
||
| // Non-"event" message types carry AWS exception details in their headers. | ||
| if msgTypeHeader := message.Headers.Get(":message-type"); msgTypeHeader != nil { |
There was a problem hiding this comment.
:message-type header check
An AWS exception frame with an empty body is unusual but not impossible. The current ordering silently continues on it rather than inspecting the :message-type / :exception-type headers, so a terminal accessDeniedException or validationException with an empty payload would be swallowed and the loop would stall waiting for more frames until the stream eventually closes.
| func newBedrockAnthropicEventReader(src io.Reader, providerName schemas.ModelProvider) *bedrockAnthropicEventReader { | ||
| return &bedrockAnthropicEventReader{ | ||
| decoder: eventstream.NewDecoder(), | ||
| src: src, | ||
| payloadBuf: make([]byte, 0, 1024*1024), // 1MB payload buffer | ||
| providerName: providerName, | ||
| } |
There was a problem hiding this comment.
payloadBuf expansion is never captured between calls
eventstream.Decoder.Decode may return a larger backing slice when the payload exceeds the initial capacity, but the result is not stored back into r.payloadBuf. Subsequent calls always pass the original 1 MB slice, so for streams with payloads larger than 1 MB the pre-allocation hint provides no benefit. The existing Converse path has the same pattern, so this is not a regression, but capturing the returned slice would make the comment accurate.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@core/providers/anthropic/anthropic.go`:
- Around line 684-685: The stream startTime is currently captured before the
HTTP handshake, inflating streaming latency; move the startTime initialization
to immediately after the activeClient.Do(req, resp) returns (i.e., after the
handshake) and pass that fresh startTime into StreamAnthropicChatEvents (and the
other new helper call referenced in the comment) instead of the earlier
startTime from HandleAnthropic*Stream so first/final-chunk latency measures only
post-handshake streaming time.
In `@core/providers/bedrock/anthropicstream.go`:
- Around line 69-75: Reorder the check so we inspect message.Headers
(specifically message.Headers.Get(":message-type") and the ":exception-type"
header) before skipping frames with an empty payload: if a non-"event"
:message-type is present (or an :exception-type header), treat it as a Bedrock
exception frame and handle/return it instead of continuing; only skip the frame
when there is no payload AND no exception/ non-"event" header. Move or adjust
the existing len(message.Payload) == 0 guard accordingly so header-only
exception frames are not dropped.
In `@core/providers/bedrock/bedrock.go`:
- Around line 1953-1973: The synthetic text-delta branch in the
streamEvent.Delta.ToolUse handling creates a BifrostResponsesStreamResponse but
omits the per-item routing fields (OutputIndex, ContentIndex, and current item
identity) used downstream to attach deltas to the active output item; fix it by
populating the same routing fields that
streamEvent.ToBifrostResponsesStream(...) would include: copy streamEvent's
OutputIndex and ContentIndex and the current item identifier (e.g.,
streamEvent.Item.ID or equivalent field present on streamEvent) into
response.ExtraFields (and any top-level fields if the normal converter sets
them), or obtain a base response from streamEvent.ToBifrostResponsesStream and
merge its routing fields into the synthesized response before calling
providerUtils.ProcessAndSendResponse.
- Around line 1528-1530: The code is repeatedly adding cached token counts into
usage.PromptTokens on every EventStream message, inflating reported prompt
tokens; to fix, stop folding CachedReadTokens/CachedWriteTokens into
usage.PromptTokens inside the per-event loop and instead perform a one-time
aggregation immediately before calling CreateBifrostChatCompletionChunkResponse
by checking usage.PromptTokensDetails and adding usage.PromptTokens +=
usage.PromptTokensDetails.CachedReadTokens +
usage.PromptTokensDetails.CachedWriteTokens once, then proceed to
CreateBifrostChatCompletionChunkResponse; remove or guard the in-loop adjustment
around usage.PromptTokensDetails so cached tokens are not added repeatedly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 7bff2855-1dd1-4781-a385-258b6eb4ce02
📒 Files selected for processing (3)
core/providers/anthropic/anthropic.gocore/providers/bedrock/anthropicstream.gocore/providers/bedrock/bedrock.go
| StreamAnthropicChatEvents(ctx, sseReader, responseChan, jsonBody, startTime, sendBackRawRequest, sendBackRawResponse, providerName, postHookRunner, postResponseConverter, logger, postHookSpanFinalizer) | ||
| }() |
There was a problem hiding this comment.
Capture the shared stream timer after the handshake.
Both new helper calls inherit startTime from HandleAnthropic*Stream, but that timestamp is still taken before activeClient.Do(req, resp). This makes first/final chunk latency include connection/header wait for both native Anthropic streams and the new Bedrock paths that reuse these helpers, so the refactor now spreads inflated latency metrics across both implementations.
Based on learnings, initialize streaming startTime after client.Do so final-chunk latency reflects post-handshake streaming time.
Also applies to: 1189-1190
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@core/providers/anthropic/anthropic.go` around lines 684 - 685, The stream
startTime is currently captured before the HTTP handshake, inflating streaming
latency; move the startTime initialization to immediately after the
activeClient.Do(req, resp) returns (i.e., after the handshake) and pass that
fresh startTime into StreamAnthropicChatEvents (and the other new helper call
referenced in the comment) instead of the earlier startTime from
HandleAnthropic*Stream so first/final-chunk latency measures only post-handshake
streaming time.
Source: Learnings
| if len(message.Payload) == 0 { | ||
| continue | ||
| } | ||
|
|
||
| // Non-"event" message types carry AWS exception details in their headers. | ||
| if msgTypeHeader := message.Headers.Get(":message-type"); msgTypeHeader != nil { | ||
| if msgType := msgTypeHeader.String(); msgType != "event" { |
There was a problem hiding this comment.
Inspect EventStream headers before skipping empty payloads.
This drops any frame whose payload is empty before checking :message-type / :exception-type. Bedrock exception frames are classified from those headers, so a header-only exception is silently ignored and the shared Anthropic loop can fall through to EOF as if the stream ended normally instead of surfacing or retrying the failure.
💡 Minimal fix
- if len(message.Payload) == 0 {
- continue
- }
-
- // Non-"event" message types carry AWS exception details in their headers.
if msgTypeHeader := message.Headers.Get(":message-type"); msgTypeHeader != nil {
if msgType := msgTypeHeader.String(); msgType != "event" {
excType := msgType
if excHeader := message.Headers.Get(":exception-type"); excHeader != nil {
if v := excHeader.String(); v != "" {
@@
return "", nil, fmt.Errorf("%s stream %s: %s", r.providerName, excType, errMsg)
}
}
+
+ if len(message.Payload) == 0 {
+ continue
+ }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@core/providers/bedrock/anthropicstream.go` around lines 69 - 75, Reorder the
check so we inspect message.Headers (specifically
message.Headers.Get(":message-type") and the ":exception-type" header) before
skipping frames with an empty payload: if a non-"event" :message-type is present
(or an :exception-type header), treat it as a Bedrock exception frame and
handle/return it instead of continuing; only skip the frame when there is no
payload AND no exception/ non-"event" header. Move or adjust the existing
len(message.Payload) == 0 guard accordingly so header-only exception frames are
not dropped.
| if usage.PromptTokensDetails != nil { | ||
| usage.PromptTokens = usage.PromptTokens + usage.PromptTokensDetails.CachedReadTokens + usage.PromptTokensDetails.CachedWriteTokens | ||
| } |
There was a problem hiding this comment.
Only fold cached tokens into PromptTokens once.
This runs inside the per-event loop, so every later EventStream message re-adds the same cached counts and inflates the final usage block. That will overstate streamed token usage for any consumer using the terminal chunk for accounting.
🩹 Suggested fix
- if usage.PromptTokensDetails != nil {
- usage.PromptTokens = usage.PromptTokens + usage.PromptTokensDetails.CachedReadTokens + usage.PromptTokensDetails.CachedWriteTokens
- }Apply that adjustment once, immediately before CreateBifrostChatCompletionChunkResponse(...).
Based on learnings: in the Bedrock provider, cached tokens should be aggregated into input tokens as a single step, and TotalTokens already includes cached token counts.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@core/providers/bedrock/bedrock.go` around lines 1528 - 1530, The code is
repeatedly adding cached token counts into usage.PromptTokens on every
EventStream message, inflating reported prompt tokens; to fix, stop folding
CachedReadTokens/CachedWriteTokens into usage.PromptTokens inside the per-event
loop and instead perform a one-time aggregation immediately before calling
CreateBifrostChatCompletionChunkResponse by checking usage.PromptTokensDetails
and adding usage.PromptTokens += usage.PromptTokensDetails.CachedReadTokens +
usage.PromptTokensDetails.CachedWriteTokens once, then proceed to
CreateBifrostChatCompletionChunkResponse; remove or guard the in-loop adjustment
around usage.PromptTokensDetails so cached tokens are not added repeatedly.
Source: Learnings
| if streamEvent.Delta != nil && streamEvent.Delta.ToolUse != nil && isAccumulatingStructuredOutput { | ||
| // Convert tool use delta to text delta | ||
| content := streamEvent.Delta.ToolUse.Input | ||
| response := &schemas.BifrostResponsesStreamResponse{ | ||
| Type: schemas.ResponsesStreamResponseTypeOutputTextDelta, | ||
| SequenceNumber: chunkIndex, | ||
| Delta: &content, | ||
| ExtraFields: schemas.BifrostResponseExtraFields{ | ||
| ChunkIndex: chunkIndex, | ||
| Latency: time.Since(lastChunkTime).Milliseconds(), | ||
| RawResponse: string(chunkPayload.Bytes), | ||
| ChunkIndex: chunkIndex, | ||
| Latency: time.Since(lastChunkTime).Milliseconds(), | ||
| }, | ||
| } | ||
| lastChunkTime = time.Now() | ||
| chunkIndex++ | ||
| providerUtils.ProcessAndSendResponse(ctx, postHookRunner, | ||
| providerUtils.GetBifrostResponseForStreamResponse(nil, nil, passthroughResp, nil, nil, nil), | ||
| responseChan, postHookSpanFinalizer) | ||
| } | ||
| for i, response := range responses { | ||
| if response != nil { | ||
| response.ExtraFields = schemas.BifrostResponseExtraFields{ | ||
| ChunkIndex: chunkIndex, | ||
| Latency: time.Since(lastChunkTime).Milliseconds(), | ||
| } | ||
| chunkIndex++ | ||
| lastChunkTime = time.Now() | ||
| // Only attach raw response to the last response of the incoming event | ||
| if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) && i == len(responses)-1 { | ||
| response.ExtraFields.RawResponse = string(chunkPayload.Bytes) | ||
| } | ||
| // Finalize the very last chunk of the stream with accumulated usage, | ||
| // raw request, total latency, and the stream-end indicator. | ||
| if isLastChunk && i == len(responses)-1 { | ||
| if response.Response == nil { | ||
| response.Response = &schemas.BifrostResponsesResponse{} | ||
| } | ||
| if usage.InputTokensDetails != nil { | ||
| usage.InputTokens = usage.InputTokens + usage.InputTokensDetails.CachedReadTokens + usage.InputTokensDetails.CachedWriteTokens | ||
| usage.TotalTokens = usage.TotalTokens + usage.InputTokensDetails.CachedReadTokens + usage.InputTokensDetails.CachedWriteTokens | ||
| } | ||
| response.Response.Usage = usage | ||
| if providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) { | ||
| providerUtils.ParseAndSetRawRequest(&response.ExtraFields, jsonData) | ||
| } | ||
| response.ExtraFields.Latency = time.Since(startTime).Milliseconds() | ||
| ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) | ||
| } | ||
| providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, nil, response, nil, nil, nil), responseChan, postHookSpanFinalizer) | ||
| lastChunkTime = time.Now() | ||
|
|
||
| if providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) { | ||
| response.ExtraFields.RawResponse = string(message.Payload) | ||
| } | ||
| } | ||
| if isLastChunk { | ||
| return | ||
| } | ||
| } else { | ||
| // Converse API path: parse Bedrock Converse-specific stream events | ||
| var streamEvent BedrockStreamEvent | ||
| if err := sonic.Unmarshal(message.Payload, &streamEvent); err != nil { | ||
| provider.logger.Debug("Failed to parse JSON from event buffer: %v, data: %s", err, string(message.Payload)) | ||
| providerUtils.ProcessAndSendError(ctx, postHookRunner, err, responseChan, provider.logger, postHookSpanFinalizer) | ||
| return | ||
| } | ||
|
|
||
| if streamEvent.Trace != nil { | ||
| streamTrace = streamEvent.Trace | ||
| providerUtils.ProcessAndSendResponse(ctx, postHookRunner, providerUtils.GetBifrostResponseForStreamResponse(nil, nil, response, nil, nil, nil), responseChan, postHookSpanFinalizer) | ||
| continue |
There was a problem hiding this comment.
Carry the normal item/content routing fields on synthetic text deltas.
These structured-output chunks bypass streamEvent.ToBifrostResponsesStream(...) and emit only Type, SequenceNumber, and Delta. That drops the per-item/per-content fields (OutputIndex, ContentIndex, and current item identity when available) that downstream Responses clients use to attach deltas to the active output item.
Mirror the routing fields the normal converter emits when synthesizing these response.output_text.delta events.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@core/providers/bedrock/bedrock.go` around lines 1953 - 1973, The synthetic
text-delta branch in the streamEvent.Delta.ToolUse handling creates a
BifrostResponsesStreamResponse but omits the per-item routing fields
(OutputIndex, ContentIndex, and current item identity) used downstream to attach
deltas to the active output item; fix it by populating the same routing fields
that streamEvent.ToBifrostResponsesStream(...) would include: copy streamEvent's
OutputIndex and ContentIndex and the current item identifier (e.g.,
streamEvent.Item.ID or equivalent field present on streamEvent) into
response.ExtraFields (and any top-level fields if the normal converter sets
them), or obtain a base response from streamEvent.ToBifrostResponsesStream and
merge its routing fields into the synthesized response before calling
providerUtils.ProcessAndSendResponse.
2 participants
Summary
Claude-on-Bedrock streaming previously duplicated the entire Anthropic SSE event loop inline, maintaining its own per-chunk state and accumulating usage independently. This PR replaces that duplicated path with a shared
bedrockAnthropicEventReaderadapter that translates Bedrock's binary AWS EventStream framing into theSSEEventReadercontract, allowing bothChatCompletionStreamandResponsesStreamfor Anthropic models on Bedrock to be driven by the canonicalStreamAnthropicChatEvents/StreamAnthropicResponsesEventsloops.Changes
HandleAnthropicChatCompletionStreamingandHandleAnthropicResponsesStreaminto standalone exported functionsStreamAnthropicChatEventsandStreamAnthropicResponsesEvents. These accept anySSEEventReader, decoupling wire framing from event processing.StreamReaderError, a sentinel error type that lets a customSSEEventReadersurface a fully-formed*schemas.BifrostErrorthrough theReadEventcontract. The shared loops detect it viaerrors.Asand forward the embedded error verbatim, preserving status codes so the retry gate behaves correctly. The native SSE reader never returns this type, so existing SSE callers are unaffected.bedrockAnthropicEventReader(incore/providers/bedrock/anthropicstream.go), which decodes AWS EventStream frames, unwraps the inner Anthropic SSE chunk from the{"bytes": ...}payload, and derives the SSE-style event type from the chunk's owntypefield. AWS exception frames are classified as retryable or terminal and surfaced as*StreamReaderErroror plain errors respectively, preserving Bedrock's retry semantics.BedrockProvider.ChatCompletionStreamandBedrockProvider.ResponsesStream, Anthropic model families now early-exit intoStreamAnthropicChatEvents/StreamAnthropicResponsesEventsvia the adapter reader. A post-response converter on the chat path stamps each chunk with a locally generated ID and the requested model ID. The Converse API path for non-Anthropic models is unchanged.anthropicChatStreamState,anthropicStreamState), duplicated usage accumulation, and the per-chunkToBifrostChatCompletionStream/ToBifrostResponsesStreamdispatch that previously lived inside the Bedrock event loop for Anthropic models.response.ID = messageIDis now assigned beforepostResponseConverterruns, giving converters (e.g. Bedrock's ID/model override) the final say on those fields.Type of change
Affected areas
How to test
go test ./...ChatCompletionStreamand verify streamed chunks arrive with correct IDs, model names, usage, and finish reasons.ResponsesStreamand verify the same.serviceUnavailableExceptionfrom Bedrock and confirm the retry gate fires as expected (retryable status code preserved).stopand content is correctly assembled.Breaking changes
Related issues
Security considerations
None. No new auth surfaces, secrets handling, or PII exposure introduced.
Checklist
docs/contributing/README.mdand followed the guidelinesSummary by CodeRabbit
Bug Fixes
Refactor