refactor(server): internalize terminal finalization state

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

Author: drew

crates/openshell-server/src/compute/mod.rs

@@ -2878,12 +2878,16 @@ impl ComputeRuntime {
         sandbox_id: &str,
         instance_id: &str,
     ) -> Result<(), String> {
-        self.set_supervisor_session_state(sandbox_id, true, Some(instance_id))
+        self.set_supervisor_session_state(sandbox_id, true, Some(instance_id), false)
             .await
     }
 
-    pub async fn supervisor_session_disconnected(&self, sandbox_id: &str) -> Result<(), String> {
-        self.set_supervisor_session_state(sandbox_id, false, None)
+    pub async fn supervisor_session_disconnected(
+        &self,
+        sandbox_id: &str,
+        terminal_delivery_finalized: bool,
+    ) -> Result<(), String> {
+        self.set_supervisor_session_state(sandbox_id, false, None, terminal_delivery_finalized)
             .await
     }
 
@@ -2892,6 +2896,7 @@ impl ComputeRuntime {
         sandbox_id: &str,
         connected: bool,
         instance_id: Option<&str>,
+        terminal_delivery_finalized: bool,
     ) -> Result<(), String> {
         let guard = self.sync_lock.lock().await;
 
@@ -2907,10 +2912,7 @@ impl ComputeRuntime {
             SandboxPhase::try_from(existing.phase()).unwrap_or(SandboxPhase::Unknown);
         if !connected
             && matches!(current_phase, SandboxPhase::Error | SandboxPhase::Completed)
-            && existing
-                .status
-                .as_ref()
-                .is_some_and(|status| status.main_process_exit_finalized)
+            && terminal_delivery_finalized
         {
             drop(guard);
             self.schedule_ephemeral_sandbox_delete(&existing);
@@ -2941,7 +2943,6 @@ impl ComputeRuntime {
                     let status = sandbox.status.get_or_insert_with(Default::default);
                     status.main_process_instance_id = instance_id.unwrap_or_default().to_string();
                     status.exit_code = None;
-                    status.main_process_exit_finalized = false;
                     sandbox.set_phase(SandboxPhase::Ready as i32);
                 } else {
                     ensure_supervisor_not_ready_status(&mut sandbox.status, &sandbox_name);
@@ -3087,22 +3088,6 @@ impl ComputeRuntime {
         {
             return Err("main-process instance does not match the terminal result".to_string());
         }
-        if status.main_process_exit_finalized {
-            return Ok(());
-        }
-        let expected_resource_version = sandbox_resource_version(&sandbox);
-        let finalized = self
-            .store
-            .update_message_cas::<Sandbox, _>(sandbox_id, expected_resource_version, |sandbox| {
-                sandbox
-                    .status
-                    .get_or_insert_with(Default::default)
-                    .main_process_exit_finalized = true;
-            })
-            .await
-            .map_err(|error| error.to_string())?;
-        self.sandbox_index.update_from_sandbox(&finalized);
-        self.sandbox_watch_bus.notify(sandbox_id);
         Ok(())
     }
 
@@ -3499,7 +3484,6 @@ fn apply_main_process_exit(sandbox: &mut Sandbox, instance_id: &str, exit_code:
     });
     status.main_process_instance_id = instance_id.to_string();
     status.exit_code = Some(exit_code);
-    status.main_process_exit_finalized = false;
     if preserve_infrastructure_error {
         return;
     }
@@ -3910,7 +3894,6 @@ fn public_status_from_driver(
         current_policy_version,
         main_process_instance_id: String::new(),
         exit_code: None,
-        main_process_exit_finalized: false,
     }
 }
 
@@ -5334,19 +5317,8 @@ mod tests {
             .await
             .unwrap();
         assert_eq!(driver.delete_calls(), 0);
-        assert!(
-            runtime
-                .store
-                .get_message::<Sandbox>("sb-1")
-                .await
-                .unwrap()
-                .unwrap()
-                .status
-                .unwrap()
-                .main_process_exit_finalized
-        );
         runtime
-            .supervisor_session_disconnected("sb-1")
+            .supervisor_session_disconnected("sb-1", true)
             .await
             .unwrap();
         tokio::time::timeout(Duration::from_secs(1), async {
@@ -7054,7 +7026,7 @@ mod tests {
         let mut watch_rx = runtime.sandbox_watch_bus.subscribe("sb-1");
 
         runtime
-            .supervisor_session_disconnected("sb-1")
+            .supervisor_session_disconnected("sb-1", false)
             .await
             .unwrap();
 
@@ -8329,7 +8301,7 @@ mod tests {
         runtime.store.put_message(&sandbox).await.unwrap();
 
         runtime
-            .supervisor_session_disconnected("sb-1")
+            .supervisor_session_disconnected("sb-1", false)
             .await
             .unwrap();
 
@@ -8573,7 +8545,7 @@ mod tests {
         // Session drops.
         runtime.supervisor_sessions.cleanup_sandbox("sb-1");
         runtime
-            .supervisor_session_disconnected("sb-1")
+            .supervisor_session_disconnected("sb-1", false)
             .await
             .unwrap();
         let stored = runtime

crates/openshell-server/src/grpc/sandbox.rs

@@ -1715,11 +1715,10 @@ fn sandbox_relay_reachable(state: &ServerState, sandbox: &Sandbox) -> bool {
     let phase = SandboxPhase::try_from(sandbox.phase()).ok();
     matches!(phase, Some(SandboxPhase::Ready))
         || (matches!(phase, Some(SandboxPhase::Completed | SandboxPhase::Error))
-            && sandbox
-                .status
-                .as_ref()
-                .is_some_and(|status| !status.main_process_exit_finalized)
-            && state.supervisor_sessions.has_session(sandbox.object_id()))
+            && state.supervisor_sessions.has_session(sandbox.object_id())
+            && !state
+                .supervisor_sessions
+                .terminal_delivery_finalized(sandbox.object_id()))
 }
 
 pub(super) async fn handle_create_ssh_session(
@@ -4009,13 +4008,12 @@ mod tests {
 
         assert!(response.is_ok());
 
-        let mut finalized = sandbox;
-        finalized
-            .status
-            .as_mut()
-            .expect("sandbox status")
-            .main_process_exit_finalized = true;
-        assert!(!sandbox_relay_reachable(&state, &finalized));
+        assert!(
+            state
+                .supervisor_sessions
+                .finalize_main_process_exit("sandbox-work")
+        );
+        assert!(!sandbox_relay_reachable(&state, &sandbox));
     }
 
     #[tokio::test]

crates/openshell-server/src/supervisor_session.rs

@@ -58,6 +58,9 @@ struct LiveSession {
     /// the old session's `tx` just before supersede could still enqueue a
     /// `RelayOpen` onto the stale stream and sit until the relay timeout.
     shutdown: oneshot::Sender<()>,
+    /// Set after the supervisor confirms that every expected foreground
+    /// attachment has closed and terminal output delivery is complete.
+    terminal_delivery_finalized: bool,
     #[allow(dead_code)]
     connected_at: Instant,
 }
@@ -125,6 +128,7 @@ impl SupervisorSessionRegistry {
                 session_id,
                 tx,
                 shutdown,
+                terminal_delivery_finalized: false,
                 connected_at: Instant::now(),
             },
         );
@@ -163,15 +167,17 @@ impl SupervisorSessionRegistry {
     /// This guards against the supersede race: an old session's task may
     /// finish long after a new session has taken its place. The old task's
     /// cleanup must not evict the new registration.
-    fn remove_if_current(&self, sandbox_id: &str, session_id: &str) -> bool {
+    fn remove_if_current(&self, sandbox_id: &str, session_id: &str) -> Option<bool> {
         let mut sessions = self.sessions.lock().unwrap();
         let is_current = sessions
             .get(sandbox_id)
             .is_some_and(|s| s.session_id == session_id);
         if is_current {
-            sessions.remove(sandbox_id);
+            return sessions
+                .remove(sandbox_id)
+                .map(|session| session.terminal_delivery_finalized);
         }
-        is_current
+        None
     }
 
     /// Look up the sender for a supervisor session, waiting up to `timeout`
@@ -210,6 +216,23 @@ impl SupervisorSessionRegistry {
         self.sessions.lock().unwrap().contains_key(sandbox_id)
     }
 
+    pub fn terminal_delivery_finalized(&self, sandbox_id: &str) -> bool {
+        self.sessions
+            .lock()
+            .unwrap()
+            .get(sandbox_id)
+            .is_some_and(|session| session.terminal_delivery_finalized)
+    }
+
+    pub fn finalize_main_process_exit(&self, sandbox_id: &str) -> bool {
+        let mut sessions = self.sessions.lock().unwrap();
+        let Some(session) = sessions.get_mut(sandbox_id) else {
+            return false;
+        };
+        session.terminal_delivery_finalized = true;
+        true
+    }
+
     pub fn is_current_session(&self, sandbox_id: &str, session_id: &str) -> bool {
         self.sessions
             .lock()
@@ -797,17 +820,17 @@ pub async fn handle_connect_supervisor(
             shutdown_rx,
         )
         .await;
-        let still_ours = state_clone
+        let terminal_finalized = state_clone
             .supervisor_sessions
             .remove_if_current(&sandbox_id_clone, &session_id);
-        if still_ours {
+        if let Some(terminal_finalized) = terminal_finalized {
             info!(sandbox_id = %sandbox_id_clone, session_id = %session_id, "supervisor session: ended");
             state_clone
                 .telemetry
                 .sandbox_session_disconnected(&sandbox_id_clone);
             if let Err(err) = state_clone
                 .compute
-                .supervisor_session_disconnected(&sandbox_id_clone)
+                .supervisor_session_disconnected(&sandbox_id_clone, terminal_finalized)
                 .await
             {
                 warn!(
@@ -874,6 +897,14 @@ pub async fn handle_finalize_main_process_exit(
         .finalize_main_process_exit(&report.sandbox_id, &report.instance_id)
         .await
         .map_err(Status::failed_precondition)?;
+    if !state
+        .supervisor_sessions
+        .finalize_main_process_exit(&report.sandbox_id)
+    {
+        return Err(Status::failed_precondition(
+            "supervisor session is not connected",
+        ));
+    }
     Ok(Response::new(
         openshell_core::proto::FinalizeMainProcessExitResponse {},
     ))
@@ -1133,7 +1164,7 @@ mod tests {
         let (tx, _rx) = mpsc::channel(1);
         registry.register("sbx".to_string(), "s1".to_string(), tx, make_shutdown());
 
-        assert!(registry.remove_if_current("sbx", "s1"));
+        assert_eq!(registry.remove_if_current("sbx", "s1"), Some(false));
         assert!(!registry.sessions.lock().unwrap().contains_key("sbx"));
     }
 
@@ -1159,7 +1190,7 @@ mod tests {
 
         // Cleanup from the old session task runs late. It must NOT evict the
         // newly registered session.
-        assert!(!registry.remove_if_current("sbx", "s-old"));
+        assert_eq!(registry.remove_if_current("sbx", "s-old"), None);
         let sessions = registry.sessions.lock().unwrap();
         assert!(
             sessions.contains_key("sbx"),
@@ -1171,7 +1202,18 @@ mod tests {
     #[test]
     fn remove_if_current_unknown_sandbox_is_noop() {
         let registry = SupervisorSessionRegistry::new();
-        assert!(!registry.remove_if_current("sbx-does-not-exist", "s1"));
+        assert_eq!(registry.remove_if_current("sbx-does-not-exist", "s1"), None);
+    }
+
+    #[test]
+    fn remove_if_current_returns_terminal_finalization_state() {
+        let registry = SupervisorSessionRegistry::new();
+        let (tx, _rx) = mpsc::channel(1);
+        registry.register("sbx".to_string(), "s1".to_string(), tx, make_shutdown());
+
+        assert!(registry.finalize_main_process_exit("sbx"));
+        assert!(registry.terminal_delivery_finalized("sbx"));
+        assert_eq!(registry.remove_if_current("sbx", "s1"), Some(true));
     }
 
     // ---- open_relay: happy path and wait semantics ----

proto/openshell.proto

@@ -918,9 +918,6 @@ message SandboxStatus {
   // Presence indicates that the canonical main process exited. Exit code 0
   // produces Completed; nonzero and signal-normalized exits produce Error.
   optional int32 exit_code = 9;
-  // True after every expected foreground main-process attachment has closed
-  // naturally and the supervisor is ready to exit.
-  bool main_process_exit_finalized = 10;
 }
 
 // User-facing sandbox condition derived from driver-native conditions.

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

@@ -63,12 +63,9 @@ func TestConverterCoversAllProtoFields_SandboxStatus(t *testing.T) {
 		"current_policy_version": true,
 		"exit_code":              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,
-	}
+	// The instance ID coordinates internal gateway/supervisor lifecycle
+	// fencing. It is exposed only through the raw protobuf API.
+	skipped := fieldSet{"main_process_instance_id": true}
 
 	assertAllFieldsCovered(t, (&pb.SandboxStatus{}).ProtoReflect().Descriptor(), handled, skipped)
 }