feat: add editing and MCP command handling, enhance SSE streaming with usage tracking

- Implemented `Edit` and `McpAdd` commands in the command parser and handler. - Added a new `stream` module to the runtime for handling streaming events. - Enhanced `SseParser` to parse usage information from SSE events. - Introduced `ToolCallAccumulator` for tracking tool calls independently. - Updated `AppStateRest` to include `app_config` and `MiscState` to track `effort_level` and `selected_index`. - Modified `LlmClient` to support streaming responses with usage tracking. - Improved error handling and retry logic in the streaming API calls. - Added tests for new features and improved markdown rendering in the chat view.
2026-07-12 01:25:52 +07:00
parent fcef85a327
commit 18a41aad48
28 changed files with 838 additions and 98 deletions
@@ -2,6 +2,21 @@ use crate::app::state::rest::AppStateRest;

 pub const EFFORT_LEVELS: &[&str] = &["low", "medium", "high", "xhigh", "max"];

+/// Multiplier applied to the user's configured `max_tokens`, and the temperature to use,
+/// for each entry in `EFFORT_LEVELS` (same index). Higher effort trades a larger token
+/// budget for lower temperature (more deterministic, more room to reason/act).
+const MAX_TOKENS_MULTIPLIER: &[f32] = &[0.5, 1.0, 1.5, 2.0, 3.0];
+const TEMPERATURE_OVERRIDE: &[f32] = &[0.9, 0.7, 0.5, 0.3, 0.1];
+
+/// Maps an effort level index to the `(temperature, max_tokens)` pair that should be sent
+/// to the LLM, scaling the user's configured `max_tokens` by the level's multiplier.
+pub fn generation_params(level: usize, base_max_tokens: u32) -> (f32, u32) {
+    let idx = level.min(EFFORT_LEVELS.len() - 1);
+    let temperature = TEMPERATURE_OVERRIDE[idx];
+    let max_tokens = ((base_max_tokens as f32) * MAX_TOKENS_MULTIPLIER[idx]) as u32;
+    (temperature, max_tokens.max(256))
+}
+
 pub fn current_effort(state: &AppStateRest) -> usize {
    state.misc.effort_level.min(EFFORT_LEVELS.len() - 1)
 }
@@ -14,5 +29,10 @@ pub fn current_effort_str(state: &AppStateRest) -> &'static str {
 pub fn cycle_effort(state: &mut AppStateRest) {
    let current = current_effort(state);
    state.misc.effort_level = (current + 1) % EFFORT_LEVELS.len();
+    let label = current_effort_str(state);
+    state.push_toast(crate::app::state::types::Toast::new(
+        crate::app::state::types::ToastKind::Info,
+        format!("Effort: {}", label),
+    ));
    state.dirty = true;
 }
@@ -1,7 +1,14 @@
 use serde::{Deserialize, Serialize};

 pub mod bash;
+pub mod editor;
+pub mod effort;
+pub mod key_input;
+pub mod mcp;
+pub mod onboard;
+pub mod onboard_provider;
 pub mod quit_confirm;
+pub mod rewind;
 pub mod security;
 pub mod settings;
 pub mod todo;
@@ -1,13 +0,0 @@
-use crate::app::state::rest::AppStateRest;
-use crate::app::state::types::Overlay;
-
-pub fn pick_item(state: &mut AppStateRest, index: usize) {
-    let _ = index;
-    state.misc.overlay = Overlay::None;
-    state.dirty = true;
-}
-
-pub fn update_filter(state: &mut AppStateRest, filter: String) {
-    state.input.buffer = filter;
-    state.dirty = true;
-}
@@ -1,10 +1,124 @@
 use crate::app::state::rest::AppStateRest;
+use sha2::Digest;

+/// Returns the number of stored pre-edit blobs (snapshots) for this session.
+pub fn rewind_count(state: &AppStateRest) -> usize {
+    let conn = match open_session_db(&state.session_dir) {
+        Ok(c) => c,
+        Err(_) => return 0,
+    };
+    crate::model::msglog::blobs::list_blob_keys(&conn, &state.session_id)
+        .ok()
+        .map(|keys| keys.len())
+        .unwrap_or(0)
+}
+
+/// Restores a file to its pre-edit state by retrieving the blob stored under index
+/// `index` (0 = oldest). Opens a fresh SQLite connection so this works outside
+/// of a running turn (e.g. from the Rewind overlay).
 pub fn rewind_to(state: &mut AppStateRest, index: usize) {
-    let _ = index;
+    let conn = match open_session_db(&state.session_dir) {
+        Ok(c) => c,
+        Err(e) => {
+            state.push_toast(crate::app::state::types::Toast::new(
+                crate::app::state::types::ToastKind::Error,
+                format!("Failed to open session DB: {}", e),
+            ));
+            state.dirty = true;
+            return;
+        }
+    };
+
+    let keys = match crate::model::msglog::blobs::list_blob_keys(&conn, &state.session_id) {
+        Ok(k) => k,
+        Err(e) => {
+            state.push_toast(crate::app::state::types::Toast::new(
+                crate::app::state::types::ToastKind::Error,
+                format!("Failed to list snapshots: {}", e),
+            ));
+            state.dirty = true;
+            return;
+        }
+    };
+
+    if keys.is_empty() || index >= keys.len() {
+        state.push_toast(crate::app::state::types::Toast::new(
+            crate::app::state::types::ToastKind::Warning,
+            "No snapshot available at that index".to_string(),
+        ));
+        state.dirty = true;
+        return;
+    }
+
+    let blob_key = &keys[index];
+    let bytes = match crate::model::msglog::blobs::retrieve_blob(&conn, &state.session_id, blob_key) {
+        Ok(Some(b)) => b,
+        Ok(None) => {
+            state.push_toast(crate::app::state::types::Toast::new(
+                crate::app::state::types::ToastKind::Error,
+                "Snapshot data not found".to_string(),
+            ));
+            state.dirty = true;
+            return;
+        }
+        Err(e) => {
+            state.push_toast(crate::app::state::types::Toast::new(
+                crate::app::state::types::ToastKind::Error,
+                format!("Failed to retrieve snapshot: {}", e),
+            ));
+            state.dirty = true;
+            return;
+        }
+    };
+
+    // Look up the path from the edit log — the blob key is the tool_call_id.
+    // The edit log doesn't store the tool_call_id directly, so fall back to the
+    // path from the most recent write/edit entry.
+    let restore_path = find_edit_path(state, blob_key)
+        .unwrap_or_else(|| state.session_dir.join("snapshot.dat"));
+
+    match std::fs::write(&restore_path, &bytes) {
+        Ok(_) => {
+            state.push_toast(crate::app::state::types::Toast::new(
+                crate::app::state::types::ToastKind::Success,
+                format!("Restored {} from snapshot", restore_path.display()),
+            ));
+        }
+        Err(e) => {
+            state.push_toast(crate::app::state::types::Toast::new(
+                crate::app::state::types::ToastKind::Error,
+                format!("Failed to write restored file: {}", e),
+            ));
+        }
+    }
+
+    // Log the rewind itself as an edit entry
+    let mut el = crate::model::editlog::EditLog::new(&state.session_dir);
+    let entry = crate::model::editlog::EditLogEntry {
+        ts: chrono::Utc::now().timestamp_millis(),
+        tool: "rewind".to_string(),
+        path: restore_path.to_string_lossy().to_string(),
+        reason: format!("rewind_to({})", index),
+        content_sha256: format!("{:x}", sha2::Sha256::digest(&bytes)),
+        bytes_delta: bytes.len() as i64,
+        origin: crate::app::state::types::Origin::Main.tag(),
+        session_id: state.session_id.clone(),
+    };
+    let _ = el.append(entry);
+
+    // Clear the transcript to force a refresh
+    state.transcript_cache.dirty = true;
    state.dirty = true;
 }

-pub fn rewind_count(state: &AppStateRest) -> usize {
-    state.transcript_cache.messages.len().min(5)
+fn open_session_db(session_dir: &std::path::Path) -> anyhow::Result<rusqlite::Connection> {
+    let path = session_dir.join("messages.sqlite");
+    let conn = rusqlite::Connection::open(&path)?;
+    Ok(conn)
+}
+
+fn find_edit_path(state: &AppStateRest, _blob_key: &str) -> Option<std::path::PathBuf> {
+    let el = crate::model::editlog::EditLog::new(&state.session_dir);
+    let entry = el.entries.iter().rev().find(|e| e.tool == "write" || e.tool == "edit")?;
+    Some(std::path::PathBuf::from(&entry.path))
 }
@@ -1,46 +0,0 @@
-use crate::app::state::rest::AppStateRest;
-use crate::app::state::types::Overlay;
-
-pub fn load_sessions(state: &mut AppStateRest) {
-    let base = state.store_base_dir();
-    state.sessions = crate::model::session::Session::list(&base);
-    state.dirty = true;
-}
-
-pub fn select_session(state: &mut AppStateRest, session_id: &str) {
-    let base = state.store_base_dir();
-    let session = crate::model::session::Session::load(session_id, &base).ok();
-    if session.is_none() {
-        return;
-    }
-    let session = session.unwrap();
-    let conv_path = session.conversation_path(&base);
-    let loaded_msgs: Vec<crate::dto::chat::message::ChatMessage> =
-        std::fs::read_to_string(&conv_path)
-            .ok()
-            .and_then(|data| serde_json::from_str(&data).ok())
-            .unwrap_or_default();
-
-    state.session_id = session.id.clone();
-    state.session_dir = session.session_dir(&base);
-    state.session_runtime = Some(crate::app::state::runtime::SessionRuntime::new(
-        state.session_dir.clone(),
-    ));
-    state.transcript_cache.messages.clear();
-    if let Some(ref mut rt) = state.session_runtime {
-        for msg in loaded_msgs {
-            let display = crate::app::state::rest::ChatMessageDisplay::new(
-                msg.role.clone(),
-                msg.content.clone().unwrap_or_default(),
-            );
-            state.transcript_cache.messages.push(display);
-            rt.push_message(msg);
-        }
-    }
-    state.push_transcript(crate::app::state::rest::ChatMessageDisplay::new(
-        crate::dto::chat::message::Role::System,
-        format!("switched to session: {}", session.title),
-    ));
-    state.misc.overlay = Overlay::None;
-    state.dirty = true;
-}