Files
optimclaw/src/agent/agent_loop.rs
T
Illia Polosukhin 9b729795fb Merge remote-tracking branch 'origin/main' into ui
# Conflicts:
#	src/channels/mod.rs
#	src/main.rs
2026-02-06 13:22:26 -08:00

1747 lines
65 KiB
Rust

//! 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<Arc<Store>>,
pub llm: Arc<dyn LlmProvider>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub workspace: Option<Arc<Workspace>>,
}
/// The main agent that coordinates all components.
pub struct Agent {
config: AgentConfig,
deps: AgentDeps,
channels: Arc<ChannelManager>,
context_manager: Arc<ContextManager>,
scheduler: Arc<Scheduler>,
router: Router,
session_manager: Arc<SessionManager>,
context_monitor: ContextMonitor,
heartbeat_config: Option<HeartbeatConfig>,
}
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<HeartbeatConfig>,
context_manager: Option<Arc<ContextManager>>,
session_manager: Option<Arc<SessionManager>>,
) -> 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<Store>> {
self.deps.store.as_ref()
}
fn llm(&self) -> &Arc<dyn LlmProvider> {
&self.deps.llm
}
fn safety(&self) -> &Arc<SafetyLayer> {
&self.deps.safety
}
fn tools(&self) -> &Arc<ToolRegistry> {
&self.deps.tools
}
fn workspace(&self) -> Option<&Arc<Workspace>> {
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::<OutgoingResponse>(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 (&notify_channel, &notify_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<Option<String>, 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(&parameters)
.unwrap_or_else(|_| parameters.to_string());
let params_truncated = if params_preview.chars().count() > 200 {
format!(
"{}...",
params_preview.chars().take(200).collect::<String>()
)
} 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<Mutex<Session>>,
thread_id: Uuid,
content: &str,
) -> Result<SubmissionResult, Error> {
// 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::<Vec<_>>()
.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<Mutex<Session>>,
thread_id: Uuid,
initial_messages: Vec<ChatMessage>,
) -> Result<AgenticLoopResult, Error> {
// 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<String, Error> {
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::<Vec<_>>()
.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<SubmissionResult, Error> {
// 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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
request_id: Option<Uuid>,
approved: bool,
always: bool,
) -> Result<SubmissionResult, Error> {
// 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<SubmissionResult, Error> {
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<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
checkpoint_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<String>,
) -> Result<String, Error> {
// 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<String>,
) -> Result<String, Error> {
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<String, Error> {
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<String>,
) -> Result<String, Error> {
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<String, Error> {
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<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
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<Option<String>, Error> {
match command {
"help" => Ok(Some(
r#"Commands:
/job <desc> - Create a job
/status [id] - Check job status
/cancel <id> - Cancel a job
/list - List all jobs
/help <job_id> - 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 <id> - Switch thread
/resume <id> - 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])
}
}