From bb621fdff2ed1f28dbaf33024dff9666c5af8e9f Mon Sep 17 00:00:00 2001 From: asepharyana Date: Sun, 12 Jul 2026 10:23:26 +0700 Subject: [PATCH] feat: enhance error handling in OAuth URL building and client creation; improve tool argument sanitization --- src/app/mcp/manager.rs | 6 ++--- src/app/runtime/actions/mod.rs | 43 +++++++++++++--------------------- src/app/subagent/engine.rs | 4 ++-- src/dto/chat/tool.rs | 11 ++++++++- src/model/app_config.rs | 17 ++++++++++---- src/model/session_lock.rs | 3 +++ src/service/oauth/loopback.rs | 8 ++++--- src/service/oauth/manager.rs | 17 +++++++++++++- src/service/provider.rs | 10 ++++++-- src/tool/fs/edit.rs | 3 +++ 10 files changed, 79 insertions(+), 43 deletions(-) diff --git a/src/app/mcp/manager.rs b/src/app/mcp/manager.rs index eb62b04..d88c04c 100644 --- a/src/app/mcp/manager.rs +++ b/src/app/mcp/manager.rs @@ -13,8 +13,8 @@ const MCP_CALL_TIMEOUT_MS: u64 = 60_000; fn mcp_static_str(s: &str) -> &'static str { static CACHE: OnceLock>> = OnceLock::new(); let mut cache = CACHE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap(); - if let Some(existing) = cache.iter().find(|e| **e == s) { - return *existing; + if let Some(&existing) = cache.iter().find(|e| **e == s) { + return existing; } let leaked: &'static str = Box::leak(s.to_string().into_boxed_str()); cache.push(leaked); @@ -163,7 +163,7 @@ fn call_via_stdio( let mut guard; let child: &mut StdioChild = if let Some(mtx) = existing_handle { guard = mtx.lock().map_err(|e| anyhow::anyhow!("MCP handle lock: {}", e))?; - &mut *guard + &mut guard } else { let mut fresh = spawn_stdio_child(command, extra_args)?; let result = fresh.call("tools/call", json!({ diff --git a/src/app/runtime/actions/mod.rs b/src/app/runtime/actions/mod.rs index c7113d3..8725296 100644 --- a/src/app/runtime/actions/mod.rs +++ b/src/app/runtime/actions/mod.rs @@ -9,8 +9,10 @@ use crate::app::state::runtime::TurnEvent; use crate::app::state::types::{Origin, Overlay, Toast, ToastKind}; use crate::dto::chat::message::{ChatMessage, Role}; -const MAX_TOOL_ONLY_TURNS: usize = usize::MAX; -const MAX_AGENT_STEPS: usize = usize::MAX; +// Step bounds intentionally left unbounded (usize::MAX) so the agent can +// continue across as many turns as needed. Each iteration still honours +// `tc.abort_flag` and the per-call LLM timeout, so a runaway loop is +// observable and cancellable from the UI. #[derive(Debug, Clone)] pub enum Action { @@ -555,19 +557,17 @@ fn generate_workspace_tree(roots: &[std::path::PathBuf]) -> String { .git_ignore(true) .build(); let mut count = 0; - for result in walker { - if let Ok(entry) = result { - let path = entry.path(); - if let Ok(rel) = path.strip_prefix(root) { - if rel.as_os_str().is_empty() { continue; } - let is_dir = entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false); - let prefix = if is_dir { "[DIR] " } else { " " }; - out.push_str(&format!(" {}{}\n", prefix, rel.display())); - count += 1; - if count > 1000 { - out.push_str(" ... (truncated)\n"); - break; - } + for entry in walker.flatten() { + let path = entry.path(); + if let Ok(rel) = path.strip_prefix(root) { + if rel.as_os_str().is_empty() { continue; } + let is_dir = entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false); + let prefix = if is_dir { "[DIR] " } else { " " }; + out.push_str(&format!(" {}{}\n", prefix, rel.display())); + count += 1; + if count > 1000 { + out.push_str(" ... (truncated)\n"); + break; } } } @@ -590,7 +590,6 @@ fn run_agent_turn( ) -> anyhow::Result<()> { let mut msgs = messages.to_vec(); let mut edits_this_turn = 0u32; - let mut tool_only_rounds = 0usize; let mut prev_shaped = false; let tree_info = generate_workspace_tree(&tc.workspace_roots); @@ -606,16 +605,7 @@ fn run_agent_turn( msgs.insert(0, sys); } - for _step in 0..MAX_AGENT_STEPS { - if tool_only_rounds >= MAX_TOOL_ONLY_TURNS { - let stop_msg = ChatMessage::user( - "Stop calling tools. Respond naturally now.".to_string(), - ); - archive_message(&tc.db, &tc.session_id, &stop_msg); - msgs.push(stop_msg); - tool_only_rounds = 0; - } - + loop { let wire_msgs = if crate::app::runtime::shortsend::should_shape(msgs.len(), prev_shaped) { let total_chars: usize = msgs.iter() .filter_map(|m| m.content.as_deref()) @@ -728,7 +718,6 @@ fn run_agent_turn( let content = response.content.clone().unwrap_or_default(); if has_tool_calls { - tool_only_rounds += 1; let tool_calls = response.tool_calls.clone().unwrap_or_default(); archive_message(&tc.db, &tc.session_id, &response); msgs.push(response); diff --git a/src/app/subagent/engine.rs b/src/app/subagent/engine.rs index c8cb3da..565cc8f 100644 --- a/src/app/subagent/engine.rs +++ b/src/app/subagent/engine.rs @@ -5,6 +5,7 @@ use crate::tool::{all_tools, tool_defs, tool_is_risky}; use super::context::SubagentContext; use super::event::SubagentEvent; +#[allow(dead_code)] pub const MAX_AGENT_STEPS: usize = usize::MAX; /// Maps a subagent's allowed tool names to concrete Tool trait objects and @@ -61,8 +62,7 @@ pub fn run_subagent(ctx: SubagentContext, tx: mpsc::Sender) -> an let (tools, tdefs) = build_subagent_tools(&ctx.allowed_tools); let tdefs_opt: Option> = if tdefs.is_empty() { None } else { Some(tdefs) }; - let max_steps = ctx.max_steps.min(MAX_AGENT_STEPS); - for step in 0..max_steps { + for step in 0..ctx.max_steps { let (api_key, model, base_url) = resolve_provider_config(); let client = crate::service::provider::LlmClient::new(api_key, model, base_url); diff --git a/src/dto/chat/tool.rs b/src/dto/chat/tool.rs index 4b69180..bdef954 100644 --- a/src/dto/chat/tool.rs +++ b/src/dto/chat/tool.rs @@ -18,7 +18,16 @@ pub struct ToolFunction { pub fn sanitize_tool_arguments(args: &Value) -> Value { match args { Value::String(s) => { - serde_json::from_str(s).unwrap_or_else(|_| args.clone()) + match serde_json::from_str::(s) { + Ok(v) => v, + Err(e) => { + eprintln!( + "warning: tool argument is a JSON string but failed to parse: {}. Using raw string.", + e + ); + args.clone() + } + } } obj @ Value::Object(_) => obj.clone(), other => other.clone(), diff --git a/src/model/app_config.rs b/src/model/app_config.rs index 7623b58..7b85321 100644 --- a/src/model/app_config.rs +++ b/src/model/app_config.rs @@ -60,10 +60,19 @@ impl AppConfig { pub fn load() -> Self { let store = super::store::Store::new(); let path = store.base_dir.join("app_config.json"); - let mut cfg: AppConfig = std::fs::read_to_string(path) - .ok() - .and_then(|s| serde_json::from_str(&s).ok()) - .unwrap_or_default(); + let mut cfg: AppConfig = match std::fs::read_to_string(&path) { + Ok(s) => match serde_json::from_str(&s) { + Ok(c) => c, + Err(e) => { + eprintln!( + "warning: failed to parse config file '{}': {}. Loading defaults.", + path.display(), e + ); + Self::default() + } + }, + Err(_) => Self::default(), + }; // Merge any default providers not present in the loaded config let defaults = Self::default(); for (name, provider) in defaults.providers { diff --git a/src/model/session_lock.rs b/src/model/session_lock.rs index 80741cb..e076423 100644 --- a/src/model/session_lock.rs +++ b/src/model/session_lock.rs @@ -32,6 +32,9 @@ impl SessionLock { } fn is_alive(&self, pid: u32) -> bool { + // 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 already validated by `try_lock`. unsafe { libc::kill(pid as i32, 0) == 0 } } } diff --git a/src/service/oauth/loopback.rs b/src/service/oauth/loopback.rs index b420e15..5046385 100644 --- a/src/service/oauth/loopback.rs +++ b/src/service/oauth/loopback.rs @@ -57,9 +57,11 @@ fn urlencoding(s: &str) -> String { let mut chars = s.chars(); while let Some(c) = chars.next() { if c == '%' { - let hi = chars.next().and_then(|c| c.to_digit(16)).unwrap_or(0); - let lo = chars.next().and_then(|c| c.to_digit(16)).unwrap_or(0); - result.push(char::from((hi * 16 + lo) as u8)); + match (chars.next().and_then(|c| c.to_digit(16)), + chars.next().and_then(|c| c.to_digit(16))) { + (Some(hi), Some(lo)) => result.push(char::from((hi * 16 + lo) as u8)), + _ => { result.push('%'); } + } } else { result.push(c); } diff --git a/src/service/oauth/manager.rs b/src/service/oauth/manager.rs index 561afa7..b1ba975 100644 --- a/src/service/oauth/manager.rs +++ b/src/service/oauth/manager.rs @@ -84,7 +84,22 @@ impl OAuthManager { } pub fn build_auth_url(&self, redirect_uri: &str, state: &str, code_challenge: &str) -> String { - let mut url = url::Url::parse(&self.config.auth_url).unwrap_or_else(|_| url::Url::parse("https://example.com").unwrap()); + // Refuse to build a URL if `auth_url` is missing or invalid. Previously this + // silently fell back to https://example.com, which produced a valid-looking + // auth URL pointing at the wrong server and leaked client credentials in + // query params. Returning an empty string signals failure to callers, who + // can prompt the user to fix the OAuth config instead of starting a flow + // against a wrong host. + let mut url = match url::Url::parse(&self.config.auth_url) { + Ok(u) if !self.config.auth_url.is_empty() => u, + _ => { + eprintln!( + "warning: OAuth auth_url is missing or invalid ('{}'); aborting build_auth_url", + self.config.auth_url + ); + return String::new(); + } + }; url.query_pairs_mut() .append_pair("response_type", "code") .append_pair("client_id", &self.config.client_id) diff --git a/src/service/provider.rs b/src/service/provider.rs index a37cefb..586eb0c 100644 --- a/src/service/provider.rs +++ b/src/service/provider.rs @@ -29,11 +29,17 @@ impl LlmClient { } else { model }; - let client = reqwest::blocking::Client::builder() + let client = match reqwest::blocking::Client::builder() .timeout(REQUEST_TIMEOUT) .connect_timeout(CONNECT_TIMEOUT) .build() - .unwrap_or_else(|_| reqwest::blocking::Client::new()); + { + Ok(c) => c, + Err(e) => { + eprintln!("warning: failed to build reqwest client with timeouts: {}. Using default client without timeouts.", e); + reqwest::blocking::Client::new() + } + }; LlmClient { client, api_key, diff --git a/src/tool/fs/edit.rs b/src/tool/fs/edit.rs index f1b548a..f50d9e6 100644 --- a/src/tool/fs/edit.rs +++ b/src/tool/fs/edit.rs @@ -56,6 +56,9 @@ impl Tool for Edit { if reason.trim().is_empty() { anyhow::bail!("reason must be a non-empty string"); } + if old.is_empty() { + anyhow::bail!("'old' must be a non-empty string; use 'write' to replace entire file contents"); + } let check_matches = check_graduated_checks(&rel, &new_str, &ctx.graduated_checks); let replace_all = args.get("replace_all").and_then(|v| v.as_bool()).unwrap_or(false); let path: PathBuf = resolve_path(&ctx.workspaces, &rel)?;