From 750a94030b42f1166f8d05f240f574d16cd2df1f Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Wed, 18 Feb 2026 00:29:53 -0800 Subject: [PATCH] fix: persist OpenAI-compatible provider and respect embeddings disable (#177) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: persist OpenAI-compatible provider and respect embeddings disable (#129) Three interrelated bugs caused the agent to ignore user choices made during onboarding when using an OpenAI-compatible LLM provider: 1. Session auth ran before DB config reload, so Config::from_env() defaulted to NearAi and attempted Clerk auth before the real backend was known. Moved session auth to after final config resolution. 2. EmbeddingsConfig::resolve() force-enabled embeddings whenever OPENAI_API_KEY was present, ignoring the user's explicit disable. Changed to respect the stored setting as source of truth. 3. LLM_BACKEND was not saved to the bootstrap .env file, so Config::from_env() always defaulted to NearAi before the DB was connected. Now saves LLM_BACKEND, LLM_BASE_URL, and OLLAMA_BASE_URL alongside the database bootstrap vars. Co-Authored-By: Claude Opus 4.6 * fix: add SAFETY comments and sanitize .env value escaping Address PR review feedback: - Add SAFETY comments to all unsafe env var manipulation in config tests (gemini-code-assist). - Escape backslashes and double quotes in save_bootstrap_env() to prevent env var injection via malicious URLs (gemini-code-assist). - Add test verifying injection attempt is neutralized. Co-Authored-By: Claude Opus 4.6 * fix: incorporate PR #138 changes (chat completions, model sorting, tool schemas) Includes all changes from bigguybobby's PR #138: - Use Chat Completions API for OpenAI-compatible providers (avoids Responses API assumptions like required tool call IDs) - Fall back to settings.selected_model when LLM_MODEL env var is unset - Update OpenAI model list (add gpt-5 family) with priority-based sorting - Add is_openai_chat_model() filter with broader exclusion patterns - Fix http tool: headers schema → array of {name,value}, body → string type, parse_headers_param() accepts both legacy object and array formats - Fix json tool: data schema → string type, parse_json_input() normalizer, validate uses strict string-only check - Add mutex-serialized config tests for env var manipulation - Update NEAR AI config comment for accuracy Co-Authored-By: Bobby (bigguybobby) Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Illia Polosukhin Co-authored-by: Claude Opus 4.6 Co-authored-by: Bobby (bigguybobby) --- src/bootstrap.rs | 33 +++++++- src/config.rs | 170 +++++++++++++++++++++++++++++++++++++- src/llm/mod.rs | 6 +- src/main.rs | 17 ++-- src/settings.rs | 34 +++++++- src/setup/README.md | 17 +++- src/setup/wizard.rs | 125 +++++++++++++++++++++++++--- src/tools/builtin/http.rs | 133 +++++++++++++++++++++++++---- src/tools/builtin/json.rs | 47 +++++++++-- 9 files changed, 527 insertions(+), 55 deletions(-) diff --git a/src/bootstrap.rs b/src/bootstrap.rs index bf7d0223..3641c790 100644 --- a/src/bootstrap.rs +++ b/src/bootstrap.rs @@ -98,7 +98,10 @@ pub fn save_bootstrap_env(vars: &[(&str, &str)]) -> std::io::Result<()> { } let mut content = String::new(); for (key, value) in vars { - content.push_str(&format!("{}=\"{}\"\n", key, value)); + // Escape backslashes and double quotes to prevent env var injection + // (e.g. a value containing `"\nINJECTED="x` would break out of quotes). + let escaped = value.replace('\\', "\\\\").replace('"', "\\\""); + content.push_str(&format!("{}=\"{}\"\n", key, escaped)); } std::fs::write(&path, content) } @@ -323,6 +326,34 @@ mod tests { assert!(content.contains("DATABASE_URL=postgres://test")); } + #[test] + fn test_save_bootstrap_env_escapes_quotes() { + let dir = tempdir().unwrap(); + let env_path = dir.path().join(".env"); + + // A malicious URL attempting to inject a second env var + let malicious = r#"http://evil.com" +INJECTED="pwned"#; + let mut content = String::new(); + let escaped = malicious.replace('\\', "\\\\").replace('"', "\\\""); + content.push_str(&format!("LLM_BASE_URL=\"{}\"\n", escaped)); + std::fs::write(&env_path, &content).unwrap(); + + let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path) + .unwrap() + .filter_map(|r| r.ok()) + .collect(); + + // Must parse as exactly one variable, not two + assert_eq!(parsed.len(), 1, "injection must not create extra vars"); + assert_eq!(parsed[0].0, "LLM_BASE_URL"); + // The value should contain the original malicious content (unescaped by dotenvy) + assert!( + parsed[0].1.contains("INJECTED"), + "value should contain the literal injection attempt, not execute it" + ); + } + #[test] fn test_ironclaw_env_path() { let path = ironclaw_env_path(); diff --git a/src/config.rs b/src/config.rs index b273c413..68951ffa 100644 --- a/src/config.rs +++ b/src/config.rs @@ -611,7 +611,7 @@ impl LlmConfig { LlmBackend::NearAi }; - // Always resolve NEAR AI config (used as fallback and for embeddings) + // Resolve NEAR AI config only when backend is NearAi (or when explicitly configured) let nearai_api_key = optional_env("NEARAI_API_KEY")?.map(SecretString::from); let api_mode = if let Some(mode_str) = optional_env("NEARAI_API_MODE")? { @@ -705,7 +705,9 @@ impl LlmConfig { hint: "Set LLM_BASE_URL when LLM_BACKEND=openai_compatible".to_string(), })?; let api_key = optional_env("LLM_API_KEY")?.map(SecretString::from); - let model = optional_env("LLM_MODEL")?.unwrap_or_else(|| "default".to_string()); + let model = optional_env("LLM_MODEL")? + .or_else(|| settings.selected_model.clone()) + .unwrap_or_else(|| "default".to_string()); Some(OpenAiCompatibleConfig { base_url, api_key, @@ -781,7 +783,7 @@ impl EmbeddingsConfig { key: "EMBEDDING_ENABLED".to_string(), message: format!("must be 'true' or 'false': {e}"), })? - .unwrap_or_else(|| settings.embeddings.enabled || openai_api_key.is_some()); + .unwrap_or(settings.embeddings.enabled); Ok(Self { enabled, @@ -1778,3 +1780,165 @@ where .transpose() .map(|opt| opt.unwrap_or(default)) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::settings::{EmbeddingsSettings, Settings}; + use std::sync::Mutex; + + /// Serializes env-mutating tests to prevent parallel races. + static ENV_MUTEX: Mutex<()> = Mutex::new(()); + + /// Clear all embedding-related env vars. + fn clear_embedding_env() { + // SAFETY: Only called under ENV_MUTEX in tests. No other threads + // observe these vars while the lock is held. + unsafe { + std::env::remove_var("EMBEDDING_ENABLED"); + std::env::remove_var("EMBEDDING_PROVIDER"); + std::env::remove_var("EMBEDDING_MODEL"); + std::env::remove_var("OPENAI_API_KEY"); + } + } + + /// Clear all openai-compatible-related env vars. + fn clear_openai_compatible_env() { + // SAFETY: Only called under ENV_MUTEX in tests. + unsafe { + std::env::remove_var("LLM_BACKEND"); + std::env::remove_var("LLM_BASE_URL"); + std::env::remove_var("LLM_MODEL"); + } + } + + #[test] + fn embeddings_disabled_not_overridden_by_openai_key() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + + clear_embedding_env(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::set_var("OPENAI_API_KEY", "sk-test-key-for-issue-129"); + } + + let settings = Settings { + embeddings: EmbeddingsSettings { + enabled: false, + ..Default::default() + }, + ..Default::default() + }; + + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert!( + !config.enabled, + "embeddings should remain disabled when settings.embeddings.enabled=false, \ + even when OPENAI_API_KEY is set (issue #129)" + ); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("OPENAI_API_KEY"); + } + } + + #[test] + fn embeddings_enabled_from_settings() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_embedding_env(); + + let settings = Settings { + embeddings: EmbeddingsSettings { + enabled: true, + ..Default::default() + }, + ..Default::default() + }; + + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert!( + config.enabled, + "embeddings should be enabled when settings say so" + ); + } + + #[test] + fn embeddings_env_override_takes_precedence() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + + clear_embedding_env(); + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("EMBEDDING_ENABLED", "true"); + } + + let settings = Settings { + embeddings: EmbeddingsSettings { + enabled: false, + ..Default::default() + }, + ..Default::default() + }; + + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert!( + config.enabled, + "EMBEDDING_ENABLED=true env var should override settings" + ); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("EMBEDDING_ENABLED"); + } + } + + #[test] + fn openai_compatible_uses_selected_model_when_llm_model_unset() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_openai_compatible_env(); + + let settings = Settings { + llm_backend: Some("openai_compatible".to_string()), + openai_compatible_base_url: Some("https://openrouter.ai/api/v1".to_string()), + selected_model: Some("openai/gpt-5.1-codex".to_string()), + ..Default::default() + }; + + let cfg = LlmConfig::resolve(&settings).expect("resolve should succeed"); + let compat = cfg + .openai_compatible + .expect("openai-compatible config should be present"); + + assert_eq!(compat.model, "openai/gpt-5.1-codex"); + } + + #[test] + fn openai_compatible_llm_model_env_overrides_selected_model() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_openai_compatible_env(); + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("LLM_MODEL", "openai/gpt-5-codex"); + } + + let settings = Settings { + llm_backend: Some("openai_compatible".to_string()), + openai_compatible_base_url: Some("https://openrouter.ai/api/v1".to_string()), + selected_model: Some("openai/gpt-5.1-codex".to_string()), + ..Default::default() + }; + + let cfg = LlmConfig::resolve(&settings).expect("resolve should succeed"); + let compat = cfg + .openai_compatible + .expect("openai-compatible config should be present"); + + assert_eq!(compat.model, "openai/gpt-5-codex"); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("LLM_MODEL"); + } + } +} diff --git a/src/llm/mod.rs b/src/llm/mod.rs index ef91f8e2..21d841f2 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -209,9 +209,11 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result anyhow::Result<()> { }; let session = create_session_manager(session_config).await; - // Session-based auth is only needed for NEAR AI backend without an API key. - // ChatCompletions mode with an API key skips session auth entirely. - if config.llm.backend == ironclaw::config::LlmBackend::NearAi - && config.llm.nearai.api_key.is_none() - { - session.ensure_authenticated().await?; - } - // Initialize tracing let env_filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("ironclaw=info,tower_http=warn")); @@ -538,6 +530,15 @@ async fn main() -> anyhow::Result<()> { } } + // Session-based auth is only needed for NEAR AI backend without an API key. + // Do this after DB-backed config reload so provider selection from onboarding + // is respected (e.g. OpenAI/OpenAI-compatible should not trigger NEAR auth). + if config.llm.backend == ironclaw::config::LlmBackend::NearAi + && config.llm.nearai.api_key.is_none() + { + session.ensure_authenticated().await?; + } + // Start managed tunnel if configured and no static URL is already set. // // The tunnel process runs in the background, exposing the local gateway diff --git a/src/settings.rs b/src/settings.rs index e8d1242f..0f12dcec 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -1054,6 +1054,37 @@ mod tests { ); } + #[test] + fn test_openai_compatible_db_map_round_trip() { + let settings = Settings { + llm_backend: Some("openai_compatible".to_string()), + openai_compatible_base_url: Some("http://my-vllm:8000/v1".to_string()), + embeddings: EmbeddingsSettings { + enabled: false, + ..Default::default() + }, + ..Default::default() + }; + + let map = settings.to_db_map(); + let restored = Settings::from_db_map(&map); + + assert_eq!( + restored.llm_backend, + Some("openai_compatible".to_string()), + "llm_backend must survive DB round-trip" + ); + assert_eq!( + restored.openai_compatible_base_url, + Some("http://my-vllm:8000/v1".to_string()), + "openai_compatible_base_url must survive DB round-trip" + ); + assert!( + !restored.embeddings.enabled, + "embeddings.enabled=false must survive DB round-trip" + ); + } + #[test] fn toml_round_trip() { let dir = tempfile::tempdir().unwrap(); @@ -1124,8 +1155,6 @@ mod tests { let mut toml_overlay = Settings::default(); toml_overlay.agent.name = "from-toml".to_string(); - // heartbeat.interval_secs stays at default (1800) in the overlay, - // so the base value (600) should be preserved. base.merge_from(&toml_overlay); @@ -1142,7 +1171,6 @@ mod tests { let overlay = Settings::default(); base.merge_from(&overlay); - // All base values preserved since overlay is entirely default assert_eq!(base.agent.name, "custom-name"); assert!(base.heartbeat.enabled); } diff --git a/src/setup/README.md b/src/setup/README.md index 5c6516ea..dda642b3 100644 --- a/src/setup/README.md +++ b/src/setup/README.md @@ -299,16 +299,26 @@ Contains only the settings needed BEFORE database connection. Written by ```env DATABASE_BACKEND="libsql" LIBSQL_PATH="/Users/name/.ironclaw/ironclaw.db" +LLM_BACKEND="openai_compatible" +LLM_BASE_URL="http://my-vllm:8000/v1" ``` -Or for PostgreSQL: +Or for PostgreSQL + NEAR AI: ```env DATABASE_BACKEND="postgres" DATABASE_URL="postgres://user:pass@localhost/ironclaw" +LLM_BACKEND="nearai" +``` + +Or for Ollama: +```env +LLM_BACKEND="ollama" +OLLAMA_BASE_URL="http://localhost:11434" ``` **Why separate?** Chicken-and-egg: you need `DATABASE_BACKEND` to know -which database to connect to, so it can't be stored in the database. +which database to connect to, and `LLM_BACKEND` to know whether to +attempt NEAR AI session auth -- neither can be stored in the database. **Layer 2: Database settings table** (everything else) @@ -339,6 +349,9 @@ Final step of the wizard: - DATABASE_URL (if postgres) - LIBSQL_PATH (if libsql) - LIBSQL_URL (if turso sync) + - LLM_BACKEND (always, when set) + - LLM_BASE_URL (if openai_compatible) + - OLLAMA_BASE_URL (if ollama) 4. Print configuration summary ``` diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 2625537e..4988abaf 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1530,6 +1530,18 @@ impl SetupWizard { env_vars.push(("LIBSQL_URL", url.clone())); } + // LLM bootstrap vars: same chicken-and-egg problem as DATABASE_BACKEND. + // Config::from_env() needs the backend before the DB is connected. + if let Some(ref backend) = self.settings.llm_backend { + env_vars.push(("LLM_BACKEND", backend.clone())); + } + if let Some(ref url) = self.settings.openai_compatible_base_url { + env_vars.push(("LLM_BASE_URL", url.clone())); + } + if let Some(ref url) = self.settings.ollama_base_url { + env_vars.push(("OLLAMA_BASE_URL", url.clone())); + } + if !env_vars.is_empty() { let pairs: Vec<(&str, &str)> = env_vars.iter().map(|(k, v)| (*k, v.as_str())).collect(); @@ -1764,8 +1776,10 @@ async fn fetch_anthropic_models(cached_key: Option<&str>) -> Vec<(String, String /// Returns `(model_id, display_label)` pairs. Falls back to static defaults on error. async fn fetch_openai_models(cached_key: Option<&str>) -> Vec<(String, String)> { let static_defaults = vec![ + ("gpt-5".into(), "GPT-5 (flagship)".into()), + ("gpt-5-mini".into(), "GPT-5 Mini (fast)".into()), + ("gpt-4.1".into(), "GPT-4.1".into()), ("gpt-4o".into(), "GPT-4o".into()), - ("gpt-4o-mini".into(), "GPT-4o Mini (fast)".into()), ("o3".into(), "o3 (reasoning)".into()), ]; @@ -1800,19 +1814,12 @@ async fn fetch_openai_models(cached_key: Option<&str>) -> Vec<(String, String)> data: Vec, } - // Prefixes that indicate chat-relevant models - let chat_prefixes = ["gpt-4", "gpt-3.5", "o1", "o3", "o4", "chatgpt"]; - match resp.json::().await { Ok(body) => { let mut models: Vec<(String, String)> = body .data .into_iter() - .filter(|m| { - chat_prefixes.iter().any(|p| m.id.starts_with(p)) - && !m.id.contains("realtime") - && !m.id.contains("audio") - }) + .filter(|m| is_openai_chat_model(&m.id)) .map(|m| { let label = m.id.clone(); (m.id, label) @@ -1821,13 +1828,74 @@ async fn fetch_openai_models(cached_key: Option<&str>) -> Vec<(String, String)> if models.is_empty() { return static_defaults; } - models.sort_by(|a, b| a.0.cmp(&b.0)); + sort_openai_models(&mut models); models } Err(_) => static_defaults, } } +fn is_openai_chat_model(model_id: &str) -> bool { + let id = model_id.to_ascii_lowercase(); + + let is_chat_family = id.starts_with("gpt-") + || id.starts_with("chatgpt-") + || id.starts_with("o1") + || id.starts_with("o3") + || id.starts_with("o4") + || id.starts_with("o5"); + + let is_non_chat_variant = id.contains("realtime") + || id.contains("audio") + || id.contains("transcribe") + || id.contains("tts") + || id.contains("embedding") + || id.contains("moderation") + || id.contains("image"); + + is_chat_family && !is_non_chat_variant +} + +fn openai_model_priority(model_id: &str) -> usize { + let id = model_id.to_ascii_lowercase(); + + const EXACT_PRIORITY: &[&str] = &[ + "gpt-5", + "gpt-5-mini", + "gpt-5-nano", + "o3", + "o4-mini", + "o1", + "gpt-4.1", + "gpt-4.1-mini", + "gpt-4o", + "gpt-4o-mini", + ]; + if let Some(pos) = EXACT_PRIORITY.iter().position(|m| id == *m) { + return pos; + } + + const PREFIX_PRIORITY: &[&str] = &[ + "gpt-5-", "o3-", "o4-", "o1-", "gpt-4.1-", "gpt-4o-", "gpt-3.5-", "chatgpt-", + ]; + if let Some(pos) = PREFIX_PRIORITY + .iter() + .position(|prefix| id.starts_with(prefix)) + { + return EXACT_PRIORITY.len() + pos; + } + + EXACT_PRIORITY.len() + PREFIX_PRIORITY.len() + 1 +} + +fn sort_openai_models(models: &mut [(String, String)]) { + models.sort_by(|a, b| { + openai_model_priority(&a.0) + .cmp(&openai_model_priority(&b.0)) + .then_with(|| a.0.cmp(&b.0)) + }); +} + /// Fetch installed models from a local Ollama instance. /// /// Returns `(model_name, display_label)` pairs. Falls back to static defaults on error. @@ -2158,12 +2226,49 @@ mod tests { let _guard = EnvGuard::clear("OPENAI_API_KEY"); let models = fetch_openai_models(None).await; assert!(!models.is_empty()); + assert_eq!(models[0].0, "gpt-5"); assert!( models.iter().any(|(id, _)| id.contains("gpt")), "static defaults should include a GPT model" ); } + #[test] + fn test_is_openai_chat_model_includes_gpt5_and_filters_non_chat_variants() { + assert!(is_openai_chat_model("gpt-5")); + assert!(is_openai_chat_model("gpt-5-mini-2026-01-01")); + assert!(is_openai_chat_model("o3-2025-04-16")); + assert!(!is_openai_chat_model("chatgpt-image-latest")); + assert!(!is_openai_chat_model("gpt-4o-realtime-preview")); + assert!(!is_openai_chat_model("gpt-4o-mini-transcribe")); + assert!(!is_openai_chat_model("text-embedding-3-large")); + } + + #[test] + fn test_sort_openai_models_prioritizes_best_models_first() { + let mut models = vec![ + ("gpt-4o-mini".to_string(), "gpt-4o-mini".to_string()), + ("gpt-5-mini".to_string(), "gpt-5-mini".to_string()), + ("o3".to_string(), "o3".to_string()), + ("gpt-4.1".to_string(), "gpt-4.1".to_string()), + ("gpt-5".to_string(), "gpt-5".to_string()), + ]; + + sort_openai_models(&mut models); + + let ordered: Vec = models.into_iter().map(|(id, _)| id).collect(); + assert_eq!( + ordered, + vec![ + "gpt-5".to_string(), + "gpt-5-mini".to_string(), + "o3".to_string(), + "gpt-4.1".to_string(), + "gpt-4o-mini".to_string(), + ] + ); + } + #[tokio::test] async fn test_fetch_ollama_models_unreachable_fallback() { // Point at a port nothing listens on diff --git a/src/tools/builtin/http.rs b/src/tools/builtin/http.rs index 1f37480c..8b4593c2 100644 --- a/src/tools/builtin/http.rs +++ b/src/tools/builtin/http.rs @@ -106,6 +106,46 @@ fn is_disallowed_ip(ip: &IpAddr) -> bool { } } +fn parse_headers_param( + headers: Option<&serde_json::Value>, +) -> Result, ToolError> { + match headers { + None => Ok(Vec::new()), + Some(serde_json::Value::Object(map)) => { + let mut out = Vec::with_capacity(map.len()); + for (k, v) in map { + let value = v.as_str().ok_or_else(|| { + ToolError::InvalidParameters(format!("header '{}' must have a string value", k)) + })?; + out.push((k.clone(), value.to_string())); + } + Ok(out) + } + Some(serde_json::Value::Array(items)) => { + let mut out = Vec::with_capacity(items.len()); + for (idx, item) in items.iter().enumerate() { + let obj = item.as_object().ok_or_else(|| { + ToolError::InvalidParameters(format!( + "headers[{}] must be an object with 'name' and 'value'", + idx + )) + })?; + let name = obj.get("name").and_then(|v| v.as_str()).ok_or_else(|| { + ToolError::InvalidParameters(format!("headers[{}].name must be a string", idx)) + })?; + let value = obj.get("value").and_then(|v| v.as_str()).ok_or_else(|| { + ToolError::InvalidParameters(format!("headers[{}].value must be a string", idx)) + })?; + out.push((name.to_string(), value.to_string())); + } + Ok(out) + } + Some(_) => Err(ToolError::InvalidParameters( + "'headers' must be an object or an array of {name, value}".to_string(), + )), + } +} + impl Default for HttpTool { fn default() -> Self { Self::new() @@ -136,12 +176,21 @@ impl Tool for HttpTool { "description": "The URL to request" }, "headers": { - "type": "object", - "additionalProperties": { "type": "string" }, - "description": "HTTP headers to include" + "type": "array", + "description": "Optional headers as a list of {name, value} objects", + "items": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "value": { "type": "string" } + }, + "required": ["name", "value"], + "additionalProperties": false + } }, "body": { - "description": "Request body (for POST/PUT/PATCH)" + "type": "string", + "description": "Request body. Use plain text or serialized JSON." }, "timeout_secs": { "type": "integer", @@ -165,14 +214,7 @@ impl Tool for HttpTool { let parsed_url = validate_url(url)?; // Parse headers - let headers: HashMap = params - .get("headers") - .and_then(|v| serde_json::from_value(v.clone()).ok()) - .unwrap_or_default(); - let headers_vec: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let headers_vec = parse_headers_param(params.get("headers"))?; // Build request let mut request = match method.to_uppercase().as_str() { @@ -190,16 +232,31 @@ impl Tool for HttpTool { }; // Add headers - for (key, value) in headers { - request = request.header(&key, &value); + for (key, value) in &headers_vec { + request = request.header(key.as_str(), value.as_str()); } // Add body if present let body_bytes = if let Some(body) = params.get("body") { - let bytes = serde_json::to_vec(body) - .map_err(|e| ToolError::InvalidParameters(format!("invalid body JSON: {}", e)))?; - request = request.json(body); - Some(bytes) + if let Some(body_str) = body.as_str() { + if let Ok(json_body) = serde_json::from_str::(body_str) { + let bytes = serde_json::to_vec(&json_body).map_err(|e| { + ToolError::InvalidParameters(format!("invalid body JSON: {}", e)) + })?; + request = request.json(&json_body); + Some(bytes) + } else { + let bytes = body_str.as_bytes().to_vec(); + request = request.body(body_str.to_string()); + Some(bytes) + } + } else { + let bytes = serde_json::to_vec(body).map_err(|e| { + ToolError::InvalidParameters(format!("invalid body JSON: {}", e)) + })?; + request = request.json(body); + Some(bytes) + } } else { None }; @@ -304,6 +361,20 @@ impl Tool for HttpTool { mod tests { use super::*; + #[test] + fn test_http_tool_schema_body_has_type() { + let tool = HttpTool::new(); + let schema = tool.parameters_schema(); + assert_eq!(schema["properties"]["body"]["type"], "string"); + } + + #[test] + fn test_http_tool_schema_headers_is_array() { + let tool = HttpTool::new(); + let schema = tool.parameters_schema(); + assert_eq!(schema["properties"]["headers"]["type"], "array"); + } + #[test] fn test_validate_url_rejects_http() { let err = validate_url("http://example.com").unwrap_err(); @@ -363,4 +434,30 @@ mod tests { // MAX_RESPONSE_SIZE should be 5 MB to prevent OOM while allowing typical API responses. assert_eq!(MAX_RESPONSE_SIZE, 5 * 1024 * 1024); } + + #[test] + fn test_parse_headers_param_accepts_object_legacy_shape() { + let headers = serde_json::json!({"Authorization": "Bearer token"}); + let parsed = parse_headers_param(Some(&headers)).unwrap(); + assert_eq!( + parsed, + vec![("Authorization".to_string(), "Bearer token".to_string())] + ); + } + + #[test] + fn test_parse_headers_param_accepts_array_shape() { + let headers = serde_json::json!([ + {"name": "Authorization", "value": "Bearer token"}, + {"name": "X-Test", "value": "1"} + ]); + let parsed = parse_headers_param(Some(&headers)).unwrap(); + assert_eq!( + parsed, + vec![ + ("Authorization".to_string(), "Bearer token".to_string()), + ("X-Test".to_string(), "1".to_string()) + ] + ); + } } diff --git a/src/tools/builtin/json.rs b/src/tools/builtin/json.rs index 5350dbb7..f6d91a0d 100644 --- a/src/tools/builtin/json.rs +++ b/src/tools/builtin/json.rs @@ -28,7 +28,8 @@ impl Tool for JsonTool { "description": "The JSON operation to perform" }, "data": { - "description": "The JSON data to operate on (string for parse, object otherwise)" + "type": "string", + "description": "JSON input string. For query/stringify/validate, pass serialized JSON." }, "path": { "type": "string", @@ -64,7 +65,8 @@ impl Tool for JsonTool { parsed } "stringify" => { - let json_str = serde_json::to_string_pretty(data).map_err(|e| { + let value = parse_json_input(data)?; + let json_str = serde_json::to_string_pretty(&value).map_err(|e| { ToolError::ExecutionFailed(format!("failed to stringify: {}", e)) })?; @@ -75,14 +77,14 @@ impl Tool for JsonTool { ToolError::InvalidParameters("missing 'path' parameter for query".to_string()) })?; - query_json(data, path)? + let value = parse_json_input(data)?; + query_json(&value, path)? } "validate" => { - let is_valid = if let Some(s) = data.as_str() { - serde_json::from_str::(s).is_ok() - } else { - true // Already a valid JSON value - }; + let is_valid = data + .as_str() + .map(|s| serde_json::from_str::(s).is_ok()) + .unwrap_or(false); serde_json::json!({ "valid": is_valid }) } @@ -102,6 +104,14 @@ impl Tool for JsonTool { } } +fn parse_json_input(data: &serde_json::Value) -> Result { + let json_str = data + .as_str() + .ok_or_else(|| ToolError::InvalidParameters("'data' must be a JSON string".to_string()))?; + serde_json::from_str(json_str) + .map_err(|e| ToolError::InvalidParameters(format!("invalid JSON input: {}", e))) +} + /// Simple JSONPath-like query implementation. fn query_json(data: &serde_json::Value, path: &str) -> Result { let mut current = data; @@ -144,6 +154,13 @@ fn query_json(data: &serde_json::Value, path: &str) -> Result