hugo/DS4Server
1
0
Fork 0
Code Issues 8 Pull Requests Actions Packages Projects Releases Wiki Activity

Support long-running agent turns

Browse Source
This commit is contained in:
Georg Bauer
2026-07-26 11:07:28 +02:00
parent 2a14b93335
commit 3c75b8f6c1
14 changed files with 432 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
migrations/20260726200000_add_system_messages/down.sql Normal file
View File

@@ -0,0 +1 @@
ALTER TABLE messages DROP COLUMN system;

1
migrations/20260726200000_add_system_messages/up.sql Normal file
View File

@@ -0,0 +1 @@
ALTER TABLE messages ADD COLUMN system BOOLEAN NOT NULL DEFAULT 0 CHECK (system IN (0, 1));

39
src/agent.rs
View File

@@ -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<Option<String>, 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 = "done<tool_call>read<arg_key>path</arg_key><arg_value>src/main.rs</arg_value></tool_call>";
let (visible, calls) = parse_tool_calls(ModelChoice::Glm52, text).unwrap();

18
src/app.rs
View File

@@ -76,6 +76,7 @@ pub(crate) struct App {
project_name_input: String,
model_download: ModelDownload,
pub(super) composer: String,
pub(super) queued_inputs: VecDeque<String>,
pub(super) conversation: Vec<ChatMessage>,
pub(super) generating: bool,
pub(super) context_used: u32,
@@ -107,6 +108,8 @@ pub(crate) struct App {
_endpoint: Option<crate::server::ServerHandle>,
error: Option<String>,
pub(super) activity: Option<String>,
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,

145
src/app/generation.rs
View File

@@ -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<String>,
pub(super) reasoning_complete: bool,
pub(super) reasoning_open: bool,
@@ -72,6 +73,7 @@ impl From<StoredMessage> 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<StoredMessage> 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::<Vec<_>>();
#[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<crate::agent::ToolCall>) -> 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::<Vec<_>>();
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,

5
src/app/projects.rs
View File

@@ -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;
}

114
src/app/view/chat.rs
View File

@@ -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::<String>() + "…";
}
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);

1
src/compaction.rs
View File

@@ -99,6 +99,7 @@ mod tests {
ChatTurn {
user,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,

134
src/database.rs
View File

@@ -118,6 +118,7 @@ pub struct StoredMessage {
pub reasoning: Option<String>,
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<String>,
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<Vec<StoredMessage>, 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<Vec<StoredMessage>, 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,

15
src/engine.rs
View File

@@ -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<String>,
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,

8
src/engine/tokenizer.rs
View File

@@ -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);

37
src/engine/validation.rs
View File

@@ -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,

1
src/schema.rs
View File

@@ -7,6 +7,7 @@ diesel::table! {
reasoning -> Nullable<Text>,
reasoning_complete -> Bool,
content -> Text,
system -> Bool,
}
}

3
src/server/tools.rs
View File

@@ -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,