feat(llm): task-scoped session affinity for prompt caching#332
Conversation
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 openai provider preset sends the key as prompt_cache_key in the
request body automatically; other providers can route it via the new
{ocr_session_key} template variable in extra_headers/extra_body, which
clients expand per request.
Closes alibaba#229
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| 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()) |
There was a problem hiding this comment.
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:
| return fmt.Sprintf("fallback-%d", time.Now().UnixNano()) | |
| return fmt.Sprintf("fallback-%d-%x", time.Now().UnixNano(), b[:4]) |
There was a problem hiding this comment.
This never runs concurrently
|
Overall this direction looks good to me. One thing I’d like to see is an explicit opt-out. Right now the built-in Should we add a config flag to disable this for a provider, e.g. FYI @lizhengfeng101 |
|
@MuoDoo @lizhengfeng101 I'd like to ask for your opinions to finalize the changes.
|
2 participants
Description
Adds provider-side session affinity for prompt caching (#229).
OCR now 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, 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_keyonce it exceeds ~15 req/min). The session-ID prefix keeps provider-side cache logs correlatable withocr sessionrecords.This PR:
adds
llm.SessionTaskKeyand context helpers (ContextWithSessionKey/SessionKeyFromContext); review/scan runs bind the real session's ID as a base key atRun, 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 historyinjects the key automatically as
prompt_cache_keyin the request body for theopenaiprovider preset (newProvider.SessionKeyBodyfield, propagated viaResolvedEndpoint); an explicitextra_bodyentry with the same name winsadds an
{ocr_session_key}template variable, expanded per request inextra_headersand (recursively)extra_bodyvalues, so providers that route prompt-cache traffic by header can be configured without OCR knowing each provider's convention:moves
extra_headers/extra_bodyapplication 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 keyAnthropic protocol gets no body injection (the API manages cache affinity server-side); the template variable still works for anthropic-protocol gateways
Type of Change
How Has This Been Tested?
make testpasses locallyAlso verified with:
go test ./...go vet ./...New tests cover:
llmloop: every round ofRunPerFilereaches the client with the same task-scoped key (TestRunPerFile_TagsRequestsWithTaskSessionKey){ocr_session_key}expansion in headers and nested body values;prompt_cache_keyinjection and userextra_bodyoverride; fallback key auto-generationopenaipreset resolvesSessionKeyBody: "prompt_cache_key"; custom providers don't, and the raw{ocr_session_key}placeholder survives resolution untouchedSessionTaskKey: deterministic, distinct per task type and scope, header-safe for non-ASCII pathsContext propagation was verified end-to-end: the async comment worker pool and background compression use
context.WithoutCancel, which preserves context values.Checklist
go fmt,go vet)Related Issues
Closes #229
🤖 Generated with Claude Code