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.
This commit is contained in:
+117
-3
@@ -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))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user