diff --git a/.changeset/session-routing-headers.md b/.changeset/session-routing-headers.md new file mode 100644 index 0000000000..decbd699ad --- /dev/null +++ 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) diff --git a/packages/engine/src/__tests__/pi-create-fn-agent.test.ts b/packages/engine/src/__tests__/pi-create-fn-agent.test.ts index 2d2b169ba6..acb7d55530 100644 --- 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 = {}) { + 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 }; + }; + 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 }; + + 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 }; + + 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 }; + + expect(result.headers).toBeUndefined(); + }); + }); + describe("skill selection", () => { beforeEach(() => { // Reset modules to ensure fresh imports for each test diff --git a/packages/engine/src/__tests__/pi-session-routing-headers.test.ts b/packages/engine/src/__tests__/pi-session-routing-headers.test.ts new file mode 100644 index 0000000000..5fd9edf4e5 --- /dev/null +++ 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; 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).getApiKeyAndHeaders).toBeUndefined(); + }); +}); diff --git a/packages/engine/src/executor.ts b/packages/engine/src/executor.ts index 795d31b7ee..a2eae72289 100644 --- 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 } : {}), }); diff --git a/packages/engine/src/pi.ts b/packages/engine/src/pi.ts index 84c2bdac7e..cad52e0130 100644 --- 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 { + 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 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