fix: extract and estimate tokens

Author: lich0821

internal/proxy/proxy.go

@@ -353,7 +353,13 @@ func (p *Proxy) handleProxy(w http.ResponseWriter, r *http.Request) {
 		isStreaming := contentType == "text/event-stream" || (claudeReq.Stream && strings.Contains(contentType, "text/event-stream"))
 
 		if resp.StatusCode == http.StatusOK && isStreaming {
-			inputTokens, outputTokens, _ := p.handleStreamingResponse(w, resp, endpoint, trans, transformerName, thinkingEnabled)
+			inputTokens, outputTokens, outputText := p.handleStreamingResponse(w, resp, endpoint, trans, transformerName, thinkingEnabled)
+
+			// Fallback: estimate tokens when usage is 0
+			if inputTokens == 0 || outputTokens == 0 {
+				inputTokens, outputTokens = p.estimateTokens(bodyBytes, outputText, inputTokens, outputTokens, endpoint.Name)
+			}
+
 			p.stats.RecordTokens(endpoint.Name, inputTokens, outputTokens)
 			p.markRequestInactive(endpoint.Name)
 			logger.Debug("[%s] Request completed successfully (streaming)", endpoint.Name)

internal/proxy/streaming.go

@@ -17,7 +17,7 @@ import (
 )
 
 // handleStreamingResponse processes streaming SSE responses
-func (p *Proxy) handleStreamingResponse(w http.ResponseWriter, resp *http.Response, endpoint config.Endpoint, trans transformer.Transformer, transformerName string, thinkingEnabled bool) (int, int, error) {
+func (p *Proxy) handleStreamingResponse(w http.ResponseWriter, resp *http.Response, endpoint config.Endpoint, trans transformer.Transformer, transformerName string, thinkingEnabled bool) (int, int, string) {
 	// Copy response headers except Content-Length (streaming response length is unknown)
 	for key, values := range resp.Header {
 		if key == "Content-Length" {
@@ -33,7 +33,7 @@ func (p *Proxy) handleStreamingResponse(w http.ResponseWriter, resp *http.Respon
 	if !ok {
 		logger.Error("[%s] ResponseWriter does not support flushing", endpoint.Name)
 		resp.Body.Close()
-		return 0, 0, nil
+		return 0, 0, ""
 	}
 
 	var streamCtx *transformer.StreamContext
@@ -110,7 +110,7 @@ func (p *Proxy) handleStreamingResponse(w http.ResponseWriter, resp *http.Respon
 	}
 
 	resp.Body.Close()
-	return inputTokens, outputTokens, nil
+	return inputTokens, outputTokens, outputText.String()
 }
 
 // transformStreamEvent transforms a single SSE event
@@ -134,7 +134,7 @@ func (p *Proxy) extractTokensFromEvent(eventData []byte, inputTokens, outputToke
 			continue
 		}
 
-		jsonData := strings.TrimPrefix(line, "data:")
+		jsonData := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
 		var event map[string]interface{}
 		if err := json.Unmarshal([]byte(jsonData), &event); err != nil {
 			continue
@@ -168,7 +168,7 @@ func (p *Proxy) extractTextFromEvent(transformedEvent []byte, outputText *string
 			continue
 		}
 
-		jsonData := strings.TrimPrefix(line, "data:")
+		jsonData := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
 		var event map[string]interface{}
 		if err := json.Unmarshal([]byte(jsonData), &event); err != nil {
 			continue

internal/proxy/utils.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/lich0821/ccNexus/internal/logger"
+	"github.com/lich0821/ccNexus/internal/tokencount"
 )
 
 // normalizeAPIUrl ensures the API URL has a protocol prefix
@@ -88,3 +89,21 @@ func cleanIncompleteToolCalls(bodyBytes []byte) ([]byte, error) {
 	req["messages"] = messages
 	return json.Marshal(req)
 }
+
+// estimateTokens estimates tokens when API doesn't provide usage
+func (p *Proxy) estimateTokens(bodyBytes []byte, outputText string, inputTokens, outputTokens int, endpointName string) (int, int) {
+	if inputTokens == 0 {
+		var req tokencount.CountTokensRequest
+		if json.Unmarshal(bodyBytes, &req) == nil {
+			inputTokens = tokencount.EstimateInputTokens(&req)
+			logger.Debug("[%s] Estimated input tokens: %d", endpointName, inputTokens)
+		}
+	}
+
+	if outputTokens == 0 && outputText != "" {
+		outputTokens = tokencount.EstimateOutputTokens(outputText)
+		logger.Debug("[%s] Estimated output tokens: %d", endpointName, outputTokens)
+	}
+
+	return inputTokens, outputTokens
+}

internal/transformer/gemini/gemini.go

@@ -410,6 +410,12 @@ func (t *GeminiTransformer) transformStreamingResponse(geminiStream []byte, ctx
 			continue
 		}
 
+		// Extract usage metadata from chunk
+		if chunk.UsageMetadata != nil {
+			ctx.InputTokens = chunk.UsageMetadata.PromptTokenCount
+			ctx.OutputTokens = chunk.UsageMetadata.CandidatesTokenCount
+		}
+
 		// Send message_start on first chunk
 		if !ctx.MessageStartSent {
 			ctx.MessageID = fmt.Sprintf("msg_%d", 0)
@@ -425,8 +431,8 @@ func (t *GeminiTransformer) transformStreamingResponse(geminiStream []byte, ctx
 					"content": []interface{}{},
 					"model":   ctx.ModelName,
 					"usage": map[string]interface{}{
-						"input_tokens":  0,
-						"output_tokens": 0,
+						"input_tokens":  ctx.InputTokens,
+						"output_tokens": ctx.OutputTokens,
 					},
 				},
 			}
@@ -610,12 +616,6 @@ func (t *GeminiTransformer) transformStreamingResponse(geminiStream []byte, ctx
 				result.WriteString("\n")
 			}
 		}
-
-		// Update usage metadata
-		if chunk.UsageMetadata != nil {
-			ctx.InputTokens = chunk.UsageMetadata.PromptTokenCount
-			ctx.OutputTokens = chunk.UsageMetadata.CandidatesTokenCount
-		}
 	}
 
 	if err := scanner.Err(); err != nil {

internal/transformer/openai/event_handler.go

@@ -24,6 +24,16 @@ func (h *MessageStartHandler) Handle(event *SSEEvent, state *StreamState) ([]*SS
 		return []*SSEEvent{event}, nil
 	}
 
+	// Extract usage from first event
+	if usage, ok := event.Data["usage"].(map[string]interface{}); ok {
+		if promptTokens, ok := usage["prompt_tokens"].(float64); ok {
+			state.InputTokens = int(promptTokens)
+		}
+		if completionTokens, ok := usage["completion_tokens"].(float64); ok {
+			state.OutputTokens = int(completionTokens)
+		}
+	}
+
 	if !state.MessageStarted {
 		if id, ok := event.Data["id"].(string); ok {
 			state.MessageID = id
@@ -46,8 +56,8 @@ func (h *MessageStartHandler) Handle(event *SSEEvent, state *StreamState) ([]*SS
 					"content": []interface{}{},
 					"model":   state.ModelName,
 					"usage": map[string]interface{}{
-						"input_tokens":  0,
-						"output_tokens": 0,
+						"input_tokens":  state.InputTokens,
+						"output_tokens": state.OutputTokens,
 					},
 				},
 			},
@@ -65,6 +75,16 @@ func (h *ContentDeltaHandler) Handle(event *SSEEvent, state *StreamState) ([]*SS
 		return nil, nil
 	}
 
+	// Extract usage from every event
+	if usage, ok := event.Data["usage"].(map[string]interface{}); ok {
+		if promptTokens, ok := usage["prompt_tokens"].(float64); ok {
+			state.InputTokens = int(promptTokens)
+		}
+		if completionTokens, ok := usage["completion_tokens"].(float64); ok {
+			state.OutputTokens = int(completionTokens)
+		}
+	}
+
 	choices, ok := event.Data["choices"].([]interface{})
 	if !ok || len(choices) == 0 {
 		return nil, nil