vgi_rpc/

server.rs

//! RPC server dispatch — reads requests, invokes handlers, writes responses.

use std::collections::HashMap;
use std::io::{Read, Write};
use std::sync::{Arc, Mutex};

use arrow_array::RecordBatch;
use arrow_cast::cast_with_options;
use arrow_schema::{Schema, SchemaRef};

use crate::errors::{Result, RpcError};
use crate::log::{LogLevel, LogMessage};
use crate::metadata::{
    CANCEL_KEY, LOG_EXTRA_KEY, LOG_LEVEL_KEY, LOG_MESSAGE_KEY, REQUEST_ID_KEY, REQUEST_VERSION,
    REQUEST_VERSION_KEY, RPC_METHOD_KEY, SERVER_ID_KEY,
};
#[cfg(feature = "shm")]
use crate::metadata::{SHM_SEGMENT_NAME_KEY, SHM_SEGMENT_SIZE_KEY};
#[cfg(feature = "shm")]
use crate::shm::{maybe_write_to_shm, resolve_shm_batch, ShmSegment};

/// Feature-off stand-in so dispatch signatures stay uniform.
#[cfg(not(feature = "shm"))]
pub(crate) struct ShmSegment;

/// Attach to a client-advertised SHM segment named in request metadata.
/// `track = false` since the client owns the lifecycle.
#[cfg(feature = "shm")]
fn maybe_attach_shm(req_md: &Metadata) -> Option<ShmSegment> {
    let name = req_md.get(SHM_SEGMENT_NAME_KEY)?;
    let size: usize = req_md.get(SHM_SEGMENT_SIZE_KEY)?.parse().ok()?;
    match ShmSegment::attach(name, size, false) {
        Ok(seg) => Some(seg),
        Err(e) => {
            tracing::warn!(target: "vgi_rpc.shm", "ignoring malformed SHM metadata ({e})");
            None
        }
    }
}

#[cfg(not(feature = "shm"))]
#[inline]
fn maybe_attach_shm(_req_md: &Metadata) -> Option<ShmSegment> {
    None
}
use crate::stream::{empty_schema, Emitted, OutputCollector, StreamResult, StreamStateKind};
use crate::wire::{empty_batch, md_get, Metadata, StreamReader, StreamWriter};

/// Serialize a parsed request batch back to a self-contained Arrow IPC
/// stream (one schema message + one record batch + EOS) for inclusion in
/// access-log `request_data`.
fn serialize_request_batch(batch: &RecordBatch) -> std::io::Result<Vec<u8>> {
    let mut buf = Vec::new();
    {
        let mut w = arrow_ipc::writer::StreamWriter::try_new(&mut buf, batch.schema_ref())
            .map_err(|e| std::io::Error::other(e.to_string()))?;
        w.write(batch)
            .map_err(|e| std::io::Error::other(e.to_string()))?;
        w.finish()
            .map_err(|e| std::io::Error::other(e.to_string()))?;
    }
    Ok(buf)
}

/// Lock a mutex, recovering the guard even if a previous holder
/// panicked. Handler code is arbitrary and *will* panic eventually; a
/// poisoned lock must not turn that into a process abort on the next
/// `.lock()`. The panic itself is surfaced to the client as an
/// `RpcError` by the `catch_unwind` wrappers in the dispatch path.
fn lock_ok<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
    m.lock().unwrap_or_else(|e| e.into_inner())
}

/// Invoke a handler closure, converting a panic into an `RpcError`
/// instead of unwinding through the serve loop (which on stdio/pipe
/// would kill the whole process). The panic message is intentionally
/// not echoed to the client.
pub(crate) fn call_guard<T>(f: impl FnOnce() -> T) -> Result<T> {
    std::panic::catch_unwind(std::panic::AssertUnwindSafe(f))
        .map_err(|_| RpcError::new("RuntimeError", "handler panicked"))
}

/// Context supplied to each handler invocation.
#[derive(Clone)]
pub struct CallContext {
    pub server_id: String,
    pub method: String,
    pub request_id: String,
    pub transport_metadata: Arc<Metadata>,
    /// Authentication state, or [`crate::AuthContext::anonymous`] when
    /// no authenticator is configured (e.g. pipe/unix transports).
    pub auth: crate::auth::AuthContext,
    /// HTTP request cookies (empty for pipe/unix). Name → value.
    pub cookies: std::collections::BTreeMap<String, String>,
    /// Coarse identifier of the bound transport. `None` until the
    /// framework has observed the transport (i.e. before the first
    /// [`RpcServer::notify_transport`] call).
    pub kind: Option<crate::transport::TransportKind>,
    pub(crate) log_sink: Arc<Mutex<Vec<LogMessage>>>,
    /// Per-tick input-batch custom metadata (updated each producer/exchange
    /// iteration). Carries e.g. `vgi_pushdown_filters` for dynamic filters.
    pub(crate) tick_metadata: Arc<Mutex<Metadata>>,
    /// Sticky-session bridge, installed by the HTTP transport when the
    /// server is sticky-enabled. `None` on pipe/unix/subprocess and on
    /// HTTP servers without sticky support — [`CallContext::open_session`]
    /// then raises a clear "not available on this transport" error.
    pub(crate) sticky: Option<Arc<dyn StickySink>>,
}

/// Bridge between [`CallContext`]'s sticky-session API and the HTTP
/// transport's per-worker session registry. Implemented by the HTTP layer
/// (see `crate::sticky`); the trait lives here so [`CallContext`] carries
/// no compile-time dependency on the `http` feature.
pub trait StickySink: Send + Sync {
    /// Whether the client opted in via `VGI-Session-Accept: true`.
    fn accept_opens(&self) -> bool;
    /// The live session state bound to this request, if any.
    fn current_state(&self) -> Option<Arc<dyn std::any::Any + Send + Sync>>;
    /// The opaque hex session id bound to this request, if any.
    fn current_session_id(&self) -> Option<String>;
    /// Register a session holding `state`; mints + stashes the response token.
    fn open(
        &self,
        state: Arc<dyn std::any::Any + Send + Sync>,
        ttl: Option<std::time::Duration>,
    ) -> Result<()>;
    /// Close the session bound to this request. Returns whether one was live.
    fn close(&self) -> Result<bool>;
}

impl CallContext {
    pub fn client_log(&self, level: LogLevel, message: impl Into<String>) {
        lock_ok(&self.log_sink).push(LogMessage::new(level, message));
    }

    pub fn client_log_with(&self, msg: LogMessage) {
        lock_ok(&self.log_sink).push(msg);
    }

    pub(crate) fn drain_logs(&self) -> Vec<LogMessage> {
        std::mem::take(&mut *lock_ok(&self.log_sink))
    }

    /// Per-tick input-batch custom metadata value (e.g. `vgi_pushdown_filters`),
    /// set by the producer/exchange loop for the current iteration.
    pub fn tick_metadata(&self, key: &str) -> Option<String> {
        lock_ok(&self.tick_metadata).get(key).cloned()
    }

    /// Build a call context for `server` serving `req`. Defaults to
    /// anonymous auth with no cookies — callers on authenticated
    /// transports (HTTP) override the two after construction or use
    /// [`CallContext::with_auth_cookies`].
    pub(crate) fn for_request(server: &RpcServer, req: &Request) -> Self {
        Self {
            server_id: server.server_id.clone(),
            method: req.method.clone(),
            request_id: req.request_id.clone(),
            transport_metadata: Arc::new(req.metadata.clone()),
            auth: crate::auth::AuthContext::anonymous(),
            cookies: std::collections::BTreeMap::new(),
            kind: server.transport_kind(),
            log_sink: Arc::new(Mutex::new(Vec::new())),
            tick_metadata: Arc::new(Mutex::new(Metadata::default())),
            sticky: None,
        }
    }

    /// Build a call context with an explicit auth context + cookie map.
    /// Only the HTTP transport constructs contexts this way; gated so the
    /// method isn't dead code (a `-D warnings` build failure) when `vgi-rpc`
    /// is compiled without the `http` feature (e.g. from `vgi-rpc-client`).
    #[cfg(feature = "http")]
    pub(crate) fn with_auth_cookies(
        server: &RpcServer,
        req: &Request,
        auth: crate::auth::AuthContext,
        cookies: std::collections::BTreeMap<String, String>,
    ) -> Self {
        Self {
            server_id: server.server_id.clone(),
            method: req.method.clone(),
            request_id: req.request_id.clone(),
            transport_metadata: Arc::new(req.metadata.clone()),
            auth,
            cookies,
            kind: server.transport_kind(),
            log_sink: Arc::new(Mutex::new(Vec::new())),
            tick_metadata: Arc::new(Mutex::new(Metadata::default())),
            sticky: None,
        }
    }

    /// Attach a sticky-session sink (HTTP transport only). No-op semantics
    /// for callers: the session API simply reports "not available" when
    /// this is never set. HTTP-only, so gated to avoid a dead-code
    /// `-D warnings` failure in non-`http` builds.
    #[cfg(feature = "http")]
    pub(crate) fn set_sticky(&mut self, sink: Arc<dyn StickySink>) {
        self.sticky = Some(sink);
    }

    // --- Sticky sessions (HTTP-only) -----------------------------------

    /// The live session state object, downcast to `T`, or `None` when no
    /// session is bound to this request (or it is not a `T`).
    ///
    /// Sticky sessions are HTTP-only; on other transports this is always
    /// `None`. Mirrors Python's `ctx.session`.
    pub fn session<T: std::any::Any + Send + Sync>(&self) -> Option<Arc<T>> {
        let state = self.sticky.as_ref()?.current_state()?;
        state.downcast::<T>().ok()
    }

    /// The opaque hex session id bound to this request, or `None`.
    /// Survives [`CallContext::close_session`] within the same request.
    pub fn session_id(&self) -> Option<String> {
        self.sticky.as_ref()?.current_session_id()
    }

    /// Register a sticky session holding `state` for subsequent requests.
    ///
    /// The framework mints a signed `VGI-Session` token and attaches it to
    /// the response; a client inside a `with_session_token()` block echoes
    /// it on subsequent requests, and the framework restores `state` as
    /// [`CallContext::session`]. `ttl` overrides the server default.
    ///
    /// Mirrors Python's `ctx.open_session`. Errors when sticky is
    /// unavailable on this transport, the client did not opt in, or a
    /// session is already bound to this request.
    pub fn open_session(
        &self,
        state: Arc<dyn std::any::Any + Send + Sync>,
        ttl: Option<std::time::Duration>,
    ) -> Result<()> {
        let sink = self.sticky.as_ref().ok_or_else(|| {
            RpcError::runtime_error("sticky sessions not available on this transport")
        })?;
        if !sink.accept_opens() {
            return Err(RpcError::runtime_error(
                "client did not opt in to sticky sessions \
                 (missing VGI-Session-Accept: true header — open the call inside \
                 an HttpConnection.with_session_token() block)",
            ));
        }
        if sink.current_state().is_some() {
            return Err(RpcError::runtime_error(
                "a sticky session is already active for this request",
            ));
        }
        sink.open(state, ttl)
    }

    /// Invalidate the sticky session bound to this request. Idempotent;
    /// mirrors Python's `ctx.close_session`.
    pub fn close_session(&self) -> Result<()> {
        let sink = self.sticky.as_ref().ok_or_else(|| {
            RpcError::runtime_error("sticky sessions not available on this transport")
        })?;
        sink.close()?;
        Ok(())
    }
}

/// A request batch parsed from the wire.
pub struct Request {
    pub method: String,
    pub request_id: String,
    pub batch: RecordBatch,
    pub metadata: Metadata,
}

impl Request {
    pub fn column(&self, name: &str) -> Option<&dyn arrow_array::Array> {
        let idx = self.batch.schema().index_of(name).ok()?;
        Some(self.batch.column(idx).as_ref())
    }

    /// Build a `Request` from a record batch carrying its own
    /// `custom_metadata`, validating the `vgi_rpc.method` /
    /// `vgi_rpc.request_version` metadata.
    ///
    /// `require_method` controls whether a missing `vgi_rpc.method` key is
    /// an error (pipe/unix transports require it; HTTP already derives the
    /// method from the URL path and may leave the key absent).
    pub(crate) fn from_read_batch(
        batch: RecordBatch,
        metadata: Metadata,
        require_method: bool,
    ) -> Result<Self> {
        let method = if require_method {
            md_get(&metadata, RPC_METHOD_KEY)
                .ok_or_else(|| {
                    RpcError::protocol_error(
                        "Missing 'vgi_rpc.method' in request batch custom_metadata.",
                    )
                })?
                .to_string()
        } else {
            md_get(&metadata, RPC_METHOD_KEY).unwrap_or("").to_string()
        };
        let version = md_get(&metadata, REQUEST_VERSION_KEY).ok_or_else(|| {
            RpcError::version_error(format!(
                "Missing 'vgi_rpc.request_version' in request batch custom_metadata. Set it to {:?}.",
                REQUEST_VERSION
            ))
        })?;
        if version != REQUEST_VERSION {
            return Err(RpcError::version_error(format!(
                "Unsupported request version {:?}, expected {:?}.",
                version, REQUEST_VERSION
            )));
        }
        if require_method && !batch.schema().fields().is_empty() && batch.num_rows() != 1 {
            return Err(RpcError::protocol_error(format!(
                "Expected 1 row in request batch, got {}",
                batch.num_rows()
            )));
        }
        let request_id = md_get(&metadata, REQUEST_ID_KEY).unwrap_or("").to_string();
        Ok(Request {
            method,
            request_id,
            batch,
            metadata,
        })
    }
}

/// Identifies the dispatch kind of a registered method.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MethodType {
    Unary,
    Producer,
    Exchange,
    /// State kind determined at runtime by handler return value.
    Dynamic,
}

/// A handler function for a unary RPC method.
pub type UnaryHandler =
    Arc<dyn Fn(&Request, &CallContext) -> Result<Option<RecordBatch>> + Send + Sync>;

/// A handler function for a streaming RPC method.
pub type StreamHandler = Arc<dyn Fn(&Request, &CallContext) -> Result<StreamResult> + Send + Sync>;

/// Fluent builder for [`RpcServer`] with describe/identity/version knobs.
#[derive(Default)]
pub struct RpcServerBuilder {
    server_id: Option<String>,
    server_version: Option<String>,
    protocol_name: Option<String>,
    protocol_version: Option<String>,
    enable_describe: bool,
    dispatch_hook: Option<Arc<dyn crate::hooks::DispatchHook>>,
    on_serve_start: Option<crate::transport::ServeStartHook>,
    #[cfg(feature = "http")]
    external_config: Option<Arc<crate::external::ExternalLocationConfig>>,
}

impl RpcServerBuilder {
    pub fn server_id(mut self, id: impl Into<String>) -> Self {
        self.server_id = Some(id.into());
        self
    }

    pub fn server_version(mut self, v: impl Into<String>) -> Self {
        self.server_version = Some(v.into());
        self
    }

    pub fn protocol_name(mut self, name: impl Into<String>) -> Self {
        self.protocol_name = Some(name.into());
        self
    }

    /// Operator-supplied free-form protocol-contract version label, reported
    /// in access-log records as ``protocol_version``. Complementary to
    /// (build) ``server_version``.
    pub fn protocol_version(mut self, v: impl Into<String>) -> Self {
        self.protocol_version = Some(v.into());
        self
    }

    pub fn enable_describe(mut self, enabled: bool) -> Self {
        self.enable_describe = enabled;
        self
    }

    pub fn with_hook(mut self, hook: Arc<dyn crate::hooks::DispatchHook>) -> Self {
        self.dispatch_hook = Some(hook);
        self
    }

    /// Register a one-shot lifecycle hook fired before the first
    /// request is dispatched on each (kind, capabilities) combination.
    /// Mirrors Python's `on_serve_start` duck-typed protocol.
    ///
    /// The hook runs synchronously on the thread that first observes
    /// the transport binding. Subsequent calls to
    /// [`RpcServer::notify_transport`] with the same `(kind, caps)`
    /// are no-ops; calls with a different combination re-fire the hook
    /// (matches Python's behaviour for test paths that re-bind).
    pub fn on_serve_start(mut self, hook: crate::transport::ServeStartHook) -> Self {
        self.on_serve_start = Some(hook);
        self
    }

    /// Enable automatic externalization of large unary results and stream
    /// output batches. Feature-gated on `http` (where the compression +
    /// fetcher deps already live).
    #[cfg(feature = "http")]
    pub fn with_external_location(mut self, cfg: crate::external::ExternalLocationConfig) -> Self {
        self.external_config = Some(Arc::new(cfg));
        self
    }

    pub fn build(self) -> RpcServer {
        RpcServer {
            methods: HashMap::new(),
            server_id: self.server_id.unwrap_or_else(crate::util::short_random_id),
            server_version: self.server_version.unwrap_or_default(),
            protocol_name: self.protocol_name.unwrap_or_default(),
            protocol_version: self.protocol_version.unwrap_or_default(),
            protocol_hash: std::sync::OnceLock::new(),
            describe_enabled: self.enable_describe,
            dispatch_hook: self.dispatch_hook,
            on_serve_start: self.on_serve_start,
            transport_state: Mutex::new(None),
            #[cfg(feature = "http")]
            external_config: self.external_config,
        }
    }
}

/// Describes one RPC method — the metadata required both for dispatch and
/// introspection via `__describe__`.
///
/// Build via [`MethodInfo::unary`] / [`MethodInfo::stream`] and attach
/// additional describe-time metadata through the builder helpers
/// (`.doc`, `.param_type`, `.param_default`, `.param_doc`, `.header_schema`).
pub struct MethodInfo {
    pub name: String,
    pub method_type: MethodType,
    /// Schema of the request parameters (one row).
    pub params_schema: SchemaRef,
    /// Schema of the unary result; empty for streams.
    pub result_schema: SchemaRef,
    /// For streams that emit a typed header, the header batch schema.
    pub header_schema: Option<SchemaRef>,
    /// Method-level docstring (the first line of Python's docstring).
    pub doc: Option<String>,
    /// Parameter type names in source order, matching the Python describe
    /// wire format ("str", "int", "list[str]", "Point", "str | None").
    pub param_types: Vec<(String, String)>,
    /// Parameter defaults; values are anything JSON-serializable.
    pub param_defaults: Vec<(String, serde_json::Value)>,
    /// Per-parameter documentation (matches the Python `param_docs_json`).
    pub param_docs: Vec<(String, String)>,
    /// Whether the method has a non-void return. `false` for streams/void.
    pub has_return: bool,
    pub unary: Option<UnaryHandler>,
    pub stream: Option<StreamHandler>,
    /// Decoder that reconstructs the method's `StreamStateKind` from a byte
    /// slice produced by `ProducerState::encode_state` /
    /// `ExchangeState::encode_state`. Required for HTTP streaming (the
    /// stateless-worker model); `None` for unary methods and for streams
    /// that will only ever run over pipe/unix.
    pub state_decoder: Option<StateDecoder>,
}

/// Decoder that reconstructs a concrete streaming state from its
/// serialized bytes, used by the HTTP transport on continuation requests.
pub type StateDecoder = Arc<dyn Fn(&[u8]) -> Result<crate::stream::StreamStateKind> + Send + Sync>;

impl MethodInfo {
    /// Start building a unary method registration.
    pub fn unary(
        name: impl Into<String>,
        params_schema: SchemaRef,
        result_schema: SchemaRef,
        handler: impl Fn(&Request, &CallContext) -> Result<Option<RecordBatch>> + Send + Sync + 'static,
    ) -> Self {
        let has_return = !result_schema.fields().is_empty();
        Self {
            name: name.into(),
            method_type: MethodType::Unary,
            params_schema,
            result_schema,
            header_schema: None,
            doc: None,
            param_types: Vec::new(),
            param_defaults: Vec::new(),
            param_docs: Vec::new(),
            has_return,
            unary: Some(Arc::new(handler)),
            stream: None,
            state_decoder: None,
        }
    }

    /// Start building a streaming method registration.
    ///
    /// **Note:** this form registers the method without a state decoder,
    /// so it will work for pipe/unix transports but HTTP continuation
    /// requests will fail. Use
    /// [`MethodInfo::producer_with_codec`] /
    /// [`MethodInfo::exchange_with_codec`] when HTTP is enabled.
    pub fn stream(
        name: impl Into<String>,
        method_type: MethodType,
        params_schema: SchemaRef,
        handler: impl Fn(&Request, &CallContext) -> Result<StreamResult> + Send + Sync + 'static,
    ) -> Self {
        assert!(
            matches!(
                method_type,
                MethodType::Producer | MethodType::Exchange | MethodType::Dynamic
            ),
            "stream methods must be Producer / Exchange / Dynamic"
        );
        Self {
            name: name.into(),
            method_type,
            params_schema,
            result_schema: empty_schema(),
            header_schema: None,
            doc: None,
            param_types: Vec::new(),
            param_defaults: Vec::new(),
            param_docs: Vec::new(),
            has_return: false,
            unary: None,
            stream: Some(Arc::new(handler)),
            state_decoder: None,
        }
    }

    /// Attach a state decoder function. See [`StateDecoder`].
    pub fn with_state_decoder(mut self, decoder: StateDecoder) -> Self {
        self.state_decoder = Some(decoder);
        self
    }

    pub fn doc(mut self, s: impl Into<String>) -> Self {
        self.doc = Some(s.into());
        self
    }

    pub fn param_type(mut self, param: impl Into<String>, ty: impl Into<String>) -> Self {
        self.param_types.push((param.into(), ty.into()));
        self
    }

    pub fn param_default(mut self, param: impl Into<String>, value: serde_json::Value) -> Self {
        self.param_defaults.push((param.into(), value));
        self
    }

    pub fn param_doc(mut self, param: impl Into<String>, doc: impl Into<String>) -> Self {
        self.param_docs.push((param.into(), doc.into()));
        self
    }

    pub fn header_schema(mut self, schema: SchemaRef) -> Self {
        self.header_schema = Some(schema);
        self
    }
}

/// The RPC server — holds method registrations and dispatches requests.
pub struct RpcServer {
    methods: HashMap<String, MethodInfo>,
    pub server_id: String,
    pub(crate) server_version: String,
    pub(crate) protocol_name: String,
    pub(crate) protocol_version: String,
    pub(crate) protocol_hash: std::sync::OnceLock<String>,
    pub(crate) describe_enabled: bool,
    pub(crate) dispatch_hook: Option<Arc<dyn crate::hooks::DispatchHook>>,
    /// Optional one-shot lifecycle hook fired on the first
    /// [`notify_transport`](Self::notify_transport) per (kind, caps).
    on_serve_start: Option<crate::transport::ServeStartHook>,
    /// Coarse identifier of the bound transport, populated by
    /// [`notify_transport`](Self::notify_transport).
    transport_state: Mutex<
        Option<(
            crate::transport::TransportKind,
            crate::transport::TransportCapabilities,
        )>,
    >,
    #[cfg(feature = "http")]
    pub(crate) external_config: Option<Arc<crate::external::ExternalLocationConfig>>,
}

impl RpcServer {
    /// Create a new `RpcServer`. For richer configuration, use [`RpcServer::builder`].
    pub fn new(server_id: impl Into<String>) -> Self {
        Self::builder().server_id(server_id).build()
    }

    /// Create a new builder.
    pub fn builder() -> RpcServerBuilder {
        RpcServerBuilder::default()
    }

    pub fn protocol_name(&self) -> &str {
        &self.protocol_name
    }

    pub fn describe_enabled(&self) -> bool {
        self.describe_enabled
    }

    pub fn server_version(&self) -> &str {
        &self.server_version
    }

    pub fn protocol_version(&self) -> &str {
        &self.protocol_version
    }

    /// SHA-256 hex digest of the canonical __describe__ payload. Computed
    /// lazily on first call and cached.
    pub fn protocol_hash(&self) -> &str {
        self.protocol_hash.get_or_init(|| {
            match crate::introspect::build_describe(
                &self.protocol_name,
                &self.methods,
                &self.server_id,
                &self.protocol_version,
            ) {
                Ok((_, md)) => md
                    .get(crate::metadata::PROTOCOL_HASH_KEY)
                    .cloned()
                    .unwrap_or_default(),
                Err(_) => String::new(),
            }
        })
    }

    #[cfg(feature = "http")]
    pub fn external_config(&self) -> Option<&Arc<crate::external::ExternalLocationConfig>> {
        self.external_config.as_ref()
    }

    /// Currently-bound [`TransportKind`](crate::transport::TransportKind),
    /// or `None` before the framework has observed a transport. Set by
    /// [`notify_transport`](Self::notify_transport).
    pub fn transport_kind(&self) -> Option<crate::transport::TransportKind> {
        lock_ok(&self.transport_state).as_ref().map(|(k, _)| *k)
    }

    /// Currently-advertised [`TransportCapabilities`](crate::transport::TransportCapabilities).
    /// Empty (all-false) before a transport is bound and for transports
    /// without extra capabilities.
    pub fn transport_capabilities(&self) -> crate::transport::TransportCapabilities {
        lock_ok(&self.transport_state)
            .as_ref()
            .map(|(_, c)| *c)
            .unwrap_or_default()
    }

    /// Bind the server to a transport, firing `on_serve_start` once per
    /// `(kind, caps)` combination. Subsequent calls with the same
    /// combination are cheap no-ops (the common case where transport
    /// glue invokes this on every request). A different combination
    /// updates the bound state and re-fires the hook — matches the
    /// Python `_notify_transport` contract.
    ///
    /// Call this from each transport entry point:
    /// - stdio / pipe `main`: once before [`serve`](Self::serve)
    /// - Unix accept loop: once per process
    /// - HTTP request handler: every request (idempotent)
    pub fn notify_transport(
        &self,
        kind: crate::transport::TransportKind,
        caps: crate::transport::TransportCapabilities,
    ) {
        let hook = {
            let mut guard = lock_ok(&self.transport_state);
            if let Some((cur_kind, cur_caps)) = guard.as_ref() {
                if *cur_kind == kind && *cur_caps == caps {
                    return;
                }
            }
            *guard = Some((kind, caps));
            self.on_serve_start.clone()
        };
        if let Some(h) = hook {
            h(kind, &caps);
        }
    }

    /// Register a method described by a [`MethodInfo`].
    pub fn register(&mut self, info: MethodInfo) {
        self.methods.insert(info.name.clone(), info);
    }

    /// Convenience wrapper for the old positional API — equivalent to
    /// `register(MethodInfo::unary(name, empty_schema(), result_schema, handler))`.
    /// Prefer [`MethodInfo::unary`] + [`RpcServer::register`] for new code.
    pub fn register_unary(
        &mut self,
        name: impl Into<String>,
        result_schema: SchemaRef,
        handler: impl Fn(&Request, &CallContext) -> Result<Option<RecordBatch>> + Send + Sync + 'static,
    ) {
        self.register(MethodInfo::unary(
            name,
            empty_schema(),
            result_schema,
            handler,
        ));
    }

    /// Convenience wrapper for the old positional API — equivalent to
    /// `register(MethodInfo::stream(name, method_type, empty_schema(), handler))`.
    /// Prefer [`MethodInfo::stream`] + [`RpcServer::register`] for new code.
    pub fn register_stream(
        &mut self,
        name: impl Into<String>,
        method_type: MethodType,
        handler: impl Fn(&Request, &CallContext) -> Result<StreamResult> + Send + Sync + 'static,
    ) {
        self.register(MethodInfo::stream(
            name,
            method_type,
            empty_schema(),
            handler,
        ));
    }

    pub fn method(&self, name: &str) -> Option<&MethodInfo> {
        self.methods.get(name)
    }

    pub fn methods(&self) -> &HashMap<String, MethodInfo> {
        &self.methods
    }

    pub fn method_names(&self) -> Vec<&str> {
        self.sorted_method_names()
    }

    /// Method names sorted alphabetically. Preferred over `methods().keys()`
    /// when order matters (introspection / describe / HTML rendering).
    pub fn sorted_method_names(&self) -> Vec<&str> {
        let mut names: Vec<_> = self.methods.keys().map(String::as_str).collect();
        names.sort();
        names
    }

    /// Run the serve loop over a single reader/writer pair (pipe or socket).
    ///
    /// Reads are **blocking with no timeout** — a peer that opens the
    /// connection and then stalls pins this thread until it sends data,
    /// EOFs, or resets. stdio/pipe has no timeout API, so that transport
    /// is trusted-peer-only (see also the SHM module docs). On a socket
    /// transport, the caller owns the stream and **should** set a read
    /// timeout (e.g. `UnixStream::set_read_timeout`) before handing it
    /// here; a `TimedOut`/`WouldBlock` error then cleanly ends the
    /// connection via the error path below.
    pub fn serve<R: Read, W: Write>(&self, mut r: R, mut w: W) {
        loop {
            match self.serve_one(&mut r, &mut w) {
                Ok(keep_going) => {
                    if !keep_going {
                        return;
                    }
                }
                Err(e) => {
                    // A frame-level error (malformed request, IO error,
                    // peer reset) ends the connection. Log it so a
                    // daemonized listener has diagnostics — silently
                    // returning made transient and hostile-input
                    // failures indistinguishable from a clean EOF.
                    tracing::warn!(
                        target: "vgi_rpc.server",
                        error = %e,
                        "serve loop terminating connection on error"
                    );
                    return;
                }
            }
        }
    }

    /// Like [`serve`], but checks `shutdown` between requests and exits
    /// cleanly when it returns `true`. Useful for daemonized pipe/unix
    /// listeners that want to drain the in-flight request before exiting
    /// on SIGTERM. Blocking reads still must terminate via EOF/peer-close
    /// — this is an *advisory* signal checked at request boundaries.
    pub fn serve_with_shutdown<R, W, F>(&self, mut r: R, mut w: W, shutdown: F)
    where
        R: Read,
        W: Write,
        F: Fn() -> bool,
    {
        loop {
            if shutdown() {
                return;
            }
            match self.serve_one(&mut r, &mut w) {
                Ok(true) => {}
                _ => return,
            }
        }
    }

    /// Handle one request. Returns `Ok(true)` to continue, `Ok(false)` on EOS/EOF.
    pub fn serve_one<R: Read, W: Write>(&self, r: &mut R, w: &mut W) -> Result<bool> {
        let result = self._serve_one(r, w);
        let _ = w.flush();
        result
    }

    fn _serve_one<R: Read, W: Write>(&self, r: &mut R, w: &mut W) -> Result<bool> {
        let req = match self.read_request(r)? {
            Some(rq) => rq,
            None => return Ok(false),
        };

        // __transport_options__ — framework transport-capability handshake,
        // handled before method dispatch (not a registered method, so it never
        // appears in `methods` / `__describe__`, and doesn't perturb the
        // protocol hash). Capabilities ride as response metadata; the response
        // batch is empty. Always available, including to version-mismatched
        // clients, since it is the negotiation they perform before `init`.
        if req.method == crate::transport_options::TRANSPORT_OPTIONS_METHOD_NAME {
            let mut md = crate::transport_options::worker_transport_metadata();
            md.insert(REQUEST_VERSION_KEY.to_string(), REQUEST_VERSION.to_string());
            md.insert(SERVER_ID_KEY.to_string(), self.server_id.clone());
            let schema = empty_schema();
            let batch = empty_batch(&schema)?;
            let mut sw = StreamWriter::new(w, &schema)?;
            sw.write(&batch, Some(&md))?;
            sw.finish()?;
            return Ok(true);
        }

        // Enforce application protocol-version compatibility: the client sends
        // its `vgi_rpc.protocol_version`; if its MAJOR differs from the
        // server's enforced version, reject (mirrors the Python framework).
        if !self.protocol_version.is_empty() {
            if let Some(client_v) = md_get(&req.metadata, crate::metadata::PROTOCOL_VERSION_KEY) {
                let major = |v: &str| v.split('.').next().unwrap_or("").to_string();
                if major(client_v) != major(&self.protocol_version) {
                    let err = RpcError::version_error(format!(
                        "protocol_version mismatch: client {:?} is incompatible with server {:?}",
                        client_v, self.protocol_version
                    ));
                    write_error_stream(w, &empty_schema(), &err, &self.server_id, &req.request_id)?;
                    return Ok(true);
                }
            }
        }

        let ctx = CallContext::for_request(self, &req);

        let stats = Arc::new(Mutex::new(crate::hooks::CallStatistics::default()));
        // Record the unary request batch as input stats (one row).
        {
            let mut s = lock_ok(&stats);
            s.input_batches = 1;
            s.input_rows = req.batch.num_rows() as u64;
        }

        // Built-in __describe__ introspection.
        if self.describe_enabled && req.method == crate::introspect::DESCRIBE_METHOD_NAME {
            match crate::introspect::build_describe(
                &self.protocol_name,
                &self.methods,
                &self.server_id,
                &self.protocol_version,
            ) {
                Ok((batch, md)) => {
                    crate::introspect::write_describe_response(w, &batch, &md)?;
                }
                Err(err) => {
                    write_error_stream(w, &empty_schema(), &err, &self.server_id, &req.request_id)?;
                }
            }
            return Ok(true);
        }

        let Some(info) = self.methods.get(&req.method) else {
            let names = self.sorted_method_names();
            let msg = format!(
                "Unknown method: '{}'. Available methods: {:?}",
                req.method, names
            );
            write_error_stream(
                w,
                &empty_schema(),
                &RpcError::attribute_error(msg),
                &self.server_id,
                &req.request_id,
            )?;
            return Ok(true);
        };

        let method_type = match info.method_type {
            MethodType::Unary => "unary",
            _ => "stream",
        };
        let mut dispatch_info =
            crate::hooks::DispatchInfo::from_request(self, &req, method_type, &ctx.auth);
        // Best-effort capture of self-contained Arrow IPC bytes of the
        // request batch for access-log `request_data`. Failures here must
        // not abort dispatch — observability is non-essential.
        if let Ok(bytes) = serialize_request_batch(&req.batch) {
            dispatch_info.request_data = bytes;
        }
        if method_type == "stream" {
            dispatch_info.stream_id = crate::access_log::random_stream_id();
        }
        let hook_token = self
            .dispatch_hook
            .as_ref()
            .map(|h| h.on_dispatch_start(&dispatch_info));

        let mut app_err: Option<RpcError> = None;
        let shm = maybe_attach_shm(&req.metadata);
        let shm_ref = shm.as_ref();
        match info.method_type {
            MethodType::Unary => {
                self.serve_unary(w, &req, info, &ctx, &stats, &mut app_err, shm_ref)?
            }
            MethodType::Producer | MethodType::Exchange | MethodType::Dynamic => {
                self.serve_stream(r, w, &req, info, &ctx, &stats, &mut app_err, shm_ref)?
            }
        }
        // `shm` (if any) is dropped here, releasing our mmap of the
        // client-owned segment without unlinking it.
        let _ = shm;

        if let Some(hook) = self.dispatch_hook.as_ref() {
            let token = hook_token.unwrap_or(0);
            let final_stats = lock_ok(&stats).clone();
            hook.on_dispatch_end(token, &dispatch_info, app_err.as_ref(), &final_stats);
        }
        Ok(true)
    }

    fn read_request<R: Read>(&self, r: &mut R) -> Result<Option<Request>> {
        let mut reader = match StreamReader::new(r) {
            Ok(r) => r,
            Err(e) => {
                // EOF at request boundary is normal
                let msg = e.message.to_lowercase();
                if msg.contains("empty ipc stream") || msg.contains("eof") {
                    return Ok(None);
                }
                return Err(e);
            }
        };
        let (batch, metadata) = match reader.read_next()? {
            Some(b) => b,
            None => return Ok(None),
        };
        reader.drain()?;
        Ok(Some(Request::from_read_batch(batch, metadata, true)?))
    }

    #[allow(clippy::too_many_arguments)]
    fn serve_unary<W: Write>(
        &self,
        w: &mut W,
        req: &Request,
        info: &MethodInfo,
        ctx: &CallContext,
        stats: &Arc<Mutex<crate::hooks::CallStatistics>>,
        app_err: &mut Option<RpcError>,
        #[cfg_attr(not(feature = "shm"), allow(unused_variables))] shm: Option<&ShmSegment>,
    ) -> Result<()> {
        // A panic in handler code is converted to an `RpcError` and
        // flows into the error-envelope path below, rather than
        // unwinding through the serve loop.
        let result = call_guard(|| (info.unary.as_ref().unwrap())(req, ctx)).and_then(|r| r);
        let logs = ctx.drain_logs();
        match result {
            Ok(maybe_batch) => {
                let mut sw = StreamWriter::new(w, &info.result_schema)?;
                for log in logs {
                    let md = build_log_metadata(&log, &self.server_id, &req.request_id);
                    sw.write(&empty_batch(&info.result_schema)?, Some(&md))?;
                }
                let out_batch = match maybe_batch {
                    Some(b) => b,
                    None => empty_batch(&info.result_schema)?,
                };
                {
                    let mut s = lock_ok(stats);
                    s.output_batches = 1;
                    s.output_rows = out_batch.num_rows() as u64;
                }
                #[cfg(feature = "shm")]
                if let Some(seg) = shm {
                    let (written, written_md) =
                        maybe_write_to_shm(out_batch.clone(), Metadata::new(), Some(seg))?;
                    if written_md.contains_key(crate::metadata::SHM_OFFSET_KEY) {
                        sw.write(&written, Some(&written_md))?;
                        sw.finish()?;
                        return Ok(());
                    }
                }
                #[cfg(feature = "http")]
                if let Some(cfg) = self.external_config.as_ref() {
                    if let Ok(Some((ptr, md))) =
                        crate::external::maybe_externalize_batch(&out_batch, None, cfg)
                    {
                        sw.write(&ptr, Some(&md))?;
                        sw.finish()?;
                        return Ok(());
                    }
                }
                #[cfg(not(feature = "shm"))]
                let _ = shm;
                sw.write(&out_batch, None)?;
                sw.finish()?;
            }
            Err(err) => {
                let mut sw = StreamWriter::new(w, &info.result_schema)?;
                for log in logs {
                    let md = build_log_metadata(&log, &self.server_id, &req.request_id);
                    sw.write(&empty_batch(&info.result_schema)?, Some(&md))?;
                }
                let md = build_error_metadata(&err, &self.server_id, &req.request_id);
                sw.write(&empty_batch(&info.result_schema)?, Some(&md))?;
                sw.finish()?;
                *app_err = Some(err);
            }
        }
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    #[allow(clippy::too_many_arguments)]
    fn serve_stream<R: Read, W: Write>(
        &self,
        r: &mut R,
        w: &mut W,
        req: &Request,
        info: &MethodInfo,
        ctx: &CallContext,
        stats: &Arc<Mutex<crate::hooks::CallStatistics>>,
        app_err: &mut Option<RpcError>,
        #[cfg_attr(not(feature = "shm"), allow(unused_variables))] shm: Option<&ShmSegment>,
    ) -> Result<()> {
        let init_result = call_guard(|| (info.stream.as_ref().unwrap())(req, ctx)).and_then(|r| r);
        let init_logs = ctx.drain_logs();
        let stream = match init_result {
            Ok(s) => s,
            Err(err) => {
                // Init error: write as unary-style error stream.
                let output_schema = info.result_schema.clone();
                let mut sw = StreamWriter::new(w, &output_schema)?;
                for log in init_logs {
                    let md = build_log_metadata(&log, &self.server_id, &req.request_id);
                    sw.write(&empty_batch(&output_schema)?, Some(&md))?;
                }
                let md = build_error_metadata(&err, &self.server_id, &req.request_id);
                sw.write(&empty_batch(&output_schema)?, Some(&md))?;
                sw.finish()?;
                // Drain any client input (ticks / exchange batches) so the transport
                // is clean for the next request.
                let _ = drain_input(r);
                *app_err = Some(err);
                return Ok(());
            }
        };

        let StreamResult {
            output_schema,
            input_schema,
            state,
            header,
            header_metadata,
        } = stream;

        // Write header as its own IPC stream if present.
        let wrote_header = header.is_some();
        if let Some(header_batch) = header {
            let mut hw = StreamWriter::new(&mut *w, header_batch.schema().as_ref())?;
            for log in &init_logs {
                let md = build_log_metadata(log, &self.server_id, &req.request_id);
                hw.write(&empty_batch(header_batch.schema().as_ref())?, Some(&md))?;
            }
            hw.write(&header_batch, header_metadata.as_ref())?;
            hw.finish()?;
        }
        let _ = w.flush();

        // Open the output stream first — the client opens the output reader
        // before the next tick is read back here, so we must make the schema
        // available without waiting on input.
        let mut out_writer = StreamWriter::new(&mut *w, output_schema.as_ref())?;
        out_writer.flush()?;

        // Open the input stream (ticks for producer, real batches for exchange).
        let mut input_reader = StreamReader::new(&mut *r)?;

        // If we didn't already write init logs into a header stream, write them now.
        if !wrote_header {
            for log in &init_logs {
                let md = build_log_metadata(log, &self.server_id, &req.request_id);
                out_writer.write(&empty_batch(output_schema.as_ref())?, Some(&md))?;
            }
        }
        let _ = header_metadata;

        let mut state = state;
        let mut cancelled = false;

        'lockstep: loop {
            let read = match input_reader.read_next() {
                Ok(x) => x,
                Err(_) => break,
            };
            let Some((input_batch, input_md)) = read else {
                break;
            };

            // Resolve SHM pointer batches before anything else — the
            // schema cast / cancel check / handler all expect the real
            // batch. Free the region as soon as it's been deserialized
            // (we copy on read, so no live borrow remains).
            #[cfg(feature = "shm")]
            let (input_batch, input_md) = {
                let resolved = resolve_shm_batch(input_batch, input_md, shm)?;
                if let (Some(off), Some(seg)) = (resolved.release_offset, shm) {
                    let _ = seg.free(off);
                }
                (resolved.batch, resolved.metadata)
            };

            {
                let mut s = lock_ok(stats);
                s.input_batches += 1;
                s.input_rows += input_batch.num_rows() as u64;
            }

            // Surface this tick's input metadata (e.g. dynamic pushdown
            // filters) to the producer/exchange handler via the context.
            *lock_ok(&ctx.tick_metadata) = input_md.clone();

            // Cancellation signal.
            if md_get(&input_md, CANCEL_KEY).is_some() {
                cancelled = true;
                match &mut state {
                    StreamStateKind::Producer(p) => p.on_cancel(ctx),
                    StreamStateKind::Exchange(e) => e.on_cancel(ctx),
                }
                break;
            }

            // Cast input schema to expected schema when required.
            let casted = match &input_schema {
                Some(expected) if input_batch.schema() != *expected => {
                    match cast_batch(&input_batch, expected) {
                        Ok(b) => b,
                        Err(e) => {
                            let md = build_error_metadata(&e, &self.server_id, &req.request_id);
                            out_writer.write(&empty_batch(output_schema.as_ref())?, Some(&md))?;
                            break 'lockstep;
                        }
                    }
                }
                _ => input_batch,
            };

            let mut out = OutputCollector::new(output_schema.clone(), input_schema.is_none());

            let iter_result = call_guard(|| match &mut state {
                StreamStateKind::Producer(p) => p.produce(&mut out, ctx),
                StreamStateKind::Exchange(e) => e.exchange(&casted, &mut out, ctx),
            })
            .and_then(|r| r);

            // Flush any iteration-level logs first (logs appended during produce/exchange).
            let iter_logs = ctx.drain_logs();
            for log in iter_logs {
                let md = build_log_metadata(&log, &self.server_id, &req.request_id);
                out_writer.write(&empty_batch(output_schema.as_ref())?, Some(&md))?;
            }

            if let Err(err) = iter_result {
                let md = build_error_metadata(&err, &self.server_id, &req.request_id);
                out_writer.write(&empty_batch(output_schema.as_ref())?, Some(&md))?;
                *app_err = Some(err);
                break;
            }

            let finished = out.finished();

            // Flush collected emitted items (logs added via OutputCollector, then batches).
            for item in out.items.drain(..) {
                match item {
                    Emitted::Log(log) => {
                        let md = build_log_metadata(&log, &self.server_id, &req.request_id);
                        out_writer.write(&empty_batch(output_schema.as_ref())?, Some(&md))?;
                    }
                    Emitted::Batch { batch, metadata } => {
                        {
                            let mut s = lock_ok(stats);
                            s.output_batches += 1;
                            s.output_rows += batch.num_rows() as u64;
                        }
                        #[cfg(feature = "shm")]
                        if let Some(seg) = shm {
                            let md_in = metadata.clone().unwrap_or_default();
                            let (written, written_md) =
                                maybe_write_to_shm(batch.clone(), md_in, Some(seg))?;
                            if written_md.contains_key(crate::metadata::SHM_OFFSET_KEY) {
                                out_writer.write(&written, Some(&written_md))?;
                                continue;
                            }
                        }
                        #[cfg(feature = "http")]
                        if let Some(cfg) = self.external_config.as_ref() {
                            match crate::external::maybe_externalize_batch(
                                &batch,
                                metadata.as_ref(),
                                cfg,
                            ) {
                                Ok(Some((ptr, md))) => {
                                    out_writer.write(&ptr, Some(&md))?;
                                    continue;
                                }
                                Ok(None) => {}
                                Err(e) => {
                                    // Externalization failed — fall through to inline write,
                                    // but record the error on the access log via app_err.
                                    *app_err = Some(e);
                                }
                            }
                        }
                        out_writer.write(&batch, metadata.as_ref())?;
                    }
                }
            }
            // The client writes a tick and then blocks reading our response;
            // we must flush after every lockstep iteration.
            out_writer.flush()?;

            if finished {
                break;
            }
        }
        let _ = cancelled;
        out_writer.finish()?;

        // Drain remaining input.
        let _ = input_reader.drain();
        Ok(())
    }
}

fn drain_input<R: Read>(r: &mut R) -> Result<()> {
    let mut rdr = StreamReader::new(r)?;
    rdr.drain()?;
    Ok(())
}

pub(crate) fn cast_batch(batch: &RecordBatch, target: &Schema) -> Result<RecordBatch> {
    if batch.num_columns() != target.fields().len() {
        return Err(RpcError::type_error(format!(
            "Input schema mismatch: expected {} fields, got {}",
            target.fields().len(),
            batch.num_columns()
        )));
    }
    let src_schema = batch.schema();
    for (i, field) in target.fields().iter().enumerate() {
        let src_name = src_schema.field(i).name();
        if src_name != field.name() {
            return Err(RpcError::type_error(format!(
                "Input schema mismatch: expected field {:?}, got {:?}",
                field.name(),
                src_name
            )));
        }
    }
    let opts = arrow_cast::CastOptions::default();
    let mut cols = Vec::with_capacity(batch.num_columns());
    for (i, field) in target.fields().iter().enumerate() {
        let src = batch.column(i);
        if src.data_type() == field.data_type() {
            cols.push(src.clone());
            continue;
        }
        let c = cast_with_options(src.as_ref(), field.data_type(), &opts)
            .map_err(|e| RpcError::type_error(format!("cast field {}: {}", field.name(), e)))?;
        cols.push(c);
    }
    RecordBatch::try_new(Arc::new(target.clone()), cols).map_err(RpcError::from)
}

pub(crate) fn build_log_metadata(msg: &LogMessage, server_id: &str, request_id: &str) -> Metadata {
    let mut md = Metadata::new();
    md.insert(LOG_LEVEL_KEY.to_string(), msg.level.as_str().to_string());
    md.insert(LOG_MESSAGE_KEY.to_string(), msg.message.clone());
    if !msg.extras.is_empty() {
        md.insert(LOG_EXTRA_KEY.to_string(), msg.extras_json());
    }
    if !server_id.is_empty() {
        md.insert(SERVER_ID_KEY.to_string(), server_id.to_string());
    }
    if !request_id.is_empty() {
        md.insert(REQUEST_ID_KEY.to_string(), request_id.to_string());
    }
    md
}

pub(crate) fn build_error_metadata(err: &RpcError, server_id: &str, request_id: &str) -> Metadata {
    let extra = serde_json::json!({
        "exception_type": err.error_type,
        "exception_message": err.message,
        "traceback": err.traceback,
    })
    .to_string();
    let mut md = Metadata::new();
    md.insert(LOG_LEVEL_KEY.to_string(), "EXCEPTION".to_string());
    md.insert(LOG_MESSAGE_KEY.to_string(), err.message.clone());
    md.insert(LOG_EXTRA_KEY.to_string(), extra);
    if !server_id.is_empty() {
        md.insert(SERVER_ID_KEY.to_string(), server_id.to_string());
    }
    if !request_id.is_empty() {
        md.insert(REQUEST_ID_KEY.to_string(), request_id.to_string());
    }
    md
}

/// Write an error as a complete single-batch IPC stream.
pub(crate) fn write_error_stream<W: Write>(
    w: &mut W,
    schema: &Schema,
    err: &RpcError,
    server_id: &str,
    request_id: &str,
) -> Result<()> {
    let mut sw = StreamWriter::new(w, schema)?;
    let md = build_error_metadata(err, server_id, request_id);
    sw.write(&empty_batch(schema)?, Some(&md))?;
    sw.finish()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use std::sync::atomic::{AtomicBool, Ordering};

    /// Frame a no-argument request for `method` as a self-contained IPC
    /// stream the pipe-transport serve loop can read.
    fn request_bytes(method: &str) -> Vec<u8> {
        let schema = empty_schema();
        let batch = empty_batch(&schema).unwrap();
        let mut buf = Vec::new();
        {
            let mut w = StreamWriter::new(&mut buf, &schema).unwrap();
            let mut md = Metadata::new();
            md.insert(RPC_METHOD_KEY.into(), method.into());
            md.insert(REQUEST_VERSION_KEY.into(), REQUEST_VERSION.into());
            md.insert(REQUEST_ID_KEY.into(), format!("req-{method}"));
            w.write(&batch, Some(&md)).unwrap();
            w.finish().unwrap();
        }
        buf
    }

    #[test]
    fn panicking_handler_yields_error_envelope_and_loop_survives() {
        let mut server = RpcServer::new("test-srv");
        server.register(MethodInfo::unary(
            "boom",
            empty_schema(),
            empty_schema(),
            |_req, _ctx| panic!("handler exploded"),
        ));
        let ran_second = Arc::new(AtomicBool::new(false));
        let flag = ran_second.clone();
        server.register(MethodInfo::unary(
            "ok",
            empty_schema(),
            empty_schema(),
            move |_req, _ctx| {
                flag.store(true, Ordering::SeqCst);
                Ok(None)
            },
        ));

        // Two back-to-back requests: the first handler panics, the
        // second must still run — the serve loop must not abort.
        let mut input = request_bytes("boom");
        input.extend(request_bytes("ok"));
        let mut output: Vec<u8> = Vec::new();
        server.serve(Cursor::new(input), &mut output);

        assert!(
            ran_second.load(Ordering::SeqCst),
            "serve loop aborted after a handler panic"
        );

        // The panic was surfaced to the client as an error envelope,
        // not a silent connection drop.
        let mut r = StreamReader::new(output.as_slice()).unwrap();
        let (_b, md) = r.read_next().unwrap().expect("error batch");
        assert_eq!(md_get(&md, LOG_LEVEL_KEY), Some("EXCEPTION"));
    }

    #[test]
    fn transport_options_reports_shm_capability_unregistered() {
        use crate::metadata::TRANSPORT_SHM_KEY;
        use crate::transport_options::{shm_available, TRANSPORT_OPTIONS_METHOD_NAME};

        let mut server = RpcServer::new("test-srv");
        server.register(MethodInfo::unary(
            "noop",
            empty_schema(),
            empty_schema(),
            |_req, _ctx| Ok(None),
        ));
        // Not a registered method — handled by pre-dispatch interception.
        assert!(!server.methods.contains_key(TRANSPORT_OPTIONS_METHOD_NAME));

        let input = request_bytes(TRANSPORT_OPTIONS_METHOD_NAME);
        let mut output: Vec<u8> = Vec::new();
        server.serve(Cursor::new(input), &mut output);

        let mut r = StreamReader::new(output.as_slice()).unwrap();
        let (_b, md) = r.read_next().unwrap().expect("transport options batch");
        let expected = if shm_available() { "true" } else { "false" };
        assert_eq!(md_get(&md, TRANSPORT_SHM_KEY), Some(expected));
        assert_eq!(md_get(&md, REQUEST_VERSION_KEY), Some(REQUEST_VERSION));
        assert_eq!(md_get(&md, SERVER_ID_KEY), Some("test-srv"));
    }
}