diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e867440..2b80fad5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Refactored OpenAI-compatible chat completion routing to use the rig adapter and `RetryProvider` composition for custom base URL usage. +- Added Ollama embeddings provider support (`EMBEDDING_PROVIDER=ollama`, `OLLAMA_BASE_URL`) in workspace embeddings. +- Added migration `V9__flexible_embedding_dimension.sql` for flexible embedding vector dimensions. + +### Changed + +- Changed default sandbox image to `ironclaw-worker:latest` in config/settings/sandbox defaults. +- Improved tool-message sanitization and provider compatibility handling across NEAR AI, rig adapter, and shared LLM provider code. + +### Fixed + +- Fixed approval-input aliases (`a`, `/approve`, `/always`, `/deny`, etc.) in submission parsing. +- Fixed multi-tool approval resume flow by preserving and replaying deferred tool calls so all prior `tool_use` IDs receive matching `tool_result` messages. +- Fixed REPL quit/exit handling to route shutdown through the agent loop for graceful termination. + ## [0.6.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.5.0...ironclaw-v0.6.0) - 2026-02-19 ### Added @@ -94,6 +111,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Bump MSRV to 1.92, add GCP deployment files ([#40](https://github.com/nearai/ironclaw/pull/40)) - Add OpenAI-compatible HTTP API (/v1/chat/completions, /v1/models) ([#31](https://github.com/nearai/ironclaw/pull/31)) + ## [0.1.3](https://github.com/nearai/ironclaw/compare/v0.1.2...v0.1.3) - 2026-02-12 ### Other diff --git a/migrations/V9__flexible_embedding_dimension.sql b/migrations/V9__flexible_embedding_dimension.sql new file mode 100644 index 00000000..e158604b --- /dev/null +++ b/migrations/V9__flexible_embedding_dimension.sql @@ -0,0 +1,43 @@ +-- Allow embedding vectors of any dimension (not just 1536). +-- This supports Ollama models (768-dim nomic-embed-text, 1024-dim mxbai-embed-large) +-- alongside OpenAI models (1536-dim text-embedding-3-small, 3072-dim text-embedding-3-large). +-- +-- NOTE: HNSW indexes require a fixed dimension, so we drop the index. +-- Exact (sequential) cosine distance search still works without the index. +-- For a personal assistant workspace the dataset is small enough that this +-- has negligible impact on query latency. + +-- Drop dependent views first +DROP VIEW IF EXISTS chunks_pending_embedding; +DROP VIEW IF EXISTS memory_documents_summary; + +DROP INDEX IF EXISTS idx_memory_chunks_embedding; + +ALTER TABLE memory_chunks + ALTER COLUMN embedding TYPE vector + USING embedding::vector; + +-- Recreate the views +CREATE VIEW memory_documents_summary AS +SELECT + d.id, + d.user_id, + d.path, + d.created_at, + d.updated_at, + COUNT(c.id) as chunk_count, + COUNT(c.embedding) as embedded_chunk_count +FROM memory_documents d +LEFT JOIN memory_chunks c ON c.document_id = d.id +GROUP BY d.id; + +CREATE VIEW chunks_pending_embedding AS +SELECT + c.id as chunk_id, + c.document_id, + d.user_id, + d.path, + LENGTH(c.content) as content_length +FROM memory_chunks c +JOIN memory_documents d ON d.id = c.document_id +WHERE c.embedding IS NULL; diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index da9ce416..2bd1c871 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -255,7 +255,10 @@ impl Agent { } // Execute each tool (with approval checking and hook interception) - for mut tc in tool_calls { + let mut idx = 0usize; + while idx < tool_calls.len() { + let mut tc = tool_calls[idx].clone(); + // Check if tool requires approval if let Some(tool) = self.tools().get(&tc.name).await && tool.requires_approval() @@ -277,7 +280,9 @@ impl Agent { } if !is_auto_approved { - // Need approval - store pending request and return + // Need approval - store pending request and return. + // Preserve remaining tool calls so they can be replayed + // after approval. let pending = PendingApproval { request_id: Uuid::new_v4(), tool_name: tc.name.clone(), @@ -285,6 +290,7 @@ impl Agent { description: tool.description().to_string(), tool_call_id: tc.id.clone(), context_messages: context_messages.clone(), + deferred_tool_calls: tool_calls[idx + 1..].to_vec(), }; return Ok(AgenticLoopResult::NeedApproval { pending }); @@ -441,6 +447,8 @@ impl Agent { &tc.name, result_content, )); + + idx += 1; } } } diff --git a/src/agent/session.rs b/src/agent/session.rs index e77149ec..c73882a3 100644 --- a/src/agent/session.rs +++ b/src/agent/session.rs @@ -16,7 +16,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::llm::ChatMessage; +use crate::llm::{ChatMessage, ToolCall}; /// A session containing one or more threads. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -148,6 +148,10 @@ pub struct PendingApproval { pub tool_call_id: String, /// Context messages at the time of the request (to resume from). pub context_messages: Vec, + /// Remaining tool calls from the same assistant message that were not + /// executed yet when approval was requested. + #[serde(default)] + pub deferred_tool_calls: Vec, } /// A conversation thread within a session. @@ -946,6 +950,7 @@ mod tests { description: "dangerous command".to_string(), tool_call_id: "call_123".to_string(), context_messages: vec![ChatMessage::user("do it")], + deferred_tool_calls: vec![], }; thread.await_approval(approval); @@ -969,6 +974,7 @@ mod tests { description: "test".to_string(), tool_call_id: "call_456".to_string(), context_messages: vec![], + deferred_tool_calls: vec![], }; thread.await_approval(approval); diff --git a/src/agent/submission.rs b/src/agent/submission.rs index a2b6b4d7..de696644 100644 --- a/src/agent/submission.rs +++ b/src/agent/submission.rs @@ -118,19 +118,19 @@ impl SubmissionParser { // Approval responses (simple yes/no/always for pending approvals) // These are short enough to check explicitly match lower.as_str() { - "yes" | "y" | "approve" | "ok" => { + "yes" | "y" | "approve" | "ok" | "/approve" | "/yes" | "/y" => { return Submission::ApprovalResponse { approved: true, always: false, }; } - "always" | "yes always" | "approve always" => { + "always" | "a" | "yes always" | "approve always" | "/always" | "/a" => { return Submission::ApprovalResponse { approved: true, always: true, }; } - "no" | "n" | "deny" | "reject" | "cancel" => { + "no" | "n" | "deny" | "reject" | "cancel" | "/deny" | "/no" | "/n" => { return Submission::ApprovalResponse { approved: false, always: false, @@ -475,6 +475,57 @@ mod tests { assert!(matches!(submission, Submission::UserInput { content } if content == "/unknown")); } + #[test] + fn test_parser_approval_response_aliases() { + // approve once + assert!(matches!( + SubmissionParser::parse("y"), + Submission::ApprovalResponse { + approved: true, + always: false + } + )); + assert!(matches!( + SubmissionParser::parse("/approve"), + Submission::ApprovalResponse { + approved: true, + always: false + } + )); + + // approve always + assert!(matches!( + SubmissionParser::parse("a"), + Submission::ApprovalResponse { + approved: true, + always: true + } + )); + assert!(matches!( + SubmissionParser::parse("/always"), + Submission::ApprovalResponse { + approved: true, + always: true + } + )); + + // deny + assert!(matches!( + SubmissionParser::parse("n"), + Submission::ApprovalResponse { + approved: false, + always: false + } + )); + assert!(matches!( + SubmissionParser::parse("/deny"), + Submission::ApprovalResponse { + approved: false, + always: false + } + )); + } + #[test] fn test_parser_json_exec_approval() { let req_id = Uuid::new_v4(); diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 195591ba..ba1a6bff 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -11,7 +11,7 @@ use uuid::Uuid; use crate::agent::Agent; use crate::agent::compaction::ContextCompactor; use crate::agent::dispatcher::{AgenticLoopResult, detect_auth_awaiting, parse_auth_result}; -use crate::agent::session::{Session, ThreadState}; +use crate::agent::session::{PendingApproval, Session, ThreadState}; use crate::agent::submission::SubmissionResult; use crate::channels::{IncomingMessage, StatusUpdate}; use crate::context::JobContext; @@ -712,6 +712,7 @@ impl Agent { // Build context including the tool result let mut context_messages = pending.context_messages; + let deferred_tool_calls = pending.deferred_tool_calls; // Record result in thread { @@ -780,6 +781,178 @@ impl Agent { result_content, )); + // Replay deferred tool calls from the same assistant message so + // every tool_use ID gets a matching tool_result before the next + // LLM call. + if !deferred_tool_calls.is_empty() { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Thinking(format!( + "Executing {} deferred tool(s)...", + deferred_tool_calls.len() + )), + &message.metadata, + ) + .await; + } + + let mut deferred_queue = std::collections::VecDeque::from(deferred_tool_calls); + while let Some(tc) = deferred_queue.pop_front() { + // Re-check approval for each deferred tool call + if let Some(tool) = self.tools().get(&tc.name).await + && tool.requires_approval() + { + let is_auto_approved = { + let sess = session.lock().await; + let mut approved = sess.is_tool_auto_approved(&tc.name); + if approved && tool.requires_approval_for(&tc.arguments) { + approved = false; + } + approved + }; + + if !is_auto_approved { + let new_pending = PendingApproval { + request_id: Uuid::new_v4(), + tool_name: tc.name.clone(), + parameters: tc.arguments.clone(), + description: tool.description().to_string(), + tool_call_id: tc.id.clone(), + context_messages: context_messages.clone(), + deferred_tool_calls: deferred_queue.iter().cloned().collect(), + }; + + let request_id = new_pending.request_id; + let tool_name = new_pending.tool_name.clone(); + let description = new_pending.description.clone(); + let parameters = new_pending.parameters.clone(); + + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.await_approval(new_pending); + } + } + + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Awaiting approval".into()), + &message.metadata, + ) + .await; + + return Ok(SubmissionResult::NeedApproval { + request_id, + tool_name, + description, + parameters, + }); + } + } + + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolStarted { + name: tc.name.clone(), + }, + &message.metadata, + ) + .await; + + let deferred_result = self + .execute_chat_tool(&tc.name, &tc.arguments, &job_ctx) + .await; + + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolCompleted { + name: tc.name.clone(), + success: deferred_result.is_ok(), + }, + &message.metadata, + ) + .await; + + if let Ok(ref output) = deferred_result + && !output.is_empty() + { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolResult { + name: tc.name.clone(), + preview: output.clone(), + }, + &message.metadata, + ) + .await; + } + + // Record in thread + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) + && let Some(turn) = thread.last_turn_mut() + { + match &deferred_result { + Ok(output) => turn.record_tool_result(serde_json::json!(output)), + Err(e) => turn.record_tool_error(e.to_string()), + } + } + } + + // Auth detection for deferred tools + if let Some((ext_name, instructions)) = + detect_auth_awaiting(&tc.name, &deferred_result) + { + let auth_data = parse_auth_result(&deferred_result); + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.enter_auth_mode(ext_name.clone()); + thread.complete_turn(&instructions); + } + } + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::AuthRequired { + extension_name: ext_name, + instructions: Some(instructions.clone()), + auth_url: auth_data.auth_url, + setup_url: auth_data.setup_url, + }, + &message.metadata, + ) + .await; + return Ok(SubmissionResult::response(instructions)); + } + + let deferred_content = match deferred_result { + Ok(output) => { + let sanitized = self.safety().sanitize_tool_output(&tc.name, &output); + self.safety().wrap_for_llm( + &tc.name, + &sanitized.content, + sanitized.was_modified, + ) + } + Err(e) => format!("Error: {}", e), + }; + + context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content)); + } + // Continue the agentic loop (a tool was already executed this turn) let result = self .run_agentic_loop(message, session.clone(), thread_id, context_messages, true) diff --git a/src/channels/repl.rs b/src/channels/repl.rs index a72547d0..812eaa52 100644 --- a/src/channels/repl.rs +++ b/src/channels/repl.rs @@ -330,7 +330,13 @@ impl Channel for ReplChannel { // Handle local REPL commands (only commands that need // immediate local handling stay here) match line.to_lowercase().as_str() { - "/quit" | "/exit" => break, + "/quit" | "/exit" => { + // Forward shutdown command so the agent loop exits even + // when other channels (e.g. web gateway) are still active. + let msg = IncomingMessage::new("repl", "default", "/quit"); + let _ = tx.blocking_send(msg); + break; + } "/help" => { print_help(); continue; diff --git a/src/channels/web/mod.rs b/src/channels/web/mod.rs index 38801e12..30bd1e7c 100644 --- a/src/channels/web/mod.rs +++ b/src/channels/web/mod.rs @@ -305,6 +305,7 @@ impl Channel for GatewayChannel { description, parameters: serde_json::to_string_pretty(¶meters) .unwrap_or_else(|_| parameters.to_string()), + thread_id, }, StatusUpdate::AuthRequired { extension_name, diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 398f2f54..fa473900 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -159,6 +159,7 @@ function connectSSE() { eventSource.addEventListener('approval_needed', (e) => { const data = JSON.parse(e.data); + if (!isCurrentThread(data.thread_id)) return; showApproval(data); }); diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index a62ddcfc..c64a8bd0 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -137,6 +137,8 @@ pub enum SseEvent { tool_name: String, description: String, parameters: String, + #[serde(skip_serializing_if = "Option::is_none")] + thread_id: Option, }, #[serde(rename = "auth_required")] AuthRequired { @@ -785,12 +787,14 @@ mod tests { tool_name: "shell".to_string(), description: "Run ls".to_string(), parameters: "{}".to_string(), + thread_id: Some("t1".to_string()), }; let ws = WsServerMessage::from_sse_event(&sse); match ws { WsServerMessage::Event { event_type, data } => { assert_eq!(event_type, "approval_needed"); assert_eq!(data["tool_name"], "shell"); + assert_eq!(data["thread_id"], "t1"); } _ => panic!("Expected Event variant"), } diff --git a/src/config/embeddings.rs b/src/config/embeddings.rs index 466f11c5..39f28f06 100644 --- a/src/config/embeddings.rs +++ b/src/config/embeddings.rs @@ -9,25 +9,48 @@ use crate::settings::Settings; pub struct EmbeddingsConfig { /// Whether embeddings are enabled. pub enabled: bool, - /// Provider to use: "openai" or "nearai" + /// Provider to use: "openai", "nearai", or "ollama" pub provider: String, /// OpenAI API key (for OpenAI provider). pub openai_api_key: Option, /// Model to use for embeddings. pub model: String, + /// Ollama base URL (for Ollama provider). Defaults to http://localhost:11434. + pub ollama_base_url: String, + /// Embedding vector dimension. Inferred from the model name when not set explicitly. + pub dimension: usize, } impl Default for EmbeddingsConfig { fn default() -> Self { + let model = "text-embedding-3-small".to_string(); + let dimension = default_dimension_for_model(&model); Self { enabled: false, provider: "openai".to_string(), openai_api_key: None, - model: "text-embedding-3-small".to_string(), + model, + ollama_base_url: "http://localhost:11434".to_string(), + dimension, } } } +/// Infer the embedding dimension from a well-known model name. +/// +/// Falls back to 1536 (OpenAI text-embedding-3-small default) for unknown models. +fn default_dimension_for_model(model: &str) -> usize { + match model { + "text-embedding-3-small" => 1536, + "text-embedding-3-large" => 3072, + "text-embedding-ada-002" => 1536, + "nomic-embed-text" => 768, + "mxbai-embed-large" => 1024, + "all-minilm" => 384, + _ => 1536, + } +} + impl EmbeddingsConfig { pub(crate) fn resolve(settings: &Settings) -> Result { let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from); @@ -38,6 +61,19 @@ impl EmbeddingsConfig { let model = optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone()); + let ollama_base_url = optional_env("OLLAMA_BASE_URL")? + .or_else(|| settings.ollama_base_url.clone()) + .unwrap_or_else(|| "http://localhost:11434".to_string()); + + let dimension = optional_env("EMBEDDING_DIMENSION")? + .map(|s| s.parse::()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "EMBEDDING_DIMENSION".to_string(), + message: format!("must be a positive integer: {e}"), + })? + .unwrap_or_else(|| default_dimension_for_model(&model)); + let enabled = optional_env("EMBEDDING_ENABLED")? .map(|s| s.parse()) .transpose() @@ -52,6 +88,8 @@ impl EmbeddingsConfig { provider, openai_api_key, model, + ollama_base_url, + dimension, }) } diff --git a/src/config/llm.rs b/src/config/llm.rs index 53b21892..2e315702 100644 --- a/src/config/llm.rs +++ b/src/config/llm.rs @@ -64,6 +64,8 @@ impl std::fmt::Display for LlmBackend { pub struct OpenAiDirectConfig { pub api_key: SecretString, pub model: String, + /// Optional base URL override (e.g. for proxies like VibeProxy). + pub base_url: Option, } /// Configuration for direct Anthropic API access. @@ -71,6 +73,8 @@ pub struct OpenAiDirectConfig { pub struct AnthropicDirectConfig { pub api_key: SecretString, pub model: String, + /// Optional base URL override (e.g. for proxies like VibeProxy). + pub base_url: Option, } /// Configuration for local Ollama. @@ -274,7 +278,12 @@ impl LlmConfig { hint: "Set OPENAI_API_KEY when LLM_BACKEND=openai".to_string(), })?; let model = optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4o".to_string()); - Some(OpenAiDirectConfig { api_key, model }) + let base_url = optional_env("OPENAI_BASE_URL")?; + Some(OpenAiDirectConfig { + api_key, + model, + base_url, + }) } else { None }; @@ -288,7 +297,12 @@ impl LlmConfig { })?; let model = optional_env("ANTHROPIC_MODEL")? .unwrap_or_else(|| "claude-sonnet-4-20250514".to_string()); - Some(AnthropicDirectConfig { api_key, model }) + let base_url = optional_env("ANTHROPIC_BASE_URL")?; + Some(AnthropicDirectConfig { + api_key, + model, + base_url, + }) } else { None }; diff --git a/src/config/sandbox.rs b/src/config/sandbox.rs index 1473507e..57a016fc 100644 --- a/src/config/sandbox.rs +++ b/src/config/sandbox.rs @@ -30,7 +30,7 @@ impl Default for SandboxModeConfig { timeout_secs: 120, memory_limit_mb: 2048, cpu_shares: 1024, - image: "ghcr.io/nearai/sandbox:latest".to_string(), + image: "ironclaw-worker:latest".to_string(), auto_pull_image: true, extra_allowed_domains: Vec::new(), } @@ -57,7 +57,7 @@ impl SandboxModeConfig { memory_limit_mb: parse_optional_env("SANDBOX_MEMORY_LIMIT_MB", 2048)?, cpu_shares: parse_optional_env("SANDBOX_CPU_SHARES", 1024)?, image: optional_env("SANDBOX_IMAGE")? - .unwrap_or_else(|| "ghcr.io/nearai/sandbox:latest".to_string()), + .unwrap_or_else(|| "ironclaw-worker:latest".to_string()), auto_pull_image: optional_env("SANDBOX_AUTO_PULL")? .map(|s| s.parse()) .transpose() diff --git a/src/llm/circuit_breaker.rs b/src/llm/circuit_breaker.rs index 12e46b30..cadf73d6 100644 --- a/src/llm/circuit_breaker.rs +++ b/src/llm/circuit_breaker.rs @@ -123,7 +123,11 @@ impl CircuitBreakerProvider { ); Ok(()) } else { - let remaining = self.config.recovery_timeout - opened_at.elapsed(); + let remaining = self + .config + .recovery_timeout + .checked_sub(opened_at.elapsed()) + .unwrap_or(Duration::ZERO); Err(LlmError::RequestFailed { provider: self.inner.model_name().to_string(), reason: format!( @@ -208,8 +212,16 @@ impl CircuitBreakerProvider { /// Returns `true` for errors that indicate the provider is degraded /// (server errors, rate limits, network failures, auth infrastructure down). /// -/// Client errors (wrong model, bad credentials, context overflow) are NOT -/// transient: they are the caller's problem, not a sign of backend trouble. +/// This answers: "should this error count toward tripping the circuit breaker?" +/// +/// Includes `SessionExpired` because repeated session failures signal backend +/// auth infrastructure trouble. +/// +/// Excludes client errors that are the caller's problem, not backend trouble: +/// `AuthFailed`, `ContextLengthExceeded`, `ModelNotAvailable`, `Json`. +/// +/// See also `retry::is_retryable()` which answers a different question: +/// "could retrying this exact request succeed?" fn is_transient(err: &LlmError) -> bool { matches!( err, @@ -219,7 +231,6 @@ fn is_transient(err: &LlmError) -> bool { | LlmError::SessionExpired { .. } | LlmError::SessionRenewalFailed { .. } | LlmError::Http(_) - | LlmError::Json(_) | LlmError::Io(_) ) } @@ -547,6 +558,9 @@ mod tests { provider: "p".into(), model: "m".into(), })); + assert!(!is_transient(&LlmError::Json( + serde_json::from_str::("bad").unwrap_err() + ))); } // -- Passthrough delegation tests -- diff --git a/src/llm/failover.rs b/src/llm/failover.rs index 57836a3f..ffd6a52e 100644 --- a/src/llm/failover.rs +++ b/src/llm/failover.rs @@ -19,34 +19,11 @@ use rust_decimal::Decimal; use crate::error::LlmError; use crate::llm::provider::{ - CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest, + CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest, ToolCompletionResponse, }; -/// Returns `true` if the error is transient and the request should be retried -/// on the next provider in the failover chain. -/// -/// Retryable: `RequestFailed`, `RateLimited`, `InvalidResponse`, -/// `SessionRenewalFailed`, `ModelNotAvailable`, `Http`, `Io`. -/// -/// `ModelNotAvailable` is retryable because the next provider in the chain may -/// offer a different model, so it's worth trying. -/// -/// Non-retryable errors (`AuthFailed`, `SessionExpired`, `ContextLengthExceeded`) -/// propagate immediately because a different provider won't fix them. -fn is_retryable(err: &LlmError) -> bool { - matches!( - err, - LlmError::RequestFailed { .. } - | LlmError::RateLimited { .. } - | LlmError::InvalidResponse { .. } - | LlmError::SessionRenewalFailed { .. } - // ModelNotAvailable is retryable: the next provider may offer a different model. - | LlmError::ModelNotAvailable { .. } - | LlmError::Http(_) - | LlmError::Io(_) - ) -} +use crate::llm::retry::is_retryable; /// Configuration for per-provider cooldown behavior. /// @@ -376,6 +353,26 @@ impl LlmProvider for FailoverProvider { Ok(all_models) } + async fn model_metadata(&self) -> Result { + self.providers[self.last_used.load(Ordering::Relaxed)] + .model_metadata() + .await + } + + fn seed_response_chain(&self, thread_id: &str, response_id: String) { + self.providers[self.last_used.load(Ordering::Relaxed)] + .seed_response_chain(thread_id, response_id); + } + + fn get_response_chain_id(&self, thread_id: &str) -> Option { + self.providers[self.last_used.load(Ordering::Relaxed)].get_response_chain_id(thread_id) + } + + fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal { + self.providers[self.last_used.load(Ordering::Relaxed)] + .calculate_cost(input_tokens, output_tokens) + } + fn effective_model_name(&self, requested_model: Option<&str>) -> String { if let Some(provider_idx) = self.take_bound_provider_for_current_task() { return self.providers[provider_idx].effective_model_name(requested_model); @@ -1111,10 +1108,6 @@ mod tests { std::io::ErrorKind::ConnectionReset, "reset" )))); - assert!(is_retryable(&LlmError::ModelNotAvailable { - provider: "p".into(), - model: "m".into(), - })); // Non-retryable assert!(!is_retryable(&LlmError::AuthFailed { @@ -1127,6 +1120,10 @@ mod tests { used: 100_000, limit: 50_000, })); + assert!(!is_retryable(&LlmError::ModelNotAvailable { + provider: "p".into(), + model: "m".into(), + })); } // Test: empty providers list returns error (not panic). diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 45b2c85a..55738ab6 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -15,7 +15,7 @@ mod nearai_chat; mod provider; mod reasoning; pub mod response_cache; -mod retry; +pub mod retry; mod rig_adapter; pub mod session; @@ -32,6 +32,7 @@ pub use reasoning::{ ToolSelection, }; pub use response_cache::{CachedProvider, ResponseCacheConfig}; +pub use retry::{RetryConfig, RetryProvider}; pub use rig_adapter::RigAdapter; pub use session::{SessionConfig, SessionManager, create_session_manager}; @@ -76,7 +77,7 @@ pub fn create_llm_provider_with_config( model = %config.model, "Using Responses API (chat-api) with session auth" ); - Ok(Arc::new(NearAiProvider::new(config.clone(), session))) + Ok(Arc::new(NearAiProvider::new(config.clone(), session)?)) } NearAiApiMode::ChatCompletions => { tracing::info!( @@ -99,15 +100,30 @@ fn create_openai_provider(config: &LlmConfig) -> Result, Ll // (Responses API). The Responses API path in rig-core panics when tool results // are sent back because ironclaw doesn't thread `call_id` through its ToolCall // type. The Chat Completions API works correctly with the existing code. - let client: openai::CompletionsClient = openai::Client::new(oai.api_key.expose_secret()) - .map_err(|e| LlmError::RequestFailed { - provider: "openai".to_string(), - reason: format!("Failed to create OpenAI client: {}", e), - })? - .completions_api(); + let client: openai::CompletionsClient = if let Some(ref base_url) = oai.base_url { + tracing::info!( + "Using OpenAI direct API (chat completions, model: {}, base_url: {})", + oai.model, + base_url, + ); + openai::Client::builder() + .base_url(base_url) + .api_key(oai.api_key.expose_secret()) + .build() + } else { + tracing::info!( + "Using OpenAI direct API (chat completions, model: {}, base_url: default)", + oai.model, + ); + openai::Client::new(oai.api_key.expose_secret()) + } + .map_err(|e| LlmError::RequestFailed { + provider: "openai".to_string(), + reason: format!("Failed to create OpenAI client: {}", e), + })? + .completions_api(); let model = client.completion_model(&oai.model); - tracing::info!("Using OpenAI direct API (model: {})", oai.model); Ok(Arc::new(RigAdapter::new(model, &oai.model))) } @@ -121,16 +137,25 @@ fn create_anthropic_provider(config: &LlmConfig) -> Result, use rig::providers::anthropic; - let client: anthropic::Client = - anthropic::Client::new(anth.api_key.expose_secret()).map_err(|e| { - LlmError::RequestFailed { - provider: "anthropic".to_string(), - reason: format!("Failed to create Anthropic client: {}", e), - } - })?; + let client: anthropic::Client = if let Some(ref base_url) = anth.base_url { + anthropic::Client::builder() + .api_key(anth.api_key.expose_secret()) + .base_url(base_url) + .build() + } else { + anthropic::Client::new(anth.api_key.expose_secret()) + } + .map_err(|e| LlmError::RequestFailed { + provider: "anthropic".to_string(), + reason: format!("Failed to create Anthropic client: {}", e), + })?; let model = client.completion_model(&anth.model); - tracing::info!("Using Anthropic direct API (model: {})", anth.model); + tracing::info!( + "Using Anthropic direct API (model: {}, base_url: {})", + anth.model, + anth.base_url.as_deref().unwrap_or("default"), + ); Ok(Arc::new(RigAdapter::new(model, &anth.model))) } @@ -199,26 +224,25 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)))), + NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)?))), NearAiApiMode::ChatCompletions => { Ok(Some(Arc::new(NearAiChatProvider::new(cheap_config)?))) } diff --git a/src/llm/nearai.rs b/src/llm/nearai.rs index de58a8be..27d90622 100644 --- a/src/llm/nearai.rs +++ b/src/llm/nearai.rs @@ -19,7 +19,6 @@ use crate::llm::provider::{ ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, }; -use crate::llm::retry::{is_retryable_status, retry_backoff_delay}; use crate::llm::session::SessionManager; /// Information about an available model from NEAR AI API. @@ -54,28 +53,34 @@ pub struct NearAiProvider { impl NearAiProvider { /// Create a new NEAR AI provider with a session manager. - pub fn new(config: NearAiConfig, session: Arc) -> Self { + pub fn new(config: NearAiConfig, session: Arc) -> Result { let client = Client::builder() .timeout(std::time::Duration::from_secs(120)) .build() - .unwrap_or_else(|_| Client::new()); + .map_err(|e| LlmError::RequestFailed { + provider: "nearai".to_string(), + reason: format!("Failed to build HTTP client: {}", e), + })?; let active_model = std::sync::RwLock::new(config.model.clone()); - Self { + Ok(Self { client, config, session, active_model, response_chains: std::sync::RwLock::new(HashMap::new()), - } + }) } /// Seed a response chain for a thread (e.g. when restoring from DB). pub fn seed_response_id(&self, thread_id: &str, response_id: String) { - let mut chains = self - .response_chains - .write() - .expect("response_chains lock poisoned"); + let mut chains = match self.response_chains.write() { + Ok(guard) => guard, + Err(poisoned) => { + tracing::warn!("response_chains lock poisoned in seed; recovering"); + poisoned.into_inner() + } + }; chains.insert( thread_id.to_string(), ChainState { @@ -87,19 +92,25 @@ impl NearAiProvider { /// Get the last response ID for a thread (for persistence). pub fn get_response_id(&self, thread_id: &str) -> Option { - let chains = self - .response_chains - .read() - .expect("response_chains lock poisoned"); + let chains = match self.response_chains.read() { + Ok(guard) => guard, + Err(poisoned) => { + tracing::warn!("response_chains lock poisoned in get; recovering"); + poisoned.into_inner() + } + }; chains.get(thread_id).map(|c| c.response_id.clone()) } /// Store a response chain state after a successful call. fn store_chain(&self, thread_id: &str, response_id: String, input_count: usize) { - let mut chains = self - .response_chains - .write() - .expect("response_chains lock poisoned"); + let mut chains = match self.response_chains.write() { + Ok(guard) => guard, + Err(poisoned) => { + tracing::warn!("response_chains lock poisoned in store; recovering"); + poisoned.into_inner() + } + }; chains.insert( thread_id.to_string(), ChainState { @@ -111,10 +122,13 @@ impl NearAiProvider { /// Clear the chain for a thread (on error / fallback). fn clear_chain(&self, thread_id: &str) { - let mut chains = self - .response_chains - .write() - .expect("response_chains lock poisoned"); + let mut chains = match self.response_chains.write() { + Ok(guard) => guard, + Err(poisoned) => { + tracing::warn!("response_chains lock poisoned in clear; recovering"); + poisoned.into_inner() + } + }; chains.remove(thread_id); } @@ -160,7 +174,10 @@ impl NearAiProvider { })?; let status = response.status(); - let response_text = response.text().await.unwrap_or_default(); + let response_text = response.text().await.map_err(|e| LlmError::RequestFailed { + provider: "nearai".to_string(), + reason: format!("Failed to read response body: {}", e), + })?; if !status.is_success() { if status.as_u16() == 401 { @@ -283,139 +300,95 @@ impl NearAiProvider { } } - /// Inner request implementation with retry logic for transient errors. + /// Inner request implementation (single attempt). /// - /// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff. - /// Does not retry on client errors (400, 401, 403, 404) or parse errors. + /// Does not retry internally — retries are handled by the external + /// `RetryProvider` wrapper in the composition chain. async fn send_request_inner Deserialize<'de>>( &self, path: &str, body: &T, ) -> Result { let url = self.api_url(path); - let max_retries = self.config.max_retries; + let token = self.session.get_token().await?; - for attempt in 0..=max_retries { - let token = self.session.get_token().await?; + tracing::debug!("Sending request to NEAR AI: {}", url); + tracing::debug!("Request body: {:?}", body); - tracing::debug!( - "Sending request to NEAR AI: {} (attempt {})", - url, - attempt + 1 - ); - tracing::debug!("Request body: {:?}", body); + let response = self + .client + .post(&url) + .header("Authorization", format!("Bearer {}", token.expose_secret())) + .header("Content-Type", "application/json") + .json(body) + .send() + .await + .map_err(|e| { + tracing::error!("NEAR AI request failed: {}", e); + LlmError::Http(e) + })?; - let response = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", token.expose_secret())) - .header("Content-Type", "application/json") - .json(body) - .send() - .await; + let status = response.status(); + let response_text = response.text().await.map_err(|e| LlmError::RequestFailed { + provider: "nearai".to_string(), + reason: format!("Failed to read response body: {}", e), + })?; - let response = match response { - Ok(r) => r, - Err(e) => { - tracing::error!("NEAR AI request failed: {}", e); - // Network errors (timeout, connection refused) are transient - if attempt < max_retries { - let delay = retry_backoff_delay(attempt); - tracing::warn!( - "NEAR AI request error (attempt {}/{}), retrying in {:?}: {}", - attempt + 1, - max_retries + 1, - delay, - e, - ); - tokio::time::sleep(delay).await; - continue; - } - return Err(e.into()); - } - }; + tracing::debug!("NEAR AI response status: {}", status); + tracing::debug!("NEAR AI response body: {}", response_text); - let status = response.status(); - let response_text = response.text().await.unwrap_or_default(); + if !status.is_success() { + let status_code = status.as_u16(); - tracing::debug!("NEAR AI response status: {}", status); - tracing::debug!("NEAR AI response body: {}", response_text); + // Check for session expiration (401 with specific message patterns) + if status_code == 401 { + let lower = response_text.to_lowercase(); + let is_session_expired = lower.contains("session") + && (lower.contains("expired") || lower.contains("invalid")); - if !status.is_success() { - let status_code = status.as_u16(); - - // Check for session expiration (401 with specific message patterns) - if status_code == 401 { - let lower = response_text.to_lowercase(); - let is_session_expired = lower.contains("session") - && (lower.contains("expired") || lower.contains("invalid")); - - if is_session_expired { - return Err(LlmError::SessionExpired { - provider: "nearai".to_string(), - }); - } - - // Generic 401 -- not retryable - return Err(LlmError::AuthFailed { + if is_session_expired { + return Err(LlmError::SessionExpired { provider: "nearai".to_string(), }); } - // Check if this is a transient error worth retrying - if is_retryable_status(status_code) && attempt < max_retries { - let delay = retry_backoff_delay(attempt); - tracing::warn!( - "NEAR AI returned HTTP {} (attempt {}/{}), retrying in {:?}", - status_code, - attempt + 1, - max_retries + 1, - delay, - ); - tokio::time::sleep(delay).await; - continue; - } - - // Non-retryable error or exhausted retries - if let Ok(error) = serde_json::from_str::(&response_text) { - if status_code == 429 { - return Err(LlmError::RateLimited { - provider: "nearai".to_string(), - retry_after: None, - }); - } - return Err(LlmError::RequestFailed { - provider: "nearai".to_string(), - reason: error.error, - }); - } - - return Err(LlmError::RequestFailed { + return Err(LlmError::AuthFailed { provider: "nearai".to_string(), - reason: format!("HTTP {}: {}", status, response_text), }); } - // Success -- parse the response - return match serde_json::from_str::(&response_text) { - Ok(parsed) => Ok(parsed), - Err(e) => { - tracing::debug!("Response is not expected JSON format: {}", e); - tracing::debug!("Will try alternative parsing in caller"); - Err(LlmError::InvalidResponse { - provider: "nearai".to_string(), - reason: format!("Parse error: {}. Raw: {}", e, response_text), - }) - } - }; + if status_code == 429 { + return Err(LlmError::RateLimited { + provider: "nearai".to_string(), + retry_after: None, + }); + } + + if let Ok(error) = serde_json::from_str::(&response_text) { + return Err(LlmError::RequestFailed { + provider: "nearai".to_string(), + reason: error.error, + }); + } + + return Err(LlmError::RequestFailed { + provider: "nearai".to_string(), + reason: format!("HTTP {}: {}", status, response_text), + }); } - // This is unreachable because the loop always returns, but the compiler - // cannot prove that. Return a generic error as a safety net. - Err(LlmError::RequestFailed { - provider: "nearai".to_string(), - reason: "retry loop exited unexpectedly".to_string(), - }) + // Success -- parse the response + match serde_json::from_str::(&response_text) { + Ok(parsed) => Ok(parsed), + Err(e) => { + tracing::debug!("Response is not expected JSON format: {}", e); + tracing::debug!("Will try alternative parsing in caller"); + Err(LlmError::InvalidResponse { + provider: "nearai".to_string(), + reason: format!("Parse error: {}. Raw: {}", e, response_text), + }) + } + } } } @@ -464,7 +437,9 @@ impl LlmProvider for NearAiProvider { async fn complete(&self, req: CompletionRequest) -> Result { let model = req.model.unwrap_or_else(|| self.active_model_name()); let thread_id = req.metadata.get("thread_id").cloned(); - let (instructions, input) = split_messages(req.messages, false); + let mut messages = req.messages; + crate::llm::provider::sanitize_tool_messages(&mut messages); + let (instructions, input) = split_messages(messages, false); let request = NearAiRequest { model, @@ -582,13 +557,20 @@ impl LlmProvider for NearAiProvider { ) -> Result { let model = req.model.unwrap_or_else(|| self.active_model_name()); let thread_id = req.metadata.get("thread_id").cloned(); + let mut messages = req.messages; + crate::llm::provider::sanitize_tool_messages(&mut messages); // Look up chaining state for this thread let chain_state = thread_id.as_ref().and_then(|tid| { - let chains = self - .response_chains - .read() - .expect("response_chains lock poisoned"); + let chains = match self.response_chains.read() { + Ok(guard) => guard, + Err(poisoned) => { + tracing::warn!( + "response_chains lock poisoned in complete_with_tools; recovering" + ); + poisoned.into_inner() + } + }; chains .get(tid) .map(|c| (c.response_id.clone(), c.input_count)) @@ -601,7 +583,7 @@ impl LlmProvider for NearAiProvider { // When chaining, only send new messages (the delta since last call). // Tool results are converted to function_call_output items. - let (instructions, all_input) = split_messages(req.messages, chaining); + let (instructions, all_input) = split_messages(messages, chaining); let input = if chaining && all_input.len() > prev_input_count { all_input[prev_input_count..].to_vec() } else { @@ -806,18 +788,25 @@ impl LlmProvider for NearAiProvider { } fn active_model_name(&self) -> String { - self.active_model - .read() - .expect("active_model lock poisoned") - .clone() + match self.active_model.read() { + Ok(guard) => guard.clone(), + Err(poisoned) => { + tracing::warn!("active_model lock poisoned while reading; continuing"); + poisoned.into_inner().clone() + } + } } fn set_model(&self, model: &str) -> Result<(), LlmError> { - let mut guard = self - .active_model - .write() - .expect("active_model lock poisoned"); - *guard = model.to_string(); + match self.active_model.write() { + Ok(mut guard) => { + *guard = model.to_string(); + } + Err(poisoned) => { + tracing::warn!("active_model lock poisoned while writing; continuing"); + *poisoned.into_inner() = model.to_string(); + } + } Ok(()) } @@ -952,7 +941,6 @@ struct NearAiTool { /// Primary response format (output array style) #[derive(Debug, Deserialize)] struct NearAiResponse { - #[allow(dead_code)] id: String, output: Vec, usage: NearAiUsage, diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index 2f97af36..68472ed8 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -16,18 +16,29 @@ use crate::llm::provider::{ ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, }; -use crate::llm::retry::{is_retryable_status, retry_backoff_delay}; /// NEAR AI Chat Completions API provider. pub struct NearAiChatProvider { client: Client, config: NearAiConfig, active_model: std::sync::RwLock, + flatten_tool_messages: bool, } impl NearAiChatProvider { /// Create a new NEAR AI chat completions provider with API key auth. + /// + /// By default this enables tool-message flattening for compatibility with + /// providers that reject `role: "tool"` messages (e.g. NEAR cloud-api). pub fn new(config: NearAiConfig) -> Result { + Self::new_with_flatten(config, true) + } + + /// Create a chat completions provider with configurable tool-message flattening. + pub fn new_with_flatten( + config: NearAiConfig, + flatten_tool_messages: bool, + ) -> Result { if config.api_key.is_none() { return Err(LlmError::AuthFailed { provider: "nearai_chat".to_string(), @@ -37,22 +48,29 @@ impl NearAiChatProvider { let client = Client::builder() .timeout(std::time::Duration::from_secs(120)) .build() - .unwrap_or_else(|_| Client::new()); + .map_err(|e| LlmError::RequestFailed { + provider: "nearai_chat".to_string(), + reason: format!("Failed to build HTTP client: {}", e), + })?; let active_model = std::sync::RwLock::new(config.model.clone()); Ok(Self { client, config, active_model, + flatten_tool_messages, }) } fn api_url(&self, path: &str) -> String { - format!( - "{}/v1/{}", - self.config.base_url, - path.trim_start_matches('/') - ) + let base = self.config.base_url.trim_end_matches('/'); + let path = path.trim_start_matches('/'); + + if base.ends_with("/v1") { + format!("{}/{}", base, path) + } else { + format!("{}/v1/{}", base, path) + } } fn api_key(&self) -> String { @@ -63,116 +81,75 @@ impl NearAiChatProvider { .unwrap_or_default() } - /// Send a request to the chat completions API with retry on transient errors. + /// Send a single request to the chat completions API. /// - /// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff. - /// Does not retry on client errors (400, 401, 403, 404) or parse errors. + /// Does not retry internally — retries are handled by the external + /// `RetryProvider` wrapper in the composition chain. async fn send_request Deserialize<'de>>( &self, body: &T, ) -> Result { let url = self.api_url("chat/completions"); - let max_retries = self.config.max_retries; - for attempt in 0..=max_retries { - tracing::debug!( - "Sending request to NEAR AI Chat: {} (attempt {})", - url, - attempt + 1, - ); + tracing::debug!("Sending request to NEAR AI Chat: {}", url); - if tracing::enabled!(tracing::Level::DEBUG) - && let Ok(json) = serde_json::to_string(body) - { - tracing::debug!("NEAR AI Chat request body: {}", json); - } + if tracing::enabled!(tracing::Level::DEBUG) + && let Ok(json) = serde_json::to_string(body) + { + tracing::debug!("NEAR AI Chat request body: {}", json); + } - let response = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key())) - .header("Content-Type", "application/json") - .json(body) - .send() - .await; + let response = self + .client + .post(&url) + .header("Authorization", format!("Bearer {}", self.api_key())) + .header("Content-Type", "application/json") + .json(body) + .send() + .await + .map_err(|e| LlmError::RequestFailed { + provider: "nearai_chat".to_string(), + reason: e.to_string(), + })?; - let response = match response { - Ok(r) => r, - Err(e) => { - tracing::error!("NEAR AI Chat request failed: {}", e); - if attempt < max_retries { - let delay = retry_backoff_delay(attempt); - tracing::warn!( - "NEAR AI Chat request error (attempt {}/{}), retrying in {:?}: {}", - attempt + 1, - max_retries + 1, - delay, - e, - ); - tokio::time::sleep(delay).await; - continue; - } - return Err(LlmError::RequestFailed { - provider: "nearai_chat".to_string(), - reason: e.to_string(), - }); - } - }; + let status = response.status(); + let response_text = response.text().await.map_err(|e| LlmError::RequestFailed { + provider: "nearai_chat".to_string(), + reason: format!("Failed to read response body: {}", e), + })?; - let status = response.status(); - let response_text = response.text().await.unwrap_or_default(); + tracing::debug!("NEAR AI Chat response status: {}", status); + tracing::debug!("NEAR AI Chat response body: {}", response_text); - tracing::debug!("NEAR AI Chat response status: {}", status); - tracing::debug!("NEAR AI Chat response body: {}", response_text); + if !status.is_success() { + let status_code = status.as_u16(); - if !status.is_success() { - let status_code = status.as_u16(); - - // Auth errors are not retryable - if status_code == 401 { - return Err(LlmError::AuthFailed { - provider: "nearai_chat".to_string(), - }); - } - - // Transient errors: retry with backoff - if is_retryable_status(status_code) && attempt < max_retries { - let delay = retry_backoff_delay(attempt); - tracing::warn!( - "NEAR AI Chat returned HTTP {} (attempt {}/{}), retrying in {:?}", - status_code, - attempt + 1, - max_retries + 1, - delay, - ); - tokio::time::sleep(delay).await; - continue; - } - - // Non-retryable or exhausted retries - if status_code == 429 { - return Err(LlmError::RateLimited { - provider: "nearai_chat".to_string(), - retry_after: None, - }); - } - return Err(LlmError::RequestFailed { + if status_code == 401 { + return Err(LlmError::AuthFailed { provider: "nearai_chat".to_string(), - reason: format!("HTTP {}: {}", status, response_text), }); } - // Success — parse the response - return serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse { + if status_code == 429 { + return Err(LlmError::RateLimited { + provider: "nearai_chat".to_string(), + retry_after: None, + }); + } + + let truncated = crate::agent::truncate_for_preview(&response_text, 512); + return Err(LlmError::RequestFailed { provider: "nearai_chat".to_string(), - reason: format!("JSON parse error: {}. Raw: {}", e, response_text), + reason: format!("HTTP {}: {}", status, truncated), }); } - // Safety net: unreachable because the loop always returns - Err(LlmError::RequestFailed { - provider: "nearai_chat".to_string(), - reason: "retry loop exited unexpectedly".to_string(), + serde_json::from_str(&response_text).map_err(|e| { + let truncated = crate::agent::truncate_for_preview(&response_text, 512); + LlmError::InvalidResponse { + provider: "nearai_chat".to_string(), + reason: format!("JSON parse error: {}. Raw: {}", e, truncated), + } }) } @@ -192,12 +169,16 @@ impl NearAiChatProvider { })?; let status = response.status(); - let response_text = response.text().await.unwrap_or_default(); + let response_text = response.text().await.map_err(|e| LlmError::RequestFailed { + provider: "nearai_chat".to_string(), + reason: format!("Failed to read response body: {}", e), + })?; if !status.is_success() { + let truncated = crate::agent::truncate_for_preview(&response_text, 512); return Err(LlmError::RequestFailed { provider: "nearai_chat".to_string(), - reason: format!("HTTP {}: {}", status, response_text), + reason: format!("HTTP {}: {}", status, truncated), }); } @@ -228,8 +209,10 @@ struct ApiModelEntry { impl LlmProvider for NearAiChatProvider { async fn complete(&self, req: CompletionRequest) -> Result { let model = req.model.unwrap_or_else(|| self.active_model_name()); + let mut raw_messages = req.messages; + crate::llm::provider::sanitize_tool_messages(&mut raw_messages); let messages: Vec = - req.messages.into_iter().map(|m| m.into()).collect(); + raw_messages.into_iter().map(|m| m.into()).collect(); let request = ChatCompletionRequest { model, @@ -261,11 +244,13 @@ impl LlmProvider for NearAiChatProvider { _ => FinishReason::Unknown, }; + let (input_tokens, output_tokens) = parse_usage(response.usage.as_ref()); + Ok(CompletionResponse { content, finish_reason, - input_tokens: response.usage.prompt_tokens, - output_tokens: response.usage.completion_tokens, + input_tokens, + output_tokens, response_id: None, }) } @@ -275,14 +260,18 @@ impl LlmProvider for NearAiChatProvider { req: ToolCompletionRequest, ) -> Result { let model = req.model.unwrap_or_else(|| self.active_model_name()); + let mut raw_messages = req.messages; + crate::llm::provider::sanitize_tool_messages(&mut raw_messages); let messages: Vec = - req.messages.into_iter().map(|m| m.into()).collect(); + raw_messages.into_iter().map(|m| m.into()).collect(); - // NEAR AI cloud-api does not support multi-turn tool calling (rejects - // any request containing role:"tool" messages with HTTP 400). Rewrite - // tool-call / tool-result pairs into plain text so the conversation - // history is preserved without using unsupported message roles. - let messages = flatten_tool_messages(messages); + // Some OpenAI-compatible providers reject `role:"tool"` messages. + // When enabled, rewrite tool-call / tool-result pairs into plain text. + let messages = if self.flatten_tool_messages { + flatten_tool_messages(messages) + } else { + messages + }; let tools: Vec = req .tools @@ -349,12 +338,14 @@ impl LlmProvider for NearAiChatProvider { } }; + let (input_tokens, output_tokens) = parse_usage(response.usage.as_ref()); + Ok(ToolCompletionResponse { content, tool_calls, finish_reason, - input_tokens: response.usage.prompt_tokens, - output_tokens: response.usage.completion_tokens, + input_tokens, + output_tokens, response_id: None, }) } @@ -384,18 +375,25 @@ impl LlmProvider for NearAiChatProvider { } fn active_model_name(&self) -> String { - self.active_model - .read() - .expect("active_model lock poisoned") - .clone() + match self.active_model.read() { + Ok(guard) => guard.clone(), + Err(poisoned) => { + tracing::warn!("active_model lock poisoned while reading; continuing"); + poisoned.into_inner().clone() + } + } } fn set_model(&self, model: &str) -> Result<(), crate::error::LlmError> { - let mut guard = self - .active_model - .write() - .expect("active_model lock poisoned"); - *guard = model.to_string(); + match self.active_model.write() { + Ok(mut guard) => { + *guard = model.to_string(); + } + Err(poisoned) => { + tracing::warn!("active_model lock poisoned while writing; continuing"); + *poisoned.into_inner() = model.to_string(); + } + } Ok(()) } } @@ -545,9 +543,11 @@ struct ChatCompletionFunction { #[derive(Debug, Deserialize)] struct ChatCompletionResponse { #[allow(dead_code)] - id: String, + #[serde(default)] + id: Option, choices: Vec, - usage: ChatCompletionUsage, + #[serde(default)] + usage: Option, } #[derive(Debug, Deserialize)] @@ -579,18 +579,90 @@ struct ChatCompletionToolCallFunction { arguments: String, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Default)] struct ChatCompletionUsage { - prompt_tokens: u32, - completion_tokens: u32, - #[allow(dead_code)] - total_tokens: u32, + #[serde(default)] + prompt_tokens: Option, + #[serde(default)] + completion_tokens: Option, + #[serde(default)] + total_tokens: Option, +} + +fn saturate_u32(val: u64) -> u32 { + val.min(u32::MAX as u64) as u32 +} + +fn parse_usage(usage: Option<&ChatCompletionUsage>) -> (u32, u32) { + let Some(u) = usage else { + return (0, 0); + }; + let input = u.prompt_tokens.map(saturate_u32).unwrap_or(0); + let output = u.completion_tokens.map(saturate_u32).unwrap_or_else(|| { + // Fall back to total - prompt if completion is missing. + match (u.total_tokens, u.prompt_tokens) { + (Some(total), Some(prompt)) => saturate_u32(total.saturating_sub(prompt)), + (Some(total), None) => saturate_u32(total), + _ => 0, + } + }); + (input, output) } #[cfg(test)] mod tests { use super::*; + fn test_nearai_config(base_url: &str) -> NearAiConfig { + NearAiConfig { + model: "test-model".to_string(), + base_url: base_url.to_string(), + auth_base_url: "https://private.near.ai".to_string(), + session_path: std::path::PathBuf::from("/tmp/session.json"), + api_mode: crate::config::NearAiApiMode::ChatCompletions, + api_key: Some(secrecy::SecretString::from("test-key".to_string())), + cheap_model: None, + fallback_model: None, + max_retries: 0, + circuit_breaker_threshold: None, + circuit_breaker_recovery_secs: 30, + response_cache_enabled: false, + response_cache_ttl_secs: 3600, + response_cache_max_entries: 1000, + failover_cooldown_secs: 300, + failover_cooldown_threshold: 3, + } + } + + #[test] + fn test_api_url_with_base_without_v1() { + let mut cfg = test_nearai_config("http://127.0.0.1:8318"); + + let provider = NearAiChatProvider::new(cfg.clone()).expect("provider"); + assert_eq!( + provider.api_url("chat/completions"), + "http://127.0.0.1:8318/v1/chat/completions" + ); + + cfg.base_url = "http://127.0.0.1:8318/".to_string(); + let provider = NearAiChatProvider::new(cfg).expect("provider"); + assert_eq!( + provider.api_url("/chat/completions"), + "http://127.0.0.1:8318/v1/chat/completions" + ); + } + + #[test] + fn test_api_url_with_base_already_v1() { + let cfg = test_nearai_config("http://127.0.0.1:8318/v1"); + + let provider = NearAiChatProvider::new(cfg).expect("provider"); + assert_eq!( + provider.api_url("chat/completions"), + "http://127.0.0.1:8318/v1/chat/completions" + ); + } + #[test] fn test_message_conversion() { let msg = ChatMessage::user("Hello"); diff --git a/src/llm/provider.rs b/src/llm/provider.rs index 1c4e8510..bf0b8981 100644 --- a/src/llm/provider.rs +++ b/src/llm/provider.rs @@ -347,3 +347,124 @@ pub trait LlmProvider: Send + Sync { input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens) } } + +/// Sanitize a message list to ensure tool_use / tool_result integrity. +/// +/// LLM APIs (especially Anthropic) require every tool_result to reference a +/// tool_call_id that exists in an immediately preceding assistant message's +/// tool_calls. Orphaned tool_results cause HTTP 400 errors. +/// +/// This function: +/// 1. Tracks all tool_call_ids emitted by assistant messages. +/// 2. Rewrites orphaned tool_result messages (whose tool_call_id has no +/// matching assistant tool_call) as user messages so the content is +/// preserved without violating the protocol. +/// +/// Call this before sending messages to any LLM provider. +pub fn sanitize_tool_messages(messages: &mut [ChatMessage]) { + use std::collections::HashSet; + + // Collect all tool_call_ids from assistant messages with tool_calls. + let mut known_ids: HashSet = HashSet::new(); + for msg in messages.iter() { + if msg.role == Role::Assistant + && let Some(ref calls) = msg.tool_calls + { + for tc in calls { + known_ids.insert(tc.id.clone()); + } + } + } + + // Rewrite orphaned tool_result messages as user messages. + for msg in messages.iter_mut() { + if msg.role != Role::Tool { + continue; + } + let is_orphaned = match &msg.tool_call_id { + Some(id) => !known_ids.contains(id), + None => true, + }; + if is_orphaned { + let tool_name = msg.name.as_deref().unwrap_or("unknown"); + tracing::debug!( + tool_call_id = ?msg.tool_call_id, + tool_name, + "Rewriting orphaned tool_result as user message", + ); + msg.role = Role::User; + msg.content = format!("[Tool `{}` returned: {}]", tool_name, msg.content); + msg.tool_call_id = None; + msg.name = None; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sanitize_preserves_valid_pairs() { + let tc = ToolCall { + id: "call_1".to_string(), + name: "echo".to_string(), + arguments: serde_json::json!({}), + }; + let mut messages = vec![ + ChatMessage::user("hello"), + ChatMessage::assistant_with_tool_calls(None, vec![tc]), + ChatMessage::tool_result("call_1", "echo", "result"), + ]; + sanitize_tool_messages(&mut messages); + assert_eq!(messages[2].role, Role::Tool); + assert_eq!(messages[2].tool_call_id, Some("call_1".to_string())); + } + + #[test] + fn test_sanitize_rewrites_orphaned_tool_result() { + let mut messages = vec![ + ChatMessage::user("hello"), + ChatMessage::assistant("I'll use a tool"), + ChatMessage::tool_result("call_missing", "search", "some result"), + ]; + sanitize_tool_messages(&mut messages); + assert_eq!(messages[2].role, Role::User); + assert!(messages[2].content.contains("[Tool `search` returned:")); + assert!(messages[2].tool_call_id.is_none()); + assert!(messages[2].name.is_none()); + } + + #[test] + fn test_sanitize_handles_no_tool_messages() { + let mut messages = vec![ + ChatMessage::system("prompt"), + ChatMessage::user("hello"), + ChatMessage::assistant("hi"), + ]; + let original_len = messages.len(); + sanitize_tool_messages(&mut messages); + assert_eq!(messages.len(), original_len); + } + + #[test] + fn test_sanitize_multiple_orphaned() { + let tc = ToolCall { + id: "call_1".to_string(), + name: "echo".to_string(), + arguments: serde_json::json!({}), + }; + let mut messages = vec![ + ChatMessage::user("test"), + ChatMessage::assistant_with_tool_calls(None, vec![tc]), + ChatMessage::tool_result("call_1", "echo", "ok"), + // These are orphaned (call_2 and call_3 have no matching assistant message) + ChatMessage::tool_result("call_2", "search", "orphan 1"), + ChatMessage::tool_result("call_3", "http", "orphan 2"), + ]; + sanitize_tool_messages(&mut messages); + assert_eq!(messages[2].role, Role::Tool); // call_1 is valid + assert_eq!(messages[3].role, Role::User); // call_2 orphaned + assert_eq!(messages[4].role, Role::User); // call_3 orphaned + } +} diff --git a/src/llm/retry.rs b/src/llm/retry.rs index 2dced0b0..75ad5da6 100644 --- a/src/llm/retry.rs +++ b/src/llm/retry.rs @@ -1,15 +1,50 @@ -//! Shared retry helpers for LLM providers. +//! Shared retry helpers and composable `RetryProvider` decorator for LLM providers. //! -//! Provides exponential backoff with jitter and retryable status classification -//! used by both `NearAiProvider` and `NearAiChatProvider`. +//! Provides: +//! - `is_retryable()` — `LlmError`-level retryability classification (shared with `failover.rs`) +//! - `retry_backoff_delay()` — exponential backoff with jitter +//! - `RetryProvider` — decorator that wraps any `LlmProvider` with automatic retries +use std::sync::Arc; use std::time::Duration; +use async_trait::async_trait; use rand::Rng; +use rust_decimal::Decimal; -/// Returns `true` if the HTTP status code is transient and worth retrying. -pub(crate) fn is_retryable_status(status: u16) -> bool { - matches!(status, 429 | 500 | 502 | 503 | 504) +use crate::error::LlmError; +use crate::llm::provider::{ + CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest, + ToolCompletionResponse, +}; + +/// Returns `true` if the `LlmError` is transient and the request should be retried. +/// +/// Used by `RetryProvider` (retry the same provider) and `FailoverProvider` +/// (try the next provider). The question is: "could this exact same request +/// succeed if we try again?" +/// +/// Retryable: `RequestFailed`, `RateLimited`, `InvalidResponse`, +/// `SessionRenewalFailed`, `Http`, `Io`. +/// +/// Non-retryable: `AuthFailed`, `SessionExpired`, `ContextLengthExceeded`, +/// `ModelNotAvailable`, `Json`. +/// - `SessionExpired` — handled by session renewal layer, not by retry +/// - `ModelNotAvailable` — the model won't appear between attempts +/// - `Json` — a serde parse bug, not a transient failure +/// +/// See also `circuit_breaker::is_transient()` which answers a different +/// question: "does this error indicate the backend is degraded?" +pub(crate) fn is_retryable(err: &LlmError) -> bool { + matches!( + err, + LlmError::RequestFailed { .. } + | LlmError::RateLimited { .. } + | LlmError::InvalidResponse { .. } + | LlmError::SessionRenewalFailed { .. } + | LlmError::Http(_) + | LlmError::Io(_) + ) } /// Calculate exponential backoff delay with random jitter. @@ -31,31 +66,183 @@ pub(crate) fn retry_backoff_delay(attempt: u32) -> Duration { Duration::from_millis(delay_ms) } +/// Configuration for the retry decorator. +#[derive(Debug, Clone)] +pub struct RetryConfig { + /// Maximum number of retry attempts (not counting the initial attempt). + /// Default: 3. + pub max_retries: u32, +} + +impl Default for RetryConfig { + fn default() -> Self { + Self { max_retries: 3 } + } +} + +/// Composable decorator that wraps any `LlmProvider` with automatic retries. +/// +/// On transient errors, sleeps using exponential backoff and retries. +/// On non-transient errors (`AuthFailed`, `ContextLengthExceeded`, `SessionExpired`), +/// returns immediately. +/// +/// Special handling for `RateLimited { retry_after }`: uses the provider-suggested +/// duration if available, otherwise falls back to standard backoff. +pub struct RetryProvider { + inner: Arc, + config: RetryConfig, +} + +impl RetryProvider { + pub fn new(inner: Arc, config: RetryConfig) -> Self { + Self { inner, config } + } +} + +#[async_trait] +impl LlmProvider for RetryProvider { + fn model_name(&self) -> &str { + self.inner.model_name() + } + + fn cost_per_token(&self) -> (Decimal, Decimal) { + self.inner.cost_per_token() + } + + async fn complete(&self, request: CompletionRequest) -> Result { + let mut last_error: Option = None; + + for attempt in 0..=self.config.max_retries { + let req = request.clone(); + match self.inner.complete(req).await { + Ok(resp) => return Ok(resp), + Err(err) => { + if !is_retryable(&err) || attempt == self.config.max_retries { + return Err(err); + } + + let delay = match &err { + LlmError::RateLimited { + retry_after: Some(duration), + .. + } => *duration, + _ => retry_backoff_delay(attempt), + }; + + tracing::warn!( + provider = %self.inner.model_name(), + attempt = attempt + 1, + max_retries = self.config.max_retries, + delay_ms = delay.as_millis() as u64, + error = %err, + "Retrying after transient error" + ); + + last_error = Some(err); + tokio::time::sleep(delay).await; + } + } + } + + Err(last_error.unwrap_or_else(|| LlmError::RequestFailed { + provider: self.inner.model_name().to_string(), + reason: "retry loop exited unexpectedly".to_string(), + })) + } + + async fn complete_with_tools( + &self, + request: ToolCompletionRequest, + ) -> Result { + let mut last_error: Option = None; + + for attempt in 0..=self.config.max_retries { + let req = request.clone(); + match self.inner.complete_with_tools(req).await { + Ok(resp) => return Ok(resp), + Err(err) => { + if !is_retryable(&err) || attempt == self.config.max_retries { + return Err(err); + } + + let delay = match &err { + LlmError::RateLimited { + retry_after: Some(duration), + .. + } => *duration, + _ => retry_backoff_delay(attempt), + }; + + tracing::warn!( + provider = %self.inner.model_name(), + attempt = attempt + 1, + max_retries = self.config.max_retries, + delay_ms = delay.as_millis() as u64, + error = %err, + "Retrying after transient error (tools)" + ); + + last_error = Some(err); + tokio::time::sleep(delay).await; + } + } + } + + Err(last_error.unwrap_or_else(|| LlmError::RequestFailed { + provider: self.inner.model_name().to_string(), + reason: "retry loop exited unexpectedly".to_string(), + })) + } + + async fn list_models(&self) -> Result, LlmError> { + self.inner.list_models().await + } + + async fn model_metadata(&self) -> Result { + self.inner.model_metadata().await + } + + fn active_model_name(&self) -> String { + self.inner.active_model_name() + } + + fn set_model(&self, model: &str) -> Result<(), LlmError> { + self.inner.set_model(model) + } + + fn seed_response_chain(&self, thread_id: &str, response_id: String) { + self.inner.seed_response_chain(thread_id, response_id) + } + + fn get_response_chain_id(&self, thread_id: &str) -> Option { + self.inner.get_response_chain_id(thread_id) + } + + fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal { + self.inner.calculate_cost(input_tokens, output_tokens) + } +} + #[cfg(test)] mod tests { use super::*; - #[test] - fn test_is_retryable_status() { - // Transient errors should be retryable - assert!(is_retryable_status(429)); - assert!(is_retryable_status(500)); - assert!(is_retryable_status(502)); - assert!(is_retryable_status(503)); - assert!(is_retryable_status(504)); + use crate::testing::StubLlm; - // Client errors should not be retryable - assert!(!is_retryable_status(400)); - assert!(!is_retryable_status(401)); - assert!(!is_retryable_status(403)); - assert!(!is_retryable_status(404)); - assert!(!is_retryable_status(422)); - - // Success codes should not be retryable - assert!(!is_retryable_status(200)); - assert!(!is_retryable_status(201)); + fn make_request() -> CompletionRequest { + CompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")]) } + fn make_tool_request() -> ToolCompletionRequest { + ToolCompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")], vec![]) + } + + fn fast_config(max_retries: u32) -> RetryConfig { + RetryConfig { max_retries } + } + + // -- Backoff delay tests -- + #[test] fn test_retry_backoff_delay_exponential_growth() { // Run multiple samples to verify the range, accounting for jitter @@ -93,4 +280,127 @@ mod tests { let delay = retry_backoff_delay(30); assert!(delay.as_millis() >= 100); } + + // -- is_retryable() classification tests -- + + #[test] + fn test_is_retryable_classification() { + // Retryable + assert!(is_retryable(&LlmError::RequestFailed { + provider: "p".into(), + reason: "err".into(), + })); + assert!(is_retryable(&LlmError::RateLimited { + provider: "p".into(), + retry_after: None, + })); + assert!(is_retryable(&LlmError::InvalidResponse { + provider: "p".into(), + reason: "bad".into(), + })); + assert!(is_retryable(&LlmError::SessionRenewalFailed { + provider: "p".into(), + reason: "timeout".into(), + })); + assert!(is_retryable(&LlmError::Io(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "reset" + )))); + + // NOT retryable + assert!(!is_retryable(&LlmError::AuthFailed { + provider: "p".into(), + })); + assert!(!is_retryable(&LlmError::SessionExpired { + provider: "p".into(), + })); + assert!(!is_retryable(&LlmError::ContextLengthExceeded { + used: 100_000, + limit: 50_000, + })); + assert!(!is_retryable(&LlmError::ModelNotAvailable { + provider: "p".into(), + model: "m".into(), + })); + } + + // -- RetryProvider tests -- + + #[tokio::test] + async fn success_on_first_attempt() { + let stub = Arc::new(StubLlm::new("ok").with_model_name("test")); + let retry = RetryProvider::new(stub.clone(), fast_config(3)); + + let resp = retry.complete(make_request()).await; + assert!(resp.is_ok()); + assert_eq!(resp.unwrap().content, "ok"); + assert_eq!(stub.calls(), 1); + } + + #[tokio::test] + async fn retries_transient_errors_then_succeeds() { + // StubLlm starts failing, then we flip it to succeed. + // With max_retries=2, it will try 3 times total. + let stub = Arc::new(StubLlm::failing("test")); + let retry = RetryProvider::new(stub.clone(), fast_config(2)); + + // Spawn a task that flips the stub to succeed after a short delay + let stub_clone = stub.clone(); + tokio::spawn(async move { + // Wait for at least 1 retry attempt (backoff is ~1s, so 1.5s should be enough) + tokio::time::sleep(Duration::from_millis(1500)).await; + stub_clone.set_failing(false); + }); + + let resp = retry.complete(make_request()).await; + assert!(resp.is_ok()); + // Should have called at least twice (first fail, then succeed after flip) + assert!(stub.calls() >= 2); + } + + #[tokio::test] + async fn non_transient_error_fails_immediately() { + let stub = Arc::new(StubLlm::failing_non_transient("test")); + let retry = RetryProvider::new(stub.clone(), fast_config(3)); + + let err = retry.complete(make_request()).await.unwrap_err(); + assert!(matches!(err, LlmError::ContextLengthExceeded { .. })); + // Should only be called once — no retries for non-transient errors + assert_eq!(stub.calls(), 1); + } + + #[tokio::test] + async fn exhausts_retries_then_returns_error() { + let stub = Arc::new(StubLlm::failing("test")); + // max_retries=0 means only the initial attempt, no retries + let retry = RetryProvider::new(stub.clone(), fast_config(0)); + + let err = retry.complete(make_request()).await.unwrap_err(); + assert!(matches!(err, LlmError::RequestFailed { .. })); + assert_eq!(stub.calls(), 1); + } + + #[tokio::test] + async fn complete_with_tools_retries_same_as_complete() { + let stub = Arc::new(StubLlm::failing_non_transient("test")); + let retry = RetryProvider::new(stub.clone(), fast_config(3)); + + let err = retry + .complete_with_tools(make_tool_request()) + .await + .unwrap_err(); + assert!(matches!(err, LlmError::ContextLengthExceeded { .. })); + assert_eq!(stub.calls(), 1); + } + + #[tokio::test] + async fn passthrough_methods_delegate_to_inner() { + let stub = Arc::new(StubLlm::new("ok").with_model_name("my-model")); + let retry = RetryProvider::new(stub, fast_config(3)); + + assert_eq!(retry.model_name(), "my-model"); + assert_eq!(retry.active_model_name(), "my-model"); + assert_eq!(retry.cost_per_token(), (Decimal::ZERO, Decimal::ZERO)); + assert_eq!(retry.calculate_cost(100, 50), Decimal::ZERO); + } } diff --git a/src/llm/rig_adapter.rs b/src/llm/rig_adapter.rs index ce2b84af..ac352120 100644 --- a/src/llm/rig_adapter.rs +++ b/src/llm/rig_adapter.rs @@ -18,6 +18,8 @@ use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value as JsonValue; +use std::collections::HashSet; + use crate::error::LlmError; use crate::llm::costs; use crate::llm::provider::{ @@ -414,7 +416,9 @@ where ); } - let (preamble, history) = convert_messages(&request.messages); + let mut messages = request.messages; + crate::llm::provider::sanitize_tool_messages(&mut messages); + let (preamble, history) = convert_messages(&messages); let rig_req = build_rig_request( preamble, @@ -459,7 +463,12 @@ where ); } - let (preamble, history) = convert_messages(&request.messages); + let known_tool_names: HashSet = + request.tools.iter().map(|t| t.name.clone()).collect(); + + let mut messages = request.messages; + crate::llm::provider::sanitize_tool_messages(&mut messages); + let (preamble, history) = convert_messages(&messages); let tools = convert_tools(&request.tools); let tool_choice = convert_tool_choice(request.tool_choice.as_deref()); @@ -481,7 +490,20 @@ where reason: e.to_string(), })?; - let (text, tool_calls, finish) = extract_response(&response.choice, &response.usage); + let (text, mut tool_calls, finish) = extract_response(&response.choice, &response.usage); + + // Normalize tool call names: some proxies prepend "proxy_" prefixes. + for tc in &mut tool_calls { + let normalized = normalize_tool_name(&tc.name, &known_tool_names); + if normalized != tc.name { + tracing::debug!( + original = %tc.name, + normalized = %normalized, + "Normalized tool call name from provider", + ); + tc.name = normalized; + } + } Ok(ToolCompletionResponse { content: text, @@ -513,6 +535,25 @@ where } } +/// Normalize a tool call name returned by an OpenAI-compatible provider. +/// +/// Some proxies (e.g. VibeProxy) prepend `proxy_` to tool names. +/// If the returned name doesn't match any known tool but stripping a +/// `proxy_` prefix yields a match, use the stripped version. +fn normalize_tool_name(name: &str, known_tools: &HashSet) -> String { + if known_tools.contains(name) { + return name.to_string(); + } + + if let Some(stripped) = name.strip_prefix("proxy_") + && known_tools.contains(stripped) + { + return stripped.to_string(); + } + + name.to_string() +} + #[cfg(test)] mod tests { use super::*; @@ -801,4 +842,33 @@ mod tests { assert_eq!(saturate_u32(u64::MAX), u32::MAX); assert_eq!(saturate_u32(u32::MAX as u64), u32::MAX); } + + // -- normalize_tool_name tests -- + + #[test] + fn test_normalize_tool_name_exact_match() { + let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]); + assert_eq!(normalize_tool_name("echo", &known), "echo"); + } + + #[test] + fn test_normalize_tool_name_proxy_prefix_match() { + let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]); + assert_eq!(normalize_tool_name("proxy_echo", &known), "echo"); + } + + #[test] + fn test_normalize_tool_name_proxy_prefix_no_match_kept() { + let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]); + assert_eq!( + normalize_tool_name("proxy_unknown", &known), + "proxy_unknown" + ); + } + + #[test] + fn test_normalize_tool_name_unknown_passthrough() { + let known = HashSet::from(["echo".to_string()]); + assert_eq!(normalize_tool_name("other_tool", &known), "other_tool"); + } } diff --git a/src/main.rs b/src/main.rs index 7d47dcc8..7f0252a9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -26,9 +26,9 @@ use ironclaw::{ hooks::{HookRegistry, bootstrap_hooks}, llm::{ CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig, - FailoverProvider, LlmProvider, ResponseCacheConfig, SessionConfig, - create_cheap_llm_provider, create_llm_provider, create_llm_provider_with_config, - create_session_manager, + FailoverProvider, LlmProvider, ResponseCacheConfig, RetryConfig, RetryProvider, + SessionConfig, create_cheap_llm_provider, create_llm_provider, + create_llm_provider_with_config, create_session_manager, }, orchestrator::{ ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore, @@ -42,7 +42,9 @@ use ironclaw::{ mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated}, wasm::{WasmToolLoader, WasmToolRuntime, load_dev_tools}, }, - workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace}, + workspace::{ + EmbeddingProvider, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings, Workspace, + }, }; #[cfg(feature = "libsql")] @@ -115,18 +117,20 @@ async fn main() -> anyhow::Result<()> { &config.llm.nearai.base_url, session, ) - .with_model(&config.embeddings.model, 1536), + .with_model(&config.embeddings.model, config.embeddings.dimension), + )), + "ollama" => Some(Arc::new( + ironclaw::workspace::OllamaEmbeddings::new( + &config.embeddings.ollama_base_url, + ) + .with_model(&config.embeddings.model, config.embeddings.dimension), )), _ => { if let Some(api_key) = config.embeddings.openai_api_key() { - let dim = match config.embeddings.model.as_str() { - "text-embedding-3-large" => 3072, - _ => 1536, - }; Some(Arc::new(ironclaw::workspace::OpenAiEmbeddings::with_model( api_key, &config.embeddings.model, - dim, + config.embeddings.dimension, ))) } else { None @@ -137,6 +141,23 @@ async fn main() -> anyhow::Result<()> { None }; + // Warn if libSQL backend is used with non-1536 embedding dimension. + // libSQL schema uses F32_BLOB(1536) which cannot be altered without a + // table rebuild, so non-1536 embeddings will cause storage failures. + if config.database.backend == ironclaw::config::DatabaseBackend::LibSql + && config.embeddings.enabled + && config.embeddings.dimension != 1536 + { + tracing::warn!( + configured_dimension = config.embeddings.dimension, + "Embedding dimension {} is not 1536. The libSQL schema uses \ + F32_BLOB(1536) which requires exactly 1536 dimensions. \ + Embedding storage will fail. Use PostgreSQL or set \ + EMBEDDING_DIMENSION=1536.", + config.embeddings.dimension + ); + } + // Create a Database-trait-backed workspace for the memory command let db: Arc = ironclaw::db::connect_from_config(&config.database) @@ -599,6 +620,22 @@ async fn main() -> anyhow::Result<()> { let llm = create_llm_provider(&config.llm, session.clone())?; tracing::info!("LLM provider initialized: {}", llm.model_name()); + // Wrap each provider with RetryProvider for automatic retries on transient errors. + // RetryProvider sits inside FailoverProvider so each provider in the failover chain + // gets its own retry attempts before the failover moves to the next provider. + let retry_config = RetryConfig { + max_retries: config.llm.nearai.max_retries, + }; + let llm: Arc = if retry_config.max_retries > 0 { + tracing::info!( + max_retries = retry_config.max_retries, + "LLM retry wrapper enabled" + ); + Arc::new(RetryProvider::new(llm, retry_config.clone())) + } else { + llm + }; + // Wrap in failover if a fallback model is configured let llm: Arc = if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() { @@ -615,6 +652,12 @@ async fn main() -> anyhow::Result<()> { fallback = %fallback.model_name(), "LLM failover enabled" ); + // Wrap fallback with retry too + let fallback: Arc = if retry_config.max_retries > 0 { + Arc::new(RetryProvider::new(fallback, retry_config.clone())) + } else { + fallback + }; let cooldown_config = CooldownConfig { cooldown_duration: std::time::Duration::from_secs( config.llm.nearai.failover_cooldown_secs, @@ -685,28 +728,39 @@ async fn main() -> anyhow::Result<()> { match config.embeddings.provider.as_str() { "nearai" => { tracing::info!( - "Embeddings enabled via NEAR AI (model: {})", - config.embeddings.model + "Embeddings enabled via NEAR AI (model: {}, dim: {})", + config.embeddings.model, + config.embeddings.dimension, ); Some(Arc::new( NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone()) - .with_model(&config.embeddings.model, 1536), + .with_model(&config.embeddings.model, config.embeddings.dimension), + )) + } + "ollama" => { + tracing::info!( + "Embeddings enabled via Ollama (model: {}, url: {}, dim: {})", + config.embeddings.model, + config.embeddings.ollama_base_url, + config.embeddings.dimension, + ); + Some(Arc::new( + OllamaEmbeddings::new(&config.embeddings.ollama_base_url) + .with_model(&config.embeddings.model, config.embeddings.dimension), )) } _ => { // Default to OpenAI for unknown providers if let Some(api_key) = config.embeddings.openai_api_key() { tracing::info!( - "Embeddings enabled via OpenAI (model: {})", - config.embeddings.model + "Embeddings enabled via OpenAI (model: {}, dim: {})", + config.embeddings.model, + config.embeddings.dimension, ); Some(Arc::new(OpenAiEmbeddings::with_model( api_key, &config.embeddings.model, - match config.embeddings.model.as_str() { - "text-embedding-3-large" => 3072, - _ => 1536, // text-embedding-3-small and ada-002 - }, + config.embeddings.dimension, ))) } else { tracing::warn!("Embeddings configured but OPENAI_API_KEY not set"); @@ -719,6 +773,21 @@ async fn main() -> anyhow::Result<()> { None }; + // Warn if libSQL backend is used with non-1536 embedding dimension. + if config.database.backend == ironclaw::config::DatabaseBackend::LibSql + && config.embeddings.enabled + && config.embeddings.dimension != 1536 + { + tracing::warn!( + configured_dimension = config.embeddings.dimension, + "Embedding dimension {} is not 1536. The libSQL schema uses \ + F32_BLOB(1536) which requires exactly 1536 dimensions. \ + Embedding storage will fail. Use PostgreSQL or set \ + EMBEDDING_DIMENSION=1536.", + config.embeddings.dimension + ); + } + // Register memory tools if database is available if let Some(ref db) = db { let mut workspace = Workspace::new_with_db("default", Arc::clone(db)); diff --git a/src/sandbox/config.rs b/src/sandbox/config.rs index 41fc36fa..3ebeb96e 100644 --- a/src/sandbox/config.rs +++ b/src/sandbox/config.rs @@ -34,7 +34,7 @@ impl Default for SandboxConfig { memory_limit_mb: 2048, cpu_shares: 1024, network_allowlist: default_allowlist(), - image: "ghcr.io/nearai/sandbox:latest".to_string(), + image: "ironclaw-worker:latest".to_string(), auto_pull_image: true, proxy_port: 0, } diff --git a/src/settings.rs b/src/settings.rs index 0f12dcec..fd4d45f9 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -462,7 +462,7 @@ fn default_sandbox_cpu_shares() -> u32 { } fn default_sandbox_image() -> String { - "ghcr.io/nearai/sandbox:latest".to_string() + "ironclaw-worker:latest".to_string() } impl Default for SandboxSettings { diff --git a/src/workspace/embeddings.rs b/src/workspace/embeddings.rs index 320eca30..42340fcb 100644 --- a/src/workspace/embeddings.rs +++ b/src/workspace/embeddings.rs @@ -354,6 +354,123 @@ impl EmbeddingProvider for NearAiEmbeddings { } } +/// Ollama embedding provider using a local Ollama instance. +/// +/// Ollama serves embedding models (e.g. `nomic-embed-text`, `mxbai-embed-large`) +/// via a REST API, typically at `http://localhost:11434`. +pub struct OllamaEmbeddings { + client: reqwest::Client, + base_url: String, + model: String, + dimension: usize, +} + +impl OllamaEmbeddings { + /// Create a new Ollama embedding provider. + /// + /// Defaults to `nomic-embed-text` (768 dimensions). + pub fn new(base_url: impl Into) -> Self { + Self { + client: reqwest::Client::new(), + base_url: base_url.into(), + model: "nomic-embed-text".to_string(), + dimension: 768, + } + } + + /// Use a specific model with a given dimension. + pub fn with_model(mut self, model: impl Into, dimension: usize) -> Self { + self.model = model.into(); + self.dimension = dimension; + self + } +} + +#[derive(Debug, Serialize)] +struct OllamaEmbedRequest<'a> { + model: &'a str, + input: &'a [String], +} + +#[derive(Debug, Deserialize)] +struct OllamaEmbedResponse { + embeddings: Vec>, +} + +#[async_trait] +impl EmbeddingProvider for OllamaEmbeddings { + fn dimension(&self) -> usize { + self.dimension + } + + fn model_name(&self) -> &str { + &self.model + } + + fn max_input_length(&self) -> usize { + // Most Ollama embedding models support 8192 tokens (~32k chars) + 32_000 + } + + async fn embed(&self, text: &str) -> Result, EmbeddingError> { + if text.len() > self.max_input_length() { + return Err(EmbeddingError::TextTooLong { + length: text.len(), + max: self.max_input_length(), + }); + } + + let embeddings = self.embed_batch(&[text.to_string()]).await?; + embeddings + .into_iter() + .next() + .ok_or_else(|| EmbeddingError::InvalidResponse("No embedding returned".to_string())) + } + + async fn embed_batch(&self, texts: &[String]) -> Result>, EmbeddingError> { + if texts.is_empty() { + return Ok(Vec::new()); + } + + let request = OllamaEmbedRequest { + model: &self.model, + input: texts, + }; + + let url = format!("{}/api/embed", self.base_url); + + let response = self.client.post(&url).json(&request).send().await?; + + let status = response.status(); + + if !status.is_success() { + let error_text = response.text().await.unwrap_or_default(); + return Err(EmbeddingError::HttpError(format!( + "Ollama returned HTTP {}: {}", + status, error_text + ))); + } + + let result: OllamaEmbedResponse = response.json().await.map_err(|e| { + EmbeddingError::InvalidResponse(format!("Failed to parse Ollama response: {}", e)) + })?; + + // Validate that returned embeddings match the configured dimension. + for (i, emb) in result.embeddings.iter().enumerate() { + if emb.len() != self.dimension { + return Err(EmbeddingError::InvalidResponse(format!( + "Ollama returned embedding of dimension {}, expected {} at index {}", + emb.len(), + self.dimension, + i + ))); + } + } + + Ok(result.embeddings) + } +} + /// A mock embedding provider for testing. /// /// Generates deterministic embeddings based on text hash. diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index 3165ecce..d7d7890e 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -50,7 +50,9 @@ mod search; pub use chunker::{ChunkConfig, chunk_document}; pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths}; -pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings}; +pub use embeddings::{ + EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings, +}; #[cfg(feature = "postgres")] pub use repository::Repository; pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};