From ab61b55b6e0cb6d8621e2697a5de5937939823ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=94=A1=E5=8F=8A?= <522caiji@gmail.com> Date: Sun, 18 Oct 2026 19:20:33 +0800 Subject: [PATCH] fix(providers): preserve prompt cache tokens and salvage terminal streams Normalize cache_read_tokens / prompt_cache_hit_tokens / prompt_tokens_details across Trae, Qoder, and WorkBuddy, and emit [DONE] when an upstream stream already finished with a terminal finish_reason. Classify Devin expired-session and Connect auth denials as auth, and surface login_required instead of cooling. --- .../prompt-cache-and-stream-salvage.md | 9 ++ internal/executor/chat.go | 31 +++++-- internal/executor/chat_test.go | 20 +++++ internal/gateway/openai_stream.go | 53 ++++++++++-- internal/gateway/responses_test.go | 17 ++++ internal/gateway/usage.go | 1 + internal/providers/devin/devin_test.go | 29 +++++++ internal/providers/devin/errors.go | 29 ++++++- internal/providers/qoder/adapter.go | 25 ++++-- internal/providers/qoder/cache_test.go | 24 ++++++ internal/providers/qoder/direct.go | 85 +++++++++++++++++-- internal/providers/qoder/direct_test.go | 23 +++++ internal/providers/trae/cache_test.go | 25 ++++++ internal/providers/trae/client.go | 19 ++++- internal/providers/trae/client_test.go | 3 + internal/providers/trae/sse.go | 52 +++++++++++- internal/providers/workbuddy/cache_test.go | 23 +++++ internal/providers/workbuddy/client.go | 18 ++-- internal/providers/workbuddy/client_test.go | 1 + internal/providers/workbuddy/sse.go | 23 +++++ internal/runtime/alias.go | 1 + internal/runtime/manager_probe.go | 12 ++- internal/store/store.go | 6 +- 23 files changed, 477 insertions(+), 52 deletions(-) create mode 100644 changelog/unreleased/prompt-cache-and-stream-salvage.md create mode 100644 internal/providers/qoder/cache_test.go create mode 100644 internal/providers/trae/cache_test.go create mode 100644 internal/providers/workbuddy/cache_test.go diff --git a/changelog/unreleased/prompt-cache-and-stream-salvage.md b/changelog/unreleased/prompt-cache-and-stream-salvage.md new file mode 100644 index 00000000..4dc0d6b8 --- /dev/null +++ b/changelog/unreleased/prompt-cache-and-stream-salvage.md @@ -0,0 +1,9 @@ +### English + +- Preserve and normalize prompt cache token metrics (`cache_read_tokens`, `prompt_cache_hit_tokens`, and `prompt_tokens_details`) across Trae, Qoder, and WorkBuddy providers in both streaming and non-streaming responses. +- Gracefully salvage upstream streams that finish with a terminal `finish_reason` before EOF without emitting literal `[DONE]`, preventing artificial 502 incomplete stream errors. + +### 中文 + +- 修复并标准化 Trae、Qoder 与 WorkBuddy 提供商在流式与非流式响应中的 Prompt 缓存 Token 指标（统一解析并透传 `cache_read_tokens`、`prompt_cache_hit_tokens` 与 `prompt_tokens_details`）。 +- 容错处理上游流式响应在带有终态 `finish_reason` 正常结束但未显式发送 `[DONE]` 的边界情况，自动补齐结束帧，避免下游误判为 502 流异常中断。 diff --git a/internal/executor/chat.go b/internal/executor/chat.go index 2070f4e9..cadf0212 100644 --- a/internal/executor/chat.go +++ b/internal/executor/chat.go @@ -981,6 +981,7 @@ func (e ChatExecutor) chatInProcessNonStreamAttempt(ctx context.Context, item It CompletionTokens: outcome.CompletionTokens, CacheReadTokens: outcome.CacheReadTokens, CacheWriteTokens: outcome.CacheWriteTokens, + CachedTokens: outcome.CacheReadTokens, UsageSource: outcome.UsageSource, ConsumedCredits: outcome.Credits, AccountID: item.ID, @@ -1075,13 +1076,14 @@ func decodeChatResult(req translate.ChatRequest, body []byte) (ChatResult, error var parsed struct { Model string `json:"model"` Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - CacheReadTokens *int `json:"cache_read_tokens"` - CacheWriteTokens *int `json:"cache_write_tokens"` - Source string `json:"source"` - Credits *float64 `json:"credits"` - PromptDetails struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CacheReadTokens *int `json:"cache_read_tokens"` + CacheWriteTokens *int `json:"cache_write_tokens"` + PromptCacheHitTokens *int `json:"prompt_cache_hit_tokens"` + Source string `json:"source"` + Credits *float64 `json:"credits"` + PromptDetails struct { CachedTokens *int `json:"cached_tokens"` } `json:"prompt_tokens_details"` } `json:"usage"` @@ -1119,6 +1121,17 @@ func decodeChatResult(req translate.ChatRequest, body []byte) (ChatResult, error if source == "" { source = "estimate" } + cacheRead := parsed.Usage.CacheReadTokens + if cacheRead == nil { + cacheRead = parsed.Usage.PromptDetails.CachedTokens + } + if cacheRead == nil { + cacheRead = parsed.Usage.PromptCacheHitTokens + } + cachedTokens := parsed.Usage.PromptDetails.CachedTokens + if cachedTokens == nil { + cachedTokens = cacheRead + } return ChatResult{ Model: model, Content: content, @@ -1127,9 +1140,9 @@ func decodeChatResult(req translate.ChatRequest, body []byte) (ChatResult, error FinishReason: finishReason, PromptTokens: parsed.Usage.PromptTokens, CompletionTokens: parsed.Usage.CompletionTokens, - CacheReadTokens: parsed.Usage.CacheReadTokens, + CacheReadTokens: cacheRead, CacheWriteTokens: parsed.Usage.CacheWriteTokens, - CachedTokens: parsed.Usage.PromptDetails.CachedTokens, + CachedTokens: cachedTokens, UsageSource: source, Credits: parsed.Usage.Credits, }, nil diff --git a/internal/executor/chat_test.go b/internal/executor/chat_test.go index 343b5bfb..3a2f1d59 100644 --- a/internal/executor/chat_test.go +++ b/internal/executor/chat_test.go @@ -41,6 +41,26 @@ func TestDecodeChatResultPreservesPromptCacheUsage(t *testing.T) { } } +func TestDecodeChatResultFallsBackToPromptCacheHitTokens(t *testing.T) { + result, err := decodeChatResult(translate.ChatRequest{Model: "deepseek-v4.1"}, []byte(`{ + "model":"deepseek-v4.1", + "choices":[{"message":{"content":"OK"},"finish_reason":"stop"}], + "usage":{ + "prompt_tokens":100, + "completion_tokens":20, + "prompt_cache_hit_tokens":80, + "source":"upstream" + } + }`)) + if err != nil { + t.Fatal(err) + } + if result.CacheReadTokens == nil || *result.CacheReadTokens != 80 || + result.CachedTokens == nil || *result.CachedTokens != 80 { + t.Fatalf("cache usage = read:%v cached:%v", result.CacheReadTokens, result.CachedTokens) + } +} + func TestChatNonStreamFailoversRateLimit(t *testing.T) { var hitsA atomic.Int32 a := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/internal/gateway/openai_stream.go b/internal/gateway/openai_stream.go index 6b3977b4..952884cf 100644 --- a/internal/gateway/openai_stream.go +++ b/internal/gateway/openai_stream.go @@ -112,6 +112,9 @@ func RelayOpenAIStream(w http.ResponseWriter, body io.Reader) (stats StreamRelay now := time.Now() stats.FirstTokenAt = &now } + if fr := SSEDeltaFinishReason(line); fr != "" { + stats.FinishReason = fr + } if usage, ok := ParseStreamUsageLine(line); ok { usage.FirstTokenAt = stats.FirstTokenAt usage.SSEEventCount = stats.SSEEventCount @@ -163,6 +166,14 @@ func RelayOpenAIStream(w http.ResponseWriter, body io.Reader) (stats StreamRelay return stats, streamErr } if !sawDone { + if stats.FinishReason != "" && (stats.FirstTokenAt != nil || stats.SSEEventCount > 0) { + sawDone = true + stats.SawDone = true + if _, err := io.WriteString(writer, "data: [DONE]\n\n"); err != nil { + return stats, &StreamRelayWriteError{err: err} + } + return stats, nil + } streamErr := executor.StreamIncompleteError() if writeErr := writeStructuredStreamError(writer, streamErr); writeErr != nil { return stats, writeErr @@ -358,6 +369,22 @@ func IsStreamClientDisconnect(err error) bool { return errors.As(err, &writeErr) } +func SSEDeltaFinishReason(line string) string { + payload := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if payload == "" || payload == "[DONE]" { + return "" + } + var parsed struct { + Choices []struct { + FinishReason string `json:"finish_reason"` + } `json:"choices"` + } + if json.Unmarshal([]byte(payload), &parsed) != nil || len(parsed.Choices) == 0 { + return "" + } + return parsed.Choices[0].FinishReason +} + func SSEDeltaHasToken(line string) bool { payload := strings.TrimSpace(strings.TrimPrefix(line, "data:")) if payload == "" || payload == "[DONE]" { @@ -413,14 +440,15 @@ func ParseStreamUsageLine(line string) (StreamRelayStats, bool) { var parsed struct { Model string `json:"model"` Usage *struct { - PromptTokens *int `json:"prompt_tokens"` - CompletionTokens *int `json:"completion_tokens"` - CacheReadTokens *int `json:"cache_read_tokens"` - CacheWriteTokens *int `json:"cache_write_tokens"` - Source string `json:"source"` - Credits *float64 `json:"credits"` - Credit *float64 `json:"credit"` - PromptDetails struct { + PromptTokens *int `json:"prompt_tokens"` + CompletionTokens *int `json:"completion_tokens"` + CacheReadTokens *int `json:"cache_read_tokens"` + CacheWriteTokens *int `json:"cache_write_tokens"` + PromptCacheHitTokens *int `json:"prompt_cache_hit_tokens"` + Source string `json:"source"` + Credits *float64 `json:"credits"` + Credit *float64 `json:"credit"` + PromptDetails struct { CachedTokens *int `json:"cached_tokens"` } `json:"prompt_tokens_details"` } `json:"usage"` @@ -436,12 +464,19 @@ func ParseStreamUsageLine(line string) (StreamRelayStats, bool) { if cacheRead == nil { cacheRead = parsed.Usage.PromptDetails.CachedTokens } + if cacheRead == nil { + cacheRead = parsed.Usage.PromptCacheHitTokens + } + cachedTokens := parsed.Usage.PromptDetails.CachedTokens + if cachedTokens == nil { + cachedTokens = cacheRead + } return StreamRelayStats{ PromptTokens: parsed.Usage.PromptTokens, CompletionTokens: parsed.Usage.CompletionTokens, CacheReadTokens: cacheRead, CacheWriteTokens: parsed.Usage.CacheWriteTokens, - CachedTokens: parsed.Usage.PromptDetails.CachedTokens, + CachedTokens: cachedTokens, UsageSource: firstNonEmpty(parsed.Usage.Source, "estimate"), Credits: credits, Model: parsed.Model, diff --git a/internal/gateway/responses_test.go b/internal/gateway/responses_test.go index 1a7a962c..d13b6b1c 100644 --- a/internal/gateway/responses_test.go +++ b/internal/gateway/responses_test.go @@ -128,6 +128,22 @@ func TestResponsesUsagePreservesCachedInputTokens(t *testing.T) { } } +func TestRelayOpenAIStreamSalvagesFinishReasonWithoutDone(t *testing.T) { + stream := "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n" + + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n" + rec := httptest.NewRecorder() + stats, err := RelayOpenAIStream(rec, strings.NewReader(stream)) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if !stats.SawDone { + t.Fatalf("expected SawDone=true") + } + if !strings.Contains(rec.Body.String(), "data: [DONE]") { + t.Fatalf("expected data: [DONE] in output: %s", rec.Body.String()) + } +} + func TestParseStreamUsageLineCacheReadFallback(t *testing.T) { for _, test := range []struct { name string @@ -137,6 +153,7 @@ func TestParseStreamUsageLineCacheReadFallback(t *testing.T) { {name: "detail only", usage: `"prompt_tokens_details":{"cached_tokens":2176}`, want: ptrInt(2176)}, {name: "explicit zero", usage: `"prompt_tokens_details":{"cached_tokens":0}`, want: ptrInt(0)}, {name: "top-level wins", usage: `"cache_read_tokens":12,"prompt_tokens_details":{"cached_tokens":2176}`, want: ptrInt(12)}, + {name: "prompt cache hit tokens", usage: `"prompt_cache_hit_tokens":80`, want: ptrInt(80)}, {name: "unknown stays absent", usage: `"prompt_tokens":16`, want: nil}, } { t.Run(test.name, func(t *testing.T) { diff --git a/internal/gateway/usage.go b/internal/gateway/usage.go index de60040f..00496f2b 100644 --- a/internal/gateway/usage.go +++ b/internal/gateway/usage.go @@ -27,6 +27,7 @@ func BuildChatUsage(res executor.ChatResult) map[string]any { } if cachedTokens != nil { out["prompt_tokens_details"] = map[string]any{"cached_tokens": *cachedTokens} + out["prompt_cache_hit_tokens"] = *cachedTokens } if res.Credits != nil { out["credits"] = *res.Credits diff --git a/internal/providers/devin/devin_test.go b/internal/providers/devin/devin_test.go index 62970d9e..b173c80a 100644 --- a/internal/providers/devin/devin_test.go +++ b/internal/providers/devin/devin_test.go @@ -744,6 +744,35 @@ func TestClassifyAuthForbiddenQuotaTrailer(t *testing.T) { } } +func TestClassifyExpiredTokenAndConnectAuthDenial(t *testing.T) { + got := Classify(401, `{"message":"token has expired"}`) + if got.Kind != accounts.KindAuth || got.Status != 401 { + t.Fatalf("expired token classify=%+v want kind=auth, status=401", got) + } + + body := "devin upstream error (permission_denied): an internal error occurred (trace ID: 3bd4ba1ff74ce0f9554a938c1157e80d)" + got = Classify(403, body) + if got.Kind != accounts.KindAuth || got.Status != 403 { + t.Fatalf("connect auth denial trailer classify=%+v want kind=auth, status=403", got) + } + + seatErr := `devin seat management error (status 403): {"code":"permission_denied","message":"an internal error occurred"}` + got = Classify(403, seatErr) + if got.Kind != accounts.KindAuth || got.Status != 403 { + t.Fatalf("seat management denial classify=%+v want kind=auth, status=403", got) + } + + got = Classify(400, "token has expired") + if got.Kind != accounts.KindAuth { + t.Fatalf("token expired text classify=%+v want kind=auth", got) + } + + got = Classify(403, "permission_denied") + if got.Kind != accounts.KindInvalidRequest { + t.Fatalf("bare permission_denied classify=%+v want kind=invalid_request", got) + } +} + func TestClassifyMCPConfigPermissionDenied(t *testing.T) { body := "devin upstream error (permission_denied): Unable to process request due to an MCP configuration issue. (trace ID: 3de0a0f4c7f0e2cd3bd8e001a46749bd)" got := Classify(403, body) diff --git a/internal/providers/devin/errors.go b/internal/providers/devin/errors.go index b3a1357b..25919b84 100644 --- a/internal/providers/devin/errors.go +++ b/internal/providers/devin/errors.go @@ -45,10 +45,7 @@ func Classify(status int, body string) providers.ClassifiedError { Status: 400, Message: firstNonEmpty(safeBody, "devin rejected MCP/hosted tools in the request"), } - case status == 401 || - strings.Contains(text, "unauthenticated") || - strings.Contains(text, "unauthorized") || - strings.Contains(text, "session dead"): + case isDevinAuthError(status, text): return providers.ClassifiedError{ Kind: accounts.KindAuth, Status: firstNonEmptyStatus(status, 401), @@ -217,6 +214,30 @@ func appendToolsDiag(message, toolsDiag string) string { return message + " | " + suffix } +func isDevinAuthError(status int, text string) bool { + if status == 401 || + strings.Contains(text, "unauthenticated") || + strings.Contains(text, "unauthorized") || + strings.Contains(text, "session dead") || + strings.Contains(text, "token has expired") || + strings.Contains(text, "token is expired") || + strings.Contains(text, "token expired") || + strings.Contains(text, "expired token") || + strings.Contains(text, "token invalid") || + strings.Contains(text, "invalid token") { + return true + } + // Codeium Connect-RPC and seat management return permission_denied with + // "an internal error occurred" when the session token is expired or invalid. + if strings.Contains(text, "permission_denied") && (strings.Contains(text, "internal error") || strings.Contains(text, "token")) { + return true + } + if status == 403 && (strings.Contains(text, "seat management error") || strings.Contains(text, "internal error")) { + return true + } + return false +} + func isDevinMCPConfigDenial(text string) bool { if text == "" { return false diff --git a/internal/providers/qoder/adapter.go b/internal/providers/qoder/adapter.go index 0362f6d3..9c736f23 100644 --- a/internal/providers/qoder/adapter.go +++ b/internal/providers/qoder/adapter.go @@ -674,12 +674,16 @@ func decodeChatOutcome(fallbackModel string, body []byte) (providers.ChatOutcome var parsed struct { Model string `json:"model"` Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - CacheReadTokens *int `json:"cache_read_tokens"` - CacheWriteTokens *int `json:"cache_write_tokens"` - Source string `json:"source"` - Credits *float64 `json:"credits"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CacheReadTokens *int `json:"cache_read_tokens"` + CacheWriteTokens *int `json:"cache_write_tokens"` + PromptCacheHitTokens *int `json:"prompt_cache_hit_tokens"` + Source string `json:"source"` + Credits *float64 `json:"credits"` + PromptDetails struct { + CachedTokens *int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` } `json:"usage"` Choices []struct { FinishReason string `json:"finish_reason"` @@ -693,11 +697,18 @@ func decodeChatOutcome(fallbackModel string, body []byte) (providers.ChatOutcome if err := json.Unmarshal(body, &parsed); err != nil { return providers.ChatOutcome{}, fmt.Errorf("decode worker response: %w", err) } + cacheRead := parsed.Usage.CacheReadTokens + if cacheRead == nil { + cacheRead = parsed.Usage.PromptDetails.CachedTokens + } + if cacheRead == nil { + cacheRead = parsed.Usage.PromptCacheHitTokens + } outcome := providers.ChatOutcome{ Model: parsed.Model, PromptTokens: parsed.Usage.PromptTokens, CompletionTokens: parsed.Usage.CompletionTokens, - CacheReadTokens: parsed.Usage.CacheReadTokens, + CacheReadTokens: cacheRead, CacheWriteTokens: parsed.Usage.CacheWriteTokens, UsageSource: parsed.Usage.Source, Credits: parsed.Usage.Credits, diff --git a/internal/providers/qoder/cache_test.go b/internal/providers/qoder/cache_test.go new file mode 100644 index 00000000..a0254173 --- /dev/null +++ b/internal/providers/qoder/cache_test.go @@ -0,0 +1,24 @@ +package qoder + +import ( + "encoding/json" + "testing" + + "github.com/caigee-cmd/cli2api/internal/providers" +) + +func TestQoderDirectPreservesPromptCacheHitTokens(t *testing.T) { + raw := json.RawMessage(`{"prompt_tokens":100,"completion_tokens":50,"prompt_tokens_details":{"cached_tokens":80}}`) + var out providers.ChatOutcome + applyUsage(&out, raw) + if out.CacheReadTokens == nil || *out.CacheReadTokens != 80 { + t.Fatalf("expected CacheReadTokens 80, got %v", out.CacheReadTokens) + } + + rawHit := json.RawMessage(`{"prompt_tokens":100,"completion_tokens":50,"prompt_cache_hit_tokens":75}`) + var outHit providers.ChatOutcome + applyUsage(&outHit, rawHit) + if outHit.CacheReadTokens == nil || *outHit.CacheReadTokens != 75 { + t.Fatalf("expected CacheReadTokens 75, got %v", outHit.CacheReadTokens) + } +} diff --git a/internal/providers/qoder/direct.go b/internal/providers/qoder/direct.go index 481ab8bf..718ef3c3 100644 --- a/internal/providers/qoder/direct.go +++ b/internal/providers/qoder/direct.go @@ -467,11 +467,15 @@ func applyUsage(out *providers.ChatOutcome, raw json.RawMessage) { return } var usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - CacheReadTokens *int `json:"cache_read_tokens"` - CacheWriteTokens *int `json:"cache_write_tokens"` - Credits *float64 `json:"credits"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CacheReadTokens *int `json:"cache_read_tokens"` + CacheWriteTokens *int `json:"cache_write_tokens"` + PromptCacheHitTokens *int `json:"prompt_cache_hit_tokens"` + Credits *float64 `json:"credits"` + PromptDetails struct { + CachedTokens *int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` } if json.Unmarshal(raw, &usage) != nil { out.UsageSource = "estimate" @@ -479,7 +483,14 @@ func applyUsage(out *providers.ChatOutcome, raw json.RawMessage) { } out.PromptTokens = usage.PromptTokens out.CompletionTokens = usage.CompletionTokens - out.CacheReadTokens = usage.CacheReadTokens + cacheRead := usage.CacheReadTokens + if cacheRead == nil { + cacheRead = usage.PromptDetails.CachedTokens + } + if cacheRead == nil { + cacheRead = usage.PromptCacheHitTokens + } + out.CacheReadTokens = cacheRead out.CacheWriteTokens = usage.CacheWriteTokens out.Credits = usage.Credits } @@ -618,6 +629,8 @@ func rewriteUpstreamSSE(resp *http.Response, model string) *http.Response { func writeOpenAIChunks(w io.Writer, src io.Reader, model string) error { var sawFinish, sawTools, sawPayload, sawError bool + var lastUsage any + var lastChunkEmittedHadUsage bool err := walkUpstream(src, func(body map[string]any) error { frame := frameFromBody(body) if frame.finish != "" { @@ -640,11 +653,21 @@ func writeOpenAIChunks(w io.Writer, src io.Reader, model string) error { "choices": []any{map[string]any{"index": 0, "delta": deltaOf(frame), "finish_reason": emptyOrNil(frame.finish)}}, } if len(frame.usage) > 0 { - var usage any - if json.Unmarshal(frame.usage, &usage) == nil { - chunk["usage"] = usage + var rawUsage any + if json.Unmarshal(frame.usage, &rawUsage) == nil { + if usage, ok := rawUsage.(map[string]any); ok { + normalizePromptCacheUsage(usage) + chunk["usage"] = usage + lastUsage = usage + } else { + chunk["usage"] = rawUsage + lastUsage = rawUsage + } } + } else if frame.finish != "" && lastUsage != nil { + chunk["usage"] = lastUsage } + lastChunkEmittedHadUsage = chunk["usage"] != nil payload, err := json.Marshal(chunk) if err != nil { return err @@ -669,6 +692,9 @@ func writeOpenAIChunks(w io.Writer, src io.Reader, model string) error { "id": "chatcmpl-qoder", "object": "chat.completion.chunk", "model": model, "choices": []any{map[string]any{"index": 0, "delta": map[string]any{}, "finish_reason": finish}}, } + if lastUsage != nil { + chunk["usage"] = lastUsage + } payload, err := json.Marshal(chunk) if err != nil { return err @@ -686,11 +712,52 @@ func writeOpenAIChunks(w io.Writer, src io.Reader, model string) error { return err } } + } else if !lastChunkEmittedHadUsage && lastUsage != nil { + chunk := map[string]any{ + "id": "chatcmpl-qoder", "object": "chat.completion.chunk", "model": model, + "choices": []any{}, + "usage": lastUsage, + } + payload, err := json.Marshal(chunk) + if err != nil { + return err + } + if _, err := fmt.Fprintf(w, "data: %s\n\n", payload); err != nil { + return err + } } _, err = io.WriteString(w, "data: [DONE]\n\n") return err } +func normalizePromptCacheUsage(usage map[string]any) { + if usage == nil { + return + } + var cachedTokens any + if ptd, ok := usage["prompt_tokens_details"].(map[string]any); ok { + cachedTokens = ptd["cached_tokens"] + } + if cachedTokens == nil { + cachedTokens = usage["cache_read_tokens"] + } + if cachedTokens == nil { + cachedTokens = usage["prompt_cache_hit_tokens"] + } + if cachedTokens == nil { + return + } + if _, ok := usage["cache_read_tokens"]; !ok { + usage["cache_read_tokens"] = cachedTokens + } + if _, ok := usage["prompt_cache_hit_tokens"]; !ok { + usage["prompt_cache_hit_tokens"] = cachedTokens + } + if _, ok := usage["prompt_tokens_details"]; !ok { + usage["prompt_tokens_details"] = map[string]any{"cached_tokens": cachedTokens} + } +} + func frameFromBody(body map[string]any) upstreamFrame { var frame upstreamFrame if errText, status, kind := upstreamError(body); errText != "" { diff --git a/internal/providers/qoder/direct_test.go b/internal/providers/qoder/direct_test.go index f7ad31f5..1e61b44a 100644 --- a/internal/providers/qoder/direct_test.go +++ b/internal/providers/qoder/direct_test.go @@ -222,3 +222,26 @@ func TestDirectChatNonStreamRejectsEmptyUpstream(t *testing.T) { t.Fatalf("err = %#v", err) } } + +func TestDirectChatStreamSalvagesTruncatedStreamWithCachedUsage(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"body\":{\"choices\":[{\"delta\":{\"content\":\"partial\"}}],\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":10,\"prompt_tokens_details\":{\"cached_tokens\":80}}}}\n\n") + })) + defer upstream.Close() + direct := newDirectFor(t, upstream) + + resp, _, err := direct.ChatStream(context.Background(), "acc", translate.ChatRequest{Model: "auto"}) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + text := string(body) + if !strings.Contains(text, `"finish_reason":"stop"`) { + t.Fatalf("missing synthesized finish_reason: %s", text) + } + if !strings.Contains(text, `"cached_tokens":80`) { + t.Fatalf("missing cached_tokens in terminal chunk: %s", text) + } +} diff --git a/internal/providers/trae/cache_test.go b/internal/providers/trae/cache_test.go new file mode 100644 index 00000000..07a0d7e3 --- /dev/null +++ b/internal/providers/trae/cache_test.go @@ -0,0 +1,25 @@ +package trae + +import ( + "testing" +) + +func TestTraeOutcomePreservesPromptCache(t *testing.T) { + aggregate := map[string]any{ + "model": "glm-5.2", + "usage": map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + "prompt_tokens_details": map[string]any{ + "cached_tokens": 80, + }, + }, + } + out, err := outcomeFromAggregate(aggregate, "fallback") + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if out.CacheReadTokens == nil || *out.CacheReadTokens != 80 { + t.Fatalf("expected CacheReadTokens 80, got %v", out.CacheReadTokens) + } +} diff --git a/internal/providers/trae/client.go b/internal/providers/trae/client.go index aabaff6a..9f6625d0 100644 --- a/internal/providers/trae/client.go +++ b/internal/providers/trae/client.go @@ -980,8 +980,14 @@ func outcomeFromAggregate(aggregate map[string]any, fallbackModel string) (provi var parsed struct { Model string `json:"model"` Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CacheReadTokens *int `json:"cache_read_tokens"` + CacheWriteTokens *int `json:"cache_write_tokens"` + PromptCacheHitTokens *int `json:"prompt_cache_hit_tokens"` + PromptDetails struct { + CachedTokens *int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` } `json:"usage"` Choices []struct { FinishReason string `json:"finish_reason"` @@ -1009,6 +1015,15 @@ func outcomeFromAggregate(aggregate map[string]any, fallbackModel string) (provi } out.PromptTokens = parsed.Usage.PromptTokens out.CompletionTokens = parsed.Usage.CompletionTokens + cacheRead := parsed.Usage.CacheReadTokens + if cacheRead == nil { + cacheRead = parsed.Usage.PromptDetails.CachedTokens + } + if cacheRead == nil { + cacheRead = parsed.Usage.PromptCacheHitTokens + } + out.CacheReadTokens = cacheRead + out.CacheWriteTokens = parsed.Usage.CacheWriteTokens return out, nil } diff --git a/internal/providers/trae/client_test.go b/internal/providers/trae/client_test.go index 6a00a73b..3e7638b6 100644 --- a/internal/providers/trae/client_test.go +++ b/internal/providers/trae/client_test.go @@ -356,6 +356,9 @@ func TestChatStreamRewritesOpenAIChunks(t *testing.T) { if !strings.Contains(text, `"object":"chat.completion.chunk"`) || !strings.Contains(text, "data: [DONE]") { t.Fatalf("stream=%s", text) } + if !strings.Contains(text, `"finish_reason":"tool_calls"`) || !strings.Contains(text, `"prompt_tokens":3`) { + t.Fatalf("terminal chunk missing finish_reason or usage: %s", text) + } if strings.Contains(text, "event: output") { t.Fatalf("must rewrite solo events, got %s", text) } diff --git a/internal/providers/trae/sse.go b/internal/providers/trae/sse.go index 6c2cd7f0..e6ad861d 100644 --- a/internal/providers/trae/sse.go +++ b/internal/providers/trae/sse.go @@ -39,6 +39,43 @@ func (e streamError) codeString() string { } } +func normalizeTraeUsage(raw []byte) any { + if len(raw) == 0 { + return nil + } + var m map[string]any + if err := json.Unmarshal(raw, &m); err != nil { + return json.RawMessage(raw) + } + var cachedTokens any + if ptd, ok := m["prompt_tokens_details"].(map[string]any); ok { + cachedTokens = ptd["cached_tokens"] + } + if cachedTokens == nil { + cachedTokens = m["cache_read_tokens"] + } + if cachedTokens == nil { + cachedTokens = m["prompt_cache_hit_tokens"] + } + if cachedTokens != nil { + if _, ok := m["cache_read_tokens"]; !ok { + m["cache_read_tokens"] = cachedTokens + } + if _, ok := m["prompt_cache_hit_tokens"]; !ok { + m["prompt_cache_hit_tokens"] = cachedTokens + } + if _, ok := m["prompt_tokens_details"]; !ok { + m["prompt_tokens_details"] = map[string]any{"cached_tokens": cachedTokens} + } + } + if _, ok := m["total_tokens"]; !ok { + pt, _ := m["prompt_tokens"].(float64) + ct, _ := m["completion_tokens"].(float64) + m["total_tokens"] = int(pt + ct) + } + return m +} + // Aggregate converts a Solo SSE stream into one OpenAI chat.completion object. func Aggregate(reader io.Reader) (map[string]any, error) { result := map[string]any{"object": "chat.completion"} @@ -78,7 +115,7 @@ func Aggregate(reader io.Reader) (map[string]any, error) { } mergeSoloToolCalls(toolCalls, &order, out.ToolCalls) case "token_usage": - usage = json.RawMessage(ev.Data) + usage = normalizeTraeUsage(ev.Data) case "done": sawDone = true var done struct { @@ -212,6 +249,7 @@ func startSoloRewrite(src io.ReadCloser, upstream io.Closer, model string) (*htt id := fmt.Sprintf("chatcmpl-%d", time.Now().UnixNano()) created := time.Now().Unix() var pendingUsage any + var lastUsage any ended := false writeChunk := func(delta map[string]any, finish string, usage any) error { chunk := map[string]any{ @@ -275,7 +313,8 @@ func startSoloRewrite(src io.ReadCloser, upstream io.Closer, model string) (*htt pendingUsage = nil return writeChunk(delta, "", usage) case "token_usage": - pendingUsage = json.RawMessage(ev.Data) + lastUsage = normalizeTraeUsage(ev.Data) + pendingUsage = lastUsage case "done": ended = true var done struct { @@ -287,6 +326,9 @@ func startSoloRewrite(src io.ReadCloser, upstream io.Closer, model string) (*htt finish = "stop" } usage := pendingUsage + if usage == nil && lastUsage != nil { + usage = lastUsage + } pendingUsage = nil if err := writeChunk(map[string]any{}, finish, usage); err != nil { return err @@ -306,7 +348,11 @@ func startSoloRewrite(src io.ReadCloser, upstream io.Closer, model string) (*htt return nil }) if !ended { - _ = writeChunk(map[string]any{}, "stop", pendingUsage) + usage := pendingUsage + if usage == nil && lastUsage != nil { + usage = lastUsage + } + _ = writeChunk(map[string]any{}, "stop", usage) _, _ = io.WriteString(pw, "data: [DONE]\n\n") } if err != nil { diff --git a/internal/providers/workbuddy/cache_test.go b/internal/providers/workbuddy/cache_test.go new file mode 100644 index 00000000..f0ad4b2d --- /dev/null +++ b/internal/providers/workbuddy/cache_test.go @@ -0,0 +1,23 @@ +package workbuddy + +import ( + "testing" +) + +func TestWorkBuddyOutcomePreservesPromptCacheHitTokens(t *testing.T) { + aggregate := map[string]any{ + "model": "deepseek-v4.1", + "usage": map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + "prompt_cache_hit_tokens": 80, + }, + } + out, err := outcomeFromAggregate(aggregate) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if out.CacheReadTokens == nil || *out.CacheReadTokens != 80 { + t.Fatalf("expected CacheReadTokens 80, got %v", out.CacheReadTokens) + } +} diff --git a/internal/providers/workbuddy/client.go b/internal/providers/workbuddy/client.go index 0ae2889e..a24b8a28 100644 --- a/internal/providers/workbuddy/client.go +++ b/internal/providers/workbuddy/client.go @@ -546,13 +546,14 @@ func outcomeFromAggregate(aggregate map[string]any) (providers.ChatOutcome, erro var parsed struct { Model string `json:"model"` Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - CacheReadTokens *int `json:"cache_read_tokens"` - CacheWriteTokens *int `json:"cache_write_tokens"` - Source string `json:"source"` - Credit *float64 `json:"credit"` - PromptDetails struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CacheReadTokens *int `json:"cache_read_tokens"` + CacheWriteTokens *int `json:"cache_write_tokens"` + PromptCacheHitTokens *int `json:"prompt_cache_hit_tokens"` + Source string `json:"source"` + Credit *float64 `json:"credit"` + PromptDetails struct { CachedTokens *int `json:"cached_tokens"` } `json:"prompt_tokens_details"` } `json:"usage"` @@ -584,6 +585,9 @@ func outcomeFromAggregate(aggregate map[string]any) (providers.ChatOutcome, erro if out.CacheReadTokens == nil { out.CacheReadTokens = parsed.Usage.PromptDetails.CachedTokens } + if out.CacheReadTokens == nil { + out.CacheReadTokens = parsed.Usage.PromptCacheHitTokens + } out.CacheWriteTokens = parsed.Usage.CacheWriteTokens out.Credits = parsed.Usage.Credit return out, nil diff --git a/internal/providers/workbuddy/client_test.go b/internal/providers/workbuddy/client_test.go index 0c0169e0..2a328994 100644 --- a/internal/providers/workbuddy/client_test.go +++ b/internal/providers/workbuddy/client_test.go @@ -1872,6 +1872,7 @@ func TestOutcomeFromAggregatePreservesCacheUsage(t *testing.T) { }{ {name: "top-level cache fields", usage: `"cache_read_tokens":12,"cache_write_tokens":3`, read: 12, write: 3}, {name: "cached tokens detail fallback", usage: `"prompt_tokens_details":{"cached_tokens":7},"cache_write_tokens":0`, read: 7, write: 0}, + {name: "prompt cache hit tokens", usage: `"prompt_cache_hit_tokens":80,"cache_write_tokens":0`, read: 80, write: 0}, } { t.Run(test.name, func(t *testing.T) { aggregate, err := Aggregate(strings.NewReader(strings.Join([]string{ diff --git a/internal/providers/workbuddy/sse.go b/internal/providers/workbuddy/sse.go index 6592c6bd..9fe64021 100644 --- a/internal/providers/workbuddy/sse.go +++ b/internal/providers/workbuddy/sse.go @@ -229,6 +229,29 @@ func sanitizeSSEPayload(payload string, sentRole *bool) (string, bool) { return payload, true } keep := chunk["usage"] != nil + if usageMap, ok := chunk["usage"].(map[string]any); ok { + var cachedTokens any + if ptd, ok := usageMap["prompt_tokens_details"].(map[string]any); ok { + cachedTokens = ptd["cached_tokens"] + } + if cachedTokens == nil { + cachedTokens = usageMap["cache_read_tokens"] + } + if cachedTokens == nil { + cachedTokens = usageMap["prompt_cache_hit_tokens"] + } + if cachedTokens != nil { + if _, ok := usageMap["cache_read_tokens"]; !ok { + usageMap["cache_read_tokens"] = cachedTokens + } + if _, ok := usageMap["prompt_cache_hit_tokens"]; !ok { + usageMap["prompt_cache_hit_tokens"] = cachedTokens + } + if _, ok := usageMap["prompt_tokens_details"]; !ok { + usageMap["prompt_tokens_details"] = map[string]any{"cached_tokens": cachedTokens} + } + } + } if choices, ok := chunk["choices"].([]any); ok { for _, raw := range choices { choice, _ := raw.(map[string]any) diff --git a/internal/runtime/alias.go b/internal/runtime/alias.go index 25340bd7..fe67f83b 100644 --- a/internal/runtime/alias.go +++ b/internal/runtime/alias.go @@ -23,6 +23,7 @@ type ( ) const ( + KindAuth = accounts.KindAuth KindUnavailable = accounts.KindUnavailable DefaultWorkBuddyCheckinTime = accounts.DefaultWorkBuddyCheckinTime backoffMaxLevel = accounts.BackoffMaxLevel diff --git a/internal/runtime/manager_probe.go b/internal/runtime/manager_probe.go index 936fbce9..68191600 100644 --- a/internal/runtime/manager_probe.go +++ b/internal/runtime/manager_probe.go @@ -112,12 +112,22 @@ func (m *Manager) refreshInProcess(ctx context.Context, item Item) error { m.resetRestartBackoff(item.ID) } status := "login_required" + kind := "" if health.Ready || health.Hot { status = "ready" } else if health.LastError != "" { status = "error" + if adapter.Classifier != nil { + classified := adapter.Classifier.Classify(0, health.LastError) + if classified.Kind != "" { + kind = classified.Kind + } + if classified.Kind == KindAuth { + status = "login_required" + } + } } - if err := m.store.Observe(ctx, item.ID, health.UID, status, health.LastError, ""); err != nil { + if err := m.store.Observe(ctx, item.ID, health.UID, status, health.LastError, kind); err != nil { return err } if health.Ready || health.Hot { diff --git a/internal/store/store.go b/internal/store/store.go index c5e2239f..b3fa129c 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -949,9 +949,13 @@ FROM account_cooldowns WHERE down_until > ? ORDER BY account_id, model`, now) func (s *Store) RecordPoolState(ctx context.Context, state accounts.PoolState) error { var cooldown any - status := "ready" if !state.DownUntil.IsZero() && time.Now().Before(state.DownUntil) { cooldown = formatTime(state.DownUntil) + } + status := "ready" + if state.LastErrorKind == accounts.KindAuth { + status = "login_required" + } else if cooldown != nil { status = "cooling" } _, err := s.db.ExecContext(ctx, `