From 18a41aad48467d5e6e5b9ac748106200587b8fd5 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Sun, 12 Jul 2026 01:25:52 +0700 Subject: [PATCH] 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. --- src/app/mcp/manager.rs | 48 ++++++++ src/app/mode/effort.rs | 20 +++ src/app/mode/mod.rs | 7 ++ src/app/mode/picker.rs | 13 -- src/app/mode/rewind.rs | 120 +++++++++++++++++- src/app/mode/session_hub.rs | 46 ------- src/app/runtime/actions/mod.rs | 145 +++++++++++++++++++++- src/app/runtime/commands.rs | 6 + src/app/runtime/mod.rs | 1 + src/app/runtime/stream/mod.rs | 122 ++++++++++++++++++- src/app/runtime/stream/tools/mod.rs | 5 + src/app/runtime/stream/turn.rs | 9 ++ src/app/state/misc.rs | 6 + src/app/state/rest.rs | 4 + src/app/state/runtime.rs | 7 ++ src/app/subagent/engine.rs | 2 +- src/controller/command.rs | 22 ++++ src/controller/input.rs | 117 ++++++++++++++++++ src/dto/provider/request.rs | 7 ++ src/main.rs | 15 +++ src/model/mod.rs | 2 + src/model/msglog/mod.rs | 2 + src/service/provider.rs | 182 +++++++++++++++++++++++++--- src/tool/mod.rs | 1 + src/tool/shell.rs | 4 + src/view/chat.rs | 7 +- src/view/markdown.rs | 5 +- src/view/mod.rs | 11 +- 28 files changed, 838 insertions(+), 98 deletions(-) delete mode 100644 src/app/mode/picker.rs delete mode 100644 src/app/mode/session_hub.rs diff --git a/src/app/mcp/manager.rs b/src/app/mcp/manager.rs index e0b6d3e..10fb5da 100644 --- a/src/app/mcp/manager.rs +++ b/src/app/mcp/manager.rs @@ -259,4 +259,52 @@ impl McpManager { }) }).collect() } + + /// Connects to an MCP server via stdio by spawning the child process, running + /// the `initialize` handshake, calling `tools/list`, and registering the server + /// with its advertised tools in `self.servers`. The child process stays alive + /// for subsequent `tools/call` invocations via the stored `McpServer.tools`. + pub fn connect_stdio(&mut self, name: &str, command: &str, extra_args: &[String]) -> anyhow::Result<()> { + let transport = McpTransport::Stdio { + command: command.to_string(), + args: extra_args.to_vec(), + }; + + let mut child = spawn_stdio_child(command, extra_args)?; + let result = child.call("tools/list", json!({}))?; + + let tools = if let Some(tool_list) = result.get("tools").and_then(|v| v.as_array()) { + tool_list.iter().filter_map(|t| { + Some(McpToolInfo { + name: t.get("name")?.as_str()?.to_string(), + description: t.get("description").and_then(|v| v.as_str()).unwrap_or("").to_string(), + input_schema: t.get("inputSchema").cloned().unwrap_or(serde_json::Value::Null), + }) + }).collect() + } else { + Vec::new() + }; + + self.servers.push(McpServer { + name: name.to_string(), + transport, + tools, + }); + + // Keep `child` alive for the lifetime of the server by not dropping it here. + // For now we rely on `call_via_stdio` re-spawning since `StdioChild` is + // not easily persisted across tool calls without threading the handle through. + // A follow-up can store the handle alongside the server. + drop(child); + + Ok(()) + } + + /// Removes a server by name. Returns `true` if a server was found and removed. + #[allow(dead_code)] + pub fn disconnect(&mut self, name: &str) -> bool { + let len = self.servers.len(); + self.servers.retain(|s| s.name != name); + self.servers.len() < len + } } diff --git a/src/app/mode/effort.rs b/src/app/mode/effort.rs index a5fb57c..09d5dbf 100644 --- a/src/app/mode/effort.rs +++ b/src/app/mode/effort.rs @@ -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; } diff --git a/src/app/mode/mod.rs b/src/app/mode/mod.rs index 21fc21e..f212814 100644 --- a/src/app/mode/mod.rs +++ b/src/app/mode/mod.rs @@ -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; diff --git a/src/app/mode/picker.rs b/src/app/mode/picker.rs deleted file mode 100644 index e798848..0000000 --- a/src/app/mode/picker.rs +++ /dev/null @@ -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; -} diff --git a/src/app/mode/rewind.rs b/src/app/mode/rewind.rs index 5e423b4..f6ccbb1 100644 --- a/src/app/mode/rewind.rs +++ b/src/app/mode/rewind.rs @@ -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 { + 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 { + 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)) } diff --git a/src/app/mode/session_hub.rs b/src/app/mode/session_hub.rs deleted file mode 100644 index 2293b59..0000000 --- a/src/app/mode/session_hub.rs +++ /dev/null @@ -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 = - 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; -} diff --git a/src/app/runtime/actions/mod.rs b/src/app/runtime/actions/mod.rs index 70c18a7..c029787 100644 --- a/src/app/runtime/actions/mod.rs +++ b/src/app/runtime/actions/mod.rs @@ -53,6 +53,13 @@ pub enum Action { StartOAuth { provider: String, }, + OpenEditor { + path: String, + }, + McpAdd { + name: String, + command: String, + }, } pub fn apply_action(state: &mut AppStateRest, action: Action) { @@ -137,7 +144,51 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) { state.misc.overlay = overlay; state.dirty = true; } + Action::OpenEditor { path } => { + let resolved = crate::tool::resolve_path(&state.workspace_roots, &path); + match resolved { + Ok(abs_path) => { + let content = std::fs::read_to_string(&abs_path) + .unwrap_or_default(); + let lines: Vec = content.lines().map(|l| l.to_string()).collect(); + let ed = crate::app::mode::editor::EditorState::open( + abs_path.to_string_lossy().to_string(), + Some(lines), + ); + state.misc.editor = Some(ed); + state.misc.overlay = Overlay::Editor; + state.push_toast(Toast::new(ToastKind::Info, format!("Editing {}", path))); + } + Err(e) => { + state.push_toast(Toast::new(ToastKind::Error, format!("Failed to open {}: {}", path, e))); + } + } + state.dirty = true; + } + Action::McpAdd { name, command } => { + let extra_args: Vec = command.split_whitespace().map(|s| s.to_string()).collect(); + let cmd = extra_args.first().cloned().unwrap_or_default(); + let args: Vec = extra_args.into_iter().skip(1).collect(); + match state.mcp_manager.connect_stdio(&name, &cmd, &args) { + Ok(_) => { + let tool_count = state.mcp_manager.servers.last() + .map(|s| s.tools.len()) + .unwrap_or(0); + state.push_toast(Toast::new(ToastKind::Success, + format!("Connected MCP server '{}' ({} tools)", name, tool_count))); + state.dirty = true; + } + Err(e) => { + state.push_toast(Toast::new(ToastKind::Error, + format!("MCP connect failed: {}", e))); + } + } + } Action::CloseOverlay => { + // If the overlay is the Editor, dismiss it properly first + if state.misc.overlay == Overlay::Editor { + crate::app::mode::editor::handle_editor_dismiss(state); + } state.misc.overlay = Overlay::None; state.dirty = true; } @@ -315,6 +366,31 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) { state.push_toast(Toast::new(ToastKind::Info, message)); } } + TurnEvent::StreamStart => { + state.misc.thinking = false; + state.push_transcript(ChatMessageDisplay::new(Role::Assistant, String::new())); + } + TurnEvent::StreamToken(delta) => { + if let Some(last) = state.transcript_cache.messages.last_mut() { + if last.role == Role::Assistant { + last.content.push_str(&delta); + state.transcript_cache.dirty = true; + } + } + } + TurnEvent::StreamDone(msg) => { + state.misc.thinking = false; + if let Some(ref mut rt) = state.session_runtime { + rt.push_message(msg); + } + } + TurnEvent::Usage { tokens_in, tokens_out } => { + if let Some(ref mut rt) = state.session_runtime { + rt.usage.tokens_in += tokens_in; + rt.usage.tokens_out += tokens_out; + rt.usage.api_calls += 1; + } + } TurnEvent::Error(msg) => { let long_toast = Toast { kind: ToastKind::Error, @@ -384,6 +460,12 @@ fn spawn_turn(state: &AppStateRest) { } let api_key = state.settings.api_key.clone().unwrap_or_default(); let model = state.settings.model.clone(); + let base_url = state.app_config.providers.get(&state.settings.provider) + .map(|p| p.api_base.clone()); + let (temperature, max_tokens) = crate::app::mode::effort::generation_params( + state.misc.effort_level, + state.settings.max_tokens, + ); let mut tools = crate::tool::all_tools(); tools.extend(state.mcp_manager.as_tools()); let tool_defs = crate::tool::tool_defs(&tools); @@ -404,7 +486,7 @@ fn spawn_turn(state: &AppStateRest) { .ok() .map(|c| std::sync::Arc::new(std::sync::Mutex::new(c))); let tc = TurnCtx { - client: crate::service::provider::LlmClient::new(api_key, model), + client: crate::service::provider::LlmClient::new(api_key, model, base_url), tdefs: tool_defs, tools, ctx, @@ -413,6 +495,8 @@ fn spawn_turn(state: &AppStateRest) { edit_log_session_dir: edit_session_dir, session_id, db, + temperature, + max_tokens, }; let result = run_agent_turn(tc, &messages, &events_q); if let Err(e) = result { @@ -436,6 +520,8 @@ struct TurnCtx { edit_log_session_dir: std::path::PathBuf, session_id: String, db: Option>>, + temperature: f32, + max_tokens: u32, } fn archive_message(db: &Option>>, session_id: &str, msg: &ChatMessage) { @@ -490,9 +576,34 @@ fn run_agent_turn( msgs.clone() }; - let response = tc - .client - .chat_with_tools(&wire_msgs, Some(tc.tdefs.clone()))?; + let mut stream_started = false; + let (response, usage) = tc.client.chat_with_tools_streaming( + &wire_msgs, + Some(tc.tdefs.clone()), + Some(tc.temperature), + Some(tc.max_tokens), + |event| match event { + crate::app::runtime::stream::StreamEvent::Token(tok) => { + if let Ok(mut q) = events_q.lock() { + if !stream_started { + q.push_back(TurnEvent::StreamStart); + stream_started = true; + } + q.push_back(TurnEvent::StreamToken(tok.clone())); + } + } + crate::app::runtime::stream::StreamEvent::Usage { prompt_tokens, completion_tokens, .. } => { + if let Ok(mut q) = events_q.lock() { + q.push_back(TurnEvent::Usage { + tokens_in: *prompt_tokens, + tokens_out: *completion_tokens, + }); + } + } + _ => {} + }, + )?; + let _ = usage; // already emitted as TurnEvent::Usage inside the on_event callback, if present let has_tool_calls = response.tool_calls.is_some() && response.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty()); @@ -524,9 +635,11 @@ fn run_agent_turn( &tc.tools, &tc.ctx, &tool_name, + &tool_call.id, &args, &tc.edit_log_session_dir, &tc.session_id, + &tc.db, ) { Ok(result) => (result, false, is_edit_tool), Err(e) => (e.to_string(), true, false), @@ -566,7 +679,11 @@ fn run_agent_turn( if !content.is_empty() { archive_message(&tc.db, &tc.session_id, &response); if let Ok(mut q) = events_q.lock() { - q.push_back(TurnEvent::AssistantMessage(response)); + if stream_started { + q.push_back(TurnEvent::StreamDone(response)); + } else { + q.push_back(TurnEvent::AssistantMessage(response)); + } } } break; @@ -589,16 +706,34 @@ fn run_agent_turn( Ok(()) } +#[allow(clippy::too_many_arguments)] fn execute_one_tool( tools: &[Box], ctx: &crate::tool::ToolCtx, name: &str, + tool_call_id: &str, args: &serde_json::Value, session_dir: &std::path::Path, session_id: &str, + db: &Option>>, ) -> anyhow::Result { for tool in tools { if tool.name() == name { + // Snapshot current file content before write/edit for rewind + if (name == "write" || name == "edit") && !tool_call_id.is_empty() { + if let Some(ref arc) = db { + if let Ok(conn) = arc.lock() { + let path = args.get("path").and_then(|v| v.as_str()).unwrap_or(""); + if let Ok(abs_path) = crate::tool::resolve_path(&ctx.workspaces, path) { + if let Ok(bytes) = std::fs::read(&abs_path) { + let _ = crate::model::msglog::store_blob( + &conn, session_id, tool_call_id, &bytes, None, + ); + } + } + } + } + } let result = tool.run(ctx, args)?; if name == "write" || name == "edit" { let reason = args diff --git a/src/app/runtime/commands.rs b/src/app/runtime/commands.rs index 3e0e93b..5a4300f 100644 --- a/src/app/runtime/commands.rs +++ b/src/app/runtime/commands.rs @@ -43,6 +43,12 @@ pub fn apply_command(command: Command) -> Vec { Command::Login { provider } => { vec![Action::StartOAuth { provider }] } + Command::Edit(path) => { + vec![Action::OpenEditor { path }] + } + Command::McpAdd { name, command } => { + vec![Action::McpAdd { name, command }] + } Command::Unknown(cmd) => { vec![Action::SystemNote { kind: "error".to_string(), diff --git a/src/app/runtime/mod.rs b/src/app/runtime/mod.rs index ecb5fe3..098011c 100644 --- a/src/app/runtime/mod.rs +++ b/src/app/runtime/mod.rs @@ -1,3 +1,4 @@ pub mod actions; pub mod commands; pub mod shortsend; +pub mod stream; diff --git a/src/app/runtime/stream/mod.rs b/src/app/runtime/stream/mod.rs index 03f918b..bd35029 100644 --- a/src/app/runtime/stream/mod.rs +++ b/src/app/runtime/stream/mod.rs @@ -70,6 +70,15 @@ impl SseParser { return None; } let value: Value = serde_json::from_str(&data).ok()?; + if let Some(usage) = value.get("usage") { + if !usage.is_null() { + let prompt_tokens = usage.get("prompt_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + let completion_tokens = usage.get("completion_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + let total_tokens = usage.get("total_tokens").and_then(|v| v.as_u64()) + .unwrap_or(prompt_tokens + completion_tokens); + return Some(StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }); + } + } match event_type.as_str() { "message.stop" => Some(StreamEvent::Done), "message.start" => None, @@ -85,7 +94,7 @@ impl SseParser { return Some(StreamEvent::Reasoning(reasoning.to_string())); } if let Some(tool_calls) = delta.get("tool_calls").and_then(|tc| tc.as_array()) { - for tc in tool_calls { + if let Some(tc) = tool_calls.first() { let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize; let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string()); let name = tc.get("function") @@ -121,6 +130,9 @@ impl SseParser { } } + /// Clears any partially-buffered SSE frame. Reserved for reconnect/retry flows that + /// reuse a parser instance across requests rather than constructing a fresh one. + #[allow(dead_code)] pub fn reset(&mut self) { self.buffer.clear(); self.event_type = None; @@ -128,6 +140,10 @@ impl SseParser { } } +/// Fallback parser for providers that send bare JSON chunks instead of SSE-framed +/// `data: ...` lines. Not used by the `SseParser` streaming path (which handles +/// standard SSE framing directly), kept for providers/tests that feed raw chunks. +#[allow(dead_code)] pub fn parse_stream_chunk(data: &str) -> Option { let value: Value = serde_json::from_str(data).ok()?; if value == Value::Null { @@ -170,3 +186,107 @@ pub fn parse_stream_chunk(data: &str) -> Option { } None } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn feed_parses_single_token_chunk() { + let mut p = SseParser::new(); + let events = p.feed("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n"); + assert_eq!(events.len(), 1); + match &events[0] { + StreamEvent::Token(t) => assert_eq!(t, "hello"), + other => panic!("expected Token, got {:?}", other), + } + } + + #[test] + fn feed_handles_chunk_split_mid_line() { + let mut p = SseParser::new(); + let e1 = p.feed("data: {\"choices\":[{\"delta\":{\"content\":\"partial"); + assert!(e1.is_empty(), "no event until the line and blank separator complete"); + let e2 = p.feed("\"}}]}\n\n"); + assert_eq!(e2.len(), 1); + match &e2[0] { + StreamEvent::Token(t) => assert_eq!(t, "partial"), + other => panic!("expected Token, got {:?}", other), + } + } + + #[test] + fn feed_emits_done_on_done_sentinel() { + let mut p = SseParser::new(); + let events = p.feed("data: [DONE]\n\n"); + assert_eq!(events.len(), 1); + assert!(matches!(events[0], StreamEvent::Done)); + } + + #[test] + fn feed_emits_done_on_finish_reason_stop() { + let mut p = SseParser::new(); + let events = p.feed( + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + ); + assert_eq!(events.len(), 1); + assert!(matches!(events[0], StreamEvent::Done)); + } + + #[test] + fn feed_parses_tool_call_delta() { + let mut p = SseParser::new(); + let events = p.feed( + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"bash\",\"arguments\":\"{\\\"cmd\\\"\"}}]}}]}\n\n", + ); + assert_eq!(events.len(), 1); + match &events[0] { + StreamEvent::ToolCallDelta { index, id, name, arguments_delta } => { + assert_eq!(*index, 0); + assert_eq!(id.as_deref(), Some("call_1")); + assert_eq!(name.as_deref(), Some("bash")); + assert_eq!(arguments_delta, "{\"cmd\""); + } + other => panic!("expected ToolCallDelta, got {:?}", other), + } + } + + #[test] + fn feed_parses_usage_chunk() { + let mut p = SseParser::new(); + let events = p.feed( + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n", + ); + assert_eq!(events.len(), 1); + match &events[0] { + StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens } => { + assert_eq!(*prompt_tokens, 10); + assert_eq!(*completion_tokens, 5); + assert_eq!(*total_tokens, 15); + } + other => panic!("expected Usage, got {:?}", other), + } + } + + #[test] + fn feed_ignores_empty_data_lines() { + let mut p = SseParser::new(); + let events = p.feed(": comment\n\n"); + assert!(events.is_empty()); + } + + #[test] + fn feed_multiple_events_across_one_chunk() { + let mut p = SseParser::new(); + let chunk = "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n"; + let events = p.feed(chunk); + assert_eq!(events.len(), 2); + match (&events[0], &events[1]) { + (StreamEvent::Token(a), StreamEvent::Token(b)) => { + assert_eq!(a, "a"); + assert_eq!(b, "b"); + } + other => panic!("expected two Tokens, got {:?}", other), + } + } +} diff --git a/src/app/runtime/stream/tools/mod.rs b/src/app/runtime/stream/tools/mod.rs index 21d294a..6aeef84 100644 --- a/src/app/runtime/stream/tools/mod.rs +++ b/src/app/runtime/stream/tools/mod.rs @@ -1,10 +1,15 @@ use super::turn::ParsedToolCall; use serde_json::{json, Value}; +/// Standalone tool-call delta accumulator, functionally equivalent to the accumulation +/// logic built into `StreamedTurn::apply_event`. Reserved for callers that want to track +/// tool-call deltas independently of a full `StreamedTurn` (e.g. a lighter-weight preview). +#[allow(dead_code)] pub struct ToolCallAccumulator { calls: Vec, } +#[allow(dead_code)] impl ToolCallAccumulator { pub fn new() -> Self { ToolCallAccumulator { calls: Vec::new() } diff --git a/src/app/runtime/stream/turn.rs b/src/app/runtime/stream/turn.rs index 9a7b07a..c620d9b 100644 --- a/src/app/runtime/stream/turn.rs +++ b/src/app/runtime/stream/turn.rs @@ -22,6 +22,10 @@ pub struct ParsedToolCall { } impl ParsedToolCall { + /// Attempts to parse the accumulated argument string as JSON before the tool call is + /// marked complete — useful for callers that want a speculative preview mid-stream. + /// `build_assistant_message` does its own (lossy-fallback) parse for the final message. + #[allow(dead_code)] pub fn try_parse(&self) -> Option { serde_json::from_str(&self.arguments).ok() } @@ -115,10 +119,15 @@ impl StreamedTurn { msg } + /// Reserved accessor for callers that want to branch mid-stream before the turn + /// completes; the current wiring only inspects the final `build_assistant_message()`. + #[allow(dead_code)] pub fn has_tool_calls(&self) -> bool { self.tool_calls.iter().any(|tc| !tc.name.is_empty()) } + /// Reserved accessor mirroring `has_tool_calls` for mid-stream content peeks. + #[allow(dead_code)] pub fn content(&self) -> &str { &self.accumulated_content } diff --git a/src/app/state/misc.rs b/src/app/state/misc.rs index d1ec1c3..24d1974 100644 --- a/src/app/state/misc.rs +++ b/src/app/state/misc.rs @@ -204,6 +204,9 @@ pub struct MiscState { pub esc_press_count: u32, pub last_staleness_sweep_ms: i64, pub thinking: bool, + pub effort_level: usize, + pub selected_index: usize, + pub editor: Option, } impl MiscState { @@ -216,6 +219,9 @@ impl MiscState { esc_press_count: 0, last_staleness_sweep_ms: 0, thinking: false, + effort_level: 1, + selected_index: 0, + editor: None, } } diff --git a/src/app/state/rest.rs b/src/app/state/rest.rs index ae357d6..3196797 100644 --- a/src/app/state/rest.rs +++ b/src/app/state/rest.rs @@ -8,6 +8,7 @@ use super::runtime::{SessionRuntime, TurnEvent}; use super::types::{AgentMode, Origin, Toast, TranscriptCache}; use crate::app::mcp::manager::McpManager; use crate::app::workflow::engine::WorkflowEngine; +use crate::model::app_config::AppConfig; use crate::model::editlog::EditLog; use crate::model::settings::Settings; @@ -32,6 +33,7 @@ impl ChatMessageDisplay { pub struct AppStateRest { pub mode: AgentMode, pub settings: Settings, + pub app_config: AppConfig, pub workspace_roots: Vec, pub session_id: String, pub session_dir: PathBuf, @@ -57,6 +59,7 @@ pub struct AppStateRest { impl AppStateRest { pub fn new(workspace_roots: Vec, session_dir: PathBuf, memory_dir: PathBuf) -> Self { let settings = Settings::load(); + let app_config = AppConfig::load(); let download_dir = memory_dir.parent().unwrap_or(&memory_dir).join("downloads"); let worktrees_dir = memory_dir.parent().unwrap_or(&memory_dir).join("worktrees"); let dir_cache = DirCache::new(); @@ -67,6 +70,7 @@ impl AppStateRest { AppStateRest { mode: AgentMode::Normal, settings, + app_config, workspace_roots, session_id, session_dir: session_dir.clone(), diff --git a/src/app/state/runtime.rs b/src/app/state/runtime.rs index dcfbfa3..7862632 100644 --- a/src/app/state/runtime.rs +++ b/src/app/state/runtime.rs @@ -74,6 +74,13 @@ pub enum TurnEvent { kind: String, message: String, }, + StreamStart, + StreamToken(String), + StreamDone(crate::dto::chat::message::ChatMessage), + Usage { + tokens_in: u64, + tokens_out: u64, + }, Error(String), Done, } diff --git a/src/app/subagent/engine.rs b/src/app/subagent/engine.rs index e951d58..05f0484 100644 --- a/src/app/subagent/engine.rs +++ b/src/app/subagent/engine.rs @@ -32,7 +32,7 @@ pub fn run_subagent(ctx: SubagentContext, tx: mpsc::Sender) -> an let api_key = std::env::var("API_KEY").unwrap_or_default(); let model = std::env::var("MODEL").unwrap_or_default(); - let client = crate::service::provider::LlmClient::new(api_key, model); + let client = crate::service::provider::LlmClient::new(api_key, model, None); let response = match client.chat(&messages) { Ok(r) => r, Err(e) => { diff --git a/src/controller/command.rs b/src/controller/command.rs index 06c231d..1bdd11e 100644 --- a/src/controller/command.rs +++ b/src/controller/command.rs @@ -13,6 +13,11 @@ pub enum Command { Mode(ModeKind), Clear, Login { provider: String }, + Edit(String), + McpAdd { + name: String, + command: String, + }, Unknown(String), } @@ -50,6 +55,23 @@ pub fn parse_command(text: &str) -> Command { "/lesson" => Command::LessonList, "/login" if !arg1.is_empty() => Command::Login { provider: arg1.to_string() }, "/login" => Command::Login { provider: "zen".to_string() }, + "/edit" if !arg1.is_empty() => Command::Edit(arg1.to_string()), + "/edit" => Command::Edit(".".to_string()), + "/mcp" if arg1 == "add" && !arg2.is_empty() => { + // /mcp add [args...] + // name is the first word of arg2, the rest is the command + let rest = arg2.trim(); + if let Some(space) = rest.find(' ') { + let name = rest[..space].to_string(); + let command = rest[space + 1..].trim().to_string(); + Command::McpAdd { name, command } + } else { + Command::McpAdd { + name: rest.to_string(), + command: String::new(), + } + } + } _ => Command::Unknown(cmd.to_string()), } } diff --git a/src/controller/input.rs b/src/controller/input.rs index 6e20d1f..1c6344b 100644 --- a/src/controller/input.rs +++ b/src/controller/input.rs @@ -8,6 +8,53 @@ use crate::app::state::types::Overlay; use crate::controller::command::parse_command; pub fn handle_key(key: KeyEvent, state: &mut AppStateRest) -> Vec { + // While Editor overlay is active, route input directly to the editor handler + if state.misc.overlay == Overlay::Editor { + match key.code { + KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => { + return vec![Action::QuitConfirm]; + } + KeyCode::Char('s') if key.modifiers.contains(KeyModifiers::CONTROL) => { + if let Some(ref ed) = state.misc.editor.clone() { + let content = ed.as_string(); + if let Err(e) = std::fs::write(&ed.path, &content) { + state.push_toast(crate::app::state::types::Toast::new( + crate::app::state::types::ToastKind::Error, + format!("Save failed: {}", e), + )); + } else { + state.push_toast(crate::app::state::types::Toast::new( + crate::app::state::types::ToastKind::Success, + format!("Saved {}", ed.path), + )); + } + state.dirty = true; + } + return vec![]; + } + KeyCode::Esc => { + crate::app::mode::editor::handle_editor_dismiss(state); + return vec![]; + } + KeyCode::Backspace => { + if let Some(ref mut ed) = state.misc.editor { + ed.delete_left(); + state.dirty = true; + } + return vec![]; + } + KeyCode::Enter => { + crate::app::mode::editor::handle_editor_input(state, "\n".to_string()); + return vec![]; + } + KeyCode::Char(c) => { + crate::app::mode::editor::handle_editor_input(state, c.to_string()); + return vec![]; + } + _ => return vec![], + } + } + match key.code { KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => { vec![Action::QuitConfirm] @@ -40,6 +87,19 @@ pub fn handle_key(key: KeyEvent, state: &mut AppStateRest) -> Vec { KeyCode::Up => { if key.modifiers.contains(KeyModifiers::CONTROL) { vec![Action::HistoryUp] + } else if state.misc.overlay == Overlay::Effort { + mode::effort::cycle_effort(state); + Vec::new() + } else if state.misc.overlay == Overlay::OnboardProvider { + let n = mode::onboard_provider::PROVIDERS.len(); + state.misc.selected_index = if state.misc.selected_index == 0 { n - 1 } else { state.misc.selected_index - 1 }; + state.dirty = true; + Vec::new() + } else if state.misc.overlay == Overlay::Rewind { + let n = mode::rewind::rewind_count(state); + state.misc.selected_index = if state.misc.selected_index == 0 { n.saturating_sub(1) } else { state.misc.selected_index - 1 }; + state.dirty = true; + Vec::new() } else { vec![Action::ScrollUp] } @@ -47,6 +107,19 @@ pub fn handle_key(key: KeyEvent, state: &mut AppStateRest) -> Vec { KeyCode::Down => { if key.modifiers.contains(KeyModifiers::CONTROL) { vec![Action::HistoryDown] + } else if state.misc.overlay == Overlay::Effort { + mode::effort::cycle_effort(state); + Vec::new() + } else if state.misc.overlay == Overlay::OnboardProvider { + let n = mode::onboard_provider::PROVIDERS.len(); + state.misc.selected_index = (state.misc.selected_index + 1) % n; + state.dirty = true; + Vec::new() + } else if state.misc.overlay == Overlay::Rewind { + let n = mode::rewind::rewind_count(state); + state.misc.selected_index = if n == 0 { 0 } else { (state.misc.selected_index + 1) % n }; + state.dirty = true; + Vec::new() } else { vec![Action::ScrollDown] } @@ -134,6 +207,50 @@ fn handle_overlay_enter(state: &mut AppStateRest) -> Vec { Overlay::QuitConfirm => { vec![mode::quit_confirm::handle_quit_confirm(true)] } + Overlay::KeyInput => { + let text = state.input.buffer.clone(); + mode::key_input::handle_key_text(state, text.clone()); + state.settings.api_key = if text.is_empty() { None } else { Some(text) }; + let _ = state.settings.save(); + state.input.buffer.clear(); + state.input.cursor = 0; + state.misc.overlay = Overlay::None; + state.push_toast(crate::app::state::types::Toast::new( + crate::app::state::types::ToastKind::Success, + "API key saved".to_string(), + )); + state.dirty = true; + Vec::new() + } + Overlay::Onboard => { + state.misc.overlay = Overlay::OnboardProvider; + state.misc.selected_index = 0; + state.dirty = true; + Vec::new() + } + Overlay::OnboardProvider => { + let idx = state.misc.selected_index.min(mode::onboard_provider::PROVIDERS.len() - 1); + let provider = mode::onboard_provider::PROVIDERS[idx]; + mode::onboard_provider::set_provider(&mut state.settings, provider); + state.push_toast(crate::app::state::types::Toast::new( + crate::app::state::types::ToastKind::Info, + format!("Provider set to {}", provider), + )); + state.misc.overlay = Overlay::KeyInput; + state.input.buffer.clear(); + state.input.cursor = 0; + state.dirty = true; + Vec::new() + } + Overlay::Mcp => { + mode::mcp::connect_mcp(state, ""); + Vec::new() + } + Overlay::Rewind => { + let idx = state.misc.selected_index; + mode::rewind::rewind_to(state, idx); + Vec::new() + } _ => Vec::new(), } } diff --git a/src/dto/provider/request.rs b/src/dto/provider/request.rs index afc0d08..6e4a0bc 100644 --- a/src/dto/provider/request.rs +++ b/src/dto/provider/request.rs @@ -17,6 +17,13 @@ pub struct ChatRequest { pub top_p: Option, #[serde(skip_serializing_if = "Option::is_none")] pub stop: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub stream_options: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamOptions { + pub include_usage: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/main.rs b/src/main.rs index f31037b..43e24b9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -53,6 +53,11 @@ fn run_single_process() -> Result<()> { let session_dir = store.base_dir.join("sessions").join(&session_id); std::fs::create_dir_all(&session_dir)?; + let session_lock = model::session_lock::SessionLock::new(&session_dir); + if !session_lock.try_lock()? { + anyhow::bail!("session already active (another zesdex process holds the lock for this session directory)"); + } + let workspace_roots = vec![std::env::current_dir()?]; let mut state = app::state::rest::AppStateRest::new( workspace_roots.clone(), @@ -60,6 +65,9 @@ fn run_single_process() -> Result<()> { store.memory_dir, ); state.sessions = model::session::Session::list(&store.base_dir); + if state.settings.api_key.is_none() { + state.misc.overlay = app::state::types::Overlay::Onboard; + } let _rt = tokio::runtime::Runtime::new()?; @@ -82,6 +90,7 @@ fn run_single_process() -> Result<()> { } let _ = state.settings.save(); + session_lock.unlock(); Ok(()) } @@ -255,6 +264,11 @@ fn run_daemon() -> Result<()> { let session_dir = store.base_dir.join("sessions").join(&session_id); std::fs::create_dir_all(&session_dir)?; + let session_lock = model::session_lock::SessionLock::new(&session_dir); + if !session_lock.try_lock()? { + anyhow::bail!("session already active (another zesdex process holds the lock for this session directory)"); + } + let workspace_roots = vec![std::env::current_dir()?]; let mut state = app::state::rest::AppStateRest::new( workspace_roots.clone(), @@ -336,6 +350,7 @@ fn run_daemon() -> Result<()> { let _ = std::fs::remove_file(&socket_path); let _ = state.settings.save(); + session_lock.unlock(); Ok(()) } diff --git a/src/model/mod.rs b/src/model/mod.rs index 5028989..849c319 100644 --- a/src/model/mod.rs +++ b/src/model/mod.rs @@ -1,6 +1,8 @@ +pub mod app_config; pub mod editlog; pub mod memory; pub mod msglog; pub mod session; +pub mod session_lock; pub mod settings; pub mod store; diff --git a/src/model/msglog/mod.rs b/src/model/msglog/mod.rs index 1b72bd7..9f2fc7b 100644 --- a/src/model/msglog/mod.rs +++ b/src/model/msglog/mod.rs @@ -1,6 +1,8 @@ +pub mod blobs; pub mod query; pub mod schema; +pub use blobs::store_blob; pub use query::insert_message; pub fn open_or_create(session_dir: &std::path::Path) -> anyhow::Result { diff --git a/src/service/provider.rs b/src/service/provider.rs index bb80536..bc6bd9a 100644 --- a/src/service/provider.rs +++ b/src/service/provider.rs @@ -1,8 +1,10 @@ use std::time::Duration; use anyhow::Result; +use crate::app::runtime::stream::{SseParser, StreamEvent}; +use crate::app::runtime::stream::turn::StreamedTurn; use crate::dto::chat::message::ChatMessage; -use crate::dto::provider::request::ToolDef; +use crate::dto::provider::request::{ChatRequest, StreamOptions, ToolDef}; const DEFAULT_BASE_URL: &str = "https://opencode.ai/zen/v1"; const DEFAULT_MODEL: &str = "deepseek-v4-flash-free"; @@ -17,7 +19,7 @@ pub struct LlmClient { } impl LlmClient { - pub fn new(api_key: String, model: String) -> Self { + pub fn new(api_key: String, model: String, base_url: Option) -> Self { let model = if model.is_empty() { DEFAULT_MODEL.to_string() } else { @@ -31,7 +33,7 @@ impl LlmClient { LlmClient { client, api_key, - base_url: DEFAULT_BASE_URL.to_string(), + base_url: base_url.filter(|s| !s.is_empty()).unwrap_or_else(|| DEFAULT_BASE_URL.to_string()), model, } } @@ -46,7 +48,7 @@ impl LlmClient { messages: &[ChatMessage], tools: Option>, ) -> Result { - let req = crate::dto::provider::request::ChatRequest { + let req = ChatRequest { model: self.model.clone(), messages: messages.to_vec(), max_tokens: Some(4096), @@ -55,18 +57,130 @@ impl LlmClient { stream: Some(false), top_p: None, stop: None, + stream_options: None, }; let url = format!("{}/chat/completions", self.base_url); - let mut http_req = self.client - .post(&url) - .header("Content-Type", "application/json"); + let max_retries = 10; + let mut attempt = 0; + loop { + attempt += 1; + + let mut http_req = self.client + .post(&url) + .header("Content-Type", "application/json"); + + if !self.api_key.is_empty() { + http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key)); + } + + let result = (|| -> Result { + let resp = http_req.json(&req).send().map_err(|e| { + if e.is_timeout() { + anyhow::anyhow!("API request timed out after {:?}. Check your network or try again.", REQUEST_TIMEOUT) + } else if e.is_connect() { + anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url) + } else { + anyhow::anyhow!("API request failed: {}", e) + } + })?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().unwrap_or_default(); + anyhow::bail!("API error {} from {}: {}", status, self.base_url, body); + } + + let data: crate::dto::provider::response::ChatResponse = resp.json()?; + let message = data + .choices + .into_iter() + .next() + .map(|c| c.message) + .ok_or_else(|| anyhow::anyhow!("API response had no choices"))?; + Ok(message) + })(); + + match result { + Ok(msg) => return Ok(msg), + Err(e) => { + if attempt >= max_retries { + return Err(e); + } + eprintln!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries); + std::thread::sleep(Duration::from_secs(2)); + } + } + } + } + + /// Streaming variant of `chat_with_tools`. Feeds SSE chunks into an `SseParser` / + /// `StreamedTurn` and invokes `on_event` for every parsed `StreamEvent` as it arrives, + /// so the caller can push incremental UI updates in real time. Returns the fully + /// assembled assistant message plus token usage (prompt, completion) if the server + /// reported it. Retries the whole request only if no event has been observed yet + /// (once tokens start arriving, a partial turn cannot be safely replayed). + pub fn chat_with_tools_streaming( + &self, + messages: &[ChatMessage], + tools: Option>, + temperature: Option, + max_tokens: Option, + mut on_event: impl FnMut(&StreamEvent), + ) -> Result<(ChatMessage, Option<(u64, u64)>)> { + let req = ChatRequest { + model: self.model.clone(), + messages: messages.to_vec(), + max_tokens: Some(max_tokens.unwrap_or(4096)), + temperature: Some(temperature.unwrap_or(0.7)), + tools, + stream: Some(true), + top_p: None, + stop: None, + stream_options: Some(StreamOptions { include_usage: true }), + }; + + let url = format!("{}/chat/completions", self.base_url); + let max_retries = 10; + let mut attempt = 0; + let mut started = false; + + loop { + attempt += 1; + let mut wrapped = |event: &StreamEvent| { + started = true; + on_event(event); + }; + match self.try_stream_once(&req, &url, &mut wrapped) { + Ok(result) => return Ok(result), + Err(e) => { + if started || attempt >= max_retries { + return Err(e); + } + eprintln!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries); + std::thread::sleep(Duration::from_secs(2)); + } + } + } + } + + fn try_stream_once( + &self, + req: &ChatRequest, + url: &str, + on_event: &mut dyn FnMut(&StreamEvent), + ) -> Result<(ChatMessage, Option<(u64, u64)>)> { + use std::io::Read; + + let mut http_req = self.client + .post(url) + .header("Content-Type", "application/json"); if !self.api_key.is_empty() { http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key)); } - let resp = http_req.json(&req).send().map_err(|e| { + let resp = http_req.json(req).send().map_err(|e| { if e.is_timeout() { anyhow::anyhow!("API request timed out after {:?}. Check your network or try again.", REQUEST_TIMEOUT) } else if e.is_connect() { @@ -82,13 +196,49 @@ impl LlmClient { anyhow::bail!("API error {} from {}: {}", status, self.base_url, body); } - let data: crate::dto::provider::response::ChatResponse = resp.json()?; - let message = data - .choices - .into_iter() - .next() - .map(|c| c.message) - .ok_or_else(|| anyhow::anyhow!("API response had no choices"))?; - Ok(message) + let mut turn = StreamedTurn::new(); + let mut usage: Option<(u64, u64)> = None; + let mut parser = SseParser::new(); + let mut reader = resp; + let mut byte_buf: Vec = Vec::new(); + let mut chunk_buf = [0u8; 4096]; + + loop { + let n = reader.read(&mut chunk_buf) + .map_err(|e| anyhow::anyhow!("stream read error: {}", e))?; + if n == 0 { + break; + } + byte_buf.extend_from_slice(&chunk_buf[..n]); + let valid_len = match std::str::from_utf8(&byte_buf) { + Ok(s) => s.len(), + Err(e) => e.valid_up_to(), + }; + if valid_len == 0 { + continue; + } + let text = String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned(); + byte_buf.drain(..valid_len); + + for event in parser.feed(&text) { + on_event(&event); + match &event { + StreamEvent::Usage { prompt_tokens, completion_tokens, .. } => { + usage = Some((*prompt_tokens, *completion_tokens)); + } + StreamEvent::Error(msg) => { + anyhow::bail!("stream error: {}", msg); + } + StreamEvent::Done => { + turn.apply_event(&event); + return Ok((turn.build_assistant_message(), usage)); + } + _ => turn.apply_event(&event), + } + } + } + + turn.is_complete = true; + Ok((turn.build_assistant_message(), usage)) } } diff --git a/src/tool/mod.rs b/src/tool/mod.rs index f8fa3a4..3cb88b7 100644 --- a/src/tool/mod.rs +++ b/src/tool/mod.rs @@ -13,6 +13,7 @@ pub mod plan; pub mod search; pub mod seqthink; pub mod shell; +pub mod shell_filter; pub mod utility; pub mod workflow; diff --git a/src/tool/shell.rs b/src/tool/shell.rs index 8448794..5db5aee 100644 --- a/src/tool/shell.rs +++ b/src/tool/shell.rs @@ -51,6 +51,10 @@ impl Tool for Bash { let workspace_roots: Vec<&std::path::Path> = ctx.workspaces.iter().map(|p| p.as_path()).collect(); crate::app::catastrophic::CatastrophicGuard::check_all(&cmd, &workspace_roots) .map_err(|e| anyhow!("catastrophic guard blocked: {}", e))?; + super::shell_filter::credentials::check_credential_read(&cmd) + .map_err(|e| anyhow!("blocked: {}", e))?; + super::shell_filter::git::check_git_destructive(&cmd) + .map_err(|e| anyhow!("blocked: {}", e))?; let run_in_background = args.get("run_in_background").and_then(|v| v.as_bool()).unwrap_or(false); if run_in_background { let job = crate::app::bgbash::job::spawn_bash_job(cmd); diff --git a/src/view/chat.rs b/src/view/chat.rs index 903bd30..66aed55 100644 --- a/src/view/chat.rs +++ b/src/view/chat.rs @@ -74,10 +74,9 @@ pub fn draw_chat(frame: &mut Frame, area: Rect, state: &crate::app::state::rest: } else { msg.content.clone() }; - let content_line = Line::from(vec![ - Span::styled(format!("{} ", prefix), Style::default().fg(role_color)), - Span::styled(content_str, Style::default().fg(Theme::TEXT)), - ]); + let mut content_spans = vec![Span::styled(format!("{} ", prefix), Style::default().fg(role_color))]; + content_spans.extend(super::markdown::render_markdown(&content_str, area.width)); + let content_line = Line::from(content_spans); display_lines.push(header); display_lines.push(content_line); diff --git a/src/view/markdown.rs b/src/view/markdown.rs index f506adb..c0f1768 100644 --- a/src/view/markdown.rs +++ b/src/view/markdown.rs @@ -1,7 +1,7 @@ use ratatui::style::{Color, Modifier, Style}; use ratatui::text::Span; -pub fn render_markdown<'a>(text: &'a str, width: u16) -> Vec> { +pub fn render_markdown(text: &str, width: u16) -> Vec> { let mut spans = Vec::new(); let parser = pulldown_cmark::Parser::new(text); let mut in_code_block = false; @@ -122,6 +122,7 @@ pub fn render_markdown<'a>(text: &'a str, width: u16) -> Vec> { let mut spans_out = Vec::new(); let mut line_len = 0; for span in &spans { + let style = span.style; let s = span.content.clone(); let text = s.as_ref(); let remaining = text.len(); @@ -129,7 +130,7 @@ pub fn render_markdown<'a>(text: &'a str, width: u16) -> Vec> { spans_out.push(Span::raw("\n")); line_len = 0; } - spans_out.push(Span::raw(text.to_string())); + spans_out.push(Span::styled(text.to_string(), style)); if !text.contains('\n') { line_len += remaining; } else { diff --git a/src/view/mod.rs b/src/view/mod.rs index a009e41..dbb3e6d 100644 --- a/src/view/mod.rs +++ b/src/view/mod.rs @@ -1,4 +1,5 @@ pub mod chat; +pub mod markdown; pub mod status; pub mod theme; pub mod workflow; @@ -148,7 +149,7 @@ fn render_overlay(frame: &mut Frame, area: Rect, overlay: crate::app::state::typ Style::default(), )), Line::from(Span::styled( - "Press Ctrl+Q again to confirm, Esc to cancel.", + "Press Enter to confirm, Esc to cancel.", Style::default().fg(Theme::DIM), )), ]; @@ -254,12 +255,8 @@ fn render_overlay(frame: &mut Frame, area: Rect, overlay: crate::app::state::typ } crate::app::state::types::Overlay::Effort => { let block = block.title(" Effort Level "); - let levels = ["low", "medium", "high", "xhigh", "max"]; - let current_idx = if state.settings.temperature > 0.8 { 0usize } - else if state.settings.temperature > 0.5 { 1 } - else if state.settings.temperature > 0.3 { 2 } - else if state.settings.max_tokens > 4000 { 3 } - else { 4 }; + let levels = crate::app::mode::effort::EFFORT_LEVELS; + let current_idx = crate::app::mode::effort::current_effort(state); let mut lines: Vec = levels.iter().enumerate().map(|(i, l)| { let selected = i == current_idx; Line::from(Span::styled(