diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 6234aa1ed..c04c2c83c 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -460,7 +460,16 @@ func (b *StreamUsageBuffer) Observe(detail usage.Detail, ok bool) { if b == nil || !ok { return } - b.detail = detail + responseServiceTier := strings.TrimSpace(detail.ResponseServiceTier) + if responseServiceTier == "" || hasNonZeroTokenUsage(detail) { + preservedTier := b.detail.ResponseServiceTier + b.detail = detail + if b.detail.ResponseServiceTier == "" { + b.detail.ResponseServiceTier = preservedTier + } + } else { + b.detail.ResponseServiceTier = responseServiceTier + } b.ok = true } @@ -482,12 +491,16 @@ func (b *StreamUsageBuffer) Detail() (usage.Detail, bool) { } func ParseCodexUsage(data []byte) (usage.Detail, bool) { + responseServiceTier := extractResponseServiceTier(data) usageNode := gjson.ParseBytes(data).Get("response.usage") if !hasOpenAIStyleUsageTokenFields(usageNode) { - return usage.Detail{}, false + if responseServiceTier == "" { + return usage.Detail{}, false + } + return usage.Detail{ResponseServiceTier: responseServiceTier}, true } detail := parseOpenAIStyleUsageNode(usageNode) - detail.ResponseServiceTier = extractResponseServiceTier(data) + detail.ResponseServiceTier = responseServiceTier return detail, true } @@ -500,12 +513,13 @@ func ParseCodexImageToolUsage(data []byte) (usage.Detail, bool) { } func ParseOpenAIUsage(data []byte) usage.Detail { + responseServiceTier := extractResponseServiceTier(data) usageNode := gjson.ParseBytes(data).Get("usage") if !hasOpenAIStyleUsageTokenFields(usageNode) { - return usage.Detail{} + return usage.Detail{ResponseServiceTier: responseServiceTier} } detail := parseOpenAIStyleUsageNode(usageNode) - detail.ResponseServiceTier = extractResponseServiceTier(data) + detail.ResponseServiceTier = responseServiceTier return detail } @@ -564,12 +578,16 @@ func ParseOpenAIStreamUsage(line []byte) (usage.Detail, bool) { if len(payload) == 0 || !gjson.ValidBytes(payload) { return usage.Detail{}, false } + responseServiceTier := extractResponseServiceTier(payload) usageNode := gjson.GetBytes(payload, "usage") if !hasOpenAIStyleUsageTokenFields(usageNode) { - return usage.Detail{}, false + if responseServiceTier == "" { + return usage.Detail{}, false + } + return usage.Detail{ResponseServiceTier: responseServiceTier}, true } detail := parseOpenAIStyleUsageNode(usageNode) - detail.ResponseServiceTier = extractResponseServiceTier(payload) + detail.ResponseServiceTier = responseServiceTier return detail, true } diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index b7107752f..0d1ff6713 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -88,6 +88,24 @@ func TestParseOpenAIUsageIgnoresNullUsage(t *testing.T) { } } +func TestParseOpenAIUsagePreservesResponseTierWithoutUsage(t *testing.T) { + t.Parallel() + + detail := ParseOpenAIUsage([]byte(`{"service_tier":"default"}`)) + if detail.ResponseServiceTier != "default" { + t.Fatalf("response service tier = %q, want default", detail.ResponseServiceTier) + } +} + +func TestParseCodexUsagePreservesResponseTierWithoutUsage(t *testing.T) { + t.Parallel() + + detail, ok := ParseCodexUsage([]byte(`{"response":{"service_tier":"default"}}`)) + if !ok || detail.ResponseServiceTier != "default" { + t.Fatalf("ParseCodexUsage() = (%+v, %v), want response tier default", detail, ok) + } +} + func TestParseOpenAIStreamUsageIgnoresNullUsage(t *testing.T) { line := []byte(`data: {"id":"chunk_1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}],"usage":null}`) if detail, ok := ParseOpenAIStreamUsage(line); ok { @@ -145,6 +163,21 @@ func TestStreamUsageBufferKeepsLastUsage(t *testing.T) { } } +func TestStreamUsageBufferPreservesTierAcrossChunks(t *testing.T) { + t.Parallel() + + var buffer StreamUsageBuffer + buffer.Observe(ParseOpenAIStreamUsage([]byte(`data: {"service_tier":"default"}`))) + buffer.Observe(ParseOpenAIStreamUsage([]byte(`data: {"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))) + detail, ok := buffer.Detail() + if !ok { + t.Fatal("Detail() ok = false, want true") + } + if detail.InputTokens != 1 || detail.OutputTokens != 1 || detail.ResponseServiceTier != "default" { + t.Fatalf("detail = %+v, want usage with response tier default", detail) + } +} + func TestStreamUsageBufferPreservesOnlyZeroUsage(t *testing.T) { var buffer StreamUsageBuffer buffer.Observe(usage.Detail{}, true)