hugo/DS4Server
1
0
Fork 0
Code Issues 7 Pull Requests Actions 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

1
migrations/20260727170000_add_message_model_content/down.sql Normal file
View File

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

1
migrations/20260727170000_add_message_model_content/up.sql Normal file
View File

@@ -0,0 +1 @@
ALTER TABLE messages ADD COLUMN model_content TEXT;

6
src/app.rs
View File

@@ -83,6 +83,7 @@ pub(crate) struct App {
pub(super) composer: String,
pub(super) queued_inputs: VecDeque<String>,
pub(super) conversation: Vec<ChatMessage>,
active_turn: Option<generation::TurnSummary>,
pub(super) a2ui: crate::a2ui::Store,
pub(super) a2ui_history: Vec<crate::a2ui::Store>,
pub(super) a2ui_history_index: Option<usize>,
@@ -406,6 +407,7 @@ impl App {
composer: String::new(),
queued_inputs: VecDeque::new(),
conversation: Vec::new(),
active_turn: None,
a2ui: crate::a2ui::Store::default(),
a2ui_history: Vec::new(),
a2ui_history_index: None,
@@ -533,6 +535,7 @@ impl App {
composer: String::new(),
queued_inputs: VecDeque::new(),
conversation: Vec::new(),
active_turn: None,
a2ui: crate::a2ui::Store::default(),
a2ui_history: Vec::new(),
a2ui_history_index: None,
@@ -1479,6 +1482,7 @@ impl App {
match loaded {
Ok((messages, a2ui)) => {
self.conversation = messages.into_iter().map(ChatMessage::from).collect();
generation::promote_legacy_turn_summaries(&mut self.conversation);
self.context_notice = None;
self.clear_a2ui();
let (history, active, errors) =
@@ -2292,6 +2296,7 @@ mod tests {
reasoning_complete: false,
reasoning_open: true,
content: String::new(),
model_content: None,
markdown: markdown::Content::new(),
transcript: text_editor::Content::new(),
a2ui_lines_processed: 0,
@@ -2322,6 +2327,7 @@ mod tests {
reasoning_complete: true,
reasoning_open: false,
content: content.into(),
model_content: None,
markdown: markdown::Content::new(),
transcript: text_editor::Content::new(),
a2ui_lines_processed: 0,

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,

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

@@ -86,6 +86,11 @@ impl App {
{
continue;
}
if message.user
&& let Some(stats) = message.generation_stats
{
messages = messages.push(generation_summary(stats));
}
let label = if message.user {
"You"
} else if message.tool {
@@ -193,9 +198,6 @@ impl App {
} else {
message_body
});
if let Some(stats) = message.generation_stats {
messages = messages.push(generation_summary(stats));
}
}
if let Some(activity) = &self.activity {
messages = messages.push(
@@ -388,6 +390,7 @@ fn generation_summary(stats: crate::app::generation::GenerationStats) -> Element
rule::horizontal(1),
]
.spacing(5)
.padding(Padding::ZERO.right(24))
.into()
}

31
src/database.rs
View File

@@ -118,6 +118,7 @@ pub struct StoredMessage {
pub reasoning: Option<String>,
pub reasoning_complete: bool,
pub content: String,
pub model_content: Option<String>,
pub system: bool,
pub compaction: bool,
pub compaction_tail_start: Option<i32>,
@@ -136,6 +137,7 @@ struct NewMessage<'a> {
reasoning: Option<&'a str>,
reasoning_complete: bool,
content: &'a str,
model_content: Option<&'a str>,
system: bool,
compaction: bool,
compaction_tail_start: Option<i32>,
@@ -376,6 +378,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content: &content,
model_content: None,
system: true,
compaction: false,
compaction_tail_start: None,
@@ -422,6 +425,7 @@ impl Database {
&mut self,
session_id: i32,
prompt: &str,
model_prompt: Option<&str>,
system_messages: &[String],
reasoning: bool,
) -> Result<Vec<StoredMessage>, String> {
@@ -438,6 +442,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content,
model_content: None,
system: true,
compaction: false,
compaction_tail_start: None,
@@ -454,6 +459,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content: prompt,
model_content: model_prompt,
system: false,
compaction: false,
compaction_tail_start: None,
@@ -468,6 +474,7 @@ impl Database {
reasoning: reasoning.then_some(""),
reasoning_complete: !reasoning,
content: "",
model_content: None,
system: false,
compaction: false,
compaction_tail_start: None,
@@ -500,6 +507,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content: result,
model_content: None,
system: false,
compaction: false,
compaction_tail_start: None,
@@ -517,6 +525,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content,
model_content: None,
system: false,
compaction: false,
compaction_tail_start: None,
@@ -535,6 +544,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content,
model_content: None,
system: true,
compaction: false,
compaction_tail_start: None,
@@ -551,6 +561,7 @@ impl Database {
reasoning: reasoning.then_some(""),
reasoning_complete: !reasoning,
content: "",
model_content: None,
system: false,
compaction: false,
compaction_tail_start: None,
@@ -642,6 +653,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content: summary,
model_content: None,
system: true,
compaction: true,
compaction_tail_start: tail_start,
@@ -659,6 +671,7 @@ impl Database {
reasoning: None,
reasoning_complete: true,
content,
model_content: None,
system: false,
compaction: false,
compaction_tail_start: None,
@@ -758,7 +771,7 @@ mod tests {
.unwrap();
let session = database.create_session(project.id, "A2UI").unwrap();
let first = database
.start_chat_turn(session.id, "First", &[], false)
.start_chat_turn(session.id, "First", None, &[], false)
.unwrap()
.pop()
.unwrap();
@@ -785,7 +798,7 @@ mod tests {
assert!(active.active_surface().is_none());
let second = database
.start_chat_turn(session.id, "Second", &[], false)
.start_chat_turn(session.id, "Second", None, &[], false)
.unwrap()
.pop()
.unwrap();
@@ -836,7 +849,13 @@ mod tests {
let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap();
let session = database.create_session(project.id, "Chat").unwrap();
let mut opening = database
.start_chat_turn(session.id, "Question", &["Date context".into()], true)
.start_chat_turn(
session.id,
"Question",
Some("Question\n\nhidden metadata"),
&["Date context".into()],
true,
)
.unwrap();
let assistant = opening.pop().unwrap();
database
@@ -875,6 +894,10 @@ mod tests {
assert_eq!(messages.len(), 7);
assert!(messages[0].system);
assert_eq!(messages[1].content, "Question");
assert_eq!(
messages[1].model_content.as_deref(),
Some("Question\n\nhidden metadata")
);
assert_eq!(messages[2].reasoning.as_deref(), Some("Reasoning"));
assert!(messages[2].reasoning_complete);
assert_eq!(messages[2].content, "Answer");
@@ -956,7 +979,7 @@ mod tests {
)
.unwrap();
reopened
.start_chat_turn(session.id, "After third compaction", &[], false)
.start_chat_turn(session.id, "After third compaction", None, &[], false)
.unwrap();
drop(reopened);

203
src/engine.rs
View File

@@ -560,6 +560,7 @@ impl Generator {
pub(crate) fn generate(
&mut self,
checkpoint: &Path,
bootstrap: Option<&Path>,
messages: &[ChatTurn],
settings: &TurnSettings,
cancelled: &AtomicBool,
@@ -578,6 +579,11 @@ impl Generator {
if checkpoint_present && !selected.found {
phase(checkpoint_rebuild_activity(selected.incompatible));
}
if !selected.found
&& let Some(directory) = bootstrap
{
self.prepare_bootstrap(directory, settings, cancelled, &mut progress, &mut phase)?;
}
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
@@ -596,6 +602,8 @@ impl Generator {
output.checkpoint_bytes = std::fs::metadata(checkpoint)
.map(|item| item.len())
.unwrap_or(0);
self.checkpoint = Some(checkpoint.to_owned());
self.resident_active = Some(checkpoint.to_owned());
Ok(output)
}
@@ -648,32 +656,8 @@ impl Generator {
let history_key =
conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
let history_tag: [u8; 32] = Sha256::digest(&history_key).into();
let mut previous_checkpoint = self.checkpoint.clone();
if self.executor.checkpoint_tag() != history_tag {
if let Some(entry) = store.find(&history_key, self.executor.context()) {
if self.select_checkpoint(&entry.checkpoint, entry.tag)?.found {
store.touch(&entry)?;
self.last_store_tokens = entry.tokens;
previous_checkpoint = Some(entry.checkpoint);
} else {
store.discard(&entry);
self.last_store_tokens = 0;
previous_checkpoint = None;
}
} else {
let key = resident_key(directory, history_tag);
let restored = self.activate_resident(key)?;
if !restored || self.executor.checkpoint_tag() != history_tag {
self.executor.reset()?;
}
self.checkpoint = None;
self.last_store_tokens = 0;
previous_checkpoint = None;
self.metrics.kv_lookup(KvLookup::Miss);
}
} else {
self.metrics.kv_lookup(KvLookup::MemoryHit);
}
let previous_checkpoint =
self.restore_cached_prefix(directory, &store, &history_key, history_tag)?;
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
@@ -731,6 +715,98 @@ impl Generator {
Ok(output)
}
fn prepare_bootstrap(
&mut self,
directory: &Path,
settings: &TurnSettings,
cancelled: &AtomicBool,
progress: &mut impl FnMut(u32, u32, Option<f32>),
phase: &mut impl FnMut(&'static str),
) -> Result<(), String> {
let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &[]);
let tag: [u8; 32] = Sha256::digest(&key).into();
let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
self.restore_cached_prefix(directory, &store, &key, tag)?;
let tokens = self.executor.model().render_history(
&settings.system_prompt,
&[],
settings.reasoning_mode,
);
if tokens.len() >= self.executor.context() as usize {
return Err(format!(
"System prompt has {} tokens, but the configured context size is {} tokens",
tokens.len(),
self.executor.context()
));
}
let reused = self.executor.align_prompt(&tokens)?;
if reused == tokens.len() {
return Ok(());
}
phase("Updating system prompt cache…");
let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?;
if completed != tokens.len() - reused {
return Err("generation cancelled while updating the system prompt cache".into());
}
self.executor.note_checkpoint_tag(tag);
if !settings.kv_cache.stores(self.executor.position(), true, 0) {
self.checkpoint = None;
self.resident_active = Some(resident_key(directory, tag));
return Ok(());
}
let checkpoint = store.checkpoint_path(&key);
self.save_checkpoint(&checkpoint, tag)?;
let retained = store.record(
&checkpoint,
&key,
tag,
self.executor.position(),
self.executor.context(),
StoreReason::Cold,
)?;
self.checkpoint = retained.then_some(checkpoint.clone());
self.resident_active = Some(if retained {
checkpoint
} else {
resident_key(directory, tag)
});
Ok(())
}
fn restore_cached_prefix(
&mut self,
directory: &Path,
store: &KvStore,
key: &[u8],
tag: [u8; 32],
) -> Result<Option<PathBuf>, String> {
if self.executor.checkpoint_tag() == tag {
self.metrics.kv_lookup(KvLookup::MemoryHit);
return Ok(self.checkpoint.clone());
}
if let Some(entry) = store.find(key, self.executor.context()) {
if self.select_checkpoint(&entry.checkpoint, entry.tag)?.found {
store.touch(&entry)?;
self.last_store_tokens = entry.tokens;
return Ok(Some(entry.checkpoint));
}
store.discard(&entry);
} else {
let key = resident_key(directory, tag);
let restored = self.activate_resident(key)?;
if !restored || self.executor.checkpoint_tag() != tag {
self.executor.reset()?;
}
self.metrics.kv_lookup(KvLookup::Miss);
}
self.checkpoint = None;
self.last_store_tokens = 0;
Ok(None)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn compact(
&mut self,
@@ -986,6 +1062,34 @@ impl Generator {
result
}
fn prefill_suffix(
&mut self,
tokens: &[i32],
reused: usize,
cancelled: &AtomicBool,
progress: &mut impl FnMut(u32, u32, Option<f32>),
) -> Result<usize, String> {
let suffix = &tokens[reused..];
if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
let context = self.executor.context();
self.executor.prefill(suffix, |used| {
progress(used, context, None);
!cancelled.load(Ordering::Relaxed)
})
} else {
let mut completed = 0;
for &token in suffix {
if cancelled.load(Ordering::Relaxed) {
break;
}
self.executor.eval(token)?;
completed += 1;
progress(self.executor.position(), self.executor.context(), None);
}
Ok(completed)
}
}
fn generate_inner(
&mut self,
messages: &[ChatTurn],
@@ -1047,24 +1151,7 @@ impl Generator {
let mut pending_utf8 = Vec::new();
let prompt_tokens = tokens.len();
let suffix = &tokens[reused..];
let completed = if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
let context = self.executor.context();
self.executor.prefill(suffix, |used| {
progress(used, context, None);
!cancelled.load(Ordering::Relaxed)
})?
} else {
let mut completed = 0;
for &token in suffix {
if cancelled.load(Ordering::Relaxed) {
break;
}
self.executor.eval(token)?;
completed += 1;
progress(self.executor.position(), self.executor.context(), None);
}
completed
};
let completed = self.prefill_suffix(&tokens, reused, cancelled, progress)?;
self.publish_execution_stats();
if completed != suffix.len() {
return Ok((
@@ -1648,6 +1735,34 @@ mod sampling_tests {
assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix));
}
#[test]
fn bootstrap_key_is_the_prefix_before_dynamic_session_context() {
let system = "System\n\nProject instructions from AGENTS.md:\n\nkeep this";
let bootstrap = conversation_key(system, ReasoningMode::High, &[]);
let messages = vec![
ChatTurn {
user: false,
tool: false,
system: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "current date and time".into(),
},
ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "hello".into(),
},
];
assert!(conversation_key(system, ReasoningMode::High, &messages).starts_with(&bootstrap));
}
#[test]
fn checkpoint_rebuilds_explain_compatibility_and_history_misses() {
assert!(checkpoint_rebuild_activity(true).contains("different model"));

31
src/runtime.rs
View File

@@ -19,7 +19,6 @@ pub(crate) struct GenerationService {
pub(crate) struct ActiveGeneration {
pub(crate) events: Receiver<GenerationEvent>,
pub(crate) cancel: Arc<AtomicBool>,
pub(crate) started_at: Instant,
}
impl Drop for ActiveGeneration {
@@ -29,7 +28,11 @@ impl Drop for ActiveGeneration {
}
pub(crate) enum CheckpointTarget {
Local(PathBuf),
Local {
checkpoint: PathBuf,
/// Shared cache for the deterministic rendered system prompt.
bootstrap: Option<PathBuf>,
},
Transient(PathBuf),
/// Same transient KV handling as [`CheckpointTarget::Transient`], but asked
/// for by the app itself (session titling) rather than by an HTTP client.
@@ -39,7 +42,7 @@ pub(crate) enum CheckpointTarget {
impl CheckpointTarget {
fn source(&self) -> WorkSource {
match self {
Self::Local(_) | Self::OneShot(_) => WorkSource::LocalChat,
Self::Local { .. } | Self::OneShot(_) => WorkSource::LocalChat,
Self::Transient(_) => WorkSource::Http,
}
}
@@ -108,7 +111,6 @@ impl GenerationService {
checkpoint: CheckpointTarget,
idle_timeout: Duration,
) -> Result<ActiveGeneration, String> {
let started_at = Instant::now();
let source = checkpoint.source();
let cancel = Arc::new(AtomicBool::new(false));
let (events, receiver) = mpsc::channel();
@@ -133,7 +135,6 @@ impl GenerationService {
Ok(ActiveGeneration {
events: receiver,
cancel,
started_at,
})
}
@@ -148,7 +149,6 @@ impl GenerationService {
checkpoint: PathBuf,
idle_timeout: Duration,
) -> Result<ActiveGeneration, String> {
let started_at = Instant::now();
let cancel = Arc::new(AtomicBool::new(false));
let (events, receiver) = mpsc::channel();
self.commands
@@ -156,7 +156,10 @@ impl GenerationService {
engine,
turn,
messages,
checkpoint: CheckpointTarget::Local(checkpoint),
checkpoint: CheckpointTarget::Local {
checkpoint,
bootstrap: None,
},
operation: Operation::Compact {
reason: reason.to_owned(),
rebuild_system_prompt,
@@ -169,7 +172,6 @@ impl GenerationService {
Ok(ActiveGeneration {
events: receiver,
cancel,
started_at,
})
}
@@ -180,7 +182,6 @@ impl GenerationService {
messages: Vec<ChatTurn>,
idle_timeout: Duration,
) -> Result<ActiveGeneration, String> {
let started_at = Instant::now();
let cancel = Arc::new(AtomicBool::new(false));
let (events, receiver) = mpsc::channel();
self.metrics.request_queued(WorkSource::LocalChat);
@@ -199,7 +200,6 @@ impl GenerationService {
Ok(ActiveGeneration {
events: receiver,
cancel,
started_at,
})
}
}
@@ -335,7 +335,7 @@ fn run_command(
rebuild_system_prompt,
} = &command.operation
{
let CheckpointTarget::Local(checkpoint) = &command.checkpoint else {
let CheckpointTarget::Local { checkpoint, .. } = &command.checkpoint else {
unreachable!("compaction checkpoints are local")
};
let result = generator.compact(
@@ -367,8 +367,12 @@ fn run_command(
return;
}
let result = match command.checkpoint {
CheckpointTarget::Local(path) => generator.generate(
&path,
CheckpointTarget::Local {
checkpoint,
bootstrap,
} => generator.generate(
&checkpoint,
bootstrap.as_deref(),
&command.messages,
&command.turn,
&command.cancel,
@@ -440,7 +444,6 @@ mod tests {
drop(ActiveGeneration {
events,
cancel: Arc::clone(&cancel),
started_at: Instant::now(),
});
assert!(cancel.load(Ordering::Relaxed));

1
src/schema.rs
View File

@@ -17,6 +17,7 @@ diesel::table! {
reasoning -> Nullable<Text>,
reasoning_complete -> Bool,
content -> Text,
model_content -> Nullable<Text>,
system -> Bool,
compaction -> Bool,
compaction_tail_start -> Nullable<Integer>,

1
src/server.rs
View File

@@ -945,7 +945,6 @@ mod tests {
let active = crate::runtime::ActiveGeneration {
events: receiver,
cancel: Arc::new(AtomicBool::new(false)),
started_at: std::time::Instant::now(),
};
let producer = thread::spawn(move || {
thread::sleep(Duration::from_millis(25));