feat(llm): task-scoped session affinity for prompt caching

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)

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