fix(translator): keep tool arguments when openai-compat upstreams send per-chunk usage

internal/translator/openai/claude/openai_claude_response.go

@@ -46,6 +46,15 @@ type ConvertOpenAIResponseToAnthropicParams struct {
 	ContentBlocksStopped bool
 	// Track if message_delta has been sent
 	MessageDeltaSent bool
+	// Usage buffered from the most recent chunk seen before the stream
+	// terminates. Upstreams that attach usage to every chunk must not end the
+	// stream on a non-terminal chunk; the buffered values are emitted when
+	// finish_reason arrives, or at [DONE] as a fallback.
+	LastUsageSeen        bool
+	LastInputTokens      int64
+	LastOutputTokens     int64
+	LastCachedTokens     int64
+	LastCacheWriteTokens int64
 	// Track if message_start has been sent
 	MessageStarted bool
 	// Track if message_stop has been sent
@@ -97,6 +106,11 @@ func ConvertOpenAIResponseToClaude(_ context.Context, _ string, originalRequestR
 			FinishReason:                "",
 			ContentBlocksStopped:        false,
 			MessageDeltaSent:            false,
+			LastUsageSeen:               false,
+			LastInputTokens:             0,
+			LastOutputTokens:            0,
+			LastCachedTokens:            0,
+			LastCacheWriteTokens:        0,
 			ToolCallBlockIndexes:        make(map[int]int),
 			TextContentBlockIndex:       -1,
 			ThinkingContentBlockIndex:   -1,
@@ -306,13 +320,31 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
 
 	// Handle usage information separately (this comes in a later chunk)
 	// Only process if usage has actual values (not null)
-	if !param.MessageDeltaSent && (param.FinishReason != "" || param.SawToolCall) {
+	if !param.MessageDeltaSent {
 		usage := root.Get("usage")
-		if usage.Exists() && usage.Type != gjson.Null {
+		hasUsage := usage.Exists() && usage.Type != gjson.Null
+		// Buffered values only exist on per-chunk-usage upstreams. Without
+		// them a usage-less finish_reason chunk must keep waiting for the
+		// usage-only chunk that OpenAI's include_usage mode sends afterwards.
+		if param.FinishReason != "" && (hasUsage || param.LastUsageSeen) {
+			// The stream is terminal: close it on this chunk. Prefer the
+			// chunk's own usage, then the buffered values, then zeros.
+			var inputTokens, outputTokens, cachedTokens, cacheWriteTokens int64
+			if hasUsage {
+				inputTokens, outputTokens, cachedTokens, cacheWriteTokens = extractOpenAIUsage(usage)
+			} else {
+				inputTokens, outputTokens, cachedTokens, cacheWriteTokens = param.LastInputTokens, param.LastOutputTokens, param.LastCachedTokens, param.LastCacheWriteTokens
+			}
 			finalizeOpenAIAnthropicContentBlocks(param, &results)
-			inputTokens, outputTokens, cachedTokens, cacheWriteTokens := extractOpenAIUsage(usage)
 			emitAnthropicMessageDelta(param, &results, inputTokens, outputTokens, cachedTokens, cacheWriteTokens)
 			emitMessageStopIfNeeded(param, &results)
+		} else if hasUsage {
+			// The stream has not terminated yet. Ending it here would drop
+			// tool call arguments that arrive in later chunks, so buffer
+			// the latest usage values instead; finish_reason or [DONE]
+			// emits them.
+			param.LastUsageSeen = true
+			param.LastInputTokens, param.LastOutputTokens, param.LastCachedTokens, param.LastCacheWriteTokens = extractOpenAIUsage(usage)
 		}
 	}
 
@@ -326,7 +358,11 @@ func convertOpenAIDoneToAnthropic(param *ConvertOpenAIResponseToAnthropicParams)
 	finalizeOpenAIAnthropicContentBlocks(param, &results)
 
 	if !param.MessageDeltaSent {
-		emitAnthropicMessageDelta(param, &results, 0, 0, 0, 0)
+		inputTokens, outputTokens, cachedTokens, cacheWriteTokens := int64(0), int64(0), int64(0), int64(0)
+		if param.LastUsageSeen {
+			inputTokens, outputTokens, cachedTokens, cacheWriteTokens = param.LastInputTokens, param.LastOutputTokens, param.LastCachedTokens, param.LastCacheWriteTokens
+		}
+		emitAnthropicMessageDelta(param, &results, inputTokens, outputTokens, cachedTokens, cacheWriteTokens)
 	}
 
 	emitMessageStopIfNeeded(param, &results)

internal/translator/openai/claude/openai_claude_response_test.go

@@ -38,7 +38,10 @@ func runStream(t *testing.T, originalReq string, chunks ...string) []sseEvent {
 		[]byte("data: [DONE]"),
 		&paramAny,
 	)...)
+	return parseSSEEvents(emitted)
+}
 
+func parseSSEEvents(emitted [][]byte) []sseEvent {
 	var events []sseEvent
 	for _, raw := range emitted {
 		s := string(raw)
@@ -740,3 +743,147 @@ func TestNonStreamingUsage_PreservesCacheWriteTokens(t *testing.T) {
 		})
 	}
 }
+
+// Upstreams that attach usage to every streaming chunk (e.g. vLLM) must not
+// have their tool call arguments dropped: the first tool chunk carries only
+// id and name, so finalizing the stream on it loses every argument delta.
+func TestStreamingTool_PerChunkUsageKeepsToolArguments(t *testing.T) {
+	events := runStream(t, streamReq,
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","function":{"name":"get_weather"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`,
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"loc\":\"Par"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":7}}`,
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"is\"}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":9}}`,
+	)
+
+	var args strings.Builder
+	inputDeltaIdx, blockStopIdx := -1, -1
+	for i, e := range events {
+		switch {
+		case e.Type == "content_block_delta" && gjson.Get(e.Payload, "delta.type").String() == "input_json_delta":
+			args.WriteString(gjson.Get(e.Payload, "delta.partial_json").String())
+			if inputDeltaIdx < 0 {
+				inputDeltaIdx = i
+			}
+		case e.Type == "content_block_stop" && blockStopIdx < 0:
+			blockStopIdx = i
+		}
+	}
+	if got := args.String(); got != `{"loc":"Paris"}` {
+		t.Fatalf("tool arguments lost: got %q, want %q", got, `{"loc":"Paris"}`)
+	}
+	if inputDeltaIdx < 0 || blockStopIdx < 0 || inputDeltaIdx > blockStopIdx {
+		t.Fatalf("input_json_delta must precede content_block_stop (delta=%d stop=%d)", inputDeltaIdx, blockStopIdx)
+	}
+	if got := countByType(events, "message_delta"); got != 1 {
+		t.Fatalf("expected exactly one message_delta, got %d (events=%+v)", got, events)
+	}
+	if got := countByType(events, "message_stop"); got != 1 {
+		t.Fatalf("expected exactly one message_stop, got %d (events=%+v)", got, events)
+	}
+	if got := lastStopReason(events); got != "tool_use" {
+		t.Fatalf("stop_reason = %q, want %q", got, "tool_use")
+	}
+	var deltaEvent *sseEvent
+	for i := range events {
+		if events[i].Type == "message_delta" {
+			deltaEvent = &events[i]
+		}
+	}
+	if deltaEvent == nil {
+		t.Fatalf("missing message_delta event")
+	}
+	if output := gjson.Get(deltaEvent.Payload, "usage.output_tokens").Int(); output != 9 {
+		t.Fatalf("output_tokens = %d, want 9 (usage from the final chunk)", output)
+	}
+}
+
+// With per-chunk usage and no finish_reason at all, the stream must still end
+// via [DONE] and report the latest buffered usage, not the first chunk's.
+func TestStreamingTool_PerChunkUsageWithoutFinishReasonUsesBufferedUsage(t *testing.T) {
+	events := runStream(t, streamReq,
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","function":{"name":"get_weather","arguments":"{\"loc\":\"Paris\"}"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`,
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{}}],"usage":{"prompt_tokens":10,"completion_tokens":8}}`,
+	)
+
+	var args strings.Builder
+	for _, e := range events {
+		if e.Type == "content_block_delta" && gjson.Get(e.Payload, "delta.type").String() == "input_json_delta" {
+			args.WriteString(gjson.Get(e.Payload, "delta.partial_json").String())
+		}
+	}
+	if got := args.String(); got != `{"loc":"Paris"}` {
+		t.Fatalf("tool arguments lost: got %q, want %q", got, `{"loc":"Paris"}`)
+	}
+	if got := countByType(events, "message_delta"); got != 1 {
+		t.Fatalf("expected exactly one message_delta, got %d (events=%+v)", got, events)
+	}
+	if got := lastStopReason(events); got != "tool_use" {
+		t.Fatalf("stop_reason = %q, want %q", got, "tool_use")
+	}
+	var deltaEvent *sseEvent
+	for i := range events {
+		if events[i].Type == "message_delta" {
+			deltaEvent = &events[i]
+		}
+	}
+	if deltaEvent == nil {
+		t.Fatalf("missing message_delta event")
+	}
+	if output := gjson.Get(deltaEvent.Payload, "usage.output_tokens").Int(); output != 8 {
+		t.Fatalf("output_tokens = %d, want 8 (latest buffered usage)", output)
+	}
+	if got := countByType(events, "message_stop"); got != 1 {
+		t.Fatalf("expected exactly one message_stop, got %d (events=%+v)", got, events)
+	}
+}
+
+// When finish_reason arrives on its own chunk (without usage), the stream is
+// terminal: the buffered usage must be flushed and message_delta/message_stop
+// emitted on that chunk, not deferred to [DONE], which compatible upstreams
+// may delay or omit.
+func TestStreamingTool_FinishReasonChunkFlushesBufferedUsage(t *testing.T) {
+	chunks := []string{
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","function":{"name":"get_weather"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`,
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"loc\":\"Paris\"}"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":7}}`,
+		`{"id":"c1","model":"m","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`,
+	}
+	var paramAny any
+	chunkEvents := make([][]sseEvent, 0, len(chunks))
+	for _, chunk := range chunks {
+		emitted := ConvertOpenAIResponseToClaude(context.Background(), "", []byte(streamReq), nil, []byte("data: "+chunk), &paramAny)
+		chunkEvents = append(chunkEvents, parseSSEEvents(emitted))
+	}
+
+	terminal := chunkEvents[2]
+	var deltaEvent *sseEvent
+	for i := range terminal {
+		if terminal[i].Type == "message_delta" {
+			deltaEvent = &terminal[i]
+		}
+	}
+	if deltaEvent == nil {
+		t.Fatalf("finish_reason chunk must emit message_delta immediately (events=%+v)", terminal)
+	}
+	if output := gjson.Get(deltaEvent.Payload, "usage.output_tokens").Int(); output != 7 {
+		t.Fatalf("output_tokens = %d, want 7 (buffered from the last usage chunk)", output)
+	}
+	if got := gjson.Get(deltaEvent.Payload, "delta.stop_reason").String(); got != "tool_use" {
+		t.Fatalf("stop_reason = %q, want tool_use", got)
+	}
+	hasStop := false
+	for _, e := range terminal {
+		if e.Type == "message_stop" {
+			hasStop = true
+		}
+	}
+	if !hasStop {
+		t.Fatalf("finish_reason chunk must emit message_stop immediately (events=%+v)", terminal)
+	}
+
+	// [DONE] must not duplicate the closing events.
+	doneEvents := parseSSEEvents(ConvertOpenAIResponseToClaude(context.Background(), "", []byte(streamReq), nil, []byte("data: [DONE]"), &paramAny))
+	for _, e := range doneEvents {
+		if e.Type == "message_delta" || e.Type == "message_stop" {
+			t.Fatalf("[DONE] must not re-emit %s (events=%+v)", e.Type, doneEvents)
+		}
+	}
+}