use std::collections::VecDeque; use std::sync::{Arc, Mutex}; use application::{ CancelAgentChat, LiveRuntimeState, LiveSessions, StructuredSessions, TerminalSessions, }; use async_trait::async_trait; use domain::input::{AgentBusyState, InputMediator, SubmitConfig}; use domain::mailbox::{AgentMailbox, MailboxError, PendingReply, Ticket, TicketId}; use domain::ports::{AgentSession, AgentSessionError, CancelCurrentTurnOutcome, ReplyStream}; use domain::{AgentId, NodeId, ProjectId, RuntimeAgentKey, SessionId}; use uuid::Uuid; struct FakeSession { id: SessionId, cancellations: Mutex>, cancel_hook: Mutex>>, sends: Mutex>, shutdowns: Mutex, } impl FakeSession { fn new(id: SessionId, outcomes: impl IntoIterator) -> Self { Self { id, cancellations: Mutex::new(outcomes.into_iter().collect()), cancel_hook: Mutex::new(None), sends: Mutex::new(Vec::new()), shutdowns: Mutex::new(0), } } fn with_cancel_hook(self, hook: impl FnOnce() + Send + 'static) -> Self { *self.cancel_hook.lock().unwrap() = Some(Box::new(hook)); self } } #[async_trait] impl AgentSession for FakeSession { fn id(&self) -> SessionId { self.id } fn conversation_id(&self) -> Option { Some("conversation-kept".to_owned()) } async fn send(&self, prompt: &str) -> Result { self.sends.lock().unwrap().push(prompt.to_owned()); Ok(Box::new(std::iter::empty())) } async fn cancel_current_turn(&self) -> Result { let outcome = self .cancellations .lock() .unwrap() .pop_front() .unwrap_or(CancelCurrentTurnOutcome::AlreadyCompleted); if let Some(hook) = self.cancel_hook.lock().unwrap().take() { hook(); } Ok(outcome) } async fn shutdown(&self) -> Result<(), AgentSessionError> { *self.shutdowns.lock().unwrap() += 1; Ok(()) } } fn id(value: u128, make: impl FnOnce(Uuid) -> T) -> T { make(Uuid::from_u128(value)) } struct TurnAccounting { state: Mutex, cancelled: Mutex>, } impl Default for TurnAccounting { fn default() -> Self { Self { state: Mutex::new(AgentBusyState::Idle), cancelled: Mutex::new(Vec::new()), } } } impl InputMediator for TurnAccounting { fn enqueue(&self, _agent: RuntimeAgentKey, _ticket: Ticket) -> PendingReply { panic!("not used") } fn bind_handle_with_submit( &self, _agent: RuntimeAgentKey, _handle: domain::ports::PtyHandle, _submit: SubmitConfig, ) { } fn preempt(&self, _agent: RuntimeAgentKey) {} fn mark_idle(&self, _agent: RuntimeAgentKey) { *self.state.lock().unwrap() = AgentBusyState::Idle; } fn busy_state(&self, _agent: RuntimeAgentKey) -> AgentBusyState { *self.state.lock().unwrap() } } impl AgentMailbox for TurnAccounting { fn enqueue(&self, _agent: RuntimeAgentKey, _ticket: Ticket) -> PendingReply { panic!("not used") } fn resolve(&self, agent: RuntimeAgentKey, _result: String) -> Result<(), MailboxError> { Err(MailboxError::NoPendingRequest(agent.agent_id)) } fn cancel_head(&self, agent: RuntimeAgentKey, ticket_id: TicketId) { self.cancelled.lock().unwrap().push((agent, ticket_id)); } } #[tokio::test] async fn cancel_keeps_session_and_accepts_a_second_prompt_without_shutdown() { let session_id = id(1, SessionId::from_uuid); let project_id = id(2, ProjectId::from_uuid); let agent_id = id(3, AgentId::from_uuid); let node_id = id(4, NodeId::from_uuid); let session = Arc::new(FakeSession::new( session_id, [CancelCurrentTurnOutcome::Cancelled], )); let sessions = Arc::new(StructuredSessions::new()); sessions.insert_in_project(project_id, session.clone(), agent_id, node_id); let live = LiveSessions::new(Arc::new(TerminalSessions::new()), Arc::clone(&sessions)); let cancel = CancelAgentChat::new(Arc::clone(&sessions)); let first_generation = sessions.begin_turn(&session_id).unwrap(); assert!(sessions.accepts_turn_output(&session_id, first_generation)); let output = cancel.execute(session_id).await.unwrap(); assert_eq!(output.outcome, CancelCurrentTurnOutcome::Cancelled); assert!(!sessions.accepts_turn_output(&session_id, first_generation)); let retained = sessions.session(&session_id).expect("session is retained"); let second_generation = sessions.begin_turn(&session_id).unwrap(); assert!(sessions.accepts_turn_output(&session_id, second_generation)); sessions.finish_turn(&session_id, first_generation); assert!( sessions.accepts_turn_output(&session_id, second_generation), "an old turn cannot retire or fence its successor" ); let _second_reply = retained.send("second prompt").await.unwrap(); assert_eq!(*session.sends.lock().unwrap(), ["second prompt"]); assert_eq!(*session.shutdowns.lock().unwrap(), 0); assert_eq!( live.runtime_status_for(project_id, &agent_id) .unwrap() .state, LiveRuntimeState::Running ); assert_eq!( sessions.meta_for_session(&session_id).unwrap().3.as_deref(), Some("conversation-kept") ); } #[tokio::test] async fn double_cancel_is_typed_and_missing_session_remains_distinct() { let session_id = id(10, SessionId::from_uuid); let session = Arc::new(FakeSession::new( session_id, [ CancelCurrentTurnOutcome::Cancelled, CancelCurrentTurnOutcome::AlreadyCancelled, ], )); let sessions = Arc::new(StructuredSessions::new()); sessions.insert_in_project( id(11, ProjectId::from_uuid), session, id(12, AgentId::from_uuid), id(13, NodeId::from_uuid), ); let cancel = CancelAgentChat::new(Arc::clone(&sessions)); sessions.begin_turn(&session_id).unwrap(); assert_eq!( cancel.execute(session_id).await.unwrap().outcome, CancelCurrentTurnOutcome::Cancelled ); assert_eq!( cancel.execute(session_id).await.unwrap().outcome, CancelCurrentTurnOutcome::AlreadyCancelled ); let error = cancel .execute(id(99, SessionId::from_uuid)) .await .unwrap_err(); assert_eq!(error.code(), "NOT_FOUND"); assert!(sessions.session(&session_id).is_some()); } #[tokio::test] async fn cancel_retires_the_exact_busy_turn_before_the_next_prompt() { let session_id = id(20, SessionId::from_uuid); let project_id = id(21, ProjectId::from_uuid); let agent_id = id(22, AgentId::from_uuid); let ticket = TicketId::from_uuid(Uuid::from_u128(23)); let accounting = Arc::new(TurnAccounting::default()); *accounting.state.lock().unwrap() = AgentBusyState::Busy { ticket, since_ms: 1, }; let session = Arc::new(FakeSession::new( session_id, [CancelCurrentTurnOutcome::Cancelled], )); let sessions = Arc::new(StructuredSessions::new()); sessions.insert_in_project(project_id, session, agent_id, id(24, NodeId::from_uuid)); let generation = sessions.begin_turn(&session_id).unwrap(); assert!(sessions.bind_turn_accounting( &session_id, generation, RuntimeAgentKey::new(project_id, agent_id), ticket, )); let input = Arc::clone(&accounting) as Arc; let mailbox = Arc::clone(&accounting) as Arc; let cancel = CancelAgentChat::new(Arc::clone(&sessions)).with_turn_accounting(input, mailbox); let output = cancel.execute(session_id).await.unwrap(); assert_eq!(output.outcome, CancelCurrentTurnOutcome::Cancelled); assert_eq!(*accounting.state.lock().unwrap(), AgentBusyState::Idle); assert_eq!( *accounting.cancelled.lock().unwrap(), [(RuntimeAgentKey::new(project_id, agent_id), ticket)] ); } #[tokio::test] async fn late_cancel_cleanup_does_not_mark_a_successor_ticket_idle() { let session_id = id(30, SessionId::from_uuid); let project_id = id(31, ProjectId::from_uuid); let agent_id = id(32, AgentId::from_uuid); let cancelled_ticket = TicketId::from_uuid(Uuid::from_u128(33)); let successor_ticket = TicketId::from_uuid(Uuid::from_u128(34)); let accounting = Arc::new(TurnAccounting::default()); *accounting.state.lock().unwrap() = AgentBusyState::Busy { ticket: cancelled_ticket, since_ms: 1, }; let accounting_during_cancel = Arc::clone(&accounting); let session = Arc::new( FakeSession::new(session_id, [CancelCurrentTurnOutcome::Cancelled]).with_cancel_hook( move || { *accounting_during_cancel.state.lock().unwrap() = AgentBusyState::Busy { ticket: successor_ticket, since_ms: 2, }; }, ), ); let sessions = Arc::new(StructuredSessions::new()); sessions.insert_in_project(project_id, session, agent_id, id(35, NodeId::from_uuid)); let generation = sessions.begin_turn(&session_id).unwrap(); assert!(sessions.bind_turn_accounting( &session_id, generation, RuntimeAgentKey::new(project_id, agent_id), cancelled_ticket, )); let input = Arc::clone(&accounting) as Arc; let mailbox = Arc::clone(&accounting) as Arc; let cancel = CancelAgentChat::new(Arc::clone(&sessions)).with_turn_accounting(input, mailbox); cancel.execute(session_id).await.unwrap(); assert_eq!( *accounting.state.lock().unwrap(), AgentBusyState::Busy { ticket: successor_ticket, since_ms: 2, } ); assert!(accounting.cancelled.lock().unwrap().is_empty()); }