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::{mpsc, RwLock, Semaphore};
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#[allow(clippy::too_many_arguments)]
41pub fn spawn(
42    mut rx: broadcast::Receiver<TaskEvent>,
43    backend: Arc<dyn ExternalBackend>,
44    tmux: Arc<TmuxManager>,
45    runner: Arc<TaskRunner>,
46    capture: Arc<CaptureService>,
47    semaphore: Arc<Semaphore>,
48    task_manager: Arc<TaskManager>,
49    weight_tx: mpsc::Sender<WeightSignal>,
50    router_arc: Arc<RwLock<Router>>,
51    dispatching: Arc<DashMap<String, String>>,
52    store: Arc<TaskStore>,
53    repo: String,
54) {
55    tokio::spawn(async move {
56        loop {
57            match rx.recv().await {
58                Ok(event) if event.new_status == "routed" && event.repo == repo => {
59                    tracing::info!(task_id = event.task_id, "event-driven dispatch triggered");
60                    // Delegate to existing dispatch logic.
61                    // The dispatching set guard inside tick_dispatch_tasks
62                    // prevents double-dispatch if the tick loop also picks this up.
63                    let (dispatch_mode, session_map) =
64                        prepare_dispatch_context(&router_arc, tmux.batch_session_active()).await;
65                    if let Err(e) = crate::engine::tick::tick_dispatch_tasks(
66                        &backend,
67                        &tmux,
68                        &repo,
69                        &runner,
70                        &capture,
71                        &semaphore,
72                        &task_manager,
73                        &weight_tx,
74                        dispatch_mode,
75                        &dispatching,
76                        &store,
77                        &session_map,
78                    )
79                    .await
80                    {
81                        tracing::warn!(?e, "event-driven dispatch failed (tick will retry)");
82                    }
83                }
84                Ok(_) => {} // Not a routed event or different repo
85                Err(broadcast::error::RecvError::Lagged(n)) => {
86                    tracing::debug!(n, "dispatch subscriber lagged, tick will catch up");
87                }
88                Err(_) => break, // Channel closed
89            }
90        }
91    });
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::engine::router::{Router, RouterConfig};
98    use std::time::Duration;
99    use tokio::sync::{oneshot, Notify};
100
101    #[tokio::test]
102    async fn prepare_dispatch_context_drops_read_lock_before_awaiting_session_map() {
103        let router_arc = Arc::new(RwLock::new(Router::new(RouterConfig::default())));
104        let (tx, rx) = oneshot::channel::<SessionMap>();
105        let waiting = Arc::new(Notify::new());
106        let waiting_for_task = Arc::clone(&waiting);
107        let router_for_task = Arc::clone(&router_arc);
108
109        let task = tokio::spawn(async move {
110            prepare_dispatch_context(&router_for_task, async move {
111                waiting_for_task.notify_one();
112                rx.await.unwrap_or_default()
113            })
114            .await
115        });
116
117        // Wait until the function has entered its awaited session-map future.
118        waiting.notified().await;
119
120        // If the read guard leaked across the await above, this write lock would block.
121        let write_guard = tokio::time::timeout(Duration::from_secs(2), router_arc.write())
122            .await
123            .expect("router write lock blocked while dispatch subscriber awaited session map");
124        drop(write_guard);
125
126        let _ = tx.send(SessionMap::new());
127        let _ = task.await.expect("prepare_dispatch_context task panicked");
128    }
129}
orch/engine/subscribers/dispatch.rs

orch/engine/subscribers/
dispatch.rs
