omni_dev/sessions/
codex_app_server.rs

1//! A **poll-only** observer of a Codex app-server's thread status.
2//!
3//! The source of exact state for Codex sessions `omni-dev codex-wrap` launches
4//! (#1910, ADR-0088), the Codex analogue of the `claude-wrap` stream tracker.
5//!
6//! Codex's app-server speaks JSON-RPC over a WebSocket; with
7//! `--listen unix://PATH` the WebSocket rides a Unix socket. Its
8//! `thread/loaded/list` names the threads in memory and `thread/read` gives
9//! each one's `status`: `idle`, `systemError`, `notLoaded`, or `active` with
10//! `activeFlags` `waitingOnApproval` / `waitingOnUserInput`. That is the state
11//! the Codex hooks can only infer.
12//!
13//! Two rules keep an observer from ever acting on a session:
14//!
15//! - **Poll, never subscribe.** The observer calls only `initialize`,
16//!   `thread/loaded/list` and `thread/read`. It never starts, resumes or
17//!   subscribes to a thread, so the server never routes a thread's approval
18//!   requests to it (verified on 0.155.1: an observer connection received no
19//!   server request while the TUI on the same server sat on an approval).
20//! - **Never answer.** If a server request does arrive, it is dropped
21//!   unanswered, so an observer cannot approve, deny or duplicate one.
22//!
23//! This module holds the client and the pure mapping; the process wiring lives
24//! in [`crate::cli::codex_wrap`].
25
26use std::collections::HashMap;
27use std::path::{Path, PathBuf};
28use std::time::Duration;
29
30use anyhow::{anyhow, bail, Context, Result};
31use futures::{SinkExt, StreamExt};
32use serde_json::{json, Value};
33use tokio::net::UnixStream;
34use tokio_tungstenite::tungstenite::Message;
35use tokio_tungstenite::WebSocketStream;
36
37use super::{Agent, ObserveRequest, SessionEvent, SessionState};
38
39/// How long one request may wait for its response.
40const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
41
42/// A connected, initialized app-server client.
43pub struct AppServerClient {
44    ws: WebSocketStream<UnixStream>,
45    next_id: u64,
46}
47
48impl AppServerClient {
49    /// Connects to the app-server listening on the Unix socket `path` and runs
50    /// the `initialize` handshake.
51    pub async fn connect(path: &Path) -> Result<Self> {
52        let stream = UnixStream::connect(path)
53            .await
54            .with_context(|| format!("failed to connect to {}", path.display()))?;
55        // The host is required by the handshake and ignored over a Unix socket.
56        let (ws, _) = tokio_tungstenite::client_async("ws://localhost/", stream)
57            .await
58            .context("app-server WebSocket handshake failed")?;
59        let mut client = Self { ws, next_id: 1 };
60        client
61            .request(
62                "initialize",
63                json!({
64                    "clientInfo": {
65                        "name": "omni-dev-codex-wrap",
66                        "version": env!("CARGO_PKG_VERSION"),
67                    },
68                }),
69            )
70            .await?;
71        client.notify("initialized").await?;
72        Ok(client)
73    }
74
75    /// Sends a notification (no response expected).
76    async fn notify(&mut self, method: &str) -> Result<()> {
77        let message = json!({ "jsonrpc": "2.0", "method": method });
78        self.ws
79            .send(Message::text(message.to_string()))
80            .await
81            .context("app-server send failed")
82    }
83
84    /// Sends one request and waits for its response, skipping notifications and
85    /// dropping any server request unanswered (see the module docs).
86    async fn request(&mut self, method: &str, params: Value) -> Result<Value> {
87        let id = self.next_id;
88        self.next_id += 1;
89        let message = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
90        self.ws
91            .send(Message::text(message.to_string()))
92            .await
93            .context("app-server send failed")?;
94        tokio::time::timeout(REQUEST_TIMEOUT, async {
95            while let Some(frame) = self.ws.next().await {
96                let text = match frame.context("app-server receive failed")? {
97                    Message::Text(text) => text,
98                    Message::Close(_) => break,
99                    _ => continue,
100                };
101                let Ok(reply) = serde_json::from_str::<Value>(text.as_str()) else {
102                    continue;
103                };
104                // A notification or a server request carries a `method`; neither
105                // is ours to handle.
106                if reply.get("method").is_some() || reply.get("id") != Some(&json!(id)) {
107                    continue;
108                }
109                if let Some(error) = reply.get("error") {
110                    bail!("app-server {method} failed: {error}");
111                }
112                return Ok(reply.get("result").cloned().unwrap_or(Value::Null));
113            }
114            Err(anyhow!("app-server closed the connection"))
115        })
116        .await
117        .map_err(|_| anyhow!("app-server {method} timed out"))?
118    }
119
120    /// The ids of the threads loaded in the server.
121    pub async fn loaded_threads(&mut self) -> Result<Vec<String>> {
122        let result = self.request("thread/loaded/list", json!({})).await?;
123        Ok(result
124            .get("data")
125            .and_then(Value::as_array)
126            .map(|ids| {
127                ids.iter()
128                    .filter_map(Value::as_str)
129                    .map(str::to_string)
130                    .collect()
131            })
132            .unwrap_or_default())
133    }
134
135    /// One thread's metadata and status, without its turns.
136    pub async fn read_thread(&mut self, id: &str) -> Result<Option<ThreadView>> {
137        let result = self
138            .request(
139                "thread/read",
140                json!({ "threadId": id, "includeTurns": false }),
141            )
142            .await?;
143        Ok(result.get("thread").and_then(ThreadView::from_json))
144    }
145
146    /// One poll of the server: the loaded thread ids, and the view of every
147    /// loaded top-level thread that could be read. A thread whose read fails
148    /// (it unloaded between the two calls, or the read timed out) is simply
149    /// missing from `views` this poll; only a failed `thread/loaded/list` is an
150    /// error.
151    pub async fn snapshot(&mut self) -> Result<Snapshot> {
152        let loaded = self.loaded_threads().await?;
153        let mut views = Vec::new();
154        for id in &loaded {
155            if let Ok(Some(view)) = self.read_thread(id).await {
156                if !view.subagent {
157                    views.push(view);
158                }
159            }
160        }
161        Ok(Snapshot { loaded, views })
162    }
163}
164
165/// One poll's result; see [`AppServerClient::snapshot`].
166#[derive(Debug, Clone, Default)]
167pub struct Snapshot {
168    /// Every thread id the server has loaded.
169    pub loaded: Vec<String>,
170    /// The readable, top-level ones among them.
171    pub views: Vec<ThreadView>,
172}
173
174/// The fields of a `Thread` the observer reads.
175#[derive(Debug, Clone, PartialEq, Eq)]
176pub struct ThreadView {
177    /// The thread id — the same `session_id` Codex's hooks report.
178    pub id: String,
179    /// The thread's working directory.
180    pub cwd: Option<PathBuf>,
181    /// The model id, when the server reports one.
182    pub model: Option<String>,
183    /// The mapped state, or `None` for a status that says nothing
184    /// (`notLoaded`, or one this build does not know).
185    pub state: Option<SessionState>,
186    /// Not a session of its own: a subagent (a parent id, or a `subAgent*`
187    /// source), or a system-spawned side thread such as the one Codex uses to
188    /// title a chat (`ephemeral`, `threadSource: "system"`). Not reported —
189    /// the user's thread already covers it.
190    pub subagent: bool,
191}
192
193impl ThreadView {
194    /// Reads a `Thread` object, tolerating missing optional fields. `None` when
195    /// it has no id.
196    fn from_json(thread: &Value) -> Option<Self> {
197        let id = thread.get("id")?.as_str()?.to_string();
198        let source = thread.get("source");
199        let subagent = thread.get("parentThreadId").is_some_and(|p| !p.is_null())
200            || thread.get("ephemeral").and_then(Value::as_bool) == Some(true)
201            || thread
202                .get("threadSource")
203                .and_then(Value::as_str)
204                .is_some_and(|s| s != "user")
205            || source.is_some_and(Value::is_object)
206            || source
207                .and_then(Value::as_str)
208                .is_some_and(|s| s.starts_with("subAgent"));
209        Some(Self {
210            id,
211            cwd: thread.get("cwd").and_then(Value::as_str).map(PathBuf::from),
212            model: thread
213                .get("model")
214                .and_then(Value::as_str)
215                .map(str::to_string),
216            state: thread.get("status").and_then(state_for_status),
217            subagent,
218        })
219    }
220}
221
222/// Maps a `ThreadStatus` onto a [`SessionState`]. `None` for `notLoaded` and
223/// any status this build does not know. `systemError` reads as idle: the turn
224/// is over and the session waits for the user.
225pub fn state_for_status(status: &Value) -> Option<SessionState> {
226    match status.get("type")?.as_str()? {
227        "idle" | "systemError" => Some(SessionState::Idle),
228        "active" => {
229            let flags: Vec<&str> = status
230                .get("activeFlags")
231                .and_then(Value::as_array)
232                .map(|flags| flags.iter().filter_map(Value::as_str).collect())
233                .unwrap_or_default();
234            Some(if flags.contains(&"waitingOnApproval") {
235                SessionState::WaitingForPermission
236            } else if flags.contains(&"waitingOnUserInput") {
237                SessionState::WaitingForInput
238            } else {
239                SessionState::Working
240            })
241        }
242        _ => None,
243    }
244}
245
246/// One thing the tracker wants the daemon to hear.
247#[derive(Debug, Clone)]
248pub enum Report {
249    /// An authoritative state for a session.
250    Observe(ObserveRequest),
251    /// The session is over.
252    End(String),
253}
254
255/// What the tracker last reported for a thread.
256#[derive(Debug, Clone)]
257struct Known {
258    state: SessionState,
259    cwd: Option<PathBuf>,
260    model: Option<String>,
261}
262
263/// Turns successive [`snapshot`](AppServerClient::snapshot)s into reports.
264///
265/// Every poll re-asserts each live session's state, not only a change: the
266/// same session also gets inferred reports from Codex's hooks and the rollout
267/// watcher, and re-asserting the exact state overrides any of them within one
268/// poll (a repeat sighting does not bump the registry, so this costs a socket
269/// round-trip, not a push). It also keeps an idle session inside the TTL. A
270/// thread that leaves the loaded set, or reports a status that says nothing,
271/// is ended; one merely unreadable this poll is left alone. Pure, so it is
272/// unit-tested without a server.
273#[derive(Debug, Default)]
274pub struct StatusTracker {
275    known: HashMap<String, Known>,
276}
277
278impl StatusTracker {
279    /// Folds in one poll.
280    pub fn update(&mut self, snapshot: &Snapshot) -> Vec<Report> {
281        let mut reports = Vec::new();
282        let mut silent = Vec::new();
283        for view in &snapshot.views {
284            let Some(state) = view.state else {
285                silent.push(view.id.clone());
286                continue;
287            };
288            let known = Known {
289                state,
290                cwd: view.cwd.clone(),
291                model: view.model.clone(),
292            };
293            reports.push(Report::Observe(observe(&view.id, &known)));
294            self.known.insert(view.id.clone(), known);
295        }
296        let gone: Vec<String> = self
297            .known
298            .keys()
299            .filter(|id| !snapshot.loaded.contains(id) || silent.contains(id))
300            .cloned()
301            .collect();
302        for id in gone {
303            self.known.remove(&id);
304            reports.push(Report::End(id));
305        }
306        reports
307    }
308
309    /// Ends every known session — the wrapper is exiting.
310    pub fn finish(&mut self) -> Vec<Report> {
311        self.known.drain().map(|(id, _)| Report::End(id)).collect()
312    }
313}
314
315/// The authoritative `observe` for a known thread.
316fn observe(id: &str, known: &Known) -> ObserveRequest {
317    ObserveRequest {
318        agent_id: None,
319        pid: None,
320        agent: Agent::Codex,
321        session_id: id.to_string(),
322        cwd: known.cwd.clone(),
323        transcript_path: None,
324        event: SessionEvent::StreamState(known.state),
325        repo: None,
326        model: known.model.clone(),
327    }
328}
329
330#[cfg(test)]
331#[allow(clippy::unwrap_used, clippy::expect_used)]
332mod tests {
333    use super::*;
334
335    fn view(id: &str, state: Option<SessionState>) -> ThreadView {
336        ThreadView {
337            id: id.to_string(),
338            cwd: Some(PathBuf::from("/work/repo")),
339            model: Some("gpt-5.5".to_string()),
340            state,
341            subagent: false,
342        }
343    }
344
345    fn summary(reports: &[Report]) -> Vec<String> {
346        reports
347            .iter()
348            .map(|r| match r {
349                Report::Observe(o) => format!("{}:{:?}", o.session_id, o.event),
350                Report::End(id) => format!("{id}:end"),
351            })
352            .collect()
353    }
354
355    #[test]
356    fn statuses_map_to_exact_states() {
357        let s = |v: Value| state_for_status(&v);
358        assert_eq!(s(json!({ "type": "idle" })), Some(SessionState::Idle));
359        assert_eq!(
360            s(json!({ "type": "systemError" })),
361            Some(SessionState::Idle)
362        );
363        assert_eq!(
364            s(json!({ "type": "active", "activeFlags": [] })),
365            Some(SessionState::Working)
366        );
367        assert_eq!(
368            s(json!({ "type": "active", "activeFlags": ["waitingOnApproval"] })),
369            Some(SessionState::WaitingForPermission)
370        );
371        assert_eq!(
372            s(json!({ "type": "active", "activeFlags": ["waitingOnUserInput"] })),
373            Some(SessionState::WaitingForInput)
374        );
375        // Approval outranks input when both are pending.
376        assert_eq!(
377            s(
378                json!({ "type": "active", "activeFlags": ["waitingOnUserInput", "waitingOnApproval"] })
379            ),
380            Some(SessionState::WaitingForPermission)
381        );
382        assert_eq!(s(json!({ "type": "notLoaded" })), None);
383        assert_eq!(s(json!({ "type": "somethingNew" })), None);
384        assert_eq!(s(json!({})), None);
385    }
386
387    #[test]
388    fn thread_views_read_the_fields_and_flag_subagents() {
389        let top = ThreadView::from_json(&json!({
390            "id": "t1",
391            "cwd": "/work/repo",
392            "model": "gpt-5.5",
393            "source": "vscode",
394            "threadSource": "user",
395            "ephemeral": false,
396            "status": { "type": "active", "activeFlags": ["waitingOnApproval"] },
397            "preview": "never read",
398        }))
399        .unwrap();
400        assert_eq!(top.id, "t1");
401        assert_eq!(top.cwd.as_deref(), Some(Path::new("/work/repo")));
402        assert_eq!(top.state, Some(SessionState::WaitingForPermission));
403        assert!(!top.subagent);
404        for sub in [
405            json!({ "id": "c", "parentThreadId": "t1", "status": { "type": "idle" } }),
406            json!({ "id": "c", "source": "subAgentThreadSpawn", "status": { "type": "idle" } }),
407            json!({ "id": "c", "source": { "subagent": "review" }, "status": { "type": "idle" } }),
408            // The chat-titling side thread observed on 0.155.1.
409            json!({ "id": "c", "ephemeral": true, "source": "vscode", "threadSource": "system",
410                    "status": { "type": "active", "activeFlags": [] } }),
411        ] {
412            assert!(ThreadView::from_json(&sub).unwrap().subagent, "{sub}");
413        }
414        assert!(ThreadView::from_json(&json!({ "cwd": "/x" })).is_none());
415    }
416
417    /// A snapshot where every view is loaded and readable.
418    fn snap(views: &[ThreadView]) -> Snapshot {
419        Snapshot {
420            loaded: views.iter().map(|v| v.id.clone()).collect(),
421            views: views.to_vec(),
422        }
423    }
424
425    #[test]
426    fn the_tracker_reasserts_every_poll_and_ends_unloaded_threads() {
427        let mut tracker = StatusTracker::default();
428        let idle = [view("t1", Some(SessionState::Idle))];
429        assert_eq!(
430            summary(&tracker.update(&snap(&idle))),
431            vec!["t1:StreamState(Idle)"]
432        );
433        // Unchanged: re-asserted, so a hook's inferred state is overridden
434        // within a poll.
435        assert_eq!(
436            summary(&tracker.update(&snap(&idle))),
437            vec!["t1:StreamState(Idle)"]
438        );
439        // A second thread (`/new`) joins; the first unloads.
440        assert_eq!(
441            summary(&tracker.update(&snap(&[view("t2", Some(SessionState::Working))]))),
442            vec!["t2:StreamState(Working)", "t1:end"]
443        );
444        // A status that says nothing counts as gone.
445        assert_eq!(
446            summary(&tracker.update(&snap(&[view("t2", None)]))),
447            vec!["t2:end"]
448        );
449        assert!(tracker.finish().is_empty());
450    }
451
452    #[test]
453    fn an_unreadable_but_loaded_thread_is_not_ended() {
454        let mut tracker = StatusTracker::default();
455        tracker.update(&snap(&[view("t1", Some(SessionState::Working))]));
456        // `thread/read` failed for t1 this poll, but it is still loaded.
457        let partial = Snapshot {
458            loaded: vec!["t1".to_string()],
459            views: Vec::new(),
460        };
461        assert!(tracker.update(&partial).is_empty());
462        let request = match &tracker.update(&snap(&[view("t1", Some(SessionState::Idle))]))[0] {
463            Report::Observe(request) => request.clone(),
464            Report::End(_) => panic!("expected an observe"),
465        };
466        assert_eq!(request.agent, Agent::Codex);
467        assert_eq!(request.model.as_deref(), Some("gpt-5.5"));
468        assert_eq!(summary(&tracker.finish()), vec!["t1:end"]);
469        assert!(tracker.finish().is_empty());
470    }
471
472    /// A scripted app-server on a Unix socket: answers `initialize`,
473    /// `thread/loaded/list` and `thread/read`, and interleaves a notification
474    /// and a server request (which the client must skip and never answer).
475    async fn fake_server(path: PathBuf, answered: std::sync::Arc<std::sync::Mutex<Vec<Value>>>) {
476        let listener = tokio::net::UnixListener::bind(&path).unwrap();
477        let (stream, _) = listener.accept().await.unwrap();
478        let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
479        while let Some(Ok(Message::Text(text))) = ws.next().await {
480            let msg: Value = serde_json::from_str(text.as_str()).unwrap();
481            answered.lock().unwrap().push(msg.clone());
482            let Some(id) = msg.get("id").cloned() else {
483                continue;
484            };
485            let result = match msg["method"].as_str().unwrap() {
486                "initialize" => json!({ "userAgent": "fake" }),
487                "thread/loaded/list" => json!({ "data": ["t1", "child"] }),
488                "thread/read" if msg["params"]["threadId"] == "t1" => json!({ "thread": {
489                    "id": "t1", "cwd": "/w", "source": "cli",
490                    "status": { "type": "active", "activeFlags": ["waitingOnApproval"] },
491                } }),
492                "thread/read" => json!({ "thread": {
493                    "id": "child", "parentThreadId": "t1", "status": { "type": "idle" },
494                } }),
495                other => panic!("unexpected {other}"),
496            };
497            let noise = [
498                json!({ "method": "thread/status/changed", "params": {} }),
499                json!({ "id": 900, "method": "item/commandExecution/requestApproval", "params": {} }),
500                json!({ "id": id, "result": result }),
501            ];
502            for frame in noise {
503                ws.send(Message::text(frame.to_string())).await.unwrap();
504            }
505        }
506    }
507
508    #[tokio::test]
509    async fn the_client_polls_and_never_answers_a_server_request() {
510        let tmp = tempfile::tempdir_in("/tmp").unwrap();
511        let path = tmp.path().join("as.sock");
512        let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
513        let server = tokio::spawn(fake_server(path.clone(), seen.clone()));
514        while !path.exists() {
515            tokio::time::sleep(Duration::from_millis(5)).await;
516        }
517        let mut client = AppServerClient::connect(&path).await.unwrap();
518        let snapshot = client.snapshot().await.unwrap();
519        assert_eq!(snapshot.loaded, ["t1", "child"]);
520        let views = snapshot.views;
521        assert_eq!(views.len(), 1, "the subagent thread is dropped");
522        assert_eq!(views[0].state, Some(SessionState::WaitingForPermission));
523        drop(client);
524        server.await.unwrap();
525        let methods: Vec<String> = seen
526            .lock()
527            .unwrap()
528            .iter()
529            .map(|m| m["method"].as_str().unwrap_or("<response>").to_string())
530            .collect();
531        assert_eq!(
532            methods,
533            [
534                "initialize",
535                "initialized",
536                "thread/loaded/list",
537                "thread/read",
538                "thread/read"
539            ],
540            "only polls were sent — no response to the server's request"
541        );
542    }
543}
omni_dev/sessions/codex_app_server.rs

omni_dev/sessions/
codex_app_server.rs
