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
This commit is contained in:
Luis Pater
2026-10-19 21:57:53 +08:00
parent 154d5144fd
commit ab9c06b42e
6 changed files with 274 additions and 17 deletions

View File

@@ -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.

View File

@@ -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")
}
}

View File

@@ -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(

View File

@@ -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,

View File

@@ -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)
}

View File

@@ -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)