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