Skip to main content

aion_server/worker/
registry.rs

1//! Connected-worker registry keyed by worker-pool address and activity type.
2
3use std::collections::{BTreeMap, BTreeSet, HashMap};
4use std::sync::{Arc, Mutex, MutexGuard};
5use std::time::{Duration, Instant};
6
7use aion_core::{ClusterEvent, InterventionCapabilities, WorkerDeathReason, WorkerTransport};
8use aion_proto::{ProtoActivityTask, ProtoRegisterWorker};
9use aion_store::{NamespaceOrigin, NamespacePlacement, NamespaceStore};
10use tokio::sync::{Notify, mpsc};
11
12use crate::cluster_publisher::ClusterEventPublisher;
13use crate::config::AutoCreate;
14use crate::error::ServerError;
15use crate::namespace::{CallerIdentity, NamespaceGuard, NamespaceMinter, NamespaceOperation};
16use crate::observability::Metrics;
17
18/// The literal task queue an empty/absent selector normalizes to.
19///
20/// A worker-pool address has two disjoint dimensions; the second one
21/// (`task_queue`) is a liveness selector, not a correctness boundary. An empty
22/// `task_queue` is normalized to this one named default pool so a producer that
23/// names no queue and a worker that advertises none both land on the same pool.
24///
25/// Re-exported from [`aion_core::DEFAULT_TASK_QUEUE`] so the server cannot drift
26/// from the canonical domain default; the name is kept stable here for existing
27/// call sites.
28pub use aion_core::DEFAULT_TASK_QUEUE;
29
30/// Server-side handle used to push activity tasks to a connected worker stream.
31pub type WorkerTaskSender = mpsc::Sender<WorkerMessage>;
32
33/// Transport through which the server delivers a dispatch to a registered worker.
34///
35/// A worker is selected the SAME way regardless of transport (`select_worker`
36/// over the `(namespace, task_queue, node)` pool key); only the delivery leg
37/// differs. The default gRPC path pushes a [`WorkerMessage`] onto the worker's
38/// stream `mpsc` ([`WorkerDelivery::Grpc`]); a liminal-connected worker is
39/// delivered to by pushing the dispatch out on its existing liminal connection
40/// ([`WorkerDelivery::Liminal`], feature-gated). This enum is the minimal
41/// transport-agnostic seam: the registry holds it on each [`WorkerHandle`], and
42/// the dispatch path reads the variant it needs. The gRPC variant carries exactly
43/// the `mpsc::Sender` it always did, so the gRPC dispatch path is unchanged.
44#[derive(Clone, Debug)]
45pub enum WorkerDelivery {
46    /// gRPC stream delivery: the dispatch path pushes a [`WorkerMessage`] onto
47    /// this `mpsc` sender, exactly as before this enum existed.
48    Grpc(WorkerTaskSender),
49    /// Liminal server-push delivery: the dispatch path pushes the serialized
50    /// dispatch out on the worker's existing liminal connection and awaits the
51    /// correlated reply. Carries the connection identity needed to address that
52    /// push.
53    #[cfg(feature = "liminal-transport")]
54    Liminal(crate::worker::liminal_transport::LiminalWorkerDelivery),
55}
56
57/// Message queued from server-side dispatch/shutdown into a worker stream writer.
58#[derive(Clone, Debug, Eq, PartialEq)]
59pub enum WorkerMessage {
60    /// Activity invocation pushed to a worker.
61    ActivityTask(Box<ProtoActivityTask>),
62    /// Graceful-shutdown notification; no new work will be dispatched.
63    DrainRequest,
64}
65
66/// Address of a worker pool: the two disjoint routing dimensions that select a
67/// pool, before an `activity_type` is matched within it.
68///
69/// `namespace` is the correctness/isolation boundary — a workflow's activities
70/// only ever reach workers in the workflow's namespace, so crossing it is a bug.
71/// `task_queue` is the pool/flavour selector within that namespace (norn /
72/// claude / cpu / gpu) — a miss is a liveness issue, never a correctness one.
73///
74/// This is a named type rather than a `(String, String)` tuple so a `node`
75/// dimension (Tier 3 affinity) can be added later without re-threading every
76/// call site.
77#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
78pub struct PoolAddress {
79    namespace: String,
80    task_queue: String,
81}
82
83impl PoolAddress {
84    /// Build a pool address, normalizing an empty `task_queue` to the named
85    /// [`DEFAULT_TASK_QUEUE`] pool. The `namespace` is the authorization
86    /// boundary and is never normalized.
87    #[must_use]
88    pub fn new(namespace: impl Into<String>, task_queue: impl Into<String>) -> Self {
89        let task_queue = task_queue.into();
90        let task_queue = if task_queue.is_empty() {
91            String::from(DEFAULT_TASK_QUEUE)
92        } else {
93            task_queue
94        };
95        Self {
96            namespace: namespace.into(),
97            task_queue,
98        }
99    }
100
101    /// The correctness/isolation boundary of this pool.
102    #[must_use]
103    pub fn namespace(&self) -> &str {
104        &self.namespace
105    }
106
107    /// The pool/flavour selector within the namespace.
108    #[must_use]
109    pub fn task_queue(&self) -> &str {
110        &self.task_queue
111    }
112}
113
114/// Registry match key: a worker-pool address plus the activity type matched
115/// within that pool. A named type (not an anonymous tuple) so the routing
116/// identity stays self-describing and extensible.
117#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
118struct ActivityKey {
119    pool: PoolAddress,
120    activity_type: String,
121}
122
123impl ActivityKey {
124    fn new(pool: PoolAddress, activity_type: impl Into<String>) -> Self {
125        Self {
126            pool,
127            activity_type: activity_type.into(),
128        }
129    }
130}
131
132type WorkerMap = HashMap<WorkerId, WorkerHandle>;
133type RegistryMap = HashMap<ActivityKey, WorkerMap>;
134
135/// Stable identifier assigned to a connected worker stream.
136#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
137pub struct WorkerId(u64);
138
139impl WorkerId {
140    /// Build a worker id from the numeric value exposed by administrative and
141    /// wire surfaces.
142    #[must_use]
143    pub const fn from_value(value: u64) -> Self {
144        Self(value)
145    }
146
147    /// Raw numeric value, as carried by the wire `RegisterAck.worker_id` so
148    /// workers can correlate their logs with the server's.
149    #[must_use]
150    pub const fn value(self) -> u64 {
151        self.0
152    }
153}
154
155/// Cloneable handle for a registered worker stream.
156///
157/// A worker serves a SET of namespaces under a single `task_queue`, so it is
158/// indexed under one `(namespace, task_queue, activity_type)` key per namespace
159/// in its set. `node` is an OPTIONAL locality affinity (a locality, not a
160/// process — many handles may share a node id) used as a within-pool filter at
161/// selection time; `None` means the worker advertised no locality.
162#[derive(Clone, Debug)]
163pub struct WorkerHandle {
164    id: WorkerId,
165    namespaces: BTreeSet<String>,
166    task_queue: String,
167    node: Option<String>,
168    activity_types: BTreeSet<String>,
169    delivery: WorkerDelivery,
170    /// The neutral mid-run intervention primitives this worker's harness advertises
171    /// support for (NOI-6). The server gates every intervention command on THIS set
172    /// and NEVER routes an unadvertised primitive. Empty = observability-only (the
173    /// default for every non-agent worker), so a normal activity worker advertises
174    /// no controls and the intervention router refuses every command for it.
175    intervention_capabilities: InterventionCapabilities,
176}
177
178impl WorkerHandle {
179    /// Worker identifier assigned by this server process.
180    #[must_use]
181    pub const fn id(&self) -> WorkerId {
182        self.id
183    }
184
185    /// Namespaces authorized for this worker stream. The worker is reachable for
186    /// a dispatch only when its set includes the workflow's namespace.
187    #[must_use]
188    pub const fn namespaces(&self) -> &BTreeSet<String> {
189        &self.namespaces
190    }
191
192    /// Task queue (pool/flavour) this worker serves within each namespace.
193    #[must_use]
194    pub fn task_queue(&self) -> &str {
195        &self.task_queue
196    }
197
198    /// Optional locality affinity this worker advertised. `None` means the
199    /// worker carries no node and is reachable only for unpinned dispatches.
200    #[must_use]
201    pub fn node(&self) -> Option<&str> {
202        self.node.as_deref()
203    }
204
205    /// Activity types advertised by this worker.
206    #[must_use]
207    pub fn activity_types(&self) -> &BTreeSet<String> {
208        &self.activity_types
209    }
210
211    /// The transport this worker is delivered to through.
212    #[must_use]
213    pub const fn delivery(&self) -> &WorkerDelivery {
214        &self.delivery
215    }
216
217    /// The neutral intervention primitives this worker's harness advertises (NOI-6).
218    ///
219    /// The intervention router gates on this set and never routes an unadvertised
220    /// primitive. Empty (the default for a plain activity worker) means the worker
221    /// is observability-only: the router refuses every intervention command for it.
222    #[must_use]
223    pub const fn intervention_capabilities(&self) -> &InterventionCapabilities {
224        &self.intervention_capabilities
225    }
226
227    /// gRPC stream sender used by the gRPC dispatch path to push work, or `None`
228    /// when this worker is delivered to over a non-gRPC transport (liminal).
229    ///
230    /// The gRPC dispatch path registers every worker with a [`WorkerDelivery::Grpc`]
231    /// delivery, so this is always `Some` for a gRPC-registered worker — the
232    /// behaviour the path relied on before delivery became transport-agnostic.
233    #[must_use]
234    pub fn sender(&self) -> Option<&WorkerTaskSender> {
235        match &self.delivery {
236            WorkerDelivery::Grpc(sender) => Some(sender),
237            #[cfg(feature = "liminal-transport")]
238            WorkerDelivery::Liminal(_) => None,
239        }
240    }
241}
242
243#[derive(Debug)]
244struct RegistryState {
245    next_worker_id: u64,
246    workers: BTreeMap<WorkerId, WorkerHandle>,
247    by_activity: RegistryMap,
248    /// Round-robin cursor per `(namespace, task_queue, activity_type)` triple, so
249    /// each pool rotates independently of every other pool.
250    rotation: HashMap<ActivityKey, usize>,
251    /// When a worker advertising a given `(pool, activity_type)` and node last
252    /// LEFT service, per advertised node. Read by [`ConnectedWorkerRegistry::pool_census`]
253    /// to answer "how old is the last compatible poller" for an address that has
254    /// none right now (R1). Node-keyed because a node-pinned dispatch's notion
255    /// of "compatible" is node-specific: a worker still serving the activity on
256    /// another node must never make a pinned address look freshly served.
257    last_departure: HashMap<ActivityKey, BTreeMap<Option<String>, Instant>>,
258    /// Workers the server currently cannot reach on the server-to-worker push
259    /// leg, and which are therefore excluded from selection while they remain
260    /// so.
261    ///
262    /// Held here rather than derived at selection time because selection must
263    /// not take the heartbeat tracker's lock: the liveness probe owns the
264    /// evidence and publishes the verdict, and selection only reads it.
265    ///
266    /// This is exclusion from DISPATCH, not deregistration. An unreachable
267    /// worker stays registered, keeps its in-flight work, and becomes eligible
268    /// again the moment a ping is answered — nothing about it is torn down on
269    /// the strength of a push failure.
270    dispatch_ineligible: BTreeSet<WorkerId>,
271}
272
273impl Default for RegistryState {
274    fn default() -> Self {
275        Self {
276            next_worker_id: 1,
277            workers: BTreeMap::new(),
278            by_activity: HashMap::new(),
279            rotation: HashMap::new(),
280            last_departure: HashMap::new(),
281            dispatch_ineligible: BTreeSet::new(),
282        }
283    }
284}
285
286/// Cloneable registry of currently connected worker streams.
287#[derive(Clone, Debug)]
288pub struct ConnectedWorkerRegistry {
289    inner: Arc<Mutex<RegistryState>>,
290    metrics: Option<Metrics>,
291    /// WS3 cluster-event publisher: emits `WorkerConnected`/`WorkerDisconnected`
292    /// topology deltas on register/deregister. `None` keeps existing
293    /// constructions (and every test) silent, exactly like `metrics`.
294    cluster_publisher: Option<ClusterEventPublisher>,
295    /// Minted-on-use hook (Control-Plane Phase 1). `None` disables minting, so
296    /// registration is byte-identical to before the registry existed; `Some`
297    /// durably records (open) or gates (closed) each authorized namespace.
298    minter: Option<NamespaceMinter>,
299    worker_arrived: Arc<Notify>,
300}
301
302impl Default for ConnectedWorkerRegistry {
303    fn default() -> Self {
304        Self {
305            inner: Arc::new(Mutex::new(RegistryState::default())),
306            metrics: None,
307            cluster_publisher: None,
308            minter: None,
309            worker_arrived: Arc::new(Notify::new()),
310        }
311    }
312}
313
314impl ConnectedWorkerRegistry {
315    /// Build a registry that records connected-worker gauge updates.
316    #[must_use]
317    pub fn with_metrics(metrics: Metrics) -> Self {
318        Self {
319            inner: Arc::new(Mutex::new(RegistryState::default())),
320            metrics: Some(metrics),
321            cluster_publisher: None,
322            minter: None,
323            worker_arrived: Arc::new(Notify::new()),
324        }
325    }
326
327    /// Attach the WS3 cluster-event publisher so worker topology changes are
328    /// pushed to the dashboard. Pure builder addition.
329    #[must_use]
330    pub fn with_cluster_publisher(mut self, publisher: ClusterEventPublisher) -> Self {
331        self.cluster_publisher = Some(publisher);
332        self
333    }
334
335    /// Install the minted-on-use namespace hook (Control-Plane Phase 1).
336    ///
337    /// After a registration is authorized and its namespace set scoped, each
338    /// authorized namespace is durably recorded ([`AutoCreate::Open`]) or gated
339    /// ([`AutoCreate::Closed`]) through `store`. Without this builder the
340    /// registry never touches the namespace registry, so registration stays
341    /// byte-identical to before the registry existed. Pure builder addition,
342    /// mirroring [`Self::with_cluster_publisher`].
343    ///
344    /// When a cluster publisher has already been attached
345    /// ([`Self::with_cluster_publisher`], called first on the boot path), it is
346    /// threaded into the minter so a first worker-mint emits the live
347    /// `namespace created` delta to the ops console (S8). Order-independence is
348    /// not assumed: callers wire the publisher before minting on the boot path.
349    #[must_use]
350    pub fn with_namespace_minting(
351        mut self,
352        store: Arc<dyn NamespaceStore>,
353        policy: AutoCreate,
354    ) -> Self {
355        let minter = NamespaceMinter::new(store, policy);
356        let minter = match &self.cluster_publisher {
357            Some(publisher) => minter.with_cluster_publisher(publisher.clone()),
358            None => minter,
359        };
360        self.minter = Some(minter);
361        self
362    }
363
364    /// Thread the boot's namespace-mint routing context into the registry's
365    /// minter, so a worker registering for a namespace whose registry shard this
366    /// node does not own mints through the shard's owner instead of being
367    /// refused `NotOwner` forever.
368    ///
369    /// The SECOND of the two minter construction sites (the first is
370    /// [`ServerState::namespace_minter`](crate::ServerState::namespace_minter),
371    /// which serves the gRPC and HTTP start seams). Called after
372    /// [`Self::with_namespace_minting`] on the boot path; a no-op when no minter
373    /// is installed, and never called at all off-cluster, so default/test
374    /// registries stay byte-identical.
375    #[must_use]
376    pub fn with_namespace_routing(mut self, routing: crate::namespace::NamespaceRouting) -> Self {
377        self.minter = self.minter.map(|minter| minter.with_routing(routing));
378        self
379    }
380
381    /// Authorize a worker registration and insert it into the connected-worker registry.
382    ///
383    /// # Errors
384    ///
385    /// Returns [`ServerError`] if namespace authorization fails or the registry lock is poisoned.
386    pub async fn accept_registration(
387        &self,
388        guard: &NamespaceGuard,
389        caller: &CallerIdentity,
390        registration: &ProtoRegisterWorker,
391        sender: WorkerTaskSender,
392    ) -> Result<WorkerRegistration, ServerError> {
393        // Verify the operation against the guard's worker-registration policy,
394        // then authorize EACH namespace in the worker's set: a worker serves a
395        // SET of correctness boundaries, so the registration is denied unless
396        // the caller is granted every one. The wire's empty `node` carries no
397        // locality affinity; a non-empty value is the worker's advertised node.
398        guard
399            .scope(caller, &NamespaceOperation::register_worker(registration))
400            .await?;
401        let namespaces = guard.scope_worker_namespaces(caller, &registration.namespaces)?;
402        // MINT HOOK (Control-Plane Phase 1). This runs strictly AFTER the
403        // per-namespace authorization above (`scope` + `scope_worker_namespaces`),
404        // so it can only ever mint a namespace the caller is already authorized
405        // for — the mint is auth-scoped by construction (CVE-2025-14986: open
406        // minting and namespace isolation only coexist when minting is
407        // auth-gated). It runs BEFORE the worker is inserted, so a `closed`
408        // rejection never leaves a half-registered worker behind.
409        self.mint_or_gate_namespaces(&namespaces).await?;
410        let node = optional_node(&registration.node);
411        // PLACEMENT-ADMISSION GATE (Control-Plane Phase 2, P2-I1). Runs strictly
412        // AFTER the mint hook (so every authorized namespace has a durable record to
413        // read a placement from) and with BOTH the worker's advertised `node` and
414        // the full authorized namespace set in scope. It rejects the WHOLE
415        // registration (Open Decision 6) when the worker's node violates any
416        // `Pinned{L}` namespace it would serve, so only L-node workers ever enter a
417        // hard-pinned namespace's pool. Auth-scoped by construction (it only ever
418        // gates a namespace already authorized above); a no-op with no minter
419        // installed, so default/test registries stay byte-identical.
420        self.enforce_pinned_placement(&namespaces, node.as_deref())
421            .await?;
422        self.register_namespaces(
423            namespaces,
424            registration.task_queue.clone(),
425            node,
426            registration.activity_types.iter(),
427            sender,
428        )
429    }
430
431    /// Apply the minted-on-use policy to an already-authorized namespace set.
432    ///
433    /// A no-op when no minter is installed (every default/test registry), so
434    /// registration stays byte-identical. With a minter, the work is delegated
435    /// to the shared [`NamespaceMinter::mint_or_gate`] — the single
436    /// transport-agnostic implementation reused by the workflow-start safety net
437    /// — with [`NamespaceOrigin::WorkerMint`] so a first mint is attributed to
438    /// worker registration. See that method for the open/closed policy, the
439    /// idempotent "namespace created" event, and the retryable `NotOwner`
440    /// surface.
441    ///
442    /// # Errors
443    ///
444    /// Returns [`ServerError::StoreBackend`] if a durable upsert/lookup fails
445    /// (including a retryable `NotOwner` fence), or [`ServerError::Namespace`]
446    /// when `closed` rejects an unknown namespace.
447    async fn mint_or_gate_namespaces(&self, namespaces: &[String]) -> Result<(), ServerError> {
448        let Some(minter) = &self.minter else {
449            return Ok(());
450        };
451        minter
452            .mint_or_gate(namespaces, NamespaceOrigin::WorkerMint)
453            .await
454    }
455
456    /// Reject the whole registration when the worker's advertised `node` violates
457    /// any `Pinned{L}` namespace it would serve (Control-Plane Phase 2, P2-I1).
458    ///
459    /// For each authorized namespace whose placement is [`NamespacePlacement::Pinned`],
460    /// the worker's advertised `node` must be `Some(n)` with `n ∈ L`; a `None` node
461    /// or an `n ∉ L` is a loud, whole-registration rejection naming the namespace,
462    /// the node, and the required set. This guarantees only L-node workers ever
463    /// serve a hard-pinned namespace's pool, which is exactly what lets the
464    /// `Some(N ∉ L)` composition case (§2.2) resolve to the correct isolation stall
465    /// at dispatch rather than needing a start-time enumeration of future nodes.
466    ///
467    /// Non-`Pinned` placements ([`NamespacePlacement::Unplaced`]/[`NamespacePlacement::Prefer`])
468    /// are UNAFFECTED — byte-identical registration. A no-op when no minter is
469    /// installed (every default/test registry), so those stay behaviour-identical:
470    /// the gate reads placement from the SAME registry record the minter/placement
471    /// endpoint writes, never a second source of truth.
472    ///
473    /// # Errors
474    ///
475    /// Returns [`ServerError::Namespace`] (placement-admission denial) when the
476    /// worker's node violates a `Pinned` namespace, or [`ServerError::StoreBackend`]
477    /// if a placement read fails at the backend.
478    async fn enforce_pinned_placement(
479        &self,
480        namespaces: &[String],
481        node: Option<&str>,
482    ) -> Result<(), ServerError> {
483        let Some(minter) = &self.minter else {
484            return Ok(());
485        };
486        for namespace in namespaces {
487            let NamespacePlacement::Pinned { nodes } = minter.placement_of(namespace).await? else {
488                continue;
489            };
490            let admitted = node.is_some_and(|n| nodes.contains(n));
491            if !admitted {
492                return Err(ServerError::placement_admission_denied(
493                    namespace, node, &nodes,
494                ));
495            }
496        }
497        Ok(())
498    }
499
500    /// Insert an already-authorized worker stream into the default task queue of
501    /// a single `namespace`, with no node affinity.
502    ///
503    /// Convenience over [`Self::register_namespaces`] for callers that serve one
504    /// namespace and do not select a task queue (notably tests of the default
505    /// pool).
506    ///
507    /// # Errors
508    ///
509    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
510    pub fn register<'a>(
511        &self,
512        namespace: impl Into<String>,
513        activity_types: impl IntoIterator<Item = &'a String>,
514        sender: WorkerTaskSender,
515    ) -> Result<WorkerRegistration, ServerError> {
516        self.register_namespaces(
517            [namespace.into()],
518            String::from(DEFAULT_TASK_QUEUE),
519            None,
520            activity_types,
521            sender,
522        )
523    }
524
525    /// Insert an already-authorized worker stream into one explicit worker pool
526    /// (single namespace + task queue), with no node affinity.
527    ///
528    /// # Errors
529    ///
530    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
531    pub fn register_pool<'a>(
532        &self,
533        pool: PoolAddress,
534        activity_types: impl IntoIterator<Item = &'a String>,
535        sender: WorkerTaskSender,
536    ) -> Result<WorkerRegistration, ServerError> {
537        let PoolAddress {
538            namespace,
539            task_queue,
540        } = pool;
541        self.register_namespaces([namespace], task_queue, None, activity_types, sender)
542    }
543
544    /// Insert an already-authorized worker stream serving a SET of namespaces
545    /// under one `task_queue`, with an optional `node` locality affinity.
546    ///
547    /// The worker is indexed under one `(namespace, task_queue, activity_type)`
548    /// key per namespace in its set, so a dispatch in any of those namespaces
549    /// can reach it. `node` is recorded on the handle and used only as a
550    /// within-pool filter at selection time — it is NOT part of [`PoolAddress`].
551    ///
552    /// # Errors
553    ///
554    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
555    pub fn register_namespaces<'a>(
556        &self,
557        namespaces: impl IntoIterator<Item = String>,
558        task_queue: impl Into<String>,
559        node: Option<String>,
560        activity_types: impl IntoIterator<Item = &'a String>,
561        sender: WorkerTaskSender,
562    ) -> Result<WorkerRegistration, ServerError> {
563        self.register_delivery(
564            namespaces,
565            task_queue,
566            node,
567            activity_types,
568            WorkerDelivery::Grpc(sender),
569        )
570    }
571
572    /// Insert an already-authorized worker serving a SET of namespaces under one
573    /// `task_queue` and optional `node`, delivered to through an explicit
574    /// [`WorkerDelivery`] transport.
575    ///
576    /// This is the transport-agnostic registration core: [`Self::register_namespaces`]
577    /// is the gRPC façade over it (it wraps the stream sender in
578    /// [`WorkerDelivery::Grpc`]). Selection (`select_worker`/`workers_for`) is
579    /// identical across transports; only the held delivery differs.
580    ///
581    /// # Errors
582    ///
583    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
584    pub fn register_delivery<'a>(
585        &self,
586        namespaces: impl IntoIterator<Item = String>,
587        task_queue: impl Into<String>,
588        node: Option<String>,
589        activity_types: impl IntoIterator<Item = &'a String>,
590        delivery: WorkerDelivery,
591    ) -> Result<WorkerRegistration, ServerError> {
592        self.register_delivery_with_capabilities(
593            namespaces,
594            task_queue,
595            node,
596            activity_types,
597            delivery,
598            InterventionCapabilities::none(),
599        )
600    }
601
602    /// Insert an already-authorized worker exactly like [`Self::register_delivery`],
603    /// additionally recording the neutral [`InterventionCapabilities`] its harness
604    /// advertises (NOI-6).
605    ///
606    /// This is the capability-carrying registration core: [`Self::register_delivery`]
607    /// is the façade over it that advertises the empty set (observability-only), so
608    /// every existing caller stays byte-identical. Selection is unchanged — the
609    /// capability set is metadata the intervention router gates on, never a routing
610    /// dimension.
611    ///
612    /// # Errors
613    ///
614    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
615    pub fn register_delivery_with_capabilities<'a>(
616        &self,
617        namespaces: impl IntoIterator<Item = String>,
618        task_queue: impl Into<String>,
619        node: Option<String>,
620        activity_types: impl IntoIterator<Item = &'a String>,
621        delivery: WorkerDelivery,
622        intervention_capabilities: InterventionCapabilities,
623    ) -> Result<WorkerRegistration, ServerError> {
624        let namespaces = namespaces.into_iter().collect::<BTreeSet<_>>();
625        let task_queue = task_queue.into();
626        let activity_types = activity_types.into_iter().cloned().collect::<BTreeSet<_>>();
627        let mut state = self.state()?;
628        let worker_id = WorkerId(state.next_worker_id);
629        state.next_worker_id = state.next_worker_id.saturating_add(1);
630
631        // Capture the node affinity for the WS3 WorkerConnected delta before the
632        // handle moves it.
633        let node_for_event = node.clone();
634        let handle = WorkerHandle {
635            id: worker_id,
636            namespaces: namespaces.clone(),
637            task_queue: task_queue.clone(),
638            node,
639            activity_types: activity_types.clone(),
640            delivery,
641            intervention_capabilities,
642        };
643
644        for namespace in &namespaces {
645            let pool = PoolAddress::new(namespace.clone(), task_queue.clone());
646            for activity_type in &activity_types {
647                let key = ActivityKey::new(pool.clone(), activity_type.clone());
648                // This address is served again, so its departure record has no
649                // remaining meaning — drop it in lockstep with the insert. The
650                // departure map is therefore bounded by the addresses real
651                // workers have served and left, never by caller-supplied
652                // dispatch strings (which never create an entry at all).
653                if let Some(by_node) = state.last_departure.get_mut(&key) {
654                    by_node.remove(&handle.node);
655                    if by_node.is_empty() {
656                        state.last_departure.remove(&key);
657                    }
658                }
659                state
660                    .by_activity
661                    .entry(key)
662                    .or_default()
663                    .insert(worker_id, handle.clone());
664            }
665        }
666        let transport = transport_of(&handle.delivery);
667        state.workers.insert(worker_id, handle);
668        drop(state);
669
670        if let Some(metrics) = &self.metrics {
671            for namespace in &namespaces {
672                metrics.worker_connected(namespace);
673            }
674        }
675
676        // WS3: one WorkerConnected delta carrying the full namespace set (the
677        // event is namespace-list-valued; the deploy-scoped cluster channel sees
678        // it whole). Edge-triggered by the real insert, never a poll.
679        if let Some(publisher) = &self.cluster_publisher {
680            let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
681            let task_queue_owned = task_queue.clone();
682            drop(publisher.emit(|meta| ClusterEvent::WorkerConnected {
683                meta,
684                worker_id: worker_id.value().to_string(),
685                namespaces: namespaces_vec,
686                task_queue: task_queue_owned,
687                transport,
688                node: node_for_event,
689            }));
690        }
691
692        self.worker_arrived.notify_waiters();
693
694        Ok(WorkerRegistration {
695            registry: self.clone(),
696            parts: Some(WorkerRegistrationParts {
697                worker_id,
698                namespaces,
699                task_queue,
700                activity_types,
701            }),
702        })
703    }
704
705    /// Wait until at least one new worker registers.
706    ///
707    /// Returns immediately if a registration occurred since the last call.
708    /// Callers should re-check the registry after waking — the newly arrived
709    /// worker may not serve the namespace or activity type the caller needs.
710    pub async fn wait_for_worker(&self) {
711        self.worker_arrived.notified().await;
712    }
713
714    /// Return a snapshot of workers registered for the
715    /// `(namespace, task_queue, activity_type)` pool, ordered by worker id and
716    /// then rotated so each call starts from the next worker in the pool. The
717    /// rotation cursor is per triple, so each pool round-robins independently.
718    ///
719    /// When `node` is `Some`, the result is filtered to workers whose advertised
720    /// node equals it — a dispatch pinned to a node reaches only workers on that
721    /// node (NODE affinity = require). When `node` is `None`, the behaviour is
722    /// exactly the unpinned pool: every worker in the `(namespace, task_queue)`
723    /// pool is a candidate regardless of locality. node is a within-pool filter,
724    /// NOT part of the pool key, so the per-triple rotation cursor is shared
725    /// across pinned and unpinned lookups of the same pool.
726    ///
727    /// The id sort matters: `by_activity` holds workers in a `HashMap`, whose
728    /// iteration order is unspecified. Sorting first makes the rotation below
729    /// the sole, deterministic source of ordering — true round-robin across
730    /// calls with the same membership, not a wobble layered on hash order.
731    ///
732    /// # Errors
733    ///
734    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
735    pub fn workers_for(
736        &self,
737        namespace: &str,
738        task_queue: &str,
739        activity_type: &str,
740        node: Option<&str>,
741    ) -> Result<Vec<WorkerHandle>, ServerError> {
742        let mut state = self.state()?;
743        let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
744        let mut workers: Vec<WorkerHandle> = state
745            .by_activity
746            .get(&key)
747            .map(|workers| {
748                workers
749                    .values()
750                    .filter(|worker| worker_matches_node(worker, node))
751                    .cloned()
752                    .collect()
753            })
754            .unwrap_or_default();
755        if workers.is_empty() {
756            return Ok(workers);
757        }
758        workers.sort_by_key(WorkerHandle::id);
759        let idx = state.rotation.entry(key).or_insert(0);
760        let start = *idx % workers.len();
761        *idx = idx.wrapping_add(1);
762        let mut rotated = Vec::with_capacity(workers.len());
763        rotated.extend_from_slice(&workers[start..]);
764        rotated.extend_from_slice(&workers[..start]);
765        Ok(rotated)
766    }
767
768    /// Return a snapshot of every connected worker stream.
769    ///
770    /// # Errors
771    ///
772    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
773    pub fn all_workers(&self) -> Result<Vec<WorkerHandle>, ServerError> {
774        let state = self.state()?;
775        Ok(state.workers.values().cloned().collect())
776    }
777
778    /// Return the handle for a worker by id, or `None` when it is not registered.
779    ///
780    /// The intervention router resolves the owning worker of a target attempt by id
781    /// (NOI-6): the attempt-owner back-index stores a [`WorkerId`], and the router
782    /// reads back the live handle to gate on its advertised capabilities and select
783    /// its delivery. A `None` result means the owner disconnected — the router
784    /// treats that as the attempt-scoped no-op.
785    ///
786    /// # Errors
787    ///
788    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
789    pub fn worker_by_id(&self, worker_id: WorkerId) -> Result<Option<WorkerHandle>, ServerError> {
790        Ok(self.state()?.workers.get(&worker_id).cloned())
791    }
792
793    /// Replace the advertised intervention capabilities of a registered worker
794    /// (NOI-6). The liminal registration frame cannot carry capabilities, so a
795    /// liminal agent worker announces them on the reserved capabilities channel
796    /// right after registering, and this applies the announcement to the live
797    /// handle the intervention router gates on. Returns `false` when the worker
798    /// is no longer registered (a disconnect racing the announcement — benign).
799    ///
800    /// # Errors
801    ///
802    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
803    pub fn set_intervention_capabilities(
804        &self,
805        worker_id: WorkerId,
806        capabilities: &InterventionCapabilities,
807    ) -> Result<bool, ServerError> {
808        let mut state = self.state()?;
809        if !state.workers.contains_key(&worker_id) {
810            return Ok(false);
811        }
812        if let Some(handle) = state.workers.get_mut(&worker_id) {
813            handle.intervention_capabilities = capabilities.clone();
814        }
815        // The selection index holds handle clones; keep them capability-consistent
816        // even though capabilities are never a routing dimension.
817        for workers in state.by_activity.values_mut() {
818            if let Some(handle) = workers.get_mut(&worker_id) {
819                handle.intervention_capabilities = capabilities.clone();
820            }
821        }
822        Ok(true)
823    }
824
825    /// Broadcast a graceful drain request to every connected worker stream.
826    /// Workers are removed from routing before any transport signal is attempted.
827    /// A liminal worker has no drain control frame, so it is force-fenced and
828    /// deregistered with an error-level identity-bearing log instead of being
829    /// silently skipped.
830    ///
831    /// # Errors
832    ///
833    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
834    pub fn broadcast_drain(&self) -> Result<usize, ServerError> {
835        let workers = self.all_workers()?;
836        let mut delivered = 0usize;
837        for worker in workers {
838            if self.drain_worker(worker.id())? {
839                delivered = delivered.saturating_add(1);
840            }
841        }
842        Ok(delivered)
843    }
844
845    /// Stop assigning work to one worker and request a graceful transport drain.
846    ///
847    /// Returns `false` if the worker is not registered or if its transport has no
848    /// drain channel. In the latter case the worker is deregistered immediately,
849    /// after an error-level log names the worker and its transport limitation.
850    ///
851    /// # Errors
852    ///
853    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
854    pub fn drain_worker(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
855        let worker = {
856            let mut state = self.state()?;
857            let Some(worker) = state.workers.get(&worker_id).cloned() else {
858                return Ok(false);
859            };
860            Self::remove_worker_from_service(&mut state, &worker);
861            worker
862        };
863        match worker.delivery() {
864            WorkerDelivery::Grpc(sender) => {
865                if sender.try_send(WorkerMessage::DrainRequest).is_ok() {
866                    tracing::info!(worker_id = worker_id.value(), "worker drain requested");
867                    Ok(true)
868                } else {
869                    tracing::error!(
870                        worker_id = worker_id.value(),
871                        "worker drain signal failed; force-deregistering closed transport"
872                    );
873                    self.deregister(worker_id)?;
874                    Ok(false)
875                }
876            }
877            #[cfg(feature = "liminal-transport")]
878            WorkerDelivery::Liminal(delivery) => {
879                tracing::error!(
880                    worker_id = worker_id.value(),
881                    connection_pid = delivery.pid(),
882                    "liminal transport has no drain control channel; worker fenced and \
883                     deregistered at drain start"
884                );
885                self.deregister(worker_id)?;
886                Ok(false)
887            }
888        }
889    }
890
891    /// Select one worker for the `(namespace, task_queue, activity_type)` pool.
892    ///
893    /// When `node` is `Some`, only workers whose advertised node equals it are
894    /// considered (NODE affinity = require); `None` considers every worker in
895    /// the pool. node is a within-pool filter, NOT part of the pool key.
896    ///
897    /// # Errors
898    ///
899    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
900    pub fn select_worker(
901        &self,
902        namespace: &str,
903        task_queue: &str,
904        activity_type: &str,
905        node: Option<&str>,
906    ) -> Result<Option<WorkerHandle>, ServerError> {
907        let state = self.state()?;
908        let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
909        Ok(state.by_activity.get(&key).and_then(|workers| {
910            workers
911                .values()
912                .filter(|worker| worker_matches_node(worker, node))
913                // Reachability is a dispatch PRECONDITION, so it is enforced
914                // here rather than discovered at push time. Selecting a worker
915                // the server cannot push to produces a dispatch that can only
916                // fail, and on the liminal transport it fails by consuming
917                // connection capacity — so an unreachable worker chosen anyway
918                // makes its own unreachability worse.
919                .filter(|worker| !state.dispatch_ineligible.contains(&worker.id))
920                .min_by_key(|worker| worker.id)
921                .cloned()
922        }))
923    }
924
925    /// Publish the liveness probe's reachability verdict: the set of workers the
926    /// server cannot currently reach on the push leg, which
927    /// [`Self::select_worker`] then skips.
928    ///
929    /// Replaces the whole set rather than toggling one worker, so the published
930    /// verdict is always exactly one round's evidence and a worker can never be
931    /// left excluded by a stale entry nobody cleared.
932    ///
933    /// # Errors
934    ///
935    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
936    pub fn set_dispatch_ineligible(
937        &self,
938        unreachable: BTreeSet<WorkerId>,
939    ) -> Result<(), ServerError> {
940        self.state()?.dispatch_ineligible = unreachable;
941        Ok(())
942    }
943
944    /// Whether a worker is currently excluded from dispatch selection for
945    /// unreachability.
946    ///
947    /// # Errors
948    ///
949    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
950    pub fn is_dispatch_ineligible(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
951        Ok(self.state()?.dispatch_ineligible.contains(&worker_id))
952    }
953
954    /// Census the live fleet for one dispatch address (R1).
955    ///
956    /// Taken under the SAME lock discipline as [`Self::select_worker`] and read
957    /// immediately after a selection miss, so the taxonomy verdict describes
958    /// the fleet selection actually saw. Three nested counts — the pool, the
959    /// activity coverage within it, the node coverage within that — are what
960    /// separate `NO_LIVE_POLLERS` from `POLLERS_INCOMPATIBLE`.
961    ///
962    /// `last_compatible_poller_age` is zero while a compatible worker is
963    /// connected, the elapsed time since the most recent compatible departure
964    /// when one has left, and `None` when this server has never had one.
965    ///
966    /// # Errors
967    ///
968    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
969    pub fn pool_census(
970        &self,
971        namespace: &str,
972        task_queue: &str,
973        activity_type: &str,
974        node: Option<&str>,
975    ) -> Result<super::queue_service::PoolCensus, ServerError> {
976        let state = self.state()?;
977        let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
978        let workers_in_pool = state
979            .workers
980            .values()
981            .filter(|worker| {
982                worker.task_queue == task_queue && worker.namespaces.contains(namespace)
983            })
984            .count();
985        let serving = state.by_activity.get(&key);
986        let workers_serving_activity = serving.map_or(0, HashMap::len);
987        let compatible_workers = serving.map_or(0, |workers| {
988            workers
989                .values()
990                .filter(|worker| worker_matches_node(worker, node))
991                .count()
992        });
993        let last_compatible_poller_age = if compatible_workers > 0 {
994            Some(Duration::ZERO)
995        } else {
996            state
997                .last_departure
998                .get(&key)
999                .and_then(|by_node| {
1000                    by_node
1001                        .iter()
1002                        .filter(|(departed_node, _)| match node {
1003                            None => true,
1004                            Some(node) => departed_node.as_deref() == Some(node),
1005                        })
1006                        .map(|(_, departed_at)| *departed_at)
1007                        .max()
1008                })
1009                .map(|departed_at| departed_at.elapsed())
1010        };
1011        Ok(super::queue_service::PoolCensus {
1012            workers_in_pool,
1013            workers_serving_activity,
1014            compatible_workers,
1015            last_compatible_poller_age,
1016        })
1017    }
1018
1019    /// Return whether a worker stream is currently registered.
1020    ///
1021    /// The activity dispatch path uses this after queuing a task to detect a
1022    /// worker whose stream tore down concurrently: a sweep that ran before
1023    /// the dispatch tracked its task can never complete it, so the dispatch
1024    /// must fail the activity itself instead of waiting forever.
1025    ///
1026    /// # Errors
1027    ///
1028    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1029    pub fn is_registered(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1030        Ok(self.state()?.workers.contains_key(&worker_id))
1031    }
1032
1033    /// Remove a worker by id from every namespace/activity index it advertised.
1034    ///
1035    /// Emits a WS3 [`WorkerDeathReason::Disconnect`] delta — the truthful default
1036    /// for a removed worker whose stream/registration went away. Callers that can
1037    /// PROVE a finer reason (a liveness-timeout sweep) call
1038    /// [`Self::deregister_with_reason`] instead, so the dashboard never sees a
1039    /// fabricated distinction.
1040    ///
1041    /// # Errors
1042    ///
1043    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1044    pub fn deregister(&self, worker_id: WorkerId) -> Result<(), ServerError> {
1045        self.deregister_with_reason(worker_id, WorkerDeathReason::Disconnect)
1046    }
1047
1048    /// Remove a worker by id, attributing the departure to an explicit
1049    /// [`WorkerDeathReason`] the caller can prove at its call site (for example a
1050    /// heartbeat sweep passes [`WorkerDeathReason::Timeout`]).
1051    ///
1052    /// # Errors
1053    ///
1054    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1055    pub fn deregister_with_reason(
1056        &self,
1057        worker_id: WorkerId,
1058        reason: WorkerDeathReason,
1059    ) -> Result<(), ServerError> {
1060        let mut state = self.state()?;
1061        let removed_namespaces = Self::remove_worker(&mut state, worker_id);
1062        drop(state);
1063
1064        let Some(namespaces) = removed_namespaces else {
1065            // Already gone: no metrics double-count, no duplicate delta.
1066            return Ok(());
1067        };
1068
1069        if let Some(metrics) = &self.metrics {
1070            for namespace in &namespaces {
1071                metrics.worker_disconnected(namespace);
1072            }
1073        }
1074        self.emit_worker_disconnected(worker_id, &namespaces, reason);
1075
1076        Ok(())
1077    }
1078
1079    /// Emit a WS3 `WorkerDisconnected` delta if a publisher is attached.
1080    fn emit_worker_disconnected(
1081        &self,
1082        worker_id: WorkerId,
1083        namespaces: &BTreeSet<String>,
1084        reason: WorkerDeathReason,
1085    ) {
1086        if let Some(publisher) = &self.cluster_publisher {
1087            let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
1088            drop(publisher.emit(|meta| ClusterEvent::WorkerDisconnected {
1089                meta,
1090                worker_id: worker_id.value().to_string(),
1091                namespaces: namespaces_vec,
1092                reason,
1093            }));
1094        }
1095    }
1096
1097    /// Remove a worker from every `(namespace, task_queue, activity_type)` index
1098    /// it advertised. Returns the namespace set it served (for metrics), or
1099    /// `None` if the worker was already gone.
1100    fn remove_worker(state: &mut RegistryState, worker_id: WorkerId) -> Option<BTreeSet<String>> {
1101        let handle = state.workers.remove(&worker_id)?;
1102        Self::remove_worker_from_service(state, &handle);
1103        Some(handle.namespaces)
1104    }
1105
1106    /// Fence a worker from every routing index while retaining its connection
1107    /// handle so an in-progress graceful drain can finish and tear down normally.
1108    fn remove_worker_from_service(state: &mut RegistryState, handle: &WorkerHandle) {
1109        let departed_at = Instant::now();
1110        for namespace in &handle.namespaces {
1111            let pool = PoolAddress::new(namespace.clone(), handle.task_queue.clone());
1112            for activity_type in &handle.activity_types {
1113                let key = ActivityKey::new(pool.clone(), activity_type.clone());
1114                // Departure time is recorded for EVERY address this worker
1115                // served, so a later census can age the last compatible poller
1116                // instead of reporting a bare "nobody" (R1).
1117                state
1118                    .last_departure
1119                    .entry(key.clone())
1120                    .or_default()
1121                    .insert(handle.node.clone(), departed_at);
1122                if let Some(workers) = state.by_activity.get_mut(&key) {
1123                    workers.remove(&handle.id);
1124                    if workers.is_empty() {
1125                        state.by_activity.remove(&key);
1126                        // Prune the round-robin cursor in lockstep: the cursor
1127                        // map is keyed on arbitrary caller-supplied strings and
1128                        // is lazily created by `workers_for`, so leaving stale
1129                        // entries behind leaks memory unboundedly on a
1130                        // never-dying server. When the last worker for a triple
1131                        // leaves, its cursor has no remaining meaning.
1132                        state.rotation.remove(&key);
1133                    }
1134                }
1135            }
1136        }
1137    }
1138
1139    fn state(&self) -> Result<MutexGuard<'_, RegistryState>, ServerError> {
1140        self.inner
1141            .lock()
1142            .map_err(|_| ServerError::lock_poisoned("connected worker registry"))
1143    }
1144}
1145
1146/// Map a held [`WorkerDelivery`] to the wire [`WorkerTransport`] discriminant for
1147/// the WS3 `WorkerConnected` delta.
1148const fn transport_of(delivery: &WorkerDelivery) -> WorkerTransport {
1149    match delivery {
1150        WorkerDelivery::Grpc(_) => WorkerTransport::Grpc,
1151        #[cfg(feature = "liminal-transport")]
1152        WorkerDelivery::Liminal(_) => WorkerTransport::Liminal,
1153    }
1154}
1155
1156/// Normalize a wire `node` string into an optional locality affinity: an empty
1157/// value (the proto3 default) carries no node, anything else is the worker's
1158/// advertised node id.
1159///
1160/// Shared with contract admission ([`super::contracts::validate_worker_contracts`])
1161/// rather than restated there: admission decides which of a package's actions a
1162/// connection owes from the same locality this registry then routes by, and two
1163/// independent normalizations of the wire default would be a place for that
1164/// agreement to drift silently.
1165pub(crate) fn optional_node(node: &str) -> Option<String> {
1166    if node.is_empty() {
1167        None
1168    } else {
1169        Some(node.to_owned())
1170    }
1171}
1172
1173/// Whether a worker satisfies an optional node filter. `None` (unpinned) matches
1174/// every worker; `Some(node)` matches only a worker advertising that exact node
1175/// (NODE affinity = require). A worker with no advertised node never matches a
1176/// pinned dispatch.
1177fn worker_matches_node(worker: &WorkerHandle, node: Option<&str>) -> bool {
1178    match node {
1179        None => true,
1180        Some(node) => worker.node() == Some(node),
1181    }
1182}
1183
1184#[derive(Clone, Debug)]
1185struct WorkerRegistrationParts {
1186    worker_id: WorkerId,
1187    namespaces: BTreeSet<String>,
1188    task_queue: String,
1189    activity_types: BTreeSet<String>,
1190}
1191
1192/// Registration token owned by the worker stream task.
1193///
1194/// Dropping the token performs best-effort cleanup for disconnect paths. Call
1195/// [`WorkerRegistration::deregister`] when the caller needs a typed poison error.
1196#[derive(Debug)]
1197pub struct WorkerRegistration {
1198    registry: ConnectedWorkerRegistry,
1199    parts: Option<WorkerRegistrationParts>,
1200}
1201
1202impl WorkerRegistration {
1203    /// Worker id assigned to this registration.
1204    #[must_use]
1205    pub fn worker_id(&self) -> Option<WorkerId> {
1206        self.parts.as_ref().map(|parts| parts.worker_id)
1207    }
1208
1209    /// Authorized namespace set for this registration.
1210    #[must_use]
1211    pub fn namespaces(&self) -> Option<&BTreeSet<String>> {
1212        self.parts.as_ref().map(|parts| &parts.namespaces)
1213    }
1214
1215    /// Task queue (pool/flavour) this registration serves within each namespace.
1216    #[must_use]
1217    pub fn task_queue(&self) -> Option<&str> {
1218        self.parts.as_ref().map(|parts| parts.task_queue.as_str())
1219    }
1220
1221    /// Activity types advertised by this registration.
1222    #[must_use]
1223    pub fn activity_types(&self) -> Option<&BTreeSet<String>> {
1224        self.parts.as_ref().map(|parts| &parts.activity_types)
1225    }
1226
1227    /// Explicitly remove this worker from the registry.
1228    ///
1229    /// # Errors
1230    ///
1231    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1232    pub fn deregister(mut self) -> Result<(), ServerError> {
1233        let Some(parts) = self.parts.take() else {
1234            return Ok(());
1235        };
1236        self.registry.deregister(parts.worker_id)
1237    }
1238}
1239
1240impl Drop for WorkerRegistration {
1241    fn drop(&mut self) {
1242        let Some(parts) = self.parts.take() else {
1243            return;
1244        };
1245        let removed_namespaces = self.registry.inner.lock().ok().and_then(|mut state| {
1246            ConnectedWorkerRegistry::remove_worker(&mut state, parts.worker_id)
1247        });
1248        if let Some(namespaces) = removed_namespaces {
1249            if let Some(metrics) = &self.registry.metrics {
1250                for namespace in &namespaces {
1251                    metrics.worker_disconnected(namespace);
1252                }
1253            }
1254            // A dropped registration token means the worker's stream/connection
1255            // went away — the truthful reason is Disconnect, not a fabricated
1256            // timeout/deregister distinction this path cannot prove.
1257            self.registry.emit_worker_disconnected(
1258                parts.worker_id,
1259                &namespaces,
1260                WorkerDeathReason::Disconnect,
1261            );
1262        }
1263    }
1264}
1265
1266#[cfg(test)]
1267mod tests {
1268    use crate::config::NamespaceMode;
1269    use crate::namespace::{NamespaceResolver, StaticScheduleNamespaces, StaticWorkflowNamespaces};
1270    use crate::worker::heartbeat::{DISPATCH_PROBATION_PINGS, HeartbeatTracker};
1271
1272    use super::*;
1273
1274    fn guard() -> NamespaceGuard {
1275        NamespaceGuard::new(NamespaceResolver::authorization_only(
1276            NamespaceMode::SharedEngine,
1277            StaticWorkflowNamespaces::default(),
1278            StaticScheduleNamespaces::default(),
1279        ))
1280    }
1281
1282    fn caller(namespace: &str) -> CallerIdentity {
1283        CallerIdentity::new("worker", [namespace.to_owned()])
1284    }
1285
1286    /// A test-expectation failure as a `ServerError`, so `Result`-returning
1287    /// tests can fail on an unexpected `None` without `panic!`/`expect`.
1288    fn test_failure(message: &str) -> ServerError {
1289        ServerError::worker_dispatch("default".to_owned(), "test".to_owned(), message.to_owned())
1290    }
1291
1292    fn registration(namespace: &str, activity_types: &[&str]) -> ProtoRegisterWorker {
1293        registration_with_queue(namespace, "", activity_types)
1294    }
1295
1296    fn registration_with_queue(
1297        namespace: &str,
1298        task_queue: &str,
1299        activity_types: &[&str],
1300    ) -> ProtoRegisterWorker {
1301        registration_full(&[namespace], task_queue, "", activity_types)
1302    }
1303
1304    fn registration_full(
1305        namespaces: &[&str],
1306        task_queue: &str,
1307        node: &str,
1308        activity_types: &[&str],
1309    ) -> ProtoRegisterWorker {
1310        ProtoRegisterWorker {
1311            namespaces: namespaces.iter().map(|value| (*value).to_owned()).collect(),
1312            activity_types: activity_types
1313                .iter()
1314                .map(|value| (*value).to_owned())
1315                .collect(),
1316            task_queue: task_queue.to_owned(),
1317            node: node.to_owned(),
1318            activities: Vec::new(),
1319            identity: String::new(),
1320        }
1321    }
1322
1323    fn multi_caller(namespaces: &[&str]) -> CallerIdentity {
1324        CallerIdentity::new("worker", namespaces.iter().map(|value| (*value).to_owned()))
1325    }
1326
1327    /// `set_intervention_capabilities` replaces a live worker's advertised set —
1328    /// the announcement path a liminal agent worker takes after its in-band
1329    /// registration (the registration frame cannot carry capabilities) — and
1330    /// reports an unknown worker as `false` (an announcement racing a
1331    /// disconnect is benign, never an error).
1332    #[tokio::test]
1333    async fn set_intervention_capabilities_updates_live_worker() -> Result<(), ServerError> {
1334        let registry = ConnectedWorkerRegistry::default();
1335        let (sender, _receiver) = mpsc::channel(1);
1336        let types = ["scout".to_owned()];
1337        let guard = registry.register_delivery_with_capabilities(
1338            ["default".to_owned()],
1339            "default",
1340            None,
1341            types.iter(),
1342            WorkerDelivery::Grpc(sender),
1343            InterventionCapabilities::none(),
1344        )?;
1345        let Some(worker_id) = guard.worker_id() else {
1346            return Err(test_failure("registration carries an id"));
1347        };
1348
1349        let announced = InterventionCapabilities {
1350            supported: vec![aion_core::InterventionPrimitive::InjectMessage],
1351        };
1352        assert!(
1353            registry.set_intervention_capabilities(worker_id, &announced)?,
1354            "a live worker's capabilities must be updatable"
1355        );
1356        let Some(handle) = registry.worker_by_id(worker_id)? else {
1357            return Err(test_failure("worker stays registered"));
1358        };
1359        assert_eq!(handle.intervention_capabilities(), &announced);
1360
1361        assert!(
1362            !registry.set_intervention_capabilities(WorkerId(u64::MAX), &announced)?,
1363            "an unknown worker reports false, never an error"
1364        );
1365        Ok(())
1366    }
1367
1368    #[tokio::test]
1369    async fn register_and_deregister_are_namespace_isolated() -> Result<(), ServerError> {
1370        let registry = ConnectedWorkerRegistry::default();
1371        let (tenant_a_tx, _tenant_a_rx) = mpsc::channel(1);
1372        let (tenant_b_tx, _tenant_b_rx) = mpsc::channel(1);
1373
1374        let tenant_a = registry
1375            .accept_registration(
1376                &guard(),
1377                &caller("tenant-a"),
1378                &registration("tenant-a", &["charge", "charge"]),
1379                tenant_a_tx,
1380            )
1381            .await?;
1382        let tenant_b = registry
1383            .accept_registration(
1384                &guard(),
1385                &caller("tenant-b"),
1386                &registration("tenant-b", &["charge"]),
1387                tenant_b_tx,
1388            )
1389            .await?;
1390
1391        let tq = DEFAULT_TASK_QUEUE;
1392        assert_eq!(
1393            registry.workers_for("tenant-a", tq, "charge", None)?.len(),
1394            1
1395        );
1396        assert_eq!(
1397            registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1398            1
1399        );
1400        assert!(
1401            registry
1402                .workers_for("tenant-a", tq, "missing", None)?
1403                .is_empty()
1404        );
1405
1406        let tenant_a_id = tenant_a.worker_id();
1407        tenant_a.deregister()?;
1408
1409        assert!(
1410            registry
1411                .workers_for("tenant-a", tq, "charge", None)?
1412                .is_empty()
1413        );
1414        assert_eq!(
1415            registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1416            1
1417        );
1418        assert_ne!(tenant_a_id, tenant_b.worker_id());
1419
1420        tenant_b.deregister()?;
1421        assert!(
1422            registry
1423                .workers_for("tenant-b", tq, "charge", None)?
1424                .is_empty()
1425        );
1426        Ok(())
1427    }
1428
1429    #[tokio::test]
1430    async fn denied_namespace_is_not_registered() -> Result<(), ServerError> {
1431        let registry = ConnectedWorkerRegistry::default();
1432        let (tx, _rx) = mpsc::channel(1);
1433        let denied = registry
1434            .accept_registration(
1435                &guard(),
1436                &caller("tenant-a"),
1437                &registration("tenant-b", &["charge"]),
1438                tx,
1439            )
1440            .await;
1441
1442        assert!(denied.is_err());
1443        assert!(
1444            registry
1445                .workers_for("tenant-b", DEFAULT_TASK_QUEUE, "charge", None)?
1446                .is_empty()
1447        );
1448        Ok(())
1449    }
1450
1451    #[tokio::test]
1452    async fn task_queues_partition_disjoint_pools_within_one_namespace() -> Result<(), ServerError>
1453    {
1454        // Same namespace + same activity_type, two DIFFERENT task queues: the
1455        // pools are disjoint, a lookup for one queue never returns the other's
1456        // worker, and round-robin holds independently per (ns, tq, type) triple.
1457        let registry = ConnectedWorkerRegistry::default();
1458        let (norn_tx, _norn_rx) = mpsc::channel(1);
1459        let (claude_a_tx, _claude_a_rx) = mpsc::channel(1);
1460        let (claude_b_tx, _claude_b_rx) = mpsc::channel(1);
1461
1462        let norn = registry
1463            .accept_registration(
1464                &guard(),
1465                &caller("local"),
1466                &registration_with_queue("local", "norn", &["dev"]),
1467                norn_tx,
1468            )
1469            .await?;
1470        // Two workers on the SAME (local, claude) pool to exercise round-robin.
1471        let claude_a = registry
1472            .accept_registration(
1473                &guard(),
1474                &caller("local"),
1475                &registration_with_queue("local", "claude", &["dev"]),
1476                claude_a_tx,
1477            )
1478            .await?;
1479        let claude_b = registry
1480            .accept_registration(
1481                &guard(),
1482                &caller("local"),
1483                &registration_with_queue("local", "claude", &["dev"]),
1484                claude_b_tx,
1485            )
1486            .await?;
1487
1488        let norn_pool = registry.workers_for("local", "norn", "dev", None)?;
1489        assert_eq!(norn_pool.len(), 1, "norn pool has exactly its one worker");
1490        let norn_id = norn.worker_id().ok_or_else(missing_id)?;
1491        assert_eq!(norn_pool[0].id(), norn_id);
1492
1493        let claude_pool = registry.workers_for("local", "claude", "dev", None)?;
1494        assert_eq!(
1495            claude_pool.len(),
1496            2,
1497            "claude pool sees only its two workers"
1498        );
1499        let claude_ids: BTreeSet<WorkerId> = claude_pool.iter().map(WorkerHandle::id).collect();
1500        assert!(
1501            !claude_ids.contains(&norn_id),
1502            "the norn worker must never appear in the claude pool"
1503        );
1504
1505        // A dispatch targeting `norn` never reaches a `claude` worker, and vice
1506        // versa: the disjoint key is the boundary.
1507        assert!(
1508            !registry
1509                .workers_for("local", "norn", "dev", None)?
1510                .iter()
1511                .any(|worker| claude_ids.contains(&worker.id()))
1512        );
1513
1514        // Round-robin per triple: the (local, claude, dev) cursor advances
1515        // independently and cycles through both claude workers, while the
1516        // (local, norn, dev) cursor keeps returning its single worker.
1517        let first = registry.workers_for("local", "claude", "dev", None)?[0].id();
1518        let second = registry.workers_for("local", "claude", "dev", None)?[0].id();
1519        assert_ne!(
1520            first, second,
1521            "claude pool round-robins across both workers"
1522        );
1523        assert_eq!(
1524            registry.workers_for("local", "norn", "dev", None)?[0].id(),
1525            norn_id,
1526            "the norn pool rotation is unaffected by claude traffic"
1527        );
1528
1529        norn.deregister()?;
1530        claude_a.deregister()?;
1531        claude_b.deregister()?;
1532        Ok(())
1533    }
1534
1535    #[tokio::test]
1536    async fn same_task_queue_in_different_namespaces_is_isolated() -> Result<(), ServerError> {
1537        // Same task_queue string, two DIFFERENT namespaces: namespace is the
1538        // correctness boundary, so the pools are isolated.
1539        let registry = ConnectedWorkerRegistry::default();
1540        let (local_tx, _local_rx) = mpsc::channel(1);
1541        let (remote_tx, _remote_rx) = mpsc::channel(1);
1542
1543        let local = registry
1544            .accept_registration(
1545                &guard(),
1546                &caller("local"),
1547                &registration_with_queue("local", "gpu", &["render"]),
1548                local_tx,
1549            )
1550            .await?;
1551        let remote = registry
1552            .accept_registration(
1553                &guard(),
1554                &caller("remote"),
1555                &registration_with_queue("remote", "gpu", &["render"]),
1556                remote_tx,
1557            )
1558            .await?;
1559
1560        let local_pool = registry.workers_for("local", "gpu", "render", None)?;
1561        let remote_pool = registry.workers_for("remote", "gpu", "render", None)?;
1562        assert_eq!(local_pool.len(), 1);
1563        assert_eq!(remote_pool.len(), 1);
1564        assert_ne!(
1565            local_pool[0].id(),
1566            remote_pool[0].id(),
1567            "a shared task_queue string does not merge two namespaces"
1568        );
1569
1570        local.deregister()?;
1571        assert!(
1572            registry
1573                .workers_for("local", "gpu", "render", None)?
1574                .is_empty(),
1575            "deregistering the local worker leaves the remote namespace untouched"
1576        );
1577        assert_eq!(
1578            registry.workers_for("remote", "gpu", "render", None)?.len(),
1579            1
1580        );
1581
1582        remote.deregister()?;
1583        Ok(())
1584    }
1585
1586    #[tokio::test]
1587    async fn worker_serving_a_namespace_set_is_reachable_in_each() -> Result<(), ServerError> {
1588        // A worker advertising {a, b} is reachable for dispatch in BOTH a and b;
1589        // a worker in {a} is NOT reachable in b.
1590        let registry = ConnectedWorkerRegistry::default();
1591        let (ab_tx, _ab_rx) = mpsc::channel(1);
1592        let (a_tx, _a_rx) = mpsc::channel(1);
1593
1594        let worker_ab = registry
1595            .accept_registration(
1596                &guard(),
1597                &multi_caller(&["a", "b"]),
1598                &registration_full(&["a", "b"], "default", "", &["dev"]),
1599                ab_tx,
1600            )
1601            .await?;
1602        let worker_a = registry
1603            .accept_registration(
1604                &guard(),
1605                &caller("a"),
1606                &registration_full(&["a"], "default", "", &["dev"]),
1607                a_tx,
1608            )
1609            .await?;
1610
1611        let in_a = registry.workers_for("a", "default", "dev", None)?;
1612        let in_b = registry.workers_for("b", "default", "dev", None)?;
1613        let both_id = worker_ab.worker_id().ok_or_else(missing_id)?;
1614        let only_a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1615
1616        // Namespace a sees BOTH workers; namespace b sees ONLY the {a, b} worker.
1617        let a_ids: BTreeSet<WorkerId> = in_a.iter().map(WorkerHandle::id).collect();
1618        assert_eq!(a_ids, BTreeSet::from([both_id, only_a_id]));
1619        assert_eq!(in_b.len(), 1, "only the {{a, b}} worker is reachable in b");
1620        assert_eq!(in_b[0].id(), both_id);
1621        assert!(
1622            !in_b.iter().any(|worker| worker.id() == only_a_id),
1623            "the {{a}}-only worker must not be reachable in b"
1624        );
1625
1626        // Deregistering the {a, b} worker removes it from BOTH buckets.
1627        worker_ab.deregister()?;
1628        assert!(
1629            registry
1630                .workers_for("b", "default", "dev", None)?
1631                .is_empty()
1632        );
1633        assert_eq!(registry.workers_for("a", "default", "dev", None)?.len(), 1);
1634
1635        worker_a.deregister()?;
1636        Ok(())
1637    }
1638
1639    #[tokio::test]
1640    async fn node_pin_filters_within_pool() -> Result<(), ServerError> {
1641        // Two workers in the same (namespace, task_queue) pool on different
1642        // nodes: unpinned round-robins across both; pinned to node N reaches
1643        // ONLY the worker(s) on N; pinned to a node with no worker finds none.
1644        let registry = ConnectedWorkerRegistry::default();
1645        let (n1_tx, _n1_rx) = mpsc::channel(1);
1646        let (n2_tx, _n2_rx) = mpsc::channel(1);
1647
1648        let on_n1 = registry
1649            .accept_registration(
1650                &guard(),
1651                &caller("ns"),
1652                &registration_full(&["ns"], "tq", "n1", &["dev"]),
1653                n1_tx,
1654            )
1655            .await?;
1656        let on_n2 = registry
1657            .accept_registration(
1658                &guard(),
1659                &caller("ns"),
1660                &registration_full(&["ns"], "tq", "n2", &["dev"]),
1661                n2_tx,
1662            )
1663            .await?;
1664        let n1_id = on_n1.worker_id().ok_or_else(missing_id)?;
1665        let n2_id = on_n2.worker_id().ok_or_else(missing_id)?;
1666
1667        // Unpinned: both workers are candidates and round-robin advances.
1668        let unpinned = registry.workers_for("ns", "tq", "dev", None)?;
1669        assert_eq!(unpinned.len(), 2, "unpinned reaches the whole pool");
1670        let first = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1671        let second = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1672        assert_ne!(first, second, "unpinned round-robins across both nodes");
1673
1674        // Pinned to n1: only the n1 worker; pinned to n2: only the n2 worker.
1675        let pinned_n1 = registry.workers_for("ns", "tq", "dev", Some("n1"))?;
1676        assert_eq!(pinned_n1.len(), 1);
1677        assert_eq!(pinned_n1[0].id(), n1_id);
1678        let pinned_n2 = registry.workers_for("ns", "tq", "dev", Some("n2"))?;
1679        assert_eq!(pinned_n2.len(), 1);
1680        assert_eq!(pinned_n2[0].id(), n2_id);
1681
1682        // select_worker honours the same filter.
1683        assert_eq!(
1684            registry
1685                .select_worker("ns", "tq", "dev", Some("n1"))?
1686                .map(|worker| worker.id()),
1687            Some(n1_id)
1688        );
1689
1690        // Pinned to a node with no worker finds no candidate (the dispatcher
1691        // then waits via the same no-worker path the existing test exercises).
1692        assert!(
1693            registry
1694                .workers_for("ns", "tq", "dev", Some("absent"))?
1695                .is_empty(),
1696            "a pin to a node with no worker yields no candidate"
1697        );
1698        assert!(
1699            registry
1700                .select_worker("ns", "tq", "dev", Some("absent"))?
1701                .is_none()
1702        );
1703
1704        on_n1.deregister()?;
1705        on_n2.deregister()?;
1706        Ok(())
1707    }
1708
1709    #[tokio::test]
1710    async fn shared_node_id_round_robins_across_workers() -> Result<(), ServerError> {
1711        // Two workers SHARING a node id in the same pool: a dispatch pinned to
1712        // that node round-robins across BOTH (node is locality, not process).
1713        let registry = ConnectedWorkerRegistry::default();
1714        let (a_tx, _a_rx) = mpsc::channel(1);
1715        let (b_tx, _b_rx) = mpsc::channel(1);
1716
1717        let worker_a = registry
1718            .accept_registration(
1719                &guard(),
1720                &caller("ns"),
1721                &registration_full(&["ns"], "tq", "shared", &["dev"]),
1722                a_tx,
1723            )
1724            .await?;
1725        let worker_b = registry
1726            .accept_registration(
1727                &guard(),
1728                &caller("ns"),
1729                &registration_full(&["ns"], "tq", "shared", &["dev"]),
1730                b_tx,
1731            )
1732            .await?;
1733        let a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1734        let b_id = worker_b.worker_id().ok_or_else(missing_id)?;
1735
1736        let pinned = registry.workers_for("ns", "tq", "dev", Some("shared"))?;
1737        assert_eq!(
1738            pinned.len(),
1739            2,
1740            "both workers on the shared node are candidates"
1741        );
1742        let pinned_ids: BTreeSet<WorkerId> = pinned.iter().map(WorkerHandle::id).collect();
1743        assert_eq!(pinned_ids, BTreeSet::from([a_id, b_id]));
1744
1745        let first = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1746        let second = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1747        assert_ne!(
1748            first, second,
1749            "a pin to a shared node round-robins across both workers on it"
1750        );
1751
1752        worker_a.deregister()?;
1753        worker_b.deregister()?;
1754        Ok(())
1755    }
1756
1757    #[tokio::test]
1758    async fn rotation_cursor_is_pruned_when_last_worker_leaves() -> Result<(), ServerError> {
1759        // The round-robin cursor is keyed on arbitrary caller-supplied strings;
1760        // it must not outlive the pool it rotates, or a never-dying server leaks
1761        // memory. After the last worker for a triple deregisters, no cursor for
1762        // that triple may remain.
1763        let registry = ConnectedWorkerRegistry::default();
1764        let (tx, _rx) = mpsc::channel(1);
1765        let worker = registry
1766            .accept_registration(
1767                &guard(),
1768                &caller("ns"),
1769                &registration_full(&["ns"], "tq", "", &["dev"]),
1770                tx,
1771            )
1772            .await?;
1773
1774        // Drive the lazy cursor insert.
1775        let _ = registry.workers_for("ns", "tq", "dev", None)?;
1776        let key = ActivityKey::new(PoolAddress::new("ns", "tq"), "dev");
1777        assert!(
1778            registry.state()?.rotation.contains_key(&key),
1779            "a lookup must have created the rotation cursor"
1780        );
1781
1782        worker.deregister()?;
1783        let state = registry.state()?;
1784        assert!(
1785            !state.rotation.contains_key(&key),
1786            "the rotation cursor must be pruned once the last worker leaves"
1787        );
1788        assert!(
1789            !state.by_activity.contains_key(&key),
1790            "the activity bucket must also be gone"
1791        );
1792        Ok(())
1793    }
1794
1795    fn missing_id() -> ServerError {
1796        ServerError::lock_poisoned("registration unexpectedly missing a worker id")
1797    }
1798
1799    // ---- Minted-on-use (Control-Plane Phase 1) -----------------------------
1800
1801    fn namespace_store() -> Arc<dyn NamespaceStore> {
1802        Arc::new(aion_store::InMemoryStore::default())
1803    }
1804
1805    fn minting_registry(
1806        store: &Arc<dyn NamespaceStore>,
1807        policy: AutoCreate,
1808    ) -> ConnectedWorkerRegistry {
1809        ConnectedWorkerRegistry::default().with_namespace_minting(Arc::clone(store), policy)
1810    }
1811
1812    #[tokio::test]
1813    async fn open_register_mints_durable_record_and_is_idempotent() -> Result<(), ServerError> {
1814        let store = namespace_store();
1815        let registry = minting_registry(&store, AutoCreate::Open);
1816
1817        // First registration mints the namespace.
1818        let (tx_one, _rx_one) = mpsc::channel(1);
1819        let first = registry
1820            .accept_registration(
1821                &guard(),
1822                &caller("orders"),
1823                &registration("orders", &["charge"]),
1824                tx_one,
1825            )
1826            .await?;
1827        let record = store
1828            .get_namespace("orders")
1829            .await?
1830            .ok_or_else(|| ServerError::namespace_denied("expected a minted record"))?;
1831        assert_eq!(record.name, "orders");
1832        assert_eq!(record.origin, NamespaceOrigin::WorkerMint);
1833
1834        // Re-registering the same namespace is idempotent: no duplicate row,
1835        // and the prior worker is unaffected.
1836        let (tx_two, _rx_two) = mpsc::channel(1);
1837        registry
1838            .accept_registration(
1839                &guard(),
1840                &caller("orders"),
1841                &registration("orders", &["refund"]),
1842                tx_two,
1843            )
1844            .await?;
1845        let all = store.list_namespaces().await?;
1846        assert_eq!(
1847            all.iter().filter(|r| r.name == "orders").count(),
1848            1,
1849            "re-register must not create a duplicate namespace row"
1850        );
1851        drop(first);
1852        Ok(())
1853    }
1854
1855    #[tokio::test]
1856    async fn open_register_mints_each_namespace_in_a_multi_namespace_worker()
1857    -> Result<(), ServerError> {
1858        let store = namespace_store();
1859        let registry = minting_registry(&store, AutoCreate::Open);
1860        let (tx, _rx) = mpsc::channel(1);
1861
1862        registry
1863            .accept_registration(
1864                &guard(),
1865                &multi_caller(&["alpha", "beta"]),
1866                &registration_full(&["alpha", "beta"], "", "", &["charge"]),
1867                tx,
1868            )
1869            .await?;
1870
1871        assert!(store.get_namespace("alpha").await?.is_some());
1872        assert!(store.get_namespace("beta").await?.is_some());
1873        Ok(())
1874    }
1875
1876    // ---- Placement admission (Control-Plane Phase 2, P2-I1) -----------------
1877
1878    /// Pre-mint `namespace` and set its placement to `Pinned{nodes}`, returning a
1879    /// minting registry over the same store so `accept_registration` reads the
1880    /// placement from the SAME durable record.
1881    async fn pinned_registry(
1882        store: &Arc<dyn NamespaceStore>,
1883        namespace: &str,
1884        nodes: &[&str],
1885    ) -> Result<ConnectedWorkerRegistry, ServerError> {
1886        store
1887            .register_namespace(namespace, NamespaceOrigin::Explicit)
1888            .await?;
1889        store
1890            .set_namespace_placement(
1891                namespace,
1892                NamespacePlacement::Pinned {
1893                    nodes: nodes.iter().map(|n| (*n).to_owned()).collect(),
1894                },
1895            )
1896            .await?;
1897        Ok(minting_registry(store, AutoCreate::Open))
1898    }
1899
1900    /// A worker on a node IN the required set registers successfully into a
1901    /// `Pinned{n1}` namespace, and is reachable in the pool.
1902    #[tokio::test]
1903    async fn pinned_admits_a_worker_on_a_required_node() -> Result<(), ServerError> {
1904        let store = namespace_store();
1905        let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1906        let (tx, _rx) = mpsc::channel(1);
1907
1908        let _registration = registry
1909            .accept_registration(
1910                &guard(),
1911                &caller("iso"),
1912                &registration_full(&["iso"], "", "n1", &["charge"]),
1913                tx,
1914            )
1915            .await?;
1916
1917        assert_eq!(
1918            registry
1919                .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", Some("n1"))?
1920                .len(),
1921            1,
1922            "an n1 worker must be admitted into the Pinned{{n1}} namespace's pool"
1923        );
1924        Ok(())
1925    }
1926
1927    /// A worker on a node NOT in the required set is rejected — the WHOLE
1928    /// registration fails (loud) and no worker is inserted. This would FAIL under
1929    /// no admission gate (the worker would join and steal Pinned dispatches).
1930    #[tokio::test]
1931    async fn pinned_rejects_a_wrong_node_worker() -> Result<(), ServerError> {
1932        let store = namespace_store();
1933        let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1934        let (tx, _rx) = mpsc::channel(1);
1935
1936        let denied = registry
1937            .accept_registration(
1938                &guard(),
1939                &caller("iso"),
1940                &registration_full(&["iso"], "", "n2", &["charge"]),
1941                tx,
1942            )
1943            .await;
1944        assert!(
1945            matches!(denied, Err(ServerError::Namespace { .. })),
1946            "a wrong-node (n2) worker must be rejected from a Pinned{{n1}} namespace"
1947        );
1948        assert!(
1949            registry
1950                .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
1951                .is_empty(),
1952            "a rejected registration must not insert a worker on any node"
1953        );
1954        Ok(())
1955    }
1956
1957    /// A worker advertising NO node (`node == ""` → `None`) is rejected from a
1958    /// `Pinned{n1}` namespace: an unlabelled worker can never satisfy a hard pin.
1959    #[tokio::test]
1960    async fn pinned_rejects_a_node_less_worker() -> Result<(), ServerError> {
1961        let store = namespace_store();
1962        let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1963        let (tx, _rx) = mpsc::channel(1);
1964
1965        let denied = registry
1966            .accept_registration(
1967                &guard(),
1968                &caller("iso"),
1969                &registration_full(&["iso"], "", "", &["charge"]),
1970                tx,
1971            )
1972            .await;
1973        assert!(
1974            matches!(denied, Err(ServerError::Namespace { .. })),
1975            "a node-less worker must be rejected from a Pinned{{n1}} namespace"
1976        );
1977        assert!(
1978            registry
1979                .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
1980                .is_empty(),
1981            "a rejected node-less registration must not insert a worker"
1982        );
1983        Ok(())
1984    }
1985
1986    /// Reject-WHOLE-registration (Open Decision 6): a worker serving BOTH a
1987    /// non-isolated namespace and a `Pinned{n1}` namespace on a wrong node is
1988    /// rejected entirely — the compliant namespace does NOT get a partial admit.
1989    #[tokio::test]
1990    async fn pinned_violation_rejects_the_whole_multi_namespace_registration()
1991    -> Result<(), ServerError> {
1992        let store = namespace_store();
1993        let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1994        let (tx, _rx) = mpsc::channel(1);
1995
1996        let denied = registry
1997            .accept_registration(
1998                &guard(),
1999                &multi_caller(&["free", "iso"]),
2000                &registration_full(&["free", "iso"], "", "n2", &["charge"]),
2001                tx,
2002            )
2003            .await;
2004        assert!(
2005            matches!(denied, Err(ServerError::Namespace { .. })),
2006            "a wrong-node worker serving a Pinned namespace fails the WHOLE registration"
2007        );
2008        assert!(
2009            registry
2010                .workers_for("free", DEFAULT_TASK_QUEUE, "charge", None)?
2011                .is_empty(),
2012            "the compliant namespace must NOT be partially admitted"
2013        );
2014        Ok(())
2015    }
2016
2017    /// Unplaced and Prefer namespaces are UNAFFECTED: a node-less worker registers
2018    /// normally (byte-identical to the pre-P2-I1 behaviour). Only Pinned gates.
2019    #[tokio::test]
2020    async fn unplaced_and_prefer_admission_is_unaffected_by_the_pinned_gate()
2021    -> Result<(), ServerError> {
2022        let store = namespace_store();
2023        // `unpl` is left Unplaced (default); `pref` is Prefer{n1}. A node-less
2024        // worker must be admitted into BOTH.
2025        store
2026            .register_namespace("pref", NamespaceOrigin::Explicit)
2027            .await?;
2028        store
2029            .set_namespace_placement(
2030                "pref",
2031                NamespacePlacement::Prefer {
2032                    nodes: ["n1".to_owned()].into_iter().collect(),
2033                },
2034            )
2035            .await?;
2036        let registry = minting_registry(&store, AutoCreate::Open);
2037
2038        let (tx_a, _rx_a) = mpsc::channel(1);
2039        let _reg_a = registry
2040            .accept_registration(
2041                &guard(),
2042                &caller("unpl"),
2043                &registration_full(&["unpl"], "", "", &["charge"]),
2044                tx_a,
2045            )
2046            .await?;
2047        let (tx_b, _rx_b) = mpsc::channel(1);
2048        let _reg_b = registry
2049            .accept_registration(
2050                &guard(),
2051                &caller("pref"),
2052                &registration_full(&["pref"], "", "", &["charge"]),
2053                tx_b,
2054            )
2055            .await?;
2056
2057        assert_eq!(
2058            registry
2059                .workers_for("unpl", DEFAULT_TASK_QUEUE, "charge", None)?
2060                .len(),
2061            1,
2062            "an Unplaced namespace admits a node-less worker unchanged"
2063        );
2064        assert_eq!(
2065            registry
2066                .workers_for("pref", DEFAULT_TASK_QUEUE, "charge", None)?
2067                .len(),
2068            1,
2069            "a Prefer namespace admits a node-less worker unchanged (only Pinned gates)"
2070        );
2071        Ok(())
2072    }
2073
2074    /// A default (no-minter) registry is byte-identical: the placement gate is a
2075    /// no-op with no minter installed, so a node-less worker registers freely even
2076    /// though there is no way to have set a placement in the first place.
2077    #[tokio::test]
2078    async fn no_minter_registry_skips_the_placement_gate() -> Result<(), ServerError> {
2079        let registry = ConnectedWorkerRegistry::default();
2080        let (tx, _rx) = mpsc::channel(1);
2081        let _registration = registry
2082            .accept_registration(
2083                &guard(),
2084                &caller("plain"),
2085                &registration_full(&["plain"], "", "", &["charge"]),
2086                tx,
2087            )
2088            .await?;
2089        assert_eq!(
2090            registry
2091                .workers_for("plain", DEFAULT_TASK_QUEUE, "charge", None)?
2092                .len(),
2093            1,
2094            "with no minter the placement gate is a no-op — registration is unchanged"
2095        );
2096        Ok(())
2097    }
2098
2099    #[tokio::test]
2100    async fn concurrent_registrations_for_a_new_namespace_create_exactly_one_record()
2101    -> Result<(), ServerError> {
2102        let store = namespace_store();
2103        let registry = minting_registry(&store, AutoCreate::Open);
2104
2105        let mut handles = Vec::new();
2106        for _ in 0..8 {
2107            let registry = registry.clone();
2108            handles.push(tokio::spawn(async move {
2109                let (tx, rx) = mpsc::channel(1);
2110                let outcome = registry
2111                    .accept_registration(
2112                        &guard(),
2113                        &caller("rush"),
2114                        &registration("rush", &["charge"]),
2115                        tx,
2116                    )
2117                    .await;
2118                // Keep the receiver alive for the duration of the registration.
2119                drop(rx);
2120                outcome.map(|registration| registration.worker_id())
2121            }));
2122        }
2123        for handle in handles {
2124            handle
2125                .await
2126                .map_err(|_| ServerError::lock_poisoned("registration task panicked"))??;
2127        }
2128
2129        let all = store.list_namespaces().await?;
2130        assert_eq!(
2131            all.iter().filter(|r| r.name == "rush").count(),
2132            1,
2133            "racing registrations must converge on exactly one durable record"
2134        );
2135        Ok(())
2136    }
2137
2138    #[tokio::test]
2139    async fn closed_rejects_unknown_namespace_and_does_not_create_it() -> Result<(), ServerError> {
2140        let store = namespace_store();
2141        let registry = minting_registry(&store, AutoCreate::Closed);
2142        let (tx, _rx) = mpsc::channel(1);
2143
2144        let denied = registry
2145            .accept_registration(
2146                &guard(),
2147                &caller("ghost"),
2148                &registration("ghost", &["charge"]),
2149                tx,
2150            )
2151            .await;
2152        assert!(
2153            matches!(denied, Err(ServerError::Namespace { .. })),
2154            "closed policy must reject an unknown namespace"
2155        );
2156        assert!(
2157            store.get_namespace("ghost").await?.is_none(),
2158            "closed policy must NOT create the namespace it rejected"
2159        );
2160        let tq = DEFAULT_TASK_QUEUE;
2161        assert!(
2162            registry
2163                .workers_for("ghost", tq, "charge", None)?
2164                .is_empty(),
2165            "a rejected registration must not insert a worker"
2166        );
2167        Ok(())
2168    }
2169
2170    #[tokio::test]
2171    async fn closed_admits_a_known_namespace() -> Result<(), ServerError> {
2172        let store = namespace_store();
2173        // Pre-mint the namespace (the POST /namespaces escape hatch's effect).
2174        store
2175            .register_namespace("known", NamespaceOrigin::Explicit)
2176            .await?;
2177        let registry = minting_registry(&store, AutoCreate::Closed);
2178        let (tx, _rx) = mpsc::channel(1);
2179
2180        // Bind the registration token: dropping it deregisters the worker.
2181        let _registration = registry
2182            .accept_registration(
2183                &guard(),
2184                &caller("known"),
2185                &registration("known", &["charge"]),
2186                tx,
2187            )
2188            .await?;
2189        let tq = DEFAULT_TASK_QUEUE;
2190        assert_eq!(
2191            registry.workers_for("known", tq, "charge", None)?.len(),
2192            1,
2193            "a known namespace must register under closed policy"
2194        );
2195        Ok(())
2196    }
2197
2198    #[tokio::test]
2199    async fn no_minter_leaves_registration_untouched() -> Result<(), ServerError> {
2200        // The default registry installs no minter: registration succeeds and
2201        // never touches any namespace registry (byte-identical legacy path).
2202        let registry = ConnectedWorkerRegistry::default();
2203        let (tx, _rx) = mpsc::channel(1);
2204        let _registration = registry
2205            .accept_registration(
2206                &guard(),
2207                &caller("orders"),
2208                &registration("orders", &["charge"]),
2209                tx,
2210            )
2211            .await?;
2212        let tq = DEFAULT_TASK_QUEUE;
2213        assert_eq!(registry.workers_for("orders", tq, "charge", None)?.len(), 1);
2214        Ok(())
2215    }
2216
2217    /// R1 census: an address no worker has ever served reports an empty fleet
2218    /// and no poller age at all — "never seen" is not "seen long ago".
2219    #[test]
2220    fn a_never_served_address_censuses_empty_with_no_poller_age() -> Result<(), ServerError> {
2221        let registry = ConnectedWorkerRegistry::default();
2222        let census = registry.pool_census("default", "general", "greet", None)?;
2223        assert_eq!(census.workers_in_pool, 0);
2224        assert_eq!(census.workers_serving_activity, 0);
2225        assert_eq!(census.compatible_workers, 0);
2226        assert_eq!(census.last_compatible_poller_age, None);
2227        assert!(!census.is_served());
2228        Ok(())
2229    }
2230
2231    /// R1 census: pool membership, activity coverage, and node coverage are
2232    /// three separate counts — that separation is what tells `NO_LIVE_POLLERS`
2233    /// apart from `POLLERS_INCOMPATIBLE`.
2234    #[test]
2235    fn the_census_separates_pool_activity_and_node_coverage() -> Result<(), ServerError> {
2236        let registry = ConnectedWorkerRegistry::default();
2237        let (tx, _rx) = mpsc::channel(1);
2238        let _worker = registry.register_namespaces(
2239            [String::from("default")],
2240            "general",
2241            Some(String::from("n1")),
2242            [String::from("greet")].iter(),
2243            tx,
2244        )?;
2245
2246        let unpinned = registry.pool_census("default", "general", "greet", None)?;
2247        assert_eq!(unpinned.workers_in_pool, 1);
2248        assert_eq!(unpinned.workers_serving_activity, 1);
2249        assert_eq!(unpinned.compatible_workers, 1);
2250        assert_eq!(unpinned.last_compatible_poller_age, Some(Duration::ZERO));
2251
2252        // Same pool, activity nobody advertises: pollers are live but
2253        // incompatible.
2254        let other_activity = registry.pool_census("default", "general", "settle", None)?;
2255        assert_eq!(other_activity.workers_in_pool, 1);
2256        assert_eq!(other_activity.workers_serving_activity, 0);
2257        assert_eq!(other_activity.compatible_workers, 0);
2258
2259        // Same activity, wrong node: coverage exists in the pool but not for
2260        // this dispatch.
2261        let wrong_node = registry.pool_census("default", "general", "greet", Some("n2"))?;
2262        assert_eq!(wrong_node.workers_serving_activity, 1);
2263        assert_eq!(wrong_node.compatible_workers, 0);
2264        assert_eq!(wrong_node.last_compatible_poller_age, None);
2265        Ok(())
2266    }
2267
2268    /// ACCEPTANCE (b) FOR #58, wiring half: a worker whose PROCESS looks alive
2269    /// but which the server cannot REACH on the push leg must lose dispatch
2270    /// eligibility — and losing it must actually remove it from selection.
2271    ///
2272    /// This is the composition the lease split exists for. The tracker half is
2273    /// pinned in `heartbeat::reachability_tests`; this pins the half that turns
2274    /// that verdict into a dispatch decision, using the SAME two types the
2275    /// liveness probe wires together — a real [`HeartbeatTracker`] census
2276    /// feeding a real [`ConnectedWorkerRegistry`].
2277    ///
2278    /// The staging is Finding B exactly: the worker's own liveness pump keeps
2279    /// beating (`record_connection_activity`) well past the window, so every
2280    /// signal that says "this process is alive" says so — and NO ping is ever
2281    /// answered, so the one signal that says "the server can reach it" is
2282    /// absent. Before the split those were one lease and the pump's weaker
2283    /// positive evidence buried the ping's negative evidence.
2284    ///
2285    /// The control is the second half of the test: the same worker, reachable,
2286    /// must still be selected. Without it a registry that selected NOBODY would
2287    /// satisfy the first assertion and prove nothing.
2288    #[test]
2289    fn a_pump_alive_worker_the_server_cannot_reach_is_not_selected() -> Result<(), ServerError> {
2290        const WINDOW: Duration = Duration::from_secs(30);
2291
2292        let registry = ConnectedWorkerRegistry::default();
2293        let tracker = HeartbeatTracker::new(WINDOW);
2294        let (tx, _rx) = mpsc::channel(1);
2295        let worker = registry.register_namespaces(
2296            [String::from("default")],
2297            "general",
2298            None,
2299            [String::from("greet")].iter(),
2300            tx,
2301        )?;
2302        let Some(worker_id) = worker.worker_id() else {
2303            return Err(test_failure("registration carries an id"));
2304        };
2305        let start = Instant::now();
2306        tracker.register_connection(worker_id, start)?;
2307
2308        // Baseline: the worker SERVES ITS PROBATION and is therefore selectable.
2309        // Registration alone grants nothing — the handshake ack is sent, never
2310        // confirmed received — so eligibility here is earned by answered pings.
2311        // If this did not hold, no worker could ever be dispatched to.
2312        for _ in 0..DISPATCH_PROBATION_PINGS {
2313            assert!(
2314                tracker.record_dispatch_reachability(worker_id, start)?,
2315                "the worker is tracked while it serves its probation"
2316            );
2317        }
2318        registry
2319            .set_dispatch_ineligible(tracker.unreachable_workers(start)?.into_iter().collect())?;
2320        assert!(
2321            registry
2322                .select_worker("default", "general", "greet", None)?
2323                .is_some(),
2324            "a worker that has answered a full run of pings must be selectable: this is the \
2325             control, and without it a registry selecting NOBODY would satisfy every assertion \
2326             below"
2327        );
2328
2329        // Now the poisoned-connection shape: the pump beats past the window and
2330        // not one ping is answered. Each unanswered probe is recorded exactly as
2331        // the liveness probe records it, because that is the trajectory a real
2332        // poisoned connection takes — staleness alone would be a shape only a
2333        // stopped probe could produce.
2334        let much_later = start + WINDOW * 4;
2335        assert!(
2336            tracker.record_connection_activity(worker_id, much_later)?,
2337            "the worker is still tracked; its process is plainly alive"
2338        );
2339        assert!(
2340            tracker.record_dispatch_unreachable(worker_id)?,
2341            "the probe fired and went unanswered"
2342        );
2343        registry.set_dispatch_ineligible(
2344            tracker
2345                .unreachable_workers(much_later)?
2346                .into_iter()
2347                .collect(),
2348        )?;
2349
2350        assert!(
2351            registry.is_dispatch_ineligible(worker_id)?,
2352            "a worker the server cannot reach must be marked ineligible however alive its \
2353             process looks"
2354        );
2355        assert!(
2356            registry
2357                .select_worker("default", "general", "greet", None)?
2358                .is_none(),
2359            "an ineligible worker must not be SELECTED: selecting one produces a dispatch that \
2360             can only fail, and on the liminal transport it fails by consuming connection \
2361             capacity — making the unreachability worse"
2362        );
2363
2364        // The recovery path: exclusion must be WITHDRAWABLE, or a single bad
2365        // round would strand a healthy worker forever. It is withdrawn by a
2366        // served probation, not by one answer — and the intermediate assertion
2367        // below pins that distinction rather than assuming it.
2368        assert!(
2369            tracker.record_dispatch_reachability(worker_id, much_later)?,
2370            "the worker is still tracked"
2371        );
2372        registry.set_dispatch_ineligible(
2373            tracker
2374                .unreachable_workers(much_later)?
2375                .into_iter()
2376                .collect(),
2377        )?;
2378        assert!(
2379            registry.is_dispatch_ineligible(worker_id)?,
2380            "one answer part-way through a fresh probation must NOT restore eligibility: a link \
2381             answering one probe in three would otherwise flap in and out of selection"
2382        );
2383
2384        for _ in 1..DISPATCH_PROBATION_PINGS {
2385            assert!(
2386                tracker.record_dispatch_reachability(worker_id, much_later)?,
2387                "the worker is still tracked"
2388            );
2389        }
2390        registry.set_dispatch_ineligible(
2391            tracker
2392                .unreachable_workers(much_later)?
2393                .into_iter()
2394                .collect(),
2395        )?;
2396        assert!(
2397            !registry.is_dispatch_ineligible(worker_id)?,
2398            "an answered ping must clear the exclusion"
2399        );
2400        assert!(
2401            registry
2402                .select_worker("default", "general", "greet", None)?
2403                .is_some(),
2404            "and the worker must be selectable again"
2405        );
2406        Ok(())
2407    }
2408
2409    /// R1 census: after the last compatible worker leaves, the address reports
2410    /// how long ago it was served — the age an operator sees on the parked
2411    /// dispatch's WARN.
2412    #[test]
2413    fn a_departed_worker_leaves_a_last_compatible_poller_age() -> Result<(), ServerError> {
2414        let registry = ConnectedWorkerRegistry::default();
2415        let (tx, _rx) = mpsc::channel(1);
2416        let worker = registry.register_namespaces(
2417            [String::from("default")],
2418            "general",
2419            None,
2420            [String::from("greet")].iter(),
2421            tx,
2422        )?;
2423        worker.deregister()?;
2424
2425        let census = registry.pool_census("default", "general", "greet", None)?;
2426        assert_eq!(census.workers_in_pool, 0);
2427        assert_eq!(census.compatible_workers, 0);
2428        let age = census
2429            .last_compatible_poller_age
2430            .ok_or_else(|| test_failure("a departed worker must leave an age behind"))?;
2431        assert!(
2432            age < Duration::from_secs(60),
2433            "the recorded departure is implausibly old: {age:?}"
2434        );
2435
2436        // A worker returning to the address clears the departure record: the
2437        // map that answers "how long ago" never accumulates live addresses.
2438        let (tx, _rx) = mpsc::channel(1);
2439        let _back = registry.register_namespaces(
2440            [String::from("default")],
2441            "general",
2442            None,
2443            [String::from("greet")].iter(),
2444            tx,
2445        )?;
2446        let served = registry.pool_census("default", "general", "greet", None)?;
2447        assert_eq!(served.last_compatible_poller_age, Some(Duration::ZERO));
2448        assert!(served.is_served());
2449        Ok(())
2450    }
2451}