1use 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};
29
30const SHELL_DATA_CHANNEL_CAPACITY: usize = 1024;
32
33const ACP_TERMINAL_LIMIT: usize = 1024;
35
36const DEFAULT_SHELL_COMMAND: &str = "sh";
39
40#[derive(Default)]
46pub struct OpenShellOptions {
47 pub command: Option<String>,
48 pub args: Vec<String>,
49 pub env: BTreeMap<String, String>,
50 pub cwd: Option<String>,
51 pub cols: Option<u16>,
52 pub rows: Option<u16>,
53}
54
55#[derive(Default)]
61pub struct ConnectTerminalOptions {
62 pub base: OpenShellOptions,
63 pub on_data: Option<OutputCallback>,
64 pub on_stderr: Option<OutputCallback>,
65}
66
67#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct ShellHandle {
70 pub shell_id: String,
71}
72
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub struct ShellData {
75 pub shell_id: String,
76 pub data: Vec<u8>,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq)]
80pub struct ShellExit {
81 pub shell_id: String,
82 pub exit_code: i32,
83}
84
85fn rejected_to_error(rejected: wire::RejectedResponse) -> ClientError {
91 ClientError::Kernel {
92 code: rejected.code,
93 message: rejected.message,
94 }
95}
96
97fn stdin_chunk(data: StdinInput) -> Vec<u8> {
101 match data {
102 StdinInput::Text(text) => text.into_bytes(),
103 StdinInput::Bytes(bytes) => bytes,
104 }
105}
106
107fn try_reserve_counter(counter: &AtomicUsize, limit: usize) -> bool {
108 counter
109 .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
110 (count < limit).then_some(count + 1)
111 })
112 .is_ok()
113}
114
115fn release_counter(counter: &AtomicUsize) {
116 let _ = counter.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| {
117 Some(count.saturating_sub(1))
118 });
119}
120
121struct AcpTerminalReservation<'a> {
122 agent: &'a AgentOs,
123 active: bool,
124}
125
126impl<'a> AcpTerminalReservation<'a> {
127 fn new(agent: &'a AgentOs) -> std::result::Result<Self, ClientError> {
128 if !try_reserve_counter(&agent.inner().acp_terminal_count, ACP_TERMINAL_LIMIT) {
129 return Err(ClientError::Sidecar(format!(
130 "acp terminal limit exceeded: at most {ACP_TERMINAL_LIMIT} terminals can be active per VM"
131 )));
132 }
133 Ok(Self {
134 agent,
135 active: true,
136 })
137 }
138
139 fn disarm(&mut self) {
140 self.active = false;
141 }
142}
143
144impl Drop for AcpTerminalReservation<'_> {
145 fn drop(&mut self) {
146 if self.active {
147 release_counter(&self.agent.inner().acp_terminal_count);
148 }
149 }
150}
151
152impl AgentOs {
153 fn vm_ownership(&self) -> wire::OwnershipScope {
155 wire::OwnershipScope::VmOwnership(wire::VmOwnership {
156 connection_id: self.connection_id().to_string(),
157 session_id: self.wire_session_id().to_string(),
158 vm_id: self.vm_id().to_string(),
159 })
160 }
161
162 pub(crate) fn finish_acp_terminal(&self, process_id: &str) {
163 if self.inner().acp_terminals.remove(process_id).is_some() {
164 release_counter(&self.inner().acp_terminal_count);
165 }
166 }
167
168 async fn start_acp_terminal(
169 &self,
170 execute: wire::ExecuteRequest,
171 ownership: wire::OwnershipScope,
172 pid_tx: tokio::sync::oneshot::Sender<std::result::Result<u32, ClientError>>,
173 process_id: &str,
174 ) -> Option<u32> {
175 {
176 let _terminal_lifecycle_guard = self.inner().acp_terminal_lifecycle_lock.lock().await;
177 if self.inner().disposed.load(Ordering::SeqCst) {
178 let error = ClientError::Sidecar(
179 "cannot connect terminal after VM shutdown has started".to_string(),
180 );
181 let _ = pid_tx.send(Err(error));
182 self.finish_acp_terminal(process_id);
183 return None;
184 }
185 }
186
187 let result = match self
188 .transport()
189 .request_wire(ownership, wire::RequestPayload::ExecuteRequest(execute))
190 .await
191 {
192 Ok(wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
193 pid,
194 ..
195 })) => pid.ok_or_else(|| {
196 ClientError::Sidecar("connect_terminal: sidecar did not return a pid".to_string())
197 }),
198 Ok(wire::ResponsePayload::RejectedResponse(rejected)) => {
199 Err(rejected_to_error(rejected))
200 }
201 Ok(other) => Err(ClientError::Sidecar(format!(
202 "unexpected response to connect_terminal: {other:?}"
203 ))),
204 Err(error) => Err(error.into()),
205 };
206
207 match result {
208 Ok(pid) => {
209 let _ = pid_tx.send(Ok(pid));
210 Some(pid)
211 }
212 Err(error) => {
213 let _ = pid_tx.send(Err(error));
214 self.finish_acp_terminal(process_id);
215 None
216 }
217 }
218 }
219}
220
221impl AgentOs {
230 pub fn open_shell(&self, mut options: OpenShellOptions) -> Result<ShellHandle> {
241 let inner = self.inner();
242 let counter = inner.shell_counter.fetch_add(1, Ordering::SeqCst) + 1;
243 let shell_id = format!("shell-{counter}");
244 let process_id = format!("shell-{}", Uuid::new_v4());
246
247 let (data_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
248 let (stderr_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
249 let (spawned_tx, _) = tokio::sync::watch::channel(false);
251 let (exit_tx, _) = tokio::sync::watch::channel(None::<i32>);
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 exit_tx: exit_tx.clone(),
263 };
264 let _ = inner.shells.insert(shell_id.clone(), entry);
266
267 let command = options
268 .command
269 .clone()
270 .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
271 options
272 .env
273 .insert(String::from("AGENTOS_EXEC_TTY"), String::from("1"));
274 if let Some(cols) = options.cols {
276 options
277 .env
278 .insert(String::from("COLUMNS"), cols.to_string());
279 }
280 if let Some(rows) = options.rows {
281 options.env.insert(String::from("LINES"), rows.to_string());
282 }
283 let execute = wire::ExecuteRequest {
284 process_id: process_id.clone(),
285 command: Some(command),
286 runtime: None,
287 entrypoint: None,
288 args: options.args.clone(),
289 env: options.env.clone().into_iter().collect(),
290 cwd: options.cwd.clone(),
291 wasm_permission_tier: None,
292 };
293
294 let agent = self.clone();
298 let ownership = self.vm_ownership();
299 let route_process_id = process_id.clone();
300 let exit_shell_id = shell_id.clone();
301 let exit_key = counter;
302 let handle = tokio::spawn(async move {
303 let mut events = agent.transport().subscribe_wire_events();
304
305 let response = match agent
306 .transport()
307 .request_wire(
308 ownership.clone(),
309 wire::RequestPayload::ExecuteRequest(execute),
310 )
311 .await
312 {
313 Ok(response) => response,
314 Err(error) => {
315 tracing::warn!(?error, shell_id = %exit_shell_id, "open_shell spawn failed");
316 agent.inner().shells.remove(&exit_shell_id);
318 agent.inner().pending_shell_exits.remove(&exit_key);
319 return;
320 }
321 };
322
323 if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
326 pid: Some(pid),
327 ..
328 }) = response
329 {
330 agent
331 .inner()
332 .shells
333 .update(&exit_shell_id, |_, existing| existing.pid = pid);
334 }
335 let _ = spawned_tx.send_replace(true);
341
342 loop {
343 let (_scope, payload) = match events.recv().await {
344 Ok(value) => value,
345 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
346 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
347 };
348 match payload {
349 EventPayload::ProcessOutputEvent(output) => {
350 if output.process_id != route_process_id {
351 continue;
352 }
353 let _ = data_tx.send(output.chunk.clone());
356 if output.channel == StreamChannel::Stderr {
357 let _ = stderr_tx.send(output.chunk);
359 }
360 }
361 EventPayload::ProcessExitedEvent(exited) => {
362 if exited.process_id == route_process_id {
363 {
367 let mut retained = agent.inner().closed_shell_exit_codes.lock();
368 retained.push_back((exit_shell_id.clone(), exited.exit_code));
369 while retained.len() > crate::CLOSED_SHELL_EXIT_CODE_RETENTION_LIMIT
370 {
371 retained.pop_front();
372 }
373 }
374 let _ = exit_tx.send(Some(exited.exit_code));
375 break;
376 }
377 }
378 EventPayload::VmLifecycleEvent(_)
379 | EventPayload::StructuredEvent(_)
380 | EventPayload::ExtEnvelope(_) => {}
381 }
382 }
383
384 agent.inner().pending_shell_exits.remove(&exit_key);
387 agent.inner().shells.remove_if(&exit_shell_id, |existing| {
388 existing.process_id == route_process_id
389 });
390 });
392
393 let _ = inner.pending_shell_exits.insert(counter, handle);
394
395 Ok(ShellHandle { shell_id })
396 }
397
398 pub(crate) fn acp_open_terminal(
405 &self,
406 options: OpenShellOptions,
407 exit_tx: tokio::sync::watch::Sender<Option<i32>>,
408 on_output: impl Fn(&[u8]) + Send + Sync + 'static,
409 ) -> Result<ShellHandle> {
410 let inner = self.inner();
411 let counter = inner.shell_counter.fetch_add(1, Ordering::SeqCst) + 1;
412 let shell_id = format!("shell-{counter}");
413 let process_id = format!("shell-{}", Uuid::new_v4());
414
415 let (data_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
416 let (stderr_tx, _) = tokio::sync::broadcast::channel(SHELL_DATA_CHANNEL_CAPACITY);
417 let (spawned_tx, _) = tokio::sync::watch::channel(false);
418
419 let entry = ShellEntry {
420 pid: 0,
421 data_tx: data_tx.clone(),
422 stderr_tx: stderr_tx.clone(),
423 process_id: process_id.clone(),
424 spawned_tx: spawned_tx.clone(),
425 exit_tx: exit_tx.clone(),
427 };
428 let _ = inner.shells.insert(shell_id.clone(), entry);
429
430 let command = options
431 .command
432 .clone()
433 .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
434 let execute = wire::ExecuteRequest {
435 process_id: process_id.clone(),
436 command: Some(command),
437 runtime: None,
438 entrypoint: None,
439 args: options.args.clone(),
440 env: options.env.clone().into_iter().collect(),
441 cwd: options.cwd.clone(),
442 wasm_permission_tier: None,
443 };
444
445 let agent = self.clone();
446 let ownership = self.vm_ownership();
447 let route_process_id = process_id.clone();
448 let exit_shell_id = shell_id.clone();
449 let exit_key = counter;
450 let on_output = std::sync::Arc::new(on_output);
451 let handle = tokio::spawn(async move {
452 let mut events = agent.transport().subscribe_wire_events();
453
454 let response = match agent
455 .transport()
456 .request_wire(
457 ownership.clone(),
458 wire::RequestPayload::ExecuteRequest(execute),
459 )
460 .await
461 {
462 Ok(response) => response,
463 Err(error) => {
464 tracing::warn!(?error, shell_id = %exit_shell_id, "acp_open_terminal spawn failed");
465 agent.inner().shells.remove(&exit_shell_id);
466 agent.inner().pending_shell_exits.remove(&exit_key);
467 let _ = exit_tx.send(Some(1));
468 return;
469 }
470 };
471
472 if let wire::ResponsePayload::ProcessStartedResponse(wire::ProcessStartedResponse {
473 pid: Some(pid),
474 ..
475 }) = response
476 {
477 agent
478 .inner()
479 .shells
480 .update(&exit_shell_id, |_, existing| existing.pid = pid);
481 }
482 let _ = spawned_tx.send_replace(true);
488
489 let mut exit_code: i32 = 0;
490 loop {
491 let (_scope, payload) = match events.recv().await {
492 Ok(value) => value,
493 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
494 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
495 };
496 match payload {
497 EventPayload::ProcessOutputEvent(output) => {
498 if output.process_id != route_process_id {
499 continue;
500 }
501 let _ = data_tx.send(output.chunk.clone());
502 if output.channel == StreamChannel::Stderr {
503 let _ = stderr_tx.send(output.chunk.clone());
504 }
505 on_output(&output.chunk);
507 }
508 EventPayload::ProcessExitedEvent(exited) => {
509 if exited.process_id == route_process_id {
510 exit_code = exited.exit_code;
511 break;
512 }
513 }
514 EventPayload::VmLifecycleEvent(_)
515 | EventPayload::StructuredEvent(_)
516 | EventPayload::ExtEnvelope(_) => {}
517 }
518 }
519
520 agent.inner().pending_shell_exits.remove(&exit_key);
521 agent.inner().shells.remove_if(&exit_shell_id, |existing| {
522 existing.process_id == route_process_id
523 });
524 let _ = exit_tx.send(Some(exit_code));
525 });
526
527 let _ = inner.pending_shell_exits.insert(counter, handle);
531 Ok(ShellHandle { shell_id })
532 }
533
534 pub(crate) fn acp_kill_terminal_shell(
538 &self,
539 shell_id: &str,
540 ) -> std::result::Result<(), ClientError> {
541 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
542 let agent = self.clone();
543 let ownership = self.vm_ownership();
544 tokio::spawn(async move {
545 wait_for_spawn(spawned_rx).await;
546 let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
547 process_id,
548 signal: String::from("SIGTERM"),
549 });
550 if let Err(error) = agent.transport().request_wire(ownership, payload).await {
551 tracing::warn!(?error, "acp_kill_terminal_shell failed");
552 }
553 });
554 Ok(())
555 }
556
557 pub async fn connect_terminal(&self, options: ConnectTerminalOptions) -> Result<u32> {
566 let ConnectTerminalOptions {
567 base,
568 on_data,
569 on_stderr,
570 } = options;
571
572 let process_id = format!("terminal-{}", Uuid::new_v4());
573 let command = base
574 .command
575 .clone()
576 .unwrap_or_else(|| DEFAULT_SHELL_COMMAND.to_string());
577 let (data_tx, _) = tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
578 let (stderr_tx, _) =
579 tokio::sync::broadcast::channel::<Vec<u8>>(SHELL_DATA_CHANNEL_CAPACITY);
580
581 if let Some(cb) = on_data {
584 install_output_callback(data_tx.clone(), cb);
585 }
586 if let Some(cb) = on_stderr {
587 install_output_callback(stderr_tx.clone(), cb);
588 }
589
590 let execute = wire::ExecuteRequest {
591 process_id: process_id.clone(),
592 command: Some(command),
593 runtime: None,
594 entrypoint: None,
595 args: base.args.clone(),
596 env: base.env.clone().into_iter().collect(),
597 cwd: base.cwd.clone(),
598 wasm_permission_tier: None,
599 };
600
601 let events = self.transport().subscribe_wire_events();
603 let ownership = self.vm_ownership();
604 let (pid_tx, pid_rx) = tokio::sync::oneshot::channel();
605 let (start_tx, start_rx) = tokio::sync::oneshot::channel::<()>();
606 let agent = self.clone();
607 let route_process_id = process_id.clone();
608 let exit_task = tokio::spawn(async move {
609 if start_rx.await.is_err() {
610 return;
611 }
612 let terminal_pid = match agent
613 .start_acp_terminal(execute, ownership, pid_tx, &route_process_id)
614 .await
615 {
616 Some(pid) => pid,
617 None => return,
618 };
619 let mut events = events;
620 loop {
621 let (_scope, payload) = match events.recv().await {
622 Ok(value) => value,
623 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
624 if terminal_process_finished(&agent, terminal_pid).await {
625 break;
626 }
627 continue;
628 }
629 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
630 };
631 match payload {
632 EventPayload::ProcessOutputEvent(output) => {
633 if output.process_id != route_process_id {
634 continue;
635 }
636 let _ = data_tx.send(output.chunk.clone());
637 if output.channel == StreamChannel::Stderr {
638 let _ = stderr_tx.send(output.chunk);
639 }
640 }
641 EventPayload::ProcessExitedEvent(exited) => {
642 if exited.process_id == route_process_id {
643 break;
644 }
645 }
646 EventPayload::VmLifecycleEvent(_)
647 | EventPayload::StructuredEvent(_)
648 | EventPayload::ExtEnvelope(_) => {}
649 }
650 }
651 agent.finish_acp_terminal(&route_process_id);
652 });
653
654 {
655 let _terminal_lifecycle_guard = self.inner().acp_terminal_lifecycle_lock.lock().await;
656 if self.inner().disposed.load(Ordering::SeqCst) {
657 exit_task.abort();
658 return Err(ClientError::Sidecar(
659 "cannot connect terminal after VM shutdown has started".to_string(),
660 )
661 .into());
662 }
663 let mut terminal_reservation = AcpTerminalReservation::new(self)?;
664 match self
665 .inner()
666 .acp_terminals
667 .insert(process_id.clone(), AcpTerminalEntry { exit_task })
668 {
669 Ok(()) => {}
670 Err((_, entry)) => {
671 entry.exit_task.abort();
672 return Err(ClientError::Sidecar(format!(
673 "terminal process id collision while tracking ACP terminal: {process_id}"
674 ))
675 .into());
676 }
677 }
678 terminal_reservation.disarm();
679 if start_tx.send(()).is_err() {
680 self.finish_acp_terminal(&process_id);
681 return Err(ClientError::Sidecar(
682 "terminal startup task ended before registration completed".to_string(),
683 )
684 .into());
685 }
686 }
687
688 pid_rx
689 .await
690 .map_err(|_| {
691 ClientError::Sidecar(
692 "terminal startup task ended before returning a pid".to_string(),
693 )
694 })?
695 .map_err(Into::into)
696 }
697
698 pub fn write_shell(
700 &self,
701 shell_id: &str,
702 data: StdinInput,
703 ) -> std::result::Result<(), ClientError> {
704 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
705 let chunk = stdin_chunk(data);
706
707 let agent = self.clone();
712 let ownership = self.vm_ownership();
713 tokio::spawn(async move {
714 wait_for_spawn(spawned_rx).await;
715 let payload = wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest {
716 process_id,
717 chunk,
718 });
719 if let Err(error) = agent.transport().request_wire(ownership, payload).await {
720 tracing::warn!(?error, "write_shell failed");
721 }
722 });
723
724 Ok(())
725 }
726
727 pub async fn write_shell_awaited(
731 &self,
732 shell_id: &str,
733 data: StdinInput,
734 ) -> std::result::Result<(), ClientError> {
735 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
736 let chunk = stdin_chunk(data);
737 tracing::debug!(shell_id, "write_shell_awaited: waiting for spawn gate");
738 wait_for_spawn(spawned_rx).await;
739 tracing::debug!(shell_id, "write_shell_awaited: issuing wire write");
740 let payload =
741 wire::RequestPayload::WriteStdinRequest(wire::WriteStdinRequest { process_id, chunk });
742 let response = self
743 .transport()
744 .request_wire(self.vm_ownership(), payload)
745 .await?;
746 tracing::debug!(shell_id, "write_shell_awaited: wire write acked");
747 match response {
748 wire::ResponsePayload::RejectedResponse(rejected) => Err(rejected_to_error(rejected)),
749 _ => Ok(()),
750 }
751 }
752
753 pub fn on_shell_data(
758 &self,
759 shell_id: &str,
760 mut handler: impl FnMut(ShellData) + Send + 'static,
761 ) -> std::result::Result<crate::stream::Subscription, ClientError> {
762 let mut rx = self
763 .inner()
764 .shells
765 .read(shell_id, |_, entry| entry.data_tx.subscribe())
766 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))?;
767 let shell_id = shell_id.to_string();
768 let task = tokio::spawn(async move {
769 loop {
770 match rx.recv().await {
771 Ok(data) => handler(ShellData {
772 shell_id: shell_id.clone(),
773 data,
774 }),
775 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
776 Err(tokio::sync::broadcast::error::RecvError::Closed) => return,
777 }
778 }
779 });
780 Ok(crate::stream::Subscription::new(move || task.abort()))
781 }
782
783 pub fn on_shell_stderr(
788 &self,
789 shell_id: &str,
790 mut handler: impl FnMut(ShellData) + Send + 'static,
791 ) -> std::result::Result<crate::stream::Subscription, ClientError> {
792 let mut rx = self
793 .inner()
794 .shells
795 .read(shell_id, |_, entry| entry.stderr_tx.subscribe())
796 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))?;
797 let shell_id = shell_id.to_string();
798 let task = tokio::spawn(async move {
799 loop {
800 match rx.recv().await {
801 Ok(data) => handler(ShellData {
802 shell_id: shell_id.clone(),
803 data,
804 }),
805 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
806 Err(tokio::sync::broadcast::error::RecvError::Closed) => return,
807 }
808 }
809 });
810 Ok(crate::stream::Subscription::new(move || task.abort()))
811 }
812
813 pub fn on_shell_exit(
814 &self,
815 shell_id: &str,
816 handler: impl FnOnce(ShellExit) + Send + 'static,
817 ) -> std::result::Result<crate::stream::Subscription, ClientError> {
818 let mut rx = self
819 .inner()
820 .shells
821 .read(shell_id, |_, entry| entry.exit_tx.subscribe())
822 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))?;
823 if let Some(exit_code) = *rx.borrow() {
824 handler(ShellExit {
825 shell_id: shell_id.to_string(),
826 exit_code,
827 });
828 return Ok(crate::stream::Subscription::noop());
829 }
830 let shell_id = shell_id.to_string();
831 let task = tokio::spawn(async move {
832 while rx.changed().await.is_ok() {
833 if let Some(exit_code) = *rx.borrow() {
834 handler(ShellExit {
835 shell_id,
836 exit_code,
837 });
838 return;
839 }
840 }
841 });
842 Ok(crate::stream::Subscription::new(move || task.abort()))
843 }
844
845 pub fn resize_shell(
849 &self,
850 shell_id: &str,
851 cols: u16,
852 rows: u16,
853 ) -> std::result::Result<(), ClientError> {
854 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
856
857 let agent = self.clone();
858 let ownership = self.vm_ownership();
859 tokio::spawn(async move {
860 wait_for_spawn(spawned_rx).await;
861 let payload = wire::RequestPayload::ResizePtyRequest(wire::ResizePtyRequest {
862 process_id,
863 cols,
864 rows,
865 });
866 if let Err(error) = agent.transport().request_wire(ownership, payload).await {
867 tracing::warn!(?error, "resize_shell failed");
868 }
869 });
870
871 Ok(())
872 }
873
874 pub async fn wait_shell(&self, shell_id: &str) -> std::result::Result<i32, ClientError> {
878 let exit_rx = self
879 .inner()
880 .shells
881 .read(shell_id, |_, entry| entry.exit_tx.subscribe());
882 let Some(mut exit_rx) = exit_rx else {
883 let retained = self.inner().closed_shell_exit_codes.lock();
885 return retained
886 .iter()
887 .rev()
888 .find(|(id, _)| id == shell_id)
889 .map(|(_, code)| *code)
890 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()));
891 };
892 loop {
893 if let Some(code) = *exit_rx.borrow_and_update() {
894 return Ok(code);
895 }
896 if exit_rx.changed().await.is_err() {
897 let retained = self.inner().closed_shell_exit_codes.lock();
900 return retained
901 .iter()
902 .rev()
903 .find(|(id, _)| id == shell_id)
904 .map(|(_, code)| *code)
905 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()));
906 }
907 }
908 }
909
910 pub fn close_shell(&self, shell_id: &str) -> std::result::Result<(), ClientError> {
913 let (process_id, spawned_rx) = self.shell_wire_handle(shell_id)?;
914
915 self.inner().shells.remove(shell_id);
918
919 let agent = self.clone();
921 let ownership = self.vm_ownership();
922 tokio::spawn(async move {
923 wait_for_spawn(spawned_rx).await;
924 let payload = wire::RequestPayload::KillProcessRequest(wire::KillProcessRequest {
925 process_id,
926 signal: String::from("SIGTERM"),
927 });
928 if let Err(error) = agent.transport().request_wire(ownership, payload).await {
929 tracing::warn!(?error, "close_shell kill failed");
930 }
931 });
932
933 Ok(())
934 }
935
936 fn shell_wire_handle(
939 &self,
940 shell_id: &str,
941 ) -> std::result::Result<(String, tokio::sync::watch::Receiver<bool>), ClientError> {
942 self.inner()
943 .shells
944 .read(shell_id, |_, entry| {
945 (entry.process_id.clone(), entry.spawned_tx.subscribe())
946 })
947 .ok_or_else(|| ClientError::ShellNotFound(shell_id.to_string()))
948 }
949}
950
951async fn wait_for_spawn(mut spawned_rx: tokio::sync::watch::Receiver<bool>) {
954 if *spawned_rx.borrow() {
955 return;
956 }
957 while spawned_rx.changed().await.is_ok() {
958 if *spawned_rx.borrow() {
959 return;
960 }
961 }
962}
963
964async fn terminal_process_finished(agent: &AgentOs, pid: u32) -> bool {
965 match agent.all_processes().await {
966 Ok(processes) => match processes.into_iter().find(|process| process.pid == pid) {
967 Some(process) => process.status != ProcessStatus::Running,
968 None => true,
969 },
970 Err(error) => {
971 tracing::warn!(?error, pid, "terminal process snapshot failed");
972 false
973 }
974 }
975}
976
977#[cfg(test)]
978mod tests {
979 use super::*;
980
981 #[test]
982 fn reserve_counter_enforces_limit_and_release_reopens_slot() {
983 let counter = AtomicUsize::new(0);
984
985 assert!(try_reserve_counter(&counter, 2));
986 assert!(try_reserve_counter(&counter, 2));
987 assert!(!try_reserve_counter(&counter, 2));
988 release_counter(&counter);
989 assert!(try_reserve_counter(&counter, 2));
990 }
991}