From bb83596df43c5fffecbcc65d00ed553cea4faf52 Mon Sep 17 00:00:00 2001 From: Keton Loo <258433097+longlixiao-ong@users.noreply.github.com> Date: Mon, 12 Oct 2026 03:44:56 +0000 Subject: [PATCH 1/4] fix: preserve prompt cache tokens across providers and salvage streams on terminal finish_reason --- .../prompt-cache-and-stream-salvage.md | 9 ++++ internal/executor/chat.go | 31 +++++++---- internal/gateway/openai_stream.go | 53 +++++++++++++++---- internal/gateway/openai_stream_test.go | 37 +++++++++++++ internal/gateway/usage.go | 1 + internal/providers/qoder/adapter.go | 25 ++++++--- internal/providers/qoder/cache_test.go | 24 +++++++++ internal/providers/qoder/direct.go | 46 +++++++++++++--- internal/providers/trae/cache_test.go | 25 +++++++++ internal/providers/trae/client.go | 19 ++++++- internal/providers/trae/sse.go | 49 ++++++++++++++++- internal/providers/workbuddy/cache_test.go | 23 ++++++++ internal/providers/workbuddy/client.go | 18 ++++--- internal/providers/workbuddy/sse.go | 23 ++++++++ 14 files changed, 340 insertions(+), 43 deletions(-) create mode 100644 changelog/unreleased/prompt-cache-and-stream-salvage.md create mode 100644 internal/gateway/openai_stream_test.go 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..4c43e1f3 --- /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 \ inish_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\）。 +- 容错处理上游流式响应在带有终态 \ inish_reason\ 正常结束但未显式发送 \[DONE]\ 的边界情况，自动补齐结束帧，避免下游误判为 502 流异常中断。 diff --git a/internal/executor/chat.go b/internal/executor/chat.go index c7062d74..c0a14cdd 100644 --- a/internal/executor/chat.go +++ b/internal/executor/chat.go @@ -967,6 +967,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, @@ -1061,13 +1062,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"` @@ -1105,6 +1107,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, @@ -1113,9 +1126,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/gateway/openai_stream.go b/internal/gateway/openai_stream.go index 6b3977b4..ffebcf7c 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 != "" { + 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 @@ -405,6 +416,22 @@ func jsonHasArray(raw json.RawMessage) bool { return len(parts) > 0 } +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 ParseStreamUsageLine(line string) (StreamRelayStats, 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/openai_stream_test.go b/internal/gateway/openai_stream_test.go new file mode 100644 index 00000000..9ec1a923 --- /dev/null +++ b/internal/gateway/openai_stream_test.go @@ -0,0 +1,37 @@ +package gateway + +import ( + "net/http/httptest" + "strings" + "testing" +) + +func TestParseStreamUsageLineWithPromptCacheHitTokens(t *testing.T) { + line := `data: {"model":"deepseek-v4.1","usage":{"prompt_tokens":100,"completion_tokens":50,"prompt_cache_hit_tokens":80}}` + stats, ok := ParseStreamUsageLine(line) + if !ok { + t.Fatalf("expected ok=true") + } + if stats.CacheReadTokens == nil || *stats.CacheReadTokens != 80 { + t.Fatalf("expected CacheReadTokens=80, got %v", stats.CacheReadTokens) + } + if stats.CachedTokens == nil || *stats.CachedTokens != 80 { + t.Fatalf("expected CachedTokens=80, got %v", stats.CachedTokens) + } +} + +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()) + } +} 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/qoder/adapter.go b/internal/providers/qoder/adapter.go index 1533121f..5d9e41ca 100644 --- a/internal/providers/qoder/adapter.go +++ b/internal/providers/qoder/adapter.go @@ -639,12 +639,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"` @@ -658,11 +662,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 2bd98c6b..e01e5caa 100644 --- a/internal/providers/qoder/direct.go +++ b/internal/providers/qoder/direct.go @@ -445,11 +445,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" @@ -457,7 +461,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,8 +629,29 @@ 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 + var usage map[string]any if json.Unmarshal(frame.usage, &usage) == nil { + 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 { + 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} + } + } chunk["usage"] = usage } } 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..de35c6c1 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/sse.go b/internal/providers/trae/sse.go index 6c2cd7f0..f731f2e2 100644 --- a/internal/providers/trae/sse.go +++ b/internal/providers/trae/sse.go @@ -40,6 +40,51 @@ func (e streamError) codeString() string { } // Aggregate converts a Solo SSE stream into one OpenAI chat.completion object. +func normalizeTraeUsage(raw []byte) any { + if len(raw) == 0 { + return nil + } + var u struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_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"` + } + if err := json.Unmarshal(raw, &u); err != nil { + return json.RawMessage(raw) + } + cacheRead := u.CacheReadTokens + if cacheRead == nil && u.PromptDetails != nil { + cacheRead = u.PromptDetails.CachedTokens + } + if cacheRead == nil { + cacheRead = u.PromptCacheHitTokens + } + total := u.TotalTokens + if total == 0 { + total = u.PromptTokens + u.CompletionTokens + } + m := map[string]any{ + "prompt_tokens": u.PromptTokens, + "completion_tokens": u.CompletionTokens, + "total_tokens": total, + } + if cacheRead != nil { + m["cache_read_tokens"] = *cacheRead + m["prompt_cache_hit_tokens"] = *cacheRead + m["prompt_tokens_details"] = map[string]any{"cached_tokens": *cacheRead} + } + if u.CacheWriteTokens != nil { + m["cache_write_tokens"] = *u.CacheWriteTokens + } + return m +} + func Aggregate(reader io.Reader) (map[string]any, error) { result := map[string]any{"object": "chat.completion"} message := map[string]any{"role": "assistant"} @@ -78,7 +123,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 { @@ -275,7 +320,7 @@ 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) + pendingUsage = normalizeTraeUsage(ev.Data) case "done": ended = true var done struct { 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/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) From ce173cb79c7e1ab2799bcf9ac899943d37ed0258 Mon Sep 17 00:00:00 2001 From: Keton Loo <258433097+longlixiao-ong@users.noreply.github.com> Date: Mon, 12 Oct 2026 05:59:36 +0000 Subject: [PATCH 2/4] fix: preserve raw usage fields in trae and harden qoder direct chunks --- deploy/docker-compose.yml | 14 +++++--- internal/gateway/openai_stream.go | 2 +- internal/providers/qoder/direct.go | 44 +++++++++++++----------- internal/providers/trae/sse.go | 54 +++++++++++++----------------- 4 files changed, 58 insertions(+), 56 deletions(-) diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index 04407342..7717af1f 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -1,8 +1,5 @@ services: qoder-api-proxy: - build: - context: .. - dockerfile: deploy/Dockerfile image: "${CLI2API_IMAGE:-ghcr.io/caigee-cmd/cli2api:latest}" container_name: qoder-api-proxy environment: @@ -24,6 +21,9 @@ services: - /run/cli2api:mode=0700 ports: - "127.0.0.1:3010:3010" + networks: + - default + - new-api-net healthcheck: test: ["CMD", "curl", "-fsS", "http://127.0.0.1:3010/health"] interval: 10s @@ -33,4 +33,10 @@ services: restart: unless-stopped volumes: - qoder-data: \ No newline at end of file + qoder-data: + +networks: + default: + new-api-net: + external: true + name: new-api-deploy_default diff --git a/internal/gateway/openai_stream.go b/internal/gateway/openai_stream.go index ffebcf7c..dbc2b9c0 100644 --- a/internal/gateway/openai_stream.go +++ b/internal/gateway/openai_stream.go @@ -166,7 +166,7 @@ func RelayOpenAIStream(w http.ResponseWriter, body io.Reader) (stats StreamRelay return stats, streamErr } if !sawDone { - if stats.FinishReason != "" { + 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 { diff --git a/internal/providers/qoder/direct.go b/internal/providers/qoder/direct.go index e01e5caa..cd173aaa 100644 --- a/internal/providers/qoder/direct.go +++ b/internal/providers/qoder/direct.go @@ -629,30 +629,34 @@ 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 map[string]any - if json.Unmarshal(frame.usage, &usage) == nil { - 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 { - if _, ok := usage["cache_read_tokens"]; !ok { - usage["cache_read_tokens"] = cachedTokens + var rawUsage any + if json.Unmarshal(frame.usage, &rawUsage) == nil { + if usage, ok := rawUsage.(map[string]any); ok { + 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 _, ok := usage["prompt_cache_hit_tokens"]; !ok { - usage["prompt_cache_hit_tokens"] = cachedTokens + if cachedTokens == nil { + cachedTokens = usage["prompt_cache_hit_tokens"] } - if _, ok := usage["prompt_tokens_details"]; !ok { - usage["prompt_tokens_details"] = map[string]any{"cached_tokens": cachedTokens} + if cachedTokens != nil { + 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} + } } + chunk["usage"] = usage + } else { + chunk["usage"] = rawUsage } - chunk["usage"] = usage } } payload, err := json.Marshal(chunk) diff --git a/internal/providers/trae/sse.go b/internal/providers/trae/sse.go index f731f2e2..9e160d69 100644 --- a/internal/providers/trae/sse.go +++ b/internal/providers/trae/sse.go @@ -44,43 +44,35 @@ func normalizeTraeUsage(raw []byte) any { if len(raw) == 0 { return nil } - var u struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_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"` - } - if err := json.Unmarshal(raw, &u); err != nil { + var m map[string]any + if err := json.Unmarshal(raw, &m); err != nil { return json.RawMessage(raw) } - cacheRead := u.CacheReadTokens - if cacheRead == nil && u.PromptDetails != nil { - cacheRead = u.PromptDetails.CachedTokens + var cachedTokens any + if ptd, ok := m["prompt_tokens_details"].(map[string]any); ok { + cachedTokens = ptd["cached_tokens"] } - if cacheRead == nil { - cacheRead = u.PromptCacheHitTokens + if cachedTokens == nil { + cachedTokens = m["cache_read_tokens"] } - total := u.TotalTokens - if total == 0 { - total = u.PromptTokens + u.CompletionTokens + if cachedTokens == nil { + cachedTokens = m["prompt_cache_hit_tokens"] } - m := map[string]any{ - "prompt_tokens": u.PromptTokens, - "completion_tokens": u.CompletionTokens, - "total_tokens": total, - } - if cacheRead != nil { - m["cache_read_tokens"] = *cacheRead - m["prompt_cache_hit_tokens"] = *cacheRead - m["prompt_tokens_details"] = map[string]any{"cached_tokens": *cacheRead} + 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 u.CacheWriteTokens != nil { - m["cache_write_tokens"] = *u.CacheWriteTokens + if _, ok := m["total_tokens"]; !ok { + pt, _ := m["prompt_tokens"].(float64) + ct, _ := m["completion_tokens"].(float64) + m["total_tokens"] = int(pt + ct) } return m } From 05fb49ed5aa3de2679e86c6d509bc6e0ef47fb19 Mon Sep 17 00:00:00 2001 From: Keton Loo <258433097+longlixiao-ong@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:52:45 +0000 Subject: [PATCH 3/4] fix: classify devin expired session and connect auth denials as kind auth --- internal/providers/devin/devin_test.go | 34 ++++++++++++++++++++++++++ internal/providers/devin/errors.go | 29 +++++++++++++++++++--- internal/runtime/alias.go | 1 + internal/runtime/manager_probe.go | 12 ++++++++- internal/store/store.go | 6 ++++- 5 files changed, 76 insertions(+), 6 deletions(-) diff --git a/internal/providers/devin/devin_test.go b/internal/providers/devin/devin_test.go index 62970d9e..29b085e0 100644 --- a/internal/providers/devin/devin_test.go +++ b/internal/providers/devin/devin_test.go @@ -744,6 +744,40 @@ func TestClassifyAuthForbiddenQuotaTrailer(t *testing.T) { } } +func TestClassifyExpiredTokenAndConnectAuthDenial(t *testing.T) { + // 1. HTTP 401 token expired + 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) + } + + // 2. Connect trailer permission_denied with internal error + 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) + } + + // 3. Seat management HTTP 403 with permission_denied and internal error + 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) + } + + // 4. Token expired in status 400 + got = Classify(400, "token has expired") + if got.Kind != accounts.KindAuth { + t.Fatalf("token expired text classify=%+v want kind=auth", got) + } + + // 5. Bare permission_denied without internal error remains invalid_request + 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/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 fcecf845..94704d5b 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -849,9 +849,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, ` From 7bd16f21732c07699f31e19af01f0d87abb50151 Mon Sep 17 00:00:00 2001 From: Keton Loo <258433097+longlixiao-ong@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:16:02 +0000 Subject: [PATCH 4/4] fix(providers): preserve usage and prompt cache tokens in streaming terminal chunks for qoder and trae --- internal/providers/qoder/direct.go | 25 +++++++++++++++++++++++++ internal/providers/qoder/direct_test.go | 23 +++++++++++++++++++++++ internal/providers/trae/client_test.go | 3 +++ internal/providers/trae/sse.go | 13 +++++++++++-- 4 files changed, 62 insertions(+), 2 deletions(-) diff --git a/internal/providers/qoder/direct.go b/internal/providers/qoder/direct.go index cd173aaa..1e32abc4 100644 --- a/internal/providers/qoder/direct.go +++ b/internal/providers/qoder/direct.go @@ -607,6 +607,9 @@ 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 != "" { @@ -654,11 +657,17 @@ func writeOpenAIChunks(w io.Writer, src io.Reader, model string) error { } } 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 @@ -683,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 @@ -700,6 +712,19 @@ 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 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/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 9e160d69..eb8fc224 100644 --- a/internal/providers/trae/sse.go +++ b/internal/providers/trae/sse.go @@ -274,6 +274,7 @@ func startSoloRewrite(src io.ReadCloser, upstream io.Closer, model string) (*htt _, err = fmt.Fprintf(pw, "data: %s\n\n", encoded) return err } + var lastUsage any err := scanSolo(src, func(ev soloEvent) error { switch ev.Name { case "metadata": @@ -312,7 +313,8 @@ func startSoloRewrite(src io.ReadCloser, upstream io.Closer, model string) (*htt pendingUsage = nil return writeChunk(delta, "", usage) case "token_usage": - pendingUsage = normalizeTraeUsage(ev.Data) + lastUsage = normalizeTraeUsage(ev.Data) + pendingUsage = lastUsage case "done": ended = true var done struct { @@ -324,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 @@ -343,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 {