MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
- 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:
asepharyana
2026-07-12 01:25:52 +07:00
parent fcef85a327
commit 18a41aad48
28 changed files with 838 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/app/runtime/actions/mod.rs
+140 -5
View File
@@ -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<String> = 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<String> = command.split_whitespace().map(|s| s.to_string()).collect();
let cmd = extra_args.first().cloned().unwrap_or_default();
let args: Vec<String> = 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<std::sync::Arc<std::sync::Mutex<rusqlite::Connection>>>,
temperature: f32,
max_tokens: u32,
}
fn archive_message(db: &Option<std::sync::Arc<std::sync::Mutex<rusqlite::Connection>>>, 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<dyn crate::tool::Tool>],
ctx: &crate::tool::ToolCtx,
name: &str,
tool_call_id: &str,
args: &serde_json::Value,
session_dir: &std::path::Path,
session_id: &str,
db: &Option<std::sync::Arc<std::sync::Mutex<rusqlite::Connection>>>,
) -> anyhow::Result<String> {
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