1#![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
112pub 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#[derive(Debug, Clone, PartialEq)]
123pub struct Output {
124 content: Vec<Content>,
125 structured: Option<serde_json::Value>,
126}
127
128impl Output {
129 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 pub fn content(content: Vec<Content>) -> Output {
140 Output {
141 content,
142 structured: None,
143 }
144 }
145}
146
147#[derive(Debug, Clone, PartialEq, Eq)]
149pub enum Content {
150 Text(String),
151 ImageFile(String),
156 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#[derive(Debug, Clone, PartialEq)]
192pub struct Answer<T> {
193 value: T,
194 content: Vec<Content>,
195}
196
197impl<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 pub fn with_image_file(self, path: impl Into<String>) -> Self {
215 self.with(Content::image_file(path))
216 }
217}
218
219pub trait Reply {
222 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
250const SEGMENT: &percent_encoding::AsciiSet = &percent_encoding::NON_ALPHANUMERIC
252 .remove(b'-')
253 .remove(b'.')
254 .remove(b'_')
255 .remove(b'~');
256
257fn 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#[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
314pub trait Server {
316 fn info(&self) -> ServerInfo<'_>;
317
318 fn tools(&self) -> Vec<serde_json::Value>;
320
321 fn call(&self, tool: &str, input: Input) -> Result<Output, ToolError>;
322}
323
324pub fn execute(stdin: &[u8], server: &impl Server) -> (Vec<u8>, u8) {
327 let stdout = stdin
329 .split_inclusive(|byte| *byte == b'\n')
330 .flat_map(|line| serve(line, server))
331 .collect();
332 (stdout, 0)
333}
334
335fn 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
349fn 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
373pub fn run(server: &impl Server) -> ExitCode {
377 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 Err(_) => return ExitCode::from(2),
392 }
393 }
394}