diff --git a/packages/core/src/session/model-request.ts b/packages/core/src/session/model-request.ts index 9a9292eb6beb..238a5b021c96 100644 --- a/packages/core/src/session/model-request.ts +++ b/packages/core/src/session/model-request.ts @@ -229,7 +229,7 @@ export const layer = Layer.effect( const entries = Object.entries(shaped.options) const generation = Object.fromEntries(entries.filter(([k]) => GENERATION_KEYS.has(k))) as GenerationOptionsFields const providerOptions = Object.fromEntries(entries.filter(([k]) => !GENERATION_KEYS.has(k))) - const root = session.fork?.sessionID ?? session.id + const affinity = session.parentID ?? session.fork?.sessionID ?? session.id const base = LLM.request({ model: model.model, http: { @@ -244,7 +244,7 @@ export const layer = Layer.effect( }, }, // TODO: Persist cache lineage so nested forks reuse the root session's cache key. - promptCacheKey: /^ses_[0-9a-f]{64}$/.test(root) ? root.slice(4) : root, + promptCacheKey: /^ses_[0-9a-f]{64}$/.test(affinity) ? affinity.slice(4) : affinity, system: shaped.system, messages: boundImages(unsupportedParts(shaped.messages, model.capabilities)), tools: Array.from(hooked, ([name, t]) => ({ ...t, name })), diff --git a/packages/core/test/session-compaction.test.ts b/packages/core/test/session-compaction.test.ts index e98bda151d5c..18a020b73726 100644 --- a/packages/core/test/session-compaction.test.ts +++ b/packages/core/test/session-compaction.test.ts @@ -397,7 +397,7 @@ it.effect("manual compaction summarizes short context instead of no-op", () => ]) expect(requests).toHaveLength(1) - expect(requests[0]?.promptCacheKey).toBe(sessionID) + expect(requests[0]?.promptCacheKey).toBe(parentID) expect(requests[0]?.http?.headers).toEqual({ "x-session-affinity": sessionID, "X-Session-Id": sessionID, diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 66f1af14af0d..f41f61e74725 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -4411,7 +4411,7 @@ describe("SessionRunnerLLM", () => { }) }) - scenario("adds the parent session header to child model requests", function* (s) { + scenario("uses parent cache affinity for child model requests", function* (s) { const parentID = Session.ID.make("ses_runner_parent") yield* s.db @@ -4423,6 +4423,7 @@ describe("SessionRunnerLLM", () => { yield* s.runPrompt("Run child request") expect(s.requests[0]?.http?.headers?.["x-parent-session-id"]).toBe(parentID) + expect(s.requests[0]?.promptCacheKey).toBe(parentID) }) scenario("runs different sessions concurrently", function* (s) {