diff --git a/src/crates/core/src/agentic/agents/definitions/custom/subagent.rs b/src/crates/core/src/agentic/agents/definitions/custom/subagent.rs index 3f5f7f0d06..6d32411791 100644 --- a/src/crates/core/src/agentic/agents/definitions/custom/subagent.rs +++ b/src/crates/core/src/agentic/agents/definitions/custom/subagent.rs @@ -53,7 +53,7 @@ impl Agent for CustomSubagent { fn system_prompt_cache_identity(&self, _model_name: Option<&str>) -> SystemPromptCacheIdentity { let prompt_hash = hex::encode(Sha256::digest(self.prompt.as_bytes())); - SystemPromptCacheIdentity::new(self.id(), format!("custom_prompt_sha256:{prompt_hash}")) + SystemPromptCacheIdentity::new(format!("custom_prompt_sha256:{prompt_hash}")) } async fn build_prompt(&self, context: &PromptBuilderContext) -> BitFunResult { diff --git a/src/crates/core/src/agentic/agents/mod.rs b/src/crates/core/src/agentic/agents/mod.rs index 47e592c100..7f31916fc1 100644 --- a/src/crates/core/src/agentic/agents/mod.rs +++ b/src/crates/core/src/agentic/agents/mod.rs @@ -10,7 +10,7 @@ mod registry; // consecutive `[N]` display IDs after the dialog turn completes. pub(crate) mod citation_renumber; -use crate::agentic::session::SystemPromptCacheIdentity; +use crate::agentic::session::{SystemPromptCacheIdentity, UserContextCacheIdentity}; use crate::agentic::tools::framework::ToolExposure; use crate::agentic::WorkspaceBinding; use crate::util::errors::{BitFunError, BitFunResult}; @@ -98,13 +98,17 @@ pub trait Agent: Send + Sync + 'static { fn system_prompt_cache_identity(&self, model_name: Option<&str>) -> SystemPromptCacheIdentity { let template_name = self.prompt_template_name(model_name).trim(); - let prompt_identity = if template_name.is_empty() { + let scope_key = if template_name.is_empty() { format!("agent:{}", self.id()) } else { format!("template:{}", template_name) }; - SystemPromptCacheIdentity::new(self.id(), prompt_identity) + SystemPromptCacheIdentity::new(scope_key) + } + + fn user_context_cache_identity(&self) -> UserContextCacheIdentity { + UserContextCacheIdentity::new(self.user_context_policy().cache_scope_key()) } fn system_reminder_template_name(&self) -> Option<&str> { @@ -180,3 +184,23 @@ pub trait Agent: Send + Sync + 'static { false } } + +#[cfg(test)] +mod tests { + use super::{Agent, AgenticMode, MultitaskMode}; + + #[test] + fn shared_template_modes_share_system_prompt_cache_identity() { + let agentic = AgenticMode::new(); + let multitask = MultitaskMode::new(); + + assert_eq!( + agentic.system_prompt_cache_identity(None), + multitask.system_prompt_cache_identity(None) + ); + assert_eq!( + agentic.user_context_cache_identity(), + multitask.user_context_cache_identity() + ); + } +} diff --git a/src/crates/core/src/agentic/agents/prompt_builder/user_context.rs b/src/crates/core/src/agentic/agents/prompt_builder/user_context.rs index 86ec656a06..9b4ef41869 100644 --- a/src/crates/core/src/agentic/agents/prompt_builder/user_context.rs +++ b/src/crates/core/src/agentic/agents/prompt_builder/user_context.rs @@ -49,6 +49,18 @@ impl UserContextPolicy { pub fn includes(&self, section: UserContextSection) -> bool { self.sections.contains(§ion) } + + pub fn cache_scope_key(&self) -> String { + if self.sections.is_empty() { + return "empty".to_string(); + } + + self.sections + .iter() + .map(UserContextSection::cache_scope_label) + .collect::>() + .join("|") + } } impl Default for UserContextPolicy { @@ -57,6 +69,17 @@ impl Default for UserContextPolicy { } } +impl UserContextSection { + fn cache_scope_label(&self) -> &'static str { + match self { + Self::WorkspaceContext => "workspace_context", + Self::WorkspaceInstructions => "workspace_instructions", + Self::WorkspaceMemoryFiles => "workspace_memory_files", + Self::ProjectLayout => "project_layout", + } + } +} + #[cfg(test)] mod tests { use super::{UserContextPolicy, UserContextSection}; @@ -85,4 +108,17 @@ mod tests { fn default_policy_is_empty() { assert!(UserContextPolicy::default().sections.is_empty()); } + + #[test] + fn cache_scope_key_preserves_section_order() { + let policy = UserContextPolicy::empty() + .with_workspace_context() + .with_workspace_instructions() + .with_workspace_memory_files(); + + assert_eq!( + policy.cache_scope_key(), + "workspace_context|workspace_instructions|workspace_memory_files" + ); + } } diff --git a/src/crates/core/src/agentic/execution/execution_engine.rs b/src/crates/core/src/agentic/execution/execution_engine.rs index f4c3697662..26e5e89db1 100644 --- a/src/crates/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/core/src/agentic/execution/execution_engine.rs @@ -707,19 +707,32 @@ impl ExecutionEngine { }; let prompt_builder = PromptBuilder::new(prompt_context); - let user_context = if let Some(cached_user_context) = - self.session_manager.cached_user_context(session_id).await + let user_context_identity = current_agent.user_context_cache_identity(); + let user_context = if let Some(cached_user_context) = self + .session_manager + .cached_user_context(session_id, &user_context_identity) + .await { - debug!("User context cache hit: session_id={}", session_id); + debug!( + "User context cache hit: session_id={}, scope_key={}", + session_id, user_context_identity.scope_key + ); Some(cached_user_context) } else { - debug!("User context cache miss: session_id={}", session_id); + debug!( + "User context cache miss: session_id={}, scope_key={}", + session_id, user_context_identity.scope_key + ); let built_user_context = prompt_builder .build_user_context_reminder(¤t_agent.user_context_policy()) .await; if let Some(ref user_context) = built_user_context { self.session_manager - .remember_user_context(session_id, user_context.clone()) + .remember_user_context( + session_id, + user_context_identity.clone(), + user_context.clone(), + ) .await; } built_user_context @@ -751,15 +764,15 @@ impl ExecutionEngine { .await { debug!( - "System prompt cache hit: session_id={}, agent_id={}, prompt_identity={}", - session_id, identity.agent_id, identity.prompt_identity + "System prompt cache hit: session_id={}, scope_key={}", + session_id, identity.scope_key ); return Ok(cached_system_prompt); } debug!( - "System prompt cache miss: session_id={}, agent_id={}, prompt_identity={}", - session_id, identity.agent_id, identity.prompt_identity + "System prompt cache miss: session_id={}, scope_key={}", + session_id, identity.scope_key ); let system_prompt = current_agent.get_system_prompt(prompt_context).await?; self.session_manager diff --git a/src/crates/core/src/agentic/session/prompt_cache.rs b/src/crates/core/src/agentic/session/prompt_cache.rs index 87c31d3aad..34333d1212 100644 --- a/src/crates/core/src/agentic/session/prompt_cache.rs +++ b/src/crates/core/src/agentic/session/prompt_cache.rs @@ -22,15 +22,26 @@ impl Default for PromptCachePolicy { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct SystemPromptCacheIdentity { - pub agent_id: String, - pub prompt_identity: String, + pub scope_key: String, } impl SystemPromptCacheIdentity { - pub fn new(agent_id: impl Into, prompt_identity: impl Into) -> Self { + pub fn new(scope_key: impl Into) -> Self { Self { - agent_id: agent_id.into(), - prompt_identity: prompt_identity.into(), + scope_key: scope_key.into(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct UserContextCacheIdentity { + pub scope_key: String, +} + +impl UserContextCacheIdentity { + pub fn new(scope_key: impl Into) -> Self { + Self { + scope_key: scope_key.into(), } } } @@ -82,12 +93,37 @@ impl CachedSystemPrompt { } } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CachedUserContext { + #[serde(flatten)] + pub text: CachedPromptText, + pub identity: UserContextCacheIdentity, +} + +impl CachedUserContext { + pub fn new(identity: UserContextCacheIdentity, content: impl Into) -> Self { + Self { + text: CachedPromptText::new(content), + identity, + } + } + + pub fn is_usable( + &self, + identity: &UserContextCacheIdentity, + ttl: Option, + now_ms: u64, + ) -> bool { + self.identity == *identity && !self.text.is_expired(ttl, now_ms) + } +} + #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct SessionPromptCache { #[serde(default, skip_serializing_if = "Option::is_none")] pub system_prompt: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub user_context: Option, + pub user_context: Option, } impl SessionPromptCache { @@ -107,7 +143,7 @@ impl SessionPromptCache { if self .user_context .as_ref() - .is_some_and(|entry| entry.is_expired(ttl, now_ms)) + .is_some_and(|entry| entry.text.is_expired(ttl, now_ms)) { self.user_context = None; changed = true; @@ -208,6 +244,7 @@ impl SessionPromptCacheStore { pub fn lookup_user_context( &self, session_id: &str, + identity: &UserContextCacheIdentity, ttl: Option, ) -> PromptCacheLookup { let now_ms = current_time_ms(); @@ -217,11 +254,14 @@ impl SessionPromptCacheStore { .and_then(|cache| cache.user_context.clone()); match cached_entry { - Some(entry) if !entry.is_expired(ttl, now_ms) => PromptCacheLookup::Hit(entry.content), - Some(_) => { + Some(entry) if entry.is_usable(identity, ttl, now_ms) => { + PromptCacheLookup::Hit(entry.text.content) + } + Some(entry) if entry.text.is_expired(ttl, now_ms) => { self.invalidate(session_id, PromptCacheScope::UserContext); PromptCacheLookup::Expired } + Some(_) => PromptCacheLookup::Miss, None => PromptCacheLookup::Miss, } } @@ -240,7 +280,7 @@ impl SessionPromptCacheStore { } } - pub fn set_user_context(&self, session_id: &str, entry: CachedPromptText) { + pub fn set_user_context(&self, session_id: &str, entry: CachedUserContext) { if let Some(mut cache) = self.session_caches.get_mut(session_id) { cache.user_context = Some(entry); } else { @@ -284,8 +324,8 @@ fn current_time_ms() -> u64 { #[cfg(test)] mod tests { use super::{ - CachedPromptText, CachedSystemPrompt, PromptCacheLookup, PromptCacheScope, - SessionPromptCacheStore, SystemPromptCacheIdentity, + CachedSystemPrompt, CachedUserContext, PromptCacheLookup, PromptCacheScope, + SessionPromptCacheStore, SystemPromptCacheIdentity, UserContextCacheIdentity, }; use std::time::Duration; @@ -296,7 +336,7 @@ mod tests { store.set_system_prompt( "session-1", CachedSystemPrompt::new( - SystemPromptCacheIdentity::new("agentic", "template:agentic_mode"), + SystemPromptCacheIdentity::new("template:agentic_mode"), "prompt-a", ), ); @@ -304,7 +344,7 @@ mod tests { assert_eq!( match store.lookup_system_prompt( "session-1", - &SystemPromptCacheIdentity::new("agentic", "template:agentic_mode"), + &SystemPromptCacheIdentity::new("template:agentic_mode"), None, ) { PromptCacheLookup::Hit(value) => Some(value), @@ -315,7 +355,7 @@ mod tests { assert!(matches!( store.lookup_system_prompt( "session-1", - &SystemPromptCacheIdentity::new("debug", "template:debug_mode"), + &SystemPromptCacheIdentity::new("template:debug_mode"), None, ), PromptCacheLookup::Miss @@ -326,10 +366,20 @@ mod tests { fn expired_user_context_is_evicted_on_read() { let store = SessionPromptCacheStore::new(); store.create_session("session-1"); - store.set_user_context("session-1", CachedPromptText::new("stale context")); + store.set_user_context( + "session-1", + CachedUserContext::new( + UserContextCacheIdentity::new("workspace_context|workspace_instructions"), + "stale context", + ), + ); assert!(matches!( - store.lookup_user_context("session-1", Some(Duration::from_millis(0))), + store.lookup_user_context( + "session-1", + &UserContextCacheIdentity::new("workspace_context|workspace_instructions"), + Some(Duration::from_millis(0)), + ), PromptCacheLookup::Expired )); assert!(store @@ -346,11 +396,17 @@ mod tests { store.set_system_prompt( "session-1", CachedSystemPrompt::new( - SystemPromptCacheIdentity::new("agentic", "template:agentic_mode"), + SystemPromptCacheIdentity::new("template:agentic_mode"), "prompt-a", ), ); - store.set_user_context("session-1", CachedPromptText::new("context")); + store.set_user_context( + "session-1", + CachedUserContext::new( + UserContextCacheIdentity::new("workspace_context"), + "context", + ), + ); assert!(store.invalidate("session-1", PromptCacheScope::All)); diff --git a/src/crates/core/src/agentic/session/session_manager.rs b/src/crates/core/src/agentic/session/session_manager.rs index 1c8466e8e1..84395c424c 100644 --- a/src/crates/core/src/agentic/session/session_manager.rs +++ b/src/crates/core/src/agentic/session/session_manager.rs @@ -9,11 +9,11 @@ use crate::agentic::core::{ use crate::agentic::image_analysis::ImageContextData; use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ - CachedPromptText, CachedSystemPrompt, EvidenceLedgerCheckpoint, EvidenceLedgerEvent, + CachedSystemPrompt, CachedUserContext, EvidenceLedgerCheckpoint, EvidenceLedgerEvent, EvidenceLedgerEventStatus, EvidenceLedgerSummary, EvidenceLedgerTargetKind, FileReadState, FileReadStateStore, PromptCacheLookup, PromptCachePolicy, PromptCacheScope, SessionContextStore, SessionEvidenceLedger, SessionPromptCache, SessionPromptCacheStore, - SystemPromptCacheIdentity, + SystemPromptCacheIdentity, UserContextCacheIdentity, }; use crate::infrastructure::ai::get_global_ai_client_factory; use crate::service::config::{ @@ -1205,12 +1205,17 @@ impl SessionManager { .await; } - pub async fn cached_user_context(&self, session_id: &str) -> Option { + pub async fn cached_user_context( + &self, + session_id: &str, + identity: &UserContextCacheIdentity, + ) -> Option { self.ensure_prompt_cache_loaded(session_id).await; - match self - .prompt_cache_store - .lookup_user_context(session_id, self.config.prompt_cache_policy.cache_ttl) - { + match self.prompt_cache_store.lookup_user_context( + session_id, + identity, + self.config.prompt_cache_policy.cache_ttl, + ) { PromptCacheLookup::Hit(user_context) => Some(user_context), PromptCacheLookup::Miss => None, PromptCacheLookup::Expired => { @@ -1224,10 +1229,15 @@ impl SessionManager { } } - pub async fn remember_user_context(&self, session_id: &str, user_context: String) { + pub async fn remember_user_context( + &self, + session_id: &str, + identity: UserContextCacheIdentity, + user_context: String, + ) { self.ensure_prompt_cache_loaded(session_id).await; self.prompt_cache_store - .set_user_context(session_id, CachedPromptText::new(user_context)); + .set_user_context(session_id, CachedUserContext::new(identity, user_context)); self.persist_prompt_cache_best_effort(session_id, "user_context_cached") .await; } @@ -3812,6 +3822,7 @@ mod tests { use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ PromptCachePolicy, PromptCacheScope, SessionContextStore, SystemPromptCacheIdentity, + UserContextCacheIdentity, }; use crate::infrastructure::PathManager; use crate::service::config::types::{ @@ -4986,7 +4997,10 @@ mod tests { ) .await .expect("session should be created"); - let identity = SystemPromptCacheIdentity::new("agentic", "template:agentic_mode"); + let identity = SystemPromptCacheIdentity::new("template:agentic_mode"); + let user_context_identity = UserContextCacheIdentity::new( + "workspace_context|workspace_instructions|workspace_memory_files|project_layout", + ); manager .remember_system_prompt( @@ -4996,7 +5010,11 @@ mod tests { ) .await; manager - .remember_user_context(&session.session_id, "cached user context".to_string()) + .remember_user_context( + &session.session_id, + user_context_identity.clone(), + "cached user context".to_string(), + ) .await; let restored_manager = test_manager(persistence_manager); @@ -5013,7 +5031,7 @@ mod tests { ); assert_eq!( restored_manager - .cached_user_context(&session.session_id) + .cached_user_context(&session.session_id, &user_context_identity) .await, Some("cached user context".to_string()) ); @@ -5037,7 +5055,10 @@ mod tests { ) .await .expect("session should be created"); - let identity = SystemPromptCacheIdentity::new("agentic", "template:agentic_mode"); + let identity = SystemPromptCacheIdentity::new("template:agentic_mode"); + let user_context_identity = UserContextCacheIdentity::new( + "workspace_context|workspace_instructions|workspace_memory_files|project_layout", + ); manager .remember_system_prompt( @@ -5047,7 +5068,11 @@ mod tests { ) .await; manager - .remember_user_context(&session.session_id, "cached user context".to_string()) + .remember_user_context( + &session.session_id, + user_context_identity.clone(), + "cached user context".to_string(), + ) .await; manager @@ -5068,7 +5093,7 @@ mod tests { ); assert_eq!( restored_manager - .cached_user_context(&session.session_id) + .cached_user_context(&session.session_id, &user_context_identity) .await, None ); @@ -5111,7 +5136,10 @@ mod tests { ) .await .expect("session should be created"); - let identity = SystemPromptCacheIdentity::new("agentic", "template:agentic_mode"); + let identity = SystemPromptCacheIdentity::new("template:agentic_mode"); + let user_context_identity = UserContextCacheIdentity::new( + "workspace_context|workspace_instructions|workspace_memory_files|project_layout", + ); manager .remember_system_prompt( @@ -5121,7 +5149,11 @@ mod tests { ) .await; manager - .remember_user_context(&session.session_id, "cached user context".to_string()) + .remember_user_context( + &session.session_id, + user_context_identity.clone(), + "cached user context".to_string(), + ) .await; assert_eq!( @@ -5131,7 +5163,9 @@ mod tests { Some("cached system prompt".to_string()) ); assert_eq!( - manager.cached_user_context(&session.session_id).await, + manager + .cached_user_context(&session.session_id, &user_context_identity) + .await, Some("cached user context".to_string()) ); @@ -5161,7 +5195,7 @@ mod tests { ); assert_eq!( restored_manager - .cached_user_context(&session.session_id) + .cached_user_context(&session.session_id, &user_context_identity) .await, None );