ai_crew_sync/
codex.rs

1//! A minimal client for the Codex app-server (#222).
2//!
3//! Codex runs a local app-server daemon (`codex app-server daemon`) that owns
4//! the user's threads and exposes them on a Unix socket. The protocol is
5//! JSON-RPC 2.0, one message per WebSocket text frame. This module speaks
6//! exactly what the wake adapter needs: `initialize`, `thread/read`,
7//! `thread/resume`, `turn/start`, and the notifications that follow. It is
8//! a client only, and only over a local Unix socket the user owns: no TCP,
9//! no TLS, no authentication of its own.
10//!
11//! The WebSocket framing is written here rather than pulled in as a crate:
12//! the adapter needs text frames, ping, pong and close over one local
13//! socket, which is a few dozen lines, and a new dependency needs approval
14//! the feature does not justify.
15
16use std::{
17    collections::VecDeque,
18    path::{Path, PathBuf},
19    time::Duration,
20};
21
22use anyhow::{Context, bail};
23use base64::Engine;
24use serde_json::{Value, json};
25use tokio::{
26    io::{AsyncReadExt, AsyncWriteExt},
27    net::UnixStream,
28};
29
30/// Largest frame accepted from the app-server. Thread reads can be large
31/// when a caller asks for turns; the adapter never does, so this is a bound
32/// against a misbehaving peer, not a limit the protocol imposes.
33const MAX_FRAME_BYTES: u64 = 64 * 1024 * 1024;
34/// How long one request may take before the adapter gives up on it.
35const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
36/// Connecting and upgrading are local and quick; a daemon that accepts and
37/// then stalls must not hang the adapter.
38const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
39
40/// Where the Codex app-server daemon listens by default: under
41/// `$CODEX_HOME` (else `~/.codex`), as `codex app-server daemon version`
42/// reports in `socketPath`.
43pub fn default_socket() -> anyhow::Result<PathBuf> {
44    let home = match std::env::var_os("CODEX_HOME") {
45        Some(h) if !h.is_empty() => PathBuf::from(h),
46        _ => crate::admin_cli::home_dir()
47            .context("neither CODEX_HOME nor HOME is set")?
48            .join(".codex"),
49    };
50    Ok(home
51        .join("app-server-control")
52        .join("app-server-control.sock"))
53}
54
55/// A connected, initialised app-server client.
56pub struct AppServer {
57    stream: UnixStream,
58    /// Bytes read past the end of the last frame.
59    buf: Vec<u8>,
60    /// Fragments of a message whose final frame has not arrived yet. Kept
61    /// here rather than on the stack so a read cancelled by a timeout loses
62    /// nothing: the next read continues the same message.
63    partial: Vec<u8>,
64    next_id: i64,
65    /// Notifications that arrived while waiting for a response, oldest first.
66    pending: VecDeque<Value>,
67}
68
69/// A JSON-RPC error the app-server returned, kept apart from transport
70/// failures so the adapter can tell "Codex refused" from "Codex is gone".
71#[derive(Debug, thiserror::Error)]
72#[error("{method}: {message}")]
73pub struct RpcError {
74    pub method: String,
75    pub code: i64,
76    pub message: String,
77}
78
79impl AppServer {
80    /// Connect to the daemon's socket, upgrade to WebSocket, and initialise
81    /// as `client_name`.
82    pub async fn connect(socket: &Path, client_name: &str) -> anyhow::Result<Self> {
83        tokio::time::timeout(CONNECT_TIMEOUT, Self::open(socket, client_name))
84            .await
85            .with_context(|| {
86                format!(
87                    "the Codex app-server at {} did not answer within {}s",
88                    socket.display(),
89                    CONNECT_TIMEOUT.as_secs()
90                )
91            })?
92    }
93
94    async fn open(socket: &Path, client_name: &str) -> anyhow::Result<Self> {
95        let stream = UnixStream::connect(socket).await.with_context(|| {
96            format!(
97                "could not reach the Codex app-server at {}; is it running? \
98                 (`codex app-server daemon start`)",
99                socket.display()
100            )
101        })?;
102        let mut client = Self {
103            stream,
104            buf: Vec::new(),
105            partial: Vec::new(),
106            next_id: 1,
107            pending: VecDeque::new(),
108        };
109        client.handshake().await?;
110        client
111            .request(
112                "initialize",
113                json!({"clientInfo": {"name": client_name, "version": env!("CARGO_PKG_VERSION")}}),
114            )
115            .await?;
116        client
117            .send_json(&json!({"jsonrpc": "2.0", "method": "initialized"}))
118            .await?;
119        Ok(client)
120    }
121
122    async fn handshake(&mut self) -> anyhow::Result<()> {
123        let key = base64::engine::general_purpose::STANDARD.encode(rand::random::<[u8; 16]>());
124        let request = format!(
125            "GET / HTTP/1.1\r\nHost: localhost\r\nUpgrade: websocket\r\n\
126             Connection: Upgrade\r\nSec-WebSocket-Key: {key}\r\n\
127             Sec-WebSocket-Version: 13\r\n\r\n"
128        );
129        self.stream.write_all(request.as_bytes()).await?;
130        let mut head = Vec::new();
131        let mut byte = [0u8; 1];
132        // The response head is short; read it byte by byte so no frame byte
133        // after it is consumed by mistake.
134        while !head.ends_with(b"\r\n\r\n") {
135            if head.len() > 16 * 1024 {
136                bail!("the Codex app-server sent an oversized handshake response");
137            }
138            let n = self.stream.read(&mut byte).await?;
139            if n == 0 {
140                bail!("the Codex app-server closed the connection during the handshake");
141            }
142            head.push(byte[0]);
143        }
144        let status = String::from_utf8_lossy(&head);
145        if !status.starts_with("HTTP/1.1 101") {
146            bail!(
147                "the Codex app-server refused the WebSocket upgrade: {}",
148                status.lines().next().unwrap_or("")
149            );
150        }
151        Ok(())
152    }
153
154    /// Send one JSON-RPC message as a masked text frame (clients mask).
155    async fn send_json(&mut self, value: &Value) -> anyhow::Result<()> {
156        let payload = serde_json::to_vec(value)?;
157        self.send_frame(0x1, &payload).await
158    }
159
160    async fn send_frame(&mut self, opcode: u8, payload: &[u8]) -> anyhow::Result<()> {
161        let mut frame = Vec::with_capacity(payload.len() + 14);
162        frame.push(0x80 | opcode);
163        let len = payload.len();
164        if len < 126 {
165            frame.push(0x80 | len as u8);
166        } else if len <= u16::MAX as usize {
167            frame.push(0x80 | 126);
168            frame.extend_from_slice(&(len as u16).to_be_bytes());
169        } else {
170            frame.push(0x80 | 127);
171            frame.extend_from_slice(&(len as u64).to_be_bytes());
172        }
173        let mask: [u8; 4] = rand::random();
174        frame.extend_from_slice(&mask);
175        frame.extend(payload.iter().enumerate().map(|(i, b)| b ^ mask[i % 4]));
176        self.stream.write_all(&frame).await?;
177        Ok(())
178    }
179
180    async fn fill(&mut self, want: usize) -> anyhow::Result<()> {
181        let mut chunk = [0u8; 16 * 1024];
182        while self.buf.len() < want {
183            let n = self.stream.read(&mut chunk).await?;
184            if n == 0 {
185                bail!("the Codex app-server closed the connection");
186            }
187            self.buf.extend_from_slice(&chunk[..n]);
188        }
189        Ok(())
190    }
191
192    /// Read the next complete message: text frames (reassembled across
193    /// continuations) are parsed as JSON; pings are answered; a close ends
194    /// the connection with an error.
195    async fn read_message(&mut self) -> anyhow::Result<Value> {
196        loop {
197            self.fill(2).await?;
198            let fin = self.buf[0] & 0x80 != 0;
199            let opcode = self.buf[0] & 0x0f;
200            let masked = self.buf[1] & 0x80 != 0;
201            let mut len = u64::from(self.buf[1] & 0x7f);
202            let mut offset = 2usize;
203            if len == 126 {
204                self.fill(4).await?;
205                len = u64::from(u16::from_be_bytes([self.buf[2], self.buf[3]]));
206                offset = 4;
207            } else if len == 127 {
208                self.fill(10).await?;
209                let mut b = [0u8; 8];
210                b.copy_from_slice(&self.buf[2..10]);
211                len = u64::from_be_bytes(b);
212                offset = 10;
213            }
214            if len > MAX_FRAME_BYTES {
215                bail!("the Codex app-server sent a {len}-byte frame; refusing it");
216            }
217            let mask_len = if masked { 4 } else { 0 };
218            let total = offset + mask_len + len as usize;
219            self.fill(total).await?;
220            let mut payload = self.buf[offset + mask_len..total].to_vec();
221            if masked {
222                let mask = [
223                    self.buf[offset],
224                    self.buf[offset + 1],
225                    self.buf[offset + 2],
226                    self.buf[offset + 3],
227                ];
228                for (i, b) in payload.iter_mut().enumerate() {
229                    *b ^= mask[i % 4];
230                }
231            }
232            self.buf.drain(..total);
233            match opcode {
234                0x0 | 0x1 => {
235                    if self.partial.len() as u64 + payload.len() as u64 > MAX_FRAME_BYTES {
236                        bail!("the Codex app-server sent an oversized message; refusing it");
237                    }
238                    self.partial.extend_from_slice(&payload);
239                    if fin {
240                        let message = std::mem::take(&mut self.partial);
241                        return serde_json::from_slice(&message)
242                            .context("the Codex app-server sent a message that is not JSON");
243                    }
244                }
245                0x8 => bail!("the Codex app-server closed the connection"),
246                0x9 => self.send_frame(0xA, &payload).await?,
247                // Pong, binary or reserved: nothing the adapter uses.
248                _ => {}
249            }
250        }
251    }
252
253    /// Send a request and wait for its response. Notifications received in
254    /// the meantime are queued for [`Self::next_notification`]. A request
255    /// the server sends us (an approval prompt) is answered with an error:
256    /// the adapter has no authority to approve anything.
257    pub async fn request(&mut self, method: &str, params: Value) -> anyhow::Result<Value> {
258        let id = self.next_id;
259        self.next_id += 1;
260        self.send_json(&json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}))
261            .await?;
262        let deadline = tokio::time::Instant::now() + REQUEST_TIMEOUT;
263        loop {
264            let message = tokio::time::timeout_at(deadline, self.read_message())
265                .await
266                .with_context(|| format!("{method}: no answer from the Codex app-server"))??;
267            if message.get("method").is_some() {
268                self.route_incoming(message).await?;
269                continue;
270            }
271            if message.get("id").and_then(Value::as_i64) != Some(id) {
272                continue;
273            }
274            if let Some(error) = message.get("error") {
275                return Err(RpcError {
276                    method: method.to_owned(),
277                    code: error.get("code").and_then(Value::as_i64).unwrap_or(0),
278                    message: error
279                        .get("message")
280                        .and_then(Value::as_str)
281                        .unwrap_or("error")
282                        .to_owned(),
283                }
284                .into());
285            }
286            return Ok(message.get("result").cloned().unwrap_or(Value::Null));
287        }
288    }
289
290    /// A notification is queued; a server request is declined.
291    async fn route_incoming(&mut self, message: Value) -> anyhow::Result<()> {
292        match message.get("id") {
293            Some(id) => {
294                let id = id.clone();
295                self.send_json(&json!({
296                    "jsonrpc": "2.0",
297                    "id": id,
298                    "error": {"code": -32601, "message": "the ai-crew-sync wake adapter approves nothing"},
299                }))
300                .await
301            }
302            None => {
303                self.pending.push_back(message);
304                Ok(())
305            }
306        }
307    }
308
309    /// The next notification, waiting up to `timeout`. `None` on timeout.
310    pub async fn next_notification(&mut self, timeout: Duration) -> anyhow::Result<Option<Value>> {
311        if let Some(n) = self.pending.pop_front() {
312            return Ok(Some(n));
313        }
314        let deadline = tokio::time::Instant::now() + timeout;
315        loop {
316            let message = match tokio::time::timeout_at(deadline, self.read_message()).await {
317                Ok(m) => m?,
318                Err(_) => return Ok(None),
319            };
320            if message.get("method").is_none() {
321                continue;
322            }
323            if message.get("id").is_some() {
324                self.route_incoming(message).await?;
325                continue;
326            }
327            return Ok(Some(message));
328        }
329    }
330}
331
332/// A thread's status as `thread/read` and `thread/status/changed` report it:
333/// `notLoaded`, `idle`, `active` or `systemError`.
334pub fn thread_status(thread_or_notification: &Value) -> Option<&str> {
335    thread_or_notification
336        .get("status")
337        .and_then(|s| s.get("type"))
338        .and_then(Value::as_str)
339}
ai_crew_sync/codex.rs

ai_crew_sync/
codex.rs
