From dc2c8a3c30f5deac96a9e50712696f8caf075244 Mon Sep 17 00:00:00 2001
From: awschmeder <awschmeder@lbl.gov>
Date: Fri, 12 Jun 2026 13:24:33 -0700
Subject: [PATCH] feat(litellm): forward taskId as X-Zoo-Session-ID request
 header

Closes: #590

LiteLLM recognizes X-<vendor>-Session-ID headers for per-conversation
request correlation in logs and spend tracking. This follows the same
convention used by Claude Code (x-claude-code-session-id) and GitHub
Copilot (x-copilot-session-id).

The header is injected only when metadata.taskId is present, passed as
the second argument to client.chat.completions.create() alongside the
existing .withResponse() pattern so no existing tests need to change.
---
 .changeset/litellm-session-id-header.md      |  5 +++
 src/api/providers/__tests__/lite-llm.spec.ts | 41 ++++++++++++++++++++
 src/api/providers/lite-llm.ts                | 12 +++++-
 3 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/litellm-session-id-header.md

diff --git a/.changeset/litellm-session-id-header.md b/.changeset/litellm-session-id-header.md
new file mode 100644
index 000000000..57b4ec7fd
--- /dev/null
+++ b/.changeset/litellm-session-id-header.md
@@ -0,0 +1,5 @@
+---
+"zoo-code": patch
+---
+
+Forward the active task ID to the LiteLLM proxy as an `X-Zoo-Session-ID` request header so individual conversations can be correlated in LiteLLM logs and spend tracking. The header is only sent when a task ID is present, and follows the `x-<vendor>-session-id` convention used by Claude Code (`x-claude-code-session-id`) and GitHub Copilot (`x-copilot-session-id`).
diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts
index df0e8b152..811910286 100644
--- a/src/api/providers/__tests__/lite-llm.spec.ts
+++ b/src/api/providers/__tests__/lite-llm.spec.ts
@@ -1115,4 +1115,45 @@ describe("LiteLLMHandler", () => {
 			expect(id1).not.toBe(id2)
 		})
 	})
+
+	describe("session ID header", () => {
+		const mockStream = {
+			async *[Symbol.asyncIterator]() {
+				yield {
+					choices: [{ delta: { content: "ok" } }],
+					usage: { prompt_tokens: 1, completion_tokens: 1 },
+				}
+			},
+		}
+
+		it("should send the X-Zoo-Session-ID header when a taskId is provided", async () => {
+			mockCreate.mockReturnValue({
+				withResponse: vi.fn().mockResolvedValue({ data: mockStream }),
+			})
+
+			const generator = handler.createMessage("system", [{ role: "user", content: "hi" }], {
+				taskId: "task-123",
+			})
+			for await (const _chunk of generator) {
+				// drain the stream
+			}
+
+			const requestHeaders = mockCreate.mock.calls[0][1]?.headers
+			expect(requestHeaders).toMatchObject({ "X-Zoo-Session-ID": "task-123" })
+		})
+
+		it("should not send the X-Zoo-Session-ID header when no taskId is provided", async () => {
+			mockCreate.mockReturnValue({
+				withResponse: vi.fn().mockResolvedValue({ data: mockStream }),
+			})
+
+			const generator = handler.createMessage("system", [{ role: "user", content: "hi" }])
+			for await (const _chunk of generator) {
+				// drain the stream
+			}
+
+			const requestHeaders = mockCreate.mock.calls[0][1]?.headers
+			expect(requestHeaders).not.toHaveProperty("X-Zoo-Session-ID")
+		})
+	})
 })
diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts
index 0b79433f3..60c41ab6e 100644
--- a/src/api/providers/lite-llm.ts
+++ b/src/api/providers/lite-llm.ts
@@ -223,8 +223,18 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa
 			requestOptions.temperature = this.options.modelTemperature ?? 0
 		}
 
+		// LiteLLM recognizes X-<vendor>-Session-ID for per-conversation request correlation,
+		// matching the convention used by Claude Code (x-claude-code-session-id) and
+		// GitHub Copilot (x-copilot-session-id).
+		const requestHeaders: Record<string, string> = {}
+		if (metadata?.taskId) {
+			requestHeaders["X-Zoo-Session-ID"] = metadata.taskId
+		}
+
 		try {
-			const { data: completion } = await this.client.chat.completions.create(requestOptions).withResponse()
+			const { data: completion } = await this.client.chat.completions
+				.create(requestOptions, { headers: requestHeaders })
+				.withResponse()
 
 			let lastUsage