myko_server/mcp/
exec.rs

1//! Tool executor abstraction for MCP transports.
2//!
3//! Two execution paths share the same dispatch core:
4//!
5//! - [`Executor::Client`] — wraps a [`MykoClient`]; talks to a remote Myko server
6//!   over WebSocket. Used by the stdio MCP binary.
7//! - [`Executor::InProcess`] — talks directly to a [`CellServerCtx`]; used by the
8//!   HTTP/WS MCP endpoints hosted inside the server.
9
10use std::{
11    sync::{Arc, Mutex},
12    time::Duration,
13};
14
15use hyphae::{Gettable, Watchable};
16use myko::{
17    client::{ConnectionStatus, MykoClient},
18    command::{CommandContext, CommandHandlerRegistration},
19    query::QueryRegistration,
20    report::ReportRegistration,
21    request::RequestContext,
22    server::CellServerCtx,
23    view::ViewRegistration,
24    wire::{WrappedCommand, WrappedQuery, WrappedReport, WrappedView},
25};
26use serde_json::{Value, json};
27use tokio::sync::oneshot;
28use uuid::Uuid;
29
30const QUERY_TIMEOUT: Duration = Duration::from_secs(5);
31const REPORT_TIMEOUT: Duration = Duration::from_secs(5);
32const COMMAND_TIMEOUT: Duration = Duration::from_secs(10);
33const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
34
35/// How MCP dispatch reaches the underlying Myko queries / reports / commands.
36#[derive(Clone)]
37pub enum Executor {
38    /// Talk to a remote Myko server over WebSocket via a `MykoClient`.
39    Client(Arc<MykoClient>),
40    /// Talk to a server hosted in the same process via its `CellServerCtx`.
41    InProcess(Arc<CellServerCtx>),
42}
43
44impl Executor {
45    /// Execute a query and return its current items as JSON.
46    pub async fn execute_query(&self, query_id: &str, args: Value) -> Result<Value, String> {
47        match self {
48            Executor::Client(client) => client_execute_query(client.clone(), query_id, args).await,
49            Executor::InProcess(ctx) => in_process_execute_query(ctx.clone(), query_id, args),
50        }
51    }
52
53    /// Execute a report and return its current output as JSON.
54    pub async fn execute_report(&self, report_id: &str, args: Value) -> Result<Value, String> {
55        match self {
56            Executor::Client(client) => {
57                client_execute_report(client.clone(), report_id, args).await
58            }
59            Executor::InProcess(ctx) => {
60                in_process_execute_report(ctx.clone(), report_id, args).await
61            }
62        }
63    }
64
65    /// Execute a view (list-typed report) and return its current items as JSON.
66    pub async fn execute_view(&self, view_id: &str, args: Value) -> Result<Value, String> {
67        match self {
68            Executor::Client(client) => client_execute_view(client.clone(), view_id, args).await,
69            Executor::InProcess(ctx) => in_process_execute_view(ctx.clone(), view_id, args),
70        }
71    }
72
73    /// Execute a command and return its result as JSON.
74    pub async fn execute_command(&self, command_id: &str, args: Value) -> Result<Value, String> {
75        match self {
76            Executor::Client(client) => {
77                client_execute_command(client.clone(), command_id, args).await
78            }
79            Executor::InProcess(ctx) => in_process_execute_command(ctx.clone(), command_id, args),
80        }
81    }
82
83    /// Status string for the built-in `connection_status` tool. Includes
84    /// server name/version (and, in-process, the host id) so a caller can
85    /// confirm which instance they hit without a separate report call.
86    pub fn connection_status(&self, info: &super::dispatch::ServerInfo) -> Value {
87        match self {
88            Executor::Client(client) => {
89                let status = client.connection_status().get();
90                let text = match &status {
91                    ConnectionStatus::Connected(addr) => format!("Connected to {}", addr),
92                    ConnectionStatus::Connecting(addr) => format!("Connecting to {}", addr),
93                    ConnectionStatus::Reconnecting(addr) => format!("Reconnecting to {}", addr),
94                    ConnectionStatus::Idle => "Idle".to_string(),
95                    ConnectionStatus::Disconnected => "Disconnected".to_string(),
96                };
97                json!({ "status": text, "name": info.name, "version": info.version })
98            }
99            Executor::InProcess(ctx) => json!({
100                "status": "In-process (always connected)",
101                "name": info.name,
102                "version": info.version,
103                "hostId": ctx.host_id,
104            }),
105        }
106    }
107}
108
109// ─────────────────────────────────────────────────────────────────────────────
110// Client-mode execution (stdio MCP path)
111// ─────────────────────────────────────────────────────────────────────────────
112
113async fn client_execute_query(
114    client: Arc<MykoClient>,
115    query_id: &str,
116    arguments: Value,
117) -> Result<Value, String> {
118    for reg in inventory::iter::<QueryRegistration> {
119        if reg.query_id == query_id {
120            let tx = Uuid::new_v4().to_string();
121            let mut query_json = arguments_object(arguments);
122            if let Some(obj) = query_json.as_object_mut() {
123                obj.insert("tx".to_string(), json!(tx));
124                obj.insert(
125                    "createdAt".to_string(),
126                    json!(chrono::Utc::now().to_rfc3339()),
127                );
128            }
129
130            let wrapped = WrappedQuery {
131                query: query_json,
132                query_id: reg.query_id.into(),
133                query_item_type: reg.query_item_type.into(),
134                window: None,
135            };
136
137            let cell = client.watch_query_raw(wrapped);
138            let (result_tx, result_rx) = oneshot::channel::<Vec<Value>>();
139            let result_tx = Arc::new(Mutex::new(Some(result_tx)));
140            let seen_initial = Arc::new(Mutex::new(false));
141            let result_tx_sub = result_tx.clone();
142            let seen_initial_sub = seen_initial.clone();
143            let _guard = cell.subscribe(move |signal| {
144                if let hyphae::Signal::Value(items) = signal {
145                    let mut seen = seen_initial_sub.lock().unwrap();
146                    if !*seen {
147                        *seen = true;
148                        return;
149                    }
150                    if let Some(tx) = result_tx_sub.lock().unwrap().take() {
151                        let _ = tx.send((**items).clone());
152                    }
153                }
154            });
155
156            return match tokio::time::timeout(QUERY_TIMEOUT, result_rx).await {
157                Ok(Ok(items)) => Ok(json!({
158                    "query_id": query_id,
159                    "item_type": reg.query_item_type,
160                    "count": items.len(),
161                    "items": items,
162                })),
163                Ok(Err(_)) => Err("Query channel closed".to_string()),
164                Err(_) => Err("Timeout waiting for query response".to_string()),
165            };
166        }
167    }
168    Err(format!("Query not found: {}", query_id))
169}
170
171async fn client_execute_view(
172    client: Arc<MykoClient>,
173    view_id: &str,
174    arguments: Value,
175) -> Result<Value, String> {
176    for reg in inventory::iter::<ViewRegistration> {
177        if reg.view_id == view_id {
178            let tx = Uuid::new_v4().to_string();
179            let mut view_json = arguments_object(arguments);
180            if let Some(obj) = view_json.as_object_mut() {
181                obj.insert("tx".to_string(), json!(tx));
182                obj.insert(
183                    "createdAt".to_string(),
184                    json!(chrono::Utc::now().to_rfc3339()),
185                );
186            }
187
188            let wrapped = WrappedView {
189                view: view_json,
190                view_id: reg.view_id.into(),
191                view_item_type: reg.view_item_type.into(),
192                window: None,
193            };
194
195            let cell = client.watch_view_raw(wrapped);
196            let (result_tx, result_rx) = oneshot::channel::<Vec<Value>>();
197            let result_tx = Arc::new(Mutex::new(Some(result_tx)));
198            let seen_initial = Arc::new(Mutex::new(false));
199            let result_tx_sub = result_tx.clone();
200            let seen_initial_sub = seen_initial.clone();
201            let _guard = cell.subscribe(move |signal| {
202                if let hyphae::Signal::Value(items) = signal {
203                    let mut seen = seen_initial_sub.lock().unwrap();
204                    if !*seen {
205                        *seen = true;
206                        return;
207                    }
208                    if let Some(tx) = result_tx_sub.lock().unwrap().take() {
209                        let _ = tx.send((**items).clone());
210                    }
211                }
212            });
213
214            return match tokio::time::timeout(QUERY_TIMEOUT, result_rx).await {
215                Ok(Ok(items)) => Ok(json!({
216                    "view_id": view_id,
217                    "item_type": reg.view_item_type,
218                    "count": items.len(),
219                    "items": items,
220                })),
221                Ok(Err(_)) => Err("View channel closed".to_string()),
222                Err(_) => Err("Timeout waiting for view response".to_string()),
223            };
224        }
225    }
226    Err(format!("View not found: {}", view_id))
227}
228
229async fn client_execute_report(
230    client: Arc<MykoClient>,
231    report_id: &str,
232    arguments: Value,
233) -> Result<Value, String> {
234    for reg in inventory::iter::<ReportRegistration> {
235        if reg.report_id == report_id {
236            let tx = Uuid::new_v4().to_string();
237            let mut report_json = arguments_object(arguments);
238            if let Some(obj) = report_json.as_object_mut() {
239                obj.insert("tx".to_string(), json!(tx));
240            }
241
242            let wrapped = WrappedReport {
243                report: report_json,
244                report_id: reg.report_id.to_string(),
245            };
246
247            let cell = client.watch_report_raw(wrapped);
248            let (result_tx, result_rx) = oneshot::channel::<Value>();
249            let result_tx = Arc::new(Mutex::new(Some(result_tx)));
250            let _guard = cell.subscribe(move |signal| {
251                if let hyphae::Signal::Value(value_opt) = signal
252                    && let Some(value) = &**value_opt
253                    && let Some(tx) = result_tx.lock().unwrap().take()
254                {
255                    let _ = tx.send(value.clone());
256                }
257            });
258
259            return match tokio::time::timeout(REPORT_TIMEOUT, result_rx).await {
260                Ok(Ok(value)) => Ok(json!({
261                    "report_id": report_id,
262                    "output_type": reg.output_type,
263                    "result": value,
264                })),
265                Ok(Err(_)) => Err("Report channel closed".to_string()),
266                Err(_) => Err("Timeout waiting for report response".to_string()),
267            };
268        }
269    }
270    Err(format!("Report not found: {}", report_id))
271}
272
273async fn client_execute_command(
274    client: Arc<MykoClient>,
275    command_id: &str,
276    arguments: Value,
277) -> Result<Value, String> {
278    let status = client.connection_status().get();
279    if !matches!(status, ConnectionStatus::Connected(_)) {
280        let (tx_connected, rx_connected) = oneshot::channel::<bool>();
281        let tx_connected = Mutex::new(Some(tx_connected));
282        let guard = client.connection_status().subscribe(move |signal| {
283            if let hyphae::Signal::Value(status) = signal
284                && let ConnectionStatus::Connected(_) = &**status
285                && let Some(sender) = tx_connected.lock().unwrap().take()
286            {
287                let _ = sender.send(true);
288            }
289        });
290
291        let connected = tokio::time::timeout(CONNECT_TIMEOUT, rx_connected)
292            .await
293            .unwrap_or(Ok(false))
294            .unwrap_or(false);
295        drop(guard);
296
297        if !connected {
298            return Err("Not connected to Myko server".to_string());
299        }
300    }
301
302    let tx = Uuid::new_v4().to_string();
303    let mut command_json = arguments_object(arguments);
304    if let Some(obj) = command_json.as_object_mut() {
305        obj.insert("tx".to_string(), json!(tx));
306    }
307
308    let wrapped = WrappedCommand {
309        command: command_json,
310        command_id: command_id.to_string(),
311    };
312
313    let result_cell = client.send_command_raw_result(wrapped);
314    let (resp_tx, resp_rx) = oneshot::channel::<Result<Value, String>>();
315    let resp_tx = Arc::new(Mutex::new(Some(resp_tx)));
316    let _guard = result_cell.subscribe(move |signal| {
317        if let hyphae::Signal::Value(result_opt) = signal
318            && let Some(result) = &**result_opt
319            && let Some(sender) = resp_tx.lock().unwrap().take()
320        {
321            let _ = sender.send(result.clone());
322        }
323    });
324
325    match tokio::time::timeout(COMMAND_TIMEOUT, resp_rx).await {
326        Ok(Ok(Ok(response))) => Ok(json!({
327            "command_id": command_id,
328            "success": true,
329            "result": response,
330        })),
331        Ok(Ok(Err(e))) => Err(e),
332        _ => Err("Timeout waiting for response".to_string()),
333    }
334}
335
336// ─────────────────────────────────────────────────────────────────────────────
337// In-process execution (HTTP/WS MCP path)
338// ─────────────────────────────────────────────────────────────────────────────
339
340fn in_process_execute_query(
341    ctx: Arc<CellServerCtx>,
342    query_id: &str,
343    arguments: Value,
344) -> Result<Value, String> {
345    let registration = inventory::iter::<QueryRegistration>
346        .into_iter()
347        .find(|r| r.query_id == query_id)
348        .ok_or_else(|| format!("Query not found: {}", query_id))?;
349
350    let query_data = ctx
351        .handler_registry
352        .get_query(query_id)
353        .ok_or_else(|| format!("Query handler not registered: {}", query_id))?;
354
355    let mut query_json = arguments_object(arguments);
356    let tx: Arc<str> = Uuid::new_v4().to_string().into();
357    if let Some(obj) = query_json.as_object_mut() {
358        obj.insert("tx".to_string(), json!(tx.as_ref()));
359        obj.insert(
360            "createdAt".to_string(),
361            json!(chrono::Utc::now().to_rfc3339()),
362        );
363    }
364
365    let parsed = (query_data.parse)(query_json)
366        .map_err(|e| format!("Failed to parse query {}: {}", query_id, e))?;
367
368    let request_context = Arc::new(RequestContext::internal(tx, ctx.host_id, "mcp"));
369
370    let cellmap = (query_data.cell_factory)(
371        parsed,
372        ctx.registry.clone(),
373        request_context,
374        Some(ctx.clone()),
375    )
376    .map_err(|e| format!("Failed to build query cell: {}", e))?;
377
378    let items: Vec<Value> = cellmap
379        .snapshot()
380        .into_iter()
381        .map(|(_, item)| serde_json::to_value(&*item).unwrap_or(Value::Null))
382        .collect();
383
384    Ok(json!({
385        "query_id": query_id,
386        "item_type": registration.query_item_type,
387        "count": items.len(),
388        "items": items,
389    }))
390}
391
392fn in_process_execute_view(
393    ctx: Arc<CellServerCtx>,
394    view_id: &str,
395    arguments: Value,
396) -> Result<Value, String> {
397    let registration = inventory::iter::<ViewRegistration>
398        .into_iter()
399        .find(|r| r.view_id == view_id)
400        .ok_or_else(|| format!("View not found: {}", view_id))?;
401
402    let view_data = ctx
403        .handler_registry
404        .get_view(view_id)
405        .ok_or_else(|| format!("View handler not registered: {}", view_id))?;
406
407    let mut view_json = arguments_object(arguments);
408    let tx: Arc<str> = Uuid::new_v4().to_string().into();
409    if let Some(obj) = view_json.as_object_mut() {
410        obj.insert("tx".to_string(), json!(tx.as_ref()));
411        obj.insert(
412            "createdAt".to_string(),
413            json!(chrono::Utc::now().to_rfc3339()),
414        );
415    }
416
417    let parsed = (view_data.parse)(view_json)
418        .map_err(|e| format!("Failed to parse view {}: {}", view_id, e))?;
419
420    let request_context = Arc::new(RequestContext::internal(tx, ctx.host_id, "mcp"));
421
422    let cellmap =
423        (view_data.cell_factory)(parsed, ctx.registry.clone(), request_context, ctx.clone())
424            .map_err(|e| format!("Failed to build view cell: {}", e))?;
425
426    let items: Vec<Value> = cellmap
427        .snapshot()
428        .into_iter()
429        .map(|(_, item)| serde_json::to_value(&*item).unwrap_or(Value::Null))
430        .collect();
431
432    Ok(json!({
433        "view_id": view_id,
434        "item_type": registration.view_item_type,
435        "count": items.len(),
436        "items": items,
437    }))
438}
439
440async fn in_process_execute_report(
441    ctx: Arc<CellServerCtx>,
442    report_id: &str,
443    arguments: Value,
444) -> Result<Value, String> {
445    let registration = inventory::iter::<ReportRegistration>
446        .into_iter()
447        .find(|r| r.report_id == report_id)
448        .ok_or_else(|| format!("Report not found: {}", report_id))?;
449
450    let report_data = ctx
451        .handler_registry
452        .get_report(report_id)
453        .ok_or_else(|| format!("Report handler not registered: {}", report_id))?;
454
455    let mut report_json = arguments_object(arguments);
456    let tx: Arc<str> = Uuid::new_v4().to_string().into();
457    if let Some(obj) = report_json.as_object_mut() {
458        obj.insert("tx".to_string(), json!(tx.as_ref()));
459    }
460
461    let parsed = (report_data.parse)(report_json)
462        .map_err(|e| format!("Failed to parse report {}: {}", report_id, e))?;
463
464    let request_context = Arc::new(RequestContext::internal(tx, ctx.host_id, "mcp"));
465
466    let cell = (report_data.cell_factory)(parsed, request_context, ctx)
467        .map_err(|e| format!("Failed to build report cell: {}", e))?;
468
469    // Subscribe to drive reactive evaluation; capture the first emission.
470    let (tx_resp, rx_resp) = oneshot::channel::<Value>();
471    let tx_resp = Arc::new(Mutex::new(Some(tx_resp)));
472    let tx_resp_sub = tx_resp.clone();
473    let _guard = cell.subscribe(move |signal| {
474        if let hyphae::Signal::Value(output) = signal
475            && let Some(sender) = tx_resp_sub.lock().unwrap().take()
476        {
477            let _ = sender.send(output.to_value());
478        }
479    });
480
481    match tokio::time::timeout(REPORT_TIMEOUT, rx_resp).await {
482        Ok(Ok(value)) => Ok(json!({
483            "report_id": report_id,
484            "output_type": registration.output_type,
485            "result": value,
486        })),
487        Ok(Err(_)) => Err("Report cell dropped before emitting".to_string()),
488        Err(_) => Err("Timeout waiting for report value".to_string()),
489    }
490}
491
492fn in_process_execute_command(
493    ctx: Arc<CellServerCtx>,
494    command_id: &str,
495    arguments: Value,
496) -> Result<Value, String> {
497    let mut command_json = arguments_object(arguments);
498    let tx: Arc<str> = Uuid::new_v4().to_string().into();
499    if let Some(obj) = command_json.as_object_mut() {
500        obj.insert("tx".to_string(), json!(tx.as_ref()));
501    }
502
503    for registration in inventory::iter::<CommandHandlerRegistration> {
504        if registration.command_id == command_id {
505            let executor = (registration.factory)();
506            let req = Arc::new(RequestContext::internal(tx.clone(), ctx.host_id, "mcp"));
507            let cmd_id: Arc<str> = Arc::from(command_id);
508            let cmd_ctx = CommandContext::new(cmd_id, req, ctx.clone());
509
510            return match executor.execute_from_value(command_json, cmd_ctx) {
511                Ok(result) => Ok(json!({
512                    "command_id": command_id,
513                    "success": true,
514                    "result": result,
515                })),
516                Err(err) => Err(err.message),
517            };
518        }
519    }
520
521    Err(format!("Command handler not found: {}", command_id))
522}
523
524// ─────────────────────────────────────────────────────────────────────────────
525// Helpers
526// ─────────────────────────────────────────────────────────────────────────────
527
528fn arguments_object(arguments: Value) -> Value {
529    if arguments.is_object() {
530        arguments
531    } else {
532        json!({})
533    }
534}
myko_server/mcp/exec.rs

myko_server/mcp/
exec.rs
