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