use std::collections::VecDeque; use std::sync::{Arc, Mutex}; use application::{ CancelAgentChat, LiveRuntimeState, LiveSessions, StructuredSessions, TerminalSessions, }; use async_trait::async_trait; use domain::ports::{AgentSession, AgentSessionError, CancelCurrentTurnOutcome, ReplyStream}; use domain::{AgentId, NodeId, ProjectId, SessionId}; use uuid::Uuid; struct FakeSession { id: SessionId, cancellations: Mutex>, sends: Mutex>, shutdowns: Mutex, } impl FakeSession { fn new(id: SessionId, outcomes: impl IntoIterator) -> Self { Self { id, cancellations: Mutex::new(outcomes.into_iter().collect()), sends: Mutex::new(Vec::new()), shutdowns: Mutex::new(0), } } } #[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 { Ok(self .cancellations .lock() .unwrap() .pop_front() .unwrap_or(CancelCurrentTurnOutcome::AlreadyCompleted)) } 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)) } #[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()); }