feat: enhance error handling in OAuth URL building and client creation; improve tool argument sanitization

This commit is contained in:
asepharyana
2026-07-12 10:23:26 +07:00
parent 0cc60c12ce
commit bb621fdff2
10 changed files with 79 additions and 43 deletions
+3 -3
View File
@@ -13,8 +13,8 @@ const MCP_CALL_TIMEOUT_MS: u64 = 60_000;
fn mcp_static_str(s: &str) -> &'static str { fn mcp_static_str(s: &str) -> &'static str {
static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new(); static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new();
let mut cache = CACHE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap(); let mut cache = CACHE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap();
if let Some(existing) = cache.iter().find(|e| **e == s) { if let Some(&existing) = cache.iter().find(|e| **e == s) {
return *existing; return existing;
} }
let leaked: &'static str = Box::leak(s.to_string().into_boxed_str()); let leaked: &'static str = Box::leak(s.to_string().into_boxed_str());
cache.push(leaked); cache.push(leaked);
@@ -163,7 +163,7 @@ fn call_via_stdio(
let mut guard; let mut guard;
let child: &mut StdioChild = if let Some(mtx) = existing_handle { let child: &mut StdioChild = if let Some(mtx) = existing_handle {
guard = mtx.lock().map_err(|e| anyhow::anyhow!("MCP handle lock: {}", e))?; guard = mtx.lock().map_err(|e| anyhow::anyhow!("MCP handle lock: {}", e))?;
&mut *guard &mut guard
} else { } else {
let mut fresh = spawn_stdio_child(command, extra_args)?; let mut fresh = spawn_stdio_child(command, extra_args)?;
let result = fresh.call("tools/call", json!({ let result = fresh.call("tools/call", json!({
+16 -27
View File
@@ -9,8 +9,10 @@ use crate::app::state::runtime::TurnEvent;
use crate::app::state::types::{Origin, Overlay, Toast, ToastKind}; use crate::app::state::types::{Origin, Overlay, Toast, ToastKind};
use crate::dto::chat::message::{ChatMessage, Role}; use crate::dto::chat::message::{ChatMessage, Role};
const MAX_TOOL_ONLY_TURNS: usize = usize::MAX; // Step bounds intentionally left unbounded (usize::MAX) so the agent can
const MAX_AGENT_STEPS: usize = usize::MAX; // 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)] #[derive(Debug, Clone)]
pub enum Action { pub enum Action {
@@ -555,19 +557,17 @@ fn generate_workspace_tree(roots: &[std::path::PathBuf]) -> String {
.git_ignore(true) .git_ignore(true)
.build(); .build();
let mut count = 0; let mut count = 0;
for result in walker { for entry in walker.flatten() {
if let Ok(entry) = result { let path = entry.path();
let path = entry.path(); if let Ok(rel) = path.strip_prefix(root) {
if let Ok(rel) = path.strip_prefix(root) { if rel.as_os_str().is_empty() { continue; }
if rel.as_os_str().is_empty() { continue; } let is_dir = entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false);
let is_dir = entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false); let prefix = if is_dir { "[DIR] " } else { " " };
let prefix = if is_dir { "[DIR] " } else { " " }; out.push_str(&format!(" {}{}\n", prefix, rel.display()));
out.push_str(&format!(" {}{}\n", prefix, rel.display())); count += 1;
count += 1; if count > 1000 {
if count > 1000 { out.push_str(" ... (truncated)\n");
out.push_str(" ... (truncated)\n"); break;
break;
}
} }
} }
} }
@@ -590,7 +590,6 @@ fn run_agent_turn(
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
let mut msgs = messages.to_vec(); let mut msgs = messages.to_vec();
let mut edits_this_turn = 0u32; let mut edits_this_turn = 0u32;
let mut tool_only_rounds = 0usize;
let mut prev_shaped = false; let mut prev_shaped = false;
let tree_info = generate_workspace_tree(&tc.workspace_roots); let tree_info = generate_workspace_tree(&tc.workspace_roots);
@@ -606,16 +605,7 @@ fn run_agent_turn(
msgs.insert(0, sys); msgs.insert(0, sys);
} }
for _step in 0..MAX_AGENT_STEPS { loop {
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;
}
let wire_msgs = if crate::app::runtime::shortsend::should_shape(msgs.len(), prev_shaped) { let wire_msgs = if crate::app::runtime::shortsend::should_shape(msgs.len(), prev_shaped) {
let total_chars: usize = msgs.iter() let total_chars: usize = msgs.iter()
.filter_map(|m| m.content.as_deref()) .filter_map(|m| m.content.as_deref())
@@ -728,7 +718,6 @@ fn run_agent_turn(
let content = response.content.clone().unwrap_or_default(); let content = response.content.clone().unwrap_or_default();
if has_tool_calls { if has_tool_calls {
tool_only_rounds += 1;
let tool_calls = response.tool_calls.clone().unwrap_or_default(); let tool_calls = response.tool_calls.clone().unwrap_or_default();
archive_message(&tc.db, &tc.session_id, &response); archive_message(&tc.db, &tc.session_id, &response);
msgs.push(response); msgs.push(response);
+2 -2
View File
@@ -5,6 +5,7 @@ use crate::tool::{all_tools, tool_defs, tool_is_risky};
use super::context::SubagentContext; use super::context::SubagentContext;
use super::event::SubagentEvent; use super::event::SubagentEvent;
#[allow(dead_code)]
pub const MAX_AGENT_STEPS: usize = usize::MAX; pub const MAX_AGENT_STEPS: usize = usize::MAX;
/// Maps a subagent's allowed tool names to concrete Tool trait objects and /// 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<SubagentEvent>) -> an
let (tools, tdefs) = build_subagent_tools(&ctx.allowed_tools); let (tools, tdefs) = build_subagent_tools(&ctx.allowed_tools);
let tdefs_opt: Option<Vec<ToolDef>> = if tdefs.is_empty() { None } else { Some(tdefs) }; let tdefs_opt: Option<Vec<ToolDef>> = if tdefs.is_empty() { None } else { Some(tdefs) };
let max_steps = ctx.max_steps.min(MAX_AGENT_STEPS); for step in 0..ctx.max_steps {
for step in 0..max_steps {
let (api_key, model, base_url) = resolve_provider_config(); let (api_key, model, base_url) = resolve_provider_config();
let client = crate::service::provider::LlmClient::new(api_key, model, base_url); let client = crate::service::provider::LlmClient::new(api_key, model, base_url);
+10 -1
View File
@@ -18,7 +18,16 @@ pub struct ToolFunction {
pub fn sanitize_tool_arguments(args: &Value) -> Value { pub fn sanitize_tool_arguments(args: &Value) -> Value {
match args { match args {
Value::String(s) => { Value::String(s) => {
serde_json::from_str(s).unwrap_or_else(|_| args.clone()) match serde_json::from_str::<Value>(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(), obj @ Value::Object(_) => obj.clone(),
other => other.clone(), other => other.clone(),
+13 -4
View File
@@ -60,10 +60,19 @@ impl AppConfig {
pub fn load() -> Self { pub fn load() -> Self {
let store = super::store::Store::new(); let store = super::store::Store::new();
let path = store.base_dir.join("app_config.json"); let path = store.base_dir.join("app_config.json");
let mut cfg: AppConfig = std::fs::read_to_string(path) let mut cfg: AppConfig = match std::fs::read_to_string(&path) {
.ok() Ok(s) => match serde_json::from_str(&s) {
.and_then(|s| serde_json::from_str(&s).ok()) Ok(c) => c,
.unwrap_or_default(); 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 // Merge any default providers not present in the loaded config
let defaults = Self::default(); let defaults = Self::default();
for (name, provider) in defaults.providers { for (name, provider) in defaults.providers {
+3
View File
@@ -32,6 +32,9 @@ impl SessionLock {
} }
fn is_alive(&self, pid: u32) -> bool { 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 } unsafe { libc::kill(pid as i32, 0) == 0 }
} }
} }
+5 -3
View File
@@ -57,9 +57,11 @@ fn urlencoding(s: &str) -> String {
let mut chars = s.chars(); let mut chars = s.chars();
while let Some(c) = chars.next() { while let Some(c) = chars.next() {
if c == '%' { if c == '%' {
let hi = chars.next().and_then(|c| c.to_digit(16)).unwrap_or(0); match (chars.next().and_then(|c| c.to_digit(16)),
let lo = chars.next().and_then(|c| c.to_digit(16)).unwrap_or(0); chars.next().and_then(|c| c.to_digit(16))) {
result.push(char::from((hi * 16 + lo) as u8)); (Some(hi), Some(lo)) => result.push(char::from((hi * 16 + lo) as u8)),
_ => { result.push('%'); }
}
} else { } else {
result.push(c); result.push(c);
} }
+16 -1
View File
@@ -84,7 +84,22 @@ impl OAuthManager {
} }
pub fn build_auth_url(&self, redirect_uri: &str, state: &str, code_challenge: &str) -> String { 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() url.query_pairs_mut()
.append_pair("response_type", "code") .append_pair("response_type", "code")
.append_pair("client_id", &self.config.client_id) .append_pair("client_id", &self.config.client_id)
+8 -2
View File
@@ -29,11 +29,17 @@ impl LlmClient {
} else { } else {
model model
}; };
let client = reqwest::blocking::Client::builder() let client = match reqwest::blocking::Client::builder()
.timeout(REQUEST_TIMEOUT) .timeout(REQUEST_TIMEOUT)
.connect_timeout(CONNECT_TIMEOUT) .connect_timeout(CONNECT_TIMEOUT)
.build() .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 { LlmClient {
client, client,
api_key, api_key,
+3
View File
@@ -56,6 +56,9 @@ impl Tool for Edit {
if reason.trim().is_empty() { if reason.trim().is_empty() {
anyhow::bail!("reason must be a non-empty string"); 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 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 replace_all = args.get("replace_all").and_then(|v| v.as_bool()).unwrap_or(false);
let path: PathBuf = resolve_path(&ctx.workspaces, &rel)?; let path: PathBuf = resolve_path(&ctx.workspaces, &rel)?;