agentos_client/
shell.rs

1//! Network (fetch) and Shell / terminal methods + supporting types.
2//!
3//! Ported from `packages/core/src/agent-os.ts` (`fetch` + shell methods) and `runtime-compat.ts`
4//! (`ShellHandle`, `OpenShellOptions`, `ConnectTerminalOptions`).
5//!
6//! Id-vs-PID is load-bearing: `open_shell` returns a synthetic `shell-N` id; `connect_terminal`
7//! returns a PID and is NOT tracked in the shells map.
8//!
9//! The native wire protocol has no PTY/winsize request, so a shell is modeled as a guest process
10//! spawned via [`ExecuteRequest`]: its `process_id` is what `write_shell`/`close_shell` address on
11//! the wire, while the public boundary keeps the synthetic `shell-N` id.
12//!
13//! Stream routing mirrors the TS PTY path: the public `data` stream (`on_shell_data`) carries stdout
14//! and stderr in the order received from the sidecar. stderr is also delivered on an optional
15//! channel-specific diagnostic tap (`on_shell_stderr` + [`OpenShellOptions::on_stderr`]); terminal
16//! renderers consume only `data` so prompts and control sequences are neither reordered nor doubled.
17
18use std::collections::BTreeMap;
19use std::sync::atomic::{AtomicUsize, Ordering};
20
21use anyhow::Result;
22use uuid::Uuid;
23
24use agentos_sidecar_client::wire::{self, EventPayload, StreamChannel};
25
26use crate::agent_os::{AcpTerminalEntry, AgentOs, ShellEntry};
27use crate::error::ClientError;
28use crate::process::{install_output_callback, OutputCallback, ProcessStatus, StdinInput};
29use crate::stream::ByteStream;
30
31/// Channel capacity for a shell's ordered terminal-data and diagnostic-stderr broadcasts.
32const SHELL_DATA_CHANNEL_CAPACITY: usize = 1024;
33
34/// Maximum active or spawning terminals created by `connect_terminal` per VM.
35const ACP_TERMINAL_LIMIT: usize = 1024;
36
37/// Default shell command used when [`OpenShellOptions::command`] is omitted (matches the kernel's
38/// PTY-backed `sh`).
39const DEFAULT_SHELL_COMMAND: &str = "sh";
40
41// ---------------------------------------------------------------------------
42// Supporting types
43// ---------------------------------------------------------------------------
44
45/// Options for `open_shell`.
46///
47/// `on_stderr` mirrors the TS `OpenShellOptions.onStderr` raw-byte callback. It is an optional
48/// stderr-only diagnostic tap; the same bytes are already present once in ordered shell data, so a
49/// terminal renderer must not consume both surfaces.
50#[derive(Default)]
51pub struct OpenShellOptions {
52    pub command: Option<String>,
53    pub args: Vec<String>,
54    pub env: BTreeMap<String, String>,
55    pub cwd: Option<String>,
56    pub cols: Option<u16>,
57    pub rows: Option<u16>,
58    pub on_stderr: Option<OutputCallback>,
59}
60
61/// Options for `connect_terminal` (extends [`OpenShellOptions`]).
62///
63/// `on_data` mirrors the TS `ConnectTerminalOptions.onData` raw-byte callback. When omitted, TS pipes
64/// shell output to host stdout; the Rust port routes it through the shell's data subscription and
65/// requires the caller to provide the sink because there is no host-process stdio to bind to.
66#[derive(Default)]
67pub struct ConnectTerminalOptions {
68    pub base: OpenShellOptions,
69    pub on_data: Option<OutputCallback>,
70}
71
72/// The synthetic shell id returned by `open_shell` (`shell-N`, NOT a pid).
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub struct ShellHandle {
75    pub shell_id: String,
76}
77
78// ---------------------------------------------------------------------------
79// Helpers
80// ---------------------------------------------------------------------------
81
82/// Map a [`RejectedResponse`] into a [`ClientError::Kernel`] so the errno `code` survives.
83fn rejected_to_error(rejected: wire::RejectedResponse) -> ClientError {
84    ClientError::Kernel {
85        code: rejected.code,
86        message: rejected.message,
87    }
88}
89
90/// Encode a [`StdinInput`] into the wire `chunk` bytes. The wire `chunk` field is bare `data`
91/// (`Vec<u8>`), so raw Binary stdin is carried verbatim (no lossy UTF-8 conversion), matching the
92/// byte-exact TS `proc.writeStdin` contract.
93fn stdin_chunk(data: StdinInput) -> Vec<u8> {
94    match data {
95        StdinInput::Text(text) => text.into_bytes(),
96        StdinInput::Bytes(bytes) => bytes,
97    }
98}
99
100fn try_reserve_counter(counter: &AtomicUsize, limit: usize) -> bool {
101    counter
102        .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
103            (count < limit).then_some(count + 1)
104        })
105        .is_ok()
106}
107
108fn release_counter(counter: &AtomicUsize) {
109    let _ = counter.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
110        Some(count.saturating_sub(1))
111    });
112}
113
114struct AcpTerminalReservation<'a> {
115    agent: &'a AgentOs,
116    active: bool,
117}
118
119impl<'a> AcpTerminalReservation<'a> {
120    fn new(agent: &'a AgentOs) -> std::result::Result<Self, ClientError> {
121        if !try_reserve_counter(&agent.inner().acp_terminal_count, ACP_TERMINAL_LIMIT) {
122            return Err(ClientError::Sidecar(format!(
123                "acp terminal limit exceeded: at most {ACP_TERMINAL_LIMIT} terminals can be active per VM"
124            )));
125        }
126        Ok(Self {
127            agent,
128            active: true,
129        })
130    }
131
132    fn disarm(&mut self) {
133        self.active = false;
134    }
135}
136
137impl Drop for AcpTerminalReservation<'_> {
138    fn drop(&mut self) {
139        if self.active {
140            release_counter(&self.agent.inner().acp_terminal_count);
141        }
142    }
143}
144
145impl AgentOs {
146    /// The VM-scoped ownership scope used for every shell/fetch wire request.
147    fn vm_ownership(&self) -> wire::OwnershipScope {
148        wire::OwnershipScope::VmOwnership(wire::VmOwnership {
149            connection_id: self.connection_id().to_string(),
150            session_id: self.wire_session_id().to_string(),
151            vm_id: self.vm_id().to_string(),
152        })
153    }
154
155    pub(crate) fn finish_acp_terminal(&self, process_id: &str) {
156        if self.inner().acp_terminals.remove(process_id).is_some() {
157            release_counter(&self.inner().acp_terminal_count);
158        }
159    }
160
161    async fn start_acp_terminal(
162        &self,
163        execute: wire::ExecuteRequest,
164        ownership: wire::OwnershipScope,
165        pid_tx: tokio::sync::oneshot::Sender<std::result::Result<u32, ClientError>>,
166        process_id: &str,
167    ) -> Option<u32> {
168        {
169            let _terminal_lifecycle_guard = self.inner().acp_terminal_lifecycle_lock.lock().await;
170            if self.inner().disposed.load(Ordering::SeqCst) {
171                let error = ClientError::Sidecar(
172                    "cannot connect terminal after VM shutdown has started".to_string(),
173                );
174                let _ = pid_tx.send(Err(error));
175                self.finish_acp_terminal(process_id);
176                return None;
177            }
178        }
179
180        let result = match self
181            .transport()
182            .request_wire(ownership, wire::RequestPayload::ExecuteRequest(execute))
183            .await
184        {
185            Ok(wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
186                pid,
187                ..
188            })) => pid.ok_or_else(|| {
189                ClientError::Sidecar("connect_terminal: sidecar did not return a pid".to_string())
190            }),
191            Ok(wire::ResponsePayload::RejectedResponse(rejected)) => {
192                Err(rejected_to_error(rejected))
193            }
194            Ok(other) => Err(ClientError::Sidecar(format!(
195                "unexpected response to connect_terminal: {other:?}"
196            ))),
197            Err(error) => Err(error.into()),
198        };
199
200        match result {
201            Ok(pid) => {
202                let _ = pid_tx.send(Ok(pid));
203                Some(pid)
204            }
205            Err(error) => {
206                let _ = pid_tx.send(Err(error));
207                self.finish_acp_terminal(process_id);
208                None
209            }
210        }
211    }
212}
213
214// ---------------------------------------------------------------------------
215// Shell / terminal
216// ---------------------------------------------------------------------------
217//
218// Note: `fetch` (the Network half of this reference section) is scaffolded in `net.rs`, which owns
219// the `impl AgentOs { fn fetch }` block. It is intentionally NOT defined here to avoid a duplicate
220// definition; the helpers below (`rejected_to_error`, `vm_ownership`) are shared by both halves.
221
222impl AgentOs {
223    /// Open a PTY-backed shell. SYNC. Returns a synthetic `shell-N` id (NOT a pid).
224    ///
225    /// The shell id and its registry entry are allocated synchronously (matching the TS sync
226    /// contract); the actual guest-process spawn, output fan-out, and exit-task registration happen
227    /// on a background task because the wire spawn is async. The exit task is tracked in the
228    /// pending-shell-exit set so `dispose` can drain it (two-phase teardown).
229    ///
230    /// Stdout and stderr are fanned into the shell's ordered `data` broadcast (`on_shell_data`).
231    /// Stderr is also fanned into a dedicated diagnostic broadcast (`on_shell_stderr` and the
232    /// [`OpenShellOptions::on_stderr`] callback); terminal renderers should consume only `data`.
233    pub fn open_shell(&self, mut options: OpenShellOptions) -> Result<ShellHandle> {
234        let inner = self.inner();
235        let counter = inner.shell_counter.fetch_add(1, Ordering::SeqCst) + 1;
236        let shell_id = format!("shell-{counter}");
237        // The wire-side process id used by write_shell/close_shell and event routing.
238        let process_id = format!("shell-{}", Uuid::new_v4());
239
240        let (data_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
241        let (stderr_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
242        // Spawn-readiness gate: write/close await this before issuing their wire request.
243        let (spawned_tx, _) = tokio::sync::watch::channel(false);
244        // Exit-code channel backing `wait_shell`.
245        let (exit_tx, _) = tokio::sync::watch::channel(None::<i32>);
246
247        // Seed any caller-provided initial stderr callback into the stderr fan-out, matching the TS
248        // initial-handler-set behavior (`stderrHandlers.add(options.onStderr)`).
249        if let Some(cb) = options.on_stderr.take() {
250            install_output_callback(stderr_tx.clone(), cb);
251        }
252
253        // Register the entry up front so write/resize/close can address it immediately, exactly like
254        // the TS map insert before the handle's async work settles.
255        let entry = ShellEntry {
256            pid: 0,
257            data_tx: data_tx.clone(),
258            stderr_tx: stderr_tx.clone(),
259            process_id: process_id.clone(),
260            spawned_tx: spawned_tx.clone(),
261            exit_tx: exit_tx.clone(),
262        };
263        // `insert` fails only if the key already exists; the monotonic counter guarantees it cannot.
264        let _ = inner.shells.insert(shell_id.clone(), entry);
265
266        let command = options
267            .command
268            .clone()
269            .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
270        options
271            .env
272            .insert(String::from("AGENTOS_EXEC_TTY"), String::from("1"));
273        // Seed the PTY winsize env exactly like the TS openShell (COLUMNS/LINES).
274        if let Some(cols) = options.cols {
275            options
276                .env
277                .insert(String::from("COLUMNS"), cols.to_string());
278        }
279        if let Some(rows) = options.rows {
280            options.env.insert(String::from("LINES"), rows.to_string());
281        }
282        let execute = wire::ExecuteRequest {
283            process_id: process_id.clone(),
284            command: Some(command),
285            runtime: None,
286            entrypoint: None,
287            args: options.args.clone(),
288            env: options.env.clone().into_iter().collect(),
289            cwd: options.cwd.clone(),
290            wasm_permission_tier: None,
291        };
292
293        // Background: subscribe to events first (so no output is missed), issue the spawn, fan
294        // stdout into the data broadcast and stderr into the stderr broadcast, and complete when the
295        // process exits.
296        let agent = self.clone();
297        let ownership = self.vm_ownership();
298        let route_process_id = process_id.clone();
299        let exit_shell_id = shell_id.clone();
300        let exit_key = counter;
301        let handle = tokio::spawn(async move {
302            let mut events = agent.transport().subscribe_wire_events();
303
304            let response = match agent
305                .transport()
306                .request_wire(
307                    ownership.clone(),
308                    wire::RequestPayload::ExecuteRequest(execute),
309                )
310                .await
311            {
312                Ok(response) => response,
313                Err(error) => {
314                    tracing::warn!(?error, shell_id = %exit_shell_id, "open_shell spawn failed");
315                    // Drop the dead entry so later shell calls report ShellNotFound rather than hang.
316                    agent.inner().shells.remove(&exit_shell_id);
317                    agent.inner().pending_shell_exits.remove(&exit_key);
318                    return;
319                }
320            };
321
322            // Record the real kernel pid on the entry (TS `ShellHandle.pid`) and release the write
323            // gate so any queued `write_shell`/`close_shell` proceed against the live spawn.
324            if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
325                pid: Some(pid),
326                ..
327            }) = response
328            {
329                agent
330                    .inner()
331                    .shells
332                    .update(&exit_shell_id, |_, existing| existing.pid = pid);
333            }
334            // send_replace, not send: `watch::Sender::send` REFUSES to store the
335            // value while no receiver exists (and the initial receiver is dropped
336            // at channel creation), which left the spawn gate permanently false
337            // for any write/resize issued after this point — they hung forever in
338            // wait_for_spawn. send_replace stores unconditionally.
339            let _ = spawned_tx.send_replace(true);
340
341            loop {
342                let (_scope, payload) = match events.recv().await {
343                    Ok(value) => value,
344                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
345                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
346                };
347                match payload {
348                    EventPayload::ProcessOutputEvent(output) => {
349                        if output.process_id != route_process_id {
350                            continue;
351                        }
352                        // Publish every PTY chunk from this single wire-event consumer so terminal
353                        // control sequences retain their original stdout/stderr order.
354                        let _ = data_tx.send(output.chunk.clone());
355                        if output.channel == StreamChannel::Stderr {
356                            // Channel identity remains available as an optional diagnostic tap.
357                            let _ = stderr_tx.send(output.chunk);
358                        }
359                    }
360                    EventPayload::ProcessExitedEvent(exited) => {
361                        if exited.process_id == route_process_id {
362                            // Record the exit code for `wait_shell`: live waiters observe the watch
363                            // update; late waiters (after the entry is dropped below) find it in the
364                            // bounded retention map, mirroring the TS closed-shell retention.
365                            {
366                                let mut retained = agent.inner().closed_shell_exit_codes.lock();
367                                retained.push_back((exit_shell_id.clone(), exited.exit_code));
368                                while retained.len() > crate::CLOSED_SHELL_EXIT_CODE_RETENTION_LIMIT
369                                {
370                                    retained.pop_front();
371                                }
372                            }
373                            let _ = exit_tx.send(Some(exited.exit_code));
374                            break;
375                        }
376                    }
377                    EventPayload::VmLifecycleEvent(_)
378                    | EventPayload::StructuredEvent(_)
379                    | EventPayload::ExtEnvelope(_) => {}
380                }
381            }
382
383            // The `.finally` equivalent: remove from both the tracking set and the shells map (only
384            // if it is still our entry, matching the TS identity check).
385            agent.inner().pending_shell_exits.remove(&exit_key);
386            agent.inner().shells.remove_if(&exit_shell_id, |existing| {
387                existing.process_id == route_process_id
388            });
389            // remove_if takes `&mut V`; the comparison only reads, which is fine.
390        });
391
392        let _ = inner.pending_shell_exits.insert(counter, handle);
393
394        Ok(ShellHandle { shell_id })
395    }
396
397    /// Open a PTY-backed terminal for the ACP `terminal/create` host request. Like [`open_shell`] it
398    /// registers a `shell-N` entry (so `write_shell`/`resize_shell`/`close_shell` address it), but the
399    /// background fan-out also (a) appends every stdout/stderr chunk to the caller's output buffer via
400    /// `on_output`, and (b) records the process exit code into `exit_tx` so `terminal/output` and
401    /// `terminal/wait_for_exit` can observe it. Mirrors the TS `_handleAcpCreateTerminal`, which builds
402    /// the terminal on top of `openShell` and tracks `output` / `exitCode` / `waitPromise`.
403    pub(crate) fn acp_open_terminal(
404        &self,
405        options: OpenShellOptions,
406        exit_tx: tokio::sync::watch::Sender<Option<i32>>,
407        on_output: impl Fn(&[u8]) + Send + Sync + 'static,
408    ) -> Result<ShellHandle> {
409        let inner = self.inner();
410        let counter = inner.shell_counter.fetch_add(1, Ordering::SeqCst) + 1;
411        let shell_id = format!("shell-{counter}");
412        let process_id = format!("shell-{}", Uuid::new_v4());
413
414        let (data_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
415        let (stderr_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
416        let (spawned_tx, _) = tokio::sync::watch::channel(false);
417
418        let entry = ShellEntry {
419            pid: 0,
420            data_tx: data_tx.clone(),
421            stderr_tx: stderr_tx.clone(),
422            process_id: process_id.clone(),
423            spawned_tx: spawned_tx.clone(),
424            // The caller-supplied exit channel doubles as the entry's `wait_shell` source.
425            exit_tx: exit_tx.clone(),
426        };
427        let _ = inner.shells.insert(shell_id.clone(), entry);
428
429        let command = options
430            .command
431            .clone()
432            .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
433        let execute = wire::ExecuteRequest {
434            process_id: process_id.clone(),
435            command: Some(command),
436            runtime: None,
437            entrypoint: None,
438            args: options.args.clone(),
439            env: options.env.clone().into_iter().collect(),
440            cwd: options.cwd.clone(),
441            wasm_permission_tier: None,
442        };
443
444        let agent = self.clone();
445        let ownership = self.vm_ownership();
446        let route_process_id = process_id.clone();
447        let exit_shell_id = shell_id.clone();
448        let exit_key = counter;
449        let on_output = std::sync::Arc::new(on_output);
450        let handle = tokio::spawn(async move {
451            let mut events = agent.transport().subscribe_wire_events();
452
453            let response = match agent
454                .transport()
455                .request_wire(
456                    ownership.clone(),
457                    wire::RequestPayload::ExecuteRequest(execute),
458                )
459                .await
460            {
461                Ok(response) => response,
462                Err(error) => {
463                    tracing::warn!(?error, shell_id = %exit_shell_id, "acp_open_terminal spawn failed");
464                    agent.inner().shells.remove(&exit_shell_id);
465                    agent.inner().pending_shell_exits.remove(&exit_key);
466                    let _ = exit_tx.send(Some(1));
467                    return;
468                }
469            };
470
471            if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
472                pid: Some(pid),
473                ..
474            }) = response
475            {
476                agent
477                    .inner()
478                    .shells
479                    .update(&exit_shell_id, |_, existing| existing.pid = pid);
480            }
481            // send_replace, not send: `watch::Sender::send` REFUSES to store the
482            // value while no receiver exists (and the initial receiver is dropped
483            // at channel creation), which left the spawn gate permanently false
484            // for any write/resize issued after this point — they hung forever in
485            // wait_for_spawn. send_replace stores unconditionally.
486            let _ = spawned_tx.send_replace(true);
487
488            let mut exit_code: i32 = 0;
489            loop {
490                let (_scope, payload) = match events.recv().await {
491                    Ok(value) => value,
492                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
493                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
494                };
495                match payload {
496                    EventPayload::ProcessOutputEvent(output) => {
497                        if output.process_id != route_process_id {
498                            continue;
499                        }
500                        let _ = data_tx.send(output.chunk.clone());
501                        if output.channel == StreamChannel::Stderr {
502                            let _ = stderr_tx.send(output.chunk.clone());
503                        }
504                        // Both channels are appended exactly once to the terminal output buffer.
505                        on_output(&output.chunk);
506                    }
507                    EventPayload::ProcessExitedEvent(exited) => {
508                        if exited.process_id == route_process_id {
509                            exit_code = exited.exit_code;
510                            break;
511                        }
512                    }
513                    EventPayload::VmLifecycleEvent(_)
514                    | EventPayload::StructuredEvent(_)
515                    | EventPayload::ExtEnvelope(_) => {}
516                }
517            }
518
519            agent.inner().pending_shell_exits.remove(&exit_key);
520            agent.inner().shells.remove_if(&exit_shell_id, |existing| {
521                existing.process_id == route_process_id
522            });
523            let _ = exit_tx.send(Some(exit_code));
524        });
525
526        // The fan-out/exit task is tracked in `pending_shell_exits` (drained by `dispose`), exactly
527        // like `open_shell`. It ends naturally when the process exits or is killed via
528        // `close_shell` / `acp_kill_terminal_shell`.
529        let _ = inner.pending_shell_exits.insert(counter, handle);
530        Ok(ShellHandle { shell_id })
531    }
532
533    /// Kill the backing process of an ACP terminal shell (SIGTERM), without removing the shell entry
534    /// or the host-terminal registry entry. Used by `terminal/kill`, which (unlike `close_shell` /
535    /// `terminal/release`) leaves the terminal addressable for output/exit queries afterward.
536    pub(crate) fn acp_kill_terminal_shell(
537        &self,
538        shell_id: &str,
539    ) -> std::result::Result<(), ClientError> {
540        let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
541        let agent = self.clone();
542        let ownership = self.vm_ownership();
543        tokio::spawn(async move {
544            wait_for_spawn(spawned_rx).await;
545            let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
546                process_id,
547                signal: String::from("SIGTERM"),
548            });
549            if let Err(error) = agent.transport().request_wire(ownership, payload).await {
550                tracing::warn!(?error, "acp_kill_terminal_shell failed");
551            }
552        });
553        Ok(())
554    }
555
556    /// Connect a terminal bound to host stdio. Returns a PID. NOT tracked in the shells map; cannot
557    /// be addressed by other shell methods. Killed during dispose via the ACP-terminal registry.
558    ///
559    /// Mirrors the TS `connectTerminal`, which routes its `onData`/`onStderr` callbacks through
560    /// `openShell`. The Rust port opens a shell, wires the caller's `on_data` to ordered terminal data
561    /// and `on_stderr` to the optional diagnostic tap, then returns the shell's pid. Host
562    /// stdin binding, terminal raw-mode, and SIGWINCH/resize forwarding are host-process concerns
563    /// that have no native wire op and are intentionally not bound here.
564    pub async fn connect_terminal(&self, options: ConnectTerminalOptions) -> Result<u32> {
565        let ConnectTerminalOptions { base, on_data } = options;
566
567        let process_id = format!("terminal-{}", Uuid::new_v4());
568        let command = base
569            .command
570            .clone()
571            .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
572        let (data_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
573        let (stderr_tx, _) =
574            tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
575
576        // onData defaults to host stdout in TS; the Rust port has no host process stdout to bind to,
577        // so it only fans out when a sink is supplied. onStderr is diagnostic and independent.
578        if let Some(cb) = on_data {
579            install_output_callback(data_tx.clone(), cb);
580        }
581        if let Some(cb) = base.on_stderr {
582            install_output_callback(stderr_tx.clone(), cb);
583        }
584
585        let execute = wire::ExecuteRequest {
586            process_id: process_id.clone(),
587            command: Some(command),
588            runtime: None,
589            entrypoint: None,
590            args: base.args.clone(),
591            env: base.env.clone().into_iter().collect(),
592            cwd: base.cwd.clone(),
593            wasm_permission_tier: None,
594        };
595
596        // Subscribe before issuing the spawn so no output is missed.
597        let events = self.transport().subscribe_wire_events();
598        let ownership = self.vm_ownership();
599        let (pid_tx, pid_rx) = tokio::sync::oneshot::channel();
600        let (start_tx, start_rx) = tokio::sync::oneshot::channel::<()>();
601        let agent = self.clone();
602        let route_process_id = process_id.clone();
603        let exit_task = tokio::spawn(async move {
604            if start_rx.await.is_err() {
605                return;
606            }
607            let terminal_pid = match agent
608                .start_acp_terminal(execute, ownership, pid_tx, &route_process_id)
609                .await
610            {
611                Some(pid) => pid,
612                None => return,
613            };
614            let mut events = events;
615            loop {
616                let (_scope, payload) = match events.recv().await {
617                    Ok(value) => value,
618                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
619                        if terminal_process_finished(&agent, terminal_pid).await {
620                            break;
621                        }
622                        continue;
623                    }
624                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
625                };
626                match payload {
627                    EventPayload::ProcessOutputEvent(output) => {
628                        if output.process_id != route_process_id {
629                            continue;
630                        }
631                        let _ = data_tx.send(output.chunk.clone());
632                        if output.channel == StreamChannel::Stderr {
633                            let _ = stderr_tx.send(output.chunk);
634                        }
635                    }
636                    EventPayload::ProcessExitedEvent(exited) => {
637                        if exited.process_id == route_process_id {
638                            break;
639                        }
640                    }
641                    EventPayload::VmLifecycleEvent(_)
642                    | EventPayload::StructuredEvent(_)
643                    | EventPayload::ExtEnvelope(_) => {}
644                }
645            }
646            agent.finish_acp_terminal(&route_process_id);
647        });
648
649        {
650            let _terminal_lifecycle_guard = self.inner().acp_terminal_lifecycle_lock.lock().await;
651            if self.inner().disposed.load(Ordering::SeqCst) {
652                exit_task.abort();
653                return Err(ClientError::Sidecar(
654                    "cannot connect terminal after VM shutdown has started".to_string(),
655                )
656                .into());
657            }
658            let mut terminal_reservation = AcpTerminalReservation::new(self)?;
659            match self
660                .inner()
661                .acp_terminals
662                .insert(process_id.clone(), AcpTerminalEntry { exit_task })
663            {
664                Ok(()) => {}
665                Err((_, entry)) => {
666                    entry.exit_task.abort();
667                    return Err(ClientError::Sidecar(format!(
668                        "terminal process id collision while tracking ACP terminal: {process_id}"
669                    ))
670                    .into());
671                }
672            }
673            terminal_reservation.disarm();
674            if start_tx.send(()).is_err() {
675                self.finish_acp_terminal(&process_id);
676                return Err(ClientError::Sidecar(
677                    "terminal startup task ended before registration completed".to_string(),
678                )
679                .into());
680            }
681        }
682
683        pid_rx
684            .await
685            .map_err(|_| {
686                ClientError::Sidecar(
687                    "terminal startup task ended before returning a pid".to_string(),
688                )
689            })?
690            .map_err(Into::into)
691    }
692
693    /// Write to a shell. SYNC fire-and-forget. Errors with [`ClientError::ShellNotFound`].
694    pub fn write_shell(
695        &self,
696        shell_id: &str,
697        data: StdinInput,
698    ) -> std::result::Result<(), ClientError> {
699        let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
700        let chunk = stdin_chunk(data);
701
702        // Fire-and-forget: the TS handle.write returns void; surface only the synchronous
703        // ShellNotFound, and dispatch the wire write in the background after the spawn lands. TS
704        // openShell is fully synchronous so the spawn is always live by the time write runs; awaiting
705        // the readiness gate reproduces that ordering and avoids dropping early input.
706        let agent = self.clone();
707        let ownership = self.vm_ownership();
708        tokio::spawn(async move {
709            wait_for_spawn(spawned_rx).await;
710            let payload = wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
711                process_id,
712                chunk,
713            });
714            if let Err(error) = agent.transport().request_wire(ownership, payload).await {
715                tracing::warn!(?error, "write_shell failed");
716            }
717        });
718
719        Ok(())
720    }
721
722    /// Write to a shell and AWAIT the wire write. Same routing as [`Self::write_shell`], but the
723    /// caller observes wire failures instead of a fire-and-forget warn — used by the actor plugin's
724    /// `writeShell` action so a failed write rejects the action.
725    pub async fn write_shell_awaited(
726        &self,
727        shell_id: &str,
728        data: StdinInput,
729    ) -> std::result::Result<(), ClientError> {
730        let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
731        let chunk = stdin_chunk(data);
732        tracing::debug!(shell_id, "write_shell_awaited: waiting for spawn gate");
733        wait_for_spawn(spawned_rx).await;
734        tracing::debug!(shell_id, "write_shell_awaited: issuing wire write");
735        let payload =
736            wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest { process_id, chunk });
737        let response = self
738            .transport()
739            .request_wire(self.vm_ownership(), payload)
740            .await?;
741        tracing::debug!(shell_id, "write_shell_awaited: wire write acked");
742        match response {
743            wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
744            _ => Ok(()),
745        }
746    }
747
748    /// Subscribe to a shell's ordered terminal data. SYNC register; multi-handler; dropping the
749    /// returned stream is the unsubscribe. Carries stdout and stderr exactly once in wire order.
750    /// Use [`Self::on_shell_stderr`] only as a channel-specific diagnostic tap, not as a second
751    /// terminal-rendering stream. Errors with [`ClientError::ShellNotFound`].
752    pub fn on_shell_data(&self, shell_id: &str) -> std::result::Result<ByteStream, ClientError> {
753        self.inner()
754            .shells
755            .read(shell_id, |_, entry| entry.data_tx.subscribe())
756            .map(ByteStream::new)
757            .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
758    }
759
760    /// Subscribe to a shell's stderr. SYNC register; multi-handler; dropping the returned stream is
761    /// the unsubscribe. This is the optional diagnostic channel backing the TS `onStderr` option;
762    /// stderr is also present once in ordered `on_shell_data`. Errors with
763    /// [`ClientError::ShellNotFound`].
764    pub fn on_shell_stderr(&self, shell_id: &str) -> std::result::Result<ByteStream, ClientError> {
765        self.inner()
766            .shells
767            .read(shell_id, |_, entry| entry.stderr_tx.subscribe())
768            .map(ByteStream::new)
769            .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
770    }
771
772    /// Resize a shell's PTY winsize. SYNC fire-and-forget, mirroring the TS `ShellHandle.resize`
773    /// (which dispatches `resizePty` in the background after the spawn lands). Errors with
774    /// [`ClientError::ShellNotFound`].
775    pub fn resize_shell(
776        &self,
777        shell_id: &str,
778        cols: u16,
779        rows: u16,
780    ) -> std::result::Result<(), ClientError> {
781        // Existence check matches the TS `if (!entry) throw Shell not found`.
782        let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
783
784        let agent = self.clone();
785        let ownership = self.vm_ownership();
786        tokio::spawn(async move {
787            wait_for_spawn(spawned_rx).await;
788            let payload = wire::RequestPayload::ResizePtyRequest(wire::ResizePtyRequest {
789                process_id,
790                cols,
791                rows,
792            });
793            if let Err(error) = agent.transport().request_wire(ownership, payload).await {
794                tracing::warn!(?error, "resize_shell failed");
795            }
796        });
797
798        Ok(())
799    }
800
801    /// Wait for a shell to exit and return its process exit code (TS `waitShell`). Resolves
802    /// immediately for a shell that already exited within the bounded retention window. Errors with
803    /// [`ClientError::ShellNotFound`] for an unknown id.
804    pub async fn wait_shell(&self, shell_id: &str) -> std::result::Result<i32, ClientError> {
805        let exit_rx = self
806            .inner()
807            .shells
808            .read(shell_id, |_, entry| entry.exit_tx.subscribe());
809        let Some(mut exit_rx) = exit_rx else {
810            // Entry already dropped: fall back to the recorded exit code (TS retention behavior).
811            let retained = self.inner().closed_shell_exit_codes.lock();
812            return retained
813                .iter()
814                .rev()
815                .find(|(id, _)| id == shell_id)
816                .map(|(_, code)| *code)
817                .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()));
818        };
819        loop {
820            if let Some(code) = *exit_rx.borrow_and_update() {
821                return Ok(code);
822            }
823            if exit_rx.changed().await.is_err() {
824                // Sender dropped without publishing a code (spawn failure / teardown): check the
825                // retention map once more before reporting the shell unknown.
826                let retained = self.inner().closed_shell_exit_codes.lock();
827                return retained
828                    .iter()
829                    .rev()
830                    .find(|(id, _)| id == shell_id)
831                    .map(|(_, code)| *code)
832                    .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()));
833            }
834        }
835    }
836
837    /// Close a shell. SYNC. `kill()` + immediate map delete; the exit task is still drained by
838    /// `dispose`. Errors with [`ClientError::ShellNotFound`].
839    pub fn close_shell(&self, shell_id: &str) -> std::result::Result<(), ClientError> {
840        let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
841
842        // Immediate map delete, exactly like the TS `_shells.delete(shellId)`; the pending-exit task
843        // remains tracked so `dispose` still drains it (two-phase teardown).
844        self.inner().shells.remove(shell_id);
845
846        // Fire-and-forget kill (SIGTERM) after the spawn lands so the kill addresses a live process.
847        let agent = self.clone();
848        let ownership = self.vm_ownership();
849        tokio::spawn(async move {
850            wait_for_spawn(spawned_rx).await;
851            let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
852                process_id,
853                signal: String::from("SIGTERM"),
854            });
855            if let Err(error) = agent.transport().request_wire(ownership, payload).await {
856                tracing::warn!(?error, "close_shell kill failed");
857            }
858        });
859
860        Ok(())
861    }
862
863    /// Look up the wire-side `process_id` and the spawn-readiness receiver for a shell id, or
864    /// [`ClientError::ShellNotFound`].
865    fn shell_wire_handle(
866        &self,
867        shell_id: &str,
868    ) -> std::result::Result<(String, tokio::sync::watch::Receiver<bool>), ClientError> {
869        self.inner()
870            .shells
871            .read(shell_id, |_, entry| {
872                (entry.process_id.clone(), entry.spawned_tx.subscribe())
873            })
874            .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
875    }
876}
877
878/// Wait until the shell's background `Execute` request has been acked (the readiness gate flips to
879/// `true`). Returns immediately if it is already ready or the sender has dropped.
880async fn wait_for_spawn(mut spawned_rx: tokio::sync::watch::Receiver<bool>) {
881    if *spawned_rx.borrow() {
882        return;
883    }
884    while spawned_rx.changed().await.is_ok() {
885        if *spawned_rx.borrow() {
886            return;
887        }
888    }
889}
890
891async fn terminal_process_finished(agent: &AgentOs, pid: u32) -> bool {
892    match agent.all_processes().await {
893        Ok(processes) => match processes.into_iter().find(|process| process.pid == pid) {
894            Some(process) => process.status != ProcessStatus::Running,
895            None => true,
896        },
897        Err(error) => {
898            tracing::warn!(?error, pid, "terminal process snapshot failed");
899            false
900        }
901    }
902}
903
904#[cfg(test)]
905mod tests {
906    use super::*;
907
908    #[test]
909    fn reserve_counter_enforces_limit_and_release_reopens_slot() {
910        let counter = AtomicUsize::new(0);
911
912        assert!(try_reserve_counter(&counter, 2));
913        assert!(try_reserve_counter(&counter, 2));
914        assert!(!try_reserve_counter(&counter, 2));
915        release_counter(&counter);
916        assert!(try_reserve_counter(&counter, 2));
917    }
918}
agentos_client/shell.rs

agentos_client/
shell.rs
