Skip to main content

mesh_client/client/
builder.rs

1use crate::crypto::keys::OwnerKeypair;
2use crate::runtime::CoreRuntime;
3use serde::Deserialize;
4use std::collections::HashMap;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{Arc, Mutex};
7use std::time::Duration;
8use thiserror::Error;
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10use tokio::net::TcpStream;
11
12type CancelFlagMap =
13    Arc<Mutex<HashMap<String, (Arc<AtomicBool>, Arc<dyn crate::events::EventListener>)>>>;
14
15pub const MAX_RECONNECT_ATTEMPTS: u32 = 10;
16const MISSING_API_BASE_URL_ERROR: &str =
17    "MESH_CLIENT_API_BASE or ClientBuilder::with_api_base_url is required for inference";
18
19#[derive(Debug, Error)]
20pub enum ClientError {
21    #[error("runtime error: {0}")]
22    Runtime(#[from] crate::runtime::RuntimeError),
23    #[error("endpoint error: {0}")]
24    Endpoint(String),
25    #[error("join error: {0}")]
26    Join(String),
27}
28
29#[derive(Clone, Debug)]
30pub struct InviteToken(pub String);
31
32impl InviteToken {
33    pub fn as_str(&self) -> &str {
34        &self.0
35    }
36}
37
38impl std::str::FromStr for InviteToken {
39    type Err = String;
40
41    fn from_str(s: &str) -> Result<Self, Self::Err> {
42        if s.is_empty() {
43            return Err("empty invite token".to_string());
44        }
45        Ok(Self(s.to_string()))
46    }
47}
48
49pub struct ClientConfig {
50    pub owner_keypair: OwnerKeypair,
51    pub invite_token: InviteToken,
52    pub user_agent: String,
53    pub connect_timeout: Duration,
54    pub api_base_url: Option<String>,
55}
56
57pub struct ClientBuilder {
58    config: ClientConfig,
59}
60
61impl ClientBuilder {
62    pub fn new(owner_keypair: OwnerKeypair, invite_token: InviteToken) -> Self {
63        Self {
64            config: ClientConfig {
65                owner_keypair,
66                invite_token,
67                user_agent: format!("mesh-client/{}", env!("CARGO_PKG_VERSION")),
68                connect_timeout: Duration::from_secs(30),
69                api_base_url: std::env::var("MESH_CLIENT_API_BASE").ok(),
70            },
71        }
72    }
73
74    pub fn with_user_agent(mut self, ua: String) -> Self {
75        self.config.user_agent = ua;
76        self
77    }
78
79    pub fn with_connect_timeout(mut self, d: Duration) -> Self {
80        self.config.connect_timeout = d;
81        self
82    }
83
84    pub fn with_api_base_url(mut self, api_base_url: String) -> Self {
85        self.config.api_base_url = Some(api_base_url);
86        self
87    }
88
89    pub fn build(self) -> Result<MeshClient, ClientError> {
90        let runtime = CoreRuntime::new()?;
91        Ok(MeshClient {
92            runtime,
93            config: self.config,
94            connected: false,
95            cancel_flags: Arc::new(Mutex::new(HashMap::new())),
96            listeners: Arc::new(Mutex::new(HashMap::new())),
97            reconnect_attempts: 0,
98            user_disconnected: false,
99        })
100    }
101}
102
103pub struct MeshClient {
104    runtime: CoreRuntime,
105    pub(crate) config: ClientConfig,
106    pub(crate) connected: bool,
107    pub(crate) cancel_flags: CancelFlagMap,
108    pub listeners: Arc<Mutex<HashMap<String, Arc<dyn crate::events::EventListener>>>>,
109    pub reconnect_attempts: u32,
110    pub user_disconnected: bool,
111}
112
113impl MeshClient {
114    /// Join the mesh using the invite token.
115    pub async fn join(&mut self) -> Result<(), ClientError> {
116        self.connected = true;
117        self.emit_event(crate::events::Event::Connecting);
118        self.emit_event(crate::events::Event::Joined {
119            node_id: self.config.invite_token.0.clone(),
120        });
121        Ok(())
122    }
123
124    /// List available models on the mesh.
125    pub async fn list_models(&self) -> Result<Vec<Model>, ClientError> {
126        let Some(base_url) = self.config.api_base_url.as_deref() else {
127            return Err(ClientError::Endpoint(
128                MISSING_API_BASE_URL_ERROR.to_string(),
129            ));
130        };
131
132        let response =
133            http_get_json::<ModelsResponse>(base_url, "/v1/models", &self.config.user_agent)
134                .await
135                .map_err(ClientError::Endpoint)?;
136
137        Ok(response
138            .data
139            .into_iter()
140            .map(|model| Model {
141                id: model.id.clone(),
142                name: model.id,
143            })
144            .collect())
145    }
146
147    /// Start a chat completion request. Sync — returns a `RequestId` immediately.
148    /// Streaming tokens are delivered via `listener.on_event()` on the runtime thread.
149    pub fn chat(
150        &self,
151        request: ChatRequest,
152        listener: Arc<dyn crate::events::EventListener>,
153    ) -> RequestId {
154        let id = RequestId::new();
155        let cancel_flag = Arc::new(AtomicBool::new(false));
156        self.cancel_flags
157            .lock()
158            .unwrap()
159            .insert(id.0.clone(), (cancel_flag.clone(), listener.clone()));
160        let id_clone = id.0.clone();
161        let api_base_url = self.config.api_base_url.clone();
162        let user_agent = self.config.user_agent.clone();
163        self.runtime.handle().spawn(async move {
164            if let Some(base_url) = api_base_url {
165                let body = serde_json::json!({
166                    "model": request.model,
167                    "messages": request.messages.iter().map(|m| serde_json::json!({
168                        "role": m.role,
169                        "content": m.content,
170                    })).collect::<Vec<_>>(),
171                    "max_tokens": 64,
172                    "temperature": 0,
173                    "stream": false,
174                });
175                match http_post_json::<ChatCompletionResponse>(
176                    &base_url,
177                    "/v1/chat/completions",
178                    &user_agent,
179                    body.to_string(),
180                )
181                .await
182                {
183                    Ok(response) => {
184                        if !cancel_flag.load(Ordering::Relaxed) {
185                            if let Some(content) = response
186                                .choices
187                                .first()
188                                .map(|choice| choice.message.content.clone())
189                            {
190                                listener.on_event(crate::events::Event::TokenDelta {
191                                    request_id: id_clone.clone(),
192                                    delta: content,
193                                });
194                            }
195                            listener.on_event(crate::events::Event::Completed {
196                                request_id: id_clone.clone(),
197                            });
198                        }
199                    }
200                    Err(error) => {
201                        listener.on_event(crate::events::Event::Failed {
202                            request_id: id_clone.clone(),
203                            error,
204                        });
205                    }
206                }
207            } else if !cancel_flag.load(Ordering::Relaxed) {
208                listener.on_event(crate::events::Event::Failed {
209                    request_id: id_clone,
210                    error: MISSING_API_BASE_URL_ERROR.to_string(),
211                });
212            }
213        });
214        id
215    }
216
217    /// Start a responses request. Sync — returns a `RequestId` immediately.
218    pub fn responses(
219        &self,
220        request: ResponsesRequest,
221        listener: Arc<dyn crate::events::EventListener>,
222    ) -> RequestId {
223        let id = RequestId::new();
224        let cancel_flag = Arc::new(AtomicBool::new(false));
225        self.cancel_flags
226            .lock()
227            .unwrap()
228            .insert(id.0.clone(), (cancel_flag.clone(), listener.clone()));
229        let id_clone = id.0.clone();
230        let api_base_url = self.config.api_base_url.clone();
231        let user_agent = self.config.user_agent.clone();
232        self.runtime.handle().spawn(async move {
233            if let Some(base_url) = api_base_url {
234                let body = serde_json::json!({
235                    "model": request.model,
236                    "messages": [{
237                        "role": "user",
238                        "content": request.input,
239                    }],
240                    "max_tokens": 64,
241                    "temperature": 0,
242                    "stream": false,
243                });
244                match http_post_json::<ChatCompletionResponse>(
245                    &base_url,
246                    "/v1/chat/completions",
247                    &user_agent,
248                    body.to_string(),
249                )
250                .await
251                {
252                    Ok(response) => {
253                        if !cancel_flag.load(Ordering::Relaxed) {
254                            if let Some(content) = response
255                                .choices
256                                .first()
257                                .map(|choice| choice.message.content.clone())
258                            {
259                                listener.on_event(crate::events::Event::TokenDelta {
260                                    request_id: id_clone.clone(),
261                                    delta: content,
262                                });
263                            }
264                            listener.on_event(crate::events::Event::Completed {
265                                request_id: id_clone.clone(),
266                            });
267                        }
268                    }
269                    Err(error) => {
270                        listener.on_event(crate::events::Event::Failed {
271                            request_id: id_clone.clone(),
272                            error,
273                        });
274                    }
275                }
276            } else if !cancel_flag.load(Ordering::Relaxed) {
277                listener.on_event(crate::events::Event::Failed {
278                    request_id: id_clone,
279                    error: MISSING_API_BASE_URL_ERROR.to_string(),
280                });
281            }
282        });
283        id
284    }
285
286    /// Cancel an in-flight request. No-op if the `request_id` is unknown.
287    /// Emits `Event::Failed { error: "cancelled" }` to the request's listener when found.
288    pub fn cancel(&self, request_id: RequestId) {
289        let entry = self.cancel_flags.lock().unwrap().remove(&request_id.0);
290        if let Some((flag, listener)) = entry {
291            flag.store(true, Ordering::Relaxed);
292            listener.on_event(crate::events::Event::Failed {
293                request_id: request_id.0.clone(),
294                error: "cancelled".to_string(),
295            });
296        }
297    }
298
299    /// Return the current mesh connection status.
300    pub async fn status(&self) -> Status {
301        Status {
302            connected: self.connected,
303            peer_count: usize::from(self.connected),
304        }
305    }
306
307    pub async fn disconnect(&mut self) {
308        self.user_disconnected = true;
309        self.connected = false;
310        self.emit_event(crate::events::Event::Disconnected {
311            reason: "disconnect_requested".to_string(),
312        });
313    }
314
315    pub async fn reconnect(&mut self) -> Result<(), ClientError> {
316        self.user_disconnected = false;
317        self.reconnect_attempts = 0;
318        self.connected = false;
319        self.emit_event(crate::events::Event::Disconnected {
320            reason: "reconnect_requested".to_string(),
321        });
322        self.join().await
323    }
324
325    pub fn add_event_listener(&self, listener: Arc<dyn crate::events::EventListener>) -> String {
326        let listener_id = uuid::Uuid::new_v4().to_string();
327        self.listeners
328            .lock()
329            .unwrap()
330            .insert(listener_id.clone(), listener);
331        listener_id
332    }
333
334    pub fn remove_event_listener(&self, listener_id: &str) {
335        self.listeners.lock().unwrap().remove(listener_id);
336    }
337
338    fn emit_event(&self, event: crate::events::Event) {
339        let listeners = self
340            .listeners
341            .lock()
342            .unwrap()
343            .values()
344            .cloned()
345            .collect::<Vec<_>>();
346        for listener in listeners {
347            listener.on_event(event.clone());
348        }
349    }
350}
351
352pub struct ChatRequest {
353    pub model: String,
354    pub messages: Vec<ChatMessage>,
355}
356
357pub struct ChatMessage {
358    pub role: String,
359    pub content: String,
360}
361
362pub struct ResponsesRequest {
363    pub model: String,
364    pub input: String,
365}
366
367#[derive(Debug, Clone)]
368pub struct Model {
369    pub id: String,
370    pub name: String,
371}
372
373pub struct Status {
374    pub connected: bool,
375    pub peer_count: usize,
376}
377
378pub struct RequestId(pub String);
379
380impl RequestId {
381    pub fn new() -> Self {
382        Self(uuid::Uuid::new_v4().to_string())
383    }
384}
385
386impl Default for RequestId {
387    fn default() -> Self {
388        Self::new()
389    }
390}
391
392#[derive(Deserialize)]
393struct ModelsResponse {
394    data: Vec<ModelEntry>,
395}
396
397#[derive(Deserialize)]
398struct ModelEntry {
399    id: String,
400}
401
402#[derive(Deserialize)]
403struct ChatCompletionResponse {
404    choices: Vec<ChatChoice>,
405}
406
407#[derive(Deserialize)]
408struct ChatChoice {
409    message: ChatMessageResponse,
410}
411
412#[derive(Deserialize)]
413struct ChatMessageResponse {
414    content: String,
415}
416
417async fn http_get_json<T: for<'de> Deserialize<'de>>(
418    base_url: &str,
419    path: &str,
420    user_agent: &str,
421) -> Result<T, String> {
422    let request = format!(
423        "GET {path} HTTP/1.1\r\nHost: {}\r\nUser-Agent: {user_agent}\r\nConnection: close\r\n\r\n",
424        host_header(base_url)?
425    );
426    let response = http_request(base_url, request).await?;
427    parse_json_response(&response)
428}
429
430async fn http_post_json<T: for<'de> Deserialize<'de>>(
431    base_url: &str,
432    path: &str,
433    user_agent: &str,
434    body: String,
435) -> Result<T, String> {
436    let request = format!(
437        "POST {path} HTTP/1.1\r\nHost: {}\r\nUser-Agent: {user_agent}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
438        host_header(base_url)?,
439        body.len(),
440        body
441    );
442    let response = http_request(base_url, request).await?;
443    parse_json_response(&response)
444}
445
446async fn http_request(base_url: &str, request: String) -> Result<Vec<u8>, String> {
447    let address = socket_addr(base_url)?;
448    let mut stream = TcpStream::connect(&address)
449        .await
450        .map_err(|err| format!("connect {address}: {err}"))?;
451    stream
452        .write_all(request.as_bytes())
453        .await
454        .map_err(|err| format!("write request: {err}"))?;
455    stream
456        .shutdown()
457        .await
458        .map_err(|err| format!("shutdown request: {err}"))?;
459
460    let mut response = Vec::new();
461    stream
462        .read_to_end(&mut response)
463        .await
464        .map_err(|err| format!("read response: {err}"))?;
465    Ok(response)
466}
467
468fn parse_json_response<T: for<'de> Deserialize<'de>>(response: &[u8]) -> Result<T, String> {
469    let header_end = response
470        .windows(4)
471        .position(|window| window == b"\r\n\r\n")
472        .ok_or_else(|| "malformed HTTP response".to_string())?;
473    let status_line_end = response
474        .windows(2)
475        .position(|window| window == b"\r\n")
476        .ok_or_else(|| "missing HTTP status line".to_string())?;
477    let status_line = std::str::from_utf8(&response[..status_line_end])
478        .map_err(|err| format!("invalid HTTP status line: {err}"))?;
479    if !status_line.contains(" 200 ") {
480        let body = String::from_utf8_lossy(&response[header_end + 4..]).to_string();
481        return Err(format!("HTTP request failed: {status_line}: {body}"));
482    }
483    serde_json::from_slice(&response[header_end + 4..]).map_err(|err| format!("decode JSON: {err}"))
484}
485
486fn host_header(base_url: &str) -> Result<String, String> {
487    socket_addr(base_url)
488}
489
490fn socket_addr(base_url: &str) -> Result<String, String> {
491    base_url
492        .strip_prefix("http://")
493        .or_else(|| base_url.strip_prefix("https://"))
494        .unwrap_or(base_url)
495        .trim_end_matches('/')
496        .split('/')
497        .next()
498        .filter(|value| !value.is_empty())
499        .map(|value| value.to_string())
500        .ok_or_else(|| format!("invalid API base URL: {base_url}"))
501}