diff --git a/core/bifrost.go b/core/bifrost.go index 4efb83ab868..6a5ebd4ab60 100644 --- a/core/bifrost.go +++ b/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)) } + 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 populateBillingAttemptStartTime(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 + } +} + +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,11 +6069,14 @@ 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) + populateBillingAttemptStartTime(ctx, recoveredResp) } drainAndAttachPluginLogs(ctx) bifrost.releaseChannelMessage(msg) @@ -6307,6 +6357,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 +7166,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,10 +7178,18 @@ 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())) @@ -7130,12 +7197,20 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas 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 +7276,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 +7302,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 diff --git a/core/bifrost_test.go b/core/bifrost_test.go index 944e236954c..40c9ecf27c9 100644 --- a/core/bifrost_test.go +++ b/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) + } +} diff --git a/core/billing_attempt_time_test.go b/core/billing_attempt_time_test.go new file mode 100644 index 00000000000..faefe59b348 --- /dev/null +++ b/core/billing_attempt_time_test.go @@ -0,0 +1,70 @@ +package bifrost + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/maximhq/bifrost/core/schemas" +) + +// TestClearBillingAttemptStartTimeMasksInheritedValue ensures a child context +// cannot expose a stale attempt timestamp inherited from its parent. +func TestClearBillingAttemptStartTimeMasksInheritedValue(t *testing.T) { + parent := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + stamp := time.Date(2026, 8, 27, 9, 0, 0, 0, time.UTC) + parent.SetBillingAttemptStartTime(stamp) + + child := schemas.NewBifrostContext(parent, schemas.NoDeadline) + if got, ok := child.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time); !ok || !got.Equal(stamp) { + t.Fatalf("child should inherit the parent timestamp, got %#v", got) + } + + child.ClearBillingAttemptStartTime() + if _, ok := child.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time); ok { + t.Fatal("cleared child context exposed an inherited billing timestamp") + } +} + +// TestNilUnaryRequestsUseIsolatedContexts pins that nil-context callers no longer +// share the mutable Bifrost-wide context used for request-scoped pricing state. +func TestNilUnaryRequestsUseIsolatedContexts(t *testing.T) { + newBifrost := func() *Bifrost { + return &Bifrost{ + ctx: schemas.NewBifrostContext(context.Background(), schemas.NoDeadline), + logger: NewNoOpLogger(), + account: NewMockAccount(), + providerMutexes: sync.Map{}, + } + } + req := &schemas.BifrostChatRequest{ + Provider: schemas.OpenAI, + Model: "test-model", + Input: []schemas.ChatMessage{{}}, + } + bifrost := newBifrost() + bifrost.bifrostRequestPool = sync.Pool{New: func() interface{} { return &schemas.BifrostRequest{} }} + bifrost.llmPlugins.Store(&[]schemas.LLMPlugin{}) + bifrost.mcpPlugins.Store(&[]schemas.MCPPlugin{}) + bifrost.pluginPipelinePool = sync.Pool{New: func() interface{} { + return &PluginPipeline{ + preHookErrors: make([]error, 0), + postHookErrors: make([]error, 0), + } + }} + bifrost.tracer.Store(&tracerWrapper{tracer: &schemas.NoOpTracer{}}) + _, err := bifrost.ChatCompletionRequest(nil, req) + if err == nil || err.Error == nil { + t.Fatalf("expected request to fail without providers, got %#v", err) + } + if bifrost.ctx.Value(schemas.BifrostContextKeyRequestStartTime) != nil { + t.Fatal("nil unary caller wrote a request start time to the shared Bifrost context") + } + if _, ok := bifrost.ctx.Value(schemas.BifrostContextKeyBillingAttemptStartTime).(time.Time); ok { + t.Fatal("nil unary caller left a billing timestamp on the shared Bifrost context") + } + if v, ok := bifrost.ctx.Value(schemas.BifrostContextKeyRequestID).(string); ok && v != "" { + t.Fatalf("nil unary caller changed the shared Bifrost request ID: %q", v) + } +} diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index d559c9d6617..a4c64c90d5c 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "strconv" + "time" ) const ( @@ -302,6 +303,7 @@ const ( BifrostContextKeyParentSpanID BifrostContextKey = "bifrost-parent-span-id" // string (parent span ID from W3C traceparent header - set by tracing middleware) BifrostContextKeyStreamStartTime BifrostContextKey = "bifrost-stream-start-time" // time.Time (start time for streaming TTFT calculation - set by bifrost) BifrostContextKeyRequestStartTime BifrostContextKey = "bifrost-request-start-time" // time.Time (whole-request start for overhead - set by bifrost) + BifrostContextKeyBillingAttemptStartTime BifrostContextKey = "bifrost-billing-attempt-start-time" // time.Time (start of the current provider attempt for time-based pricing - set by bifrost) BifrostContextKeyTracer BifrostContextKey = "bifrost-tracer" // Tracer (tracer instance for completing deferred spans - set by bifrost) BifrostContextKeyModelCatalog BifrostContextKey = "bifrost-model-catalog" // ModelInfoProvider (model pricing/capability catalog backing ctx.GetModelInfo and ctx.CalculateCost - set by bifrost) BifrostContextKeyDeferTraceCompletion BifrostContextKey = "bifrost-defer-trace-completion" // bool (signals trace completion should be deferred for streaming - set by streaming handlers) @@ -1743,7 +1745,12 @@ type BifrostResponseExtraFields struct { // serializing this response is itself overhead. The authoritative value is // stamped on the trace and logged at completion; this is only the untraced // fallback. Nil means unknown. - OverheadLatency *int64 `json:"-"` + OverheadLatency *int64 `json:"-"` + // BillingAttemptStartedAt is when the current provider attempt started. It is + // the authoritative instant for time-based pricing schedules and is deliberately + // distinct from response creation/completion time. Nil means unknown; consumers + // must not substitute CreatedAt or time.Now(). + BillingAttemptStartedAt *time.Time `json:"billing_attempt_started_at,omitempty"` ChunkIndex int `json:"chunk_index"` // used for streaming responses to identify the chunk index, will be 0 for non-streaming responses RawRequest interface{} `json:"raw_request,omitempty"` RawResponse interface{} `json:"raw_response,omitempty"` @@ -2061,4 +2068,8 @@ type BifrostErrorExtraFields struct { // the provider actually billed us for. Nil when the failure consumed no // tokens (e.g. 401/403/429 before the model ran). BilledUsage *BifrostLLMUsage `json:"billed_usage,omitempty"` + // BillingAttemptStartedAt is when the failed/cancelled provider attempt started. + // It prices BilledUsage under time-based schedules and must not be inferred from + // response creation/completion timestamps. + BillingAttemptStartedAt *time.Time `json:"billing_attempt_started_at,omitempty"` } diff --git a/core/schemas/context.go b/core/schemas/context.go index 13741055b67..cafba0b3efc 100644 --- a/core/schemas/context.go +++ b/core/schemas/context.go @@ -39,6 +39,7 @@ var reservedKeys = map[BifrostContextKey]struct{}{ BifrostContextKeyMCPHealthCheckRequest: {}, BifrostContextKeyUpstreamLatency: {}, BifrostContextKeyStreamOverhead: {}, + BifrostContextKeyBillingAttemptStartTime: {}, BifrostContextKeyRoutingInfo: {}, BifrostContextKeyMCPInboundBearer: {}, } @@ -426,6 +427,33 @@ func (bc *BifrostContext) setReservedValue(key, value any) { bc.userValues[key] = value } +// SetBillingAttemptStartTime records the provider-attempt start time. This +// internal setter bypasses the restricted-write guard because streaming +// post-hooks may hold blockRestrictedWrites while the orchestrator starts or +// retries an attempt. +func (bc *BifrostContext) SetBillingAttemptStartTime(t time.Time) { + bc.setReservedValue(BifrostContextKeyBillingAttemptStartTime, t) +} + +// ClearBillingAttemptStartTime removes a previous provider-attempt stamp. It is +// called at request entry so a short-circuit cannot inherit an attempt from an +// earlier call on a reused context. +func (bc *BifrostContext) ClearBillingAttemptStartTime() { + if bc.valueDelegate != nil { + bc.valueDelegate.ClearBillingAttemptStartTime() + return + } + bc.valuesMu.Lock() + defer bc.valuesMu.Unlock() + if bc.userValues == nil { + bc.userValues = make(map[any]any) + } + // Store an explicit nil sentinel instead of deleting the key: Value falls + // back to the parent when a key is absent locally, so this must also mask + // any timestamp inherited from a parent context. + bc.userValues[BifrostContextKeyBillingAttemptStartTime] = nil +} + // SetRoutingInfoSnapshot writes the routed-identity RoutingInfo snapshot, // bypassing the restricted-writes guard. The orchestrator needs this because a // streaming response's async per-chunk post-hooks hold blockRestrictedWrites diff --git a/core/schemas/guardraildebug.go b/core/schemas/guardraildebug.go index c5c90db6b37..d8514deba1d 100644 --- a/core/schemas/guardraildebug.go +++ b/core/schemas/guardraildebug.go @@ -1,5 +1,7 @@ package schemas +import "time" + // BifrostGuardrailDebug carries request-scoped guardrail execution metadata. type BifrostGuardrailDebug struct { JudgeCalls []BifrostGuardrailJudgeCall `json:"judge_calls,omitempty"` @@ -22,6 +24,19 @@ type BifrostGuardrailJudgeCall struct { CompletionTokens int `json:"completion_tokens,omitempty"` CompletionTokensDetails *ChatCompletionTokensDetails `json:"completion_tokens_details,omitempty"` TotalTokens int `json:"total_tokens,omitempty"` + // StartedAt is the start of this internal judge invocation. It must be priced + // independently from the parent request: output guardrails can run after the main + // attempt completes, and a schedule boundary may fall between the two calls. + StartedAt *time.Time `json:"started_at,omitempty"` +} + +// cloneTime returns an owned copy of a timestamp. +func cloneTime(value *time.Time) *time.Time { + if value == nil { + return nil + } + copy := *value + return © } // Clone returns an owned snapshot of the guardrail debug data. @@ -37,6 +52,7 @@ func (d *BifrostGuardrailDebug) Clone() *BifrostGuardrailDebug { clone.JudgeCalls[index].RuleID = cloneUint(call.RuleID) clone.JudgeCalls[index].PromptTokensDetails = cloneChatPromptTokensDetails(call.PromptTokensDetails) clone.JudgeCalls[index].CompletionTokensDetails = cloneChatCompletionTokensDetails(call.CompletionTokensDetails) + clone.JudgeCalls[index].StartedAt = cloneTime(call.StartedAt) } return clone } diff --git a/core/schemas/guardraildebug_test.go b/core/schemas/guardraildebug_test.go index 4c2037d6133..38deb2b1153 100644 --- a/core/schemas/guardraildebug_test.go +++ b/core/schemas/guardraildebug_test.go @@ -2,6 +2,7 @@ package schemas import ( "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -80,3 +81,27 @@ func TestAppendGuardrailJudgeCallRejectsEmptyUsage(t *testing.T) { _, ok := GuardrailDebugFromContext(ctx) assert.False(t, ok) } + +// TestGuardrailDebugContextClonesStartedAt verifies an internal judge call owns its +// timestamp. Output guardrails may execute after the parent attempt crosses a pricing +// schedule boundary, so the judge must retain its own invocation time. +func TestGuardrailDebugContextClonesStartedAt(t *testing.T) { + ctx := NewBifrostContext(nil, NoDeadline) + startedAt := time.Date(2026, 8, 21, 16, 0, 0, 0, time.UTC) + require.True(t, AppendGuardrailJudgeCallOnContext(ctx, BifrostGuardrailJudgeCall{ + JudgeProvider: OpenAI, + JudgeModel: "gpt-4o-mini", + TotalTokens: 10, + StartedAt: &startedAt, + })) + + first, ok := GuardrailDebugFromContext(ctx) + require.True(t, ok) + require.NotNil(t, first.JudgeCalls[0].StartedAt) + *first.JudgeCalls[0].StartedAt = startedAt.Add(time.Hour) + + second, ok := GuardrailDebugFromContext(ctx) + require.True(t, ok) + require.NotNil(t, second.JudgeCalls[0].StartedAt) + assert.True(t, second.JudgeCalls[0].StartedAt.Equal(startedAt)) +} diff --git a/core/schemas/tracer.go b/core/schemas/tracer.go index b581bcd65dd..63a4fc56379 100644 --- a/core/schemas/tracer.go +++ b/core/schemas/tracer.go @@ -13,28 +13,29 @@ type SpanHandle interface{} // StreamAccumulatorResult contains the accumulated data from streaming chunks. // This is the return type for tracer's streaming accumulation methods. type StreamAccumulatorResult struct { - RequestID string // Request ID - RequestedModel string // Original model requested by the caller - ResolvedModel string // Actual model used by the provider (equals RequestedModel when no alias mapping exists) - Provider ModelProvider // Provider used - Status string // Status of the stream - Latency int64 // Latency in milliseconds - TimeToFirstToken int64 // Time to first token in milliseconds - OutputMessage *ChatMessage // Accumulated output message - OutputMessages []ResponsesMessage // For responses API - TokenUsage *BifrostLLMUsage // Token usage - ServiceTier *BifrostServiceTier // Served tier (for example "priority", "flex", "ultrafast", or "default"); needs its own field because it lives on the response envelope, not on BifrostLLMUsage like Speed and InferenceGeo - Cost *float64 // Cost in dollars - CacheDebug *BifrostCacheDebug // Semantic cache debug info if available - GuardrailDebug *BifrostGuardrailDebug // Guardrail debug info if available - ErrorDetails *BifrostError // Error details if any - AudioOutput *BifrostSpeechResponse // For speech streaming - TranscriptionOutput *BifrostTranscriptionResponse // For transcription streaming - ImageGenerationOutput *BifrostImageGenerationResponse // For image generation streaming - PassthroughOutput *BifrostPassthroughResponse // For passthrough streaming - FinishReason *string // Finish reason - RawResponse *string // Raw response - RawRequest interface{} // Raw request + RequestID string // Request ID + RequestedModel string // Original model requested by the caller + ResolvedModel string // Actual model used by the provider (equals RequestedModel when no alias mapping exists) + Provider ModelProvider // Provider used + Status string // Status of the stream + Latency int64 // Latency in milliseconds + TimeToFirstToken int64 // Time to first token in milliseconds + OutputMessage *ChatMessage // Accumulated output message + OutputMessages []ResponsesMessage // For responses API + TokenUsage *BifrostLLMUsage // Token usage + BillingAttemptStartedAt *time.Time // Start of the provider attempt that produced this usage; nil means unknown + ServiceTier *BifrostServiceTier // Served tier (for example "priority", "flex", "ultrafast", or "default"); needs its own field because it lives on the response envelope, not on BifrostLLMUsage like Speed and InferenceGeo + Cost *float64 // Cost in dollars + CacheDebug *BifrostCacheDebug // Semantic cache debug info if available + GuardrailDebug *BifrostGuardrailDebug // Guardrail debug info if available + ErrorDetails *BifrostError // Error details if any + AudioOutput *BifrostSpeechResponse // For speech streaming + TranscriptionOutput *BifrostTranscriptionResponse // For transcription streaming + ImageGenerationOutput *BifrostImageGenerationResponse // For image generation streaming + PassthroughOutput *BifrostPassthroughResponse // For passthrough streaming + FinishReason *string // Finish reason + RawResponse *string // Raw response + RawRequest interface{} // Raw request } // Tracer defines the interface for distributed tracing in Bifrost. diff --git a/core/utils.go b/core/utils.go index 7ef655af723..feb7688dedc 100644 --- a/core/utils.go +++ b/core/utils.go @@ -352,6 +352,9 @@ func newBifrostMessageChan(message *schemas.BifrostResponse) chan *schemas.Bifro // clearCtxForFallback clears the ctx values which are not applicable for fallback requests. func clearCtxForFallback(ctx *schemas.BifrostContext) { + // A new provider attempt has not started yet. Clear any primary-attempt + // billing stamp before fallback pre-hooks or setup can short-circuit. + ctx.ClearBillingAttemptStartTime() ctx.ClearValue(schemas.BifrostContextKeyAPIKeyID) ctx.ClearValue(schemas.BifrostContextKeyAPIKeyName) ctx.ClearValue(schemas.BifrostContextKeyDirectKey) diff --git a/framework/logstore/migrations.go b/framework/logstore/migrations.go index 1d90d498a26..bf8980bea51 100644 --- a/framework/logstore/migrations.go +++ b/framework/logstore/migrations.go @@ -294,6 +294,7 @@ var logstoreMigrationSteps = []migrationStep{ {IDs: []string{"logs_add_cost_breakdown_columns"}, run: migrationAddCostBreakdownColumns}, {IDs: []string{"logs_recreate_matviews_with_cost_breakdown"}, run: migrationRecreateMatViewsWithCostBreakdown}, {IDs: []string{"logs_add_overhead_breakdown_column"}, run: migrationAddOverheadBreakdownColumn}, + {IDs: []string{"logs_add_billing_attempt_start_time_column"}, run: migrationAddBillingAttemptStartTimeColumn}, } // areThereAnyPendingMigrations returns true if there are any pending migrations to be applied. @@ -697,6 +698,33 @@ func migrationAddOverheadBreakdownColumn(ctx context.Context, db *gorm.DB, logge return nil } +// migrationAddBillingAttemptStartTimeColumn adds the attempt-start timestamp used by +// time-based pricing. It is separate from log Timestamp/CreatedAt: both describe log +// creation after the provider attempt completed, while schedule evaluation must remain +// deterministic for retries, fallbacks, and historical recomputation. +func migrationAddBillingAttemptStartTimeColumn(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { + migrationName := "logs_add_billing_attempt_start_time_column" + logger.Info("[logstore] starting migration %s", migrationName) + defer logger.Info("[logstore] finished migration %s", migrationName) + opts := *migrator.DefaultOptions + opts.UseTransaction = true + m := migrator.New(db, &opts, []*migrator.Migration{{ + ID: migrationName, + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + return addColumnIfNotExists(tx, logger, &Log{}, "billing_attempt_started_at") + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + return dropColumnIfExists(tx, logger, &Log{}, "billing_attempt_started_at") + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error while adding billing_attempt_started_at column: %s", err.Error()) + } + return nil +} + // migrationAddResponsesInputHistoryColumn adds the responses_input_history column to the logs table. func migrationAddResponsesInputHistoryColumn(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { migrationName := "logs_init_add_responses_input_history_column" @@ -2584,6 +2612,11 @@ var performanceIndexes = []performanceIndexDef{ name: "idx_logs_latency", sql: "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_logs_latency ON logs(latency)", }, + { + table: "logs", + name: "idx_logs_billing_attempt_started_at", + sql: "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_logs_billing_attempt_started_at ON logs(billing_attempt_started_at) WHERE billing_attempt_started_at IS NOT NULL", + }, { table: "logs", name: "idx_logs_total_tokens", diff --git a/framework/logstore/rdb.go b/framework/logstore/rdb.go index 829caf5389a..2c0373401b1 100644 --- a/framework/logstore/rdb.go +++ b/framework/logstore/rdb.go @@ -1264,7 +1264,7 @@ func (s *RDBLogStore) listSelectColumns() string { "prompt_tokens", "completion_tokens", "total_tokens", "cached_read_tokens", "has_object", "content_hidden", - "service_tier", "speed", "inference_geo", + "service_tier", "speed", "inference_geo", "billing_attempt_started_at", "created_at", }, ", ") @@ -1363,8 +1363,8 @@ var billingScalarColumns = []string{ "batch_debug", // Whether the payload was offloaded, and whether it can ever be fetched back. "has_object", "content_hidden", - // Served tier: scales every token rate. - "service_tier", "speed", "inference_geo", + // Served tier and schedule time: scale every token rate deterministically. + "service_tier", "speed", "inference_geo", "billing_attempt_started_at", } // billingSelectColumns returns the SELECT clause for cost recomputation. diff --git a/framework/logstore/tables.go b/framework/logstore/tables.go index 702623cdd57..bb3e27acb4a 100644 --- a/framework/logstore/tables.go +++ b/framework/logstore/tables.go @@ -332,6 +332,10 @@ type Log struct { ServiceTier *string `gorm:"type:varchar(32)" json:"service_tier,omitempty"` // OpenAI served tier, e.g. "priority", "flex", "ultrafast", or "default" Speed *string `gorm:"type:varchar(32)" json:"speed,omitempty"` // Anthropic served speed: "fast" / "standard" InferenceGeo *string `gorm:"type:varchar(32)" json:"inference_geo,omitempty"` // Anthropic data residency, e.g. "us" + // BillingAttemptStartedAt is when the provider attempt started. It is separate + // from Timestamp (log creation/completion) because time-based pricing must be + // stable across live billing and historical recomputation. + BillingAttemptStartedAt *time.Time `gorm:"index" json:"billing_attempt_started_at,omitempty"` CreatedAt time.Time `gorm:"index;not null" json:"created_at"` diff --git a/framework/modelcatalog/datasheet/cost.go b/framework/modelcatalog/datasheet/cost.go index 46064167478..f46d954dfa4 100644 --- a/framework/modelcatalog/datasheet/cost.go +++ b/framework/modelcatalog/datasheet/cost.go @@ -5,11 +5,19 @@ import ( "fmt" "strconv" "strings" + "time" "github.com/maximhq/bifrost/core/schemas" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" ) +// CostCalculationOptions carries inputs that cannot be inferred from a bare +// usage object. BillingAttemptStartedAt is authoritative: a nil value means +// unknown and intentionally falls back to base pricing. +type CostCalculationOptions struct { + BillingAttemptStartedAt *time.Time +} + // CalculateCost calculates the cost of a Bifrost response. // It handles all request types, cache and guardrail billing, and tiered pricing. // If scopes is nil, an empty LookupScopes is used; global and provider-scoped @@ -37,14 +45,15 @@ func (s *Store) CalculateCostBreakdown(result *schemas.BifrostResponse, scopes * } extraFields := result.GetExtraFields() + options := CostCalculationOptions{BillingAttemptStartedAt: extraFields.BillingAttemptStartedAt} // Handle semantic cache billing cacheDebug := extraFields.CacheDebug var requestCost *schemas.BifrostCost if cacheDebug != nil { - requestCost = s.calculateCostWithCache(result, cacheDebug, lookupScopes) + requestCost = s.calculateCostWithCache(result, cacheDebug, lookupScopes, options) } else { - requestCost = s.calculateBaseCost(result, lookupScopes) + requestCost = s.calculateBaseCost(result, lookupScopes, options) } // Handle guardrail judge-call billing @@ -85,7 +94,14 @@ func (s *Store) CalculateCostBreakdown(result *schemas.BifrostResponse, scopes * // wrapper over CalculateCostBreakdownForUsage returning only the total, so both // paths compute cost identically. func (s *Store) CalculateCostForUsage(usage *schemas.BifrostLLMUsage, provider schemas.ModelProvider, model string, requestType schemas.RequestType, scopes *LookupScopes) float64 { - breakdown := s.CalculateCostBreakdownForUsage(usage, provider, model, requestType, scopes) + return s.CalculateCostForUsageWithOptions(usage, provider, model, requestType, scopes, nil) +} + +// CalculateCostForUsageWithOptions is CalculateCostForUsage with an explicit +// billing instant. The option must come from the provider attempt that produced +// usage; callers must never synthesize it from completion or log timestamps. +func (s *Store) CalculateCostForUsageWithOptions(usage *schemas.BifrostLLMUsage, provider schemas.ModelProvider, model string, requestType schemas.RequestType, scopes *LookupScopes, options *CostCalculationOptions) float64 { + breakdown := s.CalculateCostBreakdownForUsageWithOptions(usage, provider, model, requestType, scopes, options) if breakdown == nil { return 0 } @@ -98,6 +114,12 @@ func (s *Store) CalculateCostForUsage(usage *schemas.BifrostLLMUsage, provider s // output / additional split, not just the scalar total. Returns nil when there // is no cost to record. func (s *Store) CalculateCostBreakdownForUsage(usage *schemas.BifrostLLMUsage, provider schemas.ModelProvider, model string, requestType schemas.RequestType, scopes *LookupScopes) *schemas.BifrostCost { + return s.CalculateCostBreakdownForUsageWithOptions(usage, provider, model, requestType, scopes, nil) +} + +// CalculateCostBreakdownForUsageWithOptions is the explicit-billing-time variant +// of CalculateCostBreakdownForUsage. +func (s *Store) CalculateCostBreakdownForUsageWithOptions(usage *schemas.BifrostLLMUsage, provider schemas.ModelProvider, model string, requestType schemas.RequestType, scopes *LookupScopes, options *CostCalculationOptions) *schemas.BifrostCost { if usage == nil { return nil } @@ -117,6 +139,10 @@ func (s *Store) CalculateCostBreakdownForUsage(usage *schemas.BifrostLLMUsage, p input := costInput{usage: usage} input.tier = tierFromResponse(nil, usage.Speed, usage.InferenceGeo) + var calculationOptions CostCalculationOptions + if options != nil { + calculationOptions = *options + } return s.computeCostFromInput( input, schemas.RoutingInfo{ @@ -126,6 +152,7 @@ func (s *Store) CalculateCostBreakdownForUsage(usage *schemas.BifrostLLMUsage, p }, normalizeStreamRequestType(requestType), lookupScopes, + calculationOptions, ) } @@ -171,12 +198,13 @@ func (s *Store) computeGuardrailJudgeCost(call schemas.BifrostGuardrailJudgeCall if requestType == "" { requestType = schemas.ChatCompletionRequest } - return s.CalculateCostForUsage( + return s.CalculateCostForUsageWithOptions( usage, call.JudgeProvider, call.JudgeModel, requestType, &judgeScopes, + &CostCalculationOptions{BillingAttemptStartedAt: call.StartedAt}, ) } @@ -289,7 +317,7 @@ func cloneFloat64Pointer(value *float64) *float64 { } // calculateCostWithCache handles cost calculation when semantic cache debug info is present. -func (s *Store) calculateCostWithCache(result *schemas.BifrostResponse, cacheDebug *schemas.BifrostCacheDebug, scopes LookupScopes) *schemas.BifrostCost { +func (s *Store) calculateCostWithCache(result *schemas.BifrostResponse, cacheDebug *schemas.BifrostCacheDebug, scopes LookupScopes, options CostCalculationOptions) *schemas.BifrostCost { if cacheDebug.CacheHit { // Direct cache hit — no LLM call, no cost if cacheDebug.HitType != nil && *cacheDebug.HitType == "direct" { @@ -313,7 +341,7 @@ func (s *Store) calculateCostWithCache(result *schemas.BifrostResponse, cacheDeb } // Cache miss — full LLM cost + embedding lookup cost (a sidecar additional cost) - base := s.calculateBaseCost(result, scopes) + base := s.calculateBaseCost(result, scopes, options) embeddingCost := s.computeCacheEmbeddingCost(cacheDebug, scopes) if embeddingCost == 0 { return base @@ -383,7 +411,7 @@ func computeContainerCreationCost(pricing *configstoreTables.TableModelPricing) } // calculateBaseCost extracts usage from the response and routes to the appropriate compute function. -func (s *Store) calculateBaseCost(result *schemas.BifrostResponse, scopes LookupScopes) *schemas.BifrostCost { +func (s *Store) calculateBaseCost(result *schemas.BifrostResponse, scopes LookupScopes, options CostCalculationOptions) *schemas.BifrostCost { extraFields := result.GetExtraFields() if extraFields == nil { return nil @@ -451,10 +479,10 @@ func (s *Store) calculateBaseCost(result *schemas.BifrostResponse, scopes Lookup // "model-router" deployment name on RoutingInfo.Model. if result.PassthroughResponse == nil && routingInfo.Provider == schemas.Azure && schemas.IsAzureModelRouter(routingInfo.Model) && (requestType == schemas.TextCompletionRequest || requestType == schemas.ChatCompletionRequest || requestType == schemas.ResponsesRequest) { - return s.calculateAzureModelRouterCost(result, input, routingInfo, requestType, scopes) + return s.calculateAzureModelRouterCost(result, input, routingInfo, requestType, scopes, options) } - return s.computeCostFromInput(input, routingInfo, requestType, scopes) + return s.computeCostFromInput(input, routingInfo, requestType, scopes, options) } // calculateAzureModelRouterCost bills the Model Router deployment's own @@ -462,20 +490,20 @@ func (s *Store) calculateBaseCost(result *schemas.BifrostResponse, scopes Lookup // model it actually routed to, looked up fresh under the served model name so // regular per-token/tiered pricing applies to it exactly as if it had been // called directly. -func (s *Store) calculateAzureModelRouterCost(result *schemas.BifrostResponse, input costInput, routingInfo schemas.RoutingInfo, requestType schemas.RequestType, scopes LookupScopes) *schemas.BifrostCost { +func (s *Store) calculateAzureModelRouterCost(result *schemas.BifrostResponse, input costInput, routingInfo schemas.RoutingInfo, requestType schemas.RequestType, scopes LookupScopes, options CostCalculationOptions) *schemas.BifrostCost { pricingRequestType := requestType if pricingRequestType == schemas.TextCompletionRequest { pricingRequestType = schemas.ChatCompletionRequest } - cost := s.computeCostFromInput(input, routingInfo, pricingRequestType, scopes) + cost := s.computeCostFromInput(input, routingInfo, pricingRequestType, scopes, options) if servedModel := azureModelRouterServedModel(result); servedModel != "" && servedModel != routingInfo.Model { underlyingRoutingInfo := schemas.RoutingInfo{ Provider: routingInfo.Provider, Model: servedModel, } - cost = cost.Add(s.computeCostFromInput(input, underlyingRoutingInfo, pricingRequestType, scopes)) + cost = cost.Add(s.computeCostFromInput(input, underlyingRoutingInfo, pricingRequestType, scopes, options)) } return cost @@ -503,7 +531,7 @@ func azureModelRouterServedModel(result *schemas.BifrostResponse) string { // type and routes the extracted usage to the appropriate per-modality compute // function. Shared by calculateBaseCost (response-driven) and // CalculateCostForUsage (bare-usage-driven, for failed/cancelled requests). -func (s *Store) computeCostFromInput(input costInput, routingInfo schemas.RoutingInfo, requestType schemas.RequestType, scopes LookupScopes) *schemas.BifrostCost { +func (s *Store) computeCostFromInput(input costInput, routingInfo schemas.RoutingInfo, requestType schemas.RequestType, scopes LookupScopes, options CostCalculationOptions) *schemas.BifrostCost { // When a pricing model override is set (e.g. container creates always look // up "container"), it replaces the lookup hierarchy entirely. Build a // synthetic RoutingInfo that reuses Provider but pins the model fields to @@ -566,9 +594,117 @@ func (s *Store) computeCostFromInput(input costInput, routingInfo schemas.Routin cost.InputCostDetails.RequestCost += *pricing.CostPerRequest cost.TotalCost += *pricing.CostPerRequest } + + // Time schedules multiply the final resolved rate, so they compose with + // service and context tiers rather than replacing them. They also scale the + // flat request fee because it is part of the model's final price. + cost = s.applyPricingSchedule(routingInfo, cost, options.BillingAttemptStartedAt) return cost } +// applyPricingSchedule resolves the model-level schedule and returns a scaled +// copy of cost. It never mutates the input because it can alias +// provider-supplied or caller-owned breakdowns. +func (s *Store) applyPricingSchedule(routingInfo schemas.RoutingInfo, cost *schemas.BifrostCost, at *time.Time) *schemas.BifrostCost { + if cost == nil { + return nil + } + + var aliasModelID, aliasModelName, serverSideFallbackModel string + if rka := routingInfo.ResolvedKeyAlias; rka != nil { + aliasModelID = rka.ModelID + if rka.ModelName != nil { + aliasModelName = *rka.ModelName + } + } + if routingInfo.ServerSideFallbackModel != nil { + serverSideFallbackModel = *routingInfo.ServerSideFallbackModel + } + + // Use the same candidate precedence as resolvePricing so an aliased request + // follows its canonical pricing row and its schedule together. + s.mu.RLock() + var model string + var schedule *PricingTimeSchedule + for _, candidate := range []string{serverSideFallbackModel, aliasModelName, aliasModelID, routingInfo.Model} { + if candidate == "" { + continue + } + if candidateSchedule, exists := s.pricingSchedules[candidate]; exists { + model, schedule = candidate, candidateSchedule + break + } + } + s.mu.RUnlock() + if schedule == nil { + return cost + } + + details := &schemas.PricingScheduleCostDetails{Multiplier: 1} + if at == nil { + details.TimestampAvailable = false + } else { + details.TimestampAvailable = true + evaluation, err := EvaluatePricingTimeSchedule(schedule, *at) + if err != nil { + if s.logger != nil { + s.logger.Warn("ignoring invalid pricing schedule for %s: %v", model, err) + } + } else { + details.Multiplier = evaluation.Multiplier + details.Matched = evaluation.Matched + } + } + if details.Multiplier == 1 { + cost.PricingSchedule = details + return cost + } + + scaled := scaleBifrostCost(cost, details.Multiplier) + scaled.PricingSchedule = details + return scaled +} + +// scaleBifrostCost deep-copies and multiplies every cost category. +func scaleBifrostCost(cost *schemas.BifrostCost, multiplier float64) *schemas.BifrostCost { + if cost == nil { + return nil + } + scaled := *cost + scaled.InputCost *= multiplier + scaled.OutputCost *= multiplier + scaled.AdditionalCost *= multiplier + scaled.TotalCost *= multiplier + if cost.InputCostDetails != nil { + details := *cost.InputCostDetails + details.TextCost *= multiplier + details.AudioCost *= multiplier + details.ImageCost *= multiplier + details.CachedReadCost *= multiplier + details.CachedWriteCost *= multiplier + details.RequestCost *= multiplier + scaled.InputCostDetails = &details + } + if cost.OutputCostDetails != nil { + details := *cost.OutputCostDetails + details.TextCost *= multiplier + details.AudioCost *= multiplier + details.ImageCost *= multiplier + details.ReasoningCost *= multiplier + details.CitationCost *= multiplier + details.SearchQueriesCost *= multiplier + scaled.OutputCostDetails = &details + } + if cost.AdditionalCostDetails != nil { + details := *cost.AdditionalCostDetails + details.GuardrailCost *= multiplier + details.MCPCost *= multiplier + details.SemanticCacheCost *= multiplier + scaled.AdditionalCostDetails = &details + } + return &scaled +} + // --------------------------------------------------------------------------- // Usage extraction // --------------------------------------------------------------------------- diff --git a/framework/modelcatalog/datasheet/cost_test.go b/framework/modelcatalog/datasheet/cost_test.go index ae320ac3ae7..1a367a08331 100644 --- a/framework/modelcatalog/datasheet/cost_test.go +++ b/framework/modelcatalog/datasheet/cost_test.go @@ -3,6 +3,7 @@ package datasheet import ( "encoding/json" "testing" + "time" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" @@ -5213,3 +5214,199 @@ func TestParseImageDimensions(t *testing.T) { assert.Equal(t, 0, h, bad) } } + +// pricingScheduleForTest returns a simple deterministic schedule: 0.5x between +// 16:30 and 00:30 UTC, matching the cross-midnight shape needed by DeepSeek. +func pricingScheduleForTest() *PricingTimeSchedule { + return &PricingTimeSchedule{ + Timezone: "UTC", + Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {StartTime: "16:30", EndTime: "00:30", Multiplier: 0.5}, + }, + } +} + +func TestCalculateCostSchedule_UsesAttemptStartAndComposesWithServiceTier(t *testing.T) { + tier := schemas.BifrostServiceTierPriority + startedAt := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("gpt-4o", "openai", "chat"): { + Model: "gpt-4o", + Provider: "openai", + Mode: "chat", + InputCostPerToken: new(0.000005), + OutputCostPerToken: new(0.000015), + InputCostPerTokenPriority: new(0.000010), + OutputCostPerTokenPriority: new(0.000030), + }, + }) + s.SetPricingScheduleForTest("gpt-4o", pricingScheduleForTest()) + + usage := &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500} + resp := &schemas.BifrostResponse{ + ChatResponse: &schemas.BifrostChatResponse{ + ServiceTier: &tier, + Usage: usage, + ExtraFields: schemas.BifrostResponseExtraFields{ + RequestType: schemas.ChatCompletionRequest, + RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), + BillingAttemptStartedAt: &startedAt, + }, + }, + } + + breakdown := s.CalculateCostBreakdown(resp, nil) + require.NotNil(t, breakdown) + require.NotNil(t, breakdown.PricingSchedule) + assert.True(t, breakdown.PricingSchedule.TimestampAvailable) + assert.True(t, breakdown.PricingSchedule.Matched) + assert.Equal(t, 0.5, breakdown.PricingSchedule.Multiplier) + // Priority tier first, schedule multiplier second: 0.5 * (1000*10 + 500*30) / 1e6. + assert.InDelta(t, 0.0125, breakdown.TotalCost, 1e-12) + assert.InDelta(t, 0.005, breakdown.InputCost, 1e-12) + assert.InDelta(t, 0.0075, breakdown.OutputCost, 1e-12) +} + +func TestCalculateCostSchedule_ComposesWithContextTier(t *testing.T) { + startedAt := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("long-context", "openai", "chat"): { + Model: "long-context", + Provider: "openai", + Mode: "chat", + InputCostPerToken: new(0.000005), + OutputCostPerToken: new(0.000015), + InputCostPerTokenAbove128kTokens: new(0.000010), + OutputCostPerTokenAbove128kTokens: new(0.000030), + }, + }) + s.SetPricingScheduleForTest("long-context", pricingScheduleForTest()) + + resp := makeChatResponse(schemas.OpenAI, "long-context", &schemas.BifrostLLMUsage{ + PromptTokens: 130000, + CompletionTokens: 500, + TotalTokens: 130500, + }) + resp.ChatResponse.ExtraFields.BillingAttemptStartedAt = &startedAt + + cost := s.CalculateCost(resp, nil) + // Long-context rate first, schedule multiplier second: 0.5*(130000*10 + 500*30)/1e6. + assert.InDelta(t, 0.6575, cost, 1e-12) +} + +func TestCalculateCostSchedule_MissingTimestampUsesBaseRateAndReportsIt(t *testing.T) { + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), + }) + s.SetPricingScheduleForTest("gpt-4o", pricingScheduleForTest()) + + resp := makeChatResponse(schemas.OpenAI, "gpt-4o", &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500}) + breakdown := s.CalculateCostBreakdown(resp, nil) + require.NotNil(t, breakdown) + require.NotNil(t, breakdown.PricingSchedule) + assert.False(t, breakdown.PricingSchedule.TimestampAvailable) + assert.False(t, breakdown.PricingSchedule.Matched) + assert.Equal(t, 1.0, breakdown.PricingSchedule.Multiplier) + assert.InDelta(t, 1000*0.000005+500*0.000015, breakdown.TotalCost, 1e-12) +} + +func TestCalculateCostSchedule_UnmatchedKeepsBaseRate(t *testing.T) { + startedAt := time.Date(2026, time.August, 24, 12, 0, 0, 0, time.UTC) + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), + }) + s.SetPricingScheduleForTest("gpt-4o", pricingScheduleForTest()) + + resp := makeChatResponse(schemas.OpenAI, "gpt-4o", &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500}) + resp.ChatResponse.ExtraFields.BillingAttemptStartedAt = &startedAt + breakdown := s.CalculateCostBreakdown(resp, nil) + require.NotNil(t, breakdown) + require.NotNil(t, breakdown.PricingSchedule) + assert.True(t, breakdown.PricingSchedule.TimestampAvailable) + assert.False(t, breakdown.PricingSchedule.Matched) + assert.InDelta(t, 1000*0.000005+500*0.000015, breakdown.TotalCost, 1e-12) +} + +func TestCalculateCostForUsageWithOptions_UsesExplicitBillingTime(t *testing.T) { + startedAt := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), + }) + s.SetPricingScheduleForTest("gpt-4o", pricingScheduleForTest()) + usage := &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500} + + base := s.CalculateCostBreakdownForUsage(usage, schemas.OpenAI, "gpt-4o", schemas.ChatCompletionRequest, nil) + discounted := s.CalculateCostBreakdownForUsageWithOptions( + usage, schemas.OpenAI, "gpt-4o", schemas.ChatCompletionRequest, nil, + &CostCalculationOptions{BillingAttemptStartedAt: &startedAt}, + ) + require.NotNil(t, discounted) + require.NotNil(t, discounted.PricingSchedule) + assert.InDelta(t, base.TotalCost*0.5, discounted.TotalCost, 1e-12) + assert.True(t, discounted.PricingSchedule.Matched) +} + +func TestCalculateCostSchedule_ProviderCostIsNotRescaled(t *testing.T) { + startedAt := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + providerCost := &schemas.BifrostCost{TotalCost: 1.25} + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), + }) + s.SetPricingScheduleForTest("gpt-4o", pricingScheduleForTest()) + resp := makeChatResponse(schemas.OpenAI, "gpt-4o", &schemas.BifrostLLMUsage{ + PromptTokens: 1000, CompletionTokens: 500, Cost: providerCost, + }) + resp.ChatResponse.ExtraFields.BillingAttemptStartedAt = &startedAt + + assert.Equal(t, providerCost, s.CalculateCostBreakdown(resp, nil)) +} + +func TestCalculateGuardrailCost_UsesEachJudgeStartIndependently(t *testing.T) { + offPeak := time.Date(2026, time.August, 24, 12, 0, 0, 0, time.UTC) + discounted := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("judge", "openai", "chat"): chatPricing(0.000001, 0.000002), + }) + s.SetPricingScheduleForTest("judge", pricingScheduleForTest()) + debug := &schemas.BifrostGuardrailDebug{JudgeCalls: []schemas.BifrostGuardrailJudgeCall{ + {JudgeProvider: schemas.OpenAI, JudgeModel: "judge", PromptTokens: 1000, CompletionTokens: 1000, StartedAt: &offPeak}, + {JudgeProvider: schemas.OpenAI, JudgeModel: "judge", PromptTokens: 1000, CompletionTokens: 1000, StartedAt: &discounted}, + }} + + // Base cost is 0.003 per call; only the second call is discounted. + assert.InDelta(t, 0.003+0.0015, s.CalculateGuardrailCost(debug, nil), 1e-12) +} + +func TestCalculateCostSchedule_FollowsPricingModelCandidates(t *testing.T) { + startedAt := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + s := testStoreWithPricing(map[string]configstoreTables.TableModelPricing{ + makeKey("canonical-model", "openai", "chat"): chatPricing(0.000005, 0.000015), + }) + s.SetPricingScheduleForTest("canonical-model", pricingScheduleForTest()) + + resp := &schemas.BifrostResponse{ + ChatResponse: &schemas.BifrostChatResponse{ + Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500}, + ExtraFields: schemas.BifrostResponseExtraFields{ + RequestType: schemas.ChatCompletionRequest, + RoutingInfo: schemas.RoutingInfo{ + Provider: schemas.OpenAI, + Model: "alias-model", + ResolvedKeyAlias: &schemas.ResolvedKeyAlias{ + ModelID: "wire-model", + ModelName: bifrost.Ptr("canonical-model"), + }, + }, + BillingAttemptStartedAt: &startedAt, + }, + }, + } + + breakdown := s.CalculateCostBreakdown(resp, nil) + require.NotNil(t, breakdown) + require.NotNil(t, breakdown.PricingSchedule) + assert.True(t, breakdown.PricingSchedule.Matched) + assert.Equal(t, 0.5, breakdown.PricingSchedule.Multiplier) + assert.InDelta(t, 0.5*(1000*0.000005+500*0.000015), breakdown.TotalCost, 1e-12) +} diff --git a/framework/modelcatalog/datasheet/schedule.go b/framework/modelcatalog/datasheet/schedule.go new file mode 100644 index 00000000000..ad61a76b0a7 --- /dev/null +++ b/framework/modelcatalog/datasheet/schedule.go @@ -0,0 +1,358 @@ +package datasheet + +import ( + "fmt" + "math" + "strings" + "time" +) + +// Pricing schedule calendars. The generic evaluator intentionally supports +// only deterministic calendars. Provider holiday/workday calendars require +// authoritative external data and should be separate future calendar kinds. +const ( + // PricingScheduleCalendarNone ignores the date and evaluates only time. + PricingScheduleCalendarNone = "none" + // PricingScheduleCalendarISOWeekday evaluates ISO weekday plus time. + PricingScheduleCalendarISOWeekday = "iso_weekday" +) + +// PricingTimeSchedule defines recurring time-based price multipliers. +// +// A schedule is provider-generic: DeepSeek is the first consumer, but other +// providers can describe their own timezone/day/time windows with the same +// structure. No rule is a peak/off-peak statement by itself; each rule simply +// supplies the multiplier to apply to the final resolved tier rate. +type PricingTimeSchedule struct { + // Timezone is an IANA location name (for example "Asia/Shanghai"). + Timezone string `json:"timezone,omitempty"` + // Calendar selects how dates are interpreted. Initial supported values are + // "none" and "iso_weekday". Holidays and makeup workdays are intentionally + // not modeled yet. + Calendar string `json:"calendar,omitempty"` + // Rules are recurring half-open local-time windows. + Rules []PricingTimeRule `json:"rules,omitempty"` +} + +// PricingTimeRule is one recurring schedule rule. +type PricingTimeRule struct { + // Days are ISO weekday names for iso_weekday schedules. Empty means every + // day. Values are case-insensitive and normalized to lowercase. + Days []string `json:"days,omitempty"` + // StartTime and EndTime are inclusive local wall-clock boundaries in HH:MM + // format. The window is [StartTime, EndTime). EndTime may precede StartTime, + // which creates a cross-midnight window. + StartTime string `json:"start_time,omitempty"` + EndTime string `json:"end_time,omitempty"` + // Multiplier is applied to the final resolved pricing rate. It must be + // greater than zero. + Multiplier float64 `json:"multiplier"` +} + +// PricingScheduleEvaluation reports the multiplier selected for one instant. +type PricingScheduleEvaluation struct { + Multiplier float64 + Matched bool + RuleIndex int +} + +// EvaluatePricingTimeSchedule evaluates schedule at the authoritative billing +// instant (the provider attempt start). The returned multiplier is 1 and +// Matched is false when no rule matches. Callers must not substitute another +// timestamp for at. +func EvaluatePricingTimeSchedule(schedule *PricingTimeSchedule, at time.Time) (PricingScheduleEvaluation, error) { + if schedule == nil { + return PricingScheduleEvaluation{Multiplier: 1}, nil + } + + calendar, err := normalizePricingScheduleCalendar(schedule.Calendar) + if err != nil { + return PricingScheduleEvaluation{}, err + } + location, err := time.LoadLocation(strings.TrimSpace(schedule.Timezone)) + if err != nil { + return PricingScheduleEvaluation{}, fmt.Errorf("invalid pricing schedule timezone %q: %w", schedule.Timezone, err) + } + + local := at.In(location) + for index, rule := range schedule.Rules { + matched, err := rule.matches(local, calendar) + if err != nil { + return PricingScheduleEvaluation{}, fmt.Errorf("pricing schedule rule %d: %w", index, err) + } + if matched { + return PricingScheduleEvaluation{ + Multiplier: rule.Multiplier, + Matched: true, + RuleIndex: index, + }, nil + } + } + return PricingScheduleEvaluation{Multiplier: 1}, nil +} + +// ValidatePricingTimeSchedule validates a schedule without evaluating it. +func ValidatePricingTimeSchedule(schedule *PricingTimeSchedule) error { + if schedule == nil { + return nil + } + + calendar, err := normalizePricingScheduleCalendar(schedule.Calendar) + if err != nil { + return err + } + if _, err := time.LoadLocation(strings.TrimSpace(schedule.Timezone)); err != nil { + return fmt.Errorf("invalid pricing schedule timezone %q: %w", schedule.Timezone, err) + } + if len(schedule.Rules) == 0 { + return fmt.Errorf("pricing schedule rules must not be empty") + } + + daySets := make([]map[string]struct{}, len(schedule.Rules)) + for index, rule := range schedule.Rules { + if err := rule.validate(calendar); err != nil { + return fmt.Errorf("pricing schedule rule %d: %w", index, err) + } + daySets[index] = rule.daySet() + // calendar=none ignores Dates/Days; all rules therefore share the same + // virtual day and must use the full weekday set during overlap checks. + if calendar == PricingScheduleCalendarNone { + daySets[index] = allPricingDays() + } + } + + for i := 0; i < len(schedule.Rules); i++ { + for j := i + 1; j < len(schedule.Rules); j++ { + if pricingRulesOverlap(schedule.Rules[i], schedule.Rules[j], daySets[i], daySets[j]) { + if schedule.Rules[i].Multiplier == schedule.Rules[j].Multiplier { + // Identical-multiplier overlaps are deterministic: the first + // matching rule yields the same result either way. This also + // lets a weekend full-day rule coexist with a weekday + // cross-midnight window that spills into Saturday morning. + continue + } + return fmt.Errorf("pricing schedule rules %d and %d overlap with different multipliers", i, j) + } + } + } + return nil +} + +func normalizePricingScheduleCalendar(calendar string) (string, error) { + switch value := strings.ToLower(strings.TrimSpace(calendar)); value { + case PricingScheduleCalendarNone, PricingScheduleCalendarISOWeekday: + return value, nil + default: + return "", fmt.Errorf("unsupported pricing schedule calendar %q", calendar) + } +} + +func (rule PricingTimeRule) validate(calendar string) error { + if rule.Multiplier <= 0 || math.IsNaN(rule.Multiplier) || math.IsInf(rule.Multiplier, 0) { + return fmt.Errorf("multiplier must be a finite value greater than zero") + } + if _, err := parsePricingClock(rule.StartTime); err != nil { + return fmt.Errorf("invalid start_time %q: %w", rule.StartTime, err) + } + if _, err := parsePricingClock(rule.EndTime); err != nil { + return fmt.Errorf("invalid end_time %q: %w", rule.EndTime, err) + } + if calendar == PricingScheduleCalendarISOWeekday { + for day := range rule.daySet() { + if !isPricingDay(day) { + return fmt.Errorf("unknown weekday %q", day) + } + } + } + return nil +} + +func (rule PricingTimeRule) matches(at time.Time, calendar string) (bool, error) { + if err := rule.validate(calendar); err != nil { + return false, err + } + + startMinute, startErr := parsePricingClockMinutes(rule.StartTime) + if startErr != nil { + return false, startErr + } + endMinute, endErr := parsePricingClockMinutes(rule.EndTime) + if endErr != nil { + return false, endErr + } + + currentMinute := clockMinutes(at) + if startMinute == endMinute { + days := rule.daySet() + if calendar == PricingScheduleCalendarISOWeekday && len(days) != 0 && !containsPricingDay(days, at.Weekday()) { + return false, nil + } + return true, nil + } + + effectiveAt := at + if startMinute > endMinute && currentMinute < endMinute { + // The morning tail of a wrapped window belongs to the previous weekday. + effectiveAt = at.AddDate(0, 0, -1) + } + + duration := endMinute - startMinute + if duration < 0 { + duration += 24 * 60 + } + offset := (currentMinute - startMinute + 24*60) % (24 * 60) + if offset >= duration { + return false, nil + } + + days := rule.daySet() + if calendar == PricingScheduleCalendarISOWeekday && len(days) != 0 && !containsPricingDay(days, effectiveAt.Weekday()) { + return false, nil + } + return true, nil +} + +func (rule PricingTimeRule) daySet() map[string]struct{} { + days := make(map[string]struct{}, len(rule.Days)) + for _, day := range rule.Days { + days[strings.ToLower(strings.TrimSpace(day))] = struct{}{} + } + return days +} + +func parsePricingClock(value string) (time.Duration, error) { + if len(value) != 5 || value[2] != ':' { + return 0, fmt.Errorf("must be HH:MM") + } + minutes, err := parsePricingClockMinutes(value) + if err != nil { + return 0, err + } + return time.Duration(minutes) * time.Minute, nil +} + +func parsePricingClockMinutes(value string) (int, error) { + if len(value) != 5 || value[2] != ':' || + value[0] < '0' || value[0] > '9' || + value[1] < '0' || value[1] > '9' || + value[3] < '0' || value[3] > '9' || + value[4] < '0' || value[4] > '9' { + return 0, fmt.Errorf("must be HH:MM") + } + hour := int(value[0]-'0')*10 + int(value[1]-'0') + minute := int(value[3]-'0')*10 + int(value[4]-'0') + if hour > 23 || minute > 59 { + return 0, fmt.Errorf("must be between 00:00 and 23:59") + } + return hour*60 + minute, nil +} + +func isPricingDay(day string) bool { + _, ok := allPricingDays()[day] + return ok +} + +func containsPricingDay(days map[string]struct{}, weekday time.Weekday) bool { + _, ok := days[strings.ToLower(weekday.String())] + return ok +} + +func clockMinutes(at time.Time) int { + hour, minute, _ := at.Clock() + return hour*60 + minute +} + +func pricingRuleDaysIntersect(aDays, bDays map[string]struct{}) bool { + if len(aDays) == 0 || len(bDays) == 0 { + return false + } + for day := range aDays { + if _, ok := bDays[day]; ok { + return true + } + } + return false +} + +type pricingRuleInterval struct { + days map[string]struct{} + startMinute int + endMinute int +} + +// pricingRulesOverlap reports whether two rules can match the same instant. +// Empty day sets mean every day, so they are normalized to the full weekday +// set before interval comparison. Under calendar=none every rule also shares +// the same virtual date and therefore the full weekday set. +func pricingRulesOverlap(a, b PricingTimeRule, aDays, bDays map[string]struct{}) bool { + aIntervals := normalizedPricingRuleIntervals(a, aDays) + bIntervals := normalizedPricingRuleIntervals(b, bDays) + for _, aInterval := range aIntervals { + for _, bInterval := range bIntervals { + if pricingRuleDaysIntersect(aInterval.days, bInterval.days) && + aInterval.startMinute < bInterval.endMinute && + bInterval.startMinute < aInterval.endMinute { + return true + } + } + } + return false +} + +func normalizedPricingRuleIntervals(rule PricingTimeRule, days map[string]struct{}) []pricingRuleInterval { + if len(days) == 0 { + days = allPricingDays() + } + + start, _ := parsePricingClockMinutes(rule.StartTime) + end, _ := parsePricingClockMinutes(rule.EndTime) + if start == end { + return []pricingRuleInterval{{days: days, startMinute: 0, endMinute: 24 * 60}} + } + if start < end { + return []pricingRuleInterval{{days: days, startMinute: start, endMinute: end}} + } + + // A wrapped window covers the start day's evening and the next day's + // morning. Normalizing to these two intervals makes weekday validation + // independent of the wall clock under evaluation. + return []pricingRuleInterval{ + {days: days, startMinute: start, endMinute: 24 * 60}, + {days: nextPricingDays(days), startMinute: 0, endMinute: end}, + } +} + +func allPricingDays() map[string]struct{} { + return map[string]struct{}{ + "sunday": {}, "monday": {}, "tuesday": {}, "wednesday": {}, "thursday": {}, "friday": {}, "saturday": {}, + } +} + +func nextPricingDays(days map[string]struct{}) map[string]struct{} { + nextDays := make(map[string]struct{}, len(days)) + for day := range days { + nextDays[nextPricingDay(day)] = struct{}{} + } + return nextDays +} + +func nextPricingDay(day string) string { + switch day { + case "sunday": + return "monday" + case "monday": + return "tuesday" + case "tuesday": + return "wednesday" + case "wednesday": + return "thursday" + case "thursday": + return "friday" + case "friday": + return "saturday" + case "saturday": + return "sunday" + default: + return day + } +} diff --git a/framework/modelcatalog/datasheet/schedule_test.go b/framework/modelcatalog/datasheet/schedule_test.go new file mode 100644 index 00000000000..df1bedd7293 --- /dev/null +++ b/framework/modelcatalog/datasheet/schedule_test.go @@ -0,0 +1,230 @@ +package datasheet + +import ( + "math" + "testing" + "time" +) + +func pricingScheduleAt(day string, hour, minute int) time.Time { + return time.Date(2026, time.August, mustAtoi(day), hour, minute, 0, 0, time.UTC) +} + +func mustAtoi(value string) int { + digits := map[byte]int{'0': 0, '1': 1, '2': 2, '3': 3, '4': 4, '5': 5, '6': 6, '7': 7, '8': 8, '9': 9} + if len(value) != 2 { + panic("expected two-digit day") + } + return digits[value[0]]*10 + digits[value[1]] +} + +func TestEvaluatePricingTimeSchedule(t *testing.T) { + // 2026-08-23 is a Sunday in UTC. + schedule := &PricingTimeSchedule{ + Timezone: "UTC", + Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"saturday", "sunday"}, StartTime: "00:00", EndTime: "00:00", Multiplier: 0.5}, + {Days: []string{"monday", "tuesday", "wednesday", "thursday", "friday"}, StartTime: "16:30", EndTime: "00:30", Multiplier: 0.5}, + }, + } + + tests := []struct { + name string + at time.Time + multiplier float64 + matched bool + }{ + {name: "sunday full day", at: pricingScheduleAt("23", 12, 0), multiplier: 0.5, matched: true}, + {name: "saturday full day", at: pricingScheduleAt("22", 0, 0), multiplier: 0.5, matched: true}, + {name: "monday before window", at: pricingScheduleAt("24", 16, 29), multiplier: 1}, + {name: "tuesday before window", at: pricingScheduleAt("25", 16, 29), multiplier: 1}, + {name: "weekday start is inclusive", at: pricingScheduleAt("24", 16, 30), multiplier: 0.5, matched: true}, + {name: "weekday end is exclusive", at: pricingScheduleAt("26", 0, 30), multiplier: 1}, + {name: "cross midnight before end", at: pricingScheduleAt("25", 0, 29), multiplier: 0.5, matched: true}, + {name: "weekday peak", at: pricingScheduleAt("25", 12, 0), multiplier: 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := EvaluatePricingTimeSchedule(schedule, tt.at) + if err != nil { + t.Fatal(err) + } + if got.Matched != tt.matched || got.Multiplier != tt.multiplier { + t.Fatalf("got multiplier=%f matched=%v, want multiplier=%f matched=%v", got.Multiplier, got.Matched, tt.multiplier, tt.matched) + } + }) + } +} + +func TestEvaluatePricingTimeScheduleTimezone(t *testing.T) { + schedule := &PricingTimeSchedule{ + Timezone: "Asia/Shanghai", + Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"monday"}, StartTime: "09:00", EndTime: "10:00", Multiplier: 2}, + }, + } + + at, err := time.Parse(time.RFC3339, "2026-08-24T01:30:00Z") + if err != nil { + t.Fatal(err) + } + got, err := EvaluatePricingTimeSchedule(schedule, at) + if err != nil { + t.Fatal(err) + } + if !got.Matched || got.Multiplier != 2 { + t.Fatalf("expected Shanghai 09:30 to match, got %+v", got) + } +} + +func TestEvaluatePricingTimeScheduleNilUsesBaseMultiplier(t *testing.T) { + got, err := EvaluatePricingTimeSchedule(nil, time.Now()) + if err != nil { + t.Fatal(err) + } + if got.Matched || got.Multiplier != 1 { + t.Fatalf("expected base multiplier, got %+v", got) + } +} + +func TestValidatePricingTimeSchedule(t *testing.T) { + valid := &PricingTimeSchedule{ + Timezone: "UTC", + Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"saturday", "sunday"}, StartTime: "00:00", EndTime: "00:00", Multiplier: 0.5}, + {Days: []string{"monday"}, StartTime: "09:00", EndTime: "10:00", Multiplier: 2}, + }, + } + if err := ValidatePricingTimeSchedule(valid); err != nil { + t.Fatalf("expected valid schedule: %v", err) + } + + tests := []struct { + name string + schedule PricingTimeSchedule + }{ + {name: "invalid calendar", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: "work_day", Rules: []PricingTimeRule{{Multiplier: 1}}}}, + {name: "invalid timezone", schedule: PricingTimeSchedule{Timezone: "not-a-zone", Calendar: PricingScheduleCalendarNone, Rules: []PricingTimeRule{{Multiplier: 1}}}}, + {name: "empty rules", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarNone}}, + {name: "invalid multiplier", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarNone, Rules: []PricingTimeRule{{Multiplier: 0}}}}, + {name: "NaN multiplier", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarNone, Rules: []PricingTimeRule{{Multiplier: math.NaN()}}}}, + {name: "positive infinity multiplier", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarNone, Rules: []PricingTimeRule{{Multiplier: math.Inf(1)}}}}, + {name: "invalid start", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarNone, Rules: []PricingTimeRule{{StartTime: "24:00", EndTime: "01:00", Multiplier: 1}}}}, + {name: "invalid end", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarNone, Rules: []PricingTimeRule{{StartTime: "00:00", EndTime: "1:00", Multiplier: 1}}}}, + {name: "non-digit clock", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarNone, Rules: []PricingTimeRule{{StartTime: "0::00", EndTime: "01:00", Multiplier: 1}}}}, + {name: "unknown weekday", schedule: PricingTimeSchedule{Timezone: "UTC", Calendar: PricingScheduleCalendarISOWeekday, Rules: []PricingTimeRule{{Days: []string{"funday"}, Multiplier: 1}}}}, + { + name: "calendar none overlap", + schedule: PricingTimeSchedule{ + Timezone: "UTC", Calendar: PricingScheduleCalendarNone, + Rules: []PricingTimeRule{ + {Days: []string{"monday"}, StartTime: "09:00", EndTime: "10:00", Multiplier: 1}, + {Days: []string{"tuesday"}, StartTime: "09:30", EndTime: "10:30", Multiplier: 2}, + }, + }, + }, + { + name: "overlap same day", + schedule: PricingTimeSchedule{ + Timezone: "UTC", Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {StartTime: "09:00", EndTime: "10:00", Multiplier: 1}, + {StartTime: "09:30", EndTime: "10:30", Multiplier: 2}, + }, + }, + }, + { + name: "full day overlap", + schedule: PricingTimeSchedule{ + Timezone: "UTC", Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"monday"}, StartTime: "00:00", EndTime: "00:00", Multiplier: 1}, + {StartTime: "09:00", EndTime: "10:00", Multiplier: 2}, + }, + }, + }, + { + name: "cross midnight overlap", + schedule: PricingTimeSchedule{ + Timezone: "UTC", Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"monday"}, StartTime: "22:00", EndTime: "02:00", Multiplier: 1}, + {Days: []string{"tuesday"}, StartTime: "00:00", EndTime: "03:00", Multiplier: 2}, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := ValidatePricingTimeSchedule(&tt.schedule); err == nil { + t.Fatal("expected invalid schedule") + } + }) + } +} + +func TestValidatePricingTimeScheduleAllowsIdenticalMultiplierOverlaps(t *testing.T) { + schedule := &PricingTimeSchedule{ + Timezone: "UTC", + Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"saturday", "sunday"}, StartTime: "00:00", EndTime: "00:00", Multiplier: 0.5}, + {Days: []string{"monday", "tuesday", "wednesday", "thursday", "friday"}, StartTime: "16:30", EndTime: "00:30", Multiplier: 0.5}, + }, + } + if err := ValidatePricingTimeSchedule(schedule); err != nil { + t.Fatalf("expected valid DeepSeek-style schedule: %v", err) + } +} + +func TestValidatePricingTimeScheduleAllowsAdjacentWeekdayWindows(t *testing.T) { + schedule := &PricingTimeSchedule{ + Timezone: "UTC", + Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"saturday", "sunday"}, StartTime: "00:00", EndTime: "00:00", Multiplier: 0.5}, + {Days: []string{"monday", "tuesday", "wednesday", "thursday", "friday"}, StartTime: "16:30", EndTime: "23:59", Multiplier: 0.5}, + }, + } + if err := ValidatePricingTimeSchedule(schedule); err != nil { + t.Fatalf("expected valid weekday schedule: %v", err) + } +} + +func TestEvaluatePricingTimeScheduleMondayOnlyCrossMidnight(t *testing.T) { + schedule := &PricingTimeSchedule{ + Timezone: "UTC", + Calendar: PricingScheduleCalendarISOWeekday, + Rules: []PricingTimeRule{ + {Days: []string{"monday"}, StartTime: "22:00", EndTime: "02:00", Multiplier: 2}, + }, + } + + tests := []struct { + name string + at time.Time + matched bool + }{ + {name: "monday evening", at: pricingScheduleAt("24", 23, 0), matched: true}, + {name: "tuesday wrapped tail", at: pricingScheduleAt("25", 1, 0), matched: true}, + {name: "tuesday after tail", at: pricingScheduleAt("25", 2, 0), matched: false}, + {name: "tuesday before wrapped start", at: pricingScheduleAt("25", 21, 0), matched: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := EvaluatePricingTimeSchedule(schedule, tt.at) + if err != nil { + t.Fatal(err) + } + if got.Matched != tt.matched { + t.Fatalf("got %+v, want matched=%v", got, tt.matched) + } + }) + } +} diff --git a/framework/streaming/accumulator.go b/framework/streaming/accumulator.go index e1910bd73c7..e460d9bc325 100644 --- a/framework/streaming/accumulator.go +++ b/framework/streaming/accumulator.go @@ -55,6 +55,7 @@ func (a *Accumulator) putChatStreamChunk(chunk *ChatStreamChunk) { chunk.FinishReason = nil chunk.TokenUsage = nil chunk.ServiceTier = nil + chunk.BillingAttemptStartedAt = nil chunk.RawResponse = nil a.chatStreamChunkPool.Put(chunk) } @@ -110,6 +111,7 @@ func (a *Accumulator) putResponsesStreamChunk(chunk *ResponsesStreamChunk) { chunk.FinishReason = nil chunk.TokenUsage = nil chunk.ServiceTier = nil + chunk.BillingAttemptStartedAt = nil chunk.RawResponse = nil a.responsesStreamChunkPool.Put(chunk) } @@ -130,6 +132,7 @@ func (a *Accumulator) putImageStreamChunk(chunk *ImageStreamChunk) { chunk.Cost = nil chunk.SemanticCacheDebug = nil chunk.TokenUsage = nil + chunk.BillingAttemptStartedAt = nil chunk.RawResponse = nil a.imageStreamChunkPool.Put(chunk) } diff --git a/framework/streaming/accumulator_test.go b/framework/streaming/accumulator_test.go index 3ca11e74910..88a504c3c95 100644 --- a/framework/streaming/accumulator_test.go +++ b/framework/streaming/accumulator_test.go @@ -1041,3 +1041,93 @@ func TestChatStreamingFinishReasonOnTerminalChunk(t *testing.T) { t.Fatalf("accumulated finish_reason = %q, want %q", *processed.Data.FinishReason, "stop") } } + +func TestAccumulatedStreamKeepsHighestIndexBillingAttemptStart(t *testing.T) { + logger := bifrost.NewDefaultLogger(schemas.LogLevelError) + accumulator := NewAccumulator(nil, logger) + t.Cleanup(accumulator.Cleanup) + + requestID := "billing-attempt-high-index" + earlier := time.Date(2026, time.August, 24, 12, 0, 0, 0, time.UTC) + terminal := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + + if err := accumulator.addChatStreamChunk(requestID, StreamTypeChat, &ChatStreamChunk{ + ChunkIndex: 2, + Timestamp: time.Now(), + BillingAttemptStartedAt: &terminal, + }, true); err != nil { + t.Fatal(err) + } + // Multiple plugins may replay an earlier chunk after the terminal chunk. + if err := accumulator.addChatStreamChunk(requestID, StreamTypeChat, &ChatStreamChunk{ + ChunkIndex: 0, + Timestamp: time.Now().Add(time.Second), + BillingAttemptStartedAt: &earlier, + }, false); err != nil { + t.Fatal(err) + } + + data, err := accumulator.processAccumulatedChatStreamingChunks(requestID, nil, true) + if err != nil { + t.Fatal(err) + } + if data.BillingAttemptStartedAt == nil || !data.BillingAttemptStartedAt.Equal(terminal) { + t.Fatalf("expected terminal attempt time %v, got %v", terminal, data.BillingAttemptStartedAt) + } +} + +func TestImageAccumulatedStreamKeepsHighestIndexBillingAttemptStart(t *testing.T) { + accumulator := NewAccumulator(nil, bifrost.NewDefaultLogger(schemas.LogLevelError)) + t.Cleanup(accumulator.Cleanup) + + requestID := "image-billing-attempt" + earlier := time.Date(2026, time.August, 24, 12, 0, 0, 0, time.UTC) + terminal := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + + if err := accumulator.addImageStreamChunk(requestID, &ImageStreamChunk{ + ChunkIndex: 2, + Timestamp: time.Now(), + BillingAttemptStartedAt: &terminal, + }, true); err != nil { + t.Fatal(err) + } + if err := accumulator.addImageStreamChunk(requestID, &ImageStreamChunk{ + ChunkIndex: 0, + Timestamp: time.Now().Add(time.Second), + BillingAttemptStartedAt: &earlier, + }, false); err != nil { + t.Fatal(err) + } + + data, err := accumulator.processAccumulatedImageStreamingChunks(requestID, nil, true) + if err != nil { + t.Fatal(err) + } + if data.BillingAttemptStartedAt == nil || !data.BillingAttemptStartedAt.Equal(terminal) { + t.Fatalf("expected terminal attempt time %v, got %v", terminal, data.BillingAttemptStartedAt) + } +} + +func TestImageStreamErrorCarriesBillingAttemptStart(t *testing.T) { + accumulator := NewAccumulator(nil, bifrost.NewDefaultLogger(schemas.LogLevelError)) + t.Cleanup(accumulator.Cleanup) + + requestID := "image-billing-error" + ctx := schemas.NewBifrostContext(context.Background(), time.Time{}) + ctx.SetValue(schemas.BifrostContextKeyAccumulatorID, requestID) + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + startedAt := time.Date(2026, time.August, 24, 16, 30, 0, 0, time.UTC) + bifrostErr := &schemas.BifrostError{} + bifrostErr.ExtraFields.BillingAttemptStartedAt = &startedAt + + processed, err := accumulator.processImageStreamingResponse(ctx, nil, bifrostErr) + if err != nil { + t.Fatal(err) + } + if processed == nil || processed.Data == nil { + t.Fatal("expected processed error data") + } + if processed.Data.BillingAttemptStartedAt == nil || !processed.Data.BillingAttemptStartedAt.Equal(startedAt) { + t.Fatalf("expected billing attempt time %v, got %v", startedAt, processed.Data.BillingAttemptStartedAt) + } +} diff --git a/framework/streaming/chat.go b/framework/streaming/chat.go index 9a38b945a91..088ecc1e44b 100644 --- a/framework/streaming/chat.go +++ b/framework/streaming/chat.go @@ -468,11 +468,16 @@ func (a *Accumulator) processAccumulatedChatStreamingChunks(requestID string, re // chunk, so retain the newest non-nil value rather than reading only the // highest-index chunk. tierChunkIndex := -1 + billingChunkIndex := -1 for _, streamChunk := range accumulator.ChatStreamChunks { if streamChunk.ServiceTier != nil && streamChunk.ChunkIndex > tierChunkIndex { data.ServiceTier = streamChunk.ServiceTier tierChunkIndex = streamChunk.ChunkIndex } + if streamChunk.BillingAttemptStartedAt != nil && streamChunk.ChunkIndex > billingChunkIndex { + data.BillingAttemptStartedAt = streamChunk.BillingAttemptStartedAt + billingChunkIndex = streamChunk.ChunkIndex + } } // The highest-index chunk can carry a nil finish_reason (a usage-only chunk, // or the synthetic terminal chunk the OpenAI-compatible handler appends after @@ -544,6 +549,9 @@ func (a *Accumulator) processChatStreamingResponse(ctx *schemas.BifrostContext, chunk.ErrorDetails = bifrostErr if bifrostErr != nil { chunk.FinishReason = bifrost.Ptr("error") + if bifrostErr.ExtraFields.BillingAttemptStartedAt != nil { + chunk.BillingAttemptStartedAt = bifrostErr.ExtraFields.BillingAttemptStartedAt + } } else if result != nil && result.TextCompletionResponse != nil { // Handle text completion response directly if len(result.TextCompletionResponse.Choices) > 0 { @@ -562,6 +570,9 @@ func (a *Accumulator) processChatStreamingResponse(ctx *schemas.BifrostContext, if result.TextCompletionResponse.Usage != nil && result.TextCompletionResponse.Usage.TotalTokens > 0 { chunk.TokenUsage = result.TextCompletionResponse.Usage } + if result.TextCompletionResponse.ExtraFields.BillingAttemptStartedAt != nil { + chunk.BillingAttemptStartedAt = result.TextCompletionResponse.ExtraFields.BillingAttemptStartedAt + } chunk.ChunkIndex = result.TextCompletionResponse.ExtraFields.ChunkIndex if result.TextCompletionResponse.ExtraFields.RawResponse != nil { chunk.RawResponse = bifrost.Ptr(fmt.Sprintf("%v", result.TextCompletionResponse.ExtraFields.RawResponse)) @@ -592,6 +603,9 @@ func (a *Accumulator) processChatStreamingResponse(ctx *schemas.BifrostContext, if result.ChatResponse.ServiceTier != nil { chunk.ServiceTier = new(schemas.BifrostServiceTier(*result.ChatResponse.ServiceTier)) } + if result.ChatResponse.ExtraFields.BillingAttemptStartedAt != nil { + chunk.BillingAttemptStartedAt = result.ChatResponse.ExtraFields.BillingAttemptStartedAt + } chunk.ChunkIndex = result.ChatResponse.ExtraFields.ChunkIndex if result.ChatResponse.ExtraFields.RawResponse != nil { chunk.RawResponse = bifrost.Ptr(fmt.Sprintf("%v", result.ChatResponse.ExtraFields.RawResponse)) diff --git a/framework/streaming/images.go b/framework/streaming/images.go index c999dff8257..104519e660d 100644 --- a/framework/streaming/images.go +++ b/framework/streaming/images.go @@ -169,6 +169,17 @@ func (a *Accumulator) processAccumulatedImageStreamingChunks(requestID string, b data.ImageGenerationOutput = completeImage data.ErrorDetails = bifrostErr + // Retain the attempt start from the highest-index chunk that carries one. + // Chunks can be processed by multiple plugins or arrive out of order; arrival + // order must not override the terminal attempt timestamp. + billingChunkIndex := -1 + for _, chunk := range acc.ImageStreamChunks { + if chunk.BillingAttemptStartedAt != nil && chunk.ChunkIndex > billingChunkIndex { + data.BillingAttemptStartedAt = chunk.BillingAttemptStartedAt + billingChunkIndex = chunk.ChunkIndex + } + } + // Update token usage from final chunk if available if len(acc.ImageStreamChunks) > 0 { lastChunk := acc.ImageStreamChunks[len(acc.ImageStreamChunks)-1] @@ -224,6 +235,7 @@ func (a *Accumulator) processImageStreamingResponse(ctx *schemas.BifrostContext, chunk.ErrorDetails = bifrostErr if bifrostErr != nil { chunk.FinishReason = bifrost.Ptr("error") + chunk.BillingAttemptStartedAt = bifrostErr.ExtraFields.BillingAttemptStartedAt } else if result != nil && result.ImageGenerationStreamResponse != nil { // Create a deep copy of the delta to avoid pointing to stack memory var partialImageIndex *int @@ -280,6 +292,7 @@ func (a *Accumulator) processImageStreamingResponse(ctx *schemas.BifrostContext, } chunk.SemanticCacheDebug = result.GetExtraFields().CacheDebug chunk.FinishReason = bifrost.Ptr("completed") + chunk.BillingAttemptStartedAt = result.GetExtraFields().BillingAttemptStartedAt } } diff --git a/framework/streaming/responses.go b/framework/streaming/responses.go index 6e92b897581..78f25caeae0 100644 --- a/framework/streaming/responses.go +++ b/framework/streaming/responses.go @@ -1021,11 +1021,16 @@ func (a *Accumulator) processAccumulatedResponsesStreamingChunks(requestID strin // The response envelope carrying service_tier can precede a later usage-only // event, so retain the newest non-nil tier across the stream. tierChunkIndex := -1 + billingChunkIndex := -1 for _, streamChunk := range accumulator.ResponsesStreamChunks { if streamChunk.ServiceTier != nil && streamChunk.ChunkIndex > tierChunkIndex { data.ServiceTier = streamChunk.ServiceTier tierChunkIndex = streamChunk.ChunkIndex } + if streamChunk.BillingAttemptStartedAt != nil && streamChunk.ChunkIndex > billingChunkIndex { + data.BillingAttemptStartedAt = streamChunk.BillingAttemptStartedAt + billingChunkIndex = streamChunk.ChunkIndex + } } // Accumulate raw response using strings.Builder to avoid O(n^2) string concatenation @@ -1071,6 +1076,9 @@ func (a *Accumulator) processResponsesStreamingResponse(ctx *schemas.BifrostCont if bifrostErr != nil { chunk.FinishReason = bifrost.Ptr("error") + if bifrostErr.ExtraFields.BillingAttemptStartedAt != nil { + chunk.BillingAttemptStartedAt = bifrostErr.ExtraFields.BillingAttemptStartedAt + } if bifrostErr.ExtraFields.RawResponse != nil { if rawBytes, marshalErr := sonic.Marshal(bifrostErr.ExtraFields.RawResponse); marshalErr == nil { chunk.RawResponse = bifrost.Ptr(string(rawBytes)) @@ -1096,6 +1104,9 @@ func (a *Accumulator) processResponsesStreamingResponse(ctx *schemas.BifrostCont result.ResponsesStreamResponse.Response.ServiceTier != nil { chunk.ServiceTier = new(schemas.BifrostServiceTier(*result.ResponsesStreamResponse.Response.ServiceTier)) } + if result.ResponsesStreamResponse.ExtraFields.BillingAttemptStartedAt != nil { + chunk.BillingAttemptStartedAt = result.ResponsesStreamResponse.ExtraFields.BillingAttemptStartedAt + } chunk.ChunkIndex = result.ResponsesStreamResponse.ExtraFields.ChunkIndex if isFinalChunk { if a.pricingManager != nil { diff --git a/framework/streaming/types.go b/framework/streaming/types.go index 7ce0b06b60b..0cfe3f34c18 100644 --- a/framework/streaming/types.go +++ b/framework/streaming/types.go @@ -22,30 +22,31 @@ const ( // AccumulatedData contains the accumulated data for a stream type AccumulatedData struct { - RequestID string - Model string - Status string - Stream bool - Latency int64 // in milliseconds - TimeToFirstToken int64 // Time to first token in milliseconds (streaming only) - StartTimestamp time.Time - EndTimestamp time.Time - OutputMessage *schemas.ChatMessage - OutputMessages []schemas.ResponsesMessage // For responses API - ToolCalls []schemas.ChatAssistantMessageToolCall - ErrorDetails *schemas.BifrostError - TokenUsage *schemas.BifrostLLMUsage - ServiceTier *schemas.BifrostServiceTier - CacheDebug *schemas.BifrostCacheDebug - GuardrailDebug *schemas.BifrostGuardrailDebug - Cost *float64 - AudioOutput *schemas.BifrostSpeechResponse - TranscriptionOutput *schemas.BifrostTranscriptionResponse - ImageGenerationOutput *schemas.BifrostImageGenerationResponse - PassthroughOutput *schemas.BifrostPassthroughResponse // For passthrough streaming - FinishReason *string - LogProbs *schemas.BifrostLogProbs - RawResponse *string + RequestID string + Model string + Status string + Stream bool + Latency int64 // in milliseconds + TimeToFirstToken int64 // Time to first token in milliseconds (streaming only) + StartTimestamp time.Time + EndTimestamp time.Time + OutputMessage *schemas.ChatMessage + OutputMessages []schemas.ResponsesMessage // For responses API + ToolCalls []schemas.ChatAssistantMessageToolCall + ErrorDetails *schemas.BifrostError + TokenUsage *schemas.BifrostLLMUsage + BillingAttemptStartedAt *time.Time + ServiceTier *schemas.BifrostServiceTier + CacheDebug *schemas.BifrostCacheDebug + GuardrailDebug *schemas.BifrostGuardrailDebug + Cost *float64 + AudioOutput *schemas.BifrostSpeechResponse + TranscriptionOutput *schemas.BifrostTranscriptionResponse + ImageGenerationOutput *schemas.BifrostImageGenerationResponse + PassthroughOutput *schemas.BifrostPassthroughResponse // For passthrough streaming + FinishReason *string + LogProbs *schemas.BifrostLogProbs + RawResponse *string } // AudioStreamChunk represents a single streaming chunk @@ -76,47 +77,50 @@ type TranscriptionStreamChunk struct { // ChatStreamChunk represents a single streaming chunk type ChatStreamChunk struct { - Timestamp time.Time // When chunk was received - Delta *schemas.ChatStreamResponseChoiceDelta // The actual delta content - FinishReason *string // If this is the final chunk - LogProbs *schemas.BifrostLogProbs // LogProbs if available - TokenUsage *schemas.BifrostLLMUsage // Token usage if available - ServiceTier *schemas.BifrostServiceTier // Served OpenAI tier if available - SemanticCacheDebug *schemas.BifrostCacheDebug // Semantic cache debug if available - GuardrailDebug *schemas.BifrostGuardrailDebug // Guardrail debug if available - Cost *float64 // Cost in dollars from pricing plugin - ErrorDetails *schemas.BifrostError // Error if any - ChunkIndex int // Index of the chunk in the stream - RawResponse *string // Raw response if available + Timestamp time.Time // When chunk was received + Delta *schemas.ChatStreamResponseChoiceDelta // The actual delta content + FinishReason *string // If this is the final chunk + LogProbs *schemas.BifrostLogProbs // LogProbs if available + TokenUsage *schemas.BifrostLLMUsage // Token usage if available + BillingAttemptStartedAt *time.Time // Attempt start for time-based pricing + ServiceTier *schemas.BifrostServiceTier // Served OpenAI tier if available + SemanticCacheDebug *schemas.BifrostCacheDebug // Semantic cache debug if available + GuardrailDebug *schemas.BifrostGuardrailDebug // Guardrail debug if available + Cost *float64 // Cost in dollars from pricing plugin + ErrorDetails *schemas.BifrostError // Error if any + ChunkIndex int // Index of the chunk in the stream + RawResponse *string // Raw response if available } // ResponsesStreamChunk represents a single responses streaming chunk type ResponsesStreamChunk struct { - Timestamp time.Time // When chunk was received - StreamResponse *schemas.BifrostResponsesStreamResponse // The actual stream response - FinishReason *string // If this is the final chunk - TokenUsage *schemas.BifrostLLMUsage // Token usage if available - ServiceTier *schemas.BifrostServiceTier // Served OpenAI tier if available - SemanticCacheDebug *schemas.BifrostCacheDebug // Semantic cache debug if available - GuardrailDebug *schemas.BifrostGuardrailDebug // Guardrail debug if available - Cost *float64 // Cost in dollars from pricing plugin - ErrorDetails *schemas.BifrostError // Error if any - ChunkIndex int // Index of the chunk in the stream - RawResponse *string + Timestamp time.Time // When chunk was received + StreamResponse *schemas.BifrostResponsesStreamResponse // The actual stream response + FinishReason *string // If this is the final chunk + TokenUsage *schemas.BifrostLLMUsage // Token usage if available + BillingAttemptStartedAt *time.Time // Attempt start for time-based pricing + ServiceTier *schemas.BifrostServiceTier // Served OpenAI tier if available + SemanticCacheDebug *schemas.BifrostCacheDebug // Semantic cache debug if available + GuardrailDebug *schemas.BifrostGuardrailDebug // Guardrail debug if available + Cost *float64 // Cost in dollars from pricing plugin + ErrorDetails *schemas.BifrostError // Error if any + ChunkIndex int // Index of the chunk in the stream + RawResponse *string } // ImageStreamChunk represents a single image streaming chunk type ImageStreamChunk struct { - Timestamp time.Time // When chunk was received - Delta *schemas.BifrostImageGenerationStreamResponse // The actual stream response - FinishReason *string // If this is the final chunk - ChunkIndex int // Index of the chunk in the stream - ImageIndex int // Index of the image in the stream - ErrorDetails *schemas.BifrostError // Error if any - Cost *float64 // Cost in dollars from pricing plugin - SemanticCacheDebug *schemas.BifrostCacheDebug // Semantic cache debug if available - TokenUsage *schemas.ImageUsage // Token usage if available - RawResponse *string // Raw response if available + Timestamp time.Time // When chunk was received + Delta *schemas.BifrostImageGenerationStreamResponse // The actual stream response + FinishReason *string // If this is the final chunk + ChunkIndex int // Index of the chunk in the stream + ImageIndex int // Index of the image in the stream + ErrorDetails *schemas.BifrostError // Error if any + Cost *float64 // Cost in dollars from pricing plugin + SemanticCacheDebug *schemas.BifrostCacheDebug // Semantic cache debug if available + TokenUsage *schemas.ImageUsage // Token usage if available + BillingAttemptStartedAt *time.Time // Attempt start for time-based pricing + RawResponse *string // Raw response if available } // StreamAccumulator manages accumulation of streaming chunks @@ -349,11 +353,12 @@ func (p *ProcessedStreamResponse) ToBifrostResponse() *schemas.BifrostResponse { resp.TextCompletionResponse = textResp resp.TextCompletionResponse.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.TextCompletionRequest, - Provider: p.Provider, - OriginalModelRequested: p.RequestedModel, - ResolvedModelUsed: p.ResolvedModel, - Latency: p.Data.Latency, + RequestType: schemas.TextCompletionRequest, + Provider: p.Provider, + OriginalModelRequested: p.RequestedModel, + ResolvedModelUsed: p.ResolvedModel, + Latency: p.Data.Latency, + BillingAttemptStartedAt: p.Data.BillingAttemptStartedAt, } if p.RawRequest != nil { resp.TextCompletionResponse.ExtraFields.RawRequest = p.RawRequest @@ -404,11 +409,12 @@ func (p *ProcessedStreamResponse) ToBifrostResponse() *schemas.BifrostResponse { resp.ChatResponse = chatResp resp.ChatResponse.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - Provider: p.Provider, - OriginalModelRequested: p.RequestedModel, - ResolvedModelUsed: p.ResolvedModel, - Latency: p.Data.Latency, + RequestType: schemas.ChatCompletionRequest, + Provider: p.Provider, + OriginalModelRequested: p.RequestedModel, + ResolvedModelUsed: p.ResolvedModel, + Latency: p.Data.Latency, + BillingAttemptStartedAt: p.Data.BillingAttemptStartedAt, } if p.RawRequest != nil { resp.ChatResponse.ExtraFields.RawRequest = p.RawRequest @@ -432,11 +438,12 @@ func (p *ProcessedStreamResponse) ToBifrostResponse() *schemas.BifrostResponse { responsesResp.Usage = p.Data.TokenUsage.ToResponsesResponseUsage() } responsesResp.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.ResponsesRequest, - Provider: p.Provider, - OriginalModelRequested: p.RequestedModel, - ResolvedModelUsed: p.ResolvedModel, - Latency: p.Data.Latency, + RequestType: schemas.ResponsesRequest, + Provider: p.Provider, + OriginalModelRequested: p.RequestedModel, + ResolvedModelUsed: p.ResolvedModel, + Latency: p.Data.Latency, + BillingAttemptStartedAt: p.Data.BillingAttemptStartedAt, } if p.RawRequest != nil { responsesResp.ExtraFields.RawRequest = p.RawRequest @@ -458,11 +465,12 @@ func (p *ProcessedStreamResponse) ToBifrostResponse() *schemas.BifrostResponse { } resp.SpeechResponse = speechResp resp.SpeechResponse.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.SpeechRequest, - Provider: p.Provider, - OriginalModelRequested: p.RequestedModel, - ResolvedModelUsed: p.ResolvedModel, - Latency: p.Data.Latency, + RequestType: schemas.SpeechRequest, + Provider: p.Provider, + OriginalModelRequested: p.RequestedModel, + ResolvedModelUsed: p.ResolvedModel, + Latency: p.Data.Latency, + BillingAttemptStartedAt: p.Data.BillingAttemptStartedAt, } if p.RawRequest != nil { resp.SpeechResponse.ExtraFields.RawRequest = p.RawRequest @@ -483,11 +491,12 @@ func (p *ProcessedStreamResponse) ToBifrostResponse() *schemas.BifrostResponse { } resp.TranscriptionResponse = transcriptionResp resp.TranscriptionResponse.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.TranscriptionRequest, - Provider: p.Provider, - OriginalModelRequested: p.RequestedModel, - ResolvedModelUsed: p.ResolvedModel, - Latency: p.Data.Latency, + RequestType: schemas.TranscriptionRequest, + Provider: p.Provider, + OriginalModelRequested: p.RequestedModel, + ResolvedModelUsed: p.ResolvedModel, + Latency: p.Data.Latency, + BillingAttemptStartedAt: p.Data.BillingAttemptStartedAt, } if p.RawRequest != nil { resp.TranscriptionResponse.ExtraFields.RawRequest = p.RawRequest @@ -520,11 +529,12 @@ func (p *ProcessedStreamResponse) ToBifrostResponse() *schemas.BifrostResponse { } resp.ImageGenerationResponse = imageResp resp.ImageGenerationResponse.ExtraFields = schemas.BifrostResponseExtraFields{ - RequestType: schemas.ImageGenerationRequest, - Provider: p.Provider, - OriginalModelRequested: p.RequestedModel, - ResolvedModelUsed: p.ResolvedModel, - Latency: p.Data.Latency, + RequestType: schemas.ImageGenerationRequest, + Provider: p.Provider, + OriginalModelRequested: p.RequestedModel, + ResolvedModelUsed: p.ResolvedModel, + Latency: p.Data.Latency, + BillingAttemptStartedAt: p.Data.BillingAttemptStartedAt, } if p.RawRequest != nil { resp.ImageGenerationResponse.ExtraFields.RawRequest = p.RawRequest diff --git a/framework/tracing/tracer.go b/framework/tracing/tracer.go index da466d5d2cb..d55cc88deba 100644 --- a/framework/tracing/tracer.go +++ b/framework/tracing/tracer.go @@ -785,6 +785,7 @@ func (t *Tracer) ProcessStreamingChunk(ctx *schemas.BifrostContext, traceID stri // response envelope, so it has to be copied across explicitly or the tier the // accumulator resolved across chunks is lost and the row reprices at standard // rates. + accResult.BillingAttemptStartedAt = processedResp.Data.BillingAttemptStartedAt accResult.ServiceTier = processedResp.Data.ServiceTier accResult.Cost = processedResp.Data.Cost accResult.CacheDebug = processedResp.Data.CacheDebug diff --git a/plugins/logging/costfidelity_test.go b/plugins/logging/costfidelity_test.go index f2d2161125a..7847e62c1e4 100644 --- a/plugins/logging/costfidelity_test.go +++ b/plugins/logging/costfidelity_test.go @@ -245,6 +245,29 @@ func TestDeserializeFieldsClearsDegradedFlagAfterHydration(t *testing.T) { } } +// TestApplyErrorBillingFromBilledUsagePersistsAttemptStart ensures failed or cancelled +// requests that already consumed tokens keep the attempt's start time. It cannot be +// inferred later: Timestamp records completion and BilledUsage has no temporal field. +func TestApplyErrorBillingFromBilledUsagePersistsAttemptStart(t *testing.T) { + plugin := newCostFidelityPlugin(t) + startedAt := time.Date(2026, 8, 22, 22, 30, 0, 0, time.UTC) + entry := &logstore.Log{} + + plugin.applyErrorBillingFromBilledUsage( + schemas.NewBifrostContext(context.Background(), schemas.NoDeadline), + entry, + &schemas.BifrostError{ExtraFields: schemas.BifrostErrorExtraFields{ + BilledUsage: &schemas.BifrostLLMUsage{TotalTokens: 1}, + BillingAttemptStartedAt: &startedAt, + }}, + schemas.ChatCompletionRequest, + ) + + if entry.BillingAttemptStartedAt == nil || !entry.BillingAttemptStartedAt.Equal(startedAt) { + t.Fatalf("expected billing attempt start on error log, got %v", entry.BillingAttemptStartedAt) + } +} + // TestCalculateCostForLogPreservesServedTier pins the store-independent tier leak: the // Log table had no service_tier column and BifrostLLMUsage tags Speed/InferenceGeo // `json:"-"`, so recomputation priced every row at standard rates. On OpenAI that @@ -380,6 +403,36 @@ func TestStreamingServiceTierSurvivesAccumulatorHandoff(t *testing.T) { } } +// TestStreamingBillingAttemptStartSurvivesAccumulatorHandoff pins the stream-side +// billing timestamp: it rides on response ExtraFields into the accumulator, crosses +// tracer/result conversion, and must land on the log row rather than being reconstructed +// from completion time. +func TestStreamingBillingAttemptStartSurvivesAccumulatorHandoff(t *testing.T) { + startedAt := time.Date(2026, 8, 22, 22, 30, 0, 0, time.UTC) + + processed := convertToProcessedStreamResponse(&schemas.StreamAccumulatorResult{ + RequestID: "req-stream-billing-time", + RequestedModel: "deepseek-chat", + ResolvedModel: "deepseek-chat", + Provider: schemas.OpenAI, + Status: "success", + TokenUsage: &schemas.BifrostLLMUsage{TotalTokens: 1}, + BillingAttemptStartedAt: &startedAt, + }, schemas.ChatCompletionStreamRequest) + if processed == nil || processed.Data == nil { + t.Fatal("expected a processed stream response with data") + } + if processed.Data.BillingAttemptStartedAt == nil || !processed.Data.BillingAttemptStartedAt.Equal(startedAt) { + t.Fatalf("expected billing attempt time to survive conversion, got %v", processed.Data.BillingAttemptStartedAt) + } + + entry := &logstore.Log{} + (&LoggerPlugin{}).applyStreamingOutputToEntry(entry, processed, false, false) + if entry.BillingAttemptStartedAt == nil || !entry.BillingAttemptStartedAt.Equal(startedAt) { + t.Fatalf("expected billing_attempt_started_at on log entry, got %v", entry.BillingAttemptStartedAt) + } +} + // TestCalculateCostForLogPricesOneHourCacheWrites covers the Responses-path leak: // CachedWriteTokenDetails was dropped in the conversion, so 1h cache writes billed at // the cheaper 5m rate. diff --git a/plugins/logging/main.go b/plugins/logging/main.go index f7d31949bcc..140e40c62a7 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -221,6 +221,15 @@ func (p *LoggerPlugin) contentLoggingEnabled(ctx *schemas.BifrostContext) bool { return p.resolveContentPolicy(ctx).storeContent } +// cloneTime returns an owned copy of a timestamp before storing it on a log entry. +func cloneTime(value *time.Time) *time.Time { + if value == nil { + return nil + } + copy := *value + return © +} + // applyMCPGovernanceFieldsToEntry stamps MCP log ownership from the request context. func applyMCPGovernanceFieldsToEntry(ctx *schemas.BifrostContext, entry *logstore.MCPToolLog) { if ctx == nil || entry == nil { @@ -251,7 +260,14 @@ func applyMCPGovernanceFieldsToEntry(ctx *schemas.BifrostContext, entry *logstor // when stream accumulation didn't already capture it, but cost is (re)computed // whenever it is still missing - independent of whether tokens were already // parsed, since a streaming error can populate usage without a cost. -func (p *LoggerPlugin) applyErrorBillingFromBilledUsage(ctx *schemas.BifrostContext, entry *logstore.Log, billed *schemas.BifrostLLMUsage, requestType schemas.RequestType) { +func (p *LoggerPlugin) applyErrorBillingFromBilledUsage(ctx *schemas.BifrostContext, entry *logstore.Log, bifrostErr *schemas.BifrostError, requestType schemas.RequestType) { + billed := (*schemas.BifrostLLMUsage)(nil) + if bifrostErr != nil { + billed = bifrostErr.ExtraFields.BilledUsage + if bifrostErr.ExtraFields.BillingAttemptStartedAt != nil { + entry.BillingAttemptStartedAt = cloneTime(bifrostErr.ExtraFields.BillingAttemptStartedAt) + } + } if billed == nil { return } @@ -1711,7 +1727,7 @@ func (p *LoggerPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *schemas. // processed tokens (carried on BilledUsage). Record cost + tokens so the // logs DB reflects what we were actually billed, mirroring the governance // budget. - p.applyErrorBillingFromBilledUsage(ctx, entry, bifrostErr.ExtraFields.BilledUsage, requestType) + p.applyErrorBillingFromBilledUsage(ctx, entry, bifrostErr, requestType) p.applyInternalCallCosts(ctx, entry, guardrailDebug) applyLargePayloadPreviewsToEntry(ctx, entry, contentLoggingEnabled) p.storeOrEnqueueEntry(ctx, entry, p.makePostWriteCallback(nil)) @@ -1768,7 +1784,7 @@ func (p *LoggerPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *schemas. } // A stream error can arrive with a response chunk, bypassing Path A. // Preserve provider-billed usage and the sidecar calls in that case. - p.applyErrorBillingFromBilledUsage(ctx, entry, bifrostErr.ExtraFields.BilledUsage, requestType) + p.applyErrorBillingFromBilledUsage(ctx, entry, bifrostErr, requestType) p.applyInternalCallCosts(ctx, entry, guardrailDebug) } else if streamResponse == nil { // tracer or traceID not available, or accumulator returned nil - still write what we have diff --git a/plugins/logging/operations.go b/plugins/logging/operations.go index 46cc4a074ed..30eb146c552 100644 --- a/plugins/logging/operations.go +++ b/plugins/logging/operations.go @@ -423,6 +423,9 @@ func (p *LoggerPlugin) applyStreamingOutputToEntry(entry *logstore.Log, streamRe if streamResponse.Data.ServiceTier != nil { entry.ServiceTier = new(string(*streamResponse.Data.ServiceTier)) } + if streamResponse.Data.BillingAttemptStartedAt != nil { + entry.BillingAttemptStartedAt = cloneTime(streamResponse.Data.BillingAttemptStartedAt) + } // Speed/InferenceGeo come off the usage struct (the provider sets them there for // exactly this reason). ServiceTier is accumulated separately from the streamed // response envelope above. @@ -629,6 +632,9 @@ func (p *LoggerPlugin) applyNonStreamingOutputToEntry(entry *logstore.Log, resul entry.TotalTokens = usage.TotalTokens } applyServedTierToEntry(entry, result, usage) + if result.GetExtraFields().BillingAttemptStartedAt != nil { + entry.BillingAttemptStartedAt = cloneTime(result.GetExtraFields().BillingAttemptStartedAt) + } // Extract raw request/response and output content extraFields := result.GetExtraFields() @@ -2095,12 +2101,13 @@ func (p *LoggerPlugin) calculateCostBreakdownForLog(logEntry *logstore.Log) (*sc } extraFields := schemas.BifrostResponseExtraFields{ - RequestType: requestType, - Provider: schemas.ModelProvider(logEntry.Provider), - OriginalModelRequested: originalModelRequested, - ResolvedModelUsed: logEntry.Model, - CacheDebug: cacheDebug, - GuardrailDebug: guardrailDebug, + RequestType: requestType, + BillingAttemptStartedAt: cloneTime(logEntry.BillingAttemptStartedAt), + Provider: schemas.ModelProvider(logEntry.Provider), + OriginalModelRequested: originalModelRequested, + ResolvedModelUsed: logEntry.Model, + CacheDebug: cacheDebug, + GuardrailDebug: guardrailDebug, RoutingInfo: schemas.RoutingInfo{ Provider: schemas.ModelProvider(logEntry.Provider), Model: originalModelRequested, diff --git a/plugins/logging/operations_test.go b/plugins/logging/operations_test.go index 99340fbc6ff..10b8f9fd6a5 100644 --- a/plugins/logging/operations_test.go +++ b/plugins/logging/operations_test.go @@ -907,7 +907,9 @@ func TestApplyErrorBillingFromBilledUsage_ComputesCostWhenTokensAlreadyParsed(t billed := entry.TokenUsageParsed ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - plugin.applyErrorBillingFromBilledUsage(ctx, entry, billed, schemas.ChatCompletionStreamRequest) + plugin.applyErrorBillingFromBilledUsage(ctx, entry, &schemas.BifrostError{ + ExtraFields: schemas.BifrostErrorExtraFields{BilledUsage: billed}, + }, schemas.ChatCompletionStreamRequest) if entry.Cost == nil { t.Fatal("expected cost to be computed even though token usage was already parsed") @@ -997,7 +999,9 @@ func TestApplyErrorBillingFromBilledUsage_FillsTokensAndCostWhenUnparsed(t *test } ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - plugin.applyErrorBillingFromBilledUsage(ctx, entry, billed, schemas.ChatCompletionStreamRequest) + plugin.applyErrorBillingFromBilledUsage(ctx, entry, &schemas.BifrostError{ + ExtraFields: schemas.BifrostErrorExtraFields{BilledUsage: billed}, + }, schemas.ChatCompletionStreamRequest) if entry.TokenUsageParsed == nil || entry.TotalTokens != promptTokens+completionTokens { t.Fatalf("expected tokens backfilled, got %+v", entry.TokenUsageParsed) diff --git a/plugins/logging/utils.go b/plugins/logging/utils.go index 5a1c3ff664a..5440c64d4cf 100644 --- a/plugins/logging/utils.go +++ b/plugins/logging/utils.go @@ -819,26 +819,27 @@ func convertToProcessedStreamResponse(result *schemas.StreamAccumulatorResult, r // Build accumulated data data := &streaming.AccumulatedData{ - RequestID: result.RequestID, - Model: result.RequestedModel, - Status: result.Status, - Stream: true, - Latency: result.Latency, - TimeToFirstToken: result.TimeToFirstToken, - OutputMessage: result.OutputMessage, - OutputMessages: result.OutputMessages, - ErrorDetails: result.ErrorDetails, - TokenUsage: result.TokenUsage, - ServiceTier: result.ServiceTier, - CacheDebug: result.CacheDebug, - GuardrailDebug: result.GuardrailDebug, - Cost: result.Cost, - AudioOutput: result.AudioOutput, - TranscriptionOutput: result.TranscriptionOutput, - ImageGenerationOutput: result.ImageGenerationOutput, - PassthroughOutput: result.PassthroughOutput, - FinishReason: result.FinishReason, - RawResponse: result.RawResponse, + RequestID: result.RequestID, + Model: result.RequestedModel, + Status: result.Status, + Stream: true, + Latency: result.Latency, + TimeToFirstToken: result.TimeToFirstToken, + OutputMessage: result.OutputMessage, + OutputMessages: result.OutputMessages, + ErrorDetails: result.ErrorDetails, + TokenUsage: result.TokenUsage, + BillingAttemptStartedAt: result.BillingAttemptStartedAt, + ServiceTier: result.ServiceTier, + CacheDebug: result.CacheDebug, + GuardrailDebug: result.GuardrailDebug, + Cost: result.Cost, + AudioOutput: result.AudioOutput, + TranscriptionOutput: result.TranscriptionOutput, + ImageGenerationOutput: result.ImageGenerationOutput, + PassthroughOutput: result.PassthroughOutput, + FinishReason: result.FinishReason, + RawResponse: result.RawResponse, } // Handle tool calls if present