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        pid: None,
319        agent: Agent::Codex,
320        session_id: id.to_string(),
321        cwd: known.cwd.clone(),
322        transcript_path: None,
323        event: SessionEvent::StreamState(known.state),
324        repo: None,
325        model: known.model.clone(),
326    }
327}
328
329#[cfg(test)]
330#[allow(clippy::unwrap_used, clippy::expect_used)]
331mod tests {
332    use super::*;
333
334    fn view(id: &str, state: Option<SessionState>) -> ThreadView {
335        ThreadView {
336            id: id.to_string(),
337            cwd: Some(PathBuf::from("/work/repo")),
338            model: Some("gpt-5.5".to_string()),
339            state,
340            subagent: false,
341        }
342    }
343
344    fn summary(reports: &[Report]) -> Vec<String> {
345        reports
346            .iter()
347            .map(|r| match r {
348                Report::Observe(o) => format!("{}:{:?}", o.session_id, o.event),
349                Report::End(id) => format!("{id}:end"),
350            })
351            .collect()
352    }
353
354    #[test]
355    fn statuses_map_to_exact_states() {
356        let s = |v: Value| state_for_status(&v);
357        assert_eq!(s(json!({ "type": "idle" })), Some(SessionState::Idle));
358        assert_eq!(
359            s(json!({ "type": "systemError" })),
360            Some(SessionState::Idle)
361        );
362        assert_eq!(
363            s(json!({ "type": "active", "activeFlags": [] })),
364            Some(SessionState::Working)
365        );
366        assert_eq!(
367            s(json!({ "type": "active", "activeFlags": ["waitingOnApproval"] })),
368            Some(SessionState::WaitingForPermission)
369        );
370        assert_eq!(
371            s(json!({ "type": "active", "activeFlags": ["waitingOnUserInput"] })),
372            Some(SessionState::WaitingForInput)
373        );
374        // Approval outranks input when both are pending.
375        assert_eq!(
376            s(
377                json!({ "type": "active", "activeFlags": ["waitingOnUserInput", "waitingOnApproval"] })
378            ),
379            Some(SessionState::WaitingForPermission)
380        );
381        assert_eq!(s(json!({ "type": "notLoaded" })), None);
382        assert_eq!(s(json!({ "type": "somethingNew" })), None);
383        assert_eq!(s(json!({})), None);
384    }
385
386    #[test]
387    fn thread_views_read_the_fields_and_flag_subagents() {
388        let top = ThreadView::from_json(&json!({
389            "id": "t1",
390            "cwd": "/work/repo",
391            "model": "gpt-5.5",
392            "source": "vscode",
393            "threadSource": "user",
394            "ephemeral": false,
395            "status": { "type": "active", "activeFlags": ["waitingOnApproval"] },
396            "preview": "never read",
397        }))
398        .unwrap();
399        assert_eq!(top.id, "t1");
400        assert_eq!(top.cwd.as_deref(), Some(Path::new("/work/repo")));
401        assert_eq!(top.state, Some(SessionState::WaitingForPermission));
402        assert!(!top.subagent);
403        for sub in [
404            json!({ "id": "c", "parentThreadId": "t1", "status": { "type": "idle" } }),
405            json!({ "id": "c", "source": "subAgentThreadSpawn", "status": { "type": "idle" } }),
406            json!({ "id": "c", "source": { "subagent": "review" }, "status": { "type": "idle" } }),
407            // The chat-titling side thread observed on 0.155.1.
408            json!({ "id": "c", "ephemeral": true, "source": "vscode", "threadSource": "system",
409                    "status": { "type": "active", "activeFlags": [] } }),
410        ] {
411            assert!(ThreadView::from_json(&sub).unwrap().subagent, "{sub}");
412        }
413        assert!(ThreadView::from_json(&json!({ "cwd": "/x" })).is_none());
414    }
415
416    /// A snapshot where every view is loaded and readable.
417    fn snap(views: &[ThreadView]) -> Snapshot {
418        Snapshot {
419            loaded: views.iter().map(|v| v.id.clone()).collect(),
420            views: views.to_vec(),
421        }
422    }
423
424    #[test]
425    fn the_tracker_reasserts_every_poll_and_ends_unloaded_threads() {
426        let mut tracker = StatusTracker::default();
427        let idle = [view("t1", Some(SessionState::Idle))];
428        assert_eq!(
429            summary(&tracker.update(&snap(&idle))),
430            vec!["t1:StreamState(Idle)"]
431        );
432        // Unchanged: re-asserted, so a hook's inferred state is overridden
433        // within a poll.
434        assert_eq!(
435            summary(&tracker.update(&snap(&idle))),
436            vec!["t1:StreamState(Idle)"]
437        );
438        // A second thread (`/new`) joins; the first unloads.
439        assert_eq!(
440            summary(&tracker.update(&snap(&[view("t2", Some(SessionState::Working))]))),
441            vec!["t2:StreamState(Working)", "t1:end"]
442        );
443        // A status that says nothing counts as gone.
444        assert_eq!(
445            summary(&tracker.update(&snap(&[view("t2", None)]))),
446            vec!["t2:end"]
447        );
448        assert!(tracker.finish().is_empty());
449    }
450
451    #[test]
452    fn an_unreadable_but_loaded_thread_is_not_ended() {
453        let mut tracker = StatusTracker::default();
454        tracker.update(&snap(&[view("t1", Some(SessionState::Working))]));
455        // `thread/read` failed for t1 this poll, but it is still loaded.
456        let partial = Snapshot {
457            loaded: vec!["t1".to_string()],
458            views: Vec::new(),
459        };
460        assert!(tracker.update(&partial).is_empty());
461        let request = match &tracker.update(&snap(&[view("t1", Some(SessionState::Idle))]))[0] {
462            Report::Observe(request) => request.clone(),
463            Report::End(_) => panic!("expected an observe"),
464        };
465        assert_eq!(request.agent, Agent::Codex);
466        assert_eq!(request.model.as_deref(), Some("gpt-5.5"));
467        assert_eq!(summary(&tracker.finish()), vec!["t1:end"]);
468        assert!(tracker.finish().is_empty());
469    }
470
471    /// A scripted app-server on a Unix socket: answers `initialize`,
472    /// `thread/loaded/list` and `thread/read`, and interleaves a notification
473    /// and a server request (which the client must skip and never answer).
474    async fn fake_server(path: PathBuf, answered: std::sync::Arc<std::sync::Mutex<Vec<Value>>>) {
475        let listener = tokio::net::UnixListener::bind(&path).unwrap();
476        let (stream, _) = listener.accept().await.unwrap();
477        let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
478        while let Some(Ok(Message::Text(text))) = ws.next().await {
479            let msg: Value = serde_json::from_str(text.as_str()).unwrap();
480            answered.lock().unwrap().push(msg.clone());
481            let Some(id) = msg.get("id").cloned() else {
482                continue;
483            };
484            let result = match msg["method"].as_str().unwrap() {
485                "initialize" => json!({ "userAgent": "fake" }),
486                "thread/loaded/list" => json!({ "data": ["t1", "child"] }),
487                "thread/read" if msg["params"]["threadId"] == "t1" => json!({ "thread": {
488                    "id": "t1", "cwd": "/w", "source": "cli",
489                    "status": { "type": "active", "activeFlags": ["waitingOnApproval"] },
490                } }),
491                "thread/read" => json!({ "thread": {
492                    "id": "child", "parentThreadId": "t1", "status": { "type": "idle" },
493                } }),
494                other => panic!("unexpected {other}"),
495            };
496            let noise = [
497                json!({ "method": "thread/status/changed", "params": {} }),
498                json!({ "id": 900, "method": "item/commandExecution/requestApproval", "params": {} }),
499                json!({ "id": id, "result": result }),
500            ];
501            for frame in noise {
502                ws.send(Message::text(frame.to_string())).await.unwrap();
503            }
504        }
505    }
506
507    #[tokio::test]
508    async fn the_client_polls_and_never_answers_a_server_request() {
509        let tmp = tempfile::tempdir_in("/tmp").unwrap();
510        let path = tmp.path().join("as.sock");
511        let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
512        let server = tokio::spawn(fake_server(path.clone(), seen.clone()));
513        while !path.exists() {
514            tokio::time::sleep(Duration::from_millis(5)).await;
515        }
516        let mut client = AppServerClient::connect(&path).await.unwrap();
517        let snapshot = client.snapshot().await.unwrap();
518        assert_eq!(snapshot.loaded, ["t1", "child"]);
519        let views = snapshot.views;
520        assert_eq!(views.len(), 1, "the subagent thread is dropped");
521        assert_eq!(views[0].state, Some(SessionState::WaitingForPermission));
522        drop(client);
523        server.await.unwrap();
524        let methods: Vec<String> = seen
525            .lock()
526            .unwrap()
527            .iter()
528            .map(|m| m["method"].as_str().unwrap_or("<response>").to_string())
529            .collect();
530        assert_eq!(
531            methods,
532            [
533                "initialize",
534                "initialized",
535                "thread/loaded/list",
536                "thread/read",
537                "thread/read"
538            ],
539            "only polls were sent — no response to the server's request"
540        );
541    }
542}
omni_dev/sessions/codex_app_server.rs

omni_dev/sessions/
codex_app_server.rs
