diff --git a/architecture/gateway.md b/architecture/gateway.md index 62eceed3eb..e1fb2d4837 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -792,6 +792,26 @@ alone — instead an unanswered keepalive on a wedged or orphaned relay closes t channel and returns the exec with an error. Once a command reports its exit status, the gateway also bounds how long it waits for the trailing channel close. +Interactive exec treats normal request-stream EOF as the end of stdin and resize +input. The gateway sends SSH EOF while keeping the output channel open until +command completion. Input errors terminate the operation rather than masquerading +as normal EOF. The input and output pumps are owned by the exec operation, so +timeout or response abandonment cannot leave a detached stdin task behind. +The pumps share polling fairly, and request processing yields cooperatively even +for ignored resize messages, so sustained input cannot monopolize the operation. + +Go and TypeScript interactive-exec helpers distinguish process exit from stream +completion. They consume the final gRPC status before reporting success and retain +an observed exit code if transport completion fails. Callers must drain output +concurrently with waiting for completion. + +TypeScript starts interactive exec eagerly and uses a bounded output queue between +the background receiver and the consumer. Cancellation wakes a receiver blocked +on that queue. Go exposes input closure through an optional session capability, +preserving the original interface for existing implementations. TypeScript also +preserves its original session interface; SDK-created sessions expose lifecycle +controls through an extended interface. + `ForwardTcp` is the client-facing byte stream for SSH and service forwarding. The first frame is a `TcpForwardInit` that carries the sandbox ID, an authorization token from `CreateSshSession`, and an explicit target: diff --git a/crates/openshell-server/src/grpc/interactive_exec_tests.rs b/crates/openshell-server/src/grpc/interactive_exec_tests.rs new file mode 100644 index 0000000000..63be73ff59 --- /dev/null +++ b/crates/openshell-server/src/grpc/interactive_exec_tests.rs @@ -0,0 +1,537 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Exercise the gateway's actual SSH input/output bridge with a controllable peer. + +use super::*; +use openshell_core::proto::{exec_sandbox_event, exec_sandbox_input}; +use russh::server::{Auth, ChannelOpenHandle, Handler, Msg, Session}; +use std::time::Duration; + +struct ExecPeer { + channel: Option>, + echo_tx: Option>>, + input: Arc>>, + eof: Arc, + release_output: Arc, + output_task: Option>, +} + +impl Drop for ExecPeer { + fn drop(&mut self) { + if let Some(task) = self.output_task.take() { + task.abort(); + } + } +} + +impl Handler for ExecPeer { + type Error = russh::Error; + + async fn auth_none(&mut self, _user: &str) -> Result { + Ok(Auth::Accept) + } + + async fn channel_open_session( + &mut self, + channel: russh::Channel, + reply: ChannelOpenHandle, + _session: &mut Session, + ) -> Result<(), Self::Error> { + reply.accept().await; + self.channel = Some(channel); + Ok(()) + } + + async fn exec_request( + &mut self, + channel: russh::ChannelId, + data: &[u8], + session: &mut Session, + ) -> Result<(), Self::Error> { + session.channel_success(channel)?; + if data == b"duplex" { + let channel = self.channel.take().unwrap(); + let release = self.release_output.clone(); + let (echo_tx, mut echo_rx) = mpsc::unbounded_channel::>(); + self.echo_tx = Some(echo_tx); + self.output_task = Some(tokio::spawn(async move { + let (reader, writer) = channel.split(); + // The handler receives stdin independently of output flow + // control, as a real process with separate I/O pumps would. + drop(reader); + while let Some(data) = echo_rx.recv().await { + writer.data_bytes(data.clone()).await.unwrap(); + writer.extended_data_bytes(1, data).await.unwrap(); + } + release.notified().await; + writer.exit_status(7).await.unwrap(); + writer.close().await.unwrap(); + })); + } else { + self.channel.take(); + } + if data == b"early" { + session.exit_status_request(channel, 0)?; + session.close(channel)?; + return Ok(()); + } + // Signal that SSH setup is complete without depending on input EOF. + session.data(channel, b"ready".to_vec())?; + Ok(()) + } + + async fn data( + &mut self, + _channel: russh::ChannelId, + data: &[u8], + _session: &mut Session, + ) -> Result<(), Self::Error> { + self.input.lock().unwrap().extend_from_slice(data); + if let Some(tx) = &self.echo_tx { + tx.send(data.to_vec()).unwrap(); + } + Ok(()) + } + + async fn channel_eof( + &mut self, + channel: russh::ChannelId, + session: &mut Session, + ) -> Result<(), Self::Error> { + self.eof.store(true, Ordering::SeqCst); + self.echo_tx.take(); + if self.output_task.is_some() { + return Ok(()); + } + let release = self.release_output.clone(); + let handle = session.handle(); + self.output_task = Some(tokio::spawn(async move { + release.notified().await; + if handle.data(channel, b"after eof".to_vec()).await.is_err() { + return; + } + let _ = handle + .extended_data(channel, 1, b"stderr after eof".to_vec()) + .await; + let _ = handle.exit_status_request(channel, 7).await; + let _ = handle.eof(channel).await; + let _ = handle.close(channel).await; + })); + Ok(()) + } + + async fn channel_close( + &mut self, + _channel: russh::ChannelId, + _session: &mut Session, + ) -> Result<(), Self::Error> { + // A client CLOSE before release must really prevent later output. + if let Some(task) = self.output_task.take() { + task.abort(); + } + Ok(()) + } +} + +struct Fixture { + port: u16, + input: Arc>>, + eof: Arc, + release_output: Arc, + server: tokio::task::JoinHandle<()>, +} + +impl Drop for Fixture { + fn drop(&mut self) { + self.server.abort(); + } +} + +impl Fixture { + async fn new() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let mut config = russh::server::Config::default(); + config + .keys + .push(russh::keys::ssh_key::private::Ed25519Keypair::from_seed(&rand::random()).into()); + let input = Arc::new(std::sync::Mutex::new(Vec::new())); + let eof = Arc::new(AtomicBool::new(false)); + let release_output = Arc::new(tokio::sync::Notify::new()); + let handler = ExecPeer { + channel: None, + echo_tx: None, + input: input.clone(), + eof: eof.clone(), + release_output: release_output.clone(), + output_task: None, + }; + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + set_tcp_nodelay_best_effort(&socket); + if let Ok(session) = russh::server::run_stream(Arc::new(config), socket, handler).await + { + let _ = session.await; + } + }); + Self { + port, + input, + eof, + release_output, + server, + } + } +} + +async fn ready(rx: &mut mpsc::Receiver>) { + let event = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!( + matches!(event.payload, Some(exec_sandbox_event::Payload::Stdout(s)) if s.data == b"ready") + ); +} + +#[tokio::test] +async fn interactive_exec_drains_stdout_and_stderr_after_input_eof() { + let fixture = Fixture::new().await; + let (input_tx, input_rx) = mpsc::channel(2); + let (output_tx, mut output_rx) = mpsc::channel(2); + let exec = run_interactive_exec_with_russh( + fixture.port, + "test", + ReceiverStream::new(input_rx), + false, + false, + 0, + 0, + output_tx, + ); + let exercise = async { + ready(&mut output_rx).await; + input_tx + .send(Ok(ExecSandboxInput { + payload: Some(exec_sandbox_input::Payload::Stdin(b"input".to_vec())), + })) + .await + .unwrap(); + drop(input_tx); + while !fixture.eof.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + // Keep the peer silent briefly after EOF so a premature SSH CLOSE is + // processed before output is released. + tokio::time::sleep(Duration::from_millis(50)).await; + fixture.release_output.notify_one(); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + while let Some(event) = output_rx.recv().await { + match event.unwrap().payload.unwrap() { + exec_sandbox_event::Payload::Stdout(s) => stdout.extend(s.data), + exec_sandbox_event::Payload::Stderr(s) => stderr.extend(s.data), + other @ exec_sandbox_event::Payload::Exit(_) => { + panic!("unexpected event: {other:?}") + } + } + } + assert_eq!(stdout, b"after eof"); + assert_eq!(stderr, b"stderr after eof"); + assert_eq!(*fixture.input.lock().unwrap(), b"input"); + }; + let (result, ()) = tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(exec, exercise) + }) + .await + .unwrap(); + assert_eq!(result.unwrap(), 7); +} + +#[tokio::test] +async fn interactive_exec_makes_progress_in_both_directions_before_eof() { + const CHUNKS: usize = 128; + const CHUNK_SIZE: usize = 64 * 1024; + const BATCH: usize = 4; + let fixture = Fixture::new().await; + let (input_tx, input_rx) = mpsc::channel(16); + let (output_tx, mut output_rx) = mpsc::channel(2); + let drained = tokio::sync::Notify::new(); + let progress = std::cell::Cell::new((0, 0)); + let exec = run_interactive_exec_with_russh( + fixture.port, + "duplex", + ReceiverStream::new(input_rx), + false, + false, + 0, + 0, + output_tx, + ); + let writer = async { + for chunk in 1..=CHUNKS { + input_tx + .send(Ok(ExecSandboxInput { + payload: Some(exec_sandbox_input::Payload::Stdin(vec![b'x'; CHUNK_SIZE])), + })) + .await + .unwrap(); + if chunk % BATCH == 0 { + drained.notified().await; + } + } + // Keep the request stream open until BOTH output streams have drained. + // Bound in-flight data to exercise sustained interactive traffic without + // saturating both ends of the fixture's SSH transport simultaneously. + drop(input_tx); + }; + let reader = async { + ready(&mut output_rx).await; + let mut stdout = 0; + let mut stderr = 0; + let mut acknowledged = 0; + while stdout < CHUNKS * CHUNK_SIZE || stderr < CHUNKS * CHUNK_SIZE { + let bytes = match output_rx.recv().await.unwrap().unwrap().payload.unwrap() { + exec_sandbox_event::Payload::Stdout(s) => { + stdout += s.data.len(); + s.data + } + exec_sandbox_event::Payload::Stderr(s) => { + stderr += s.data.len(); + s.data + } + event @ exec_sandbox_event::Payload::Exit(_) => { + panic!("unexpected event: {event:?}") + } + }; + assert!(bytes.iter().all(|b| *b == b'x')); + progress.set((stdout, stderr)); + assert!(!fixture.eof.load(Ordering::SeqCst)); + if stdout.min(stderr) >= acknowledged + BATCH * CHUNK_SIZE { + acknowledged += BATCH * CHUNK_SIZE; + drained.notify_one(); + } + } + assert_eq!(stdout, CHUNKS * CHUNK_SIZE); + assert_eq!(stderr, CHUNKS * CHUNK_SIZE); + fixture.release_output.notify_one(); + while output_rx.recv().await.is_some() {} + }; + let (result, (), ()) = tokio::time::timeout(Duration::from_secs(30), async { + tokio::join!(exec, writer, reader) + }) + .await + .unwrap_or_else(|_| { + panic!( + "duplex stalled: input={}, output={:?}, eof={}", + fixture.input.lock().unwrap().len(), + progress.get(), + fixture.eof.load(Ordering::SeqCst) + ) + }); + assert_eq!(result.unwrap(), 7); + assert_eq!(fixture.input.lock().unwrap().len(), CHUNKS * CHUNK_SIZE); +} + +#[tokio::test] +async fn interactive_exec_ready_resize_stream_does_not_starve_output() { + use futures::StreamExt; + use std::sync::atomic::AtomicUsize; + + // Finite to make a regression fail rather than wedge the runtime forever. + // These frames have no SSH write await because this session has no PTY. + const FRAMES: usize = 100_000; + let fixture = Fixture::new().await; + let consumed = AtomicUsize::new(0); + let input = futures::stream::repeat_with(|| { + consumed.fetch_add(1, Ordering::SeqCst); + Ok(ExecSandboxInput { + payload: Some(exec_sandbox_input::Payload::Resize( + openshell_core::proto::ExecSandboxWindowResize::default(), + )), + }) + }) + .take(FRAMES); + let (output_tx, mut output_rx) = mpsc::channel(2); + let exec = + run_interactive_exec_with_russh(fixture.port, "test", input, false, false, 0, 0, output_tx); + let reader = async { + ready(&mut output_rx).await; + assert!( + consumed.load(Ordering::SeqCst) < FRAMES, + "output must be delivered before the continuously ready input ends" + ); + drop(output_rx); + }; + let (result, ()) = + tokio::time::timeout(Duration::from_secs(5), async { tokio::join!(exec, reader) }) + .await + .unwrap(); + assert_eq!(result.unwrap_err().code(), tonic::Code::Cancelled); +} + +#[tokio::test] +async fn interactive_exec_input_error_is_not_graceful_eof() { + for message in [ + Err(Status::cancelled("input cancelled")), + Ok(ExecSandboxInput { payload: None }), + Ok(ExecSandboxInput { + payload: Some(exec_sandbox_input::Payload::Start( + ExecSandboxRequest::default(), + )), + }), + ] { + let expected = message + .as_ref() + .err() + .map_or(tonic::Code::InvalidArgument, Status::code); + let fixture = Fixture::new().await; + let (input_tx, input_rx) = mpsc::channel(1); + let (output_tx, mut output_rx) = mpsc::channel(1); + let exec = run_interactive_exec_with_russh( + fixture.port, + "test", + ReceiverStream::new(input_rx), + false, + false, + 0, + 0, + output_tx, + ); + let exercise = async { + ready(&mut output_rx).await; + input_tx.send(message).await.unwrap(); + while output_rx.recv().await.is_some() {} + }; + let (result, ()) = tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(exec, exercise) + }) + .await + .unwrap(); + assert_eq!(result.unwrap_err().code(), expected); + assert!(!fixture.eof.load(Ordering::SeqCst)); + } +} + +#[tokio::test] +async fn interactive_exec_response_drop_cancels_idle_input() { + let fixture = Fixture::new().await; + let (input_tx, input_rx) = mpsc::channel(1); + let (output_tx, mut output_rx) = mpsc::channel(1); + let exec = run_interactive_exec_with_russh( + fixture.port, + "test", + ReceiverStream::new(input_rx), + false, + false, + 0, + 0, + output_tx, + ); + let exercise = async { + ready(&mut output_rx).await; + drop(output_rx); + }; + let (result, ()) = tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(exec, exercise) + }) + .await + .unwrap(); + assert_eq!(result.unwrap_err().code(), tonic::Code::Cancelled); + assert!(input_tx.is_closed(), "stdin receiver must not outlive exec"); + assert!(!fixture.eof.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn interactive_exec_parent_abort_drops_input_receiver() { + let fixture = Fixture::new().await; + let (input_tx, input_rx) = mpsc::channel(1); + let (output_tx, mut output_rx) = mpsc::channel(1); + let exec = tokio::spawn(run_interactive_exec_with_russh( + fixture.port, + "test", + ReceiverStream::new(input_rx), + false, + false, + 0, + 0, + output_tx, + )); + ready(&mut output_rx).await; + // An operation timeout drops the same future as this task abort. + exec.abort(); + assert!(exec.await.unwrap_err().is_cancelled()); + assert!( + input_tx.is_closed(), + "aborted exec must not leave a stdin task alive" + ); +} + +#[tokio::test] +async fn interactive_exec_response_drop_unblocks_full_output_queue() { + let fixture = Fixture::new().await; + let (input_tx, input_rx) = mpsc::channel(1); + let (output_tx, mut output_rx) = mpsc::channel(1); + // Use an empty request stream: EOF releases the peer's output once notified. + drop(input_tx); + let exec = run_interactive_exec_with_russh( + fixture.port, + "test", + ReceiverStream::new(input_rx), + false, + false, + 0, + 0, + output_tx, + ); + let exercise = async { + ready(&mut output_rx).await; + fixture.release_output.notify_one(); + while output_rx.is_empty() { + tokio::task::yield_now().await; + } + // stdout fills the one-slot queue while stderr still needs delivery. + drop(output_rx); + }; + let (result, ()) = tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(exec, exercise) + }) + .await + .unwrap(); + assert_eq!(result.unwrap_err().code(), tonic::Code::Cancelled); +} + +#[tokio::test] +async fn interactive_exec_early_exit_does_not_wait_for_stdin() { + let fixture = Fixture::new().await; + let (input_tx, input_rx) = mpsc::channel(1); + input_tx + .send(Ok(ExecSandboxInput { + payload: Some(exec_sandbox_input::Payload::Stdin(vec![ + b'x'; + 4 * 1024 * 1024 + ])), + })) + .await + .unwrap(); + let (output_tx, _output_rx) = mpsc::channel(1); + let result = tokio::time::timeout( + Duration::from_secs(5), + run_interactive_exec_with_russh( + fixture.port, + "early", + ReceiverStream::new(input_rx), + false, + false, + 0, + 0, + output_tx, + ), + ) + .await + .unwrap(); + assert_eq!(result.unwrap(), 0); + assert!(input_tx.is_closed()); +} diff --git a/crates/openshell-server/src/grpc/sandbox.rs b/crates/openshell-server/src/grpc/sandbox.rs index ea470a7408..1611890ba5 100644 --- a/crates/openshell-server/src/grpc/sandbox.rs +++ b/crates/openshell-server/src/grpc/sandbox.rs @@ -76,6 +76,10 @@ const TCP_FORWARD_CHUNK_SIZE: usize = 64 * 1024; const NO_LOGIN_SHELL_ENV: (&str, &str) = ("OPENSHELL_NO_LOGIN_SHELL", "1"); const MAX_TEMPLATES_PER_WORKSPACE: u32 = 1000; +#[cfg(test)] +#[path = "interactive_exec_tests.rs"] +mod interactive_exec_tests; + #[derive(Debug)] pub struct WatchSandboxStream { receiver: ReceiverStream>, @@ -2816,7 +2820,7 @@ async fn stream_interactive_exec_over_relay( )), })) .await; - let _ = proxy_task.await; + finish_interactive_exec_proxy(proxy_task).await; return Ok(()); } } else { @@ -2826,12 +2830,12 @@ async fn stream_interactive_exec_over_relay( let exit_code = match exec_result { Ok(code) => code, Err(status) => { - let _ = proxy_task.await; + finish_interactive_exec_proxy(proxy_task).await; return Err(status); } }; - let _ = proxy_task.await; + finish_interactive_exec_proxy(proxy_task).await; let _ = tx .send(Ok(ExecSandboxEvent { @@ -2844,17 +2848,28 @@ async fn stream_interactive_exec_over_relay( Ok(()) } +async fn finish_interactive_exec_proxy(mut task: tokio::task::JoinHandle<()>) { + if tokio::time::timeout(EXEC_POST_EXIT_CLOSE_TIMEOUT, &mut task) + .await + .is_err() + { + task.abort(); + let _ = task.await; + } +} + #[allow(clippy::too_many_arguments)] async fn run_interactive_exec_with_russh( local_proxy_port: u16, command: &str, - mut input_stream: tonic::Streaming, + mut input_stream: impl futures::Stream> + Unpin, request_tty: bool, no_login_shell: bool, cols: u32, rows: u32, tx: mpsc::Sender>, ) -> Result { + use futures::StreamExt; use openshell_core::proto::exec_sandbox_input::Payload; use russh::ChannelMsg; @@ -2930,83 +2945,125 @@ async fn run_interactive_exec_with_russh( let (mut read_half, write_half) = channel.split(); - let stdin_task = tokio::spawn(async move { - while let Ok(Some(msg)) = input_stream.message().await { + // Keep both directions independently polled, but owned by this operation. + // A detached stdin task would survive operation timeout and retain the SSH + // channel. Normal request EOF closes stdin only, never the output channel. + let input = async { + while let Some(msg) = input_stream.next().await { + // Even ignored non-PTY resize frames must cooperate with output + // and cancellation when the request stream stays continuously ready. + tokio::task::consume_budget().await; + let msg = msg?; match msg.payload { Some(Payload::Stdin(data)) => { - if write_half.data(std::io::Cursor::new(data)).await.is_err() { - break; - } + write_half + .data(std::io::Cursor::new(data)) + .await + .map_err(|_| { + Status::unavailable("exec relay failed while writing stdin") + })?; } Some(Payload::Resize(resize)) => { if request_tty { - let _ = write_half + write_half .window_change(resize.cols, resize.rows, 0, 0) - .await; + .await + .map_err(|_| Status::unavailable("exec relay failed while resizing"))?; } } - Some(Payload::Start(_)) | None => {} + Some(Payload::Start(_)) | None => { + return Err(Status::invalid_argument( + "expected stdin or resize after exec start", + )); + } } } - let _ = write_half.eof().await; - let _ = write_half.close().await; - }); + write_half + .eof() + .await + .map_err(|_| Status::unavailable("exec relay failed while closing stdin")) + }; - let mut exit_code: Option = None; - loop { - // Bound the post-ExitStatus wait against a lost Close. - let msg = if exit_code.is_some() { - match tokio::time::timeout(EXEC_POST_EXIT_CLOSE_TIMEOUT, read_half.wait()).await { - Ok(Some(msg)) => msg, - Ok(None) | Err(_) => break, - } - } else { - match read_half.wait().await { - Some(msg) => msg, - None => break, - } - }; - match msg { - ChannelMsg::Data { data } => { - let event = Ok(ExecSandboxEvent { - payload: Some(openshell_core::proto::exec_sandbox_event::Payload::Stdout( - ExecSandboxStdout { - data: data.to_vec(), - }, - )), - }); - if tx.send(event).await.is_err() { - break; + let output = async { + let mut exit_code: Option = None; + loop { + // Bound the post-ExitStatus wait against a lost Close. + let msg = if exit_code.is_some() { + match tokio::time::timeout(EXEC_POST_EXIT_CLOSE_TIMEOUT, read_half.wait()).await { + Ok(Some(msg)) => msg, + Ok(None) | Err(_) => break, } - } - ChannelMsg::ExtendedData { data, .. } => { - let event = Ok(ExecSandboxEvent { - payload: Some(openshell_core::proto::exec_sandbox_event::Payload::Stderr( - ExecSandboxStderr { - data: data.to_vec(), - }, - )), - }); - if tx.send(event).await.is_err() { - break; + } else { + match read_half.wait().await { + Some(msg) => msg, + None => break, } + }; + match msg { + ChannelMsg::Data { data } => { + let event = Ok(ExecSandboxEvent { + payload: Some(openshell_core::proto::exec_sandbox_event::Payload::Stdout( + ExecSandboxStdout { + data: data.to_vec(), + }, + )), + }); + if tx.send(event).await.is_err() { + break; + } + } + ChannelMsg::ExtendedData { data, .. } => { + let event = Ok(ExecSandboxEvent { + payload: Some(openshell_core::proto::exec_sandbox_event::Payload::Stderr( + ExecSandboxStderr { + data: data.to_vec(), + }, + )), + }); + if tx.send(event).await.is_err() { + break; + } + } + ChannelMsg::ExitStatus { exit_status } => { + let converted = i32::try_from(exit_status).unwrap_or(i32::MAX); + exit_code = Some(converted); + } + ChannelMsg::Close => break, + _ => {} } - ChannelMsg::ExitStatus { exit_status } => { - let converted = i32::try_from(exit_status).unwrap_or(i32::MAX); - exit_code = Some(converted); - } - ChannelMsg::Close => break, - _ => {} } - } - stdin_task.abort(); + exec_loop_result(exit_code) + }; - let _ = client - .disconnect(russh::Disconnect::ByApplication, "exec complete", "en") - .await; + let result = { + tokio::pin!(input, output); + let exchange = async { + tokio::select! { + result = &mut input => { + result?; + output.await + } + result = &mut output => result, + } + }; + tokio::select! { + biased; + () = tx.closed() => Err(Status::cancelled("exec response stream closed")), + result = exchange => result, + } + }; - exec_loop_result(exit_code) + // EOF above deliberately leaves this channel open until output completes. + // Bound cleanup even if the SSH peer is no longer making progress. + let _ = tokio::time::timeout(EXEC_POST_EXIT_CLOSE_TIMEOUT, write_half.close()).await; + let _ = tokio::time::timeout( + EXEC_POST_EXIT_CLOSE_TIMEOUT, + client.disconnect(russh::Disconnect::ByApplication, "exec complete", "en"), + ) + .await; + + result } /// Create a localhost SSH proxy that bridges to a relay `DuplexStream`. diff --git a/docs/sandboxes/manage-sandboxes.mdx b/docs/sandboxes/manage-sandboxes.mdx index 86968a9910..366b05499f 100644 --- a/docs/sandboxes/manage-sandboxes.mdx +++ b/docs/sandboxes/manage-sandboxes.mdx @@ -319,6 +319,12 @@ OpenShell-managed SSH include file instead of cluttering your main ## Execute a Command in a Sandbox +For raw `ExecSandboxInteractive` clients, ending the request stream closes stdin +and terminal resize input. Continue reading the response to drain stdout/stderr +and receive the exit event and final gRPC status. Input EOF does not immediately +close the SSH output channel and is distinct from cancelling the RPC. With a PTY, +input closure is not equivalent to sending a terminal Ctrl-D keystroke. + Run a one-shot command inside a running sandbox without opening an interactive shell: ```shell diff --git a/e2e/python/test_sandbox_api.py b/e2e/python/test_sandbox_api.py index 9b48152961..a0a8eeb2f3 100644 --- a/e2e/python/test_sandbox_api.py +++ b/e2e/python/test_sandbox_api.py @@ -198,6 +198,49 @@ def requests(): assert b"TTT" in stdout + stderr +def test_interactive_exec_drains_output_after_request_eof( + sandbox: Callable[..., Sandbox], + sandbox_client: SandboxClient, +) -> None: + with sandbox(delete_on_exit=True) as sb: + + def requests(): + yield openshell_pb2.ExecSandboxInput( + start=openshell_pb2.ExecSandboxRequest( + sandbox_id=sb.id, + command=[ + "/bin/sh", + "-c", + "input=$(cat); sleep 0.2; " + "printf '%s' \"$input\"; printf 'drained-stderr' >&2; exit 7", + ], + tty=False, + execution_timeout=duration_pb2.Duration(seconds=20), + ) + ) + yield openshell_pb2.ExecSandboxInput(stdin=b"drained-stdout") + # End requests before the command emits output. The receive + # direction must survive long enough to drain both output streams. + + stdout = bytearray() + stderr = bytearray() + exit_codes: list[int] = [] + for event in sandbox_client._stub.ExecSandboxInteractive( + requests(), timeout=30 + ): + assert not exit_codes, "received another event after terminal exit" + payload = event.WhichOneof("payload") + if payload == "stdout": + stdout.extend(event.stdout.data) + elif payload == "stderr": + stderr.extend(event.stderr.data) + elif payload == "exit": + exit_codes.append(event.exit.exit_code) + assert stdout == b"drained-stdout" + assert stderr == b"drained-stderr" + assert exit_codes == [7] + + def test_list_scoped_and_for_all_workspaces( sandbox_client: SandboxClient, workspace_client: WorkspaceClient, diff --git a/sdk/go/README.md b/sdk/go/README.md index 6d695c6b4b..8e82aa2c34 100644 --- a/sdk/go/README.md +++ b/sdk/go/README.md @@ -278,6 +278,18 @@ consumers import a single package. See the [Architecture](https://ro14nd.de/open ## Features +Use `CloseInteractiveInput(session)` to close stdin and resize input while keeping +output readable. SDK sessions implement the optional `InteractiveSessionControl` +interface (`CloseWrite()` and `Cancel()`); the original `InteractiveSession` +interface remains unchanged for existing mocks and wrappers. Input closure returns +`ErrorUnimplemented` for sessions without that capability and leaves them open. +`CancelInteractive(session)` uses `Cancel()` when available and otherwise calls +`Close()`. SDK close/cancel operations are idempotent; writes and resizes after +input closure return `io.ErrClosedPipe`. Drain `Read` concurrently with waiting for `ExitCode()`. +`ExitCode()` waits for final gRPC status and returns any observed process exit code +alongside a later stream error. An exit event alone does not establish successful +stream completion. + | Feature | Interface | Docs | |---------|-----------|------| | Sandbox lifecycle (create, get, list, delete, watch, wait) | `SandboxInterface` | [Sandboxes](https://ro14nd.de/openshell-sdk-go/api/sandboxes.html) | diff --git a/sdk/go/docs/src/api/exec.md b/sdk/go/docs/src/api/exec.md index 0d088a6999..72ff49b347 100644 --- a/sdk/go/docs/src/api/exec.md +++ b/sdk/go/docs/src/api/exec.md @@ -139,8 +139,22 @@ type InteractiveSession interface { | `Read` | Reads output from the process into the provided buffer. | | `Write` | Sends input to the process. | | `Resize` | Updates the terminal dimensions (columns and rows). | -| `ExitCode` | Returns the process exit code after the session ends. | -| `Close` | Closes the session and releases resources. | +| `ExitCode` | Waits for final gRPC status; returns any observed exit code alongside a later stream error. | +| `Close` | Aborts the RPC and releases resources. | + +SDK sessions also implement the optional `InteractiveSessionControl` interface, +which embeds `InteractiveSession` and adds `CloseWrite() error` and `Cancel() error`. +Existing custom implementations only need the original methods above. + +Call `CloseInteractiveInput(session)` to end stdin and resize input without +cancelling output. It invokes `CloseWrite()` when supported; otherwise it returns +`ErrorUnimplemented` without closing the session. Call `CancelInteractive(session)` +to invoke `Cancel()` when supported, falling back to `Close()` for legacy sessions. +The SDK's close and cancel methods are idempotent. + +Drain `Read` concurrently with waiting for `ExitCode`; bounded output buffers can +otherwise prevent completion. After input closure, writes and resizes return +`io.ErrClosedPipe`. Request EOF is not equivalent to a terminal Ctrl-D keystroke. ## ExecChunk diff --git a/sdk/go/openshell/v1/exec.go b/sdk/go/openshell/v1/exec.go index 217a1dc9b3..f30975d15c 100644 --- a/sdk/go/openshell/v1/exec.go +++ b/sdk/go/openshell/v1/exec.go @@ -27,10 +27,40 @@ type InteractiveSession interface { Read(p []byte) (int, error) Write(p []byte) (int, error) Resize(cols, rows uint32) error + // ExitCode waits for final stream completion. An observed exit code is + // returned alongside any later transport error. Drain Read concurrently. ExitCode() (int, error) Close() error } +// InteractiveSessionControl adds optional input closure and cancellation to +// InteractiveSession without requiring existing implementations to add methods. +// Sessions returned by this SDK implement both interfaces. +type InteractiveSessionControl interface { + InteractiveSession + // CloseWrite ends stdin and resize input without cancelling output. + CloseWrite() error + // Cancel aborts the RPC. Close retains the same full-close behavior. + Cancel() error +} + +// CloseInteractiveInput ends input while preserving output when supported. +// Unsupported sessions return ErrorUnimplemented and are left open. +func CloseInteractiveInput(session InteractiveSession) error { + if closer, ok := session.(interface{ CloseWrite() error }); ok { + return closer.CloseWrite() + } + return &StatusError{Code: ErrorUnimplemented, Message: "interactive session does not support closing input independently"} +} + +// CancelInteractive aborts a session, falling back to the original Close contract. +func CancelInteractive(session InteractiveSession) error { + if canceler, ok := session.(interface{ Cancel() error }); ok { + return canceler.Cancel() + } + return session.Close() +} + // ExecInterface defines command execution operations on sandboxes. // Methods accept a sandbox name and resolve it to an ID internally. type ExecInterface interface { diff --git a/sdk/go/openshell/v1/exec_client.go b/sdk/go/openshell/v1/exec_client.go index c74fe7a5d3..fdfe4a8630 100644 --- a/sdk/go/openshell/v1/exec_client.go +++ b/sdk/go/openshell/v1/exec_client.go @@ -11,6 +11,7 @@ import ( "github.com/NVIDIA/OpenShell/sdk/go/openshell/v1/internal/converter" pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1" "google.golang.org/grpc" + "google.golang.org/grpc/status" ) type execClient struct { @@ -177,19 +178,21 @@ func (s *execStream) Close() error { // interactiveSession wraps a bidirectional streaming RPC into the InteractiveSession interface. // A background goroutine owns the Recv loop and routes events to dataCh (for Read) -// and exitCh (for ExitCode), preventing concurrent Recv calls on the stream. +// and publishes the final exit/status through done, preventing concurrent Recv calls. type interactiveSession struct { - stream grpc.BidiStreamingClient[pb.ExecSandboxInput, pb.ExecSandboxEvent] - cancel context.CancelFunc - sendMu sync.Mutex - dataCh chan []byte - exitCh chan int - done chan struct{} - errOnce sync.Once - err error - buf []byte - - exitMu sync.Mutex + stream grpc.BidiStreamingClient[pb.ExecSandboxInput, pb.ExecSandboxEvent] + cancel context.CancelFunc + sendMu sync.Mutex + inputClosed bool + inputCloseErr error + closeOnce sync.Once + closeErr error + dataCh chan []byte + done chan struct{} + errOnce sync.Once + err error + buf []byte + exitCode int hasExitCode bool } @@ -199,7 +202,6 @@ func newInteractiveSession(ctx context.Context, cancel context.CancelFunc, strea stream: stream, cancel: cancel, dataCh: make(chan []byte, 64), - exitCh: make(chan int, 1), done: make(chan struct{}), } go s.readLoop(ctx) @@ -211,6 +213,7 @@ func (s *interactiveSession) setErr(err error) { } func (s *interactiveSession) readLoop(ctx context.Context) { + defer s.cancel() defer close(s.dataCh) defer close(s.done) for { @@ -218,10 +221,18 @@ func (s *interactiveSession) readLoop(ctx context.Context) { if err != nil { if err != io.EOF { s.setErr(converter.FromGRPCError(err)) + } else if ctx.Err() != nil { + s.setErr(converter.FromGRPCError(status.FromContextError(ctx.Err()).Err())) + } else if !s.hasExitCode { + s.setErr(&StatusError{Code: ErrorInternal, Message: "stream ended without exit event"}) } return } + if s.hasExitCode { + s.setErr(&StatusError{Code: ErrorInternal, Message: "received event after exit"}) + return + } chunk, code, convErr := converter.ExecChunkFromEvent(ev) if convErr != nil { s.setErr(convErr) @@ -229,15 +240,14 @@ func (s *interactiveSession) readLoop(ctx context.Context) { } // nil chunk with no error means exit event if chunk == nil { - select { - case s.exitCh <- code: - default: - } - return + s.exitCode = code + s.hasExitCode = true + continue } select { case s.dataCh <- chunk.Data: case <-ctx.Done(): + s.setErr(converter.FromGRPCError(status.FromContextError(ctx.Err()).Err())) return } } @@ -267,6 +277,9 @@ func (s *interactiveSession) Read(p []byte) (int, error) { func (s *interactiveSession) Write(p []byte) (int, error) { s.sendMu.Lock() defer s.sendMu.Unlock() + if s.inputClosed { + return 0, io.ErrClosedPipe + } err := s.stream.Send(&pb.ExecSandboxInput{ Payload: &pb.ExecSandboxInput_Stdin{Stdin: p}, }) @@ -279,6 +292,9 @@ func (s *interactiveSession) Write(p []byte) (int, error) { func (s *interactiveSession) Resize(cols, rows uint32) error { s.sendMu.Lock() defer s.sendMu.Unlock() + if s.inputClosed { + return io.ErrClosedPipe + } err := s.stream.Send(&pb.ExecSandboxInput{ Payload: &pb.ExecSandboxInput_Resize{ Resize: &pb.ExecSandboxWindowResize{ @@ -294,41 +310,33 @@ func (s *interactiveSession) Resize(cols, rows uint32) error { } func (s *interactiveSession) ExitCode() (int, error) { - s.exitMu.Lock() + <-s.done if s.hasExitCode { - code := s.exitCode - s.exitMu.Unlock() - return code, nil - } - s.exitMu.Unlock() - - select { - case code := <-s.exitCh: - s.exitMu.Lock() - s.exitCode = code - s.hasExitCode = true - s.exitMu.Unlock() - return code, nil - case <-s.done: - select { - case code := <-s.exitCh: - s.exitMu.Lock() - s.exitCode = code - s.hasExitCode = true - s.exitMu.Unlock() - return code, nil - default: - if s.err != nil { - return -1, s.err - } - return -1, &StatusError{Code: ErrorInternal, Message: "stream ended without exit event"} - } + return s.exitCode, s.err + } + return -1, s.err +} + +func (s *interactiveSession) CloseWrite() error { + s.sendMu.Lock() + defer s.sendMu.Unlock() + if !s.inputClosed { + s.inputClosed = true + s.inputCloseErr = s.stream.CloseSend() } + return s.inputCloseErr +} + +func (s *interactiveSession) Cancel() error { + s.closeOnce.Do(func() { + // Cancel before taking sendMu so a blocked Write can release it. + s.cancel() + s.closeErr = s.CloseWrite() + <-s.done + }) + return s.closeErr } func (s *interactiveSession) Close() error { - s.cancel() - err := s.stream.CloseSend() - <-s.done - return err + return s.Cancel() } diff --git a/sdk/go/openshell/v1/exec_client_test.go b/sdk/go/openshell/v1/exec_client_test.go index e82550e9bb..d0fa0c4b83 100644 --- a/sdk/go/openshell/v1/exec_client_test.go +++ b/sdk/go/openshell/v1/exec_client_test.go @@ -85,6 +85,8 @@ type mockExecServer struct { interactiveErr error interactiveWaitInput bool interactiveBlock bool + interactiveWaitEOF bool + interactiveFinalErr error receivedInputs []*pb.ExecSandboxInput } @@ -116,6 +118,8 @@ func (s *mockExecServer) ExecSandboxInteractive(stream grpc.BidiStreamingServer[ s.mu.Lock() interactiveErr := s.interactiveErr interactiveBlock := s.interactiveBlock + waitEOF := s.interactiveWaitEOF + finalErr := s.interactiveFinalErr events := make([]*pb.ExecSandboxEvent, len(s.interactiveEvents)) copy(events, s.interactiveEvents) s.mu.Unlock() @@ -136,6 +140,17 @@ func (s *mockExecServer) ExecSandboxInteractive(stream grpc.BidiStreamingServer[ <-stream.Context().Done() return stream.Context().Err() } + if waitEOF { + for { + _, recvErr := stream.Recv() + if recvErr == io.EOF { + break + } + if recvErr != nil { + return recvErr + } + } + } s.mu.Lock() waitInput := s.interactiveWaitInput @@ -170,7 +185,7 @@ func (s *mockExecServer) ExecSandboxInteractive(stream grpc.BidiStreamingServer[ return err } } - return nil + return finalErr } func setupExecTest(t *testing.T, mock *mockExecServer) (*execClient, func()) { @@ -407,6 +422,86 @@ func TestExecInteractive_CloseCancelsReceiveStream(t *testing.T) { } } +func TestExecInteractive_CloseWriteDrainsOutput(t *testing.T) { + mock := newMockExecServer() + mock.interactiveWaitEOF = true + mock.interactiveEvents = []*pb.ExecSandboxEvent{ + {Payload: &pb.ExecSandboxEvent_Stdout{Stdout: &pb.ExecSandboxStdout{Data: []byte("after EOF")}}}, + {Payload: &pb.ExecSandboxEvent_Exit{Exit: &pb.ExecSandboxExit{ExitCode: 7}}}, + } + client, cleanup := setupExecTest(t, mock) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) + defer cancel() + session, err := client.Interactive(ctx, "default", "test-sandbox", []string{"cat"}, 0, 0) + require.NoError(t, err) + defer func() { _ = session.Close() }() + require.Implements(t, (*InteractiveSessionControl)(nil), session) + require.NoError(t, CloseInteractiveInput(session)) + require.NoError(t, CloseInteractiveInput(session)) + _, err = session.Write([]byte("late")) + require.ErrorIs(t, err, io.ErrClosedPipe) + require.ErrorIs(t, session.Resize(80, 24), io.ErrClosedPipe) + output, err := io.ReadAll(session) + require.NoError(t, err) + require.Equal(t, "after EOF", string(output)) + code, err := session.ExitCode() + require.NoError(t, err) + require.Equal(t, 7, code) +} + +func TestExecInteractive_RetainsExitOnTerminalFailure(t *testing.T) { + for _, finalErr := range []error{status.Error(codes.Unavailable, "connection lost"), status.Error(codes.Canceled, "cancelled")} { + t.Run(status.Code(finalErr).String(), func(t *testing.T) { + mock := newMockExecServer() + mock.interactiveEvents = []*pb.ExecSandboxEvent{ + {Payload: &pb.ExecSandboxEvent_Exit{Exit: &pb.ExecSandboxExit{ExitCode: 7}}}, + } + mock.interactiveFinalErr = finalErr + client, cleanup := setupExecTest(t, mock) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) + defer cancel() + session, err := client.Interactive(ctx, "default", "test-sandbox", []string{"true"}, 0, 0) + require.NoError(t, err) + defer func() { _ = session.Close() }() + code, err := session.ExitCode() + require.Error(t, err) + require.Equal(t, 7, code) + var sdkErr *StatusError + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, status.Code(finalErr), status.Code(sdkErr.Cause)) + codeAgain, errAgain := session.ExitCode() + require.Equal(t, code, codeAgain) + require.Equal(t, err, errAgain) + }) + } +} + +func TestExecInteractive_RejectsEventsAfterExit(t *testing.T) { + for name, event := range map[string]*pb.ExecSandboxEvent{ + "duplicate exit": {Payload: &pb.ExecSandboxEvent_Exit{Exit: &pb.ExecSandboxExit{ExitCode: 1}}}, + "late output": {Payload: &pb.ExecSandboxEvent_Stdout{Stdout: &pb.ExecSandboxStdout{Data: []byte("late")}}}, + } { + t.Run(name, func(t *testing.T) { + mock := newMockExecServer() + mock.interactiveEvents = []*pb.ExecSandboxEvent{ + {Payload: &pb.ExecSandboxEvent_Exit{Exit: &pb.ExecSandboxExit{ExitCode: 0}}}, event, + } + client, cleanup := setupExecTest(t, mock) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) + defer cancel() + session, err := client.Interactive(ctx, "default", "test-sandbox", []string{"true"}, 0, 0) + require.NoError(t, err) + defer func() { _ = session.Close() }() + code, err := session.ExitCode() + require.Equal(t, 0, code) + require.ErrorContains(t, err, "after exit") + }) + } +} + func TestExecInteractive_Write(t *testing.T) { mock := newMockExecServer() mock.interactiveWaitInput = true diff --git a/sdk/go/openshell/v1/exec_control_test.go b/sdk/go/openshell/v1/exec_control_test.go new file mode 100644 index 0000000000..3e2392a821 --- /dev/null +++ b/sdk/go/openshell/v1/exec_control_test.go @@ -0,0 +1,56 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package v1 + +import ( + "errors" + "io" + "testing" + + "github.com/stretchr/testify/require" +) + +// Deliberately implements only the original interface: downstream mocks must +// continue compiling without adding the optional lifecycle methods. +type legacyInteractiveSession struct { + closed bool + err error +} + +var _ InteractiveSession = (*legacyInteractiveSession)(nil) + +func (*legacyInteractiveSession) Read([]byte) (int, error) { return 0, io.EOF } +func (*legacyInteractiveSession) Write(p []byte) (int, error) { return len(p), nil } +func (*legacyInteractiveSession) Resize(uint32, uint32) error { return nil } +func (*legacyInteractiveSession) ExitCode() (int, error) { return 0, nil } +func (s *legacyInteractiveSession) Close() error { s.closed = true; return s.err } + +type controlledInteractiveSession struct { + legacyInteractiveSession + inputClosed bool + cancelled bool +} + +func (s *controlledInteractiveSession) CloseWrite() error { s.inputClosed = true; return s.err } +func (s *controlledInteractiveSession) Cancel() error { s.cancelled = true; return s.err } + +func TestInteractiveControlLegacyCompatibility(t *testing.T) { + s := &legacyInteractiveSession{err: errors.New("close error")} + var statusErr *StatusError + require.ErrorAs(t, CloseInteractiveInput(s), &statusErr) + require.Equal(t, ErrorUnimplemented, statusErr.Code) + require.False(t, s.closed, "unsupported input closure must not cancel the session") + require.ErrorIs(t, CancelInteractive(s), s.err) + require.True(t, s.closed) +} + +func TestInteractiveControlDispatch(t *testing.T) { + s := &controlledInteractiveSession{legacyInteractiveSession: legacyInteractiveSession{err: errors.New("control error")}} + require.Implements(t, (*InteractiveSessionControl)(nil), s) + require.ErrorIs(t, CloseInteractiveInput(s), s.err) + require.True(t, s.inputClosed) + require.ErrorIs(t, CancelInteractive(s), s.err) + require.True(t, s.cancelled) + require.False(t, s.closed, "explicit cancellation takes precedence over fallback") +} diff --git a/sdk/typescript/README.md b/sdk/typescript/README.md index 52122ccfb6..c33fe5000f 100644 --- a/sdk/typescript/README.md +++ b/sdk/typescript/README.md @@ -124,7 +124,20 @@ for await (const event of client.sandbox.execStream(name, ['pytest', '-q'])) { } ``` -`execInteractive` is the TTY + stdin transport primitive. Drive it by consuming `output`, which yields the same chunk/exit events; `done` resolves with the exit code once the stream reaches its exit event and rejects if it ends without one. It ships raw bytes only; raw mode, signal forwarding, and SIGWINCH stay with the caller. +`execInteractive` is the TTY + stdin transport primitive. Consume `output` concurrently with awaiting `done`. The helper yields its terminal exit event and resolves `done` only after receiving an exit event and successful final gRPC status. If transport completion fails, `done` rejects and `session.exitCode` retains any observed process exit code. Duplicate exit events and output after exit are errors. + +The RPC starts immediately when the session is created, even before you consume +`output`. A background receiver observes transport errors and buffers up to 1 MiB +of output in 64 KiB chunks. Once the queue fills, receiving waits for the caller +to drain it; cancellation interrupts that wait. The bound excludes the transport's +current response frame and its own buffers. + +Call `closeInput()` (or its compatibility alias `close()`) to end stdin and resize input while continuing to receive output. Later writes and resizes throw. Call `cancel()` to abort the RPC. Both close and cancel are idempotent. The helper ships raw bytes only; raw mode, signal forwarding, and SIGWINCH stay with the caller. Input closure is not a terminal Ctrl-D keystroke. + +`execInteractive()` returns `ExecInteractiveSessionControl`, which extends the +original `ExecInteractiveSession` with `closeInput()`, `cancel()`, and `exitCode`. +Existing mocks and wrappers can keep implementing the original interface without +adding those members. Use the extended type when a wrapper exposes the new controls. ```ts const session = await client.sandbox.execInteractive(name, ['bash']) @@ -282,6 +295,7 @@ mise run sdk:ts:lint # Biome: lint + format check (read-only) mise run sdk:ts:typecheck # tsc --noEmit mise run sdk:ts:test # Vitest unit tests with an 80% line-coverage gate mise run sdk:ts:build # emit dist/ +mise run e2e:sdk:ts:exec # public interactive helper against an isolated Docker gateway ``` Formatting and linting are handled by [Biome](https://biomejs.dev) (`biome.json`): 2-space indent, single quotes, semicolons, 120-column width. Generated `src/gen/` is excluded. `sdk:ts:lint` runs in CI as part of `sdk:ts:ci`. diff --git a/sdk/typescript/e2e/interactive-exec.mjs b/sdk/typescript/e2e/interactive-exec.mjs new file mode 100644 index 0000000000..eaeb94f5f1 --- /dev/null +++ b/sdk/typescript/e2e/interactive-exec.mjs @@ -0,0 +1,67 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Run through mise run e2e:sdk:ts:exec, which supplies an isolated gateway. +import assert from 'node:assert/strict'; +import { readFileSync } from 'node:fs'; +import { join } from 'node:path'; +import { test } from 'node:test'; +import { SandboxClient } from '../dist/index.js'; + +test('public interactive helper drains output after input EOF and verifies completion', { + timeout: 300_000, +}, async () => { + assert.ok(process.env.XDG_CONFIG_HOME, 'run with the Docker gateway wrapper'); + assert.ok(process.env.OPENSHELL_GATEWAY, 'run with the Docker gateway wrapper'); + const dir = join(process.env.XDG_CONFIG_HOME, 'openshell', 'gateways', process.env.OPENSHELL_GATEWAY); + const metadata = JSON.parse(readFileSync(join(dir, 'metadata.json'), 'utf8')); + const client = await SandboxClient.connect({ + gateway: metadata.gateway_endpoint, + ...(metadata.gateway_endpoint.startsWith('https:') + ? { + caCert: readFileSync(join(dir, 'mtls', 'ca.crt')), + clientCert: readFileSync(join(dir, 'mtls', 'tls.crt')), + clientKey: readFileSync(join(dir, 'mtls', 'tls.key')), + } + : {}), + }); + const name = `ts-eof-${Date.now().toString(36)}`; + await client.create({ + name, + image: process.env.OPENSHELL_E2E_DOCKER_SANDBOX_IMAGE ?? 'ghcr.io/nvidia/openshell-community/sandboxes/base:latest', + }); + try { + await client.waitReady(name, 180); + const session = await client.execInteractive( + name, + ['/bin/sh', '-c', 'input=$(cat); printf "stdout:%s" "$input"; printf "stderr:drained" >&2; exit 7'], + { tty: false, timeoutSecs: 30, signal: AbortSignal.timeout(45_000) }, + ); + try { + // cat cannot finish until closeInput reaches the sandbox. All command + // output therefore proves that request EOF left the response open. + session.write(Buffer.from('input-before-eof')); + session.closeInput(); + session.closeInput(); // The public control remains idempotent. + const stdout = []; + const stderr = []; + const exits = []; + for await (const event of session.output) { + assert.equal(exits.length, 0, 'exit must be the final application event'); + if ('type' in event) exits.push(event.exitCode); + else (event.stream === 'stdout' ? stdout : stderr).push(event.data); + } + assert.equal(Buffer.concat(stdout).toString(), 'stdout:input-before-eof'); + assert.equal(Buffer.concat(stderr).toString(), 'stderr:drained'); + assert.deepEqual(exits, [7]); + // A nonzero process exit is not a failed transport. done resolves only + // after the helper has verified the final gRPC status. + assert.equal(await session.done, 7); + assert.equal(session.exitCode, 7); + } finally { + session.cancel(); + } + } finally { + await client.delete(name); + } +}); diff --git a/sdk/typescript/src/client.test.ts b/sdk/typescript/src/client.test.ts index efdd6aabcc..ec9426a6d0 100644 --- a/sdk/typescript/src/client.test.ts +++ b/sdk/typescript/src/client.test.ts @@ -24,6 +24,7 @@ import { } from './client.js'; import { OpenShell, SandboxPhase, ServiceStatus } from './gen/openshell_pb.js'; import { PolicySource, SettingScope } from './gen/sandbox_pb.js'; +import type { ExecInteractiveSession, ExecInteractiveSessionControl } from './index.js'; function client(impl: Partial>): SandboxClient { const transport: Transport = createRouterTransport((router) => { @@ -910,6 +911,35 @@ describe('Pushable', () => { }); describe('execInteractive', () => { + it('preserves legacy session implementations and exposes SDK lifecycle controls', async () => { + // These are exactly the original required members, checked through the + // public package exports so downstream wrappers can keep their old types. + const legacy: ExecInteractiveSession = { + output: (async function* () { + yield { type: 'exit' as const, exitCode: 0 }; + })(), + write() {}, + resize() {}, + close() {}, + done: Promise.resolve(0), + }; + const wrap = (session: ExecInteractiveSession): ExecInteractiveSession => session; + expect(wrap(legacy)).toBe(legacy); + + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* () { + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + }, + }); + const controlled: ExecInteractiveSessionControl = await sandbox.execInteractive('sb', ['true']); + expect(wrap(controlled)).toBe(controlled); + controlled.closeInput(); + expect(await controlled.done).toBe(0); + expect(controlled.exitCode).toBe(0); + controlled.cancel(); + }); + it('sends start first with tty/cols/rows, streams output, and resolves done', async () => { const cases: string[] = []; let started: { tty?: boolean; cols?: number; rows?: number; sandboxId?: string } | undefined; @@ -963,6 +993,179 @@ describe('execInteractive', () => { }); describe('exec done settlement', () => { + it('starts the command and observes completion before output is consumed', async () => { + let started = false; + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* () { + started = true; + yield { payload: { case: 'stdout', value: { data: enc('started') } } }; + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + }, + }); + const session = await sandbox.execInteractive('sb', ['bash']); + await expect.poll(() => started).toBe(true); + expect(await session.done).toBe(0); + const events = []; + for await (const event of session.output) events.push(event); + expect(events).toEqual([ + { stream: 'stdout', data: Buffer.from('started') }, + { type: 'exit', exitCode: 0 }, + ]); + }); + + it('observes early transport failures without output consumption', async () => { + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: () => { + throw new ConnectError('early failure', Code.Unavailable); + }, + }); + const session = await sandbox.execInteractive('sb', ['bash']); + await expect(session.done).rejects.toMatchObject({ connectCode: Code.Unavailable }); + const iterator = session.output[Symbol.asyncIterator](); + await expect(iterator.next()).rejects.toMatchObject({ connectCode: Code.Unavailable }); + }); + + it('cancels a receiver blocked by output backpressure', async () => { + let released = false; + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* (_requests, ctx) { + ctx.signal.addEventListener('abort', () => { + released = true; + }); + // More than the SDK queue budget, in one transport event. + yield { payload: { case: 'stdout', value: { data: new Uint8Array(4 * 1024 * 1024) } } }; + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + }, + }); + const session = await sandbox.execInteractive('sb', ['bash']); + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(session.exitCode).toBeUndefined(); + session.cancel(); + await expect(session.done).rejects.toMatchObject({ code: 'canceled' }); + await expect.poll(() => released).toBe(true); + await expect(session.output[Symbol.asyncIterator]().next()).rejects.toMatchObject({ code: 'canceled' }); + }); + + it('drains output larger than the queue budget without losing bytes', async () => { + const data = Buffer.alloc(2 * 1024 * 1024 + 7, 'x'); + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* () { + yield { payload: { case: 'stdout', value: { data } } }; + yield { payload: { case: 'stderr', value: { data: enc('last') } } }; + yield { payload: { case: 'exit', value: { exitCode: 3 } } }; + }, + }); + const session = await sandbox.execInteractive('sb', ['bash']); + const chunks: Buffer[] = []; + for await (const event of session.output) { + if ('type' in event) expect(await session.done).toBe(3); + else chunks.push(event.data); + } + expect(Buffer.concat(chunks)).toEqual(Buffer.concat([data, Buffer.from('last')])); + }); + + it('retains the exit code but rejects completion on a later transport error', async () => { + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* () { + yield { payload: { case: 'exit', value: { exitCode: 7 } } }; + throw new ConnectError('connection lost after exit', Code.Unavailable); + }, + }); + const session = await sandbox.execInteractive('sb', ['bash']); + await expect( + (async () => { + for await (const _event of session.output) { + /* drain through trailers */ + } + })(), + ).rejects.toMatchObject({ connectCode: Code.Unavailable }); + await expect(session.done).rejects.toMatchObject({ connectCode: Code.Unavailable }); + expect(session.exitCode).toBe(7); + }); + + it.each(['stdout', 'exit'] as const)('rejects %s after an exit event', async (payload) => { + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* () { + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + if (payload === 'stdout') yield { payload: { case: 'stdout', value: { data: enc('late') } } }; + else yield { payload: { case: 'exit', value: { exitCode: 1 } } }; + }, + }); + const session = await sandbox.execInteractive('sb', ['bash']); + await expect( + (async () => { + for await (const _event of session.output) { + /* drain */ + } + })(), + ).rejects.toThrow('after exit'); + await expect(session.done).rejects.toThrow('after exit'); + expect(session.exitCode).toBe(0); + }); + + it('closes input idempotently and rejects later stdin and resize', async () => { + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* (requests) { + for await (const _input of requests) { + /* wait for request EOF */ + } + yield { payload: { case: 'stdout', value: { data: enc('drained') } } }; + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + }, + }); + const session = await sandbox.execInteractive('sb', ['bash']); + session.closeInput(); + session.close(); + expect(() => session.write(Buffer.from('late'))).toThrow('input is closed'); + expect(() => session.resize(80, 24)).toThrow('input is closed'); + const output = []; + for await (const event of session.output) output.push(event); + expect(output).toHaveLength(2); + expect(await session.done).toBe(0); + }); + + it('cancel settles completion even before output is consumed', async () => { + const sandbox = client({ getSandbox: () => readySandbox('sb', 'sb-id') }); + const session = await sandbox.execInteractive('sb', ['bash']); + session.cancel(); + session.cancel(); + await expect(session.done).rejects.toMatchObject({ code: 'canceled' }); + expect(() => session.write(Buffer.from('late'))).toThrow('input is closed'); + }); + + it('external cancellation settles completion before output is consumed', async () => { + const controller = new AbortController(); + const sandbox = client({ getSandbox: () => readySandbox('sb', 'sb-id') }); + const session = await sandbox.execInteractive('sb', ['bash'], { signal: controller.signal }); + controller.abort(); + await expect(session.done).rejects.toMatchObject({ code: 'canceled' }); + expect(() => session.resize(80, 24)).toThrow('input is closed'); + }); + + it('external cancellation settles completion while output is paused', async () => { + const controller = new AbortController(); + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id'), + execSandboxInteractive: async function* () { + yield { payload: { case: 'stdout', value: { data: enc('partial') } } }; + yield { payload: { case: 'exit', value: { exitCode: 0 } } }; + }, + }); + const session = await sandbox.execInteractive('sb', ['bash'], { signal: controller.signal }); + const iterator = session.output[Symbol.asyncIterator](); + await iterator.next(); + controller.abort(); + await expect(session.done).rejects.toMatchObject({ code: 'canceled' }); + await iterator.return?.(); + }); + it('resolves done even when the consumer breaks right after the exit event', async () => { const sandbox = client({ getSandbox: () => readySandbox('sb', 'sb-id'), @@ -976,7 +1179,7 @@ describe('exec done settlement', () => { for await (const event of session.output) { if ('type' in event) break; // break on exit: the generator never resumes } - // Without settling `done` before the exit yield, this would hang forever. + // The public exit is yielded only after successful terminal status. expect(await session.done).toBe(3); }); diff --git a/sdk/typescript/src/client.ts b/sdk/typescript/src/client.ts index b90ae6026b..eec8c92568 100644 --- a/sdk/typescript/src/client.ts +++ b/sdk/typescript/src/client.ts @@ -286,16 +286,28 @@ export interface ExecInteractiveOptions extends SandboxWorkspaceOptions { // The transport half of an interactive exec: raw stdin/stdout/stderr plus // resize, with no terminal glue. Drive it by consuming `output`, which yields -// chunks then a terminal exit event; `done` resolves with the exit code once -// the stream reaches that exit event and rejects if the stream ends without one. +// chunks then a terminal exit event; `done` resolves only after an exit event +// and successful RPC completion. Consume output concurrently with awaiting done. export interface ExecInteractiveSession { output: AsyncIterable; write(data: Buffer): void; resize(cols: number, rows: number): void; + /** Close stdin and resize input while preserving output. */ close(): void; done: Promise; } +/** Lifecycle controls available on SDK-created sessions. The base interface + * retains its original members for existing custom sessions and wrappers. */ +export interface ExecInteractiveSessionControl extends ExecInteractiveSession { + /** Close stdin and resize input, preserving output until completion. */ + closeInput(): void; + /** Cancel the RPC and stop receiving output. */ + cancel(): void; + /** Observed process exit, retained even if final RPC completion fails. */ + readonly exitCode: number | undefined; +} + /** Cancellation for the poll-based wait helpers. */ export interface WaitOptions extends SandboxWorkspaceOptions { /** Abort the wait (and the in-flight poll RPC) early. */ @@ -687,6 +699,60 @@ export class Pushable implements AsyncIterable { } } +// Single producer/consumer queue. Closing wakes both directions, including a +// producer blocked by backpressure. Successful completion drains queued values. +class ExecOutputQueue { + private readonly values: ExecStreamEvent[] = []; + private ended = false; + private error: unknown; + private reader?: () => void; + private writer?: () => void; + + async push(value: ExecStreamEvent): Promise { + // The terminal exit carries no output bytes and must never block cleanup + // after done has already settled and its cancellation listener is removed. + while (!('type' in value) && this.values.length >= 16 && !this.ended) { + await new Promise((resolve) => { + this.writer = resolve; + }); + } + if (this.ended) throw this.error ?? new SdkError('canceled', 'exec output closed'); + this.values.push(value); + this.reader?.(); + this.reader = undefined; + } + + end(error?: unknown, discard = false): void { + if (discard) this.values.length = 0; + if (!this.ended) { + this.ended = true; + this.error = error; + } + this.reader?.(); + this.writer?.(); + this.reader = undefined; + this.writer = undefined; + } + + async *[Symbol.asyncIterator](): AsyncGenerator { + for (;;) { + const value = this.values.shift(); + if (value !== undefined) { + this.writer?.(); + this.writer = undefined; + yield value; + } else if (this.ended) { + if (this.error !== undefined) throw this.error; + return; + } else { + await new Promise((resolve) => { + this.reader = resolve; + }); + } + } + } +} + /** One response page from a list operation. */ export interface Page { readonly items: T[]; @@ -1084,7 +1150,7 @@ export class SandboxClient { name: string, command: string[], options?: ExecInteractiveOptions | null, - ): Promise { + ): Promise { let sandboxId: string; try { sandboxId = ( @@ -1113,7 +1179,19 @@ export class SandboxClient { }, }); - const stream = this.grpc.execSandboxInteractive(input, { signal: options?.signal }); + const controller = new AbortController(); + const signal = options?.signal ? AbortSignal.any([options.signal, controller.signal]) : controller.signal; + const grpc = this.grpc; + const queue = new ExecOutputQueue(); + let inputClosed = false; + let exitCode: number | undefined; + const closeInput = (): void => { + inputClosed = true; + input.end(); + }; + const assertInputOpen = (): void => { + if (inputClosed || signal.aborted) throw new SdkError('io', 'exec input is closed'); + }; let resolveDone!: (code: number) => void; let rejectDone!: (err: unknown) => void; const done = new Promise((resolve, reject) => { @@ -1124,72 +1202,113 @@ export class SandboxClient { // keeps an unobserved rejection from surfacing as an unhandledRejection; // real awaiters still receive it through their own handler. void done.catch(() => {}); - // Settle exactly once. The exit code wins; error/abandonment only apply - // when no exit was observed. + // The process exit and the terminal transport status are separate outcomes. let settled = false; const settleExit = (code: number): void => { if (settled) return; settled = true; + signal.removeEventListener('abort', onAbort); resolveDone(code); }; const settleError = (err: unknown): void => { if (settled) return; settled = true; + signal.removeEventListener('abort', onAbort); rejectDone(err); }; - async function* output(): AsyncGenerator { - let sawExit = false; + const onAbort = (): void => { + closeInput(); + const error = new SdkError('canceled', 'exec cancelled'); + queue.end(error, true); + settleError(error); + }; + signal.addEventListener('abort', onAbort, { once: true }); + if (signal.aborted) onAbort(); + + async function receive(): Promise { try { + if (signal.aborted) throw new SdkError('canceled', 'exec cancelled'); + // Start and observe the transport immediately, independently of output + // consumption. Backpressure bounds the queue to 16 chunks of 64 KiB. + const stream = grpc.execSandboxInteractive(input, { signal }); for await (const event of stream) { + if (signal.aborted) throw new SdkError('canceled', 'exec cancelled'); + if (exitCode !== undefined) { + throw new SdkError('rpc', 'ExecSandboxInteractive received an event after exit'); + } switch (event.payload.case) { case 'stdout': - yield { - stream: 'stdout', - data: Buffer.from(event.payload.value.data), - }; - break; case 'stderr': - yield { - stream: 'stderr', - data: Buffer.from(event.payload.value.data), - }; + for (let offset = 0; offset < event.payload.value.data.length; offset += 64 * 1024) { + await queue.push({ + stream: event.payload.case, + data: Buffer.from(event.payload.value.data.subarray(offset, offset + 64 * 1024)), + }); + } break; case 'exit': - sawExit = true; - // Settle `done` before yielding: a consumer that breaks on the - // exit event abandons the generator at the yield, so anything - // after it would never run. - settleExit(event.payload.value.exitCode); - yield { type: 'exit', exitCode: event.payload.value.exitCode }; + exitCode = event.payload.value.exitCode; + closeInput(); break; + default: + throw new SdkError('rpc', 'ExecSandboxInteractive received an empty or unknown event'); } } - if (!sawExit) { + if (exitCode === undefined) { throw new SdkError('rpc', 'ExecSandboxInteractive stream ended without an exit event'); } + if (signal.aborted) throw new SdkError('canceled', 'exec cancelled before completion'); + // Delay the public exit event until trailers have been consumed. A + // caller can still break on exit without losing the terminal status. + settleExit(exitCode); + await queue.push({ type: 'exit', exitCode }); + queue.end(); } catch (e) { const err = e instanceof SdkError ? e : fromConnect(e); settleError(err); - throw err; + queue.end(err); + } finally { + closeInput(); + controller.abort(); + } + } + + // receive catches transport failures even when nobody consumes output/done. + const receiving = receive(); + async function* output(): AsyncGenerator { + try { + yield* queue; } finally { - input.end(); - // Consumer abandoned the stream before an exit event (early break or - // return): settle `done` so it can never hang. - settleError(new SdkError('rpc', 'exec output abandoned before exit')); + const error = new SdkError('rpc', 'exec output abandoned before completion'); + settleError(error); + queue.end(error, true); + closeInput(); + controller.abort(); + await receiving; } } return { output: output(), write(data: Buffer): void { + assertInputOpen(); input.push({ payload: { case: 'stdin', value: new Uint8Array(data) } }); }, resize(cols: number, rows: number): void { + assertInputOpen(); input.push({ payload: { case: 'resize', value: { cols, rows } } }); }, - close(): void { - input.end(); + closeInput, + close: closeInput, + cancel(): void { + closeInput(); + queue.end(new SdkError('canceled', 'exec cancelled'), true); + controller.abort(); + settleError(new SdkError('canceled', 'exec cancelled')); + }, + get exitCode(): number | undefined { + return exitCode; }, done, }; diff --git a/sdk/typescript/src/index.ts b/sdk/typescript/src/index.ts index d858cd75b6..4668f2a357 100644 --- a/sdk/typescript/src/index.ts +++ b/sdk/typescript/src/index.ts @@ -12,6 +12,7 @@ export type { ExecExitEvent, ExecInteractiveOptions, ExecInteractiveSession, + ExecInteractiveSessionControl, ExecOptions, ExecResult, ExecStreamChunk, diff --git a/tasks/typescript.toml b/tasks/typescript.toml index 57277e05ba..ed07379c63 100644 --- a/tasks/typescript.toml +++ b/tasks/typescript.toml @@ -59,6 +59,11 @@ depends = ["sdk:ts:proto"] dir = "sdk/typescript" run = "npm test" +["e2e:sdk:ts:exec"] +description = "Test public TypeScript interactive exec against a Docker-backed gateway" +depends = ["sdk:ts:build"] +run = "e2e/with-docker-gateway.sh node --test sdk/typescript/e2e/interactive-exec.mjs" + ["sdk:ts:ci"] description = "TypeScript SDK checks (proto lint, Biome lint, codegen, typecheck, test, build)" depends = [