feat(llm): task-scoped session affinity for prompt caching - #332
Conversation
![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.
|
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>
5a6fe95 to
f2e6d72
Compare
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>
This comment was marked as outdated.
This comment was marked as outdated.
f2e6d72 to
870332f
Compare
|
Simplified the configuration interface to use only the Users can opt in prompt-caching with |
870332f to
cc6a342
Compare
|
@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. |
cc6a342 to
50e401b
Compare
50e401b to
0f62851
Compare
|
rebased |
66b03d3 to
422f0d2
Compare
|
Rebased again. @lizhengfeng101 Can I get a review for this? I think this is a pretty important feature for cost management. While providers like DeepSeek that offer automatic caching are popular right now, there are still times when I need to use providers like Fireworks AI that don't. |
css521
left a comment
css521
left a comment
There was a problem hiding this comment.
Thanks for working on this. The task-scoped key derivation and opt-in template approach look reasonable, but I found one blocking gap.
The new expansion is wired into OpenAIClient and AnthropicClient, but not OpenAIResponsesClient. That client still applies ExtraHeaders at construction time and iterates the raw ExtraBody map per request without consulting SessionKeyFromContext. With the supported openai-responses protocol, both:
prompt_cache_key: "{ocr_session_key}"X-Session-Affinity: {ocr_session_key}
are therefore sent literally. I reproduced this with an HTTP test server. This collapses all runs/tasks onto the same literal affinity key (or can make a gateway reject the request), instead of providing the task-scoped affinity promised by this PR.
Could you please mirror the per-request key resolution and header/body expansion in OpenAIResponsesClient, initialize its fallback key, and add regression coverage through NewLLMClient using ProtocolOpenAIResponses? Since this client already maps ChatRequest.SessionID to PromptCacheKey, the precedence between that typed field and extra_body.prompt_cache_key should also be made explicit and tested.
Non-blocking: main was force-pushed after the latest CI run and the branch still contains the old #827 merge commit, so a rebase and fresh CI run would also be helpful.
|
@cometkim nice work！ Please rebase main |
|
I'm on it. Maybe I made some mistakes while rebasing this multiple times 😅 Btw, I just noticed that Alibaba Cloud's Model Studio also supports implicit prompt caching, but it charges different pricing than the explicit one! |
f884059 to
98604c2
Compare
|
@css521 rebased with fixes. I added a test suite for the responses API client to match other two clients. |
|
Could you also sync the ru docs ( |
…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>
98604c2 to
9e6c5a9
Compare
|
@yingjiexu2002 done by Claude. The content looks good when cross-checked with Google Translate. |
|
Missing This PR adds the "Session affinity for prompt caching" section to three of the four maintained doc locales —
The anchor already exists in the Russian file: Since the whole feature is opt-in and discoverable only through this doc — nothing is sent unless the user writes Suggested addition to ### Аффинити сессии для кеширования промптов
OCR выводит ключ аффинити кеша промптов для каждого диалога с LLM,
ограниченный областью сессии ревью и задачи внутри неё
(`<ID сессии>-<тип задачи>-<хеш области>`). Кеши промптов сопоставляются
по префиксам, поэтому ключи на уровне диалога удерживают каждый растущий
диалог (например, цикл инструментов при ревью одного файла) на одном и том
же узле кеша, вместо того чтобы стягивать весь запуск к одному «горячему»
ключу; префикс с ID сессии позволяет сопоставлять журналы кеша на стороне
поставщика с записями `ocr session`.
Чтобы включить это, вставьте шаблонную переменную `{ocr_session_key}`
в значения `extra_headers` или `extra_body` там, где её ожидает ваш
поставщик — OCR подставляет ключ диалога в каждый запрос, а без такой
настройки не отправляет ничего:
```bash
# OpenAI: поле prompt_cache_key в теле запроса
ocr config set providers.openai.extra_body '{"prompt_cache_key": "{ocr_session_key}"}'
# Шлюзы с маршрутизацией по заголовкам
ocr config set custom_providers.my-gateway.extra_headers "x-session-affinity={ocr_session_key}"
```One thing to note about that draft: the second example deliberately omits the vendor list that the |
Cloudflare uses the I agree that not mentioning specific vendor names here would be better. |
|
Edited the docs based on the section written in English as canonical: ### 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
(`<session-id>-<task-type>-<scope-hash>`). 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:
# By OpenAI-style request body field (e.g. prompt_cache_key)
ocr config set providers.openai.extra_body '{"prompt_cache_key": "{ocr_session_key}"}'
# By HTTP header (e.g. x-session-affinity)
ocr config set custom_providers.my-gateway.extra_headers "x-session-affinity={ocr_session_key}"
|
5 participants
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:
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 historyexpands the
{ocr_session_key}template variable per request inextra_headersvalues and (recursively)extra_bodyvalues, so any provider's convention can be expressed with existing config fields — no new config surface: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 keyrequests 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
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 (including the OpenAIprompt_cache_keyrecipe viaextra_body){ocr_session_key}placeholder survives endpoint 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