From a158eee1b0729ae6e2a202efb17ffbd02e710c94 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Tue, 17 Feb 2026 11:35:56 -0800 Subject: [PATCH] feat: 10 infrastructure improvements from zeroclaw (#126) * refactor: break up agent_loop.rs into four focused modules Split the monolithic 2835-line agent_loop.rs into: - agent_loop.rs (722L): Agent struct, event loop, message dispatch - dispatcher.rs (635L): Agentic tool loop, tool execution, auth detection - commands.rs (484L): System commands, job handlers, heartbeat, summarize - thread_ops.rs (1059L): Thread lifecycle, approval, undo/redo, persistence Each module gets its own impl Agent block. Agent fields changed to pub(super) so sibling modules in the agent package can access them. All 16 existing tests pass in their new locations. Inspired by ZeroClaw's agent module split (agent.rs, loop_.rs, dispatcher.rs, prompt.rs, memory_loader.rs). Co-Authored-By: Claude Opus 4.6 * feat: add cost caps and guardrails for autonomous agent spending Daily budget (MAX_COST_PER_DAY_CENTS) and hourly action rate (MAX_ACTIONS_PER_HOUR) limits prevent runaway agents from burning through API credits, especially in daemon/heartbeat modes. - CostGuard with pre-flight check and post-call recording - Sliding window for hourly rate, midnight-UTC daily reset - 80% threshold warning, atomic fast-path for exceeded budget - Wired into dispatcher loop (check before LLM call, record after) Co-Authored-By: Claude Opus 4.6 * feat: add circuit breaker on LLM providers Wraps LlmProvider with a Closed/Open/HalfOpen state machine that trips after consecutive transient failures, preventing request storms against a degraded backend. Automatically probes for recovery. - CircuitBreakerProvider implements LlmProvider (drop-in wrapper) - Transient error classification (server, rate-limit, network, auth infra) - Client errors (wrong model, context overflow) don't trip the breaker - Configurable via CIRCUIT_BREAKER_THRESHOLD and CIRCUIT_BREAKER_RECOVERY_SECS - Composes with existing FailoverProvider (circuit breaker wraps failover) - 12 tests covering full state machine and error classification Co-Authored-By: Claude Opus 4.6 * feat: add tunnel abstraction for remote access Trait-based tunnel system with lifecycle management (start/stop/health) for exposing the agent to the internet through external tunnel binaries. Five providers: - Cloudflare Tunnel (cloudflared, Zero Trust token auth) - Tailscale (serve for tailnet, funnel for public) - ngrok (with optional custom domain) - Custom (arbitrary command with {host}/{port} placeholders) - None (local-only, no external exposure) Config via TUNNEL_PROVIDER + provider-specific env vars. Extends existing TunnelConfig with optional managed provider alongside the static TUNNEL_URL path. Factory, shared process management, and 37 tests covering all providers and edge cases. Co-Authored-By: Claude Opus 4.6 * feat: add OS service management (launchd/systemd) Adds `ironclaw service {install,start,stop,status,uninstall}` for running the agent as a background daemon. macOS uses launchd plists under ~/Library/LaunchAgents, Linux uses systemd user units. Co-Authored-By: Claude Opus 4.6 * feat: add observability trait system with noop, log, and multi backends Introduces an Observer trait for recording agent lifecycle events and metrics, with pluggable backends. The noop backend compiles to zero overhead, log backend uses tracing, and multi fans out to multiple observers. Configured via OBSERVABILITY_BACKEND env var. Co-Authored-By: Claude Opus 4.6 * feat: add in-memory LLM response cache with TTL and LRU eviction CachedProvider wraps any LlmProvider and caches complete() responses keyed by SHA-256(model + messages). Tool-calling requests are never cached since they trigger side effects. Configurable via RESPONSE_CACHE_ENABLED, RESPONSE_CACHE_TTL_SECS, and RESPONSE_CACHE_MAX_ENTRIES env vars. Co-Authored-By: Claude Opus 4.6 * feat: add memory hygiene with cadence-gated daily log cleanup Adds workspace::hygiene module that automatically deletes daily log documents older than a configurable retention period (default 30 days). Runs on a 12-hour cadence tracked via a local state file to avoid redundant passes. Best-effort design: failures are logged, never fatal. Co-Authored-By: Claude Opus 4.6 * feat: add doctor diagnostics command for active health probing Probes external dependencies (Docker, cloudflared, ngrok, tailscale), validates NEAR AI session, checks database connectivity, and verifies workspace directory. Complements the passive `status` command. Co-Authored-By: Claude Opus 4.6 * feat: add structured TOML config file support Adds ~/.ironclaw/config.toml as a configuration layer between env vars and database settings. Priority: env var > TOML file > DB > defaults. - `ironclaw config init` generates a commented config.toml from current settings - `ironclaw --config path/to/config.toml` loads a custom config file - Settings.merge_from() only overlays non-default values from the TOML file - `ironclaw config path` now shows TOML file status Co-Authored-By: Claude Opus 4.6 * fix: address codex review findings - apply_toml_overlay now returns Result and errors on explicit missing or invalid config paths (was log-only, violating the documented contract that explicit paths are fatal) - custom tunnel url_pattern is now used to filter extracted URLs, not just as a gate for scanning stdout - systemd ExecStart path is now quoted to handle spaces in paths Co-Authored-By: Claude Opus 4.6 * fix: address PR review feedback - Cache key now includes max_tokens, temperature, and stop_sequences so different request parameters produce distinct keys - to_cents() uses .trunc() + parse:: instead of f64 intermediary, avoiding precision loss for large values - Tailscale public URL no longer includes local port (serve/funnel expose on standard HTTPS port 443) Co-Authored-By: Claude Opus 4.6 * feat: wire up tunnel lifecycle and fix audit findings Connect the tunnel module to the rest of the application so that setting TUNNEL_PROVIDER actually starts a managed tunnel at boot and stops it on shutdown. Previously create_tunnel() was never called outside tests. Changes: - Expand TunnelSettings with provider credential fields (settings.rs) - TunnelConfig::resolve() falls back to DB settings when env vars unset - Start tunnel at boot, stop on shutdown, show URL in boot screen - Setup wizard collects provider-specific credentials (ngrok, cloudflare, tailscale, custom, static URL) - Fix public_url() returning None under lock contention (SharedUrl) - Fix local_host parameter ignored by cloudflare/ngrok/tailscale - Fix tailscale silent fallback to "localhost" on bad JSON - Fix ngrok globally mutating config via add-authtoken (use env var) - Add 10s timeout to tailscale status --json Co-Authored-By: Claude Opus 4.6 * fix: address PR review comments - Document split_whitespace limitation in CustomTunnel doc comment - Remove unnecessary quotes from systemd ExecStart directive Co-Authored-By: Claude Opus 4.6 * fix: address PR review feedback (round 3) - doctor: missing libSQL DB on fresh install is Pass, not Fail - service: quote ExecStart path for systemd space handling Co-Authored-By: Claude Opus 4.6 * fix: correct cost guard doc comment (LLM calls, not LLM/tool) Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- Cargo.lock | 1 + Cargo.toml | 1 + src/agent/agent_loop.rs | 2160 +---------------------------------- src/agent/commands.rs | 484 ++++++++ src/agent/cost_guard.rs | 339 ++++++ src/agent/dispatcher.rs | 635 ++++++++++ src/agent/mod.rs | 4 + src/agent/thread_ops.rs | 1063 +++++++++++++++++ src/boot_screen.rs | 20 + src/cli/config.rs | 89 ++ src/cli/doctor.rs | 287 +++++ src/cli/mod.rs | 13 + src/cli/service.rs | 37 + src/config.rs | 185 ++- src/lib.rs | 3 + src/llm/circuit_breaker.rs | 674 +++++++++++ src/llm/mod.rs | 11 +- src/llm/response_cache.rs | 526 +++++++++ src/llm/session.rs | 6 +- src/main.rs | 148 ++- src/observability/log.rs | 181 +++ src/observability/mod.rs | 105 ++ src/observability/multi.rs | 139 +++ src/observability/noop.rs | 60 + src/observability/traits.rs | 143 +++ src/service.rs | 357 ++++++ src/settings.rs | 299 ++++- src/setup/channels.rs | 148 ++- src/setup/wizard.rs | 25 +- src/tunnel/cloudflare.rs | 140 +++ src/tunnel/custom.rs | 245 ++++ src/tunnel/mod.rs | 329 ++++++ src/tunnel/ngrok.rs | 142 +++ src/tunnel/none.rs | 62 + src/tunnel/tailscale.rs | 140 +++ src/workspace/hygiene.rs | 247 ++++ src/workspace/mod.rs | 1 + 37 files changed, 7272 insertions(+), 2177 deletions(-) create mode 100644 src/agent/commands.rs create mode 100644 src/agent/cost_guard.rs create mode 100644 src/agent/dispatcher.rs create mode 100644 src/agent/thread_ops.rs create mode 100644 src/cli/doctor.rs create mode 100644 src/cli/service.rs create mode 100644 src/llm/circuit_breaker.rs create mode 100644 src/llm/response_cache.rs create mode 100644 src/observability/log.rs create mode 100644 src/observability/mod.rs create mode 100644 src/observability/multi.rs create mode 100644 src/observability/noop.rs create mode 100644 src/observability/traits.rs create mode 100644 src/service.rs create mode 100644 src/tunnel/cloudflare.rs create mode 100644 src/tunnel/custom.rs create mode 100644 src/tunnel/mod.rs create mode 100644 src/tunnel/ngrok.rs create mode 100644 src/tunnel/none.rs create mode 100644 src/tunnel/tailscale.rs create mode 100644 src/workspace/hygiene.rs diff --git a/Cargo.lock b/Cargo.lock index 61307200..c5929bce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2544,6 +2544,7 @@ dependencies = [ "tokio-stream", "tokio-test", "tokio-tungstenite 0.26.2", + "toml", "tower 0.5.3", "tower-http 0.6.8", "tracing", diff --git a/Cargo.toml b/Cargo.toml index 1e45e3fd..16af4d0c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } # Configuration dotenvy = "0.15" +toml = "0.8" # Core types uuid = { version = "1", features = ["v4", "serde"] } diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index b1015fe8..09110ee8 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -1,29 +1,31 @@ //! Main agent loop. +//! +//! Contains the `Agent` struct, `AgentDeps`, and the core event loop (`run`). +//! The heavy lifting is delegated to sibling modules: +//! +//! - `dispatcher` - Tool dispatch (agentic loop, tool execution) +//! - `commands` - System commands and job handlers +//! - `thread_ops` - Thread/session operations (user input, undo, approval, persistence) use std::sync::Arc; use futures::StreamExt; -use tokio::sync::Mutex; -use uuid::Uuid; -use crate::agent::compaction::ContextCompactor; use crate::agent::context_monitor::ContextMonitor; use crate::agent::heartbeat::spawn_heartbeat; use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker}; use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair}; -use crate::agent::session::{PendingApproval, Session, ThreadState}; use crate::agent::session_manager::SessionManager; use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult}; -use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, MessageIntent, Router, Scheduler}; +use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, Router, Scheduler}; use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate}; use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig}; use crate::context::ContextManager; -use crate::context::JobContext; use crate::db::Database; use crate::error::Error; use crate::extensions::ExtensionManager; use crate::hooks::HookRegistry; -use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult}; +use crate::llm::LlmProvider; use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; use crate::workspace::Workspace; @@ -51,17 +53,6 @@ pub(crate) fn truncate_for_preview(output: &str, max_chars: usize) -> String { } } -/// Result of the agentic loop execution. -enum AgenticLoopResult { - /// Completed with a response. - Response(String), - /// A tool requires approval before continuing. - NeedApproval { - /// The pending approval request to store. - pending: PendingApproval, - }, -} - /// Core dependencies for the agent. /// /// Bundles the shared components to reduce argument count. @@ -76,20 +67,22 @@ pub struct AgentDeps { pub workspace: Option>, pub extension_manager: Option>, pub hooks: Arc, + /// Cost enforcement guardrails (daily budget, hourly rate limits). + pub cost_guard: Arc, } /// The main agent that coordinates all components. pub struct Agent { - config: AgentConfig, - deps: AgentDeps, - channels: Arc, - context_manager: Arc, - scheduler: Arc, - router: Router, - session_manager: Arc, - context_monitor: ContextMonitor, - heartbeat_config: Option, - routine_config: Option, + pub(super) config: AgentConfig, + pub(super) deps: AgentDeps, + pub(super) channels: Arc, + pub(super) context_manager: Arc, + pub(super) scheduler: Arc, + pub(super) router: Router, + pub(super) session_manager: Arc, + pub(super) context_monitor: ContextMonitor, + pub(super) heartbeat_config: Option, + pub(super) routine_config: Option, } impl Agent { @@ -136,35 +129,40 @@ impl Agent { } // Convenience accessors - fn store(&self) -> Option<&Arc> { + + pub(super) fn store(&self) -> Option<&Arc> { self.deps.store.as_ref() } - fn llm(&self) -> &Arc { + pub(super) fn llm(&self) -> &Arc { &self.deps.llm } /// Get the cheap/fast LLM provider, falling back to the main one. - fn cheap_llm(&self) -> &Arc { + pub(super) fn cheap_llm(&self) -> &Arc { self.deps.cheap_llm.as_ref().unwrap_or(&self.deps.llm) } - fn safety(&self) -> &Arc { + pub(super) fn safety(&self) -> &Arc { &self.deps.safety } - fn tools(&self) -> &Arc { + pub(super) fn tools(&self) -> &Arc { &self.deps.tools } - fn workspace(&self) -> Option<&Arc> { + pub(super) fn workspace(&self) -> Option<&Arc> { self.deps.workspace.as_ref() } - fn hooks(&self) -> &Arc { + pub(super) fn hooks(&self) -> &Arc { &self.deps.hooks } + pub(super) fn cost_guard(&self) -> &Arc { + &self.deps.cost_guard + } + /// Run the agent main loop. pub async fn run(self) -> Result<(), Error> { // Start channels @@ -660,2102 +658,10 @@ impl Agent { } } } - - /// Hydrate a historical thread from DB into memory if not already present. - /// - /// Called before `resolve_thread` so that the session manager finds the - /// thread on lookup instead of creating a new one. - /// - /// Creates an in-memory thread with the exact UUID the frontend sent, - /// even when the conversation has zero messages (e.g. a brand-new - /// assistant thread). Without this, `resolve_thread` would mint a - /// fresh UUID and all messages would land in the wrong conversation. - async fn maybe_hydrate_thread(&self, message: &IncomingMessage, external_thread_id: &str) { - // Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs) - let thread_uuid = match Uuid::parse_str(external_thread_id) { - Ok(id) => id, - Err(_) => return, - }; - - // Check if already in memory - let session = self - .session_manager - .get_or_create_session(&message.user_id) - .await; - { - let sess = session.lock().await; - if sess.threads.contains_key(&thread_uuid) { - return; - } - } - - // Load history from DB (may be empty for a newly created thread). - let mut chat_messages: Vec = Vec::new(); - let msg_count; - - if let Some(store) = self.store() { - let db_messages = store - .list_conversation_messages(thread_uuid) - .await - .unwrap_or_default(); - msg_count = db_messages.len(); - chat_messages = db_messages - .iter() - .filter_map(|m| match m.role.as_str() { - "user" => Some(ChatMessage::user(&m.content)), - "assistant" => Some(ChatMessage::assistant(&m.content)), - _ => None, - }) - .collect(); - } else { - msg_count = 0; - } - - // Create thread with the historical ID and restore messages - let session_id = { - let sess = session.lock().await; - sess.id - }; - - let mut thread = crate::agent::session::Thread::with_id(thread_uuid, session_id); - if !chat_messages.is_empty() { - thread.restore_from_messages(chat_messages); - } - - // Restore response chain from conversation metadata - if let Some(store) = self.store() - && let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await - && let Some(rid) = metadata - .get("last_response_id") - .and_then(|v| v.as_str()) - .map(String::from) - { - thread.last_response_id = Some(rid.clone()); - self.llm() - .seed_response_chain(&thread_uuid.to_string(), rid); - tracing::debug!("Restored response chain for thread {}", thread_uuid); - } - - // Insert into session and register with session manager - { - let mut sess = session.lock().await; - sess.threads.insert(thread_uuid, thread); - sess.active_thread = Some(thread_uuid); - sess.last_active_at = chrono::Utc::now(); - } - - self.session_manager - .register_thread( - &message.user_id, - &message.channel, - thread_uuid, - Arc::clone(&session), - ) - .await; - - tracing::debug!( - "Hydrated thread {} from DB ({} messages)", - thread_uuid, - msg_count - ); - } - - async fn process_user_input( - &self, - message: &IncomingMessage, - session: Arc>, - thread_id: Uuid, - content: &str, - ) -> Result { - // First check thread state without holding lock during I/O - let thread_state = { - let sess = session.lock().await; - let thread = sess - .threads - .get(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - thread.state - }; - - // Check thread state - match thread_state { - ThreadState::Processing => { - return Ok(SubmissionResult::error( - "Turn in progress. Use /interrupt to cancel.", - )); - } - ThreadState::AwaitingApproval => { - return Ok(SubmissionResult::error( - "Waiting for approval. Use /interrupt to cancel.", - )); - } - ThreadState::Completed => { - return Ok(SubmissionResult::error( - "Thread completed. Use /thread new.", - )); - } - ThreadState::Idle | ThreadState::Interrupted => { - // Can proceed - } - } - - // Safety validation for user input - let validation = self.safety().validate_input(content); - if !validation.is_valid { - let details = validation - .errors - .iter() - .map(|e| format!("{}: {}", e.field, e.message)) - .collect::>() - .join("; "); - return Ok(SubmissionResult::error(format!( - "Input rejected by safety validation: {}", - details - ))); - } - - let violations = self.safety().check_policy(content); - if violations - .iter() - .any(|rule| rule.action == crate::safety::PolicyAction::Block) - { - return Ok(SubmissionResult::error("Input rejected by safety policy.")); - } - - // Handle explicit commands (starting with /) directly - // Everything else goes through the normal agentic loop with tools - let temp_message = IncomingMessage { - content: content.to_string(), - ..message.clone() - }; - - if let Some(intent) = self.router.route_command(&temp_message) { - // Explicit command like /status, /job, /list - handle directly - return self.handle_job_or_command(intent, message).await; - } - - // Natural language goes through the agentic loop - // Job tools (create_job, list_jobs, etc.) are in the tool registry - - // Auto-compact if needed BEFORE adding new turn - { - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - let messages = thread.messages(); - if let Some(strategy) = self.context_monitor.suggest_compaction(&messages) { - let pct = self.context_monitor.usage_percent(&messages); - tracing::info!("Context at {:.1}% capacity, auto-compacting", pct); - - // Notify the user that compaction is happening - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Status(format!( - "Context at {:.0}% capacity, compacting...", - pct - )), - &message.metadata, - ) - .await; - - let compactor = ContextCompactor::new(self.llm().clone()); - if let Err(e) = compactor - .compact(thread, strategy, self.workspace().map(|w| w.as_ref())) - .await - { - tracing::warn!("Auto-compaction failed: {}", e); - } - } - } - - // Create checkpoint before turn - let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; - { - let sess = session.lock().await; - let thread = sess - .threads - .get(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - let mut mgr = undo_mgr.lock().await; - mgr.checkpoint( - thread.turn_number(), - thread.messages(), - format!("Before turn {}", thread.turn_number()), - ); - } - - // Start the turn and get messages - let turn_messages = { - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - thread.start_turn(content); - thread.messages() - }; - - // Send thinking status - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Thinking("Processing...".into()), - &message.metadata, - ) - .await; - - // Run the agentic tool execution loop - let result = self - .run_agentic_loop(message, session.clone(), thread_id, turn_messages, false) - .await; - - // Re-acquire lock and check if interrupted - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - if thread.state == ThreadState::Interrupted { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Status("Interrupted".into()), - &message.metadata, - ) - .await; - return Ok(SubmissionResult::Interrupted); - } - - // Complete, fail, or request approval - match result { - Ok(AgenticLoopResult::Response(response)) => { - // Hook: TransformResponse — allow hooks to modify or reject the final response - let response = { - let event = crate::hooks::HookEvent::ResponseTransform { - user_id: message.user_id.clone(), - thread_id: thread_id.to_string(), - response: response.clone(), - }; - match self.hooks().run(&event).await { - Err(crate::hooks::HookError::Rejected { reason }) => { - format!("[Response filtered: {}]", reason) - } - Err(err) => { - format!("[Response blocked by hook policy: {}]", err) - } - Ok(crate::hooks::HookOutcome::Continue { - modified: Some(new_response), - }) => new_response, - _ => response, // fail-open: use original - } - }; - - thread.complete_turn(&response); - self.persist_response_chain(thread); - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Status("Done".into()), - &message.metadata, - ) - .await; - - // Fire-and-forget: persist turn to DB - self.persist_turn(thread_id, &message.user_id, content, Some(&response)); - - Ok(SubmissionResult::response(response)) - } - Ok(AgenticLoopResult::NeedApproval { pending }) => { - // Store pending approval in thread and update state - let request_id = pending.request_id; - let tool_name = pending.tool_name.clone(); - let description = pending.description.clone(); - let parameters = pending.parameters.clone(); - thread.await_approval(pending); - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Status("Awaiting approval".into()), - &message.metadata, - ) - .await; - Ok(SubmissionResult::NeedApproval { - request_id, - tool_name, - description, - parameters, - }) - } - Err(e) => { - thread.fail_turn(e.to_string()); - - // Persist the user message even on failure - self.persist_turn(thread_id, &message.user_id, content, None); - - Ok(SubmissionResult::error(e.to_string())) - } - } - } - - /// Fire-and-forget: persist a turn (user message + optional assistant response) to the DB. - fn persist_turn( - &self, - thread_id: Uuid, - user_id: &str, - user_input: &str, - response: Option<&str>, - ) { - let store = match self.store() { - Some(s) => Arc::clone(s), - None => return, - }; - - let user_id = user_id.to_string(); - let user_input = user_input.to_string(); - let response = response.map(String::from); - - tokio::spawn(async move { - if let Err(e) = store - .ensure_conversation(thread_id, "gateway", &user_id, None) - .await - { - tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e); - return; - } - - if let Err(e) = store - .add_conversation_message(thread_id, "user", &user_input) - .await - { - tracing::warn!("Failed to persist user message: {}", e); - return; - } - - if let Some(ref resp) = response - && let Err(e) = store - .add_conversation_message(thread_id, "assistant", resp) - .await - { - tracing::warn!("Failed to persist assistant message: {}", e); - } - }); - } - - /// Sync the provider's response chain ID to the thread and DB metadata. - /// - /// Call after a successful agentic loop to persist the latest - /// `previous_response_id` so chaining survives restarts. - fn persist_response_chain(&self, thread: &mut crate::agent::session::Thread) { - let tid = thread.id.to_string(); - let response_id = match self.llm().get_response_chain_id(&tid) { - Some(rid) => rid, - None => return, - }; - - // Update in-memory thread - thread.last_response_id = Some(response_id.clone()); - - // Fire-and-forget DB write - let store = match self.store() { - Some(s) => Arc::clone(s), - None => return, - }; - let thread_id = thread.id; - tokio::spawn(async move { - let val = serde_json::json!(response_id); - if let Err(e) = store - .update_conversation_metadata_field(thread_id, "last_response_id", &val) - .await - { - tracing::warn!( - "Failed to persist response chain for thread {}: {}", - thread_id, - e - ); - } - }); - } - - /// Run the agentic loop: call LLM, execute tools, repeat until text response. - /// - /// Returns `AgenticLoopResult::Response` on completion, or - /// `AgenticLoopResult::NeedApproval` if a tool requires user approval. - /// - /// When `resume_after_tool` is true the loop already knows a tool was - /// executed earlier in this turn (e.g. an approved tool), so it won't - /// force the LLM to use tools if it responds with text. - async fn run_agentic_loop( - &self, - message: &IncomingMessage, - session: Arc>, - thread_id: Uuid, - initial_messages: Vec, - resume_after_tool: bool, - ) -> Result { - // Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.) - let system_prompt = if let Some(ws) = self.workspace() { - match ws.system_prompt().await { - Ok(prompt) if !prompt.is_empty() => Some(prompt), - Ok(_) => None, - Err(e) => { - tracing::debug!("Could not load workspace system prompt: {}", e); - None - } - } - } else { - None - }; - - let mut reasoning = Reasoning::new(self.llm().clone(), self.safety().clone()); - if let Some(prompt) = system_prompt { - reasoning = reasoning.with_system_prompt(prompt); - } - - // Build context with messages that we'll mutate during the loop - let mut context_messages = initial_messages; - - // Create a JobContext for tool execution (chat doesn't have a real job) - let job_ctx = JobContext::with_user(&message.user_id, "chat", "Interactive chat session"); - - const MAX_TOOL_ITERATIONS: usize = 10; - let mut iteration = 0; - let mut tools_executed = resume_after_tool; - - loop { - iteration += 1; - if iteration > MAX_TOOL_ITERATIONS { - return Err(crate::error::LlmError::InvalidResponse { - provider: "agent".to_string(), - reason: format!("Exceeded maximum tool iterations ({})", MAX_TOOL_ITERATIONS), - } - .into()); - } - - // Check if interrupted - { - let sess = session.lock().await; - if let Some(thread) = sess.threads.get(&thread_id) - && thread.state == ThreadState::Interrupted - { - return Err(crate::error::JobError::ContextError { - id: thread_id, - reason: "Interrupted".to_string(), - } - .into()); - } - } - - // Refresh tool definitions each iteration so newly built tools become visible - let tool_defs = self.tools().tool_definitions().await; - - // Call LLM with current context - let context = ReasoningContext::new() - .with_messages(context_messages.clone()) - .with_tools(tool_defs) - .with_metadata({ - let mut m = std::collections::HashMap::new(); - m.insert("thread_id".to_string(), thread_id.to_string()); - m - }); - - let output = reasoning.respond_with_tools(&context).await?; - - // Track token usage for budget enforcement - tracing::debug!( - "LLM call used {} input + {} output tokens", - output.usage.input_tokens, - output.usage.output_tokens - ); - - match output.result { - RespondResult::Text(text) => { - // If no tools have been executed yet, prompt the LLM to use tools - // This handles the case where the model explains what it will do - // instead of actually calling tools - if !tools_executed && iteration < 3 { - tracing::debug!( - "No tools executed yet (iteration {}), prompting for tool use", - iteration - ); - context_messages.push(ChatMessage::assistant(&text)); - context_messages.push(ChatMessage::user( - "Please proceed and use the available tools to complete this task.", - )); - continue; - } - - // Tools have been executed or we've tried multiple times, return response - return Ok(AgenticLoopResult::Response(text)); - } - RespondResult::ToolCalls { - tool_calls, - content, - } => { - tools_executed = true; - - // Add the assistant message with tool_calls to context. - // OpenAI protocol requires this before tool-result messages. - context_messages.push(ChatMessage::assistant_with_tool_calls( - content, - tool_calls.clone(), - )); - - // Execute tools and add results to context - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Thinking(format!( - "Executing {} tool(s)...", - tool_calls.len() - )), - &message.metadata, - ) - .await; - - // Record tool calls in the thread - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) - && let Some(turn) = thread.last_turn_mut() - { - for tc in &tool_calls { - turn.record_tool_call(&tc.name, tc.arguments.clone()); - } - } - } - - // Execute each tool (with approval checking and hook interception) - for mut tc in tool_calls { - // Check if tool requires approval - if let Some(tool) = self.tools().get(&tc.name).await - && tool.requires_approval() - { - // Check if auto-approved for this session - let mut is_auto_approved = { - let sess = session.lock().await; - sess.is_tool_auto_approved(&tc.name) - }; - - // Let the tool inspect the specific parameters and - // override auto-approval (e.g. destructive shell commands). - if is_auto_approved && tool.requires_approval_for(&tc.arguments) { - tracing::info!( - tool = %tc.name, - "Tool requires explicit approval for these parameters despite auto-approve" - ); - is_auto_approved = false; - } - - if !is_auto_approved { - // Need approval - store pending request and return - let 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(), - }; - - return Ok(AgenticLoopResult::NeedApproval { pending }); - } - } - - // Hook: BeforeToolCall — allow hooks to modify or reject tool calls - { - let event = crate::hooks::HookEvent::ToolCall { - tool_name: tc.name.clone(), - parameters: tc.arguments.clone(), - user_id: message.user_id.clone(), - context: "chat".to_string(), - }; - match self.hooks().run(&event).await { - Err(crate::hooks::HookError::Rejected { reason }) => { - context_messages.push(ChatMessage::tool_result( - &tc.id, - &tc.name, - format!("Tool call rejected by hook: {}", reason), - )); - continue; - } - Err(err) => { - context_messages.push(ChatMessage::tool_result( - &tc.id, - &tc.name, - format!("Tool call blocked by hook policy: {}", err), - )); - continue; - } - Ok(crate::hooks::HookOutcome::Continue { - modified: Some(new_params), - }) => match serde_json::from_str(&new_params) { - Ok(parsed) => tc.arguments = parsed, - Err(e) => { - tracing::warn!( - tool = %tc.name, - "Hook returned non-JSON modification for ToolCall, ignoring: {}", - e - ); - } - }, - _ => {} // Continue, fail-open errors already logged - } - } - - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolStarted { - name: tc.name.clone(), - }, - &message.metadata, - ) - .await; - - let tool_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: tool_result.is_ok(), - }, - &message.metadata, - ) - .await; - - if let Ok(ref output) = tool_result - && !output.is_empty() - { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolResult { - name: tc.name.clone(), - preview: output.clone(), - }, - &message.metadata, - ) - .await; - } - - // Record result 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 &tool_result { - Ok(output) => { - turn.record_tool_result(serde_json::json!(output)); - } - Err(e) => { - turn.record_tool_error(e.to_string()); - } - } - } - } - - // If tool_auth returned awaiting_token, enter auth mode - // and short-circuit: return the instructions directly so - // the LLM doesn't get a chance to hallucinate tool calls. - if let Some((ext_name, instructions)) = - detect_auth_awaiting(&tc.name, &tool_result) - { - let auth_data = parse_auth_result(&tool_result); - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.enter_auth_mode(ext_name.clone()); - } - } - 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(AgenticLoopResult::Response(instructions)); - } - - // Add tool result to context for next LLM call - let result_content = match tool_result { - Ok(output) => { - // Sanitize output before showing to LLM - 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, - result_content, - )); - } - } - } - } - } - - /// Execute a tool for chat (without full job context). - async fn execute_chat_tool( - &self, - tool_name: &str, - params: &serde_json::Value, - job_ctx: &JobContext, - ) -> Result { - let tool = - self.tools() - .get(tool_name) - .await - .ok_or_else(|| crate::error::ToolError::NotFound { - name: tool_name.to_string(), - })?; - - // Validate tool parameters - let validation = self.safety().validator().validate_tool_params(params); - if !validation.is_valid { - let details = validation - .errors - .iter() - .map(|e| format!("{}: {}", e.field, e.message)) - .collect::>() - .join("; "); - return Err(crate::error::ToolError::InvalidParameters { - name: tool_name.to_string(), - reason: format!("Invalid tool parameters: {}", details), - } - .into()); - } - - tracing::debug!( - tool = %tool_name, - params = %params, - "Tool call started" - ); - - // Execute with per-tool timeout - let timeout = tool.execution_timeout(); - let start = std::time::Instant::now(); - let result = tokio::time::timeout(timeout, async { - tool.execute(params.clone(), job_ctx).await - }) - .await; - let elapsed = start.elapsed(); - - match &result { - Ok(Ok(output)) => { - let result_str = serde_json::to_string(&output.result) - .unwrap_or_else(|_| "".to_string()); - tracing::debug!( - tool = %tool_name, - elapsed_ms = elapsed.as_millis() as u64, - result = %result_str, - "Tool call succeeded" - ); - } - Ok(Err(e)) => { - tracing::debug!( - tool = %tool_name, - elapsed_ms = elapsed.as_millis() as u64, - error = %e, - "Tool call failed" - ); - } - Err(_) => { - tracing::debug!( - tool = %tool_name, - elapsed_ms = elapsed.as_millis() as u64, - timeout_secs = timeout.as_secs(), - "Tool call timed out" - ); - } - } - - let result = result - .map_err(|_| crate::error::ToolError::Timeout { - name: tool_name.to_string(), - timeout, - })? - .map_err(|e| crate::error::ToolError::ExecutionFailed { - name: tool_name.to_string(), - reason: e.to_string(), - })?; - - // Convert result to string - serde_json::to_string_pretty(&result.result).map_err(|e| { - crate::error::ToolError::ExecutionFailed { - name: tool_name.to_string(), - reason: format!("Failed to serialize result: {}", e), - } - .into() - }) - } - - /// Handle job-related intents without turn tracking. - async fn handle_job_or_command( - &self, - intent: MessageIntent, - message: &IncomingMessage, - ) -> Result { - // Send thinking status for non-trivial operations - if let MessageIntent::CreateJob { .. } = &intent { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Thinking("Processing...".into()), - &message.metadata, - ) - .await; - } - - let response = match intent { - MessageIntent::CreateJob { - title, - description, - category, - } => { - self.handle_create_job(&message.user_id, title, description, category) - .await? - } - MessageIntent::CheckJobStatus { job_id } => { - self.handle_check_status(&message.user_id, job_id).await? - } - MessageIntent::CancelJob { job_id } => { - self.handle_cancel_job(&message.user_id, &job_id).await? - } - MessageIntent::ListJobs { filter } => { - self.handle_list_jobs(&message.user_id, filter).await? - } - MessageIntent::HelpJob { job_id } => { - self.handle_help_job(&message.user_id, &job_id).await? - } - MessageIntent::Command { command, args } => { - match self.handle_command(&command, &args).await? { - Some(s) => s, - None => return Ok(SubmissionResult::Ok { message: None }), // Shutdown signal - } - } - _ => "Unknown intent".to_string(), - }; - Ok(SubmissionResult::response(response)) - } - - async fn process_undo( - &self, - session: Arc>, - thread_id: Uuid, - ) -> Result { - // Lock session first, then undo manager -- consistent with process_user_input - // to avoid potential deadlocks. - let mut sess = session.lock().await; - let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; - let mut mgr = undo_mgr.lock().await; - - if !mgr.can_undo() { - return Ok(SubmissionResult::ok_with_message("Nothing to undo.")); - } - - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - // Save current state to redo, get previous checkpoint - let current_messages = thread.messages(); - let current_turn = thread.turn_number(); - - if let Some(checkpoint) = mgr.undo(current_turn, current_messages) { - let turn_number = checkpoint.turn_number; - let undo_count = mgr.undo_count(); - // Restore thread from checkpoint - thread.restore_from_messages(checkpoint.messages); - Ok(SubmissionResult::ok_with_message(format!( - "Undone to turn {}. {} undo(s) remaining.", - turn_number, undo_count - ))) - } else { - Ok(SubmissionResult::error("Undo failed.")) - } - } - - async fn process_redo( - &self, - session: Arc>, - thread_id: Uuid, - ) -> Result { - // Lock session first, then undo manager -- consistent with process_user_input - // to avoid potential deadlocks. - let mut sess = session.lock().await; - let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; - let mut mgr = undo_mgr.lock().await; - - if !mgr.can_redo() { - return Ok(SubmissionResult::ok_with_message("Nothing to redo.")); - } - - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - let current_messages = thread.messages(); - let current_turn = thread.turn_number(); - - if let Some(checkpoint) = mgr.redo(current_turn, current_messages) { - thread.restore_from_messages(checkpoint.messages); - Ok(SubmissionResult::ok_with_message(format!( - "Redone to turn {}.", - checkpoint.turn_number - ))) - } else { - Ok(SubmissionResult::error("Redo failed.")) - } - } - - async fn process_interrupt( - &self, - session: Arc>, - thread_id: Uuid, - ) -> Result { - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - match thread.state { - ThreadState::Processing | ThreadState::AwaitingApproval => { - thread.interrupt(); - Ok(SubmissionResult::ok_with_message("Interrupted.")) - } - _ => Ok(SubmissionResult::ok_with_message("Nothing to interrupt.")), - } - } - - async fn process_compact( - &self, - session: Arc>, - thread_id: Uuid, - ) -> Result { - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - let messages = thread.messages(); - let usage = self.context_monitor.usage_percent(&messages); - let strategy = self - .context_monitor - .suggest_compaction(&messages) - .unwrap_or( - crate::agent::context_monitor::CompactionStrategy::Summarize { keep_recent: 5 }, - ); - - let compactor = ContextCompactor::new(self.llm().clone()); - match compactor - .compact(thread, strategy, self.workspace().map(|w| w.as_ref())) - .await - { - Ok(result) => { - let mut msg = format!( - "Compacted: {} turns removed, {} → {} tokens (was {:.1}% full)", - result.turns_removed, result.tokens_before, result.tokens_after, usage - ); - if result.summary_written { - msg.push_str(", summary saved to workspace"); - } - Ok(SubmissionResult::ok_with_message(msg)) - } - Err(e) => Ok(SubmissionResult::error(format!("Compaction failed: {}", e))), - } - } - - async fn process_clear( - &self, - session: Arc>, - thread_id: Uuid, - ) -> Result { - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - thread.turns.clear(); - thread.state = ThreadState::Idle; - - // Clear undo history too - let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; - undo_mgr.lock().await.clear(); - - Ok(SubmissionResult::ok_with_message("Thread cleared.")) - } - - /// Process an approval or rejection of a pending tool execution. - async fn process_approval( - &self, - message: &IncomingMessage, - session: Arc>, - thread_id: Uuid, - request_id: Option, - approved: bool, - always: bool, - ) -> Result { - // Get thread state and pending approval - let (_thread_state, pending) = { - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - if thread.state != ThreadState::AwaitingApproval { - return Ok(SubmissionResult::error("No pending approval request.")); - } - - let pending = thread.take_pending_approval(); - (thread.state, pending) - }; - - let pending = match pending { - Some(p) => p, - None => return Ok(SubmissionResult::error("No pending approval request.")), - }; - - // Verify request ID if provided - if let Some(req_id) = request_id - && req_id != pending.request_id - { - // Put it back and return error - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.await_approval(pending); - } - return Ok(SubmissionResult::error( - "Request ID mismatch. Use the correct request ID.", - )); - } - - if approved { - // If always, add to auto-approved set - if always { - let mut sess = session.lock().await; - sess.auto_approve_tool(&pending.tool_name); - tracing::info!( - "Auto-approved tool '{}' for session {}", - pending.tool_name, - sess.id - ); - } - - // Reset thread state to processing - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.state = ThreadState::Processing; - } - } - - // Execute the approved tool and continue the loop - let job_ctx = - JobContext::with_user(&message.user_id, "chat", "Interactive chat session"); - - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolStarted { - name: pending.tool_name.clone(), - }, - &message.metadata, - ) - .await; - - let tool_result = self - .execute_chat_tool(&pending.tool_name, &pending.parameters, &job_ctx) - .await; - - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolCompleted { - name: pending.tool_name.clone(), - success: tool_result.is_ok(), - }, - &message.metadata, - ) - .await; - - if let Ok(ref output) = tool_result - && !output.is_empty() - { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolResult { - name: pending.tool_name.clone(), - preview: output.clone(), - }, - &message.metadata, - ) - .await; - } - - // Build context including the tool result - let mut context_messages = pending.context_messages; - - // Record result 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 &tool_result { - Ok(output) => { - turn.record_tool_result(serde_json::json!(output)); - } - Err(e) => { - turn.record_tool_error(e.to_string()); - } - } - } - } - - // If tool_auth returned awaiting_token, enter auth mode and - // return instructions directly (skip agentic loop continuation). - if let Some((ext_name, instructions)) = - detect_auth_awaiting(&pending.tool_name, &tool_result) - { - let auth_data = parse_auth_result(&tool_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)); - } - - // Add tool result to context - let result_content = match tool_result { - Ok(output) => { - let sanitized = self - .safety() - .sanitize_tool_output(&pending.tool_name, &output); - self.safety().wrap_for_llm( - &pending.tool_name, - &sanitized.content, - sanitized.was_modified, - ) - } - Err(e) => format!("Error: {}", e), - }; - - context_messages.push(ChatMessage::tool_result( - &pending.tool_call_id, - &pending.tool_name, - result_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) - .await; - - // Handle the result - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - - match result { - Ok(AgenticLoopResult::Response(response)) => { - thread.complete_turn(&response); - self.persist_response_chain(thread); - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Status("Done".into()), - &message.metadata, - ) - .await; - Ok(SubmissionResult::response(response)) - } - Ok(AgenticLoopResult::NeedApproval { - pending: new_pending, - }) => { - 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(); - thread.await_approval(new_pending); - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Status("Awaiting approval".into()), - &message.metadata, - ) - .await; - Ok(SubmissionResult::NeedApproval { - request_id, - tool_name, - description, - parameters, - }) - } - Err(e) => { - thread.fail_turn(e.to_string()); - Ok(SubmissionResult::error(e.to_string())) - } - } - } else { - // Rejected - clear approval and return to idle - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.clear_pending_approval(); - } - } - - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Status("Rejected".into()), - &message.metadata, - ) - .await; - - Ok(SubmissionResult::response(format!( - "Tool '{}' was rejected. The agent will not execute this tool.\n\n\ - You can continue the conversation or try a different approach.", - pending.tool_name - ))) - } - } - - /// Handle an auth token submitted while the thread is in auth mode. - /// - /// The token goes directly to the extension manager's credential store, - /// completely bypassing logging, turn creation, history, and compaction. - async fn process_auth_token( - &self, - message: &IncomingMessage, - pending: &crate::agent::session::PendingAuth, - token: &str, - session: Arc>, - thread_id: Uuid, - ) -> Result, Error> { - let token = token.trim(); - - // Clear auth mode regardless of outcome - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.pending_auth = None; - } - } - - let ext_mgr = match self.deps.extension_manager.as_ref() { - Some(mgr) => mgr, - None => return Ok(Some("Extension manager not available.".to_string())), - }; - - match ext_mgr.auth(&pending.extension_name, Some(token)).await { - Ok(result) if result.status == "authenticated" => { - tracing::info!( - "Extension '{}' authenticated via auth mode", - pending.extension_name - ); - - // Auto-activate so tools are available immediately after auth - match ext_mgr.activate(&pending.extension_name).await { - Ok(activate_result) => { - let tool_count = activate_result.tools_loaded.len(); - let tool_list = if activate_result.tools_loaded.is_empty() { - String::new() - } else { - format!("\n\nTools: {}", activate_result.tools_loaded.join(", ")) - }; - let msg = format!( - "{} authenticated and activated ({} tools loaded).{}", - pending.extension_name, tool_count, tool_list - ); - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::AuthCompleted { - extension_name: pending.extension_name.clone(), - success: true, - message: msg.clone(), - }, - &message.metadata, - ) - .await; - Ok(Some(msg)) - } - Err(e) => { - tracing::warn!( - "Extension '{}' authenticated but activation failed: {}", - pending.extension_name, - e - ); - let msg = format!( - "{} authenticated successfully, but activation failed: {}. \ - Try activating manually.", - pending.extension_name, e - ); - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::AuthCompleted { - extension_name: pending.extension_name.clone(), - success: true, - message: msg.clone(), - }, - &message.metadata, - ) - .await; - Ok(Some(msg)) - } - } - } - Ok(result) => { - // Invalid token, re-enter auth mode - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.enter_auth_mode(pending.extension_name.clone()); - } - } - let msg = result - .instructions - .clone() - .unwrap_or_else(|| "Invalid token. Please try again.".to_string()); - // Re-emit AuthRequired so web UI re-shows the card - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::AuthRequired { - extension_name: pending.extension_name.clone(), - instructions: Some(msg.clone()), - auth_url: result.auth_url, - setup_url: result.setup_url, - }, - &message.metadata, - ) - .await; - Ok(Some(msg)) - } - Err(e) => { - let msg = format!( - "Authentication failed for {}: {}", - pending.extension_name, e - ); - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::AuthCompleted { - extension_name: pending.extension_name.clone(), - success: false, - message: msg.clone(), - }, - &message.metadata, - ) - .await; - Ok(Some(msg)) - } - } - } - - async fn process_new_thread( - &self, - message: &IncomingMessage, - ) -> Result { - let session = self - .session_manager - .get_or_create_session(&message.user_id) - .await; - let mut sess = session.lock().await; - let thread = sess.create_thread(); - let thread_id = thread.id; - Ok(SubmissionResult::ok_with_message(format!( - "New thread: {}", - thread_id - ))) - } - - async fn process_switch_thread( - &self, - message: &IncomingMessage, - target_thread_id: Uuid, - ) -> Result { - let session = self - .session_manager - .get_or_create_session(&message.user_id) - .await; - let mut sess = session.lock().await; - - if sess.switch_thread(target_thread_id) { - Ok(SubmissionResult::ok_with_message(format!( - "Switched to thread {}", - target_thread_id - ))) - } else { - Ok(SubmissionResult::error("Thread not found.")) - } - } - - async fn process_resume( - &self, - session: Arc>, - thread_id: Uuid, - checkpoint_id: Uuid, - ) -> Result { - let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; - let mut mgr = undo_mgr.lock().await; - - if let Some(checkpoint) = mgr.restore(checkpoint_id) { - let mut sess = session.lock().await; - let thread = sess - .threads - .get_mut(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - thread.restore_from_messages(checkpoint.messages); - Ok(SubmissionResult::ok_with_message(format!( - "Resumed from checkpoint: {}", - checkpoint.description - ))) - } else { - Ok(SubmissionResult::error("Checkpoint not found.")) - } - } - - async fn handle_create_job( - &self, - user_id: &str, - title: String, - description: String, - category: Option, - ) -> Result { - // Create job context - let job_id = self - .context_manager - .create_job_for_user(user_id, &title, &description) - .await?; - - // Update category if provided - if let Some(cat) = category { - self.context_manager - .update_context(job_id, |ctx| { - ctx.category = Some(cat); - }) - .await?; - } - - // Persist new job to database (fire-and-forget) - if let Some(store) = self.store() - && let Ok(ctx) = self.context_manager.get_context(job_id).await - { - let store = store.clone(); - tokio::spawn(async move { - if let Err(e) = store.save_job(&ctx).await { - tracing::warn!("Failed to persist new job {}: {}", job_id, e); - } - }); - } - - // Schedule for execution - self.scheduler.schedule(job_id).await?; - - Ok(format!( - "Created job: {}\nID: {}\n\nThe job has been scheduled and is now running.", - title, job_id - )) - } - - async fn handle_check_status( - &self, - user_id: &str, - job_id: Option, - ) -> Result { - match job_id { - Some(id) => { - let uuid = Uuid::parse_str(&id) - .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; - - let ctx = self.context_manager.get_context(uuid).await?; - if ctx.user_id != user_id { - return Err(crate::error::JobError::NotFound { id: uuid }.into()); - } - - Ok(format!( - "Job: {}\nStatus: {:?}\nCreated: {}\nStarted: {}\nActual cost: {}", - ctx.title, - ctx.state, - ctx.created_at.format("%Y-%m-%d %H:%M:%S"), - ctx.started_at - .map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string()) - .unwrap_or_else(|| "Not started".to_string()), - ctx.actual_cost - )) - } - None => { - // Show summary of all jobs - let summary = self.context_manager.summary_for(user_id).await; - Ok(format!( - "Jobs summary:\n Total: {}\n In Progress: {}\n Completed: {}\n Failed: {}\n Stuck: {}", - summary.total, - summary.in_progress, - summary.completed, - summary.failed, - summary.stuck - )) - } - } - } - - async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result { - let uuid = Uuid::parse_str(job_id) - .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; - - let ctx = self.context_manager.get_context(uuid).await?; - if ctx.user_id != user_id { - return Err(crate::error::JobError::NotFound { id: uuid }.into()); - } - - self.scheduler.stop(uuid).await?; - - Ok(format!("Job {} has been cancelled.", job_id)) - } - - async fn handle_list_jobs( - &self, - user_id: &str, - _filter: Option, - ) -> Result { - let jobs = self.context_manager.all_jobs_for(user_id).await; - - if jobs.is_empty() { - return Ok("No jobs found.".to_string()); - } - - let mut output = String::from("Jobs:\n"); - for job_id in jobs { - if let Ok(ctx) = self.context_manager.get_context(job_id).await - && ctx.user_id == user_id - { - output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state)); - } - } - - Ok(output) - } - - async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result { - let uuid = Uuid::parse_str(job_id) - .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; - - let ctx = self.context_manager.get_context(uuid).await?; - if ctx.user_id != user_id { - return Err(crate::error::JobError::NotFound { id: uuid }.into()); - } - - if ctx.state == crate::context::JobState::Stuck { - // Attempt recovery - self.context_manager - .update_context(uuid, |ctx| ctx.attempt_recovery()) - .await? - .map_err(|s| crate::error::JobError::ContextError { - id: uuid, - reason: s, - })?; - - // Reschedule - self.scheduler.schedule(uuid).await?; - - Ok(format!( - "Job {} was stuck. Attempting recovery (attempt #{}).", - job_id, - ctx.repair_attempts + 1 - )) - } else { - Ok(format!( - "Job {} is not stuck (current state: {:?}). No help needed.", - job_id, ctx.state - )) - } - } - - /// Trigger a manual heartbeat check. - async fn process_heartbeat(&self) -> Result { - let Some(workspace) = self.workspace() else { - return Ok(SubmissionResult::error( - "Heartbeat requires a workspace (database must be connected).", - )); - }; - - let runner = crate::agent::HeartbeatRunner::new( - crate::agent::HeartbeatConfig::default(), - workspace.clone(), - self.llm().clone(), - ); - - match runner.check_heartbeat().await { - crate::agent::HeartbeatResult::Ok => Ok(SubmissionResult::ok_with_message( - "Heartbeat: all clear, nothing needs attention.", - )), - crate::agent::HeartbeatResult::NeedsAttention(msg) => Ok(SubmissionResult::response( - format!("Heartbeat findings:\n\n{}", msg), - )), - crate::agent::HeartbeatResult::Skipped => Ok(SubmissionResult::ok_with_message( - "Heartbeat skipped: no HEARTBEAT.md checklist found in workspace.", - )), - crate::agent::HeartbeatResult::Failed(err) => Ok(SubmissionResult::error(format!( - "Heartbeat failed: {}", - err - ))), - } - } - - /// Summarize the current thread's conversation. - async fn process_summarize( - &self, - session: Arc>, - thread_id: Uuid, - ) -> Result { - let messages = { - let sess = session.lock().await; - let thread = sess - .threads - .get(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - thread.messages() - }; - - if messages.is_empty() { - return Ok(SubmissionResult::ok_with_message( - "Nothing to summarize (empty thread).", - )); - } - - // Build a summary prompt with the conversation - let mut context = Vec::new(); - context.push(ChatMessage::system( - "Summarize the conversation so far in 3-5 concise bullet points. \ - Focus on decisions made, actions taken, and key outcomes. \ - Be brief and factual.", - )); - // Include the conversation messages (truncate to last 20 to avoid context overflow) - let start = if messages.len() > 20 { - messages.len() - 20 - } else { - 0 - }; - context.extend_from_slice(&messages[start..]); - context.push(ChatMessage::user("Summarize this conversation.")); - - let request = crate::llm::CompletionRequest::new(context) - .with_max_tokens(512) - .with_temperature(0.3); - - match self.llm().complete(request).await { - Ok(response) => Ok(SubmissionResult::response(format!( - "Thread Summary:\n\n{}", - response.content.trim() - ))), - Err(e) => Ok(SubmissionResult::error(format!("Summarize failed: {}", e))), - } - } - - /// Suggest next steps based on the current thread. - async fn process_suggest( - &self, - session: Arc>, - thread_id: Uuid, - ) -> Result { - let messages = { - let sess = session.lock().await; - let thread = sess - .threads - .get(&thread_id) - .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - thread.messages() - }; - - if messages.is_empty() { - return Ok(SubmissionResult::ok_with_message( - "Nothing to suggest from (empty thread).", - )); - } - - let mut context = Vec::new(); - context.push(ChatMessage::system( - "Based on the conversation so far, suggest 2-4 concrete next steps the user could take. \ - Be actionable and specific. Format as a numbered list.", - )); - let start = if messages.len() > 20 { - messages.len() - 20 - } else { - 0 - }; - context.extend_from_slice(&messages[start..]); - context.push(ChatMessage::user("What should I do next?")); - - let request = crate::llm::CompletionRequest::new(context) - .with_max_tokens(512) - .with_temperature(0.5); - - match self.llm().complete(request).await { - Ok(response) => Ok(SubmissionResult::response(format!( - "Suggested Next Steps:\n\n{}", - response.content.trim() - ))), - Err(e) => Ok(SubmissionResult::error(format!("Suggest failed: {}", e))), - } - } - - /// Handle system commands that bypass thread-state checks entirely. - async fn handle_system_command( - &self, - command: &str, - args: &[String], - ) -> Result { - match command { - "help" => Ok(SubmissionResult::response(concat!( - "System:\n", - " /help Show this help\n", - " /model [name] Show or switch the active model\n", - " /version Show version info\n", - " /tools List available tools\n", - " /debug Toggle debug mode\n", - " /ping Connectivity check\n", - "\n", - "Jobs:\n", - " /job Create a new job\n", - " /status [id] Check job status\n", - " /cancel Cancel a job\n", - " /list List all jobs\n", - "\n", - "Session:\n", - " /undo Undo last turn\n", - " /redo Redo undone turn\n", - " /compact Compress context window\n", - " /clear Clear current thread\n", - " /interrupt Stop current operation\n", - " /new New conversation thread\n", - " /thread Switch to thread\n", - " /resume Resume from checkpoint\n", - "\n", - "Agent:\n", - " /heartbeat Run heartbeat check\n", - " /summarize Summarize current thread\n", - " /suggest Suggest next steps\n", - "\n", - " /quit Exit", - ))), - - "ping" => Ok(SubmissionResult::response("pong!")), - - "version" => Ok(SubmissionResult::response(format!( - "{} v{}", - env!("CARGO_PKG_NAME"), - env!("CARGO_PKG_VERSION") - ))), - - "tools" => { - let tools = self.tools().list().await; - Ok(SubmissionResult::response(format!( - "Available tools: {}", - tools.join(", ") - ))) - } - - "debug" => { - // Debug toggle is handled client-side in the REPL. - // For non-REPL channels, just acknowledge. - Ok(SubmissionResult::ok_with_message( - "Debug toggle is handled by your client.", - )) - } - - "model" => { - if args.is_empty() { - // Show current model - let name = self.llm().active_model_name(); - Ok(SubmissionResult::response(format!( - "Active model: {}", - name - ))) - } else { - let requested = &args[0]; - - // Validate the model exists - match self.llm().list_models().await { - Ok(models) if !models.is_empty() => { - if !models.iter().any(|m| m == requested) { - return Ok(SubmissionResult::error(format!( - "Unknown model: {}. Available models:\n {}", - requested, - models.join("\n ") - ))); - } - } - Ok(_) => { - // Empty model list, can't validate but try anyway - } - Err(e) => { - tracing::warn!("Could not fetch model list for validation: {}", e); - // Proceed anyway, the provider will error on the next call if invalid - } - } - - match self.llm().set_model(requested) { - Ok(()) => Ok(SubmissionResult::response(format!( - "Switched model to: {}", - requested - ))), - Err(e) => Ok(SubmissionResult::error(format!( - "Failed to switch model: {}", - e - ))), - } - } - } - - _ => Ok(SubmissionResult::error(format!( - "Unknown command: {}. Try /help", - command - ))), - } - } - - /// Handle legacy command routing from the Router (job commands that go through - /// process_user_input -> router -> handle_job_or_command -> here). - async fn handle_command( - &self, - command: &str, - args: &[String], - ) -> Result, Error> { - // System commands are now handled directly via Submission::SystemCommand, - // but the router may still send us unknown /commands. - match self.handle_system_command(command, args).await? { - SubmissionResult::Response { content } => Ok(Some(content)), - SubmissionResult::Ok { message } => Ok(message), - SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))), - _ => Ok(None), - } - } -} - -/// Parsed auth result fields for emitting StatusUpdate::AuthRequired. -struct ParsedAuthData { - auth_url: Option, - setup_url: Option, -} - -/// Extract auth_url and setup_url from a tool_auth result JSON string. -fn parse_auth_result(result: &Result) -> ParsedAuthData { - let parsed = result - .as_ref() - .ok() - .and_then(|s| serde_json::from_str::(s).ok()); - ParsedAuthData { - auth_url: parsed - .as_ref() - .and_then(|v| v.get("auth_url")) - .and_then(|v| v.as_str()) - .map(|s| s.to_string()), - setup_url: parsed - .as_ref() - .and_then(|v| v.get("setup_url")) - .and_then(|v| v.as_str()) - .map(|s| s.to_string()), - } -} - -/// Check if a tool_auth result indicates the extension is awaiting a token. -/// -/// Returns `Some((extension_name, instructions))` if the tool result contains -/// `awaiting_token: true`, meaning the thread should enter auth mode. -fn detect_auth_awaiting( - tool_name: &str, - result: &Result, -) -> Option<(String, String)> { - if tool_name != "tool_auth" && tool_name != "tool_activate" { - return None; - } - let output = result.as_ref().ok()?; - let parsed: serde_json::Value = serde_json::from_str(output).ok()?; - if parsed.get("awaiting_token") != Some(&serde_json::Value::Bool(true)) { - return None; - } - let name = parsed.get("name")?.as_str()?.to_string(); - let instructions = parsed - .get("instructions") - .and_then(|v| v.as_str()) - .unwrap_or("Please provide your API token/key.") - .to_string(); - Some((name, instructions)) } #[cfg(test)] mod tests { - use crate::error::Error; - - use super::detect_auth_awaiting; - - #[test] - fn test_detect_auth_awaiting_positive() { - let result: Result = Ok(serde_json::json!({ - "name": "telegram", - "kind": "WasmTool", - "awaiting_token": true, - "status": "awaiting_token", - "instructions": "Please provide your Telegram Bot API token." - }) - .to_string()); - - let detected = detect_auth_awaiting("tool_auth", &result); - assert!(detected.is_some()); - let (name, instructions) = detected.unwrap(); - assert_eq!(name, "telegram"); - assert!(instructions.contains("Telegram Bot API")); - } - - #[test] - fn test_detect_auth_awaiting_not_awaiting() { - let result: Result = Ok(serde_json::json!({ - "name": "telegram", - "kind": "WasmTool", - "awaiting_token": false, - "status": "authenticated" - }) - .to_string()); - - assert!(detect_auth_awaiting("tool_auth", &result).is_none()); - } - - #[test] - fn test_detect_auth_awaiting_wrong_tool() { - let result: Result = Ok(serde_json::json!({ - "name": "telegram", - "awaiting_token": true, - }) - .to_string()); - - assert!(detect_auth_awaiting("tool_list", &result).is_none()); - } - - #[test] - fn test_detect_auth_awaiting_error_result() { - let result: Result = - Err(crate::error::ToolError::NotFound { name: "x".into() }.into()); - assert!(detect_auth_awaiting("tool_auth", &result).is_none()); - } - - #[test] - fn test_detect_auth_awaiting_default_instructions() { - let result: Result = Ok(serde_json::json!({ - "name": "custom_tool", - "awaiting_token": true, - "status": "awaiting_token" - }) - .to_string()); - - let (_, instructions) = detect_auth_awaiting("tool_auth", &result).unwrap(); - assert_eq!(instructions, "Please provide your API token/key."); - } - - #[test] - fn test_detect_auth_awaiting_tool_activate() { - let result: Result = Ok(serde_json::json!({ - "name": "slack", - "kind": "McpServer", - "awaiting_token": true, - "status": "awaiting_token", - "instructions": "Provide your Slack Bot token." - }) - .to_string()); - - let detected = detect_auth_awaiting("tool_activate", &result); - assert!(detected.is_some()); - let (name, instructions) = detected.unwrap(); - assert_eq!(name, "slack"); - assert!(instructions.contains("Slack Bot")); - } - - #[test] - fn test_detect_auth_awaiting_tool_activate_not_awaiting() { - let result: Result = Ok(serde_json::json!({ - "name": "slack", - "tools_loaded": ["slack_post_message"], - "message": "Activated" - }) - .to_string()); - - assert!(detect_auth_awaiting("tool_activate", &result).is_none()); - } - - // --- truncate_for_preview tests --- - use super::truncate_for_preview; #[test] diff --git a/src/agent/commands.rs b/src/agent/commands.rs new file mode 100644 index 00000000..5525a528 --- /dev/null +++ b/src/agent/commands.rs @@ -0,0 +1,484 @@ +//! System commands and job handlers for the agent. +//! +//! Extracted from `agent_loop.rs` to isolate the /help, /model, /status, +//! and other command processing from the core agent loop. + +use std::sync::Arc; + +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::agent::session::Session; +use crate::agent::submission::SubmissionResult; +use crate::agent::{Agent, MessageIntent}; +use crate::channels::{IncomingMessage, StatusUpdate}; +use crate::error::Error; +use crate::llm::ChatMessage; + +impl Agent { + /// Handle job-related intents without turn tracking. + pub(super) async fn handle_job_or_command( + &self, + intent: MessageIntent, + message: &IncomingMessage, + ) -> Result { + // Send thinking status for non-trivial operations + if let MessageIntent::CreateJob { .. } = &intent { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Thinking("Processing...".into()), + &message.metadata, + ) + .await; + } + + let response = match intent { + MessageIntent::CreateJob { + title, + description, + category, + } => { + self.handle_create_job(&message.user_id, title, description, category) + .await? + } + MessageIntent::CheckJobStatus { job_id } => { + self.handle_check_status(&message.user_id, job_id).await? + } + MessageIntent::CancelJob { job_id } => { + self.handle_cancel_job(&message.user_id, &job_id).await? + } + MessageIntent::ListJobs { filter } => { + self.handle_list_jobs(&message.user_id, filter).await? + } + MessageIntent::HelpJob { job_id } => { + self.handle_help_job(&message.user_id, &job_id).await? + } + MessageIntent::Command { command, args } => { + match self.handle_command(&command, &args).await? { + Some(s) => s, + None => return Ok(SubmissionResult::Ok { message: None }), // Shutdown signal + } + } + _ => "Unknown intent".to_string(), + }; + Ok(SubmissionResult::response(response)) + } + + async fn handle_create_job( + &self, + user_id: &str, + title: String, + description: String, + category: Option, + ) -> Result { + // Create job context + let job_id = self + .context_manager + .create_job_for_user(user_id, &title, &description) + .await?; + + // Update category if provided + if let Some(cat) = category { + self.context_manager + .update_context(job_id, |ctx| { + ctx.category = Some(cat); + }) + .await?; + } + + // Persist new job to database (fire-and-forget) + if let Some(store) = self.store() + && let Ok(ctx) = self.context_manager.get_context(job_id).await + { + let store = store.clone(); + tokio::spawn(async move { + if let Err(e) = store.save_job(&ctx).await { + tracing::warn!("Failed to persist new job {}: {}", job_id, e); + } + }); + } + + // Schedule for execution + self.scheduler.schedule(job_id).await?; + + Ok(format!( + "Created job: {}\nID: {}\n\nThe job has been scheduled and is now running.", + title, job_id + )) + } + + async fn handle_check_status( + &self, + user_id: &str, + job_id: Option, + ) -> Result { + match job_id { + Some(id) => { + let uuid = Uuid::parse_str(&id) + .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; + + let ctx = self.context_manager.get_context(uuid).await?; + if ctx.user_id != user_id { + return Err(crate::error::JobError::NotFound { id: uuid }.into()); + } + + Ok(format!( + "Job: {}\nStatus: {:?}\nCreated: {}\nStarted: {}\nActual cost: {}", + ctx.title, + ctx.state, + ctx.created_at.format("%Y-%m-%d %H:%M:%S"), + ctx.started_at + .map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string()) + .unwrap_or_else(|| "Not started".to_string()), + ctx.actual_cost + )) + } + None => { + // Show summary of all jobs + let summary = self.context_manager.summary_for(user_id).await; + Ok(format!( + "Jobs summary:\n Total: {}\n In Progress: {}\n Completed: {}\n Failed: {}\n Stuck: {}", + summary.total, + summary.in_progress, + summary.completed, + summary.failed, + summary.stuck + )) + } + } + } + + async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result { + let uuid = Uuid::parse_str(job_id) + .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; + + let ctx = self.context_manager.get_context(uuid).await?; + if ctx.user_id != user_id { + return Err(crate::error::JobError::NotFound { id: uuid }.into()); + } + + self.scheduler.stop(uuid).await?; + + Ok(format!("Job {} has been cancelled.", job_id)) + } + + async fn handle_list_jobs( + &self, + user_id: &str, + _filter: Option, + ) -> Result { + let jobs = self.context_manager.all_jobs_for(user_id).await; + + if jobs.is_empty() { + return Ok("No jobs found.".to_string()); + } + + let mut output = String::from("Jobs:\n"); + for job_id in jobs { + if let Ok(ctx) = self.context_manager.get_context(job_id).await + && ctx.user_id == user_id + { + output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state)); + } + } + + Ok(output) + } + + async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result { + let uuid = Uuid::parse_str(job_id) + .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; + + let ctx = self.context_manager.get_context(uuid).await?; + if ctx.user_id != user_id { + return Err(crate::error::JobError::NotFound { id: uuid }.into()); + } + + if ctx.state == crate::context::JobState::Stuck { + // Attempt recovery + self.context_manager + .update_context(uuid, |ctx| ctx.attempt_recovery()) + .await? + .map_err(|s| crate::error::JobError::ContextError { + id: uuid, + reason: s, + })?; + + // Reschedule + self.scheduler.schedule(uuid).await?; + + Ok(format!( + "Job {} was stuck. Attempting recovery (attempt #{}).", + job_id, + ctx.repair_attempts + 1 + )) + } else { + Ok(format!( + "Job {} is not stuck (current state: {:?}). No help needed.", + job_id, ctx.state + )) + } + } + + /// Trigger a manual heartbeat check. + pub(super) async fn process_heartbeat(&self) -> Result { + let Some(workspace) = self.workspace() else { + return Ok(SubmissionResult::error( + "Heartbeat requires a workspace (database must be connected).", + )); + }; + + let runner = crate::agent::HeartbeatRunner::new( + crate::agent::HeartbeatConfig::default(), + workspace.clone(), + self.llm().clone(), + ); + + match runner.check_heartbeat().await { + crate::agent::HeartbeatResult::Ok => Ok(SubmissionResult::ok_with_message( + "Heartbeat: all clear, nothing needs attention.", + )), + crate::agent::HeartbeatResult::NeedsAttention(msg) => Ok(SubmissionResult::response( + format!("Heartbeat findings:\n\n{}", msg), + )), + crate::agent::HeartbeatResult::Skipped => Ok(SubmissionResult::ok_with_message( + "Heartbeat skipped: no HEARTBEAT.md checklist found in workspace.", + )), + crate::agent::HeartbeatResult::Failed(err) => Ok(SubmissionResult::error(format!( + "Heartbeat failed: {}", + err + ))), + } + } + + /// Summarize the current thread's conversation. + pub(super) async fn process_summarize( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let messages = { + let sess = session.lock().await; + let thread = sess + .threads + .get(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.messages() + }; + + if messages.is_empty() { + return Ok(SubmissionResult::ok_with_message( + "Nothing to summarize (empty thread).", + )); + } + + // Build a summary prompt with the conversation + let mut context = Vec::new(); + context.push(ChatMessage::system( + "Summarize the conversation so far in 3-5 concise bullet points. \ + Focus on decisions made, actions taken, and key outcomes. \ + Be brief and factual.", + )); + // Include the conversation messages (truncate to last 20 to avoid context overflow) + let start = if messages.len() > 20 { + messages.len() - 20 + } else { + 0 + }; + context.extend_from_slice(&messages[start..]); + context.push(ChatMessage::user("Summarize this conversation.")); + + let request = crate::llm::CompletionRequest::new(context) + .with_max_tokens(512) + .with_temperature(0.3); + + match self.llm().complete(request).await { + Ok(response) => Ok(SubmissionResult::response(format!( + "Thread Summary:\n\n{}", + response.content.trim() + ))), + Err(e) => Ok(SubmissionResult::error(format!("Summarize failed: {}", e))), + } + } + + /// Suggest next steps based on the current thread. + pub(super) async fn process_suggest( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let messages = { + let sess = session.lock().await; + let thread = sess + .threads + .get(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.messages() + }; + + if messages.is_empty() { + return Ok(SubmissionResult::ok_with_message( + "Nothing to suggest from (empty thread).", + )); + } + + let mut context = Vec::new(); + context.push(ChatMessage::system( + "Based on the conversation so far, suggest 2-4 concrete next steps the user could take. \ + Be actionable and specific. Format as a numbered list.", + )); + let start = if messages.len() > 20 { + messages.len() - 20 + } else { + 0 + }; + context.extend_from_slice(&messages[start..]); + context.push(ChatMessage::user("What should I do next?")); + + let request = crate::llm::CompletionRequest::new(context) + .with_max_tokens(512) + .with_temperature(0.5); + + match self.llm().complete(request).await { + Ok(response) => Ok(SubmissionResult::response(format!( + "Suggested Next Steps:\n\n{}", + response.content.trim() + ))), + Err(e) => Ok(SubmissionResult::error(format!("Suggest failed: {}", e))), + } + } + + /// Handle system commands that bypass thread-state checks entirely. + pub(super) async fn handle_system_command( + &self, + command: &str, + args: &[String], + ) -> Result { + match command { + "help" => Ok(SubmissionResult::response(concat!( + "System:\n", + " /help Show this help\n", + " /model [name] Show or switch the active model\n", + " /version Show version info\n", + " /tools List available tools\n", + " /debug Toggle debug mode\n", + " /ping Connectivity check\n", + "\n", + "Jobs:\n", + " /job Create a new job\n", + " /status [id] Check job status\n", + " /cancel Cancel a job\n", + " /list List all jobs\n", + "\n", + "Session:\n", + " /undo Undo last turn\n", + " /redo Redo undone turn\n", + " /compact Compress context window\n", + " /clear Clear current thread\n", + " /interrupt Stop current operation\n", + " /new New conversation thread\n", + " /thread Switch to thread\n", + " /resume Resume from checkpoint\n", + "\n", + "Agent:\n", + " /heartbeat Run heartbeat check\n", + " /summarize Summarize current thread\n", + " /suggest Suggest next steps\n", + "\n", + " /quit Exit", + ))), + + "ping" => Ok(SubmissionResult::response("pong!")), + + "version" => Ok(SubmissionResult::response(format!( + "{} v{}", + env!("CARGO_PKG_NAME"), + env!("CARGO_PKG_VERSION") + ))), + + "tools" => { + let tools = self.tools().list().await; + Ok(SubmissionResult::response(format!( + "Available tools: {}", + tools.join(", ") + ))) + } + + "debug" => { + // Debug toggle is handled client-side in the REPL. + // For non-REPL channels, just acknowledge. + Ok(SubmissionResult::ok_with_message( + "Debug toggle is handled by your client.", + )) + } + + "model" => { + if args.is_empty() { + // Show current model + let name = self.llm().active_model_name(); + Ok(SubmissionResult::response(format!( + "Active model: {}", + name + ))) + } else { + let requested = &args[0]; + + // Validate the model exists + match self.llm().list_models().await { + Ok(models) if !models.is_empty() => { + if !models.iter().any(|m| m == requested) { + return Ok(SubmissionResult::error(format!( + "Unknown model: {}. Available models:\n {}", + requested, + models.join("\n ") + ))); + } + } + Ok(_) => { + // Empty model list, can't validate but try anyway + } + Err(e) => { + tracing::warn!("Could not fetch model list for validation: {}", e); + // Proceed anyway, the provider will error on the next call if invalid + } + } + + match self.llm().set_model(requested) { + Ok(()) => Ok(SubmissionResult::response(format!( + "Switched model to: {}", + requested + ))), + Err(e) => Ok(SubmissionResult::error(format!( + "Failed to switch model: {}", + e + ))), + } + } + } + + _ => Ok(SubmissionResult::error(format!( + "Unknown command: {}. Try /help", + command + ))), + } + } + + /// Handle legacy command routing from the Router (job commands that go through + /// process_user_input -> router -> handle_job_or_command -> here). + pub(super) async fn handle_command( + &self, + command: &str, + args: &[String], + ) -> Result, Error> { + // System commands are now handled directly via Submission::SystemCommand, + // but the router may still send us unknown /commands. + match self.handle_system_command(command, args).await? { + SubmissionResult::Response { content } => Ok(Some(content)), + SubmissionResult::Ok { message } => Ok(message), + SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))), + _ => Ok(None), + } + } +} diff --git a/src/agent/cost_guard.rs b/src/agent/cost_guard.rs new file mode 100644 index 00000000..2ddeae58 --- /dev/null +++ b/src/agent/cost_guard.rs @@ -0,0 +1,339 @@ +//! Cost enforcement guardrails for the agent. +//! +//! Tracks LLM spending and action rates, enforcing configurable limits +//! to prevent runaway agents from burning through API credits. Especially +//! important for daemon/heartbeat modes where the agent acts autonomously. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Instant; + +use rust_decimal::Decimal; +use rust_decimal_macros::dec; +use tokio::sync::Mutex; + +use crate::llm::costs; + +/// Configuration for cost guardrails. +#[derive(Debug, Clone, Default)] +pub struct CostGuardConfig { + /// Maximum spend per day in cents (e.g. 10000 = $100). None = unlimited. + pub max_cost_per_day_cents: Option, + /// Maximum LLM calls per hour. None = unlimited. + pub max_actions_per_hour: Option, +} + +/// Error returned when a cost limit is exceeded. +#[derive(Debug, Clone)] +pub enum CostLimitExceeded { + /// Daily spending cap reached. + DailyBudget { spent_cents: u64, limit_cents: u64 }, + /// Hourly action rate limit reached. + HourlyRate { actions: u64, limit: u64 }, +} + +impl std::fmt::Display for CostLimitExceeded { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::DailyBudget { + spent_cents, + limit_cents, + } => write!( + f, + "Daily cost limit exceeded: spent ${:.2} of ${:.2} allowed", + *spent_cents as f64 / 100.0, + *limit_cents as f64 / 100.0 + ), + Self::HourlyRate { actions, limit } => write!( + f, + "Hourly action limit exceeded: {} actions of {} allowed per hour", + actions, limit + ), + } + } +} + +/// Tracks costs and action rates, enforcing configurable limits. +/// +/// Thread-safe; designed to be shared via `Arc`. +pub struct CostGuard { + config: CostGuardConfig, + + /// Running cost total for the current day (in USD, not cents). + daily_cost: Mutex, + + /// Sliding window of action timestamps for rate limiting. + action_window: Mutex>, + + /// Flag set when daily budget is exceeded to short-circuit checks. + budget_exceeded: AtomicBool, +} + +struct DailyCost { + total: Decimal, + /// Day boundary (midnight UTC) for resetting the counter. + reset_date: chrono::NaiveDate, +} + +impl CostGuard { + pub fn new(config: CostGuardConfig) -> Self { + Self { + config, + daily_cost: Mutex::new(DailyCost { + total: Decimal::ZERO, + reset_date: chrono::Utc::now().date_naive(), + }), + action_window: Mutex::new(VecDeque::new()), + budget_exceeded: AtomicBool::new(false), + } + } + + /// Check whether the next action is allowed under the configured limits. + /// + /// Call this BEFORE making an LLM call. Does NOT record the action yet, + /// call `record_action` after the action completes. + pub async fn check_allowed(&self) -> Result<(), CostLimitExceeded> { + // Fast path: if budget already blown, skip the lock + if self.budget_exceeded.load(Ordering::Relaxed) { + let daily = self.daily_cost.lock().await; + let spent_cents = to_cents(daily.total); + return Err(CostLimitExceeded::DailyBudget { + spent_cents, + limit_cents: self.config.max_cost_per_day_cents.unwrap_or(0), + }); + } + + // Check daily budget + if let Some(limit_cents) = self.config.max_cost_per_day_cents { + let daily = self.daily_cost.lock().await; + let spent_cents = to_cents(daily.total); + if spent_cents >= limit_cents { + self.budget_exceeded.store(true, Ordering::Relaxed); + return Err(CostLimitExceeded::DailyBudget { + spent_cents, + limit_cents, + }); + } + } + + // Check hourly rate + if let Some(limit) = self.config.max_actions_per_hour { + let mut window = self.action_window.lock().await; + let cutoff = Instant::now() - std::time::Duration::from_secs(3600); + // Drain expired entries + while window.front().is_some_and(|t| *t < cutoff) { + window.pop_front(); + } + let count = window.len() as u64; + if count >= limit { + return Err(CostLimitExceeded::HourlyRate { + actions: count, + limit, + }); + } + } + + Ok(()) + } + + /// Record a completed LLM action: its token costs and the action timestamp. + /// + /// Call this AFTER an LLM call completes so that costs are tracked. + pub async fn record_llm_call( + &self, + model: &str, + input_tokens: u32, + output_tokens: u32, + ) -> Decimal { + let (input_rate, output_rate) = + costs::model_cost(model).unwrap_or_else(costs::default_cost); + let cost = + input_rate * Decimal::from(input_tokens) + output_rate * Decimal::from(output_tokens); + + // Update daily cost (reset if new day) + { + let mut daily = self.daily_cost.lock().await; + let today = chrono::Utc::now().date_naive(); + if today != daily.reset_date { + daily.total = Decimal::ZERO; + daily.reset_date = today; + self.budget_exceeded.store(false, Ordering::Relaxed); + tracing::info!("Cost guard: daily counter reset for {}", today); + } + daily.total += cost; + + // Check if we just crossed the threshold + if let Some(limit_cents) = self.config.max_cost_per_day_cents { + let spent_cents = to_cents(daily.total); + if spent_cents >= limit_cents { + self.budget_exceeded.store(true, Ordering::Relaxed); + tracing::warn!( + "Daily cost limit reached: ${:.2} of ${:.2}", + daily.total, + Decimal::from(limit_cents) / dec!(100) + ); + } + // Warn at 80% threshold + let warn_threshold = limit_cents * 80 / 100; + if spent_cents >= warn_threshold && spent_cents < limit_cents { + tracing::warn!( + "Approaching daily cost limit: ${:.2} of ${:.2} ({}%)", + daily.total, + Decimal::from(limit_cents) / dec!(100), + spent_cents * 100 / limit_cents + ); + } + } + } + + // Record action in sliding window + { + let mut window = self.action_window.lock().await; + window.push_back(Instant::now()); + } + + cost + } + + /// Current daily spend in USD (as Decimal). + pub async fn daily_spend(&self) -> Decimal { + let daily = self.daily_cost.lock().await; + let today = chrono::Utc::now().date_naive(); + if today != daily.reset_date { + Decimal::ZERO + } else { + daily.total + } + } + + /// Number of actions in the current hourly window. + pub async fn actions_this_hour(&self) -> u64 { + let mut window = self.action_window.lock().await; + let cutoff = Instant::now() - std::time::Duration::from_secs(3600); + while window.front().is_some_and(|t| *t < cutoff) { + window.pop_front(); + } + window.len() as u64 + } +} + +/// Convert a Decimal USD amount to whole cents (truncated). +fn to_cents(usd: Decimal) -> u64 { + let cents = (usd * dec!(100)).trunc(); + cents.to_string().parse::().unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_unlimited_allows_everything() { + let guard = CostGuard::new(CostGuardConfig::default()); + + // No limits set, should always be allowed + assert!(guard.check_allowed().await.is_ok()); + + // Record a big call, still allowed + guard.record_llm_call("gpt-4o", 100_000, 100_000).await; + assert!(guard.check_allowed().await.is_ok()); + } + + #[tokio::test] + async fn test_daily_budget_enforcement() { + let guard = CostGuard::new(CostGuardConfig { + max_cost_per_day_cents: Some(1), // $0.01 limit + max_actions_per_hour: None, + }); + + // First call allowed + assert!(guard.check_allowed().await.is_ok()); + + // Record a call that costs more than $0.01 + // gpt-4o: input=$0.0000025/tok, output=$0.00001/tok + // 10000 input + 10000 output = $0.025 + $0.10 = $0.125 + guard.record_llm_call("gpt-4o", 10_000, 10_000).await; + + // Now should be blocked + let result = guard.check_allowed().await; + assert!(result.is_err()); + match result.unwrap_err() { + CostLimitExceeded::DailyBudget { limit_cents, .. } => { + assert_eq!(limit_cents, 1); + } + other => panic!("Expected DailyBudget, got {:?}", other), + } + } + + #[tokio::test] + async fn test_hourly_rate_enforcement() { + let guard = CostGuard::new(CostGuardConfig { + max_cost_per_day_cents: None, + max_actions_per_hour: Some(3), + }); + + // First 3 actions allowed + for _ in 0..3 { + assert!(guard.check_allowed().await.is_ok()); + guard.record_llm_call("gpt-4o", 10, 10).await; + } + + // 4th should be blocked + let result = guard.check_allowed().await; + assert!(result.is_err()); + match result.unwrap_err() { + CostLimitExceeded::HourlyRate { actions, limit } => { + assert_eq!(actions, 3); + assert_eq!(limit, 3); + } + other => panic!("Expected HourlyRate, got {:?}", other), + } + } + + #[tokio::test] + async fn test_daily_spend_tracking() { + let guard = CostGuard::new(CostGuardConfig::default()); + + assert_eq!(guard.daily_spend().await, Decimal::ZERO); + + let cost = guard.record_llm_call("gpt-4o", 1000, 500).await; + assert!(cost > Decimal::ZERO); + assert_eq!(guard.daily_spend().await, cost); + } + + #[tokio::test] + async fn test_actions_this_hour() { + let guard = CostGuard::new(CostGuardConfig::default()); + + assert_eq!(guard.actions_this_hour().await, 0); + + guard.record_llm_call("gpt-4o", 10, 10).await; + guard.record_llm_call("gpt-4o", 10, 10).await; + + assert_eq!(guard.actions_this_hour().await, 2); + } + + #[test] + fn test_to_cents() { + assert_eq!(to_cents(dec!(1.50)), 150); + assert_eq!(to_cents(dec!(0.01)), 1); + assert_eq!(to_cents(Decimal::ZERO), 0); + } + + #[test] + fn test_cost_limit_display() { + let budget = CostLimitExceeded::DailyBudget { + spent_cents: 1050, + limit_cents: 1000, + }; + assert!(budget.to_string().contains("$10.50")); + assert!(budget.to_string().contains("$10.00")); + + let rate = CostLimitExceeded::HourlyRate { + actions: 101, + limit: 100, + }; + assert!(rate.to_string().contains("101 actions")); + assert!(rate.to_string().contains("100 allowed")); + } +} diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs new file mode 100644 index 00000000..49cca0af --- /dev/null +++ b/src/agent/dispatcher.rs @@ -0,0 +1,635 @@ +//! Tool dispatch logic for the agent. +//! +//! Extracted from `agent_loop.rs` to keep the core agentic tool execution +//! loop (LLM call -> tool calls -> repeat) in its own focused module. + +use std::sync::Arc; + +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::agent::Agent; +use crate::agent::session::{PendingApproval, Session, ThreadState}; +use crate::channels::{IncomingMessage, StatusUpdate}; +use crate::context::JobContext; +use crate::error::Error; +use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult}; + +/// Result of the agentic loop execution. +pub(super) enum AgenticLoopResult { + /// Completed with a response. + Response(String), + /// A tool requires approval before continuing. + NeedApproval { + /// The pending approval request to store. + pending: PendingApproval, + }, +} + +impl Agent { + /// Run the agentic loop: call LLM, execute tools, repeat until text response. + /// + /// Returns `AgenticLoopResult::Response` on completion, or + /// `AgenticLoopResult::NeedApproval` if a tool requires user approval. + /// + /// When `resume_after_tool` is true the loop already knows a tool was + /// executed earlier in this turn (e.g. an approved tool), so it won't + /// force the LLM to use tools if it responds with text. + pub(super) async fn run_agentic_loop( + &self, + message: &IncomingMessage, + session: Arc>, + thread_id: Uuid, + initial_messages: Vec, + resume_after_tool: bool, + ) -> Result { + // Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.) + let system_prompt = if let Some(ws) = self.workspace() { + match ws.system_prompt().await { + Ok(prompt) if !prompt.is_empty() => Some(prompt), + Ok(_) => None, + Err(e) => { + tracing::debug!("Could not load workspace system prompt: {}", e); + None + } + } + } else { + None + }; + + let mut reasoning = Reasoning::new(self.llm().clone(), self.safety().clone()); + if let Some(prompt) = system_prompt { + reasoning = reasoning.with_system_prompt(prompt); + } + + // Build context with messages that we'll mutate during the loop + let mut context_messages = initial_messages; + + // Create a JobContext for tool execution (chat doesn't have a real job) + let job_ctx = JobContext::with_user(&message.user_id, "chat", "Interactive chat session"); + + const MAX_TOOL_ITERATIONS: usize = 10; + let mut iteration = 0; + let mut tools_executed = resume_after_tool; + + loop { + iteration += 1; + if iteration > MAX_TOOL_ITERATIONS { + return Err(crate::error::LlmError::InvalidResponse { + provider: "agent".to_string(), + reason: format!("Exceeded maximum tool iterations ({})", MAX_TOOL_ITERATIONS), + } + .into()); + } + + // Check if interrupted + { + let sess = session.lock().await; + if let Some(thread) = sess.threads.get(&thread_id) + && thread.state == ThreadState::Interrupted + { + return Err(crate::error::JobError::ContextError { + id: thread_id, + reason: "Interrupted".to_string(), + } + .into()); + } + } + + // Enforce cost guardrails before the LLM call + if let Err(limit) = self.cost_guard().check_allowed().await { + return Err(crate::error::LlmError::InvalidResponse { + provider: "agent".to_string(), + reason: limit.to_string(), + } + .into()); + } + + // Refresh tool definitions each iteration so newly built tools become visible + let tool_defs = self.tools().tool_definitions().await; + + // Call LLM with current context + let context = ReasoningContext::new() + .with_messages(context_messages.clone()) + .with_tools(tool_defs) + .with_metadata({ + let mut m = std::collections::HashMap::new(); + m.insert("thread_id".to_string(), thread_id.to_string()); + m + }); + + let output = reasoning.respond_with_tools(&context).await?; + + // Record cost and track token usage + let model_name = self.llm().active_model_name(); + let call_cost = self + .cost_guard() + .record_llm_call( + &model_name, + output.usage.input_tokens, + output.usage.output_tokens, + ) + .await; + tracing::debug!( + "LLM call used {} input + {} output tokens (${:.6})", + output.usage.input_tokens, + output.usage.output_tokens, + call_cost, + ); + + match output.result { + RespondResult::Text(text) => { + // If no tools have been executed yet, prompt the LLM to use tools + // This handles the case where the model explains what it will do + // instead of actually calling tools + if !tools_executed && iteration < 3 { + tracing::debug!( + "No tools executed yet (iteration {}), prompting for tool use", + iteration + ); + context_messages.push(ChatMessage::assistant(&text)); + context_messages.push(ChatMessage::user( + "Please proceed and use the available tools to complete this task.", + )); + continue; + } + + // Tools have been executed or we've tried multiple times, return response + return Ok(AgenticLoopResult::Response(text)); + } + RespondResult::ToolCalls { + tool_calls, + content, + } => { + tools_executed = true; + + // Add the assistant message with tool_calls to context. + // OpenAI protocol requires this before tool-result messages. + context_messages.push(ChatMessage::assistant_with_tool_calls( + content, + tool_calls.clone(), + )); + + // Execute tools and add results to context + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Thinking(format!( + "Executing {} tool(s)...", + tool_calls.len() + )), + &message.metadata, + ) + .await; + + // Record tool calls in the thread + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) + && let Some(turn) = thread.last_turn_mut() + { + for tc in &tool_calls { + turn.record_tool_call(&tc.name, tc.arguments.clone()); + } + } + } + + // Execute each tool (with approval checking and hook interception) + for mut tc in tool_calls { + // Check if tool requires approval + if let Some(tool) = self.tools().get(&tc.name).await + && tool.requires_approval() + { + // Check if auto-approved for this session + let mut is_auto_approved = { + let sess = session.lock().await; + sess.is_tool_auto_approved(&tc.name) + }; + + // Override auto-approval for destructive parameters + // (e.g. `rm -rf`, `git push --force` in shell commands). + if is_auto_approved && tool.requires_approval_for(&tc.arguments) { + tracing::info!( + tool = %tc.name, + "Parameters require explicit approval despite auto-approve" + ); + is_auto_approved = false; + } + + if !is_auto_approved { + // Need approval - store pending request and return + let 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(), + }; + + return Ok(AgenticLoopResult::NeedApproval { pending }); + } + } + + // Hook: BeforeToolCall — allow hooks to modify or reject tool calls + { + let event = crate::hooks::HookEvent::ToolCall { + tool_name: tc.name.clone(), + parameters: tc.arguments.clone(), + user_id: message.user_id.clone(), + context: "chat".to_string(), + }; + match self.hooks().run(&event).await { + Err(crate::hooks::HookError::Rejected { reason }) => { + context_messages.push(ChatMessage::tool_result( + &tc.id, + &tc.name, + format!("Tool call rejected by hook: {}", reason), + )); + continue; + } + Err(err) => { + context_messages.push(ChatMessage::tool_result( + &tc.id, + &tc.name, + format!("Tool call blocked by hook policy: {}", err), + )); + continue; + } + Ok(crate::hooks::HookOutcome::Continue { + modified: Some(new_params), + }) => match serde_json::from_str(&new_params) { + Ok(parsed) => tc.arguments = parsed, + Err(e) => { + tracing::warn!( + tool = %tc.name, + "Hook returned non-JSON modification for ToolCall, ignoring: {}", + e + ); + } + }, + _ => {} // Continue, fail-open errors already logged + } + } + + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolStarted { + name: tc.name.clone(), + }, + &message.metadata, + ) + .await; + + let tool_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: tool_result.is_ok(), + }, + &message.metadata, + ) + .await; + + if let Ok(ref output) = tool_result + && !output.is_empty() + { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolResult { + name: tc.name.clone(), + preview: output.clone(), + }, + &message.metadata, + ) + .await; + } + + // Record result 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 &tool_result { + Ok(output) => { + turn.record_tool_result(serde_json::json!(output)); + } + Err(e) => { + turn.record_tool_error(e.to_string()); + } + } + } + } + + // If tool_auth returned awaiting_token, enter auth mode + // and short-circuit: return the instructions directly so + // the LLM doesn't get a chance to hallucinate tool calls. + if let Some((ext_name, instructions)) = + detect_auth_awaiting(&tc.name, &tool_result) + { + let auth_data = parse_auth_result(&tool_result); + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.enter_auth_mode(ext_name.clone()); + } + } + 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(AgenticLoopResult::Response(instructions)); + } + + // Add tool result to context for next LLM call + let result_content = match tool_result { + Ok(output) => { + // Sanitize output before showing to LLM + 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, + result_content, + )); + } + } + } + } + } + + /// Execute a tool for chat (without full job context). + pub(super) async fn execute_chat_tool( + &self, + tool_name: &str, + params: &serde_json::Value, + job_ctx: &JobContext, + ) -> Result { + let tool = + self.tools() + .get(tool_name) + .await + .ok_or_else(|| crate::error::ToolError::NotFound { + name: tool_name.to_string(), + })?; + + // Validate tool parameters + let validation = self.safety().validator().validate_tool_params(params); + if !validation.is_valid { + let details = validation + .errors + .iter() + .map(|e| format!("{}: {}", e.field, e.message)) + .collect::>() + .join("; "); + return Err(crate::error::ToolError::InvalidParameters { + name: tool_name.to_string(), + reason: format!("Invalid tool parameters: {}", details), + } + .into()); + } + + tracing::debug!( + tool = %tool_name, + params = %params, + "Tool call started" + ); + + // Execute with per-tool timeout + let timeout = tool.execution_timeout(); + let start = std::time::Instant::now(); + let result = tokio::time::timeout(timeout, async { + tool.execute(params.clone(), job_ctx).await + }) + .await; + let elapsed = start.elapsed(); + + match &result { + Ok(Ok(output)) => { + let result_str = serde_json::to_string(&output.result) + .unwrap_or_else(|_| "".to_string()); + tracing::debug!( + tool = %tool_name, + elapsed_ms = elapsed.as_millis() as u64, + result = %result_str, + "Tool call succeeded" + ); + } + Ok(Err(e)) => { + tracing::debug!( + tool = %tool_name, + elapsed_ms = elapsed.as_millis() as u64, + error = %e, + "Tool call failed" + ); + } + Err(_) => { + tracing::debug!( + tool = %tool_name, + elapsed_ms = elapsed.as_millis() as u64, + timeout_secs = timeout.as_secs(), + "Tool call timed out" + ); + } + } + + let result = result + .map_err(|_| crate::error::ToolError::Timeout { + name: tool_name.to_string(), + timeout, + })? + .map_err(|e| crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: e.to_string(), + })?; + + // Convert result to string + serde_json::to_string_pretty(&result.result).map_err(|e| { + crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: format!("Failed to serialize result: {}", e), + } + .into() + }) + } +} + +/// Parsed auth result fields for emitting StatusUpdate::AuthRequired. +pub(super) struct ParsedAuthData { + pub(super) auth_url: Option, + pub(super) setup_url: Option, +} + +/// Extract auth_url and setup_url from a tool_auth result JSON string. +pub(super) fn parse_auth_result(result: &Result) -> ParsedAuthData { + let parsed = result + .as_ref() + .ok() + .and_then(|s| serde_json::from_str::(s).ok()); + ParsedAuthData { + auth_url: parsed + .as_ref() + .and_then(|v| v.get("auth_url")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()), + setup_url: parsed + .as_ref() + .and_then(|v| v.get("setup_url")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()), + } +} + +/// Check if a tool_auth result indicates the extension is awaiting a token. +/// +/// Returns `Some((extension_name, instructions))` if the tool result contains +/// `awaiting_token: true`, meaning the thread should enter auth mode. +pub(super) fn detect_auth_awaiting( + tool_name: &str, + result: &Result, +) -> Option<(String, String)> { + if tool_name != "tool_auth" && tool_name != "tool_activate" { + return None; + } + let output = result.as_ref().ok()?; + let parsed: serde_json::Value = serde_json::from_str(output).ok()?; + if parsed.get("awaiting_token") != Some(&serde_json::Value::Bool(true)) { + return None; + } + let name = parsed.get("name")?.as_str()?.to_string(); + let instructions = parsed + .get("instructions") + .and_then(|v| v.as_str()) + .unwrap_or("Please provide your API token/key.") + .to_string(); + Some((name, instructions)) +} + +#[cfg(test)] +mod tests { + use crate::error::Error; + + use super::detect_auth_awaiting; + + #[test] + fn test_detect_auth_awaiting_positive() { + let result: Result = Ok(serde_json::json!({ + "name": "telegram", + "kind": "WasmTool", + "awaiting_token": true, + "status": "awaiting_token", + "instructions": "Please provide your Telegram Bot API token." + }) + .to_string()); + + let detected = detect_auth_awaiting("tool_auth", &result); + assert!(detected.is_some()); + let (name, instructions) = detected.unwrap(); + assert_eq!(name, "telegram"); + assert!(instructions.contains("Telegram Bot API")); + } + + #[test] + fn test_detect_auth_awaiting_not_awaiting() { + let result: Result = Ok(serde_json::json!({ + "name": "telegram", + "kind": "WasmTool", + "awaiting_token": false, + "status": "authenticated" + }) + .to_string()); + + assert!(detect_auth_awaiting("tool_auth", &result).is_none()); + } + + #[test] + fn test_detect_auth_awaiting_wrong_tool() { + let result: Result = Ok(serde_json::json!({ + "name": "telegram", + "awaiting_token": true, + }) + .to_string()); + + assert!(detect_auth_awaiting("tool_list", &result).is_none()); + } + + #[test] + fn test_detect_auth_awaiting_error_result() { + let result: Result = + Err(crate::error::ToolError::NotFound { name: "x".into() }.into()); + assert!(detect_auth_awaiting("tool_auth", &result).is_none()); + } + + #[test] + fn test_detect_auth_awaiting_default_instructions() { + let result: Result = Ok(serde_json::json!({ + "name": "custom_tool", + "awaiting_token": true, + "status": "awaiting_token" + }) + .to_string()); + + let (_, instructions) = detect_auth_awaiting("tool_auth", &result).unwrap(); + assert_eq!(instructions, "Please provide your API token/key."); + } + + #[test] + fn test_detect_auth_awaiting_tool_activate() { + let result: Result = Ok(serde_json::json!({ + "name": "slack", + "kind": "McpServer", + "awaiting_token": true, + "status": "awaiting_token", + "instructions": "Provide your Slack Bot token." + }) + .to_string()); + + let detected = detect_auth_awaiting("tool_activate", &result); + assert!(detected.is_some()); + let (name, instructions) = detected.unwrap(); + assert_eq!(name, "slack"); + assert!(instructions.contains("Slack Bot")); + } + + #[test] + fn test_detect_auth_awaiting_tool_activate_not_awaiting() { + let result: Result = Ok(serde_json::json!({ + "name": "slack", + "tools_loaded": ["slack_post_message"], + "message": "Activated" + }) + .to_string()); + + assert!(detect_auth_awaiting("tool_activate", &result).is_none()); + } +} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index a667b17d..5e1bf64c 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -11,8 +11,11 @@ //! - Context compaction for long conversations mod agent_loop; +mod commands; pub mod compaction; pub mod context_monitor; +pub mod cost_guard; +mod dispatcher; mod heartbeat; mod router; pub mod routine; @@ -23,6 +26,7 @@ pub mod session; mod session_manager; pub mod submission; pub mod task; +mod thread_ops; pub mod undo; pub mod worker; diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs new file mode 100644 index 00000000..195591ba --- /dev/null +++ b/src/agent/thread_ops.rs @@ -0,0 +1,1063 @@ +//! Thread and session operations for the agent. +//! +//! Extracted from `agent_loop.rs` to isolate thread management (user input +//! processing, undo/redo, approval, auth, persistence) from the core loop. + +use std::sync::Arc; + +use tokio::sync::Mutex; +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::submission::SubmissionResult; +use crate::channels::{IncomingMessage, StatusUpdate}; +use crate::context::JobContext; +use crate::error::Error; +use crate::llm::ChatMessage; + +impl Agent { + /// Hydrate a historical thread from DB into memory if not already present. + /// + /// Called before `resolve_thread` so that the session manager finds the + /// thread on lookup instead of creating a new one. + /// + /// Creates an in-memory thread with the exact UUID the frontend sent, + /// even when the conversation has zero messages (e.g. a brand-new + /// assistant thread). Without this, `resolve_thread` would mint a + /// fresh UUID and all messages would land in the wrong conversation. + pub(super) async fn maybe_hydrate_thread( + &self, + message: &IncomingMessage, + external_thread_id: &str, + ) { + // Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs) + let thread_uuid = match Uuid::parse_str(external_thread_id) { + Ok(id) => id, + Err(_) => return, + }; + + // Check if already in memory + let session = self + .session_manager + .get_or_create_session(&message.user_id) + .await; + { + let sess = session.lock().await; + if sess.threads.contains_key(&thread_uuid) { + return; + } + } + + // Load history from DB (may be empty for a newly created thread). + let mut chat_messages: Vec = Vec::new(); + let msg_count; + + if let Some(store) = self.store() { + let db_messages = store + .list_conversation_messages(thread_uuid) + .await + .unwrap_or_default(); + msg_count = db_messages.len(); + chat_messages = db_messages + .iter() + .filter_map(|m| match m.role.as_str() { + "user" => Some(ChatMessage::user(&m.content)), + "assistant" => Some(ChatMessage::assistant(&m.content)), + _ => None, + }) + .collect(); + } else { + msg_count = 0; + } + + // Create thread with the historical ID and restore messages + let session_id = { + let sess = session.lock().await; + sess.id + }; + + let mut thread = crate::agent::session::Thread::with_id(thread_uuid, session_id); + if !chat_messages.is_empty() { + thread.restore_from_messages(chat_messages); + } + + // Restore response chain from conversation metadata + if let Some(store) = self.store() + && let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await + && let Some(rid) = metadata + .get("last_response_id") + .and_then(|v| v.as_str()) + .map(String::from) + { + thread.last_response_id = Some(rid.clone()); + self.llm() + .seed_response_chain(&thread_uuid.to_string(), rid); + tracing::debug!("Restored response chain for thread {}", thread_uuid); + } + + // Insert into session and register with session manager + { + let mut sess = session.lock().await; + sess.threads.insert(thread_uuid, thread); + sess.active_thread = Some(thread_uuid); + sess.last_active_at = chrono::Utc::now(); + } + + self.session_manager + .register_thread( + &message.user_id, + &message.channel, + thread_uuid, + Arc::clone(&session), + ) + .await; + + tracing::debug!( + "Hydrated thread {} from DB ({} messages)", + thread_uuid, + msg_count + ); + } + + pub(super) async fn process_user_input( + &self, + message: &IncomingMessage, + session: Arc>, + thread_id: Uuid, + content: &str, + ) -> Result { + // First check thread state without holding lock during I/O + let thread_state = { + let sess = session.lock().await; + let thread = sess + .threads + .get(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.state + }; + + // Check thread state + match thread_state { + ThreadState::Processing => { + return Ok(SubmissionResult::error( + "Turn in progress. Use /interrupt to cancel.", + )); + } + ThreadState::AwaitingApproval => { + return Ok(SubmissionResult::error( + "Waiting for approval. Use /interrupt to cancel.", + )); + } + ThreadState::Completed => { + return Ok(SubmissionResult::error( + "Thread completed. Use /thread new.", + )); + } + ThreadState::Idle | ThreadState::Interrupted => { + // Can proceed + } + } + + // Safety validation for user input + let validation = self.safety().validate_input(content); + if !validation.is_valid { + let details = validation + .errors + .iter() + .map(|e| format!("{}: {}", e.field, e.message)) + .collect::>() + .join("; "); + return Ok(SubmissionResult::error(format!( + "Input rejected by safety validation: {}", + details + ))); + } + + let violations = self.safety().check_policy(content); + if violations + .iter() + .any(|rule| rule.action == crate::safety::PolicyAction::Block) + { + return Ok(SubmissionResult::error("Input rejected by safety policy.")); + } + + // Handle explicit commands (starting with /) directly + // Everything else goes through the normal agentic loop with tools + let temp_message = IncomingMessage { + content: content.to_string(), + ..message.clone() + }; + + if let Some(intent) = self.router.route_command(&temp_message) { + // Explicit command like /status, /job, /list - handle directly + return self.handle_job_or_command(intent, message).await; + } + + // Natural language goes through the agentic loop + // Job tools (create_job, list_jobs, etc.) are in the tool registry + + // Auto-compact if needed BEFORE adding new turn + { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + let messages = thread.messages(); + if let Some(strategy) = self.context_monitor.suggest_compaction(&messages) { + let pct = self.context_monitor.usage_percent(&messages); + tracing::info!("Context at {:.1}% capacity, auto-compacting", pct); + + // Notify the user that compaction is happening + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status(format!( + "Context at {:.0}% capacity, compacting...", + pct + )), + &message.metadata, + ) + .await; + + let compactor = ContextCompactor::new(self.llm().clone()); + if let Err(e) = compactor + .compact(thread, strategy, self.workspace().map(|w| w.as_ref())) + .await + { + tracing::warn!("Auto-compaction failed: {}", e); + } + } + } + + // Create checkpoint before turn + let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; + { + let sess = session.lock().await; + let thread = sess + .threads + .get(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + let mut mgr = undo_mgr.lock().await; + mgr.checkpoint( + thread.turn_number(), + thread.messages(), + format!("Before turn {}", thread.turn_number()), + ); + } + + // Start the turn and get messages + let turn_messages = { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.start_turn(content); + thread.messages() + }; + + // Send thinking status + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Thinking("Processing...".into()), + &message.metadata, + ) + .await; + + // Run the agentic tool execution loop + let result = self + .run_agentic_loop(message, session.clone(), thread_id, turn_messages, false) + .await; + + // Re-acquire lock and check if interrupted + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + if thread.state == ThreadState::Interrupted { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Interrupted".into()), + &message.metadata, + ) + .await; + return Ok(SubmissionResult::Interrupted); + } + + // Complete, fail, or request approval + match result { + Ok(AgenticLoopResult::Response(response)) => { + // Hook: TransformResponse — allow hooks to modify or reject the final response + let response = { + let event = crate::hooks::HookEvent::ResponseTransform { + user_id: message.user_id.clone(), + thread_id: thread_id.to_string(), + response: response.clone(), + }; + match self.hooks().run(&event).await { + Err(crate::hooks::HookError::Rejected { reason }) => { + format!("[Response filtered: {}]", reason) + } + Err(err) => { + format!("[Response blocked by hook policy: {}]", err) + } + Ok(crate::hooks::HookOutcome::Continue { + modified: Some(new_response), + }) => new_response, + _ => response, // fail-open: use original + } + }; + + thread.complete_turn(&response); + self.persist_response_chain(thread); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Done".into()), + &message.metadata, + ) + .await; + + // Fire-and-forget: persist turn to DB + self.persist_turn(thread_id, &message.user_id, content, Some(&response)); + + Ok(SubmissionResult::response(response)) + } + Ok(AgenticLoopResult::NeedApproval { pending }) => { + // Store pending approval in thread and update state + let request_id = pending.request_id; + let tool_name = pending.tool_name.clone(); + let description = pending.description.clone(); + let parameters = pending.parameters.clone(); + thread.await_approval(pending); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Awaiting approval".into()), + &message.metadata, + ) + .await; + Ok(SubmissionResult::NeedApproval { + request_id, + tool_name, + description, + parameters, + }) + } + Err(e) => { + thread.fail_turn(e.to_string()); + + // Persist the user message even on failure + self.persist_turn(thread_id, &message.user_id, content, None); + + Ok(SubmissionResult::error(e.to_string())) + } + } + } + + /// Fire-and-forget: persist a turn (user message + optional assistant response) to the DB. + pub(super) fn persist_turn( + &self, + thread_id: Uuid, + user_id: &str, + user_input: &str, + response: Option<&str>, + ) { + let store = match self.store() { + Some(s) => Arc::clone(s), + None => return, + }; + + let user_id = user_id.to_string(); + let user_input = user_input.to_string(); + let response = response.map(String::from); + + tokio::spawn(async move { + if let Err(e) = store + .ensure_conversation(thread_id, "gateway", &user_id, None) + .await + { + tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e); + return; + } + + if let Err(e) = store + .add_conversation_message(thread_id, "user", &user_input) + .await + { + tracing::warn!("Failed to persist user message: {}", e); + return; + } + + if let Some(ref resp) = response + && let Err(e) = store + .add_conversation_message(thread_id, "assistant", resp) + .await + { + tracing::warn!("Failed to persist assistant message: {}", e); + } + }); + } + + /// Sync the provider's response chain ID to the thread and DB metadata. + /// + /// Call after a successful agentic loop to persist the latest + /// `previous_response_id` so chaining survives restarts. + pub(super) fn persist_response_chain(&self, thread: &mut crate::agent::session::Thread) { + let tid = thread.id.to_string(); + let response_id = match self.llm().get_response_chain_id(&tid) { + Some(rid) => rid, + None => return, + }; + + // Update in-memory thread + thread.last_response_id = Some(response_id.clone()); + + // Fire-and-forget DB write + let store = match self.store() { + Some(s) => Arc::clone(s), + None => return, + }; + let thread_id = thread.id; + tokio::spawn(async move { + let val = serde_json::json!(response_id); + if let Err(e) = store + .update_conversation_metadata_field(thread_id, "last_response_id", &val) + .await + { + tracing::warn!( + "Failed to persist response chain for thread {}: {}", + thread_id, + e + ); + } + }); + } + + pub(super) async fn process_undo( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; + let mut mgr = undo_mgr.lock().await; + + if !mgr.can_undo() { + return Ok(SubmissionResult::ok_with_message("Nothing to undo.")); + } + + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + // Save current state to redo, get previous checkpoint + let current_messages = thread.messages(); + let current_turn = thread.turn_number(); + + if let Some(checkpoint) = mgr.undo(current_turn, current_messages) { + // Extract values before consuming the reference + let turn_number = checkpoint.turn_number; + let messages = checkpoint.messages.clone(); + let undo_count = mgr.undo_count(); + // Restore thread from checkpoint + thread.restore_from_messages(messages); + Ok(SubmissionResult::ok_with_message(format!( + "Undone to turn {}. {} undo(s) remaining.", + turn_number, undo_count + ))) + } else { + Ok(SubmissionResult::error("Undo failed.")) + } + } + + pub(super) async fn process_redo( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; + let mut mgr = undo_mgr.lock().await; + + if !mgr.can_redo() { + return Ok(SubmissionResult::ok_with_message("Nothing to redo.")); + } + + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + let current_messages = thread.messages(); + let current_turn = thread.turn_number(); + + if let Some(checkpoint) = mgr.redo(current_turn, current_messages) { + thread.restore_from_messages(checkpoint.messages); + Ok(SubmissionResult::ok_with_message(format!( + "Redone to turn {}.", + checkpoint.turn_number + ))) + } else { + Ok(SubmissionResult::error("Redo failed.")) + } + } + + pub(super) async fn process_interrupt( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + match thread.state { + ThreadState::Processing | ThreadState::AwaitingApproval => { + thread.interrupt(); + Ok(SubmissionResult::ok_with_message("Interrupted.")) + } + _ => Ok(SubmissionResult::ok_with_message("Nothing to interrupt.")), + } + } + + pub(super) async fn process_compact( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + let messages = thread.messages(); + let usage = self.context_monitor.usage_percent(&messages); + let strategy = self + .context_monitor + .suggest_compaction(&messages) + .unwrap_or( + crate::agent::context_monitor::CompactionStrategy::Summarize { keep_recent: 5 }, + ); + + let compactor = ContextCompactor::new(self.llm().clone()); + match compactor + .compact(thread, strategy, self.workspace().map(|w| w.as_ref())) + .await + { + Ok(result) => { + let mut msg = format!( + "Compacted: {} turns removed, {} → {} tokens (was {:.1}% full)", + result.turns_removed, result.tokens_before, result.tokens_after, usage + ); + if result.summary_written { + msg.push_str(", summary saved to workspace"); + } + Ok(SubmissionResult::ok_with_message(msg)) + } + Err(e) => Ok(SubmissionResult::error(format!("Compaction failed: {}", e))), + } + } + + pub(super) async fn process_clear( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.turns.clear(); + thread.state = ThreadState::Idle; + + // Clear undo history too + let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; + undo_mgr.lock().await.clear(); + + Ok(SubmissionResult::ok_with_message("Thread cleared.")) + } + + /// Process an approval or rejection of a pending tool execution. + pub(super) async fn process_approval( + &self, + message: &IncomingMessage, + session: Arc>, + thread_id: Uuid, + request_id: Option, + approved: bool, + always: bool, + ) -> Result { + // Get thread state and pending approval + let (_thread_state, pending) = { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + if thread.state != ThreadState::AwaitingApproval { + return Ok(SubmissionResult::error("No pending approval request.")); + } + + let pending = thread.take_pending_approval(); + (thread.state, pending) + }; + + let pending = match pending { + Some(p) => p, + None => return Ok(SubmissionResult::error("No pending approval request.")), + }; + + // Verify request ID if provided + if let Some(req_id) = request_id + && req_id != pending.request_id + { + // Put it back and return error + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.await_approval(pending); + } + return Ok(SubmissionResult::error( + "Request ID mismatch. Use the correct request ID.", + )); + } + + if approved { + // If always, add to auto-approved set + if always { + let mut sess = session.lock().await; + sess.auto_approve_tool(&pending.tool_name); + tracing::info!( + "Auto-approved tool '{}' for session {}", + pending.tool_name, + sess.id + ); + } + + // Reset thread state to processing + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.state = ThreadState::Processing; + } + } + + // Execute the approved tool and continue the loop + let job_ctx = + JobContext::with_user(&message.user_id, "chat", "Interactive chat session"); + + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolStarted { + name: pending.tool_name.clone(), + }, + &message.metadata, + ) + .await; + + let tool_result = self + .execute_chat_tool(&pending.tool_name, &pending.parameters, &job_ctx) + .await; + + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolCompleted { + name: pending.tool_name.clone(), + success: tool_result.is_ok(), + }, + &message.metadata, + ) + .await; + + if let Ok(ref output) = tool_result + && !output.is_empty() + { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ToolResult { + name: pending.tool_name.clone(), + preview: output.clone(), + }, + &message.metadata, + ) + .await; + } + + // Build context including the tool result + let mut context_messages = pending.context_messages; + + // Record result 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 &tool_result { + Ok(output) => { + turn.record_tool_result(serde_json::json!(output)); + } + Err(e) => { + turn.record_tool_error(e.to_string()); + } + } + } + } + + // If tool_auth returned awaiting_token, enter auth mode and + // return instructions directly (skip agentic loop continuation). + if let Some((ext_name, instructions)) = + detect_auth_awaiting(&pending.tool_name, &tool_result) + { + let auth_data = parse_auth_result(&tool_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)); + } + + // Add tool result to context + let result_content = match tool_result { + Ok(output) => { + let sanitized = self + .safety() + .sanitize_tool_output(&pending.tool_name, &output); + self.safety().wrap_for_llm( + &pending.tool_name, + &sanitized.content, + sanitized.was_modified, + ) + } + Err(e) => format!("Error: {}", e), + }; + + context_messages.push(ChatMessage::tool_result( + &pending.tool_call_id, + &pending.tool_name, + result_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) + .await; + + // Handle the result + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + + match result { + Ok(AgenticLoopResult::Response(response)) => { + thread.complete_turn(&response); + self.persist_response_chain(thread); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Done".into()), + &message.metadata, + ) + .await; + Ok(SubmissionResult::response(response)) + } + Ok(AgenticLoopResult::NeedApproval { + pending: new_pending, + }) => { + 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(); + thread.await_approval(new_pending); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Awaiting approval".into()), + &message.metadata, + ) + .await; + Ok(SubmissionResult::NeedApproval { + request_id, + tool_name, + description, + parameters, + }) + } + Err(e) => { + thread.fail_turn(e.to_string()); + Ok(SubmissionResult::error(e.to_string())) + } + } + } else { + // Rejected - clear approval and return to idle + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.clear_pending_approval(); + } + } + + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Rejected".into()), + &message.metadata, + ) + .await; + + Ok(SubmissionResult::response(format!( + "Tool '{}' was rejected. The agent will not execute this tool.\n\n\ + You can continue the conversation or try a different approach.", + pending.tool_name + ))) + } + } + + /// Handle an auth token submitted while the thread is in auth mode. + /// + /// The token goes directly to the extension manager's credential store, + /// completely bypassing logging, turn creation, history, and compaction. + pub(super) async fn process_auth_token( + &self, + message: &IncomingMessage, + pending: &crate::agent::session::PendingAuth, + token: &str, + session: Arc>, + thread_id: Uuid, + ) -> Result, Error> { + let token = token.trim(); + + // Clear auth mode regardless of outcome + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.pending_auth = None; + } + } + + let ext_mgr = match self.deps.extension_manager.as_ref() { + Some(mgr) => mgr, + None => return Ok(Some("Extension manager not available.".to_string())), + }; + + match ext_mgr.auth(&pending.extension_name, Some(token)).await { + Ok(result) if result.status == "authenticated" => { + tracing::info!( + "Extension '{}' authenticated via auth mode", + pending.extension_name + ); + + // Auto-activate so tools are available immediately after auth + match ext_mgr.activate(&pending.extension_name).await { + Ok(activate_result) => { + let tool_count = activate_result.tools_loaded.len(); + let tool_list = if activate_result.tools_loaded.is_empty() { + String::new() + } else { + format!("\n\nTools: {}", activate_result.tools_loaded.join(", ")) + }; + let msg = format!( + "{} authenticated and activated ({} tools loaded).{}", + pending.extension_name, tool_count, tool_list + ); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::AuthCompleted { + extension_name: pending.extension_name.clone(), + success: true, + message: msg.clone(), + }, + &message.metadata, + ) + .await; + Ok(Some(msg)) + } + Err(e) => { + tracing::warn!( + "Extension '{}' authenticated but activation failed: {}", + pending.extension_name, + e + ); + let msg = format!( + "{} authenticated successfully, but activation failed: {}. \ + Try activating manually.", + pending.extension_name, e + ); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::AuthCompleted { + extension_name: pending.extension_name.clone(), + success: true, + message: msg.clone(), + }, + &message.metadata, + ) + .await; + Ok(Some(msg)) + } + } + } + Ok(result) => { + // Invalid token, re-enter auth mode + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.enter_auth_mode(pending.extension_name.clone()); + } + } + let msg = result + .instructions + .clone() + .unwrap_or_else(|| "Invalid token. Please try again.".to_string()); + // Re-emit AuthRequired so web UI re-shows the card + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::AuthRequired { + extension_name: pending.extension_name.clone(), + instructions: Some(msg.clone()), + auth_url: result.auth_url, + setup_url: result.setup_url, + }, + &message.metadata, + ) + .await; + Ok(Some(msg)) + } + Err(e) => { + let msg = format!( + "Authentication failed for {}: {}", + pending.extension_name, e + ); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::AuthCompleted { + extension_name: pending.extension_name.clone(), + success: false, + message: msg.clone(), + }, + &message.metadata, + ) + .await; + Ok(Some(msg)) + } + } + } + + pub(super) async fn process_new_thread( + &self, + message: &IncomingMessage, + ) -> Result { + let session = self + .session_manager + .get_or_create_session(&message.user_id) + .await; + let mut sess = session.lock().await; + let thread = sess.create_thread(); + let thread_id = thread.id; + Ok(SubmissionResult::ok_with_message(format!( + "New thread: {}", + thread_id + ))) + } + + pub(super) async fn process_switch_thread( + &self, + message: &IncomingMessage, + target_thread_id: Uuid, + ) -> Result { + let session = self + .session_manager + .get_or_create_session(&message.user_id) + .await; + let mut sess = session.lock().await; + + if sess.switch_thread(target_thread_id) { + Ok(SubmissionResult::ok_with_message(format!( + "Switched to thread {}", + target_thread_id + ))) + } else { + Ok(SubmissionResult::error("Thread not found.")) + } + } + + pub(super) async fn process_resume( + &self, + session: Arc>, + thread_id: Uuid, + checkpoint_id: Uuid, + ) -> Result { + let undo_mgr = self.session_manager.get_undo_manager(thread_id).await; + let mut mgr = undo_mgr.lock().await; + + if let Some(checkpoint) = mgr.restore(checkpoint_id) { + let mut sess = session.lock().await; + let thread = sess + .threads + .get_mut(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.restore_from_messages(checkpoint.messages); + Ok(SubmissionResult::ok_with_message(format!( + "Resumed from checkpoint: {}", + checkpoint.description + ))) + } else { + Ok(SubmissionResult::error("Checkpoint not found.")) + } + } +} diff --git a/src/boot_screen.rs b/src/boot_screen.rs index 57b082eb..881c0f1a 100644 --- a/src/boot_screen.rs +++ b/src/boot_screen.rs @@ -23,6 +23,10 @@ pub struct BootInfo { pub claude_code_enabled: bool, pub routines_enabled: bool, pub channels: Vec, + /// Public URL from a managed tunnel (e.g., "https://abc.ngrok.io"). + pub tunnel_url: Option, + /// Provider name for the managed tunnel (e.g., "ngrok"). + pub tunnel_provider: Option, } /// Print the boot screen to stdout. @@ -116,6 +120,16 @@ pub fn print_boot_screen(info: &BootInfo) { println!(" {dim}gateway{reset} {yellow_underline}{url}{reset}"); } + // Tunnel URL + if let Some(ref url) = info.tunnel_url { + let provider_tag = info + .tunnel_provider + .as_deref() + .map(|p| format!(" {dim}({p}){reset}")) + .unwrap_or_default(); + println!(" {dim}tunnel{reset} {yellow_underline}{url}{reset}{provider_tag}"); + } + println!(); println!("{border}"); println!(); @@ -151,6 +165,8 @@ mod tests { "gateway".to_string(), "telegram".to_string(), ], + tunnel_url: Some("https://abc123.ngrok.io".to_string()), + tunnel_provider: Some("ngrok".to_string()), }; // Should not panic print_boot_screen(&info); @@ -176,6 +192,8 @@ mod tests { claude_code_enabled: false, routines_enabled: false, channels: vec![], + tunnel_url: None, + tunnel_provider: None, }; // Should not panic print_boot_screen(&info); @@ -201,6 +219,8 @@ mod tests { claude_code_enabled: false, routines_enabled: false, channels: vec!["repl".to_string()], + tunnel_url: None, + tunnel_provider: None, }; // Should not panic print_boot_screen(&info); diff --git a/src/cli/config.rs b/src/cli/config.rs index f91e9241..fc1312f6 100644 --- a/src/cli/config.rs +++ b/src/cli/config.rs @@ -11,6 +11,17 @@ use crate::settings::Settings; #[derive(Subcommand, Debug, Clone)] pub enum ConfigCommand { + /// Generate a default config.toml file + Init { + /// Output path (default: ~/.ironclaw/config.toml) + #[arg(short, long)] + output: Option, + + /// Overwrite existing file + #[arg(long)] + force: bool, + }, + /// List all settings and their current values List { /// Show only settings matching this prefix (e.g., "agent", "heartbeat") @@ -62,6 +73,7 @@ pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> { let db_ref = db.as_deref(); match cmd { + ConfigCommand::Init { output, force } => init_toml(db_ref, output, force).await, ConfigCommand::List { filter } => list_settings(db_ref, filter).await, ConfigCommand::Get { path } => get_setting(db_ref, &path).await, ConfigCommand::Set { path, value } => set_setting(db_ref, &path, &value).await, @@ -188,6 +200,36 @@ async fn reset_setting(store: Option<&dyn crate::db::Database>, path: &str) -> a Ok(()) } +/// Generate a default TOML config file. +async fn init_toml( + store: Option<&dyn crate::db::Database>, + output: Option, + force: bool, +) -> anyhow::Result<()> { + let path = output.unwrap_or_else(Settings::default_toml_path); + + if path.exists() && !force { + anyhow::bail!( + "Config file already exists: {}\nUse --force to overwrite.", + path.display() + ); + } + + // Start from current settings (DB or defaults) so the generated file + // reflects the user's existing configuration. + let settings = load_settings(store).await; + + settings + .save_toml(&path) + .map_err(|e| anyhow::anyhow!("{}", e))?; + + println!("Config file written to {}", path.display()); + println!(); + println!("Edit the file to customize settings."); + println!("Priority: env var > config.toml > database > defaults"); + Ok(()) +} + /// Show the settings storage info. fn show_path(has_db: bool) -> anyhow::Result<()> { if has_db { @@ -200,6 +242,18 @@ fn show_path(has_db: bool) -> anyhow::Result<()> { crate::bootstrap::ironclaw_env_path().display() ); + let toml_path = Settings::default_toml_path(); + let toml_status = if toml_path.exists() { + "found" + } else { + "not found (run `ironclaw config init` to create)" + }; + println!( + "TOML config: {} ({})", + toml_path.display(), + toml_status + ); + Ok(()) } @@ -230,4 +284,39 @@ mod tests { settings.reset("agent.name").unwrap(); assert_eq!(settings.agent.name, "ironclaw"); } + + #[tokio::test] + async fn init_toml_creates_file() { + let dir = tempdir().unwrap(); + let path = dir.path().join("config.toml"); + + init_toml(None, Some(path.clone()), false).await.unwrap(); + assert!(path.exists()); + + let content = std::fs::read_to_string(&path).unwrap(); + assert!(content.contains("[agent]")); + } + + #[tokio::test] + async fn init_toml_refuses_overwrite_without_force() { + let dir = tempdir().unwrap(); + let path = dir.path().join("config.toml"); + std::fs::write(&path, "existing").unwrap(); + + let result = init_toml(None, Some(path.clone()), false).await; + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("already exists")); + } + + #[tokio::test] + async fn init_toml_force_overwrites() { + let dir = tempdir().unwrap(); + let path = dir.path().join("config.toml"); + std::fs::write(&path, "old content").unwrap(); + + init_toml(None, Some(path.clone()), true).await.unwrap(); + + let content = std::fs::read_to_string(&path).unwrap(); + assert!(content.contains("[agent]")); + } } diff --git a/src/cli/doctor.rs b/src/cli/doctor.rs new file mode 100644 index 00000000..6746a8e1 --- /dev/null +++ b/src/cli/doctor.rs @@ -0,0 +1,287 @@ +//! `ironclaw doctor` - active health diagnostics. +//! +//! Probes external dependencies and validates configuration to surface +//! problems before they bite during normal operation. Each check reports +//! pass/fail with actionable guidance on failures. + +use std::path::PathBuf; + +/// Run all diagnostic checks and print results. +pub async fn run_doctor_command() -> anyhow::Result<()> { + println!("IronClaw Doctor"); + println!("===============\n"); + + let mut passed = 0u32; + let mut failed = 0u32; + + // ── Configuration checks ────────────────────────────────── + + check( + "NEAR AI session", + check_nearai_session().await, + &mut passed, + &mut failed, + ); + + check( + "Database backend", + check_database().await, + &mut passed, + &mut failed, + ); + + check( + "Workspace directory", + check_workspace_dir(), + &mut passed, + &mut failed, + ); + + // ── External binary checks ──────────────────────────────── + + check( + "Docker", + check_binary("docker", &["--version"]), + &mut passed, + &mut failed, + ); + + check( + "cloudflared", + check_binary("cloudflared", &["--version"]), + &mut passed, + &mut failed, + ); + + check( + "ngrok", + check_binary("ngrok", &["version"]), + &mut passed, + &mut failed, + ); + + check( + "tailscale", + check_binary("tailscale", &["version"]), + &mut passed, + &mut failed, + ); + + // ── Summary ─────────────────────────────────────────────── + + println!(); + println!(" {passed} passed, {failed} failed"); + + if failed > 0 { + println!("\n Some checks failed. This is normal if you don't use those features."); + } + + Ok(()) +} + +// ── Individual checks ─────────────────────────────────────── + +fn check(name: &str, result: CheckResult, passed: &mut u32, failed: &mut u32) { + match result { + CheckResult::Pass(detail) => { + *passed += 1; + println!(" [pass] {name}: {detail}"); + } + CheckResult::Fail(detail) => { + *failed += 1; + println!(" [FAIL] {name}: {detail}"); + } + CheckResult::Skip(reason) => { + println!(" [skip] {name}: {reason}"); + } + } +} + +enum CheckResult { + Pass(String), + Fail(String), + Skip(String), +} + +async fn check_nearai_session() -> CheckResult { + // Check if session file exists + let session_path = crate::llm::session::default_session_path(); + if !session_path.exists() { + // Check for API key mode + if std::env::var("NEARAI_API_KEY").is_ok() { + return CheckResult::Pass("API key configured".into()); + } + return CheckResult::Fail(format!( + "session file not found at {}. Run `ironclaw onboard`", + session_path.display() + )); + } + + // Verify the session file is readable and non-empty + match std::fs::read_to_string(&session_path) { + Ok(content) if content.trim().is_empty() => { + CheckResult::Fail("session file is empty".into()) + } + Ok(_) => CheckResult::Pass(format!("session found ({})", session_path.display())), + Err(e) => CheckResult::Fail(format!("cannot read session file: {e}")), + } +} + +async fn check_database() -> CheckResult { + let backend = std::env::var("DATABASE_BACKEND") + .ok() + .unwrap_or_else(|| "postgres".into()); + + match backend.as_str() { + "libsql" | "turso" | "sqlite" => { + let path = std::env::var("LIBSQL_PATH") + .map(PathBuf::from) + .unwrap_or_else(|_| crate::config::default_libsql_path()); + + if path.exists() { + CheckResult::Pass(format!("libSQL database exists ({})", path.display())) + } else { + CheckResult::Pass(format!( + "libSQL database not found at {} (will be created on first run)", + path.display() + )) + } + } + _ => { + if std::env::var("DATABASE_URL").is_ok() { + // Try to connect + match try_pg_connect().await { + Ok(()) => CheckResult::Pass("PostgreSQL connected".into()), + Err(e) => CheckResult::Fail(format!("PostgreSQL connection failed: {e}")), + } + } else { + CheckResult::Fail("DATABASE_URL not set".into()) + } + } + } +} + +#[cfg(feature = "postgres")] +async fn try_pg_connect() -> Result<(), String> { + let url = std::env::var("DATABASE_URL").map_err(|_| "DATABASE_URL not set".to_string())?; + + let config = deadpool_postgres::Config { + url: Some(url), + ..Default::default() + }; + let pool = config + .create_pool( + Some(deadpool_postgres::Runtime::Tokio1), + tokio_postgres::NoTls, + ) + .map_err(|e| format!("pool error: {e}"))?; + + let client = tokio::time::timeout(std::time::Duration::from_secs(5), pool.get()) + .await + .map_err(|_| "connection timeout (5s)".to_string())? + .map_err(|e| format!("{e}"))?; + + client + .execute("SELECT 1", &[]) + .await + .map_err(|e| format!("{e}"))?; + + Ok(()) +} + +#[cfg(not(feature = "postgres"))] +async fn try_pg_connect() -> Result<(), String> { + Err("postgres feature not compiled in".into()) +} + +fn check_workspace_dir() -> CheckResult { + let dir = dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw"); + + if dir.exists() { + if dir.is_dir() { + CheckResult::Pass(format!("{}", dir.display())) + } else { + CheckResult::Fail(format!("{} exists but is not a directory", dir.display())) + } + } else { + CheckResult::Pass(format!("{} will be created on first run", dir.display())) + } +} + +fn check_binary(name: &str, args: &[&str]) -> CheckResult { + match std::process::Command::new(name) + .args(args) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .output() + { + Ok(output) => { + let version = String::from_utf8_lossy(&output.stdout); + let version = version.trim(); + // Some tools print version to stderr + let version = if version.is_empty() { + let stderr = String::from_utf8_lossy(&output.stderr); + stderr.trim().lines().next().unwrap_or("").to_string() + } else { + version.lines().next().unwrap_or("").to_string() + }; + + if output.status.success() { + CheckResult::Pass(version) + } else { + CheckResult::Fail(format!("exited with {}", output.status)) + } + } + Err(_) => CheckResult::Skip(format!("{name} not found in PATH")), + } +} + +#[cfg(test)] +mod tests { + use crate::cli::doctor::*; + + #[test] + fn check_binary_finds_sh() { + match check_binary("sh", &["-c", "echo ok"]) { + CheckResult::Pass(_) => {} + other => panic!("expected Pass for sh, got: {}", format_result(&other)), + } + } + + #[test] + fn check_binary_skips_nonexistent() { + match check_binary("__ironclaw_nonexistent_binary__", &["--version"]) { + CheckResult::Skip(_) => {} + other => panic!( + "expected Skip for nonexistent binary, got: {}", + format_result(&other) + ), + } + } + + #[test] + fn check_workspace_dir_does_not_panic() { + let result = check_workspace_dir(); + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[tokio::test] + async fn check_nearai_session_does_not_panic() { + let result = check_nearai_session().await; + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + fn format_result(r: &CheckResult) -> String { + match r { + CheckResult::Pass(s) => format!("Pass({s})"), + CheckResult::Fail(s) => format!("Fail({s})"), + CheckResult::Skip(s) => format!("Skip({s})"), + } + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 06715d8c..ce193013 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -7,23 +7,29 @@ //! - Managing WASM tools (`tool install`, `tool list`, `tool remove`) //! - Managing MCP servers (`mcp add`, `mcp auth`, `mcp list`, `mcp test`) //! - Querying workspace memory (`memory search`, `memory read`, `memory write`) +//! - Managing OS service (`service install`, `service start`, `service stop`) +//! - Active health diagnostics (`doctor`) //! - Checking system health (`status`) mod config; +mod doctor; mod mcp; pub mod memory; pub mod oauth_defaults; mod pairing; +mod service; pub mod status; mod tool; pub use config::{ConfigCommand, run_config_command}; +pub use doctor::run_doctor_command; pub use mcp::{McpCommand, run_mcp_command}; pub use memory::MemoryCommand; #[cfg(feature = "postgres")] pub use memory::run_memory_command; pub use memory::run_memory_command_with_db; pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store}; +pub use service::{ServiceCommand, run_service_command}; pub use status::run_status_command; pub use tool::{ToolCommand, run_tool_command}; @@ -96,6 +102,13 @@ pub enum Command { #[command(subcommand)] Pairing(PairingCommand), + /// Manage OS service (launchd / systemd) + #[command(subcommand)] + Service(ServiceCommand), + + /// Probe external dependencies and validate configuration + Doctor, + /// Show system health and diagnostics Status, diff --git a/src/cli/service.rs b/src/cli/service.rs new file mode 100644 index 00000000..100f472f --- /dev/null +++ b/src/cli/service.rs @@ -0,0 +1,37 @@ +//! CLI subcommand definitions for `ironclaw service`. + +use clap::Subcommand; + +use crate::service::ServiceAction; + +#[derive(Subcommand, Debug, Clone)] +pub enum ServiceCommand { + /// Install the OS service (launchd on macOS, systemd on Linux). + Install, + /// Start the installed service. + Start, + /// Stop the running service. + Stop, + /// Show service status. + Status, + /// Uninstall the OS service and remove the unit file. + Uninstall, +} + +impl ServiceCommand { + /// Convert the CLI variant into the domain action. + pub fn to_action(&self) -> ServiceAction { + match self { + ServiceCommand::Install => ServiceAction::Install, + ServiceCommand::Start => ServiceAction::Start, + ServiceCommand::Stop => ServiceAction::Stop, + ServiceCommand::Status => ServiceAction::Status, + ServiceCommand::Uninstall => ServiceAction::Uninstall, + } + } +} + +/// Run the service command. +pub fn run_service_command(cmd: &ServiceCommand) -> anyhow::Result<()> { + crate::service::handle_command(&cmd.to_action()) +} diff --git a/src/config.rs b/src/config.rs index b5bd55cb..c2819d54 100644 --- a/src/config.rs +++ b/src/config.rs @@ -39,22 +39,32 @@ pub struct Config { pub routines: RoutineConfig, pub sandbox: SandboxModeConfig, pub claude_code: ClaudeCodeConfig, + pub observability: crate::observability::ObservabilityConfig, } impl Config { /// Load configuration from environment variables and the database. /// - /// Priority: env var > DB settings > default. + /// Priority: env var > TOML config file > DB settings > default. /// This is the primary way to load config after DB is connected. pub async fn from_db( store: &dyn crate::db::Database, user_id: &str, + ) -> Result { + Self::from_db_with_toml(store, user_id, None).await + } + + /// Load from DB with an optional TOML config file overlay. + pub async fn from_db_with_toml( + store: &dyn crate::db::Database, + user_id: &str, + toml_path: Option<&std::path::Path>, ) -> Result { let _ = dotenvy::dotenv(); crate::bootstrap::load_ironclaw_env(); // Load all settings from DB into a Settings struct - let db_settings = match store.get_all_settings(user_id).await { + let mut db_settings = match store.get_all_settings(user_id).await { Ok(map) => Settings::from_db_map(&map), Err(e) => { tracing::warn!("Failed to load settings from DB, using defaults: {}", e); @@ -62,6 +72,9 @@ impl Config { } }; + // Overlay TOML config file (values win over DB settings) + Self::apply_toml_overlay(&mut db_settings, toml_path)?; + Self::build(&db_settings).await } @@ -74,12 +87,63 @@ impl Config { /// Loads both `./.env` (standard, higher priority) and `~/.ironclaw/.env` /// (lower priority) via dotenvy, which never overwrites existing vars. pub async fn from_env() -> Result { + Self::from_env_with_toml(None).await + } + + /// Load from env with an optional TOML config file overlay. + pub async fn from_env_with_toml( + toml_path: Option<&std::path::Path>, + ) -> Result { let _ = dotenvy::dotenv(); crate::bootstrap::load_ironclaw_env(); - let settings = Settings::load(); + let mut settings = Settings::load(); + + // Overlay TOML config file (values win over JSON settings) + Self::apply_toml_overlay(&mut settings, toml_path)?; + Self::build(&settings).await } + /// Load and merge a TOML config file into settings. + /// + /// If `explicit_path` is `Some`, loads from that path (errors are fatal). + /// If `None`, tries the default path `~/.ironclaw/config.toml` (missing + /// file is silently ignored). + fn apply_toml_overlay( + settings: &mut Settings, + explicit_path: Option<&std::path::Path>, + ) -> Result<(), ConfigError> { + let path = explicit_path + .map(std::path::PathBuf::from) + .unwrap_or_else(Settings::default_toml_path); + + match Settings::load_toml(&path) { + Ok(Some(toml_settings)) => { + settings.merge_from(&toml_settings); + tracing::debug!("Loaded TOML config from {}", path.display()); + } + Ok(None) => { + if explicit_path.is_some() { + return Err(ConfigError::ParseError(format!( + "Config file not found: {}", + path.display() + ))); + } + } + Err(e) => { + if explicit_path.is_some() { + return Err(ConfigError::ParseError(format!( + "Failed to load config file {}: {}", + path.display(), + e + ))); + } + tracing::warn!("Failed to load default config file: {}", e); + } + } + Ok(()) + } + /// Build config from settings (shared by from_env and from_db). async fn build(settings: &Settings) -> Result { Ok(Self { @@ -97,6 +161,9 @@ impl Config { routines: RoutineConfig::resolve()?, sandbox: SandboxModeConfig::resolve()?, claude_code: ClaudeCodeConfig::resolve()?, + observability: crate::observability::ObservabilityConfig { + backend: std::env::var("OBSERVABILITY_BACKEND").unwrap_or_else(|_| "none".into()), + }, }) } } @@ -105,10 +172,21 @@ impl Config { /// /// Used by channels and tools that need public webhook endpoints. /// The tunnel URL is shared across all channels (Telegram, Slack, etc.). +/// +/// Two modes: +/// - **Static URL** (`TUNNEL_URL`): set the public URL directly (manual tunnel) +/// - **Managed provider** (`TUNNEL_PROVIDER`): lifecycle-managed tunnel process +/// +/// When a managed provider is configured _and_ no static URL is set, +/// the gateway starts the tunnel on boot and populates `public_url`. #[derive(Debug, Clone, Default)] pub struct TunnelConfig { /// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io"). + /// Set statically via `TUNNEL_URL` or populated at runtime by a managed tunnel. pub public_url: Option, + /// Provider configuration for lifecycle-managed tunnels. + /// `None` when using a static URL or no tunnel at all. + pub provider: Option, } impl TunnelConfig { @@ -125,12 +203,65 @@ impl TunnelConfig { }); } - Ok(Self { public_url }) + // Resolve managed tunnel provider config. + // Priority: env var > settings > default (none). + let provider_name = optional_env("TUNNEL_PROVIDER")? + .or_else(|| settings.tunnel.provider.clone()) + .unwrap_or_default(); + + let provider = if provider_name.is_empty() || provider_name == "none" { + None + } else { + Some(crate::tunnel::TunnelProviderConfig { + provider: provider_name.clone(), + cloudflare: optional_env("TUNNEL_CF_TOKEN")? + .or_else(|| settings.tunnel.cf_token.clone()) + .map(|token| crate::tunnel::CloudflareTunnelConfig { token }), + tailscale: Some(crate::tunnel::TailscaleTunnelConfig { + funnel: optional_env("TUNNEL_TS_FUNNEL") + .ok() + .flatten() + .map(|s| s == "true" || s == "1") + .unwrap_or(settings.tunnel.ts_funnel), + hostname: optional_env("TUNNEL_TS_HOSTNAME") + .ok() + .flatten() + .or_else(|| settings.tunnel.ts_hostname.clone()), + }), + ngrok: optional_env("TUNNEL_NGROK_TOKEN")? + .or_else(|| settings.tunnel.ngrok_token.clone()) + .map(|auth_token| crate::tunnel::NgrokTunnelConfig { + auth_token, + domain: optional_env("TUNNEL_NGROK_DOMAIN") + .ok() + .flatten() + .or_else(|| settings.tunnel.ngrok_domain.clone()), + }), + custom: optional_env("TUNNEL_CUSTOM_COMMAND")? + .or_else(|| settings.tunnel.custom_command.clone()) + .map(|start_command| crate::tunnel::CustomTunnelConfig { + start_command, + health_url: optional_env("TUNNEL_CUSTOM_HEALTH_URL") + .ok() + .flatten() + .or_else(|| settings.tunnel.custom_health_url.clone()), + url_pattern: optional_env("TUNNEL_CUSTOM_URL_PATTERN") + .ok() + .flatten() + .or_else(|| settings.tunnel.custom_url_pattern.clone()), + }), + }) + }; + + Ok(Self { + public_url, + provider, + }) } - /// Check if a tunnel is configured. + /// Check if a tunnel is configured (static URL or managed provider). pub fn is_enabled(&self) -> bool { - self.public_url.is_some() + self.public_url.is_some() || self.provider.is_some() } /// Get the webhook URL for a given path. @@ -419,6 +550,19 @@ pub struct NearAiConfig { /// With the default of 3, the provider makes up to 4 total attempts /// (1 initial + 3 retries) before giving up. pub max_retries: u32, + /// Consecutive transient failures before the circuit breaker opens. + /// None = disabled (default). E.g. 5 means after 5 consecutive failures + /// all requests are rejected until recovery timeout elapses. + pub circuit_breaker_threshold: Option, + /// How long (seconds) the circuit stays open before allowing a probe (default: 30). + pub circuit_breaker_recovery_secs: u64, + /// Enable in-memory response caching for `complete()` calls. + /// Saves tokens on repeated prompts within a session. Default: false. + pub response_cache_enabled: bool, + /// TTL in seconds for cached responses (default: 3600 = 1 hour). + pub response_cache_ttl_secs: u64, + /// Max cached responses before LRU eviction (default: 1000). + pub response_cache_max_entries: usize, /// Cooldown duration in seconds for the failover provider (default: 300). /// When a provider accumulates enough consecutive failures it is skipped /// for this many seconds. @@ -485,6 +629,17 @@ impl LlmConfig { api_key: nearai_api_key, fallback_model: optional_env("NEARAI_FALLBACK_MODEL")?, max_retries: parse_optional_env("NEARAI_MAX_RETRIES", 3)?, + circuit_breaker_threshold: optional_env("CIRCUIT_BREAKER_THRESHOLD")? + .map(|s| s.parse()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "CIRCUIT_BREAKER_THRESHOLD".to_string(), + message: format!("must be a positive integer: {e}"), + })?, + circuit_breaker_recovery_secs: parse_optional_env("CIRCUIT_BREAKER_RECOVERY_SECS", 30)?, + response_cache_enabled: parse_optional_env("RESPONSE_CACHE_ENABLED", false)?, + response_cache_ttl_secs: parse_optional_env("RESPONSE_CACHE_TTL_SECS", 3600)?, + response_cache_max_entries: parse_optional_env("RESPONSE_CACHE_MAX_ENTRIES", 1000)?, failover_cooldown_secs: parse_optional_env("LLM_FAILOVER_COOLDOWN_SECS", 300)?, failover_cooldown_threshold: parse_optional_env("LLM_FAILOVER_THRESHOLD", 3)?, }; @@ -755,6 +910,10 @@ pub struct AgentConfig { pub session_idle_timeout: Duration, /// Allow chat to use filesystem/shell tools directly (bypass sandbox). pub allow_local_tools: bool, + /// Maximum daily LLM spend in cents (e.g. 10000 = $100). None = unlimited. + pub max_cost_per_day_cents: Option, + /// Maximum LLM/tool actions per hour. None = unlimited. + pub max_actions_per_hour: Option, } impl AgentConfig { @@ -833,6 +992,20 @@ impl AgentConfig { message: format!("must be 'true' or 'false': {e}"), })? .unwrap_or(false), + max_cost_per_day_cents: optional_env("MAX_COST_PER_DAY_CENTS")? + .map(|s| s.parse()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "MAX_COST_PER_DAY_CENTS".to_string(), + message: format!("must be a positive integer: {e}"), + })?, + max_actions_per_hour: optional_env("MAX_ACTIONS_PER_HOUR")? + .map(|s| s.parse()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "MAX_ACTIONS_PER_HOUR".to_string(), + message: format!("must be a positive integer: {e}"), + })?, }) } } diff --git a/src/lib.rs b/src/lib.rs index 742dadbc..0abcdb5b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -53,15 +53,18 @@ pub mod extensions; pub mod history; pub mod hooks; pub mod llm; +pub mod observability; pub mod orchestrator; pub mod pairing; pub mod safety; pub mod sandbox; pub mod secrets; +pub mod service; pub mod settings; pub mod setup; pub mod tools; pub mod tracing_fmt; +pub mod tunnel; pub mod util; pub mod worker; pub mod workspace; diff --git a/src/llm/circuit_breaker.rs b/src/llm/circuit_breaker.rs new file mode 100644 index 00000000..7ac39970 --- /dev/null +++ b/src/llm/circuit_breaker.rs @@ -0,0 +1,674 @@ +//! Circuit breaker for LLM providers. +//! +//! Wraps any `LlmProvider` with a state machine that trips open after +//! consecutive transient failures, preventing request storms against a +//! degraded backend. Automatically probes for recovery via half-open state. +//! +//! ```text +//! Closed ──(failures >= threshold)──► Open +//! ▲ │ +//! │ (recovery timeout) +//! │ ▼ +//! └──(probe succeeds)──── HalfOpen ──(probe fails)──► Open +//! ``` + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use rust_decimal::Decimal; +use tokio::sync::Mutex; + +use crate::error::LlmError; +use crate::llm::provider::{ + CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest, + ToolCompletionResponse, +}; + +/// Configuration for the circuit breaker. +#[derive(Debug, Clone)] +pub struct CircuitBreakerConfig { + /// Consecutive transient failures before the circuit opens. + pub failure_threshold: u32, + /// How long the circuit stays open before allowing a probe. + pub recovery_timeout: Duration, + /// Successful probes needed in half-open to close the circuit. + pub half_open_successes_needed: u32, +} + +impl Default for CircuitBreakerConfig { + fn default() -> Self { + Self { + failure_threshold: 5, + recovery_timeout: Duration::from_secs(30), + half_open_successes_needed: 2, + } + } +} + +/// Circuit breaker states. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CircuitState { + /// Normal operation; tracking consecutive failures. + Closed, + /// Rejecting all calls; waiting for recovery timeout to elapse. + Open, + /// Allowing probe calls to test whether the backend recovered. + HalfOpen, +} + +/// Internal mutable state. +struct BreakerState { + state: CircuitState, + consecutive_failures: u32, + opened_at: Option, + half_open_successes: u32, +} + +impl BreakerState { + fn new() -> Self { + Self { + state: CircuitState::Closed, + consecutive_failures: 0, + opened_at: None, + half_open_successes: 0, + } + } +} + +/// Wraps an `LlmProvider` with circuit breaker protection. +/// +/// Tracks consecutive transient failures. After `failure_threshold` failures +/// the circuit opens and all requests are rejected for `recovery_timeout`. +/// After that timeout a probe call is allowed through (half-open); if it +/// succeeds the circuit closes, otherwise it reopens. +pub struct CircuitBreakerProvider { + inner: Arc, + state: Mutex, + config: CircuitBreakerConfig, +} + +impl CircuitBreakerProvider { + pub fn new(inner: Arc, config: CircuitBreakerConfig) -> Self { + Self { + inner, + state: Mutex::new(BreakerState::new()), + config, + } + } + + /// Current circuit state (for observability / health checks). + pub async fn circuit_state(&self) -> CircuitState { + self.state.lock().await.state + } + + /// Number of consecutive failures recorded so far. + pub async fn consecutive_failures(&self) -> u32 { + self.state.lock().await.consecutive_failures + } + + /// Pre-flight: is a call allowed right now? + async fn check_allowed(&self) -> Result<(), LlmError> { + let mut state = self.state.lock().await; + match state.state { + CircuitState::Closed | CircuitState::HalfOpen => Ok(()), + CircuitState::Open => { + if let Some(opened_at) = state.opened_at { + if opened_at.elapsed() >= self.config.recovery_timeout { + state.state = CircuitState::HalfOpen; + state.half_open_successes = 0; + tracing::info!( + provider = self.inner.model_name(), + "Circuit breaker: Open -> HalfOpen, allowing probe" + ); + Ok(()) + } else { + let remaining = self.config.recovery_timeout - opened_at.elapsed(); + Err(LlmError::RequestFailed { + provider: self.inner.model_name().to_string(), + reason: format!( + "Circuit breaker open ({} consecutive failures, \ + recovery in {:.0}s)", + state.consecutive_failures, + remaining.as_secs_f64() + ), + }) + } + } else { + // opened_at should always be Some when Open; recover gracefully + state.state = CircuitState::Closed; + Ok(()) + } + } + } + } + + /// Record a successful call. + async fn record_success(&self) { + let mut state = self.state.lock().await; + match state.state { + CircuitState::Closed => { + state.consecutive_failures = 0; + } + CircuitState::HalfOpen => { + state.half_open_successes += 1; + if state.half_open_successes >= self.config.half_open_successes_needed { + state.state = CircuitState::Closed; + state.consecutive_failures = 0; + state.opened_at = None; + tracing::info!( + provider = self.inner.model_name(), + "Circuit breaker: HalfOpen -> Closed (recovered)" + ); + } + } + CircuitState::Open => { + // Shouldn't get here (check_allowed blocks Open), but recover + state.state = CircuitState::Closed; + state.consecutive_failures = 0; + state.opened_at = None; + } + } + } + + /// Record a failed call; only transient errors count toward the threshold. + async fn record_failure(&self, err: &LlmError) { + if !is_transient(err) { + return; + } + + let mut state = self.state.lock().await; + match state.state { + CircuitState::Closed => { + state.consecutive_failures += 1; + if state.consecutive_failures >= self.config.failure_threshold { + state.state = CircuitState::Open; + state.opened_at = Some(Instant::now()); + tracing::warn!( + provider = self.inner.model_name(), + failures = state.consecutive_failures, + "Circuit breaker: Closed -> Open" + ); + } + } + CircuitState::HalfOpen => { + state.state = CircuitState::Open; + state.opened_at = Some(Instant::now()); + state.half_open_successes = 0; + tracing::warn!( + provider = self.inner.model_name(), + "Circuit breaker: HalfOpen -> Open (probe failed)" + ); + } + CircuitState::Open => {} + } + } +} + +/// 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. +fn is_transient(err: &LlmError) -> bool { + matches!( + err, + LlmError::RequestFailed { .. } + | LlmError::RateLimited { .. } + | LlmError::InvalidResponse { .. } + | LlmError::SessionExpired { .. } + | LlmError::SessionRenewalFailed { .. } + | LlmError::Http(_) + | LlmError::Json(_) + | LlmError::Io(_) + ) +} + +#[async_trait] +impl LlmProvider for CircuitBreakerProvider { + 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 { + self.check_allowed().await?; + match self.inner.complete(request).await { + Ok(resp) => { + self.record_success().await; + Ok(resp) + } + Err(err) => { + self.record_failure(&err).await; + Err(err) + } + } + } + + async fn complete_with_tools( + &self, + request: ToolCompletionRequest, + ) -> Result { + self.check_allowed().await?; + match self.inner.complete_with_tools(request).await { + Ok(resp) => { + self.record_success().await; + Ok(resp) + } + Err(err) => { + self.record_failure(&err).await; + Err(err) + } + } + } + + 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::*; + + use std::sync::atomic::{AtomicBool, Ordering}; + + use crate::llm::provider::{CompletionResponse, FinishReason, ToolCompletionResponse}; + + /// A test stub that either always succeeds or always fails with a + /// configurable error. The `should_fail` flag can be flipped at + /// runtime for half-open recovery tests. + struct StubProvider { + name: String, + should_fail: AtomicBool, + error_kind: StubError, + } + + #[derive(Clone)] + enum StubError { + Transient, + NonTransient, + } + + impl StubProvider { + fn always_ok(name: &str) -> Arc { + Arc::new(Self { + name: name.to_string(), + should_fail: AtomicBool::new(false), + error_kind: StubError::Transient, + }) + } + + fn always_fail(name: &str) -> Arc { + Arc::new(Self { + name: name.to_string(), + should_fail: AtomicBool::new(true), + error_kind: StubError::Transient, + }) + } + + fn always_fail_non_transient(name: &str) -> Arc { + Arc::new(Self { + name: name.to_string(), + should_fail: AtomicBool::new(true), + error_kind: StubError::NonTransient, + }) + } + + fn set_failing(&self, fail: bool) { + self.should_fail.store(fail, Ordering::Relaxed); + } + + fn make_error(&self) -> LlmError { + match self.error_kind { + StubError::Transient => LlmError::RequestFailed { + provider: self.name.clone(), + reason: "server error".to_string(), + }, + StubError::NonTransient => LlmError::ContextLengthExceeded { + used: 100_000, + limit: 50_000, + }, + } + } + + fn ok_response() -> CompletionResponse { + CompletionResponse { + content: "ok".to_string(), + input_tokens: 10, + output_tokens: 5, + finish_reason: FinishReason::Stop, + response_id: None, + } + } + + fn ok_tool_response() -> ToolCompletionResponse { + ToolCompletionResponse { + content: Some("ok".to_string()), + tool_calls: vec![], + input_tokens: 10, + output_tokens: 5, + finish_reason: FinishReason::Stop, + response_id: None, + } + } + } + + #[async_trait] + impl LlmProvider for StubProvider { + fn model_name(&self) -> &str { + &self.name + } + + fn cost_per_token(&self) -> (Decimal, Decimal) { + (Decimal::ZERO, Decimal::ZERO) + } + + async fn complete( + &self, + _request: CompletionRequest, + ) -> Result { + if self.should_fail.load(Ordering::Relaxed) { + Err(self.make_error()) + } else { + Ok(Self::ok_response()) + } + } + + async fn complete_with_tools( + &self, + _request: ToolCompletionRequest, + ) -> Result { + if self.should_fail.load(Ordering::Relaxed) { + Err(self.make_error()) + } else { + Ok(Self::ok_tool_response()) + } + } + } + + 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(threshold: u32) -> CircuitBreakerConfig { + CircuitBreakerConfig { + failure_threshold: threshold, + recovery_timeout: Duration::from_millis(50), + half_open_successes_needed: 1, + } + } + + // -- State machine tests -- + + #[tokio::test] + async fn closed_allows_calls_and_resets_on_success() { + let stub = StubProvider::always_ok("test"); + let cb = CircuitBreakerProvider::new(stub, fast_config(3)); + + let resp = cb.complete(make_request()).await; + assert!(resp.is_ok()); + assert_eq!(cb.circuit_state().await, CircuitState::Closed); + assert_eq!(cb.consecutive_failures().await, 0); + } + + #[tokio::test] + async fn failures_accumulate_then_trip_to_open() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new(stub, fast_config(3)); + + // First 2 failures: still closed + for i in 0..2 { + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Closed); + assert_eq!(cb.consecutive_failures().await, i + 1); + } + + // 3rd failure: trips to open + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Open); + } + + #[tokio::test] + async fn open_rejects_immediately() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new( + stub, + CircuitBreakerConfig { + failure_threshold: 1, + recovery_timeout: Duration::from_secs(60), + half_open_successes_needed: 1, + }, + ); + + // Trip the breaker + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Open); + + // Next call should fail with circuit breaker message + let err = cb.complete(make_request()).await.unwrap_err(); + match err { + LlmError::RequestFailed { reason, .. } => { + assert!( + reason.contains("Circuit breaker open"), + "Expected circuit breaker message, got: {}", + reason + ); + } + other => panic!("Expected RequestFailed, got: {:?}", other), + } + } + + #[tokio::test] + async fn recovery_timeout_transitions_to_half_open() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new(stub, fast_config(1)); + + // Trip to open + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Open); + + // Wait for recovery timeout + tokio::time::sleep(Duration::from_millis(60)).await; + + // Next call should transition to half-open (and fail, since stub fails) + let _ = cb.complete(make_request()).await; + // Failed probe sends it back to Open + assert_eq!(cb.circuit_state().await, CircuitState::Open); + } + + #[tokio::test] + async fn half_open_success_closes_circuit() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(1)); + + // Trip to open + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Open); + + // Wait for recovery, then make the stub succeed + tokio::time::sleep(Duration::from_millis(60)).await; + stub.set_failing(false); + + // Probe should succeed, closing the circuit + let resp = cb.complete(make_request()).await; + assert!(resp.is_ok()); + assert_eq!(cb.circuit_state().await, CircuitState::Closed); + assert_eq!(cb.consecutive_failures().await, 0); + } + + #[tokio::test] + async fn half_open_failure_reopens_circuit() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new(stub, fast_config(1)); + + // Trip to open + let _ = cb.complete(make_request()).await; + + // Wait for recovery timeout + tokio::time::sleep(Duration::from_millis(60)).await; + + // Probe fails (stub still failing) + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Open); + } + + #[tokio::test] + async fn non_transient_errors_do_not_trip_breaker() { + let stub = StubProvider::always_fail_non_transient("test"); + let cb = CircuitBreakerProvider::new(stub, fast_config(1)); + + // ContextLengthExceeded is not transient; breaker should stay closed + for _ in 0..5 { + let _ = cb.complete(make_request()).await; + } + assert_eq!(cb.circuit_state().await, CircuitState::Closed); + assert_eq!(cb.consecutive_failures().await, 0); + } + + #[tokio::test] + async fn success_resets_failure_count() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(3)); + + // Accumulate 2 failures + let _ = cb.complete(make_request()).await; + let _ = cb.complete(make_request()).await; + assert_eq!(cb.consecutive_failures().await, 2); + + // One success resets the counter + stub.set_failing(false); + let resp = cb.complete(make_request()).await; + assert!(resp.is_ok()); + assert_eq!(cb.consecutive_failures().await, 0); + } + + #[tokio::test] + async fn complete_with_tools_uses_same_breaker_logic() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new(stub, fast_config(2)); + + let _ = cb.complete_with_tools(make_tool_request()).await; + let _ = cb.complete_with_tools(make_tool_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Open); + } + + #[tokio::test] + async fn multiple_half_open_successes_needed() { + let stub = StubProvider::always_fail("test"); + let cb = CircuitBreakerProvider::new( + stub.clone(), + CircuitBreakerConfig { + failure_threshold: 1, + recovery_timeout: Duration::from_millis(50), + half_open_successes_needed: 3, + }, + ); + + // Trip to open + let _ = cb.complete(make_request()).await; + + // Wait and flip to succeed + tokio::time::sleep(Duration::from_millis(60)).await; + stub.set_failing(false); + + // First probe: half-open, success but not enough yet + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen); + + // Second probe: still half-open + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen); + + // Third probe: closes + let _ = cb.complete(make_request()).await; + assert_eq!(cb.circuit_state().await, CircuitState::Closed); + } + + // -- Error classification tests -- + + #[test] + fn transient_classification() { + // Transient + assert!(is_transient(&LlmError::RequestFailed { + provider: "p".into(), + reason: "err".into(), + })); + assert!(is_transient(&LlmError::RateLimited { + provider: "p".into(), + retry_after: None, + })); + assert!(is_transient(&LlmError::InvalidResponse { + provider: "p".into(), + reason: "bad".into(), + })); + assert!(is_transient(&LlmError::SessionExpired { + provider: "p".into(), + })); + assert!(is_transient(&LlmError::SessionRenewalFailed { + provider: "p".into(), + reason: "timeout".into(), + })); + assert!(is_transient(&LlmError::Io(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "reset" + )))); + + // NOT transient + assert!(!is_transient(&LlmError::AuthFailed { + provider: "p".into(), + })); + assert!(!is_transient(&LlmError::ContextLengthExceeded { + used: 100_000, + limit: 50_000, + })); + assert!(!is_transient(&LlmError::ModelNotAvailable { + provider: "p".into(), + model: "m".into(), + })); + } + + // -- Passthrough delegation tests -- + + #[tokio::test] + async fn passthrough_methods_delegate_to_inner() { + let stub = StubProvider::always_ok("my-model"); + let cb = CircuitBreakerProvider::new(stub, fast_config(3)); + + assert_eq!(cb.model_name(), "my-model"); + assert_eq!(cb.active_model_name(), "my-model"); + assert_eq!(cb.cost_per_token(), (Decimal::ZERO, Decimal::ZERO)); + assert_eq!(cb.calculate_cost(100, 50), Decimal::ZERO); + } +} diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 290e9751..198d57c7 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -7,16 +7,19 @@ //! - **Ollama**: Local model inference //! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API -mod costs; +pub mod circuit_breaker; +pub mod costs; pub mod failover; mod nearai; mod nearai_chat; mod provider; mod reasoning; +pub mod response_cache; mod retry; mod rig_adapter; pub mod session; +pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider}; pub use failover::{CooldownConfig, FailoverProvider}; pub use nearai::{ModelInfo, NearAiProvider}; pub use nearai_chat::NearAiChatProvider; @@ -28,6 +31,7 @@ pub use reasoning::{ ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, TokenUsage, ToolSelection, }; +pub use response_cache::{CachedProvider, ResponseCacheConfig}; pub use rig_adapter::RigAdapter; pub use session::{SessionConfig, SessionManager, create_session_manager}; @@ -235,6 +239,11 @@ mod tests { api_key: None, fallback_model: None, max_retries: 3, + 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, } diff --git a/src/llm/response_cache.rs b/src/llm/response_cache.rs new file mode 100644 index 00000000..cd91cabf --- /dev/null +++ b/src/llm/response_cache.rs @@ -0,0 +1,526 @@ +//! In-memory LLM response cache with TTL and LRU eviction. +//! +//! Wraps any [`LlmProvider`] and caches [`complete()`] responses keyed +//! by a SHA-256 hash of the messages and model name. Tool-calling +//! requests are never cached since they can trigger side effects. +//! +//! ```text +//! ┌──────────────────────────────────────────────────┐ +//! │ CachedProvider │ +//! │ complete() ──► cache lookup ──► hit? return │ +//! │ miss? call inner │ +//! │ store response │ +//! │ │ +//! │ complete_with_tools() ──► always call inner │ +//! └──────────────────────────────────────────────────┘ +//! ``` + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use rust_decimal::Decimal; +use sha2::{Digest, Sha256}; +use tokio::sync::Mutex; + +use crate::error::LlmError; +use crate::llm::provider::{ + CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest, + ToolCompletionResponse, +}; + +/// Configuration for the response cache. +#[derive(Debug, Clone)] +pub struct ResponseCacheConfig { + /// Time-to-live for cache entries. + pub ttl: Duration, + /// Maximum number of cached entries before LRU eviction. + pub max_entries: usize, +} + +impl Default for ResponseCacheConfig { + fn default() -> Self { + Self { + ttl: Duration::from_secs(3600), // 1 hour + max_entries: 1000, + } + } +} + +struct CacheEntry { + response: CompletionResponse, + created_at: Instant, + last_accessed: Instant, + hit_count: u64, +} + +/// LLM provider wrapper that caches `complete()` responses. +/// +/// Tool completion requests are always forwarded without caching since +/// tool calls can have side effects that should not be replayed. +pub struct CachedProvider { + inner: Arc, + cache: Mutex>, + config: ResponseCacheConfig, +} + +impl CachedProvider { + /// Wrap an existing provider with response caching. + pub fn new(inner: Arc, config: ResponseCacheConfig) -> Self { + Self { + inner, + cache: Mutex::new(HashMap::new()), + config, + } + } + + /// Number of entries currently in the cache. + pub async fn len(&self) -> usize { + self.cache.lock().await.len() + } + + /// Whether the cache is empty. + pub async fn is_empty(&self) -> bool { + self.cache.lock().await.is_empty() + } + + /// Total cache hits across all entries. + pub async fn total_hits(&self) -> u64 { + self.cache.lock().await.values().map(|e| e.hit_count).sum() + } + + /// Clear all cached entries. + pub async fn clear(&self) { + self.cache.lock().await.clear(); + } +} + +/// Build a deterministic cache key from a completion request. +/// +/// Hashes the model name, messages, and response-affecting parameters +/// (max_tokens, temperature, stop_sequences) via SHA-256. Two requests +/// with identical content and parameters produce the same key. +fn cache_key(model: &str, request: &CompletionRequest) -> String { + let mut hasher = Sha256::new(); + hasher.update(model.as_bytes()); + hasher.update(b"|"); + + // Messages are Serialize, so we can deterministically hash them. + // serde_json produces stable output for the same input structure. + if let Ok(json) = serde_json::to_string(&request.messages) { + hasher.update(json.as_bytes()); + } + + // Include response-affecting parameters so different temperatures, + // max_tokens, or stop sequences produce distinct cache keys. + hasher.update(b"|"); + if let Some(max_tokens) = request.max_tokens { + hasher.update(max_tokens.to_le_bytes()); + } + hasher.update(b"|"); + if let Some(temp) = request.temperature { + hasher.update(temp.to_le_bytes()); + } + hasher.update(b"|"); + if let Some(ref stops) = request.stop_sequences { + for s in stops { + hasher.update(s.as_bytes()); + hasher.update(b"\x00"); + } + } + + format!("{:x}", hasher.finalize()) +} + +#[async_trait] +impl LlmProvider for CachedProvider { + 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 key = cache_key(self.inner.model_name(), &request); + let now = Instant::now(); + + // Check cache + { + let mut guard = self.cache.lock().await; + if let Some(entry) = guard.get_mut(&key) { + if now.duration_since(entry.created_at) < self.config.ttl { + entry.last_accessed = now; + entry.hit_count += 1; + tracing::debug!(hits = entry.hit_count, "response cache hit"); + return Ok(entry.response.clone()); + } + // Expired, remove it + guard.remove(&key); + } + } + + // Cache miss, call the real provider + let response = self.inner.complete(request).await?; + + // Store in cache + { + let mut guard = self.cache.lock().await; + + // Evict expired entries + guard.retain(|_, entry| now.duration_since(entry.created_at) < self.config.ttl); + + // LRU eviction if over capacity + while guard.len() >= self.config.max_entries { + let oldest_key = guard + .iter() + .min_by_key(|(_, entry)| entry.last_accessed) + .map(|(k, _)| k.clone()); + + if let Some(k) = oldest_key { + guard.remove(&k); + } else { + break; + } + } + + guard.insert( + key, + CacheEntry { + response: response.clone(), + created_at: now, + last_accessed: now, + hit_count: 0, + }, + ); + } + + Ok(response) + } + + async fn complete_with_tools( + &self, + request: ToolCompletionRequest, + ) -> Result { + // Never cache tool calls; they can trigger side effects. + self.inner.complete_with_tools(request).await + } + + 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) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; + + use crate::llm::provider::{ChatMessage, FinishReason}; + use crate::llm::response_cache::*; + + /// Controllable stub provider for testing cache behavior. + struct StubProvider { + call_count: AtomicU32, + should_fail: AtomicBool, + } + + impl StubProvider { + fn new() -> Self { + Self { + call_count: AtomicU32::new(0), + should_fail: AtomicBool::new(false), + } + } + + fn calls(&self) -> u32 { + self.call_count.load(Ordering::Relaxed) + } + } + + #[async_trait] + impl LlmProvider for StubProvider { + fn model_name(&self) -> &str { + "stub-model" + } + + fn cost_per_token(&self) -> (Decimal, Decimal) { + (Decimal::ZERO, Decimal::ZERO) + } + + async fn complete( + &self, + _request: CompletionRequest, + ) -> Result { + self.call_count.fetch_add(1, Ordering::Relaxed); + if self.should_fail.load(Ordering::Relaxed) { + return Err(LlmError::RequestFailed { + provider: "stub".into(), + reason: "forced failure".into(), + }); + } + Ok(CompletionResponse { + content: "cached response".into(), + input_tokens: 10, + output_tokens: 5, + finish_reason: FinishReason::Stop, + response_id: None, + }) + } + + async fn complete_with_tools( + &self, + _request: ToolCompletionRequest, + ) -> Result { + self.call_count.fetch_add(1, Ordering::Relaxed); + Ok(ToolCompletionResponse { + content: Some("tool response".into()), + tool_calls: vec![], + input_tokens: 10, + output_tokens: 5, + finish_reason: FinishReason::Stop, + response_id: None, + }) + } + } + + fn simple_request() -> CompletionRequest { + CompletionRequest { + messages: vec![ChatMessage::user("hello")], + max_tokens: None, + temperature: None, + stop_sequences: None, + metadata: Default::default(), + } + } + + fn different_request() -> CompletionRequest { + CompletionRequest { + messages: vec![ChatMessage::user("goodbye")], + max_tokens: None, + temperature: None, + stop_sequences: None, + metadata: Default::default(), + } + } + + #[test] + fn cache_key_is_deterministic() { + let req = simple_request(); + let k1 = cache_key("model-a", &req); + let k2 = cache_key("model-a", &req); + assert_eq!(k1, k2); + assert_eq!(k1.len(), 64); // SHA-256 hex + } + + #[test] + fn cache_key_varies_by_model() { + let req = simple_request(); + let k1 = cache_key("model-a", &req); + let k2 = cache_key("model-b", &req); + assert_ne!(k1, k2); + } + + #[test] + fn cache_key_varies_by_messages() { + let k1 = cache_key("model-a", &simple_request()); + let k2 = cache_key("model-a", &different_request()); + assert_ne!(k1, k2); + } + + #[test] + fn cache_key_varies_by_temperature() { + let mut req_a = simple_request(); + req_a.temperature = Some(0.0); + let mut req_b = simple_request(); + req_b.temperature = Some(1.0); + assert_ne!(cache_key("m", &req_a), cache_key("m", &req_b)); + } + + #[test] + fn cache_key_varies_by_max_tokens() { + let mut req_a = simple_request(); + req_a.max_tokens = Some(100); + let mut req_b = simple_request(); + req_b.max_tokens = Some(500); + assert_ne!(cache_key("m", &req_a), cache_key("m", &req_b)); + } + + #[tokio::test] + async fn cache_hit_avoids_provider_call() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new( + stub.clone(), + ResponseCacheConfig { + ttl: Duration::from_secs(60), + max_entries: 100, + }, + ); + + // First call: cache miss + let r1 = cached.complete(simple_request()).await.unwrap(); + assert_eq!(stub.calls(), 1); + assert_eq!(r1.content, "cached response"); + + // Second call: cache hit + let r2 = cached.complete(simple_request()).await.unwrap(); + assert_eq!(stub.calls(), 1); // still 1 + assert_eq!(r2.content, "cached response"); + + assert_eq!(cached.total_hits().await, 1); + } + + #[tokio::test] + async fn different_messages_get_different_entries() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); + + cached.complete(simple_request()).await.unwrap(); + cached.complete(different_request()).await.unwrap(); + + assert_eq!(stub.calls(), 2); + assert_eq!(cached.len().await, 2); + } + + #[tokio::test] + async fn expired_entries_are_evicted() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new( + stub.clone(), + ResponseCacheConfig { + ttl: Duration::from_millis(1), + max_entries: 100, + }, + ); + + cached.complete(simple_request()).await.unwrap(); + assert_eq!(stub.calls(), 1); + + // Wait for TTL to expire + tokio::time::sleep(Duration::from_millis(10)).await; + + // Should be a cache miss now + cached.complete(simple_request()).await.unwrap(); + assert_eq!(stub.calls(), 2); + } + + #[tokio::test] + async fn lru_eviction_removes_oldest() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new( + stub.clone(), + ResponseCacheConfig { + ttl: Duration::from_secs(60), + max_entries: 2, + }, + ); + + // Fill cache with 2 entries + cached.complete(simple_request()).await.unwrap(); + cached.complete(different_request()).await.unwrap(); + assert_eq!(cached.len().await, 2); + + // Add a third: should evict the oldest + let third = CompletionRequest { + messages: vec![ChatMessage::user("third")], + max_tokens: None, + temperature: None, + stop_sequences: None, + metadata: Default::default(), + }; + cached.complete(third).await.unwrap(); + assert_eq!(cached.len().await, 2); + assert_eq!(stub.calls(), 3); + } + + #[tokio::test] + async fn tool_calls_are_never_cached() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); + + let req = ToolCompletionRequest { + messages: vec![ChatMessage::user("use tool")], + tools: vec![], + max_tokens: None, + temperature: None, + tool_choice: None, + metadata: Default::default(), + }; + + cached.complete_with_tools(req.clone()).await.unwrap(); + cached.complete_with_tools(req).await.unwrap(); + + // Both should have called through + assert_eq!(stub.calls(), 2); + assert!(cached.is_empty().await); + } + + #[tokio::test] + async fn provider_errors_are_not_cached() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new( + stub.clone(), + ResponseCacheConfig { + ttl: Duration::from_secs(60), + max_entries: 100, + }, + ); + + stub.should_fail.store(true, Ordering::Relaxed); + let result = cached.complete(simple_request()).await; + assert!(result.is_err()); + assert!(cached.is_empty().await); + + // After fixing the provider, should succeed and cache + stub.should_fail.store(false, Ordering::Relaxed); + cached.complete(simple_request()).await.unwrap(); + assert_eq!(cached.len().await, 1); + } + + #[tokio::test] + async fn clear_empties_cache() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); + + cached.complete(simple_request()).await.unwrap(); + assert_eq!(cached.len().await, 1); + + cached.clear().await; + assert!(cached.is_empty().await); + } + + #[test] + fn default_config_is_reasonable() { + let cfg = ResponseCacheConfig::default(); + assert_eq!(cfg.ttl, Duration::from_secs(3600)); + assert_eq!(cfg.max_entries, 1000); + } + + #[tokio::test] + async fn delegates_model_name() { + let stub = Arc::new(StubProvider::new()); + let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); + assert_eq!(cached.model_name(), "stub-model"); + } +} diff --git a/src/llm/session.rs b/src/llm/session.rs index 4e96deed..b9932c21 100644 --- a/src/llm/session.rs +++ b/src/llm/session.rs @@ -432,9 +432,9 @@ impl SessionManager { .get_setting(&user_id, "nearai.session_token") .await .map_err(|e| LlmError::SessionRenewalFailed { - provider: "nearai".to_string(), - reason: format!("DB query failed: {}", e), - })? { + provider: "nearai".to_string(), + reason: format!("DB query failed: {}", e), + })? { value } else { tracing::warn!( diff --git a/src/main.rs b/src/main.rs index b4f8578c..9939fe7f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,15 +17,18 @@ use ironclaw::{ web::log_layer::{LogBroadcaster, WebLogLayer}, }, cli::{ - Cli, Command, run_mcp_command, run_pairing_command, run_status_command, run_tool_command, + Cli, Command, run_mcp_command, run_pairing_command, run_service_command, + run_status_command, run_tool_command, }, config::Config, context::ContextManager, extensions::ExtensionManager, hooks::HookRegistry, llm::{ - CooldownConfig, FailoverProvider, LlmProvider, SessionConfig, create_cheap_llm_provider, - create_llm_provider, create_llm_provider_with_config, create_session_manager, + CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig, + FailoverProvider, LlmProvider, ResponseCacheConfig, SessionConfig, + create_cheap_llm_provider, create_llm_provider, create_llm_provider_with_config, + create_session_manager, }, orchestrator::{ ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore, @@ -152,6 +155,24 @@ async fn main() -> anyhow::Result<()> { return run_pairing_command(pairing_cmd.clone()).map_err(|e| anyhow::anyhow!("{}", e)); } + Some(Command::Service(service_cmd)) => { + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), + ) + .init(); + + return run_service_command(service_cmd); + } + Some(Command::Doctor) => { + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), + ) + .init(); + + return ironclaw::cli::run_doctor_command().await; + } Some(Command::Status) => { tracing_subscriber::fmt() .with_env_filter( @@ -286,8 +307,9 @@ async fn main() -> anyhow::Result<()> { wizard.run().await?; } - // Load initial config from env + disk (before DB is available) - let mut config = match Config::from_env().await { + // Load initial config from env + disk + optional TOML (before DB is available) + let toml_path = cli.config.as_deref(); + let mut config = match Config::from_env_with_toml(toml_path).await { Ok(c) => c, Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => { eprintln!("Configuration error: Missing required setting '{}'", key); @@ -429,7 +451,7 @@ async fn main() -> anyhow::Result<()> { } // Reload config from DB now that we have a connection. - match Config::from_db(db.as_ref(), "default").await { + match Config::from_db_with_toml(db.as_ref(), "default", toml_path).await { Ok(db_config) => { config = db_config; tracing::info!("Configuration reloaded from database"); @@ -504,7 +526,7 @@ async fn main() -> anyhow::Result<()> { // Re-resolve LlmConfig now that secrets overlay has been populated if let Some(ref db_ref) = db { - match Config::from_db(db_ref.as_ref(), "default").await { + match Config::from_db_with_toml(db_ref.as_ref(), "default", toml_path).await { Ok(refreshed) => { config = refreshed; tracing::debug!("LlmConfig re-resolved after secret injection"); @@ -516,6 +538,62 @@ async fn main() -> anyhow::Result<()> { } } + // Start managed tunnel if configured and no static URL is already set. + // + // The tunnel process runs in the background, exposing the local gateway + // port to the internet. The resulting public URL is injected into + // config.tunnel.public_url so channels and extensions pick it up. + let active_tunnel: Option> = + if config.tunnel.public_url.is_some() { + tracing::info!( + "Static tunnel URL in use: {}", + config.tunnel.public_url.as_deref().unwrap_or("?") + ); + None + } else if let Some(ref provider_config) = config.tunnel.provider { + let gateway_port = config + .channels + .gateway + .as_ref() + .map(|g| g.port) + .unwrap_or(3000); + let gateway_host = config + .channels + .gateway + .as_ref() + .map(|g| g.host.as_str()) + .unwrap_or("127.0.0.1"); + + match ironclaw::tunnel::create_tunnel(provider_config) { + Ok(Some(tunnel)) => { + tracing::info!( + "Starting {} tunnel on {}:{}...", + tunnel.name(), + gateway_host, + gateway_port + ); + match tunnel.start(gateway_host, gateway_port).await { + Ok(url) => { + tracing::info!("Tunnel started: {}", url); + config.tunnel.public_url = Some(url); + Some(tunnel) + } + Err(e) => { + tracing::error!("Failed to start tunnel: {}", e); + None + } + } + } + Ok(None) => None, + Err(e) => { + tracing::error!("Failed to create tunnel: {}", e); + None + } + } + } else { + None + }; + // Initialize LLM provider (clone session so we can reuse it for embeddings) let llm = create_llm_provider(&config.llm, session.clone())?; tracing::info!("LLM provider initialized: {}", llm.model_name()); @@ -550,6 +628,42 @@ async fn main() -> anyhow::Result<()> { llm }; + // Wrap in circuit breaker if configured + let llm: Arc = + if let Some(threshold) = config.llm.nearai.circuit_breaker_threshold { + let cb_config = CircuitBreakerConfig { + failure_threshold: threshold, + recovery_timeout: std::time::Duration::from_secs( + config.llm.nearai.circuit_breaker_recovery_secs, + ), + ..CircuitBreakerConfig::default() + }; + tracing::info!( + threshold, + recovery_secs = config.llm.nearai.circuit_breaker_recovery_secs, + "LLM circuit breaker enabled" + ); + Arc::new(CircuitBreakerProvider::new(llm, cb_config)) + } else { + llm + }; + + // Wrap in response cache if configured + let llm: Arc = if config.llm.nearai.response_cache_enabled { + let rc_config = ResponseCacheConfig { + ttl: std::time::Duration::from_secs(config.llm.nearai.response_cache_ttl_secs), + max_entries: config.llm.nearai.response_cache_max_entries, + }; + tracing::info!( + ttl_secs = config.llm.nearai.response_cache_ttl_secs, + max_entries = config.llm.nearai.response_cache_max_entries, + "LLM response cache enabled" + ); + Arc::new(CachedProvider::new(llm, rc_config)) + } else { + llm + }; + // Initialize cheap LLM provider for lightweight tasks (heartbeat, evaluation) let cheap_llm = create_cheap_llm_provider(&config.llm, session.clone())?; if let Some(ref cheap) = cheap_llm { @@ -1220,6 +1334,12 @@ async fn main() -> anyhow::Result<()> { let boot_cheap_model = cheap_llm.as_ref().map(|c| c.model_name().to_string()); // Create and run the agent + let cost_guard = Arc::new(ironclaw::agent::cost_guard::CostGuard::new( + ironclaw::agent::cost_guard::CostGuardConfig { + max_cost_per_day_cents: config.agent.max_cost_per_day_cents, + max_actions_per_hour: config.agent.max_actions_per_hour, + }, + )); let deps = AgentDeps { store: db, llm, @@ -1229,6 +1349,7 @@ async fn main() -> anyhow::Result<()> { workspace, extension_manager, hooks, + cost_guard, }; let agent = Agent::new( config.agent.clone(), @@ -1270,6 +1391,11 @@ async fn main() -> anyhow::Result<()> { claude_code_enabled: config.claude_code.enabled, routines_enabled: config.routines.enabled, channels: channel_names, + tunnel_url: active_tunnel + .as_ref() + .and_then(|t| t.public_url()) + .or_else(|| config.tunnel.public_url.clone()), + tunnel_provider: active_tunnel.as_ref().map(|t| t.name().to_string()), }; ironclaw::boot_screen::print_boot_screen(&boot_info); } @@ -1282,6 +1408,14 @@ async fn main() -> anyhow::Result<()> { server.shutdown().await; } + // Stop managed tunnel if one was started + if let Some(tunnel) = active_tunnel { + tracing::info!("Stopping {} tunnel...", tunnel.name()); + if let Err(e) = tunnel.stop().await { + tracing::warn!("Failed to stop tunnel cleanly: {}", e); + } + } + tracing::info!("Agent shutdown complete"); Ok(()) } diff --git a/src/observability/log.rs b/src/observability/log.rs new file mode 100644 index 00000000..a476ea09 --- /dev/null +++ b/src/observability/log.rs @@ -0,0 +1,181 @@ +//! Tracing-based observer that emits structured log events. +//! +//! Uses the existing `tracing` infrastructure so events appear alongside +//! normal application logs, with no extra dependencies. Good for local +//! development and debugging. + +use crate::observability::traits::{Observer, ObserverEvent, ObserverMetric}; + +/// Observer that logs events and metrics via `tracing`. +pub struct LogObserver; + +impl Observer for LogObserver { + fn record_event(&self, event: &ObserverEvent) { + match event { + ObserverEvent::AgentStart { provider, model } => { + tracing::info!(provider, model, "observer: agent.start"); + } + ObserverEvent::LlmRequest { + provider, + model, + message_count, + } => { + tracing::info!(provider, model, message_count, "observer: llm.request"); + } + ObserverEvent::LlmResponse { + provider, + model, + duration, + success, + error_message, + } => { + tracing::info!( + provider, + model, + duration_ms = duration.as_millis() as u64, + success, + error = error_message.as_deref().unwrap_or(""), + "observer: llm.response" + ); + } + ObserverEvent::ToolCallStart { tool } => { + tracing::info!(tool, "observer: tool.start"); + } + ObserverEvent::ToolCallEnd { + tool, + duration, + success, + } => { + tracing::info!( + tool, + duration_ms = duration.as_millis() as u64, + success, + "observer: tool.end" + ); + } + ObserverEvent::TurnComplete => { + tracing::info!("observer: turn.complete"); + } + ObserverEvent::ChannelMessage { channel, direction } => { + tracing::info!(channel, direction, "observer: channel.message"); + } + ObserverEvent::HeartbeatTick => { + tracing::debug!("observer: heartbeat.tick"); + } + ObserverEvent::AgentEnd { + duration, + tokens_used, + } => { + tracing::info!( + duration_secs = duration.as_secs_f64(), + tokens_used = tokens_used.unwrap_or(0), + "observer: agent.end" + ); + } + ObserverEvent::Error { component, message } => { + tracing::warn!(component, error = message.as_str(), "observer: error"); + } + } + } + + fn record_metric(&self, metric: &ObserverMetric) { + match metric { + ObserverMetric::RequestLatency(d) => { + tracing::debug!( + latency_ms = d.as_millis() as u64, + "observer: metric.request_latency" + ); + } + ObserverMetric::TokensUsed(n) => { + tracing::debug!(tokens = n, "observer: metric.tokens_used"); + } + ObserverMetric::ActiveJobs(n) => { + tracing::debug!(active_jobs = n, "observer: metric.active_jobs"); + } + ObserverMetric::QueueDepth(n) => { + tracing::debug!(queue_depth = n, "observer: metric.queue_depth"); + } + } + } + + fn name(&self) -> &str { + "log" + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use crate::observability::log::LogObserver; + use crate::observability::traits::*; + + #[test] + fn name_is_log() { + assert_eq!(LogObserver.name(), "log"); + } + + #[test] + fn record_event_does_not_panic() { + let obs = LogObserver; + obs.record_event(&ObserverEvent::AgentStart { + provider: "nearai".into(), + model: "test".into(), + }); + obs.record_event(&ObserverEvent::LlmRequest { + provider: "nearai".into(), + model: "test".into(), + message_count: 5, + }); + obs.record_event(&ObserverEvent::LlmResponse { + provider: "nearai".into(), + model: "test".into(), + duration: Duration::from_millis(150), + success: true, + error_message: None, + }); + obs.record_event(&ObserverEvent::LlmResponse { + provider: "nearai".into(), + model: "test".into(), + duration: Duration::from_millis(1500), + success: false, + error_message: Some("timeout".into()), + }); + obs.record_event(&ObserverEvent::ToolCallStart { + tool: "shell".into(), + }); + obs.record_event(&ObserverEvent::ToolCallEnd { + tool: "shell".into(), + duration: Duration::from_millis(20), + success: true, + }); + obs.record_event(&ObserverEvent::TurnComplete); + obs.record_event(&ObserverEvent::ChannelMessage { + channel: "tui".into(), + direction: "inbound".into(), + }); + obs.record_event(&ObserverEvent::HeartbeatTick); + obs.record_event(&ObserverEvent::AgentEnd { + duration: Duration::from_secs(30), + tokens_used: Some(2500), + }); + obs.record_event(&ObserverEvent::Error { + component: "llm".into(), + message: "connection refused".into(), + }); + } + + #[test] + fn record_metric_does_not_panic() { + let obs = LogObserver; + obs.record_metric(&ObserverMetric::RequestLatency(Duration::from_millis(200))); + obs.record_metric(&ObserverMetric::TokensUsed(1000)); + obs.record_metric(&ObserverMetric::ActiveJobs(5)); + obs.record_metric(&ObserverMetric::QueueDepth(12)); + } + + #[test] + fn flush_does_not_panic() { + LogObserver.flush(); + } +} diff --git a/src/observability/mod.rs b/src/observability/mod.rs new file mode 100644 index 00000000..fa0c667a --- /dev/null +++ b/src/observability/mod.rs @@ -0,0 +1,105 @@ +//! Observability subsystem: trait-based event and metric recording. +//! +//! Provides a pluggable [`Observer`] trait with multiple backends: +//! +//! | Backend | Description | +//! |---------|-------------| +//! | `noop` | Zero overhead, discards everything (default) | +//! | `log` | Emits structured events via `tracing` | +//! | `multi` | Fan-out to multiple backends simultaneously | +//! +//! The [`create_observer`] factory builds the right backend from +//! [`ObservabilityConfig`]. Future backends (OpenTelemetry, Prometheus) +//! can be added by implementing [`Observer`]. + +mod log; +mod multi; +mod noop; +pub mod traits; + +pub use self::log::LogObserver; +pub use self::multi::MultiObserver; +pub use self::noop::NoopObserver; +pub use self::traits::{Observer, ObserverEvent, ObserverMetric}; + +/// Configuration for the observability backend. +#[derive(Debug, Clone)] +pub struct ObservabilityConfig { + /// Backend name: "none", "noop", "log". + pub backend: String, +} + +impl Default for ObservabilityConfig { + fn default() -> Self { + Self { + backend: "none".into(), + } + } +} + +/// Create an observer from configuration. +/// +/// Returns a [`NoopObserver`] for "none"/"noop" (or unknown values), +/// and a [`LogObserver`] for "log". +pub fn create_observer(config: &ObservabilityConfig) -> Box { + match config.backend.as_str() { + "log" => Box::new(LogObserver), + _ => Box::new(NoopObserver), + } +} + +#[cfg(test)] +mod tests { + use crate::observability::*; + + #[test] + fn default_config_is_none() { + let cfg = ObservabilityConfig::default(); + assert_eq!(cfg.backend, "none"); + } + + #[test] + fn factory_returns_noop_for_none() { + let cfg = ObservabilityConfig { + backend: "none".into(), + }; + let obs = create_observer(&cfg); + assert_eq!(obs.name(), "noop"); + } + + #[test] + fn factory_returns_noop_for_empty() { + let cfg = ObservabilityConfig { + backend: String::new(), + }; + let obs = create_observer(&cfg); + assert_eq!(obs.name(), "noop"); + } + + #[test] + fn factory_returns_noop_for_unknown() { + let cfg = ObservabilityConfig { + backend: "prometheus".into(), + }; + let obs = create_observer(&cfg); + assert_eq!(obs.name(), "noop"); + } + + #[test] + fn factory_returns_log_for_log() { + let cfg = ObservabilityConfig { + backend: "log".into(), + }; + let obs = create_observer(&cfg); + assert_eq!(obs.name(), "log"); + } + + #[test] + fn factory_returns_noop_for_noop() { + let cfg = ObservabilityConfig { + backend: "noop".into(), + }; + let obs = create_observer(&cfg); + assert_eq!(obs.name(), "noop"); + } +} diff --git a/src/observability/multi.rs b/src/observability/multi.rs new file mode 100644 index 00000000..2b3e5d9f --- /dev/null +++ b/src/observability/multi.rs @@ -0,0 +1,139 @@ +//! Fan-out observer that dispatches to multiple backends. +//! +//! Useful for combining backends, e.g. log + OpenTelemetry simultaneously. + +use crate::observability::traits::{Observer, ObserverEvent, ObserverMetric}; + +/// Dispatches events and metrics to all inner observers. +pub struct MultiObserver { + observers: Vec>, +} + +impl MultiObserver { + /// Create from a list of observers. If the list is empty the result + /// behaves like a noop. + pub fn new(observers: Vec>) -> Self { + Self { observers } + } +} + +impl Observer for MultiObserver { + fn record_event(&self, event: &ObserverEvent) { + for obs in &self.observers { + obs.record_event(event); + } + } + + fn record_metric(&self, metric: &ObserverMetric) { + for obs in &self.observers { + obs.record_metric(metric); + } + } + + fn flush(&self) { + for obs in &self.observers { + obs.flush(); + } + } + + fn name(&self) -> &str { + "multi" + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use crate::observability::multi::MultiObserver; + use crate::observability::traits::*; + + /// Test observer that counts calls via shared atomic counters. + struct CountingObserver { + events: Arc, + metrics: Arc, + flushes: Arc, + } + + impl CountingObserver { + fn new() -> (Self, Arc, Arc, Arc) { + let events = Arc::new(AtomicUsize::new(0)); + let metrics = Arc::new(AtomicUsize::new(0)); + let flushes = Arc::new(AtomicUsize::new(0)); + ( + Self { + events: Arc::clone(&events), + metrics: Arc::clone(&metrics), + flushes: Arc::clone(&flushes), + }, + events, + metrics, + flushes, + ) + } + } + + impl Observer for CountingObserver { + fn record_event(&self, _event: &ObserverEvent) { + self.events.fetch_add(1, Ordering::Relaxed); + } + fn record_metric(&self, _metric: &ObserverMetric) { + self.metrics.fetch_add(1, Ordering::Relaxed); + } + fn flush(&self) { + self.flushes.fetch_add(1, Ordering::Relaxed); + } + fn name(&self) -> &str { + "counting" + } + } + + #[test] + fn name_is_multi() { + let multi = MultiObserver::new(vec![]); + assert_eq!(multi.name(), "multi"); + } + + #[test] + fn empty_multi_does_not_panic() { + let multi = MultiObserver::new(vec![]); + multi.record_event(&ObserverEvent::TurnComplete); + multi.record_metric(&ObserverMetric::TokensUsed(100)); + multi.flush(); + } + + #[test] + fn dispatches_to_all_observers() { + let (a, a_events, a_metrics, a_flushes) = CountingObserver::new(); + let (b, b_events, b_metrics, b_flushes) = CountingObserver::new(); + + let multi = MultiObserver::new(vec![Box::new(a), Box::new(b)]); + + multi.record_event(&ObserverEvent::TurnComplete); + multi.record_event(&ObserverEvent::HeartbeatTick); + multi.record_metric(&ObserverMetric::TokensUsed(50)); + multi.flush(); + + assert_eq!(a_events.load(Ordering::Relaxed), 2); + assert_eq!(a_metrics.load(Ordering::Relaxed), 1); + assert_eq!(a_flushes.load(Ordering::Relaxed), 1); + assert_eq!(b_events.load(Ordering::Relaxed), 2); + assert_eq!(b_metrics.load(Ordering::Relaxed), 1); + assert_eq!(b_flushes.load(Ordering::Relaxed), 1); + } + + #[test] + fn single_observer_works() { + let (obs, events, _, _) = CountingObserver::new(); + + let multi = MultiObserver::new(vec![Box::new(obs)]); + multi.record_event(&ObserverEvent::AgentEnd { + duration: Duration::from_secs(1), + tokens_used: None, + }); + + assert_eq!(events.load(Ordering::Relaxed), 1); + } +} diff --git a/src/observability/noop.rs b/src/observability/noop.rs new file mode 100644 index 00000000..ba5c41fa --- /dev/null +++ b/src/observability/noop.rs @@ -0,0 +1,60 @@ +//! Zero-overhead no-op observer. +//! +//! Default backend when observability is disabled. All methods compile to +//! nothing, so there is zero runtime cost. + +use crate::observability::traits::{Observer, ObserverEvent, ObserverMetric}; + +/// Observer that discards all events and metrics. +pub struct NoopObserver; + +impl Observer for NoopObserver { + #[inline(always)] + fn record_event(&self, _event: &ObserverEvent) {} + + #[inline(always)] + fn record_metric(&self, _metric: &ObserverMetric) {} + + fn name(&self) -> &str { + "noop" + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use crate::observability::traits::*; + + use crate::observability::noop::NoopObserver; + + #[test] + fn name_is_noop() { + assert_eq!(NoopObserver.name(), "noop"); + } + + #[test] + fn record_event_does_not_panic() { + let obs = NoopObserver; + obs.record_event(&ObserverEvent::TurnComplete); + obs.record_event(&ObserverEvent::HeartbeatTick); + obs.record_event(&ObserverEvent::AgentStart { + provider: "x".into(), + model: "y".into(), + }); + } + + #[test] + fn record_metric_does_not_panic() { + let obs = NoopObserver; + obs.record_metric(&ObserverMetric::TokensUsed(100)); + obs.record_metric(&ObserverMetric::RequestLatency(Duration::from_millis(50))); + obs.record_metric(&ObserverMetric::ActiveJobs(2)); + obs.record_metric(&ObserverMetric::QueueDepth(0)); + } + + #[test] + fn flush_does_not_panic() { + NoopObserver.flush(); + } +} diff --git a/src/observability/traits.rs b/src/observability/traits.rs new file mode 100644 index 00000000..94316b7b --- /dev/null +++ b/src/observability/traits.rs @@ -0,0 +1,143 @@ +//! Core observer trait and event/metric types. + +use std::time::Duration; + +/// Provider-agnostic observer for agent lifecycle events and metrics. +/// +/// Implementations can log to tracing, export to OpenTelemetry, write to +/// Prometheus, or do nothing at all. The agent records events at key +/// lifecycle points and the observer decides what to do with them. +/// +/// Thread-safe and cheaply cloneable behind `Arc`. +pub trait Observer: Send + Sync { + /// Record a discrete lifecycle event. + fn record_event(&self, event: &ObserverEvent); + + /// Record a numeric metric sample. + fn record_metric(&self, metric: &ObserverMetric); + + /// Flush any buffered data (e.g. OTLP batch exporter). No-op by default. + fn flush(&self) {} + + /// Human-readable backend name (e.g. "noop", "log", "otel"). + fn name(&self) -> &str; +} + +/// Discrete lifecycle events the agent can emit. +#[derive(Debug, Clone)] +pub enum ObserverEvent { + /// Agent started processing. + AgentStart { provider: String, model: String }, + + /// An LLM request was sent. + LlmRequest { + provider: String, + model: String, + message_count: usize, + }, + + /// An LLM response was received. + LlmResponse { + provider: String, + model: String, + duration: Duration, + success: bool, + error_message: Option, + }, + + /// A tool call is about to start. + ToolCallStart { tool: String }, + + /// A tool call finished. + ToolCallEnd { + tool: String, + duration: Duration, + success: bool, + }, + + /// One reasoning turn completed. + TurnComplete, + + /// A message was sent or received on a channel. + ChannelMessage { channel: String, direction: String }, + + /// The heartbeat system ran a tick. + HeartbeatTick, + + /// Agent finished processing. + AgentEnd { + duration: Duration, + tokens_used: Option, + }, + + /// An error occurred in a component. + Error { component: String, message: String }, +} + +/// Numeric metric samples. +#[derive(Debug, Clone)] +pub enum ObserverMetric { + /// Latency of a single request (histogram-style). + RequestLatency(Duration), + /// Cumulative tokens consumed. + TokensUsed(u64), + /// Current number of active jobs (gauge). + ActiveJobs(u64), + /// Current message queue depth (gauge). + QueueDepth(u64), +} + +#[cfg(test)] +mod tests { + use crate::observability::traits::*; + + #[test] + fn event_variants_are_constructible() { + let _ = ObserverEvent::AgentStart { + provider: "nearai".into(), + model: "test".into(), + }; + let _ = ObserverEvent::LlmRequest { + provider: "nearai".into(), + model: "test".into(), + message_count: 3, + }; + let _ = ObserverEvent::LlmResponse { + provider: "nearai".into(), + model: "test".into(), + duration: Duration::from_millis(100), + success: true, + error_message: None, + }; + let _ = ObserverEvent::ToolCallStart { + tool: "echo".into(), + }; + let _ = ObserverEvent::ToolCallEnd { + tool: "echo".into(), + duration: Duration::from_millis(5), + success: true, + }; + let _ = ObserverEvent::TurnComplete; + let _ = ObserverEvent::ChannelMessage { + channel: "tui".into(), + direction: "inbound".into(), + }; + let _ = ObserverEvent::HeartbeatTick; + let _ = ObserverEvent::AgentEnd { + duration: Duration::from_secs(10), + tokens_used: Some(1500), + }; + let _ = ObserverEvent::Error { + component: "llm".into(), + message: "timeout".into(), + }; + } + + #[test] + fn metric_variants_are_constructible() { + let _ = ObserverMetric::RequestLatency(Duration::from_millis(200)); + let _ = ObserverMetric::TokensUsed(500); + let _ = ObserverMetric::ActiveJobs(3); + let _ = ObserverMetric::QueueDepth(10); + } +} diff --git a/src/service.rs b/src/service.rs new file mode 100644 index 00000000..d86d36ae --- /dev/null +++ b/src/service.rs @@ -0,0 +1,357 @@ +//! OS service management for running IronClaw as a daemon. +//! +//! Generates and manages platform-native service definitions: +//! - **macOS**: launchd plist at `~/Library/LaunchAgents/com.ironclaw.daemon.plist` +//! - **Linux**: systemd user unit at `~/.config/systemd/user/ironclaw.service` +//! +//! The installed service runs `ironclaw run` (the default agent mode) and is +//! configured to restart automatically on failure. + +use std::path::PathBuf; +use std::process::Command; + +use anyhow::{Context, Result, bail}; + +const SERVICE_LABEL: &str = "com.ironclaw.daemon"; +const SYSTEMD_UNIT: &str = "ironclaw.service"; + +// ── Public dispatch ───────────────────────────────────────────── + +/// Route a service subcommand to the appropriate handler. +pub fn handle_command(command: &ServiceAction) -> Result<()> { + match command { + ServiceAction::Install => install(), + ServiceAction::Start => start(), + ServiceAction::Stop => stop(), + ServiceAction::Status => status(), + ServiceAction::Uninstall => uninstall(), + } +} + +/// The five service lifecycle actions. +#[derive(Debug, Clone)] +pub enum ServiceAction { + Install, + Start, + Stop, + Status, + Uninstall, +} + +// ── Install ───────────────────────────────────────────────────── + +fn install() -> Result<()> { + if cfg!(target_os = "macos") { + install_macos() + } else if cfg!(target_os = "linux") { + install_linux() + } else { + bail!("Service management is only supported on macOS and Linux"); + } +} + +fn install_macos() -> Result<()> { + let file = macos_plist_path()?; + if let Some(parent) = file.parent() { + std::fs::create_dir_all(parent)?; + } + + let exe = std::env::current_exe().context("failed to resolve current executable")?; + let logs_dir = ironclaw_logs_dir()?; + std::fs::create_dir_all(&logs_dir)?; + + let stdout = logs_dir.join("daemon.stdout.log"); + let stderr = logs_dir.join("daemon.stderr.log"); + + let plist = format!( + r#" + + + + Label + {label} + ProgramArguments + + {exe} + run + + RunAtLoad + + KeepAlive + + StandardOutPath + {stdout} + StandardErrorPath + {stderr} + + +"#, + label = SERVICE_LABEL, + exe = xml_escape(&exe.display().to_string()), + stdout = xml_escape(&stdout.display().to_string()), + stderr = xml_escape(&stderr.display().to_string()), + ); + + std::fs::write(&file, plist)?; + println!("Installed launchd service: {}", file.display()); + println!(" Start with: ironclaw service start"); + Ok(()) +} + +fn install_linux() -> Result<()> { + let file = linux_unit_path()?; + if let Some(parent) = file.parent() { + std::fs::create_dir_all(parent)?; + } + + let exe = std::env::current_exe().context("failed to resolve current executable")?; + let unit = format!( + "[Unit]\n\ + Description=IronClaw daemon\n\ + After=network.target\n\ + \n\ + [Service]\n\ + Type=simple\n\ + ExecStart=\"{exe}\" run\n\ + Restart=always\n\ + RestartSec=3\n\ + \n\ + [Install]\n\ + WantedBy=default.target\n", + exe = exe.display(), + ); + + std::fs::write(&file, unit)?; + run_checked(Command::new("systemctl").args(["--user", "daemon-reload"])).ok(); + run_checked(Command::new("systemctl").args(["--user", "enable", SYSTEMD_UNIT])).ok(); + println!("Installed systemd user service: {}", file.display()); + println!(" Start with: ironclaw service start"); + Ok(()) +} + +// ── Start ─────────────────────────────────────────────────────── + +fn start() -> Result<()> { + if cfg!(target_os = "macos") { + let plist = macos_plist_path()?; + if !plist.exists() { + bail!("Service not installed. Run `ironclaw service install` first."); + } + run_checked(Command::new("launchctl").arg("load").arg("-w").arg(&plist))?; + run_checked(Command::new("launchctl").arg("start").arg(SERVICE_LABEL))?; + println!("Service started"); + Ok(()) + } else if cfg!(target_os = "linux") { + run_checked(Command::new("systemctl").args(["--user", "daemon-reload"]))?; + run_checked(Command::new("systemctl").args(["--user", "start", SYSTEMD_UNIT]))?; + println!("Service started"); + Ok(()) + } else { + bail!("Service management is only supported on macOS and Linux"); + } +} + +// ── Stop ──────────────────────────────────────────────────────── + +fn stop() -> Result<()> { + if cfg!(target_os = "macos") { + let plist = macos_plist_path()?; + run_checked(Command::new("launchctl").arg("stop").arg(SERVICE_LABEL)).ok(); + run_checked( + Command::new("launchctl") + .arg("unload") + .arg("-w") + .arg(&plist), + ) + .ok(); + println!("Service stopped"); + Ok(()) + } else if cfg!(target_os = "linux") { + run_checked(Command::new("systemctl").args(["--user", "stop", SYSTEMD_UNIT])).ok(); + println!("Service stopped"); + Ok(()) + } else { + bail!("Service management is only supported on macOS and Linux"); + } +} + +// ── Status ────────────────────────────────────────────────────── + +fn status() -> Result<()> { + if cfg!(target_os = "macos") { + let out = run_capture(Command::new("launchctl").arg("list"))?; + let running = out.lines().any(|line| line.contains(SERVICE_LABEL)); + println!( + "Service: {}", + if running { + "running/loaded" + } else { + "not loaded" + } + ); + println!("Unit: {}", macos_plist_path()?.display()); + Ok(()) + } else if cfg!(target_os = "linux") { + let state = + run_capture(Command::new("systemctl").args(["--user", "is-active", SYSTEMD_UNIT])) + .unwrap_or_else(|_| "unknown".into()); + println!("Service state: {}", state.trim()); + println!("Unit: {}", linux_unit_path()?.display()); + Ok(()) + } else { + bail!("Service management is only supported on macOS and Linux"); + } +} + +// ── Uninstall ─────────────────────────────────────────────────── + +fn uninstall() -> Result<()> { + // Stop first (ignore errors, service might not be running) + stop().ok(); + + if cfg!(target_os = "macos") { + let file = macos_plist_path()?; + if file.exists() { + std::fs::remove_file(&file) + .with_context(|| format!("failed to remove {}", file.display()))?; + } + println!("Service uninstalled ({})", file.display()); + Ok(()) + } else if cfg!(target_os = "linux") { + let file = linux_unit_path()?; + if file.exists() { + std::fs::remove_file(&file) + .with_context(|| format!("failed to remove {}", file.display()))?; + } + run_checked(Command::new("systemctl").args(["--user", "daemon-reload"])).ok(); + println!("Service uninstalled ({})", file.display()); + Ok(()) + } else { + bail!("Service management is only supported on macOS and Linux"); + } +} + +// ── Path helpers ──────────────────────────────────────────────── + +fn macos_plist_path() -> Result { + let home = dirs::home_dir().context("could not find home directory")?; + Ok(home + .join("Library") + .join("LaunchAgents") + .join(format!("{SERVICE_LABEL}.plist"))) +} + +fn linux_unit_path() -> Result { + let home = dirs::home_dir().context("could not find home directory")?; + Ok(home + .join(".config") + .join("systemd") + .join("user") + .join(SYSTEMD_UNIT)) +} + +fn ironclaw_logs_dir() -> Result { + let home = dirs::home_dir().context("could not find home directory")?; + Ok(home.join(".ironclaw").join("logs")) +} + +// ── Shell helpers ─────────────────────────────────────────────── + +fn run_checked(command: &mut Command) -> Result<()> { + let output = command.output().context("failed to spawn command")?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + bail!("command failed: {}", stderr.trim()); + } + Ok(()) +} + +fn run_capture(command: &mut Command) -> Result { + let output = command.output().context("failed to spawn command")?; + let mut text = String::from_utf8_lossy(&output.stdout).to_string(); + if text.trim().is_empty() { + text = String::from_utf8_lossy(&output.stderr).to_string(); + } + Ok(text) +} + +fn xml_escape(raw: &str) -> String { + raw.replace('&', "&") + .replace('<', "<") + .replace('>', ">") + .replace('"', """) + .replace('\'', "'") +} + +// ── Tests ─────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use crate::service::*; + + #[test] + fn xml_escape_handles_reserved_chars() { + let escaped = xml_escape("<&>\"' and text"); + assert_eq!(escaped, "<&>"' and text"); + } + + #[test] + fn xml_escape_passes_through_plain_text() { + assert_eq!(xml_escape("hello world"), "hello world"); + } + + #[test] + fn run_capture_reads_stdout() { + let out = run_capture(Command::new("sh").args(["-c", "echo hello"])) + .expect("stdout capture should succeed"); + assert_eq!(out.trim(), "hello"); + } + + #[test] + fn run_capture_falls_back_to_stderr() { + let out = run_capture(Command::new("sh").args(["-c", "echo warn 1>&2"])) + .expect("stderr capture should succeed"); + assert_eq!(out.trim(), "warn"); + } + + #[test] + fn run_checked_errors_on_non_zero_exit() { + let err = run_checked(Command::new("sh").args(["-c", "exit 17"])) + .expect_err("non-zero exit should error"); + assert!(err.to_string().contains("command failed")); + } + + #[test] + fn run_checked_succeeds_on_zero_exit() { + assert!(run_checked(Command::new("sh").args(["-c", "exit 0"])).is_ok()); + } + + #[cfg(target_os = "macos")] + #[test] + fn macos_plist_path_has_expected_suffix() { + let path = macos_plist_path().unwrap(); + let s = path.to_string_lossy(); + assert!( + s.ends_with("Library/LaunchAgents/com.ironclaw.daemon.plist"), + "unexpected path: {s}" + ); + } + + #[cfg(target_os = "linux")] + #[test] + fn linux_unit_path_has_expected_suffix() { + let path = linux_unit_path().unwrap(); + let s = path.to_string_lossy(); + assert!( + s.ends_with(".config/systemd/user/ironclaw.service"), + "unexpected path: {s}" + ); + } + + #[test] + fn logs_dir_under_ironclaw() { + let path = ironclaw_logs_dir().unwrap(); + let s = path.to_string_lossy(); + assert!(s.ends_with(".ironclaw/logs"), "unexpected path: {s}"); + } +} diff --git a/src/settings.rs b/src/settings.rs index 60b14cd2..e8d1242f 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -149,11 +149,52 @@ impl Default for EmbeddingsSettings { /// Tunnel settings for public webhook endpoints. /// /// The tunnel URL is shared across all channels that need webhooks. +/// Two modes: +/// - **Static URL**: `public_url` set directly (manual tunnel management). +/// - **Managed provider**: `provider` is set and the agent starts/stops the +/// tunnel process automatically at boot/shutdown. #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct TunnelSettings { /// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io"). + /// When set without a provider, treated as a static (externally managed) URL. #[serde(default)] pub public_url: Option, + + /// Managed tunnel provider: "ngrok", "cloudflare", "tailscale", "custom". + #[serde(default)] + pub provider: Option, + + /// Cloudflare tunnel token. + #[serde(default)] + pub cf_token: Option, + + /// ngrok auth token. + #[serde(default)] + pub ngrok_token: Option, + + /// ngrok custom domain (paid plans). + #[serde(default)] + pub ngrok_domain: Option, + + /// Use Tailscale Funnel (public) instead of Serve (tailnet-only). + #[serde(default)] + pub ts_funnel: bool, + + /// Tailscale hostname override. + #[serde(default)] + pub ts_hostname: Option, + + /// Shell command for custom tunnel (with `{port}` / `{host}` placeholders). + #[serde(default)] + pub custom_command: Option, + + /// Health check URL for custom tunnel. + #[serde(default)] + pub custom_health_url: Option, + + /// Substring pattern to extract URL from custom tunnel stdout. + #[serde(default)] + pub custom_url_pattern: Option, } /// Channel-specific settings. @@ -585,6 +626,83 @@ impl Settings { } } + /// Default TOML config file path (~/.ironclaw/config.toml). + pub fn default_toml_path() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("config.toml") + } + + /// Load settings from a TOML file. + /// + /// Returns `None` if the file doesn't exist. Returns an error only + /// if the file exists but can't be parsed. + pub fn load_toml(path: &std::path::Path) -> Result, String> { + let data = match std::fs::read_to_string(path) { + Ok(d) => d, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(e) => return Err(format!("failed to read {}: {}", path.display(), e)), + }; + + let settings: Self = toml::from_str(&data) + .map_err(|e| format!("invalid TOML in {}: {}", path.display(), e))?; + Ok(Some(settings)) + } + + /// Write a well-commented TOML config file with current settings. + pub fn save_toml(&self, path: &std::path::Path) -> Result<(), String> { + let raw = toml::to_string_pretty(self) + .map_err(|e| format!("failed to serialize settings: {}", e))?; + + let content = format!( + "# IronClaw configuration file.\n\ + #\n\ + # Priority: env var > this file > database settings > defaults.\n\ + # Uncomment and edit values to override defaults.\n\ + # Run `ironclaw config init` to regenerate this file.\n\ + #\n\ + # Documentation: https://github.com/nearai/ironclaw\n\ + \n\ + {raw}" + ); + + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .map_err(|e| format!("failed to create {}: {}", parent.display(), e))?; + } + + std::fs::write(path, content) + .map_err(|e| format!("failed to write {}: {}", path.display(), e)) + } + + /// Merge values from `other` into `self`, preferring `other` for + /// fields that differ from the default. + /// + /// This enables layering: load DB/JSON settings as the base, then + /// overlay TOML values on top. Only fields that the TOML file + /// explicitly changed (i.e. differ from Default) are applied. + pub fn merge_from(&mut self, other: &Self) { + let default_json = match serde_json::to_value(Self::default()) { + Ok(v) => v, + Err(_) => return, + }; + let other_json = match serde_json::to_value(other) { + Ok(v) => v, + Err(_) => return, + }; + let mut self_json = match serde_json::to_value(&*self) { + Ok(v) => v, + Err(_) => return, + }; + + merge_non_default(&mut self_json, &other_json, &default_json); + + if let Ok(merged) = serde_json::from_value(self_json) { + *self = merged; + } + } + /// Get a setting value by dotted path (e.g., "agent.max_parallel_jobs"). pub fn get(&self, path: &str) -> Option { let json = serde_json::to_value(self).ok()?; @@ -766,9 +884,40 @@ fn collect_settings( } } +/// Recursively merge `other` into `target`, but only for fields where +/// `other` differs from `defaults`. This means only explicitly-set values +/// in the TOML file override the base settings. +fn merge_non_default( + target: &mut serde_json::Value, + other: &serde_json::Value, + defaults: &serde_json::Value, +) { + match (target, other, defaults) { + ( + serde_json::Value::Object(t), + serde_json::Value::Object(o), + serde_json::Value::Object(d), + ) => { + for (key, other_val) in o { + let default_val = d.get(key).cloned().unwrap_or(serde_json::Value::Null); + if let Some(target_val) = t.get_mut(key) { + merge_non_default(target_val, other_val, &default_val); + } else if other_val != &default_val { + t.insert(key.clone(), other_val.clone()); + } + } + } + (target, other, defaults) => { + if other != defaults { + *target = other.clone(); + } + } + } +} + #[cfg(test)] mod tests { - use super::*; + use crate::settings::*; #[test] fn test_db_map_round_trip() { @@ -904,4 +1053,152 @@ mod tests { Some("http://my-vllm:8000/v1".to_string()) ); } + + #[test] + fn toml_round_trip() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("config.toml"); + + let mut settings = Settings::default(); + settings.agent.name = "toml-bot".to_string(); + settings.heartbeat.enabled = true; + settings.heartbeat.interval_secs = 900; + + settings.save_toml(&path).unwrap(); + let loaded = Settings::load_toml(&path).unwrap().unwrap(); + + assert_eq!(loaded.agent.name, "toml-bot"); + assert!(loaded.heartbeat.enabled); + assert_eq!(loaded.heartbeat.interval_secs, 900); + } + + #[test] + fn toml_missing_file_returns_none() { + let result = Settings::load_toml(std::path::Path::new("/tmp/nonexistent_config.toml")); + assert!(result.unwrap().is_none()); + } + + #[test] + fn toml_invalid_content_returns_error() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("bad.toml"); + std::fs::write(&path, "this is not valid toml [[[").unwrap(); + + let result = Settings::load_toml(&path); + assert!(result.is_err()); + } + + #[test] + fn toml_partial_config_uses_defaults() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("partial.toml"); + + // Only set agent name, everything else should be default + std::fs::write(&path, "[agent]\nname = \"partial-bot\"\n").unwrap(); + + let loaded = Settings::load_toml(&path).unwrap().unwrap(); + assert_eq!(loaded.agent.name, "partial-bot"); + // Defaults preserved + assert_eq!(loaded.agent.max_parallel_jobs, 5); + assert!(!loaded.heartbeat.enabled); + } + + #[test] + fn toml_header_comment_present() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("config.toml"); + + Settings::default().save_toml(&path).unwrap(); + let content = std::fs::read_to_string(&path).unwrap(); + + assert!(content.starts_with("# IronClaw configuration file.")); + assert!(content.contains("[agent]")); + assert!(content.contains("[heartbeat]")); + } + + #[test] + fn merge_only_overrides_non_default_values() { + let mut base = Settings::default(); + base.agent.name = "from-db".to_string(); + base.heartbeat.interval_secs = 600; + + 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); + + assert_eq!(base.agent.name, "from-toml"); + assert_eq!(base.heartbeat.interval_secs, 600); + } + + #[test] + fn merge_preserves_base_when_overlay_is_default() { + let mut base = Settings::default(); + base.agent.name = "custom-name".to_string(); + base.heartbeat.enabled = true; + + 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); + } + + #[test] + fn toml_creates_parent_dirs() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("nested").join("deep").join("config.toml"); + + Settings::default().save_toml(&path).unwrap(); + assert!(path.exists()); + } + + #[test] + fn default_toml_path_under_ironclaw() { + let path = Settings::default_toml_path(); + assert!(path.to_string_lossy().contains(".ironclaw")); + assert!(path.to_string_lossy().ends_with("config.toml")); + } + + #[test] + fn tunnel_settings_round_trip() { + let settings = Settings { + tunnel: TunnelSettings { + provider: Some("ngrok".to_string()), + ngrok_token: Some("tok_abc123".to_string()), + ngrok_domain: Some("my.ngrok.dev".to_string()), + ..Default::default() + }, + ..Default::default() + }; + + // JSON round-trip + let json = serde_json::to_string(&settings).unwrap(); + let restored: Settings = serde_json::from_str(&json).unwrap(); + assert_eq!(restored.tunnel.provider, Some("ngrok".to_string())); + assert_eq!(restored.tunnel.ngrok_token, Some("tok_abc123".to_string())); + assert_eq!( + restored.tunnel.ngrok_domain, + Some("my.ngrok.dev".to_string()) + ); + assert!(restored.tunnel.public_url.is_none()); + + // DB map round-trip + let map = settings.to_db_map(); + let from_db = Settings::from_db_map(&map); + assert_eq!(from_db.tunnel.provider, Some("ngrok".to_string())); + assert_eq!(from_db.tunnel.ngrok_token, Some("tok_abc123".to_string())); + + // get/set round-trip + let mut s = Settings::default(); + s.set("tunnel.provider", "cloudflare").unwrap(); + s.set("tunnel.cf_token", "cf_tok_xyz").unwrap(); + s.set("tunnel.ts_funnel", "true").unwrap(); + assert_eq!(s.tunnel.provider, Some("cloudflare".to_string())); + assert_eq!(s.tunnel.cf_token, Some("cf_tok_xyz".to_string())); + assert!(s.tunnel.ts_funnel); + } } diff --git a/src/setup/channels.rs b/src/setup/channels.rs index 34811358..5b0f66bf 100644 --- a/src/setup/channels.rs +++ b/src/setup/channels.rs @@ -18,6 +18,7 @@ use crate::secrets::{CreateSecretParams, SecretsStore}; use crate::settings::{Settings, TunnelSettings}; use crate::setup::prompts::{ confirm, input, optional_input, print_error, print_info, print_success, secret_input, + select_one, }; /// Typed errors for channel setup flows. @@ -356,12 +357,20 @@ async fn bind_telegram_owner_flow( /// Set up a tunnel for exposing the agent to the internet. /// /// This is shared across all channels that need webhook endpoints. -/// Returns the tunnel URL if configured. -pub fn setup_tunnel(settings: &Settings) -> Result, ChannelSetupError> { - if let Some(ref url) = settings.tunnel.public_url { - print_info(&format!("Existing tunnel configured: {}", url)); +/// Returns a `TunnelSettings` with provider config (managed tunnel) +/// or a static URL. +pub fn setup_tunnel(settings: &Settings) -> Result { + // Show existing config + let has_existing = settings.tunnel.public_url.is_some() || settings.tunnel.provider.is_some(); + if has_existing { + if let Some(ref url) = settings.tunnel.public_url { + print_info(&format!("Existing static tunnel URL: {}", url)); + } + if let Some(ref provider) = settings.tunnel.provider { + print_info(&format!("Existing managed provider: {}", provider)); + } if !confirm("Change tunnel configuration?", false)? { - return Ok(Some(url.clone())); + return Ok(settings.tunnel.clone()); } } @@ -369,24 +378,121 @@ pub fn setup_tunnel(settings: &Settings) -> Result, ChannelSetupE print_info("Tunnel Configuration (for webhook endpoints):"); print_info("A tunnel exposes your local agent to the internet, enabling:"); print_info(" - Instant Telegram message delivery (instead of polling)"); - print_info(" - Future: Slack, Discord, GitHub webhooks"); - print_info(""); - print_info("Supported tunnel providers:"); - print_info(" - ngrok: ngrok http 8080"); - print_info(" - Cloudflare: cloudflared tunnel --url http://localhost:8080"); - print_info(" - localtunnel: lt --port 8080"); - print_info(""); - print_info("Security note: Webhook endpoints don't use tunnel-level auth."); - print_info("Security comes from provider-specific secrets (e.g., Telegram webhook secret)."); + print_info(" - Slack, Discord, GitHub webhooks"); println!(); if !confirm("Configure a tunnel?", false)? { - return Ok(None); + return Ok(TunnelSettings::default()); } + let options = &[ + "ngrok - managed tunnel, starts automatically", + "Cloudflare - cloudflared tunnel, starts automatically", + "Tailscale - Tailscale Funnel/Serve, starts automatically", + "Custom - your own tunnel command", + "Static URL - you manage the tunnel yourself", + ]; + + let choice = select_one("Select tunnel provider:", options)?; + + match choice { + 0 => setup_tunnel_ngrok(), + 1 => setup_tunnel_cloudflare(), + 2 => setup_tunnel_tailscale(), + 3 => setup_tunnel_custom(), + 4 => setup_tunnel_static(), + _ => Ok(TunnelSettings::default()), + } +} + +fn setup_tunnel_ngrok() -> Result { + print_info("Get your auth token from: https://dashboard.ngrok.com/get-started/your-authtoken"); + println!(); + + let token = secret_input("ngrok auth token")?; + let domain = optional_input("Custom domain", Some("leave empty for auto-assigned"))?; + + print_success("ngrok configured. Tunnel will start automatically at boot."); + + Ok(TunnelSettings { + provider: Some("ngrok".to_string()), + ngrok_token: Some(token.expose_secret().to_string()), + ngrok_domain: domain, + ..Default::default() + }) +} + +fn setup_tunnel_cloudflare() -> Result { + print_info("Get your tunnel token from the Cloudflare Zero Trust dashboard:"); + print_info(" https://one.dash.cloudflare.com/ > Networks > Tunnels"); + println!(); + + let token = secret_input("Cloudflare tunnel token")?; + + print_success("Cloudflare tunnel configured. Tunnel will start automatically at boot."); + + Ok(TunnelSettings { + provider: Some("cloudflare".to_string()), + cf_token: Some(token.expose_secret().to_string()), + ..Default::default() + }) +} + +fn setup_tunnel_tailscale() -> Result { + let funnel = confirm("Use Tailscale Funnel (public internet)?", true)?; + let hostname = optional_input("Hostname override", Some("leave empty for auto-detect"))?; + + let mode = if funnel { + "Funnel (public)" + } else { + "Serve (tailnet-only)" + }; + print_success(&format!("Tailscale {} configured.", mode)); + + Ok(TunnelSettings { + provider: Some("tailscale".to_string()), + ts_funnel: funnel, + ts_hostname: hostname, + ..Default::default() + }) +} + +fn setup_tunnel_custom() -> Result { + print_info("Enter a shell command to start your tunnel."); + print_info("Use {port} and {host} as placeholders."); + print_info("Example: bore local {port} --to bore.pub"); + println!(); + + let command = input("Tunnel command")?; + if command.is_empty() { + return Err(ChannelSetupError::Validation( + "Tunnel command cannot be empty".to_string(), + )); + } + + let health_url = optional_input("Health check URL", Some("optional"))?; + let url_pattern = optional_input( + "URL pattern (substring to match in stdout)", + Some("optional"), + )?; + + print_success("Custom tunnel configured."); + + Ok(TunnelSettings { + provider: Some("custom".to_string()), + custom_command: Some(command), + custom_health_url: health_url, + custom_url_pattern: url_pattern, + ..Default::default() + }) +} + +fn setup_tunnel_static() -> Result { + print_info("Enter the public URL of your externally managed tunnel."); + println!(); + let tunnel_url = input("Tunnel URL (e.g., https://abc123.ngrok.io)")?; - // Validate URL format if !tunnel_url.starts_with("https://") { print_error("URL must start with https:// (webhooks require HTTPS)"); return Err(ChannelSetupError::Validation( @@ -394,15 +500,15 @@ pub fn setup_tunnel(settings: &Settings) -> Result, ChannelSetupE )); } - // Remove trailing slash if present let tunnel_url = tunnel_url.trim_end_matches('/').to_string(); - print_success(&format!("Tunnel URL configured: {}", tunnel_url)); - print_info(""); + print_success(&format!("Static tunnel URL configured: {}", tunnel_url)); print_info("Make sure your tunnel is running before starting the agent."); - print_info("You can also set TUNNEL_URL environment variable to override."); - Ok(Some(tunnel_url)) + Ok(TunnelSettings { + public_url: Some(tunnel_url), + ..Default::default() + }) } /// Set up Telegram webhook secret for signature validation. diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index ff4e0566..718989a7 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1022,6 +1022,11 @@ impl SetupWizard { api_key: None, fallback_model: None, max_retries: 3, + 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, }, @@ -1254,11 +1259,8 @@ impl SetupWizard { async fn step_channels(&mut self) -> Result<(), SetupError> { // First, configure tunnel (shared across all channels that need webhooks) match setup_tunnel(&self.settings) { - Ok(Some(url)) => { - self.settings.tunnel.public_url = Some(url); - } - Ok(None) => { - self.settings.tunnel.public_url = None; + Ok(tunnel_settings) => { + self.settings.tunnel = tunnel_settings; } Err(e) => { print_info(&format!("Tunnel setup skipped: {}", e)); @@ -1609,9 +1611,14 @@ impl SetupWizard { } if let Some(ref tunnel_url) = self.settings.tunnel.public_url { - println!(" Tunnel: {}", tunnel_url); + println!(" Tunnel: {} (static)", tunnel_url); + } else if let Some(ref provider) = self.settings.tunnel.provider { + println!(" Tunnel: {} (managed, starts at boot)", provider); } + let has_tunnel = + self.settings.tunnel.public_url.is_some() || self.settings.tunnel.provider.is_some(); + println!(" Channels:"); println!(" - CLI/TUI: enabled"); @@ -1621,11 +1628,7 @@ impl SetupWizard { } for channel_name in &self.settings.channels.wasm_channels { - let mode = if self.settings.tunnel.public_url.is_some() { - "webhook" - } else { - "polling" - }; + let mode = if has_tunnel { "webhook" } else { "polling" }; println!( " - {}: enabled ({})", capitalize_first(channel_name), diff --git a/src/tunnel/cloudflare.rs b/src/tunnel/cloudflare.rs new file mode 100644 index 00000000..38f0cd97 --- /dev/null +++ b/src/tunnel/cloudflare.rs @@ -0,0 +1,140 @@ +//! Cloudflare Tunnel via the `cloudflared` binary. + +use anyhow::{Result, bail}; +use tokio::io::AsyncBufReadExt; +use tokio::process::Command; + +use crate::tunnel::{ + SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process, + new_shared_url, +}; + +/// Wraps `cloudflared` with token-based auth from the Zero Trust dashboard. +pub struct CloudflareTunnel { + token: String, + proc: SharedProcess, + url: SharedUrl, +} + +impl CloudflareTunnel { + pub fn new(token: String) -> Self { + Self { + token, + proc: new_shared_process(), + url: new_shared_url(), + } + } +} + +#[async_trait::async_trait] +impl Tunnel for CloudflareTunnel { + fn name(&self) -> &str { + "cloudflare" + } + + async fn start(&self, local_host: &str, local_port: u16) -> Result { + let origin = format!("http://{local_host}:{local_port}"); + let mut child = Command::new("cloudflared") + .args([ + "tunnel", + "--no-autoupdate", + "run", + "--token", + &self.token, + "--url", + &origin, + ]) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn()?; + + // cloudflared prints the public URL on stderr + let stderr = child + .stderr + .take() + .ok_or_else(|| anyhow::anyhow!("Failed to capture cloudflared stderr"))?; + + let mut reader = tokio::io::BufReader::new(stderr).lines(); + let mut public_url = String::new(); + + let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(30); + while tokio::time::Instant::now() < deadline { + let line = + tokio::time::timeout(tokio::time::Duration::from_secs(5), reader.next_line()).await; + + match line { + Ok(Ok(Some(l))) => { + tracing::debug!("cloudflared: {l}"); + if let Some(idx) = l.find("https://") { + let url_part = &l[idx..]; + let end = url_part + .find(|c: char| c.is_whitespace()) + .unwrap_or(url_part.len()); + public_url = url_part[..end].to_string(); + break; + } + } + Ok(Ok(None)) => break, + Ok(Err(e)) => bail!("Error reading cloudflared output: {e}"), + Err(_) => {} // line timeout, keep waiting + } + } + + if public_url.is_empty() { + child.kill().await.ok(); + bail!("cloudflared did not produce a public URL within 30s. Is the token valid?"); + } + + if let Ok(mut guard) = self.url.write() { + *guard = Some(public_url.clone()); + } + + let mut guard = self.proc.lock().await; + *guard = Some(TunnelProcess { child }); + + Ok(public_url) + } + + async fn stop(&self) -> Result<()> { + if let Ok(mut guard) = self.url.write() { + *guard = None; + } + kill_shared(&self.proc).await + } + + async fn health_check(&self) -> bool { + let guard = self.proc.lock().await; + guard.as_ref().is_some_and(|tp| tp.child.id().is_some()) + } + + fn public_url(&self) -> Option { + self.url.read().ok().and_then(|guard| guard.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn constructor_stores_token() { + let tunnel = CloudflareTunnel::new("cf-token".into()); + assert_eq!(tunnel.token, "cf-token"); + } + + #[test] + fn public_url_none_before_start() { + assert!(CloudflareTunnel::new("tok".into()).public_url().is_none()); + } + + #[tokio::test] + async fn stop_without_start_is_ok() { + assert!(CloudflareTunnel::new("tok".into()).stop().await.is_ok()); + } + + #[tokio::test] + async fn health_false_before_start() { + assert!(!CloudflareTunnel::new("tok".into()).health_check().await); + } +} diff --git a/src/tunnel/custom.rs b/src/tunnel/custom.rs new file mode 100644 index 00000000..888cb698 --- /dev/null +++ b/src/tunnel/custom.rs @@ -0,0 +1,245 @@ +//! Custom tunnel via an arbitrary shell command. + +use anyhow::{Result, bail}; +use tokio::io::AsyncBufReadExt; +use tokio::process::Command; + +use crate::tunnel::{ + SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process, + new_shared_url, +}; + +/// Bring-your-own tunnel binary. +/// +/// `start_command` supports `{port}` and `{host}` placeholders. +/// If `url_pattern` is set, stdout is scanned for a URL matching that +/// substring. If `health_url` is set, health checks poll that endpoint. +/// +/// **Note:** The command is split on whitespace, so quoted arguments like +/// `--arg "hello world"` won't work. Each token must be a single word. +/// +/// Examples: +/// - `bore local {port} --to bore.pub` +/// - `ssh -R 80:localhost:{port} serveo.net` +pub struct CustomTunnel { + start_command: String, + health_url: Option, + url_pattern: Option, + proc: SharedProcess, + url: SharedUrl, +} + +impl CustomTunnel { + pub fn new( + start_command: String, + health_url: Option, + url_pattern: Option, + ) -> Self { + Self { + start_command, + health_url, + url_pattern, + proc: new_shared_process(), + url: new_shared_url(), + } + } +} + +#[async_trait::async_trait] +impl Tunnel for CustomTunnel { + fn name(&self) -> &str { + "custom" + } + + async fn start(&self, local_host: &str, local_port: u16) -> Result { + let cmd = self + .start_command + .replace("{port}", &local_port.to_string()) + .replace("{host}", local_host); + + let parts: Vec<&str> = cmd.split_whitespace().collect(); + if parts.is_empty() { + bail!("Custom tunnel start_command is empty"); + } + + let mut child = Command::new(parts[0]) + .args(&parts[1..]) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn()?; + + let mut public_url = format!("http://{local_host}:{local_port}"); + + if self.url_pattern.is_some() + && let Some(stdout) = child.stdout.take() + { + let mut reader = tokio::io::BufReader::new(stdout).lines(); + let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(15); + + while tokio::time::Instant::now() < deadline { + let line = + tokio::time::timeout(tokio::time::Duration::from_secs(3), reader.next_line()) + .await; + + match line { + Ok(Ok(Some(l))) => { + tracing::debug!("custom-tunnel: {l}"); + if let Some(url) = extract_url(&l) { + let matches_pattern = self + .url_pattern + .as_ref() + .is_none_or(|pat| url.contains(pat.as_str())); + if matches_pattern { + public_url = url; + break; + } + } + } + Ok(Ok(None) | Err(_)) => break, + Err(_) => {} + } + } + } + + if let Ok(mut guard) = self.url.write() { + *guard = Some(public_url.clone()); + } + + let mut guard = self.proc.lock().await; + *guard = Some(TunnelProcess { child }); + + Ok(public_url) + } + + async fn stop(&self) -> Result<()> { + if let Ok(mut guard) = self.url.write() { + *guard = None; + } + kill_shared(&self.proc).await + } + + async fn health_check(&self) -> bool { + if let Some(ref url) = self.health_url { + return reqwest::Client::new() + .get(url) + .timeout(std::time::Duration::from_secs(5)) + .send() + .await + .is_ok(); + } + + let guard = self.proc.lock().await; + guard.as_ref().is_some_and(|tp| tp.child.id().is_some()) + } + + fn public_url(&self) -> Option { + self.url.read().ok().and_then(|guard| guard.clone()) + } +} + +/// Extract the first `https://` or `http://` URL from a line of text. +fn extract_url(line: &str) -> Option { + let idx = line.find("https://").or_else(|| line.find("http://"))?; + let url_part = &line[idx..]; + let end = url_part + .find(|c: char| c.is_whitespace()) + .unwrap_or(url_part.len()); + Some(url_part[..end].to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn empty_command_returns_error() { + let tunnel = CustomTunnel::new(" ".into(), None, None); + let result = tunnel.start("127.0.0.1", 8080).await; + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("start_command is empty") + ); + } + + #[tokio::test] + async fn start_without_pattern_returns_local() { + let tunnel = CustomTunnel::new("sleep 1".into(), None, None); + let url = tunnel.start("127.0.0.1", 4455).await.unwrap(); + assert_eq!(url, "http://127.0.0.1:4455"); + tunnel.stop().await.unwrap(); + } + + #[tokio::test] + async fn start_with_pattern_extracts_url() { + let tunnel = CustomTunnel::new( + "echo https://public.example".into(), + None, + Some("public.example".into()), + ); + let url = tunnel.start("localhost", 9999).await.unwrap(); + assert_eq!(url, "https://public.example"); + tunnel.stop().await.unwrap(); + } + + #[tokio::test] + async fn pattern_filters_non_matching_urls() { + // The command outputs two lines: first a non-matching URL, then a matching one. + // The pattern filter should skip the first and grab the second. + // No shell quoting needed; Command passes args directly to the binary. + let tunnel = CustomTunnel::new( + r"printf http://internal:1234\nhttps://real.tunnel.io/abc\n".into(), + None, + Some("tunnel.io".into()), + ); + let url = tunnel.start("localhost", 9999).await.unwrap(); + assert_eq!(url, "https://real.tunnel.io/abc"); + tunnel.stop().await.unwrap(); + } + + #[tokio::test] + async fn replaces_host_and_port_placeholders() { + let tunnel = CustomTunnel::new( + "echo http://{host}:{port}".into(), + None, + Some("http://".into()), + ); + let url = tunnel.start("10.1.2.3", 4321).await.unwrap(); + assert_eq!(url, "http://10.1.2.3:4321"); + tunnel.stop().await.unwrap(); + } + + #[tokio::test] + async fn health_with_unreachable_url_is_false() { + let tunnel = CustomTunnel::new( + "sleep 1".into(), + Some("http://127.0.0.1:9/healthz".into()), + None, + ); + assert!(!tunnel.health_check().await); + } + + #[test] + fn extract_url_finds_https() { + assert_eq!( + extract_url("tunnel ready at https://foo.bar.com/path more text"), + Some("https://foo.bar.com/path".to_string()) + ); + } + + #[test] + fn extract_url_finds_http() { + assert_eq!( + extract_url("url=http://localhost:8080"), + Some("http://localhost:8080".to_string()) + ); + } + + #[test] + fn extract_url_none_when_absent() { + assert_eq!(extract_url("no url here"), None); + } +} diff --git a/src/tunnel/mod.rs b/src/tunnel/mod.rs new file mode 100644 index 00000000..5551f2ed --- /dev/null +++ b/src/tunnel/mod.rs @@ -0,0 +1,329 @@ +//! Tunnel abstraction for exposing the agent to the internet. +//! +//! Wraps external tunnel binaries (cloudflared, ngrok, tailscale, etc.) behind +//! a common trait. The gateway starts a tunnel after binding its local port +//! and stops it on shutdown. +//! +//! Supported providers: +//! - **cloudflare** - Zero Trust tunnels via `cloudflared` +//! - **tailscale** - `tailscale serve` (tailnet) or `tailscale funnel` (public) +//! - **ngrok** - instant public URLs via `ngrok` +//! - **custom** - any command with `{host}`/`{port}` placeholders +//! - **none** - local-only, no external exposure + +mod cloudflare; +mod custom; +mod ngrok; +mod none; +mod tailscale; + +pub use cloudflare::CloudflareTunnel; +pub use custom::CustomTunnel; +pub use ngrok::NgrokTunnel; +pub use none::NoneTunnel; +pub use tailscale::TailscaleTunnel; + +use std::sync::Arc; + +use anyhow::{Result, bail}; +use tokio::sync::Mutex; + +/// Lock-free URL storage. Uses `std::sync::RwLock` so `public_url()` (sync) +/// never returns a spurious `None` due to async lock contention. +pub(crate) type SharedUrl = Arc>>; + +pub(crate) fn new_shared_url() -> SharedUrl { + Arc::new(std::sync::RwLock::new(None)) +} + +// ── Tunnel trait ───────────────────────────────────────────────── + +/// Provider-agnostic tunnel with lifecycle management. +/// +/// Implementations wrap an external tunnel binary. The gateway calls +/// `start()` after binding its local port and `stop()` on shutdown. +#[async_trait::async_trait] +pub trait Tunnel: Send + Sync { + /// Human-readable provider name (e.g. "cloudflare", "tailscale"). + fn name(&self) -> &str; + + /// Start the tunnel exposing `local_host:local_port` externally. + /// Returns the public URL on success. + async fn start(&self, local_host: &str, local_port: u16) -> Result; + + /// Stop the tunnel process gracefully. + async fn stop(&self) -> Result<()>; + + /// Check if the tunnel process is still alive. + async fn health_check(&self) -> bool; + + /// Return the public URL if the tunnel is running, `None` otherwise. + fn public_url(&self) -> Option; +} + +// ── Shared child-process handle ────────────────────────────────── + +/// Wraps a spawned tunnel child process. +pub(crate) struct TunnelProcess { + pub child: tokio::process::Child, +} + +pub(crate) type SharedProcess = Arc>>; + +pub(crate) fn new_shared_process() -> SharedProcess { + Arc::new(Mutex::new(None)) +} + +/// Kill a shared tunnel process if running. +pub(crate) async fn kill_shared(proc: &SharedProcess) -> Result<()> { + let mut guard = proc.lock().await; + if let Some(ref mut tp) = *guard { + tp.child.kill().await.ok(); + tp.child.wait().await.ok(); + } + *guard = None; + Ok(()) +} + +// ── Configuration types ────────────────────────────────────────── + +/// Provider-specific config for Cloudflare tunnels. +#[derive(Debug, Clone, Default)] +pub struct CloudflareTunnelConfig { + /// Token from the Cloudflare Zero Trust dashboard. + pub token: String, +} + +/// Provider-specific config for Tailscale tunnels. +#[derive(Debug, Clone, Default)] +pub struct TailscaleTunnelConfig { + /// Use `tailscale funnel` (public) instead of `tailscale serve` (tailnet). + pub funnel: bool, + /// Override the hostname (default: auto-detect from `tailscale status`). + pub hostname: Option, +} + +/// Provider-specific config for ngrok tunnels. +#[derive(Debug, Clone, Default)] +pub struct NgrokTunnelConfig { + /// ngrok auth token (required). + pub auth_token: String, + /// Custom domain (requires ngrok paid plan). + pub domain: Option, +} + +/// Provider-specific config for custom tunnel commands. +#[derive(Debug, Clone, Default)] +pub struct CustomTunnelConfig { + /// Shell command with `{port}` and `{host}` placeholders. + pub start_command: String, + /// HTTP endpoint to poll for health checks. + pub health_url: Option, + /// Substring to match in stdout for URL extraction. + pub url_pattern: Option, +} + +/// Full tunnel configuration. +#[derive(Debug, Clone, Default)] +pub struct TunnelProviderConfig { + /// Provider name: "none", "cloudflare", "tailscale", "ngrok", "custom". + pub provider: String, + pub cloudflare: Option, + pub tailscale: Option, + pub ngrok: Option, + pub custom: Option, +} + +// ── Factory ────────────────────────────────────────────────────── + +/// Create a tunnel from config. Returns `None` for provider "none" or empty. +pub fn create_tunnel(config: &TunnelProviderConfig) -> Result>> { + match config.provider.as_str() { + "none" | "" => Ok(None), + + "cloudflare" => { + let cf = config.cloudflare.as_ref().ok_or_else(|| { + anyhow::anyhow!("TUNNEL_PROVIDER=cloudflare but no TUNNEL_CF_TOKEN configured") + })?; + Ok(Some(Box::new(CloudflareTunnel::new(cf.token.clone())))) + } + + "tailscale" => { + let ts = config.tailscale.as_ref().cloned().unwrap_or_default(); + Ok(Some(Box::new(TailscaleTunnel::new(ts.funnel, ts.hostname)))) + } + + "ngrok" => { + let ng = config.ngrok.as_ref().ok_or_else(|| { + anyhow::anyhow!("TUNNEL_PROVIDER=ngrok but no TUNNEL_NGROK_TOKEN configured") + })?; + Ok(Some(Box::new(NgrokTunnel::new( + ng.auth_token.clone(), + ng.domain.clone(), + )))) + } + + "custom" => { + let cu = config.custom.as_ref().ok_or_else(|| { + anyhow::anyhow!("TUNNEL_PROVIDER=custom but no TUNNEL_CUSTOM_COMMAND configured") + })?; + Ok(Some(Box::new(CustomTunnel::new( + cu.start_command.clone(), + cu.health_url.clone(), + cu.url_pattern.clone(), + )))) + } + + other => bail!( + "Unknown tunnel provider: \"{other}\". Valid: none, cloudflare, tailscale, ngrok, custom" + ), + } +} + +// ── Tests ──────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use tokio::process::Command; + + fn assert_tunnel_err(cfg: &TunnelProviderConfig, needle: &str) { + match create_tunnel(cfg) { + Err(e) => assert!( + e.to_string().contains(needle), + "Expected error containing \"{needle}\", got: {e}" + ), + Ok(_) => panic!("Expected error containing \"{needle}\", but got Ok"), + } + } + + #[test] + fn factory_none_returns_none() { + let cfg = TunnelProviderConfig::default(); + assert!(create_tunnel(&cfg).unwrap().is_none()); + } + + #[test] + fn factory_empty_returns_none() { + let cfg = TunnelProviderConfig { + provider: String::new(), + ..Default::default() + }; + assert!(create_tunnel(&cfg).unwrap().is_none()); + } + + #[test] + fn factory_unknown_provider_errors() { + let cfg = TunnelProviderConfig { + provider: "wireguard".into(), + ..Default::default() + }; + assert_tunnel_err(&cfg, "Unknown tunnel provider"); + } + + #[test] + fn factory_cloudflare_missing_config_errors() { + let cfg = TunnelProviderConfig { + provider: "cloudflare".into(), + ..Default::default() + }; + assert_tunnel_err(&cfg, "TUNNEL_CF_TOKEN"); + } + + #[test] + fn factory_cloudflare_with_config_ok() { + let cfg = TunnelProviderConfig { + provider: "cloudflare".into(), + cloudflare: Some(CloudflareTunnelConfig { + token: "test-token".into(), + }), + ..Default::default() + }; + let t = create_tunnel(&cfg).unwrap().unwrap(); + assert_eq!(t.name(), "cloudflare"); + } + + #[test] + fn factory_tailscale_defaults_ok() { + let cfg = TunnelProviderConfig { + provider: "tailscale".into(), + ..Default::default() + }; + let t = create_tunnel(&cfg).unwrap().unwrap(); + assert_eq!(t.name(), "tailscale"); + } + + #[test] + fn factory_ngrok_missing_config_errors() { + let cfg = TunnelProviderConfig { + provider: "ngrok".into(), + ..Default::default() + }; + assert_tunnel_err(&cfg, "TUNNEL_NGROK_TOKEN"); + } + + #[test] + fn factory_ngrok_with_config_ok() { + let cfg = TunnelProviderConfig { + provider: "ngrok".into(), + ngrok: Some(NgrokTunnelConfig { + auth_token: "tok".into(), + domain: None, + }), + ..Default::default() + }; + let t = create_tunnel(&cfg).unwrap().unwrap(); + assert_eq!(t.name(), "ngrok"); + } + + #[test] + fn factory_custom_missing_config_errors() { + let cfg = TunnelProviderConfig { + provider: "custom".into(), + ..Default::default() + }; + assert_tunnel_err(&cfg, "TUNNEL_CUSTOM_COMMAND"); + } + + #[test] + fn factory_custom_with_config_ok() { + let cfg = TunnelProviderConfig { + provider: "custom".into(), + custom: Some(CustomTunnelConfig { + start_command: "echo tunnel".into(), + health_url: None, + url_pattern: None, + }), + ..Default::default() + }; + let t = create_tunnel(&cfg).unwrap().unwrap(); + assert_eq!(t.name(), "custom"); + } + + #[tokio::test] + async fn kill_shared_no_process_is_ok() { + let proc = new_shared_process(); + assert!(kill_shared(&proc).await.is_ok()); + assert!(proc.lock().await.is_none()); + } + + #[tokio::test] + async fn kill_shared_terminates_child() { + let proc = new_shared_process(); + + let child = Command::new("sleep") + .arg("30") + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .spawn() + .expect("sleep should spawn"); + + { + let mut guard = proc.lock().await; + *guard = Some(TunnelProcess { child }); + } + + kill_shared(&proc).await.unwrap(); + assert!(proc.lock().await.is_none()); + } +} diff --git a/src/tunnel/ngrok.rs b/src/tunnel/ngrok.rs new file mode 100644 index 00000000..2b0e0df9 --- /dev/null +++ b/src/tunnel/ngrok.rs @@ -0,0 +1,142 @@ +//! ngrok tunnel via the `ngrok` binary. + +use anyhow::{Result, bail}; +use tokio::io::AsyncBufReadExt; +use tokio::process::Command; + +use crate::tunnel::{ + SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process, + new_shared_url, +}; + +/// Wraps `ngrok` with optional custom domain support (paid plan). +pub struct NgrokTunnel { + auth_token: String, + domain: Option, + proc: SharedProcess, + url: SharedUrl, +} + +impl NgrokTunnel { + pub fn new(auth_token: String, domain: Option) -> Self { + Self { + auth_token, + domain, + proc: new_shared_process(), + url: new_shared_url(), + } + } +} + +#[async_trait::async_trait] +impl Tunnel for NgrokTunnel { + fn name(&self) -> &str { + "ngrok" + } + + async fn start(&self, local_host: &str, local_port: u16) -> Result { + let mut args = vec!["http".to_string(), format!("{local_host}:{local_port}")]; + if let Some(ref domain) = self.domain { + args.push("--domain".into()); + args.push(domain.clone()); + } + args.extend(["--log", "stdout", "--log-format", "logfmt"].map(String::from)); + + let mut child = Command::new("ngrok") + .args(&args) + .env("NGROK_AUTHTOKEN", &self.auth_token) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn()?; + + let stdout = child + .stdout + .take() + .ok_or_else(|| anyhow::anyhow!("Failed to capture ngrok stdout"))?; + + let mut reader = tokio::io::BufReader::new(stdout).lines(); + let mut public_url = String::new(); + + let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(15); + while tokio::time::Instant::now() < deadline { + let line = + tokio::time::timeout(tokio::time::Duration::from_secs(3), reader.next_line()).await; + + match line { + Ok(Ok(Some(l))) => { + tracing::debug!("ngrok: {l}"); + // ngrok logfmt: url=https://xxxx.ngrok-free.app + if let Some(idx) = l.find("url=https://") { + let url_start = idx + 4; // skip "url=" + let url_part = &l[url_start..]; + let end = url_part + .find(|c: char| c.is_whitespace()) + .unwrap_or(url_part.len()); + public_url = url_part[..end].to_string(); + break; + } + } + Ok(Ok(None)) => break, + Ok(Err(e)) => bail!("Error reading ngrok output: {e}"), + Err(_) => {} + } + } + + if public_url.is_empty() { + child.kill().await.ok(); + bail!("ngrok did not produce a public URL within 15s. Is the auth token valid?"); + } + + if let Ok(mut guard) = self.url.write() { + *guard = Some(public_url.clone()); + } + + let mut guard = self.proc.lock().await; + *guard = Some(TunnelProcess { child }); + + Ok(public_url) + } + + async fn stop(&self) -> Result<()> { + if let Ok(mut guard) = self.url.write() { + *guard = None; + } + kill_shared(&self.proc).await + } + + async fn health_check(&self) -> bool { + let guard = self.proc.lock().await; + guard.as_ref().is_some_and(|tp| tp.child.id().is_some()) + } + + fn public_url(&self) -> Option { + self.url.read().ok().and_then(|guard| guard.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn constructor_stores_domain() { + let tunnel = NgrokTunnel::new("tok".into(), Some("my.ngrok.app".into())); + assert_eq!(tunnel.domain.as_deref(), Some("my.ngrok.app")); + } + + #[test] + fn public_url_none_before_start() { + assert!(NgrokTunnel::new("tok".into(), None).public_url().is_none()); + } + + #[tokio::test] + async fn stop_without_start_is_ok() { + assert!(NgrokTunnel::new("tok".into(), None).stop().await.is_ok()); + } + + #[tokio::test] + async fn health_false_before_start() { + assert!(!NgrokTunnel::new("tok".into(), None).health_check().await); + } +} diff --git a/src/tunnel/none.rs b/src/tunnel/none.rs new file mode 100644 index 00000000..169c6948 --- /dev/null +++ b/src/tunnel/none.rs @@ -0,0 +1,62 @@ +//! No-op tunnel for local-only access. + +use anyhow::Result; + +use crate::tunnel::Tunnel; + +/// No-op tunnel, no external exposure. `public_url()` always returns `None`. +pub struct NoneTunnel; + +#[async_trait::async_trait] +impl Tunnel for NoneTunnel { + fn name(&self) -> &str { + "none" + } + + async fn start(&self, local_host: &str, local_port: u16) -> Result { + Ok(format!("http://{local_host}:{local_port}")) + } + + async fn stop(&self) -> Result<()> { + Ok(()) + } + + async fn health_check(&self) -> bool { + true + } + + fn public_url(&self) -> Option { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn name_is_none() { + assert_eq!(NoneTunnel.name(), "none"); + } + + #[tokio::test] + async fn start_returns_local_url() { + let url = NoneTunnel.start("127.0.0.1", 7788).await.unwrap(); + assert_eq!(url, "http://127.0.0.1:7788"); + } + + #[tokio::test] + async fn stop_is_noop() { + assert!(NoneTunnel.stop().await.is_ok()); + } + + #[tokio::test] + async fn health_is_always_true() { + assert!(NoneTunnel.health_check().await); + } + + #[test] + fn public_url_is_always_none() { + assert!(NoneTunnel.public_url().is_none()); + } +} diff --git a/src/tunnel/tailscale.rs b/src/tunnel/tailscale.rs new file mode 100644 index 00000000..04e35575 --- /dev/null +++ b/src/tunnel/tailscale.rs @@ -0,0 +1,140 @@ +//! Tailscale tunnel via `tailscale serve` or `tailscale funnel`. + +use anyhow::{Result, bail}; +use tokio::process::Command; + +use crate::tunnel::{ + SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process, + new_shared_url, +}; + +/// Uses `tailscale serve` (tailnet-only) or `tailscale funnel` (public). +/// +/// Requires Tailscale installed and authenticated (`tailscale up`). +pub struct TailscaleTunnel { + funnel: bool, + hostname: Option, + proc: SharedProcess, + url: SharedUrl, +} + +impl TailscaleTunnel { + pub fn new(funnel: bool, hostname: Option) -> Self { + Self { + funnel, + hostname, + proc: new_shared_process(), + url: new_shared_url(), + } + } +} + +#[async_trait::async_trait] +impl Tunnel for TailscaleTunnel { + fn name(&self) -> &str { + "tailscale" + } + + async fn start(&self, local_host: &str, local_port: u16) -> Result { + let subcommand = if self.funnel { "funnel" } else { "serve" }; + + let hostname = if let Some(ref h) = self.hostname { + h.clone() + } else { + let output = tokio::time::timeout( + tokio::time::Duration::from_secs(10), + Command::new("tailscale") + .args(["status", "--json"]) + .output(), + ) + .await + .map_err(|_| anyhow::anyhow!("tailscale status --json timed out after 10s"))??; + + if !output.status.success() { + bail!( + "tailscale status failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + let status: serde_json::Value = serde_json::from_slice(&output.stdout) + .map_err(|e| anyhow::anyhow!("Failed to parse tailscale status JSON: {e}"))?; + status["Self"]["DNSName"] + .as_str() + .ok_or_else(|| anyhow::anyhow!("tailscale status missing Self.DNSName field"))? + .trim_end_matches('.') + .to_string() + }; + + let target = format!("http://{local_host}:{local_port}"); + let child = Command::new("tailscale") + .args([subcommand, &target]) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn()?; + + let public_url = format!("https://{hostname}"); + + if let Ok(mut guard) = self.url.write() { + *guard = Some(public_url.clone()); + } + + let mut guard = self.proc.lock().await; + *guard = Some(TunnelProcess { child }); + + Ok(public_url) + } + + async fn stop(&self) -> Result<()> { + let subcommand = if self.funnel { "funnel" } else { "serve" }; + if let Err(e) = Command::new("tailscale") + .args([subcommand, "reset"]) + .output() + .await + { + tracing::warn!("tailscale {subcommand} reset failed: {e}"); + } + + if let Ok(mut guard) = self.url.write() { + *guard = None; + } + kill_shared(&self.proc).await + } + + async fn health_check(&self) -> bool { + let guard = self.proc.lock().await; + guard.as_ref().is_some_and(|tp| tp.child.id().is_some()) + } + + fn public_url(&self) -> Option { + self.url.read().ok().and_then(|guard| guard.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn constructor_stores_hostname_and_mode() { + let tunnel = TailscaleTunnel::new(true, Some("myhost.ts.net".into())); + assert!(tunnel.funnel); + assert_eq!(tunnel.hostname.as_deref(), Some("myhost.ts.net")); + } + + #[test] + fn public_url_none_before_start() { + assert!(TailscaleTunnel::new(false, None).public_url().is_none()); + } + + #[tokio::test] + async fn health_false_before_start() { + assert!(!TailscaleTunnel::new(false, None).health_check().await); + } + + #[tokio::test] + async fn stop_without_start_is_ok() { + assert!(TailscaleTunnel::new(false, None).stop().await.is_ok()); + } +} diff --git a/src/workspace/hygiene.rs b/src/workspace/hygiene.rs new file mode 100644 index 00000000..d269232b --- /dev/null +++ b/src/workspace/hygiene.rs @@ -0,0 +1,247 @@ +//! Memory hygiene: automatic cleanup of stale workspace documents. +//! +//! Runs on a configurable cadence and deletes daily log entries older +//! than the retention period. Identity files (`IDENTITY.md`, `SOUL.md`, +//! etc.) are never touched. +//! +//! ```text +//! ┌─────────────────────────────────────────────┐ +//! │ Hygiene Pass │ +//! │ │ +//! │ 1. Check cadence (skip if ran recently) │ +//! │ 2. List daily/ documents │ +//! │ 3. Delete those older than retention_days │ +//! │ 4. Log summary │ +//! └─────────────────────────────────────────────┘ +//! ``` + +use std::path::PathBuf; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +use crate::workspace::Workspace; + +/// Configuration for workspace hygiene. +#[derive(Debug, Clone)] +pub struct HygieneConfig { + /// Whether hygiene is enabled at all. + pub enabled: bool, + /// Documents in `daily/` older than this many days are deleted. + pub retention_days: u32, + /// Minimum hours between hygiene passes. + pub cadence_hours: u32, + /// Directory to store state file (default: `~/.ironclaw`). + pub state_dir: PathBuf, +} + +impl Default for HygieneConfig { + fn default() -> Self { + let state_dir = dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw"); + + Self { + enabled: true, + retention_days: 30, + cadence_hours: 12, + state_dir, + } + } +} + +/// Persisted state for tracking hygiene cadence. +#[derive(Debug, Serialize, Deserialize)] +struct HygieneState { + last_run: DateTime, +} + +/// Summary of what a hygiene pass cleaned up. +#[derive(Debug, Default)] +pub struct HygieneReport { + /// Number of daily log documents deleted. + pub daily_logs_deleted: u32, + /// Whether the run was skipped (cadence not yet elapsed). + pub skipped: bool, +} + +impl HygieneReport { + /// True if any cleanup work was done. + pub fn had_work(&self) -> bool { + self.daily_logs_deleted > 0 + } +} + +/// Run a hygiene pass if the cadence has elapsed. +/// +/// This is best-effort: failures are logged but never propagate. The +/// agent should not crash because cleanup failed. +pub async fn run_if_due(workspace: &Workspace, config: &HygieneConfig) -> HygieneReport { + if !config.enabled { + return HygieneReport { + skipped: true, + ..Default::default() + }; + } + + let state_file = config.state_dir.join("memory_hygiene_state.json"); + + // Check cadence + if let Some(state) = load_state(&state_file) { + let elapsed = Utc::now().signed_duration_since(state.last_run); + let cadence = chrono::Duration::hours(i64::from(config.cadence_hours)); + if elapsed < cadence { + tracing::debug!( + hours_since_last = elapsed.num_hours(), + cadence_hours = config.cadence_hours, + "memory hygiene: skipping (cadence not elapsed)" + ); + return HygieneReport { + skipped: true, + ..Default::default() + }; + } + } + + tracing::info!( + retention_days = config.retention_days, + "memory hygiene: starting cleanup pass" + ); + + let mut report = HygieneReport::default(); + + // Delete old daily logs + match cleanup_daily_logs(workspace, config.retention_days).await { + Ok(count) => report.daily_logs_deleted = count, + Err(e) => tracing::warn!("memory hygiene: failed to clean daily logs: {e}"), + } + + if report.had_work() { + tracing::info!( + daily_logs_deleted = report.daily_logs_deleted, + "memory hygiene: cleanup complete" + ); + } else { + tracing::debug!("memory hygiene: nothing to clean"); + } + + // Save state (best-effort) + save_state(&state_file); + + report +} + +/// Delete daily log documents older than `retention_days`. +async fn cleanup_daily_logs( + workspace: &Workspace, + retention_days: u32, +) -> Result { + let cutoff = Utc::now() - chrono::Duration::days(i64::from(retention_days)); + let entries = workspace.list("daily/").await?; + + let mut deleted = 0u32; + for entry in entries { + if entry.is_directory { + continue; + } + + // Check if the document is old enough to delete + if let Some(updated_at) = entry.updated_at + && updated_at < cutoff + { + let path = if entry.path.starts_with("daily/") { + entry.path.clone() + } else { + format!("daily/{}", entry.path) + }; + + if let Err(e) = workspace.delete(&path).await { + tracing::warn!(path, "memory hygiene: failed to delete: {e}"); + } else { + tracing::debug!(path, "memory hygiene: deleted old daily log"); + deleted += 1; + } + } + } + + Ok(deleted) +} + +fn state_path_dir(state_file: &std::path::Path) -> Option<&std::path::Path> { + state_file.parent() +} + +fn load_state(path: &std::path::Path) -> Option { + let data = std::fs::read_to_string(path).ok()?; + serde_json::from_str(&data).ok() +} + +fn save_state(path: &std::path::Path) { + let state = HygieneState { + last_run: Utc::now(), + }; + if let Some(dir) = state_path_dir(path) { + std::fs::create_dir_all(dir).ok(); + } + if let Ok(json) = serde_json::to_string_pretty(&state) + && let Err(e) = std::fs::write(path, json) + { + tracing::warn!("memory hygiene: failed to save state: {e}"); + } +} + +#[cfg(test)] +mod tests { + use crate::workspace::hygiene::*; + + #[test] + fn default_config_is_reasonable() { + let cfg = HygieneConfig::default(); + assert!(cfg.enabled); + assert_eq!(cfg.retention_days, 30); + assert_eq!(cfg.cadence_hours, 12); + } + + #[test] + fn report_defaults_to_no_work() { + let report = HygieneReport::default(); + assert!(!report.had_work()); + assert!(!report.skipped); + } + + #[test] + fn report_had_work_when_deleted() { + let report = HygieneReport { + daily_logs_deleted: 3, + skipped: false, + }; + assert!(report.had_work()); + } + + #[test] + fn load_state_returns_none_for_missing_file() { + assert!(load_state(std::path::Path::new("/tmp/nonexistent_hygiene.json")).is_none()); + } + + #[test] + fn save_and_load_state_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("hygiene_state.json"); + + save_state(&path); + let state = load_state(&path).expect("state should be loadable after save"); + + // Should be within the last second + let elapsed = Utc::now().signed_duration_since(state.last_run); + assert!(elapsed.num_seconds() < 2); + } + + #[test] + fn save_state_creates_parent_dirs() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("nested").join("deep").join("state.json"); + + save_state(&path); + assert!(path.exists()); + } +} diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index a6afb070..3165ecce 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -43,6 +43,7 @@ mod chunker; mod document; mod embeddings; +pub mod hygiene; #[cfg(feature = "postgres")] mod repository; mod search;