feat: stream SSE tokens in-flight with delta content fallback

internal/pricing/pricing.go

@@ -137,9 +137,11 @@ func ExtractUsageFromResponse(provider string, responseBody []byte, logger *slog
 	}
 }
 
-// AccumulateStreamingUsage extracts usage from SSE chunks and returns final token counts.
+// AccumulateStreamingUsage extracts usage from a fully buffered SSE response.
+// When no usage field is found, it falls back to accumulating delta.content byte length
+// before resorting to whole-body-length estimation.
 func AccumulateStreamingUsage(provider string, sseChunks []byte, logger *slog.Logger) Usage {
-	var usage Usage
+	var inputTokens, outputTokens, deltaContentLen int
 	scanner := bufio.NewScanner(bytes.NewReader(sseChunks))
 	scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
 
@@ -148,35 +150,36 @@ func AccumulateStreamingUsage(provider string, sseChunks []byte, logger *slog.Lo
 		if !strings.HasPrefix(line, "data:") {
 			continue
 		}
-
 		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
 		if payload == "" || payload == "[DONE]" {
 			continue
 		}
-
-		inputTokens, outputTokens, ok := extractUsage([]byte(payload))
-		if !ok {
+		raw := []byte(payload)
+		if input, output, ok := extractUsage(raw); ok {
+			if input > 0 {
+				inputTokens = input
+			}
+			if output > 0 {
+				outputTokens = output
+			}
 			continue
 		}
-		if inputTokens > 0 {
-			usage.InputTokens = inputTokens
-		}
-		if outputTokens > 0 {
-			usage.OutputTokens = outputTokens
+		if content, ok := extractDeltaContent(raw); ok {
+			deltaContentLen += len(content)
 		}
 	}
 
-	if usage.InputTokens > 0 || usage.OutputTokens > 0 {
-		return usage
+	if inputTokens > 0 || outputTokens > 0 {
+		return Usage{InputTokens: inputTokens, OutputTokens: outputTokens}
+	}
+	if deltaContentLen > 0 {
+		tokens := estimateTokensFromBodyLength(deltaContentLen)
+		warnMissingUsage(logger, provider, deltaContentLen, true)
+		return Usage{OutputTokens: tokens, Estimated: true}
 	}
-
 	estimated := estimateTokensFromBodyLength(len(sseChunks))
 	warnMissingUsage(logger, provider, len(sseChunks), true)
-	return Usage{
-		InputTokens:  estimated,
-		OutputTokens: estimated,
-		Estimated:    true,
-	}
+	return Usage{InputTokens: estimated, OutputTokens: estimated, Estimated: true}
 }
 
 func extractUsage(rawJSON []byte) (int, int, bool) {

internal/pricing/pricing_test.go

@@ -318,8 +318,9 @@ func TestAccumulateStreamingUsage_Fixtures(t *testing.T) {
 				t.Fatalf("OutputTokens = %d, want %d", got.OutputTokens, tt.wantOutput)
 			}
 			if tt.wantEstimated {
-				if got.InputTokens < 1 || got.OutputTokens < 1 {
-					t.Fatalf("fallback estimated usage must be >=1 token, got input=%d output=%d", got.InputTokens, got.OutputTokens)
+				// Delta-content fallback sets OutputTokens only; body-length fallback sets both.
+				if got.OutputTokens < 1 {
+					t.Fatalf("fallback estimated output tokens must be >=1, got input=%d output=%d", got.InputTokens, got.OutputTokens)
 				}
 			}
 			if tt.wantWarnSubstr != "" && !strings.Contains(logs.String(), tt.wantWarnSubstr) {

internal/pricing/streaming.go

@@ -0,0 +1,101 @@
+package pricing
+
+import (
+	"encoding/json"
+	"log/slog"
+	"strings"
+)
+
+// SSEAccumulator processes SSE response chunks in-flight for zero-copy token accounting.
+// Feed each response chunk through Write as it arrives; call Usage once the stream ends.
+//
+//nolint:govet // field order kept for readability over alignment.
+type SSEAccumulator struct {
+	inputTokens     int
+	outputTokens    int
+	deltaContentLen int             // byte length of accumulated delta.content strings
+	partialLine     strings.Builder // incomplete SSE line waiting for \n
+}
+
+// NewSSEAccumulator returns a ready-to-use accumulator.
+func NewSSEAccumulator() *SSEAccumulator {
+	return &SSEAccumulator{}
+}
+
+// Write processes a response chunk in-flight. Safe to call multiple times.
+// Chunks may contain partial SSE lines; state is carried across calls.
+func (a *SSEAccumulator) Write(chunk []byte) {
+	for _, b := range chunk {
+		if b == '\n' {
+			a.processLine(strings.TrimSpace(a.partialLine.String()))
+			a.partialLine.Reset()
+		} else {
+			a.partialLine.WriteByte(b)
+		}
+	}
+}
+
+// Usage returns the final token counts after the stream has been fully consumed.
+// Priority: usage field from any chunk > delta content estimation > minimum fallback.
+func (a *SSEAccumulator) Usage(provider string, logger *slog.Logger) Usage {
+	// Flush any partial line that arrived without a trailing newline.
+	if a.partialLine.Len() > 0 {
+		a.processLine(strings.TrimSpace(a.partialLine.String()))
+		a.partialLine.Reset()
+	}
+	if a.inputTokens > 0 || a.outputTokens > 0 {
+		return Usage{InputTokens: a.inputTokens, OutputTokens: a.outputTokens}
+	}
+	if a.deltaContentLen > 0 {
+		tokens := estimateTokensFromBodyLength(a.deltaContentLen)
+		warnMissingUsage(logger, provider, a.deltaContentLen, true)
+		return Usage{OutputTokens: tokens, Estimated: true}
+	}
+	warnMissingUsage(logger, provider, 0, true)
+	return Usage{InputTokens: 1, OutputTokens: 1, Estimated: true}
+}
+
+func (a *SSEAccumulator) processLine(line string) {
+	if !strings.HasPrefix(line, "data:") {
+		return
+	}
+	payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
+	if payload == "" || payload == "[DONE]" {
+		return
+	}
+	raw := []byte(payload)
+
+	// First try the usage field (final chunk from most providers).
+	if input, output, ok := extractUsage(raw); ok {
+		if input > 0 {
+			a.inputTokens = input
+		}
+		if output > 0 {
+			a.outputTokens = output
+		}
+		return
+	}
+
+	// Fall back to accumulating delta content bytes.
+	if content, ok := extractDeltaContent(raw); ok {
+		a.deltaContentLen += len(content)
+	}
+}
+
+// extractDeltaContent pulls the delta.content string from an OpenAI streaming chunk.
+func extractDeltaContent(raw []byte) (string, bool) {
+	var chunk struct {
+		Choices []struct {
+			Delta struct {
+				Content string `json:"content"`
+			} `json:"delta"`
+		} `json:"choices"`
+	}
+	if err := json.Unmarshal(raw, &chunk); err != nil {
+		return "", false
+	}
+	if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" {
+		return chunk.Choices[0].Delta.Content, true
+	}
+	return "", false
+}

internal/pricing/streaming_test.go

@@ -0,0 +1,171 @@
+package pricing
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestSSEAccumulator_WithUsageField(t *testing.T) {
+	t.Parallel()
+
+	acc := NewSSEAccumulator()
+	// Chunk with delta content followed by final chunk with usage.
+	acc.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n"))
+	acc.Write([]byte("data: {\"usage\":{\"prompt_tokens\":30,\"completion_tokens\":15}}\n\n"))
+	acc.Write([]byte("data: [DONE]\n\n"))
+
+	logger, _ := newTestLogger()
+	got := acc.Usage("openai", logger)
+
+	if got.Estimated {
+		t.Fatal("Estimated = true, want false (usage field present)")
+	}
+	if got.InputTokens != 30 {
+		t.Fatalf("InputTokens = %d, want 30", got.InputTokens)
+	}
+	if got.OutputTokens != 15 {
+		t.Fatalf("OutputTokens = %d, want 15", got.OutputTokens)
+	}
+}
+
+func TestSSEAccumulator_DeltaContentFallback(t *testing.T) {
+	t.Parallel()
+
+	acc := NewSSEAccumulator()
+	// Two chunks with content, no usage field.
+	acc.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n"))
+	acc.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"world\"}}]}\n\n"))
+	acc.Write([]byte("data: [DONE]\n\n"))
+
+	logger, logs := newTestLogger()
+	got := acc.Usage("openai", logger)
+
+	if !got.Estimated {
+		t.Fatal("Estimated = false, want true (no usage field, delta content only)")
+	}
+	// "Hello " (6) + "world" (5) = 11 bytes → max(1, 11/4) = 2 output tokens
+	if got.OutputTokens < 1 {
+		t.Fatalf("OutputTokens = %d, want >= 1", got.OutputTokens)
+	}
+	if got.InputTokens != 0 {
+		t.Fatalf("InputTokens = %d, want 0 (unknown for delta-only fallback)", got.InputTokens)
+	}
+	if !strings.Contains(logs.String(), "usage data missing") {
+		t.Fatalf("expected warning log, got %q", logs.String())
+	}
+}
+
+func TestSSEAccumulator_MalformedChunksPassThrough(t *testing.T) {
+	t.Parallel()
+
+	acc := NewSSEAccumulator()
+	// Malformed JSON in data line — should not panic or error.
+	acc.Write([]byte("data: {not valid json}\n\n"))
+	acc.Write([]byte("data: \n\n"))
+	acc.Write([]byte(": comment line\n\n"))
+	acc.Write([]byte("data: [DONE]\n\n"))
+
+	logger, logs := newTestLogger()
+	got := acc.Usage("openai", logger)
+
+	// No usage found, no delta content → minimum fallback.
+	if !got.Estimated {
+		t.Fatal("Estimated = false, want true (no parseable data)")
+	}
+	if got.OutputTokens < 1 {
+		t.Fatalf("OutputTokens = %d, want >= 1 (minimum fallback)", got.OutputTokens)
+	}
+	if !strings.Contains(logs.String(), "usage data missing") {
+		t.Fatalf("expected warning log, got %q", logs.String())
+	}
+}
+
+func TestSSEAccumulator_DoneTriggersRecording(t *testing.T) {
+	t.Parallel()
+
+	acc := NewSSEAccumulator()
+	// Usage arrives before [DONE] — should be captured regardless of [DONE].
+	acc.Write([]byte("data: {\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":3}}\n\n"))
+	acc.Write([]byte("data: [DONE]\n\n"))
+
+	logger, _ := newTestLogger()
+	got := acc.Usage("openai", logger)
+
+	if got.Estimated {
+		t.Fatal("Estimated = true, want false")
+	}
+	if got.InputTokens != 5 || got.OutputTokens != 3 {
+		t.Fatalf("tokens = (%d, %d), want (5, 3)", got.InputTokens, got.OutputTokens)
+	}
+}
+
+func TestSSEAccumulator_PartialLineSpansChunks(t *testing.T) {
+	t.Parallel()
+
+	acc := NewSSEAccumulator()
+	// Split "data: {...}\n" across two Write calls.
+	acc.Write([]byte("data: {\"usage\":{\"prompt_tokens\":10,"))
+	acc.Write([]byte("\"completion_tokens\":7}}\n\n"))
+
+	logger, _ := newTestLogger()
+	got := acc.Usage("openai", logger)
+
+	if got.Estimated {
+		t.Fatal("Estimated = true, want false (usage spread across two chunks)")
+	}
+	if got.InputTokens != 10 || got.OutputTokens != 7 {
+		t.Fatalf("tokens = (%d, %d), want (10, 7)", got.InputTokens, got.OutputTokens)
+	}
+}
+
+func TestExtractDeltaContent(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name   string
+		raw    string
+		want   string
+		wantOK bool
+	}{
+		{
+			name:   "extracts content from delta",
+			raw:    `{"choices":[{"delta":{"content":"Hello"}}]}`,
+			want:   "Hello",
+			wantOK: true,
+		},
+		{
+			name:   "empty content returns false",
+			raw:    `{"choices":[{"delta":{"content":""}}]}`,
+			wantOK: false,
+		},
+		{
+			name:   "no choices returns false",
+			raw:    `{"choices":[]}`,
+			wantOK: false,
+		},
+		{
+			name:   "usage-only chunk returns false",
+			raw:    `{"usage":{"prompt_tokens":10,"completion_tokens":5}}`,
+			wantOK: false,
+		},
+		{
+			name:   "malformed json returns false",
+			raw:    `{not json}`,
+			wantOK: false,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			got, ok := extractDeltaContent([]byte(tt.raw))
+			if ok != tt.wantOK {
+				t.Fatalf("ok = %v, want %v", ok, tt.wantOK)
+			}
+			if ok && got != tt.want {
+				t.Fatalf("content = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}