feat(server): cap in-flight relay channels per sandbox and globally

Adds two `ResourceExhausted`-returning guards on `open_relay` to bound the
`pending_relays` map against runaway or abusive callers:

- `MAX_PENDING_RELAYS = 256` — upper bound across all sandboxes. Caps the
  memory a caller can pin by calling `open_relay` in a loop while no
  supervisor ever claims (or the supervisor is hung).
- `MAX_PENDING_RELAYS_PER_SANDBOX = 32` — per-sandbox ceiling so one noisy
  tenant can't consume the entire global budget. Sits above the existing
  SSH-tunnel per-sandbox cap (20) so tunnel-specific limits still fire
  first for that caller.

Both checks and the `pending_relays` insert happen under a single lock
hold so concurrent callers can't each observe "under the cap" and both
insert past it. Adds a `sandbox_id` field on `PendingRelay` so the
per-sandbox count is a single filter over the map without extra indexes.

Tests:

- Two unit tests in `supervisor_session.rs` — assert the global cap and
  the per-sandbox cap both return `ResourceExhausted` with the right
  message, and a cap-hit on one sandbox doesn't leak onto others.
- One integration test in `supervisor_relay_integration.rs` — bursts 64
  concurrent `open_relay` calls at a single sandbox and asserts exactly
  32 succeed, exactly 32 are rejected with the per-sandbox message, and
  a different sandbox still accepts new relays.

Reaper behaviour is unchanged; the cap makes the map bounded, so the
existing `HashMap::retain` pass stays cheap under any load.

Author: pimlock

crates/openshell-server/src/supervisor_session.rs

@@ -25,6 +25,16 @@ const RELAY_PENDING_TIMEOUT: Duration = Duration::from_secs(10);
 const SESSION_WAIT_INITIAL_BACKOFF: Duration = Duration::from_millis(100);
 /// Maximum backoff between session-availability polls in `wait_for_session`.
 const SESSION_WAIT_MAX_BACKOFF: Duration = Duration::from_secs(2);
+/// Upper bound on unclaimed relay channels across all sandboxes. Caps the
+/// memory a misbehaving caller can pin by calling `open_relay` repeatedly
+/// while the supervisor never claims (or isn't responding). Sized generously
+/// so normal bursts pass through; exceeding it returns `ResourceExhausted`.
+const MAX_PENDING_RELAYS: usize = 256;
+/// Upper bound on concurrent unclaimed relay channels for a single sandbox.
+/// Enforces the same shape per sandbox so one misbehaving sandbox can't
+/// consume the entire global budget. Sits above the SSH-tunnel per-sandbox
+/// cap (20) so tunnel-specific limits still fire first for that caller.
+const MAX_PENDING_RELAYS_PER_SANDBOX: usize = 32;
 
 // ---------------------------------------------------------------------------
 // Session registry
@@ -56,6 +66,7 @@ pub struct SupervisorSessionRegistry {
 
 struct PendingRelay {
     sender: RelayStreamSender,
+    sandbox_id: String,
     created_at: Instant,
 }
 
@@ -186,13 +197,31 @@ impl SupervisorSessionRegistry {
         let channel_id = Uuid::new_v4().to_string();
 
         // Register the pending relay before sending RelayOpen to avoid a race.
+        // Both caps are checked and the insert happens under a single lock hold
+        // so two concurrent calls can't both observe "under the cap" and then
+        // both insert past it.
         let (relay_tx, relay_rx) = oneshot::channel();
         {
             let mut pending = self.pending_relays.lock().unwrap();
+            if pending.len() >= MAX_PENDING_RELAYS {
+                return Err(Status::resource_exhausted(format!(
+                    "gateway relay capacity reached ({MAX_PENDING_RELAYS} in flight)"
+                )));
+            }
+            let per_sandbox = pending
+                .values()
+                .filter(|p| p.sandbox_id == sandbox_id)
+                .count();
+            if per_sandbox >= MAX_PENDING_RELAYS_PER_SANDBOX {
+                return Err(Status::resource_exhausted(format!(
+                    "per-sandbox relay limit reached ({MAX_PENDING_RELAYS_PER_SANDBOX} in flight for {sandbox_id})"
+                )));
+            }
             pending.insert(
                 channel_id.clone(),
                 PendingRelay {
                     sender: relay_tx,
+                    sandbox_id: sandbox_id.to_string(),
                     created_at: Instant::now(),
                 },
             );
@@ -731,6 +760,76 @@ mod tests {
         assert!(registry.pending_relays.lock().unwrap().is_empty());
     }
 
+    #[tokio::test]
+    async fn open_relay_rejects_when_global_cap_reached() {
+        let registry = SupervisorSessionRegistry::new();
+        let (tx, _rx) = mpsc::channel::<GatewayMessage>(8);
+        registry.register("sbx-a".to_string(), "s-a".to_string(), tx.clone());
+        registry.register("sbx-b".to_string(), "s-b".to_string(), tx);
+
+        // Pre-seed pending_relays to exactly the global cap, split across two
+        // sandboxes so neither hits the per-sandbox cap first.
+        {
+            let mut pending = registry.pending_relays.lock().unwrap();
+            for i in 0..MAX_PENDING_RELAYS {
+                let (oneshot_tx, _) = oneshot::channel();
+                let sandbox_id = if i % 2 == 0 { "sbx-a" } else { "sbx-b" };
+                pending.insert(
+                    format!("channel-{i}"),
+                    PendingRelay {
+                        sender: oneshot_tx,
+                        sandbox_id: sandbox_id.to_string(),
+                        created_at: Instant::now(),
+                    },
+                );
+            }
+        }
+
+        let err = registry
+            .open_relay("sbx-a", Duration::from_millis(50))
+            .await
+            .expect_err("open_relay should reject once global cap is reached");
+        assert_eq!(err.code(), tonic::Code::ResourceExhausted);
+        assert!(err.message().contains("gateway relay capacity"));
+    }
+
+    #[tokio::test]
+    async fn open_relay_rejects_when_per_sandbox_cap_reached() {
+        let registry = SupervisorSessionRegistry::new();
+        let (tx, _rx) = mpsc::channel::<GatewayMessage>(8);
+        registry.register("sbx".to_string(), "s".to_string(), tx);
+
+        {
+            let mut pending = registry.pending_relays.lock().unwrap();
+            for i in 0..MAX_PENDING_RELAYS_PER_SANDBOX {
+                let (oneshot_tx, _) = oneshot::channel();
+                pending.insert(
+                    format!("channel-{i}"),
+                    PendingRelay {
+                        sender: oneshot_tx,
+                        sandbox_id: "sbx".to_string(),
+                        created_at: Instant::now(),
+                    },
+                );
+            }
+        }
+
+        let err = registry
+            .open_relay("sbx", Duration::from_millis(50))
+            .await
+            .expect_err("open_relay should reject when per-sandbox cap is reached");
+        assert_eq!(err.code(), tonic::Code::ResourceExhausted);
+        assert!(err.message().contains("per-sandbox relay limit"));
+
+        // A different sandbox still has headroom.
+        let (tx2, _rx2) = mpsc::channel::<GatewayMessage>(8);
+        registry.register("sbx-other".to_string(), "s-other".to_string(), tx2);
+        registry
+            .open_relay("sbx-other", Duration::from_millis(50))
+            .await
+            .expect("different sandbox should still accept new relays");
+    }
+
     #[tokio::test]
     async fn open_relay_uses_newest_session_after_supersede() {
         let registry = SupervisorSessionRegistry::new();
@@ -785,6 +884,7 @@ mod tests {
             "ch-1".to_string(),
             PendingRelay {
                 sender: relay_tx,
+                sandbox_id: "sbx-test".to_string(),
                 created_at: Instant::now(),
             },
         );
@@ -802,6 +902,7 @@ mod tests {
             "ch-old".to_string(),
             PendingRelay {
                 sender: relay_tx,
+                sandbox_id: "sbx-test".to_string(),
                 created_at: Instant::now() - Duration::from_secs(60),
             },
         );
@@ -829,6 +930,7 @@ mod tests {
             "ch-1".to_string(),
             PendingRelay {
                 sender: relay_tx,
+                sandbox_id: "sbx-test".to_string(),
                 created_at: Instant::now(),
             },
         );
@@ -847,6 +949,7 @@ mod tests {
             "ch-io".to_string(),
             PendingRelay {
                 sender: relay_tx,
+                sandbox_id: "sbx-test".to_string(),
                 created_at: Instant::now(),
             },
         );
@@ -877,6 +980,7 @@ mod tests {
             "ch-old".to_string(),
             PendingRelay {
                 sender: relay_tx,
+                sandbox_id: "sbx-test".to_string(),
                 created_at: Instant::now() - Duration::from_secs(60),
             },
         );
@@ -899,6 +1003,7 @@ mod tests {
             "ch-fresh".to_string(),
             PendingRelay {
                 sender: relay_tx,
+                sandbox_id: "sbx-test".to_string(),
                 created_at: Instant::now(),
             },
         );

crates/openshell-server/tests/supervisor_relay_integration.rs

@@ -306,7 +306,15 @@ fn register_session(
     registry: &SupervisorSessionRegistry,
     sandbox_id: &str,
 ) -> mpsc::Receiver<GatewayMessage> {
-    let (tx, rx) = mpsc::channel(8);
+    register_session_with_capacity(registry, sandbox_id, 8)
+}
+
+fn register_session_with_capacity(
+    registry: &SupervisorSessionRegistry,
+    sandbox_id: &str,
+    capacity: usize,
+) -> mpsc::Receiver<GatewayMessage> {
+    let (tx, rx) = mpsc::channel(capacity);
     registry.register(sandbox_id.to_string(), "sess-1".to_string(), tx);
     rx
 }
@@ -505,3 +513,54 @@ async fn concurrent_relays_multiplex_independently() {
     assert_eq!(&buf_a, b"stream-A");
     assert_eq!(&buf_b, b"stream-B");
 }
+
+/// Bursts more `open_relay` calls than the per-sandbox cap allows in parallel
+/// and asserts the registry enforces the ceiling cleanly. A well-behaved
+/// caller inside the cap still succeeds; overflow calls return `ResourceExhausted`
+/// rather than racing the pending map into an inconsistent state.
+#[tokio::test]
+async fn open_relay_enforces_per_sandbox_cap_under_concurrent_burst() {
+    let registry = Arc::new(SupervisorSessionRegistry::new());
+    let _channel = spawn_gateway(Arc::clone(&registry)).await;
+    // Oversized mpsc so the session doesn't backpressure the burst — the cap,
+    // not the channel, is what we're testing.
+    let _session_rx = register_session_with_capacity(&registry, "sbx", 256);
+
+    // Fire 64 concurrent opens. Per-sandbox cap is 32, global cap is 256,
+    // so exactly 32 should succeed and 32 should be rejected with
+    // `ResourceExhausted` carrying the per-sandbox message.
+    let mut handles = Vec::with_capacity(64);
+    for _ in 0..64 {
+        let r = Arc::clone(&registry);
+        handles.push(tokio::spawn(async move {
+            r.open_relay("sbx", Duration::from_secs(1)).await
+        }));
+    }
+
+    let mut ok = 0usize;
+    let mut exhausted = 0usize;
+    for h in handles {
+        match h.await.expect("task joined") {
+            Ok(_pair) => ok += 1,
+            Err(status) if status.code() == tonic::Code::ResourceExhausted => {
+                assert!(
+                    status.message().contains("per-sandbox relay limit"),
+                    "expected per-sandbox error message, got: {}",
+                    status.message()
+                );
+                exhausted += 1;
+            }
+            Err(other) => panic!("unexpected open_relay error: {other:?}"),
+        }
+    }
+    assert_eq!(ok, 32, "exactly per-sandbox cap should succeed");
+    assert_eq!(exhausted, 32, "overflow should be rejected, not dropped");
+
+    // A different sandbox still has headroom — the per-sandbox cap doesn't
+    // leak onto unrelated tenants.
+    let _other_rx = register_session_with_capacity(&registry, "sbx-other", 8);
+    registry
+        .open_relay("sbx-other", Duration::from_secs(1))
+        .await
+        .expect("other sandbox should not be affected by sbx cap");
+}