From fd22dae853aed233bb987eef0958bca06c2eca03 Mon Sep 17 00:00:00 2001 From: tejas ghatte Date: Mon, 25 May 2026 14:14:42 +0530 Subject: [PATCH] fix: passthrough budgets --- core/bifrost.go | 3 + core/providers/anthropic/anthropic.go | 98 ++--- core/providers/anthropic/passthrough_usage.go | 188 ++++++++ .../anthropic/passthrough_usage_test.go | 167 ++++++++ core/providers/azure/azure.go | 106 ++--- .../providers/azure/passthrough_usage_test.go | 54 +++ core/providers/gemini/gemini.go | 106 ++--- core/providers/gemini/passthrough_usage.go | 364 ++++++++++++++++ .../gemini/passthrough_usage_test.go | 186 ++++++++ core/providers/openai/openai.go | 91 ++-- core/providers/openai/passthrough_usage.go | 400 ++++++++++++++++++ .../openai/passthrough_usage_test.go | 292 +++++++++++++ core/providers/utils/passthrough_stream.go | 196 +++++++++ core/providers/utils/utils.go | 1 - core/providers/vertex/vertex.go | 112 ++--- core/schemas/bifrost.go | 1 + core/schemas/passthrough.go | 40 +- framework/modelcatalog/pricing.go | 145 ++++++- framework/streaming/passthrough.go | 27 +- framework/streaming/types.go | 7 +- plugins/governance/main.go | 7 +- plugins/logging/main.go | 8 + plugins/logging/operations.go | 4 + .../bifrost-http/integrations/router.go | 20 +- 24 files changed, 2243 insertions(+), 380 deletions(-) create mode 100644 core/providers/anthropic/passthrough_usage.go create mode 100644 core/providers/anthropic/passthrough_usage_test.go create mode 100644 core/providers/azure/passthrough_usage_test.go create mode 100644 core/providers/gemini/passthrough_usage.go create mode 100644 core/providers/gemini/passthrough_usage_test.go create mode 100644 core/providers/openai/passthrough_usage.go create mode 100644 core/providers/openai/passthrough_usage_test.go create mode 100644 core/providers/utils/passthrough_stream.go diff --git a/core/bifrost.go b/core/bifrost.go index a5d839cbc70..c9b2f04dbd0 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -6448,6 +6448,9 @@ func (bifrost *Bifrost) handleProviderRequest(provider schemas.Provider, config if bifrostError != nil { return nil, bifrostError } + if passthroughResponse != nil { + passthroughResponse.Path = req.BifrostRequest.PassthroughRequest.Path + } response.PassthroughResponse = passthroughResponse default: _, model, _ := req.BifrostRequest.GetRequestFields() diff --git a/core/providers/anthropic/anthropic.go b/core/providers/anthropic/anthropic.go index 21bf81db78d..c3dde7bc3ac 100644 --- a/core/providers/anthropic/anthropic.go +++ b/core/providers/anthropic/anthropic.go @@ -2605,6 +2605,11 @@ func (provider *AnthropicProvider) Passthrough( return nil, providerUtils.NewBifrostOperationError("failed to decode response body", err) } + var passthroughUsage *schemas.BifrostPassthroughUsage + if resp.StatusCode() >= 200 && resp.StatusCode() < 300 { + passthroughUsage = ExtractAnthropicPassthroughUsage(req.Path, req.Body, body) + } + bifrostResponse := &schemas.BifrostPassthroughResponse{ StatusCode: resp.StatusCode(), Headers: headers, @@ -2612,7 +2617,9 @@ func (provider *AnthropicProvider) Passthrough( ExtraFields: schemas.BifrostResponseExtraFields{ Latency: latency.Milliseconds(), ProviderResponseHeaders: headers, + PassthroughPath: req.Path, }, + PassthroughUsage: passthroughUsage, } return bifrostResponse, nil @@ -2690,71 +2697,32 @@ func (provider *AnthropicProvider) PassthroughStream( ) } - // Wrap reader with idle timeout to detect stalled streams. providerUtils.SetStreamIdleTimeoutIfEmpty(ctx, provider.networkConfig.StreamIdleTimeoutInSeconds) - rawBodyStream := bodyStream - bodyStream, stopIdleTimeout := providerUtils.NewIdleTimeoutReader(bodyStream, rawBodyStream, providerUtils.GetStreamIdleTimeout(ctx), ctx) - - // Cancellation must close the raw stream to unblock reads. - stopCancellation := providerUtils.SetupStreamCancellation(ctx, rawBodyStream, provider.logger) - - extraFields := schemas.BifrostResponseExtraFields{ - ProviderResponseHeaders: headers, - } - statusCode := resp.StatusCode() - - ch := make(chan *schemas.BifrostStreamChunk, schemas.DefaultStreamBufferSize) - go func() { - defer providerUtils.EnsureStreamFinalizerCalled(ctx, postHookSpanFinalizer) - defer func() { - if ctx.Err() == context.Canceled { - providerUtils.HandleStreamCancellation(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, req.Body) - } else if ctx.Err() == context.DeadlineExceeded { - providerUtils.HandleStreamTimeout(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, req.Body) - } - close(ch) - }() - defer providerUtils.ReleaseStreamingResponse(ctx, resp) - defer stopIdleTimeout() - defer stopCancellation() - - buf := make([]byte, 4096) - for { - n, readErr := bodyStream.Read(buf) - if n > 0 { - chunk := make([]byte, n) - copy(chunk, buf[:n]) - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - Body: chunk, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - } - if readErr == io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - return - } - if readErr != nil { - if ctx.Err() != nil { - return // let defer handle cancel/timeout + strippedPath := req.Path + if idx := strings.IndexByte(strippedPath, '?'); idx >= 0 { + strippedPath = strippedPath[:idx] + } + var messagesUsage *AnthropicPassthroughStreamUsage + if strings.HasSuffix(strippedPath, "/messages") { + messagesUsage = &AnthropicPassthroughStreamUsage{} + } + return providerUtils.StreamPassthrough( + ctx, postHookRunner, postHookSpanFinalizer, resp, bodyStream, + providerUtils.PassthroughStreamParams{ + StatusCode: resp.StatusCode(), + Headers: headers, + Path: req.Path, + RawRequest: req.Body, + CancellationBody: req.Body, + StartTime: startTime, + Logger: provider.logger, + HasUsage: HasAnthropicPassthroughUsage, + Observe: func(event []byte) *schemas.BifrostPassthroughUsage { + if messagesUsage != nil { + return messagesUsage.ObserveEvent(event) } - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendError(ctx, postHookRunner, readErr, ch, provider.logger, postHookSpanFinalizer) - return - } - } - }() - return ch, nil + return ExtractAnthropicPassthroughUsage(req.Path, req.Body, event) + }, + }, + ), nil } diff --git a/core/providers/anthropic/passthrough_usage.go b/core/providers/anthropic/passthrough_usage.go new file mode 100644 index 00000000000..1996e4521e2 --- /dev/null +++ b/core/providers/anthropic/passthrough_usage.go @@ -0,0 +1,188 @@ +package anthropic + +import ( + "strings" + + "github.com/bytedance/sonic" + providerUtils "github.com/maximhq/bifrost/core/providers/utils" + "github.com/maximhq/bifrost/core/schemas" +) + +// ExtractAnthropicPassthroughUsage extracts usage from a passthrough response payload. path is +// the stripped request path; body is a single SSE data event (streaming) or the full response +// body (non-streaming). Streaming /messages usage is assembled per-event by +// AnthropicPassthroughStreamUsage, so here /messages only ever sees a plain JSON body. +func ExtractAnthropicPassthroughUsage(path string, _, body []byte) *schemas.BifrostPassthroughUsage { + if idx := strings.IndexByte(path, '?'); idx >= 0 { + path = path[:idx] + } + + switch { + case strings.HasSuffix(path, "/messages"): + return extractAnthropicMessagesUsage(body) + case strings.HasSuffix(path, "/complete"): + return extractAnthropicCompleteUsage(body) + } + return nil +} + +func HasAnthropicPassthroughUsage(event []byte) bool { + return providerUtils.GetJSONField(event, "usage").Exists() || + providerUtils.GetJSONField(event, "message.usage").Exists() +} + +// buildAnthropicPassthroughUsage converts AnthropicUsage directly into BifrostPassthroughUsage. +func buildAnthropicPassthroughUsage(au *AnthropicUsage) *schemas.BifrostPassthroughUsage { + if au == nil { + return nil + } + totalInput := au.InputTokens + au.CacheReadInputTokens + au.CacheCreationInputTokens + total := totalInput + au.OutputTokens + if total == 0 { + return nil + } + + usage := &schemas.BifrostLLMUsage{ + PromptTokens: totalInput, + CompletionTokens: au.OutputTokens, + TotalTokens: total, + } + + if au.CacheReadInputTokens > 0 || au.CacheCreationInputTokens > 0 { + details := &schemas.ChatPromptTokensDetails{ + CachedReadTokens: au.CacheReadInputTokens, + CachedWriteTokens: au.CacheCreationInputTokens, + } + if au.CacheCreation.Ephemeral5mInputTokens > 0 || au.CacheCreation.Ephemeral1hInputTokens > 0 { + details.CachedWriteTokenDetails = &schemas.ChatCachedWriteTokenDetails{ + CachedWriteTokens5m: au.CacheCreation.Ephemeral5mInputTokens, + CachedWriteTokens1h: au.CacheCreation.Ephemeral1hInputTokens, + } + } + usage.PromptTokensDetails = details + } + + if au.ServerToolUse != nil && au.ServerToolUse.WebSearchRequests > 0 { + n := au.ServerToolUse.WebSearchRequests + usage.CompletionTokensDetails = &schemas.ChatCompletionTokensDetails{ + NumSearchQueries: &n, + } + } + + u := &schemas.BifrostPassthroughUsage{LLMUsage: usage} + if au.ServiceTier != nil { + t := MapAnthropicServiceTierToBifrost(*au.ServiceTier) + u.ServiceTier = &t + } + return u +} + +// AnthropicPassthroughStreamUsage incrementally merges /v1/messages stream usage across events +// without retaining the response body. Anthropic splits usage: message_start nests it under +// message.usage (input, cache tokens incl. 5m/1h split, service_tier), while message_delta has +// it at the top level (final output). Taking the max of each field across events combines them +// order-independently — the same merge the native Anthropic stream does (anthropic.go). +type AnthropicPassthroughStreamUsage struct { + combined AnthropicUsage + seen bool +} + +// ObserveEvent merges one framed SSE data payload's usage into the running total and returns +// the running usage (nil until any usage-bearing event is seen). +func (a *AnthropicPassthroughStreamUsage) ObserveEvent(event []byte) *schemas.BifrostPassthroughUsage { + var evt AnthropicStreamEvent + if err := sonic.Unmarshal(event, &evt); err != nil { + return a.usage() + } + // message_delta carries usage at the top level; message_start nests it under message.usage. + var u *AnthropicUsage + if evt.Usage != nil { + u = evt.Usage + } else if evt.Message != nil && evt.Message.Usage != nil { + u = evt.Message.Usage + } + if u == nil { + return a.usage() + } + + a.seen = true + c := &a.combined + if u.InputTokens > c.InputTokens { + c.InputTokens = u.InputTokens + } + if u.OutputTokens > c.OutputTokens { + c.OutputTokens = u.OutputTokens + } + if u.CacheReadInputTokens > c.CacheReadInputTokens { + c.CacheReadInputTokens = u.CacheReadInputTokens + } + if u.CacheCreationInputTokens > c.CacheCreationInputTokens { + c.CacheCreationInputTokens = u.CacheCreationInputTokens + } + if u.CacheCreation.Ephemeral5mInputTokens > c.CacheCreation.Ephemeral5mInputTokens { + c.CacheCreation.Ephemeral5mInputTokens = u.CacheCreation.Ephemeral5mInputTokens + } + if u.CacheCreation.Ephemeral1hInputTokens > c.CacheCreation.Ephemeral1hInputTokens { + c.CacheCreation.Ephemeral1hInputTokens = u.CacheCreation.Ephemeral1hInputTokens + } + if u.ServerToolUse != nil { + if c.ServerToolUse == nil { + c.ServerToolUse = &AnthropicServerToolUseUsage{} + } + if u.ServerToolUse.WebSearchRequests > c.ServerToolUse.WebSearchRequests { + c.ServerToolUse.WebSearchRequests = u.ServerToolUse.WebSearchRequests + } + } + if u.ServiceTier != nil { + c.ServiceTier = u.ServiceTier + } + return a.usage() +} + +func (a *AnthropicPassthroughStreamUsage) usage() *schemas.BifrostPassthroughUsage { + if !a.seen { + return nil + } + return buildAnthropicPassthroughUsage(&a.combined) +} + +// extractAnthropicMessagesUsage parses usage from a /v1/messages response body. Streaming usage +// is assembled per-event by AnthropicPassthroughStreamUsage, so this only sees a plain JSON +// (non-streaming) body, which carries the full usage block at the top level. +func extractAnthropicMessagesUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + var resp AnthropicMessageResponse + if err := sonic.Unmarshal(body, &resp); err != nil || resp.Usage == nil { + return nil + } + return buildAnthropicPassthroughUsage(resp.Usage) +} + +// extractAnthropicCompleteUsage handles the legacy /v1/complete endpoint. +func extractAnthropicCompleteUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + var resp struct { + Usage *struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } `json:"usage"` + } + if err := sonic.Unmarshal(body, &resp); err != nil || resp.Usage == nil { + return nil + } + total := resp.Usage.InputTokens + resp.Usage.OutputTokens + if total == 0 { + return nil + } + return &schemas.BifrostPassthroughUsage{ + LLMUsage: &schemas.BifrostLLMUsage{ + PromptTokens: resp.Usage.InputTokens, + CompletionTokens: resp.Usage.OutputTokens, + TotalTokens: total, + }, + } +} diff --git a/core/providers/anthropic/passthrough_usage_test.go b/core/providers/anthropic/passthrough_usage_test.go new file mode 100644 index 00000000000..065cc33f86a --- /dev/null +++ b/core/providers/anthropic/passthrough_usage_test.go @@ -0,0 +1,167 @@ +package anthropic_test + +import ( + "testing" + + "github.com/maximhq/bifrost/core/providers/anthropic" + "github.com/maximhq/bifrost/core/schemas" +) + +func TestExtractAnthropicPassthroughUsage(t *testing.T) { + tests := []struct { + name string + path string + body string + check func(t *testing.T, u *schemas.BifrostPassthroughUsage) + }{ + { + name: "messages non-stream usage + service tier mapping", + path: "/v1/messages", + body: `{"usage":{"input_tokens":66,"output_tokens":26,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"service_tier":"standard"}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + anthropicMustLLM(t, u, 66, 26, 92) + // Anthropic "standard" normalizes to the neutral "default". + if u.ServiceTier == nil || *u.ServiceTier != schemas.BifrostServiceTierDefault { + t.Fatalf("service tier = %v, want default", u.ServiceTier) + } + }, + }, + { + name: "messages with cache 5m/1h breakdown", + path: "/v1/messages", + body: `{"usage":{"input_tokens":10,"output_tokens":5,"cache_read_input_tokens":3,"cache_creation_input_tokens":7,"cache_creation":{"ephemeral_5m_input_tokens":4,"ephemeral_1h_input_tokens":3}}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + // PromptTokens = input + cache_read + cache_creation = 10 + 3 + 7 = 20 + anthropicMustLLM(t, u, 20, 5, 25) + d := u.LLMUsage.PromptTokensDetails + if d == nil || d.CachedReadTokens != 3 || d.CachedWriteTokens != 7 { + t.Fatalf("cache tokens = %+v", d) + } + if d.CachedWriteTokenDetails == nil || d.CachedWriteTokenDetails.CachedWriteTokens1h != 3 || + d.CachedWriteTokenDetails.CachedWriteTokens5m != 4 { + t.Fatalf("5m/1h breakdown = %+v", d.CachedWriteTokenDetails) + } + }, + }, + { + name: "messages with web search server tool", + path: "/v1/messages", + body: `{"usage":{"input_tokens":10,"output_tokens":5,"server_tool_use":{"web_search_requests":2}}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil || u.LLMUsage.CompletionTokensDetails == nil || + u.LLMUsage.CompletionTokensDetails.NumSearchQueries == nil || + *u.LLMUsage.CompletionTokensDetails.NumSearchQueries != 2 { + t.Fatalf("web search requests = %+v, want 2", u) + } + }, + }, + { + name: "messages with priority tier", + path: "/v1/messages", + body: `{"usage":{"input_tokens":1,"output_tokens":1,"service_tier":"priority"}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ServiceTier == nil || *u.ServiceTier != schemas.BifrostServiceTierPriority { + t.Fatalf("service tier = %v, want priority", u.ServiceTier) + } + }, + }, + { + name: "legacy complete endpoint", + path: "/v1/complete", + body: `{"completion":"hi","usage":{"input_tokens":8,"output_tokens":4}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + anthropicMustLLM(t, u, 8, 4, 12) + }, + }, + { + name: "messages zero usage -> nil", + path: "/v1/messages", + body: `{"usage":{"input_tokens":0,"output_tokens":0}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u != nil { + t.Fatalf("expected nil, got %+v", u) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + u := anthropic.ExtractAnthropicPassthroughUsage(tt.path, nil, []byte(tt.body)) + tt.check(t, u) + }) + } +} + +// TestAnthropicPassthroughStreamUsage exercises the per-event max-merge accumulator used for +// streaming /messages: input/cache/tier come from message_start, final output from message_delta. +func TestAnthropicPassthroughStreamUsage(t *testing.T) { + t.Run("merges message_start + message_delta", func(t *testing.T) { + acc := &anthropic.AnthropicPassthroughStreamUsage{} + + // Non-usage event before any usage -> still nil. + if u := acc.ObserveEvent([]byte(`{"type":"content_block_start","index":0}`)); u != nil { + t.Fatalf("expected nil before any usage event, got %+v", u) + } + + // message_start: input + cache 5m/1h + service_tier (output is a placeholder here). + acc.ObserveEvent([]byte(`{"type":"message_start","message":{"usage":{"input_tokens":66,"cache_read_input_tokens":2,"cache_creation_input_tokens":5,"cache_creation":{"ephemeral_5m_input_tokens":2,"ephemeral_1h_input_tokens":3},"output_tokens":1,"service_tier":"standard"}}}`)) + // content delta carries no usage; must not disturb the running totals. + acc.ObserveEvent([]byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}`)) + // message_delta: final output tokens. + u := acc.ObserveEvent([]byte(`{"type":"message_delta","usage":{"output_tokens":26}}`)) + + // PromptTokens = input(66) + cache_read(2) + cache_creation(5) = 73; output = 26. + anthropicMustLLM(t, u, 73, 26, 99) + d := u.LLMUsage.PromptTokensDetails + if d == nil || d.CachedWriteTokenDetails == nil || d.CachedWriteTokenDetails.CachedWriteTokens1h != 3 { + t.Fatalf("1h cache split lost in merge: %+v", d) + } + // service_tier from message_start, normalized. + if u.ServiceTier == nil || *u.ServiceTier != schemas.BifrostServiceTierDefault { + t.Fatalf("service tier = %v, want default", u.ServiceTier) + } + }) + + t.Run("server tool use from message_delta", func(t *testing.T) { + acc := &anthropic.AnthropicPassthroughStreamUsage{} + acc.ObserveEvent([]byte(`{"type":"message_start","message":{"usage":{"input_tokens":10,"output_tokens":1}}}`)) + u := acc.ObserveEvent([]byte(`{"type":"message_delta","usage":{"output_tokens":4,"server_tool_use":{"web_search_requests":3}}}`)) + if u == nil || u.LLMUsage == nil || u.LLMUsage.CompletionTokensDetails == nil || + u.LLMUsage.CompletionTokensDetails.NumSearchQueries == nil || + *u.LLMUsage.CompletionTokensDetails.NumSearchQueries != 3 { + t.Fatalf("web search requests = %+v, want 3", u) + } + }) +} + +func TestHasAnthropicPassthroughUsage(t *testing.T) { + tests := []struct { + name string + event string + want bool + }{ + {"message_start (nested message.usage)", `{"type":"message_start","message":{"usage":{"input_tokens":1}}}`, true}, + {"message_delta (top-level usage)", `{"type":"message_delta","usage":{"output_tokens":1}}`, true}, + {"content_block_delta", `{"type":"content_block_delta","delta":{"text":"x"}}`, false}, + {"ping", `{"type":"ping"}`, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := anthropic.HasAnthropicPassthroughUsage([]byte(tt.event)); got != tt.want { + t.Fatalf("HasAnthropicPassthroughUsage = %v, want %v", got, tt.want) + } + }) + } +} + +func anthropicMustLLM(t *testing.T, u *schemas.BifrostPassthroughUsage, prompt, completion, total int) { + t.Helper() + if u == nil || u.LLMUsage == nil { + t.Fatalf("expected LLMUsage, got %+v", u) + } + if u.LLMUsage.PromptTokens != prompt || u.LLMUsage.CompletionTokens != completion || u.LLMUsage.TotalTokens != total { + t.Fatalf("LLMUsage = {prompt:%d completion:%d total:%d}, want {%d %d %d}", + u.LLMUsage.PromptTokens, u.LLMUsage.CompletionTokens, u.LLMUsage.TotalTokens, prompt, completion, total) + } +} diff --git a/core/providers/azure/azure.go b/core/providers/azure/azure.go index 57050e58ace..790dffad165 100644 --- a/core/providers/azure/azure.go +++ b/core/providers/azure/azure.go @@ -3546,6 +3546,11 @@ func (provider *AzureProvider) Passthrough( return nil, providerUtils.NewBifrostOperationError("failed to decode response body", err) } + var passthroughUsage *schemas.BifrostPassthroughUsage + if resp.StatusCode() >= 200 && resp.StatusCode() < 300 { + passthroughUsage = extractAzurePassthroughUsage(req.Method, req.Path, req.Body, body, req.Model) + } + bifrostResponse := &schemas.BifrostPassthroughResponse{ StatusCode: resp.StatusCode(), Headers: headers, @@ -3553,7 +3558,9 @@ func (provider *AzureProvider) Passthrough( ExtraFields: schemas.BifrostResponseExtraFields{ Latency: latency.Milliseconds(), ProviderResponseHeaders: headers, + PassthroughPath: req.Path, }, + PassthroughUsage: passthroughUsage, } return bifrostResponse, nil @@ -3628,68 +3635,34 @@ func (provider *AzureProvider) PassthroughStream( return nil, providerUtils.NewBifrostOperationError("provider returned an empty stream body", fmt.Errorf("provider returned an empty stream body")) } - bodyStream, stopIdleTimeout := providerUtils.NewIdleTimeoutReader(rawBodyStream, rawBodyStream, providerUtils.GetStreamIdleTimeout(ctx), ctx) - stopCancellation := providerUtils.SetupStreamCancellation(ctx, rawBodyStream, provider.logger) - - extraFields := schemas.BifrostResponseExtraFields{ - ProviderResponseHeaders: headers, - } - statusCode := resp.StatusCode() - - ch := make(chan *schemas.BifrostStreamChunk, schemas.DefaultStreamBufferSize) - go func() { - defer providerUtils.EnsureStreamFinalizerCalled(ctx, postHookSpanFinalizer) - defer func() { - if ctx.Err() == context.Canceled { - providerUtils.HandleStreamCancellation(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } else if ctx.Err() == context.DeadlineExceeded { - providerUtils.HandleStreamTimeout(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } - close(ch) - }() - defer providerUtils.ReleaseStreamingResponse(ctx, resp) - defer stopIdleTimeout() - defer stopCancellation() - - buf := make([]byte, 4096) - for { - n, readErr := bodyStream.Read(buf) - if n > 0 { - chunk := make([]byte, n) - copy(chunk, buf[:n]) - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - Body: chunk, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - } - if readErr == io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - return - } - if readErr != nil { - if ctx.Err() != nil { - return + var anthropicUsage *anthropic.AnthropicPassthroughStreamUsage + if schemas.IsAnthropicModel(req.Model) { + anthropicUsage = &anthropic.AnthropicPassthroughStreamUsage{} + } + return providerUtils.StreamPassthrough( + ctx, postHookRunner, postHookSpanFinalizer, resp, rawBodyStream, + providerUtils.PassthroughStreamParams{ + StatusCode: resp.StatusCode(), + Headers: headers, + Path: req.Path, + RawRequest: req.Body, + CancellationBody: providerUtils.PassthroughJSONBody(fasthttpReq, req.Body), + StartTime: startTime, + Logger: provider.logger, + HasUsage: func(event []byte) bool { + if anthropicUsage != nil { + return anthropic.HasAnthropicPassthroughUsage(event) } - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendError(ctx, postHookRunner, readErr, ch, provider.logger, postHookSpanFinalizer) - return - } - } - }() - return ch, nil + return openai.HasOpenAIPassthroughUsage(event) + }, + Observe: func(event []byte) *schemas.BifrostPassthroughUsage { + if anthropicUsage != nil { + return anthropicUsage.ObserveEvent(event) + } + return openai.ExtractOpenAIPassthroughUsage(req.Method, req.Path, req.Body, event) + }, + }, + ), nil } // buildPassthroughURL constructs the full Azure URL for a passthrough request. @@ -3729,3 +3702,14 @@ func (provider *AzureProvider) buildPassthroughURL(key schemas.Key, path, rawQue } return fullURL } + +// extractAzurePassthroughUsage dispatches usage extraction by the upstream API the +// passthrough request targets. Azure serves both OpenAI and Azure-hosted Anthropic models, +// so Anthropic routes (e.g. /messages) must use the Anthropic extractor — otherwise their +// usage is dropped and budgets/logging stay wrong. +func extractAzurePassthroughUsage(method, path string, reqBody, body []byte, model string) *schemas.BifrostPassthroughUsage { + if schemas.IsAnthropicModel(model) { + return anthropic.ExtractAnthropicPassthroughUsage(path, reqBody, body) + } + return openai.ExtractOpenAIPassthroughUsage(method, path, reqBody, body) +} diff --git a/core/providers/azure/passthrough_usage_test.go b/core/providers/azure/passthrough_usage_test.go new file mode 100644 index 00000000000..b7f0657725e --- /dev/null +++ b/core/providers/azure/passthrough_usage_test.go @@ -0,0 +1,54 @@ +package azure + +import ( + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +// TestExtractAzurePassthroughUsage verifies Azure dispatches usage extraction by upstream model: +// Anthropic models use the Anthropic extractor (its usage shape), everything else uses OpenAI's. +func TestExtractAzurePassthroughUsage(t *testing.T) { + tests := []struct { + name string + method string + path string + body string + model string + check func(t *testing.T, u *schemas.BifrostPassthroughUsage) + }{ + { + name: "anthropic model routes to anthropic extractor", + method: "POST", + path: "/v1/messages", + body: `{"usage":{"input_tokens":66,"output_tokens":26}}`, + model: "claude-sonnet-4-5", + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil || u.LLMUsage.PromptTokens != 66 || + u.LLMUsage.CompletionTokens != 26 || u.LLMUsage.TotalTokens != 92 { + t.Fatalf("anthropic dispatch usage = %+v", u) + } + }, + }, + { + name: "openai model routes to openai extractor", + method: "POST", + path: "/chat/completions", + body: `{"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`, + model: "gpt-4o", + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil || u.LLMUsage.PromptTokens != 10 || + u.LLMUsage.CompletionTokens != 5 || u.LLMUsage.TotalTokens != 15 { + t.Fatalf("openai dispatch usage = %+v", u) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + u := extractAzurePassthroughUsage(tt.method, tt.path, nil, []byte(tt.body), tt.model) + tt.check(t, u) + }) + } +} diff --git a/core/providers/gemini/gemini.go b/core/providers/gemini/gemini.go index bdab2575333..dcdea471d1a 100644 --- a/core/providers/gemini/gemini.go +++ b/core/providers/gemini/gemini.go @@ -4135,6 +4135,11 @@ func (provider *GeminiProvider) Passthrough( return nil, providerUtils.NewBifrostOperationError("failed to decode response body", err) } + var passthroughUsage *schemas.BifrostPassthroughUsage + if resp.StatusCode() >= 200 && resp.StatusCode() < 300 { + passthroughUsage = ExtractGeminiPassthroughUsage(req.Path, req.Body, body) + } + bifrostResponse := &schemas.BifrostPassthroughResponse{ StatusCode: resp.StatusCode(), Headers: headers, @@ -4142,7 +4147,9 @@ func (provider *GeminiProvider) Passthrough( ExtraFields: schemas.BifrostResponseExtraFields{ Latency: latency.Milliseconds(), ProviderResponseHeaders: headers, + PassthroughPath: req.Path, }, + PassthroughUsage: passthroughUsage, } return bifrostResponse, nil @@ -4220,87 +4227,22 @@ func (provider *GeminiProvider) PassthroughStream( ) } - // Wrap reader with idle timeout to detect stalled streams. providerUtils.SetStreamIdleTimeoutIfEmpty(ctx, provider.networkConfig.StreamIdleTimeoutInSeconds) - rawBodyStream := bodyStream - bodyStream, stopIdleTimeout := providerUtils.NewIdleTimeoutReader(bodyStream, rawBodyStream, providerUtils.GetStreamIdleTimeout(ctx), ctx) - - // Cancellation must close the raw stream to unblock reads. - stopCancellation := providerUtils.SetupStreamCancellation(ctx, rawBodyStream, provider.logger) - - extraFields := schemas.BifrostResponseExtraFields{ - ProviderResponseHeaders: headers, - } - statusCode := resp.StatusCode() - - ch := make(chan *schemas.BifrostStreamChunk, schemas.DefaultStreamBufferSize) - go func() { - defer providerUtils.EnsureStreamFinalizerCalled(ctx, postHookSpanFinalizer) - defer func() { - if ctx.Err() == context.Canceled { - providerUtils.HandleStreamCancellation(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } else if ctx.Err() == context.DeadlineExceeded { - providerUtils.HandleStreamTimeout(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } - close(ch) - }() - defer providerUtils.ReleaseStreamingResponse(ctx, resp) - defer stopIdleTimeout() - defer stopCancellation() - - terminalDetector := &providerUtils.StreamTerminalDetector{} - buf := make([]byte, 4096) - for { - n, readErr := bodyStream.Read(buf) - if n > 0 { - chunk := make([]byte, n) - copy(chunk, buf[:n]) - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - Body: chunk, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - - if terminalDetector.ObserveChunk(chunk) { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - return - } - } - if readErr == io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - return - } - if readErr != nil { - if ctx.Err() != nil { - return // let defer handle cancel/timeout - } - if readErr != io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendError(ctx, postHookRunner, readErr, ch, provider.logger, postHookSpanFinalizer) - } - return - } - } - }() - return ch, nil + return providerUtils.StreamPassthrough( + ctx, postHookRunner, postHookSpanFinalizer, resp, bodyStream, + providerUtils.PassthroughStreamParams{ + StatusCode: resp.StatusCode(), + Headers: headers, + Path: req.Path, + RawRequest: req.Body, + CancellationBody: providerUtils.PassthroughJSONBody(fasthttpReq, req.Body), + StartTime: startTime, + UseTerminalDetector: true, + Logger: provider.logger, + HasUsage: HasGeminiPassthroughUsage, + Observe: func(event []byte) *schemas.BifrostPassthroughUsage { + return ExtractGeminiPassthroughUsage(req.Path, req.Body, event) + }, + }, + ), nil } diff --git a/core/providers/gemini/passthrough_usage.go b/core/providers/gemini/passthrough_usage.go new file mode 100644 index 00000000000..75a209f3be7 --- /dev/null +++ b/core/providers/gemini/passthrough_usage.go @@ -0,0 +1,364 @@ +package gemini + +import ( + "strconv" + "strings" + + "github.com/bytedance/sonic" + providerUtils "github.com/maximhq/bifrost/core/providers/utils" + "github.com/maximhq/bifrost/core/schemas" +) + +// ExtractGeminiPassthroughUsage extracts usage from a completed Gemini/Vertex +// passthrough response. Handles both SSE streaming (last event) and plain JSON +// (non-streaming) for all billable endpoint types. +// +// :generateContent handles all modalities — text, speech, transcription, and non-Imagen +// image generation. Output modality detection routes to the correct BifrostPassthroughUsage +// shape so the pricing engine uses the appropriate cost function. +// +// :predict is Imagen priced per-image. :predictLongRunning is Veo priced per-second. +// /interactions paths use the Interactions API usage shape. +func ExtractGeminiPassthroughUsage(path string, reqBody, body []byte) *schemas.BifrostPassthroughUsage { + if idx := strings.IndexByte(path, '?'); idx >= 0 { + path = path[:idx] + } + + // Interactions API uses /interactions paths — no colon suffix. + if strings.Contains(path, "/interactions") { + return extractGeminiInteractionsUsage(body) + } + + colonIdx := strings.LastIndexByte(path, ':') + action := "" + if colonIdx >= 0 { + action = path[colonIdx+1:] + } + + switch action { + case "generateContent", "streamGenerateContent": + return extractGeminiGenerateContentUsage(body) + + case "embedContent", "batchEmbedContents": + return extractGeminiEmbeddingUsage(body) + + case "predict": + return extractGeminiPredictUsage(reqBody, body) + + case "predictLongRunning": + // Veo video generation — priced per second of video. + return extractGeminiVeoUsage(reqBody) + } + + // Unknown action (e.g. countTokens, generateVideos) — try usageMetadata as best-effort. + return extractGeminiGenerateContentUsage(body) +} + +func HasGeminiPassthroughUsage(event []byte) bool { + return providerUtils.GetJSONField(event, "usageMetadata").Exists() || + providerUtils.GetJSONField(event, "usage").Exists() || + providerUtils.GetJSONField(event, "interaction.usage").Exists() +} + +// ---- :generateContent / :streamGenerateContent ---- + +type geminiPassthroughResp struct { + UsageMetadata *GenerateContentResponseUsageMetadata `json:"usageMetadata"` +} + +// extractGeminiGenerateContentUsage routes to the correct BifrostPassthroughUsage shape +// based on output modality from the response's candidatesTokensDetails: +// +// - IMAGE tokens in output → ImageUsage + LLMUsage → ImageGenerationRequest → computeImageCost +// - AUDIO tokens in output → LLMUsage with CompletionTokensDetails.AudioTokens → ResponsesRequest → computeTextCost audio differential +// - TEXT / default → LLMUsage → ResponsesRequest → computeTextCost +// +// Gemini TTS bills by input tokens (not chars), so AUDIO output stays in the ResponsesRequest +// path where computeTextCost applies the OutputCostPerAudioToken rate differential. +func extractGeminiGenerateContentUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + + var resp geminiPassthroughResp + if err := sonic.Unmarshal(body, &resp); err != nil || resp.UsageMetadata == nil { + return nil + } + + // ConvertGeminiUsageMetadataToResponsesUsage handles thinking tokens, cached content, + // and per-modality breakdowns (text, audio, image) for all :generateContent request types. + ru := ConvertGeminiUsageMetadataToResponsesUsage(resp.UsageMetadata) + if ru == nil || ru.TotalTokens == 0 { + return nil + } + + // IMAGE output → ImageUsage routes to computeImageCost via ImageGenerationRequest. + // LLMUsage is also set so the logging plugin can display tokens in/out. + if ru.OutputTokensDetails != nil && ru.OutputTokensDetails.ImageTokens != nil && *ru.OutputTokensDetails.ImageTokens > 0 { + imageUsage := &schemas.ImageUsage{ + InputTokens: ru.InputTokens, + OutputTokens: ru.OutputTokens, + TotalTokens: ru.TotalTokens, + OutputTokensDetails: &schemas.ImageTokenDetails{ + ImageTokens: *ru.OutputTokensDetails.ImageTokens, + }, + } + if ru.InputTokensDetails != nil && (ru.InputTokensDetails.TextTokens > 0 || ru.InputTokensDetails.ImageTokens > 0) { + imageUsage.InputTokensDetails = &schemas.ImageTokenDetails{ + TextTokens: ru.InputTokensDetails.TextTokens, + ImageTokens: ru.InputTokensDetails.ImageTokens, + } + } + return &schemas.BifrostPassthroughUsage{ + ImageUsage: imageUsage, + LLMUsage: &schemas.BifrostLLMUsage{ + PromptTokens: ru.InputTokens, + CompletionTokens: ru.OutputTokens, + TotalTokens: ru.TotalTokens, + }, + } + } + + // TEXT / AUDIO / default → LLMUsage with full modality details for computeTextCost. + // For AUDIO output, CompletionTokensDetails.AudioTokens is set so computeTextCost applies + // the OutputCostPerAudioToken rate differential: cost = tokens * (audioRate - textRate). + usage := &schemas.BifrostLLMUsage{ + PromptTokens: ru.InputTokens, + CompletionTokens: ru.OutputTokens, + TotalTokens: ru.TotalTokens, + } + if ru.InputTokensDetails != nil { + usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ + CachedReadTokens: ru.InputTokensDetails.CachedReadTokens, + TextTokens: ru.InputTokensDetails.TextTokens, + AudioTokens: ru.InputTokensDetails.AudioTokens, + ImageTokens: ru.InputTokensDetails.ImageTokens, + } + } + if ru.OutputTokensDetails != nil { + usage.CompletionTokensDetails = &schemas.ChatCompletionTokensDetails{ + ReasoningTokens: ru.OutputTokensDetails.ReasoningTokens, + AudioTokens: ru.OutputTokensDetails.AudioTokens, + } + } + + return &schemas.BifrostPassthroughUsage{LLMUsage: usage} +} + +// ---- :embedContent / :batchEmbedContents ---- + +func extractGeminiEmbeddingUsage(body []byte) *schemas.BifrostPassthroughUsage { + // Embeddings are never streamed, so body is plain JSON. + if len(body) == 0 { + return nil + } + + var resp geminiPassthroughResp + if err := sonic.Unmarshal(body, &resp); err != nil || resp.UsageMetadata == nil { + return nil + } + + m := resp.UsageMetadata + total := int(m.TotalTokenCount) + prompt := int(m.PromptTokenCount) + if total == 0 && prompt == 0 { + return nil + } + if total == 0 { + total = prompt + } + return &schemas.BifrostPassthroughUsage{ + LLMUsage: &schemas.BifrostLLMUsage{ + PromptTokens: prompt, + TotalTokens: total, + }, + } +} + +// ---- /interactions (Interactions API) ---- +// Non-streaming: usage sits at top level. +// Streaming InteractionCompletedEvent: usage is nested under "interaction". +// Fields: total_input_tokens, total_output_tokens, total_thought_tokens (reasoning), +// total_cached_tokens. service_tier "standard" is the default and is not forwarded. + +type geminiInteractionsUsage struct { + TotalTokens int `json:"total_tokens"` + InputTokens int `json:"total_input_tokens"` + OutputTokens int `json:"total_output_tokens"` + ThoughtTokens int `json:"total_thought_tokens"` + CachedTokens int `json:"total_cached_tokens"` +} + +type geminiInteractionsWrapper struct { + Usage *geminiInteractionsUsage `json:"usage"` + ServiceTier *string `json:"service_tier"` + // Streaming InteractionCompletedEvent nests the completed interaction object. + Interaction *struct { + Usage *geminiInteractionsUsage `json:"usage"` + ServiceTier *string `json:"service_tier"` + } `json:"interaction"` +} + +func extractGeminiInteractionsUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + + var w geminiInteractionsWrapper + if err := sonic.Unmarshal(body, &w); err != nil { + return nil + } + + // Streaming takes priority: nested under "interaction" with a non-zero total. + u, tier := w.Usage, w.ServiceTier + if w.Interaction != nil && w.Interaction.Usage != nil && w.Interaction.Usage.TotalTokens > 0 { + u, tier = w.Interaction.Usage, w.Interaction.ServiceTier + } + if u == nil || u.TotalTokens == 0 { + return nil + } + + usage := &schemas.BifrostLLMUsage{ + PromptTokens: u.InputTokens, + CompletionTokens: u.OutputTokens + u.ThoughtTokens, + TotalTokens: u.TotalTokens, + } + if u.CachedTokens > 0 { + usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ + CachedReadTokens: u.CachedTokens, + } + } + if u.ThoughtTokens > 0 { + usage.CompletionTokensDetails = &schemas.ChatCompletionTokensDetails{ + ReasoningTokens: u.ThoughtTokens, + } + } + + result := &schemas.BifrostPassthroughUsage{LLMUsage: usage} + // "standard" is the default tier — only forward non-standard values. + if tier != nil && *tier != "" && *tier != "standard" { + t := schemas.BifrostServiceTier(*tier) + result.ServiceTier = &t + } + return result +} + +// ---- :predictLongRunning (Veo video generation) ---- + +func extractGeminiVeoUsage(reqBody []byte) *schemas.BifrostPassthroughUsage { + // Default matches the native Gemini/Vertex path (schemas.DefaultVideoDuration). + secs := 8 + if d, err := strconv.Atoi(schemas.DefaultVideoDuration); err == nil { + secs = d + } + if len(reqBody) > 0 { + if d := providerUtils.GetJSONField(reqBody, "parameters.durationSeconds"); d.Exists() && d.Int() > 0 { + secs = int(d.Int()) + } + } + return &schemas.BifrostPassthroughUsage{VideoSeconds: &secs} +} + +// ---- :predict dispatch (Vertex/Gemini prediction endpoint) ---- +// The :predict action is shared by embeddings and Imagen image generation, distinguished by +// the per-prediction structure: +// +// - embedding (text or multimodal): predictions[] carry an `embeddings` block (text, with +// statistics.token_count) or modality vectors (textEmbedding/imageEmbedding/videoEmbeddings, +// multimodal — no token count) → token usage +// - Imagen image gen: predictions[].bytesBase64Encoded → per-image count +// +// An embedding response carries only LLMUsage, so it resolves to EmbeddingRequest via +// detectPassthroughRequestType's :predict mapping. Imagen sets ImageUsage and is resolved by +// its usage shape before that fallback, so the fallback only ever classifies embeddings. +type geminiPredictResponse struct { + Predictions []struct { + Embeddings *struct { + Statistics *struct { + TokenCount int `json:"token_count"` + } `json:"statistics"` + } `json:"embeddings"` + // Multimodal embedding modality vectors — used only to recognize the response as an + // embedding (multimodal responses carry no token count to bill). + TextEmbedding []float64 `json:"textEmbedding"` + ImageEmbedding []float64 `json:"imageEmbedding"` + VideoEmbeddings []any `json:"videoEmbeddings"` + } `json:"predictions"` +} + +func extractGeminiPredictUsage(reqBody, body []byte) *schemas.BifrostPassthroughUsage { + if len(body) > 0 { + var resp geminiPredictResponse + if err := sonic.Unmarshal(body, &resp); err == nil { + if u := resp.embeddingUsage(); u != nil { + return u + } + } + } + // Default: Imagen image generation (priced per image). + return extractGeminiImagenUsage(reqBody, body) +} + +// embeddingUsage returns token usage when the :predict response is an embedding response +// (text or multimodal), else nil. It returns non-nil for any embedding response — even when no +// token count is present (multimodal) — so an embedding is billed as such rather than misrouted +// to the Imagen per-image path. +func (r *geminiPredictResponse) embeddingUsage() *schemas.BifrostPassthroughUsage { + total, isEmbedding := 0, false + for i := range r.Predictions { + p := r.Predictions[i] + if p.Embeddings != nil { + isEmbedding = true + if p.Embeddings.Statistics != nil { + total += p.Embeddings.Statistics.TokenCount + } + } + if len(p.TextEmbedding) > 0 || len(p.ImageEmbedding) > 0 || len(p.VideoEmbeddings) > 0 { + isEmbedding = true + } + } + if !isEmbedding { + return nil + } + return &schemas.BifrostPassthroughUsage{ + LLMUsage: &schemas.BifrostLLMUsage{PromptTokens: total, TotalTokens: total}, + } +} + +// ---- :predict (Imagen) ---- +// Imagen is priced per image. Extract count from predictions in the response, +// with the requested sampleCount from the request body as a fallback. + +func extractGeminiImagenUsage(reqBody, body []byte) *schemas.BifrostPassthroughUsage { + u := &schemas.BifrostPassthroughUsage{ + ImageUsage: &schemas.ImageUsage{}, + } + + // Request body: sampleCount is the requested number of images. + if len(reqBody) > 0 { + var req GeminiImagenRequest + if err := sonic.Unmarshal(reqBody, &req); err == nil && + req.Parameters.SampleCount != nil && *req.Parameters.SampleCount > 0 { + if u.ImageUsage.OutputTokensDetails == nil { + u.ImageUsage.OutputTokensDetails = &schemas.ImageTokenDetails{} + } + u.ImageUsage.OutputTokensDetails.NImages = *req.Parameters.SampleCount + } + } + + // Response body: actual delivered predictions (may be fewer than requested). + if len(body) > 0 { + var resp GeminiImagenResponse + if err := sonic.Unmarshal(body, &resp); err == nil && len(resp.Predictions) > 0 { + if u.ImageUsage.OutputTokensDetails == nil { + u.ImageUsage.OutputTokensDetails = &schemas.ImageTokenDetails{} + } + u.ImageUsage.OutputTokensDetails.NImages = len(resp.Predictions) + } + } + + if u.ImageUsage.OutputTokensDetails == nil || u.ImageUsage.OutputTokensDetails.NImages == 0 { + return nil + } + return u +} diff --git a/core/providers/gemini/passthrough_usage_test.go b/core/providers/gemini/passthrough_usage_test.go new file mode 100644 index 00000000000..c278861ae87 --- /dev/null +++ b/core/providers/gemini/passthrough_usage_test.go @@ -0,0 +1,186 @@ +package gemini_test + +import ( + "testing" + + "github.com/maximhq/bifrost/core/providers/gemini" + "github.com/maximhq/bifrost/core/schemas" +) + +func TestExtractGeminiPassthroughUsage(t *testing.T) { + tests := []struct { + name string + path string + reqBody string + body string + check func(t *testing.T, u *schemas.BifrostPassthroughUsage) + }{ + { + name: "generateContent text with thinking tokens", + path: "/v1beta/models/gemini-2.5-flash:generateContent", + body: `{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":30,"totalTokenCount":199,"thoughtsTokenCount":159,"promptTokensDetails":[{"modality":"TEXT","tokenCount":10}]}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + // thinking tokens fold into completion: 30 + 159 = 189 + geminiMustLLM(t, u, 10, 189, 199) + if u.LLMUsage.CompletionTokensDetails == nil || u.LLMUsage.CompletionTokensDetails.ReasoningTokens != 159 { + t.Fatalf("reasoning tokens = %+v, want 159", u.LLMUsage.CompletionTokensDetails) + } + }, + }, + { + name: "generateContent image output -> ImageUsage", + path: "/v1beta/models/gemini-2.0-flash:generateContent", + body: `{"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":200,"totalTokenCount":205,"candidatesTokensDetails":[{"modality":"IMAGE","tokenCount":200}]}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ImageUsage == nil || u.ImageUsage.OutputTokensDetails == nil || + u.ImageUsage.OutputTokensDetails.ImageTokens != 200 { + t.Fatalf("image usage = %+v", u) + } + }, + }, + { + name: "embedContent", + path: "/v1beta/models/text-embedding-004:embedContent", + body: `{"usageMetadata":{"promptTokenCount":7,"totalTokenCount":7}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil || u.LLMUsage.PromptTokens != 7 || u.LLMUsage.TotalTokens != 7 { + t.Fatalf("embedding usage = %+v", u) + } + }, + }, + { + name: "predict text embedding (token_count)", + path: "/v1/projects/p/locations/l/publishers/google/models/text-embedding-005:predict", + body: `{"predictions":[{"embeddings":{"statistics":{"token_count":6}}}]}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil || u.LLMUsage.PromptTokens != 6 || u.LLMUsage.TotalTokens != 6 { + t.Fatalf("text embedding usage = %+v", u) + } + if u.ImageUsage != nil { + t.Fatalf("text embedding must not be billed as image: %+v", u.ImageUsage) + } + }, + }, + { + name: "predict multimodal embedding -> embedding (not image)", + path: "/v1/projects/p/locations/l/publishers/google/models/multimodalembedding@001:predict", + body: `{"predictions":[{"textEmbedding":[0.1,0.2],"imageEmbedding":[0.3]}]}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil { + t.Fatalf("multimodal embedding usage = %+v", u) + } + if u.ImageUsage != nil { + t.Fatalf("multimodal embedding must not be billed as image: %+v", u.ImageUsage) + } + }, + }, + { + name: "predict imagen -> per-image count from predictions", + path: "/v1/projects/p/locations/l/publishers/google/models/imagen-3.0-generate-002:predict", + body: `{"predictions":[{"bytesBase64Encoded":"aaa"},{"bytesBase64Encoded":"bbb"}]}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ImageUsage == nil || u.ImageUsage.OutputTokensDetails == nil || + u.ImageUsage.OutputTokensDetails.NImages != 2 { + t.Fatalf("imagen NImages = %+v, want 2", u) + } + }, + }, + { + name: "predict imagen -> sampleCount fallback when no predictions", + path: "/v1/projects/p/locations/l/publishers/google/models/imagen-3.0-generate-002:predict", + reqBody: `{"parameters":{"sampleCount":3}}`, + body: `{}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ImageUsage == nil || u.ImageUsage.OutputTokensDetails == nil || + u.ImageUsage.OutputTokensDetails.NImages != 3 { + t.Fatalf("imagen sampleCount fallback NImages = %+v, want 3", u) + } + }, + }, + { + name: "predictLongRunning veo -> seconds from request", + path: "/v1beta/models/veo-3.1-generate-preview:predictLongRunning", + reqBody: `{"parameters":{"durationSeconds":6}}`, + body: `{"name":"models/veo/operations/abc"}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.VideoSeconds == nil || *u.VideoSeconds != 6 { + t.Fatalf("veo seconds = %+v, want 6", u) + } + }, + }, + { + name: "predictLongRunning veo -> default seconds", + path: "/v1beta/models/veo-3.1-generate-preview:predictLongRunning", + reqBody: `{}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.VideoSeconds == nil || *u.VideoSeconds != 8 { + t.Fatalf("veo default seconds = %+v, want 8", u) + } + }, + }, + { + name: "interactions top-level usage (standard tier dropped)", + path: "/v1beta/interactions", + body: `{"usage":{"total_tokens":104,"total_input_tokens":7,"total_output_tokens":27,"total_thought_tokens":70,"total_cached_tokens":0},"service_tier":"standard"}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + // output + thought folded into completion: 27 + 70 = 97 + geminiMustLLM(t, u, 7, 97, 104) + if u.LLMUsage.CompletionTokensDetails == nil || u.LLMUsage.CompletionTokensDetails.ReasoningTokens != 70 { + t.Fatalf("reasoning tokens = %+v, want 70", u.LLMUsage.CompletionTokensDetails) + } + if u.ServiceTier != nil { + t.Fatalf("standard tier should be dropped, got %v", *u.ServiceTier) + } + }, + }, + { + name: "interactions nested (streaming completed) + non-standard tier", + path: "/v1beta/interactions", + body: `{"interaction":{"usage":{"total_tokens":50,"total_input_tokens":10,"total_output_tokens":40},"service_tier":"priority"}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + geminiMustLLM(t, u, 10, 40, 50) + if u.ServiceTier == nil || *u.ServiceTier != "priority" { + t.Fatalf("service tier = %v, want priority", u.ServiceTier) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + u := gemini.ExtractGeminiPassthroughUsage(tt.path, []byte(tt.reqBody), []byte(tt.body)) + tt.check(t, u) + }) + } +} + +func TestHasGeminiPassthroughUsage(t *testing.T) { + tests := []struct { + name string + event string + want bool + }{ + {"usageMetadata", `{"usageMetadata":{"totalTokenCount":5}}`, true}, + {"interactions top-level usage", `{"usage":{"total_tokens":5}}`, true}, + {"interactions nested", `{"interaction":{"usage":{"total_tokens":5}}}`, true}, + {"content chunk (no usage)", `{"candidates":[{"content":{"parts":[{"text":"hi"}]}}]}`, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := gemini.HasGeminiPassthroughUsage([]byte(tt.event)); got != tt.want { + t.Fatalf("HasGeminiPassthroughUsage = %v, want %v", got, tt.want) + } + }) + } +} + +func geminiMustLLM(t *testing.T, u *schemas.BifrostPassthroughUsage, prompt, completion, total int) { + t.Helper() + if u == nil || u.LLMUsage == nil { + t.Fatalf("expected LLMUsage, got %+v", u) + } + if u.LLMUsage.PromptTokens != prompt || u.LLMUsage.CompletionTokens != completion || u.LLMUsage.TotalTokens != total { + t.Fatalf("LLMUsage = {prompt:%d completion:%d total:%d}, want {%d %d %d}", + u.LLMUsage.PromptTokens, u.LLMUsage.CompletionTokens, u.LLMUsage.TotalTokens, prompt, completion, total) + } +} diff --git a/core/providers/openai/openai.go b/core/providers/openai/openai.go index 04eb2c8260c..3be4c5f03d0 100644 --- a/core/providers/openai/openai.go +++ b/core/providers/openai/openai.go @@ -6974,6 +6974,11 @@ func (provider *OpenAIProvider) Passthrough( return nil, providerUtils.NewBifrostOperationError("failed to decode response body", err) } + var passthroughUsage *schemas.BifrostPassthroughUsage + if resp.StatusCode() >= 200 && resp.StatusCode() < 300 { + passthroughUsage = ExtractOpenAIPassthroughUsage(req.Method, req.Path, req.Body, body) + } + bifrostResponse := &schemas.BifrostPassthroughResponse{ StatusCode: resp.StatusCode(), Headers: headers, @@ -6981,7 +6986,9 @@ func (provider *OpenAIProvider) Passthrough( ExtraFields: schemas.BifrostResponseExtraFields{ Latency: latency.Milliseconds(), ProviderResponseHeaders: headers, + PassthroughPath: req.Path, }, + PassthroughUsage: passthroughUsage, } return bifrostResponse, nil @@ -7063,71 +7070,21 @@ func (provider *OpenAIProvider) PassthroughStream( fmt.Errorf("provider returned an empty stream body")) } - // Wrap reader with idle timeout to detect stalled streams. - bodyStream, stopIdleTimeout := providerUtils.NewIdleTimeoutReader(rawBodyStream, rawBodyStream, providerUtils.GetStreamIdleTimeout(ctx), ctx) - - // Cancellation must close the raw stream to unblock reads. - stopCancellation := providerUtils.SetupStreamCancellation(ctx, rawBodyStream, provider.logger) - - extraFields := schemas.BifrostResponseExtraFields{ - ProviderResponseHeaders: headers, - } - statusCode := resp.StatusCode() - - ch := make(chan *schemas.BifrostStreamChunk, schemas.DefaultStreamBufferSize) - go func() { - defer providerUtils.EnsureStreamFinalizerCalled(ctx, postHookSpanFinalizer) - defer func() { - if ctx.Err() == context.Canceled { - providerUtils.HandleStreamCancellation(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } else if ctx.Err() == context.DeadlineExceeded { - providerUtils.HandleStreamTimeout(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } - close(ch) - }() - defer providerUtils.ReleaseStreamingResponse(ctx, resp) - defer stopIdleTimeout() - defer stopCancellation() - - buf := make([]byte, 4096) - for { - n, readErr := bodyStream.Read(buf) - if n > 0 { - chunk := make([]byte, n) - copy(chunk, buf[:n]) - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - Body: chunk, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - } - if readErr == io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - return - } - if readErr != nil { - if ctx.Err() != nil { - return // let defer handle cancel/timeout - } - if readErr != io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(startTime).Milliseconds() - providerUtils.ProcessAndSendError(ctx, postHookRunner, readErr, ch, provider.logger, postHookSpanFinalizer) - } - return - } - } - }() - return ch, nil + // Forward raw chunks to the client and extract usage incrementally per SSE event — + return providerUtils.StreamPassthrough( + ctx, postHookRunner, postHookSpanFinalizer, resp, rawBodyStream, + providerUtils.PassthroughStreamParams{ + StatusCode: resp.StatusCode(), + Headers: headers, + Path: req.Path, + RawRequest: req.Body, + CancellationBody: providerUtils.PassthroughJSONBody(fasthttpReq, req.Body), + StartTime: startTime, + Logger: provider.logger, + HasUsage: HasOpenAIPassthroughUsage, + Observe: func(event []byte) *schemas.BifrostPassthroughUsage { + return ExtractOpenAIPassthroughUsage(req.Method, req.Path, req.Body, event) + }, + }, + ), nil } diff --git a/core/providers/openai/passthrough_usage.go b/core/providers/openai/passthrough_usage.go new file mode 100644 index 00000000000..38b4e28af7d --- /dev/null +++ b/core/providers/openai/passthrough_usage.go @@ -0,0 +1,400 @@ +package openai + +import ( + "bytes" + "mime/multipart" + "strconv" + "strings" + + "github.com/bytedance/sonic" + providerUtils "github.com/maximhq/bifrost/core/providers/utils" + "github.com/maximhq/bifrost/core/schemas" +) + +// ExtractOpenAIPassthroughUsage extracts usage from a passthrough response payload. method is the +// HTTP method (used to bill only generation routes); path is the stripped request path; reqBody is +// the original request body (needed for speech char count and image/video parameters); body is a +// single SSE data event (streaming) or the full response body (non-streaming). +func ExtractOpenAIPassthroughUsage(method, path string, reqBody, body []byte) *schemas.BifrostPassthroughUsage { + if idx := strings.IndexByte(path, '?'); idx >= 0 { + path = path[:idx] + } + + switch { + case strings.HasSuffix(path, "/chat/completions"), + strings.HasSuffix(path, "/completions"): + return extractOAIChatUsage(body) + + case strings.HasSuffix(path, "/responses"): + return extractOAIResponsesUsage(body) + + case strings.HasSuffix(path, "/embeddings"): + return extractOAIEmbeddingUsage(body) + + case strings.HasSuffix(path, "/audio/speech"): + return extractOAISpeechUsage(reqBody) + + case strings.HasSuffix(path, "/audio/transcriptions"), + strings.HasSuffix(path, "/audio/translations"): + return extractOAITranscriptionUsage(body) + + case strings.HasSuffix(path, "/images/generations"), + strings.HasSuffix(path, "/images/edits"), + strings.HasSuffix(path, "/images/variations"): + return extractOAIImageUsage(reqBody, body) + + case strings.Contains(path, "/video"): + if strings.EqualFold(method, "POST") { + return extractOAIVideoUsage(reqBody) + } + return nil + + case strings.HasSuffix(path, "/containers"): + // Collection path serves both create (POST, billable) and list (GET, free); + // extractOAIContainerUsage disambiguates by response shape. Retrieve/delete use + // /containers/{id} and never match this suffix. + return extractOAIContainerUsage(body) + } + + return nil +} + +func HasOpenAIPassthroughUsage(event []byte) bool { + return providerUtils.GetJSONField(event, "usage").Exists() || + providerUtils.GetJSONField(event, "response.usage").Exists() +} + +// ---- video generation ---- +const openAIVideoDefaultSeconds = 4 + +func extractOAIVideoUsage(reqBody []byte) *schemas.BifrostPassthroughUsage { + secs := openAIVideoDefaultSeconds + if len(reqBody) > 0 { + // JSON body: OpenAI documents `seconds` as a top-level request field. gjson .Float() + // handles both the numeric and string forms the API accepts. + if v := providerUtils.GetJSONField(reqBody, "seconds"); v.Exists() && v.Float() > 0 { + secs = int(v.Float()) + } else if form := parseMultipartFormValues(reqBody); form != nil { + // Multipart body (binary asset upload): `seconds` rides as a form field. + if v := firstFormValue(form, "seconds"); v != "" { + if f, parseErr := strconv.ParseFloat(v, 64); parseErr == nil && f > 0 { + secs = int(f) + } + } + } + } + return &schemas.BifrostPassthroughUsage{VideoSeconds: &secs} +} + +// parseMultipartFormValues sniffs the multipart boundary from the first line of body and returns +// the parsed form values, or nil when body is not multipart/form-data (e.g. JSON). OpenAI sends +// /v1/images/{edits,variations} and binary-asset video requests as multipart; their scalar +// params (seconds, size, quality, n) ride along as form fields. +func parseMultipartFormValues(body []byte) map[string][]string { + if len(body) == 0 { + return nil + } + firstLine, _, _ := bytes.Cut(body, []byte("\n")) + boundary := strings.TrimRight(strings.TrimPrefix(string(firstLine), "--"), "\r") + if boundary == "" { + return nil + } + mr := multipart.NewReader(bytes.NewReader(body), boundary) + form, err := mr.ReadForm(32 << 20) + if err != nil { + return nil + } + // ReadForm spills parts over maxMemory to temp files; clean them up. form.Value is held in + // memory and stays valid after RemoveAll (which only purges spilled file parts). + defer form.RemoveAll() + return form.Value +} + +func firstFormValue(form map[string][]string, key string) string { + if v := form[key]; len(v) > 0 { + return v[0] + } + return "" +} + +// ---- chat / text completions ---- +// BifrostLLMUsage is OpenAI-compatible so we can unmarshal directly. + +type oaiChatUsageWrapper struct { + Usage *schemas.BifrostLLMUsage `json:"usage"` + ServiceTier *string `json:"service_tier"` +} + +func extractOAIChatUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + var w oaiChatUsageWrapper + if err := sonic.Unmarshal(body, &w); err != nil || w.Usage == nil || w.Usage.TotalTokens == 0 { + return nil + } + u := &schemas.BifrostPassthroughUsage{LLMUsage: w.Usage} + if w.ServiceTier != nil { + t := schemas.BifrostServiceTier(*w.ServiceTier) + u.ServiceTier = &t + } + return u +} + +// ---- responses API ---- +// A single wrapper handles both response formats in one unmarshal pass: +// - streaming: "response.completed" event nests usage under "response" +// - non-streaming: usage sits at the top level +type oaiResponsesWrapper struct { + Response *struct { + Usage *schemas.ResponsesResponseUsage `json:"usage"` + ServiceTier *string `json:"service_tier"` + } `json:"response"` + Usage *schemas.ResponsesResponseUsage `json:"usage"` + ServiceTier *string `json:"service_tier"` +} + +func extractOAIResponsesUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + + var w oaiResponsesWrapper + if err := sonic.Unmarshal(body, &w); err != nil { + return nil + } + + // Streaming takes priority: nested under "response" with a non-zero total. + ru, tier := w.Usage, w.ServiceTier + if w.Response != nil && w.Response.Usage != nil && w.Response.Usage.TotalTokens > 0 { + ru, tier = w.Response.Usage, w.Response.ServiceTier + } + if ru == nil || ru.TotalTokens == 0 { + return nil + } + return buildOAIResponsesUsage(ru, tier) +} + +func buildOAIResponsesUsage(ru *schemas.ResponsesResponseUsage, serviceTier *string) *schemas.BifrostPassthroughUsage { + usage := &schemas.BifrostLLMUsage{ + PromptTokens: ru.InputTokens, + CompletionTokens: ru.OutputTokens, + TotalTokens: ru.TotalTokens, + } + if ru.InputTokensDetails != nil { + usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ + CachedReadTokens: ru.InputTokensDetails.CachedReadTokens, + CachedWriteTokens: ru.InputTokensDetails.CachedWriteTokens, + } + } + if ru.OutputTokensDetails != nil { + usage.CompletionTokensDetails = &schemas.ChatCompletionTokensDetails{ + ReasoningTokens: ru.OutputTokensDetails.ReasoningTokens, + } + if ru.OutputTokensDetails.NumSearchQueries != nil { + usage.CompletionTokensDetails.NumSearchQueries = ru.OutputTokensDetails.NumSearchQueries + } + } + u := &schemas.BifrostPassthroughUsage{LLMUsage: usage} + if serviceTier != nil { + t := schemas.BifrostServiceTier(*serviceTier) + u.ServiceTier = &t + } + return u +} + +// ---- embeddings ---- +// Embeddings are not typically streamed; body is plain JSON. + +func extractOAIEmbeddingUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + var w oaiChatUsageWrapper + if err := sonic.Unmarshal(body, &w); err != nil || w.Usage == nil || w.Usage.TotalTokens == 0 { + return nil + } + return &schemas.BifrostPassthroughUsage{LLMUsage: w.Usage} +} + +// ---- speech (TTS) ---- +// Response is binary audio; pricing is based on input character count from the request. + +func extractOAISpeechUsage(reqBody []byte) *schemas.BifrostPassthroughUsage { + if len(reqBody) == 0 { + return nil + } + var req OpenAISpeechRequest + if err := sonic.Unmarshal(reqBody, &req); err != nil || req.Input == "" { + return nil + } + return &schemas.BifrostPassthroughUsage{ + AudioInputChars: len([]rune(req.Input)), + } +} + +// ---- transcription / translation ---- + +type oaiTranscriptionResponseWrapper struct { + Usage *schemas.TranscriptionUsage `json:"usage"` + Duration float64 `json:"duration"` // seconds fallback for older models +} + +func extractOAITranscriptionUsage(body []byte) *schemas.BifrostPassthroughUsage { + var r oaiTranscriptionResponseWrapper + if err := sonic.Unmarshal(body, &r); err != nil { + return nil + } + u := &schemas.BifrostPassthroughUsage{} + if r.Usage != nil && r.Usage.TotalTokens != nil && *r.Usage.TotalTokens > 0 { + promptTokens := 0 + if r.Usage.InputTokens != nil { + promptTokens = *r.Usage.InputTokens + } + u.LLMUsage = &schemas.BifrostLLMUsage{ + PromptTokens: promptTokens, + TotalTokens: *r.Usage.TotalTokens, + } + if r.Usage.InputTokenDetails != nil { + u.AudioTokenDetails = &schemas.TranscriptionUsageInputTokenDetails{ + AudioTokens: r.Usage.InputTokenDetails.AudioTokens, + TextTokens: r.Usage.InputTokenDetails.TextTokens, + } + } + u.AudioSeconds = r.Usage.Seconds + } else if r.Duration > 0 { + secs := int(r.Duration) + u.AudioSeconds = &secs + } + if u.LLMUsage == nil && u.AudioSeconds == nil { + return nil + } + return u +} + +// ---- image generation / edit / variation ---- +// Size, Quality, N come from the request body; usage/data count from the response. + +func extractOAIImageUsage(reqBody, body []byte) *schemas.BifrostPassthroughUsage { + u := &schemas.BifrostPassthroughUsage{} + + // Request body: size, quality, n. /v1/images/{edits,variations} are sent as + // multipart/form-data (binary image upload) with these as form fields; /v1/images/generations + // (and JSON-mode edits) carry them as a JSON OpenAIImageGenerationRequest. + if len(reqBody) > 0 { + var size, quality string + var n int + if form := parseMultipartFormValues(reqBody); form != nil { + size = firstFormValue(form, "size") + quality = firstFormValue(form, "quality") + if v := firstFormValue(form, "n"); v != "" { + if parsed, err := strconv.Atoi(v); err == nil { + n = parsed + } + } + } else { + var req OpenAIImageGenerationRequest + if err := sonic.Unmarshal(reqBody, &req); err == nil { + if req.Size != nil { + size = *req.Size + } + if req.Quality != nil { + quality = *req.Quality + } + if req.N != nil { + n = *req.N + } + } + } + if size != "" { + u.ImageSize = size + } + if quality != "" { + u.ImageQuality = quality + } + if n > 0 { + if u.ImageUsage == nil { + u.ImageUsage = &schemas.ImageUsage{} + } + if u.ImageUsage.OutputTokensDetails == nil { + u.ImageUsage.OutputTokensDetails = &schemas.ImageTokenDetails{} + } + u.ImageUsage.OutputTokensDetails.NImages = n + } + } + + // Response body: use OpenAIImageStreamResponse (streaming SSE event) or fall back to + // plain JSON for non-streaming passthrough routes. + if len(body) > 0 { + var resp OpenAIImageStreamResponse + if err := sonic.Unmarshal(body, &resp); err == nil { + if resp.Usage != nil { + u.ImageUsage = resp.Usage + } + if resp.Size != "" && u.ImageSize == "" { + u.ImageSize = resp.Size + } + if resp.Quality != "" && u.ImageQuality == "" { + u.ImageQuality = resp.Quality + } + } + // Mirror the native path (populateOutputImageCount): count delivered images from + // the `data` array when the request didn't specify n and no token usage was + // returned (e.g. DALL·E, which has no usage block). + if dataLen := int(providerUtils.GetJSONField(body, "data.#").Int()); dataLen > 0 { + if u.ImageUsage == nil { + u.ImageUsage = &schemas.ImageUsage{} + } + if u.ImageUsage.OutputTokensDetails == nil { + u.ImageUsage.OutputTokensDetails = &schemas.ImageTokenDetails{} + } + if u.ImageUsage.OutputTokensDetails.NImages == 0 { + u.ImageUsage.OutputTokensDetails.NImages = dataLen + } + } + } + + if u.ImageUsage == nil { + u.ImageUsage = &schemas.ImageUsage{} + } + // Populate LLMUsage from image token counts so logs show token totals. + if u.ImageUsage.TotalTokens > 0 { + u.LLMUsage = &schemas.BifrostLLMUsage{ + PromptTokens: u.ImageUsage.InputTokens, + CompletionTokens: u.ImageUsage.OutputTokens, + TotalTokens: u.ImageUsage.TotalTokens, + } + } + return u +} + +// ---- containers (code interpreter sessions) ---- +// Only the create call (POST /v1/containers) is billable — a flat per-session fee priced +// under the synthetic "container-{memory_limit}" model key (falling back to "container"). +// The collection path also serves list (GET /v1/containers), so disambiguate by response +// shape: a create returns a single {"object":"container", "id":...} object, while a list +// returns {"object":"list", "data":[...]}. Retrieve/delete hit /containers/{id} and never +// reach this extractor. Containers are never streamed, so body is plain JSON. + +func extractOAIContainerUsage(body []byte) *schemas.BifrostPassthroughUsage { + if len(body) == 0 { + return nil + } + var resp struct { + Object string `json:"object"` + ID string `json:"id"` + MemoryLimit string `json:"memory_limit"` + } + if err := sonic.Unmarshal(body, &resp); err != nil { + return nil + } + // Bill only a created container (single object), not list/other shapes. + if resp.Object != "container" || resp.ID == "" { + return nil + } + identifier := "container" + if resp.MemoryLimit != "" { + identifier = "container-" + resp.MemoryLimit + } + return &schemas.BifrostPassthroughUsage{ContainerIdentifier: identifier} +} diff --git a/core/providers/openai/passthrough_usage_test.go b/core/providers/openai/passthrough_usage_test.go new file mode 100644 index 00000000000..c1bdae8995b --- /dev/null +++ b/core/providers/openai/passthrough_usage_test.go @@ -0,0 +1,292 @@ +package openai_test + +import ( + "bytes" + "mime/multipart" + "testing" + + "github.com/maximhq/bifrost/core/providers/openai" + "github.com/maximhq/bifrost/core/schemas" +) + +// multipartBody builds a multipart/form-data body with the given fields, matching what the +// video passthrough extractor sniffs (boundary on the first line). +func multipartBody(t *testing.T, fields map[string]string) []byte { + t.Helper() + var buf bytes.Buffer + w := multipart.NewWriter(&buf) + for k, v := range fields { + if err := w.WriteField(k, v); err != nil { + t.Fatalf("WriteField: %v", err) + } + } + if err := w.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + return buf.Bytes() +} + +func TestExtractOpenAIPassthroughUsage(t *testing.T) { + tests := []struct { + name string + method string + path string + reqBody string + body string + check func(t *testing.T, u *schemas.BifrostPassthroughUsage) + }{ + { + name: "chat completions usage + service tier", + path: "/v1/chat/completions", + body: `{"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15},"service_tier":"default"}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + mustLLM(t, u, 10, 5, 15) + if u.ServiceTier == nil || *u.ServiceTier != schemas.BifrostServiceTierDefault { + t.Fatalf("service tier = %v, want default", u.ServiceTier) + } + }, + }, + { + name: "chat completions zero usage -> nil", + path: "/v1/chat/completions", + body: `{"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}}`, + check: mustNil, + }, + { + name: "chat completions content delta (no usage) -> nil", + path: "/v1/chat/completions", + body: `{"choices":[{"delta":{"content":"hi"}}]}`, + check: mustNil, + }, + { + name: "responses top-level usage with reasoning + search queries", + path: "/v1/responses", + body: `{"usage":{"input_tokens":20,"output_tokens":8,"total_tokens":28,"output_tokens_details":{"reasoning_tokens":3,"num_search_queries":2}}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + mustLLM(t, u, 20, 8, 28) + d := u.LLMUsage.CompletionTokensDetails + if d == nil || d.ReasoningTokens != 3 { + t.Fatalf("reasoning tokens = %v, want 3", d) + } + if d.NumSearchQueries == nil || *d.NumSearchQueries != 2 { + t.Fatalf("num search queries = %v, want 2", d.NumSearchQueries) + } + }, + }, + { + name: "responses nested response.usage", + path: "/v1/responses", + body: `{"response":{"usage":{"input_tokens":20,"output_tokens":8,"total_tokens":28}}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + mustLLM(t, u, 20, 8, 28) + }, + }, + { + name: "embeddings", + path: "/v1/embeddings", + body: `{"usage":{"prompt_tokens":12,"total_tokens":12}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil || u.LLMUsage.PromptTokens != 12 || u.LLMUsage.TotalTokens != 12 { + t.Fatalf("embeddings usage = %+v", u) + } + }, + }, + { + name: "speech char count from request", + path: "/v1/audio/speech", + reqBody: `{"input":"héllo"}`, // 5 runes + body: "", + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.AudioInputChars != 5 { + t.Fatalf("audio input chars = %+v, want 5", u) + } + }, + }, + { + name: "transcription token usage", + path: "/v1/audio/transcriptions", + body: `{"usage":{"type":"tokens","input_tokens":4,"total_tokens":10,"input_token_details":{"audio_tokens":3,"text_tokens":1},"seconds":2}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.LLMUsage == nil || u.LLMUsage.PromptTokens != 4 || u.LLMUsage.TotalTokens != 10 { + t.Fatalf("transcription usage = %+v", u) + } + if u.AudioTokenDetails == nil || u.AudioTokenDetails.AudioTokens != 3 { + t.Fatalf("audio token details = %+v", u.AudioTokenDetails) + } + if u.AudioSeconds == nil || *u.AudioSeconds != 2 { + t.Fatalf("audio seconds = %v, want 2", u.AudioSeconds) + } + }, + }, + { + name: "transcription duration fallback", + path: "/v1/audio/transcriptions", + body: `{"duration":3}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.AudioSeconds == nil || *u.AudioSeconds != 3 { + t.Fatalf("audio seconds = %+v, want 3", u) + } + }, + }, + { + name: "image generation response usage", + path: "/v1/images/generations", + body: `{"usage":{"input_tokens":5,"output_tokens":100,"total_tokens":105}}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ImageUsage == nil || u.ImageUsage.TotalTokens != 105 { + t.Fatalf("image usage = %+v", u) + } + if u.LLMUsage == nil || u.LLMUsage.TotalTokens != 105 { + t.Fatalf("image llm usage = %+v", u.LLMUsage) + } + }, + }, + { + name: "image variation count from request n", + path: "/v1/images/variations", + reqBody: `{"n":2}`, + body: `{"data":[{"b64_json":"x"},{"b64_json":"y"}]}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ImageUsage == nil || u.ImageUsage.OutputTokensDetails == nil || + u.ImageUsage.OutputTokensDetails.NImages != 2 { + t.Fatalf("image NImages = %+v", u) + } + }, + }, + { + name: "video seconds default (no body)", + path: "/v1/videos", + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.VideoSeconds == nil || *u.VideoSeconds != 4 { + t.Fatalf("video seconds = %+v, want default 4", u) + } + }, + }, + { + name: "container create with memory limit", + path: "/v1/containers", + body: `{"object":"container","id":"cntr_1","memory_limit":"1g"}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ContainerIdentifier != "container-1g" { + t.Fatalf("container id = %+v, want container-1g", u) + } + }, + }, + { + name: "container create without memory limit", + path: "/v1/containers", + body: `{"object":"container","id":"cntr_1"}`, + check: func(t *testing.T, u *schemas.BifrostPassthroughUsage) { + if u == nil || u.ContainerIdentifier != "container" { + t.Fatalf("container id = %+v, want container", u) + } + }, + }, + { + name: "container list -> nil (not billable)", + path: "/v1/containers", + body: `{"object":"list","data":[{"object":"container","id":"c1"}]}`, + check: mustNil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + method := tt.method + if method == "" { + method = "POST" + } + u := openai.ExtractOpenAIPassthroughUsage(method, tt.path, []byte(tt.reqBody), []byte(tt.body)) + tt.check(t, u) + }) + } +} + +func TestExtractOpenAIPassthroughUsage_VideoMultipartSeconds(t *testing.T) { + reqBody := multipartBody(t, map[string]string{"seconds": "6"}) + u := openai.ExtractOpenAIPassthroughUsage("POST", "/v1/videos", reqBody, nil) + if u == nil || u.VideoSeconds == nil || *u.VideoSeconds != 6 { + t.Fatalf("video seconds = %+v, want 6", u) + } +} + +// Video usage is billable only on POST. GET (list/retrieve) and DELETE must not accrue +// per-second video usage even though the path contains "/videos". +func TestExtractOpenAIPassthroughUsage_VideoNonPOSTNotBilled(t *testing.T) { + for _, method := range []string{"GET", "DELETE"} { + t.Run(method, func(t *testing.T) { + if u := openai.ExtractOpenAIPassthroughUsage(method, "/v1/videos", nil, nil); u != nil { + t.Fatalf("%s /v1/videos = %+v, want nil", method, u) + } + }) + } +} + +// JSON create-video bodies carry `seconds` as a top-level field (numeric or string). +func TestExtractOpenAIPassthroughUsage_VideoJSONSeconds(t *testing.T) { + for _, body := range []string{`{"seconds":12}`, `{"seconds":"12"}`} { + u := openai.ExtractOpenAIPassthroughUsage("POST", "/v1/videos", []byte(body), nil) + if u == nil || u.VideoSeconds == nil || *u.VideoSeconds != 12 { + t.Fatalf("video seconds for %s = %+v, want 12", body, u) + } + } +} + +// /v1/images/variations is multipart/form-data; size/quality/n ride as form fields. +func TestExtractOpenAIPassthroughUsage_ImageMultipartParams(t *testing.T) { + reqBody := multipartBody(t, map[string]string{"size": "1024x1024", "quality": "high", "n": "3"}) + u := openai.ExtractOpenAIPassthroughUsage("POST", "/v1/images/variations", reqBody, nil) + if u == nil { + t.Fatal("expected usage, got nil") + } + if u.ImageSize != "1024x1024" { + t.Fatalf("image size = %q, want 1024x1024", u.ImageSize) + } + if u.ImageQuality != "high" { + t.Fatalf("image quality = %q, want high", u.ImageQuality) + } + if u.ImageUsage == nil || u.ImageUsage.OutputTokensDetails == nil || + u.ImageUsage.OutputTokensDetails.NImages != 3 { + t.Fatalf("image NImages = %+v, want 3", u) + } +} + +func TestHasOpenAIPassthroughUsage(t *testing.T) { + tests := []struct { + name string + event string + want bool + }{ + {"top-level usage", `{"usage":{"total_tokens":5}}`, true}, + {"nested response.usage", `{"response":{"usage":{"total_tokens":5}}}`, true}, + {"content delta", `{"choices":[{"delta":{"content":"hi"}}]}`, false}, + {"ping", `{"type":"ping"}`, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := openai.HasOpenAIPassthroughUsage([]byte(tt.event)); got != tt.want { + t.Fatalf("HasOpenAIPassthroughUsage = %v, want %v", got, tt.want) + } + }) + } +} + +// ---- shared assertion helpers ---- + +func mustNil(t *testing.T, u *schemas.BifrostPassthroughUsage) { + t.Helper() + if u != nil { + t.Fatalf("expected nil usage, got %+v", u) + } +} + +func mustLLM(t *testing.T, u *schemas.BifrostPassthroughUsage, prompt, completion, total int) { + t.Helper() + if u == nil || u.LLMUsage == nil { + t.Fatalf("expected LLMUsage, got %+v", u) + } + if u.LLMUsage.PromptTokens != prompt || u.LLMUsage.CompletionTokens != completion || u.LLMUsage.TotalTokens != total { + t.Fatalf("LLMUsage = {prompt:%d completion:%d total:%d}, want {%d %d %d}", + u.LLMUsage.PromptTokens, u.LLMUsage.CompletionTokens, u.LLMUsage.TotalTokens, prompt, completion, total) + } +} diff --git a/core/providers/utils/passthrough_stream.go b/core/providers/utils/passthrough_stream.go new file mode 100644 index 00000000000..3e7aacf9899 --- /dev/null +++ b/core/providers/utils/passthrough_stream.go @@ -0,0 +1,196 @@ +package utils + +import ( + "bytes" + "context" + "io" + "time" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/valyala/fasthttp" +) + +// PassthroughStreamParams configures StreamPassthrough. +type PassthroughStreamParams struct { + StatusCode int + Headers map[string]string + Path string + // RawRequest is attached to the final chunk only (the streaming accumulator reads it there). + RawRequest []byte + // CancellationBody is forwarded to cancellation/timeout handlers. + CancellationBody []byte + StartTime time.Time + // UseTerminalDetector finalizes the stream early when a terminal marker (finishReason) + // appears in a framed event — for providers (Gemini/Vertex) that emit it before the HTTP + // body closes. + UseTerminalDetector bool + Logger schemas.Logger + // HasUsage is an optional cheap gjson presence check: it returns true only when an event + // carries a usage field worth parsing. When set, Observe is skipped (no full unmarshal) for + // events that fail it — the common case on long streams (content deltas, pings). When nil, + // every event is passed to Observe. + HasUsage func(event []byte) bool + // Observe is called once per complete SSE data event (JSON payload) as it streams, and + // returns the running usage (nil when the event adds nothing). The last non-nil value is + // attached to the final chunk. Implementations populate usage directly from the event — + // no full response body is retained. + Observe func(event []byte) *schemas.BifrostPassthroughUsage +} + +// StreamPassthrough runs the shared passthrough streaming loop. It forwards each raw upstream +// chunk to the client unchanged (byte-exact, unbounded — forwarding never depends on usage +// parsing), and in parallel frames complete SSE events into a bounded buffer, feeding each to +// params.Observe to build usage incrementally. On a terminal marker or EOF it emits the final +// chunk carrying RawRequest + the observed usage. No full response body is accumulated. +// +// This owns the idle-timeout wrapper, cancellation hookup, response release, and goroutine. +func StreamPassthrough( + ctx *schemas.BifrostContext, + postHookRunner schemas.PostHookRunner, + postHookSpanFinalizer func(context.Context), + resp *fasthttp.Response, + rawBodyStream io.Reader, + params PassthroughStreamParams, +) chan *schemas.BifrostStreamChunk { + // Wrap reader with idle timeout to detect stalled streams. + bodyStream, stopIdleTimeout := NewIdleTimeoutReader(rawBodyStream, rawBodyStream, GetStreamIdleTimeout(ctx), ctx) + // Cancellation must close the raw stream to unblock reads. + stopCancellation := SetupStreamCancellation(ctx, rawBodyStream, params.Logger) + + extraFields := schemas.BifrostResponseExtraFields{ + ProviderResponseHeaders: params.Headers, + PassthroughPath: params.Path, + } + + ch := make(chan *schemas.BifrostStreamChunk, schemas.DefaultStreamBufferSize) + go func() { + defer EnsureStreamFinalizerCalled(ctx, postHookSpanFinalizer) + defer func() { + if ctx.Err() == context.Canceled { + HandleStreamCancellation(ctx, postHookRunner, ch, params.Logger, postHookSpanFinalizer, params.CancellationBody) + } else if ctx.Err() == context.DeadlineExceeded { + HandleStreamTimeout(ctx, postHookRunner, ch, params.Logger, postHookSpanFinalizer, params.CancellationBody) + } + close(ch) + }() + defer ReleaseStreamingResponse(ctx, resp) + defer stopIdleTimeout() + defer stopCancellation() + + var pending bytes.Buffer + var usage *schemas.BifrostPassthroughUsage + + success := params.StatusCode >= 200 && params.StatusCode < 300 + + observe := func(payload []byte) (terminal bool) { + if len(bytes.TrimSpace(payload)) == 0 { + return false + } + // Cheap gjson gate: only fully parse events that actually carry usage. + if success && params.Observe != nil && (params.HasUsage == nil || params.HasUsage(payload)) { + if u := params.Observe(payload); u != nil { + usage = u + } + } + return params.UseTerminalDetector && isTerminalSSEPayload(payload) + } + + // drainFrames extracts every complete SSE event currently buffered and observes each. + // Returns true when a terminal event is seen. + drainFrames := func() bool { + for { + data := pending.Bytes() + idx, delimLen := findFirstSSEFrameDelimiter(data) + if idx < 0 { + break + } + frame := append([]byte(nil), data[:idx]...) + pending.Next(idx + delimLen) + if observe(extractSSEDataPayload(frame)) { + return true + } + } + // Bound the buffer. A single SSE event can be large — e.g. an + // image_generation.completed event carries the full base64 image with `usage` at + // its tail, so the whole event must be buffered to read usage. Match the native SSE + // scanner's per-line ceiling (sseMaxBufSize). Only when an undelimited event exceeds + // that do we drop to the last frame boundary (or reset) to stay bounded. + if pending.Len() > sseMaxBufSize { + drain := pending.Bytes() + if idx, delimLen := findLastSSEFrameDelimiter(drain); idx >= 0 { + pending.Next(idx + delimLen) + } else { + pending.Reset() + } + } + return false + } + + finalize := func() { + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + extraFields.Latency = time.Since(params.StartTime).Milliseconds() + extraFields.RawRequest = params.RawRequest + ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ + PassthroughResponse: &schemas.BifrostPassthroughResponse{ + StatusCode: params.StatusCode, + Headers: params.Headers, + ExtraFields: extraFields, + PassthroughUsage: usage, + }, + }, ch, postHookSpanFinalizer) + } + + buf := make([]byte, 4096) + for { + n, readErr := bodyStream.Read(buf) + if n > 0 { + chunk := make([]byte, n) + copy(chunk, buf[:n]) + // Forward the raw chunk to the client unchanged. + ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ + PassthroughResponse: &schemas.BifrostPassthroughResponse{ + StatusCode: params.StatusCode, + Headers: params.Headers, + Body: chunk, + ExtraFields: extraFields, + }, + }, ch, postHookSpanFinalizer) + + pending.Write(chunk) + if drainFrames() { + finalize() + return + } + } + if readErr == io.EOF { + // Flush a trailing event not terminated by a delimiter. + if pending.Len() > 0 { + observe(extractSSEDataPayload(pending.Bytes())) + } + finalize() + return + } + if readErr != nil { + if ctx.Err() != nil { + return // let defer handle cancel/timeout + } + ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) + extraFields.Latency = time.Since(params.StartTime).Milliseconds() + ProcessAndSendError(ctx, postHookRunner, readErr, ch, params.Logger, postHookSpanFinalizer) + return + } + } + }() + + return ch +} + +// isTerminalSSEPayload reports whether a framed SSE data payload signals stream completion +// via a finishReason/usage terminal marker. ([DONE] is handled by the SSE readers as EOF.) +func isTerminalSSEPayload(payload []byte) bool { + p := bytes.TrimSpace(payload) + if len(p) == 0 { + return false + } + return hasFinishReasonMarker(p) +} diff --git a/core/providers/utils/utils.go b/core/providers/utils/utils.go index 5ff1018518c..a05675d945f 100644 --- a/core/providers/utils/utils.go +++ b/core/providers/utils/utils.go @@ -518,7 +518,6 @@ var providerResponseFilterHeaders = map[string]bool{ "server": true, "alt-svc": true, "strict-transport-security": true, - "content-type": true, "access-control-allow-origin": true, "access-control-allow-methods": true, "access-control-allow-headers": true, diff --git a/core/providers/vertex/vertex.go b/core/providers/vertex/vertex.go index 26e978b2636..fbfbebcfedb 100644 --- a/core/providers/vertex/vertex.go +++ b/core/providers/vertex/vertex.go @@ -7,7 +7,6 @@ import ( "encoding/hex" "errors" "fmt" - "io" "net/http" "net/url" "regexp" @@ -3033,6 +3032,11 @@ func (provider *VertexProvider) Passthrough( return nil, providerUtils.NewBifrostOperationError("failed to decode response body", err) } + var passthroughUsage *schemas.BifrostPassthroughUsage + if resp.StatusCode() >= 200 && resp.StatusCode() < 300 { + passthroughUsage = gemini.ExtractGeminiPassthroughUsage(req.Path, req.Body, body) + } + bifrostResponse := &schemas.BifrostPassthroughResponse{ StatusCode: resp.StatusCode(), Headers: headers, @@ -3040,7 +3044,10 @@ func (provider *VertexProvider) Passthrough( ExtraFields: schemas.BifrostResponseExtraFields{ Latency: latency.Milliseconds(), ProviderResponseHeaders: headers, + PassthroughPath: req.Path, + RawRequest: req.Body, }, + PassthroughUsage: passthroughUsage, } return bifrostResponse, nil @@ -3180,91 +3187,22 @@ func (provider *VertexProvider) PassthroughStream( fmt.Errorf("provider returned an empty stream body")) } - // Set stream idle timeout from provider config. providerUtils.SetStreamIdleTimeoutIfEmpty(ctx, provider.networkConfig.StreamIdleTimeoutInSeconds) - - // Wrap body with idle timeout to detect stalled streams. - rawBodyStream := bodyStream - bodyStream, stopIdleTimeout := providerUtils.NewIdleTimeoutReader(bodyStream, rawBodyStream, providerUtils.GetStreamIdleTimeout(ctx), ctx) - - // Cancellation must close the raw stream to unblock reads. - stopCancellation := providerUtils.SetupStreamCancellation(ctx, rawBodyStream, provider.logger) - - extraFields := schemas.BifrostResponseExtraFields{} - statusCode := resp.StatusCode() - - ch := make(chan *schemas.BifrostStreamChunk, schemas.DefaultStreamBufferSize) - go func() { - defer providerUtils.EnsureStreamFinalizerCalled(ctx, postHookSpanFinalizer) - defer func() { - if ctx.Err() == context.Canceled { - providerUtils.HandleStreamCancellation(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } else if ctx.Err() == context.DeadlineExceeded { - providerUtils.HandleStreamTimeout(ctx, postHookRunner, ch, provider.logger, postHookSpanFinalizer, providerUtils.PassthroughJSONBody(fasthttpReq, req.Body)) - } - close(ch) - }() - defer providerUtils.ReleaseStreamingResponse(ctx, resp) - defer stopIdleTimeout() - defer stopCancellation() - streamStart := time.Now() - - terminalDetector := &providerUtils.StreamTerminalDetector{} - buf := make([]byte, 4096) - for { - n, readErr := bodyStream.Read(buf) - if n > 0 { - chunk := make([]byte, n) - copy(chunk, buf[:n]) - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - Body: chunk, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - - // Vertex streamGenerateContent passthrough can emit terminal markers - // (finishReason) before the underlying HTTP body is closed. - // Finalize as success once this appears to avoid hanging clients. - if terminalDetector.ObserveChunk(chunk) { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(streamStart).Milliseconds() - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - return - } - } - if readErr == io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(streamStart).Milliseconds() - providerUtils.ProcessAndSendResponse(ctx, postHookRunner, &schemas.BifrostResponse{ - PassthroughResponse: &schemas.BifrostPassthroughResponse{ - StatusCode: statusCode, - Headers: headers, - ExtraFields: extraFields, - }, - }, ch, postHookSpanFinalizer) - return - } - if readErr != nil { - if ctx.Err() != nil { - return // let defer handle cancel/timeout - } - if readErr != io.EOF { - ctx.SetValue(schemas.BifrostContextKeyStreamEndIndicator, true) - extraFields.Latency = time.Since(streamStart).Milliseconds() - providerUtils.ProcessAndSendError(ctx, postHookRunner, readErr, ch, provider.logger, postHookSpanFinalizer) - } - return - } - } - }() - return ch, nil + return providerUtils.StreamPassthrough( + ctx, postHookRunner, postHookSpanFinalizer, resp, bodyStream, + providerUtils.PassthroughStreamParams{ + StatusCode: resp.StatusCode(), + Headers: headers, + Path: req.Path, + RawRequest: req.Body, + CancellationBody: providerUtils.PassthroughJSONBody(fasthttpReq, req.Body), + StartTime: time.Now(), + UseTerminalDetector: true, + Logger: provider.logger, + HasUsage: gemini.HasGeminiPassthroughUsage, + Observe: func(event []byte) *schemas.BifrostPassthroughUsage { + return gemini.ExtractGeminiPassthroughUsage(req.Path, req.Body, event) + }, + }, + ), nil } diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index fadfa088122..242c90e91bf 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -1432,6 +1432,7 @@ type BifrostResponseExtraFields struct { ConvertedRequestType RequestType `json:"converted_request_type,omitempty"` DroppedCompatPluginParams []string `json:"dropped_compat_plugin_params,omitempty"` // params dropped by the compat plugin based on model catalog ProviderResponseHeaders map[string]string `json:"provider_response_headers,omitempty"` // HTTP response headers from the provider (filtered to exclude transport-level headers) + PassthroughPath string `json:"passthrough_path,omitempty"` // Stripped provider path for passthrough requests, e.g. "/v1/chat/completions" } type BifrostMCPResponseExtraFields struct { diff --git a/core/schemas/passthrough.go b/core/schemas/passthrough.go index 2da0028301b..ed743da3865 100644 --- a/core/schemas/passthrough.go +++ b/core/schemas/passthrough.go @@ -10,11 +10,42 @@ type BifrostPassthroughRequest struct { SafeHeaders map[string]string // client headers, auth already stripped } +// BifrostPassthroughUsage carries usage data extracted by the provider at stream +// completion. The pricing module converts this into cost using the existing compute +// functions — no new pricing logic is required. +type BifrostPassthroughUsage struct { + // Text / chat / responses / embeddings + LLMUsage *BifrostLLMUsage + ServiceTier *BifrostServiceTier // "priority" | "flex" | nil (default) + + // Image generation / edit / variation + ImageUsage *ImageUsage + ImageSize string // e.g. "1024x1024" + ImageQuality string // "low" | "medium" | "high" | "auto" + + // Speech TTS — character count from request body `input` field + AudioInputChars int + + // Transcription — token details or raw seconds as duration fallback + AudioSeconds *int + AudioTokenDetails *TranscriptionUsageInputTokenDetails + + // Video generation + VideoSeconds *int + + // Container creation (code interpreter session) — synthetic pricing identifier, + // e.g. "container-1g", or "container" when no memory limit is reported. Maps to + // costInput.containerIdentifierString for the flat per-session fee. + ContainerIdentifier string +} + type BifrostPassthroughResponse struct { - StatusCode int - Headers map[string]string - Body []byte - ExtraFields BifrostResponseExtraFields + StatusCode int + Headers map[string]string + Body []byte + ExtraFields BifrostResponseExtraFields + Path string // stripped provider path, e.g. "/v1/chat/completions" + PassthroughUsage *BifrostPassthroughUsage // usage extracted by the provider for billing — set on the unary response (non-streaming) or the final streaming chunk; nil when no billable usage could be extracted } type PassthroughLogParams struct { @@ -22,4 +53,5 @@ type PassthroughLogParams struct { Path string `json:"path"` // stripped path, e.g. "/v1/fine-tuning/jobs" RawQuery string `json:"raw_query"` // raw query string, no "?" StatusCode int `json:"status_code"` + Model string `json:"model,omitempty"` // model extracted from path or request body } diff --git a/framework/modelcatalog/pricing.go b/framework/modelcatalog/pricing.go index f7dbfb2e169..eb80ed833d7 100644 --- a/framework/modelcatalog/pricing.go +++ b/framework/modelcatalog/pricing.go @@ -299,7 +299,7 @@ func (mc *ModelCatalog) calculateBaseCost(result *schemas.BifrostResponse, scope resolvedModelUsed := extraFields.ResolvedModelUsed requestType := extraFields.RequestType - // Extract usage data from the response + // Extract usage data from the response (passthrough and native paths unified) input := extractCostInput(result) // If provider already computed cost, use it @@ -312,8 +312,13 @@ func (mc *ModelCatalog) calculateBaseCost(result *schemas.BifrostResponse, scope return 0 } - // Normalize stream request types to their base type for pricing lookup - requestType = normalizeStreamRequestType(requestType) + if result.PassthroughResponse != nil { + // Infer request type from usage fields + path; passthrough bypasses stream normalization. + requestType = inferPassthroughRequestType(extraFields.Provider, extraFields.PassthroughPath, result.PassthroughResponse.PassthroughUsage) + } else { + // Normalize stream request types to their base type for pricing lookup + requestType = normalizeStreamRequestType(requestType) + } // When a pricing model override is set, use it in place of the actual requested/resolved // model names during pricing lookup (e.g. container creates always look up "container"). @@ -362,6 +367,9 @@ func extractCostInput(result *schemas.BifrostResponse) costInput { var input costInput switch { + case result.PassthroughResponse != nil && result.PassthroughResponse.PassthroughUsage != nil: + return passthroughUsageToCostInput(result.PassthroughResponse.PassthroughUsage) + case result.TextCompletionResponse != nil && result.TextCompletionResponse.Usage != nil: input.usage = result.TextCompletionResponse.Usage @@ -1335,3 +1343,134 @@ func (mc *ModelCatalog) UpsertModelPricingAttributes(ctx context.Context, model } return rows, nil } + +// --------------------------------------------------------------------------- +// Passthrough pricing helpers +// --------------------------------------------------------------------------- + +// detectPassthroughRequestType maps a provider + stripped path to a RequestType. +func detectPassthroughRequestType(provider schemas.ModelProvider, path string) schemas.RequestType { + if idx := strings.IndexByte(path, '?'); idx >= 0 { + path = path[:idx] + } + path = strings.TrimRight(path, "/") + switch provider { + case schemas.OpenAI, schemas.Azure: + switch { + case strings.HasSuffix(path, "/chat/completions"): + return schemas.ChatCompletionRequest + case strings.HasSuffix(path, "/completions"): + return schemas.TextCompletionRequest + case strings.HasSuffix(path, "/embeddings"): + return schemas.EmbeddingRequest + case strings.HasSuffix(path, "/responses"): + return schemas.ResponsesRequest + case strings.HasSuffix(path, "/images/generations"): + return schemas.ImageGenerationRequest + case strings.HasSuffix(path, "/images/edits"): + return schemas.ImageEditRequest + case strings.HasSuffix(path, "/images/variations"): + return schemas.ImageVariationRequest + case strings.HasSuffix(path, "/audio/speech"): + return schemas.SpeechRequest + case strings.HasSuffix(path, "/audio/transcriptions"), + strings.HasSuffix(path, "/audio/translations"): + return schemas.TranscriptionRequest + case strings.HasSuffix(path, "/containers"): + return schemas.ContainerCreateRequest + case strings.Contains(path, "/video"): + return schemas.VideoGenerationRequest + default: + return schemas.ChatCompletionRequest + } + case schemas.Gemini, schemas.Vertex: + // Interactions API paths carry no colon action suffix. + if strings.Contains(path, "/interactions") { + return schemas.ResponsesRequest + } + colonIdx := strings.LastIndexByte(path, ':') + if colonIdx < 0 { + return schemas.ChatCompletionRequest + } + switch path[colonIdx+1:] { + case "generateContent", "streamGenerateContent": + return schemas.ResponsesRequest + case "embedContent", "batchEmbedContents": + return schemas.EmbeddingRequest + case "generateImages": + return schemas.ImageGenerationRequest + case "predict": + return schemas.EmbeddingRequest + case "predictLongRunning": + return schemas.VideoGenerationRequest + default: + return schemas.ChatCompletionRequest + } + case schemas.Anthropic: + switch { + case strings.HasSuffix(path, "/messages"): + return schemas.ResponsesRequest + case strings.HasSuffix(path, "/complete"): + return schemas.TextCompletionRequest + default: + return schemas.ResponsesRequest + } + default: + return schemas.ChatCompletionRequest + } +} + +// inferPassthroughRequestType determines the request type from usage fields (primary) +// and falls back to path detection for text/embedding/responses where LLMUsage is ambiguous. +func inferPassthroughRequestType(provider schemas.ModelProvider, path string, su *schemas.BifrostPassthroughUsage) schemas.RequestType { + if su != nil { + if su.ContainerIdentifier != "" { + return schemas.ContainerCreateRequest + } + if su.ImageUsage != nil { + return schemas.ImageGenerationRequest + } + if su.AudioInputChars > 0 { + return schemas.SpeechRequest + } + if su.AudioTokenDetails != nil || su.AudioSeconds != nil { + return schemas.TranscriptionRequest + } + if su.VideoSeconds != nil { + return schemas.VideoGenerationRequest + } + } + return detectPassthroughRequestType(provider, path) +} + +// passthroughUsageToCostInput converts BifrostPassthroughUsage into costInput. +func passthroughUsageToCostInput(su *schemas.BifrostPassthroughUsage) costInput { + var input costInput + if su.LLMUsage != nil { + input.usage = su.LLMUsage + } + if su.ServiceTier != nil { + input.tier = tierFromString(su.ServiceTier) + } + if su.ImageUsage != nil { + input.imageUsage = su.ImageUsage + input.imageSize = su.ImageSize + input.imageQuality = su.ImageQuality + } + if su.AudioInputChars > 0 { + input.audioTextInputChars = su.AudioInputChars + } + if su.AudioSeconds != nil { + input.audioSeconds = su.AudioSeconds + } + if su.AudioTokenDetails != nil { + input.audioTokenDetails = su.AudioTokenDetails + } + if su.VideoSeconds != nil { + input.videoSeconds = su.VideoSeconds + } + if su.ContainerIdentifier != "" { + input.containerIdentifierString = su.ContainerIdentifier + } + return input +} diff --git a/framework/streaming/passthrough.go b/framework/streaming/passthrough.go index e898e680454..1cdad9870d1 100644 --- a/framework/streaming/passthrough.go +++ b/framework/streaming/passthrough.go @@ -41,6 +41,11 @@ func (a *Accumulator) processPassthroughStreamingResponse(ctx *schemas.BifrostCo maps.Copy(accumulator.PassthroughHeaders, result.PassthroughResponse.Headers) } + // Save path from the first chunk that carries it (set once in ExtraFields by the provider) + if accumulator.PassthroughPath == "" && result != nil && result.PassthroughResponse != nil { + accumulator.PassthroughPath = result.PassthroughResponse.ExtraFields.PassthroughPath + } + // Accumulate the body bytes from this chunk if result != nil && result.PassthroughResponse != nil && len(result.PassthroughResponse.Body) > 0 { // Make a copy of the body bytes to avoid referencing pooled memory @@ -67,14 +72,25 @@ func (a *Accumulator) processPassthroughStreamingResponse(ctx *schemas.BifrostCo accumulator.FinalTimestamp = time.Now() } + // PassthroughUsage is set by the provider on the final EOF chunk before any + // plugin runs — read it from result rather than re-extracting here. + var passthroughUsage *schemas.BifrostPassthroughUsage + if result != nil && result.PassthroughResponse != nil { + passthroughUsage = result.PassthroughResponse.PassthroughUsage + } + // Build the accumulated passthrough response passthroughResp := &schemas.BifrostPassthroughResponse{ - StatusCode: accumulator.PassthroughStatusCode, - Headers: accumulator.PassthroughHeaders, - Body: accumulator.PassthroughBody, + StatusCode: accumulator.PassthroughStatusCode, + Headers: accumulator.PassthroughHeaders, + Body: accumulator.PassthroughBody, + Path: accumulator.PassthroughPath, + PassthroughUsage: passthroughUsage, } - // Build accumulated data with the passthrough response + // Build accumulated data with the passthrough response. + // Populate TokenUsage from PassthroughUsage.LLMUsage so applyStreamingOutputToEntry sets + // entry.TokenUsageParsed via the standard streaming token path. data := &AccumulatedData{ RequestID: requestID, Model: requestedModel, @@ -84,6 +100,9 @@ func (a *Accumulator) processPassthroughStreamingResponse(ctx *schemas.BifrostCo EndTimestamp: accumulator.FinalTimestamp, PassthroughOutput: passthroughResp, } + if passthroughUsage != nil && passthroughUsage.LLMUsage != nil { + data.TokenUsage = passthroughUsage.LLMUsage + } // Set error status if there was an error if bifrostErr != nil { diff --git a/framework/streaming/types.go b/framework/streaming/types.go index 16779891d75..c6739633314 100644 --- a/framework/streaming/types.go +++ b/framework/streaming/types.go @@ -141,9 +141,10 @@ type StreamAccumulator struct { TerminalErrorChunkIndex int // Passthrough streaming accumulation - PassthroughBody []byte // Accumulated body bytes from passthrough streaming chunks - PassthroughStatusCode int // Status code from passthrough response - PassthroughHeaders map[string]string // Headers from passthrough response + PassthroughBody []byte // Accumulated body bytes from passthrough streaming chunks + PassthroughStatusCode int // Status code from passthrough response + PassthroughHeaders map[string]string // Headers from passthrough response + PassthroughPath string // Stripped provider path, e.g. "/v1/chat/completions" IsComplete bool FinalTimestamp time.Time diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 5658a19f819..09e4ecb0971 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -1819,8 +1819,13 @@ func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provi tokensUsed = *result.TranscriptionResponse.Usage.TotalTokens case result.TranscriptionStreamResponse != nil && result.TranscriptionStreamResponse.Usage != nil && result.TranscriptionStreamResponse.Usage.TotalTokens != nil: tokensUsed = *result.TranscriptionStreamResponse.Usage.TotalTokens + case result.PassthroughResponse != nil: + if su := result.PassthroughResponse.PassthroughUsage; su != nil && su.LLMUsage != nil { + tokensUsed = su.LLMUsage.TotalTokens + } } } + // Create usage update for tracker (business logic) usageUpdate := &UsageUpdate{ VirtualKey: virtualKey, @@ -1833,7 +1838,7 @@ func (p *GovernancePlugin) postHookWorker(result *schemas.BifrostResponse, provi UserID: userID, IsStreaming: isStreaming, IsFinalChunk: isFinalChunk, - HasUsageData: tokensUsed > 0, + HasUsageData: tokensUsed > 0 || cost > 0, } // Queue usage update asynchronously using tracker diff --git a/plugins/logging/main.go b/plugins/logging/main.go index 7e9688c63c1..b94570e564a 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -675,6 +675,7 @@ func (p *LoggerPlugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.Bifr Method: req.PassthroughRequest.Method, Path: req.PassthroughRequest.Path, RawQuery: req.PassthroughRequest.RawQuery, + Model: req.PassthroughRequest.Model, } if len(req.PassthroughRequest.Body) > 0 { ct := strings.ToLower(req.PassthroughRequest.SafeHeaders["content-type"]) @@ -1046,6 +1047,13 @@ func (p *LoggerPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *schemas. if isPassthroughErrorResponse(result) { entry.Status = "error" } + // Compute cost for streaming passthrough using StreamUsage set by the accumulator. + if entry.Cost == nil && p.pricingManager != nil && result.PassthroughResponse.PassthroughUsage != nil { + pricingScopes := modelcatalog.PricingLookupScopesFromContext(ctx, string(entry.Provider)) + if cost := p.pricingManager.CalculateCost(result, pricingScopes); cost > 0 { + entry.Cost = &cost + } + } } applyLargePayloadPreviewsToEntry(ctx, entry, contentLoggingEnabled) if tracer != nil && traceID != "" { diff --git a/plugins/logging/operations.go b/plugins/logging/operations.go index ccbcc6b3f42..4421ad72a7d 100644 --- a/plugins/logging/operations.go +++ b/plugins/logging/operations.go @@ -478,6 +478,10 @@ func (p *LoggerPlugin) applyNonStreamingOutputToEntry(entry *logstore.Log, resul } else { usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } + case result.PassthroughResponse != nil: + if su := result.PassthroughResponse.PassthroughUsage; su != nil { + usage = su.LLMUsage + } } if usage != nil { entry.TokenUsageParsed = usage diff --git a/transports/bifrost-http/integrations/router.go b/transports/bifrost-http/integrations/router.go index 96db94f5c3c..60d6b52e68c 100644 --- a/transports/bifrost-http/integrations/router.go +++ b/transports/bifrost-http/integrations/router.go @@ -3054,7 +3054,22 @@ func (g *GenericRouter) handlePassthroughStream( ctx.SetUserValue(schemas.BifrostContextKeyDeferTraceCompletion, true) ctx.SetStatusCode(passthroughResp.StatusCode) - ctx.SetContentType("text/event-stream") + // Preserve the upstream Content-Type. Passthrough streams aren't always SSE — e.g. + // Vertex/Gemini :streamGenerateContent without ?alt=sse returns an incrementally-delivered + // JSON array with Content-Type: application/json. Forcing text/event-stream mislabels that + // stream, so clients that dispatch on content-type run an SSE parser over a non-SSE body and + // hang. Fall back to text/event-stream only when the upstream didn't provide a Content-Type. + contentType := "" + for k, v := range passthroughResp.Headers { + if strings.EqualFold(k, "content-type") { + contentType = v + break + } + } + if contentType == "" { + contentType = "text/event-stream" + } + ctx.SetContentType(contentType) ctx.Response.Header.Set("Cache-Control", "no-cache") ctx.Response.Header.Set("Connection", "keep-alive") ctx.Response.Header.Set("X-Accel-Buffering", "no") @@ -3063,7 +3078,8 @@ func (g *GenericRouter) handlePassthroughStream( case "connection", "transfer-encoding", "content-length", "content-type", "cache-control", "x-accel-buffering", "set-cookie", "proxy-authenticate", "www-authenticate": - // drop — streaming invariants are set explicitly above; upstream must not override them + // drop — streaming invariants are set explicitly above (Content-Type is set from the + // upstream value before this loop); upstream must not override them here default: ctx.Response.Header.Set(k, v) }