feat(#1675): add X-Session-Id and X-Session-Affinity routing headers to LLM requests (#1736)

Closes #1675 ## Summary Adds `X-Session-Id` and `X-Session-Affinity` request headers to all outbound LLM chat completion requests. These headers are widely understood by LLM gateways, proxies, and observability tooling: - **Gateway sticky routing** — keep consecutive requests from one conversation on the same backend or cache instance - **Observability trace grouping** — tools like Langfuse and Arize group individually stateless API calls into a single cohesive multi-turn chat trace - **Memory/proxy middleware** — fetch and append conversation history matching the session id Both headers carry the same stable identifier: the **task id** when available (stable across pause/resume), otherwise the **pi session id** for non-task sessions (chat, summarizer, reviewer). ## Implementation The headers are injected by wrapping `modelRegistry.getApiKeyAndHeaders` — the single chokepoint pi-coding-agent uses to resolve per-request auth and headers for both the main stream and compaction. Verified against the pi-coding-agent source (`sdk.js:171`, `agent-session.js:145/162/1490`) that all HTTP-based provider paths route through this method. The wrapper merges routing headers into the resolved output, preserving provider-specific headers and never disturbing API-key resolution. The `modelRegistry` is created fresh per `createFnAgent` call, so the mutation is session-scoped. ## Changes - **`packages/engine/src/pi.ts`** — `buildSessionRoutingHeaders()` and `attachSessionRoutingHeaders()` helpers; wiring call inside `createFnAgent` resolving `sessionRoutingId = options.taskId ?? piSessionId`. Warns (instead of silently no-oping) if the pi API changes and `getApiKeyAndHeaders` is absent. - **`packages/engine/src/executor.ts`** — propagate `taskId` to four secondary task-scoped sessions (retry, verification-fix, workflow-step, child-agent) that previously fell back to a per-instance pi id, fragmenting per-task observability grouping. - **`packages/engine/src/__tests__/pi-session-routing-headers.test.ts`** — unit tests for the helpers (header shape, merge, failed-auth passthrough, absent-method no-op). - **`packages/engine/src/__tests__/pi-create-fn-agent.test.ts`** — end-to-end wiring tests asserting the `taskId ?? piSessionId` precedence inside `createFnAgent`. - **`.changeset/session-routing-headers.md`** — `@runfusion/fusion: minor` (new feature). ## Test plan - [x] Engine typecheck clean - [x] ESLint clean - [x] `pi-session-routing-headers.test.ts` (5 tests) - [x] `pi-create-fn-agent.test.ts` wiring tests (3 new, 79 total) - [x] Engine-core gate suite (610 tests) - [x] Executor + pi test suites (946 tests)  --- <a href="https://stagereview.app/Runfusion/Fusion/pull/1736"> <picture> <source media="(prefers-color-scheme: dark)" srcset="https://stagereview.app/assets/gh-open-in-stage-dark.svg"> <img src="https://stagereview.app/assets/gh-open-in-stage-light.svg" alt="Open in Stage"> </picture> </a>   ## Summary by CodeRabbit * **New Features** * Added session-routing HTTP headers (`X-Session-Id`, `X-Session-Affinity`) to all outbound LLM chat completion requests. * Uses a stable routing identifier derived from `taskId` when available (preserving consistency across pause/resume), otherwise falls back to the current session id. * Ensures follow-up and spawned sessions reuse the same routing id for consistent sticky routing and trace correlation. * **Tests** * Added coverage to verify header construction, merge behavior, and fallback/no-op cases.
2026-06-23 18:03:37 -07:00
parent 4745f3bd79 0ebc9ea252
commit afa388a865
5 changed files with 245 additions and 1 deletions
--- a/.changeset/session-routing-headers.md
+++ b/.changeset/session-routing-headers.md
@@ -0,0 +1,5 @@
+---
+"@runfusion/fusion": minor
+---
+
+Add `X-Session-Id` and `X-Session-Affinity` request headers to all LLM chat completion requests. These let LLM gateways sticky-route consecutive requests from the same conversation to the same backend, and let observability tools (Langfuse, Arize, etc.) group the otherwise-stateless API calls of a session into a single multi-turn trace. Both headers carry the same stable identifier — the task id when available (stable across pause/resume), otherwise the pi session id. (#1675)
--- a/packages/engine/src/tests/pi-create-fn-agent.test.ts
+++ b/packages/engine/src/tests/pi-create-fn-agent.test.ts
@@ -15,6 +15,11 @@ const findMock = vi.fn();
 const getAllMock = vi.fn(() => [] as any[]);
 const registerProviderMock = vi.fn();
 const refreshMock = vi.fn();
+// FNXC:SessionRouting 2026-06-24-11:30:
+// #1675: capture model-registry auth resolution + session id so the wiring
+// test can assert X-Session-Id/X-Session-Affinity precedence end-to-end.
+const getApiKeyAndHeadersMock = vi.fn(async () => ({ ok: true, apiKey: undefined, headers: undefined }));
+const sessionManagerGetSessionIdMock = vi.fn(() => undefined);
 const settingsManagerCreateMock = vi.fn(() => ({ kind: "settings-manager-create" }));
 const settingsManagerInMemoryMock = vi.fn(() => ({ kind: "settings-manager" }));
 const setFallbackResolverMock = vi.fn();
@@ -138,9 +143,12 @@ vi.mock("@earendil-works/pi-coding-agent", () => ({
    refresh() {
      return refreshMock();
    }
+    getApiKeyAndHeaders() {
+      return getApiKeyAndHeadersMock();
+    }
  },
  SessionManager: {
-    inMemory: () => ({ kind: "session-manager" }),
+    inMemory: () => ({ kind: "session-manager", getSessionId: sessionManagerGetSessionIdMock }),
  },
  SettingsManager: {
    create: settingsManagerCreateMock,
@@ -1024,6 +1032,9 @@ describe("createFnAgent", () => {
    realpathSyncNativeMock.mockImplementation((path: PathLike) => String(path));
    readCustomProvidersMock.mockReturnValue([]);
    findMock.mockImplementation((provider: string, modelId: string) => ({ provider, id: modelId }));
+    // #1675: re-establish default auth + session-id mock returns after clearAllMocks.
+    getApiKeyAndHeadersMock.mockResolvedValue({ ok: true, apiKey: undefined, headers: undefined });
+    sessionManagerGetSessionIdMock.mockReturnValue(undefined);
    createBashToolMock.mockClear();
    createAgentSessionMock.mockResolvedValue({
      session: {
@@ -1921,6 +1932,62 @@ describe("createFnAgent", () => {
    warnSpy.mockRestore();
  });

+  // FNXC:SessionRouting 2026-06-24-11:30:
+  // #1675: createFnAgent must resolve sessionRoutingId = taskId ?? piSessionId and
+  // wrap the registry's getApiKeyAndHeaders so outbound requests carry routing
+  // headers. These assert the wiring precedence end-to-end, not just the helper.
+  describe("session routing headers wiring (#1675)", () => {
+    const anyModel = { provider: "anthropic", id: "claude" } as never;
+
+    async function createAndCaptureRegistry(overrides: Record<string, unknown> = {}) {
+      const { createFnAgent } = await import("../pi.js");
+      await createFnAgent({
+        cwd: "/tmp",
+        systemPrompt: "test",
+        tools: "readonly",
+        ...overrides,
+      });
+      const sessionOptions = createAgentSessionMock.mock.calls.at(-1)?.[0] as {
+        modelRegistry: { getApiKeyAndHeaders: (model: unknown) => Promise<unknown> };
+      };
+      return sessionOptions.modelRegistry;
+    }
+
+    it("uses taskId as the routing id when provided", async () => {
+      const registry = await createAndCaptureRegistry({ taskId: "FN-7788" });
+
+      const result = await registry.getApiKeyAndHeaders(anyModel) as { ok: boolean; headers?: Record<string, string> };
+
+      expect(result.ok).toBe(true);
+      expect(result.headers).toEqual({
+        "X-Session-Id": "FN-7788",
+        "X-Session-Affinity": "FN-7788",
+      });
+    });
+
+    it("falls back to the pi session id when taskId is absent", async () => {
+      sessionManagerGetSessionIdMock.mockReturnValue("pi-session-abc");
+      const registry = await createAndCaptureRegistry();
+
+      const result = await registry.getApiKeyAndHeaders(anyModel) as { ok: boolean; headers?: Record<string, string> };
+
+      expect(result.headers).toEqual({
+        "X-Session-Id": "pi-session-abc",
+        "X-Session-Affinity": "pi-session-abc",
+      });
+    });
+
+    it("does not wrap getApiKeyAndHeaders when neither taskId nor a session id is available", async () => {
+      // getApiKeyAndHeadersMock returns { ok: true, headers: undefined }; if the
+      // wrapper were applied, headers would be populated with X-Session-*.
+      const registry = await createAndCaptureRegistry();
+
+      const result = await registry.getApiKeyAndHeaders(anyModel) as { ok: boolean; headers?: Record<string, string> };
+
+      expect(result.headers).toBeUndefined();
+    });
+  });
+
  describe("skill selection", () => {
    beforeEach(() => {
      // Reset modules to ensure fresh imports for each test
--- a/packages/engine/src/tests/pi-session-routing-headers.test.ts
+++ b/packages/engine/src/tests/pi-session-routing-headers.test.ts
@@ -0,0 +1,83 @@
+import { describe, it, expect } from "vitest";
+import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
+import { attachSessionRoutingHeaders, buildSessionRoutingHeaders } from "../pi.js";
+
+// FNXC:SessionRouting 2026-06-23-16:40:
+// Issue #1675: chat completion requests must carry X-Session-Id and
+// X-Session-Affinity so LLM gateways can sticky-route and observability tools
+// can group the stateless API calls of one conversation into a single trace.
+
+describe("buildSessionRoutingHeaders", () => {
+  it("emits X-Session-Id and X-Session-Affinity with the same identifier", () => {
+    expect(buildSessionRoutingHeaders("sess-123")).toEqual({
+      "X-Session-Id": "sess-123",
+      "X-Session-Affinity": "sess-123",
+    });
+  });
+});
+
+describe("attachSessionRoutingHeaders", () => {
+  // Minimal stand-in for the bits of ModelRegistry the wrapper touches.
+  function makeRegistry(
+    resolve: (model: unknown) => Promise<{ ok: boolean; apiKey?: string; headers?: Record<string, string>; error?: string }>,
+  ): ModelRegistry {
+    return { getApiKeyAndHeaders: resolve } as unknown as ModelRegistry;
+  }
+
+  const anyModel = { provider: "anthropic", id: "claude" } as never;
+
+  it("merges the routing headers into resolved request headers", async () => {
+    const registry = makeRegistry(async () => ({ ok: true, apiKey: "sk-live", headers: undefined }));
+    attachSessionRoutingHeaders(registry, "sess-abc");
+
+    const result = await registry.getApiKeyAndHeaders(anyModel);
+
+    expect(result).toEqual({
+      ok: true,
+      apiKey: "sk-live",
+      headers: {
+        "X-Session-Id": "sess-abc",
+        "X-Session-Affinity": "sess-abc",
+      },
+    });
+  });
+
+  it("preserves the resolved apiKey and any provider-specific headers", async () => {
+    const registry = makeRegistry(async () => ({
+      ok: true,
+      apiKey: "sk-custom",
+      headers: { "HTTP-Referer": "https://example.com", "X-Title": "Fusion" },
+    }));
+    attachSessionRoutingHeaders(registry, "sess-xyz");
+
+    const result = await registry.getApiKeyAndHeaders(anyModel);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok auth result");
+    expect(result.apiKey).toBe("sk-custom");
+    expect(result.headers).toEqual({
+      "HTTP-Referer": "https://example.com",
+      "X-Title": "Fusion",
+      "X-Session-Id": "sess-xyz",
+      "X-Session-Affinity": "sess-xyz",
+    });
+  });
+
+  it("does not alter failed auth resolutions", async () => {
+    const registry = makeRegistry(async () => ({ ok: false, error: "No API key found" }));
+    attachSessionRoutingHeaders(registry, "sess-fail");
+
+    const result = await registry.getApiKeyAndHeaders(anyModel);
+
+    expect(result).toEqual({ ok: false, error: "No API key found" });
+  });
+
+  it("no-ops without throwing when getApiKeyAndHeaders is absent", () => {
+    // If a future pi-coding-agent rename removes the method, the wrapper must not
+    // break session creation. It leaves the registry untouched and warns instead.
+    const registry = {} as ModelRegistry;
+
+    expect(() => attachSessionRoutingHeaders(registry, "sess-none")).not.toThrow();
+    expect((registry as unknown as Record<string, unknown>).getApiKeyAndHeaders).toBeUndefined();
+  });
+});
--- a/packages/engine/src/executor.ts
+++ b/packages/engine/src/executor.ts
@@ -9113,6 +9113,11 @@ export class TaskExecutor {
                  // mirroring the primary execute-seam session above.
                  actionGateContext: this.buildActionGateContext(task.id, identityAgent, settings.defaultAgentPermissionPolicy),
                  permanentAgentGating: this.buildPermanentAgentGatingContext(task.id, identityAgent, settings.defaultAgentPermissionPolicy),
+                  // FNXC:SessionRouting 2026-06-24-11:20:
+                  // #1675: propagate task id so retry-session requests carry the same
+                  // X-Session-Id/X-Session-Affinity as the primary session, keeping the
+                  // task's LLM requests grouped under one stable routing/observability id.
+                  taskId: task.id,
                });
                retrySession = createdRetrySession.session;
                if (createdRetrySession.sessionFile) {
@@ -11940,6 +11945,10 @@ Do not refactor, rename broadly, or make opportunistic improvements.
        runAuditor: createRunAuditor(this.store, this.getRunContextFor(task.id)),
        settings,
        taskEnv: extraEnv,
+        // FNXC:SessionRouting 2026-06-24-11:20:
+        // #1675: propagate task id so verification-fix requests carry the same
+        // X-Session-Id/X-Session-Affinity as the primary session.
+        taskId: task.id,
        ...(skillContext?.skillSelectionContext ? { skillSelection: skillContext.skillSelectionContext } : {}),
      });

@@ -13127,6 +13136,10 @@ You have access to the file system to review changes.${verdictBlock}`;
        runAuditor: createRunAuditor(this.store, this.getRunContextFor(task.id)),
        settings,
        taskEnv: stepEnv,
+        // FNXC:SessionRouting 2026-06-24-11:20:
+        // #1675: propagate task id so workflow-step requests carry the same
+        // X-Session-Id/X-Session-Affinity as the primary session.
+        taskId: task.id,
        // Skill selection: assigned-agent / role-fallback skills, plus the step's
        // own named skill (U1) made discoverable via additionalSkillPaths.
        ...(effectiveSkillSelection ? { skillSelection: effectiveSkillSelection } : {}),
@@ -15888,6 +15901,10 @@ Child agent: ${agent.id} (${name})`;
            runAuditor: createRunAuditor(this.store, this.getRunContextFor(taskId)),
            settings,
            taskEnv,
+            // FNXC:SessionRouting 2026-06-24-11:20:
+            // #1675: propagate task id so child-agent requests carry the same
+            // X-Session-Id/X-Session-Affinity as the parent task session.
+            taskId,
            // Skill selection: use assigned agent skills if available, otherwise role fallback
            ...(skillContext.skillSelectionContext ? { skillSelection: skillContext.skillSelectionContext } : {}),
          });
--- a/packages/engine/src/pi.ts
+++ b/packages/engine/src/pi.ts
@@ -1895,6 +1895,64 @@ export function wrapToolsWithActionGate(
  });
 }

+/**
+ * FNXC:SessionRouting 2026-06-23-16:40:
+ * Outbound LLM chat completion requests must carry `X-Session-Id` and
+ * `X-Session-Affinity` headers (GitHub issue #1675). These are widely
+ * understood by LLM gateways, proxies, and observability tooling:
+ *  - Gateways/routers use them for sticky routing, keeping consecutive requests
+ *    from one conversation on the same backend or cache instance.
+ *  - Observability tools (e.g. Langfuse, Arize) use them to group individually
+ *    stateless API calls into a single cohesive multi-turn chat trace.
+ *  - Memory/proxy middleware uses them to fetch and append conversation history.
+ *
+ * Both headers carry the same stable identifier so sticky-routing affinity and
+ * trace grouping refer to the same session. Builds the header pair for a given
+ * session id.
+ */
+export function buildSessionRoutingHeaders(sessionId: string): Record<string, string> {
+  return {
+    "X-Session-Id": sessionId,
+    "X-Session-Affinity": sessionId,
+  };
+}
+
+/**
+ * FNXC:SessionRouting 2026-06-23-16:40:
+ * Merge the session-routing headers into every header set the model registry
+ * resolves for outbound LLM requests (#1675). `getApiKeyAndHeaders` is the
+ * single point pi-coding-agent uses to resolve per-request auth and headers
+ * (for the main stream and compaction alike), so wrapping it applies the
+ * headers to every HTTP-based provider path (built-in, custom, and
+ * HTTP-streaming extension providers). Subprocess-based providers that make
+ * their own outbound HTTP calls inside a child process (e.g. CLI bridges) are
+ * outside this seam and do not inherit the headers.
+ * Operating on the resolved output (rather than re-registering providers)
+ * preserves provider-specific headers and never disturbs API-key resolution.
+ */
+export function attachSessionRoutingHeaders(modelRegistry: ModelRegistry, sessionId: string): void {
+  // FNXC:SessionRouting 2026-06-23-16:46:
+  // Auxiliary feature: never let header injection break session creation. If a
+  // future pi-coding-agent rename removes getApiKeyAndHeaders, warn (rather than
+  // silently no-op) so the degraded routing/observability headers are detectable.
+  if (typeof modelRegistry.getApiKeyAndHeaders !== "function") {
+    piLog.warn("[pi] session-routing headers not attached: ModelRegistry.getApiKeyAndHeaders is not a function (pi API changed?)");
+    return;
+  }
+  const routingHeaders = buildSessionRoutingHeaders(sessionId);
+  const resolveAuth = modelRegistry.getApiKeyAndHeaders.bind(modelRegistry);
+  modelRegistry.getApiKeyAndHeaders = async (model) => {
+    const result = await resolveAuth(model);
+    if (!result.ok) {
+      return result;
+    }
+    return {
+      ...result,
+      headers: { ...result.headers, ...routingHeaders },
+    };
+  };
+}
+
 /**
 * Create a pi agent session configured for fn.
 * Reuses the user's existing pi auth and model configuration.
@@ -2098,6 +2156,20 @@ export async function createFnAgent(options: AgentOptions): Promise<AgentResult>
  const sessionManager = options.sessionManager ?? SessionManager.inMemory();
  normalizeSessionHistoryEntries(sessionManager as unknown as SessionManagerLike);

+  // FNXC:SessionRouting 2026-06-23-16:40:
+  // Tag every outbound LLM chat completion request with stable session-routing
+  // headers (X-Session-Id / X-Session-Affinity) for gateway sticky routing and
+  // observability trace grouping (#1675). Prefer the task id, which is stable
+  // across pause/resume (each resume spins up a fresh SessionManager), and fall
+  // back to the pi session id for non-task sessions (chat, summarizer, reviewer).
+  const piSessionId = typeof sessionManager.getSessionId === "function"
+    ? sessionManager.getSessionId()
+    : undefined;
+  const sessionRoutingId = options.taskId ?? piSessionId;
+  if (sessionRoutingId) {
+    attachSessionRoutingHeaders(modelRegistry, sessionRoutingId);
+  }
+
  const createSessionWithModel = async (modelOverride?: typeof selectedModel) => {
    // pi-coding-agent 0.68+: `tools` is a string[] allowlist of tool names, not
    // Tool instances. We need boundary-wrapped versions of the built-ins, so we