mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-08-25 16:02:11 +00:00
feat: send prompt_cache_key on the openai responses api (#10507)
* feat: send prompt_cache_key on the openai responses api * fix: bound prompt_cache_key to the provider limit and scope it to retryable paths * fix: keep a digest suffix when bounding long frontend cache keys * docs: attach the cache-key doc block to the function it describes
This commit is contained in:
@@ -866,6 +866,7 @@ mod tests {
|
||||
user_message: "hello",
|
||||
attachments: None,
|
||||
has_websearch: false,
|
||||
prompt_cache_key: None,
|
||||
};
|
||||
|
||||
AnthropicQueryBuilder::new(AIProvider::Anthropic, AIPlatform::Standard)
|
||||
|
||||
@@ -218,6 +218,12 @@ pub struct ResponsesApiRequest<'a> {
|
||||
pub max_output_tokens: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub text: Option<ResponsesApiTextFormat>,
|
||||
/// From `gpt-5.6` on, the prefix hash alone no longer reliably matches a cached
|
||||
/// prefix: the key is combined with it to route the request. Omitting it costs the
|
||||
/// read discount and, since these models bill cache writes, pays to re-write the
|
||||
/// prefix on every miss.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_cache_key: Option<&'a str>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Debug)]
|
||||
@@ -434,6 +440,7 @@ impl OpenAIQueryBuilder {
|
||||
.map(|effort| ResponsesApiReasoning { effort: effort.to_string() }),
|
||||
max_output_tokens: args.max_tokens,
|
||||
text,
|
||||
prompt_cache_key: args.prompt_cache_key,
|
||||
};
|
||||
|
||||
serde_json::to_string(&request)
|
||||
@@ -490,6 +497,8 @@ impl OpenAIQueryBuilder {
|
||||
reasoning: None, // Image generation models don't take a reasoning effort
|
||||
max_output_tokens: args.max_tokens,
|
||||
text: None, // No structured output for image generation
|
||||
// A one-shot image prompt has no reusable prefix to route to a cache
|
||||
prompt_cache_key: None,
|
||||
};
|
||||
|
||||
serde_json::to_string(&request)
|
||||
@@ -609,6 +618,7 @@ mod tests {
|
||||
use crate::query_builder::QueryBuilder;
|
||||
|
||||
const SYSTEM_PROMPT: &str = "You are a helpful assistant";
|
||||
const PROMPT_CACHE_KEY: &str = "test-workspace:f/agent:step_1";
|
||||
|
||||
fn client() -> AuthedClient {
|
||||
AuthedClient::new(
|
||||
@@ -641,6 +651,7 @@ mod tests {
|
||||
user_message: "hello",
|
||||
attachments: None,
|
||||
has_websearch: false,
|
||||
prompt_cache_key: Some(PROMPT_CACHE_KEY),
|
||||
};
|
||||
|
||||
OpenAIQueryBuilder::new(AIProvider::OpenAI)
|
||||
@@ -712,4 +723,16 @@ mod tests {
|
||||
|
||||
assert!(request.get("instructions").is_none());
|
||||
}
|
||||
|
||||
/// `gpt-5.6` and later only match a cached prefix reliably when the request carries
|
||||
/// the routing key, so dropping it from the body silently forfeits the cache.
|
||||
#[tokio::test]
|
||||
async fn sends_the_prompt_cache_key() {
|
||||
let messages = vec![message("user", "hi")];
|
||||
|
||||
let body = build_text_body(&messages, None).await;
|
||||
let request: serde_json::Value = serde_json::from_str(&body).unwrap();
|
||||
|
||||
assert_eq!(request["prompt_cache_key"], PROMPT_CACHE_KEY);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,6 +22,11 @@ pub struct BuildRequestArgs<'a> {
|
||||
pub user_message: &'a str,
|
||||
pub attachments: Option<&'a [S3Object]>,
|
||||
pub has_websearch: bool,
|
||||
/// Routing hint for the provider's prompt cache. Requests sharing a prompt prefix
|
||||
/// must reuse one key to land on the same cache, so it is derived from what pins
|
||||
/// the prefix (the step), never from the request. `None` retries a key the
|
||||
/// endpoint rejected.
|
||||
pub prompt_cache_key: Option<&'a str>,
|
||||
}
|
||||
|
||||
/// Response from AI provider
|
||||
|
||||
@@ -11,6 +11,7 @@ use crate::worker_flow::{get_previous_job_result, get_transform_context};
|
||||
use async_recursion::async_recursion;
|
||||
use regex::Regex;
|
||||
use serde_json::value::RawValue;
|
||||
use sha2::Digest;
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
use uuid::Uuid;
|
||||
#[cfg(feature = "bedrock")]
|
||||
@@ -799,6 +800,24 @@ pub async fn handle_ai_agent_job(
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAI rejects a `prompt_cache_key` over 64 characters
|
||||
/// (`Invalid 'prompt_cache_key': string too long`), and a runnable path alone can pass
|
||||
/// that. Fold an over-long key into a digest of itself: same step still yields the same
|
||||
/// key across runs, which is the whole property that routes them to one cache.
|
||||
fn bounded_prompt_cache_key(raw: &str) -> String {
|
||||
const MAX_LEN: usize = 64;
|
||||
if raw.len() <= MAX_LEN {
|
||||
return raw.to_string();
|
||||
}
|
||||
let suffix = hex::encode(&sha2::Sha256::digest(raw.as_bytes())[..16]);
|
||||
// Keep a readable head so a key stays traceable to its workspace in provider logs.
|
||||
let mut head = MAX_LEN - suffix.len() - 1;
|
||||
while head > 0 && !raw.is_char_boundary(head) {
|
||||
head -= 1;
|
||||
}
|
||||
format!("{}:{}", &raw[..head], suffix)
|
||||
}
|
||||
|
||||
#[async_recursion]
|
||||
pub async fn run_agent(
|
||||
// connection
|
||||
@@ -844,9 +863,10 @@ pub async fn run_agent(
|
||||
{
|
||||
query_builder = create_chat_completions_query_builder(&credentials);
|
||||
}
|
||||
// Both outlive the iteration that discovers them: a request shape or a route the
|
||||
// These outlive the iteration that discovers them: a request shape or a route the
|
||||
// endpoint rejected once stays rejected for the whole step.
|
||||
let mut include_usage = true;
|
||||
let mut include_prompt_cache_key = true;
|
||||
|
||||
// Initialize messages
|
||||
let mut messages =
|
||||
@@ -864,6 +884,17 @@ pub async fn run_agent(
|
||||
let effective_flow_step_id: Option<&str> =
|
||||
flow_step_id_override.or(job.flow_step_id.as_deref());
|
||||
|
||||
// Keyed on the step, not the run: every run of this step opens with the same system
|
||||
// prompt and tool definitions, and each agent-loop iteration extends the previous
|
||||
// one's prefix. Above ~15 requests/minute one key starts missing again, which is a
|
||||
// reason to split it further, never to make it per-run.
|
||||
let prompt_cache_key = bounded_prompt_cache_key(&format!(
|
||||
"{}:{}:{}",
|
||||
job.workspace_id,
|
||||
job.runnable_path(),
|
||||
effective_flow_step_id.unwrap_or_default()
|
||||
));
|
||||
|
||||
// Fetch flow context for input transforms context, chat and memory
|
||||
let mut flow_context = get_flow_context(db, job).await;
|
||||
|
||||
@@ -1146,7 +1177,7 @@ pub async fn run_agent(
|
||||
}
|
||||
} else {
|
||||
// For all other providers, use the HTTP client approach
|
||||
let build_args = BuildRequestArgs {
|
||||
let mut build_args = BuildRequestArgs {
|
||||
messages: &messages,
|
||||
tools: tool_defs.as_deref(),
|
||||
model: args.provider.get_model(),
|
||||
@@ -1159,6 +1190,7 @@ pub async fn run_agent(
|
||||
user_message: args.user_message.as_deref().unwrap_or(""),
|
||||
attachments: args.user_attachments.as_deref(),
|
||||
has_websearch,
|
||||
prompt_cache_key: include_prompt_cache_key.then_some(prompt_cache_key.as_str()),
|
||||
};
|
||||
|
||||
// A worker cannot run the client credentials exchange, so an OAuth resource
|
||||
@@ -1208,9 +1240,10 @@ pub async fn run_agent(
|
||||
};
|
||||
|
||||
// An endpoint can reject the request shape rather than the model:
|
||||
// `stream_options`, which not every OpenAI-compatible provider accepts, and
|
||||
// the route itself, when an Azure resource is outside the Responses API's
|
||||
// model/region matrix. Each is retried once with that part dropped.
|
||||
// `stream_options` and `prompt_cache_key`, which not every OpenAI-compatible
|
||||
// gateway accepts, and the route itself, when an Azure resource is outside
|
||||
// the Responses API's model/region matrix. Each is retried once with that
|
||||
// part dropped.
|
||||
// Set where the route is found to be absent, and read once the fallback has
|
||||
// answered: a rejection it did not resolve says nothing about the deployment.
|
||||
let mut rerouted_by_a_route_rejection = false;
|
||||
@@ -1258,6 +1291,13 @@ pub async fn run_agent(
|
||||
|| text.contains("include_usage")
|
||||
|| text.contains("Additional properties are not allowed"));
|
||||
|
||||
// An OpenAI-compatible gateway that validates the body strictly
|
||||
// names the offending field, whether it calls it an unrecognized
|
||||
// argument or an unexpected additional property.
|
||||
let rejects_prompt_cache_key = build_args.prompt_cache_key.is_some()
|
||||
&& status.as_u16() == 400
|
||||
&& text.contains("prompt_cache_key");
|
||||
|
||||
// Only the first call of the step may re-route: an endpoint that
|
||||
// does not serve this API rejects that one already, whereas a
|
||||
// rejection once the conversation is under way is about the
|
||||
@@ -1272,6 +1312,15 @@ pub async fn run_agent(
|
||||
"Retrying request without stream_options due to provider incompatibility"
|
||||
);
|
||||
include_usage = false;
|
||||
} else if rejects_prompt_cache_key {
|
||||
// Checked before the route fallback: the endpoint serves this
|
||||
// route, it just refuses one optional field, and re-routing
|
||||
// the whole step over that would give up far more.
|
||||
tracing::info!(
|
||||
"Retrying request without prompt_cache_key due to provider incompatibility"
|
||||
);
|
||||
include_prompt_cache_key = false;
|
||||
build_args.prompt_cache_key = None;
|
||||
} else if route_unserved {
|
||||
tracing::info!(
|
||||
"Endpoint rejected the request ({}), falling back to chat/completions",
|
||||
@@ -1664,6 +1713,41 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// Over 64 characters OpenAI rejects the key outright, which costs a wasted round
|
||||
/// trip per run and silently leaves that step with no prompt caching at all.
|
||||
#[test]
|
||||
fn prompt_cache_key_stays_within_the_provider_bound() {
|
||||
let long = format!("my-workspace:f/{}/agent:step_12", "nested_folder".repeat(8));
|
||||
assert!(long.len() > 64);
|
||||
|
||||
let bounded = bounded_prompt_cache_key(&long);
|
||||
|
||||
assert!(
|
||||
bounded.len() <= 64,
|
||||
"got {} chars: {bounded}",
|
||||
bounded.len()
|
||||
);
|
||||
// Stable for the same step, or every run would land on a different cache.
|
||||
assert_eq!(bounded, bounded_prompt_cache_key(&long));
|
||||
assert_ne!(
|
||||
bounded,
|
||||
bounded_prompt_cache_key(&long.replace("step_12", "step_13"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_cache_key_passes_short_keys_through_unchanged() {
|
||||
let short = "admins:f/agent/step:a";
|
||||
assert_eq!(bounded_prompt_cache_key(short), short);
|
||||
}
|
||||
|
||||
/// Truncation on a byte index would panic mid-character.
|
||||
#[test]
|
||||
fn prompt_cache_key_truncates_on_a_char_boundary() {
|
||||
let long = format!("workspace:f/{}/agent:step", "é".repeat(80));
|
||||
assert!(bounded_prompt_cache_key(&long).len() <= 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overlay_tool_inputs_binds_matching_flowmodule_tool_only() {
|
||||
fn js(expr: &str) -> InputTransform {
|
||||
|
||||
@@ -10,6 +10,7 @@ const mocks = vi.hoisted(() => ({
|
||||
providerSupportsWebSearch: vi.fn(),
|
||||
getOpenAIResponsesCompletion: vi.fn(),
|
||||
parseOpenAIResponsesCompletion: vi.fn(),
|
||||
buildPromptCacheKey: vi.fn(),
|
||||
getAnthropicCompletion: vi.fn(),
|
||||
parseAnthropicCompletion: vi.fn(),
|
||||
resolveRequestReasoning: vi.fn(),
|
||||
@@ -29,7 +30,8 @@ vi.mock('../reasoningRegistry', () => ({
|
||||
|
||||
vi.mock('./openai-responses', () => ({
|
||||
getOpenAIResponsesCompletion: mocks.getOpenAIResponsesCompletion,
|
||||
parseOpenAIResponsesCompletion: mocks.parseOpenAIResponsesCompletion
|
||||
parseOpenAIResponsesCompletion: mocks.parseOpenAIResponsesCompletion,
|
||||
buildPromptCacheKey: mocks.buildPromptCacheKey
|
||||
}))
|
||||
|
||||
vi.mock('./anthropic', () => ({
|
||||
@@ -423,6 +425,76 @@ describe('runChatLoop reasoning summary fallback', () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe('runChatLoop prompt cache key fallback', () => {
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks()
|
||||
mocks.providerSupportsWebSearch.mockReturnValue(false)
|
||||
mocks.resolveRequestReasoning.mockReturnValue(undefined)
|
||||
mocks.resolveEffectiveReasoning.mockReturnValue(undefined)
|
||||
mocks.parseOpenAIResponsesCompletion.mockResolvedValue({
|
||||
shouldContinue: false,
|
||||
tokenUsage
|
||||
})
|
||||
mocks.parseOpenAICompletion.mockResolvedValue({
|
||||
shouldContinue: false,
|
||||
tokenUsage
|
||||
})
|
||||
})
|
||||
|
||||
it('retries once without the key when the endpoint rejects it, and caches that', async () => {
|
||||
const workspace = `workspace-${randomUUID()}`
|
||||
const modelProvider: ReasoningProviderModel = { provider: 'openai', model: 'gpt-5.6' }
|
||||
const promptCacheKey = `${workspace}:openai:gpt-5.6:chat`
|
||||
mocks.buildPromptCacheKey.mockReturnValue(promptCacheKey)
|
||||
|
||||
mocks.getOpenAIResponsesCompletion
|
||||
.mockRejectedValueOnce(
|
||||
Object.assign(new Error('Unrecognized request argument supplied: prompt_cache_key'), {
|
||||
status: 400,
|
||||
param: 'prompt_cache_key'
|
||||
})
|
||||
)
|
||||
.mockResolvedValue({})
|
||||
|
||||
await runChatLoop(createConfig({ workspace, modelProvider }))
|
||||
|
||||
expect(mocks.getOpenAIResponsesCompletion).toHaveBeenCalledTimes(2)
|
||||
expect(mocks.getOpenAIResponsesCompletion.mock.calls[0][3]).toEqual(
|
||||
expect.objectContaining({ promptCacheKey })
|
||||
)
|
||||
expect(mocks.getOpenAIResponsesCompletion.mock.calls[1][3]).toEqual(
|
||||
expect.objectContaining({ promptCacheKey: undefined })
|
||||
)
|
||||
// The rejection belongs to the endpoint, so a later run skips the key outright
|
||||
// rather than paying the failed round-trip again.
|
||||
expect(mocks.getCompletion).not.toHaveBeenCalled()
|
||||
|
||||
await runChatLoop(createConfig({ workspace, modelProvider }))
|
||||
|
||||
expect(mocks.getOpenAIResponsesCompletion).toHaveBeenCalledTimes(3)
|
||||
expect(mocks.getOpenAIResponsesCompletion.mock.calls[2][3]).toEqual(
|
||||
expect.objectContaining({ promptCacheKey: undefined })
|
||||
)
|
||||
})
|
||||
|
||||
it('does not treat an unrelated 400 as a prompt cache key rejection', async () => {
|
||||
const workspace = `workspace-${randomUUID()}`
|
||||
const modelProvider: ReasoningProviderModel = { provider: 'openai', model: 'gpt-5.6' }
|
||||
mocks.buildPromptCacheKey.mockReturnValue(`${workspace}:openai:gpt-5.6:chat`)
|
||||
|
||||
mocks.getOpenAIResponsesCompletion.mockRejectedValue(
|
||||
Object.assign(new Error('context_length_exceeded'), { status: 400 })
|
||||
)
|
||||
mocks.getCompletion.mockResolvedValue({})
|
||||
|
||||
await runChatLoop(createConfig({ workspace, modelProvider }))
|
||||
|
||||
// Falls through to the Completions API instead of burning a retry on the key.
|
||||
expect(mocks.getOpenAIResponsesCompletion).toHaveBeenCalledTimes(1)
|
||||
expect(mocks.getCompletion).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
})
|
||||
|
||||
describe('runChatLoop lastIterationUsage', () => {
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks()
|
||||
|
||||
@@ -14,7 +14,11 @@ import {
|
||||
import { getAnthropicCompletion, parseAnthropicCompletion } from './anthropic'
|
||||
import { modelSupportsVision, usesAnthropicMessagesApi } from '../modelConfig'
|
||||
import { boundImagePartBytes, stripImagePartsFromMessages } from './imageUtils'
|
||||
import { getOpenAIResponsesCompletion, parseOpenAIResponsesCompletion } from './openai-responses'
|
||||
import {
|
||||
buildPromptCacheKey,
|
||||
getOpenAIResponsesCompletion,
|
||||
parseOpenAIResponsesCompletion
|
||||
} from './openai-responses'
|
||||
import type { Tool, ToolCallbacks } from './shared'
|
||||
import { sanitizeToolCallArguments } from './toolCallArguments'
|
||||
import { addChatTokenUsage, emptyChatTokenUsage, type ChatTokenUsage } from './tokenUsage'
|
||||
@@ -130,6 +134,11 @@ const WEB_SEARCH_UNAVAILABLE_STATUS_CODES = new Set([400, 403, 404])
|
||||
const unsupportedReasoningSummaryCache = new Set<string>()
|
||||
const REASONING_SUMMARY_UNAVAILABLE_STATUS_CODES = new Set([400, 403])
|
||||
|
||||
// A gateway that validates the request body strictly rejects `prompt_cache_key`
|
||||
// outright, so it is a property of the endpoint the credentials point at, not of the
|
||||
// model. Same in-memory reasoning as above: a reload re-probes.
|
||||
const unsupportedPromptCacheKeyCache = new Set<string>()
|
||||
|
||||
function getWebSearchCacheKey(workspace: string, modelProvider: ReasoningProviderModel): string {
|
||||
return [workspace, modelProvider.provider, modelProvider.model].join(':')
|
||||
}
|
||||
@@ -141,6 +150,15 @@ function getReasoningSummaryCacheKey(
|
||||
return [workspace, modelProvider.provider].join(':')
|
||||
}
|
||||
|
||||
// Keyed without the model, like the reasoning-summary probe: a body-validation refusal
|
||||
// belongs to the endpoint, so switching models must not re-pay the failed round trip.
|
||||
function getPromptCacheKeySupportKey(
|
||||
workspace: string,
|
||||
modelProvider: ReasoningProviderModel
|
||||
): string {
|
||||
return [workspace, modelProvider.provider].join(':')
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === 'object' && value !== null
|
||||
}
|
||||
@@ -251,6 +269,25 @@ function shouldRetryWithoutReasoningSummary(err: unknown): boolean {
|
||||
)
|
||||
}
|
||||
|
||||
// An OpenAI-compatible gateway that validates the body strictly names the offending
|
||||
// field, whether it calls it an unrecognized argument or an unexpected additional
|
||||
// property.
|
||||
function shouldRetryWithoutPromptCacheKey(err: unknown): boolean {
|
||||
const status = getErrorStatus(err)
|
||||
if (status !== undefined && status !== 400) {
|
||||
return false
|
||||
}
|
||||
if (getErrorParam(err) === 'prompt_cache_key') {
|
||||
return true
|
||||
}
|
||||
return getErrorText(err).includes('prompt_cache_key')
|
||||
}
|
||||
|
||||
function markPromptCacheKeyUnsupported(cacheKey: string, err: unknown) {
|
||||
unsupportedPromptCacheKeyCache.add(cacheKey)
|
||||
console.warn('prompt_cache_key rejected; retrying without prompt caching hints:', err)
|
||||
}
|
||||
|
||||
function markReasoningSummaryUnsupported(
|
||||
cacheKey: string,
|
||||
err: unknown,
|
||||
@@ -360,6 +397,13 @@ export async function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResul
|
||||
resolveEffectiveReasoning(modelProvider) !== undefined &&
|
||||
!unsupportedReasoningSummaryCache.has(reasoningSummaryCacheKey)
|
||||
|
||||
// One key for the whole chat surface: every iteration opens with the same
|
||||
// system prompt and tool definitions, and each one extends the previous
|
||||
// iteration's prefix, so they all belong on the same cache.
|
||||
const promptCacheKey = buildPromptCacheKey('chat', modelProvider, workspace)
|
||||
const promptCacheSupportKey = getPromptCacheKeySupportKey(workspace, modelProvider)
|
||||
let usePromptCacheKey = !unsupportedPromptCacheKeyCache.has(promptCacheSupportKey)
|
||||
|
||||
const runOpenAIResponses = async (useWebSearch: boolean): Promise<boolean> => {
|
||||
const completion = await getOpenAIResponsesCompletion(
|
||||
messageParams,
|
||||
@@ -370,7 +414,8 @@ export async function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResul
|
||||
openaiClient: clients.openai,
|
||||
webSearch: useWebSearch,
|
||||
reasoningEffort,
|
||||
reasoningSummary
|
||||
reasoningSummary,
|
||||
promptCacheKey: usePromptCacheKey ? promptCacheKey : undefined
|
||||
}
|
||||
)
|
||||
const continueCompletion = await parseOpenAIResponsesCompletion(
|
||||
@@ -389,9 +434,10 @@ export async function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResul
|
||||
let useCompletionsApi = skipResponsesApi
|
||||
if (!skipResponsesApi) {
|
||||
// Retry the Responses call disabling whichever optional feature the
|
||||
// provider rejected (reasoning summary, web search) in the order the
|
||||
// errors arrive — a turn can hit both, either one first. Each retry
|
||||
// permanently disables one feature, so this loops at most twice.
|
||||
// provider rejected (reasoning summary, web search, prompt cache key) in
|
||||
// the order the errors arrive — a turn can hit several, any one first.
|
||||
// Each retry permanently disables one feature, so this loops at most
|
||||
// once per feature.
|
||||
let useWebSearch = webSearch
|
||||
let outcome: 'break' | 'continue' | undefined
|
||||
let fallbackError: unknown
|
||||
@@ -409,6 +455,9 @@ export async function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResul
|
||||
} else if (useWebSearch && shouldRetryWithoutWebSearch(err)) {
|
||||
markWebSearchUnsupported(webSearchCacheKey, err, config.onWebSearchUnavailable)
|
||||
useWebSearch = false
|
||||
} else if (usePromptCacheKey && shouldRetryWithoutPromptCacheKey(err)) {
|
||||
markPromptCacheKeyUnsupported(promptCacheSupportKey, err)
|
||||
usePromptCacheKey = false
|
||||
} else {
|
||||
fallbackError = err
|
||||
break
|
||||
|
||||
@@ -1,18 +1,28 @@
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
import { openAIWebSearchDetails, toResponsesContent } from './openai-responses'
|
||||
import {
|
||||
buildPromptCacheKey,
|
||||
getOpenAIResponsesCompletion,
|
||||
openAIWebSearchDetails,
|
||||
toResponsesContent
|
||||
} from './openai-responses'
|
||||
|
||||
const mocks = vi.hoisted(() => ({
|
||||
getProviderAndCompletionConfig: vi.fn(),
|
||||
applyReasoningToConfig: vi.fn()
|
||||
}))
|
||||
|
||||
// openai-responses.ts pulls in the chat client/registry layer at import time; the
|
||||
// helper under test is pure, so stub those side-effecting modules away.
|
||||
vi.mock('../lib', () => ({
|
||||
createOpenAIProxyClient: vi.fn(),
|
||||
getAiProxyBaseURL: vi.fn(),
|
||||
getProviderAndCompletionConfig: vi.fn(),
|
||||
getProviderAndCompletionConfig: mocks.getProviderAndCompletionConfig,
|
||||
providerSupportsWebSearch: vi.fn(),
|
||||
workspaceAIClients: {}
|
||||
}))
|
||||
|
||||
vi.mock('../reasoningRegistry', () => ({
|
||||
applyReasoningToConfig: vi.fn()
|
||||
applyReasoningToConfig: mocks.applyReasoningToConfig
|
||||
}))
|
||||
|
||||
vi.mock('./shared', () => ({
|
||||
@@ -38,6 +48,92 @@ describe('toResponsesContent', () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe('buildPromptCacheKey', () => {
|
||||
it('composes workspace, provider, model and surface', () => {
|
||||
expect(buildPromptCacheKey('chat', { provider: 'openai', model: 'gpt-5.6' }, 'admins')).toBe(
|
||||
'admins:openai:gpt-5.6:chat'
|
||||
)
|
||||
})
|
||||
|
||||
// Over 64 characters OpenAI rejects the key outright, and an Azure deployment name
|
||||
// is user-chosen, so the model segment can be arbitrarily long.
|
||||
it('stays within the provider length bound for a long deployment name', () => {
|
||||
const key = buildPromptCacheKey(
|
||||
'chat',
|
||||
{ provider: 'azure_openai', model: 'our-very-long-production-deployment-name-for-gpt-5-6' },
|
||||
'some-customer-workspace'
|
||||
)
|
||||
|
||||
expect(key.length).toBeLessThanOrEqual(64)
|
||||
expect(key.startsWith('some-customer-workspace:azure_openai:')).toBe(true)
|
||||
})
|
||||
|
||||
// A max-length workspace on the longest provider name leaves no room for the model
|
||||
// or surface, so only the digest can keep those keys apart.
|
||||
it('keeps distinct models and surfaces apart when the workspace fills the bound', () => {
|
||||
const workspace = 'w'.repeat(50)
|
||||
const keys = [
|
||||
buildPromptCacheKey('chat', { provider: 'azure_openai', model: 'gpt-5.6' }, workspace),
|
||||
buildPromptCacheKey('chat', { provider: 'azure_openai', model: 'gpt-5.1' }, workspace),
|
||||
buildPromptCacheKey('script', { provider: 'azure_openai', model: 'gpt-5.6' }, workspace)
|
||||
]
|
||||
|
||||
for (const key of keys) {
|
||||
expect(key.length).toBeLessThanOrEqual(64)
|
||||
}
|
||||
expect(new Set(keys).size).toBe(3)
|
||||
// Stable, or every turn would land on a different cache.
|
||||
expect(keys[0]).toBe(
|
||||
buildPromptCacheKey('chat', { provider: 'azure_openai', model: 'gpt-5.6' }, workspace)
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe('getOpenAIResponsesCompletion prompt cache key', () => {
|
||||
function stubClient() {
|
||||
const stream = vi.fn().mockReturnValue({})
|
||||
return { client: { responses: { stream } } as any, stream }
|
||||
}
|
||||
|
||||
function stubConfig() {
|
||||
mocks.getProviderAndCompletionConfig.mockReturnValue({
|
||||
provider: 'openai',
|
||||
config: { model: 'gpt-5.6' }
|
||||
})
|
||||
// The reasoning layer rebuilds the config object; keep it a pass-through so the
|
||||
// assertion is about the body this module produces.
|
||||
mocks.applyReasoningToConfig.mockImplementation((config) => config)
|
||||
}
|
||||
|
||||
// `gpt-5.6` and later only match a cached prefix reliably when the request carries
|
||||
// the routing key, so dropping it from the body silently forfeits the cache.
|
||||
it('puts the caller key on the request body', async () => {
|
||||
stubConfig()
|
||||
const { client, stream } = stubClient()
|
||||
|
||||
await getOpenAIResponsesCompletion([], new AbortController(), undefined, {
|
||||
openaiClient: client,
|
||||
promptCacheKey: 'admins:openai:gpt-5.6:chat'
|
||||
})
|
||||
|
||||
expect(stream.mock.calls[0][0]).toEqual(
|
||||
expect.objectContaining({ prompt_cache_key: 'admins:openai:gpt-5.6:chat' })
|
||||
)
|
||||
})
|
||||
|
||||
it('omits the field entirely when the caller withdrew the key', async () => {
|
||||
stubConfig()
|
||||
const { client, stream } = stubClient()
|
||||
|
||||
await getOpenAIResponsesCompletion([], new AbortController(), undefined, {
|
||||
openaiClient: client,
|
||||
promptCacheKey: undefined
|
||||
})
|
||||
|
||||
expect(stream.mock.calls[0][0]).not.toHaveProperty('prompt_cache_key')
|
||||
})
|
||||
})
|
||||
|
||||
describe('openAIWebSearchDetails', () => {
|
||||
it('prefers the queries array over the deprecated singular query', () => {
|
||||
expect(
|
||||
|
||||
@@ -179,13 +179,54 @@ function convertMessagesToResponsesInput(messages: ChatCompletionMessageParam[])
|
||||
}
|
||||
}
|
||||
|
||||
/** OpenAI rejects a key over 64 characters, and an Azure deployment name is user-chosen. */
|
||||
const MAX_PROMPT_CACHE_KEY_LENGTH = 64
|
||||
|
||||
/** FNV-1a. A routing key needs to be stable and distinct, not cryptographic. */
|
||||
function shortHash(value: string): string {
|
||||
let h = 0x811c9dc5
|
||||
for (let i = 0; i < value.length; i++) {
|
||||
h ^= value.charCodeAt(i)
|
||||
h = Math.imul(h, 0x01000193) >>> 0
|
||||
}
|
||||
return h.toString(16).padStart(8, '0')
|
||||
}
|
||||
|
||||
/**
|
||||
* Routing key for the provider's prompt cache. Built only from what fixes the prompt
|
||||
* prefix (the surface, plus the model) and never from anything per-request, since
|
||||
* requests sharing a prefix must reuse one key to land on the same cache. Workspace
|
||||
* splits traffic across the ~15 requests/minute one key sustains before it starts
|
||||
* missing again.
|
||||
*/
|
||||
export function buildPromptCacheKey(
|
||||
surface: string,
|
||||
modelProvider: { provider: string; model: string },
|
||||
workspace: string
|
||||
): string {
|
||||
const key = [workspace, modelProvider.provider, modelProvider.model, surface].join(':')
|
||||
if (key.length <= MAX_PROMPT_CACHE_KEY_LENGTH) {
|
||||
return key
|
||||
}
|
||||
// Same shape as the backend's `bounded_prompt_cache_key`: a readable head keeps the
|
||||
// key traceable, and the digest carries every distinction the head lost. Truncating
|
||||
// alone would collapse a long workspace's models and surfaces onto one key.
|
||||
const suffix = shortHash(key)
|
||||
return `${key.slice(0, MAX_PROMPT_CACHE_KEY_LENGTH - suffix.length - 1)}:${suffix}`
|
||||
}
|
||||
|
||||
function convertCompletionConfigToResponsesConfig(
|
||||
config: ChatCompletionCreateParams
|
||||
config: ChatCompletionCreateParams,
|
||||
promptCacheKey?: string
|
||||
): Record<string, any> {
|
||||
const responsesConfig: Record<string, any> = {
|
||||
model: config.model
|
||||
}
|
||||
|
||||
if (promptCacheKey) {
|
||||
responsesConfig.prompt_cache_key = promptCacheKey
|
||||
}
|
||||
|
||||
// Map max_tokens or max_completion_tokens to max_output_tokens
|
||||
if ('max_completion_tokens' in config && config.max_completion_tokens) {
|
||||
responsesConfig.max_output_tokens = config.max_completion_tokens
|
||||
@@ -223,6 +264,7 @@ export async function getOpenAIResponsesCompletion(
|
||||
webSearch?: boolean
|
||||
reasoningEffort?: string
|
||||
reasoningSummary?: boolean
|
||||
promptCacheKey?: string
|
||||
}
|
||||
) {
|
||||
const { provider, config } = getProviderAndCompletionConfig({
|
||||
@@ -233,7 +275,7 @@ export async function getOpenAIResponsesCompletion(
|
||||
})
|
||||
const { instructions, input } = convertMessagesToResponsesInput(messages)
|
||||
const responsesConfig = applyReasoningToConfig(
|
||||
convertCompletionConfigToResponsesConfig(config),
|
||||
convertCompletionConfigToResponsesConfig(config, options?.promptCacheKey),
|
||||
'responses',
|
||||
options?.reasoningEffort
|
||||
)
|
||||
@@ -291,6 +333,10 @@ export async function* getOpenAIResponsesCompletionStream(
|
||||
forceModelProvider: options?.forceModelProvider
|
||||
})
|
||||
const { instructions, input } = convertMessagesToResponsesInput(messages)
|
||||
// No prompt cache key here: a rejected key has to be retried without it, and this is
|
||||
// an async generator, so the caller's try/catch never sees the failure (invoking a
|
||||
// generator runs none of its body). One-shot generations have no repeated prefix to
|
||||
// route anyway; the chat loop is the surface that does, and it can retry.
|
||||
const responsesConfig = applyReasoningToConfig(
|
||||
convertCompletionConfigToResponsesConfig(config),
|
||||
'responses',
|
||||
@@ -575,6 +621,7 @@ export async function getNonStreamingOpenAIResponsesCompletion(
|
||||
})
|
||||
|
||||
const { instructions, input } = convertMessagesToResponsesInput(messages)
|
||||
// No prompt cache key, for the same reason as the streaming variant above.
|
||||
const responsesConfig = convertCompletionConfigToResponsesConfig(config)
|
||||
|
||||
const fetchOptions: {
|
||||
|
||||
Reference in New Issue
Block a user