adds pause resume flows for streaming calls

Author: akshaydeo

AGENTS.md

@@ -245,6 +245,13 @@ ctx.WithValue(key, value)    // Chainable variant
 
 **Gotcha**: `BlockRestrictedWrites()` silently drops writes to reserved keys. This prevents plugins from accidentally overwriting internal state.
 
+**Hard rule — never store stream-sized data in `BifrostContext`.** Context holds small handles only: IDs, durations, booleans, interface pointers. Any per-request state that scales with stream content (chunk buffers, accumulated payloads, replay queues, large per-request slices/maps) must live in a top-level manager keyed by `RequestID`, not in `ctx`. Reference implementations:
+
+- `framework/streaming.Accumulator` — owns a `sync.Map` of per-stream `StreamAccumulator` entries keyed by `RequestID`. Only `BifrostContextKeyAccumulatorID` (the ID string) is stored on the context; the chunk buffers live in the manager. The pause/resume gate (`gate.go`) extends the same per-stream entry with a state machine — again, **no buffer in ctx**.
+- The `Tracer` interface (in ctx as a small pointer) is the access path for plugins/providers to reach managers without putting bulky data on the context itself.
+
+When in doubt: if your new ctx key would hold a slice/map that grows with request content, route the storage through a manager and keep only the ID in ctx.
+
 ---
 
 ## Core Patterns

core/bifrost.go

@@ -4982,9 +4982,9 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
 					// Guarded send: if the consumer abandons outputStream (client
 					// disconnect, ctx cancel), drain the upstream shortCircuit.Stream
 					// so its producer can exit cleanly instead of blocking on its send.
-					select {
-					case outputStream <- streamResponse:
-					case <-ctx.Done():
+					// GateSendChunk routes through the pause/resume gate when a plugin
+					// has engaged it; otherwise it's a bare ctx-guarded channel send.
+					if !providerUtils.GateSendChunk(ctx, streamResponse, outputStream) {
 						for range shortCircuit.Stream {
 						}
 						return

core/providers/utils/utils.go

@@ -1876,6 +1876,35 @@ func BuildClientStreamChunk(ctx context.Context, processedResponse *schemas.Bifr
 	return streamResponse
 }
 
+// GateSendChunk routes a stream chunk through the tracer's pause/resume/end
+// gate ONLY when a plugin has engaged the gate for this stream (via
+// ctx.PauseStream/ResumeStream/EndStream, which sets BifrostContextKeyStreamGated).
+// Streams that never engage the gate (the overwhelmingly common case) take the
+// fast path: a direct channel send with ctx.Done() guard — same code as before
+// the gate was introduced, no extra lookups or locks.
+func GateSendChunk(ctx *schemas.BifrostContext, chunk *schemas.BifrostStreamChunk, responseChan chan *schemas.BifrostStreamChunk) bool {
+	if gated, _ := ctx.Value(schemas.BifrostContextKeyStreamGated).(bool); gated {
+		isFinal := false
+		if v := ctx.Value(schemas.BifrostContextKeyStreamEndIndicator); v != nil {
+			if b, ok := v.(bool); ok {
+				isFinal = b
+			}
+		}
+		isHardErr := chunk != nil && chunk.BifrostError != nil && chunk.BifrostError.IsBifrostError
+		if tracer, ok := ctx.Value(schemas.BifrostContextKeyTracer).(schemas.Tracer); ok && tracer != nil {
+			if traceID, ok := ctx.Value(schemas.BifrostContextKeyTraceID).(string); ok && traceID != "" {
+				return tracer.GateSend(traceID, chunk, isFinal, isHardErr, responseChan, ctx)
+			}
+		}
+	}
+	select {
+	case responseChan <- chunk:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}
+
 // ProcessAndSendResponse handles post-hook processing and sends the response to the channel.
 // This utility reduces code duplication across streaming implementations by encapsulating
 // the common pattern of running post hooks, handling errors, and sending responses with
@@ -1910,9 +1939,7 @@ func ProcessAndSendResponse(
 
 	streamResponse := BuildClientStreamChunk(ctx, processedResponse, processedError)
 
-	select {
-	case responseChan <- streamResponse:
-	case <-ctx.Done():
+	if !GateSendChunk(ctx, streamResponse, responseChan) {
 		return
 	}
 
@@ -1952,10 +1979,7 @@ func ProcessAndSendBifrostError(
 
 	streamResponse := BuildClientStreamChunk(ctx, processedResponse, processedError)
 
-	select {
-	case responseChan <- streamResponse:
-	case <-ctx.Done():
-	}
+	GateSendChunk(ctx, streamResponse, responseChan)
 
 	// Check if this is the final chunk and complete deferred span with post-processed data
 	if isFinalChunk := ctx.Value(schemas.BifrostContextKeyStreamEndIndicator); isFinalChunk != nil {
@@ -2208,10 +2232,7 @@ func ProcessAndSendError(
 		streamResponse.BifrostError = processedError
 	}
 
-	select {
-	case responseChan <- streamResponse:
-	case <-ctx.Done():
-	}
+	GateSendChunk(ctx, streamResponse, responseChan)
 }
 
 // CreateBifrostTextCompletionChunkResponse creates a bifrost text completion chunk response.

core/schemas/bifrost.go

@@ -194,6 +194,7 @@ const (
 	BifrostContextKeyNumberOfRetries                     BifrostContextKey = "bifrost-number-of-retries"             // int (to store the number of retries (set by bifrost - DO NOT SET THIS MANUALLY))
 	BifrostContextKeyFallbackIndex                       BifrostContextKey = "bifrost-fallback-index"                // int (to store the fallback index (set by bifrost - DO NOT SET THIS MANUALLY)) 0 for primary, 1 for first fallback, etc.
 	BifrostContextKeyStreamEndIndicator                  BifrostContextKey = "bifrost-stream-end-indicator"          // bool (set by bifrost - DO NOT SET THIS MANUALLY))
+	BifrostContextKeyStreamGated                         BifrostContextKey = "bifrost-stream-gated"                  // bool (set by ctx.PauseStream/ResumeStream/EndStream when a plugin first engages the pause/resume gate; provider helpers use this as a fast-path check to skip Tracer.GateSend on streams that never engage the gate)
 	BifrostContextKeyStreamIdleTimeout                   BifrostContextKey = "bifrost-stream-idle-timeout"           // time.Duration (per-chunk idle timeout for streaming)
 	BifrostContextKeySkipKeySelection                    BifrostContextKey = "bifrost-skip-key-selection"            // bool (will pass an empty key to the provider)
 	BifrostContextKeyExtraHeaders                        BifrostContextKey = "bifrost-extra-headers"                 // map[string][]string

core/schemas/context.go

@@ -354,6 +354,57 @@ func (bc *BifrostContext) GetParentCtxWithUserValues() context.Context {
 	return parentCtx
 }
 
+// PauseStream marks the active streaming response associated with this context
+// as paused. While paused, chunks continue to flow through PostLLMHook (so
+// plugins can still inspect them), but they are buffered instead of delivered
+// to the client. Buffered chunks are flushed in order when ResumeStream is
+// called. Idempotent. No-op if no Tracer or trace ID is present in ctx.
+//
+// Calling this method engages the pause/resume gate for the stream: provider
+// send sites switch from a direct channel send to Tracer.GateSend. Streams that
+// never call Pause/Resume/End pay no extra cost.
+func (bc *BifrostContext) PauseStream() {
+	tr, _ := bc.Value(BifrostContextKeyTracer).(Tracer)
+	tid, _ := bc.Value(BifrostContextKeyTraceID).(string)
+	if tr == nil || tid == "" {
+		return
+	}
+	bc.SetValue(BifrostContextKeyStreamGated, true)
+	tr.PauseStream(tid)
+}
+
+// ResumeStream resumes a previously paused stream. Buffered chunks are flushed
+// to the client in order, then live streaming continues. Idempotent. No-op if
+// no Tracer or trace ID is present in ctx.
+//
+// Engages the pause/resume gate (see PauseStream).
+func (bc *BifrostContext) ResumeStream() {
+	tr, _ := bc.Value(BifrostContextKeyTracer).(Tracer)
+	tid, _ := bc.Value(BifrostContextKeyTraceID).(string)
+	if tr == nil || tid == "" {
+		return
+	}
+	bc.SetValue(BifrostContextKeyStreamGated, true)
+	tr.ResumeStream(tid)
+}
+
+// EndStream terminates the active streaming response. Any buffered chunks are
+// flushed first; if err is non-nil it is then delivered as a final error chunk.
+// After EndStream returns, all further provider chunks for this stream are
+// dropped (PostLLMHook still fires, but no client delivery happens). Idempotent.
+// No-op if no Tracer or trace ID is present in ctx.
+//
+// Engages the pause/resume gate (see PauseStream).
+func (bc *BifrostContext) EndStream(err *BifrostError) {
+	tr, _ := bc.Value(BifrostContextKeyTracer).(Tracer)
+	tid, _ := bc.Value(BifrostContextKeyTraceID).(string)
+	if tr == nil || tid == "" {
+		return
+	}
+	bc.SetValue(BifrostContextKeyStreamGated, true)
+	tr.EndStream(tid, err)
+}
+
 // AppendRoutingEngineLog appends a routing engine log entry to the context.
 // Parameters:
 //   - ctx: The Bifrost context

core/schemas/tracer.go

@@ -114,6 +114,32 @@ type Tracer interface {
 	// The ctx parameter must contain the stream end indicator for proper final chunk detection.
 	ProcessStreamingChunk(traceID string, isFinalChunk bool, result *BifrostResponse, err *BifrostError) *StreamAccumulatorResult
 
+	// PauseStream marks the streaming response identified by traceID as paused.
+	// While paused, post-processed chunks are buffered (not delivered) but plugin
+	// hooks continue to fire. Idempotent. No-op if no active stream is found.
+	PauseStream(traceID string)
+
+	// ResumeStream resumes a previously paused stream. Buffered chunks are flushed
+	// to the client in order, then live streaming continues. Idempotent.
+	ResumeStream(traceID string)
+
+	// EndStream terminates the stream. If err is non-nil, it is delivered to the
+	// client as a final error chunk after any buffered chunks are flushed. After
+	// EndStream, all further chunks for this stream are dropped (post-hooks still
+	// run but no client delivery happens). Idempotent.
+	EndStream(traceID string, err *BifrostError)
+
+	// GateSend is called by stream producers (provider helpers) instead of writing
+	// directly to the response channel. It implements the pause/resume/end gate:
+	//   - Active state: chunk is forwarded to ch (with ctx.Done() guard)
+	//   - Paused state: chunk is buffered for later replay
+	//   - Ended state:  chunk is dropped
+	// Final chunks (isFinal) and hard provider errors (isHardErr) bypass the gate
+	// and force-flush + transition to Ended.
+	// Returns true if the chunk was handled (delivered or buffered), false if the
+	// caller should stop sending (ctx done or stream ended).
+	GateSend(traceID string, chunk *BifrostStreamChunk, isFinal, isHardErr bool, ch chan *BifrostStreamChunk, ctx *BifrostContext) bool
+
 	// AttachPluginLogs appends plugin log entries to the trace identified by traceID.
 	// Thread-safe. Should be called after plugin hooks complete, before trace completion.
 	AttachPluginLogs(traceID string, logs []PluginLogEntry)
@@ -187,6 +213,30 @@ func (n *NoOpTracer) ProcessStreamingChunk(_ string, _ bool, _ *BifrostResponse,
 	return nil
 }
 
+// PauseStream does nothing.
+func (n *NoOpTracer) PauseStream(_ string) {}
+
+// ResumeStream does nothing.
+func (n *NoOpTracer) ResumeStream(_ string) {}
+
+// EndStream does nothing.
+func (n *NoOpTracer) EndStream(_ string, _ *BifrostError) {}
+
+// GateSend forwards the chunk directly to the channel with ctx.Done() guard.
+// NoOpTracer has no gate state, so this is a pure passthrough.
+func (n *NoOpTracer) GateSend(_ string, chunk *BifrostStreamChunk, _ bool, _ bool, ch chan *BifrostStreamChunk, ctx *BifrostContext) bool {
+	if ctx == nil {
+		ch <- chunk
+		return true
+	}
+	select {
+	case ch <- chunk:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}
+
 // AttachPluginLogs does nothing.
 func (n *NoOpTracer) AttachPluginLogs(_ string, _ []PluginLogEntry) {}
 

framework/streaming/accumulator.go

@@ -158,7 +158,10 @@ func (a *Accumulator) createStreamAccumulator(requestID string) *StreamAccumulat
 		mu:                         sync.Mutex{},
 		Timestamp:                  now,
 		StartTimestamp:             now, // Set default StartTimestamp for proper TTFT/latency calculation
+		gateState:                  StreamStateActive,
+		gatePausedAt:               -1,
 	}
+	sc.gateCond = sync.NewCond(&sc.mu)
 	a.streamAccumulators.Store(requestID, sc)
 	return sc
 }
@@ -193,7 +196,10 @@ func (a *Accumulator) getOrCreateStreamAccumulator(requestID string) *StreamAccu
 		mu:                         sync.Mutex{},
 		Timestamp:                  now,
 		StartTimestamp:             now,
+		gateState:                  StreamStateActive,
+		gatePausedAt:               -1,
 	}
+	newAcc.gateCond = sync.NewCond(&newAcc.mu)
 
 	// LoadOrStore atomically: if key exists, return existing; else store new
 	actual, _ := a.streamAccumulators.LoadOrStore(requestID, newAcc)
@@ -361,6 +367,15 @@ func (a *Accumulator) cleanupStreamAccumulator(requestID string) {
 	if accumulator, exists := a.streamAccumulators.Load(requestID); exists {
 		acc := accumulator.(*StreamAccumulator)
 
+		// Force-end the gate so any running flusher goroutine wakes up and
+		// exits. Caller already holds acc.mu, so we can mutate gate fields.
+		if acc.gateState != StreamStateEnded {
+			acc.gateState = StreamStateEnded
+			if acc.gateCond != nil {
+				acc.gateCond.Broadcast()
+			}
+		}
+
 		// Return all chunks to the pool before deleting
 		for _, chunk := range acc.ChatStreamChunks {
 			a.putChatStreamChunk(chunk)