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 history via RPC sidecar.
266    pub async fn fetch_loop_history(&self, loop_id: &str) -> Result<Map<String, Value>> {
267        self.ensure_rpc_connected().await?;
268        self.rpc_client.loop_history_fetch(loop_id).await
269    }
270
271    /// Stream turn chunks until idle / stream.end.
272    ///
273    /// `max_wait` of `None` means no absolute deadline.
274    pub async fn iter_turn_chunks(&self, max_wait: Option<Duration>) -> Result<Vec<TurnChunk>> {
275        let _guard = self.read_lock.lock().await;
276        self.iter_turn_chunks_locked(max_wait).await
277    }
278
279    async fn iter_turn_chunks_locked(&self, max_wait: Option<Duration>) -> Result<Vec<TurnChunk>> {
280        *self.last_turn_end_state.lock().await = String::new();
281        *self.last_turn_error_message.lock().await = String::new();
282        *self.last_turn_cancel_seen.lock().await = false;
283        *self.turn_event_stats.lock().await = TurnEventStats::new();
284
285        let mut out = Vec::new();
286        let mut query_started = false;
287        let mut expected_loop_id = self.loop_id().await;
288        let mut stream_payload_seen = false;
289        let mut turn_progress_seen = false;
290        let mut cancel_seen = false;
291        let absolute_deadline = max_wait.map(|d| tokio::time::Instant::now() + d);
292
293        let _ = self.client.peel_stale_pending_control_events().await;
294
295        loop {
296            if let Some(deadline) = absolute_deadline {
297                if tokio::time::Instant::now() > deadline {
298                    let err = format!(
299                        "turn timed out after {:?} (loop={})",
300                        max_wait.unwrap_or_default(),
301                        expected_loop_id
302                    );
303                    *self.last_turn_error_message.lock().await = err.clone();
304                    return Err(Error::msg(err));
305                }
306            }
307
308            let ev = self
309                .client
310                .read_event_with_timeout(Duration::from_millis(250))
311                .await?;
312            let Some(ev) = ev else {
313                if query_started && !self.client.is_connection_alive() {
314                    *self.last_turn_end_state.lock().await = "connection_lost".into();
315                    return Err(Error::msg("daemon connection lost"));
316                }
317                // Idle timeout on read — keep waiting unless absolute deadline.
318                continue;
319            };
320
321            let mut frame = ev;
322            let mut event_type = frame
323                .get("type")
324                .and_then(|v| v.as_str())
325                .unwrap_or("")
326                .to_string();
327            if event_type == "next" {
328                frame = unwrap_next_frame(&frame);
329                event_type = frame
330                    .get("type")
331                    .and_then(|v| v.as_str())
332                    .unwrap_or("")
333                    .to_string();
334            }
335
336            let event_loop_id = frame
337                .get("loop_id")
338                .and_then(|v| v.as_str())
339                .unwrap_or("")
340                .to_string();
341            if !expected_loop_id.is_empty()
342                && !event_loop_id.is_empty()
343                && event_loop_id != expected_loop_id
344            {
345                continue;
346            }
347
348            if event_type == "error" {
349                let msg = frame
350                    .get("error")
351                    .and_then(|e| e.get("message"))
352                    .and_then(|m| m.as_str())
353                    .or_else(|| frame.get("message").and_then(|m| m.as_str()))
354                    .unwrap_or("daemon error")
355                    .to_string();
356                *self.last_turn_error_message.lock().await = msg.clone();
357                return Err(Error::msg(msg));
358            }
359
360            if event_type == "status" {
361                if let Some(lid) = frame.get("loop_id").and_then(|v| v.as_str()) {
362                    if !lid.is_empty() {
363                        *self.loop_id.lock().await = lid.to_string();
364                        expected_loop_id = lid.to_string();
365                    }
366                }
367                let state = frame.get("state").and_then(|v| v.as_str()).unwrap_or("");
368                match state {
369                    "running" => query_started = true,
370                    "stopped" if query_started => {
371                        *self.last_turn_end_state.lock().await = state.into();
372                        self.drain_after_idle(&expected_loop_id, &mut out).await;
373                        return Ok(out);
374                    }
375                    "idle" if query_started => {
376                        if !stream_payload_seen && !cancel_seen {
377                            continue;
378                        }
379                        *self.last_turn_end_state.lock().await = state.into();
380                        self.drain_after_idle(&expected_loop_id, &mut out).await;
381                        return Ok(out);
382                    }
383                    _ => {}
384                }
385                continue;
386            }
387
388            if event_type == "command_response" {
389                let content = frame.get("content").and_then(|v| v.as_str()).unwrap_or("");
390                if content.contains("Cancellation requested") {
391                    cancel_seen = true;
392                    *self.last_turn_cancel_seen.lock().await = true;
393                }
394                continue;
395            }
396
397            if event_type != "event" {
398                continue;
399            }
400
401            let data = frame.get("data").cloned().unwrap_or(Value::Null);
402            let namespace = frame
403                .get("namespace")
404                .cloned()
405                .unwrap_or(Value::Array(vec![]));
406            let mode = frame
407                .get("mode")
408                .and_then(|v| v.as_str())
409                .unwrap_or("")
410                .to_string();
411
412            let ns_slice: Vec<Value> = match &namespace {
413                Value::Array(a) => a.clone(),
414                _ => vec![],
415            };
416            if (self.early_drop_fn)(&ns_slice, &mode, &data) {
417                self.turn_event_stats.lock().await.filtered_early += 1;
418                continue;
419            }
420
421            if mode == "custom"
422                && is_turn_end_custom_data(&data)
423                && (!query_started || !turn_progress_seen)
424            {
425                continue;
426            }
427
428            stream_payload_seen = true;
429            if is_turn_progress_chunk(&mode, &data) {
430                turn_progress_seen = true;
431            }
432
433            out.push(TurnChunk {
434                namespace,
435                mode: mode.clone(),
436                data: data.clone(),
437            });
438
439            if mode == "custom" && is_turn_end_custom_data(&data) {
440                let custom_type = data.get("type").and_then(|v| v.as_str()).unwrap_or("");
441                *self.last_turn_end_state.lock().await = if custom_type == STREAM_END {
442                    "stream_end".into()
443                } else {
444                    "completed".into()
445                };
446                self.drain_after_idle(&expected_loop_id, &mut out).await;
447                return Ok(out);
448            }
449        }
450    }
451
452    async fn drain_after_idle(&self, expected_loop_id: &str, out: &mut Vec<TurnChunk>) {
453        let deadline = tokio::time::Instant::now() + self.opts.post_idle_drain;
454        while tokio::time::Instant::now() < deadline {
455            let ev = match self
456                .client
457                .read_event_with_timeout(Duration::from_millis(250))
458                .await
459            {
460                Ok(Some(e)) => e,
461                _ => return,
462            };
463            let mut frame = ev;
464            let mut event_type = frame
465                .get("type")
466                .and_then(|v| v.as_str())
467                .unwrap_or("")
468                .to_string();
469            if event_type == "next" {
470                frame = unwrap_next_frame(&frame);
471                event_type = frame
472                    .get("type")
473                    .and_then(|v| v.as_str())
474                    .unwrap_or("")
475                    .to_string();
476            }
477            let event_loop_id = frame.get("loop_id").and_then(|v| v.as_str()).unwrap_or("");
478            if !expected_loop_id.is_empty()
479                && !event_loop_id.is_empty()
480                && event_loop_id != expected_loop_id
481            {
482                continue;
483            }
484            if event_type == "error" {
485                return;
486            }
487            if event_type == "status" {
488                if let Some(lid) = frame.get("loop_id").and_then(|v| v.as_str()) {
489                    if !lid.is_empty() {
490                        *self.loop_id.lock().await = lid.to_string();
491                    }
492                }
493                continue;
494            }
495            if event_type != "event" {
496                continue;
497            }
498            let data = frame.get("data").cloned().unwrap_or(Value::Null);
499            let namespace = frame
500                .get("namespace")
501                .cloned()
502                .unwrap_or(Value::Array(vec![]));
503            let mode = frame
504                .get("mode")
505                .and_then(|v| v.as_str())
506                .unwrap_or("")
507                .to_string();
508            let ns_slice: Vec<Value> = match &namespace {
509                Value::Array(a) => a.clone(),
510                _ => vec![],
511            };
512            if (self.early_drop_fn)(&ns_slice, &mode, &data) {
513                self.turn_event_stats.lock().await.filtered_early += 1;
514                continue;
515            }
516            self.turn_event_stats.lock().await.post_idle_drained += 1;
517            out.push(TurnChunk {
518                namespace,
519                mode,
520                data,
521            });
522        }
523    }
524}
soothe_client/appkit/daemon_session.rs

soothe_client/appkit/
daemon_session.rs
