Skip to main content

subako_function_sdk/
lib.rs

1//! The guest side of a Subako function: a WASI preview 1 program the worker
2//! runs with `argv = [package]` as an MCP server over MCP's stdio transport,
3//! a JSON-RPC 2.0 message per line. It speaks the stateless MCP revision
4//! `2026-07-28`, so every request carries its own `_meta` and there is no
5//! handshake.
6//!
7//! On upload, core asks the module `server/discover` and `tools/list`; on a
8//! call, the worker sends one `tools/call`. Either way stdin closes once the
9//! last request is answered. The guest answers each request as it arrives and
10//! exits 0 when stdin ends.
11//!
12//! A call's result is MCP's `CallToolResult`, with `content` the model reads,
13//! `structuredContent` beside it, and `isError` set on a failure the model
14//! should read. An image in the session volume travels as a `resource_link`
15//! to `volume:///<path beneath /volume>`, which core reads and shows the
16//! model. A JSON-RPC error, a line that is no message, or any other exit is a
17//! crash.
18//!
19//! A package is written as typed functions. Each one's name, doc comment, and
20//! input and output types become the tool `tools/list` names, and the program
21//! prints the package's `subako-package.json` (its ABI, name, and
22//! description) when run with `--subako-manifest`:
23//!
24//! ```no_run
25//! use schemars::JsonSchema;
26//! use serde::{Deserialize, Serialize};
27//! use subako_function_sdk::{ToolError, function, package};
28//!
29//! #[derive(Deserialize, JsonSchema)]
30//! struct GreetInput {
31//!     /// Who to greet.
32//!     name: String,
33//! }
34//!
35//! #[derive(Serialize, JsonSchema)]
36//! struct Greeting {
37//!     message: String,
38//! }
39//!
40//! /// Greet someone by name.
41//! #[function]
42//! fn greet(input: GreetInput) -> Result<Greeting, ToolError> {
43//!     Ok(Greeting { message: format!("Hello, {}!", input.name) })
44//! }
45//!
46//! package! {
47//!     name: "hello",
48//!     description: "Says hello.",
49//!     functions: [greet],
50//! }
51//! ```
52//!
53//! The same crate built for the host prints the same manifest, so packaging
54//! needs no WebAssembly runtime: `cargo run -- --subako-manifest`.
55//!
56//! A function that shows the model an image returns an [`Answer`]: its JSON,
57//! and the image after it.
58//!
59//! ```no_run
60//! # use schemars::JsonSchema;
61//! # use serde::{Deserialize, Serialize};
62//! # use subako_function_sdk::{Answer, ToolError, function};
63//! # #[derive(Deserialize, JsonSchema)]
64//! # struct PlotInput { title: String }
65//! #[derive(Serialize, JsonSchema)]
66//! struct Plot {
67//!     path: String,
68//! }
69//!
70//! /// Plot a chart and show it.
71//! #[function]
72//! fn plot(input: PlotInput) -> Result<Answer<Plot>, ToolError> {
73//!     let path = "/volume/tmp/chart.png".to_string();
74//!     // ... write the chart to `path` ...
75//!     Ok(Answer::new(Plot { path: path.clone() }).with_image_file(path))
76//! }
77//! ```
78
79// A panic in a guest surfaces as an opaque trap instead of a message the model
80// can read, so every failure has to travel as a `ToolError`.
81//
82// Duplicated per crate: `[workspace.lints]` can't express the test exemption.
83#![cfg_attr(
84    not(test),
85    deny(
86        clippy::unwrap_used,
87        clippy::expect_used,
88        clippy::panic,
89        clippy::unreachable,
90        clippy::todo,
91        clippy::unimplemented,
92        clippy::indexing_slicing,
93    )
94)]
95
96pub mod http;
97mod package;
98mod stdio;
99
100pub use package::{ABI, Function, MANIFEST_FLAG, Manifest, Package};
101pub use stdio::ServerInfo;
102
103pub use subako_function_sdk_macros::{function, package};
104
105use std::io::{BufRead as _, Write as _};
106use std::process::ExitCode;
107
108use schemars::JsonSchema;
109use serde::Serialize;
110use serde::de::DeserializeOwned;
111
112/// A call's input, held until the tool says which shape it takes.
113pub struct Input(serde_json::Value);
114
115impl Input {
116    pub fn parse<T: DeserializeOwned>(self) -> Result<T, ToolError> {
117        serde_json::from_value(self.0).map_err(|e| ToolError::new(format!("invalid input: {e}")))
118    }
119}
120
121/// A call's result: the blocks the model reads, and the JSON beside them.
122#[derive(Debug, Clone, PartialEq)]
123pub struct Output {
124    content: Vec<Content>,
125    structured: Option<serde_json::Value>,
126}
127
128impl Output {
129    /// `value` as the JSON beside the result and, as text, the model's read.
130    pub fn json<T: Serialize>(value: &T) -> Result<Output, ToolError> {
131        let value = to_json(value)?;
132        Ok(Output {
133            content: vec![Content::Text(value.to_string())],
134            structured: Some(value),
135        })
136    }
137
138    /// Blocks the model reads, with no JSON beside them.
139    pub fn content(content: Vec<Content>) -> Output {
140        Output {
141            content,
142            structured: None,
143        }
144    }
145}
146
147/// One block of a result the model reads.
148#[derive(Debug, Clone, PartialEq, Eq)]
149pub enum Content {
150    Text(String),
151    /// An image in the session volume, by its absolute path under
152    /// `/volume`. It travels as MCP's `resource_link` to a `volume:///` URI;
153    /// Core reads it after the call returns and shows it to the model, so it
154    /// must still be there then.
155    ImageFile(String),
156    /// An image's bytes, base64-encoded, as MCP carries one. It counts
157    /// against the call's stdout, so a large image goes as an
158    /// [`ImageFile`](Content::ImageFile) instead.
159    Image {
160        mime_type: String,
161        data: String,
162    },
163}
164
165impl Content {
166    pub fn text(text: impl Into<String>) -> Self {
167        Content::Text(text.into())
168    }
169
170    pub fn image_file(path: impl Into<String>) -> Self {
171        Content::ImageFile(path.into())
172    }
173
174    fn to_json(&self) -> serde_json::Value {
175        match self {
176            Content::Text(text) => serde_json::json!({"type": "text", "text": text}),
177            Content::ImageFile(path) => {
178                let name = path.rsplit('/').next().unwrap_or(path);
179                serde_json::json!({"type": "resource_link", "uri": volume_uri(path), "name": name})
180            }
181            Content::Image { mime_type, data } => {
182                serde_json::json!({"type": "image", "mimeType": mime_type, "data": data})
183            }
184        }
185    }
186}
187
188/// A function's typed result with blocks after it: the model reads `value`
189/// as JSON text first, then each block in order. Its output schema is
190/// `T`'s.
191#[derive(Debug, Clone, PartialEq)]
192pub struct Answer<T> {
193    value: T,
194    content: Vec<Content>,
195}
196
197/// An answer starts from its value; each `with` method adds a block after
198/// what it already holds.
199impl<T> Answer<T> {
200    pub fn new(value: T) -> Self {
201        Answer {
202            value,
203            content: Vec::new(),
204        }
205    }
206
207    pub fn with(mut self, content: Content) -> Self {
208        self.content.push(content);
209        self
210    }
211
212    /// Adds the image at `path` in the session volume, as
213    /// [`Content::image_file`].
214    pub fn with_image_file(self, path: impl Into<String>) -> Self {
215        self.with(Content::image_file(path))
216    }
217}
218
219/// What a function may return: a JSON value, or an [`Answer`] carrying blocks
220/// beside one.
221pub trait Reply {
222    /// The schema of the JSON the result carries.
223    fn schema(contract: schemars::generate::Contract) -> serde_json::Value;
224    fn into_output(self) -> Result<Output, ToolError>;
225}
226
227impl<T: Serialize + JsonSchema> Reply for T {
228    fn schema(contract: schemars::generate::Contract) -> serde_json::Value {
229        package::schema::<T>(contract)
230    }
231
232    fn into_output(self) -> Result<Output, ToolError> {
233        Output::json(&self)
234    }
235}
236
237impl<T: Serialize + JsonSchema> Reply for Answer<T> {
238    fn schema(contract: schemars::generate::Contract) -> serde_json::Value {
239        package::schema::<T>(contract)
240    }
241
242    fn into_output(self) -> Result<Output, ToolError> {
243        Output::json(&self.value).map(|output| Output {
244            content: output.content.into_iter().chain(self.content).collect(),
245            ..output
246        })
247    }
248}
249
250/// The characters a URI path segment keeps as they are.
251const SEGMENT: &percent_encoding::AsciiSet = &percent_encoding::NON_ALPHANUMERIC
252    .remove(b'-')
253    .remove(b'.')
254    .remove(b'_')
255    .remove(b'~');
256
257/// The URI naming `path`: `volume:///` and its path beneath `/volume`,
258/// percent-encoded, when it is in the volume. Core shows anything else as a
259/// line naming the link.
260// CODESYNC(volume-uri)
261fn volume_uri(path: &str) -> String {
262    let encode = |path: &str| {
263        path.split('/')
264            .map(|segment| percent_encoding::utf8_percent_encode(segment, SEGMENT).to_string())
265            .collect::<Vec<_>>()
266            .join("/")
267    };
268    match path.strip_prefix("/volume/") {
269        Some(inside) => format!("volume:///{}", encode(inside)),
270        None => format!("file://{}", encode(path)),
271    }
272}
273
274fn to_json<T: Serialize>(value: &T) -> Result<serde_json::Value, ToolError> {
275    serde_json::to_value(value).map_err(|e| ToolError::new(format!("output is not JSON: {e}")))
276}
277
278/// A failure the model should read: the message is all it is told.
279#[derive(Debug, Clone, PartialEq, Eq)]
280pub struct ToolError(String);
281
282impl ToolError {
283    pub fn new(message: impl Into<String>) -> Self {
284        ToolError(message.into())
285    }
286
287    pub fn unknown_tool(tool: &str) -> Self {
288        ToolError(format!("unknown tool: {tool}"))
289    }
290
291    pub fn message(&self) -> &str {
292        &self.0
293    }
294}
295
296impl From<String> for ToolError {
297    fn from(message: String) -> Self {
298        ToolError(message)
299    }
300}
301
302impl From<&str> for ToolError {
303    fn from(message: &str) -> Self {
304        ToolError(message.to_string())
305    }
306}
307
308impl From<std::io::Error> for ToolError {
309    fn from(error: std::io::Error) -> Self {
310        ToolError(error.to_string())
311    }
312}
313
314/// What a guest serves: who it is, its tools, and how to call one.
315pub trait Server {
316    fn info(&self) -> ServerInfo<'_>;
317
318    /// Each tool as MCP's `Tool`, always in the same order.
319    fn tools(&self) -> Vec<serde_json::Value>;
320
321    fn call(&self, tool: &str, input: Input) -> Result<Output, ToolError>;
322}
323
324/// Every exchange as a pure function of stdin, so a host test drives exactly
325/// what the guest runs: the bytes are every line the guest writes.
326pub fn execute(stdin: &[u8], server: &impl Server) -> (Vec<u8>, u8) {
327    // Line by line as `run` reads them, each with its newline.
328    let stdout = stdin
329        .split_inclusive(|byte| *byte == b'\n')
330        .flat_map(|line| serve(line, server))
331        .collect();
332    (stdout, 0)
333}
334
335/// Answers the message on `line`, if it asks for an answer.
336fn serve(line: &[u8], server: &impl Server) -> Vec<u8> {
337    match stdio::read(line) {
338        stdio::Incoming::Discover { id } => stdio::discover_line(&id, &server.info()),
339        stdio::Incoming::ListTools { id } => stdio::tools_line(&id, &server.info(), server.tools()),
340        stdio::Incoming::Call(call) => {
341            let result = server.call(&call.name, Input(call.arguments));
342            stdio::result_line(&call.id, &server.info(), result_fields(result).into_iter())
343        }
344        stdio::Incoming::Ignored => Vec::new(),
345        stdio::Incoming::Refused(refusal) => refusal.line(),
346    }
347}
348
349/// The result as the fields of MCP's `CallToolResult`.
350fn result_fields(result: Result<Output, ToolError>) -> Vec<(&'static str, serde_json::Value)> {
351    match result {
352        Ok(Output {
353            content,
354            structured,
355        }) => {
356            let content: Vec<_> = content.iter().map(Content::to_json).collect();
357            let content = ("content", serde_json::json!(content));
358            match structured {
359                Some(structured) => vec![content, ("structuredContent", structured)],
360                None => vec![content],
361            }
362        }
363        Err(ToolError(message)) => vec![
364            (
365                "content",
366                serde_json::json!([{"type": "text", "text": message}]),
367            ),
368            ("isError", serde_json::json!(true)),
369        ],
370    }
371}
372
373/// The guest's `main`: answers each line of stdin as it arrives, until stdin
374/// ends.
375// COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
376pub fn run(server: &impl Server) -> ExitCode {
377    // Neither handle is held while a tool runs, so the tool may use them too.
378    loop {
379        let mut line = Vec::new();
380        let read = std::io::stdin().lock().read_until(b'\n', &mut line);
381        let (answer, more) = match read {
382            Ok(0) => return ExitCode::SUCCESS,
383            Ok(_) => (serve(&line, server), true),
384            Err(e) => (stdio::Refusal::unreadable(&e).line(), false),
385        };
386        let mut stdout = std::io::stdout().lock();
387        match stdout.write_all(&answer).and_then(|()| stdout.flush()) {
388            Ok(()) if more => {}
389            Ok(()) => return ExitCode::SUCCESS,
390            // stdout is the only channel back; a full one is reported by exit code.
391            Err(_) => return ExitCode::from(2),
392        }
393    }
394}