Skip to main content

codex_app_server_client/
transport.rs

1//! Wire transport: the app-server's stdio mode is newline-delimited JSON
2//! (JSONL) with the bare JSON-RPC 2.0 message shape and the `"jsonrpc":"2.0"`
3//! header omitted (see <https://developers.openai.com/codex/app-server>).
4//! This module only frames/deframes lines; message typing lives in
5//! [`crate::protocol`] and dispatch lives in [`crate::client`].
6
7use std::process::Stdio;
8
9use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
10use tokio::process::{Child, ChildStdin, ChildStdout, Command};
11
12use crate::protocol::RequestId;
13use crate::{Error, Result};
14
15/// Hard cap on a single NDJSON line's size. Without this, a single huge or
16/// unterminated line from a buggy or malicious app-server peer would grow
17/// `read_line`'s buffer without bound. 64 MiB comfortably covers legitimate
18/// large payloads (e.g. `fs/readFile` results, big diffs) while still being a
19/// real, finite bound.
20pub const MAX_LINE_BYTES: usize = 64 * 1024 * 1024;
21
22/// A reply to a server->client request, queued for the writer task.
23pub(crate) enum OutgoingReply {
24    Result {
25        id: RequestId,
26        result: serde_json::Value,
27    },
28    Error {
29        id: RequestId,
30        code: i64,
31        message: String,
32        data: Option<serde_json::Value>,
33    },
34}
35
36impl OutgoingReply {
37    pub(crate) fn into_line(self) -> Result<String> {
38        let value = match self {
39            OutgoingReply::Result { id, result } => {
40                serde_json::json!({ "id": id, "result": result })
41            }
42            OutgoingReply::Error {
43                id,
44                code,
45                message,
46                data,
47            } => {
48                serde_json::json!({ "id": id, "error": { "code": code, "message": message, "data": data } })
49            }
50        };
51        Ok(serde_json::to_string(&value)?)
52    }
53}
54
55/// Spawns `command app-server [extra_args...]` with stdio piped, ready for
56/// [`crate::CodexAppServerClient::connect`].
57pub(crate) fn spawn_app_server(
58    command: &str,
59    extra_args: &[String],
60) -> Result<(ChildStdin, BufReader<ChildStdout>, Child)> {
61    let mut cmd = Command::new(command);
62    cmd.arg("app-server");
63    cmd.args(extra_args);
64    cmd.stdin(Stdio::piped());
65    cmd.stdout(Stdio::piped());
66    cmd.stderr(Stdio::inherit());
67    cmd.kill_on_drop(true);
68
69    let mut child = cmd.spawn().map_err(|source| Error::Spawn {
70        command: command.to_string(),
71        source,
72    })?;
73
74    let stdin = child.stdin.take().expect("piped stdin");
75    let stdout = child.stdout.take().expect("piped stdout");
76    Ok((stdin, BufReader::new(stdout), child))
77}
78
79pub(crate) async fn write_line<W>(writer: &mut W, line: &str) -> std::io::Result<()>
80where
81    W: tokio::io::AsyncWrite + Unpin,
82{
83    writer.write_all(line.as_bytes()).await?;
84    writer.write_all(b"\n").await?;
85    writer.flush().await
86}
87
88/// Reads one `\n`-terminated line into `buf` (cleared first, but its
89/// allocation is reused rather than replaced - callers keep one persistent
90/// `buf` across the whole read loop so this stays allocation-free after the
91/// first few calls), enforcing [`MAX_LINE_BYTES`]. Returns the number of
92/// bytes read (0 on clean EOF with nothing left to read), or an error if the
93/// line is invalid UTF-8 or would exceed the cap - in either case the caller
94/// should treat the connection as dead rather than try to resynchronize
95/// mid-line.
96pub(crate) async fn read_line<R>(reader: &mut R, buf: &mut String) -> std::io::Result<usize>
97where
98    R: tokio::io::AsyncBufRead + Unpin,
99{
100    read_line_capped(reader, buf, MAX_LINE_BYTES).await
101}
102
103/// [`read_line`]'s implementation, parameterized over the cap so tests can
104/// exercise the boundary condition without allocating [`MAX_LINE_BYTES`]
105/// worth of memory.
106async fn read_line_capped<R>(
107    reader: &mut R,
108    buf: &mut String,
109    max_bytes: usize,
110) -> std::io::Result<usize>
111where
112    R: tokio::io::AsyncBufRead + Unpin,
113{
114    let mut bytes = std::mem::take(buf).into_bytes();
115    bytes.clear(); // drops content, keeps the allocated capacity
116    loop {
117        let available = reader.fill_buf().await?;
118        if available.is_empty() {
119            break; // EOF
120        }
121        if let Some(pos) = available.iter().position(|&b| b == b'\n') {
122            // The cap must be enforced here too: `bytes` can already be close
123            // to max_bytes from prior chunks (each individually under the
124            // cap), and this chunk - up to one BufReader-internal-buffer's
125            // worth of bytes - could push the *line-terminated* total over it
126            // even though a newline was found. Without this check the cap
127            // could be overshot by up to one buffer's worth per line.
128            if bytes.len() + pos + 1 > max_bytes {
129                reader.consume(pos + 1);
130                return Err(std::io::Error::new(
131                    std::io::ErrorKind::InvalidData,
132                    format!("NDJSON line exceeded the {max_bytes}-byte cap"),
133                ));
134            }
135            bytes.extend_from_slice(&available[..=pos]);
136            reader.consume(pos + 1);
137            break;
138        }
139        let n = available.len();
140        if bytes.len() + n > max_bytes {
141            reader.consume(n);
142            return Err(std::io::Error::new(
143                std::io::ErrorKind::InvalidData,
144                format!("NDJSON line exceeded the {max_bytes}-byte cap"),
145            ));
146        }
147        bytes.extend_from_slice(available);
148        reader.consume(n);
149    }
150    if bytes.is_empty() {
151        return Ok(0);
152    }
153    *buf = String::from_utf8(bytes)
154        .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?;
155    Ok(buf.len())
156}
157
158// Re-export the trait bounds so callers of `CodexAppServerClient::connect` don't
159// need to import tokio themselves for common cases.
160pub use tokio::io::{AsyncBufRead, AsyncWrite};
161
162#[cfg(unix)]
163pub(crate) fn split_unix_stream(
164    stream: tokio::net::UnixStream,
165) -> (
166    tokio::net::unix::OwnedWriteHalf,
167    BufReader<tokio::net::unix::OwnedReadHalf>,
168) {
169    let (read_half, write_half) = stream.into_split();
170    (write_half, BufReader::new(read_half))
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use tokio::io::BufReader;
177
178    /// Finding #3 regression: a chunk that both crosses `max_bytes` *and*
179    /// contains the terminating newline must still be rejected. Before the
180    /// fix, only the no-newline-found branch checked the cap, so a
181    /// newline-terminated final chunk could overshoot it by up to one
182    /// `fill_buf` chunk's worth of bytes.
183    #[tokio::test]
184    async fn read_line_capped_rejects_an_oversized_line_even_when_newline_terminated() {
185        let mut reader = BufReader::new(b"0123456789\n".as_slice());
186        let mut buf = String::new();
187
188        let err = read_line_capped(&mut reader, &mut buf, 5)
189            .await
190            .expect_err("an 11-byte line must be rejected under a 5-byte cap");
191
192        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
193        assert!(err.to_string().contains("5-byte cap"), "{err}");
194    }
195
196    /// A newline-terminated line at exactly the cap is still accepted (the
197    /// check is `>`, not `>=`).
198    #[tokio::test]
199    async fn read_line_capped_accepts_a_line_exactly_at_the_cap() {
200        let mut reader = BufReader::new(b"01234\n".as_slice());
201        let mut buf = String::new();
202
203        let n = read_line_capped(&mut reader, &mut buf, 6)
204            .await
205            .expect("a 6-byte line (including the newline) fits a 6-byte cap");
206
207        assert_eq!(n, 6);
208        assert_eq!(buf, "01234\n");
209    }
210}