Skip to main content

subc_client_rs/
lib.rs

1#![forbid(unsafe_code)]
2
3pub mod consumer;
4pub use consumer::{
5    CallError, CallOptions, CatalogList, CloseRouteOptions, ConnectionState, ConsumerError,
6    ConsumerOptions, PushEvent, RetryBackoff, RoutePollResult, SubcConsumer, SubscribeOptions,
7    Subscription, SubscriptionClosed,
8};
9
10use std::{
11    collections::HashMap,
12    env,
13    error::Error,
14    ffi::OsString,
15    fmt,
16    future::Future,
17    io,
18    path::{Path, PathBuf},
19    pin::Pin,
20    sync::{
21        atomic::{AtomicU64, Ordering},
22        Arc, Mutex,
23    },
24    time::Duration,
25};
26
27pub use async_trait::async_trait;
28pub use subc_control::{CatalogEntry, ConsumerIdentity};
29use subc_protocol::{
30    manifest::ModuleManifest,
31    session::{
32        ModuleControlRequest, ModuleControlRequestFromModule, ModuleControlResponse,
33        ModuleControlResponseToModule, MODULE_CONTROL_OP_HEALTH_CHECK,
34        MODULE_TO_SUBC_OP_CATALOG_UPDATE,
35    },
36    BindIdentity, ErrorBody, Flags, Frame, FrameBuildError, FrameType, ModuleHelloAckBody,
37    ModuleHelloBody, Principal, Priority, RouteTarget, PROTOCOL_VERSION, SUBC_LAUNCH_NONCE_ENV,
38    SUBC_MODULE_ID_ENV,
39};
40pub use subc_protocol::{
41    manifest::{ExecutionMode, ProviderRole, Tool},
42    session::{HealthReport, HealthStatus},
43    AdmissionClass,
44};
45use subc_transport::{
46    authenticate_client, connection_file, read_frame, write_frame, AuthError, ConnectionFileError,
47    FrameIoError,
48};
49use tokio::{
50    io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufWriter},
51    net::TcpStream,
52    sync::{mpsc, oneshot, Semaphore},
53    time::timeout,
54};
55use tokio_util::sync::{CancellationToken, WaitForCancellationFuture};
56
57const AUTH_DEADLINE: Duration = Duration::from_secs(2);
58const CATALOG_UPDATE_TIMEOUT: Duration = Duration::from_secs(10);
59const EGRESS_BUFFER: usize = 64;
60const HANDLER_TASK_CAPACITY: usize = 64;
61const HELLO_CORR: u64 = 1;
62static NEXT_MODULE_CONNECTION_TOKEN: AtomicU64 = AtomicU64::new(1);
63
64type RequestKey = (u16, u32, u64);
65type InFlight = Arc<Mutex<HashMap<RequestKey, CancellationToken>>>;
66
67/// Immutable identity of one route binding on one live connection.
68///
69/// Only `channel` and `epoch` are serialized. The private connection token prevents
70/// work retained from an earlier connection from acting on a later connection that
71/// happens to reuse the same wire pair.
72#[derive(Clone, Copy, PartialEq, Eq, Hash)]
73pub struct RouteHandle {
74    pub channel: u16,
75    pub epoch: u32,
76    connection_token: u64,
77}
78
79impl RouteHandle {
80    pub(crate) fn new(channel: u16, epoch: u32, connection_token: u64) -> Self {
81        Self {
82            channel,
83            epoch,
84            connection_token,
85        }
86    }
87
88    pub(crate) fn connection_token(self) -> u64 {
89        self.connection_token
90    }
91}
92
93impl fmt::Debug for RouteHandle {
94    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95        f.debug_struct("RouteHandle")
96            .field("channel", &self.channel)
97            .field("epoch", &self.epoch)
98            .finish_non_exhaustive()
99    }
100}
101type CatalogUpdateReply = oneshot::Sender<Result<(), CatalogUpdateError>>;
102type CatalogUpdateWaiter = oneshot::Receiver<Result<(), CatalogUpdateError>>;
103type CatalogUpdateRequest = (u64, mpsc::Sender<Frame>, CatalogUpdateWaiter);
104
105/// Future returned by [`serve_with_handle`] that runs the module until GOODBYE or EOF.
106pub type ModuleServeFuture = Pin<Box<dyn Future<Output = Result<(), SubcModuleError>> + Send>>;
107
108#[derive(Clone)]
109struct RequestDispatcher {
110    in_flight: InFlight,
111    permits: Arc<Semaphore>,
112}
113
114impl RequestDispatcher {
115    fn new() -> Self {
116        Self {
117            in_flight: Arc::new(Mutex::new(HashMap::new())),
118            permits: Arc::new(Semaphore::new(HANDLER_TASK_CAPACITY)),
119        }
120    }
121}
122
123/// Cloneable handle for module-originated control RPCs on channel 0.
124#[derive(Clone)]
125pub struct ModuleHandle {
126    shared: Arc<ModuleHandleShared>,
127}
128
129struct ModuleHandleShared {
130    negotiated_ver: u8,
131    supports_catalog_update: bool,
132    connection_token: u64,
133    live_routes: Mutex<HashMap<u16, RouteHandle>>,
134    dropped_route_frames: AtomicU64,
135    close_token: CancellationToken,
136    inner: Mutex<ModuleHandleState>,
137}
138
139struct ModuleHandleState {
140    writer: Option<mpsc::Sender<Frame>>,
141    next_corr: Option<u64>,
142    pending_catalog_updates: HashMap<u64, CatalogUpdateReply>,
143    closed: bool,
144}
145
146impl ModuleHandle {
147    fn new(
148        ack: &ModuleHelloAckBody,
149        writer: mpsc::Sender<Frame>,
150        connection_token: u64,
151        close_token: CancellationToken,
152    ) -> Self {
153        Self {
154            shared: Arc::new(ModuleHandleShared {
155                negotiated_ver: ack.negotiated_ver,
156                supports_catalog_update: ack
157                    .subc_ops
158                    .iter()
159                    .any(|op| op == MODULE_TO_SUBC_OP_CATALOG_UPDATE),
160                connection_token,
161                live_routes: Mutex::new(HashMap::new()),
162                dropped_route_frames: AtomicU64::new(0),
163                close_token,
164                inner: Mutex::new(ModuleHandleState {
165                    writer: Some(writer),
166                    next_corr: Some(HELLO_CORR + 1),
167                    pending_catalog_updates: HashMap::new(),
168                    closed: false,
169                }),
170            }),
171        }
172    }
173
174    /// Ask the daemon to replace this module's advertised provider roles in place.
175    ///
176    /// The returned result resolves when the daemon ACKs the update, rejects it with
177    /// a typed channel-0 Error frame, the request times out, or the connection dies.
178    pub async fn catalog_update(
179        &self,
180        provides: Vec<ProviderRole>,
181    ) -> Result<(), CatalogUpdateError> {
182        if !self.shared.supports_catalog_update {
183            return Err(CatalogUpdateError::NotSupported);
184        }
185
186        let body = serde_json::to_vec(&ModuleControlRequestFromModule::CatalogUpdate { provides })
187            .map_err(|err| {
188                CatalogUpdateError::Protocol(format!(
189                    "failed to encode catalog.update request body: {err}"
190                ))
191            })?;
192        let (corr, writer, rx) = self.shared.begin_catalog_update()?;
193        let frame = Frame::build_with_version(
194            self.shared.negotiated_ver,
195            FrameType::Request,
196            control_flags(),
197            0,
198            0,
199            corr,
200            body,
201        )
202        .map_err(|err| {
203            self.shared.remove_pending_catalog_update(corr);
204            CatalogUpdateError::Protocol(format!(
205                "failed to build catalog.update request frame: {err}"
206            ))
207        })?;
208
209        if writer.send(frame).await.is_err() {
210            self.shared.remove_pending_catalog_update(corr);
211            return Err(CatalogUpdateError::ConnectionClosed);
212        }
213
214        match timeout(CATALOG_UPDATE_TIMEOUT, rx).await {
215            Ok(Ok(result)) => result,
216            Ok(Err(_)) => Err(CatalogUpdateError::ConnectionClosed),
217            Err(_) => {
218                self.shared.remove_pending_catalog_update(corr);
219                Err(CatalogUpdateError::Timeout)
220            }
221        }
222    }
223
224    /// Emit an uncorrelated Push on a live route.
225    pub async fn push(
226        &self,
227        handle: &RouteHandle,
228        body: Vec<u8>,
229        admission_class: Option<AdmissionClass>,
230    ) -> Result<(), SubcModuleError> {
231        self.validate_route(*handle)?;
232        let writer = self
233            .shared
234            .lock_inner()
235            .writer
236            .clone()
237            .ok_or(SubcModuleError::WriterClosed)?;
238        let frame = Frame::build_with_version(
239            self.shared.negotiated_ver,
240            FrameType::Push,
241            data_flags().with_admission_class(admission_class.unwrap_or(AdmissionClass::Normal)),
242            handle.channel,
243            handle.epoch,
244            0,
245            body,
246        )
247        .map_err(SubcModuleError::FrameBuild)?;
248        send_outbound(&writer, frame).await
249    }
250
251    /// Number of unknown or stale route frames silently dropped by endpoint validation.
252    pub fn dropped_route_frames(&self) -> u64 {
253        self.shared.dropped_route_frames.load(Ordering::Relaxed)
254    }
255
256    fn validate_route(&self, handle: RouteHandle) -> Result<(), SubcModuleError> {
257        if handle.connection_token() != self.shared.connection_token {
258            return Err(SubcModuleError::StaleRouteHandle(handle));
259        }
260        let routes = self
261            .shared
262            .live_routes
263            .lock()
264            .map_err(|_| SubcModuleError::InFlightPoisoned)?;
265        if routes.get(&handle.channel) == Some(&handle) {
266            Ok(())
267        } else {
268            Err(SubcModuleError::StaleRouteHandle(handle))
269        }
270    }
271
272    fn install_route(&self, handle: RouteHandle) -> Result<(), SubcModuleError> {
273        self.shared
274            .live_routes
275            .lock()
276            .map_err(|_| SubcModuleError::InFlightPoisoned)?
277            .insert(handle.channel, handle);
278        Ok(())
279    }
280
281    fn installed_route(&self, channel: u16) -> Result<Option<RouteHandle>, SubcModuleError> {
282        Ok(self
283            .shared
284            .live_routes
285            .lock()
286            .map_err(|_| SubcModuleError::InFlightPoisoned)?
287            .get(&channel)
288            .copied())
289    }
290
291    fn remove_route(&self, handle: RouteHandle) -> Result<bool, SubcModuleError> {
292        let mut routes = self
293            .shared
294            .live_routes
295            .lock()
296            .map_err(|_| SubcModuleError::InFlightPoisoned)?;
297        if routes.get(&handle.channel) == Some(&handle) {
298            routes.remove(&handle.channel);
299            Ok(true)
300        } else {
301            Ok(false)
302        }
303    }
304
305    fn validate_ingress(&self, channel: u16, epoch: u32) -> Result<bool, SubcModuleError> {
306        let handle = RouteHandle::new(channel, epoch, self.shared.connection_token);
307        let valid = self
308            .shared
309            .live_routes
310            .lock()
311            .map_err(|_| SubcModuleError::InFlightPoisoned)?
312            .get(&channel)
313            == Some(&handle);
314        if !valid {
315            self.shared
316                .dropped_route_frames
317                .fetch_add(1, Ordering::Relaxed);
318        }
319        Ok(valid)
320    }
321
322    fn route_handle(&self, channel: u16, epoch: u32) -> RouteHandle {
323        RouteHandle::new(channel, epoch, self.shared.connection_token)
324    }
325
326    fn handle_control_reply(&self, frame: Frame) -> bool {
327        let Some(reply) = self.shared.take_pending_catalog_update(frame.header.corr) else {
328            return false;
329        };
330        let result = match frame.header.ty {
331            FrameType::Response => {
332                match serde_json::from_slice::<ModuleControlResponseToModule>(&frame.body) {
333                    Ok(ModuleControlResponseToModule::CatalogUpdate {}) => Ok(()),
334                    Err(err) => Err(CatalogUpdateError::Protocol(format!(
335                        "invalid catalog.update response body: {err}"
336                    ))),
337                }
338            }
339            FrameType::Error => match serde_json::from_slice::<ErrorBody>(&frame.body) {
340                Ok(body) => Err(match body.code.as_str() {
341                    "catalog_update_frozen_field" => CatalogUpdateError::FrozenField(body),
342                    "not_registered" => CatalogUpdateError::NotRegistered(body),
343                    _ => CatalogUpdateError::Rejected(body),
344                }),
345                Err(err) => Err(CatalogUpdateError::Protocol(format!(
346                    "invalid catalog.update error body: {err}"
347                ))),
348            },
349            ty => Err(CatalogUpdateError::Protocol(format!(
350                "unexpected catalog.update terminal frame: {ty:?}"
351            ))),
352        };
353        let _ = reply.send(result);
354        true
355    }
356
357    fn close_connection(&self) {
358        self.shared.close_connection();
359    }
360}
361
362impl ModuleHandleShared {
363    fn begin_catalog_update(&self) -> Result<CatalogUpdateRequest, CatalogUpdateError> {
364        let mut inner = self.lock_inner();
365        if inner.closed {
366            return Err(CatalogUpdateError::ConnectionClosed);
367        }
368        let Some(writer) = inner.writer.clone() else {
369            inner.closed = true;
370            self.close_token.cancel();
371            drop(inner);
372            self.clear_live_routes();
373            return Err(CatalogUpdateError::ConnectionClosed);
374        };
375        let Some(corr) = next_module_control_corr(&mut inner) else {
376            inner.closed = true;
377            inner.writer = None;
378            let pending = inner
379                .pending_catalog_updates
380                .drain()
381                .map(|(_, reply)| reply)
382                .collect::<Vec<_>>();
383            self.close_token.cancel();
384            drop(inner);
385            self.clear_live_routes();
386            for reply in pending {
387                let _ = reply.send(Err(CatalogUpdateError::ConnectionClosed));
388            }
389            return Err(CatalogUpdateError::ConnectionClosed);
390        };
391        let (tx, rx) = oneshot::channel();
392        inner.pending_catalog_updates.insert(corr, tx);
393        Ok((corr, writer, rx))
394    }
395
396    fn take_pending_catalog_update(&self, corr: u64) -> Option<CatalogUpdateReply> {
397        self.lock_inner().pending_catalog_updates.remove(&corr)
398    }
399
400    fn remove_pending_catalog_update(&self, corr: u64) {
401        self.lock_inner().pending_catalog_updates.remove(&corr);
402    }
403
404    fn close_connection(&self) {
405        let pending = {
406            let mut inner = self.lock_inner();
407            if inner.closed {
408                return;
409            }
410            inner.closed = true;
411            inner.writer = None;
412            self.close_token.cancel();
413            inner
414                .pending_catalog_updates
415                .drain()
416                .map(|(_, reply)| reply)
417                .collect::<Vec<_>>()
418        };
419        self.clear_live_routes();
420        for reply in pending {
421            let _ = reply.send(Err(CatalogUpdateError::ConnectionClosed));
422        }
423    }
424
425    fn clear_live_routes(&self) {
426        if let Ok(mut routes) = self.live_routes.lock() {
427            routes.clear();
428        }
429    }
430
431    fn lock_inner(&self) -> std::sync::MutexGuard<'_, ModuleHandleState> {
432        self.inner
433            .lock()
434            .unwrap_or_else(|poisoned| poisoned.into_inner())
435    }
436}
437
438fn next_module_control_corr(inner: &mut ModuleHandleState) -> Option<u64> {
439    let corr = inner.next_corr?;
440    inner.next_corr = corr.checked_add(1);
441    Some(corr)
442}
443
444/// Errors returned by [`ModuleHandle::catalog_update`].
445#[derive(Debug)]
446pub enum CatalogUpdateError {
447    NotSupported,
448    FrozenField(ErrorBody),
449    NotRegistered(ErrorBody),
450    Rejected(ErrorBody),
451    Timeout,
452    ConnectionClosed,
453    Protocol(String),
454}
455
456impl fmt::Display for CatalogUpdateError {
457    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
458        match self {
459            Self::NotSupported => write!(
460                f,
461                "daemon HELLO_ACK did not advertise catalog.update support"
462            ),
463            Self::FrozenField(body) => {
464                write!(f, "catalog.update rejected frozen field: {}", body.message)
465            }
466            Self::NotRegistered(body) => write!(
467                f,
468                "catalog.update requires a registered module: {}",
469                body.message
470            ),
471            Self::Rejected(body) => write!(
472                f,
473                "catalog.update rejected by subc: {} ({})",
474                body.code, body.message
475            ),
476            Self::Timeout => write!(f, "catalog.update timed out waiting for an ACK"),
477            Self::ConnectionClosed => {
478                write!(f, "subc connection closed before catalog.update completed")
479            }
480            Self::Protocol(message) => write!(f, "catalog.update protocol error: {message}"),
481        }
482    }
483}
484
485impl Error for CatalogUpdateError {}
486
487/// Trait implemented by a module for its business logic. The serve functions in
488/// this crate own all wire-protocol plumbing.
489#[async_trait]
490pub trait ModuleHandler: Send + Sync + 'static {
491    /// Handle a data-plane request on a route channel. Return a unary response, a
492    /// typed error, or stream interim events via [`RequestCtx::emit`] and return
493    /// [`HandlerOutcome::Streamed`]. Each request runs in its own task so one slow
494    /// handler cannot head-of-line-block another route.
495    async fn handle(&self, ctx: RequestCtx, body: Vec<u8>) -> HandlerOutcome;
496
497    /// Called once after HELLO_ACK so the module can inspect the ack body for
498    /// negotiated capabilities and any storage descriptor supplied by the daemon.
499    async fn on_hello_ack(&self, _ack: &ModuleHelloAckBody) {}
500
501    /// Decide a route.bind. This hook is decision-only and must not emit route traffic.
502    async fn on_bind(&self, _req: &RouteBindRequest) -> BindDecision {
503        BindDecision::accept()
504    }
505
506    /// Called after an accepted bind ACK is queued and the handle is installed.
507    async fn on_bound(&self, _handle: &RouteHandle) {}
508
509    /// Return cheap in-memory health for the module.
510    ///
511    /// THE DEFAULT ASSERTS HEALTH ON BEHALF OF A MODULE THAT NEVER WROTE ANY.
512    /// A module that has not implemented this is indistinguishable on the wire
513    /// from one that measured itself and found nothing wrong -- and the daemon
514    /// acts on the difference, since a healthy report suppresses escalation
515    /// while an absent implementation means nothing was ever checked.
516    ///
517    /// It stays a default because health is genuinely optional: a module that
518    /// advertises no health capability is never probed, so the value is unread
519    /// for those. The hazard is the module that DOES advertise health and
520    /// inherits this -- it answers "ok" forever, including while wedged.
521    ///
522    /// Per Health-Path-Rule v3 an implementation must derive its status
523    /// mechanically from signals the dispatch path stamps (a monotonic
524    /// heartbeat, oldest-queued age), never from its own opinion, and must not
525    /// take a blocking lock, touch disk, or spawn a subprocess on this path.
526    /// A health reply that execs queues behind the host's slowest shared
527    /// resource -- which is exactly the resource degraded under the conditions
528    /// being probed.
529    async fn health(&self) -> HealthReport {
530        // SAY THAT NOBODY MEASURED, rather than that everything is fine.
531        //
532        // The status stays Ok because a module advertising no health capability
533        // is never probed, and one that advertises health but has nothing to
534        // report is not unhealthy. What changes is that the report now
535        // IDENTIFIES ITSELF as the inherited default, so an operator reading
536        // `ck health <module>` can tell "measured, nothing wrong" from "nobody
537        // wrote a health path" -- which were previously the same bytes.
538        //
539        // `detail` is carried verbatim by the daemon and rendered for humans;
540        // nothing parses it, so this is display-only and cannot change any
541        // supervision decision.
542        HealthReport {
543            detail: Some("no health implementation; inherited default".to_string()),
544            ..HealthReport::ok()
545        }
546    }
547
548    /// A route was torn down, rejected, or abandoned before its bind ACK was queued.
549    async fn on_route_gone(&self, _handle: &RouteHandle) {}
550}
551
552/// The terminal result of a module request handler.
553#[derive(Debug, Clone, PartialEq, Eq)]
554pub enum HandlerOutcome {
555    /// Send a Response frame carrying these bytes.
556    Response(Vec<u8>),
557    /// Send an Error frame carrying an [`ErrorBody`] with this code and message.
558    Error { code: String, message: String },
559    /// The handler emitted stream data with [`RequestCtx::emit`]; the serve code
560    /// sends the StreamEnd terminal frame.
561    Streamed,
562}
563
564/// Per-request context. Retains the full route handle and correlation id, provides
565/// interim stream emission, and exposes a cancellation signal.
566#[derive(Clone)]
567pub struct RequestCtx {
568    handle: RouteHandle,
569    corr: u64,
570    ver: u8,
571    egress: mpsc::Sender<Frame>,
572    module_handle: ModuleHandle,
573    cancelled: CancellationToken,
574}
575
576impl RequestCtx {
577    /// Full route handle retained from ingress.
578    pub fn route_handle(&self) -> RouteHandle {
579        self.handle
580    }
581
582    /// Correlation id for this request.
583    pub fn corr(&self) -> u64 {
584        self.corr
585    }
586
587    /// Emit an interim StreamData frame on this request's `(channel, corr)`. Once
588    /// the request is cancelled or its route is gone, late emits are dropped.
589    pub async fn emit(&self, body: Vec<u8>) -> Result<(), SubcModuleError> {
590        self.emit_with_admission(body, None).await
591    }
592
593    /// Emit StreamData with an explicit admission class. `None` means NORMAL.
594    pub async fn emit_with_admission(
595        &self,
596        body: Vec<u8>,
597        admission_class: Option<AdmissionClass>,
598    ) -> Result<(), SubcModuleError> {
599        self.module_handle.validate_route(self.handle)?;
600        if self.cancelled.is_cancelled() {
601            return Ok(());
602        }
603        self.send_frame(
604            FrameType::StreamData,
605            data_flags().with_admission_class(admission_class.unwrap_or(AdmissionClass::Normal)),
606            body,
607        )
608        .await
609    }
610
611    /// Completes when the other side sends Cancel for this request or the route
612    /// is torn down.
613    pub fn cancelled(&self) -> WaitForCancellationFuture<'_> {
614        self.cancelled.cancelled()
615    }
616
617    /// Return a cloneable cancellation token for code that prefers token polling.
618    pub fn cancellation_token(&self) -> CancellationToken {
619        self.cancelled.clone()
620    }
621
622    async fn send_frame(
623        &self,
624        frame_type: FrameType,
625        flags: Flags,
626        body: Vec<u8>,
627    ) -> Result<(), SubcModuleError> {
628        self.module_handle.validate_route(self.handle)?;
629        let frame = Frame::build_with_version(
630            self.ver,
631            frame_type,
632            flags,
633            self.handle.channel,
634            self.handle.epoch,
635            self.corr,
636            body,
637        )
638        .map_err(SubcModuleError::FrameBuild)?;
639        send_outbound(&self.egress, frame).await
640    }
641}
642
643/// Route-bind request delivered on channel 0.
644#[derive(Debug, Clone)]
645pub struct RouteBindRequest {
646    pub handle: RouteHandle,
647    pub target: RouteTarget,
648    pub identity: BindIdentity,
649    pub principal: Option<Principal>,
650    /// Consumer-declared reverse-request capabilities for this bind. This is a
651    /// declaration, not a verified privilege; providers treat an absent field as
652    /// no reverse-request capability. Known MCP method-family values today are
653    /// "elicitation", "sampling", and "roots".
654    pub consumer_capabilities: Option<Vec<String>>,
655    /// Opaque admission facts relayed by subc from its configured carrier.
656    pub admission_facts: Option<serde_json::Value>,
657}
658
659/// Decision returned by [`ModuleHandler::on_bind`].
660#[derive(Debug, Clone)]
661pub struct BindDecision {
662    kind: BindDecisionKind,
663}
664
665impl BindDecision {
666    /// Accept the route.bind request.
667    pub fn accept() -> Self {
668        Self {
669            kind: BindDecisionKind::Accept,
670        }
671    }
672
673    /// Reject the route.bind request with a typed Error frame.
674    pub fn reject(code: impl Into<String>, message: impl Into<String>) -> Self {
675        Self {
676            kind: BindDecisionKind::Reject {
677                code: code.into(),
678                message: message.into(),
679            },
680        }
681    }
682}
683
684#[derive(Debug, Clone)]
685enum BindDecisionKind {
686    Accept,
687    Reject { code: String, message: String },
688}
689
690/// Run a module to completion. Reads `--subc <connection-file>` from args, uses
691/// `SUBC_MODULE_ID` when set by the process that launched the module, connects,
692/// authenticates, sends HELLO, waits for HELLO_ACK, then serves frames until
693/// GOODBYE or clean EOF.
694pub async fn serve<H>(mut manifest: ModuleManifest, handler: H) -> Result<(), SubcModuleError>
695where
696    H: ModuleHandler,
697{
698    let connection_file = parse_subc_arg(env::args_os().skip(1))?;
699    if let Some(module_id) = module_id_from_env()? {
700        manifest.module_id = module_id;
701    }
702    serve_with(&connection_file, manifest, handler).await
703}
704
705/// Run a module with an explicit connection-file path. The manifest is sent as
706/// provided; callers that need a nonstandard module id should set it before calling.
707pub async fn serve_with<H>(
708    connection_file: &Path,
709    manifest: ModuleManifest,
710    handler: H,
711) -> Result<(), SubcModuleError>
712where
713    H: ModuleHandler,
714{
715    let (_handle, serve_future) = serve_with_handle(connection_file, manifest, handler).await?;
716    serve_future.await
717}
718
719/// Connect, register the module, and return a cloneable handle plus the future that
720/// must be awaited or spawned to keep serving the connection.
721pub async fn serve_with_handle<H>(
722    connection_file: &Path,
723    manifest: ModuleManifest,
724    handler: H,
725) -> Result<(ModuleHandle, ModuleServeFuture), SubcModuleError>
726where
727    H: ModuleHandler,
728{
729    let stream = connect_to_subc(connection_file).await?;
730    let (mut read_half, write_half) = tokio::io::split(stream);
731    let (tx, rx) = mpsc::channel::<Frame>(EGRESS_BUFFER);
732    let writer = tokio::spawn(drain_writer(write_half, rx));
733    let handler = Arc::new(handler);
734
735    send_hello(&tx, manifest).await?;
736    let ack = expect_hello_ack(&mut read_half).await?;
737    handler.on_hello_ack(&ack).await;
738
739    let connection_token = NEXT_MODULE_CONNECTION_TOKEN
740        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |token| {
741            token.checked_add(1)
742        })
743        .map_err(|_| SubcModuleError::ConnectionTokenExhausted)?;
744    let close_token = CancellationToken::new();
745    let handle = ModuleHandle::new(&ack, tx.clone(), connection_token, close_token);
746    let serve_handle = handle.clone();
747    let serve_future = Box::pin(async move {
748        // Connection loss ends this serve future. Module serving retains no
749        // reconnect task or in-flight reconnect gate; a supervisor that needs
750        // recovery starts a fresh serve_with_handle invocation.
751        let loop_result =
752            module_loop(read_half, tx, Arc::clone(&handler), serve_handle.clone()).await;
753        serve_handle.close_connection();
754
755        let writer_result = writer.await.map_err(SubcModuleError::WriterTask);
756        match (loop_result, writer_result) {
757            (Err(loop_err), _) => Err(loop_err),
758            (Ok(()), Ok(Ok(()))) => Ok(()),
759            // The read loop already saw the daemon go away; the writer failing
760            // to flush its remaining frames to that dead socket (BrokenPipe on
761            // Unix, ConnectionReset on Windows) is part of the same terminal,
762            // not a distinct fault.
763            (Ok(()), Ok(Err(FrameIoError::Io(err))))
764                if matches!(
765                    err.kind(),
766                    std::io::ErrorKind::ConnectionReset
767                        | std::io::ErrorKind::ConnectionAborted
768                        | std::io::ErrorKind::BrokenPipe
769                ) =>
770            {
771                Ok(())
772            }
773            (Ok(()), Ok(Err(writer_err))) => Err(SubcModuleError::FrameIo(writer_err)),
774            (Ok(()), Err(join_err)) => Err(join_err),
775        }
776    });
777    Ok((handle, serve_future))
778}
779
780async fn module_loop<R, H>(
781    mut reader: R,
782    egress: mpsc::Sender<Frame>,
783    handler: Arc<H>,
784    module_handle: ModuleHandle,
785) -> Result<(), SubcModuleError>
786where
787    R: AsyncRead + Unpin,
788    H: ModuleHandler,
789{
790    let dispatcher = RequestDispatcher::new();
791    loop {
792        let read = tokio::select! {
793            () = module_handle.shared.close_token.cancelled() => return Ok(()),
794            read = read_frame(&mut reader) => read,
795        };
796        let frame = match read {
797            Ok(Some(frame)) => frame,
798            // Clean EOF: the daemon closed the connection.
799            Ok(None) => return Ok(()),
800            // A reset/abort on the read path also means the daemon is gone. On
801            // Unix a killed daemon closes the socket with FIN (clean EOF above),
802            // but Windows sends RST on process death, surfacing here as
803            // ConnectionReset. Both are the same "serve until the daemon goes
804            // away" terminal, so normalize to a clean exit for a
805            // platform-independent serve() contract.
806            Err(FrameIoError::Io(err))
807                if matches!(
808                    err.kind(),
809                    std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted
810                ) =>
811            {
812                return Ok(());
813            }
814            Err(err) => return Err(SubcModuleError::FrameIo(err)),
815        };
816        if !handle_frame(
817            frame,
818            &egress,
819            Arc::clone(&handler),
820            dispatcher.clone(),
821            module_handle.clone(),
822        )
823        .await?
824        {
825            return Ok(());
826        }
827    }
828}
829
830async fn handle_frame<H>(
831    frame: Frame,
832    egress: &mpsc::Sender<Frame>,
833    handler: Arc<H>,
834    dispatcher: RequestDispatcher,
835    module_handle: ModuleHandle,
836) -> Result<bool, SubcModuleError>
837where
838    H: ModuleHandler,
839{
840    if frame.header.channel != 0
841        && !module_handle.validate_ingress(frame.header.channel, frame.header.epoch)?
842    {
843        return Ok(true);
844    }
845    match frame.header.ty {
846        FrameType::Ping if frame.header.channel == 0 => {
847            let pong = Frame::build_with_version(
848                frame.header.ver,
849                FrameType::Pong,
850                frame.header.flags,
851                0,
852                0,
853                frame.header.corr,
854                Vec::new(),
855            )
856            .map_err(SubcModuleError::FrameBuild)?;
857            send_outbound(egress, pong).await?;
858            Ok(true)
859        }
860        FrameType::Goodbye if frame.header.channel == 0 => Ok(false),
861        FrameType::Goodbye => {
862            let handle = module_handle.route_handle(frame.header.channel, frame.header.epoch);
863            if module_handle.remove_route(handle)? {
864                cancel_handle(&dispatcher.in_flight, handle)?;
865                handler.on_route_gone(&handle).await;
866            }
867            Ok(true)
868        }
869        FrameType::Response if frame.header.channel == 0 => {
870            let _ = module_handle.handle_control_reply(frame);
871            Ok(true)
872        }
873        FrameType::Error if frame.header.channel == 0 => {
874            let _ = module_handle.handle_control_reply(frame);
875            Ok(true)
876        }
877        FrameType::Cancel => {
878            handle_cancel(frame, &dispatcher.in_flight)?;
879            Ok(true)
880        }
881        FrameType::Request if frame.header.channel == 0 => {
882            handle_control_request(frame, egress, handler, dispatcher, module_handle.clone())
883                .await?;
884            Ok(true)
885        }
886        FrameType::Request => {
887            spawn_data_request(frame, egress.clone(), handler, dispatcher, module_handle)?;
888            Ok(true)
889        }
890        _ => Ok(true),
891    }
892}
893
894fn spawn_data_request<H>(
895    frame: Frame,
896    egress: mpsc::Sender<Frame>,
897    handler: Arc<H>,
898    dispatcher: RequestDispatcher,
899    module_handle: ModuleHandle,
900) -> Result<(), SubcModuleError>
901where
902    H: ModuleHandler,
903{
904    let handle = module_handle.route_handle(frame.header.channel, frame.header.epoch);
905    let corr = frame.header.corr;
906    let cancellation = CancellationToken::new();
907    {
908        let mut guard = lock_in_flight(&dispatcher.in_flight)?;
909        guard.insert((handle.channel, handle.epoch, corr), cancellation.clone());
910    }
911
912    let ctx = RequestCtx {
913        handle,
914        corr,
915        ver: frame.header.ver,
916        egress,
917        module_handle,
918        cancelled: cancellation,
919    };
920    let body = frame.body;
921    let in_flight = Arc::clone(&dispatcher.in_flight);
922    let permits = Arc::clone(&dispatcher.permits);
923    tokio::spawn(async move {
924        let Ok(_permit) = permits.acquire_owned().await else {
925            // A closed dispatcher means connection teardown will release every route credit.
926            if let Ok(mut guard) = in_flight.lock() {
927                guard.remove(&(handle.channel, handle.epoch, corr));
928            }
929            return;
930        };
931        if ctx.cancelled.is_cancelled() {
932            let _ = send_handler_outcome(
933                &ctx,
934                HandlerOutcome::Error {
935                    code: "cancelled".to_string(),
936                    message: "request cancelled".to_string(),
937                },
938            )
939            .await;
940            if let Ok(mut guard) = in_flight.lock() {
941                guard.remove(&(handle.channel, handle.epoch, corr));
942            }
943            return;
944        }
945        let outcome = handler.handle(ctx.clone(), body).await;
946        let _ = send_handler_outcome(&ctx, outcome).await;
947        if let Ok(mut guard) = in_flight.lock() {
948            guard.remove(&(handle.channel, handle.epoch, corr));
949        }
950    });
951    Ok(())
952}
953
954fn spawn_health_request<H>(
955    frame: Frame,
956    egress: mpsc::Sender<Frame>,
957    handler: Arc<H>,
958    dispatcher: RequestDispatcher,
959) -> Result<(), SubcModuleError>
960where
961    H: ModuleHandler,
962{
963    let channel = frame.header.channel;
964    let epoch = frame.header.epoch;
965    let corr = frame.header.corr;
966    let ver = frame.header.ver;
967    let cancellation = CancellationToken::new();
968    {
969        let mut guard = lock_in_flight(&dispatcher.in_flight)?;
970        guard.insert((channel, epoch, corr), cancellation.clone());
971    }
972
973    let in_flight = Arc::clone(&dispatcher.in_flight);
974    let permits = Arc::clone(&dispatcher.permits);
975    tokio::spawn(async move {
976        let Ok(_permit) = permits.acquire_owned().await else {
977            if let Ok(mut guard) = in_flight.lock() {
978                guard.remove(&(channel, epoch, corr));
979            }
980            return;
981        };
982        if !cancellation.is_cancelled() {
983            let report = handler.health().await;
984            let response = ModuleControlResponse::from(report);
985            if let Ok(body) = serde_json::to_vec(&response) {
986                if let Ok(frame) = Frame::build_with_version(
987                    ver,
988                    FrameType::Response,
989                    control_flags(),
990                    channel,
991                    epoch,
992                    corr,
993                    body,
994                ) {
995                    let _ = send_outbound(&egress, frame).await;
996                }
997            }
998        }
999        if let Ok(mut guard) = in_flight.lock() {
1000            guard.remove(&(channel, epoch, corr));
1001        }
1002    });
1003    Ok(())
1004}
1005
1006async fn send_handler_outcome(
1007    ctx: &RequestCtx,
1008    outcome: HandlerOutcome,
1009) -> Result<(), SubcModuleError> {
1010    match outcome {
1011        HandlerOutcome::Response(body) => {
1012            ctx.send_frame(FrameType::Response, data_flags(), body)
1013                .await
1014        }
1015        HandlerOutcome::Error { code, message } => {
1016            let body =
1017                serde_json::to_vec(&ErrorBody { code, message }).map_err(SubcModuleError::Json)?;
1018            ctx.send_frame(FrameType::Error, data_flags(), body).await
1019        }
1020        HandlerOutcome::Streamed => {
1021            ctx.send_frame(FrameType::StreamEnd, data_flags(), Vec::new())
1022                .await
1023        }
1024    }
1025}
1026
1027fn handle_cancel(frame: Frame, in_flight: &InFlight) -> Result<(), SubcModuleError> {
1028    let cancellation = {
1029        let guard = lock_in_flight(in_flight)?;
1030        guard
1031            .get(&(frame.header.channel, frame.header.epoch, frame.header.corr))
1032            .cloned()
1033    };
1034    if let Some(cancellation) = cancellation {
1035        cancellation.cancel();
1036    }
1037    Ok(())
1038}
1039
1040fn cancel_handle(in_flight: &InFlight, handle: RouteHandle) -> Result<(), SubcModuleError> {
1041    let cancelled = {
1042        let mut guard = lock_in_flight(in_flight)?;
1043        let keys = guard
1044            .keys()
1045            .copied()
1046            .filter(|(channel, epoch, _)| *channel == handle.channel && *epoch == handle.epoch)
1047            .collect::<Vec<_>>();
1048        keys.into_iter()
1049            .filter_map(|key| guard.remove(&key))
1050            .collect::<Vec<_>>()
1051    };
1052    for cancellation in cancelled {
1053        cancellation.cancel();
1054    }
1055    Ok(())
1056}
1057
1058async fn handle_control_request<H>(
1059    frame: Frame,
1060    egress: &mpsc::Sender<Frame>,
1061    handler: Arc<H>,
1062    dispatcher: RequestDispatcher,
1063    module_handle: ModuleHandle,
1064) -> Result<(), SubcModuleError>
1065where
1066    H: ModuleHandler,
1067{
1068    let request = serde_json::from_slice::<ModuleControlRequest>(&frame.body)
1069        .map_err(SubcModuleError::Json)?;
1070    match request {
1071        ModuleControlRequest::RouteBind {
1072            route_channel,
1073            epoch,
1074            target,
1075            identity,
1076            principal,
1077            consumer_capabilities,
1078            admission_facts,
1079        } => {
1080            // Implicit-replace rule (wire spec 3.3.0): the daemon never rebinds a live
1081            // channel, but its route-gone GOODBYE to modules is best-effort, so a bind
1082            // can arrive for a channel this endpoint still believes installed. A
1083            // strictly higher epoch proves the daemon freed the old binding: tear the
1084            // stale install down locally and proceed. Equal or lower epoch is a
1085            // protocol violation the daemon cannot produce: reject the bind.
1086            if let Some(stale) = module_handle.installed_route(route_channel)? {
1087                if epoch <= stale.epoch {
1088                    let body = serde_json::to_vec(&ErrorBody {
1089                        code: "route_rejected".to_string(),
1090                        message: format!(
1091                            "route.bind epoch {epoch} does not supersede installed epoch {} on channel {route_channel}",
1092                            stale.epoch
1093                        ),
1094                    })
1095                    .map_err(SubcModuleError::Json)?;
1096                    let reject = Frame::build_with_version(
1097                        frame.header.ver,
1098                        FrameType::Error,
1099                        control_flags(),
1100                        0,
1101                        0,
1102                        frame.header.corr,
1103                        body,
1104                    )
1105                    .map_err(SubcModuleError::FrameBuild)?;
1106                    send_outbound(egress, reject).await?;
1107                    return Ok(());
1108                }
1109                if module_handle.remove_route(stale)? {
1110                    cancel_handle(&dispatcher.in_flight, stale)?;
1111                    handler.on_route_gone(&stale).await;
1112                }
1113            }
1114            let handle = module_handle.route_handle(route_channel, epoch);
1115            let req = RouteBindRequest {
1116                handle,
1117                target,
1118                identity,
1119                principal,
1120                consumer_capabilities,
1121                admission_facts,
1122            };
1123            let decision = handler.on_bind(&req).await;
1124            match decision.kind {
1125                BindDecisionKind::Accept => {
1126                    let response = match serde_json::to_vec(&ModuleControlResponse::RouteBindAck {})
1127                        .map_err(SubcModuleError::Json)
1128                        .and_then(|body| {
1129                            Frame::build_with_version(
1130                                frame.header.ver,
1131                                FrameType::Response,
1132                                control_flags(),
1133                                0,
1134                                0,
1135                                frame.header.corr,
1136                                body,
1137                            )
1138                            .map_err(SubcModuleError::FrameBuild)
1139                        }) {
1140                        Ok(response) => response,
1141                        Err(err) => {
1142                            handler.on_route_gone(&handle).await;
1143                            return Err(err);
1144                        }
1145                    };
1146                    if let Err(err) = send_outbound(egress, response).await {
1147                        handler.on_route_gone(&handle).await;
1148                        return Err(err);
1149                    }
1150                    if let Err(err) = module_handle.install_route(handle) {
1151                        handler.on_route_gone(&handle).await;
1152                        return Err(err);
1153                    }
1154                    handler.on_bound(&handle).await;
1155                }
1156                BindDecisionKind::Reject { code, message } => {
1157                    let result = serde_json::to_vec(&ErrorBody { code, message })
1158                        .map_err(SubcModuleError::Json)
1159                        .and_then(|body| {
1160                            Frame::build_with_version(
1161                                frame.header.ver,
1162                                FrameType::Error,
1163                                control_flags(),
1164                                0,
1165                                0,
1166                                frame.header.corr,
1167                                body,
1168                            )
1169                            .map_err(SubcModuleError::FrameBuild)
1170                        });
1171                    let result = match result {
1172                        Ok(response) => send_outbound(egress, response).await,
1173                        Err(err) => Err(err),
1174                    };
1175                    handler.on_route_gone(&handle).await;
1176                    result?;
1177                }
1178            }
1179        }
1180        ModuleControlRequest::HealthCheck {} => {
1181            spawn_health_request(frame, egress.clone(), handler, dispatcher)?;
1182        }
1183    }
1184    Ok(())
1185}
1186
1187async fn send_hello(
1188    egress: &mpsc::Sender<Frame>,
1189    manifest: ModuleManifest,
1190) -> Result<(), SubcModuleError> {
1191    let body = serde_json::to_vec(&ModuleHelloBody {
1192        manifest,
1193        protocol_ver: PROTOCOL_VERSION,
1194        control_ops: Some(vec![MODULE_CONTROL_OP_HEALTH_CHECK.to_string()]),
1195        launch_nonce: env::var(SUBC_LAUNCH_NONCE_ENV)
1196            .ok()
1197            .filter(|value| !value.is_empty()),
1198    })
1199    .map_err(SubcModuleError::Json)?;
1200    let frame = Frame::build(FrameType::Hello, control_flags(), 0, 0, HELLO_CORR, body)
1201        .map_err(SubcModuleError::FrameBuild)?;
1202    send_outbound(egress, frame).await
1203}
1204
1205async fn expect_hello_ack<R>(reader: &mut R) -> Result<ModuleHelloAckBody, SubcModuleError>
1206where
1207    R: AsyncRead + Unpin,
1208{
1209    let Some(frame) = read_frame(reader).await.map_err(SubcModuleError::FrameIo)? else {
1210        return Err(SubcModuleError::ConnectionClosedBeforeHelloAck);
1211    };
1212    match frame.header.ty {
1213        FrameType::HelloAck => serde_json::from_slice(&frame.body).map_err(SubcModuleError::Json),
1214        FrameType::Error => {
1215            let body =
1216                serde_json::from_slice::<ErrorBody>(&frame.body).map_err(SubcModuleError::Json)?;
1217            Err(SubcModuleError::HelloRejected { body })
1218        }
1219        ty => Err(SubcModuleError::UnexpectedHelloAck { ty }),
1220    }
1221}
1222
1223async fn connect_to_subc(connection_file_path: &Path) -> Result<TcpStream, SubcModuleError> {
1224    let conn = connection_file::read_for_client(connection_file_path).map_err(|source| {
1225        SubcModuleError::ConnectionFile {
1226            path: connection_file_path.to_path_buf(),
1227            source,
1228        }
1229    })?;
1230    let endpoint = conn
1231        .endpoints
1232        .first()
1233        .ok_or_else(|| SubcModuleError::NoEndpoint {
1234            path: connection_file_path.to_path_buf(),
1235        })?;
1236    let endpoint_label = format!("{}:{}", endpoint.host, endpoint.port);
1237    let mut stream = TcpStream::connect(&endpoint_label)
1238        .await
1239        .map_err(|source| SubcModuleError::Connect {
1240            path: connection_file_path.to_path_buf(),
1241            endpoint: endpoint_label.clone(),
1242            source,
1243        })?;
1244    // This socket carries the module's replies back to the daemon, so Nagle here
1245    // delays every response rather than every request -- the same cost on the
1246    // return leg. Both ends of the hop have to disable it for either to help.
1247    //
1248    // The result is deliberately dropped rather than logged: this crate takes no
1249    // logging dependency, and the only ways setting a socket option on a
1250    // just-connected stream fail leave the socket unusable, which the handshake on
1251    // the very next line reports as a typed Auth error. Swallowing it here would
1252    // hide nothing that stays hidden.
1253    let _ = stream.set_nodelay(true);
1254    authenticate_client(&mut stream, &conn, AUTH_DEADLINE)
1255        .await
1256        .map_err(|source| SubcModuleError::Auth {
1257            path: connection_file_path.to_path_buf(),
1258            endpoint: endpoint_label,
1259            source,
1260        })?;
1261    Ok(stream)
1262}
1263
1264async fn drain_writer<W>(write_half: W, mut rx: mpsc::Receiver<Frame>) -> Result<(), FrameIoError>
1265where
1266    W: AsyncWrite + Unpin,
1267{
1268    let mut writer = BufWriter::new(write_half);
1269    while let Some(frame) = rx.recv().await {
1270        write_frame(&mut writer, &frame).await?;
1271        while let Ok(frame) = rx.try_recv() {
1272            write_frame(&mut writer, &frame).await?;
1273        }
1274        writer.flush().await.map_err(FrameIoError::Io)?;
1275    }
1276    writer.flush().await.map_err(FrameIoError::Io)?;
1277    Ok(())
1278}
1279
1280async fn send_outbound(egress: &mpsc::Sender<Frame>, frame: Frame) -> Result<(), SubcModuleError> {
1281    egress
1282        .send(frame)
1283        .await
1284        .map_err(|_| SubcModuleError::WriterClosed)
1285}
1286
1287fn parse_subc_arg(args: impl IntoIterator<Item = OsString>) -> Result<PathBuf, SubcModuleError> {
1288    let mut args = args.into_iter();
1289    while let Some(arg) = args.next() {
1290        if arg == "--subc" {
1291            let value = args.next().ok_or(SubcModuleError::MissingSubcValue)?;
1292            return Ok(PathBuf::from(value));
1293        }
1294        if let Some(raw) = arg.to_str().and_then(|arg| arg.strip_prefix("--subc=")) {
1295            if raw.is_empty() {
1296                return Err(SubcModuleError::MissingSubcValue);
1297            }
1298            return Ok(PathBuf::from(raw));
1299        }
1300    }
1301    Err(SubcModuleError::MissingSubcArg)
1302}
1303
1304fn module_id_from_env() -> Result<Option<String>, SubcModuleError> {
1305    match env::var(SUBC_MODULE_ID_ENV) {
1306        Ok(value) if !value.trim().is_empty() => Ok(Some(value)),
1307        Ok(_) => Err(SubcModuleError::EmptyModuleIdEnv),
1308        Err(env::VarError::NotPresent) => Ok(None),
1309        Err(env::VarError::NotUnicode(value)) => {
1310            Err(SubcModuleError::NonUnicodeModuleIdEnv { value })
1311        }
1312    }
1313}
1314
1315fn lock_in_flight(
1316    in_flight: &InFlight,
1317) -> Result<std::sync::MutexGuard<'_, HashMap<RequestKey, CancellationToken>>, SubcModuleError> {
1318    in_flight
1319        .lock()
1320        .map_err(|_| SubcModuleError::InFlightPoisoned)
1321}
1322
1323fn control_flags() -> Flags {
1324    Flags::new(false, Priority::Passive, false)
1325}
1326
1327fn data_flags() -> Flags {
1328    Flags::new(false, Priority::Interactive, false)
1329}
1330
1331#[derive(Debug)]
1332pub enum SubcModuleError {
1333    MissingSubcArg,
1334    MissingSubcValue,
1335    EmptyModuleIdEnv,
1336    NonUnicodeModuleIdEnv {
1337        value: OsString,
1338    },
1339    ConnectionFile {
1340        path: PathBuf,
1341        source: ConnectionFileError,
1342    },
1343    NoEndpoint {
1344        path: PathBuf,
1345    },
1346    Connect {
1347        path: PathBuf,
1348        endpoint: String,
1349        source: io::Error,
1350    },
1351    Auth {
1352        path: PathBuf,
1353        endpoint: String,
1354        source: AuthError,
1355    },
1356    FrameIo(FrameIoError),
1357    FrameBuild(FrameBuildError),
1358    Json(serde_json::Error),
1359    WriterClosed,
1360    StaleRouteHandle(RouteHandle),
1361    ConnectionTokenExhausted,
1362    WriterTask(tokio::task::JoinError),
1363    InFlightPoisoned,
1364    ConnectionClosedBeforeHelloAck,
1365    UnexpectedHelloAck {
1366        ty: FrameType,
1367    },
1368    HelloRejected {
1369        body: ErrorBody,
1370    },
1371}
1372
1373impl fmt::Display for SubcModuleError {
1374    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1375        match self {
1376            Self::MissingSubcArg => write!(f, "missing required --subc <connection-file> argument"),
1377            Self::MissingSubcValue => write!(f, "--subc requires a connection-file path value"),
1378            Self::EmptyModuleIdEnv => write!(f, "{SUBC_MODULE_ID_ENV} must not be empty when set"),
1379            Self::NonUnicodeModuleIdEnv { value } => write!(
1380                f,
1381                "{SUBC_MODULE_ID_ENV} must be valid UTF-8, got '{}'",
1382                value.to_string_lossy()
1383            ),
1384            Self::ConnectionFile { path, source } => write!(
1385                f,
1386                "failed to read subc connection file '{}': {source}",
1387                path.display()
1388            ),
1389            Self::NoEndpoint { path } => write!(
1390                f,
1391                "subc connection file '{}' has no endpoints",
1392                path.display()
1393            ),
1394            Self::Connect {
1395                path,
1396                endpoint,
1397                source,
1398            } => write!(
1399                f,
1400                "failed to connect to subc endpoint {endpoint} from '{}': {source}",
1401                path.display()
1402            ),
1403            Self::Auth {
1404                path,
1405                endpoint,
1406                source,
1407            } => write!(
1408                f,
1409                "failed to authenticate to subc endpoint {endpoint} from '{}': {source}",
1410                path.display()
1411            ),
1412            Self::FrameIo(err) => write!(f, "frame I/O error: {err}"),
1413            Self::FrameBuild(err) => write!(f, "frame build error: {err}"),
1414            Self::Json(err) => write!(f, "JSON error: {err}"),
1415            Self::WriterClosed => write!(f, "module writer task closed"),
1416            Self::StaleRouteHandle(handle) => write!(f, "stale route handle: {handle:?}"),
1417            Self::ConnectionTokenExhausted => write!(f, "module connection token exhausted"),
1418            Self::WriterTask(err) => write!(f, "module writer task failed: {err}"),
1419            Self::InFlightPoisoned => write!(f, "in-flight registry lock poisoned"),
1420            Self::ConnectionClosedBeforeHelloAck => write!(f, "connection closed before HELLO_ACK"),
1421            Self::UnexpectedHelloAck { ty } => write!(f, "expected HELLO_ACK, got {ty:?}"),
1422            Self::HelloRejected { body } => write!(
1423                f,
1424                "HELLO rejected by subc: {} ({})",
1425                body.code, body.message
1426            ),
1427        }
1428    }
1429}
1430
1431impl Error for SubcModuleError {
1432    fn source(&self) -> Option<&(dyn Error + 'static)> {
1433        match self {
1434            Self::ConnectionFile { source, .. } => Some(source),
1435            Self::Connect { source, .. } => Some(source),
1436            Self::Auth { source, .. } => Some(source),
1437            Self::FrameIo(err) => Some(err),
1438            Self::FrameBuild(err) => Some(err),
1439            Self::Json(err) => Some(err),
1440            Self::WriterTask(err) => Some(err),
1441            Self::MissingSubcArg
1442            | Self::MissingSubcValue
1443            | Self::EmptyModuleIdEnv
1444            | Self::NonUnicodeModuleIdEnv { .. }
1445            | Self::NoEndpoint { .. }
1446            | Self::WriterClosed
1447            | Self::StaleRouteHandle(_)
1448            | Self::ConnectionTokenExhausted
1449            | Self::InFlightPoisoned
1450            | Self::ConnectionClosedBeforeHelloAck
1451            | Self::UnexpectedHelloAck { .. }
1452            | Self::HelloRejected { .. } => None,
1453        }
1454    }
1455}
1456
1457impl From<serde_json::Error> for SubcModuleError {
1458    fn from(err: serde_json::Error) -> Self {
1459        Self::Json(err)
1460    }
1461}
1462
1463#[cfg(test)]
1464mod tests {
1465    use std::sync::atomic::{AtomicUsize, Ordering};
1466
1467    use serde_json::json;
1468    use subc_protocol::manifest::{Concurrency, ExecutionMode, IdentityScope, Tool};
1469    use tokio::{sync::Notify, time::timeout};
1470
1471    use super::*;
1472
1473    struct EchoHandler;
1474
1475    #[async_trait]
1476    impl ModuleHandler for EchoHandler {
1477        async fn handle(&self, _ctx: RequestCtx, body: Vec<u8>) -> HandlerOutcome {
1478            HandlerOutcome::Response(body)
1479        }
1480    }
1481
1482    struct BlockingHandler {
1483        entered: Arc<AtomicUsize>,
1484        release: Arc<Notify>,
1485    }
1486
1487    #[async_trait]
1488    impl ModuleHandler for BlockingHandler {
1489        async fn handle(&self, _ctx: RequestCtx, _body: Vec<u8>) -> HandlerOutcome {
1490            self.entered.fetch_add(1, Ordering::SeqCst);
1491            self.release.notified().await;
1492            HandlerOutcome::Streamed
1493        }
1494    }
1495
1496    fn health_request(corr: u64) -> Frame {
1497        Frame::build(
1498            FrameType::Request,
1499            control_flags(),
1500            0,
1501            0,
1502            corr,
1503            serde_json::to_vec(&ModuleControlRequest::HealthCheck {}).unwrap(),
1504        )
1505        .unwrap()
1506    }
1507
1508    fn data_request(channel: u16, corr: u64) -> Frame {
1509        Frame::build(
1510            FrameType::Request,
1511            data_flags(),
1512            channel,
1513            1,
1514            corr,
1515            b"opaque".to_vec(),
1516        )
1517        .unwrap()
1518    }
1519
1520    fn catalog_update_response(corr: u64) -> Frame {
1521        Frame::build(
1522            FrameType::Response,
1523            control_flags(),
1524            0,
1525            0,
1526            corr,
1527            serde_json::to_vec(&ModuleControlResponseToModule::CatalogUpdate {}).unwrap(),
1528        )
1529        .unwrap()
1530    }
1531
1532    fn test_module_handle(subc_ops: &[&str]) -> (ModuleHandle, mpsc::Receiver<Frame>) {
1533        let (tx, rx) = mpsc::channel(4);
1534        let ack = ModuleHelloAckBody {
1535            negotiated_ver: PROTOCOL_VERSION,
1536            subc_ops: subc_ops.iter().map(|op| (*op).to_string()).collect(),
1537            subc_capabilities: Vec::new(),
1538            storage: None,
1539        };
1540        (ModuleHandle::new(&ack, tx, 1, CancellationToken::new()), rx)
1541    }
1542
1543    fn test_provider_role(tool_names: &[&str]) -> ProviderRole {
1544        ProviderRole::ToolProvider {
1545            tools: tool_names
1546                .iter()
1547                .map(|name| Tool {
1548                    name: (*name).to_string(),
1549                    description: None,
1550                    execution_mode: ExecutionMode::Pure,
1551                    schema: json!({"type": "object"}),
1552                })
1553                .collect(),
1554            identity_scope: vec![IdentityScope::Project],
1555            concurrency: Concurrency::ModuleManaged,
1556            emits_push: false,
1557            sub_supervises: false,
1558        }
1559    }
1560
1561    #[tokio::test]
1562    async fn catalog_update_fails_fast_when_hello_ack_does_not_advertise_support() {
1563        let (handle, mut rx) = test_module_handle(&[]);
1564
1565        let error = handle
1566            .catalog_update(vec![test_provider_role(&["a"])])
1567            .await
1568            .unwrap_err();
1569        assert!(matches!(error, CatalogUpdateError::NotSupported));
1570        assert!(timeout(Duration::from_millis(75), rx.recv()).await.is_err());
1571    }
1572
1573    #[tokio::test]
1574    async fn catalog_update_demuxes_multiple_in_flight_requests() {
1575        let (handle, mut rx) = test_module_handle(&[MODULE_TO_SUBC_OP_CATALOG_UPDATE]);
1576        let first_handle = handle.clone();
1577        let second_handle = handle.clone();
1578        let first = tokio::spawn(async move {
1579            first_handle
1580                .catalog_update(vec![test_provider_role(&["a"])])
1581                .await
1582        });
1583        let second = tokio::spawn(async move {
1584            second_handle
1585                .catalog_update(vec![test_provider_role(&["b"])])
1586                .await
1587        });
1588
1589        let first_frame = timeout(Duration::from_secs(1), rx.recv())
1590            .await
1591            .unwrap()
1592            .unwrap();
1593        let second_frame = timeout(Duration::from_secs(1), rx.recv())
1594            .await
1595            .unwrap()
1596            .unwrap();
1597        assert_eq!(first_frame.header.ty, FrameType::Request);
1598        assert_eq!(first_frame.header.channel, 0);
1599        assert_eq!(second_frame.header.ty, FrameType::Request);
1600        assert_eq!(second_frame.header.channel, 0);
1601        assert_ne!(first_frame.header.corr, second_frame.header.corr);
1602
1603        assert!(handle.handle_control_reply(catalog_update_response(second_frame.header.corr)));
1604        assert!(handle.handle_control_reply(catalog_update_response(first_frame.header.corr)));
1605        assert!(first.await.unwrap().is_ok());
1606        assert!(second.await.unwrap().is_ok());
1607    }
1608
1609    #[tokio::test]
1610    async fn default_health_check_answers_ok() {
1611        let (tx, mut rx) = mpsc::channel(4);
1612        let handler = Arc::new(EchoHandler);
1613        let dispatcher = RequestDispatcher::new();
1614        let (module_handle, _unused_rx) = test_module_handle(&[]);
1615
1616        assert!(
1617            handle_frame(health_request(77), &tx, handler, dispatcher, module_handle)
1618                .await
1619                .unwrap()
1620        );
1621
1622        let response = timeout(Duration::from_secs(1), rx.recv())
1623            .await
1624            .unwrap()
1625            .unwrap();
1626        assert_eq!(response.header.ty, FrameType::Response);
1627        assert_eq!(response.header.channel, 0);
1628        assert_eq!(response.header.corr, 77);
1629
1630        // Assert the PROPERTIES that matter rather than byte-equality with
1631        // HealthReport::ok(). Comparing against the constructor made this test a
1632        // restatement of the implementation: it reddened on any change to the
1633        // default without saying which property had broken.
1634        let parsed = serde_json::from_slice::<ModuleControlResponse>(&response.body).unwrap();
1635        let ModuleControlResponse::HealthCheck { status, detail, .. } = parsed else {
1636            panic!("expected a health.check response");
1637        };
1638        // A module that never implemented health is not UNHEALTHY -- the daemon
1639        // must not escalate on it.
1640        assert_eq!(status, HealthStatus::Ok);
1641        // ...but the report must SAY that nobody measured, so an operator can
1642        // tell it from a real all-clear. Substring rather than exact text: the
1643        // wording is for humans and nothing parses it.
1644        assert!(
1645            detail
1646                .as_deref()
1647                .is_some_and(|d| d.contains("no health implementation")),
1648            "the inherited default must identify itself, got {detail:?}"
1649        );
1650    }
1651
1652    #[tokio::test]
1653    async fn health_check_waits_behind_saturated_request_dispatcher() {
1654        let (tx, mut rx) = mpsc::channel(4);
1655        let entered = Arc::new(AtomicUsize::new(0));
1656        let release = Arc::new(Notify::new());
1657        let handler = Arc::new(BlockingHandler {
1658            entered: Arc::clone(&entered),
1659            release: Arc::clone(&release),
1660        });
1661        let dispatcher = RequestDispatcher::new();
1662        let (module_handle, _unused_rx) = test_module_handle(&[]);
1663        module_handle
1664            .install_route(RouteHandle::new(7, 1, 1))
1665            .unwrap();
1666
1667        for corr in 0..HANDLER_TASK_CAPACITY as u64 {
1668            handle_frame(
1669                data_request(7, corr + 1),
1670                &tx,
1671                Arc::clone(&handler),
1672                dispatcher.clone(),
1673                module_handle.clone(),
1674            )
1675            .await
1676            .unwrap();
1677        }
1678
1679        timeout(Duration::from_secs(1), async {
1680            while entered.load(Ordering::SeqCst) < HANDLER_TASK_CAPACITY {
1681                tokio::task::yield_now().await;
1682            }
1683        })
1684        .await
1685        .unwrap();
1686
1687        handle_frame(health_request(900), &tx, handler, dispatcher, module_handle)
1688            .await
1689            .unwrap();
1690        assert!(
1691            timeout(Duration::from_millis(75), rx.recv()).await.is_err(),
1692            "health.check must share the same saturated request dispatch capacity as data requests"
1693        );
1694
1695        release.notify_waiters();
1696    }
1697
1698    struct CorrBlockingHandler {
1699        entered: Arc<Mutex<Vec<u64>>>,
1700        release_first: Arc<Semaphore>,
1701    }
1702
1703    #[async_trait]
1704    impl ModuleHandler for CorrBlockingHandler {
1705        async fn handle(&self, ctx: RequestCtx, _body: Vec<u8>) -> HandlerOutcome {
1706            let corr = ctx.corr();
1707            self.entered.lock().unwrap().push(corr);
1708            if corr == 1 {
1709                self.release_first.acquire().await.unwrap().forget();
1710            }
1711            HandlerOutcome::Response(Vec::new())
1712        }
1713    }
1714
1715    #[tokio::test]
1716    async fn cancelled_capacity_queued_data_request_emits_terminal_and_skips_handler() {
1717        let (tx, mut rx) = mpsc::channel(4);
1718        let entered = Arc::new(Mutex::new(Vec::new()));
1719        let release_first = Arc::new(Semaphore::new(0));
1720        let handler = Arc::new(CorrBlockingHandler {
1721            entered: Arc::clone(&entered),
1722            release_first: Arc::clone(&release_first),
1723        });
1724        let dispatcher = RequestDispatcher {
1725            in_flight: Arc::new(Mutex::new(HashMap::new())),
1726            permits: Arc::new(Semaphore::new(1)),
1727        };
1728        let (module_handle, _unused_rx) = test_module_handle(&[]);
1729        module_handle
1730            .install_route(RouteHandle::new(7, 1, 1))
1731            .unwrap();
1732
1733        handle_frame(
1734            data_request(7, 1),
1735            &tx,
1736            Arc::clone(&handler),
1737            dispatcher.clone(),
1738            module_handle.clone(),
1739        )
1740        .await
1741        .unwrap();
1742        timeout(Duration::from_secs(1), async {
1743            while entered.lock().unwrap().as_slice() != [1] {
1744                tokio::task::yield_now().await;
1745            }
1746        })
1747        .await
1748        .unwrap();
1749
1750        handle_frame(
1751            data_request(7, 2),
1752            &tx,
1753            Arc::clone(&handler),
1754            dispatcher.clone(),
1755            module_handle.clone(),
1756        )
1757        .await
1758        .unwrap();
1759        timeout(Duration::from_secs(1), async {
1760            while !dispatcher
1761                .in_flight
1762                .lock()
1763                .unwrap()
1764                .contains_key(&(7, 1, 2))
1765            {
1766                tokio::task::yield_now().await;
1767            }
1768        })
1769        .await
1770        .unwrap();
1771        handle_frame(
1772            Frame::build(FrameType::Cancel, data_flags(), 7, 1, 2, Vec::new()).unwrap(),
1773            &tx,
1774            Arc::clone(&handler),
1775            dispatcher.clone(),
1776            module_handle,
1777        )
1778        .await
1779        .unwrap();
1780        assert!(
1781            dispatcher
1782                .in_flight
1783                .lock()
1784                .unwrap()
1785                .get(&(7, 1, 2))
1786                .unwrap()
1787                .is_cancelled(),
1788            "cancel must land while the second request waits for handler capacity"
1789        );
1790
1791        release_first.add_permits(1);
1792        timeout(Duration::from_secs(1), async {
1793            while !dispatcher.in_flight.lock().unwrap().is_empty() {
1794                tokio::task::yield_now().await;
1795            }
1796        })
1797        .await
1798        .unwrap();
1799
1800        assert_eq!(*entered.lock().unwrap(), vec![1]);
1801        let response = timeout(Duration::from_secs(1), rx.recv())
1802            .await
1803            .unwrap()
1804            .unwrap();
1805        assert_eq!(response.header.ty, FrameType::Response);
1806        assert_eq!(response.header.corr, 1);
1807
1808        let cancelled = timeout(Duration::from_secs(1), rx.recv())
1809            .await
1810            .unwrap()
1811            .unwrap();
1812        assert_eq!(cancelled.header.ty, FrameType::Error);
1813        assert_eq!(cancelled.header.channel, 7);
1814        assert_eq!(cancelled.header.epoch, 1);
1815        assert_eq!(cancelled.header.corr, 2);
1816        assert_eq!(
1817            serde_json::from_slice::<ErrorBody>(&cancelled.body).unwrap(),
1818            ErrorBody {
1819                code: "cancelled".to_string(),
1820                message: "request cancelled".to_string(),
1821            }
1822        );
1823        assert!(timeout(Duration::from_millis(50), rx.recv()).await.is_err());
1824    }
1825
1826    #[tokio::test]
1827    async fn cancelled_terminal_is_not_sent_after_route_teardown() {
1828        let (tx, mut rx) = mpsc::channel(1);
1829        let (module_handle, _unused_rx) = test_module_handle(&[]);
1830        let handle = RouteHandle::new(7, 1, 1);
1831        module_handle.install_route(handle).unwrap();
1832        let ctx = RequestCtx {
1833            handle,
1834            corr: 2,
1835            ver: PROTOCOL_VERSION,
1836            egress: tx,
1837            module_handle: module_handle.clone(),
1838            cancelled: CancellationToken::new(),
1839        };
1840        assert!(module_handle.remove_route(handle).unwrap());
1841
1842        let result = send_handler_outcome(
1843            &ctx,
1844            HandlerOutcome::Error {
1845                code: "cancelled".to_string(),
1846                message: "request cancelled".to_string(),
1847            },
1848        )
1849        .await;
1850
1851        assert!(matches!(
1852            result,
1853            Err(SubcModuleError::StaleRouteHandle(stale)) if stale == handle
1854        ));
1855        assert!(rx.try_recv().is_err());
1856    }
1857
1858    struct CountingHandler {
1859        calls: Arc<AtomicUsize>,
1860    }
1861
1862    #[async_trait]
1863    impl ModuleHandler for CountingHandler {
1864        async fn handle(&self, _ctx: RequestCtx, _body: Vec<u8>) -> HandlerOutcome {
1865            self.calls.fetch_add(1, Ordering::SeqCst);
1866            HandlerOutcome::Response(Vec::new())
1867        }
1868    }
1869
1870    #[tokio::test]
1871    async fn endpoint_validation_drops_stale_request_before_handler_dispatch() {
1872        let (tx, mut rx) = mpsc::channel(4);
1873        let calls = Arc::new(AtomicUsize::new(0));
1874        let handler = Arc::new(CountingHandler {
1875            calls: Arc::clone(&calls),
1876        });
1877        let dispatcher = RequestDispatcher::new();
1878        let (module_handle, _unused_rx) = test_module_handle(&[]);
1879        module_handle
1880            .install_route(RouteHandle::new(7, 2, 1))
1881            .unwrap();
1882        let stale = Frame::build(FrameType::Request, data_flags(), 7, 1, 55, Vec::new()).unwrap();
1883
1884        assert!(
1885            handle_frame(stale, &tx, handler, dispatcher, module_handle.clone())
1886                .await
1887                .unwrap()
1888        );
1889        tokio::task::yield_now().await;
1890        assert_eq!(calls.load(Ordering::SeqCst), 0);
1891        assert_eq!(module_handle.dropped_route_frames(), 1);
1892        assert!(rx.try_recv().is_err());
1893    }
1894
1895    struct BindOrderingHandler {
1896        module_handle: ModuleHandle,
1897        bind_emit_rejected: Arc<AtomicUsize>,
1898        bound: Arc<AtomicUsize>,
1899        cleanup: Arc<AtomicUsize>,
1900    }
1901
1902    #[async_trait]
1903    impl ModuleHandler for BindOrderingHandler {
1904        async fn handle(&self, _ctx: RequestCtx, _body: Vec<u8>) -> HandlerOutcome {
1905            HandlerOutcome::Response(Vec::new())
1906        }
1907
1908        async fn on_bind(&self, req: &RouteBindRequest) -> BindDecision {
1909            if matches!(
1910                self.module_handle
1911                    .push(&req.handle, b"too-early".to_vec(), None)
1912                    .await,
1913                Err(SubcModuleError::StaleRouteHandle(_))
1914            ) {
1915                self.bind_emit_rejected.fetch_add(1, Ordering::SeqCst);
1916            }
1917            BindDecision::accept()
1918        }
1919
1920        async fn on_bound(&self, handle: &RouteHandle) {
1921            self.bound.fetch_add(1, Ordering::SeqCst);
1922            self.module_handle
1923                .push(handle, b"bound".to_vec(), Some(AdmissionClass::Expedite))
1924                .await
1925                .unwrap();
1926        }
1927
1928        async fn on_route_gone(&self, _handle: &RouteHandle) {
1929            self.cleanup.fetch_add(1, Ordering::SeqCst);
1930        }
1931    }
1932
1933    fn route_bind_frame(channel: u16, epoch: u32, corr: u64) -> Frame {
1934        let body = serde_json::to_vec(&ModuleControlRequest::RouteBind {
1935            route_channel: channel,
1936            epoch,
1937            target: RouteTarget::ToolProvider {
1938                module_id: "provider".to_string(),
1939            },
1940            identity: BindIdentity {
1941                project_root: PathBuf::from("/tmp/project"),
1942                harness: "test".to_string(),
1943                session: "bind".to_string(),
1944            },
1945            principal: None,
1946            consumer_capabilities: None,
1947            admission_facts: None,
1948        })
1949        .unwrap();
1950        Frame::build(FrameType::Request, control_flags(), 0, 0, corr, body).unwrap()
1951    }
1952
1953    #[tokio::test]
1954    async fn on_bound_runs_only_after_ack_queue_and_handle_install() {
1955        let (module_handle, mut rx) = test_module_handle(&[]);
1956        let tx = module_handle.shared.lock_inner().writer.clone().unwrap();
1957        let bind_emit_rejected = Arc::new(AtomicUsize::new(0));
1958        let bound = Arc::new(AtomicUsize::new(0));
1959        let cleanup = Arc::new(AtomicUsize::new(0));
1960        let handler = Arc::new(BindOrderingHandler {
1961            module_handle: module_handle.clone(),
1962            bind_emit_rejected: Arc::clone(&bind_emit_rejected),
1963            bound: Arc::clone(&bound),
1964            cleanup: Arc::clone(&cleanup),
1965        });
1966
1967        assert!(handle_frame(
1968            route_bind_frame(8, 4, 90),
1969            &tx,
1970            handler,
1971            RequestDispatcher::new(),
1972            module_handle.clone(),
1973        )
1974        .await
1975        .unwrap());
1976        assert_eq!(bind_emit_rejected.load(Ordering::SeqCst), 1);
1977        assert_eq!(bound.load(Ordering::SeqCst), 1);
1978        assert_eq!(cleanup.load(Ordering::SeqCst), 0);
1979
1980        let ack = rx.recv().await.unwrap();
1981        let push = rx.recv().await.unwrap();
1982        assert_eq!(ack.header.ty, FrameType::Response);
1983        assert_eq!(ack.header.channel, 0);
1984        assert_eq!(push.header.ty, FrameType::Push);
1985        assert_eq!((push.header.channel, push.header.epoch), (8, 4));
1986        assert_eq!(
1987            push.header.flags.admission_class(),
1988            Some(AdmissionClass::Expedite)
1989        );
1990
1991        let captured = RouteHandle::new(8, 4, 1);
1992        assert!(module_handle.remove_route(captured).unwrap());
1993        let stale = module_handle
1994            .push(&captured, Vec::new(), None)
1995            .await
1996            .unwrap_err();
1997        assert!(matches!(stale, SubcModuleError::StaleRouteHandle(_)));
1998        assert!(rx.try_recv().is_err());
1999    }
2000
2001    struct RejectingHandler {
2002        bound: Arc<AtomicUsize>,
2003        cleanup: Arc<AtomicUsize>,
2004    }
2005
2006    #[async_trait]
2007    impl ModuleHandler for RejectingHandler {
2008        async fn handle(&self, _ctx: RequestCtx, _body: Vec<u8>) -> HandlerOutcome {
2009            HandlerOutcome::Response(Vec::new())
2010        }
2011
2012        async fn on_bind(&self, _req: &RouteBindRequest) -> BindDecision {
2013            BindDecision::reject("no", "rejected")
2014        }
2015
2016        async fn on_bound(&self, _handle: &RouteHandle) {
2017            self.bound.fetch_add(1, Ordering::SeqCst);
2018        }
2019
2020        async fn on_route_gone(&self, _handle: &RouteHandle) {
2021            self.cleanup.fetch_add(1, Ordering::SeqCst);
2022        }
2023    }
2024
2025    #[tokio::test]
2026    async fn rejected_bind_cleans_up_without_installing_or_calling_on_bound() {
2027        let (module_handle, mut rx) = test_module_handle(&[]);
2028        let tx = module_handle.shared.lock_inner().writer.clone().unwrap();
2029        let bound = Arc::new(AtomicUsize::new(0));
2030        let cleanup = Arc::new(AtomicUsize::new(0));
2031        let handler = Arc::new(RejectingHandler {
2032            bound: Arc::clone(&bound),
2033            cleanup: Arc::clone(&cleanup),
2034        });
2035        handle_frame(
2036            route_bind_frame(6, 3, 91),
2037            &tx,
2038            handler,
2039            RequestDispatcher::new(),
2040            module_handle.clone(),
2041        )
2042        .await
2043        .unwrap();
2044        assert_eq!(rx.recv().await.unwrap().header.ty, FrameType::Error);
2045        assert_eq!(bound.load(Ordering::SeqCst), 0);
2046        assert_eq!(cleanup.load(Ordering::SeqCst), 1);
2047        assert!(matches!(
2048            module_handle.validate_route(RouteHandle::new(6, 3, 1)),
2049            Err(SubcModuleError::StaleRouteHandle(_))
2050        ));
2051    }
2052
2053    struct RebindCountingHandler {
2054        bound: Arc<AtomicUsize>,
2055        cleanup: Arc<AtomicUsize>,
2056    }
2057
2058    #[async_trait]
2059    impl ModuleHandler for RebindCountingHandler {
2060        async fn handle(&self, _ctx: RequestCtx, _body: Vec<u8>) -> HandlerOutcome {
2061            HandlerOutcome::Response(Vec::new())
2062        }
2063
2064        async fn on_bound(&self, _handle: &RouteHandle) {
2065            self.bound.fetch_add(1, Ordering::SeqCst);
2066        }
2067
2068        async fn on_route_gone(&self, _handle: &RouteHandle) {
2069            self.cleanup.fetch_add(1, Ordering::SeqCst);
2070        }
2071    }
2072
2073    // Wire spec 3.3.0: a bind on an installed channel with a strictly higher epoch
2074    // replaces the stale install (the daemon freed that binding; its route-gone
2075    // GOODBYE is best-effort and can be dropped), firing the replaced install's
2076    // route-gone teardown. Equal-or-lower epoch is a protocol violation: rejected,
2077    // installed route untouched.
2078    #[tokio::test]
2079    async fn rebind_on_installed_channel_replaces_on_higher_epoch_only() {
2080        let (module_handle, mut rx) = test_module_handle(&[]);
2081        let tx = module_handle.shared.lock_inner().writer.clone().unwrap();
2082        let bound = Arc::new(AtomicUsize::new(0));
2083        let cleanup = Arc::new(AtomicUsize::new(0));
2084        let handler = Arc::new(RebindCountingHandler {
2085            bound: Arc::clone(&bound),
2086            cleanup: Arc::clone(&cleanup),
2087        });
2088
2089        // Install epoch 4 on channel 8.
2090        handle_frame(
2091            route_bind_frame(8, 4, 90),
2092            &tx,
2093            Arc::clone(&handler),
2094            RequestDispatcher::new(),
2095            module_handle.clone(),
2096        )
2097        .await
2098        .unwrap();
2099        assert_eq!(rx.recv().await.unwrap().header.ty, FrameType::Response);
2100        assert_eq!(
2101            (bound.load(Ordering::SeqCst), cleanup.load(Ordering::SeqCst)),
2102            (1, 0)
2103        );
2104
2105        // Same epoch: rejected, install untouched, no teardown fired.
2106        handle_frame(
2107            route_bind_frame(8, 4, 91),
2108            &tx,
2109            Arc::clone(&handler),
2110            RequestDispatcher::new(),
2111            module_handle.clone(),
2112        )
2113        .await
2114        .unwrap();
2115        let reject = rx.recv().await.unwrap();
2116        assert_eq!(reject.header.ty, FrameType::Error);
2117        assert_eq!(
2118            (bound.load(Ordering::SeqCst), cleanup.load(Ordering::SeqCst)),
2119            (1, 0)
2120        );
2121        module_handle
2122            .validate_route(RouteHandle::new(8, 4, 1))
2123            .expect("epoch-4 install must survive the rejected rebind");
2124
2125        // Lower epoch: same rejection shape.
2126        handle_frame(
2127            route_bind_frame(8, 3, 92),
2128            &tx,
2129            Arc::clone(&handler),
2130            RequestDispatcher::new(),
2131            module_handle.clone(),
2132        )
2133        .await
2134        .unwrap();
2135        assert_eq!(rx.recv().await.unwrap().header.ty, FrameType::Error);
2136        assert_eq!(
2137            (bound.load(Ordering::SeqCst), cleanup.load(Ordering::SeqCst)),
2138            (1, 0)
2139        );
2140
2141        // Strictly higher epoch: implicit replace — stale install torn down
2142        // (route-gone fired exactly once), new epoch installed and bound.
2143        handle_frame(
2144            route_bind_frame(8, 5, 93),
2145            &tx,
2146            Arc::clone(&handler),
2147            RequestDispatcher::new(),
2148            module_handle.clone(),
2149        )
2150        .await
2151        .unwrap();
2152        assert_eq!(rx.recv().await.unwrap().header.ty, FrameType::Response);
2153        assert_eq!(
2154            (bound.load(Ordering::SeqCst), cleanup.load(Ordering::SeqCst)),
2155            (2, 1)
2156        );
2157        assert!(matches!(
2158            module_handle.validate_route(RouteHandle::new(8, 4, 1)),
2159            Err(SubcModuleError::StaleRouteHandle(_))
2160        ));
2161        module_handle
2162            .validate_route(RouteHandle::new(8, 5, 1))
2163            .expect("epoch-5 install must be live after implicit replace");
2164    }
2165
2166    #[test]
2167    fn module_control_corr_is_monotonic_and_exhausts_without_wrap() {
2168        let (module_handle, _rx) = test_module_handle(&[MODULE_TO_SUBC_OP_CATALOG_UPDATE]);
2169        let mut inner = module_handle.shared.lock_inner();
2170        inner.next_corr = Some(u64::MAX);
2171        assert_eq!(next_module_control_corr(&mut inner), Some(u64::MAX));
2172        assert_eq!(next_module_control_corr(&mut inner), None);
2173    }
2174}