moves sse hearbeats to a common structure to reuse

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

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

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

transports/bifrost-http/handlers/skills_serving.go

@@ -1321,14 +1321,31 @@ 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")
 	if !ok {
 		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))

transports/bifrost-http/integrations/router.go

@@ -2737,23 +2737,45 @@ 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 {
 				traceCompleter(nil)
 			}
 		}()
 
-		// 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)
 			}