From bde1c9094de6c4d16bd287c0dfe75faf46da0e84 Mon Sep 17 00:00:00 2001 From: caiji <522caiji@gmail.com> Date: Fri, 13 Nov 2026 17:34:09 +0800 Subject: [PATCH] fix(gateway): cache affinity and Anthropic message rectification Port the cache-stability and message-format fixes from the reference 2api gateways and cc-switch: - Recognise conversation ids from headers and body metadata aliases (Sliverkiss/workbuddy2api #35, qoder2api-hub session affinity), so turns stay on one account instead of rotating and losing the upstream prompt cache. - Add opt-in CLI2API_MSGSTAT_SHARDS (workbuddy2api-hub WB_MSGSTAT_SHARDS): one (key, model, system prefix, tools) tuple maps to a small stable account set instead of being scattered by the first user message. - Add opt-in CLI2API_AFFINITY_DEBUG diagnostics separating route-level session affinity from real upstream cache-token hits. - Rectify historical Anthropic thinking / redacted_thinking blocks and stale signatures on /v1/messages (cc-switch Rectifier). Disable with CLI2API_THINKING_RECTIFIER=off. - Skip WAF fingerprint regexes for bodies with no known trigger, avoiding needless CPU work on large histories and keeping cached prefixes stable. No SQL migration changes; diagnostics are structured logs. New env switches are documented in deploy/README.md. Add a bilingual changelog fragment. --- .../cache-affinity-and-message-rectifier.md | 15 ++ deploy/README.md | 14 ++ internal/executor/prepare.go | 52 +++-- internal/executor/session_key.go | 173 ++++++++++++++++ internal/executor/session_key_test.go | 108 ++++++++++ internal/gateway/anthropic.go | 62 +++++- internal/gateway/openai.go | 72 ++++++- internal/gateway/prepare.go | 1 + internal/gateway/responses.go | 3 + internal/gateway/session_context.go | 108 ++++++++++ internal/gateway/session_context_test.go | 53 +++++ internal/providers/workbuddy/payload.go | 66 ++++++ .../workbuddy/waf_fingerprint_test.go | 28 +++ internal/translate/session_seed.go | 27 +++ internal/translate/thinking_rectifier.go | 188 ++++++++++++++++++ internal/translate/thinking_rectifier_test.go | 79 ++++++++ 16 files changed, 1026 insertions(+), 23 deletions(-) create mode 100644 changelog/unreleased/cache-affinity-and-message-rectifier.md create mode 100644 internal/executor/session_key.go create mode 100644 internal/executor/session_key_test.go create mode 100644 internal/gateway/session_context.go create mode 100644 internal/gateway/session_context_test.go create mode 100644 internal/translate/thinking_rectifier.go create mode 100644 internal/translate/thinking_rectifier_test.go diff --git a/changelog/unreleased/cache-affinity-and-message-rectifier.md b/changelog/unreleased/cache-affinity-and-message-rectifier.md new file mode 100644 index 00000000..c419ff3d --- /dev/null +++ b/changelog/unreleased/cache-affinity-and-message-rectifier.md @@ -0,0 +1,15 @@ +### English + +- Recognise client conversation ids from headers and body metadata aliases (`X-Conversation-Id`, `X-Session-Id`, `X-Thread-Id`, `conversation_id`, `conversationId`, `metadata.*`, ...), so chat turns stay on one account and keep the upstream prompt cache warm instead of rotating through the pool. +- Add opt-in `CLI2API_MSGSTAT_SHARDS` to bound one-shot Messages calls to a small, stable set of accounts per (API key, model, system prefix, tools), matching `WB_MSGSTAT_SHARDS` in workbuddy2api-hub. +- Add opt-in `CLI2API_AFFINITY_DEBUG` diagnostics that separate route-level session affinity from real upstream cache-token hits, including prompt/cache tokens and hit rate. +- Rectify historical Anthropic `thinking` / `redacted_thinking` blocks and stale signatures on `/v1/messages` (cc-switch-style minimal cleanup) so OpenAI-compatible upstreams no longer receive invalid signed blocks. Disable with `CLI2API_THINKING_RECTIFIER=off`. +- Skip WAF fingerprint regexes for request bodies with no known trigger, avoiding unnecessary per-request CPU work on large histories and keeping cached prefixes stable. + +### 中文 + +- 支持从请求头和 body 元数据别名识别会话 ID（`X-Conversation-Id`、`X-Session-Id`、`X-Thread-Id`、`conversation_id`、`conversationId`、`metadata.*` 等），让同一对话固定在同一账号上，避免轮转导致上游 prompt 缓存失效。 +- 新增可选 `CLI2API_MSGSTAT_SHARDS`，按（API Key、模型、system 前缀、tools）把单发 Messages 请求收敛到少量固定账号，对齐 workbuddy2api-hub 的 `WB_MSGSTAT_SHARDS`。 +- 新增可选 `CLI2API_AFFINITY_DEBUG` 诊断日志，把「路由级会话亲和」与「上游真实缓存 token 命中」分开记录，包含 prompt/cache tokens 与命中率。 +- `/v1/messages` 入站按 cc-switch 思路最小清理历史 `thinking` / `redacted_thinking` 块和失效签名，避免把无效签名转发给 OpenAI 兼容上游；可用 `CLI2API_THINKING_RECTIFIER=off` 关闭。 +- 对不含任何已知指纹的请求体跳过 WAF 正则清洗，减少长历史请求的 CPU 开销，并避免清洗改动缓存前缀。 diff --git a/deploy/README.md b/deploy/README.md index 67d091a4..c12952b3 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -240,3 +240,17 @@ Keep a separate database backup before upgrading. Keep `deploy/.env` private: Docker Desktop mode stores an updater token there. Do not remove `qoder-data` or run `docker compose down -v` unless you intend to delete the accounts and credentials. + + +## Cache affinity and message rectification + +These switches follow the reference 2api gateways and are off by default except +the Anthropic thinking rectifier, which is on because OpenAI-compatible +upstreams do not validate Anthropic thinking signatures. + +| Variable | Default | Description | +|---|---|---| +| `CLI2API_MSGSTAT_SHARDS` | `0` | Bound one-shot `/v1/messages` calls to N account buckets per (API key, model, system prefix, tools). Mirrors `WB_MSGSTAT_SHARDS`; `0` keeps the historical content-seed routing. | +| `CLI2API_AFFINITY_DEBUG` | `0` | Log session-affinity source and cache-token diagnostics per request. | +| `CLI2API_THINKING_RECTIFIER` | `on` | Set `off` to replay historical Anthropic thinking blocks unchanged. | +| `X-CLI2API-Session` / `X-Conversation-Id` / `X-Session-Id` / `X-Thread-Id` | — | Explicit conversation id. Body `conversation_id` / `conversationId` / `session_id` / `metadata.*` aliases are also recognised so client turns stay on one account and keep the upstream prompt cache warm. | diff --git a/internal/executor/prepare.go b/internal/executor/prepare.go index facd737b..1045bb7e 100644 --- a/internal/executor/prepare.go +++ b/internal/executor/prepare.go @@ -6,6 +6,7 @@ import ( "encoding/hex" "encoding/json" "log" + "os" "strconv" "strings" "time" @@ -59,6 +60,7 @@ type PrepareInput struct { Identity auth.Identity PreferAccount string SessionHeader string + ConversationID string CrossProviderPool bool Now time.Time NewRequestID func() string @@ -146,10 +148,18 @@ func (e ChatExecutor) Prepare(in PrepareInput) (PreparedRequest, error) { }) } ctx := WithAllowedProviders(WithRequestID(in.Context, requestID), in.Identity.AllowedProviders) - sessionKey := SessionKeyFor(in.SessionHeader, in.Identity, request) + sessionKey := sessionKeyFor(in.SessionHeader, in.ConversationID, in.Identity, request) if sessionKey == "" && in.NativeRequest != nil { sessionKey = NativeSessionKeyFor(in.NativeRequest, in.Identity) } + // No caller-supplied conversation: hub-style sharding keeps one + // (key, model, system, tools) tuple on a small account set. The content + // seed still runs when sharding is disabled (default). + if strings.TrimSpace(in.SessionHeader) == "" && strings.TrimSpace(in.ConversationID) == "" { + if shard := ShardFallbackKey(in.Identity, request); shard != "" { + sessionKey = shard + } + } if sessionKey != "" { ctx = WithSessionKey(ctx, sessionKey) ctx = WithConversation(ctx, sessionKey) @@ -160,6 +170,18 @@ func (e ChatExecutor) Prepare(in PrepareInput) (PreparedRequest, error) { if requestedReasoning != "" { log.Printf("request reasoning request_id=%q model=%q level=%q", requestID, request.Model, requestedReasoning) } + if sessionKey != "" && sessionAffinityDebugOn() { + source := "content-seed" + if strings.TrimSpace(in.SessionHeader) != "" { + source = "header" + } else if strings.TrimSpace(in.ConversationID) != "" { + source = "conversation" + } else if MessageStatShards() > 0 { + source = "shard" + } + log.Printf("session affinity request_id=%q source=%q key=%q model=%q tools=%t msgs=%d", + requestID, source, shortSessionKey(sessionKey), request.Model, len(request.Tools) > 0, len(request.Messages)) + } return PreparedRequest{ Context: ctx, RequestID: requestID, @@ -312,22 +334,7 @@ func NativeSessionKeyFor(req *translate.NativeResponsesRequest, identity auth.Id } func SessionKeyFor(header string, identity auth.Identity, req translate.ChatRequest) string { - raw := strings.TrimSpace(header) - kind := "content" - if raw != "" { - kind = "header" - } else { - raw = translate.ContentSessionSeed(req) - } - if raw == "" { - return "" - } - namespace := "console" - if identity.Kind == auth.KindKey && identity.KeyID != "" { - namespace = "key:" + identity.KeyID - } - sum := sha256.Sum256([]byte(namespace + "\x00" + kind + "\x00" + raw)) - return hex.EncodeToString(sum[:]) + return sessionKeyFor(header, "", identity, req) } // RequestedReasoningLevel surfaces the reasoning level the client asked for, @@ -371,3 +378,14 @@ func RequestedReasoningLevel(req translate.ChatRequest) string { } return "" } + +// sessionAffinityDebugOn mirrors qoder2api-hub's QD_AFFINITY_DEBUG: routing +// diagnostics are opt-in so normal request paths stay quiet. +func sessionAffinityDebugOn() bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv("CLI2API_AFFINITY_DEBUG"))) { + case "1", "true", "on", "yes": + return true + default: + return false + } +} diff --git a/internal/executor/session_key.go b/internal/executor/session_key.go new file mode 100644 index 00000000..429dd20e --- /dev/null +++ b/internal/executor/session_key.go @@ -0,0 +1,173 @@ +package executor + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "os" + "strconv" + "strings" + + "github.com/caigee-cmd/cli2api/internal/auth" + "github.com/caigee-cmd/cli2api/internal/translate" +) + +// sessionIDAliases mirrors the reference gateways' conversation-key lookup +// (Sliverkiss/workbuddy2api #35, qoder2api-hub session affinity). Clients use +// several spellings for the same conversation id; recognising only one makes +// turns rotate accounts and silently destroys the upstream prompt cache. +var sessionIDAliases = []string{ + "conversation_id", "conversationId", "conversation-id", + "session_id", "sessionId", "session-id", + "thread_id", "threadId", + "conversation", "session", +} + +// sessionKeySources carries every id a caller can use to pin one chat to one +// account. HTTP headers win over body metadata; the gateway's own header stays +// first for compatibility. +type sessionKeySources struct { + Headers []string + BodyJSON []byte + SessionID string +} + +// extractConversationKey returns the first non-empty, bounded conversation id +// found in headers, then in the request body (top level and metadata), then in +// the native Responses session seed. Unknown keys are ignored; no id is ever +// invented. Values are capped so a hostile client cannot pin unbounded memory. +func extractConversationKey(src sessionKeySources) string { + for _, value := range src.Headers { + if id := cleanConversationID(value); id != "" { + return id + } + } + if len(src.BodyJSON) > 0 && len(src.BodyJSON) <= 64<<20 { + var body map[string]json.RawMessage + if json.Unmarshal(src.BodyJSON, &body) == nil { + if id := conversationIDFromBody(body); id != "" { + return id + } + } + } + if id := cleanConversationID(src.SessionID); id != "" { + return id + } + return "" +} + +func conversationIDFromBody(body map[string]json.RawMessage) string { + for _, key := range sessionIDAliases { + if id, ok := rawConversationString(body[key]); ok { + return id + } + } + for _, metaKey := range []string{"metadata", "meta", "context"} { + var meta map[string]json.RawMessage + if json.Unmarshal(body[metaKey], &meta) != nil { + continue + } + for _, key := range sessionIDAliases { + if id, ok := rawConversationString(meta[key]); ok { + return id + } + } + } + return "" +} + +func rawConversationString(raw json.RawMessage) (string, bool) { + if len(raw) == 0 { + return "", false + } + var value string + if json.Unmarshal(raw, &value) == nil { + if id := cleanConversationID(value); id != "" { + return id, true + } + } + return "", false +} + +func cleanConversationID(value string) string { + value = strings.TrimSpace(value) + if value == "" || len(value) > 512 { + return "" + } + return value +} + +// sessionKeyFor is the shared implementation; PrepareInput carries an explicit +// conversation id in addition to the legacy header. +func sessionKeyFor(header, conversation string, identity auth.Identity, req translate.ChatRequest) string { + raw := strings.TrimSpace(header) + kind := "content" + if raw == "" { + raw = strings.TrimSpace(conversation) + } + if raw != "" { + kind = "header" + } else { + raw = translate.ContentSessionSeed(req) + } + if raw == "" { + return "" + } + namespace := "console" + if identity.Kind == auth.KindKey && identity.KeyID != "" { + namespace = "key:" + identity.KeyID + } + sum := sha256.Sum256([]byte(namespace + "\x00" + kind + "\x00" + raw)) + return hex.EncodeToString(sum[:]) +} + +// ShardKeyFor implements workbuddy2api-hub's WB_MSGSTAT_SHARDS idea with a pure +// function: one (key, model, tools, system-prefix) tuple maps to a small, stable +// set of accounts instead of being scattered by the first user text. The caller +// owns the account selection; shards is the maximum number of buckets. +func ShardKeyFor(identity auth.Identity, req translate.ChatRequest, shards int) string { + if shards <= 0 { + return "" + } + namespace := "console" + if identity.Kind == auth.KindKey && identity.KeyID != "" { + namespace = "key:" + identity.KeyID + } + seed := translate.PrefixSessionSeed(req) + if seed == "" { + return "" + } + sum := sha256.Sum256([]byte(namespace + "\x00shard\x00" + seed)) + bucket := int(sum[0])%shards + return namespace + "\x00" + hex.EncodeToString(sum[:8]) + "\x00" + string(rune('0'+bucket)) +} + + +func shortSessionKey(key string) string { + if len(key) <= 12 { + return key + } + return key[:12] +} + + +// MessageStatShards mirrors workbuddy2api-hub's WB_MSGSTAT_SHARDS. Default 0 +// keeps historical routing unchanged; a positive value bounds one-shot +// /v1/messages calls to N account buckets per (key, model, system, tools). +func MessageStatShards() int { + raw := strings.TrimSpace(os.Getenv("CLI2API_MSGSTAT_SHARDS")) + if raw == "" { + return 0 + } + n, err := strconv.Atoi(raw) + if err != nil || n < 1 || n > 64 { + return 0 + } + return n +} + +// ShardFallbackKey returns a stable, non-empty affinity key only when sharding +// is enabled. It never overrides an explicit conversation id. +func ShardFallbackKey(identity auth.Identity, req translate.ChatRequest) string { + return ShardKeyFor(identity, req, MessageStatShards()) +} diff --git a/internal/executor/session_key_test.go b/internal/executor/session_key_test.go new file mode 100644 index 00000000..4d6bbc25 --- /dev/null +++ b/internal/executor/session_key_test.go @@ -0,0 +1,108 @@ +package executor + +import ( + "testing" + + "github.com/caigee-cmd/cli2api/internal/auth" + "github.com/caigee-cmd/cli2api/internal/providers" + "github.com/caigee-cmd/cli2api/internal/translate" +) + +func TestSessionKeyForConversationAliases(t *testing.T) { + req := translate.ChatRequest{Model: "hy4free", Messages: []translate.ChatMessage{{Role: "user", Content: "hi"}}} + identity := auth.ConsoleIdentity() + base := SessionKeyFor("", identity, req) + explicit := sessionKeyFor("", "chat-42", identity, req) + if explicit == "" || explicit == base { + t.Fatalf("conversation id must override content seed: base=%q explicit=%q", base, explicit) + } + if got := sessionKeyFor("header-1", "chat-42", identity, req); got == explicit { + t.Fatal("header must win over conversation id") + } + if got := sessionKeyFor("header-1", "chat-42", identity, req); got != sessionKeyFor("header-1", "ignored", identity, req) { + t.Fatal("conversation id must not change a header-pinned key") + } +} + +func TestShardKeyStableAcrossTurns(t *testing.T) { + identity := auth.ConsoleIdentity() + first := translate.ChatRequest{Model: "hy4free", Messages: []translate.ChatMessage{ + {Role: "system", Content: "stable system"}, + {Role: "user", Content: "first question"}, + }} + later := translate.ChatRequest{Model: "hy4free", Messages: []translate.ChatMessage{ + {Role: "system", Content: "stable system"}, + {Role: "user", Content: "first question"}, + {Role: "assistant", Content: "answer"}, + {Role: "user", Content: "follow up"}, + }} + a := ShardKeyFor(identity, first, 4) + b := ShardKeyFor(identity, later, 4) + if a == "" || a != b { + t.Fatalf("shard broke across turns: %q vs %q", a, b) + } + if ShardKeyFor(identity, first, 0) != "" { + t.Fatal("disabled sharding must return empty") + } + other := translate.ChatRequest{Model: "hy3", Messages: first.Messages} + if ShardKeyFor(identity, other, 4) == a { + t.Fatal("different model must not share a shard") + } +} + +func TestPrepareUsesConversationID(t *testing.T) { + ex := NewChatExecutor(NewPool(nil, nil), "") + req := translate.ChatRequest{Model: "glm-5.2", Messages: []translate.ChatMessage{{Role: "user", Content: "hi"}}} + plain, err := ex.Prepare(PrepareInput{Request: req, Identity: auth.ConsoleIdentity(), CrossProviderPool: true}) + if err != nil { + t.Fatal(err) + } + withID, err := ex.Prepare(PrepareInput{Request: req, Identity: auth.ConsoleIdentity(), ConversationID: "conv-9", CrossProviderPool: true}) + if err != nil { + t.Fatal(err) + } + if sessionKeyFromContext(plain.Context) == sessionKeyFromContext(withID.Context) { + t.Fatal("conversation id must change the affinity key") + } + if providers.ConversationFromContext(withID.Context) == "" { + t.Fatal("conversation id must reach provider headers") + } +} + + +func TestMessageStatShardsEnvAndFallback(t *testing.T) { + if MessageStatShards() != 0 { + t.Fatal("sharding must be off by default") + } + t.Setenv("CLI2API_MSGSTAT_SHARDS", "4") + if MessageStatShards() != 4 { + t.Fatal("env not applied") + } + req := translate.ChatRequest{Model: "hy4free", Messages: []translate.ChatMessage{ + {Role: "system", Content: "stable"}, + {Role: "user", Content: "one"}, + }} + later := translate.ChatRequest{Model: "hy4free", Messages: []translate.ChatMessage{ + {Role: "system", Content: "stable"}, + {Role: "user", Content: "one"}, + {Role: "assistant", Content: "answer"}, + {Role: "user", Content: "two"}, + }} + a := ShardFallbackKey(auth.ConsoleIdentity(), req) + b := ShardFallbackKey(auth.ConsoleIdentity(), later) + if a == "" || a != b { + t.Fatalf("shard fallback broke across turns: %q %q", a, b) + } + ex := NewChatExecutor(NewPool(nil, nil), "") + prepared, err := ex.Prepare(PrepareInput{Request: req, Identity: auth.ConsoleIdentity(), CrossProviderPool: true}) + if err != nil { + t.Fatal(err) + } + if sessionKeyFromContext(prepared.Context) != a { + t.Fatal("Prepare did not use shard fallback") + } + t.Setenv("CLI2API_MSGSTAT_SHARDS", "bad") + if MessageStatShards() != 0 { + t.Fatal("invalid env must disable sharding") + } +} diff --git a/internal/gateway/anthropic.go b/internal/gateway/anthropic.go index a516ed81..f1c97637 100644 --- a/internal/gateway/anthropic.go +++ b/internal/gateway/anthropic.go @@ -5,7 +5,9 @@ import ( "encoding/json" "errors" "fmt" + "strings" "net/http" + "os" "time" "github.com/caigee-cmd/cli2api/internal/accounts" @@ -24,14 +26,35 @@ func (h *Handler) HandleAnthropicMessages(w http.ResponseWriter, r *http.Request writeAnthropicRequestBodyError(w, body.err) return } - if err := json.NewDecoder(body).Decode(&source); err != nil { - if isRequestBodyTooLarge(err) { - writeAnthropicRequestBodyError(w, err) + rawBody, readErr := readConversationBody(body, maxChatRequestBytes) + if readErr != nil { + if isRequestBodyTooLarge(readErr) { + writeAnthropicRequestBodyError(w, readErr) return } + writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", readErr.Error()) + return + } + rectified := rawBody + var rectify translate.RectifyResult + if thinkingRectifierOn() { + var rectifyErr error + rectified, rectify, rectifyErr = translate.RectifyAnthropicRequest(rawBody) + if rectifyErr != nil { + writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", rectifyErr.Error()) + return + } + if rectify.Applied { + w.Header().Set("X-CLI2API-Rectified", rectifySummary(rectify)) + } + } + if err := json.Unmarshal(rectified, &source); err != nil { writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) return } + if id := conversationIDFromRequest(r, rawBody); id != "" { + r = r.WithContext(WithConversationID(r.Context(), id)) + } request, err := translate.TranslateAnthropicMessages(source) if err != nil { writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) @@ -173,3 +196,36 @@ func anthropicStopReason(finishReason string, toolCalls []proxyToolCall) string } return "end_turn" } + + +// rectifySummary is a bounded, non-sensitive diagnostic header. It never +// includes message content or signatures. +func rectifySummary(result translate.RectifyResult) string { + parts := make([]string, 0, 4) + if result.RemovedThinkingBlocks > 0 { + parts = append(parts, fmt.Sprintf("thinking=%d", result.RemovedThinkingBlocks)) + } + if result.RemovedRedactedThinkingBlocks > 0 { + parts = append(parts, fmt.Sprintf("redacted=%d", result.RemovedRedactedThinkingBlocks)) + } + if result.RemovedSignatureFields > 0 { + parts = append(parts, fmt.Sprintf("signature=%d", result.RemovedSignatureFields)) + } + if result.RemovedTopLevelThinking { + parts = append(parts, "top-level=1") + } + return strings.Join(parts, ",") +} + + +// thinkingRectifierOn mirrors cc-switch's configurable rectifier. Default on: +// OpenAI-compatible upstreams have no signature validation, so historical +// Anthropic thinking signatures are always invalid there. +func thinkingRectifierOn() bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv("CLI2API_THINKING_RECTIFIER"))) { + case "0", "false", "off", "no": + return false + default: + return true + } +} diff --git a/internal/gateway/openai.go b/internal/gateway/openai.go index b91297e1..52907a0e 100644 --- a/internal/gateway/openai.go +++ b/internal/gateway/openai.go @@ -4,6 +4,9 @@ import ( "context" "encoding/json" "errors" + "log" + "os" + "strings" "net/http" "time" @@ -30,14 +33,22 @@ func (h *Handler) HandleChatCompletions(w http.ResponseWriter, r *http.Request) writeRequestBodyError(w, body.err) return } - if err := json.NewDecoder(body).Decode(&req); err != nil { - if isRequestBodyTooLarge(err) { - writeRequestBodyError(w, err) + rawBody, readErr := readConversationBody(body, maxChatRequestBytes) + if readErr != nil { + if isRequestBodyTooLarge(readErr) { + writeRequestBodyError(w, readErr) return } + writeErr(w, http.StatusBadRequest, "invalid_request", readErr.Error()) + return + } + if err := json.Unmarshal(rawBody, &req); err != nil { writeErr(w, http.StatusBadRequest, "invalid_request", err.Error()) return } + if id := conversationIDFromRequest(r, rawBody); id != "" { + r = r.WithContext(WithConversationID(r.Context(), id)) + } if len(req.Messages) == 0 { writeErr(w, http.StatusBadRequest, "invalid_request", "messages required") return @@ -282,6 +293,7 @@ func (h *Handler) finishRequestLog(requestID string, started time.Time, req tran entry.ErrorMessage = classified.Message } h.Recorder.Finish(entry) + logCacheDiagnostics(entry, routing) if stats != nil && entry.Credits != nil { h.Recorder.UsageDetail(accounts.RequestUsageDetail{ RequestID: requestID, @@ -344,3 +356,57 @@ func (h *Handler) recordStreamDiagnostic(requestID string, response *http.Respon } func ptrInt(value int) *int { return &value } + + +// logCacheDiagnostics mirrors qoder2api-hub's split between route-level session +// affinity and real upstream prompt-cache token hits. A bound session can still +// miss the cache; without both numbers an operator cannot tell which one broke. +// Opt-in so high-throughput gateways do not pay for per-request logging. +func logCacheDiagnostics(entry accounts.RequestLog, routing string) { + if !affinityDebugOn() || entry.Status == accounts.RequestStatusStreaming { + return + } + prompt, cache := 0, 0 + if entry.PromptTokens != nil { + prompt = *entry.PromptTokens + } + if entry.CacheReadTokens != nil { + cache = *entry.CacheReadTokens + } + hit := 0.0 + if prompt > 0 { + hit = float64(cache) * 100 / float64(prompt) + } + log.Printf("cache diagnostics request_id=%q model=%q account=%q routing=%q affinity=%q prompt=%d cache_read=%d hit=%.1f%% ttfb_ms=%v status=%q", + entry.ID, entry.RequestedModel, entry.AccountID, routing, affinityLabel(routing), prompt, cache, hit, + optionalInt(entry.TTFBMs), entry.Status) +} + +func affinityLabel(routing string) string { + switch routing { + case "sticky": + return "memory-prefix" + case "sticky_escape": + return "escape" + case "pin": + return "pinned" + default: + return "random" + } +} + +func optionalInt(value *int) any { + if value == nil { + return "n/a" + } + return *value +} + +func affinityDebugOn() bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv("CLI2API_AFFINITY_DEBUG"))) { + case "1", "true", "on", "yes": + return true + default: + return false + } +} diff --git a/internal/gateway/prepare.go b/internal/gateway/prepare.go index ca802ef7..6cce6438 100644 --- a/internal/gateway/prepare.go +++ b/internal/gateway/prepare.go @@ -75,6 +75,7 @@ func (h *Handler) prepareExecution(r *http.Request, request translate.ChatReques PreferAccount: prefer, SessionHeader: sessionHeader, + ConversationID: ConversationIDFromContext(ctx), CrossProviderPool: h.crossProviderPoolOn(), ModelContexts: modelContexts, Catalogs: catalogs, diff --git a/internal/gateway/responses.go b/internal/gateway/responses.go index bb325b74..51612bd2 100644 --- a/internal/gateway/responses.go +++ b/internal/gateway/responses.go @@ -39,6 +39,9 @@ func (h *Handler) HandleResponses(w http.ResponseWriter, r *http.Request) { writeErr(w, http.StatusRequestEntityTooLarge, "invalid_request", "request body too large") return } + if id := conversationIDFromRequest(r, body); id != "" { + r = r.WithContext(WithConversationID(r.Context(), id)) + } native, err := translate.ParseNativeResponses(body) if err != nil { writeErr(w, http.StatusBadRequest, "invalid_request", err.Error()) diff --git a/internal/gateway/session_context.go b/internal/gateway/session_context.go new file mode 100644 index 00000000..0d2e57c5 --- /dev/null +++ b/internal/gateway/session_context.go @@ -0,0 +1,108 @@ +package gateway + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" +) + +type conversationContextKey struct{} + +// WithConversationID records a client-supplied conversation/session id found +// in the decoded request. Headers win; body aliases are the compatibility path +// for clients that never send X-CLI2API-Session. +func WithConversationID(ctx context.Context, id string) context.Context { + id = strings.TrimSpace(id) + if id == "" || len(id) > 512 { + return ctx + } + return context.WithValue(ctx, conversationContextKey{}, id) +} + +func ConversationIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + id, _ := ctx.Value(conversationContextKey{}).(string) + return strings.TrimSpace(id) +} + +var conversationHeaders = []string{ + "X-CLI2API-Session", + "X-Conversation-Id", "X-Conversation-ID", + "X-Session-Id", "X-Session-ID", + "X-Thread-Id", "X-Thread-ID", +} + +var conversationBodyKeys = []string{ + "conversation_id", "conversationId", "conversation-id", + "session_id", "sessionId", "session-id", + "thread_id", "threadId", "conversation", "session", +} + +// conversationIDFromRequest reads the conversation id without consuming the +// body. Callers pass the already-decoded bytes when available. +func conversationIDFromRequest(r *http.Request, body []byte) string { + if r != nil { + for _, name := range conversationHeaders { + if id := cleanID(r.Header.Get(name)); id != "" { + return id + } + } + } + if len(body) == 0 || len(body) > 64<<20 { + return "" + } + var top map[string]json.RawMessage + if json.Unmarshal(body, &top) != nil { + return "" + } + for _, key := range conversationBodyKeys { + if id, ok := stringField(top[key]); ok { + return id + } + } + for _, metaKey := range []string{"metadata", "meta", "context"} { + var meta map[string]json.RawMessage + if json.Unmarshal(top[metaKey], &meta) != nil { + continue + } + for _, key := range conversationBodyKeys { + if id, ok := stringField(meta[key]); ok { + return id + } + } + } + return "" +} + +func stringField(raw json.RawMessage) (string, bool) { + if len(raw) == 0 { + return "", false + } + var value string + if json.Unmarshal(raw, &value) != nil { + return "", false + } + if id := cleanID(value); id != "" { + return id, true + } + return "", false +} + +func cleanID(value string) string { + value = strings.TrimSpace(value) + if value == "" || len(value) > 512 { + return "" + } + return value +} + +// readConversationBody consumes a request body reader and returns the bytes so +// handlers can both decode JSON and derive the conversation id. It is only used +// for the request body, never the upstream response stream. +func readConversationBody(reader io.Reader, limit int64) ([]byte, error) { + return io.ReadAll(io.LimitReader(reader, limit)) +} diff --git a/internal/gateway/session_context_test.go b/internal/gateway/session_context_test.go new file mode 100644 index 00000000..adb81b1b --- /dev/null +++ b/internal/gateway/session_context_test.go @@ -0,0 +1,53 @@ +package gateway + +import ( + "net/http/httptest" + "testing" +) + +func TestConversationIDFromHeaderAliases(t *testing.T) { + for _, name := range []string{"X-CLI2API-Session", "X-Conversation-Id", "X-Session-ID", "X-Thread-Id"} { + req := httptest.NewRequest("POST", "/v1/chat/completions", nil) + req.Header.Set(name, "chat-42") + if got := conversationIDFromRequest(req, nil); got != "chat-42" { + t.Fatalf("%s = %q", name, got) + } + } +} + +func TestConversationIDFromBodyAliases(t *testing.T) { + for _, body := range []string{ + `{"conversation_id":"a"}`, + `{"conversationId":"b"}`, + `{"session_id":"c"}`, + `{"sessionId":"d"}`, + `{"thread_id":"e"}`, + `{"metadata":{"conversation_id":"f"}}`, + `{"meta":{"conversationId":"g"}}`, + `{"context":{"sessionId":"h"}}`, + } { + req := httptest.NewRequest("POST", "/v1/chat/completions", nil) + got := conversationIDFromRequest(req, []byte(body)) + if got == "" { + t.Fatalf("no id from %s", body) + } + } +} + +func TestConversationIDHeaderWinsAndBoundsValue(t *testing.T) { + req := httptest.NewRequest("POST", "/v1/messages", nil) + req.Header.Set("X-Conversation-ID", "header-wins") + if got := conversationIDFromRequest(req, []byte(`{"metadata":{"conversation_id":"body"}}`)); got != "header-wins" { + t.Fatalf("got %q", got) + } + if got := conversationIDFromRequest(req, []byte(`{"conversation_id":"`+string(make([]byte, 600))+`"}`)); got != "header-wins" { + t.Fatalf("oversize value leaked: %q", got) + } +} + +func TestWithConversationIDRoundTrip(t *testing.T) { + ctx := WithConversationID(httptest.NewRequest("GET", "/", nil).Context(), " abc ") + if got := ConversationIDFromContext(ctx); got != "abc" { + t.Fatalf("got %q", got) + } +} diff --git a/internal/providers/workbuddy/payload.go b/internal/providers/workbuddy/payload.go index 9a728bdb..ddf6abbe 100644 --- a/internal/providers/workbuddy/payload.go +++ b/internal/providers/workbuddy/payload.go @@ -437,6 +437,10 @@ func rewriteIdentitySentences(text string) string { // (empty content is code 11151); ensureLeadingSystem then restores a // non-empty leading system prompt if needed. func sanitizeWAFKeywords(body map[string]any) { + if !bodyMayNeedWAFRewrite(body) { + return + } + raw, ok := body["messages"] if !ok { return @@ -527,3 +531,65 @@ func sanitizeBillingHeaders(text string) string { } return wafBillingHdrBareRe.ReplaceAllString(text, "anthropic-header") } + + +// bodyMayNeedWAFRewrite is a cheap prefilter. Most requests have none of the +// client fingerprints the WAF rules rewrite; running a dozen regexes over a +// multi-hundred-kilobyte history on every turn is pure CPU and, worse, can +// change the cached prefix between otherwise identical turns. +func bodyMayNeedWAFRewrite(body map[string]any) bool { + raw, ok := body["messages"] + if !ok { + return false + } + list, ok := raw.([]any) + if !ok { + return false + } + for _, item := range list { + message, ok := item.(map[string]any) + if !ok { + continue + } + for _, key := range []string{"content", "reasoning_content"} { + if textMayNeedWAFRewrite(message[key]) { + return true + } + } + } + return false +} + +func textMayNeedWAFRewrite(value any) bool { + switch typed := value.(type) { + case string: + return wafTextMayNeedRewrite(typed) + case []any: + for _, part := range typed { + object, ok := part.(map[string]any) + if !ok { + continue + } + if text, ok := object["text"].(string); ok && wafTextMayNeedRewrite(text) { + return true + } + } + } + return false +} + +func wafTextMayNeedRewrite(text string) bool { + if text == "" { + return false + } + lower := strings.ToLower(text) + for _, needle := range []string{ + "main branch", "powered by", "built on", "built by", "created by", "developed by", "made by", "trained by", + "deepseek", "x-anthropic-billing", "cc_", "tokens prefixed with @", + } { + if strings.Contains(lower, needle) { + return true + } + } + return false +} diff --git a/internal/providers/workbuddy/waf_fingerprint_test.go b/internal/providers/workbuddy/waf_fingerprint_test.go index da586b15..081911a9 100644 --- a/internal/providers/workbuddy/waf_fingerprint_test.go +++ b/internal/providers/workbuddy/waf_fingerprint_test.go @@ -314,3 +314,31 @@ func fmtContent(v any) string { s, _ := v.(string) return s } + + +func TestBodyMayNeedWAFRewritePrefilter(t *testing.T) { + plain := map[string]any{"messages": []any{ + map[string]any{"role": "user", "content": "write a function"}, + map[string]any{"role": "assistant", "content": "done"}, + }} + if bodyMayNeedWAFRewrite(plain) { + t.Fatal("plain history must skip WAF regexes") + } + for _, text := range []string{ + "Main branch (target for pull requests): main", + "You are an agent powered by DeepSeek Harness", + "x-anthropic-billing-header: cc_version=1;", + "Tokens prefixed with @ are paths the user explicitly referenced.", + } { + body := map[string]any{"messages": []any{map[string]any{"role": "system", "content": text}}} + if !bodyMayNeedWAFRewrite(body) { + t.Fatalf("prefilter missed %q", text) + } + } + parts := map[string]any{"messages": []any{ + map[string]any{"role": "system", "content": []any{map[string]any{"type": "text", "text": "DeepSeek Harness"}}}, + }} + if !bodyMayNeedWAFRewrite(parts) { + t.Fatal("prefilter missed content parts") + } +} diff --git a/internal/translate/session_seed.go b/internal/translate/session_seed.go index a6074468..e22599fb 100644 --- a/internal/translate/session_seed.go +++ b/internal/translate/session_seed.go @@ -260,3 +260,30 @@ func compactSessionJSON(raw json.RawMessage) string { } return compact.String() } + + +// PrefixSessionSeed is the cache-affine part of ContentSessionSeed: model, +// tools and the leading system/developer prefix only. It deliberately excludes +// the first user message so workbuddy2api-hub-style sharding does not scatter +// one (key, model, system, tools) tuple across the whole account pool. +func PrefixSessionSeed(req ChatRequest) string { + parts := []string{ + strings.ToLower(strings.TrimSpace(req.Model)), + compactSessionJSON(req.Tools), + } + for _, message := range req.Messages { + role := strings.ToLower(strings.TrimSpace(message.Role)) + if role != "system" && role != "developer" { + break + } + for _, part := range sessionContentParts(message.Content) { + parts = append(parts, role, part.Type, part.Text, part.Image) + } + } + raw, err := json.Marshal(parts) + if err != nil { + return "" + } + sum := sha256.Sum256(raw) + return hex.EncodeToString(sum[:]) +} diff --git a/internal/translate/thinking_rectifier.go b/internal/translate/thinking_rectifier.go new file mode 100644 index 00000000..ba4add29 --- /dev/null +++ b/internal/translate/thinking_rectifier.go @@ -0,0 +1,188 @@ +package translate + +import ( + "encoding/json" + "fmt" + "strings" +) + +// RectifyResult mirrors cc-switch's thinking rectifier counters. The gateway +// only ever removes fields that OpenAI-compatible upstreams cannot validate; +// it never invents signatures or rewrites model output. +type RectifyResult struct { + Applied bool `json:"applied"` + RemovedThinkingBlocks int `json:"removed_thinking_blocks,omitempty"` + RemovedRedactedThinkingBlocks int `json:"removed_redacted_thinking_blocks,omitempty"` + RemovedSignatureFields int `json:"removed_signature_fields,omitempty"` + RemovedTopLevelThinking bool `json:"removed_top_level_thinking,omitempty"` +} + +// ShouldRectifyThinking mirrors cc-switch's error-triggered detection. It is +// deliberately narrow: only upstream errors that name an invalid/missing +// thinking signature or a malformed thinking block sequence qualify. +func ShouldRectifyThinking(message string) bool { + lower := strings.ToLower(strings.TrimSpace(message)) + if lower == "" { + return false + } + if strings.Contains(lower, "invalid") && strings.Contains(lower, "signature") && strings.Contains(lower, "thinking") { + return true + } + if strings.Contains(lower, "thought signature") && (strings.Contains(lower, "not valid") || strings.Contains(lower, "invalid")) { + return true + } + if strings.Contains(lower, "must start with a thinking block") { + return true + } + if strings.Contains(lower, "expected") && (strings.Contains(lower, "thinking") || strings.Contains(lower, "redacted_thinking")) && strings.Contains(lower, "tool_use") { + return true + } + if strings.Contains(lower, "signature") && (strings.Contains(lower, "field required") || strings.Contains(lower, "extra inputs are not permitted")) { + return true + } + if (strings.Contains(lower, "thinking") || strings.Contains(lower, "redacted_thinking")) && strings.Contains(lower, "cannot be modified") { + return true + } + if strings.Contains(lower, "budget_tokens") && strings.Contains(lower, "thinking") && (strings.Contains(lower, "1024") || strings.Contains(lower, "input should be")) { + return true + } + return false +} + +// RectifyAnthropicRequest applies the minimal cc-switch cleanup to a decoded +// Anthropic Messages body. Historical thinking blocks carry signatures that +// belong to a different provider; replaying them upstream causes 400 errors or +// visible scaffolding in clients. Plain text, tool calls and tool +// results are preserved byte-for-byte where possible. +func RectifyAnthropicRequest(body []byte) ([]byte, RectifyResult, error) { + var root map[string]json.RawMessage + if err := json.Unmarshal(body, &root); err != nil { + return nil, RectifyResult{}, err + } + result := RectifyResult{} + messagesRaw := root["messages"] + var messages []json.RawMessage + if len(messagesRaw) > 0 && string(messagesRaw) != "null" { + if err := json.Unmarshal(messagesRaw, &messages); err != nil { + return nil, result, fmt.Errorf("messages: %w", err) + } + } + changed := false + for i, raw := range messages { + var message map[string]json.RawMessage + if err := json.Unmarshal(raw, &message); err != nil { + continue + } + contentRaw := message["content"] + var blocks []json.RawMessage + if len(contentRaw) == 0 || json.Unmarshal(contentRaw, &blocks) != nil { + continue + } + kept := make([]json.RawMessage, 0, len(blocks)) + messageChanged := false + for _, blockRaw := range blocks { + var block map[string]json.RawMessage + if json.Unmarshal(blockRaw, &block) != nil { + kept = append(kept, blockRaw) + continue + } + kind := rawMapString(block, "type") + if kind == "thinking" { + result.RemovedThinkingBlocks++ + messageChanged = true + continue + } + if kind == "redacted_thinking" { + result.RemovedRedactedThinkingBlocks++ + messageChanged = true + continue + } + if _, exists := block["signature"]; exists { + delete(block, "signature") + result.RemovedSignatureFields++ + messageChanged = true + encoded, err := json.Marshal(block) + if err != nil { + return nil, result, err + } + kept = append(kept, encoded) + continue + } + kept = append(kept, blockRaw) + } + if messageChanged { + encoded, err := json.Marshal(kept) + if err != nil { + return nil, result, err + } + message["content"] = encoded + encodedMessage, err := json.Marshal(message) + if err != nil { + return nil, result, err + } + messages[i] = encodedMessage + changed = true + } + } + // cc-switch fallback: when thinking is enabled and the last assistant turn + // starts with a tool_use instead of a thinking block, drop the top-level + // switch rather than replaying an invalid sequence. + if shouldRemoveTopLevelThinking(root, messages) { + delete(root, "thinking") + result.RemovedTopLevelThinking = true + changed = true + } + if changed { + encodedMessages, err := json.Marshal(messages) + if err != nil { + return nil, result, err + } + root["messages"] = encodedMessages + out, err := json.Marshal(root) + if err != nil { + return nil, result, err + } + result.Applied = true + return out, result, nil + } + return body, result, nil +} + +func shouldRemoveTopLevelThinking(root map[string]json.RawMessage, messages []json.RawMessage) bool { + var thinking map[string]json.RawMessage + if json.Unmarshal(root["thinking"], &thinking) != nil { + return false + } + if rawMapString(thinking, "type") != "enabled" { + return false + } + for i := len(messages) - 1; i >= 0; i-- { + var message map[string]json.RawMessage + if json.Unmarshal(messages[i], &message) != nil { + continue + } + if rawMapString(message, "role") != "assistant" { + continue + } + var blocks []json.RawMessage + if json.Unmarshal(message["content"], &blocks) != nil || len(blocks) == 0 { + return false + } + var first map[string]json.RawMessage + if json.Unmarshal(blocks[0], &first) != nil { + return false + } + kind := rawMapString(first, "type") + if kind == "thinking" || kind == "redacted_thinking" { + return false + } + for _, raw := range blocks { + var block map[string]json.RawMessage + if json.Unmarshal(raw, &block) == nil && rawMapString(block, "type") == "tool_use" { + return true + } + } + return false + } + return false +} diff --git a/internal/translate/thinking_rectifier_test.go b/internal/translate/thinking_rectifier_test.go new file mode 100644 index 00000000..4d36adad --- /dev/null +++ b/internal/translate/thinking_rectifier_test.go @@ -0,0 +1,79 @@ +package translate + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestShouldRectifyThinkingMatchesKnownErrors(t *testing.T) { + for _, message := range []string{ + "Invalid 'signature' in 'thinking' block", + "Thought signature is not valid", + "messages.1.content.0.thinking.signature: Field required", + "Expected `thinking` or `redacted_thinking`, but found `tool_use`", + "thinking.budget_tokens: Input should be greater than or equal to 1024", + } { + if !ShouldRectifyThinking(message) { + t.Fatalf("missed %q", message) + } + } + for _, message := range []string{"", "rate limited", "connection reset", "invalid api key"} { + if ShouldRectifyThinking(message) { + t.Fatalf("false positive %q", message) + } + } +} + +func TestRectifyAnthropicRequestDropsHistoricalThinking(t *testing.T) { + body := []byte(`{"model":"hy3","thinking":{"type":"enabled","budget_tokens":1024},"max_tokens":64, + "messages":[ + {"role":"user","content":[{"type":"text","text":"hello"}]}, + {"role":"assistant","content":[ + {"type":"thinking","thinking":"secret chain","signature":"sig-from-other-provider"}, + {"type":"text","text":"answer"}, + {"type":"tool_use","id":"call_1","name":"Read","input":{}} + ]} + ]}`) + out, result, err := RectifyAnthropicRequest(body) + if err != nil { + t.Fatal(err) + } + if !result.Applied || result.RemovedThinkingBlocks != 1 || !result.RemovedTopLevelThinking { + t.Fatalf("result = %+v", result) + } + if strings.Contains(string(out), "thinking") && strings.Contains(string(out), "secret chain") { + t.Fatalf("thinking content leaked: %s", out) + } + if !strings.Contains(string(out), "tool_use") || !strings.Contains(string(out), "answer") { + t.Fatalf("normal content lost: %s", out) + } + var source AnthropicMessagesRequest + if err := json.Unmarshal(out, &source); err != nil { + t.Fatal(err) + } + chat, err := TranslateAnthropicMessages(source) + if err != nil { + t.Fatal(err) + } + if len(chat.Messages) != 2 { + t.Fatalf("messages = %d", len(chat.Messages)) + } +} + +func TestRectifyAnthropicRequestIsIdempotentAndPreservesPlainBodies(t *testing.T) { + plain := []byte(`{"model":"hy3","messages":[{"role":"user","content":"hi"}]}`) + out, result, err := RectifyAnthropicRequest(plain) + if err != nil || result.Applied || string(out) != string(plain) { + t.Fatalf("plain body changed: %+v %s", result, out) + } + withSignature := []byte(`{"model":"hy3","messages":[{"role":"assistant","content":[{"type":"text","text":"ok","signature":"stale"}]}]}`) + first, r1, err := RectifyAnthropicRequest(withSignature) + if err != nil || !r1.Applied || r1.RemovedSignatureFields != 1 { + t.Fatalf("first = %+v", r1) + } + second, r2, err := RectifyAnthropicRequest(first) + if err != nil || r2.Applied || string(second) != string(first) { + t.Fatalf("not idempotent: %+v", r2) + } +}