use std::sync::Arc; use async_trait::async_trait; use serde_json::{Value, json}; use crate::agent::AgentCoordinator; use crate::storage::agent_run::AgentRunRecord; use crate::tools::traits::{DelegationPolicy, Tool, ToolExecutionContext, ToolOutput, ToolResult}; const RESULT_PREVIEW_CHARS: usize = 2_000; /// Scoped inspection and control of durable Agent runs. Authorization is /// derived from the caller's ToolExecutionContext (session for ROOT, tree /// position for named Agents); run IDs are never credentials. pub struct AgentTaskTool { coordinator: Arc, } impl AgentTaskTool { pub fn new(coordinator: Arc) -> Self { Self { coordinator } } } #[async_trait] impl Tool for AgentTaskTool { fn name(&self) -> &str { "agent_task" } fn description(&self) -> &str { "Inspect or control delegated Agent runs: get reads one run, list shows the session's runs, get_result returns the full terminal result, cancel stops a non-terminal run." } fn delegation_policy(&self) -> DelegationPolicy { DelegationPolicy::RuntimeInjected } fn parameters_schema(&self) -> Value { json!({ "type": "object", "properties": { "action": { "type": "string", "enum": ["get", "list", "get_result", "cancel"], "description": "Operation to perform on Agent runs" }, "run_id": { "type": "string", "description": "Target run identifier for get/get_result/cancel" }, "cursor_created_at": { "type": "integer", "description": "Pagination cursor: created_at of the last run seen" }, "cursor_id": { "type": "string", "description": "Pagination cursor: id of the last run seen" }, "limit": { "type": "integer", "minimum": 1, "maximum": 100, "description": "Maximum number of runs to list (default 20)" } }, "required": ["action"] }) } fn read_only(&self) -> bool { false } async fn execute(&self, args: Value) -> anyhow::Result { self.execute_with_context(&ToolExecutionContext::default(), args) .await .map(|output| output.result) } async fn execute_with_context( &self, context: &ToolExecutionContext, args: Value, ) -> anyhow::Result { let action = args .get("action") .and_then(Value::as_str) .unwrap_or_default(); let result = match action { "get" => self.handle_get(context, &args).await, "list" => self.handle_list(context, &args).await, "get_result" => self.handle_get_result(context, &args).await, "cancel" => self.handle_cancel(context, &args).await, other => Ok(ToolResult { success: false, output: String::new(), error: Some(format!( "unknown agent_task action '{other}'; supported: get, list, get_result, cancel" )), }), }; Ok(result?.into()) } } impl AgentTaskTool { fn run_id<'a>(&self, args: &'a Value) -> anyhow::Result<&'a str> { args.get("run_id") .and_then(Value::as_str) .filter(|value| !value.trim().is_empty()) .ok_or_else(|| anyhow::anyhow!("missing required parameter: run_id")) } async fn handle_get( &self, context: &ToolExecutionContext, args: &Value, ) -> anyhow::Result { let run_id = self.run_id(args)?; match self.coordinator.get_run(context, run_id).await { Ok(Some(run)) => Ok(success(run_projection(&run, true))), Ok(None) => Ok(failure(format!("run not found: {run_id}"))), Err(error) => Ok(failure(error.to_string())), } } async fn handle_list( &self, context: &ToolExecutionContext, args: &Value, ) -> anyhow::Result { let cursor = match ( args.get("cursor_created_at").and_then(Value::as_i64), args.get("cursor_id").and_then(Value::as_str), ) { (Some(created_at), Some(id)) => Some((created_at, id.to_string())), (None, None) => None, _ => { return Ok(failure( "cursor requires both cursor_created_at and cursor_id", )); } }; let limit = args.get("limit").and_then(Value::as_i64).unwrap_or(20); match self.coordinator.list_runs(context, cursor, limit).await { Ok(runs) => { let payload: Vec = runs.iter().map(|run| run_projection(run, false)).collect(); Ok(success(json!({ "runs": payload }))) } Err(error) => Ok(failure(error.to_string())), } } async fn handle_get_result( &self, context: &ToolExecutionContext, args: &Value, ) -> anyhow::Result { let run_id = self.run_id(args)?; match self.coordinator.get_result(context, run_id).await { Ok(Some(run)) => Ok(success(json!({ "run_id": run.id, "status": run.status.as_str(), "result": run.result, "error": run.error, "tool_calls": run.tool_calls_count, "iterations": run.iterations, "finished_at": run.finished_at }))), Ok(None) => Ok(failure(format!( "run {run_id} is not terminal or does not exist" ))), Err(error) => Ok(failure(error.to_string())), } } async fn handle_cancel( &self, context: &ToolExecutionContext, args: &Value, ) -> anyhow::Result { let run_id = self.run_id(args)?; match self .coordinator .cancel_run(context, run_id, "cancelled via agent_task") .await { Ok(true) => Ok(success(json!({ "run_id": run_id, "status": "cancelled" }))), Ok(false) => Ok(failure(format!( "cannot cancel run {run_id}; it is terminal or does not exist" ))), Err(error) => Ok(failure(error.to_string())), } } } fn run_projection(run: &AgentRunRecord, include_result_preview: bool) -> Value { let mut value = json!({ "run_id": run.id, "group_id": run.group_id, "agent_id": run.agent_id, "status": run.status.as_str(), "mode": run.mode.as_str(), "depth": run.depth, "parent_run_id": run.parent_run_id, "provider_profile": run.provider_profile, "model_id": run.model_id, "tool_calls": run.tool_calls_count, "iterations": run.iterations, "created_at": run.created_at, "started_at": run.started_at, "finished_at": run.finished_at, "error": run.error, }); if include_result_preview { value["result_preview"] = json!(run.result.as_deref().map(preview)); value["result_truncated"] = json!( run.result .as_deref() .is_some_and(|result| result.chars().count() > RESULT_PREVIEW_CHARS) ); } value } fn preview(value: &str) -> String { if value.chars().count() <= RESULT_PREVIEW_CHARS { value.to_string() } else { let cut = value.floor_char_boundary(RESULT_PREVIEW_CHARS); format!("{}...", &value[..cut]) } } fn success(value: Value) -> ToolResult { ToolResult { success: true, output: value.to_string(), error: None, } } fn failure(error: impl Into) -> ToolResult { ToolResult { success: false, output: String::new(), error: Some(error.into()), } }