harn_vm/
mcp_client_request.rs

1//! Shared MCP server-to-client request plumbing.
2//!
3//! Harn-served MCP handlers use this bus when they need to issue a JSON-RPC
4//! request to the connected client and wait for that client's response.
5
6use std::cell::RefCell;
7use std::collections::HashMap;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::sync::{Arc, Mutex};
10
11use serde_json::{json, Value as JsonValue};
12use tokio::sync::{mpsc, oneshot};
13
14use crate::value::{VmError, VmValue};
15
16/// Outbound message sink — typically wraps the MCP transport's writer
17/// (stdout for stdio servers, an SSE channel for HTTP servers).
18pub type OutboundSender = mpsc::UnboundedSender<JsonValue>;
19
20/// Per-connection bus that owns in-flight server-to-client requests.
21///
22/// Cheap to clone (every clone shares the same pending map), so the transport
23/// layer can hand a copy to its reader task while the dispatch loop installs
24/// the same bus thread-locally for tool/resource/prompt handlers.
25#[derive(Clone)]
26pub struct ClientRequestBus {
27    outbound: OutboundSender,
28    pending: Arc<Mutex<HashMap<String, oneshot::Sender<JsonValue>>>>,
29    next_id: Arc<AtomicU64>,
30}
31
32impl ClientRequestBus {
33    pub fn new(outbound: OutboundSender) -> Self {
34        Self {
35            outbound,
36            pending: Arc::new(Mutex::new(HashMap::new())),
37            next_id: Arc::new(AtomicU64::new(1)),
38        }
39    }
40
41    /// Try to dispatch `msg` as a response to a pending request. Returns `true`
42    /// when the message was a JSON-RPC response whose id matches an in-flight
43    /// request. Otherwise the caller should treat `msg` as a new inbound
44    /// request.
45    pub fn route_response(&self, msg: &JsonValue) -> bool {
46        if msg.get("method").is_some() {
47            return false;
48        }
49        if msg.get("result").is_none() && msg.get("error").is_none() {
50            return false;
51        }
52        let Some(id) = msg.get("id") else {
53            return false;
54        };
55        let id_key = canonical_id(id);
56        let mut pending = self
57            .pending
58            .lock()
59            .expect("client-request pending poisoned");
60        if let Some(tx) = pending.remove(&id_key) {
61            let _ = tx.send(msg.clone());
62            true
63        } else {
64            false
65        }
66    }
67
68    pub async fn request(
69        &self,
70        id_prefix: &str,
71        method: &str,
72        params: JsonValue,
73        error_prefix: &str,
74    ) -> Result<JsonValue, VmError> {
75        let id_seq = self.next_id.fetch_add(1, Ordering::Relaxed);
76        let id = format!("harn-{id_prefix}-{id_seq}");
77        let (tx, rx) = oneshot::channel();
78        self.pending
79            .lock()
80            .expect("client-request pending poisoned")
81            .insert(id.clone(), tx);
82
83        let request = json!({
84            "jsonrpc": "2.0",
85            "id": id,
86            "method": method,
87            "params": params,
88        });
89
90        if self.outbound.send(request).is_err() {
91            self.pending
92                .lock()
93                .expect("client-request pending poisoned")
94                .remove(&id);
95            return Err(VmError::Runtime(format!(
96                "{error_prefix}: transport closed before request could be sent"
97            )));
98        }
99
100        let response = match rx.await {
101            Ok(value) => value,
102            Err(_) => {
103                return Err(VmError::Runtime(format!(
104                    "{error_prefix}: transport dropped before client responded"
105                )));
106            }
107        };
108
109        if let Some(error) = response.get("error") {
110            let message = error
111                .get("message")
112                .and_then(|value| value.as_str())
113                .unwrap_or("unknown error");
114            let code = error
115                .get("code")
116                .and_then(|value| value.as_i64())
117                .unwrap_or(-1);
118            return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
119                format!("{error_prefix}: client error ({code}): {message}"),
120            ))));
121        }
122
123        Ok(response.get("result").cloned().unwrap_or(JsonValue::Null))
124    }
125}
126
127fn canonical_id(value: &JsonValue) -> String {
128    if let Some(s) = value.as_str() {
129        return s.to_string();
130    }
131    if let Some(n) = value.as_i64() {
132        return n.to_string();
133    }
134    if let Some(n) = value.as_u64() {
135        return n.to_string();
136    }
137    value.to_string()
138}
139
140thread_local! {
141    static CURRENT_BUS: RefCell<Option<ClientRequestBus>> = const { RefCell::new(None) };
142}
143
144pub fn install_bus(bus: Option<ClientRequestBus>) -> Option<ClientRequestBus> {
145    CURRENT_BUS.with(|cell| std::mem::replace(&mut *cell.borrow_mut(), bus))
146}
147
148pub fn current_bus() -> Option<ClientRequestBus> {
149    CURRENT_BUS.with(|cell| cell.borrow().clone())
150}
151
152#[cfg(test)]
153mod tests {
154    use serde_json::json;
155
156    use super::*;
157
158    #[test]
159    fn canonical_id_handles_strings_numbers_and_other() {
160        assert_eq!(canonical_id(&json!("a")), "a");
161        assert_eq!(canonical_id(&json!(42)), "42");
162        assert_eq!(canonical_id(&json!(true)), "true");
163    }
164}
harn_vm/mcp_client_request.rs

harn_vm/
mcp_client_request.rs
