Refactor WebSocket connector and session management

- Simplified `WebSocketConnectorImpl` by removing the `raw_payload` field and related logic, delegating payload framing to the `WsCodec`.
- Introduced `WsDispatch` and `WsSession` for handling WebSocket connections, separating concerns for dispatching messages and managing session state.
- Added `WsServerConnection` to manage WebSocket connections on the server side, including multi-topic subscription handling.
- Retired the previous session management logic in favor of a more modular approach using `run_session`.
- Implemented a new `SnapshotProvider` trait for late-join functionality, allowing clients to receive the current state of topics upon subscription.
- Enhanced the transport layer with `WsDialer` and `WsClientConnection` for client-side WebSocket handling.
- Updated server state management to include shared dispatch and client manager.
- Added tests for multi-topic subscription handling and ensured compatibility with existing message protocols.

Author: lxsaah

aimdb-client/tests/pump_client.rs

@@ -79,6 +79,11 @@ async fn pump_client_mirrors_record_both_directions() {
         .with_connector(AimxClientConnector::new(&sock).with_config(ClientConfig {
             reconnect: true,
             reconnect_delay: Duration::from_millis(50),
+            max_reconnect_delay: Duration::from_millis(50),
+            max_reconnect_attempts: 0,
+            keepalive_interval: None,
+            max_offline_queue: usize::MAX,
+            topic_routed_subs: false,
             sends_hello: false,
         }));
     cb.configure::<Msg>("cfg", |reg| {

aimdb-core/src/session/aimx/codec.rs

@@ -176,6 +176,11 @@ impl EnvelopeCodec for AimxCodec {
                 write_frame(out, &frame)
             }
             Outbound::Pong => write_frame(out, &Frame::tagged("pong")),
+            // AimX has no explicit subscribe ack (the client owns the id; events
+            // carry it back). `run_session` only emits this when `acks_subscribe`
+            // is set, which the AimX server leaves off — so this is unreachable on
+            // the AimX wire.
+            Outbound::Subscribed { .. } => Err(CodecError::Malformed),
         }
     }
 

aimdb-core/src/session/aimx/dispatch.rs

@@ -102,13 +102,19 @@ where
         })
     }
 
-    fn subscribe(&mut self, topic: &str) -> Result<BoxStream<'static, Payload>, RpcError> {
+    fn subscribe<'a>(
+        &'a mut self,
+        topic: &'a str,
+    ) -> BoxFut<'a, Result<BoxStream<'static, Payload>, RpcError>> {
         // The engine owns the subscription lifecycle (keyed by request id) and
         // the per-connection cap (SessionLimits); no `generate_subscription_id`
-        // / `max_subs` bookkeeping here.
-        let stream =
-            crate::remote::stream::stream_record_updates(&self.db, topic).map_err(map_db_err)?;
-        Ok(Box::pin(stream.map(|v| to_payload(&v))))
+        // / `max_subs` bookkeeping here. AimX has no async authorization, so this
+        // is a trivial wrapper.
+        Box::pin(async move {
+            let stream = crate::remote::stream::stream_record_updates(&self.db, topic)
+                .map_err(map_db_err)?;
+            Ok(Box::pin(stream.map(|v| to_payload(&v))) as BoxStream<'static, Payload>)
+        })
     }
 
     fn write<'a>(
@@ -410,6 +416,8 @@ where
             max_subs_per_connection: config.max_subs_per_connection,
         },
         reads_hello: false,
+        // AimX's subscribe ack stays implicit (events flow); no explicit ack frame.
+        acks_subscribe: false,
     };
     let dispatch = Arc::new(AimxDispatch::new(db, config));
     let listener = UdsListener::new(listener);

aimdb-core/src/session/client.rs

@@ -33,8 +33,29 @@ use crate::{AimDb, RuntimeAdapter};
 pub struct ClientConfig {
     /// Redial after a dropped/failed connection instead of ending the engine.
     pub reconnect: bool,
-    /// Delay before each redial when `reconnect` is set.
+    /// Base delay before the first redial when `reconnect` is set. Subsequent
+    /// redials grow this exponentially, capped at [`max_reconnect_delay`](Self::max_reconnect_delay).
     pub reconnect_delay: Duration,
+    /// Upper bound for the exponential reconnect backoff. Defaults to
+    /// [`reconnect_delay`](Self::reconnect_delay) (i.e. no escalation — a fixed
+    /// delay, preserving the pre-Phase-4 behavior).
+    pub max_reconnect_delay: Duration,
+    /// Maximum redial attempts before the engine gives up. `0` = unlimited
+    /// (the default).
+    pub max_reconnect_attempts: usize,
+    /// If set, send a keepalive `Ping` on this interval while a connection is
+    /// idle. `None` (default) disables keepalive.
+    pub keepalive_interval: Option<Duration>,
+    /// Cap on caller commands buffered while disconnected; the oldest are dropped
+    /// past this bound. Defaults to `usize::MAX` (effectively unbounded — the
+    /// pre-Phase-4 behavior).
+    pub max_offline_queue: usize,
+    /// Key the subscription demux by **topic** instead of the engine request id.
+    /// `false` (default, AimX-style) — events carry the request id back, demux by
+    /// id. `true` (WS-style) — the wire pushes data keyed by topic with no id, so
+    /// the codec's `decode_outbound` returns the topic as `Event.sub` and the
+    /// engine routes by topic.
+    pub topic_routed_subs: bool,
     /// Send a Ping handshake on connect and wait for the Pong before accepting
     /// caller commands (the proactive "handshake-as-caller"). Mirrors the
     /// server's `reads_hello`; a real protocol swaps Ping/Pong for its Hello.
@@ -46,11 +67,31 @@ impl Default for ClientConfig {
         Self {
             reconnect: true,
             reconnect_delay: Duration::from_millis(200),
+            max_reconnect_delay: Duration::from_millis(200),
+            max_reconnect_attempts: 0,
+            keepalive_interval: None,
+            max_offline_queue: usize::MAX,
+            topic_routed_subs: false,
             sends_hello: false,
         }
     }
 }
 
+/// Exponential backoff for the `attempt`-th redial (1-based), capped at
+/// [`ClientConfig::max_reconnect_delay`]. Defaults collapse this to a fixed
+/// `reconnect_delay` (max == base), preserving pre-Phase-4 behavior.
+fn backoff_delay(config: &ClientConfig, attempt: usize) -> Duration {
+    let base = config.reconnect_delay;
+    let cap = config.max_reconnect_delay.max(base);
+    let shift = attempt.saturating_sub(1).min(16) as u32;
+    base.saturating_mul(1u32 << shift).min(cap)
+}
+
+/// Bound the offline backlog: drop the oldest buffered commands beyond `cap`.
+fn bound_offline_queue(cmd_rx: &mut mpsc::UnboundedReceiver<ClientCmd>, cap: usize) {
+    while cmd_rx.len() > cap && cmd_rx.try_recv().is_ok() {}
+}
+
 /// A cheap-clone handle to a running [`run_client`] engine — the caller-facing
 /// RPC surface. Every method funnels a command to the engine, which owns the
 /// pending-call map and the wire.
@@ -165,33 +206,62 @@ async fn client_loop<D, C>(
     D: Dialer,
     C: EnvelopeCodec,
 {
+    // Consecutive failed attempts since the last successful connection; drives
+    // exponential backoff and the optional attempt cap.
+    let mut attempt: usize = 0;
     loop {
         let conn = match dialer.connect().await {
-            Ok(conn) => conn,
+            Ok(conn) => {
+                attempt = 0;
+                conn
+            }
             Err(_e) => {
                 #[cfg(feature = "tracing")]
                 tracing::warn!("client dial failed: {:?}", _e);
-                if config.reconnect {
-                    tokio::time::sleep(config.reconnect_delay).await;
-                    continue;
+                match reconnect_after(&mut attempt, &config, &mut cmd_rx).await {
+                    true => continue,
+                    false => return,
                 }
-                return;
             }
         };
 
         match drive_connection(conn, &codec, &mut cmd_rx, &config).await {
             Ended::HandlesDropped => return,
             Ended::Disconnected => {
-                if config.reconnect {
-                    tokio::time::sleep(config.reconnect_delay).await;
-                    continue;
+                match reconnect_after(&mut attempt, &config, &mut cmd_rx).await {
+                    true => continue,
+                    false => return,
                 }
-                return;
             }
         }
     }
 }
 
+/// Decide whether to redial: honor `reconnect`, the attempt cap, the offline-queue
+/// bound, and the exponential backoff sleep. Returns `true` to retry, `false` to
+/// stop the engine.
+async fn reconnect_after(
+    attempt: &mut usize,
+    config: &ClientConfig,
+    cmd_rx: &mut mpsc::UnboundedReceiver<ClientCmd>,
+) -> bool {
+    if !config.reconnect {
+        return false;
+    }
+    *attempt += 1;
+    if config.max_reconnect_attempts != 0 && *attempt >= config.max_reconnect_attempts {
+        #[cfg(feature = "tracing")]
+        tracing::warn!(
+            "client giving up after {} reconnect attempts",
+            config.max_reconnect_attempts
+        );
+        return false;
+    }
+    bound_offline_queue(cmd_rx, config.max_offline_queue);
+    tokio::time::sleep(backoff_delay(config, *attempt)).await;
+    true
+}
+
 /// Drive one dialed [`Connection`]: optional handshake, then `biased` demux of
 /// server frames (resolve `Reply` by `id`, route `Event`/`Snapshot` to their
 /// subscription channels) interleaved with caller commands. Pending state is
@@ -229,6 +299,9 @@ where
         }
     }
 
+    // Optional keepalive ticker — `None` parks the arm forever (see below).
+    let mut keepalive = config.keepalive_interval.map(tokio::time::interval);
+
     loop {
         tokio::select! {
             biased;
@@ -269,10 +342,31 @@ where
                         }
                     }
                     Ok(Outbound::Pong) => {}
+                    // Explicit subscribe ack (WS). Informational — the local
+                    // event sink already exists from the Subscribe command, so
+                    // there is nothing to route; just confirm liveness.
+                    Ok(Outbound::Subscribed { .. }) => {}
                     Err(_e) => continue, // skip a malformed frame, keep the connection
                 }
             }
 
+            // ---- keepalive: send a Ping when the ticker fires --------------
+            // With no interval configured the arm parks on `pending()` forever,
+            // so it never wins the `select!`.
+            _ = async {
+                match keepalive.as_mut() {
+                    Some(i) => { i.tick().await; }
+                    None => std::future::pending::<()>().await,
+                }
+            } => {
+                out.clear();
+                if codec.encode_inbound(Inbound::Ping, &mut out).is_ok()
+                    && conn.send(&out).await.is_err()
+                {
+                    return Ended::Disconnected;
+                }
+            }
+
             // ---- caller commands from ClientHandle -------------------------
             cmd = cmd_rx.recv() => {
                 let cmd = match cmd {
@@ -299,7 +393,14 @@ where
                     ClientCmd::Subscribe { topic, events } => {
                         let id = next_id;
                         next_id += 1;
-                        subs.insert(id.to_string(), events);
+                        // Topic-routed (WS): the wire pushes data keyed by topic,
+                        // so demux by topic; id-routed (AimX): events echo the id.
+                        let key = if config.topic_routed_subs {
+                            topic.clone()
+                        } else {
+                            id.to_string()
+                        };
+                        subs.insert(key, events);
                         out.clear();
                         let sent = codec
                             .encode_inbound(Inbound::Subscribe { id, topic }, &mut out)

aimdb-core/src/session/mod.rs

@@ -80,22 +80,78 @@ pub type TransportResult<T> = Result<T, TransportError>;
 // Supporting types (stubs — sufficient for the signatures to compile)
 // ===========================================================================
 
-/// Remote-peer metadata carried by a [`Connection`] (remote addr, headers,
-/// pre-resolved auth).
+/// Remote-peer metadata carried by a [`Connection`] (remote addr, pre-resolved
+/// auth).
 ///
-/// Opaque placeholder. The concrete fields — and whether one shape carries both
-/// AimX `SecurityPolicy` and WS `Permissions` — are **deferred to Phase 4**.
-#[derive(Debug, Clone, Default)]
+/// **Phase 4 (resolved — doc 037 auth-context gate).** One shape serves both
+/// connectors: a neutral [`peer_addr`](Self::peer_addr) plus a type-erased
+/// [`ext`](Self::ext) slot the connector fills with its own resolved identity
+/// (WS stuffs `ClientInfo`/`Permissions` at the HTTP upgrade; AimX stuffs its
+/// `SecurityPolicy`). Core stays connector-agnostic; each side downcasts `ext`.
+#[derive(Clone, Default)]
 #[non_exhaustive]
-pub struct PeerInfo {}
+pub struct PeerInfo {
+    /// Remote address, if the transport exposes one.
+    pub peer_addr: Option<String>,
+    /// Connector-resolved identity, type-erased so core need not know the
+    /// connector's auth types. Downcast with [`ext_as`](Self::ext_as).
+    pub ext: Option<Arc<dyn core::any::Any + Send + Sync>>,
+}
+
+impl PeerInfo {
+    /// Attach a connector-resolved identity (consumed by [`Dispatch::authenticate`]).
+    pub fn with_ext(mut self, ext: Arc<dyn core::any::Any + Send + Sync>) -> Self {
+        self.ext = Some(ext);
+        self
+    }
+
+    /// Downcast the [`ext`](Self::ext) identity to a concrete connector type.
+    pub fn ext_as<T: core::any::Any + Send + Sync>(&self) -> Option<Arc<T>> {
+        self.ext.clone()?.downcast::<T>().ok()
+    }
+}
+
+impl core::fmt::Debug for PeerInfo {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("PeerInfo")
+            .field("peer_addr", &self.peer_addr)
+            .field("ext", &self.ext.as_ref().map(|_| "<opaque>"))
+            .finish()
+    }
+}
 
 /// The authenticated session context threaded through [`Dispatch`] calls.
 ///
-/// Minimal/opaque placeholder. Auth fields are **deferred to Phase 4**
-/// (the auth-context shape gate).
-#[derive(Debug, Clone, Default)]
+/// **Phase 4 (resolved — doc 037 auth-context gate).** Carries the resolved
+/// principal as a type-erased [`ext`](Self::ext) that [`Dispatch::open`] threads
+/// into the per-connection [`Session`] for per-operation authorization
+/// (`authorize_subscribe`/`authorize_write`). AimX leaves it `None`.
+#[derive(Clone, Default)]
 #[non_exhaustive]
-pub struct SessionCtx {}
+pub struct SessionCtx {
+    /// The resolved principal, type-erased. Downcast with [`ext_as`](Self::ext_as).
+    pub ext: Option<Arc<dyn core::any::Any + Send + Sync>>,
+}
+
+impl SessionCtx {
+    /// Build a context carrying a connector-resolved principal.
+    pub fn with_ext(ext: Arc<dyn core::any::Any + Send + Sync>) -> Self {
+        Self { ext: Some(ext) }
+    }
+
+    /// Downcast the [`ext`](Self::ext) principal to a concrete connector type.
+    pub fn ext_as<T: core::any::Any + Send + Sync>(&self) -> Option<Arc<T>> {
+        self.ext.clone()?.downcast::<T>().ok()
+    }
+}
+
+impl core::fmt::Debug for SessionCtx {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("SessionCtx")
+            .field("ext", &self.ext.as_ref().map(|_| "<opaque>"))
+            .finish()
+    }
+}
 
 /// Engine-local bounds for a session (consumed by the Phase 2 engines, not by
 /// the contracts here).
@@ -222,6 +278,16 @@ pub enum Outbound<'a> {
         /// Unparsed record value.
         data: Payload,
     },
+    /// An explicit acknowledgement that a subscription opened. Emitted by
+    /// [`run_session`](super::server::run_session) only when
+    /// [`SessionConfig::acks_subscribe`](super::server::SessionConfig) is set
+    /// (WS needs it; AimX's ack is implicit, so it leaves the flag off and never
+    /// emits this). The `sub` is the subscription's routing id — the same value
+    /// that tags its [`Event`](Outbound::Event)s.
+    Subscribed {
+        /// Subscription id that was opened.
+        sub: &'a str,
+    },
     /// Keepalive response.
     Pong,
 }
@@ -314,9 +380,25 @@ pub trait Session: Send {
     /// Defaulted to [`RpcError::NotFound`] so a dispatch with no streaming
     /// surface need not implement it (doc 037 § the server-port refinement —
     /// the stream is side-neutral, so it is defaulted here for symmetry).
-    fn subscribe(&mut self, topic: &str) -> Result<BoxStream<'static, Payload>, RpcError> {
+    ///
+    /// Async (Phase 4): opening a subscription may need to *await* per-operation
+    /// authorization (e.g. WS `authorize_subscribe`); the engine awaits it.
+    fn subscribe<'a>(
+        &'a mut self,
+        topic: &'a str,
+    ) -> BoxFut<'a, Result<BoxStream<'static, Payload>, RpcError>> {
         let _ = topic;
-        Err(RpcError::NotFound)
+        Box::pin(async { Err(RpcError::NotFound) })
+    }
+
+    /// Late-join snapshot: the current value for `topic`, emitted by
+    /// [`run_session`](super::server::run_session) as an [`Outbound::Snapshot`]
+    /// right after a successful [`subscribe`](Session::subscribe) and before the
+    /// first event. Defaulted to `None` (no snapshot) — AimX inherits this; WS
+    /// overrides it from its `SnapshotProvider`.
+    fn snapshot(&mut self, topic: &str) -> Option<Payload> {
+        let _ = topic;
+        None
     }
 
     /// Fire-and-forget write: no reply. Routes through the existing
@@ -461,7 +543,10 @@ mod tests {
         ) -> BoxFut<'a, Result<Payload, RpcError>> {
             unimplemented!()
         }
-        fn subscribe(&mut self, _topic: &str) -> Result<BoxStream<'static, Payload>, RpcError> {
+        fn subscribe<'a>(
+            &'a mut self,
+            _topic: &'a str,
+        ) -> BoxFut<'a, Result<BoxStream<'static, Payload>, RpcError>> {
             unimplemented!()
         }
         fn write<'a>(