soothe_client/appkit/
daemon_session.rs

1//! Dual-socket DaemonSession for one conversation with turn streaming.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use serde_json::{json, Map, Value};
7use tokio::sync::Mutex;
8
9use crate::client::{unwrap_next_frame, Client, SendInputOptions};
10use crate::errors::{Error, Result};
11use crate::session::{bootstrap_loop_session, connect_with_retries, BootstrapOptions};
12use crate::stream_terminal::{is_turn_end_custom_data, is_turn_progress_chunk, STREAM_END};
13
14use super::chunk_filter::should_drop_stream_chunk_early;
15use super::observability::TurnEventStats;
16
17/// Default post-idle drain window (Go `DefaultPostIdleDrain`).
18pub const DEFAULT_POST_IDLE_DRAIN: Duration = Duration::from_millis(500);
19
20/// Filters non-actionable stream chunks before yield (Go `EarlyDropFn`).
21pub type EarlyDropFn = Arc<dyn Fn(&[Value], &str, &Value) -> bool + Send + Sync>;
22
23/// Options for constructing a DaemonSession.
24#[derive(Clone)]
25pub struct DaemonSessionOptions {
26    /// Workspace path.
27    pub workspace: Option<String>,
28    /// Stream delivery mode.
29    pub stream_delivery: String,
30    /// Post-idle drain window.
31    pub post_idle_drain: Duration,
32    /// Optional early-drop filter (defaults to [`should_drop_stream_chunk_early`]).
33    pub early_drop_fn: Option<EarlyDropFn>,
34}
35
36impl Default for DaemonSessionOptions {
37    fn default() -> Self {
38        Self {
39            workspace: None,
40            stream_delivery: "adaptive".into(),
41            post_idle_drain: DEFAULT_POST_IDLE_DRAIN,
42            early_drop_fn: None,
43        }
44    }
45}
46
47impl std::fmt::Debug for DaemonSessionOptions {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        f.debug_struct("DaemonSessionOptions")
50            .field("workspace", &self.workspace)
51            .field("stream_delivery", &self.stream_delivery)
52            .field("post_idle_drain", &self.post_idle_drain)
53            .field(
54                "early_drop_fn",
55                &self.early_drop_fn.as_ref().map(|_| "<fn>"),
56            )
57            .finish()
58    }
59}
60
61/// Options for `send_turn`.
62#[derive(Debug, Clone, Default)]
63pub struct SendTurnOptions {
64    /// Autonomous mode.
65    pub autonomous: bool,
66    /// Max iterations.
67    pub max_iterations: Option<u32>,
68    /// Preferred subagent.
69    pub preferred_subagent: Option<String>,
70    /// Model override.
71    pub model: Option<String>,
72    /// Model params.
73    pub model_params: Option<Value>,
74    /// Attachments.
75    pub attachments: Option<Value>,
76    /// Clarification mode.
77    pub clarification_mode: Option<String>,
78    /// Clarification answer.
79    pub clarification_answer: bool,
80    /// Intent hint.
81    pub intent_hint: Option<String>,
82}
83
84/// One streamed turn chunk.
85#[derive(Debug, Clone)]
86pub struct TurnChunk {
87    /// Namespace path.
88    pub namespace: Value,
89    /// Mode (`messages`, `custom`, …).
90    pub mode: String,
91    /// Payload data.
92    pub data: Value,
93}
94
95/// Dual-socket session: stream + lazy RPC sidecar.
96pub struct DaemonSession {
97    opts: DaemonSessionOptions,
98    client: Client,
99    rpc_client: Client,
100    rpc_connected: Mutex<bool>,
101    loop_id: Mutex<String>,
102    read_lock: Mutex<()>,
103    early_drop_fn: EarlyDropFn,
104    /// Per-turn stream filtering counters.
105    pub turn_event_stats: Mutex<TurnEventStats>,
106    /// Last turn end state label.
107    pub last_turn_end_state: Mutex<String>,
108    /// Whether cancel was observed for the last turn.
109    pub last_turn_cancel_seen: Mutex<bool>,
110    /// Last turn error message.
111    pub last_turn_error_message: Mutex<String>,
112}
113
114impl DaemonSession {
115    /// Create a session for `ws_url`.
116    pub fn new(ws_url: impl Into<String>, opts: Option<DaemonSessionOptions>) -> Self {
117        let ws_url = ws_url.into();
118        let opts = opts.unwrap_or_default();
119        let early_drop_fn = opts.early_drop_fn.clone().unwrap_or_else(|| {
120            Arc::new(|ns: &[Value], mode: &str, data: &Value| {
121                should_drop_stream_chunk_early(ns, mode, data)
122            })
123        });
124        Self {
125            client: Client::new(&ws_url),
126            rpc_client: Client::new(&ws_url),
127            rpc_connected: Mutex::new(false),
128            loop_id: Mutex::new(String::new()),
129            read_lock: Mutex::new(()),
130            early_drop_fn,
131            turn_event_stats: Mutex::new(TurnEventStats::new()),
132            last_turn_end_state: Mutex::new(String::new()),
133            last_turn_cancel_seen: Mutex::new(false),
134            last_turn_error_message: Mutex::new(String::new()),
135            opts,
136        }
137    }
138
139    /// Stream socket.
140    pub fn stream_client(&self) -> &Client {
141        &self.client
142    }
143
144    /// RPC sidecar socket (lazy-connected for list/cards/history).
145    pub fn rpc_client(&self) -> &Client {
146        &self.rpc_client
147    }
148
149    /// Active loop id.
150    pub async fn loop_id(&self) -> String {
151        self.loop_id.lock().await.clone()
152    }
153
154    /// Connect and bootstrap a loop (or resume).
155    pub async fn connect(&self, resume_loop_id: Option<&str>) -> Result<Map<String, Value>> {
156        connect_with_retries(&self.client, 40, Duration::from_millis(250)).await?;
157        self.bootstrap_loop(resume_loop_id).await
158    }
159
160    async fn bootstrap_loop(&self, resume_loop_id: Option<&str>) -> Result<Map<String, Value>> {
161        let mut boot = BootstrapOptions::new();
162        boot.resume_loop_id = resume_loop_id.map(|s| s.to_string());
163        boot.workspace = self.opts.workspace.clone();
164        boot.stream_delivery = self.opts.stream_delivery.clone();
165        let ready = bootstrap_loop_session(&self.client, boot, None).await?;
166        if let Some(lid) = ready.get("loop_id").and_then(|v| v.as_str()) {
167            *self.loop_id.lock().await = lid.to_string();
168        }
169        Ok(ready)
170    }
171
172    /// Start a fresh loop on the stream socket.
173    pub async fn new_loop(&self) -> Result<Map<String, Value>> {
174        self.bootstrap_loop(None).await
175    }
176
177    /// Switch to an existing loop.
178    pub async fn switch_loop(&self, loop_id: &str) -> Result<Map<String, Value>> {
179        self.bootstrap_loop(Some(loop_id)).await
180    }
181
182    /// Reconnect + reattach, or fresh bootstrap if stale.
183    pub async fn ensure_connected(&self) -> Result<()> {
184        if self.client.is_connection_alive() {
185            return Ok(());
186        }
187        connect_with_retries(&self.client, 40, Duration::from_millis(250)).await?;
188        let lid = self.loop_id().await;
189        if lid.is_empty() {
190            self.bootstrap_loop(None).await?;
191            return Ok(());
192        }
193        match self.client.reattach_and_probe(&lid).await {
194            Ok(()) => Ok(()),
195            Err(Error::StaleLoop(_)) | Err(_) => {
196                // Close RPC sidecar; fresh bootstrap.
197                let _ = self.rpc_client.close().await;
198                *self.rpc_connected.lock().await = false;
199                self.bootstrap_loop(None).await?;
200                Ok(())
201            }
202        }
203    }
204
205    /// Close both sockets.
206    pub async fn close(&self) -> Result<()> {
207        let _ = self.client.close().await;
208        let _ = self.rpc_client.close().await;
209        *self.rpc_connected.lock().await = false;
210        Ok(())
211    }
212
213    /// Notify disconnect (loops keep running server-side).
214    pub async fn detach(&self) -> Result<()> {
215        self.client.notify("disconnect", Map::new()).await
216    }
217
218    /// Send a user turn on the stream socket.
219    pub async fn send_turn(&self, text: &str, opts: Option<SendTurnOptions>) -> Result<()> {
220        let loop_id = self.loop_id().await;
221        if loop_id.is_empty() {
222            return Err(Error::msg("no active loop session"));
223        }
224        let opts = opts.unwrap_or_default();
225        let input = SendInputOptions {
226            loop_id: Some(loop_id),
227            autonomous: opts.autonomous,
228            max_iterations: opts.max_iterations,
229            preferred_subagent: opts.preferred_subagent,
230            model: opts.model,
231            model_params: opts.model_params,
232            attachments: opts.attachments,
233            clarification_mode: opts.clarification_mode,
234            clarification_answer: opts.clarification_answer,
235            intent_hint: opts.intent_hint,
236            ..Default::default()
237        };
238        self.client.send_input(text, input).await
239    }
240
241    /// Cancel active turn via `/cancel`.
242    pub async fn cancel_active_turn(&self) -> Result<()> {
243        let mut params = Map::new();
244        params.insert("cmd".into(), json!("/cancel"));
245        self.client.notify("slash_command", params).await
246    }
247
248    async fn ensure_rpc_connected(&self) -> Result<()> {
249        let mut flag = self.rpc_connected.lock().await;
250        if *flag && self.rpc_client.is_connected() {
251            return Ok(());
252        }
253        connect_with_retries(&self.rpc_client, 5, Duration::from_millis(250)).await?;
254        *flag = true;
255        Ok(())
256    }
257
258    /// List loops via RPC sidecar.
259    pub async fn list_loops(&self, limit: u32) -> Result<Map<String, Value>> {
260        self.ensure_rpc_connected().await?;
261        let lim = if limit == 0 { 20 } else { limit };
262        self.rpc_client.loop_list(lim).await
263    }
264
265    /// Fetch cards via RPC sidecar.
266    pub async fn fetch_loop_cards(&self, loop_id: &str) -> Result<Map<String, Value>> {
267        self.ensure_rpc_connected().await?;
268        self.rpc_client.loop_cards_fetch(loop_id).await
269    }
270
271    /// Fetch history via RPC sidecar.
272    pub async fn fetch_loop_history(&self, loop_id: &str) -> Result<Map<String, Value>> {
273        self.ensure_rpc_connected().await?;
274        self.rpc_client.loop_history_fetch(loop_id).await
275    }
276
277    /// Stream turn chunks until idle / stream.end.
278    ///
279    /// `max_wait` of `None` means no absolute deadline.
280    pub async fn iter_turn_chunks(&self, max_wait: Option<Duration>) -> Result<Vec<TurnChunk>> {
281        let _guard = self.read_lock.lock().await;
282        self.iter_turn_chunks_locked(max_wait).await
283    }
284
285    async fn iter_turn_chunks_locked(&self, max_wait: Option<Duration>) -> Result<Vec<TurnChunk>> {
286        *self.last_turn_end_state.lock().await = String::new();
287        *self.last_turn_error_message.lock().await = String::new();
288        *self.last_turn_cancel_seen.lock().await = false;
289        *self.turn_event_stats.lock().await = TurnEventStats::new();
290
291        let mut out = Vec::new();
292        let mut query_started = false;
293        let mut expected_loop_id = self.loop_id().await;
294        let mut stream_payload_seen = false;
295        let mut turn_progress_seen = false;
296        let mut cancel_seen = false;
297        let absolute_deadline = max_wait.map(|d| tokio::time::Instant::now() + d);
298
299        let _ = self.client.peel_stale_pending_control_events().await;
300
301        loop {
302            if let Some(deadline) = absolute_deadline {
303                if tokio::time::Instant::now() > deadline {
304                    let err = format!(
305                        "turn timed out after {:?} (loop={})",
306                        max_wait.unwrap_or_default(),
307                        expected_loop_id
308                    );
309                    *self.last_turn_error_message.lock().await = err.clone();
310                    return Err(Error::msg(err));
311                }
312            }
313
314            let ev = self
315                .client
316                .read_event_with_timeout(Duration::from_millis(250))
317                .await?;
318            let Some(ev) = ev else {
319                if query_started && !self.client.is_connection_alive() {
320                    *self.last_turn_end_state.lock().await = "connection_lost".into();
321                    return Err(Error::msg("daemon connection lost"));
322                }
323                // Idle timeout on read — keep waiting unless absolute deadline.
324                continue;
325            };
326
327            let mut frame = ev;
328            let mut event_type = frame
329                .get("type")
330                .and_then(|v| v.as_str())
331                .unwrap_or("")
332                .to_string();
333            if event_type == "next" {
334                frame = unwrap_next_frame(&frame);
335                event_type = frame
336                    .get("type")
337                    .and_then(|v| v.as_str())
338                    .unwrap_or("")
339                    .to_string();
340            }
341
342            let event_loop_id = frame
343                .get("loop_id")
344                .and_then(|v| v.as_str())
345                .unwrap_or("")
346                .to_string();
347            if !expected_loop_id.is_empty()
348                && !event_loop_id.is_empty()
349                && event_loop_id != expected_loop_id
350            {
351                continue;
352            }
353
354            if event_type == "error" {
355                let msg = frame
356                    .get("error")
357                    .and_then(|e| e.get("message"))
358                    .and_then(|m| m.as_str())
359                    .or_else(|| frame.get("message").and_then(|m| m.as_str()))
360                    .unwrap_or("daemon error")
361                    .to_string();
362                *self.last_turn_error_message.lock().await = msg.clone();
363                return Err(Error::msg(msg));
364            }
365
366            if event_type == "status" {
367                if let Some(lid) = frame.get("loop_id").and_then(|v| v.as_str()) {
368                    if !lid.is_empty() {
369                        *self.loop_id.lock().await = lid.to_string();
370                        expected_loop_id = lid.to_string();
371                    }
372                }
373                let state = frame.get("state").and_then(|v| v.as_str()).unwrap_or("");
374                match state {
375                    "running" => query_started = true,
376                    "stopped" if query_started => {
377                        *self.last_turn_end_state.lock().await = state.into();
378                        self.drain_after_idle(&expected_loop_id, &mut out).await;
379                        return Ok(out);
380                    }
381                    "idle" if query_started => {
382                        if !stream_payload_seen && !cancel_seen {
383                            continue;
384                        }
385                        *self.last_turn_end_state.lock().await = state.into();
386                        self.drain_after_idle(&expected_loop_id, &mut out).await;
387                        return Ok(out);
388                    }
389                    _ => {}
390                }
391                continue;
392            }
393
394            if event_type == "command_response" {
395                let content = frame.get("content").and_then(|v| v.as_str()).unwrap_or("");
396                if content.contains("Cancellation requested") {
397                    cancel_seen = true;
398                    *self.last_turn_cancel_seen.lock().await = true;
399                }
400                continue;
401            }
402
403            if event_type != "event" {
404                continue;
405            }
406
407            let data = frame.get("data").cloned().unwrap_or(Value::Null);
408            let namespace = frame
409                .get("namespace")
410                .cloned()
411                .unwrap_or(Value::Array(vec![]));
412            let mode = frame
413                .get("mode")
414                .and_then(|v| v.as_str())
415                .unwrap_or("")
416                .to_string();
417
418            let ns_slice: Vec<Value> = match &namespace {
419                Value::Array(a) => a.clone(),
420                _ => vec![],
421            };
422            if (self.early_drop_fn)(&ns_slice, &mode, &data) {
423                self.turn_event_stats.lock().await.filtered_early += 1;
424                continue;
425            }
426
427            if mode == "custom"
428                && is_turn_end_custom_data(&data)
429                && (!query_started || !turn_progress_seen)
430            {
431                continue;
432            }
433
434            stream_payload_seen = true;
435            if is_turn_progress_chunk(&mode, &data) {
436                turn_progress_seen = true;
437            }
438
439            out.push(TurnChunk {
440                namespace,
441                mode: mode.clone(),
442                data: data.clone(),
443            });
444
445            if mode == "custom" && is_turn_end_custom_data(&data) {
446                let custom_type = data.get("type").and_then(|v| v.as_str()).unwrap_or("");
447                *self.last_turn_end_state.lock().await = if custom_type == STREAM_END {
448                    "stream_end".into()
449                } else {
450                    "completed".into()
451                };
452                self.drain_after_idle(&expected_loop_id, &mut out).await;
453                return Ok(out);
454            }
455        }
456    }
457
458    async fn drain_after_idle(&self, expected_loop_id: &str, out: &mut Vec<TurnChunk>) {
459        let deadline = tokio::time::Instant::now() + self.opts.post_idle_drain;
460        while tokio::time::Instant::now() < deadline {
461            let ev = match self
462                .client
463                .read_event_with_timeout(Duration::from_millis(250))
464                .await
465            {
466                Ok(Some(e)) => e,
467                _ => return,
468            };
469            let mut frame = ev;
470            let mut event_type = frame
471                .get("type")
472                .and_then(|v| v.as_str())
473                .unwrap_or("")
474                .to_string();
475            if event_type == "next" {
476                frame = unwrap_next_frame(&frame);
477                event_type = frame
478                    .get("type")
479                    .and_then(|v| v.as_str())
480                    .unwrap_or("")
481                    .to_string();
482            }
483            let event_loop_id = frame.get("loop_id").and_then(|v| v.as_str()).unwrap_or("");
484            if !expected_loop_id.is_empty()
485                && !event_loop_id.is_empty()
486                && event_loop_id != expected_loop_id
487            {
488                continue;
489            }
490            if event_type == "error" {
491                return;
492            }
493            if event_type == "status" {
494                if let Some(lid) = frame.get("loop_id").and_then(|v| v.as_str()) {
495                    if !lid.is_empty() {
496                        *self.loop_id.lock().await = lid.to_string();
497                    }
498                }
499                continue;
500            }
501            if event_type != "event" {
502                continue;
503            }
504            let data = frame.get("data").cloned().unwrap_or(Value::Null);
505            let namespace = frame
506                .get("namespace")
507                .cloned()
508                .unwrap_or(Value::Array(vec![]));
509            let mode = frame
510                .get("mode")
511                .and_then(|v| v.as_str())
512                .unwrap_or("")
513                .to_string();
514            let ns_slice: Vec<Value> = match &namespace {
515                Value::Array(a) => a.clone(),
516                _ => vec![],
517            };
518            if (self.early_drop_fn)(&ns_slice, &mode, &data) {
519                self.turn_event_stats.lock().await.filtered_early += 1;
520                continue;
521            }
522            self.turn_event_stats.lock().await.post_idle_drained += 1;
523            out.push(TurnChunk {
524                namespace,
525                mode,
526                data,
527            });
528        }
529    }
530}
soothe_client/appkit/daemon_session.rs

soothe_client/appkit/
daemon_session.rs
