blomios/IdeA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(runtime): isolate agent state by project (#101)

Browse Source
This commit is contained in:
Blomios
2026-07-25 22:14:02 +02:00
parent 6a87c4635f
commit 6e98fd89f7
52 changed files with 1894 additions and 805 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
crates/application/tests/agent_wake.rs
View File

@ -7,7 +7,7 @@ use domain::background_task::{
BackgroundTask, BackgroundTaskKind, BackgroundTaskResult, BackgroundTaskState,
BackgroundTaskWakePolicy,
};
use domain::ids::{AgentId, ProjectId, SessionId, TaskId};
use domain::ids::{AgentId, ProjectId, RuntimeAgentKey, SessionId, TaskId};
use domain::inbox::{
AgentInbox, AgentInboxSnapshot, InboxError, InboxItem, InboxItemKind, InboxReceipt,
InboxReceiptStatus, InboxSource,
@ -30,6 +30,10 @@ fn agent(n: u128) -> AgentId {
AgentId::from_uuid(id(n))
}
fn runtime_key(agent_id: AgentId) -> RuntimeAgentKey {
RuntimeAgentKey::new(ProjectId::from_uuid(id(100)), agent_id)
}
fn task_id(n: u128) -> TaskId {
TaskId::from_uuid(id(n))
}
@ -93,44 +97,46 @@ fn completed_task(
#[derive(Default)]
struct FakeInbox {
queues: Mutex<HashMap<AgentId, VecDeque<InboxItem>>>,
queues: Mutex<HashMap<RuntimeAgentKey, VecDeque<InboxItem>>>,
}
impl AgentInbox for FakeInbox {
fn enqueue_message(
&self,
agent_id: AgentId,
agent: RuntimeAgentKey,
item: InboxItem,
) -> Result<InboxReceipt, InboxError> {
let mut queues = self.queues.lock().unwrap();
let queue = queues.entry(agent_id).or_default();
let queue = queues.entry(agent).or_default();
let item_id = item.id;
queue.push_back(item);
Ok(InboxReceipt {
item_id,
agent_id,
agent_id: agent.agent_id,
runtime_key: agent,
depth: queue.len(),
status: InboxReceiptStatus::Queued,
})
}
fn dequeue_next(&self, agent_id: AgentId) -> Option<InboxItem> {
fn dequeue_next(&self, agent: RuntimeAgentKey) -> Option<InboxItem> {
self.queues
.lock()
.unwrap()
.entry(agent_id)
.entry(agent)
.or_default()
.pop_front()
}
fn snapshot(&self, agent_id: AgentId) -> AgentInboxSnapshot {
fn snapshot(&self, agent: RuntimeAgentKey) -> AgentInboxSnapshot {
let queues = self.queues.lock().unwrap();
let items = queues
.get(&agent_id)
.get(&agent)
.map(|queue| queue.iter().cloned().collect::<Vec<_>>())
.unwrap_or_default();
AgentInboxSnapshot {
agent_id,
agent_id: agent.agent_id,
runtime_key: agent,
depth: items.len(),
items,
}
@ -139,14 +145,14 @@ impl AgentInbox for FakeInbox {
#[derive(Default)]
struct SharedTurnState {
busy: Mutex<HashMap<AgentId, AgentBusyState>>,
tickets: Mutex<HashMap<AgentId, VecDeque<Ticket>>>,
busy: Mutex<HashMap<RuntimeAgentKey, AgentBusyState>>,
tickets: Mutex<HashMap<RuntimeAgentKey, VecDeque<Ticket>>>,
}
impl SharedTurnState {
fn force_busy(&self, agent_id: AgentId, ticket_id: TicketId) {
self.busy.lock().unwrap().insert(
agent_id,
runtime_key(agent_id),
AgentBusyState::Busy {
ticket: ticket_id,
since_ms: 1,
@ -158,20 +164,20 @@ impl SharedTurnState {
self.tickets
.lock()
.unwrap()
.get(&agent_id)
.get(&runtime_key(agent_id))
.map(VecDeque::len)
.unwrap_or_default()
}
}
impl InputMediator for SharedTurnState {
fn enqueue(&self, agent_id: AgentId, ticket: Ticket) -> PendingReply {
self.enqueue_silent(agent_id, ticket)
fn enqueue(&self, agent: RuntimeAgentKey, ticket: Ticket) -> PendingReply {
self.enqueue_silent(agent, ticket)
}
fn enqueue_silent(&self, agent_id: AgentId, ticket: Ticket) -> PendingReply {
fn enqueue_silent(&self, agent: RuntimeAgentKey, ticket: Ticket) -> PendingReply {
self.busy.lock().unwrap().insert(
agent_id,
agent,
AgentBusyState::Busy {
ticket: ticket.id,
since_ms: 1,
@ -180,43 +186,43 @@ impl InputMediator for SharedTurnState {
self.tickets
.lock()
.unwrap()
.entry(agent_id)
.entry(agent)
.or_default()
.push_back(ticket);
PendingReply::new(Box::pin(async { Err(MailboxError::Cancelled) }))
}
fn preempt(&self, _agent: AgentId) {}
fn preempt(&self, _agent: RuntimeAgentKey) {}
fn mark_idle(&self, agent_id: AgentId) {
fn mark_idle(&self, agent: RuntimeAgentKey) {
self.busy
.lock()
.unwrap()
.insert(agent_id, AgentBusyState::Idle);
.insert(agent, AgentBusyState::Idle);
}
fn busy_state(&self, agent_id: AgentId) -> AgentBusyState {
fn busy_state(&self, agent: RuntimeAgentKey) -> AgentBusyState {
self.busy
.lock()
.unwrap()
.get(&agent_id)
.get(&agent)
.copied()
.unwrap_or(AgentBusyState::Idle)
}
}
impl AgentMailbox for SharedTurnState {
fn enqueue(&self, agent_id: AgentId, ticket: Ticket) -> PendingReply {
<Self as InputMediator>::enqueue(self, agent_id, ticket)
fn enqueue(&self, agent: RuntimeAgentKey, ticket: Ticket) -> PendingReply {
<Self as InputMediator>::enqueue(self, agent, ticket)
}
fn resolve(&self, _agent: AgentId, _result: String) -> Result<(), MailboxError> {
fn resolve(&self, _agent: RuntimeAgentKey, _result: String) -> Result<(), MailboxError> {
Ok(())
}
fn cancel_head(&self, agent_id: AgentId, ticket_id: TicketId) {
fn cancel_head(&self, agent: RuntimeAgentKey, ticket_id: TicketId) {
let mut tickets = self.tickets.lock().unwrap();
if let Some(queue) = tickets.get_mut(&agent_id) {
if let Some(queue) = tickets.get_mut(&agent) {
if queue.front().is_some_and(|ticket| ticket.id == ticket_id) {
queue.pop_front();
}
@ -395,7 +401,10 @@ async fn wake_if_idle_starts_turn_with_background_completion_prompt() {
let sessions = Arc::new(FakeSessionProvider::with_session(session.clone()));
tasks.insert(completed_task(&project, owner, task_id, "build finished"));
inbox
.enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
.enqueue_message(
runtime_key(owner),
completion_item(owner, task_id, ticket(20)),
)
.unwrap();
service(inbox, turns, tasks, sessions)
@ -427,7 +436,10 @@ async fn owner_busy_does_not_start_concurrent_wake_and_keeps_item_queued() {
let sessions = Arc::new(FakeSessionProvider::with_session(session.clone()));
tasks.insert(completed_task(&project, owner, task_id, "done"));
inbox
.enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
.enqueue_message(
runtime_key(owner),
completion_item(owner, task_id, ticket(20)),
)
.unwrap();
turns.force_busy(owner, ticket(99));
@ -442,7 +454,7 @@ async fn owner_busy_does_not_start_concurrent_wake_and_keeps_item_queued() {
assert_eq!(err, WakeError::AgentBusy { agent_id: owner });
assert!(session.prompts.lock().unwrap().is_empty());
assert_eq!(inbox.snapshot(owner).depth, 1);
assert_eq!(inbox.snapshot(runtime_key(owner)).depth, 1);
}
#[tokio::test]
@ -456,7 +468,10 @@ async fn absent_session_is_launched_or_reattached_by_provider() {
let sessions = Arc::new(FakeSessionProvider::default());
tasks.insert(completed_task(&project, owner, task_id, "done"));
inbox
.enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
.enqueue_message(
runtime_key(owner),
completion_item(owner, task_id, ticket(20)),
)
.unwrap();
service(inbox, turns, tasks, sessions.clone())
@ -485,7 +500,10 @@ async fn completion_is_marked_delivered_after_successful_wake() {
)));
tasks.insert(completed_task(&project, owner, task_id, "done"));
inbox
.enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
.enqueue_message(
runtime_key(owner),
completion_item(owner, task_id, ticket(20)),
)
.unwrap();
service(inbox, turns, tasks.clone(), sessions)
@ -512,7 +530,10 @@ async fn completion_is_marked_delivered_once_send_is_accepted_even_if_drain_fail
let sessions = Arc::new(FakeSessionProvider::with_session(session.clone()));
tasks.insert(completed_task(&project, owner, task_id, "done"));
inbox
.enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
.enqueue_message(
runtime_key(owner),
completion_item(owner, task_id, ticket(20)),
)
.unwrap();
let err = service(inbox, turns, tasks.clone(), sessions)
@ -543,10 +564,16 @@ async fn wake_drains_exactly_one_item_per_turn() {
tasks.insert(completed_task(&project, owner, first, "first"));
tasks.insert(completed_task(&project, owner, second, "second"));
inbox
.enqueue_message(owner, completion_item(owner, first, ticket(20)))
.enqueue_message(
runtime_key(owner),
completion_item(owner, first, ticket(20)),
)
.unwrap();
inbox
.enqueue_message(owner, completion_item(owner, second, ticket(21)))
.enqueue_message(
runtime_key(owner),
completion_item(owner, second, ticket(21)),
)
.unwrap();
service(inbox.clone(), turns.clone(), tasks, sessions)
@ -559,6 +586,6 @@ async fn wake_drains_exactly_one_item_per_turn() {
.unwrap();
assert_eq!(session.prompts.lock().unwrap().len(), 1);
assert_eq!(inbox.snapshot(owner).depth, 1);
assert_eq!(inbox.snapshot(runtime_key(owner)).depth, 1);
assert_eq!(turns.ticket_depth(owner), 0);
}