tiny-agent 0.3.0

mod default_tools;

pub use default_tools::{register_default_tools, register_skill_tools, register_spawn_agent};

use super::sandbox::Sandbox;
use crate::messages::MessageChannel;
use crate::messages::MessageKind;
use crate::shared::{ToolDefinition, UserInteraction};
use crate::tasks::{TaskBuffer, TaskManager};
use schemars::{JsonSchema, schema_for};
use serde::de::DeserializeOwned;
use serde_json::{Map, Value, json};
use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc, time::Duration};
use tokio_util::sync::CancellationToken;

/// 工具到达让步窗口后转为后台任务的执行策略。
#[derive(Clone, Debug)]
pub struct AsyncToolConfig {
    pub yield_after: Duration,
}

impl AsyncToolConfig {
    pub fn new(yield_after: Duration) -> Self {
        Self { yield_after }
    }
}

/// 工具注册元信息。旧的 `register(...)` 仍可用;需要配置异步策略时使用
/// `ToolSpec::new(...).asynchronous(...)` 再交给 `register_tool(...)`。
#[derive(Clone, Debug)]
pub struct ToolSpec {
    pub name: String,
    pub description: String,
    pub read_only: bool,
    pub asynchronous: Option<AsyncToolConfig>,
}

impl ToolSpec {
    pub fn new(name: impl Into<String>, description: impl Into<String>, read_only: bool) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            read_only,
            asynchronous: None,
        }
    }

    pub fn asynchronous(mut self, config: AsyncToolConfig) -> Self {
        self.asynchronous = Some(config);
        self
    }
}

/// 工具执行期间可发出的进度通道。异步工具的 progress 会被 `check(id)` 按增量读回;
/// 同时也会推给 `MessageChannel` 供前端实时展示。
#[derive(Clone, Default)]
pub struct ToolProgress {
    messages: MessageChannel,
    session_id: String,
    buffer: Option<TaskBuffer>,
}

impl ToolProgress {
    pub(crate) fn new(
        messages: MessageChannel,
        session_id: String,
        buffer: Option<TaskBuffer>,
    ) -> Self {
        Self {
            messages,
            session_id,
            buffer,
        }
    }

    pub fn emit_text(&self, chunk: impl AsRef<str>) {
        let chunk = chunk.as_ref();
        if let Some(buffer) = &self.buffer {
            buffer.append(chunk);
        }
        self.messages.send(
            &self.session_id,
            MessageKind::ToolOutput {
                chunk: chunk.to_string(),
            },
        );
    }
}

/// 调用工具时注入的运行时上下文。
///
/// 大多数工具用不到它(用 [`ToolRegistry::register`] 注册的 handler 拿不到);
/// 需要上报工具进度时用 [`ToolRegistry::register_with_ctx`] 注册,并通过
/// [`ToolCtx::progress`] 发出增量输出。这样实时 UI 与 `check(id)` 看到的内容保持一致。
#[derive(Clone, Default)]
pub struct ToolCtx {
    messages: MessageChannel,
    session_id: String,
    tasks: TaskManager,
    /// 当前这轮 `run` 的取消令牌。工具若自己派生子 agent,应把它(或其 `child_token()`)
    /// 透传下去,这样用户取消父任务时子 agent 也能**干净地**一并取消(走 Interrupted/存档),
    /// 而不是被父 future 的 drop 硬掐断。
    cancellation: CancellationToken,
    /// 当前工具调用的进度通道。
    pub progress: ToolProgress,
}

impl ToolCtx {
    pub(crate) fn new(
        messages: MessageChannel,
        session_id: String,
        tasks: TaskManager,
        cancellation: CancellationToken,
    ) -> Self {
        Self {
            messages,
            session_id,
            tasks,
            cancellation,
            progress: ToolProgress::default(),
        }
    }

    /// 当前会话 id。工具需要打标签或组织输出时可读取;进度输出请走 [`ToolCtx::progress`]。
    pub fn session_id(&self) -> &str {
        &self.session_id
    }

    /// 当前这轮 `run` 的取消令牌。派生子 agent 的工具应用它(或 `child_token()`)替代
    /// `CancellationToken::new()`,以便取消能向下传播。
    pub fn cancellation_token(&self) -> CancellationToken {
        self.cancellation.clone()
    }

    pub(crate) fn tasks(&self) -> &TaskManager {
        &self.tasks
    }
}

type BoxFuture = Pin<Box<dyn Future<Output = Result<ToolOutcome, String>> + Send>>;
type BoxedCall = Arc<dyn Fn(Arc<dyn Sandbox>, Value, ToolCtx) -> BoxFuture + Send + Sync>;

#[derive(Debug, Clone)]
pub enum ToolOutcome {
    Completed(Value),
    NeedsUserInteraction(UserInteraction),
}

impl From<Value> for ToolOutcome {
    fn from(value: Value) -> Self {
        Self::Completed(value)
    }
}

#[derive(Clone)]
pub struct ToolEntry {
    pub name: String,
    pub description: String,
    pub schema: Value,
    /// 是否只读、无副作用。只读工具之间可以并发执行(见 `core::tool_execution` 的分段调度);
    /// 任何会改动 sandbox 状态的工具都必须保持 `false`,以免与并发读发生 read-write race。
    pub read_only: bool,
    pub asynchronous: Option<AsyncToolConfig>,
    call: BoxedCall,
}

#[derive(Clone)]
pub struct ToolRegistry {
    tools: HashMap<String, ToolEntry>,
}

impl ToolRegistry {
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// 注册一个工具。`read_only` 标记它是否无副作用 —— 只读工具之间可被并发调度
    /// (见 `core::tool_execution` 的分段执行),任何会改动 sandbox 的工具都必须传 `false`。
    /// handler 统一返回 `ToolOutcome`(普通完成用 `Value.into()`,需要用户交互用 `NeedsUserInteraction`)。
    pub fn register<A, F, Fut>(
        &mut self,
        name: &str,
        description: &str,
        read_only: bool,
        handler: F,
    ) where
        A: DeserializeOwned + JsonSchema + Send + 'static,
        F: Fn(Arc<dyn Sandbox>, A) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
    {
        self.register_tool(ToolSpec::new(name, description, read_only), handler);
    }

    pub fn register_tool<A, F, Fut>(&mut self, spec: ToolSpec, handler: F)
    where
        A: DeserializeOwned + JsonSchema + Send + 'static,
        F: Fn(Arc<dyn Sandbox>, A) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
    {
        let schema = serde_json::to_value(schema_for!(A)).unwrap();
        let handler = Arc::new(handler);
        let call: BoxedCall =
            Arc::new(move |sandbox: Arc<dyn Sandbox>, v: Value, _ctx: ToolCtx| {
                let handler = handler.clone();
                Box::pin(async move {
                    let args: A =
                        serde_json::from_value(v).map_err(|e| format!("参数不合法: {e}"))?;
                    handler(sandbox, args).await
                })
            });
        self.insert(spec, schema, call);
    }

    fn insert(&mut self, spec: ToolSpec, schema: Value, call: BoxedCall) {
        self.tools.insert(
            spec.name.clone(),
            ToolEntry {
                name: spec.name,
                description: spec.description,
                schema,
                read_only: spec.read_only,
                asynchronous: spec.asynchronous,
                call,
            },
        );
    }

    /// 与 [`register`](Self::register) 相同,但 handler 额外收到一个 [`ToolCtx`]。
    /// 工具需要上报运行进度时应使用 `ctx.progress`,这样 UI 实时输出与 `check(id)`
    /// 增量读取会保持同源。
    pub fn register_with_ctx<A, F, Fut>(
        &mut self,
        name: &str,
        description: &str,
        read_only: bool,
        handler: F,
    ) where
        A: DeserializeOwned + JsonSchema + Send + 'static,
        F: Fn(Arc<dyn Sandbox>, A, ToolCtx) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
    {
        self.register_with_ctx_tool(ToolSpec::new(name, description, read_only), handler);
    }

    pub fn register_with_ctx_tool<A, F, Fut>(&mut self, spec: ToolSpec, handler: F)
    where
        A: DeserializeOwned + JsonSchema + Send + 'static,
        F: Fn(Arc<dyn Sandbox>, A, ToolCtx) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ToolOutcome, String>> + Send + 'static,
    {
        let schema = serde_json::to_value(schema_for!(A)).unwrap();
        let handler = Arc::new(handler);
        let call: BoxedCall = Arc::new(move |sandbox: Arc<dyn Sandbox>, v: Value, ctx: ToolCtx| {
            let handler = handler.clone();
            Box::pin(async move {
                let args: A = serde_json::from_value(v).map_err(|e| format!("参数不合法: {e}"))?;
                handler(sandbox, args, ctx).await
            })
        });
        self.insert(spec, schema, call);
    }

    /// 工具是否声明为只读。未知工具按非只读处理(最保守,强制串行)。
    pub fn is_read_only(&self, name: &str) -> bool {
        self.tools.get(name).is_some_and(|t| t.read_only)
    }

    /// 是否注册了某个工具。
    pub fn contains(&self, name: &str) -> bool {
        self.tools.contains_key(name)
    }

    /// 已注册的全部工具名(无序)。
    pub fn names(&self) -> Vec<String> {
        self.tools.keys().cloned().collect()
    }

    /// 派生一个只含给定名字的子集 registry,完整保留各工具的 callable/schema/异步策略。
    /// 未知名字静默忽略——调用方若需校验可先用 [`contains`](Self::contains)。
    /// 用于按能力边界裁剪某个(子)agent 可见的工具集。
    pub fn subset<S: AsRef<str>>(&self, names: &[S]) -> ToolRegistry {
        let mut tools = HashMap::new();
        for name in names {
            let name = name.as_ref();
            if let Some(entry) = self.tools.get(name) {
                tools.insert(name.to_string(), entry.clone());
            }
        }
        ToolRegistry { tools }
    }

    /// 按名字调度执行,执行时注入该会话绑定的 sandbox 与 [`ToolCtx`]。
    pub async fn call(
        &self,
        name: &str,
        args: Value,
        sandbox: Arc<dyn Sandbox>,
        ctx: ToolCtx,
    ) -> Result<ToolOutcome, String> {
        let tool = self
            .tools
            .get(name)
            .ok_or_else(|| format!("工具不存在: {name}"))?;
        let mut ctx = ctx;
        ctx.progress = ToolProgress::new(ctx.messages.clone(), ctx.session_id.clone(), None);

        let Some(config) = tool.asynchronous.clone() else {
            return (tool.call)(sandbox, args, ctx).await;
        };

        let buffer = TaskBuffer::new();
        ctx.progress = ToolProgress::new(
            ctx.messages.clone(),
            ctx.session_id.clone(),
            Some(buffer.clone()),
        );
        call_asynchronous_tool(
            tool.name.clone(),
            tool.call.clone(),
            sandbox,
            args,
            ctx,
            buffer,
            config,
        )
        .await
    }

    /// 将工具 struct 转成 `ToolDefinition` 列表
    pub fn definitions(&self) -> Vec<ToolDefinition> {
        self.tools
            .values()
            .map(|t| ToolDefinition {
                name: t.name.clone(),
                desc: t.description.clone(),
                arguments: t.schema.clone(),
            })
            .collect()
    }
}

async fn call_asynchronous_tool(
    tool_name: String,
    call: BoxedCall,
    sandbox: Arc<dyn Sandbox>,
    args: Value,
    ctx: ToolCtx,
    buffer: TaskBuffer,
    config: AsyncToolConfig,
) -> Result<ToolOutcome, String> {
    let tasks = ctx.tasks.clone();
    let mut handle = tokio::spawn({
        let buffer = buffer.clone();
        async move {
            let result = match (call)(sandbox, args, ctx).await {
                Ok(ToolOutcome::Completed(value)) => Ok(value),
                Ok(ToolOutcome::NeedsUserInteraction(_)) => Err(
                    "asynchronous tools cannot request user interaction after yielding".to_string(),
                ),
                Err(error) => Err(error),
            };
            buffer.finish(result);
        }
    });

    tokio::select! {
        join_result = &mut handle => {
            if let Err(err) = join_result {
                return Err(format!("tool task failed: {err}"));
            }
            let output = buffer.drain_new();
            let (_finished, result, error) = buffer.status_parts();
            if let Some(error) = error {
                return Err(error);
            }
            Ok(merge_tool_result(
                "completed",
                output,
                result.unwrap_or_else(|| json!({})),
            ).into())
        }
        _ = tokio::time::sleep(config.yield_after) => {
            let output = buffer.drain_new();
            let id = tasks.register(tool_name.clone(), handle, buffer);
            Ok(json!({
                "id": id,
                "tool_name": tool_name,
                "status": "running",
                "output": output,
                "note": "工具仍在后台运行;用 check 工具按 id 回来查看新增输出与最终状态,用 kill 终止。",
            }).into())
        }
    }
}

pub(crate) fn merge_tool_result(status: &str, output: String, result: Value) -> Value {
    match result {
        Value::Object(mut object) => {
            object
                .entry("status".to_string())
                .or_insert_with(|| Value::String(status.to_string()));
            object
                .entry("output".to_string())
                .or_insert_with(|| Value::String(output));
            Value::Object(object)
        }
        other => {
            let mut object = Map::new();
            object.insert("status".to_string(), Value::String(status.to_string()));
            object.insert("output".to_string(), Value::String(output));
            object.insert("result".to_string(), other);
            Value::Object(object)
        }
    }
}

impl Default for ToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod subset_tests {
    use super::*;

    fn full() -> ToolRegistry {
        let mut r = ToolRegistry::new();
        register_default_tools(&mut r);
        r
    }

    #[test]
    fn subset_keeps_only_named_tools_and_preserves_callable() {
        let r = full();
        let sub = r.subset(&["read_file", "bash"]);
        assert!(sub.contains("read_file"));
        assert!(sub.contains("bash"));
        assert!(!sub.contains("write_file"));
        let mut names = sub.names();
        names.sort();
        assert_eq!(names, vec!["bash".to_string(), "read_file".to_string()]);
        // bash 的异步策略要随子集一起保留(callable/schema/async 都来自原 entry 的 clone)。
        assert!(sub.definitions().iter().any(|d| d.name == "bash"));
    }

    #[test]
    fn subset_ignores_unknown_names() {
        let sub = full().subset(&["read_file", "does_not_exist"]);
        assert!(sub.contains("read_file"));
        assert!(!sub.contains("does_not_exist"));
        assert_eq!(sub.names().len(), 1);
    }
}