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 anyhow::Result;
8use async_trait::async_trait;
9use oxi_agent::{AgentTool, AgentToolResult, ToolError};
10use serde_json::Value;
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
41#[async_trait]
42impl AgentTool for WasmTool {
43    fn name(&self) -> &str {
44        &self.tool_name
45    }
46    fn label(&self) -> &str {
47        &self.description
48    }
49    fn description(&self) -> &str {
50        &self.description
51    }
52    fn parameters_schema(&self) -> Value {
53        self.schema.clone()
54    }
55
56    async fn execute(
57        &self,
58        _tool_call_id: &str,
59        params: Value,
60        _signal: Option<oneshot::Receiver<()>>,
61        _ctx: &oxi_agent::ToolContext,
62    ) -> Result<AgentToolResult, ToolError> {
63        let manager = self.manager.clone();
64        let tool_name = self.tool_name.clone();
65
66        let result = tokio::task::spawn_blocking(move || manager.execute_tool(&tool_name, params))
67            .await
68            .map_err(|e| format!("WASM execution panicked: {}", e))?;
69
70        match result {
71            Ok(value) => {
72                let success = value
73                    .get("success")
74                    .and_then(|v| v.as_bool())
75                    .unwrap_or(true);
76                let output = value
77                    .get("output")
78                    .and_then(|v| v.as_str())
79                    .unwrap_or("(no output)")
80                    .to_string();
81
82                if success {
83                    let mut tool_result = AgentToolResult::success(output);
84                    if let Some(meta) = value.get("metadata").cloned() {
85                        tool_result = tool_result.with_metadata(meta);
86                    }
87                    Ok(tool_result)
88                } else {
89                    Ok(AgentToolResult::error(output))
90                }
91            }
92            Err(e) => Ok(AgentToolResult::error(format!(
93                "WASM tool '{}' error: {}",
94                self.tool_name, e
95            ))),
96        }
97    }
98}