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

Add persistent local agent tools

Browse Source
This commit is contained in:
Georg Bauer
2026-07-25 14:57:17 +02:00
parent 668d8b787e
commit 4a45735072
20 changed files with 3139 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1309
src/agent.rs Normal file
View File

File diff suppressed because it is too large Load Diff

26
src/app.rs
View File

@@ -32,7 +32,7 @@ use std::fs;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc::{self, TryRecvError};
use std::sync::{Arc, RwLock};
use std::sync::{Arc, Mutex, RwLock};
use std::thread;
use std::time::{Duration, Instant};
@@ -84,6 +84,10 @@ pub(crate) struct App {
#[cfg(target_os = "macos")]
active_generation: Option<ActiveGeneration>,
#[cfg(target_os = "macos")]
agent_tools: Option<(i32, Arc<Mutex<crate::agent::Tools>>)>,
#[cfg(target_os = "macos")]
active_tools: Option<crate::agent::ActiveTools>,
#[cfg(target_os = "macos")]
active_titling: Option<generation::TitleRequest>,
#[cfg(target_os = "macos")]
runtime_preferences: Arc<RwLock<AppPreferences>>,
@@ -249,6 +253,10 @@ impl App {
#[cfg(target_os = "macos")]
active_generation: None,
#[cfg(target_os = "macos")]
agent_tools: None,
#[cfg(target_os = "macos")]
active_tools: None,
#[cfg(target_os = "macos")]
active_titling: None,
#[cfg(target_os = "macos")]
runtime_preferences,
@@ -326,6 +334,10 @@ impl App {
#[cfg(target_os = "macos")]
active_generation: None,
#[cfg(target_os = "macos")]
agent_tools: None,
#[cfg(target_os = "macos")]
active_tools: None,
#[cfg(target_os = "macos")]
active_titling: None,
#[cfg(target_os = "macos")]
runtime_preferences,
@@ -643,12 +655,15 @@ impl App {
self.start_generation();
return scroll_chat_to_end();
}
Message::StopGeneration =>
{
Message::StopGeneration => {
#[cfg(target_os = "macos")]
if let Some(active) = &self.active_generation {
active.cancel.store(true, Ordering::Relaxed);
}
#[cfg(target_os = "macos")]
if let Some(active) = &self.active_tools {
active.cancel.store(true, Ordering::Relaxed);
}
}
Message::GenerationTick => {
#[cfg(target_os = "macos")]
@@ -1007,6 +1022,10 @@ impl Drop for App {
if let Some(active) = &self.active_generation {
active.cancel.store(true, Ordering::Relaxed);
}
#[cfg(target_os = "macos")]
if let Some(active) = &self.active_tools {
active.cancel.store(true, Ordering::Relaxed);
}
}
}
@@ -1130,6 +1149,7 @@ mod tests {
let mut message = ChatMessage {
id: 1,
user: false,
tool: false,
reasoning: Some(String::new()),
reasoning_complete: false,
reasoning_open: true,

148
src/app/generation.rs
View File

@@ -22,6 +22,7 @@ pub(super) struct TitleRequest {
pub(crate) struct ChatMessage {
pub(super) id: i32,
pub(super) user: bool,
pub(super) tool: bool,
pub(super) reasoning: Option<String>,
pub(super) reasoning_complete: bool,
pub(super) reasoning_open: bool,
@@ -40,11 +41,12 @@ impl ChatMessage {
}
pub(super) fn refresh_markdown(&mut self) {
if !self.user {
if !self.user && !self.tool {
let visible = crate::agent::visible_content(&self.content);
let content = if self.reasoning.is_some() {
self.content.trim_start()
visible.trim_start()
} else {
&self.content
visible
};
self.markdown = markdown::parse(content).collect();
}
@@ -56,6 +58,7 @@ impl From<StoredMessage> for ChatMessage {
let mut message = Self {
id: message.id,
user: message.user,
tool: message.tool,
reasoning: message.reasoning,
reasoning_complete: message.reasoning_complete,
reasoning_open: false,
@@ -88,13 +91,15 @@ impl App {
crate::settings::effective_settings(model, &generation, &runtime, &models_path())
})
});
let effective = match effective {
let mut effective = match effective {
Ok(settings) => settings,
Err(error) => {
self.error = Some(error);
return;
}
};
effective.turn.system_prompt =
crate::agent::system_prompt(model, &effective.turn.system_prompt);
let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct;
#[cfg(target_os = "macos")]
let mut messages = self
@@ -102,6 +107,7 @@ impl App {
.iter()
.map(|message| ChatTurn {
user: message.user,
tool: message.tool,
skip_previous_eos: false,
reasoning: message.reasoning.clone(),
reasoning_complete: message.reasoning_complete,
@@ -111,6 +117,7 @@ impl App {
#[cfg(target_os = "macos")]
messages.push(ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -189,6 +196,31 @@ impl App {
}
pub(super) fn poll_generation(&mut self) -> bool {
#[cfg(target_os = "macos")]
if let Some(active) = &self.active_tools {
match crate::agent::try_tool_result(active) {
Ok(Some(result)) => {
let cancelled = active.cancel.load(Ordering::Relaxed);
self.active_tools = None;
if cancelled {
self.generating = false;
return false;
}
if let Err(error) = self.continue_after_tool_result(&result) {
self.generating = false;
self.error = Some(error);
}
return true;
}
Ok(None) => return false,
Err(error) => {
self.active_tools = None;
self.generating = false;
self.error = Some(error);
return false;
}
}
}
#[cfg(target_os = "macos")]
let Some(active) = &mut self.active_generation else {
self.generating = false;
@@ -221,9 +253,32 @@ impl App {
context_changed = true;
}
Ok(GenerationEvent::Finished(result)) => {
self.generating = false;
if let Err(error) = result {
self.error = Some(error);
match result {
Ok(_) => {
let model = ModelChoice::from_id(&self.preferences.selected_model)
.unwrap_or_default();
let content = self
.conversation
.last()
.map(|message| message.content.clone())
.unwrap_or_default();
match crate::agent::parse_tool_calls(model, &content) {
Ok((_, calls)) if !calls.is_empty() => {
if let Err(error) = self.start_agent_tools(calls) {
self.generating = false;
self.error = Some(error);
}
}
Ok(_) => self.generating = false,
Err(error) => {
self.active_tools = Some(crate::agent::error_async(error));
}
}
}
Err(error) => {
self.generating = false;
self.error = Some(error);
}
}
self.active_generation = None;
break;
@@ -286,6 +341,83 @@ impl App {
false
}
#[cfg(target_os = "macos")]
fn start_agent_tools(&mut self, calls: Vec<crate::agent::ToolCall>) -> Result<(), String> {
let session_id = self
.selected_session
.ok_or_else(|| "The active session is unavailable.".to_owned())?;
if self.agent_tools.as_ref().map(|(id, _)| *id) != Some(session_id) {
let project_id = self
.selected_project
.ok_or_else(|| "The active project is unavailable.".to_owned())?;
let root = self
.projects
.iter()
.find(|project| project.project.id == project_id)
.map(|project| PathBuf::from(&project.project.path))
.ok_or_else(|| "The active project is unavailable.".to_owned())?;
let tools = crate::agent::Tools::new(&root, self.preferences.context_tokens)?;
self.agent_tools = Some((session_id, Arc::new(Mutex::new(tools))));
}
let tools = Arc::clone(&self.agent_tools.as_ref().unwrap().1);
self.active_tools = Some(crate::agent::execute_async(tools, calls));
Ok(())
}
#[cfg(target_os = "macos")]
fn continue_after_tool_result(&mut self, result: &str) -> Result<(), String> {
let session_id = self
.selected_session
.ok_or_else(|| "The active session is unavailable.".to_owned())?;
let model = ModelChoice::from_id(&self.preferences.selected_model)
.ok_or_else(|| "The selected model is not supported.".to_owned())?;
let generation = self.preferences.generation()?;
let runtime = self.preferences.runtime()?;
let mut effective =
crate::settings::effective_settings(model, &generation, &runtime, &models_path())?;
effective.turn.system_prompt =
crate::agent::system_prompt(model, &effective.turn.system_prompt);
let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct;
let saved = self
.database
.as_mut()
.ok_or_else(|| "The project database is unavailable.".to_owned())?
.continue_tool_turn(session_id, result, assistant_reasoning)
.map_err(|error| format!("Could not save the tool turn: {error}"))?;
self.conversation.push(ChatMessage::from(saved.0));
let messages = self
.conversation
.iter()
.map(|message| ChatTurn {
user: message.user,
tool: message.tool,
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.preferences.idle_timeout_minutes.max(1) as u64 * 60);
self.active_generation = Some(
self.generation_service
.as_ref()
.ok_or_else(|| "The model runtime is unavailable.".to_owned())?
.generate(
effective.engine,
effective.turn,
messages,
CheckpointTarget::Local(session_checkpoint_path(session_id)),
idle_timeout,
)?,
);
self.tokens_per_second = None;
Ok(())
}
/// Asks the model for a session title without opening a session for it: the
/// turn runs against the shared transient KV cache and only its text is kept.
///
@@ -318,6 +450,7 @@ impl App {
.filter(|message| !message.content.trim().is_empty())
.map(|message| ChatTurn {
user: message.user,
tool: message.tool,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -329,6 +462,7 @@ impl App {
}
messages.push(ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,

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

@@ -57,7 +57,13 @@ impl App {
} else {
let markdown_style = markdown::Style::from_palette(app_theme().palette());
for (index, message) in self.conversation.iter().enumerate() {
let label = if message.user { "You" } else { "DS4" };
let label = if message.user {
"You"
} else if message.tool {
"Tool"
} else {
"DS4"
};
let active = self.generating && index + 1 == self.conversation.len();
let mut body = column![text(label).size(11)].spacing(5);
if let Some(reasoning) = &message.reasoning {
@@ -89,11 +95,11 @@ impl App {
}
}
if !message.content.is_empty() {
if message.user || message.markdown.is_empty() {
if message.user || message.tool || message.markdown.is_empty() {
let content = if message.reasoning.is_some() {
message.content.trim_start()
crate::agent::visible_content(&message.content).trim_start()
} else {
&message.content
crate::agent::visible_content(&message.content)
};
body = body.push(text(content).size(14));
} else {
@@ -109,6 +115,14 @@ impl App {
} else if active && message.reasoning.is_none() {
body = body.push(text("Loading model…").size(14));
}
if !message.user
&& !message.tool
&& let Some(model) = ModelChoice::from_id(&self.preferences.selected_model)
{
for summary in crate::agent::tool_summaries(model, &message.content) {
body = body.push(text(summary).size(13).color(muted_text()));
}
}
let user = message.user;
messages = messages.push(
container(body)

49
src/database.rs
View File

@@ -347,6 +347,7 @@ pub struct StoredMessage {
pub id: i32,
pub session_id: i32,
pub user: bool,
pub tool: bool,
pub reasoning: Option<String>,
pub reasoning_complete: bool,
pub content: String,
@@ -357,6 +358,7 @@ pub struct StoredMessage {
struct NewMessage<'a> {
session_id: i32,
user: bool,
tool: bool,
reasoning: Option<&'a str>,
reasoning_complete: bool,
content: &'a str,
@@ -621,6 +623,7 @@ impl Database {
.values(NewMessage {
session_id,
user: true,
tool: false,
reasoning: None,
reasoning_complete: true,
content: prompt,
@@ -631,6 +634,7 @@ impl Database {
.values(NewMessage {
session_id,
user: false,
tool: false,
reasoning: reasoning.then_some(""),
reasoning_complete: !reasoning,
content: "",
@@ -642,6 +646,41 @@ impl Database {
.map_err(|error: diesel::result::Error| error.to_string())
}
pub fn continue_tool_turn(
&mut self,
session_id: i32,
result: &str,
reasoning: bool,
) -> Result<(StoredMessage, 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 assistant = diesel::insert_into(messages::table)
.values(NewMessage {
session_id,
user: false,
tool: false,
reasoning: reasoning.then_some(""),
reasoning_complete: !reasoning,
content: "",
})
.returning(StoredMessage::as_returning())
.get_result(connection)?;
Ok((tool, assistant))
})
.map_err(|error: diesel::result::Error| error.to_string())
}
pub fn update_message(
&mut self,
id: i32,
@@ -832,6 +871,9 @@ mod tests {
database
.update_message(assistant.id, Some("Reasoning"), true, "Answer")
.unwrap();
database
.continue_tool_turn(session.id, "Tool result", false)
.unwrap();
database
.update_session_context(session.id, 1_234, 65_536, Some(12.5))
.unwrap();
@@ -843,12 +885,17 @@ 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(), 2);
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);
reopened.delete_session(session.id).unwrap();
assert!(reopened.load_messages(session.id).unwrap().is_empty());
drop(reopened);

8
src/engine.rs
View File

@@ -319,6 +319,7 @@ pub(crate) struct Generator {
#[derive(Clone)]
pub(crate) struct ChatTurn {
pub(crate) user: bool,
pub(crate) tool: bool,
pub(crate) skip_previous_eos: bool,
pub(crate) reasoning: Option<String>,
pub(crate) reasoning_complete: bool,
@@ -571,6 +572,7 @@ impl Generator {
let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
let mut generated = ChatTurn {
user: false,
tool: false,
skip_previous_eos: false,
reasoning: reasoning.then(String::new),
reasoning_complete: !reasoning,
@@ -896,7 +898,7 @@ fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn
output.extend_from_slice(value.as_bytes());
}
let mut output = b"DS4Server chat checkpoint v2".to_vec();
let mut output = b"DS4Server chat checkpoint v3".to_vec();
text(&mut output, system);
output.push(match reasoning {
ReasoningMode::Direct => 0,
@@ -905,6 +907,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.skip_previous_eos));
match &message.reasoning {
Some(reasoning) => {
@@ -1057,6 +1060,7 @@ mod sampling_tests {
fn split_utf8_token_bytes_are_joined_before_decoding() {
let mut generated = ChatTurn {
user: false,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -1076,6 +1080,7 @@ mod sampling_tests {
fn checkpoint_tag_covers_the_canonical_chat_state() {
let mut messages = vec![ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -1102,6 +1107,7 @@ mod sampling_tests {
let prefix = conversation_key("System", ReasoningMode::High, &messages);
messages.push(ChatTurn {
user: false,
tool: false,
skip_previous_eos: false,
reasoning: Some("because".into()),
reasoning_complete: true,

27
src/engine/tokenizer.rs
View File

@@ -195,6 +195,7 @@ impl Tokenizer {
system_prompt,
&[ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -235,6 +236,32 @@ impl Tokenizer {
output.extend(self.tokenize_rendered(system_prompt));
}
for message in messages {
if message.tool {
if self.family == ModelFamily::Glm {
output.push(self.observation);
output.extend(self.tokenize_rendered("<tool_response>"));
output.extend(
self.tokenize_rendered(
&message
.content
.replace("</tool_response>", "&lt;/tool_response>"),
),
);
output.extend(self.tokenize_rendered("</tool_response>"));
} else {
output.push(self.user);
output.extend(self.tokenize("<tool_result>"));
output.extend(
self.tokenize(
&message
.content
.replace("</tool_result>", "&lt;/tool_result>"),
),
);
output.extend(self.tokenize("</tool_result>"));
}
continue;
}
output.push(if message.user {
self.user
} else {

26
src/engine/validation.rs
View File

@@ -904,6 +904,7 @@ mod tests {
&[
ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -911,6 +912,7 @@ mod tests {
},
ChatTurn {
user: false,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -918,6 +920,7 @@ mod tests {
},
ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -938,12 +941,33 @@ mod tests {
model.render_continuation("Hello", ReasoningMode::Direct, true),
[128_803, 19_923, 128_804, 128_822]
);
let tool_turn = ChatTurn {
user: false,
tool: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
};
let wrapped_user = ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "<tool_result>Hello</tool_result>".into(),
};
assert_eq!(
model.render_conversation("", &[tool_turn], ReasoningMode::Direct),
model.render_conversation("", &[wrapped_user], ReasoningMode::Direct)
);
assert_eq!(
model.render_conversation(
"",
&[
ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -951,6 +975,7 @@ mod tests {
},
ChatTurn {
user: false,
tool: false,
skip_previous_eos: false,
reasoning: Some("Hello".into()),
reasoning_complete: true,
@@ -958,6 +983,7 @@ mod tests {
},
ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,

1
src/main.rs
View File

@@ -1,3 +1,4 @@
mod agent;
mod app;
mod database;
mod engine;

1
src/schema.rs
View File

@@ -3,6 +3,7 @@ diesel::table! {
id -> Integer,
session_id -> Integer,
user -> Bool,
tool -> Bool,
reasoning -> Nullable<Text>,
reasoning_complete -> Bool,
content -> Text,

43
src/server.rs
View File

@@ -312,6 +312,49 @@ impl Drop for ConnectionSlot {
}
}
pub(crate) fn parse_dsml_tool_calls(text: &str) -> Result<(String, Vec<(String, Value)>), String> {
let has_tool_marker = TOOL_SYNTAXES
.iter()
.any(|syntax| text.contains(syntax.tool_start));
let mut projector = ToolProjector::new();
let events = projector.push(text, true, "");
let mut content = String::new();
let mut calls = Vec::<(String, String, bool)>::new();
for event in events {
match event {
ToolProjectionEvent::Text(text) => content.push_str(&text),
ToolProjectionEvent::Start { index, name, .. } => {
if calls.len() == index {
calls.push((name, String::new(), false));
}
}
ToolProjectionEvent::Arguments { index, fragment } => {
if let Some((_, arguments, _)) = calls.get_mut(index) {
arguments.push_str(&fragment);
}
}
ToolProjectionEvent::End { index } => {
if let Some((_, _, complete)) = calls.get_mut(index) {
*complete = true;
}
}
}
}
let calls = calls
.into_iter()
.filter(|(_, _, complete)| *complete)
.map(|(name, arguments, _)| {
serde_json::from_str(&arguments)
.map(|arguments| (name, arguments))
.map_err(|error| format!("invalid DSML tool arguments: {error}"))
})
.collect::<Result<Vec<_>, _>>()?;
if has_tool_marker && calls.is_empty() {
return Err("invalid or incomplete DSML tool call".into());
}
Ok((content, calls))
}
fn handle(mut stream: TcpStream, state: &State) {
let _ = stream.set_write_timeout(Some(HTTP_IO_TIMEOUT));
let started = Instant::now();

9
src/server/tools.rs
View File

@@ -132,7 +132,11 @@ impl ToolProjector {
self.state = ToolProjectionState::Failed;
break;
};
let id = random_tool_id(prefix);
let id = if prefix.is_empty() {
String::new()
} else {
random_tool_id(prefix)
};
self.ids.push(id.clone());
events.push(ToolProjectionEvent::Start {
index: self.index,
@@ -370,6 +374,7 @@ pub(super) fn render_messages(
}
"user" => turns.push(ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
@@ -388,6 +393,7 @@ pub(super) fn render_messages(
} else {
turns.push(ChatTurn {
user: true,
tool: false,
skip_previous_eos: protocol == Protocol::Responses,
reasoning: None,
reasoning_complete: true,
@@ -403,6 +409,7 @@ pub(super) fn render_messages(
let reasoning = content_text(&message.reasoning_content);
turns.push(ChatTurn {
user: false,
tool: false,
skip_previous_eos: false,
reasoning: (preserve_reasoning && !reasoning.is_empty()).then_some(reasoning),
reasoning_complete: true,