Skip to main content

omni_dev/daemon/
registry.rs

1//! The [`ServiceRegistry`]: the daemon's set of hosted services, plus routing.
2
3use std::sync::Arc;
4
5use anyhow::{anyhow, Result};
6use serde_json::Value;
7
8use crate::github_rate_limit::{RateLimitCache, RateLimitSnapshot};
9
10use super::service::{DaemonService, ServiceStatus, ServiceStream};
11
12/// Holds the daemon's registered services and routes control-socket envelopes
13/// to them by [`name`](DaemonService::name).
14#[derive(Clone, Default)]
15pub struct ServiceRegistry {
16    services: Vec<Arc<dyn DaemonService>>,
17    /// The GitHub API rate-limit monitor's cache (#1375), shared with the poller
18    /// that fills it (hosted by the worktrees service). `None` until the daemon
19    /// wires it in via [`set_github_rate_limit`](Self::set_github_rate_limit), so
20    /// the built-in `status` op can report machine-wide GitHub budget usage as a
21    /// top-level field without downcasting a specific service.
22    github_rate_limit: Option<Arc<RateLimitCache>>,
23}
24
25impl ServiceRegistry {
26    /// Creates an empty registry.
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    /// Adds a service. Lookups match by [`DaemonService::name`] and the first
32    /// registration wins; a second service sharing a name would be dead code
33    /// for routing and would double-count in status/menu iteration, so it is
34    /// rejected with a warning rather than silently kept.
35    pub fn register(&mut self, service: Arc<dyn DaemonService>) {
36        let name = service.name();
37        if self.services.iter().any(|s| s.name() == name) {
38            tracing::warn!("ignoring duplicate registration of daemon service `{name}`");
39            return;
40        }
41        self.services.push(service);
42    }
43
44    /// Returns the registered service with the given name, if any.
45    pub fn get(&self, name: &str) -> Option<&Arc<dyn DaemonService>> {
46        self.services.iter().find(|s| s.name() == name)
47    }
48
49    /// All registered services, in registration order.
50    pub fn services(&self) -> &[Arc<dyn DaemonService>] {
51        &self.services
52    }
53
54    /// Wires in the GitHub rate-limit monitor's cache (#1375), so the built-in
55    /// `status` op can surface budget usage. Called once at daemon start with the
56    /// same `Arc` the worktrees service's poller writes.
57    pub fn set_github_rate_limit(&mut self, cache: Arc<RateLimitCache>) {
58        self.github_rate_limit = Some(cache);
59    }
60
61    /// The latest GitHub rate-limit snapshot, or `None` when the monitor is
62    /// unwired or has not polled successfully yet.
63    #[must_use]
64    pub fn github_rate_limit(&self) -> Option<RateLimitSnapshot> {
65        self.github_rate_limit.as_ref().and_then(|c| c.get())
66    }
67
68    /// Routes an operation to the named service, erroring if no such service is
69    /// registered.
70    pub async fn dispatch(&self, service: &str, op: &str, payload: Value) -> Result<Value> {
71        let svc = self
72            .get(service)
73            .ok_or_else(|| anyhow!("unknown service: {service}"))?;
74        svc.handle(op, payload).await
75    }
76
77    /// Opens a push subscription on the named service for a streaming `op`, or
78    /// `None` when the service is unknown or does not stream that op — in which
79    /// case the caller falls back to the normal [`dispatch`](Self::dispatch)
80    /// request→reply path (#1267).
81    pub fn subscribe(
82        &self,
83        service: &str,
84        op: &str,
85        payload: &Value,
86    ) -> Option<Box<dyn ServiceStream>> {
87        self.get(service)?.subscribe(op, payload)
88    }
89
90    /// Collects status from every service, in registration order.
91    pub async fn statuses(&self) -> Vec<ServiceStatus> {
92        let mut out = Vec::with_capacity(self.services.len());
93        for svc in &self.services {
94            out.push(svc.status().await);
95        }
96        out
97    }
98
99    /// Gracefully shuts down every service, in registration order.
100    pub async fn shutdown_all(&self) {
101        for svc in &self.services {
102            svc.shutdown().await;
103        }
104    }
105}
106
107#[cfg(test)]
108#[allow(clippy::unwrap_used, clippy::expect_used)]
109mod tests {
110    use super::*;
111    use crate::daemon::services::echo::EchoService;
112    use serde_json::json;
113
114    #[tokio::test]
115    async fn routes_known_service_and_rejects_unknown() {
116        let mut registry = ServiceRegistry::new();
117        assert!(registry.services().is_empty());
118        registry.register(Arc::new(EchoService));
119
120        assert!(registry.get("echo").is_some());
121        assert!(registry.get("missing").is_none());
122
123        // Routed op reaches the service; an unknown service is an error.
124        assert_eq!(
125            registry
126                .dispatch("echo", "echo", json!({ "x": 1 }))
127                .await
128                .unwrap(),
129            json!({ "x": 1 })
130        );
131        let err = registry
132            .dispatch("missing", "echo", Value::Null)
133            .await
134            .unwrap_err();
135        assert!(err.to_string().contains("unknown service"));
136
137        // Aggregation iterates every registered service.
138        assert_eq!(registry.statuses().await.len(), 1);
139        registry.shutdown_all().await;
140    }
141
142    #[test]
143    fn duplicate_registration_is_ignored() {
144        let mut registry = ServiceRegistry::new();
145        registry.register(Arc::new(EchoService));
146        registry.register(Arc::new(EchoService));
147        // The second registration shares `echo`'s name, so it is dropped: the
148        // first wins on lookup and iteration is not double-counted.
149        assert_eq!(registry.services().len(), 1);
150    }
151}