140 lines
4.1 KiB
Rust
140 lines
4.1 KiB
Rust
//! 工具抽象接口与执行上下文。
|
|
|
|
use std::sync::Arc;
|
|
|
|
use async_trait::async_trait;
|
|
use serde_json::Value;
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
use crate::tools::error::ToolError;
|
|
use crate::tools::permission::Permission;
|
|
|
|
/// 工具执行上下文 —— 携带每次执行的运行时信息。
|
|
///
|
|
/// 字段在 Phase 2 即注入 `execute()` 签名中，防止后续扩展时出现
|
|
/// breaking change。后续阶段可扩展字段（如 `progress`、`shared_state`），
|
|
/// 但已有工具实现无需修改。
|
|
#[derive(Debug)]
|
|
pub struct ToolContext<'a> {
|
|
/// 当前对话/会话 ID，用于关联性追踪。
|
|
pub session_id: &'a str,
|
|
/// 链路追踪 ID，用于跨工具调用的耗时分布。
|
|
pub trace_id: &'a str,
|
|
/// 取消令牌，用于优雅取消正在执行的工具。
|
|
pub cancellation_token: CancellationToken,
|
|
}
|
|
|
|
impl<'a> ToolContext<'a> {
|
|
/// 创建一个新的工具执行上下文。
|
|
pub fn new(session_id: &'a str, trace_id: &'a str) -> Self {
|
|
Self {
|
|
session_id,
|
|
trace_id,
|
|
cancellation_token: CancellationToken::new(),
|
|
}
|
|
}
|
|
|
|
/// 创建一个使用给定取消令牌的上下文。
|
|
pub fn with_cancellation_token(
|
|
session_id: &'a str,
|
|
trace_id: &'a str,
|
|
token: CancellationToken,
|
|
) -> Self {
|
|
Self {
|
|
session_id,
|
|
trace_id,
|
|
cancellation_token: token,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 工具抽象接口 —— 所有工具（自定义或 MCP）最终都实现此 trait。
|
|
#[async_trait]
|
|
pub trait BaseTool: Send + Sync {
|
|
/// 工具名称（唯一标识，用于 LLM 的 tool_calls.name 匹配）。
|
|
fn name(&self) -> &str;
|
|
|
|
/// 工具描述（LLM 据此决定是否调用此工具）。
|
|
fn description(&self) -> &str;
|
|
|
|
/// 工具参数定义（JSON Schema 格式，传递给 LLM 的 tool.parameters）。
|
|
fn parameters(&self) -> Value;
|
|
|
|
/// 声明工具所需的权限列表。
|
|
fn required_permissions(&self) -> Vec<Permission> {
|
|
Vec::new()
|
|
}
|
|
|
|
/// 执行工具调用。
|
|
///
|
|
/// `ctx` 携带执行上下文（session_id、trace_id、cancellation_token），
|
|
/// 工具实现可在执行期间检查 `ctx.cancellation_token` 来支持优雅取消。
|
|
async fn execute(&self, args: Value, ctx: &ToolContext<'_>) -> Result<Value, ToolError>;
|
|
}
|
|
|
|
/// 为 `Arc<dyn BaseTool>` 提供 `Send + Sync` 包装，便于在 `Vec<Arc<dyn BaseTool>>` 中使用。
|
|
pub type ToolRef = Arc<dyn BaseTool>;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use serde_json::json;
|
|
|
|
struct EchoTool;
|
|
|
|
#[async_trait]
|
|
impl BaseTool for EchoTool {
|
|
fn name(&self) -> &str {
|
|
"echo"
|
|
}
|
|
|
|
fn description(&self) -> &str {
|
|
"回显输入"
|
|
}
|
|
|
|
fn parameters(&self) -> Value {
|
|
json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"text": { "type": "string" }
|
|
},
|
|
"required": ["text"]
|
|
})
|
|
}
|
|
|
|
async fn execute(&self, args: Value, _ctx: &ToolContext<'_>) -> Result<Value, ToolError> {
|
|
Ok(args)
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_mock_tool_execute() {
|
|
let tool = EchoTool;
|
|
let ctx = ToolContext::new("session-1", "trace-1");
|
|
let result = tool.execute(json!({"text": "hello"}), &ctx).await.unwrap();
|
|
assert_eq!(result, json!({"text": "hello"}));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_default_permissions_empty() {
|
|
let tool = EchoTool;
|
|
assert!(tool.required_permissions().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_context_creation() {
|
|
let ctx = ToolContext::new("s1", "t1");
|
|
assert_eq!(ctx.session_id, "s1");
|
|
assert_eq!(ctx.trace_id, "t1");
|
|
assert!(!ctx.cancellation_token.is_cancelled());
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_context_cancellation() {
|
|
let token = CancellationToken::new();
|
|
token.cancel();
|
|
let ctx = ToolContext::with_cancellation_token("s1", "t1", token);
|
|
assert!(ctx.cancellation_token.is_cancelled());
|
|
}
|
|
}
|