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

Finish long-running agent parity

Browse Source
This commit is contained in:
Georg Bauer
2026-07-26 13:41:53 +02:00
parent 65c9cbfc45
commit 171b041ba6
15 changed files with 896 additions and 324 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/agent.rs
View File

@@ -14,6 +14,7 @@ use std::thread;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
const MAX_FILE_BYTES: u64 = 16 * 1024 * 1024;
pub(crate) const COMPACTION_OBSERVATION_PREFIX: &str = "Bash job update after context compaction.";
#[repr(C)]
struct WebConfig {
@@ -297,7 +298,7 @@ impl Tools {
}
running.sort();
Some(format!(
"Bash job update after context compaction. Running jobs still need explicit bash_status or bash_stop if relevant.\n{}",
"{COMPACTION_OBSERVATION_PREFIX} Running jobs still need explicit bash_status or bash_stop if relevant.\n{}",
running.concat()
))
}
@@ -988,10 +989,10 @@ pub(crate) fn system_prompt(model: ModelChoice, extra: &str) -> String {
}
}
pub(crate) fn system_prompt_reminder(model: ModelChoice, extra: &str) -> String {
pub(crate) fn system_prompt_reminder(model: ModelChoice) -> String {
format!(
"[System prompt reminder follows.]\n{}\n[End system prompt reminder.]",
system_prompt(model, extra)
system_prompt(model, "")
)
}
@@ -1256,7 +1257,7 @@ mod tests {
}
assert!(prompt.ends_with("extra"));
assert!(
system_prompt_reminder(ModelChoice::DeepSeekV4Flash, "extra")
system_prompt_reminder(ModelChoice::DeepSeekV4Flash)
.contains("[System prompt reminder follows.]")
);
assert!(datetime_context().starts_with("Current local date and time at session start:"));

18
src/app.rs
View File

@@ -97,6 +97,8 @@ pub(crate) struct App {
#[cfg(target_os = "macos")]
active_compaction: Option<generation::CompactionRequest>,
#[cfg(target_os = "macos")]
active_tool_check: Option<generation::ToolResultCheck>,
#[cfg(target_os = "macos")]
agent_tools: Option<(i32, Arc<Mutex<crate::agent::Tools>>)>,
#[cfg(target_os = "macos")]
active_tools: Option<crate::agent::ActiveTools>,
@@ -110,6 +112,7 @@ pub(crate) struct App {
pub(super) activity: Option<String>,
stop_requested: bool,
system_prompt_seen_at: u32,
manual_compaction_queued: bool,
#[cfg(target_os = "macos")]
skip_compaction_once: bool,
}
@@ -307,6 +310,7 @@ impl App {
active_generation: None,
#[cfg(target_os = "macos")]
active_compaction: None,
active_tool_check: None,
#[cfg(target_os = "macos")]
agent_tools: None,
#[cfg(target_os = "macos")]
@@ -330,6 +334,7 @@ impl App {
activity: None,
stop_requested: false,
system_prompt_seen_at: 0,
manual_compaction_queued: false,
#[cfg(target_os = "macos")]
skip_compaction_once: false,
}
@@ -404,6 +409,7 @@ impl App {
active_generation: None,
#[cfg(target_os = "macos")]
active_compaction: None,
active_tool_check: None,
#[cfg(target_os = "macos")]
agent_tools: None,
#[cfg(target_os = "macos")]
@@ -421,6 +427,7 @@ impl App {
activity: None,
stop_requested: false,
system_prompt_seen_at: 0,
manual_compaction_queued: false,
#[cfg(target_os = "macos")]
skip_compaction_once: false,
}
@@ -790,6 +797,10 @@ impl App {
if let Some(compaction) = &self.active_compaction {
compaction.active.cancel.store(true, Ordering::Relaxed);
}
#[cfg(target_os = "macos")]
if let Some(check) = &self.active_tool_check {
check.active.cancel.store(true, Ordering::Relaxed);
}
}
Message::GenerationTick => {
#[cfg(target_os = "macos")]
@@ -945,8 +956,9 @@ impl App {
}
Message::CompactSession(session_id) => {
self.session_menu = None;
if self.generating || self.selected_session != Some(session_id) {
self.error = Some("Open an idle session before compacting it.".into());
if !self.can_compact_session(session_id) {
self.error =
Some("Open an idle session with new chat before compacting it.".into());
} else {
#[cfg(target_os = "macos")]
if let Err(error) = self.start_compaction(
@@ -1617,6 +1629,8 @@ mod tests {
user: false,
tool: false,
system: false,
compaction: false,
compaction_tail_start: None,
reasoning: Some(String::new()),
reasoning_complete: false,
reasoning_open: true,

602
src/app/generation.rs
View File

@@ -29,6 +29,22 @@ pub(super) enum PendingContinuation {
pub(super) struct CompactionRequest {
pub(super) active: ActiveGeneration,
pending: PendingContinuation,
message_ids: Vec<i32>,
}
#[cfg(target_os = "macos")]
#[derive(Clone, Copy)]
enum ToolCheckStage {
Initial,
AfterCompaction,
BoundedError,
}
#[cfg(target_os = "macos")]
pub(super) struct ToolResultCheck {
pub(super) active: ActiveGeneration,
result: String,
stage: ToolCheckStage,
}
#[derive(Clone, Debug)]
@@ -37,6 +53,8 @@ pub(crate) struct ChatMessage {
pub(super) user: bool,
pub(super) tool: bool,
pub(super) system: bool,
pub(super) compaction: bool,
pub(super) compaction_tail_start: Option<i32>,
pub(super) reasoning: Option<String>,
pub(super) reasoning_complete: bool,
pub(super) reasoning_open: bool,
@@ -74,6 +92,8 @@ impl From<StoredMessage> for ChatMessage {
user: message.user,
tool: message.tool,
system: message.system,
compaction: message.compaction,
compaction_tail_start: message.compaction_tail_start,
reasoning: message.reasoning,
reasoning_complete: message.reasoning_complete,
reasoning_open: false,
@@ -85,7 +105,75 @@ impl From<StoredMessage> for ChatMessage {
}
}
#[cfg(target_os = "macos")]
fn chat_turn(message: &ChatMessage) -> ChatTurn {
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(),
}
}
fn queued_prompt(messages: impl IntoIterator<Item = String>) -> Option<String> {
let messages = messages.into_iter().collect::<Vec<_>>();
match messages.as_slice() {
[] => None,
[message] => Some(message.clone()),
_ => Some(
messages
.iter()
.enumerate()
.map(|(index, message)| format!("Queued user message {}:\n{message}", index + 1))
.collect::<Vec<_>>()
.join("\n\n"),
),
}
}
fn compacted_context_start(messages: &[ChatMessage]) -> usize {
let Some(marker_index) = messages.iter().rposition(|message| message.compaction) else {
return 0;
};
messages[marker_index]
.compaction_tail_start
.and_then(|id| {
messages[..marker_index]
.iter()
.position(|message| message.id == id)
})
.unwrap_or(marker_index + 1)
}
fn has_chat_after_last_compaction(messages: &[ChatMessage]) -> bool {
let start = messages
.iter()
.rposition(|message| message.compaction)
.map_or(0, |index| index + 1);
messages[start..].iter().any(|message| {
!message.system
&& !message.compaction
&& !message
.content
.starts_with(crate::agent::COMPACTION_OBSERVATION_PREFIX)
&& (!message.content.trim().is_empty()
|| message
.reasoning
.as_deref()
.is_some_and(|reasoning| !reasoning.trim().is_empty()))
})
}
impl App {
pub(super) fn can_compact_session(&self, session_id: i32) -> bool {
!self.generating
&& self.selected_session == Some(session_id)
&& has_chat_after_last_compaction(&self.conversation)
}
pub(super) fn start_generation(&mut self) {
if self.selected_project.is_none() {
return;
@@ -94,6 +182,19 @@ impl App {
if prompt.is_empty() {
return;
}
#[cfg(target_os = "macos")]
if prompt == "/compact" {
self.composer.clear();
if self.generating {
self.manual_compaction_queued = true;
self.activity = Some("Compaction queued for the next safe point…".into());
} else if let Err(error) =
self.start_compaction(PendingContinuation::None, "manual /compact request")
{
self.error = Some(error);
}
return;
}
if self.generating {
self.queued_inputs.push_back(prompt);
self.composer.clear();
@@ -109,16 +210,6 @@ impl App {
return;
}
#[cfg(target_os = "macos")]
if prompt == "/compact" {
self.composer.clear();
if let Err(error) =
self.start_compaction(PendingContinuation::None, "manual /compact request")
{
self.error = Some(error);
}
return;
}
#[cfg(target_os = "macos")]
if !std::mem::take(&mut self.skip_compaction_once)
&& crate::compaction::should_compact(self.context_used, self.context_limit)
{
@@ -163,24 +254,13 @@ impl App {
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,
));
injected_system.extend(self.system_prompt_reminders(model));
}
#[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,
content: message.content.clone(),
})
.model_chat_messages()
.into_iter()
.map(chat_turn)
.collect::<Vec<_>>();
#[cfg(target_os = "macos")]
messages.extend(injected_system.iter().map(|content| ChatTurn {
@@ -290,12 +370,24 @@ impl App {
crate::agent::prompt_reminder_due(self.context_used, self.system_prompt_seen_at)
}
fn system_prompt_reminders(&self, model: ModelChoice) -> Vec<String> {
let mut reminders = vec![crate::agent::system_prompt_reminder(model)];
if !self.config.generation.system_prompt.trim().is_empty() {
reminders.push(self.config.generation.system_prompt.clone());
}
reminders
}
pub(super) fn poll_generation(&mut self) -> bool {
#[cfg(target_os = "macos")]
if self.active_compaction.is_some() {
return self.poll_compaction();
}
#[cfg(target_os = "macos")]
if self.active_tool_check.is_some() {
return self.poll_tool_result_check();
}
#[cfg(target_os = "macos")]
if let Some(active) = &self.active_tools {
match crate::agent::try_tool_result(active) {
Ok(Some(result)) => {
@@ -304,10 +396,21 @@ impl App {
if cancelled {
self.generating = false;
self.activity = Some("Stopped".into());
self.start_next_queued();
return false;
}
if let Err(error) = self.continue_after_tool_result(&result) {
let continuation = if self.manual_compaction_queued {
self.manual_compaction_queued = false;
self.start_compaction(
PendingContinuation::Tool(result),
"queued manual compaction",
)
} else {
self.start_tool_result_check(result, ToolCheckStage::Initial)
};
if let Err(error) = continuation {
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some(error);
}
return true;
@@ -337,6 +440,9 @@ impl App {
loop {
match active.events.try_recv() {
Ok(GenerationEvent::Loading) => {}
Ok(GenerationEvent::Activity(activity)) => {
self.activity = Some(activity.into());
}
Ok(GenerationEvent::Compacted(_)) => {
self.generating = false;
self.error =
@@ -358,6 +464,16 @@ impl App {
self.context_used = used;
self.context_limit = limit;
self.tokens_per_second = tokens_per_second;
if !self.stop_requested {
self.activity = Some(
if tokens_per_second.is_some() {
"Generating…"
} else {
"Reading conversation…"
}
.into(),
);
}
context_changed = true;
}
Ok(GenerationEvent::Finished(result)) => {
@@ -365,6 +481,8 @@ impl App {
Ok(_) if self.stop_requested => {
self.generating = false;
self.activity = Some("Stopped".into());
start_queued =
!self.queued_inputs.is_empty() || self.manual_compaction_queued;
}
Ok(_) => {
let model = self.config.model;
@@ -383,7 +501,8 @@ impl App {
Ok(_) => {
self.generating = false;
self.activity = None;
start_queued = !self.queued_inputs.is_empty();
start_queued = !self.queued_inputs.is_empty()
|| self.manual_compaction_queued;
}
Err(error) => {
self.active_tools = Some(crate::agent::error_async(error));
@@ -403,6 +522,13 @@ impl App {
self.active_generation = None;
break;
}
Ok(GenerationEvent::Measured(_)) => {
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some(
"The model runtime returned an unexpected context measurement.".into(),
);
}
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => {
self.generating = false;
@@ -467,7 +593,17 @@ impl App {
}
fn start_next_queued(&mut self) {
if let Some(prompt) = self.queued_inputs.pop_front() {
#[cfg(target_os = "macos")]
if std::mem::take(&mut self.manual_compaction_queued) {
if let Err(error) =
self.start_compaction(PendingContinuation::None, "queued manual compaction")
{
self.activity = Some("Failed".into());
self.error = Some(error);
}
return;
}
if let Some(prompt) = queued_prompt(self.queued_inputs.drain(..)) {
self.composer = prompt;
self.start_generation();
}
@@ -493,6 +629,7 @@ impl App {
}
let tools = Arc::clone(&self.agent_tools.as_ref().unwrap().1);
self.active_tools = Some(crate::agent::execute_async(tools, calls));
self.activity = Some("Running tools…".into());
Ok(())
}
@@ -501,18 +638,7 @@ impl App {
let session_id = self
.selected_session
.ok_or_else(|| "The active session is unavailable.".to_owned())?;
if !std::mem::take(&mut self.skip_compaction_once)
&& crate::compaction::tool_result_needs_compaction(
self.context_used,
self.context_limit,
result,
)
{
return self.start_compaction(
PendingContinuation::Tool(result.to_owned()),
"context pressure before tool continuation",
);
}
self.skip_compaction_once = false;
let model = self.config.model;
let mut effective = crate::settings::effective_settings(
model,
@@ -527,10 +653,12 @@ impl App {
self.compaction_summary(),
);
let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct;
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 queued = queued_prompt(self.queued_inputs.drain(..));
let reminders = if self.system_prompt_reminder_due() {
self.system_prompt_reminders(model)
} else {
Vec::new()
};
let mut saved = self
.database
.as_mut()
@@ -538,13 +666,12 @@ impl App {
.continue_tool_turn(
session_id,
result,
&queued,
reminder.as_deref(),
queued.as_deref(),
&reminders,
assistant_reasoning,
)
.map_err(|error| format!("Could not save the tool turn: {error}"))?;
self.queued_inputs.clear();
if reminder.is_some() {
if !reminders.is_empty() {
self.system_prompt_seen_at = self.context_used;
}
let mut assistant = ChatMessage::from(saved.pop().unwrap());
@@ -552,17 +679,9 @@ impl App {
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(),
})
.model_chat_messages()
.into_iter()
.map(chat_turn)
.collect();
assistant.reasoning_open = assistant_reasoning;
self.conversation.push(assistant);
@@ -585,6 +704,163 @@ impl App {
Ok(())
}
#[cfg(target_os = "macos")]
fn start_tool_result_check(
&mut self,
result: String,
stage: ToolCheckStage,
) -> Result<(), String> {
if matches!(stage, ToolCheckStage::Initial)
&& crate::compaction::should_compact(self.context_used, self.context_limit)
{
return self.start_compaction(
PendingContinuation::Tool(result),
"soft limit before tool continuation",
);
}
let model = self.config.model;
let mut effective = crate::settings::effective_settings(
model,
&self.config.generation,
&self.config.runtime,
&models_path(),
)?;
effective.turn.system_prompt =
crate::agent::system_prompt(model, &effective.turn.system_prompt);
effective.turn.system_prompt = crate::compaction::summary_system_prompt(
&effective.turn.system_prompt,
self.compaction_summary(),
);
let mut messages = self
.model_chat_messages()
.into_iter()
.map(chat_turn)
.collect::<Vec<_>>();
messages.push(ChatTurn {
user: false,
tool: true,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: result.clone(),
});
let idle_timeout = Duration::from_secs(self.config.idle_timeout_minutes.max(1) as u64 * 60);
let active = self
.generation_service
.as_ref()
.ok_or_else(|| "The model runtime is unavailable.".to_owned())?
.measure_context(effective.engine, effective.turn, messages, idle_timeout)?;
self.active_tool_check = Some(ToolResultCheck {
active,
result,
stage,
});
self.activity = Some("Checking tool result context…".into());
Ok(())
}
#[cfg(target_os = "macos")]
fn poll_tool_result_check(&mut self) -> bool {
let event = self
.active_tool_check
.as_mut()
.unwrap()
.active
.events
.try_recv();
match event {
Ok(GenerationEvent::Loading) => {
self.activity = Some("Loading model…".into());
false
}
Ok(GenerationEvent::Measured(Ok(projected))) => {
let check = self.active_tool_check.take().unwrap();
if self.manual_compaction_queued && matches!(check.stage, ToolCheckStage::Initial) {
self.manual_compaction_queued = false;
if let Err(error) = self.start_compaction(
PendingContinuation::Tool(check.result),
"queued manual compaction",
) {
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some(error);
}
return true;
}
let reserve = if matches!(check.stage, ToolCheckStage::BoundedError) {
16
} else {
crate::compaction::tool_result_reserve(self.context_limit)
};
if crate::compaction::tool_result_fits(projected, self.context_limit, reserve) {
if let Err(error) = self.continue_after_tool_result(&check.result) {
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some(error);
}
} else {
match check.stage {
ToolCheckStage::Initial => {
if let Err(error) = self.start_compaction(
PendingContinuation::Tool(check.result),
"tool result would exceed context",
) {
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some(error);
}
}
ToolCheckStage::AfterCompaction => {
let error = crate::compaction::bounded_tool_error(
projected,
self.context_limit,
reserve,
);
if let Err(error) =
self.start_tool_result_check(error, ToolCheckStage::BoundedError)
{
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some(error);
}
}
ToolCheckStage::BoundedError => {
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some("context full after compaction".into());
}
}
}
true
}
Ok(GenerationEvent::Measured(Err(error))) => {
self.active_tool_check = None;
self.generating = false;
if self.stop_requested {
self.activity = Some("Stopped".into());
} else {
self.activity = Some("Failed".into());
self.error = Some(error);
}
false
}
Ok(GenerationEvent::Activity(activity)) => {
self.activity = Some(activity.into());
false
}
Ok(_) => false,
Err(TryRecvError::Empty) => false,
Err(TryRecvError::Disconnected) => {
self.active_tool_check = None;
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some("The model runtime stopped unexpectedly.".into());
false
}
}
}
#[cfg(target_os = "macos")]
pub(super) fn start_compaction(
&mut self,
@@ -600,23 +876,14 @@ impl App {
)?;
effective.turn.system_prompt =
crate::agent::system_prompt(model, &effective.turn.system_prompt);
let rebuild_system_prompt = effective.turn.system_prompt.clone();
effective.turn.system_prompt = crate::compaction::summary_system_prompt(
&effective.turn.system_prompt,
self.compaction_summary(),
);
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 model_messages = self.model_chat_messages();
let message_ids = model_messages.iter().map(|message| message.id).collect();
let messages = model_messages.into_iter().map(chat_turn).collect();
let idle_timeout = Duration::from_secs(self.config.idle_timeout_minutes.max(1) as u64 * 60);
let active = self
.generation_service
@@ -627,13 +894,19 @@ impl App {
effective.turn,
messages,
reason,
rebuild_system_prompt,
session_compaction_checkpoint_path(
self.selected_session
.ok_or_else(|| "The active session is unavailable.".to_owned())?,
),
idle_timeout,
)?;
self.active_compaction = Some(CompactionRequest { active, pending });
self.active_compaction = Some(CompactionRequest {
active,
pending,
message_ids,
});
self.manual_compaction_queued = false;
self.generating = true;
self.activity = Some("Compacting durable task state…".into());
self.tokens_per_second = None;
@@ -648,6 +921,7 @@ impl App {
loop {
match request.active.events.try_recv() {
Ok(GenerationEvent::Loading) => self.activity = Some("Loading model…".into()),
Ok(GenerationEvent::Activity(activity)) => self.activity = Some(activity.into()),
Ok(GenerationEvent::Context {
used,
limit,
@@ -661,10 +935,12 @@ impl App {
let request = self.active_compaction.take().unwrap();
match result {
Ok(compacted) => {
if let Err(error) = self.apply_compaction(&compacted) {
if let Err(error) =
self.apply_compaction(&compacted, &request.message_ids)
{
let _ = fs::remove_file(&compacted.checkpoint);
self.generating = false;
self.activity = None;
self.activity = Some("Failed".into());
self.error = Some(error);
return false;
}
@@ -680,14 +956,12 @@ impl App {
self.start_generation();
}
PendingContinuation::Tool(result) => {
let result = crate::compaction::bounded_tool_result(
self.context_used,
self.context_limit,
if let Err(error) = self.start_tool_result_check(
result,
);
self.skip_compaction_once = true;
if let Err(error) = self.continue_after_tool_result(&result) {
ToolCheckStage::AfterCompaction,
) {
self.generating = false;
self.activity = Some("Failed".into());
self.error = Some(error);
}
}
@@ -696,6 +970,13 @@ impl App {
}
Err(error) => {
self.generating = false;
if let PendingContinuation::User(prompt) = request.pending {
if self.composer.trim().is_empty() {
self.composer = prompt;
} else {
self.queued_inputs.push_front(prompt);
}
}
if self.stop_requested {
self.activity = Some("Stopped".into());
} else {
@@ -706,12 +987,14 @@ impl App {
}
}
}
Ok(GenerationEvent::Finished(_)) | Ok(GenerationEvent::Chunk { .. }) => {}
Ok(GenerationEvent::Finished(_))
| Ok(GenerationEvent::Chunk { .. })
| Ok(GenerationEvent::Measured(_)) => {}
Err(TryRecvError::Empty) => return false,
Err(TryRecvError::Disconnected) => {
self.active_compaction = None;
self.generating = false;
self.activity = None;
self.activity = Some("Failed".into());
self.error = Some("The model runtime stopped unexpectedly.".into());
return false;
}
@@ -723,51 +1006,37 @@ impl App {
fn apply_compaction(
&mut self,
compacted: &crate::engine::CompactionOutput,
message_ids: &[i32],
) -> Result<(), String> {
let session_id = self
.selected_session
.ok_or_else(|| "The active session is unavailable.".to_owned())?;
let mut tail = compacted
.tail
.iter()
.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(),
})
.collect::<Vec<_>>();
if let Some((tools_session, tools)) = &self.agent_tools
let running_jobs = if let Some((tools_session, tools)) = &self.agent_tools
&& *tools_session == session_id
&& let Some(observation) = tools
{
tools
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.compaction_observation()
{
tail.push(crate::database::MessageDraft {
user: false,
tool: true,
system: false,
reasoning: None,
reasoning_complete: true,
content: observation,
});
}
} else {
None
};
let tail_start = message_ids.get(compacted.tail_start).copied();
let messages = self
.database
.as_mut()
.ok_or_else(|| "The project database is unavailable.".to_owned())?
.replace_with_compacted_transcript(
.record_compaction(
session_id,
&compacted.summary,
&tail,
tail_start,
running_jobs.as_deref(),
compacted.context_tokens,
self.context_limit,
)
.map_err(|error| format!("Could not save compacted conversation: {error}"))?;
self.conversation = messages.into_iter().map(ChatMessage::from).collect();
self.conversation
.extend(messages.into_iter().map(ChatMessage::from));
if let Some(session) = self
.projects
.iter_mut()
@@ -780,15 +1049,24 @@ impl App {
session.last_tokens_per_second = None;
}
let final_checkpoint = session_checkpoint_path(session_id);
fs::rename(&compacted.checkpoint, &final_checkpoint).map_err(|error| {
format!(
if let Err(error) = fs::rename(&compacted.checkpoint, &final_checkpoint) {
let _ = fs::remove_file(&compacted.checkpoint);
self.error = Some(format!(
"The compacted conversation was saved, but its checkpoint could not be promoted: {error}. It will rebuild on next use."
)
})?;
));
}
Ok(())
}
pub(super) fn compaction_summary(&self) -> Option<&str> {
if let Some(summary) = self
.conversation
.iter()
.rfind(|message| message.compaction)
.map(|message| message.content.as_str())
{
return Some(summary);
}
let session_id = self.selected_session?;
self.projects
.iter()
@@ -797,6 +1075,14 @@ impl App {
.and_then(|session| session.compacted_summary.as_deref())
}
fn model_chat_messages(&self) -> Vec<&ChatMessage> {
let start = compacted_context_start(&self.conversation);
self.conversation[start..]
.iter()
.filter(|message| !message.compaction)
.collect()
}
/// 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.
///
@@ -826,7 +1112,7 @@ impl App {
// pass summarize the user's opening message on its own.
let mut messages = stored
.into_iter()
.filter(|message| !message.content.trim().is_empty())
.filter(|message| !message.compaction && !message.content.trim().is_empty())
.map(|message| ChatTurn {
user: message.user,
tool: message.tool,
@@ -980,3 +1266,89 @@ pub(super) fn session_title(reply: &str) -> Option<String> {
None => title.to_owned(),
})
}
#[cfg(test)]
mod tests {
use super::{
ChatMessage, compacted_context_start, has_chat_after_last_compaction, queued_prompt,
};
#[test]
fn queued_guidance_is_one_reference_style_user_turn() {
assert_eq!(queued_prompt(Vec::new()), None);
assert_eq!(queued_prompt(["one".to_owned()]), Some("one".to_owned()));
assert_eq!(
queued_prompt(["one".to_owned(), "two".to_owned()]),
Some("Queued user message 1:\none\n\nQueued user message 2:\ntwo".to_owned())
);
}
#[test]
fn last_compaction_selects_its_tail_without_hiding_history() {
let message = |id: i32, compaction: bool, tail: Option<i32>| ChatMessage {
id,
user: !compaction,
tool: false,
system: compaction,
compaction,
compaction_tail_start: tail,
reasoning: None,
reasoning_complete: true,
reasoning_open: false,
content: format!("message {id}"),
markdown: Vec::new(),
};
let history = vec![
message(1, false, None),
message(2, false, None),
message(3, true, Some(2)),
message(4, false, None),
message(5, true, Some(4)),
message(6, false, None),
message(7, true, Some(6)),
message(8, false, None),
];
assert_eq!(history.len(), 8);
assert_eq!(compacted_context_start(&history), 5);
assert_eq!(history[compacted_context_start(&history)].id, 6);
}
#[test]
fn forced_compaction_requires_new_visible_chat() {
let message =
|id: i32, user: bool, tool: bool, system: bool, compaction: bool| ChatMessage {
id,
user,
tool,
system,
compaction,
compaction_tail_start: None,
reasoning: None,
reasoning_complete: true,
reasoning_open: false,
content: format!("message {id}"),
markdown: Vec::new(),
};
let mut history = vec![
message(1, true, false, false, false),
message(2, false, false, true, true),
ChatMessage {
content: format!(
"{} running job",
crate::agent::COMPACTION_OBSERVATION_PREFIX
),
..message(3, false, true, false, false)
},
message(4, false, false, true, false),
];
assert!(!has_chat_after_last_compaction(&history));
history.push(message(5, false, true, false, false));
assert!(has_chat_after_last_compaction(&history));
history.push(message(6, false, false, true, true));
assert!(!has_chat_after_last_compaction(&history));
history.push(message(7, true, false, false, false));
assert!(has_chat_after_last_compaction(&history));
}
}

9
src/app/view.rs
View File

@@ -441,13 +441,18 @@ impl App {
/// on the state the session is in.
fn session_menu_panel<'a>(&self, session: &'a Session) -> Element<'a, Message> {
let state = session.state();
let compact = menu_action(ICON_SPARK, "Compact context");
let compact = if self.can_compact_session(session.id) {
compact.on_press(Message::CompactSession(session.id))
} else {
compact
};
let mut actions = column![
menu_action(ICON_NEW_SESSION, "Rename session")
.on_press(Message::StartRenameSession(session.id)),
menu_action(ICON_SPARK, "Retitle with AI")
.on_press(Message::RetitleSession(session.id)),
menu_action(ICON_SPARK, "Compact context")
.on_press(Message::CompactSession(session.id)),
compact,
menu_action(ICON_ARCHIVE, "Rebuild context on next use")
.on_press(Message::RebuildSessionContext(session.id)),
]

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

@@ -56,21 +56,22 @@ impl App {
);
} else {
let markdown_style = markdown::Style::from_palette(app_theme().palette());
if let Some(summary) = self.compaction_summary() {
messages = messages.push(
container(
column![
text("Compacted task state").size(11),
text(summary).size(13).color(muted_text()),
]
.spacing(5),
)
.padding(14)
.width(Length::Fill)
.style(preference_group_style),
);
}
for (index, message) in self.conversation.iter().enumerate() {
if message.compaction {
messages = messages.push(
container(
column![
text("Context compacted").size(11),
text(&message.content).size(13).color(muted_text()),
]
.spacing(5),
)
.padding(14)
.width(Length::Fill)
.style(preference_group_style),
);
continue;
}
if message.system {
continue;
}
@@ -149,6 +150,14 @@ impl App {
.style(move |theme| chat_message_style(theme, user)),
);
}
if let Some(activity) = &self.activity {
messages = messages.push(
container(text(activity).size(13).color(muted_text()))
.padding(12)
.width(Length::Fill)
.style(preference_group_style),
);
}
}
// Tab walks every focusable widget of every window, so a composer
// left behind an open dialog would take a turn in that dialog's

61
src/compaction.rs
View File

@@ -2,6 +2,7 @@ use crate::engine::ChatTurn;
pub(crate) const SUMMARY_MAX_TOKENS: i32 = 4096;
pub(crate) const TOOL_RESULT_RESERVE_TOKENS: u32 = 1024;
const MIN_SUMMARY_TOKENS: u32 = 256;
const SOFT_PERCENT: u32 = 85;
const MIN_FREE_TOKENS: u32 = 8192;
@@ -21,24 +22,23 @@ pub(crate) fn should_compact(used: u32, limit: u32) -> bool {
|| limit.saturating_sub(used) <= MIN_FREE_TOKENS.min(limit / 8)
}
pub(crate) fn tool_result_needs_compaction(used: u32, limit: u32, result: &str) -> bool {
should_compact(used, limit)
|| used
.saturating_add(result.len().min(u32::MAX as usize) as u32)
.saturating_add(TOOL_RESULT_RESERVE_TOKENS)
>= limit
pub(crate) fn tool_result_reserve(limit: u32) -> u32 {
TOOL_RESULT_RESERVE_TOKENS.min((limit / 8).max(16))
}
pub(crate) fn bounded_tool_result(used: u32, limit: u32, result: String) -> String {
if used
.saturating_add(result.len().min(u32::MAX as usize) as u32)
.saturating_add(TOOL_RESULT_RESERVE_TOKENS)
< limit
{
result
} else {
"Tool error: the result is too large for the remaining context after compaction. Retry with a smaller read/search/bash output.\n".into()
}
pub(crate) fn tool_result_fits(projected: u32, limit: u32, reserve: u32) -> bool {
limit > 0 && projected.saturating_add(reserve) < limit
}
pub(crate) fn bounded_tool_error(projected: u32, limit: u32, reserve: u32) -> String {
format!(
"Tool error: tool result still does not fit after context compaction (projected_prompt={projected} tokens, ctx={limit}, reserve={reserve}). Retry with a smaller read/search/bash output.\n"
)
}
pub(crate) fn summary_budget(prompt: u32, limit: u32) -> Option<i32> {
let room = limit.saturating_sub(prompt).saturating_sub(1);
(room >= MIN_SUMMARY_TOKENS).then_some(room.min(SUMMARY_MAX_TOKENS as u32) as i32)
}
pub(crate) fn tail_budget(context: u32) -> u32 {
@@ -132,11 +132,30 @@ mod tests {
}
#[test]
fn oversized_tool_result_becomes_a_bounded_retry_error() {
let result = "x".repeat(4_000);
assert!(tool_result_needs_compaction(6_000, 10_000, &result));
let error = bounded_tool_result(4_000, 5_000, result);
fn a_new_summary_replaces_the_previous_rebuild_summary() {
let current = summary_system_prompt("tools", Some("old state"));
assert!(current.contains("old state"));
let rebuilt = summary_system_prompt("tools", Some("new state"));
assert!(rebuilt.contains("new state"));
assert!(!rebuilt.contains("old state"));
}
#[test]
fn hard_trigger_and_retry_reserve_match_the_reference() {
assert_eq!(tool_result_reserve(4096), 512);
assert_eq!(tool_result_reserve(65_536), 1024);
assert!(tool_result_fits(3000, 4096, 512));
assert!(!tool_result_fits(3584, 4096, 512));
let error = bounded_tool_error(5000, 4096, 512);
assert!(error.starts_with("Tool error:"));
assert!(error.len() < 160);
assert!(error.contains("projected_prompt=5000 tokens"));
assert!(error.len() < 256);
}
#[test]
fn summary_budget_keeps_answer_room_and_rejects_exhausted_contexts() {
assert_eq!(summary_budget(1000, 8192), Some(SUMMARY_MAX_TOKENS));
assert_eq!(summary_budget(7900, 8192), Some(291));
assert_eq!(summary_budget(7936, 8192), None);
}
}

162
src/database.rs
View File

@@ -119,6 +119,8 @@ pub struct StoredMessage {
pub reasoning_complete: bool,
pub content: String,
pub system: bool,
pub compaction: bool,
pub compaction_tail_start: Option<i32>,
}
#[derive(Insertable)]
@@ -131,15 +133,8 @@ struct NewMessage<'a> {
reasoning_complete: bool,
content: &'a str,
system: bool,
}
pub struct MessageDraft {
pub user: bool,
pub tool: bool,
pub reasoning: Option<String>,
pub reasoning_complete: bool,
pub content: String,
pub system: bool,
compaction: bool,
compaction_tail_start: Option<i32>,
}
#[derive(Debug)]
@@ -338,6 +333,8 @@ impl Database {
reasoning_complete: true,
content,
system: true,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?,
@@ -352,6 +349,8 @@ impl Database {
reasoning_complete: true,
content: prompt,
system: false,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?;
@@ -364,6 +363,8 @@ impl Database {
reasoning_complete: !reasoning,
content: "",
system: false,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?;
@@ -377,13 +378,13 @@ impl Database {
&mut self,
session_id: i32,
result: &str,
queued_users: &[String],
reminder: Option<&str>,
queued_user: Option<&str>,
system_messages: &[String],
reasoning: bool,
) -> Result<Vec<StoredMessage>, String> {
self.connection
.transaction(|connection| {
let mut stored = Vec::with_capacity(queued_users.len() + 3);
let mut stored = Vec::with_capacity(system_messages.len() + 3);
stored.push(
diesel::insert_into(messages::table)
.values(NewMessage {
@@ -394,11 +395,13 @@ impl Database {
reasoning_complete: true,
content: result,
system: false,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?,
);
for content in queued_users {
if let Some(content) = queued_user {
stored.push(
diesel::insert_into(messages::table)
.values(NewMessage {
@@ -409,12 +412,14 @@ impl Database {
reasoning_complete: true,
content,
system: false,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?,
);
}
if let Some(content) = reminder {
for content in system_messages {
stored.push(
diesel::insert_into(messages::table)
.values(NewMessage {
@@ -425,6 +430,8 @@ impl Database {
reasoning_complete: true,
content,
system: true,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?,
@@ -439,6 +446,8 @@ impl Database {
reasoning_complete: !reasoning,
content: "",
system: false,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?;
@@ -466,11 +475,12 @@ impl Database {
.map_err(|error| error.to_string())
}
pub fn replace_with_compacted_transcript(
pub fn record_compaction(
&mut self,
session_id: i32,
summary: &str,
tail: &[MessageDraft],
tail_start: Option<i32>,
running_jobs: Option<&str>,
context_used: u32,
context_limit: u32,
) -> Result<Vec<StoredMessage>, String> {
@@ -480,8 +490,6 @@ impl Database {
.map_err(|_| "Context limit is too large to save".to_owned())?;
self.connection
.transaction(|connection| {
diesel::delete(messages::table.filter(messages::session_id.eq(session_id)))
.execute(connection)?;
diesel::update(sessions::table.find(session_id))
.set((
sessions::compacted_summary.eq(Some(summary)),
@@ -490,18 +498,36 @@ impl Database {
sessions::last_tokens_per_second.eq(None::<f32>),
))
.execute(connection)?;
let mut stored = Vec::with_capacity(tail.len());
for message in tail {
let mut stored = Vec::with_capacity(2);
stored.push(
diesel::insert_into(messages::table)
.values(NewMessage {
session_id,
user: false,
tool: false,
reasoning: None,
reasoning_complete: true,
content: summary,
system: true,
compaction: true,
compaction_tail_start: tail_start,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?,
);
if let Some(content) = running_jobs {
stored.push(
diesel::insert_into(messages::table)
.values(NewMessage {
session_id,
user: message.user,
tool: message.tool,
reasoning: message.reasoning.as_deref(),
reasoning_complete: message.reasoning_complete,
content: &message.content,
system: message.system,
user: false,
tool: true,
reasoning: None,
reasoning_complete: true,
content,
system: false,
compaction: false,
compaction_tail_start: None,
})
.returning(StoredMessage::as_returning())
.get_result(connection)?,
@@ -586,7 +612,7 @@ mod tests {
}
#[test]
fn chat_messages_survive_reopen_and_follow_session_deletion() {
fn chat_and_compaction_history_survive_reopen_and_session_deletion() {
let id = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
@@ -606,8 +632,8 @@ mod tests {
.continue_tool_turn(
session.id,
"Tool result",
&["Queued correction".into()],
Some("Tool reminder"),
Some("Queued correction"),
&["Tool reminder".into()],
false,
)
.unwrap();
@@ -636,36 +662,86 @@ mod tests {
assert_eq!(messages[5].content, "Tool reminder");
assert!(!messages[6].user);
assert!(!messages[6].tool);
let compacted = reopened
.replace_with_compacted_transcript(
let first = reopened
.record_compaction(
session.id,
"Keep the active task.",
&[MessageDraft {
user: true,
tool: false,
reasoning: None,
reasoning_complete: true,
content: "Recent question".into(),
system: false,
}],
"First durable state.",
Some(messages[4].id),
Some("bash job=1 status=running"),
321,
32_768,
)
.unwrap();
assert_eq!(compacted.len(), 1);
assert_eq!(first.len(), 2);
assert!(first[0].compaction);
assert_eq!(first[0].compaction_tail_start, Some(messages[4].id));
drop(reopened);
let mut reopened = Database::open(&path).unwrap();
assert_eq!(
reopened.load_projects().unwrap()[0].sessions[0]
.compacted_summary
.as_deref(),
Some("Keep the active task.")
Some("First durable state.")
);
assert_eq!(
reopened.load_projects().unwrap()[0].sessions[0].context_used,
321
);
assert_eq!(reopened.load_messages(session.id).unwrap().len(), 1);
let messages = reopened.load_messages(session.id).unwrap();
assert_eq!(messages.len(), 9);
assert_eq!(messages[1].content, "Question");
assert!(messages[7].compaction);
assert_eq!(messages[7].content, "First durable state.");
assert_eq!(messages[8].content, "bash job=1 status=running");
reopened
.continue_tool_turn(session.id, "Reloaded tool result", None, &[], false)
.unwrap();
let continued = reopened.load_messages(session.id).unwrap();
let second_tail = continued[9].id;
reopened
.record_compaction(
session.id,
"Second durable state.",
Some(second_tail),
None,
222,
32_768,
)
.unwrap();
reopened
.continue_tool_turn(session.id, "Final tool result", None, &[], false)
.unwrap();
let continued = reopened.load_messages(session.id).unwrap();
let third_tail = continued[12].id;
reopened
.record_compaction(
session.id,
"Third durable state.",
Some(third_tail),
None,
111,
32_768,
)
.unwrap();
reopened
.start_chat_turn(session.id, "After third compaction", &[], false)
.unwrap();
drop(reopened);
let mut reopened = Database::open(&path).unwrap();
let history = reopened.load_messages(session.id).unwrap();
assert_eq!(history.len(), 17);
assert_eq!(history[1].content, "Question");
let markers = history
.iter()
.filter(|message| message.compaction)
.collect::<Vec<_>>();
assert_eq!(markers.len(), 3);
assert_eq!(markers[0].content, "First durable state.");
assert_eq!(markers[1].content, "Second durable state.");
assert_eq!(markers[2].content, "Third durable state.");
assert_eq!(markers[2].compaction_tail_start, Some(third_tail));
assert_eq!(history[15].content, "After third compaction");
reopened.delete_session(session.id).unwrap();
assert!(reopened.load_messages(session.id).unwrap().is_empty());
drop(reopened);

41
src/engine.rs
View File

@@ -364,7 +364,7 @@ pub(crate) struct GenerationOutput {
#[cfg(target_os = "macos")]
pub(crate) struct CompactionOutput {
pub(crate) summary: String,
pub(crate) tail: Vec<ChatTurn>,
pub(crate) tail_start: usize,
pub(crate) context_tokens: u32,
pub(crate) checkpoint: PathBuf,
}
@@ -530,19 +530,23 @@ impl Generator {
Ok(output)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn compact(
&mut self,
messages: &[ChatTurn],
settings: &TurnSettings,
rebuild_system_prompt: &str,
reason: &str,
checkpoint: &Path,
cancelled: &AtomicBool,
mut progress: impl FnMut(u32, u32, Option<f32>),
mut phase: impl FnMut(&'static str),
) -> Result<CompactionOutput, String> {
let _ = std::fs::remove_file(checkpoint);
self.executor.reset()?;
self.checkpoint = None;
let result = (|| {
phase("Compacting durable task state…");
let mut private_messages = messages.to_vec();
private_messages.push(ChatTurn {
user: true,
@@ -555,15 +559,24 @@ impl Generator {
});
let mut private_settings = settings.clone();
private_settings.reasoning_mode = ReasoningMode::Direct;
private_settings.max_generated_tokens = crate::compaction::SUMMARY_MAX_TOKENS;
let private_prompt = self.executor.model().render_conversation(
&private_settings.system_prompt,
&private_messages,
private_settings.reasoning_mode,
);
private_settings.max_generated_tokens = crate::compaction::summary_budget(
private_prompt.len().min(u32::MAX as usize) as u32,
self.executor.context(),
)
.ok_or_else(|| "not enough context left to request compaction summary".to_owned())?;
private_settings.temperature = 0.0;
private_settings.stops.extend([
private_settings.stops = vec![
"<｜DSML｜".into(),
"<DSML｜".into(),
"<tool_call>".into(),
"<think>".into(),
"</think>".into(),
]);
];
let (output, prompt_complete) = self.generate_inner(
&private_messages,
&private_settings,
@@ -580,6 +593,7 @@ impl Generator {
}
// The private request must never become the rebuilt session prefix.
phase("Rebuilding compacted context…");
self.executor.reset()?;
self.checkpoint = None;
@@ -609,7 +623,7 @@ impl Generator {
);
let tail = messages[start..].to_vec();
let rebuilt_system =
crate::compaction::summary_system_prompt(&settings.system_prompt, Some(&summary));
crate::compaction::summary_system_prompt(rebuild_system_prompt, Some(&summary));
let history_tokens = self.executor.model().render_history(
&rebuilt_system,
&tail,
@@ -627,10 +641,11 @@ impl Generator {
return Err("context compaction interrupted during rebuild".into());
}
let tag = conversation_tag(&rebuilt_system, settings.reasoning_mode, &tail);
phase("Saving compacted context…");
self.save_checkpoint(checkpoint, tag)?;
Ok(CompactionOutput {
summary,
tail,
tail_start: start,
context_tokens: history_tokens.len() as u32,
checkpoint: checkpoint.to_owned(),
})
@@ -645,6 +660,20 @@ impl Generator {
result
}
pub(crate) fn rendered_history_tokens(
&self,
messages: &[ChatTurn],
settings: &TurnSettings,
) -> Result<u32, String> {
u32::try_from(
self.executor
.model()
.render_history(&settings.system_prompt, messages, settings.reasoning_mode)
.len(),
)
.map_err(|_| "rendered conversation is too large".to_owned())
}
fn select_checkpoint(
&mut self,
checkpoint: &Path,

132
src/runtime.rs
View File

@@ -46,6 +46,7 @@ impl CheckpointTarget {
pub(crate) enum GenerationEvent {
Loading,
Activity(&'static str),
Chunk {
reasoning: bool,
content: String,
@@ -57,6 +58,23 @@ pub(crate) enum GenerationEvent {
},
Finished(Result<GenerationOutput, String>),
Compacted(Result<CompactionOutput, String>),
Measured(Result<u32, String>),
}
enum Operation {
Generate,
Compact {
reason: String,
rebuild_system_prompt: String,
},
Measure,
}
#[derive(Clone, Copy)]
enum ResponseKind {
Generation,
Compaction,
Measurement,
}
struct Command {
@@ -64,7 +82,7 @@ struct Command {
turn: TurnSettings,
messages: Vec<ChatTurn>,
checkpoint: CheckpointTarget,
compact_reason: Option<String>,
operation: Operation,
idle_timeout: Duration,
cancel: Arc<AtomicBool>,
events: Sender<GenerationEvent>,
@@ -100,7 +118,7 @@ impl GenerationService {
turn,
messages,
checkpoint,
compact_reason: None,
operation: Operation::Generate,
idle_timeout,
cancel: Arc::clone(&cancel),
events,
@@ -116,14 +134,46 @@ impl GenerationService {
})
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn compact(
&self,
engine: EngineSettings,
turn: TurnSettings,
messages: Vec<ChatTurn>,
reason: &str,
rebuild_system_prompt: String,
checkpoint: PathBuf,
idle_timeout: Duration,
) -> Result<ActiveGeneration, String> {
let cancel = Arc::new(AtomicBool::new(false));
let (events, receiver) = mpsc::channel();
self.commands
.send(Command {
engine,
turn,
messages,
checkpoint: CheckpointTarget::Local(checkpoint),
operation: Operation::Compact {
reason: reason.to_owned(),
rebuild_system_prompt,
},
idle_timeout,
cancel: Arc::clone(&cancel),
events,
})
.map_err(|_| "The model runtime stopped unexpectedly.".to_owned())?;
Ok(ActiveGeneration {
events: receiver,
cancel,
})
}
pub(crate) fn measure_context(
&self,
engine: EngineSettings,
turn: TurnSettings,
messages: Vec<ChatTurn>,
idle_timeout: Duration,
) -> Result<ActiveGeneration, String> {
let cancel = Arc::new(AtomicBool::new(false));
let (events, receiver) = mpsc::channel();
@@ -133,8 +183,8 @@ impl GenerationService {
engine,
turn,
messages,
checkpoint: CheckpointTarget::Local(checkpoint),
compact_reason: Some(reason.to_owned()),
checkpoint: CheckpointTarget::OneShot(PathBuf::new()),
operation: Operation::Measure,
idle_timeout,
cancel: Arc::clone(&cancel),
events,
@@ -157,7 +207,11 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
let request_started = Instant::now();
let source = command.checkpoint.source();
let events = command.events.clone();
metrics.request_started(source);
let response = response_kind(&command.operation);
let tracked = !matches!(command.operation, Operation::Measure);
if tracked {
metrics.request_started(source);
}
if let Err(error) = catch_runtime_panic(|| {
run_command(
command,
@@ -169,11 +223,13 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
&mut idle_timeout,
);
}) {
metrics.request_failed(request_started.elapsed());
if tracked {
metrics.request_failed(request_started.elapsed());
}
if loaded.take().is_some() {
metrics.unloaded();
}
let _ = events.send(GenerationEvent::Finished(Err(error)));
let _ = events.send(error_event(response, error));
}
}
Err(mpsc::RecvTimeoutError::Timeout) => {
@@ -197,12 +253,16 @@ fn run_command(
last_used: &mut Instant,
idle_timeout: &mut Duration,
) {
let response = response_kind(&command.operation);
let tracked = !matches!(command.operation, Operation::Measure);
*idle_timeout = command.idle_timeout;
if command.cancel.load(Ordering::Relaxed) {
metrics.request_failed(request_started.elapsed());
let _ = command.events.send(GenerationEvent::Finished(
Err("generation cancelled".into()),
));
if tracked {
metrics.request_failed(request_started.elapsed());
}
let _ = command
.events
.send(error_event(response, "generation cancelled".into()));
return;
}
if loaded
@@ -228,8 +288,10 @@ fn run_command(
Some((command.engine.clone(), generator))
}
Err(error) => {
metrics.request_failed(request_started.elapsed());
let _ = command.events.send(GenerationEvent::Finished(Err(error)));
if tracked {
metrics.request_failed(request_started.elapsed());
}
let _ = command.events.send(error_event(response, error));
None
}
};
@@ -261,17 +323,25 @@ fn run_command(
tokens_per_second,
});
};
if let Some(reason) = &command.compact_reason {
if let Operation::Compact {
reason,
rebuild_system_prompt,
} = &command.operation
{
let CheckpointTarget::Local(checkpoint) = &command.checkpoint else {
unreachable!("compaction checkpoints are local")
};
let result = generator.compact(
&command.messages,
&command.turn,
rebuild_system_prompt,
reason,
checkpoint,
&command.cancel,
&mut progress,
|activity| {
let _ = command.events.send(GenerationEvent::Activity(activity));
},
);
match &result {
Ok(_) => {
@@ -283,6 +353,12 @@ fn run_command(
*last_used = Instant::now();
return;
}
if matches!(command.operation, Operation::Measure) {
let result = generator.rendered_history_tokens(&command.messages, &command.turn);
let _ = command.events.send(GenerationEvent::Measured(result));
*last_used = Instant::now();
return;
}
let result = match command.checkpoint {
CheckpointTarget::Local(path) => generator.generate(
&path,
@@ -320,6 +396,22 @@ fn run_command(
}
}
fn response_kind(operation: &Operation) -> ResponseKind {
match operation {
Operation::Generate => ResponseKind::Generation,
Operation::Compact { .. } => ResponseKind::Compaction,
Operation::Measure => ResponseKind::Measurement,
}
}
fn error_event(kind: ResponseKind, error: String) -> GenerationEvent {
match kind {
ResponseKind::Generation => GenerationEvent::Finished(Err(error)),
ResponseKind::Compaction => GenerationEvent::Compacted(Err(error)),
ResponseKind::Measurement => GenerationEvent::Measured(Err(error)),
}
}
fn catch_runtime_panic<T>(operation: impl FnOnce() -> T) -> Result<T, String> {
// Generator owns GPU handles and is not unwind-safe; callers discard it on error.
std::panic::catch_unwind(std::panic::AssertUnwindSafe(operation))
@@ -351,4 +443,16 @@ mod tests {
);
assert_eq!(catch_runtime_panic(|| 42), Ok(42));
}
#[test]
fn operation_failures_use_the_matching_event() {
assert!(matches!(
error_event(ResponseKind::Compaction, "stopped".into()),
GenerationEvent::Compacted(Err(error)) if error == "stopped"
));
assert!(matches!(
error_event(ResponseKind::Measurement, "stopped".into()),
GenerationEvent::Measured(Err(error)) if error == "stopped"
));
}
}

2
src/schema.rs
View File

@@ -8,6 +8,8 @@ diesel::table! {
reasoning_complete -> Bool,
content -> Text,
system -> Bool,
compaction -> Bool,
compaction_tail_start -> Nullable<Integer>,
}
}

1
src/server/response.rs
View File

@@ -749,6 +749,7 @@ fn stream_response_with_keepalive(
"The model runtime returned an unexpected compaction event.".into(),
));
}
GenerationEvent::Activity(_) | GenerationEvent::Measured(_) => {}
GenerationEvent::Loading => {}
GenerationEvent::Context {
tokens_per_second, ..