From 57e98b6d86256a90e2047909be931e265fc0c4b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=8C=E4=B8=AB=E8=AE=B2=E6=A2=B5?= Date: Wed, 16 Sep 2026 23:23:32 +0800 Subject: [PATCH] feat(ai): session-level prompt cache key for openai responses Closes #9232 --- .../src/components/editor/AiAssistant.vue | 1 + apps/desktop/src/lib/ai/ai.ts | 5 ++ apps/desktop/src/lib/backend/tauri.ts | 2 + crates/dbx-core/src/agent_loop.rs | 28 +++++++++ crates/dbx-core/src/ai.rs | 59 +++++++++++++++++++ crates/dbx-web/src/routes/ai.rs | 1 + src-tauri/src/commands/ai.rs | 1 + 7 files changed, 97 insertions(+) diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index b3a906023c..8aa5e24583 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -3151,6 +3151,7 @@ async function send() { confirmedConnectionId: confirmedTargetConnId, confirmedDatabase: confirmedTargetDb, confirmedSchema: confirmedTargetSchema, + promptCacheKey: conversationId.value || undefined, }, history, (event: AgentEvent) => { diff --git a/apps/desktop/src/lib/ai/ai.ts b/apps/desktop/src/lib/ai/ai.ts index f7505a7f41..46ba8f6f0a 100644 --- a/apps/desktop/src/lib/ai/ai.ts +++ b/apps/desktop/src/lib/ai/ai.ts @@ -132,6 +132,8 @@ export interface AiRequestInput { confirmedConnectionId?: string; confirmedDatabase?: string; confirmedSchema?: string; + /** Stable per-conversation key forwarded to the Responses API. */ + promptCacheKey?: string; } export interface AiNamespaceSelection { @@ -194,6 +196,7 @@ export async function runAiAction(input: AiRequestInput, history?: api.AiMessage messages, taskContract, maxTokens, + promptCacheKey: input.promptCacheKey, }); } @@ -209,6 +212,7 @@ export async function runAiStream(input: AiRequestInput, history: api.AiMessage[ messages, taskContract, maxTokens, + promptCacheKey: input.promptCacheKey, }, (chunk) => { if (!chunk.done) { @@ -232,6 +236,7 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa messages, taskContract, maxTokens, + promptCacheKey: input.promptCacheKey, }, input.context.connectionId, input.context.database, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 0e70c7f453..4f8cb02588 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -498,6 +498,8 @@ export interface AiCompletionRequest { messages: AiMessage[]; taskContract?: AiTaskContract; maxTokens?: number; + /** Stable per-conversation key used by the Responses API prompt cache. */ + promptCacheKey?: string; } export interface AiModelInfo { diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs index 1b9f925403..5bad9efbb5 100644 --- a/crates/dbx-core/src/agent_loop.rs +++ b/crates/dbx-core/src/agent_loop.rs @@ -85,6 +85,9 @@ pub struct AgentLoopContext { /// Turn limit for this run, already clamped by the settings layer. /// Callers that have no user setting should pass [`DEFAULT_MAX_AGENT_TURNS`]. pub max_agent_turns: u32, + /// Stable per-conversation key forwarded to providers that support prompt + /// caching (OpenAI Responses API). `None` disables the field entirely. + pub prompt_cache_key: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -283,6 +286,7 @@ pub async fn run_agent_loop( &on_event, cancelled, false, + agent_ctx.prompt_cache_key.as_deref(), ) .await, CompactResult::Cancelled @@ -306,6 +310,7 @@ pub async fn run_agent_loop( &tools, max_tokens, task_contract.clone(), + agent_ctx.prompt_cache_key.clone(), ); // Stream the LLM response, collecting text and tool_calls. @@ -372,6 +377,7 @@ pub async fn run_agent_loop( &on_event, cancelled, true, + agent_ctx.prompt_cache_key.as_deref(), ) .await; match compacted { @@ -659,6 +665,7 @@ fn build_tool_request( _tools: &[ToolDefinition], // Tools are injected in ai::stream_with_tools, not via AiCompletionRequest. max_tokens: Option, task_contract: Option, + prompt_cache_key: Option, ) -> AiCompletionRequest { // Note: tools are passed via the body, not via AiCompletionRequest. // The actual injection happens in stream_with_tools. @@ -668,6 +675,7 @@ fn build_tool_request( messages: messages.to_vec(), task_contract, max_tokens: max_tokens.or(Some(4096)), + prompt_cache_key, } } @@ -920,6 +928,7 @@ async fn run_agent_loop_text_only( messages: messages.to_vec(), task_contract: task_contract.cloned(), max_tokens: max_tokens.or(Some(4096)), + prompt_cache_key: agent_ctx.prompt_cache_key.clone(), }; for attempt in 0..=MAX_CONTRACT_REPAIR_ATTEMPTS { @@ -1127,6 +1136,7 @@ async fn maybe_compact( on_event: &(impl Fn(AgentEvent) + Send + Sync), cancelled: &Notify, force: bool, + prompt_cache_key: Option<&str>, ) -> CompactResult { let window = effective_context_window(config); let budget = prompt_budget(window, max_tokens); @@ -1183,6 +1193,7 @@ async fn maybe_compact( }], task_contract: None, max_tokens: Some(1024), + prompt_cache_key: prompt_cache_key.map(str::to_string), }; let summary = match cancelled.notified().now_or_never() { @@ -1516,6 +1527,23 @@ mod tests { assert_eq!(effective_context_window(&config), 65_536); } + #[test] + fn build_tool_request_carries_prompt_cache_key() { + let config: AiConfig = serde_json::from_value(serde_json::json!({ + "provider": "openai", + "model": "test-model", + })) + .unwrap(); + + // Every agent turn rebuilds the request here, so this is the join that + // used to drop the conversation key before it reached the provider. + let with_key = build_tool_request(&config, "system", &[], &[], Some(128), None, Some("conv-1".to_string())); + assert_eq!(with_key.prompt_cache_key.as_deref(), Some("conv-1")); + + let without_key = build_tool_request(&config, "system", &[], &[], Some(128), None, None); + assert!(without_key.prompt_cache_key.is_none()); + } + #[test] fn identifies_explicit_tool_rejection_errors() { for error in [ diff --git a/crates/dbx-core/src/ai.rs b/crates/dbx-core/src/ai.rs index 044f042b2d..7cbe673ea4 100644 --- a/crates/dbx-core/src/ai.rs +++ b/crates/dbx-core/src/ai.rs @@ -571,6 +571,19 @@ pub struct AiCompletionRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub task_contract: Option, pub max_tokens: Option, + /// Stable key for reusing the provider's prompt cache within one conversation. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_cache_key: Option, +} + +fn apply_prompt_cache_key(body: &mut serde_json::Value, request: &AiCompletionRequest) { + // Only OpenAI's Responses API documents this field. Other providers may + // expose a Responses-compatible route with a strict request schema. + if request.config.api_style == AiApiStyle::Responses && is_openai_api_config(&request.config) { + if let Some(key) = request.prompt_cache_key.as_deref().filter(|key| !key.is_empty()) { + body["prompt_cache_key"] = json!(key); + } + } } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -2402,6 +2415,7 @@ pub async fn call_responses_api(client: &reqwest::Client, request: AiCompletionR "max_output_tokens": responses_max_output_tokens(request.max_tokens, &request.config), }); crate::ai_effort::apply_runtime_effort(&mut body, &request.config); + apply_prompt_cache_key(&mut body, &request); let res = client .post(resolve_endpoint(&request.config)) @@ -3959,6 +3973,7 @@ async fn stream_responses_api( "stream": true, }); crate::ai_effort::apply_runtime_effort(&mut body, &request.config); + apply_prompt_cache_key(&mut body, request); let endpoint = resolve_endpoint(&request.config); let config = request.config.clone(); @@ -4661,6 +4676,7 @@ async fn stream_responses_with_tools( "stream": true, }); crate::ai_effort::apply_runtime_effort(&mut body, &request.config); + apply_prompt_cache_key(&mut body, request); let endpoint = resolve_endpoint(&request.config); let config = request.config.clone(); @@ -5462,6 +5478,7 @@ mod tests { }], task_contract: None, max_tokens: Some(64), + prompt_cache_key: None, } } @@ -5498,6 +5515,7 @@ mod tests { }], task_contract: None, max_tokens: Some(64), + prompt_cache_key: None, } } @@ -8386,6 +8404,7 @@ mod tests { }], task_contract: None, max_tokens: Some(64), + prompt_cache_key: None, }; let client = build_ai_http_client(&config, 10).unwrap(); let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); @@ -9217,6 +9236,46 @@ mod tests { .unwrap() } + fn prompt_cache_test_request(config: AiConfig, key: Option<&str>) -> AiCompletionRequest { + AiCompletionRequest { + config, + system_prompt: "Be concise.".to_string(), + messages: Vec::new(), + task_contract: None, + max_tokens: Some(64), + prompt_cache_key: key.map(str::to_string), + } + } + + #[test] + fn prompt_cache_key_is_injected_only_for_openai_responses() { + let body = |config: AiConfig, key: Option<&str>| { + let mut body = serde_json::json!({ "model": "test-model" }); + super::apply_prompt_cache_key(&mut body, &prompt_cache_test_request(config, key)); + body + }; + let key_of = |value: &serde_json::Value| value.get("prompt_cache_key").cloned(); + + let mut responses = test_config(AiProvider::Openai); + responses.api_style = AiApiStyle::Responses; + assert_eq!(key_of(&body(responses, Some("conv-42"))), Some(serde_json::Value::from("conv-42"))); + + // The field is documented only for OpenAI's Responses API. Chat + // completions, other providers, and empty keys must stay untouched. + let mut completions = test_config(AiProvider::Openai); + completions.api_style = AiApiStyle::Completions; + assert_eq!(key_of(&body(completions, Some("conv-42"))), None); + + let mut compatible = test_config(AiProvider::OpenaiCompatible); + compatible.api_style = AiApiStyle::Responses; + assert_eq!(key_of(&body(compatible, Some("conv-42"))), None); + + let mut responses = test_config(AiProvider::Openai); + responses.api_style = AiApiStyle::Responses; + assert_eq!(key_of(&body(responses.clone(), Some(""))), None); + assert_eq!(key_of(&body(responses, None)), None); + } + #[test] fn merge_global_max_retries_sets_for_api_providers() { for provider in &[ diff --git a/crates/dbx-web/src/routes/ai.rs b/crates/dbx-web/src/routes/ai.rs index d338afaeb5..c9fc01783b 100644 --- a/crates/dbx-web/src/routes/ai.rs +++ b/crates/dbx-web/src/routes/ai.rs @@ -460,6 +460,7 @@ pub async fn ai_agent_stream( cli_mcp_server_command: None, sql_permissions, max_agent_turns, + prompt_cache_key: request.prompt_cache_key.clone(), }; let sid = session_id.clone(); diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs index c7ac535b45..03f8e2e026 100644 --- a/src-tauri/src/commands/ai.rs +++ b/src-tauri/src/commands/ai.rs @@ -394,6 +394,7 @@ pub async fn ai_agent_stream( cli_mcp_server_command, sql_permissions, max_agent_turns, + prompt_cache_key: request.prompt_cache_key.clone(), }; let is_agent_mode = mode.as_deref() == Some("agent");