dynamo_runtime/pipeline/network/tcp/

server.rs

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use socket2::{Domain, SockAddr, SockRef, Socket, Type};
use std::{
    collections::{HashMap, HashSet},
    net::{IpAddr, SocketAddr},
    os::fd::{AsFd, FromRawFd},
    sync::Arc,
    time::Duration,
};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::time::Instant;
use tokio_rustls::TlsAcceptor;

/// Tombstone lifetime. Bridges the `register()` → `associate_instance()`
/// window (sub-millisecond in practice); 5s bounds the set by recent worker
/// churn rather than process lifetime, since etcd lease IDs are unique per
/// restart and never get cleared by an `Added` event for the same identity.
const TOMBSTONE_TTL: Duration = Duration::from_secs(5);

/// Bound best-effort worker notification after a response handshake is cancelled.
const HANDSHAKE_CANCEL_CLEANUP_TIMEOUT: Duration = Duration::from_secs(1);

use bytes::Bytes;
use derive_builder::Builder;
use futures::{SinkExt, StreamExt};
use parking_lot::Mutex;

use serde::{Deserialize, Serialize};
use tokio::{
    io::AsyncWriteExt,
    sync::{mpsc, oneshot},
    time,
};
use tokio_util::codec::{FramedRead, FramedWrite};

use super::{
    CallHomeHandshake, ControlMessage, PendingConnections, RegisteredStream, StreamOptions,
    StreamReceiver, StreamSender, TcpStreamConnectionInfo, TwoPartCodec, bind_listener,
};
use crate::discovery::EndpointInstanceId;
use crate::engine::AsyncEngineContext;
use crate::error::{DynamoError, ErrorType};
use crate::pipeline::{
    PipelineError,
    network::{
        ResponseService, ResponseStreamPrologue, StreamPrologueError,
        codec::{TwoPartMessage, TwoPartMessageType},
        tcp::StreamType,
    },
};
use crate::utils::ip_resolver::{resolve_host_or_interface, resolve_local_host};
use anyhow::{Result, anyhow as error};

pub use crate::utils::ip_resolver::{DefaultIpResolver, IpResolver};

#[allow(dead_code)]
type ResponseType = TwoPartMessage;

#[derive(Debug, Serialize, Deserialize, Clone, Builder, Default)]
pub struct ServerOptions {
    #[builder(default = "0")]
    pub port: u16,

    /// IP literal, bracketed IPv6 literal, wildcard, or network interface name.
    /// When unset, Dynamo selects a local address automatically.
    /// The field name is retained for source compatibility.
    #[builder(default)]
    pub interface: Option<String>,
}

impl ServerOptions {
    pub fn builder() -> ServerOptionsBuilder {
        ServerOptionsBuilder::default()
    }
}

/// A [`TcpStreamServer`] is a TCP service that listens on a port for incoming response connections.
/// A Response connection is a connection that is established by a client with the intention of sending
/// specific data back to the server.
pub struct TcpStreamServer {
    address: String,
    state: Arc<Mutex<State>>,
}

// pub struct TcpStreamReceiver {
//     address: TcpStreamConnectionInfo,
//     state: Arc<Mutex<State>>,
//     rx: mpsc::Receiver<ResponseType>,
// }

#[allow(dead_code)]
struct RequestedSendConnection {
    context: Arc<dyn AsyncEngineContext>,
    connection: oneshot::Sender<Result<StreamSender, StreamPrologueError>>,
    /// Capacity of the per-stream mpsc buffer between the socket task and the
    /// engine producer; carried from the registration [`StreamOptions`].
    send_buffer_count: usize,
}

struct RequestedRecvConnection {
    context: Arc<dyn AsyncEngineContext>,
    connection: oneshot::Sender<Result<StreamReceiver, StreamPrologueError>>,
    /// Remains addressable after the worker's TCP call-home so discovery
    /// removal can interrupt a stalled response prologue.
    cancellation: tokio_util::sync::CancellationToken,
    /// Capacity of the per-stream mpsc buffer between the socket task and the
    /// engine consumer; carried from the registration [`StreamOptions`].
    send_buffer_count: usize,
    defer_cancellation_until_prologue: bool,
}

/// Build the per-stream data-plane mpsc channel that bridges the socket task
/// and the engine producer/consumer. The capacity is driven by the
/// registration options ([`StreamOptions::send_buffer_count`]) rather than a
/// hard-coded constant; both `process_request_stream` and
/// `process_response_stream` size their channel through this helper. See #10293.
fn data_plane_channel<T>(send_buffer_count: usize) -> (mpsc::Sender<T>, mpsc::Receiver<T>) {
    // `tokio::sync::mpsc::channel` panics on a capacity of 0. Now that the value
    // is caller-configurable via `StreamOptions::send_buffer_count`, clamp to at
    // least 1 so a misconfigured `0` degrades to a minimal buffer instead of
    // panicking the connection handler task.
    mpsc::channel(send_buffer_count.max(1))
}

// /// When registering a new TcpStream on the server, the registration method will return a [`Connections`] object.
// /// This [`Connections`] object will have two [`oneshot::Receiver`] objects, one for the [`TcpStreamSender`] and one for the [`TcpStreamReceiver`].
// /// The [`Connections`] object can be awaited to get the [`TcpStreamSender`] and [`TcpStreamReceiver`] objects; these objects will
// /// be made available when the matching Client has connected to the server.
// pub struct Connections {
//     pub address: TcpStreamConnectionInfo,

//     /// The [`oneshot::Receiver`] for the [`TcpStreamSender`]. Awaiting this object will return the [`TcpStreamSender`] object once
//     /// the client has connected to the server.
//     pub sender: Option<oneshot::Receiver<StreamSender>>,

//     /// The [`oneshot::Receiver`] for the [`TcpStreamReceiver`]. Awaiting this object will return the [`TcpStreamReceiver`] object once
//     /// the client has connected to the server.
//     pub receiver: Option<oneshot::Receiver<StreamReceiver>>,
// }

#[derive(Default)]
struct State {
    tx_subjects: HashMap<String, RequestedSendConnection>,
    rx_subjects: HashMap<String, RequestedRecvConnection>,
    /// Response registrations remain cancellable after `rx_subjects` is
    /// consumed by call-home and until the response prologue is accepted.
    rx_cancellations: HashMap<String, tokio_util::sync::CancellationToken>,
    /// subject UUID -> EndpointInstanceId. Full 4-field key isolates services
    /// that share an endpoint name across namespaces/components.
    subject_instance: HashMap<String, EndpointInstanceId>,
    /// EndpointInstanceId -> tagged subject UUIDs, for batch cancellation on
    /// removal. The `StreamType` tag tells `cancel_instance_streams` which
    /// of `rx_subjects` / `tx_subjects` holds the registration so both halves
    /// of a bidirectional session get dropped together.
    instance_subjects: HashMap<EndpointInstanceId, HashSet<(StreamType, String)>>,
    /// Tombstones (instance -> insertion time) close the
    /// `cancel_instance_streams` vs `associate_instance` race; entries expire
    /// after [`TOMBSTONE_TTL`].
    removed_instances: HashMap<EndpointInstanceId, Instant>,
    handle: Option<tokio::task::JoinHandle<Result<()>>>,
}

struct ResponseRegistrationGuard {
    state: Arc<Mutex<State>>,
    subject: String,
    is_active: bool,
}

impl ResponseRegistrationGuard {
    /// Serialize prologue acceptance with cancellation under the registration lock.
    /// A removed cancellation entry means cancellation already won.
    fn accept(mut self) -> bool {
        self.is_active = false;
        TcpStreamServer::finish_response_stream(&self.state, &self.subject)
    }
}

impl Drop for ResponseRegistrationGuard {
    fn drop(&mut self) {
        if self.is_active {
            TcpStreamServer::finish_response_stream(&self.state, &self.subject);
        }
    }
}

/// Drop tombstones older than [`TOMBSTONE_TTL`]. Called lazily on every
/// `associate_instance` / `cancel_instance_streams` to bound the set size.
fn prune_tombstones(tombstones: &mut HashMap<EndpointInstanceId, Instant>, now: Instant) {
    tombstones.retain(|_, ts| now.saturating_duration_since(*ts) < TOMBSTONE_TTL);
}

impl TcpStreamServer {
    pub fn local_address(&self) -> Result<SocketAddr> {
        Ok(self.address.parse()?)
    }

    pub fn options_builder() -> ServerOptionsBuilder {
        ServerOptionsBuilder::default()
    }

    pub async fn new(options: ServerOptions) -> Result<Arc<Self>, PipelineError> {
        Self::new_with_resolver(options, DefaultIpResolver).await
    }

    pub async fn new_with_resolver<R: IpResolver>(
        options: ServerOptions,
        resolver: R,
    ) -> Result<Arc<Self>, PipelineError> {
        let resolved_host = if let Some(host_or_interface) = options.interface.as_deref() {
            resolve_host_or_interface(host_or_interface, &resolver).map_err(|error| {
                PipelineError::Generic(format!(
                    "Failed to resolve configured TCP host '{host_or_interface}': {error}"
                ))
            })?
        } else {
            resolve_local_host(&resolver).map_err(|error| {
                PipelineError::Generic(format!("Failed to resolve local IP address: {error}"))
            })?
        };

        if resolved_host.used_loopback_fallback() {
            tracing::warn!(
                bind_ip = %resolved_host.bind_ip(),
                advertise_ip = %resolved_host.advertise_ip(),
                "No usable local IP address found; advertising loopback"
            );
        }

        let state = Arc::new(Mutex::new(State::default()));

        // Build TLS acceptor from environment if cert+key paths are configured.
        let tls_acceptor = Self::build_tls_acceptor().map_err(|e| {
            PipelineError::Generic(format!("Failed to build TCP TLS acceptor: {}", e))
        })?;

        let bind_address = SocketAddr::new(resolved_host.bind_ip(), options.port);
        let local_address = Self::start(bind_address, state.clone(), tls_acceptor)
            .await
            .map_err(|error| {
                PipelineError::Generic(format!("Failed to start TcpStreamServer: {error}"))
            })?;
        let advertised_address =
            SocketAddr::new(resolved_host.advertise_ip(), local_address.port());
        let address = advertised_address.to_string();

        tracing::debug!(%local_address, %advertised_address, "tcp transport service started");

        Ok(Arc::new(Self { address, state }))
    }

    /// Associate one or both halves of a registration with a backend instance.
    ///
    /// `recv_subject` is the response-stream subject (always present on TCP);
    /// `send_subject` is the request-stream subject, set only for
    /// bidirectional sessions. Tracking the send half here is what lets
    /// [`Self::cancel_instance_streams`] drop the request-stream
    /// `tx_subjects` oneshot directly when discovery removes the worker,
    /// instead of relying on the cascade from the recv-side cancellation.
    ///
    /// Returns `false` if the instance is already tombstoned, in which case
    /// both subjects are cancelled immediately and the caller should skip
    /// `send_request` and fail with a migratable `Disconnected` error.
    pub async fn associate_instance(
        &self,
        recv_subject: &str,
        send_subject: Option<&str>,
        id: &EndpointInstanceId,
    ) -> bool {
        let mut state = self.state.lock();
        let now = Instant::now();
        prune_tombstones(&mut state.removed_instances, now);
        if state.removed_instances.contains_key(id) {
            // Instance was already removed -- cancel immediately.
            tracing::warn!(
                recv_subject,
                send_subject,
                namespace = %id.namespace,
                component = %id.component,
                endpoint = %id.endpoint,
                instance_id = id.instance_id,
                "Cancelling subject immediately: instance already removed (tombstoned)"
            );
            state.rx_subjects.remove(recv_subject);
            if let Some(token) = state.rx_cancellations.remove(recv_subject) {
                token.cancel();
            }
            if let Some(s) = send_subject {
                state.tx_subjects.remove(s);
            }
            return false;
        }
        state
            .subject_instance
            .insert(recv_subject.to_string(), id.clone());
        if let Some(s) = send_subject {
            state.subject_instance.insert(s.to_string(), id.clone());
        }
        let entry = state.instance_subjects.entry(id.clone()).or_default();
        entry.insert((StreamType::Response, recv_subject.to_string()));
        if let Some(s) = send_subject {
            entry.insert((StreamType::Request, s.to_string()));
        }
        true
    }

    /// Associate a request-only callback registration with a backend instance.
    /// QUIC response mode still uses TCP for bidirectional request callbacks.
    pub async fn associate_request_instance(&self, subject: &str, id: &EndpointInstanceId) -> bool {
        let mut state = self.state.lock();
        let now = Instant::now();
        prune_tombstones(&mut state.removed_instances, now);
        if state.removed_instances.contains_key(id) {
            state.tx_subjects.remove(subject);
            return false;
        }
        state
            .subject_instance
            .insert(subject.to_string(), id.clone());
        state
            .instance_subjects
            .entry(id.clone())
            .or_default()
            .insert((StreamType::Request, subject.to_string()));
        true
    }

    /// Cancel one pending response-stream registration. Drops the
    /// `oneshot::Sender` so the waiting receiver resolves with `RecvError`.
    pub async fn cancel_recv_stream(&self, subject: &str) {
        let mut state = self.state.lock();
        state.rx_subjects.remove(subject);
        if let Some(token) = state.rx_cancellations.remove(subject) {
            token.cancel();
        }
        if let Some(key) = state.subject_instance.remove(subject)
            && let Some(subjects) = state.instance_subjects.get_mut(&key)
        {
            subjects.remove(&(StreamType::Response, subject.to_string()));
            if subjects.is_empty() {
                state.instance_subjects.remove(&key);
            }
        }
    }

    /// Cancel one pending request-stream registration. Parallel to
    /// [`Self::cancel_recv_stream`]: drops the `tx_subjects` entry and, if
    /// the subject was associated with an instance, clears its
    /// `(StreamType::Request, _)` tag from `instance_subjects` so the per-
    /// instance bookkeeping stays consistent.
    pub async fn cancel_send_stream(&self, subject: &str) {
        let mut state = self.state.lock();
        state.tx_subjects.remove(subject);
        if let Some(key) = state.subject_instance.remove(subject)
            && let Some(subjects) = state.instance_subjects.get_mut(&key)
        {
            subjects.remove(&(StreamType::Request, subject.to_string()));
            if subjects.is_empty() {
                state.instance_subjects.remove(&key);
            }
        }
    }

    /// Cancel all pending stream handshakes for an instance — both response-
    /// side and request-side halves of any bidirectional sessions tracked by
    /// `associate_instance` — and tombstone the id so any racing associate
    /// for the same id cancels too. Returns the number of streams cancelled.
    pub async fn cancel_instance_streams(&self, id: &EndpointInstanceId) -> usize {
        let mut state = self.state.lock();
        let now = Instant::now();
        prune_tombstones(&mut state.removed_instances, now);
        state.removed_instances.insert(id.clone(), now);
        let subjects = match state.instance_subjects.remove(id) {
            Some(subjects) => subjects,
            None => return 0,
        };
        let count = subjects.len();
        for (kind, subject) in &subjects {
            match kind {
                StreamType::Response => {
                    state.rx_subjects.remove(subject);
                    if let Some(token) = state.rx_cancellations.remove(subject) {
                        token.cancel();
                    }
                }
                StreamType::Request => {
                    state.tx_subjects.remove(subject);
                }
            }
            state.subject_instance.remove(subject);
        }
        count
    }

    /// Drop the tombstone for an instance that has reappeared in discovery,
    /// so future subjects for that identity are tracked normally.
    pub async fn clear_instance_tombstone(&self, id: &EndpointInstanceId) {
        let mut state = self.state.lock();
        state.removed_instances.remove(id);
    }

    /// Build a TLS acceptor from env vars if `DYN_TCP_TLS_CERT_PATH` and
    /// `DYN_TCP_TLS_KEY_PATH` are set. Validation and diagnostics are shared with
    /// the request-plane server via
    /// [`crate::tls_utils::server_tls_acceptor_config`]; a client CA enables mTLS.
    fn build_tls_acceptor() -> anyhow::Result<Option<TlsAcceptor>> {
        use crate::config::environment_names::tcp_response_stream::tls as env;
        let cert_path = std::env::var(env::DYN_TCP_TLS_CERT_PATH).ok();
        let key_path = std::env::var(env::DYN_TCP_TLS_KEY_PATH).ok();
        let client_ca = std::env::var(env::DYN_TCP_TLS_CLIENT_CA_CERT_PATH).ok();
        Ok(crate::tls_utils::server_tls_acceptor_config(
            "TCP server",
            cert_path.as_deref().map(std::path::Path::new),
            key_path.as_deref().map(std::path::Path::new),
            client_ca.as_deref().map(std::path::Path::new),
        )?
        .map(|config| TlsAcceptor::from(Arc::new(config))))
    }

    async fn start(
        bind_address: SocketAddr,
        state: Arc<Mutex<State>>,
        tls_acceptor: Option<TlsAcceptor>,
    ) -> Result<SocketAddr> {
        let state_clone = state.clone();
        let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<Result<SocketAddr>>();
        {
            let mut guard = state.lock();
            if guard.handle.is_some() {
                panic!("TcpStreamServer already started");
            }
            guard.handle = Some(tokio::spawn(tcp_listener(
                bind_address,
                state_clone,
                tls_acceptor,
                ready_tx,
            )));
        }
        ready_rx.await?
    }

    fn insert_request_stream(&self, subject: String, connection: RequestedSendConnection) {
        self.state.lock().tx_subjects.insert(subject, connection);
    }

    fn insert_response_stream(&self, subject: String, connection: RequestedRecvConnection) {
        let cancellation = connection.cancellation.clone();
        let mut state = self.state.lock();
        state.rx_subjects.insert(subject.clone(), connection);
        state.rx_cancellations.insert(subject, cancellation);
    }

    fn take_request_stream(state: &Mutex<State>, subject: &str) -> Option<RequestedSendConnection> {
        let mut state = state.lock();
        let connection = state.tx_subjects.remove(subject);
        if let Some(key) = state.subject_instance.remove(subject)
            && let Some(subjects) = state.instance_subjects.get_mut(&key)
        {
            subjects.remove(&(StreamType::Request, subject.to_string()));
            if subjects.is_empty() {
                state.instance_subjects.remove(&key);
            }
        }
        connection
    }

    fn take_response_stream(
        state: &Mutex<State>,
        subject: &str,
    ) -> Option<RequestedRecvConnection> {
        state.lock().rx_subjects.remove(subject)
    }

    /// Finish tracking a response registration after its prologue is accepted
    /// or its connection handler exits. Returns whether cancellation had not
    /// already removed the registration.
    fn finish_response_stream(state: &Mutex<State>, subject: &str) -> bool {
        let mut state = state.lock();
        let is_pending = state.rx_cancellations.remove(subject).is_some();
        if let Some(key) = state.subject_instance.remove(subject)
            && let Some(subjects) = state.instance_subjects.get_mut(&key)
        {
            subjects.remove(&(StreamType::Response, subject.to_string()));
            if subjects.is_empty() {
                state.instance_subjects.remove(&key);
            }
        }
        is_pending
    }
}

// todo - possible rename ResponseService to ResponseServer
#[async_trait::async_trait]
impl ResponseService for TcpStreamServer {
    /// Register a new subject and sender with the response subscriber
    /// Produces an RAII object that will deregister the subject when dropped
    ///
    /// we need to register both data in and data out entries
    /// there might be forward pipeline that want to consume the data out stream
    /// and there might be a response stream that wants to consume the data in stream
    /// on registration, we need to specific if we want data-in, data-out or both
    /// this will map to the type of service that is runniing, i.e. Single or Many In //
    /// Single or Many Out
    ///
    /// todo(ryan) - return a connection object that can be awaited. when successfully connected,
    /// can ask for the sender and receiver
    ///
    /// OR
    ///
    /// we make it into register sender and register receiver, both would return a connection object
    /// and when a connection is established, we'd get the respective sender or receiver
    ///
    /// the registration probably needs to be done in one-go, so we should use a builder object for
    /// requesting a receiver and optional sender
    async fn register(&self, options: StreamOptions) -> PendingConnections {
        // oneshot channels to pass back the sender and receiver objects

        let address = self.address.clone();
        tracing::debug!("Registering new TcpStream on {address}");

        let send_stream = if options.enable_request_stream {
            let sender_subject = uuid::Uuid::new_v4().to_string();
            let registry_subject = sender_subject.clone();

            let (pending_sender_tx, pending_sender_rx) = oneshot::channel();

            let connection_info = RequestedSendConnection {
                context: options.context.clone(),
                connection: pending_sender_tx,
                send_buffer_count: options.send_buffer_count,
            };

            let cleanup_subject = sender_subject.clone();
            let cleanup_state = self.state.clone();
            let registered_stream = RegisteredStream::new(
                TcpStreamConnectionInfo {
                    address: address.clone(),
                    subject: sender_subject,
                    context: options.context.id().to_string(),
                    stream_type: StreamType::Request,
                }
                .into(),
                pending_sender_rx,
            )
            .with_cleanup(move || {
                // Drop is sync; fire-and-forget the lock acquisition.
                tokio::spawn(async move {
                    let mut state = cleanup_state.lock();
                    state.tx_subjects.remove(&cleanup_subject);
                    if let Some(key) = state.subject_instance.remove(&cleanup_subject)
                        && let Some(subjects) = state.instance_subjects.get_mut(&key)
                    {
                        subjects.remove(&(StreamType::Request, cleanup_subject.clone()));
                        if subjects.is_empty() {
                            state.instance_subjects.remove(&key);
                        }
                    }
                });
            });

            self.insert_request_stream(registry_subject, connection_info);

            Some(registered_stream)
        } else {
            None
        };

        let recv_stream = if options.enable_response_stream {
            let (pending_recver_tx, pending_recver_rx) = oneshot::channel();
            let receiver_subject = uuid::Uuid::new_v4().to_string();
            let registry_subject = receiver_subject.clone();

            let connection_info = RequestedRecvConnection {
                context: options.context.clone(),
                connection: pending_recver_tx,
                cancellation: tokio_util::sync::CancellationToken::new(),
                send_buffer_count: options.send_buffer_count,
                defer_cancellation_until_prologue: options.defer_cancellation_until_prologue,
            };

            let cleanup_subject = receiver_subject.clone();
            let cleanup_state = self.state.clone();
            let registered_stream = RegisteredStream::new(
                TcpStreamConnectionInfo {
                    address: address.clone(),
                    subject: receiver_subject,
                    context: options.context.id().to_string(),
                    stream_type: StreamType::Response,
                }
                .into(),
                pending_recver_rx,
            )
            .with_cleanup(move || {
                // Drop is sync; fire-and-forget the lock acquisition.
                tokio::spawn(async move {
                    let mut state = cleanup_state.lock();
                    state.rx_subjects.remove(&cleanup_subject);
                    if let Some(token) = state.rx_cancellations.remove(&cleanup_subject) {
                        token.cancel();
                    }
                    if let Some(key) = state.subject_instance.remove(&cleanup_subject)
                        && let Some(subjects) = state.instance_subjects.get_mut(&key)
                    {
                        subjects.remove(&(StreamType::Response, cleanup_subject.clone()));
                        if subjects.is_empty() {
                            state.instance_subjects.remove(&key);
                        }
                    }
                });
            });

            self.insert_response_stream(registry_subject, connection_info);

            Some(registered_stream)
        } else {
            None
        };

        PendingConnections {
            send_stream,
            recv_stream,
        }
    }
}

/// First retry delay applied after an `AcceptFailure::Exhaustion`.
const ACCEPT_BACKOFF_INITIAL_DELAY: Duration = Duration::from_millis(5);
/// Ceiling the retry delay saturates at, so the listener keeps polling often
/// enough to notice recovery while no longer spinning.
const ACCEPT_BACKOFF_MAX_DELAY: Duration = Duration::from_secs(1);
/// Minimum wall-clock interval between two emitted exhaustion summaries.
const ACCEPT_BACKOFF_LOG_INTERVAL: Duration = Duration::from_secs(5);

/// How the listener loop must treat a failed `accept()`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AcceptFailure {
    /// The process or the host is out of file descriptors or kernel memory
    /// — see `AcceptBackoff::classify` for the exact errno set. Retrying
    /// immediately cannot succeed, so the loop must back off.
    Exhaustion,
    /// Anything else — a per-connection error that the client is expected to
    /// retry. Keeps the historical warn-and-retry-immediately behavior.
    Ordinary,
}

/// What the caller should do about one exhaustion failure.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct AcceptBackoffAction {
    /// How long to sleep before calling `accept()` again.
    delay: Duration,
    /// `Some(n)` when the caller should emit one summary line, where `n` is the
    /// number of failures suppressed since the previous emission. `None` means
    /// this failure is inside the current rate-limit window and must stay silent.
    log_suppressed: Option<u64>,
}

/// Bounded exponential retry policy for the accept loop, kept deliberately pure:
/// it reads no clock, performs no I/O and never sleeps. The caller supplies the
/// current `Instant` and performs the sleep, which is what lets the retry
/// schedule and the log-rate decision be unit-tested without reproducing the
/// host's file-descriptor limit. Vocabulary mirrors `transports::etcd::connector`'s
/// `BackoffState`.
#[derive(Debug)]
struct AcceptBackoff {
    initial_delay: Duration,
    max_delay: Duration,
    log_interval: Duration,
    current_delay: Duration,
    suppressed: u64,
    last_log_at: Option<std::time::Instant>,
    /// Whether the loop is currently inside an exhaustion episode — the flag
    /// `record_success` checks to tell a real recovery from an ordinary
    /// steady-state accept.
    in_backoff: bool,
}

impl Default for AcceptBackoff {
    fn default() -> Self {
        Self {
            initial_delay: ACCEPT_BACKOFF_INITIAL_DELAY,
            max_delay: ACCEPT_BACKOFF_MAX_DELAY,
            log_interval: ACCEPT_BACKOFF_LOG_INTERVAL,
            current_delay: ACCEPT_BACKOFF_INITIAL_DELAY,
            suppressed: 0,
            last_log_at: None,
            in_backoff: false,
        }
    }
}

impl AcceptBackoff {
    /// Classify an `accept()` error. Only descriptor or kernel-memory exhaustion
    /// gets the backoff treatment; everything else stays on the pre-existing
    /// path so genuinely unexpected failures are not hidden or slowed down.
    fn classify(err: &std::io::Error) -> AcceptFailure {
        #[cfg(unix)]
        {
            // `std::io::ErrorKind` has no stable variant for these errnos, so
            // the raw value is the portable-with-cfg way to recognize them.
            // `EMFILE`/`ENFILE` are descriptor exhaustion; `ENOBUFS`/`ENOMEM`
            // are kernel out-of-memory for the socket. All four make an
            // immediate retry deterministic, so they get the backoff path.
            if matches!(
                err.raw_os_error(),
                Some(libc::EMFILE) | Some(libc::ENFILE) | Some(libc::ENOBUFS) | Some(libc::ENOMEM)
            ) {
                return AcceptFailure::Exhaustion;
            }
        }
        #[cfg(not(unix))]
        {
            let _ = err;
        }
        AcceptFailure::Ordinary
    }

    /// Record one exhaustion failure and return the delay to sleep plus the
    /// rate-limited logging decision. The delay doubles per consecutive failure
    /// and saturates at `max_delay`.
    fn record_exhaustion(&mut self, now: std::time::Instant) -> AcceptBackoffAction {
        self.in_backoff = true;

        let delay = self.current_delay;
        self.current_delay = self.current_delay.saturating_mul(2).min(self.max_delay);

        let due = match self.last_log_at {
            None => true,
            Some(prev) => now.saturating_duration_since(prev) >= self.log_interval,
        };

        let log_suppressed = if due {
            self.last_log_at = Some(now);
            Some(std::mem::take(&mut self.suppressed))
        } else {
            self.suppressed += 1;
            None
        };

        AcceptBackoffAction {
            delay,
            log_suppressed,
        }
    }

    /// Record a successful accept. `Some(suppressed)` only when this success
    /// both ended an exhaustion episode and the shared log-rate window allows
    /// a line — so a listener flapping at the descriptor ceiling cannot emit
    /// an unbounded stream of recovery lines, and a suppressed recovery keeps
    /// its count for the next emission rather than discarding it. The clock
    /// arrives as a closure, not an `Instant`, so the steady-state accept
    /// (every ordinary connection) hits the early return below without
    /// reading a clock at all.
    fn record_success(&mut self, now: impl FnOnce() -> std::time::Instant) -> Option<u64> {
        if !self.in_backoff {
            return None;
        }
        self.current_delay = self.initial_delay;
        self.in_backoff = false;
        let now = now();

        let due = match self.last_log_at {
            None => true,
            Some(prev) => now.saturating_duration_since(prev) >= self.log_interval,
        };
        if !due {
            return None;
        }
        self.last_log_at = Some(now);
        Some(std::mem::take(&mut self.suppressed))
    }
}

/// Apply the accept-loop error policy to one failed `accept()`. Returns the
/// delay that was actually slept — `Duration::ZERO` for ordinary errors, which
/// keep the historical immediate-retry behavior.
async fn handle_accept_error(err: &std::io::Error, backoff: &mut AcceptBackoff) -> Duration {
    match AcceptBackoff::classify(err) {
        AcceptFailure::Ordinary => {
            // the client should retry, so we don't need to abort
            tracing::warn!(error = %err, "failed to accept tcp connection");
            // Gated like the sibling in `handle_connection`: an unconditional
            // stderr write here doubles the log volume of any accept-error storm.
            #[cfg(debug_assertions)]
            eprintln!("failed to accept tcp connection: {}", err);
            Duration::ZERO
        }
        AcceptFailure::Exhaustion => {
            crate::metrics::transport_metrics::TCP_ACCEPT_BACKOFF_TOTAL.inc();
            let action = backoff.record_exhaustion(std::time::Instant::now());
            if let Some(suppressed) = action.log_suppressed {
                tracing::warn!(
                    error = %err,
                    retry_delay_ms = action.delay.as_millis() as u64,
                    suppressed_failures = suppressed,
                    "tcp accept failed: out of file descriptors or kernel memory; backing off before retry"
                );
            }
            time::sleep(action.delay).await;
            action.delay
        }
    }
}

// Type aliases for boxed split halves used throughout the nested handlers below.
type BoxRead = Box<dyn tokio::io::AsyncRead + Unpin + Send>;
type BoxWrite = Box<dyn tokio::io::AsyncWrite + Unpin + Send>;

/// Unblock the requester before attempting any network cleanup. A stalled
/// worker must not keep either the requester or this socket task alive.
async fn cancel_response_handshake(
    connection: oneshot::Sender<Result<StreamReceiver, StreamPrologueError>>,
    registration_guard: Option<ResponseRegistrationGuard>,
    mut writer: FramedWrite<BoxWrite, TwoPartCodec>,
) {
    drop(registration_guard);
    drop(connection);
    let cleanup = async move {
        if let Ok(bytes) = serde_json::to_vec(&ControlMessage::Kill) {
            let _ = writer.send(TwoPartMessage::from_header(bytes.into())).await;
        }
        let _ = writer.into_inner().shutdown().await;
    };
    if time::timeout(HANDSHAKE_CANCEL_CLEANUP_TIMEOUT, cleanup)
        .await
        .is_err()
    {
        tracing::debug!("timed out notifying worker of cancelled response handshake");
    }
}

// this method listens on a tcp port for incoming connections
// new connections are expected to send a protocol specific handshake
// for us to determine the subject they are interested in, in this case,
// we expect the first message to be [`FirstMessage`] from which we find
// the sender, then we spawn a task to forward all bytes from the tcp stream
// to the sender
async fn tcp_listener(
    addr: SocketAddr,
    state: Arc<Mutex<State>>,
    tls_acceptor: Option<TlsAcceptor>,
    read_tx: tokio::sync::oneshot::Sender<Result<SocketAddr>>,
) -> Result<()> {
    let listener = bind_listener(addr)
        .map_err(|e| anyhow::anyhow!("Failed to start TcpListender on {}: {}", addr, e));

    let listener = match listener {
        Ok(listener) => {
            let addr = listener
                .local_addr()
                .map_err(|e| anyhow::anyhow!("Failed get SocketAddr: {:?}", e))
                .unwrap();

            read_tx.send(Ok(addr)).expect("Failed to send ready signal");

            listener
        }
        Err(e) => {
            read_tx.send(Err(e)).expect("Failed to send ready signal");
            return Err(anyhow::anyhow!("Failed to start TcpListender on {}", addr));
        }
    };

    let mut accept_backoff = AcceptBackoff::default();

    loop {
        // todo - add instrumentation
        // todo - add counter for all accepted connections
        // todo - add gauge for all inflight connections
        // todo - add counter for incoming bytes
        // todo - add counter for outgoing bytes
        let (stream, _addr) = match listener.accept().await {
            Ok((stream, _addr)) => {
                if let Some(suppressed) = accept_backoff.record_success(std::time::Instant::now) {
                    tracing::warn!(
                        suppressed_failures = suppressed,
                        "tcp accept recovered from resource exhaustion"
                    );
                }
                (stream, _addr)
            }
            Err(e) => {
                handle_accept_error(&e, &mut accept_backoff).await;
                continue;
            }
        };

        match stream.set_nodelay(true) {
            Ok(_) => (),
            Err(e) => {
                tracing::warn!("failed to set tcp stream to nodelay: {e}");
            }
        }

        match SockRef::from(&stream).set_linger(Some(std::time::Duration::from_secs(0))) {
            Ok(_) => (),
            Err(e) => {
                tracing::warn!("failed to set tcp stream to linger: {e}");
            }
        }

        // Spawn per-connection so the accept loop is never blocked by a slow
        // TLS handshake. The handshake is bounded by DYN_TCP_TLS_HANDSHAKE_TIMEOUT_SECS (default 3s).
        let state_clone = state.clone();
        let tls_acceptor_clone = tls_acceptor.clone();
        tokio::spawn(async move {
            let (reader, writer) = if let Some(ref tls) = tls_acceptor_clone {
                match tokio::time::timeout(
                    crate::tls_utils::handshake_timeout(),
                    tls.accept(stream),
                )
                .await
                {
                    Ok(Ok(tls_stream)) => {
                        let (r, w) = tokio::io::split(tls_stream);
                        (Box::new(r) as BoxRead, Box::new(w) as BoxWrite)
                    }
                    Ok(Err(e)) => {
                        tracing::warn!("TLS handshake failed: {e}");
                        return;
                    }
                    Err(_) => {
                        tracing::warn!("TLS handshake timed out");
                        return;
                    }
                }
            } else {
                let (r, w) = tokio::io::split(stream);
                (Box::new(r) as BoxRead, Box::new(w) as BoxWrite)
            };
            handle_connection(reader, writer, state_clone).await;
        });
    }

    // #[instrument(level = "trace"), skip(state)]
    // todo - clone before spawn and trace process_stream
    async fn handle_connection(reader: BoxRead, writer: BoxWrite, state: Arc<Mutex<State>>) {
        let result = process_stream(reader, writer, state).await;
        match result {
            Ok(_) => tracing::trace!("successfully processed tcp connection"),
            Err(e) => {
                tracing::warn!("failed to handle tcp connection: {e}");
                #[cfg(debug_assertions)]
                eprintln!("failed to handle tcp connection: {}", e);
            }
        }
    }

    /// This method is responsible for the internal tcp stream handshake
    /// The handshake will specialize the stream as a request/sender or response/receiver stream
    async fn process_stream(
        read_half: BoxRead,
        write_half: BoxWrite,
        state: Arc<Mutex<State>>,
    ) -> Result<()> {
        // attach the codec to the reader and writer to get framed readers and writers
        let mut framed_reader = FramedRead::new(read_half, TwoPartCodec::default());
        let framed_writer = FramedWrite::new(write_half, TwoPartCodec::default());

        // the internal tcp [`CallHomeHandshake`] connects the socket to the requester
        // here we await this first message as a raw bytes two part message
        let first_message =
            tokio::time::timeout(std::time::Duration::from_secs(10), framed_reader.next())
                .await
                .map_err(|_| error!("Timed out waiting for CallHomeHandshake"))?
                .ok_or(error!("Connection closed without a ControlMessage"))??;

        // we await on the raw bytes which should come in as a header only message
        // todo - improve error handling - check for no data
        let handshake: CallHomeHandshake = match first_message.header() {
            Some(header) => serde_json::from_slice(header).map_err(|e| {
                error!(
                    "Failed to deserialize the first message as a valid `CallHomeHandshake`: {e}",
                )
            })?,
            None => {
                return Err(error!("Expected ControlMessage, got DataMessage"));
            }
        };

        // branch here to handle sender stream or receiver stream
        match handshake.stream_type {
            StreamType::Request => {
                process_request_stream(handshake.subject, state, framed_reader, framed_writer).await
            }
            StreamType::Response => {
                process_response_stream(handshake.subject, state, framed_reader, framed_writer)
                    .await
            }
        }
    }

    /// Symmetric to [`process_response_stream`] for the upstream→downstream
    /// data direction: deliver the [`StreamSender`] half registered by the
    /// upstream to whoever awaits it, then pump every frame the upstream pushes
    /// into the now-connected TCP socket.
    ///
    /// One difference is that the request stream is **unidirectional**:
    /// the upstream writes data + one closing control message, and the
    /// downstream is not expected to reply: downstream response or inference
    /// error should be returned through response stream. We therefore drop the
    /// read half, on fatal error, downstream should drop the request stream.
    async fn process_request_stream(
        subject: String,
        state: Arc<Mutex<State>>,
        reader: FramedRead<BoxRead, TwoPartCodec>,
        writer: FramedWrite<BoxWrite, TwoPartCodec>,
    ) -> Result<()> {
        // Request stream is unidirectional; we don't read from the downstream.
        drop(reader);

        let request_stream = TcpStreamServer::take_request_stream(&state, &subject).ok_or_else(|| {
            error!(
                "Subject not found: {}; downstream subscriber specified a subject unknown to the upstream publisher",
                subject
            )
        })?;

        let RequestedSendConnection {
            context,
            connection,
            send_buffer_count,
        } = request_stream;

        // Buffer size is driven by the registration options
        // ([`StreamOptions::send_buffer_count`]) rather than hard-coded; the
        // same applies to `process_response_stream`. See #10293.
        let (request_tx, request_rx) = data_plane_channel(send_buffer_count);

        if connection
            .send(Ok(crate::pipeline::network::StreamSender {
                tx: request_tx,
                // Request streams don't carry a downstream-prologue today; the
                // upstream may begin sending immediately.
                prologue: None,
            }))
            .is_err()
        {
            return Err(error!(
                "The requester of the request stream has been dropped before the connection was established"
            ));
        }

        request_stream_send_handler(writer, request_rx, context).await;
        Ok(())
    }

    /// Pump frames the upstream queued on its `StreamSender` into the TCP socket.
    /// The closing control message depends on why the loop exited:
    /// - `context.killed()` → [`ControlMessage::Kill`] (hard cancel notification)
    /// - `context.stopped()` → [`ControlMessage::Stop`] (graceful cancel notification)
    /// - `request_rx` returns `None` → [`ControlMessage::Sentinel`] (clean EOS)
    /// - write error → no control message, the socket is already broken
    ///
    /// The downstream `handle_request_reader` matches on the received variant
    /// and reacts accordingly.
    async fn request_stream_send_handler(
        mut framed_writer: FramedWrite<BoxWrite, TwoPartCodec>,
        mut request_rx: mpsc::Receiver<TwoPartMessage>,
        context: Arc<dyn AsyncEngineContext>,
    ) {
        // Construct the cancellation futures once. Recreating them for every frame clones
        // the context's watch receivers and repeatedly registers/drops Tokio notifications.
        let killed = context.killed();
        let stopped = context.stopped();
        tokio::pin!(killed, stopped);

        let closing_msg: Option<ControlMessage> = loop {
            tokio::select! {
                biased;

                _ = &mut killed => {
                    tracing::trace!("context kill received in request-stream send handler");
                    break Some(ControlMessage::Kill);
                }

                _ = &mut stopped => {
                    tracing::trace!("context stop received in request-stream send handler");
                    break Some(ControlMessage::Stop);
                }

                msg = request_rx.recv() => {
                    match msg {
                        Some(msg) => {
                            if let Err(e) = framed_writer.send(msg).await {
                                tracing::trace!(
                                    "failed to send request-stream frame to downstream: {:?}",
                                    e
                                );
                                break None;
                            }
                        }
                        None => {
                            tracing::trace!("upstream request-stream sender closed; sending sentinel");
                            break Some(ControlMessage::Sentinel);
                        }
                    }
                }
            }
        };

        if let Some(ctrl) = closing_msg
            && let Ok(bytes) = serde_json::to_vec(&ctrl)
            && let Err(err) = framed_writer
                .send(TwoPartMessage::from_header(bytes.into()))
                .await
        {
            tracing::trace!(?err, ?ctrl, "request-stream closing-frame send failed");
        }

        let mut inner = framed_writer.into_inner();
        if let Err(err) = inner.flush().await {
            tracing::trace!(?err, "request-stream socket flush failed");
        }
        if let Err(err) = inner.shutdown().await {
            tracing::trace!(?err, "request-stream socket shutdown failed");
        }
    }

    async fn process_response_stream(
        subject: String,
        state: Arc<Mutex<State>>,
        mut reader: FramedRead<BoxRead, TwoPartCodec>,
        mut writer: FramedWrite<BoxWrite, TwoPartCodec>,
    ) -> Result<()> {
        let response_stream = TcpStreamServer::take_response_stream(&state, &subject).ok_or_else(|| {
            error!("Subject not found: {}; upstream publisher specified a subject unknown to the downsteam subscriber", subject)
        })?;
        let registration_guard = ResponseRegistrationGuard {
            state,
            subject,
            is_active: true,
        };

        // unwrap response_stream
        let RequestedRecvConnection {
            context,
            mut connection,
            cancellation,
            send_buffer_count,
            defer_cancellation_until_prologue,
        } = response_stream;

        // the [`Prologue`]
        // there must be a second control message it indicate the other segment's generate method was successful
        // No timeout here: the worker sends the prologue only after generate() setup completes,
        // which can take arbitrarily long (model load, queue delay, cold start).
        let prologue = tokio::select! {
            biased;
            _ = cancellation.cancelled() => {
                cancel_response_handshake(connection, Some(registration_guard), writer).await;
                return Ok(());
            }
            _ = context.killed(), if !defer_cancellation_until_prologue => Err(ControlMessage::Kill),
            _ = connection.closed() => Err(ControlMessage::Kill),
            _ = context.stopped(), if !defer_cancellation_until_prologue => Err(ControlMessage::Stop),
            prologue = reader.next() => Ok(prologue),
        };
        let prologue = match prologue {
            Ok(prologue) => {
                prologue.ok_or(error!("Connection closed without a ControlMessage"))??
            }
            Err(control) => {
                drop(registration_guard);
                let message = "Request cancelled before response stream was established";
                let error = DynamoError::builder()
                    .error_type(ErrorType::Cancelled)
                    .message(message)
                    .build();
                let _ = connection.send(Err(StreamPrologueError::new(message, error)));
                writer
                    .send(TwoPartMessage::from_header(
                        serde_json::to_vec(&control)?.into(),
                    ))
                    .await?;
                writer.close().await?;
                return Ok(());
            }
        };

        // deserialize prologue
        let prologue = match prologue.into_message_type() {
            TwoPartMessageType::HeaderOnly(header) => {
                match serde_json::from_slice::<ResponseStreamPrologue>(&header) {
                    Ok(prologue) => prologue,
                    Err(e) => {
                        // Notify the requester as the sibling arm does. Returning on
                        // `?` alone drops the oneshot un-sent, and the requester then
                        // reports a bare disconnect that names neither the worker's
                        // failure nor this one.
                        let msg = format!("malformed prologue: {e}");
                        let _ =
                            connection.send(Err(StreamPrologueError::from_message(msg.clone())));
                        return Err(error!(msg));
                    }
                }
            }
            _ => {
                // Worker sent a non-HeaderOnly frame in the prologue slot
                // (protocol violation, version skew, corruption). Notify the
                // requester so the generate call chain fails cleanly, then
                // return Err so the connection task ends without panicking.
                let msg = "malformed prologue: expected HeaderOnly ControlMessage";
                let _ = connection.send(Err(StreamPrologueError::from_message(msg)));
                return Err(error!(msg));
            }
        };

        // await the control message of GTG or Error, if error, then connection.send(Err(String)), which should fail the
        // generate call chain
        //
        // note: this second control message might be delayed, but the expensive part of setting up the connection
        // is both complete and ready for data flow; awaiting here is not a performance hit or problem and it allows
        // us to trace the initial setup time vs the time to prologue
        if let Some(error) = prologue.error {
            let returned = error!("Received error prologue: {error}");
            // Forward the worker's typed error so the requesting side can classify
            // the failure instead of parsing the message. An older worker sends none.
            let _ = connection.send(Err(StreamPrologueError {
                message: error,
                typed_error: prologue.typed_error,
            }));
            return Err(returned);
        }

        // Discovery removal prevents new work but does not abort an established
        // response stream. From this point normal transport failure or the
        // request context owns cancellation.
        if !registration_guard.accept() {
            cancel_response_handshake(connection, None, writer).await;
            return Ok(());
        }

        // Buffer size is driven by the registration options
        // ([`StreamOptions::send_buffer_count`]) rather than hard-coded; the
        // same applies to `process_request_stream`. See #10293.
        let (response_tx, response_rx) = data_plane_channel(send_buffer_count);

        if connection
            .send(Ok(crate::pipeline::network::StreamReceiver {
                rx: response_rx,
            }))
            .is_err()
        {
            return Err(error!(
                "The requester of the stream has been dropped before the connection was established"
            ));
        }

        let (control_tx, control_rx) = mpsc::channel::<ControlMessage>(1);

        // sender task
        // issues control messages to the sender and when finished shuts down the socket
        // this should be the last task to finish and must
        let send_task = tokio::spawn(network_send_handler(writer, control_rx));

        // forward task
        let recv_task = tokio::spawn(network_receive_handler(
            reader,
            response_tx,
            control_tx,
            context.clone(),
        ));

        // check the results of each of the tasks
        let (monitor_result, forward_result) = tokio::join!(send_task, recv_task);

        monitor_result?;
        forward_result?;

        Ok(())
    }

    async fn network_receive_handler(
        mut framed_reader: FramedRead<BoxRead, TwoPartCodec>,
        response_tx: mpsc::Sender<Bytes>,
        control_tx: mpsc::Sender<ControlMessage>,
        context: Arc<dyn AsyncEngineContext>,
    ) {
        // These futures stay pending across frames. Constructing them inside the loop clones
        // watch receivers and registers/drops notifications for every streamed token.
        let response_closed = response_tx.closed();
        let killed = context.killed();
        let stopped = context.stopped();
        tokio::pin!(response_closed, killed, stopped);

        // loop over reading the tcp stream and checking if the writer is closed
        let mut can_stop = true;
        loop {
            tokio::select! {
                biased;

                _ = &mut response_closed => {
                    tracing::trace!("response channel closed before the client finished writing data");
                    let _ = control_tx.send(ControlMessage::Kill).await;
                    break;
                }

                _ = &mut killed => {
                    tracing::trace!("context kill signal received; shutting down");
                    let _ = control_tx.send(ControlMessage::Kill).await;
                    break;
                }

                _ = &mut stopped, if can_stop => {
                    tracing::trace!("context stop signal received; shutting down");
                    // `stopped` is now complete; keep this branch disabled because polling
                    // the same completed async future again would panic.
                    can_stop = false;
                    let _ = control_tx.send(ControlMessage::Stop).await;
                }

                msg = framed_reader.next() => {
                    match msg {
                        Some(Ok(msg)) => {
                            let (header, data) = msg.into_parts();

                            // received a control message
                            if !header.is_empty() {
                                match process_control_message(header) {
                                    Ok(ControlAction::Continue) => {}
                                    Ok(ControlAction::Shutdown) => {
                                        if !data.is_empty() {
                                            // Sentinel-with-data is a protocol
                                            // violation; kill this stream, don't
                                            // assert!() the process down.
                                            tracing::warn!(
                                                data_len = data.len(),
                                                "client sent Sentinel with data (protocol violation); killing stream"
                                            );
                                            let _ = control_tx.send(ControlMessage::Kill).await;
                                            break;
                                        }
                                        tracing::trace!("received sentinel message; shutting down");
                                        break;
                                    }
                                    Err(e) => {
                                        // Malformed control message — kill only
                                        // this stream.
                                        tracing::warn!(err = ?e, "malformed control message, closing connection");
                                        let _ = control_tx.send(ControlMessage::Kill).await;
                                        break;
                                    }
                                }
                            }

                            if !data.is_empty()
                                && let Err(err) = response_tx.send(data).await {
                                    tracing::debug!(?err, "forwarding body/data to response channel failed");
                                    let _ = control_tx.send(ControlMessage::Kill).await;
                                    break;
                                };
                        }
                        Some(Err(e)) => {
                            // TCP RST or decode error from worker — kill only
                            // this stream.
                            tracing::warn!(err = ?e, "tcp stream read error from worker, closing connection");
                            let _ = control_tx.send(ControlMessage::Kill).await;
                            break;
                        }
                        None => {
                            // this is allowed but we try to avoid it
                            // the logic is that the client will tell us when its is done and the server
                            // will close the connection naturally when the sentinel message is received
                            // the client closing early represents a transport error outside the control of the
                            // transport library
                            tracing::trace!("tcp stream was closed by client");
                            break;
                        }
                    }
                }

            }
        }
    }

    async fn network_send_handler(
        socket_tx: FramedWrite<BoxWrite, TwoPartCodec>,
        control_rx: mpsc::Receiver<ControlMessage>,
    ) {
        let mut socket_tx = socket_tx;
        let mut control_rx = control_rx;

        while let Some(control_msg) = control_rx.recv().await {
            // Sentinel is a worker→frontend message; receiving one here means
            // a producer is buggy. Skip rather than asserting — a stream-level
            // bug must not panic the worker.
            if matches!(control_msg, ControlMessage::Sentinel) {
                tracing::warn!("received sentinel on send-side control channel; dropping");
                continue;
            }
            let bytes = match serde_json::to_vec(&control_msg) {
                Ok(b) => b,
                Err(e) => {
                    // Closed enum of small variants; serialization shouldn't
                    // fail. If it ever does, log and skip rather than panic.
                    tracing::warn!(err = ?e, ?control_msg, "failed to serialize control message");
                    continue;
                }
            };
            let message = TwoPartMessage::from_header(bytes.into());
            match socket_tx.send(message).await {
                Ok(_) => tracing::debug!(?control_msg, "issued control message"),
                Err(e) => {
                    tracing::debug!(err = ?e, ?control_msg, "failed to send control message")
                }
            }
        }

        let mut inner = socket_tx.into_inner();
        if let Err(e) = inner.flush().await {
            tracing::debug!("failed to flush socket: {e}");
        }
        if let Err(e) = inner.shutdown().await {
            tracing::debug!("failed to shutdown socket: {e}");
        }
    }
}

enum ControlAction {
    Continue,
    Shutdown,
}

fn process_control_message(message: Bytes) -> Result<ControlAction> {
    match serde_json::from_slice::<ControlMessage>(&message)? {
        ControlMessage::Sentinel => {
            // the client issued a sentinel message
            // it has finished writing data and is now awaiting the server to close the connection
            tracing::trace!("sentinel received; shutting down");
            Ok(ControlAction::Shutdown)
        }
        ControlMessage::Kill | ControlMessage::Stop => {
            // Worker→frontend control direction only carries Sentinel. Kill/Stop
            // here is a protocol violation; the caller turns this Err into a
            // stream-local Kill rather than a process-fatal event.
            anyhow::bail!("unexpected control message on response stream");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::AsyncEngineContextProvider;
    use crate::error::{BackendError, DynamoError, ErrorType};
    use crate::pipeline::Context;
    use crate::pipeline::network::DEFAULT_SEND_BUFFER_COUNT;
    use crate::pipeline::network::tcp::client::TcpClient;
    use crate::tls_utils::test_certs::self_signed_pair;
    use crate::utils::ip_resolver::test_support::{ProbeOutcome, StubResolver};
    use tokio::io::{AsyncWriteExt, ReadHalf, WriteHalf};
    use tokio::net::TcpStream;

    #[test]
    fn build_tls_acceptor_no_env_vars_is_plaintext() {
        if crate::test_utils::run_isolated(
            concat!(
                module_path!(),
                "::build_tls_acceptor_no_env_vars_is_plaintext"
            ),
            &[],
        ) {
            return;
        }

        // Also clear the client-CA var: ambient it would turn this into an error
        // (client CA without a server cert/key) instead of plaintext.
        temp_env::with_vars_unset(
            [
                "DYN_TCP_TLS_CERT_PATH",
                "DYN_TCP_TLS_KEY_PATH",
                "DYN_TCP_TLS_CLIENT_CA_CERT_PATH",
            ],
            || {
                assert!(TcpStreamServer::build_tls_acceptor().unwrap().is_none());
            },
        );
    }

    #[test]
    fn build_tls_acceptor_partial_config_errors() {
        if crate::test_utils::run_isolated(
            concat!(module_path!(), "::build_tls_acceptor_partial_config_errors"),
            &[],
        ) {
            return;
        }

        let (cert, key) = self_signed_pair();
        let cert_str = cert.path().to_str().unwrap();
        let key_str = key.path().to_str().unwrap();
        // only cert
        temp_env::with_vars(
            [
                ("DYN_TCP_TLS_CERT_PATH", Some(cert_str)),
                ("DYN_TCP_TLS_KEY_PATH", None),
            ],
            || assert!(TcpStreamServer::build_tls_acceptor().is_err()),
        );
        // only key
        temp_env::with_vars(
            [
                ("DYN_TCP_TLS_CERT_PATH", None),
                ("DYN_TCP_TLS_KEY_PATH", Some(key_str)),
            ],
            || assert!(TcpStreamServer::build_tls_acceptor().is_err()),
        );
    }

    #[test]
    fn build_tls_acceptor_both_paths_is_tls() {
        if crate::test_utils::run_isolated(
            concat!(module_path!(), "::build_tls_acceptor_both_paths_is_tls"),
            &[],
        ) {
            return;
        }

        let (cert, key) = self_signed_pair();
        temp_env::with_vars(
            [
                ("DYN_TCP_TLS_CERT_PATH", Some(cert.path().to_str().unwrap())),
                ("DYN_TCP_TLS_KEY_PATH", Some(key.path().to_str().unwrap())),
            ],
            || assert!(TcpStreamServer::build_tls_acceptor().unwrap().is_some()),
        );
    }

    #[test]
    fn build_tls_acceptor_with_client_ca_is_mtls() {
        if crate::test_utils::run_isolated(
            concat!(
                module_path!(),
                "::build_tls_acceptor_with_client_ca_is_mtls"
            ),
            &[],
        ) {
            return;
        }

        // A client CA turns the response-stream server into an mTLS acceptor.
        let (cert, key) = self_signed_pair();
        temp_env::with_vars(
            [
                ("DYN_TCP_TLS_CERT_PATH", Some(cert.path().to_str().unwrap())),
                ("DYN_TCP_TLS_KEY_PATH", Some(key.path().to_str().unwrap())),
                (
                    "DYN_TCP_TLS_CLIENT_CA_CERT_PATH",
                    Some(cert.path().to_str().unwrap()),
                ),
            ],
            || assert!(TcpStreamServer::build_tls_acceptor().unwrap().is_some()),
        );
    }

    #[test]
    fn build_tls_acceptor_client_ca_without_server_identity_errors() {
        if crate::test_utils::run_isolated(
            concat!(
                module_path!(),
                "::build_tls_acceptor_client_ca_without_server_identity_errors"
            ),
            &[],
        ) {
            return;
        }

        let (cert, _key) = self_signed_pair();
        temp_env::with_vars(
            [
                ("DYN_TCP_TLS_CERT_PATH", None),
                ("DYN_TCP_TLS_KEY_PATH", None),
                (
                    "DYN_TCP_TLS_CLIENT_CA_CERT_PATH",
                    Some(cert.path().to_str().unwrap()),
                ),
            ],
            || assert!(TcpStreamServer::build_tls_acceptor().is_err()),
        );
    }

    async fn registered_tcp_info(server: &TcpStreamServer) -> TcpStreamConnectionInfo {
        let context = Context::new(());
        let stream_options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();

        let pending_connection = server.register(stream_options).await;
        let connection_info = pending_connection
            .recv_stream
            .as_ref()
            .unwrap()
            .connection_info
            .clone();
        connection_info.try_into().unwrap()
    }

    #[tokio::test]
    async fn wildcard_bind_advertises_concrete_ipv4_with_bound_port() {
        for host in ["0.0.0.0", "::ffff:0.0.0.0", " [::ffff:0.0.0.0] "] {
            let mut resolver = StubResolver::not_found();
            resolver
                .interfaces
                .push(("eth0", "192.0.2.20".parse().unwrap()));
            let options = ServerOptions::builder()
                .port(0)
                .interface(Some(host.to_string()))
                .build()
                .unwrap();
            let server = TcpStreamServer::new_with_resolver(options, resolver)
                .await
                .unwrap();

            let tcp_info = registered_tcp_info(&server).await;
            let socket_addr: SocketAddr = tcp_info.address.parse().unwrap();
            assert_eq!(socket_addr.ip(), "192.0.2.20".parse::<IpAddr>().unwrap());
            assert_ne!(socket_addr.port(), 0);
            assert!(!socket_addr.ip().is_unspecified());
            // The advertised port must belong to the real wildcard listener.
            let _connection = tokio::net::TcpStream::connect(SocketAddr::new(
                std::net::Ipv4Addr::LOCALHOST.into(),
                socket_addr.port(),
            ))
            .await
            .unwrap();
        }
    }

    #[tokio::test]
    async fn real_bracketed_ipv6_host_binds_and_registers() {
        if let Err(error) = std::net::TcpListener::bind("[::1]:0") {
            eprintln!("Skipping IPv6 loopback bind test: {error}");
            return;
        }

        let options = ServerOptions::builder()
            .port(0)
            .interface(Some("[::1]".to_string()))
            .build()
            .unwrap();
        let server = TcpStreamServer::new_with_resolver(options, StubResolver::not_found())
            .await
            .unwrap();

        let tcp_info = registered_tcp_info(&server).await;
        let socket_addr: SocketAddr = tcp_info.address.parse().unwrap();
        assert_eq!(socket_addr.ip(), IpAddr::V6(std::net::Ipv6Addr::LOCALHOST));
        assert_ne!(socket_addr.port(), 0);
        assert!(
            tcp_info.address.starts_with("[::1]:"),
            "{}",
            tcp_info.address
        );
    }

    #[tokio::test]
    async fn tcp_server_preserves_interface_enumeration_error() {
        let mut resolver = StubResolver::not_found();
        resolver.interface_error = Some(ProbeOutcome::Platform("test-platform"));
        let options = ServerOptions::builder().port(0).build().unwrap();
        let error = TcpStreamServer::new_with_resolver(options, resolver)
            .await
            .err()
            .expect("an interface enumeration failure must be returned");

        let error = error.to_string();
        assert!(
            error.contains("failed to enumerate network interfaces"),
            "{error}"
        );
        assert!(error.contains("test-platform"), "{error}");
    }

    #[tokio::test]
    async fn test_tcp_stream_server_default_behavior() {
        // Test that TcpStreamServer::new works with default options
        // This verifies normal operation when IP detection succeeds
        let options = ServerOptions::default();
        let result = TcpStreamServer::new(options).await;

        assert!(
            result.is_ok(),
            "TcpStreamServer::new should succeed with default options"
        );

        let server = result.unwrap();

        let tcp_info = registered_tcp_info(&server).await;
        let socket_addr = tcp_info.address.parse::<std::net::SocketAddr>().unwrap();

        // Should have a valid port assigned
        assert!(
            socket_addr.port() > 0,
            "Server should be assigned a valid port number"
        );

        println!(
            "Server created successfully with address: {}",
            tcp_info.address
        );
    }

    /// The data-plane channel helper sizes the mpsc buffer from
    /// `send_buffer_count` — this is the value `process_request_stream` /
    /// `process_response_stream` feed it. `max_capacity()` reflects the
    /// channel's configured buffer, so a custom value and the default both
    /// reach the channel. Guards against regressing back to a hard-coded 64.
    #[test]
    fn data_plane_channel_capacity_matches_send_buffer_count() {
        let (tx, _rx) = data_plane_channel::<()>(7);
        assert_eq!(tx.max_capacity(), 7);

        let (tx, _rx) = data_plane_channel::<()>(DEFAULT_SEND_BUFFER_COUNT);
        assert_eq!(tx.max_capacity(), 64);

        // A misconfigured 0 must clamp to 1, not panic (mpsc::channel(0) panics).
        let (tx, _rx) = data_plane_channel::<()>(0);
        assert_eq!(tx.max_capacity(), 1);
    }

    /// `register` must thread `StreamOptions::send_buffer_count` through to the
    /// stored `RequestedSendConnection` / `RequestedRecvConnection` (the
    /// registration structs `process_*_stream` later destructure to size the
    /// channel). Verified here against the real registration path.
    #[tokio::test]
    async fn register_threads_send_buffer_count_into_connection_structs() {
        let server = TcpStreamServer::new(ServerOptions::default())
            .await
            .expect("server");
        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(true)
            .enable_response_stream(true)
            .send_buffer_count(7)
            .build()
            .unwrap();

        let _pending = server.register(options).await;

        let state = server.state.lock();
        assert_eq!(state.tx_subjects.len(), 1, "one request stream registered");
        assert_eq!(state.rx_subjects.len(), 1, "one response stream registered");
        assert!(
            state.tx_subjects.values().all(|c| c.send_buffer_count == 7),
            "send_buffer_count must reach RequestedSendConnection"
        );
        assert!(
            state.rx_subjects.values().all(|c| c.send_buffer_count == 7),
            "send_buffer_count must reach RequestedRecvConnection"
        );
    }

    #[tokio::test]
    async fn test_tcp_stream_server_fallback_to_loopback() {
        // Test fallback behavior using a mock resolver that always fails
        // This guarantees the fallback logic is triggered

        let options = ServerOptions::builder().port(0).build().unwrap();

        // Use the failing resolver to force the fallback
        let result = TcpStreamServer::new_with_resolver(options, StubResolver::not_found()).await;
        assert!(
            result.is_ok(),
            "Server creation should succeed with fallback even when IP detection fails"
        );

        let server = result.unwrap();

        let tcp_info = registered_tcp_info(&server).await;
        let socket_addr = tcp_info.address.parse::<std::net::SocketAddr>().unwrap();

        // With the failing resolver, fallback should ALWAYS be used
        let ip = socket_addr.ip();
        assert!(
            ip.is_loopback(),
            "Should use loopback when IP detection fails"
        );

        // Verify it's specifically 127.0.0.1 (the fallback value from the patch)
        assert_eq!(
            ip,
            std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
            "Fallback should use exactly 127.0.0.1, got: {}",
            ip
        );

        println!("SUCCESS: Fallback to 127.0.0.1 was confirmed: {}", ip);

        // The server should work with the fallback IP
        assert!(socket_addr.port() > 0, "Server should have a valid port");
    }

    #[tokio::test]
    async fn configured_ip_literals_bind_and_format_addresses() {
        let ipv6_available = std::net::TcpListener::bind("[::1]:0").is_ok();

        for (host, expected_ip) in [
            ("127.0.0.1", "127.0.0.1".parse::<IpAddr>().unwrap()),
            ("::1", "::1".parse::<IpAddr>().unwrap()),
            (" 127.0.0.1 ", "127.0.0.1".parse::<IpAddr>().unwrap()),
            (" [::1]\t", "::1".parse::<IpAddr>().unwrap()),
            ("::ffff:127.0.0.1", "127.0.0.1".parse::<IpAddr>().unwrap()),
            (" lo ", "127.0.0.1".parse::<IpAddr>().unwrap()),
            (" lo:1 ", "127.0.0.1".parse::<IpAddr>().unwrap()),
        ] {
            if expected_ip.is_ipv6() && !ipv6_available {
                eprintln!("skipping IPv6 bind because this host does not support IPv6 loopback");
                continue;
            }

            let mut resolver = StubResolver::not_found();
            resolver
                .interfaces
                .push(("lo:1", "127.0.0.1".parse().unwrap()));
            let server = TcpStreamServer::new_with_resolver(
                ServerOptions {
                    port: 0,
                    interface: Some(host.to_string()),
                },
                resolver,
            )
            .await
            .unwrap();
            let context = Context::new(());
            let pending = server
                .register(
                    StreamOptions::builder()
                        .context(context.context())
                        .enable_request_stream(false)
                        .enable_response_stream(true)
                        .build()
                        .unwrap(),
                )
                .await;
            let connection_info = pending.recv_stream.unwrap().connection_info;
            let tcp_info: TcpStreamConnectionInfo = connection_info.try_into().unwrap();
            let address = tcp_info.address.parse::<SocketAddr>().unwrap();

            assert_eq!(address, server.local_address().unwrap());
            assert_eq!(address.ip(), expected_ip);
            assert_ne!(address.port(), 0);
            if expected_ip.is_ipv6() {
                assert!(tcp_info.address.starts_with('['));
            }
        }
    }

    /// Create a test server using the failing IP resolver (falls back to loopback).
    async fn test_server() -> Arc<TcpStreamServer> {
        TcpStreamServer::new_with_resolver(
            ServerOptions::builder().port(0).build().unwrap(),
            StubResolver::not_found(),
        )
        .await
        .unwrap()
    }

    /// Helper: register a response stream and extract its subject string.
    async fn register_and_get_subject(
        server: &TcpStreamServer,
    ) -> (
        String,
        tokio::sync::oneshot::Receiver<Result<super::StreamReceiver, StreamPrologueError>>,
    ) {
        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();

        let pending = server.register(options).await;
        let recv_stream = pending.recv_stream.unwrap();
        let (conn_info, provider) = recv_stream.into_parts();
        let tcp_info: TcpStreamConnectionInfo = conn_info.try_into().unwrap();
        (tcp_info.subject, provider)
    }

    /// Convenience constructor so tests don't repeat the struct literal.
    fn make_eid(
        namespace: &str,
        component: &str,
        endpoint: &str,
        instance_id: u64,
    ) -> EndpointInstanceId {
        EndpointInstanceId {
            namespace: namespace.to_string(),
            component: component.to_string(),
            endpoint: endpoint.to_string(),
            instance_id,
        }
    }

    /// Helper: register a bidirectional pair (both request + response halves)
    /// and return both subjects + their providers.
    async fn register_and_get_bidi_subjects(
        server: &TcpStreamServer,
    ) -> (
        String,
        tokio::sync::oneshot::Receiver<Result<super::StreamSender, StreamPrologueError>>,
        String,
        tokio::sync::oneshot::Receiver<Result<super::StreamReceiver, StreamPrologueError>>,
    ) {
        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(true)
            .enable_response_stream(true)
            .build()
            .unwrap();

        let pending = server.register(options).await;
        let send_stream = pending.send_stream.unwrap();
        let recv_stream = pending.recv_stream.unwrap();
        let (send_info, send_provider) = send_stream.into_parts();
        let (recv_info, recv_provider) = recv_stream.into_parts();
        let send_tcp_info: TcpStreamConnectionInfo = send_info.try_into().unwrap();
        let recv_tcp_info: TcpStreamConnectionInfo = recv_info.try_into().unwrap();
        (
            send_tcp_info.subject,
            send_provider,
            recv_tcp_info.subject,
            recv_provider,
        )
    }

    /// `cancel_instance_streams` must drop the request-stream oneshot too,
    /// not just the response-stream one. Without the tagged tracker this test
    /// would hang on `send_provider.await` because the tx_subjects entry
    /// would leak past instance removal.
    #[tokio::test]
    async fn test_cancel_instance_streams_drops_both_bidi_halves() {
        let server = test_server().await;
        let (send_subj, send_provider, recv_subj, recv_provider) =
            register_and_get_bidi_subjects(&server).await;

        let id = make_eid("ns", "comp", "generate", 7);
        assert!(
            server
                .associate_instance(&recv_subj, Some(&send_subj), &id)
                .await,
            "fresh instance must not be tombstoned"
        );

        let cancelled = server.cancel_instance_streams(&id).await;
        assert_eq!(cancelled, 2, "both request + response halves must count");

        assert!(
            recv_provider.await.is_err(),
            "recv provider should resolve with RecvError"
        );
        assert!(
            send_provider.await.is_err(),
            "send provider should resolve with RecvError after instance cancellation"
        );
    }

    /// Pre-tombstoning an instance must drop both halves of a later
    /// `associate_instance(recv, Some(send), id)` call, not just the recv.
    #[tokio::test]
    async fn test_associate_instance_tombstone_cancels_both_bidi_halves() {
        let server = test_server().await;
        let id = make_eid("ns", "comp", "generate", 8);
        // Pre-tombstone the instance.
        server.cancel_instance_streams(&id).await;

        let (send_subj, send_provider, recv_subj, recv_provider) =
            register_and_get_bidi_subjects(&server).await;

        assert!(
            !server
                .associate_instance(&recv_subj, Some(&send_subj), &id)
                .await,
            "tombstoned instance must reject association"
        );

        assert!(recv_provider.await.is_err());
        assert!(send_provider.await.is_err());
    }

    #[tokio::test]
    async fn test_cancel_instance_streams_unblocks_receiver() {
        let server = test_server().await;

        let (subject, provider) = register_and_get_subject(&server).await;

        let id = make_eid("ns", "comp", "generate", 42);
        assert!(server.associate_instance(&subject, None, &id).await);

        let cancelled = server.cancel_instance_streams(&id).await;
        assert_eq!(cancelled, 1);

        // The oneshot receiver should now resolve with an error (sender dropped)
        let result = provider.await;
        assert!(result.is_err(), "Expected RecvError after cancellation");
    }

    #[tokio::test]
    async fn test_cancel_instance_streams_multiple_subjects() {
        let server = test_server().await;

        let (subj1, prov1) = register_and_get_subject(&server).await;
        let (subj2, prov2) = register_and_get_subject(&server).await;
        let (subj3, prov3) = register_and_get_subject(&server).await;

        let id10 = make_eid("ns", "comp", "generate", 10);
        let id20 = make_eid("ns", "comp", "generate", 20);

        // Associate first two with instance 10, third with instance 20
        assert!(server.associate_instance(&subj1, None, &id10).await);
        assert!(server.associate_instance(&subj2, None, &id10).await);
        assert!(server.associate_instance(&subj3, None, &id20).await);

        // Cancel instance 10 -- should cancel 2 subjects
        let cancelled = server.cancel_instance_streams(&id10).await;
        assert_eq!(cancelled, 2);

        assert!(prov1.await.is_err());
        assert!(prov2.await.is_err());

        // Instance 20 should be unaffected -- cancel it separately
        let cancelled = server.cancel_instance_streams(&id20).await;
        assert_eq!(cancelled, 1);
        assert!(prov3.await.is_err());
    }

    #[tokio::test]
    async fn test_cancel_instance_streams_nonexistent_instance() {
        let server = test_server().await;

        let id = make_eid("ns", "comp", "generate", 999);
        let cancelled = server.cancel_instance_streams(&id).await;
        assert_eq!(cancelled, 0);
    }

    #[tokio::test]
    async fn test_cancel_recv_stream_cleans_up_instance_tracking() {
        let server = test_server().await;

        let (subject, _provider) = register_and_get_subject(&server).await;
        let id = make_eid("ns", "comp", "generate", 42);
        assert!(server.associate_instance(&subject, None, &id).await);

        // Cancel the individual subject
        server.cancel_recv_stream(&subject).await;

        // Instance should have no remaining subjects
        let cancelled = server.cancel_instance_streams(&id).await;
        assert_eq!(
            cancelled, 0,
            "Instance tracking should have been cleaned up"
        );
    }

    #[tokio::test]
    async fn test_registered_stream_drop_runs_cleanup() {
        let server = test_server().await;

        // Register a response stream but DON'T call into_parts -- just drop it
        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();

        let pending = server.register(options).await;
        let recv_stream = pending.recv_stream.unwrap();

        // Get the subject before dropping
        let tcp_info: TcpStreamConnectionInfo =
            recv_stream.connection_info.clone().try_into().unwrap();
        let subject = tcp_info.subject.clone();

        // Verify it's in rx_subjects
        {
            let state = server.state.lock();
            assert!(state.rx_subjects.contains_key(&subject));
        }

        // Drop the RegisteredStream -- RAII cleanup should fire
        drop(recv_stream);

        // Give the spawned cleanup task a moment to run
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // Verify it's been removed from rx_subjects
        {
            let state = server.state.lock();
            assert!(
                !state.rx_subjects.contains_key(&subject),
                "RAII cleanup should have removed the rx_subjects entry"
            );
        }
    }

    #[tokio::test]
    async fn test_registered_stream_into_parts_disarms_cleanup() {
        let server = test_server().await;

        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();

        let pending = server.register(options).await;
        let recv_stream = pending.recv_stream.unwrap();

        let tcp_info: TcpStreamConnectionInfo =
            recv_stream.connection_info.clone().try_into().unwrap();
        let subject = tcp_info.subject.clone();

        // Call into_parts to disarm the cleanup
        let (_conn_info, _provider) = recv_stream.into_parts();

        // Give any potential cleanup a moment to run
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // The entry should still be in rx_subjects (cleanup was disarmed)
        {
            let state = server.state.lock();
            assert!(
                state.rx_subjects.contains_key(&subject),
                "into_parts() should disarm the RAII cleanup"
            );
        }
    }

    #[tokio::test]
    async fn test_dropped_registered_stream_wait_cleans_instance_indexes() {
        let server = test_server().await;
        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();

        let pending = server.register(options).await;
        let recv_stream = pending.recv_stream.unwrap();
        let tcp_info: TcpStreamConnectionInfo =
            recv_stream.connection_info.clone().try_into().unwrap();
        let subject = tcp_info.subject;
        let instance = make_eid("ns", "comp", "generate", 42);
        assert!(server.associate_instance(&subject, None, &instance).await);

        let mut wait = Box::pin(recv_stream.wait());
        assert!(futures::poll!(wait.as_mut()).is_pending());
        drop(wait);

        tokio::time::timeout(std::time::Duration::from_secs(1), async {
            loop {
                let cleaned = {
                    let state = server.state.lock();
                    !state.rx_subjects.contains_key(&subject)
                        && !state.rx_cancellations.contains_key(&subject)
                        && !state.subject_instance.contains_key(&subject)
                        && !state.instance_subjects.contains_key(&instance)
                };
                if cleaned {
                    break;
                }
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("cancelled wait did not clean response registration indexes");
    }

    #[tokio::test]
    async fn test_associate_after_cancel_is_immediately_cancelled() {
        // Simulates the race: cancel_instance_streams fires before associate_instance.
        let server = test_server().await;

        let id = make_eid("ns", "comp", "generate", 42);

        // Cancel BEFORE any subject is registered (tombstone).
        let cancelled = server.cancel_instance_streams(&id).await;
        assert_eq!(cancelled, 0);

        // Now register a subject and try to associate it with the tombstoned instance.
        let (subject, provider) = register_and_get_subject(&server).await;
        let associated = server.associate_instance(&subject, None, &id).await;

        // associate_instance should return false when the instance is tombstoned.
        assert!(
            !associated,
            "associate_instance on a tombstoned instance should return false"
        );

        // The provider should resolve with an error because associate_instance
        // found the tombstone and immediately cancelled the subject.
        let result = provider.await;
        assert!(
            result.is_err(),
            "Late associate_instance on a tombstoned instance should immediately cancel"
        );
    }

    #[tokio::test]
    async fn test_clear_tombstone_allows_new_associations() {
        let server = test_server().await;

        let id = make_eid("ns", "comp", "generate", 42);

        server.cancel_instance_streams(&id).await;
        server.clear_instance_tombstone(&id).await;

        // Now associate should work normally (subject NOT cancelled).
        let (subject, _provider) = register_and_get_subject(&server).await;
        assert!(server.associate_instance(&subject, None, &id).await);

        // Subject should be tracked, not cancelled.
        let cancelled = server.cancel_instance_streams(&id).await;
        assert_eq!(
            cancelled, 1,
            "After clearing tombstone, subjects should be tracked normally"
        );
    }

    #[tokio::test]
    async fn test_cancel_does_not_affect_sibling_endpoint() {
        // Regression: cancelling "generate" must not cancel "prefill" subjects
        // that share the same instance_id (same backend runtime).
        let server = test_server().await;

        let (gen_subj, gen_prov) = register_and_get_subject(&server).await;
        let (pre_subj, pre_prov) = register_and_get_subject(&server).await;

        let gen_id = make_eid("ns", "comp", "generate", 42);
        let pre_id = make_eid("ns", "comp", "prefill", 42);

        assert!(server.associate_instance(&gen_subj, None, &gen_id).await);
        assert!(server.associate_instance(&pre_subj, None, &pre_id).await);

        // Cancel only the "generate" endpoint's subjects.
        let cancelled = server.cancel_instance_streams(&gen_id).await;
        assert_eq!(
            cancelled, 1,
            "Only the generate subject should be cancelled"
        );
        assert!(gen_prov.await.is_err());

        // prefill must still be tracked.
        let still_pending = server.cancel_instance_streams(&pre_id).await;
        assert_eq!(still_pending, 1, "prefill subject should still be tracked");
        assert!(pre_prov.await.is_err());
    }

    #[tokio::test]
    async fn test_tombstone_is_endpoint_scoped() {
        // Tombstoning "generate" must not prevent new associations on "prefill"
        // for the same instance_id.
        let server = test_server().await;

        let gen_id = make_eid("ns", "comp", "generate", 42);
        let pre_id = make_eid("ns", "comp", "prefill", 42);

        server.cancel_instance_streams(&gen_id).await;

        // A new subject for "generate" should be rejected.
        let (gen_subj, gen_prov) = register_and_get_subject(&server).await;
        assert!(
            !server.associate_instance(&gen_subj, None, &gen_id).await,
            "generate should be tombstoned"
        );
        assert!(gen_prov.await.is_err());

        // A new subject for "prefill" with the same instance_id should be accepted.
        let (pre_subj, _pre_prov) = register_and_get_subject(&server).await;
        assert!(
            server.associate_instance(&pre_subj, None, &pre_id).await,
            "prefill tombstone is independent; subject should be tracked"
        );
        let count = server.cancel_instance_streams(&pre_id).await;
        assert_eq!(count, 1, "prefill subject should be tracked normally");
    }

    #[tokio::test]
    async fn test_cancel_does_not_affect_different_component() {
        // Regression: two services with different (namespace, component) but the
        // same endpoint name and the same pod-backed instance_id must not interfere,
        // even though they share a single TcpStreamServer runtime.
        let server = test_server().await;

        let (subj_a, prov_a) = register_and_get_subject(&server).await;
        let (subj_b, prov_b) = register_and_get_subject(&server).await;

        // Same endpoint name + instance_id, different namespace/component.
        let id_a = make_eid("ns-a", "comp-a", "generate", 42);
        let id_b = make_eid("ns-b", "comp-b", "generate", 42);

        assert!(server.associate_instance(&subj_a, None, &id_a).await);
        assert!(server.associate_instance(&subj_b, None, &id_b).await);

        // Cancel service A -- only subj_a should be affected.
        let cancelled = server.cancel_instance_streams(&id_a).await;
        assert_eq!(cancelled, 1, "Only service-A subject should be cancelled");
        assert!(prov_a.await.is_err());

        // Service B subject must still be pending.
        let still_tracked = server.cancel_instance_streams(&id_b).await;
        assert_eq!(still_tracked, 1, "Service-B subject should be unaffected");
        assert!(prov_b.await.is_err());
    }

    #[tokio::test(start_paused = true)]
    async fn test_tombstone_expires_after_ttl() {
        // After TOMBSTONE_TTL elapses, a previously-tombstoned identity must
        // accept new associations again, AND the entry must be physically
        // pruned from `removed_instances` so the set remains bounded.
        let server = test_server().await;

        let id = make_eid("ns", "comp", "generate", 42);

        // Tombstone the identity.
        server.cancel_instance_streams(&id).await;
        {
            let state = server.state.lock();
            assert!(state.removed_instances.contains_key(&id));
        }

        // Advance past the TTL.
        tokio::time::advance(TOMBSTONE_TTL + Duration::from_secs(1)).await;

        // associate_instance for the same identity should now succeed (no
        // longer tombstoned). Any new subject must be tracked normally.
        let (subject, _provider) = register_and_get_subject(&server).await;
        assert!(
            server.associate_instance(&subject, None, &id).await,
            "tombstone older than TTL should not block association"
        );

        // The expired tombstone must have been pruned (lazy pruning fires on
        // every associate_instance/cancel_instance_streams call).
        {
            let state = server.state.lock();
            assert!(
                !state.removed_instances.contains_key(&id),
                "expired tombstone should be pruned, not retained"
            );
        }
    }

    #[tokio::test(start_paused = true)]
    async fn test_tombstone_within_ttl_blocks_associate() {
        // Regression net for the original tombstone fix: a tombstone younger
        // than TTL must still cancel late-arriving associate_instance() calls.
        let server = test_server().await;

        let id = make_eid("ns", "comp", "generate", 42);
        server.cancel_instance_streams(&id).await;

        // Advance only a small fraction of the TTL.
        tokio::time::advance(Duration::from_secs(1)).await;

        let (subject, provider) = register_and_get_subject(&server).await;
        assert!(
            !server.associate_instance(&subject, None, &id).await,
            "tombstone within TTL must still block association"
        );
        assert!(provider.await.is_err());
    }

    #[tokio::test(start_paused = true)]
    async fn test_tombstone_lazy_prune_on_cancel() {
        // Old tombstones must be pruned on the next cancel_instance_streams
        // call, regardless of which identity is being tombstoned.
        let server = test_server().await;

        let id_old = make_eid("ns", "comp", "generate", 1);
        let id_new = make_eid("ns", "comp", "generate", 2);

        server.cancel_instance_streams(&id_old).await;
        tokio::time::advance(TOMBSTONE_TTL + Duration::from_secs(1)).await;
        server.cancel_instance_streams(&id_new).await;

        let state = server.state.lock();
        assert!(
            !state.removed_instances.contains_key(&id_old),
            "old tombstone should be pruned by the next cancel_instance_streams call"
        );
        assert!(
            state.removed_instances.contains_key(&id_new),
            "fresh tombstone should be retained"
        );
        assert_eq!(state.removed_instances.len(), 1);
    }

    #[tokio::test]
    async fn test_clear_tombstone_only_affects_named_identity() {
        // Documents the monotonic-lease invariant: `clear_instance_tombstone`
        // for one EndpointInstanceId must not touch a sibling entry. With etcd
        // lease IDs this defensive code rarely fires (new lease = new
        // EndpointInstanceId), but the per-key scope must hold.
        let server = test_server().await;

        let id_a = make_eid("ns", "comp", "generate", 1);
        let id_b = make_eid("ns", "comp", "generate", 2);

        server.cancel_instance_streams(&id_a).await;
        server.clear_instance_tombstone(&id_b).await;

        let state = server.state.lock();
        assert!(
            state.removed_instances.contains_key(&id_a),
            "clearing a different identity must not remove id_a's tombstone"
        );
    }

    #[tokio::test]
    async fn test_tombstone_scoped_to_full_identity() {
        // A tombstone on (ns-a, comp-a, generate, 42) must not block
        // associations on (ns-b, comp-b, generate, 42).
        let server = test_server().await;

        let id_a = make_eid("ns-a", "comp-a", "generate", 42);
        let id_b = make_eid("ns-b", "comp-b", "generate", 42);

        // Tombstone only service A.
        server.cancel_instance_streams(&id_a).await;

        // Service A is tombstoned — new association is rejected.
        let (subj_a, prov_a) = register_and_get_subject(&server).await;
        assert!(!server.associate_instance(&subj_a, None, &id_a).await);
        assert!(prov_a.await.is_err());

        // Service B with same endpoint name + instance_id must be accepted.
        let (subj_b, _prov_b) = register_and_get_subject(&server).await;
        assert!(
            server.associate_instance(&subj_b, None, &id_b).await,
            "Different namespace/component must not be tombstoned"
        );
        assert_eq!(server.cancel_instance_streams(&id_b).await, 1);
    }

    type TestFramedRead = FramedRead<ReadHalf<TcpStream>, TwoPartCodec>;
    type TestFramedWrite = FramedWrite<WriteHalf<TcpStream>, TwoPartCodec>;
    type TestResponseStream = (TestFramedRead, TestFramedWrite, StreamReceiver);

    /// Stand up a TcpStreamServer, register a response stream, connect a
    /// client, drive the handshake + prologue, and return the client-side
    /// framed reader/writer along with the receiver.
    async fn open_registered_response_stream() -> TestResponseStream {
        let options = ServerOptions::builder().port(0).build().unwrap();
        let server = TcpStreamServer::new_with_resolver(options, StubResolver::not_found())
            .await
            .unwrap();
        let context = Context::new(());
        let stream_options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();
        let pending_connection = server.register(stream_options).await;
        let registered_stream = pending_connection.recv_stream.unwrap();
        let (connection_info, stream_provider) = registered_stream.into_parts();
        let tcp_info: TcpStreamConnectionInfo = connection_info.try_into().unwrap();

        let stream = TcpStream::connect(&tcp_info.address).await.unwrap();
        let (read_half, write_half) = tokio::io::split(stream);
        let framed_reader = FramedRead::new(read_half, TwoPartCodec::default());
        let mut framed_writer = FramedWrite::new(write_half, TwoPartCodec::default());

        let handshake = CallHomeHandshake {
            subject: tcp_info.subject,
            stream_type: StreamType::Response,
        };
        framed_writer
            .send(TwoPartMessage::from_header(
                serde_json::to_vec(&handshake).unwrap().into(),
            ))
            .await
            .unwrap();
        framed_writer
            .send(TwoPartMessage::from_header(
                serde_json::to_vec(&ResponseStreamPrologue {
                    error: None,
                    typed_error: None,
                })
                .unwrap()
                .into(),
            ))
            .await
            .unwrap();

        // SAFETY (test-only): healthy localhost handshake always resolves all
        // three layers; a panic here means the harness is broken.
        let receiver = tokio::time::timeout(std::time::Duration::from_secs(1), stream_provider)
            .await
            .expect("server should establish response stream within timeout")
            .expect("stream provider should not be dropped")
            .expect("response stream should be accepted");

        (framed_reader, framed_writer, receiver)
    }

    async fn recv_control_message(framed_reader: &mut TestFramedRead) -> ControlMessage {
        // SAFETY (test-only): a misbehaving server in any of these layers is
        // exactly the harness failure we want surfaced as a test panic.
        let message = tokio::time::timeout(std::time::Duration::from_secs(1), framed_reader.next())
            .await
            .expect("server should send a control message within timeout")
            .expect("server should not close before sending control")
            .expect("control message should decode");
        let (header, data) = message.optional_parts();
        assert!(data.is_none(), "control message should not contain data");
        serde_json::from_slice(header.expect("control header missing").as_ref()).unwrap()
    }

    /// Sending an unexpected control message (Stop or Kill from the data
    /// direction) is a protocol violation. The server's
    /// network_receive_handler must reply with ControlMessage::Kill on
    /// that stream alone, not panic.
    #[tokio::test]
    async fn test_tcp_stream_server_sends_kill_on_unexpected_control_message() {
        let (mut framed_reader, mut framed_writer, _receiver) =
            open_registered_response_stream().await;

        framed_writer
            .send(TwoPartMessage::from_header(
                serde_json::to_vec(&ControlMessage::Stop).unwrap().into(),
            ))
            .await
            .unwrap();

        assert_eq!(
            recv_control_message(&mut framed_reader).await,
            ControlMessage::Kill,
            "unexpected control message should kill only this stream"
        );
    }

    #[tokio::test]
    async fn test_tcp_stream_server_sends_kill_when_prologue_wait_is_cancelled() {
        let server = test_server().await;
        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();
        let pending = server.register(options).await;
        let registered_stream = pending.recv_stream.unwrap();
        let (connection_info, stream_provider) = registered_stream.into_parts();
        let tcp_info: TcpStreamConnectionInfo = connection_info.try_into().unwrap();
        let subject = tcp_info.subject.clone();
        let instance = make_eid("ns", "comp", "generate", 42);
        assert!(server.associate_instance(&subject, None, &instance).await);

        let stream = TcpStream::connect(&tcp_info.address).await.unwrap();
        let (read_half, write_half) = tokio::io::split(stream);
        let mut framed_reader = FramedRead::new(read_half, TwoPartCodec::default());
        let mut framed_writer = FramedWrite::new(write_half, TwoPartCodec::default());
        let handshake = CallHomeHandshake {
            subject: subject.clone(),
            stream_type: StreamType::Response,
        };
        framed_writer
            .send(TwoPartMessage::from_header(
                serde_json::to_vec(&handshake).unwrap().into(),
            ))
            .await
            .unwrap();

        tokio::time::timeout(std::time::Duration::from_secs(1), async {
            while server.state.lock().rx_subjects.contains_key(&subject) {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("server did not accept response call-home");

        assert_eq!(server.cancel_instance_streams(&instance).await, 1);
        assert!(
            time::timeout(Duration::from_secs(1), stream_provider)
                .await
                .expect("worker removal must unblock the response handshake")
                .is_err()
        );
        assert_eq!(
            recv_control_message(&mut framed_reader).await,
            ControlMessage::Kill,
            "cancelling the prologue wait should stop worker generation"
        );
        let state = server.state.lock();
        assert!(!state.rx_subjects.contains_key(&subject));
        assert!(!state.rx_cancellations.contains_key(&subject));
        assert!(!state.subject_instance.contains_key(&subject));
        assert!(!state.instance_subjects.contains_key(&instance));
    }

    /// Cover both orderings at the boundary between a pending handshake and an
    /// established response. Reading a prologue alone must not defeat removal.
    #[tokio::test]
    async fn response_prologue_acceptance_is_ordered_with_worker_removal() {
        let server = test_server().await;
        for removal_first in [true, false] {
            let instance = make_eid("ns", "comp", "generate", u64::from(removal_first));
            let (subject, _provider) = register_and_get_subject(&server).await;
            assert!(server.associate_instance(&subject, None, &instance).await);
            let pending = TcpStreamServer::take_response_stream(&server.state, &subject).unwrap();
            let guard = ResponseRegistrationGuard {
                state: server.state.clone(),
                subject: subject.clone(),
                is_active: true,
            };

            if removal_first {
                assert_eq!(server.cancel_instance_streams(&instance).await, 1);
                assert!(
                    !guard.accept(),
                    "a cancelled handshake must not be accepted"
                );
                assert!(pending.cancellation.is_cancelled());
            } else {
                assert!(guard.accept());
                assert_eq!(server.cancel_instance_streams(&instance).await, 0);
                assert!(!pending.cancellation.is_cancelled());
            }

            let state = server.state.lock();
            assert!(!state.rx_cancellations.contains_key(&subject));
            assert!(!state.subject_instance.contains_key(&subject));
            assert!(!state.instance_subjects.contains_key(&instance));
        }
    }

    /// A peer that never reads Kill must not delay the requester or retain the
    /// handshake registration. Its socket cleanup must also have a deadline.
    #[tokio::test(start_paused = true)]
    async fn cancelled_handshake_unblocks_before_stalled_cleanup() {
        let server = test_server().await;
        let instance = make_eid("ns", "comp", "generate", 42);
        let (subject, provider) = register_and_get_subject(&server).await;
        assert!(server.associate_instance(&subject, None, &instance).await);
        let pending = TcpStreamServer::take_response_stream(&server.state, &subject).unwrap();
        let guard = ResponseRegistrationGuard {
            state: server.state.clone(),
            subject: subject.clone(),
            is_active: true,
        };
        // One byte cannot hold a framed Kill; retaining the unread peer stalls send().
        let (socket, _peer) = tokio::io::duplex(1);
        let writer = FramedWrite::new(Box::new(socket) as BoxWrite, TwoPartCodec::default());
        let cleanup = tokio::spawn(cancel_response_handshake(
            pending.connection,
            Some(guard),
            writer,
        ));

        assert!(
            time::timeout(Duration::from_millis(50), provider)
                .await
                .expect("requester must be released before network cleanup")
                .is_err()
        );
        assert!(
            !cleanup.is_finished(),
            "peer should still be blocking the Kill write"
        );
        {
            let state = server.state.lock();
            assert!(!state.rx_cancellations.contains_key(&subject));
            assert!(!state.subject_instance.contains_key(&subject));
            assert!(!state.instance_subjects.contains_key(&instance));
        }
        time::advance(HANDSHAKE_CANCEL_CLEANUP_TIMEOUT).await;
        time::timeout(Duration::from_millis(50), cleanup)
            .await
            .expect("network cleanup must stop at its deadline")
            .unwrap();
    }

    /// A framing/decode error from the worker side is unrecoverable for
    /// this stream but must not panic the worker. Server should send Kill
    /// and tear down only this connection.
    #[tokio::test]
    async fn test_tcp_stream_server_sends_kill_on_read_error() {
        let (mut framed_reader, framed_writer, _receiver) = open_registered_response_stream().await;

        let mut raw_writer = framed_writer.into_inner();
        raw_writer.write_all(&[0u8; 8]).await.unwrap();
        raw_writer.shutdown().await.unwrap();

        assert_eq!(
            recv_control_message(&mut framed_reader).await,
            ControlMessage::Kill,
            "framing read error should kill only this stream"
        );
    }

    /// Sentinel is supposed to be header-only. A misbehaving client that
    /// attaches a data payload must not panic the worker via assert!().
    #[tokio::test]
    async fn test_tcp_stream_server_sends_kill_on_sentinel_with_data() {
        let (mut framed_reader, mut framed_writer, _receiver) =
            open_registered_response_stream().await;

        let header = serde_json::to_vec(&ControlMessage::Sentinel)
            .unwrap()
            .into();
        framed_writer
            .send(TwoPartMessage::from_parts(
                header,
                Bytes::from_static(b"unexpected payload"),
            ))
            .await
            .unwrap();

        assert_eq!(
            recv_control_message(&mut framed_reader).await,
            ControlMessage::Kill,
            "Sentinel with data should kill only this stream"
        );
    }

    /// The prologue must be a HeaderOnly frame. A non-HeaderOnly prologue
    /// (data-only or mixed) must surface as Err to the requester rather
    /// than panic the worker.
    #[tokio::test]
    async fn test_tcp_stream_server_returns_error_on_invalid_prologue() {
        let options = ServerOptions::builder().port(0).build().unwrap();
        let server = TcpStreamServer::new_with_resolver(options, StubResolver::not_found())
            .await
            .unwrap();
        let context = Context::new(());
        let stream_options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();
        let pending_connection = server.register(stream_options).await;
        let registered_stream = pending_connection.recv_stream.unwrap();
        let (connection_info, stream_provider) = registered_stream.into_parts();
        let tcp_info: TcpStreamConnectionInfo = connection_info.try_into().unwrap();

        let stream = TcpStream::connect(&tcp_info.address).await.unwrap();
        let (_read_half, write_half) = tokio::io::split(stream);
        let mut framed_writer = FramedWrite::new(write_half, TwoPartCodec::default());

        let handshake = CallHomeHandshake {
            subject: tcp_info.subject,
            stream_type: StreamType::Response,
        };
        framed_writer
            .send(TwoPartMessage::from_header(
                serde_json::to_vec(&handshake).unwrap().into(),
            ))
            .await
            .unwrap();

        // Send a data-only frame in the prologue slot.
        framed_writer
            .send(TwoPartMessage::from_data(Bytes::from_static(
                b"not a prologue",
            )))
            .await
            .unwrap();

        let outcome = tokio::time::timeout(std::time::Duration::from_secs(1), stream_provider)
            .await
            .expect("stream provider should resolve quickly")
            .expect("stream provider channel should not be dropped");
        // StreamReceiver doesn't impl Debug, so we can't use `.expect_err`.
        match outcome {
            Err(err) => assert!(
                err.contains("malformed prologue"),
                "expected malformed-prologue error, got: {err}"
            ),
            Ok(_) => panic!("invalid prologue should produce an error, but got Ok"),
        }
    }

    /// A prologue from a newer worker may carry an `ErrorType` this build does
    /// not know; its legacy error must still reach the requester.
    #[tokio::test]
    async fn test_unknown_typed_error_preserves_the_legacy_prologue_error() {
        let options = ServerOptions::builder().port(0).build().unwrap();
        let server = TcpStreamServer::new_with_resolver(options, StubResolver::not_found())
            .await
            .unwrap();
        let context = Context::new(());
        let stream_options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();
        let pending_connection = server.register(stream_options).await;
        let (connection_info, stream_provider) =
            pending_connection.recv_stream.unwrap().into_parts();
        let tcp_info: TcpStreamConnectionInfo = connection_info.try_into().unwrap();

        let stream = TcpStream::connect(&tcp_info.address).await.unwrap();
        let (_read_half, write_half) = tokio::io::split(stream);
        let mut framed_writer = FramedWrite::new(write_half, TwoPartCodec::default());

        let handshake = CallHomeHandshake {
            subject: tcp_info.subject,
            stream_type: StreamType::Response,
        };
        framed_writer
            .send(TwoPartMessage::from_header(
                serde_json::to_vec(&handshake).unwrap().into(),
            ))
            .await
            .unwrap();

        // Correctly framed HeaderOnly, but the typed error names a variant this
        // build does not know.
        framed_writer
            .send(TwoPartMessage::from_header(Bytes::from_static(
                br#"{"error":"Generate Error: boom","typed_error":{"error_type":"VariantFromTheFuture","message":"boom"}}"#,
            )))
            .await
            .unwrap();

        let outcome = tokio::time::timeout(std::time::Duration::from_secs(5), stream_provider)
            .await
            .expect("stream provider should resolve quickly")
            .expect("the oneshot must be notified, not dropped");

        // `StreamReceiver` is not `Debug`, so match instead of `expect_err`.
        match outcome {
            Err(err) => {
                assert_eq!(err.message, "Generate Error: boom");
                assert!(err.typed_error.is_none());
            }
            Ok(_) => panic!("an error prologue must not yield a usable stream"),
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn test_concurrent_response_registration_and_call_home() {
        const STREAMS: usize = 128;

        let result = time::timeout(Duration::from_secs(20), async {
            let server = test_server().await;
            let mut pending_streams = Vec::with_capacity(STREAMS);
            let mut client_tasks = Vec::with_capacity(STREAMS);

            for idx in 0..STREAMS {
                let context = Context::new(());
                let options = StreamOptions::builder()
                    .context(context.context())
                    .enable_request_stream(false)
                    .enable_response_stream(true)
                    .build()
                    .unwrap();

                let pending = server.register(options).await;
                let registered_stream = pending.recv_stream.unwrap();
                let (connection_info, stream_provider) = registered_stream.into_parts();
                let client_context =
                    Context::with_id_and_metadata((), context.id().to_string(), Default::default());
                let payload = Bytes::from(format!("payload-{idx}"));

                pending_streams.push((idx, payload.clone(), stream_provider));
                client_tasks.push(tokio::spawn(async move {
                    let mut sender = TcpClient::create_response_stream(
                        client_context.context(),
                        connection_info,
                        None,
                    )
                    .await
                    .unwrap();
                    sender.send_prologue(None).await.unwrap();
                    sender.send(payload).await.unwrap();
                }));
            }

            for task in client_tasks {
                task.await.unwrap();
            }

            for (idx, expected, stream_provider) in pending_streams {
                let mut stream = stream_provider.await.unwrap().unwrap();
                let actual = stream.rx.recv().await.unwrap();
                assert_eq!(actual, expected, "payload mismatch for stream {idx}");
            }
        })
        .await;

        assert!(
            result.is_ok(),
            "concurrent response registration and call-home timed out"
        );
    }

    #[tokio::test]
    async fn test_response_stream_cancellation_before_prologue() {
        temp_env::async_with_vars(
            [
                ("DYN_TCP_TLS_CERT_PATH", None::<&str>),
                ("DYN_TCP_TLS_KEY_PATH", None),
                ("DYN_TCP_TLS_CLIENT_CA_CERT_PATH", None),
                ("DYN_TCP_TLS_CA_CERT_PATH", None),
                ("DYN_TCP_TLS_INSECURE", None),
                ("DYN_TCP_TLS_CLIENT_CERT_PATH", None),
                ("DYN_TCP_TLS_CLIENT_KEY_PATH", None),
            ],
            async {
                let server = test_server().await;
                for action in [Some(ControlMessage::Stop), Some(ControlMessage::Kill), None] {
                    tokio::time::timeout(Duration::from_secs(5), async {
                        let context = Context::new(()).context();
                        let options = StreamOptions::builder()
                            .context(context.clone())
                            .enable_request_stream(false)
                            .enable_response_stream(true)
                            .build()
                            .unwrap();
                        let pending = server.register(options).await;
                        let (connection_info, stream_provider) =
                            pending.recv_stream.unwrap().into_parts();
                        let worker_context = Context::with_id_and_metadata(
                            (),
                            context.id().to_string(),
                            Default::default(),
                        )
                        .context();
                        let sender = TcpClient::create_response_stream(
                            worker_context.clone(),
                            connection_info,
                            None,
                        )
                        .await
                        .unwrap();

                        let stream_provider = match action {
                            Some(ControlMessage::Stop) => {
                                context.stop_generating();
                                Some(stream_provider)
                            }
                            Some(ControlMessage::Kill) => {
                                context.kill();
                                Some(stream_provider)
                            }
                            None => {
                                drop(stream_provider);
                                assert!(!context.is_stopped());
                                None
                            }
                            Some(ControlMessage::Sentinel) => unreachable!(),
                        };
                        if let Some(stream_provider) = stream_provider {
                            let failure = match stream_provider.await.unwrap() {
                                Err(failure) => failure,
                                Ok(_) => {
                                    panic!("cancellation must not establish a response stream")
                                }
                            };
                            assert_eq!(
                                failure.typed_error.as_ref().map(|error| error.error_type()),
                                Some(ErrorType::Cancelled)
                            );
                        }
                        worker_context.stopped().await;
                        sender.tx.closed().await;
                        assert_eq!(
                            worker_context.is_killed(),
                            action != Some(ControlMessage::Stop)
                        );
                    })
                    .await
                    .unwrap_or_else(|_| {
                        panic!("cancellation before prologue timed out: {action:?}")
                    });
                }
            },
        )
        .await;
    }

    /// A worker that refuses a request before producing any response bytes must
    /// keep its error type all the way to the requesting side.
    #[tokio::test]
    async fn test_typed_prologue_error_survives_to_requester() {
        let server = test_server().await;
        let context = Context::new(());
        let options = StreamOptions::builder()
            .context(context.context())
            .enable_request_stream(false)
            .enable_response_stream(true)
            .build()
            .unwrap();

        let pending = server.register(options).await;
        let (connection_info, stream_provider) = pending.recv_stream.unwrap().into_parts();
        let client_context =
            Context::with_id_and_metadata((), context.id().to_string(), Default::default());

        let worker_error = DynamoError::builder()
            .error_type(ErrorType::Backend(BackendError::InvalidArgument))
            .message("multimodal input is not supported by this backend")
            .build();

        let mut sender =
            TcpClient::create_response_stream(client_context.context(), connection_info, None)
                .await
                .unwrap();
        sender
            .send_prologue_typed(Some(StreamPrologueError::new(
                "Generate Error: multimodal input is not supported by this backend",
                worker_error,
            )))
            .await
            .unwrap();

        let outcome = tokio::time::timeout(std::time::Duration::from_secs(5), stream_provider)
            .await
            .expect("stream provider should resolve quickly")
            .expect("stream provider channel should not be dropped");

        // `StreamReceiver` is not `Debug`, so match instead of `expect_err`.
        let prologue_error = match outcome {
            Err(err) => err,
            Ok(_) => panic!("an error prologue must not yield a usable stream"),
        };
        assert_eq!(
            prologue_error.typed_error.as_ref().map(|e| e.error_type()),
            Some(ErrorType::Backend(BackendError::InvalidArgument)),
            "the worker's error type must survive the prologue round trip"
        );
    }

    // ==================== request_stream_send_handler integration tests ====================
    //
    // These exercise the closing-message contract of `request_stream_send_handler`
    // end-to-end: register a request stream, dial it as a raw client (so we can
    // inspect frames directly), then trigger each of the exit branches and
    // assert which ControlMessage arrives on the wire.

    use futures::SinkExt;

    /// Register a request stream and dial it with a raw client. Returns the
    /// framed reader on the raw client side, the StreamSender held by the
    /// upstream, and the upstream's engine context (so the test can drive
    /// kill / stop externally).
    async fn register_and_dial_request_stream(
        server: &TcpStreamServer,
    ) -> (
        FramedRead<tokio::io::ReadHalf<TcpStream>, TwoPartCodec>,
        super::StreamSender,
        Arc<dyn AsyncEngineContext>,
    ) {
        let upstream_ctx = Context::new(()).context();
        let options = StreamOptions::builder()
            .context(upstream_ctx.clone())
            .enable_request_stream(true)
            .enable_response_stream(false)
            .build()
            .unwrap();

        let pending = server.register(options).await;
        let send_stream = pending.send_stream.unwrap();
        let (conn_info, send_provider) = send_stream.into_parts();
        let tcp_info: TcpStreamConnectionInfo = conn_info.try_into().unwrap();

        let raw = TcpStream::connect(&tcp_info.address).await.unwrap();
        let (read_half, write_half) = tokio::io::split(raw);
        let framed_reader = FramedRead::new(read_half, TwoPartCodec::default());
        let mut framed_writer = FramedWrite::new(write_half, TwoPartCodec::default());

        let handshake = super::CallHomeHandshake {
            subject: tcp_info.subject.clone(),
            stream_type: StreamType::Request,
        };
        let handshake_bytes = serde_json::to_vec(&handshake).unwrap();
        framed_writer
            .send(TwoPartMessage::from_header(handshake_bytes.into()))
            .await
            .unwrap();
        drop(framed_writer);

        let sender = send_provider.await.unwrap().unwrap();
        (framed_reader, sender, upstream_ctx)
    }

    /// Pull frames off the raw client reader until the first `ControlMessage`
    /// arrives, ignoring any DataOnly frames before it. Returns the variant.
    async fn next_control_message(
        reader: &mut FramedRead<tokio::io::ReadHalf<TcpStream>, TwoPartCodec>,
    ) -> ControlMessage {
        loop {
            let frame = reader
                .next()
                .await
                .expect("socket closed before control message arrived")
                .expect("decode error");
            if let Some(header) = frame.header() {
                return serde_json::from_slice::<ControlMessage>(header)
                    .expect("invalid control message bytes");
            }
            // DataOnly frame — skip and keep reading.
        }
    }

    /// Dropping the upstream's StreamSender drains `request_rx` and the server
    /// emits [`ControlMessage::Sentinel`] as the closing frame.
    #[tokio::test]
    async fn test_request_stream_sends_sentinel_on_clean_drop() {
        let server = test_server().await;
        let (mut reader, sender, _ctx) = register_and_dial_request_stream(&server).await;

        drop(sender);

        let ctrl = next_control_message(&mut reader).await;
        assert!(
            matches!(ctrl, ControlMessage::Sentinel),
            "clean drain should emit Sentinel, got {ctrl:?}"
        );
    }

    /// `context.kill()` makes the server emit [`ControlMessage::Kill`] before
    /// shutting down the write half.
    #[tokio::test]
    async fn test_request_stream_sends_kill_on_context_killed() {
        let server = test_server().await;
        let (mut reader, _sender, ctx) = register_and_dial_request_stream(&server).await;

        ctx.kill();

        let ctrl = next_control_message(&mut reader).await;
        assert!(
            matches!(ctrl, ControlMessage::Kill),
            "context.kill() should emit Kill, got {ctrl:?}"
        );
    }

    /// `context.stop()` makes the server emit [`ControlMessage::Stop`] before
    /// shutting down the write half.
    #[tokio::test]
    async fn test_request_stream_sends_stop_on_context_stopped() {
        let server = test_server().await;
        let (mut reader, _sender, ctx) = register_and_dial_request_stream(&server).await;

        ctx.stop();

        let ctrl = next_control_message(&mut reader).await;
        assert!(
            matches!(ctrl, ControlMessage::Stop),
            "context.stop() should emit Stop, got {ctrl:?}"
        );
    }

    // ---- accept-loop backoff under resource exhaustion (issue #11822) ----

    #[cfg(unix)]
    fn emfile_error() -> std::io::Error {
        std::io::Error::from_raw_os_error(libc::EMFILE)
    }

    #[cfg(unix)]
    #[test]
    fn accept_backoff_classifies_all_exhaustion_errnos() {
        for errno in [libc::EMFILE, libc::ENFILE, libc::ENOBUFS, libc::ENOMEM] {
            let err = std::io::Error::from_raw_os_error(errno);
            assert_eq!(
                AcceptBackoff::classify(&err),
                AcceptFailure::Exhaustion,
                "errno {errno} must classify as exhaustion"
            );
        }
        // Ordinary errors stay on the immediate-retry path.
        let ordinary = std::io::Error::from(std::io::ErrorKind::ConnectionAborted);
        assert_eq!(AcceptBackoff::classify(&ordinary), AcceptFailure::Ordinary,);
    }

    #[cfg(unix)]
    #[test]
    fn accept_backoff_grows_resets_and_rate_limits() {
        let mut backoff = AcceptBackoff::default();
        let now = std::time::Instant::now();

        // Delay doubles per consecutive failure and saturates at the ceiling.
        let delays: Vec<Duration> = (0..12)
            .map(|_| backoff.record_exhaustion(now).delay)
            .collect();
        assert!(delays[0] > Duration::ZERO);
        let sat = delays
            .iter()
            .position(|d| *d == ACCEPT_BACKOFF_MAX_DELAY)
            .expect("delay reaches the ceiling");
        for pair in delays[..=sat].windows(2) {
            assert!(
                pair[1] > pair[0],
                "delay must grow: {:?} -> {:?}",
                pair[0],
                pair[1]
            );
        }
        assert!(delays[sat..].iter().all(|d| *d == ACCEPT_BACKOFF_MAX_DELAY));

        // A successful accept resets the schedule to the initial delay.
        backoff.record_success(|| now);
        assert_eq!(
            backoff.record_exhaustion(now).delay,
            ACCEPT_BACKOFF_INITIAL_DELAY,
        );

        // The summary is rate-limited: one emission per interval, carrying the
        // count of failures suppressed since the previous one.
        let t0 = std::time::Instant::now();
        let mut backoff = AcceptBackoff::default();
        let emitted: Vec<Option<u64>> = (0..50)
            .map(|_| backoff.record_exhaustion(t0).log_suppressed)
            .collect();
        assert_eq!(emitted.iter().filter(|e| e.is_some()).count(), 1);
        assert_eq!(emitted[0], Some(0));

        let t1 = t0 + ACCEPT_BACKOFF_LOG_INTERVAL + Duration::from_millis(1);
        assert_eq!(
            backoff.record_exhaustion(t1).log_suppressed,
            Some(49),
            "next emission reports every failure suppressed since the last one",
        );

        // Recovery shares the rate-limit window with the failure summary, so a
        // flapping listener cannot emit a recovery line per accept.
        let mut backoff = AcceptBackoff::default();
        let t0 = std::time::Instant::now();
        assert_eq!(backoff.record_exhaustion(t0).log_suppressed, Some(0));
        backoff.record_exhaustion(t0);
        backoff.record_exhaustion(t0);
        assert_eq!(
            backoff.record_success(|| t0),
            None,
            "recovery inside the log window must not emit",
        );
        backoff.record_exhaustion(t0);
        assert_eq!(
            backoff.record_success(|| t0 + ACCEPT_BACKOFF_LOG_INTERVAL),
            Some(3),
            "after the window elapses the recovery emits with the rolled-forward count",
        );
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn accept_backoff_socket_recovers_after_injected_exhaustion() {
        use crate::metrics::transport_metrics::TCP_ACCEPT_BACKOFF_TOTAL;

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .expect("bind ephemeral listener");
        let addr = listener.local_addr().expect("local_addr");

        let mut backoff = AcceptBackoff::default();
        let counter_before = TCP_ACCEPT_BACKOFF_TOTAL.get();

        let started = std::time::Instant::now();
        let mut expected_total = Duration::ZERO;
        for _ in 0..3 {
            expected_total += handle_accept_error(&emfile_error(), &mut backoff).await;
        }
        assert!(expected_total >= ACCEPT_BACKOFF_INITIAL_DELAY * 3);
        assert!(
            started.elapsed() >= expected_total,
            "the accept loop must sleep, not spin; elapsed {:?}",
            started.elapsed(),
        );
        assert!(TCP_ACCEPT_BACKOFF_TOTAL.get() >= counter_before + 3.0);

        let client = tokio::spawn(async move { tokio::net::TcpStream::connect(addr).await });
        let (accepted, _peer) = tokio::time::timeout(Duration::from_secs(5), listener.accept())
            .await
            .expect("listener should still accept after backing off")
            .expect("accept should succeed");
        let _client = client.await.expect("client task").expect("client connect");
        drop(accepted);

        backoff.record_success(std::time::Instant::now);
        assert_eq!(
            backoff.record_exhaustion(std::time::Instant::now()).delay,
            ACCEPT_BACKOFF_INITIAL_DELAY,
        );
    }
}