fix(sandbox): stabilize terminal attachment checks

Signed-off-by: Drew Newberry <anewberry@nvidia.com>

Author: drew

crates/openshell-cli/src/run.rs

@@ -1054,7 +1054,7 @@ pub async fn sandbox_create(
             if detach {
                 return Ok(0);
             }
-            let connect_result = crate::ssh::sandbox_connect_without_exec(
+            let connect_result = crate::ssh::sandbox_connect_terminal_main(
                 &effective_server,
                 &sandbox_name,
                 &effective_tls,

crates/openshell-cli/src/ssh.rs

@@ -29,6 +29,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream;
 use tokio::process::{Child, Command as TokioCommand};
 use tokio_stream::wrappers::ReceiverStream;
+use tonic::Code;
 
 /// Time budget for the local listener to become reachable after `ssh` starts.
 /// This is a user-visible readiness deadline for both foreground and background
@@ -39,6 +40,10 @@ const FORWARD_LISTENER_PROBE_INTERVAL: Duration = Duration::from_millis(50);
 /// Per-attempt connect timeout, so one hung probe cannot consume the whole
 /// grace period.
 const FORWARD_LISTENER_CONNECT_TIMEOUT: Duration = Duration::from_millis(200);
+/// Time budget for the supervisor relay to register after a fast canonical
+/// command has already reported its terminal result.
+const TERMINAL_RELAY_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(5);
+const TERMINAL_RELAY_REGISTRATION_INTERVAL: Duration = Duration::from_millis(50);
 
 #[derive(Clone, Copy, Debug)]
 pub enum Editor {
@@ -79,6 +84,7 @@ async fn ssh_session_config(
     name: &str,
     tls: &TlsOptions,
     workspace: &str,
+    terminal_relay_registration_timeout: Option<Duration>,
 ) -> Result<SshSessionConfig> {
     let mut client = grpc_client(server, tls).await?;
 
@@ -94,12 +100,26 @@ async fn ssh_session_config(
         .sandbox
         .ok_or_else(|| miette::miette!("sandbox not found"))?;
 
-    let response = client
-        .create_ssh_session(CreateSshSessionRequest {
-            sandbox_id: sandbox.object_id().to_string(),
-        })
-        .await
-        .into_diagnostic()?;
+    let relay_registration_deadline =
+        terminal_relay_registration_timeout.map(|timeout| tokio::time::Instant::now() + timeout);
+    let response = loop {
+        match client
+            .create_ssh_session(CreateSshSessionRequest {
+                sandbox_id: sandbox.object_id().to_string(),
+            })
+            .await
+        {
+            Ok(response) => break response,
+            Err(status)
+                if status.code() == Code::FailedPrecondition
+                    && relay_registration_deadline
+                        .is_some_and(|deadline| tokio::time::Instant::now() < deadline) =>
+            {
+                tokio::time::sleep(TERMINAL_RELAY_REGISTRATION_INTERVAL).await;
+            }
+            Err(status) => return Err(status).into_diagnostic(),
+        }
+    };
     let session = response.into_inner();
     validate_ssh_session_response(&session)
         .map_err(|err| miette::miette!("gateway returned invalid SSH session response: {err}"))?;
@@ -257,8 +277,16 @@ async fn sandbox_connect_with_mode(
     tls: &TlsOptions,
     replace_process: bool,
     workspace: &str,
+    terminal_relay_registration_timeout: Option<Duration>,
 ) -> Result<i32> {
-    let session = ssh_session_config(server, name, tls, workspace).await?;
+    let session = ssh_session_config(
+        server,
+        name,
+        tls,
+        workspace,
+        terminal_relay_registration_timeout,
+    )
+    .await?;
 
     let mut command = ssh_base_command(&session.proxy_command);
     if session.main_terminal {
@@ -290,7 +318,7 @@ pub async fn sandbox_connect(
     tls: &TlsOptions,
     workspace: &str,
 ) -> Result<i32> {
-    sandbox_connect_with_mode(server, name, tls, true, workspace).await
+    sandbox_connect_with_mode(server, name, tls, true, workspace, None).await
 }
 
 pub(crate) async fn sandbox_connect_without_exec(
@@ -299,7 +327,24 @@ pub(crate) async fn sandbox_connect_without_exec(
     tls: &TlsOptions,
     workspace: &str,
 ) -> Result<i32> {
-    sandbox_connect_with_mode(server, name, tls, false, workspace).await
+    sandbox_connect_with_mode(server, name, tls, false, workspace, None).await
+}
+
+pub(crate) async fn sandbox_connect_terminal_main(
+    server: &str,
+    name: &str,
+    tls: &TlsOptions,
+    workspace: &str,
+) -> Result<i32> {
+    sandbox_connect_with_mode(
+        server,
+        name,
+        tls,
+        false,
+        workspace,
+        Some(TERMINAL_RELAY_REGISTRATION_TIMEOUT),
+    )
+    .await
 }
 
 pub async fn sandbox_connect_editor(
@@ -310,7 +355,7 @@ pub async fn sandbox_connect_editor(
     tls: &TlsOptions,
     workspace: &str,
 ) -> Result<()> {
-    let session = ssh_session_config(server, name, tls, workspace).await?;
+    let session = ssh_session_config(server, name, tls, workspace, None).await?;
     let workspace_root = discover_workspace_root(&session).await?;
 
     let host_alias = host_alias(name, workspace);
@@ -339,7 +384,7 @@ pub async fn sandbox_forward(
 ) -> Result<()> {
     openshell_core::forward::check_port_available(spec)?;
 
-    let session = ssh_session_config(server, name, tls, workspace).await?;
+    let session = ssh_session_config(server, name, tls, workspace, None).await?;
 
     let mut command = TokioCommand::from(ssh_base_command(&session.proxy_command));
     command
@@ -554,7 +599,7 @@ async fn sandbox_exec_with_mode(
         return Err(miette::miette!("no command provided"));
     }
 
-    let session = ssh_session_config(server, name, tls, workspace).await?;
+    let session = ssh_session_config(server, name, tls, workspace, None).await?;
     let mut ssh = ssh_base_command(&session.proxy_command);
 
     if tty {
@@ -767,7 +812,7 @@ async fn ssh_tar_upload(
     tls: &TlsOptions,
     workspace: &str,
 ) -> Result<()> {
-    let session = ssh_session_config(server, name, tls, workspace).await?;
+    let session = ssh_session_config(server, name, tls, workspace, None).await?;
 
     let dest_dir = dest_dir.unwrap_or(".");
     let escaped_dest = shell_escape(dest_dir);
@@ -1200,7 +1245,7 @@ pub async fn sandbox_sync_down(
     tls: &TlsOptions,
     workspace: &str,
 ) -> Result<()> {
-    let session = ssh_session_config(server, name, tls, workspace).await?;
+    let session = ssh_session_config(server, name, tls, workspace, None).await?;
     let sandbox_path = resolve_sandbox_source_path(&session, sandbox_path).await?;
     let kind = probe_sandbox_source_kind(&session, &sandbox_path).await?;
 
@@ -1482,7 +1527,7 @@ pub async fn sandbox_ssh_proxy_by_name(
     tls: &TlsOptions,
     workspace: &str,
 ) -> Result<()> {
-    let session = ssh_session_config(server, name, tls, workspace).await?;
+    let session = ssh_session_config(server, name, tls, workspace, None).await?;
     sandbox_ssh_proxy(
         &session.gateway_url,
         &session.sandbox_id,

crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs

@@ -51,6 +51,9 @@ struct SandboxState {
     vm_error_with_observed_exit: Arc<AtomicBool>,
     vm_slow_progress_before_ready: Arc<AtomicBool>,
     vm_log_churn_before_ready: Arc<AtomicBool>,
+    terminal_before_relay: Arc<AtomicBool>,
+    ssh_session_failures_remaining: Arc<AtomicUsize>,
+    ssh_session_requests: Arc<AtomicUsize>,
     global_settings: Arc<Mutex<HashMap<String, SettingValue>>>,
     gateway_config_requests: Arc<AtomicUsize>,
 }
@@ -244,6 +247,19 @@ impl OpenShell for TestOpenShell {
         &self,
         request: tonic::Request<CreateSshSessionRequest>,
     ) -> Result<Response<CreateSshSessionResponse>, Status> {
+        self.state
+            .ssh_session_requests
+            .fetch_add(1, Ordering::SeqCst);
+        if self
+            .state
+            .ssh_session_failures_remaining
+            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
+                remaining.checked_sub(1)
+            })
+            .is_ok()
+        {
+            return Err(Status::failed_precondition("sandbox is not ready"));
+        }
         let sandbox_id = request.into_inner().sandbox_id;
         Ok(Response::new(CreateSshSessionResponse {
             sandbox_id,
@@ -426,6 +442,7 @@ impl OpenShell for TestOpenShell {
             .vm_slow_progress_before_ready
             .load(Ordering::SeqCst);
         let vm_log_churn_before_ready = self.state.vm_log_churn_before_ready.load(Ordering::SeqCst);
+        let terminal_before_relay = self.state.terminal_before_relay.load(Ordering::SeqCst);
 
         tokio::spawn(async move {
             let mut provisioning = Sandbox {
@@ -462,6 +479,12 @@ impl OpenShell for TestOpenShell {
             }
             let mut ready = provisioning.clone();
             ready.set_phase(SandboxPhase::Ready as i32);
+            let mut completed = provisioning.clone();
+            completed.status = Some(SandboxStatus {
+                exit_code: Some(0),
+                ..SandboxStatus::default()
+            });
+            completed.set_phase(SandboxPhase::Completed as i32);
 
             let _ = tx
                 .send(Ok(SandboxStreamEvent {
@@ -511,6 +534,14 @@ impl OpenShell for TestOpenShell {
                     .await;
                 return;
             }
+            if terminal_before_relay {
+                let _ = tx
+                    .send(Ok(SandboxStreamEvent {
+                        payload: Some(sandbox_stream_event::Payload::Sandbox(completed)),
+                    }))
+                    .await;
+                return;
+            }
             if vm_slow_progress_before_ready {
                 tokio::time::sleep(Duration::from_millis(600)).await;
                 let _ = tx
@@ -1713,6 +1744,50 @@ async fn sandbox_create_times_out_when_only_logs_arrive() {
     assert!(err.to_string().contains("sandbox provisioning timed out"));
 }
 
+#[tokio::test]
+async fn sandbox_create_retries_terminal_attachment_until_relay_registers() {
+    let server = run_server().await;
+    server
+        .openshell
+        .state
+        .terminal_before_relay
+        .store(true, Ordering::SeqCst);
+    server
+        .openshell
+        .state
+        .ssh_session_failures_remaining
+        .store(1, Ordering::SeqCst);
+    let fake_ssh_dir = tempfile::tempdir().unwrap();
+    let xdg_dir = tempfile::tempdir().unwrap();
+    let _env = test_env(&fake_ssh_dir, &xdg_dir);
+    let tls = test_tls(&server);
+    install_fake_ssh(&fake_ssh_dir);
+
+    let exit_code = run::sandbox_create(
+        &server.endpoint,
+        "openshell",
+        run::SandboxCreateConfig {
+            name: Some("fast-command"),
+            command: &["echo".into(), "OK".into()],
+            ..test_config()
+        },
+        "default",
+        &tls,
+    )
+    .await
+    .expect("sandbox create should wait for the declared terminal attachment relay");
+
+    assert_eq!(exit_code, 0);
+    assert_eq!(
+        server
+            .openshell
+            .state
+            .ssh_session_requests
+            .load(Ordering::SeqCst),
+        2
+    );
+}
+
 #[tokio::test]
 async fn sandbox_create_deletes_command_sessions_with_no_keep() {
     let server = run_server().await;

sdk/go/openshell/v1/internal/converter/coverage_test.go

@@ -63,9 +63,12 @@ func TestConverterCoversAllProtoFields_SandboxStatus(t *testing.T) {
 		"current_policy_version": true,
 		"exit_code":              true,
 	}
-	// The instance ID is an internal gateway/supervisor fencing token exposed
-	// only through the raw protobuf API.
-	skipped := fieldSet{"main_process_instance_id": true}
+	// These fields coordinate internal gateway/supervisor lifecycle fencing and
+	// terminal delivery. They are exposed only through the raw protobuf API.
+	skipped := fieldSet{
+		"main_process_instance_id":    true,
+		"main_process_exit_finalized": true,
+	}
 
 	assertAllFieldsCovered(t, (&pb.SandboxStatus{}).ProtoReflect().Descriptor(), handled, skipped)
 }