Skip to content

feat(chat): add prompt cache guard for mode switching #927

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wsp1911 merged 1 commit into from
May 28, 2026
Merged

feat(chat): add prompt cache guard for mode switching #927

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions src/apps/desktop/src/api/agentic_api.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,12 @@ pub struct GetSessionRequest {
pub struct SessionResponse {
pub session_id: String,
pub session_name: String,
/// Current/default mode selection for the next dialog turn.
pub agent_type: String,
/// Mode of the last surviving user dialog turn in session history.
pub last_user_dialog_agent_type: Option<String>,
/// Mode of the most recent user submission accepted by the scheduler.
pub last_submitted_agent_type: Option<String>,
pub state: String,
pub turn_count: usize,
pub created_at: u64,
Expand Down Expand Up @@ -1459,6 +1464,8 @@ pub async fn list_sessions(
session_id: summary.session_id,
session_name: summary.session_name,
agent_type: summary.agent_type,
last_user_dialog_agent_type: summary.last_user_dialog_agent_type,
last_submitted_agent_type: summary.last_submitted_agent_type,
state: format!("{:?}", summary.state),
turn_count: summary.turn_count,
created_at: system_time_to_unix_secs(summary.created_at),
Expand Down Expand Up @@ -1522,6 +1529,7 @@ pub async fn get_available_modes(state: State<'_, AppState>) -> Result<Vec<ModeI
is_readonly: info.is_readonly,
tool_count: info.tool_count,
default_tools: info.default_tools,
prompt_cache_scope_key: info.prompt_cache_scope_key,
})
.collect();

Expand All @@ -1542,6 +1550,7 @@ pub struct ModeInfoDTO {
pub is_readonly: bool,
pub tool_count: usize,
pub default_tools: Vec<String>,
pub prompt_cache_scope_key: String,
}

fn assistant_bootstrap_outcome_to_response(
Expand Down Expand Up @@ -1600,6 +1609,8 @@ fn session_to_response(session: Session) -> SessionResponse {
session_id: session.session_id,
session_name: session.session_name,
agent_type: session.agent_type,
last_user_dialog_agent_type: session.last_user_dialog_agent_type,
last_submitted_agent_type: session.last_submitted_agent_type,
state: format!("{:?}", session.state),
turn_count: session.dialog_turn_ids.len(),
created_at: system_time_to_unix_secs(session.created_at),
Expand Down
10 changes: 10 additions & 0 deletions src/crates/core/src/agentic/agents/registry/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,11 @@ pub struct AgentInfo {
pub is_review: bool,
pub tool_count: usize,
pub default_tools: Vec<String>,
/// Combined prompt-cache compatibility key for frontend mode-switch guards.
///
/// Modes that share this key can reuse the same session-level prompt cache
/// for the next accepted submission.
pub prompt_cache_scope_key: String,
#[serde(default)]
pub default_enabled: bool,
#[serde(default = "default_true")]
Expand Down Expand Up @@ -182,6 +187,11 @@ impl AgentInfo {
is_review: is_review_agent_entry(entry),
tool_count: default_tools.len(),
default_tools,
prompt_cache_scope_key: format!(
"{}||{}",
agent.system_prompt_cache_identity(None).scope_key,
agent.user_context_cache_identity().scope_key
),
default_enabled: true,
effective_enabled: true,
override_state: None,
Expand Down
70 changes: 47 additions & 23 deletions src/crates/core/src/agentic/coordination/coordinator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1257,6 +1257,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
session_id: session_id.to_string(),
session_name: "Recovered Session".to_string(),
agent_type: "agentic".to_string(),
last_user_dialog_agent_type: None,
last_submitted_agent_type: None,
created_by: None,
session_kind: SessionKind::Standard,
model_name: "default".to_string(),
Expand Down Expand Up @@ -1391,7 +1393,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
}
}

(crate::service::session::TurnStatus::Completed, final_response)
(
crate::service::session::TurnStatus::Completed,
final_response,
)
}

async fn persist_cancelled_dialog_turn(
Expand All @@ -1401,7 +1406,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
session_id: &str,
turn_id: &str,
) -> crate::service::session::TurnStatus {
info!("Dialog turn cancelled: session={}, turn={}", session_id, turn_id);
info!(
"Dialog turn cancelled: session={}, turn={}",
session_id, turn_id
);

// The execution engine only emits DialogTurnCancelled when cancellation is
// detected between rounds. If cancellation interrupted streaming mid-round,
Expand All @@ -1422,7 +1430,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
);
}

if let Err(error) = session_manager.cancel_dialog_turn(session_id, turn_id).await {
if let Err(error) = session_manager
.cancel_dialog_turn(session_id, turn_id)
.await
{
error!(
"Failed to cancel dialog turn in persistence: session_id={}, turn_id={}, error={}",
session_id, turn_id, error
Expand Down Expand Up @@ -2812,19 +2823,17 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
)
.await
{
Ok(execution_result) => {
Some(
Self::persist_completed_dialog_turn(
session_manager.as_ref(),
scheduler_notify_tx.as_ref(),
&session_id_clone,
&turn_id_clone,
&execution_result,
)
.await
.0,
Ok(execution_result) => Some(
Self::persist_completed_dialog_turn(
session_manager.as_ref(),
scheduler_notify_tx.as_ref(),
&session_id_clone,
&turn_id_clone,
&execution_result,
)
}
.await
.0,
),
Err(e) => {
if matches!(&e, BitFunError::Cancelled(_)) {
Some(
Expand Down Expand Up @@ -4241,14 +4250,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet

// Persist turn lifecycle before cleaning up the hidden subagent runtime.
let (workspace_turn_status, response_text) = match result {
Ok(exec_result) => Self::persist_completed_dialog_turn(
self.session_manager.as_ref(),
None,
&session_id,
&dialog_turn_id,
&exec_result,
)
.await,
Ok(exec_result) => {
Self::persist_completed_dialog_turn(
self.session_manager.as_ref(),
None,
&session_id,
&dialog_turn_id,
&exec_result,
)
.await
}
Err(e) => {
let turn_status = if matches!(&e, BitFunError::Cancelled(_)) {
Self::persist_cancelled_dialog_turn(
Expand Down Expand Up @@ -4836,6 +4847,19 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet
.await
}

/// Update the session-level prompt-cache guard mode for the latest
/// scheduler-accepted user submission.
pub async fn update_last_submitted_agent_type(
&self,
session_id: &str,
agent_type: &str,
) -> BitFunResult<()> {
let normalized = Self::normalize_agent_type(agent_type);
self.session_manager
.update_last_submitted_agent_type(session_id, &normalized)
.await
}

/// Emit event
async fn emit_event(&self, event: AgenticEvent) {
let _ = self
Expand Down
24 changes: 24 additions & 0 deletions src/crates/core/src/agentic/coordination/scheduler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -419,6 +419,8 @@ impl DialogScheduler {
match state {
None => {
let tid = self.start_turn(&session_id, &queued_turn).await?;
self.record_last_submitted_agent_type(&session_id, &queued_turn.agent_type)
.await;
Ok(DialogSubmitOutcome::Started {
session_id,
turn_id: tid,
Expand All @@ -428,6 +430,8 @@ impl DialogScheduler {
Some(SessionState::Error { .. }) => {
self.clear_queue(&session_id);
let tid = self.start_turn(&session_id, &queued_turn).await?;
self.record_last_submitted_agent_type(&session_id, &queued_turn.agent_type)
.await;
Ok(DialogSubmitOutcome::Started {
session_id,
turn_id: tid,
Expand All @@ -443,6 +447,8 @@ impl DialogScheduler {

if queue_non_empty {
self.enqueue(&session_id, queued_turn.clone())?;
self.record_last_submitted_agent_type(&session_id, &queued_turn.agent_type)
.await;
let started_tid = self.try_start_next_queued(&session_id).await?;
let outcome = match started_tid {
Some(tid) if tid == resolved_turn_id => DialogSubmitOutcome::Started {
Expand All @@ -457,6 +463,8 @@ impl DialogScheduler {
Ok(outcome)
} else {
let tid = self.start_turn(&session_id, &queued_turn).await?;
self.record_last_submitted_agent_type(&session_id, &queued_turn.agent_type)
.await;
Ok(DialogSubmitOutcome::Started {
session_id,
turn_id: tid,
Expand All @@ -466,7 +474,10 @@ impl DialogScheduler {

Some(SessionState::Processing { .. }) => {
let may_preempt = Self::user_message_may_preempt(&queued_turn.policy);
let accepted_agent_type = queued_turn.agent_type.clone();
self.enqueue(&session_id, queued_turn)?;
self.record_last_submitted_agent_type(&session_id, &accepted_agent_type)
.await;
if may_preempt {
self.round_yield_flags.request_yield(&session_id);
}
Expand All @@ -478,6 +489,19 @@ impl DialogScheduler {
}
}

async fn record_last_submitted_agent_type(&self, session_id: &str, agent_type: &str) {
if let Err(error) = self
.coordinator
.update_last_submitted_agent_type(session_id, agent_type)
.await
{
warn!(
"Failed to record last submitted agent type: session_id={}, agent_type={}, error={}",
session_id, agent_type, error
);
}
}

/// Number of messages currently queued for a session.
pub fn queue_depth(&self, session_id: &str) -> usize {
self.queues.get(session_id).map(|q| q.len()).unwrap_or(0)
Expand Down
28 changes: 24 additions & 4 deletions src/crates/core/src/agentic/core/session.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,25 @@ pub struct Session {
pub session_id: String,
pub session_name: String,
/// Current/default mode selection for the session.
/// This reflects what the next turn should run with by default, not
/// necessarily what the last surviving history turn used.
///
/// This is the mode the next dialog turn should run with by default. It is
/// not required to match either the last surviving history turn or the last
/// message submission accepted by the scheduler.
pub agent_type: String,
/// Cached mode of the last surviving user dialog turn in history.
/// `previous_agent_type` reminders should read this instead of `agent_type`
/// so rollback-to-empty can still be treated as a fresh mode entry.
///
/// Reminder builders use this value for `previous_agent_type` so
/// first-entry vs ongoing mode prompts follow the surviving transcript
/// after rollbacks or turn truncation.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_user_dialog_agent_type: Option<String>,
/// Mode of the most recent user submission accepted by the scheduler.
///
/// Unlike `last_user_dialog_agent_type`, this value is not rewound by
/// history rollback. It tracks session-level prompt-cache compatibility for
/// the next accepted submission.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_submitted_agent_type: Option<String>,
#[serde(
default,
skip_serializing_if = "Option::is_none",
Expand Down Expand Up @@ -80,6 +91,7 @@ impl Session {
session_name,
agent_type,
last_user_dialog_agent_type: None,
last_submitted_agent_type: None,
created_by: None,
kind: SessionKind::Standard,
snapshot_session_id: None,
Expand All @@ -105,6 +117,7 @@ impl Session {
session_name,
agent_type,
last_user_dialog_agent_type: None,
last_submitted_agent_type: None,
created_by: None,
kind: SessionKind::Standard,
snapshot_session_id: None,
Expand Down Expand Up @@ -175,7 +188,14 @@ impl Default for SessionConfig {
pub struct SessionSummary {
pub session_id: String,
pub session_name: String,
/// Current/default mode selection for the session.
pub agent_type: String,
/// Mode of the last surviving user dialog turn in the session history.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_user_dialog_agent_type: Option<String>,
/// Mode of the most recent user submission accepted by the scheduler.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_submitted_agent_type: Option<String>,
#[serde(
default,
skip_serializing_if = "Option::is_none",
Expand Down
18 changes: 17 additions & 1 deletion src/crates/core/src/agentic/persistence/manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,11 @@ struct StoredSessionStateFile {
// on persisted dialog turns via `DialogTurnData.agent_type`.
#[serde(default, skip_serializing_if = "Option::is_none")]
last_user_dialog_agent_type: Option<String>,
// Session-level prompt-cache guard state. This records the most recent user
// submission accepted by the scheduler and intentionally does not rewind on
// history rollback.
#[serde(default, skip_serializing_if = "Option::is_none")]
last_submitted_agent_type: Option<String>,
compression_state: CompressionState,
runtime_state: SessionState,
}
Expand Down Expand Up @@ -867,6 +872,8 @@ impl PersistenceManager {
session_id: session.session_id.clone(),
session_name: session.session_name.clone(),
agent_type: session.agent_type.clone(),
last_user_dialog_agent_type: session.last_user_dialog_agent_type.clone(),
last_submitted_agent_type: session.last_submitted_agent_type.clone(),
created_by: session
.created_by
.clone()
Expand Down Expand Up @@ -1934,6 +1941,7 @@ impl PersistenceManager {
config: session.config.clone(),
snapshot_session_id: session.snapshot_session_id.clone(),
last_user_dialog_agent_type: session.last_user_dialog_agent_type.clone(),
last_submitted_agent_type: session.last_submitted_agent_type.clone(),
compression_state: session.compression_state.clone(),
runtime_state: Self::sanitize_runtime_state(&session.state),
};
Expand Down Expand Up @@ -2005,7 +2013,12 @@ impl PersistenceManager {
agent_type: metadata.agent_type.clone(),
last_user_dialog_agent_type: stored_state
.as_ref()
.and_then(|value| value.last_user_dialog_agent_type.clone()),
.and_then(|value| value.last_user_dialog_agent_type.clone())
.or_else(|| metadata.last_user_dialog_agent_type.clone()),
last_submitted_agent_type: stored_state
.as_ref()
.and_then(|value| value.last_submitted_agent_type.clone())
.or_else(|| metadata.last_submitted_agent_type.clone()),
created_by: metadata.created_by.clone(),
kind: metadata.session_kind,
snapshot_session_id: stored_state
Expand Down Expand Up @@ -2042,6 +2055,7 @@ impl PersistenceManager {
},
snapshot_session_id: None,
last_user_dialog_agent_type: None,
last_submitted_agent_type: None,
compression_state: CompressionState::default(),
runtime_state: SessionState::Idle,
});
Expand Down Expand Up @@ -2085,6 +2099,8 @@ impl PersistenceManager {
session_id: metadata.session_id,
session_name: metadata.session_name,
agent_type: metadata.agent_type,
last_user_dialog_agent_type: metadata.last_user_dialog_agent_type,
last_submitted_agent_type: metadata.last_submitted_agent_type,
created_by: metadata.created_by,
kind: metadata.session_kind,
turn_count: metadata.turn_count,
Expand Down
Loading
Loading