Skip to main content

soothe_client/
client.rs

1//! Protocol-1 WebSocket transport client with mux and delivery_ack.
2
3use 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::heartbeat::HeartbeatTracker;
18use crate::protocol::{
19    decode_message, expand_wire_messages, new_connection_init, new_notification, new_ping,
20    new_pong, new_request, new_subscribe, new_unsubscribe, Envelope,
21};
22use crate::stream_terminal::{
23    extract_loop_id_from_inbound, inbound_needs_delivery_ack, stale_pending_frame_label,
24};
25
26const MAX_INBOUND: usize = 20_000;
27const DEFAULT_MAX_FRAME: usize = 10 * 1024 * 1024;
28
29/// Options for Client construction.
30#[derive(Debug, Clone)]
31pub struct ClientConfig {
32    /// WebSocket URL.
33    pub url: String,
34    /// Max inbound queue depth before priority drop.
35    pub max_inbound: usize,
36    /// Max frame size hint (informational).
37    pub max_frame_size: usize,
38}
39
40impl Default for ClientConfig {
41    fn default() -> Self {
42        Self {
43            url: "ws://127.0.0.1:8765".into(),
44            max_inbound: MAX_INBOUND,
45            max_frame_size: DEFAULT_MAX_FRAME,
46        }
47    }
48}
49
50/// Options for `send_input` / loop_input notification.
51#[derive(Debug, Clone, Default)]
52pub struct SendInputOptions {
53    /// Loop id (required for multi-loop).
54    pub loop_id: Option<String>,
55    /// Autonomous mode.
56    pub autonomous: bool,
57    /// Max iterations when autonomous.
58    pub max_iterations: Option<u32>,
59    /// Preferred subagent.
60    pub preferred_subagent: Option<String>,
61    /// Model override.
62    pub model: Option<String>,
63    /// Model params.
64    pub model_params: Option<Value>,
65    /// Router profile.
66    pub router_profile: Option<String>,
67    /// Attachments array.
68    pub attachments: Option<Value>,
69    /// Intent hint.
70    pub intent_hint: Option<String>,
71    /// Response schema.
72    pub response_schema: Option<Value>,
73    /// Response schema name.
74    pub response_schema_name: Option<String>,
75    /// Strict schema.
76    pub response_schema_strict: Option<bool>,
77    /// Clarification mode.
78    pub clarification_mode: Option<String>,
79    /// Clarification answer flag.
80    pub clarification_answer: bool,
81    /// Clarification answers payload.
82    pub clarification_answers: Option<Value>,
83}
84
85type RpcWaiter = oneshot::Sender<std::result::Result<Value, DaemonError>>;
86type StreamDegradedCallback = Arc<dyn Fn(u64, String) + Send + Sync>;
87
88struct Shared {
89    url: String,
90    max_inbound: usize,
91    write_tx: Mutex<Option<mpsc::UnboundedSender<Message>>>,
92    rpc_waiters: Mutex<HashMap<String, RpcWaiter>>,
93    inbound: Mutex<VecDeque<Value>>,
94    inbound_notify: Notify,
95    connected: AtomicBool,
96    reader_alive: AtomicBool,
97    disconnect_cause: Mutex<Option<DisconnectCause>>,
98    disconnect_notify: Notify,
99    inbound_dropped: AtomicU64,
100    delivery_seq: Mutex<HashMap<String, u64>>,
101    heartbeat_interval_ms: Mutex<Option<u64>>,
102    stream_degraded_cb: std::sync::Mutex<Option<StreamDegradedCallback>>,
103    degraded_notified: AtomicBool,
104    heartbeat_tracker: std::sync::Mutex<Option<Arc<HeartbeatTracker>>>,
105}
106
107/// Long-lived protocol-1 WebSocket client.
108#[derive(Clone)]
109pub struct Client {
110    shared: Arc<Shared>,
111}
112
113impl Client {
114    /// Create a client for `url`.
115    pub fn new(url: impl Into<String>) -> Self {
116        Self::with_config(ClientConfig {
117            url: url.into(),
118            ..Default::default()
119        })
120    }
121
122    /// Create with explicit config.
123    pub fn with_config(cfg: ClientConfig) -> Self {
124        Self {
125            shared: Arc::new(Shared {
126                url: cfg.url,
127                max_inbound: cfg.max_inbound,
128                write_tx: Mutex::new(None),
129                rpc_waiters: Mutex::new(HashMap::new()),
130                inbound: Mutex::new(VecDeque::new()),
131                inbound_notify: Notify::new(),
132                connected: AtomicBool::new(false),
133                reader_alive: AtomicBool::new(false),
134                disconnect_cause: Mutex::new(None),
135                disconnect_notify: Notify::new(),
136                inbound_dropped: AtomicU64::new(0),
137                delivery_seq: Mutex::new(HashMap::new()),
138                heartbeat_interval_ms: Mutex::new(None),
139                stream_degraded_cb: std::sync::Mutex::new(None),
140                degraded_notified: AtomicBool::new(false),
141                heartbeat_tracker: std::sync::Mutex::new(None),
142            }),
143        }
144    }
145
146    /// Daemon URL.
147    pub fn url(&self) -> &str {
148        &self.shared.url
149    }
150
151    /// Whether the socket is connected.
152    pub fn is_connected(&self) -> bool {
153        self.shared.connected.load(Ordering::SeqCst)
154    }
155
156    /// Whether the reader task is alive.
157    pub fn is_connection_alive(&self) -> bool {
158        self.shared.reader_alive.load(Ordering::SeqCst)
159            && self.shared.connected.load(Ordering::SeqCst)
160    }
161
162    /// Count of dropped inbound frames under backpressure.
163    pub fn inbound_dropped(&self) -> u64 {
164        self.shared.inbound_dropped.load(Ordering::SeqCst)
165    }
166
167    /// Register a hook invoked on the first inbound overflow drop.
168    pub fn set_stream_degraded_callback(&self, cb: Option<Arc<dyn Fn(u64, String) + Send + Sync>>) {
169        *self
170            .shared
171            .stream_degraded_cb
172            .lock()
173            .expect("stream_degraded lock") = cb;
174        self.shared.degraded_notified.store(false, Ordering::SeqCst);
175    }
176
177    /// Enable heartbeat tracking with the default 15s alive threshold.
178    pub fn enable_heartbeat_tracking(&self) -> Arc<HeartbeatTracker> {
179        self.enable_heartbeat_tracking_with_threshold(Duration::from_secs(15))
180    }
181
182    /// Enable heartbeat tracking with a custom alive threshold.
183    pub fn enable_heartbeat_tracking_with_threshold(
184        &self,
185        threshold: Duration,
186    ) -> Arc<HeartbeatTracker> {
187        let tracker = Arc::new(HeartbeatTracker::with_threshold(threshold));
188        *self
189            .shared
190            .heartbeat_tracker
191            .lock()
192            .expect("heartbeat lock") = Some(tracker.clone());
193        tracker
194    }
195
196    /// Disable heartbeat tracking.
197    pub fn disable_heartbeat_tracking(&self) {
198        *self
199            .shared
200            .heartbeat_tracker
201            .lock()
202            .expect("heartbeat lock") = None;
203    }
204
205    /// Current heartbeat tracker, if enabled.
206    pub fn heartbeat_tracker(&self) -> Option<Arc<HeartbeatTracker>> {
207        self.shared
208            .heartbeat_tracker
209            .lock()
210            .expect("heartbeat lock")
211            .clone()
212    }
213
214    /// Whether the tracked daemon is considered alive (true if tracking disabled).
215    pub fn is_daemon_alive(&self) -> bool {
216        match self
217            .shared
218            .heartbeat_tracker
219            .lock()
220            .expect("heartbeat lock")
221            .as_ref()
222        {
223            Some(t) => t.get_health().is_alive,
224            None => true,
225        }
226    }
227
228    /// Disconnect cause if disconnected.
229    pub async fn disconnect_cause(&self) -> Option<DisconnectCause> {
230        *self.shared.disconnect_cause.lock().await
231    }
232
233    /// Wait until disconnected.
234    pub async fn wait_disconnected(&self) -> DisconnectCause {
235        loop {
236            if let Some(c) = *self.shared.disconnect_cause.lock().await {
237                return c;
238            }
239            self.shared.disconnect_notify.notified().await;
240        }
241    }
242
243    /// Dial + handshake (`connection_init` / `connection_ack` ready).
244    pub async fn connect(&self) -> Result<()> {
245        if self.is_connected() {
246            return Ok(());
247        }
248        let (ws, _) = connect_async(&self.shared.url).await.map_err(|e| {
249            ConnectionError::new(
250                &self.shared.url,
251                1,
252                Box::new(e) as Box<dyn std::error::Error + Send + Sync>,
253            )
254        })?;
255        let (mut write, mut read) = ws.split();
256        let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
257        {
258            let mut slot = self.shared.write_tx.lock().await;
259            *slot = Some(tx);
260        }
261        self.shared.connected.store(true, Ordering::SeqCst);
262        self.shared.reader_alive.store(true, Ordering::SeqCst);
263        *self.shared.disconnect_cause.lock().await = None;
264
265        let shared_r2 = self.shared.clone();
266        let client_for_hb = self.clone();
267        tokio::spawn(async move {
268            while let Some(msg) = rx.recv().await {
269                if write.send(msg).await.is_err() {
270                    break;
271                }
272            }
273            let _ = write.close().await;
274        });
275
276        tokio::spawn(async move {
277            while let Some(item) = read.next().await {
278                match item {
279                    Ok(Message::Text(text)) => {
280                        if let Ok(raw) = decode_message(&text) {
281                            for msg in expand_wire_messages(raw) {
282                                shared_r2.route_inbound(msg).await;
283                            }
284                        }
285                    }
286                    Ok(Message::Ping(data)) => {
287                        let _ = shared_r2.send_raw(Message::Pong(data)).await;
288                    }
289                    Ok(Message::Close(_)) => break,
290                    Ok(_) => {}
291                    Err(_) => break,
292                }
293            }
294            shared_r2.mark_disconnected(DisconnectCause::Unclean).await;
295        });
296
297        // Handshake
298        self.send_envelope(new_connection_init()).await?;
299        let ack = self.wait_connection_ack(Duration::from_secs(15)).await?;
300        if let Some(interval) = ack
301            .get("result")
302            .and_then(|r| r.get("heartbeat_interval_ms"))
303            .and_then(|v| v.as_u64())
304        {
305            *self.shared.heartbeat_interval_ms.lock().await = Some(interval);
306            if interval > 0 {
307                let hb_client = client_for_hb;
308                tokio::spawn(async move {
309                    let period = Duration::from_millis(interval.max(5_000));
310                    while hb_client.is_connection_alive() {
311                        sleep(period).await;
312                        if !hb_client.is_connection_alive() {
313                            break;
314                        }
315                        let _ = hb_client.send_envelope(new_ping()).await;
316                    }
317                });
318            }
319        }
320        Ok(())
321    }
322
323    async fn wait_connection_ack(&self, overall: Duration) -> Result<Value> {
324        let deadline = tokio::time::Instant::now() + overall;
325        loop {
326            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
327            if remaining.is_zero() {
328                return Err(TimeoutError::new("connection_ack", format!("{overall:?}")).into());
329            }
330            let msg = self
331                .read_event_with_timeout(remaining.min(Duration::from_secs(2)))
332                .await?;
333            let Some(msg) = msg else {
334                continue;
335            };
336            let msg_type = msg.get("type").and_then(|v| v.as_str()).unwrap_or("");
337            if msg_type == "status" {
338                continue;
339            }
340            if msg_type == "error" {
341                let err = parse_daemon_error(&msg);
342                // Retryable readiness codes.
343                if matches!(err.code, -32003..=-32001) {
344                    sleep(Duration::from_millis(50)).await;
345                    self.send_envelope(new_connection_init()).await?;
346                    continue;
347                }
348                return Err(err.into());
349            }
350            if msg_type == "connection_ack" {
351                let state = msg
352                    .get("result")
353                    .and_then(|r| r.get("readiness_state"))
354                    .and_then(|v| v.as_str())
355                    .unwrap_or("");
356                if state == "starting" || state == "warming" {
357                    sleep(Duration::from_millis(50)).await;
358                    self.send_envelope(new_connection_init()).await?;
359                    continue;
360                }
361                if state == "ready" || state.is_empty() {
362                    return Ok(msg);
363                }
364                return Err(Error::protocol(format!(
365                    "connection_ack readiness_state={state}"
366                )));
367            }
368        }
369    }
370
371    /// Close the connection.
372    pub async fn close(&self) -> Result<()> {
373        if self.is_connected() {
374            let _ = self.notify("disconnect", Map::new()).await;
375        }
376        {
377            let mut tx = self.shared.write_tx.lock().await;
378            *tx = None;
379        }
380        self.shared.mark_disconnected(DisconnectCause::Clean).await;
381        Ok(())
382    }
383
384    /// Re-dial and re-handshake after a connection drop.
385    ///
386    /// Does not re-establish loop subscriptions; follow with
387    /// [`Self::reattach_and_probe`] to resume a loop session.
388    pub async fn reconnect(&self) -> Result<()> {
389        {
390            let mut tx = self.shared.write_tx.lock().await;
391            *tx = None;
392        }
393        self.shared.connected.store(false, Ordering::SeqCst);
394        self.shared.reader_alive.store(false, Ordering::SeqCst);
395        *self.shared.disconnect_cause.lock().await = None;
396        self.shared.rpc_waiters.lock().await.clear();
397        self.shared.inbound.lock().await.clear();
398        self.shared.degraded_notified.store(false, Ordering::SeqCst);
399        self.connect().await
400    }
401
402    /// Send a raw envelope.
403    pub async fn send_envelope(&self, env: Envelope) -> Result<()> {
404        let text = env.to_wire_json()?;
405        self.shared.send_raw(Message::Text(text.into())).await
406    }
407
408    /// Fire-and-forget notification.
409    pub async fn notify(&self, method: &str, params: Map<String, Value>) -> Result<()> {
410        self.send_envelope(new_notification(method, params)).await
411    }
412
413    /// RPC request correlated by id.
414    pub async fn request(
415        &self,
416        method: &str,
417        params: Map<String, Value>,
418        req_timeout: Duration,
419    ) -> Result<Map<String, Value>> {
420        let env = new_request(method, params);
421        let id = env.id.clone().unwrap_or_default();
422        let (tx, rx) = oneshot::channel();
423        self.shared.rpc_waiters.lock().await.insert(id.clone(), tx);
424        if let Err(e) = self.send_envelope(env).await {
425            self.shared.rpc_waiters.lock().await.remove(&id);
426            return Err(e);
427        }
428        match timeout(req_timeout, rx).await {
429            Ok(Ok(Ok(Value::Object(m)))) => Ok(m),
430            Ok(Ok(Ok(other))) => {
431                let mut m = Map::new();
432                m.insert("result".into(), other);
433                Ok(m)
434            }
435            Ok(Ok(Err(de))) => Err(de.into()),
436            Ok(Err(_)) => Err(Error::protocol("rpc waiter dropped")),
437            Err(_) => {
438                self.shared.rpc_waiters.lock().await.remove(&id);
439                Err(TimeoutError::new(method, format!("{req_timeout:?}")).into())
440            }
441        }
442    }
443
444    /// Legacy-shaped request map (`type` → method) like Go RequestResponse.
445    pub async fn request_response(
446        &self,
447        payload: Map<String, Value>,
448        fallback_method: &str,
449        req_timeout: Duration,
450    ) -> Result<Map<String, Value>> {
451        let method = payload
452            .get("type")
453            .and_then(|v| v.as_str())
454            .unwrap_or(fallback_method)
455            .to_string();
456        let mut params = Map::new();
457        for (k, v) in payload {
458            if k == "type" || k == "request_id" {
459                continue;
460            }
461            params.insert(k, v);
462        }
463        self.request(&method, params, req_timeout).await
464    }
465
466    /// Subscribe; returns subscription id.
467    pub async fn subscribe(
468        &self,
469        method: &str,
470        params: Map<String, Value>,
471        req_timeout: Duration,
472    ) -> Result<String> {
473        let env = new_subscribe(method, params);
474        let id = env.id.clone().unwrap_or_default();
475        // Subscribe confirmation may arrive as next/complete; we treat send success as ok
476        // and rely on subsequent events. Still register a short waiter for error.
477        let (tx, rx) = oneshot::channel();
478        self.shared.rpc_waiters.lock().await.insert(id.clone(), tx);
479        self.send_envelope(env).await?;
480        // Don't block long — many daemons only send `next` without a response.
481        match timeout(req_timeout.min(Duration::from_millis(500)), rx).await {
482            Ok(Ok(Err(de))) => return Err(de.into()),
483            Ok(Ok(Ok(_))) => {}
484            _ => {
485                self.shared.rpc_waiters.lock().await.remove(&id);
486            }
487        }
488        Ok(id)
489    }
490
491    /// Unsubscribe by id.
492    pub async fn unsubscribe(&self, id: &str) -> Result<()> {
493        self.send_envelope(new_unsubscribe(id)).await
494    }
495
496    /// Read next inbound app event (blocks).
497    pub async fn read_event(&self) -> Result<Option<Value>> {
498        loop {
499            {
500                let mut q = self.shared.inbound.lock().await;
501                if let Some(v) = q.pop_front() {
502                    return Ok(Some(v));
503                }
504            }
505            if !self.is_connection_alive() {
506                return Ok(None);
507            }
508            self.shared.inbound_notify.notified().await;
509        }
510    }
511
512    /// Read with timeout; `None` on timeout.
513    pub async fn read_event_with_timeout(&self, dur: Duration) -> Result<Option<Value>> {
514        match timeout(dur, self.read_event()).await {
515            Ok(r) => r,
516            Err(_) => Ok(None),
517        }
518    }
519
520    /// Clear pending inbound events.
521    pub async fn clear_pending_events(&self) {
522        self.shared.inbound.lock().await.clear();
523    }
524
525    /// Peel stale pending control frames at turn start.
526    pub async fn peel_stale_pending_control_events(&self) -> Vec<String> {
527        let mut labels = Vec::new();
528        let mut kept = VecDeque::new();
529        let mut q = self.shared.inbound.lock().await;
530        while let Some(ev) = q.pop_front() {
531            if let Some(label) = stale_pending_frame_label(&ev) {
532                labels.push(label);
533            } else {
534                kept.push_back(ev);
535            }
536        }
537        *q = kept;
538        labels
539    }
540
541    /// Notify `loop_input`.
542    pub async fn send_input(&self, text: &str, opts: SendInputOptions) -> Result<()> {
543        let mut params = Map::new();
544        params.insert("content".into(), json!(text));
545        if let Some(lid) = opts.loop_id {
546            params.insert("loop_id".into(), json!(lid));
547        }
548        if opts.autonomous {
549            params.insert("autonomous".into(), json!(true));
550        }
551        if let Some(n) = opts.max_iterations {
552            params.insert("max_iterations".into(), json!(n));
553        }
554        if let Some(v) = opts.preferred_subagent {
555            params.insert("preferred_subagent".into(), json!(v));
556        }
557        if let Some(v) = opts.model {
558            params.insert("model".into(), json!(v));
559        }
560        if let Some(v) = opts.model_params {
561            params.insert("model_params".into(), v);
562        }
563        if let Some(v) = opts.router_profile {
564            params.insert("router_profile".into(), json!(v));
565        }
566        if let Some(v) = opts.attachments {
567            params.insert("attachments".into(), v);
568        }
569        if let Some(v) = opts.intent_hint {
570            params.insert("intent_hint".into(), json!(v));
571        }
572        if let Some(v) = opts.response_schema {
573            params.insert("response_schema".into(), v);
574        }
575        if let Some(v) = opts.response_schema_name {
576            params.insert("response_schema_name".into(), json!(v));
577        }
578        if let Some(v) = opts.response_schema_strict {
579            params.insert("response_schema_strict".into(), json!(v));
580        }
581        if let Some(v) = opts.clarification_mode {
582            params.insert("clarification_mode".into(), json!(v));
583        }
584        if opts.clarification_answer {
585            params.insert("clarification_answer".into(), json!(true));
586        }
587        if let Some(v) = opts.clarification_answers {
588            params.insert("clarification_answers".into(), v);
589        }
590        self.notify("loop_input", params).await
591    }
592
593    /// `loop_new` RPC.
594    pub async fn loop_new(&self, params: Map<String, Value>) -> Result<Map<String, Value>> {
595        self.request("loop_new", params, Duration::from_secs(30))
596            .await
597    }
598
599    /// `loop_list`.
600    pub async fn loop_list(&self, limit: u32) -> Result<Map<String, Value>> {
601        let mut params = Map::new();
602        params.insert("limit".into(), json!(limit));
603        self.request("loop_list", params, Duration::from_secs(15))
604            .await
605    }
606
607    /// `loop_get`.
608    pub async fn loop_get(&self, loop_id: &str) -> Result<Map<String, Value>> {
609        let mut params = Map::new();
610        params.insert("loop_id".into(), json!(loop_id));
611        self.request("loop_get", params, Duration::from_secs(15))
612            .await
613    }
614
615    /// `loop_reattach`.
616    pub async fn loop_reattach(&self, loop_id: &str) -> Result<Map<String, Value>> {
617        let mut params = Map::new();
618        params.insert("loop_id".into(), json!(loop_id));
619        self.request("loop_reattach", params, Duration::from_secs(15))
620            .await
621    }
622
623    /// Subscribe to `loop_events`.
624    pub async fn loop_subscribe(&self, loop_id: &str, stream_delivery: &str) -> Result<String> {
625        let mut params = Map::new();
626        params.insert("loop_id".into(), json!(loop_id));
627        params.insert("stream_delivery".into(), json!(stream_delivery));
628        params.insert("wire_tier".into(), json!("full"));
629        self.subscribe("loop_events", params, Duration::from_secs(30))
630            .await
631    }
632
633    /// `loop_cards_fetch`.
634    pub async fn loop_cards_fetch(&self, loop_id: &str) -> Result<Map<String, Value>> {
635        let mut params = Map::new();
636        params.insert("loop_id".into(), json!(loop_id));
637        self.request("loop_cards_fetch", params, Duration::from_secs(30))
638            .await
639    }
640
641    /// `loop_history_fetch`.
642    pub async fn loop_history_fetch(&self, loop_id: &str) -> Result<Map<String, Value>> {
643        let mut params = Map::new();
644        params.insert("loop_id".into(), json!(loop_id));
645        self.request("loop_history_fetch", params, Duration::from_secs(30))
646            .await
647    }
648
649    /// `loop_messages`.
650    pub async fn loop_messages(
651        &self,
652        loop_id: &str,
653        limit: u32,
654        offset: u32,
655    ) -> Result<Map<String, Value>> {
656        let mut params = Map::new();
657        params.insert("loop_id".into(), json!(loop_id));
658        params.insert("limit".into(), json!(limit));
659        params.insert("offset".into(), json!(offset));
660        self.request("loop_messages", params, Duration::from_secs(10))
661            .await
662    }
663
664    /// `loop_state_get`.
665    pub async fn loop_state_get(&self, loop_id: &str) -> Result<Map<String, Value>> {
666        let mut params = Map::new();
667        params.insert("loop_id".into(), json!(loop_id));
668        self.request("loop_state_get", params, Duration::from_secs(30))
669            .await
670    }
671
672    /// `loop_state_update`.
673    pub async fn loop_state_update(
674        &self,
675        loop_id: &str,
676        state: Map<String, Value>,
677    ) -> Result<Map<String, Value>> {
678        let mut params = Map::new();
679        params.insert("loop_id".into(), json!(loop_id));
680        params.insert("state".into(), Value::Object(state));
681        self.request("loop_state_update", params, Duration::from_secs(30))
682            .await
683    }
684
685    /// `loop_tree`.
686    pub async fn loop_tree(
687        &self,
688        loop_id: &str,
689        format: Option<&str>,
690    ) -> Result<Map<String, Value>> {
691        let mut params = Map::new();
692        params.insert("loop_id".into(), json!(loop_id));
693        if let Some(f) = format {
694            if !f.is_empty() {
695                params.insert("format".into(), json!(f));
696            }
697        }
698        self.request("loop_tree", params, Duration::from_secs(15))
699            .await
700    }
701
702    /// `loop_prune`.
703    pub async fn loop_prune(
704        &self,
705        loop_id: &str,
706        retention_days: Option<i32>,
707        dry_run: bool,
708    ) -> Result<Map<String, Value>> {
709        let mut params = Map::new();
710        params.insert("loop_id".into(), json!(loop_id));
711        if let Some(d) = retention_days {
712            if d > 0 {
713                params.insert("retention_days".into(), json!(d));
714            }
715        }
716        if dry_run {
717            params.insert("dry_run".into(), json!(true));
718        }
719        self.request("loop_prune", params, Duration::from_secs(30))
720            .await
721    }
722
723    /// `loop_delete`.
724    pub async fn loop_delete(&self, loop_id: &str) -> Result<Map<String, Value>> {
725        let mut params = Map::new();
726        params.insert("loop_id".into(), json!(loop_id));
727        self.request("loop_delete", params, Duration::from_secs(10))
728            .await
729    }
730
731    /// Detach from a loop by unsubscribing (`subscription_id` from `loop_subscribe`).
732    pub async fn loop_detach(&self, subscription_id: &str) -> Result<()> {
733        self.unsubscribe(subscription_id).await
734    }
735
736    /// Authenticate with access/secret keys.
737    pub async fn authenticate(
738        &self,
739        access_key: &str,
740        secret_key: &str,
741    ) -> Result<Map<String, Value>> {
742        let mut params = Map::new();
743        params.insert("access_key".into(), json!(access_key));
744        params.insert("secret_key".into(), json!(secret_key));
745        self.request("auth", params, Duration::from_secs(15)).await
746    }
747
748    /// Refresh auth token.
749    pub async fn refresh_auth_token(&self, refresh_token: &str) -> Result<Map<String, Value>> {
750        let mut params = Map::new();
751        params.insert("refresh_token".into(), json!(refresh_token));
752        self.request("auth_refresh", params, Duration::from_secs(15))
753            .await
754    }
755
756    /// `job_create` on this long-lived connection.
757    pub async fn job_create(
758        &self,
759        goal: &str,
760        workspace: Option<&str>,
761    ) -> Result<Map<String, Value>> {
762        let mut params = Map::new();
763        params.insert("goal".into(), json!(goal));
764        if let Some(ws) = workspace {
765            params.insert("workspace".into(), json!(ws));
766        }
767        self.request("job_create", params, Duration::from_secs(30))
768            .await
769    }
770
771    /// `job_status`.
772    pub async fn job_status(&self, job_id: &str) -> Result<Map<String, Value>> {
773        let mut params = Map::new();
774        params.insert("job_id".into(), json!(job_id));
775        self.request("job_status", params, Duration::from_secs(15))
776            .await
777    }
778
779    /// `job_pause`.
780    pub async fn job_pause(&self, job_id: &str) -> Result<Map<String, Value>> {
781        let mut params = Map::new();
782        params.insert("job_id".into(), json!(job_id));
783        self.request("job_pause", params, Duration::from_secs(15))
784            .await
785    }
786
787    /// `job_resume`.
788    pub async fn job_resume(&self, job_id: &str) -> Result<Map<String, Value>> {
789        let mut params = Map::new();
790        params.insert("job_id".into(), json!(job_id));
791        self.request("job_resume", params, Duration::from_secs(15))
792            .await
793    }
794
795    /// `job_cancel`.
796    pub async fn job_cancel(&self, job_id: &str) -> Result<Map<String, Value>> {
797        let mut params = Map::new();
798        params.insert("job_id".into(), json!(job_id));
799        self.request("job_cancel", params, Duration::from_secs(15))
800            .await
801    }
802
803    /// `job_dag`.
804    pub async fn job_dag(&self, job_id: &str) -> Result<Map<String, Value>> {
805        let mut params = Map::new();
806        params.insert("job_id".into(), json!(job_id));
807        self.request("job_dag", params, Duration::from_secs(15))
808            .await
809    }
810
811    /// `job_guidance`.
812    pub async fn job_guidance(
813        &self,
814        job_id: &str,
815        content: &str,
816        goal_id: Option<&str>,
817    ) -> Result<Map<String, Value>> {
818        let mut params = Map::new();
819        params.insert("job_id".into(), json!(job_id));
820        params.insert("content".into(), json!(content));
821        if let Some(g) = goal_id {
822            params.insert("goal_id".into(), json!(g));
823        }
824        self.request("job_guidance", params, Duration::from_secs(30))
825            .await
826    }
827
828    /// `autopilot_status`.
829    pub async fn autopilot_status(&self) -> Result<Map<String, Value>> {
830        self.request("autopilot_status", Map::new(), Duration::from_secs(15))
831            .await
832    }
833
834    /// `autopilot_submit`.
835    pub async fn autopilot_submit(
836        &self,
837        description: &str,
838        priority: i32,
839        workspace: Option<&str>,
840    ) -> Result<Map<String, Value>> {
841        let mut params = Map::new();
842        params.insert("description".into(), json!(description));
843        params.insert("priority".into(), json!(priority));
844        if let Some(ws) = workspace {
845            params.insert("workspace".into(), json!(ws));
846        }
847        self.request("autopilot_submit", params, Duration::from_secs(30))
848            .await
849    }
850
851    /// `autopilot_list_goals`.
852    pub async fn autopilot_list_goals(&self) -> Result<Map<String, Value>> {
853        self.request("autopilot_list_goals", Map::new(), Duration::from_secs(15))
854            .await
855    }
856
857    /// `autopilot_get_goal`.
858    pub async fn autopilot_get_goal(&self, goal_id: &str) -> Result<Map<String, Value>> {
859        let mut params = Map::new();
860        params.insert("goal_id".into(), json!(goal_id));
861        self.request("autopilot_get_goal", params, Duration::from_secs(15))
862            .await
863    }
864
865    /// `autopilot_cancel_goal`.
866    pub async fn autopilot_cancel_goal(&self, goal_id: &str) -> Result<Map<String, Value>> {
867        let mut params = Map::new();
868        params.insert("goal_id".into(), json!(goal_id));
869        self.request("autopilot_cancel_goal", params, Duration::from_secs(15))
870            .await
871    }
872
873    /// `autopilot_cancel_all`.
874    pub async fn autopilot_cancel_all(&self) -> Result<Map<String, Value>> {
875        self.request("autopilot_cancel_all", Map::new(), Duration::from_secs(15))
876            .await
877    }
878
879    /// `autopilot_wake`.
880    pub async fn autopilot_wake(&self) -> Result<Map<String, Value>> {
881        self.request("autopilot_wake", Map::new(), Duration::from_secs(15))
882            .await
883    }
884
885    /// `autopilot_dream`.
886    pub async fn autopilot_dream(&self) -> Result<Map<String, Value>> {
887        self.request("autopilot_dream", Map::new(), Duration::from_secs(15))
888            .await
889    }
890
891    /// `autopilot_resume`.
892    pub async fn autopilot_resume(&self, goal_id: &str) -> Result<Map<String, Value>> {
893        let mut params = Map::new();
894        params.insert("goal_id".into(), json!(goal_id));
895        self.request("autopilot_resume", params, Duration::from_secs(15))
896            .await
897    }
898
899    /// `autopilot_list_jobs`.
900    pub async fn autopilot_list_jobs(&self) -> Result<Map<String, Value>> {
901        self.request("autopilot_list_jobs", Map::new(), Duration::from_secs(15))
902            .await
903    }
904
905    /// `autopilot_get_job`.
906    pub async fn autopilot_get_job(&self, job_id: &str) -> Result<Map<String, Value>> {
907        let mut params = Map::new();
908        params.insert("job_id".into(), json!(job_id));
909        self.request("autopilot_get_job", params, Duration::from_secs(15))
910            .await
911    }
912
913    /// Subscribe to `autopilot_events` (long-lived worker stream).
914    pub async fn autopilot_subscribe(&self) -> Result<String> {
915        self.subscribe("autopilot_events", Map::new(), Duration::from_secs(15))
916            .await
917    }
918
919    /// Unsubscribe from an autopilot events subscription.
920    pub async fn autopilot_unsubscribe(&self, subscription_id: &str) -> Result<()> {
921        self.unsubscribe(subscription_id).await
922    }
923
924    /// `cron_add`.
925    pub async fn cron_add(&self, text: &str, priority: Option<i32>) -> Result<Map<String, Value>> {
926        let mut params = Map::new();
927        params.insert("text".into(), json!(text));
928        if let Some(p) = priority {
929            params.insert("priority".into(), json!(p));
930        }
931        self.request("cron_add", params, Duration::from_secs(15))
932            .await
933    }
934
935    /// `cron_list`.
936    pub async fn cron_list(&self, status: Option<&str>) -> Result<Map<String, Value>> {
937        let mut params = Map::new();
938        if let Some(s) = status {
939            params.insert("status".into(), json!(s));
940        }
941        self.request("cron_list", params, Duration::from_secs(15))
942            .await
943    }
944
945    /// `cron_show`.
946    pub async fn cron_show(&self, job_id: &str) -> Result<Map<String, Value>> {
947        let mut params = Map::new();
948        params.insert("job_id".into(), json!(job_id));
949        self.request("cron_show", params, Duration::from_secs(15))
950            .await
951    }
952
953    /// `cron_cancel`.
954    pub async fn cron_cancel(&self, job_id: &str) -> Result<Map<String, Value>> {
955        let mut params = Map::new();
956        params.insert("job_id".into(), json!(job_id));
957        self.request("cron_cancel", params, Duration::from_secs(15))
958            .await
959    }
960
961    /// `memory_stats`.
962    pub async fn memory_stats(&self, mode: &str) -> Result<Map<String, Value>> {
963        let mut params = Map::new();
964        params.insert("mode".into(), json!(mode));
965        self.request("memory_stats", params, Duration::from_secs(15))
966            .await
967    }
968
969    /// `skills_list`.
970    pub async fn list_skills(&self) -> Result<Map<String, Value>> {
971        self.request("skills_list", Map::new(), Duration::from_secs(15))
972            .await
973    }
974
975    /// `models_list`.
976    pub async fn list_models(&self) -> Result<Map<String, Value>> {
977        self.request("models_list", Map::new(), Duration::from_secs(15))
978            .await
979    }
980
981    /// `mcp_status`.
982    pub async fn mcp_status(&self) -> Result<Map<String, Value>> {
983        self.request("mcp_status", Map::new(), Duration::from_secs(15))
984            .await
985    }
986
987    /// `daemon_status`.
988    pub async fn fetch_daemon_status(&self) -> Result<Map<String, Value>> {
989        self.request("daemon_status", Map::new(), Duration::from_secs(10))
990            .await
991    }
992
993    /// `invoke_skill` on this connection (stream socket for turn enqueue).
994    pub async fn invoke_skill(&self, skill: &str, args: &str) -> Result<Map<String, Value>> {
995        let mut params = Map::new();
996        params.insert("skill".into(), json!(skill));
997        if !args.is_empty() {
998            params.insert("args".into(), json!(args));
999        }
1000        self.request("invoke_skill", params, Duration::from_secs(120))
1001            .await
1002    }
1003
1004    /// Reattach + subscribe + `loop_get` probe.
1005    pub async fn reattach_and_probe(&self, loop_id: &str) -> Result<()> {
1006        self.loop_reattach(loop_id).await?;
1007        self.loop_subscribe(loop_id, "adaptive").await?;
1008        match self.loop_get(loop_id).await {
1009            Ok(_) => Ok(()),
1010            Err(Error::Daemon(de)) if de.code == -32200 => {
1011                Err(StaleLoopError::new(loop_id, Some(Box::new(de))).into())
1012            }
1013            Err(e) => Err(StaleLoopError::new(loop_id, Some(Box::new(e))).into()),
1014        }
1015    }
1016}
1017
1018impl Shared {
1019    async fn send_raw(&self, msg: Message) -> Result<()> {
1020        let tx = self.write_tx.lock().await;
1021        let Some(tx) = tx.as_ref() else {
1022            return Err(Error::protocol("not connected"));
1023        };
1024        tx.send(msg)
1025            .map_err(|_| Error::protocol("write channel closed"))?;
1026        Ok(())
1027    }
1028
1029    async fn mark_disconnected(&self, cause: DisconnectCause) {
1030        self.connected.store(false, Ordering::SeqCst);
1031        self.reader_alive.store(false, Ordering::SeqCst);
1032        {
1033            let mut slot = self.write_tx.lock().await;
1034            *slot = None;
1035        }
1036        {
1037            let mut c = self.disconnect_cause.lock().await;
1038            if c.is_none() {
1039                *c = Some(cause);
1040            }
1041        }
1042        // Fail pending RPCs.
1043        let mut waiters = self.rpc_waiters.lock().await;
1044        for (_, tx) in waiters.drain() {
1045            let _ = tx.send(Err(DaemonError::new(-1, "disconnected")));
1046        }
1047        self.inbound_notify.notify_waiters();
1048        self.disconnect_notify.notify_waiters();
1049    }
1050
1051    async fn route_inbound(&self, msg: Value) {
1052        // Auto pong for app-level ping.
1053        if msg.get("type").and_then(|v| v.as_str()) == Some("ping") {
1054            let _ = self
1055                .send_raw(Message::Text(
1056                    new_pong().to_wire_json().unwrap_or_default().into(),
1057                ))
1058                .await;
1059            return;
1060        }
1061
1062        if msg.get("type").and_then(|v| v.as_str()) == Some("pong") {
1063            if let Some(tracker) = self.heartbeat_tracker.lock().expect("hb").as_ref() {
1064                tracker.note_pong();
1065            }
1066            return;
1067        }
1068
1069        // Daemon heartbeat custom events (namespace / type heuristics).
1070        let ns = msg
1071            .get("namespace")
1072            .or_else(|| msg.get("type"))
1073            .and_then(|v| v.as_str())
1074            .unwrap_or("");
1075        if ns.contains("heartbeat") {
1076            if let Some(tracker) = self.heartbeat_tracker.lock().expect("hb").as_ref() {
1077                tracker.update(msg.get("data"));
1078            }
1079        }
1080
1081        if inbound_needs_delivery_ack(&msg) {
1082            let loop_id = extract_loop_id_from_inbound(&msg);
1083            if !loop_id.is_empty() {
1084                let seq = {
1085                    let mut map = self.delivery_seq.lock().await;
1086                    let e = map.entry(loop_id.clone()).or_insert(0);
1087                    *e += 1;
1088                    *e
1089                };
1090                let mut params = Map::new();
1091                params.insert("loop_id".into(), json!(loop_id));
1092                params.insert("seq".into(), json!(seq));
1093                let _ = self
1094                    .send_raw(Message::Text(
1095                        new_notification("delivery_ack", params)
1096                            .to_wire_json()
1097                            .unwrap_or_default()
1098                            .into(),
1099                    ))
1100                    .await;
1101            }
1102        }
1103
1104        let msg_type = msg.get("type").and_then(|v| v.as_str()).unwrap_or("");
1105        let id = msg
1106            .get("id")
1107            .and_then(|v| v.as_str())
1108            .unwrap_or("")
1109            .to_string();
1110
1111        if matches!(msg_type, "response" | "error") && !id.is_empty() {
1112            if let Some(waiter) = self.rpc_waiters.lock().await.remove(&id) {
1113                let result = if msg_type == "error" {
1114                    Err(parse_daemon_error(&msg))
1115                } else {
1116                    Ok(msg
1117                        .get("result")
1118                        .cloned()
1119                        .unwrap_or(Value::Object(Map::new())))
1120                };
1121                let _ = waiter.send(result);
1122                return;
1123            }
1124        }
1125
1126        // Push to inbound queue with priority backpressure.
1127        let critical = is_critical_inbound(&msg);
1128        let mut q = self.inbound.lock().await;
1129        if q.len() >= self.max_inbound {
1130            if critical {
1131                if let Some(pos) = q.iter().position(|m| !is_critical_inbound(m)) {
1132                    q.remove(pos);
1133                    self.note_inbound_drop("priority_evict");
1134                } else {
1135                    self.note_inbound_drop("queue_full_critical");
1136                    return;
1137                }
1138            } else {
1139                self.note_inbound_drop("queue_full");
1140                return;
1141            }
1142        }
1143        q.push_back(msg);
1144        self.inbound_notify.notify_one();
1145    }
1146
1147    fn note_inbound_drop(&self, reason: &str) {
1148        let n = self.inbound_dropped.fetch_add(1, Ordering::SeqCst) + 1;
1149        if !self.degraded_notified.swap(true, Ordering::SeqCst) {
1150            if let Some(cb) = self.stream_degraded_cb.lock().expect("cb").as_ref() {
1151                cb(n, reason.to_string());
1152            }
1153        }
1154    }
1155}
1156
1157fn is_critical_inbound(msg: &Value) -> bool {
1158    let t = msg.get("type").and_then(|v| v.as_str()).unwrap_or("");
1159    matches!(
1160        t,
1161        "status" | "error" | "complete" | "connection_ack" | "response"
1162    ) || inbound_needs_delivery_ack(msg)
1163}
1164
1165fn parse_daemon_error(msg: &Value) -> DaemonError {
1166    if let Some(err) = msg.get("error") {
1167        let code = err.get("code").and_then(|c| c.as_i64()).unwrap_or(-1);
1168        let message = err
1169            .get("message")
1170            .and_then(|m| m.as_str())
1171            .unwrap_or("daemon error")
1172            .to_string();
1173        let data = err.get("data").cloned();
1174        let mut de = DaemonError::new(code, message);
1175        if let Some(d) = data {
1176            de = de.with_data(d);
1177        }
1178        return de;
1179    }
1180    DaemonError::new(
1181        -1,
1182        msg.get("message")
1183            .and_then(|m| m.as_str())
1184            .unwrap_or("daemon error"),
1185    )
1186}
1187
1188/// Re-export unwrap helper for appkit.
1189pub use crate::protocol::unwrap_next as unwrap_next_frame;
1190
1191#[cfg(test)]
1192mod tests {
1193    use super::*;
1194
1195    #[test]
1196    fn send_input_options_default() {
1197        let o = SendInputOptions::default();
1198        assert!(!o.autonomous);
1199        assert!(o.loop_id.is_none());
1200    }
1201}