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

orch/engine/subscribers/
dispatch.rs
