feat: propagate token usage

cmd/ocgo/main.go

@@ -60,13 +60,18 @@ type ATool struct {
 }
 
 type OAIRequest struct {
-	Model       string       `json:"model"`
+	Model         string            `json:"model"`
-	Messages    []OAIMessage `json:"messages"`
+	Messages      []OAIMessage      `json:"messages"`
-	Stream      bool         `json:"stream,omitempty"`
+	Stream        bool              `json:"stream,omitempty"`
-	MaxTokens   int          `json:"max_tokens,omitempty"`
+	StreamOptions *OAIStreamOptions `json:"stream_options,omitempty"`
-	Temperature *float64     `json:"temperature,omitempty"`
+	MaxTokens     int               `json:"max_tokens,omitempty"`
-	TopP        *float64     `json:"top_p,omitempty"`
+	Temperature   *float64          `json:"temperature,omitempty"`
-	Tools       []OAITool    `json:"tools,omitempty"`
+	TopP          *float64          `json:"top_p,omitempty"`
+	Tools         []OAITool         `json:"tools,omitempty"`
+}
+
+type OAIStreamOptions struct {
+	IncludeUsage bool `json:"include_usage"`
 }
 
 type ResponsesRequest struct {
@@ -482,14 +487,14 @@ func prepareChatBody(body []byte) ([]byte, error) {
 	if json.Unmarshal(body, &req) != nil {
 		return body, nil
 	}
+	changed := requestStreamingUsage(req)
 	model, _ := req["model"].(string)
-	if !rawChatBodyHasImages(req) {
+	if rawChatBodyHasImages(req) {
-		return body, nil
+		if !modelSupportsImages(model) {
+			return nil, unsupportedImageModelError(model)
+		}
+		changed = stripRawChatImageDetails(req) || changed
 	}
-	if !modelSupportsImages(model) {
-		return nil, unsupportedImageModelError(model)
-	}
-	changed := stripRawChatImageDetails(req)
 	if !changed {
 		return body, nil
 	}
@@ -500,6 +505,23 @@ func prepareChatBody(body []byte) ([]byte, error) {
 	return out, nil
 }
 
+func requestStreamingUsage(req map[string]any) bool {
+	streaming, _ := req["stream"].(bool)
+	if !streaming {
+		return false
+	}
+	options, ok := req["stream_options"].(map[string]any)
+	if !ok {
+		options = map[string]any{}
+		req["stream_options"] = options
+	}
+	if enabled, _ := options["include_usage"].(bool); enabled {
+		return false
+	}
+	options["include_usage"] = true
+	return true
+}
+
 func rawChatBodyHasImages(req map[string]any) bool {
 	messages, _ := req["messages"].([]any)
 	for _, item := range messages {
@@ -564,7 +586,7 @@ func convertRequest(ar AnthropicRequest) OAIRequest {
 	if model == "" || strings.HasPrefix(model, "claude-") {
 		model = "kimi-k2.6"
 	}
-	out := OAIRequest{Model: model, Stream: ar.Stream, MaxTokens: ar.MaxTokens, Temperature: ar.Temperature, TopP: ar.TopP}
+	out := OAIRequest{Model: model, Stream: ar.Stream, StreamOptions: streamUsageOptions(ar.Stream), MaxTokens: ar.MaxTokens, Temperature: ar.Temperature, TopP: ar.TopP}
 	if sys := systemText(ar.System); sys != "" {
 		out.Messages = append(out.Messages, OAIMessage{Role: "system", Content: sys})
 	}
@@ -582,7 +604,7 @@ func responsesToChat(rr ResponsesRequest) OAIRequest {
 	if model == "" {
 		model = "kimi-k2.6"
 	}
-	out := OAIRequest{Model: model, Stream: rr.Stream, MaxTokens: rr.MaxTokens, Temperature: rr.Temperature, TopP: rr.TopP}
+	out := OAIRequest{Model: model, Stream: rr.Stream, StreamOptions: streamUsageOptions(rr.Stream), MaxTokens: rr.MaxTokens, Temperature: rr.Temperature, TopP: rr.TopP}
 	if rr.Instructions != "" {
 		out.Messages = append(out.Messages, OAIMessage{Role: "system", Content: rr.Instructions})
 	}
@@ -595,6 +617,13 @@ func responsesToChat(rr ResponsesRequest) OAIRequest {
 	return out
 }
 
+func streamUsageOptions(streaming bool) *OAIStreamOptions {
+	if !streaming {
+		return nil
+	}
+	return &OAIStreamOptions{IncludeUsage: true}
+}
+
 func requestHasImages(or OAIRequest) bool {
 	for _, m := range or.Messages {
 		if contentHasImage(m.Content) {
@@ -927,6 +956,98 @@ func blockText(raw json.RawMessage) string {
 	return b.String()
 }
 
+type tokenUsage struct {
+	InputTokens       int
+	OutputTokens      int
+	TotalTokens       int
+	CachedInputTokens int
+	Present           bool
+}
+
+func usageFromJSON(raw json.RawMessage) tokenUsage {
+	var fields map[string]any
+	if len(raw) == 0 || json.Unmarshal(raw, &fields) != nil {
+		return tokenUsage{}
+	}
+	return usageFromFields(fields)
+}
+
+func usageFromFields(fields map[string]any) tokenUsage {
+	if len(fields) == 0 {
+		return tokenUsage{}
+	}
+	u := tokenUsage{Present: true}
+	u.InputTokens = intField(fields, "prompt_tokens")
+	if u.InputTokens == 0 {
+		u.InputTokens = intField(fields, "input_tokens")
+	}
+	u.OutputTokens = intField(fields, "completion_tokens")
+	if u.OutputTokens == 0 {
+		u.OutputTokens = intField(fields, "output_tokens")
+	}
+	u.TotalTokens = intField(fields, "total_tokens")
+	if u.TotalTokens == 0 && (u.InputTokens > 0 || u.OutputTokens > 0) {
+		u.TotalTokens = u.InputTokens + u.OutputTokens
+	}
+	u.CachedInputTokens = cachedTokens(fields)
+	return u
+}
+
+func intField(fields map[string]any, name string) int {
+	v, ok := fields[name]
+	if !ok {
+		return 0
+	}
+	switch n := v.(type) {
+	case float64:
+		return int(n)
+	case int:
+		return n
+	case json.Number:
+		i, _ := n.Int64()
+		return int(i)
+	}
+	return 0
+}
+
+func cachedTokens(fields map[string]any) int {
+	for _, key := range []string{"prompt_tokens_details", "input_tokens_details"} {
+		if nested, ok := fields[key].(map[string]any); ok {
+			if n := intField(nested, "cached_tokens"); n > 0 {
+				return n
+			}
+		}
+	}
+	return intField(fields, "cached_tokens")
+}
+
+func anthropicUsage(u tokenUsage) map[string]int {
+	usage := map[string]int{"input_tokens": u.InputTokens, "output_tokens": u.OutputTokens}
+	if u.CachedInputTokens > 0 {
+		usage["cache_read_input_tokens"] = u.CachedInputTokens
+	}
+	return usage
+}
+
+func anthropicDeltaUsage(u tokenUsage) map[string]int {
+	usage := map[string]int{"output_tokens": u.OutputTokens}
+	if u.InputTokens > 0 {
+		usage["input_tokens"] = u.InputTokens
+	}
+	if u.CachedInputTokens > 0 {
+		usage["cache_read_input_tokens"] = u.CachedInputTokens
+	}
+	return usage
+}
+
+func responsesUsage(u tokenUsage) map[string]any {
+	usage := map[string]any{"input_tokens": u.InputTokens, "output_tokens": u.OutputTokens, "total_tokens": u.TotalTokens}
+	if u.CachedInputTokens > 0 {
+		usage["input_tokens_details"] = map[string]int{"cached_tokens": u.CachedInputTokens}
+	}
+	return usage
+}
+
 func streamAnthropic(w http.ResponseWriter, body io.Reader, model string) {
 	w.Header().Set("Content-Type", "text/event-stream")
 	flusher, _ := w.(http.Flusher)
@@ -940,6 +1061,7 @@ func streamAnthropic(w http.ResponseWriter, body io.Reader, model string) {
 	toolIndexes := map[int]int{}
 	var tools []streamedResponseToolCall
 	var reasoning strings.Builder
+	usage := tokenUsage{}
 	s := bufio.NewScanner(body)
 	for s.Scan() {
 		line := strings.TrimSpace(s.Text())
@@ -951,6 +1073,9 @@ func streamAnthropic(w http.ResponseWriter, body io.Reader, model string) {
 			break
 		}
 		chunk := parseOpenAIStreamChunk([]byte(data))
+		if chunk.Usage.Present {
+			usage = chunk.Usage
+		}
 		if chunk.ReasoningContent != "" {
 			reasoning.WriteString(chunk.ReasoningContent)
 		}
@@ -1014,7 +1139,8 @@ func streamAnthropic(w http.ResponseWriter, body io.Reader, model string) {
 	if len(tools) > 0 {
 		stopReason = "tool_use"
 	}
-	fmt.Fprintf(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":%q,\"stop_sequence\":null},\"usage\":{\"output_tokens\":0}}\n\n", stopReason)
+	usageJSON, _ := json.Marshal(anthropicDeltaUsage(usage))
+	fmt.Fprintf(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":%q,\"stop_sequence\":null},\"usage\":%s}\n\n", stopReason, usageJSON)
 	fmt.Fprint(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
 }
 
@@ -1040,6 +1166,7 @@ func writeAnthropicResponse(w http.ResponseWriter, body io.Reader, model string)
 				Content string `json:"content"`
 			} `json:"message"`
 		} `json:"choices"`
+		Usage json.RawMessage `json:"usage"`
 	}
 	_ = json.NewDecoder(body).Decode(&v)
 	text := ""
@@ -1047,7 +1174,7 @@ func writeAnthropicResponse(w http.ResponseWriter, body io.Reader, model string)
 		text = v.Choices[0].Message.Content
 	}
 	w.Header().Set("Content-Type", "application/json")
-	_ = json.NewEncoder(w).Encode(map[string]any{"id": "ocgo", "type": "message", "role": "assistant", "model": model, "content": []map[string]string{{"type": "text", "text": text}}, "stop_reason": "end_turn", "usage": map[string]int{"input_tokens": 0, "output_tokens": 0}})
+	_ = json.NewEncoder(w).Encode(map[string]any{"id": "ocgo", "type": "message", "role": "assistant", "model": model, "content": []map[string]string{{"type": "text", "text": text}}, "stop_reason": "end_turn", "usage": anthropicUsage(usageFromJSON(v.Usage))})
 }
 
 func streamResponses(w http.ResponseWriter, body io.Reader, model string) {
@@ -1064,6 +1191,7 @@ func streamResponses(w http.ResponseWriter, body io.Reader, model string) {
 	nextOutputIndex := 0
 	var text strings.Builder
 	var reasoning strings.Builder
+	usage := tokenUsage{}
 	toolIndexes := map[int]int{}
 	var tools []streamedResponseToolCall
 	s := bufio.NewScanner(body)
@@ -1077,6 +1205,9 @@ func streamResponses(w http.ResponseWriter, body io.Reader, model string) {
 			break
 		}
 		chunk := parseOpenAIStreamChunk([]byte(data))
+		if chunk.Usage.Present {
+			usage = chunk.Usage
+		}
 		if chunk.ReasoningContent != "" {
 			reasoning.WriteString(chunk.ReasoningContent)
 		}
@@ -1144,7 +1275,7 @@ func streamResponses(w http.ResponseWriter, body io.Reader, model string) {
 		writeResponseEvent(w, "response.output_item.done", map[string]any{"type": "response.output_item.done", "output_index": tool.OutputIndex, "item": item})
 		output = append(output, item)
 	}
-	writeResponseEvent(w, "response.completed", map[string]any{"type": "response.completed", "response": map[string]any{"id": id, "object": "response", "model": model, "status": "completed", "output": output}})
+	writeResponseEvent(w, "response.completed", map[string]any{"type": "response.completed", "response": map[string]any{"id": id, "object": "response", "model": model, "status": "completed", "output": output, "usage": responsesUsage(usage)}})
 }
 
 type streamedResponseToolCall struct {
@@ -1163,6 +1294,7 @@ type openAIStreamChunk struct {
 	Content          string
 	ReasoningContent string
 	ToolCalls        []openAIStreamToolCall
+	Usage            tokenUsage
 }
 
 func parseOpenAIStreamChunk(data []byte) openAIStreamChunk {
@@ -1181,13 +1313,16 @@ func parseOpenAIStreamChunk(data []byte) openAIStreamChunk {
 				} `json:"tool_calls"`
 			} `json:"delta"`
 		} `json:"choices"`
+		Usage json.RawMessage `json:"usage"`
 	}
 	_ = json.Unmarshal(data, &v)
+	out := openAIStreamChunk{Usage: usageFromJSON(v.Usage)}
 	if len(v.Choices) == 0 {
-		return openAIStreamChunk{}
+		return out
 	}
 	delta := v.Choices[0].Delta
-	out := openAIStreamChunk{Content: delta.Content, ReasoningContent: delta.ReasoningContent}
+	out.Content = delta.Content
+	out.ReasoningContent = delta.ReasoningContent
 	for _, tc := range delta.ToolCalls {
 		out.ToolCalls = append(out.ToolCalls, openAIStreamToolCall{Index: tc.Index, ID: tc.ID, Name: tc.Function.Name, Arguments: tc.Function.Arguments})
 	}
@@ -1208,6 +1343,7 @@ func writeResponsesResponse(w http.ResponseWriter, body io.Reader, model string)
 				ToolCalls        []OAIToolCall `json:"tool_calls"`
 			} `json:"message"`
 		} `json:"choices"`
+		Usage json.RawMessage `json:"usage"`
 	}
 	_ = json.NewDecoder(body).Decode(&v)
 	text := ""
@@ -1225,7 +1361,7 @@ func writeResponsesResponse(w http.ResponseWriter, body io.Reader, model string)
 		output = append(output, map[string]any{"id": "msg_ocgo", "type": "message", "role": "assistant", "content": []map[string]string{{"type": "output_text", "text": text}}})
 	}
 	w.Header().Set("Content-Type", "application/json")
-	_ = json.NewEncoder(w).Encode(map[string]any{"id": "resp_ocgo", "object": "response", "created_at": time.Now().Unix(), "model": model, "status": "completed", "output": output, "usage": map[string]int{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}})
+	_ = json.NewEncoder(w).Encode(map[string]any{"id": "resp_ocgo", "object": "response", "created_at": time.Now().Unix(), "model": model, "status": "completed", "output": output, "usage": responsesUsage(usageFromJSON(v.Usage))})
 }
 
 func countTokens(w http.ResponseWriter, r *http.Request) {

cmd/ocgo/main_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
@@ -223,6 +224,39 @@ func TestRawChatImageRejectsUnsupportedModel(t *testing.T) {
 	}
 }
 
+func TestRawChatStreamRequestsUsage(t *testing.T) {
+	body, err := prepareChatBody([]byte(`{"model":"kimi-k2.6","stream":true,"stream_options":{"foo":"bar"},"messages":[{"role":"user","content":"hello"}]}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	var req map[string]any
+	if err := json.Unmarshal(body, &req); err != nil {
+		t.Fatal(err)
+	}
+	options, ok := req["stream_options"].(map[string]any)
+	if !ok {
+		t.Fatalf("missing stream options in %s", string(body))
+	}
+	if options["include_usage"] != true || options["foo"] != "bar" {
+		t.Fatalf("bad stream options: %+v", options)
+	}
+}
+
+func TestConvertedStreamingRequestsAskForUsage(t *testing.T) {
+	anthropic := convertRequest(AnthropicRequest{Model: "kimi-k2.6", Stream: true, Messages: []AMessage{{Role: "user", Content: []byte(`hello`)}}})
+	if anthropic.StreamOptions == nil || !anthropic.StreamOptions.IncludeUsage {
+		t.Fatalf("anthropic conversion should request stream usage: %+v", anthropic.StreamOptions)
+	}
+	responses := responsesToChat(ResponsesRequest{Model: "kimi-k2.6", Stream: true, Input: []byte(`"hello"`)})
+	if responses.StreamOptions == nil || !responses.StreamOptions.IncludeUsage {
+		t.Fatalf("responses conversion should request stream usage: %+v", responses.StreamOptions)
+	}
+	plain := responsesToChat(ResponsesRequest{Model: "kimi-k2.6", Input: []byte(`"hello"`)})
+	if plain.StreamOptions != nil {
+		t.Fatalf("non-streaming conversion should not set stream options: %+v", plain.StreamOptions)
+	}
+}
+
 func TestAnthropicContentPreservesImages(t *testing.T) {
 	messages := contentToOpenAI(AMessage{Role: "user", Content: []byte(`[{"type":"text","text":"what is this?"},{"type":"image","source":{"type":"base64","media_type":"image/jpeg","data":"abc"}}]`)})
 	if len(messages) != 1 {
@@ -289,6 +323,88 @@ func TestParseWindowsNetstatPIDMatchesIPv6(t *testing.T) {
 	}
 }
 
+func TestWriteAnthropicResponseIncludesUsage(t *testing.T) {
+	body := strings.NewReader(`{"choices":[{"message":{"content":"done"}}],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16,"prompt_tokens_details":{"cached_tokens":4}}}`)
+	w := httptest.NewRecorder()
+	writeAnthropicResponse(w, body, "kimi-k2.6")
+	var out map[string]any
+	if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil {
+		t.Fatal(err)
+	}
+	usage, ok := out["usage"].(map[string]any)
+	if !ok {
+		t.Fatalf("missing usage: %+v", out)
+	}
+	if usage["input_tokens"] != float64(11) || usage["output_tokens"] != float64(5) || usage["cache_read_input_tokens"] != float64(4) {
+		t.Fatalf("bad anthropic usage: %+v", usage)
+	}
+}
+
+func TestWriteResponsesResponseIncludesUsage(t *testing.T) {
+	body := strings.NewReader(`{"choices":[{"message":{"content":"done"}}],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10,"input_tokens_details":{"cached_tokens":2}}}`)
+	w := httptest.NewRecorder()
+	writeResponsesResponse(w, body, "kimi-k2.6")
+	var out map[string]any
+	if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil {
+		t.Fatal(err)
+	}
+	usage, ok := out["usage"].(map[string]any)
+	if !ok {
+		t.Fatalf("missing usage: %+v", out)
+	}
+	if usage["input_tokens"] != float64(7) || usage["output_tokens"] != float64(3) || usage["total_tokens"] != float64(10) {
+		t.Fatalf("bad responses usage: %+v", usage)
+	}
+	details, ok := usage["input_tokens_details"].(map[string]any)
+	if !ok || details["cached_tokens"] != float64(2) {
+		t.Fatalf("bad cached details: %+v", usage["input_tokens_details"])
+	}
+}
+
+func TestParseOpenAIStreamChunkReadsUsageOnly(t *testing.T) {
+	chunk := parseOpenAIStreamChunk([]byte(`{"choices":[],"usage":{"prompt_tokens":8,"completion_tokens":4,"total_tokens":12}}`))
+	if !chunk.Usage.Present || chunk.Usage.InputTokens != 8 || chunk.Usage.OutputTokens != 4 || chunk.Usage.TotalTokens != 12 {
+		t.Fatalf("bad stream usage: %+v", chunk.Usage)
+	}
+	if chunk.Content != "" || len(chunk.ToolCalls) != 0 {
+		t.Fatalf("usage-only chunk should not include deltas: %+v", chunk)
+	}
+}
+
+func TestStreamAnthropicIncludesFinalUsage(t *testing.T) {
+	body := strings.NewReader(strings.Join([]string{
+		`data: {"choices":[{"delta":{"content":"hi"}}]}`,
+		`data: {"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}`,
+		`data: [DONE]`,
+		``,
+	}, "\n\n"))
+	w := httptest.NewRecorder()
+	streamAnthropic(w, body, "kimi-k2.6")
+	out := w.Body.String()
+	for _, want := range []string{`"input_tokens":7`, `"output_tokens":3`} {
+		if !strings.Contains(out, want) {
+			t.Fatalf("missing %q in:\n%s", want, out)
+		}
+	}
+}
+
+func TestStreamResponsesIncludesCompletedUsage(t *testing.T) {
+	body := strings.NewReader(strings.Join([]string{
+		`data: {"choices":[{"delta":{"content":"hi"}}]}`,
+		`data: {"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}`,
+		`data: [DONE]`,
+		``,
+	}, "\n\n"))
+	w := httptest.NewRecorder()
+	streamResponses(w, body, "kimi-k2.6")
+	out := w.Body.String()
+	for _, want := range []string{`event: response.completed`, `"input_tokens":7`, `"output_tokens":3`, `"total_tokens":10`} {
+		if !strings.Contains(out, want) {
+			t.Fatalf("missing %q in:\n%s", want, out)
+		}
+	}
+}
+
 func TestStreamAnthropicForwardsToolCalls(t *testing.T) {
 	reasoningContentCache.Lock()
 	reasoningContentCache.byCallID = map[string]string{}