//! Shared dispatch loop for `alknet/call` connections. //! //! Both [`CallAdapter`]'s accept path and [`crate::client::CallClient`]'s //! connect path produce a [`CallConnection`] and hand it to the same dispatch //! loop here (ADR-017 §1): the loop reads `EventEnvelope` frames off accepted //! bidirectional streams, dispatches `call.requested` events against the //! operation registry, and writes the response back on the same stream. The //! connection-establishment half differs (accept vs dial); the dispatch half //! is shared. //! //! See `docs/architecture/crates/call/call-protocol.md` and //! `docs/architecture/crates/call/client-and-adapters.md` for the spec. use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; use alknet_core::auth::{AuthToken, Identity, IdentityProvider}; use alknet_core::types::StreamError; use futures::stream::StreamExt; use serde_json::Value; use tokio::task::JoinHandle; use tracing::{debug, warn}; use super::abort::AbortCascade; use super::connection::CallConnection; use super::wire::{ CallError, EventEnvelope, FrameFramedReader, FrameFramedWriter, ResponseEnvelope, EVENT_ABORTED, EVENT_REQUESTED, }; use crate::protocol::adapter::SessionOverlaySource; use crate::registry::context::{AbortPolicy, OperationContext, ScopedPeerEnv}; use crate::registry::env::{LocalOperationEnv, OperationEnv, PeerCompositeEnv}; use crate::registry::registration::{OperationRegistry, ResponseStream}; use crate::registry::spec::OperationType; const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30); const SWEEPER_INTERVAL: Duration = Duration::from_secs(10); /// Outcome of dispatching a `call.requested` event. The dispatcher branches on /// the registered operation's `op_type` (ADR-049 §6): `Query`/`Mutation` produce /// a single [`ResponseEnvelope`] (`Once`), `Subscription` produces a /// [`ResponseStream`] (`Stream`) that `handle_stream` pumps to the wire. /// /// This enum is the branch point the spec describes ("branches on `op_type` in /// `handle_stream`"): `dispatch` returns it and `handle_stream` matches on it, /// keeping the Once path (one frame, no `call.completed`) and the Stream path /// (each envelope → frame, `call.completed` on natural end) visibly distinct. pub enum DispatchResult { Once(ResponseEnvelope), Stream(ResponseStream), } impl std::fmt::Debug for DispatchResult { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { DispatchResult::Once(env) => f.debug_tuple("Once").field(env).finish(), DispatchResult::Stream(_) => { f.debug_tuple("Stream").field(&"").finish() } } } } /// Shared dispatcher for an established `CallConnection`. Constructed by /// both `CallAdapter` (accept path) and `CallClient` (connect path) and used /// to run the dispatch loop. Holds no per-connection state; the /// `CallConnection` is passed into `run_loop`. pub struct Dispatcher { pub registry: Arc, pub identity_provider: Arc, pub session_source: Option>, pub default_timeout: Duration, } impl Dispatcher { pub fn new( registry: Arc, identity_provider: Arc, ) -> Self { Self { registry, identity_provider, session_source: None, default_timeout: DEFAULT_TIMEOUT, } } pub fn with_session_source( mut self, source: Arc, ) -> Self { self.session_source = Some(source); self } pub fn with_timeout(mut self, timeout: Duration) -> Self { self.default_timeout = timeout; self } fn strip_leading_slash(operation_id: &str) -> &str { operation_id.strip_prefix('/').unwrap_or(operation_id) } pub(crate) fn resolve_identity( &self, connection_identity: Option, payload: &Value, ) -> Option { let auth_token = payload.get("auth_token").and_then(|v| v.as_str()); match auth_token { Some(token_str) => { let token = AuthToken { raw: token_str.as_bytes().to_vec(), }; match self.identity_provider.resolve_from_token(&token) { Some(identity) => Some(identity), None => connection_identity, } } None => connection_identity, } } pub fn compose_root_env( &self, connection: &CallConnection, context: &OperationContext, ) -> Arc { let base: Arc = Arc::new(LocalOperationEnv::new(Arc::clone(&self.registry))); let session = self .session_source .as_ref() .and_then(|s| s.overlay_for(context)); let mut env = PeerCompositeEnv::new(base); if let Some(session) = session { env = env.with_session(session); } if let Some(peer_id) = connection.identity().map(|identity| identity.id.clone()) { env.attach_peer(peer_id, connection.overlay_env()); } Arc::new(env) } pub(crate) fn build_root_context( &self, request_id: String, operation_name: &str, identity: Option, forwarded_for: Option, connection: &CallConnection, ) -> OperationContext { let registration = self.registry.registration(operation_name); let (composition_authority, capabilities, scoped_env) = match registration { Some(r) => ( r.composition_authority.clone(), r.capabilities.clone(), r.scoped_env.clone().unwrap_or_else(ScopedPeerEnv::empty), ), None => ( None, alknet_core::types::Capabilities::new(), ScopedPeerEnv::empty(), ), }; let stub_env: Arc = Arc::new(LocalOperationEnv::new(Arc::clone(&self.registry))); let mut context = OperationContext { request_id, parent_request_id: None, identity: identity.clone(), handler_identity: composition_authority, forwarded_for, capabilities, metadata: HashMap::new(), deadline: Some(Instant::now() + self.default_timeout), scoped_env, env: stub_env, abort_policy: AbortPolicy::default(), internal: false, }; context.env = self.compose_root_env(connection, &context); context } pub async fn dispatch_requested( &self, connection: &Arc, request_id: String, payload: Value, ) -> ResponseEnvelope { match self.dispatch(connection, request_id, payload).await { DispatchResult::Once(envelope) => envelope, DispatchResult::Stream(mut stream) => stream.next().await.unwrap_or_else(|| { ResponseEnvelope::error( String::new(), CallError::internal( "dispatch_requested called on a Subscription op; use the streaming path", ), ) }), } } /// Dispatch a `call.requested` event, branching on the registered /// operation's `op_type` (ADR-049 §6). `Query`/`Mutation` → `invoke()` → /// [`DispatchResult::Once`]; `Subscription` → `invoke_streaming()` → /// [`DispatchResult::Stream`]. Unknown ops and ACL failures resolve via /// the registry's own envelope/error paths (Once for `invoke`, a single /// error envelope for `invoke_streaming`). /// /// For the streaming branch the root context's deadline is cleared /// (`deadline: None`): subscriptions are long-running and unbounded — the /// 30s request/response deadline does not apply (ADR-049 §6, call-protocol /// Timeouts). The Once branch keeps the deadline from `build_root_context`. pub async fn dispatch( &self, connection: &Arc, request_id: String, payload: Value, ) -> DispatchResult { let operation_id = payload .get("operationId") .and_then(|v| v.as_str()) .unwrap_or(""); let operation_name = Self::strip_leading_slash(operation_id).to_string(); let connection_identity = connection.identity().cloned(); let identity = self.resolve_identity(connection_identity, &payload); let forwarded_for = payload .get("forwarded_for") .and_then(|v| serde_json::from_value::(v.clone()).ok()); let input = payload.get("input").cloned().unwrap_or(Value::Null); let is_subscription = self .registry .registration(&operation_name) .map(|r| r.spec.op_type == OperationType::Subscription) .unwrap_or(false); let mut context = self.build_root_context( request_id.clone(), &operation_name, identity, forwarded_for, connection, ); if is_subscription { context.deadline = None; let stream = self .registry .invoke_streaming(&operation_name, input, context); DispatchResult::Stream(stream) } else { let envelope = self.registry.invoke(&operation_name, input, context).await; DispatchResult::Once(envelope) } } pub async fn handle_abort(&self, connection: &Arc, request_id: &str) { let mut pending = connection.pending().lock(); let mut cascade = AbortCascade::new(&mut pending); let aborted = cascade.cascade_abort(request_id, AbortPolicy::AbortDependents); pending.handle_aborted(request_id); if !aborted.is_empty() { debug!(count = aborted.len(), "abort cascade evicted descendants"); } } pub(crate) async fn handle_stream( &self, connection: Arc, send: alknet_core::types::SendStream, recv: alknet_core::types::RecvStream, ) { let mut reader = FrameFramedReader::new(recv); let mut writer = FrameFramedWriter::new(send); loop { let envelope = match reader.read_frame().await { Ok(env) => env, Err(super::wire::FrameError::ConnectionClosed) => break, Err(err) => { warn!(error = %err, "stream frame read error; closing stream"); break; } }; match envelope.r#type.as_str() { EVENT_REQUESTED => { let request_id = envelope.id.clone(); let payload = envelope.payload.clone(); match self .dispatch(&connection, request_id.clone(), payload) .await { DispatchResult::Once(response) => { let event: EventEnvelope = response.into(); if let Err(err) = writer.write_frame(&event).await { warn!(error = %err, "failed to write response frame; closing stream"); break; } } DispatchResult::Stream(stream) => { self.pump_stream(&mut writer, &request_id, stream).await; } } } EVENT_ABORTED => { let request_id = envelope.id.clone(); self.handle_abort(&connection, &request_id).await; } other => { debug!(event_type = %other, id = %envelope.id, "ignoring non-requested/non-aborted event on inbound stream"); } } } } /// Pump a subscription's [`ResponseStream`] to the wire: each /// [`ResponseEnvelope`] becomes an [`EventEnvelope`] frame (`call.responded` /// for `Ok`, `call.error` for `Err`). On natural stream end (the stream /// returned `None` without the last item being an `Err`), write a /// `call.completed` frame. An `Err` envelope is terminal — the stream /// ends after it and we do NOT write `call.completed` (ADR-049 §6). /// /// If a frame write fails the pump stops early; the stream is dropped on /// return, releasing the handler's resources via `Drop` (ADR-016). The /// pump is cancellable: it runs inside the `handle_stream` task, so a /// `call.aborted` for this request ID (handled by `handle_abort` on /// another stream) or connection close cancels the task and drops the /// stream. pub(crate) async fn pump_stream( &self, writer: &mut super::wire::FrameFramedWriter, request_id: &str, mut stream: ResponseStream, ) { let mut last_was_error = false; while let Some(envelope) = stream.next().await { last_was_error = envelope.result.is_err(); let event: EventEnvelope = envelope.into(); if let Err(err) = writer.write_frame(&event).await { warn!(error = %err, "failed to write streaming frame; closing stream"); return; } } if !last_was_error { let completed = EventEnvelope::completed(request_id); if let Err(err) = writer.write_frame(&completed).await { warn!(error = %err, "failed to write call.completed"); } } } /// Run the shared dispatch loop over an established `CallConnection`: /// spawn the pending-entry sweeper, accept bidirectional streams until the /// connection closes, dispatch each stream via `handle_stream`, and fail /// outstanding pending requests on close. Returns when the connection is /// closed (accept loop yields `ConnectionClosed`/`StreamClosed`/`Timeout`). pub async fn run_loop(self, connection: Arc) { let pending = Arc::clone(connection.pending()); let quic = match connection.connection() { Some(c) => Arc::clone(c), None => { warn!("run_loop called with an overlay-only CallConnection; returning"); return; } }; let sweeper_pending = Arc::clone(&pending); let sweeper_handle: JoinHandle<()> = tokio::spawn(async move { let mut interval = tokio::time::interval(SWEEPER_INTERVAL); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; let evicted = sweeper_pending.lock().evict_expired(); if !evicted.is_empty() { debug!( count = evicted.len(), "sweeper evicted expired pending entries" ); } } }); loop { match quic.accept_bi().await { Ok((send, recv)) => { let conn = Arc::clone(&connection); let dispatcher = self.clone(); tokio::spawn(async move { dispatcher.handle_stream(conn, send, recv).await; }); } Err(StreamError::ConnectionClosed) => break, Err(StreamError::StreamClosed) => break, Err(StreamError::Timeout) => break, Err(err) => { warn!(error = %err, "accept_bi error; stopping accept loop"); break; } } } let failed = pending .lock() .fail_all(CallError::internal("connection closed")); if !failed.is_empty() { debug!( count = failed.len(), "failed pending requests on connection close" ); } sweeper_handle.abort(); } } impl Clone for Dispatcher { fn clone(&self) -> Self { Self { registry: Arc::clone(&self.registry), identity_provider: Arc::clone(&self.identity_provider), session_source: self.session_source.clone(), default_timeout: self.default_timeout, } } } #[cfg(test)] mod tests { use super::*; use crate::protocol::wire::{EVENT_COMPLETED, EVENT_ERROR, EVENT_RESPONDED}; use crate::registry::registration::{ make_handler, make_streaming_handler, HandlerKind, HandlerRegistration, OperationProvenance, }; use crate::registry::spec::{AccessControl, OperationSpec, OperationType, Visibility}; use alknet_core::auth::{AuthToken, Identity, IdentityProvider}; use alknet_core::types::{Capabilities, MockConnection}; use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::sync::Mutex as StdMutex; struct StubConnection { alpn: &'static [u8], addr: Option, closed: StdMutex>, } impl MockConnection for StubConnection { fn remote_alpn(&self) -> &[u8] { self.alpn } fn remote_addr(&self) -> Option { self.addr } fn close(&self, code: u32, reason: &str) { *self.closed.lock().unwrap() = Some((code, reason.to_string())); } } fn stub_connection() -> alknet_core::types::Connection { alknet_core::types::Connection::from_mock(Arc::new(StubConnection { alpn: b"alknet/call", addr: Some(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4321)), closed: StdMutex::new(None), })) } struct StaticIdentityProvider { tokens: StdMutex>, } impl StaticIdentityProvider { fn new() -> Self { Self { tokens: StdMutex::new(HashMap::new()), } } fn with_token(self, token: &str, identity: Identity) -> Self { self.tokens .lock() .unwrap() .insert(token.to_string(), identity); self } } impl IdentityProvider for StaticIdentityProvider { fn resolve_from_fingerprint(&self, _fp: &str) -> Option { None } fn resolve_from_token(&self, token: &AuthToken) -> Option { let token_str = String::from_utf8_lossy(&token.raw); self.tokens.lock().unwrap().get(token_str.as_ref()).cloned() } } fn identity_with_scopes(id: &str, scopes: &[&str]) -> Identity { Identity { id: id.to_string(), scopes: scopes.iter().map(|s| s.to_string()).collect(), resources: HashMap::new(), } } fn external_spec(name: &str, acl: AccessControl) -> OperationSpec { OperationSpec::new( name, OperationType::Query, Visibility::External, serde_json::json!({}), serde_json::json!({}), vec![], acl, ) } fn registry_with(name: &str, visibility: Visibility, acl: AccessControl) -> OperationRegistry { let mut registry = OperationRegistry::new(); registry .register(HandlerRegistration::new( OperationSpec::new( name, OperationType::Query, visibility, serde_json::json!({}), serde_json::json!({}), vec![], acl, ), HandlerKind::Once(make_handler(|input, context| async move { ResponseEnvelope::ok(context.request_id, input) })), OperationProvenance::Local, None, None, Capabilities::new(), )) .unwrap(); registry } fn dispatcher() -> Dispatcher { Dispatcher::new( Arc::new(OperationRegistry::new()), Arc::new(StaticIdentityProvider::new()), ) } #[tokio::test] async fn dispatch_authorized_peer_dispatches_and_populates_capabilities() { let caps = Capabilities::new().with_api_key("google", "k".to_string()); let mut registry = OperationRegistry::new(); let handler = make_handler(|_input, context| async move { let has_google = context.capabilities.get("google").is_some(); ResponseEnvelope::ok( context.request_id, serde_json::json!({ "has_google": has_google }), ) }); registry .register(HandlerRegistration::new( external_spec("admin/run", AccessControl::default()), HandlerKind::Once(handler), OperationProvenance::Local, None, None, caps, )) .unwrap(); let registry = Arc::new(registry); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/admin/run", "input": {}, }); let response = dp .dispatch_requested(&conn, "req-1".to_string(), payload) .await; let out = response.result.expect("dispatch ok"); assert_eq!(out["has_google"], Value::Bool(true)); } #[tokio::test] async fn dispatch_unauthorized_peer_returns_forbidden_capabilities_never_populated() { let caps = Capabilities::new().with_api_key("google", "k".to_string()); let mut registry = OperationRegistry::new(); let handler = make_handler(|_input, context| async move { let has_google = context.capabilities.get("google").is_some(); ResponseEnvelope::ok( context.request_id, serde_json::json!({ "has_google": has_google }), ) }); registry .register(HandlerRegistration::new( external_spec( "admin/run", AccessControl { required_scopes: vec!["admin".to_string()], ..Default::default() }, ), HandlerKind::Once(handler), OperationProvenance::Local, None, None, caps, )) .unwrap(); let registry = Arc::new(registry); let provider: Arc = Arc::new( StaticIdentityProvider::new() .with_token("alk_user", identity_with_scopes("regular-user", &["user"])), ); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/admin/run", "input": {}, "auth_token": "alk_user", }); let response = dp .dispatch_requested(&conn, "req-2".to_string(), payload) .await; match response.result { Err(e) => { assert_eq!(e.code, "FORBIDDEN"); assert!(e.message.contains("admin")); } other => panic!("expected FORBIDDEN, got {other:?}"), } } #[tokio::test] async fn dispatch_internal_op_from_wire_returns_not_found_before_acl() { let registry = Arc::new(registry_with( "secret/op", Visibility::Internal, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/secret/op", "input": {}, }); let response = dp .dispatch_requested(&conn, "req-3".to_string(), payload) .await; match response.result { Err(e) => { assert_eq!(e.code, "NOT_FOUND"); assert!(e.message.contains("secret/op")); } other => panic!("expected NOT_FOUND, got {other:?}"), } } #[tokio::test] async fn dispatch_connection_with_no_identity_produces_no_peer_id_in_env() { let registry = Arc::new(registry_with( "fs/readFile", Visibility::External, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = CallConnection::new(stub_connection()); let context = dp.build_root_context("req-4".to_string(), "fs/readFile", None, None, &conn); assert!( context.identity.is_none(), "no connection identity → context.identity is None" ); assert!( context.env.peer_ids().is_empty(), "no peer overlay attached when connection has no identity" ); } #[tokio::test] async fn dispatch_connection_with_identity_attaches_peer_overlay_keyed_by_identity_id() { let registry = Arc::new(registry_with( "fs/readFile", Visibility::External, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = CallConnection::new(stub_connection()); conn.connection() .expect("quic connection present") .set_identity(identity_with_scopes("worker-a", &[])) .expect("identity not yet set"); let context = dp.build_root_context("req-5".to_string(), "fs/readFile", None, None, &conn); assert_eq!( context.env.peer_ids(), vec!["worker-a".to_string()], "PeerId for connection comes from connection.identity().id" ); } #[tokio::test] async fn dispatch_extract_forwarded_for_from_payload_into_context() { let mut registry = OperationRegistry::new(); let handler = make_handler(|_input, context| async move { let forwarded_id = context.forwarded_for.as_ref().map(|i| i.id.clone()); ResponseEnvelope::ok( context.request_id, serde_json::json!({ "forwarded_for_id": forwarded_id }), ) }); registry .register(HandlerRegistration::new( external_spec("fs/readFile", AccessControl::default()), HandlerKind::Once(handler), OperationProvenance::Local, None, None, Capabilities::new(), )) .unwrap(); let registry = Arc::new(registry); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/fs/readFile", "input": {}, "forwarded_for": { "id": "alice", "scopes": ["fs:read"], "resources": {} }, }); let response = dp .dispatch_requested(&conn, "req-6".to_string(), payload) .await; let out = response.result.expect("ok"); assert_eq!(out["forwarded_for_id"], Value::String("alice".into())); } #[tokio::test] async fn dispatch_without_forwarded_for_field_is_none() { let mut registry = OperationRegistry::new(); let handler = make_handler(|_input, context| async move { let present = context.forwarded_for.is_some(); ResponseEnvelope::ok( context.request_id, serde_json::json!({ "present": present }), ) }); registry .register(HandlerRegistration::new( external_spec("fs/readFile", AccessControl::default()), HandlerKind::Once(handler), OperationProvenance::Local, None, None, Capabilities::new(), )) .unwrap(); let registry = Arc::new(registry); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/fs/readFile", "input": {}, }); let response = dp .dispatch_requested(&conn, "req-7".to_string(), payload) .await; let out = response.result.expect("ok"); assert_eq!(out["present"], Value::Bool(false)); } #[tokio::test] async fn dispatch_default_access_control_dispatches_to_any_peer() { let registry = Arc::new(registry_with( "echo/run", Visibility::External, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/echo/run", "input": { "msg": "hi" }, }); let response = dp .dispatch_requested(&conn, "req-8".to_string(), payload) .await; assert_eq!(response.result, Ok(serde_json::json!({ "msg": "hi" }))); } #[test] fn dispatcher_helper_compiles_with_full_signature() { let _dp = dispatcher(); } // --- non-QUIC (overlay-only) dispatch path ---------------------------- fn overlay_only_connection(identity: Identity) -> Arc { Arc::new(CallConnection::new_overlay_only(identity)) } #[tokio::test] async fn dispatch_requested_works_with_overlay_only_connection() { let registry = Arc::new(registry_with( "echo/run", Visibility::External, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = overlay_only_connection(identity_with_scopes("ws-peer", &[])); let payload = serde_json::json!({ "operationId": "/echo/run", "input": { "msg": "hello" }, }); let response = dp .dispatch_requested(&conn, "ws-req-1".to_string(), payload) .await; assert_eq!(response.request_id, "ws-req-1"); assert_eq!(response.result, Ok(serde_json::json!({ "msg": "hello" }))); } #[tokio::test] async fn dispatch_requested_overlay_only_attaches_peer_keyed_by_stored_identity() { let mut registry = OperationRegistry::new(); let handler = make_handler(|_input, context| async move { let peer_ids = context.env.peer_ids(); ResponseEnvelope::ok( context.request_id, serde_json::json!({ "peer_ids": peer_ids }), ) }); registry .register(HandlerRegistration::new( external_spec("fs/readFile", AccessControl::default()), HandlerKind::Once(handler), OperationProvenance::Local, None, None, Capabilities::new(), )) .unwrap(); let registry = Arc::new(registry); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = overlay_only_connection(identity_with_scopes("ws-peer", &[])); let payload = serde_json::json!({ "operationId": "/fs/readFile", "input": {}, }); let response = dp .dispatch_requested(&conn, "ws-req-2".to_string(), payload) .await; let out = response.result.expect("ok"); assert_eq!(out["peer_ids"], serde_json::json!(["ws-peer"])); } #[tokio::test] async fn dispatch_requested_overlay_only_unknown_op_returns_not_found() { let registry = Arc::new(OperationRegistry::new()); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = overlay_only_connection(identity_with_scopes("ws-peer", &[])); let payload = serde_json::json!({ "operationId": "/no/such/op", "input": {}, }); let response = dp .dispatch_requested(&conn, "ws-req-3".to_string(), payload) .await; match response.result { Err(e) => assert_eq!(e.code, "NOT_FOUND"), other => panic!("expected NOT_FOUND, got {other:?}"), } } #[tokio::test] async fn handle_abort_works_with_overlay_only_connection() { let registry = Arc::new(registry_with( "echo/run", Visibility::External, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = overlay_only_connection(identity_with_scopes("ws-peer", &[])); let parent_id = "ws-abort-root".to_string(); let child_id = "ws-abort-child".to_string(); { let mut pending = conn.pending().lock(); pending.register_call( parent_id.clone(), Instant::now() + Duration::from_secs(30), None, ); pending.register_call( child_id.clone(), Instant::now() + Duration::from_secs(30), Some(parent_id.clone()), ); } assert!(conn.pending().lock().contains(&parent_id)); assert!(conn.pending().lock().contains(&child_id)); dp.handle_abort(&conn, &parent_id).await; assert!( !conn.pending().lock().contains(&parent_id), "parent entry removed after abort" ); assert!( !conn.pending().lock().contains(&child_id), "child aborted by cascade" ); } #[tokio::test] async fn handle_abort_unknown_request_id_is_noop_for_overlay_only() { let registry = Arc::new(OperationRegistry::new()); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = overlay_only_connection(identity_with_scopes("ws-peer", &[])); dp.handle_abort(&conn, "totally-unknown").await; assert!(conn.pending().lock().is_empty()); } #[tokio::test] async fn overlay_only_full_dispatch_round_trip_returns_response_envelope() { let registry = Arc::new(registry_with( "echo/run", Visibility::External, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = overlay_only_connection(identity_with_scopes("ws-peer", &[])); let payload = serde_json::json!({ "operationId": "/echo/run", "input": { "v": 42 }, }); let request_id = "ws-roundtrip-1".to_string(); let response = dp .dispatch_requested(&conn, request_id.clone(), payload) .await; assert!(response.result.is_ok()); let envelope: EventEnvelope = response.into(); assert_eq!(envelope.r#type, EVENT_RESPONDED); assert_eq!(envelope.id, "ws-roundtrip-1"); assert_eq!( envelope.payload.get("output"), Some(&serde_json::json!({ "v": 42 })) ); } // --- streaming dispatch branch (ADR-049 §6) --------------------------- fn subscription_spec(name: &str, acl: AccessControl) -> OperationSpec { OperationSpec::new( name, OperationType::Subscription, Visibility::External, serde_json::json!({}), serde_json::json!({}), vec![], acl, ) } fn encode_frame(envelope: &EventEnvelope) -> Vec { let body = serde_json::to_vec(envelope).expect("serialize envelope"); let mut buf = (body.len() as u32).to_be_bytes().to_vec(); buf.extend_from_slice(&body); buf } async fn read_all_frames( reader: &mut (impl tokio::io::AsyncRead + Unpin), ) -> Vec { let mut buf = Vec::new(); use tokio::io::AsyncReadExt; let _ = reader.read_to_end(&mut buf).await; let mut frames = Vec::new(); let mut cursor = std::io::Cursor::new(buf); loop { let mut len_buf = [0u8; 4]; match tokio::io::AsyncReadExt::read_exact(&mut cursor, &mut len_buf).await { Ok(_) => {} Err(_) => break, } let len = u32::from_be_bytes(len_buf) as usize; let mut body = vec![0u8; len]; if tokio::io::AsyncReadExt::read_exact(&mut cursor, &mut body) .await .is_err() { break; } let envelope: EventEnvelope = serde_json::from_slice(&body).expect("deserialize written frame"); frames.push(envelope); } frames } fn registry_with_subscription( name: &str, handler: crate::registry::registration::StreamingHandler, ) -> Arc { let mut registry = OperationRegistry::new(); registry .register(HandlerRegistration::new( subscription_spec(name, AccessControl::default()), HandlerKind::Stream(handler), OperationProvenance::Local, None, None, Capabilities::new(), )) .unwrap(); Arc::new(registry) } #[tokio::test] async fn dispatch_subscription_returns_stream_result() { let handler = make_streaming_handler(|input, ctx| { futures::stream::iter(vec![ ResponseEnvelope::ok(ctx.request_id.clone(), input.clone()), ResponseEnvelope::ok(ctx.request_id.clone(), serde_json::json!({"done": true})), ]) }); let registry = registry_with_subscription("events/stream", handler); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/events/stream", "input": { "v": 1 }, }); match dp.dispatch(&conn, "sub-1".to_string(), payload).await { DispatchResult::Stream(mut stream) => { use futures::stream::StreamExt; let first = stream.next().await.expect("first envelope"); assert_eq!(first.request_id, "sub-1"); assert_eq!(first.result, Ok(serde_json::json!({ "v": 1 }))); let second = stream.next().await.expect("second envelope"); assert_eq!(second.result, Ok(serde_json::json!({ "done": true }))); assert!( stream.next().await.is_none(), "stream ends after two values" ); } other => panic!("expected Stream, got {other:?}"), } } #[tokio::test] async fn dispatch_subscription_clears_deadline_to_none() { let handler = make_streaming_handler(|_input, ctx| { let deadline = ctx.deadline; futures::stream::iter(vec![ResponseEnvelope::ok( ctx.request_id.clone(), serde_json::json!({ "deadline_is_none": deadline.is_none() }), )]) }); let registry = registry_with_subscription("events/stream", handler); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/events/stream", "input": {}, }); match dp.dispatch(&conn, "sub-dl".to_string(), payload).await { DispatchResult::Stream(mut stream) => { use futures::stream::StreamExt; let env = stream.next().await.expect("one envelope"); let out = env.result.expect("ok"); assert_eq!(out["deadline_is_none"], Value::Bool(true)); } other => panic!("expected Stream, got {other:?}"), } } #[tokio::test] async fn dispatch_query_keeps_deadline_some() { let mut registry = OperationRegistry::new(); let handler = make_handler(|_input, ctx| async move { let deadline_is_some = ctx.deadline.is_some(); ResponseEnvelope::ok( ctx.request_id.clone(), serde_json::json!({ "deadline_is_some": deadline_is_some }), ) }); registry .register(HandlerRegistration::new( external_spec("echo/run", AccessControl::default()), HandlerKind::Once(handler), OperationProvenance::Local, None, None, Capabilities::new(), )) .unwrap(); let registry = Arc::new(registry); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let payload = serde_json::json!({ "operationId": "/echo/run", "input": {}, }); match dp.dispatch(&conn, "q-1".to_string(), payload).await { DispatchResult::Once(env) => { let out = env.result.expect("ok"); assert_eq!(out["deadline_is_some"], Value::Bool(true)); } other => panic!("expected Once, got {other:?}"), } } #[tokio::test] async fn handle_stream_subscription_pumps_each_frame_then_completed() { let handler = make_streaming_handler(|input, ctx| { let first = input.clone(); let rid = ctx.request_id.clone(); futures::stream::iter(vec![ ResponseEnvelope::ok(rid.clone(), first), ResponseEnvelope::ok(rid.clone(), serde_json::json!({"n": 2})), ResponseEnvelope::ok(rid, serde_json::json!({"n": 3})), ]) }); let registry = registry_with_subscription("events/stream", handler); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let request = EventEnvelope::requested( "sub-pump-1", serde_json::json!({ "operationId": "/events/stream", "input": { "n": 1 }, }), ); let recv = tokio::io::BufReader::new(std::io::Cursor::new(encode_frame(&request))); let (send, mut sink) = tokio::io::duplex(8 * 1024); let send = alknet_core::types::SendStream::from_mock(send); let recv = alknet_core::types::RecvStream::from_mock(recv); dp.handle_stream(conn, send, recv).await; let frames = read_all_frames(&mut sink).await; assert_eq!(frames.len(), 4, "3 responded + 1 completed"); for (i, f) in frames[..3].iter().enumerate() { assert_eq!(f.r#type, EVENT_RESPONDED, "frame {i} is call.responded"); assert_eq!(f.id, "sub-pump-1"); } assert_eq!(frames[3].r#type, EVENT_COMPLETED); assert_eq!(frames[3].id, "sub-pump-1"); assert_eq!(frames[3].payload, serde_json::json!({})); } #[tokio::test] async fn handle_stream_subscription_error_is_terminal_no_completed() { let handler = make_streaming_handler(|_input, ctx| { let rid = ctx.request_id.clone(); futures::stream::iter(vec![ ResponseEnvelope::ok(rid.clone(), serde_json::json!({"ok": true})), ResponseEnvelope::error(rid.clone(), CallError::internal("boom")), ]) }); let registry = registry_with_subscription("events/stream", handler); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let request = EventEnvelope::requested( "sub-err-1", serde_json::json!({ "operationId": "/events/stream", "input": {}, }), ); let recv = tokio::io::BufReader::new(std::io::Cursor::new(encode_frame(&request))); let (send, mut sink) = tokio::io::duplex(8 * 1024); let send = alknet_core::types::SendStream::from_mock(send); let recv = alknet_core::types::RecvStream::from_mock(recv); dp.handle_stream(conn, send, recv).await; let frames = read_all_frames(&mut sink).await; assert_eq!(frames.len(), 2, "1 responded + 1 error, no completed"); assert_eq!(frames[0].r#type, EVENT_RESPONDED); assert_eq!(frames[1].r#type, EVENT_ERROR); assert_eq!(frames[1].id, "sub-err-1"); assert_eq!( frames[1].payload.get("code"), Some(&Value::String("INTERNAL".into())) ); } #[tokio::test] async fn handle_stream_query_dispatch_unchanged_one_frame_no_completed() { let registry = Arc::new(registry_with( "echo/run", Visibility::External, AccessControl::default(), )); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let request = EventEnvelope::requested( "q-pump-1", serde_json::json!({ "operationId": "/echo/run", "input": { "msg": "hi" }, }), ); let recv = tokio::io::BufReader::new(std::io::Cursor::new(encode_frame(&request))); let (send, mut sink) = tokio::io::duplex(8 * 1024); let send = alknet_core::types::SendStream::from_mock(send); let recv = alknet_core::types::RecvStream::from_mock(recv); dp.handle_stream(conn, send, recv).await; let frames = read_all_frames(&mut sink).await; assert_eq!(frames.len(), 1, "query: exactly one frame, no completed"); assert_eq!(frames[0].r#type, EVENT_RESPONDED); assert_eq!(frames[0].id, "q-pump-1"); assert_eq!( frames[0].payload.get("output"), Some(&serde_json::json!({ "msg": "hi" })) ); } #[tokio::test] async fn handle_stream_subscription_unknown_op_yields_single_error_no_completed() { let registry = Arc::new(OperationRegistry::new()); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let request = EventEnvelope::requested( "sub-missing-1", serde_json::json!({ "operationId": "/no/such/stream", "input": {}, }), ); let recv = tokio::io::BufReader::new(std::io::Cursor::new(encode_frame(&request))); let (send, mut sink) = tokio::io::duplex(8 * 1024); let send = alknet_core::types::SendStream::from_mock(send); let recv = alknet_core::types::RecvStream::from_mock(recv); dp.handle_stream(conn, send, recv).await; let frames = read_all_frames(&mut sink).await; assert_eq!(frames.len(), 1, "unknown op: single error, no completed"); assert_eq!(frames[0].r#type, EVENT_ERROR); assert_eq!(frames[0].id, "sub-missing-1"); assert_eq!( frames[0].payload.get("code"), Some(&Value::String("NOT_FOUND".into())) ); } #[tokio::test] async fn handle_stream_aborted_for_streaming_request_drops_stream() { use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc as StdArc; let dropped = StdArc::new(AtomicBool::new(false)); let dropped_clone = StdArc::clone(&dropped); let handler = make_streaming_handler(move |_input, ctx| { let rid = ctx.request_id.clone(); let flag = StdArc::clone(&dropped_clone); struct DropGuard(StdArc); impl Drop for DropGuard { fn drop(&mut self) { self.0.store(true, Ordering::SeqCst); } } let guard = DropGuard(StdArc::clone(&flag)); futures::stream::poll_fn(move |_cx| { if flag.load(Ordering::SeqCst) { return std::task::Poll::Ready(None); } std::task::Poll::Ready(Some(ResponseEnvelope::ok( rid.clone(), serde_json::json!({"tick": 1}), ))) }) .map(move |env| { let _keep_guard = &guard; env }) }); let registry = registry_with_subscription("events/stream", handler); let provider: Arc = Arc::new(StaticIdentityProvider::new()); let dp = Dispatcher::new(registry, provider); let conn = Arc::new(CallConnection::new(stub_connection())); let request = EventEnvelope::requested( "sub-abort-1", serde_json::json!({ "operationId": "/events/stream", "input": {}, }), ); let recv = tokio::io::BufReader::new(std::io::Cursor::new(encode_frame(&request))); let (send, _sink) = tokio::io::duplex(8 * 1024); let send = alknet_core::types::SendStream::from_mock(send); let recv = alknet_core::types::RecvStream::from_mock(recv); let conn_clone = Arc::clone(&conn); let dp_clone = dp.clone(); let handle = tokio::spawn(async move { dp_clone.handle_stream(conn_clone, send, recv).await; }); tokio::time::sleep(std::time::Duration::from_millis(50)).await; dp.handle_abort(&conn, "sub-abort-1").await; assert!( !conn.pending().lock().contains("sub-abort-1"), "abort removes the pending entry" ); handle.abort(); let _ = handle.await; assert!( dropped.load(Ordering::SeqCst), "stream future dropped → Drop guard released handler resources" ); } }