Skip to main content

unb_runtime/
core_runtime.rs

1use std::collections::HashMap;
2use std::collections::VecDeque;
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::Arc;
7
8use unb_core::{
9    ClientDelivery as CoreClientDelivery, ClientOperationId, CoreEffect, CoreInput, EffectId,
10    Envelope, ProtocolCore, RetirementReason, SendResult, SessionId,
11};
12use bytes::Bytes;
13use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore};
14use web_time::Instant;
15
16use crate::cancellation::CancellationToken;
17use crate::client::{ClientDelivery, ClientSession};
18use crate::error::WsError;
19use crate::transport::{Pipe, PipeReader, PipeWriter};
20use crate::wire::{Directive, Wire};
21use crate::{BodyStream, WireBody};
22
23const CORE_QUEUE: usize = 256;
24const SESSION_QUEUE: usize = 64;
25const OUTBOUND_QUEUE: usize = 8_192;
26const OUTBOUND_BYTE_BUDGET: usize = 64 * 1024 * 1024;
27const BODY_WINDOW_BYTES: usize = 256 * 1024;
28const WRITE_BATCH: usize = 8;
29const RELAY_BACKPRESSURE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
30const TERMINAL_SEND_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
31const CLIENT_DELIVERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
32const ESTABLISHMENT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
33const PENDING_MODE_WRITES_MAX: usize = 128;
34const OPERATION_WRITE_QUEUE: usize = OUTBOUND_QUEUE;
35const OPERATION_ACTIVE_MAX: usize = 128;
36const OPERATION_BODY_CHUNK_BYTES: usize = 64 * 1024;
37
38#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
39fn clock_now() -> Instant {
40    n0_future::time::Instant::now().into_std()
41}
42
43#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
44async fn sleep_to(deadline: Instant) {
45    n0_future::time::sleep_until(n0_future::time::Instant::from_std(deadline)).await
46}
47
48#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
49fn clock_now() -> Instant {
50    n0_future::time::Instant::now()
51}
52
53#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
54async fn sleep_to(deadline: Instant) {
55    n0_future::time::sleep_until(deadline).await
56}
57
58#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
59pub use tokio::runtime::Handle as RuntimeHandle;
60
61#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
62#[derive(Clone, Debug)]
63pub struct RuntimeHandle;
64
65#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
66impl RuntimeHandle {
67    pub fn current() -> Self {
68        RuntimeHandle
69    }
70
71    pub fn spawn<F>(&self, future: F)
72    where
73        F: Future<Output = ()> + 'static,
74    {
75        n0_future::task::spawn(future);
76    }
77}
78
79pub type EffectFuture = Pin<Box<dyn Future<Output = Option<CoreInput>> + Send>>;
80
81pub trait EffectExecutor: Send + Sync + 'static {
82    fn execute(&self, effect: CoreEffect, handle: ProtocolCoreHandle) -> EffectFuture;
83}
84
85#[derive(Clone, Debug)]
86pub enum SessionOutcome {
87    Established,
88    Retired(RetirementReason),
89}
90
91pub trait SessionHandler: Send + 'static {
92    fn deliver(&mut self, envelope: Envelope) -> impl Future<Output = ()> + Send;
93
94    fn stream_closed(&mut self, _operation: ClientOperationId) -> impl Future<Output = ()> + Send {
95        async {}
96    }
97}
98
99impl SessionHandler for () {
100    async fn deliver(&mut self, _envelope: Envelope) {}
101}
102
103impl SessionHandler for mpsc::Sender<Envelope> {
104    async fn deliver(&mut self, envelope: Envelope) {
105        let _ = self.send(envelope).await;
106    }
107}
108
109impl SessionHandler for tokio::sync::broadcast::Sender<Envelope> {
110    async fn deliver(&mut self, envelope: Envelope) {
111        let _ = self.send(envelope);
112    }
113}
114
115type TokenBodies = Arc<std::sync::Mutex<HashMap<(SessionId, String), WireBody>>>;
116
117type BodyPumps = Arc<std::sync::Mutex<HashMap<(SessionId, String), (String, CancellationToken)>>>;
118
119const SESSION_BODY_TOKENS_MAX: usize = OUTBOUND_QUEUE;
120const SESSION_BODY_BYTES_MAX: usize = OUTBOUND_BYTE_BUDGET;
121
122fn mint_body_token() -> String {
123    use std::hash::{BuildHasher, Hasher};
124    type TokenKeys = (
125        std::collections::hash_map::RandomState,
126        std::collections::hash_map::RandomState,
127    );
128    static KEYS: std::sync::OnceLock<TokenKeys> = std::sync::OnceLock::new();
129    static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
130    let (low, high) = KEYS.get_or_init(TokenKeys::default);
131    let serial = NEXT.fetch_add(1, Ordering::Relaxed);
132    let mut a = low.build_hasher();
133    a.write_u64(serial);
134    let mut b = high.build_hasher();
135    b.write_u64(serial);
136    format!("b{:016x}{:016x}", a.finish(), b.finish())
137}
138
139fn application_core_input(
140    handle: &ProtocolCoreHandle,
141    session: &SessionId,
142    mut envelope: Envelope,
143) -> Result<CoreInput, unb_core::CoreError> {
144    if envelope.kind == unb_core::Kind::Discover {
145        return Ok(CoreInput::FrameReceived {
146            session: session.clone(),
147            envelope,
148        });
149    }
150    if !(envelope.kind.is_application_request()
151        || (envelope.corr.is_some()
152            && (envelope.kind.is_application_response()
153                || envelope.kind == unb_core::Kind::Cancel)))
154    {
155        return Ok(CoreInput::FrameReceived {
156            session: session.clone(),
157            envelope,
158        });
159    }
160    if envelope.body_token.is_some() && !envelope.payload.is_empty() {
161        return Err(unb_core::CoreError::Malformed(
162            "application frame cannot carry both inline and referenced bodies".into(),
163        ));
164    }
165    if envelope.body_token.is_none()
166        && envelope.kind != unb_core::Kind::Error
167        && !envelope.payload.is_empty()
168    {
169        let payload = std::mem::take(&mut envelope.payload);
170        envelope.body_token = Some(
171            handle
172                .register_body(session, WireBody::Bytes(payload))?
173                .to_string(),
174        );
175    }
176    Ok(CoreInput::ApplicationFrameReceived {
177        session: session.clone(),
178        frame: unb_core::ApplicationFrame::from_envelope(&envelope)?,
179    })
180}
181
182fn register_outbound_body(
183    handle: &ProtocolCoreHandle,
184    session: &SessionId,
185    payload: Bytes,
186    body: Option<BodyStream>,
187) -> Result<Option<unb_core::BodyId>, unb_core::CoreError> {
188    if body.is_some() && !payload.is_empty() {
189        return Err(unb_core::CoreError::Malformed(
190            "an operation cannot provide both unary bytes and a streaming body".into(),
191        ));
192    }
193    let body = match body {
194        Some(body) => Some(WireBody::Stream(body)),
195        None if !payload.is_empty() => Some(WireBody::Bytes(payload)),
196        None => None,
197    };
198    body.map(|body| handle.register_body(session, body))
199        .transpose()
200}
201
202#[derive(Clone)]
203pub struct ProtocolCoreHandle {
204    inputs: mpsc::Sender<ActorInput>,
205    cancellation: CancellationToken,
206    token_bodies: TokenBodies,
207    body_pumps: BodyPumps,
208}
209
210impl ProtocolCoreHandle {
211    pub fn spawn(
212        core: ProtocolCore,
213        executor: Arc<dyn EffectExecutor>,
214        cancellation: CancellationToken,
215        runtime: &RuntimeHandle,
216    ) -> Self {
217        let (inputs, receiver) = mpsc::channel(CORE_QUEUE);
218        let handle = Self {
219            inputs,
220            cancellation,
221            token_bodies: Arc::new(std::sync::Mutex::new(HashMap::new())),
222            body_pumps: Arc::new(std::sync::Mutex::new(HashMap::new())),
223        };
224        runtime.spawn(CoreRuntime::new(core, receiver, handle.clone(), executor).run());
225        handle
226    }
227
228    pub fn claim_body(&self, session: &SessionId, token: &str) -> Option<WireBody> {
229        self.token_bodies
230            .lock()
231            .expect("body token registry")
232            .remove(&(session.clone(), token.to_string()))
233    }
234
235    pub fn claim_body_stream(&self, session: &SessionId, token: &str) -> Option<BodyStream> {
236        self.claim_body(session, token).map(wire_body_stream)
237    }
238
239    pub fn register_body(
240        &self,
241        session: &SessionId,
242        body: WireBody,
243    ) -> Result<unb_core::BodyId, unb_core::CoreError> {
244        let mut registry = self.token_bodies.lock().expect("body token registry");
245        let owned = registry
246            .keys()
247            .filter(|(owner, _)| owner == session)
248            .count();
249        if owned >= SESSION_BODY_TOKENS_MAX {
250            return Err(unb_core::CoreError::EffectQueueFull);
251        }
252        let incoming = match &body {
253            WireBody::Bytes(payload) => payload.len(),
254            WireBody::Stream(_) => 0,
255        };
256        let retained = registry
257            .iter()
258            .filter_map(|((owner, _), body)| {
259                (owner == session).then_some(match body {
260                    WireBody::Bytes(payload) => payload.len(),
261                    WireBody::Stream(_) => 0,
262                })
263            })
264            .sum::<usize>();
265        if retained.saturating_add(incoming) > SESSION_BODY_BYTES_MAX {
266            return Err(unb_core::CoreError::BodyRegistryFull(
267                SESSION_BODY_BYTES_MAX,
268            ));
269        }
270        let token = mint_body_token();
271        registry.insert((session.clone(), token.clone()), body);
272        Ok(unb_core::BodyId::from(token))
273    }
274
275    fn register_body_pump(
276        &self,
277        session: &SessionId,
278        corr: &str,
279        token: &str,
280        cancellation: CancellationToken,
281    ) {
282        self.body_pumps.lock().expect("body pump registry").insert(
283            (session.clone(), corr.to_string()),
284            (token.to_string(), cancellation),
285        );
286    }
287
288    fn cancel_body_pump(&self, session: &SessionId, corr: &str) {
289        let cancelled = self
290            .body_pumps
291            .lock()
292            .expect("body pump registry")
293            .remove(&(session.clone(), corr.to_string()));
294        if let Some((_, cancellation)) = cancelled {
295            cancellation.cancel();
296        }
297    }
298
299    fn release_body_pump(&self, session: &SessionId, corr: &str, token: &str) {
300        let mut pumps = self.body_pumps.lock().expect("body pump registry");
301        let key = (session.clone(), corr.to_string());
302        if pumps.get(&key).is_some_and(|(owner, _)| owner == token) {
303            pumps.remove(&key);
304        }
305    }
306
307    pub async fn submit(&self, input: CoreInput) -> Result<(), WsError> {
308        self.inputs
309            .send(ActorInput::Core(input))
310            .await
311            .map_err(|_| WsError::Gone)
312    }
313
314    pub async fn submit_batch(&self, inputs: Vec<CoreInput>) -> Result<(), WsError> {
315        self.inputs
316            .send(ActorInput::CoreBatch(inputs))
317            .await
318            .map_err(|_| WsError::Gone)
319    }
320
321    pub async fn install<F>(&self, input: CoreInput, publish: F) -> Result<(), WsError>
322    where
323        F: FnOnce() + Send + 'static,
324    {
325        let (reply, response) = oneshot::channel();
326        self.inputs
327            .send(ActorInput::Install {
328                input,
329                publish: Box::new(publish),
330                reply,
331            })
332            .await
333            .map_err(|_| WsError::Gone)?;
334        response.await.map_err(|_| WsError::Gone)?
335    }
336
337    pub async fn attach<H: SessionHandler>(
338        &self,
339        session: SessionId,
340        transport: Pipe,
341        expected_peer: Option<String>,
342        handler: H,
343    ) -> Result<Arc<Wire>, WsError> {
344        self.attach_with_ceiling(
345            session,
346            transport,
347            expected_peer,
348            handler,
349            unb_transport::DEFAULT_MAX_FRAME_SIZE,
350        )
351        .await
352    }
353
354    pub async fn attach_with_ceiling<H: SessionHandler>(
355        &self,
356        session: SessionId,
357        transport: Pipe,
358        expected_peer: Option<String>,
359        handler: H,
360        ws_collect_ceiling: usize,
361    ) -> Result<Arc<Wire>, WsError> {
362        let cancellation = self.cancellation.child_token();
363        let (wire, directives, outcome, routes) = Wire::channel(cancellation.clone());
364        let (reply, response) = oneshot::channel();
365        self.inputs
366            .send(ActorInput::Attach {
367                session,
368                transport,
369                expected_peer,
370                directives,
371                handler: handler_channel(handler),
372                client: wire.client_session(),
373                cancellation,
374                establish_peer: true,
375                outcome,
376                routes,
377                ws_collect_ceiling,
378                reply,
379            })
380            .await
381            .map_err(|_| WsError::Gone)?;
382        response.await.map_err(|_| WsError::Gone)??;
383        Ok(wire)
384    }
385
386    pub fn shutdown(&self) {
387        self.cancellation.cancel();
388    }
389
390    pub(crate) fn open(transport: Pipe, runtime: &RuntimeHandle) -> (Arc<Wire>, Self) {
391        struct StandaloneExecutor;
392
393        impl EffectExecutor for StandaloneExecutor {
394            fn execute(&self, _effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
395                Box::pin(async { None })
396            }
397        }
398
399        let cancellation = CancellationToken::new();
400        let handle = Self::spawn(
401            ProtocolCore::new("client"),
402            Arc::new(StandaloneExecutor),
403            cancellation.clone(),
404            runtime,
405        );
406        let session_cancellation = cancellation.child_token();
407        let (wire, directives, outcome, routes) =
408            Wire::standalone_channel(session_cancellation.clone());
409        let (reply, _response) = oneshot::channel();
410        handle
411            .inputs
412            .try_send(ActorInput::Attach {
413                session: SessionId::from("client-session"),
414                transport,
415                expected_peer: None,
416                directives,
417                handler: handler_channel(wire.observation_sender()),
418                client: wire.client_session(),
419                cancellation: session_cancellation,
420                establish_peer: false,
421                outcome,
422                routes,
423                ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
424                reply,
425            })
426            .expect("new protocol core input queue");
427        (wire, handle)
428    }
429}
430
431fn wire_body_stream(body: WireBody) -> BodyStream {
432    match body {
433        WireBody::Bytes(bytes) => Box::pin(futures_util::stream::once(async move { Ok(bytes) })),
434        WireBody::Stream(stream) => stream,
435    }
436}
437
438enum ActorInput {
439    Core(CoreInput),
440    CoreBatch(Vec<CoreInput>),
441    Install {
442        input: CoreInput,
443        publish: Box<dyn FnOnce() + Send>,
444        reply: oneshot::Sender<Result<(), WsError>>,
445    },
446    Directive(SessionId, Directive),
447    Attach {
448        session: SessionId,
449        transport: Pipe,
450        expected_peer: Option<String>,
451        directives: mpsc::Receiver<Directive>,
452        handler: mpsc::Sender<HandlerInput>,
453        client: ClientSession,
454        cancellation: CancellationToken,
455        establish_peer: bool,
456        outcome: tokio::sync::watch::Sender<Option<SessionOutcome>>,
457        routes: tokio::sync::watch::Sender<bool>,
458        ws_collect_ceiling: usize,
459        reply: oneshot::Sender<Result<(), WsError>>,
460    },
461    Write {
462        session: SessionId,
463        envelope: Envelope,
464        completion: Option<EffectId>,
465        relay: bool,
466    },
467    TransportModeSelected {
468        session: SessionId,
469        mode: unb_transport::TransportMode,
470    },
471}
472
473struct SessionIo {
474    outbound: mpsc::Sender<WriteRequest>,
475    byte_budget: Arc<Semaphore>,
476    ws_collect_ceiling: usize,
477    bodies: crate::transport::SessionStreams,
478    handler: mpsc::Sender<HandlerInput>,
479    client: ClientSession,
480    cancellation: CancellationToken,
481    reservation_pending: Arc<AtomicBool>,
482    establish_peer: bool,
483    deadline: tokio::sync::watch::Sender<Option<Instant>>,
484    routes: tokio::sync::watch::Sender<bool>,
485    transport_mode: Option<unb_transport::TransportMode>,
486    pending_mode_writes: VecDeque<PendingModeWrite>,
487    operation_lanes: OperationLanes,
488    operation_slots: Arc<Semaphore>,
489}
490
491struct PendingModeWrite {
492    envelope: Envelope,
493    completion: Option<EffectId>,
494    relay: bool,
495}
496
497type OperationLanes = Arc<std::sync::Mutex<HashMap<String, OperationLane>>>;
498
499#[derive(Clone)]
500struct OperationLane {
501    outbound: Option<mpsc::Sender<OperationWriteRequest>>,
502    control: Option<unb_transport::operation::OperationControl>,
503}
504
505struct OperationWriteRequest {
506    envelope: Envelope,
507    completion: Option<EffectId>,
508    relay: bool,
509    body: Option<BodyStream>,
510    _budget: OwnedSemaphorePermit,
511}
512
513struct WriteRequest {
514    envelope: Envelope,
515    completion: Option<EffectId>,
516    relay: bool,
517    _budget: OwnedSemaphorePermit,
518}
519
520impl WriteRequest {
521    fn weight(envelope: &Envelope) -> u32 {
522        let head = envelope.subject.len()
523            + envelope.id.len()
524            + envelope.corr.as_ref().map_or(0, String::len)
525            + envelope.path.iter().map(String::len).sum::<usize>()
526            + envelope
527                .headers
528                .iter()
529                .map(|(name, value)| name.len() + value.as_str().map_or(16, str::len))
530                .sum::<usize>();
531        (envelope.payload.len() + head).clamp(1, OUTBOUND_BYTE_BUDGET) as u32
532    }
533}
534
535#[derive(Debug)]
536pub(crate) struct MeteredChunk {
537    pub(crate) bytes: Bytes,
538    _window: OwnedSemaphorePermit,
539    _budget: OwnedSemaphorePermit,
540}
541
542pub(crate) struct MeteredBody {
543    stream: BodyStream,
544    budget: Arc<Semaphore>,
545    window: Arc<Semaphore>,
546    pending: Option<Bytes>,
547    cancellation: CancellationToken,
548}
549
550impl MeteredBody {
551    pub(crate) fn new(
552        stream: BodyStream,
553        budget: Arc<Semaphore>,
554        cancellation: CancellationToken,
555    ) -> MeteredBody {
556        MeteredBody {
557            stream,
558            budget,
559            window: Arc::new(Semaphore::new(BODY_WINDOW_BYTES)),
560            pending: None,
561            cancellation,
562        }
563    }
564
565    pub(crate) async fn next_chunk(&mut self) -> Result<Option<MeteredChunk>, WsError> {
566        let chunk = match self.pending.take() {
567            Some(chunk) => chunk,
568            None => {
569                let next = tokio::select! {
570                    biased;
571                    () = self.cancellation.cancelled() => return Ok(None),
572                    next = std::future::poll_fn(|cx| self.stream.as_mut().poll_next(cx)) => next,
573                };
574                match next {
575                    Some(Ok(chunk)) => chunk,
576                    Some(Err(error)) => return Err(error.into()),
577                    None => return Ok(None),
578                }
579            }
580        };
581        let len = chunk.len();
582        if len > unb_transport::DEFAULT_MAX_FRAME_SIZE {
583            return Err(unb_core::CoreError::BodyTooLarge(
584                unb_transport::DEFAULT_MAX_FRAME_SIZE,
585            )
586            .into());
587        }
588        self.pending = Some(chunk);
589        let window_weight = len.clamp(1, BODY_WINDOW_BYTES) as u32;
590        let budget_weight = len.max(1) as u32;
591        let acquire = async {
592            let window = self
593                .window
594                .clone()
595                .acquire_many_owned(window_weight)
596                .await
597                .map_err(|_| unb_core::CoreError::BodyStalled)?;
598            let budget = self
599                .budget
600                .clone()
601                .acquire_many_owned(budget_weight)
602                .await
603                .map_err(|_| unb_core::CoreError::BodyStalled)?;
604            Ok::<_, unb_core::CoreError>((window, budget))
605        };
606        let acquired = tokio::select! {
607            biased;
608            () = self.cancellation.cancelled() => return Ok(None),
609            acquired = n0_future::time::timeout(RELAY_BACKPRESSURE_TIMEOUT, acquire) => acquired,
610        };
611        match acquired {
612            Ok(Ok((window, budget))) => Ok(Some(MeteredChunk {
613                bytes: self.pending.take().expect("pending chunk present"),
614                _window: window,
615                _budget: budget,
616            })),
617            Ok(Err(error)) => Err(error.into()),
618            Err(_) => Err(unb_core::CoreError::BodyStalled.into()),
619        }
620    }
621}
622
623struct CoreRuntime {
624    core: ProtocolCore,
625    inputs: mpsc::Receiver<ActorInput>,
626    handle: ProtocolCoreHandle,
627    executor: Arc<dyn EffectExecutor>,
628    sessions: HashMap<SessionId, SessionIo>,
629    outcomes: HashMap<SessionId, tokio::sync::watch::Sender<Option<SessionOutcome>>>,
630    client_starts: HashMap<SessionId, PendingClientStart>,
631    pending_bodies: std::sync::Mutex<HashMap<(SessionId, String), BodyStream>>,
632}
633
634struct PendingClientStart {
635    sender: mpsc::Sender<ClientDelivery>,
636    reply: oneshot::Sender<Result<ClientOperationId, unb_core::CoreError>>,
637}
638
639impl CoreRuntime {
640    fn new(
641        core: ProtocolCore,
642        inputs: mpsc::Receiver<ActorInput>,
643        handle: ProtocolCoreHandle,
644        executor: Arc<dyn EffectExecutor>,
645    ) -> Self {
646        Self {
647            core,
648            inputs,
649            handle,
650            executor,
651            sessions: HashMap::new(),
652            outcomes: HashMap::new(),
653            client_starts: HashMap::new(),
654            pending_bodies: std::sync::Mutex::new(HashMap::new()),
655        }
656    }
657
658    async fn run(mut self) {
659        loop {
660            let input = tokio::select! {
661                biased;
662                () = self.handle.cancellation.cancelled() => return,
663                input = self.inputs.recv() => input,
664            };
665            let Some(input) = input else { return };
666            match input {
667                ActorInput::Core(input) => {
668                    let _ = self.core.handle(clock_now(), input);
669                }
670                ActorInput::CoreBatch(inputs) => {
671                    let now = clock_now();
672                    for input in inputs {
673                        let _ = self.core.handle(now, input);
674                    }
675                }
676                ActorInput::Install {
677                    input,
678                    publish,
679                    reply,
680                } => {
681                    let result = self.core.handle(clock_now(), input).map_err(WsError::from);
682                    if result.is_ok() {
683                        publish();
684                    }
685                    let _ = reply.send(result);
686                }
687                ActorInput::Directive(session, directive) => {
688                    self.directive(&session, directive).await;
689                }
690                ActorInput::Attach {
691                    session,
692                    transport,
693                    expected_peer,
694                    directives,
695                    handler,
696                    client,
697                    cancellation,
698                    establish_peer,
699                    outcome,
700                    routes,
701                    ws_collect_ceiling,
702                    reply,
703                } => {
704                    let result = self
705                        .attach(
706                            session,
707                            transport,
708                            expected_peer,
709                            directives,
710                            handler,
711                            client,
712                            cancellation,
713                            establish_peer,
714                            outcome,
715                            routes,
716                            ws_collect_ceiling,
717                        )
718                        .await;
719                    let _ = reply.send(result);
720                }
721                ActorInput::Write {
722                    session,
723                    envelope,
724                    completion,
725                    relay,
726                } => {
727                    self.write(&session, envelope, completion, relay).await;
728                }
729                ActorInput::TransportModeSelected { session, mode } => {
730                    let pending = if let Some(io) = self.sessions.get_mut(&session) {
731                        io.transport_mode = Some(mode);
732                        std::mem::take(&mut io.pending_mode_writes)
733                    } else {
734                        VecDeque::new()
735                    };
736                    if mode == unb_transport::TransportMode::OperationStreamsV1 {
737                        self.start_operation_acceptor(&session);
738                    }
739                    for pending in pending {
740                        self.write(
741                            &session,
742                            pending.envelope,
743                            pending.completion,
744                            pending.relay,
745                        )
746                        .await;
747                    }
748                }
749            }
750            self.drain().await;
751        }
752    }
753
754    #[allow(clippy::too_many_arguments)]
755    async fn attach(
756        &mut self,
757        session: SessionId,
758        transport: Pipe,
759        expected_peer: Option<String>,
760        directives: mpsc::Receiver<Directive>,
761        handler: mpsc::Sender<HandlerInput>,
762        client: ClientSession,
763        cancellation: CancellationToken,
764        establish_peer: bool,
765        outcome: tokio::sync::watch::Sender<Option<SessionOutcome>>,
766        routes: tokio::sync::watch::Sender<bool>,
767        ws_collect_ceiling: usize,
768    ) -> Result<(), WsError> {
769        let initiator = transport.initiator();
770        let bodies = transport.streams();
771        let (reader, writer) = transport.split();
772        let (outbound, outbound_rx) = mpsc::channel(OUTBOUND_QUEUE);
773        let (deadline, deadline_rx) = tokio::sync::watch::channel(None);
774        self.core.handle(
775            clock_now(),
776            CoreInput::SessionOpened {
777                session: session.clone(),
778                initiator,
779                establish_peer,
780                expected_peer,
781            },
782        )?;
783        self.sessions.insert(
784            session.clone(),
785            SessionIo {
786                outbound,
787                byte_budget: Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET)),
788                ws_collect_ceiling,
789                bodies: bodies.clone(),
790                handler,
791                client,
792                cancellation: cancellation.clone(),
793                reservation_pending: Arc::new(AtomicBool::new(false)),
794                establish_peer,
795                deadline,
796                routes,
797                transport_mode: None,
798                pending_mode_writes: VecDeque::new(),
799                operation_lanes: Arc::new(std::sync::Mutex::new(HashMap::new())),
800                operation_slots: Arc::new(Semaphore::new(OPERATION_ACTIVE_MAX)),
801            },
802        );
803        self.outcomes.insert(session.clone(), outcome);
804        spawn_session_deadline(
805            session.clone(),
806            deadline_rx,
807            self.handle.clone(),
808            cancellation.clone(),
809        );
810        spawn_directives(
811            session.clone(),
812            directives,
813            self.handle.clone(),
814            cancellation.clone(),
815        );
816        spawn_reader(
817            session.clone(),
818            reader,
819            self.handle.clone(),
820            cancellation.clone(),
821            bodies,
822        );
823        spawn_writer(
824            session.clone(),
825            writer,
826            outbound_rx,
827            self.handle.clone(),
828            cancellation,
829        );
830        if establish_peer {
831            spawn_establishment_timeout(session, self.handle.clone());
832        }
833        Ok(())
834    }
835
836    async fn directive(&mut self, session: &SessionId, directive: Directive) {
837        match directive {
838            Directive::StartClientOperation {
839                subject,
840                kind,
841                payload,
842                hops,
843                headers,
844                body,
845                timeout,
846                sender,
847                reply,
848            } => {
849                let input = if kind == unb_core::Kind::Discover {
850                    if body.is_some() {
851                        Err(unb_core::CoreError::Malformed(
852                            "discovery does not accept a streaming body".into(),
853                        ))
854                    } else {
855                        serde_json::from_slice(&payload)
856                            .map(unb_core::OperationInput::Discovery)
857                            .map_err(|error| unb_core::CoreError::Malformed(error.to_string()))
858                    }
859                } else {
860                    register_outbound_body(&self.handle, session, payload, body)
861                        .map(unb_core::OperationInput::Body)
862                };
863                let input = match input {
864                    Ok(input) => input,
865                    Err(error) => {
866                        let _ = reply.send(Err(error));
867                        return;
868                    }
869                };
870                let registered_body = match &input {
871                    unb_core::OperationInput::Body(body) => body.clone(),
872                    unb_core::OperationInput::Discovery(_) => None,
873                };
874                let result = self.core.handle(
875                    clock_now(),
876                    CoreInput::StartClientOperation {
877                        session: session.clone(),
878                        subject,
879                        kind,
880                        input,
881                        hops,
882                        headers,
883                        timeout,
884                    },
885                );
886                if let Err(error) = result {
887                    if let Some(body) = &registered_body {
888                        let _ = self.handle.claim_body(session, body.as_str());
889                    }
890                    let _ = reply.send(Err(error));
891                } else {
892                    self.client_starts
893                        .insert(session.clone(), PendingClientStart { sender, reply });
894                }
895            }
896            Directive::OpenStream {
897                subject,
898                kind,
899                payload,
900                hops,
901                headers,
902                body,
903                opened,
904                reply,
905            } => {
906                let result = match register_outbound_body(&self.handle, session, payload, body) {
907                    Ok(body) => {
908                        let result = self.core.open_stream_body(
909                            session,
910                            &subject,
911                            kind,
912                            body.clone(),
913                            hops,
914                            headers,
915                        );
916                        if result.is_err() {
917                            if let Some(body) = &body {
918                                let _ = self.handle.claim_body(session, body.as_str());
919                            }
920                        }
921                        result
922                    }
923                    Err(error) => Err(error),
924                };
925                if let (Ok(corr), Some(opened)) = (&result, opened) {
926                    opened(corr.clone()).await;
927                }
928                let _ = reply.send(result);
929            }
930            Directive::Send { corr, payload } => {
931                match register_outbound_body(&self.handle, session, payload, None) {
932                    Ok(body) => {
933                        if self.core.send_body(session, &corr, body.clone()).is_err() {
934                            if let Some(body) = &body {
935                                let _ = self.handle.claim_body(session, body.as_str());
936                            }
937                        }
938                    }
939                    Err(error) => {
940                        let _ = self.core.fail(
941                            session,
942                            &corr,
943                            unb_core::ErrorCode::Busy,
944                            &error.to_string(),
945                        );
946                    }
947                }
948            }
949            Directive::Respond {
950                corr,
951                payload,
952                headers,
953                body,
954            } => match register_outbound_body(&self.handle, session, payload, body) {
955                Ok(body) => {
956                    if self
957                        .core
958                        .respond_body(session, &corr, body.clone(), headers)
959                        .is_err()
960                    {
961                        if let Some(body) = &body {
962                            let _ = self.handle.claim_body(session, body.as_str());
963                        }
964                    }
965                }
966                Err(error) => {
967                    let _ = self.core.fail(
968                        session,
969                        &corr,
970                        unb_core::ErrorCode::Busy,
971                        &error.to_string(),
972                    );
973                }
974            },
975            Directive::Fail {
976                corr,
977                code,
978                message,
979            } => {
980                let _ = self.core.fail(session, &corr, code, &message);
981            }
982            Directive::Cancel { corr } => {
983                let _ = self.core.cancel(session, &corr);
984            }
985            Directive::CancelClientOperation { operation } => {
986                let _ = self.core.handle(
987                    clock_now(),
988                    CoreInput::CancelClientOperation {
989                        session: session.clone(),
990                        operation,
991                    },
992                );
993            }
994            Directive::Control { kind, payload } => {
995                let _ = self.core.control(session, kind, payload);
996            }
997        }
998    }
999
1000    async fn drain(&mut self) {
1001        while let Some(effect) = self.core.poll_effect() {
1002            match effect {
1003                CoreEffect::SendFrame { session, envelope } => {
1004                    self.write(&session, envelope, None, false).await
1005                }
1006                CoreEffect::HandshakeEstablished { session, .. } => {
1007                    if let Some(bodies) =
1008                        self.sessions.get(&session).and_then(|io| io.bodies.clone())
1009                    {
1010                        let inputs = self.handle.inputs.clone();
1011                        let mode_session = session.clone();
1012                        n0_future::task::spawn(async move {
1013                            let mode = bodies.negotiate().await;
1014                            let _ = inputs
1015                                .send(ActorInput::TransportModeSelected {
1016                                    session: mode_session,
1017                                    mode,
1018                                })
1019                                .await;
1020                        });
1021                    }
1022                    if self
1023                        .sessions
1024                        .get(&session)
1025                        .is_some_and(|io| !io.establish_peer)
1026                    {
1027                        if let Some(outcome) = self.outcomes.get(&session) {
1028                            let _ = outcome.send(Some(SessionOutcome::Established));
1029                        }
1030                    }
1031                }
1032                CoreEffect::DeliverClient {
1033                    session,
1034                    operation,
1035                    delivery,
1036                } => {
1037                    let body = match &delivery {
1038                        CoreClientDelivery::Item(frame) | CoreClientDelivery::Terminal(frame) => {
1039                            frame
1040                                .body
1041                                .as_ref()
1042                                .and_then(|body| self.handle.claim_body(&session, body.as_str()))
1043                        }
1044                        _ => None,
1045                    };
1046                    let terminal = matches!(
1047                        delivery,
1048                        CoreClientDelivery::Terminal(_)
1049                            | CoreClientDelivery::Cancelled
1050                            | CoreClientDelivery::TimedOut
1051                            | CoreClientDelivery::SessionClosed
1052                    );
1053                    let client = self.sessions.get(&session).map(|io| io.client.clone());
1054                    if let Some(client) = client {
1055                        let delivered = n0_future::time::timeout(
1056                            CLIENT_DELIVERY_TIMEOUT,
1057                            client.deliver(&operation, delivery, body),
1058                        )
1059                        .await
1060                        .is_ok();
1061                        if !delivered {
1062                            client.abandon(&operation);
1063                            let _ = self.core.handle(
1064                                clock_now(),
1065                                CoreInput::CancelClientOperation {
1066                                    session: session.clone(),
1067                                    operation: operation.clone(),
1068                                },
1069                            );
1070                        }
1071                    }
1072                    if terminal {
1073                        self.handler(&session, HandlerInput::StreamClosed(operation))
1074                            .await;
1075                    }
1076                }
1077                CoreEffect::RegisterClient { session, operation } => {
1078                    let Some(start) = self.client_starts.remove(&session) else {
1079                        continue;
1080                    };
1081                    let client = self.sessions.get(&session).map(|io| io.client.clone());
1082                    let result = if let Some(client) = &client {
1083                        client
1084                            .register(operation.clone(), start.sender)
1085                            .map(|()| operation)
1086                            .map_err(|error| {
1087                                unb_core::CoreError::UnknownStream(error.to_string())
1088                            })
1089                    } else {
1090                        Err(unb_core::CoreError::UnknownSession(session.to_string()))
1091                    };
1092                    if let Err(Ok(operation)) = start.reply.send(result) {
1093                        if let Some(client) = client {
1094                            client.abandon(&operation);
1095                        }
1096                        let _ = self.core.handle(
1097                            clock_now(),
1098                            CoreInput::CancelClientOperation { session, operation },
1099                        );
1100                    }
1101                }
1102                CoreEffect::Deliver {
1103                    session,
1104                    mut envelope,
1105                } => {
1106                    if let Some(token) = envelope.body_token.take() {
1107                        match self.handle.claim_body(&session, &token) {
1108                            Some(WireBody::Bytes(payload)) => {
1109                                envelope.payload = payload;
1110                            }
1111                            Some(body @ WireBody::Stream(_)) => {
1112                                let ceiling = self
1113                                    .sessions
1114                                    .get(&session)
1115                                    .map_or(unb_transport::DEFAULT_MAX_FRAME_SIZE, |io| {
1116                                        io.ws_collect_ceiling
1117                                    });
1118                                match body.collect_to(ceiling).await {
1119                                    Ok(payload) => envelope.payload = payload,
1120                                    Err(error) => {
1121                                        envelope.kind = unb_core::Kind::Error;
1122                                        envelope.payload =
1123                                            Envelope::encode_payload(&serde_json::json!({
1124                                                "code": unb_core::ErrorCode::PayloadTooLarge,
1125                                                "message": error.to_string(),
1126                                            }));
1127                                    }
1128                                }
1129                            }
1130                            None => {}
1131                        }
1132                    }
1133                    self.handler(&session, HandlerInput::Deliver(envelope))
1134                        .await
1135                }
1136                CoreEffect::ReleaseBody { session, body } => {
1137                    let _ = self.handle.claim_body(&session, body.as_str());
1138                }
1139                CoreEffect::StreamClosed { session, operation } => {
1140                    self.handler(&session, HandlerInput::StreamClosed(operation))
1141                        .await
1142                }
1143                CoreEffect::CloseTransport {
1144                    session,
1145                    code,
1146                    message,
1147                } => {
1148                    let _ = (code, message);
1149                    if let Some(io) = self.sessions.get(&session) {
1150                        io.cancellation.cancel();
1151                    }
1152                }
1153                CoreEffect::ScheduleSessionDeadline { session, deadline } => {
1154                    if let Some(io) = self.sessions.get(&session) {
1155                        io.deadline.send_replace(Some(deadline));
1156                    }
1157                }
1158                CoreEffect::RouteExportAcked { session } => {
1159                    if let Some(io) = self.sessions.get(&session) {
1160                        io.routes.send_replace(true);
1161                    }
1162                }
1163                CoreEffect::ScheduleClientDeadline {
1164                    session,
1165                    operation,
1166                    deadline,
1167                } => spawn_client_timeout(
1168                    session.clone(),
1169                    operation,
1170                    deadline,
1171                    self.handle.clone(),
1172                    self.sessions
1173                        .get(&session)
1174                        .map(|io| io.cancellation.clone()),
1175                ),
1176                CoreEffect::Send {
1177                    effect,
1178                    session,
1179                    frame,
1180                } => {
1181                    if frame.head.kind == unb_core::Kind::Error {
1182                        let _ = self.core.handle(
1183                            clock_now(),
1184                            CoreInput::SendCompleted {
1185                                effect,
1186                                result: SendResult::Written,
1187                            },
1188                        );
1189                        self.write(&session, frame.into_envelope(), None, false)
1190                            .await;
1191                    } else {
1192                        self.write(&session, frame.into_envelope(), Some(effect), false)
1193                            .await;
1194                    }
1195                }
1196                CoreEffect::SendProtocol {
1197                    effect,
1198                    session,
1199                    envelope,
1200                } => {
1201                    self.write(&session, envelope, Some(effect), false).await;
1202                }
1203                CoreEffect::ContinueDiscovery { stream } => {
1204                    let _ = self
1205                        .core
1206                        .handle(clock_now(), CoreInput::ContinueDiscovery { stream });
1207                }
1208                CoreEffect::SessionRetired { session, reason } => {
1209                    #[cfg(feature = "observability")]
1210                    {
1211                        metrics::counter!("unb_sessions_retired", "reason" => format!("{reason:?}"))
1212                            .increment(1);
1213                        tracing::debug!(session = %session, reason = ?reason, "session retired");
1214                    }
1215                    if let Some(outcome) = self.outcomes.remove(&session) {
1216                        let _ = outcome.send(Some(SessionOutcome::Retired(reason)));
1217                    }
1218                    if let Some(io) = self.sessions.remove(&session) {
1219                        io.cancellation.cancel();
1220                    }
1221                    self.pending_bodies
1222                        .lock()
1223                        .expect("pending body registry")
1224                        .retain(|(owner, _), _| owner != &session);
1225                    self.handle
1226                        .token_bodies
1227                        .lock()
1228                        .expect("body token registry")
1229                        .retain(|(owner, _), _| owner != &session);
1230                    self.execute(CoreEffect::SessionRetired { session, reason });
1231                }
1232                CoreEffect::SessionEstablished { session, peer } => {
1233                    let effect = CoreEffect::SessionEstablished {
1234                        session: session.clone(),
1235                        peer,
1236                    };
1237                    if let Some(input) = self.executor.execute(effect, self.handle.clone()).await {
1238                        let _ = self.core.handle(clock_now(), input);
1239                    }
1240                    if let Some(outcome) = self.outcomes.get(&session) {
1241                        let _ = outcome.send(Some(SessionOutcome::Established));
1242                    }
1243                }
1244                CoreEffect::ForwardRelay {
1245                    effect,
1246                    source,
1247                    target,
1248                    mut frame,
1249                    ..
1250                } => {
1251                    if let Some(body_id) = frame.body.take() {
1252                        let Some(body) = self.handle.claim_body(&source.session, body_id.as_str())
1253                        else {
1254                            let _ = self.core.handle(
1255                                clock_now(),
1256                                CoreInput::RelayForwardCompleted {
1257                                    effect,
1258                                    result: SendResult::Refused {
1259                                        code: unb_core::ErrorCode::Protocol,
1260                                        message: "relay body capability is unavailable".into(),
1261                                    },
1262                                },
1263                            );
1264                            continue;
1265                        };
1266                        match self.handle.register_body(&target.session, body) {
1267                            Ok(target_body) => frame.body = Some(target_body),
1268                            Err(error) => {
1269                                let _ = self.core.handle(
1270                                    clock_now(),
1271                                    CoreInput::RelayForwardCompleted {
1272                                        effect,
1273                                        result: SendResult::Refused {
1274                                            code: unb_core::ErrorCode::Busy,
1275                                            message: error.to_string(),
1276                                        },
1277                                    },
1278                                );
1279                                continue;
1280                            }
1281                        }
1282                    }
1283                    self.write(&target.session, frame.into_envelope(), Some(effect), true)
1284                        .await;
1285                }
1286                effect @ (CoreEffect::RouteSnapshotApplied { .. }
1287                | CoreEffect::RouteDeltaApplied { .. }
1288                | CoreEffect::RouteSessionWithdrawn { .. }) => {
1289                    if let Some(input) = self.executor.execute(effect, self.handle.clone()).await {
1290                        let _ = self.core.handle(clock_now(), input);
1291                    }
1292                }
1293                effect => self.execute(effect),
1294            }
1295        }
1296    }
1297
1298    fn execute(&self, effect: CoreEffect) {
1299        let executor = self.executor.clone();
1300        let handle = self.handle.clone();
1301        n0_future::task::spawn(async move {
1302            if let Some(input) = executor.execute(effect, handle.clone()).await {
1303                let _ = handle.submit(input).await;
1304            }
1305        });
1306    }
1307
1308    async fn write(
1309        &mut self,
1310        session: &SessionId,
1311        mut envelope: Envelope,
1312        completion: Option<EffectId>,
1313        relay: bool,
1314    ) {
1315        if is_operation_envelope(&envelope) {
1316            let mode = self.sessions.get(session).and_then(|io| io.transport_mode);
1317            match mode {
1318                None if self
1319                    .sessions
1320                    .get(session)
1321                    .is_some_and(|io| io.bodies.is_some()) =>
1322                {
1323                    let queued = self.sessions.get_mut(session).is_some_and(|io| {
1324                        if io.pending_mode_writes.len() >= PENDING_MODE_WRITES_MAX {
1325                            false
1326                        } else {
1327                            io.pending_mode_writes.push_back(PendingModeWrite {
1328                                envelope,
1329                                completion,
1330                                relay,
1331                            });
1332                            true
1333                        }
1334                    });
1335                    if !queued {
1336                        if let Some(effect) = completion {
1337                            self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1338                        }
1339                    }
1340                    return;
1341                }
1342                Some(unb_transport::TransportMode::OperationStreamsV1) => {
1343                    self.write_operation(session, envelope, completion, relay)
1344                        .await;
1345                    return;
1346                }
1347                _ => {}
1348            }
1349        }
1350        if envelope.kind == unb_core::Kind::Cancel {
1351            if let Some(corr) = &envelope.corr {
1352                self.handle.cancel_body_pump(session, corr);
1353            }
1354        }
1355        let body = match (&envelope.body_token, &envelope.corr) {
1356            (Some(token), _) => self.handle.claim_body(session, token),
1357            (None, Some(corr))
1358                if envelope.kind.is_application_request()
1359                    || envelope.kind == unb_core::Kind::Response =>
1360            {
1361                let pending = self
1362                    .pending_bodies
1363                    .lock()
1364                    .expect("pending body registry")
1365                    .remove(&(session.clone(), corr.clone()));
1366                if let Some(stream) = pending {
1367                    envelope.body_token = Some(mint_body_token());
1368                    Some(WireBody::Stream(stream))
1369                } else {
1370                    None
1371                }
1372            }
1373            _ => None,
1374        };
1375        let Some(io) = self.sessions.get(session) else {
1376            if let Some(effect) = completion {
1377                self.complete_send(effect, relay, SendResult::Closed);
1378            }
1379            return;
1380        };
1381        if let Some(WireBody::Bytes(payload)) = body {
1382            envelope.body_token = None;
1383            envelope.payload = payload;
1384        } else if let Some(WireBody::Stream(stream)) = body {
1385            if let Some(peer) = io.bodies.clone().filter(|bodies| bodies.can_open()) {
1386                let token = envelope.body_token.clone().expect("tokenized body");
1387                let pump_cancellation = io.cancellation.child_token();
1388                if let Some(corr) = &envelope.corr {
1389                    self.handle.register_body_pump(
1390                        session,
1391                        corr,
1392                        &token,
1393                        pump_cancellation.clone(),
1394                    );
1395                }
1396                spawn_body_pump(
1397                    peer,
1398                    token,
1399                    stream,
1400                    io.byte_budget.clone(),
1401                    pump_cancellation,
1402                    self.handle.clone(),
1403                    envelope.corr.clone().map(|corr| (session.clone(), corr)),
1404                );
1405            } else {
1406                envelope.body_token = None;
1407                spawn_body_collect(
1408                    self.handle.clone(),
1409                    session.clone(),
1410                    envelope,
1411                    stream,
1412                    io.ws_collect_ceiling,
1413                    completion,
1414                    relay,
1415                    io.cancellation.clone(),
1416                );
1417                return;
1418            }
1419        }
1420        let envelope = match io
1421            .byte_budget
1422            .clone()
1423            .try_acquire_many_owned(WriteRequest::weight(&envelope))
1424        {
1425            Ok(budget) => {
1426                let request = WriteRequest {
1427                    envelope,
1428                    completion,
1429                    relay,
1430                    _budget: budget,
1431                };
1432                match io.outbound.try_send(request) {
1433                    Ok(()) => {
1434                        if let Some(effect) = completion {
1435                            self.complete_send(effect, relay, SendResult::Reserved);
1436                        }
1437                        return;
1438                    }
1439                    Err(mpsc::error::TrySendError::Closed(_)) => {
1440                        if let Some(effect) = completion {
1441                            self.complete_send(effect, relay, SendResult::Closed);
1442                        }
1443                        return;
1444                    }
1445                    Err(mpsc::error::TrySendError::Full(request)) => {
1446                        #[cfg(feature = "observability")]
1447                        {
1448                            metrics::counter!("unb_write_queue_full").increment(1);
1449                            tracing::warn!(session = %session, "outbound write queue is full");
1450                        }
1451                        request.envelope
1452                    }
1453                }
1454            }
1455            Err(_) => {
1456                #[cfg(feature = "observability")]
1457                {
1458                    metrics::counter!("unb_write_budget_exhausted").increment(1);
1459                    tracing::warn!(session = %session, "outbound byte budget is exhausted");
1460                }
1461                envelope
1462            }
1463        };
1464        if completion.is_none()
1465            || io
1466                .reservation_pending
1467                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
1468                .is_ok()
1469        {
1470            let timeout = if matches!(
1471                envelope.kind,
1472                unb_core::Kind::Response
1473                    | unb_core::Kind::Error
1474                    | unb_core::Kind::Cancel
1475            ) {
1476                TERMINAL_SEND_TIMEOUT
1477            } else {
1478                RELAY_BACKPRESSURE_TIMEOUT
1479            };
1480            reserve_write(
1481                io.outbound.clone(),
1482                io.byte_budget.clone(),
1483                envelope,
1484                completion,
1485                relay,
1486                self.handle.clone(),
1487                io.cancellation.clone(),
1488                completion.map(|_| io.reservation_pending.clone()),
1489                timeout,
1490            );
1491        } else if let Some(effect) = completion {
1492            self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1493        }
1494    }
1495
1496    fn start_operation_acceptor(&self, session: &SessionId) {
1497        let Some(io) = self.sessions.get(session) else {
1498            return;
1499        };
1500        let Some(bodies) = io.bodies.clone() else {
1501            return;
1502        };
1503        spawn_operation_acceptor(
1504            session.clone(),
1505            bodies,
1506            io.operation_lanes.clone(),
1507            io.operation_slots.clone(),
1508            io.byte_budget.clone(),
1509            self.handle.clone(),
1510            io.cancellation.clone(),
1511        );
1512    }
1513
1514    async fn write_operation(
1515        &mut self,
1516        session: &SessionId,
1517        mut envelope: Envelope,
1518        completion: Option<EffectId>,
1519        relay: bool,
1520    ) {
1521        let Some(corr) = envelope.corr.clone() else {
1522            if let Some(effect) = completion {
1523                self.complete_send(
1524                    effect,
1525                    relay,
1526                    SendResult::WriteFailed("missing unb-corr".into()),
1527                );
1528            }
1529            return;
1530        };
1531        let body = match &envelope.body_token {
1532            Some(token) => self.handle.claim_body(session, token),
1533            None if envelope.kind.is_application_request()
1534                || envelope.kind == unb_core::Kind::Response =>
1535            {
1536                self.pending_bodies
1537                    .lock()
1538                    .expect("pending body registry")
1539                    .remove(&(session.clone(), corr.clone()))
1540                    .map(WireBody::Stream)
1541            }
1542            None => None,
1543        };
1544        let body = match body {
1545            Some(WireBody::Bytes(payload)) => {
1546                envelope.payload = payload;
1547                None
1548            }
1549            Some(WireBody::Stream(stream)) => Some(stream),
1550            None => None,
1551        };
1552        envelope.body_token = None;
1553        let Some(io) = self.sessions.get(session) else {
1554            if let Some(effect) = completion {
1555                self.complete_send(effect, relay, SendResult::Closed);
1556            }
1557            return;
1558        };
1559        if envelope.kind == unb_core::Kind::Cancel {
1560            let lane = io
1561                .operation_lanes
1562                .lock()
1563                .expect("operation lane registry")
1564                .remove(&corr);
1565            if let Some(control) = lane.and_then(|lane| lane.control) {
1566                let _ = control.reset(unb_transport::operation::RESET_CANCELLED);
1567                let _ = control.stop(unb_transport::operation::RESET_CANCELLED);
1568            }
1569            if let Some(effect) = completion {
1570                self.complete_send(effect, relay, SendResult::Written);
1571            }
1572            return;
1573        }
1574        if envelope.kind.is_application_request() {
1575            spawn_outbound_operation(
1576                session.clone(),
1577                corr,
1578                envelope,
1579                body,
1580                completion,
1581                relay,
1582                io.bodies
1583                    .clone()
1584                    .expect("operation mode has stream facility"),
1585                io.operation_lanes.clone(),
1586                io.operation_slots.clone(),
1587                io.byte_budget.clone(),
1588                self.handle.clone(),
1589                io.cancellation.clone(),
1590            );
1591            return;
1592        }
1593        let lane = io
1594            .operation_lanes
1595            .lock()
1596            .expect("operation lane registry")
1597            .get(&corr)
1598            .cloned();
1599        let Some(sender) = lane.and_then(|lane| lane.outbound) else {
1600            if let Some(effect) = completion {
1601                self.complete_send(effect, relay, SendResult::Written);
1602            }
1603            return;
1604        };
1605        let budget = match io
1606            .byte_budget
1607            .clone()
1608            .try_acquire_many_owned(WriteRequest::weight(&envelope))
1609        {
1610            Ok(budget) => budget,
1611            Err(_) => {
1612                if let Some(effect) = completion {
1613                    self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1614                }
1615                return;
1616            }
1617        };
1618        let request = OperationWriteRequest {
1619            envelope,
1620            completion,
1621            relay,
1622            body,
1623            _budget: budget,
1624        };
1625        match sender.try_send(request) {
1626            Ok(()) => {
1627                if let Some(effect) = completion {
1628                    self.complete_send(effect, relay, SendResult::Reserved);
1629                }
1630            }
1631            Err(_) => {
1632                if let Some(effect) = completion {
1633                    self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1634                }
1635            }
1636        }
1637    }
1638
1639    fn complete_send(&mut self, effect: EffectId, relay: bool, result: SendResult) {
1640        let _ = self
1641            .core
1642            .handle(clock_now(), send_completion(effect, relay, result));
1643    }
1644
1645    async fn handler(&self, session: &SessionId, event: HandlerInput) {
1646        if let Some(io) = self.sessions.get(session) {
1647            let _ = io.handler.send(event).await;
1648        }
1649    }
1650}
1651
1652enum HandlerInput {
1653    Deliver(Envelope),
1654    StreamClosed(ClientOperationId),
1655}
1656
1657fn handler_channel<H: SessionHandler>(mut handler: H) -> mpsc::Sender<HandlerInput> {
1658    let (events, mut receiver) = mpsc::channel(SESSION_QUEUE);
1659    n0_future::task::spawn(async move {
1660        while let Some(event) = receiver.recv().await {
1661            match event {
1662                HandlerInput::Deliver(envelope) => handler.deliver(envelope).await,
1663                HandlerInput::StreamClosed(operation) => handler.stream_closed(operation).await,
1664            }
1665        }
1666    });
1667    events
1668}
1669
1670fn spawn_directives(
1671    session: SessionId,
1672    mut directives: mpsc::Receiver<Directive>,
1673    handle: ProtocolCoreHandle,
1674    cancellation: CancellationToken,
1675) {
1676    n0_future::task::spawn(async move {
1677        loop {
1678            let directive = tokio::select! {
1679                biased;
1680                () = cancellation.cancelled() => return,
1681                directive = directives.recv() => directive,
1682            };
1683            let Some(directive) = directive else { return };
1684            if handle
1685                .inputs
1686                .send(ActorInput::Directive(session.clone(), directive))
1687                .await
1688                .is_err()
1689            {
1690                return;
1691            }
1692        }
1693    });
1694}
1695
1696fn is_operation_envelope(envelope: &Envelope) -> bool {
1697    envelope.corr.is_some()
1698        && (envelope.kind.is_application_request()
1699            || envelope.kind.is_application_response()
1700            || envelope.kind == unb_core::Kind::Cancel)
1701}
1702
1703fn spawn_operation_acceptor(
1704    session: SessionId,
1705    bodies: Arc<unb_transport::ConnectionStreams>,
1706    lanes: OperationLanes,
1707    operation_slots: Arc<Semaphore>,
1708    byte_budget: Arc<Semaphore>,
1709    handle: ProtocolCoreHandle,
1710    cancellation: CancellationToken,
1711) {
1712    n0_future::task::spawn(async move {
1713        loop {
1714            let operation = tokio::select! {
1715                biased;
1716                () = cancellation.cancelled() => return,
1717                operation = bodies.accept_operation() => operation,
1718            };
1719            let Ok(operation) = operation else {
1720                return;
1721            };
1722            let Ok(operation_slot) = operation_slots.clone().try_acquire_owned() else {
1723                if let Some(control) = operation.control {
1724                    let _ = control.reset(unb_transport::operation::RESET_BUSY);
1725                    let _ = control.stop(unb_transport::operation::RESET_BUSY);
1726                }
1727                continue;
1728            };
1729            let corr = operation.corr.clone();
1730            let (outbound, receiver) = mpsc::channel(OPERATION_WRITE_QUEUE);
1731            let lane = OperationLane {
1732                outbound: Some(outbound),
1733                control: operation.control.clone(),
1734            };
1735            let admitted = {
1736                let mut lanes = lanes.lock().expect("operation lane registry");
1737                if lanes.contains_key(&corr) {
1738                    false
1739                } else {
1740                    lanes.insert(corr.clone(), lane);
1741                    true
1742                }
1743            };
1744            if !admitted {
1745                if let Some(control) = operation.control {
1746                    let _ = control.reset(unb_transport::operation::RESET_PROTOCOL);
1747                }
1748                continue;
1749            }
1750            spawn_accepted_operation(
1751                session.clone(),
1752                operation,
1753                receiver,
1754                lanes.clone(),
1755                byte_budget.clone(),
1756                handle.clone(),
1757                cancellation.child_token(),
1758                operation_slot,
1759            );
1760        }
1761    });
1762}
1763
1764#[allow(clippy::too_many_arguments)]
1765fn spawn_outbound_operation(
1766    session: SessionId,
1767    corr: String,
1768    envelope: Envelope,
1769    body: Option<BodyStream>,
1770    completion: Option<EffectId>,
1771    relay: bool,
1772    bodies: Arc<unb_transport::ConnectionStreams>,
1773    lanes: OperationLanes,
1774    operation_slots: Arc<Semaphore>,
1775    byte_budget: Arc<Semaphore>,
1776    handle: ProtocolCoreHandle,
1777    cancellation: CancellationToken,
1778) {
1779    n0_future::task::spawn(async move {
1780        let open_byte_budget = byte_budget.clone();
1781        let open = async {
1782            let operation_slot = operation_slots
1783                .acquire_owned()
1784                .await
1785                .map_err(|_| unb_transport::TransportError::Closed)?;
1786            let budget = open_byte_budget
1787                .acquire_many_owned(WriteRequest::weight(&envelope))
1788                .await
1789                .map_err(|_| unb_transport::TransportError::Closed)?;
1790            let operation = bodies
1791                .open_operation(&corr)
1792                .await?
1793                .ok_or(unb_transport::TransportError::Closed)?;
1794            Ok::<_, unb_transport::TransportError>((operation, budget, operation_slot))
1795        };
1796        let opened = tokio::select! {
1797            biased;
1798            () = cancellation.cancelled() => return,
1799            opened = n0_future::time::timeout(RELAY_BACKPRESSURE_TIMEOUT, open) => opened,
1800        };
1801        let Ok(Ok((operation, _budget, _operation_slot))) = opened else {
1802            if let Some(effect) = completion {
1803                let _ = handle
1804                    .submit(send_completion(
1805                        effect,
1806                        relay,
1807                        SendResult::ReservationTimedOut,
1808                    ))
1809                    .await;
1810            }
1811            return;
1812        };
1813        lanes.lock().expect("operation lane registry").insert(
1814            corr.clone(),
1815            OperationLane {
1816                outbound: None,
1817                control: operation.control.clone(),
1818            },
1819        );
1820        if let Some(effect) = completion {
1821            if handle
1822                .submit(send_completion(effect, relay, SendResult::Reserved))
1823                .await
1824                .is_err()
1825            {
1826                return;
1827            }
1828        }
1829        let mut writer = unb_transport::operation::OperationRecordWriter::new(operation.write);
1830        let sent = write_operation_envelope(
1831            &mut writer,
1832            envelope,
1833            body,
1834            byte_budget.clone(),
1835            cancellation.clone(),
1836        )
1837        .await;
1838        let sent = match sent {
1839            Ok(()) => writer.finish().await,
1840            Err(error) => Err(error),
1841        };
1842        if let Some(effect) = completion {
1843            let _ = handle
1844                .submit(send_completion(effect, relay, SendResult::Written))
1845                .await;
1846        }
1847        if let Err(error) = sent {
1848            if !operation_peer_stopped_cancelled(&error) {
1849                if let Some(control) = &operation.control {
1850                    let _ = control.reset(error.reset_code());
1851                }
1852                let _ = submit_operation_end(
1853                    &handle,
1854                    &session,
1855                    &corr,
1856                    OperationDirection::Opening,
1857                    operation_outcome(&error),
1858                )
1859                .await;
1860                lanes.lock().expect("operation lane registry").remove(&corr);
1861                return;
1862            }
1863        }
1864        let result = read_operation_direction(
1865            &session,
1866            &corr,
1867            operation.read,
1868            OperationDirection::Return,
1869            &handle,
1870            &cancellation,
1871        )
1872        .await;
1873        if let Err(error) = result {
1874            if let Some(control) = &operation.control {
1875                let _ = control.reset(error.reset_code());
1876                let _ = control.stop(error.reset_code());
1877            }
1878            let _ = submit_operation_end(
1879                &handle,
1880                &session,
1881                &corr,
1882                OperationDirection::Return,
1883                operation_outcome(&error),
1884            )
1885            .await;
1886        }
1887        lanes.lock().expect("operation lane registry").remove(&corr);
1888    });
1889}
1890
1891#[allow(clippy::too_many_arguments)]
1892fn spawn_accepted_operation(
1893    session: SessionId,
1894    operation: unb_transport::operation::OperationIo,
1895    receiver: mpsc::Receiver<OperationWriteRequest>,
1896    lanes: OperationLanes,
1897    byte_budget: Arc<Semaphore>,
1898    handle: ProtocolCoreHandle,
1899    cancellation: CancellationToken,
1900    operation_slot: OwnedSemaphorePermit,
1901) {
1902    let corr = operation.corr.clone();
1903    let control = operation.control.clone();
1904    let read_corr = corr.clone();
1905    let read_lanes = lanes.clone();
1906    let read_handle = handle.clone();
1907    let read_cancellation = cancellation.clone();
1908    let writer_session = session.clone();
1909    n0_future::task::spawn(async move {
1910        let result = read_operation_direction(
1911            &session,
1912            &read_corr,
1913            operation.read,
1914            OperationDirection::Opening,
1915            &read_handle,
1916            &read_cancellation,
1917        )
1918        .await;
1919        if let Err(error) = result {
1920            if operation_body_consumer_closed(&error) {
1921                if let Some(control) = control {
1922                    let _ = control.stop(unb_transport::operation::RESET_CANCELLED);
1923                }
1924                return;
1925            }
1926            if let Some(control) = control {
1927                let _ = control.reset(error.reset_code());
1928                let _ = control.stop(error.reset_code());
1929            }
1930            let _ = submit_operation_end(
1931                &read_handle,
1932                &session,
1933                &read_corr,
1934                OperationDirection::Opening,
1935                operation_outcome(&error),
1936            )
1937            .await;
1938            read_lanes
1939                .lock()
1940                .expect("operation lane registry")
1941                .remove(&read_corr);
1942        }
1943    });
1944    spawn_operation_writer(
1945        writer_session,
1946        corr,
1947        operation.write,
1948        operation.control,
1949        receiver,
1950        lanes,
1951        byte_budget,
1952        handle,
1953        cancellation,
1954        operation_slot,
1955    );
1956}
1957
1958#[allow(clippy::too_many_arguments)]
1959fn spawn_operation_writer(
1960    session: SessionId,
1961    corr: String,
1962    write: Box<dyn tokio::io::AsyncWrite + Send + Unpin>,
1963    control: Option<unb_transport::operation::OperationControl>,
1964    mut receiver: mpsc::Receiver<OperationWriteRequest>,
1965    lanes: OperationLanes,
1966    byte_budget: Arc<Semaphore>,
1967    handle: ProtocolCoreHandle,
1968    cancellation: CancellationToken,
1969    _operation_slot: OwnedSemaphorePermit,
1970) {
1971    n0_future::task::spawn(async move {
1972        let mut writer = unb_transport::operation::OperationRecordWriter::new(write);
1973        loop {
1974            let request = tokio::select! {
1975                biased;
1976                () = cancellation.cancelled() => break,
1977                request = receiver.recv() => request,
1978            };
1979            let Some(request) = request else {
1980                break;
1981            };
1982            let terminal = matches!(
1983                request.envelope.kind,
1984                unb_core::Kind::Response | unb_core::Kind::Error
1985            );
1986            let result = write_operation_envelope(
1987                &mut writer,
1988                request.envelope,
1989                request.body,
1990                byte_budget.clone(),
1991                cancellation.clone(),
1992            )
1993            .await;
1994            let result = if result.is_ok() && terminal {
1995                writer.finish().await
1996            } else {
1997                result
1998            };
1999            if let Some(effect) = request.completion {
2000                let _ = handle
2001                    .submit(send_completion(effect, request.relay, SendResult::Written))
2002                    .await;
2003            }
2004            if result.is_err() || terminal {
2005                if let Err(error) = result {
2006                    if let Some(control) = &control {
2007                        let _ = control.reset(error.reset_code());
2008                    }
2009                    let _ = submit_operation_end(
2010                        &handle,
2011                        &session,
2012                        &corr,
2013                        OperationDirection::Return,
2014                        operation_outcome(&error),
2015                    )
2016                    .await;
2017                }
2018                break;
2019            }
2020        }
2021        lanes.lock().expect("operation lane registry").remove(&corr);
2022    });
2023}
2024
2025async fn write_operation_envelope<W>(
2026    writer: &mut unb_transport::operation::OperationRecordWriter<W>,
2027    mut envelope: Envelope,
2028    body: Option<BodyStream>,
2029    byte_budget: Arc<Semaphore>,
2030    cancellation: CancellationToken,
2031) -> Result<(), unb_transport::operation::OperationCodecError>
2032where
2033    W: tokio::io::AsyncWrite + Unpin,
2034{
2035    let Some(body) = body else {
2036        return writer
2037            .write_record(
2038                unb_transport::operation::OperationRecordKind::CompleteMessage,
2039                &envelope.encode(),
2040            )
2041            .await;
2042    };
2043    envelope.body_token = None;
2044    envelope.payload = Bytes::new();
2045    writer
2046        .write_record(
2047            unb_transport::operation::OperationRecordKind::Head,
2048            &envelope.encode(),
2049        )
2050        .await?;
2051    let mut body = MeteredBody::new(body, byte_budget, cancellation.clone());
2052    loop {
2053        let chunk = body.next_chunk().await.map_err(|error| {
2054            unb_transport::operation::OperationCodecError::Io(error.to_string())
2055        })?;
2056        let Some(chunk) = chunk else {
2057            if cancellation.is_cancelled() {
2058                return Err(unb_transport::operation::OperationCodecError::Io(
2059                    "operation body cancelled".into(),
2060                ));
2061            }
2062            break;
2063        };
2064        for piece in chunk.bytes.chunks(OPERATION_BODY_CHUNK_BYTES) {
2065            writer
2066                .write_record(
2067                    unb_transport::operation::OperationRecordKind::BodyChunk,
2068                    piece,
2069                )
2070                .await?;
2071        }
2072    }
2073    writer
2074        .write_record(
2075            unb_transport::operation::OperationRecordKind::BodyEnd,
2076            &[],
2077        )
2078        .await
2079}
2080
2081#[derive(Clone, Copy)]
2082enum OperationDirection {
2083    Opening,
2084    Return,
2085}
2086
2087async fn read_operation_direction(
2088    session: &SessionId,
2089    corr: &str,
2090    read: Box<dyn tokio::io::AsyncRead + Send + Unpin>,
2091    direction: OperationDirection,
2092    handle: &ProtocolCoreHandle,
2093    cancellation: &CancellationToken,
2094) -> Result<(), unb_transport::operation::OperationCodecError> {
2095    let mut reader = unb_transport::operation::OperationRecordReader::new(read);
2096    let mut records = 0usize;
2097    let mut terminal = false;
2098    let mut body_sender: Option<mpsc::Sender<Result<Bytes, unb_core::CoreError>>> = None;
2099    loop {
2100        let kind = tokio::select! {
2101            biased;
2102            () = cancellation.cancelled() => {
2103                if let Some(sender) = body_sender.take() {
2104                    let _ = sender.send(Err(unb_core::CoreError::Malformed(
2105                        "operation body cancelled".into(),
2106                    ))).await;
2107                }
2108                return Ok(())
2109            },
2110            record = reader.read_record() => record,
2111        };
2112        let (kind, payload) = match kind {
2113            Ok(record) => record,
2114            Err(error) => {
2115                if let Some(sender) = body_sender.take() {
2116                    let _ = sender
2117                        .send(Err(unb_core::CoreError::Malformed(error.to_string())))
2118                        .await;
2119                }
2120                return Err(error);
2121            }
2122        };
2123        match kind {
2124            unb_transport::operation::OperationRecordKind::CompleteMessage => {
2125                if body_sender.is_some() {
2126                    return Err(unb_transport::operation::OperationCodecError::Protocol(
2127                        "complete message arrived before body end".into(),
2128                    ));
2129                }
2130                let envelope = decode_operation_message(payload, corr, direction, records)?;
2131                terminal |= matches!(
2132                    envelope.kind,
2133                    unb_core::Kind::Response | unb_core::Kind::Error
2134                );
2135                records += 1;
2136                let input = application_core_input(handle, session, envelope).map_err(|error| {
2137                    unb_transport::operation::OperationCodecError::Protocol(error.to_string())
2138                })?;
2139                handle.submit(input).await.map_err(|error| {
2140                    unb_transport::operation::OperationCodecError::Io(error.to_string())
2141                })?;
2142            }
2143            unb_transport::operation::OperationRecordKind::Head => {
2144                if body_sender.is_some() {
2145                    return Err(unb_transport::operation::OperationCodecError::Protocol(
2146                        "operation head arrived before body end".into(),
2147                    ));
2148                }
2149                let mut envelope = decode_operation_message(payload, corr, direction, records)?;
2150                let (token, sender) = register_operation_body(handle, session)?;
2151                envelope.body_token = Some(token);
2152                terminal |= matches!(
2153                    envelope.kind,
2154                    unb_core::Kind::Response | unb_core::Kind::Error
2155                );
2156                records += 1;
2157                body_sender = Some(sender);
2158                let input = application_core_input(handle, session, envelope).map_err(|error| {
2159                    unb_transport::operation::OperationCodecError::Protocol(error.to_string())
2160                })?;
2161                handle.submit(input).await.map_err(|error| {
2162                    unb_transport::operation::OperationCodecError::Io(error.to_string())
2163                })?;
2164            }
2165            unb_transport::operation::OperationRecordKind::BodyChunk => {
2166                let Some(sender) = &body_sender else {
2167                    return Err(unb_transport::operation::OperationCodecError::Protocol(
2168                        "body chunk arrived without a head".into(),
2169                    ));
2170                };
2171                sender.send(Ok(payload)).await.map_err(|_| {
2172                    unb_transport::operation::OperationCodecError::Io(
2173                        "operation body consumer closed".into(),
2174                    )
2175                })?;
2176            }
2177            unb_transport::operation::OperationRecordKind::BodyEnd => {
2178                let Some(sender) = body_sender.take() else {
2179                    return Err(unb_transport::operation::OperationCodecError::Protocol(
2180                        "body end arrived without a head".into(),
2181                    ));
2182                };
2183                drop(sender);
2184            }
2185            unb_transport::operation::OperationRecordKind::DirectionEnd => {
2186                if body_sender.is_some() {
2187                    return Err(
2188                        unb_transport::operation::OperationCodecError::Truncated(
2189                            "direction ended before body end".into(),
2190                        ),
2191                    );
2192                }
2193                if records == 0 || matches!(direction, OperationDirection::Return) && !terminal {
2194                    return Err(unb_transport::operation::OperationCodecError::Protocol(
2195                        "operation direction ended before its required message".into(),
2196                    ));
2197                }
2198                reader.expect_fin().await?;
2199                submit_operation_end(
2200                    handle,
2201                    session,
2202                    corr,
2203                    direction,
2204                    unb_core::OperationStreamOutcome::Clean,
2205                )
2206                .await
2207                .map_err(|error| {
2208                    unb_transport::operation::OperationCodecError::Io(error.to_string())
2209                })?;
2210                return Ok(());
2211            }
2212        }
2213    }
2214}
2215
2216fn register_operation_body(
2217    handle: &ProtocolCoreHandle,
2218    session: &SessionId,
2219) -> Result<
2220    (String, mpsc::Sender<Result<Bytes, unb_core::CoreError>>),
2221    unb_transport::operation::OperationCodecError,
2222> {
2223    let token = mint_body_token();
2224    let (sender, receiver) = mpsc::channel(8);
2225    let stream: BodyStream = Box::pin(futures_util::stream::unfold(
2226        receiver,
2227        |mut receiver| async move { receiver.recv().await.map(|chunk| (chunk, receiver)) },
2228    ));
2229    let mut registry = handle.token_bodies.lock().expect("body token registry");
2230    let owned = registry
2231        .keys()
2232        .filter(|(owner, _)| owner == session)
2233        .count();
2234    if owned >= SESSION_BODY_TOKENS_MAX {
2235        return Err(unb_transport::operation::OperationCodecError::Protocol(
2236            "session operation body limit reached".into(),
2237        ));
2238    }
2239    registry.insert((session.clone(), token.clone()), WireBody::Stream(stream));
2240    Ok((token, sender))
2241}
2242
2243fn operation_outcome(
2244    error: &unb_transport::operation::OperationCodecError,
2245) -> unb_core::OperationStreamOutcome {
2246    match error {
2247        unb_transport::operation::OperationCodecError::Protocol(message) => {
2248            unb_core::OperationStreamOutcome::Protocol(message.clone())
2249        }
2250        unb_transport::operation::OperationCodecError::Truncated(_) => {
2251            unb_core::OperationStreamOutcome::Truncated
2252        }
2253        unb_transport::operation::OperationCodecError::Oversized { size, max } => {
2254            unb_core::OperationStreamOutcome::Protocol(format!(
2255                "operation record of {size} bytes exceeds max {max}"
2256            ))
2257        }
2258        unb_transport::operation::OperationCodecError::Io(message) => {
2259            unb_core::OperationStreamOutcome::Transport(message.clone())
2260        }
2261    }
2262}
2263
2264fn operation_body_consumer_closed(
2265    error: &unb_transport::operation::OperationCodecError,
2266) -> bool {
2267    matches!(
2268        error,
2269        unb_transport::operation::OperationCodecError::Io(message)
2270            if message == "operation body consumer closed"
2271    )
2272}
2273
2274fn operation_peer_stopped_cancelled(
2275    error: &unb_transport::operation::OperationCodecError,
2276) -> bool {
2277    matches!(
2278        error,
2279        unb_transport::operation::OperationCodecError::Io(message)
2280            if message.contains(&unb_transport::operation::RESET_CANCELLED.to_string())
2281    )
2282}
2283
2284async fn submit_operation_end(
2285    handle: &ProtocolCoreHandle,
2286    session: &SessionId,
2287    corr: &str,
2288    direction: OperationDirection,
2289    outcome: unb_core::OperationStreamOutcome,
2290) -> Result<(), WsError> {
2291    handle
2292        .submit(CoreInput::OperationStreamEnded {
2293            session: session.clone(),
2294            corr: unb_core::CorrelationId::from(corr.to_owned()),
2295            direction: match direction {
2296                OperationDirection::Opening => unb_core::OperationStreamDirection::Opening,
2297                OperationDirection::Return => unb_core::OperationStreamDirection::Return,
2298            },
2299            outcome,
2300        })
2301        .await
2302}
2303
2304fn decode_operation_message(
2305    payload: Bytes,
2306    corr: &str,
2307    direction: OperationDirection,
2308    records: usize,
2309) -> Result<Envelope, unb_transport::operation::OperationCodecError> {
2310    let envelope = Envelope::decode(payload).map_err(|error| {
2311        unb_transport::operation::OperationCodecError::Protocol(error.to_string())
2312    })?;
2313    if envelope.corr.as_deref() != Some(corr) {
2314        return Err(unb_transport::operation::OperationCodecError::Protocol(
2315            "operation record changed unb-corr".into(),
2316        ));
2317    }
2318    match direction {
2319        OperationDirection::Opening if records != 0 || !envelope.kind.is_application_request() => {
2320            Err(unb_transport::operation::OperationCodecError::Protocol(
2321                "opening direction requires exactly one application opener".into(),
2322            ))
2323        }
2324        OperationDirection::Return if !envelope.kind.is_application_response() => {
2325            Err(unb_transport::operation::OperationCodecError::Protocol(
2326                "return direction requires application responses".into(),
2327            ))
2328        }
2329        _ => Ok(envelope),
2330    }
2331}
2332
2333fn spawn_body_pump(
2334    peer: std::sync::Arc<unb_transport::ConnectionStreams>,
2335    token: String,
2336    stream: BodyStream,
2337    budget: Arc<Semaphore>,
2338    cancellation: CancellationToken,
2339    handle: ProtocolCoreHandle,
2340    registration: Option<(SessionId, String)>,
2341) {
2342    n0_future::task::spawn(async move {
2343        let pump = async {
2344            let Ok(mut writer) = peer.open_body(&token).await else {
2345                return;
2346            };
2347            let mut metered = MeteredBody::new(stream, budget, cancellation.clone());
2348            loop {
2349                match metered.next_chunk().await {
2350                    Ok(Some(chunk)) => {
2351                        if writer.write_chunk(&chunk.bytes).await.is_err() {
2352                            return;
2353                        }
2354                    }
2355                    Ok(None) => {
2356                        if cancellation.is_cancelled() {
2357                            writer.abort().await;
2358                        } else {
2359                            let _ = writer.finish().await;
2360                        }
2361                        return;
2362                    }
2363                    Err(_) => {
2364                        writer.abort().await;
2365                        return;
2366                    }
2367                }
2368            }
2369        };
2370        pump.await;
2371        if let Some((session, corr)) = registration {
2372            handle.release_body_pump(&session, &corr, &token);
2373        }
2374    });
2375}
2376
2377#[allow(clippy::too_many_arguments)]
2378fn spawn_body_collect(
2379    handle: ProtocolCoreHandle,
2380    session: SessionId,
2381    mut envelope: Envelope,
2382    stream: BodyStream,
2383    ceiling: usize,
2384    completion: Option<EffectId>,
2385    relay: bool,
2386    cancellation: CancellationToken,
2387) {
2388    n0_future::task::spawn(async move {
2389        let collected = tokio::select! {
2390            biased;
2391            () = cancellation.cancelled() => {
2392                if let Some(effect) = completion {
2393                    let _ = handle
2394                        .submit(send_completion(effect, relay, SendResult::Cancelled))
2395                        .await;
2396                }
2397                return;
2398            }
2399            collected = WireBody::Stream(stream).collect_to(ceiling) => collected,
2400        };
2401        match collected {
2402            Ok(payload) => {
2403                envelope.payload = payload;
2404                let _ = handle
2405                    .inputs
2406                    .send(ActorInput::Write {
2407                        session,
2408                        envelope,
2409                        completion,
2410                        relay,
2411                    })
2412                    .await;
2413            }
2414            Err(error) => {
2415                let (code, message) = match error {
2416                    unb_core::CoreError::BodyTooLarge(_) => (
2417                        unb_core::ErrorCode::PayloadTooLarge,
2418                        "streaming body exceeds this route's WS collect ceiling; raise the \
2419                         ceiling or reach this node over WebTransport"
2420                            .to_string(),
2421                    ),
2422                    other => (unb_core::ErrorCode::Protocol, other.to_string()),
2423                };
2424                if let Some(effect) = completion {
2425                    let refused = SendResult::Refused { code, message };
2426                    let _ = handle.submit(send_completion(effect, relay, refused)).await;
2427                } else if envelope.kind.is_application_request() {
2428                    if let Some(corr) = envelope.corr.clone() {
2429                        let outcome = if code == unb_core::ErrorCode::PayloadTooLarge {
2430                            unb_core::OperationStreamOutcome::PayloadTooLarge(message)
2431                        } else {
2432                            unb_core::OperationStreamOutcome::Protocol(message)
2433                        };
2434                        let _ = handle
2435                            .submit(CoreInput::OperationStreamEnded {
2436                                session,
2437                                corr: unb_core::CorrelationId::from(corr),
2438                                direction: unb_core::OperationStreamDirection::Opening,
2439                                outcome,
2440                            })
2441                            .await;
2442                    }
2443                } else if envelope.kind == unb_core::Kind::Response {
2444                    envelope.kind = unb_core::Kind::Error;
2445                    envelope.payload = Envelope::encode_payload(&serde_json::json!({
2446                        "code": code,
2447                        "message": message,
2448                    }));
2449                    let _ = handle
2450                        .inputs
2451                        .send(ActorInput::Write {
2452                            session,
2453                            envelope,
2454                            completion: None,
2455                            relay: false,
2456                        })
2457                        .await;
2458                }
2459            }
2460        }
2461    });
2462}
2463
2464fn spawn_reader(
2465    session: SessionId,
2466    mut reader: PipeReader,
2467    handle: ProtocolCoreHandle,
2468    cancellation: CancellationToken,
2469    bodies: crate::transport::SessionStreams,
2470) {
2471    n0_future::task::spawn(async move {
2472        loop {
2473            let received = tokio::select! {
2474                biased;
2475                () = cancellation.cancelled() => {
2476                    let _ = handle.submit(CoreInput::SessionClosed { session }).await;
2477                    return;
2478                },
2479                received = reader.recv() => received,
2480            };
2481            let envelope = match received {
2482                Ok(Some((envelope, _body))) => {
2483                    if let Some(token) = &envelope.body_token {
2484                        if !(envelope.kind.is_application_request()
2485                            || envelope.kind == unb_core::Kind::Response)
2486                        {
2487                            let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2488                            return;
2489                        }
2490                        {
2491                            let Some(peer) = bodies.clone() else {
2492                                let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2493                                return;
2494                            };
2495                            let lazy_token = token.clone();
2496                            let accepted = futures_util::stream::once(async move {
2497                                peer.accept_body(&lazy_token, RELAY_BACKPRESSURE_TIMEOUT)
2498                                    .await
2499                            });
2500                            let stream: BodyStream =
2501                                Box::pin(futures_util::StreamExt::flat_map(accepted, |outcome| {
2502                                    match outcome {
2503                                        Ok(reader) => futures_util::StreamExt::boxed(
2504                                            futures_util::StreamExt::map(reader, |chunk| {
2505                                                chunk.map_err(|error| {
2506                                                    unb_core::CoreError::Malformed(
2507                                                        error.to_string(),
2508                                                    )
2509                                                })
2510                                            }),
2511                                        ),
2512                                        Err(error) => futures_util::StreamExt::boxed(
2513                                            futures_util::stream::once(async move {
2514                                                Err(unb_core::CoreError::Malformed(
2515                                                    error.to_string(),
2516                                                ))
2517                                            }),
2518                                        ),
2519                                    }
2520                                }));
2521                            let rejected = {
2522                                let mut registry =
2523                                    handle.token_bodies.lock().expect("body token registry");
2524                                let key = (session.clone(), token.clone());
2525                                let owned = registry
2526                                    .keys()
2527                                    .filter(|(owner, _)| owner == &session)
2528                                    .count();
2529                                if registry.contains_key(&key) || owned >= SESSION_BODY_TOKENS_MAX {
2530                                    true
2531                                } else {
2532                                    registry.insert(key, WireBody::Stream(stream));
2533                                    false
2534                                }
2535                            };
2536                            if rejected {
2537                                let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2538                                return;
2539                            }
2540                        }
2541                    }
2542                    envelope
2543                }
2544                Ok(None) => {
2545                    let _ = handle.submit(CoreInput::SessionClosed { session }).await;
2546                    return;
2547                }
2548                Err(_) => {
2549                    #[cfg(feature = "observability")]
2550                    {
2551                        metrics::counter!("unb_transport_failures").increment(1);
2552                        tracing::warn!(session = %session, "transport read failed");
2553                    }
2554                    let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2555                    return;
2556                }
2557            };
2558            if envelope.kind == unb_core::Kind::Cancel {
2559                if let Some(corr) = &envelope.corr {
2560                    handle.cancel_body_pump(&session, corr);
2561                }
2562            }
2563            let input = match application_core_input(&handle, &session, envelope) {
2564                Ok(input) => input,
2565                Err(_) => {
2566                    let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2567                    return;
2568                }
2569            };
2570            if handle.submit(input).await.is_err() {
2571                return;
2572            }
2573        }
2574    });
2575}
2576
2577fn spawn_writer(
2578    session: SessionId,
2579    mut writer: PipeWriter,
2580    mut outbound: mpsc::Receiver<WriteRequest>,
2581    handle: ProtocolCoreHandle,
2582    cancellation: CancellationToken,
2583) {
2584    n0_future::task::spawn(async move {
2585        loop {
2586            let request = tokio::select! {
2587                biased;
2588                () = cancellation.cancelled() => {
2589                    let farewell = async {
2590                        while let Ok(request) = outbound.try_recv() {
2591                            if writer.feed(request.envelope).await.is_err() {
2592                                break;
2593                            }
2594                        }
2595                        let _ = writer.flush().await;
2596                        let _ = writer.close().await;
2597                    };
2598                    let _ = n0_future::time::timeout(TERMINAL_SEND_TIMEOUT, farewell).await;
2599                    return;
2600                },
2601                request = outbound.recv() => request,
2602            };
2603            let Some(request) = request else {
2604                let _ = handle
2605                    .submit(CoreInput::SessionClosed {
2606                        session: session.clone(),
2607                    })
2608                    .await;
2609                return;
2610            };
2611            let mut batch = Vec::with_capacity(WRITE_BATCH);
2612            batch.push(request);
2613            while batch.len() < WRITE_BATCH {
2614                match outbound.try_recv() {
2615                    Ok(request) => batch.push(request),
2616                    Err(_) => break,
2617                }
2618            }
2619            let completions = batch
2620                .iter()
2621                .filter_map(|request| request.completion.map(|effect| (effect, request.relay)))
2622                .collect::<Vec<_>>();
2623            let result = async {
2624                for request in batch {
2625                    writer.feed(request.envelope).await?;
2626                }
2627                writer.flush().await
2628            }
2629            .await;
2630            let send_result = result.as_ref().map_or_else(
2631                |error| SendResult::WriteFailed(error.to_string()),
2632                |_| SendResult::Written,
2633            );
2634            let completions = completions
2635                .into_iter()
2636                .map(|(effect, relay)| send_completion(effect, relay, send_result.clone()))
2637                .collect::<Vec<_>>();
2638            if !completions.is_empty()
2639                && handle
2640                    .inputs
2641                    .send(ActorInput::CoreBatch(completions))
2642                    .await
2643                    .is_err()
2644            {
2645                return;
2646            }
2647            if result.is_err() {
2648                let _ = handle
2649                    .submit(CoreInput::SessionClosed {
2650                        session: session.clone(),
2651                    })
2652                    .await;
2653                cancellation.cancel();
2654                let _ = writer.close().await;
2655                return;
2656            }
2657        }
2658    });
2659}
2660
2661#[allow(clippy::too_many_arguments)]
2662fn reserve_write(
2663    sender: mpsc::Sender<WriteRequest>,
2664    byte_budget: Arc<Semaphore>,
2665    envelope: Envelope,
2666    completion: Option<EffectId>,
2667    relay: bool,
2668    handle: ProtocolCoreHandle,
2669    cancellation: CancellationToken,
2670    reservation_pending: Option<Arc<AtomicBool>>,
2671    timeout: std::time::Duration,
2672) {
2673    n0_future::task::spawn(async move {
2674        let effect = completion;
2675        let acquire = async {
2676            let budget = byte_budget
2677                .acquire_many_owned(WriteRequest::weight(&envelope))
2678                .await
2679                .map_err(|_| ())?;
2680            let slot = sender.reserve_owned().await.map_err(|_| ())?;
2681            Ok::<_, ()>((budget, slot))
2682        };
2683        let reserved = tokio::select! {
2684            biased;
2685            () = cancellation.cancelled() => {
2686                if let Some(pending) = &reservation_pending {
2687                    pending.store(false, Ordering::Release);
2688                }
2689                if let Some(effect) = effect {
2690                    let _ = handle.submit(send_completion(effect, relay, SendResult::Cancelled)).await;
2691                }
2692                return;
2693            }
2694            reserved = n0_future::time::timeout(timeout, acquire) => reserved,
2695        };
2696        if let Some(pending) = &reservation_pending {
2697            pending.store(false, Ordering::Release);
2698        }
2699        match reserved {
2700            Ok(Ok((budget, slot))) => {
2701                if let Some(effect) = effect {
2702                    if handle
2703                        .submit(send_completion(effect, relay, SendResult::Reserved))
2704                        .await
2705                        .is_err()
2706                    {
2707                        return;
2708                    }
2709                }
2710                slot.send(WriteRequest {
2711                    envelope,
2712                    completion,
2713                    relay,
2714                    _budget: budget,
2715                });
2716            }
2717            Ok(Err(())) => {
2718                if let Some(effect) = effect {
2719                    let _ = handle
2720                        .submit(send_completion(effect, relay, SendResult::Closed))
2721                        .await;
2722                }
2723            }
2724            Err(_) => {
2725                if let Some(effect) = effect {
2726                    let _ = handle
2727                        .submit(send_completion(
2728                            effect,
2729                            relay,
2730                            SendResult::ReservationTimedOut,
2731                        ))
2732                        .await;
2733                }
2734            }
2735        }
2736    });
2737}
2738
2739fn send_completion(effect: EffectId, relay: bool, result: SendResult) -> CoreInput {
2740    if relay {
2741        CoreInput::RelayForwardCompleted { effect, result }
2742    } else {
2743        CoreInput::SendCompleted { effect, result }
2744    }
2745}
2746
2747fn spawn_session_deadline(
2748    session: SessionId,
2749    mut deadlines: tokio::sync::watch::Receiver<Option<Instant>>,
2750    handle: ProtocolCoreHandle,
2751    cancellation: CancellationToken,
2752) {
2753    n0_future::task::spawn(async move {
2754        loop {
2755            let deadline = *deadlines.borrow_and_update();
2756            match deadline {
2757                Some(deadline) if clock_now() < deadline => {
2758                    tokio::select! {
2759                        biased;
2760                        () = cancellation.cancelled() => return,
2761                        () = sleep_to(deadline) => {}
2762                    }
2763                }
2764                Some(_) => {
2765                    if handle
2766                        .submit(CoreInput::SessionTimeout {
2767                            session: session.clone(),
2768                        })
2769                        .await
2770                        .is_err()
2771                    {
2772                        return;
2773                    }
2774                    tokio::select! {
2775                        biased;
2776                        () = cancellation.cancelled() => return,
2777                        changed = deadlines.changed() => {
2778                            if changed.is_err() {
2779                                return;
2780                            }
2781                        }
2782                    }
2783                }
2784                None => {
2785                    tokio::select! {
2786                        biased;
2787                        () = cancellation.cancelled() => return,
2788                        changed = deadlines.changed() => {
2789                            if changed.is_err() {
2790                                return;
2791                            }
2792                        }
2793                    }
2794                }
2795            }
2796        }
2797    });
2798}
2799
2800fn spawn_client_timeout(
2801    session: SessionId,
2802    operation: ClientOperationId,
2803    deadline: Instant,
2804    handle: ProtocolCoreHandle,
2805    cancellation: Option<CancellationToken>,
2806) {
2807    let Some(cancellation) = cancellation else {
2808        return;
2809    };
2810    n0_future::task::spawn(async move {
2811        tokio::select! {
2812            biased;
2813            () = cancellation.cancelled() => {}
2814            () = sleep_to(deadline) => {
2815                let _ = handle.submit(CoreInput::ClientOperationTimeout { session, operation }).await;
2816            }
2817        }
2818    });
2819}
2820
2821fn spawn_establishment_timeout(session: SessionId, handle: ProtocolCoreHandle) {
2822    n0_future::task::spawn(async move {
2823        n0_future::time::sleep(ESTABLISHMENT_TIMEOUT).await;
2824        let _ = handle
2825            .submit(CoreInput::EstablishmentTimeout { session })
2826            .await;
2827    });
2828}
2829
2830#[cfg(test)]
2831mod tests {
2832    use super::*;
2833    use unb_core::{DiscoverPlan, Kind, RetirementReason, PROTOCOL_VERSION};
2834    use bytes::Bytes;
2835    use std::sync::atomic::{AtomicBool, Ordering};
2836
2837    struct RecordingExecutor {
2838        effects: mpsc::Sender<CoreEffect>,
2839    }
2840
2841    impl EffectExecutor for RecordingExecutor {
2842        fn execute(&self, effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
2843            let effects = self.effects.clone();
2844            Box::pin(async move {
2845                let _ = effects.send(effect).await;
2846                None
2847            })
2848        }
2849    }
2850
2851    struct PublicationExecutor {
2852        published: Arc<AtomicBool>,
2853        effects: mpsc::Sender<bool>,
2854    }
2855
2856    struct ClosingHandler {
2857        delivered: mpsc::Sender<Envelope>,
2858        closed: mpsc::Sender<ClientOperationId>,
2859    }
2860
2861    impl SessionHandler for ClosingHandler {
2862        async fn deliver(&mut self, envelope: Envelope) {
2863            let _ = self.delivered.send(envelope).await;
2864        }
2865
2866        async fn stream_closed(&mut self, operation: ClientOperationId) {
2867            let _ = self.closed.send(operation).await;
2868        }
2869    }
2870
2871    impl EffectExecutor for PublicationExecutor {
2872        fn execute(&self, effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
2873            let published = self.published.clone();
2874            let effects = self.effects.clone();
2875            Box::pin(async move {
2876                if matches!(effect, CoreEffect::SessionRetired { .. }) {
2877                    let _ = effects.send(published.load(Ordering::SeqCst)).await;
2878                }
2879                None
2880            })
2881        }
2882    }
2883
2884    fn local_pipe() -> (Pipe, mpsc::Sender<Envelope>, mpsc::Receiver<Envelope>) {
2885        let (incoming, rx) = mpsc::channel(16);
2886        let (tx, outgoing) = mpsc::channel(16);
2887        (
2888            Pipe::Local {
2889                rx,
2890                tx,
2891                initiator: true,
2892            },
2893            incoming,
2894            outgoing,
2895        )
2896    }
2897
2898    fn frame(kind: Kind, corr: Option<String>, payload: serde_json::Value) -> Envelope {
2899        Envelope {
2900            v: PROTOCOL_VERSION,
2901            id: format!("{kind:?}"),
2902            subject: String::new(),
2903            kind,
2904            corr,
2905            seq: None,
2906            hops: None,
2907            body_token: None,
2908            payload: Envelope::encode_payload(&payload),
2909            path: Vec::new(),
2910            headers: Default::default(),
2911        }
2912    }
2913
2914    async fn establish(incoming: &mpsc::Sender<Envelope>, outgoing: &mut mpsc::Receiver<Envelope>) {
2915        assert_eq!(outgoing.recv().await.unwrap().kind, Kind::Hello);
2916        incoming
2917            .send(frame(
2918                Kind::Welcome,
2919                None,
2920                serde_json::json!({ "version": 1 }),
2921            ))
2922            .await
2923            .unwrap();
2924    }
2925
2926    fn request(subject: &str, body: &'static [u8]) -> http::Request<Bytes> {
2927        http::Request::post(format!("/{subject}"))
2928            .body(Bytes::from_static(body))
2929            .unwrap()
2930    }
2931
2932    async fn attach_client(handle: &ProtocolCoreHandle, name: &str, pipe: Pipe) -> Arc<Wire> {
2933        attach_client_with_ceiling(
2934            handle,
2935            name,
2936            pipe,
2937            unb_transport::DEFAULT_MAX_FRAME_SIZE,
2938        )
2939        .await
2940    }
2941
2942    async fn attach_client_with_ceiling(
2943        handle: &ProtocolCoreHandle,
2944        name: &str,
2945        pipe: Pipe,
2946        ws_collect_ceiling: usize,
2947    ) -> Arc<Wire> {
2948        let session_cancellation = handle.cancellation.child_token();
2949        let (wire, directives, outcome, routes) =
2950            Wire::standalone_channel(session_cancellation.clone());
2951        let (reply, response) = oneshot::channel();
2952        handle
2953            .inputs
2954            .send(ActorInput::Attach {
2955                session: SessionId::from(name),
2956                transport: pipe,
2957                expected_peer: None,
2958                directives,
2959                handler: handler_channel(wire.observation_sender()),
2960                client: wire.client_session(),
2961                cancellation: session_cancellation,
2962                establish_peer: false,
2963                outcome,
2964                routes,
2965                ws_collect_ceiling,
2966                reply,
2967            })
2968            .await
2969            .unwrap();
2970        response.await.unwrap().unwrap();
2971        wire
2972    }
2973
2974    fn spawn_core() -> ProtocolCoreHandle {
2975        ProtocolCoreHandle::spawn(
2976            ProtocolCore::new("node"),
2977            Arc::new(RecordingExecutor {
2978                effects: mpsc::channel(1).0,
2979            }),
2980            CancellationToken::new(),
2981            &tokio::runtime::Handle::current(),
2982        )
2983    }
2984
2985    fn stalled_pipe() -> (Pipe, mpsc::Sender<Envelope>, mpsc::Receiver<Envelope>) {
2986        let (incoming, rx) = mpsc::channel(16);
2987        let (tx, outgoing) = mpsc::channel(1);
2988        (
2989            Pipe::Local {
2990                rx,
2991                tx,
2992                initiator: true,
2993            },
2994            incoming,
2995            outgoing,
2996        )
2997    }
2998
2999    fn chunk_stream(chunks: Vec<Result<Bytes, unb_core::CoreError>>) -> BodyStream {
3000        Box::pin(futures_util::stream::iter(chunks))
3001    }
3002
3003    #[tokio::test(start_paused = true)]
3004    async fn a_metered_body_stall_past_the_deadline_is_busy_shaped() {
3005        let budget = Arc::new(Semaphore::new(4));
3006        let hold = budget.clone().try_acquire_many_owned(4).unwrap();
3007        let mut body = MeteredBody::new(
3008            chunk_stream(vec![Ok(Bytes::from_static(b"x"))]),
3009            budget,
3010            CancellationToken::new(),
3011        );
3012        let outcome = body.next_chunk().await;
3013        assert!(
3014            matches!(
3015                outcome,
3016                Err(WsError::Core(unb_core::CoreError::BodyStalled))
3017            ),
3018            "budget starvation past the deadline must be BUSY-shaped: {outcome:?}"
3019        );
3020        drop(hold);
3021    }
3022
3023    #[tokio::test]
3024    async fn a_window_saturated_body_leaves_a_sibling_pump_undisturbed() {
3025        let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
3026        let mut saturated = MeteredBody::new(
3027            chunk_stream(vec![
3028                Ok(Bytes::from(vec![0u8; BODY_WINDOW_BYTES])),
3029                Ok(Bytes::from_static(b"tail")),
3030            ]),
3031            budget.clone(),
3032            CancellationToken::new(),
3033        );
3034        let held = saturated.next_chunk().await.unwrap().unwrap();
3035        assert!(
3036            tokio::time::timeout(std::time::Duration::from_millis(50), saturated.next_chunk())
3037                .await
3038                .is_err(),
3039            "a full per-body window must block the next chunk"
3040        );
3041        let mut sibling = MeteredBody::new(
3042            chunk_stream(vec![Ok(Bytes::from_static(b"y"))]),
3043            budget,
3044            CancellationToken::new(),
3045        );
3046        let chunk = tokio::time::timeout(std::time::Duration::from_secs(1), sibling.next_chunk())
3047            .await
3048            .expect("a sibling body must progress while another window is saturated")
3049            .unwrap()
3050            .unwrap();
3051        assert_eq!(chunk.bytes, Bytes::from_static(b"y"));
3052        drop(held);
3053        let tail = tokio::time::timeout(std::time::Duration::from_secs(1), saturated.next_chunk())
3054            .await
3055            .expect("draining the window must unblock the body")
3056            .unwrap()
3057            .unwrap();
3058        assert_eq!(tail.bytes, Bytes::from_static(b"tail"));
3059    }
3060
3061    #[test]
3062    fn body_tokens_are_unique_and_carry_entropy() {
3063        let first = mint_body_token();
3064        let second = mint_body_token();
3065        assert_ne!(first, second);
3066        assert!(
3067            first.len() >= 33,
3068            "a guessable serial token invites body substitution: {first}"
3069        );
3070    }
3071
3072    #[tokio::test]
3073    async fn an_oversized_source_chunk_is_rejected_before_retention() {
3074        let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
3075        let mut body = MeteredBody::new(
3076            chunk_stream(vec![Ok(Bytes::from(vec![
3077                0u8;
3078                unb_transport::DEFAULT_MAX_FRAME_SIZE
3079                    + 1
3080            ]))]),
3081            budget.clone(),
3082            CancellationToken::new(),
3083        );
3084        let outcome = body.next_chunk().await;
3085        assert!(
3086            matches!(
3087                outcome,
3088                Err(WsError::Core(unb_core::CoreError::BodyTooLarge(_)))
3089            ),
3090            "a chunk beyond the frame cap must be refused before retention: {outcome:?}"
3091        );
3092        assert_eq!(
3093            budget.available_permits(),
3094            OUTBOUND_BYTE_BUDGET,
3095            "a refused chunk must not consume budget"
3096        );
3097    }
3098
3099    #[tokio::test]
3100    async fn concurrent_bodies_stay_within_the_session_budget() {
3101        let budget = Arc::new(Semaphore::new(100));
3102        let mut first = MeteredBody::new(
3103            chunk_stream(vec![Ok(Bytes::from(vec![0u8; 60]))]),
3104            budget.clone(),
3105            CancellationToken::new(),
3106        );
3107        let mut second = MeteredBody::new(
3108            chunk_stream(vec![Ok(Bytes::from(vec![0u8; 60]))]),
3109            budget.clone(),
3110            CancellationToken::new(),
3111        );
3112        let held = first.next_chunk().await.unwrap().unwrap();
3113        assert_eq!(budget.available_permits(), 40);
3114        assert!(
3115            tokio::time::timeout(std::time::Duration::from_millis(50), second.next_chunk())
3116                .await
3117                .is_err(),
3118            "combined queued bytes must not exceed the session budget"
3119        );
3120        drop(held);
3121        let chunk = tokio::time::timeout(std::time::Duration::from_secs(1), second.next_chunk())
3122            .await
3123            .expect("released budget must admit the waiting body")
3124            .unwrap()
3125            .unwrap();
3126        assert_eq!(chunk.bytes.len(), 60);
3127    }
3128
3129    #[tokio::test(start_paused = true)]
3130    async fn a_steady_body_streams_past_the_terminal_deadline_without_timing_out() {
3131        let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
3132        let chunks: Vec<Result<Bytes, unb_core::CoreError>> =
3133            (0..8).map(|_| Ok(Bytes::from_static(b"chunk"))).collect();
3134        let mut body = MeteredBody::new(chunk_stream(chunks), budget, CancellationToken::new());
3135        for _ in 0..8 {
3136            let chunk = body
3137                .next_chunk()
3138                .await
3139                .expect("a steadily progressing body never times out")
3140                .expect("chunks remain");
3141            drop(chunk);
3142            tokio::time::advance(std::time::Duration::from_secs(6)).await;
3143        }
3144        assert!(body.next_chunk().await.unwrap().is_none());
3145    }
3146
3147    #[tokio::test]
3148    async fn pipe_recv_yields_whole_message_envelopes_with_no_body_stream() {
3149        let mut envelope = frame(
3150            Kind::Request,
3151            Some("c1".into()),
3152            serde_json::json!({"n": 1}),
3153        );
3154        envelope.subject = "echo".into();
3155        let encoded = envelope.encode();
3156        let pipe = Pipe::Piped {
3157            pipe: Box::pin(ScriptedPipe {
3158                frames: std::collections::VecDeque::from([Ok(encoded.clone())]),
3159            }),
3160            initiator: true,
3161        };
3162        let (mut reader, _writer) = pipe.split();
3163        let (received, body) = reader.recv().await.unwrap().unwrap();
3164        assert!(body.is_none(), "whole-message carriers collect the body");
3165        assert_eq!(received, Envelope::decode(encoded).unwrap());
3166
3167        let (incoming, rx) = mpsc::channel(4);
3168        let (tx, _outgoing) = mpsc::channel(4);
3169        let pipe = Pipe::Local {
3170            rx,
3171            tx,
3172            initiator: true,
3173        };
3174        let (mut reader, _writer) = pipe.split();
3175        incoming.send(envelope.clone()).await.unwrap();
3176        let (received, body) = reader.recv().await.unwrap().unwrap();
3177        assert!(body.is_none());
3178        assert_eq!(received, envelope);
3179    }
3180
3181    #[test]
3182    fn write_weight_counts_head_and_payload_and_clamps_to_the_whole_budget() {
3183        let mut envelope = frame(Kind::Request, None, serde_json::Value::Null);
3184        envelope.id = String::new();
3185        envelope.payload = Bytes::new();
3186        assert_eq!(WriteRequest::weight(&envelope), 1);
3187        envelope
3188            .headers
3189            .insert("origin".into(), serde_json::Value::String("x".repeat(64)));
3190        assert_eq!(WriteRequest::weight(&envelope), 70);
3191        envelope.payload = Bytes::from(vec![0u8; OUTBOUND_BYTE_BUDGET + 1]);
3192        assert_eq!(WriteRequest::weight(&envelope), OUTBOUND_BYTE_BUDGET as u32);
3193    }
3194
3195    #[tokio::test]
3196    async fn a_byte_saturated_outbound_ends_with_busy_and_recovers_on_drain() {
3197        let handle = spawn_core();
3198        let (pipe, incoming, mut outgoing) = stalled_pipe();
3199        let wire = attach_client(&handle, "stalled", pipe).await;
3200        establish(&incoming, &mut outgoing).await;
3201        let client = wire.client_session();
3202        let payload = Bytes::from(vec![0u8; 16 * 1024 * 1024]);
3203
3204        let mut flood = Vec::new();
3205        for _ in 0..6 {
3206            flood.push(
3207                client
3208                    .start(
3209                        "flood",
3210                        Kind::Request,
3211                        payload.clone(),
3212                        None,
3213                        Default::default(),
3214                    )
3215                    .await
3216                    .unwrap(),
3217            );
3218        }
3219        let mut rejected = client
3220            .start(
3221                "flood",
3222                Kind::Request,
3223                payload.clone(),
3224                None,
3225                Default::default(),
3226            )
3227            .await
3228            .unwrap();
3229
3230        let refused = rejected.operation().as_str().to_owned();
3231        let refusal = tokio::time::timeout(std::time::Duration::from_secs(2), rejected.next())
3232            .await
3233            .expect("an over-budget write must be refused promptly, not silently queued");
3234        assert!(
3235            refusal.is_err(),
3236            "queued outbound bytes past the budget must refuse the write: {refusal:?}"
3237        );
3238
3239        let mut saw_busy = false;
3240        loop {
3241            match tokio::time::timeout(std::time::Duration::from_millis(300), outgoing.recv()).await
3242            {
3243                Ok(Some(envelope)) => {
3244                    if envelope.kind == Kind::Error
3245                        && envelope.corr.as_deref() == Some(refused.as_str())
3246                    {
3247                        assert_eq!(envelope.payload_json()["code"], "BUSY");
3248                        saw_busy = true;
3249                    }
3250                }
3251                Ok(None) => panic!("the stalled pipe closed while draining"),
3252                Err(_) => break,
3253            }
3254        }
3255        assert!(saw_busy, "the refused write must surface as a BUSY error");
3256
3257        let probe = client
3258            .start(
3259                "probe",
3260                Kind::Request,
3261                payload.clone(),
3262                None,
3263                Default::default(),
3264            )
3265            .await
3266            .unwrap();
3267        let opening = tokio::time::timeout(std::time::Duration::from_secs(5), async {
3268            loop {
3269                let envelope = outgoing.recv().await.unwrap();
3270                if envelope.subject == "probe" {
3271                    break envelope;
3272                }
3273            }
3274        })
3275        .await
3276        .expect("draining the writer must release budget for later writes");
3277        incoming
3278            .send(frame(
3279                Kind::Response,
3280                opening.corr,
3281                serde_json::json!({ "ok": true }),
3282            ))
3283            .await
3284            .unwrap();
3285        let mut probe = probe;
3286        let response = tokio::time::timeout(std::time::Duration::from_secs(5), probe.next())
3287            .await
3288            .expect("a post-drain write must complete end to end")
3289            .unwrap()
3290            .unwrap();
3291        assert_eq!(response.payload_json(), serde_json::json!({ "ok": true }));
3292        handle.shutdown();
3293    }
3294
3295    #[tokio::test]
3296    async fn a_byte_saturated_session_leaves_a_sibling_session_undisturbed() {
3297        let handle = spawn_core();
3298        let (saturated_pipe, stalled_incoming, mut stalled_outgoing) = stalled_pipe();
3299        let stalled = attach_client(&handle, "stalled", saturated_pipe).await;
3300        establish(&stalled_incoming, &mut stalled_outgoing).await;
3301        let (live_pipe, live_incoming, mut live_outgoing) = local_pipe();
3302        let live = attach_client(&handle, "live", live_pipe).await;
3303        establish(&live_incoming, &mut live_outgoing).await;
3304
3305        let payload = Bytes::from(vec![0u8; 16 * 1024 * 1024]);
3306        let mut flood = Vec::new();
3307        for _ in 0..6 {
3308            flood.push(
3309                stalled
3310                    .client_session()
3311                    .start(
3312                        "flood",
3313                        Kind::Request,
3314                        payload.clone(),
3315                        None,
3316                        Default::default(),
3317                    )
3318                    .await
3319                    .unwrap(),
3320            );
3321        }
3322
3323        let mut sibling = live
3324            .client_session()
3325            .start(
3326                "sibling",
3327                Kind::Request,
3328                Bytes::new(),
3329                None,
3330                Default::default(),
3331            )
3332            .await
3333            .unwrap();
3334        let opening = tokio::time::timeout(std::time::Duration::from_secs(1), live_outgoing.recv())
3335            .await
3336            .expect("a sibling session must keep writing while another is byte-saturated")
3337            .unwrap();
3338        assert_eq!(opening.subject, "sibling");
3339        live_incoming
3340            .send(frame(
3341                Kind::Response,
3342                opening.corr,
3343                serde_json::json!({ "ok": true }),
3344            ))
3345            .await
3346            .unwrap();
3347        let response = tokio::time::timeout(std::time::Duration::from_secs(2), sibling.next())
3348            .await
3349            .expect("a sibling response must be delivered while another session is saturated")
3350            .unwrap()
3351            .unwrap();
3352        assert_eq!(response.payload_json(), serde_json::json!({ "ok": true }));
3353        handle.shutdown();
3354    }
3355
3356    #[tokio::test]
3357    async fn terminal_client_delivery_closes_handler_stream_once() {
3358        let cancellation = CancellationToken::new();
3359        let handle = ProtocolCoreHandle::spawn(
3360            ProtocolCore::new("node"),
3361            Arc::new(RecordingExecutor {
3362                effects: mpsc::channel(1).0,
3363            }),
3364            cancellation,
3365            &tokio::runtime::Handle::current(),
3366        );
3367        let (pipe, incoming, mut outgoing) = local_pipe();
3368        let (delivered, _deliveries) = mpsc::channel(4);
3369        let (closed, mut closures) = mpsc::channel(4);
3370        let session_cancellation = handle.cancellation.child_token();
3371        let (wire, directives, outcome, routes) =
3372            Wire::standalone_channel(session_cancellation.clone());
3373        let (reply, response) = oneshot::channel();
3374        handle
3375            .inputs
3376            .send(ActorInput::Attach {
3377                session: SessionId::from("session"),
3378                transport: pipe,
3379                expected_peer: None,
3380                directives,
3381                handler: handler_channel(ClosingHandler { delivered, closed }),
3382                client: wire.client_session(),
3383                cancellation: session_cancellation,
3384                establish_peer: false,
3385                outcome,
3386                routes,
3387                ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
3388                reply,
3389            })
3390            .await
3391            .unwrap();
3392        response.await.unwrap().unwrap();
3393        establish(&incoming, &mut outgoing).await;
3394        let mut stream = wire
3395            .client_session()
3396            .start(
3397                "echo",
3398                Kind::Request,
3399                Bytes::new(),
3400                None,
3401                Default::default(),
3402            )
3403            .await
3404            .unwrap();
3405        let opening = outgoing.recv().await.unwrap();
3406
3407        incoming
3408            .send(frame(
3409                Kind::Response,
3410                opening.corr,
3411                serde_json::json!({ "ok": true }),
3412            ))
3413            .await
3414            .unwrap();
3415
3416        assert!(stream.next().await.unwrap().is_some());
3417        assert_eq!(closures.recv().await, Some(stream.operation().clone()));
3418        assert!(
3419            tokio::time::timeout(std::time::Duration::from_millis(25), closures.recv())
3420                .await
3421                .is_err()
3422        );
3423        handle.shutdown();
3424    }
3425
3426    #[tokio::test]
3427    async fn generic_terminal_delivery_closes_handler_stream_once() {
3428        let cancellation = CancellationToken::new();
3429        let handle = ProtocolCoreHandle::spawn(
3430            ProtocolCore::new("node"),
3431            Arc::new(RecordingExecutor {
3432                effects: mpsc::channel(1).0,
3433            }),
3434            cancellation,
3435            &tokio::runtime::Handle::current(),
3436        );
3437        let (pipe, incoming, mut outgoing) = local_pipe();
3438        let (delivered, mut deliveries) = mpsc::channel(4);
3439        let (closed, mut closures) = mpsc::channel(4);
3440        let session_cancellation = handle.cancellation.child_token();
3441        let (wire, directives, outcome, routes) =
3442            Wire::standalone_channel(session_cancellation.clone());
3443        let (reply, response) = oneshot::channel();
3444        handle
3445            .inputs
3446            .send(ActorInput::Attach {
3447                session: SessionId::from("session"),
3448                transport: pipe,
3449                expected_peer: None,
3450                directives,
3451                handler: handler_channel(ClosingHandler { delivered, closed }),
3452                client: wire.client_session(),
3453                cancellation: session_cancellation,
3454                establish_peer: false,
3455                outcome,
3456                routes,
3457                ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
3458                reply,
3459            })
3460            .await
3461            .unwrap();
3462        response.await.unwrap().unwrap();
3463        establish(&incoming, &mut outgoing).await;
3464        let corr = wire
3465            .open_stream("echo", Kind::Request, serde_json::Value::Null)
3466            .await
3467            .unwrap();
3468        let opening = outgoing.recv().await.unwrap();
3469        assert_eq!(opening.corr.as_deref(), Some(corr.as_str()));
3470
3471        incoming
3472            .send(frame(
3473                Kind::Response,
3474                Some(corr.clone()),
3475                serde_json::json!({ "ok": true }),
3476            ))
3477            .await
3478            .unwrap();
3479
3480        assert_eq!(
3481            deliveries.recv().await.unwrap().corr.as_deref(),
3482            Some(corr.as_str())
3483        );
3484        assert_eq!(closures.recv().await, Some(ClientOperationId::from(corr)));
3485        assert!(
3486            tokio::time::timeout(std::time::Duration::from_millis(25), closures.recv())
3487                .await
3488                .is_err()
3489        );
3490        handle.shutdown();
3491    }
3492
3493    #[tokio::test]
3494    async fn immediate_loopback_response_finds_registered_client_operation() {
3495        let (pipe, incoming, mut outgoing) = local_pipe();
3496        let wire = Wire::open(pipe);
3497        establish(&incoming, &mut outgoing).await;
3498        let loopback = tokio::spawn(async move {
3499            let opening = outgoing.recv().await.unwrap();
3500            let corr = opening.corr.unwrap();
3501            incoming
3502                .send(frame(
3503                    Kind::Response,
3504                    Some(corr),
3505                    serde_json::json!({ "ok": true }),
3506                ))
3507                .await
3508                .unwrap();
3509        });
3510
3511        let mut operation = wire
3512            .client_session()
3513            .start(
3514                "echo",
3515                Kind::Request,
3516                Bytes::new(),
3517                None,
3518                Default::default(),
3519            )
3520            .await
3521            .unwrap();
3522
3523        loopback.await.unwrap();
3524        assert_eq!(
3525            operation
3526                .try_next()
3527                .unwrap()
3528                .unwrap()
3529                .unwrap()
3530                .payload_json(),
3531            serde_json::json!({ "ok": true })
3532        );
3533        assert_eq!(operation.try_next().unwrap(), Some(None));
3534    }
3535
3536    #[tokio::test]
3537    async fn opening_send_failure_cleans_up_client_operation() {
3538        let (pipe, incoming, mut outgoing) = local_pipe();
3539        let wire = Wire::open(pipe);
3540        establish(&incoming, &mut outgoing).await;
3541        drop(outgoing);
3542
3543        let mut operation = wire
3544            .client_session()
3545            .start(
3546                "echo",
3547                Kind::Request,
3548                Bytes::new(),
3549                None,
3550                Default::default(),
3551            )
3552            .await
3553            .unwrap();
3554
3555        assert!(matches!(
3556            operation.next().await,
3557            Err(crate::ClientError::SessionClosed(_))
3558        ));
3559    }
3560
3561    #[tokio::test]
3562    async fn fetch_projects_response_and_ignores_late_terminal_frame() {
3563        let (pipe, incoming, mut outgoing) = local_pipe();
3564        let wire = Wire::open(pipe);
3565        establish(&incoming, &mut outgoing).await;
3566        let responder = tokio::spawn(async move {
3567            let opening = outgoing.recv().await.unwrap();
3568            let corr = opening.corr.unwrap();
3569            let mut response = frame(
3570                Kind::Response,
3571                Some(corr.clone()),
3572                serde_json::json!({ "ok": true }),
3573            );
3574            response.headers.insert("x-source".into(), "owner".into());
3575            incoming.send(response.clone()).await.unwrap();
3576            incoming.send(response).await.unwrap();
3577        });
3578
3579        let response = wire
3580            .client_session()
3581            .fetch(
3582                request("echo", b"request"),
3583                std::time::Duration::from_secs(1),
3584            )
3585            .await
3586            .unwrap();
3587
3588        assert_eq!(response.status(), http::StatusCode::OK);
3589        assert_eq!(response.headers()["x-source"], "owner");
3590        assert_eq!(response.body(), &Bytes::from_static(br#"{"ok":true}"#));
3591        responder.await.unwrap();
3592    }
3593
3594    #[tokio::test]
3595    async fn subscribe_and_discovery_are_isolated() {
3596        let (pipe, incoming, mut outgoing) = local_pipe();
3597        let wire = Wire::open(pipe);
3598        establish(&incoming, &mut outgoing).await;
3599        let client = wire.client_session();
3600        let mut subscription = client
3601            .subscribe(request("events", b""), None)
3602            .await
3603            .unwrap();
3604        let mut discovery = client
3605            .discover(DiscoverPlan::decode(&Bytes::new()).unwrap())
3606            .await
3607            .unwrap();
3608        let first = outgoing.recv().await.unwrap();
3609        let second = outgoing.recv().await.unwrap();
3610        let (subscribe_corr, discover_corr) = if first.kind == Kind::Subscribe {
3611            (first.corr.unwrap(), second.corr.unwrap())
3612        } else {
3613            (second.corr.unwrap(), first.corr.unwrap())
3614        };
3615
3616        incoming
3617            .send(frame(
3618                Kind::Event,
3619                Some(discover_corr.clone()),
3620                serde_json::json!({ "discover": 1 }),
3621            ))
3622            .await
3623            .unwrap();
3624        incoming
3625            .send(frame(
3626                Kind::Event,
3627                Some(subscribe_corr.clone()),
3628                serde_json::json!({ "event": 1 }),
3629            ))
3630            .await
3631            .unwrap();
3632        incoming
3633            .send(frame(
3634                Kind::Response,
3635                Some(discover_corr),
3636                serde_json::json!({ "done": true }),
3637            ))
3638            .await
3639            .unwrap();
3640
3641        assert_eq!(
3642            subscription.next().await.unwrap().unwrap().payload_json()["event"],
3643            1
3644        );
3645        assert_eq!(
3646            discovery.next().await.unwrap().unwrap().payload_json()["discover"],
3647            1
3648        );
3649        assert_eq!(
3650            discovery.next().await.unwrap().unwrap().payload_json()["done"],
3651            true
3652        );
3653    }
3654
3655    #[tokio::test]
3656    async fn dropping_a_saturated_operation_frees_its_session_for_idle_delivery() {
3657        let (pipe, incoming, mut outgoing) = local_pipe();
3658        let wire = Wire::open(pipe);
3659        establish(&incoming, &mut outgoing).await;
3660        let client = wire.client_session();
3661        let busy = client.subscribe(request("busy", b""), None).await.unwrap();
3662        let mut idle = client.subscribe(request("idle", b""), None).await.unwrap();
3663        let busy_open = outgoing.recv().await.unwrap();
3664        let idle_open = outgoing.recv().await.unwrap();
3665
3666        let burst = tokio::spawn(async move {
3667            for sequence in 0..257 {
3668                incoming
3669                    .send(frame(
3670                        Kind::Event,
3671                        busy_open.corr.clone(),
3672                        serde_json::json!({ "sequence": sequence }),
3673                    ))
3674                    .await
3675                    .unwrap();
3676            }
3677            incoming
3678                .send(frame(
3679                    Kind::Response,
3680                    idle_open.corr,
3681                    serde_json::json!({ "done": true }),
3682                ))
3683                .await
3684                .unwrap();
3685            incoming
3686        });
3687        tokio::time::sleep(std::time::Duration::from_millis(25)).await;
3688
3689        drop(busy);
3690
3691        let _incoming = tokio::time::timeout(std::time::Duration::from_secs(5), burst)
3692            .await
3693            .expect("dropping the saturated stream must unblock its session")
3694            .unwrap();
3695        assert_eq!(
3696            idle.next().await.unwrap().unwrap().payload_json()["done"],
3697            true
3698        );
3699    }
3700
3701    #[tokio::test]
3702    async fn out_of_order_unary_responses_reach_their_operations() {
3703        let (pipe, incoming, mut outgoing) = local_pipe();
3704        let wire = Wire::open(pipe);
3705        establish(&incoming, &mut outgoing).await;
3706        let client = wire.client_session();
3707        let mut first = client
3708            .start(
3709                "first",
3710                Kind::Request,
3711                Bytes::new(),
3712                None,
3713                Default::default(),
3714            )
3715            .await
3716            .unwrap();
3717        let mut second = client
3718            .start(
3719                "second",
3720                Kind::Request,
3721                Bytes::new(),
3722                None,
3723                Default::default(),
3724            )
3725            .await
3726            .unwrap();
3727        let first_open = outgoing.recv().await.unwrap();
3728        let second_open = outgoing.recv().await.unwrap();
3729
3730        incoming
3731            .send(frame(
3732                Kind::Response,
3733                second_open.corr,
3734                serde_json::json!({ "response": "second" }),
3735            ))
3736            .await
3737            .unwrap();
3738        incoming
3739            .send(frame(
3740                Kind::Response,
3741                first_open.corr,
3742                serde_json::json!({ "response": "first" }),
3743            ))
3744            .await
3745            .unwrap();
3746
3747        assert_eq!(
3748            first.next().await.unwrap().unwrap().payload_json()["response"],
3749            "first"
3750        );
3751        assert_eq!(
3752            second.next().await.unwrap().unwrap().payload_json()["response"],
3753            "second"
3754        );
3755    }
3756
3757    #[tokio::test]
3758    async fn mixed_unary_subscription_and_discovery_are_isolated() {
3759        let (pipe, incoming, mut outgoing) = local_pipe();
3760        let wire = Wire::open(pipe);
3761        establish(&incoming, &mut outgoing).await;
3762        let client = wire.client_session();
3763        let mut unary = client
3764            .start(
3765                "unary",
3766                Kind::Request,
3767                Bytes::new(),
3768                None,
3769                Default::default(),
3770            )
3771            .await
3772            .unwrap();
3773        let mut subscription = client
3774            .subscribe(request("events", b""), None)
3775            .await
3776            .unwrap();
3777        let mut discovery = client
3778            .discover(DiscoverPlan::decode(&Bytes::new()).unwrap())
3779            .await
3780            .unwrap();
3781        let mut unary_corr = None;
3782        let mut subscription_corr = None;
3783        let mut discovery_corr = None;
3784        for _ in 0..3 {
3785            let opening = outgoing.recv().await.unwrap();
3786            match opening.kind {
3787                Kind::Request => unary_corr = opening.corr,
3788                Kind::Subscribe => subscription_corr = opening.corr,
3789                Kind::Discover => discovery_corr = opening.corr,
3790                kind => panic!("unexpected opening kind {kind:?}"),
3791            }
3792        }
3793
3794        for (kind, corr, value) in [
3795            (Kind::Event, discovery_corr, "discover"),
3796            (Kind::Response, unary_corr, "unary"),
3797            (Kind::Event, subscription_corr, "subscription"),
3798        ] {
3799            incoming
3800                .send(frame(kind, corr, serde_json::json!({ "owner": value })))
3801                .await
3802                .unwrap();
3803        }
3804
3805        assert_eq!(
3806            unary.next().await.unwrap().unwrap().payload_json()["owner"],
3807            "unary"
3808        );
3809        assert_eq!(
3810            subscription.next().await.unwrap().unwrap().payload_json()["owner"],
3811            "subscription"
3812        );
3813        assert_eq!(
3814            discovery.next().await.unwrap().unwrap().payload_json()["owner"],
3815            "discover"
3816        );
3817    }
3818
3819    #[tokio::test]
3820    async fn cancelling_one_operation_does_not_complete_another() {
3821        let (pipe, incoming, mut outgoing) = local_pipe();
3822        let wire = Wire::open(pipe);
3823        establish(&incoming, &mut outgoing).await;
3824        let client = wire.client_session();
3825        let cancelled = client
3826            .subscribe(request("cancelled", b""), None)
3827            .await
3828            .unwrap();
3829        let mut live = client.subscribe(request("live", b""), None).await.unwrap();
3830        let cancelled_open = outgoing.recv().await.unwrap();
3831        let live_open = outgoing.recv().await.unwrap();
3832
3833        drop(cancelled);
3834        let cancel = outgoing.recv().await.unwrap();
3835        assert_eq!(cancel.kind, Kind::Cancel);
3836        assert_eq!(cancel.corr, cancelled_open.corr);
3837        incoming
3838            .send(frame(
3839                Kind::Event,
3840                live_open.corr,
3841                serde_json::json!({ "live": true }),
3842            ))
3843            .await
3844            .unwrap();
3845        assert_eq!(
3846            live.next().await.unwrap().unwrap().payload_json()["live"],
3847            true
3848        );
3849    }
3850
3851    #[tokio::test]
3852    async fn stream_drop_sends_exactly_one_cancel() {
3853        let (pipe, incoming, mut outgoing) = local_pipe();
3854        let wire = Wire::open(pipe);
3855        establish(&incoming, &mut outgoing).await;
3856        let stream = wire
3857            .client_session()
3858            .subscribe(request("events", b""), None)
3859            .await
3860            .unwrap();
3861        let opening = outgoing.recv().await.unwrap();
3862        let corr = opening.corr.unwrap();
3863
3864        drop(stream);
3865
3866        let cancel = tokio::time::timeout(std::time::Duration::from_secs(1), outgoing.recv())
3867            .await
3868            .unwrap()
3869            .unwrap();
3870        assert_eq!(cancel.kind, Kind::Cancel);
3871        assert_eq!(cancel.corr.as_deref(), Some(corr.as_str()));
3872        assert!(
3873            tokio::time::timeout(std::time::Duration::from_millis(25), outgoing.recv())
3874                .await
3875                .is_err()
3876        );
3877    }
3878
3879    #[tokio::test]
3880    async fn operation_timeout_is_typed_and_cancels_once() {
3881        let (pipe, incoming, mut outgoing) = local_pipe();
3882        let wire = Wire::open(pipe);
3883        establish(&incoming, &mut outgoing).await;
3884        let mut stream = wire
3885            .client_session()
3886            .subscribe(
3887                request("events", b""),
3888                Some(std::time::Duration::from_millis(10)),
3889            )
3890            .await
3891            .unwrap();
3892        let opening = outgoing.recv().await.unwrap();
3893
3894        assert!(matches!(
3895            stream.next().await,
3896            Err(crate::ClientError::Timeout(_))
3897        ));
3898        let cancel = outgoing.recv().await.unwrap();
3899        assert_eq!(cancel.kind, Kind::Cancel);
3900        assert_eq!(cancel.corr, opening.corr);
3901        drop(stream);
3902        assert!(
3903            tokio::time::timeout(std::time::Duration::from_millis(25), outgoing.recv())
3904                .await
3905                .is_err()
3906        );
3907    }
3908
3909    #[tokio::test]
3910    async fn session_close_completes_every_operation_once() {
3911        let (pipe, incoming, mut outgoing) = local_pipe();
3912        let wire = Wire::open(pipe);
3913        establish(&incoming, &mut outgoing).await;
3914        let client = wire.client_session();
3915        let mut first = client.subscribe(request("first", b""), None).await.unwrap();
3916        let mut second = client
3917            .subscribe(request("second", b""), None)
3918            .await
3919            .unwrap();
3920        outgoing.recv().await.unwrap();
3921        outgoing.recv().await.unwrap();
3922
3923        drop(incoming);
3924
3925        assert!(matches!(
3926            first.next().await,
3927            Err(crate::ClientError::SessionClosed(_))
3928        ));
3929        assert!(matches!(
3930            second.next().await,
3931            Err(crate::ClientError::SessionClosed(_))
3932        ));
3933        assert_eq!(first.next().await.unwrap(), None);
3934        assert_eq!(second.next().await.unwrap(), None);
3935    }
3936
3937    #[tokio::test]
3938    async fn a_session_the_core_rejects_closes_its_transport() {
3939        let (pipe, incoming, mut outgoing) = local_pipe();
3940        let _wire = Wire::open(pipe);
3941        assert_eq!(outgoing.recv().await.unwrap().kind, Kind::Hello);
3942        incoming
3943            .send(frame(
3944                Kind::Welcome,
3945                None,
3946                serde_json::json!({ "version": 999 }),
3947            ))
3948            .await
3949            .unwrap();
3950        let mut saw_error = false;
3951        let closed = loop {
3952            match tokio::time::timeout(std::time::Duration::from_millis(500), outgoing.recv()).await
3953            {
3954                Ok(Some(frame)) => {
3955                    if frame.kind == Kind::Error {
3956                        saw_error = true;
3957                    }
3958                }
3959                Ok(None) => break true,
3960                Err(_) => break false,
3961            }
3962        };
3963        assert!(
3964            saw_error,
3965            "the peer must receive the close-reason error frame before the transport closes"
3966        );
3967        assert!(
3968            closed,
3969            "core deciding a session invalid must close the transport, not leave it open"
3970        );
3971    }
3972
3973    struct ScriptedPipe {
3974        frames: std::collections::VecDeque<Result<Bytes, unb_transport::TransportError>>,
3975    }
3976
3977    impl futures_util::Stream for ScriptedPipe {
3978        type Item = Result<Bytes, unb_transport::TransportError>;
3979
3980        fn poll_next(
3981            mut self: Pin<&mut Self>,
3982            _cx: &mut std::task::Context<'_>,
3983        ) -> std::task::Poll<Option<Self::Item>> {
3984            std::task::Poll::Ready(self.frames.pop_front())
3985        }
3986    }
3987
3988    impl futures_util::Sink<Bytes> for ScriptedPipe {
3989        type Error = unb_transport::TransportError;
3990
3991        fn poll_ready(
3992            self: Pin<&mut Self>,
3993            _cx: &mut std::task::Context<'_>,
3994        ) -> std::task::Poll<Result<(), Self::Error>> {
3995            std::task::Poll::Ready(Ok(()))
3996        }
3997
3998        fn start_send(self: Pin<&mut Self>, _item: Bytes) -> Result<(), Self::Error> {
3999            Ok(())
4000        }
4001
4002        fn poll_flush(
4003            self: Pin<&mut Self>,
4004            _cx: &mut std::task::Context<'_>,
4005        ) -> std::task::Poll<Result<(), Self::Error>> {
4006            std::task::Poll::Ready(Ok(()))
4007        }
4008
4009        fn poll_close(
4010            self: Pin<&mut Self>,
4011            _cx: &mut std::task::Context<'_>,
4012        ) -> std::task::Poll<Result<(), Self::Error>> {
4013            std::task::Poll::Ready(Ok(()))
4014        }
4015    }
4016
4017    #[tokio::test]
4018    async fn a_malformed_frame_retires_the_session_with_a_typed_reason() {
4019        let cancellation = CancellationToken::new();
4020        let (effect_tx, mut effects) = mpsc::channel(64);
4021        let handle = ProtocolCoreHandle::spawn(
4022            ProtocolCore::new("node"),
4023            Arc::new(RecordingExecutor { effects: effect_tx }),
4024            cancellation,
4025            &tokio::runtime::Handle::current(),
4026        );
4027        let pipe = Pipe::Piped {
4028            pipe: Box::pin(ScriptedPipe {
4029                frames: std::collections::VecDeque::from([Ok(Bytes::from_static(
4030                    b"\xff\xff\xff\xffgarbage",
4031                ))]),
4032            }),
4033            initiator: true,
4034        };
4035        let (events, _events_rx) = mpsc::channel(16);
4036        handle
4037            .attach(SessionId::from("session"), pipe, None, events)
4038            .await
4039            .unwrap();
4040        let reason = tokio::time::timeout(std::time::Duration::from_secs(2), async {
4041            loop {
4042                if let CoreEffect::SessionRetired { reason, .. } = effects.recv().await.unwrap() {
4043                    break reason;
4044                }
4045            }
4046        })
4047        .await
4048        .expect("a malformed frame must retire the session");
4049        assert_eq!(reason, RetirementReason::TransportFailed);
4050        handle.shutdown();
4051    }
4052
4053    #[tokio::test]
4054    async fn a_transport_read_error_retires_the_session_with_a_typed_reason() {
4055        let cancellation = CancellationToken::new();
4056        let (effect_tx, mut effects) = mpsc::channel(64);
4057        let handle = ProtocolCoreHandle::spawn(
4058            ProtocolCore::new("node"),
4059            Arc::new(RecordingExecutor { effects: effect_tx }),
4060            cancellation,
4061            &tokio::runtime::Handle::current(),
4062        );
4063        let pipe = Pipe::Piped {
4064            pipe: Box::pin(ScriptedPipe {
4065                frames: std::collections::VecDeque::from([Err(
4066                    unb_transport::TransportError::Io("connection reset".into()),
4067                )]),
4068            }),
4069            initiator: true,
4070        };
4071        let (events, _events_rx) = mpsc::channel(16);
4072        handle
4073            .attach(SessionId::from("session"), pipe, None, events)
4074            .await
4075            .unwrap();
4076        let reason = tokio::time::timeout(std::time::Duration::from_secs(2), async {
4077            loop {
4078                if let CoreEffect::SessionRetired { reason, .. } = effects.recv().await.unwrap() {
4079                    break reason;
4080                }
4081            }
4082        })
4083        .await
4084        .expect("a transport read error must retire the session");
4085        assert_eq!(reason, RetirementReason::TransportFailed);
4086        handle.shutdown();
4087    }
4088
4089    #[tokio::test]
4090    async fn routes_two_sessions_and_submits_transport_closure() {
4091        let cancellation = CancellationToken::new();
4092        let (effect_tx, mut effects) = mpsc::channel(16);
4093        let handle = ProtocolCoreHandle::spawn(
4094            ProtocolCore::new("node"),
4095            Arc::new(RecordingExecutor { effects: effect_tx }),
4096            cancellation,
4097            &tokio::runtime::Handle::current(),
4098        );
4099        let (first_pipe, first_incoming, mut first_outgoing) = local_pipe();
4100        let (second_pipe, _second_incoming, mut second_outgoing) = local_pipe();
4101        let (first_events, _first_events_rx) = mpsc::channel(16);
4102        let (second_events, _second_events_rx) = mpsc::channel(16);
4103
4104        let first = handle
4105            .attach(SessionId::from("first"), first_pipe, None, first_events)
4106            .await
4107            .unwrap();
4108        let second = handle
4109            .attach(SessionId::from("second"), second_pipe, None, second_events)
4110            .await
4111            .unwrap();
4112
4113        assert!(first_outgoing.recv().await.is_some());
4114        assert!(second_outgoing.recv().await.is_some());
4115        assert!(Arc::ptr_eq(&first, &first));
4116        assert!(!Arc::ptr_eq(&first, &second));
4117
4118        drop(first_incoming);
4119        let retired = loop {
4120            let effect = effects.recv().await.unwrap();
4121            if matches!(
4122                effect,
4123                CoreEffect::SessionRetired {
4124                    ref session,
4125                    reason: RetirementReason::SessionClosed,
4126                } if session == &SessionId::from("first")
4127            ) {
4128                break effect;
4129            }
4130        };
4131        assert!(matches!(
4132            retired,
4133            CoreEffect::SessionRetired {
4134                reason: RetirementReason::SessionClosed,
4135                ..
4136            }
4137        ));
4138
4139        second.shutdown();
4140        handle.shutdown();
4141    }
4142
4143    #[tokio::test]
4144    async fn install_publishes_after_acceptance_and_before_effect_execution() {
4145        let cancellation = CancellationToken::new();
4146        let published = Arc::new(AtomicBool::new(false));
4147        let (effect_tx, mut effects) = mpsc::channel(1);
4148        let handle = ProtocolCoreHandle::spawn(
4149            ProtocolCore::new("node"),
4150            Arc::new(PublicationExecutor {
4151                published: published.clone(),
4152                effects: effect_tx,
4153            }),
4154            cancellation,
4155            &tokio::runtime::Handle::current(),
4156        );
4157        let (pipe, _incoming, mut outgoing) = local_pipe();
4158        let (events, _events_rx) = mpsc::channel(16);
4159        handle
4160            .attach(SessionId::from("session"), pipe, None, events)
4161            .await
4162            .unwrap();
4163        assert!(outgoing.recv().await.is_some());
4164
4165        handle
4166            .install(
4167                CoreInput::SessionClosed {
4168                    session: SessionId::from("session"),
4169                },
4170                {
4171                    let published = published.clone();
4172                    move || published.store(true, Ordering::SeqCst)
4173                },
4174            )
4175            .await
4176            .unwrap();
4177
4178        assert!(effects.recv().await.unwrap());
4179        handle.shutdown();
4180    }
4181
4182    #[tokio::test]
4183    async fn install_failure_does_not_publish() {
4184        let cancellation = CancellationToken::new();
4185        let (effect_tx, _effects) = mpsc::channel(1);
4186        let handle = ProtocolCoreHandle::spawn(
4187            ProtocolCore::new("node"),
4188            Arc::new(RecordingExecutor { effects: effect_tx }),
4189            cancellation,
4190            &tokio::runtime::Handle::current(),
4191        );
4192        let input = CoreInput::SessionOpened {
4193            session: SessionId::from("duplicate"),
4194            initiator: true,
4195            establish_peer: false,
4196            expected_peer: None,
4197        };
4198        handle.submit(input.clone()).await.unwrap();
4199        let published = Arc::new(AtomicBool::new(false));
4200
4201        let result = handle
4202            .install(input, {
4203                let published = published.clone();
4204                move || published.store(true, Ordering::SeqCst)
4205            })
4206            .await;
4207
4208        assert!(matches!(result, Err(WsError::Core(_))));
4209        assert!(!published.load(Ordering::SeqCst));
4210        handle.shutdown();
4211    }
4212
4213    #[test]
4214    fn operation_records_cannot_change_their_bound_correlation() {
4215        let response = frame(
4216            Kind::Response,
4217            Some("s2".into()),
4218            serde_json::json!({"ok": true}),
4219        );
4220        let error =
4221            decode_operation_message(response.encode(), "s1", OperationDirection::Return, 0)
4222                .unwrap_err();
4223        assert!(matches!(
4224            error,
4225            unb_transport::operation::OperationCodecError::Protocol(_)
4226        ));
4227        assert_eq!(
4228            error.reset_code(),
4229            unb_transport::operation::RESET_PROTOCOL
4230        );
4231    }
4232
4233    #[tokio::test]
4234    async fn streaming_operation_body_uses_head_chunks_and_no_wire_body_token() {
4235        let mut envelope = frame(Kind::Response, Some("s1".into()), serde_json::json!(null));
4236        envelope.body_token = Some("legacy-token-must-not-leak".into());
4237        let body: BodyStream = Box::pin(futures_util::stream::iter([
4238            Ok(Bytes::from_static(b"first")),
4239            Ok(Bytes::from_static(b"second")),
4240        ]));
4241        let (write, read) = tokio::io::duplex(4096);
4242        let sender = async move {
4243            let mut writer = unb_transport::operation::OperationRecordWriter::new(write);
4244            write_operation_envelope(
4245                &mut writer,
4246                envelope,
4247                Some(body),
4248                Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET)),
4249                CancellationToken::new(),
4250            )
4251            .await
4252            .unwrap();
4253            writer.finish().await.unwrap();
4254        };
4255        let receiver = async move {
4256            let mut reader = unb_transport::operation::OperationRecordReader::new(read);
4257            let (kind, payload) = reader.read_record().await.unwrap();
4258            assert_eq!(
4259                kind,
4260                unb_transport::operation::OperationRecordKind::Head
4261            );
4262            let head = Envelope::decode(payload).unwrap();
4263            assert_eq!(head.corr.as_deref(), Some("s1"));
4264            assert!(head.body_token.is_none());
4265            assert!(head.payload.is_empty());
4266            let mut collected = Vec::new();
4267            loop {
4268                let (kind, payload) = reader.read_record().await.unwrap();
4269                match kind {
4270                    unb_transport::operation::OperationRecordKind::BodyChunk => {
4271                        collected.extend_from_slice(&payload)
4272                    }
4273                    unb_transport::operation::OperationRecordKind::BodyEnd => break,
4274                    other => panic!("unexpected body record {other:?}"),
4275                }
4276            }
4277            assert_eq!(collected, b"firstsecond");
4278            let (kind, payload) = reader.read_record().await.unwrap();
4279            assert_eq!(
4280                kind,
4281                unb_transport::operation::OperationRecordKind::DirectionEnd
4282            );
4283            assert!(payload.is_empty());
4284            reader.expect_fin().await.unwrap();
4285        };
4286        tokio::join!(sender, receiver);
4287    }
4288
4289    #[tokio::test]
4290    async fn cancelling_a_client_operation_releases_a_late_response_body() {
4291        let handle = spawn_core();
4292        let (pipe, incoming, mut outgoing) = local_pipe();
4293        let wire = attach_client(&handle, "cancel-body", pipe).await;
4294        establish(&incoming, &mut outgoing).await;
4295
4296        let call = wire
4297            .client_session()
4298            .start(
4299                "echo",
4300                Kind::Request,
4301                Bytes::from_static(b"request"),
4302                None,
4303                Default::default(),
4304            )
4305            .await
4306            .unwrap();
4307        let corr = call.operation().as_str().to_owned();
4308        assert_eq!(
4309            outgoing.recv().await.unwrap().corr.as_deref(),
4310            Some(corr.as_str())
4311        );
4312        drop(call);
4313        let cancel = outgoing.recv().await.unwrap();
4314        assert_eq!(cancel.kind, Kind::Cancel);
4315        assert_eq!(cancel.corr.as_deref(), Some(corr.as_str()));
4316
4317        incoming
4318            .send(frame(
4319                Kind::Response,
4320                Some(corr),
4321                serde_json::json!({ "late": true }),
4322            ))
4323            .await
4324            .unwrap();
4325        tokio::time::timeout(std::time::Duration::from_secs(1), async {
4326            loop {
4327                let retained = handle
4328                    .token_bodies
4329                    .lock()
4330                    .expect("body token registry")
4331                    .keys()
4332                    .any(|(owner, _)| owner.as_str() == "cancel-body");
4333                if !retained {
4334                    break;
4335                }
4336                tokio::task::yield_now().await;
4337            }
4338        })
4339        .await
4340        .expect("late response body must be released");
4341        handle.shutdown();
4342    }
4343
4344    #[tokio::test]
4345    async fn session_retirement_releases_unclaimed_bodies() {
4346        let handle = spawn_core();
4347        let (pipe, incoming, mut outgoing) = local_pipe();
4348        let _wire = attach_client(&handle, "retired-body", pipe).await;
4349        establish(&incoming, &mut outgoing).await;
4350        handle
4351            .register_body(
4352                &SessionId::from("retired-body"),
4353                WireBody::Bytes(Bytes::from_static(b"retained")),
4354            )
4355            .unwrap();
4356        assert!(handle
4357            .token_bodies
4358            .lock()
4359            .expect("body token registry")
4360            .keys()
4361            .any(|(owner, _)| owner.as_str() == "retired-body"));
4362
4363        drop(incoming);
4364        tokio::time::timeout(std::time::Duration::from_secs(1), async {
4365            loop {
4366                let retained = handle
4367                    .token_bodies
4368                    .lock()
4369                    .expect("body token registry")
4370                    .keys()
4371                    .any(|(owner, _)| owner.as_str() == "retired-body");
4372                if !retained {
4373                    break;
4374                }
4375                tokio::task::yield_now().await;
4376            }
4377        })
4378        .await
4379        .expect("session retirement must release body capabilities");
4380        handle.shutdown();
4381    }
4382}