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 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 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 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 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 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 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}