tenant.openagents/omega

1use anyhow::Context as _;
2use collections::HashMap;
3use futures::{FutureExt, Stream, StreamExt as _, future::BoxFuture, lock::Mutex};
4use gpui::BackgroundExecutor;
5use std::{pin::Pin, sync::Arc};
6
7use crate::{
8    transport::Transport,
9    types::{Implementation, InitializeResponse, ProtocolVersion, ServerCapabilities},
10};
11
12pub fn create_fake_transport(
13    name: impl Into<String>,
14    executor: BackgroundExecutor,
15) -> FakeTransport {
16    let name = name.into();
17    FakeTransport::new(executor).on_request::<crate::types::requests::Initialize, _>(
18        move |_params| {
19            let name = name.clone();
20            async move { create_initialize_response(name.clone()) }
21        },
22    )
23}
24
25fn create_initialize_response(server_name: String) -> InitializeResponse {
26    InitializeResponse {
27        protocol_version: ProtocolVersion(crate::types::LATEST_PROTOCOL_VERSION.to_string()),
28        server_info: Implementation {
29            name: server_name,
30            title: None,
31            version: "1.0.0".to_string(),
32            description: None,
33        },
34        capabilities: ServerCapabilities::default(),
35        meta: None,
36    }
37}
38
39pub struct FakeTransport {
40    request_handlers: HashMap<
41        &'static str,
42        Arc<dyn Send + Sync + Fn(serde_json::Value) -> BoxFuture<'static, serde_json::Value>>,
43    >,
44    tx: futures::channel::mpsc::UnboundedSender<String>,
45    rx: Arc<Mutex<futures::channel::mpsc::UnboundedReceiver<String>>>,
46    executor: BackgroundExecutor,
47}
48
49impl FakeTransport {
50    pub fn new(executor: BackgroundExecutor) -> Self {
51        let (tx, rx) = futures::channel::mpsc::unbounded();
52        Self {
53            request_handlers: Default::default(),
54            tx,
55            rx: Arc::new(Mutex::new(rx)),
56            executor,
57        }
58    }
59
60    pub fn on_request<T, Fut>(
61        mut self,
62        handler: impl 'static + Send + Sync + Fn(T::Params) -> Fut,
63    ) -> Self
64    where
65        T: crate::types::Request,
66        Fut: 'static + Send + Future<Output = T::Response>,
67    {
68        self.request_handlers.insert(
69            T::METHOD,
70            Arc::new(move |value| {
71                let params = value
72                    .get("params")
73                    .cloned()
74                    .unwrap_or(serde_json::Value::Null);
75                let params: T::Params =
76                    serde_json::from_value(params).expect("Invalid parameters received");
77                let response = handler(params);
78                async move { serde_json::to_value(response.await).unwrap() }.boxed()
79            }),
80        );
81        self
82    }
83}
84
85#[async_trait::async_trait]
86impl Transport for FakeTransport {
87    async fn send(&self, message: String) -> anyhow::Result<()> {
88        if let Ok(msg) = serde_json::from_str::<serde_json::Value>(&message) {
89            let id = msg.get("id").and_then(|id| id.as_u64()).unwrap_or(0);
90
91            if let Some(method) = msg.get("method") {
92                let method = method.as_str().expect("Invalid method received");
93                if let Some(handler) = self.request_handlers.get(method) {
94                    let payload = handler(msg).await;
95                    let response = serde_json::json!({
96                        "jsonrpc": "2.0",
97                        "id": id,
98                        "result": payload
99                    });
100                    self.tx
101                        .unbounded_send(response.to_string())
102                        .context("sending a message")?;
103                } else {
104                    log::debug!("No handler registered for MCP request '{method}'");
105                }
106            }
107        }
108        Ok(())
109    }
110
111    fn receive(&self) -> Pin<Box<dyn Stream<Item = String> + Send>> {
112        let rx = self.rx.clone();
113        let executor = self.executor.clone();
114        Box::pin(futures::stream::unfold(rx, move |rx| {
115            let executor = executor.clone();
116            async move {
117                let mut rx_guard = rx.lock().await;
118                executor.simulate_random_delay().await;
119                if let Some(message) = rx_guard.next().await {
120                    drop(rx_guard);
121                    Some((message, rx))
122                } else {
123                    None
124                }
125            }
126        }))
127    }
128
129    fn receive_err(&self) -> Pin<Box<dyn Stream<Item = String> + Send>> {
130        Box::pin(futures::stream::empty())
131    }
132}
133