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;