diff --git a/internal/agent/agent.go b/internal/agent/agent.go index ba287c1d..e75d4966 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -188,6 +188,13 @@ func New(args Args) *Agent { // Run executes the full review pipeline: parse diffs -> plan per file -> LLM tool-loop -> collect comments. func (a *Agent) Run(ctx context.Context) ([]model.LlmComment, error) { + // Base prompt-cache affinity key for any LLM request in this run that a + // task doesn't re-scope. Each task conversation (plan, per-file main + // loop, compression, ...) refines it with llm.SessionTaskKey where it + // starts, so affinity keys stay per-conversation — the granularity + // provider prompt caches actually reuse prefixes at. + ctx = llm.ContextWithSessionKey(ctx, a.SessionID()) + // Step 1: Parse diffs ctx, diffSpan := telemetry.StartSpan(ctx, "diff.parse") if err := a.loadDiffs(ctx); err != nil { @@ -670,6 +677,8 @@ func (a *Agent) executeReviewFilter(ctx context.Context, d model.Diff, newPath s fs := a.session.GetOrCreateFileSession(newPath) rec := fs.AppendTaskRecord(session.ReviewFilterTask, messages) + ctx = llm.ContextWithSessionKey(ctx, + llm.SessionTaskKey(a.session.SessionID, string(session.ReviewFilterTask), newPath)) startTime := time.Now() _, llmSpan := telemetry.StartLLMSpan(ctx, a.args.Model) @@ -893,6 +902,8 @@ func (a *Agent) executePlanPhase(ctx context.Context, newPath, rawDiff, changeFi fs := a.session.GetOrCreateFileSession(newPath) rec := fs.AppendTaskRecord(session.PlanTask, messages) + ctx = llm.ContextWithSessionKey(ctx, + llm.SessionTaskKey(a.session.SessionID, string(session.PlanTask), newPath)) startTime := time.Now() _, llmSpan := telemetry.StartLLMSpan(ctx, a.args.Model) diff --git a/internal/llm/client.go b/internal/llm/client.go index 1ecc01e2..5d429097 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -189,6 +189,14 @@ type ClientConfig struct { Timeout time.Duration // Request timeout ExtraBody map[string]any // Vendor-specific fields merged into every request body ExtraHeaders map[string]string // Extra HTTP headers sent with every request + // SessionKey is the fallback prompt-cache affinity key for requests + // whose context carries none (see ContextWithSessionKey). Review and + // scan runs tag every request context with the real session's ID, so + // this fallback only serves session-less callers such as `ocr llm + // test`. Auto-generated when empty. The effective key replaces + // SessionKeyTemplateVar in ExtraHeaders/ExtraBody values; embedding + // that placeholder is how users opt into session affinity. + SessionKey string } // --- Factory --- @@ -297,10 +305,15 @@ type OpenAIClient struct { } // NewOpenAIClient creates a new OpenAI-compatible LLM client. +// ExtraHeaders are applied per request (not baked into the SDK client) so +// SessionKeyTemplateVar can expand to the session key each request carries. func NewOpenAIClient(cfg ClientConfig) *OpenAIClient { if cfg.Timeout <= 0 { cfg.Timeout = 5 * time.Minute } + if cfg.SessionKey == "" { + cfg.SessionKey = NewSessionKey() + } baseURL := strings.TrimRight(cfg.URL, "/") if !strings.HasSuffix(baseURL, "/chat/completions") { cfg.URL = baseURL + "/chat/completions" @@ -315,9 +328,6 @@ func NewOpenAIClient(cfg ClientConfig) *OpenAIClient { openaiopt.WithHeader("User-Agent", userAgent("")), openaiopt.WithRequestTimeout(cfg.Timeout), } - for k, v := range cfg.ExtraHeaders { - opts = append(opts, openaiopt.WithHeader(k, v)) - } return &OpenAIClient{ cfg: cfg, @@ -344,8 +354,16 @@ func (c *OpenAIClient) CompletionsWithCtx(ctx context.Context, req ChatRequest) params := c.buildOpenAIParams(model, req) + sessionKey := c.cfg.SessionKey + if k := SessionKeyFromContext(ctx); k != "" { + sessionKey = k + } + var opts []openaiopt.RequestOption - for k, v := range c.cfg.ExtraBody { + for k, v := range expandSessionKeyInHeaders(c.cfg.ExtraHeaders, sessionKey) { + opts = append(opts, openaiopt.WithHeader(k, v)) + } + for k, v := range expandSessionKeyInBody(c.cfg.ExtraBody, sessionKey) { opts = append(opts, openaiopt.WithJSONSet(k, v)) } if stream, ok := c.cfg.ExtraBody["stream"].(bool); ok && stream { @@ -570,10 +588,17 @@ type AnthropicClient struct { } // NewAnthropicClient creates a new Anthropic Messages API client. +// The Anthropic API manages prompt-cache affinity server-side, so unlike the +// OpenAI client no session key body field is injected; the key is still +// available to ExtraHeaders/ExtraBody via SessionKeyTemplateVar, applied per +// request so it can expand to the session key each request carries. func NewAnthropicClient(cfg ClientConfig) *AnthropicClient { if cfg.Timeout <= 0 { cfg.Timeout = 5 * time.Minute } + if cfg.SessionKey == "" { + cfg.SessionKey = NewSessionKey() + } if !strings.HasSuffix(cfg.URL, "/v1/messages") && !strings.HasSuffix(cfg.URL, "/v1/messages/") { baseURL := strings.TrimRight(cfg.URL, "/") if !strings.HasSuffix(baseURL, "/v1/messages") { @@ -608,10 +633,6 @@ func NewAnthropicClient(cfg ClientConfig) *AnthropicClient { ) } - for k, v := range cfg.ExtraHeaders { - opts = append(opts, option.WithHeader(k, v)) - } - return &AnthropicClient{ cfg: cfg, sdk: anthropic.NewClient(opts...), @@ -630,8 +651,16 @@ func (c *AnthropicClient) CompletionsWithCtx(ctx context.Context, req ChatReques return nil, err } + sessionKey := c.cfg.SessionKey + if k := SessionKeyFromContext(ctx); k != "" { + sessionKey = k + } + var opts []option.RequestOption - for k, v := range c.cfg.ExtraBody { + for k, v := range expandSessionKeyInHeaders(c.cfg.ExtraHeaders, sessionKey) { + opts = append(opts, option.WithHeader(k, v)) + } + for k, v := range expandSessionKeyInBody(c.cfg.ExtraBody, sessionKey) { opts = append(opts, option.WithJSONSet(k, v)) } diff --git a/internal/llm/client_test.go b/internal/llm/client_test.go index 354219a5..b368d6c9 100644 --- a/internal/llm/client_test.go +++ b/internal/llm/client_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http" "net/http/httptest" "strings" @@ -1041,6 +1042,297 @@ func TestAnthropicClient_NoExtraHeadersWhenEmpty(t *testing.T) { } } +// newOpenAITestServer returns an httptest server that captures each request's +// headers and JSON body and responds with a minimal valid chat completion. +func newOpenAITestServer(gotHeaders *http.Header, gotBody *map[string]any) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if gotHeaders != nil { + *gotHeaders = r.Header.Clone() + } + if gotBody != nil { + raw, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(raw, gotBody) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "id":"chatcmpl-test", + "object":"chat.completion", + "model":"gpt-test", + "choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}], + "usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2} + }`)) + })) +} + +func TestOpenAIClient_SessionKeyExpandedInExtraHeaders(t *testing.T) { + var gotHeaders http.Header + server := newOpenAITestServer(&gotHeaders, nil) + defer server.Close() + + client := NewOpenAIClient(ClientConfig{ + URL: server.URL + "/v1", + APIKey: "test-key", + Model: "gpt-test", + SessionKey: "sess-abc", + ExtraHeaders: map[string]string{ + "x-session-affinity": "{ocr_session_key}", + }, + }) + + _, err := client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if got := gotHeaders.Get("X-Session-Affinity"); got != "sess-abc" { + t.Errorf("x-session-affinity = %q, want %q", got, "sess-abc") + } +} + +func TestOpenAIClient_SessionKeyExpandedInExtraBody(t *testing.T) { + var gotBody map[string]any + server := newOpenAITestServer(nil, &gotBody) + defer server.Close() + + // The documented recipe for OpenAI prompt caching: route the session key + // into prompt_cache_key via extra_body. + client := NewOpenAIClient(ClientConfig{ + URL: server.URL + "/v1", + APIKey: "test-key", + Model: "gpt-test", + SessionKey: "sess-abc", + ExtraBody: map[string]any{"prompt_cache_key": "{ocr_session_key}"}, + }) + + _, err := client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if got := gotBody["prompt_cache_key"]; got != "sess-abc" { + t.Errorf("prompt_cache_key = %v, want %q", got, "sess-abc") + } +} + +func TestOpenAIClient_NoInjectionWithoutPlaceholder(t *testing.T) { + var gotBody map[string]any + server := newOpenAITestServer(nil, &gotBody) + defer server.Close() + + // Without an {ocr_session_key} placeholder anywhere, OCR must not add + // any session-related field to the request. + client := NewOpenAIClient(ClientConfig{ + URL: server.URL + "/v1", + APIKey: "test-key", + Model: "gpt-test", + }) + + _, err := client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if _, ok := gotBody["prompt_cache_key"]; ok { + t.Errorf("prompt_cache_key = %v, want absent without explicit configuration", gotBody["prompt_cache_key"]) + } +} + +func TestOpenAIClient_SessionKeyGeneratedWhenEmpty(t *testing.T) { + var gotHeaders http.Header + server := newOpenAITestServer(&gotHeaders, nil) + defer server.Close() + + client := NewOpenAIClient(ClientConfig{ + URL: server.URL + "/v1", + APIKey: "test-key", + Model: "gpt-test", + ExtraHeaders: map[string]string{ + "x-session-affinity": "{ocr_session_key}", + }, + }) + + _, err := client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + got := gotHeaders.Get("X-Session-Affinity") + if got == "" || got == "{ocr_session_key}" { + t.Errorf("x-session-affinity = %q, want an auto-generated session key", got) + } +} + +func TestAnthropicClient_SessionKeyExpandedInExtraHeadersAndBody(t *testing.T) { + var gotAffinity string + var gotBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAffinity = r.Header.Get("x-session-affinity") + raw, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(raw, &gotBody) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "id":"msg_test", + "type":"message", + "role":"assistant", + "model":"claude-test", + "content":[{"type":"text","text":"ok"}], + "stop_reason":"end_turn", + "usage":{"input_tokens":1,"output_tokens":1} + }`)) + })) + defer server.Close() + + client := NewAnthropicClient(ClientConfig{ + URL: server.URL + "/v1/messages", + APIKey: "test-key", + Model: "claude-test", + SessionKey: "sess-xyz", + ExtraHeaders: map[string]string{ + "x-session-affinity": "{ocr_session_key}", + }, + ExtraBody: map[string]any{"cache_key": "{ocr_session_key}"}, + }) + + _, err := client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if gotAffinity != "sess-xyz" { + t.Errorf("x-session-affinity = %q, want %q", gotAffinity, "sess-xyz") + } + if got := gotBody["cache_key"]; got != "sess-xyz" { + t.Errorf("cache_key = %v, want %q", got, "sess-xyz") + } +} + +func TestNewLLMClient_ExpandsSessionKeyInExtraBody(t *testing.T) { + var gotBody map[string]any + server := newOpenAITestServer(nil, &gotBody) + defer server.Close() + + client := NewLLMClient(ResolvedEndpoint{ + URL: server.URL + "/v1", + Token: "test-key", + Model: "gpt-test", + Protocol: "openai", + ExtraBody: map[string]any{"prompt_cache_key": "{ocr_session_key}"}, + }) + + ctx := ContextWithSessionKey(context.Background(), "sess-from-run") + _, err := client.CompletionsWithCtx(ctx, ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if got := gotBody["prompt_cache_key"]; got != "sess-from-run" { + t.Errorf("prompt_cache_key = %v, want %q", got, "sess-from-run") + } +} + +func TestOpenAIClient_ContextSessionKeyOverridesFallback(t *testing.T) { + var gotHeaders http.Header + var gotBody map[string]any + server := newOpenAITestServer(&gotHeaders, &gotBody) + defer server.Close() + + client := NewOpenAIClient(ClientConfig{ + URL: server.URL + "/v1", + APIKey: "test-key", + Model: "gpt-test", + SessionKey: "fallback-key", + ExtraBody: map[string]any{"prompt_cache_key": "{ocr_session_key}"}, + ExtraHeaders: map[string]string{ + "x-session-affinity": "{ocr_session_key}", + }, + }) + + // A context-carried session key (the real OCR session's ID) must win + // over the client's construction-time fallback. + ctx := ContextWithSessionKey(context.Background(), "real-session-id") + _, err := client.CompletionsWithCtx(ctx, ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if got := gotHeaders.Get("X-Session-Affinity"); got != "real-session-id" { + t.Errorf("x-session-affinity = %q, want %q", got, "real-session-id") + } + if got := gotBody["prompt_cache_key"]; got != "real-session-id" { + t.Errorf("prompt_cache_key = %v, want %q", got, "real-session-id") + } + + // Without a context key the fallback applies. + _, err = client.CompletionsWithCtx(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if got := gotHeaders.Get("X-Session-Affinity"); got != "fallback-key" { + t.Errorf("x-session-affinity = %q, want %q", got, "fallback-key") + } +} + +func TestAnthropicClient_ContextSessionKeyOverridesFallback(t *testing.T) { + var gotAffinity string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAffinity = r.Header.Get("x-session-affinity") + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "id":"msg_test", + "type":"message", + "role":"assistant", + "model":"claude-test", + "content":[{"type":"text","text":"ok"}], + "stop_reason":"end_turn", + "usage":{"input_tokens":1,"output_tokens":1} + }`)) + })) + defer server.Close() + + client := NewAnthropicClient(ClientConfig{ + URL: server.URL + "/v1/messages", + APIKey: "test-key", + Model: "claude-test", + SessionKey: "fallback-key", + ExtraHeaders: map[string]string{ + "x-session-affinity": "{ocr_session_key}", + }, + }) + + ctx := ContextWithSessionKey(context.Background(), "real-session-id") + _, err := client.CompletionsWithCtx(ctx, ChatRequest{ + Messages: []Message{{Role: "user", Content: "ping"}}, + MaxTokens: 64, + }) + if err != nil { + t.Fatalf("CompletionsWithCtx: %v", err) + } + if gotAffinity != "real-session-id" { + t.Errorf("x-session-affinity = %q, want %q", gotAffinity, "real-session-id") + } +} + // Verify the SDK constant is accessible (compile-time check). var _ anthropic.CacheControlEphemeralParam = anthropic.NewCacheControlEphemeralParam() diff --git a/internal/llm/resolver_test.go b/internal/llm/resolver_test.go index c9953768..6d5cc328 100644 --- a/internal/llm/resolver_test.go +++ b/internal/llm/resolver_test.go @@ -1190,6 +1190,38 @@ func TestResolveEndpoint_EnvExtraHeadersMergedWithConfigFile(t *testing.T) { } } +func TestResolveEndpoint_SessionKeyPlaceholderPreserved(t *testing.T) { + clearAllEnv(t) + + cfg := configFile{ + Provider: "my-gateway", + CustomProviders: map[string]providerEntryConfig{ + "my-gateway": { + APIKey: "sk-test", + URL: "https://gateway.example.com/v1", + Protocol: "openai", + Model: "some-model", + ExtraHeaders: map[string]string{"x-session-affinity": "{ocr_session_key}"}, + }, + }, + } + data, _ := json.Marshal(cfg) + cfgPath := filepath.Join(t.TempDir(), "config.json") + if err := os.WriteFile(cfgPath, data, 0644); err != nil { + t.Fatalf("write config: %v", err) + } + + ep, err := ResolveEndpoint(cfgPath) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // The template placeholder must survive resolution untouched; expansion + // happens in the client, per request. + if v := ep.ExtraHeaders["x-session-affinity"]; v != "{ocr_session_key}" { + t.Errorf("ExtraHeaders[\"x-session-affinity\"] = %q, want raw placeholder", v) + } +} + func TestResolveEndpoint_LegacyLlmExtraHeaders(t *testing.T) { clearAllEnv(t) diff --git a/internal/llm/sessionkey.go b/internal/llm/sessionkey.go new file mode 100644 index 00000000..985fb6fe --- /dev/null +++ b/internal/llm/sessionkey.go @@ -0,0 +1,128 @@ +package llm + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "fmt" + "io" + "strings" + "time" +) + +// SessionKeyTemplateVar is the placeholder users can embed in extra_headers +// and extra_body values. Clients replace it with the run's session key, so +// providers that route prompt-cache traffic by an explicit key (e.g. +// x-session-affinity headers on Cloudflare/Fireworks/Mistral-style gateways) +// can be configured without OCR knowing each provider's convention. +const SessionKeyTemplateVar = "{ocr_session_key}" + +// sessionKeyCtxKey is the context key carrying the run's session key. +type sessionKeyCtxKey struct{} + +// ContextWithSessionKey returns a context carrying the given session key. +// Review and scan runs bind their session history's SessionID at the top of +// Run as a base key, and each task refines it with SessionTaskKey where its +// conversation starts, so every LLM request is tagged with the real OCR +// session's key at the granularity prompt caches actually work at. +func ContextWithSessionKey(ctx context.Context, key string) context.Context { + if key == "" { + return ctx + } + return context.WithValue(ctx, sessionKeyCtxKey{}, key) +} + +// SessionTaskKey derives the prompt-cache affinity key for one task +// conversation within a session. Prompt caches match on prefixes, and OCR's +// task types (plan, main tool-loop, compression, dedup, ...) use unrelated +// prompts — routing a whole run under one key would pin every concurrent +// per-file conversation to one cache node with no shared prefix to reuse, +// and hot keys degrade anyway (OpenAI reroutes a prompt_cache_key once it +// exceeds roughly 15 requests/minute). Scoping the key to one conversation +// — (session, task type, file or batch) — keeps each conversation's growing +// prefix on a consistent node instead. +// +// The scope (usually a file path) is hashed so the key stays header-safe; +// the session key and task type stay readable so provider-side logs can be +// correlated with OCR session records. +func SessionTaskKey(sessionKey, taskType, scope string) string { + if taskType == "" && scope == "" { + return sessionKey + } + if scope == "" { + return sessionKey + "-" + taskType + } + sum := sha256.Sum256([]byte(scope)) + return fmt.Sprintf("%s-%s-%x", sessionKey, taskType, sum[:8]) +} + +// SessionKeyFromContext returns the session key carried by ctx, or "" when +// none was set. +func SessionKeyFromContext(ctx context.Context) string { + key, _ := ctx.Value(sessionKeyCtxKey{}).(string) + return key +} + +// NewSessionKey returns a fresh UUIDv4 session key. Clients use it as a +// per-client fallback for requests whose context carries no session key +// (e.g. `ocr llm test`, which has no review session). +func NewSessionKey() string { + b := make([]byte, 16) + if _, err := io.ReadFull(rand.Reader, b); err != nil { + // Fallback — extremely unlikely but keeps things working without panics. + return fmt.Sprintf("fallback-%d", time.Now().UnixNano()) + } + b[6] = (b[6] & 0x0f) | 0x40 // version 4 + b[8] = (b[8] & 0x3f) | 0x80 // variant 1 + return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x", + b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) +} + +// expandSessionKeyInHeaders returns a copy of headers with every occurrence +// of SessionKeyTemplateVar in the values replaced by key. The input map is +// never mutated; nil input returns nil. +func expandSessionKeyInHeaders(headers map[string]string, key string) map[string]string { + if len(headers) == 0 { + return headers + } + out := make(map[string]string, len(headers)) + for k, v := range headers { + out[k] = strings.ReplaceAll(v, SessionKeyTemplateVar, key) + } + return out +} + +// expandSessionKeyInBody returns a copy of body with SessionKeyTemplateVar +// replaced by key in every string value, recursing into nested maps and +// slices. The input is never mutated; nil input returns nil. +func expandSessionKeyInBody(body map[string]any, key string) map[string]any { + if len(body) == 0 { + return body + } + out := make(map[string]any, len(body)) + for k, v := range body { + out[k] = expandSessionKeyValue(v, key) + } + return out +} + +func expandSessionKeyValue(v any, key string) any { + switch val := v.(type) { + case string: + return strings.ReplaceAll(val, SessionKeyTemplateVar, key) + case map[string]any: + out := make(map[string]any, len(val)) + for k, nested := range val { + out[k] = expandSessionKeyValue(nested, key) + } + return out + case []any: + out := make([]any, len(val)) + for i, nested := range val { + out[i] = expandSessionKeyValue(nested, key) + } + return out + default: + return v + } +} diff --git a/internal/llm/sessionkey_test.go b/internal/llm/sessionkey_test.go new file mode 100644 index 00000000..cd1a3a7b --- /dev/null +++ b/internal/llm/sessionkey_test.go @@ -0,0 +1,133 @@ +package llm + +import ( + "context" + "regexp" + "strings" + "testing" +) + +var uuidV4Re = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$`) + +func TestNewSessionKey(t *testing.T) { + k1 := NewSessionKey() + k2 := NewSessionKey() + if !uuidV4Re.MatchString(k1) { + t.Errorf("NewSessionKey() = %q, not a UUIDv4", k1) + } + if k1 == k2 { + t.Errorf("NewSessionKey() returned duplicate key %q", k1) + } +} + +func TestSessionKeyContext(t *testing.T) { + ctx := context.Background() + if got := SessionKeyFromContext(ctx); got != "" { + t.Errorf("SessionKeyFromContext(empty) = %q, want \"\"", got) + } + if got := SessionKeyFromContext(ContextWithSessionKey(ctx, "key-1")); got != "key-1" { + t.Errorf("SessionKeyFromContext = %q, want %q", got, "key-1") + } + // Empty keys are not stored. + if got := SessionKeyFromContext(ContextWithSessionKey(ctx, "")); got != "" { + t.Errorf("SessionKeyFromContext(empty key) = %q, want \"\"", got) + } + // A nested key overrides the outer one. + nested := ContextWithSessionKey(ContextWithSessionKey(ctx, "outer"), "inner") + if got := SessionKeyFromContext(nested); got != "inner" { + t.Errorf("SessionKeyFromContext(nested) = %q, want %q", got, "inner") + } +} + +func TestSessionTaskKey(t *testing.T) { + if got := SessionTaskKey("sess", "", ""); got != "sess" { + t.Errorf("SessionTaskKey(sess,,) = %q, want %q", got, "sess") + } + if got := SessionTaskKey("sess", "plan_task", ""); got != "sess-plan_task" { + t.Errorf("SessionTaskKey without scope = %q, want %q", got, "sess-plan_task") + } + + k1 := SessionTaskKey("sess", "main_task", "a/b.go") + k2 := SessionTaskKey("sess", "main_task", "a/b.go") + if k1 != k2 { + t.Errorf("SessionTaskKey not deterministic: %q != %q", k1, k2) + } + if !strings.HasPrefix(k1, "sess-main_task-") { + t.Errorf("SessionTaskKey = %q, want prefix %q", k1, "sess-main_task-") + } + if k1 == SessionTaskKey("sess", "main_task", "a/c.go") { + t.Error("different scopes must produce different keys") + } + if k1 == SessionTaskKey("sess", "plan_task", "a/b.go") { + t.Error("different task types must produce different keys") + } + // Non-ASCII scopes must still yield a header-safe key. + for _, r := range SessionTaskKey("sess", "main_task", "文件/경로.go") { + if r <= ' ' || r > '~' { + t.Errorf("SessionTaskKey produced non header-safe rune %q", r) + } + } +} + +func TestExpandSessionKeyInHeaders(t *testing.T) { + in := map[string]string{ + "x-session-affinity": "{ocr_session_key}", + "X-Tenant": "tenant-{ocr_session_key}-suffix", + "X-Static": "unchanged", + } + out := expandSessionKeyInHeaders(in, "key-123") + + if got := out["x-session-affinity"]; got != "key-123" { + t.Errorf("x-session-affinity = %q, want %q", got, "key-123") + } + if got := out["X-Tenant"]; got != "tenant-key-123-suffix" { + t.Errorf("X-Tenant = %q, want %q", got, "tenant-key-123-suffix") + } + if got := out["X-Static"]; got != "unchanged" { + t.Errorf("X-Static = %q, want %q", got, "unchanged") + } + // The input map must not be mutated. + if in["x-session-affinity"] != "{ocr_session_key}" { + t.Error("input map was mutated") + } + if expandSessionKeyInHeaders(nil, "key") != nil { + t.Error("nil input should return nil") + } +} + +func TestExpandSessionKeyInBody(t *testing.T) { + in := map[string]any{ + "prompt_cache_key": "{ocr_session_key}", + "count": 42, + "nested": map[string]any{ + "session": "{ocr_session_key}", + }, + "list": []any{"{ocr_session_key}", 1, true}, + } + out := expandSessionKeyInBody(in, "key-456") + + if got := out["prompt_cache_key"]; got != "key-456" { + t.Errorf("prompt_cache_key = %v, want %q", got, "key-456") + } + if got := out["count"]; got != 42 { + t.Errorf("count = %v, want 42", got) + } + nested, ok := out["nested"].(map[string]any) + if !ok || nested["session"] != "key-456" { + t.Errorf("nested.session = %v, want %q", out["nested"], "key-456") + } + list, ok := out["list"].([]any) + if !ok || list[0] != "key-456" || list[1] != 1 || list[2] != true { + t.Errorf("list = %v, want [key-456 1 true]", out["list"]) + } + // The input map must not be mutated. + if in["prompt_cache_key"] != "{ocr_session_key}" { + t.Error("input map was mutated") + } + if in["nested"].(map[string]any)["session"] != "{ocr_session_key}" { + t.Error("nested input map was mutated") + } + if expandSessionKeyInBody(nil, "key") != nil { + t.Error("nil input should return nil") + } +} diff --git a/internal/llmloop/compression.go b/internal/llmloop/compression.go index 5a801a6f..10b165d1 100644 --- a/internal/llmloop/compression.go +++ b/internal/llmloop/compression.go @@ -224,6 +224,9 @@ func (r *Runner) runCompression(ctx context.Context, msgs []llm.Message, filePat compressionMsgs = append(compressionMsgs, llm.NewTextMessage(m.Role, content)) } + ctx = llm.ContextWithSessionKey(ctx, + llm.SessionTaskKey(r.deps.Session.SessionID, string(session.MemoryCompressionTask), filePath)) + startTime := time.Now() resp, err := r.deps.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{ Model: r.deps.Model, diff --git a/internal/llmloop/loop.go b/internal/llmloop/loop.go index b474fef0..d618ba72 100644 --- a/internal/llmloop/loop.go +++ b/internal/llmloop/loop.go @@ -147,6 +147,13 @@ func (r *Runner) CollectPendingComments() []model.LlmComment { // are aggregated on the Runner across all files. The returned bool is true // only when the model explicitly calls task_done. func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath string) (bool, error) { + // Every round of this loop re-sends the growing conversation, so each + // request is a prefix extension of the previous one — exactly what + // provider prompt caches reuse. Scope the affinity key to this file's + // main-task conversation so every round routes to the same cache node. + ctx = llm.ContextWithSessionKey(ctx, + llm.SessionTaskKey(r.deps.Session.SessionID, string(session.MainTask), newPath)) + toolReqCount := r.deps.Template.MaxToolRequestTimes const maxConsecutiveEmptyRounds = 3 consecutiveEmptyRounds := 0 @@ -353,7 +360,9 @@ func (r *Runner) executeToolCall(ctx context.Context, newPath string, call llm.T if d != nil { if !diff.ResolveComment(cm, d) && r.deps.Template.ReLocationTask != nil { rlStart := time.Now() - _, resp, msgs := diff.ReLocateComment(rctx, cm, d, r.deps.LLMClient, r.deps.Template.ReLocationTask, r.deps.Model, r.deps.Template.MaxTokens) + rlCtx := llm.ContextWithSessionKey(rctx, + llm.SessionTaskKey(r.deps.Session.SessionID, string(session.ReLocationTask), cm.Path)) + _, resp, msgs := diff.ReLocateComment(rlCtx, cm, d, r.deps.LLMClient, r.deps.Template.ReLocationTask, r.deps.Model, r.deps.Template.MaxTokens) if msgs != nil { fs := r.deps.Session.GetOrCreateFileSession(cm.Path) rlRec := fs.AppendTaskRecord(session.ReLocationTask, msgs) diff --git a/internal/llmloop/loop_test.go b/internal/llmloop/loop_test.go index b2dd9e68..1010ffe8 100644 --- a/internal/llmloop/loop_test.go +++ b/internal/llmloop/loop_test.go @@ -15,9 +15,12 @@ import ( type fakeClient struct { responses []*llm.ChatResponse calls int + // sessionKeys records llm.SessionKeyFromContext for each call. + sessionKeys []string } -func (f *fakeClient) CompletionsWithCtx(_ context.Context, _ llm.ChatRequest) (*llm.ChatResponse, error) { +func (f *fakeClient) CompletionsWithCtx(ctx context.Context, _ llm.ChatRequest) (*llm.ChatResponse, error) { + f.sessionKeys = append(f.sessionKeys, llm.SessionKeyFromContext(ctx)) if f.calls >= len(f.responses) { content := "" return &llm.ChatResponse{ @@ -118,6 +121,30 @@ func TestRunPerFile_TaskDoneImmediately(t *testing.T) { } } +func TestRunPerFile_TagsRequestsWithTaskSessionKey(t *testing.T) { + client := &fakeClient{responses: []*llm.ChatResponse{ + fileReadToolCallResponse("call_1", `{"path":"main.go"}`), + taskDoneResponse(), + }} + deps := newTestDeps(client) + runner := NewRunner(deps) + + msgs := []llm.Message{llm.NewTextMessage("user", "review this file")} + if _, err := runner.RunPerFile(context.Background(), msgs, "main.go"); err != nil { + t.Fatalf("RunPerFile: %v", err) + } + + want := llm.SessionTaskKey(deps.Session.SessionID, string(session.MainTask), "main.go") + if len(client.sessionKeys) != 2 { + t.Fatalf("expected 2 recorded session keys, got %d", len(client.sessionKeys)) + } + for i, got := range client.sessionKeys { + if got != want { + t.Errorf("call %d session key = %q, want %q", i, got, want) + } + } +} + func TestRunPerFile_ToolCallThenDone(t *testing.T) { client := &fakeClient{responses: []*llm.ChatResponse{ fileReadToolCallResponse("call_1", `{"path":"main.go"}`), diff --git a/internal/scan/agent.go b/internal/scan/agent.go index 1e1670e9..ec3ea5ae 100644 --- a/internal/scan/agent.go +++ b/internal/scan/agent.go @@ -202,6 +202,13 @@ func (a *Agent) Run(ctx context.Context) ([]model.LlmComment, error) { return nil, fmt.Errorf("scan template MAIN_TASK is missing or empty") } + // Base prompt-cache affinity key for any LLM request in this run that a + // task doesn't re-scope. Each task conversation (plan, per-file main + // loop, dedup, summary) refines it with llm.SessionTaskKey where it + // starts, so affinity keys stay per-conversation — the granularity + // provider prompt caches actually reuse prefixes at. + ctx = llm.ContextWithSessionKey(ctx, a.SessionID()) + ctx, scanSpan := telemetry.StartSpan(ctx, "scan.enumerate") provider := NewProvider(a.args.RepoDir, a.args.Paths, a.args.GitRunner, a.args.MaxFileSizeBytes) items, err := provider.Enumerate(ctx) @@ -595,6 +602,8 @@ func (a *Agent) maybeRunPlan(ctx context.Context, it model.ScanItem, rule string fs := a.session.GetOrCreateFileSession(it.Path) rec := fs.AppendTaskRecord(session.PlanTask, messages) + ctx = llm.ContextWithSessionKey(ctx, + llm.SessionTaskKey(a.session.SessionID, string(session.PlanTask), it.Path)) startTime := time.Now() resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{ @@ -648,6 +657,8 @@ func (a *Agent) maybeRunProjectSummary(ctx context.Context, comments []model.Llm const pathKey = "__scan_project_summary__" fs := a.session.GetOrCreateFileSession(pathKey) rec := fs.AppendTaskRecord(session.MemoryCompressionTask, messages) // reuse existing task type + ctx = llm.ContextWithSessionKey(ctx, + llm.SessionTaskKey(a.session.SessionID, string(session.MemoryCompressionTask), pathKey)) startTime := time.Now() resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{ @@ -723,6 +734,8 @@ func (a *Agent) maybeRunDedup(ctx context.Context, batchIdx, batchStart int) { pathKey := fmt.Sprintf("__scan_dedup_batch_%d__", batchIdx) fs := a.session.GetOrCreateFileSession(pathKey) rec := fs.AppendTaskRecord(session.MemoryCompressionTask, messages) // reuse existing task type; no scan-specific type to invent + ctx = llm.ContextWithSessionKey(ctx, + llm.SessionTaskKey(a.session.SessionID, string(session.MemoryCompressionTask), pathKey)) startTime := time.Now() resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{ diff --git a/pages/src/content/docs/en/configuration.md b/pages/src/content/docs/en/configuration.md index 63b5d777..58fc9150 100644 --- a/pages/src/content/docs/en/configuration.md +++ b/pages/src/content/docs/en/configuration.md @@ -153,6 +153,29 @@ without patching the source: ocr config set providers.anthropic.extra_body '{"thinking":{"type":"disabled"}}' ``` +### Session affinity for prompt caching + +OCR derives a prompt-cache affinity key for every LLM conversation, scoped +to the review session and the task within it +(`--`). Prompt caches match on prefixes, +so per-conversation keys keep each growing conversation (such as a file's +review tool-loop) on a consistent cache node instead of pinning the whole +run to one hot key; the session-ID prefix lets provider-side cache logs be +correlated with `ocr session` records. + +To opt in, embed the `{ocr_session_key}` template variable in +`extra_headers` or `extra_body` values wherever your provider expects the +key — OCR substitutes the conversation's key per request and sends nothing +otherwise: + +```bash +# OpenAI: prompt_cache_key request body field +ocr config set providers.openai.extra_body '{"prompt_cache_key": "{ocr_session_key}"}' + +# Header-routed gateways (Cloudflare, Fireworks, Mistral commonly use x-session-affinity) +ocr config set custom_providers.my-gateway.extra_headers "x-session-affinity={ocr_session_key}" +``` + ## Configuring the review language `language` determines which language review comments are written in; diff --git a/pages/src/content/docs/ja/configuration.md b/pages/src/content/docs/ja/configuration.md index da268426..23cf9395 100644 --- a/pages/src/content/docs/ja/configuration.md +++ b/pages/src/content/docs/ja/configuration.md @@ -151,6 +151,27 @@ ocr config set providers..url http://127.0.0.1:15721/v1 ocr config set providers.anthropic.extra_body '{"thinking":{"type":"disabled"}}' ``` +### プロンプトキャッシュのセッションアフィニティ + +OCR はすべての LLM 会話ごとに、レビューセッションとその中のタスクにスコープされた +プロンプトキャッシュ・アフィニティキー(`<セッションID>-<タスク種別>-<スコープハッシュ>`)を +導出します。プロンプトキャッシュはプレフィックス単位でマッチするため、会話ごとのキーは、 +実行全体を 1 つのホットキーに集中させず、成長していく各会話(例:ファイルごとの +レビューツールループ)を一貫したキャッシュノードに保ちます。キーのセッション ID +プレフィックスにより、プロバイダー側のキャッシュログを `ocr session` の記録と照合できます。 + +利用するには、プロバイダーがキーを期待する場所の `extra_headers` または `extra_body` の +値に `{ocr_session_key}` テンプレート変数を埋め込みます。OCR がリクエストごとにその会話の +キーに置換し、設定がなければ何も送信しません。 + +```bash +# OpenAI: prompt_cache_key リクエストボディフィールド +ocr config set providers.openai.extra_body '{"prompt_cache_key": "{ocr_session_key}"}' + +# ヘッダーでルーティングするゲートウェイ(Cloudflare、Fireworks、Mistral は通常 x-session-affinity を使用) +ocr config set custom_providers.my-gateway.extra_headers "x-session-affinity={ocr_session_key}" +``` + ## レビュー言語を設定する `language` はレビューコメントをどの言語で出力するかを決めます。未設定の場合は diff --git a/pages/src/content/docs/zh/configuration.md b/pages/src/content/docs/zh/configuration.md index abbe443c..0aef96f1 100644 --- a/pages/src/content/docs/zh/configuration.md +++ b/pages/src/content/docs/zh/configuration.md @@ -139,6 +139,20 @@ ocr config set providers..url http://127.0.0.1:15721/v1 ocr config set providers.anthropic.extra_body '{"thinking":{"type":"disabled"}}' ``` +### 提示词缓存的会话亲和性 + +OCR 为每个 LLM 对话派生一个提示词缓存亲和性密钥,作用域为评审会话及其中的任务(`<会话 ID>-<任务类型>-<作用域哈希>`)。提示词缓存按前缀匹配,因此按对话划分密钥可让每个不断增长的对话(如某个文件的评审工具循环)稳定路由到同一缓存节点,而不是把整次运行压在一个热点密钥上;密钥中的会话 ID 前缀可将供应商侧缓存日志与 `ocr session` 记录对应。 + +要启用,请在供应商期望的位置——`extra_headers` 或 `extra_body` 的值中——嵌入 `{ocr_session_key}` 模板变量。OCR 会在每个请求中将其替换为该对话的密钥;未配置时不会发送任何内容: + +```bash +# OpenAI:prompt_cache_key 请求体字段 +ocr config set providers.openai.extra_body '{"prompt_cache_key": "{ocr_session_key}"}' + +# 通过请求头路由的网关(Cloudflare、Fireworks、Mistral 通常使用 x-session-affinity) +ocr config set custom_providers.my-gateway.extra_headers "x-session-affinity={ocr_session_key}" +``` + ## 配置评审语言 `language` 决定评审评论用哪种语言输出,未设置时默认英文: