328 lines
11 KiB
Rust
328 lines
11 KiB
Rust
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))
|
|
}
|
|
|
|
#[test]
|
|
fn custom_conversation_binding_keeps_logical_id_and_consumes_resume_once() {
|
|
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 conversation_id = domain::ConversationId::from_uuid(Uuid::from_u128(5));
|
|
let sessions = StructuredSessions::new();
|
|
sessions.insert_in_project(
|
|
project_id,
|
|
Arc::new(FakeSession::new(session_id, [])),
|
|
agent_id,
|
|
node_id,
|
|
);
|
|
|
|
sessions.bind_custom_conversation(
|
|
session_id,
|
|
conversation_id,
|
|
Some("bounded handoff".to_owned()),
|
|
);
|
|
|
|
assert_eq!(
|
|
sessions.custom_conversation_id(&session_id),
|
|
Some(conversation_id)
|
|
);
|
|
assert_eq!(
|
|
sessions
|
|
.take_pending_resume_preamble(&session_id)
|
|
.as_deref(),
|
|
Some("bounded handoff")
|
|
);
|
|
assert_eq!(sessions.take_pending_resume_preamble(&session_id), None);
|
|
}
|
|
|
|
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());
|
|
}
|