language_model: Allow explicit prompt cache keys (#64460)
## Objective Related conversations can’t share an OpenAI cache key without also sharing a thread ID. We currently derive `prompt_cache_key` directly from `thread_id`, tying cache reuse to conversation identity. ## Solution Add an optional `prompt_cache_key` to `LanguageModelRequest` and use it in both OpenAI Chat Completions and Responses conversions. We still fall back to `thread_id` when no key is supplied, so existing callers keep their current behavior. This lets callers share a cache key across parent and child conversations while keeping their conversation and transport identities separate. ## Testing Added regression coverage for explicit keys, legacy fallback, provider capability gating, and serialization compatibility. ## Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content adheres to Zed's UI standards - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Release Notes: - N/A
This commit is contained in:
1 parent
31c3c86010
commit
c5d6f0b2ca
10 files changed
+113
-2
No files matched your search
@@ -4148,6 +4148,7 @@ impl Thread {
|
||||
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: Some(self.id.to_string()),
|
||||
prompt_cache_key: None,
|
||||
prompt_id: Some(self.prompt_id.to_string()),
|
||||
intent: Some(completion_intent),
|
||||
messages,
|
||||
|
||||
@@ -546,6 +546,7 @@ impl CodegenAlternative {
|
||||
|
||||
LanguageModelRequest {
|
||||
thread_id: Some(session_id),
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: Some(CompletionIntent::InlineAssist),
|
||||
tools,
|
||||
@@ -629,6 +630,7 @@ impl CodegenAlternative {
|
||||
|
||||
LanguageModelRequest {
|
||||
thread_id: Some(session_id),
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: Some(CompletionIntent::InlineAssist),
|
||||
tools: Vec::new(),
|
||||
|
||||
@@ -265,6 +265,7 @@ impl TerminalInlineAssistant {
|
||||
|
||||
LanguageModelRequest {
|
||||
thread_id: Some(session_id),
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: Some(CompletionIntent::TerminalInlineAssist),
|
||||
messages: vec![request_message],
|
||||
|
||||
@@ -895,6 +895,7 @@ mod tests {
|
||||
},
|
||||
],
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -1008,6 +1009,7 @@ mod tests {
|
||||
},
|
||||
],
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -1078,6 +1080,7 @@ mod tests {
|
||||
reasoning_details: None,
|
||||
}],
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -1130,6 +1133,7 @@ mod tests {
|
||||
reasoning_details: None,
|
||||
}],
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -1200,6 +1204,7 @@ mod tests {
|
||||
reasoning_details: None,
|
||||
}],
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -1262,6 +1267,7 @@ mod tests {
|
||||
},
|
||||
],
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -1313,6 +1319,7 @@ mod tests {
|
||||
}],
|
||||
thinking_effort: None,
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
|
||||
@@ -1091,6 +1091,7 @@ fn into_copilot_responses(
|
||||
.and_then(|_| request.effective_max_output_tokens(model.max_output_tokens()));
|
||||
let LanguageModelRequest {
|
||||
thread_id: _,
|
||||
prompt_cache_key: _,
|
||||
prompt_id: _,
|
||||
intent: _,
|
||||
messages,
|
||||
|
||||
@@ -4150,6 +4150,7 @@ impl GitPanel {
|
||||
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: Some(CompletionIntent::GenerateGitCommitMessage),
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
|
||||
@@ -469,6 +469,11 @@ pub enum CompletionIntent {
|
||||
#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
|
||||
pub struct LanguageModelRequest {
|
||||
pub thread_id: Option<String>,
|
||||
/// Provider cache affinity, independent of conversation and transport session identity.
|
||||
///
|
||||
/// Supporting providers fall back to `thread_id` when absent for legacy callers.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_cache_key: Option<String>,
|
||||
pub prompt_id: Option<String>,
|
||||
pub intent: Option<CompletionIntent>,
|
||||
pub messages: Vec<LanguageModelRequestMessage>,
|
||||
@@ -545,6 +550,33 @@ pub struct LanguageModelResponseMessage {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn request_prompt_cache_key_serialization() -> serde_json::Result<()> {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: Some("thread".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let mut serialized = serde_json::to_value(&request)?;
|
||||
assert!(serialized.get("prompt_cache_key").is_none());
|
||||
assert_eq!(
|
||||
serde_json::from_value::<LanguageModelRequest>(serialized.clone())?,
|
||||
request
|
||||
);
|
||||
|
||||
serialized["prompt_cache_key"] = serde_json::json!("cache-affinity");
|
||||
let explicit: LanguageModelRequest = serde_json::from_value(serialized.clone())?;
|
||||
assert_eq!(explicit.prompt_cache_key.as_deref(), Some("cache-affinity"));
|
||||
assert_eq!(explicit.thread_id, request.thread_id);
|
||||
assert_eq!(serde_json::to_value(explicit)?, serialized);
|
||||
|
||||
serialized["prompt_cache_key"] = serde_json::Value::Null;
|
||||
assert_eq!(
|
||||
serde_json::from_value::<LanguageModelRequest>(serialized)?,
|
||||
request
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_output_limit_serialization() -> serde_json::Result<()> {
|
||||
let request = LanguageModelRequest::default();
|
||||
|
||||
@@ -877,6 +877,7 @@ mod tests {
|
||||
tools: vec![],
|
||||
tool_choice: None,
|
||||
thread_id: Some("abcdef".into()),
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -925,6 +926,7 @@ mod tests {
|
||||
tools: vec![],
|
||||
tool_choice: None,
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
@@ -969,6 +971,7 @@ mod tests {
|
||||
tool_choice: None,
|
||||
temperature: None,
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
stop: vec![],
|
||||
|
||||
@@ -876,6 +876,7 @@ mod tests {
|
||||
thinking_effort: None,
|
||||
speed: None,
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
compact_at_tokens: None,
|
||||
@@ -1019,6 +1020,7 @@ mod tests {
|
||||
thinking_effort: None,
|
||||
speed: None,
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
compact_at_tokens: None,
|
||||
@@ -1086,6 +1088,7 @@ mod tests {
|
||||
thinking_effort: None,
|
||||
speed: None,
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
compact_at_tokens: None,
|
||||
|
||||
@@ -216,7 +216,7 @@ pub fn into_open_ai(
|
||||
None
|
||||
},
|
||||
prompt_cache_key: if supports_prompt_cache_key {
|
||||
request.thread_id
|
||||
request.prompt_cache_key.or(request.thread_id)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
@@ -268,6 +268,7 @@ pub fn into_open_ai_response(
|
||||
|
||||
let LanguageModelRequest {
|
||||
thread_id,
|
||||
prompt_cache_key,
|
||||
prompt_id: _,
|
||||
intent: _,
|
||||
messages,
|
||||
@@ -404,7 +405,7 @@ pub fn into_open_ai_response(
|
||||
}),
|
||||
tools,
|
||||
prompt_cache_key: if supports_prompt_cache_key {
|
||||
thread_id
|
||||
prompt_cache_key.or(thread_id)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
@@ -1524,6 +1525,51 @@ mod tests {
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn prompt_cache_key_respects_override_fallback_and_capability() -> Result<()> {
|
||||
for (explicit, thread, supported, expected) in [
|
||||
(Some("cache"), Some("thread"), true, Some("cache")),
|
||||
(Some("cache"), None, true, Some("cache")),
|
||||
(None, Some("thread"), true, Some("thread")),
|
||||
(None, None, true, None),
|
||||
(Some("cache"), Some("thread"), false, None),
|
||||
(None, Some("thread"), false, None),
|
||||
] {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: thread.map(str::to_owned),
|
||||
prompt_cache_key: explicit.map(str::to_owned),
|
||||
..Default::default()
|
||||
};
|
||||
let chat = into_open_ai(
|
||||
request.clone(),
|
||||
"gpt-5",
|
||||
true,
|
||||
supported,
|
||||
None,
|
||||
ChatCompletionMaxTokensParameter::MaxCompletionTokens,
|
||||
None,
|
||||
false,
|
||||
)?;
|
||||
let response = into_open_ai_response(
|
||||
request,
|
||||
"gpt-5",
|
||||
true,
|
||||
supported,
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
&OPEN_AI_PROVIDER_ID,
|
||||
)?;
|
||||
assert_eq!(chat.prompt_cache_key.as_deref(), expected);
|
||||
assert_eq!(response.prompt_cache_key.as_deref(), expected);
|
||||
assert_eq!(
|
||||
response.into_compact_request().prompt_cache_key.as_deref(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn map_response_events(events: Vec<ResponsesStreamEvent>) -> Vec<LanguageModelCompletionEvent> {
|
||||
block_on(async {
|
||||
OpenAiResponseEventMapper::new(OPEN_AI_PROVIDER_ID)
|
||||
@@ -1795,6 +1841,7 @@ mod tests {
|
||||
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: Some("thread-123".into()),
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![
|
||||
@@ -2001,6 +2048,7 @@ mod tests {
|
||||
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2095,6 +2143,7 @@ mod tests {
|
||||
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2187,6 +2236,7 @@ mod tests {
|
||||
fn into_open_ai_response_replays_reasoning_without_encrypted_content() {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2268,6 +2318,7 @@ mod tests {
|
||||
fn into_open_ai_response_omits_reasoning_when_thinking_is_disabled_and_none_is_unsupported() {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2362,6 +2413,7 @@ mod tests {
|
||||
] {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2415,6 +2467,7 @@ mod tests {
|
||||
] {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2461,6 +2514,7 @@ mod tests {
|
||||
fn into_open_ai_can_send_max_tokens_parameter() -> Result<()> {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2538,6 +2592,7 @@ mod tests {
|
||||
fn into_open_ai_response_sends_none_reasoning_when_thinking_is_disabled() -> Result<()> {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2580,6 +2635,7 @@ mod tests {
|
||||
fn into_open_ai_response_uses_default_effort_when_selected_effort_is_none() -> Result<()> {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2624,6 +2680,7 @@ mod tests {
|
||||
fn into_open_ai_response_replays_assistant_phase() {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -2719,6 +2776,7 @@ mod tests {
|
||||
});
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![
|
||||
@@ -2798,6 +2856,7 @@ mod tests {
|
||||
fn into_open_ai_response_replays_reasoning_details_but_not_thinking_text() {
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![LanguageModelRequestMessage {
|
||||
@@ -4098,6 +4157,7 @@ mod tests {
|
||||
};
|
||||
let request = LanguageModelRequest {
|
||||
thread_id: None,
|
||||
prompt_cache_key: None,
|
||||
prompt_id: None,
|
||||
intent: None,
|
||||
messages: vec![
|
||||
|
||||
Reference in new issue
Block a user