From b864092dad4355c9c2b0c56cac22beac9f87e244 Mon Sep 17 00:00:00 2001 From: Alexander Date: Tue, 14 Apr 2026 13:55:06 +0200 Subject: [PATCH] fix(stream): extract input tokens from message_start event message_delta only contains output_tokens. Input tokens are in the message_start event under message.usage.input_tokens. This was causing input token counts to be near-zero for all streaming requests. --- internal/proxy/handler.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/internal/proxy/handler.go b/internal/proxy/handler.go index 146d92e..dd3a519 100644 --- a/internal/proxy/handler.go +++ b/internal/proxy/handler.go @@ -278,11 +278,13 @@ func handleStream(c *gin.Context, upstream *UpstreamClient, san *Sanitizer, pool c.Writer.WriteString(line + "\n") flusher.Flush() - // Extract token usage from message_delta event if len(line) > 5 && line[:5] == "data:" { data := line[5:] - if gjson.Get(data, "type").String() == "message_delta" { - inputTokens = gjson.Get(data, "usage.input_tokens").Int() + eventType := gjson.Get(data, "type").String() + switch eventType { + case "message_start": + inputTokens = gjson.Get(data, "message.usage.input_tokens").Int() + case "message_delta": outputTokens = gjson.Get(data, "usage.output_tokens").Int() } }