[fix]: framework/tracing - emit BilledUsage token and cost attributes on failed requests

framework/tracing/llmspan.go

@@ -148,6 +148,60 @@ func PopulateErrorAttributes(err *schemas.BifrostError) map[string]any {
 		attrs[schemas.AttrHTTPResponseStatusCode] = *err.StatusCode
 	}
 
+	// Usage the provider billed us for even though the request failed or was
+	// cancelled (see BifrostError.ExtraFields.BilledUsage). Governance and the
+	// logging plugin already charge for this; without emitting it here every
+	// span-based consumer — the otel plugin and the BigQuery / Datadog / Kafka /
+	// Pub-Sub connectors — records zero tokens for the same request.
+	if u := err.ExtraFields.BilledUsage; u != nil {
+		// Everything below is gated on > 0: attachBilledUsageFromContext
+		// (core/providers/utils) attaches a BilledUsage when *any* of tokens,
+		// details or cost is set, so a details-only usage reaches here with
+		// zero totals; emitting those would stamp explicit zeros on the span
+		// where the success paths only ever see provider-reported values.
+		if u.PromptTokens > 0 {
+			attrs[schemas.AttrInputTokens] = u.PromptTokens
+		}
+		if u.CompletionTokens > 0 {
+			attrs[schemas.AttrOutputTokens] = u.CompletionTokens
+		}
+		if u.TotalTokens > 0 {
+			attrs[schemas.AttrTotalTokens] = u.TotalTokens
+		}
+
+		if d := u.PromptTokensDetails; d != nil {
+			// The nested cached-write detail keys are namespaced per request
+			// type on the success paths (input_token_details.* for Responses,
+			// prompt_token_details.* for chat); mirror that here instead of
+			// emitting both families, which would break the otel plugin's
+			// assumption that the two namespaces are mutually exclusive.
+			isResponses := err.ExtraFields.RequestType == schemas.ResponsesRequest ||
+				err.ExtraFields.RequestType == schemas.ResponsesStreamRequest
+			if d.CachedReadTokens > 0 {
+				attrs[schemas.AttrUsageCacheReadInputTokens] = d.CachedReadTokens
+			}
+			if d.CachedWriteTokens > 0 {
+				attrs[schemas.AttrUsageCacheCreationInputTokens] = d.CachedWriteTokens
+			}
+			if wd := d.CachedWriteTokenDetails; wd != nil {
+				if wd.CachedWriteTokens5m > 0 {
+					if isResponses {
+						attrs[schemas.AttrInputTokenDetailsCachedWrite5m] = wd.CachedWriteTokens5m
+					} else {
+						attrs[schemas.AttrPromptTokenDetailsCachedWrite5m] = wd.CachedWriteTokens5m
+					}
+				}
+				if wd.CachedWriteTokens1h > 0 {
+					if isResponses {
+						attrs[schemas.AttrInputTokenDetailsCachedWrite1h] = wd.CachedWriteTokens1h
+					} else {
+						attrs[schemas.AttrPromptTokenDetailsCachedWrite1h] = wd.CachedWriteTokens1h
+					}
+				}
+			}
+		}
+	}
+
 	return attrs
 }
 

framework/tracing/llmspan_test.go

@@ -176,3 +176,162 @@ func TestPopulateRequestExtraParamsSerializesStructuredValues(t *testing.T) {
 		})
 	}
 }
+
+func TestPopulateErrorAttributesEmitsBilledUsage(t *testing.T) {
+	msg := "stream cancelled by client"
+	bifrostErr := &schemas.BifrostError{
+		Error: &schemas.ErrorField{Message: msg},
+	}
+	bifrostErr.ExtraFields.RequestType = schemas.ChatCompletionStreamRequest
+	bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{
+		PromptTokens:     1200,
+		CompletionTokens: 34,
+		TotalTokens:      1234,
+		PromptTokensDetails: &schemas.ChatPromptTokensDetails{
+			CachedReadTokens:  1000,
+			CachedWriteTokens: 200,
+			CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{
+				CachedWriteTokens5m: 120,
+				CachedWriteTokens1h: 80,
+			},
+		},
+	}
+
+	attrs := PopulateErrorAttributes(bifrostErr)
+
+	for key, want := range map[string]any{
+		schemas.AttrInputTokens:                     1200,
+		schemas.AttrOutputTokens:                    34,
+		schemas.AttrTotalTokens:                     1234,
+		schemas.AttrUsageCacheReadInputTokens:       1000,
+		schemas.AttrUsageCacheCreationInputTokens:   200,
+		schemas.AttrPromptTokenDetailsCachedWrite5m: 120,
+		schemas.AttrPromptTokenDetailsCachedWrite1h: 80,
+	} {
+		if got := attrs[key]; got != want {
+			t.Errorf("attribute %s = %v, want %v", key, got, want)
+		}
+	}
+	// A failed chat span must not carry the Responses namespace: the otel
+	// plugin treats the two 5m/1h families as mutually exclusive per request.
+	for _, key := range []string{
+		schemas.AttrInputTokenDetailsCachedWrite5m,
+		schemas.AttrInputTokenDetailsCachedWrite1h,
+	} {
+		if _, ok := attrs[key]; ok {
+			t.Errorf("Responses-namespace attribute %s present on a chat span", key)
+		}
+	}
+}
+
+func TestPopulateErrorAttributesUsesResponsesNamespace(t *testing.T) {
+	bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "responses stream cancelled"}}
+	bifrostErr.ExtraFields.RequestType = schemas.ResponsesStreamRequest
+	bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{
+		PromptTokensDetails: &schemas.ChatPromptTokensDetails{
+			CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{
+				CachedWriteTokens5m: 120,
+				CachedWriteTokens1h: 80,
+			},
+		},
+	}
+
+	attrs := PopulateErrorAttributes(bifrostErr)
+
+	for key, want := range map[string]any{
+		schemas.AttrInputTokenDetailsCachedWrite5m: 120,
+		schemas.AttrInputTokenDetailsCachedWrite1h: 80,
+	} {
+		if got := attrs[key]; got != want {
+			t.Errorf("attribute %s = %v, want %v", key, got, want)
+		}
+	}
+	for _, key := range []string{
+		schemas.AttrPromptTokenDetailsCachedWrite5m,
+		schemas.AttrPromptTokenDetailsCachedWrite1h,
+	} {
+		if _, ok := attrs[key]; ok {
+			t.Errorf("chat-namespace attribute %s present on a Responses span", key)
+		}
+	}
+}
+
+func TestPopulateErrorAttributesWithoutBilledUsageEmitsNoTokens(t *testing.T) {
+	msg := "401 before the model ran"
+	bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: msg}}
+
+	attrs := PopulateErrorAttributes(bifrostErr)
+
+	for _, key := range []string{schemas.AttrInputTokens, schemas.AttrOutputTokens, schemas.AttrTotalTokens} {
+		if _, ok := attrs[key]; ok {
+			t.Errorf("attribute %s present for a request that consumed no tokens", key)
+		}
+	}
+}
+
+func TestPopulateErrorAttributesEmitsCacheWriteDetailsWithoutAggregate(t *testing.T) {
+	bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "stream failed during cache creation"}}
+	bifrostErr.ExtraFields.RequestType = schemas.ChatCompletionStreamRequest
+	bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{
+		PromptTokensDetails: &schemas.ChatPromptTokensDetails{
+			CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{
+				CachedWriteTokens5m: 120,
+				CachedWriteTokens1h: 80,
+			},
+		},
+	}
+
+	attrs := PopulateErrorAttributes(bifrostErr)
+
+	for key, want := range map[string]any{
+		schemas.AttrPromptTokenDetailsCachedWrite5m: 120,
+		schemas.AttrPromptTokenDetailsCachedWrite1h: 80,
+	} {
+		if got := attrs[key]; got != want {
+			t.Errorf("attribute %s = %v, want %v", key, got, want)
+		}
+	}
+	// Zero-valued aggregates and totals stay absent: this BilledUsage carries
+	// only cache-write details, so emitting the totals would stamp explicit
+	// zeros on the span.
+	for _, key := range []string{
+		schemas.AttrUsageCacheCreationInputTokens,
+		schemas.AttrInputTokens,
+		schemas.AttrOutputTokens,
+		schemas.AttrTotalTokens,
+	} {
+		if _, ok := attrs[key]; ok {
+			t.Errorf("zero-valued attribute %s is present", key)
+		}
+	}
+}
+
+// A cancelled stream reaches PopulateLLMResponseAttributes with BOTH a non-nil
+// accumulated response and a non-nil error (see core/providers/utils). The
+// accumulated response is missing the final usage chunk, so the error's
+// BilledUsage must win. This mirrors the merge order in Tracer.
+func TestErrorAttributesOverrideAccumulatedResponseTokens(t *testing.T) {
+	partial := &schemas.BifrostResponse{
+		ChatResponse: &schemas.BifrostChatResponse{
+			Usage: &schemas.BifrostLLMUsage{PromptTokens: 0, CompletionTokens: 0, TotalTokens: 0},
+		},
+	}
+	bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "client cancelled the stream"}}
+	bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{
+		PromptTokens:     4096,
+		CompletionTokens: 128,
+		TotalTokens:      4224,
+	}
+
+	attrs := PopulateResponseAttributes(partial)
+	for k, v := range PopulateErrorAttributes(bifrostErr) {
+		attrs[k] = v
+	}
+
+	if got := attrs[schemas.AttrInputTokens]; got != 4096 {
+		t.Errorf("%s = %v, want 4096 from BilledUsage", schemas.AttrInputTokens, got)
+	}
+	if got := attrs[schemas.AttrTotalTokens]; got != 4224 {
+		t.Errorf("%s = %v, want 4224 from BilledUsage", schemas.AttrTotalTokens, got)
+	}
+}

framework/tracing/tracer.go

@@ -451,7 +451,19 @@ func (t *Tracer) PopulateLLMResponseAttributes(ctx *schemas.BifrostContext, hand
 		return
 	}
 	respAttrs := PopulateResponseAttributes(resp)
+	// A cancelled stream arrives here with an accumulated response whose usage
+	// is missing the final chunk, so its aggregate token counts read zero. When
+	// the error carries the authoritative BilledUsage, drop those zeros from
+	// the response side: PopulateErrorAttributes gates its own emissions on
+	// > 0, so a zero stamped here would survive the merge and turn "not
+	// recorded" into a false zero on the span.
+	billed := err != nil && err.ExtraFields.BilledUsage != nil
 	for k, v := range respAttrs {
+		if billed && (k == schemas.AttrInputTokens || k == schemas.AttrOutputTokens || k == schemas.AttrTotalTokens) {
+			if n, ok := v.(int); ok && n == 0 {
+				continue
+			}
+		}
 		if k == schemas.AttrFinishReasons {
 			// Spec: gen_ai.response.finish_reasons (string[]) belongs on the GenAI (llm.call) span.
 			span.SetAttribute(schemas.AttrFinishReasons, v)
@@ -481,8 +493,40 @@ func (t *Tracer) PopulateLLMResponseAttributes(ctx *schemas.BifrostContext, hand
 		span.SetAttribute(schemas.AttrBifrostRoutingEngineUsed, strings.Join(engines, ","))
 	}
 
-	// Populate cost attribute using pricing manager
-	if t.pricingManager != nil && resp != nil {
+	// Populate cost attribute using pricing manager. BilledUsage wins when it is
+	// present: it is what the provider actually charged for a failed or cancelled
+	// turn. A cancelled stream still yields a non-nil accumulated response (see
+	// providers/utils, which passes both accumulatedResp and err), but that
+	// response is missing the final usage chunk, so pricing it would report 0.
+	if t.pricingManager != nil && err != nil && err.ExtraFields.BilledUsage != nil {
+		// Core calls BifrostError.PopulateExtraFields around RunPostLLMHooks, so
+		// Provider / RequestType / the model fields are always set here.
+		ef := err.ExtraFields
+		model := ef.ResolvedModelUsed
+		if model == "" {
+			model = ef.OriginalModelRequested
+		}
+		cost := t.pricingManager.CalculateCostForUsage(
+			ef.BilledUsage,
+			ef.Provider,
+			model,
+			ef.RequestType,
+			modelcatalog.PricingLookupScopesFromContext(ctx, string(ef.Provider)),
+		)
+		// When the catalog cannot price the model, fall back to the cost the
+		// provider itself reported (deep-copied into BilledUsage by
+		// attachBilledUsageFromContext) rather than discarding it.
+		if cost == 0 && ef.BilledUsage.Cost != nil {
+			cost = ef.BilledUsage.Cost.TotalCost
+		}
+		// Guarded write: a resp == nil failure emitted no cost attribute before
+		// this path existed, and consumers rely on distinguishing "no cost
+		// recorded" from a genuine zero (see plugins/logging, which guards the
+		// same way).
+		if cost > 0 {
+			span.SetAttribute(schemas.AttrUsageCost, cost)
+		}
+	} else if t.pricingManager != nil && resp != nil {
 		cost := t.pricingManager.CalculateCost(resp, modelcatalog.PricingLookupScopesFromContext(ctx, string(resp.GetExtraFields().Provider)))
 		span.SetAttribute(schemas.AttrUsageCost, cost)
 	}

framework/tracing/tracer_test.go

@@ -417,6 +417,50 @@ func TestTracer_SetAttribute(t *testing.T) {
 	}
 }
 
+// A cancelled stream reaches PopulateLLMResponseAttributes with an accumulated
+// response whose usage exists but reads zero (the final usage chunk never
+// arrived) and an error carrying the authoritative BilledUsage. The response
+// side's zero aggregates must not survive onto the span: PopulateErrorAttributes
+// gates its emissions on > 0, so a details-only BilledUsage would otherwise
+// leave a false zero in gen_ai.usage.*.
+func TestTracer_PopulateLLMResponseAttributesDropsZeroAggregatesWhenBilled(t *testing.T) {
+	store := NewTraceStore(5*time.Minute, nil)
+	defer store.Stop()
+
+	tracer := NewTracer(store, nil, nil)
+	defer tracer.Stop()
+
+	traceID := tracer.CreateTrace("")
+	ctx := context.WithValue(context.Background(), schemas.BifrostContextKeyTraceID, traceID)
+	_, handle := tracer.StartSpan(ctx, "llm.call", schemas.SpanKindLLMCall)
+
+	resp := &schemas.BifrostResponse{
+		ChatResponse: &schemas.BifrostChatResponse{
+			Usage: &schemas.BifrostLLMUsage{PromptTokens: 0, CompletionTokens: 0, TotalTokens: 0},
+		},
+	}
+	bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "client cancelled the stream"}}
+	bifrostErr.ExtraFields.RequestType = schemas.ChatCompletionStreamRequest
+	bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{
+		PromptTokensDetails: &schemas.ChatPromptTokensDetails{
+			CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{CachedWriteTokens5m: 120},
+		},
+	}
+
+	bctx := schemas.NewBifrostContext(context.Background(), time.Time{})
+	tracer.PopulateLLMResponseAttributes(bctx, handle, resp, bifrostErr)
+
+	span := store.GetTrace(traceID).RootSpan
+	for _, key := range []string{schemas.AttrInputTokens, schemas.AttrOutputTokens, schemas.AttrTotalTokens} {
+		if v, ok := span.Attributes[key]; ok {
+			t.Errorf("zero-valued response aggregate %s = %v survived onto the billed failed span", key, v)
+		}
+	}
+	if got := span.Attributes[schemas.AttrPromptTokenDetailsCachedWrite5m]; got != 120 {
+		t.Errorf("attribute %s = %v, want 120", schemas.AttrPromptTokenDetailsCachedWrite5m, got)
+	}
+}
+
 func TestTracer_GetSpanHandleByID_RootSpan(t *testing.T) {
 	store := NewTraceStore(5*time.Minute, nil)
 	defer store.Stop()