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, Socket, Type};
use std::{
    collections::{HashMap, HashSet},
    net::{IpAddr, SocketAddr, TcpListener},
    os::fd::{AsFd, FromRawFd},
    sync::Arc,
    time::Duration,
};
use tokio::sync::Mutex;
use tokio::time::Instant;

/// 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);

use bytes::Bytes;
use derive_builder::Builder;
use futures::{SinkExt, StreamExt};
use local_ip_address::{Error, list_afinet_netifas, local_ip, local_ipv6};

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,
};
use crate::discovery::EndpointInstanceId;
use crate::engine::AsyncEngineContext;
use crate::pipeline::{
    PipelineError,
    network::{
        ResponseService, ResponseStreamPrologue,
        codec::{TwoPartMessage, TwoPartMessageType},
        tcp::StreamType,
    },
};
use anyhow::{Context, Result, anyhow as error};

// Trait for IP address resolution - allows dependency injection for testing
pub trait IpResolver {
    fn local_ip(&self) -> Result<std::net::IpAddr, Error>;
    fn local_ipv6(&self) -> Result<std::net::IpAddr, Error>;
}

// Default implementation using the real local_ip_address crate
pub struct DefaultIpResolver;

impl IpResolver for DefaultIpResolver {
    fn local_ip(&self) -> Result<std::net::IpAddr, Error> {
        local_ip()
    }

    fn local_ipv6(&self) -> Result<std::net::IpAddr, Error> {
        local_ipv6()
    }
}

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

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

    #[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 {
    local_ip: String,
    local_port: u16,
    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, String>>,
    /// 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, String>>,
    /// Capacity of the per-stream mpsc buffer between the socket task and the
    /// engine consumer; carried from the registration [`StreamOptions`].
    send_buffer_count: usize,
}

/// 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>,
    /// 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<()>>>,
}

/// 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 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 local_ip = match options.interface {
            Some(interface) => {
                let interfaces: HashMap<String, std::net::IpAddr> =
                    list_afinet_netifas()?.into_iter().collect();

                interfaces
                    .get(&interface)
                    .ok_or(PipelineError::Generic(format!(
                        "Interface not found: {}",
                        interface
                    )))?
                    .to_string()
            }
            None => {
                let resolved_ip = resolver.local_ip().or_else(|err| match err {
                    Error::LocalIpAddressNotFound => resolver.local_ipv6(),
                    _ => Err(err),
                });

                match resolved_ip {
                    Ok(addr) => addr,
                    // Only fall back to loopback when no routable IP exists at all;
                    // propagate other resolver errors (I/O, platform) so
                    // misconfigured hosts fail fast instead of silently binding
                    // to 127.0.0.1.
                    Err(Error::LocalIpAddressNotFound) => {
                        tracing::warn!(
                            "No routable local IP address found; falling back to 127.0.0.1"
                        );
                        IpAddr::from([127, 0, 0, 1])
                    }
                    Err(err) => {
                        return Err(PipelineError::Generic(format!(
                            "Failed to resolve local IP address: {err}"
                        )));
                    }
                }
                .to_string()
            }
        };

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

        let local_port = Self::start(local_ip.clone(), options.port, state.clone())
            .await
            .map_err(|e| {
                PipelineError::Generic(format!("Failed to start TcpStreamServer: {}", e))
            })?;

        tracing::debug!("tcp transport service on {local_ip}:{local_port}");

        Ok(Arc::new(Self {
            local_ip,
            local_port,
            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().await;
        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(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
    }

    /// 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().await;
        state.rx_subjects.remove(subject);
        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().await;
        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 streams 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().await;
        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);
                }
                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().await;
        state.removed_instances.remove(id);
    }

    #[allow(clippy::await_holding_lock)]
    async fn start(local_ip: String, local_port: u16, state: Arc<Mutex<State>>) -> Result<u16> {
        let addr = format!("{}:{}", local_ip, local_port);
        let state_clone = state.clone();
        let mut guard = state.lock().await;
        if guard.handle.is_some() {
            panic!("TcpStreamServer already started");
        }
        let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<Result<u16>>();
        let handle = tokio::spawn(tcp_listener(addr, state_clone, ready_tx));
        guard.handle = Some(handle);
        drop(guard);
        let local_port = ready_rx.await??;
        Ok(local_port)
    }
}

// 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 = format!("{}:{}", self.local_ip, self.local_port);
        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 (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 mut state = self.state.lock().await;
            state
                .tx_subjects
                .insert(sender_subject.clone(), connection_info);

            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().await;
                    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);
                        }
                    }
                });
            });

            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 connection_info = RequestedRecvConnection {
                context: options.context.clone(),
                connection: pending_recver_tx,
                send_buffer_count: options.send_buffer_count,
            };

            let mut state = self.state.lock().await;
            state
                .rx_subjects
                .insert(receiver_subject.clone(), connection_info);

            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().await;
                    state.rx_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::Response, cleanup_subject.clone()));
                        if subjects.is_empty() {
                            state.instance_subjects.remove(&key);
                        }
                    }
                });
            });

            Some(registered_stream)
        } else {
            None
        };

        PendingConnections {
            send_stream,
            recv_stream,
        }
    }
}

// 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: String,
    state: Arc<Mutex<State>>,
    read_tx: tokio::sync::oneshot::Sender<Result<u16>>,
) -> Result<()> {
    let listener = tokio::net::TcpListener::bind(&addr)
        .await
        .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.port()))
                .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));
        }
    };

    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)) => (stream, _addr),
            Err(e) => {
                // the client should retry, so we don't need to abort
                tracing::warn!("failed to accept tcp connection: {e}");
                eprintln!("failed to accept tcp connection: {}", e);
                continue;
            }
        };

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

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

        tokio::spawn(handle_connection(stream, state.clone()));
    }

    // #[instrument(level = "trace"), skip(state)]
    // todo - clone before spawn and trace process_stream
    async fn handle_connection(stream: tokio::net::TcpStream, state: Arc<Mutex<State>>) {
        let result = process_stream(stream, 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(stream: tokio::net::TcpStream, state: Arc<Mutex<State>>) -> Result<()> {
        // split the socket in to a reader and writer
        let (read_half, write_half) = tokio::io::split(stream);

        // 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 = framed_reader
            .next()
            .await
            .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<tokio::io::ReadHalf<tokio::net::TcpStream>, TwoPartCodec>,
        writer: FramedWrite<tokio::io::WriteHalf<tokio::net::TcpStream>, TwoPartCodec>,
    ) -> Result<()> {
        // Request stream is unidirectional; we don't read from the downstream.
        drop(reader);

        let request_stream = {
            let mut guard = state.lock().await;
            let conn = guard.tx_subjects.remove(&subject).ok_or(error!(
                "Subject not found: {}; downstream subscriber specified a subject unknown to the upstream publisher",
                subject
            ))?;
            if let Some(key) = guard.subject_instance.remove(&subject)
                && let Some(subjects) = guard.instance_subjects.get_mut(&key)
            {
                subjects.remove(&(StreamType::Request, subject.clone()));
                if subjects.is_empty() {
                    guard.instance_subjects.remove(&key);
                }
            }
            conn
        };

        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<tokio::io::WriteHalf<tokio::net::TcpStream>, TwoPartCodec>,
        mut request_rx: mpsc::Receiver<TwoPartMessage>,
        context: Arc<dyn AsyncEngineContext>,
    ) {
        let closing_msg: Option<ControlMessage> = loop {
            tokio::select! {
                biased;

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

                _ = context.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<tokio::io::ReadHalf<tokio::net::TcpStream>, TwoPartCodec>,
        writer: FramedWrite<tokio::io::WriteHalf<tokio::net::TcpStream>, TwoPartCodec>,
    ) -> Result<()> {
        let response_stream = {
            let mut guard = state.lock().await;
            let conn = guard
                .rx_subjects
                .remove(&subject)
                .ok_or(error!("Subject not found: {}; upstream publisher specified a subject unknown to the downsteam subscriber", subject))?;
            if let Some(key) = guard.subject_instance.remove(&subject)
                && let Some(subjects) = guard.instance_subjects.get_mut(&key)
            {
                subjects.remove(&(StreamType::Response, subject.clone()));
                if subjects.is_empty() {
                    guard.instance_subjects.remove(&key);
                }
            }
            conn
        };

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

        // the [`Prologue`]
        // there must be a second control message it indicate the other segment's generate method was successful
        let prologue = reader
            .next()
            .await
            .ok_or(error!("Connection closed without a ControlMessge"))??;

        // deserialize prologue
        let prologue = match prologue.into_message_type() {
            TwoPartMessageType::HeaderOnly(header) => {
                let prologue: ResponseStreamPrologue = serde_json::from_slice(&header)
                    .map_err(|e| error!("Failed to deserialize ControlMessage: {}", e))?;
                prologue
            }
            _ => {
                // 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(msg.to_string()));
                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 _ = connection.send(Err(error.clone()));
            return Err(error!("Received error prologue: {}", error));
        }

        // 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<tokio::io::ReadHalf<tokio::net::TcpStream>, TwoPartCodec>,
        response_tx: mpsc::Sender<Bytes>,
        control_tx: mpsc::Sender<ControlMessage>,
        context: Arc<dyn AsyncEngineContext>,
    ) {
        // loop over reading the tcp stream and checking if the writer is closed
        let mut can_stop = true;
        loop {
            tokio::select! {
                biased;

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

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

                _ = context.stopped(), if can_stop => {
                    tracing::trace!("context stop signal received; shutting down");
                    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<tokio::io::WriteHalf<tokio::net::TcpStream>, 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::pipeline::Context;
    use crate::pipeline::network::DEFAULT_SEND_BUFFER_COUNT;
    use tokio::io::{AsyncWriteExt, ReadHalf, WriteHalf};
    use tokio::net::TcpStream;

    // Mock resolver that always fails to simulate the fallback scenario
    struct FailingIpResolver;

    impl IpResolver for FailingIpResolver {
        fn local_ip(&self) -> Result<std::net::IpAddr, Error> {
            Err(Error::LocalIpAddressNotFound)
        }

        fn local_ipv6(&self) -> Result<std::net::IpAddr, Error> {
            Err(Error::LocalIpAddressNotFound)
        }
    }

    #[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();

        // Verify the server can be used by registering a stream
        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;

        // Verify connection info is available and valid
        let connection_info = pending_connection
            .recv_stream
            .as_ref()
            .unwrap()
            .connection_info
            .clone();

        let tcp_info: TcpStreamConnectionInfo = connection_info.try_into().unwrap();
        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().await;
        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, FailingIpResolver).await;
        assert!(
            result.is_ok(),
            "Server creation should succeed with fallback even when IP detection fails"
        );

        let server = result.unwrap();

        // Get the actual bound address by registering a stream
        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();

        let tcp_info: TcpStreamConnectionInfo = connection_info.try_into().unwrap();
        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");
    }

    /// 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(),
            FailingIpResolver,
        )
        .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, String>>,
    ) {
        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, String>>,
        String,
        tokio::sync::oneshot::Receiver<Result<super::StreamReceiver, String>>,
    ) {
        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().await;
            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().await;
            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().await;
            assert!(
                state.rx_subjects.contains_key(&subject),
                "into_parts() should disarm the RAII cleanup"
            );
        }
    }

    #[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().await;
            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().await;
            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().await;
        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().await;
        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, FailingIpResolver)
            .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 })
                    .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"
        );
    }

    /// 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, FailingIpResolver)
            .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"),
        }
    }

    // ==================== 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:?}"
        );
    }
}