agentmux_srv/backend/blockcontroller/
app_server.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Provider-neutral App Server process and JSONL/JSON-RPC transport.
5//!
6//! This module deliberately stops at lifecycle and handshake. Thread, turn,
7//! rendering, and interactive-request semantics live in the protocol adapter
8//! added by later rollout slices. No provider selects this controller yet.
9
10use std::collections::{HashMap, VecDeque};
11use std::process::{ExitStatus, Stdio};
12use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering};
13use std::sync::{Arc, Mutex};
14use std::time::Duration;
15
16use serde::Serialize;
17use serde_json::{Map, Value};
18use thiserror::Error;
19use tokio::io::{
20    AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader,
21};
22use tokio::process::Command;
23use tokio::sync::{mpsc, oneshot, watch, Mutex as AsyncMutex};
24
25const INITIALIZING: u8 = 1;
26const INITIALIZED: u8 = 2;
27const INITIALIZE_FAILED: u8 = 3;
28
29/// Conservative resource limits at the untrusted child-process boundary.
30#[derive(Debug, Clone, Copy)]
31pub struct AppServerLimits {
32    pub max_frame_bytes: usize,
33    pub max_pending_requests: usize,
34    pub outbound_queue_capacity: usize,
35    pub inbound_queue_capacity: usize,
36    pub max_stderr_bytes: usize,
37}
38
39impl Default for AppServerLimits {
40    fn default() -> Self {
41        Self {
42            max_frame_bytes: 4 * 1024 * 1024,
43            max_pending_requests: 128,
44            outbound_queue_capacity: 128,
45            inbound_queue_capacity: 256,
46            max_stderr_bytes: 64 * 1024,
47        }
48    }
49}
50
51#[derive(Debug, Clone, PartialEq, Eq, Hash)]
52pub enum RpcId {
53    Number(String),
54    String(String),
55    Null,
56}
57
58impl RpcId {
59    pub(crate) fn from_value(value: &Value) -> Result<Self, AppServerError> {
60        match value {
61            Value::Number(number) if number.is_i64() || number.is_u64() => {
62                Ok(Self::Number(number.to_string()))
63            }
64            Value::String(value) => Ok(Self::String(value.clone())),
65            Value::Null => Ok(Self::Null),
66            _ => Err(AppServerError::Protocol(
67                "JSON-RPC id must be an integer, string, or null".to_string(),
68            )),
69        }
70    }
71
72    fn outgoing_u64(&self) -> Option<u64> {
73        match self {
74            Self::Number(value) => value.parse().ok(),
75            Self::String(_) | Self::Null => None,
76        }
77    }
78
79    fn to_value(&self) -> Value {
80        match self {
81            Self::Number(value) => serde_json::from_str(value).unwrap_or(Value::Null),
82            Self::String(value) => Value::String(value.clone()),
83            Self::Null => Value::Null,
84        }
85    }
86}
87
88#[derive(Debug, Clone, PartialEq)]
89pub enum AppServerIncoming {
90    Notification {
91        method: String,
92        params: Value,
93    },
94    Request {
95        id: RpcId,
96        method: String,
97        params: Value,
98    },
99}
100
101#[derive(Debug, Clone, PartialEq, Error)]
102pub enum AppServerError {
103    #[error("failed to spawn App Server: {0}")]
104    Spawn(String),
105    #[error("App Server I/O failed: {0}")]
106    Io(String),
107    #[error("App Server protocol violation: {0}")]
108    Protocol(String),
109    #[error("App Server request limit reached ({0})")]
110    PendingRequestLimit(usize),
111    #[error("App Server request timed out after {0:?}")]
112    Timeout(Duration),
113    #[error("App Server request was rejected ({code}): {message}")]
114    Remote {
115        code: i64,
116        message: String,
117        data: Option<Value>,
118    },
119    #[error("App Server transport is closed")]
120    Closed,
121    #[error("App Server initialize may only be attempted once")]
122    InitializeAlreadyAttempted,
123}
124
125#[derive(Debug, Clone, PartialEq, Eq)]
126pub enum AppServerExit {
127    Clean {
128        code: Option<i32>,
129        stderr_tail: String,
130    },
131    ForcedAfterTimeout {
132        stderr_tail: String,
133    },
134    BeforeInitialize {
135        code: Option<i32>,
136        stderr_tail: String,
137    },
138    NonZero {
139        code: Option<i32>,
140        stderr_tail: String,
141    },
142    ProtocolViolation {
143        message: String,
144        stderr_tail: String,
145    },
146    UnexpectedEof {
147        stderr_tail: String,
148    },
149}
150
151type PendingSender = oneshot::Sender<Result<Value, AppServerError>>;
152
153struct SharedState {
154    pending: Mutex<HashMap<u64, PendingSender>>,
155    failure: Mutex<Option<AppServerError>>,
156    closing: AtomicBool,
157    closed: AtomicBool,
158}
159
160struct PendingGuard {
161    shared: Arc<SharedState>,
162    id: u64,
163}
164
165impl Drop for PendingGuard {
166    fn drop(&mut self) {
167        self.shared.pending.lock().unwrap().remove(&self.id);
168    }
169}
170
171enum Outbound {
172    Frame {
173        bytes: Vec<u8>,
174        written: oneshot::Sender<Result<(), AppServerError>>,
175    },
176    Close {
177        closed: oneshot::Sender<Result<(), AppServerError>>,
178    },
179}
180
181pub struct AppServerTransport {
182    outbound_tx: mpsc::Sender<Outbound>,
183    incoming_rx: AsyncMutex<mpsc::Receiver<AppServerIncoming>>,
184    reader_task: AsyncMutex<Option<tokio::task::JoinHandle<()>>>,
185    shared: Arc<SharedState>,
186    next_request_id: AtomicU64,
187    initialize_state: AtomicU8,
188    max_frame_bytes: usize,
189    max_pending_requests: usize,
190}
191
192impl AppServerTransport {
193    pub fn new<R, W>(reader: R, writer: W, limits: AppServerLimits) -> Self
194    where
195        R: AsyncRead + Unpin + Send + 'static,
196        W: AsyncWrite + Unpin + Send + 'static,
197    {
198        let shared = Arc::new(SharedState {
199            pending: Mutex::new(HashMap::new()),
200            failure: Mutex::new(None),
201            closing: AtomicBool::new(false),
202            closed: AtomicBool::new(false),
203        });
204        let (outbound_tx, outbound_rx) = mpsc::channel(limits.outbound_queue_capacity.max(1));
205        let (incoming_tx, incoming_rx) = mpsc::channel(limits.inbound_queue_capacity.max(1));
206
207        tokio::spawn(run_writer(writer, outbound_rx, shared.clone()));
208        let reader_task = tokio::spawn(run_reader(
209            reader,
210            incoming_tx,
211            shared.clone(),
212            limits.max_frame_bytes,
213        ));
214
215        Self {
216            outbound_tx,
217            incoming_rx: AsyncMutex::new(incoming_rx),
218            reader_task: AsyncMutex::new(Some(reader_task)),
219            shared,
220            next_request_id: AtomicU64::new(1),
221            initialize_state: AtomicU8::new(0),
222            max_frame_bytes: limits.max_frame_bytes,
223            max_pending_requests: limits.max_pending_requests,
224        }
225    }
226
227    pub fn is_initialized(&self) -> bool {
228        self.initialize_state.load(Ordering::Acquire) == INITIALIZED
229    }
230
231    pub fn failure(&self) -> Option<AppServerError> {
232        self.shared.failure.lock().unwrap().clone()
233    }
234
235    pub async fn next_incoming(&self) -> Option<AppServerIncoming> {
236        self.incoming_rx.lock().await.recv().await
237    }
238
239    /// Wait for `run_reader` to finish processing whatever was already
240    /// buffered on stdout. The reader terminates on its own shortly after
241    /// the child's stdout pipe closes (EOF on process exit), but a caller
242    /// classifying the exit right after `Child::wait()` returns can race
243    /// ahead of it: `child.wait()` only reports the OS-level exit, not
244    /// whether the reader has finished parsing the final bytes. Without
245    /// this, a malformed/oversized final frame can be misclassified as
246    /// `UnexpectedEof`/`NonZero`/`BeforeInitialize` instead of
247    /// `ProtocolViolation`, because `self.failure()` hadn't been set yet.
248    async fn wait_for_reader_to_finish(&self) {
249        if let Some(task) = self.reader_task.lock().await.take() {
250            let _ = task.await;
251        }
252    }
253
254    pub async fn request(
255        &self,
256        method: &str,
257        params: Value,
258        request_timeout: Duration,
259    ) -> Result<Value, AppServerError> {
260        if let Some(error) = self.failure() {
261            return Err(error);
262        }
263        if self.shared.closed.load(Ordering::Acquire) {
264            return Err(AppServerError::Closed);
265        }
266
267        let id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
268        let (response_tx, response_rx) = oneshot::channel();
269        {
270            let mut pending = self.shared.pending.lock().unwrap();
271            if pending.len() >= self.max_pending_requests {
272                return Err(AppServerError::PendingRequestLimit(
273                    self.max_pending_requests,
274                ));
275            }
276            pending.insert(id, response_tx);
277        }
278        let _pending_guard = PendingGuard {
279            shared: self.shared.clone(),
280            id,
281        };
282
283        let frame = serde_json::json!({
284            "id": id,
285            "method": method,
286            "params": params,
287        });
288
289        let operation = async {
290            self.write_value(frame).await?;
291            response_rx.await.unwrap_or(Err(AppServerError::Closed))
292        };
293
294        match tokio::time::timeout(request_timeout, operation).await {
295            Ok(result) => result,
296            Err(_) => Err(AppServerError::Timeout(request_timeout)),
297        }
298    }
299
300    pub async fn notify(
301        &self,
302        method: &str,
303        params: Value,
304        write_timeout: Duration,
305    ) -> Result<(), AppServerError> {
306        let mut frame = Map::from_iter([("method".to_string(), Value::String(method.to_string()))]);
307        if !params.is_null() {
308            frame.insert("params".to_string(), params);
309        }
310        tokio::time::timeout(write_timeout, self.write_value(Value::Object(frame)))
311            .await
312            .map_err(|_| AppServerError::Timeout(write_timeout))?
313    }
314
315    /// Reply to a server-initiated request. The caller must have an explicit
316    /// policy/UI decision; this primitive never invents an approval.
317    #[cfg(test)]
318    pub async fn respond(
319        &self,
320        id: RpcId,
321        result: Result<Value, ServerResponseError>,
322        write_timeout: Duration,
323    ) -> Result<(), AppServerError> {
324        self.respond_tracked(id, result, write_timeout, &WriteOutcomeTracker::new())
325            .await
326    }
327
328    /// Same as `respond`, but also reports via `outcome` whether it would
329    /// be safe to retry after a non-success result — including a caller
330    /// whose own future is cancelled while this is still awaiting the
331    /// write-confirmation ack, which never produces a return value at all
332    /// (see `WriteOutcomeTracker`, `write_value_tracked`). A caller that
333    /// mistakes "timed out/cancelled waiting for confirmation" for "never
334    /// sent" and retries risks a second, duplicate response for the same
335    /// id (ReAgent P1, PR #3212 6th review).
336    pub async fn respond_tracked(
337        &self,
338        id: RpcId,
339        result: Result<Value, ServerResponseError>,
340        write_timeout: Duration,
341        outcome: &WriteOutcomeTracker,
342    ) -> Result<(), AppServerError> {
343        let mut frame = Map::new();
344        frame.insert("id".to_string(), id.to_value());
345        match result {
346            Ok(result) => {
347                frame.insert("result".to_string(), result);
348            }
349            Err(error) => {
350                frame.insert(
351                    "error".to_string(),
352                    serde_json::to_value(error)
353                        .map_err(|error| AppServerError::Protocol(error.to_string()))?,
354                );
355            }
356        }
357        tokio::time::timeout(
358            write_timeout,
359            self.write_value_tracked(Value::Object(frame), outcome),
360        )
361        .await
362        .map_err(|_| AppServerError::Timeout(write_timeout))?
363    }
364
365    pub async fn initialize(&self, timeout: Duration) -> Result<Value, AppServerError> {
366        if self
367            .initialize_state
368            .compare_exchange(0, INITIALIZING, Ordering::AcqRel, Ordering::Acquire)
369            .is_err()
370        {
371            return Err(AppServerError::InitializeAlreadyAttempted);
372        }
373
374        let params = InitializeParams {
375            client_info: ClientInfo {
376                name: "agentmux",
377                title: Some("AgentMux"),
378                version: env!("CARGO_PKG_VERSION"),
379            },
380            capabilities: InitializeCapabilities {
381                experimental_api: false,
382            },
383        };
384        let params = serde_json::to_value(params)
385            .map_err(|error| AppServerError::Protocol(error.to_string()))?;
386
387        let result = match self.request("initialize", params, timeout).await {
388            Ok(result) => result,
389            Err(error) => {
390                self.initialize_state
391                    .store(INITIALIZE_FAILED, Ordering::Release);
392                return Err(error);
393            }
394        };
395        if let Err(error) = self.notify("initialized", Value::Null, timeout).await {
396            self.initialize_state
397                .store(INITIALIZE_FAILED, Ordering::Release);
398            return Err(error);
399        }
400        self.initialize_state.store(INITIALIZED, Ordering::Release);
401        Ok(result)
402    }
403
404    pub async fn shutdown(&self, timeout: Duration) -> Result<(), AppServerError> {
405        self.shared.closing.store(true, Ordering::Release);
406        let operation = async {
407            let (closed_tx, closed_rx) = oneshot::channel();
408            self.outbound_tx
409                .send(Outbound::Close { closed: closed_tx })
410                .await
411                .map_err(|_| AppServerError::Closed)?;
412            closed_rx.await.unwrap_or(Err(AppServerError::Closed))
413        };
414        let result = tokio::time::timeout(timeout, operation)
415            .await
416            .map_err(|_| AppServerError::Timeout(timeout))?;
417        fail_transport(&self.shared, AppServerError::Closed);
418        result
419    }
420
421    async fn write_value(&self, value: Value) -> Result<(), AppServerError> {
422        self.write_value_tracked(value, &WriteOutcomeTracker::new())
423            .await
424    }
425
426    /// See `WriteOutcomeTracker`. `outcome` moves to `OUTCOME_UNCONFIRMED`
427    /// the instant the frame is handed to `run_writer` — before awaiting
428    /// its write-confirmation ack — and to `CONFIRMED_UNDELIVERED` if that
429    /// ack later reports a genuine failure, rather than staying stuck at
430    /// "unconfirmed" forever for an outcome that in fact became known.
431    async fn write_value_tracked(
432        &self,
433        value: Value,
434        outcome: &WriteOutcomeTracker,
435    ) -> Result<(), AppServerError> {
436        if let Some(error) = self.failure() {
437            return Err(error);
438        }
439        if self.shared.closing.load(Ordering::Acquire) || self.shared.closed.load(Ordering::Acquire)
440        {
441            return Err(AppServerError::Closed);
442        }
443        let mut bytes = serde_json::to_vec(&value)
444            .map_err(|error| AppServerError::Protocol(error.to_string()))?;
445        if bytes.len() > self.max_frame_bytes {
446            return Err(AppServerError::Protocol(format!(
447                "outbound frame is {} bytes; limit is {}",
448                bytes.len(),
449                self.max_frame_bytes
450            )));
451        }
452        bytes.push(b'\n');
453        let (written_tx, written_rx) = oneshot::channel();
454        self.outbound_tx
455            .send(Outbound::Frame {
456                bytes,
457                written: written_tx,
458            })
459            .await
460            .map_err(|_| AppServerError::Closed)?;
461        // From here on the frame is irrevocably queued for writing by
462        // run_writer, independent of `written_rx` — dropping this future
463        // (a timeout or cancellation) cannot un-send it, so from this
464        // point a caller must not treat "didn't observe a confirmed
465        // outcome" the same as "never sent."
466        outcome.mark_queued();
467        let result = written_rx.await.unwrap_or(Err(AppServerError::Closed));
468        if result.is_err() {
469            // The writer itself confirmed delivery failed — genuinely
470            // never sent after all, so it's safe to retry again.
471            outcome.mark_confirmed_undelivered();
472        }
473        result
474    }
475}
476
477/// Tracks, for a single tracked outbound write, whether retrying after a
478/// non-success outcome risks sending the same frame twice. A plain
479/// success/failure `Result` isn't enough for this: once a frame is handed
480/// to `run_writer` it will be delivered independent of anything that
481/// happens afterward to the caller's future, including that future being
482/// dropped/cancelled before it ever observes an outcome — a case a Result
483/// alone can't represent, but which a caller's own `Drop` glue can still
484/// consult here (ReAgent P1, PR #3212 6th review).
485#[derive(Debug)]
486pub struct WriteOutcomeTracker(AtomicU8);
487
488impl WriteOutcomeTracker {
489    const NOT_QUEUED: u8 = 0;
490    const OUTCOME_UNCONFIRMED: u8 = 1;
491    const CONFIRMED_UNDELIVERED: u8 = 2;
492
493    pub fn new() -> Self {
494        Self(AtomicU8::new(Self::NOT_QUEUED))
495    }
496
497    fn mark_queued(&self) {
498        self.0.store(Self::OUTCOME_UNCONFIRMED, Ordering::Release);
499    }
500
501    fn mark_confirmed_undelivered(&self) {
502        self.0
503            .store(Self::CONFIRMED_UNDELIVERED, Ordering::Release);
504    }
505
506    /// True if the frame was never queued, or was queued and its delivery
507    /// is now confirmed to have failed — in both cases nothing was, or
508    /// ever will be, delivered, so sending again cannot produce a
509    /// duplicate. False once queued but unconfirmed, which also covers a
510    /// confirmed success (obviously never safe to retry either).
511    pub fn safe_to_retry(&self) -> bool {
512        matches!(
513            self.0.load(Ordering::Acquire),
514            Self::NOT_QUEUED | Self::CONFIRMED_UNDELIVERED
515        )
516    }
517}
518
519#[derive(Debug, Clone, Serialize)]
520pub struct ServerResponseError {
521    pub code: i64,
522    pub message: String,
523    #[serde(skip_serializing_if = "Option::is_none")]
524    pub data: Option<Value>,
525}
526
527#[derive(Serialize)]
528#[serde(rename_all = "camelCase")]
529struct InitializeParams<'a> {
530    client_info: ClientInfo<'a>,
531    capabilities: InitializeCapabilities,
532}
533
534#[derive(Serialize)]
535struct ClientInfo<'a> {
536    name: &'a str,
537    title: Option<&'a str>,
538    version: &'a str,
539}
540
541#[derive(Serialize)]
542#[serde(rename_all = "camelCase")]
543struct InitializeCapabilities {
544    experimental_api: bool,
545}
546
547async fn run_writer<W>(
548    mut writer: W,
549    mut outbound_rx: mpsc::Receiver<Outbound>,
550    shared: Arc<SharedState>,
551) where
552    W: AsyncWrite + Unpin,
553{
554    while let Some(outbound) = outbound_rx.recv().await {
555        match outbound {
556            Outbound::Frame { bytes, written } => {
557                let result = async {
558                    writer.write_all(&bytes).await?;
559                    writer.flush().await
560                }
561                .await
562                .map_err(|error: std::io::Error| AppServerError::Io(error.to_string()));
563                let _ = written.send(result.clone());
564                if let Err(error) = result {
565                    fail_transport(&shared, error);
566                    break;
567                }
568            }
569            Outbound::Close { closed } => {
570                let result = writer
571                    .shutdown()
572                    .await
573                    .map_err(|error| AppServerError::Io(error.to_string()));
574                let _ = closed.send(result.clone());
575                if let Err(error) = result {
576                    fail_transport(&shared, error);
577                }
578                break;
579            }
580        }
581    }
582    shared.closed.store(true, Ordering::Release);
583}
584
585async fn run_reader<R>(
586    reader: R,
587    incoming_tx: mpsc::Sender<AppServerIncoming>,
588    shared: Arc<SharedState>,
589    max_frame_bytes: usize,
590) where
591    R: AsyncRead + Unpin,
592{
593    let mut reader = BufReader::new(reader);
594    loop {
595        let frame = match read_frame(&mut reader, max_frame_bytes).await {
596            Ok(Some(frame)) => frame,
597            Ok(None) if shared.closing.load(Ordering::Acquire) => break,
598            Ok(None) => {
599                fail_transport(&shared, AppServerError::Closed);
600                break;
601            }
602            Err(error) => {
603                fail_transport(&shared, error);
604                break;
605            }
606        };
607
608        match classify_frame(&frame) {
609            Ok(WireMessage::Response { id, result }) => {
610                let Some(id) = id.outgoing_u64() else {
611                    tracing::warn!(?id, "ignoring App Server response with non-client id");
612                    continue;
613                };
614                let pending = shared.pending.lock().unwrap().remove(&id);
615                if let Some(pending) = pending {
616                    let _ = pending.send(result);
617                } else {
618                    tracing::warn!(
619                        id,
620                        "ignoring duplicate, unknown, or cancelled App Server response"
621                    );
622                }
623            }
624            Ok(WireMessage::Incoming(incoming)) => {
625                if incoming_tx.try_send(incoming).is_err() {
626                    fail_transport(
627                        &shared,
628                        AppServerError::Protocol(
629                            "inbound App Server event queue exceeded its bound".to_string(),
630                        ),
631                    );
632                    break;
633                }
634            }
635            Err(error) => {
636                fail_transport(&shared, error);
637                break;
638            }
639        }
640    }
641}
642
643async fn read_frame<R>(
644    reader: &mut R,
645    max_frame_bytes: usize,
646) -> Result<Option<Vec<u8>>, AppServerError>
647where
648    R: AsyncBufRead + Unpin,
649{
650    let mut frame = Vec::new();
651    loop {
652        let available = reader
653            .fill_buf()
654            .await
655            .map_err(|error| AppServerError::Io(error.to_string()))?;
656        if available.is_empty() {
657            return if frame.is_empty() {
658                Ok(None)
659            } else {
660                Ok(Some(frame))
661            };
662        }
663
664        let newline = available.iter().position(|byte| *byte == b'\n');
665        let consumed = newline.map_or(available.len(), |index| index + 1);
666        let payload_bytes = newline.map_or(consumed, |index| index);
667        if frame.len().saturating_add(payload_bytes) > max_frame_bytes {
668            return Err(AppServerError::Protocol(format!(
669                "inbound frame exceeds {max_frame_bytes} bytes"
670            )));
671        }
672        frame.extend_from_slice(&available[..payload_bytes]);
673        reader.consume(consumed);
674
675        if newline.is_some() {
676            if frame.last() == Some(&b'\r') {
677                frame.pop();
678            }
679            return Ok(Some(frame));
680        }
681    }
682}
683
684enum WireMessage {
685    Response {
686        id: RpcId,
687        result: Result<Value, AppServerError>,
688    },
689    Incoming(AppServerIncoming),
690}
691
692fn classify_frame(frame: &[u8]) -> Result<WireMessage, AppServerError> {
693    if frame.is_empty() {
694        return Err(AppServerError::Protocol(
695            "stdout contained an empty frame".to_string(),
696        ));
697    }
698    let value: Value = serde_json::from_slice(frame)
699        .map_err(|error| AppServerError::Protocol(format!("invalid JSON frame: {error}")))?;
700    let object = value
701        .as_object()
702        .ok_or_else(|| AppServerError::Protocol("JSON-RPC frame must be an object".to_string()))?;
703    validate_jsonrpc_marker(object)?;
704
705    if let Some(method) = object.get("method") {
706        let method = method.as_str().ok_or_else(|| {
707            AppServerError::Protocol("JSON-RPC method must be a string".to_string())
708        })?;
709        let params = object.get("params").cloned().unwrap_or(Value::Null);
710        return if let Some(id) = object.get("id") {
711            Ok(WireMessage::Incoming(AppServerIncoming::Request {
712                id: RpcId::from_value(id)?,
713                method: method.to_string(),
714                params,
715            }))
716        } else {
717            Ok(WireMessage::Incoming(AppServerIncoming::Notification {
718                method: method.to_string(),
719                params,
720            }))
721        };
722    }
723
724    let id = object
725        .get("id")
726        .ok_or_else(|| AppServerError::Protocol("response is missing an id".to_string()))?;
727    let id = RpcId::from_value(id)?;
728    let has_result = object.contains_key("result");
729    let has_error = object.contains_key("error");
730    if has_result == has_error {
731        return Err(AppServerError::Protocol(
732            "response must contain exactly one of result or error".to_string(),
733        ));
734    }
735    let result = if has_result {
736        Ok(object.get("result").cloned().unwrap_or(Value::Null))
737    } else {
738        Err(parse_remote_error(object.get("error").unwrap())?)
739    };
740    Ok(WireMessage::Response { id, result })
741}
742
743fn validate_jsonrpc_marker(object: &Map<String, Value>) -> Result<(), AppServerError> {
744    match object.get("jsonrpc") {
745        None => Ok(()), // Codex App Server intentionally omits the JSON-RPC header.
746        Some(Value::String(version)) if version == "2.0" => Ok(()),
747        Some(_) => Err(AppServerError::Protocol(
748            "jsonrpc, when present, must equal \"2.0\"".to_string(),
749        )),
750    }
751}
752
753fn parse_remote_error(value: &Value) -> Result<AppServerError, AppServerError> {
754    let error = value
755        .as_object()
756        .ok_or_else(|| AppServerError::Protocol("response error must be an object".to_string()))?;
757    let code = error.get("code").and_then(Value::as_i64).ok_or_else(|| {
758        AppServerError::Protocol("response error code must be an integer".to_string())
759    })?;
760    let message = error
761        .get("message")
762        .and_then(Value::as_str)
763        .ok_or_else(|| {
764            AppServerError::Protocol("response error message must be a string".to_string())
765        })?;
766    Ok(AppServerError::Remote {
767        code,
768        message: message.to_string(),
769        data: error.get("data").cloned(),
770    })
771}
772
773fn fail_transport(shared: &SharedState, error: AppServerError) {
774    let failure = {
775        let mut failure = shared.failure.lock().unwrap();
776        if failure.is_none() {
777            *failure = Some(error.clone());
778        }
779        failure.clone().unwrap()
780    };
781    shared.closed.store(true, Ordering::Release);
782    let pending = std::mem::take(&mut *shared.pending.lock().unwrap());
783    for (_, sender) in pending {
784        let _ = sender.send(Err(failure.clone()));
785    }
786}
787
788struct StderrTail {
789    bytes: VecDeque<u8>,
790    max_bytes: usize,
791}
792
793impl StderrTail {
794    fn push(&mut self, chunk: &[u8]) {
795        self.bytes.extend(chunk);
796        while self.bytes.len() > self.max_bytes {
797            self.bytes.pop_front();
798        }
799    }
800
801    fn snapshot(&self) -> String {
802        let bytes: Vec<u8> = self.bytes.iter().copied().collect();
803        String::from_utf8_lossy(&bytes).into_owned()
804    }
805}
806
807pub struct AppServerProcess {
808    // Arc, not a bare AppServerTransport: the controller runtime hands this
809    // same transport to both its own event-loop task and the
810    // CodexAppServerSession it constructs, which itself requires
811    // Arc<AppServerTransport> so it can be shared across the caller and its
812    // own notification-reading task.
813    pub transport: Arc<AppServerTransport>,
814    // A single dedicated reaper task (spawned once, below) owns `Child`
815    // exclusively for its whole lifetime and is the only thing that ever
816    // calls `child.wait()`. This exists because `wait_for_exit()` and
817    // `shutdown()` can legitimately run concurrently (one waiting
818    // passively for natural exit, the other trying to force-stop) and both
819    // need to observe the same exit — but `Child::wait()`/`Child::kill()`
820    // both require `&mut Child`, so whichever call took the lock first
821    // used to block the other for as long as the child kept running,
822    // including `shutdown()`'s own force-kill step (ReAgent P1, PR #3209).
823    // `watch` (not `Notify`/`oneshot`) specifically because it's safe for
824    // any number of callers to subscribe at any time, including after the
825    // child has already exited — a late subscriber just sees the value
826    // that's already there instead of missing a wakeup that already fired.
827    exit_status: watch::Receiver<Option<Result<ExitStatus, String>>>,
828    stderr_tail: Arc<Mutex<StderrTail>>,
829    stderr_task: AsyncMutex<Option<tokio::task::JoinHandle<()>>>,
830    shutdown_requested: AtomicBool,
831    pid: Option<u32>,
832}
833
834impl AppServerProcess {
835    pub fn spawn(mut command: Command, limits: AppServerLimits) -> Result<Self, AppServerError> {
836        command.kill_on_drop(true);
837        command.stdin(Stdio::piped());
838        command.stdout(Stdio::piped());
839        command.stderr(Stdio::piped());
840        let mut child = command
841            .spawn()
842            .map_err(|error| AppServerError::Spawn(error.to_string()))?;
843        let pid = child.id();
844        let stdin = child
845            .stdin
846            .take()
847            .ok_or_else(|| AppServerError::Spawn("stdin was not captured".to_string()))?;
848        let stdout = child
849            .stdout
850            .take()
851            .ok_or_else(|| AppServerError::Spawn("stdout was not captured".to_string()))?;
852        let mut stderr = child
853            .stderr
854            .take()
855            .ok_or_else(|| AppServerError::Spawn("stderr was not captured".to_string()))?;
856
857        let stderr_tail = Arc::new(Mutex::new(StderrTail {
858            bytes: VecDeque::new(),
859            max_bytes: limits.max_stderr_bytes,
860        }));
861        let stderr_tail_task = stderr_tail.clone();
862        let stderr_task = tokio::spawn(async move {
863            let mut buffer = [0_u8; 4096];
864            loop {
865                match stderr.read(&mut buffer).await {
866                    Ok(0) => break,
867                    Ok(read) => stderr_tail_task.lock().unwrap().push(&buffer[..read]),
868                    Err(error) => {
869                        tracing::warn!(%error, "App Server stderr read failed");
870                        break;
871                    }
872                }
873            }
874        });
875
876        // Both stdout and stderr readers exist before initialize can write.
877        let transport = Arc::new(AppServerTransport::new(stdout, stdin, limits));
878
879        let (exit_tx, exit_rx) = watch::channel(None);
880        tokio::spawn(async move {
881            let result = child
882                .wait()
883                .await
884                .map_err(|error| error.to_string());
885            // `child` (and its `kill_on_drop`) drops here, after it has
886            // already exited — a no-op, not a race with anyone else, since
887            // this task is the only place that ever touches `Child`.
888            let _ = exit_tx.send(Some(result));
889        });
890
891        Ok(Self {
892            transport,
893            exit_status: exit_rx,
894            stderr_tail,
895            stderr_task: AsyncMutex::new(Some(stderr_task)),
896            shutdown_requested: AtomicBool::new(false),
897            pid,
898        })
899    }
900
901    pub fn pid(&self) -> Option<u32> {
902        self.pid
903    }
904
905    /// Await the reaper task's exit status. Race-free for any number of
906    /// concurrent/late callers — see `exit_status`'s doc comment on why
907    /// this exists instead of each caller calling `Child::wait()` itself.
908    async fn await_exit(&self) -> Result<ExitStatus, AppServerError> {
909        let mut rx = self.exit_status.clone();
910        loop {
911            if let Some(result) = rx.borrow().clone() {
912                return result.map_err(AppServerError::Io);
913            }
914            if rx.changed().await.is_err() {
915                return Err(AppServerError::Closed);
916            }
917        }
918    }
919
920    pub async fn wait_for_exit(&self) -> Result<AppServerExit, AppServerError> {
921        let status = self.await_exit().await?;
922        self.transport.wait_for_reader_to_finish().await;
923        self.finish_stderr().await;
924        Ok(self.classify_exit(status))
925    }
926
927    pub async fn shutdown(&self, timeout: Duration) -> Result<AppServerExit, AppServerError> {
928        self.shutdown_requested.store(true, Ordering::Release);
929        let _ = self.transport.shutdown(timeout).await;
930        let status = match tokio::time::timeout(timeout, self.await_exit()).await {
931            Ok(result) => Some(result?),
932            Err(_) => {
933                // Force-kill by pid rather than `Child::kill()` — the
934                // reaper task spawned in `spawn()` owns the `Child` value
935                // exclusively for its whole lifetime (that's what fixes the
936                // concurrent-shutdown deadlock this replaces), so nothing
937                // outside that task can get `&mut Child` to call `.kill()`
938                // on it. `kill_pid` needs only the raw pid, not the handle.
939                if let Some(pid) = self.pid {
940                    let _ = agentmux_common::process::kill_pid(pid);
941                }
942                // The kill above makes the reaper's own wait() return
943                // shortly; discard the resulting status like the pre-fix
944                // code did (`ForcedAfterTimeout` below reports the timeout,
945                // not whatever exit code a killed process happens to have).
946                let _ = self.await_exit().await;
947                None
948            }
949        };
950        self.transport.wait_for_reader_to_finish().await;
951        self.finish_stderr().await;
952        match status {
953            Some(status) => Ok(self.classify_exit(status)),
954            None => Ok(AppServerExit::ForcedAfterTimeout {
955                stderr_tail: self.stderr_snapshot(),
956            }),
957        }
958    }
959
960    async fn finish_stderr(&self) {
961        if let Some(task) = self.stderr_task.lock().await.take() {
962            let _ = task.await;
963        }
964    }
965
966    fn stderr_snapshot(&self) -> String {
967        self.stderr_tail.lock().unwrap().snapshot()
968    }
969
970    fn classify_exit(&self, status: ExitStatus) -> AppServerExit {
971        let stderr_tail = self.stderr_snapshot();
972        if let Some(AppServerError::Protocol(message)) = self.transport.failure() {
973            return AppServerExit::ProtocolViolation {
974                message,
975                stderr_tail,
976            };
977        }
978        if !self.transport.is_initialized() {
979            return AppServerExit::BeforeInitialize {
980                code: status.code(),
981                stderr_tail,
982            };
983        }
984        if !status.success() {
985            return AppServerExit::NonZero {
986                code: status.code(),
987                stderr_tail,
988            };
989        }
990        if self.shutdown_requested.load(Ordering::Acquire) {
991            AppServerExit::Clean {
992                code: status.code(),
993                stderr_tail,
994            }
995        } else {
996            AppServerExit::UnexpectedEof { stderr_tail }
997        }
998    }
999}
1000
1001#[cfg(test)]
1002mod tests {
1003    use serde_json::json;
1004    use tokio::io::{duplex, AsyncBufReadExt, AsyncWriteExt, BufReader};
1005
1006    use super::*;
1007
1008    #[tokio::test]
1009    async fn correlates_out_of_order_responses_and_preserves_notifications() {
1010        let (client_io, server_io) = duplex(8192);
1011        let (client_reader, client_writer) = tokio::io::split(client_io);
1012        let (server_reader, mut server_writer) = tokio::io::split(server_io);
1013        let transport = Arc::new(AppServerTransport::new(
1014            client_reader,
1015            client_writer,
1016            AppServerLimits::default(),
1017        ));
1018
1019        let server = tokio::spawn(async move {
1020            let mut lines = BufReader::new(server_reader).lines();
1021            let first: Value =
1022                serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap();
1023            let second: Value =
1024                serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap();
1025            let (first_id, second_id) = if first["method"] == "one" {
1026                (&first["id"], &second["id"])
1027            } else {
1028                (&second["id"], &first["id"])
1029            };
1030            let frames = format!(
1031                "{{\"method\":\"thread/status/changed\",\"params\":{{\"ok\":true}}}}\r\n{{\"id\":{},\"result\":\"second\"}}\n{{\"id\":{},\"result\":\"first\"}}\n",
1032                second_id, first_id
1033            );
1034            server_writer
1035                .write_all(&frames.as_bytes()[..17])
1036                .await
1037                .unwrap();
1038            server_writer
1039                .write_all(&frames.as_bytes()[17..])
1040                .await
1041                .unwrap();
1042        });
1043
1044        let first = {
1045            let transport = transport.clone();
1046            tokio::spawn(async move {
1047                transport
1048                    .request("one", Value::Null, Duration::from_secs(1))
1049                    .await
1050            })
1051        };
1052        let second = {
1053            let transport = transport.clone();
1054            tokio::spawn(async move {
1055                transport
1056                    .request("two", Value::Null, Duration::from_secs(1))
1057                    .await
1058            })
1059        };
1060
1061        assert_eq!(first.await.unwrap().unwrap(), Value::String("first".into()));
1062        assert_eq!(
1063            second.await.unwrap().unwrap(),
1064            Value::String("second".into())
1065        );
1066        assert_eq!(
1067            transport.next_incoming().await,
1068            Some(AppServerIncoming::Notification {
1069                method: "thread/status/changed".to_string(),
1070                params: serde_json::json!({ "ok": true }),
1071            })
1072        );
1073        server.await.unwrap();
1074    }
1075
1076    #[tokio::test]
1077    async fn remote_errors_resolve_only_the_matching_request() {
1078        let (client_io, server_io) = duplex(4096);
1079        let (client_reader, client_writer) = tokio::io::split(client_io);
1080        let (server_reader, mut server_writer) = tokio::io::split(server_io);
1081        let transport =
1082            AppServerTransport::new(client_reader, client_writer, AppServerLimits::default());
1083        let server = tokio::spawn(async move {
1084            let request = BufReader::new(server_reader)
1085                .lines()
1086                .next_line()
1087                .await
1088                .unwrap()
1089                .unwrap();
1090            let request: Value = serde_json::from_str(&request).unwrap();
1091            server_writer
1092                .write_all(
1093                    format!(
1094                        "{{\"id\":{},\"error\":{{\"code\":-32001,\"message\":\"no thread\",\"data\":{{\"retry\":false}}}}}}\n",
1095                        request["id"]
1096                    )
1097                    .as_bytes(),
1098                )
1099                .await
1100                .unwrap();
1101        });
1102
1103        assert_eq!(
1104            transport
1105                .request("thread/resume", Value::Null, Duration::from_secs(1))
1106                .await,
1107            Err(AppServerError::Remote {
1108                code: -32001,
1109                message: "no thread".to_string(),
1110                data: Some(serde_json::json!({ "retry": false })),
1111            })
1112        );
1113        server.await.unwrap();
1114    }
1115
1116    #[tokio::test]
1117    async fn responds_to_server_requests_without_auto_approving_them() {
1118        let (client_io, server_io) = duplex(4096);
1119        let (client_reader, client_writer) = tokio::io::split(client_io);
1120        let (server_reader, mut server_writer) = tokio::io::split(server_io);
1121        let transport =
1122            AppServerTransport::new(client_reader, client_writer, AppServerLimits::default());
1123        server_writer
1124            .write_all(
1125                br#"{"id":"approval-1","method":"item/commandExecution/requestApproval","params":{"command":"echo hi"}}
1126"#,
1127            )
1128            .await
1129            .unwrap();
1130        let request = transport.next_incoming().await.unwrap();
1131        let AppServerIncoming::Request { id, method, .. } = request else {
1132            panic!("expected server request");
1133        };
1134        assert_eq!(method, "item/commandExecution/requestApproval");
1135        transport
1136            .respond(id, Ok(json!({"decision":"accept"})), Duration::from_secs(1))
1137            .await
1138            .unwrap();
1139        let response = BufReader::new(server_reader)
1140            .lines()
1141            .next_line()
1142            .await
1143            .unwrap()
1144            .unwrap();
1145        let response: Value = serde_json::from_str(&response).unwrap();
1146        assert_eq!(
1147            response,
1148            json!({"id":"approval-1","result":{"decision":"accept"}})
1149        );
1150    }
1151
1152    #[test]
1153    fn preserves_string_numeric_and_null_server_request_ids() {
1154        let cases = [
1155            (
1156                br#"{"id":"approval-7","method":"item/commandExecution/requestApproval"}"#
1157                    .as_slice(),
1158                RpcId::String("approval-7".to_string()),
1159            ),
1160            (
1161                br#"{"id":42,"method":"item/fileChange/requestApproval"}"#.as_slice(),
1162                RpcId::Number("42".to_string()),
1163            ),
1164            (
1165                br#"{"id":null,"method":"tool/requestUserInput"}"#.as_slice(),
1166                RpcId::Null,
1167            ),
1168        ];
1169        for (frame, expected_id) in cases {
1170            match classify_frame(frame).unwrap() {
1171                WireMessage::Incoming(AppServerIncoming::Request { id, .. }) => {
1172                    assert_eq!(id, expected_id)
1173                }
1174                _ => panic!("expected a server request"),
1175            }
1176        }
1177
1178        assert!(matches!(
1179            classify_frame(br#"{"result":{}}"#),
1180            Err(AppServerError::Protocol(message)) if message.contains("missing an id")
1181        ));
1182        assert!(matches!(
1183            classify_frame(br#"{"id":1,"result":{},"error":{"code":-1,"message":"bad"}}"#),
1184            Err(AppServerError::Protocol(message)) if message.contains("exactly one")
1185        ));
1186    }
1187
1188    #[tokio::test]
1189    async fn bounds_pending_requests_and_releases_cancelled_slots() {
1190        let limits = AppServerLimits {
1191            max_pending_requests: 1,
1192            ..AppServerLimits::default()
1193        };
1194        let (client_io, server_io) = duplex(4096);
1195        let (client_reader, client_writer) = tokio::io::split(client_io);
1196        let (server_reader, _server_writer) = tokio::io::split(server_io);
1197        let transport = Arc::new(AppServerTransport::new(
1198            client_reader,
1199            client_writer,
1200            limits,
1201        ));
1202        let (observed_tx, observed_rx) = oneshot::channel();
1203        let server = tokio::spawn(async move {
1204            let mut lines = BufReader::new(server_reader).lines();
1205            let _ = lines.next_line().await.unwrap().unwrap();
1206            let _ = observed_tx.send(());
1207            while lines.next_line().await.unwrap().is_some() {}
1208        });
1209        let first = {
1210            let transport = transport.clone();
1211            tokio::spawn(async move {
1212                transport
1213                    .request("one", Value::Null, Duration::from_secs(5))
1214                    .await
1215            })
1216        };
1217        observed_rx.await.unwrap();
1218
1219        assert_eq!(
1220            transport
1221                .request("two", Value::Null, Duration::from_secs(1))
1222                .await,
1223            Err(AppServerError::PendingRequestLimit(1))
1224        );
1225        first.abort();
1226        assert!(first.await.unwrap_err().is_cancelled());
1227        assert!(transport.shared.pending.lock().unwrap().is_empty());
1228
1229        assert_eq!(
1230            transport
1231                .request("timeout", Value::Null, Duration::from_millis(20))
1232                .await,
1233            Err(AppServerError::Timeout(Duration::from_millis(20)))
1234        );
1235        assert!(transport.shared.pending.lock().unwrap().is_empty());
1236        server.abort();
1237    }
1238
1239    #[tokio::test]
1240    async fn malformed_and_oversized_frames_fail_closed() {
1241        let limits = AppServerLimits {
1242            max_frame_bytes: 16,
1243            ..AppServerLimits::default()
1244        };
1245        let (client_io, mut server_io) = duplex(4096);
1246        let (client_reader, client_writer) = tokio::io::split(client_io);
1247        let transport = AppServerTransport::new(client_reader, client_writer, limits);
1248        server_io.write_all(b"not-json\n").await.unwrap();
1249        assert_eq!(transport.next_incoming().await, None);
1250        let error = transport
1251            .request("after-error", Value::Null, Duration::from_secs(1))
1252            .await
1253            .unwrap_err();
1254        assert!(matches!(error, AppServerError::Protocol(_)));
1255
1256        let (client_io, mut server_io) = duplex(4096);
1257        let (client_reader, client_writer) = tokio::io::split(client_io);
1258        let transport = AppServerTransport::new(client_reader, client_writer, limits);
1259        server_io
1260            .write_all(b"{\"way\":\"too long for bound\"}\n")
1261            .await
1262            .unwrap();
1263        assert_eq!(transport.next_incoming().await, None);
1264        let error = transport
1265            .request("after-error", Value::Null, Duration::from_secs(1))
1266            .await
1267            .unwrap_err();
1268        assert!(matches!(error, AppServerError::Protocol(_)));
1269    }
1270
1271    #[tokio::test]
1272    async fn initializes_and_gracefully_stops_a_fake_process() {
1273        let process = AppServerProcess::spawn(fake_command("happy"), AppServerLimits::default())
1274            .expect("spawn fake App Server");
1275        let response = process
1276            .transport
1277            .initialize(Duration::from_secs(2))
1278            .await
1279            .expect("initialize");
1280        assert_eq!(response["serverInfo"]["name"], "fake-app-server");
1281        assert!(process.transport.is_initialized());
1282        assert_eq!(
1283            process.transport.next_incoming().await,
1284            Some(AppServerIncoming::Notification {
1285                method: "server/ready".to_string(),
1286                params: serde_json::json!({ "fixture": true }),
1287            })
1288        );
1289
1290        let exit = process.shutdown(Duration::from_secs(2)).await.unwrap();
1291        match exit {
1292            AppServerExit::Clean { stderr_tail, .. } => {
1293                assert!(stderr_tail.contains("fake server started"));
1294                assert!(stderr_tail.contains("fake server observed EOF"));
1295            }
1296            other => panic!("expected clean exit, got {other:?}"),
1297        }
1298    }
1299
1300    /// ReAgent P1, PR #3209: `wait_for_exit()` used to hold the `AsyncMutex`
1301    /// guarding `child` for the entire duration of `Child::wait()`, so a
1302    /// concurrent `shutdown()` (e.g. spawn_event_loop's passive EOF-wait
1303    /// racing a caller-initiated stop) could never acquire that same lock
1304    /// to run its own force-kill step — it would just block on the lock
1305    /// forever, defeating shutdown's own timeout entirely. Every assertion
1306    /// here is wrapped in a bounded `tokio::time::timeout` specifically so
1307    /// a regression fails this test in seconds instead of hanging the
1308    /// whole suite.
1309    #[tokio::test]
1310    async fn shutdown_can_force_kill_while_wait_for_exit_is_concurrently_awaiting() {
1311        let process = Arc::new(
1312            AppServerProcess::spawn(fake_command("ignores-shutdown"), AppServerLimits::default())
1313                .expect("spawn fake App Server"),
1314        );
1315        process
1316            .transport
1317            .initialize(Duration::from_secs(2))
1318            .await
1319            .expect("initialize");
1320
1321        // Starts waiting BEFORE shutdown() is called, mirroring the real
1322        // race: something else already passively waiting when a stop comes in.
1323        let waiter_process = process.clone();
1324        let waiter = tokio::spawn(async move { waiter_process.wait_for_exit().await });
1325        tokio::time::sleep(Duration::from_millis(50)).await;
1326
1327        let shutdown_result = tokio::time::timeout(
1328            Duration::from_secs(5),
1329            process.shutdown(Duration::from_millis(200)),
1330        )
1331        .await
1332        .expect("shutdown() must not hang waiting on the same lock wait_for_exit() holds")
1333        .unwrap();
1334        assert!(
1335            matches!(shutdown_result, AppServerExit::ForcedAfterTimeout { .. }),
1336            "expected a forced kill after the 200ms grace period (this fixture ignores \
1337             graceful shutdown by design), got {shutdown_result:?}"
1338        );
1339
1340        // The waiter and shutdown() don't necessarily agree on the exit
1341        // CLASSIFICATION — shutdown() knows it forced the kill and reports
1342        // that deliberately (discarding the real exit code); wait_for_exit()
1343        // is a passive observer that just sees the killed process's actual
1344        // (non-zero) exit code and classifies it normally. What matters
1345        // here is that it resolves at all, promptly, instead of hanging —
1346        // that's the deadlock this fix removes.
1347        let waiter_result = tokio::time::timeout(Duration::from_secs(5), waiter)
1348            .await
1349            .expect("the concurrent wait_for_exit() must resolve once shutdown() kills the process")
1350            .unwrap()
1351            .unwrap();
1352        assert!(
1353            matches!(waiter_result, AppServerExit::NonZero { .. }),
1354            "expected the killed process's real non-zero exit to surface here, got {waiter_result:?}"
1355        );
1356    }
1357
1358    #[tokio::test]
1359    async fn classifies_exit_before_initialize_with_bounded_stderr() {
1360        let limits = AppServerLimits {
1361            max_stderr_bytes: 24,
1362            ..AppServerLimits::default()
1363        };
1364        let process = AppServerProcess::spawn(fake_command("exit-before-init"), limits)
1365            .expect("spawn fake App Server");
1366        let exit = process.wait_for_exit().await.unwrap();
1367        match exit {
1368            AppServerExit::BeforeInitialize { code, stderr_tail } => {
1369                assert_eq!(code, Some(23));
1370                assert!(stderr_tail.len() <= 24);
1371                assert!(stderr_tail.ends_with("before initialize\n"));
1372            }
1373            other => panic!("expected pre-initialize exit, got {other:?}"),
1374        }
1375    }
1376
1377    #[tokio::test]
1378    async fn classifies_post_initialize_nonzero_exit_and_unexpected_eof() {
1379        let process =
1380            AppServerProcess::spawn(fake_command("exit-after-init"), AppServerLimits::default())
1381                .expect("spawn fake App Server");
1382        process
1383            .transport
1384            .initialize(Duration::from_secs(2))
1385            .await
1386            .unwrap();
1387        match process.wait_for_exit().await.unwrap() {
1388            AppServerExit::NonZero { code, stderr_tail } => {
1389                assert_eq!(code, Some(27));
1390                assert!(stderr_tail.contains("failed after initialize"));
1391            }
1392            other => panic!("expected post-initialize failure, got {other:?}"),
1393        }
1394
1395        let process =
1396            AppServerProcess::spawn(fake_command("eof-after-init"), AppServerLimits::default())
1397                .expect("spawn fake App Server");
1398        process
1399            .transport
1400            .initialize(Duration::from_secs(2))
1401            .await
1402            .unwrap();
1403        assert!(matches!(
1404            process.wait_for_exit().await.unwrap(),
1405            AppServerExit::UnexpectedEof { .. }
1406        ));
1407    }
1408
1409    #[tokio::test]
1410    async fn classifies_malformed_stdout_as_protocol_violation() {
1411        let process = AppServerProcess::spawn(
1412            fake_command("protocol-violation"),
1413            AppServerLimits::default(),
1414        )
1415        .expect("spawn fake App Server");
1416        assert!(matches!(
1417            process.transport.initialize(Duration::from_secs(2)).await,
1418            Err(AppServerError::Protocol(_))
1419        ));
1420        match process.shutdown(Duration::from_secs(2)).await.unwrap() {
1421            AppServerExit::ProtocolViolation {
1422                message,
1423                stderr_tail,
1424            } => {
1425                assert!(message.contains("invalid JSON"));
1426                assert!(stderr_tail.contains("protocol violation"));
1427            }
1428            other => panic!("expected protocol violation, got {other:?}"),
1429        }
1430    }
1431
1432    fn fake_command(mode: &str) -> Command {
1433        let mut command = Command::new(fake_server_binary());
1434        command.env("AGENTMUX_FAKE_APP_SERVER_MODE", mode);
1435        command
1436    }
1437
1438    fn fake_server_binary() -> &'static std::path::PathBuf {
1439        crate::test_support::fake_app_server_binary()
1440    }
1441
1442    /// ReAgent P1, PR #3212 (6th review): a failure/timeout/cancellation
1443    /// while `write_value_tracked` is still awaiting the failure() check
1444    /// or the `outbound_tx.send(...)` step itself means the frame was
1445    /// never handed to `run_writer` at all -- genuinely safe to retry.
1446    #[tokio::test]
1447    async fn write_outcome_tracker_stays_safe_to_retry_when_the_write_never_starts() {
1448        let (client_io, _server_io) = duplex(4096);
1449        let (client_reader, client_writer) = tokio::io::split(client_io);
1450        let transport =
1451            AppServerTransport::new(client_reader, client_writer, AppServerLimits::default());
1452        fail_transport(&transport.shared, AppServerError::Closed);
1453
1454        let tracker = WriteOutcomeTracker::new();
1455        let result = transport
1456            .write_value_tracked(json!({"x": 1}), &tracker)
1457            .await;
1458        assert!(result.is_err());
1459        assert!(tracker.safe_to_retry());
1460    }
1461
1462    /// See write_outcome_tracker_stays_safe_to_retry_when_the_write_never_starts
1463    /// -- once the peer is fully gone, run_writer's actual write fails and
1464    /// reports that failure back through the confirmation ack, so the
1465    /// caller learns for certain the frame was never delivered and it
1466    /// becomes safe to retry again.
1467    #[tokio::test]
1468    async fn write_outcome_tracker_becomes_safe_to_retry_on_a_confirmed_write_failure() {
1469        let (client_io, server_io) = duplex(4096);
1470        drop(server_io);
1471        let (client_reader, client_writer) = tokio::io::split(client_io);
1472        let transport =
1473            AppServerTransport::new(client_reader, client_writer, AppServerLimits::default());
1474
1475        let tracker = WriteOutcomeTracker::new();
1476        let result = transport
1477            .write_value_tracked(json!({"x": 1}), &tracker)
1478            .await;
1479        assert!(result.is_err(), "the write should fail once the peer is gone");
1480        assert!(
1481            tracker.safe_to_retry(),
1482            "a confirmed write failure means the frame was genuinely never delivered"
1483        );
1484    }
1485
1486    /// ReAgent P1, PR #3212 (6th review): once the frame is handed to
1487    /// run_writer, it will be delivered independent of anything that
1488    /// happens to this call afterward. A caller whose future is
1489    /// cancelled while still waiting on the confirmation ack must not
1490    /// treat that the same as "never sent" -- a 1-byte duplex buffer with
1491    /// nothing reading the other side guarantees the write stays stuck
1492    /// there deterministically, rather than racing a timing window.
1493    #[tokio::test]
1494    async fn write_outcome_tracker_stays_unconfirmed_when_cancelled_after_queueing() {
1495        let (client_io, _server_io) = duplex(1);
1496        let (client_reader, client_writer) = tokio::io::split(client_io);
1497        let transport = Arc::new(AppServerTransport::new(
1498            client_reader,
1499            client_writer,
1500            AppServerLimits::default(),
1501        ));
1502        let tracker = Arc::new(WriteOutcomeTracker::new());
1503
1504        let task_transport = transport.clone();
1505        let task_tracker = tracker.clone();
1506        let handle = tokio::spawn(async move {
1507            task_transport
1508                .write_value_tracked(json!({"x": 1}), &task_tracker)
1509                .await
1510        });
1511
1512        for _ in 0..100 {
1513            if !tracker.safe_to_retry() {
1514                break;
1515            }
1516            tokio::task::yield_now().await;
1517        }
1518        assert!(
1519            !tracker.safe_to_retry(),
1520            "the frame should already be queued (unconfirmed) while the write is stuck"
1521        );
1522
1523        handle.abort();
1524        let _ = handle.await;
1525
1526        assert!(
1527            !tracker.safe_to_retry(),
1528            "a cancelled write whose outcome was never confirmed must stay unsafe to retry"
1529        );
1530    }
1531}
agentmux_srv/backend/blockcontroller/app_server.rs

agentmux_srv/backend/blockcontroller/
app_server.rs
