subako-function-sdk 0.2.0

//! The guest side of a Subako function: a WASI preview 1 program the worker
//! runs with `argv = [package, tool]`, the tool's input JSON on stdin, and the
//! output JSON expected on stdout. Exit 0 is a result; exit 1 carries
//! `{"error": ...}`, which the model reads as a failed call.
//!
//! A package is written as typed functions. Each one's name, doc comment, and
//! input and output types become its tool, and the program prints the
//! package's `subako-package.json` 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 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;

pub use package::{Function, MANIFEST_FLAG, Manifest, ManifestFunction, Package};
pub use subako_function_sdk_macros::{function, package};

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

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, already JSON.
pub struct Output(serde_json::Value);

impl Output {
    pub fn json<T: Serialize>(value: &T) -> Result<Output, ToolError> {
        serde_json::to_value(value)
            .map(Output)
            .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())
    }
}

/// The whole call as a pure function of its arguments and stdin, so a host
/// test drives exactly what the guest runs.
pub fn execute(
    args: &[String],
    stdin: &[u8],
    route: impl FnOnce(&str, Input) -> Result<Output, ToolError>,
) -> (Vec<u8>, u8) {
    let result = match (args.get(1), serde_json::from_slice(stdin)) {
        (None, _) => Err(ToolError::new("no tool named on the command line")),
        (Some(_), Err(e)) => Err(ToolError::new(format!("input is not JSON: {e}"))),
        (Some(tool), Ok(value)) => route(tool, Input(value)),
    };
    match result {
        Ok(Output(value)) => (encode(&value), 0),
        Err(ToolError(message)) => (encode(&serde_json::json!({ "error": message })), 1),
    }
}

/// The guest's `main`: reads argv and stdin, routes, writes stdout.
// COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
pub fn run(route: impl FnOnce(&str, Input) -> Result<Output, ToolError>) -> ExitCode {
    let args: Vec<String> = std::env::args().collect();
    let mut stdin = Vec::new();
    if let Err(e) = std::io::stdin().read_to_end(&mut stdin) {
        let (out, code) = (
            encode(&serde_json::json!({ "error": format!("stdin unreadable: {e}") })),
            1,
        );
        return finish(&out, code);
    }
    let (out, code) = execute(&args, &stdin, route);
    finish(&out, code)
}

// COVERAGE(untestable): guest code runs under wasmtime, where the host's llvm-cov cannot see it.
fn finish(out: &[u8], code: u8) -> ExitCode {
    let mut stdout = std::io::stdout().lock();
    match stdout.write_all(out).and_then(|()| stdout.flush()) {
        Ok(()) => ExitCode::from(code),
        // stdout is the only channel back; a full one is reported by exit code.
        Err(_) => ExitCode::from(2),
    }
}

fn encode(value: &serde_json::Value) -> Vec<u8> {
    // A `Value` always serializes.
    serde_json::to_vec(value).unwrap_or_default()
}