codex_codes/
client_async.rs

1//! Asynchronous multi-turn client for the Codex app-server.
2//!
3//! Spawns `codex app-server --listen stdio://` and communicates over
4//! newline-delimited JSON-RPC. The connection stays open for multiple
5//! turns until explicitly shut down.
6//!
7//! # Lifecycle
8//!
9//! 1. Create a client with [`AsyncClient::start`] (spawns and initializes the app-server)
10//! 2. Call [`AsyncClient::thread_start`] to create a conversation session
11//! 3. Call [`AsyncClient::turn_start`] to send user input
12//! 4. Consume [`AsyncClient::next_message`] to stream notifications
13//! 5. Handle approval requests via [`AsyncClient::respond`]
14//! 6. Repeat steps 3-5 for follow-up turns
15//! 7. The client kills the app-server on [`Drop`]
16//!
17//! # Example
18//!
19//! ```ignore
20//! use codex_codes::{AsyncClient, ThreadStartParams, TurnStartParams, UserInput, ServerMessage};
21//!
22//! let mut client = AsyncClient::start().await?;
23//! let thread = client.thread_start(&ThreadStartParams::default()).await?;
24//!
25//! client.turn_start(&TurnStartParams {
26//!     thread_id: thread.thread_id().to_string(),
27//!     input: vec![UserInput::Text { text: "Hello!".into() }],
28//!     model: None,
29//!     reasoning_effort: None,
30//!     sandbox_policy: None,
31//! }).await?;
32//!
33//! while let Some(msg) = client.next_message().await? {
34//!     match msg {
35//!         ServerMessage::Notification(n) => {
36//!             if let codex_codes::Notification::TurnCompleted(_) = n { break; }
37//!         }
38//!         ServerMessage::Request { id, .. } => {
39//!             client.respond(id, &serde_json::json!({"decision": "accept"})).await?;
40//!         }
41//!     }
42//! }
43//! ```
44
45use crate::cli::AppServerBuilder;
46use crate::error::{Error, ParseError, Result};
47use crate::jsonrpc::{
48    JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, RequestId,
49};
50use crate::messages::{Notification, ServerMessage, ServerRequest};
51use crate::protocol::{
52    ClientInfo, InitializeParams, InitializeResponse, ThreadArchiveParams, ThreadArchiveResponse,
53    ThreadDeleteParams, ThreadDeleteResponse, ThreadForkParams, ThreadForkResponse,
54    ThreadItemsListParams, ThreadItemsListResponse, ThreadResumeParams, ThreadResumeResponse,
55    ThreadRevertParams, ThreadRevertResponse, ThreadStartParams, ThreadStartResponse,
56    ThreadTurnsListParams, ThreadTurnsListResponse, TurnInterruptParams, TurnInterruptResponse,
57    TurnStartParams, TurnStartResponse, TurnSteerParams, TurnSteerResponse,
58};
59use crate::protocol_generated::types::{
60    CancelLoginAccountParams, CancelLoginAccountResponse, GetAccountParams,
61    GetAccountRateLimitsParams, GetAccountRateLimitsResponse, GetAccountResponse,
62    GetAccountTokenUsageResponse, LoginAccountParams, LoginAccountResponse, LogoutAccountResponse,
63};
64use log::{debug, error, warn};
65use serde::de::DeserializeOwned;
66use serde::Serialize;
67use std::collections::VecDeque;
68use std::sync::atomic::{AtomicI64, Ordering};
69use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
70use tokio::process::Child;
71
72/// Buffer size for reading stdout (10MB).
73const STDOUT_BUFFER_SIZE: usize = 10 * 1024 * 1024;
74
75/// Asynchronous multi-turn client for the Codex app-server.
76///
77/// Communicates with a long-lived `codex app-server` process via
78/// newline-delimited JSON-RPC over stdio. Manages request/response
79/// correlation and buffers incoming notifications that arrive while
80/// waiting for RPC responses.
81///
82/// Reads performed by [`AsyncClient::request`] and [`AsyncClient::next_message`]
83/// are cancellation-safe. If either future is dropped while a JSON line is
84/// only partially available, the partial frame is retained and completed by
85/// the next read operation.
86///
87/// The client automatically kills the app-server process when dropped.
88pub struct AsyncClient {
89    child: Child,
90    writer: BufWriter<tokio::process::ChildStdin>,
91    reader: BufReader<tokio::process::ChildStdout>,
92    /// Bytes read from stdout for the current, not-yet-decoded frame.
93    inbound_frame: Vec<u8>,
94    /// Handle to the background task draining the child's stderr pipe.
95    /// Kept alive for the lifetime of the client; the task exits on EOF
96    /// when the child is killed.
97    _stderr_drain: tokio::task::JoinHandle<()>,
98    next_id: AtomicI64,
99    /// Buffered incoming messages (notifications/server requests) that arrived
100    /// while waiting for a response to a client request.
101    buffered: VecDeque<ServerMessage>,
102}
103
104impl AsyncClient {
105    /// Create a client from an existing Tokio child process.
106    ///
107    /// The child's stdin, stdout, and stderr must all be piped. This does not
108    /// perform the app-server `initialize` handshake or a Codex version check.
109    /// Use [`AppServerBuilder::build_command`] to retain the SDK's command-line
110    /// and stdio configuration while customizing how the process is spawned.
111    pub fn new(mut child: Child) -> Result<Self> {
112        let stdin = child
113            .stdin
114            .take()
115            .ok_or_else(|| Error::Protocol("Failed to get stdin".to_string()))?;
116        let stdout = child
117            .stdout
118            .take()
119            .ok_or_else(|| Error::Protocol("Failed to get stdout".to_string()))?;
120        let stderr = child
121            .stderr
122            .take()
123            .ok_or_else(|| Error::Protocol("Failed to get stderr".to_string()))?;
124
125        // The app-server emits ~200 KB/s of tracing to stderr. Without an
126        // active reader, the ~64 KB kernel pipe fills almost instantly and
127        // the child blocks. Drain in the background and route lines through
128        // the `log` crate (see [`crate::stderr_drain`]).
129        let stderr_drain = crate::stderr_drain::spawn_async(stderr);
130
131        Ok(Self {
132            child,
133            writer: BufWriter::new(stdin),
134            reader: BufReader::with_capacity(STDOUT_BUFFER_SIZE, stdout),
135            inbound_frame: Vec::new(),
136            _stderr_drain: stderr_drain,
137            next_id: AtomicI64::new(1),
138            buffered: VecDeque::new(),
139        })
140    }
141
142    /// Start an app-server with default settings.
143    ///
144    /// Spawns `codex app-server --listen stdio://`, performs the required
145    /// `initialize` handshake, and returns a connected client ready for
146    /// `thread_start()`.
147    ///
148    /// # Errors
149    ///
150    /// Returns an error if the `codex` CLI is not installed, the version is
151    /// incompatible, the process fails to start, or the initialization
152    /// handshake fails.
153    pub async fn start() -> Result<Self> {
154        Self::start_with(AppServerBuilder::new()).await
155    }
156
157    /// Start an app-server with a custom [`AppServerBuilder`].
158    ///
159    /// Performs the required `initialize` handshake before returning.
160    /// Use this to configure the binary path, working directory, environment,
161    /// or CLI arguments.
162    ///
163    /// # Errors
164    ///
165    /// Returns an error if the process fails to start, stdio pipes
166    /// cannot be established, or the initialization handshake fails.
167    pub async fn start_with(builder: AppServerBuilder) -> Result<Self> {
168        let mut client = Self::spawn(builder).await?;
169        client
170            .initialize(&InitializeParams {
171                client_info: ClientInfo {
172                    name: "codex-codes".to_string(),
173                    version: env!("CARGO_PKG_VERSION").to_string(),
174                    title: None,
175                },
176                capabilities: None,
177            })
178            .await?;
179        Ok(client)
180    }
181
182    /// Spawn an app-server without performing the `initialize` handshake.
183    ///
184    /// Use this if you need to send a custom [`InitializeParams`] (e.g., with
185    /// specific capabilities). You **must** call [`AsyncClient::initialize`]
186    /// before any other requests.
187    pub async fn spawn(builder: AppServerBuilder) -> Result<Self> {
188        crate::version::check_codex_version_async().await?;
189        Self::new(builder.spawn().await?)
190    }
191
192    /// Send a JSON-RPC request and wait for the matching response.
193    ///
194    /// Any notifications or server requests that arrive before the response
195    /// are buffered and can be retrieved via [`AsyncClient::next_message`].
196    /// Dropping this future during a partial inbound frame preserves that frame
197    /// for the next call to `request` or [`AsyncClient::next_message`].
198    ///
199    /// # Errors
200    ///
201    /// - [`Error::JsonRpc`] if the server returns a JSON-RPC error
202    /// - [`Error::ServerClosed`] if the connection drops before a response arrives
203    /// - [`Error::Json`] if response deserialization fails
204    pub async fn request<P: Serialize, R: DeserializeOwned>(
205        &mut self,
206        method: &str,
207        params: &P,
208    ) -> Result<R> {
209        let id = RequestId::Integer(self.next_id.fetch_add(1, Ordering::Relaxed));
210
211        let req = JsonRpcRequest {
212            id: id.clone(),
213            method: method.to_string(),
214            params: Some(serde_json::to_value(params).map_err(Error::Json)?),
215        };
216
217        self.send_raw(&req).await?;
218
219        // Read lines until we get a response matching our id
220        loop {
221            let msg = self.read_message().await?;
222            match msg {
223                JsonRpcMessage::Response(resp) if resp.id == id => {
224                    let result: R = serde_json::from_value(resp.result).map_err(Error::Json)?;
225                    return Ok(result);
226                }
227                JsonRpcMessage::Error(err) if err.id == id => {
228                    return Err(Error::JsonRpc {
229                        code: err.error.code,
230                        message: err.error.message,
231                    });
232                }
233                // Buffer notifications and server requests
234                JsonRpcMessage::Notification(notif) => {
235                    let typed = Notification::from_envelope(&notif.method, notif.params)
236                        .map_err(Error::Json)?;
237                    self.buffered.push_back(ServerMessage::Notification(typed));
238                }
239                JsonRpcMessage::Request(req) => {
240                    let typed = ServerRequest::from_envelope(&req.method, req.params)
241                        .map_err(Error::Json)?;
242                    self.buffered.push_back(ServerMessage::Request {
243                        id: req.id,
244                        request: typed,
245                    });
246                }
247                // Response/error for a different id — unexpected
248                JsonRpcMessage::Response(resp) => {
249                    warn!(
250                        "[CLIENT] Unexpected response for id={}, expected id={}",
251                        resp.id, id
252                    );
253                }
254                JsonRpcMessage::Error(err) => {
255                    warn!(
256                        "[CLIENT] Unexpected error for id={}, expected id={}",
257                        err.id, id
258                    );
259                }
260            }
261        }
262    }
263
264    /// Start a new thread (conversation session).
265    ///
266    /// A thread must be created before any turns can be started. The returned
267    /// [`ThreadStartResponse`] contains the `thread_id` needed for subsequent calls.
268    pub async fn thread_start(
269        &mut self,
270        params: &ThreadStartParams,
271    ) -> Result<ThreadStartResponse> {
272        self.request(crate::protocol::methods::THREAD_START, params)
273            .await
274    }
275
276    /// Resume a previously persisted thread by id.
277    ///
278    /// Replays the thread's history so turns can continue where they left off.
279    pub async fn thread_resume(
280        &mut self,
281        params: &ThreadResumeParams,
282    ) -> Result<ThreadResumeResponse> {
283        self.request(crate::protocol::methods::THREAD_RESUME, params)
284            .await
285    }
286
287    /// Fork an existing thread into a new independent thread.
288    pub async fn thread_fork(&mut self, params: &ThreadForkParams) -> Result<ThreadForkResponse> {
289        self.request(crate::protocol::methods::THREAD_FORK, params)
290            .await
291    }
292
293    /// Start a new turn within a thread.
294    ///
295    /// Sends user input to the agent. After calling this, use [`AsyncClient::next_message`]
296    /// to stream notifications until `turn/completed` arrives.
297    pub async fn turn_start(&mut self, params: &TurnStartParams) -> Result<TurnStartResponse> {
298        self.request(crate::protocol::methods::TURN_START, params)
299            .await
300    }
301
302    /// Steer an active turn with additional user input (`turn/steer`) —
303    /// appends to the running turn instead of starting a new one.
304    pub async fn turn_steer(&mut self, params: &TurnSteerParams) -> Result<TurnSteerResponse> {
305        self.request(crate::protocol::methods::TURN_STEER, params)
306            .await
307    }
308
309    /// Page through a thread's items in canonical order (`thread/items/list`).
310    /// Follow `next_cursor` for subsequent pages (0.148 upstream).
311    pub async fn thread_items_list(
312        &mut self,
313        params: &ThreadItemsListParams,
314    ) -> Result<ThreadItemsListResponse> {
315        self.request(crate::protocol::methods::THREAD_ITEMS_LIST, params)
316            .await
317    }
318
319    /// Page through a thread's turns (`thread/turns/list`); defaults to
320    /// newest-first with summary item detail (0.148 upstream).
321    pub async fn thread_turns_list(
322        &mut self,
323        params: &ThreadTurnsListParams,
324    ) -> Result<ThreadTurnsListResponse> {
325        self.request(crate::protocol::methods::THREAD_TURNS_LIST, params)
326            .await
327    }
328
329    /// Replace a paginated thread's durable history with the prefix before
330    /// one turn (`thread/revert`). Does not revert local file changes
331    /// (0.148 upstream).
332    pub async fn thread_revert(
333        &mut self,
334        params: &ThreadRevertParams,
335    ) -> Result<ThreadRevertResponse> {
336        self.request(crate::protocol::methods::THREAD_REVERT, params)
337            .await
338    }
339
340    /// Interrupt an active turn.
341    pub async fn turn_interrupt(
342        &mut self,
343        params: &TurnInterruptParams,
344    ) -> Result<TurnInterruptResponse> {
345        self.request(crate::protocol::methods::TURN_INTERRUPT, params)
346            .await
347    }
348
349    /// Archive a thread.
350    pub async fn thread_archive(
351        &mut self,
352        params: &ThreadArchiveParams,
353    ) -> Result<ThreadArchiveResponse> {
354        self.request(crate::protocol::methods::THREAD_ARCHIVE, params)
355            .await
356    }
357
358    /// Delete a thread.
359    pub async fn thread_delete(
360        &mut self,
361        params: &ThreadDeleteParams,
362    ) -> Result<ThreadDeleteResponse> {
363        self.request(crate::protocol::methods::THREAD_DELETE, params)
364            .await
365    }
366
367    /// Perform the `initialize` handshake with the app-server.
368    ///
369    /// Sends `initialize` with the given params and then sends the
370    /// `initialized` notification. This must be the first request after
371    /// spawning the process.
372    pub async fn initialize(&mut self, params: &InitializeParams) -> Result<InitializeResponse> {
373        let resp: InitializeResponse = self
374            .request(crate::protocol::methods::INITIALIZE, params)
375            .await?;
376        self.send_notification(crate::protocol::methods::INITIALIZED)
377            .await?;
378        Ok(resp)
379    }
380
381    /// Respond to a server-to-client request (e.g., approval flow).
382    ///
383    /// When the server sends a [`ServerMessage::Request`], it expects a response.
384    /// Use this method with the request's `id` and a result payload. For command
385    /// approval, pass a [`CommandExecutionApprovalResponse`](crate::CommandExecutionApprovalResponse).
386    /// For file change approval, pass a [`FileChangeApprovalResponse`](crate::FileChangeApprovalResponse).
387    pub async fn respond<R: Serialize>(&mut self, id: RequestId, result: &R) -> Result<()> {
388        let resp = JsonRpcResponse {
389            id,
390            result: serde_json::to_value(result).map_err(Error::Json)?,
391        };
392        self.send_raw(&resp).await
393    }
394
395    /// Respond to a server-to-client request with an error.
396    pub async fn respond_error(&mut self, id: RequestId, code: i64, message: &str) -> Result<()> {
397        let err = JsonRpcError {
398            id,
399            error: crate::jsonrpc::JsonRpcErrorData {
400                code,
401                message: message.to_string(),
402                data: None,
403            },
404        };
405        self.send_raw(&err).await
406    }
407
408    /// Read the next incoming server message (notification or server request).
409    ///
410    /// Returns buffered messages first (from notifications that arrived during
411    /// an [`AsyncClient::request`] call), then reads from the wire.
412    ///
413    /// Returns `Ok(None)` when the app-server closes the connection (EOF).
414    /// Dropping this future during a partial inbound frame preserves that frame
415    /// for the next call to `next_message` or [`AsyncClient::request`].
416    ///
417    /// # Typical notification methods
418    ///
419    /// | Method | Meaning |
420    /// |--------|---------|
421    /// | `turn/started` | Agent began processing |
422    /// | `item/agentMessage/delta` | Streaming text chunk |
423    /// | `item/commandExecution/outputDelta` | Command output chunk |
424    /// | `item/started` / `item/completed` | Item lifecycle |
425    /// | `turn/completed` | Agent finished the turn |
426    /// | `error` | Server-side error |
427    pub async fn next_message(&mut self) -> Result<Option<ServerMessage>> {
428        // Drain buffered messages first
429        if let Some(msg) = self.buffered.pop_front() {
430            return Ok(Some(msg));
431        }
432
433        // Read from the wire
434        loop {
435            let msg = match self.read_message_opt().await? {
436                Some(m) => m,
437                None => return Ok(None),
438            };
439
440            match msg {
441                JsonRpcMessage::Notification(notif) => {
442                    let JsonRpcNotification { method, params } = notif;
443                    let typed =
444                        Notification::from_envelope(&method, params.clone()).map_err(|e| {
445                            Error::Deserialization(ParseError::from_envelope(method, params, e))
446                        })?;
447                    return Ok(Some(ServerMessage::Notification(typed)));
448                }
449                JsonRpcMessage::Request(req) => {
450                    let JsonRpcRequest { id, method, params } = req;
451                    let typed =
452                        ServerRequest::from_envelope(&method, params.clone()).map_err(|e| {
453                            Error::Deserialization(ParseError::from_envelope(method, params, e))
454                        })?;
455                    return Ok(Some(ServerMessage::Request { id, request: typed }));
456                }
457                // Unexpected responses without a pending request
458                JsonRpcMessage::Response(resp) => {
459                    warn!(
460                        "[CLIENT] Unexpected response (no pending request): id={}",
461                        resp.id
462                    );
463                }
464                JsonRpcMessage::Error(err) => {
465                    warn!(
466                        "[CLIENT] Unexpected error (no pending request): id={} code={}",
467                        err.id, err.error.code
468                    );
469                }
470            }
471        }
472    }
473
474    /// Return an async event stream over [`ServerMessage`]s.
475    ///
476    /// Wraps [`AsyncClient::next_message`] in a stream-like API. Call
477    /// [`EventStream::next`] in a loop, or [`EventStream::collect`] to
478    /// gather all messages until EOF.
479    pub fn events(&mut self) -> EventStream<'_> {
480        EventStream { client: self }
481    }
482
483    /// Get the process ID.
484    pub fn pid(&self) -> Option<u32> {
485        self.child.id()
486    }
487
488    // ── Account / auth methods ─────────────────────────────────────────
489
490    /// `account/read` — the active account (plan, email, auth mode), or
491    /// `account: null` when logged out.
492    pub async fn account_read(&mut self, params: &GetAccountParams) -> Result<GetAccountResponse> {
493        self.request(crate::protocol::methods::ACCOUNT_READ, params)
494            .await
495    }
496
497    /// `account/login/start` — begin a login. The params select the mode
498    /// (`apiKey` completes immediately; `chatgpt` returns an auth URL to
499    /// open; `chatgptDeviceCode` returns a user code + verification URL).
500    /// Browser/device modes complete asynchronously: watch for the
501    /// `account/login/completed` notification, or cancel with
502    /// [`account_login_cancel`](Self::account_login_cancel).
503    pub async fn account_login_start(
504        &mut self,
505        params: &LoginAccountParams,
506    ) -> Result<LoginAccountResponse> {
507        self.request(crate::protocol::methods::ACCOUNT_LOGIN_START, params)
508            .await
509    }
510
511    /// `account/login/cancel` — abort an in-flight browser/device login.
512    pub async fn account_login_cancel(
513        &mut self,
514        params: &CancelLoginAccountParams,
515    ) -> Result<CancelLoginAccountResponse> {
516        self.request(crate::protocol::methods::ACCOUNT_LOGIN_CANCEL, params)
517            .await
518    }
519
520    /// `account/logout` — remove the stored credential.
521    pub async fn account_logout(&mut self) -> Result<LogoutAccountResponse> {
522        self.request(
523            crate::protocol::methods::ACCOUNT_LOGOUT,
524            &serde_json::json!({}),
525        )
526        .await
527    }
528
529    /// `account/rateLimits/read` — current rate-limit windows.
530    ///
531    /// `params` declares the client's usage-read capabilities;
532    /// `GetAccountRateLimitsParams::default()` serializes as `{}` and matches
533    /// the pre-capability wire shape.
534    pub async fn account_rate_limits_read(
535        &mut self,
536        params: GetAccountRateLimitsParams,
537    ) -> Result<GetAccountRateLimitsResponse> {
538        self.request(crate::protocol::methods::ACCOUNT_RATELIMITS_READ, &params)
539            .await
540    }
541
542    /// `account/usage/read` — token-usage summary for the account.
543    pub async fn account_usage_read(&mut self) -> Result<GetAccountTokenUsageResponse> {
544        self.request(
545            crate::protocol::methods::ACCOUNT_USAGE_READ,
546            &serde_json::json!({}),
547        )
548        .await
549    }
550
551    /// Check if the child process is still running.
552    pub fn is_alive(&mut self) -> bool {
553        self.child.try_wait().ok().flatten().is_none()
554    }
555
556    /// Shut down the app-server process.
557    ///
558    /// Consumes the client. If you don't call this explicitly, the
559    /// [`Drop`] implementation will kill the process automatically.
560    pub async fn shutdown(mut self) -> Result<()> {
561        debug!("[CLIENT] Shutting down");
562        self.child.kill().await.map_err(Error::Io)?;
563        Ok(())
564    }
565
566    // -- internal --
567
568    async fn send_notification(&mut self, method: &str) -> Result<()> {
569        let notif = JsonRpcNotification {
570            method: method.to_string(),
571            params: None,
572        };
573        self.send_raw(&notif).await
574    }
575
576    async fn send_raw<T: Serialize>(&mut self, msg: &T) -> Result<()> {
577        let json = serde_json::to_string(msg).map_err(Error::Json)?;
578        debug!("[CLIENT] Sending: {}", json);
579        self.writer
580            .write_all(json.as_bytes())
581            .await
582            .map_err(Error::Io)?;
583        self.writer.write_all(b"\n").await.map_err(Error::Io)?;
584        self.writer.flush().await.map_err(Error::Io)?;
585        Ok(())
586    }
587
588    async fn read_message(&mut self) -> Result<JsonRpcMessage> {
589        self.read_message_opt().await?.ok_or(Error::ServerClosed)
590    }
591
592    async fn read_message_opt(&mut self) -> Result<Option<JsonRpcMessage>> {
593        loop {
594            // `read_until` is cancellation-safe: bytes consumed from `reader`
595            // are appended to the persistent buffer before the await can be
596            // cancelled. Do not clear that buffer until a full frame exists.
597            let bytes_read = self
598                .reader
599                .read_until(b'\n', &mut self.inbound_frame)
600                .await
601                .map_err(Error::Io)?;
602
603            if bytes_read == 0 {
604                debug!("[CLIENT] Stream closed (EOF)");
605                if self.inbound_frame.is_empty() {
606                    return Ok(None);
607                }
608            }
609
610            if !self.inbound_frame.ends_with(b"\n") && bytes_read != 0 {
611                continue;
612            }
613
614            let line = match std::str::from_utf8(&self.inbound_frame) {
615                Ok(line) => line,
616                Err(error) => {
617                    let error = std::io::Error::new(std::io::ErrorKind::InvalidData, error);
618                    self.inbound_frame.clear();
619                    return Err(Error::Io(error));
620                }
621            };
622            let trimmed = line.trim();
623            if trimmed.is_empty() {
624                self.inbound_frame.clear();
625                continue;
626            }
627
628            debug!("[CLIENT] Received: {}", trimmed);
629
630            let decoded = serde_json::from_str::<JsonRpcMessage>(trimmed);
631            match decoded {
632                Ok(msg) => {
633                    self.inbound_frame.clear();
634                    return Ok(Some(msg));
635                }
636                Err(e) => {
637                    warn!(
638                        "[CLIENT] Failed to deserialize message. \
639                         Please report this at https://github.com/meawoppl/rust-code-agent-sdks/issues"
640                    );
641                    warn!("[CLIENT] Parse error: {}", e);
642                    warn!("[CLIENT] Raw: {}", trimmed);
643                    let parse_error = ParseError::from_line(trimmed, e);
644                    self.inbound_frame.clear();
645                    return Err(Error::Deserialization(parse_error));
646                }
647            }
648        }
649    }
650}
651
652impl Drop for AsyncClient {
653    fn drop(&mut self) {
654        if self.is_alive() {
655            if let Err(e) = self.child.start_kill() {
656                error!("Failed to kill app-server process on drop: {}", e);
657            }
658        }
659    }
660}
661
662/// Async stream of [`ServerMessage`]s from an [`AsyncClient`].
663pub struct EventStream<'a> {
664    client: &'a mut AsyncClient,
665}
666
667impl EventStream<'_> {
668    /// Get the next server message.
669    pub async fn next(&mut self) -> Option<Result<ServerMessage>> {
670        match self.client.next_message().await {
671            Ok(Some(msg)) => Some(Ok(msg)),
672            Ok(None) => None,
673            Err(e) => Some(Err(e)),
674        }
675    }
676
677    /// Collect all remaining messages.
678    pub async fn collect(mut self) -> Result<Vec<ServerMessage>> {
679        let mut msgs = Vec::new();
680        while let Some(result) = self.next().await {
681            msgs.push(result?);
682        }
683        Ok(msgs)
684    }
685}
686
687#[cfg(test)]
688mod tests {
689    use super::*;
690    use std::process::Stdio;
691    use tokio::process::Command;
692
693    #[cfg(unix)]
694    fn scripted_client(script: &str) -> AsyncClient {
695        let mut command = Command::new("sh");
696        command
697            .arg("-c")
698            .arg(script)
699            .stdin(Stdio::piped())
700            .stdout(Stdio::piped())
701            .stderr(Stdio::piped());
702        AsyncClient::new(command.spawn().expect("spawn scripted app-server"))
703            .expect("construct async client")
704    }
705
706    fn unknown_method(message: ServerMessage) -> String {
707        match message {
708            ServerMessage::Notification(Notification::Unknown { method, .. }) => method,
709            other => panic!("expected unknown notification, got {other:?}"),
710        }
711    }
712
713    #[test]
714    fn test_buffer_size() {
715        assert_eq!(STDOUT_BUFFER_SIZE, 10 * 1024 * 1024);
716    }
717
718    #[cfg(unix)]
719    #[tokio::test]
720    async fn cancelled_next_message_resumes_partial_frame_exactly_once() {
721        const PARTIAL: &[u8] = br#"{"method":"test/first","params":{"part":"#;
722        let mut client = scripted_client(
723            r#"printf '%s' '{"method":"test/first","params":{"part":'; IFS= read -r release; printf '%s\n' '1}}'; printf '%s\n' '{"method":"test/second","params":{}}'"#,
724        );
725
726        assert_eq!(
727            client
728                .reader
729                .fill_buf()
730                .await
731                .expect("buffer partial frame"),
732            PARTIAL
733        );
734        let mut pending_read = Box::pin(client.next_message());
735        tokio::select! {
736            biased;
737            result = &mut pending_read => panic!("partial frame completed unexpectedly: {result:?}"),
738            _ = async {} => {}
739        }
740        drop(pending_read);
741        // A pipe may split one write across reads, so require the bytes that
742        // reached the decoder to be a nonempty prefix, not the entire write.
743        assert!(!client.inbound_frame.is_empty());
744        assert!(PARTIAL.starts_with(&client.inbound_frame));
745        client
746            .writer
747            .write_all(b"release\n")
748            .await
749            .expect("release remaining frame");
750        client.writer.flush().await.expect("flush release");
751
752        let first = client
753            .next_message()
754            .await
755            .expect("resume first frame")
756            .expect("first message");
757        let second = client
758            .next_message()
759            .await
760            .expect("read second frame")
761            .expect("second message");
762
763        assert_eq!(unknown_method(first), "test/first");
764        assert_eq!(unknown_method(second), "test/second");
765        assert!(client.next_message().await.expect("read EOF").is_none());
766    }
767
768    #[cfg(unix)]
769    #[tokio::test]
770    async fn cancelled_request_preserves_shared_decoder_framing() {
771        const PARTIAL: &[u8] = br#"{"id":1,"result":{"abandoned":"#;
772        let mut client = scripted_client(
773            r#"printf '%s' '{"id":1,"result":{"abandoned":'; IFS= read -r first; IFS= read -r release; printf '%s\n' 'true}}' '{"method":"test/between","params":{}}'; IFS= read -r second; printf '%s\n' '{"id":2,"result":{"ok":true}}' '{"method":"test/after","params":{}}'"#,
774        );
775
776        assert_eq!(
777            client
778                .reader
779                .fill_buf()
780                .await
781                .expect("buffer partial response"),
782            PARTIAL
783        );
784        let params = serde_json::json!({});
785        let mut pending_request =
786            Box::pin(client.request::<_, serde_json::Value>("test/abandoned", &params));
787        tokio::select! {
788            biased;
789            result = &mut pending_request => panic!("partial response completed unexpectedly: {result:?}"),
790            _ = async {} => {}
791        }
792        drop(pending_request);
793        // A pipe may split one write across reads, so require the bytes that
794        // reached the decoder to be a nonempty prefix, not the entire write.
795        assert!(!client.inbound_frame.is_empty());
796        assert!(PARTIAL.starts_with(&client.inbound_frame));
797        client
798            .writer
799            .write_all(b"release\n")
800            .await
801            .expect("release remaining response");
802        client.writer.flush().await.expect("flush release");
803
804        let response: serde_json::Value = client
805            .request("test/resumed", &serde_json::json!({}))
806            .await
807            .expect("second request should resume the shared decoder");
808        assert_eq!(response, serde_json::json!({"ok": true}));
809
810        let between = client
811            .next_message()
812            .await
813            .expect("read buffered notification")
814            .expect("between message");
815        let after = client
816            .next_message()
817            .await
818            .expect("read trailing notification")
819            .expect("after message");
820        assert_eq!(unknown_method(between), "test/between");
821        assert_eq!(unknown_method(after), "test/after");
822        assert!(client.next_message().await.expect("read EOF").is_none());
823    }
824
825    #[cfg(unix)]
826    #[tokio::test]
827    async fn next_message_resumes_notification_partially_read_by_cancelled_request() {
828        const PARTIAL: &[u8] = br#"{"method":"test/during-request","params":{"part":"#;
829        let mut client = scripted_client(
830            r#"printf '%s' '{"method":"test/during-request","params":{"part":'; IFS= read -r request; IFS= read -r release; printf '%s\n' '1}}' '{"id":1,"result":{}}'"#,
831        );
832
833        assert_eq!(
834            client
835                .reader
836                .fill_buf()
837                .await
838                .expect("buffer partial notification"),
839            PARTIAL
840        );
841        let params = serde_json::json!({});
842        let mut pending_request =
843            Box::pin(client.request::<_, serde_json::Value>("test/abandoned", &params));
844        tokio::select! {
845            biased;
846            result = &mut pending_request => panic!("partial notification completed unexpectedly: {result:?}"),
847            _ = async {} => {}
848        }
849        drop(pending_request);
850        // A pipe may split one write across reads, so require the bytes that
851        // reached the decoder to be a nonempty prefix, not the entire write.
852        assert!(!client.inbound_frame.is_empty());
853        assert!(PARTIAL.starts_with(&client.inbound_frame));
854        client
855            .writer
856            .write_all(b"release\n")
857            .await
858            .expect("release remaining notification");
859        client.writer.flush().await.expect("flush release");
860
861        let notification = client
862            .next_message()
863            .await
864            .expect("resume partial notification")
865            .expect("notification");
866        assert_eq!(unknown_method(notification), "test/during-request");
867        assert!(client.next_message().await.expect("read EOF").is_none());
868    }
869
870    #[cfg(unix)]
871    #[tokio::test]
872    async fn request_resumes_notification_partially_read_by_cancelled_next_message() {
873        const PARTIAL: &[u8] = br#"{"method":"test/before-request","params":{"part":"#;
874        let mut client = scripted_client(
875            r#"printf '%s' '{"method":"test/before-request","params":{"part":'; IFS= read -r release; printf '%s\n' '1}}'; IFS= read -r request; printf '%s\n' '{"id":1,"result":{"ok":true}}'"#,
876        );
877
878        assert_eq!(
879            client
880                .reader
881                .fill_buf()
882                .await
883                .expect("buffer partial notification"),
884            PARTIAL
885        );
886        let mut pending_read = Box::pin(client.next_message());
887        tokio::select! {
888            biased;
889            result = &mut pending_read => panic!("partial notification completed unexpectedly: {result:?}"),
890            _ = async {} => {}
891        }
892        drop(pending_read);
893        // A pipe may split one write across reads, so require the bytes that
894        // reached the decoder to be a nonempty prefix, not the entire write.
895        assert!(!client.inbound_frame.is_empty());
896        assert!(PARTIAL.starts_with(&client.inbound_frame));
897        client
898            .writer
899            .write_all(b"release\n")
900            .await
901            .expect("release remaining notification");
902        client.writer.flush().await.expect("flush release");
903
904        let response: serde_json::Value = client
905            .request("test/resumed", &serde_json::json!({}))
906            .await
907            .expect("request should resume the shared decoder");
908        assert_eq!(response, serde_json::json!({"ok": true}));
909
910        let notification = client
911            .next_message()
912            .await
913            .expect("read buffered notification")
914            .expect("notification");
915        assert_eq!(unknown_method(notification), "test/before-request");
916        assert!(client.next_message().await.expect("read EOF").is_none());
917    }
918
919    #[cfg(unix)]
920    #[tokio::test]
921    async fn request_directly_resumes_response_partially_read_by_cancelled_next_message() {
922        const PARTIAL: &[u8] = br#"{"id":1,"result":{"ok":"#;
923        // The fixture knows the client's initial request id, but withholds the
924        // rest of the response until it has received that request.
925        let mut client = scripted_client(
926            r#"printf '%s' '{"id":1,"result":{"ok":'; IFS= read -r request; printf '%s\n' 'true}}'"#,
927        );
928
929        assert_eq!(
930            client
931                .reader
932                .fill_buf()
933                .await
934                .expect("buffer partial response"),
935            PARTIAL
936        );
937        let mut pending_read = Box::pin(client.next_message());
938        tokio::select! {
939            biased;
940            result = &mut pending_read => panic!("partial response completed unexpectedly: {result:?}"),
941            _ = async {} => {}
942        }
943        drop(pending_read);
944        // A pipe may split one write across reads, so require the bytes that
945        // reached the decoder to be a nonempty prefix, not the entire write.
946        assert!(!client.inbound_frame.is_empty());
947        assert!(PARTIAL.starts_with(&client.inbound_frame));
948        assert!(client.buffered.is_empty());
949
950        let response: serde_json::Value = client
951            .request("test/resumed", &serde_json::json!({}))
952            .await
953            .expect("request should directly resume the partial response");
954        assert_eq!(response, serde_json::json!({"ok": true}));
955        assert!(client.buffered.is_empty());
956        assert!(client.next_message().await.expect("read EOF").is_none());
957    }
958}
codex_codes/client_async.rs

codex_codes/
client_async.rs
