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, tool]`, the tool's input JSON on stdin, and the
3//! output JSON expected on stdout. Exit 0 is a result; exit 1 carries
4//! `{"error": ...}`, which the model reads as a failed call.
5//!
6//! A package is written as typed functions. Each one's name, doc comment, and
7//! input and output types become its tool, and the program prints the
8//! package's `subako-package.json` when run with `--subako-manifest`:
9//!
10//! ```no_run
11//! use schemars::JsonSchema;
12//! use serde::{Deserialize, Serialize};
13//! use subako_function_sdk::{ToolError, function, package};
14//!
15//! #[derive(Deserialize, JsonSchema)]
16//! struct GreetInput {
17//!     /// Who to greet.
18//!     name: String,
19//! }
20//!
21//! #[derive(Serialize, JsonSchema)]
22//! struct Greeting {
23//!     message: String,
24//! }
25//!
26//! /// Greet someone by name.
27//! #[function]
28//! fn greet(input: GreetInput) -> Result<Greeting, ToolError> {
29//!     Ok(Greeting { message: format!("Hello, {}!", input.name) })
30//! }
31//!
32//! package! {
33//!     name: "hello",
34//!     description: "Says hello.",
35//!     functions: [greet],
36//! }
37//! ```
38//!
39//! The same crate built for the host prints the same manifest, so packaging
40//! needs no WebAssembly runtime: `cargo run -- --subako-manifest`.
41
42// A panic in a guest surfaces as an opaque trap instead of a message the model
43// can read, so every failure has to travel as a `ToolError`.
44//
45// Duplicated per crate: `[workspace.lints]` can't express the test exemption.
46#![cfg_attr(
47    not(test),
48    deny(
49        clippy::unwrap_used,
50        clippy::expect_used,
51        clippy::panic,
52        clippy::unreachable,
53        clippy::todo,
54        clippy::unimplemented,
55        clippy::indexing_slicing,
56    )
57)]
58
59pub mod http;
60mod package;
61
62pub use package::{Function, MANIFEST_FLAG, Manifest, ManifestFunction, Package};
63pub use subako_function_sdk_macros::{function, package};
64
65use std::io::{Read as _, Write as _};
66use std::process::ExitCode;
67
68use serde::Serialize;
69use serde::de::DeserializeOwned;
70
71/// A call's input, held until the tool says which shape it takes.
72pub struct Input(serde_json::Value);
73
74impl Input {
75    pub fn parse<T: DeserializeOwned>(self) -> Result<T, ToolError> {
76        serde_json::from_value(self.0).map_err(|e| ToolError::new(format!("invalid input: {e}")))
77    }
78}
79
80/// A call's result, already JSON.
81pub struct Output(serde_json::Value);
82
83impl Output {
84    pub fn json<T: Serialize>(value: &T) -> Result<Output, ToolError> {
85        serde_json::to_value(value)
86            .map(Output)
87            .map_err(|e| ToolError::new(format!("output is not JSON: {e}")))
88    }
89}
90
91/// A failure the model should read: the message is all it is told.
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct ToolError(String);
94
95impl ToolError {
96    pub fn new(message: impl Into<String>) -> Self {
97        ToolError(message.into())
98    }
99
100    pub fn unknown_tool(tool: &str) -> Self {
101        ToolError(format!("unknown tool: {tool}"))
102    }
103
104    pub fn message(&self) -> &str {
105        &self.0
106    }
107}
108
109impl From<String> for ToolError {
110    fn from(message: String) -> Self {
111        ToolError(message)
112    }
113}
114
115impl From<&str> for ToolError {
116    fn from(message: &str) -> Self {
117        ToolError(message.to_string())
118    }
119}
120
121impl From<std::io::Error> for ToolError {
122    fn from(error: std::io::Error) -> Self {
123        ToolError(error.to_string())
124    }
125}
126
127/// The whole call as a pure function of its arguments and stdin, so a host
128/// test drives exactly what the guest runs.
129pub fn execute(
130    args: &[String],
131    stdin: &[u8],
132    route: impl FnOnce(&str, Input) -> Result<Output, ToolError>,
133) -> (Vec<u8>, u8) {
134    let result = match (args.get(1), serde_json::from_slice(stdin)) {
135        (None, _) => Err(ToolError::new("no tool named on the command line")),
136        (Some(_), Err(e)) => Err(ToolError::new(format!("input is not JSON: {e}"))),
137        (Some(tool), Ok(value)) => route(tool, Input(value)),
138    };
139    match result {
140        Ok(Output(value)) => (encode(&value), 0),
141        Err(ToolError(message)) => (encode(&serde_json::json!({ "error": message })), 1),
142    }
143}
144
145/// The guest's `main`: reads argv and stdin, routes, writes stdout.
146// COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
147pub fn run(route: impl FnOnce(&str, Input) -> Result<Output, ToolError>) -> ExitCode {
148    let args: Vec<String> = std::env::args().collect();
149    let mut stdin = Vec::new();
150    if let Err(e) = std::io::stdin().read_to_end(&mut stdin) {
151        let (out, code) = (
152            encode(&serde_json::json!({ "error": format!("stdin unreadable: {e}") })),
153            1,
154        );
155        return finish(&out, code);
156    }
157    let (out, code) = execute(&args, &stdin, route);
158    finish(&out, code)
159}
160
161// COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
162fn finish(out: &[u8], code: u8) -> ExitCode {
163    let mut stdout = std::io::stdout().lock();
164    match stdout.write_all(out).and_then(|()| stdout.flush()) {
165        Ok(()) => ExitCode::from(code),
166        // stdout is the only channel back; a full one is reported by exit code.
167        Err(_) => ExitCode::from(2),
168    }
169}
170
171fn encode(value: &serde_json::Value) -> Vec<u8> {
172    // A `Value` always serializes.
173    serde_json::to_vec(value).unwrap_or_default()
174}
subako_function_sdk/lib.rs

subako_function_sdk/
lib.rs
