From 4b2caf43a8a07d4099079d6ec5b1e46e2f5d1059 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Apr 2026 16:55:16 +0200 Subject: [PATCH] fix(ai): add completions session affinity compat closes #3430 --- packages/ai/README.md | 1 + packages/ai/src/types.ts | 2 + .../openai-completions-prompt-cache.test.ts | 141 +++++++++++++----- ...nai-completions-tool-result-images.test.ts | 1 + 4 files changed, 107 insertions(+), 38 deletions(-) diff --git a/packages/ai/README.md b/packages/ai/README.md index a54f0664..8c1531c5 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -851,6 +851,7 @@ interface OpenAICompletionsCompat { supportsReasoningEffort?: boolean; // Whether provider supports `reasoning_effort` (default: true) supportsUsageInStreaming?: boolean; // Whether provider supports `stream_options: { include_usage: true }` (default: true) supportsStrictMode?: boolean; // Whether provider supports `strict` in tool definitions (default: true) + sendSessionAffinityHeaders?: boolean; // Whether to send `session_id`, `x-client-request-id`, and `x-session-affinity` from `sessionId` when caching is enabled (default: false) maxTokensField?: 'max_completion_tokens' | 'max_tokens'; // Which field name to use (default: max_completion_tokens) requiresToolResultName?: boolean; // Whether tool results require the `name` field (default: false) requiresAssistantAfterToolResult?: boolean; // Whether tool results must be followed by an assistant message (default: false) diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 79281c6c..9319aacb 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -291,6 +291,8 @@ export interface OpenAICompletionsCompat { zaiToolStream?: boolean; /** Whether the provider supports the `strict` field in tool definitions. Default: true. */ supportsStrictMode?: boolean; + /** Whether to send known session-affinity headers (`session_id`, `x-client-request-id`, `x-session-affinity`) from `options.sessionId` when caching is enabled. Default: false. */ + sendSessionAffinityHeaders?: boolean; } /** Compatibility settings for OpenAI Responses APIs. */ diff --git a/packages/ai/test/openai-completions-prompt-cache.test.ts b/packages/ai/test/openai-completions-prompt-cache.test.ts index c86b9ef0..11c9ae94 100644 --- a/packages/ai/test/openai-completions-prompt-cache.test.ts +++ b/packages/ai/test/openai-completions-prompt-cache.test.ts @@ -1,16 +1,30 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { getModel } from "../src/models.js"; import { streamOpenAICompletions } from "../src/providers/openai-completions.js"; +import type { Model } from "../src/types.js"; + +interface FakeOpenAIClientOptions { + apiKey: string; + baseURL: string; + dangerouslyAllowBrowser: boolean; + defaultHeaders?: Record; +} + +interface CapturedCompletionsPayload { + prompt_cache_key?: string; + prompt_cache_retention?: "24h" | "in-memory" | null; +} const mockState = vi.hoisted(() => ({ - lastParams: undefined as unknown, + lastParams: undefined as CapturedCompletionsPayload | undefined, + lastClientOptions: undefined as FakeOpenAIClientOptions | undefined, })); vi.mock("openai", () => { class FakeOpenAI { chat = { completions: { - create: (params: unknown) => { + create: (params: CapturedCompletionsPayload) => { mockState.lastParams = params; const stream = { async *[Symbol.asyncIterator]() { @@ -39,6 +53,10 @@ vi.mock("openai", () => { }, }, }; + + constructor(options: FakeOpenAIClientOptions) { + mockState.lastClientOptions = options; + } } return { default: FakeOpenAI }; @@ -49,6 +67,7 @@ describe("openai-completions prompt caching", () => { beforeEach(() => { mockState.lastParams = undefined; + mockState.lastClientOptions = undefined; delete process.env.PI_CACHE_RETENTION; }); @@ -60,10 +79,23 @@ describe("openai-completions prompt caching", () => { } }); - async function capturePayload(options?: { cacheRetention?: "none" | "short" | "long"; sessionId?: string }) { + function createModel(overrides: Partial> = {}): Model<"openai-completions"> { const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini"); - const model = { ...baseModel, api: "openai-completions" } as const; + return { + ...(baseModel as Omit, "api">), + api: "openai-completions", + ...overrides, + }; + } + async function captureRequest( + options?: { + cacheRetention?: "none" | "short" | "long"; + sessionId?: string; + headers?: Record; + }, + model: Model<"openai-completions"> = createModel(), + ) { await streamOpenAICompletions( model, { @@ -73,61 +105,94 @@ describe("openai-completions prompt caching", () => { { apiKey: "test-key", ...options }, ).result(); - return mockState.lastParams as { prompt_cache_key?: string; prompt_cache_retention?: "24h" | "in-memory" | null }; + return { + payload: mockState.lastParams, + headers: mockState.lastClientOptions?.defaultHeaders ?? {}, + }; } it("sets prompt_cache_key for direct OpenAI requests when caching is enabled", async () => { - const payload = await capturePayload({ sessionId: "session-123" }); + const { payload } = await captureRequest({ sessionId: "session-123" }); - expect(payload.prompt_cache_key).toBe("session-123"); - expect(payload.prompt_cache_retention).toBeUndefined(); + expect(payload?.prompt_cache_key).toBe("session-123"); + expect(payload?.prompt_cache_retention).toBeUndefined(); }); it("sets prompt_cache_retention to 24h for direct OpenAI requests when cacheRetention is long", async () => { - const payload = await capturePayload({ cacheRetention: "long", sessionId: "session-456" }); + const { payload } = await captureRequest({ cacheRetention: "long", sessionId: "session-456" }); - expect(payload.prompt_cache_key).toBe("session-456"); - expect(payload.prompt_cache_retention).toBe("24h"); + expect(payload?.prompt_cache_key).toBe("session-456"); + expect(payload?.prompt_cache_retention).toBe("24h"); }); it("omits prompt cache fields when cacheRetention is none", async () => { - const payload = await capturePayload({ cacheRetention: "none", sessionId: "session-789" }); + const { payload } = await captureRequest({ cacheRetention: "none", sessionId: "session-789" }); - expect(payload.prompt_cache_key).toBeUndefined(); - expect(payload.prompt_cache_retention).toBeUndefined(); + expect(payload?.prompt_cache_key).toBeUndefined(); + expect(payload?.prompt_cache_retention).toBeUndefined(); }); it("omits prompt cache fields for non-OpenAI base URLs", async () => { - const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini"); - const model = { - ...baseModel, - api: "openai-completions", + const model = createModel({ baseUrl: "https://proxy.example.com/v1", - } as const; + }); + const { payload } = await captureRequest({ cacheRetention: "long", sessionId: "session-proxy" }, model); - await streamOpenAICompletions( - model, - { - systemPrompt: "sys", - messages: [{ role: "user", content: "hi", timestamp: Date.now() }], - }, - { apiKey: "test-key", cacheRetention: "long", sessionId: "session-proxy" }, - ).result(); - - const payload = mockState.lastParams as { - prompt_cache_key?: string; - prompt_cache_retention?: "24h" | "in-memory" | null; - }; - - expect(payload.prompt_cache_key).toBeUndefined(); - expect(payload.prompt_cache_retention).toBeUndefined(); + expect(payload?.prompt_cache_key).toBeUndefined(); + expect(payload?.prompt_cache_retention).toBeUndefined(); }); it("uses PI_CACHE_RETENTION for direct OpenAI requests", async () => { process.env.PI_CACHE_RETENTION = "long"; - const payload = await capturePayload({ sessionId: "session-env" }); + const { payload } = await captureRequest({ sessionId: "session-env" }); - expect(payload.prompt_cache_key).toBe("session-env"); - expect(payload.prompt_cache_retention).toBe("24h"); + expect(payload?.prompt_cache_key).toBe("session-env"); + expect(payload?.prompt_cache_retention).toBe("24h"); + }); + + it("sends known session-affinity headers when compat.sendSessionAffinityHeaders is enabled", async () => { + const model = createModel({ + baseUrl: "https://proxy.example.com/v1", + compat: { sendSessionAffinityHeaders: true }, + }); + const { headers } = await captureRequest({ sessionId: "session-affinity" }, model); + + expect(headers.session_id).toBe("session-affinity"); + expect(headers["x-client-request-id"]).toBe("session-affinity"); + expect(headers["x-session-affinity"]).toBe("session-affinity"); + }); + + it("omits session-affinity headers when cacheRetention is none", async () => { + const model = createModel({ + baseUrl: "https://proxy.example.com/v1", + compat: { sendSessionAffinityHeaders: true }, + }); + const { headers } = await captureRequest({ cacheRetention: "none", sessionId: "session-affinity" }, model); + + expect(headers.session_id).toBeUndefined(); + expect(headers["x-client-request-id"]).toBeUndefined(); + expect(headers["x-session-affinity"]).toBeUndefined(); + }); + + it("lets explicit headers override generated session-affinity headers", async () => { + const model = createModel({ + baseUrl: "https://proxy.example.com/v1", + compat: { sendSessionAffinityHeaders: true }, + }); + const { headers } = await captureRequest( + { + sessionId: "session-affinity", + headers: { + session_id: "override-session", + "x-client-request-id": "override-request", + "x-session-affinity": "override-affinity", + }, + }, + model, + ); + + expect(headers.session_id).toBe("override-session"); + expect(headers["x-client-request-id"]).toBe("override-request"); + expect(headers["x-session-affinity"]).toBe("override-affinity"); }); }); diff --git a/packages/ai/test/openai-completions-tool-result-images.test.ts b/packages/ai/test/openai-completions-tool-result-images.test.ts index 6100547a..acab74c0 100644 --- a/packages/ai/test/openai-completions-tool-result-images.test.ts +++ b/packages/ai/test/openai-completions-tool-result-images.test.ts @@ -34,6 +34,7 @@ const compat: Required = { vercelGatewayRouting: {}, zaiToolStream: false, supportsStrictMode: true, + sendSessionAffinityHeaders: false, }; function buildToolResult(toolCallId: string, timestamp: number): ToolResultMessage {