Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
128
backend/internal/runtime/executor/helps/cache_helpers.go
Normal file
128
backend/internal/runtime/executor/helps/cache_helpers.go
Normal 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)
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
2048
backend/internal/runtime/executor/helps/claude_bip39_words.txt
Normal file
2048
backend/internal/runtime/executor/helps/claude_bip39_words.txt
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
110
backend/internal/runtime/executor/helps/claude_code_session.go
Normal file
110
backend/internal/runtime/executor/helps/claude_code_session.go
Normal 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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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")
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
634
backend/internal/runtime/executor/helps/claude_device_profile.go
Normal file
634
backend/internal/runtime/executor/helps/claude_device_profile.go
Normal 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)
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
137
backend/internal/runtime/executor/helps/claude_diagnostics.go
Normal file
137
backend/internal/runtime/executor/helps/claude_diagnostics.go
Normal 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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
387
backend/internal/runtime/executor/helps/claude_input_tokens.go
Normal file
387
backend/internal/runtime/executor/helps/claude_input_tokens.go
Normal 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")
|
||||
}
|
||||
|
|
@ -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,
|
||||
¶m,
|
||||
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,
|
||||
¶m,
|
||||
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()
|
||||
}
|
||||
144
backend/internal/runtime/executor/helps/claude_mcp_alias.go
Normal file
144
backend/internal/runtime/executor/helps/claude_mcp_alias.go
Normal 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
|
||||
}
|
||||
272
backend/internal/runtime/executor/helps/claude_mcp_alias_test.go
Normal file
272
backend/internal/runtime/executor/helps/claude_mcp_alias_test.go
Normal 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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
249
backend/internal/runtime/executor/helps/claude_ratelimit.go
Normal file
249
backend/internal/runtime/executor/helps/claude_ratelimit.go
Normal 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
|
||||
}
|
||||
193
backend/internal/runtime/executor/helps/claude_ratelimit_test.go
Normal file
193
backend/internal/runtime/executor/helps/claude_ratelimit_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
17
backend/internal/runtime/executor/helps/claude_upstream.go
Normal file
17
backend/internal/runtime/executor/helps/claude_upstream.go
Normal 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"
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
214
backend/internal/runtime/executor/helps/cloak_obfuscate.go
Normal file
214
backend/internal/runtime/executor/helps/cloak_obfuscate.go
Normal 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
|
||||
}
|
||||
69
backend/internal/runtime/executor/helps/cloak_utils.go
Normal file
69
backend/internal/runtime/executor/helps/cloak_utils.go
Normal 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)
|
||||
}
|
||||
194
backend/internal/runtime/executor/helps/codex_input_ids.go
Normal file
194
backend/internal/runtime/executor/helps/codex_input_ids.go
Normal 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
|
||||
}
|
||||
318
backend/internal/runtime/executor/helps/codex_input_ids_test.go
Normal file
318
backend/internal/runtime/executor/helps/codex_input_ids_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
127
backend/internal/runtime/executor/helps/codex_multi_agent_v2.go
Normal file
127
backend/internal/runtime/executor/helps/codex_multi_agent_v2.go
Normal 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)
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
66
backend/internal/runtime/executor/helps/derived_session.go
Normal file
66
backend/internal/runtime/executor/helps/derived_session.go
Normal 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)
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
155
backend/internal/runtime/executor/helps/home_refresh.go
Normal file
155
backend/internal/runtime/executor/helps/home_refresh.go
Normal 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
|
||||
}
|
||||
}
|
||||
178
backend/internal/runtime/executor/helps/home_refresh_test.go
Normal file
178
backend/internal/runtime/executor/helps/home_refresh_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
761
backend/internal/runtime/executor/helps/logging_helpers.go
Normal file
761
backend/internal/runtime/executor/helps/logging_helpers.go
Normal 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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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: ®istry.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: ®istry.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: ®istry.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: ®istry.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)
|
||||
}
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
1003
backend/internal/runtime/executor/helps/payload_helpers.go
Normal file
1003
backend/internal/runtime/executor/helps/payload_helpers.go
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
81
backend/internal/runtime/executor/helps/payload_mutations.go
Normal file
81
backend/internal/runtime/executor/helps/payload_mutations.go
Normal 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, ']')
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
79
backend/internal/runtime/executor/helps/proxy_helpers.go
Normal file
79
backend/internal/runtime/executor/helps/proxy_helpers.go
Normal 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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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,
|
||||
¶m,
|
||||
nil,
|
||||
)
|
||||
chunks2 := TranslateStreamWithClaudeInputTokens(
|
||||
ctx,
|
||||
sdktranslator.FormatOpenAI,
|
||||
sdktranslator.FormatOpenAIResponse,
|
||||
"deepseek-v4-flash",
|
||||
reqBody,
|
||||
translatedReq,
|
||||
chunk2,
|
||||
¶m,
|
||||
nil,
|
||||
)
|
||||
chunks3 := TranslateStreamWithClaudeInputTokens(
|
||||
ctx,
|
||||
sdktranslator.FormatOpenAI,
|
||||
sdktranslator.FormatOpenAIResponse,
|
||||
"deepseek-v4-flash",
|
||||
reqBody,
|
||||
translatedReq,
|
||||
chunk3,
|
||||
¶m,
|
||||
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")
|
||||
}
|
||||
}
|
||||
148
backend/internal/runtime/executor/helps/session_id_cache.go
Normal file
148
backend/internal/runtime/executor/helps/session_id_cache.go
Normal 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)
|
||||
}
|
||||
178
backend/internal/runtime/executor/helps/session_id_cache_test.go
Normal file
178
backend/internal/runtime/executor/helps/session_id_cache_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
64
backend/internal/runtime/executor/helps/thinking.go
Normal file
64
backend/internal/runtime/executor/helps/thinking.go
Normal 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
|
||||
}
|
||||
|
|
@ -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"
|
||||
)
|
||||
100
backend/internal/runtime/executor/helps/thinking_test.go
Normal file
100
backend/internal/runtime/executor/helps/thinking_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
236
backend/internal/runtime/executor/helps/token_helpers.go
Normal file
236
backend/internal/runtime/executor/helps/token_helpers.go
Normal 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)
|
||||
}
|
||||
}
|
||||
125
backend/internal/runtime/executor/helps/transport_cache.go
Normal file
125
backend/internal/runtime/executor/helps/transport_cache.go
Normal 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)
|
||||
}
|
||||
172
backend/internal/runtime/executor/helps/transport_cache_test.go
Normal file
172
backend/internal/runtime/executor/helps/transport_cache_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
1170
backend/internal/runtime/executor/helps/usage_helpers.go
Normal file
1170
backend/internal/runtime/executor/helps/usage_helpers.go
Normal file
File diff suppressed because it is too large
Load diff
753
backend/internal/runtime/executor/helps/usage_helpers_test.go
Normal file
753
backend/internal/runtime/executor/helps/usage_helpers_test.go
Normal 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"
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
150
backend/internal/runtime/executor/helps/user_id_cache.go
Normal file
150
backend/internal/runtime/executor/helps/user_id_cache.go
Normal 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)
|
||||
}
|
||||
196
backend/internal/runtime/executor/helps/user_id_cache_test.go
Normal file
196
backend/internal/runtime/executor/helps/user_id_cache_test.go
Normal 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")
|
||||
}
|
||||
}
|
||||
407
backend/internal/runtime/executor/helps/utls_client.go
Normal file
407
backend/internal/runtime/executor/helps/utls_client.go
Normal 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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
641
backend/internal/runtime/executor/helps/utls_client_test.go
Normal file
641
backend/internal/runtime/executor/helps/utls_client_test.go
Normal 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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue