diff --git a/PLAN.md b/PLAN.md index 062a299..c7b2b43 100644 --- a/PLAN.md +++ b/PLAN.md @@ -23,106 +23,36 @@ execution targets one self-contained Mac. continuation, malformed DSML recovery, transient prefix-aware KV reuse, disconnect/queue cancellation, and shared model scheduling. The full automated C/Rust differential corpus remains open. -- Local sessions persist user, assistant, reasoning, and tool roles. The agent - executes the complete `ds4_agent.c` starting tool set and reinjects the same - model-specific tool contract when sessions are opened or continued. -- The remaining baseline gaps are long-running agent behavior and tool safety. - SSD streaming, speculative decoding, steering, GLM 5.2 execution, and - DeepSeek V4 Pro execution are not implemented in the Rust executor. Some - related catalog, validation, and preference plumbing already exists but must - not be treated as runtime support. +- Local sessions now match the long-running `ds4_agent.c` baseline: the full + starting tool set, unlimited tool rounds, queued user guidance between tool + rounds, session date/time context, periodic tool-contract reminders, + cooperative Stop, and explicit activity/failure states are implemented. +- Context compaction uses the reference soft and exact token-counted hard + triggers, private live-model summaries, bounded summary and tool-result + retries, a recent verbatim tail, running-job observations, and compatible KV + checkpoints. Every summary is a visible durable history marker carrying its + tail boundary; the full chat remains scrollable, while missing or + incompatible KV state rebuilds from the last marker, its tail, and later + messages. Manual compact is available after new chat following the latest + marker, alongside a checkpoint-discard/rebuild action. +- Focused coverage exercises triggers, summary bounds and sanitizing, tail + selection, queued guidance, checkpoint identity, running jobs, durable + compaction markers, relaunch, and continued tool work after rebuild. +- The next baseline gap is tool hardening and safety. SSD streaming, + speculative decoding, steering, GLM 5.2 execution, and DeepSeek V4 Pro + execution are not implemented in the Rust executor. Related catalog, + validation, and preference plumbing must not be treated as runtime support. ## Delivery order -1. **Next:** long-running agent stability and remaining `ds4_agent.c` parity. -2. Tool hardening, approvals, and productive tool presentation. -3. Remaining DS4 execution technology, starting with SSD streaming, then +1. **Next:** tool hardening, approvals, and productive tool presentation. +2. Remaining DS4 execution technology, starting with SSD streaming, then speculative decoding and the other Metal/runtime parity work. -4. Additional model execution: GLM 5.2 and DeepSeek V4 Pro. -5. Product completion, exhaustive parity verification, and distribution. -6. Optional extensions: Dev Brain and A2UI. +3. Additional model execution: GLM 5.2 and DeepSeek V4 Pro. +4. Product completion, exhaustive parity verification, and distribution. +5. Optional extensions: Dev Brain and A2UI. -## 1. Next — long-running agent stability and `ds4_agent.c` parity - -Goal: a local agent session must be able to run long read/edit/test loops, -cross the context limit repeatedly, survive interruption or relaunch, and -continue with the same durable task state as `ds4-agent`. - -### Context compaction - -- Port the reference soft trigger: compact before a user turn or tool - continuation at 85% context use, or when at most 8192 tokens remain, with the - free-token threshold capped to one eighth of small contexts. -- Port the hard trigger: before appending a tool result that would leave - insufficient answer room, compact once and retry. If it still does not fit, - return a bounded tool error that tells the model to request less output. -- Use the live model to generate an internal durable task-state summary. The - prompt and generated summary are private compaction work, never ordinary - user/assistant messages, and may not execute tools or retain thinking/DSML - control markup. -- Preserve goals, constraints, files touched, commands and important results, - decisions, known failures, and next steps. Prefer reloadable paths, ranges, - and commands over copying bulky data into the summary. -- Rebuild the model context exactly as the reference does: current system/tool - contract, durable summary, then a recent verbatim tail. Keep up to 10% of the - configured context as the tail, capped at 50000 tokens, and align it to a - user-turn boundary when possible. -- Generate at most the reference summary budget, stop at model control or tool - markers, and never let the private compaction exchange become the reusable - session prefix. - -### Durable transition and recovery - -- Treat transcript replacement and the new KV checkpoint as one logical - transition. Persist the compacted semantic transcript and checkpoint - metadata only after the rebuilt prefix is valid. -- On cancellation, summary failure, prefill failure, or application exit, keep - the previous durable transcript, invalidate any KV state contaminated by the - private compaction prompt, and make the next turn rebuild safely. -- Reopen a compacted session with the same summary and recent verbatim turns. - If its checkpoint is absent or incompatible, rebuild it from persisted - messages without changing the visible conversation. -- Preserve relevant live tool state across compaction. In particular, append a - compact observation for running shell jobs so the model can still inspect or - stop them after the context rebuild. -- Add a user-visible action equivalent to `/compact`, plus the reference - strip/rebuild behavior: discard a session KV payload without discarding its - transcript, then rebuild on the next use. -- Bind checkpoint compatibility to the model identity, quantization, context, - rendered transcript, and payload ABI. A model/configuration change must - rebuild rather than reuse an invalid prefix. - -### Long-turn behavior - -- Accept user input while an assistant/tool loop is active. Queue it visibly - and inject it after the current tool result, before the next assistant - continuation, matching `ds4_agent.c` instead of starting a competing turn. -- Preserve the reference rule that there is no arbitrary maximum tool-round - count. Completion, Stop, context pressure, or a real error ends the loop. -- Match the reference date/time context injection and periodic system/tool - prompt reminder so long or reopened sessions do not drift away from the tool - contract. -- Keep Stop cooperative across summarization, compacted-prefix prefill, - generation, and active tool work. A stop must always leave a transcript that - can be reopened. - -### Presentation and verification - -- Show explicit `Compacting`, rebuilding/prefill, queued-input, stopped, and - failed states without blocking the Iced event loop. Keep the compacted - summary inspectable without presenting the private prompt as user history. -- Add focused tests for soft and hard triggers, tail selection, tool-result - retry, cancellation rollback, checkpoint invalidation, queued input between - tool rounds, running-job preservation, and reopen after compaction. -- Add a reference fixture that runs a long tool loop through compaction, saves, - relaunches, and continues without losing the active task. - -Exit criterion: repeat the reference `ds4-agent` long-context scenarios, -including compaction forced by a large tool result, stop during compaction, and -restart after a successful compaction. The same task state, recent turns, tool -contract, and running-job awareness must remain available. - -## 2. Tool hardening and safety +## 1. Next — tool hardening and safety Goal: make the existing tool set safe and clear enough for productive daily use without weakening its ability to inspect, edit, build, and test a project. @@ -159,13 +89,13 @@ Exit criterion: use the agent for a real inspect/edit/test cycle while every side effect is visible, risky actions require consent, Stop works at every stage, and no file tool can escape the selected project. -## 3. DS4 execution technology parity +## 2. DS4 execution technology parity Goal: finish the model-independent Metal/runtime capabilities in `ds4.c` before adding larger model families. Every capability must be shared by local chat and the HTTP endpoint through the single process-wide model owner. -### 3.1 SSD streaming — first runtime priority +### 2.1 SSD streaming — first runtime priority SSD streaming is the capacity prerequisite for larger models and therefore comes before GLM 5.2 and DeepSeek V4 Pro execution. @@ -186,7 +116,7 @@ comes before GLM 5.2 and DeepSeek V4 Pro execution. the cache and I/O layer model-aware so later GLM/Pro milestones add policy and graph support rather than a second streaming subsystem. -### 3.2 Speculative decoding: legacy MTP and DSpark +### 2.2 Speculative decoding: legacy MTP and DSpark - Load and validate the optional Flash legacy-MTP or DSpark support GGUF without treating either as a standalone model. Preserve exact support-kind and target @@ -206,7 +136,7 @@ comes before GLM 5.2 and DeepSeek V4 Pro execution. - GLM's in-model MTP path belongs to the GLM milestone, but it should reuse the verifier/session machinery established here. -### 3.3 Remaining Metal execution controls +### 2.3 Remaining Metal execution controls - Port directional steering files and exact FFN/attention application, including DS4 defaults, validation, zero-scale behavior, and checkpoint/model @@ -218,7 +148,7 @@ comes before GLM 5.2 and DeepSeek V4 Pro execution. - Add hardware-backed token/activation fixtures for each mode and keep the ordinary resident Flash path unchanged when optional features are off. -### 3.4 Single-machine server batching +### 2.4 Single-machine server batching - Port DS4's resident multi-session batching and server scheduling only after the serialized path remains the correctness oracle. Preserve per-request @@ -233,7 +163,7 @@ resident, SSD-streamed, MTP, DSpark, steering, and batched-server configurations, with optional modes off producing the same baseline behavior as today. -## 4. Additional model execution +## 3. Additional model execution Start these only after the shared capacity and execution technology above is stable. Catalog entries, settings, tokenizer work, or GGUF validation alone do @@ -269,7 +199,7 @@ Exit criterion: each advertised model passes the same local-agent, checkpoint, HTTP, SSD-capacity, cancellation, and deterministic token-output matrix as DeepSeek V4 Flash. -## 5. Product completion and verification +## 4. Product completion and verification ### Reference parity and regression coverage @@ -301,7 +231,7 @@ Exit criterion: a notarized build can be installed on a clean supported Mac, run the full local-agent and endpoint smoke matrix, restart into its previous sessions, and update without losing projects, transcripts, models, or KV data. -## 6. Optional future extensions +## 5. Optional future extensions These are not DS4 baseline parity and must not delay the milestones above. diff --git a/README.md b/README.md index af42b8b..a27bf65 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,12 @@ persisted as transcript roles and automatically continue the same model turn. File tools stay inside the selected project. Web tools ask before starting a visible Chrome profile. +Long sessions compact automatically while retaining the complete scrollable +chat. Each compaction appears in history with its durable summary, and a missing +or incompatible KV checkpoint rebuilds from the latest summary and subsequent +chat. An idle session's `…` menu can force compaction after new chat has been +added since the latest marker, or discard its checkpoint for a clean rebuild. + The app also listens on `127.0.0.1:4000` by default for Models, Chat Completions, Completions, Anthropic Messages, and Responses APIs. The listener, port, and opt-in CORS are configurable in Preferences. The diff --git a/migrations/20260726220000_add_compaction_messages/down.sql b/migrations/20260726220000_add_compaction_messages/down.sql new file mode 100644 index 0000000..404a809 --- /dev/null +++ b/migrations/20260726220000_add_compaction_messages/down.sql @@ -0,0 +1,2 @@ +ALTER TABLE messages DROP COLUMN compaction_tail_start; +ALTER TABLE messages DROP COLUMN compaction; diff --git a/migrations/20260726220000_add_compaction_messages/up.sql b/migrations/20260726220000_add_compaction_messages/up.sql new file mode 100644 index 0000000..445fab3 --- /dev/null +++ b/migrations/20260726220000_add_compaction_messages/up.sql @@ -0,0 +1,2 @@ +ALTER TABLE messages ADD COLUMN compaction BOOLEAN NOT NULL DEFAULT FALSE; +ALTER TABLE messages ADD COLUMN compaction_tail_start INTEGER; diff --git a/src/agent.rs b/src/agent.rs index b9b4c2d..9d5dbd6 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -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:")); diff --git a/src/app.rs b/src/app.rs index f7e5c2a..3b90ae5 100644 --- a/src/app.rs +++ b/src/app.rs @@ -97,6 +97,8 @@ pub(crate) struct App { #[cfg(target_os = "macos")] active_compaction: Option, #[cfg(target_os = "macos")] + active_tool_check: Option, + #[cfg(target_os = "macos")] agent_tools: Option<(i32, Arc>)>, #[cfg(target_os = "macos")] active_tools: Option, @@ -110,6 +112,7 @@ pub(crate) struct App { pub(super) activity: Option, 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, diff --git a/src/app/generation.rs b/src/app/generation.rs index e675e05..db5f6f6 100644 --- a/src/app/generation.rs +++ b/src/app/generation.rs @@ -29,6 +29,22 @@ pub(super) enum PendingContinuation { pub(super) struct CompactionRequest { pub(super) active: ActiveGeneration, pending: PendingContinuation, + message_ids: Vec, +} + +#[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, pub(super) reasoning: Option, pub(super) reasoning_complete: bool, pub(super) reasoning_open: bool, @@ -74,6 +92,8 @@ impl From 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 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) -> Option { + let messages = messages.into_iter().collect::>(); + 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::>() + .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::>(); #[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 { + 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::>(); - 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::>(); + 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::>(); - 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 { 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| 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)); + } +} diff --git a/src/app/view.rs b/src/app/view.rs index d107cd0..81478ec 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -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)), ] diff --git a/src/app/view/chat.rs b/src/app/view/chat.rs index e447874..fa8d38c 100644 --- a/src/app/view/chat.rs +++ b/src/app/view/chat.rs @@ -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 diff --git a/src/compaction.rs b/src/compaction.rs index d0028da..af529df 100644 --- a/src/compaction.rs +++ b/src/compaction.rs @@ -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 { + 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); } } diff --git a/src/database.rs b/src/database.rs index 978a7ae..806aace 100644 --- a/src/database.rs +++ b/src/database.rs @@ -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, } #[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, - pub reasoning_complete: bool, - pub content: String, - pub system: bool, + compaction: bool, + compaction_tail_start: Option, } #[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, 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, + running_jobs: Option<&str>, context_used: u32, context_limit: u32, ) -> Result, 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::), )) .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::>(); + 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); diff --git a/src/engine.rs b/src/engine.rs index 8e06d95..e20a80f 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -364,7 +364,7 @@ pub(crate) struct GenerationOutput { #[cfg(target_os = "macos")] pub(crate) struct CompactionOutput { pub(crate) summary: String, - pub(crate) tail: Vec, + 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), + mut phase: impl FnMut(&'static str), ) -> Result { 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(), "".into(), "".into(), "".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::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, diff --git a/src/runtime.rs b/src/runtime.rs index 74f9eb7..46fe5f9 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -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), Compacted(Result), + Measured(Result), +} + +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, checkpoint: CheckpointTarget, - compact_reason: Option, + operation: Operation, idle_timeout: Duration, cancel: Arc, events: Sender, @@ -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, reason: &str, + rebuild_system_prompt: String, checkpoint: PathBuf, idle_timeout: Duration, + ) -> Result { + 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, + idle_timeout: Duration, ) -> Result { 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, metrics: Arc) { 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, metrics: Arc) { &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(operation: impl FnOnce() -> T) -> Result { // 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" + )); + } } diff --git a/src/schema.rs b/src/schema.rs index 4950ea3..8908f9e 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -8,6 +8,8 @@ diesel::table! { reasoning_complete -> Bool, content -> Text, system -> Bool, + compaction -> Bool, + compaction_tail_start -> Nullable, } } diff --git a/src/server/response.rs b/src/server/response.rs index 02e0135..fd786fe 100644 --- a/src/server/response.rs +++ b/src/server/response.rs @@ -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, ..