subako-function-sdk 0.2.2

//! The guest side of a Subako function: a WASI preview 1 program the worker
//! runs with `argv = [package]` as an MCP server over MCP's stdio transport,
//! a JSON-RPC 2.0 message per line. It speaks the stateless MCP revision
//! `2026-07-28`, so every request carries its own `_meta` and there is no
//! handshake.
//!
//! On upload, core asks the module `server/discover` and `tools/list`; on a
//! call, the worker sends one `tools/call`. Either way stdin closes once the
//! last request is answered. The guest answers each request as it arrives and
//! exits 0 when stdin ends.
//!
//! A call's result is MCP's `CallToolResult`, with `content` the model reads,
//! `structuredContent` beside it, and `isError` set on a failure the model
//! should read. An image in the session volume travels as a `resource_link`
//! to `volume:///<path beneath /volume>`, which core reads and shows the
//! model. A JSON-RPC error, a line that is no message, or any other exit is a
//! crash.
//!
//! A package is written as typed functions. Each one's name, doc comment, and
//! input and output types become the tool `tools/list` names, and the program
//! prints the package's `subako-package.json` (its ABI, name, and
//! description) when run with `--subako-manifest`:
//!
//! ```no_run
//! use schemars::JsonSchema;
//! use serde::{Deserialize, Serialize};
//! use subako_function_sdk::{ToolError, function, package};
//!
//! #[derive(Deserialize, JsonSchema)]
//! struct GreetInput {
//!     /// Who to greet.
//!     name: String,
//! }
//!
//! #[derive(Serialize, JsonSchema)]
//! struct Greeting {
//!     message: String,
//! }
//!
//! /// Greet someone by name.
//! #[function]
//! fn greet(input: GreetInput) -> Result<Greeting, ToolError> {
//!     Ok(Greeting { message: format!("Hello, {}!", input.name) })
//! }
//!
//! package! {
//!     name: "hello",
//!     description: "Says hello.",
//!     functions: [greet],
//! }
//! ```
//!
//! The same crate built for the host prints the same manifest, so packaging
//! needs no WebAssembly runtime: `cargo run -- --subako-manifest`.
//!
//! A function that shows the model an image returns an [`Answer`]: its JSON,
//! and the image after it.
//!
//! ```no_run
//! # use schemars::JsonSchema;
//! # use serde::{Deserialize, Serialize};
//! # use subako_function_sdk::{Answer, ToolError, function};
//! # #[derive(Deserialize, JsonSchema)]
//! # struct PlotInput { title: String }
//! #[derive(Serialize, JsonSchema)]
//! struct Plot {
//!     path: String,
//! }
//!
//! /// Plot a chart and show it.
//! #[function]
//! fn plot(input: PlotInput) -> Result<Answer<Plot>, ToolError> {
//!     let path = "/volume/tmp/chart.png".to_string();
//!     // ... write the chart to `path` ...
//!     Ok(Answer::new(Plot { path: path.clone() }).with_image_file(path))
//! }
//! ```

// A panic in a guest surfaces as an opaque trap instead of a message the model
// can read, so every failure has to travel as a `ToolError`.
//
// Duplicated per crate: `[workspace.lints]` can't express the test exemption.
#![cfg_attr(
    not(test),
    deny(
        clippy::unwrap_used,
        clippy::expect_used,
        clippy::panic,
        clippy::unreachable,
        clippy::todo,
        clippy::unimplemented,
        clippy::indexing_slicing,
    )
)]

pub mod http;
mod package;
mod stdio;

pub use package::{ABI, Function, MANIFEST_FLAG, Manifest, Package};
pub use stdio::ServerInfo;

pub use subako_function_sdk_macros::{function, package};

use std::io::{BufRead as _, Write as _};
use std::process::ExitCode;

use schemars::JsonSchema;
use serde::Serialize;
use serde::de::DeserializeOwned;

/// A call's input, held until the tool says which shape it takes.
pub struct Input(serde_json::Value);

impl Input {
    pub fn parse<T: DeserializeOwned>(self) -> Result<T, ToolError> {
        serde_json::from_value(self.0).map_err(|e| ToolError::new(format!("invalid input: {e}")))
    }
}

/// A call's result: the blocks the model reads, and the JSON beside them.
#[derive(Debug, Clone, PartialEq)]
pub struct Output {
    content: Vec<Content>,
    structured: Option<serde_json::Value>,
}

impl Output {
    /// `value` as the JSON beside the result and, as text, the model's read.
    pub fn json<T: Serialize>(value: &T) -> Result<Output, ToolError> {
        let value = to_json(value)?;
        Ok(Output {
            content: vec![Content::Text(value.to_string())],
            structured: Some(value),
        })
    }

    /// Blocks the model reads, with no JSON beside them.
    pub fn content(content: Vec<Content>) -> Output {
        Output {
            content,
            structured: None,
        }
    }
}

/// One block of a result the model reads.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Content {
    Text(String),
    /// An image in the session volume, by its absolute path under
    /// `/volume`. It travels as MCP's `resource_link` to a `volume:///` URI;
    /// Core reads it after the call returns and shows it to the model, so it
    /// must still be there then.
    ImageFile(String),
    /// An image's bytes, base64-encoded, as MCP carries one. It counts
    /// against the call's stdout, so a large image goes as an
    /// [`ImageFile`](Content::ImageFile) instead.
    Image {
        mime_type: String,
        data: String,
    },
}

impl Content {
    pub fn text(text: impl Into<String>) -> Self {
        Content::Text(text.into())
    }

    pub fn image_file(path: impl Into<String>) -> Self {
        Content::ImageFile(path.into())
    }

    fn to_json(&self) -> serde_json::Value {
        match self {
            Content::Text(text) => serde_json::json!({"type": "text", "text": text}),
            Content::ImageFile(path) => {
                let name = path.rsplit('/').next().unwrap_or(path);
                serde_json::json!({"type": "resource_link", "uri": volume_uri(path), "name": name})
            }
            Content::Image { mime_type, data } => {
                serde_json::json!({"type": "image", "mimeType": mime_type, "data": data})
            }
        }
    }
}

/// A function's typed result with blocks after it: the model reads `value`
/// as JSON text first, then each block in order. Its output schema is
/// `T`'s.
#[derive(Debug, Clone, PartialEq)]
pub struct Answer<T> {
    value: T,
    content: Vec<Content>,
}

/// An answer starts from its value; each `with` method adds a block after
/// what it already holds.
impl<T> Answer<T> {
    pub fn new(value: T) -> Self {
        Answer {
            value,
            content: Vec::new(),
        }
    }

    pub fn with(mut self, content: Content) -> Self {
        self.content.push(content);
        self
    }

    /// Adds the image at `path` in the session volume, as
    /// [`Content::image_file`].
    pub fn with_image_file(self, path: impl Into<String>) -> Self {
        self.with(Content::image_file(path))
    }
}

/// What a function may return: a JSON value, or an [`Answer`] carrying blocks
/// beside one.
pub trait Reply {
    /// The schema of the JSON the result carries.
    fn schema(contract: schemars::generate::Contract) -> serde_json::Value;
    fn into_output(self) -> Result<Output, ToolError>;
}

impl<T: Serialize + JsonSchema> Reply for T {
    fn schema(contract: schemars::generate::Contract) -> serde_json::Value {
        package::schema::<T>(contract)
    }

    fn into_output(self) -> Result<Output, ToolError> {
        Output::json(&self)
    }
}

impl<T: Serialize + JsonSchema> Reply for Answer<T> {
    fn schema(contract: schemars::generate::Contract) -> serde_json::Value {
        package::schema::<T>(contract)
    }

    fn into_output(self) -> Result<Output, ToolError> {
        Output::json(&self.value).map(|output| Output {
            content: output.content.into_iter().chain(self.content).collect(),
            ..output
        })
    }
}

/// The characters a URI path segment keeps as they are.
const SEGMENT: &percent_encoding::AsciiSet = &percent_encoding::NON_ALPHANUMERIC
    .remove(b'-')
    .remove(b'.')
    .remove(b'_')
    .remove(b'~');

/// The URI naming `path`: `volume:///` and its path beneath `/volume`,
/// percent-encoded, when it is in the volume. Core shows anything else as a
/// line naming the link.
// CODESYNC(volume-uri)
fn volume_uri(path: &str) -> String {
    let encode = |path: &str| {
        path.split('/')
            .map(|segment| percent_encoding::utf8_percent_encode(segment, SEGMENT).to_string())
            .collect::<Vec<_>>()
            .join("/")
    };
    match path.strip_prefix("/volume/") {
        Some(inside) => format!("volume:///{}", encode(inside)),
        None => format!("file://{}", encode(path)),
    }
}

fn to_json<T: Serialize>(value: &T) -> Result<serde_json::Value, ToolError> {
    serde_json::to_value(value).map_err(|e| ToolError::new(format!("output is not JSON: {e}")))
}

/// A failure the model should read: the message is all it is told.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolError(String);

impl ToolError {
    pub fn new(message: impl Into<String>) -> Self {
        ToolError(message.into())
    }

    pub fn unknown_tool(tool: &str) -> Self {
        ToolError(format!("unknown tool: {tool}"))
    }

    pub fn message(&self) -> &str {
        &self.0
    }
}

impl From<String> for ToolError {
    fn from(message: String) -> Self {
        ToolError(message)
    }
}

impl From<&str> for ToolError {
    fn from(message: &str) -> Self {
        ToolError(message.to_string())
    }
}

impl From<std::io::Error> for ToolError {
    fn from(error: std::io::Error) -> Self {
        ToolError(error.to_string())
    }
}

/// What a guest serves: who it is, its tools, and how to call one.
pub trait Server {
    fn info(&self) -> ServerInfo<'_>;

    /// Each tool as MCP's `Tool`, always in the same order.
    fn tools(&self) -> Vec<serde_json::Value>;

    fn call(&self, tool: &str, input: Input) -> Result<Output, ToolError>;
}

/// Every exchange as a pure function of stdin, so a host test drives exactly
/// what the guest runs: the bytes are every line the guest writes.
pub fn execute(stdin: &[u8], server: &impl Server) -> (Vec<u8>, u8) {
    // Line by line as `run` reads them, each with its newline.
    let stdout = stdin
        .split_inclusive(|byte| *byte == b'\n')
        .flat_map(|line| serve(line, server))
        .collect();
    (stdout, 0)
}

/// Answers the message on `line`, if it asks for an answer.
fn serve(line: &[u8], server: &impl Server) -> Vec<u8> {
    match stdio::read(line) {
        stdio::Incoming::Discover { id } => stdio::discover_line(&id, &server.info()),
        stdio::Incoming::ListTools { id } => stdio::tools_line(&id, &server.info(), server.tools()),
        stdio::Incoming::Call(call) => {
            let result = server.call(&call.name, Input(call.arguments));
            stdio::result_line(&call.id, &server.info(), result_fields(result).into_iter())
        }
        stdio::Incoming::Ignored => Vec::new(),
        stdio::Incoming::Refused(refusal) => refusal.line(),
    }
}

/// The result as the fields of MCP's `CallToolResult`.
fn result_fields(result: Result<Output, ToolError>) -> Vec<(&'static str, serde_json::Value)> {
    match result {
        Ok(Output {
            content,
            structured,
        }) => {
            let content: Vec<_> = content.iter().map(Content::to_json).collect();
            let content = ("content", serde_json::json!(content));
            match structured {
                Some(structured) => vec![content, ("structuredContent", structured)],
                None => vec![content],
            }
        }
        Err(ToolError(message)) => vec![
            (
                "content",
                serde_json::json!([{"type": "text", "text": message}]),
            ),
            ("isError", serde_json::json!(true)),
        ],
    }
}

/// The guest's `main`: answers each line of stdin as it arrives, until stdin
/// ends.
// COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
pub fn run(server: &impl Server) -> ExitCode {
    // Neither handle is held while a tool runs, so the tool may use them too.
    loop {
        let mut line = Vec::new();
        let read = std::io::stdin().lock().read_until(b'\n', &mut line);
        let (answer, more) = match read {
            Ok(0) => return ExitCode::SUCCESS,
            Ok(_) => (serve(&line, server), true),
            Err(e) => (stdio::Refusal::unreadable(&e).line(), false),
        };
        let mut stdout = std::io::stdout().lock();
        match stdout.write_all(&answer).and_then(|()| stdout.flush()) {
            Ok(()) if more => {}
            Ok(()) => return ExitCode::SUCCESS,
            // stdout is the only channel back; a full one is reported by exit code.
            Err(_) => return ExitCode::from(2),
        }
    }
}