fix(server): prevent unbounded bus entry growth for sandbox IDs (#138)

Closes NVIDIA#27

Add remove() methods to TracingLogBus, PlatformEventBus, and
SandboxWatchBus to clean up entries when sandboxes are deleted.
Wire cleanup into both handle_deleted (K8s reconciler) and
delete_sandbox (gRPC handler). Reorder watch_sandbox to validate
sandbox existence before subscribing to buses, preventing entries
for non-existent IDs. Add one-time sandbox validation at stream
open in push_sandbox_logs.

Co-authored-by: John Myers <johntmyers@users.noreply.github.com>

Author: johntmyers

architecture/gateway.md

@@ -340,6 +340,10 @@ All buses use `tokio::sync::broadcast` channels keyed by sandbox ID. Buffer size
 
 Broadcast lag is translated to `Status::resource_exhausted` via `broadcast_to_status()`.
 
+**Cleanup:** Each bus exposes a `remove(sandbox_id)` method that drops the broadcast sender (closing active receivers with `RecvError::Closed`) and frees internal map entries. Cleanup is wired into both the `handle_deleted` reconciler (Kubernetes watcher) and the `delete_sandbox` gRPC handler to prevent unbounded memory growth from accumulated entries for deleted sandboxes.
+
+**Validation:** `WatchSandbox` validates that the sandbox exists before subscribing to any bus, preventing entries from being created for non-existent IDs. `PushSandboxLogs` validates sandbox existence once on the first batch of the stream.
+
 ## Remote Exec via SSH
 
 The `ExecSandbox` RPC (`crates/navigator-server/src/grpc.rs`) executes a command inside a sandbox pod over SSH and streams stdout/stderr/exit back to the client.

crates/navigator-server/src/grpc.rs

@@ -205,7 +205,23 @@ impl Navigator for NavigatorService {
 
         // Spawn producer task.
         tokio::spawn(async move {
-            // Subscribe to all buses BEFORE reading the initial snapshot to avoid
+            // Validate that the sandbox exists BEFORE subscribing to any buses.
+            // This prevents creating bus entries for non-existent sandbox IDs.
+            match state.store.get_message::<Sandbox>(&sandbox_id).await {
+                Ok(Some(_)) => {} // sandbox exists, proceed
+                Ok(None) => {
+                    let _ = tx.send(Err(Status::not_found("sandbox not found"))).await;
+                    return;
+                }
+                Err(e) => {
+                    let _ = tx
+                        .send(Err(Status::internal(format!("fetch sandbox failed: {e}"))))
+                        .await;
+                    return;
+                }
+            }
+
+            // Subscribe to all buses BEFORE reading the snapshot to avoid
             // missing notifications that fire between the snapshot read and subscribe.
             let mut status_rx = if follow_status {
                 Some(state.sandbox_watch_bus.subscribe(&sandbox_id))
@@ -228,7 +244,8 @@ impl Navigator for NavigatorService {
                 None
             };
 
-            // Always start with a snapshot if present.
+            // Re-read the snapshot now that we have subscriptions active
+            // (avoids missing notifications between validate and subscribe).
             match state.store.get_message::<Sandbox>(&sandbox_id).await {
                 Ok(Some(sandbox)) => {
                     state.sandbox_index.update_from_sandbox(&sandbox);
@@ -253,6 +270,7 @@ impl Navigator for NavigatorService {
                     }
                 }
                 Ok(None) => {
+                    // Sandbox was deleted between validate and subscribe — end stream.
                     let _ = tx.send(Err(Status::not_found("sandbox not found"))).await;
                     return;
                 }
@@ -481,6 +499,11 @@ impl Navigator for NavigatorService {
             warn!(sandbox_id = %id, error = %e, "Failed to clean up store after delete");
         }
 
+        // Clean up bus entries to prevent unbounded memory growth.
+        self.state.tracing_log_bus.remove(&id);
+        self.state.tracing_log_bus.platform_event_bus.remove(&id);
+        self.state.sandbox_watch_bus.remove(&id);
+
         info!(
             sandbox_id = %id,
             sandbox_name = %sandbox.name,
@@ -1147,6 +1170,7 @@ impl Navigator for NavigatorService {
         request: Request<tonic::Streaming<PushSandboxLogsRequest>>,
     ) -> Result<Response<PushSandboxLogsResponse>, Status> {
         let mut stream = request.into_inner();
+        let mut validated = false;
 
         while let Some(batch) = stream
             .message()
@@ -1157,6 +1181,20 @@ impl Navigator for NavigatorService {
                 continue;
             }
 
+            // Validate sandbox existence once at stream open (first batch).
+            // Subsequent batches trust the validated sandbox_id. If the sandbox
+            // is deleted mid-stream, bus remove() drops the sender and publish
+            // silently discards via `let _ = tx.send(...)`.
+            if !validated {
+                self.state
+                    .store
+                    .get_message::<Sandbox>(&batch.sandbox_id)
+                    .await
+                    .map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))?
+                    .ok_or_else(|| Status::not_found("sandbox not found"))?;
+                validated = true;
+            }
+
             // Cap lines per batch to prevent abuse.
             for log in batch.logs.into_iter().take(100) {
                 let mut log = log;

crates/navigator-server/src/lib.rs

@@ -128,6 +128,7 @@ pub async fn run_server(config: Config, tracing_log_bus: TracingLogBus) -> Resul
         state.sandbox_client.clone(),
         state.sandbox_index.clone(),
         state.sandbox_watch_bus.clone(),
+        state.tracing_log_bus.clone(),
     );
     spawn_kube_event_tailer(state.clone());
 

crates/navigator-server/src/sandbox/mod.rs

@@ -217,6 +217,7 @@ pub fn spawn_sandbox_watcher(
     client: SandboxClient,
     index: crate::sandbox_index::SandboxIndex,
     watch_bus: crate::sandbox_watch::SandboxWatchBus,
+    tracing_log_bus: crate::tracing_bus::TracingLogBus,
 ) {
     let namespace = client.namespace().to_string();
     info!(namespace = %namespace, "Starting sandbox watcher");
@@ -240,7 +241,9 @@ pub fn spawn_sandbox_watcher(
                     Event::Deleted(obj) => {
                         let obj_name = obj.metadata.name.clone().unwrap_or_default();
                         debug!(sandbox_name = %obj_name, "Received Deleted event from Kubernetes");
-                        if let Err(err) = handle_deleted(&store, &index, &watch_bus, obj).await {
+                        if let Err(err) =
+                            handle_deleted(&store, &index, &watch_bus, &tracing_log_bus, obj).await
+                        {
                             warn!(sandbox_name = %obj_name, error = %err, "Failed to delete sandbox record");
                         }
                     }
@@ -363,6 +366,7 @@ async fn handle_deleted(
     store: &Store,
     index: &crate::sandbox_index::SandboxIndex,
     watch_bus: &crate::sandbox_watch::SandboxWatchBus,
+    tracing_log_bus: &crate::tracing_bus::TracingLogBus,
     obj: DynamicObject,
 ) -> Result<(), String> {
     let id = sandbox_id_from_object(&obj)?;
@@ -373,6 +377,12 @@ async fn handle_deleted(
     debug!(sandbox_id = %id, deleted, "Deleted sandbox record");
     index.remove_sandbox(&id);
     watch_bus.notify(&id);
+
+    // Clean up bus entries to prevent unbounded memory growth.
+    tracing_log_bus.remove(&id);
+    tracing_log_bus.platform_event_bus.remove(&id);
+    watch_bus.remove(&id);
+
     Ok(())
 }
 

crates/navigator-server/src/sandbox_watch.rs

@@ -57,6 +57,15 @@ impl SandboxWatchBus {
     pub fn subscribe(&self, sandbox_id: &str) -> broadcast::Receiver<()> {
         self.sender_for(sandbox_id).subscribe()
     }
+
+    /// Remove the bus entry for the given sandbox id.
+    ///
+    /// This drops the broadcast sender, closing any active receivers with
+    /// `RecvError::Closed`.
+    pub fn remove(&self, sandbox_id: &str) {
+        let mut inner = self.inner.lock().expect("sandbox watch bus lock poisoned");
+        inner.remove(sandbox_id);
+    }
 }
 
 /// Spawn a background Kubernetes Event tailer.
@@ -160,6 +169,53 @@ fn map_kube_event_to_platform(
     ))
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn sandbox_watch_bus_remove_cleans_up() {
+        let bus = SandboxWatchBus::new();
+        let sandbox_id = "sb-1";
+
+        let mut rx = bus.subscribe(sandbox_id);
+
+        // Notify and receive
+        bus.notify(sandbox_id);
+        assert!(rx.try_recv().is_ok());
+
+        // Remove
+        bus.remove(sandbox_id);
+
+        // Receiver should be closed
+        match rx.try_recv() {
+            Err(broadcast::error::TryRecvError::Closed) => {} // expected
+            other => panic!("expected Closed, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn sandbox_watch_bus_subscribe_after_remove_creates_fresh_channel() {
+        let bus = SandboxWatchBus::new();
+        let sandbox_id = "sb-2";
+
+        let _old_rx = bus.subscribe(sandbox_id);
+        bus.remove(sandbox_id);
+
+        // New subscription should work
+        let mut new_rx = bus.subscribe(sandbox_id);
+        bus.notify(sandbox_id);
+        assert!(new_rx.try_recv().is_ok());
+    }
+
+    #[test]
+    fn sandbox_watch_bus_remove_nonexistent_is_noop() {
+        let bus = SandboxWatchBus::new();
+        // Should not panic
+        bus.remove("nonexistent");
+    }
+}
+
 /// Helper to translate broadcast lag into a gRPC status.
 pub fn broadcast_to_status(err: broadcast::error::RecvError) -> Status {
     match err {

crates/navigator-server/src/tracing_bus.rs

@@ -75,6 +75,16 @@ impl TracingLogBus {
         self.sender_for(sandbox_id).subscribe()
     }
 
+    /// Remove all bus entries for the given sandbox id.
+    ///
+    /// This drops the broadcast sender (closing any active receivers with
+    /// `RecvError::Closed`) and frees the tail buffer.
+    pub fn remove(&self, sandbox_id: &str) {
+        let mut inner = self.inner.lock().expect("tracing bus lock poisoned");
+        inner.per_id.remove(sandbox_id);
+        inner.tails.remove(sandbox_id);
+    }
+
     pub fn tail(&self, sandbox_id: &str, max: usize) -> Vec<SandboxStreamEvent> {
         let inner = self.inner.lock().expect("tracing bus lock poisoned");
         inner
@@ -186,6 +196,129 @@ fn current_time_ms() -> Option<i64> {
     i64::try_from(now.as_millis()).ok()
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn make_log_event(sandbox_id: &str, message: &str) -> SandboxLogLine {
+        SandboxLogLine {
+            sandbox_id: sandbox_id.to_string(),
+            timestamp_ms: 1000,
+            level: "INFO".to_string(),
+            target: "test".to_string(),
+            message: message.to_string(),
+            source: "gateway".to_string(),
+            fields: HashMap::new(),
+        }
+    }
+
+    #[test]
+    fn tracing_log_bus_remove_cleans_up_all_maps() {
+        let bus = TracingLogBus::new();
+        let sandbox_id = "sb-1";
+
+        // Create entries via subscribe and publish
+        let _rx = bus.subscribe(sandbox_id);
+        bus.publish_external(make_log_event(sandbox_id, "hello"));
+
+        // Verify entries exist
+        assert_eq!(bus.tail(sandbox_id, 10).len(), 1);
+
+        // Remove
+        bus.remove(sandbox_id);
+
+        // Verify entries are gone
+        assert!(bus.tail(sandbox_id, 10).is_empty());
+    }
+
+    #[test]
+    fn tracing_log_bus_subscribe_after_remove_creates_fresh_channel() {
+        let bus = TracingLogBus::new();
+        let sandbox_id = "sb-2";
+
+        // Create and remove
+        bus.publish_external(make_log_event(sandbox_id, "old message"));
+        bus.remove(sandbox_id);
+
+        // Subscribe again — should get a fresh channel with no history
+        let mut rx = bus.subscribe(sandbox_id);
+        assert!(bus.tail(sandbox_id, 10).is_empty());
+
+        // New publish should reach the new subscriber
+        bus.publish_external(make_log_event(sandbox_id, "new message"));
+        let evt = rx.try_recv().expect("should receive new event");
+        assert!(evt.payload.is_some());
+    }
+
+    #[test]
+    fn tracing_log_bus_remove_closes_active_receivers() {
+        let bus = TracingLogBus::new();
+        let sandbox_id = "sb-3";
+
+        let mut rx = bus.subscribe(sandbox_id);
+
+        // Remove drops the sender
+        bus.remove(sandbox_id);
+
+        // Existing receiver should get Closed error
+        match rx.try_recv() {
+            Err(broadcast::error::TryRecvError::Closed) => {} // expected
+            other => panic!("expected Closed, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn tracing_log_bus_remove_nonexistent_is_noop() {
+        let bus = TracingLogBus::new();
+        // Should not panic
+        bus.remove("nonexistent");
+    }
+
+    #[test]
+    fn platform_event_bus_remove_cleans_up() {
+        let bus = PlatformEventBus::new();
+        let sandbox_id = "sb-4";
+
+        let mut rx = bus.subscribe(sandbox_id);
+
+        // Publish an event
+        let evt = SandboxStreamEvent { payload: None };
+        bus.publish(sandbox_id, evt);
+        assert!(rx.try_recv().is_ok());
+
+        // Remove
+        bus.remove(sandbox_id);
+
+        // Receiver should be closed
+        match rx.try_recv() {
+            Err(broadcast::error::TryRecvError::Closed) => {} // expected
+            other => panic!("expected Closed, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn platform_event_bus_subscribe_after_remove_creates_fresh_channel() {
+        let bus = PlatformEventBus::new();
+        let sandbox_id = "sb-5";
+
+        let _old_rx = bus.subscribe(sandbox_id);
+        bus.remove(sandbox_id);
+
+        // New subscription should work
+        let mut new_rx = bus.subscribe(sandbox_id);
+        let evt = SandboxStreamEvent { payload: None };
+        bus.publish(sandbox_id, evt);
+        assert!(new_rx.try_recv().is_ok());
+    }
+
+    #[test]
+    fn platform_event_bus_remove_nonexistent_is_noop() {
+        let bus = PlatformEventBus::new();
+        // Should not panic
+        bus.remove("nonexistent");
+    }
+}
+
 /// Separate bus for platform event stream events.
 ///
 /// This keeps platform events isolated from tracing capture.
@@ -220,4 +353,12 @@ impl PlatformEventBus {
         let tx = self.sender_for(sandbox_id);
         let _ = tx.send(event);
     }
+
+    /// Remove the bus entry for the given sandbox id.
+    ///
+    /// This drops the broadcast sender, closing any active receivers.
+    pub(crate) fn remove(&self, sandbox_id: &str) {
+        let mut inner = self.inner.lock().expect("platform event bus lock poisoned");
+        inner.remove(sandbox_id);
+    }
 }