fix: persist OpenAI-compatible provider and respect embeddings disable (#177)

* 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 <[email protected]>

* 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 <[email protected]>

* 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) <[email protected]>
Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Illia Polosukhin <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: Bobby (bigguybobby) <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-18 08:29:53 +00:00
committed by GitHub
co-authored by Illia Polosukhin Claude Opus 4.6 Bobby
parent c3340c60ef
commit 750a94030b
9 changed files with 527 additions and 55 deletions
+32 -1
View File
@@ -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();
+167 -3
View File
@@ -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");
}
}
}
+4 -2
View File
@@ -209,9 +209,11 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result<Arc<dyn LlmPr
reason: format!("Failed to create OpenAI-compatible client: {}", e),
})?;
let model = client.completion_model(&compat.model);
// OpenAI-compatible providers (e.g. OpenRouter) are most reliable on Chat Completions.
// This avoids Responses-API-specific assumptions such as required tool call IDs.
let model = client.completions_api().completion_model(&compat.model);
tracing::info!(
"Using OpenAI-compatible endpoint (base_url: {}, model: {})",
"Using OpenAI-compatible endpoint via Chat Completions API (base_url: {}, model: {})",
compat.base_url,
compat.model
);
+9 -8
View File
@@ -330,14 +330,6 @@ async fn main() -> 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
+31 -3
View File
@@ -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);
}
+15 -2
View File
@@ -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
```
+115 -10
View File
@@ -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<ModelEntry>,
}
// Prefixes that indicate chat-relevant models
let chat_prefixes = ["gpt-4", "gpt-3.5", "o1", "o3", "o4", "chatgpt"];
match resp.json::<ModelsResponse>().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<String> = 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
+115 -18
View File
@@ -106,6 +106,46 @@ fn is_disallowed_ip(ip: &IpAddr) -> bool {
}
}
fn parse_headers_param(
headers: Option<&serde_json::Value>,
) -> Result<Vec<(String, String)>, 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<String, String> = 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::<serde_json::Value>(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())
]
);
}
}
+39 -8
View File
@@ -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::<serde_json::Value>(s).is_ok()
} else {
true // Already a valid JSON value
};
let is_valid = data
.as_str()
.map(|s| serde_json::from_str::<serde_json::Value>(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<serde_json::Value, ToolError> {
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<serde_json::Value, ToolError> {
let mut current = data;
@@ -144,6 +154,13 @@ fn query_json(data: &serde_json::Value, path: &str) -> Result<serde_json::Value,
mod tests {
use super::*;
#[test]
fn test_json_tool_schema_data_has_type() {
let tool = JsonTool;
let schema = tool.parameters_schema();
assert_eq!(schema["properties"]["data"]["type"], "string");
}
#[test]
fn test_query_json() {
let data = serde_json::json!({
@@ -166,4 +183,18 @@ mod tests {
serde_json::json!(3)
);
}
#[test]
fn test_parse_json_input_accepts_valid_json_string() {
let input = serde_json::json!("{\"ok\":true}");
let parsed = parse_json_input(&input).unwrap();
assert_eq!(parsed, serde_json::json!({"ok": true}));
}
#[test]
fn test_parse_json_input_rejects_invalid_json_string() {
let input = serde_json::json!("{not valid json}");
let err = parse_json_input(&input).unwrap_err();
assert!(err.to_string().contains("invalid JSON input"));
}
}