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#[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
105pub 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#[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 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 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 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#[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#[async_trait]
490pub trait ModuleHandler: Send + Sync + 'static {
491 async fn handle(&self, ctx: RequestCtx, body: Vec<u8>) -> HandlerOutcome;
496
497 async fn on_hello_ack(&self, _ack: &ModuleHelloAckBody) {}
500
501 async fn on_bind(&self, _req: &RouteBindRequest) -> BindDecision {
503 BindDecision::accept()
504 }
505
506 async fn on_bound(&self, _handle: &RouteHandle) {}
508
509 async fn health(&self) -> HealthReport {
530 HealthReport {
543 detail: Some("no health implementation; inherited default".to_string()),
544 ..HealthReport::ok()
545 }
546 }
547
548 async fn on_route_gone(&self, _handle: &RouteHandle) {}
550}
551
552#[derive(Debug, Clone, PartialEq, Eq)]
554pub enum HandlerOutcome {
555 Response(Vec<u8>),
557 Error { code: String, message: String },
559 Streamed,
562}
563
564#[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 pub fn route_handle(&self) -> RouteHandle {
579 self.handle
580 }
581
582 pub fn corr(&self) -> u64 {
584 self.corr
585 }
586
587 pub async fn emit(&self, body: Vec<u8>) -> Result<(), SubcModuleError> {
590 self.emit_with_admission(body, None).await
591 }
592
593 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 pub fn cancelled(&self) -> WaitForCancellationFuture<'_> {
614 self.cancelled.cancelled()
615 }
616
617 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#[derive(Debug, Clone)]
645pub struct RouteBindRequest {
646 pub handle: RouteHandle,
647 pub target: RouteTarget,
648 pub identity: BindIdentity,
649 pub principal: Option<Principal>,
650 pub consumer_capabilities: Option<Vec<String>>,
655 pub admission_facts: Option<serde_json::Value>,
657}
658
659#[derive(Debug, Clone)]
661pub struct BindDecision {
662 kind: BindDecisionKind,
663}
664
665impl BindDecision {
666 pub fn accept() -> Self {
668 Self {
669 kind: BindDecisionKind::Accept,
670 }
671 }
672
673 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
690pub 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
705pub 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
719pub 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 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 (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 Ok(None) => return Ok(()),
800 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 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 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 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 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 assert_eq!(status, HealthStatus::Ok);
1641 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 #[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 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 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 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 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}