OXYGENHUB
# 835ceb9b3ea7 oxygen / crates / oxy-core / src / mcp.rs
12.0 KiB · Rust · last changed oxy mcp: an MCP server for a repository, signing as its configured identity (closes #159) 4 days ago source blame history raw
1
//! A Model Context Protocol server over stdio (oxygen#159): JSON-RPC 2.0,
2
//! one message per line, as MCP's stdio transport has it. This is only the
3
//! protocol: `initialize`, `ping`, `tools/list` and `tools/call`, with the
4
//! tools themselves supplied by the caller (`oxy mcp`, and `oh mcp` in the
5
//! hub, which reaches this module through its pinned oxygen).
6
//!
7
//! Nothing is written to stdout but protocol messages: a tool's own
8
//! output goes into its result.
9
10
use serde_json::{json, Value};
11
use std::io::{BufRead, Write};
12
13
/// The protocol versions this server speaks, newest first. A client asking
14
/// for one of them gets it back; any other gets the newest.
15
pub const PROTOCOL_VERSIONS: &[&str] = &["2025-06-18", "2025-03-26", "2024-11-05"];
16
17
/// One tool as `tools/list` describes it.
18
#[derive(Clone, Debug)]
19
pub struct Tool {
20
    pub name: String,
21
    pub description: String,
22
    /// A JSON Schema object for the arguments.
23
    pub input_schema: Value,
24
    /// Reads only: no change to the repository, the remote or the hub.
25
    pub read_only: bool,
26
    /// May lose work or history (undo, abandon, deletes).
27
    pub destructive: bool,
28
}
29
30
impl Tool {
31
    fn describe(&self) -> Value {
32
        json!({
33
            "name": self.name,
34
            "description": self.description,
35
            "inputSchema": self.input_schema,
36
            "annotations": {
37
                "readOnlyHint": self.read_only,
38
                "destructiveHint": self.destructive,
39
            },
40
        })
41
    }
42
}
43
44
/// What a tool call produced: text for the model to read, the same as JSON
45
/// when there is some, and whether it failed. A failed tool is a result, not
46
/// a protocol error, so the model sees why.
47
#[derive(Clone, Debug, Default, PartialEq)]
48
pub struct ToolOutput {
49
    pub text: String,
50
    pub structured: Option<Value>,
51
    pub is_error: bool,
52
}
53
54
impl ToolOutput {
55
    pub fn ok(text: impl Into<String>) -> ToolOutput {
56
        ToolOutput {
57
            text: text.into(),
58
            structured: None,
59
            is_error: false,
60
        }
61
    }
62
63
    pub fn error(text: impl Into<String>) -> ToolOutput {
64
        ToolOutput {
65
            text: text.into(),
66
            structured: None,
67
            is_error: true,
68
        }
69
    }
70
71
    fn to_json(&self) -> Value {
72
        let mut out = json!({
73
            "content": [{ "type": "text", "text": self.text }],
74
            "isError": self.is_error,
75
        });
76
        // Structured content must be an object; an array is wrapped.
77
        if let Some(s) = &self.structured {
78
            let obj = if s.is_object() {
79
                s.clone()
80
            } else {
81
                json!({ "result": s })
82
            };
83
            out["structuredContent"] = obj;
84
        }
85
        out
86
    }
87
}
88
89
/// The tools a server offers.
90
pub trait Tools {
91
    fn list(&self) -> Vec<Tool>;
92
    /// Run a tool. `None` when no tool has that name.
93
    fn call(&self, name: &str, args: &Value) -> Option<ToolOutput>;
94
}
95
96
/// Who the server is, for `initialize`.
97
#[derive(Clone, Debug)]
98
pub struct ServerInfo {
99
    pub name: String,
100
    pub version: String,
101
    /// Shown to the model: what the server is for and how to use it safely.
102
    pub instructions: String,
103
}
104
105
const PARSE_ERROR: i64 = -32700;
106
const INVALID_REQUEST: i64 = -32600;
107
const METHOD_NOT_FOUND: i64 = -32601;
108
const INVALID_PARAMS: i64 = -32602;
109
110
fn error(id: Value, code: i64, message: &str) -> Value {
111
    json!({ "jsonrpc": "2.0", "id": id, "error": { "code": code, "message": message } })
112
}
113
114
fn result(id: Value, result: Value) -> Value {
115
    json!({ "jsonrpc": "2.0", "id": id, "result": result })
116
}
117
118
/// Answer one message; `None` for a notification (no reply is due).
119
pub fn handle(info: &ServerInfo, tools: &dyn Tools, message: &Value) -> Option<Value> {
120
    let Some(obj) = message.as_object() else {
121
        return Some(error(
122
            Value::Null,
123
            INVALID_REQUEST,
124
            "a message is a JSON object",
125
        ));
126
    };
127
    let method = obj.get("method").and_then(Value::as_str);
128
    let id = obj.get("id").cloned();
129
    let Some(method) = method else {
130
        // A response to something we never sent, or nonsense: a request
131
        // without a method gets an error, anything else is ignored.
132
        return id.map(|id| error(id, INVALID_REQUEST, "a request needs a method"));
133
    };
134
    // Notifications (initialized, cancelled, …) carry no id and get no reply.
135
    let id = id?;
136
    let params = obj.get("params").cloned().unwrap_or(Value::Null);
137
    Some(match method {
138
        "initialize" => {
139
            let asked = params.get("protocolVersion").and_then(Value::as_str);
140
            let version = asked
141
                .filter(|v| PROTOCOL_VERSIONS.contains(v))
142
                .unwrap_or(PROTOCOL_VERSIONS[0]);
143
            result(
144
                id,
145
                json!({
146
                    "protocolVersion": version,
147
                    "capabilities": { "tools": { "listChanged": false } },
148
                    "serverInfo": { "name": info.name, "version": info.version },
149
                    "instructions": info.instructions,
150
                }),
151
            )
152
        }
153
        "ping" => result(id, json!({})),
154
        "tools/list" => result(
155
            id,
156
            json!({ "tools": tools.list().iter().map(Tool::describe).collect::<Vec<_>>() }),
157
        ),
158
        "tools/call" => {
159
            let Some(name) = params.get("name").and_then(Value::as_str) else {
160
                return Some(error(
161
                    id,
162
                    INVALID_PARAMS,
163
                    "tools/call needs the tool's name",
164
                ));
165
            };
166
            let args = params
167
                .get("arguments")
168
                .cloned()
169
                .unwrap_or_else(|| json!({}));
170
            if !args.is_object() {
171
                return Some(error(
172
                    id,
173
                    INVALID_PARAMS,
174
                    "a tool's arguments are a JSON object",
175
                ));
176
            }
177
            match tools.call(name, &args) {
178
                Some(out) => result(id, out.to_json()),
179
                None => error(id, INVALID_PARAMS, &format!("no tool named {name}")),
180
            }
181
        }
182
        other => error(id, METHOD_NOT_FOUND, &format!("{other} is not supported")),
183
    })
184
}
185
186
/// Serve until the input ends: read a message per line, write a reply per
187
/// request. A line that is not JSON gets a parse error; blank lines are
188
/// skipped.
189
pub fn serve(
190
    info: &ServerInfo,
191
    tools: &dyn Tools,
192
    input: impl BufRead,
193
    mut output: impl Write,
194
) -> std::io::Result<()> {
195
    for line in input.lines() {
196
        let line = line?;
197
        if line.trim().is_empty() {
198
            continue;
199
        }
200
        let reply = match serde_json::from_str::<Value>(&line) {
201
            Ok(Value::Array(batch)) => {
202
                let replies: Vec<Value> = batch
203
                    .iter()
204
                    .filter_map(|m| handle(info, tools, m))
205
                    .collect();
206
                (!replies.is_empty()).then_some(Value::Array(replies))
207
            }
208
            Ok(message) => handle(info, tools, &message),
209
            Err(e) => Some(error(Value::Null, PARSE_ERROR, &format!("not JSON: {e}"))),
210
        };
211
        if let Some(reply) = reply {
212
            writeln!(output, "{reply}")?;
213
            output.flush()?;
214
        }
215
    }
216
    Ok(())
217
}
218
219
/// A string argument, if given and a string.
220
pub fn arg_str<'a>(args: &'a Value, key: &str) -> Option<&'a str> {
221
    args.get(key).and_then(Value::as_str)
222
}
223
224
/// A list of strings, if given; a single string counts as a list of one.
225
pub fn arg_strings(args: &Value, key: &str) -> Result<Vec<String>, String> {
226
    match args.get(key) {
227
        None | Some(Value::Null) => Ok(vec![]),
228
        Some(Value::String(s)) => Ok(vec![s.clone()]),
229
        Some(Value::Array(a)) => a
230
            .iter()
231
            .map(|v| {
232
                v.as_str()
233
                    .map(str::to_string)
234
                    .ok_or_else(|| format!("{key} must be a list of strings"))
235
            })
236
            .collect(),
237
        Some(_) => Err(format!("{key} must be a list of strings")),
238
    }
239
}
240
241
#[cfg(test)]
242
mod tests {
243
    use super::*;
244
245
    struct Echo;
246
247
    impl Tools for Echo {
248
        fn list(&self) -> Vec<Tool> {
249
            vec![Tool {
250
                name: "echo".into(),
251
                description: "says it back".into(),
252
                input_schema: json!({"type": "object"}),
253
                read_only: true,
254
                destructive: false,
255
            }]
256
        }
257
        fn call(&self, name: &str, args: &Value) -> Option<ToolOutput> {
258
            (name == "echo").then(|| ToolOutput {
259
                text: args.to_string(),
260
                structured: Some(json!([1, 2])),
261
                is_error: false,
262
            })
263
        }
264
    }
265
266
    fn info() -> ServerInfo {
267
        ServerInfo {
268
            name: "t".into(),
269
            version: "1".into(),
270
            instructions: "be careful".into(),
271
        }
272
    }
273
274
    fn talk(lines: &str) -> Vec<Value> {
275
        let mut out = Vec::new();
276
        serve(&info(), &Echo, lines.as_bytes(), &mut out).unwrap();
277
        String::from_utf8(out)
278
            .unwrap()
279
            .lines()
280
            .map(|l| serde_json::from_str(l).unwrap())
281
            .collect()
282
    }
283
284
    #[test]
285
    fn a_session_initializes_lists_and_calls() {
286
        let r = talk(concat!(
287
            r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"c","version":"0"}}}"#,
288
            "\n",
289
            r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#,
290
            "\n\n",
291
            r#"{"jsonrpc":"2.0","id":2,"method":"tools/list"}"#,
292
            "\n",
293
            r#"{"jsonrpc":"2.0","id":"three","method":"tools/call","params":{"name":"echo","arguments":{"a":1}}}"#,
294
            "\n",
295
        ));
296
        assert_eq!(r.len(), 3, "the notification gets no reply: {r:?}");
297
        assert_eq!(
298
            r[0]["result"]["protocolVersion"], "2025-03-26",
299
            "a version we speak is kept"
300
        );
301
        assert_eq!(r[0]["result"]["instructions"], "be careful");
302
        assert_eq!(r[1]["result"]["tools"][0]["name"], "echo");
303
        assert_eq!(
304
            r[1]["result"]["tools"][0]["annotations"]["readOnlyHint"],
305
            true
306
        );
307
        assert_eq!(r[2]["id"], "three");
308
        assert_eq!(r[2]["result"]["content"][0]["text"], r#"{"a":1}"#);
309
        assert_eq!(
310
            r[2]["result"]["structuredContent"],
311
            json!({"result": [1, 2]}),
312
            "an array is wrapped in an object"
313
        );
314
        assert_eq!(r[2]["result"]["isError"], false);
315
    }
316
317
    #[test]
318
    fn errors_are_protocol_errors_and_unknown_versions_get_the_newest() {
319
        let r = talk(concat!(
320
            "not json\n",
321
            r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"1999-01-01"}}"#,
322
            "\n",
323
            r#"{"jsonrpc":"2.0","id":2,"method":"resources/list"}"#,
324
            "\n",
325
            r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"nope"}}"#,
326
            "\n",
327
            r#"{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"echo","arguments":[1]}}"#,
328
            "\n",
329
            r#"[{"jsonrpc":"2.0","id":5,"method":"ping"},{"jsonrpc":"2.0","method":"notifications/cancelled"}]"#,
330
            "\n",
331
        ));
332
        assert_eq!(r[0]["error"]["code"], PARSE_ERROR);
333
        assert_eq!(r[1]["result"]["protocolVersion"], PROTOCOL_VERSIONS[0]);
334
        assert_eq!(r[2]["error"]["code"], METHOD_NOT_FOUND);
335
        assert_eq!(r[3]["error"]["code"], INVALID_PARAMS);
336
        assert_eq!(r[4]["error"]["code"], INVALID_PARAMS);
337
        assert_eq!(
338
            r[5],
339
            json!([{"jsonrpc": "2.0", "id": 5, "result": {}}]),
340
            "a batch is answered as a batch"
341
        );
342
    }
343
344
    #[test]
345
    fn string_arguments_and_lists() {
346
        let a = json!({"one": "x", "many": ["a", "b"], "bad": [1]});
347
        assert_eq!(arg_str(&a, "one"), Some("x"));
348
        assert_eq!(arg_strings(&a, "one").unwrap(), vec!["x"]);
349
        assert_eq!(arg_strings(&a, "many").unwrap(), vec!["a", "b"]);
350
        assert_eq!(arg_strings(&a, "missing").unwrap(), Vec::<String>::new());
351
        assert!(arg_strings(&a, "bad").is_err());
352
    }
353
}