subako-function-sdk 0.2.3

//! A package as a value: its functions, the manifest they describe, and the
//! `main` that serves them.

use std::io::Write as _;
use std::process::ExitCode;

use schemars::JsonSchema;
use schemars::generate::{Contract, SchemaSettings};
use serde::Serialize;
use serde::de::DeserializeOwned;

use crate::{Input, Output, Reply, Server, ServerInfo, ToolError};

/// The argument that prints the manifest instead of calling a function.
pub const MANIFEST_FLAG: &str = "--subako-manifest";

/// One function a package serves as a tool.
pub struct Function {
    name: String,
    description: String,
    input_schema: fn(Contract) -> serde_json::Value,
    output_schema: fn(Contract) -> serde_json::Value,
    call: Box<dyn Fn(Input) -> Result<Output, ToolError>>,
}

impl Function {
    /// `f` as the tool `name`. The schemas come from `I` and `O`.
    pub fn new<I, O, E>(
        name: impl Into<String>,
        description: impl Into<String>,
        f: impl Fn(I) -> Result<O, E> + 'static,
    ) -> Self
    where
        I: DeserializeOwned + JsonSchema,
        O: Reply,
        E: Into<ToolError>,
    {
        Function {
            name: name.into(),
            description: description.into(),
            input_schema: schema::<I>,
            output_schema: O::schema,
            call: Box::new(move |input| f(input.parse()?).map_err(Into::into)?.into_output()),
        }
    }

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

/// The function ABI this SDK's guests speak: MCP 2026-07-28 over stdio.
// CODESYNC(function-abi)
pub const ABI: u32 = 1;

/// What `subako-package.json` holds: the ABI the module speaks, and the
/// package's name and description. Its tools are what `tools/list` answers.
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct Manifest {
    pub abi: u32,
    pub name: String,
    pub description: String,
}

/// A package: its name, its description, and the functions it serves.
pub struct Package {
    name: String,
    version: String,
    description: String,
    functions: Vec<Function>,
}

impl Package {
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        functions: Vec<Function>,
    ) -> Self {
        Package {
            name: name.into(),
            version: "0.0.0".to_string(),
            description: description.into(),
            functions,
        }
    }

    /// The version each result's `serverInfo` reports; `package!` sets the
    /// crate's own.
    pub fn with_version(self, version: impl Into<String>) -> Self {
        Package {
            version: version.into(),
            ..self
        }
    }

    pub fn functions(&self) -> &[Function] {
        &self.functions
    }

    pub fn manifest(&self) -> Manifest {
        Manifest {
            abi: ABI,
            name: self.name.clone(),
            description: self.description.clone(),
        }
    }

    /// Calls the function named `tool`.
    pub fn call(&self, tool: &str, input: Input) -> Result<Output, ToolError> {
        match self.functions.iter().find(|f| f.name == tool) {
            Some(function) => (function.call)(input),
            None => Err(ToolError::unknown_tool(tool)),
        }
    }

    /// The whole process as a pure function of its arguments and stdin: the
    /// manifest under [`MANIFEST_FLAG`], otherwise a call.
    pub fn execute(&self, args: &[String], stdin: &[u8]) -> (Vec<u8>, u8) {
        match args.get(1).map(String::as_str) {
            Some(MANIFEST_FLAG) => (self.manifest_json(), 0),
            _ => crate::execute(stdin, self),
        }
    }

    /// The program's `main`. The manifest is printed without reading stdin,
    /// so `module --subako-manifest` needs no input.
    // COVERAGE(untestable): reads the real process's arguments and stdio; `execute` is its testable half.
    pub fn run(&self) -> ExitCode {
        let args: Vec<String> = std::env::args().collect();
        match args.get(1).map(String::as_str) {
            Some(MANIFEST_FLAG) => {
                let mut stdout = std::io::stdout().lock();
                match stdout
                    .write_all(&self.manifest_json())
                    .and_then(|()| stdout.flush())
                {
                    Ok(()) => ExitCode::SUCCESS,
                    Err(_) => ExitCode::from(2),
                }
            }
            _ => crate::run(self),
        }
    }

    fn manifest_json(&self) -> Vec<u8> {
        // A manifest is plain data, so it always serializes.
        let mut json = serde_json::to_vec_pretty(&self.manifest()).unwrap_or_default();
        json.push(b'\n');
        json
    }
}

impl Server for Package {
    fn info(&self) -> ServerInfo<'_> {
        ServerInfo {
            name: &self.name,
            version: &self.version,
        }
    }

    fn tools(&self) -> Vec<serde_json::Value> {
        self.functions
            .iter()
            .map(|f| {
                serde_json::json!({
                    "name": f.name,
                    "description": f.description,
                    "inputSchema": (f.input_schema)(Contract::Deserialize),
                    "outputSchema": (f.output_schema)(Contract::Serialize),
                })
            })
            .collect()
    }

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

/// `T`'s JSON Schema as `contract` reads or writes it, with every subschema
/// inline so a model reads one self-contained object.
pub(crate) fn schema<T: JsonSchema>(contract: Contract) -> serde_json::Value {
    let mut schema = SchemaSettings::draft2020_12()
        .with(|s| {
            s.inline_subschemas = true;
            s.meta_schema = None;
            s.contract = contract;
        })
        .into_generator()
        .into_root_schema_for::<T>();
    schema.remove("title");
    schema.to_value()
}