mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-10-24 16:45:45 +08:00
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:
85
internal/cache/claude_web_search_result_cache.go
vendored
85
internal/cache/claude_web_search_result_cache.go
vendored
@@ -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.
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user