Add projects

This commit is contained in:
Alois 2026-08-24 00:10:41 +02:00
commit 8b607dd700
Signed by: alois
SSH key fingerprint: SHA256:GBzT2DXvAuGV9XIV5W3WrzVpjU54FThmxHXdbz95J24
1802 changed files with 503346 additions and 2 deletions

View file

@ -0,0 +1,104 @@
package helps
import (
"context"
"fmt"
"net/http"
"net/url"
"strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
func isAntigravityVertexSearchRedirect(rawURL string) bool {
parsed, err := url.Parse(rawURL)
if err != nil {
return false
}
return parsed.Scheme == "https" &&
parsed.Host == "vertexaisearch.cloud.google.com" &&
strings.HasPrefix(parsed.Path, "/grounding-api-redirect/")
}
func resolveAntigravityGroundingURL(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, rawURL string) string {
if !isAntigravityVertexSearchRedirect(rawURL) {
return rawURL
}
client := NewProxyAwareHTTPClient(ctx, cfg, auth, 0)
client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
return http.ErrUseLastResponse
}
req, errReq := http.NewRequestWithContext(ctx, http.MethodHead, rawURL, nil)
if errReq != nil {
log.WithError(errReq).Debug("antigravity grounding url: create redirect request failed")
return rawURL
}
resp, errDo := client.Do(req)
if errDo != nil {
log.WithError(errDo).Debug("antigravity grounding url: resolve redirect failed")
return rawURL
}
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
log.WithError(errClose).Debug("antigravity grounding url: close redirect response failed")
}
}()
if resp.StatusCode < http.StatusMultipleChoices || resp.StatusCode >= http.StatusBadRequest {
return rawURL
}
location := strings.TrimSpace(resp.Header.Get("Location"))
if location == "" {
return rawURL
}
parsed, errParse := url.Parse(location)
if errParse != nil || parsed.Scheme != "https" || parsed.Host == "" {
return rawURL
}
return location
}
// ResolveAntigravityGroundingURLs replaces Vertex Search redirect URLs in grounding chunks with their target URLs.
func ResolveAntigravityGroundingURLs(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, payload []byte) []byte {
if len(payload) == 0 {
return payload
}
basePath := "response.candidates.0.groundingMetadata.groundingChunks"
chunks := gjson.GetBytes(payload, basePath)
if !chunks.IsArray() {
basePath = "candidates.0.groundingMetadata.groundingChunks"
chunks = gjson.GetBytes(payload, basePath)
}
if !chunks.IsArray() {
return payload
}
output := payload
resolved := map[string]string{}
for i, chunk := range chunks.Array() {
uri := strings.TrimSpace(chunk.Get("web.uri").String())
if uri == "" {
continue
}
resolvedURI, ok := resolved[uri]
if !ok {
resolvedURI = resolveAntigravityGroundingURL(ctx, cfg, auth, uri)
resolved[uri] = resolvedURI
}
if resolvedURI == uri {
continue
}
updated, errSet := sjson.SetBytes(output, fmt.Sprintf("%s.%d.web.uri", basePath, i), resolvedURI)
if errSet != nil {
log.WithError(errSet).Debug("antigravity grounding url: set resolved url failed")
continue
}
output = updated
}
return output
}

View file

@ -0,0 +1,66 @@
package helps
import (
"context"
"io"
"net/http"
"strings"
"testing"
"github.com/tidwall/gjson"
)
type groundingURLRoundTripper func(*http.Request) (*http.Response, error)
func (f groundingURLRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
func TestResolveAntigravityGroundingURLsResolvesVertexRedirects(t *testing.T) {
t.Parallel()
const redirectURL = "https://vertexaisearch.cloud.google.com/grounding-api-redirect/example-token"
const resolvedURL = "https://example.com/weather"
var sawRedirectRequest bool
ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", groundingURLRoundTripper(func(req *http.Request) (*http.Response, error) {
if req.Method != http.MethodHead {
t.Fatalf("method = %s, want HEAD", req.Method)
}
if req.URL.String() != redirectURL {
t.Fatalf("url = %s, want %s", req.URL.String(), redirectURL)
}
sawRedirectRequest = true
return &http.Response{
StatusCode: http.StatusFound,
Header: http.Header{
"Location": []string{resolvedURL},
},
Body: io.NopCloser(strings.NewReader("")),
}, nil
}))
input := []byte(`{
"response": {
"candidates": [{
"groundingMetadata": {
"groundingChunks": [
{"web": {"uri": "` + redirectURL + `", "title": "Weather"}},
{"web": {"uri": "https://already.example/source", "title": "Existing"}}
]
}
}]
}
}`)
output := ResolveAntigravityGroundingURLs(ctx, nil, nil, input)
if !sawRedirectRequest {
t.Fatal("expected resolver to request the vertex redirect")
}
if got := gjson.GetBytes(output, "response.candidates.0.groundingMetadata.groundingChunks.0.web.uri").String(); got != resolvedURL {
t.Fatalf("resolved uri = %q, want %q; output=%s", got, resolvedURL, output)
}
if got := gjson.GetBytes(output, "response.candidates.0.groundingMetadata.groundingChunks.1.web.uri").String(); got != "https://already.example/source" {
t.Fatalf("non-vertex uri = %q", got)
}
}

View file

@ -0,0 +1,128 @@
package helps
import (
"context"
"sync"
"time"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
)
type CodexCache struct {
ID string
Expire time.Time
}
// codexCacheMap stores prompt cache IDs keyed by model+user_id.
// Protected by codexCacheMu. Entries expire after 1 hour.
var (
codexCacheMap = make(map[string]CodexCache)
codexCacheMu sync.RWMutex
)
// codexCacheCleanupInterval controls how often expired entries are purged.
const codexCacheCleanupInterval = 15 * time.Minute
// codexCacheCleanupOnce ensures the background cleanup goroutine starts only once.
var codexCacheCleanupOnce sync.Once
// startCodexCacheCleanup launches a background goroutine that periodically
// removes expired entries from codexCacheMap to prevent memory leaks.
func startCodexCacheCleanup() {
go func() {
ticker := time.NewTicker(codexCacheCleanupInterval)
defer ticker.Stop()
for range ticker.C {
purgeExpiredCodexCache()
}
}()
}
// purgeExpiredCodexCache removes entries that have expired.
func purgeExpiredCodexCache() {
now := time.Now()
codexCacheMu.Lock()
defer codexCacheMu.Unlock()
for key, cache := range codexCacheMap {
if cache.Expire.Before(now) {
delete(codexCacheMap, key)
}
}
}
// GetCodexCache retrieves a cached entry, returning ok=false if not found or expired.
func GetCodexCache(key string) (CodexCache, bool) {
cache, ok, err := GetCodexCacheRequired(context.Background(), key)
if err == nil {
return cache, ok
}
return CodexCache{}, false
}
// GetCodexCacheRequired retrieves a cached entry for request-time paths.
func GetCodexCacheRequired(ctx context.Context, key string) (CodexCache, bool, error) {
var homeCache CodexCache
homeMode, found, errGet := homekv.KVGetJSONRequired(ctx, key, &homeCache)
if homeMode {
if errGet != nil || !found {
return CodexCache{}, false, errGet
}
if homeCache.Expire.Before(time.Now()) {
_, _, _ = homekv.KVDelRequired(ctx, key)
return CodexCache{}, false, nil
}
return homeCache, true, nil
}
codexCacheCleanupOnce.Do(startCodexCacheCleanup)
codexCacheMu.RLock()
cache, ok := codexCacheMap[key]
codexCacheMu.RUnlock()
if !ok || cache.Expire.Before(time.Now()) {
return CodexCache{}, false, nil
}
return cache, true, nil
}
// SetCodexCache stores a cache entry.
func SetCodexCache(key string, cache CodexCache) {
SetCodexCacheBestEffort(context.Background(), key, cache)
}
// SetCodexCacheRequired stores a cache entry for request-time paths.
func SetCodexCacheRequired(ctx context.Context, key string, cache CodexCache) error {
ttl := time.Until(cache.Expire)
if ttl <= 0 {
return nil
}
if _, homeMode, _ := homekv.CurrentKVClient(); homeMode {
_, errSet := homekv.KVSetJSONRequired(ctx, key, cache, ttl)
return errSet
}
codexCacheCleanupOnce.Do(startCodexCacheCleanup)
codexCacheMu.Lock()
codexCacheMap[key] = cache
codexCacheMu.Unlock()
return nil
}
// SetCodexCacheBestEffort stores a cache entry without failing completed responses.
func SetCodexCacheBestEffort(ctx context.Context, key string, cache CodexCache) bool {
ttl := time.Until(cache.Expire)
if ttl <= 0 {
return false
}
if _, homeMode, _ := homekv.CurrentKVClient(); homeMode {
return homekv.KVSetJSONBestEffort(ctx, key, cache, ttl)
}
codexCacheCleanupOnce.Do(startCodexCacheCleanup)
codexCacheMu.Lock()
codexCacheMap[key] = cache
codexCacheMu.Unlock()
return true
}
// CodexPromptCacheKey builds the Home KV key for a model/user prompt cache.
func CodexPromptCacheKey(modelName string, userScope string) string {
return "cpa:codex:prompt-cache:" + homekv.HashKeyPart(modelName) + ":" + homekv.HashKeyPart(userScope)
}

View file

@ -0,0 +1,27 @@
package helps
import (
"context"
"strings"
"testing"
"time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
)
func TestSetCodexCacheRequiredHomeUnavailableReturnsError(t *testing.T) {
homekv.SetCurrent(homekv.New(config.HomeConfig{Enabled: false}))
t.Cleanup(homekv.ClearCurrent)
errSet := SetCodexCacheRequired(context.Background(), "cpa:codex:prompt-cache:test", CodexCache{
ID: "cache-id",
Expire: time.Now().Add(time.Hour),
})
if errSet == nil {
t.Fatal("SetCodexCacheRequired() error = nil, want home kv unavailable error")
}
if !strings.Contains(errSet.Error(), "home kv store unavailable") {
t.Fatalf("SetCodexCacheRequired() error = %v, want home kv store unavailable", errSet)
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,66 @@
package helps
import (
"strings"
"github.com/tidwall/gjson"
)
var defaultClaudeBuiltinToolNames = []string{
"web_search",
"code_execution",
"text_editor",
"computer",
}
func newClaudeBuiltinToolRegistry() map[string]bool {
registry := make(map[string]bool, len(defaultClaudeBuiltinToolNames))
for _, name := range defaultClaudeBuiltinToolNames {
registry[name] = true
}
return registry
}
// IsClaudeServerToolType reports whether a typed declaration is a recognized
// Anthropic-operated tool. Client-defined type:"custom" declarations are not
// server tools and must remain eligible for MCP aliasing.
func IsClaudeServerToolType(toolType string) bool {
toolType = strings.ToLower(strings.TrimSpace(toolType))
for _, prefix := range []string{
"advisor_",
"agent_toolset_",
"bash_",
"code_execution_",
"computer_",
"memory_",
"text_editor_",
"tool_search_tool_",
"web_fetch_",
"web_search_",
} {
if strings.HasPrefix(toolType, prefix) {
return true
}
}
return false
}
func AugmentClaudeBuiltinToolRegistry(body []byte, registry map[string]bool) map[string]bool {
if registry == nil {
registry = newClaudeBuiltinToolRegistry()
}
tools := gjson.GetBytes(body, "tools")
if !tools.Exists() || !tools.IsArray() {
return registry
}
tools.ForEach(func(_, tool gjson.Result) bool {
if !IsClaudeServerToolType(tool.Get("type").String()) {
return true
}
if name := tool.Get("name").String(); name != "" {
registry[name] = true
}
return true
})
return registry
}

View file

@ -0,0 +1,56 @@
package helps
import "testing"
func TestClaudeBuiltinToolRegistry_DefaultSeedFallback(t *testing.T) {
registry := AugmentClaudeBuiltinToolRegistry(nil, nil)
for _, name := range defaultClaudeBuiltinToolNames {
if !registry[name] {
t.Fatalf("default builtin %q missing from fallback registry", name)
}
}
}
func TestClaudeBuiltinToolRegistry_AugmentsKnownTypedBuiltinsFromBody(t *testing.T) {
registry := AugmentClaudeBuiltinToolRegistry([]byte(`{
"tools": [
{"type": "web_search_20250305", "name": "web_search"},
{"type": "custom", "name": "client_custom"},
{"type": "custom_builtin_20250401", "name": "unknown_typed"},
{"name": "Read"}
]
}`), nil)
if !registry["web_search"] {
t.Fatal("expected known typed builtin web_search in registry")
}
for _, name := range []string{"client_custom", "unknown_typed", "Read"} {
if registry[name] {
t.Fatalf("expected client tool %q to stay out of builtin registry", name)
}
}
}
func TestIsClaudeServerToolType(t *testing.T) {
for _, toolType := range []string{
"web_search_20250305",
"code_execution_20250522",
"tool_search_tool_regex_20251119",
"advisor_20260301",
"agent_toolset_20260401",
"bash_20250124",
"text_editor_20250728",
"memory_20250818",
"computer_20241022",
"web_fetch_20260209",
} {
if !IsClaudeServerToolType(toolType) {
t.Fatalf("IsClaudeServerToolType(%q) = false, want true", toolType)
}
}
for _, toolType := range []string{"", "custom", "custom_builtin_20250401"} {
if IsClaudeServerToolType(toolType) {
t.Fatalf("IsClaudeServerToolType(%q) = true, want false", toolType)
}
}
}

View file

@ -0,0 +1,102 @@
package helps
import (
"crypto/sha256"
"encoding/hex"
"fmt"
"strings"
"github.com/google/uuid"
claudeauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
)
// Stable identity seeds for fingerprint-profile=claude-code-cli on non-OAuth credentials.
// Real OAuth credentials keep their stored account/device pool; this only fills gaps
// so ApplyClaudeCredentialMetadata can run as the single identity algorithm.
var claudeCLIIdentityNamespace = uuid.MustParse("6ba7b812-9dad-11d1-80b4-00c04fd430c8")
func stableClaudeCLIDeviceID(seed string) string {
sum := sha256.Sum256([]byte("cpa-claude-code-cli-device|" + seed))
return hex.EncodeToString(sum[:])
}
// StableClaudeCLIDeviceID returns a deterministic device ID derived from a seed.
func StableClaudeCLIDeviceID(seed string) string {
return stableClaudeCLIDeviceID(seed)
}
func stableClaudeCLIAccountUUID(seed string) string {
return uuid.NewSHA1(claudeCLIIdentityNamespace, []byte("cpa-claude-code-cli-account|"+seed)).String()
}
// StableClaudeCLIAccountUUID returns a deterministic UUIDv5 account ID derived from a seed.
func StableClaudeCLIAccountUUID(seed string) string {
return stableClaudeCLIAccountUUID(seed)
}
// ClaudeCLIAuthIdentitySeed returns a stable credential identity that does not
// rotate with delegated-provider access tokens.
func ClaudeCLIAuthIdentitySeed(auth *cliproxyauth.Auth) string {
if auth != nil {
if id := strings.TrimSpace(auth.ID); id != "" {
return "auth-id|" + id
}
if index := strings.TrimSpace(auth.Index); index != "" {
return "auth-index|" + index
}
if fileName := strings.TrimSpace(auth.FileName); fileName != "" {
return "auth-file|" + fileName
}
}
return ""
}
// PrepareClaudeCLIFingerprintAuth returns the auth object that should receive
// ApplyClaudeCredentialMetadata. Synthesized API-key / delegated-provider
// identity is written to a clone so the shared credential metadata map is not
// mutated on the request path.
func PrepareClaudeCLIFingerprintAuth(auth *cliproxyauth.Auth, seed string, synthesizeMissing bool) (*cliproxyauth.Auth, error) {
if auth == nil {
return nil, fmt.Errorf("auth is nil")
}
if !synthesizeMissing {
return auth, nil
}
local := auth.Clone()
if err := EnsureClaudeCLIFingerprintIdentity(local, seed, true); err != nil {
return nil, err
}
return local, nil
}
// EnsureClaudeCLIFingerprintIdentity prepares auth.Metadata so the shared
// ApplyClaudeCredentialMetadata path can run.
//
// When synthesizeMissing is false (real OAuth), this is a no-op: missing account
// or device data must surface as credential errors.
// When synthesizeMissing is true (fingerprint-profile=claude-code-cli on API keys),
// missing account_uuid / device pool are filled with stable values derived from seed.
// Callers that hold a shared Auth must use PrepareClaudeCLIFingerprintAuth instead.
func EnsureClaudeCLIFingerprintIdentity(auth *cliproxyauth.Auth, seed string, synthesizeMissing bool) error {
if auth == nil {
return fmt.Errorf("auth is nil")
}
if !synthesizeMissing {
return nil
}
seed = strings.TrimSpace(seed)
if seed == "" {
seed = "anonymous"
}
if ClaudeCredentialAccountUUID(auth) == "" {
claudeauth.StoreMetadataString(&auth.Metadata, "account_uuid", stableClaudeCLIAccountUUID(seed))
}
if !claudeauth.HasCanonicalDeviceIDPool(claudeauth.ReadDeviceIDPool(&auth.Metadata)) {
claudeauth.StoreDeviceIDPool(&auth.Metadata, []string{stableClaudeCLIDeviceID(seed)})
}
if _, _, errPool := claudeauth.EnsureDeviceIDPoolFor(&auth.Metadata); errPool != nil {
return fmt.Errorf("ensure device pool: %w", errPool)
}
return nil
}

View file

@ -0,0 +1,175 @@
package helps
import (
"sync"
"testing"
claudeauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
"github.com/tidwall/gjson"
)
func TestEnsureClaudeCLIFingerprintIdentitySynthesizesStableSources(t *testing.T) {
t.Parallel()
auth := &cliproxyauth.Auth{}
if err := EnsureClaudeCLIFingerprintIdentity(auth, "key-a", true); err != nil {
t.Fatalf("EnsureClaudeCLIFingerprintIdentity() error = %v", err)
}
account := ClaudeCredentialAccountUUID(auth)
if account == "" {
t.Fatal("account_uuid is empty")
}
deviceIDs, _, errPool := claudeauth.EnsureDeviceIDPoolFor(&auth.Metadata)
if errPool != nil {
t.Fatalf("EnsureDeviceIDPoolFor() error = %v", errPool)
}
if len(deviceIDs) != 1 || deviceIDs[0] != stableClaudeCLIDeviceID("key-a") {
t.Fatalf("device pool = %#v, want stable single device", deviceIDs)
}
// Second call must not rotate identity.
if err := EnsureClaudeCLIFingerprintIdentity(auth, "key-a", true); err != nil {
t.Fatalf("second EnsureClaudeCLIFingerprintIdentity() error = %v", err)
}
if got := ClaudeCredentialAccountUUID(auth); got != account {
t.Fatalf("account_uuid changed: %q vs %q", got, account)
}
const sessionID = "11111111-2222-4333-8444-555555555555"
updated, deviceID, errApply := ApplyClaudeCredentialMetadata([]byte(`{"messages":[]}`), auth, sessionID)
if errApply != nil {
t.Fatalf("ApplyClaudeCredentialMetadata() error = %v", errApply)
}
if deviceID != deviceIDs[0] {
t.Fatalf("selected device = %q, want %q", deviceID, deviceIDs[0])
}
userID := gjson.GetBytes(updated, "metadata.user_id").String()
if !IsValidUserID(userID) {
t.Fatalf("user_id = %q, want valid", userID)
}
if got := gjson.Get(userID, "account_uuid").String(); got != account {
t.Fatalf("user_id account = %q, want %q", got, account)
}
if got := gjson.Get(userID, "session_id").String(); got != sessionID {
t.Fatalf("user_id session = %q, want %q", got, sessionID)
}
}
func TestClaudeCLIAuthIdentitySeedPrefersStableAuthIdentity(t *testing.T) {
t.Parallel()
tests := []struct {
name string
auth *cliproxyauth.Auth
want string
}{
{name: "auth ID", auth: &cliproxyauth.Auth{ID: "kimi-auth"}, want: "auth-id|kimi-auth"},
{name: "auth index", auth: &cliproxyauth.Auth{Index: "kimi-index"}, want: "auth-index|kimi-index"},
{name: "auth file", auth: &cliproxyauth.Auth{FileName: "kimi.json"}, want: "auth-file|kimi.json"},
{name: "missing identity", auth: &cliproxyauth.Auth{}, want: ""},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
if got := ClaudeCLIAuthIdentitySeed(tt.auth); got != tt.want {
t.Fatalf("ClaudeCLIAuthIdentitySeed() = %q, want %q", got, tt.want)
}
})
}
}
func TestPrepareClaudeCLIFingerprintAuthDoesNotMutateSharedMetadata(t *testing.T) {
t.Parallel()
shared := &cliproxyauth.Auth{
ID: "kimi-shared",
Metadata: map[string]any{
"access_token": "token-1",
},
}
prepared, errPrepare := PrepareClaudeCLIFingerprintAuth(shared, ClaudeCLIAuthIdentitySeed(shared), true)
if errPrepare != nil {
t.Fatalf("PrepareClaudeCLIFingerprintAuth() error = %v", errPrepare)
}
if prepared == shared {
t.Fatal("PrepareClaudeCLIFingerprintAuth() returned the shared auth")
}
if ClaudeCredentialAccountUUID(shared) != "" {
t.Fatalf("shared account_uuid = %q, want empty", ClaudeCredentialAccountUUID(shared))
}
if ClaudeCredentialAccountUUID(prepared) == "" {
t.Fatal("prepared account_uuid is empty")
}
if _, ok := shared.Metadata[claudeauth.ClaudeDeviceIDsMetadataKey]; ok {
t.Fatalf("shared metadata gained device pool: %#v", shared.Metadata)
}
}
func TestPrepareClaudeCLIFingerprintAuthIsolatesUnlockedMetadataReaders(t *testing.T) {
shared := &cliproxyauth.Auth{
ID: "kimi-race",
Metadata: map[string]any{
"access_token": "token-1",
"refresh_token": "refresh-1",
},
}
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
for range 200 {
prepared, errPrepare := PrepareClaudeCLIFingerprintAuth(shared, ClaudeCLIAuthIdentitySeed(shared), true)
if errPrepare != nil {
t.Errorf("PrepareClaudeCLIFingerprintAuth() error = %v", errPrepare)
return
}
if ClaudeCredentialAccountUUID(prepared) == "" {
t.Error("prepared account_uuid is empty")
return
}
}
}()
go func() {
defer wg.Done()
for range 200 {
// Same unlocked read Kimi OpenAI-compat requests perform via kimiCreds.
_ = shared.Metadata["access_token"].(string)
}
}()
wg.Wait()
}
func TestEnsureClaudeCLIFingerprintIdentityNoopWithoutSynthesize(t *testing.T) {
t.Parallel()
auth := &cliproxyauth.Auth{}
if err := EnsureClaudeCLIFingerprintIdentity(auth, "key-a", false); err != nil {
t.Fatalf("EnsureClaudeCLIFingerprintIdentity() error = %v", err)
}
if ClaudeCredentialAccountUUID(auth) != "" {
t.Fatal("expected no synthesized account without synthesizeMissing")
}
}
func TestEnsureClaudeCLIFingerprintIdentityPreservesExistingOAuthSources(t *testing.T) {
t.Parallel()
auth := &cliproxyauth.Auth{Metadata: map[string]any{
"account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa",
claudeauth.ClaudeDeviceIDsMetadataKey: []string{"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"},
}}
if err := EnsureClaudeCLIFingerprintIdentity(auth, "key-a", true); err != nil {
t.Fatalf("EnsureClaudeCLIFingerprintIdentity() error = %v", err)
}
if got := ClaudeCredentialAccountUUID(auth); got != "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" {
t.Fatalf("account_uuid = %q, want preserved", got)
}
deviceIDs, _, errPool := claudeauth.EnsureDeviceIDPoolFor(&auth.Metadata)
if errPool != nil {
t.Fatalf("EnsureDeviceIDPoolFor() error = %v", errPool)
}
if deviceIDs[0] != "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" {
t.Fatalf("device pool mutated: %#v", deviceIDs)
}
}

View file

@ -0,0 +1,517 @@
package helps
import (
"bytes"
"encoding/json"
"net/http"
"regexp"
"sort"
"strings"
"github.com/google/uuid"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/tidwall/gjson"
)
const (
// claudeAnthropicVersion is the only Anthropic-Version Claude Code sends.
claudeAnthropicVersion = "2023-06-01"
// claudeDefaultStainlessTimeout is the X-Stainless-Timeout every measured
// native helper sends. It is deliberately NOT read from
// claude-header-defaults.timeout: applyClaudeHeaders routes a confirmed client
// through misc.EnsureHeader, which prefers the incoming header and only falls
// back to the configured value when the caller sent none. A confirmed helper
// therefore always forwards its own 600, so comparing against the operator
// value would make any non-600 configuration reject every genuine helper.
claudeDefaultStainlessTimeout = "600"
)
var (
claudeCodeUserAgentPattern = regexp.MustCompile(`(?i)^claude-cli/`)
claudeCodeUserAgentDetailsPattern = regexp.MustCompile(`(?i)^claude-cli/\S+\s+\(external,\s*([^,)]+)(?:,\s*agent-sdk/([^,)]+))?`)
claudeCodeNativeUserAgentPattern = regexp.MustCompile(`(?i)^claude-cli/[0-9]+\.[0-9]+\.[0-9]+\s+\(external,\s*[^,)]+(?:,\s*agent-sdk/[0-9]+\.[0-9]+\.[0-9]+)?\)$`)
)
var claudeCodeSubclientByEntrypoint = map[string]string{
"cli": "claude-code-cli",
"mcp": "claude-code-mcp",
"bench": "claude-code-bench",
"sdk-cli": "claude-code-cli-sdk",
"sdk-ts": "claude-code-sdk-ts",
"sdk-py": "claude-code-sdk-py",
"claude-vscode": "claude-code-vscode",
"claude-code-github-action": "claude-code-gh-action",
"local-agent": "claude-local-agent",
"local_agent": "claude-local-agent",
"claude-desktop": "claude-desktop",
"claude-desktop-3p": "claude-desktop-3p",
"remote": "claude-remote",
"remote_baku": "claude-remote-baku",
"remote_cowork": "claude-remote-cowork",
"remote_trigger": "claude-remote-trigger",
"remote_desktop": "claude-remote-desktop",
"remote_mobile": "claude-remote-mobile",
"claude_in_slack": "claude-in-slack",
"claude-in-slack": "claude-in-slack",
"claude-in-teams": "claude-in-teams",
"claude-security": "claude-security",
"ssh-remote": "claude-ssh-remote",
"claude-coworker": "claude-coworker",
"claude-coworker-terminal": "claude-coworker-terminal",
}
// Only product surfaces with verified 2.1.220 wire behavior are eligible for
// pass-through. Other first-party-looking entrypoints are cloaked until their
// CPA-reachable request shape has been captured and reviewed.
var nativeClaudeEntrypoints = map[string]bool{
"cli": true,
"sdk-cli": true,
"claude-vscode": true,
}
type claudeCodeHelperShape uint8
const (
claudeCodeHelperShapeNone claudeCodeHelperShape = iota
claudeCodeHelperShapeMinimal
claudeCodeHelperShapeStructured
claudeCodeHelperModel = "claude-haiku-4-5-20251001"
)
// These are the six exact beta sequences observed across 14 markerless native
// Claude Code 2.1.220 Haiku helper requests. Keeping the allowlist exact avoids
// turning the helper exception into a generic no-claude-code-beta bypass.
var measuredClaudeCodeHelperBetaProfiles = map[string]claudeCodeHelperShape{
claudeCodeHelperBetaProfile(true): claudeCodeHelperShapeMinimal,
claudeCodeHelperBetaProfile(false): claudeCodeHelperShapeMinimal,
claudeCodeHelperBetaProfile(true,
"advisor-tool-2026-03-01",
"structured-outputs-2025-12-15",
"cache-diagnosis-2026-04-07",
): claudeCodeHelperShapeStructured,
claudeCodeHelperBetaProfile(true,
"structured-outputs-2025-12-15",
"fallback-credit-2026-06-01",
): claudeCodeHelperShapeStructured,
claudeCodeHelperBetaProfile(true,
"structured-outputs-2025-12-15",
): claudeCodeHelperShapeStructured,
claudeCodeHelperBetaProfile(false,
"structured-outputs-2025-12-15",
): claudeCodeHelperShapeStructured,
}
// ClaudeCodeRequestDetection records the strong signals and first-party
// subclient identity used to distinguish an official Claude Code request from
// a client that only copied its User-Agent.
type ClaudeCodeRequestDetection struct {
Confirmed bool
StrongSignals bool
NativeClient bool
XAppCLI bool
UserAgent bool
BetasPresent bool
MetadataUserID bool
HelperProfile bool
Entrypoint string
Subclient string
AgentSDKVersion string
}
// DetectClaudeCodeRequest first mirrors CCH's strong-signal contract, then
// applies CPA's native-client policy. Standard Messages requests require all
// four strong signals; count_tokens omits metadata.user_id. A separate narrow
// profile recognizes measured native Haiku helper requests that intentionally
// omit claude-code-20250219. Generic sdk-ts/sdk-py Agent SDK entrypoints remain
// unconfirmed and receive CLI cloaking.
func DetectClaudeCodeRequest(headers http.Header, payload []byte, countTokens bool, configs ...*config.Config) ClaudeCodeRequestDetection {
var cfg *config.Config
if len(configs) > 0 {
cfg = configs[0]
}
userAgent := headerValue(headers, "User-Agent")
entrypoint, agentSDKVersion := parseClaudeCodeUserAgentDetails(userAgent)
detection := ClaudeCodeRequestDetection{
XAppCLI: headerValue(headers, "X-App") == "cli",
UserAgent: plausibleClaudeCodeUserAgent(userAgent, cfg),
BetasPresent: headerContainsClaudeCodeBeta(headers),
Entrypoint: entrypoint,
Subclient: claudeCodeSubclientByEntrypoint[entrypoint],
AgentSDKVersion: agentSDKVersion,
}
metadataUserID := gjson.GetBytes(payload, "metadata.user_id")
detection.MetadataUserID = metadataUserID.Exists() && metadataUserID.Type == gjson.String && isValidUserID(metadataUserID.String())
detection.NativeClient = nativeClaudeEntrypoints[entrypoint]
standardSignals := detection.XAppCLI && detection.UserAgent && detection.BetasPresent && (countTokens || detection.MetadataUserID)
detection.HelperProfile = detection.NativeClient && matchesMeasuredClaudeCodeHelperProfile(headers, payload, countTokens, detection, cfg)
detection.StrongSignals = standardSignals || detection.HelperProfile
detection.Confirmed = detection.StrongSignals && detection.NativeClient
return detection
}
func claudeCodeHelperBetaProfile(redactThinking bool, trailing ...string) string {
betas := []string{"oauth-2025-04-20", "interleaved-thinking-2025-05-14"}
if redactThinking {
betas = append(betas, "redact-thinking-2026-02-12")
}
betas = append(betas,
"thinking-token-count-2026-05-13",
"context-management-2025-06-27",
"prompt-caching-scope-2026-01-05",
)
betas = append(betas, trailing...)
return strings.Join(betas, ",")
}
func matchesMeasuredClaudeCodeHelperProfile(
headers http.Header,
payload []byte,
countTokens bool,
detection ClaudeCodeRequestDetection,
cfg *config.Config,
) bool {
if countTokens ||
detection.Entrypoint != "cli" ||
detection.BetasPresent ||
!detection.XAppCLI ||
!detection.UserAgent ||
!detection.MetadataUserID {
return false
}
shape := measuredClaudeCodeHelperBetaProfiles[normalizedClaudeBetaHeader(headers)]
if shape == claudeCodeHelperShapeNone || measuredClaudeCodeHelperBodyShape(payload) != shape {
return false
}
if !measuredClaudeCodeHelperHeadersMatch(headers, cfg, shape) {
return false
}
return measuredClaudeCodeHelperSessionMatches(headers, payload)
}
// normalizedClaudeBetaHeader joins every Anthropic-Beta value in wire order.
// Values() is tried first so canonical headers keep a deterministic order; the
// case-insensitive fallback only exists for hand-built header maps that store a
// non-canonical key, where ranging the map alone would be order-dependent.
func normalizedClaudeBetaHeader(headers http.Header) string {
if headers == nil {
return ""
}
values := headers.Values("Anthropic-Beta")
if len(values) == 0 {
keys := make([]string, 0, 2)
for key := range headers {
if strings.EqualFold(key, "Anthropic-Beta") {
keys = append(keys, key)
}
}
sort.Strings(keys)
for _, key := range keys {
values = append(values, headers[key]...)
}
}
betas := make([]string, 0, 12)
for _, value := range values {
for _, beta := range strings.Split(value, ",") {
if beta = strings.TrimSpace(beta); beta != "" {
betas = append(betas, beta)
}
}
}
return strings.Join(betas, ",")
}
// measuredClaudeCodeHelperHeadersMatch validates the helper transport envelope.
//
// Platform and software-version headers are deliberately NOT compared for
// equality. The device-profile pipeline this detector feeds already pins OS/Arch
// to the configured baseline and replaces a non-baseline software tuple instead
// of rejecting it, so demanding equality here would classify a genuine Claude
// Code helper from Windows/Linux, or from a different Node or SDK build, as a
// foreign client and cloak it. Values that carry real discriminating power - the
// exact beta allowlist, the body shape, the billing CCH and the session binding -
// stay strict.
func measuredClaudeCodeHelperHeadersMatch(headers http.Header, cfg *config.Config, shape claudeCodeHelperShape) bool {
profile := defaultClaudeDeviceProfile(cfg)
expected := map[string]string{
"Accept": "application/json",
"Content-Type": "application/json",
"X-Stainless-Lang": "js",
"X-Stainless-Runtime": "node",
"X-Stainless-Retry-Count": "0",
"X-Stainless-Timeout": claudeDefaultStainlessTimeout,
"Anthropic-Version": claudeAnthropicVersion,
"Anthropic-Dangerous-Direct-Browser-Access": "true",
}
for name, want := range expected {
if headerValue(headers, name) != want {
return false
}
}
// Presence is still required: the native SDK always sends these.
for _, name := range []string{
"X-Stainless-Package-Version",
"X-Stainless-Runtime-Version",
"X-Stainless-OS",
"X-Stainless-Arch",
} {
if headerValue(headers, name) == "" {
return false
}
}
candidate := ClaudeDeviceProfile{
UserAgent: headerValue(headers, "User-Agent"),
PackageVersion: headerValue(headers, "X-Stainless-Package-Version"),
RuntimeVersion: headerValue(headers, "X-Stainless-Runtime-Version"),
}
if version, ok := parseClaudeCLIVersion(candidate.UserAgent); ok {
candidate.version = version
candidate.hasVersion = true
}
if !meetsClaudeDeviceProfileBaseline(candidate, profile) {
return false
}
if async := headerValue(headers, "X-Stainless-Async"); (shape == claudeCodeHelperShapeStructured && async != "async") ||
(shape == claudeCodeHelperShapeMinimal && async != "") {
return false
}
compression := headerValue(headers, "Accept-Encoding")
if (shape == claudeCodeHelperShapeStructured && compression != "gzip, deflate, br, zstd") ||
(shape == claudeCodeHelperShapeMinimal && compression != "gzip") {
return false
}
requestID := headerValue(headers, "X-Client-Request-Id")
_, errRequestID := uuid.Parse(requestID)
return errRequestID == nil
}
func measuredClaudeCodeHelperSessionMatches(headers http.Header, payload []byte) bool {
metadata := gjson.GetBytes(payload, "metadata")
if !metadata.IsObject() || !claudeJSONObjectHasKeys([]byte(metadata.Raw), []string{"user_id"}) {
return false
}
userID := metadata.Get("user_id")
if userID.Type != gjson.String || !isValidUserID(userID.String()) {
return false
}
// The native metadata builder is
// {...extraMetadata, device_id, account_uuid, session_id, ...parentSessionId && {parent_session_id}}
// in 2.1.220, 2.1.221 and 2.1.227 alike, so parent_session_id is a legitimate
// optional trailing key for sub-agent and forked sessions. Rejecting it would
// cloak the helper requests those sessions issue.
identityRaw := []byte(userID.String())
if !claudeJSONObjectHasKeys(identityRaw, []string{"device_id", "account_uuid", "session_id"}) &&
!claudeJSONObjectHasKeys(identityRaw, []string{"device_id", "account_uuid", "session_id", "parent_session_id"}) {
return false
}
return headerValue(headers, ClaudeCodeSessionHeader) == gjson.GetBytes(identityRaw, "session_id").String()
}
func measuredClaudeCodeHelperBodyShape(payload []byte) claudeCodeHelperShape {
minimalKeys := []string{"model", "max_tokens", "messages", "metadata"}
structuredKeys := []string{"model", "messages", "system", "tools", "metadata", "max_tokens", "thinking", "temperature", "output_config", "stream"}
shape := claudeCodeHelperShapeNone
switch {
case claudeJSONObjectHasKeys(payload, minimalKeys):
shape = claudeCodeHelperShapeMinimal
case claudeJSONObjectHasKeys(payload, structuredKeys):
shape = claudeCodeHelperShapeStructured
default:
return claudeCodeHelperShapeNone
}
maxTokens := gjson.GetBytes(payload, "max_tokens")
if gjson.GetBytes(payload, "model").String() != claudeCodeHelperModel ||
maxTokens.Type != gjson.Number {
return claudeCodeHelperShapeNone
}
messages := gjson.GetBytes(payload, "messages")
if !messages.IsArray() || len(messages.Array()) != 1 {
return claudeCodeHelperShapeNone
}
message := messages.Get("0")
if !claudeJSONObjectHasKeys([]byte(message.Raw), []string{"role", "content"}) ||
message.Get("role").String() != "user" {
return claudeCodeHelperShapeNone
}
if shape == claudeCodeHelperShapeMinimal {
if maxTokens.Raw != "1" || message.Get("content").Type != gjson.String {
return claudeCodeHelperShapeNone
}
return shape
}
content := message.Get("content")
if !content.IsArray() || len(content.Array()) != 1 {
return claudeCodeHelperShapeNone
}
contentBlock := content.Get("0")
if !claudeJSONObjectHasKeys([]byte(contentBlock.Raw), []string{"type", "text"}) ||
contentBlock.Get("type").String() != "text" {
return claudeCodeHelperShapeNone
}
if !measuredClaudeCodeHelperSystemMatches(gjson.GetBytes(payload, "system")) {
return claudeCodeHelperShapeNone
}
if tools := gjson.GetBytes(payload, "tools"); !tools.IsArray() || len(tools.Array()) != 0 {
return claudeCodeHelperShapeNone
}
thinking := gjson.GetBytes(payload, "thinking")
outputConfig := gjson.GetBytes(payload, "output_config")
if !claudeJSONObjectHasKeys([]byte(thinking.Raw), []string{"type"}) ||
thinking.Get("type").String() != "disabled" {
return claudeCodeHelperShapeNone
}
format := outputConfig.Get("format")
schema := format.Get("schema")
properties := schema.Get("properties")
titleProperty := properties.Get("title")
required := schema.Get("required")
additionalProperties := schema.Get("additionalProperties")
if !claudeJSONObjectHasKeys([]byte(outputConfig.Raw), []string{"format"}) ||
!claudeJSONObjectHasKeys([]byte(format.Raw), []string{"type", "schema"}) ||
format.Get("type").String() != "json_schema" ||
!claudeJSONObjectHasKeys([]byte(schema.Raw), []string{"type", "properties", "required", "additionalProperties"}) ||
schema.Get("type").String() != "object" ||
!claudeJSONObjectHasKeys([]byte(properties.Raw), []string{"title"}) ||
!claudeJSONObjectHasKeys([]byte(titleProperty.Raw), []string{"type"}) ||
titleProperty.Get("type").String() != "string" ||
!required.IsArray() || len(required.Array()) != 1 || required.Get("0").String() != "title" ||
additionalProperties.Type != gjson.False {
return claudeCodeHelperShapeNone
}
temperature := gjson.GetBytes(payload, "temperature")
if maxTokens.Raw != "32000" ||
temperature.Raw != "1" ||
gjson.GetBytes(payload, "stream").Type != gjson.True {
return claudeCodeHelperShapeNone
}
return shape
}
func measuredClaudeCodeHelperSystemMatches(system gjson.Result) bool {
if !system.IsArray() || len(system.Array()) != 3 {
return false
}
for _, block := range system.Array() {
if !claudeJSONObjectHasKeys([]byte(block.Raw), []string{"type", "text"}) || block.Get("type").String() != "text" {
return false
}
}
billing := system.Get("0.text").String()
identity := system.Get("1.text").String()
return strings.HasPrefix(billing, "x-anthropic-billing-header:") && measuredClaudeBillingCCH(billing) && strings.HasPrefix(identity, "You are Claude Code")
}
// measuredClaudeBillingCCH validates the five lowercase hexadecimal characters the
// native billing header carries. It duplicates isLowerHex in
// internal/runtime/executor/claude_signing.go because the signing side lives in the
// package that imports this one; keep the two definitions in step.
func measuredClaudeBillingCCH(billing string) bool {
marker := strings.Index(billing, " cch=")
if marker < 0 {
return false
}
valueStart := marker + len(" cch=")
valueEnd := valueStart + 5
if valueEnd >= len(billing) || billing[valueEnd] != ';' {
return false
}
for _, character := range billing[valueStart:valueEnd] {
decimal := character >= '0' && character <= '9'
lowerHex := character >= 'a' && character <= 'f'
if !decimal && !lowerHex {
return false
}
}
return true
}
func claudeJSONObjectHasKeys(raw []byte, want []string) bool {
if !json.Valid(raw) {
return false
}
decoder := json.NewDecoder(bytes.NewReader(raw))
opening, errOpening := decoder.Token()
if errOpening != nil || opening != json.Delim('{') {
return false
}
keyIndex := 0
for decoder.More() {
token, errToken := decoder.Token()
if errToken != nil {
return false
}
key, okKey := token.(string)
if !okKey || keyIndex >= len(want) || key != want[keyIndex] {
return false
}
keyIndex++
var value json.RawMessage
if errValue := decoder.Decode(&value); errValue != nil {
return false
}
}
closing, errClosing := decoder.Token()
return errClosing == nil && closing == json.Delim('}') && keyIndex == len(want)
}
func plausibleClaudeCodeUserAgent(userAgent string, cfg *config.Config) bool {
userAgent = strings.TrimSpace(userAgent)
if !claudeCodeUserAgentPattern.MatchString(userAgent) || !claudeCodeNativeUserAgentPattern.MatchString(userAgent) {
return false
}
candidate, okCandidate := parseClaudeCLIVersion(userAgent)
baseline, okBaseline := parseClaudeCLIVersion(defaultClaudeDeviceProfile(cfg).UserAgent)
return okCandidate && okBaseline && plausibleClaudeCLIVersion(candidate, baseline)
}
func parseClaudeCodeUserAgentDetails(userAgent string) (entrypoint, agentSDKVersion string) {
matches := claudeCodeUserAgentDetailsPattern.FindStringSubmatch(strings.TrimSpace(userAgent))
if len(matches) < 2 {
return "", ""
}
entrypoint = strings.ToLower(strings.TrimSpace(matches[1]))
if len(matches) >= 3 {
agentSDKVersion = strings.TrimSpace(matches[2])
}
return entrypoint, agentSDKVersion
}
func headerValue(headers http.Header, name string) string {
if headers == nil {
return ""
}
if value := headers.Get(name); value != "" {
return value
}
for key, values := range headers {
if !strings.EqualFold(key, name) || len(values) == 0 {
continue
}
return values[0]
}
return ""
}
func headerContainsClaudeCodeBeta(headers http.Header) bool {
if headers == nil {
return false
}
for key, values := range headers {
if !strings.EqualFold(key, "Anthropic-Beta") {
continue
}
for _, value := range values {
for _, beta := range strings.Split(value, ",") {
if strings.TrimSpace(beta) == "claude-code-20250219" {
return true
}
}
}
}
return false
}

View file

@ -0,0 +1,530 @@
package helps
import (
"encoding/json"
"net/http"
"strings"
"testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
)
const validClaudeCodeMetadataUserID = `{"device_id":"0000000000000000000000000000000000000000000000000000000000000000","account_uuid":"aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa","session_id":"11111111-2222-4333-8444-555555555555"}`
func claudeCodeDetectionPayload(userID string) []byte {
encodedUserID, _ := json.Marshal(userID)
return []byte(`{"metadata":{"user_id":` + string(encodedUserID) + `}}`)
}
func confirmedClaudeCodeHeaders() http.Header {
return http.Header{
"User-Agent": {"claude-cli/2.1.220 (external, cli)"},
"X-App": {"cli"},
"Anthropic-Beta": {"claude-code-20250219,interleaved-thinking-2025-05-14"},
}
}
func measuredClaudeCodeHelperHeaders(betaProfile string, structured bool) http.Header {
profile := defaultClaudeDeviceProfile(&config.Config{})
headers := http.Header{
"Accept": {"application/json"},
"Accept-Encoding": {"gzip"},
"Content-Type": {"application/json"},
"User-Agent": {profile.UserAgent},
"X-App": {"cli"},
"Anthropic-Beta": {betaProfile},
"Anthropic-Version": {"2023-06-01"},
"Anthropic-Dangerous-Direct-Browser-Access": {"true"},
"X-Claude-Code-Session-Id": {"11111111-2222-4333-8444-555555555555"},
"X-Client-Request-Id": {"66666666-7777-4888-8999-aaaaaaaaaaaa"},
"X-Stainless-Lang": {"js"},
"X-Stainless-Runtime": {"node"},
"X-Stainless-Package-Version": {profile.PackageVersion},
"X-Stainless-Runtime-Version": {profile.RuntimeVersion},
"X-Stainless-OS": {profile.OS},
"X-Stainless-Arch": {profile.Arch},
"X-Stainless-Retry-Count": {"0"},
"X-Stainless-Timeout": {"600"},
}
if structured {
headers.Set("Accept-Encoding", "gzip, deflate, br, zstd")
headers.Set("X-Stainless-Async", "async")
}
canonical := make(http.Header, len(headers))
for name, values := range headers {
for _, value := range values {
canonical.Add(name, value)
}
}
return canonical
}
func measuredClaudeCodeMinimalHelperPayload() []byte {
encodedUserID, _ := json.Marshal(validClaudeCodeMetadataUserID)
return []byte(`{"model":"claude-haiku-4-5-20251001","max_tokens":1,"messages":[{"role":"user","content":"helper probe"}],"metadata":{"user_id":` + string(encodedUserID) + `}}`)
}
func measuredClaudeCodeStructuredHelperPayload() []byte {
encodedUserID, _ := json.Marshal(validClaudeCodeMetadataUserID)
return []byte(`{"model":"claude-haiku-4-5-20251001","messages":[{"role":"user","content":[{"type":"text","text":"helper probe"}]}],"system":[{"type":"text","text":"x-anthropic-billing-header: cc_version=2.1.220; cc_entrypoint=cli; cch=00000;"},{"type":"text","text":"You are Claude Code, Anthropic's official CLI for Claude."},{"type":"text","text":"Return a short title."}],"tools":[],"metadata":{"user_id":` + string(encodedUserID) + `},"max_tokens":32000,"thinking":{"type":"disabled"},"temperature":1,"output_config":{"format":{"type":"json_schema","schema":{"type":"object","properties":{"title":{"type":"string"}},"required":["title"],"additionalProperties":false}}},"stream":true}`)
}
func TestDetectClaudeCodeRequestRequiresAllFourMessageSignals(t *testing.T) {
payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID)
detection := DetectClaudeCodeRequest(confirmedClaudeCodeHeaders(), payload, false)
if !detection.Confirmed || !detection.StrongSignals || !detection.NativeClient {
t.Fatalf("detection = %#v, want native CLI confirmed", detection)
}
if !detection.XAppCLI || !detection.UserAgent || !detection.BetasPresent || !detection.MetadataUserID {
t.Fatalf("detection signals = %#v, want all present", detection)
}
}
func TestDetectClaudeCodeRequestAcceptsConfiguredMeasuredBaseline(t *testing.T) {
headers := confirmedClaudeCodeHeaders()
headers.Set("User-Agent", "claude-cli/2.2.0 (external, cli)")
payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID)
if detection := DetectClaudeCodeRequest(headers, payload, false); detection.Confirmed {
t.Fatalf("default detection = %#v, want unconfigured 2.2.0 rejected", detection)
}
cfg := &config.Config{ClaudeHeaderDefaults: config.ClaudeHeaderDefaults{
UserAgent: "claude-cli/2.2.0 (external, cli)",
PackageVersion: "0.95.0",
RuntimeVersion: "v26.4.0",
}}
if detection := DetectClaudeCodeRequest(headers, payload, false, cfg); !detection.Confirmed {
t.Fatalf("configured detection = %#v, want measured baseline confirmed", detection)
}
}
func TestDetectClaudeCodeRequestRejectsEachMissingMessageSignal(t *testing.T) {
payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID)
for _, test := range []struct {
name string
headers http.Header
body []byte
}{
{name: "x-app", headers: http.Header{"User-Agent": {"claude-cli/2.1.220 (external, cli)"}, "Anthropic-Beta": {"claude-code-20250219"}}, body: payload},
{name: "user-agent", headers: http.Header{"User-Agent": {"curl/8.7.1"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, body: payload},
{name: "betas", headers: http.Header{"User-Agent": {"claude-cli/2.1.220 (external, cli)"}, "X-App": {"cli"}}, body: payload},
{name: "metadata", headers: confirmedClaudeCodeHeaders(), body: []byte(`{"messages":[]}`)},
} {
t.Run(test.name, func(t *testing.T) {
if detection := DetectClaudeCodeRequest(test.headers, test.body, false); detection.Confirmed {
t.Fatalf("detection = %#v, want unconfirmed", detection)
}
})
}
}
func TestDetectClaudeCodeRequestClassifiesEntrypoints(t *testing.T) {
payload := claudeCodeDetectionPayload(validClaudeCodeMetadataUserID)
for _, test := range []struct {
name string
userAgent string
entrypoint string
subclient string
agentSDKVersion string
native bool
}{
{name: "cli", userAgent: "claude-cli/2.1.220 (external, cli)", entrypoint: "cli", subclient: "claude-code-cli", native: true},
{name: "vscode-agent-sdk", userAgent: "claude-cli/2.1.220 (external, claude-vscode, agent-sdk/0.3.220)", entrypoint: "claude-vscode", subclient: "claude-code-vscode", agentSDKVersion: "0.3.220", native: true},
{name: "sdk-cli", userAgent: "claude-cli/2.1.220 (external, sdk-cli)", entrypoint: "sdk-cli", subclient: "claude-code-cli-sdk", native: true},
{name: "sdk-ts", userAgent: "claude-cli/2.1.220 (external, sdk-ts, agent-sdk/0.3.220)", entrypoint: "sdk-ts", subclient: "claude-code-sdk-ts", agentSDKVersion: "0.3.220"},
{name: "sdk-py", userAgent: "claude-cli/2.1.220 (external, sdk-py, agent-sdk/0.1.0)", entrypoint: "sdk-py", subclient: "claude-code-sdk-py", agentSDKVersion: "0.1.0"},
{name: "desktop", userAgent: "claude-cli/2.1.220 (external, claude-desktop)", entrypoint: "claude-desktop", subclient: "claude-desktop"},
{name: "desktop-third-party-inference", userAgent: "claude-cli/2.1.220 (external, claude-desktop-3p)", entrypoint: "claude-desktop-3p", subclient: "claude-desktop-3p"},
{name: "remote", userAgent: "claude-cli/2.1.220 (external, remote)", entrypoint: "remote", subclient: "claude-remote"},
{name: "github-action", userAgent: "claude-cli/2.1.220 (external, claude-code-github-action)", entrypoint: "claude-code-github-action", subclient: "claude-code-gh-action"},
{name: "unknown", userAgent: "claude-cli/2.1.220 (external, copied-client)", entrypoint: "copied-client"},
} {
t.Run(test.name, func(t *testing.T) {
headers := confirmedClaudeCodeHeaders()
headers.Set("User-Agent", test.userAgent)
detection := DetectClaudeCodeRequest(headers, payload, false)
if !detection.StrongSignals {
t.Fatalf("detection = %#v, want all CCH strong signals", detection)
}
if detection.Confirmed != test.native || detection.NativeClient != test.native {
t.Fatalf("detection = %#v, want native/confirmed %t", detection, test.native)
}
if detection.Entrypoint != test.entrypoint || detection.Subclient != test.subclient || detection.AgentSDKVersion != test.agentSDKVersion {
t.Fatalf("detection identity = %#v, want entrypoint %q subclient %q agent SDK %q", detection, test.entrypoint, test.subclient, test.agentSDKVersion)
}
})
}
}
func TestDetectClaudeCodeCountTokensAllowsMissingMetadata(t *testing.T) {
headers := confirmedClaudeCodeHeaders()
headers.Set("User-Agent", "claude-cli/2.1.220 (external, claude-vscode, agent-sdk/0.3.220)")
detection := DetectClaudeCodeRequest(headers, []byte(`{"messages":[]}`), true)
if !detection.Confirmed {
t.Fatalf("detection = %#v, want confirmed", detection)
}
if detection.MetadataUserID {
t.Fatalf("metadata signal = true, want false: %#v", detection)
}
if detection.Subclient != "claude-code-vscode" || detection.AgentSDKVersion != "0.3.220" {
t.Fatalf("count_tokens identity = %#v, want VSCode Agent SDK", detection)
}
}
func TestDetectClaudeCodeRequestRecognizesMeasuredHaikuHelpers(t *testing.T) {
tests := []struct {
name string
beta string
structured bool
payload []byte
}{
{
name: "minimal with redact thinking",
beta: claudeCodeHelperBetaProfile(true),
payload: measuredClaudeCodeMinimalHelperPayload(),
},
{
name: "minimal without redact thinking",
beta: claudeCodeHelperBetaProfile(false),
payload: measuredClaudeCodeMinimalHelperPayload(),
},
{
name: "structured title helper with advisor",
beta: claudeCodeHelperBetaProfile(true, "advisor-tool-2026-03-01", "structured-outputs-2025-12-15", "cache-diagnosis-2026-04-07"),
structured: true,
payload: measuredClaudeCodeStructuredHelperPayload(),
},
{
name: "structured title helper with fallback credit",
beta: claudeCodeHelperBetaProfile(true, "structured-outputs-2025-12-15", "fallback-credit-2026-06-01"),
structured: true,
payload: measuredClaudeCodeStructuredHelperPayload(),
},
{
name: "structured title helper with lowercase hex CCH",
beta: claudeCodeHelperBetaProfile(true, "structured-outputs-2025-12-15"),
structured: true,
payload: []byte(strings.Replace(string(measuredClaudeCodeStructuredHelperPayload()), "cch=00000", "cch=7ee87", 1)),
},
{
name: "structured title helper without redact thinking",
beta: claudeCodeHelperBetaProfile(false, "structured-outputs-2025-12-15"),
structured: true,
payload: measuredClaudeCodeStructuredHelperPayload(),
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
detection := DetectClaudeCodeRequest(
measuredClaudeCodeHelperHeaders(test.beta, test.structured),
test.payload,
false,
)
if !detection.Confirmed || !detection.StrongSignals || !detection.NativeClient || !detection.HelperProfile {
t.Fatalf("detection = %#v, want confirmed measured helper", detection)
}
if detection.BetasPresent {
t.Fatalf("claude-code beta signal = true, want helper profile to remain separate: %#v", detection)
}
})
}
}
func TestDetectClaudeCodeRequestRejectsMalformedStructuredHaikuHelpers(t *testing.T) {
basePayload := string(measuredClaudeCodeStructuredHelperPayload())
beta := claudeCodeHelperBetaProfile(true, "structured-outputs-2025-12-15")
for _, test := range []struct {
name string
payload string
}{
{name: "non-hex CCH", payload: strings.Replace(basePayload, "cch=00000", "cch=ghijk", 1)},
{name: "uppercase CCH", payload: strings.Replace(basePayload, "cch=00000", "cch=7EE87", 1)},
{name: "wrong token cap", payload: strings.Replace(basePayload, `"max_tokens":32000`, `"max_tokens":32001`, 1)},
{name: "open schema", payload: strings.Replace(basePayload, `"additionalProperties":false`, `"additionalProperties":true`, 1)},
} {
t.Run(test.name, func(t *testing.T) {
detection := DetectClaudeCodeRequest(measuredClaudeCodeHelperHeaders(beta, true), []byte(test.payload), false)
if detection.Confirmed || detection.HelperProfile {
t.Fatalf("detection = %#v, want malformed structured helper rejected", detection)
}
})
}
}
func TestDetectClaudeCodeRequestRejectsNearMissHaikuHelpers(t *testing.T) {
minimalPayload := string(measuredClaudeCodeMinimalHelperPayload())
tests := []struct {
name string
mutate func(http.Header)
payload string
countTokens bool
}{
{
name: "unexpected beta profile",
mutate: func(headers http.Header) {
headers.Set("Anthropic-Beta", headers.Get("Anthropic-Beta")+",unknown-beta")
},
payload: minimalPayload,
},
{
name: "missing stainless package",
mutate: func(headers http.Header) {
headers.Del("X-Stainless-Package-Version")
},
payload: minimalPayload,
},
{
name: "wrong compression profile",
mutate: func(headers http.Header) {
headers.Set("Accept-Encoding", "gzip, deflate, br, zstd")
},
payload: minimalPayload,
},
{
name: "mismatched session header",
mutate: func(headers http.Header) {
headers.Set("X-Claude-Code-Session-Id", "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee")
},
payload: minimalPayload,
},
{
name: "invalid request id",
mutate: func(headers http.Header) {
headers.Set("X-Client-Request-Id", "not-a-uuid")
},
payload: minimalPayload,
},
{
name: "unexpected async mode",
mutate: func(headers http.Header) {
headers.Set("X-Stainless-Async", "async")
},
payload: minimalPayload,
},
{
name: "wrong helper model",
payload: strings.Replace(minimalPayload, claudeCodeHelperModel, "claude-sonnet-4-6", 1),
},
{
name: "wrong helper token cap",
payload: strings.Replace(minimalPayload, `"max_tokens":1`, `"max_tokens":2`, 1),
},
{
name: "extra root key",
payload: strings.TrimSuffix(minimalPayload, "}") + `,"tools":[]}`,
},
{
name: "cache marker content shape",
payload: strings.Replace(minimalPayload, `"content":"helper probe"`, `"content":[{"type":"text","text":"helper probe","cache_control":{"type":"ephemeral","ttl":"1h"}}]`, 1),
},
{
name: "count tokens endpoint",
payload: minimalPayload,
countTokens: true,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
headers := measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false)
if test.mutate != nil {
test.mutate(headers)
}
detection := DetectClaudeCodeRequest(headers, []byte(test.payload), test.countTokens)
if detection.Confirmed || detection.HelperProfile {
t.Fatalf("detection = %#v, want helper near miss rejected", detection)
}
})
}
}
func TestDetectClaudeCodeRequestRejectsMalformedNativeSignals(t *testing.T) {
tests := []struct {
name string
headers http.Header
userID string
}{
{name: "legacy metadata", headers: confirmedClaudeCodeHeaders(), userID: "user_abc_account__session_session"},
{name: "short device", headers: confirmedClaudeCodeHeaders(), userID: `{"device_id":"abc","account_uuid":"","session_id":"11111111-2222-4333-8444-555555555555"}`},
{name: "uppercase device", headers: confirmedClaudeCodeHeaders(), userID: `{"device_id":"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA","account_uuid":"","session_id":"11111111-2222-4333-8444-555555555555"}`},
{name: "invalid session", headers: confirmedClaudeCodeHeaders(), userID: `{"device_id":"0000000000000000000000000000000000000000000000000000000000000000","account_uuid":"","session_id":"session"}`},
{name: "malformed user agent", headers: http.Header{"User-Agent": {"claude-cli/not-a-version (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, userID: validClaudeCodeMetadataUserID},
{name: "unmeasured next-minor user agent", headers: http.Header{"User-Agent": {"claude-cli/2.2.0 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, userID: validClaudeCodeMetadataUserID},
{name: "implausible future user agent", headers: http.Header{"User-Agent": {"claude-cli/999.0.0 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"claude-code-20250219"}}, userID: validClaudeCodeMetadataUserID},
{name: "unrelated beta", headers: http.Header{"User-Agent": {"claude-cli/2.1.220 (external, cli)"}, "X-App": {"cli"}, "Anthropic-Beta": {"anything"}}, userID: validClaudeCodeMetadataUserID},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if detection := DetectClaudeCodeRequest(test.headers, claudeCodeDetectionPayload(test.userID), false); detection.Confirmed {
t.Fatalf("detection = %#v, want malformed signal to use local profile", detection)
}
})
}
}
// Recovered from the native metadata builder in 2.1.220, 2.1.221 and 2.1.227:
//
// {...extraMetadata, device_id, account_uuid, session_id, ...parentSessionId && {parent_session_id}}
//
// parent_session_id is therefore a legitimate optional trailing key that sub-agent
// and forked sessions attach, and it must not disqualify a helper request.
func TestDetectClaudeCodeRequestAcceptsHelperSubagentParentSessionID(t *testing.T) {
identity := `{"device_id":"0000000000000000000000000000000000000000000000000000000000000000","account_uuid":"aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa","session_id":"11111111-2222-4333-8444-555555555555","parent_session_id":"99999999-8888-4777-8666-555555555555"}`
encoded, _ := json.Marshal(identity)
payload := []byte(`{"model":"claude-haiku-4-5-20251001","max_tokens":1,"messages":[{"role":"user","content":"helper probe"}],"metadata":{"user_id":` + string(encoded) + `}}`)
detection := DetectClaudeCodeRequest(
measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false),
payload,
false,
)
if !detection.Confirmed || !detection.HelperProfile {
t.Fatalf("detection = %#v, want a confirmed sub-agent helper", detection)
}
}
func TestDetectClaudeCodeRequestRejectsHelperIdentityWithUnknownKeys(t *testing.T) {
identity := `{"device_id":"0000000000000000000000000000000000000000000000000000000000000000","account_uuid":"aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa","session_id":"11111111-2222-4333-8444-555555555555","spoofed":"x"}`
encoded, _ := json.Marshal(identity)
payload := []byte(`{"model":"claude-haiku-4-5-20251001","max_tokens":1,"messages":[{"role":"user","content":"helper probe"}],"metadata":{"user_id":` + string(encoded) + `}}`)
detection := DetectClaudeCodeRequest(
measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false),
payload,
false,
)
if detection.HelperProfile {
t.Fatalf("detection = %#v, want an unknown identity key to disqualify the helper profile", detection)
}
}
// The surrounding device-profile pipeline pins OS/Arch to the configured baseline
// rather than rejecting a foreign platform, so a genuine Windows or Linux helper
// must still be recognized instead of being cloaked.
func TestDetectClaudeCodeRequestAcceptsHelperFromNonBaselinePlatform(t *testing.T) {
for _, platform := range []struct{ os, arch string }{
{"Windows", "x64"},
{"Linux", "x64"},
{"MacOS", "x64"},
} {
t.Run(platform.os+"/"+platform.arch, func(t *testing.T) {
headers := measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false)
headers.Set("X-Stainless-OS", platform.os)
headers.Set("X-Stainless-Arch", platform.arch)
detection := DetectClaudeCodeRequest(headers, measuredClaudeCodeMinimalHelperPayload(), false)
if !detection.Confirmed || !detection.HelperProfile {
t.Fatalf("detection = %#v, want a confirmed helper on a non-baseline platform", detection)
}
})
}
}
func TestDetectClaudeCodeRequestRejectsHelperWithoutPlatformHeaders(t *testing.T) {
for _, name := range []string{
"X-Stainless-OS",
"X-Stainless-Arch",
"X-Stainless-Package-Version",
"X-Stainless-Runtime-Version",
} {
t.Run("missing "+name, func(t *testing.T) {
headers := measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false)
headers.Del(name)
detection := DetectClaudeCodeRequest(headers, measuredClaudeCodeMinimalHelperPayload(), false)
if detection.HelperProfile {
t.Fatalf("detection = %#v, want a missing %s to disqualify the helper profile", detection, name)
}
})
}
}
func TestDetectClaudeCodeRequestRejectsHelperWithForeignSoftwareTuple(t *testing.T) {
for name, value := range map[string]string{
"X-Stainless-Package-Version": "0.0.1",
"X-Stainless-Runtime-Version": "v0.0.1",
} {
t.Run(name, func(t *testing.T) {
headers := measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false)
headers.Set(name, value)
detection := DetectClaudeCodeRequest(headers, measuredClaudeCodeMinimalHelperPayload(), false)
if detection.HelperProfile {
t.Fatalf("detection = %#v, want a foreign %s to disqualify the helper profile", detection, name)
}
})
}
}
func TestNormalizedClaudeBetaHeaderIsDeterministic(t *testing.T) {
canonical := http.Header{}
canonical.Add("Anthropic-Beta", "oauth-2025-04-20")
canonical.Add("Anthropic-Beta", "interleaved-thinking-2025-05-14")
if got, want := normalizedClaudeBetaHeader(canonical), "oauth-2025-04-20,interleaved-thinking-2025-05-14"; got != want {
t.Fatalf("canonical join = %q, want %q", got, want)
}
// Two non-canonical spellings in one map used to be joined in Go map order.
nonCanonical := http.Header{
"anthropic-beta": {"oauth-2025-04-20"},
"ANTHROPIC-BETA": {"interleaved-thinking-2025-05-14"},
}
first := normalizedClaudeBetaHeader(nonCanonical)
for i := 0; i < 50; i++ {
if got := normalizedClaudeBetaHeader(nonCanonical); got != first {
t.Fatalf("non-canonical join is order-dependent: %q then %q", first, got)
}
}
if !strings.Contains(first, "oauth-2025-04-20") || !strings.Contains(first, "interleaved-thinking-2025-05-14") {
t.Fatalf("non-canonical join lost values: %q", first)
}
if got := normalizedClaudeBetaHeader(nil); got != "" {
t.Fatalf("nil header join = %q, want empty", got)
}
}
// A confirmed helper is routed through misc.EnsureHeader, so CPA forwards the
// helper's own X-Stainless-Timeout and never the operator default. Keying the
// detector on claude-header-defaults.timeout instead of the measured constant
// therefore rejected every genuine helper whenever that value was customized.
func TestMeasuredHelperProfileIgnoresConfiguredStainlessTimeout(t *testing.T) {
headers := measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false)
payload := measuredClaudeCodeMinimalHelperPayload()
if got := headers.Get("X-Stainless-Timeout"); got != claudeDefaultStainlessTimeout {
t.Fatalf("measured helper timeout = %q, want %q", got, claudeDefaultStainlessTimeout)
}
withTimeout := func(timeout string) *config.Config {
cfg := &config.Config{}
cfg.ClaudeHeaderDefaults.Timeout = timeout
return cfg
}
for _, test := range []struct {
name string
cfg *config.Config
}{
{name: "nil config"},
{name: "unset", cfg: &config.Config{}},
{name: "measured default", cfg: withTimeout(claudeDefaultStainlessTimeout)},
{name: "shorter operator default", cfg: withTimeout("300")},
{name: "longer operator default", cfg: withTimeout("900")},
} {
t.Run(test.name, func(t *testing.T) {
detection := DetectClaudeCodeRequest(headers, payload, false, test.cfg)
if !detection.HelperProfile || !detection.Confirmed {
t.Fatalf("detection = %#v, want confirmed helper regardless of configured timeout", detection)
}
})
}
// The measured constant stays the only accepted value, so a caller that does not
// send it is still disqualified even when the operator default happens to match.
t.Run("foreign timeout stays rejected", func(t *testing.T) {
foreign := measuredClaudeCodeHelperHeaders(claudeCodeHelperBetaProfile(true), false)
foreign.Set("X-Stainless-Timeout", "900")
if detection := DetectClaudeCodeRequest(foreign, payload, false, withTimeout("900")); detection.HelperProfile {
t.Fatalf("detection = %#v, want a non-measured timeout to disqualify the helper profile", detection)
}
})
}

View file

@ -0,0 +1,110 @@
package helps
import (
"context"
"net/http"
"regexp"
"strings"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/tidwall/gjson"
)
const (
ClaudeCodeSessionHeader = "X-Claude-Code-Session-Id"
ClaudeCodeAgentHeader = "X-Claude-Code-Agent-Id"
ClaudeCodeMainAgentID = "main"
)
var claudeCodeSessionSuffixPattern = regexp.MustCompile(`_session_([a-f0-9-]+)$`)
// ExtractClaudeCodeSessionID resolves a Claude Code session ID, preferring X-Claude-Code-Session-Id over payload metadata.
func ExtractClaudeCodeSessionID(ctx context.Context, payload []byte, headers http.Header) string {
if sessionID := claudeCodeHeader(ctx, headers, ClaudeCodeSessionHeader); sessionID != "" {
return sessionID
}
return extractClaudeCodeSessionIDFromPayload(payload)
}
// ExtractClaudeCodeAgentID resolves the Claude Code agent ID and uses a stable sentinel for the root agent.
func ExtractClaudeCodeAgentID(ctx context.Context, headers http.Header) string {
if agentID := claudeCodeHeader(ctx, headers, ClaudeCodeAgentHeader); agentID != "" {
return agentID
}
return ClaudeCodeMainAgentID
}
// ClaudeCodeExecutionScope returns the stable root-session and agent identity used by Codex execution state.
func ClaudeCodeExecutionScope(ctx context.Context, payload []byte, headers http.Header) (string, bool) {
sessionID := ExtractClaudeCodeSessionID(ctx, payload, headers)
if sessionID == "" {
return "", false
}
return "claude:" + sessionID + ":agent:" + ExtractClaudeCodeAgentID(ctx, headers), true
}
func claudeCodeHeader(ctx context.Context, headers http.Header, name string) string {
if value := headerValueCaseInsensitive(headers, name); value != "" {
return value
}
if ctx != nil {
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
return headerValueCaseInsensitive(ginCtx.Request.Header, name)
}
}
return ""
}
// HeaderValueCaseInsensitive returns the first non-empty header value matching name case-insensitively.
func HeaderValueCaseInsensitive(headers http.Header, name string) string {
return headerValueCaseInsensitive(headers, name)
}
func headerValueCaseInsensitive(headers http.Header, name string) string {
if headers == nil {
return ""
}
if value := strings.TrimSpace(headers.Get(name)); value != "" {
return value
}
for key, values := range headers {
if !strings.EqualFold(key, name) {
continue
}
for _, value := range values {
if value = strings.TrimSpace(value); value != "" {
return value
}
}
}
return ""
}
func extractClaudeCodeSessionIDFromPayload(payload []byte) string {
if len(payload) == 0 {
return ""
}
userID := gjson.GetBytes(payload, "metadata.user_id").String()
if userID == "" {
return ""
}
if matches := claudeCodeSessionSuffixPattern.FindStringSubmatch(userID); len(matches) >= 2 {
return matches[1]
}
if len(userID) > 0 && userID[0] == '{' {
return strings.TrimSpace(gjson.Get(userID, "session_id").String())
}
return ""
}
// ClaudeCodePromptCache derives a deterministic upstream prompt_cache_key for one Claude Code agent.
func ClaudeCodePromptCache(ctx context.Context, modelName string, payload []byte, headers http.Header) (CodexCache, bool, error) {
modelName = strings.TrimSpace(modelName)
executionScope, ok := ClaudeCodeExecutionScope(ctx, payload, headers)
if modelName == "" || !ok {
return CodexCache{}, false, nil
}
identity := strings.Join([]string{"cli-proxy-api:codex:claude-code", modelName, executionScope}, "\x00")
return CodexCache{ID: uuid.NewSHA1(uuid.NameSpaceOID, []byte(identity)).String()}, true, nil
}

View file

@ -0,0 +1,122 @@
package helps
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
)
func TestExtractClaudeCodeSessionIDFromPayloadJSON(t *testing.T) {
payload := []byte(`{"metadata":{"user_id":"{\"device_id\":\"d\",\"session_id\":\"cache-session-1\"}"}}`)
got := ExtractClaudeCodeSessionID(context.Background(), payload, nil)
if got != "cache-session-1" {
t.Fatalf("ExtractClaudeCodeSessionID() = %q, want cache-session-1", got)
}
}
func TestExtractClaudeCodeSessionIDFromHeader(t *testing.T) {
recorder := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(recorder)
ginCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
ginCtx.Request.Header.Set(ClaudeCodeSessionHeader, "header-session-1")
ctx := context.WithValue(context.Background(), "gin", ginCtx)
got := ExtractClaudeCodeSessionID(ctx, []byte(`{"model":"gpt-5.4"}`), nil)
if got != "header-session-1" {
t.Fatalf("ExtractClaudeCodeSessionID() = %q, want header-session-1", got)
}
}
func TestClaudeCodePromptCacheStableAcrossRequests(t *testing.T) {
ctx := context.Background()
payload := []byte(`{"metadata":{"user_id":"{\"session_id\":\"cache-session-2\"}"}}`)
first, ok, err := ClaudeCodePromptCache(ctx, "grok-composer-2.5-fast", payload, nil)
if err != nil {
t.Fatalf("ClaudeCodePromptCache first error: %v", err)
}
if !ok || first.ID == "" {
t.Fatalf("ClaudeCodePromptCache first = %#v, ok=%v, want cached id", first, ok)
}
second, ok, err := ClaudeCodePromptCache(ctx, "grok-composer-2.5-fast", payload, nil)
if err != nil {
t.Fatalf("ClaudeCodePromptCache second error: %v", err)
}
if !ok || second.ID != first.ID {
t.Fatalf("second cache id = %q, want %q", second.ID, first.ID)
}
}
func TestExtractClaudeCodeSessionIDPrefersHeaderOverPayload(t *testing.T) {
payload := []byte(`{"metadata":{"user_id":"{"session_id":"payload-session"}"}}`)
headers := http.Header{}
headers.Set(ClaudeCodeSessionHeader, "header-session")
got := ExtractClaudeCodeSessionID(context.Background(), payload, headers)
if got != "header-session" {
t.Fatalf("ExtractClaudeCodeSessionID() = %q, want header-session", got)
}
}
func TestClaudeCodeExecutionScopeAcceptsLowercaseHeaderMapKeys(t *testing.T) {
headers := http.Header{
"x-claude-code-session-id": []string{"lower-session"},
"x-claude-code-agent-id": []string{"lower-agent"},
}
scope, ok := ClaudeCodeExecutionScope(context.Background(), nil, headers)
if !ok || scope != "claude:lower-session:agent:lower-agent" {
t.Fatalf("lowercase header scope = %q, %v", scope, ok)
}
}
func TestClaudeCodeExecutionScopeIsolatesAgents(t *testing.T) {
rootHeaders := http.Header{}
rootHeaders.Set(ClaudeCodeSessionHeader, "session-agents")
childAHeaders := rootHeaders.Clone()
childAHeaders.Set(ClaudeCodeAgentHeader, "agent-a")
childBHeaders := rootHeaders.Clone()
childBHeaders.Set(ClaudeCodeAgentHeader, "agent-b")
rootScope, ok := ClaudeCodeExecutionScope(context.Background(), nil, rootHeaders)
if !ok || rootScope != "claude:session-agents:agent:main" {
t.Fatalf("root scope = %q, %v", rootScope, ok)
}
childAScope, ok := ClaudeCodeExecutionScope(context.Background(), nil, childAHeaders)
if !ok || childAScope != "claude:session-agents:agent:agent-a" {
t.Fatalf("child A scope = %q, %v", childAScope, ok)
}
childBScope, ok := ClaudeCodeExecutionScope(context.Background(), nil, childBHeaders)
if !ok || childBScope != "claude:session-agents:agent:agent-b" {
t.Fatalf("child B scope = %q, %v", childBScope, ok)
}
if rootScope == childAScope || childAScope == childBScope || rootScope == childBScope {
t.Fatalf("agent scopes are not isolated: root=%q a=%q b=%q", rootScope, childAScope, childBScope)
}
}
func TestClaudeCodePromptCacheDeterministicAndAgentScoped(t *testing.T) {
rootHeaders := http.Header{}
rootHeaders.Set(ClaudeCodeSessionHeader, "session-cache-agents")
childHeaders := rootHeaders.Clone()
childHeaders.Set(ClaudeCodeAgentHeader, "agent-a")
rootFirst, ok, errFirst := ClaudeCodePromptCache(context.Background(), "gpt-5.4", nil, rootHeaders)
if errFirst != nil || !ok {
t.Fatalf("root first cache = %#v, %v, %v", rootFirst, ok, errFirst)
}
rootSecond, ok, errSecond := ClaudeCodePromptCache(context.Background(), "gpt-5.4", nil, rootHeaders)
if errSecond != nil || !ok || rootSecond.ID != rootFirst.ID {
t.Fatalf("root second cache = %#v, %v, %v; want ID %q", rootSecond, ok, errSecond, rootFirst.ID)
}
child, ok, errChild := ClaudeCodePromptCache(context.Background(), "gpt-5.4", nil, childHeaders)
if errChild != nil || !ok || child.ID == rootFirst.ID {
t.Fatalf("child cache = %#v, %v, %v; root ID %q", child, ok, errChild, rootFirst.ID)
}
otherModel, ok, errModel := ClaudeCodePromptCache(context.Background(), "gpt-5.5", nil, rootHeaders)
if errModel != nil || !ok || otherModel.ID == rootFirst.ID {
t.Fatalf("other model cache = %#v, %v, %v; root ID %q", otherModel, ok, errModel, rootFirst.ID)
}
}

View file

@ -0,0 +1,457 @@
package helps
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"strings"
"github.com/google/uuid"
claudeauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
"github.com/tidwall/sjson"
)
// ClaudeAgentSessionUUID maps the downstream agent conversation to one stable UUID,
// preserving native Claude Code session signals.
func ClaudeAgentSessionUUID(headers http.Header, originalPayload, translatedPayload []byte, metadataSets ...map[string]any) string {
return claudeAgentSessionUUID(headers, originalPayload, translatedPayload, metadataSets...)
}
// ClaudeAgentSessionUUIDForRequest preserves Claude-specific session signals only
// for a confirmed native caller. Other callers use protocol session fields,
// execution metadata, or the stable derived conversation root.
func ClaudeAgentSessionUUIDForRequest(headers http.Header, originalPayload, translatedPayload []byte, confirmedClaudeCode bool, metadataSets ...map[string]any) string {
if !confirmedClaudeCode {
headers = headers.Clone()
for key := range headers {
if strings.EqualFold(key, "X-Claude-Code-Session-Id") {
delete(headers, key)
}
}
originalPayload = withoutClaudeMetadataUserID(originalPayload)
translatedPayload = withoutClaudeMetadataUserID(translatedPayload)
}
return claudeAgentSessionUUID(headers, originalPayload, translatedPayload, metadataSets...)
}
func claudeAgentSessionUUID(headers http.Header, originalPayload, translatedPayload []byte, metadataSets ...map[string]any) string {
metadata := mergeClaudeSessionMetadata(metadataSets...)
identity := cliproxyauth.ExtractSessionID(headers, originalPayload, metadata)
if identity == "" && len(translatedPayload) > 0 {
identity = cliproxyauth.ExtractSessionID(headers, translatedPayload, metadata)
}
if identity == "" {
return uuid.NewString()
}
if strings.HasPrefix(identity, "claude:") {
if parsed, errParse := uuid.Parse(strings.TrimPrefix(identity, "claude:")); errParse == nil {
return parsed.String()
}
}
if parsed, errParse := uuid.Parse(identity); errParse == nil {
return parsed.String()
}
stableInput := "cli-proxy-api\x00claude\x00agent-conversation\x00" + identity
return uuid.NewSHA1(uuid.NameSpaceOID, []byte(stableInput)).String()
}
func withoutClaudeMetadataUserID(payload []byte) []byte {
if len(payload) == 0 {
return payload
}
updated, errDelete := sjson.DeleteBytes(payload, "metadata.user_id")
if errDelete != nil {
return payload
}
return updated
}
func mergeClaudeSessionMetadata(metadataSets ...map[string]any) map[string]any {
var merged map[string]any
for _, metadata := range metadataSets {
if len(metadata) == 0 {
continue
}
if merged == nil {
merged = make(map[string]any)
}
for key, value := range metadata {
if _, exists := merged[key]; !exists {
merged[key] = value
}
}
}
return merged
}
type claudeCredentialDevicePoolKVClient interface {
KVGet(context.Context, string) ([]byte, bool, error)
KVSet(context.Context, string, []byte, homekv.KVSetOptions) (bool, error)
}
var currentClaudeCredentialDevicePoolKVClient = func() (claudeCredentialDevicePoolKVClient, bool, error) {
client, homeMode, errClient := homekv.CurrentKVClient()
return client, homeMode, errClient
}
// EnsureClaudeCredentialDevicePoolRequired initializes a credential pool locally,
// or coordinates it through Home KV when the selected auth is a remote dispatch clone.
func EnsureClaudeCredentialDevicePoolRequired(ctx context.Context, auth *cliproxyauth.Auth) ([]string, error) {
if auth == nil {
return nil, fmt.Errorf("ensure Claude credential device pool: auth is nil")
}
rawCredentialDeviceIDs := claudeauth.ReadDeviceIDPool(&auth.Metadata)
if claudeauth.HasCanonicalDeviceIDPool(rawCredentialDeviceIDs) {
return claudeauth.NormalizeDeviceIDPool(rawCredentialDeviceIDs), nil
}
credentialCandidate := claudeauth.NormalizeDeviceIDPool(rawCredentialDeviceIDs)
client, homeMode, errClient := currentClaudeCredentialDevicePoolKVClient()
if !homeMode {
deviceIDs, _, errEnsure := claudeauth.EnsureDeviceIDPoolFor(&auth.Metadata)
return deviceIDs, errEnsure
}
if errClient != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: Home KV client: %w", errClient)
}
identity := strings.TrimSpace(auth.EnsureIndex())
if identity == "" {
identity = strings.TrimSpace(auth.ID)
}
if identity == "" {
return nil, fmt.Errorf("ensure Claude credential device pool: credential identity is empty")
}
key := "cpa:claude:credential-device-pool:" + homekv.HashKeyPart(identity)
if raw, found, errGet := client.KVGet(ctx, key); errGet != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: Home KV get: %w", errGet)
} else if found {
var stored []string
if errUnmarshal := json.Unmarshal(raw, &stored); errUnmarshal == nil {
if deviceIDs := claudeauth.NormalizeDeviceIDPool(stored); len(deviceIDs) == claudeauth.ClaudeDevicePoolSize {
if !claudeauth.HasCanonicalDeviceIDPool(stored) {
canonicalRaw, errMarshal := json.Marshal(deviceIDs)
if errMarshal != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: marshal canonical Home KV value: %w", errMarshal)
}
written, errSet := client.KVSet(ctx, key, canonicalRaw, homekv.KVSetOptions{XX: true})
if errSet != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: canonicalize Home KV value: %w", errSet)
}
if !written {
return nil, fmt.Errorf("ensure Claude credential device pool: canonical Home KV value was not written")
}
}
claudeauth.StoreDeviceIDPool(&auth.Metadata, deviceIDs)
return deviceIDs, nil
}
}
}
deviceIDs := credentialCandidate
if len(deviceIDs) != claudeauth.ClaudeDevicePoolSize {
var errGenerate error
deviceIDs, errGenerate = claudeauth.GenerateDeviceIDPool()
if errGenerate != nil {
return nil, errGenerate
}
}
raw, errMarshal := json.Marshal(deviceIDs)
if errMarshal != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: marshal Home KV value: %w", errMarshal)
}
if _, errSet := client.KVSet(ctx, key, raw, homekv.KVSetOptions{NX: true}); errSet != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: Home KV set: %w", errSet)
}
raw, found, errGet := client.KVGet(ctx, key)
if errGet != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: Home KV reread: %w", errGet)
}
if !found {
return nil, fmt.Errorf("ensure Claude credential device pool: Home KV value missing after set")
}
var stored []string
if errUnmarshal := json.Unmarshal(raw, &stored); errUnmarshal != nil {
return nil, fmt.Errorf("ensure Claude credential device pool: decode Home KV value: %w", errUnmarshal)
}
deviceIDs = claudeauth.NormalizeDeviceIDPool(stored)
if len(deviceIDs) != claudeauth.ClaudeDevicePoolSize {
return nil, fmt.Errorf("ensure Claude credential device pool: Home KV pool has %d entries, want %d", len(deviceIDs), claudeauth.ClaudeDevicePoolSize)
}
claudeauth.StoreDeviceIDPool(&auth.Metadata, deviceIDs)
return deviceIDs, nil
}
// ClaudeCredentialAccountUUID returns the selected upstream credential's account UUID.
func ClaudeCredentialAccountUUID(auth *cliproxyauth.Auth) string {
if auth == nil {
return ""
}
for _, key := range []string{"account_uuid", "accountUuid"} {
value := strings.TrimSpace(claudeauth.ReadMetadataString(&auth.Metadata, key))
if value != "" {
return value
}
}
return ""
}
type claudeCredentialMetadataRequestError struct {
cause error
}
func (e *claudeCredentialMetadataRequestError) Error() string {
if e == nil || e.cause == nil {
return ""
}
return e.cause.Error()
}
func (e *claudeCredentialMetadataRequestError) Unwrap() error {
if e == nil {
return nil
}
return e.cause
}
func (e *claudeCredentialMetadataRequestError) StatusCode() int {
if e == nil {
return 0
}
return http.StatusBadRequest
}
func (e *claudeCredentialMetadataRequestError) IsRequestScoped() bool {
return e != nil
}
func newClaudeCredentialMetadataRequestError(err error) error {
if err == nil {
return nil
}
return &claudeCredentialMetadataRequestError{cause: err}
}
// ApplyClaudeCredentialMetadata rewrites the identity exception shared by native and cloaked OAuth requests.
func ApplyClaudeCredentialMetadata(payload []byte, auth *cliproxyauth.Auth, sessionID string) ([]byte, string, error) {
if auth == nil {
return nil, "", fmt.Errorf("apply Claude credential metadata: auth is nil")
}
metadata, metadataPresent, errMetadata := uniqueClaudeJSONObjectMember(payload, "metadata")
if errMetadata != nil {
return nil, "", newClaudeCredentialMetadataRequestError(fmt.Errorf("apply Claude credential metadata: %w", errMetadata))
}
var existing string
if metadataPresent {
trimmedMetadata := bytes.TrimSpace(metadata)
if len(trimmedMetadata) >= 2 && trimmedMetadata[0] == '{' {
userID, userIDPresent, errUserID := uniqueClaudeJSONObjectMember(trimmedMetadata, "user_id")
if errUserID != nil {
return nil, "", newClaudeCredentialMetadataRequestError(fmt.Errorf("apply Claude credential metadata: metadata: %w", errUserID))
}
if userIDPresent && json.Unmarshal(userID, &existing) != nil {
existing = ""
}
}
}
deviceIDs, _, errDeviceIDs := claudeauth.EnsureDeviceIDPoolFor(&auth.Metadata)
if errDeviceIDs != nil {
return nil, "", errDeviceIDs
}
deviceID, errDeviceID := claudeauth.SelectDeviceID(deviceIDs, sessionID)
if errDeviceID != nil {
return nil, "", errDeviceID
}
accountUUID := ClaudeCredentialAccountUUID(auth)
if accountUUID == "" {
return nil, "", fmt.Errorf("apply Claude credential metadata: account UUID is empty")
}
encoded, errIdentity := rebuildClaudeMetadataUserID(existing, deviceID, accountUUID, sessionID)
if errIdentity != nil {
return nil, "", newClaudeCredentialMetadataRequestError(fmt.Errorf("apply Claude credential metadata: %w", errIdentity))
}
updated, errSet := sjson.SetBytes(payload, "metadata.user_id", string(encoded))
if errSet != nil {
return nil, "", fmt.Errorf("set Claude credential metadata: %w", errSet)
}
return updated, deviceID, nil
}
type claudeJSONMember struct {
key string
value json.RawMessage
}
func uniqueClaudeJSONObjectMember(raw []byte, target string) ([]byte, bool, error) {
raw = bytes.TrimSpace(raw)
if !json.Valid(raw) || len(raw) < 2 || raw[0] != '{' {
return nil, false, fmt.Errorf("request must be a JSON object")
}
position := 1
found := false
var value []byte
for {
position = skipClaudeJSONWhitespace(raw, position)
if position >= len(raw) {
return nil, false, fmt.Errorf("unterminated JSON object")
}
if raw[position] == '}' {
break
}
keyStart := position
keyEnd := skipClaudeJSONString(raw, keyStart)
var key string
if errUnmarshal := json.Unmarshal(raw[keyStart:keyEnd], &key); errUnmarshal != nil {
return nil, false, fmt.Errorf("decode JSON object key: %w", errUnmarshal)
}
position = skipClaudeJSONWhitespace(raw, keyEnd)
if position >= len(raw) || raw[position] != ':' {
return nil, false, fmt.Errorf("JSON object key %q is missing a value", key)
}
position = skipClaudeJSONWhitespace(raw, position+1)
valueStart := position
position = skipClaudeJSONValue(raw, position)
if key == target {
if found {
return nil, false, fmt.Errorf("duplicate JSON object key %q", target)
}
found = true
value = raw[valueStart:position]
}
position = skipClaudeJSONWhitespace(raw, position)
if position < len(raw) && raw[position] == ',' {
position++
continue
}
if position >= len(raw) || raw[position] != '}' {
return nil, false, fmt.Errorf("JSON object key %q has an invalid terminator", key)
}
}
return value, found, nil
}
func skipClaudeJSONWhitespace(raw []byte, position int) int {
for position < len(raw) {
switch raw[position] {
case ' ', '\t', '\r', '\n':
position++
default:
return position
}
}
return position
}
func skipClaudeJSONString(raw []byte, position int) int {
if position >= len(raw) || raw[position] != '"' {
return position
}
position++
for position < len(raw) {
switch raw[position] {
case '\\':
position += 2
case '"':
return position + 1
default:
position++
}
}
return position
}
func skipClaudeJSONValue(raw []byte, position int) int {
if position >= len(raw) {
return position
}
switch raw[position] {
case '"':
return skipClaudeJSONString(raw, position)
case '{', '[':
stack := []byte{raw[position]}
position++
for position < len(raw) && len(stack) > 0 {
switch raw[position] {
case '"':
position = skipClaudeJSONString(raw, position)
continue
case '{', '[':
stack = append(stack, raw[position])
case '}', ']':
stack = stack[:len(stack)-1]
}
position++
}
return position
default:
for position < len(raw) {
switch raw[position] {
case ',', '}', ']', ' ', '\t', '\r', '\n':
return position
default:
position++
}
}
return position
}
}
func rebuildClaudeMetadataUserID(existing, deviceID, accountUUID, sessionID string) ([]byte, error) {
extras := make([]claudeJSONMember, 0)
rawExisting := []byte(strings.TrimSpace(existing))
if json.Valid(rawExisting) && len(rawExisting) >= 2 && rawExisting[0] == '{' {
decoder := json.NewDecoder(bytes.NewReader(rawExisting))
_, _ = decoder.Token()
seen := make(map[string]bool)
for decoder.More() {
token, errToken := decoder.Token()
if errToken != nil {
return nil, errToken
}
key, ok := token.(string)
if !ok {
return nil, fmt.Errorf("metadata.user_id contains a non-string key")
}
if seen[key] {
return nil, fmt.Errorf("metadata.user_id contains duplicate key %q", key)
}
seen[key] = true
var value json.RawMessage
if errDecode := decoder.Decode(&value); errDecode != nil {
return nil, errDecode
}
switch key {
case "device_id", "account_uuid", "session_id":
default:
extras = append(extras, claudeJSONMember{key: key, value: value})
}
}
}
var output bytes.Buffer
output.WriteString(`{"device_id":`)
writeClaudeJSONQuoted(&output, deviceID)
output.WriteString(`,"account_uuid":`)
writeClaudeJSONQuoted(&output, accountUUID)
output.WriteString(`,"session_id":`)
writeClaudeJSONQuoted(&output, sessionID)
for _, extra := range extras {
output.WriteByte(',')
writeClaudeJSONQuoted(&output, extra.key)
output.WriteByte(':')
output.Write(extra.value)
}
output.WriteByte('}')
return output.Bytes(), nil
}
func writeClaudeJSONQuoted(output *bytes.Buffer, value string) {
encoded, _ := json.Marshal(value)
output.Write(encoded)
}

View file

@ -0,0 +1,103 @@
package helps
import (
"errors"
"sync"
"testing"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
)
// TestApplyClaudeCredentialMetadataConcurrentSharedAuth pins the invariant that a
// single *Auth shared by concurrent requests is safe to use. Before the device
// pool accessors were introduced these paths initialized and wrote auth.Metadata
// outside claudeDevicePoolMu, which aborts the process with "concurrent map
// writes" rather than failing a request. Run with -race.
func TestApplyClaudeCredentialMetadataConcurrentSharedAuth(t *testing.T) {
auth := &cliproxyauth.Auth{
ID: "shared-credential",
Metadata: map[string]any{"account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"},
}
payload := []byte(`{"model":"claude-opus-4-6","messages":[{"role":"user","content":"hi"}]}`)
const goroutines = 32
var wg sync.WaitGroup
errs := make(chan error, goroutines)
start := make(chan struct{})
for i := range goroutines {
wg.Add(1)
go func(i int) {
defer wg.Done()
<-start
sessionID := "session-" + string(rune('a'+i%26))
if _, _, err := ApplyClaudeCredentialMetadata(payload, auth, sessionID); err != nil {
errs <- err
return
}
// Concurrent readers of the same map must be safe too.
_ = ClaudeCredentialAccountUUID(auth)
}(i)
}
close(start)
wg.Wait()
close(errs)
for err := range errs {
t.Fatalf("ApplyClaudeCredentialMetadata on shared auth: %v", err)
}
if auth.Metadata == nil {
t.Fatal("expected metadata to be initialized")
}
}
// TestEnsureClaudeCredentialDevicePoolConcurrentSharedAuth covers the local
// (non Home KV) branch of the pool bootstrap on a shared credential.
func TestEnsureClaudeCredentialDevicePoolConcurrentSharedAuth(t *testing.T) {
auth := &cliproxyauth.Auth{ID: "shared-credential"}
const goroutines = 32
var wg sync.WaitGroup
results := make(chan string, goroutines)
errs := make(chan error, goroutines)
start := make(chan struct{})
for range goroutines {
wg.Add(1)
go func() {
defer wg.Done()
<-start
deviceIDs, err := EnsureClaudeCredentialDevicePoolRequired(t.Context(), auth)
if err != nil {
errs <- err
return
}
if len(deviceIDs) == 0 {
errs <- errEmptyPool
return
}
results <- deviceIDs[0]
}()
}
close(start)
wg.Wait()
close(errs)
close(results)
for err := range errs {
t.Fatalf("EnsureClaudeCredentialDevicePoolRequired on shared auth: %v", err)
}
// Every caller must agree on the pool; a racing bootstrap would hand out
// different device IDs to different requests on the same credential.
seen := make(map[string]struct{})
for deviceID := range results {
seen[deviceID] = struct{}{}
}
if len(seen) != 1 {
t.Fatalf("device pool bootstrap was not stable: got %d distinct device IDs, want 1", len(seen))
}
}
var errEmptyPool = errors.New("device pool is empty")

View file

@ -0,0 +1,236 @@
package helps
import (
"bytes"
"context"
"encoding/json"
"errors"
"net/http"
"strings"
"testing"
claudeauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
"github.com/tidwall/gjson"
)
type fakeClaudeCredentialDevicePoolKV struct {
values map[string][]byte
setOpts []homekv.KVSetOptions
}
func (fake *fakeClaudeCredentialDevicePoolKV) KVGet(_ context.Context, key string) ([]byte, bool, error) {
value, found := fake.values[key]
return bytes.Clone(value), found, nil
}
func (fake *fakeClaudeCredentialDevicePoolKV) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) {
_, found := fake.values[key]
if (opts.NX && found) || (opts.XX && !found) {
return false, nil
}
fake.values[key] = bytes.Clone(value)
fake.setOpts = append(fake.setOpts, opts)
return true, nil
}
func TestClaudeAgentSessionUUIDPreservesNativeSession(t *testing.T) {
const sessionID = "11111111-2222-4333-8444-555555555555"
got := ClaudeAgentSessionUUIDForRequest(http.Header{"X-Claude-Code-Session-Id": {sessionID}}, nil, nil, true)
if got != sessionID {
t.Fatalf("ClaudeAgentSessionUUIDForRequest() = %q, want native session %q", got, sessionID)
}
}
func TestClaudeAgentSessionUUIDIgnoresUnconfirmedClaudeSignals(t *testing.T) {
const nativeSessionID = "11111111-2222-4333-8444-555555555555"
metadata := map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "non-native-conversation"}
got := ClaudeAgentSessionUUIDForRequest(
http.Header{"X-Claude-Code-Session-Id": {nativeSessionID}},
[]byte(`{"metadata":{"user_id":"{\"device_id\":\"0000000000000000000000000000000000000000000000000000000000000000\",\"session_id\":\"11111111-2222-4333-8444-555555555555\"}"}}`),
nil,
false,
metadata,
)
if got == nativeSessionID {
t.Fatalf("ClaudeAgentSessionUUIDForRequest() = native session %q for unconfirmed caller", got)
}
if repeated := ClaudeAgentSessionUUIDForRequest(nil, nil, nil, false, metadata); repeated != got {
t.Fatalf("derived session changed: first=%q repeated=%q", got, repeated)
}
}
func TestClaudeAgentSessionUUIDUsesExecutionAndDerivedIdentity(t *testing.T) {
tests := []struct {
name string
metadata map[string]any
}{
{
name: "execution session",
metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "agent-run-1"},
},
{
name: "derived session",
metadata: map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:conversation-root"},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
first := ClaudeAgentSessionUUID(nil, nil, nil, test.metadata)
second := ClaudeAgentSessionUUID(nil, nil, nil, test.metadata)
if first == "" || first != second {
t.Fatalf("session UUIDs = %q and %q, want equal non-empty values", first, second)
}
})
}
}
func TestEnsureClaudeCredentialDevicePoolRequiredMigratesHomeKVToOne(t *testing.T) {
auth := &cliproxyauth.Auth{ID: "legacy-five-device-credential", Metadata: map[string]any{}}
key := "cpa:claude:credential-device-pool:" + homekv.HashKeyPart(auth.EnsureIndex())
legacy := []string{
"0000000000000000000000000000000000000000000000000000000000000000",
"1111111111111111111111111111111111111111111111111111111111111111",
"2222222222222222222222222222222222222222222222222222222222222222",
"3333333333333333333333333333333333333333333333333333333333333333",
"4444444444444444444444444444444444444444444444444444444444444444",
}
rawLegacy, errMarshal := json.Marshal(legacy)
if errMarshal != nil {
t.Fatalf("marshal legacy device pool: %v", errMarshal)
}
fake := &fakeClaudeCredentialDevicePoolKV{values: map[string][]byte{key: rawLegacy}}
previousClient := currentClaudeCredentialDevicePoolKVClient
currentClaudeCredentialDevicePoolKVClient = func() (claudeCredentialDevicePoolKVClient, bool, error) {
return fake, true, nil
}
t.Cleanup(func() { currentClaudeCredentialDevicePoolKVClient = previousClient })
deviceIDs, errEnsure := EnsureClaudeCredentialDevicePoolRequired(context.Background(), auth)
if errEnsure != nil {
t.Fatalf("EnsureClaudeCredentialDevicePoolRequired() error = %v", errEnsure)
}
want := []string{legacy[0]}
if len(deviceIDs) != 1 || deviceIDs[0] != want[0] {
t.Fatalf("device IDs = %#v, want %#v", deviceIDs, want)
}
if len(fake.setOpts) != 1 || !fake.setOpts[0].XX || fake.setOpts[0].NX || fake.setOpts[0].EX != 0 || fake.setOpts[0].PX != 0 {
t.Fatalf("Home KV set options = %#v, want one persistent XX rewrite", fake.setOpts)
}
var stored []string
if errUnmarshal := json.Unmarshal(fake.values[key], &stored); errUnmarshal != nil {
t.Fatalf("decode canonical Home KV pool: %v", errUnmarshal)
}
if len(stored) != 1 || stored[0] != want[0] {
t.Fatalf("Home KV device IDs = %#v, want %#v", stored, want)
}
if !claudeauth.HasCanonicalDeviceIDPool(auth.Metadata[claudeauth.ClaudeDeviceIDsMetadataKey]) {
t.Fatalf("auth metadata device pool = %#v, want canonical single device", auth.Metadata[claudeauth.ClaudeDeviceIDsMetadataKey])
}
}
func TestApplyClaudeCredentialMetadataUsesCredentialDeviceAndPreservesExtras(t *testing.T) {
deviceIDs := []string{
"0000000000000000000000000000000000000000000000000000000000000000",
}
auth := &cliproxyauth.Auth{Metadata: map[string]any{
"account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa",
claudeauth.ClaudeDeviceIDsMetadataKey: deviceIDs,
}}
const sessionID = "11111111-2222-4333-8444-555555555555"
body := []byte(`{"messages":[{"role":"user","content":"x"}],"metadata":{"user_id":"{\"device_id\":\"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff\",\"account_uuid\":\"downstream-account\",\"session_id\":\"downstream-session\",\"parent_session_id\":\"parent-1\",\"extra\":true}"}}`)
updated, selectedDevice, errApply := ApplyClaudeCredentialMetadata(body, auth, sessionID)
if errApply != nil {
t.Fatalf("ApplyClaudeCredentialMetadata() error = %v", errApply)
}
userID := gjson.GetBytes(updated, "metadata.user_id").String()
if got := gjson.Get(userID, "device_id").String(); got != selectedDevice {
t.Fatalf("device_id = %q, want selected %q", got, selectedDevice)
}
if got := gjson.Get(userID, "account_uuid").String(); got != "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" {
t.Fatalf("account_uuid = %q, want credential account", got)
}
if got := gjson.Get(userID, "session_id").String(); got != sessionID {
t.Fatalf("session_id = %q, want %q", got, sessionID)
}
if got := gjson.Get(userID, "parent_session_id").String(); got != "parent-1" {
t.Fatalf("parent_session_id = %q, want preserved", got)
}
if !gjson.Get(userID, "extra").Bool() {
t.Fatal("extra metadata was not preserved")
}
wantPrefix := `{"device_id":"` + selectedDevice + `","account_uuid":"aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa","session_id":"` + sessionID + `"`
if !strings.HasPrefix(userID, wantPrefix) {
t.Fatalf("metadata.user_id = %q, want credential identity fields first", userID)
}
}
func TestApplyClaudeCredentialMetadataRejectsDuplicateIdentityContainers(t *testing.T) {
auth := &cliproxyauth.Auth{Metadata: map[string]any{
"account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa",
claudeauth.ClaudeDeviceIDsMetadataKey: []string{
"0000000000000000000000000000000000000000000000000000000000000000",
},
}}
const sessionID = "11111111-2222-4333-8444-555555555555"
tests := []struct {
name string
body string
}{
{
name: "invalid request JSON",
body: `{"messages":[],"metadata":`,
},
{
name: "duplicate top-level metadata",
body: `{"messages":[],"metadata":{"user_id":"{}"},"metadata":{"user_id":"{}"}}`,
},
{
name: "duplicate metadata user ID",
body: `{"messages":[],"metadata":{"user_id":"{}","user_id":"{}"}}`,
},
{
name: "duplicate encoded account UUID",
body: `{"messages":[],"metadata":{"user_id":"{\"account_uuid\":\"first\",\"account_uuid\":\"last\"}"}}`,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
_, _, errApply := ApplyClaudeCredentialMetadata([]byte(test.body), auth, sessionID)
if errApply == nil {
t.Fatal("ApplyClaudeCredentialMetadata() error = nil, want duplicate-key rejection")
}
var requestErr cliproxyexecutor.RequestScopedError
if !errors.As(errApply, &requestErr) || requestErr == nil || !requestErr.IsRequestScoped() {
t.Fatalf("ApplyClaudeCredentialMetadata() error = %T %v, want request-scoped", errApply, errApply)
}
var statusErr interface{ StatusCode() int }
if !errors.As(errApply, &statusErr) || statusErr.StatusCode() != http.StatusBadRequest {
t.Fatalf("ApplyClaudeCredentialMetadata() error = %T %v, want HTTP 400", errApply, errApply)
}
})
}
}
func TestApplyClaudeCredentialMetadataRequiresAccountUUID(t *testing.T) {
auth := &cliproxyauth.Auth{Metadata: map[string]any{
claudeauth.ClaudeDeviceIDsMetadataKey: []string{
"0000000000000000000000000000000000000000000000000000000000000000",
},
}}
_, _, errApply := ApplyClaudeCredentialMetadata(
[]byte(`{"messages":[]}`),
auth,
"11111111-2222-4333-8444-555555555555",
)
if errApply == nil {
t.Fatal("ApplyClaudeCredentialMetadata() error = nil, want missing account UUID rejection")
}
var requestErr cliproxyexecutor.RequestScopedError
if errors.As(errApply, &requestErr) && requestErr != nil && requestErr.IsRequestScoped() {
t.Fatalf("missing credential identity error = %T %v, want credential-scoped", errApply, errApply)
}
}

View file

@ -0,0 +1,634 @@
package helps
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"net/http"
"regexp"
"runtime"
"strconv"
"strings"
"sync"
"time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
)
const (
defaultClaudeFingerprintUserAgent = "claude-cli/2.1.220 (external, cli)"
defaultClaudeFingerprintPackageVersion = "0.94.0"
defaultClaudeFingerprintRuntimeVersion = "v26.3.0"
defaultClaudeFingerprintOS = "MacOS"
defaultClaudeFingerprintArch = "arm64"
claudeDeviceProfileTTL = 7 * 24 * time.Hour
claudeDeviceProfileLockTTL = 5 * time.Second
claudeDeviceProfileCleanupPeriod = time.Hour
)
var (
claudeCLIVersionPattern = regexp.MustCompile(`^claude-cli/(\d+)\.(\d+)\.(\d+)`)
claudePackageVersionPattern = regexp.MustCompile(`^[0-9]+\.[0-9]+\.[0-9]+$`)
claudeRuntimeVersionPattern = regexp.MustCompile(`^v[0-9]+\.[0-9]+\.[0-9]+$`)
claudeDeviceProfileCache = make(map[string]claudeDeviceProfileCacheEntry)
claudeDeviceProfileCacheMu sync.RWMutex
claudeDeviceProfileCacheCleanupOnce sync.Once
ClaudeDeviceProfileBeforeCandidateStore func(ClaudeDeviceProfile)
)
type claudeDeviceProfileKVClient interface {
KVGet(ctx context.Context, key string) ([]byte, bool, error)
KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error)
KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error)
KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error)
}
var currentClaudeDeviceProfileKVClient = func() (claudeDeviceProfileKVClient, bool, error) {
return homekv.CurrentKVClient()
}
type claudeCLIVersion struct {
major int
minor int
patch int
}
func (v claudeCLIVersion) Compare(other claudeCLIVersion) int {
switch {
case v.major != other.major:
if v.major > other.major {
return 1
}
return -1
case v.minor != other.minor:
if v.minor > other.minor {
return 1
}
return -1
case v.patch != other.patch:
if v.patch > other.patch {
return 1
}
return -1
default:
return 0
}
}
type ClaudeDeviceProfile struct {
UserAgent string
PackageVersion string
RuntimeVersion string
OS string
Arch string
version claudeCLIVersion
hasVersion bool
}
type claudeDeviceProfileCacheEntry struct {
profile ClaudeDeviceProfile
expire time.Time
}
type claudeDeviceProfileKVValue struct {
UserAgent string `json:"user_agent"`
PackageVersion string `json:"package_version"`
RuntimeVersion string `json:"runtime_version"`
OS string `json:"os"`
Arch string `json:"arch"`
}
func ClaudeDeviceProfileStabilizationEnabled(cfg *config.Config) bool {
if cfg == nil || cfg.ClaudeHeaderDefaults.StabilizeDeviceProfile == nil {
return false
}
return *cfg.ClaudeHeaderDefaults.StabilizeDeviceProfile
}
func ResetClaudeDeviceProfileCache() {
claudeDeviceProfileCacheMu.Lock()
claudeDeviceProfileCache = make(map[string]claudeDeviceProfileCacheEntry)
claudeDeviceProfileCacheMu.Unlock()
}
func MapStainlessOS() string {
return mapStainlessOS()
}
func MapStainlessArch() string {
return mapStainlessArch()
}
func defaultClaudeDeviceProfile(cfg *config.Config) ClaudeDeviceProfile {
hdrDefault := func(cfgVal, fallback string) string {
if strings.TrimSpace(cfgVal) != "" {
return strings.TrimSpace(cfgVal)
}
return fallback
}
var hd config.ClaudeHeaderDefaults
if cfg != nil {
hd = cfg.ClaudeHeaderDefaults
}
profile := ClaudeDeviceProfile{
UserAgent: hdrDefault(hd.UserAgent, defaultClaudeFingerprintUserAgent),
PackageVersion: hdrDefault(hd.PackageVersion, defaultClaudeFingerprintPackageVersion),
RuntimeVersion: hdrDefault(hd.RuntimeVersion, defaultClaudeFingerprintRuntimeVersion),
OS: hdrDefault(hd.OS, defaultClaudeFingerprintOS),
Arch: hdrDefault(hd.Arch, defaultClaudeFingerprintArch),
}
if version, ok := parseClaudeCLIVersion(profile.UserAgent); ok {
profile.version = version
profile.hasVersion = true
}
return profile
}
// mapStainlessOS maps runtime.GOOS to Stainless SDK OS names.
func mapStainlessOS() string {
switch runtime.GOOS {
case "darwin":
return "MacOS"
case "windows":
return "Windows"
case "linux":
return "Linux"
case "freebsd":
return "FreeBSD"
default:
return "Other::" + runtime.GOOS
}
}
// mapStainlessArch maps runtime.GOARCH to Stainless SDK architecture names.
func mapStainlessArch() string {
switch runtime.GOARCH {
case "amd64":
return "x64"
case "arm64":
return "arm64"
case "386":
return "x86"
default:
return "other::" + runtime.GOARCH
}
}
func parseClaudeCLIVersion(userAgent string) (claudeCLIVersion, bool) {
matches := claudeCLIVersionPattern.FindStringSubmatch(strings.TrimSpace(userAgent))
if len(matches) != 4 {
return claudeCLIVersion{}, false
}
major, err := strconv.Atoi(matches[1])
if err != nil {
return claudeCLIVersion{}, false
}
minor, err := strconv.Atoi(matches[2])
if err != nil {
return claudeCLIVersion{}, false
}
patch, err := strconv.Atoi(matches[3])
if err != nil {
return claudeCLIVersion{}, false
}
return claudeCLIVersion{major: major, minor: minor, patch: patch}, true
}
func shouldUpgradeClaudeDeviceProfile(candidate, current ClaudeDeviceProfile) bool {
if candidate.UserAgent == "" || !candidate.hasVersion {
return false
}
if current.UserAgent == "" || !current.hasVersion {
return true
}
return candidate.version.Compare(current.version) > 0
}
func plausibleClaudeCLIVersion(candidate, baseline claudeCLIVersion) bool {
return candidate.Compare(baseline) == 0
}
func meetsClaudeDeviceProfileBaseline(candidate, baseline ClaudeDeviceProfile) bool {
if candidate.UserAgent == "" || !candidate.hasVersion {
return false
}
if baseline.UserAgent == "" || !baseline.hasVersion {
return false
}
return plausibleClaudeCLIVersion(candidate.version, baseline.version) &&
candidate.PackageVersion == baseline.PackageVersion &&
candidate.RuntimeVersion == baseline.RuntimeVersion
}
func pinClaudeDeviceProfilePlatform(profile, baseline ClaudeDeviceProfile) ClaudeDeviceProfile {
profile.OS = baseline.OS
profile.Arch = baseline.Arch
return profile
}
// normalizeClaudeDeviceProfile pins stabilized profiles to the configured platform
// and replaces any software tuple that does not exactly match the measured baseline.
func normalizeClaudeDeviceProfile(profile, baseline ClaudeDeviceProfile) ClaudeDeviceProfile {
profile = pinClaudeDeviceProfilePlatform(profile, baseline)
if !meetsClaudeDeviceProfileBaseline(profile, baseline) {
profile.UserAgent = baseline.UserAgent
profile.PackageVersion = baseline.PackageVersion
profile.RuntimeVersion = baseline.RuntimeVersion
profile.version = baseline.version
profile.hasVersion = baseline.hasVersion
}
return profile
}
func extractClaudeDeviceProfile(headers http.Header, cfg *config.Config) (ClaudeDeviceProfile, bool) {
if headers == nil {
return ClaudeDeviceProfile{}, false
}
userAgent := strings.TrimSpace(headers.Get("User-Agent"))
version, ok := parseClaudeCLIVersion(userAgent)
if !ok || !claudeCodeNativeUserAgentPattern.MatchString(userAgent) {
return ClaudeDeviceProfile{}, false
}
baseline := defaultClaudeDeviceProfile(cfg)
packageVersion := firstNonEmptyHeader(headers, "X-Stainless-Package-Version", baseline.PackageVersion)
if !claudePackageVersionPattern.MatchString(packageVersion) {
packageVersion = baseline.PackageVersion
}
runtimeVersion := firstNonEmptyHeader(headers, "X-Stainless-Runtime-Version", baseline.RuntimeVersion)
if !claudeRuntimeVersionPattern.MatchString(runtimeVersion) {
runtimeVersion = baseline.RuntimeVersion
}
profile := ClaudeDeviceProfile{
UserAgent: userAgent,
PackageVersion: packageVersion,
RuntimeVersion: runtimeVersion,
OS: firstNonEmptyHeader(headers, "X-Stainless-Os", baseline.OS),
Arch: firstNonEmptyHeader(headers, "X-Stainless-Arch", baseline.Arch),
version: version,
hasVersion: true,
}
return profile, true
}
func firstNonEmptyHeader(headers http.Header, name, fallback string) string {
if headers == nil {
return fallback
}
if value := strings.TrimSpace(headers.Get(name)); value != "" {
return value
}
return fallback
}
func claudeDeviceProfileScopeKey(auth *cliproxyauth.Auth, apiKey string) string {
switch {
case auth != nil && strings.TrimSpace(auth.ID) != "":
return "auth:" + strings.TrimSpace(auth.ID)
case strings.TrimSpace(apiKey) != "":
return "api_key:" + strings.TrimSpace(apiKey)
default:
return "global"
}
}
// claudeDeviceProfileSubclientScope keeps first-party clients with distinct
// wire identities from replacing one another in a credential's stabilized
// profile. The CLI retains the legacy base scope for cache compatibility.
func claudeDeviceProfileSubclientScope(profile ClaudeDeviceProfile) string {
entrypoint, _ := parseClaudeCodeUserAgentDetails(profile.UserAgent)
if entrypoint == "" || entrypoint == "cli" {
return ""
}
if nativeClaudeEntrypoints[entrypoint] {
return entrypoint
}
return "other"
}
func claudeDeviceProfileScopedKey(auth *cliproxyauth.Auth, apiKey string, profile ClaudeDeviceProfile) string {
key := claudeDeviceProfileScopeKey(auth, apiKey)
if subclient := claudeDeviceProfileSubclientScope(profile); subclient != "" {
key += "|subclient:" + subclient
}
return key
}
func claudeDeviceProfileCacheKey(auth *cliproxyauth.Auth, apiKey string, profile ClaudeDeviceProfile) string {
sum := sha256.Sum256([]byte(claudeDeviceProfileScopedKey(auth, apiKey, profile)))
return hex.EncodeToString(sum[:])
}
func claudeDeviceProfileKVKey(auth *cliproxyauth.Auth, apiKey string, profile ClaudeDeviceProfile) string {
return "cpa:claude:device-profile:" + homekv.HashKeyPart(claudeDeviceProfileScopedKey(auth, apiKey, profile))
}
func claudeDeviceProfileLockKVKey(auth *cliproxyauth.Auth, apiKey string, profile ClaudeDeviceProfile) string {
return "cpa:claude:device-profile-lock:" + homekv.HashKeyPart(claudeDeviceProfileScopedKey(auth, apiKey, profile))
}
func startClaudeDeviceProfileCacheCleanup() {
go func() {
ticker := time.NewTicker(claudeDeviceProfileCleanupPeriod)
defer ticker.Stop()
for range ticker.C {
purgeExpiredClaudeDeviceProfiles()
}
}()
}
func purgeExpiredClaudeDeviceProfiles() {
now := time.Now()
claudeDeviceProfileCacheMu.Lock()
for key, entry := range claudeDeviceProfileCache {
if !entry.expire.After(now) {
delete(claudeDeviceProfileCache, key)
}
}
claudeDeviceProfileCacheMu.Unlock()
}
func ResolveClaudeDeviceProfile(auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) ClaudeDeviceProfile {
profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, apiKey, headers, cfg)
if errProfile != nil {
return defaultClaudeDeviceProfile(cfg)
}
return profile
}
// ResolveClaudeDeviceProfileRequired resolves a stable Claude Code device profile for request-time paths.
func ResolveClaudeDeviceProfileRequired(ctx context.Context, auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) (ClaudeDeviceProfile, error) {
client, homeMode, errClient := currentClaudeDeviceProfileKVClient()
if homeMode {
if errClient != nil {
return ClaudeDeviceProfile{}, errClient
}
return resolveClaudeDeviceProfileHome(ctx, client, auth, apiKey, headers, cfg)
}
return resolveClaudeDeviceProfileLocal(auth, apiKey, headers, cfg), nil
}
func resolveClaudeDeviceProfileLocal(auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) ClaudeDeviceProfile {
claudeDeviceProfileCacheCleanupOnce.Do(startClaudeDeviceProfileCacheCleanup)
now := time.Now()
baseline := defaultClaudeDeviceProfile(cfg)
candidate, hasCandidate := extractClaudeDeviceProfile(headers, cfg)
if hasCandidate {
candidate = pinClaudeDeviceProfilePlatform(candidate, baseline)
}
if hasCandidate && !meetsClaudeDeviceProfileBaseline(candidate, baseline) {
hasCandidate = false
}
cacheProfile := ClaudeDeviceProfile{}
if hasCandidate {
cacheProfile = candidate
}
cacheKey := claudeDeviceProfileCacheKey(auth, apiKey, cacheProfile)
claudeDeviceProfileCacheMu.RLock()
entry, hasCached := claudeDeviceProfileCache[cacheKey]
cachedValid := hasCached && entry.expire.After(now) && entry.profile.UserAgent != ""
claudeDeviceProfileCacheMu.RUnlock()
if hasCandidate {
if ClaudeDeviceProfileBeforeCandidateStore != nil {
ClaudeDeviceProfileBeforeCandidateStore(candidate)
}
claudeDeviceProfileCacheMu.Lock()
entry, hasCached = claudeDeviceProfileCache[cacheKey]
cachedValid = hasCached && entry.expire.After(now) && entry.profile.UserAgent != ""
if cachedValid {
entry.profile = normalizeClaudeDeviceProfile(entry.profile, baseline)
}
if cachedValid && !shouldUpgradeClaudeDeviceProfile(candidate, entry.profile) {
entry.expire = now.Add(claudeDeviceProfileTTL)
claudeDeviceProfileCache[cacheKey] = entry
claudeDeviceProfileCacheMu.Unlock()
return entry.profile
}
claudeDeviceProfileCache[cacheKey] = claudeDeviceProfileCacheEntry{
profile: candidate,
expire: now.Add(claudeDeviceProfileTTL),
}
claudeDeviceProfileCacheMu.Unlock()
return candidate
}
if cachedValid {
claudeDeviceProfileCacheMu.Lock()
entry = claudeDeviceProfileCache[cacheKey]
if entry.expire.After(now) && entry.profile.UserAgent != "" {
entry.profile = normalizeClaudeDeviceProfile(entry.profile, baseline)
entry.expire = now.Add(claudeDeviceProfileTTL)
claudeDeviceProfileCache[cacheKey] = entry
claudeDeviceProfileCacheMu.Unlock()
return entry.profile
}
claudeDeviceProfileCacheMu.Unlock()
}
return baseline
}
func resolveClaudeDeviceProfileHome(ctx context.Context, client claudeDeviceProfileKVClient, auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) (ClaudeDeviceProfile, error) {
baseline := defaultClaudeDeviceProfile(cfg)
candidate, hasCandidate := extractClaudeDeviceProfile(headers, cfg)
if hasCandidate {
candidate = pinClaudeDeviceProfilePlatform(candidate, baseline)
}
if hasCandidate && !meetsClaudeDeviceProfileBaseline(candidate, baseline) {
hasCandidate = false
}
cacheProfile := ClaudeDeviceProfile{}
if hasCandidate {
cacheProfile = candidate
}
valueKey := claudeDeviceProfileKVKey(auth, apiKey, cacheProfile)
if !hasCandidate {
return readClaudeDeviceProfileFromHome(ctx, client, valueKey, baseline)
}
lockKey := claudeDeviceProfileLockKVKey(auth, apiKey, cacheProfile)
gotLock, errLock := client.KVSetNX(ctx, lockKey, []byte("1"), claudeDeviceProfileLockTTL)
if errLock != nil {
return ClaudeDeviceProfile{}, errLock
}
if ClaudeDeviceProfileBeforeCandidateStore != nil {
ClaudeDeviceProfileBeforeCandidateStore(candidate)
}
cached, found, errRead := readClaudeDeviceProfileValueFromHome(ctx, client, valueKey, baseline)
if errRead != nil {
return ClaudeDeviceProfile{}, errRead
}
if found && !shouldUpgradeClaudeDeviceProfile(candidate, cached) {
if _, errExpire := client.KVExpire(ctx, valueKey, claudeDeviceProfileTTL); errExpire != nil {
return ClaudeDeviceProfile{}, errExpire
}
return cached, nil
}
if !gotLock {
if found {
return cached, nil
}
return ClaudeDeviceProfile{}, fmt.Errorf("home kv device profile lock not acquired and profile missing")
}
if errWrite := writeClaudeDeviceProfileToHome(ctx, client, valueKey, candidate); errWrite != nil {
return ClaudeDeviceProfile{}, errWrite
}
return candidate, nil
}
func readClaudeDeviceProfileFromHome(ctx context.Context, client claudeDeviceProfileKVClient, key string, baseline ClaudeDeviceProfile) (ClaudeDeviceProfile, error) {
profile, found, errRead := readClaudeDeviceProfileValueFromHome(ctx, client, key, baseline)
if errRead != nil {
return ClaudeDeviceProfile{}, errRead
}
if !found {
return baseline, nil
}
if _, errExpire := client.KVExpire(ctx, key, claudeDeviceProfileTTL); errExpire != nil {
return ClaudeDeviceProfile{}, errExpire
}
return profile, nil
}
func readClaudeDeviceProfileValueFromHome(ctx context.Context, client claudeDeviceProfileKVClient, key string, baseline ClaudeDeviceProfile) (ClaudeDeviceProfile, bool, error) {
raw, found, errGet := client.KVGet(ctx, key)
if errGet != nil || !found {
return ClaudeDeviceProfile{}, false, errGet
}
var value claudeDeviceProfileKVValue
if errUnmarshal := json.Unmarshal(raw, &value); errUnmarshal != nil {
return ClaudeDeviceProfile{}, false, errUnmarshal
}
profile := value.ToProfile()
if strings.TrimSpace(profile.UserAgent) == "" {
return ClaudeDeviceProfile{}, false, nil
}
return normalizeClaudeDeviceProfile(profile, baseline), true, nil
}
func writeClaudeDeviceProfileToHome(ctx context.Context, client claudeDeviceProfileKVClient, key string, profile ClaudeDeviceProfile) error {
raw, errMarshal := json.Marshal(claudeDeviceProfileKVValueFromProfile(profile))
if errMarshal != nil {
return errMarshal
}
written, errSet := client.KVSet(ctx, key, raw, homekv.KVSetOptions{EX: claudeDeviceProfileTTL})
if errSet != nil {
return errSet
}
if !written {
return fmt.Errorf("home kv device profile write skipped")
}
return nil
}
func claudeDeviceProfileKVValueFromProfile(profile ClaudeDeviceProfile) claudeDeviceProfileKVValue {
return claudeDeviceProfileKVValue{
UserAgent: profile.UserAgent,
PackageVersion: profile.PackageVersion,
RuntimeVersion: profile.RuntimeVersion,
OS: profile.OS,
Arch: profile.Arch,
}
}
func (value claudeDeviceProfileKVValue) ToProfile() ClaudeDeviceProfile {
profile := ClaudeDeviceProfile{
UserAgent: strings.TrimSpace(value.UserAgent),
PackageVersion: strings.TrimSpace(value.PackageVersion),
RuntimeVersion: strings.TrimSpace(value.RuntimeVersion),
OS: strings.TrimSpace(value.OS),
Arch: strings.TrimSpace(value.Arch),
}
if version, ok := parseClaudeCLIVersion(profile.UserAgent); ok {
profile.version = version
profile.hasVersion = true
}
return profile
}
func ApplyClaudeDeviceProfileHeaders(r *http.Request, profile ClaudeDeviceProfile) {
if r == nil {
return
}
for _, headerName := range []string{
"User-Agent",
"X-Stainless-Package-Version",
"X-Stainless-Runtime-Version",
"X-Stainless-Os",
"X-Stainless-Arch",
} {
r.Header.Del(headerName)
}
r.Header.Set("User-Agent", profile.UserAgent)
r.Header.Set("X-Stainless-Package-Version", profile.PackageVersion)
r.Header.Set("X-Stainless-Runtime-Version", profile.RuntimeVersion)
r.Header.Set("X-Stainless-Os", profile.OS)
r.Header.Set("X-Stainless-Arch", profile.Arch)
}
// DefaultClaudeVersion returns the version string (e.g. "2.1.220") from the
// current baseline device profile. It extracts the version from the User-Agent.
func DefaultClaudeVersion(cfg *config.Config) string {
profile := defaultClaudeDeviceProfile(cfg)
if version, ok := parseClaudeCLIVersion(profile.UserAgent); ok {
return strconv.Itoa(version.major) + "." + strconv.Itoa(version.minor) + "." + strconv.Itoa(version.patch)
}
return "2.1.220"
}
func ApplyClaudeDefaultDeviceProfileHeaders(r *http.Request, cfg *config.Config) {
ApplyClaudeDeviceProfileHeaders(r, defaultClaudeDeviceProfile(cfg))
}
func ApplyClaudeLegacyDeviceHeaders(r *http.Request, ginHeaders http.Header, cfg *config.Config, confirmedClaudeCode bool) {
if r == nil {
return
}
profile := defaultClaudeDeviceProfile(cfg)
miscEnsure := func(name, fallback string, valid func(string) bool) {
if current := strings.TrimSpace(r.Header.Get(name)); current != "" && (valid == nil || valid(current)) {
return
}
if incoming := strings.TrimSpace(ginHeaders.Get(name)); incoming != "" && (valid == nil || valid(incoming)) {
r.Header.Set(name, incoming)
return
}
r.Header.Set(name, fallback)
}
if confirmedClaudeCode {
miscEnsure("X-Stainless-Runtime-Version", profile.RuntimeVersion, func(value string) bool { return value == profile.RuntimeVersion })
miscEnsure("X-Stainless-Package-Version", profile.PackageVersion, func(value string) bool { return value == profile.PackageVersion })
miscEnsure("X-Stainless-Os", mapStainlessOS(), nil)
miscEnsure("X-Stainless-Arch", mapStainlessArch(), nil)
if clientUA := strings.TrimSpace(ginHeaders.Get("User-Agent")); plausibleClaudeCodeUserAgent(clientUA, cfg) {
r.Header.Set("User-Agent", clientUA)
return
}
}
// Unconfirmed clients must not leak a copied or third-party software profile
// into the upstream Claude Code SDK fingerprint.
r.Header.Set("X-Stainless-Runtime-Version", profile.RuntimeVersion)
r.Header.Set("X-Stainless-Package-Version", profile.PackageVersion)
r.Header.Set("X-Stainless-Os", profile.OS)
r.Header.Set("X-Stainless-Arch", profile.Arch)
r.Header.Set("User-Agent", profile.UserAgent)
}

View file

@ -0,0 +1,400 @@
package helps
import (
"context"
"encoding/json"
"errors"
"net/http"
"testing"
"time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
)
type fakeClaudeDeviceProfileKVClient struct {
values map[string][]byte
getErr error
setErr error
setNXErr error
expireErr error
setNXResult bool
getCount int
setCount int
setNXCount int
expireCount int
lastSetTTL time.Duration
lastSetNXTTL time.Duration
lastExpireTTL time.Duration
}
func newFakeClaudeDeviceProfileKVClient() *fakeClaudeDeviceProfileKVClient {
return &fakeClaudeDeviceProfileKVClient{
values: make(map[string][]byte),
setNXResult: true,
}
}
func (c *fakeClaudeDeviceProfileKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) {
c.getCount++
if c.getErr != nil {
return nil, false, c.getErr
}
value, ok := c.values[key]
if !ok {
return nil, false, nil
}
return append([]byte(nil), value...), true, nil
}
func (c *fakeClaudeDeviceProfileKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) {
c.setCount++
c.lastSetTTL = opts.EX
if c.setErr != nil {
return false, c.setErr
}
c.values[key] = append([]byte(nil), value...)
return true, nil
}
func (c *fakeClaudeDeviceProfileKVClient) KVSetNX(_ context.Context, key string, value []byte, ttl time.Duration) (bool, error) {
c.setNXCount++
c.lastSetNXTTL = ttl
if c.setNXErr != nil {
return false, c.setNXErr
}
if _, ok := c.values[key]; ok {
return false, nil
}
if c.setNXResult {
c.values[key] = append([]byte(nil), value...)
return true, nil
}
return false, nil
}
func (c *fakeClaudeDeviceProfileKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) {
c.expireCount++
c.lastExpireTTL = ttl
if c.expireErr != nil {
return false, c.expireErr
}
return true, nil
}
func useFakeClaudeDeviceProfileKVClient(t *testing.T, client *fakeClaudeDeviceProfileKVClient, homeMode bool, errClient error) {
t.Helper()
previous := currentClaudeDeviceProfileKVClient
currentClaudeDeviceProfileKVClient = func() (claudeDeviceProfileKVClient, bool, error) {
return client, homeMode, errClient
}
t.Cleanup(func() {
currentClaudeDeviceProfileKVClient = previous
})
}
func mustClaudeDeviceProfileJSON(t *testing.T, value claudeDeviceProfileKVValue) []byte {
t.Helper()
raw, errMarshal := json.Marshal(value)
if errMarshal != nil {
t.Fatalf("marshal device profile: %v", errMarshal)
}
return raw
}
func claudeDeviceHeaders(userAgent string) http.Header {
return http.Header{
"User-Agent": {userAgent},
"X-Stainless-Package-Version": {defaultClaudeFingerprintPackageVersion},
"X-Stainless-Runtime-Version": {defaultClaudeFingerprintRuntimeVersion},
"X-Stainless-Os": {"Windows"},
"X-Stainless-Arch": {"x64"},
}
}
func TestResolveClaudeDeviceProfileLocalUsesBaselineForInvalidSignals(t *testing.T) {
ResetClaudeDeviceProfileCache()
auth := &cliproxyauth.Auth{ID: "auth-invalid-signals"}
headers := claudeDeviceHeaders("claude-cli/999.0.0 (external, cli)")
headers.Set("X-Stainless-Package-Version", "999.0.0")
headers.Set("X-Stainless-Runtime-Version", "v999.0.0")
profile := resolveClaudeDeviceProfileLocal(auth, "api-key", headers, nil)
baseline := defaultClaudeDeviceProfile(nil)
if profile.UserAgent != baseline.UserAgent || profile.PackageVersion != baseline.PackageVersion || profile.RuntimeVersion != baseline.RuntimeVersion {
t.Fatalf("invalid profile = %#v, want local baseline %#v", profile, baseline)
}
}
func TestApplyClaudeLegacyDeviceHeadersReplacesInvalidNativeSoftwareSignals(t *testing.T) {
request, errRequest := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages", nil)
if errRequest != nil {
t.Fatal(errRequest)
}
incoming := claudeDeviceHeaders("claude-cli/999.0.0 (external, cli)")
incoming.Set("X-Stainless-Package-Version", "999.0.0")
incoming.Set("X-Stainless-Runtime-Version", "v999.0.0")
ApplyClaudeLegacyDeviceHeaders(request, incoming, nil, true)
baseline := defaultClaudeDeviceProfile(nil)
if got := request.Header.Get("User-Agent"); got != baseline.UserAgent {
t.Fatalf("User-Agent = %q, want local baseline %q", got, baseline.UserAgent)
}
if got := request.Header.Get("X-Stainless-Package-Version"); got != baseline.PackageVersion {
t.Fatalf("X-Stainless-Package-Version = %q, want %q", got, baseline.PackageVersion)
}
if got := request.Header.Get("X-Stainless-Runtime-Version"); got != baseline.RuntimeVersion {
t.Fatalf("X-Stainless-Runtime-Version = %q, want %q", got, baseline.RuntimeVersion)
}
}
func TestApplyClaudeLegacyDeviceHeadersAcceptsConfiguredMeasuredBaseline(t *testing.T) {
request, errRequest := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages", nil)
if errRequest != nil {
t.Fatal(errRequest)
}
cfg := &config.Config{ClaudeHeaderDefaults: config.ClaudeHeaderDefaults{
UserAgent: "claude-cli/2.2.0 (external, cli)",
PackageVersion: "0.95.0",
RuntimeVersion: "v26.4.0",
OS: "MacOS",
Arch: "arm64",
}}
incoming := claudeDeviceHeaders("claude-cli/2.2.0 (external, cli)")
incoming.Set("X-Stainless-Package-Version", "0.95.0")
incoming.Set("X-Stainless-Runtime-Version", "v26.4.0")
ApplyClaudeLegacyDeviceHeaders(request, incoming, cfg, true)
if got := request.Header.Get("User-Agent"); got != "claude-cli/2.2.0 (external, cli)" {
t.Fatalf("User-Agent = %q, want configured measured baseline", got)
}
if got := request.Header.Get("X-Stainless-Package-Version"); got != "0.95.0" {
t.Fatalf("X-Stainless-Package-Version = %q, want 0.95.0", got)
}
if got := request.Header.Get("X-Stainless-Runtime-Version"); got != "v26.4.0" {
t.Fatalf("X-Stainless-Runtime-Version = %q, want v26.4.0", got)
}
}
func TestResolveClaudeDeviceProfileRequiredHomeReadWithoutCandidate(t *testing.T) {
client := newFakeClaudeDeviceProfileKVClient()
auth := &cliproxyauth.Auth{ID: "auth-1"}
key := claudeDeviceProfileKVKey(auth, "api-key", ClaudeDeviceProfile{})
client.values[key] = mustClaudeDeviceProfileJSON(t, claudeDeviceProfileKVValue{
UserAgent: "claude-cli/2.2.0 (external, cli)",
PackageVersion: "0.80.0",
RuntimeVersion: "v24.4.0",
OS: "Windows",
Arch: "x64",
})
useFakeClaudeDeviceProfileKVClient(t, client, true, nil)
profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", nil, nil)
if errProfile != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile)
}
if profile.UserAgent != defaultClaudeFingerprintUserAgent {
t.Fatalf("UserAgent = %q, want local baseline %q for unmeasured cached profile", profile.UserAgent, defaultClaudeFingerprintUserAgent)
}
if profile.OS != defaultClaudeFingerprintOS || profile.Arch != defaultClaudeFingerprintArch {
t.Fatalf("platform = %s/%s, want baseline pinned %s/%s", profile.OS, profile.Arch, defaultClaudeFingerprintOS, defaultClaudeFingerprintArch)
}
if client.expireCount != 1 || client.lastExpireTTL != claudeDeviceProfileTTL {
t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, claudeDeviceProfileTTL)
}
}
func TestResolveClaudeDeviceProfileRequiredHomeCandidateLocksRereadsAndWrites(t *testing.T) {
client := newFakeClaudeDeviceProfileKVClient()
auth := &cliproxyauth.Auth{ID: "auth-1"}
useFakeClaudeDeviceProfileKVClient(t, client, true, nil)
profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders(defaultClaudeFingerprintUserAgent), nil)
if errProfile != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile)
}
if profile.UserAgent != defaultClaudeFingerprintUserAgent {
t.Fatalf("UserAgent = %q, want candidate %q", profile.UserAgent, defaultClaudeFingerprintUserAgent)
}
if client.setNXCount != 1 || client.lastSetNXTTL != claudeDeviceProfileLockTTL {
t.Fatalf("KVSetNX count/ttl = %d/%v, want 1/%v", client.setNXCount, client.lastSetNXTTL, claudeDeviceProfileLockTTL)
}
if client.getCount != 1 {
t.Fatalf("KVGet count = %d, want re-read after lock", client.getCount)
}
if client.setCount != 1 || client.lastSetTTL != claudeDeviceProfileTTL {
t.Fatalf("KVSet count/ttl = %d/%v, want 1/%v", client.setCount, client.lastSetTTL, claudeDeviceProfileTTL)
}
}
func TestResolveClaudeDeviceProfileRequiredHomeSeparatesVSCodeAgentSDKFromCLI(t *testing.T) {
client := newFakeClaudeDeviceProfileKVClient()
auth := &cliproxyauth.Auth{ID: "auth-home-subclient-isolation"}
useFakeClaudeDeviceProfileKVClient(t, client, true, nil)
cliProfile, errCLI := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders(defaultClaudeFingerprintUserAgent), nil)
if errCLI != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() CLI error = %v", errCLI)
}
vscodeUA := "claude-cli/2.1.220 (external, claude-vscode, agent-sdk/0.3.220)"
vscodeProfile, errVSCode := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders(vscodeUA), nil)
if errVSCode != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() VSCode error = %v", errVSCode)
}
if cliProfile.UserAgent != defaultClaudeFingerprintUserAgent {
t.Fatalf("CLI UserAgent = %q, want CLI profile", cliProfile.UserAgent)
}
if vscodeProfile.UserAgent != vscodeUA {
t.Fatalf("VSCode UserAgent = %q, want %q", vscodeProfile.UserAgent, vscodeUA)
}
if client.setCount != 2 {
t.Fatalf("KVSet count = %d, want separate CLI and VSCode profiles", client.setCount)
}
cliKey := claudeDeviceProfileKVKey(auth, "api-key", cliProfile)
vscodeKey := claudeDeviceProfileKVKey(auth, "api-key", vscodeProfile)
if cliKey == vscodeKey {
t.Fatalf("CLI and VSCode KV keys are equal: %q", cliKey)
}
if _, ok := client.values[cliKey]; !ok {
t.Fatalf("CLI profile missing from KV key %q", cliKey)
}
if _, ok := client.values[vscodeKey]; !ok {
t.Fatalf("VSCode profile missing from KV key %q", vscodeKey)
}
}
func TestResolveClaudeDeviceProfileRequiredHomeNormalizesUnmeasuredCachedProfile(t *testing.T) {
client := newFakeClaudeDeviceProfileKVClient()
auth := &cliproxyauth.Auth{ID: "auth-1"}
key := claudeDeviceProfileKVKey(auth, "api-key", ClaudeDeviceProfile{})
client.values[key] = mustClaudeDeviceProfileJSON(t, claudeDeviceProfileKVValue{
UserAgent: "claude-cli/2.4.0 (external, cli)",
PackageVersion: "0.90.0",
RuntimeVersion: "v24.5.0",
OS: "Windows",
Arch: "x64",
})
useFakeClaudeDeviceProfileKVClient(t, client, true, nil)
profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders("claude-cli/2.3.0 (external, cli)"), nil)
if errProfile != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile)
}
if profile.UserAgent != defaultClaudeFingerprintUserAgent {
t.Fatalf("UserAgent = %q, want local baseline %q", profile.UserAgent, defaultClaudeFingerprintUserAgent)
}
if client.setCount != 0 {
t.Fatalf("KVSet count = %d, want no downgrade write", client.setCount)
}
if client.expireCount != 1 {
t.Fatalf("KVExpire count = %d, want cached refresh", client.expireCount)
}
}
func TestResolveClaudeDeviceProfileRequiredHomeFailures(t *testing.T) {
for _, tc := range []struct {
name string
headers http.Header
client *fakeClaudeDeviceProfileKVClient
}{
{name: "read", client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}},
{name: "lock", headers: claudeDeviceHeaders(defaultClaudeFingerprintUserAgent), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: true, setNXErr: errors.New("lock failed")}},
{name: "lock-miss", headers: claudeDeviceHeaders(defaultClaudeFingerprintUserAgent), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: false}},
{name: "reread", headers: claudeDeviceHeaders(defaultClaudeFingerprintUserAgent), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: true, getErr: errors.New("re-read failed")}},
{name: "write", headers: claudeDeviceHeaders(defaultClaudeFingerprintUserAgent), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: true, setErr: errors.New("write failed")}},
} {
t.Run(tc.name, func(t *testing.T) {
useFakeClaudeDeviceProfileKVClient(t, tc.client, true, nil)
if _, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), &cliproxyauth.Auth{ID: "auth-1"}, "api-key", tc.headers, nil); errProfile == nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() error = nil, want error")
}
})
}
}
func TestResolveClaudeDeviceProfilePreservesConfirmedClientAtBaselineVersion(t *testing.T) {
ResetClaudeDeviceProfileCache()
client := newFakeClaudeDeviceProfileKVClient()
useFakeClaudeDeviceProfileKVClient(t, client, false, nil)
auth := &cliproxyauth.Auth{ID: "auth-baseline-entrypoint"}
headers := claudeDeviceHeaders("claude-cli/2.1.220 (external, cli)")
headers.Set("X-Stainless-Package-Version", "0.94.0")
headers.Set("X-Stainless-Runtime-Version", "v26.3.0")
profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", headers, nil)
if errProfile != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile)
}
if profile.UserAgent != "claude-cli/2.1.220 (external, cli)" {
t.Fatalf("UserAgent = %q, want confirmed cli entrypoint preserved", profile.UserAgent)
}
if profile.PackageVersion != "0.94.0" || profile.RuntimeVersion != "v26.3.0" {
t.Fatalf("software profile = %s/%s, want 0.94.0/v26.3.0", profile.PackageVersion, profile.RuntimeVersion)
}
}
func TestResolveClaudeDeviceProfileSeparatesVSCodeAgentSDKFromCLI(t *testing.T) {
ResetClaudeDeviceProfileCache()
client := newFakeClaudeDeviceProfileKVClient()
useFakeClaudeDeviceProfileKVClient(t, client, false, nil)
auth := &cliproxyauth.Auth{ID: "auth-subclient-isolation"}
cliHeaders := claudeDeviceHeaders("claude-cli/2.1.220 (external, cli)")
cliHeaders.Set("X-Stainless-Package-Version", "0.94.0")
cliHeaders.Set("X-Stainless-Runtime-Version", "v26.3.0")
cliProfile, errCLI := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", cliHeaders, nil)
if errCLI != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() CLI error = %v", errCLI)
}
vscodeUA := "claude-cli/2.1.220 (external, claude-vscode, agent-sdk/0.3.220)"
vscodeHeaders := claudeDeviceHeaders(vscodeUA)
vscodeHeaders.Set("X-Stainless-Package-Version", "0.94.0")
vscodeHeaders.Set("X-Stainless-Runtime-Version", "v26.3.0")
vscodeProfile, errVSCode := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", vscodeHeaders, nil)
if errVSCode != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() VSCode error = %v", errVSCode)
}
if cliProfile.UserAgent != "claude-cli/2.1.220 (external, cli)" {
t.Fatalf("CLI UserAgent = %q, want CLI profile", cliProfile.UserAgent)
}
if vscodeProfile.UserAgent != vscodeUA {
t.Fatalf("VSCode UserAgent = %q, want %q", vscodeProfile.UserAgent, vscodeUA)
}
cliProfileAgain, errCLIAgain := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", cliHeaders, nil)
if errCLIAgain != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() second CLI error = %v", errCLIAgain)
}
if cliProfileAgain.UserAgent != cliProfile.UserAgent {
t.Fatalf("second CLI UserAgent = %q, want isolated cached %q", cliProfileAgain.UserAgent, cliProfile.UserAgent)
}
}
func TestResolveClaudeDeviceProfileRequiredNonHomeKeepsLocalCache(t *testing.T) {
ResetClaudeDeviceProfileCache()
client := newFakeClaudeDeviceProfileKVClient()
useFakeClaudeDeviceProfileKVClient(t, client, false, nil)
auth := &cliproxyauth.Auth{ID: "auth-1"}
cfg := &config.Config{}
first, errFirst := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders(defaultClaudeFingerprintUserAgent), cfg)
if errFirst != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() first error = %v", errFirst)
}
second, errSecond := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", nil, cfg)
if errSecond != nil {
t.Fatalf("ResolveClaudeDeviceProfileRequired() second error = %v", errSecond)
}
if second.UserAgent != first.UserAgent {
t.Fatalf("cached UserAgent = %q, want %q", second.UserAgent, first.UserAgent)
}
if client.getCount != 0 || client.setCount != 0 || client.setNXCount != 0 {
t.Fatalf("KV calls = get %d set %d setnx %d, want all zero", client.getCount, client.setCount, client.setNXCount)
}
}

View file

@ -0,0 +1,137 @@
package helps
import (
"crypto/sha256"
"encoding/hex"
"sort"
"strings"
"sync"
"time"
)
const (
claudeDiagnosticsTTL = time.Hour
claudeDiagnosticsCleanupPeriod = 15 * time.Minute
claudeDiagnosticsMaxEntries = 4096
claudeDiagnosticsEvictBatchSize = 256
)
type claudeDiagnosticsEntry struct {
previousMessageID string
minimumSequence uint64
committedSequence uint64
lastAccess uint64
expiresAt time.Time
}
var claudeDiagnosticsState = struct {
sync.Mutex
entries map[string]claudeDiagnosticsEntry
lastCleanup time.Time
nextSequence uint64
nextAccess uint64
}{entries: make(map[string]claudeDiagnosticsEntry)}
// BeginClaudeDiagnostics starts one request generation for a stable credential
// identity and Claude conversation. It returns the last successfully completed
// upstream message ID, if any. Only a SHA-256 digest of the credential identity
// and session is retained as the cache key, so access-token rotation does not
// interrupt continuity.
func BeginClaudeDiagnostics(credentialIdentity, sessionID string) (key string, sequence uint64, previousMessageID string) {
credentialIdentity = strings.TrimSpace(credentialIdentity)
sessionID = strings.TrimSpace(sessionID)
if credentialIdentity == "" || sessionID == "" {
return "", 0, ""
}
digest := sha256.Sum256([]byte(credentialIdentity + "\x00" + sessionID))
key = hex.EncodeToString(digest[:])
now := time.Now()
claudeDiagnosticsState.Lock()
defer claudeDiagnosticsState.Unlock()
cleanupClaudeDiagnosticsLocked(now)
entry, found := claudeDiagnosticsState.entries[key]
newGeneration := !found || (!entry.expiresAt.IsZero() && now.After(entry.expiresAt))
if newGeneration && !found {
evictClaudeDiagnosticsLocked()
}
claudeDiagnosticsState.nextSequence++
sequence = claudeDiagnosticsState.nextSequence
if newGeneration {
entry = claudeDiagnosticsEntry{minimumSequence: sequence}
}
claudeDiagnosticsState.nextAccess++
entry.lastAccess = claudeDiagnosticsState.nextAccess
entry.expiresAt = now.Add(claudeDiagnosticsTTL)
claudeDiagnosticsState.entries[key] = entry
return key, sequence, entry.previousMessageID
}
// CommitClaudeDiagnostics advances continuity only after a response completes.
// A response from an older concurrently-started request cannot overwrite a
// newer committed generation, including after TTL expiry or capacity eviction.
func CommitClaudeDiagnostics(key string, sequence uint64, messageID string) {
key = strings.TrimSpace(key)
messageID = strings.TrimSpace(messageID)
if key == "" || sequence == 0 || messageID == "" {
return
}
now := time.Now()
claudeDiagnosticsState.Lock()
defer claudeDiagnosticsState.Unlock()
entry, ok := claudeDiagnosticsState.entries[key]
if !ok || sequence < entry.minimumSequence || sequence < entry.committedSequence {
return
}
claudeDiagnosticsState.nextAccess++
entry.previousMessageID = messageID
entry.committedSequence = sequence
entry.lastAccess = claudeDiagnosticsState.nextAccess
entry.expiresAt = now.Add(claudeDiagnosticsTTL)
claudeDiagnosticsState.entries[key] = entry
}
func cleanupClaudeDiagnosticsLocked(now time.Time) {
if !claudeDiagnosticsState.lastCleanup.IsZero() && now.Sub(claudeDiagnosticsState.lastCleanup) < claudeDiagnosticsCleanupPeriod {
return
}
for key, entry := range claudeDiagnosticsState.entries {
if !entry.expiresAt.IsZero() && now.After(entry.expiresAt) {
delete(claudeDiagnosticsState.entries, key)
}
}
claudeDiagnosticsState.lastCleanup = now
}
func evictClaudeDiagnosticsLocked() {
if len(claudeDiagnosticsState.entries) < claudeDiagnosticsMaxEntries {
return
}
type candidate struct {
key string
lastAccess uint64
}
candidates := make([]candidate, 0, len(claudeDiagnosticsState.entries))
for key, entry := range claudeDiagnosticsState.entries {
candidates = append(candidates, candidate{key: key, lastAccess: entry.lastAccess})
}
sort.Slice(candidates, func(i, j int) bool {
return candidates[i].lastAccess < candidates[j].lastAccess
})
count := min(claudeDiagnosticsEvictBatchSize, len(candidates))
for _, candidate := range candidates[:count] {
delete(claudeDiagnosticsState.entries, candidate.key)
}
}
func resetClaudeDiagnosticsForTest() {
claudeDiagnosticsState.Lock()
defer claudeDiagnosticsState.Unlock()
claudeDiagnosticsState.entries = make(map[string]claudeDiagnosticsEntry)
claudeDiagnosticsState.lastCleanup = time.Time{}
claudeDiagnosticsState.nextSequence = 0
claudeDiagnosticsState.nextAccess = 0
}

View file

@ -0,0 +1,102 @@
package helps
import (
"fmt"
"testing"
"time"
)
func TestClaudeDiagnosticsTracksCompletedMessagePerCredentialSession(t *testing.T) {
resetClaudeDiagnosticsForTest()
defer resetClaudeDiagnosticsForTest()
key, sequence, previous := BeginClaudeDiagnostics("credential-a", "session-a")
if key == "" || sequence != 1 || previous != "" {
t.Fatalf("first begin = %q/%d/%q, want key/1/empty", key, sequence, previous)
}
CommitClaudeDiagnostics(key, sequence, "msg_first")
_, secondSequence, previous := BeginClaudeDiagnostics("credential-a", "session-a")
if secondSequence != 2 || previous != "msg_first" {
t.Fatalf("second begin = %d/%q, want 2/msg_first", secondSequence, previous)
}
_, _, otherSession := BeginClaudeDiagnostics("credential-a", "session-b")
_, _, otherCredential := BeginClaudeDiagnostics("credential-b", "session-a")
if otherSession != "" || otherCredential != "" {
t.Fatalf("diagnostics leaked across identity: session=%q credential=%q", otherSession, otherCredential)
}
}
func TestClaudeDiagnosticsRejectsExpiredGenerationCommit(t *testing.T) {
resetClaudeDiagnosticsForTest()
defer resetClaudeDiagnosticsForTest()
key, expiredSequence, _ := BeginClaudeDiagnostics("credential", "session")
claudeDiagnosticsState.Lock()
entry := claudeDiagnosticsState.entries[key]
entry.expiresAt = time.Now().Add(-time.Second)
claudeDiagnosticsState.entries[key] = entry
claudeDiagnosticsState.Unlock()
newKey, currentSequence, previous := BeginClaudeDiagnostics("credential", "session")
if newKey != key || currentSequence <= expiredSequence || previous != "" {
t.Fatalf("new generation = %q/%d/%q, want same key/new sequence/empty", newKey, currentSequence, previous)
}
CommitClaudeDiagnostics(newKey, currentSequence, "msg_current")
CommitClaudeDiagnostics(key, expiredSequence, "msg_expired")
_, _, previous = BeginClaudeDiagnostics("credential", "session")
if previous != "msg_current" {
t.Fatalf("previous message = %q, want current generation", previous)
}
}
func TestClaudeDiagnosticsCacheEvictsOldestEntriesWithinCapacity(t *testing.T) {
resetClaudeDiagnosticsForTest()
defer resetClaudeDiagnosticsForTest()
firstKey, firstSequence, _ := BeginClaudeDiagnostics("credential", "session-0")
var newestKey string
for index := 1; index <= claudeDiagnosticsMaxEntries; index++ {
newestKey, _, _ = BeginClaudeDiagnostics("credential", fmt.Sprintf("session-%d", index))
}
claudeDiagnosticsState.Lock()
entryCount := len(claudeDiagnosticsState.entries)
_, firstFound := claudeDiagnosticsState.entries[firstKey]
_, newestFound := claudeDiagnosticsState.entries[newestKey]
claudeDiagnosticsState.Unlock()
if entryCount > claudeDiagnosticsMaxEntries {
t.Fatalf("cache entries = %d, want at most %d", entryCount, claudeDiagnosticsMaxEntries)
}
if firstFound {
t.Fatal("oldest diagnostics entry was not evicted")
}
if !newestFound {
t.Fatal("newest diagnostics entry was evicted")
}
newKey, newSequence, _ := BeginClaudeDiagnostics("credential", "session-0")
if newKey != firstKey || newSequence <= firstSequence {
t.Fatalf("recreated generation = %q/%d, want same key after sequence %d", newKey, newSequence, firstSequence)
}
CommitClaudeDiagnostics(newKey, newSequence, "msg_recreated")
CommitClaudeDiagnostics(firstKey, firstSequence, "msg_evicted")
_, _, previous := BeginClaudeDiagnostics("credential", "session-0")
if previous != "msg_recreated" {
t.Fatalf("previous message = %q, want recreated generation", previous)
}
}
func TestClaudeDiagnosticsRejectsLateOlderCommit(t *testing.T) {
resetClaudeDiagnosticsForTest()
defer resetClaudeDiagnosticsForTest()
key, first, _ := BeginClaudeDiagnostics("credential", "session")
_, second, _ := BeginClaudeDiagnostics("credential", "session")
CommitClaudeDiagnostics(key, second, "msg_newer")
CommitClaudeDiagnostics(key, first, "msg_older")
_, _, previous := BeginClaudeDiagnostics("credential", "session")
if previous != "msg_newer" {
t.Fatalf("previous message = %q, want newer completed generation", previous)
}
}

View file

@ -0,0 +1,387 @@
package helps
import (
"bytes"
"context"
"encoding/json"
"fmt"
"strings"
"sync"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
"github.com/tiktoken-go/tokenizer"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
)
var (
claudeInputTokenizerOnce sync.Once
claudeInputTokenizerCodec tokenizer.Codec
claudeInputTokenizerErr error
)
// ClaudeInputTokenState tracks the one-time input token update for a translated Claude stream.
type ClaudeInputTokenState struct {
upstreamFormat sdktranslator.Format
responseFormat sdktranslator.Format
originalRequest []byte
codec tokenizer.Codec
handled bool
}
// NewClaudeInputTokenState creates request-scoped state for translated Claude input token usage.
func NewClaudeInputTokenState(sourceFormat, upstreamFormat, responseFormat sdktranslator.Format, originalRequest []byte) *ClaudeInputTokenState {
enabled := sourceFormat == sdktranslator.FormatClaude &&
upstreamFormat != sdktranslator.FormatClaude &&
responseFormat == sdktranslator.FormatClaude
return &ClaudeInputTokenState{
upstreamFormat: upstreamFormat,
responseFormat: responseFormat,
originalRequest: originalRequest,
handled: !enabled,
}
}
// TranslateStreamWithClaudeInputTokens translates a stream chunk and estimates Claude message_start input usage once.
func TranslateStreamWithClaudeInputTokens(
ctx context.Context,
upstreamFormat, responseFormat sdktranslator.Format,
model string,
originalRequestRawJSON, requestRawJSON, rawJSON []byte,
param *any,
state *ClaudeInputTokenState,
) [][]byte {
chunks := sdktranslator.TranslateStream(
ctx,
upstreamFormat,
responseFormat,
model,
originalRequestRawJSON,
requestRawJSON,
rawJSON,
param,
)
if responseFormat == sdktranslator.FormatOpenAIResponse {
for i, chunk := range chunks {
chunks[i] = EnsureResponsesUsageDetails(chunk)
}
}
if state == nil {
return chunks
}
return state.apply(ctx, chunks)
}
func claudeInputTokenizer() (tokenizer.Codec, error) {
claudeInputTokenizerOnce.Do(func() {
claudeInputTokenizerCodec, claudeInputTokenizerErr = tokenizer.Get(tokenizer.O200kBase)
})
return claudeInputTokenizerCodec, claudeInputTokenizerErr
}
// CountClaudeInputTokens estimates tokens for a Claude request with the O200kBase tokenizer.
func CountClaudeInputTokens(payload []byte) (int64, error) {
enc, err := claudeInputTokenizer()
if err != nil {
return 0, fmt.Errorf("initialize O200kBase tokenizer: %w", err)
}
count, err := countClaudeInputTokens(enc, payload)
if err != nil {
return 0, fmt.Errorf("count Claude input tokens: %w", err)
}
return count, nil
}
func countClaudeInputTokens(enc tokenizer.Codec, payload []byte) (int64, error) {
if enc == nil {
return 0, fmt.Errorf("encoder is nil")
}
segments, err := collectClaudeInputTokenSegments(payload)
if err != nil {
return 0, err
}
if len(segments) == 0 {
return 0, nil
}
count, err := enc.Count(strings.Join(segments, "\n"))
if err != nil {
return 0, err
}
return int64(count), nil
}
func collectClaudeInputTokenSegments(payload []byte) ([]string, error) {
if len(bytes.TrimSpace(payload)) == 0 {
return nil, nil
}
if !gjson.ValidBytes(payload) {
return nil, fmt.Errorf("invalid Claude request JSON")
}
root := gjson.ParseBytes(payload)
segments := make([]string, 0, 32)
collectClaudeSystemTokenSegments(root.Get("system"), &segments)
collectClaudeMessageTokenSegments(root.Get("messages"), &segments)
collectClaudeToolTokenSegments(root.Get("tools"), &segments)
collectClaudeToolChoiceTokenSegments(root.Get("tool_choice"), &segments)
return segments, nil
}
func collectClaudeSystemTokenSegments(system gjson.Result, segments *[]string) {
if system.Type == gjson.String {
appendClaudeTokenString(segments, system.String())
return
}
if !system.IsArray() {
return
}
system.ForEach(func(_, part gjson.Result) bool {
if part.Type == gjson.String {
appendClaudeTokenString(segments, part.String())
} else if part.Get("type").String() == "text" {
appendClaudeTokenString(segments, part.Get("text").String())
}
return true
})
}
func collectClaudeMessageTokenSegments(messages gjson.Result, segments *[]string) {
if !messages.IsArray() {
return
}
messages.ForEach(func(_, message gjson.Result) bool {
appendClaudeTokenString(segments, message.Get("role").String())
collectClaudeContentTokenSegments(message.Get("content"), segments)
return true
})
}
func collectClaudeContentTokenSegments(content gjson.Result, segments *[]string) {
if !content.Exists() {
return
}
if content.Type == gjson.String {
appendClaudeTokenString(segments, content.String())
return
}
if content.IsArray() {
content.ForEach(func(_, part gjson.Result) bool {
collectClaudeContentTokenSegments(part, segments)
return true
})
return
}
if !content.IsObject() {
return
}
switch content.Get("type").String() {
case "text":
appendClaudeTokenString(segments, content.Get("text").String())
case "thinking":
appendClaudeTokenString(segments, content.Get("thinking").String())
case "document":
collectClaudeDocumentTokenSegments(content, segments)
case "tool_use", "server_tool_use", "mcp_tool_use":
appendClaudeTokenString(segments, content.Get("id").String())
appendClaudeTokenString(segments, content.Get("name").String())
appendClaudeTokenJSON(segments, content.Get("input"))
case "tool_result", "mcp_tool_result", "web_search_tool_result", "web_fetch_tool_result", "code_execution_tool_result", "bash_code_execution_tool_result", "text_editor_code_execution_tool_result":
appendClaudeTokenString(segments, content.Get("tool_use_id").String())
appendClaudeTokenString(segments, content.Get("tool_call_id").String())
collectClaudeContentTokenSegments(content.Get("content"), segments)
case "web_search_result", "search_result":
if source := content.Get("source"); source.Type == gjson.String {
appendClaudeTokenString(segments, source.String())
}
appendClaudeTokenString(segments, content.Get("title").String())
appendClaudeTokenString(segments, content.Get("url").String())
appendClaudeTokenString(segments, content.Get("page_age").String())
collectClaudeContentTokenSegments(content.Get("content"), segments)
case "web_fetch_result":
appendClaudeTokenString(segments, content.Get("url").String())
appendClaudeTokenString(segments, content.Get("retrieved_at").String())
collectClaudeContentTokenSegments(content.Get("content"), segments)
case "code_execution_result", "bash_code_execution_result", "text_editor_code_execution_result":
appendClaudeTokenString(segments, content.Get("stdout").String())
appendClaudeTokenString(segments, content.Get("stderr").String())
appendClaudeTokenString(segments, content.Get("return_code").String())
collectClaudeContentTokenSegments(content.Get("content"), segments)
collectClaudeContentTokenSegments(content.Get("output"), segments)
case "tool_reference":
appendClaudeTokenString(segments, content.Get("tool_name").String())
case "image", "input_audio", "audio", "video", "redacted_thinking":
return
case "":
appendClaudeTokenJSON(segments, content)
default:
appendClaudeTokenString(segments, content.Get("text").String())
}
}
func collectClaudeDocumentTokenSegments(document gjson.Result, segments *[]string) {
source := document.Get("source")
if source.Get("type").String() != "text" {
return
}
appendClaudeTokenString(segments, document.Get("title").String())
appendClaudeTokenString(segments, document.Get("context").String())
appendClaudeTokenString(segments, source.Get("data").String())
appendClaudeTokenString(segments, source.Get("content").String())
}
func collectClaudeToolTokenSegments(tools gjson.Result, segments *[]string) {
if !tools.IsArray() {
return
}
tools.ForEach(func(_, tool gjson.Result) bool {
appendClaudeTokenString(segments, tool.Get("type").String())
appendClaudeTokenString(segments, tool.Get("name").String())
appendClaudeTokenString(segments, tool.Get("description").String())
appendClaudeTokenJSON(segments, tool.Get("input_schema"))
return true
})
}
func collectClaudeToolChoiceTokenSegments(toolChoice gjson.Result, segments *[]string) {
if !toolChoice.Exists() {
return
}
if toolChoice.Type == gjson.String {
appendClaudeTokenString(segments, toolChoice.String())
return
}
appendClaudeTokenString(segments, toolChoice.Get("type").String())
appendClaudeTokenString(segments, toolChoice.Get("name").String())
}
func appendClaudeTokenString(segments *[]string, value string) {
if segments == nil {
return
}
if trimmed := strings.TrimSpace(value); trimmed != "" {
*segments = append(*segments, trimmed)
}
}
func appendClaudeTokenJSON(segments *[]string, value gjson.Result) {
if !value.Exists() {
return
}
if value.Type == gjson.String {
appendClaudeTokenString(segments, value.String())
return
}
raw := strings.TrimSpace(value.Raw)
if raw == "" {
return
}
var compact bytes.Buffer
if err := json.Compact(&compact, []byte(raw)); err == nil {
appendClaudeTokenString(segments, compact.String())
return
}
appendClaudeTokenString(segments, raw)
}
func (state *ClaudeInputTokenState) apply(ctx context.Context, chunks [][]byte) [][]byte {
if state == nil || state.handled {
return chunks
}
for i := range chunks {
updated, found := state.applyChunk(ctx, chunks[i])
if !found {
continue
}
state.handled = true
chunks[i] = updated
break
}
return chunks
}
func (state *ClaudeInputTokenState) applyChunk(ctx context.Context, chunk []byte) ([]byte, bool) {
for lineStart := 0; lineStart < len(chunk); {
lineEnd := bytes.IndexByte(chunk[lineStart:], '\n')
if lineEnd < 0 {
lineEnd = len(chunk)
} else {
lineEnd += lineStart
}
contentEnd := lineEnd
if contentEnd > lineStart && chunk[contentEnd-1] == '\r' {
contentEnd--
}
line := chunk[lineStart:contentEnd]
trimmedLeft := bytes.TrimLeft(line, " \t")
if bytes.HasPrefix(trimmedLeft, []byte("data:")) {
payloadOffset := len(line) - len(trimmedLeft) + len("data:")
for payloadOffset < len(line) && (line[payloadOffset] == ' ' || line[payloadOffset] == '\t') {
payloadOffset++
}
payloadEnd := len(line)
for payloadEnd > payloadOffset && (line[payloadEnd-1] == ' ' || line[payloadEnd-1] == '\t') {
payloadEnd--
}
payload := line[payloadOffset:payloadEnd]
if gjson.GetBytes(payload, "type").String() == "message_start" {
inputTokens := gjson.GetBytes(payload, "message.usage.input_tokens")
if inputTokens.Exists() && inputTokens.Int() != 0 {
return chunk, true
}
count, err := state.estimate()
if err != nil {
state.logEstimateError(ctx, err)
return chunk, true
}
if count == 0 {
return chunk, true
}
updatedPayload, errSet := sjson.SetBytes(payload, "message.usage.input_tokens", count)
if errSet != nil {
state.logEstimateError(ctx, fmt.Errorf("set message_start usage: %w", errSet))
return chunk, true
}
payloadStart := lineStart + payloadOffset
payloadStop := lineStart + payloadEnd
updated := make([]byte, 0, len(chunk)+len(updatedPayload)-len(payload))
updated = append(updated, chunk[:payloadStart]...)
updated = append(updated, updatedPayload...)
updated = append(updated, chunk[payloadStop:]...)
return updated, true
}
}
if lineEnd == len(chunk) {
break
}
lineStart = lineEnd + 1
}
return chunk, false
}
func (state *ClaudeInputTokenState) estimate() (int64, error) {
enc := state.codec
if enc == nil {
var err error
enc, err = claudeInputTokenizer()
if err != nil {
return 0, fmt.Errorf("initialize O200kBase tokenizer: %w", err)
}
}
count, err := countClaudeInputTokens(enc, state.originalRequest)
if err != nil {
return 0, fmt.Errorf("count Claude input tokens: %w", err)
}
return count, nil
}
func (state *ClaudeInputTokenState) logEstimateError(ctx context.Context, err error) {
LogWithRequestID(ctx).WithFields(log.Fields{
"upstream_format": state.upstreamFormat.String(),
"response_format": state.responseFormat.String(),
}).WithError(err).Warn("failed to estimate Claude input tokens")
}

View file

@ -0,0 +1,443 @@
package helps
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"strings"
"sync"
"testing"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tiktoken-go/tokenizer"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
)
type failingClaudeInputCodec struct{}
func (failingClaudeInputCodec) GetName() string {
return "failing"
}
func (failingClaudeInputCodec) Count(string) (int, error) {
return 0, errors.New("count failed")
}
func (failingClaudeInputCodec) Encode(string) ([]uint, []string, error) {
return nil, nil, errors.New("encode failed")
}
func (failingClaudeInputCodec) Decode([]uint) (string, error) {
return "", errors.New("decode failed")
}
func TestCollectClaudeInputTokenSegments(t *testing.T) {
payload := []byte(`{
"model":"claude-test",
"system":[
{"type":"text","text":"Follow repository rules.","cache_control":{"type":"ephemeral"}},
{"type":"image","source":{"type":"base64","media_type":"image/png","data":"ignored-system-image"}}
],
"messages":[
{"role":"user","content":[
{"type":"text","text":"Review the implementation."},
{"type":"document","source":{"type":"text","data":"Reference document text."}},
{"type":"image","source":{"type":"base64","media_type":"image/png","data":"ignored-image"}}
]},
{"role":"assistant","content":[
{"type":"thinking","thinking":"Inspect the relevant files.","signature":"ignored-signature"},
{"type":"tool_use","id":"toolu_1","name":"read_file","input":{"path":"main.go"}}
]},
{"role":"user","content":[
{"type":"tool_result","tool_use_id":"toolu_1","content":[
{"type":"text","text":"package main"},
{"type":"image","source":{"type":"base64","data":"ignored-tool-image"}}
]}
]}
],
"tools":[{
"name":"read_file",
"description":"Reads a repository file.",
"input_schema":{"type":"object","properties":{"path":{"type":"string"}}},
"cache_control":{"type":"ephemeral"}
}],
"tool_choice":{"type":"tool","name":"read_file"},
"metadata":{"user_id":"ignored-metadata"},
"max_tokens":4096,
"stream":true
}`)
got, err := collectClaudeInputTokenSegments(payload)
if err != nil {
t.Fatalf("collectClaudeInputTokenSegments() error = %v", err)
}
want := []string{
"Follow repository rules.",
"user",
"Review the implementation.",
"Reference document text.",
"assistant",
"Inspect the relevant files.",
"toolu_1",
"read_file",
`{"path":"main.go"}`,
"user",
"toolu_1",
"package main",
"read_file",
"Reads a repository file.",
`{"type":"object","properties":{"path":{"type":"string"}}}`,
"tool",
"read_file",
}
if fmt.Sprint(got) != fmt.Sprint(want) {
t.Fatalf("segments = %#v, want %#v", got, want)
}
}
func TestCollectClaudeInputTokenSegmentsIncludesKnownToolResults(t *testing.T) {
payload := []byte(`{
"messages":[{"role":"user","content":[
{"type":"web_search_tool_result","tool_use_id":"ws_tool_1","content":[
{"type":"web_search_result","source":"Search source","title":"Search result title","url":"https://search.example/result","page_age":"1 day","encrypted_content":"ignored-secret"}
]},
{"type":"web_fetch_tool_result","tool_use_id":"fetch_tool_1","content":{
"type":"web_fetch_result","url":"https://docs.example/page","retrieved_at":"2026-07-22T00:00:00Z","content":{
"type":"document","title":"Fetched document","source":{"type":"text","data":"Fetched body"}
}
}},
{"type":"bash_code_execution_tool_result","tool_use_id":"bash_tool_1","content":{
"type":"bash_code_execution_result","stdout":"command output","stderr":"command error","return_code":1,
"content":[{"type":"text","text":"additional output"}]
}},
{"type":"tool_result","tool_use_id":"toolu_1","content":[
{"type":"tool_reference","tool_name":"proxy_mcp__nia__manage_resource"}
]}
]}]
}`)
segments, err := collectClaudeInputTokenSegments(payload)
if err != nil {
t.Fatalf("collectClaudeInputTokenSegments() error = %v", err)
}
joined := "\n" + strings.Join(segments, "\n") + "\n"
for _, want := range []string{
"ws_tool_1",
"Search source",
"Search result title",
"https://search.example/result",
"1 day",
"fetch_tool_1",
"https://docs.example/page",
"2026-07-22T00:00:00Z",
"Fetched document",
"Fetched body",
"bash_tool_1",
"command output",
"command error",
"1",
"additional output",
"toolu_1",
"proxy_mcp__nia__manage_resource",
} {
if !strings.Contains(joined, "\n"+want+"\n") {
t.Errorf("segments do not contain %q: %#v", want, segments)
}
}
if strings.Contains(joined, "ignored-secret") {
t.Fatalf("segments contain encrypted content: %#v", segments)
}
}
func TestCountClaudeInputTokensExcludesMultimediaAndControlFields(t *testing.T) {
enc, err := tokenizer.Get(tokenizer.O200kBase)
if err != nil {
t.Fatalf("tokenizer.Get() error = %v", err)
}
base := []byte(`{
"system":"System text.",
"messages":[{"role":"user","content":[{"type":"text","text":"User text."}]}],
"tools":[{"name":"lookup","description":"Looks up data.","input_schema":{"type":"object"}}]
}`)
withExcludedFields := []byte(`{
"model":"claude-test",
"system":"System text.",
"messages":[{"role":"user","content":[
{"type":"text","text":"User text."},
{"type":"image","source":{"type":"base64","media_type":"image/png","data":"very-large-image-data"}},
{"type":"input_audio","source":{"type":"base64","data":"very-large-audio-data"}},
{"type":"video","source":{"type":"url","url":"https://example.com/video.mp4"}},
{"type":"document","source":{"type":"base64","media_type":"application/pdf","data":"very-large-pdf-data"}}
]}],
"tools":[{"name":"lookup","description":"Looks up data.","input_schema":{"type":"object"},"cache_control":{"type":"ephemeral"}}],
"metadata":{"large_wrapper":"ignored"},
"max_tokens":8192,
"temperature":0.8,
"top_p":0.9,
"thinking":{"type":"enabled","budget_tokens":4096},
"stream":true
}`)
baseCount, errBase := countClaudeInputTokens(enc, base)
if errBase != nil {
t.Fatalf("countClaudeInputTokens(base) error = %v", errBase)
}
excludedCount, errExcluded := countClaudeInputTokens(enc, withExcludedFields)
if errExcluded != nil {
t.Fatalf("countClaudeInputTokens(withExcludedFields) error = %v", errExcluded)
}
if excludedCount != baseCount {
t.Fatalf("count with excluded fields = %d, want %d", excludedCount, baseCount)
}
}
func TestTranslateStreamWithClaudeInputTokensPatchesMessageStartOnce(t *testing.T) {
upstreamFormat := sdktranslator.Format("claude-input-token-test-upstream")
sdktranslator.Register(sdktranslator.FormatClaude, upstreamFormat, nil, sdktranslator.ResponseTransform{
Stream: func(_ context.Context, _ string, _, _, rawJSON []byte, _ *any) [][]byte {
return [][]byte{rawJSON}
},
})
originalRequest := []byte(`{"system":"System text.","messages":[{"role":"user","content":"Hello."}]}`)
state := NewClaudeInputTokenState(sdktranslator.FormatClaude, upstreamFormat, sdktranslator.FormatClaude, originalRequest)
var param any
combined := []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0,\"output_tokens\":0}}}\n\n" +
"event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0}\n\n")
got := TranslateStreamWithClaudeInputTokens(
context.Background(),
upstreamFormat,
sdktranslator.FormatClaude,
"claude-test",
originalRequest,
nil,
combined,
&param,
state,
)
if tokens := messageStartInputTokens(got); tokens <= 0 {
t.Fatalf("message_start input_tokens = %d, want positive estimate; output = %q", tokens, joinClaudeInputChunks(got))
}
if !state.handled {
t.Fatal("state.handled = false, want true after message_start")
}
if !strings.Contains(joinClaudeInputChunks(got), `"type":"content_block_start"`) {
t.Fatalf("combined non-target event was not preserved: %q", joinClaudeInputChunks(got))
}
secondStart := []byte(`event: message_start
data: {"type":"message_start","message":{"usage":{"input_tokens":0}}}
`)
gotSecond := TranslateStreamWithClaudeInputTokens(
context.Background(),
upstreamFormat,
sdktranslator.FormatClaude,
"claude-test",
originalRequest,
nil,
secondStart,
&param,
state,
)
if tokens := messageStartInputTokens(gotSecond); tokens != 0 {
t.Fatalf("second message_start input_tokens = %d, want 0 after state handled", tokens)
}
}
func TestClaudeInputTokenStatePreservesCRLFAndNonTargetEvents(t *testing.T) {
originalRequest := []byte(`{"messages":[{"role":"user","content":"Hello."}]}`)
state := NewClaudeInputTokenState(sdktranslator.FormatClaude, sdktranslator.FormatOpenAI, sdktranslator.FormatClaude, originalRequest)
chunk := []byte("event: message_start\r\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0,\"output_tokens\":0}}} \r\n\r\n" +
"event: ping\r\ndata: {\"type\":\"ping\",\"value\":\"keep\"}\r\n\r\n")
got := state.apply(context.Background(), [][]byte{chunk})
tokens := messageStartInputTokens(got)
if tokens <= 0 {
t.Fatalf("input_tokens = %d, want positive estimate", tokens)
}
want := fmt.Sprintf("event: message_start\r\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":%d,\"output_tokens\":0}}} \r\n\r\n"+
"event: ping\r\ndata: {\"type\":\"ping\",\"value\":\"keep\"}\r\n\r\n", tokens)
if joined := joinClaudeInputChunks(got); joined != want {
t.Fatalf("output bytes changed unexpectedly:\n got: %q\nwant: %q", joined, want)
}
}
func TestClaudeInputTokenStatePatchesMissingAndPreservesNonZero(t *testing.T) {
originalRequest := []byte(`{"messages":[{"role":"user","content":"Hello."}]}`)
t.Run("missing", func(t *testing.T) {
state := NewClaudeInputTokenState(sdktranslator.FormatClaude, sdktranslator.FormatOpenAI, sdktranslator.FormatClaude, originalRequest)
chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"output_tokens\":0}}}\n\n")}
got := state.apply(context.Background(), chunks)
if tokens := messageStartInputTokens(got); tokens <= 0 {
t.Fatalf("input_tokens = %d, want positive estimate", tokens)
}
})
t.Run("non-zero", func(t *testing.T) {
state := NewClaudeInputTokenState(sdktranslator.FormatClaude, sdktranslator.FormatOpenAI, sdktranslator.FormatClaude, []byte(`not valid json`))
chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":73}}}\n\n")}
got := state.apply(context.Background(), chunks)
if tokens := messageStartInputTokens(got); tokens != 73 {
t.Fatalf("input_tokens = %d, want preserved value 73", tokens)
}
if !state.handled {
t.Fatal("state.handled = false, want true")
}
})
}
func TestClaudeInputTokenStateSkipsUnsupportedFlows(t *testing.T) {
originalRequest := []byte(`{"messages":[{"role":"user","content":"Hello."}]}`)
testCases := []struct {
name string
sourceFormat sdktranslator.Format
upstreamFormat sdktranslator.Format
responseFormat sdktranslator.Format
}{
{name: "non-Claude source", sourceFormat: sdktranslator.FormatOpenAI, upstreamFormat: sdktranslator.FormatGemini, responseFormat: sdktranslator.FormatClaude},
{name: "Claude passthrough", sourceFormat: sdktranslator.FormatClaude, upstreamFormat: sdktranslator.FormatClaude, responseFormat: sdktranslator.FormatClaude},
{name: "non-Claude response", sourceFormat: sdktranslator.FormatClaude, upstreamFormat: sdktranslator.FormatOpenAI, responseFormat: sdktranslator.FormatOpenAI},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
state := NewClaudeInputTokenState(tc.sourceFormat, tc.upstreamFormat, tc.responseFormat, originalRequest)
chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0}}}\n\n")}
got := state.apply(context.Background(), chunks)
if tokens := messageStartInputTokens(got); tokens != 0 {
t.Fatalf("input_tokens = %d, want unchanged 0", tokens)
}
if !state.handled {
t.Fatal("state.handled = false, want disabled flow handled at initialization")
}
})
}
}
func TestClaudeInputTokenStateCountErrorKeepsZero(t *testing.T) {
originalLogOutput := log.StandardLogger().Out
log.SetOutput(io.Discard)
defer log.SetOutput(originalLogOutput)
state := NewClaudeInputTokenState(
sdktranslator.FormatClaude,
sdktranslator.FormatOpenAI,
sdktranslator.FormatClaude,
[]byte(`{"messages":[{"role":"user","content":"Hello."}]}`),
)
state.codec = failingClaudeInputCodec{}
chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0}}}\n\n")}
got := state.apply(context.Background(), chunks)
if tokens := messageStartInputTokens(got); tokens != 0 {
t.Fatalf("input_tokens = %d, want fallback 0", tokens)
}
if !state.handled {
t.Fatal("state.handled = false, want true after failed estimate")
}
}
func TestClaudeInputTokenStateInvalidJSONKeepsZeroWithoutLoggingRequest(t *testing.T) {
originalLogOutput := log.StandardLogger().Out
var logOutput bytes.Buffer
log.SetOutput(&logOutput)
defer log.SetOutput(originalLogOutput)
const sensitiveRequest = `{"messages":["sensitive-original-request"`
state := NewClaudeInputTokenState(
sdktranslator.FormatClaude,
sdktranslator.FormatOpenAI,
sdktranslator.FormatClaude,
[]byte(sensitiveRequest),
)
chunks := [][]byte{[]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":0}}}\n\n")}
got := state.apply(context.Background(), chunks)
if tokens := messageStartInputTokens(got); tokens != 0 {
t.Fatalf("input_tokens = %d, want fallback 0", tokens)
}
if !state.handled {
t.Fatal("state.handled = false, want true after invalid JSON")
}
if !strings.Contains(logOutput.String(), "failed to estimate Claude input tokens") {
t.Fatalf("warning not logged: %q", logOutput.String())
}
if strings.Contains(logOutput.String(), "sensitive-original-request") {
t.Fatalf("warning leaked original request: %q", logOutput.String())
}
}
func TestClaudeInputTokenizerConcurrentCount(t *testing.T) {
first, errFirst := claudeInputTokenizer()
if errFirst != nil {
t.Fatalf("claudeInputTokenizer() error = %v", errFirst)
}
second, errSecond := claudeInputTokenizer()
if errSecond != nil {
t.Fatalf("claudeInputTokenizer() second error = %v", errSecond)
}
if first != second {
t.Fatal("claudeInputTokenizer() returned different codec instances")
}
const workers = 32
const iterations = 50
var wg sync.WaitGroup
errs := make(chan error, workers)
for worker := 0; worker < workers; worker++ {
worker := worker
wg.Add(1)
go func() {
defer wg.Done()
for iteration := 0; iteration < iterations; iteration++ {
payload := []byte(fmt.Sprintf(`{"messages":[{"role":"user","content":"worker %d iteration %d 你好"}]}`, worker, iteration))
count, errCount := countClaudeInputTokens(first, payload)
if errCount != nil {
errs <- errCount
return
}
if count <= 0 {
errs <- fmt.Errorf("non-positive count: %d", count)
return
}
}
}()
}
wg.Wait()
close(errs)
for err := range errs {
t.Error(err)
}
}
func messageStartInputTokens(chunks [][]byte) int64 {
for _, chunk := range chunks {
for _, line := range strings.Split(string(chunk), "\n") {
trimmed := strings.TrimSpace(line)
if !strings.HasPrefix(trimmed, "data:") {
continue
}
payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
if gjson.Get(payload, "type").String() == "message_start" {
return gjson.Get(payload, "message.usage.input_tokens").Int()
}
}
}
return 0
}
func joinClaudeInputChunks(chunks [][]byte) string {
var builder strings.Builder
for _, chunk := range chunks {
builder.Write(chunk)
}
return builder.String()
}

View file

@ -0,0 +1,144 @@
package helps
import (
"crypto/hmac"
"crypto/sha256"
"encoding/binary"
"strings"
log "github.com/sirupsen/logrus"
)
// IsClaudeMCPToolName reports whether name follows Claude Code's MCP tool
// convention and contains only characters accepted by Anthropic tool names.
func IsClaudeMCPToolName(name string) bool {
if len(name) == 0 || len(name) > 64 || !strings.HasPrefix(name, "mcp__") {
return false
}
rest := strings.TrimPrefix(name, "mcp__")
separator := strings.Index(rest, "__")
if separator <= 0 || separator+2 >= len(rest) {
return false
}
for _, char := range name {
if (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') ||
(char >= '0' && char <= '9') || char == '_' || char == '-' {
continue
}
return false
}
return true
}
// ClaudeMCPAliasWordCount is the BIP-39 English dictionary size used for the
// virtual server pair and the one-word tool ID.
func ClaudeMCPAliasWordCount() int {
return len(claudeMCPAliasEnglishWords)
}
// ClaudeMCPToolAlias derives a Claude Code-style MCP tool name. Aliases from
// one caller share a virtual server component. The tool component combines a
// stable keyed ID with a truncated semantic suffix so the model can distinguish
// tools by name while the request-local symbol table restores the exact original.
// A higher attempt linearly probes the next word when a collision must be avoided.
// Server and tool IDs use BIP-39 English words so weak models are less likely
// to drift high-entropy Base32 fragments.
func ClaudeMCPToolAlias(secret, original string, attempt uint32) string {
toolDigest := claudeMCPAliasDigest(secret, "tool", original)
return claudeMCPAliasFor(
claudeMCPAliasServerComponent(secret),
claudeMCPAliasWord(toolDigest[:], 0, attempt),
original,
)
}
// AllocateClaudeMCPToolAlias picks an alias that is not already reserved.
// Attempts are capped at the wordlist size so names that sanitize to the same
// suffix cannot spin forever. ok is false only when every one-word tool ID for
// this semantic is already reserved.
func AllocateClaudeMCPToolAlias(secret, original string, reserved map[string]bool) (string, bool) {
words := claudeMCPAliasEnglishWords
totalWords := len(words)
if totalWords == 0 {
log.Error("claude oauth mcp alias: embedded BIP-39 wordlist is empty, tool aliasing is disabled")
return "", false
}
server := claudeMCPAliasServerComponent(secret)
toolDigest := claudeMCPAliasDigest(secret, "tool", original)
baseIndex := int(binary.BigEndian.Uint16(toolDigest[0:2])) % totalWords
for attempt := 0; attempt < totalWords; attempt++ {
alias := claudeMCPAliasFor(server, words[(baseIndex+attempt)%totalWords], original)
if reserved != nil && reserved[alias] {
continue
}
return alias, true
}
return "", false
}
// claudeMCPAliasFor assembles the final alias for one server/tool word pair.
// Both the single-shot and the allocating entry point must build names here so
// the two cannot drift apart.
func claudeMCPAliasFor(server, toolID, original string) string {
prefix := "mcp__" + server + "__" + toolID + "_"
maxSemanticLen := 64 - len(prefix)
if maxSemanticLen < 1 {
maxSemanticLen = 1
}
return prefix + claudeMCPToolSemanticSuffix(original, maxSemanticLen)
}
// claudeMCPAliasServerComponent derives the caller-stable two-word virtual
// server shared by every alias generated for one credential.
func claudeMCPAliasServerComponent(secret string) string {
serverDigest := claudeMCPAliasDigest(secret, "server", "")
return claudeMCPAliasWord(serverDigest[:], 0, 0) + "_" + claudeMCPAliasWord(serverDigest[:], 2, 0)
}
func claudeMCPAliasWord(digest []byte, offset int, attempt uint32) string {
words := claudeMCPAliasEnglishWords
if len(words) == 0 || offset < 0 || offset+2 > len(digest) {
return "tool"
}
base := int(binary.BigEndian.Uint16(digest[offset : offset+2]))
return words[(base+int(attempt))%len(words)]
}
func claudeMCPToolSemanticSuffix(original string, maxLength int) string {
var semantic strings.Builder
semantic.Grow(min(len(original), maxLength))
pendingSeparator := false
for _, char := range original {
valid := (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') ||
(char >= '0' && char <= '9') || char == '_' || char == '-'
if !valid {
pendingSeparator = semantic.Len() > 0
continue
}
if pendingSeparator && semantic.Len()+1 < maxLength {
semantic.WriteByte('_')
}
pendingSeparator = false
if semantic.Len() >= maxLength {
break
}
semantic.WriteRune(char)
}
result := strings.Trim(semantic.String(), "_-")
if result == "" {
return "tool"
}
return result
}
func claudeMCPAliasDigest(secret, purpose, original string) [sha256.Size]byte {
mac := hmac.New(sha256.New, []byte(secret))
_, _ = mac.Write([]byte("cpa-claude-mcp-alias-v2\x00"))
_, _ = mac.Write([]byte(purpose))
_, _ = mac.Write([]byte{0})
_, _ = mac.Write([]byte(original))
var digest [sha256.Size]byte
copy(digest[:], mac.Sum(nil))
return digest
}

View file

@ -0,0 +1,272 @@
package helps
import (
"fmt"
"regexp"
"strings"
"testing"
)
func TestIsClaudeMCPToolName(t *testing.T) {
for _, name := range []string{
"mcp__context7__query-docs",
"mcp__amber_cedar__quiet_harbor",
"mcp__server__tool__variant",
} {
if !IsClaudeMCPToolName(name) {
t.Fatalf("IsClaudeMCPToolName(%q) = false, want true", name)
}
}
for _, name := range []string{
"context7__query-docs",
"mcp____query-docs",
"mcp__context7__",
"mcp__context7__query.docs",
"mcp__context7__" + strings.Repeat("x", 64),
} {
if IsClaudeMCPToolName(name) {
t.Fatalf("IsClaudeMCPToolName(%q) = true, want false", name)
}
}
}
func TestClaudeMCPToolAlias(t *testing.T) {
first := ClaudeMCPToolAlias("credential-secret", "search_web", 0)
if second := ClaudeMCPToolAlias("credential-secret", "search_web", 0); second != first {
t.Fatalf("alias is not deterministic: %q != %q", first, second)
}
caseDistinct := ClaudeMCPToolAlias("credential-secret", "Search_Web", 0)
if first == caseDistinct {
t.Fatalf("case-distinct names produced the same initial alias: %q", first)
}
retry := ClaudeMCPToolAlias("credential-secret", "search_web", 1)
if first == retry {
t.Fatalf("collision retry did not change alias: %q", first)
}
if !IsClaudeMCPToolName(first) {
t.Fatalf("generated alias %q is not a valid MCP tool name", first)
}
if !strings.HasSuffix(first, "_search_web") {
t.Fatalf("generated alias %q does not preserve the semantic suffix", first)
}
if matched, _ := regexp.MatchString(`^mcp__[a-z]+_[a-z]+__[a-z]+_search_web$`, first); !matched {
t.Fatalf("generated alias %q does not contain word-based IDs plus semantics", first)
}
assertClaudeMCPAliasWords(t, first)
server := strings.Split(first, "__")[1]
if got := strings.Split(caseDistinct, "__")[1]; got != server {
t.Fatalf("case-distinct tool server = %q, want shared caller server %q", got, server)
}
if got := strings.Split(retry, "__")[1]; got != server {
t.Fatalf("retry server = %q, want shared caller server %q", got, server)
}
if got := strings.Split(ClaudeMCPToolAlias("other-caller", "search_web", 0), "__")[1]; got == server {
t.Fatalf("different caller unexpectedly shared server %q", server)
}
}
func TestClaudeMCPToolAlias_SemanticSuffixIsSafeAndBounded(t *testing.T) {
tests := []struct {
name string
original string
wantSuffix string
}{
{name: "invalid separators", original: "browser.open URL", wantSuffix: "_browser_open_URL"},
{name: "unicode mixed", original: "search.网页/tool with spaces", wantSuffix: "_search_tool_with_spaces"},
{name: "unicode only", original: "搜索网页", wantSuffix: "_tool"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
alias := ClaudeMCPToolAlias("credential-secret", tt.original, 0)
if !IsClaudeMCPToolName(alias) {
t.Fatalf("generated alias %q is not a valid MCP tool name", alias)
}
if len(alias) > 64 {
t.Fatalf("generated alias length = %d, want <= 64: %q", len(alias), alias)
}
if !strings.HasSuffix(alias, tt.wantSuffix) {
t.Fatalf("generated alias %q does not end in %q", alias, tt.wantSuffix)
}
})
}
const original = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
alias := ClaudeMCPToolAlias("credential-secret", original, 0)
underscore := strings.LastIndex(alias, "_")
if underscore < 0 {
t.Fatalf("generated alias %q has no semantic separator", alias)
}
prefixLen := underscore + 1
wantSemanticLen := 64 - prefixLen
if wantSemanticLen < 1 {
wantSemanticLen = 1
}
if got := alias[prefixLen:]; got != strings.Repeat("a", wantSemanticLen) {
t.Fatalf("semantic suffix = %q, want %d a's", got, wantSemanticLen)
}
if len(alias) != 64 {
t.Fatalf("generated alias length = %d, want 64: %q", len(alias), alias)
}
}
func TestClaudeMCPToolAlias_Strict64CharLimitUnderAllWordCombinations(t *testing.T) {
for i := 0; i < ClaudeMCPAliasWordCount(); i++ {
secret := fmt.Sprintf("test-secret-%d", i)
original := strings.Repeat(fmt.Sprintf("tool_%d_long_name_", i), 50)
alias := ClaudeMCPToolAlias(secret, original, uint32(i))
if len(alias) > 64 {
t.Fatalf("alias length %d exceeds Anthropic 64-char limit: %q", len(alias), alias)
}
if !IsClaudeMCPToolName(alias) {
t.Fatalf("alias %q is not a valid MCP tool name", alias)
}
assertClaudeMCPAliasWords(t, alias)
}
}
func TestAllocateClaudeMCPToolAlias_StopsWhenAttemptsExhausted(t *testing.T) {
const secret = "exhaust-space"
const original = "tool.name"
reserved := make(map[string]bool, ClaudeMCPAliasWordCount())
for attempt := 0; attempt < ClaudeMCPAliasWordCount(); attempt++ {
reserved[ClaudeMCPToolAlias(secret, original, uint32(attempt))] = true
}
if _, ok := AllocateClaudeMCPToolAlias(secret, original, reserved); ok {
t.Fatal("allocate succeeded after every attempt alias was reserved")
}
if alias, ok := AllocateClaudeMCPToolAlias(secret, original, nil); !ok || alias == "" {
t.Fatal("allocate failed with an empty reserved set")
}
}
func TestClaudeMCPToolAlias_ProbesAllWordsWithoutDuplicates(t *testing.T) {
const secret = "test-secret"
const original = "tool.name"
totalWords := ClaudeMCPAliasWordCount()
seen := make(map[string]bool, totalWords)
for attempt := 0; attempt < totalWords; attempt++ {
alias := ClaudeMCPToolAlias(secret, original, uint32(attempt))
parts := strings.Split(alias, "__")
toolID, _, _ := strings.Cut(parts[2], "_")
if seen[toolID] {
t.Fatalf("attempt %d generated duplicate toolID %q", attempt, toolID)
}
seen[toolID] = true
}
if len(seen) != totalWords {
t.Fatalf("covered %d words in %d attempts, want 100%% (%d words)", len(seen), totalWords, totalWords)
}
}
func TestAllocateClaudeMCPToolAlias_AllocatesEveryDistinctWord(t *testing.T) {
const secret = "allocate-full-space"
const original = "tool.name"
totalWords := ClaudeMCPAliasWordCount()
reserved := make(map[string]bool, totalWords)
for i := 0; i < totalWords; i++ {
alias, ok := AllocateClaudeMCPToolAlias(secret, original, reserved)
if !ok {
t.Fatalf("failed to allocate at step %d with %d words reserved", i, len(reserved))
}
if reserved[alias] {
t.Fatalf("allocated duplicate alias %q at step %d", alias, i)
}
reserved[alias] = true
}
if len(reserved) != totalWords {
t.Fatalf("reserved count = %d, want %d", len(reserved), totalWords)
}
if _, ok := AllocateClaudeMCPToolAlias(secret, original, reserved); ok {
t.Fatal("allocate succeeded when all 2048 words are reserved")
}
}
func BenchmarkAllocateClaudeMCPToolAlias_Collision(b *testing.B) {
const secret = "test-secret"
const original = "tool.name"
reserved := make(map[string]bool)
for attempt := 0; attempt < 100; attempt++ {
reserved[ClaudeMCPToolAlias(secret, original, uint32(attempt))] = true
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _ = AllocateClaudeMCPToolAlias(secret, original, reserved)
}
}
func assertClaudeMCPAliasWords(t *testing.T, alias string) {
t.Helper()
parts := strings.Split(alias, "__")
if len(parts) != 3 {
t.Fatalf("alias %q does not have mcp/server/tool parts", alias)
}
serverWords := strings.Split(parts[1], "_")
if len(serverWords) != 2 {
t.Fatalf("alias %q server %q is not two BIP-39 words", alias, parts[1])
}
toolID, _, ok := strings.Cut(parts[2], "_")
if !ok {
t.Fatalf("alias %q tool component %q has no semantic suffix", alias, parts[2])
}
allowed := make(map[string]struct{}, len(claudeMCPAliasEnglishWords))
for _, word := range claudeMCPAliasEnglishWords {
allowed[word] = struct{}{}
}
for _, word := range append(append([]string{}, serverWords...), toolID) {
if _, exists := allowed[word]; !exists {
t.Fatalf("alias %q uses non-BIP39 word %q", alias, word)
}
}
}
func TestClaudeMCPAliasWordlistIntegrity(t *testing.T) {
// The wordlist is embedded, so a truncated or reordered file would silently
// disable aliasing (AllocateClaudeMCPToolAlias returns false for every tool)
// instead of failing loudly. Pin the exact BIP-39 English dictionary.
if got := ClaudeMCPAliasWordCount(); got != 2048 {
t.Fatalf("wordlist size = %d, want the 2048-word BIP-39 English dictionary", got)
}
if got := claudeMCPAliasEnglishWords[0]; got != "abandon" {
t.Fatalf("first word = %q, want %q", got, "abandon")
}
if got := claudeMCPAliasEnglishWords[2047]; got != "zoo" {
t.Fatalf("last word = %q, want %q", got, "zoo")
}
seen := make(map[string]struct{}, len(claudeMCPAliasEnglishWords))
for _, word := range claudeMCPAliasEnglishWords {
if _, duplicate := seen[word]; duplicate {
t.Fatalf("duplicate word %q would shrink the usable alias space", word)
}
seen[word] = struct{}{}
if word == "" || len(word) > 8 {
t.Fatalf("word %q is outside the 1..8 character budget assumed by the 64-char alias limit", word)
}
for _, char := range word {
if char < 'a' || char > 'z' {
t.Fatalf("word %q contains a non-lowercase-ASCII rune %q", word, char)
}
}
}
}
func TestAllocateClaudeMCPToolAliasMatchesSingleShotConstruction(t *testing.T) {
// Both entry points must build identical names; the exhaustion tests above
// use ClaudeMCPToolAlias to seed the reserved set, so any drift between the
// two would make them silently stop testing the production path.
const secret = "shared-construction"
for _, original := range []string{"Bash", "read_file", strings.Repeat("long_tool_name_", 9)} {
reserved := make(map[string]bool, ClaudeMCPAliasWordCount())
for attempt := 0; attempt < ClaudeMCPAliasWordCount(); attempt++ {
allocated, ok := AllocateClaudeMCPToolAlias(secret, original, reserved)
if !ok {
t.Fatalf("original %q: allocation exhausted at attempt %d", original, attempt)
}
if want := ClaudeMCPToolAlias(secret, original, uint32(attempt)); allocated != want {
t.Fatalf("original %q attempt %d: allocated %q, single-shot %q", original, attempt, allocated, want)
}
reserved[allocated] = true
}
}
}

View file

@ -0,0 +1,12 @@
package helps
import (
_ "embed"
"strings"
)
//go:embed claude_bip39_words.txt
var rawBIP39EnglishWords string
// claudeMCPAliasEnglishWords contains the standard BIP-39 English wordlist (2048 words).
var claudeMCPAliasEnglishWords = strings.Fields(rawBIP39EnglishWords)

View file

@ -0,0 +1,249 @@
package helps
import (
cryptorand "crypto/rand"
"math/big"
"net/http"
"strconv"
"strings"
"time"
log "github.com/sirupsen/logrus"
)
const (
defaultClaudeRateLimitFuzzMinSeconds = 1
defaultClaudeRateLimitFuzzMaxSeconds = 30
)
// ClaudeHeadersIndicateUnifiedRateLimitRejection reports whether response headers explicitly
// declare an Anthropic shared 5h or 7d rate-limit rejection. A Fable-only 7d_oi rejection
// remains model-scoped when both shared windows are explicitly allowed.
func ClaudeHeadersIndicateUnifiedRateLimitRejection(headers http.Header) bool {
if headers == nil {
return false
}
unifiedStatus := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-Status")))
status5h := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-5h-Status")))
if status5h == "rejected" {
return true
}
status7d := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-7d-Status")))
if status7d == "rejected" {
return true
}
if unifiedStatus != "rejected" {
return false
}
status7dOI := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-7d_oi-Status")))
return !isFableOnlyRejection(status5h, status7d, status7dOI)
}
func isFableOnlyRejection(status5h, status7d, status7dOI string) bool {
return status5h == "allowed" && status7d == "allowed" && status7dOI == "rejected"
}
// ParseClaudeRateLimitReset inspects Anthropic response headers for shared and Fable-specific
// unified rate-limit and standard Retry-After reset information, returning the conservative cooldown
// duration including a bounded non-negative random grace period.
// If no valid future reset information is present, it returns nil.
func ParseClaudeRateLimitReset(headers http.Header, now time.Time) *time.Duration {
return parseClaudeRateLimitResetWithFuzz(headers, now, defaultClaudeRateLimitFuzzMinSeconds, defaultClaudeRateLimitFuzzMaxSeconds)
}
func parseClaudeRateLimitResetWithFuzz(headers http.Header, now time.Time, minFuzzSec, maxFuzzSec int) *time.Duration {
if headers == nil {
return nil
}
unifiedStatus := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-Status")))
status5h := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-5h-Status")))
status7d := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-7d-Status")))
status7dOI := strings.ToLower(strings.TrimSpace(getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-7d_oi-Status")))
fableOnlyRejection := isFableOnlyRejection(status5h, status7d, status7dOI)
var candidateDeadlines []time.Time
var rejectedWindows []string
if unifiedStatus == "rejected" {
rejectedWindows = append(rejectedWindows, "unified")
}
if status5h == "rejected" {
rejectedWindows = append(rejectedWindows, "5h")
}
if status7d == "rejected" {
rejectedWindows = append(rejectedWindows, "7d")
}
if status7dOI == "rejected" {
rejectedWindows = append(rejectedWindows, "7d_oi")
}
// 1. Retry-After header
if rawRetryAfter := getHeaderCaseInsensitive(headers, "Retry-After"); rawRetryAfter != "" {
if !containsString(rejectedWindows, "retry-after") {
rejectedWindows = append(rejectedWindows, "retry-after")
}
if t, ok := parseRetryAfterHeader(rawRetryAfter, now); ok && t.After(now) {
candidateDeadlines = append(candidateDeadlines, t)
}
}
// 2. 5-hour window reset (only when rejected)
if status5h == "rejected" {
if raw := getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-5h-Reset"); raw != "" {
if t, ok := parseUnixOrTimestamp(raw); ok && t.After(now) {
candidateDeadlines = append(candidateDeadlines, t)
}
}
}
// 3. 7-day window reset (only when rejected)
if status7d == "rejected" {
if raw := getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-7d-Reset"); raw != "" {
if t, ok := parseUnixOrTimestamp(raw); ok && t.After(now) {
candidateDeadlines = append(candidateDeadlines, t)
}
}
}
// 4. Fable-specific 7-day window reset (only when rejected and not a Fable-only rejection)
if status7dOI == "rejected" && !fableOnlyRejection {
if raw := getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-7d_oi-Reset"); raw != "" {
if t, ok := parseUnixOrTimestamp(raw); ok && t.After(now) {
candidateDeadlines = append(candidateDeadlines, t)
}
}
}
// 5. Unified reset header:
unifiedRejected := !fableOnlyRejection && (unifiedStatus == "rejected" || status5h == "rejected" || status7d == "rejected" || status7dOI == "rejected" ||
(unifiedStatus == "" && status5h != "allowed" && status7d != "allowed"))
if unifiedRejected {
if raw := getHeaderCaseInsensitive(headers, "Anthropic-Ratelimit-Unified-Reset"); raw != "" {
if !containsString(rejectedWindows, "unified") {
rejectedWindows = append(rejectedWindows, "unified")
}
if t, ok := parseUnixOrTimestamp(raw); ok && t.After(now) {
candidateDeadlines = append(candidateDeadlines, t)
}
}
}
if len(candidateDeadlines) == 0 {
if len(rejectedWindows) > 0 {
log.WithFields(log.Fields{
"rejected_windows": strings.Join(rejectedWindows, ","),
"status": "fallback_exponential_backoff",
}).Info("Anthropic rate limit window rejected; falling back to generic exponential backoff")
}
return nil
}
// Pick the latest applicable deadline across rejected windows
var latestDeadline time.Time
for _, deadline := range candidateDeadlines {
if deadline.After(latestDeadline) {
latestDeadline = deadline
}
}
if latestDeadline.IsZero() || !latestDeadline.After(now) {
if len(rejectedWindows) > 0 {
log.WithFields(log.Fields{
"rejected_windows": strings.Join(rejectedWindows, ","),
"status": "fallback_exponential_backoff",
}).Info("Anthropic rate limit window rejected; falling back to generic exponential backoff")
}
return nil
}
baseDuration := latestDeadline.Sub(now)
fuzz := randomClaudeFuzzDuration(minFuzzSec, maxFuzzSec)
effectiveDuration := baseDuration + fuzz
log.WithFields(log.Fields{
"rejected_windows": strings.Join(rejectedWindows, ","),
"effective_cooldown": effectiveDuration.String(),
"base_cooldown": baseDuration.String(),
"fuzz": fuzz.String(),
"deadline": latestDeadline.Format(time.RFC3339),
}).Info("parsed Anthropic rate limit reset headers")
return &effectiveDuration
}
func containsString(list []string, target string) bool {
for _, item := range list {
if item == target {
return true
}
}
return false
}
func getHeaderCaseInsensitive(h http.Header, target string) string {
if h == nil {
return ""
}
if val := h.Get(target); val != "" {
return val
}
for k, v := range h {
if strings.EqualFold(k, target) && len(v) > 0 {
return v[0]
}
}
return ""
}
func parseUnixOrTimestamp(raw string) (time.Time, bool) {
raw = strings.TrimSpace(raw)
if raw == "" {
return time.Time{}, false
}
if sec, err := strconv.ParseFloat(raw, 64); err == nil && sec > 0 {
secInt := int64(sec)
nsec := int64((sec - float64(secInt)) * 1e9)
return time.Unix(secInt, nsec), true
}
if t, err := time.Parse(time.RFC3339, raw); err == nil {
return t, true
}
if t, err := http.ParseTime(raw); err == nil {
return t, true
}
return time.Time{}, false
}
func parseRetryAfterHeader(raw string, now time.Time) (time.Time, bool) {
raw = strings.TrimSpace(raw)
if raw == "" {
return time.Time{}, false
}
if sec, err := strconv.ParseFloat(raw, 64); err == nil && sec > 0 {
d := time.Duration(sec * float64(time.Second))
return now.Add(d), true
}
if t, err := http.ParseTime(raw); err == nil {
return t, true
}
if t, err := time.Parse(time.RFC3339, raw); err == nil {
return t, true
}
return time.Time{}, false
}
func randomClaudeFuzzDuration(minSec, maxSec int) time.Duration {
if maxSec <= minSec {
if minSec < 0 {
return 0
}
return time.Duration(minSec) * time.Second
}
nBig, err := cryptorand.Int(cryptorand.Reader, big.NewInt(int64(maxSec-minSec+1)))
if err != nil {
return time.Duration(minSec) * time.Second
}
return time.Duration(minSec+int(nBig.Int64())) * time.Second
}

View file

@ -0,0 +1,193 @@
package helps
import (
"net/http"
"strconv"
"testing"
"time"
)
func TestParseClaudeRateLimitReset_AllCases(t *testing.T) {
now := time.Now()
t.Run("nil headers returns nil", func(t *testing.T) {
if got := ParseClaudeRateLimitReset(nil, now); got != nil {
t.Fatalf("expected nil, got %v", got)
}
})
t.Run("empty headers returns nil", func(t *testing.T) {
h := make(http.Header)
if got := ParseClaudeRateLimitReset(h, now); got != nil {
t.Fatalf("expected nil, got %v", got)
}
})
t.Run("retry-after only seconds", func(t *testing.T) {
h := make(http.Header)
h.Set("Retry-After", "60")
got := parseClaudeRateLimitResetWithFuzz(h, now, 0, 0)
if got == nil {
t.Fatal("expected non-nil RetryAfter")
}
if *got != 60*time.Second {
t.Fatalf("expected 60s, got %v", *got)
}
})
t.Run("retry-after HTTP date", func(t *testing.T) {
h := make(http.Header)
futureTime := now.Add(90 * time.Second).UTC().Truncate(time.Second)
h.Set("Retry-After", futureTime.Format(http.TimeFormat))
got := parseClaudeRateLimitResetWithFuzz(h, now, 0, 0)
if got == nil {
t.Fatal("expected non-nil RetryAfter")
}
if *got < 89*time.Second || *got > 91*time.Second {
t.Fatalf("expected ~90s, got %v", *got)
}
})
t.Run("5h rejected and 7d allowed with unified reset", func(t *testing.T) {
h := make(http.Header)
// Missing Anthropic-Ratelimit-Unified-Status, 5h is rejected, 7d is allowed
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(now.Add(5*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-7d-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-Reset", strconv.FormatInt(now.Add(5*time.Hour).Unix(), 10))
got := parseClaudeRateLimitResetWithFuzz(h, now, 0, 0)
if got == nil {
t.Fatal("expected non-nil RetryAfter")
}
if *got < 5*time.Hour-5*time.Second || *got > 5*time.Hour+5*time.Second {
t.Fatalf("expected ~5h, got %v", *got)
}
})
t.Run("7d rejected and 5h allowed", func(t *testing.T) {
h := make(http.Header)
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(now.Add(5*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
got := parseClaudeRateLimitResetWithFuzz(h, now, 0, 0)
if got == nil {
t.Fatal("expected non-nil RetryAfter")
}
if *got < 7*24*time.Hour-5*time.Second || *got > 7*24*time.Hour+5*time.Second {
t.Fatalf("expected ~7d, got %v", *got)
}
})
t.Run("both 5h and 7d rejected chooses longest", func(t *testing.T) {
h := make(http.Header)
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(now.Add(5*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
got := parseClaudeRateLimitResetWithFuzz(h, now, 0, 0)
if got == nil {
t.Fatal("expected non-nil RetryAfter")
}
if *got < 7*24*time.Hour-5*time.Second || *got > 7*24*time.Hour+5*time.Second {
t.Fatalf("expected ~7d, got %v", *got)
}
})
t.Run("all allowed returns nil", func(t *testing.T) {
h := make(http.Header)
h.Set("Anthropic-Ratelimit-Unified-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(now.Add(5*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-7d-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
got := ParseClaudeRateLimitReset(h, now)
if got != nil {
t.Fatalf("expected nil for allowed status, got %v", got)
}
})
t.Run("fable-only rejection with 7d_oi reset and retry-after uses retry-after only", func(t *testing.T) {
h := make(http.Header)
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-7d-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-7d_oi-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-7d_oi-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
h.Set("Retry-After", "60")
got := parseClaudeRateLimitResetWithFuzz(h, now, 0, 0)
if got == nil {
t.Fatal("expected non-nil RetryAfter")
}
if *got != 60*time.Second {
t.Fatalf("expected 60s from Retry-After, got %v", *got)
}
})
t.Run("fable-only rejection with 7d_oi reset only returns nil for exponential backoff", func(t *testing.T) {
h := make(http.Header)
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-7d-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-7d_oi-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-7d_oi-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
got := ParseClaudeRateLimitReset(h, now)
if got != nil {
t.Fatalf("expected nil for fable-only rejection without retry-after, got %v", *got)
}
})
t.Run("non-fable combined rejection with 7d_oi reset keeps longer duration", func(t *testing.T) {
h := make(http.Header)
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(now.Add(5*time.Hour).Unix(), 10))
h.Set("Anthropic-Ratelimit-Unified-7d-Status", "allowed")
h.Set("Anthropic-Ratelimit-Unified-7d_oi-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-7d_oi-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10))
got := parseClaudeRateLimitResetWithFuzz(h, now, 0, 0)
if got == nil {
t.Fatal("expected non-nil RetryAfter")
}
if *got < 7*24*time.Hour-5*time.Second || *got > 7*24*time.Hour+5*time.Second {
t.Fatalf("expected ~7d, got %v", *got)
}
})
t.Run("past timestamp returns nil", func(t *testing.T) {
h := make(http.Header)
h.Set("Anthropic-Ratelimit-Unified-5h-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(now.Add(-5*time.Hour).Unix(), 10))
got := ParseClaudeRateLimitReset(h, now)
if got != nil {
t.Fatalf("expected nil for past reset, got %v", got)
}
})
t.Run("fuzz is bounded and non-negative", func(t *testing.T) {
h := make(http.Header)
h.Set("Retry-After", "100")
for i := 0; i < 50; i++ {
got := ParseClaudeRateLimitReset(h, now)
if got == nil {
t.Fatal("expected non-nil")
}
diff := *got - 100*time.Second
if diff < 1*time.Second || diff > 30*time.Second {
t.Fatalf("fuzz %v out of bounds [1s, 30s]", diff)
}
}
})
}

View file

@ -0,0 +1,17 @@
package helps
import (
"net/url"
"strings"
)
// IsAnthropicUpstreamURL reports whether a resolved request targets Anthropic's
// first-party API origin. Claude-specific body, header, HTTP, and TLS behavior
// must all use this gate so they cannot drift onto custom ports or userinfo URLs.
func IsAnthropicUpstreamURL(u *url.URL) bool {
if u == nil || u.User != nil || !strings.EqualFold(u.Scheme, "https") || !strings.EqualFold(u.Hostname(), "api.anthropic.com") {
return false
}
port := u.Port()
return port == "" || port == "443"
}

View file

@ -0,0 +1,38 @@
package helps
import (
"net/url"
"testing"
)
func TestIsAnthropicUpstreamURL(t *testing.T) {
testCases := []struct {
name string
targetURL string
want bool
}{
{name: "default HTTPS port", targetURL: "https://api.anthropic.com/v1/messages", want: true},
{name: "explicit HTTPS port", targetURL: "https://api.anthropic.com:443/v1/messages", want: true},
{name: "case insensitive host", targetURL: "https://API.ANTHROPIC.COM/v1/messages", want: true},
{name: "HTTP", targetURL: "http://api.anthropic.com/v1/messages", want: false},
{name: "custom port", targetURL: "https://api.anthropic.com:8443/v1/messages", want: false},
{name: "userinfo", targetURL: "https://caller@api.anthropic.com/v1/messages", want: false},
{name: "lookalike host", targetURL: "https://api.anthropic.com.example/v1/messages", want: false},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
parsed, errParse := url.Parse(testCase.targetURL)
if errParse != nil {
t.Fatal(errParse)
}
if got := IsAnthropicUpstreamURL(parsed); got != testCase.want {
t.Fatalf("IsAnthropicUpstreamURL(%q) = %t, want %t", testCase.targetURL, got, testCase.want)
}
})
}
if IsAnthropicUpstreamURL(nil) {
t.Fatal("IsAnthropicUpstreamURL(nil) = true")
}
}

View file

@ -0,0 +1,214 @@
package helps
import (
"regexp"
"sort"
"strings"
"unicode/utf8"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// zeroWidthSpace is the Unicode zero-width space character used for obfuscation.
const zeroWidthSpace = "\u200B"
// SensitiveWordMatcher holds the compiled regex for matching sensitive words.
type SensitiveWordMatcher struct {
regex *regexp.Regexp
}
// BuildSensitiveWordMatcher compiles a regex from the word list.
// Words are sorted by length (longest first) for proper matching.
func BuildSensitiveWordMatcher(words []string) *SensitiveWordMatcher {
if len(words) == 0 {
return nil
}
// Filter and normalize words
var validWords []string
for _, w := range words {
w = strings.TrimSpace(w)
if utf8.RuneCountInString(w) >= 2 && !strings.Contains(w, zeroWidthSpace) {
validWords = append(validWords, w)
}
}
if len(validWords) == 0 {
return nil
}
// Sort by length (longest first) for proper matching
sort.Slice(validWords, func(i, j int) bool {
return len(validWords[i]) > len(validWords[j])
})
// Escape and join
escaped := make([]string, len(validWords))
for i, w := range validWords {
escaped[i] = regexp.QuoteMeta(w)
}
pattern := "(?i)" + strings.Join(escaped, "|")
re, err := regexp.Compile(pattern)
if err != nil {
return nil
}
return &SensitiveWordMatcher{regex: re}
}
// obfuscateWord inserts a zero-width space after the first grapheme.
func obfuscateWord(word string) string {
if strings.Contains(word, zeroWidthSpace) {
return word
}
// Get first rune
r, size := utf8.DecodeRuneInString(word)
if r == utf8.RuneError || size >= len(word) {
return word
}
return string(r) + zeroWidthSpace + word[size:]
}
// obfuscateText replaces all sensitive words in the text.
func (m *SensitiveWordMatcher) obfuscateText(text string) string {
if m == nil || m.regex == nil {
return text
}
return m.regex.ReplaceAllStringFunc(text, obfuscateWord)
}
// ObfuscateSensitiveWords processes the payload and obfuscates sensitive words
// in system blocks and message content.
func ObfuscateSensitiveWords(payload []byte, matcher *SensitiveWordMatcher) []byte {
if matcher == nil || matcher.regex == nil {
return payload
}
// Obfuscate in system blocks
payload = obfuscateSystemBlocks(payload, matcher)
// Obfuscate in messages
payload = obfuscateMessages(payload, matcher)
return payload
}
// ObfuscateSensitiveWordsInSystemInstruction obfuscates sensitive words in an Antigravity system instruction.
func ObfuscateSensitiveWordsInSystemInstruction(payload []byte, matcher *SensitiveWordMatcher) []byte {
if matcher == nil || matcher.regex == nil {
return payload
}
for _, path := range []string{"request.systemInstruction", "request.system_instruction"} {
instruction := gjson.GetBytes(payload, path)
if !instruction.Exists() {
continue
}
if instruction.Type == gjson.String {
text := instruction.String()
if obfuscated := matcher.obfuscateText(text); obfuscated != text {
payload, _ = sjson.SetBytes(payload, path, obfuscated)
}
continue
}
parts := instruction.Get("parts")
if !parts.IsArray() {
continue
}
parts.ForEach(func(key, part gjson.Result) bool {
if part.Get("text").Type != gjson.String {
return true
}
text := part.Get("text").String()
if obfuscated := matcher.obfuscateText(text); obfuscated != text {
payload, _ = sjson.SetBytes(payload, path+".parts."+key.String()+".text", obfuscated)
}
return true
})
}
return payload
}
// obfuscateSystemBlocks obfuscates sensitive words in system blocks.
func obfuscateSystemBlocks(payload []byte, matcher *SensitiveWordMatcher) []byte {
system := gjson.GetBytes(payload, "system")
if !system.Exists() {
return payload
}
if system.IsArray() {
modified := false
system.ForEach(func(key, value gjson.Result) bool {
if value.Get("type").String() == "text" {
text := value.Get("text").String()
obfuscated := matcher.obfuscateText(text)
if obfuscated != text {
path := "system." + key.String() + ".text"
payload, _ = sjson.SetBytes(payload, path, obfuscated)
modified = true
}
}
return true
})
if modified {
return payload
}
} else if system.Type == gjson.String {
text := system.String()
obfuscated := matcher.obfuscateText(text)
if obfuscated != text {
payload, _ = sjson.SetBytes(payload, "system", obfuscated)
}
}
return payload
}
// obfuscateMessages obfuscates sensitive words in message content.
func obfuscateMessages(payload []byte, matcher *SensitiveWordMatcher) []byte {
messages := gjson.GetBytes(payload, "messages")
if !messages.Exists() || !messages.IsArray() {
return payload
}
messages.ForEach(func(msgKey, msg gjson.Result) bool {
content := msg.Get("content")
if !content.Exists() {
return true
}
msgPath := "messages." + msgKey.String()
if content.Type == gjson.String {
// Simple string content
text := content.String()
obfuscated := matcher.obfuscateText(text)
if obfuscated != text {
payload, _ = sjson.SetBytes(payload, msgPath+".content", obfuscated)
}
} else if content.IsArray() {
// Array of content blocks
content.ForEach(func(blockKey, block gjson.Result) bool {
if block.Get("type").String() == "text" {
text := block.Get("text").String()
obfuscated := matcher.obfuscateText(text)
if obfuscated != text {
path := msgPath + ".content." + blockKey.String() + ".text"
payload, _ = sjson.SetBytes(payload, path, obfuscated)
}
}
return true
})
}
return true
})
return payload
}

View file

@ -0,0 +1,69 @@
package helps
import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"regexp"
"github.com/google/uuid"
)
var claudeMetadataDeviceIDPattern = regexp.MustCompile(`^[a-f0-9]{64}$`)
type claudeMetadataUserID struct {
DeviceID string `json:"device_id"`
AccountUUID string `json:"account_uuid"`
SessionID string `json:"session_id"`
}
// generateFakeUserID generates metadata.user_id in the JSON string format used
// by Claude Code 2.1.78 and newer.
func generateFakeUserID() string {
return generateFakeUserIDWithSessionID(uuid.New().String())
}
func generateFakeUserIDWithSessionID(sessionID string) string {
if _, errParse := uuid.Parse(sessionID); errParse != nil {
sessionID = uuid.New().String()
}
hexBytes := make([]byte, 32)
_, _ = rand.Read(hexBytes)
value, _ := json.Marshal(claudeMetadataUserID{
DeviceID: hex.EncodeToString(hexBytes),
AccountUUID: "",
SessionID: sessionID,
})
return string(value)
}
// isValidUserID checks the Claude Code 2.1.220 metadata.user_id shape.
func isValidUserID(userID string) bool {
var value claudeMetadataUserID
if errUnmarshal := json.Unmarshal([]byte(userID), &value); errUnmarshal != nil {
return false
}
if !claudeMetadataDeviceIDPattern.MatchString(value.DeviceID) {
return false
}
if _, errParse := uuid.Parse(value.SessionID); errParse != nil {
return false
}
if value.AccountUUID == "" {
return true
}
_, errParse := uuid.Parse(value.AccountUUID)
return errParse == nil
}
func GenerateFakeUserID() string {
return generateFakeUserID()
}
func GenerateFakeUserIDWithSessionID(sessionID string) string {
return generateFakeUserIDWithSessionID(sessionID)
}
func IsValidUserID(userID string) bool {
return isValidUserID(userID)
}

View file

@ -0,0 +1,194 @@
package helps
import (
"crypto/sha256"
"encoding/hex"
"strconv"
"strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
const (
codexInputItemIDLimit = 64
codexMessageItemIDPrefix = "msg"
codexReasoningItemIDPrefix = "rs"
codexFunctionCallItemIDPrefix = "fc"
codexCustomToolCallItemIDPrefix = "ctc"
codexCustomToolCallOutputItemIDPrefix = "ctco"
codexInputItemIDOccupied uint8 = 1 << 0
codexInputItemIDPreserved uint8 = 1 << 1
)
// SanitizeCodexInputItemIDs normalizes supported input item IDs for Codex, removes encrypted
// reasoning items whose IDs exceed the Codex limit, and deterministically shortens
// other overlong input item IDs.
func SanitizeCodexInputItemIDs(body []byte) []byte {
input := util.GetGJSONBytesNoCopy(body, "input")
if !input.IsArray() {
return body
}
items := input.Array()
idStates := make(map[string]uint8, len(items))
for _, item := range items {
if shouldDropCodexEncryptedReasoningItem(item) {
continue
}
itemID := item.Get("id")
if itemID.Type != gjson.String {
continue
}
originalID := itemID.String()
id := normalizeCodexInputItemID(item, originalID)
state := idStates[id]
if id == originalID {
state |= codexInputItemIDPreserved
}
if len([]rune(id)) <= codexInputItemIDLimit {
state |= codexInputItemIDOccupied
}
if state != 0 {
idStates[id] = state
}
}
var mapped map[string]string
var collisionMapped map[string]string
rebuilt := make([]string, 0, len(items))
changed := false
for _, item := range items {
if shouldDropCodexEncryptedReasoningItem(item) {
changed = true
continue
}
raw := item.Raw
itemID := item.Get("id")
if itemID.Type == gjson.String {
originalID := itemID.String()
id := normalizeCodexInputItemID(item, originalID)
if id != originalID && idStates[id]&codexInputItemIDPreserved != 0 {
collisionID, ok := collisionMapped[id]
if !ok {
for attempt := 0; ; attempt++ {
collisionID = codexInputItemIDWithHashSuffix(id, attempt)
if idStates[collisionID]&codexInputItemIDOccupied != 0 {
continue
}
if collisionMapped == nil {
collisionMapped = make(map[string]string)
}
collisionMapped[id] = collisionID
idStates[collisionID] |= codexInputItemIDOccupied
break
}
}
id = collisionID
}
if len([]rune(id)) > codexInputItemIDLimit {
shortened, ok := mapped[id]
if !ok {
shortened = shortenCodexInputItemID(id)
for attempt := 1; ; attempt++ {
if idStates[shortened]&codexInputItemIDOccupied == 0 {
break
}
shortened = shortenCodexInputItemIDWithAttempt(id, attempt)
}
if mapped == nil {
mapped = make(map[string]string)
}
mapped[id] = shortened
idStates[shortened] |= codexInputItemIDOccupied
}
id = shortened
}
if id != originalID {
next, errSet := sjson.SetBytes([]byte(raw), "id", id)
if errSet == nil {
raw = string(next)
changed = true
}
}
}
rebuilt = append(rebuilt, raw)
}
if !changed {
return body
}
updated, errSet := sjson.SetRawBytes(body, "input", []byte("["+strings.Join(rebuilt, ",")+"]"))
if errSet != nil {
return body
}
return updated
}
func normalizeCodexInputItemID(item gjson.Result, id string) string {
var prefix string
switch item.Get("type").String() {
case "message":
prefix = codexMessageItemIDPrefix
case "reasoning":
prefix = codexReasoningItemIDPrefix
case "function_call":
prefix = codexFunctionCallItemIDPrefix
case "custom_tool_call":
prefix = codexCustomToolCallItemIDPrefix
case "custom_tool_call_output":
prefix = codexCustomToolCallOutputItemIDPrefix
default:
return id
}
if id == "" || strings.HasPrefix(id, prefix) {
return id
}
return prefix + "_" + id
}
func shouldDropCodexEncryptedReasoningItem(item gjson.Result) bool {
if item.Get("type").String() != "reasoning" {
return false
}
itemID := item.Get("id")
if itemID.Type != gjson.String || len([]rune(itemID.String())) <= codexInputItemIDLimit {
return false
}
encryptedContent := item.Get("encrypted_content")
return encryptedContent.Type == gjson.String && encryptedContent.String() != ""
}
func shortenCodexInputItemID(id string) string {
return shortenCodexInputItemIDWithAttempt(id, 0)
}
func shortenCodexInputItemIDWithAttempt(id string, attempt int) string {
runes := []rune(id)
if len(runes) <= codexInputItemIDLimit {
return id
}
return codexInputItemIDWithHashSuffixRunes(id, runes, attempt)
}
func codexInputItemIDWithHashSuffix(id string, attempt int) string {
return codexInputItemIDWithHashSuffixRunes(id, []rune(id), attempt)
}
func codexInputItemIDWithHashSuffixRunes(id string, runes []rune, attempt int) string {
hashInput := id
if attempt > 0 {
hashInput += "\x00" + strconv.Itoa(attempt)
}
sum := sha256.Sum256([]byte(hashInput))
suffix := "_" + hex.EncodeToString(sum[:8])
prefixLength := codexInputItemIDLimit - len(suffix)
if len(runes) < prefixLength {
prefixLength = len(runes)
}
return string(runes[:prefixLength]) + suffix
}

View file

@ -0,0 +1,318 @@
package helps
import (
"fmt"
"strings"
"testing"
"github.com/tidwall/gjson"
)
var benchmarkSanitizeCodexInputItemIDsOutput []byte
func TestSanitizeCodexInputItemIDsBoundaries(t *testing.T) {
id64 := strings.Repeat("a", 64)
id65 := strings.Repeat("b", 65)
unicode65 := strings.Repeat("界", 65)
body := []byte(`{"input":[{"id":"` + id64 + `"},{"id":"` + id65 + `"},{"id":"` + unicode65 + `"}]}`)
got := SanitizeCodexInputItemIDs(body)
if actual := gjson.GetBytes(got, "input.0.id").String(); actual != id64 {
t.Fatalf("64-character ID changed: %q", actual)
}
for _, path := range []string{"input.1.id", "input.2.id"} {
actual := gjson.GetBytes(got, path).String()
if len([]rune(actual)) != 64 {
t.Fatalf("%s length = %d, want 64: %q", path, len([]rune(actual)), actual)
}
}
}
func TestSanitizeCodexInputItemIDsNormalizesMessageIDs(t *testing.T) {
const invalidID = "item_74ec40c883248ebb4885ec84"
body := []byte(`{"input":[` +
`{"type":"message","id":"` + invalidID + `","role":"user"},` +
`{"type":"message","id":"msg-1","role":"assistant"},` +
`{"type":"function_call","id":"item_call","call_id":"call-1"}` +
`]}`)
first := SanitizeCodexInputItemIDs(body)
second := SanitizeCodexInputItemIDs(body)
if got := gjson.GetBytes(first, "input.0.id").String(); got != "msg_"+invalidID {
t.Fatalf("message ID = %q, want msg-prefixed ID", got)
}
if got := gjson.GetBytes(first, "input.1.id").String(); got != "msg-1" {
t.Fatalf("valid message ID changed: %q", got)
}
if got := gjson.GetBytes(first, "input.2.id").String(); got != "fc_item_call" {
t.Fatalf("function_call ID was not normalized: %q", got)
}
if string(first) != string(second) {
t.Fatalf("message ID normalization is not deterministic: first=%s second=%s", first, second)
}
}
func TestSanitizeCodexInputItemIDsNormalizesResponseItemIDs(t *testing.T) {
const (
messageID = "item_message"
reasoningID = "item_reasoning"
functionCallID = "item_function_call"
functionCallOutputID = "item_function_call_output"
)
body := []byte(`{"input":[` +
`{"type":"message","id":"` + messageID + `"},` +
`{"type":"reasoning","id":"` + reasoningID + `"},` +
`{"type":"function_call","id":"` + functionCallID + `","call_id":"call-1"},` +
`{"type":"function_call_output","id":"` + functionCallOutputID + `","call_id":"call-1"},` +
`{"type":"reasoning","id":"rs-existing"},` +
`{"type":"function_call","id":"fc-existing","call_id":"call-2"},` +
`{"type":"message","id":"msg-existing"}` +
`]}`)
got := SanitizeCodexInputItemIDs(body)
want := []string{
"msg_" + messageID,
"rs_" + reasoningID,
"fc_" + functionCallID,
functionCallOutputID,
"rs-existing",
"fc-existing",
"msg-existing",
}
for index, expected := range want {
path := fmt.Sprintf("input.%d.id", index)
if actual := gjson.GetBytes(got, path).String(); actual != expected {
t.Fatalf("%s = %q, want %q; payload=%s", path, actual, expected, got)
}
}
if second := SanitizeCodexInputItemIDs(body); string(second) != string(got) {
t.Fatalf("normalization is not deterministic: first=%s second=%s", got, second)
}
}
func TestSanitizeCodexInputItemIDsAvoidsNormalizationCollisions(t *testing.T) {
for _, testCase := range []struct {
name string
itemType string
prefix string
}{
{name: "message", itemType: "message", prefix: "msg_"},
{name: "reasoning", itemType: "reasoning", prefix: "rs_"},
{name: "function call", itemType: "function_call", prefix: "fc_"},
{name: "custom tool call", itemType: "custom_tool_call", prefix: "ctc_"},
{name: "custom tool call output", itemType: "custom_tool_call_output", prefix: "ctco_"},
} {
for _, idCase := range []struct {
name string
invalidID string
}{
{name: "short", invalidID: "item_collision"},
{name: "overlong", invalidID: strings.Repeat("x", codexInputItemIDLimit-len([]rune(testCase.prefix))+1)},
} {
prefixedID := testCase.prefix + idCase.invalidID
for _, order := range []struct {
name string
ids [2]string
prefixedIndex int
}{
{name: "local first", ids: [2]string{idCase.invalidID, prefixedID}, prefixedIndex: 1},
{name: "prefixed first", ids: [2]string{prefixedID, idCase.invalidID}, prefixedIndex: 0},
} {
t.Run(testCase.name+"/"+idCase.name+"/"+order.name, func(t *testing.T) {
body := []byte(fmt.Sprintf(`{"input":[{"type":%q,"id":%q},{"type":%q,"id":%q}]}`, testCase.itemType, order.ids[0], testCase.itemType, order.ids[1]))
first := SanitizeCodexInputItemIDs(body)
second := SanitizeCodexInputItemIDs(body)
normalizedAgain := SanitizeCodexInputItemIDs(first)
ids := [2]string{
gjson.GetBytes(first, "input.0.id").String(),
gjson.GetBytes(first, "input.1.id").String(),
}
if ids[0] == ids[1] {
t.Fatalf("distinct IDs collided after normalization: %q; payload=%s", ids[0], first)
}
for index, id := range ids {
if !strings.HasPrefix(id, testCase.prefix) {
t.Fatalf("input.%d.id = %q, want prefix %q", index, id, testCase.prefix)
}
if len([]rune(id)) > codexInputItemIDLimit {
t.Fatalf("input.%d.id length = %d, want at most %d: %q", index, len([]rune(id)), codexInputItemIDLimit, id)
}
}
if len([]rune(prefixedID)) <= codexInputItemIDLimit && ids[order.prefixedIndex] != prefixedID {
t.Fatalf("existing valid ID changed: got %q want %q", ids[order.prefixedIndex], prefixedID)
}
if string(first) != string(second) {
t.Fatalf("collision resolution is not deterministic: first=%s second=%s", first, second)
}
if string(first) != string(normalizedAgain) {
t.Fatalf("collision resolution is not idempotent: first=%s normalized_again=%s", first, normalizedAgain)
}
})
}
}
}
}
func TestSanitizeCodexInputItemIDsNormalizesCustomToolCallIDs(t *testing.T) {
const invalidID = "item_44e13caebc1ddf25f1337cbe"
body := []byte(`{"input":[{"type":"custom_tool_call","id":"` + invalidID + `","call_id":"call-1","name":"lookup","input":"{}"}]}`)
got := SanitizeCodexInputItemIDs(body)
if actual := gjson.GetBytes(got, "input.0.id").String(); actual != "ctc_"+invalidID {
t.Fatalf("custom_tool_call ID = %q, want ctc-prefixed ID", actual)
}
}
func TestSanitizeCodexInputItemIDsNormalizesCustomToolCallOutputIDs(t *testing.T) {
const (
invalidID = "item_44e13caebc1ddf25f1337cbe_output"
validID = "ctco-existing"
)
body := []byte(`{"input":[` +
`{"type":"custom_tool_call_output","id":"` + invalidID + `","call_id":"call-1","output":"done"},` +
`{"type":"custom_tool_call_output","id":"` + validID + `","call_id":"call-2","output":"done"}` +
`]}`)
first := SanitizeCodexInputItemIDs(body)
second := SanitizeCodexInputItemIDs(body)
normalizedAgain := SanitizeCodexInputItemIDs(first)
if actual := gjson.GetBytes(first, "input.0.id").String(); actual != "ctco_"+invalidID {
t.Fatalf("custom_tool_call_output ID = %q, want ctco-prefixed ID", actual)
}
if actual := gjson.GetBytes(first, "input.1.id").String(); actual != validID {
t.Fatalf("valid custom_tool_call_output ID changed: %q", actual)
}
if string(first) != string(second) {
t.Fatalf("custom_tool_call_output ID normalization is not deterministic: first=%s second=%s", first, second)
}
if string(first) != string(normalizedAgain) {
t.Fatalf("custom_tool_call_output ID normalization is not idempotent: first=%s normalized_again=%s", first, normalizedAgain)
}
}
func TestSanitizeCodexInputItemIDsDropsOverlongEncryptedReasoningItem(t *testing.T) {
longReasoningID := "rs_" + strings.Repeat("a", 64)
shortReasoningID := "rs_" + strings.Repeat("b", 48)
longCallID := strings.Repeat("call-item-", 8)
body := []byte(`{"input":[` +
`{"type":"message","id":"msg-1","role":"user","content":"before"},` +
`{"type":"reasoning","id":"` + longReasoningID + `","encrypted_content":"gAAAA-encrypted","summary":[{"type":"summary_text","text":"drop me"}]},` +
`{"type":"reasoning","id":"` + shortReasoningID + `","encrypted_content":"gAAAA-encrypted","summary":[]},` +
`{"type":"function_call","id":"` + longCallID + `","call_id":"call-1","name":"lookup","arguments":"{}"}` +
`]}`)
got := SanitizeCodexInputItemIDs(body)
input := gjson.GetBytes(got, "input").Array()
if len(input) != 3 {
t.Fatalf("input length = %d, want 3: %s", len(input), got)
}
if gotID := input[0].Get("id").String(); gotID != "msg-1" {
t.Fatalf("input.0.id = %q, want msg-1", gotID)
}
if gotID := input[1].Get("id").String(); gotID != shortReasoningID {
t.Fatalf("short encrypted reasoning id changed: %q", gotID)
}
if gotID := input[2].Get("id").String(); gotID == longCallID || len([]rune(gotID)) != 64 {
t.Fatalf("ordinary overlong id was not shortened: %q", gotID)
}
}
func TestSanitizeCodexInputItemIDsShortensOverlongReasoningWithoutEncryptedContent(t *testing.T) {
longReasoningID := "rs_" + strings.Repeat("a", 64)
for _, testCase := range []struct {
name string
encryptedContent string
}{
{name: "missing"},
{name: "empty", encryptedContent: `,"encrypted_content":""`},
{name: "null", encryptedContent: `,"encrypted_content":null`},
} {
t.Run(testCase.name, func(t *testing.T) {
body := []byte(`{"input":[{"type":"reasoning","id":"` + longReasoningID + `"` + testCase.encryptedContent + `,"summary":[]}]}`)
got := SanitizeCodexInputItemIDs(body)
input := gjson.GetBytes(got, "input").Array()
if len(input) != 1 {
t.Fatalf("input length = %d, want 1: %s", len(input), got)
}
gotID := input[0].Get("id").String()
if gotID == longReasoningID || len([]rune(gotID)) != 64 {
t.Fatalf("overlong reasoning id was not shortened: %q", gotID)
}
})
}
}
func TestSanitizeCodexInputItemIDsAvoidsExistingIDCollision(t *testing.T) {
longID := strings.Repeat("grok-item-", 10)
collidingValidID := shortenCodexInputItemID(longID)
body := []byte(`{"input":[{"id":"` + longID + `"},{"id":"` + collidingValidID + `"}]}`)
first := SanitizeCodexInputItemIDs(body)
second := SanitizeCodexInputItemIDs(body)
shortened := gjson.GetBytes(first, "input.0.id").String()
if shortened == collidingValidID {
t.Fatalf("shortened ID collided with an existing valid ID: %q", shortened)
}
if len([]rune(shortened)) > 64 {
t.Fatalf("shortened ID length = %d, want at most 64", len([]rune(shortened)))
}
if actual := gjson.GetBytes(first, "input.1.id").String(); actual != collidingValidID {
t.Fatalf("existing valid ID changed: %q", actual)
}
if actual := gjson.GetBytes(second, "input.0.id").String(); actual != shortened {
t.Fatalf("collision resolution is not deterministic: first=%q second=%q", shortened, actual)
}
}
func TestSanitizeCodexInputItemIDsLeavesUnsupportedPayloadsUnchanged(t *testing.T) {
for _, body := range [][]byte{
[]byte(`not-json`),
[]byte(`{"input":{"id":"item-1"}}`),
[]byte(`{"input":[1,{"id":2},{"id":"item-1"}]}`),
} {
if got := string(SanitizeCodexInputItemIDs(body)); got != string(body) {
t.Fatalf("payload changed: got=%q want=%q", got, body)
}
}
}
func BenchmarkSanitizeCodexInputItemIDsLargeNoopPayload(b *testing.B) {
body := []byte(`{"input":[{"type":"message","id":"msg_1","role":"user","content":"` + strings.Repeat("x", 8<<20) + `"}]}`)
b.ReportAllocs()
b.SetBytes(int64(len(body)))
b.ResetTimer()
for b.Loop() {
benchmarkSanitizeCodexInputItemIDsOutput = SanitizeCodexInputItemIDs(body)
}
}
func BenchmarkSanitizeCodexInputItemIDsLargeHistory(b *testing.B) {
var payload strings.Builder
payload.Grow(64 << 10)
payload.WriteString(`{"input":[`)
for index := range 1000 {
if index > 0 {
payload.WriteByte(',')
}
fmt.Fprintf(&payload, `{"type":"message","id":"msg_%d","role":"user","content":"x"}`, index)
}
payload.WriteString(`]}`)
body := []byte(payload.String())
b.ReportAllocs()
b.SetBytes(int64(len(body)))
b.ResetTimer()
for b.Loop() {
benchmarkSanitizeCodexInputItemIDsOutput = SanitizeCodexInputItemIDs(body)
}
}

View file

@ -0,0 +1,127 @@
package helps
import (
"context"
"net/http"
multiagentv2 "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/optimize-multi-agent-v2"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
openaichatclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/claude/openai/chat-completions"
responsesclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/claude/openai/responses"
codexclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/codex/claude"
geminiclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/gemini/claude"
interactionsclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/interactions/claude"
openaiclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/openai/claude"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
)
// RewriteCodexSpawnAgentDescription optimizes spawn_agent definitions for
// official Codex clients when multi-agent v2 optimization is enabled.
func RewriteCodexSpawnAgentDescription(ctx context.Context, headers http.Header, payload []byte, cfg *config.Config) []byte {
return multiagentv2.RewriteCodexSpawnAgentDescription(ctx, headers, payload, cfg)
}
// RewriteCodexMultiAgentV2Input converts official Codex multi-agent input into
// standard Responses API messages when multi-agent v2 optimization is enabled.
func RewriteCodexMultiAgentV2Input(ctx context.Context, headers http.Header, payload []byte, cfg *config.Config) []byte {
return multiagentv2.RewriteCodexMultiAgentV2Input(ctx, headers, payload, cfg)
}
// TranslateRequestWithCodexMultiAgentV2 normalizes official Codex multi-agent
// input before translating it to a non-Codex target protocol.
func TranslateRequestWithCodexMultiAgentV2(ctx context.Context, headers http.Header, cfg *config.Config, from, to sdktranslator.Format, model string, payload []byte, stream bool) []byte {
return multiagentv2.TranslateRequestWithCodexMultiAgentV2(ctx, headers, cfg, from, to, model, payload, stream)
}
// TranslateRequestPairWithCodexMultiAgentV2 translates the untouched baseline
// payload and the working payload that later stages mutate in place. Executors
// normally assign the original payload to the request before translating, so both
// translations would rescan the same bytes and produce the same result. Built-in
// request translation is deterministic, so that case is translated once and
// duplicated when no plugin hooks are installed. Hooks retain two invocations
// because they may have request-scoped output or side effects. This removes a
// full extra pass over payloads that can reach tens of megabytes.
func TranslateRequestPairWithCodexMultiAgentV2(ctx context.Context, headers http.Header, cfg *config.Config, from, to sdktranslator.Format, model string, originalPayload, requestPayload []byte, stream bool) (original, working []byte) {
original = TranslateRequestWithCodexMultiAgentV2(ctx, headers, cfg, from, to, model, originalPayload, stream)
if sameByteSlice(originalPayload, requestPayload) && !sdktranslator.HasPluginHooks() {
// The caller mutates the working copy, so it must not share the baseline array.
return original, append([]byte(nil), original...)
}
return original, TranslateRequestWithCodexMultiAgentV2(ctx, headers, cfg, from, to, model, requestPayload, stream)
}
// sameByteSlice reports whether both slices describe the same bytes of the same
// backing array. It compares identity rather than content so the check stays
// constant time on large payloads.
func sameByteSlice(a, b []byte) bool {
if len(a) != len(b) {
return false
}
if len(a) == 0 {
return true
}
return &a[0] == &b[0]
}
// TranslateRequestWithAPIKeyModelCompatibility applies compatibility-aware
// request translators when a configured API-key model enables compatibility mode.
func TranslateRequestWithAPIKeyModelCompatibility(ctx context.Context, headers http.Header, cfg *config.Config, from, to sdktranslator.Format, model string, payload []byte, stream, isCompat bool) []byte {
if !isCompat {
return TranslateRequestWithCodexMultiAgentV2(ctx, headers, cfg, from, to, model, payload, stream)
}
if from == sdktranslator.FormatOpenAIResponse && to != sdktranslator.FormatCodex && to != sdktranslator.FormatOpenAIResponse {
payload = multiagentv2.RewriteCodexMultiAgentV2Input(ctx, headers, payload, cfg)
}
var translated []byte
switch {
case from == sdktranslator.FormatClaude && to == sdktranslator.FormatCodex:
translated = codexclaude.ConvertClaudeRequestToCodexWithCompat(model, payload, stream)
case from == sdktranslator.FormatClaude && to == sdktranslator.FormatGemini:
translated = geminiclaude.ConvertClaudeRequestToGeminiWithCompat(model, payload, stream)
case from == sdktranslator.FormatClaude && to == sdktranslator.FormatInteractions:
translated = interactionsclaude.ConvertClaudeRequestToInteractionsWithCompat(model, payload, stream)
case from == sdktranslator.FormatClaude && to == sdktranslator.FormatOpenAI:
translated = openaiclaude.ConvertClaudeRequestToOpenAIWithCompat(model, payload, stream)
case from == sdktranslator.FormatOpenAI && to == sdktranslator.FormatClaude:
translated = openaichatclaude.ConvertOpenAIRequestToClaudeWithCompat(model, payload, stream)
case from == sdktranslator.FormatOpenAIResponse && to == sdktranslator.FormatClaude:
translated = responsesclaude.ConvertOpenAIResponsesRequestToClaudeWithCompat(model, payload, stream)
default:
return TranslateRequestWithCodexMultiAgentV2(ctx, headers, cfg, from, to, model, payload, stream)
}
summaryConfig := thinking.ExtractSummaryConfig(payload, from.String())
return thinking.ApplySummaryConfigForModel(translated, to.String(), model, summaryConfig)
}
// HasCodexMultiAgentV2NamespaceConflict reports whether the request defines
// the reserved optimized namespace, which must remain untouched.
func HasCodexMultiAgentV2NamespaceConflict(payload []byte) bool {
return multiagentv2.HasCodexMultiAgentV2NamespaceConflict(payload)
}
// OptimizeCodexMultiAgentV2Request rewrites an eligible spawn_agent request and
// reports whether the collaboration namespace was renamed for upstream use.
func OptimizeCodexMultiAgentV2Request(ctx context.Context, headers http.Header, payload []byte, cfg *config.Config) ([]byte, bool) {
return multiagentv2.OptimizeCodexMultiAgentV2Request(ctx, headers, payload, cfg)
}
// OptimizeCodexMultiAgentV2RequestForAuth applies the standard Codex MultiAgentV2
// request optimization and, when the selected codex-api-key model has is-compat
// enabled, also converts agent_message items into portable message/user input.
func OptimizeCodexMultiAgentV2RequestForAuth(ctx context.Context, headers http.Header, payload []byte, cfg *config.Config, auth *cliproxyauth.Auth, model string) ([]byte, bool) {
updated, optimized := multiagentv2.OptimizeCodexMultiAgentV2Request(ctx, headers, payload, cfg)
if cliproxyauth.CodexAPIKeyModelIsCompat(cfg, auth, model) {
updated = multiagentv2.RewriteCodexMultiAgentV2Input(ctx, headers, updated, cfg)
}
return updated, optimized
}
// RestoreCodexMultiAgentV2Response restores optimized collaboration namespace
// values before an upstream response is translated and returned to the client.
func RestoreCodexMultiAgentV2Response(payload []byte, optimized bool) []byte {
return multiagentv2.RestoreCodexMultiAgentV2Response(payload, optimized)
}

View file

@ -0,0 +1,182 @@
package helps
import (
"bytes"
"context"
"fmt"
"net/http"
"strings"
"testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
type pairRequestPluginHooks struct {
calls int64
}
func (h *pairRequestPluginHooks) NormalizeRequest(_ context.Context, _, _ sdktranslator.Format, _ string, body []byte, _ bool) []byte {
h.calls++
updated, _ := sjson.SetBytes(body, "plugin_call", h.calls)
return updated
}
func (*pairRequestPluginHooks) TranslateRequest(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, bool) ([]byte, bool) {
return nil, false
}
func (*pairRequestPluginHooks) NormalizeResponseBefore(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, []byte, []byte, bool) []byte {
return nil
}
func (*pairRequestPluginHooks) TranslateResponse(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, []byte, []byte, bool) ([]byte, bool) {
return nil, false
}
func (*pairRequestPluginHooks) NormalizeResponseAfter(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, []byte, []byte, bool) []byte {
return nil
}
func geminiToolHistoryPayload(turns int) []byte {
contents := []string{`{"role":"user","parts":[{"text":"start"}]}`}
for i := 0; i < turns; i++ {
contents = append(contents,
fmt.Sprintf(`{"role":"user","parts":[{"text":"ask %d"}]}`, i),
fmt.Sprintf(`{"role":"model","parts":[{"text":"think %d"},{"thoughtSignature":"sig-%d","functionCall":{"id":"c%d","name":"read_file","args":{"path":"a%d.go"}}}]}`, i, i, i, i),
fmt.Sprintf(`{"role":"user","parts":[{"functionResponse":{"id":"c%d","name":"read_file","response":{"content":"data %d"}}}]}`, i, i),
fmt.Sprintf(`{"role":"model","parts":[{"text":"answer %d"}]}`, i))
}
return []byte(fmt.Sprintf(
`{"contents":[%s],"tools":[{"functionDeclarations":[{"name":"read_file","description":"read a file","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}]}],"generationConfig":{"temperature":1}}`,
strings.Join(contents, ",")))
}
// TestTranslateRequestPairMatchesSeparateTranslations pins the reuse fast path to
// the behavior of translating both payloads independently.
func TestTranslateRequestPairMatchesSeparateTranslations(t *testing.T) {
from := sdktranslator.FormatGemini
to := sdktranslator.FromString("antigravity")
cfg := &config.Config{}
const model = "gemini-3.6-flash-high"
for _, turns := range []int{0, 1, 5, 20} {
payload := geminiToolHistoryPayload(turns)
// Same bytes in a different backing array forces the translate-twice branch.
detached := append([]byte(nil), payload...)
want := TranslateRequestWithCodexMultiAgentV2(context.Background(), http.Header{}, cfg, from, to, model, payload, true)
reuseBase, reuseWork := TranslateRequestPairWithCodexMultiAgentV2(
context.Background(), http.Header{}, cfg, from, to, model, payload, payload, true)
twiceBase, twiceWork := TranslateRequestPairWithCodexMultiAgentV2(
context.Background(), http.Header{}, cfg, from, to, model, payload, detached, true)
for name, got := range map[string][]byte{
"reuse baseline": reuseBase,
"reuse working": reuseWork,
"twice baseline": twiceBase,
"twice working": twiceWork,
} {
if !bytes.Equal(want, got) {
t.Fatalf("turns=%d: %s translation differs from a standalone translation", turns, name)
}
}
if len(reuseBase) > 0 && &reuseBase[0] == &reuseWork[0] {
t.Fatalf("turns=%d: working copy aliases the baseline; later in-place edits would corrupt it", turns)
}
// The caller mutates the working copy, so the baseline must stay intact.
baselineBefore := append([]byte(nil), reuseBase...)
reuseWork[0] = 'X'
if !bytes.Equal(baselineBefore, reuseBase) {
t.Fatalf("turns=%d: mutating the working copy changed the baseline", turns)
}
}
}
// TestTranslateRequestPairTranslatesDistinctPayloads guards the case where the
// executor really does hand over two different requests.
func TestTranslateRequestPairTranslatesDistinctPayloads(t *testing.T) {
from := sdktranslator.FormatGemini
to := sdktranslator.FromString("antigravity")
cfg := &config.Config{}
const model = "gemini-3.6-flash-high"
original := geminiToolHistoryPayload(2)
request := geminiToolHistoryPayload(4)
base, work := TranslateRequestPairWithCodexMultiAgentV2(
context.Background(), http.Header{}, cfg, from, to, model, original, request, true)
wantBase := TranslateRequestWithCodexMultiAgentV2(context.Background(), http.Header{}, cfg, from, to, model, original, true)
wantWork := TranslateRequestWithCodexMultiAgentV2(context.Background(), http.Header{}, cfg, from, to, model, request, true)
if !bytes.Equal(wantBase, base) {
t.Fatal("baseline translation differs for distinct payloads")
}
if !bytes.Equal(wantWork, work) {
t.Fatal("working translation differs for distinct payloads")
}
if bytes.Equal(base, work) {
t.Fatal("distinct payloads produced identical translations; the reuse path was taken by mistake")
}
}
func TestTranslateRequestPairPreservesPluginHookInvocations(t *testing.T) {
hooks := &pairRequestPluginHooks{}
sdktranslator.SetPluginHooks(hooks)
t.Cleanup(func() { sdktranslator.SetPluginHooks(nil) })
payload := geminiToolHistoryPayload(1)
base, work := TranslateRequestPairWithCodexMultiAgentV2(
context.Background(),
http.Header{},
&config.Config{},
sdktranslator.FormatGemini,
sdktranslator.FromString("antigravity"),
"gemini-3.6-flash-high",
payload,
payload,
true,
)
if hooks.calls != 2 {
t.Fatalf("plugin hook calls = %d, want 2", hooks.calls)
}
if got := gjson.GetBytes(base, "plugin_call").Int(); got != 1 {
t.Fatalf("baseline plugin_call = %d, want 1", got)
}
if got := gjson.GetBytes(work, "plugin_call").Int(); got != 2 {
t.Fatalf("working plugin_call = %d, want 2", got)
}
}
func TestSameByteSlice(t *testing.T) {
buf := []byte("payload")
cases := []struct {
name string
a, b []byte
want bool
}{
{"identical slice", buf, buf, true},
{"same array same length", buf[:3], buf[:3], true},
{"equal bytes different array", buf, append([]byte(nil), buf...), false},
{"different length", buf, buf[:3], false},
{"both nil", nil, nil, true},
{"nil and empty", nil, []byte{}, true},
{"nil and non-empty", nil, buf, false},
{"offset alias", buf, buf[1:], false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := sameByteSlice(tc.a, tc.b); got != tc.want {
t.Fatalf("sameByteSlice() = %v, want %v", got, tc.want)
}
})
}
}

View file

@ -0,0 +1,66 @@
package helps
import (
"crypto/sha256"
"encoding/binary"
"strconv"
"strings"
"github.com/google/uuid"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
cliproxysession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session"
)
// DerivedSessionID returns the first context-derived session identity in metadata order.
func DerivedSessionID(metadataSets ...map[string]any) string {
for _, metadata := range metadataSets {
if derivedID := cliproxysession.DerivedID(metadata); derivedID != "" {
return derivedID
}
}
return ""
}
// DerivedSessionUUID maps a derived session identity to a provider-scoped stable UUID.
func DerivedSessionUUID(provider string, metadataSets ...map[string]any) string {
return stableProviderSessionUUID(provider, "derived-session", DerivedSessionID(metadataSets...))
}
// ProviderSessionUUID prefers a long-lived execution session and falls back to the derived identity.
func ProviderSessionUUID(provider string, metadataSets ...map[string]any) string {
for _, metadata := range metadataSets {
if executionID := metadataString(metadata, cliproxyexecutor.ExecutionSessionMetadataKey); executionID != "" {
return stableProviderSessionUUID(provider, "execution-session", executionID)
}
}
return DerivedSessionUUID(provider, metadataSets...)
}
func stableProviderSessionUUID(provider string, kind string, identityValue string) string {
provider = strings.ToLower(strings.TrimSpace(provider))
identityValue = strings.TrimSpace(identityValue)
if provider == "" || identityValue == "" {
return ""
}
identity := strings.Join([]string{"cli-proxy-api", provider, kind, identityValue}, "\x00")
return uuid.NewSHA1(uuid.NameSpaceOID, []byte(identity)).String()
}
// DerivedAntigravitySessionID maps a derived session identity to Antigravity's negative decimal format.
func DerivedAntigravitySessionID(metadataSets ...map[string]any) string {
derivedID := DerivedSessionID(metadataSets...)
if derivedID == "" {
return ""
}
sum := sha256.Sum256([]byte("cli-proxy-api:antigravity:derived-session\x00" + derivedID))
value := int64(binary.BigEndian.Uint64(sum[:8])) & 0x7FFFFFFFFFFFFFFF
return "-" + strconv.FormatInt(value, 10)
}
func metadataString(metadata map[string]any, key string) string {
if metadata == nil {
return ""
}
value, _ := metadata[key].(string)
return strings.TrimSpace(value)
}

View file

@ -0,0 +1,69 @@
package helps
import (
"regexp"
"testing"
"github.com/google/uuid"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
)
func TestDerivedSessionProviderMappings(t *testing.T) {
t.Parallel()
metadata := map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:test-root"}
codexID := DerivedSessionUUID("codex", metadata)
xaiID := DerivedSessionUUID("xai", metadata)
if _, errParse := uuid.Parse(codexID); errParse != nil {
t.Fatalf("Codex mapping %q is not a UUID: %v", codexID, errParse)
}
if _, errParse := uuid.Parse(xaiID); errParse != nil {
t.Fatalf("xAI mapping %q is not a UUID: %v", xaiID, errParse)
}
if codexID == xaiID {
t.Fatalf("provider namespaces produced the same UUID: %q", codexID)
}
if repeated := DerivedSessionUUID("codex", metadata); repeated != codexID {
t.Fatalf("Codex mapping is not stable: first=%q repeated=%q", codexID, repeated)
}
antigravityID := DerivedAntigravitySessionID(metadata)
if matched := regexp.MustCompile(`^-[0-9]+$`).MatchString(antigravityID); !matched {
t.Fatalf("Antigravity mapping = %q, want negative decimal", antigravityID)
}
if repeated := DerivedAntigravitySessionID(metadata); repeated != antigravityID {
t.Fatalf("Antigravity mapping is not stable: first=%q repeated=%q", antigravityID, repeated)
}
}
func TestProviderSessionUUIDPrefersExecutionSession(t *testing.T) {
t.Parallel()
first := map[string]any{
cliproxyexecutor.ExecutionSessionMetadataKey: "connection-1",
cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:first-root",
}
second := map[string]any{
cliproxyexecutor.ExecutionSessionMetadataKey: "connection-1",
cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:second-root",
}
firstID := ProviderSessionUUID("codex", first)
secondID := ProviderSessionUUID("codex", second)
if firstID == "" || firstID != secondID {
t.Fatalf("execution session did not stabilize provider UUID: first=%q second=%q", firstID, secondID)
}
if firstID == DerivedSessionUUID("codex", first) {
t.Fatalf("provider UUID did not prefer execution session: %q", firstID)
}
}
func TestDerivedSessionProviderMappingsRequireIdentity(t *testing.T) {
t.Parallel()
if got := DerivedSessionUUID("codex", nil); got != "" {
t.Fatalf("DerivedSessionUUID() = %q, want empty", got)
}
if got := DerivedAntigravitySessionID(nil); got != "" {
t.Fatalf("DerivedAntigravitySessionID() = %q, want empty", got)
}
}

View file

@ -0,0 +1,39 @@
package helps
import (
translatorcommon "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/common"
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
var emptyGeminiUserTurnJSON = []byte(`{"role":"user","parts":[{"text":""}]}`)
// EnsureGeminiLeadingUserContent ensures that the contents array at the given path
// starts with a user turn when sending to Gemini/Antigravity upstreams.
func EnsureGeminiLeadingUserContent(payload []byte, path string) []byte {
firstRole := gjson.GetBytes(payload, path+".0.role")
if firstRole.String() != "model" {
return payload
}
contents := util.GetGJSONBytesNoCopy(payload, path)
if !contents.IsArray() {
return payload
}
contentArray := contents.Array()
if len(contentArray) == 0 {
return payload
}
contentItems := make([][]byte, 0, len(contentArray)+1)
contentItems = append(contentItems, emptyGeminiUserTurnJSON)
for _, content := range contentArray {
contentItems = append(contentItems, []byte(content.Raw))
}
out, errSet := sjson.SetRawBytes(payload, path, translatorcommon.JoinRawArray(contentItems))
if errSet != nil {
return payload
}
return out
}

View file

@ -0,0 +1,99 @@
package helps
import (
"strings"
"testing"
"github.com/tidwall/gjson"
)
var leadingGeminiUserContentOutput []byte
func TestEnsureGeminiLeadingUserContentReusesLargeValidPayload(t *testing.T) {
input := []byte(`{"contents":[{"role":"user","parts":[{"inlineData":{"mimeType":"video/mp4","data":"` + strings.Repeat("A", 4<<20) + `"}}]}]}`)
output := EnsureGeminiLeadingUserContent(input, "contents")
if &output[0] != &input[0] {
t.Fatal("valid request should reuse the input payload")
}
result := testing.Benchmark(func(b *testing.B) {
for b.Loop() {
leadingGeminiUserContentOutput = EnsureGeminiLeadingUserContent(input, "contents")
}
})
if allocated := result.AllocedBytesPerOp(); allocated >= 1<<20 {
t.Fatalf("valid 4 MiB request allocated %d bytes/op, want less than 1 MiB", allocated)
}
}
func TestEnsureGeminiLeadingUserContent(t *testing.T) {
tests := []struct {
name string
inputJSON string
path string
wantRoles string
wantLeadingEmpty bool
}{
{
name: "user first is unchanged",
inputJSON: `{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`,
path: "contents",
wantRoles: "user",
},
{
name: "leading model functionCall gets empty user",
inputJSON: `{"contents":[{"role":"model","parts":[{"functionCall":{"name":"run"}}]},{"role":"user","parts":[{"functionResponse":{"name":"run"}}]}]}`,
path: "contents",
wantRoles: "user,model,user",
wantLeadingEmpty: true,
},
{
name: "leading model text gets empty user and preserves following turns",
inputJSON: `{"contents":[{"role":"model","parts":[{"text":"answer"}]},{"role":"user","parts":[{"text":"continue"}]}]}`,
path: "contents",
wantRoles: "user,model,user",
wantLeadingEmpty: true,
},
{
name: "nested contents are normalized",
inputJSON: `{"request":{"contents":[{"role":"model","parts":[{"text":"answer"}]},{"role":"user","parts":[{"text":"continue"}]}]}}`,
path: "request.contents",
wantRoles: "request.user,model,user",
wantLeadingEmpty: true,
},
{
name: "empty contents are unchanged",
inputJSON: `{"contents":[]}`,
path: "contents",
wantRoles: "",
},
{
name: "missing contents are unchanged",
inputJSON: `{"model":"test"}`,
path: "contents",
wantRoles: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
out := EnsureGeminiLeadingUserContent([]byte(tt.inputJSON), tt.path)
contents := gjson.GetBytes(out, tt.path).Array()
roles := make([]string, 0, len(contents))
for _, content := range contents {
roles = append(roles, content.Get("role").String())
}
expectedRoles := strings.TrimPrefix(tt.wantRoles, "request.")
if got := strings.Join(roles, ","); got != expectedRoles {
t.Fatalf("roles = %q, want %q; output=%s", got, expectedRoles, out)
}
if tt.wantLeadingEmpty {
text := gjson.GetBytes(out, tt.path+".0.parts.0.text")
if !text.Exists() || text.String() != "" {
t.Fatalf("leading empty user part missing; output=%s", out)
}
}
})
}
}

View file

@ -0,0 +1,155 @@
package helps
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/home"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
)
type homeStatusErr struct {
code int
msg string
}
func (e homeStatusErr) Error() string {
if e.msg != "" {
return e.msg
}
return fmt.Sprintf("status %d", e.code)
}
func (e homeStatusErr) StatusCode() int { return e.code }
type homeErrorEnvelope struct {
Error *homeErrorDetail `json:"error"`
}
type homeRefreshAuthEnvelope struct {
Auth cliproxyauth.Auth `json:"auth"`
AuthIndex string `json:"auth_index"`
}
type homeErrorDetail struct {
Type string `json:"type"`
Message string `json:"message"`
Code string `json:"code,omitempty"`
}
type homeRefreshClient interface {
HeartbeatOK() bool
GetRefreshAuth(ctx context.Context, authIndex string, accessTokenSHA256 string) ([]byte, error)
}
var currentHomeRefreshClient = func() homeRefreshClient {
return home.Current()
}
// RefreshAuthViaHome replaces local refresh logic when home control plane integration is enabled.
// It returns (updatedAuth, true, nil) when home refresh succeeds; (nil, true, err) when home is
// enabled but refresh fails; and (nil, false, nil) when home is disabled.
func RefreshAuthViaHome(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, bool, error) {
if cfg == nil || !cfg.Home.Enabled {
return nil, false, nil
}
if ctx == nil {
ctx = context.Background()
}
if auth == nil {
return nil, true, homeStatusErr{code: http.StatusInternalServerError, msg: "home refresh: auth is nil"}
}
client := currentHomeRefreshClient()
if client == nil || !client.HeartbeatOK() {
return nil, true, homeStatusErr{code: http.StatusServiceUnavailable, msg: "home control center unavailable"}
}
authIndex := strings.TrimSpace(auth.Index)
if authIndex == "" {
authIndex = strings.TrimSpace(auth.EnsureIndex())
}
if authIndex == "" {
return nil, true, homeStatusErr{code: http.StatusBadGateway, msg: "home refresh: auth_index is empty"}
}
raw, err := client.GetRefreshAuth(ctx, authIndex, authAccessTokenSHA256(auth))
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return nil, true, err
}
return nil, true, homeStatusErr{code: http.StatusServiceUnavailable, msg: "home refresh temporarily unavailable"}
}
var env homeErrorEnvelope
if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal == nil && env.Error != nil {
code := strings.TrimSpace(env.Error.Type)
if code == "" {
code = strings.TrimSpace(env.Error.Code)
}
statusCode := statusFromHomeErrorCode(code)
message := "credential refresh temporarily unavailable"
switch statusCode {
case http.StatusUnauthorized:
message = "credential unauthorized"
case http.StatusNotFound:
message = "credential refresh target not found"
}
return nil, true, homeStatusErr{code: statusCode, msg: message}
}
updated, returnedIndex, errParse := parseHomeRefreshAuth(raw)
if errParse != nil {
return nil, true, homeStatusErr{code: http.StatusBadGateway, msg: "home returned invalid auth payload"}
}
if updated.Disabled || updated.Status == cliproxyauth.StatusDisabled {
return nil, true, homeStatusErr{code: http.StatusUnauthorized, msg: "credential unauthorized"}
}
if returnedIndex != "" {
authIndex = returnedIndex
}
updated.Index = authIndex
updated.EnsureIndex()
return updated, true, nil
}
func authAccessTokenSHA256(auth *cliproxyauth.Auth) string {
return cliproxyauth.AccessTokenSHA256(auth)
}
func parseHomeRefreshAuth(raw []byte) (*cliproxyauth.Auth, string, error) {
var rawObject map[string]json.RawMessage
if errUnmarshal := json.Unmarshal(raw, &rawObject); errUnmarshal != nil {
return nil, "", errUnmarshal
}
if _, ok := rawObject["auth"]; ok {
var envelope homeRefreshAuthEnvelope
if errUnmarshal := json.Unmarshal(raw, &envelope); errUnmarshal != nil {
return nil, "", errUnmarshal
}
return &envelope.Auth, strings.TrimSpace(envelope.AuthIndex), nil
}
var updated cliproxyauth.Auth
if errUnmarshal := json.Unmarshal(raw, &updated); errUnmarshal != nil {
return nil, "", errUnmarshal
}
return &updated, "", nil
}
func statusFromHomeErrorCode(code string) int {
switch strings.ToLower(strings.TrimSpace(code)) {
case "authentication_error", "unauthorized", "invalid_grant", "refresh_token_expired", "refresh_token_revoked", "refresh_token_reused":
return http.StatusUnauthorized
case "model_not_found":
return http.StatusNotFound
case "auth_not_found", "auth_unavailable", "refresh_temporarily_unavailable", "refresh_unsupported", "home_unavailable":
return http.StatusServiceUnavailable
default:
return http.StatusServiceUnavailable
}
}

View file

@ -0,0 +1,178 @@
package helps
import (
"context"
"encoding/json"
"errors"
"net/http"
"strings"
"sync/atomic"
"testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
)
func TestStatusFromHomeErrorCodeMapsAuthenticationErrorToUnauthorized(t *testing.T) {
if got := statusFromHomeErrorCode("authentication_error"); got != http.StatusUnauthorized {
t.Fatalf("statusFromHomeErrorCode(authentication_error) = %d, want %d", got, http.StatusUnauthorized)
}
if got := statusFromHomeErrorCode("unauthorized"); got != http.StatusUnauthorized {
t.Fatalf("statusFromHomeErrorCode(unauthorized) = %d, want %d", got, http.StatusUnauthorized)
}
for _, code := range []string{"auth_not_found", "auth_unavailable", "refresh_temporarily_unavailable", "refresh_unsupported"} {
if got := statusFromHomeErrorCode(code); got != http.StatusServiceUnavailable {
t.Fatalf("statusFromHomeErrorCode(%s) = %d, want %d", code, got, http.StatusServiceUnavailable)
}
}
}
type fakeHomeRefreshClient struct {
calls atomic.Int32
authIndex string
accessTokenHash string
raw []byte
err error
}
func (c *fakeHomeRefreshClient) HeartbeatOK() bool {
return true
}
func (c *fakeHomeRefreshClient) GetRefreshAuth(_ context.Context, authIndex string, accessTokenHash string) ([]byte, error) {
c.calls.Add(1)
c.authIndex = authIndex
c.accessTokenHash = accessTokenHash
return c.raw, c.err
}
func TestRefreshAuthViaHomePreservesContextErrors(t *testing.T) {
client := &fakeHomeRefreshClient{err: context.DeadlineExceeded}
oldCurrentHomeRefreshClient := currentHomeRefreshClient
currentHomeRefreshClient = func() homeRefreshClient { return client }
t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient })
cfg := &config.Config{Home: config.HomeConfig{Enabled: true}}
auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: "codex"}
_, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth)
if !handled || !errors.Is(errRefresh, context.DeadlineExceeded) {
t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want true/context.DeadlineExceeded", handled, errRefresh)
}
}
func TestRefreshAuthViaHomeMapsTransportFailureToRedacted503(t *testing.T) {
client := &fakeHomeRefreshClient{err: errors.New("dial failed with provider-secret")}
oldCurrentHomeRefreshClient := currentHomeRefreshClient
currentHomeRefreshClient = func() homeRefreshClient { return client }
t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient })
cfg := &config.Config{Home: config.HomeConfig{Enabled: true}}
auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: "codex"}
_, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth)
statusErr, okStatus := errRefresh.(interface{ StatusCode() int })
if !handled || !okStatus || statusErr.StatusCode() != http.StatusServiceUnavailable {
t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want redacted 503", handled, errRefresh)
}
if strings.Contains(errRefresh.Error(), "provider-secret") {
t.Fatalf("refresh error leaked transport detail: %v", errRefresh)
}
}
func TestRefreshAuthViaHomeRedactsLegacyErrorEnvelope(t *testing.T) {
client := &fakeHomeRefreshClient{raw: []byte(`{"error":{"type":"error","message":"provider response: refresh_token=provider-secret"}}`)}
oldCurrentHomeRefreshClient := currentHomeRefreshClient
currentHomeRefreshClient = func() homeRefreshClient { return client }
t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient })
cfg := &config.Config{Home: config.HomeConfig{Enabled: true}}
auth := &cliproxyauth.Auth{ID: "home-auth", Index: "home-auth", Provider: "codex"}
_, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth)
statusErr, okStatus := errRefresh.(interface{ StatusCode() int })
if !handled || !okStatus || statusErr.StatusCode() != http.StatusServiceUnavailable {
t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want redacted 503", handled, errRefresh)
}
if strings.Contains(errRefresh.Error(), "provider-secret") {
t.Fatalf("refresh error leaked legacy Home detail: %v", errRefresh)
}
}
func TestAuthAccessTokenSHA256SupportsKnownMetadataShapes(t *testing.T) {
want := authAccessTokenSHA256(&cliproxyauth.Auth{Metadata: map[string]any{"access_token": "same-token"}})
cases := map[string]*cliproxyauth.Auth{
"camel case": {Metadata: map[string]any{"accessToken": "same-token"}},
"nested any map": {Metadata: map[string]any{"token": map[string]any{"access_token": "same-token"}}},
"nested string map": {Metadata: map[string]any{"Token": map[string]string{"accessToken": "same-token"}}},
}
for name, auth := range cases {
t.Run(name, func(t *testing.T) {
if got := authAccessTokenSHA256(auth); got == "" || got != want {
t.Fatalf("token hash = %q, want %q", got, want)
}
})
}
}
func TestRefreshAuthViaHomeAcceptsAuthEnvelope(t *testing.T) {
raw, errMarshal := json.Marshal(struct {
Auth cliproxyauth.Auth `json:"auth"`
AuthIndex string `json:"auth_index"`
}{
Auth: cliproxyauth.Auth{
ID: "home-auth-1",
Provider: "antigravity",
Metadata: map[string]any{
"access_token": "new-access-token",
},
},
AuthIndex: "home-index-1",
})
if errMarshal != nil {
t.Fatalf("marshal home envelope: %v", errMarshal)
}
client := &fakeHomeRefreshClient{raw: raw}
oldCurrentHomeRefreshClient := currentHomeRefreshClient
currentHomeRefreshClient = func() homeRefreshClient {
return client
}
t.Cleanup(func() {
currentHomeRefreshClient = oldCurrentHomeRefreshClient
})
cfg := &config.Config{Home: config.HomeConfig{Enabled: true}}
auth := &cliproxyauth.Auth{
ID: "home-auth-1",
Provider: "antigravity",
Index: "home-index-1",
Metadata: map[string]any{
"access_token": "old-access-token",
"refresh_token": "refresh-token",
},
}
updated, handled, err := RefreshAuthViaHome(context.Background(), cfg, auth)
if err != nil {
t.Fatalf("RefreshAuthViaHome error: %v", err)
}
if !handled {
t.Fatal("RefreshAuthViaHome handled = false, want true")
}
if got := client.calls.Load(); got != 1 {
t.Fatalf("home refresh calls = %d, want 1", got)
}
if client.authIndex != "home-index-1" {
t.Fatalf("home refresh auth_index = %q, want home-index-1", client.authIndex)
}
if client.accessTokenHash != authAccessTokenSHA256(auth) {
t.Fatalf("home refresh access token hash = %q, want %q", client.accessTokenHash, authAccessTokenSHA256(auth))
}
if updated == nil {
t.Fatal("updated auth = nil")
}
if got := updated.Metadata["access_token"]; got != "new-access-token" {
t.Fatalf("updated access_token = %q, want new-access-token", got)
}
if updated.Index != "home-index-1" {
t.Fatalf("updated auth_index = %q, want home-index-1", updated.Index)
}
}

View file

@ -0,0 +1,80 @@
package helps
import (
"fmt"
"regexp"
"strconv"
"strings"
"time"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// DeleteJSONField removes a top-level or nested JSON field from a payload.
func DeleteJSONField(body []byte, key string) []byte {
if key == "" || len(body) == 0 {
return body
}
updated, err := sjson.DeleteBytes(body, key)
if err != nil {
return body
}
return updated
}
// ParseRetryDelay extracts the retry delay from a Google API 429 error response.
func ParseRetryDelay(errorBody []byte) (*time.Duration, error) {
details := gjson.GetBytes(errorBody, "error.details")
if details.Exists() && details.IsArray() {
for _, detail := range details.Array() {
if detail.Get("@type").String() != "type.googleapis.com/google.rpc.RetryInfo" {
continue
}
retryDelay := detail.Get("retryDelay").String()
if retryDelay == "" {
continue
}
duration, err := time.ParseDuration(retryDelay)
if err != nil {
return nil, fmt.Errorf("failed to parse duration")
}
return &duration, nil
}
for _, detail := range details.Array() {
if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" {
continue
}
quotaResetDelay := detail.Get("metadata.quotaResetDelay").String()
if quotaResetDelay == "" {
continue
}
duration, err := time.ParseDuration(quotaResetDelay)
if err == nil {
return &duration, nil
}
}
}
message := gjson.GetBytes(errorBody, "error.message").String()
if message != "" {
re := regexp.MustCompile(`after\s+(\d+)s\.?`)
if matches := re.FindStringSubmatch(message); len(matches) > 1 {
seconds, err := strconv.Atoi(matches[1])
if err == nil {
duration := time.Duration(seconds) * time.Second
return &duration, nil
}
}
reHuman := regexp.MustCompile(`after\s+((?:\d+h)?(?:\d+m)?(?:\d+s)?)\.?`)
if matches := reHuman.FindStringSubmatch(strings.ToLower(message)); len(matches) > 1 {
duration, err := time.ParseDuration(matches[1])
if err == nil && duration > 0 {
return &duration, nil
}
}
}
return nil, fmt.Errorf("no RetryInfo found")
}

View file

@ -0,0 +1,761 @@
package helps
import (
"bytes"
"context"
"fmt"
"html"
"net/http"
"net/url"
"sort"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
)
const (
apiAttemptsKey = "API_UPSTREAM_ATTEMPTS"
apiRequestKey = "API_REQUEST"
apiResponseKey = "API_RESPONSE"
apiWebsocketTimelineKey = "API_WEBSOCKET_TIMELINE"
deferredAPIRequestBytesKey = "DEFERRED_API_REQUEST_BYTES"
creditsUsedKey = "__antigravity_credits_used__"
maxDeferredAPIRequestBodyBytes = 32 << 20 // 32 MiB
)
// UpstreamRequestLog captures the outbound upstream request details for logging.
type UpstreamRequestLog struct {
URL string
Method string
Headers http.Header
Body []byte
Provider string
AuthID string
AuthLabel string
AuthType string
AuthValue string
}
type upstreamAttempt struct {
index int
request string
response *strings.Builder
responseSource *logging.FileBodySource
responseIntroWritten bool
statusWritten bool
headersWritten bool
bodyStarted bool
bodyHasContent bool
prevWasSSEEvent bool
errorWritten bool
}
func requestLogCaptureEnabled(cfg *config.Config) bool {
return cfg != nil && cfg.RequestLog && !cfg.CommercialMode
}
// RecordAPIRequest stores the upstream request metadata in Gin context for request logging.
func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) {
if cfg == nil || cfg.CommercialMode {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
if !cfg.RequestLog {
deferAPIRequest(ginCtx, info)
return
}
attempts := getAttempts(ginCtx)
index := len(attempts) + 1
builder := newAPIRequestLogBuilder(index, info, time.Now())
requestText := ""
if source, ok := apiRequestSource(ginCtx); ok {
if errWrite := source.AppendBytes([]byte(builder.String())); errWrite == nil {
if len(info.Body) > 0 {
if errBody := source.AppendBytes(info.Body); errBody != nil {
log.WithError(errBody).Warn("failed to append api request body log part")
}
} else if errEmpty := source.AppendBytes([]byte("<empty>")); errEmpty != nil {
log.WithError(errEmpty).Warn("failed to append empty api request log part")
}
if errEnd := source.AppendBytes([]byte("\n\n")); errEnd != nil {
log.WithError(errEnd).Warn("failed to append api request log terminator")
}
} else {
log.WithError(errWrite).Warn("failed to append api request log part")
if len(info.Body) > 0 {
builder.WriteString(string(info.Body))
} else {
builder.WriteString("<empty>")
}
builder.WriteString("\n\n")
requestText = builder.String()
}
} else {
if len(info.Body) > 0 {
builder.WriteString(string(info.Body))
} else {
builder.WriteString("<empty>")
}
builder.WriteString("\n\n")
requestText = builder.String()
}
attempt := &upstreamAttempt{
index: index,
request: requestText,
response: &strings.Builder{},
responseSource: apiResponseSourceOrNil(ginCtx),
}
attempts = append(attempts, attempt)
ginCtx.Set(apiAttemptsKey, attempts)
if requestText != "" {
updateAggregatedRequest(ginCtx, attempts)
}
}
func deferAPIRequest(ginCtx *gin.Context, info UpstreamRequestLog) {
if ginCtx == nil {
return
}
var requests []logging.DeferredAPIRequest
if value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey); exists {
requests, _ = value.([]logging.DeferredAPIRequest)
}
index := len(requests) + 1
capturedInfo := info
capturedAt := time.Now()
capturedBytes, _ := ginCtx.Get(deferredAPIRequestBytesKey)
bytesUsed, _ := capturedBytes.(int)
remaining := maxDeferredAPIRequestBodyBytes - bytesUsed
if remaining < 0 {
remaining = 0
}
captureLength := len(info.Body)
if captureLength > remaining {
captureLength = remaining
}
capturedInfo.Body = bytes.Clone(info.Body[:captureLength])
bodyEmpty := len(info.Body) == 0
bodyTruncated := captureLength < len(info.Body)
ginCtx.Set(deferredAPIRequestBytesKey, bytesUsed+captureLength)
requests = append(requests, func() []byte {
builder := newAPIRequestLogBuilder(index, capturedInfo, capturedAt)
if bodyEmpty {
builder.WriteString("<empty>")
} else {
builder.Write(capturedInfo.Body)
if bodyTruncated {
builder.WriteString(fmt.Sprintf("\n[API REQUEST BODY TRUNCATED: captured first %d bytes]", captureLength))
}
}
builder.WriteString("\n\n")
return []byte(builder.String())
})
ginCtx.Set(logging.DeferredAPIRequestContextKey, requests)
}
func newAPIRequestLogBuilder(index int, info UpstreamRequestLog, timestamp time.Time) *strings.Builder {
builder := &strings.Builder{}
builder.WriteString(fmt.Sprintf("=== API REQUEST %d ===\n", index))
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano)))
if info.URL != "" {
builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL))
} else {
builder.WriteString("Upstream URL: <unknown>\n")
}
if info.Method != "" {
builder.WriteString(fmt.Sprintf("HTTP Method: %s\n", info.Method))
}
if auth := formatAuthInfo(info); auth != "" {
builder.WriteString(fmt.Sprintf("Auth: %s\n", auth))
}
builder.WriteString("\nHeaders:\n")
writeHeaders(builder, info.Headers)
builder.WriteString("\nBody:\n")
return builder
}
// RecordAPIResponseMetadata captures upstream response status/header information for the latest attempt.
func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status int, headers http.Header) {
logging.SetResponseHeaders(ctx, headers)
if !requestLogCaptureEnabled(cfg) {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
attempts, attempt := ensureAttempt(ginCtx)
ensureResponseIntro(ginCtx, attempt)
if status > 0 && !attempt.statusWritten {
writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("Status: %d\n", status)))
attempt.statusWritten = true
}
if !attempt.headersWritten {
builder := &strings.Builder{}
builder.WriteString("Headers:\n")
writeHeaders(builder, headers)
writeAttemptResponse(ginCtx, attempt, []byte(builder.String()))
attempt.headersWritten = true
writeAttemptResponse(ginCtx, attempt, []byte("\n"))
}
updateAggregatedResponseIfMemoryBacked(ginCtx, attempts)
}
// RecordAPIResponseError adds an error entry for the latest attempt when no HTTP response is available.
func RecordAPIResponseError(ctx context.Context, cfg *config.Config, err error) {
if !requestLogCaptureEnabled(cfg) || err == nil {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
attempts, attempt := ensureAttempt(ginCtx)
ensureResponseIntro(ginCtx, attempt)
if attempt.bodyStarted && !attempt.bodyHasContent {
// Ensure body does not stay empty marker if error arrives first.
attempt.bodyStarted = false
}
if attempt.errorWritten {
writeAttemptResponse(ginCtx, attempt, []byte("\n"))
}
writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("Error: %s\n", err.Error())))
attempt.errorWritten = true
updateAggregatedResponseIfMemoryBacked(ginCtx, attempts)
}
// AppendAPIResponseChunk appends an upstream response chunk to Gin context for request logging.
func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byte) {
if !requestLogCaptureEnabled(cfg) {
return
}
data := bytes.TrimSpace(chunk)
if len(data) == 0 {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
attempts, attempt := ensureAttempt(ginCtx)
ensureResponseIntro(ginCtx, attempt)
if !attempt.headersWritten {
builder := &strings.Builder{}
builder.WriteString("Headers:\n")
writeHeaders(builder, nil)
writeAttemptResponse(ginCtx, attempt, []byte(builder.String()))
attempt.headersWritten = true
writeAttemptResponse(ginCtx, attempt, []byte("\n"))
}
if !attempt.bodyStarted {
writeAttemptResponse(ginCtx, attempt, []byte("Body:\n"))
attempt.bodyStarted = true
}
currentChunkIsSSEEvent := bytes.HasPrefix(data, []byte("event:"))
currentChunkIsSSEData := bytes.HasPrefix(data, []byte("data:"))
if attempt.bodyHasContent {
separator := "\n\n"
if attempt.prevWasSSEEvent && currentChunkIsSSEData {
separator = "\n"
}
writeAttemptResponse(ginCtx, attempt, []byte(separator))
}
writeAttemptResponse(ginCtx, attempt, data)
attempt.bodyHasContent = true
attempt.prevWasSSEEvent = currentChunkIsSSEEvent
updateAggregatedResponseIfMemoryBacked(ginCtx, attempts)
}
// RecordAPIWebsocketRequest stores an upstream websocket request event in Gin context.
func RecordAPIWebsocketRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) {
if !requestLogCaptureEnabled(cfg) {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
builder := &strings.Builder{}
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
builder.WriteString("Event: api.websocket.request\n")
if info.URL != "" {
builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL))
}
if auth := formatAuthInfo(info); auth != "" {
builder.WriteString(fmt.Sprintf("Auth: %s\n", auth))
}
builder.WriteString("Headers:\n")
writeHeaders(builder, info.Headers)
builder.WriteString("\nBody:\n")
if len(info.Body) > 0 {
builder.Write(info.Body)
} else {
builder.WriteString("<empty>")
}
builder.WriteString("\n")
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
}
// RecordAPIWebsocketHandshake stores the upstream websocket handshake response metadata.
func RecordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, status int, headers http.Header) {
logging.SetResponseHeaders(ctx, headers)
if !requestLogCaptureEnabled(cfg) {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
builder := &strings.Builder{}
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
builder.WriteString("Event: api.websocket.handshake\n")
if status > 0 {
builder.WriteString(fmt.Sprintf("Status: %d\n", status))
}
builder.WriteString("Headers:\n")
writeHeaders(builder, headers)
builder.WriteString("\n")
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
}
// RecordAPIWebsocketUpgradeRejection stores a rejected websocket upgrade as an HTTP attempt.
func RecordAPIWebsocketUpgradeRejection(ctx context.Context, cfg *config.Config, info UpstreamRequestLog, status int, headers http.Header, body []byte) {
logging.SetResponseHeaders(ctx, headers)
if !requestLogCaptureEnabled(cfg) {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
RecordAPIRequest(ctx, cfg, info)
RecordAPIResponseMetadata(ctx, cfg, status, headers)
AppendAPIResponseChunk(ctx, cfg, body)
}
// WebsocketUpgradeRequestURL converts a websocket URL back to its HTTP handshake URL for logging.
func WebsocketUpgradeRequestURL(rawURL string) string {
trimmedURL := strings.TrimSpace(rawURL)
if trimmedURL == "" {
return ""
}
parsed, err := url.Parse(trimmedURL)
if err != nil {
return trimmedURL
}
switch strings.ToLower(parsed.Scheme) {
case "ws":
parsed.Scheme = "http"
case "wss":
parsed.Scheme = "https"
}
return parsed.String()
}
// AppendAPIWebsocketResponse stores an upstream websocket response frame in Gin context.
func AppendAPIWebsocketResponse(ctx context.Context, cfg *config.Config, payload []byte) {
if !requestLogCaptureEnabled(cfg) {
return
}
data := bytes.TrimSpace(payload)
if len(data) == 0 {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
markAPIResponseTimestamp(ginCtx)
builder := &strings.Builder{}
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
builder.WriteString("Event: api.websocket.response\n")
builder.Write(data)
builder.WriteString("\n")
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
}
// RecordAPIWebsocketError stores an upstream websocket error event in Gin context.
func RecordAPIWebsocketError(ctx context.Context, cfg *config.Config, stage string, err error) {
if !requestLogCaptureEnabled(cfg) || err == nil {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
}
markAPIResponseTimestamp(ginCtx)
builder := &strings.Builder{}
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
builder.WriteString("Event: api.websocket.error\n")
if trimmed := strings.TrimSpace(stage); trimmed != "" {
builder.WriteString(fmt.Sprintf("Stage: %s\n", trimmed))
}
builder.WriteString(fmt.Sprintf("Error: %s\n", err.Error()))
appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))
}
func ginContextFrom(ctx context.Context) *gin.Context {
ginCtx, _ := ctx.Value("gin").(*gin.Context)
return ginCtx
}
func getAttempts(ginCtx *gin.Context) []*upstreamAttempt {
if ginCtx == nil {
return nil
}
if value, exists := ginCtx.Get(apiAttemptsKey); exists {
if attempts, ok := value.([]*upstreamAttempt); ok {
return attempts
}
}
return nil
}
func ensureAttempt(ginCtx *gin.Context) ([]*upstreamAttempt, *upstreamAttempt) {
attempts := getAttempts(ginCtx)
if len(attempts) == 0 {
attempt := &upstreamAttempt{
index: 1,
response: &strings.Builder{},
responseSource: apiResponseSourceOrNil(ginCtx),
}
if source, ok := apiRequestSource(ginCtx); ok {
if errWrite := source.AppendBytes([]byte("=== API REQUEST 1 ===\n<missing>\n\n")); errWrite != nil {
log.WithError(errWrite).Warn("failed to append missing api request log part")
attempt.request = "=== API REQUEST 1 ===\n<missing>\n\n"
}
} else {
attempt.request = "=== API REQUEST 1 ===\n<missing>\n\n"
}
attempts = []*upstreamAttempt{attempt}
ginCtx.Set(apiAttemptsKey, attempts)
if attempt.request != "" {
updateAggregatedRequest(ginCtx, attempts)
}
}
return attempts, attempts[len(attempts)-1]
}
func ensureResponseIntro(ginCtx *gin.Context, attempt *upstreamAttempt) {
if attempt == nil || attempt.response == nil || attempt.responseIntroWritten {
return
}
writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("=== API RESPONSE %d ===\n", attempt.index)))
writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano))))
writeAttemptResponse(ginCtx, attempt, []byte("\n"))
attempt.responseIntroWritten = true
}
func writeAttemptResponse(ginCtx *gin.Context, attempt *upstreamAttempt, payload []byte) {
if attempt == nil || len(payload) == 0 {
return
}
if attempt.responseSource == nil {
attempt.responseSource = apiResponseSourceOrNil(ginCtx)
}
if attempt.responseSource != nil {
if errWrite := attempt.responseSource.AppendBytes(payload); errWrite == nil {
if ginCtx != nil {
ginCtx.Set(logging.APIResponseCapturedContextKey, true)
}
return
} else {
log.WithError(errWrite).Warn("failed to append api response log part")
attempt.responseSource = nil
}
}
if attempt.response == nil {
attempt.response = &strings.Builder{}
}
attempt.response.Write(payload)
}
func updateAggregatedRequest(ginCtx *gin.Context, attempts []*upstreamAttempt) {
if ginCtx == nil {
return
}
var builder strings.Builder
for _, attempt := range attempts {
builder.WriteString(attempt.request)
}
ginCtx.Set(apiRequestKey, []byte(builder.String()))
}
func updateAggregatedResponseIfMemoryBacked(ginCtx *gin.Context, attempts []*upstreamAttempt) {
if apiResponseSourceOrNil(ginCtx) != nil {
return
}
updateAggregatedResponse(ginCtx, attempts)
}
func updateAggregatedResponse(ginCtx *gin.Context, attempts []*upstreamAttempt) {
if ginCtx == nil {
return
}
var builder strings.Builder
for idx, attempt := range attempts {
if attempt == nil || attempt.response == nil {
continue
}
responseText := attempt.response.String()
if responseText == "" {
continue
}
builder.WriteString(responseText)
if !strings.HasSuffix(responseText, "\n") {
builder.WriteString("\n")
}
if idx < len(attempts)-1 {
builder.WriteString("\n")
}
}
ginCtx.Set(apiResponseKey, []byte(builder.String()))
}
func apiRequestSource(ginCtx *gin.Context) (*logging.FileBodySource, bool) {
return fileBodySourceFromGin(ginCtx, logging.APIRequestSourceContextKey)
}
func apiResponseSourceOrNil(ginCtx *gin.Context) *logging.FileBodySource {
source, ok := fileBodySourceFromGin(ginCtx, logging.APIResponseSourceContextKey)
if !ok {
return nil
}
return source
}
func appendAPIWebsocketTimeline(ginCtx *gin.Context, chunk []byte) {
if ginCtx == nil {
return
}
data := bytes.TrimSpace(chunk)
if len(data) == 0 {
return
}
if source, ok := apiWebsocketTimelineSource(ginCtx); ok {
if errAppend := source.AppendPart(data); errAppend == nil {
return
} else {
log.WithError(errAppend).Warn("failed to append api websocket timeline log part")
}
}
if existing, exists := ginCtx.Get(apiWebsocketTimelineKey); exists {
if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 {
combined := make([]byte, 0, len(existingBytes)+len(data)+2)
combined = append(combined, existingBytes...)
if !bytes.HasSuffix(existingBytes, []byte("\n")) {
combined = append(combined, '\n')
}
combined = append(combined, '\n')
combined = append(combined, data...)
ginCtx.Set(apiWebsocketTimelineKey, combined)
return
}
}
ginCtx.Set(apiWebsocketTimelineKey, bytes.Clone(data))
}
func apiWebsocketTimelineSource(ginCtx *gin.Context) (*logging.FileBodySource, bool) {
return fileBodySourceFromGin(ginCtx, logging.APIWebsocketTimelineSourceContextKey)
}
func fileBodySourceFromGin(ginCtx *gin.Context, key string) (*logging.FileBodySource, bool) {
if ginCtx == nil {
return nil, false
}
value, exists := ginCtx.Get(key)
if !exists {
return nil, false
}
source, ok := value.(*logging.FileBodySource)
return source, ok && source != nil
}
func markAPIResponseTimestamp(ginCtx *gin.Context) {
if ginCtx == nil {
return
}
if _, exists := ginCtx.Get("API_RESPONSE_TIMESTAMP"); exists {
return
}
ginCtx.Set("API_RESPONSE_TIMESTAMP", time.Now())
}
func writeHeaders(builder *strings.Builder, headers http.Header) {
if builder == nil {
return
}
if len(headers) == 0 {
builder.WriteString("<none>\n")
return
}
keys := make([]string, 0, len(headers))
for key := range headers {
keys = append(keys, key)
}
sort.Strings(keys)
for _, key := range keys {
values := headers[key]
if len(values) == 0 {
builder.WriteString(fmt.Sprintf("%s:\n", key))
continue
}
for _, value := range values {
masked := util.MaskSensitiveHeaderValue(key, value)
builder.WriteString(fmt.Sprintf("%s: %s\n", key, masked))
}
}
}
func formatAuthInfo(info UpstreamRequestLog) string {
var parts []string
if trimmed := strings.TrimSpace(info.Provider); trimmed != "" {
parts = append(parts, fmt.Sprintf("provider=%s", trimmed))
}
if trimmed := strings.TrimSpace(info.AuthID); trimmed != "" {
parts = append(parts, fmt.Sprintf("auth_id=%s", trimmed))
}
if trimmed := strings.TrimSpace(info.AuthLabel); trimmed != "" {
parts = append(parts, fmt.Sprintf("label=%s", trimmed))
}
authType := strings.ToLower(strings.TrimSpace(info.AuthType))
authValue := strings.TrimSpace(info.AuthValue)
switch authType {
case "api_key":
if authValue != "" {
parts = append(parts, fmt.Sprintf("type=api_key value=%s", util.HideAPIKey(authValue)))
} else {
parts = append(parts, "type=api_key")
}
case "oauth":
parts = append(parts, "type=oauth")
default:
if authType != "" {
if authValue != "" {
parts = append(parts, fmt.Sprintf("type=%s value=%s", authType, authValue))
} else {
parts = append(parts, fmt.Sprintf("type=%s", authType))
}
}
}
return strings.Join(parts, ", ")
}
func SummarizeErrorBody(contentType string, body []byte) string {
isHTML := strings.Contains(strings.ToLower(contentType), "text/html")
if !isHTML {
trimmed := bytes.TrimSpace(bytes.ToLower(body))
if bytes.HasPrefix(trimmed, []byte("<!doctype html")) || bytes.HasPrefix(trimmed, []byte("<html")) {
isHTML = true
}
}
if isHTML {
if title := extractHTMLTitle(body); title != "" {
return title
}
return "[html body omitted]"
}
// Try to extract error message from JSON response
if message := extractJSONErrorMessage(body); message != "" {
return message
}
return string(body)
}
func extractHTMLTitle(body []byte) string {
lower := bytes.ToLower(body)
start := bytes.Index(lower, []byte("<title"))
if start == -1 {
return ""
}
gt := bytes.IndexByte(lower[start:], '>')
if gt == -1 {
return ""
}
start += gt + 1
end := bytes.Index(lower[start:], []byte("</title>"))
if end == -1 {
return ""
}
title := string(body[start : start+end])
title = html.UnescapeString(title)
title = strings.TrimSpace(title)
if title == "" {
return ""
}
return strings.Join(strings.Fields(title), " ")
}
// extractJSONErrorMessage attempts to extract error.message from JSON error responses
func extractJSONErrorMessage(body []byte) string {
result := gjson.GetBytes(body, "error.message")
if result.Exists() && result.String() != "" {
return result.String()
}
return ""
}
// logWithRequestID returns a logrus Entry with request_id field populated from context.
// If no request ID is found in context, it returns the standard logger.
func LogWithRequestID(ctx context.Context) *log.Entry {
if ctx == nil {
return log.NewEntry(log.StandardLogger())
}
requestID := logging.GetRequestID(ctx)
if requestID == "" {
return log.NewEntry(log.StandardLogger())
}
return log.WithField("request_id", requestID)
}
// MarkCreditsUsed flags the request as having used AI credits for billing.
func MarkCreditsUsed(ctx context.Context) {
ginCtx := ginContextFrom(ctx)
if ginCtx != nil {
ginCtx.Set(creditsUsedKey, true)
}
}
// CreditsUsed returns true if the request used AI credits.
func CreditsUsed(ctx context.Context) bool {
ginCtx := ginContextFrom(ctx)
if ginCtx != nil {
if val, exists := ginCtx.Get(creditsUsedKey); exists {
if b, ok := val.(bool); ok {
return b
}
}
}
return false
}

View file

@ -0,0 +1,55 @@
package helps
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
)
func TestRecordAPIRequestClonesDeferredBodyWhenRequestLogDisabled(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(recorder)
ctx := context.WithValue(context.Background(), "gin", ginCtx)
body := []byte(`{"model":"original"}`)
RecordAPIRequest(ctx, &config.Config{}, UpstreamRequestLog{
URL: "https://api.example.com/v1/responses",
Method: http.MethodPost,
Body: body,
})
body[10] = 'X'
value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey)
if !exists {
t.Fatal("deferred API request was not captured")
}
requests, ok := value.([]logging.DeferredAPIRequest)
if !ok || len(requests) != 1 {
t.Fatalf("deferred API requests = %#v, want one request", value)
}
captured := string(requests[0]())
if !strings.Contains(captured, `{"model":"original"}`) {
t.Fatalf("captured API request = %q, want original body", captured)
}
}
func TestRecordAPIResponseMetadataStoresHeadersWhenRequestLogDisabled(t *testing.T) {
ctx := logging.WithResponseHeadersHolder(context.Background())
headers := http.Header{}
headers.Add("X-Upstream-Request-Id", "upstream-req-1")
RecordAPIResponseMetadata(ctx, &config.Config{}, http.StatusOK, headers)
headers.Set("X-Upstream-Request-Id", "mutated")
got := logging.GetResponseHeaders(ctx)
if got.Get("X-Upstream-Request-Id") != "upstream-req-1" {
t.Fatalf("response header = %q, want %q", got.Get("X-Upstream-Request-Id"), "upstream-req-1")
}
}

View file

@ -0,0 +1,28 @@
package helps
import (
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
)
// APIKeyModelIsCompat reports whether the selected API-key model enables
// compatibility handling for Claude thinking blocks.
func APIKeyModelIsCompat(req cliproxyexecutor.Request) bool {
modelInfo, ok := cliproxyauth.ResolvedAPIKeyModelInfo(req)
return ok && modelInfo != nil && modelInfo.IsCompat
}
// ApplyRequestThinking preserves the registry lookup path unless the auth
// manager bound an exact configured API-key model definition to this attempt.
func ApplyRequestThinking(body []byte, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, fromFormat, toFormat, provider string) ([]byte, error) {
originalSource := opts.OriginalRequest
if len(originalSource) == 0 {
originalSource = req.Payload
}
summaryConfig := translatedRequestSummaryConfig(body, req.Payload, originalSource, req.Model, fromFormat, toFormat)
if modelInfo, ok := cliproxyauth.ResolvedAPIKeyModelInfo(req); ok {
return thinking.ApplyThinkingWithModelInfoAndSummary(body, originalSource, req.Model, fromFormat, toFormat, provider, modelInfo, summaryConfig)
}
return thinking.ApplyThinkingWithSummary(body, req.Model, fromFormat, toFormat, provider, summaryConfig)
}

View file

@ -0,0 +1,233 @@
package helps_test
import (
"context"
"net/http"
"testing"
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
helps "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/claude"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
"github.com/tidwall/gjson"
)
type configuredThinkingExecutor struct {
seenModel string
resolved bool
translateRequest bool
translatedBody []byte
}
func (*configuredThinkingExecutor) Identifier() string { return "claude" }
func (e *configuredThinkingExecutor) Execute(_ context.Context, _ *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
e.seenModel = req.Model
modelInfo, resolved := cliproxyauth.ResolvedAPIKeyModelInfo(req)
e.resolved = resolved && modelInfo != nil
body := []byte(`{"thinking":{"type":"adaptive"},"output_config":{"effort":"low"}}`)
if e.translateRequest {
body = sdktranslator.TranslateRequest(opts.SourceFormat, sdktranslator.FormatClaude, req.Model, req.Payload, opts.Stream)
e.translatedBody = append(e.translatedBody[:0], body...)
}
out, err := helps.ApplyRequestThinking(body, req, opts, opts.SourceFormat.String(), "claude", "claude")
return cliproxyexecutor.Response{Payload: out}, err
}
func (e *configuredThinkingExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
response, err := e.Execute(ctx, auth, req, opts)
if err != nil {
return nil, err
}
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
chunks <- cliproxyexecutor.StreamChunk{Payload: response.Payload}
close(chunks)
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
}
func (*configuredThinkingExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
return auth, nil
}
func (e *configuredThinkingExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
return e.Execute(ctx, auth, req, opts)
}
func (*configuredThinkingExecutor) HttpRequest(context.Context, *cliproxyauth.Auth, *http.Request) (*http.Response, error) {
return nil, nil
}
func TestApplyRequestThinkingUsesExactClaudeModeForSummaryOnlyRequest(t *testing.T) {
manager := cliproxyauth.NewManager(nil, nil, nil)
manager.SetConfig(&internalconfig.Config{
SDKConfig: internalconfig.SDKConfig{ForceModelPrefix: true},
ClaudeKey: []internalconfig.ClaudeKey{{
APIKey: "summary-selected-key",
Prefix: "summary-tenant",
Models: []internalconfig.ClaudeModel{{
Name: "summary-shared-upstream",
Alias: "summary-public-model",
Thinking: &registry.ThinkingSupport{
Min: 1024,
Max: 16000,
},
}},
}},
})
executor := &configuredThinkingExecutor{translateRequest: true}
manager.RegisterExecutor(executor)
auth := &cliproxyauth.Auth{
ID: "summary-selected-auth",
Provider: "claude",
Prefix: "summary-tenant",
Attributes: map[string]string{
cliproxyauth.AttributeAuthKind: cliproxyauth.AuthKindAPIKey,
cliproxyauth.AttributeAPIKey: "summary-selected-key",
cliproxyauth.AttributeSource: "config:claude[0]",
},
}
modelRegistry := registry.GetGlobalRegistry()
modelRegistry.RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{
ID: "summary-tenant/summary-public-model", Type: "claude",
}})
modelRegistry.RegisterClient("summary-unrelated-auth", auth.Provider, []*registry.ModelInfo{{
ID: "summary-shared-upstream", Type: "claude",
Thinking: &registry.ThinkingSupport{Levels: []string{"high"}},
}})
t.Cleanup(func() {
modelRegistry.UnregisterClient(auth.ID)
modelRegistry.UnregisterClient("summary-unrelated-auth")
})
if registered, errRegister := manager.Register(t.Context(), auth); errRegister != nil {
t.Fatalf("Register() error = %v", errRegister)
} else if registered == nil {
t.Fatal("Register() returned nil auth")
}
original := []byte(`{"model":"summary-tenant/summary-public-model","reasoning":{"summary":"auto"},"input":"hi"}`)
response, errExecute := manager.Execute(t.Context(), []string{"claude"}, cliproxyexecutor.Request{
Model: "summary-tenant/summary-public-model",
Payload: original,
Format: sdktranslator.FormatOpenAIResponse,
}, cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatOpenAIResponse,
OriginalRequest: original,
})
if errExecute != nil {
t.Fatalf("Execute() error = %v", errExecute)
}
if got := gjson.GetBytes(executor.translatedBody, "thinking.type").String(); got != "adaptive" {
t.Fatalf("pre-executor thinking.type = %q, want global adaptive trigger; body=%s", got, executor.translatedBody)
}
if got := gjson.GetBytes(response.Payload, "thinking.type").String(); got != "enabled" {
t.Fatalf("thinking.type = %q, want exact manual mode; body=%s", got, response.Payload)
}
if got := gjson.GetBytes(response.Payload, "thinking.budget_tokens").Int(); got != 1024 {
t.Fatalf("thinking.budget_tokens = %d, want exact minimum 1024; body=%s", got, response.Payload)
}
if got := gjson.GetBytes(response.Payload, "thinking.display").String(); got != "summarized" {
t.Fatalf("thinking.display = %q, want summarized; body=%s", got, response.Payload)
}
if gjson.GetBytes(response.Payload, "output_config.effort").Exists() {
t.Fatalf("manual thinking retained adaptive effort: %s", response.Payload)
}
}
func TestApplyRequestThinkingUsesSelectedPrefixedAPIKeyModel(t *testing.T) {
manager := cliproxyauth.NewManager(nil, nil, nil)
manager.SetConfig(&internalconfig.Config{
SDKConfig: internalconfig.SDKConfig{ForceModelPrefix: true},
ClaudeKey: []internalconfig.ClaudeKey{{
APIKey: "selected-key",
Prefix: "tenant",
Models: []internalconfig.ClaudeModel{{
Name: "shared-upstream", Alias: "public-model",
Thinking: &registry.ThinkingSupport{Levels: []string{"high"}},
}},
}},
})
executor := &configuredThinkingExecutor{}
manager.RegisterExecutor(executor)
auth := &cliproxyauth.Auth{
ID: "selected-auth",
Provider: "claude",
Prefix: "tenant",
Attributes: map[string]string{
cliproxyauth.AttributeAuthKind: cliproxyauth.AuthKindAPIKey,
cliproxyauth.AttributeAPIKey: "selected-key",
cliproxyauth.AttributeSource: "config:claude[0]",
},
}
modelRegistry := registry.GetGlobalRegistry()
modelRegistry.RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "tenant/public-model", Type: "claude"}})
modelRegistry.RegisterClient("unrelated-auth", auth.Provider, []*registry.ModelInfo{{
ID: "shared-upstream", Type: "claude",
Thinking: &registry.ThinkingSupport{Levels: []string{"max"}},
}})
t.Cleanup(func() {
modelRegistry.UnregisterClient(auth.ID)
modelRegistry.UnregisterClient("unrelated-auth")
})
ctx := t.Context()
registered, errRegister := manager.Register(ctx, auth)
if errRegister != nil {
t.Fatalf("Register() error = %v", errRegister)
}
if registered == nil {
t.Fatal("Register() returned nil auth")
}
original := []byte(`{"model":"tenant/public-model","reasoning_effort":"max","messages":[{"role":"user","content":"hello"}]}`)
req := cliproxyexecutor.Request{
Model: "tenant/public-model",
Payload: original,
Format: sdktranslator.FormatOpenAI,
}
opts := cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatOpenAI,
OriginalRequest: original,
}
assertResponse := func(path string, payload []byte) {
t.Helper()
if executor.seenModel != "shared-upstream" {
t.Fatalf("%s executor model = %q, want shared-upstream", path, executor.seenModel)
}
if !executor.resolved {
t.Fatalf("%s request did not receive selected model capabilities", path)
}
if got := gjson.GetBytes(payload, "output_config.effort").String(); got != "high" {
t.Fatalf("%s output effort = %q, want selected credential capability high; body=%s", path, got, payload)
}
}
response, errExecute := manager.Execute(ctx, []string{"claude"}, req, opts)
if errExecute != nil {
t.Fatalf("Execute() error = %v", errExecute)
}
assertResponse("execute", response.Payload)
countResponse, errCount := manager.ExecuteCount(ctx, []string{"claude"}, req, opts)
if errCount != nil {
t.Fatalf("ExecuteCount() error = %v", errCount)
}
assertResponse("count", countResponse.Payload)
streamResult, errStream := manager.ExecuteStream(ctx, []string{"claude"}, req, opts)
if errStream != nil {
t.Fatalf("ExecuteStream() error = %v", errStream)
}
var streamPayload []byte
for chunk := range streamResult.Chunks {
if chunk.Err != nil {
t.Fatalf("ExecuteStream() chunk error = %v", chunk.Err)
}
streamPayload = append(streamPayload, chunk.Payload...)
}
assertResponse("stream", streamPayload)
}

View file

@ -0,0 +1,162 @@
package helps
import (
"fmt"
"strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
const openAIToolResultImageOmittedText = "[image omitted: unsupported by upstream]"
// ShouldNormalizeOpenAIToolResultsForModel reports whether the selected model
// explicitly excludes image input through its input-modalities configuration.
func ShouldNormalizeOpenAIToolResultsForModel(compat *config.OpenAICompatibility, upstreamModel, requestedModel string) bool {
if compat == nil {
return false
}
if normalize, matched := openAICompatibilityModelExcludesImages(compat.Models, upstreamModel); matched {
return normalize
}
normalize, _ := openAICompatibilityModelExcludesImages(compat.Models, requestedModel)
return normalize
}
// NormalizeOpenAIToolResultsTextOnly converts tool message content to strings.
// Text parts are preserved and image parts are replaced with a short marker.
func NormalizeOpenAIToolResultsTextOnly(payload []byte) []byte {
messages := gjson.GetBytes(payload, "messages")
if !messages.Exists() || !messages.IsArray() {
return payload
}
out := payload
messageIndex := 0
messages.ForEach(func(_, message gjson.Result) bool {
if message.Get("role").String() == "tool" {
content := message.Get("content")
if content.Exists() && content.Type != gjson.String {
path := fmt.Sprintf("messages.%d.content", messageIndex)
if updated, errSet := sjson.SetBytes(out, path, flattenOpenAIToolResultContent(content)); errSet == nil {
out = updated
}
}
}
messageIndex++
return true
})
return out
}
func openAICompatibilityModelExcludesImages(models []config.OpenAICompatibilityModel, model string) (bool, bool) {
model = normalizeOpenAICompatibilityModelName(model)
if model == "" {
return false, false
}
for i := range models {
if strings.EqualFold(model, normalizeOpenAICompatibilityModelName(models[i].Name)) {
return inputModalitiesExcludeImages(models[i].InputModalities), true
}
}
matched := false
excludesImages := true
for i := range models {
if !strings.EqualFold(model, normalizeOpenAICompatibilityModelName(models[i].Alias)) {
continue
}
matched = true
if !inputModalitiesExcludeImages(models[i].InputModalities) {
excludesImages = false
}
}
return excludesImages && matched, matched
}
func inputModalitiesExcludeImages(modalities []string) bool {
if len(modalities) == 0 {
return false
}
hasText := false
for _, rawModality := range modalities {
switch strings.ToLower(strings.TrimSpace(rawModality)) {
case "image":
return false
case "text":
hasText = true
}
}
return hasText
}
func normalizeOpenAICompatibilityModelName(model string) string {
model = strings.TrimSpace(model)
if model == "" {
return ""
}
return strings.TrimSpace(thinking.ParseSuffix(model).ModelName)
}
func flattenOpenAIToolResultContent(content gjson.Result) string {
if content.Type == gjson.String {
return content.String()
}
if content.IsArray() {
parts := make([]string, 0, 4)
content.ForEach(func(_, item gjson.Result) bool {
if part, ok := openAIToolResultPartText(item); ok {
parts = append(parts, part)
}
return true
})
return strings.Join(parts, "\n\n")
}
if content.IsObject() {
if isOpenAIImageToolResultPart(content) {
return openAIToolResultImageOmittedText
}
if text := content.Get("text"); text.Type == gjson.String {
return text.String()
}
}
return content.Raw
}
func openAIToolResultPartText(item gjson.Result) (string, bool) {
if item.Type == gjson.String {
return item.String(), true
}
if item.IsObject() {
if isOpenAIImageToolResultPart(item) {
return openAIToolResultImageOmittedText, true
}
if text := item.Get("text"); text.Type == gjson.String {
return text.String(), true
}
}
if item.Raw == "" {
return "", false
}
return item.Raw, true
}
func isOpenAIImageToolResultPart(item gjson.Result) bool {
if !item.IsObject() {
return false
}
switch strings.ToLower(strings.TrimSpace(item.Get("type").String())) {
case "image", "image_url", "input_image":
return true
}
return item.Get("image_url").Exists() || item.Get("input_image").Exists()
}

View file

@ -0,0 +1,111 @@
package helps
import (
"testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/tidwall/gjson"
)
func TestNormalizeOpenAIToolResultsTextOnly(t *testing.T) {
input := []byte(`{"messages":[
{"role":"assistant","content":[{"type":"text","text":"before"}]},
{"role":"tool","tool_call_id":"call_1","content":[
{"type":"text","text":"image inspected"},
{"type":"image_url","image_url":{"url":"data:image/png;base64,AA=="}}
]},
{"role":"tool","tool_call_id":"call_2","content":"already text"},
{"role":"user","content":[{"type":"image_url","image_url":{"url":"https://example.com/user.png"}}]}
]}`)
got := NormalizeOpenAIToolResultsTextOnly(input)
toolContent := gjson.GetBytes(got, "messages.1.content")
if toolContent.Type != gjson.String {
t.Fatalf("tool content type = %s, want string", toolContent.Type)
}
if toolContent.String() != "image inspected\n\n"+openAIToolResultImageOmittedText {
t.Fatalf("tool content = %q", toolContent.String())
}
if gotContent := gjson.GetBytes(got, "messages.2.content"); gotContent.String() != "already text" {
t.Fatalf("existing string tool content = %q", gotContent.String())
}
if !gjson.GetBytes(got, "messages.0.content").IsArray() {
t.Fatal("assistant content array was unexpectedly changed")
}
if !gjson.GetBytes(got, "messages.3.content").IsArray() {
t.Fatal("non-tool content array was unexpectedly changed")
}
}
func TestNormalizeOpenAIToolResultsTextOnlyImageAndUnknownContent(t *testing.T) {
tests := []struct {
name string
input string
want string
}{
{
name: "image-only array",
input: `{"messages":[{"role":"tool","content":[{"type":"image_url","image_url":{"url":"https://example.com/image.png"}}]}]}`,
want: openAIToolResultImageOmittedText,
},
{
name: "image object",
input: `{"messages":[{"role":"tool","content":{"type":"image","source":{"type":"base64","data":"AA=="}}}]}`,
want: openAIToolResultImageOmittedText,
},
{
name: "unknown object",
input: `{"messages":[{"role":"tool","content":[{"type":"custom","value":1}]}]}`,
want: `{"type":"custom","value":1}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := NormalizeOpenAIToolResultsTextOnly([]byte(tt.input))
if content := gjson.GetBytes(got, "messages.0.content").String(); content != tt.want {
t.Fatalf("tool content = %q, want %q", content, tt.want)
}
})
}
}
func TestShouldNormalizeOpenAIToolResultsForModel(t *testing.T) {
compat := &config.OpenAICompatibility{Models: []config.OpenAICompatibilityModel{
{Name: "upstream-text", Alias: "alias-text", InputModalities: []string{"text"}},
{Name: "upstream-multimodal", Alias: "alias-multimodal", InputModalities: []string{"text", "image"}},
{Name: "upstream-unspecified", Alias: "alias-unspecified"},
{Name: "upstream-uppercase", Alias: "alias-uppercase", InputModalities: []string{"TEXT"}},
{Name: "pool-text", Alias: "shared-alias", InputModalities: []string{"text"}},
{Name: "pool-image", Alias: "shared-alias", InputModalities: []string{"text", "image"}},
}}
tests := []struct {
name string
upstreamModel string
requestedModel string
want bool
}{
{name: "upstream text", upstreamModel: "upstream-text", want: true},
{name: "upstream suffix", upstreamModel: "upstream-text(high)", want: true},
{name: "requested alias", upstreamModel: "unknown", requestedModel: "alias-text", want: true},
{name: "multimodal", upstreamModel: "upstream-multimodal", want: false},
{name: "unspecified", upstreamModel: "upstream-unspecified", want: false},
{name: "case insensitive modality", upstreamModel: "upstream-uppercase", want: true},
{name: "mixed alias pool", upstreamModel: "unknown", requestedModel: "shared-alias", want: false},
{name: "unknown", upstreamModel: "unknown", requestedModel: "missing", want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := ShouldNormalizeOpenAIToolResultsForModel(compat, tt.upstreamModel, tt.requestedModel); got != tt.want {
t.Fatalf("normalize = %t, want %t", got, tt.want)
}
})
}
if ShouldNormalizeOpenAIToolResultsForModel(nil, "upstream-text", "alias-text") {
t.Fatal("nil compatibility config unexpectedly enabled normalization")
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,340 @@
package helps
import (
"net/http"
"testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/tidwall/gjson"
)
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolsEntry(t *testing.T) {
cfg := &config.Config{
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
}
payload := []byte(`{"tools":[{"type":"image_generation","output_format":"png"},{"type":"function","name":"f1"}]}`)
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", "")
tools := gjson.GetBytes(out, "tools")
if !tools.Exists() || !tools.IsArray() {
t.Fatalf("expected tools array, got %v", tools.Type)
}
arr := tools.Array()
if len(arr) != 1 {
t.Fatalf("expected 1 tool after removal, got %d", len(arr))
}
if got := arr[0].Get("type").String(); got != "function" {
t.Fatalf("expected remaining tool type=function, got %q", got)
}
}
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolsEntryWithRoot(t *testing.T) {
cfg := &config.Config{
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
}
payload := []byte(`{"request":{"tools":[{"type":"image_generation"},{"type":"web_search"}]}}`)
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "antigravity", "request", payload, nil, "", "")
tools := gjson.GetBytes(out, "request.tools")
if !tools.Exists() || !tools.IsArray() {
t.Fatalf("expected request.tools array, got %v", tools.Type)
}
arr := tools.Array()
if len(arr) != 1 {
t.Fatalf("expected 1 tool after removal, got %d", len(arr))
}
if got := arr[0].Get("type").String(); got != "web_search" {
t.Fatalf("expected remaining tool type=web_search, got %q", got)
}
}
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolChoiceByType(t *testing.T) {
cfg := &config.Config{
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
}
payload := []byte(`{"tools":[{"type":"image_generation"},{"type":"function","name":"f1"}],"tool_choice":{"type":"image_generation"}}`)
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", "")
if gjson.GetBytes(out, "tool_choice").Exists() {
t.Fatalf("expected tool_choice to be removed")
}
}
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_RemovesToolChoiceByNameWithRoot(t *testing.T) {
cfg := &config.Config{
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
}
payload := []byte(`{"request":{"tools":[{"type":"image_generation"},{"type":"web_search"}],"tool_choice":{"type":"tool","name":"image_generation"}}}`)
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "antigravity", "request", payload, nil, "", "")
if gjson.GetBytes(out, "request.tool_choice").Exists() {
t.Fatalf("expected request.tool_choice to be removed")
}
}
func TestApplyPayloadConfigWithRoot_DisableImageGenerationChat_KeepsImageGenerationOnImagesEndpoints(t *testing.T) {
cfg := &config.Config{
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationChat},
}
payload := []byte(`{"tools":[{"type":"image_generation"},{"type":"function","name":"f1"}],"tool_choice":{"type":"image_generation"}}`)
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", "/v1/images/generations")
tools := gjson.GetBytes(out, "tools")
if !tools.Exists() || !tools.IsArray() {
t.Fatalf("expected tools array, got %v", tools.Type)
}
arr := tools.Array()
if len(arr) != 2 {
t.Fatalf("expected 2 tools (no removal), got %d", len(arr))
}
if !gjson.GetBytes(out, "tool_choice").Exists() {
t.Fatalf("expected tool_choice to be kept on images endpoint")
}
}
func TestApplyPayloadConfigWithRoot_DisableImageGenerationPassthrough_KeepsPayloadUnchanged(t *testing.T) {
cfg := &config.Config{
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationPassthrough},
}
payload := []byte(`{"tools":[{"type":"image_generation"},{"type":"function","name":"f1"}],"tool_choice":{"type":"image_generation"}}`)
// Passthrough must never inject or strip image_generation. The payload is forwarded as-is on
// non-images endpoints, and /v1/images/* endpoints behave like "chat" (also no removal).
for _, requestPath := range []string{"", "/v1/responses", "/v1/images/generations"} {
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", requestPath)
tools := gjson.GetBytes(out, "tools")
if !tools.Exists() || !tools.IsArray() {
t.Fatalf("path %q: expected tools array, got %v", requestPath, tools.Type)
}
if got := len(tools.Array()); got != 2 {
t.Fatalf("path %q: expected 2 tools (no removal), got %d", requestPath, got)
}
if got := tools.Array()[0].Get("type").String(); got != "image_generation" {
t.Fatalf("path %q: expected image_generation tool to be kept, got %q", requestPath, got)
}
if !gjson.GetBytes(out, "tool_choice").Exists() {
t.Fatalf("path %q: expected tool_choice to be kept", requestPath)
}
}
}
func TestApplyPayloadConfigWithRoot_DisableImageGeneration_PayloadOverrideCanRestoreImageGeneration(t *testing.T) {
cfg := &config.Config{
SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll},
Payload: config.PayloadConfig{
OverrideRaw: []config.PayloadRule{
{
Models: []config.PayloadModelRule{
{Name: "gpt-5.4", Protocol: "openai-response"},
},
Params: map[string]any{
"tools": `[{"type":"image_generation"},{"type":"function","name":"f1"}]`,
"tool_choice": `{"type":"image_generation"}`,
},
},
},
},
}
payload := []byte(`{"tools":[{"type":"image_generation"},{"type":"function","name":"f1"}],"tool_choice":{"type":"image_generation"}}`)
out := ApplyPayloadConfigWithRoot(cfg, "gpt-5.4", "openai-response", "", payload, nil, "", "")
tools := gjson.GetBytes(out, "tools")
if !tools.Exists() || !tools.IsArray() {
t.Fatalf("expected tools array, got %v", tools.Type)
}
arr := tools.Array()
if len(arr) != 2 {
t.Fatalf("expected 2 tools after payload override, got %d", len(arr))
}
if got := arr[0].Get("type").String(); got != "image_generation" {
t.Fatalf("expected first tool type=image_generation, got %q", got)
}
if !gjson.GetBytes(out, "tool_choice").Exists() {
t.Fatalf("expected tool_choice to be restored by payload override")
}
}
func TestApplyPayloadConfigWithRequest_HeaderGateRequiresWildcardMatch(t *testing.T) {
cfg := &config.Config{
Payload: config.PayloadConfig{
Override: []config.PayloadRule{
{
Models: []config.PayloadModelRule{
{
Name: "gpt-*",
Protocol: "openai",
Headers: map[string]string{
"X-Client-Tier": "tenant-*-region-*",
},
},
},
Params: map[string]any{
"metadata.enabled": true,
},
},
},
},
}
payload := []byte(`{"model":"gpt-5.4"}`)
headers := http.Header{}
headers.Set("X-Client-Tier", "tenant-alpha-region-us")
out := ApplyPayloadConfigWithRequest(cfg, "gpt-5.4", "openai", "responses", "", payload, nil, "", "", headers)
if !gjson.GetBytes(out, "metadata.enabled").Bool() {
t.Fatalf("expected header-matched payload rule to apply, payload=%s", string(out))
}
headers.Set("X-Client-Tier", "tenant-alpha")
out = ApplyPayloadConfigWithRequest(cfg, "gpt-5.4", "openai", "responses", "", payload, nil, "", "", headers)
if gjson.GetBytes(out, "metadata.enabled").Exists() {
t.Fatalf("expected header-mismatched payload rule to be skipped, payload=%s", string(out))
}
}
func TestApplyPayloadConfigWithRequest_FromProtocolGateUsesSourceProtocol(t *testing.T) {
cfg := &config.Config{
Payload: config.PayloadConfig{
Override: []config.PayloadRule{
{
Models: []config.PayloadModelRule{
{Name: "gpt-*", Protocol: "openai", FromProtocol: "responses"},
},
Params: map[string]any{
"metadata.source": "responses",
},
},
{
Models: []config.PayloadModelRule{
{Name: "gpt-*", Protocol: "openai", FromProtocol: "openai"},
},
Params: map[string]any{
"metadata.source": "openai",
},
},
},
},
}
payload := []byte(`{"model":"gpt-5.4"}`)
out := ApplyPayloadConfigWithRequest(cfg, "gpt-5.4", "openai", "openai-response", "", payload, nil, "", "", nil)
if got := gjson.GetBytes(out, "metadata.source").String(); got != "responses" {
t.Fatalf("metadata.source = %q, want responses; payload=%s", got, string(out))
}
out = ApplyPayloadConfigWithRequest(cfg, "gpt-5.4", "openai", "openai", "", payload, nil, "", "", nil)
if got := gjson.GetBytes(out, "metadata.source").String(); got != "openai" {
t.Fatalf("metadata.source = %q, want openai; payload=%s", got, string(out))
}
}
func TestApplyPayloadConfigWithRequest_PayloadConditionsNarrowRule(t *testing.T) {
cfg := &config.Config{
Payload: config.PayloadConfig{
Override: []config.PayloadRule{
{
Models: []config.PayloadModelRule{
{
Name: "gpt-*",
Match: []map[string]any{
{"metadata.client": "codex"},
{"tools.#(type==\"web_search\").enabled": true},
},
NotMatch: []map[string]any{
{"metadata.mode": "dev"},
},
Exist: []string{
"tools.#(type==\"web_search\").type",
},
NotExist: []string{
"metadata.missing",
"metadata.null_value",
},
},
},
Params: map[string]any{
"metadata.applied": true,
},
},
},
},
}
payload := []byte(`{"model":"gpt-5.4","metadata":{"client":"codex","mode":"prod","null_value":null},"tools":[{"type":"function"},{"type":"web_search","enabled":true}]}`)
out := ApplyPayloadConfigWithRequest(cfg, "gpt-5.4", "openai", "responses", "", payload, nil, "", "", nil)
if !gjson.GetBytes(out, "metadata.applied").Bool() {
t.Fatalf("expected payload condition-matched rule to apply, payload=%s", string(out))
}
}
func TestApplyPayloadConfigWithRequest_PayloadConditionsSkipRule(t *testing.T) {
testCases := []struct {
name string
model config.PayloadModelRule
}{
{
name: "match mismatch",
model: config.PayloadModelRule{
Name: "gpt-*",
Match: []map[string]any{{"metadata.client": "codex"}},
},
},
{
name: "not-match matched",
model: config.PayloadModelRule{
Name: "gpt-*",
NotMatch: []map[string]any{{"metadata.mode": "dev"}},
},
},
{
name: "exist missing",
model: config.PayloadModelRule{
Name: "gpt-*",
Exist: []string{"metadata.missing"},
},
},
{
name: "exist null",
model: config.PayloadModelRule{
Name: "gpt-*",
Exist: []string{"metadata.null_value"},
},
},
{
name: "not-exist present",
model: config.PayloadModelRule{
Name: "gpt-*",
NotExist: []string{"metadata.client"},
},
},
}
payload := []byte(`{"model":"gpt-5.4","metadata":{"client":"other","mode":"dev","null_value":null}}`)
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
cfg := &config.Config{
Payload: config.PayloadConfig{
Override: []config.PayloadRule{
{
Models: []config.PayloadModelRule{tc.model},
Params: map[string]any{
"metadata.applied": true,
},
},
},
},
}
out := ApplyPayloadConfigWithRequest(cfg, "gpt-5.4", "openai", "responses", "", payload, nil, "", "", nil)
if gjson.GetBytes(out, "metadata.applied").Exists() {
t.Fatalf("expected payload condition-mismatched rule to be skipped, payload=%s", string(out))
}
})
}
}

View file

@ -0,0 +1,81 @@
package helps
import (
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// SetStringIfDifferent updates path only when its value is not already the
// canonical JSON string. Values with another JSON type are still normalized.
func SetStringIfDifferent(payload []byte, path, value string) []byte {
current := gjson.GetBytes(payload, path)
if current.Type == gjson.String && current.String() == value {
return payload
}
updated, errSet := sjson.SetBytes(payload, path, value)
if errSet != nil {
return payload
}
return updated
}
// SetBoolIfDifferent updates path only when its value is not already the
// canonical JSON boolean. Values with another JSON type are still normalized.
func SetBoolIfDifferent(payload []byte, path string, value bool) []byte {
current := gjson.GetBytes(payload, path)
if (value && current.Type == gjson.True) || (!value && current.Type == gjson.False) {
return payload
}
updated, errSet := sjson.SetBytes(payload, path, value)
if errSet != nil {
return payload
}
return updated
}
// SetRawIfDifferent updates path only when the existing raw JSON is identical.
func SetRawIfDifferent(payload []byte, path string, value []byte) []byte {
current := gjson.GetBytes(payload, path)
if current.Exists() && len(current.Indexes) == 0 && current.Raw == string(value) {
return payload
}
updated, errSet := sjson.SetRawBytes(payload, path, value)
if errSet != nil {
return payload
}
return updated
}
// JoinRawJSONArray joins validated raw JSON array items without re-encoding them.
func JoinRawJSONArray(items [][]byte) []byte {
size := len(items) + 1
for _, item := range items {
size += len(item)
}
out := make([]byte, 0, size)
out = append(out, '[')
for index, item := range items {
if index > 0 {
out = append(out, ',')
}
out = append(out, item...)
}
return append(out, ']')
}
// JoinRawJSONStrings joins raw JSON array items held as strings.
func JoinRawJSONStrings(items []string) []byte {
size := len(items) + 1
for _, item := range items {
size += len(item)
}
out := make([]byte, 0, size)
out = append(out, '[')
for index, item := range items {
if index > 0 {
out = append(out, ',')
}
out = append(out, item...)
}
return append(out, ']')
}

View file

@ -0,0 +1,282 @@
package helps
import (
"bytes"
"encoding/json"
"strings"
"testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/tidwall/gjson"
)
type countingPayloadMarshaler struct {
calls *int
value string
}
func (m countingPayloadMarshaler) MarshalJSON() ([]byte, error) {
*m.calls = *m.calls + 1
return json.Marshal(m.value)
}
func TestSetStringIfDifferentReusesCanonicalValue(t *testing.T) {
input := []byte(`{"model":"gpt-test","messages":[]}`)
output := SetStringIfDifferent(input, "model", "gpt-test")
if &output[0] != &input[0] {
t.Fatal("canonical string caused a payload copy")
}
}
func TestSetStringIfDifferentNormalizesWrongType(t *testing.T) {
input := []byte(`{"model":123}`)
original := bytes.Clone(input)
output := SetStringIfDifferent(input, "model", "123")
model := gjson.GetBytes(output, "model")
if model.Type != gjson.String || model.String() != "123" {
t.Fatalf("model = %s, want string 123", model.Raw)
}
if !bytes.Equal(input, original) {
t.Fatal("input payload was modified in place")
}
}
func TestSetBoolIfDifferentReusesCanonicalValue(t *testing.T) {
input := []byte(`{"stream":true,"input":[]}`)
output := SetBoolIfDifferent(input, "stream", true)
if &output[0] != &input[0] {
t.Fatal("canonical boolean caused a payload copy")
}
}
func TestSetBoolIfDifferentNormalizesWrongType(t *testing.T) {
input := []byte(`{"stream":"true"}`)
output := SetBoolIfDifferent(input, "stream", true)
if stream := gjson.GetBytes(output, "stream"); stream.Type != gjson.True {
t.Fatalf("stream = %s, want boolean true", stream.Raw)
}
}
func TestSetRawIfDifferentReusesIdenticalRawValue(t *testing.T) {
input := []byte(`{"metadata":{"source":"executor"},"input":[]}`)
output := SetRawIfDifferent(input, "metadata", []byte(`{"source":"executor"}`))
if &output[0] != &input[0] {
t.Fatal("identical raw value caused a payload copy")
}
}
func TestSetRawIfDifferentUpdatesDifferentRawValue(t *testing.T) {
input := []byte(`{"metadata":"executor"}`)
output := SetRawIfDifferent(input, "metadata", []byte(`{"source":"executor"}`))
metadata := gjson.GetBytes(output, "metadata")
if !metadata.IsObject() || metadata.Get("source").String() != "executor" {
t.Fatalf("metadata = %s, want object", metadata.Raw)
}
}
func TestApplyPayloadConfigReusesCanonicalOverrides(t *testing.T) {
cfg := &config.Config{Payload: config.PayloadConfig{
Override: []config.PayloadRule{{
Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}},
Params: map[string]any{"stream": true, "model": "gpt-test"},
}},
OverrideRaw: []config.PayloadRule{{
Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}},
Params: map[string]any{"metadata": `{"source":"executor"}`},
}},
}}
input := []byte(`{"model":"gpt-test","stream":true,"metadata":{"source":"executor"},"messages":[]}`)
output := ApplyPayloadConfigWithRoot(cfg, "gpt-test", "openai", "", input, nil, "", "")
if &output[0] != &input[0] {
t.Fatal("canonical payload overrides caused a payload copy")
}
}
func TestApplyPayloadConfigWithRequestTrackedReportsContextManagementTouches(t *testing.T) {
const automatic = `{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]}`
modelRules := []config.PayloadModelRule{{Name: "claude-opus-5", Protocol: "claude"}}
originalWithoutContextManagement := []byte(`{"model":"claude-opus-5"}`)
for _, test := range []struct {
name string
payload string
original []byte
payloadConfig config.PayloadConfig
wantTouched bool
}{
{
name: "default",
payload: `{"model":"claude-opus-5"}`,
original: originalWithoutContextManagement,
payloadConfig: config.PayloadConfig{Default: []config.PayloadRule{{
Models: modelRules,
Params: map[string]any{"context_management": map[string]any{"edits": []any{map[string]any{"type": "default"}}}},
}}},
wantTouched: true,
},
{
name: "raw default",
payload: `{"model":"claude-opus-5"}`,
original: originalWithoutContextManagement,
payloadConfig: config.PayloadConfig{DefaultRaw: []config.PayloadRule{{
Models: modelRules,
Params: map[string]any{"context_management": `{"edits":[{"type":"raw_default"}]}`},
}}},
wantTouched: true,
},
{
name: "canonical descendant override",
payload: `{"model":"claude-opus-5","context_management":` + automatic + `}`,
payloadConfig: config.PayloadConfig{Override: []config.PayloadRule{{
Models: modelRules,
Params: map[string]any{"context_management.edits.0.keep": "all"},
}}},
wantTouched: true,
},
{
name: "identical raw override",
payload: `{"model":"claude-opus-5","context_management":` + automatic + `}`,
payloadConfig: config.PayloadConfig{OverrideRaw: []config.PayloadRule{{
Models: modelRules,
Params: map[string]any{"context_management": automatic},
}}},
wantTouched: true,
},
{
name: "filter already absent",
payload: `{"model":"claude-opus-5"}`,
payloadConfig: config.PayloadConfig{Filter: []config.PayloadFilterRule{{
Models: modelRules,
Params: []string{"context_management"},
}}},
wantTouched: true,
},
{
name: "unrelated override",
payload: `{"model":"claude-opus-5"}`,
payloadConfig: config.PayloadConfig{Override: []config.PayloadRule{{
Models: modelRules,
Params: map[string]any{"thinking.type": "enabled"},
}}},
},
{
name: "nonmatching override",
payload: `{"model":"claude-opus-5"}`,
payloadConfig: config.PayloadConfig{Override: []config.PayloadRule{{
Models: []config.PayloadModelRule{{Name: "other-model", Protocol: "claude"}},
Params: map[string]any{"context_management": map[string]any{"edits": []any{}}},
}}},
},
{
name: "default skipped for caller owned field",
payload: `{"model":"claude-opus-5","context_management":{"edits":[{"type":"caller"}]}}`,
original: []byte(`{"model":"claude-opus-5","context_management":{"edits":[{"type":"caller"}]}}`),
payloadConfig: config.PayloadConfig{Default: []config.PayloadRule{{
Models: modelRules,
Params: map[string]any{"context_management": map[string]any{"edits": []any{map[string]any{"type": "default"}}}},
}}},
},
} {
t.Run(test.name, func(t *testing.T) {
cfg := &config.Config{Payload: test.payloadConfig}
_, touched := ApplyPayloadConfigWithRequestTracked(cfg, "claude-opus-5", "claude", "claude", "", []byte(test.payload), test.original, "claude-opus-5", "", nil, "context_management")
if touched != test.wantTouched {
t.Fatalf("context_management touched = %t, want %t", touched, test.wantTouched)
}
})
}
}
func TestApplyPayloadConfigProjectionOverrideWritesEveryMatch(t *testing.T) {
cfg := &config.Config{Payload: config.PayloadConfig{
Override: []config.PayloadRule{{
Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}},
Params: map[string]any{"items.#.value": []any{1, 2}},
}},
}}
input := []byte(`{"items":[{"value":1},{"value":2}]}`)
output := ApplyPayloadConfigWithRoot(cfg, "gpt-test", "openai", "", input, nil, "", "")
for _, path := range []string{"items.0.value", "items.1.value"} {
if got := gjson.GetBytes(output, path).Raw; got != `[1,2]` {
t.Fatalf("%s = %s, want [1,2]", path, got)
}
}
}
func TestApplyPayloadConfigProjectionOverrideRawWritesEveryMatch(t *testing.T) {
cfg := &config.Config{Payload: config.PayloadConfig{
OverrideRaw: []config.PayloadRule{{
Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}},
Params: map[string]any{"items.#.value": `[1,2]`},
}},
}}
input := []byte(`{"items":[{"value":1},{"value":2}]}`)
output := ApplyPayloadConfigWithRoot(cfg, "gpt-test", "openai", "", input, nil, "", "")
for _, path := range []string{"items.0.value", "items.1.value"} {
if got := gjson.GetBytes(output, path).Raw; got != `[1,2]` {
t.Fatalf("%s = %s, want [1,2]", path, got)
}
}
}
func TestApplyPayloadConfigNormalizesByteSliceOverride(t *testing.T) {
cfg := &config.Config{Payload: config.PayloadConfig{
Override: []config.PayloadRule{{
Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}},
Params: map[string]any{"value": []byte("abc")},
}},
}}
input := []byte(`{"value":"YWJj"}`)
output := ApplyPayloadConfigWithRoot(cfg, "gpt-test", "openai", "", input, nil, "", "")
value := gjson.GetBytes(output, "value")
if value.Type != gjson.String || value.String() != "abc" {
t.Fatalf("value = %s, want string abc", value.Raw)
}
}
func TestSetPayloadValueIfDifferentUsesSJSONNumberEncoding(t *testing.T) {
input := []byte(`{"value":1.2}`)
output := setPayloadValueIfDifferent(input, "value", float32(1.2))
if got := gjson.GetBytes(output, "value").Raw; got != "1.2000000476837158" {
t.Fatalf("value = %s, want sjson float32 encoding", got)
}
canonical := []byte(`{"value":1.2000000476837158}`)
reused := setPayloadValueIfDifferent(canonical, "value", float32(1.2))
if &reused[0] != &canonical[0] {
t.Fatal("canonical float32 encoding caused a payload copy")
}
}
func TestSetPayloadValueIfDifferentCallsMarshalerOnce(t *testing.T) {
for _, input := range [][]byte{[]byte(`{"value":"old"}`), []byte(`{"value":"new"}`)} {
calls := 0
value := countingPayloadMarshaler{calls: &calls, value: "new"}
output := setPayloadValueIfDifferent(input, "value", value)
if calls != 1 {
t.Fatalf("MarshalJSON calls = %d, want 1", calls)
}
if got := gjson.GetBytes(output, "value").String(); got != "new" {
t.Fatalf("value = %q, want new", got)
}
}
}
func TestRemoveToolTypeReusesArrayWithoutMatch(t *testing.T) {
input := []byte(`{"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]}`)
output := removeToolTypeFromToolsArray(input, "tools", "image_generation")
if &output[0] != &input[0] {
t.Fatal("tool filtering without a match caused a payload copy")
}
}
var benchmarkPayloadMutationOutput []byte
func BenchmarkSetStringIfDifferentLargeCanonicalPayload(b *testing.B) {
input := []byte(`{"model":"gpt-test","messages":[{"role":"user","content":"` + strings.Repeat("x", 8<<20) + `"}]}`)
b.ReportAllocs()
b.SetBytes(int64(len(input)))
b.ResetTimer()
for b.Loop() {
benchmarkPayloadMutationOutput = SetStringIfDifferent(input, "model", "gpt-test")
}
}

View file

@ -0,0 +1,79 @@
package helps
import (
"context"
"net/http"
"strings"
"time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
"github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil"
log "github.com/sirupsen/logrus"
)
// NewProxyAwareHTTPClient creates an HTTP client with proper proxy configuration priority:
// 1. Use auth.ProxyURL if configured (highest priority)
// 2. Use cfg.ProxyURL if auth proxy is not configured
// 3. Use RoundTripper from context if neither are configured
//
// Parameters:
// - ctx: The context containing optional RoundTripper
// - cfg: The application configuration
// - auth: The authentication information
// - timeout: The client timeout (0 means no timeout)
//
// Returns:
// - *http.Client: An HTTP client with configured proxy or transport
func NewProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
httpClient := &http.Client{}
if timeout > 0 {
httpClient.Timeout = timeout
}
// Priority 1: Use auth.ProxyURL if configured
var proxyURL string
if auth != nil {
proxyURL = strings.TrimSpace(auth.ProxyURL)
}
// Priority 2: Use cfg.ProxyURL if auth proxy is not configured
if proxyURL == "" && cfg != nil {
proxyURL = strings.TrimSpace(cfg.ProxyURL)
}
// If we have a proxy URL configured, set up the transport
if proxyURL != "" {
transport := buildProxyTransport(proxyURL)
if transport != nil {
httpClient.Transport = transport
return httpClient
}
// If proxy setup failed, log and fall through to context RoundTripper
log.Debugf("failed to setup proxy from URL: %s, falling back to context transport", proxyutil.Redact(proxyURL))
}
// Priority 3: Use RoundTripper from context (typically from RoundTripperFor)
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
return httpClient
}
// buildProxyTransport creates an HTTP transport configured for the given proxy URL.
// It supports SOCKS5, HTTP, and HTTPS proxy protocols.
//
// Parameters:
// - proxyURL: The proxy URL string (e.g., "socks5://user:pass@host:port", "http://host:port")
//
// Returns:
// - *http.Transport: A configured transport, or nil if the proxy URL is invalid
func buildProxyTransport(proxyURL string) *http.Transport {
transport, _, errBuild := proxyutil.BuildHTTPTransport(proxyURL)
if errBuild != nil {
log.Errorf("%v", errBuild)
return nil
}
return transport
}

View file

@ -0,0 +1,30 @@
package helps
import (
"context"
"net/http"
"testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
)
func TestNewProxyAwareHTTPClientDirectBypassesGlobalProxy(t *testing.T) {
t.Parallel()
client := NewProxyAwareHTTPClient(
context.Background(),
&config.Config{SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"}},
&cliproxyauth.Auth{ProxyURL: "direct"},
0,
)
transport, ok := client.Transport.(*http.Transport)
if !ok {
t.Fatalf("transport type = %T, want *http.Transport", client.Transport)
}
if transport.Proxy != nil {
t.Fatal("expected direct transport to disable proxy function")
}
}

View file

@ -0,0 +1,108 @@
package helps
import (
"bytes"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// EnsureResponsesUsageDetails ensures that Responses usage objects contain output_tokens_details
// (defaulting reasoning_tokens to 0) and input_tokens_details (defaulting cached_tokens to 0).
// It supports plain JSON payloads, single-line SSE data: lines, and multi-line SSE frames (e.g. event: ...\ndata: ...).
func EnsureResponsesUsageDetails(payload []byte) []byte {
if len(payload) == 0 {
return payload
}
trimmed := bytes.TrimSpace(payload)
if len(trimmed) == 0 {
return payload
}
// 1. JSON-first: If trimmed payload starts with '{', process as a plain JSON object.
if trimmed[0] == '{' {
if gjson.GetBytes(trimmed, "object").String() == "response.compaction" {
return payload
}
updated := trimmed
updated = ensureUsageDetailsAt(updated, "response.usage")
updated = ensureUsageDetailsAt(updated, "usage")
if bytes.Equal(updated, trimmed) {
return payload
}
return updated
}
// 2. SSE frames: Scan lines for data: prefixed lines and patch their JSON payloads.
if bytes.Contains(payload, []byte("data:")) {
lines := bytes.Split(payload, []byte("\n"))
modified := false
for i, line := range lines {
trimmedLine := bytes.TrimSpace(line)
if !bytes.HasPrefix(trimmedLine, []byte("data:")) {
continue
}
prefixLen := len("data:")
if bytes.HasPrefix(line, []byte("data: ")) {
prefixLen = len("data: ")
} else if bytes.HasPrefix(line, []byte("data:")) {
prefixLen = len("data:")
}
dataPayload := bytes.TrimSpace(line[prefixLen:])
if len(dataPayload) == 0 || dataPayload[0] != '{' {
continue
}
if gjson.GetBytes(dataPayload, "object").String() == "response.compaction" {
continue
}
updated := dataPayload
updated = ensureUsageDetailsAt(updated, "response.usage")
updated = ensureUsageDetailsAt(updated, "usage")
if !bytes.Equal(updated, dataPayload) {
newPrefix := bytes.Clone(line[:prefixLen])
lines[i] = append(newPrefix, updated...)
modified = true
}
}
if modified {
return bytes.Join(lines, []byte("\n"))
}
return payload
}
return payload
}
func ensureUsageDetailsAt(jsonBody []byte, path string) []byte {
usageNode := gjson.GetBytes(jsonBody, path)
if !usageNode.Exists() || !usageNode.IsObject() {
return jsonBody
}
outputDetails := usageNode.Get("output_tokens_details")
if !outputDetails.Exists() {
jsonBody, _ = sjson.SetBytes(jsonBody, path+".output_tokens_details.reasoning_tokens", 0)
} else if outputDetails.Type == gjson.Null || !outputDetails.IsObject() {
jsonBody, _ = sjson.SetRawBytes(jsonBody, path+".output_tokens_details", []byte(`{"reasoning_tokens":0}`))
} else {
reasoning := outputDetails.Get("reasoning_tokens")
if !reasoning.Exists() || reasoning.Type == gjson.Null {
jsonBody, _ = sjson.SetBytes(jsonBody, path+".output_tokens_details.reasoning_tokens", 0)
}
}
inputDetails := usageNode.Get("input_tokens_details")
if !inputDetails.Exists() {
jsonBody, _ = sjson.SetBytes(jsonBody, path+".input_tokens_details.cached_tokens", 0)
} else if inputDetails.Type == gjson.Null || !inputDetails.IsObject() {
jsonBody, _ = sjson.SetRawBytes(jsonBody, path+".input_tokens_details", []byte(`{"cached_tokens":0}`))
} else {
cached := inputDetails.Get("cached_tokens")
if !cached.Exists() || cached.Type == gjson.Null {
jsonBody, _ = sjson.SetBytes(jsonBody, path+".input_tokens_details.cached_tokens", 0)
}
}
return jsonBody
}

View file

@ -0,0 +1,210 @@
package helps
import (
"bytes"
"context"
"testing"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
"github.com/tidwall/gjson"
)
func TestEnsureResponsesUsageDetails_NonStreamJSON(t *testing.T) {
raw := []byte(`{"id":"resp_1","object":"response","status":"completed","usage":{"input_tokens":84,"output_tokens":16,"total_tokens":100}}`)
got := EnsureResponsesUsageDetails(raw)
if !gjson.GetBytes(got, "usage.output_tokens_details").Exists() {
t.Fatalf("expected usage.output_tokens_details to exist, got %s", string(got))
}
if gjson.GetBytes(got, "usage.output_tokens_details.reasoning_tokens").Int() != 0 {
t.Fatalf("expected usage.output_tokens_details.reasoning_tokens == 0, got %d", gjson.GetBytes(got, "usage.output_tokens_details.reasoning_tokens").Int())
}
if !gjson.GetBytes(got, "usage.input_tokens_details").Exists() {
t.Fatalf("expected usage.input_tokens_details to exist, got %s", string(got))
}
if gjson.GetBytes(got, "usage.input_tokens_details.cached_tokens").Int() != 0 {
t.Fatalf("expected usage.input_tokens_details.cached_tokens == 0, got %d", gjson.GetBytes(got, "usage.input_tokens_details.cached_tokens").Int())
}
}
func TestEnsureResponsesUsageDetails_NonStreamJSONWithDataSubstring(t *testing.T) {
raw := []byte(`{"id":"resp_1","object":"response","status":"completed","output":[{"type":"message","content":[{"type":"text","text":"data:image/png;base64,iVBORw0KGgoAAAANSUhEUg"}]}],"usage":{"input_tokens":84,"output_tokens":16,"total_tokens":100}}`)
got := EnsureResponsesUsageDetails(raw)
if !gjson.GetBytes(got, "usage.output_tokens_details").Exists() {
t.Fatalf("expected usage.output_tokens_details to exist, got %s", string(got))
}
if gjson.GetBytes(got, "usage.output_tokens_details.reasoning_tokens").Int() != 0 {
t.Fatalf("expected usage.output_tokens_details.reasoning_tokens == 0, got %d", gjson.GetBytes(got, "usage.output_tokens_details.reasoning_tokens").Int())
}
if !gjson.GetBytes(got, "usage.input_tokens_details").Exists() {
t.Fatalf("expected usage.input_tokens_details to exist, got %s", string(got))
}
if gjson.GetBytes(got, "usage.input_tokens_details.cached_tokens").Int() != 0 {
t.Fatalf("expected usage.input_tokens_details.cached_tokens == 0, got %d", gjson.GetBytes(got, "usage.input_tokens_details.cached_tokens").Int())
}
}
func TestEnsureResponsesUsageDetails_SSEData(t *testing.T) {
raw := []byte(`data: {"type":"response.completed","response":{"id":"resp_1","usage":{"input_tokens":10,"output_tokens":4,"total_tokens":14}}}`)
got := EnsureResponsesUsageDetails(raw)
if !bytes.HasPrefix(got, []byte("data: ")) {
t.Fatalf("expected data: prefix preserved, got %s", string(got))
}
jsonBody := bytes.TrimPrefix(got, []byte("data: "))
if !gjson.GetBytes(jsonBody, "response.usage.output_tokens_details").Exists() {
t.Fatalf("expected response.usage.output_tokens_details to exist, got %s", string(got))
}
if gjson.GetBytes(jsonBody, "response.usage.output_tokens_details.reasoning_tokens").Int() != 0 {
t.Fatalf("expected reasoning_tokens == 0, got %d", gjson.GetBytes(jsonBody, "response.usage.output_tokens_details.reasoning_tokens").Int())
}
if !gjson.GetBytes(jsonBody, "response.usage.input_tokens_details").Exists() {
t.Fatalf("expected response.usage.input_tokens_details to exist, got %s", string(got))
}
if gjson.GetBytes(jsonBody, "response.usage.input_tokens_details.cached_tokens").Int() != 0 {
t.Fatalf("expected cached_tokens == 0, got %d", gjson.GetBytes(jsonBody, "response.usage.input_tokens_details.cached_tokens").Int())
}
}
func TestEnsureResponsesUsageDetails_SSEEventDataMultiLine(t *testing.T) {
raw := []byte("event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":84,\"output_tokens\":16,\"total_tokens\":100}}}\n\n")
got := EnsureResponsesUsageDetails(raw)
if !bytes.HasPrefix(got, []byte("event: response.completed\n")) {
t.Fatalf("expected event header preserved, got %s", string(got))
}
for _, line := range bytes.Split(got, []byte("\n")) {
if bytes.HasPrefix(line, []byte("data: ")) {
jsonBody := bytes.TrimPrefix(line, []byte("data: "))
if !gjson.GetBytes(jsonBody, "response.usage.output_tokens_details").Exists() {
t.Fatalf("expected response.usage.output_tokens_details to exist in multi-line frame, got %s", string(got))
}
if gjson.GetBytes(jsonBody, "response.usage.output_tokens_details.reasoning_tokens").Int() != 0 {
t.Fatalf("expected reasoning_tokens == 0, got %d", gjson.GetBytes(jsonBody, "response.usage.output_tokens_details.reasoning_tokens").Int())
}
if !gjson.GetBytes(jsonBody, "response.usage.input_tokens_details").Exists() {
t.Fatalf("expected response.usage.input_tokens_details to exist in multi-line frame, got %s", string(got))
}
if gjson.GetBytes(jsonBody, "response.usage.input_tokens_details.cached_tokens").Int() != 0 {
t.Fatalf("expected cached_tokens == 0, got %d", gjson.GetBytes(jsonBody, "response.usage.input_tokens_details.cached_tokens").Int())
}
}
}
}
func TestEnsureResponsesUsageDetails_PreservesExistingDetails(t *testing.T) {
raw := []byte(`data: {"type":"response.completed","response":{"id":"resp_1","usage":{"input_tokens":10,"input_tokens_details":{"cached_tokens":3},"output_tokens":4,"output_tokens_details":{"reasoning_tokens":2},"total_tokens":14}}}`)
got := EnsureResponsesUsageDetails(raw)
jsonBody := bytes.TrimPrefix(got, []byte("data: "))
if gjson.GetBytes(jsonBody, "response.usage.output_tokens_details.reasoning_tokens").Int() != 2 {
t.Fatalf("expected reasoning_tokens == 2, got %d", gjson.GetBytes(jsonBody, "response.usage.output_tokens_details.reasoning_tokens").Int())
}
if gjson.GetBytes(jsonBody, "response.usage.input_tokens_details.cached_tokens").Int() != 3 {
t.Fatalf("expected cached_tokens == 3, got %d", gjson.GetBytes(jsonBody, "response.usage.input_tokens_details.cached_tokens").Int())
}
}
func TestEnsureResponsesUsageDetails_HandlesNullOrEmptyDetails(t *testing.T) {
raw := []byte(`{"id":"resp_1","usage":{"input_tokens":10,"input_tokens_details":null,"output_tokens":4,"output_tokens_details":{},"total_tokens":14}}`)
got := EnsureResponsesUsageDetails(raw)
if gjson.GetBytes(got, "usage.output_tokens_details.reasoning_tokens").Int() != 0 {
t.Fatalf("expected reasoning_tokens == 0, got %d", gjson.GetBytes(got, "usage.output_tokens_details.reasoning_tokens").Int())
}
if gjson.GetBytes(got, "usage.input_tokens_details.cached_tokens").Int() != 0 {
t.Fatalf("expected cached_tokens == 0, got %d", gjson.GetBytes(got, "usage.input_tokens_details.cached_tokens").Int())
}
}
func TestEnsureResponsesUsageDetails_NonJSONAndDone(t *testing.T) {
cases := [][]byte{
[]byte("data: [DONE]"),
[]byte("[DONE]"),
[]byte(": keepalive"),
[]byte(""),
[]byte(`{"type":"response.output_item.added"}`),
}
for _, c := range cases {
got := EnsureResponsesUsageDetails(c)
if !bytes.Equal(got, c) {
t.Fatalf("expected unchanged for %q, got %q", string(c), string(got))
}
}
}
func TestTranslateStreamWithClaudeInputTokens_OpenAICompatTranslation_PatchesResponsesUsage(t *testing.T) {
ctx := context.Background()
reqBody := []byte(`{"model":"deepseek-v4-flash","input":"hi","stream":true}`)
translatedReq := []byte(`{"model":"deepseek-v4-flash","messages":[{"role":"user","content":"hi"}],"stream":true,"stream_options":{"include_usage":true}}`)
chunk1 := []byte(`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}`)
chunk2 := []byte(`data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":84,"completion_tokens":16,"total_tokens":100}}`)
chunk3 := []byte(`data: [DONE]`)
var param any
_ = TranslateStreamWithClaudeInputTokens(
ctx,
sdktranslator.FormatOpenAI,
sdktranslator.FormatOpenAIResponse,
"deepseek-v4-flash",
reqBody,
translatedReq,
chunk1,
&param,
nil,
)
chunks2 := TranslateStreamWithClaudeInputTokens(
ctx,
sdktranslator.FormatOpenAI,
sdktranslator.FormatOpenAIResponse,
"deepseek-v4-flash",
reqBody,
translatedReq,
chunk2,
&param,
nil,
)
chunks3 := TranslateStreamWithClaudeInputTokens(
ctx,
sdktranslator.FormatOpenAI,
sdktranslator.FormatOpenAIResponse,
"deepseek-v4-flash",
reqBody,
translatedReq,
chunk3,
&param,
nil,
)
allChunks := append(chunks2, chunks3...)
foundCompleted := false
for _, ch := range allChunks {
for _, line := range bytes.Split(ch, []byte("\n")) {
if bytes.HasPrefix(line, []byte("data: ")) {
payload := bytes.TrimPrefix(line, []byte("data: "))
if gjson.GetBytes(payload, "type").String() == "response.completed" {
foundCompleted = true
if !gjson.GetBytes(payload, "response.usage.output_tokens_details").Exists() {
t.Fatalf("expected output_tokens_details to exist on translated response.completed: %s", string(ch))
}
if gjson.GetBytes(payload, "response.usage.output_tokens_details.reasoning_tokens").Int() != 0 {
t.Fatalf("expected reasoning_tokens == 0, got %d", gjson.GetBytes(payload, "response.usage.output_tokens_details.reasoning_tokens").Int())
}
if !gjson.GetBytes(payload, "response.usage.input_tokens_details").Exists() {
t.Fatalf("expected input_tokens_details to exist on translated response.completed: %s", string(ch))
}
if gjson.GetBytes(payload, "response.usage.input_tokens_details.cached_tokens").Int() != 0 {
t.Fatalf("expected cached_tokens == 0, got %d", gjson.GetBytes(payload, "response.usage.input_tokens_details.cached_tokens").Int())
}
}
}
}
}
if !foundCompleted {
t.Fatalf("did not find response.completed chunk in stream translation output")
}
}

View file

@ -0,0 +1,148 @@
package helps
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"strings"
"sync"
"time"
"github.com/google/uuid"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
)
type sessionIDCacheEntry struct {
value string
expire time.Time
}
var (
sessionIDCache = make(map[string]sessionIDCacheEntry)
sessionIDCacheMu sync.RWMutex
sessionIDCacheCleanupOnce sync.Once
)
type claudeIDKVClient interface {
KVGet(ctx context.Context, key string) ([]byte, bool, error)
KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error)
KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error)
}
var currentClaudeIDKVClient = func() (claudeIDKVClient, bool, error) {
return homekv.CurrentKVClient()
}
const (
sessionIDTTL = time.Hour
sessionIDCacheCleanupPeriod = 15 * time.Minute
)
func startSessionIDCacheCleanup() {
go func() {
ticker := time.NewTicker(sessionIDCacheCleanupPeriod)
defer ticker.Stop()
for range ticker.C {
purgeExpiredSessionIDs()
}
}()
}
func purgeExpiredSessionIDs() {
now := time.Now()
sessionIDCacheMu.Lock()
for key, entry := range sessionIDCache {
if !entry.expire.After(now) {
delete(sessionIDCache, key)
}
}
sessionIDCacheMu.Unlock()
}
func sessionIDCacheKey(apiKey string) string {
sum := sha256.Sum256([]byte(apiKey))
return hex.EncodeToString(sum[:])
}
// CachedSessionID returns a stable session UUID per apiKey, refreshing the TTL on each access.
func CachedSessionID(apiKey string) string {
value, errValue := CachedSessionIDRequired(context.Background(), apiKey)
if errValue == nil && value != "" {
return value
}
return uuid.New().String()
}
// CachedSessionIDRequired returns a stable session UUID per apiKey for request-time paths.
func CachedSessionIDRequired(ctx context.Context, apiKey string) (string, error) {
if apiKey == "" {
return uuid.New().String(), nil
}
client, homeMode, errClient := currentClaudeIDKVClient()
if homeMode {
if errClient != nil {
return "", errClient
}
key := claudeSessionIDKVKey(apiKey)
raw, found, errGet := client.KVGet(ctx, key)
if errGet != nil {
return "", errGet
}
if found && strings.TrimSpace(string(raw)) != "" {
if _, errExpire := client.KVExpire(ctx, key, sessionIDTTL); errExpire != nil {
return "", errExpire
}
return strings.TrimSpace(string(raw)), nil
}
newID := uuid.New().String()
if _, errSet := client.KVSetNX(ctx, key, []byte(newID), sessionIDTTL); errSet != nil {
return "", errSet
}
raw, found, errGet = client.KVGet(ctx, key)
if errGet != nil {
return "", errGet
}
if found && strings.TrimSpace(string(raw)) != "" {
return strings.TrimSpace(string(raw)), nil
}
return "", fmt.Errorf("home kv session id missing after set")
}
sessionIDCacheCleanupOnce.Do(startSessionIDCacheCleanup)
key := sessionIDCacheKey(apiKey)
now := time.Now()
sessionIDCacheMu.RLock()
entry, ok := sessionIDCache[key]
valid := ok && entry.value != "" && entry.expire.After(now)
sessionIDCacheMu.RUnlock()
if valid {
sessionIDCacheMu.Lock()
entry = sessionIDCache[key]
if entry.value != "" && entry.expire.After(now) {
entry.expire = now.Add(sessionIDTTL)
sessionIDCache[key] = entry
sessionIDCacheMu.Unlock()
return entry.value, nil
}
sessionIDCacheMu.Unlock()
}
newID := uuid.New().String()
sessionIDCacheMu.Lock()
entry, ok = sessionIDCache[key]
if !ok || entry.value == "" || !entry.expire.After(now) {
entry.value = newID
}
entry.expire = now.Add(sessionIDTTL)
sessionIDCache[key] = entry
sessionIDCacheMu.Unlock()
return entry.value, nil
}
func claudeSessionIDKVKey(apiKey string) string {
return "cpa:claude:session-id:" + homekv.HashKeyPart(apiKey)
}

View file

@ -0,0 +1,178 @@
package helps
import (
"context"
"errors"
"testing"
"time"
"github.com/google/uuid"
)
func resetSessionIDCache() {
sessionIDCacheMu.Lock()
sessionIDCache = make(map[string]sessionIDCacheEntry)
sessionIDCacheMu.Unlock()
}
type fakeClaudeIDKVClient struct {
values map[string][]byte
getErr error
setErr error
expireErr error
setNoPersist bool
getCount int
setCount int
expireCount int
lastSetTTL time.Duration
lastExpireTTL time.Duration
}
func newFakeClaudeIDKVClient() *fakeClaudeIDKVClient {
return &fakeClaudeIDKVClient{values: make(map[string][]byte)}
}
func (c *fakeClaudeIDKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) {
c.getCount++
if c.getErr != nil {
return nil, false, c.getErr
}
value, ok := c.values[key]
if !ok {
return nil, false, nil
}
return append([]byte(nil), value...), true, nil
}
func (c *fakeClaudeIDKVClient) KVSetNX(_ context.Context, key string, value []byte, ttl time.Duration) (bool, error) {
c.setCount++
c.lastSetTTL = ttl
if c.setErr != nil {
return false, c.setErr
}
if _, ok := c.values[key]; ok {
return false, nil
}
if !c.setNoPersist {
c.values[key] = append([]byte(nil), value...)
}
return true, nil
}
func (c *fakeClaudeIDKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) {
c.expireCount++
c.lastExpireTTL = ttl
if c.expireErr != nil {
return false, c.expireErr
}
return true, nil
}
func useFakeClaudeIDKVClient(t *testing.T, client *fakeClaudeIDKVClient, homeMode bool, errClient error) {
t.Helper()
previous := currentClaudeIDKVClient
currentClaudeIDKVClient = func() (claudeIDKVClient, bool, error) {
return client, homeMode, errClient
}
t.Cleanup(func() {
currentClaudeIDKVClient = previous
})
}
func TestCachedSessionIDRequiredHomeReusesKVAcrossLocalCacheReset(t *testing.T) {
resetSessionIDCache()
client := newFakeClaudeIDKVClient()
useFakeClaudeIDKVClient(t, client, true, nil)
first, errFirst := CachedSessionIDRequired(context.Background(), "api-key-1")
if errFirst != nil {
t.Fatalf("CachedSessionIDRequired() first error = %v", errFirst)
}
resetSessionIDCache()
second, errSecond := CachedSessionIDRequired(context.Background(), "api-key-1")
if errSecond != nil {
t.Fatalf("CachedSessionIDRequired() second error = %v", errSecond)
}
if first != second {
t.Fatalf("session id = %q then %q, want same Home KV value", first, second)
}
if _, errParse := uuid.Parse(first); errParse != nil {
t.Fatalf("session id %q is not a UUID: %v", first, errParse)
}
if client.setCount != 1 {
t.Fatalf("KVSetNX count = %d, want 1", client.setCount)
}
if client.expireCount != 1 || client.lastExpireTTL != sessionIDTTL {
t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, sessionIDTTL)
}
if client.lastSetTTL != sessionIDTTL {
t.Fatalf("KVSetNX ttl = %v, want %v", client.lastSetTTL, sessionIDTTL)
}
}
func TestCachedSessionIDRequiredEmptyAPIKeyDoesNotUseHomeKV(t *testing.T) {
client := newFakeClaudeIDKVClient()
useFakeClaudeIDKVClient(t, client, true, nil)
value, errValue := CachedSessionIDRequired(context.Background(), "")
if errValue != nil {
t.Fatalf("CachedSessionIDRequired(empty) error = %v", errValue)
}
if _, errParse := uuid.Parse(value); errParse != nil {
t.Fatalf("session id %q is not a UUID: %v", value, errParse)
}
if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 {
t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount)
}
}
func TestCachedSessionIDRequiredHomeKVFailures(t *testing.T) {
for _, tc := range []struct {
name string
client *fakeClaudeIDKVClient
}{
{name: "get", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}},
{name: "set", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), setErr: errors.New("set failed")}},
{name: "expire", client: &fakeClaudeIDKVClient{values: map[string][]byte{
claudeSessionIDKVKey("api-key-1"): []byte(uuid.New().String()),
}, expireErr: errors.New("expire failed")}},
} {
t.Run(tc.name, func(t *testing.T) {
useFakeClaudeIDKVClient(t, tc.client, true, nil)
if _, errValue := CachedSessionIDRequired(context.Background(), "api-key-1"); errValue == nil {
t.Fatalf("CachedSessionIDRequired() error = nil, want error")
}
})
}
}
func TestCachedSessionIDRequiredHomeRequiresReadAfterSet(t *testing.T) {
client := newFakeClaudeIDKVClient()
client.setNoPersist = true
useFakeClaudeIDKVClient(t, client, true, nil)
if _, errValue := CachedSessionIDRequired(context.Background(), "api-key-1"); errValue == nil {
t.Fatalf("CachedSessionIDRequired() error = nil, want missing-after-set error")
}
}
func TestCachedSessionIDRequiredNonHomeModeUsesLocalMap(t *testing.T) {
resetSessionIDCache()
client := newFakeClaudeIDKVClient()
useFakeClaudeIDKVClient(t, client, false, nil)
first, errFirst := CachedSessionIDRequired(context.Background(), "api-key-1")
if errFirst != nil {
t.Fatalf("CachedSessionIDRequired() first error = %v", errFirst)
}
second, errSecond := CachedSessionIDRequired(context.Background(), "api-key-1")
if errSecond != nil {
t.Fatalf("CachedSessionIDRequired() second error = %v", errSecond)
}
if first != second {
t.Fatalf("session id = %q then %q, want local cache reuse", first, second)
}
if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 {
t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount)
}
}

View file

@ -0,0 +1,64 @@
package helps
import (
"strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
)
// ApplyThinkingWithSourcePayload preserves summary visibility from the original
// client payload while applying thinking configuration to its translated target
// payload. currentSourcePayload is the payload that was translated, while
// originalSourcePayload retains intent removed by an earlier interceptor.
func ApplyThinkingWithSourcePayload(body, currentSourcePayload, originalSourcePayload []byte, model, fromFormat, toFormat, providerKey string) ([]byte, error) {
summary := translatedRequestSummaryConfig(body, currentSourcePayload, originalSourcePayload, model, fromFormat, toFormat)
return thinking.ApplyThinkingWithSummary(body, model, fromFormat, toFormat, providerKey, summary)
}
// translatedRequestSummaryConfig gives the translated target payload precedence
// so a plugin request normalizer can remove or rewrite a canonical summary field.
// The original source is consulted only when the payload that was translated no
// longer carries the inbound intent, or when the target could not represent that
// intent until model-aware thinking is applied later (notably Claude).
func translatedRequestSummaryConfig(body, currentSourcePayload, originalSourcePayload []byte, model, fromFormat, toFormat string) thinking.SummaryConfig {
fromFormat = strings.ToLower(strings.TrimSpace(fromFormat))
toFormat = strings.ToLower(strings.TrimSpace(toFormat))
var targetSummary thinking.SummaryConfig
if fromFormat == toFormat {
targetSummary = thinking.ExtractSummaryConfig(body, toFormat)
} else {
targetSummary = thinking.ExtractExplicitSummaryConfig(body, toFormat)
}
if targetSummary.Mode != thinking.SummaryUnspecified {
return targetSummary
}
currentSummary := thinking.ExtractSummaryConfig(currentSourcePayload, fromFormat)
originalSummary := thinking.ExtractSummaryConfig(originalSourcePayload, fromFormat)
if currentSummary.Mode == thinking.SummaryUnspecified {
return originalSummary
}
from := sdktranslator.FromString(fromFormat)
to := sdktranslator.FromString(toFormat)
if !sdktranslator.HasRequestTransformer(from, to) {
// A missing translation must remain source-shaped. Same-format requests
// were handled by targetSummary above, including explicit native aliases.
return thinking.SummaryConfig{}
}
candidate := thinking.ApplySummaryConfigForModel(body, toFormat, model, currentSummary)
if thinking.ExtractExplicitSummaryConfig(candidate, toFormat).Mode != thinking.SummaryUnspecified {
// Registry translation applied this field before plugin normalization. If
// it is absent now but can be represented on the normalized body, the
// normalizer deliberately removed it and must remain authoritative.
return thinking.SummaryConfig{}
}
// Some intents cannot be represented until the final model-aware pass. For
// example, Claude display is invalid on disabled thinking, but a suffix can
// subsequently activate adaptive thinking. Preserve the source in that case.
return currentSummary
}

View file

@ -0,0 +1,12 @@
package helps
import (
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/antigravity"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/claude"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/codex"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/gemini"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/interactions"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/kimi"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/openai"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/xai"
)

View file

@ -0,0 +1,100 @@
package helps_test
import (
"context"
"testing"
helps "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking/provider/gemini"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
type summaryRemovingPluginHooks struct {
t *testing.T
}
func (h *summaryRemovingPluginHooks) NormalizeRequest(_ context.Context, _, _ sdktranslator.Format, _ string, body []byte, _ bool) []byte {
h.t.Helper()
const path = "generationConfig.thinkingConfig.includeThoughts"
if !gjson.GetBytes(body, path).Bool() {
h.t.Fatalf("request normalizer did not receive enabled summary: %s", body)
}
out, _ := sjson.DeleteBytes(body, path)
return out
}
func (*summaryRemovingPluginHooks) TranslateRequest(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, bool) ([]byte, bool) {
return nil, false
}
func (*summaryRemovingPluginHooks) NormalizeResponseBefore(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, []byte, []byte, bool) []byte {
return nil
}
func (*summaryRemovingPluginHooks) TranslateResponse(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, []byte, []byte, bool) ([]byte, bool) {
return nil, false
}
func (*summaryRemovingPluginHooks) NormalizeResponseAfter(context.Context, sdktranslator.Format, sdktranslator.Format, string, []byte, []byte, []byte, bool) []byte {
return nil
}
func TestApplyThinkingWithSourcePayloadPreservesNormalizerSummaryRemoval(t *testing.T) {
hooks := &summaryRemovingPluginHooks{t: t}
sdktranslator.SetPluginHooks(hooks)
t.Cleanup(func() { sdktranslator.SetPluginHooks(nil) })
source := []byte(`{"model":"gemini-3.6-flash","reasoning":{"effort":"high","summary":"auto"},"input":"hi"}`)
translated := sdktranslator.TranslateRequest(
sdktranslator.FormatOpenAIResponse,
sdktranslator.FormatGemini,
"gemini-3.6-flash",
source,
false,
)
const summaryPath = "generationConfig.thinkingConfig.includeThoughts"
if gjson.GetBytes(translated, summaryPath).Exists() {
t.Fatalf("request normalizer did not remove summary: %s", translated)
}
out, err := helps.ApplyThinkingWithSourcePayload(
translated,
source,
source,
"gemini-3.6-flash",
sdktranslator.FormatOpenAIResponse.String(),
sdktranslator.FormatGemini.String(),
"gemini",
)
if err != nil {
t.Fatalf("ApplyThinkingWithSourcePayload() error = %v", err)
}
if gjson.GetBytes(out, summaryPath).Exists() {
t.Fatalf("executor restored summary removed by request normalizer: %s", out)
}
}
func TestApplyThinkingWithSourcePayloadPreservesOriginalOnlySummary(t *testing.T) {
currentSource := []byte(`{"model":"gemini-3.6-flash","input":"hi"}`)
originalSource := []byte(`{"model":"gemini-3.6-flash","reasoning":{"summary":null},"input":"hi"}`)
body := []byte(`{"generationConfig":{"thinkingConfig":{"thinkingLevel":"high"}}}`)
out, err := helps.ApplyThinkingWithSourcePayload(
body,
currentSource,
originalSource,
"gemini-3.6-flash",
sdktranslator.FormatOpenAIResponse.String(),
sdktranslator.FormatGemini.String(),
"gemini",
)
if err != nil {
t.Fatalf("ApplyThinkingWithSourcePayload() error = %v", err)
}
if include := gjson.GetBytes(out, "generationConfig.thinkingConfig.includeThoughts"); !include.Exists() || include.Bool() {
t.Fatalf("original disabled summary was not preserved: %s", out)
}
}

View file

@ -0,0 +1,236 @@
package helps
import (
"fmt"
"strings"
"github.com/tidwall/gjson"
"github.com/tiktoken-go/tokenizer"
)
// TokenizerForModel returns a tokenizer codec suitable for an OpenAI-style model id.
func TokenizerForModel(model string) (tokenizer.Codec, error) {
sanitized := strings.ToLower(strings.TrimSpace(model))
switch {
case sanitized == "":
return tokenizer.Get(tokenizer.Cl100kBase)
case strings.HasPrefix(sanitized, "gpt-5"):
return tokenizer.ForModel(tokenizer.GPT5)
case strings.HasPrefix(sanitized, "gpt-5.1"):
return tokenizer.ForModel(tokenizer.GPT5)
case strings.HasPrefix(sanitized, "gpt-4.1"):
return tokenizer.ForModel(tokenizer.GPT41)
case strings.HasPrefix(sanitized, "gpt-4o"):
return tokenizer.ForModel(tokenizer.GPT4o)
case strings.HasPrefix(sanitized, "gpt-4"):
return tokenizer.ForModel(tokenizer.GPT4)
case strings.HasPrefix(sanitized, "gpt-3.5"), strings.HasPrefix(sanitized, "gpt-3"):
return tokenizer.ForModel(tokenizer.GPT35Turbo)
case strings.HasPrefix(sanitized, "o1"):
return tokenizer.ForModel(tokenizer.O1)
case strings.HasPrefix(sanitized, "o3"):
return tokenizer.ForModel(tokenizer.O3)
case strings.HasPrefix(sanitized, "o4"):
return tokenizer.ForModel(tokenizer.O4Mini)
default:
return tokenizer.Get(tokenizer.O200kBase)
}
}
// CountOpenAIChatTokens approximates prompt tokens for OpenAI chat completions payloads.
func CountOpenAIChatTokens(enc tokenizer.Codec, payload []byte) (int64, error) {
if enc == nil {
return 0, fmt.Errorf("encoder is nil")
}
if len(payload) == 0 {
return 0, nil
}
root := gjson.ParseBytes(payload)
segments := make([]string, 0, 32)
collectOpenAIMessages(root.Get("messages"), &segments)
collectOpenAITools(root.Get("tools"), &segments)
collectOpenAIFunctions(root.Get("functions"), &segments)
collectOpenAIToolChoice(root.Get("tool_choice"), &segments)
collectOpenAIResponseFormat(root.Get("response_format"), &segments)
addIfNotEmpty(&segments, root.Get("input").String())
addIfNotEmpty(&segments, root.Get("prompt").String())
joined := strings.TrimSpace(strings.Join(segments, "\n"))
if joined == "" {
return 0, nil
}
count, err := enc.Count(joined)
if err != nil {
return 0, err
}
return int64(count), nil
}
// BuildOpenAIUsageJSON returns a minimal usage structure understood by downstream translators.
func BuildOpenAIUsageJSON(count int64) []byte {
return []byte(fmt.Sprintf(`{"usage":{"prompt_tokens":%d,"completion_tokens":0,"total_tokens":%d}}`, count, count))
}
func collectOpenAIMessages(messages gjson.Result, segments *[]string) {
if !messages.Exists() || !messages.IsArray() {
return
}
messages.ForEach(func(_, message gjson.Result) bool {
addIfNotEmpty(segments, message.Get("role").String())
addIfNotEmpty(segments, message.Get("name").String())
collectOpenAIContent(message.Get("content"), segments)
collectOpenAIToolCalls(message.Get("tool_calls"), segments)
collectOpenAIFunctionCall(message.Get("function_call"), segments)
return true
})
}
func collectOpenAIContent(content gjson.Result, segments *[]string) {
if !content.Exists() {
return
}
if content.Type == gjson.String {
addIfNotEmpty(segments, content.String())
return
}
if content.IsArray() {
content.ForEach(func(_, part gjson.Result) bool {
partType := part.Get("type").String()
switch partType {
case "text", "input_text", "output_text":
addIfNotEmpty(segments, part.Get("text").String())
case "image_url":
addIfNotEmpty(segments, part.Get("image_url.url").String())
case "input_audio", "output_audio", "audio":
addIfNotEmpty(segments, part.Get("id").String())
case "tool_result":
addIfNotEmpty(segments, part.Get("name").String())
collectOpenAIContent(part.Get("content"), segments)
default:
if part.IsArray() {
collectOpenAIContent(part, segments)
return true
}
if part.Type == gjson.JSON {
addIfNotEmpty(segments, part.Raw)
return true
}
addIfNotEmpty(segments, part.String())
}
return true
})
return
}
if content.Type == gjson.JSON {
addIfNotEmpty(segments, content.Raw)
}
}
func collectOpenAIToolCalls(calls gjson.Result, segments *[]string) {
if !calls.Exists() || !calls.IsArray() {
return
}
calls.ForEach(func(_, call gjson.Result) bool {
addIfNotEmpty(segments, call.Get("id").String())
addIfNotEmpty(segments, call.Get("type").String())
function := call.Get("function")
if function.Exists() {
addIfNotEmpty(segments, function.Get("name").String())
addIfNotEmpty(segments, function.Get("description").String())
addIfNotEmpty(segments, function.Get("arguments").String())
if params := function.Get("parameters"); params.Exists() {
addIfNotEmpty(segments, params.Raw)
}
}
return true
})
}
func collectOpenAIFunctionCall(call gjson.Result, segments *[]string) {
if !call.Exists() {
return
}
addIfNotEmpty(segments, call.Get("name").String())
addIfNotEmpty(segments, call.Get("arguments").String())
}
func collectOpenAITools(tools gjson.Result, segments *[]string) {
if !tools.Exists() {
return
}
if tools.IsArray() {
tools.ForEach(func(_, tool gjson.Result) bool {
appendToolPayload(tool, segments)
return true
})
return
}
appendToolPayload(tools, segments)
}
func collectOpenAIFunctions(functions gjson.Result, segments *[]string) {
if !functions.Exists() || !functions.IsArray() {
return
}
functions.ForEach(func(_, function gjson.Result) bool {
addIfNotEmpty(segments, function.Get("name").String())
addIfNotEmpty(segments, function.Get("description").String())
if params := function.Get("parameters"); params.Exists() {
addIfNotEmpty(segments, params.Raw)
}
return true
})
}
func collectOpenAIToolChoice(choice gjson.Result, segments *[]string) {
if !choice.Exists() {
return
}
if choice.Type == gjson.String {
addIfNotEmpty(segments, choice.String())
return
}
addIfNotEmpty(segments, choice.Raw)
}
func collectOpenAIResponseFormat(format gjson.Result, segments *[]string) {
if !format.Exists() {
return
}
addIfNotEmpty(segments, format.Get("type").String())
addIfNotEmpty(segments, format.Get("name").String())
if schema := format.Get("json_schema"); schema.Exists() {
addIfNotEmpty(segments, schema.Raw)
}
if schema := format.Get("schema"); schema.Exists() {
addIfNotEmpty(segments, schema.Raw)
}
}
func appendToolPayload(tool gjson.Result, segments *[]string) {
if !tool.Exists() {
return
}
addIfNotEmpty(segments, tool.Get("type").String())
addIfNotEmpty(segments, tool.Get("name").String())
addIfNotEmpty(segments, tool.Get("description").String())
if function := tool.Get("function"); function.Exists() {
addIfNotEmpty(segments, function.Get("name").String())
addIfNotEmpty(segments, function.Get("description").String())
if params := function.Get("parameters"); params.Exists() {
addIfNotEmpty(segments, params.Raw)
}
}
}
func addIfNotEmpty(segments *[]string, value string) {
if segments == nil {
return
}
if trimmed := strings.TrimSpace(value); trimmed != "" {
*segments = append(*segments, trimmed)
}
}

View file

@ -0,0 +1,125 @@
package helps
import (
"container/list"
"errors"
"net/http"
"sync"
)
// DefaultTransportCacheCapacity bounds how many transports a TransportCache keeps
// alive at once. Every cached transport owns an independent connection pool, so an
// unbounded cache would let idle sockets and the goroutines managing them grow
// without limit whenever keys churn, for example when a credential's proxy is
// rotated through the management API or when an SDK embedder supplies a freshly
// built base transport per request.
const DefaultTransportCacheCapacity = 64
// TransportCache memoizes HTTP transports under a comparable key using a bounded
// LRU. Evicting an entry closes its idle connections so neither the pool nor its
// background goroutines outlive the cache entry.
//
// The key type is generic so callers can mix value identity (a normalized proxy
// URL) with pointer identity (a base transport supplied by the caller) without the
// cache retaining either beyond the LRU window.
type TransportCache[K comparable] struct {
mu sync.Mutex
capacity int
// order keeps the most recently used entry at the front.
order *list.List
items map[K]*list.Element
}
type transportCacheEntry[K comparable] struct {
key K
transport *http.Transport
}
// NewTransportCache returns a cache holding at most capacity transports. A
// non-positive capacity falls back to DefaultTransportCacheCapacity.
func NewTransportCache[K comparable](capacity int) *TransportCache[K] {
if capacity <= 0 {
capacity = DefaultTransportCacheCapacity
}
return &TransportCache[K]{
capacity: capacity,
order: list.New(),
items: make(map[K]*list.Element, capacity),
}
}
// Get returns the transport cached under key, calling build on the first use of
// that key. Concurrent callers observe the same instance.
//
// A build error is propagated without being cached, so a later call can retry and
// a failed lookup never occupies a cache slot. build must not call back into the
// same cache.
func (c *TransportCache[K]) Get(key K, build func() (*http.Transport, error)) (*http.Transport, error) {
if c == nil {
return nil, errors.New("transport cache: nil cache")
}
if build == nil {
return nil, errors.New("transport cache: nil build function")
}
c.mu.Lock()
defer c.mu.Unlock()
if element, ok := c.items[key]; ok {
c.order.MoveToFront(element)
return element.Value.(*transportCacheEntry[K]).transport, nil
}
transport, errBuild := build()
if errBuild != nil {
return nil, errBuild
}
if transport == nil {
return nil, errors.New("transport cache: build returned no transport")
}
c.items[key] = c.order.PushFront(&transportCacheEntry[K]{key: key, transport: transport})
c.evictLocked()
return transport, nil
}
// evictLocked drops least recently used entries until the cache fits its capacity.
// Closing idle connections is what actually releases the evicted pool; in-flight
// requests still holding the transport are unaffected because CloseIdleConnections
// only reaps connections that are currently idle.
func (c *TransportCache[K]) evictLocked() {
for c.order.Len() > c.capacity {
oldest := c.order.Back()
if oldest == nil {
return
}
c.order.Remove(oldest)
entry := oldest.Value.(*transportCacheEntry[K])
delete(c.items, entry.key)
entry.transport.CloseIdleConnections()
}
}
// Len reports how many transports the cache currently holds.
func (c *TransportCache[K]) Len() int {
if c == nil {
return 0
}
c.mu.Lock()
defer c.mu.Unlock()
return c.order.Len()
}
// Purge drops every entry and closes the idle connections it was holding.
func (c *TransportCache[K]) Purge() {
if c == nil {
return
}
c.mu.Lock()
defer c.mu.Unlock()
for element := c.order.Front(); element != nil; element = element.Next() {
element.Value.(*transportCacheEntry[K]).transport.CloseIdleConnections()
}
c.order.Init()
c.items = make(map[K]*list.Element, c.capacity)
}

View file

@ -0,0 +1,172 @@
package helps
import (
"errors"
"net/http"
"sync"
"testing"
)
type cacheKey struct {
scope string
proxy string
}
func TestTransportCacheReusesEntriesPerKey(t *testing.T) {
cache := NewTransportCache[cacheKey](8)
builds := 0
build := func() (*http.Transport, error) {
builds++
return &http.Transport{}, nil
}
first, errFirst := cache.Get(cacheKey{"auth-a", "p1"}, build)
if errFirst != nil {
t.Fatalf("Get() error = %v", errFirst)
}
second, errSecond := cache.Get(cacheKey{"auth-a", "p1"}, build)
if errSecond != nil {
t.Fatalf("Get() second error = %v", errSecond)
}
if first == nil || first != second {
t.Fatalf("expected one cached transport, got %p and %p", first, second)
}
if builds != 1 {
t.Fatalf("build called %d times, want 1", builds)
}
otherProxy, _ := cache.Get(cacheKey{"auth-a", "p2"}, build)
if otherProxy == first {
t.Fatal("distinct proxies must not share a transport")
}
otherScope, _ := cache.Get(cacheKey{"auth-b", "p1"}, build)
if otherScope == first {
t.Fatal("distinct credential scopes must not share a transport")
}
if got := cache.Len(); got != 3 {
t.Fatalf("cache Len() = %d, want 3", got)
}
}
// TestTransportCacheBoundsEntries is the regression test for unbounded pool growth:
// every cached transport owns a connection pool, so churning keys must evict.
func TestTransportCacheBoundsEntries(t *testing.T) {
const capacity = 4
cache := NewTransportCache[cacheKey](capacity)
for i := 0; i < 100; i++ {
key := cacheKey{"auth", string(rune('a' + i%97))}
if _, err := cache.Get(key, func() (*http.Transport, error) { return &http.Transport{}, nil }); err != nil {
t.Fatalf("Get() error = %v", err)
}
if got := cache.Len(); got > capacity {
t.Fatalf("cache grew to %d entries, want at most %d", got, capacity)
}
}
}
// TestTransportCacheEvictsLeastRecentlyUsed proves recency is honoured, so a hot
// credential is not evicted by a burst of one-off keys.
func TestTransportCacheEvictsLeastRecentlyUsed(t *testing.T) {
cache := NewTransportCache[cacheKey](2)
build := func() (*http.Transport, error) { return &http.Transport{}, nil }
hot, _ := cache.Get(cacheKey{"hot", ""}, build)
cache.Get(cacheKey{"cold", ""}, build)
// Touch hot so cold becomes the least recently used entry.
if again, _ := cache.Get(cacheKey{"hot", ""}, build); again != hot {
t.Fatal("expected the hot entry to still be cached")
}
cache.Get(cacheKey{"new", ""}, build)
if again, _ := cache.Get(cacheKey{"hot", ""}, build); again != hot {
t.Fatal("the most recently used entry must survive eviction")
}
}
// TestTransportCacheDoesNotCacheBuildFailures ensures a transient failure neither
// occupies a cache slot nor becomes permanent.
func TestTransportCacheDoesNotCacheBuildFailures(t *testing.T) {
cache := NewTransportCache[cacheKey](4)
key := cacheKey{"auth", "broken"}
if _, err := cache.Get(key, func() (*http.Transport, error) { return nil, errors.New("boom") }); err == nil {
t.Fatal("expected the build error to be propagated")
}
if got := cache.Len(); got != 0 {
t.Fatalf("a failed build must not occupy a cache slot, Len() = %d", got)
}
// A build returning (nil, nil) must be reported rather than cached as usable.
if _, err := cache.Get(key, func() (*http.Transport, error) { return nil, nil }); err == nil {
t.Fatal("expected an error when build returns no transport")
}
transport, err := cache.Get(key, func() (*http.Transport, error) { return &http.Transport{}, nil })
if err != nil || transport == nil {
t.Fatalf("retry after failure must succeed, got (%p, %v)", transport, err)
}
}
func TestTransportCacheConcurrentCallersShareOneInstance(t *testing.T) {
cache := NewTransportCache[cacheKey](8)
key := cacheKey{"auth-concurrent", "socks5://127.0.0.1:1080"}
const callers = 32
results := make([]*http.Transport, callers)
var wg sync.WaitGroup
wg.Add(callers)
for i := 0; i < callers; i++ {
go func(index int) {
defer wg.Done()
results[index], _ = cache.Get(key, func() (*http.Transport, error) { return &http.Transport{}, nil })
}(i)
}
wg.Wait()
for i := 1; i < callers; i++ {
if results[i] != results[0] {
t.Fatalf("caller %d observed a different transport (%p vs %p)", i, results[i], results[0])
}
}
}
func TestTransportCachePurgeAndNilSafety(t *testing.T) {
cache := NewTransportCache[cacheKey](4)
build := func() (*http.Transport, error) { return &http.Transport{}, nil }
cache.Get(cacheKey{"a", ""}, build)
cache.Get(cacheKey{"b", ""}, build)
if got := cache.Len(); got != 2 {
t.Fatalf("Len() = %d, want 2", got)
}
cache.Purge()
if got := cache.Len(); got != 0 {
t.Fatalf("Len() after Purge() = %d, want 0", got)
}
// The cache stays usable after a purge.
if transport, err := cache.Get(cacheKey{"a", ""}, build); err != nil || transport == nil {
t.Fatalf("Get() after Purge() = (%p, %v)", transport, err)
}
var nilCache *TransportCache[cacheKey]
if _, err := nilCache.Get(cacheKey{}, build); err == nil {
t.Fatal("expected an error from a nil cache")
}
if got := nilCache.Len(); got != 0 {
t.Fatalf("nil cache Len() = %d, want 0", got)
}
nilCache.Purge() // must not panic
if _, err := cache.Get(cacheKey{"nil-build", ""}, nil); err == nil {
t.Fatal("expected an error for a nil build function")
}
}
func TestNewTransportCacheDefaultsCapacity(t *testing.T) {
for _, capacity := range []int{0, -1} {
cache := NewTransportCache[cacheKey](capacity)
if cache.capacity != DefaultTransportCacheCapacity {
t.Fatalf("NewTransportCache(%d).capacity = %d, want %d", capacity, cache.capacity, DefaultTransportCacheCapacity)
}
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,753 @@
package helps
import (
"context"
"errors"
"io"
"net/http"
"net/url"
"strings"
"testing"
"time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror"
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
)
func TestParseOpenAIUsageChatCompletions(t *testing.T) {
data := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":6,"total_tokens":16,"prompt_tokens_details":{"cached_tokens":4},"completion_tokens_details":{"reasoning_tokens":5}}}`)
detail := ParseOpenAIUsage(data)
if detail.InputTokens != 10 {
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 10)
}
if detail.OutputTokens != 6 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 6)
}
if detail.TotalTokens != 16 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 16)
}
if detail.CachedTokens != 4 {
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 4)
}
if detail.CacheReadTokens != 4 {
t.Fatalf("cache read tokens = %d, want %d", detail.CacheReadTokens, 4)
}
if detail.ReasoningTokens != 5 {
t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 5)
}
if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
if detail.TokenBreakdown.Input.UncachedTokens != 6 || detail.TokenBreakdown.Output.NonReasoningTokens != 1 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseOpenAIUsageResponses(t *testing.T) {
data := []byte(`{"service_tier":"default","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30,"input_tokens_details":{"cached_tokens":7},"output_tokens_details":{"reasoning_tokens":9}}}`)
detail := ParseOpenAIUsage(data)
if detail.InputTokens != 10 {
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 10)
}
if detail.OutputTokens != 20 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 20)
}
if detail.TotalTokens != 30 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 30)
}
if detail.CachedTokens != 7 {
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 7)
}
if detail.CacheReadTokens != 7 {
t.Fatalf("cache read tokens = %d, want %d", detail.CacheReadTokens, 7)
}
if detail.ReasoningTokens != 9 {
t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 9)
}
if detail.ResponseServiceTier != "default" {
t.Fatalf("response service tier = %q, want default", detail.ResponseServiceTier)
}
if detail.TokenBreakdown.Input.UncachedTokens != 3 || detail.TokenBreakdown.Output.NonReasoningTokens != 11 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseOpenAIUsageTotalOnlyIsUnclassified(t *testing.T) {
detail := ParseOpenAIUsage([]byte(`{"usage":{"total_tokens":42}}`))
if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityUnclassified ||
detail.TotalTokens != 42 || detail.TokenBreakdown.UnclassifiedTokens != 42 {
t.Fatalf("detail = %+v", detail)
}
}
func TestParseOpenAIUsagePartialBucketsPreserveKnownTokens(t *testing.T) {
detail := ParseOpenAIUsage([]byte(`{"usage":{"input_tokens":10,"total_tokens":15}}`))
if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityUnclassified ||
detail.TokenBreakdown.Input.TotalTokens != 10 || detail.TokenBreakdown.UnclassifiedTokens != 5 {
t.Fatalf("detail = %+v", detail)
}
}
func TestParseOpenAIUsageExplicitZeroBucketsRemainInconsistent(t *testing.T) {
detail := ParseOpenAIUsage([]byte(`{"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":42}}`))
if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityInconsistent {
t.Fatalf("detail = %+v", detail)
}
}
func TestParseCodexUsageIncludesCacheWriteTokens(t *testing.T) {
data := []byte(`{"response":{"service_tier":"priority","usage":{"input_tokens":100,"output_tokens":20,"total_tokens":120,"input_tokens_details":{"cached_tokens":30,"cache_write_tokens":40}}}}`)
detail, ok := ParseCodexUsage(data)
if !ok {
t.Fatal("ParseCodexUsage() ok = false, want true")
}
if detail.InputTokens != 100 {
t.Fatalf("input tokens = %d, want 100", detail.InputTokens)
}
if detail.OutputTokens != 20 {
t.Fatalf("output tokens = %d, want 20", detail.OutputTokens)
}
if detail.CachedTokens != 30 {
t.Fatalf("cached tokens = %d, want 30", detail.CachedTokens)
}
if detail.CacheReadTokens != 30 {
t.Fatalf("cache read tokens = %d, want 30", detail.CacheReadTokens)
}
if detail.CacheCreationTokens != 40 {
t.Fatalf("cache creation tokens = %d, want 40", detail.CacheCreationTokens)
}
if detail.TotalTokens != 120 {
t.Fatalf("total tokens = %d, want 120", detail.TotalTokens)
}
if detail.ResponseServiceTier != "priority" {
t.Fatalf("response service tier = %q, want priority", detail.ResponseServiceTier)
}
if detail.TokenBreakdown.Input.UncachedTokens != 30 || detail.TokenBreakdown.Input.CacheWriteTokens != 40 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseOpenAIUsageNormalizesCacheCreationAlias(t *testing.T) {
data := []byte(`{"usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12,"input_tokens_details":{"cache_creation_tokens":4}}}`)
detail := ParseOpenAIUsage(data)
if detail.CacheCreationTokens != 4 {
t.Fatalf("cache creation tokens = %d, want 4", detail.CacheCreationTokens)
}
}
func TestParseOpenAIUsageIgnoresNullUsage(t *testing.T) {
data := []byte(`{"usage":null}`)
detail := ParseOpenAIUsage(data)
if detail != (usage.Detail{}) {
t.Fatalf("detail = %+v, want zero detail", detail)
}
}
func TestParseOpenAIUsagePreservesResponseTierWithoutUsage(t *testing.T) {
t.Parallel()
detail := ParseOpenAIUsage([]byte(`{"service_tier":"default"}`))
if detail.ResponseServiceTier != "default" {
t.Fatalf("response service tier = %q, want default", detail.ResponseServiceTier)
}
}
func TestParseCodexUsagePreservesResponseTierWithoutUsage(t *testing.T) {
t.Parallel()
detail, ok := ParseCodexUsage([]byte(`{"response":{"service_tier":"default"}}`))
if !ok || detail.ResponseServiceTier != "default" {
t.Fatalf("ParseCodexUsage() = (%+v, %v), want response tier default", detail, ok)
}
}
func TestParseOpenAIStreamUsageIgnoresNullUsage(t *testing.T) {
line := []byte(`data: {"id":"chunk_1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}],"usage":null}`)
if detail, ok := ParseOpenAIStreamUsage(line); ok {
t.Fatalf("ParseOpenAIStreamUsage() = (%+v, true), want false for null usage", detail)
}
}
func TestParseOpenAIStreamUsageResponsesFields(t *testing.T) {
line := []byte(`data: {"id":"chunk_1","object":"chat.completion.chunk","service_tier":"flex","choices":[],"usage":{"input_tokens":8,"output_tokens":5,"total_tokens":13,"input_tokens_details":{"cached_tokens":3},"output_tokens_details":{"reasoning_tokens":2}}}`)
detail, ok := ParseOpenAIStreamUsage(line)
if !ok {
t.Fatal("ParseOpenAIStreamUsage() ok = false, want true")
}
if detail.InputTokens != 8 {
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 8)
}
if detail.OutputTokens != 5 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 5)
}
if detail.TotalTokens != 13 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 13)
}
if detail.CachedTokens != 3 {
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 3)
}
if detail.CacheReadTokens != 3 {
t.Fatalf("cache read tokens = %d, want %d", detail.CacheReadTokens, 3)
}
if detail.ReasoningTokens != 2 {
t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 2)
}
if detail.ResponseServiceTier != "flex" {
t.Fatalf("response service tier = %q, want flex", detail.ResponseServiceTier)
}
}
func TestStreamUsageBufferKeepsLastUsage(t *testing.T) {
var buffer StreamUsageBuffer
buffer.Observe(usage.Detail{}, true)
buffer.Observe(usage.Detail{InputTokens: 1, OutputTokens: 1, TotalTokens: 2}, false)
buffer.Observe(usage.Detail{InputTokens: 39320, OutputTokens: 26, TotalTokens: 39346, CachedTokens: 33280}, true)
detail, ok := buffer.Detail()
if !ok {
t.Fatal("buffer detail ok = false, want true")
}
if detail.InputTokens != 39320 {
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 39320)
}
if detail.OutputTokens != 26 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 26)
}
if detail.TotalTokens != 39346 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 39346)
}
if detail.CachedTokens != 33280 {
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 33280)
}
}
func TestStreamUsageBufferPreservesTierAcrossChunks(t *testing.T) {
t.Parallel()
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"default"}`))
buffer.ObserveOpenAIStream([]byte(`data: {"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))
detail, ok := buffer.Detail()
if !ok {
t.Fatal("Detail() ok = false, want true")
}
if detail.InputTokens != 1 || detail.OutputTokens != 1 || detail.ResponseServiceTier != "default" {
t.Fatalf("detail = %+v, want usage with response tier default", detail)
}
}
func TestStreamUsageBufferObserveOpenAIStreamStateTransitions(t *testing.T) {
t.Parallel()
t.Run("same chunk", func(t *testing.T) {
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"flex","usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`))
detail, ok := buffer.Detail()
if !ok || detail.InputTokens != 2 || detail.ResponseServiceTier != "flex" {
t.Fatalf("detail = %+v ok=%v", detail, ok)
}
})
t.Run("usage before tier", func(t *testing.T) {
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream([]byte(`data: {"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`))
buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"default"}`))
detail, ok := buffer.Detail()
if !ok || detail.InputTokens != 2 || detail.ResponseServiceTier != "default" {
t.Fatalf("detail = %+v ok=%v", detail, ok)
}
})
t.Run("final usage tier overrides early tier", func(t *testing.T) {
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"default"}`))
buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"priority","usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`))
detail, ok := buffer.Detail()
if !ok || detail.ResponseServiceTier != "priority" {
t.Fatalf("detail = %+v ok=%v", detail, ok)
}
})
t.Run("irrelevant and invalid chunks do not change state", func(t *testing.T) {
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream([]byte(`data: {"content":"the word \"usage\" appears here"}`))
buffer.ObserveOpenAIStream([]byte(`data: {"usage":`))
buffer.ObserveOpenAIStream([]byte(`data: {"usage":null}`))
if detail, ok := buffer.Detail(); ok {
t.Fatalf("detail = %+v ok=true, want empty buffer", detail)
}
})
t.Run("zero token usage is retained", func(t *testing.T) {
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream([]byte(`data: {"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}`))
if _, ok := buffer.Detail(); !ok {
t.Fatal("Detail() ok = false, want true")
}
})
}
func TestStreamUsageBufferPreservesOnlyZeroUsage(t *testing.T) {
var buffer StreamUsageBuffer
buffer.Observe(usage.Detail{}, true)
detail, ok := buffer.Detail()
if !ok {
t.Fatal("buffer detail ok = false, want true")
}
if detail != (usage.Detail{}) {
t.Fatalf("detail = %+v, want zero detail", detail)
}
}
func TestParseClaudeUsageIncludesCacheTokensInTotal(t *testing.T) {
data := []byte(`{"usage":{"input_tokens":3085,"output_tokens":253,"cache_read_input_tokens":7,"cache_creation_input_tokens":19514}}`)
detail := ParseClaudeUsage(data)
if detail.InputTokens != 3085 {
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 3085)
}
if detail.OutputTokens != 253 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 253)
}
if detail.CacheReadTokens != 7 {
t.Fatalf("cache read tokens = %d, want %d", detail.CacheReadTokens, 7)
}
if detail.CacheCreationTokens != 19514 {
t.Fatalf("cache creation tokens = %d, want %d", detail.CacheCreationTokens, 19514)
}
if detail.CachedTokens != 7 {
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 7)
}
if detail.TotalTokens != 22859 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 22859)
}
if detail.TokenBreakdown.Input.TotalTokens != 22606 || detail.TokenBreakdown.Input.UncachedTokens != 3085 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseClaudeUsageFallsBackCachedTokensToCacheCreation(t *testing.T) {
data := []byte(`{"usage":{"input_tokens":3085,"output_tokens":253,"cache_creation_input_tokens":19514}}`)
detail := ParseClaudeUsage(data)
if detail.CachedTokens != 19514 {
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 19514)
}
if detail.TotalTokens != 22852 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 22852)
}
}
func TestParseClaudeUsagePreservesThinkingTokensAsReasoningSubset(t *testing.T) {
// Sanitized shape from local Anthropic request logs under ~/.config/cpa/logs.
data := []byte(`{"usage":{"input_tokens":2,"cache_creation_input_tokens":831,"cache_read_input_tokens":44225,"output_tokens":244,"output_tokens_details":{"thinking_tokens":40}}}`)
detail := ParseClaudeUsage(data)
if detail.OutputTokens != 244 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 244)
}
if detail.ReasoningTokens != 40 {
t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 40)
}
if detail.TotalTokens != 45302 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 45302)
}
if !detail.TokenBreakdown.Valid() ||
detail.TokenBreakdown.Output.TotalTokens != 244 ||
detail.TokenBreakdown.Output.NonReasoningTokens != 204 ||
detail.TokenBreakdown.Output.ReasoningTokens != 40 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseClaudeStreamUsagePreservesThinkingTokensAsReasoningSubset(t *testing.T) {
line := []byte(`data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":2,"cache_creation_input_tokens":831,"cache_read_input_tokens":44225,"output_tokens":244,"output_tokens_details":{"thinking_tokens":40}}}`)
detail, ok := ParseClaudeStreamUsage(line)
if !ok {
t.Fatal("expected stream usage to parse")
}
if detail.OutputTokens != 244 || detail.ReasoningTokens != 40 || detail.TotalTokens != 45302 {
t.Fatalf("stream usage detail = %+v", detail)
}
if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Output.NonReasoningTokens != 204 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseClaudeUsageFallsBackToTopLevelThinkingTokens(t *testing.T) {
data := []byte(`{"usage":{"input_tokens":3,"output_tokens":10,"thinking_tokens":4}}`)
detail := ParseClaudeUsage(data)
if detail.OutputTokens != 10 || detail.ReasoningTokens != 4 || detail.TotalTokens != 13 {
t.Fatalf("detail = %+v", detail)
}
if detail.TokenBreakdown.Output.NonReasoningTokens != 6 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseGeminiUsageNormalizesCachedContent(t *testing.T) {
detail := ParseGeminiUsage([]byte(`{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"cachedContentTokenCount":4,"totalTokenCount":12}}`))
if detail.CachedTokens != 4 {
t.Fatalf("cached tokens = %d, want 4", detail.CachedTokens)
}
if detail.CacheReadTokens != 4 {
t.Fatalf("cache read tokens = %d, want 4", detail.CacheReadTokens)
}
if detail.TokenBreakdown.Input.UncachedTokens != 6 || detail.TokenBreakdown.TotalTokens != 12 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseGeminiUsageIncludesToolUsePromptTokens(t *testing.T) {
detail := ParseGeminiUsage([]byte(`{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"thoughtsTokenCount":3,"toolUsePromptTokenCount":5,"totalTokenCount":20}}`))
if detail.InputTokens != 15 || detail.TotalTokens != 20 {
t.Fatalf("detail = %+v", detail)
}
if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete ||
detail.TokenBreakdown.Input.UncachedTokens != 15 || detail.TokenBreakdown.Output.ReasoningTokens != 3 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseGeminiStreamUsageSkipsZeroPlaceholder(t *testing.T) {
lines := [][]byte{
[]byte(`data: {"usageMetadata":{"promptTokenCount":0,"candidatesTokenCount":0,"thoughtsTokenCount":0,"totalTokenCount":0}}`),
[]byte(`data: {"usageMetadata":{"promptTokenCount":17984,"candidatesTokenCount":2668,"thoughtsTokenCount":1028,"totalTokenCount":21680}}`),
}
accepted := make([]usage.Detail, 0, len(lines))
for _, line := range lines {
detail, ok := ParseGeminiStreamUsage(line)
if ok {
accepted = append(accepted, detail)
}
}
if len(accepted) != 1 {
t.Fatalf("accepted usage count = %d, want 1", len(accepted))
}
detail := accepted[0]
if detail.InputTokens != 17984 || detail.OutputTokens != 2668 || detail.ReasoningTokens != 1028 || detail.TotalTokens != 21680 {
t.Fatalf("accepted usage detail = %+v", detail)
}
}
func TestParseGeminiUsageRejectsInvalidToolUseSums(t *testing.T) {
tests := map[string]string{
"negative": `{"usageMetadata":{"promptTokenCount":10,"toolUsePromptTokenCount":-1,"totalTokenCount":10}}`,
"overflow": `{"usageMetadata":{"promptTokenCount":9223372036854775807,"toolUsePromptTokenCount":1,"totalTokenCount":9223372036854775807}}`,
}
for name, payload := range tests {
t.Run(name, func(t *testing.T) {
detail := ParseGeminiUsage([]byte(payload))
if detail.InputTokens < 0 || !detail.TokenBreakdown.Valid() ||
detail.TokenBreakdown.Quality != usage.TokenAccountingQualityInconsistent {
t.Fatalf("detail = %+v", detail)
}
})
}
}
func TestParseInteractionsUsage(t *testing.T) {
detail := ParseInteractionsUsage([]byte(`{"usage":{"input_tokens":3,"output_tokens":4,"reasoning_tokens":5,"cached_tokens":2}}`))
if detail.InputTokens != 3 {
t.Fatalf("input tokens = %d, want 3", detail.InputTokens)
}
if detail.OutputTokens != 4 {
t.Fatalf("output tokens = %d, want 4", detail.OutputTokens)
}
if detail.ReasoningTokens != 5 {
t.Fatalf("reasoning tokens = %d, want 5", detail.ReasoningTokens)
}
if detail.TotalTokens != 12 {
t.Fatalf("total tokens = %d, want 12", detail.TotalTokens)
}
if detail.CachedTokens != 2 {
t.Fatalf("cached tokens = %d, want 2", detail.CachedTokens)
}
if detail.CacheReadTokens != 2 {
t.Fatalf("cache read tokens = %d, want 2", detail.CacheReadTokens)
}
if detail.TokenBreakdown.Input.UncachedTokens != 1 || detail.TokenBreakdown.Output.TotalTokens != 9 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestNormalizeUsageDetailTotalDoesNotDoubleCountReasoning(t *testing.T) {
detail := normalizeUsageDetailTotal(usage.Detail{
InputTokens: 100,
OutputTokens: 30,
ReasoningTokens: 12,
}, "openai", "")
if detail.TotalTokens != 130 {
t.Fatalf("total tokens = %d, want 130", detail.TotalTokens)
}
if detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete || detail.TokenBreakdown.Output.ReasoningTokens != 12 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseInteractionsUsageNormalizesCacheWriteAlias(t *testing.T) {
detail := ParseInteractionsUsage([]byte(`{"usage":{"input_tokens":3,"cache_write_tokens":2}}`))
if detail.CacheCreationTokens != 2 {
t.Fatalf("cache creation tokens = %d, want 2", detail.CacheCreationTokens)
}
}
func TestParseInteractionsUsageIncludesToolUseTokens(t *testing.T) {
detail := ParseInteractionsUsage([]byte(`{"usage":{"total_input_tokens":2,"total_output_tokens":6,"total_thought_tokens":3,"total_tool_use_tokens":4,"total_tokens":15}}`))
if detail.InputTokens != 6 || detail.OutputTokens != 6 || detail.ReasoningTokens != 3 || detail.TotalTokens != 15 {
t.Fatalf("detail = %+v", detail)
}
if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete ||
detail.TokenBreakdown.Input.UncachedTokens != 6 || detail.TokenBreakdown.Output.TotalTokens != 9 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseInteractionsStreamUsage(t *testing.T) {
detail, ok := ParseInteractionsStreamUsage([]byte(`{"type":"interaction.completed","interaction":{"usage":{"input_tokens":2,"output_tokens":6,"total_tokens":8}}}`))
if !ok {
t.Fatal("ParseInteractionsStreamUsage() ok = false, want true")
}
if detail.TotalTokens != 8 {
t.Fatalf("total tokens = %d, want 8", detail.TotalTokens)
}
}
func TestParseInteractionsStreamUsageOfficialMetadata(t *testing.T) {
detail, ok := ParseInteractionsStreamUsage([]byte(`data: {"event_type":"finish","metadata":{"total_usage":{"total_input_tokens":2,"total_output_tokens":6,"total_thought_tokens":3,"total_cached_tokens":1,"total_tokens":11}}}`))
if !ok {
t.Fatal("ParseInteractionsStreamUsage() ok = false, want true")
}
if detail.InputTokens != 2 {
t.Fatalf("input tokens = %d, want 2", detail.InputTokens)
}
if detail.OutputTokens != 6 {
t.Fatalf("output tokens = %d, want 6", detail.OutputTokens)
}
if detail.ReasoningTokens != 3 {
t.Fatalf("reasoning tokens = %d, want 3", detail.ReasoningTokens)
}
if detail.CachedTokens != 1 {
t.Fatalf("cached tokens = %d, want 1", detail.CachedTokens)
}
if detail.CacheReadTokens != 1 {
t.Fatalf("cache read tokens = %d, want 1", detail.CacheReadTokens)
}
if detail.TotalTokens != 11 {
t.Fatalf("total tokens = %d, want 11", detail.TotalTokens)
}
}
func TestUsageReporterBuildRecordIncludesLatency(t *testing.T) {
reporter := &UsageReporter{
provider: "openai",
model: "gpt-5.4",
requestedAt: time.Now().Add(-1500 * time.Millisecond),
}
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
if record.Latency < time.Second {
t.Fatalf("latency = %v, want >= 1s", record.Latency)
}
if record.Latency > 3*time.Second {
t.Fatalf("latency = %v, want <= 3s", record.Latency)
}
}
func TestUsageReporterTrackHTTPClientStartsTTFTBeforeRoundTrip(t *testing.T) {
delay := 40 * time.Millisecond
reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil)
client := reporter.TrackHTTPClient(&http.Client{
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
time.Sleep(delay)
return &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader("ok")),
Request: req,
}, nil
}),
})
req, errNewRequest := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.invalid/v1/chat/completions", strings.NewReader("{}"))
if errNewRequest != nil {
t.Fatalf("NewRequestWithContext() error = %v", errNewRequest)
}
resp, errDo := client.Do(req)
if errDo != nil {
t.Fatalf("Do() error = %v", errDo)
}
if _, errRead := io.ReadAll(resp.Body); errRead != nil {
t.Fatalf("ReadAll() error = %v", errRead)
}
if errClose := resp.Body.Close(); errClose != nil {
t.Fatalf("response body close error = %v", errClose)
}
if got := reporter.ttftDuration(); got < delay {
t.Fatalf("ttft = %v, want >= %v", got, delay)
}
}
func TestUsageReporterBuildRecordIncludesRequestedModelAlias(t *testing.T) {
ctx := usage.WithRequestedModelAlias(context.Background(), "client-gpt")
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
if record.Model != "gpt-5.4" {
t.Fatalf("model = %q, want %q", record.Model, "gpt-5.4")
}
if record.Alias != "client-gpt" {
t.Fatalf("alias = %q, want %q", record.Alias, "client-gpt")
}
}
func TestNewExecutorUsageReporterIncludesExecutorType(t *testing.T) {
reporter := NewExecutorUsageReporter(context.Background(), &TestUsageExecutor{}, "gpt-5.4", nil)
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
if record.Provider != "test-provider" {
t.Fatalf("provider = %q, want %q", record.Provider, "test-provider")
}
if record.ExecutorType != "TestUsageExecutor" {
t.Fatalf("executor type = %q, want %q", record.ExecutorType, "TestUsageExecutor")
}
}
func TestUsageReporterBuildRecordIncludesReasoningEffort(t *testing.T) {
ctx := usage.WithReasoningEffort(context.Background(), "medium")
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
if record.ReasoningEffort != "medium" {
t.Fatalf("reasoning effort = %q, want %q", record.ReasoningEffort, "medium")
}
}
func TestUsageReporterBuildRecordIncludesServiceTier(t *testing.T) {
ctx := usage.WithServiceTier(context.Background(), "auto")
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
record := reporter.buildRecord(usage.Detail{TotalTokens: 3, ResponseServiceTier: "default"}, false)
if record.ServiceTier != "auto" {
t.Fatalf("service tier = %q, want %q", record.ServiceTier, "auto")
}
if record.ResponseServiceTier != "default" {
t.Fatalf("response service tier = %q, want default", record.ResponseServiceTier)
}
}
func TestUsageReporterBuildRecordDefaultsGenerateTrue(t *testing.T) {
reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil)
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
if !usage.GenerateEnabled(record.Generate) {
t.Fatalf("generate = %v, want true", usage.GenerateEnabled(record.Generate))
}
}
func TestUsageReporterBuildRecordIncludesGenerateFalse(t *testing.T) {
ctx := usage.WithGenerate(context.Background(), false)
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
if usage.GenerateEnabled(record.Generate) {
t.Fatalf("generate = %v, want false", usage.GenerateEnabled(record.Generate))
}
}
func TestUsageReporterSetTranslatedReasoningEffortPreservesClientServiceTier(t *testing.T) {
ctx := usage.WithServiceTier(context.Background(), "auto")
reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil)
reporter.SetTranslatedReasoningEffort([]byte(`{"service_tier":"priority"}`), "openai")
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
if record.ServiceTier != "auto" {
t.Fatalf("service tier = %q, want %q", record.ServiceTier, "auto")
}
}
func TestUsageReporterBuildAdditionalModelRecordSkipsZeroTokens(t *testing.T) {
reporter := &UsageReporter{
provider: "codex",
model: "gpt-5.4",
requestedAt: time.Now(),
}
if _, ok := reporter.buildAdditionalModelRecord("gpt-image-2", usage.Detail{}); ok {
t.Fatalf("expected all-zero token usage to be skipped")
}
if _, ok := reporter.buildAdditionalModelRecord("gpt-image-2", usage.Detail{InputTokens: 2}); !ok {
t.Fatalf("expected non-zero input token usage to be recorded")
}
if _, ok := reporter.buildAdditionalModelRecord("gpt-image-2", usage.Detail{CachedTokens: 2}); !ok {
t.Fatalf("expected non-zero cached token usage to be recorded")
}
}
func TestFailFromErrorsMapsContextStatuses(t *testing.T) {
tests := []struct {
name string
err error
want int
}{
{name: "canceled", err: context.Canceled, want: clienterror.StatusClientClosedRequest},
{name: "deadline", err: context.DeadlineExceeded, want: http.StatusGatewayTimeout},
{
name: "url error wraps canceled",
err: &url.Error{Op: "Post", URL: "https://example.com", Err: context.Canceled},
want: clienterror.StatusClientClosedRequest,
},
{name: "plain error", err: errors.New("boom"), want: 0},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
fail := failFromErrors(tc.err)
if fail.StatusCode != tc.want {
t.Fatalf("StatusCode = %d, want %d; body=%q", fail.StatusCode, tc.want, fail.Body)
}
if strings.TrimSpace(fail.Body) == "" {
t.Fatalf("expected non-empty failure body")
}
})
}
if fail := failFromErrors(nil, nil); fail.StatusCode != 0 || fail.Body != "" {
t.Fatalf("failFromErrors(nil) = %+v, want empty failure", fail)
}
}
func TestStreamUsageBufferPublishFailure(t *testing.T) {
var buffer StreamUsageBuffer
buffer.Observe(usage.Detail{InputTokens: 10, OutputTokens: 5, TotalTokens: 15}, true)
reporter := &UsageReporter{
provider: "openai",
model: "gpt-5.4",
}
record := reporter.buildRecord(buffer.detail, true, failFromErrors(context.Canceled))
if !record.Failed {
t.Fatal("expected record to be marked failed")
}
if record.Fail.StatusCode != clienterror.StatusClientClosedRequest {
t.Fatalf("Fail.StatusCode = %d, want %d", record.Fail.StatusCode, clienterror.StatusClientClosedRequest)
}
if record.Detail.TotalTokens != 15 {
t.Fatalf("Detail.TotalTokens = %d, want 15", record.Detail.TotalTokens)
}
}
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
type TestUsageExecutor struct{}
func (TestUsageExecutor) Identifier() string {
return "test-provider"
}

View file

@ -0,0 +1,31 @@
package helps
import "testing"
var (
benchmarkOpenAIContentChunk = []byte(`data: {"choices":[{"delta":{"content":"hello"}}]}`)
benchmarkOpenAITierChunk = []byte(`data: {"service_tier":"default","choices":[]}`)
benchmarkOpenAIUsageChunk = []byte(`data: {"usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}`)
)
func BenchmarkStreamUsageBufferObserveOpenAIStreamContentChunk(b *testing.B) {
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream(benchmarkOpenAITierChunk)
b.ReportAllocs()
b.ResetTimer()
for index := 0; index < b.N; index++ {
buffer.ObserveOpenAIStream(benchmarkOpenAIContentChunk)
}
}
func BenchmarkStreamUsageBufferObserveOpenAIStream100Chunks(b *testing.B) {
b.ReportAllocs()
for index := 0; index < b.N; index++ {
var buffer StreamUsageBuffer
buffer.ObserveOpenAIStream(benchmarkOpenAITierChunk)
for chunk := 0; chunk < 98; chunk++ {
buffer.ObserveOpenAIStream(benchmarkOpenAIContentChunk)
}
buffer.ObserveOpenAIStream(benchmarkOpenAIUsageChunk)
}
}

View file

@ -0,0 +1,150 @@
package helps
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"strings"
"sync"
"time"
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
)
type userIDCacheEntry struct {
value string
expire time.Time
}
var (
userIDCache = make(map[string]userIDCacheEntry)
userIDCacheMu sync.RWMutex
userIDCacheCleanupOnce sync.Once
)
const (
userIDTTL = time.Hour
userIDCacheCleanupPeriod = 15 * time.Minute
)
func startUserIDCacheCleanup() {
go func() {
ticker := time.NewTicker(userIDCacheCleanupPeriod)
defer ticker.Stop()
for range ticker.C {
purgeExpiredUserIDs()
}
}()
}
func purgeExpiredUserIDs() {
now := time.Now()
userIDCacheMu.Lock()
for key, entry := range userIDCache {
if !entry.expire.After(now) {
delete(userIDCache, key)
}
}
userIDCacheMu.Unlock()
}
func userIDCacheKey(apiKey string) string {
sum := sha256.Sum256([]byte(apiKey))
return hex.EncodeToString(sum[:])
}
func CachedUserID(apiKey string) string {
value, errValue := CachedUserIDRequired(context.Background(), apiKey)
if errValue == nil && value != "" {
return value
}
return generateFakeUserID()
}
// CachedUserIDRequired returns a stable fake user ID per apiKey for request-time paths.
func CachedUserIDRequired(ctx context.Context, apiKey string) (string, error) {
newUserID := func() (string, error) {
sessionID, errSessionID := CachedSessionIDRequired(ctx, apiKey)
if errSessionID != nil {
return "", errSessionID
}
return generateFakeUserIDWithSessionID(sessionID), nil
}
if apiKey == "" {
return newUserID()
}
client, homeMode, errClient := currentClaudeIDKVClient()
if homeMode {
if errClient != nil {
return "", errClient
}
key := claudeUserIDKVKey(apiKey)
raw, found, errGet := client.KVGet(ctx, key)
if errGet != nil {
return "", errGet
}
if found && isValidUserID(strings.TrimSpace(string(raw))) {
if _, errExpire := client.KVExpire(ctx, key, userIDTTL); errExpire != nil {
return "", errExpire
}
return strings.TrimSpace(string(raw)), nil
}
newID, errNewID := newUserID()
if errNewID != nil {
return "", errNewID
}
if _, errSet := client.KVSetNX(ctx, key, []byte(newID), userIDTTL); errSet != nil {
return "", errSet
}
raw, found, errGet = client.KVGet(ctx, key)
if errGet != nil {
return "", errGet
}
if found && isValidUserID(strings.TrimSpace(string(raw))) {
return strings.TrimSpace(string(raw)), nil
}
return "", fmt.Errorf("home kv user id missing after set")
}
userIDCacheCleanupOnce.Do(startUserIDCacheCleanup)
key := userIDCacheKey(apiKey)
now := time.Now()
userIDCacheMu.RLock()
entry, ok := userIDCache[key]
valid := ok && entry.value != "" && entry.expire.After(now) && isValidUserID(entry.value)
userIDCacheMu.RUnlock()
if valid {
userIDCacheMu.Lock()
entry = userIDCache[key]
if entry.value != "" && entry.expire.After(now) && isValidUserID(entry.value) {
entry.expire = now.Add(userIDTTL)
userIDCache[key] = entry
userIDCacheMu.Unlock()
return entry.value, nil
}
userIDCacheMu.Unlock()
}
newID, errNewID := newUserID()
if errNewID != nil {
return "", errNewID
}
userIDCacheMu.Lock()
entry, ok = userIDCache[key]
if !ok || entry.value == "" || !entry.expire.After(now) || !isValidUserID(entry.value) {
entry.value = newID
}
entry.expire = now.Add(userIDTTL)
userIDCache[key] = entry
userIDCacheMu.Unlock()
return entry.value, nil
}
func claudeUserIDKVKey(apiKey string) string {
return "cpa:claude:user-id:" + homekv.HashKeyPart(apiKey)
}

View file

@ -0,0 +1,196 @@
package helps
import (
"context"
"encoding/json"
"errors"
"testing"
"time"
)
func resetUserIDCache() {
userIDCacheMu.Lock()
userIDCache = make(map[string]userIDCacheEntry)
userIDCacheMu.Unlock()
}
func TestGenerateFakeUserIDUsesClaudeCode220JSONShape(t *testing.T) {
userID := GenerateFakeUserID()
if !IsValidUserID(userID) {
t.Fatalf("user ID %q is not valid", userID)
}
var value claudeMetadataUserID
if errUnmarshal := json.Unmarshal([]byte(userID), &value); errUnmarshal != nil {
t.Fatalf("unmarshal user ID: %v", errUnmarshal)
}
if value.AccountUUID != "" {
t.Fatalf("account_uuid = %q, want empty", value.AccountUUID)
}
}
func TestCachedUserIDUsesCachedClaudeSessionID(t *testing.T) {
resetUserIDCache()
resetSessionIDCache()
const key = "api-key-shared-session"
sessionID := CachedSessionID(key)
userID := CachedUserID(key)
var value claudeMetadataUserID
if errUnmarshal := json.Unmarshal([]byte(userID), &value); errUnmarshal != nil {
t.Fatalf("unmarshal user ID: %v", errUnmarshal)
}
if value.SessionID != sessionID {
t.Fatalf("metadata session_id = %q, header session ID = %q", value.SessionID, sessionID)
}
}
func TestCachedUserID_ReusesWithinTTL(t *testing.T) {
resetUserIDCache()
first := CachedUserID("api-key-1")
second := CachedUserID("api-key-1")
if first == "" {
t.Fatal("expected generated user_id to be non-empty")
}
if first != second {
t.Fatalf("expected cached user_id to be reused, got %q and %q", first, second)
}
}
func TestCachedUserID_ExpiresAfterTTL(t *testing.T) {
resetUserIDCache()
expiredID := CachedUserID("api-key-expired")
cacheKey := userIDCacheKey("api-key-expired")
userIDCacheMu.Lock()
userIDCache[cacheKey] = userIDCacheEntry{
value: expiredID,
expire: time.Now().Add(-time.Minute),
}
userIDCacheMu.Unlock()
newID := CachedUserID("api-key-expired")
if newID == expiredID {
t.Fatalf("expected expired user_id to be replaced, got %q", newID)
}
if newID == "" {
t.Fatal("expected regenerated user_id to be non-empty")
}
}
func TestCachedUserID_IsScopedByAPIKey(t *testing.T) {
resetUserIDCache()
first := CachedUserID("api-key-1")
second := CachedUserID("api-key-2")
if first == second {
t.Fatalf("expected different API keys to have different user_ids, got %q", first)
}
}
func TestCachedUserID_RenewsTTLOnHit(t *testing.T) {
resetUserIDCache()
key := "api-key-renew"
id := CachedUserID(key)
cacheKey := userIDCacheKey(key)
soon := time.Now()
userIDCacheMu.Lock()
userIDCache[cacheKey] = userIDCacheEntry{
value: id,
expire: soon.Add(2 * time.Second),
}
userIDCacheMu.Unlock()
if refreshed := CachedUserID(key); refreshed != id {
t.Fatalf("expected cached user_id to be reused before expiry, got %q", refreshed)
}
userIDCacheMu.RLock()
entry := userIDCache[cacheKey]
userIDCacheMu.RUnlock()
if entry.expire.Sub(soon) < 30*time.Minute {
t.Fatalf("expected TTL to renew, got %v remaining", entry.expire.Sub(soon))
}
}
func TestCachedUserIDRequiredHomeReusesKVAcrossLocalCacheReset(t *testing.T) {
resetUserIDCache()
client := newFakeClaudeIDKVClient()
useFakeClaudeIDKVClient(t, client, true, nil)
first, errFirst := CachedUserIDRequired(context.Background(), "api-key-1")
if errFirst != nil {
t.Fatalf("CachedUserIDRequired() first error = %v", errFirst)
}
resetUserIDCache()
second, errSecond := CachedUserIDRequired(context.Background(), "api-key-1")
if errSecond != nil {
t.Fatalf("CachedUserIDRequired() second error = %v", errSecond)
}
if first != second {
t.Fatalf("user id = %q then %q, want same Home KV value", first, second)
}
if !IsValidUserID(first) {
t.Fatalf("user id %q is not valid", first)
}
if client.setCount != 2 {
t.Fatalf("KVSetNX count = %d, want 2 (session and user ID)", client.setCount)
}
if client.expireCount != 1 || client.lastExpireTTL != userIDTTL {
t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, userIDTTL)
}
if client.lastSetTTL != userIDTTL {
t.Fatalf("KVSetNX ttl = %v, want %v", client.lastSetTTL, userIDTTL)
}
}
func TestCachedUserIDRequiredEmptyAPIKeyDoesNotUseHomeKV(t *testing.T) {
client := newFakeClaudeIDKVClient()
useFakeClaudeIDKVClient(t, client, true, nil)
value, errValue := CachedUserIDRequired(context.Background(), "")
if errValue != nil {
t.Fatalf("CachedUserIDRequired(empty) error = %v", errValue)
}
if !IsValidUserID(value) {
t.Fatalf("user id %q is not valid", value)
}
if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 {
t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount)
}
}
func TestCachedUserIDRequiredHomeKVFailures(t *testing.T) {
for _, tc := range []struct {
name string
client *fakeClaudeIDKVClient
}{
{name: "get", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}},
{name: "set", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), setErr: errors.New("set failed")}},
{name: "expire", client: &fakeClaudeIDKVClient{values: map[string][]byte{
claudeUserIDKVKey("api-key-1"): []byte(GenerateFakeUserID()),
}, expireErr: errors.New("expire failed")}},
} {
t.Run(tc.name, func(t *testing.T) {
useFakeClaudeIDKVClient(t, tc.client, true, nil)
if _, errValue := CachedUserIDRequired(context.Background(), "api-key-1"); errValue == nil {
t.Fatalf("CachedUserIDRequired() error = nil, want error")
}
})
}
}
func TestCachedUserIDRequiredHomeRequiresReadAfterSet(t *testing.T) {
client := newFakeClaudeIDKVClient()
client.setNoPersist = true
useFakeClaudeIDKVClient(t, client, true, nil)
if _, errValue := CachedUserIDRequired(context.Background(), "api-key-1"); errValue == nil {
t.Fatalf("CachedUserIDRequired() error = nil, want missing-after-set error")
}
}

View file

@ -0,0 +1,407 @@
package helps
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
"strings"
"sync"
"time"
tls "github.com/refraction-networking/utls"
internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/httpwire"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
"github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil"
log "github.com/sirupsen/logrus"
"golang.org/x/net/http2"
"golang.org/x/net/proxy"
)
// utlsRoundTripper implements http.RoundTripper using a Chrome fingerprint for
// providers that require a browser-like TLS and HTTP/2 transport. Each request
// gets a dedicated connection that is closed with the response body.
type utlsRoundTripper struct {
dialer proxy.Dialer
}
type closeConnectionBody struct {
io.ReadCloser
closeConnection func() error
once sync.Once
err error
}
func (b *closeConnectionBody) Close() error {
if b == nil {
return nil
}
b.once.Do(func() {
var errConnection error
if b.closeConnection != nil {
errConnection = b.closeConnection()
}
var errBody error
if b.ReadCloser != nil {
errBody = b.ReadCloser.Close()
}
b.err = errors.Join(errBody, errConnection)
})
return b.err
}
func newUtlsRoundTripper(proxyURL string) *utlsRoundTripper {
var dialer proxy.Dialer = proxy.Direct
if proxyURL != "" {
proxyDialer, mode, errBuild := proxyutil.BuildDialer(proxyURL)
if errBuild != nil {
log.Errorf("utls: failed to configure proxy dialer for %q: %v", proxyutil.Redact(proxyURL), errBuild)
} else if mode != proxyutil.ModeInherit && proxyDialer != nil {
dialer = proxyDialer
}
}
return &utlsRoundTripper{dialer: dialer}
}
func (t *utlsRoundTripper) createConnection(ctx context.Context, host, addr string) (*http2.ClientConn, error) {
contextDialer, ok := t.dialer.(proxy.ContextDialer)
if !ok {
return nil, fmt.Errorf("utls: dialer does not support context cancellation")
}
conn, errDial := contextDialer.DialContext(ctx, "tcp", addr)
if errDial != nil {
return nil, fmt.Errorf("utls: dial upstream: %w", errDial)
}
tlsConfig := &tls.Config{ServerName: host}
tlsConn := tls.UClient(conn, tlsConfig, tls.HelloChrome_Auto)
if errHandshake := tlsConn.HandshakeContext(ctx); errHandshake != nil {
if errors.Is(errHandshake, context.Canceled) || errors.Is(errHandshake, context.DeadlineExceeded) {
return nil, fmt.Errorf("utls: TLS handshake: %w", errHandshake)
}
if errClose := conn.Close(); errClose != nil {
return nil, fmt.Errorf("utls: TLS handshake: %w; close connection: %v", errHandshake, errClose)
}
return nil, fmt.Errorf("utls: TLS handshake: %w", errHandshake)
}
tr := &http2.Transport{}
h2Conn, errClientConn := tr.NewClientConn(tlsConn)
if errClientConn != nil {
if errClose := tlsConn.Close(); errClose != nil {
return nil, fmt.Errorf("utls: initialize HTTP/2 connection: %w; close TLS connection: %v", errClientConn, errClose)
}
return nil, fmt.Errorf("utls: initialize HTTP/2 connection: %w", errClientConn)
}
return h2Conn, nil
}
func (t *utlsRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
hostname := req.URL.Hostname()
port := req.URL.Port()
if port == "" {
port = "443"
}
addr := net.JoinHostPort(hostname, port)
h2Conn, err := t.createConnection(req.Context(), hostname, addr)
if err != nil {
return nil, err
}
resp, err := h2Conn.RoundTrip(req)
if err != nil {
if errClose := h2Conn.Close(); errClose != nil {
log.Debugf("utls: close connection after round trip failure: %v", errClose)
}
return nil, err
}
if resp == nil {
if errClose := h2Conn.Close(); errClose != nil {
log.Debugf("utls: close connection after empty response: %v", errClose)
}
return nil, fmt.Errorf("utls: upstream returned an empty response")
}
if resp.Body == nil {
resp.Body = http.NoBody
}
resp.Body = &closeConnectionBody{
ReadCloser: resp.Body,
closeConnection: h2Conn.Close,
}
return resp, nil
}
// claudeCodeSessionCacheCapacity bounds the per-transport TLS session cache for
// the Anthropic inference plane.
const claudeCodeSessionCacheCapacity = 32
// newClaudeCodeTLSConfig builds the uTLS config for one inference-plane dial.
//
// OmitEmptyPsk keeps the pre_shared_key extension silent until a session is
// cached, so an unresumed ClientHello stays byte-identical to the captured
// native handshake. PreferSkipResumptionOnNilExtension turns uTLS's HelloCustom
// "resume without the matching extension" panic into a skipped resumption.
func newClaudeCodeTLSConfig(host string, sessionCache tls.ClientSessionCache) *tls.Config {
return &tls.Config{
ServerName: host,
ClientSessionCache: sessionCache,
OmitEmptyPsk: true,
PreferSkipResumptionOnNilExtension: true,
}
}
// claudeCodeTLSClientHelloSpec reproduces the deterministic Node/OpenSSL
// ClientHello emitted by Claude Code 2.1.220 on macOS arm64. Keep this spec in
// sync with a fresh native capture whenever the advertised Claude Code version
// changes.
func claudeCodeTLSClientHelloSpec() *tls.ClientHelloSpec {
return &tls.ClientHelloSpec{
CipherSuites: []uint16{
tls.TLS_AES_128_GCM_SHA256,
tls.TLS_AES_256_GCM_SHA384,
tls.TLS_CHACHA20_POLY1305_SHA256,
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
tls.TLS_RSA_WITH_AES_128_GCM_SHA256,
tls.TLS_RSA_WITH_AES_256_GCM_SHA384,
tls.TLS_RSA_WITH_AES_128_CBC_SHA,
tls.TLS_RSA_WITH_AES_256_CBC_SHA,
},
CompressionMethods: []uint8{0},
Extensions: []tls.TLSExtension{
&tls.SNIExtension{},
&tls.ExtendedMasterSecretExtension{},
&tls.RenegotiationInfoExtension{Renegotiation: tls.RenegotiateOnceAsClient},
&tls.SupportedCurvesExtension{Curves: []tls.CurveID{tls.X25519, tls.CurveP256, tls.CurveP384}},
&tls.SupportedPointsExtension{SupportedPoints: []byte{0}},
&tls.SessionTicketExtension{},
&tls.ALPNExtension{AlpnProtocols: []string{"http/1.1"}},
&tls.StatusRequestExtension{},
&tls.SignatureAlgorithmsExtension{SupportedSignatureAlgorithms: []tls.SignatureScheme{
tls.ECDSAWithP256AndSHA256,
tls.PSSWithSHA256,
tls.PKCS1WithSHA256,
tls.ECDSAWithP384AndSHA384,
tls.PSSWithSHA384,
tls.PKCS1WithSHA384,
tls.PSSWithSHA512,
tls.PKCS1WithSHA512,
tls.PKCS1WithSHA1,
}},
&tls.SCTExtension{},
&tls.KeyShareExtension{KeyShares: []tls.KeyShare{{Group: tls.X25519}}},
&tls.PSKKeyExchangeModesExtension{Modes: []uint8{tls.PskModeDHE}},
&tls.SupportedVersionsExtension{Versions: []uint16{tls.VersionTLS13, tls.VersionTLS12}},
&tls.UtlsPaddingExtension{GetPaddingLen: tls.BoringPaddingStyle},
// pre_shared_key MUST be the final extension (RFC 8446 4.2.11), after
// padding. It contributes zero bytes until a cached session exists.
&tls.UtlsPreSharedKeyExtension{},
},
}
}
const claudeCodeRoundTripperCacheCapacity = 64
var claudeCodeRoundTripperCache = internalcache.NewBoundedLRU[string, http.RoundTripper](
claudeCodeRoundTripperCacheCapacity,
func(_ string, roundTripper http.RoundTripper) {
if transport, ok := roundTripper.(interface{ CloseIdleConnections() }); ok {
transport.CloseIdleConnections()
}
},
)
var claudeCodeMessagesHeaderOrder = []string{
"Accept",
"Authorization",
"Content-Type",
"User-Agent",
"X-Claude-Code-Session-Id",
"X-Stainless-Arch",
"X-Stainless-Lang",
"X-Stainless-OS",
"X-Stainless-Package-Version",
"X-Stainless-Retry-Count",
"X-Stainless-Runtime",
"X-Stainless-Runtime-Version",
"X-Stainless-Timeout",
"anthropic-beta",
"anthropic-dangerous-direct-browser-access",
"anthropic-version",
"x-app",
"x-client-request-id",
"Connection",
"Host",
"Accept-Encoding",
"Content-Length",
}
var claudeCodeCountTokensHeaderOrder = []string{
"Accept",
"Authorization",
"Content-Type",
"User-Agent",
"X-Claude-Code-Session-Id",
"X-Stainless-Arch",
"X-Stainless-Lang",
"X-Stainless-OS",
"X-Stainless-Package-Version",
"X-Stainless-Retry-Count",
"X-Stainless-Runtime",
"X-Stainless-Runtime-Version",
"anthropic-beta",
"anthropic-dangerous-direct-browser-access",
"anthropic-version",
"x-app",
"x-client-request-id",
"Connection",
"Host",
"Accept-Encoding",
"Content-Length",
}
func claudeCodeRequestHeaderOrder(_, requestTarget string) []string {
if strings.HasPrefix(requestTarget, "/v1/messages/count_tokens") {
return claudeCodeCountTokensHeaderOrder
}
return claudeCodeMessagesHeaderOrder
}
func cachedClaudeCodeRoundTripper(proxyURL string) http.RoundTripper {
return claudeCodeRoundTripperCache.GetOrAdd(proxyURL, func() http.RoundTripper {
return newClaudeCodeRoundTripper(proxyURL)
})
}
func newClaudeCodeRoundTripper(proxyURL string) http.RoundTripper {
// The cache is scoped to this round tripper, which is already keyed by proxy,
// so resumption never crosses proxy boundaries.
sessionCache := tls.NewLRUClientSessionCache(claudeCodeSessionCacheCapacity)
var dialer proxy.Dialer = proxy.Direct
if proxyURL != "" {
proxyDialer, mode, errBuild := proxyutil.BuildDialer(proxyURL)
if errBuild != nil {
log.Errorf("claude tls: failed to configure proxy dialer for %q: %v", proxyutil.Redact(proxyURL), errBuild)
} else if mode != proxyutil.ModeInherit && proxyDialer != nil {
dialer = proxyDialer
}
}
transport := &http.Transport{
ForceAttemptHTTP2: false,
DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
var (
conn net.Conn
err error
)
if contextDialer, ok := dialer.(proxy.ContextDialer); ok {
conn, err = contextDialer.DialContext(ctx, network, addr)
} else {
conn, err = dialer.Dial(network, addr)
}
if err != nil {
return nil, fmt.Errorf("claude tls: dial upstream: %w", err)
}
host, _, errSplit := net.SplitHostPort(addr)
if errSplit != nil {
if errClose := conn.Close(); errClose != nil {
log.Debugf("claude tls: close failed connection: %v", errClose)
}
return nil, fmt.Errorf("claude tls: split upstream address: %w", errSplit)
}
tlsConn := tls.UClient(conn, newClaudeCodeTLSConfig(host, sessionCache), tls.HelloCustom)
if errPreset := tlsConn.ApplyPreset(claudeCodeTLSClientHelloSpec()); errPreset != nil {
if errClose := tlsConn.Close(); errClose != nil {
log.Debugf("claude tls: close connection after preset failure: %v", errClose)
}
return nil, fmt.Errorf("claude tls: apply Claude Code ClientHello: %w", errPreset)
}
if errHandshake := tlsConn.HandshakeContext(ctx); errHandshake != nil {
if errClose := tlsConn.Close(); errClose != nil {
log.Debugf("claude tls: close connection after handshake failure: %v", errClose)
}
return nil, fmt.Errorf("claude tls: handshake upstream: %w", errHandshake)
}
return httpwire.NewOrderedRequestConn(tlsConn, claudeCodeRequestHeaderOrder), nil
},
}
return transport
}
// fallbackRoundTripper uses provider-specific TLS fingerprints for protected
// HTTPS hosts and falls back to the standard transport for all other requests.
type fallbackRoundTripper struct {
anthropic http.RoundTripper
chrome http.RoundTripper
fallback http.RoundTripper
}
func (f *fallbackRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if IsAnthropicUpstreamURL(req.URL) {
return f.anthropic.RoundTrip(req)
}
if req.URL.Scheme == "https" && strings.EqualFold(req.URL.Hostname(), "chatgpt.com") {
return f.chrome.RoundTrip(req)
}
return f.fallback.RoundTrip(req)
}
// NewUtlsHTTPClient creates an HTTP client using provider-specific TLS
// fingerprints for protected hosts. It uses Claude Code's Node/OpenSSL profile
// for Anthropic and a Chrome profile for ChatGPT, with a standard-transport
// fallback for other hosts.
func NewUtlsHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
var proxyURL string
if auth != nil {
proxyURL = strings.TrimSpace(auth.ProxyURL)
}
if proxyURL == "" && cfg != nil {
proxyURL = strings.TrimSpace(cfg.ProxyURL)
}
var ctxRoundTripper http.RoundTripper
if ctx != nil {
ctxRoundTripper, _ = ctx.Value("cliproxy.roundtripper").(http.RoundTripper)
}
var chromeRT http.RoundTripper = newUtlsRoundTripper(proxyURL)
var anthropicRT http.RoundTripper = cachedClaudeCodeRoundTripper(proxyURL)
var standardTransport http.RoundTripper = http.DefaultTransport
if proxyURL != "" {
if transport := buildProxyTransport(proxyURL); transport != nil {
standardTransport = transport
}
} else if ctxRoundTripper != nil {
chromeRT = ctxRoundTripper
anthropicRT = ctxRoundTripper
standardTransport = ctxRoundTripper
}
client := &http.Client{
Transport: &fallbackRoundTripper{
anthropic: anthropicRT,
chrome: chromeRT,
fallback: standardTransport,
},
}
if timeout > 0 {
client.Timeout = timeout
}
return client
}

View file

@ -0,0 +1,136 @@
package helps
import (
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
gotls "crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"errors"
"io"
"math/big"
"net"
"testing"
"time"
tls "github.com/refraction-networking/utls"
)
// newResumptionTestCertificate mints a short-lived self-signed leaf for the
// loopback TLS server used by the resumption test.
func newResumptionTestCertificate(t *testing.T) gotls.Certificate {
t.Helper()
key, errKey := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if errKey != nil {
t.Fatalf("generate test key: %v", errKey)
}
template := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "api.anthropic.com"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(time.Hour),
DNSNames: []string{"api.anthropic.com"},
KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
IsCA: true,
}
der, errCreate := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
if errCreate != nil {
t.Fatalf("create test certificate: %v", errCreate)
}
leaf, errParse := x509.ParseCertificate(der)
if errParse != nil {
t.Fatalf("parse test certificate: %v", errParse)
}
return gotls.Certificate{Certificate: [][]byte{der}, PrivateKey: key, Leaf: leaf}
}
// TestClaudeCodeTLSSessionResumptionCompletesHandshake proves the Claude Code
// inference ClientHello can actually resume: the spec places pre_shared_key
// after the padding extension, so a malformed ordering or padding interaction
// would surface here as a handshake failure rather than a silent regression.
func TestClaudeCodeTLSSessionResumptionCompletesHandshake(t *testing.T) {
certificate := newResumptionTestCertificate(t)
roots := x509.NewCertPool()
roots.AddCert(certificate.Leaf)
listener, errListen := net.Listen("tcp", "127.0.0.1:0")
if errListen != nil {
t.Fatalf("listen: %v", errListen)
}
t.Cleanup(func() {
if errClose := listener.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) {
t.Errorf("close listener: %v", errClose)
}
})
serverConfig := &gotls.Config{
Certificates: []gotls.Certificate{certificate},
MinVersion: gotls.VersionTLS13,
}
go func() {
for {
raw, errAccept := listener.Accept()
if errAccept != nil {
return
}
go func(conn net.Conn) {
server := gotls.Server(conn, serverConfig)
if errHandshake := server.Handshake(); errHandshake != nil {
_ = conn.Close()
return
}
// The greeting flushes the post-handshake NewSessionTicket
// messages the client needs in order to resume.
_, _ = server.Write([]byte("ok\n"))
_, _ = server.Read(make([]byte, 8))
_ = server.Close()
}(raw)
}
}()
sessionCache := tls.NewLRUClientSessionCache(claudeCodeSessionCacheCapacity)
dial := func(round int) (resumed bool, helloLength int) {
raw, errDial := net.Dial("tcp", listener.Addr().String())
if errDial != nil {
t.Fatalf("round %d dial: %v", round, errDial)
}
defer func() {
if errClose := raw.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) {
t.Errorf("round %d close: %v", round, errClose)
}
}()
config := newClaudeCodeTLSConfig("api.anthropic.com", sessionCache)
config.RootCAs = roots
conn := tls.UClient(raw, config, tls.HelloCustom)
if errPreset := conn.ApplyPreset(claudeCodeTLSClientHelloSpec()); errPreset != nil {
t.Fatalf("round %d apply preset: %v", round, errPreset)
}
if errHandshake := conn.Handshake(); errHandshake != nil {
t.Fatalf("round %d handshake: %v", round, errHandshake)
}
helloLength = len(conn.HandshakeState.Hello.Raw)
if _, errRead := conn.Read(make([]byte, 8)); errRead != nil && !errors.Is(errRead, io.EOF) {
t.Fatalf("round %d read: %v", round, errRead)
}
_, _ = conn.Write([]byte("bye\n"))
return conn.ConnectionState().DidResume, helloLength
}
firstResumed, firstLength := dial(1)
if firstResumed {
t.Fatal("first handshake reported resumption without a cached session")
}
secondResumed, secondLength := dial(2)
if !secondResumed {
t.Fatal("second handshake did not resume, so the session cache is not effective")
}
// The padding extension absorbs the pre_shared_key bytes, so a resumed
// ClientHello keeps the same BoringSSL padding boundary as a fresh one.
if firstLength != secondLength {
t.Fatalf("resumed ClientHello length = %d, want %d to match the fresh handshake", secondLength, firstLength)
}
}

View file

@ -0,0 +1,641 @@
package helps
import (
"bytes"
"context"
"crypto/md5"
"encoding/binary"
"encoding/hex"
"errors"
"fmt"
"io"
"net"
"net/http"
"os"
"reflect"
"strconv"
"strings"
"sync/atomic"
"testing"
"time"
tls "github.com/refraction-networking/utls"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
)
type utlsClientRoundTripFunc func(*http.Request) (*http.Response, error)
func (f utlsClientRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
type trackedReadCloser struct {
io.Reader
closeCount int
closeErr error
onClose func()
}
func (r *trackedReadCloser) Close() error {
r.closeCount++
if r.onClose != nil {
r.onClose()
}
return r.closeErr
}
type contextDialerFunc func(context.Context, string, string) (net.Conn, error)
func (f contextDialerFunc) Dial(network, addr string) (net.Conn, error) {
return f(context.Background(), network, addr)
}
func (f contextDialerFunc) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
return f(ctx, network, addr)
}
type trackedNetConn struct {
net.Conn
closeCount atomic.Int32
}
func (c *trackedNetConn) Close() error {
c.closeCount.Add(1)
return c.Conn.Close()
}
func TestCloseConnectionBodyClosesConnectionBeforeBodyOnce(t *testing.T) {
bodyErr := errors.New("body close failed")
connectionErr := errors.New("connection close failed")
var closeOrder []string
body := &trackedReadCloser{
Reader: strings.NewReader("response"),
closeErr: bodyErr,
onClose: func() {
closeOrder = append(closeOrder, "body")
},
}
connectionCloseCount := 0
wrapped := &closeConnectionBody{
ReadCloser: body,
closeConnection: func() error {
connectionCloseCount++
closeOrder = append(closeOrder, "connection")
return connectionErr
},
}
payload, errRead := io.ReadAll(wrapped)
if errRead != nil {
t.Fatal(errRead)
}
if got, want := string(payload), "response"; got != want {
t.Fatalf("response body = %q, want %q", got, want)
}
errClose := wrapped.Close()
if !errors.Is(errClose, bodyErr) {
t.Fatalf("close error = %v, want body close error", errClose)
}
if !errors.Is(errClose, connectionErr) {
t.Fatalf("close error = %v, want connection close error", errClose)
}
if errCloseAgain := wrapped.Close(); errCloseAgain != errClose {
t.Fatalf("second close error = %v, want %v", errCloseAgain, errClose)
}
if body.closeCount != 1 {
t.Fatalf("body close count = %d, want 1", body.closeCount)
}
if connectionCloseCount != 1 {
t.Fatalf("connection close count = %d, want 1", connectionCloseCount)
}
if want := []string{"connection", "body"}; !reflect.DeepEqual(closeOrder, want) {
t.Fatalf("close order = %v, want %v", closeOrder, want)
}
}
func TestUtlsRoundTripperDialUsesRequestContext(t *testing.T) {
dialStarted := make(chan struct{})
roundTripper := &utlsRoundTripper{dialer: contextDialerFunc(func(ctx context.Context, _, _ string) (net.Conn, error) {
close(dialStarted)
<-ctx.Done()
return nil, ctx.Err()
})}
ctx, cancel := context.WithCancel(t.Context())
req, errRequest := http.NewRequestWithContext(ctx, http.MethodGet, "https://chatgpt.com/backend-api/codex/responses", nil)
if errRequest != nil {
t.Fatal(errRequest)
}
roundTripDone := make(chan error, 1)
go func() {
resp, errRoundTrip := roundTripper.RoundTrip(req)
if resp != nil && resp.Body != nil {
errRoundTrip = errors.Join(errRoundTrip, resp.Body.Close())
}
roundTripDone <- errRoundTrip
}()
select {
case <-dialStarted:
case <-time.After(time.Second):
t.Fatal("dial did not start")
}
cancel()
select {
case errRoundTrip := <-roundTripDone:
if !errors.Is(errRoundTrip, context.Canceled) {
t.Fatalf("RoundTrip error = %v, want context canceled", errRoundTrip)
}
case <-time.After(time.Second):
t.Fatal("RoundTrip did not stop after context cancellation")
}
}
func TestUtlsRoundTripperHandshakeUsesRequestContext(t *testing.T) {
clientConn, serverConn := net.Pipe()
t.Cleanup(func() {
if errClose := clientConn.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) && !errors.Is(errClose, io.ErrClosedPipe) {
t.Errorf("close client connection: %v", errClose)
}
if errClose := serverConn.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) && !errors.Is(errClose, io.ErrClosedPipe) {
t.Errorf("close server connection: %v", errClose)
}
})
trackedConn := &trackedNetConn{Conn: clientConn}
dialDone := make(chan struct{})
roundTripper := &utlsRoundTripper{dialer: contextDialerFunc(func(context.Context, string, string) (net.Conn, error) {
close(dialDone)
return trackedConn, nil
})}
ctx, cancel := context.WithCancel(t.Context())
connectionDone := make(chan error, 1)
go func() {
h2Conn, errConnect := roundTripper.createConnection(ctx, "chatgpt.com", "chatgpt.com:443")
if h2Conn != nil {
errConnect = errors.Join(errConnect, h2Conn.Close())
}
connectionDone <- errConnect
}()
select {
case <-dialDone:
case <-time.After(time.Second):
t.Fatal("dial did not complete")
}
cancel()
select {
case errConnect := <-connectionDone:
if !errors.Is(errConnect, context.Canceled) {
t.Fatalf("createConnection error = %v, want context canceled", errConnect)
}
case <-time.After(time.Second):
t.Fatal("TLS handshake did not stop after context cancellation")
}
if got := trackedConn.closeCount.Load(); got != 1 {
t.Fatalf("connection close count = %d, want 1", got)
}
}
type claudeCodeTLSFingerprintFixture struct {
ClientHelloLength int
JA3 string
JA3MD5 string
ALPN []string
HTTPVersion string
CipherSuites []uint16
ExtensionTypes []uint16
ExtensionLengths [][2]int
SupportedGroups []uint16
PointFormats []uint8
SignatureAlgorithms []uint16
SupportedVersions []uint16
KeyShareGroups []uint16
}
func TestClaudeCodeTLSClientHelloSpecMatches220Capture(t *testing.T) {
t.Parallel()
fixture := claudeCodeTLSFingerprintFixture{
ClientHelloLength: 508,
JA3: "771,4865-4866-4867-49195-49199-49196-49200-52393-52392-49161-49171-49162-49172-156-157-47-53,0-23-65281-10-11-35-16-5-13-18-51-45-43-21,29-23-24,0",
JA3MD5: "d871d02cecbde59abbf8f4806134addf",
ALPN: []string{"http/1.1"},
HTTPVersion: "HTTP/1.1",
CipherSuites: []uint16{4865, 4866, 4867, 49195, 49199, 49196, 49200, 52393, 52392, 49161, 49171, 49162, 49172, 156, 157, 47, 53},
ExtensionTypes: []uint16{0, 23, 65281, 10, 11, 35, 16, 5, 13, 18, 51, 45, 43, 21},
ExtensionLengths: [][2]int{
{0, 22}, {23, 0}, {65281, 1}, {10, 8}, {11, 2}, {35, 0}, {16, 11},
{5, 5}, {13, 20}, {18, 0}, {51, 38}, {45, 2}, {43, 5}, {21, 231},
},
SupportedGroups: []uint16{29, 23, 24},
PointFormats: []uint8{0},
SignatureAlgorithms: []uint16{1027, 2052, 1025, 1283, 2053, 1281, 2054, 1537, 513},
SupportedVersions: []uint16{772, 771},
KeyShareGroups: []uint16{29},
}
record := captureClaudeCodeClientHello(t)
if got := len(record) - 9; got != fixture.ClientHelloLength {
t.Fatalf("ClientHello length = %d, want %d", got, fixture.ClientHelloLength)
}
if got := parseClientHelloExtensionLengths(t, record); !reflect.DeepEqual(got, fixture.ExtensionLengths) {
t.Fatalf("extension lengths = %v, want %v", got, fixture.ExtensionLengths)
}
spec, errFingerprint := (&tls.Fingerprinter{}).FingerprintClientHello(record)
if errFingerprint != nil {
t.Fatal(errFingerprint)
}
actual := summarizeClaudeCodeClientHelloSpec(t, spec)
if !reflect.DeepEqual(actual.CipherSuites, fixture.CipherSuites) {
t.Fatalf("cipher suites = %v, want %v", actual.CipherSuites, fixture.CipherSuites)
}
if !reflect.DeepEqual(actual.ExtensionTypes, fixture.ExtensionTypes) {
t.Fatalf("extension types = %v, want %v", actual.ExtensionTypes, fixture.ExtensionTypes)
}
if !reflect.DeepEqual(actual.ALPN, fixture.ALPN) {
t.Fatalf("ALPN = %v, want %v", actual.ALPN, fixture.ALPN)
}
if !reflect.DeepEqual(actual.SupportedGroups, fixture.SupportedGroups) {
t.Fatalf("supported groups = %v, want %v", actual.SupportedGroups, fixture.SupportedGroups)
}
if !reflect.DeepEqual(actual.PointFormats, fixture.PointFormats) {
t.Fatalf("point formats = %v, want %v", actual.PointFormats, fixture.PointFormats)
}
if !reflect.DeepEqual(actual.SignatureAlgorithms, fixture.SignatureAlgorithms) {
t.Fatalf("signature algorithms = %v, want %v", actual.SignatureAlgorithms, fixture.SignatureAlgorithms)
}
if !reflect.DeepEqual(actual.SupportedVersions, fixture.SupportedVersions) {
t.Fatalf("supported versions = %v, want %v", actual.SupportedVersions, fixture.SupportedVersions)
}
if !reflect.DeepEqual(actual.KeyShareGroups, fixture.KeyShareGroups) {
t.Fatalf("key share groups = %v, want %v", actual.KeyShareGroups, fixture.KeyShareGroups)
}
if actual.JA3 != fixture.JA3 || actual.JA3MD5 != fixture.JA3MD5 {
t.Fatalf("JA3 = %q (%s), want %q (%s)", actual.JA3, actual.JA3MD5, fixture.JA3, fixture.JA3MD5)
}
transport, ok := newClaudeCodeRoundTripper("").(*http.Transport)
if !ok {
t.Fatalf("Claude Code transport type = %T, want *http.Transport", newClaudeCodeRoundTripper(""))
}
if transport.ForceAttemptHTTP2 {
t.Fatal("Claude Code transport must not force HTTP/2")
}
if fixture.HTTPVersion != "HTTP/1.1" {
t.Fatalf("fixture HTTP version = %q, want HTTP/1.1", fixture.HTTPVersion)
}
}
func TestClaudeCodeTLSResumptionIsWireSafe(t *testing.T) {
t.Parallel()
// RFC 8446 4.2.11 requires pre_shared_key to be the final extension, after
// the padding extension.
spec := claudeCodeTLSClientHelloSpec()
last := spec.Extensions[len(spec.Extensions)-1]
if _, ok := last.(*tls.UtlsPreSharedKeyExtension); !ok {
t.Fatalf("last inference extension = %T, want *tls.UtlsPreSharedKeyExtension", last)
}
if _, ok := spec.Extensions[len(spec.Extensions)-2].(*tls.UtlsPaddingExtension); !ok {
t.Fatalf("extension before pre_shared_key = %T, want *tls.UtlsPaddingExtension", spec.Extensions[len(spec.Extensions)-2])
}
// Without OmitEmptyPsk uTLS refuses to marshal an empty PSK, and without
// PreferSkipResumptionOnNilExtension a HelloCustom resumption attempt panics.
cfg := newClaudeCodeTLSConfig("api.anthropic.com", tls.NewLRUClientSessionCache(claudeCodeSessionCacheCapacity))
if cfg.ClientSessionCache == nil {
t.Fatal("ClientSessionCache = nil, want a session cache so resumption is possible")
}
if !cfg.OmitEmptyPsk {
t.Fatal("OmitEmptyPsk = false, want true so an unresumed ClientHello stays byte-identical")
}
if !cfg.PreferSkipResumptionOnNilExtension {
t.Fatal("PreferSkipResumptionOnNilExtension = false, want true to avoid a HelloCustom resumption panic")
}
}
func TestClaudeCodeRequestHeaderOrderMatchesNative220Capture(t *testing.T) {
t.Parallel()
if got, want := claudeCodeRequestHeaderOrder(http.MethodPost, "/v1/messages?beta=true"), claudeCodeMessagesHeaderOrder; !reflect.DeepEqual(got, want) {
t.Fatalf("Messages header order = %v, want %v", got, want)
}
if got, want := claudeCodeRequestHeaderOrder(http.MethodPost, "/v1/messages/count_tokens?beta=true"), claudeCodeCountTokensHeaderOrder; !reflect.DeepEqual(got, want) {
t.Fatalf("count_tokens header order = %v, want %v", got, want)
}
for _, name := range claudeCodeCountTokensHeaderOrder {
if name == "X-Stainless-Timeout" {
t.Fatal("count_tokens header order unexpectedly contains X-Stainless-Timeout")
}
}
}
func TestCachedClaudeCodeRoundTripperReusesTransport(t *testing.T) {
t.Parallel()
const proxyURL = "http://127.0.0.1:29653"
first := cachedClaudeCodeRoundTripper(proxyURL)
second := cachedClaudeCodeRoundTripper(proxyURL)
if first != second {
t.Fatal("Claude Code transport cache returned different transports for one proxy")
}
}
func TestCachedClaudeCodeRoundTripperBoundsProxyCardinality(t *testing.T) {
firstProxy := fmt.Sprintf("http://127.0.0.1:%d", 30000)
first := cachedClaudeCodeRoundTripper(firstProxy)
for index := 1; index <= claudeCodeRoundTripperCacheCapacity; index++ {
cachedClaudeCodeRoundTripper(fmt.Sprintf("http://127.0.0.1:%d", 30000+index))
}
if got := claudeCodeRoundTripperCache.Len(); got > claudeCodeRoundTripperCacheCapacity {
t.Fatalf("transport cache entries = %d, want at most %d", got, claudeCodeRoundTripperCacheCapacity)
}
if recreated := cachedClaudeCodeRoundTripper(firstProxy); recreated == first {
t.Fatal("least recently used proxy transport was not evicted")
}
}
func TestClaudeCodeTLSClientHelloCapture(t *testing.T) {
proxyURL := os.Getenv("CPA_TLS_FP_PROXY")
if proxyURL == "" {
t.Skip("CPA_TLS_FP_PROXY is not set")
}
client := NewUtlsHTTPClient(t.Context(), nil, &cliproxyauth.Auth{ProxyURL: proxyURL}, 0)
req, errRequest := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://api.anthropic.com/v1/messages", bytes.NewBufferString(`{"model":"claude-opus-4-6","max_tokens":1,"messages":[{"role":"user","content":"x"}]}`))
if errRequest != nil {
t.Fatal(errRequest)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-api-key", "dummy-tls-fingerprint")
resp, errDo := client.Do(req)
if errDo != nil {
t.Fatal(errDo)
}
if errClose := resp.Body.Close(); errClose != nil {
t.Fatal(errClose)
}
}
func TestFallbackRoundTripperSelectsProviderFingerprint(t *testing.T) {
t.Parallel()
route := func(label string) http.RoundTripper {
return utlsClientRoundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"X-Test-Route": []string{label}},
Body: io.NopCloser(strings.NewReader("{}")),
Request: req,
}, nil
})
}
roundTripper := &fallbackRoundTripper{
anthropic: route("anthropic"),
chrome: route("chrome"),
fallback: route("fallback"),
}
tests := []struct {
name string
url string
want string
}{
{name: "Anthropic HTTPS", url: "https://api.anthropic.com/v1/messages", want: "anthropic"},
{name: "Anthropic explicit HTTPS port", url: "https://api.anthropic.com:443/v1/messages", want: "anthropic"},
{name: "Anthropic custom port", url: "https://api.anthropic.com:8443/v1/messages", want: "fallback"},
{name: "Anthropic userinfo", url: "https://caller@api.anthropic.com/v1/messages", want: "fallback"},
{name: "Anthropic lookalike", url: "https://api.anthropic.com.example/v1/messages", want: "fallback"},
{name: "ChatGPT HTTPS", url: "https://chatgpt.com/backend-api/codex/responses", want: "chrome"},
{name: "Other HTTPS", url: "https://example.com/v1/messages", want: "fallback"},
{name: "Anthropic HTTP", url: "http://api.anthropic.com/v1/messages", want: "fallback"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req, errRequest := http.NewRequest(http.MethodGet, tt.url, nil)
if errRequest != nil {
t.Fatal(errRequest)
}
resp, errRoundTrip := roundTripper.RoundTrip(req)
if errRoundTrip != nil {
t.Fatal(errRoundTrip)
}
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
t.Errorf("close response body: %v", errClose)
}
}()
if got := resp.Header.Get("X-Test-Route"); got != tt.want {
t.Fatalf("route = %q, want %q", got, tt.want)
}
})
}
}
func TestNewUtlsHTTPClientUsesContextRoundTripperForProtectedHost(t *testing.T) {
t.Parallel()
for _, targetURL := range []string{
"https://api.anthropic.com/v1/messages",
"https://chatgpt.com/backend-api/codex/responses",
} {
t.Run(targetURL, func(t *testing.T) {
called := false
ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", utlsClientRoundTripFunc(func(req *http.Request) (*http.Response, error) {
called = true
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader("{}")),
Request: req,
}, nil
}))
client := NewUtlsHTTPClient(ctx, nil, nil, 0)
resp, err := client.Get(targetURL)
if err != nil {
t.Fatalf("client.Get returned error: %v", err)
}
if errClose := resp.Body.Close(); errClose != nil {
t.Fatalf("response body close returned error: %v", errClose)
}
if !called {
t.Fatal("expected context RoundTripper to handle protected host request")
}
})
}
}
type claudeCodeClientHelloSummary struct {
CipherSuites []uint16
ExtensionTypes []uint16
ALPN []string
SupportedGroups []uint16
PointFormats []uint8
SignatureAlgorithms []uint16
SupportedVersions []uint16
KeyShareGroups []uint16
JA3 string
JA3MD5 string
}
func captureClaudeCodeClientHello(t *testing.T) []byte {
t.Helper()
clientConn, serverConn := net.Pipe()
t.Cleanup(func() {
if errClose := clientConn.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) {
t.Errorf("close client pipe: %v", errClose)
}
if errClose := serverConn.Close(); errClose != nil && !errors.Is(errClose, net.ErrClosed) {
t.Errorf("close server pipe: %v", errClose)
}
})
// Use the production config so the captured bytes reflect the real dial path,
// including the resumption settings.
cfg := newClaudeCodeTLSConfig("api.anthropic.com", tls.NewLRUClientSessionCache(claudeCodeSessionCacheCapacity))
tlsConn := tls.UClient(clientConn, cfg, tls.HelloCustom)
if errPreset := tlsConn.ApplyPreset(claudeCodeTLSClientHelloSpec()); errPreset != nil {
t.Fatal(errPreset)
}
handshakeDone := make(chan error, 1)
go func() {
handshakeDone <- tlsConn.Handshake()
}()
if errDeadline := serverConn.SetReadDeadline(time.Now().Add(5 * time.Second)); errDeadline != nil {
t.Fatal(errDeadline)
}
header := make([]byte, 5)
if _, errRead := io.ReadFull(serverConn, header); errRead != nil {
t.Fatal(errRead)
}
payload := make([]byte, int(binary.BigEndian.Uint16(header[3:5])))
if _, errRead := io.ReadFull(serverConn, payload); errRead != nil {
t.Fatal(errRead)
}
if errClose := serverConn.Close(); errClose != nil {
t.Fatal(errClose)
}
select {
case <-handshakeDone:
case <-time.After(5 * time.Second):
t.Fatal("uTLS handshake did not exit after the capture connection closed")
}
return append(header, payload...)
}
func parseClientHelloExtensionLengths(t *testing.T, record []byte) [][2]int {
t.Helper()
if len(record) < 9 || record[0] != 22 || record[5] != 1 {
t.Fatalf("invalid TLS ClientHello record")
}
body := record[9:]
offset := 2 + 32
if offset >= len(body) {
t.Fatal("truncated ClientHello random")
}
sessionLength := int(body[offset])
offset += 1 + sessionLength
if offset+2 > len(body) {
t.Fatal("truncated ClientHello cipher suites")
}
cipherLength := int(binary.BigEndian.Uint16(body[offset : offset+2]))
offset += 2 + cipherLength
if offset >= len(body) {
t.Fatal("truncated ClientHello compression methods")
}
compressionLength := int(body[offset])
offset += 1 + compressionLength
if offset+2 > len(body) {
t.Fatal("truncated ClientHello extensions")
}
extensionsLength := int(binary.BigEndian.Uint16(body[offset : offset+2]))
offset += 2
end := offset + extensionsLength
if end > len(body) {
t.Fatal("truncated ClientHello extension data")
}
lengths := make([][2]int, 0)
for offset+4 <= end {
extensionType := int(binary.BigEndian.Uint16(body[offset : offset+2]))
extensionLength := int(binary.BigEndian.Uint16(body[offset+2 : offset+4]))
lengths = append(lengths, [2]int{extensionType, extensionLength})
offset += 4 + extensionLength
}
if offset != end {
t.Fatal("misaligned ClientHello extension data")
}
return lengths
}
func summarizeClaudeCodeClientHelloSpec(t *testing.T, spec *tls.ClientHelloSpec) claudeCodeClientHelloSummary {
t.Helper()
summary := claudeCodeClientHelloSummary{CipherSuites: append([]uint16(nil), spec.CipherSuites...)}
for _, extension := range spec.Extensions {
switch ext := extension.(type) {
case *tls.SNIExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 0)
case *tls.ExtendedMasterSecretExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 23)
case *tls.RenegotiationInfoExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 65281)
case *tls.SupportedCurvesExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 10)
for _, curve := range ext.Curves {
summary.SupportedGroups = append(summary.SupportedGroups, uint16(curve))
}
case *tls.SupportedPointsExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 11)
summary.PointFormats = append(summary.PointFormats, ext.SupportedPoints...)
case *tls.SessionTicketExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 35)
case *tls.ALPNExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 16)
summary.ALPN = append(summary.ALPN, ext.AlpnProtocols...)
case *tls.StatusRequestExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 5)
case *tls.SignatureAlgorithmsExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 13)
for _, algorithm := range ext.SupportedSignatureAlgorithms {
summary.SignatureAlgorithms = append(summary.SignatureAlgorithms, uint16(algorithm))
}
case *tls.SCTExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 18)
case *tls.KeyShareExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 51)
for _, keyShare := range ext.KeyShares {
summary.KeyShareGroups = append(summary.KeyShareGroups, uint16(keyShare.Group))
}
case *tls.PSKKeyExchangeModesExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 45)
case *tls.SupportedVersionsExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 43)
summary.SupportedVersions = append(summary.SupportedVersions, ext.Versions...)
case *tls.UtlsPaddingExtension:
summary.ExtensionTypes = append(summary.ExtensionTypes, 21)
default:
t.Fatalf("unexpected ClientHello extension type %T", extension)
}
}
cipherStrings := make([]string, 0, len(summary.CipherSuites))
for _, cipher := range summary.CipherSuites {
cipherStrings = append(cipherStrings, strconv.Itoa(int(cipher)))
}
extensionStrings := make([]string, 0, len(summary.ExtensionTypes))
for _, extensionType := range summary.ExtensionTypes {
extensionStrings = append(extensionStrings, strconv.Itoa(int(extensionType)))
}
groupStrings := make([]string, 0, len(summary.SupportedGroups))
for _, group := range summary.SupportedGroups {
groupStrings = append(groupStrings, strconv.Itoa(int(group)))
}
pointStrings := make([]string, 0, len(summary.PointFormats))
for _, point := range summary.PointFormats {
pointStrings = append(pointStrings, strconv.Itoa(int(point)))
}
summary.JA3 = fmt.Sprintf("771,%s,%s,%s,%s", strings.Join(cipherStrings, "-"), strings.Join(extensionStrings, "-"), strings.Join(groupStrings, "-"), strings.Join(pointStrings, "-"))
digest := md5.Sum([]byte(summary.JA3)) // #nosec G401 -- JA3 requires MD5.
summary.JA3MD5 = hex.EncodeToString(digest[:])
return summary
}

View file

@ -0,0 +1,86 @@
package helps
import (
"strings"
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// StripVertexOpenAIResponsesToolCallIDs removes OpenAI Responses call IDs that
// Vertex rejects in Gemini functionCall/functionResponse payloads.
func StripVertexOpenAIResponsesToolCallIDs(payload []byte, sourceFormat string) []byte {
if !strings.EqualFold(strings.TrimSpace(sourceFormat), "openai-response") {
return payload
}
contents := util.GetGJSONBytesNoCopy(payload, "contents")
if !contents.IsArray() || !vertexContentsHaveToolCallIDs(contents) {
return payload
}
contentsChanged := false
contentItems := make([][]byte, 0, int(contents.Get("#").Int()))
contents.ForEach(func(_, content gjson.Result) bool {
parts := content.Get("parts")
if !parts.IsArray() {
contentItems = append(contentItems, []byte(content.Raw))
return true
}
partsChanged := false
partItems := make([][]byte, 0, int(parts.Get("#").Int()))
parts.ForEach(func(_, part gjson.Result) bool {
partJSON := []byte(part.Raw)
for _, path := range []string{"functionCall.id", "functionResponse.id"} {
if !part.Get(path).Exists() {
continue
}
updated, errDelete := sjson.DeleteBytes(partJSON, path)
if errDelete == nil {
partJSON = updated
partsChanged = true
}
}
partItems = append(partItems, partJSON)
return true
})
contentJSON := []byte(content.Raw)
if partsChanged {
updated, errSet := sjson.SetRawBytes(contentJSON, "parts", JoinRawJSONArray(partItems))
if errSet == nil {
contentJSON = updated
contentsChanged = true
}
}
contentItems = append(contentItems, contentJSON)
return true
})
if !contentsChanged {
return payload
}
updated, errSet := sjson.SetRawBytes(payload, "contents", JoinRawJSONArray(contentItems))
if errSet != nil {
return payload
}
return updated
}
func vertexContentsHaveToolCallIDs(contents gjson.Result) bool {
hasIDs := false
contents.ForEach(func(_, content gjson.Result) bool {
parts := content.Get("parts")
if !parts.IsArray() {
return true
}
parts.ForEach(func(_, part gjson.Result) bool {
hasIDs = part.Get("functionCall.id").Exists() || part.Get("functionResponse.id").Exists()
return !hasIDs
})
return !hasIDs
})
return hasIDs
}

View file

@ -0,0 +1,45 @@
package helps
import (
"strings"
"testing"
"github.com/tidwall/gjson"
)
func TestStripVertexToolCallIDsReusesPayloadWithoutIDs(t *testing.T) {
input := []byte(`{"contents":[{"role":"model","parts":[{"functionCall":{"name":"lookup","args":{"id":9007199254740993}}}]}]}`)
output := StripVertexOpenAIResponsesToolCallIDs(input, "openai-response")
if &output[0] != &input[0] {
t.Fatal("payload without tool call IDs was copied")
}
}
func TestStripVertexToolCallIDsRebuildsContentsOnce(t *testing.T) {
input := []byte(`{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"lookup","args":{"id":9007199254740993}}}]},{"role":"user","parts":[{"functionResponse":{"id":"call_1","name":"lookup","response":{"id":"keep"}}}]}]}`)
output := StripVertexOpenAIResponsesToolCallIDs(input, "openai-response")
if gjson.GetBytes(output, "contents.0.parts.0.functionCall.id").Exists() {
t.Fatal("functionCall.id was not removed")
}
if gjson.GetBytes(output, "contents.1.parts.0.functionResponse.id").Exists() {
t.Fatal("functionResponse.id was not removed")
}
if got := gjson.GetBytes(output, "contents.1.parts.0.functionResponse.response.id").String(); got != "keep" {
t.Fatalf("nested response id = %q, want keep", got)
}
if got := gjson.GetBytes(output, "contents.0.parts.0.functionCall.args.id").Raw; got != "9007199254740993" {
t.Fatalf("large integer = %s, want exact original value", got)
}
}
var benchmarkVertexPayloadOutput []byte
func BenchmarkStripVertexToolCallIDsLargeNoopPayload(b *testing.B) {
input := []byte(`{"contents":[{"role":"user","parts":[{"text":"` + strings.Repeat("x", 8<<20) + `"}]}]}`)
b.ReportAllocs()
b.SetBytes(int64(len(input)))
b.ResetTimer()
for b.Loop() {
benchmarkVertexPayloadOutput = StripVertexOpenAIResponsesToolCallIDs(input, "openai-response")
}
}