//! Tests for the Lot D agent-level controlled actions (`AttachLiveAgent`, //! `StopLiveAgent`) over the live-session registries. //! //! Every port is faked in-memory so the use cases run without a real PTY or //! agent-session backend: [`FakePty`] records `kill`s, [`FakeSession`] records //! `shutdown`s. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use async_trait::async_trait; use application::{ AgentRestartScheduled, AgentSelfRestartScheduler, AppError, AttachLiveAgent, AttachLiveAgentInput, CellAgentPresentation, CloseTerminal, LiveRuntimeState, LiveSessionKind, LiveSessions, RestartCellAgentRuntime, RestartCellAgentRuntimeInput, StopAgentExecutionCascade, StopLiveAgent, StopLiveAgentInput, StructuredSessions, TerminalSessions, }; use domain::ports::{ AgentSession, AgentSessionError, ExitStatus, OutputStream, PtyError, PtyHandle, PtyPort, ReplyStream, SpawnSpec, }; use domain::Project; use domain::{ AgentId, LayoutId, NodeId, ProjectId, ProjectPath, PtySize, RemoteRef, SessionId, SessionKind, TerminalSession, }; use uuid::Uuid; // --------------------------------------------------------------------------- // ids / fixtures // --------------------------------------------------------------------------- fn aid(n: u128) -> AgentId { AgentId::from_uuid(Uuid::from_u128(n)) } fn sid(n: u128) -> SessionId { SessionId::from_uuid(Uuid::from_u128(n)) } fn nid(n: u128) -> NodeId { NodeId::from_uuid(Uuid::from_u128(n)) } fn project() -> Project { Project::new( ProjectId::from_uuid(Uuid::from_u128(1)), "demo", ProjectPath::new("/tmp/idea-workstate-actions").unwrap(), RemoteRef::local(), 1_700_000_000_000, ) .unwrap() } // --------------------------------------------------------------------------- // Fakes // --------------------------------------------------------------------------- /// A recording [`PtyPort`] whose `kill` records the killed session id. #[derive(Default)] struct FakePty { kills: Mutex>, } impl FakePty { fn kills(&self) -> Vec { self.kills.lock().unwrap().clone() } } #[async_trait] impl PtyPort for FakePty { async fn spawn(&self, _spec: SpawnSpec, _size: PtySize) -> Result { unreachable!("Lot D never spawns") } fn write(&self, _handle: &PtyHandle, _data: &[u8]) -> Result<(), PtyError> { Ok(()) } fn resize(&self, _handle: &PtyHandle, _size: PtySize) -> Result<(), PtyError> { Ok(()) } fn subscribe_output(&self, _handle: &PtyHandle) -> Result { Ok(Box::new(std::iter::empty())) } fn scrollback(&self, _handle: &PtyHandle) -> Result, PtyError> { Ok(Vec::new()) } async fn wait(&self, _handle: &PtyHandle) -> Result { Ok(ExitStatus { code: Some(0) }) } fn try_wait(&self, _handle: &PtyHandle) -> Result, PtyError> { Ok(Some(ExitStatus { code: Some(0) })) } async fn kill(&self, handle: &PtyHandle) -> Result { self.kills.lock().unwrap().push(handle.session_id); Ok(ExitStatus { code: Some(0) }) } } /// A fake structured [`AgentSession`] recording stop/cancel calls. struct FakeSession { id: SessionId, shutdown_called: Arc, cancel_called: Arc, } #[async_trait] impl AgentSession for FakeSession { fn id(&self) -> SessionId { self.id } fn conversation_id(&self) -> Option { None } async fn send(&self, _prompt: &str) -> Result { Ok(Box::new(std::iter::empty())) } async fn cancel_current_turn(&self) -> Result<(), AgentSessionError> { self.cancel_called.store(true, Ordering::SeqCst); Ok(()) } async fn shutdown(&self) -> Result<(), AgentSessionError> { self.shutdown_called.store(true, Ordering::SeqCst); Ok(()) } } #[derive(Default)] struct FakeCascade { calls: Mutex>, } #[derive(Default)] struct FakeRestartScheduler { calls: Mutex>, } impl AgentSelfRestartScheduler for FakeRestartScheduler { fn schedule( &self, _project: &Project, _agent: AgentId, _reason: Option, ) -> Result { unreachable!("cell restart must use the cell-aware scheduler seam") } fn schedule_cell( &self, project: &Project, agent: AgentId, node: NodeId, _reason: Option, ) -> Result { self.calls.lock().unwrap().push((project.id, agent, node)); Ok(AgentRestartScheduled { restart_id: "restart-cell-test".to_owned(), }) } } #[async_trait] impl StopAgentExecutionCascade for FakeCascade { async fn stop_agent_execution_children(&self, project: &Project, agent_id: AgentId) { self.calls.lock().unwrap().push((project.id, agent_id)); } } // --------------------------------------------------------------------------- // wiring helpers // --------------------------------------------------------------------------- struct Fixture { attach: AttachLiveAgent, stop: StopLiveAgent, pty: Arc, structured: Arc, pty_port: Arc, restart: RestartCellAgentRuntime, restart_scheduler: Arc, } fn fixture() -> Fixture { let pty = Arc::new(TerminalSessions::new()); let structured = Arc::new(StructuredSessions::new()); let live = Arc::new(LiveSessions::new(Arc::clone(&pty), Arc::clone(&structured))); let pty_port = Arc::new(FakePty::default()); let close = Arc::new(CloseTerminal::new( Arc::clone(&pty_port) as Arc, Arc::clone(&pty), )); let restart_scheduler = Arc::new(FakeRestartScheduler::default()); Fixture { attach: AttachLiveAgent::new(Arc::clone(&live)), stop: StopLiveAgent::new(Arc::clone(&live), close), pty, structured, pty_port, restart: RestartCellAgentRuntime::new( Arc::clone(&live), Arc::clone(&restart_scheduler) as Arc, ), restart_scheduler, } } fn restart_input( agent_id: AgentId, node_id: NodeId, presentation: CellAgentPresentation, ) -> RestartCellAgentRuntimeInput { RestartCellAgentRuntimeInput { project: project(), node_id, agent_id, presentation, layout_id: Some(LayoutId::from_uuid(Uuid::from_u128(500))), } } fn insert_pty( sessions: &TerminalSessions, session_id: SessionId, agent_id: AgentId, node_id: NodeId, ) { sessions.insert_in_project( project().id, PtyHandle { session_id }, TerminalSession::starting( session_id, node_id, ProjectPath::new("/tmp/idea-workstate-actions").unwrap(), SessionKind::Agent { agent_id }, PtySize { rows: 24, cols: 80 }, ), ); } fn insert_structured( sessions: &StructuredSessions, session_id: SessionId, agent_id: AgentId, node_id: NodeId, ) -> (Arc, Arc) { let shutdown_called = Arc::new(AtomicBool::new(false)); let cancel_called = Arc::new(AtomicBool::new(false)); sessions.insert_in_project( project().id, Arc::new(FakeSession { id: session_id, shutdown_called: Arc::clone(&shutdown_called), cancel_called: Arc::clone(&cancel_called), }), agent_id, node_id, ); (shutdown_called, cancel_called) } // --------------------------------------------------------------------------- // AttachLiveAgent // --------------------------------------------------------------------------- #[test] fn attach_pty_rebinds_node_without_changing_session() { let f = fixture(); let a = aid(10); insert_pty(&f.pty, sid(1), a, nid(100)); let out = f .attach .execute(AttachLiveAgentInput { project: project(), agent_id: a, node_id: nid(200), }) .unwrap(); assert_eq!(out.agent_id, a); assert_eq!(out.session_id, sid(1), "session id is preserved on rebind"); assert_eq!(out.node_id, nid(200), "view rebound to the new node"); assert_eq!(out.kind, LiveSessionKind::Pty); // The registry reflects the new host node, same session. assert_eq!( f.pty.node_for_agent_in_project(project().id, &a), Some(nid(200)) ); assert_eq!( f.pty.session_for_agent_in_project(project().id, &a), Some(sid(1)) ); } #[test] fn attach_structured_rebinds_node() { let f = fixture(); let a = aid(10); let _flags = insert_structured(&f.structured, sid(2), a, nid(100)); let out = f .attach .execute(AttachLiveAgentInput { project: project(), agent_id: a, node_id: nid(300), }) .unwrap(); assert_eq!(out.session_id, sid(2)); assert_eq!(out.node_id, nid(300)); assert_eq!(out.kind, LiveSessionKind::Structured); assert_eq!( f.structured.node_for_agent_in_project(project().id, &a), Some(nid(300)) ); } #[test] fn attach_unknown_agent_is_not_found() { let f = fixture(); let err = f .attach .execute(AttachLiveAgentInput { project: project(), agent_id: aid(999), node_id: nid(1), }) .expect_err("absent agent ⇒ NOT_FOUND"); assert!( matches!(err, application::AppError::NotFound(_)), "got {err:?}" ); } #[test] fn attach_same_node_is_idempotent() { let f = fixture(); let a = aid(10); insert_pty(&f.pty, sid(1), a, nid(100)); let first = f .attach .execute(AttachLiveAgentInput { project: project(), agent_id: a, node_id: nid(100), }) .unwrap(); let second = f .attach .execute(AttachLiveAgentInput { project: project(), agent_id: a, node_id: nid(100), }) .unwrap(); // Re-attaching to the same node is a no-op: same session, same node, no spawn. assert_eq!(first.session_id, second.session_id); assert_eq!(first.node_id, nid(100)); assert_eq!(second.node_id, nid(100)); assert_eq!(f.pty.len(), 1, "no extra session created"); } // --------------------------------------------------------------------------- // RestartCellAgentRuntime // --------------------------------------------------------------------------- #[test] fn restart_cell_accepts_tui_and_preserves_exact_binding() { let f = fixture(); let agent = aid(10); let node = nid(100); insert_pty(&f.pty, sid(1), agent, node); let out = f .restart .execute(restart_input(agent, node, CellAgentPresentation::Tui)) .unwrap(); assert_eq!(out.restart_id, "restart-cell-test"); assert_eq!(out.agent_id, agent); assert_eq!(out.node_id, node); assert_eq!(out.kind, LiveSessionKind::Pty); assert_eq!(out.presentation, CellAgentPresentation::Tui); assert_eq!( *f.restart_scheduler.calls.lock().unwrap(), vec![(project().id, agent, node)] ); } #[test] fn restart_cell_accepts_structured_custom_cli_and_headless_hints() { for presentation in [ CellAgentPresentation::CustomCli, CellAgentPresentation::Headless, ] { let f = fixture(); let agent = aid(10); let node = nid(100); insert_structured(&f.structured, sid(2), agent, node); let out = f .restart .execute(restart_input(agent, node, presentation)) .unwrap(); assert_eq!(out.kind, LiveSessionKind::Structured); assert_eq!(out.presentation, presentation); assert_eq!(f.restart_scheduler.calls.lock().unwrap().len(), 1); } } #[test] fn restart_cell_rejects_stale_agent_or_node_without_scheduling() { let f = fixture(); let agent = aid(10); insert_pty(&f.pty, sid(1), agent, nid(100)); let error = f .restart .execute(restart_input( agent, nid(999), CellAgentPresentation::Unknown, )) .unwrap_err(); assert_eq!(error.code(), "AGENT_NOT_BOUND_TO_CELL"); assert!(f.restart_scheduler.calls.lock().unwrap().is_empty()); } #[test] fn restart_cell_rejects_missing_runtime_without_scheduling() { let f = fixture(); let error = f .restart .execute(restart_input( aid(999), nid(100), CellAgentPresentation::Unknown, )) .unwrap_err(); assert_eq!(error.code(), "RUNTIME_NOT_FOUND"); assert!(f.restart_scheduler.calls.lock().unwrap().is_empty()); } #[test] fn restart_cell_rejects_agent_visible_in_multiple_live_cells() { let f = fixture(); let agent = aid(10); insert_pty(&f.pty, sid(1), agent, nid(100)); insert_structured(&f.structured, sid(2), agent, nid(200)); let error = f .restart .execute(restart_input(agent, nid(100), CellAgentPresentation::Tui)) .unwrap_err(); assert_eq!(error.code(), "CONFLICT_AGENT_VISIBLE_ELSEWHERE"); assert!(f.restart_scheduler.calls.lock().unwrap().is_empty()); } // --------------------------------------------------------------------------- // StopLiveAgent // --------------------------------------------------------------------------- #[tokio::test] async fn stop_pty_kills_and_removes_session() { let f = fixture(); let a = aid(10); insert_pty(&f.pty, sid(1), a, nid(100)); let out = f .stop .execute(StopLiveAgentInput { project: project(), agent_id: a, reason: None, }) .await .unwrap(); assert_eq!(out.agent_id, a); assert!(!out.already_stopped); assert_eq!(out.stopped[0].session_id, sid(1)); assert_eq!(out.stopped[0].node_id, nid(100)); assert_eq!(out.stopped[0].kind, LiveSessionKind::Pty); assert_eq!(out.stopped[0].state, LiveRuntimeState::Stopped); // Delegated to the close primitive: process killed and registry emptied. assert_eq!(f.pty_port.kills(), vec![sid(1)]); assert!(f.pty.is_empty(), "live session removed from the registry"); assert_eq!(f.pty.session_for_agent_in_project(project().id, &a), None); } #[tokio::test] async fn stop_structured_shuts_down_and_removes_session() { let f = fixture(); let a = aid(10); let (shutdown_called, cancel_called) = insert_structured(&f.structured, sid(2), a, nid(100)); let out = f .stop .execute(StopLiveAgentInput { project: project(), agent_id: a, reason: None, }) .await .unwrap(); assert_eq!(out.stopped[0].session_id, sid(2)); assert_eq!(out.stopped[0].kind, LiveSessionKind::Structured); assert_eq!(out.stopped[0].state, LiveRuntimeState::Stopped); assert!(!cancel_called.load(Ordering::SeqCst)); assert!( shutdown_called.load(Ordering::SeqCst), "session.shutdown() must be called" ); assert!(f.structured.is_empty(), "live session is removed"); assert_eq!( f.structured .session_id_for_agent_in_project(project().id, &a), None ); // No PTY was touched. assert!(f.pty_port.kills().is_empty()); } #[tokio::test] async fn stop_structured_runs_execution_cascade_best_effort() { let f = fixture(); let a = aid(10); let (shutdown_called, cancel_called) = insert_structured(&f.structured, sid(2), a, nid(100)); let cascade = Arc::new(FakeCascade::default()); let stop = f .stop .with_execution_cascade(Arc::clone(&cascade) as Arc); let out = stop .execute(StopLiveAgentInput { project: project(), agent_id: a, reason: None, }) .await .unwrap(); assert_eq!(out.stopped[0].session_id, sid(2)); assert_eq!(*cascade.calls.lock().unwrap(), vec![(project().id, a)]); assert!( !cancel_called.load(Ordering::SeqCst), "session.cancel_current_turn() is not the runtime stop primitive" ); assert!( shutdown_called.load(Ordering::SeqCst), "session.shutdown() must be called" ); assert!(f.structured.is_empty(), "live session is removed"); } #[tokio::test] async fn stop_removes_every_runtime_for_the_agent_across_registries() { let f = fixture(); let agent = aid(10); insert_pty(&f.pty, sid(1), agent, nid(100)); let (shutdown_called, _) = insert_structured(&f.structured, sid(2), agent, nid(200)); let out = f .stop .execute(StopLiveAgentInput { project: project(), agent_id: agent, reason: Some("user requested".to_owned()), }) .await .unwrap(); assert_eq!(out.stopped.len(), 2); assert!(f.pty.is_empty()); assert!(f.structured.is_empty()); assert_eq!(f.pty_port.kills(), vec![sid(1)]); assert!(shutdown_called.load(Ordering::SeqCst)); } #[tokio::test] async fn stop_pty_is_idempotent_after_success() { let f = fixture(); let a = aid(10); insert_pty(&f.pty, sid(1), a, nid(100)); let first = f .stop .execute(StopLiveAgentInput { project: project(), agent_id: a, reason: None, }) .await .unwrap(); let second = f .stop .execute(StopLiveAgentInput { project: project(), agent_id: a, reason: None, }) .await .unwrap(); assert!(!first.already_stopped); assert!(second.already_stopped); assert!(second.stopped.is_empty()); assert_eq!(f.pty_port.kills(), vec![sid(1)], "kill is not repeated"); } #[tokio::test] async fn stop_unknown_agent_is_idempotent_success() { let f = fixture(); let out = f .stop .execute(StopLiveAgentInput { project: project(), agent_id: aid(999), reason: None, }) .await .expect("absent agent is an idempotent success"); assert!(out.already_stopped); assert!(out.stopped.is_empty()); assert!(f.pty_port.kills().is_empty(), "nothing killed"); }