diff --git a/crates/server/src/router/ws/agent/mod.rs b/crates/server/src/router/ws/agent/mod.rs index 8fe372b5..b601b33f 100644 --- a/crates/server/src/router/ws/agent/mod.rs +++ b/crates/server/src/router/ws/agent/mod.rs @@ -281,11 +281,6 @@ async fn handle_agent_ws( { crate::service::agent_manager::cleanup_disconnected_docker_state(&state, &server_id).await; } - // Drop any temporary-grant countdowns so a disconnected agent does not leave - // stale grants ticking in the REST DTO / browser view. - state - .agent_manager - .update_temporary_grants(&server_id, vec![]); write_task.abort(); tracing::info!("Agent {server_id} disconnected"); } diff --git a/crates/server/src/service/agent_manager.rs b/crates/server/src/service/agent_manager.rs index 5b964c9e..94eeb1fa 100644 --- a/crates/server/src/service/agent_manager.rs +++ b/crates/server/src/service/agent_manager.rs @@ -433,6 +433,7 @@ impl AgentManager { fn finish_connection_removal(&self, server_id: &str) { self.agent_local_capabilities.remove(server_id); + self.temporary_grants.remove(server_id); self.remove_docker_log_sessions_for_server(server_id); self.clear_docker_caches(server_id); @@ -1266,6 +1267,29 @@ mod tests { assert!(mgr.get_docker_containers("s1").is_some()); } + #[test] + fn test_remove_connection_if_current_scopes_temporary_grant_cleanup() { + let (mgr, _rx) = make_manager(); + let (tx1, _) = mpsc::channel(1); + let (tx2, _) = mpsc::channel(1); + let first_connection_id = mgr.add_connection("s1".into(), "Srv".into(), tx1, test_addr()); + let second_connection_id = mgr.add_connection("s1".into(), "Srv".into(), tx2, test_addr()); + mgr.update_temporary_grants( + "s1", + vec![TemporaryGrant { + cap: "terminal".into(), + granted_at: 1, + expires_at: 100, + }], + ); + + assert!(!mgr.remove_connection_if_current("s1", first_connection_id)); + assert_eq!(mgr.get_temporary_grants("s1").len(), 1); + + assert!(mgr.remove_connection_if_current("s1", second_connection_id)); + assert!(mgr.get_temporary_grants("s1").is_empty()); + } + #[test] fn test_stale_connection_candidates_do_not_remove_newer_connection() { let (mgr, _rx) = make_manager(); diff --git a/crates/server/tests/ws_agent_dispatch2.rs b/crates/server/tests/ws_agent_dispatch2.rs index d112cd27..cb7ddaa8 100644 --- a/crates/server/tests/ws_agent_dispatch2.rs +++ b/crates/server/tests/ws_agent_dispatch2.rs @@ -30,7 +30,8 @@ //! - DockerStats (unsolicited) -> stats visible via GET /docker/stats //! - SystemInfo (second, changed caps) -> mirror update + capabilities_changed //! - reconnect / superseded connection -> second connect wins; the first -//! socket's next frame stops its read loop (end-to-end, not the unit arm) +//! socket's next frame stops its read loop without clearing the current +//! connection's temporary grants (end-to-end, not the unit arm) //! - SecurityEvent (ssh_login, FULL evidence) -> persisted, queryable //! - CapabilityDenied (capability=terminal, with session_id) -> the //! terminal-session unregister arm; connection survives @@ -481,7 +482,7 @@ async fn test_second_system_info_updates_capability_mirror() { // =========================================================================== #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn test_reconnect_supersedes_first_connection() { +async fn test_superseded_connection_cleanup_preserves_current_temporary_grants() { let (base_url, _tmp) = start_test_server().await; let client = http_client(); login_admin(&client, &base_url).await; @@ -501,6 +502,51 @@ async fn test_reconnect_supersedes_first_connection() { send_system_info(&mut sink_b, &mut reader_b, "recon-b", Some(CAP_DEFAULT)).await; drain_first_connect_pushes(&mut reader_b, 300).await; + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system clock before epoch") + .as_secs() as i64; + let expires_at = now + 3600; + send_agent_frame( + &mut sink_b, + json!({ + "type": "capabilities_changed", + "msg_id": "recon-b-grant", + "capabilities": CAP_DEFAULT | CAP_TERMINAL, + "temporary": [{ + "cap": "terminal", + "granted_at": now, + "expires_at": expires_at + }], + "changes": [] + }), + ) + .await; + + let mut grant_visible = false; + for _ in 0..20 { + let resp = client + .get(format!("{base_url}/api/servers/{server_id}")) + .send() + .await + .expect("GET server failed"); + assert_eq!(resp.status(), 200); + let body: Value = resp.json().await.expect("parse server"); + grant_visible = body["data"]["temporary"] + .as_array() + .is_some_and(|grants| { + grants.iter().any(|grant| { + grant["cap"].as_str() == Some("terminal") + && grant["expires_at"].as_i64() == Some(expires_at) + }) + }); + if grant_visible { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + assert!(grant_visible, "the current connection should report its temporary grant"); + // The FIRST socket is now superseded: handle_current_connection_frame returns // false on its next frame, which breaks its read loop and closes the socket. // Sending a frame on A and then draining its stream must reach end-of-stream. @@ -532,6 +578,24 @@ async fn test_reconnect_supersedes_first_connection() { "the superseded first connection should be torn down once it sends a frame" ); + let resp = client + .get(format!("{base_url}/api/servers/{server_id}")) + .send() + .await + .expect("GET server failed"); + assert_eq!(resp.status(), 200); + let body: Value = resp.json().await.expect("parse server"); + let temporary = body["data"]["temporary"] + .as_array() + .expect("temporary grants should be an array"); + assert!( + temporary.iter().any(|grant| { + grant["cap"].as_str() == Some("terminal") + && grant["expires_at"].as_i64() == Some(expires_at) + }), + "superseded connection cleanup should preserve the current connection's temporary grants" + ); + // The SECOND connection is the live one: it still gets Acked. send_system_info(&mut sink_b, &mut reader_b, "recon-b-still-live", Some(CAP_DEFAULT)).await;