1use std::collections::BTreeMap;
22use std::sync::atomic::{AtomicUsize, Ordering};
23
24use anyhow::Result;
25use uuid::Uuid;
26
27use secure_exec_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
34const SHELL_DATA_CHANNEL_CAPACITY: usize = 1024;
36
37const ACP_TERMINAL_LIMIT: usize = 1024;
39
40const DEFAULT_SHELL_COMMAND: &str = "sh";
43
44#[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#[derive(Default)]
70pub struct ConnectTerminalOptions {
71 pub base: OpenShellOptions,
72 pub on_data: Option<OutputCallback>,
73}
74
75#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct ShellHandle {
78 pub shell_id: String,
79}
80
81fn rejected_to_error(rejected: wire::RejectedResponse) -> ClientError {
87 ClientError::Kernel {
88 code: rejected.code,
89 message: rejected.message,
90 }
91}
92
93fn 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 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
217impl AgentOs {
226 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 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 let (spawned_tx, _) = tokio::sync::watch::channel(false);
247
248 if let Some(cb) = options.on_stderr.take() {
251 install_output_callback(stderr_tx.clone(), cb);
252 }
253
254 let entry = ShellEntry {
257 pid: 0,
258 data_tx: data_tx.clone(),
259 stderr_tx: stderr_tx.clone(),
260 process_id: process_id.clone(),
261 spawned_tx: spawned_tx.clone(),
262 };
263 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 let execute = wire::ExecuteRequest {
271 process_id: process_id.clone(),
272 command: Some(command),
273 runtime: None,
274 entrypoint: None,
275 args: options.args.clone(),
276 env: options.env.clone().into_iter().collect(),
277 cwd: options.cwd.clone(),
278 wasm_permission_tier: None,
279 };
280
281 let agent = self.clone();
285 let ownership = self.vm_ownership();
286 let route_process_id = process_id.clone();
287 let exit_shell_id = shell_id.clone();
288 let exit_key = counter;
289 let handle = tokio::spawn(async move {
290 let mut events = agent.transport().subscribe_wire_events();
291
292 let response = match agent
293 .transport()
294 .request_wire(
295 ownership.clone(),
296 wire::RequestPayload::ExecuteRequest(execute),
297 )
298 .await
299 {
300 Ok(response) => response,
301 Err(error) => {
302 tracing::warn!(?error, shell_id = %exit_shell_id, "open_shell spawn failed");
303 agent.inner().shells.remove(&exit_shell_id);
305 agent.inner().pending_shell_exits.remove(&exit_key);
306 return;
307 }
308 };
309
310 if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
313 pid: Some(pid),
314 ..
315 }) = response
316 {
317 agent
318 .inner()
319 .shells
320 .update(&exit_shell_id, |_, existing| existing.pid = pid);
321 }
322 let _ = spawned_tx.send(true);
323
324 loop {
325 let (_scope, payload) = match events.recv().await {
326 Ok(value) => value,
327 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
328 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
329 };
330 match payload {
331 EventPayload::ProcessOutputEvent(output) => {
332 if output.process_id != route_process_id {
333 continue;
334 }
335 match output.channel {
337 StreamChannel::Stdout => {
338 let _ = data_tx.send(output.chunk);
339 }
340 StreamChannel::Stderr => {
341 let _ = stderr_tx.send(output.chunk);
342 }
343 }
344 }
345 EventPayload::ProcessExitedEvent(exited) => {
346 if exited.process_id == route_process_id {
347 break;
348 }
349 }
350 EventPayload::VmLifecycleEvent(_)
351 | EventPayload::StructuredEvent(_)
352 | EventPayload::ExtEnvelope(_) => {}
353 }
354 }
355
356 agent.inner().pending_shell_exits.remove(&exit_key);
359 agent.inner().shells.remove_if(&exit_shell_id, |existing| {
360 existing.process_id == route_process_id
361 });
362 });
364
365 let _ = inner.pending_shell_exits.insert(counter, handle);
366
367 Ok(ShellHandle { shell_id })
368 }
369
370 pub(crate) fn acp_open_terminal(
377 &self,
378 options: OpenShellOptions,
379 exit_tx: tokio::sync::watch::Sender<Option<i32>>,
380 on_output: impl Fn(&[u8]) + Send + Sync + 'static,
381 ) -> Result<ShellHandle> {
382 let inner = self.inner();
383 let counter = inner.shell_counter.fetch_add(1, Ordering::SeqCst) + 1;
384 let shell_id = format!("shell-{counter}");
385 let process_id = format!("shell-{}", Uuid::new_v4());
386
387 let (data_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
388 let (stderr_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
389 let (spawned_tx, _) = tokio::sync::watch::channel(false);
390
391 let entry = ShellEntry {
392 pid: 0,
393 data_tx: data_tx.clone(),
394 stderr_tx: stderr_tx.clone(),
395 process_id: process_id.clone(),
396 spawned_tx: spawned_tx.clone(),
397 };
398 let _ = inner.shells.insert(shell_id.clone(), entry);
399
400 let command = options
401 .command
402 .clone()
403 .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
404 let execute = wire::ExecuteRequest {
405 process_id: process_id.clone(),
406 command: Some(command),
407 runtime: None,
408 entrypoint: None,
409 args: options.args.clone(),
410 env: options.env.clone().into_iter().collect(),
411 cwd: options.cwd.clone(),
412 wasm_permission_tier: None,
413 };
414
415 let agent = self.clone();
416 let ownership = self.vm_ownership();
417 let route_process_id = process_id.clone();
418 let exit_shell_id = shell_id.clone();
419 let exit_key = counter;
420 let on_output = std::sync::Arc::new(on_output);
421 let handle = tokio::spawn(async move {
422 let mut events = agent.transport().subscribe_wire_events();
423
424 let response = match agent
425 .transport()
426 .request_wire(
427 ownership.clone(),
428 wire::RequestPayload::ExecuteRequest(execute),
429 )
430 .await
431 {
432 Ok(response) => response,
433 Err(error) => {
434 tracing::warn!(?error, shell_id = %exit_shell_id, "acp_open_terminal spawn failed");
435 agent.inner().shells.remove(&exit_shell_id);
436 agent.inner().pending_shell_exits.remove(&exit_key);
437 let _ = exit_tx.send(Some(1));
438 return;
439 }
440 };
441
442 if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
443 pid: Some(pid),
444 ..
445 }) = response
446 {
447 agent
448 .inner()
449 .shells
450 .update(&exit_shell_id, |_, existing| existing.pid = pid);
451 }
452 let _ = spawned_tx.send(true);
453
454 let mut exit_code: i32 = 0;
455 loop {
456 let (_scope, payload) = match events.recv().await {
457 Ok(value) => value,
458 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
459 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
460 };
461 match payload {
462 EventPayload::ProcessOutputEvent(output) => {
463 if output.process_id != route_process_id {
464 continue;
465 }
466 on_output(&output.chunk);
469 }
470 EventPayload::ProcessExitedEvent(exited) => {
471 if exited.process_id == route_process_id {
472 exit_code = exited.exit_code;
473 break;
474 }
475 }
476 EventPayload::VmLifecycleEvent(_)
477 | EventPayload::StructuredEvent(_)
478 | EventPayload::ExtEnvelope(_) => {}
479 }
480 }
481
482 agent.inner().pending_shell_exits.remove(&exit_key);
483 agent.inner().shells.remove_if(&exit_shell_id, |existing| {
484 existing.process_id == route_process_id
485 });
486 let _ = exit_tx.send(Some(exit_code));
487 });
488
489 let _ = inner.pending_shell_exits.insert(counter, handle);
493 Ok(ShellHandle { shell_id })
494 }
495
496 pub(crate) fn acp_kill_terminal_shell(
500 &self,
501 shell_id: &str,
502 ) -> std::result::Result<(), ClientError> {
503 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
504 let agent = self.clone();
505 let ownership = self.vm_ownership();
506 tokio::spawn(async move {
507 wait_for_spawn(spawned_rx).await;
508 let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
509 process_id,
510 signal: String::from("SIGTERM"),
511 });
512 if let Err(error) = agent.transport().request_wire(ownership, payload).await {
513 tracing::warn!(?error, "acp_kill_terminal_shell failed");
514 }
515 });
516 Ok(())
517 }
518
519 pub async fn connect_terminal(&self, options: ConnectTerminalOptions) -> Result<u32> {
528 let ConnectTerminalOptions { base, on_data } = options;
529
530 let process_id = format!("terminal-{}", Uuid::new_v4());
531 let command = base
532 .command
533 .clone()
534 .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
535 let (data_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
536 let (stderr_tx, _) =
537 tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
538
539 if let Some(cb) = on_data {
543 install_output_callback(data_tx.clone(), cb);
544 }
545 if let Some(cb) = base.on_stderr {
546 install_output_callback(stderr_tx.clone(), cb);
547 }
548
549 let execute = wire::ExecuteRequest {
550 process_id: process_id.clone(),
551 command: Some(command),
552 runtime: None,
553 entrypoint: None,
554 args: base.args.clone(),
555 env: base.env.clone().into_iter().collect(),
556 cwd: base.cwd.clone(),
557 wasm_permission_tier: None,
558 };
559
560 let events = self.transport().subscribe_wire_events();
562 let ownership = self.vm_ownership();
563 let (pid_tx, pid_rx) = tokio::sync::oneshot::channel();
564 let (start_tx, start_rx) = tokio::sync::oneshot::channel::<()>();
565 let agent = self.clone();
566 let route_process_id = process_id.clone();
567 let exit_task = tokio::spawn(async move {
568 if start_rx.await.is_err() {
569 return;
570 }
571 let terminal_pid = match agent
572 .start_acp_terminal(execute, ownership, pid_tx, &route_process_id)
573 .await
574 {
575 Some(pid) => pid,
576 None => return,
577 };
578 let mut events = events;
579 loop {
580 let (_scope, payload) = match events.recv().await {
581 Ok(value) => value,
582 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
583 if terminal_process_finished(&agent, terminal_pid).await {
584 break;
585 }
586 continue;
587 }
588 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
589 };
590 match payload {
591 EventPayload::ProcessOutputEvent(output) => {
592 if output.process_id != route_process_id {
593 continue;
594 }
595 match output.channel {
596 StreamChannel::Stdout => {
597 let _ = data_tx.send(output.chunk);
598 }
599 StreamChannel::Stderr => {
600 let _ = stderr_tx.send(output.chunk);
601 }
602 }
603 }
604 EventPayload::ProcessExitedEvent(exited) => {
605 if exited.process_id == route_process_id {
606 break;
607 }
608 }
609 EventPayload::VmLifecycleEvent(_)
610 | EventPayload::StructuredEvent(_)
611 | EventPayload::ExtEnvelope(_) => {}
612 }
613 }
614 agent.finish_acp_terminal(&route_process_id);
615 });
616
617 {
618 let _terminal_lifecycle_guard = self.inner().acp_terminal_lifecycle_lock.lock().await;
619 if self.inner().disposed.load(Ordering::SeqCst) {
620 exit_task.abort();
621 return Err(ClientError::Sidecar(
622 "cannot connect terminal after VM shutdown has started".to_string(),
623 )
624 .into());
625 }
626 let mut terminal_reservation = AcpTerminalReservation::new(self)?;
627 match self
628 .inner()
629 .acp_terminals
630 .insert(process_id.clone(), AcpTerminalEntry { exit_task })
631 {
632 Ok(()) => {}
633 Err((_, entry)) => {
634 entry.exit_task.abort();
635 return Err(ClientError::Sidecar(format!(
636 "terminal process id collision while tracking ACP terminal: {process_id}"
637 ))
638 .into());
639 }
640 }
641 terminal_reservation.disarm();
642 if start_tx.send(()).is_err() {
643 self.finish_acp_terminal(&process_id);
644 return Err(ClientError::Sidecar(
645 "terminal startup task ended before registration completed".to_string(),
646 )
647 .into());
648 }
649 }
650
651 pid_rx
652 .await
653 .map_err(|_| {
654 ClientError::Sidecar(
655 "terminal startup task ended before returning a pid".to_string(),
656 )
657 })?
658 .map_err(Into::into)
659 }
660
661 pub fn write_shell(
663 &self,
664 shell_id: &str,
665 data: StdinInput,
666 ) -> std::result::Result<(), ClientError> {
667 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
668 let chunk = stdin_chunk(data);
669
670 let agent = self.clone();
675 let ownership = self.vm_ownership();
676 tokio::spawn(async move {
677 wait_for_spawn(spawned_rx).await;
678 let payload = wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
679 process_id,
680 chunk,
681 });
682 if let Err(error) = agent.transport().request_wire(ownership, payload).await {
683 tracing::warn!(?error, "write_shell failed");
684 }
685 });
686
687 Ok(())
688 }
689
690 pub fn on_shell_data(&self, shell_id: &str) -> std::result::Result<ByteStream, ClientError> {
694 self.inner()
695 .shells
696 .read(shell_id, |_, entry| entry.data_tx.subscribe())
697 .map(ByteStream::new)
698 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
699 }
700
701 pub fn on_shell_stderr(&self, shell_id: &str) -> std::result::Result<ByteStream, ClientError> {
705 self.inner()
706 .shells
707 .read(shell_id, |_, entry| entry.stderr_tx.subscribe())
708 .map(ByteStream::new)
709 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
710 }
711
712 pub fn resize_shell(
718 &self,
719 shell_id: &str,
720 cols: u16,
721 rows: u16,
722 ) -> std::result::Result<(), ClientError> {
723 let _ = self.shell_wire_handle(shell_id)?;
725 tracing::warn!(
726 shell_id = %shell_id,
727 cols,
728 rows,
729 "resize_shell has no native winsize wire op; resize is a no-op"
730 );
731 Ok(())
732 }
733
734 pub fn close_shell(&self, shell_id: &str) -> std::result::Result<(), ClientError> {
737 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
738
739 self.inner().shells.remove(shell_id);
742
743 let agent = self.clone();
745 let ownership = self.vm_ownership();
746 tokio::spawn(async move {
747 wait_for_spawn(spawned_rx).await;
748 let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
749 process_id,
750 signal: String::from("SIGTERM"),
751 });
752 if let Err(error) = agent.transport().request_wire(ownership, payload).await {
753 tracing::warn!(?error, "close_shell kill failed");
754 }
755 });
756
757 Ok(())
758 }
759
760 fn shell_wire_handle(
763 &self,
764 shell_id: &str,
765 ) -> std::result::Result<(String, tokio::sync::watch::Receiver<bool>), ClientError> {
766 self.inner()
767 .shells
768 .read(shell_id, |_, entry| {
769 (entry.process_id.clone(), entry.spawned_tx.subscribe())
770 })
771 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
772 }
773}
774
775async fn wait_for_spawn(mut spawned_rx: tokio::sync::watch::Receiver<bool>) {
778 if *spawned_rx.borrow() {
779 return;
780 }
781 while spawned_rx.changed().await.is_ok() {
782 if *spawned_rx.borrow() {
783 return;
784 }
785 }
786}
787
788async fn terminal_process_finished(agent: &AgentOs, pid: u32) -> bool {
789 match agent.all_processes().await {
790 Ok(processes) => match processes.into_iter().find(|process| process.pid == pid) {
791 Some(process) => process.status != ProcessStatus::Running,
792 None => true,
793 },
794 Err(error) => {
795 tracing::warn!(?error, pid, "terminal process snapshot failed");
796 false
797 }
798 }
799}
800
801#[cfg(test)]
802mod tests {
803 use super::*;
804
805 #[test]
806 fn reserve_counter_enforces_limit_and_release_reopens_slot() {
807 let counter = AtomicUsize::new(0);
808
809 assert!(try_reserve_counter(&counter, 2));
810 assert!(try_reserve_counter(&counter, 2));
811 assert!(!try_reserve_counter(&counter, 2));
812 release_counter(&counter);
813 assert!(try_reserve_counter(&counter, 2));
814 }
815}