use std::collections::BTreeMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use anyhow::Result;
use uuid::Uuid;
use secure_exec_client::wire::{self, EventPayload, StreamChannel};
use crate::agent_os::{AcpTerminalEntry, AgentOs, ShellEntry};
use crate::error::ClientError;
use crate::process::{install_output_callback, OutputCallback, ProcessStatus, StdinInput};
use crate::stream::ByteStream;
const SHELL_DATA_CHANNEL_CAPACITY: usize = 1024;
const ACP_TERMINAL_LIMIT: usize = 1024;
const DEFAULT_SHELL_COMMAND: &str = "sh";
#[derive(Default)]
pub struct OpenShellOptions {
pub command: Option<String>,
pub args: Vec<String>,
pub env: BTreeMap<String, String>,
pub cwd: Option<String>,
pub cols: Option<u16>,
pub rows: Option<u16>,
pub on_stderr: Option<OutputCallback>,
}
#[derive(Default)]
pub struct ConnectTerminalOptions {
pub base: OpenShellOptions,
pub on_data: Option<OutputCallback>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ShellHandle {
pub shell_id: String,
}
fn rejected_to_error(rejected: wire::RejectedResponse) -> ClientError {
ClientError::Kernel {
code: rejected.code,
message: rejected.message,
}
}
fn stdin_chunk(data: StdinInput) -> Vec<u8> {
match data {
StdinInput::Text(text) => text.into_bytes(),
StdinInput::Bytes(bytes) => bytes,
}
}
fn try_reserve_counter(counter: &AtomicUsize, limit: usize) -> bool {
counter
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
(count < limit).then_some(count + 1)
})
.is_ok()
}
fn release_counter(counter: &AtomicUsize) {
let _ = counter.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
Some(count.saturating_sub(1))
});
}
struct AcpTerminalReservation<'a> {
agent: &'a AgentOs,
active: bool,
}
impl<'a> AcpTerminalReservation<'a> {
fn new(agent: &'a AgentOs) -> std::result::Result<Self, ClientError> {
if !try_reserve_counter(&agent.inner().acp_terminal_count, ACP_TERMINAL_LIMIT) {
return Err(ClientError::Sidecar(format!(
"acp terminal limit exceeded: at most {ACP_TERMINAL_LIMIT} terminals can be active per VM"
)));
}
Ok(Self {
agent,
active: true,
})
}
fn disarm(&mut self) {
self.active = false;
}
}
impl Drop for AcpTerminalReservation<'_> {
fn drop(&mut self) {
if self.active {
release_counter(&self.agent.inner().acp_terminal_count);
}
}
}
impl AgentOs {
fn vm_ownership(&self) -> wire::OwnershipScope {
wire::OwnershipScope::VmOwnership(wire::VmOwnership {
connection_id: self.connection_id().to_string(),
session_id: self.wire_session_id().to_string(),
vm_id: self.vm_id().to_string(),
})
}
pub(crate) fn finish_acp_terminal(&self, process_id: &str) {
if self.inner().acp_terminals.remove(process_id).is_some() {
release_counter(&self.inner().acp_terminal_count);
}
}
async fn start_acp_terminal(
&self,
execute: wire::ExecuteRequest,
ownership: wire::OwnershipScope,
pid_tx: tokio::sync::oneshot::Sender<std::result::Result<u32, ClientError>>,
process_id: &str,
) -> Option<u32> {
{
let _terminal_lifecycle_guard = self.inner().acp_terminal_lifecycle_lock.lock().await;
if self.inner().disposed.load(Ordering::SeqCst) {
let error = ClientError::Sidecar(
"cannot connect terminal after VM shutdown has started".to_string(),
);
let _ = pid_tx.send(Err(error));
self.finish_acp_terminal(process_id);
return None;
}
}
let result = match self
.transport()
.request_wire(ownership, wire::RequestPayload::ExecuteRequest(execute))
.await
{
Ok(wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
pid,
..
})) => pid.ok_or_else(|| {
ClientError::Sidecar("connect_terminal: sidecar did not return a pid".to_string())
}),
Ok(wire::ResponsePayload::RejectedResponse(rejected)) => {
Err(rejected_to_error(rejected))
}
Ok(other) => Err(ClientError::Sidecar(format!(
"unexpected response to connect_terminal: {other:?}"
))),
Err(error) => Err(error.into()),
};
match result {
Ok(pid) => {
let _ = pid_tx.send(Ok(pid));
Some(pid)
}
Err(error) => {
let _ = pid_tx.send(Err(error));
self.finish_acp_terminal(process_id);
None
}
}
}
}
impl AgentOs {
pub fn open_shell(&self, mut options: OpenShellOptions) -> Result<ShellHandle> {
let inner = self.inner();
let counter = inner.shell_counter.fetch_add(1, Ordering::SeqCst) + 1;
let shell_id = format!("shell-{counter}");
let process_id = format!("shell-{}", Uuid::new_v4());
let (data_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
let (stderr_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
let (spawned_tx, _) = tokio::sync::watch::channel(false);
let (exit_tx, _) = tokio::sync::watch::channel(None::<i32>);
if let Some(cb) = options.on_stderr.take() {
install_output_callback(stderr_tx.clone(), cb);
}
let entry = ShellEntry {
pid: 0,
data_tx: data_tx.clone(),
stderr_tx: stderr_tx.clone(),
process_id: process_id.clone(),
spawned_tx: spawned_tx.clone(),
exit_tx: exit_tx.clone(),
};
let _ = inner.shells.insert(shell_id.clone(), entry);
let command = options
.command
.clone()
.unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
options
.env
.insert(String::from("AGENTOS_EXEC_TTY"), String::from("1"));
if let Some(cols) = options.cols {
options
.env
.insert(String::from("COLUMNS"), cols.to_string());
}
if let Some(rows) = options.rows {
options.env.insert(String::from("LINES"), rows.to_string());
}
let execute = wire::ExecuteRequest {
process_id: process_id.clone(),
command: Some(command),
runtime: None,
entrypoint: None,
args: options.args.clone(),
env: options.env.clone().into_iter().collect(),
cwd: options.cwd.clone(),
wasm_permission_tier: None,
};
let agent = self.clone();
let ownership = self.vm_ownership();
let route_process_id = process_id.clone();
let exit_shell_id = shell_id.clone();
let exit_key = counter;
let handle = tokio::spawn(async move {
let mut events = agent.transport().subscribe_wire_events();
let response = match agent
.transport()
.request_wire(
ownership.clone(),
wire::RequestPayload::ExecuteRequest(execute),
)
.await
{
Ok(response) => response,
Err(error) => {
tracing::warn!(?error, shell_id = %exit_shell_id, "open_shell spawn failed");
agent.inner().shells.remove(&exit_shell_id);
agent.inner().pending_shell_exits.remove(&exit_key);
return;
}
};
if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
pid: Some(pid),
..
}) = response
{
agent
.inner()
.shells
.update(&exit_shell_id, |_, existing| existing.pid = pid);
}
let _ = spawned_tx.send_replace(true);
loop {
let (_scope, payload) = match events.recv().await {
Ok(value) => value,
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
};
match payload {
EventPayload::ProcessOutputEvent(output) => {
if output.process_id != route_process_id {
continue;
}
match output.channel {
StreamChannel::Stdout => {
let _ = data_tx.send(output.chunk);
}
StreamChannel::Stderr => {
let _ = stderr_tx.send(output.chunk);
}
}
}
EventPayload::ProcessExitedEvent(exited) => {
if exited.process_id == route_process_id {
{
let mut retained = agent.inner().closed_shell_exit_codes.lock();
retained.push_back((exit_shell_id.clone(), exited.exit_code));
while retained.len() > crate::CLOSED_SHELL_EXIT_CODE_RETENTION_LIMIT
{
retained.pop_front();
}
}
let _ = exit_tx.send(Some(exited.exit_code));
break;
}
}
EventPayload::VmLifecycleEvent(_)
| EventPayload::StructuredEvent(_)
| EventPayload::ExtEnvelope(_) => {}
}
}
agent.inner().pending_shell_exits.remove(&exit_key);
agent.inner().shells.remove_if(&exit_shell_id, |existing| {
existing.process_id == route_process_id
});
});
let _ = inner.pending_shell_exits.insert(counter, handle);
Ok(ShellHandle { shell_id })
}
pub(crate) fn acp_open_terminal(
&self,
options: OpenShellOptions,
exit_tx: tokio::sync::watch::Sender<Option<i32>>,
on_output: impl Fn(&[u8]) + Send + Sync + 'static,
) -> Result<ShellHandle> {
let inner = self.inner();
let counter = inner.shell_counter.fetch_add(1, Ordering::SeqCst) + 1;
let shell_id = format!("shell-{counter}");
let process_id = format!("shell-{}", Uuid::new_v4());
let (data_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
let (stderr_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
let (spawned_tx, _) = tokio::sync::watch::channel(false);
let entry = ShellEntry {
pid: 0,
data_tx: data_tx.clone(),
stderr_tx: stderr_tx.clone(),
process_id: process_id.clone(),
spawned_tx: spawned_tx.clone(),
exit_tx: exit_tx.clone(),
};
let _ = inner.shells.insert(shell_id.clone(), entry);
let command = options
.command
.clone()
.unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
let execute = wire::ExecuteRequest {
process_id: process_id.clone(),
command: Some(command),
runtime: None,
entrypoint: None,
args: options.args.clone(),
env: options.env.clone().into_iter().collect(),
cwd: options.cwd.clone(),
wasm_permission_tier: None,
};
let agent = self.clone();
let ownership = self.vm_ownership();
let route_process_id = process_id.clone();
let exit_shell_id = shell_id.clone();
let exit_key = counter;
let on_output = std::sync::Arc::new(on_output);
let handle = tokio::spawn(async move {
let mut events = agent.transport().subscribe_wire_events();
let response = match agent
.transport()
.request_wire(
ownership.clone(),
wire::RequestPayload::ExecuteRequest(execute),
)
.await
{
Ok(response) => response,
Err(error) => {
tracing::warn!(?error, shell_id = %exit_shell_id, "acp_open_terminal spawn failed");
agent.inner().shells.remove(&exit_shell_id);
agent.inner().pending_shell_exits.remove(&exit_key);
let _ = exit_tx.send(Some(1));
return;
}
};
if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
pid: Some(pid),
..
}) = response
{
agent
.inner()
.shells
.update(&exit_shell_id, |_, existing| existing.pid = pid);
}
let _ = spawned_tx.send_replace(true);
let mut exit_code: i32 = 0;
loop {
let (_scope, payload) = match events.recv().await {
Ok(value) => value,
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
};
match payload {
EventPayload::ProcessOutputEvent(output) => {
if output.process_id != route_process_id {
continue;
}
on_output(&output.chunk);
}
EventPayload::ProcessExitedEvent(exited) => {
if exited.process_id == route_process_id {
exit_code = exited.exit_code;
break;
}
}
EventPayload::VmLifecycleEvent(_)
| EventPayload::StructuredEvent(_)
| EventPayload::ExtEnvelope(_) => {}
}
}
agent.inner().pending_shell_exits.remove(&exit_key);
agent.inner().shells.remove_if(&exit_shell_id, |existing| {
existing.process_id == route_process_id
});
let _ = exit_tx.send(Some(exit_code));
});
let _ = inner.pending_shell_exits.insert(counter, handle);
Ok(ShellHandle { shell_id })
}
pub(crate) fn acp_kill_terminal_shell(
&self,
shell_id: &str,
) -> std::result::Result<(), ClientError> {
let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
let agent = self.clone();
let ownership = self.vm_ownership();
tokio::spawn(async move {
wait_for_spawn(spawned_rx).await;
let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
process_id,
signal: String::from("SIGTERM"),
});
if let Err(error) = agent.transport().request_wire(ownership, payload).await {
tracing::warn!(?error, "acp_kill_terminal_shell failed");
}
});
Ok(())
}
pub async fn connect_terminal(&self, options: ConnectTerminalOptions) -> Result<u32> {
let ConnectTerminalOptions { base, on_data } = options;
let process_id = format!("terminal-{}", Uuid::new_v4());
let command = base
.command
.clone()
.unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
let (data_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
let (stderr_tx, _) =
tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
if let Some(cb) = on_data {
install_output_callback(data_tx.clone(), cb);
}
if let Some(cb) = base.on_stderr {
install_output_callback(stderr_tx.clone(), cb);
}
let execute = wire::ExecuteRequest {
process_id: process_id.clone(),
command: Some(command),
runtime: None,
entrypoint: None,
args: base.args.clone(),
env: base.env.clone().into_iter().collect(),
cwd: base.cwd.clone(),
wasm_permission_tier: None,
};
let events = self.transport().subscribe_wire_events();
let ownership = self.vm_ownership();
let (pid_tx, pid_rx) = tokio::sync::oneshot::channel();
let (start_tx, start_rx) = tokio::sync::oneshot::channel::<()>();
let agent = self.clone();
let route_process_id = process_id.clone();
let exit_task = tokio::spawn(async move {
if start_rx.await.is_err() {
return;
}
let terminal_pid = match agent
.start_acp_terminal(execute, ownership, pid_tx, &route_process_id)
.await
{
Some(pid) => pid,
None => return,
};
let mut events = events;
loop {
let (_scope, payload) = match events.recv().await {
Ok(value) => value,
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
if terminal_process_finished(&agent, terminal_pid).await {
break;
}
continue;
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
};
match payload {
EventPayload::ProcessOutputEvent(output) => {
if output.process_id != route_process_id {
continue;
}
match output.channel {
StreamChannel::Stdout => {
let _ = data_tx.send(output.chunk);
}
StreamChannel::Stderr => {
let _ = stderr_tx.send(output.chunk);
}
}
}
EventPayload::ProcessExitedEvent(exited) => {
if exited.process_id == route_process_id {
break;
}
}
EventPayload::VmLifecycleEvent(_)
| EventPayload::StructuredEvent(_)
| EventPayload::ExtEnvelope(_) => {}
}
}
agent.finish_acp_terminal(&route_process_id);
});
{
let _terminal_lifecycle_guard = self.inner().acp_terminal_lifecycle_lock.lock().await;
if self.inner().disposed.load(Ordering::SeqCst) {
exit_task.abort();
return Err(ClientError::Sidecar(
"cannot connect terminal after VM shutdown has started".to_string(),
)
.into());
}
let mut terminal_reservation = AcpTerminalReservation::new(self)?;
match self
.inner()
.acp_terminals
.insert(process_id.clone(), AcpTerminalEntry { exit_task })
{
Ok(()) => {}
Err((_, entry)) => {
entry.exit_task.abort();
return Err(ClientError::Sidecar(format!(
"terminal process id collision while tracking ACP terminal: {process_id}"
))
.into());
}
}
terminal_reservation.disarm();
if start_tx.send(()).is_err() {
self.finish_acp_terminal(&process_id);
return Err(ClientError::Sidecar(
"terminal startup task ended before registration completed".to_string(),
)
.into());
}
}
pid_rx
.await
.map_err(|_| {
ClientError::Sidecar(
"terminal startup task ended before returning a pid".to_string(),
)
})?
.map_err(Into::into)
}
pub fn write_shell(
&self,
shell_id: &str,
data: StdinInput,
) -> std::result::Result<(), ClientError> {
let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
let chunk = stdin_chunk(data);
let agent = self.clone();
let ownership = self.vm_ownership();
tokio::spawn(async move {
wait_for_spawn(spawned_rx).await;
let payload = wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
process_id,
chunk,
});
if let Err(error) = agent.transport().request_wire(ownership, payload).await {
tracing::warn!(?error, "write_shell failed");
}
});
Ok(())
}
pub async fn write_shell_awaited(
&self,
shell_id: &str,
data: StdinInput,
) -> std::result::Result<(), ClientError> {
let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
let chunk = stdin_chunk(data);
tracing::debug!(shell_id, "write_shell_awaited: waiting for spawn gate");
wait_for_spawn(spawned_rx).await;
tracing::debug!(shell_id, "write_shell_awaited: issuing wire write");
let payload =
wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest { process_id, chunk });
let response = self
.transport()
.request_wire(self.vm_ownership(), payload)
.await?;
tracing::debug!(shell_id, "write_shell_awaited: wire write acked");
match response {
wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
_ => Ok(()),
}
}
pub fn on_shell_data(&self, shell_id: &str) -> std::result::Result<ByteStream, ClientError> {
self.inner()
.shells
.read(shell_id, |_, entry| entry.data_tx.subscribe())
.map(ByteStream::new)
.ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
}
pub fn on_shell_stderr(&self, shell_id: &str) -> std::result::Result<ByteStream, ClientError> {
self.inner()
.shells
.read(shell_id, |_, entry| entry.stderr_tx.subscribe())
.map(ByteStream::new)
.ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
}
pub fn resize_shell(
&self,
shell_id: &str,
cols: u16,
rows: u16,
) -> std::result::Result<(), ClientError> {
let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
let agent = self.clone();
let ownership = self.vm_ownership();
tokio::spawn(async move {
wait_for_spawn(spawned_rx).await;
let payload = wire::RequestPayload::ResizePtyRequest(wire::ResizePtyRequest {
process_id,
cols,
rows,
});
if let Err(error) = agent.transport().request_wire(ownership, payload).await {
tracing::warn!(?error, "resize_shell failed");
}
});
Ok(())
}
pub async fn wait_shell(&self, shell_id: &str) -> std::result::Result<i32, ClientError> {
let exit_rx = self
.inner()
.shells
.read(shell_id, |_, entry| entry.exit_tx.subscribe());
let Some(mut exit_rx) = exit_rx else {
let retained = self.inner().closed_shell_exit_codes.lock();
return retained
.iter()
.rev()
.find(|(id, _)| id == shell_id)
.map(|(_, code)| *code)
.ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()));
};
loop {
if let Some(code) = *exit_rx.borrow_and_update() {
return Ok(code);
}
if exit_rx.changed().await.is_err() {
let retained = self.inner().closed_shell_exit_codes.lock();
return retained
.iter()
.rev()
.find(|(id, _)| id == shell_id)
.map(|(_, code)| *code)
.ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()));
}
}
}
pub fn close_shell(&self, shell_id: &str) -> std::result::Result<(), ClientError> {
let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
self.inner().shells.remove(shell_id);
let agent = self.clone();
let ownership = self.vm_ownership();
tokio::spawn(async move {
wait_for_spawn(spawned_rx).await;
let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
process_id,
signal: String::from("SIGTERM"),
});
if let Err(error) = agent.transport().request_wire(ownership, payload).await {
tracing::warn!(?error, "close_shell kill failed");
}
});
Ok(())
}
fn shell_wire_handle(
&self,
shell_id: &str,
) -> std::result::Result<(String, tokio::sync::watch::Receiver<bool>), ClientError> {
self.inner()
.shells
.read(shell_id, |_, entry| {
(entry.process_id.clone(), entry.spawned_tx.subscribe())
})
.ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
}
}
async fn wait_for_spawn(mut spawned_rx: tokio::sync::watch::Receiver<bool>) {
if *spawned_rx.borrow() {
return;
}
while spawned_rx.changed().await.is_ok() {
if *spawned_rx.borrow() {
return;
}
}
}
async fn terminal_process_finished(agent: &AgentOs, pid: u32) -> bool {
match agent.all_processes().await {
Ok(processes) => match processes.into_iter().find(|process| process.pid == pid) {
Some(process) => process.status != ProcessStatus::Running,
None => true,
},
Err(error) => {
tracing::warn!(?error, pid, "terminal process snapshot failed");
false
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reserve_counter_enforces_limit_and_release_reopens_slot() {
let counter = AtomicUsize::new(0);
assert!(try_reserve_counter(&counter, 2));
assert!(try_reserve_counter(&counter, 2));
assert!(!try_reserve_counter(&counter, 2));
release_counter(&counter);
assert!(try_reserve_counter(&counter, 2));
}
}