orch/engine/subscribers/
dispatch.rs

1//! Reacts to Routed events — dispatches agent immediately.
2
3use crate::backends::ExternalBackend;
4use crate::channels::capture::CaptureService;
5use crate::engine::events::TaskEvent;
6use crate::engine::router::Router;
7use crate::engine::runner::{TaskRunner, WeightSignal};
8use crate::engine::tasks::TaskManager;
9use crate::store::TaskStore;
10use crate::tmux::TmuxManager;
11use dashmap::DashMap;
12use std::sync::Arc;
13use tokio::sync::broadcast;
14use tokio::sync::{RwLock, Semaphore, mpsc};
15
16type SessionMap = std::collections::HashMap<String, bool>;
17
18async fn prepare_dispatch_context<F>(
19    router_arc: &Arc<RwLock<Router>>,
20    session_map_fut: F,
21) -> (crate::engine::tick::DispatchMode, SessionMap)
22where
23    F: std::future::Future<Output = SessionMap>,
24{
25    // Keep lock scope minimal: snapshot router state, then drop the guard
26    // before awaiting session/dispatch work.
27    let dispatch_mode = {
28        let router_guard = router_arc.read().await;
29        crate::engine::tick::dispatch_mode_from_router(&router_guard)
30    };
31    let session_map = session_map_fut.await;
32    (dispatch_mode, session_map)
33}
34
35/// Spawn a task that listens for Routed events and dispatches agents.
36///
37/// This mirrors the logic in `tick_dispatch_tasks` but triggers instantly
38/// instead of waiting for the next tick. The dispatching set prevents
39/// double-dispatch if the tick loop picks up the same task.
40#[expect(
41    clippy::too_many_arguments,
42    reason = "the arguments have no natural grouping"
43)]
44pub fn spawn(
45    mut rx: broadcast::Receiver<TaskEvent>,
46    backend: Arc<dyn ExternalBackend>,
47    tmux: Arc<TmuxManager>,
48    runner: Arc<TaskRunner>,
49    capture: Arc<CaptureService>,
50    semaphore: Arc<Semaphore>,
51    task_manager: Arc<TaskManager>,
52    weight_tx: mpsc::Sender<WeightSignal>,
53    router_arc: Arc<RwLock<Router>>,
54    dispatching: Arc<DashMap<String, String>>,
55    store: Arc<TaskStore>,
56    repo: String,
57) {
58    tokio::spawn(async move {
59        loop {
60            match rx.recv().await {
61                Ok(event) if event.new_status == "routed" && event.repo == repo => {
62                    tracing::info!(task_id = event.task_id, "event-driven dispatch triggered");
63                    // Delegate to existing dispatch logic.
64                    // The dispatching set guard inside tick_dispatch_tasks
65                    // prevents double-dispatch if the tick loop also picks this up.
66                    let (dispatch_mode, session_map) =
67                        prepare_dispatch_context(&router_arc, tmux.batch_session_active()).await;
68                    if let Err(e) = crate::engine::tick::tick_dispatch_tasks(
69                        &backend,
70                        &tmux,
71                        &repo,
72                        &runner,
73                        &capture,
74                        &semaphore,
75                        &task_manager,
76                        &weight_tx,
77                        dispatch_mode,
78                        &dispatching,
79                        &store,
80                        &session_map,
81                    )
82                    .await
83                    {
84                        tracing::warn!(?e, "event-driven dispatch failed (tick will retry)");
85                    }
86                }
87                Ok(_) => {} // Not a routed event or different repo
88                Err(broadcast::error::RecvError::Lagged(n)) => {
89                    tracing::debug!(n, "dispatch subscriber lagged, tick will catch up");
90                }
91                Err(_) => break, // Channel closed
92            }
93        }
94    });
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use crate::engine::router::{Router, RouterConfig};
101    use std::time::Duration;
102    use tokio::sync::{Notify, oneshot};
103
104    #[tokio::test]
105    async fn prepare_dispatch_context_drops_read_lock_before_awaiting_session_map() {
106        let router_arc = Arc::new(RwLock::new(Router::new(RouterConfig::default())));
107        let (tx, rx) = oneshot::channel::<SessionMap>();
108        let waiting = Arc::new(Notify::new());
109        let waiting_for_task = Arc::clone(&waiting);
110        let router_for_task = Arc::clone(&router_arc);
111
112        let task = tokio::spawn(async move {
113            prepare_dispatch_context(&router_for_task, async move {
114                waiting_for_task.notify_one();
115                rx.await.unwrap_or_default()
116            })
117            .await
118        });
119
120        // Wait until the function has entered its awaited session-map future.
121        waiting.notified().await;
122
123        // If the read guard leaked across the await above, this write lock would block.
124        let write_guard = tokio::time::timeout(Duration::from_secs(2), router_arc.write())
125            .await
126            .expect("router write lock blocked while dispatch subscriber awaited session map");
127        drop(write_guard);
128
129        let _ = tx.send(SessionMap::new());
130        let _ = task.await.expect("prepare_dispatch_context task panicked");
131    }
132}
orch/engine/subscribers/dispatch.rs

orch/engine/subscribers/
dispatch.rs
