fix: passthrough budgets

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()

core/providers/anthropic/anthropic.go

@@ -2605,14 +2605,21 @@ 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,
 		Body:       body,
 		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
 }

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,
+		},
+	}
+}