PicoBot/src/tools/agent_task.rs

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<AgentCoordinator>,
}

impl AgentTaskTool {
    pub fn new(coordinator: Arc<AgentCoordinator>) -> 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<ToolResult> {
        self.execute_with_context(&ToolExecutionContext::default(), args)
            .await
            .map(|output| output.result)
    }

    async fn execute_with_context(
        &self,
        context: &ToolExecutionContext,
        args: Value,
    ) -> anyhow::Result<ToolOutput> {
        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<ToolResult> {
        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<ToolResult> {
        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<Value> =
                    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<ToolResult> {
        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<ToolResult> {
        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<String>) -> ToolResult {
    ToolResult {
        success: false,
        output: String::new(),
        error: Some(error.into()),
    }
}