From 3c75b8f6c1da4433b11dcd1472e39a583750da3d Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 26 Jul 2026 11:07:28 +0200 Subject: [PATCH] Support long-running agent turns --- .../down.sql | 1 + .../20260726200000_add_system_messages/up.sql | 1 + src/agent.rs | 39 +++++ src/app.rs | 18 +++ src/app/generation.rs | 145 ++++++++++++++++-- src/app/projects.rs | 5 + src/app/view/chat.rs | 114 ++++++++------ src/compaction.rs | 1 + src/database.rs | 134 ++++++++++++---- src/engine.rs | 15 +- src/engine/tokenizer.rs | 8 + src/engine/validation.rs | 37 +++++ src/schema.rs | 1 + src/server/tools.rs | 3 + 14 files changed, 432 insertions(+), 90 deletions(-) create mode 100644 migrations/20260726200000_add_system_messages/down.sql create mode 100644 migrations/20260726200000_add_system_messages/up.sql diff --git a/migrations/20260726200000_add_system_messages/down.sql b/migrations/20260726200000_add_system_messages/down.sql new file mode 100644 index 0000000..bf84c01 --- /dev/null +++ b/migrations/20260726200000_add_system_messages/down.sql @@ -0,0 +1 @@ +ALTER TABLE messages DROP COLUMN system; diff --git a/migrations/20260726200000_add_system_messages/up.sql b/migrations/20260726200000_add_system_messages/up.sql new file mode 100644 index 0000000..3afdfcb --- /dev/null +++ b/migrations/20260726200000_add_system_messages/up.sql @@ -0,0 +1 @@ +ALTER TABLE messages ADD COLUMN system BOOLEAN NOT NULL DEFAULT 0 CHECK (system IN (0, 1)); diff --git a/src/agent.rs b/src/agent.rs index c64766e..b9b4c2d 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -988,6 +988,37 @@ pub(crate) fn system_prompt(model: ModelChoice, extra: &str) -> String { } } +pub(crate) fn system_prompt_reminder(model: ModelChoice, extra: &str) -> String { + format!( + "[System prompt reminder follows.]\n{}\n[End system prompt reminder.]", + system_prompt(model, extra) + ) +} + +pub(crate) fn prompt_reminder_due(used: u32, last: u32) -> bool { + used.saturating_sub(last) >= 50_000 +} + +pub(crate) fn datetime_context() -> String { + let when = Command::new("/bin/date") + .arg("+%Y-%m-%d %H:%M:%S %Z") + .output() + .ok() + .filter(|output| output.status.success()) + .map(|output| String::from_utf8_lossy(&output.stdout).trim().to_owned()) + .filter(|output| !output.is_empty()) + .unwrap_or_else(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + .to_string() + }); + format!( + "Current local date and time at session start: {when}. Use this only when date or time matters." + ) +} + pub(crate) fn try_tool_result(active: &ActiveTools) -> Result, String> { match active.results.try_recv() { Ok(result) => Ok(Some(result)), @@ -1224,6 +1255,14 @@ mod tests { assert!(prompt.contains(&format!("\"name\":\"{name}\""))); } assert!(prompt.ends_with("extra")); + assert!( + system_prompt_reminder(ModelChoice::DeepSeekV4Flash, "extra") + .contains("[System prompt reminder follows.]") + ); + assert!(datetime_context().starts_with("Current local date and time at session start:")); + assert!(!prompt_reminder_due(49_999, 0)); + assert!(prompt_reminder_due(50_000, 0)); + assert!(!prompt_reminder_due(80_000, 50_000)); let text = "donereadpathsrc/main.rs"; let (visible, calls) = parse_tool_calls(ModelChoice::Glm52, text).unwrap(); diff --git a/src/app.rs b/src/app.rs index ba2b6f8..f7e5c2a 100644 --- a/src/app.rs +++ b/src/app.rs @@ -76,6 +76,7 @@ pub(crate) struct App { project_name_input: String, model_download: ModelDownload, pub(super) composer: String, + pub(super) queued_inputs: VecDeque, pub(super) conversation: Vec, pub(super) generating: bool, pub(super) context_used: u32, @@ -107,6 +108,8 @@ pub(crate) struct App { _endpoint: Option, error: Option, pub(super) activity: Option, + stop_requested: bool, + system_prompt_seen_at: u32, #[cfg(target_os = "macos")] skip_compaction_once: bool, } @@ -285,6 +288,7 @@ impl App { project_name_input: String::new(), model_download: ModelDownload::Idle, composer: String::new(), + queued_inputs: VecDeque::new(), conversation: Vec::new(), generating: false, context_used: 0, @@ -324,6 +328,8 @@ impl App { } }, activity: None, + stop_requested: false, + system_prompt_seen_at: 0, #[cfg(target_os = "macos")] skip_compaction_once: false, } @@ -379,6 +385,7 @@ impl App { project_name_input: String::new(), model_download: ModelDownload::Idle, composer: String::new(), + queued_inputs: VecDeque::new(), conversation: Vec::new(), generating: false, context_used: 0, @@ -412,6 +419,8 @@ impl App { None => error, }), activity: None, + stop_requested: false, + system_prompt_seen_at: 0, #[cfg(target_os = "macos")] skip_compaction_once: false, } @@ -767,6 +776,8 @@ impl App { return scroll_chat_to_end(); } Message::StopGeneration => { + self.stop_requested = true; + self.activity = Some("Stopping…".into()); #[cfg(target_os = "macos")] if let Some(active) = &self.active_generation { active.cancel.store(true, Ordering::Relaxed); @@ -863,6 +874,8 @@ impl App { self.selected_session = None; self.conversation.clear(); self.composer.clear(); + self.queued_inputs.clear(); + self.system_prompt_seen_at = 0; self.context_used = 0; self.tokens_per_second = None; } @@ -1015,6 +1028,8 @@ impl App { self.composer.clear(); self.remember_project(project_id); self.selected_session = Some(session_id); + self.system_prompt_seen_at = 0; + self.queued_inputs.clear(); let (used, limit, tokens_per_second) = saved_context.unwrap_or_default(); self.context_used = used.max(0) as u32; self.context_limit = if limit > 0 { @@ -1056,6 +1071,8 @@ impl App { self.selected_session = None; self.conversation.clear(); self.composer.clear(); + self.queued_inputs.clear(); + self.system_prompt_seen_at = 0; self.context_used = 0; self.tokens_per_second = None; } @@ -1599,6 +1616,7 @@ mod tests { id: 1, user: false, tool: false, + system: false, reasoning: Some(String::new()), reasoning_complete: false, reasoning_open: true, diff --git a/src/app/generation.rs b/src/app/generation.rs index 28da7f5..e675e05 100644 --- a/src/app/generation.rs +++ b/src/app/generation.rs @@ -36,6 +36,7 @@ pub(crate) struct ChatMessage { pub(super) id: i32, pub(super) user: bool, pub(super) tool: bool, + pub(super) system: bool, pub(super) reasoning: Option, pub(super) reasoning_complete: bool, pub(super) reasoning_open: bool, @@ -72,6 +73,7 @@ impl From for ChatMessage { id: message.id, user: message.user, tool: message.tool, + system: message.system, reasoning: message.reasoning, reasoning_complete: message.reasoning_complete, reasoning_open: false, @@ -85,13 +87,27 @@ impl From for ChatMessage { impl App { pub(super) fn start_generation(&mut self) { - if self.generating || self.selected_project.is_none() { + if self.selected_project.is_none() { return; } let prompt = self.composer.trim().to_owned(); if prompt.is_empty() { return; } + if self.generating { + self.queued_inputs.push_back(prompt); + self.composer.clear(); + self.activity = Some(format!( + "{} queued input{}", + self.queued_inputs.len(), + if self.queued_inputs.len() == 1 { + "" + } else { + "s" + } + )); + return; + } #[cfg(target_os = "macos")] if prompt == "/compact" { self.composer.clear(); @@ -136,12 +152,30 @@ impl App { ); let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct; #[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.push(crate::agent::datetime_context()); + } + #[cfg(target_os = "macos")] + let reminder_injected = self.system_prompt_reminder_due(); + #[cfg(target_os = "macos")] + if reminder_injected { + injected_system.push(crate::agent::system_prompt_reminder( + model, + &self.config.generation.system_prompt, + )); + } + #[cfg(target_os = "macos")] let mut messages = self .conversation .iter() .map(|message| ChatTurn { user: message.user, tool: message.tool, + system: message.system, skip_previous_eos: false, reasoning: message.reasoning.clone(), reasoning_complete: message.reasoning_complete, @@ -149,9 +183,20 @@ impl App { }) .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(), + })); + #[cfg(target_os = "macos")] messages.push(ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -161,7 +206,6 @@ impl App { #[cfg(target_os = "macos")] { // A draft session only reaches the database once there is a turn to store. - let opening_turn = self.selected_session.is_none(); let session_id = match self.selected_session { Some(session_id) => session_id, None => { @@ -184,13 +228,21 @@ impl App { let Some(database) = &mut self.database else { return; }; - let saved = match database.start_chat_turn(session_id, &prompt, assistant_reasoning) { + let mut saved = match database.start_chat_turn( + session_id, + &prompt, + &injected_system, + assistant_reasoning, + ) { Ok(turn) => turn, Err(error) => { self.error = Some(format!("Could not save the chat turn: {error}")); return; } }; + 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); self.active_generation = match service.generate( @@ -207,13 +259,18 @@ impl App { return; } }; - let user = ChatMessage::from(saved.0); - let mut assistant = ChatMessage::from(saved.1); + let mut assistant = ChatMessage::from(saved.pop().unwrap()); + let user = ChatMessage::from(saved.pop().unwrap()); + for message in saved { + self.conversation.push(ChatMessage::from(message)); + } assistant.reasoning_open = assistant_reasoning; self.composer.clear(); self.conversation.push(user); self.conversation.push(assistant); self.generating = true; + self.stop_requested = false; + self.activity = Some("Loading model…".into()); self.tokens_per_second = None; self.error = None; if opening_turn { @@ -229,6 +286,10 @@ impl App { } } + fn system_prompt_reminder_due(&self) -> bool { + crate::agent::prompt_reminder_due(self.context_used, self.system_prompt_seen_at) + } + pub(super) fn poll_generation(&mut self) -> bool { #[cfg(target_os = "macos")] if self.active_compaction.is_some() { @@ -242,6 +303,7 @@ impl App { self.active_tools = None; if cancelled { self.generating = false; + self.activity = Some("Stopped".into()); return false; } if let Err(error) = self.continue_after_tool_result(&result) { @@ -254,6 +316,7 @@ impl App { Err(error) => { self.active_tools = None; self.generating = false; + self.activity = Some("Failed".into()); self.error = Some(error); return false; } @@ -269,6 +332,8 @@ impl App { #[cfg(target_os = "macos")] let mut context_changed = false; #[cfg(target_os = "macos")] + let mut start_queued = false; + #[cfg(target_os = "macos")] loop { match active.events.try_recv() { Ok(GenerationEvent::Loading) => {} @@ -297,6 +362,10 @@ impl App { } Ok(GenerationEvent::Finished(result)) => { match result { + Ok(_) if self.stop_requested => { + self.generating = false; + self.activity = Some("Stopped".into()); + } Ok(_) => { let model = self.config.model; let content = self @@ -311,14 +380,23 @@ impl App { self.error = Some(error); } } - Ok(_) => self.generating = false, + Ok(_) => { + self.generating = false; + self.activity = None; + start_queued = !self.queued_inputs.is_empty(); + } Err(error) => { self.active_tools = Some(crate::agent::error_async(error)); } } } + Err(_) if self.stop_requested => { + self.generating = false; + self.activity = Some("Stopped".into()); + } Err(error) => { self.generating = false; + self.activity = Some("Failed".into()); self.error = Some(error); } } @@ -330,6 +408,7 @@ impl App { self.generating = false; self.active_generation = None; self.error = Some("The model runtime stopped unexpectedly.".into()); + self.activity = Some("Failed".into()); break; } } @@ -378,11 +457,22 @@ impl App { } } #[cfg(target_os = "macos")] + if start_queued { + self.start_next_queued(); + } + #[cfg(target_os = "macos")] return transcript_changed; #[cfg(not(target_os = "macos"))] false } + fn start_next_queued(&mut self) { + if let Some(prompt) = self.queued_inputs.pop_front() { + self.composer = prompt; + self.start_generation(); + } + } + #[cfg(target_os = "macos")] fn start_agent_tools(&mut self, calls: Vec) -> Result<(), String> { let session_id = self @@ -437,26 +527,43 @@ impl App { self.compaction_summary(), ); let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct; - let saved = self + let queued = self.queued_inputs.iter().cloned().collect::>(); + let reminder = self.system_prompt_reminder_due().then(|| { + crate::agent::system_prompt_reminder(model, &self.config.generation.system_prompt) + }); + let mut saved = self .database .as_mut() .ok_or_else(|| "The project database is unavailable.".to_owned())? - .continue_tool_turn(session_id, result, assistant_reasoning) + .continue_tool_turn( + session_id, + result, + &queued, + reminder.as_deref(), + assistant_reasoning, + ) .map_err(|error| format!("Could not save the tool turn: {error}"))?; - self.conversation.push(ChatMessage::from(saved.0)); + self.queued_inputs.clear(); + if reminder.is_some() { + self.system_prompt_seen_at = self.context_used; + } + let mut assistant = ChatMessage::from(saved.pop().unwrap()); + for message in saved { + self.conversation.push(ChatMessage::from(message)); + } let messages = self .conversation .iter() .map(|message| ChatTurn { user: message.user, tool: message.tool, + system: message.system, skip_previous_eos: false, reasoning: message.reasoning.clone(), reasoning_complete: message.reasoning_complete, content: message.content.clone(), }) .collect(); - let mut assistant = ChatMessage::from(saved.1); assistant.reasoning_open = assistant_reasoning; self.conversation.push(assistant); let idle_timeout = Duration::from_secs(self.config.idle_timeout_minutes.max(1) as u64 * 60); @@ -473,6 +580,8 @@ impl App { )?, ); self.tokens_per_second = None; + self.activity = Some("Continuing after tools…".into()); + self.stop_requested = false; Ok(()) } @@ -501,6 +610,7 @@ impl App { .map(|message| ChatTurn { user: message.user, tool: message.tool, + system: message.system, skip_previous_eos: false, reasoning: message.reasoning.clone(), reasoning_complete: message.reasoning_complete, @@ -559,10 +669,11 @@ impl App { return false; } self.context_used = compacted.context_tokens; + self.system_prompt_seen_at = compacted.context_tokens; self.generating = false; self.activity = None; match request.pending { - PendingContinuation::None => {} + PendingContinuation::None => self.start_next_queued(), PendingContinuation::User(prompt) => { self.composer = prompt; self.skip_compaction_once = true; @@ -585,8 +696,12 @@ impl App { } Err(error) => { self.generating = false; - self.activity = None; - self.error = Some(error); + if self.stop_requested { + self.activity = Some("Stopped".into()); + } else { + self.activity = Some("Failed".into()); + self.error = Some(error); + } return false; } } @@ -618,6 +733,7 @@ impl App { .map(|message| crate::database::MessageDraft { user: message.user, tool: message.tool, + system: message.system, reasoning: message.reasoning.clone(), reasoning_complete: message.reasoning_complete, content: message.content.clone(), @@ -633,6 +749,7 @@ impl App { tail.push(crate::database::MessageDraft { user: false, tool: true, + system: false, reasoning: None, reasoning_complete: true, content: observation, @@ -713,6 +830,7 @@ impl App { .map(|message| ChatTurn { user: message.user, tool: message.tool, + system: message.system, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -725,6 +843,7 @@ impl App { messages.push(ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, diff --git a/src/app/projects.rs b/src/app/projects.rs index 5b55872..3a377bd 100644 --- a/src/app/projects.rs +++ b/src/app/projects.rs @@ -49,6 +49,7 @@ impl App { Ok(project) => { self.remember_project(project.id); self.selected_session = None; + self.system_prompt_seen_at = 0; self.pending_project_path = None; self.project_name_input.clear(); self.error = None; @@ -74,6 +75,8 @@ impl App { self.selected_session = None; self.conversation.clear(); self.composer.clear(); + self.queued_inputs.clear(); + self.system_prompt_seen_at = 0; self.context_used = 0; self.context_limit = self.config.generation.context_tokens.max(0) as u32; self.tokens_per_second = None; @@ -84,6 +87,8 @@ impl App { if self.drafts.remove(&project_id).is_some() && self.draft_selected(project_id) { self.conversation.clear(); self.composer.clear(); + self.queued_inputs.clear(); + self.system_prompt_seen_at = 0; self.context_used = 0; self.tokens_per_second = None; } diff --git a/src/app/view/chat.rs b/src/app/view/chat.rs index abf7fba..e447874 100644 --- a/src/app/view/chat.rs +++ b/src/app/view/chat.rs @@ -71,6 +71,9 @@ impl App { ); } for (index, message) in self.conversation.iter().enumerate() { + if message.system { + continue; + } let label = if message.user { "You" } else if message.tool { @@ -165,13 +168,20 @@ impl App { .width(Length::Fill) .into() } else { - text_input("Ask DS4Server anything…", &self.composer) - .id(composer_id()) - .on_input(Message::ComposerChanged) - .on_submit(Message::SubmitPrompt) - .padding(12) - .size(14) - .into() + text_input( + if self.generating { + "Add guidance to the queue…" + } else { + "Ask DS4Server anything…" + }, + &self.composer, + ) + .id(composer_id()) + .on_input(Message::ComposerChanged) + .on_submit(Message::SubmitPrompt) + .padding(12) + .size(14) + .into() }; let action = if self.generating { action_button(text("Stop").size(12)).on_press(Message::StopGeneration) @@ -187,50 +197,60 @@ impl App { } else { self.context_used.min(self.context_limit) as f32 / self.context_limit as f32 }; + let mut composer_content = column![composer].spacing(6); + for queued in &self.queued_inputs { + let mut queued = queued.replace('\n', " "); + if queued.chars().count() > 120 { + queued = queued.chars().take(119).collect::() + "…"; + } + composer_content = composer_content.push( + text(format!("Queued · {queued}")) + .size(12) + .color(muted_text()), + ); + } + composer_content = + composer_content.push( + row![ + icon(ICON_PAPERCLIP, 19), + tooltip( + context_pie(context_fraction, 19), + container( + text(format!( + "{} / {} tokens ({:.0}%)", + self.context_used, + self.context_limit, + context_fraction * 100.0 + )) + .size(12) + ) + .padding(10) + .style(preference_group_style), + tooltip::Position::Top, + ) + .gap(6), + text(self.tokens_per_second.map_or_else( + || "— tok/s".to_owned(), + |speed| format!("{speed:.1} tok/s") + )) + .size(11) + .color(muted_text()), + Space::with_width(Length::Fill), + icon(ICON_MODEL, 16), + text(self.config.model.to_string()).size(12), + action, + ] + .spacing(6) + .align_y(Alignment::Center), + ); let conversation = column![ scrollable(messages) .id(chat_scroll_id()) .height(Length::Fill), - container( - column![ - composer, - row![ - icon(ICON_PAPERCLIP, 19), - tooltip( - context_pie(context_fraction, 19), - container( - text(format!( - "{} / {} tokens ({:.0}%)", - self.context_used, - self.context_limit, - context_fraction * 100.0 - )) - .size(12) - ) - .padding(10) - .style(preference_group_style), - tooltip::Position::Top, - ) - .gap(6), - text(self.tokens_per_second.map_or_else( - || "— tok/s".to_owned(), - |speed| format!("{speed:.1} tok/s") - )) - .size(11) - .color(muted_text()), - Space::with_width(Length::Fill), - icon(ICON_MODEL, 16), - text(self.config.model.to_string()).size(12), - action, - ] - .spacing(6) - .align_y(Alignment::Center), - ] - .spacing(8), - ) - .padding(16) - .width(Length::Fill) - .style(overview_style), + container(composer_content,) + .padding(16) + .width(Length::Fill) + .style(overview_style), ] .height(Length::Fill) .spacing(8); diff --git a/src/compaction.rs b/src/compaction.rs index 3713bbb..d0028da 100644 --- a/src/compaction.rs +++ b/src/compaction.rs @@ -99,6 +99,7 @@ mod tests { ChatTurn { user, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, diff --git a/src/database.rs b/src/database.rs index d161404..978a7ae 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 system: bool, } #[derive(Insertable)] @@ -129,6 +130,7 @@ struct NewMessage<'a> { reasoning: Option<&'a str>, reasoning_complete: bool, content: &'a str, + system: bool, } pub struct MessageDraft { @@ -137,6 +139,7 @@ pub struct MessageDraft { pub reasoning: Option, pub reasoning_complete: bool, pub content: String, + pub system: bool, } #[derive(Debug)] @@ -318,10 +321,28 @@ impl Database { &mut self, session_id: i32, prompt: &str, + system_messages: &[String], reasoning: bool, - ) -> Result<(StoredMessage, StoredMessage), String> { + ) -> Result, String> { self.connection .transaction(|connection| { + let mut stored = Vec::with_capacity(system_messages.len() + 2); + for content in system_messages { + stored.push( + diesel::insert_into(messages::table) + .values(NewMessage { + session_id, + user: false, + tool: false, + reasoning: None, + reasoning_complete: true, + content, + system: true, + }) + .returning(StoredMessage::as_returning()) + .get_result(connection)?, + ); + } let user = diesel::insert_into(messages::table) .values(NewMessage { session_id, @@ -330,6 +351,7 @@ impl Database { reasoning: None, reasoning_complete: true, content: prompt, + system: false, }) .returning(StoredMessage::as_returning()) .get_result(connection)?; @@ -341,10 +363,12 @@ impl Database { reasoning: reasoning.then_some(""), reasoning_complete: !reasoning, content: "", + system: false, }) .returning(StoredMessage::as_returning()) .get_result(connection)?; - Ok((user, assistant)) + stored.extend([user, assistant]); + Ok(stored) }) .map_err(|error: diesel::result::Error| error.to_string()) } @@ -353,21 +377,59 @@ impl Database { &mut self, session_id: i32, result: &str, + queued_users: &[String], + reminder: Option<&str>, reasoning: bool, - ) -> Result<(StoredMessage, StoredMessage), String> { + ) -> Result, String> { self.connection .transaction(|connection| { - let tool = diesel::insert_into(messages::table) - .values(NewMessage { - session_id, - user: false, - tool: true, - reasoning: None, - reasoning_complete: true, - content: result, - }) - .returning(StoredMessage::as_returning()) - .get_result(connection)?; + let mut stored = Vec::with_capacity(queued_users.len() + 3); + stored.push( + diesel::insert_into(messages::table) + .values(NewMessage { + session_id, + user: false, + tool: true, + reasoning: None, + reasoning_complete: true, + content: result, + system: false, + }) + .returning(StoredMessage::as_returning()) + .get_result(connection)?, + ); + for content in queued_users { + stored.push( + diesel::insert_into(messages::table) + .values(NewMessage { + session_id, + user: true, + tool: false, + reasoning: None, + reasoning_complete: true, + content, + system: false, + }) + .returning(StoredMessage::as_returning()) + .get_result(connection)?, + ); + } + if let Some(content) = reminder { + stored.push( + diesel::insert_into(messages::table) + .values(NewMessage { + session_id, + user: false, + tool: false, + reasoning: None, + reasoning_complete: true, + content, + system: true, + }) + .returning(StoredMessage::as_returning()) + .get_result(connection)?, + ); + } let assistant = diesel::insert_into(messages::table) .values(NewMessage { session_id, @@ -376,10 +438,12 @@ impl Database { reasoning: reasoning.then_some(""), reasoning_complete: !reasoning, content: "", + system: false, }) .returning(StoredMessage::as_returning()) .get_result(connection)?; - Ok((tool, assistant)) + stored.push(assistant); + Ok(stored) }) .map_err(|error: diesel::result::Error| error.to_string()) } @@ -437,6 +501,7 @@ impl Database { reasoning: message.reasoning.as_deref(), reasoning_complete: message.reasoning_complete, content: &message.content, + system: message.system, }) .returning(StoredMessage::as_returning()) .get_result(connection)?, @@ -530,14 +595,21 @@ mod tests { let mut database = Database::open(&path).unwrap(); let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap(); let session = database.create_session(project.id, "Chat").unwrap(); - let (_, assistant) = database - .start_chat_turn(session.id, "Question", true) + let mut opening = database + .start_chat_turn(session.id, "Question", &["Date context".into()], true) .unwrap(); + let assistant = opening.pop().unwrap(); database .update_message(assistant.id, Some("Reasoning"), true, "Answer") .unwrap(); database - .continue_tool_turn(session.id, "Tool result", false) + .continue_tool_turn( + session.id, + "Tool result", + &["Queued correction".into()], + Some("Tool reminder"), + false, + ) .unwrap(); database .update_session_context(session.id, 1_234, 65_536, Some(12.5)) @@ -550,17 +622,20 @@ mod tests { assert_eq!(projects[0].sessions[0].context_limit, 65_536); assert_eq!(projects[0].sessions[0].last_tokens_per_second, Some(12.5)); let messages = reopened.load_messages(session.id).unwrap(); - assert_eq!(messages.len(), 4); - assert!(messages[0].user); - assert!(!messages[0].tool); - assert_eq!(messages[0].content, "Question"); - assert_eq!(messages[1].reasoning.as_deref(), Some("Reasoning")); - assert!(messages[1].reasoning_complete); - assert_eq!(messages[1].content, "Answer"); - assert!(messages[2].tool); - assert_eq!(messages[2].content, "Tool result"); - assert!(!messages[3].user); - assert!(!messages[3].tool); + assert_eq!(messages.len(), 7); + assert!(messages[0].system); + assert_eq!(messages[1].content, "Question"); + assert_eq!(messages[2].reasoning.as_deref(), Some("Reasoning")); + assert!(messages[2].reasoning_complete); + assert_eq!(messages[2].content, "Answer"); + assert!(messages[3].tool); + assert_eq!(messages[3].content, "Tool result"); + assert!(messages[4].user); + assert_eq!(messages[4].content, "Queued correction"); + assert!(messages[5].system); + assert_eq!(messages[5].content, "Tool reminder"); + assert!(!messages[6].user); + assert!(!messages[6].tool); let compacted = reopened .replace_with_compacted_transcript( session.id, @@ -571,6 +646,7 @@ mod tests { reasoning: None, reasoning_complete: true, content: "Recent question".into(), + system: false, }], 321, 32_768, diff --git a/src/engine.rs b/src/engine.rs index abd16a2..d8f0cc2 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -344,6 +344,7 @@ pub(crate) struct Generator { pub(crate) struct ChatTurn { pub(crate) user: bool, pub(crate) tool: bool, + pub(crate) system: bool, pub(crate) skip_previous_eos: bool, pub(crate) reasoning: Option, pub(crate) reasoning_complete: bool, @@ -540,6 +541,7 @@ impl Generator { private_messages.push(ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -743,6 +745,7 @@ impl Generator { let mut generated = ChatTurn { user: false, tool: false, + system: false, skip_previous_eos: false, reasoning: reasoning.then(String::new), reasoning_complete: !reasoning, @@ -1068,7 +1071,7 @@ fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn output.extend_from_slice(value.as_bytes()); } - let mut output = b"DS4Server chat checkpoint v3".to_vec(); + let mut output = b"DS4Server chat checkpoint v4".to_vec(); text(&mut output, system); output.push(match reasoning { ReasoningMode::Direct => 0, @@ -1078,6 +1081,7 @@ fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn for message in messages { output.push(u8::from(message.user)); output.push(u8::from(message.tool)); + output.push(u8::from(message.system)); output.push(u8::from(message.skip_previous_eos)); match &message.reasoning { Some(reasoning) => { @@ -1231,6 +1235,7 @@ mod sampling_tests { let mut generated = ChatTurn { user: false, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -1251,6 +1256,7 @@ mod sampling_tests { let mut messages = vec![ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -1269,6 +1275,12 @@ mod sampling_tests { tag, conversation_tag("System", ReasoningMode::Direct, &messages) ); + messages[0].system = true; + assert_ne!( + tag, + conversation_tag("System", ReasoningMode::High, &messages) + ); + messages[0].system = false; messages[0].skip_previous_eos = true; assert_ne!( tag, @@ -1278,6 +1290,7 @@ mod sampling_tests { messages.push(ChatTurn { user: false, tool: false, + system: false, skip_previous_eos: false, reasoning: Some("because".into()), reasoning_complete: true, diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index 9fe0489..a17efa3 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -196,6 +196,7 @@ impl Tokenizer { &[ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -255,6 +256,13 @@ impl Tokenizer { output.extend(self.tokenize_rendered(system_prompt)); } for (index, message) in messages.iter().enumerate() { + if message.system { + if self.family == ModelFamily::Glm { + output.push(self.system); + } + output.extend(self.tokenize_rendered(&message.content)); + continue; + } if message.tool { if self.family == ModelFamily::Glm { output.push(self.observation); diff --git a/src/engine/validation.rs b/src/engine/validation.rs index 5dd8d33..484b6f0 100644 --- a/src/engine/validation.rs +++ b/src/engine/validation.rs @@ -905,6 +905,7 @@ mod tests { ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -913,6 +914,7 @@ mod tests { ChatTurn { user: false, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -921,6 +923,7 @@ mod tests { ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -941,6 +944,33 @@ mod tests { model.render_continuation("Hello", ReasoningMode::Direct, true), [128_803, 19_923, 128_804, 128_822] ); + assert_eq!( + model.render_conversation( + "", + &[ + ChatTurn { + user: false, + tool: false, + system: true, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "Hello".into(), + }, + ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: "Hello".into(), + }, + ], + ReasoningMode::Direct, + ), + [0, 19_923, 128_803, 19_923, 128_804, 128_822] + ); assert_eq!( model.render_history( "", @@ -948,6 +978,7 @@ mod tests { ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -956,6 +987,7 @@ mod tests { ChatTurn { user: false, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -969,6 +1001,7 @@ mod tests { let tool_turn = ChatTurn { user: false, tool: true, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -977,6 +1010,7 @@ mod tests { let wrapped_user = ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -993,6 +1027,7 @@ mod tests { ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -1001,6 +1036,7 @@ mod tests { ChatTurn { user: false, tool: false, + system: false, skip_previous_eos: false, reasoning: Some("Hello".into()), reasoning_complete: true, @@ -1009,6 +1045,7 @@ mod tests { ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, diff --git a/src/schema.rs b/src/schema.rs index 731a280..4950ea3 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -7,6 +7,7 @@ diesel::table! { reasoning -> Nullable, reasoning_complete -> Bool, content -> Text, + system -> Bool, } } diff --git a/src/server/tools.rs b/src/server/tools.rs index 035ebd6..199bfa8 100644 --- a/src/server/tools.rs +++ b/src/server/tools.rs @@ -375,6 +375,7 @@ pub(super) fn render_messages( "user" => turns.push(ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: false, reasoning: None, reasoning_complete: true, @@ -394,6 +395,7 @@ pub(super) fn render_messages( turns.push(ChatTurn { user: true, tool: false, + system: false, skip_previous_eos: protocol == Protocol::Responses, reasoning: None, reasoning_complete: true, @@ -410,6 +412,7 @@ pub(super) fn render_messages( turns.push(ChatTurn { user: false, tool: false, + system: false, skip_previous_eos: false, reasoning: (preserve_reasoning && !reasoning.is_empty()).then_some(reasoning), reasoning_complete: true,