//! Main agent loop. use std::sync::Arc; /// Escape special characters for Telegram's legacy Markdown. /// /// In Telegram's Markdown mode, these characters have special meaning: /// - `_` starts/ends italic /// - `*` starts/ends bold /// - `` ` `` starts/ends code /// - `[` starts a link /// /// We escape them with backslash so dynamic content doesn't break formatting. fn escape_telegram_markdown(s: &str) -> String { let mut result = String::with_capacity(s.len()); for c in s.chars() { match c { '_' | '*' | '`' | '[' => { result.push('\\'); result.push(c); } _ => result.push(c), } } result } 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::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}; use crate::context::ContextManager; use crate::context::JobContext; use crate::error::Error; use crate::history::Store; use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult}; use crate::safety::SafetyLayer; use crate::tools::ToolRegistry; use crate::workspace::Workspace; /// 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>, } /// 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, } 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, 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, } } // 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 { // Route notification to configured channel/user, or broadcast to all match (¬ify_channel, ¬ify_user) { (Some(channel), Some(user)) => { // Send to specific channel and user if let Err(e) = channels.broadcast(channel, user, response.clone()).await { tracing::warn!( "Failed to send heartbeat to {}/{}: {}", channel, user, e ); } else { tracing::debug!( "Heartbeat notification sent to {}/{}", channel, user ); } } (None, Some(user)) => { // Broadcast to all channels for this user 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 ); } } } _ => { // No explicit target, broadcast to all channels // for the default user so notifications actually // reach someone instead of vanishing into logs. let results = channels.broadcast_all("default", 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 }; // Main message loop tracing::info!("Agent {} ready and listening", self.config.name); while let Some(message) = message_stream.next().await { match self.handle_message(&message).await { Ok(Some(response)) => { let _ = self .channels .respond(&message, OutgoingResponse::text(response)) .await; } Ok(None) => { // Shutdown signal received tracing::info!("Shutdown signal received, exiting..."); break; } Err(e) => { tracing::error!("Error handling message: {}", e); let _ = self .channels .respond(&message, OutgoingResponse::text(format!("Error: {}", e))) .await; } } } // Cleanup tracing::info!("Agent shutting down..."); repair_handle.abort(); pruning_handle.abort(); if let Some(handle) = heartbeat_handle { handle.abort(); } self.scheduler.stop_all().await; self.channels.shutdown_all().await?; Ok(()) } async fn handle_message(&self, message: &IncomingMessage) -> Result, Error> { tracing::debug!( "Received message from {} on {}: {}", message.user_id, message.channel, truncate(&message.content, 100) ); // Parse submission type first let submission = SubmissionParser::parse(&message.content); // Resolve session and thread let (session, thread_id) = self .session_manager .resolve_thread( &message.user_id, &message.channel, message.thread_id.as_deref(), ) .await; // Process based on submission type let result = match submission { Submission::UserInput { content } => { self.process_user_input(message, session, thread_id, &content) .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::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, } => { // Format approval request for user let params_preview = serde_json::to_string_pretty(¶meters) .unwrap_or_else(|_| parameters.to_string()); let params_truncated = if params_preview.chars().count() > 200 { format!( "{}...", params_preview.chars().take(200).collect::() ) } else { params_preview }; // Escape Markdown special chars in dynamic values to avoid breaking // Telegram's Markdown parser (underscores, asterisks, backticks, brackets) let tool_name_escaped = escape_telegram_markdown(&tool_name); let description_escaped = escape_telegram_markdown(&description); // Params go inside a code block, so no escaping needed there Ok(Some(format!( "🔒 Tool requires approval:\n\n\ *Tool:* {}\n\ *Description:* {}\n\ *Parameters:*\n```\n{}\n```\n\n\ Reply with:\n\ • yes or approve to allow this tool\n\ • always to always allow this tool in this session\n\ • no or deny to reject\n\n\ Request ID: {}", tool_name_escaped, description_escaped, params_truncated, request_id ))) } } } 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) .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); let _ = self .channels .send_status( &message.channel, StatusUpdate::Status("Done".into()), &message.metadata, ) .await; 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()); Ok(SubmissionResult::error(e.to_string())) } } } /// 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. async fn run_agentic_loop( &self, message: &IncomingMessage, session: Arc>, thread_id: Uuid, initial_messages: Vec, ) -> 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 = false; loop { iteration += 1; if iteration > MAX_TOOL_ITERATIONS { return Err(crate::error::LlmError::InvalidResponse { provider: "nearai".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) { if 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); let result = reasoning.respond_with_tools(&context).await?; match 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) => { tools_executed = true; // 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) { if 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 { if tool.requires_approval() { // Check if auto-approved for this session let is_auto_approved = { let sess = session.lock().await; sess.is_tool_auto_approved(&tc.name) }; 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 tool_result = self .execute_chat_tool(&tc.name, &tc.arguments, &job_ctx) .await; // Record result in thread { let mut sess = session.lock().await; if let Some(thread) = sess.threads.get_mut(&thread_id) { if 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()); } } } } } // 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()); } // Execute with timeout let result = tokio::time::timeout(std::time::Duration::from_secs(60), async { tool.execute(params.clone(), job_ctx).await }) .await .map_err(|_| crate::error::ToolError::Timeout { name: tool_name.to_string(), timeout: std::time::Duration::from_secs(60), })? .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 { if 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 tool_result = self .execute_chat_tool(&pending.tool_name, &pending.parameters, &job_ctx) .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) { if 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()); } } } } } // 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 let result = self .run_agentic_loop(message, session.clone(), thread_id, context_messages) .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); 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 ))) } } 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() { if 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 { if 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))), } } async fn handle_command( &self, command: &str, _args: &[String], ) -> Result, Error> { match command { "help" => Ok(Some( r#"Commands: /job - Create a job /status [id] - Check job status /cancel - Cancel a job /list - List all jobs /help - Help a stuck job /undo - Undo last turn /redo - Redo undone turn /compact - Compress context /clear - Clear thread /interrupt - Stop current turn /thread new - New thread /thread - Switch thread /resume - Resume checkpoint /heartbeat - Run heartbeat check now /summarize - Summarize current thread /suggest - Suggest next steps /quit - Exit"# .to_string(), )), "ping" => Ok(Some("pong!".to_string())), "version" => Ok(Some(format!( "{} v{}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION") ))), "tools" => { let tools = self.tools().list().await; Ok(Some(format!("Available tools: {}", tools.join(", ")))) } "quit" | "exit" | "shutdown" => { // Signal shutdown - return None to indicate no response needed Ok(None) } _ => Ok(Some(format!("Unknown command: {}. Try /help", command))), } } } fn truncate(s: &str, max_len: usize) -> String { if s.len() <= max_len { s.to_string() } else { format!("{}...", &s[..max_len]) } }