From 8bb3154ab84ef452e79cea86b351e3be70b40821 Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Tue, 4 Aug 2026 15:39:09 -0700 Subject: [PATCH] moves sse hearbeats to a common structure to reuse --- core/utils.go | 25 +++++- core/utils_test.go | 77 ++++++++++++++++ transports/bifrost-http/handlers/inference.go | 65 ++++---------- .../bifrost-http/handlers/skills_serving.go | 19 +++- .../bifrost-http/integrations/router.go | 70 +++++++++++++-- .../integrations/router_heartbeat_test.go | 88 +++++++++++++++++++ transports/bifrost-http/lib/streamreader.go | 57 ++++++++++++ .../bifrost-http/lib/streamreader_test.go | 71 +++++++++++++++ transports/config.schema.json | 6 ++ 9 files changed, 416 insertions(+), 62 deletions(-) create mode 100644 transports/bifrost-http/integrations/router_heartbeat_test.go diff --git a/core/utils.go b/core/utils.go index 84353e9fbc0..b8d8078c625 100644 --- a/core/utils.go +++ b/core/utils.go @@ -621,6 +621,19 @@ func RedactSensitiveString(s string) string { return s[:4] + "[REDACTED]" + s[len(s)-4:] } +// externalURLDNSLookupTimeout bounds the DNS resolution in ValidateExternalURL. The +// package-level net.LookupIP has no context/deadline of its own, so a hung or blackholed +// resolver (rather than a fast refusal) can block the caller indefinitely -- this timeout +// closes that gap regardless of the caller's own retry/timeout logic, since those only +// bound the HTTP request that follows resolution, not resolution itself. +const externalURLDNSLookupTimeout = 5 * time.Second + +// lookupIPAddr is a seam over (&net.Resolver{}).LookupIPAddr so tests can substitute a +// resolver that blocks until its context is canceled, proving externalURLDNSLookupTimeout +// actually cuts off an in-flight lookup rather than only a lookup whose context had +// already expired before it started. +var lookupIPAddr = (&net.Resolver{}).LookupIPAddr + // ValidateExternalURL validates a URL for security concerns (SSRF protection). // When allowPrivateNetwork is true, RFC 1918 private IPs are permitted (for k8s/LAN deployments). // Link-local addresses (169.254.x.x, fe80::) are always blocked regardless of allowPrivateNetwork. @@ -642,11 +655,19 @@ func ValidateExternalURL(urlStr string, allowPrivateNetwork bool) error { if hostname == "" { return fmt.Errorf("URL must have a hostname") } - // Resolve hostname to IP addresses - ips, err := net.LookupIP(hostname) + // Resolve hostname to IP addresses. Bounded via net.Resolver.LookupIPAddr (not the + // package-level net.LookupIP, which has no way to accept a deadline) so a stalled + // resolver fails fast instead of hanging the caller forever. + lookupCtx, cancel := context.WithTimeout(context.Background(), externalURLDNSLookupTimeout) + defer cancel() + addrs, err := lookupIPAddr(lookupCtx, hostname) if err != nil { return fmt.Errorf("failed to resolve hostname: %w", err) } + ips := make([]net.IP, len(addrs)) + for i, addr := range addrs { + ips[i] = addr.IP + } for _, ip := range ips { if ip.IsLoopback() { continue diff --git a/core/utils_test.go b/core/utils_test.go index 15d010e719b..787654ce54e 100644 --- a/core/utils_test.go +++ b/core/utils_test.go @@ -1,9 +1,13 @@ package bifrost import ( + "context" + "errors" "net" "strings" "testing" + "testing/synctest" + "time" "github.com/maximhq/bifrost/core/network" ) @@ -201,6 +205,79 @@ func TestValidateExternalURL(t *testing.T) { } } +// TestValidateExternalURLBoundsDNSLookup guards against regressing to the package-level +// net.LookupIP (which has no context/deadline of its own, and previously let a stalled +// resolver block the caller indefinitely -- see externalURLDNSLookupTimeout's doc). A +// hostname under the .invalid TLD (reserved by RFC 2606 to never resolve) still exercises +// the real resolution path; the assertion is on wall-clock bound, not on the specific +// error, since what matters is that resolution can never run unbounded. +func TestValidateExternalURLBoundsDNSLookup(t *testing.T) { + start := time.Now() + err := ValidateExternalURL("https://this-host-does-not-exist-12345.invalid", false) + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected an error resolving a .invalid hostname, got nil") + } + // Generous margin above externalURLDNSLookupTimeout (5s): proves resolution is bounded + // by that timeout rather than by an unrelated, much longer OS/runtime resolver ceiling. + if elapsed > 8*time.Second { + t.Errorf("ValidateExternalURL took %v to fail on an unresolvable host, want well under externalURLDNSLookupTimeout + margin", elapsed) + } +} + +// TestValidateExternalURLBoundsInFlightDNSLookup complements +// TestValidateExternalURLBoundsDNSLookup: a .invalid hostname NXDOMAINs almost immediately, +// so it never proves the timeout can cut off a lookup that's actually in flight -- the exact +// scenario externalURLDNSLookupTimeout exists to guard against (a hung/blackholed resolver). +// This substitutes a resolver that blocks until its context is canceled, so the only way +// ValidateExternalURL returns is via externalURLDNSLookupTimeout firing. Runs inside a +// synctest bubble -- context.WithTimeout is explicitly supported there (see the +// testing/synctest package docs' Context.WithTimeout example) -- so the 5s timeout resolves +// against a deterministic fake clock instead of a real 5-second wait. +func TestValidateExternalURLBoundsInFlightDNSLookup(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + original := lookupIPAddr + defer func() { lookupIPAddr = original }() + + lookupIPAddr = func(ctx context.Context, host string) ([]net.IPAddr, error) { + <-ctx.Done() + return nil, ctx.Err() + } + + err := ValidateExternalURL("https://api.openai.com", false) + if err == nil { + t.Fatal("expected an error when the resolver blocks past the timeout, got nil") + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected error to wrap context.DeadlineExceeded (proving externalURLDNSLookupTimeout fired), got %v", err) + } + }) +} + +// TestResolverLookupIPAddrRespectsContextDeadline pins down the stdlib behavior +// ValidateExternalURL's fix relies on: unlike the package-level net.LookupIP (which +// always uses context.Background() internally and ignores any external deadline), +// net.Resolver.LookupIPAddr honors a context deadline and returns promptly once it +// elapses, even against a real hostname that would otherwise resolve. +func TestResolverLookupIPAddrRespectsContextDeadline(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancel() + // Let the deadline definitely elapse before the lookup starts. + time.Sleep(time.Millisecond) + + start := time.Now() + _, err := (&net.Resolver{}).LookupIPAddr(ctx, "api.openai.com") + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected LookupIPAddr to fail against an already-expired context, got nil error") + } + if elapsed > time.Second { + t.Errorf("LookupIPAddr took %v to return after an already-expired deadline, want near-instant", elapsed) + } +} + func TestIsLocalhost(t *testing.T) { tests := []struct { hostname string diff --git a/transports/bifrost-http/handlers/inference.go b/transports/bifrost-http/handlers/inference.go index 6dc27c53b18..ab0e98d71ba 100644 --- a/transports/bifrost-http/handlers/inference.go +++ b/transports/bifrost-http/handlers/inference.go @@ -1897,12 +1897,6 @@ func (h *CompletionHandler) handleStreamingTranscriptionRequest(ctx *fasthttp.Re h.handleStreamingResponse(ctx, bifrostCtx, schemas.TranscriptionStreamRequest, getStream, cancel) } -// streamHeartbeatInterval is how often handleStreamingResponse's heartbeat goroutine -// probes the downstream connection with a no-op SSE comment. See that goroutine's -// comment for why this exists: disconnect detection is otherwise purely reactive to -// real-data write failures, which a fast/bursty upstream can outrun entirely. -const streamHeartbeatInterval = 100 * time.Millisecond - // handleStreamingResponse is a generic function to handle streaming responses using Server-Sent Events (SSE) // The cancel function is called ONLY when client disconnects are detected via write errors. // Bifrost handles cleanup internally for normal completion and errors, so we only cancel @@ -2009,25 +2003,23 @@ func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, bi transportLogs = logs } - // heartbeatDone asks the keep-alive goroutine below to stop once this producer - // goroutine exits, via any path (normal completion, disconnect, interceptor error). - // heartbeatExited is closed BY that goroutine right before it returns -- the deferred - // cleanup below waits on it before calling reader.Done(), so eventCh is never closed - // while the heartbeat goroutine could still be mid-send on it (see reader.Close() call - // below for why closing heartbeatDone alone isn't enough). - heartbeatDone := make(chan struct{}) - heartbeatExited := make(chan struct{}) + // Client-disconnect detection is otherwise purely reactive: cancel() only fires when + // a downstream SendEvent write actually fails, which only happens when this loop + // attempts one. If the upstream provider delivers its whole response in a few + // large/fast chunks, this loop may never attempt another write during the window + // where the client has already disconnected, so the disconnect goes undetected and + // the request logs as a false success (observed live: Vertex's streamGenerateContent + // delivers fewer, larger deltas than direct Gemini for the same prompt, giving the + // write-failure detector too few chances to fire before the stream finished). + // A periodic no-op heartbeat forces an extra write attempt during otherwise-idle + // gaps, closing that window without touching fasthttp's connection internals. + heartbeatDone, heartbeatExited := lib.StartSSEHeartbeat(lib.DefaultSSEHeartbeatInterval, reader.SendHeartbeat, cancel) defer func() { - close(heartbeatDone) - // heartbeatDone only unblocks the heartbeat goroutine if it's waiting at its - // outer select; if it's currently blocked inside SendHeartbeat -> Send's - // `eventCh <- event` case (channel full, nothing reading), closing heartbeatDone - // does nothing -- that goroutine isn't looking at it. reader.Close() closes - // closeCh instead, which Send's inner select already watches, so it safely - // unblocks a pending send without racing eventCh's close below. - reader.Close() - <-heartbeatExited + // Must run before reader.Done(): closing eventCh while the heartbeat goroutine + // could still be mid-send on it panics ("send on closed channel"). See + // lib.StopSSEHeartbeat's doc for the full ordering rationale. + lib.StopSSEHeartbeat(reader, heartbeatDone, heartbeatExited) schemas.ReleaseHTTPRequest(httpReq) // Fallback: on early-return paths (client disconnect, interceptor error) // we never reached the pre-[DONE] invocation, so run it now. Any error is @@ -2044,33 +2036,6 @@ func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, bi } }() - // Client-disconnect detection above is purely reactive: cancel() only fires when - // a downstream SendEvent write actually fails, which only happens when this loop - // attempts one. If the upstream provider delivers its whole response in a few - // large/fast chunks, this loop may never attempt another write during the window - // where the client has already disconnected, so the disconnect goes undetected and - // the request logs as a false success (observed live: Vertex's streamGenerateContent - // delivers fewer, larger deltas than direct Gemini for the same prompt, giving the - // write-failure detector too few chances to fire before the stream finished). - // A periodic no-op heartbeat forces an extra write attempt during otherwise-idle - // gaps, closing that window without touching fasthttp's connection internals. - go func() { - defer close(heartbeatExited) - ticker := time.NewTicker(streamHeartbeatInterval) - defer ticker.Stop() - for { - select { - case <-ticker.C: - if !reader.SendHeartbeat() { - cancel() // Disconnect discovered via the heartbeat, not real data. - return - } - case <-heartbeatDone: - return - } - } - }() - var includeEventType bool var skipDoneMarker bool // Set once a hard error frame reaches the client. An error frame is a diff --git a/transports/bifrost-http/handlers/skills_serving.go b/transports/bifrost-http/handlers/skills_serving.go index 435dc9c9997..17f248b503c 100644 --- a/transports/bifrost-http/handlers/skills_serving.go +++ b/transports/bifrost-http/handlers/skills_serving.go @@ -1321,6 +1321,10 @@ func buildSkillFilePath(skillName string, file *tables.TableSkillFile) string { return path.Join(skillName, file.Path) } +// lookupSkillByPathParamTimeout bounds the DB lookup in lookupSkillByPathParam. It must +// never be derived from the request's *fasthttp.RequestCtx (see that function's comment). +const lookupSkillByPathParamTimeout = 10 * time.Second + // lookupSkillByPathParam extracts the skill-name path parameter and fetches the skill. func (h *SkillsServingHandler) lookupSkillByPathParam(ctx *fasthttp.RequestCtx) (*tables.TableSkill, bool) { name, ok := decodeStringPathParam(ctx, "skill-name", "skill name") @@ -1328,7 +1332,20 @@ func (h *SkillsServingHandler) lookupSkillByPathParam(ctx *fasthttp.RequestCtx) return nil, false } - skill, err := h.store.GetSkillByName(ctx, name) + // Must not pass ctx (a *fasthttp.RequestCtx) directly as the context.Context here: its + // Done() returns a server-wide channel (fasthttp.Server.done), closed only on server + // shutdown -- not per-request (fasthttp's own documented tradeoff, since allocating a + // channel per request is expensive). GetSkillByName's nested Preload("Files")/ + // Preload("Files.Blob") queries make database/sql spawn an internal cancellation-watcher + // goroutine per query whenever ctx.Done() != nil, and that goroutine reads + // RequestCtx.s.done unsynchronized against Server.Shutdown()'s write of s.done = nil -- + // a data race confirmed under -race. Deriving from context.Background() instead (as + // allSkillsZipDownload/genericZipDownload already do for their streaming bodies) avoids + // ever handing the raw RequestCtx to anything that watches Done() asynchronously. + lookupCtx, cancel := context.WithTimeout(context.Background(), lookupSkillByPathParamTimeout) + defer cancel() + + skill, err := h.store.GetSkillByName(lookupCtx, name) if err != nil { if errors.Is(err, configstore.ErrNotFound) { SendError(ctx, fasthttp.StatusNotFound, fmt.Sprintf("skill %q not found", name)) diff --git a/transports/bifrost-http/integrations/router.go b/transports/bifrost-http/integrations/router.go index 1cfbdcff556..0744c1a1f36 100644 --- a/transports/bifrost-http/integrations/router.go +++ b/transports/bifrost-http/integrations/router.go @@ -2737,10 +2737,38 @@ func (g *GenericRouter) handleStreaming(ctx *fasthttp.RequestCtx, bifrostCtx *sc // Producer goroutine: processes the stream channel, formats events, sends to reader go func() { - // Separate defers ensure each cleanup runs even if an earlier one panics (LIFO order) - defer reader.Done() - defer schemas.ReleaseHTTPRequest(httpReq) + // Create encoder for AWS Event Stream if needed + var eventStreamEncoder *eventstream.Encoder + if config.Type == RouteConfigTypeBedrock { + eventStreamEncoder = eventstream.NewEncoder() + } + + // Client-disconnect detection below is otherwise purely reactive: cancel() only + // fires when a downstream write actually fails, which only happens when this loop + // attempts one. A periodic no-op heartbeat forces an extra write attempt during + // otherwise-idle gaps, closing the window where a fast/bursty upstream finishes + // before a disconnect is ever discovered (see lib.StartSSEHeartbeat's doc). + // + // Bedrock uses binary AWS EventStream framing, not SSE, and has no safe no-op frame: + // real AWS Bedrock never emits synthetic frames, and unlike botocore, other official + // AWS SDKs (e.g. Go SDK v2) don't silently drop an unmodeled `:event-type` -- they + // surface it to the caller as a typed union member (types.UnknownUnionMember). So + // Bedrock streams stay on purely reactive (write-failure-based) disconnect detection. + var heartbeatDone chan struct{} + var heartbeatExited <-chan struct{} + if config.Type != RouteConfigTypeBedrock { + heartbeatDone, heartbeatExited = lib.StartSSEHeartbeat(lib.DefaultSSEHeartbeatInterval, reader.SendHeartbeat, cancel) + } + defer func() { + // Must run before reader.Done(): closing eventCh while the heartbeat goroutine + // could still be mid-send on it panics ("send on closed channel"). See + // lib.StopSSEHeartbeat's doc for the full ordering rationale. + if heartbeatDone != nil { + lib.StopSSEHeartbeat(reader, heartbeatDone, heartbeatExited) + } + schemas.ReleaseHTTPRequest(httpReq) + reader.Done() // Complete the trace after streaming finishes // This ensures all spans (including llm.call) are properly ended before the trace is sent to OTEL if traceCompleter != nil { @@ -2748,12 +2776,6 @@ func (g *GenericRouter) handleStreaming(ctx *fasthttp.RequestCtx, bifrostCtx *sc } }() - // Create encoder for AWS Event Stream if needed - var eventStreamEncoder *eventstream.Encoder - if config.Type == RouteConfigTypeBedrock { - eventStreamEncoder = eventstream.NewEncoder() - } - // sendConvertedStreamError converts a sanitized BifrostError through the // integration's error converter and emits it in the route's native SSE // format. Callers decide how to terminate the stream afterwards, with @@ -3291,6 +3313,18 @@ func (g *GenericRouter) handlePassthroughNonStream( ctx.Response.SetBody(resp.Body) } +// passthroughHeartbeatEligible reports whether a resolved passthrough response +// content-type is safe for the SSE-comment heartbeat. handlePassthroughStream proxies +// raw upstream bytes 1:1, and content-type isn't always SSE -- e.g. Vertex/Gemini's +// non-alt=sse mode returns an incrementally-delivered JSON array with +// Content-Type: application/json. The media type is compared exactly (case-insensitively, +// ignoring parameters like "; charset=utf-8") so lookalikes such as +// "text/event-stream+json" or "text/event-streaming" aren't treated as SSE. +func passthroughHeartbeatEligible(contentType string) bool { + mediaType, _, _ := strings.Cut(contentType, ";") + return strings.EqualFold(strings.TrimSpace(mediaType), "text/event-stream") +} + func (g *GenericRouter) handlePassthroughStream( ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, @@ -3387,8 +3421,26 @@ func (g *GenericRouter) handlePassthroughStream( reader := lib.NewSSEStreamReader() ctx.Response.SetBodyStream(reader, -1) + // This path proxies raw upstream bytes 1:1, and content-type isn't always SSE (see the + // resolution above) -- e.g. Vertex/Gemini's non-alt=sse mode returns an incrementally + // delivered JSON array. Injecting an SSE comment or any other bytes into a non-SSE stream + // whose framing we don't control would corrupt it: unlike an SSE comment line or an + // unrecognized Bedrock EventStream event-type, there is no protocol guarantee that extra + // bytes here are safely ignorable. So the heartbeat only runs when the resolved + // content-type is actually SSE. + var heartbeatDone chan struct{} + var heartbeatExited <-chan struct{} + if passthroughHeartbeatEligible(contentType) { + heartbeatDone, heartbeatExited = lib.StartSSEHeartbeat(lib.DefaultSSEHeartbeatInterval, reader.SendHeartbeat, cancel) + } + go func() { defer func() { + if heartbeatDone != nil { + // Must run before reader.Done(): closing eventCh while the heartbeat + // goroutine could still be mid-send on it panics ("send on closed channel"). + lib.StopSSEHeartbeat(reader, heartbeatDone, heartbeatExited) + } if traceCompleter != nil { traceCompleter(nil) } diff --git a/transports/bifrost-http/integrations/router_heartbeat_test.go b/transports/bifrost-http/integrations/router_heartbeat_test.go new file mode 100644 index 00000000000..0e74e11bb2d --- /dev/null +++ b/transports/bifrost-http/integrations/router_heartbeat_test.go @@ -0,0 +1,88 @@ +package integrations + +import ( + "io" + "testing" + "testing/synctest" + "time" + + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/transports/bifrost-http/lib" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/valyala/fasthttp" +) + +// A fast/bursty upstream can finish a stream before the reactive (write-failure-based) +// disconnect detector ever gets a second chance to fire (see lib.StartSSEHeartbeat's doc, +// and the inference.go handleStreamingResponse feature this mirrors). This test confirms +// handleStreaming's default (SSE) branch carries that same heartbeat coverage. Bedrock's +// route branch intentionally has no heartbeat -- see the comment at its call site in +// handleStreaming for why a synthetic AWS EventStream frame isn't safe there. + +// Test_handleStreamingSSESendsHeartbeatDuringIdleGap verifies the default (SSE) route +// branch of handleStreaming emits SendHeartbeat's comment frame while the stream channel +// is held open, so a client-disconnect during a long idle gap is discovered proactively. +// Runs inside a synctest bubble so the sleep below advances the bubble's fake clock +// deterministically instead of real wall-clock time -- a busy CI worker delaying the +// heartbeat ticker or reader goroutine can no longer make this flaky, unlike a real +// time.Sleep. +func Test_handleStreamingSSESendsHeartbeatDuringIdleGap(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + stream := make(chan *schemas.BifrostStreamChunk) + router := NewGenericRouter(nil, &mockHandlerStore{}, nil, nil, bifrost.NewNoOpLogger()) + ctx := &fasthttp.RequestCtx{} + router.handleStreaming(ctx, nil, RouteConfig{}, stream, func() {}) + + bodyStream := ctx.Response.BodyStream() + type readResult struct { + body string + err error + } + readDone := make(chan readResult, 1) + go func() { + b, err := io.ReadAll(bodyStream) + readDone <- readResult{body: string(b), err: err} + }() + + // > 1 tick at lib.DefaultSSEHeartbeatInterval before the stream ends. Computed + // relative to the constant (rather than hardcoded) so this stays correct if the + // default interval changes. + time.Sleep(2*lib.DefaultSSEHeartbeatInterval + time.Millisecond) + close(stream) + + result := <-readDone + require.NoError(t, result.err) + assert.Contains(t, result.body, ": heartbeat\n\n", "expected at least one heartbeat comment frame during the idle gap") + }) +} + +// Test_passthroughHeartbeatEligible pins down the content-type gate handlePassthroughStream +// uses to decide whether the heartbeat is safe to inject: only when the resolved +// content-type is actually SSE. Injecting anything into a non-SSE passthrough body (e.g. +// Vertex/Gemini's raw incrementally-delivered JSON array) would corrupt a framing this +// path doesn't control, and lookalike media types (e.g. "text/event-stream+json") must not +// be mistaken for SSE either. +func Test_passthroughHeartbeatEligible(t *testing.T) { + cases := []struct { + name string + contentType string + want bool + }{ + {"plain SSE", "text/event-stream", true}, + {"SSE with charset param", "text/event-stream; charset=utf-8", true}, + {"SSE uppercase", "TEXT/EVENT-STREAM", true}, + {"SSE with surrounding space before param", "text/event-stream ; charset=utf-8", true}, + {"raw JSON passthrough", "application/json", false}, + {"empty content-type", "", false}, + {"binary passthrough", "application/vnd.amazon.eventstream", false}, + {"prefix-collision suffix", "text/event-stream+json", false}, + {"prefix-collision longer type", "text/event-streaming", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, passthroughHeartbeatEligible(tc.contentType)) + }) + } +} diff --git a/transports/bifrost-http/lib/streamreader.go b/transports/bifrost-http/lib/streamreader.go index d05f0b5d1ac..484d31ecfc1 100644 --- a/transports/bifrost-http/lib/streamreader.go +++ b/transports/bifrost-http/lib/streamreader.go @@ -3,6 +3,7 @@ package lib import ( "io" "sync" + "time" ) // SSEStreamReader is an io.ReadCloser that delivers one event per Read call, @@ -130,3 +131,59 @@ func (r *SSEStreamReader) SendHeartbeat() bool { func (r *SSEStreamReader) Done() { close(r.eventCh) } + +// DefaultSSEHeartbeatInterval is how often a heartbeat goroutine started with +// StartSSEHeartbeat probes the downstream connection by default. This helper now serves +// every idle inference, routed streaming, and SSE passthrough stream, so the interval +// trades off disconnect-detection latency against downstream write volume: at N +// concurrently idle streams, a shorter interval means up to N/interval extra writes per +// second carrying no provider data. One second keeps that volume low while still closing +// the disconnect-detection gap described in SendHeartbeat's doc; short intervals belong +// only in focused lifecycle tests, not this default. +const DefaultSSEHeartbeatInterval = time.Second + +// StartSSEHeartbeat launches a goroutine that calls send on every tick of interval, +// purely to force a downstream write attempt during otherwise-idle gaps. Client-disconnect +// detection is otherwise reactive: it only fires when a producer actually attempts a +// write, which a fast/bursty producer may not do again during the exact window a client +// disconnects (see SendHeartbeat's doc for the full rationale). If send returns false +// (the reader was closed, i.e. a disconnect was discovered), onDisconnect is called once +// and the goroutine exits. +// +// Callers MUST shut this down via StopSSEHeartbeat before calling reader.Done() -- never +// by closing the returned done channel directly. See StopSSEHeartbeat for why. +func StartSSEHeartbeat(interval time.Duration, send func() bool, onDisconnect func()) (done chan struct{}, exited <-chan struct{}) { + doneCh := make(chan struct{}) + exitedCh := make(chan struct{}) + go func() { + defer close(exitedCh) + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + if !send() { + onDisconnect() + return + } + case <-doneCh: + return + } + } + }() + return doneCh, exitedCh +} + +// StopSSEHeartbeat performs the safe shutdown sequence for a heartbeat started with +// StartSSEHeartbeat. Closing done alone only unblocks the heartbeat goroutine if it's +// waiting at its outer select; if it's currently blocked inside send (e.g. Send's +// `eventCh <- event` case, buffer full, nothing reading), closing done does nothing -- +// that goroutine isn't looking at it. reader.Close() closes closeCh instead, which +// Send's inner select already watches, safely unblocking a pending send without racing +// eventCh's close. Callers must call this BEFORE reader.Done(): closing eventCh while +// the heartbeat goroutine could still be mid-send on it panics ("send on closed channel"). +func StopSSEHeartbeat(reader *SSEStreamReader, done chan struct{}, exited <-chan struct{}) { + close(done) + reader.Close() + <-exited +} diff --git a/transports/bifrost-http/lib/streamreader_test.go b/transports/bifrost-http/lib/streamreader_test.go index 3ddb7aea57a..2ccc9ba1b57 100644 --- a/transports/bifrost-http/lib/streamreader_test.go +++ b/transports/bifrost-http/lib/streamreader_test.go @@ -4,6 +4,7 @@ import ( "fmt" "io" "sync" + "sync/atomic" "testing" "testing/synctest" "time" @@ -809,3 +810,73 @@ func TestSSEStreamReaderCloseUnblocksProducer(t *testing.T) { default: } } + +// TestStartSSEHeartbeatFiresPeriodically verifies the extracted heartbeat goroutine +// (used by both handleStreamingResponse and router.go's handleStreaming) sends more than +// one frame while a reader is actively drained, and that StopSSEHeartbeat followed by +// Done() shuts it down cleanly. +// Runs inside a synctest bubble so the sleep below advances the bubble's fake clock +// deterministically instead of real wall-clock time -- a busy CI worker delaying the +// ticker or reader goroutine can no longer make this flaky, unlike a real time.Sleep. +func TestStartSSEHeartbeatFiresPeriodically(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + r := NewSSEStreamReader() + + var received atomic.Int32 + drainDone := make(chan struct{}) + go func() { + defer close(drainDone) + buf := make([]byte, 4096) + for { + n, err := r.Read(buf) + if n > 0 { + received.Add(1) + } + if err != nil { + return + } + } + }() + + // onDisconnect is intentionally a no-op here, not an assertion that it never fires: + // StopSSEHeartbeat's reader.Close() call can legitimately unblock a heartbeat that's + // mid-Send (buffer full, drain lagging) with a false return, the same as a real + // disconnect would -- see TestSSEStreamReaderDoneWhileSendBlockedDoesNotPanic. In + // production onDisconnect is cancel(), which is idempotent and safe to call during an + // already-ending stream, so this is benign, not a bug. + done, exited := StartSSEHeartbeat(5*time.Millisecond, r.SendHeartbeat, func() {}) + + time.Sleep(60 * time.Millisecond) // several ticks at the 5ms interval + + StopSSEHeartbeat(r, done, exited) + r.Done() + <-drainDone + + if got := received.Load(); got < 2 { + t.Errorf("expected at least 2 heartbeat frames in 60ms at a 5ms interval, got %d", got) + } + }) +} + +// TestStartSSEHeartbeatCallsOnDisconnectAfterClose verifies the goroutine detects a closed +// reader (SendHeartbeat returning false) and invokes onDisconnect exactly once, mirroring +// how handleStreamingResponse's cancel() is wired to heartbeat-discovered disconnects. +func TestStartSSEHeartbeatCallsOnDisconnectAfterClose(t *testing.T) { + r := NewSSEStreamReader() + r.Close() // SendHeartbeat will return false on the very first tick + + var onDisconnectCalls atomic.Int32 + _, exited := StartSSEHeartbeat(5*time.Millisecond, r.SendHeartbeat, func() { + onDisconnectCalls.Add(1) + }) + + select { + case <-exited: + case <-time.After(time.Second): + t.Fatal("heartbeat goroutine never exited after the reader was closed") + } + + if got := onDisconnectCalls.Load(); got != 1 { + t.Errorf("onDisconnect calls = %d, want exactly 1", got) + } +} diff --git a/transports/config.schema.json b/transports/config.schema.json index 63262950a21..cf421a390f9 100644 --- a/transports/config.schema.json +++ b/transports/config.schema.json @@ -681,6 +681,12 @@ "type": "boolean", "description": "Snap the customer's budget and rate-limit reset windows to clean calendar boundaries (day, week, month, year)", "default": false + }, + "virtual_key_count": { + "type": "integer", + "description": "Computed count of virtual keys associated with this customer", + "minimum": 0, + "readOnly": true } }, "required": ["id", "name"],