1use std::collections::{HashMap, VecDeque};
4use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
5use std::sync::Arc;
6use std::time::Duration;
7
8use futures_util::{SinkExt, StreamExt};
9use serde_json::{json, Map, Value};
10use tokio::sync::{mpsc, oneshot, Mutex, Notify};
11use tokio::time::{sleep, timeout};
12use tokio_tungstenite::{connect_async, tungstenite::Message};
13
14use crate::errors::{
15 ConnectionError, DaemonError, DisconnectCause, Error, Result, StaleLoopError, TimeoutError,
16};
17use crate::protocol::{
18 decode_message, expand_wire_messages, new_connection_init, new_notification, new_ping,
19 new_pong, new_request, new_subscribe, new_unsubscribe, Envelope,
20};
21use crate::stream_terminal::{
22 extract_loop_id_from_inbound, inbound_needs_delivery_ack, stale_pending_frame_label,
23};
24
25const MAX_INBOUND: usize = 20_000;
26const DEFAULT_MAX_FRAME: usize = 10 * 1024 * 1024;
27
28#[derive(Debug, Clone)]
30pub struct ClientConfig {
31 pub url: String,
33 pub max_inbound: usize,
35 pub max_frame_size: usize,
37}
38
39impl Default for ClientConfig {
40 fn default() -> Self {
41 Self {
42 url: "ws://127.0.0.1:8765".into(),
43 max_inbound: MAX_INBOUND,
44 max_frame_size: DEFAULT_MAX_FRAME,
45 }
46 }
47}
48
49#[derive(Debug, Clone, Default)]
51pub struct SendInputOptions {
52 pub loop_id: Option<String>,
54 pub autonomous: bool,
56 pub max_iterations: Option<u32>,
58 pub preferred_subagent: Option<String>,
60 pub model: Option<String>,
62 pub model_params: Option<Value>,
64 pub router_profile: Option<String>,
66 pub attachments: Option<Value>,
68 pub intent_hint: Option<String>,
70 pub response_schema: Option<Value>,
72 pub response_schema_name: Option<String>,
74 pub response_schema_strict: Option<bool>,
76 pub clarification_mode: Option<String>,
78 pub clarification_answer: bool,
80 pub clarification_answers: Option<Value>,
82}
83
84type RpcWaiter = oneshot::Sender<std::result::Result<Value, DaemonError>>;
85
86struct Shared {
87 url: String,
88 max_inbound: usize,
89 write_tx: Mutex<Option<mpsc::UnboundedSender<Message>>>,
90 rpc_waiters: Mutex<HashMap<String, RpcWaiter>>,
91 inbound: Mutex<VecDeque<Value>>,
92 inbound_notify: Notify,
93 connected: AtomicBool,
94 reader_alive: AtomicBool,
95 disconnect_cause: Mutex<Option<DisconnectCause>>,
96 disconnect_notify: Notify,
97 inbound_dropped: AtomicU64,
98 delivery_seq: Mutex<HashMap<String, u64>>,
99 heartbeat_interval_ms: Mutex<Option<u64>>,
100}
101
102#[derive(Clone)]
104pub struct Client {
105 shared: Arc<Shared>,
106}
107
108impl Client {
109 pub fn new(url: impl Into<String>) -> Self {
111 Self::with_config(ClientConfig {
112 url: url.into(),
113 ..Default::default()
114 })
115 }
116
117 pub fn with_config(cfg: ClientConfig) -> Self {
119 Self {
120 shared: Arc::new(Shared {
121 url: cfg.url,
122 max_inbound: cfg.max_inbound,
123 write_tx: Mutex::new(None),
124 rpc_waiters: Mutex::new(HashMap::new()),
125 inbound: Mutex::new(VecDeque::new()),
126 inbound_notify: Notify::new(),
127 connected: AtomicBool::new(false),
128 reader_alive: AtomicBool::new(false),
129 disconnect_cause: Mutex::new(None),
130 disconnect_notify: Notify::new(),
131 inbound_dropped: AtomicU64::new(0),
132 delivery_seq: Mutex::new(HashMap::new()),
133 heartbeat_interval_ms: Mutex::new(None),
134 }),
135 }
136 }
137
138 pub fn url(&self) -> &str {
140 &self.shared.url
141 }
142
143 pub fn is_connected(&self) -> bool {
145 self.shared.connected.load(Ordering::SeqCst)
146 }
147
148 pub fn is_connection_alive(&self) -> bool {
150 self.shared.reader_alive.load(Ordering::SeqCst)
151 && self.shared.connected.load(Ordering::SeqCst)
152 }
153
154 pub fn inbound_dropped(&self) -> u64 {
156 self.shared.inbound_dropped.load(Ordering::SeqCst)
157 }
158
159 pub async fn disconnect_cause(&self) -> Option<DisconnectCause> {
161 *self.shared.disconnect_cause.lock().await
162 }
163
164 pub async fn wait_disconnected(&self) -> DisconnectCause {
166 loop {
167 if let Some(c) = *self.shared.disconnect_cause.lock().await {
168 return c;
169 }
170 self.shared.disconnect_notify.notified().await;
171 }
172 }
173
174 pub async fn connect(&self) -> Result<()> {
176 if self.is_connected() {
177 return Ok(());
178 }
179 let (ws, _) = connect_async(&self.shared.url).await.map_err(|e| {
180 ConnectionError::new(
181 &self.shared.url,
182 1,
183 Box::new(e) as Box<dyn std::error::Error + Send + Sync>,
184 )
185 })?;
186 let (mut write, mut read) = ws.split();
187 let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
188 {
189 let mut slot = self.shared.write_tx.lock().await;
190 *slot = Some(tx);
191 }
192 self.shared.connected.store(true, Ordering::SeqCst);
193 self.shared.reader_alive.store(true, Ordering::SeqCst);
194 *self.shared.disconnect_cause.lock().await = None;
195
196 let shared_r2 = self.shared.clone();
197 let client_for_hb = self.clone();
198 tokio::spawn(async move {
199 while let Some(msg) = rx.recv().await {
200 if write.send(msg).await.is_err() {
201 break;
202 }
203 }
204 let _ = write.close().await;
205 });
206
207 tokio::spawn(async move {
208 while let Some(item) = read.next().await {
209 match item {
210 Ok(Message::Text(text)) => {
211 if let Ok(raw) = decode_message(&text) {
212 for msg in expand_wire_messages(raw) {
213 shared_r2.route_inbound(msg).await;
214 }
215 }
216 }
217 Ok(Message::Ping(data)) => {
218 let _ = shared_r2.send_raw(Message::Pong(data)).await;
219 }
220 Ok(Message::Close(_)) => break,
221 Ok(_) => {}
222 Err(_) => break,
223 }
224 }
225 shared_r2.mark_disconnected(DisconnectCause::Unclean).await;
226 });
227
228 self.send_envelope(new_connection_init()).await?;
230 let ack = self.wait_connection_ack(Duration::from_secs(15)).await?;
231 if let Some(interval) = ack
232 .get("result")
233 .and_then(|r| r.get("heartbeat_interval_ms"))
234 .and_then(|v| v.as_u64())
235 {
236 *self.shared.heartbeat_interval_ms.lock().await = Some(interval);
237 if interval > 0 {
238 let hb_client = client_for_hb;
239 tokio::spawn(async move {
240 let period = Duration::from_millis(interval.max(5_000));
241 while hb_client.is_connection_alive() {
242 sleep(period).await;
243 if !hb_client.is_connection_alive() {
244 break;
245 }
246 let _ = hb_client.send_envelope(new_ping()).await;
247 }
248 });
249 }
250 }
251 Ok(())
252 }
253
254 async fn wait_connection_ack(&self, overall: Duration) -> Result<Value> {
255 let deadline = tokio::time::Instant::now() + overall;
256 loop {
257 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
258 if remaining.is_zero() {
259 return Err(TimeoutError::new("connection_ack", format!("{overall:?}")).into());
260 }
261 let msg = self
262 .read_event_with_timeout(remaining.min(Duration::from_secs(2)))
263 .await?;
264 let Some(msg) = msg else {
265 continue;
266 };
267 let msg_type = msg.get("type").and_then(|v| v.as_str()).unwrap_or("");
268 if msg_type == "status" {
269 continue;
270 }
271 if msg_type == "error" {
272 let err = parse_daemon_error(&msg);
273 if matches!(err.code, -32003..=-32001) {
275 sleep(Duration::from_millis(50)).await;
276 self.send_envelope(new_connection_init()).await?;
277 continue;
278 }
279 return Err(err.into());
280 }
281 if msg_type == "connection_ack" {
282 let state = msg
283 .get("result")
284 .and_then(|r| r.get("readiness_state"))
285 .and_then(|v| v.as_str())
286 .unwrap_or("");
287 if state == "starting" || state == "warming" {
288 sleep(Duration::from_millis(50)).await;
289 self.send_envelope(new_connection_init()).await?;
290 continue;
291 }
292 if state == "ready" || state.is_empty() {
293 return Ok(msg);
294 }
295 return Err(Error::protocol(format!(
296 "connection_ack readiness_state={state}"
297 )));
298 }
299 }
300 }
301
302 pub async fn close(&self) -> Result<()> {
304 if self.is_connected() {
305 let _ = self.notify("disconnect", Map::new()).await;
306 }
307 {
308 let mut tx = self.shared.write_tx.lock().await;
309 *tx = None;
310 }
311 self.shared.mark_disconnected(DisconnectCause::Clean).await;
312 Ok(())
313 }
314
315 pub async fn send_envelope(&self, env: Envelope) -> Result<()> {
317 let text = env.to_wire_json()?;
318 self.shared.send_raw(Message::Text(text.into())).await
319 }
320
321 pub async fn notify(&self, method: &str, params: Map<String, Value>) -> Result<()> {
323 self.send_envelope(new_notification(method, params)).await
324 }
325
326 pub async fn request(
328 &self,
329 method: &str,
330 params: Map<String, Value>,
331 req_timeout: Duration,
332 ) -> Result<Map<String, Value>> {
333 let env = new_request(method, params);
334 let id = env.id.clone().unwrap_or_default();
335 let (tx, rx) = oneshot::channel();
336 self.shared.rpc_waiters.lock().await.insert(id.clone(), tx);
337 if let Err(e) = self.send_envelope(env).await {
338 self.shared.rpc_waiters.lock().await.remove(&id);
339 return Err(e);
340 }
341 match timeout(req_timeout, rx).await {
342 Ok(Ok(Ok(Value::Object(m)))) => Ok(m),
343 Ok(Ok(Ok(other))) => {
344 let mut m = Map::new();
345 m.insert("result".into(), other);
346 Ok(m)
347 }
348 Ok(Ok(Err(de))) => Err(de.into()),
349 Ok(Err(_)) => Err(Error::protocol("rpc waiter dropped")),
350 Err(_) => {
351 self.shared.rpc_waiters.lock().await.remove(&id);
352 Err(TimeoutError::new(method, format!("{req_timeout:?}")).into())
353 }
354 }
355 }
356
357 pub async fn request_response(
359 &self,
360 payload: Map<String, Value>,
361 fallback_method: &str,
362 req_timeout: Duration,
363 ) -> Result<Map<String, Value>> {
364 let method = payload
365 .get("type")
366 .and_then(|v| v.as_str())
367 .unwrap_or(fallback_method)
368 .to_string();
369 let mut params = Map::new();
370 for (k, v) in payload {
371 if k == "type" || k == "request_id" {
372 continue;
373 }
374 params.insert(k, v);
375 }
376 self.request(&method, params, req_timeout).await
377 }
378
379 pub async fn subscribe(
381 &self,
382 method: &str,
383 params: Map<String, Value>,
384 req_timeout: Duration,
385 ) -> Result<String> {
386 let env = new_subscribe(method, params);
387 let id = env.id.clone().unwrap_or_default();
388 let (tx, rx) = oneshot::channel();
391 self.shared.rpc_waiters.lock().await.insert(id.clone(), tx);
392 self.send_envelope(env).await?;
393 match timeout(req_timeout.min(Duration::from_millis(500)), rx).await {
395 Ok(Ok(Err(de))) => return Err(de.into()),
396 Ok(Ok(Ok(_))) => {}
397 _ => {
398 self.shared.rpc_waiters.lock().await.remove(&id);
399 }
400 }
401 Ok(id)
402 }
403
404 pub async fn unsubscribe(&self, id: &str) -> Result<()> {
406 self.send_envelope(new_unsubscribe(id)).await
407 }
408
409 pub async fn read_event(&self) -> Result<Option<Value>> {
411 loop {
412 {
413 let mut q = self.shared.inbound.lock().await;
414 if let Some(v) = q.pop_front() {
415 return Ok(Some(v));
416 }
417 }
418 if !self.is_connection_alive() {
419 return Ok(None);
420 }
421 self.shared.inbound_notify.notified().await;
422 }
423 }
424
425 pub async fn read_event_with_timeout(&self, dur: Duration) -> Result<Option<Value>> {
427 match timeout(dur, self.read_event()).await {
428 Ok(r) => r,
429 Err(_) => Ok(None),
430 }
431 }
432
433 pub async fn clear_pending_events(&self) {
435 self.shared.inbound.lock().await.clear();
436 }
437
438 pub async fn peel_stale_pending_control_events(&self) -> Vec<String> {
440 let mut labels = Vec::new();
441 let mut kept = VecDeque::new();
442 let mut q = self.shared.inbound.lock().await;
443 while let Some(ev) = q.pop_front() {
444 if let Some(label) = stale_pending_frame_label(&ev) {
445 labels.push(label);
446 } else {
447 kept.push_back(ev);
448 }
449 }
450 *q = kept;
451 labels
452 }
453
454 pub async fn send_input(&self, text: &str, opts: SendInputOptions) -> Result<()> {
456 let mut params = Map::new();
457 params.insert("content".into(), json!(text));
458 if let Some(lid) = opts.loop_id {
459 params.insert("loop_id".into(), json!(lid));
460 }
461 if opts.autonomous {
462 params.insert("autonomous".into(), json!(true));
463 }
464 if let Some(n) = opts.max_iterations {
465 params.insert("max_iterations".into(), json!(n));
466 }
467 if let Some(v) = opts.preferred_subagent {
468 params.insert("preferred_subagent".into(), json!(v));
469 }
470 if let Some(v) = opts.model {
471 params.insert("model".into(), json!(v));
472 }
473 if let Some(v) = opts.model_params {
474 params.insert("model_params".into(), v);
475 }
476 if let Some(v) = opts.router_profile {
477 params.insert("router_profile".into(), json!(v));
478 }
479 if let Some(v) = opts.attachments {
480 params.insert("attachments".into(), v);
481 }
482 if let Some(v) = opts.intent_hint {
483 params.insert("intent_hint".into(), json!(v));
484 }
485 if let Some(v) = opts.response_schema {
486 params.insert("response_schema".into(), v);
487 }
488 if let Some(v) = opts.response_schema_name {
489 params.insert("response_schema_name".into(), json!(v));
490 }
491 if let Some(v) = opts.response_schema_strict {
492 params.insert("response_schema_strict".into(), json!(v));
493 }
494 if let Some(v) = opts.clarification_mode {
495 params.insert("clarification_mode".into(), json!(v));
496 }
497 if opts.clarification_answer {
498 params.insert("clarification_answer".into(), json!(true));
499 }
500 if let Some(v) = opts.clarification_answers {
501 params.insert("clarification_answers".into(), v);
502 }
503 self.notify("loop_input", params).await
504 }
505
506 pub async fn loop_new(&self, params: Map<String, Value>) -> Result<Map<String, Value>> {
508 self.request("loop_new", params, Duration::from_secs(30))
509 .await
510 }
511
512 pub async fn loop_list(&self, limit: u32) -> Result<Map<String, Value>> {
514 let mut params = Map::new();
515 params.insert("limit".into(), json!(limit));
516 self.request("loop_list", params, Duration::from_secs(15))
517 .await
518 }
519
520 pub async fn loop_get(&self, loop_id: &str) -> Result<Map<String, Value>> {
522 let mut params = Map::new();
523 params.insert("loop_id".into(), json!(loop_id));
524 self.request("loop_get", params, Duration::from_secs(15))
525 .await
526 }
527
528 pub async fn loop_reattach(&self, loop_id: &str) -> Result<Map<String, Value>> {
530 let mut params = Map::new();
531 params.insert("loop_id".into(), json!(loop_id));
532 self.request("loop_reattach", params, Duration::from_secs(15))
533 .await
534 }
535
536 pub async fn loop_subscribe(&self, loop_id: &str, stream_delivery: &str) -> Result<String> {
538 let mut params = Map::new();
539 params.insert("loop_id".into(), json!(loop_id));
540 params.insert("stream_delivery".into(), json!(stream_delivery));
541 params.insert("wire_tier".into(), json!("full"));
542 self.subscribe("loop_events", params, Duration::from_secs(30))
543 .await
544 }
545
546 pub async fn loop_cards_fetch(&self, loop_id: &str) -> Result<Map<String, Value>> {
548 let mut params = Map::new();
549 params.insert("loop_id".into(), json!(loop_id));
550 self.request("loop_cards_fetch", params, Duration::from_secs(30))
551 .await
552 }
553
554 pub async fn loop_history_fetch(&self, loop_id: &str) -> Result<Map<String, Value>> {
556 let mut params = Map::new();
557 params.insert("loop_id".into(), json!(loop_id));
558 self.request("loop_history_fetch", params, Duration::from_secs(30))
559 .await
560 }
561
562 pub async fn loop_messages(
564 &self,
565 loop_id: &str,
566 limit: u32,
567 offset: u32,
568 ) -> Result<Map<String, Value>> {
569 let mut params = Map::new();
570 params.insert("loop_id".into(), json!(loop_id));
571 params.insert("limit".into(), json!(limit));
572 params.insert("offset".into(), json!(offset));
573 self.request("loop_messages", params, Duration::from_secs(10))
574 .await
575 }
576
577 pub async fn loop_state_get(&self, loop_id: &str) -> Result<Map<String, Value>> {
579 let mut params = Map::new();
580 params.insert("loop_id".into(), json!(loop_id));
581 self.request("loop_state_get", params, Duration::from_secs(30))
582 .await
583 }
584
585 pub async fn list_skills(&self) -> Result<Map<String, Value>> {
587 self.request("skills_list", Map::new(), Duration::from_secs(15))
588 .await
589 }
590
591 pub async fn list_models(&self) -> Result<Map<String, Value>> {
593 self.request("models_list", Map::new(), Duration::from_secs(15))
594 .await
595 }
596
597 pub async fn mcp_status(&self) -> Result<Map<String, Value>> {
599 self.request("mcp_status", Map::new(), Duration::from_secs(15))
600 .await
601 }
602
603 pub async fn fetch_daemon_status(&self) -> Result<Map<String, Value>> {
605 self.request("daemon_status", Map::new(), Duration::from_secs(10))
606 .await
607 }
608
609 pub async fn invoke_skill(&self, skill: &str, args: &str) -> Result<Map<String, Value>> {
611 let mut params = Map::new();
612 params.insert("skill".into(), json!(skill));
613 if !args.is_empty() {
614 params.insert("args".into(), json!(args));
615 }
616 self.request("invoke_skill", params, Duration::from_secs(120))
617 .await
618 }
619
620 pub async fn reattach_and_probe(&self, loop_id: &str) -> Result<()> {
622 self.loop_reattach(loop_id).await?;
623 self.loop_subscribe(loop_id, "adaptive").await?;
624 match self.loop_get(loop_id).await {
625 Ok(_) => Ok(()),
626 Err(Error::Daemon(de)) if de.code == -32200 => {
627 Err(StaleLoopError::new(loop_id, Some(Box::new(de))).into())
628 }
629 Err(e) => Err(StaleLoopError::new(loop_id, Some(Box::new(e))).into()),
630 }
631 }
632}
633
634impl Shared {
635 async fn send_raw(&self, msg: Message) -> Result<()> {
636 let tx = self.write_tx.lock().await;
637 let Some(tx) = tx.as_ref() else {
638 return Err(Error::protocol("not connected"));
639 };
640 tx.send(msg)
641 .map_err(|_| Error::protocol("write channel closed"))?;
642 Ok(())
643 }
644
645 async fn mark_disconnected(&self, cause: DisconnectCause) {
646 self.connected.store(false, Ordering::SeqCst);
647 self.reader_alive.store(false, Ordering::SeqCst);
648 {
649 let mut slot = self.write_tx.lock().await;
650 *slot = None;
651 }
652 {
653 let mut c = self.disconnect_cause.lock().await;
654 if c.is_none() {
655 *c = Some(cause);
656 }
657 }
658 let mut waiters = self.rpc_waiters.lock().await;
660 for (_, tx) in waiters.drain() {
661 let _ = tx.send(Err(DaemonError::new(-1, "disconnected")));
662 }
663 self.inbound_notify.notify_waiters();
664 self.disconnect_notify.notify_waiters();
665 }
666
667 async fn route_inbound(&self, msg: Value) {
668 if msg.get("type").and_then(|v| v.as_str()) == Some("ping") {
670 let _ = self
671 .send_raw(Message::Text(
672 new_pong().to_wire_json().unwrap_or_default().into(),
673 ))
674 .await;
675 return;
676 }
677
678 if inbound_needs_delivery_ack(&msg) {
679 let loop_id = extract_loop_id_from_inbound(&msg);
680 if !loop_id.is_empty() {
681 let seq = {
682 let mut map = self.delivery_seq.lock().await;
683 let e = map.entry(loop_id.clone()).or_insert(0);
684 *e += 1;
685 *e
686 };
687 let mut params = Map::new();
688 params.insert("loop_id".into(), json!(loop_id));
689 params.insert("seq".into(), json!(seq));
690 let _ = self
691 .send_raw(Message::Text(
692 new_notification("delivery_ack", params)
693 .to_wire_json()
694 .unwrap_or_default()
695 .into(),
696 ))
697 .await;
698 }
699 }
700
701 let msg_type = msg.get("type").and_then(|v| v.as_str()).unwrap_or("");
702 let id = msg
703 .get("id")
704 .and_then(|v| v.as_str())
705 .unwrap_or("")
706 .to_string();
707
708 if matches!(msg_type, "response" | "error") && !id.is_empty() {
709 if let Some(waiter) = self.rpc_waiters.lock().await.remove(&id) {
710 let result = if msg_type == "error" {
711 Err(parse_daemon_error(&msg))
712 } else {
713 Ok(msg
714 .get("result")
715 .cloned()
716 .unwrap_or(Value::Object(Map::new())))
717 };
718 let _ = waiter.send(result);
719 return;
720 }
721 }
722
723 let critical = is_critical_inbound(&msg);
725 let mut q = self.inbound.lock().await;
726 if q.len() >= self.max_inbound {
727 if critical {
728 if let Some(pos) = q.iter().position(|m| !is_critical_inbound(m)) {
730 q.remove(pos);
731 self.inbound_dropped.fetch_add(1, Ordering::SeqCst);
732 } else {
733 self.inbound_dropped.fetch_add(1, Ordering::SeqCst);
734 return;
735 }
736 } else {
737 self.inbound_dropped.fetch_add(1, Ordering::SeqCst);
738 return;
739 }
740 }
741 q.push_back(msg);
742 self.inbound_notify.notify_one();
743 }
744}
745
746fn is_critical_inbound(msg: &Value) -> bool {
747 let t = msg.get("type").and_then(|v| v.as_str()).unwrap_or("");
748 matches!(
749 t,
750 "status" | "error" | "complete" | "connection_ack" | "response"
751 ) || inbound_needs_delivery_ack(msg)
752}
753
754fn parse_daemon_error(msg: &Value) -> DaemonError {
755 if let Some(err) = msg.get("error") {
756 let code = err.get("code").and_then(|c| c.as_i64()).unwrap_or(-1);
757 let message = err
758 .get("message")
759 .and_then(|m| m.as_str())
760 .unwrap_or("daemon error")
761 .to_string();
762 let data = err.get("data").cloned();
763 let mut de = DaemonError::new(code, message);
764 if let Some(d) = data {
765 de = de.with_data(d);
766 }
767 return de;
768 }
769 DaemonError::new(
770 -1,
771 msg.get("message")
772 .and_then(|m| m.as_str())
773 .unwrap_or("daemon error"),
774 )
775}
776
777pub use crate::protocol::unwrap_next as unwrap_next_frame;
779
780#[cfg(test)]
781mod tests {
782 use super::*;
783
784 #[test]
785 fn send_input_options_default() {
786 let o = SendInputOptions::default();
787 assert!(!o.autonomous);
788 assert!(o.loop_id.is_none());
789 }
790}