From aa1b587b104b6dd3dd73cf5ff479a1e708dc3ddb Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 20 Apr 2026 15:37:29 +0200 Subject: [PATCH] fix(ai): add direct OpenAI completions prompt caching closes #3426 --- packages/ai/CHANGELOG.md | 1 + .../ai/src/providers/openai-completions.ts | 15 ++ .../openai-completions-prompt-cache.test.ts | 133 ++++++++++++++++++ 3 files changed, 149 insertions(+) create mode 100644 packages/ai/test/openai-completions-prompt-cache.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index ab79a21f..a5347748 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -5,6 +5,7 @@ ### Fixed - Fixed OpenRouter Meta tests by switching `meta-llama/llama-4-maverick` to `meta-llama/llama-4-scout` to avoid type-check failures from model-catalog drift. +- Fixed direct OpenAI Chat Completions requests to map `sessionId` and `cacheRetention` to OpenAI prompt caching fields, sending `prompt_cache_key` when caching is enabled and `prompt_cache_retention: "24h"` for direct `api.openai.com` requests with long retention ([#3426](https://github.com/badlogic/pi-mono/issues/3426)) ## [0.67.68] - 2026-04-17 diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 77e76b18..645da92a 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -12,6 +12,7 @@ import { getEnvApiKey } from "../env-api-keys.js"; import { calculateCost, supportsXhigh } from "../models.js"; import type { AssistantMessage, + CacheRetention, Context, Message, Model, @@ -58,6 +59,16 @@ export interface OpenAICompletionsOptions extends StreamOptions { reasoningEffort?: "minimal" | "low" | "medium" | "high" | "xhigh"; } +function resolveCacheRetention(cacheRetention?: CacheRetention): CacheRetention { + if (cacheRetention) { + return cacheRetention; + } + if (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long") { + return "long"; + } + return "short"; +} + export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenAICompletionsOptions> = ( model: Model<"openai-completions">, context: Context, @@ -377,10 +388,14 @@ function buildParams(model: Model<"openai-completions">, context: Context, optio const messages = convertMessages(model, context, compat); maybeAddOpenRouterAnthropicCacheControl(model, messages); + const cacheRetention = resolveCacheRetention(options?.cacheRetention); const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { model: model.id, messages, stream: true, + prompt_cache_key: + model.baseUrl.includes("api.openai.com") && cacheRetention !== "none" ? options?.sessionId : undefined, + prompt_cache_retention: model.baseUrl.includes("api.openai.com") && cacheRetention === "long" ? "24h" : undefined, }; if (compat.supportsUsageInStreaming !== false) { diff --git a/packages/ai/test/openai-completions-prompt-cache.test.ts b/packages/ai/test/openai-completions-prompt-cache.test.ts new file mode 100644 index 00000000..c86b9ef0 --- /dev/null +++ b/packages/ai/test/openai-completions-prompt-cache.test.ts @@ -0,0 +1,133 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { getModel } from "../src/models.js"; +import { streamOpenAICompletions } from "../src/providers/openai-completions.js"; + +const mockState = vi.hoisted(() => ({ + lastParams: undefined as unknown, +})); + +vi.mock("openai", () => { + class FakeOpenAI { + chat = { + completions: { + create: (params: unknown) => { + mockState.lastParams = params; + const stream = { + async *[Symbol.asyncIterator]() { + yield { + choices: [{ delta: {}, finish_reason: "stop" }], + usage: { + prompt_tokens: 1, + completion_tokens: 1, + prompt_tokens_details: { cached_tokens: 0 }, + completion_tokens_details: { reasoning_tokens: 0 }, + }, + }; + }, + }; + const promise = Promise.resolve(stream) as Promise & { + withResponse: () => Promise<{ + data: typeof stream; + response: { status: number; headers: Headers }; + }>; + }; + promise.withResponse = async () => ({ + data: stream, + response: { status: 200, headers: new Headers() }, + }); + return promise; + }, + }, + }; + } + + return { default: FakeOpenAI }; +}); + +describe("openai-completions prompt caching", () => { + const originalEnv = process.env.PI_CACHE_RETENTION; + + beforeEach(() => { + mockState.lastParams = undefined; + delete process.env.PI_CACHE_RETENTION; + }); + + afterEach(() => { + if (originalEnv === undefined) { + delete process.env.PI_CACHE_RETENTION; + } else { + process.env.PI_CACHE_RETENTION = originalEnv; + } + }); + + async function capturePayload(options?: { cacheRetention?: "none" | "short" | "long"; sessionId?: string }) { + const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini"); + const model = { ...baseModel, api: "openai-completions" } as const; + + await streamOpenAICompletions( + model, + { + systemPrompt: "sys", + messages: [{ role: "user", content: "hi", timestamp: Date.now() }], + }, + { apiKey: "test-key", ...options }, + ).result(); + + return mockState.lastParams as { prompt_cache_key?: string; prompt_cache_retention?: "24h" | "in-memory" | null }; + } + + it("sets prompt_cache_key for direct OpenAI requests when caching is enabled", async () => { + const payload = await capturePayload({ sessionId: "session-123" }); + + 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" }); + + 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" }); + + 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", + baseUrl: "https://proxy.example.com/v1", + } as const; + + 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(); + }); + + it("uses PI_CACHE_RETENTION for direct OpenAI requests", async () => { + process.env.PI_CACHE_RETENTION = "long"; + const payload = await capturePayload({ sessionId: "session-env" }); + + expect(payload.prompt_cache_key).toBe("session-env"); + expect(payload.prompt_cache_retention).toBe("24h"); + }); +});