Skip to main content

unlit_ecs/mcp/
host.rs

1//! The world side of the MCP server.
2//!
3//! A [World] is not [Send], and the work that reads it must happen on the
4//! thread that owns it. The MCP transport, on the other hand, runs on a tokio
5//! thread. The two meet at a command channel: a tool builds an [McpCommand]
6//! holding a closure and the [WorldId] of the world it should run against, the
7//! thread that owns that world resolves the id and calls
8//! [McpCommand::dispatch], and the [Output] travels back on the command's reply
9//! channel.
10//!
11//! Nothing here assumes how many worlds there are or what a world contains. A
12//! crate that defines tools of its own submits a
13//! [ToolSource](crate::mcp::ToolSource) and keeps its tool bodies next to the
14//! types they read; this module carries only the contract every tool shares.
15
16use std::panic::AssertUnwindSafe;
17use std::path::Path;
18
19use schemars::JsonSchema;
20use serde::Deserialize;
21use serde_json::Value;
22
23use crate::{Entity, World};
24
25/// What a finished job produced.
26///
27/// Most tools answer with structured JSON that the transport serialises
28/// through the value's reflection. Image tools answer with an encoded image
29/// that the transport sends as an MCP image content block, so a client
30/// renders the picture instead of decoding JSON; audio tools answer the same
31/// way, with a sound in an audio content block, so a client plays or decodes
32/// the mix instead of unbase64-ing a JSON field.
33pub enum Output {
34    /// A structured JSON result.
35    Json(Value),
36    /// An encoded image.
37    Image(EncodedImage),
38    /// An encoded sound.
39    Audio(EncodedAudio),
40}
41
42impl From<Value> for Output {
43    fn from(value: Value) -> Self {
44        Self::Json(value)
45    }
46}
47
48/// An image the transport sends as an image content block.
49pub struct EncodedImage {
50    /// The encoded bytes, base64-encoded.
51    pub data: String,
52    /// The format the bytes are in, as a MIME type.
53    pub mime_type: String,
54    /// The image's metadata — its size — as structured JSON.
55    pub structured: Value,
56}
57
58/// A sound the transport sends as an audio content block.
59pub struct EncodedAudio {
60    /// The encoded bytes, base64-encoded.
61    pub data: String,
62    /// The format the bytes are in, as a MIME type.
63    pub mime_type: String,
64    /// The sound's metadata — its length, rate and channel count — as
65    /// structured JSON.
66    pub structured: Value,
67}
68
69/// A unit of work for the thread that owns a world.
70///
71/// The closure captures plain data (tool arguments) and runs with exclusive
72/// access to one world. It must be [Send] because it travels across the
73/// channel; the world it touches never does.
74pub type Job = Box<dyn FnOnce(&mut World) -> Result<Output, String> + Send + 'static>;
75
76/// Which world a tool or a command runs against.
77///
78/// A host carries several worlds: the one its own thread holds directly, and
79/// any number it reaches through a host entity carrying a [World] component.
80/// A [World] is not [Send], so a command
81/// never names one directly — it names the entity that carries it, and the
82/// thread that owns that entity resolves the id to the world it holds.
83#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, JsonSchema, facet::Facet)]
84#[repr(u8)]
85pub enum WorldId {
86    /// The world the thread receiving the command owns directly.
87    #[default]
88    Host,
89    /// The world carried by the host entity with these bits.
90    Handle(u64),
91}
92
93impl WorldId {
94    /// The id of the world carried by `entity`.
95    #[must_use]
96    pub const fn handle(entity: Entity) -> Self {
97        Self::Handle(entity.to_bits())
98    }
99
100    /// The host entity carrying the world, or `None` for the host world.
101    #[must_use]
102    pub const fn entity(self) -> Option<Entity> {
103        match self {
104            Self::Host => None,
105            Self::Handle(bits) => Some(Entity::from_bits(bits)),
106        }
107    }
108}
109
110/// One world a host reports, with the name it was spawned under.
111///
112/// `list_worlds` answers with these: the id is what every other tool names
113/// the world by, and the name is what tells a person which world it is. The
114/// host supplies both, because only the host knows the worlds it carries.
115#[derive(Debug, Clone, Copy, PartialEq, Eq, facet::Facet)]
116pub struct WorldEntry {
117    /// The id every other tool names this world by.
118    pub world: WorldId,
119    /// The name the world was spawned under.
120    pub name: &'static str,
121}
122
123/// What a command asks the host to do.
124enum Request {
125    /// Run a job against one world.
126    Job(Job),
127    /// Answer with the worlds the host carries. Only the host knows them: the
128    /// worlds on other threads are not components this crate can recognise, so
129    /// the server asks rather than keeping a list of its own.
130    Worlds,
131}
132
133/// A request together with the channel its result is sent back on.
134pub struct McpCommand {
135    request: Request,
136    reply: async_channel::Sender<Result<Output, String>>,
137    /// The world a [`Request::Job`] runs against.
138    world: WorldId,
139}
140
141impl McpCommand {
142    /// Build a command that runs a job, and the receiver its result will
143    /// arrive on.
144    ///
145    /// The command starts out aimed at the world the receiving thread owns
146    /// directly; redirect it with [`with_world`](Self::with_world).
147    pub fn new(
148        job: impl FnOnce(&mut World) -> Result<Output, String> + Send + 'static,
149    ) -> (Self, async_channel::Receiver<Result<Output, String>>) {
150        Self::build(Request::Job(Box::new(job)))
151    }
152
153    /// Build the command the `list_worlds` tool sends.
154    ///
155    /// The host answers it with [`answer_worlds`](Self::answer_worlds); the
156    /// server holds no world list of its own.
157    #[must_use]
158    pub fn list_worlds() -> (Self, async_channel::Receiver<Result<Output, String>>) {
159        Self::build(Request::Worlds)
160    }
161
162    fn build(request: Request) -> (Self, async_channel::Receiver<Result<Output, String>>) {
163        let (reply, receiver) = async_channel::bounded(1);
164        (
165            Self {
166                request,
167                reply,
168                world: WorldId::Host,
169            },
170            receiver,
171        )
172    }
173
174    /// Direct the command at `world`.
175    #[must_use]
176    pub fn with_world(mut self, world: WorldId) -> Self {
177        self.world = world;
178        self
179    }
180
181    /// The world a job runs against.
182    #[must_use]
183    pub const fn world(&self) -> WorldId {
184        self.world
185    }
186
187    /// Whether this is the `list_worlds` request, which the host answers with
188    /// [`answer_worlds`](Self::answer_worlds) instead of dispatching a job.
189    #[must_use]
190    pub const fn wants_worlds(&self) -> bool {
191        matches!(self.request, Request::Worlds)
192    }
193
194    /// Answer the `list_worlds` request with the worlds the host carries.
195    ///
196    /// The host names every world a tool can reach — the host world itself and
197    /// each world it routes commands to, including ones on other threads —
198    /// each with the name it was spawned under.
199    pub fn answer_worlds(self, worlds: Vec<WorldEntry>) {
200        let result = to_value(&worlds).map(Output::from);
201        let _ = self.reply.send_blocking(result);
202    }
203
204    /// Answer the command with an error without running it.
205    ///
206    /// The host calls this when the world cannot be resolved, so the tool call
207    /// still gets a readable message rather than a dropped channel.
208    pub fn fail(self, message: String) {
209        let _ = self.reply.send_blocking(Err(message));
210    }
211
212    /// Run the command's job against `world` and send its result back.
213    ///
214    /// A [`list_worlds`](Self::list_worlds) request has no job; the host is
215    /// expected to answer it with [`answer_worlds`](Self::answer_worlds)
216    /// after checking [`wants_worlds`](Self::wants_worlds).
217    ///
218    /// A job that panics is caught and answered as an error: a world usually
219    /// shares its thread with others — the host's own, or sibling sub-worlds —
220    /// and an unwind out of a job would kill them all with it. The world the
221    /// job died in may still be inconsistent; the error says so.
222    ///
223    /// The send is best-effort: a tool call that has already been cancelled has
224    /// no receiver left, and the result is simply dropped.
225    pub fn dispatch(self, world: &mut World) {
226        let result = match self.request {
227            Request::Job(job) => catch_job(world, job),
228            Request::Worlds => Err("the host did not answer list_worlds".to_owned()),
229        };
230        let _ = self.reply.send_blocking(result);
231    }
232}
233
234/// Run one tool job, catching a panic into an error reply.
235///
236/// The world outlives the unwind only as a memory: what a job that panicked
237/// mid-way left behind is not promised to be consistent, so the error says the
238/// world may need a restart rather than pretending the tool merely failed.
239fn catch_job(world: &mut World, job: Job) -> Result<Output, String> {
240    match std::panic::catch_unwind(AssertUnwindSafe(move || job(world))) {
241        Ok(result) => result,
242        Err(payload) => Err(format!(
243            "the tool panicked and the world it ran in may be inconsistent: {}",
244            panic_message(payload)
245        )),
246    }
247}
248
249/// The message a panic payload carries, with a fallback for foreign payloads.
250fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
251    payload
252        .downcast_ref::<&'static str>()
253        .map(|message| (*message).to_owned())
254        .or_else(|| payload.downcast_ref::<String>().cloned())
255        .unwrap_or_else(|| "an opaque panic payload".to_owned())
256}
257
258/// Serialise a value through the reflection it derives.
259///
260/// The information structs the original crates expose carry their own
261/// reflection, so the JSON this transport speaks is derived from their fields
262/// rather than written down a second time here.
263pub fn to_value<T: facet::Facet<'static>>(value: &T) -> Result<Value, String> {
264    let text = facet_json::to_string(value).map_err(|error| error.to_string())?;
265    serde_json::from_str(&text).map_err(|error| error.to_string())
266}
267
268/// A spawned entity, as the tools report it.
269#[derive(facet::Facet)]
270pub struct SpawnedEntity {
271    /// The spawned entity.
272    #[facet(opaque, proxy = crate::EntityProxy)]
273    pub entity: Entity,
274    /// How the entity is called, so a reply read by a person names it.
275    pub name: &'static str,
276    /// The scope it was spawned in, so a caller can use it as one.
277    #[facet(opaque, proxy = crate::EntityProxy)]
278    pub scope: Entity,
279}
280
281/// What a despawn removed.
282#[derive(facet::Facet)]
283pub struct Despawned {
284    /// How many entities the despawn took with it.
285    pub removed: usize,
286    /// How the last of them was called.
287    pub name: &'static str,
288}
289
290/// A file a save tool wrote.
291#[derive(facet::Facet)]
292pub struct SavedFile {
293    /// The absolute path of the file that was written.
294    pub path: String,
295    /// How many bytes were written.
296    pub size: usize,
297}
298
299pub use crate::naming::{extension_of, with_extension};
300
301/// The name a save writes when its caller named none: `stem` with a fresh
302/// random id, so a save never overwrites the file an earlier one wrote.
303fn generated_name(stem: &str, extension: &str) -> String {
304    format!("{stem}-{}.{extension}", uuid::Uuid::new_v4().simple())
305}
306
307/// The name a save writes: the caller's own, given `extension` when it
308/// carries none, or a fresh `stem` name when the caller named none.
309pub fn saved_name(file_name: Option<&str>, stem: &str, extension: &str) -> String {
310    with_extension(file_name, extension).unwrap_or_else(|| generated_name(stem, extension))
311}
312
313/// Write `bytes` to `directory/file_name`, creating the directory.
314pub fn write_file(directory: &str, file_name: &str, bytes: &[u8]) -> Result<SavedFile, String> {
315    let path = Path::new(directory).join(file_name);
316    if let Some(parent) = path.parent() {
317        std::fs::create_dir_all(parent).map_err(|error| error.to_string())?;
318    }
319    std::fs::write(&path, bytes).map_err(|error| error.to_string())?;
320    // Report an absolute path, so a caller does not have to know the process's
321    // working directory to find what was written.
322    let path = std::fs::canonicalize(&path).unwrap_or(path);
323    Ok(SavedFile {
324        path: path.display().to_string(),
325        size: bytes.len(),
326    })
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use crate::bundle;
333
334    /// The JSON a tool argument carries for a world.
335    #[derive(Debug, Deserialize)]
336    struct Arg {
337        #[serde(default)]
338        world: WorldId,
339    }
340
341    #[test]
342    fn a_world_id_parses_from_json() {
343        let host: Arg = serde_json::from_str("{}").unwrap();
344        assert_eq!(host.world, WorldId::Host);
345        let host: Arg = serde_json::from_str(r#"{"world": "Host"}"#).unwrap();
346        assert_eq!(host.world, WorldId::Host);
347        let handle: Arg = serde_json::from_str(r#"{"world": {"Handle": 7}}"#).unwrap();
348        assert_eq!(handle.world, WorldId::Handle(7));
349    }
350
351    #[test]
352    fn a_world_id_maps_to_and_from_an_entity() {
353        let entity = Entity::from_bits(0x0000_0002_0000_0001);
354        assert_eq!(
355            WorldId::handle(entity),
356            WorldId::Handle(0x0000_0002_0000_0001)
357        );
358        assert_eq!(WorldId::handle(entity).entity(), Some(entity));
359        assert_eq!(WorldId::Host.entity(), None);
360        assert_eq!(WorldId::default(), WorldId::Host);
361    }
362
363    #[test]
364    fn a_list_worlds_command_carries_no_job() {
365        let (command, receiver) = McpCommand::list_worlds();
366        assert!(command.wants_worlds());
367        assert_eq!(command.world(), WorldId::Host);
368        command.answer_worlds(vec![
369            WorldEntry {
370                world: WorldId::Host,
371                name: "host",
372            },
373            WorldEntry {
374                world: WorldId::Handle(3),
375                name: "render",
376            },
377        ]);
378        let answer = receiver.recv_blocking().unwrap().unwrap();
379        let Output::Json(value) = answer else {
380            panic!("list_worlds answers with JSON");
381        };
382        assert_eq!(
383            value,
384            serde_json::json!([
385                {"world": "Host", "name": "host"},
386                {"world": {"Handle": 3}, "name": "render"}
387            ])
388        );
389    }
390
391    #[test]
392    fn a_job_command_dispatches_against_a_world() {
393        let (command, receiver) =
394            McpCommand::new(|world| Ok(Output::Json(serde_json::json!(world.root_scope().len()))));
395        assert!(!command.wants_worlds());
396        let mut world = World::new();
397        world.root_scope_mut().spawn(bundle![1u32]);
398        command.dispatch(&mut world);
399        let answer = receiver.recv_blocking().unwrap().unwrap();
400        let Output::Json(value) = answer else {
401            panic!("a job answers with JSON");
402        };
403        assert_eq!(value, serde_json::json!(1));
404    }
405
406    #[test]
407    fn a_panicking_job_answers_an_error_and_the_thread_survives() {
408        let (command, receiver) = McpCommand::new(|_| {
409            panic!("a tool blew up");
410        });
411        let mut world = World::new();
412        command.dispatch(&mut world);
413        let answer = receiver.recv_blocking().unwrap();
414        let Err(message) = answer else {
415            panic!("a panicking job answers with an error");
416        };
417        assert_eq!(
418            message,
419            "the tool panicked and the world it ran in may be inconsistent: a tool blew up"
420        );
421
422        // The same thread answers the next command: the unwind stopped at the
423        // command's edge instead of killing the world's owner.
424        let (command, receiver) =
425            McpCommand::new(|world| Ok(Output::Json(serde_json::json!(world.root_scope().len()))));
426        command.dispatch(&mut world);
427        let answer = receiver.recv_blocking().unwrap().unwrap();
428        let Output::Json(value) = answer else {
429            panic!("a job answers with JSON");
430        };
431        assert_eq!(value, serde_json::json!(0));
432    }
433
434    #[test]
435    fn a_generated_name_carries_a_fresh_id() {
436        let name = generated_name("texture", "webp");
437        assert!(name.starts_with("texture-"), "{name}");
438        assert!(name.ends_with(".webp"), "{name}");
439        assert_ne!(name, generated_name("texture", "webp"));
440    }
441
442    #[test]
443    fn a_saved_name_prefers_the_callers_own() {
444        assert_eq!(saved_name(Some("a"), "buffer", "bin"), "a.bin");
445        assert_eq!(saved_name(Some("a.png"), "buffer", "bin"), "a.png");
446        let generated = saved_name(None, "screenshot", "webp");
447        assert!(generated.starts_with("screenshot-"), "{generated}");
448        assert!(generated.ends_with(".webp"), "{generated}");
449    }
450}