Skip to main content

oxi/extensions/
wasm_tool.rs

1//! AgentTool wrapper for WASM extension tools.
2//!
3//! Each tool exported by a WASM extension is wrapped in a `WasmTool` that
4//! implements the `AgentTool` trait. Tool execution is delegated to the
5//! WASM extension via `WasmExtensionManager::execute_tool`.
6
7use oxi_agent::{AgentTool, AgentToolResult, ToolContext, ToolError};
8use serde_json::Value;
9use std::sync::Arc;
10use tokio::sync::oneshot;
11
12use super::wasm::WasmExtensionManager;
13
14/// An `AgentTool` backed by a WASM extension.
15pub struct WasmTool {
16    manager: Arc<WasmExtensionManager>,
17    tool_name: String,
18    description: String,
19    schema: Value,
20}
21
22impl WasmTool {
23    /// Create a new WasmTool.
24    pub fn new(
25        manager: Arc<WasmExtensionManager>,
26        tool_name: String,
27        description: String,
28        schema: Value,
29    ) -> Self {
30        Self {
31            manager,
32            tool_name,
33            description,
34            schema,
35        }
36    }
37}
38
39#[async_trait::async_trait]
40impl AgentTool for WasmTool {
41    fn name(&self) -> &str {
42        &self.tool_name
43    }
44    fn label(&self) -> &str {
45        &self.description
46    }
47    fn description(&self) -> &str {
48        &self.description
49    }
50    fn parameters_schema(&self) -> Value {
51        self.schema.clone()
52    }
53
54    async fn execute(
55        &self,
56        _tool_call_id: &str,
57        params: Value,
58        _signal: Option<oneshot::Receiver<()>>,
59        _ctx: &ToolContext,
60    ) -> Result<AgentToolResult, ToolError> {
61        let manager = self.manager.clone();
62        let tool_name = self.tool_name.clone();
63
64        let result = tokio::task::spawn_blocking(move || manager.execute_tool(&tool_name, params))
65            .await
66            .map_err(|e| format!("WASM execution panicked: {}", e))?;
67
68        match result {
69            Ok(value) => {
70                let success = value
71                    .get("success")
72                    .and_then(|v| v.as_bool())
73                    .unwrap_or(true);
74                let output = value
75                    .get("output")
76                    .and_then(|v| v.as_str())
77                    .unwrap_or("(no output)")
78                    .to_string();
79
80                if success {
81                    let mut tool_result = AgentToolResult::success(output);
82                    if let Some(meta) = value.get("metadata").cloned() {
83                        tool_result = tool_result.with_metadata(meta);
84                    }
85                    Ok(tool_result)
86                } else {
87                    Ok(AgentToolResult::error(output))
88                }
89            }
90            Err(e) => Ok(AgentToolResult::error(format!(
91                "WASM tool '{}' error: {}",
92                self.tool_name, e
93            ))),
94        }
95    }
96}