//! Main agent loop. 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::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::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult}; use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; use crate::workspace::Workspace; /// Collapse a tool output string into a single-line preview for display. pub(crate) fn truncate_for_preview(output: &str, max_chars: usize) -> String { let collapsed: String = output .chars() .take(max_chars + 50) .map(|c| if c == '\n' { ' ' } else { c }) .collect::() .split_whitespace() .collect::>() .join(" "); // char_indices gives us byte offsets at char boundaries, so the slice is always valid UTF-8. if collapsed.chars().count() > max_chars { let byte_offset = collapsed .char_indices() .nth(max_chars) .map(|(i, _)| i) .unwrap_or(collapsed.len()); format!("{}...", &collapsed[..byte_offset]) } else { collapsed } } /// 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. pub struct AgentDeps { pub store: Option>, pub llm: Arc, pub safety: Arc, pub tools: Arc, pub workspace: Option>, pub extension_manager: Option>, } /// 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, } impl Agent { /// Create a new agent. /// /// Optionally accepts pre-created `ContextManager` and `SessionManager` for sharing /// with external components (job tools, web gateway). Creates new ones if not provided. pub fn new( config: AgentConfig, deps: AgentDeps, channels: ChannelManager, heartbeat_config: Option, routine_config: Option, context_manager: Option>, session_manager: Option>, ) -> Self { let context_manager = context_manager .unwrap_or_else(|| Arc::new(ContextManager::new(config.max_parallel_jobs))); let session_manager = session_manager.unwrap_or_else(|| Arc::new(SessionManager::new())); let scheduler = Arc::new(Scheduler::new( config.clone(), context_manager.clone(), deps.llm.clone(), deps.safety.clone(), deps.tools.clone(), deps.store.clone(), )); Self { config, deps, channels: Arc::new(channels), context_manager, scheduler, router: Router::new(), session_manager, context_monitor: ContextMonitor::new(), heartbeat_config, routine_config, } } // Convenience accessors fn store(&self) -> Option<&Arc> { self.deps.store.as_ref() } fn llm(&self) -> &Arc { &self.deps.llm } fn safety(&self) -> &Arc { &self.deps.safety } fn tools(&self) -> &Arc { &self.deps.tools } fn workspace(&self) -> Option<&Arc> { self.deps.workspace.as_ref() } /// Run the agent main loop. pub async fn run(self) -> Result<(), Error> { // Start channels let mut message_stream = self.channels.start_all().await?; // Start self-repair task with notification forwarding let repair = Arc::new(DefaultSelfRepair::new( self.context_manager.clone(), self.config.stuck_threshold, self.config.max_repair_attempts, )); let repair_interval = self.config.repair_check_interval; let repair_channels = self.channels.clone(); let repair_handle = tokio::spawn(async move { loop { tokio::time::sleep(repair_interval).await; // Check stuck jobs let stuck_jobs = repair.detect_stuck_jobs().await; for job in stuck_jobs { tracing::info!("Attempting to repair stuck job {}", job.job_id); let result = repair.repair_stuck_job(&job).await; let notification = match &result { Ok(RepairResult::Success { message }) => { tracing::info!("Repair succeeded: {}", message); Some(format!( "Job {} was stuck for {}s, recovery succeeded: {}", job.job_id, job.stuck_duration.as_secs(), message )) } Ok(RepairResult::Failed { message }) => { tracing::error!("Repair failed: {}", message); Some(format!( "Job {} was stuck for {}s, recovery failed permanently: {}", job.job_id, job.stuck_duration.as_secs(), message )) } Ok(RepairResult::ManualRequired { message }) => { tracing::warn!("Manual intervention needed: {}", message); Some(format!( "Job {} needs manual intervention: {}", job.job_id, message )) } Ok(RepairResult::Retry { message }) => { tracing::warn!("Repair needs retry: {}", message); None // Don't spam the user on retries } Err(e) => { tracing::error!("Repair error: {}", e); None } }; if let Some(msg) = notification { let response = OutgoingResponse::text(format!("Self-Repair: {}", msg)); let _ = repair_channels.broadcast_all("default", response).await; } } // Check broken tools let broken_tools = repair.detect_broken_tools().await; for tool in broken_tools { tracing::info!("Attempting to repair broken tool: {}", tool.name); match repair.repair_broken_tool(&tool).await { Ok(RepairResult::Success { message }) => { let response = OutgoingResponse::text(format!( "Self-Repair: Tool '{}' repaired: {}", tool.name, message )); let _ = repair_channels.broadcast_all("default", response).await; } Ok(result) => { tracing::info!("Tool repair result: {:?}", result); } Err(e) => { tracing::error!("Tool repair error: {}", e); } } } } }); // Spawn session pruning task let session_mgr = self.session_manager.clone(); let session_idle_timeout = self.config.session_idle_timeout; let pruning_handle = tokio::spawn(async move { let mut interval = tokio::time::interval(std::time::Duration::from_secs(600)); // Every 10 min interval.tick().await; // Skip immediate first tick loop { interval.tick().await; session_mgr.prune_stale_sessions(session_idle_timeout).await; } }); // Spawn heartbeat if enabled let heartbeat_handle = if let Some(ref hb_config) = self.heartbeat_config { if hb_config.enabled { if let Some(workspace) = self.workspace() { let config = AgentHeartbeatConfig::default() .with_interval(std::time::Duration::from_secs(hb_config.interval_secs)); // Set up notification channel let (notify_tx, mut notify_rx) = tokio::sync::mpsc::channel::(16); // Spawn notification forwarder that routes through channel manager let notify_channel = hb_config.notify_channel.clone(); let notify_user = hb_config.notify_user.clone(); let channels = self.channels.clone(); tokio::spawn(async move { while let Some(response) = notify_rx.recv().await { let user = notify_user.as_deref().unwrap_or("default"); // Try the configured channel first, fall back to // broadcasting on all channels. let targeted_ok = if let Some(ref channel) = notify_channel { channels .broadcast(channel, user, response.clone()) .await .is_ok() } else { false }; if !targeted_ok { let results = channels.broadcast_all(user, response).await; for (ch, result) in results { if let Err(e) = result { tracing::warn!( "Failed to broadcast heartbeat to {}: {}", ch, e ); } } } } }); tracing::info!( "Heartbeat enabled with {}s interval", hb_config.interval_secs ); Some(spawn_heartbeat( config, workspace.clone(), self.llm().clone(), Some(notify_tx), )) } else { tracing::warn!("Heartbeat enabled but no workspace available"); None } } else { None } } else { None }; // Spawn routine engine if enabled let routine_handle = if let Some(ref rt_config) = self.routine_config { if rt_config.enabled { if let (Some(store), Some(workspace)) = (self.store(), self.workspace()) { // Set up notification channel (same pattern as heartbeat) let (notify_tx, mut notify_rx) = tokio::sync::mpsc::channel::(32); let engine = Arc::new(RoutineEngine::new( rt_config.clone(), Arc::clone(store), self.llm().clone(), Arc::clone(workspace), notify_tx, )); // Register routine tools self.deps .tools .register_routine_tools(Arc::clone(store), Arc::clone(&engine)); // Load initial event cache engine.refresh_event_cache().await; // Spawn notification forwarder let channels = self.channels.clone(); tokio::spawn(async move { while let Some(response) = notify_rx.recv().await { let user = response .metadata .get("notify_user") .and_then(|v| v.as_str()) .unwrap_or("default") .to_string(); let results = channels.broadcast_all(&user, response).await; for (ch, result) in results { if let Err(e) = result { tracing::warn!( "Failed to broadcast routine notification to {}: {}", ch, e ); } } } }); // Spawn cron ticker let cron_interval = std::time::Duration::from_secs(rt_config.cron_check_interval_secs); let cron_handle = spawn_cron_ticker(Arc::clone(&engine), cron_interval); // Store engine reference for event trigger checking // Safety: we're in run() which takes self, no other reference exists let engine_ref = Arc::clone(&engine); // SAFETY: self is consumed by run(), we can smuggle the engine in // via a local to use in the message loop below. tracing::info!( "Routines enabled: cron ticker every {}s, max {} concurrent", rt_config.cron_check_interval_secs, rt_config.max_concurrent_routines ); Some((cron_handle, engine_ref)) } else { tracing::warn!("Routines enabled but store/workspace not available"); None } } else { None } } else { None }; // Extract engine ref for use in message loop let routine_engine_for_loop = routine_handle.as_ref().map(|(_, e)| Arc::clone(e)); // Main message loop tracing::info!("Agent {} ready and listening", self.config.name); loop { let message = tokio::select! { biased; _ = tokio::signal::ctrl_c() => { tracing::info!("Ctrl+C received, shutting down..."); break; } msg = message_stream.next() => { match msg { Some(m) => m, None => { tracing::info!("All channel streams ended, shutting down..."); break; } } } }; match self.handle_message(&message).await { Ok(Some(response)) if !response.is_empty() => { let _ = self .channels .respond(&message, OutgoingResponse::text(response)) .await; } Ok(Some(_)) => { // Empty response, nothing to send (e.g. approval handled via send_status) } Ok(None) => { // Shutdown signal received (/quit, /exit, /shutdown) tracing::info!("Shutdown command received, exiting..."); break; } Err(e) => { tracing::error!("Error handling message: {}", e); let _ = self .channels .respond(&message, OutgoingResponse::text(format!("Error: {}", e))) .await; } } // Check event triggers (cheap in-memory regex, fires async if matched) if let Some(ref engine) = routine_engine_for_loop { let fired = engine.check_event_triggers(&message).await; if fired > 0 { tracing::debug!("Fired {} event-triggered routines", fired); } } } // Cleanup tracing::info!("Agent shutting down..."); repair_handle.abort(); pruning_handle.abort(); if let Some(handle) = heartbeat_handle { handle.abort(); } if let Some((cron_handle, _)) = routine_handle { cron_handle.abort(); } self.scheduler.stop_all().await; self.channels.shutdown_all().await?; Ok(()) } async fn handle_message(&self, message: &IncomingMessage) -> Result, Error> { // Parse submission type first let submission = SubmissionParser::parse(&message.content); // Hydrate thread from DB if it's a historical thread not in memory if let Some(ref external_thread_id) = message.thread_id { self.maybe_hydrate_thread(message, external_thread_id).await; } // Resolve session and thread let (session, thread_id) = self .session_manager .resolve_thread( &message.user_id, &message.channel, message.thread_id.as_deref(), ) .await; // Auth mode interception: if the thread is awaiting a token, route // the message directly to the credential store. Nothing touches // logs, turns, history, or compaction. let pending_auth = { let sess = session.lock().await; sess.threads .get(&thread_id) .and_then(|t| t.pending_auth.clone()) }; if let Some(pending) = pending_auth { match &submission { Submission::UserInput { content } => { return self .process_auth_token(message, &pending, content, session, thread_id) .await; } _ => { // Any control submission (interrupt, undo, etc.) cancels auth mode let mut sess = session.lock().await; if let Some(thread) = sess.threads.get_mut(&thread_id) { thread.pending_auth = None; } // Fall through to normal handling } } } tracing::debug!( "Received message from {} on {} ({} chars)", message.user_id, message.channel, message.content.len() ); // Process based on submission type let result = match submission { Submission::UserInput { content } => { self.process_user_input(message, session, thread_id, &content) .await } Submission::SystemCommand { command, args } => { self.handle_system_command(&command, &args).await } Submission::Undo => self.process_undo(session, thread_id).await, Submission::Redo => self.process_redo(session, thread_id).await, Submission::Interrupt => self.process_interrupt(session, thread_id).await, Submission::Compact => self.process_compact(session, thread_id).await, Submission::Clear => self.process_clear(session, thread_id).await, Submission::NewThread => self.process_new_thread(message).await, Submission::Heartbeat => self.process_heartbeat().await, Submission::Summarize => self.process_summarize(session, thread_id).await, Submission::Suggest => self.process_suggest(session, thread_id).await, Submission::Quit => return Ok(None), Submission::SwitchThread { thread_id: target } => { self.process_switch_thread(message, target).await } Submission::Resume { checkpoint_id } => { self.process_resume(session, thread_id, checkpoint_id).await } Submission::ExecApproval { request_id, approved, always, } => { self.process_approval( message, session, thread_id, Some(request_id), approved, always, ) .await } Submission::ApprovalResponse { approved, always } => { self.process_approval(message, session, thread_id, None, approved, always) .await } }; // Convert SubmissionResult to response string match result? { SubmissionResult::Response { content } => Ok(Some(content)), SubmissionResult::Ok { message } => Ok(message), SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))), SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())), SubmissionResult::NeedApproval { request_id, tool_name, description, parameters, } => { // Each channel renders the approval prompt via send_status. // Web gateway shows an inline card, REPL prints a formatted prompt, etc. let _ = self .channels .send_status( &message.channel, StatusUpdate::ApprovalNeeded { request_id: request_id.to_string(), tool_name, description, parameters, }, &message.metadata, ) .await; // Empty string signals the caller to skip respond() (no duplicate text) Ok(Some(String::new())) } } } /// 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)) => { 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) for 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) }; // For shell commands, override auto-approval for // destructive patterns that should always require // explicit per-invocation approval. if is_auto_approved && tc.name == "shell" && let Some(cmd) = tc .arguments .get("command") .and_then(|c| c.as_str().map(String::from)) .or_else(|| { tc.arguments .as_str() .and_then(|s| { serde_json::from_str::(s).ok() }) .and_then(|v| { v.get("command") .and_then(|c| c.as_str().map(String::from)) }) }) && crate::tools::builtin::shell::requires_explicit_approval(&cmd) { tracing::info!( "Shell command '{}' requires explicit approval despite auto-approve", cmd.chars().take(80).collect::() ); 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 }); } } 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 { 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.")) } } 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.")); } if let Some(checkpoint) = mgr.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 }))?; 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] fn test_truncate_short_input() { assert_eq!(truncate_for_preview("hello", 10), "hello"); } #[test] fn test_truncate_empty_input() { assert_eq!(truncate_for_preview("", 10), ""); } #[test] fn test_truncate_exact_length() { assert_eq!(truncate_for_preview("hello", 5), "hello"); } #[test] fn test_truncate_over_limit() { let result = truncate_for_preview("hello world, this is long", 10); assert!(result.ends_with("...")); // "hello worl" = 10 chars + "..." assert_eq!(result, "hello worl..."); } #[test] fn test_truncate_collapses_newlines() { let result = truncate_for_preview("line1\nline2\nline3", 100); assert!(!result.contains('\n')); assert_eq!(result, "line1 line2 line3"); } #[test] fn test_truncate_collapses_whitespace() { let result = truncate_for_preview("hello world", 100); assert_eq!(result, "hello world"); } #[test] fn test_truncate_multibyte_utf8() { // Each emoji is 4 bytes. Truncating at char boundary must not panic. let input = "😀😁😂ðŸĪĢ😃😄😅😆😉😊"; let result = truncate_for_preview(input, 5); assert!(result.ends_with("...")); // First 5 chars = 5 emoji assert_eq!(result, "😀😁😂ðŸĪĢ😃..."); } #[test] fn test_truncate_cjk_characters() { // CJK chars are 3 bytes each in UTF-8. let input = "ä― åĨ―äļ–į•Œæĩ‹čŊ•æ•°æŪåūˆé•ŋįš„å­—įŽĶäļē"; let result = truncate_for_preview(input, 4); assert_eq!(result, "ä― åĨ―äļ–į•Œ..."); } #[test] fn test_truncate_mixed_multibyte_and_ascii() { let input = "hello äļ–į•Œ foo"; let result = truncate_for_preview(input, 8); // 'h','e','l','l','o',' ','äļ–','į•Œ' = 8 chars assert_eq!(result, "hello äļ–į•Œ..."); } }