ag_agent/
app_server_transport.rs

1//! Shared stdio JSON-RPC transport utilities for app-server protocols.
2//!
3//! Provides low-level helpers for NDJSON-over-stdio communication used by
4//! persistent app-server backends such as Codex app-server. Each helper is
5//! protocol-agnostic — it operates on raw JSON values and async stdio handles
6//! without knowledge of specific method names or event shapes.
7
8use std::time::Duration;
9
10use serde_json::Value;
11use tokio::io::{AsyncWriteExt, BufReader, Lines};
12
13use crate::app_server::AppServerError;
14
15/// Typed error returned by shared app-server transport operations.
16///
17/// Covers the low-level stdio communication failures that can occur when
18/// writing JSON-RPC payloads to a child process or reading responses from
19/// its stdout stream.
20#[derive(Debug, thiserror::Error)]
21pub enum AppServerTransportError {
22    /// An IO error occurred during app-server stdio communication.
23    #[error("{context}: {source}")]
24    Io {
25        /// Human-readable description of the operation that failed.
26        context: String,
27        /// Underlying IO error.
28        #[source]
29        source: std::io::Error,
30    },
31
32    /// The app-server process terminated before sending the expected response.
33    #[error("App-server terminated before sending expected response")]
34    ProcessTerminated,
35
36    /// Timed out waiting for a JSON-RPC response from the app-server.
37    #[error(
38        "Timed out waiting for app-server response `{response_id}` after {timeout_seconds} seconds"
39    )]
40    Timeout {
41        /// The JSON-RPC request identifier that was being awaited.
42        response_id: String,
43        /// Number of seconds elapsed before the timeout fired.
44        timeout_seconds: u64,
45    },
46}
47
48/// Default timeout for initialization handshakes and session creation.
49///
50/// App-server cold starts can take materially longer than a typical
51/// request/response round trip while the runtime initializes tools and model
52/// state, so the shared startup window stays measured in minutes rather than
53/// seconds to avoid aborting healthy app-server bootstraps.
54pub const STARTUP_TIMEOUT: Duration = Duration::from_mins(5);
55
56/// Default timeout for a single prompt turn.
57///
58/// App-server turns may legitimately run for long periods while agents plan,
59/// execute tools, and compact context, so the shared turn window is aligned
60/// with the long-running Codex behavior instead of the shorter bootstrap
61/// timeout.
62pub const TURN_TIMEOUT: Duration = Duration::from_hours(4);
63
64/// Writes one JSON-RPC payload as a newline-delimited line to `stdin`.
65///
66/// # Errors
67///
68/// Returns an error when the write or flush to stdin fails.
69pub async fn write_json_line(
70    stdin: &mut tokio::process::ChildStdin,
71    payload: &Value,
72) -> Result<(), AppServerTransportError> {
73    let serialized_payload = payload.to_string();
74
75    stdin
76        .write_all(serialized_payload.as_bytes())
77        .await
78        .map_err(|source| AppServerTransportError::Io {
79            context: "Failed writing to app-server stdin".to_string(),
80            source,
81        })?;
82    stdin
83        .write_all(b"\n")
84        .await
85        .map_err(|source| AppServerTransportError::Io {
86            context: "Failed writing newline to app-server stdin".to_string(),
87            source,
88        })?;
89    stdin
90        .flush()
91        .await
92        .map_err(|source| AppServerTransportError::Io {
93            context: "Failed flushing app-server stdin".to_string(),
94            source,
95        })
96}
97
98/// Reads stdout lines until a JSON-RPC response carrying `response_id` arrives.
99///
100/// Non-matching lines (notifications, other responses) are silently skipped.
101/// Times out after [`STARTUP_TIMEOUT`].
102///
103/// # Errors
104///
105/// Returns an error when the read times out or the child process terminates
106/// before a matching response is received.
107pub async fn wait_for_response_line<R>(
108    stdout_lines: &mut Lines<BufReader<R>>,
109    response_id: &str,
110) -> Result<String, AppServerTransportError>
111where
112    R: tokio::io::AsyncRead + Unpin,
113{
114    tokio::time::timeout(STARTUP_TIMEOUT, async {
115        loop {
116            let stdout_line = stdout_lines
117                .next_line()
118                .await
119                .map_err(|source| AppServerTransportError::Io {
120                    context: "Failed reading app-server stdout".to_string(),
121                    source,
122                })?
123                .ok_or(AppServerTransportError::ProcessTerminated)?;
124
125            let Ok(response_value) = serde_json::from_str::<Value>(&stdout_line) else {
126                continue;
127            };
128            if response_id_matches(&response_value, response_id) {
129                return Ok(stdout_line);
130            }
131        }
132    })
133    .await
134    .map_err(|_| AppServerTransportError::Timeout {
135        response_id: response_id.to_string(),
136        timeout_seconds: STARTUP_TIMEOUT.as_secs(),
137    })?
138}
139
140/// Returns whether a JSON-RPC response line carries the expected `id`.
141pub fn response_id_matches(response_value: &Value, response_id: &str) -> bool {
142    response_value
143        .get("id")
144        .and_then(Value::as_str)
145        .is_some_and(|line_id| line_id == response_id)
146}
147
148/// Extracts a top-level `error.message` string from a JSON-RPC error response.
149pub fn extract_json_error_message(response_value: &Value) -> Option<String> {
150    response_value
151        .get("error")
152        .and_then(|error| error.get("message"))
153        .and_then(Value::as_str)
154        .map(ToString::to_string)
155}
156
157/// Gracefully shuts down a child process by closing stdin, waiting briefly,
158/// then killing if the process has not exited.
159pub async fn shutdown_child(child: &mut tokio::process::Child) {
160    // Closing stdin signals the child to exit cleanly.
161    drop(child.stdin.take());
162
163    if tokio::time::timeout(Duration::from_secs(1), child.wait())
164        .await
165        .is_err()
166    {
167        // Best-effort: process may have already exited.
168        let _ = child.kill().await;
169        // Best-effort: process may have already exited.
170        let _ = child.wait().await;
171    }
172}
173
174/// Spawns one app-server child process with piped stdin/stdout and hidden
175/// stderr, returning the child plus owned stdio handles.
176///
177/// Runtime bootstraps require line-delimited JSON-RPC over stdin/stdout, no
178/// interactive stderr stream, and `kill_on_drop(true)` so abandoned runtimes
179/// do not leak.
180///
181/// # Errors
182///
183/// Returns a provider error when the command cannot be spawned or either
184/// required stdio pipe is unavailable.
185pub fn spawn_runtime_command(
186    command: std::process::Command,
187    runtime_name: &str,
188) -> Result<
189    (
190        tokio::process::Child,
191        tokio::process::ChildStdin,
192        tokio::process::ChildStdout,
193    ),
194    AppServerError,
195> {
196    let mut command = tokio::process::Command::from(command);
197    command
198        .stdin(std::process::Stdio::piped())
199        .stdout(std::process::Stdio::piped())
200        .stderr(std::process::Stdio::null())
201        .kill_on_drop(true);
202
203    let mut child = command.spawn().map_err(|error| {
204        AppServerError::Provider(format!("Failed to spawn `{runtime_name}`: {error}"))
205    })?;
206    let stdin = child
207        .stdin
208        .take()
209        .ok_or_else(|| AppServerError::Provider(format!("{runtime_name} stdin is unavailable")))?;
210    let stdout = child
211        .stdout
212        .take()
213        .ok_or_else(|| AppServerError::Provider(format!("{runtime_name} stdout is unavailable")))?;
214
215    Ok((child, stdin, stdout))
216}
217
218#[cfg(test)]
219mod tests {
220    use std::process::Stdio;
221
222    use tokio::io::AsyncBufReadExt;
223
224    use super::*;
225
226    /// Spawns a simple echo process that mirrors stdin to stdout for transport
227    /// write tests.
228    fn spawn_cat_process() -> tokio::process::Child {
229        let mut command = tokio::process::Command::new("cat");
230        command
231            .stdin(Stdio::piped())
232            .stdout(Stdio::piped())
233            .stderr(Stdio::null())
234            .kill_on_drop(true);
235
236        command.spawn().expect("failed to spawn `cat`")
237    }
238
239    #[test]
240    fn response_id_matches_returns_true_for_matching_string_id() {
241        // Arrange
242        let response_value = serde_json::json!({"id": "init-123", "result": {}});
243
244        // Act / Assert
245        assert!(response_id_matches(&response_value, "init-123"));
246    }
247
248    #[test]
249    fn response_id_matches_returns_false_for_different_id() {
250        // Arrange
251        let response_value = serde_json::json!({"id": "init-123", "result": {}});
252
253        // Act / Assert
254        assert!(!response_id_matches(&response_value, "init-456"));
255    }
256
257    #[test]
258    fn response_id_matches_returns_false_when_id_is_missing() {
259        // Arrange
260        let response_value = serde_json::json!({"method": "session/update", "params": {}});
261
262        // Act / Assert
263        assert!(!response_id_matches(&response_value, "init-123"));
264    }
265
266    #[test]
267    fn response_id_matches_returns_false_for_integer_id() {
268        // Arrange
269        let response_value = serde_json::json!({"id": 1, "result": {}});
270
271        // Act / Assert
272        assert!(!response_id_matches(&response_value, "1"));
273    }
274
275    #[test]
276    fn extract_json_error_message_returns_message_string() {
277        // Arrange
278        let response_value = serde_json::json!({
279            "id": "req-1",
280            "error": {"code": -32600, "message": "Invalid request"}
281        });
282
283        // Act
284        let message = extract_json_error_message(&response_value);
285
286        // Assert
287        assert_eq!(message, Some("Invalid request".to_string()));
288    }
289
290    #[test]
291    fn extract_json_error_message_returns_none_without_error() {
292        // Arrange
293        let response_value = serde_json::json!({"id": "req-1", "result": {}});
294
295        // Act
296        let message = extract_json_error_message(&response_value);
297
298        // Assert
299        assert_eq!(message, None);
300    }
301
302    #[test]
303    fn extract_json_error_message_returns_none_without_message_field() {
304        // Arrange
305        let response_value = serde_json::json!({
306            "id": "req-1",
307            "error": {"code": -32600}
308        });
309
310        // Act
311        let message = extract_json_error_message(&response_value);
312
313        // Assert
314        assert_eq!(message, None);
315    }
316
317    /// Verifies `write_json_line()` serializes one compact JSON line followed
318    /// by a newline.
319    #[tokio::test]
320    async fn write_json_line_writes_serialized_payload_with_newline() {
321        // Arrange
322        let mut child = spawn_cat_process();
323        let mut stdin = child.stdin.take().expect("`cat` stdin should be piped");
324        let stdout = child.stdout.take().expect("`cat` stdout should be piped");
325        let payload = serde_json::json!({
326            "id": "req-1",
327            "method": "initialize",
328            "params": {"value": 1}
329        });
330
331        // Act
332        write_json_line(&mut stdin, &payload)
333            .await
334            .expect("write should succeed");
335        drop(stdin);
336        let echoed_line = BufReader::new(stdout)
337            .lines()
338            .next_line()
339            .await
340            .expect("stdout read should succeed")
341            .expect("echoed payload line should exist");
342
343        // Assert
344        assert_eq!(echoed_line, payload.to_string());
345        shutdown_child(&mut child).await;
346    }
347
348    /// Verifies `wait_for_response_line()` skips unrelated or invalid lines
349    /// until the matching response id arrives.
350    #[tokio::test]
351    async fn wait_for_response_line_skips_invalid_and_non_matching_lines() {
352        // Arrange
353        let (reader, mut writer) = tokio::io::duplex(512);
354        let writer_task = tokio::spawn(async move {
355            writer
356                .write_all(
357                    b"not-json\n{\"id\":\"other\",\"result\":{}}\n{\"id\":\"req-1\",\"result\":{\"ok\":true}}\n",
358                )
359                .await
360                .expect("test writer should succeed");
361        });
362        let mut stdout_lines = BufReader::new(reader).lines();
363
364        // Act
365        let response_line = wait_for_response_line(&mut stdout_lines, "req-1")
366            .await
367            .expect("matching response should be returned");
368
369        // Assert
370        assert_eq!(response_line, "{\"id\":\"req-1\",\"result\":{\"ok\":true}}");
371        writer_task.await.expect("writer task should finish");
372    }
373
374    /// Verifies `wait_for_response_line()` reports early process termination
375    /// when the stream ends before the expected response arrives.
376    #[tokio::test]
377    async fn wait_for_response_line_returns_error_when_stream_ends() {
378        // Arrange
379        let (reader, mut writer) = tokio::io::duplex(256);
380        let writer_task = tokio::spawn(async move {
381            writer
382                .write_all(b"{\"id\":\"other\",\"result\":{}}\n")
383                .await
384                .expect("test writer should succeed");
385            drop(writer);
386        });
387        let mut stdout_lines = BufReader::new(reader).lines();
388
389        // Act
390        let response_result = wait_for_response_line(&mut stdout_lines, "req-1").await;
391
392        // Assert
393        assert!(
394            matches!(
395                response_result,
396                Err(AppServerTransportError::ProcessTerminated)
397            ),
398            "expected ProcessTerminated, got: {response_result:?}"
399        );
400        writer_task.await.expect("writer task should finish");
401    }
402
403    #[test]
404    fn io_error_display_includes_context_and_source() {
405        // Arrange
406        let error = AppServerTransportError::Io {
407            context: "Failed writing to app-server stdin".to_string(),
408            source: std::io::Error::new(std::io::ErrorKind::BrokenPipe, "pipe closed"),
409        };
410
411        // Act
412        let display = error.to_string();
413
414        // Assert
415        assert_eq!(display, "Failed writing to app-server stdin: pipe closed");
416    }
417
418    #[test]
419    fn process_terminated_display_message() {
420        // Arrange
421        let error = AppServerTransportError::ProcessTerminated;
422
423        // Act / Assert
424        assert_eq!(
425            error.to_string(),
426            "App-server terminated before sending expected response"
427        );
428    }
429
430    #[test]
431    fn timeout_display_includes_response_id_and_seconds() {
432        // Arrange
433        let error = AppServerTransportError::Timeout {
434            response_id: "init-123".to_string(),
435            timeout_seconds: 300,
436        };
437
438        // Act / Assert
439        assert_eq!(
440            error.to_string(),
441            "Timed out waiting for app-server response `init-123` after 300 seconds"
442        );
443    }
444
445    /// Verifies `shutdown_child()` closes stdin and waits for a cooperative
446    /// child process to exit.
447    #[tokio::test]
448    async fn shutdown_child_exits_cleanly_after_closing_stdin() {
449        // Arrange
450        let mut child = spawn_cat_process();
451
452        // Act
453        shutdown_child(&mut child).await;
454        let exit_status = child.wait().await.expect("child wait should succeed");
455
456        // Assert
457        assert!(exit_status.success());
458    }
459}
ag_agent/app_server_transport.rs

ag_agent/
app_server_transport.rs
