diff --git a/migrations/20260727170000_add_message_model_content/down.sql b/migrations/20260727170000_add_message_model_content/down.sql new file mode 100644 index 0000000..2596679 --- /dev/null +++ b/migrations/20260727170000_add_message_model_content/down.sql @@ -0,0 +1 @@ +ALTER TABLE messages DROP COLUMN model_content; diff --git a/migrations/20260727170000_add_message_model_content/up.sql b/migrations/20260727170000_add_message_model_content/up.sql new file mode 100644 index 0000000..65b2372 --- /dev/null +++ b/migrations/20260727170000_add_message_model_content/up.sql @@ -0,0 +1 @@ +ALTER TABLE messages ADD COLUMN model_content TEXT; diff --git a/src/app.rs b/src/app.rs index b3b6f9a..acc8f1f 100644 --- a/src/app.rs +++ b/src/app.rs @@ -83,6 +83,7 @@ pub(crate) struct App { pub(super) composer: String, pub(super) queued_inputs: VecDeque, pub(super) conversation: Vec, + active_turn: Option, pub(super) a2ui: crate::a2ui::Store, pub(super) a2ui_history: Vec, pub(super) a2ui_history_index: Option, @@ -406,6 +407,7 @@ impl App { composer: String::new(), queued_inputs: VecDeque::new(), conversation: Vec::new(), + active_turn: None, a2ui: crate::a2ui::Store::default(), a2ui_history: Vec::new(), a2ui_history_index: None, @@ -533,6 +535,7 @@ impl App { composer: String::new(), queued_inputs: VecDeque::new(), conversation: Vec::new(), + active_turn: None, a2ui: crate::a2ui::Store::default(), a2ui_history: Vec::new(), a2ui_history_index: None, @@ -1479,6 +1482,7 @@ impl App { match loaded { Ok((messages, a2ui)) => { self.conversation = messages.into_iter().map(ChatMessage::from).collect(); + generation::promote_legacy_turn_summaries(&mut self.conversation); self.context_notice = None; self.clear_a2ui(); let (history, active, errors) = @@ -2292,6 +2296,7 @@ mod tests { reasoning_complete: false, reasoning_open: true, content: String::new(), + model_content: None, markdown: markdown::Content::new(), transcript: text_editor::Content::new(), a2ui_lines_processed: 0, @@ -2322,6 +2327,7 @@ mod tests { reasoning_complete: true, reasoning_open: false, content: content.into(), + model_content: None, markdown: markdown::Content::new(), transcript: text_editor::Content::new(), a2ui_lines_processed: 0, diff --git a/src/app/generation.rs b/src/app/generation.rs index 39959f1..4eaae88 100644 --- a/src/app/generation.rs +++ b/src/app/generation.rs @@ -5,6 +5,7 @@ const TITLE_INSTRUCTION: &str = "Give this conversation a short title of at most Reply with the title alone: no quotes, no trailing period, no explanation."; const TITLE_MAX_TOKENS: i32 = 48; const TITLE_MAX_CHARS: usize = 60; +const AGENTS_PREFIX: &str = "Project instructions from AGENTS.md:\n\n"; /// A one-shot generation that produces a session title. It runs against the /// transient KV cache, so it never creates or touches a stored session. @@ -60,6 +61,7 @@ pub(crate) struct ChatMessage { pub(super) reasoning_complete: bool, pub(super) reasoning_open: bool, pub(super) content: String, + pub(super) model_content: Option, pub(super) markdown: markdown::Content, pub(super) transcript: text_editor::Content, pub(super) a2ui_lines_processed: usize, @@ -68,7 +70,7 @@ pub(crate) struct ChatMessage { pub(super) a2ui_open_urls: Vec, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub(super) struct GenerationStats { pub(super) duration_ms: u64, pub(super) input_tokens: u32, @@ -76,6 +78,62 @@ pub(super) struct GenerationStats { pub(super) output_tokens: u32, } +impl GenerationStats { + fn add(&mut self, other: Self) { + self.duration_ms = self.duration_ms.saturating_add(other.duration_ms); + self.input_tokens = self.input_tokens.saturating_add(other.input_tokens); + self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens); + self.output_tokens = self.output_tokens.saturating_add(other.output_tokens); + } +} + +pub(super) struct TurnSummary { + started_at: Instant, + message_id: Option, + stats: GenerationStats, +} + +impl TurnSummary { + fn new() -> Self { + Self { + started_at: Instant::now(), + message_id: None, + stats: GenerationStats::default(), + } + } + + fn record(&mut self, input_tokens: u32, cached_tokens: u32, output_tokens: u32) { + self.stats.add(GenerationStats { + duration_ms: 0, + input_tokens, + cached_tokens, + output_tokens, + }); + } + + fn finish(mut self) -> GenerationStats { + self.stats.duration_ms = + u64::try_from(self.started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + self.stats + } +} + +pub(super) fn promote_legacy_turn_summaries(messages: &mut [ChatMessage]) { + let mut owner = None; + for index in 0..messages.len() { + if messages[index].user { + owner = Some((index, messages[index].generation_stats.is_some())); + } else if let Some(stats) = messages[index].generation_stats.take() + && let Some((owner, false)) = owner + { + messages[owner] + .generation_stats + .get_or_insert_default() + .add(stats); + } + } +} + const TOOL_PROTOCOL_CORRECTION: &str = "Tool protocol correction (no tool was executed): you emitted a complete tool call inside private reasoning. End reasoning, then emit the tool call again as assistant content using the required syntax; do not merely discuss it."; const EMPTY_RESPONSE_CORRECTION: &str = "Protocol correction: your previous response was empty. Continue the task now with either a valid tool call or a final answer."; @@ -168,6 +226,7 @@ impl From for ChatMessage { reasoning_complete: message.reasoning_complete, reasoning_open: false, content: message.content, + model_content: message.model_content, markdown: iced::widget::markdown::Content::new(), transcript: iced::widget::text_editor::Content::new(), a2ui_lines_processed: 0, @@ -241,7 +300,10 @@ fn chat_turn(message: &ChatMessage) -> ChatTurn { skip_previous_eos: false, reasoning: message.reasoning.clone(), reasoning_complete: message.reasoning_complete, - content: message.content.clone(), + content: message + .model_content + .clone() + .unwrap_or_else(|| message.content.clone()), } } @@ -294,16 +356,51 @@ fn has_chat_after_last_compaction(messages: &[ChatMessage]) -> bool { }) } +fn project_agents(path: &Path) -> Result, String> { + let path = path.join("AGENTS.md"); + match fs::read_to_string(&path) { + Ok(content) if content.is_empty() => Ok(None), + Ok(content) => Ok(Some(content)), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(format!("Could not read {}: {error}", path.display())), + } +} + +#[cfg(target_os = "macos")] +fn title_context(messages: impl IntoIterator) -> Vec { + let mut started = false; + messages + .into_iter() + .filter(|message| { + if !started { + started = message.user && !message.tool && !message.system; + } + started && !message.system && !message.content.trim().is_empty() + }) + .collect() +} + impl App { - fn chat_system_prompt(&self, model: ModelChoice, prompt: &str) -> String { + fn chat_system_prompt(&self, model: ModelChoice, prompt: &str, agents: Option<&str>) -> String { let mut prompt = crate::agent::system_prompt(model, prompt); if self.config.a2ui_enabled { prompt.push_str("\n\n"); prompt.push_str(crate::a2ui::SYSTEM_PROMPT); } + if let Some(agents) = agents { + prompt.push_str("\n\n"); + prompt.push_str(agents); + } prompt } + fn session_agents_prompt(&self) -> Option<&str> { + self.conversation + .iter() + .find(|message| message.system && message.content.starts_with(AGENTS_PREFIX)) + .map(|message| message.content.as_str()) + } + pub(super) fn can_compact_session(&self, session_id: i32) -> bool { !self.generating && self.selected_session == Some(session_id) @@ -345,6 +442,28 @@ impl App { )); return; } + #[cfg(target_os = "macos")] + let opening_turn = self.selected_session.is_none(); + #[cfg(target_os = "macos")] + let agents = if opening_turn { + let Some(project) = self + .projects + .iter() + .find(|project| Some(project.project.id) == self.selected_project) + else { + self.error = Some("The selected project is unavailable.".into()); + return; + }; + match project_agents(Path::new(&project.project.path)) { + Ok(agents) => agents.map(|agents| format!("{AGENTS_PREFIX}{agents}")), + Err(error) => { + self.error = Some(error); + return; + } + } + } else { + None + }; self.a2ui_auto_switch_pending = true; self.context_notice = None; #[cfg(target_os = "macos")] @@ -353,10 +472,12 @@ impl App { if !std::mem::take(&mut self.skip_compaction_once) && crate::compaction::should_compact(self.context_used, self.context_limit) { + self.active_turn.get_or_insert_with(TurnSummary::new); if let Err(error) = self.start_compaction( PendingContinuation::User(prompt), "soft limit before user turn", ) { + self.active_turn = None; self.error = Some(error); } return; @@ -376,7 +497,7 @@ impl App { } }; effective.turn.system_prompt = - self.chat_system_prompt(model, &effective.turn.system_prompt); + self.chat_system_prompt(model, &effective.turn.system_prompt, agents.as_deref()); effective.turn.system_prompt = crate::compaction::summary_system_prompt( &effective.turn.system_prompt, self.compaction_summary(), @@ -398,11 +519,10 @@ impl App { prompt.clone() }; #[cfg(target_os = "macos")] - let opening_turn = self.selected_session.is_none(); - #[cfg(target_os = "macos")] let mut injected_system = Vec::new(); #[cfg(target_os = "macos")] if opening_turn { + injected_system.extend(agents.clone()); injected_system.push(crate::agent::datetime_context()); } #[cfg(target_os = "macos")] @@ -418,15 +538,20 @@ impl App { .map(chat_turn) .collect::>(); #[cfg(target_os = "macos")] - messages.extend(injected_system.iter().map(|content| ChatTurn { - user: false, - tool: false, - system: true, - skip_previous_eos: false, - reasoning: None, - reasoning_complete: true, - content: content.clone(), - })); + messages.extend( + injected_system + .iter() + .skip(usize::from(opening_turn && agents.is_some())) + .map(|content| ChatTurn { + user: false, + tool: false, + system: true, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: content.clone(), + }), + ); #[cfg(target_os = "macos")] messages.push(ChatTurn { user: true, @@ -435,7 +560,7 @@ impl App { skip_previous_eos: false, reasoning: None, reasoning_complete: true, - content: model_prompt, + content: model_prompt.clone(), }); #[cfg(target_os = "macos")] @@ -466,6 +591,7 @@ impl App { let mut saved = match database.start_chat_turn( session_id, &prompt, + (model_prompt != prompt).then_some(model_prompt.as_str()), &injected_system, assistant_reasoning, ) { @@ -475,22 +601,38 @@ impl App { return; } }; + let user_id = saved[saved.len() - 2].id; + self.active_turn + .get_or_insert_with(TurnSummary::new) + .message_id = Some(user_id); if reminder_injected { self.system_prompt_seen_at = self.context_used; } let idle_timeout = Duration::from_secs(self.config.idle_timeout_minutes.max(1) as u64 * 60); + let checkpoint = if opening_turn { + CheckpointTarget::Local { + checkpoint: session_checkpoint_path(session_id), + bootstrap: Some(transient_cache_path()), + } + } else { + CheckpointTarget::Local { + checkpoint: session_checkpoint_path(session_id), + bootstrap: None, + } + }; self.active_generation = match service.generate( effective.engine, effective.turn, messages, - CheckpointTarget::Local(session_checkpoint_path(session_id)), + checkpoint, idle_timeout, ) { Ok(active) => Some(active), Err(error) => { self.generation_service = None; self.error = Some(error); + self.finish_turn_summary(); return; } }; @@ -533,10 +675,21 @@ impl App { if !self.config.generation.system_prompt.trim().is_empty() { reminders.push(self.config.generation.system_prompt.clone()); } + if let Some(agents) = self.session_agents_prompt() { + reminders.push(agents.to_owned()); + } reminders } pub(super) fn poll_generation(&mut self) -> bool { + let changed = self.poll_generation_step(); + if !self.generating { + self.finish_turn_summary(); + } + changed + } + + fn poll_generation_step(&mut self) -> bool { #[cfg(target_os = "macos")] if self.active_compaction.is_some() { return self.poll_compaction(); @@ -701,29 +854,14 @@ impl App { context_changed = true; } Ok(GenerationEvent::Finished(result)) => { - if let Ok(output) = &result { - let stats = GenerationStats { - duration_ms: u64::try_from(active.started_at.elapsed().as_millis()) - .unwrap_or(u64::MAX), - input_tokens: output.prompt_tokens, - cached_tokens: output.cached_tokens, - output_tokens: output.completion_tokens, - }; - if let Some(message) = self.conversation.last_mut() { - message.generation_stats = Some(stats); - if let Some(database) = &mut self.database - && let Err(error) = database.update_message_generation_stats( - message.id, - stats.duration_ms, - stats.input_tokens, - stats.cached_tokens, - stats.output_tokens, - ) - { - self.error = - Some(format!("Could not save generation summary: {error}")); - } - } + if let Ok(output) = &result + && let Some(turn) = &mut self.active_turn + { + turn.record( + output.prompt_tokens, + output.cached_tokens, + output.completion_tokens, + ); } match result { Ok(_) if self.stop_requested => { @@ -966,11 +1104,41 @@ impl App { return; } if let Some(prompt) = queued_prompt(self.queued_inputs.drain(..)) { + self.finish_turn_summary(); self.composer = prompt; self.start_generation(); } } + fn finish_turn_summary(&mut self) { + self.context_notice = None; + let Some(turn) = self.active_turn.take() else { + return; + }; + let Some(message_id) = turn.message_id else { + return; + }; + let stats = turn.finish(); + if let Some(message) = self + .conversation + .iter_mut() + .find(|message| message.id == message_id) + { + message.generation_stats = Some(stats); + } + if let Some(database) = &mut self.database + && let Err(error) = database.update_message_generation_stats( + message_id, + stats.duration_ms, + stats.input_tokens, + stats.cached_tokens, + stats.output_tokens, + ) + { + self.error = Some(format!("Could not save turn summary: {error}")); + } + } + #[cfg(target_os = "macos")] fn start_agent_tools(&mut self, calls: Vec) -> Result<(), String> { let session_id = self @@ -1013,8 +1181,11 @@ impl App { &self.config.runtime, &models_path(), )?; - effective.turn.system_prompt = - self.chat_system_prompt(model, &effective.turn.system_prompt); + effective.turn.system_prompt = self.chat_system_prompt( + model, + &effective.turn.system_prompt, + self.session_agents_prompt(), + ); effective.turn.system_prompt = crate::compaction::summary_system_prompt( &effective.turn.system_prompt, self.compaction_summary(), @@ -1065,13 +1236,17 @@ impl App { effective.engine, effective.turn, messages, - CheckpointTarget::Local(session_checkpoint_path(session_id)), + CheckpointTarget::Local { + checkpoint: session_checkpoint_path(session_id), + bootstrap: None, + }, idle_timeout, )?, ); self.tokens_per_second = None; self.activity = Some("Continuing after tools…".into()); self.stop_requested = false; + self.generating = true; Ok(()) } @@ -1110,8 +1285,11 @@ impl App { &self.config.runtime, &models_path(), )?; - effective.turn.system_prompt = - self.chat_system_prompt(model, &effective.turn.system_prompt); + effective.turn.system_prompt = self.chat_system_prompt( + model, + &effective.turn.system_prompt, + self.session_agents_prompt(), + ); effective.turn.system_prompt = crate::compaction::summary_system_prompt( &effective.turn.system_prompt, self.compaction_summary(), @@ -1141,6 +1319,7 @@ impl App { result, stage, }); + self.generating = true; self.activity = Some("Checking tool result context…".into()); Ok(()) } @@ -1259,8 +1438,11 @@ impl App { &self.config.runtime, &models_path(), )?; - effective.turn.system_prompt = - self.chat_system_prompt(model, &effective.turn.system_prompt); + effective.turn.system_prompt = self.chat_system_prompt( + model, + &effective.turn.system_prompt, + self.session_agents_prompt(), + ); let rebuild_system_prompt = effective.turn.system_prompt.clone(); effective.turn.system_prompt = crate::compaction::summary_system_prompt( &effective.turn.system_prompt, @@ -1464,7 +1646,10 @@ impl App { let start = compacted_context_start(&self.conversation); self.conversation[start..] .iter() - .filter(|message| !message.compaction) + .filter(|message| { + !message.compaction + && !(message.system && message.content.starts_with(AGENTS_PREFIX)) + }) .collect() } @@ -1495,9 +1680,9 @@ impl App { // The assistant row is inserted empty when a turn starts; an empty turn // is never useful context, and skipping it is what lets the automatic // pass summarize the user's opening message on its own. - let mut messages = stored + let messages = stored .into_iter() - .filter(|message| !message.compaction && !message.content.trim().is_empty()) + .filter(|message| !message.compaction) .map(|message| ChatTurn { user: message.user, tool: message.tool, @@ -1506,8 +1691,8 @@ impl App { reasoning: None, reasoning_complete: true, content: message.content, - }) - .collect::>(); + }); + let mut messages = title_context(messages); if messages.is_empty() { return Err("This session has no messages to summarize yet.".into()); } @@ -1655,10 +1840,12 @@ pub(super) fn session_title(reply: &str) -> Option { #[cfg(test)] mod tests { use super::{ - ChatMessage, TOOL_PROTOCOL_CORRECTION, compacted_context_start, correction_already_sent, - has_chat_after_last_compaction, has_misplaced_tool_call, is_empty_response, queued_prompt, - sync_a2ui_message, + ChatMessage, TOOL_PROTOCOL_CORRECTION, TurnSummary, chat_turn, compacted_context_start, + correction_already_sent, has_chat_after_last_compaction, has_misplaced_tool_call, + is_empty_response, project_agents, promote_legacy_turn_summaries, queued_prompt, + sync_a2ui_message, title_context, }; + use crate::engine::ChatTurn; use crate::model::ModelChoice; fn assistant(reasoning: Option<&str>, content: &str) -> ChatMessage { @@ -1674,6 +1861,7 @@ mod tests { reasoning_complete: false, reasoning_open: false, content: content.to_owned(), + model_content: None, markdown: iced::widget::markdown::Content::new(), transcript: iced::widget::text_editor::Content::new(), a2ui_lines_processed: 0, @@ -1683,6 +1871,109 @@ mod tests { } } + fn turn(user: bool, system: bool, content: &str) -> ChatTurn { + ChatTurn { + user, + tool: false, + system, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: content.to_owned(), + } + } + + #[test] + fn project_agents_is_optional_and_preserves_the_file() { + let directory = + std::env::temp_dir().join(format!("ds4-server-agents-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + assert_eq!(project_agents(&directory).unwrap(), None); + + let instructions = "# Instructions\n\nKeep this exact.\n"; + std::fs::write(directory.join("AGENTS.md"), instructions).unwrap(); + assert_eq!( + project_agents(&directory).unwrap().as_deref(), + Some(instructions) + ); + std::fs::remove_dir_all(directory).unwrap(); + } + + #[test] + fn turn_summary_accumulates_model_continuations() { + let mut summary = TurnSummary::new(); + summary.record(100, 80, 20); + summary.record(140, 120, 30); + + let stats = summary.finish(); + assert_eq!(stats.input_tokens, 240); + assert_eq!(stats.cached_tokens, 200); + assert_eq!(stats.output_tokens, 50); + } + + #[test] + fn model_history_uses_persisted_content_without_changing_the_transcript() { + let mut user = assistant(None, "visible question"); + user.user = true; + user.model_content = Some("visible question\n\nhidden metadata".into()); + + assert_eq!( + chat_turn(&user).content, + "visible question\n\nhidden metadata" + ); + assert_eq!(user.content, "visible question"); + } + + #[test] + fn legacy_generation_rows_collapse_into_their_user_turn() { + let mut user = assistant(None, "question"); + user.user = true; + let mut first = assistant(None, "tool call"); + first.generation_stats = Some(super::GenerationStats { + duration_ms: 1_000, + input_tokens: 100, + cached_tokens: 80, + output_tokens: 20, + }); + let mut second = assistant(None, "answer"); + second.generation_stats = Some(super::GenerationStats { + duration_ms: 2_000, + input_tokens: 140, + cached_tokens: 120, + output_tokens: 30, + }); + let mut messages = [user, first, second]; + + promote_legacy_turn_summaries(&mut messages); + + assert_eq!( + messages[0].generation_stats, + Some(super::GenerationStats { + duration_ms: 3_000, + input_tokens: 240, + cached_tokens: 200, + output_tokens: 50, + }) + ); + assert!(messages[1].generation_stats.is_none()); + assert!(messages[2].generation_stats.is_none()); + } + + #[test] + fn title_context_starts_at_the_first_user_and_ignores_system_rows() { + let messages = title_context([ + turn(false, true, "project AGENTS.md"), + turn(false, true, "date and time"), + turn(true, false, "actual request"), + turn(false, true, "system reminder"), + turn(false, false, "answer"), + ]); + + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].content, "actual request"); + assert_eq!(messages[1].content, "answer"); + } + #[test] fn complete_tool_call_in_reasoning_gets_one_correction() { let call = r#"<｜DSML｜tool_calls> @@ -1736,6 +2027,7 @@ mod tests { reasoning_open: false, content: "### Core / Setup\n\n| File | Lines |\n|---|---:|\n| `src/app.rs` | **1,750** |\n| `src/engine.rs` | 2,400 |\n\n### Summary\n\nDone." .to_owned(), + model_content: None, markdown: iced::widget::markdown::Content::new(), transcript: iced::widget::text_editor::Content::new(), a2ui_lines_processed: 0, @@ -1802,6 +2094,7 @@ mod tests { reasoning_complete: true, reasoning_open: false, content: format!("message {id}"), + model_content: None, markdown: iced::widget::markdown::Content::new(), transcript: iced::widget::text_editor::Content::new(), a2ui_lines_processed: 0, @@ -1840,6 +2133,7 @@ mod tests { reasoning_complete: true, reasoning_open: false, content: format!("message {id}"), + model_content: None, markdown: iced::widget::markdown::Content::new(), transcript: iced::widget::text_editor::Content::new(), a2ui_lines_processed: 0, diff --git a/src/app/view/chat.rs b/src/app/view/chat.rs index 6dcd06f..b680a51 100644 --- a/src/app/view/chat.rs +++ b/src/app/view/chat.rs @@ -86,6 +86,11 @@ impl App { { continue; } + if message.user + && let Some(stats) = message.generation_stats + { + messages = messages.push(generation_summary(stats)); + } let label = if message.user { "You" } else if message.tool { @@ -193,9 +198,6 @@ impl App { } else { message_body }); - if let Some(stats) = message.generation_stats { - messages = messages.push(generation_summary(stats)); - } } if let Some(activity) = &self.activity { messages = messages.push( @@ -388,6 +390,7 @@ fn generation_summary(stats: crate::app::generation::GenerationStats) -> Element rule::horizontal(1), ] .spacing(5) + .padding(Padding::ZERO.right(24)) .into() } diff --git a/src/database.rs b/src/database.rs index 2b15d52..4e171a2 100644 --- a/src/database.rs +++ b/src/database.rs @@ -118,6 +118,7 @@ pub struct StoredMessage { pub reasoning: Option, pub reasoning_complete: bool, pub content: String, + pub model_content: Option, pub system: bool, pub compaction: bool, pub compaction_tail_start: Option, @@ -136,6 +137,7 @@ struct NewMessage<'a> { reasoning: Option<&'a str>, reasoning_complete: bool, content: &'a str, + model_content: Option<&'a str>, system: bool, compaction: bool, compaction_tail_start: Option, @@ -376,6 +378,7 @@ impl Database { reasoning: None, reasoning_complete: true, content: &content, + model_content: None, system: true, compaction: false, compaction_tail_start: None, @@ -422,6 +425,7 @@ impl Database { &mut self, session_id: i32, prompt: &str, + model_prompt: Option<&str>, system_messages: &[String], reasoning: bool, ) -> Result, String> { @@ -438,6 +442,7 @@ impl Database { reasoning: None, reasoning_complete: true, content, + model_content: None, system: true, compaction: false, compaction_tail_start: None, @@ -454,6 +459,7 @@ impl Database { reasoning: None, reasoning_complete: true, content: prompt, + model_content: model_prompt, system: false, compaction: false, compaction_tail_start: None, @@ -468,6 +474,7 @@ impl Database { reasoning: reasoning.then_some(""), reasoning_complete: !reasoning, content: "", + model_content: None, system: false, compaction: false, compaction_tail_start: None, @@ -500,6 +507,7 @@ impl Database { reasoning: None, reasoning_complete: true, content: result, + model_content: None, system: false, compaction: false, compaction_tail_start: None, @@ -517,6 +525,7 @@ impl Database { reasoning: None, reasoning_complete: true, content, + model_content: None, system: false, compaction: false, compaction_tail_start: None, @@ -535,6 +544,7 @@ impl Database { reasoning: None, reasoning_complete: true, content, + model_content: None, system: true, compaction: false, compaction_tail_start: None, @@ -551,6 +561,7 @@ impl Database { reasoning: reasoning.then_some(""), reasoning_complete: !reasoning, content: "", + model_content: None, system: false, compaction: false, compaction_tail_start: None, @@ -642,6 +653,7 @@ impl Database { reasoning: None, reasoning_complete: true, content: summary, + model_content: None, system: true, compaction: true, compaction_tail_start: tail_start, @@ -659,6 +671,7 @@ impl Database { reasoning: None, reasoning_complete: true, content, + model_content: None, system: false, compaction: false, compaction_tail_start: None, @@ -758,7 +771,7 @@ mod tests { .unwrap(); let session = database.create_session(project.id, "A2UI").unwrap(); let first = database - .start_chat_turn(session.id, "First", &[], false) + .start_chat_turn(session.id, "First", None, &[], false) .unwrap() .pop() .unwrap(); @@ -785,7 +798,7 @@ mod tests { assert!(active.active_surface().is_none()); let second = database - .start_chat_turn(session.id, "Second", &[], false) + .start_chat_turn(session.id, "Second", None, &[], false) .unwrap() .pop() .unwrap(); @@ -836,7 +849,13 @@ mod tests { let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap(); let session = database.create_session(project.id, "Chat").unwrap(); let mut opening = database - .start_chat_turn(session.id, "Question", &["Date context".into()], true) + .start_chat_turn( + session.id, + "Question", + Some("Question\n\nhidden metadata"), + &["Date context".into()], + true, + ) .unwrap(); let assistant = opening.pop().unwrap(); database @@ -875,6 +894,10 @@ mod tests { assert_eq!(messages.len(), 7); assert!(messages[0].system); assert_eq!(messages[1].content, "Question"); + assert_eq!( + messages[1].model_content.as_deref(), + Some("Question\n\nhidden metadata") + ); assert_eq!(messages[2].reasoning.as_deref(), Some("Reasoning")); assert!(messages[2].reasoning_complete); assert_eq!(messages[2].content, "Answer"); @@ -956,7 +979,7 @@ mod tests { ) .unwrap(); reopened - .start_chat_turn(session.id, "After third compaction", &[], false) + .start_chat_turn(session.id, "After third compaction", None, &[], false) .unwrap(); drop(reopened); diff --git a/src/engine.rs b/src/engine.rs index 40babf4..bb9abe2 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -560,6 +560,7 @@ impl Generator { pub(crate) fn generate( &mut self, checkpoint: &Path, + bootstrap: Option<&Path>, messages: &[ChatTurn], settings: &TurnSettings, cancelled: &AtomicBool, @@ -578,6 +579,11 @@ impl Generator { if checkpoint_present && !selected.found { phase(checkpoint_rebuild_activity(selected.incompatible)); } + if !selected.found + && let Some(directory) = bootstrap + { + self.prepare_bootstrap(directory, settings, cancelled, &mut progress, &mut phase)?; + } let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress); self.publish_execution_stats(); let (mut output, prompt_complete) = result?; @@ -596,6 +602,8 @@ impl Generator { output.checkpoint_bytes = std::fs::metadata(checkpoint) .map(|item| item.len()) .unwrap_or(0); + self.checkpoint = Some(checkpoint.to_owned()); + self.resident_active = Some(checkpoint.to_owned()); Ok(output) } @@ -648,32 +656,8 @@ impl Generator { let history_key = conversation_key(&settings.system_prompt, settings.reasoning_mode, history); let history_tag: [u8; 32] = Sha256::digest(&history_key).into(); - let mut previous_checkpoint = self.checkpoint.clone(); - if self.executor.checkpoint_tag() != history_tag { - if let Some(entry) = store.find(&history_key, self.executor.context()) { - if self.select_checkpoint(&entry.checkpoint, entry.tag)?.found { - store.touch(&entry)?; - self.last_store_tokens = entry.tokens; - previous_checkpoint = Some(entry.checkpoint); - } else { - store.discard(&entry); - self.last_store_tokens = 0; - previous_checkpoint = None; - } - } else { - let key = resident_key(directory, history_tag); - let restored = self.activate_resident(key)?; - if !restored || self.executor.checkpoint_tag() != history_tag { - self.executor.reset()?; - } - self.checkpoint = None; - self.last_store_tokens = 0; - previous_checkpoint = None; - self.metrics.kv_lookup(KvLookup::Miss); - } - } else { - self.metrics.kv_lookup(KvLookup::MemoryHit); - } + let previous_checkpoint = + self.restore_cached_prefix(directory, &store, &history_key, history_tag)?; let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress); self.publish_execution_stats(); @@ -731,6 +715,98 @@ impl Generator { Ok(output) } + fn prepare_bootstrap( + &mut self, + directory: &Path, + settings: &TurnSettings, + cancelled: &AtomicBool, + progress: &mut impl FnMut(u32, u32, Option), + phase: &mut impl FnMut(&'static str), + ) -> Result<(), String> { + let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &[]); + let tag: [u8; 32] = Sha256::digest(&key).into(); + let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?; + self.restore_cached_prefix(directory, &store, &key, tag)?; + + let tokens = self.executor.model().render_history( + &settings.system_prompt, + &[], + settings.reasoning_mode, + ); + if tokens.len() >= self.executor.context() as usize { + return Err(format!( + "System prompt has {} tokens, but the configured context size is {} tokens", + tokens.len(), + self.executor.context() + )); + } + let reused = self.executor.align_prompt(&tokens)?; + if reused == tokens.len() { + return Ok(()); + } + + phase("Updating system prompt cache…"); + let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?; + if completed != tokens.len() - reused { + return Err("generation cancelled while updating the system prompt cache".into()); + } + self.executor.note_checkpoint_tag(tag); + if !settings.kv_cache.stores(self.executor.position(), true, 0) { + self.checkpoint = None; + self.resident_active = Some(resident_key(directory, tag)); + return Ok(()); + } + + let checkpoint = store.checkpoint_path(&key); + self.save_checkpoint(&checkpoint, tag)?; + let retained = store.record( + &checkpoint, + &key, + tag, + self.executor.position(), + self.executor.context(), + StoreReason::Cold, + )?; + self.checkpoint = retained.then_some(checkpoint.clone()); + self.resident_active = Some(if retained { + checkpoint + } else { + resident_key(directory, tag) + }); + Ok(()) + } + + fn restore_cached_prefix( + &mut self, + directory: &Path, + store: &KvStore, + key: &[u8], + tag: [u8; 32], + ) -> Result, String> { + if self.executor.checkpoint_tag() == tag { + self.metrics.kv_lookup(KvLookup::MemoryHit); + return Ok(self.checkpoint.clone()); + } + if let Some(entry) = store.find(key, self.executor.context()) { + if self.select_checkpoint(&entry.checkpoint, entry.tag)?.found { + store.touch(&entry)?; + self.last_store_tokens = entry.tokens; + return Ok(Some(entry.checkpoint)); + } + store.discard(&entry); + } else { + let key = resident_key(directory, tag); + let restored = self.activate_resident(key)?; + if !restored || self.executor.checkpoint_tag() != tag { + self.executor.reset()?; + } + self.metrics.kv_lookup(KvLookup::Miss); + } + self.checkpoint = None; + self.last_store_tokens = 0; + Ok(None) + } + #[allow(clippy::too_many_arguments)] pub(crate) fn compact( &mut self, @@ -986,6 +1062,34 @@ impl Generator { result } + fn prefill_suffix( + &mut self, + tokens: &[i32], + reused: usize, + cancelled: &AtomicBool, + progress: &mut impl FnMut(u32, u32, Option), + ) -> Result { + let suffix = &tokens[reused..]; + if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 { + let context = self.executor.context(); + self.executor.prefill(suffix, |used| { + progress(used, context, None); + !cancelled.load(Ordering::Relaxed) + }) + } else { + let mut completed = 0; + for &token in suffix { + if cancelled.load(Ordering::Relaxed) { + break; + } + self.executor.eval(token)?; + completed += 1; + progress(self.executor.position(), self.executor.context(), None); + } + Ok(completed) + } + } + fn generate_inner( &mut self, messages: &[ChatTurn], @@ -1047,24 +1151,7 @@ impl Generator { let mut pending_utf8 = Vec::new(); let prompt_tokens = tokens.len(); let suffix = &tokens[reused..]; - let completed = if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 { - let context = self.executor.context(); - self.executor.prefill(suffix, |used| { - progress(used, context, None); - !cancelled.load(Ordering::Relaxed) - })? - } else { - let mut completed = 0; - for &token in suffix { - if cancelled.load(Ordering::Relaxed) { - break; - } - self.executor.eval(token)?; - completed += 1; - progress(self.executor.position(), self.executor.context(), None); - } - completed - }; + let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?; self.publish_execution_stats(); if completed != suffix.len() { return Ok(( @@ -1648,6 +1735,34 @@ mod sampling_tests { assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix)); } + #[test] + fn bootstrap_key_is_the_prefix_before_dynamic_session_context() { + let system = "System\n\nProject instructions from AGENTS.md:\n\nkeep this"; + let bootstrap = conversation_key(system, ReasoningMode::High, &[]); + let messages = vec![ + ChatTurn { + user: false, + tool: false, + system: true, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "current date and time".into(), + }, + ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "hello".into(), + }, + ]; + + assert!(conversation_key(system, ReasoningMode::High, &messages).starts_with(&bootstrap)); + } + #[test] fn checkpoint_rebuilds_explain_compatibility_and_history_misses() { assert!(checkpoint_rebuild_activity(true).contains("different model")); diff --git a/src/runtime.rs b/src/runtime.rs index 3cb0345..26648b4 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -19,7 +19,6 @@ pub(crate) struct GenerationService { pub(crate) struct ActiveGeneration { pub(crate) events: Receiver, pub(crate) cancel: Arc, - pub(crate) started_at: Instant, } impl Drop for ActiveGeneration { @@ -29,7 +28,11 @@ impl Drop for ActiveGeneration { } pub(crate) enum CheckpointTarget { - Local(PathBuf), + Local { + checkpoint: PathBuf, + /// Shared cache for the deterministic rendered system prompt. + bootstrap: Option, + }, Transient(PathBuf), /// Same transient KV handling as [`CheckpointTarget::Transient`], but asked /// for by the app itself (session titling) rather than by an HTTP client. @@ -39,7 +42,7 @@ pub(crate) enum CheckpointTarget { impl CheckpointTarget { fn source(&self) -> WorkSource { match self { - Self::Local(_) | Self::OneShot(_) => WorkSource::LocalChat, + Self::Local { .. } | Self::OneShot(_) => WorkSource::LocalChat, Self::Transient(_) => WorkSource::Http, } } @@ -108,7 +111,6 @@ impl GenerationService { checkpoint: CheckpointTarget, idle_timeout: Duration, ) -> Result { - let started_at = Instant::now(); let source = checkpoint.source(); let cancel = Arc::new(AtomicBool::new(false)); let (events, receiver) = mpsc::channel(); @@ -133,7 +135,6 @@ impl GenerationService { Ok(ActiveGeneration { events: receiver, cancel, - started_at, }) } @@ -148,7 +149,6 @@ impl GenerationService { checkpoint: PathBuf, idle_timeout: Duration, ) -> Result { - let started_at = Instant::now(); let cancel = Arc::new(AtomicBool::new(false)); let (events, receiver) = mpsc::channel(); self.commands @@ -156,7 +156,10 @@ impl GenerationService { engine, turn, messages, - checkpoint: CheckpointTarget::Local(checkpoint), + checkpoint: CheckpointTarget::Local { + checkpoint, + bootstrap: None, + }, operation: Operation::Compact { reason: reason.to_owned(), rebuild_system_prompt, @@ -169,7 +172,6 @@ impl GenerationService { Ok(ActiveGeneration { events: receiver, cancel, - started_at, }) } @@ -180,7 +182,6 @@ impl GenerationService { messages: Vec, idle_timeout: Duration, ) -> Result { - let started_at = Instant::now(); let cancel = Arc::new(AtomicBool::new(false)); let (events, receiver) = mpsc::channel(); self.metrics.request_queued(WorkSource::LocalChat); @@ -199,7 +200,6 @@ impl GenerationService { Ok(ActiveGeneration { events: receiver, cancel, - started_at, }) } } @@ -335,7 +335,7 @@ fn run_command( rebuild_system_prompt, } = &command.operation { - let CheckpointTarget::Local(checkpoint) = &command.checkpoint else { + let CheckpointTarget::Local { checkpoint, .. } = &command.checkpoint else { unreachable!("compaction checkpoints are local") }; let result = generator.compact( @@ -367,8 +367,12 @@ fn run_command( return; } let result = match command.checkpoint { - CheckpointTarget::Local(path) => generator.generate( - &path, + CheckpointTarget::Local { + checkpoint, + bootstrap, + } => generator.generate( + &checkpoint, + bootstrap.as_deref(), &command.messages, &command.turn, &command.cancel, @@ -440,7 +444,6 @@ mod tests { drop(ActiveGeneration { events, cancel: Arc::clone(&cancel), - started_at: Instant::now(), }); assert!(cancel.load(Ordering::Relaxed)); diff --git a/src/schema.rs b/src/schema.rs index bbf0e82..9eac92c 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -17,6 +17,7 @@ diesel::table! { reasoning -> Nullable, reasoning_complete -> Bool, content -> Text, + model_content -> Nullable, system -> Bool, compaction -> Bool, compaction_tail_start -> Nullable, diff --git a/src/server.rs b/src/server.rs index 7fcc8ca..9c39996 100644 --- a/src/server.rs +++ b/src/server.rs @@ -945,7 +945,6 @@ mod tests { let active = crate::runtime::ActiveGeneration { events: receiver, cancel: Arc::new(AtomicBool::new(false)), - started_at: std::time::Instant::now(), }; let producer = thread::spawn(move || { thread::sleep(Duration::from_millis(25));