feat(llm): task-scoped session affinity for prompt caching by cometkim · Pull Request #332 · alibaba/open-code-review

cometkim · 2026-07-09T07:08:51Z

Description

Adds provider-side session affinity for prompt caching (#229) via a {ocr_session_key} template variable. Writing the placeholder is the opt-in — OCR never invents a parameter or header name, and sends nothing session-related unless configured.

OCR derives a prompt-cache affinity key for every LLM conversation, scoped to the review session and the task within it:

<session-id>-<task-type>-<scope-hash>

Prompt caches match on prefixes, and OCR's task types (plan, per-file main tool-loop, compression, dedup, filter, relocation) use unrelated prompts — while each file's main tool-loop re-sends a growing conversation prefix every round, which is where cache hits actually come from. Scoping the key per task conversation keeps each conversation on a consistent cache node instead of pinning a whole run to one hot key (e.g. OpenAI reroutes a prompt_cache_key once it exceeds ~15 req/min). The session-ID prefix keeps provider-side cache logs correlatable with ocr session records.

This PR:

adds llm.SessionTaskKey and context helpers (ContextWithSessionKey / SessionKeyFromContext); review/scan runs bind the real session's ID as a base key at Run, and each task conversation refines it where it starts (llmloop.RunPerFile, plan, review filter, compression, dedup, project summary, relocation) — using the same (session, task type, path) triple those sites already record into session history

expands the {ocr_session_key} template variable per request in extra_headers values and (recursively) extra_body values, so any provider's convention can be expressed with existing config fields — no new config surface:

# 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, ...)
ocr config set custom_providers.my-gateway.extra_headers "x-session-affinity={ocr_session_key}"

moves extra_headers/extra_body application from client construction to per-request SDK options so the template variable can expand to the key each request's context carries; session-less callers (ocr llm test) fall back to a per-client generated key
requests without the placeholder are byte-for-byte unchanged — no behavior change for existing configurations, and nothing is sent to gateways that reject unknown fields

Type of Change

Bug fix (non-breaking change that fixes an issue)
New feature (non-breaking change that adds functionality)
Breaking change (fix or feature that would cause existing functionality to change)
Refactoring (no functional changes)
Documentation update
CI / Build / Tooling

How Has This Been Tested?

make test passes locally
Manual testing (describe below)

Also verified with:

go test ./...
go vet ./...

New tests cover:

no injection without the placeholder: requests carry no session-related fields unless explicitly configured
llmloop: every round of RunPerFile reaches the client with the same task-scoped key (TestRunPerFile_TagsRequestsWithTaskSessionKey)
clients: context key takes precedence over the client fallback for both protocols; {ocr_session_key} expansion in headers and nested body values (including the OpenAI prompt_cache_key recipe via extra_body)
resolver: the raw {ocr_session_key} placeholder survives endpoint resolution untouched
SessionTaskKey: deterministic, distinct per task type and scope, header-safe for non-ASCII paths

Context propagation was verified end-to-end: the async comment worker pool and background compression use context.WithoutCancel, which preserves context values.

Checklist

My code follows the project's coding style (go fmt, go vet)
I have performed a self-review of my code
I have added tests that prove my fix is effective or my feature works
New and existing unit tests pass locally with my changes
I have updated the documentation accordingly (if applicable)
I have signed the CLA

Related Issues

Closes #229

🤖 Generated with Claude Code

github-actions

🔍 OpenCodeReview found 1 issue(s) in this PR.

✅ 1 posted as inline comment(s)
📝 0 posted as summary

github-actions · 2026-07-09T07:12:46Z

+	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())


The fallback key generation uses only time.Now().UnixNano(), which can produce identical values when called concurrently within the same nanosecond (e.g., multiple clients initializing simultaneously). While crypto/rand failure is extremely rare, the fallback should still avoid collisions. Consider adding an atomic counter or mixing in a monotonic source to guarantee uniqueness even in the fallback path.

Suggestion:

Suggested change

return fmt.Sprintf("fallback-%d", time.Now().UnixNano())

return fmt.Sprintf("fallback-%d-%x", time.Now().UnixNano(), b[:4])

This never runs concurrently

MuoDoo · 2026-07-09T11:52:12Z

Overall this direction looks good to me.

One thing I’d like to see is an explicit opt-out. Right now the built-in openai provider always sends prompt_cache_key. That’s fine for the real OpenAI API, but if someone points providers.openai.url at an OpenAI-compatible gateway that rejects unknown body fields, this could break the main LLM path. The workaround is to switch to custom_providers, but that’s not very obvious.

Should we add a config flag to disable this for a provider, e.g. providers.<name>.session_affinity = false or similar?

FYI @lizhengfeng101

cometkim · 2026-07-09T13:37:29Z

@MuoDoo @lizhengfeng101 I'd like to ask for your opinions to finalize the changes.

Since this is a breaking change, would it be a good idea to even introduce it as an opt-in feature?
And would it be better to use a provider-oriented term like prompt_cache?

Per review feedback on alibaba#332: the openai preset unconditionally sending prompt_cache_key could break OpenAI-compatible gateways (pointed at via providers.openai.url) that reject unknown body fields. Gate the whole mechanism behind an explicit opt-in: session_affinity on provider entries and the legacy llm block (also settable via ocr config set and OCR_LLM_SESSION_AFFINITY). When off (the default), no key is injected and {ocr_session_key} placeholders pass through verbatim. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

cometkim · 2026-07-13T07:32:35Z

Simplified the configuration interface to use only the {ocr_session_key} template variable.

Users can opt in prompt-caching with extra_body or extra_headers. No extra config or env is needed for this.

cometkim · 2026-07-15T05:33:46Z

@MuoDoo Done rebasing. Since this is now an explicit opt-in via extra body or headers, I think it is safe even without other flags.

…e variable Derive a prompt-cache affinity key per LLM conversation, scoped to the review session and the task within it (<session-id>-<task-type>-<hash>). Review/scan runs bind the session ID into the request context and each task conversation refines it where it starts, so every request carries the real OCR session's key at per-conversation granularity — the granularity provider prompt caches reuse prefixes at. Embedding the {ocr_session_key} placeholder in extra_headers or extra_body values is the opt-in: clients expand it per request, and requests without it are unchanged. OCR never enforces a parameter or header name, so any provider convention works with existing config fields, e.g.: extra_body: {"prompt_cache_key": "{ocr_session_key}"} (OpenAI) extra_headers: x-session-affinity={ocr_session_key} (gateways) Closes alibaba#229 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

github-actions Bot reviewed Jul 9, 2026

View reviewed changes

lizhengfeng101 requested a review from MuoDoo July 9, 2026 09:10

lizhengfeng101 requested a review from css521 July 10, 2026 02:41

cometkim mentioned this pull request Jul 13, 2026

feat(mcp): support remote MCP servers via Streamable HTTP #360

Open

cometkim force-pushed the x-session-affinity branch from 5a6fe95 to f2e6d72 Compare July 13, 2026 07:07

This comment was marked as outdated.

Sign in to view

cometkim force-pushed the x-session-affinity branch from f2e6d72 to 870332f Compare July 13, 2026 07:28

cometkim force-pushed the x-session-affinity branch from 870332f to cc6a342 Compare July 15, 2026 05:31

cometkim force-pushed the x-session-affinity branch from cc6a342 to 50e401b Compare July 22, 2026 03:28

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

feat(llm): task-scoped session affinity for prompt caching#332
cometkim wants to merge 1 commit into
alibaba:mainfrom
cometkim:x-session-affinity

cometkim commented Jul 9, 2026 •

edited

Loading

Uh oh!

github-actions Bot left a comment

Uh oh!

github-actions Bot Jul 9, 2026

Uh oh!

cometkim Jul 9, 2026

Uh oh!

MuoDoo commented Jul 9, 2026

Uh oh!

cometkim commented Jul 9, 2026

Uh oh!

This comment was marked as outdated.

cometkim commented Jul 13, 2026

Uh oh!

cometkim commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	return fmt.Sprintf("fallback-%d", time.Now().UnixNano())
	return fmt.Sprintf("fallback-%d-%x", time.Now().UnixNano(), b[:4])

Uh oh!

Conversation

cometkim commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of Change

How Has This Been Tested?

Checklist

Related Issues

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

cometkim Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

MuoDoo commented Jul 9, 2026

Uh oh!

cometkim commented Jul 9, 2026

Uh oh!

This comment was marked as outdated.

cometkim commented Jul 13, 2026

Uh oh!

cometkim commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cometkim commented Jul 9, 2026 •

edited

Loading