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

agentos_client/
shell.rs
