IdeA/crates/application/tests/cancel_agent_chat.rs

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<VecDeque<CancelCurrentTurnOutcome>>,
    cancel_hook: Mutex<Option<Box<dyn FnOnce() + Send>>>,
    sends: Mutex<Vec<String>>,
    shutdowns: Mutex<usize>,
}

impl FakeSession {
    fn new(id: SessionId, outcomes: impl IntoIterator<Item = CancelCurrentTurnOutcome>) -> 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<String> {
        Some("conversation-kept".to_owned())
    }

    async fn send(&self, prompt: &str) -> Result<ReplyStream, AgentSessionError> {
        self.sends.lock().unwrap().push(prompt.to_owned());
        Ok(Box::new(std::iter::empty()))
    }

    async fn cancel_current_turn(&self) -> Result<CancelCurrentTurnOutcome, AgentSessionError> {
        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<T>(value: u128, make: impl FnOnce(Uuid) -> T) -> T {
    make(Uuid::from_u128(value))
}

struct TurnAccounting {
    state: Mutex<AgentBusyState>,
    cancelled: Mutex<Vec<(RuntimeAgentKey, TicketId)>>,
}

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<dyn InputMediator>;
    let mailbox = Arc::clone(&accounting) as Arc<dyn AgentMailbox>;
    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<dyn InputMediator>;
    let mailbox = Arc::clone(&accounting) as Arc<dyn AgentMailbox>;
    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());
}