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

Reuse canonical chat context

Browse Source
This commit is contained in:
Georg Bauer
2026-07-27 18:35:52 +02:00
parent ee7e5696a5
commit 4adb49a4d1
10 changed files with 567 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

404
src/app/generation.rs
View File

@@ -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<String>,
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<String>,
}
#[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<i32>,
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<StoredMessage> 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<Option<String>, 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<Item = ChatTurn>) -> Vec<ChatTurn> {
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::<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(),
}));
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<crate::agent::ToolCall>) -> 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::<Vec<_>>();
});
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<String> {
#[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,