feat: add generic pricing time schedules

core/bifrost.go

@@ -797,9 +797,11 @@ func (bifrost *Bifrost) makeChatCompletionRequest(ctx *schemas.BifrostContext, r
 
 // ChatCompletionRequest sends a chat completion request to the specified provider.
 func (bifrost *Bifrost) ChatCompletionRequest(ctx *schemas.BifrostContext, req *schemas.BifrostChatRequest) (*schemas.BifrostChatResponse, *schemas.BifrostError) {
-	// If ctx is nil, use the bifrost context (defensive check for mcp agent mode)
+	// Isolate nil-context callers so concurrent requests cannot share mutable
+	// per-request state such as the billing attempt timestamp.
 	if ctx == nil {
-		ctx = bifrost.ctx
+		ctx = schemas.NewBifrostContext(bifrost.ctx, schemas.NoDeadline)
+		defer ctx.Cancel()
 	}
 
 	response, err := bifrost.makeChatCompletionRequest(ctx, req)
@@ -899,9 +901,11 @@ func (bifrost *Bifrost) makeResponsesRequest(ctx *schemas.BifrostContext, req *s
 
 // ResponsesRequest sends a responses request to the specified provider.
 func (bifrost *Bifrost) ResponsesRequest(ctx *schemas.BifrostContext, req *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) {
-	// If ctx is nil, use the bifrost context (defensive check for mcp agent mode)
+	// Isolate nil-context callers so concurrent requests cannot share mutable
+	// per-request state such as the billing attempt timestamp.
 	if ctx == nil {
-		ctx = bifrost.ctx
+		ctx = schemas.NewBifrostContext(bifrost.ctx, schemas.NoDeadline)
+		defer ctx.Cancel()
 	}
 
 	response, err := bifrost.makeResponsesRequest(ctx, req)
@@ -5187,6 +5191,33 @@ func populateLatencyExtraFields(ctx *schemas.BifrostContext, resp *schemas.Bifro
 	if start, ok := ctx.Value(schemas.BifrostContextKeyRequestStartTime).(time.Time); ok {
 		resp.PopulateOverheadLatency(ctx, time.Since(start))
 	}
+	populateBillingAttemptResponseExtraFields(ctx, resp)
+}
+
+// populateBillingAttemptResponseExtraFields stamps the current attempt's start
+// time on a response. Post-hooks may replace responses without copying metadata,
+// so callers also use this between hooks to preserve the authoritative value.
+func populateBillingAttemptResponseExtraFields(ctx *schemas.BifrostContext, resp *schemas.BifrostResponse) {
+	if resp == nil {
+		return
+	}
+	if start, ok := ctx.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time); ok {
+		startCopy := start
+		resp.GetExtraFields().BillingAttemptStartedAt = &startCopy
+	}
+}
+
+// populateBillingAttemptExtraFields stamps the current attempt's start time on
+// errors carrying billed usage. It is separate from response latency population
+// because errors may be produced before a BifrostResponse exists.
+func populateBillingAttemptExtraFields(ctx *schemas.BifrostContext, bifrostErr *schemas.BifrostError) {
+	if bifrostErr == nil {
+		return
+	}
+	if start, ok := ctx.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time); ok {
+		startCopy := start
+		bifrostErr.ExtraFields.BillingAttemptStartedAt = &startCopy
+	}
 }
 
 // handleRequest handles the request to the provider based on the request type
@@ -5197,13 +5228,19 @@ func (bifrost *Bifrost) handleRequest(ctx *schemas.BifrostContext, req *schemas.
 	defer bifrost.releaseBifrostRequest(req)
 	provider, model, fallbacks := req.GetRequestFields()
 
-	// Handle nil context early to prevent blocking
+	// Handle nil context early. Do not reuse bifrost.ctx: nil-ctx callers would
+	// share mutable request state (request ID, retry counters, billing attempt
+	// time) across concurrent requests.
 	if ctx == nil {
-		ctx = bifrost.ctx
+		ctx = schemas.NewBifrostContext(bifrost.ctx, schemas.NoDeadline)
+		defer ctx.Cancel()
 	}
 
-	// Reset first: bifrost.ctx is shared across every nil-ctx caller.
+	// Reset first in case the caller reuses a BifrostContext.
 	ctx.ResetUpstreamLatency()
+	// Clear any previous attempt stamp before pre-hooks. Without this, a cache
+	// hit or plugin short-circuit could inherit a timestamp from an earlier call.
+	ctx.ClearBillingAttemptStartTime()
 	// Whole-request start for top-down overhead (total minus upstream). On ctx so
 	// tryRequest can stamp the response before post-hooks, where logging reads it.
 	ctx.SetValue(schemas.BifrostContextKeyRequestStartTime, time.Now())
@@ -5343,17 +5380,25 @@ func (bifrost *Bifrost) handleRequest(ctx *schemas.BifrostContext, req *schemas.
 // It handles plugin hooks, request validation, response processing, and fallback providers.
 // If the primary provider fails, it will try each fallback provider in order until one succeeds.
 // It is the wrapper for all streaming public API methods.
-func (bifrost *Bifrost) handleStreamRequest(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) {
+func (bifrost *Bifrost) handleStreamRequest(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (stream chan *schemas.BifrostStreamChunk, bifrostErr *schemas.BifrostError) {
 	defer bifrost.releaseBifrostRequest(req)
 	provider, model, fallbacks := req.GetRequestFields()
 
-	// Handle nil context early to prevent blocking
+	// Streaming callers must own the context lifecycle. Unlike unary requests,
+	// the returned stream can outlive this method, so Bifrost cannot create an
+	// internal context on the caller's behalf: there would be no way for an
+	// abandoned consumer to cancel the provider stream and release its workers.
 	if ctx == nil {
-		ctx = bifrost.ctx
+		bifrostErr := newBifrostErrorFromMsg("context is required for streaming requests")
+		bifrostErr.PopulateExtraFields(req.RequestType, provider, model, model)
+		return nil, bifrostErr
 	}
 
 	ctx.ResetUpstreamLatency()
 	ctx.ResetStreamOverhead()
+	// Clear any previous attempt stamp before pre-hooks. Without this, a cache
+	// hit or plugin short-circuit could inherit a timestamp from an earlier call.
+	ctx.ClearBillingAttemptStartTime()
 	// Whole-request start for overhead on the streaming short-circuit path; the
 	// normal path derives its total from the final chunk.
 	ctx.SetValue(schemas.BifrostContextKeyRequestStartTime, time.Now())
@@ -5701,8 +5746,10 @@ func (bifrost *Bifrost) tryRequest(ctx *schemas.BifrostContext, req *schemas.Bif
 		if ph := bifrost.startCoreSpan(msg.Context, "pipeline-post"); ph.h != nil {
 			defer bifrost.endCoreSpan(ph)
 		}
+		populateBillingAttemptExtraFields(msg.Context, bifrostErrPtr)
 		resp, bifrostErrPtr = pipeline.RunPostLLMHooks(msg.Context, nil, bifrostErrPtr, pluginCount)
 		if bifrostErrPtr != nil {
+			populateBillingAttemptExtraFields(msg.Context, bifrostErrPtr)
 			bifrostErrPtr.PopulateExtraFields(req.RequestType, provider, model, model)
 		} else if resp != nil {
 			resp.PopulateExtraFields(req.RequestType, provider, model, model)
@@ -6022,8 +6069,10 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
 		// Marking final chunk
 		ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true)
 		// On error we will complete post-hooks
+		populateBillingAttemptExtraFields(ctx, &bifrostErrVal)
 		recoveredResp, recoveredErr := pipeline.RunPostLLMHooks(ctx, nil, &bifrostErrVal, len(*bifrost.llmPlugins.Load()))
 		if recoveredErr != nil {
+			populateBillingAttemptExtraFields(ctx, recoveredErr)
 			recoveredErr.PopulateExtraFields(req.RequestType, provider, model, model)
 		} else if recoveredResp != nil {
 			recoveredResp.PopulateExtraFields(req.RequestType, provider, model, model)
@@ -6307,6 +6356,11 @@ func executeRequestWithRetries[T any](
 		}
 
 		logger.Debug("attempting %s request for provider %s", requestType, providerKey)
+		// Record the provider attempt's start before dispatch. Retry and fallback
+		// attempts overwrite this value, so time-based pricing always uses the
+		// attempt that actually produced the billed usage, never completion time.
+		attemptStartedAt := time.Now()
+		ctx.SetBillingAttemptStartTime(attemptStartedAt)
 
 		// Start span for LLM call (or retry attempt)
 		tracer, ok := ctx.Value(schemas.BifrostContextKeyTracer).(schemas.Tracer)
@@ -7111,6 +7165,10 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 				// a concurrent request can then reuse it and overwrite RequestType.
 				// Reading req.RequestType inside the closure would observe the new request's type.
 				attemptRequestType := req.RequestType
+				// Snapshot the attempt start for the same reason: this post-hook runs
+				// asynchronously, and a retry/fallback may already have overwritten the
+				// shared context with a later attempt's timestamp.
+				attemptBillingStartedAt, _ := req.Context.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time)
 				pipeline := bifrost.getPluginPipeline()
 				postHookRunner := func(ctx *schemas.BifrostContext, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError) {
 					// Populate extra fields before RunPostLLMHooks so plugins (e.g. logging)
@@ -7119,23 +7177,39 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 					// reference would let a later retry's alias bleed into this attempt's chunks.
 					if result != nil {
 						result.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel)
+						if !attemptBillingStartedAt.IsZero() {
+							startedAt := attemptBillingStartedAt
+							result.GetExtraFields().BillingAttemptStartedAt = &startedAt
+						}
 						result.PopulateRoutingInfo(perAttemptRoutingInfo)
 					}
 					if err != nil {
 						err.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel)
+						if !attemptBillingStartedAt.IsZero() {
+							startedAt := attemptBillingStartedAt
+							err.ExtraFields.BillingAttemptStartedAt = &startedAt
+						}
 						err.PopulateRoutingInfo(perAttemptRoutingInfo)
 					}
 					resp, bifrostErr := pipeline.RunPostLLMHooks(ctx, result, err, len(*bifrost.llmPlugins.Load()))
 					if IsFinalChunk(ctx) {
 						drainAndAttachPluginLogs(ctx)
 					}
 					if bifrostErr != nil {
+						if !attemptBillingStartedAt.IsZero() {
+							startedAt := attemptBillingStartedAt
+							bifrostErr.ExtraFields.BillingAttemptStartedAt = &startedAt
+						}
 						bifrostErr.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel)
 						bifrostErr.PopulateRoutingInfo(perAttemptRoutingInfo)
 						return nil, bifrostErr
 					} else if resp != nil {
 						resp.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel)
 						resp.PopulateRoutingInfo(perAttemptRoutingInfo)
+						if !attemptBillingStartedAt.IsZero() {
+							startedAt := attemptBillingStartedAt
+							resp.GetExtraFields().BillingAttemptStartedAt = &startedAt
+						}
 					}
 					return resp, nil
 				}
@@ -7201,6 +7275,10 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 		if bifrostError != nil {
 			bifrostError.PopulateExtraFields(req.RequestType, provider.GetProviderKey(), originalModelRequested, resolvedModel)
 			bifrostError.PopulateRoutingInfo(attemptRoutingInfo)
+			if start, ok := req.Context.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time); ok {
+				startCopy := start
+				bifrostError.ExtraFields.BillingAttemptStartedAt = &startCopy
+			}
 
 			// Send error with context awareness to prevent deadlock
 			deliveryTimer.Reset(5 * time.Second)
@@ -7223,6 +7301,10 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 			if result != nil {
 				result.PopulateExtraFields(req.RequestType, provider.GetProviderKey(), originalModelRequested, resolvedModel)
 				result.PopulateRoutingInfo(attemptRoutingInfo)
+				if start, ok := req.Context.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time); ok {
+					startCopy := start
+					result.GetExtraFields().BillingAttemptStartedAt = &startCopy
+				}
 			}
 			if IsStreamRequestType(req.RequestType) {
 				// Send stream with context awareness to prevent deadlock
@@ -7921,6 +8003,12 @@ func (p *PluginPipeline) RunPostLLMHooks(ctx *schemas.BifrostContext, resp *sche
 			// Restore the parent so the next plugin is a sibling, not chained under this one.
 			ctx.SetValue(schemas.BifrostContextKeySpanID, prevSpanID)
 		}
+		// A post-hook may replace either object without copying ExtraFields. Restore
+		// the authoritative attempt timestamp before the next (outer) post-hook runs,
+		// so logging and pricing plugins never observe an unstamped replacement.
+		populateBillingAttemptResponseExtraFields(ctx, resp)
+		populateBillingAttemptExtraFields(ctx, bifrostErr)
+
 		// If a plugin recovers from an error (sets bifrostErr to nil and sets resp), allow that
 		// If a plugin invalidates a response (sets resp to nil and sets bifrostErr), allow that
 	}

core/bifrost_test.go

@@ -3253,3 +3253,40 @@ func TestExecuteRequestWithRetries_EmptyStreamReturnsClosedChannel(t *testing.T)
 		t.Errorf("Expected range over empty stream to yield 0 chunks, got %d", count)
 	}
 }
+
+// TestHandleStreamRequest_RequiresContext pins the streaming lifecycle contract:
+// the caller must own the context so it can cancel an abandoned stream.
+func TestHandleStreamRequest_RequiresContext(t *testing.T) {
+	bifrost := &Bifrost{
+		bifrostRequestPool: sync.Pool{New: func() interface{} { return &schemas.BifrostRequest{} }},
+		logger:             NewNoOpLogger(),
+	}
+	req := &schemas.BifrostRequest{
+		RequestType: schemas.ChatCompletionStreamRequest,
+		ChatRequest: &schemas.BifrostChatRequest{
+			Provider: schemas.OpenAI,
+			Model:    "test-model",
+			Input:    []schemas.ChatMessage{{}},
+		},
+	}
+
+	stream, bifrostErr := bifrost.handleStreamRequest(nil, req)
+	if stream != nil {
+		t.Fatal("expected no stream for a nil streaming context")
+	}
+	if bifrostErr == nil || bifrostErr.Error == nil {
+		t.Fatalf("expected a detailed error, got %#v", bifrostErr)
+	}
+	if bifrostErr.Error.Message != "context is required for streaming requests" {
+		t.Fatalf("unexpected error message: %q", bifrostErr.Error.Message)
+	}
+	if bifrostErr.ExtraFields.RequestType != schemas.ChatCompletionStreamRequest {
+		t.Fatalf("unexpected request type: %v", bifrostErr.ExtraFields.RequestType)
+	}
+	if bifrostErr.ExtraFields.Provider != schemas.OpenAI {
+		t.Fatalf("unexpected provider: %v", bifrostErr.ExtraFields.Provider)
+	}
+	if bifrostErr.ExtraFields.OriginalModelRequested != "test-model" || bifrostErr.ExtraFields.ResolvedModelUsed != "test-model" {
+		t.Fatalf("unexpected model metadata: %#v", bifrostErr.ExtraFields)
+	}
+}