diff --git a/packages/core/src/session/compaction.ts b/packages/core/src/session/compaction.ts index ea4cf04aaa..f714633295 100644 --- a/packages/core/src/session/compaction.ts +++ b/packages/core/src/session/compaction.ts @@ -202,6 +202,7 @@ export const make = (dependencies: Dependencies) => { .stream( LLM.request({ model: input.model, + http: input.request.http, messages: [Message.user(summaryPrompt)], tools: [], generation: { maxTokens: summaryOutput }, diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 72c761e10d..874086a06b 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -204,6 +204,13 @@ const layer = Layer.effect( const promptCacheKey = /^ses_[0-9a-f]{64}$/.test(session.id) ? session.id.slice(4) : session.id const request = LLM.request({ model, + http: { + headers: { + "x-session-affinity": session.id, + "X-Session-Id": session.id, + ...(session.parentID ? { "x-parent-session-id": session.parentID } : {}), + }, + }, providerOptions: { openai: { promptCacheKey } }, system: [agent.info?.system, system.baseline] .filter((part): part is string => part !== undefined && part.length > 0) diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 5b40258b2f..cc58b43b29 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -1101,6 +1101,16 @@ describe("SessionRunnerLLM", () => { yield* session.resume(sessionID) expect(requests).toHaveLength(2) + expect(requests.map((request) => request.http?.headers)).toEqual([ + { + "x-session-affinity": sessionID, + "X-Session-Id": sessionID, + }, + { + "x-session-affinity": sessionID, + "X-Session-Id": sessionID, + }, + ]) expect(userTexts(requests[0])[0]).toContain("## Objective") expect(userTexts(requests[1])).toHaveLength(1) expect(userTexts(requests[1])[0]).toContain("\n## Objective\n- Preserve the task\n") @@ -2508,6 +2518,43 @@ describe("SessionRunnerLLM", () => { }), ) + it.effect("adds session correlation headers to model requests", () => + Effect.gen(function* () { + yield* setup + const session = yield* SessionV2.Service + yield* session.prompt({ sessionID, prompt: Prompt.make({ text: "Run correlated request" }), resume: false }) + + requests.length = 0 + yield* session.resume(sessionID) + + expect(requests[0]?.http?.headers).toEqual({ + "x-session-affinity": sessionID, + "X-Session-Id": sessionID, + }) + }), + ) + + it.effect("adds the parent session header to child model requests", () => + Effect.gen(function* () { + yield* setup + const session = yield* SessionV2.Service + const parentID = SessionV2.ID.make("ses_runner_parent") + const { db } = yield* Database.Service + yield* db + .update(SessionTable) + .set({ parent_id: parentID }) + .where(eq(SessionTable.id, sessionID)) + .run() + .pipe(Effect.orDie) + yield* session.prompt({ sessionID, prompt: Prompt.make({ text: "Run child request" }), resume: false }) + + requests.length = 0 + yield* session.resume(sessionID) + + expect(requests[0]?.http?.headers?.["x-parent-session-id"]).toBe(parentID) + }), + ) + it.effect("bounds 64-character session prompt cache keys", () => Effect.gen(function* () { yield* setup