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::ports::{AgentSession, AgentSessionError, CancelCurrentTurnOutcome, ReplyStream};
use domain::{AgentId, NodeId, ProjectId, SessionId};
use uuid::Uuid;

struct FakeSession {
    id: SessionId,
    cancellations: Mutex<VecDeque<CancelCurrentTurnOutcome>>,
    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()),
            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<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> {
        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<T>(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());
}