From efcd191f9698023944f9096a86c1efa47e412b6b Mon Sep 17 00:00:00 2001 From: asepharyana Date: Mon, 20 Jul 2026 16:29:32 +0700 Subject: [PATCH] feat(agent): implement agent execution engine and turn handling with background processing --- apps/domain/src/auth/session_lock.rs | 53 ++-- apps/domain/src/core/provider.rs | 8 +- apps/gateway/src/main.rs | 12 +- apps/infrastructure/src/agent/mod.rs | 6 + apps/infrastructure/src/agent/runner.rs | 216 +++++++++++++++++ apps/infrastructure/src/ipc/client.rs | 6 +- apps/infrastructure/src/ipc/conn.rs | 7 +- apps/infrastructure/src/ipc/server.rs | 2 +- apps/infrastructure/src/lib.rs | 1 + apps/infrastructure/src/llm/provider.rs | 11 +- .../src/persistence/iam/session_lock_repo.rs | 35 ++- apps/interfaces/daemon/src/handler.rs | 39 +++ apps/interfaces/tui/src/turn.rs | 228 ++---------------- apps/interfaces/tui/src/view/sidebar.rs | 8 +- apps/interfaces/ws/src/lib.rs | 67 ++++- 15 files changed, 452 insertions(+), 247 deletions(-) create mode 100644 apps/infrastructure/src/agent/mod.rs create mode 100644 apps/infrastructure/src/agent/runner.rs diff --git a/apps/domain/src/auth/session_lock.rs b/apps/domain/src/auth/session_lock.rs index 3d45b8b..c42246f 100644 --- a/apps/domain/src/auth/session_lock.rs +++ b/apps/domain/src/auth/session_lock.rs @@ -104,36 +104,55 @@ impl SessionLock { let _ = fs::remove_file(&self.path); } - /// Check whether a process with the given PID is currently alive. + /// Check whether a process with the given PID is currently alive and + /// belongs to the same binary (mitigating PID-reuse races). + /// + /// Strategy (Unix): + /// 1. Resolve `/proc//exe` — if it doesn't match our own binary, + /// the PID either belongs to another process or is reused — return false. + /// 2. Send `kill(pid, 0)` to verify the process is still alive. + /// 3. Re-check `/proc//exe` to close the TOCTOU window between + /// step 1 and step 2 (PID reuse after exe check, before kill). /// - /// Uses `kill(pid, 0)` on Unix via the `nix` or `libc` crate in production; - /// here we provide a best-effort check using the process table. /// On non-Unix platforms this always returns `true` (conservative). fn is_alive(pid: u32) -> bool { - // On Unix, signal 0 checks process existence without sending a signal. #[cfg(unix)] { - // SAFETY: `libc::kill(pid, 0)` does not send a signal; it only checks - // whether the process exists and the caller has permission to signal it. - // The integer argument is a PID validated by `try_lock`. + // Resolve our own executable path once. + let self_exe = match std::fs::read_link("/proc/self/exe") { + Ok(exe) => exe, + Err(_) => return false, + }; + let pid_signed: i32 = match pid.try_into() { Ok(p) => p, Err(_) => return false, }; + + let proc_exe = std::path::PathBuf::from(format!("/proc/{pid}/exe")); + + // Phase 1: read /proc//exe and compare with self_exe. + let target = match std::fs::read_link(&proc_exe) { + Ok(t) => t, + Err(_) => return false, + }; + if target != self_exe { + return false; + } + + // Phase 2: verify the process is still alive. + // SAFETY: `libc::kill(pid, 0)` does not send a signal; it only checks + // whether the process exists and the caller has permission to signal it. if unsafe { libc::kill(pid_signed, 0) != 0 } { return false; } - // Extra check: verify the PID belongs to a zesdex process via - // /proc//exe to mitigate the PID-reuse race. - let proc_exe = std::path::PathBuf::from(format!("/proc/{pid}/exe")); - if let Ok(target) = std::fs::read_link(&proc_exe) { - if let Ok(exe) = std::env::current_exe() { - if target != exe { - return false; - } - } + + // Phase 3: re-check /proc//exe to detect PID reuse between + // Phase 1 and Phase 2. + match std::fs::read_link(&proc_exe) { + Ok(recheck) if recheck == self_exe => true, + _ => false, } - true } #[cfg(not(unix))] diff --git a/apps/domain/src/core/provider.rs b/apps/domain/src/core/provider.rs index 210f7e3..296bb29 100644 --- a/apps/domain/src/core/provider.rs +++ b/apps/domain/src/core/provider.rs @@ -333,8 +333,9 @@ impl SseParser { if let Some(tool_calls) = d.get("tool_calls").and_then(|tc| tc.as_array()) { + const MAX_TOOL_CALLS: usize = 64; for tc in tool_calls { - let index = + let raw_index = tc.get("index").and_then(Value::as_u64).unwrap_or_else( || { tracing::warn!( @@ -344,7 +345,10 @@ impl SseParser { 0 }, ); - let index = usize::try_from(index).unwrap_or(0); + // Clamp index to prevent out-of-bounds / memory exhaustion + let index = usize::try_from(raw_index) + .unwrap_or(0) + .min(MAX_TOOL_CALLS.saturating_sub(1)); let id = tc .get("id") .and_then(|i| i.as_str()) diff --git a/apps/gateway/src/main.rs b/apps/gateway/src/main.rs index 72efd2e..b90635e 100644 --- a/apps/gateway/src/main.rs +++ b/apps/gateway/src/main.rs @@ -137,9 +137,19 @@ fn run_api_server(port: u16) -> anyhow::Result<()> { rt.block_on(async { let store = zesdex_domain::core::Store::new(); store.ensure_dirs()?; + + // Load JWT secret from environment variable with a secure default warning + let jwt_secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| { + tracing::warn!( + "JWT_SECRET environment variable not set; using insecure default. \ + Set JWT_SECRET to a secure random value in production." + ); + "dev-secret".to_string() + }); + let state = zesdex_api::ApiState::new( store.base_dir.clone(), - "dev-secret", + jwt_secret, "", "deepseek-v4-flash-free", Some("https://opencode.ai/zen/v1".to_string()), diff --git a/apps/infrastructure/src/agent/mod.rs b/apps/infrastructure/src/agent/mod.rs new file mode 100644 index 0000000..fac884c --- /dev/null +++ b/apps/infrastructure/src/agent/mod.rs @@ -0,0 +1,6 @@ +//! Agent execution engine — orchestrates LLM streaming, system prompt assembly, +//! tool execution, and turn event emitting on background threads. + +pub mod runner; + +pub use runner::{spawn_agent_turn, AgentTurnParams}; diff --git a/apps/infrastructure/src/agent/runner.rs b/apps/infrastructure/src/agent/runner.rs new file mode 100644 index 0000000..998f4c8 --- /dev/null +++ b/apps/infrastructure/src/agent/runner.rs @@ -0,0 +1,216 @@ +//! Agent turn engine runner — handles LLM API calls, system prompt assembly, +//! and tool execution loops on background threads. + +use std::collections::VecDeque; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use tracing::{debug, info, warn}; + +use zesdex_domain::core::tool_call::sanitize_tool_arguments; +use zesdex_domain::core::ChatMessage; +use crate::llm::provider::LlmClient; +use crate::tools::{all_tools, tool_defs, ToolCtx}; +use crate::TurnEvent; + +/// Owned parameters required to spawn and execute an agent turn on a background thread. +pub struct AgentTurnParams { + pub messages: Vec, + pub session_dir: PathBuf, + pub workspace_roots: Vec, + pub turn_events: Arc>>, + pub in_flight: Arc, + pub abort: Arc, + pub api_key: String, + pub model: String, + pub api_base: Option, +} + +/// Spawns an agent turn on a background OS thread. +#[tracing::instrument(skip(params))] +pub fn spawn_agent_turn(mut params: AgentTurnParams) { + info!( + "spawning agent turn with {} messages (model: {})", + params.messages.len(), + params.model + ); + + std::thread::spawn(move || { + run_turn(&mut params); + }); +} + +/// The core agent turn loop — LLM call → tool execution → repeat. +#[tracing::instrument(skip(params))] +fn run_turn(params: &mut AgentTurnParams) { + let client = LlmClient::new( + params.api_key.clone(), + params.model.clone(), + params.api_base.clone(), + ); + + let tools = all_tools(); + let defs = tool_defs(&tools); + + let mut sys_prompt = "You are Zesdex, an AI coding assistant. You have access to various tools via native function calling to help the user. \ + For complex problems, use `seq_think` to reason step-by-step. \ + For any non-trivial tasks, you MUST prioritize creating a structured plan (using `plan_enter`) and a list of TODOs (using `todowrite`) BEFORE executing any other tools or modifying files. \ + For large multi-step operations or delegating tasks, you MUST prioritize using `workflow_run` (to run a yaml workflow script) or `hive_mind` (to orchestrate multiple agents) to complete the task efficiently. \ + Use other tools when necessary. If a tool returns an error, fix the issue before retrying. \ + Respond conversationally, concisely, and helpfully.".to_string(); + + if let Some(root) = params.workspace_roots.first() { + let tree = crate::utils::build_workspace_tree(root, 800); + let rich_ctx = crate::utils::build_rich_context(root); + + sys_prompt.push_str("\n\nWorkspace structure:\n```\n"); + sys_prompt.push_str(&tree); + sys_prompt.push_str("\n```\n\n"); + + sys_prompt.push_str(&rich_ctx); + } + + let sys_msg = ChatMessage::system(sys_prompt); + + let tool_ctx = ToolCtx::builder() + .session_dir(params.session_dir.clone()) + .workspaces(params.workspace_roots.clone()) + .turn_events(Arc::clone(¶ms.turn_events)) + .build(); + + for iteration in 0..50 { + if params.abort.load(Ordering::SeqCst) { + params.abort.store(false, Ordering::SeqCst); + push_event( + ¶ms.turn_events, + TurnEvent::SystemNote { + kind: "info".into(), + message: "Turn aborted by user".into(), + }, + ); + break; + } + + debug!("agent turn iteration {iteration}"); + + // Stream the LLM response + push_event(¶ms.turn_events, TurnEvent::StreamStart); + + let mut req_messages = params.messages.clone(); + req_messages.insert(0, sys_msg.clone()); + + let result = client.chat_with_tools_streaming( + &req_messages, + Some(defs.clone()), + Some(0.7), + Some(4096), + |event| { + if params.abort.load(Ordering::SeqCst) { + return false; + } + match event { + zesdex_domain::core::StreamEvent::Token(s) => { + push_event(¶ms.turn_events, TurnEvent::StreamToken(s.clone())); + } + zesdex_domain::core::StreamEvent::Reasoning(s) => { + push_event(¶ms.turn_events, TurnEvent::StreamReasoning(s.clone())); + } + _ => {} + } + true + }, + Some(¶ms.abort), + ); + + match result { + Ok((assistant_msg, usage)) => { + let content = assistant_msg.content.clone().unwrap_or_default(); + let tool_calls = assistant_msg.tool_calls.clone().unwrap_or_default(); + + push_event( + ¶ms.turn_events, + TurnEvent::StreamDone(assistant_msg.clone()), + ); + + if let Some((tokens_in, tokens_out)) = usage { + push_event( + ¶ms.turn_events, + TurnEvent::Usage { + tokens_in, + tokens_out, + }, + ); + } + + if tool_calls.is_empty() { + params.messages.push(ChatMessage::assistant(Some(content))); + break; + } + + params.messages.push(assistant_msg); + + for tc in &tool_calls { + let name = &tc.function.name; + let args = sanitize_tool_arguments(&tc.function.arguments); + + debug!("executing tool: {name}"); + + let output = if let Some(tool) = tools.iter().find(|t| t.name() == name) { + match tool.run(&tool_ctx, &args) { + Ok(o) => o, + Err(e) => format!("Error: {e}"), + } + } else { + format!("Unknown tool: {name}") + }; + + let is_error = output.starts_with("Error:"); + + push_event( + ¶ms.turn_events, + TurnEvent::ToolResult { + tool_call_id: tc.id.clone(), + tool_name: name.clone(), + output: output.clone(), + is_error, + path: None, + }, + ); + + params + .messages + .push(ChatMessage::tool(tc.id.clone(), output.clone())); + } + } + Err(e) => { + warn!("LLM call failed: {e}"); + push_event( + ¶ms.turn_events, + TurnEvent::Error(format!("LLM error: {e}")), + ); + break; + } + } + } + + // Propagate accumulated messages back to caller so the next turn starts + // with full history (assistant replies + tool results). + push_event( + ¶ms.turn_events, + TurnEvent::Compacted(params.messages.clone()), + ); + push_event(¶ms.turn_events, TurnEvent::Done); + mark_done(¶ms.in_flight); +} + +fn push_event(queue: &Arc>>, event: TurnEvent) { + if let Ok(mut q) = queue.lock() { + q.push_back(event); + } +} + +/// Mark the turn as done using lock-free atomic store. +fn mark_done(flag: &Arc) { + flag.store(false, Ordering::SeqCst); +} diff --git a/apps/infrastructure/src/ipc/client.rs b/apps/infrastructure/src/ipc/client.rs index ca95b68..bdf1541 100644 --- a/apps/infrastructure/src/ipc/client.rs +++ b/apps/infrastructure/src/ipc/client.rs @@ -12,7 +12,7 @@ pub struct IpcClient { impl IpcClient { pub fn connect_unix(path: &str) -> anyhow::Result { let stream = UnixStream::connect(path)?; - let conn = crate::ipc::conn::Connection::new(stream); + let conn = crate::ipc::conn::Connection::new(stream)?; Ok(Self { conn: Mutex::new(conn), }) @@ -22,7 +22,7 @@ impl IpcClient { let mut guard = self .conn .lock() - .expect("IpcClient mutex poisoned"); + .map_err(|e| anyhow::anyhow!("IpcClient mutex poisoned: {e}"))?; guard.send(msg) } @@ -30,7 +30,7 @@ impl IpcClient { let mut guard = self .conn .lock() - .expect("IpcClient mutex poisoned"); + .map_err(|e| anyhow::anyhow!("IpcClient mutex poisoned: {e}"))?; guard.receive() } } diff --git a/apps/infrastructure/src/ipc/conn.rs b/apps/infrastructure/src/ipc/conn.rs index 4dbe879..d444a42 100644 --- a/apps/infrastructure/src/ipc/conn.rs +++ b/apps/infrastructure/src/ipc/conn.rs @@ -1,6 +1,7 @@ //! Connection wrapper around a Unix socket stream, //! pairing a buffered reader with a raw writer. +use anyhow::Context; use std::io::BufReader; use std::os::unix::net::UnixStream; @@ -11,14 +12,14 @@ pub struct Connection { } impl Connection { - pub fn new(stream: UnixStream) -> Self { + pub fn new(stream: UnixStream) -> anyhow::Result { let reader = BufReader::new( stream .try_clone() - .expect("UnixStream::try_clone should never fail on Linux"), + .context("failed to clone Unix stream for IPC reader")?, ); let writer = stream; - Self { reader, writer } + Ok(Self { reader, writer }) } pub fn send(&mut self, msg: &T) -> anyhow::Result<()> { diff --git a/apps/infrastructure/src/ipc/server.rs b/apps/infrastructure/src/ipc/server.rs index fb16b12..2b30c7c 100644 --- a/apps/infrastructure/src/ipc/server.rs +++ b/apps/infrastructure/src/ipc/server.rs @@ -21,6 +21,6 @@ impl IpcServer { pub fn accept(&self) -> anyhow::Result { let (stream, _addr) = self.listener.accept()?; - Ok(crate::ipc::conn::Connection::new(stream)) + crate::ipc::conn::Connection::new(stream) } } diff --git a/apps/infrastructure/src/lib.rs b/apps/infrastructure/src/lib.rs index 620f5f5..6d7b4e4 100644 --- a/apps/infrastructure/src/lib.rs +++ b/apps/infrastructure/src/lib.rs @@ -26,6 +26,7 @@ //! └── middleware/ — Axum HTTP middleware (auth, cors, rate-limit) //! ``` +pub mod agent; pub mod auth; pub mod bgbash; pub mod guard; diff --git a/apps/infrastructure/src/llm/provider.rs b/apps/infrastructure/src/llm/provider.rs index ca4e7f1..e649243 100644 --- a/apps/infrastructure/src/llm/provider.rs +++ b/apps/infrastructure/src/llm/provider.rs @@ -324,7 +324,12 @@ impl LlmClient { name, arguments_delta, } => { - while self.tool_calls.len() <= *index { + // Cap the index to prevent memory exhaustion from + // maliciously large indices. + const MAX_TOOL_CALLS: usize = 64; + let index = usize::min(*index, MAX_TOOL_CALLS.saturating_sub(1)); + + while self.tool_calls.len() <= index { self.tool_calls.push(zesdex_domain::core::ToolCall { id: String::new(), type_: "function".to_string(), @@ -334,8 +339,8 @@ impl LlmClient { }, }); } - - let tc = &mut self.tool_calls[*index]; + + let tc = &mut self.tool_calls[index]; if let Some(ref id_val) = id { tc.id = id_val.clone(); diff --git a/apps/infrastructure/src/persistence/iam/session_lock_repo.rs b/apps/infrastructure/src/persistence/iam/session_lock_repo.rs index c1275fa..32647db 100644 --- a/apps/infrastructure/src/persistence/iam/session_lock_repo.rs +++ b/apps/infrastructure/src/persistence/iam/session_lock_repo.rs @@ -75,17 +75,36 @@ impl SessionLockRepository for FileSystemSessionLockRepository { Ok(p) => p, Err(_) => return false, }; + + // Resolve our own executable path once. + let self_exe = match std::fs::read_link("/proc/self/exe") { + Ok(exe) => exe, + Err(_) => return false, + }; + + let proc_exe = std::path::PathBuf::from(format!("/proc/{pid}/exe")); + + // Phase 1: read /proc//exe and compare with self_exe. + let target = match std::fs::read_link(&proc_exe) { + Ok(t) => t, + Err(_) => return false, + }; + if target != self_exe { + return false; + } + + // Phase 2: verify the process is still alive. + // SAFETY: `libc::kill(pid, 0)` does not send a signal; it only checks + // whether the process exists and the caller has permission to signal it. if unsafe { libc::kill(pid_signed, 0) != 0 } { return false; } - let proc_exe = std::path::PathBuf::from(format!("/proc/{pid}/exe")); - if let Ok(target) = std::fs::read_link(&proc_exe) { - if let Ok(exe) = std::env::current_exe() { - if target != exe { - return false; - } - } + + // Phase 3: re-check /proc//exe to detect PID reuse between + // Phase 1 and Phase 2. + match std::fs::read_link(&proc_exe) { + Ok(recheck) if recheck == self_exe => true, + _ => false, } - true } } diff --git a/apps/interfaces/daemon/src/handler.rs b/apps/interfaces/daemon/src/handler.rs index 789e54b..31da271 100644 --- a/apps/interfaces/daemon/src/handler.rs +++ b/apps/interfaces/daemon/src/handler.rs @@ -292,6 +292,45 @@ fn handle_submit_input(state: &mut AppStateRest, text: String) { state.input.cursor = 0; state.input.history_idx = None; state.dirty = true; + + // Resolve LLM provider credentials + let provider_name = &state.settings.provider; + let provider_cfg = state.app_config.providers.get(provider_name).cloned(); + let mut api_key = String::new(); + if let Some(key) = state.settings.api_keys.get(provider_name) { + api_key = key.clone(); + } else if let Some(ref cfg) = provider_cfg { + if let Some(ref default_key) = cfg.default_api_key { + api_key = default_key.clone(); + } + if api_key.is_empty() { + if let Some(ref env_name) = cfg.api_key_env { + if let Ok(val) = std::env::var(env_name) { + api_key = val; + } + } + } + } + + let messages = state + .session_runtime + .as_ref() + .map(|rt| rt.messages.clone()) + .unwrap_or_default(); + + let params = zesdex_infrastructure::agent::AgentTurnParams { + messages, + session_dir: state.session_dir.clone(), + workspace_roots: state.workspace_roots.clone(), + turn_events: state.turn_events.clone(), + in_flight: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), + abort: state.abort_flag.clone(), + api_key, + model: state.settings.model.clone(), + api_base: provider_cfg.map(|cfg| cfg.api_base.clone()), + }; + + zesdex_infrastructure::agent::spawn_agent_turn(params); } fn handle_delete_char(state: &mut AppStateRest) { diff --git a/apps/interfaces/tui/src/turn.rs b/apps/interfaces/tui/src/turn.rs index 1008203..c17593b 100644 --- a/apps/interfaces/tui/src/turn.rs +++ b/apps/interfaces/tui/src/turn.rs @@ -1,31 +1,18 @@ -//! Agent turn engine — runs LLM + tool execution on a background thread. -//! -//! Flow: push user message → spawn OS thread → loop: call blocking LLM -//! client → execute tool calls → push TurnEvents → repeat until done. +//! TUI agent turn interface adapter — delegates execution to `zesdex_infrastructure::agent`. -use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; -use std::collections::VecDeque; +use std::sync::atomic::Ordering; +use tracing::info; -use tracing::{debug, info, warn}; - -use zesdex_domain::core::tool_call::sanitize_tool_arguments; use zesdex_domain::core::ChatMessage; -use zesdex_infrastructure::llm::provider::LlmClient; -use zesdex_infrastructure::tools::{all_tools, tool_defs, ToolCtx}; -use zesdex_infrastructure::TurnEvent; +use zesdex_infrastructure::agent::{spawn_agent_turn as backend_spawn_agent_turn, AgentTurnParams}; use crate::state::AppStateRest; -/// Spawn an agent turn on a background OS thread. -/// -/// Flow: compare-exchange the in-flight flag → snapshot state fields → -/// clone session runtime messages → push user message → spawn OS thread -/// that runs `run_turn`. +/// Spawn an agent turn on a background thread by delegating to `zesdex-infrastructure`. #[tracing::instrument(skip(state))] pub fn spawn_agent_turn(state: &mut AppStateRest, text: String) { // compare_exchange: only mark in-flight if not already running - if state.turn_in_flight_flag + if state + .turn_in_flight_flag .compare_exchange(false, true, Ordering::SeqCst, Ordering::Relaxed) .is_err() { @@ -41,7 +28,7 @@ pub fn spawn_agent_turn(state: &mut AppStateRest, text: String) { // Resolve LLM provider configuration from settings let provider_name = &state.settings.provider; let provider_cfg = state.app_config.providers.get(provider_name).cloned(); - + let mut api_key = String::new(); if let Some(key) = state.settings.api_keys.get(provider_name) { api_key = key.clone(); @@ -72,190 +59,19 @@ pub fn spawn_agent_turn(state: &mut AppStateRest, text: String) { rt.messages = messages.clone(); } - info!("spawning agent turn with {} messages (model: {})", messages.len(), model); + info!("delegating agent turn to infrastructure engine (model: {})", model); - std::thread::spawn(move || { - let params = TurnParams { - messages: &mut messages, - session_dir: &session_dir, - workspace_roots: &workspace_roots, - turn_events: &turn_events, - in_flight: &in_flight, - abort: &abort, - api_key, - model, - api_base, - }; - run_turn(params); - }); -} - -/// Parameters for a turn, grouped to avoid too-many-arguments lint. -/// -/// Holds all the references and owned values that `run_turn` needs: -/// message history, turn-event queue, abort/in-flight flags, API credentials, -/// and environment paths. -struct TurnParams<'a> { - messages: &'a mut Vec, - session_dir: &'a Path, - workspace_roots: &'a [PathBuf], - turn_events: &'a Arc>>, - in_flight: &'a Arc, - abort: &'a Arc, - api_key: String, - model: String, - api_base: Option, -} - -/// The core agent turn — LLM call → tool execution → repeat. -/// -/// Flow: build `LlmClient` → compile tools → prepend system message → -/// loop (max 50 iterations): abort check → stream LLM response → -/// push events → execute tool calls → push results → break on -/// no tool calls or error → emit final `Compacted` + `Done`. -#[tracing::instrument(skip(params))] -fn run_turn(params: TurnParams) { - let client = LlmClient::new(params.api_key, params.model, params.api_base); - - let tools = all_tools(); - let defs = tool_defs(&tools); - - let mut sys_prompt = "You are Zesdex, an AI coding assistant. You have access to various tools via native function calling to help the user. \ - For complex problems, use `seq_think` to reason step-by-step. \ - For any non-trivial tasks, you MUST prioritize creating a structured plan (using `plan_enter`) and a list of TODOs (using `todowrite`) BEFORE executing any other tools or modifying files. \ - For large multi-step operations or delegating tasks, you MUST prioritize using `workflow_run` (to run a yaml workflow script) or `hive_mind` (to orchestrate multiple agents) to complete the task efficiently. \ - Use other tools when necessary. If a tool returns an error, fix the issue before retrying. \ - Respond conversationally, concisely, and helpfully.".to_string(); - - if let Some(root) = params.workspace_roots.first() { - let tree = zesdex_infrastructure::utils::build_workspace_tree(root, 800); - let rich_ctx = zesdex_infrastructure::utils::build_rich_context(root); - - sys_prompt.push_str("\n\nWorkspace structure:\n```\n"); - sys_prompt.push_str(&tree); - sys_prompt.push_str("\n```\n\n"); - - sys_prompt.push_str(&rich_ctx); - } - - let sys_msg = ChatMessage::system(sys_prompt); - params.messages.insert(0, sys_msg); - - let tool_ctx = ToolCtx::builder() - .session_dir(params.session_dir.to_path_buf()) - .workspaces(params.workspace_roots.to_vec()) - .turn_events(Arc::clone(params.turn_events)) - .build(); - - for iteration in 0..50 { - if params.abort.load(Ordering::SeqCst) { - params.abort.store(false, Ordering::SeqCst); - push_event(params.turn_events, TurnEvent::SystemNote { - kind: "info".into(), - message: "Turn aborted by user".into(), - }); - break; - } - - debug!("agent turn iteration {iteration}"); - - // Stream the LLM response - push_event(params.turn_events, TurnEvent::StreamStart); - - let result = client.chat_with_tools_streaming( - params.messages, - Some(defs.clone()), - Some(0.7), - Some(4096), - |event| { - if params.abort.load(Ordering::SeqCst) { - return false; - } - match event { - zesdex_domain::core::StreamEvent::Token(s) => { - push_event(params.turn_events, TurnEvent::StreamToken(s.clone())); - } - zesdex_domain::core::StreamEvent::Reasoning(s) => { - push_event(params.turn_events, TurnEvent::StreamReasoning(s.clone())); - } - _ => {} - } - true - }, - Some(params.abort), - ); - - match result { - Ok((assistant_msg, usage)) => { - let content = assistant_msg.content.clone().unwrap_or_default(); - let tool_calls = assistant_msg.tool_calls.clone().unwrap_or_default(); - - push_event(params.turn_events, TurnEvent::StreamDone(assistant_msg.clone())); - - if let Some((tokens_in, tokens_out)) = usage { - push_event(params.turn_events, TurnEvent::Usage { - tokens_in, - tokens_out, - }); - } - - if tool_calls.is_empty() { - params.messages.push(ChatMessage::assistant(Some(content))); - break; - } - - params.messages.push(assistant_msg); - - for tc in &tool_calls { - let name = &tc.function.name; - let args = sanitize_tool_arguments(&tc.function.arguments); - - debug!("executing tool: {name}"); - - let output = if let Some(tool) = tools.iter().find(|t| t.name() == name) { - match tool.run(&tool_ctx, &args) { - Ok(o) => o, - Err(e) => format!("Error: {e}"), - } - } else { - format!("Unknown tool: {name}") - }; - - let is_error = output.starts_with("Error:"); - - push_event(params.turn_events, TurnEvent::ToolResult { - tool_call_id: tc.id.clone(), - tool_name: name.clone(), - output: output.clone(), - is_error, - path: None, - }); - - params.messages.push(ChatMessage::tool(tc.id.clone(), output.clone())); - } - } - Err(e) => { - warn!("LLM call failed: {e}"); - push_event(params.turn_events, TurnEvent::Error(format!("LLM error: {e}"))); - break; - } - } - } - - // Propagate accumulated messages back to session_runtime so the next - // turn starts with the full history (assistant replies + tool results). - push_event(params.turn_events, TurnEvent::Compacted(params.messages.clone())); - push_event(params.turn_events, TurnEvent::Done); - mark_done(params.in_flight); -} - -fn push_event(queue: &Arc>>, event: TurnEvent) { - if let Ok(mut q) = queue.lock() { - q.push_back(event); - } -} - -/// Mark the turn as done using lock-free atomic store. -fn mark_done(flag: &Arc) { - flag.store(false, Ordering::SeqCst); + let params = AgentTurnParams { + messages, + session_dir, + workspace_roots, + turn_events, + in_flight, + abort, + api_key, + model, + api_base, + }; + + backend_spawn_agent_turn(params); } diff --git a/apps/interfaces/tui/src/view/sidebar.rs b/apps/interfaces/tui/src/view/sidebar.rs index d861b1f..53bcb84 100644 --- a/apps/interfaces/tui/src/view/sidebar.rs +++ b/apps/interfaces/tui/src/view/sidebar.rs @@ -97,7 +97,13 @@ fn draw_usage_widget(frame: &mut Frame, area: Rect, state: &crate::state::AppSta let now_ms = chrono::Utc::now().timestamp_millis(); let summary = compute_usage_summary(&rt.usage, rt.session_start, now_ms); let max_tokens = crate::state::resolve_context_window(&state.app_config, &state.settings); - let current_tokens = state.cached_token_count; + let current_tokens = if rt.usage.last_tokens_in > 0 { + // Actual context window used by the LLM (includes system prompt + tree) + rt.usage.last_tokens_in as usize + } else { + // Fallback for brand new sessions before the first API call + state.cached_token_count + }; let mut items = vec![ Line::from(Span::styled( diff --git a/apps/interfaces/ws/src/lib.rs b/apps/interfaces/ws/src/lib.rs index af26aae..a477e93 100644 --- a/apps/interfaces/ws/src/lib.rs +++ b/apps/interfaces/ws/src/lib.rs @@ -68,10 +68,73 @@ async fn handle_socket(mut socket: WebSocket, state: Arc) { while let Some(Ok(msg)) = receiver.next().await { match msg { Message::Text(text) => { - // Convert Utf8Bytes -> String for JSON serialisation let text_str = text.to_string(); info!("Received WS message: {text_str}"); - // Echo back for now + + if let Ok(val) = serde_json::from_str::(&text_str) { + if val.get("type").and_then(|v| v.as_str()) == Some("prompt") { + if let Some(prompt) = val.get("message").and_then(|v| v.as_str()) { + let session_dir = std::env::current_dir().unwrap_or_default(); + let workspace_roots = vec![session_dir.clone()]; + let turn_events = Arc::new(std::sync::Mutex::new(std::collections::VecDeque::new())); + let in_flight = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let abort = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let api_key = std::env::var("OPENAI_API_KEY").unwrap_or_default(); + let model = val.get("model").and_then(|v| v.as_str()).unwrap_or("gpt-4o").to_string(); + + let params = zesdex_infrastructure::agent::AgentTurnParams { + messages: vec![zesdex_domain::core::ChatMessage::user(prompt)], + session_dir, + workspace_roots, + turn_events: turn_events.clone(), + in_flight, + abort, + api_key, + model, + api_base: None, + }; + + zesdex_infrastructure::agent::spawn_agent_turn(params); + + let tx_clone = tx.clone(); + tokio::spawn(async move { + let mut done = false; + while !done { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + let events: Vec<_> = { + if let Ok(mut q) = turn_events.lock() { + q.drain(..).collect() + } else { + vec![] + } + }; + for ev in events { + match ev { + zesdex_infrastructure::TurnEvent::StreamToken(tok) => { + let json = serde_json::json!({ "type": "token", "content": tok }); + let _ = tx_clone.send(json.to_string()); + } + zesdex_infrastructure::TurnEvent::Done => { + let json = serde_json::json!({ "type": "done" }); + let _ = tx_clone.send(json.to_string()); + done = true; + } + zesdex_infrastructure::TurnEvent::Error(err) => { + let json = serde_json::json!({ "type": "error", "message": err }); + let _ = tx_clone.send(json.to_string()); + done = true; + } + _ => {} + } + } + } + }); + continue; + } + } + } + + // Fallback echo let response = serde_json::json!({ "type": "echo", "data": text_str