From ab9c06b42ecc2db6212eaca7a4f226710dbd43e6 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 19 Oct 2026 21:57:53 +0800 Subject: [PATCH] feat(cache): add Home KV support to Claude web search result cache - Store and retrieve cached Claude web search results via Home KV with sliding TTL expiration when Home mode is enabled. - Propagate request context through Claude web search result recording and restoration helpers in execution and token-counting paths. Closes: #6780 --- .../cache/claude_web_search_result_cache.go | 85 ++++++++- .../claude_web_search_result_cache_test.go | 168 ++++++++++++++++++ .../executor/claude_executor_execute.go | 4 +- .../executor/claude_executor_stream.go | 4 +- .../executor/claude_executor_tokens.go | 4 +- .../helps/claude_web_search_results.go | 26 ++- 6 files changed, 274 insertions(+), 17 deletions(-) diff --git a/internal/cache/claude_web_search_result_cache.go b/internal/cache/claude_web_search_result_cache.go index a37cd46d66..6573881e67 100644 --- a/internal/cache/claude_web_search_result_cache.go +++ b/internal/cache/claude_web_search_result_cache.go @@ -3,11 +3,15 @@ package cache import ( "bytes" "container/list" + "context" "crypto/sha256" "encoding/hex" "strings" "sync" "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" + log "github.com/sirupsen/logrus" ) const ( @@ -29,6 +33,16 @@ var ( claudeWebSearchResultNow = time.Now ) +type claudeWebSearchResultKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) +} + +var currentClaudeWebSearchResultKVClient = func() (claudeWebSearchResultKVClient, bool, error) { + return homekv.CurrentKVClient() +} + func claudeWebSearchResultKey(sessionKey, toolUseID string) (string, bool) { sessionKey = strings.TrimSpace(sessionKey) toolUseID = strings.TrimSpace(toolUseID) @@ -39,17 +53,46 @@ func claudeWebSearchResultKey(sessionKey, toolUseID string) (string, bool) { return hex.EncodeToString(sum[:]), true } +func claudeWebSearchResultKVKey(sessionKey, toolUseID string) string { + key, ok := claudeWebSearchResultKey(sessionKey, toolUseID) + if !ok { + return "" + } + return "cpa:claude:web-search-result:" + key +} + // CacheClaudeWebSearchResult stores the raw content of one web_search_tool_result // block for a session. Memory is bounded by the maximum entry count (limit); // least recently used entries are evicted when the entry limit is exceeded. func CacheClaudeWebSearchResult(sessionKey, toolUseID string, content []byte, ttl time.Duration, limit int) { + CacheClaudeWebSearchResultBestEffort(context.Background(), sessionKey, toolUseID, content, ttl, limit) +} + +// CacheClaudeWebSearchResultBestEffort stores the raw content of one +// web_search_tool_result block for completed response paths. +func CacheClaudeWebSearchResultBestEffort(ctx context.Context, sessionKey, toolUseID string, content []byte, ttl time.Duration, limit int) bool { key, ok := claudeWebSearchResultKey(sessionKey, toolUseID) if !ok || len(content) == 0 { - return + return false } if ttl <= 0 { ttl = defaultClaudeWebSearchResultTTL } + if ctx == nil { + ctx = context.Background() + } + if client, homeMode, errClient := currentClaudeWebSearchResultKVClient(); homeMode { + if errClient != nil { + log.Errorf("home kv best-effort claude web search result set failed prefix=cpa:claude:*: %v", errClient) + return false + } + written, errSet := client.KVSet(ctx, "cpa:claude:web-search-result:"+key, bytes.Clone(content), homekv.KVSetOptions{EX: ttl}) + if errSet != nil { + log.Errorf("home kv best-effort claude web search result set failed prefix=cpa:claude:*: %v", errSet) + return false + } + return written + } if limit <= 0 { limit = defaultClaudeWebSearchResultLimit } @@ -79,27 +122,59 @@ func CacheClaudeWebSearchResult(sessionKey, toolUseID string, content []byte, tt } removeClaudeWebSearchResultLocked(oldest) } + return true } // GetClaudeWebSearchResult returns a copy of the stored content and extends its // lifetime. Expired entries are removed and reported as missing. func GetClaudeWebSearchResult(sessionKey, toolUseID string, ttl ...time.Duration) ([]byte, bool) { + content, ok, errGet := GetClaudeWebSearchResultRequired(context.Background(), sessionKey, toolUseID, ttl...) + if errGet != nil { + return nil, false + } + return content, ok +} + +// GetClaudeWebSearchResultRequired retrieves a copy of the stored content for +// request-time paths and extends its lifetime. +func GetClaudeWebSearchResultRequired(ctx context.Context, sessionKey, toolUseID string, ttl ...time.Duration) ([]byte, bool, error) { key, ok := claudeWebSearchResultKey(sessionKey, toolUseID) if !ok { - return nil, false + return nil, false, nil + } + if ctx == nil { + ctx = context.Background() + } + if client, homeMode, errClient := currentClaudeWebSearchResultKVClient(); homeMode { + if errClient != nil { + return nil, false, errClient + } + kvKey := "cpa:claude:web-search-result:" + key + raw, found, errGet := client.KVGet(ctx, kvKey) + if errGet != nil || !found || len(raw) == 0 { + return nil, false, errGet + } + effectiveTTL := defaultClaudeWebSearchResultTTL + if len(ttl) > 0 && ttl[0] > 0 { + effectiveTTL = ttl[0] + } + if _, errExpire := client.KVExpire(ctx, kvKey, effectiveTTL); errExpire != nil { + log.Warnf("home kv claude web search result expire failed prefix=cpa:claude:*: %v", errExpire) + } + return bytes.Clone(raw), true, nil } claudeWebSearchResultMu.Lock() defer claudeWebSearchResultMu.Unlock() element, exists := claudeWebSearchResultIndex[key] if !exists { - return nil, false + return nil, false, nil } entry := element.Value.(*claudeWebSearchResultEntry) now := claudeWebSearchResultNow() if !now.Before(entry.expiresAt) { removeClaudeWebSearchResultLocked(element) - return nil, false + return nil, false, nil } effectiveTTL := entry.ttl if len(ttl) > 0 && ttl[0] > 0 { @@ -111,7 +186,7 @@ func GetClaudeWebSearchResult(sessionKey, toolUseID string, ttl ...time.Duration } entry.expiresAt = now.Add(effectiveTTL) claudeWebSearchResultOrder.MoveToFront(element) - return bytes.Clone(entry.content), true + return bytes.Clone(entry.content), true, nil } // ClearClaudeWebSearchResultCache removes every stored entry. diff --git a/internal/cache/claude_web_search_result_cache_test.go b/internal/cache/claude_web_search_result_cache_test.go index 472850f8b1..51141ad5ee 100644 --- a/internal/cache/claude_web_search_result_cache_test.go +++ b/internal/cache/claude_web_search_result_cache_test.go @@ -2,12 +2,74 @@ package cache import ( "bytes" + "context" + "errors" "strconv" "strings" "testing" "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" ) +type fakeClaudeWebSearchResultKVClient struct { + values map[string][]byte + getErr error + setErr error + expireErr error + getCount int + setCount int + expireCount int + lastSetTTL time.Duration + lastExpireTTL time.Duration +} + +func newFakeClaudeWebSearchResultKVClient() *fakeClaudeWebSearchResultKVClient { + return &fakeClaudeWebSearchResultKVClient{values: make(map[string][]byte)} +} + +func (c *fakeClaudeWebSearchResultKVClient) 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 bytes.Clone(value), true, nil +} + +func (c *fakeClaudeWebSearchResultKVClient) 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] = bytes.Clone(value) + return true, nil +} + +func (c *fakeClaudeWebSearchResultKVClient) 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 useFakeClaudeWebSearchResultKVClient(t *testing.T, client *fakeClaudeWebSearchResultKVClient, homeMode bool, errClient error) { + t.Helper() + previous := currentClaudeWebSearchResultKVClient + currentClaudeWebSearchResultKVClient = func() (claudeWebSearchResultKVClient, bool, error) { + return client, homeMode, errClient + } + t.Cleanup(func() { + currentClaudeWebSearchResultKVClient = previous + }) +} + func useClaudeWebSearchResultClock(t *testing.T, now *time.Time) { t.Helper() ClearClaudeWebSearchResultCache() @@ -124,3 +186,109 @@ func TestClaudeWebSearchResultCachePurgesExpiredBeforeEvictingLiveEntries(t *tes t.Fatal("new entry missing") } } + +func TestClaudeWebSearchResultCacheHomeModeDoesNotUseLocalCache(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + useClaudeWebSearchResultClock(t, &now) + + CacheClaudeWebSearchResult("session-local", "srvtoolu_local", []byte(`[{"url":"local"}]`), 30*time.Minute, 16) + + homekv.SetCurrent(&homekv.Client{}) + t.Cleanup(homekv.ClearCurrent) + + if got, ok := GetClaudeWebSearchResult("session-local", "srvtoolu_local", 30*time.Minute); ok { + t.Fatalf("GetClaudeWebSearchResult() = %q, true; want Home miss without local fallback", got) + } + + CacheClaudeWebSearchResult("session-home", "srvtoolu_home", []byte(`[{"url":"home"}]`), 30*time.Minute, 16) + homekv.ClearCurrent() + + if got, ok := GetClaudeWebSearchResult("session-home", "srvtoolu_home", 30*time.Minute); ok { + t.Fatalf("local cache = %q, true; want empty after Home write failure", got) + } +} + +func TestClaudeWebSearchResultCacheHomeWriteReadAndSlidingExpire(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + useClaudeWebSearchResultClock(t, &now) + + client := newFakeClaudeWebSearchResultKVClient() + useFakeClaudeWebSearchResultKVClient(t, client, true, nil) + + content := []byte(`[{"type":"web_search_result","url":"https://example.com","encrypted_content":"enc"}]`) + if !CacheClaudeWebSearchResultBestEffort(context.Background(), "session-home", "srvtoolu_1", content, 45*time.Minute, 16) { + t.Fatal("CacheClaudeWebSearchResultBestEffort() = false, want true") + } + if claudeWebSearchResultOrder.Len() != 0 { + t.Fatalf("local cache len = %d, want 0 in Home mode", claudeWebSearchResultOrder.Len()) + } + kvKey := claudeWebSearchResultKVKey("session-home", "srvtoolu_1") + if client.setCount != 1 || client.lastSetTTL != 45*time.Minute { + t.Fatalf("KVSet count/ttl = %d/%v, want 1/45m", client.setCount, client.lastSetTTL) + } + if !bytes.Equal(client.values[kvKey], content) { + t.Fatalf("KV stored value = %q, want %q", client.values[kvKey], content) + } + + got, ok, errGet := GetClaudeWebSearchResultRequired(context.Background(), "session-home", "srvtoolu_1", 30*time.Minute) + if errGet != nil || !ok || !bytes.Equal(got, content) { + t.Fatalf("GetClaudeWebSearchResultRequired() = %q, %v, %v; want %q, true, nil", got, ok, errGet, content) + } + if client.expireCount != 1 || client.lastExpireTTL != 30*time.Minute { + t.Fatalf("KVExpire count/ttl = %d/%v, want 1/30m", client.expireCount, client.lastExpireTTL) + } + got[0] = 'X' + if !bytes.Equal(client.values[kvKey], content) { + t.Fatalf("caller mutation altered Home KV fake state: %q", client.values[kvKey]) + } + + // Expire error should not discard a successfully fetched result. + client.expireErr = errors.New("expire failed") + gotAfterExpireErr, okAfterExpireErr, errAfterExpireErr := GetClaudeWebSearchResultRequired(context.Background(), "session-home", "srvtoolu_1") + if errAfterExpireErr != nil || !okAfterExpireErr || !bytes.Equal(gotAfterExpireErr, content) { + t.Fatalf("GetClaudeWebSearchResultRequired() with expireErr = %q, %v, %v; want %q, true, nil", gotAfterExpireErr, okAfterExpireErr, errAfterExpireErr, content) + } + if client.expireCount != 2 || client.lastExpireTTL != defaultClaudeWebSearchResultTTL { + t.Fatalf("KVExpire default ttl = %d/%v, want 2/%v", client.expireCount, client.lastExpireTTL, defaultClaudeWebSearchResultTTL) + } +} + +func TestClaudeWebSearchResultCacheHomeRejectsEmptyScopeAndPropagatesErrors(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + useClaudeWebSearchResultClock(t, &now) + + client := newFakeClaudeWebSearchResultKVClient() + useFakeClaudeWebSearchResultKVClient(t, client, true, nil) + + if CacheClaudeWebSearchResultBestEffort(context.Background(), "", "srvtoolu_1", []byte(`[1]`), time.Minute, 16) { + t.Fatal("empty session key should be rejected") + } + if CacheClaudeWebSearchResultBestEffort(context.Background(), "s", "", []byte(`[1]`), time.Minute, 16) { + t.Fatal("empty toolUseID should be rejected") + } + if CacheClaudeWebSearchResultBestEffort(context.Background(), "s", "srvtoolu_1", nil, time.Minute, 16) { + t.Fatal("empty content should be rejected") + } + if _, ok, errGet := GetClaudeWebSearchResultRequired(context.Background(), "", "srvtoolu_1"); errGet != nil || ok { + t.Fatalf("empty session key Get = %v, %v; want false, nil", ok, errGet) + } + if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 { + t.Fatalf("KV calls = get %d set %d expire %d, want all 0", client.getCount, client.setCount, client.expireCount) + } + + client.setErr = errors.New("set failed") + if CacheClaudeWebSearchResultBestEffort(context.Background(), "s", "srvtoolu_1", []byte(`[1]`), 0, 16) { + t.Fatal("CacheClaudeWebSearchResultBestEffort() = true on setErr, want false") + } + if client.lastSetTTL != defaultClaudeWebSearchResultTTL { + t.Fatalf("zero ttl defaulted to %v, want %v", client.lastSetTTL, defaultClaudeWebSearchResultTTL) + } + + client.getErr = errors.New("get failed") + if _, ok, errGet := GetClaudeWebSearchResultRequired(context.Background(), "s", "srvtoolu_1"); errGet == nil || ok { + t.Fatalf("GetClaudeWebSearchResultRequired() on getErr = %v, %v; want false, error", ok, errGet) + } + if _, ok := GetClaudeWebSearchResult("s", "srvtoolu_1"); ok { + t.Fatal("GetClaudeWebSearchResult() on getErr = true, want false") + } +} diff --git a/internal/runtime/executor/claude_executor_execute.go b/internal/runtime/executor/claude_executor_execute.go index b6803635b8..3d0806d43a 100644 --- a/internal/runtime/executor/claude_executor_execute.go +++ b/internal/runtime/executor/claude_executor_execute.go @@ -91,7 +91,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r return resp, err } webSearchSessionKey, webSearchTTL, webSearchLimit := e.claudeWebSearchResultCacheScope(ctx, from, req, opts, originalPayload) - body = helps.RestoreClaudeWebSearchResults(body, webSearchSessionKey, webSearchTTL) + body = helps.RestoreClaudeWebSearchResultsContext(ctx, body, webSearchSessionKey, webSearchTTL) if rebuildMidSystemMessageEnabled(e.cfg, auth) { body = rebuildMidSystemMessagesToTopLevel(body) } @@ -396,7 +396,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r } } data = e.restoreResponseModel(data, req.Model) - helps.RecordClaudeWebSearchResults(webSearchSessionKey, data, webSearchTTL, webSearchLimit) + helps.RecordClaudeWebSearchResultsContext(ctx, webSearchSessionKey, data, webSearchTTL, webSearchLimit) cacheClaudeThinkingReplayResponse(ctx, replayScope, data) var param any out := sdktranslator.TranslateNonStream( diff --git a/internal/runtime/executor/claude_executor_stream.go b/internal/runtime/executor/claude_executor_stream.go index 594012027e..5f8b0cdc9e 100644 --- a/internal/runtime/executor/claude_executor_stream.go +++ b/internal/runtime/executor/claude_executor_stream.go @@ -94,7 +94,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A return nil, err } webSearchSessionKey, webSearchTTL, webSearchLimit := e.claudeWebSearchResultCacheScope(ctx, from, req, opts, originalPayload) - body = helps.RestoreClaudeWebSearchResults(body, webSearchSessionKey, webSearchTTL) + body = helps.RestoreClaudeWebSearchResultsContext(ctx, body, webSearchSessionKey, webSearchTTL) if rebuildMidSystemMessageEnabled(e.cfg, auth) { body = rebuildMidSystemMessagesToTopLevel(body) } @@ -476,7 +476,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A return } line = e.restoreResponseModel(restoredLine, req.Model) - helps.RecordClaudeWebSearchResults(webSearchSessionKey, line, webSearchTTL, webSearchLimit) + helps.RecordClaudeWebSearchResultsContext(ctx, webSearchSessionKey, line, webSearchTTL, webSearchLimit) chunks := sdktranslator.TranslateStream( ctx, to, diff --git a/internal/runtime/executor/claude_executor_tokens.go b/internal/runtime/executor/claude_executor_tokens.go index 8d2c9c0773..29eda087a0 100644 --- a/internal/runtime/executor/claude_executor_tokens.go +++ b/internal/runtime/executor/claude_executor_tokens.go @@ -57,7 +57,7 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut originalPayload = opts.OriginalRequest } webSearchSessionKey, webSearchTTL, _ := e.claudeWebSearchResultCacheScope(ctx, from, req, opts, originalPayload) - body = helps.RestoreClaudeWebSearchResults(body, webSearchSessionKey, webSearchTTL) + body = helps.RestoreClaudeWebSearchResultsContext(ctx, body, webSearchSessionKey, webSearchTTL) if rebuildMidSystemMessageEnabled(e.cfg, auth) { body = rebuildMidSystemMessagesToTopLevel(body) } @@ -176,7 +176,7 @@ func (e *ClaudeExecutor) countTokensUpstream(ctx context.Context, auth *cliproxy return cliproxyexecutor.Response{}, errThinking } webSearchSessionKey, webSearchTTL, _ := e.claudeWebSearchResultCacheScope(ctx, from, req, opts, originalPayload) - body = helps.RestoreClaudeWebSearchResults(body, webSearchSessionKey, webSearchTTL) + body = helps.RestoreClaudeWebSearchResultsContext(ctx, body, webSearchSessionKey, webSearchTTL) if rebuildMidSystemMessageEnabled(e.cfg, auth) { body = rebuildMidSystemMessagesToTopLevel(body) } diff --git a/internal/runtime/executor/helps/claude_web_search_results.go b/internal/runtime/executor/helps/claude_web_search_results.go index 27179ffe6d..97757dcdda 100644 --- a/internal/runtime/executor/helps/claude_web_search_results.go +++ b/internal/runtime/executor/helps/claude_web_search_results.go @@ -2,6 +2,7 @@ package helps import ( "bytes" + "context" "fmt" "strings" "time" @@ -19,12 +20,18 @@ var claudeWebSearchToolResultMarker = []byte(`"web_search_tool_result"`) // or SSE text (one or more lines); the result block arrives whole in its // content_block_start event. func RecordClaudeWebSearchResults(sessionKey string, data []byte, ttl time.Duration, limit int) { + RecordClaudeWebSearchResultsContext(context.Background(), sessionKey, data, ttl, limit) +} + +// RecordClaudeWebSearchResultsContext stores the web_search_tool_result blocks +// of a Claude response for sessionKey using ctx for Home KV operations. +func RecordClaudeWebSearchResultsContext(ctx context.Context, sessionKey string, data []byte, ttl time.Duration, limit int) { if strings.TrimSpace(sessionKey) == "" || !bytes.Contains(data, claudeWebSearchToolResultMarker) { return } if content := gjson.GetBytes(data, "content"); content.IsArray() { content.ForEach(func(_, block gjson.Result) bool { - recordClaudeWebSearchResultBlock(sessionKey, block, ttl, limit) + recordClaudeWebSearchResultBlock(ctx, sessionKey, block, ttl, limit) return true }) return @@ -36,12 +43,12 @@ func RecordClaudeWebSearchResults(sessionKey string, data []byte, ttl time.Durat } event := gjson.ParseBytes(bytes.TrimSpace(line[len("data:"):])) if event.Get("type").String() == "content_block_start" { - recordClaudeWebSearchResultBlock(sessionKey, event.Get("content_block"), ttl, limit) + recordClaudeWebSearchResultBlock(ctx, sessionKey, event.Get("content_block"), ttl, limit) } } } -func recordClaudeWebSearchResultBlock(sessionKey string, block gjson.Result, ttl time.Duration, limit int) { +func recordClaudeWebSearchResultBlock(ctx context.Context, sessionKey string, block gjson.Result, ttl time.Duration, limit int) { if block.Get("type").String() != "web_search_tool_result" { return } @@ -49,7 +56,7 @@ func recordClaudeWebSearchResultBlock(sessionKey string, block gjson.Result, ttl if !content.Exists() || claudeWebSearchResultContentEmpty(content) { return } - cache.CacheClaudeWebSearchResult(sessionKey, block.Get("tool_use_id").String(), []byte(content.Raw), ttl, limit) + cache.CacheClaudeWebSearchResultBestEffort(ctx, sessionKey, block.Get("tool_use_id").String(), []byte(content.Raw), ttl, limit) } // RestoreClaudeWebSearchResults fills assistant web_search_tool_result blocks @@ -57,6 +64,13 @@ func recordClaudeWebSearchResultBlock(sessionKey string, block gjson.Result, ttl // Blocks that already carry content, and blocks without a stored entry, are // left unchanged. func RestoreClaudeWebSearchResults(body []byte, sessionKey string, ttl ...time.Duration) []byte { + return RestoreClaudeWebSearchResultsContext(context.Background(), body, sessionKey, ttl...) +} + +// RestoreClaudeWebSearchResultsContext fills assistant web_search_tool_result +// blocks whose content is missing or empty with the content stored for +// sessionKey using ctx for Home KV operations. +func RestoreClaudeWebSearchResultsContext(ctx context.Context, body []byte, sessionKey string, ttl ...time.Duration) []byte { if strings.TrimSpace(sessionKey) == "" || !bytes.Contains(body, claudeWebSearchToolResultMarker) { return body } @@ -73,8 +87,8 @@ func RestoreClaudeWebSearchResults(body []byte, sessionKey string, ttl ...time.D if block.Get("type").String() != "web_search_tool_result" || !claudeWebSearchResultContentEmpty(block.Get("content")) { continue } - cached, ok := cache.GetClaudeWebSearchResult(sessionKey, block.Get("tool_use_id").String(), ttl...) - if !ok { + cached, ok, errGet := cache.GetClaudeWebSearchResultRequired(ctx, sessionKey, block.Get("tool_use_id").String(), ttl...) + if errGet != nil || !ok { continue } updated, errSet := sjson.SetRawBytes(body, fmt.Sprintf("messages.%d.content.%d.content", i, j), cached)