moves sse hearbeats to a common structure to reuse - #5850
Merged
akshaydeo merged 1 commit intoAug 5, 2026
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
🚧 Files skipped from review as they are similar to previous changes (9)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change centralizes SSE heartbeat management, adds heartbeats to eligible streaming routes, bounds DNS and skill lookups with timeouts, and adds a read-only virtual key count schema field. ChangesSSE heartbeat and disconnect handling
Bounded external operations
Customer configuration schema
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant StreamingRoute
participant SSEHeartbeat
participant StreamReader
participant Client
StreamingRoute->>SSEHeartbeat: StartSSEHeartbeat
SSEHeartbeat->>StreamReader: Send heartbeat
StreamReader->>Client: Write SSE heartbeat
Client-->>StreamReader: Accept or disconnect
StreamingRoute->>SSEHeartbeat: StopSSEHeartbeat
SSEHeartbeat->>StreamReader: Close and await exit
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
akshaydeo
marked this pull request as ready for review
August 4, 2026 22:39
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 `@transports/bifrost-http/integrations/router.go`:
- Around line 3365-3373: Update passthroughHeartbeatEligible to parse or split
the Content-Type at parameters and compare the complete media type
case-insensitively against text/event-stream, rejecting prefix-collision values
such as text/event-stream+json and text/event-streaming. Extend
Test_passthroughHeartbeatEligible with uppercase, parameterized, and
prefix-collision cases.
- Around line 3106-3150: Remove the synthetic Bedrock heartbeat implementation,
including bedrockHeartbeatEventType and sendBedrockEventStreamHeartbeat, and
remove all call sites that emit this event. Preserve the normal Bedrock
EventStream frames and avoid sending any unknown event type that AWS SDK for Go
v2 exposes as types.UnknownUnionMember.
In `@transports/bifrost-http/lib/streamreader_test.go`:
- Around line 843-852: Replace fixed sleeps with channel-based delivery
synchronization: in transports/bifrost-http/lib/streamreader_test.go lines
843-852, wait for two observed heartbeat frames before calling StopSSEHeartbeat,
using a timeout only as the failure bound; in
transports/bifrost-http/integrations/router_heartbeat_test.go lines 41-46, wait
for the reader to receive an SSE comment before closing stream; and at lines
65-74, wait for a complete Bedrock EventStream frame before closing stream.
In `@transports/bifrost-http/lib/streamreader.go`:
- Around line 135-140: Increase DefaultSSEHeartbeatInterval from 100
milliseconds to a less frequent default such as one second to reduce unnecessary
writes across idle inference, routed streaming, and SSE passthrough streams.
Keep any shorter heartbeat intervals limited to focused lifecycle tests rather
than changing transport configuration.
🪄 Autofix
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: CHILL
Plan: Pro Plus
Run ID: 00cd5e83-a131-4559-ae59-e7d71ca48106
📒 Files selected for processing (5)
transports/bifrost-http/handlers/inference.gotransports/bifrost-http/integrations/router.gotransports/bifrost-http/integrations/router_heartbeat_test.gotransports/bifrost-http/lib/streamreader.gotransports/bifrost-http/lib/streamreader_test.go
akshaydeo
force-pushed
the
08-04-moves_sse_hearbeats_to_a_common_structure_to_reuse
branch
from
August 5, 2026 00:46
24699bf to
f52cd43
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
transports/bifrost-http/handlers/skills_serving.go (1)
1324-1348: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a regression test for the detached lookup context.
In
transports/bifrost-http/handlers/skills_serving_test.go, add a table-driven test that records the context passed toGetSkillByNameand asserts a deadline nearlookupSkillByPathParamTimeout. This prevents reintroducing*fasthttp.RequestCtxinto the database lookup.🤖 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 `@transports/bifrost-http/handlers/skills_serving.go` around lines 1324 - 1348, Add a table-driven regression test in the skills serving handler tests that invokes lookupSkillByPathParam with a mock store recording the context passed to GetSkillByName, then assert the context has a deadline approximately lookupSkillByPathParamTimeout from invocation and is not the request context. Cover the successful lookup path while preserving existing handler behavior.Source: Coding guidelines
transports/bifrost-http/integrations/router_heartbeat_test.go (1)
61-80: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd route-level passthrough stream coverage.
Test_passthroughHeartbeatEligibletests only the predicate. No test invokeshandlePassthroughStream. Add coverage for SSE and non-SSE responses and assert the proxied response body. This catches call-site regressions that inject heartbeats into non-SSE bodies or omit them from SSE.🤖 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 `@transports/bifrost-http/integrations/router_heartbeat_test.go` around lines 61 - 80, Add route-level tests that invoke handlePassthroughStream for both SSE and non-SSE responses, asserting the proxied response body remains unchanged for non-SSE content and receives the expected heartbeat behavior for SSE content. Keep Test_passthroughHeartbeatEligible focused on predicate cases and use the existing router/test helpers and response setup.core/utils.go (2)
657-664: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winUse
net.DefaultResolverfor validation.Use
net.DefaultResolver.LookupIPAddrso validation andConfigureDialerhonor the same resolver configuration.🤖 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/utils.go` around lines 657 - 664, Update the hostname validation lookup to call net.DefaultResolver.LookupIPAddr instead of constructing a new net.Resolver, while preserving the existing lookup context, error wrapping, and IP extraction behavior.
655-656: 🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy liftPropagate caller cancellation into DNS validation.
ValidateExternalURLignores existing request contexts and can delay cancellation for up to five seconds. Add a context-aware helper that derivescontext.WithTimeoutfrom the caller context, and pass available contexts from handlers and framework loaders. Keep a context-free wrapper for background callers.🤖 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/utils.go` around lines 655 - 656, Update ValidateExternalURL and its DNS lookup path to accept and propagate a caller context when creating the timeout, adding a context-aware helper while retaining a context-free wrapper for background callers. Update available call sites in handlers and framework loaders to pass their request or loader contexts, using the existing context-aware symbol throughout.Source: Coding guidelines
🤖 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/utils_test.go`:
- Around line 206-225: Update TestValidateExternalURLBoundsDNSLookup and the
ValidateExternalURL DNS-resolution path to use an injectable resolver seam. Add
a fake resolver that blocks until its context is canceled, then assert the
active lookup receives cancellation before the configured timeout; retain the
.invalid hostname test for immediate NXDOMAIN behavior and keep expired-context
coverage separate.
In `@transports/bifrost-http/integrations/router_heartbeat_test.go`:
- Around line 39-51: The goroutine in the heartbeat test discards the error
returned by io.ReadAll, allowing partial data to produce a false pass. Update
the readDone result and the goroutine around io.ReadAll to preserve both body
bytes and the read error, then assert the error is nil before checking for the
heartbeat frame.
---
Nitpick comments:
In `@core/utils.go`:
- Around line 657-664: Update the hostname validation lookup to call
net.DefaultResolver.LookupIPAddr instead of constructing a new net.Resolver,
while preserving the existing lookup context, error wrapping, and IP extraction
behavior.
- Around line 655-656: Update ValidateExternalURL and its DNS lookup path to
accept and propagate a caller context when creating the timeout, adding a
context-aware helper while retaining a context-free wrapper for background
callers. Update available call sites in handlers and framework loaders to pass
their request or loader contexts, using the existing context-aware symbol
throughout.
In `@transports/bifrost-http/handlers/skills_serving.go`:
- Around line 1324-1348: Add a table-driven regression test in the skills
serving handler tests that invokes lookupSkillByPathParam with a mock store
recording the context passed to GetSkillByName, then assert the context has a
deadline approximately lookupSkillByPathParamTimeout from invocation and is not
the request context. Cover the successful lookup path while preserving existing
handler behavior.
In `@transports/bifrost-http/integrations/router_heartbeat_test.go`:
- Around line 61-80: Add route-level tests that invoke handlePassthroughStream
for both SSE and non-SSE responses, asserting the proxied response body remains
unchanged for non-SSE content and receives the expected heartbeat behavior for
SSE content. Keep Test_passthroughHeartbeatEligible focused on predicate cases
and use the existing router/test helpers and response setup.
🪄 Autofix
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: CHILL
Plan: Pro Plus
Run ID: 555cd487-8e3e-47bc-986b-30b341041729
📒 Files selected for processing (9)
core/utils.gocore/utils_test.gotransports/bifrost-http/handlers/inference.gotransports/bifrost-http/handlers/skills_serving.gotransports/bifrost-http/integrations/router.gotransports/bifrost-http/integrations/router_heartbeat_test.gotransports/bifrost-http/lib/streamreader.gotransports/bifrost-http/lib/streamreader_test.gotransports/config.schema.json
🚧 Files skipped from review as they are similar to previous changes (4)
- transports/bifrost-http/lib/streamreader.go
- transports/bifrost-http/handlers/inference.go
- transports/bifrost-http/lib/streamreader_test.go
- transports/bifrost-http/integrations/router.go
akshaydeo
force-pushed
the
08-04-moves_sse_hearbeats_to_a_common_structure_to_reuse
branch
from
August 5, 2026 01:15
f52cd43 to
8bb3154
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
Author
Merge activity
|
akshaydeo
deleted the
08-04-moves_sse_hearbeats_to_a_common_structure_to_reuse
branch
August 5, 2026 01:39
This was referenced Aug 5, 2026
[Bug]: SSE streams can be severed by an intermediary idle timeout during long provider silence
#5212
Closed
2 tasks
This was referenced Aug 6, 2026
atharvamhaske
pushed a commit
to atharvamhaske/bifrost
that referenced
this pull request
Aug 13, 2026
## Summary Client-disconnect detection in streaming paths was purely reactive: `cancel()` only fired when a downstream write actually failed, which only happened when the producer loop attempted a write. Fast or bursty upstream providers (notably Vertex's `streamGenerateContent`, which delivers fewer, larger deltas than direct Gemini) could finish an entire stream before the write-failure detector ever got a second chance to fire, causing disconnected clients to be logged as false successes. This PR extracts the SSE heartbeat goroutine from `handleStreamingResponse` into a reusable `lib.StartSSEHeartbeat` / `lib.StopSSEHeartbeat` pair, and wires it into `handleStreaming` (both SSE and Bedrock branches) and `handlePassthroughStream` (SSE content-types only). Closes maximhq#5010 ## Changes - **`lib/streamreader.go`**: Added `StartSSEHeartbeat`, `StopSSEHeartbeat`, and `DefaultSSEHeartbeatInterval`. The heartbeat goroutine calls a caller-supplied `send` function on each tick; if `send` returns false (reader closed, i.e. disconnect discovered), it calls `onDisconnect` once and exits. `StopSSEHeartbeat` enforces the correct shutdown order — `close(done)` → `reader.Close()` → `<-exited` — before the caller calls `reader.Done()`, preventing a "send on closed channel" panic. - **`handlers/inference.go`**: Replaced the inline heartbeat goroutine in `handleStreamingResponse` with calls to `StartSSEHeartbeat` / `StopSSEHeartbeat`. - **`integrations/router.go`**: - `handleStreaming`: Added heartbeat coverage for both the SSE and Bedrock branches. The Bedrock branch uses a dedicated `*eventstream.Encoder` for the heartbeat goroutine (separate from the producer loop's encoder) because `eventstream.Encoder` reuses internal scratch buffers and is not safe for concurrent use. The heartbeat frame uses a synthetic `:event-type` (`bifrostHeartbeat`) that botocore's event-stream parser silently drops, making it the EventStream-binary equivalent of an SSE comment. - `handlePassthroughStream`: Heartbeat is only started when the resolved content-type is `text/event-stream`. Injecting SSE comment bytes into a non-SSE passthrough body (e.g. Vertex/Gemini's raw incrementally-delivered JSON array) would corrupt framing this path doesn't control. - Added `passthroughHeartbeatEligible` and `sendBedrockEventStreamHeartbeat` helpers. - **`lib/streamreader_test.go`**: Tests for `StartSSEHeartbeat` firing periodically and calling `onDisconnect` exactly once on a closed reader. - **`integrations/router_heartbeat_test.go`**: Tests confirming the SSE and Bedrock branches of `handleStreaming` emit heartbeat frames during idle gaps, that `bedrockHeartbeatEventType` never collides with real `BedrockStreamEvent.ToEncodedEvents()` event types, that concurrent use of separate encoders is race-detector clean, and that `passthroughHeartbeatEligible` gates correctly on content-type. ## Type of change - [x] Bug fix - [x] Refactor ## Affected areas - [x] Transports (HTTP) - [x] Providers/Integrations ## How to test ```sh go test ./transports/bifrost-http/lib/... ./transports/bifrost-http/integrations/... -race ``` The `router_heartbeat_test.go` tests use real timers to confirm heartbeat frames are emitted during idle gaps. Run with `-race` to validate the Bedrock encoder-isolation test (`Test_bedrockEventStreamHeartbeatUsesSeparateEncoderSafely`), which is specifically designed to be caught by the race detector if a single encoder were shared between the producer and heartbeat goroutines. ## Breaking changes - [x] No ## Security considerations None. The heartbeat frames are no-op probes (SSE comment lines or unrecognized EventStream event-types) that are invisible to conforming clients. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Client-disconnect detection in streaming paths was purely reactive:
cancel()only fired when a downstream write actually failed, which only happened when the producer loop attempted a write. Fast or bursty upstream providers (notably Vertex'sstreamGenerateContent, which delivers fewer, larger deltas than direct Gemini) could finish an entire stream before the write-failure detector ever got a second chance to fire, causing disconnected clients to be logged as false successes.This PR extracts the SSE heartbeat goroutine from
handleStreamingResponseinto a reusablelib.StartSSEHeartbeat/lib.StopSSEHeartbeatpair, and wires it intohandleStreaming(both SSE and Bedrock branches) andhandlePassthroughStream(SSE content-types only).Closes #5010
Changes
lib/streamreader.go: AddedStartSSEHeartbeat,StopSSEHeartbeat, andDefaultSSEHeartbeatInterval. The heartbeat goroutine calls a caller-suppliedsendfunction on each tick; ifsendreturns false (reader closed, i.e. disconnect discovered), it callsonDisconnectonce and exits.StopSSEHeartbeatenforces the correct shutdown order —close(done)→reader.Close()→<-exited— before the caller callsreader.Done(), preventing a "send on closed channel" panic.handlers/inference.go: Replaced the inline heartbeat goroutine inhandleStreamingResponsewith calls toStartSSEHeartbeat/StopSSEHeartbeat.integrations/router.go:handleStreaming: Added heartbeat coverage for both the SSE and Bedrock branches. The Bedrock branch uses a dedicated*eventstream.Encoderfor the heartbeat goroutine (separate from the producer loop's encoder) becauseeventstream.Encoderreuses internal scratch buffers and is not safe for concurrent use. The heartbeat frame uses a synthetic:event-type(bifrostHeartbeat) that botocore's event-stream parser silently drops, making it the EventStream-binary equivalent of an SSE comment.handlePassthroughStream: Heartbeat is only started when the resolved content-type istext/event-stream. Injecting SSE comment bytes into a non-SSE passthrough body (e.g. Vertex/Gemini's raw incrementally-delivered JSON array) would corrupt framing this path doesn't control.passthroughHeartbeatEligibleandsendBedrockEventStreamHeartbeathelpers.lib/streamreader_test.go: Tests forStartSSEHeartbeatfiring periodically and callingonDisconnectexactly once on a closed reader.integrations/router_heartbeat_test.go: Tests confirming the SSE and Bedrock branches ofhandleStreamingemit heartbeat frames during idle gaps, thatbedrockHeartbeatEventTypenever collides with realBedrockStreamEvent.ToEncodedEvents()event types, that concurrent use of separate encoders is race-detector clean, and thatpassthroughHeartbeatEligiblegates correctly on content-type.Type of change
Affected areas
How to test
go test ./transports/bifrost-http/lib/... ./transports/bifrost-http/integrations/... -raceThe
router_heartbeat_test.gotests use real timers to confirm heartbeat frames are emitted during idle gaps. Run with-raceto validate the Bedrock encoder-isolation test (Test_bedrockEventStreamHeartbeatUsesSeparateEncoderSafely), which is specifically designed to be caught by the race detector if a single encoder were shared between the producer and heartbeat goroutines.Breaking changes
Security considerations
None. The heartbeat frames are no-op probes (SSE comment lines or unrecognized EventStream event-types) that are invisible to conforming clients.
Checklist
docs/contributing/README.mdand followed the guidelines