fix: harden connection lifecycle, connect budgets, and follow-up paths

crates/openwire-core/src/body.rs

@@ -141,7 +141,7 @@ impl Body for RequestBody {
     fn is_end_stream(&self) -> bool {
         match &self.inner {
             RequestBodyInner::Empty => true,
-            RequestBodyInner::Replayable { emitted, .. } => *emitted,
+            RequestBodyInner::Replayable { bytes, emitted } => *emitted || bytes.is_empty(),
             RequestBodyInner::Streaming { inner } => inner.is_end_stream(),
         }
     }

crates/openwire-core/src/context.rs

@@ -53,6 +53,8 @@ struct CallContextInner {
     created_at: Instant,
     deadline: Option<Instant>,
     connection_established: AtomicBool,
+    /// When set, response-body Drop discards the connection (decode/body errors).
+    body_force_discard: AtomicBool,
     tls_alpn_preference: TlsAlpnPreference,
 }
 
@@ -86,6 +88,7 @@ impl CallContext {
                 created_at,
                 deadline,
                 connection_established: AtomicBool::new(false),
+                body_force_discard: AtomicBool::new(false),
                 tls_alpn_preference,
             }),
         }
@@ -113,6 +116,14 @@ impl CallContext {
         &self.inner.listener
     }
 
+    pub fn mark_body_force_discard(&self) {
+        self.inner.body_force_discard.store(true, Ordering::Release);
+    }
+
+    pub fn body_force_discard(&self) -> bool {
+        self.inner.body_force_discard.load(Ordering::Acquire)
+    }
+
     pub fn created_at(&self) -> Instant {
         self.inner.created_at
     }

crates/openwire-tungstenite/src/lib.rs

@@ -73,10 +73,15 @@ impl WebSocketEngine for TungsteniteEngine {
 
             // BoxConnection (hyper::rt::Read+Write) → tokio AsyncRead+Write.
             let tokio_io = TokioIo::new(io);
+            let ws_config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig {
+                max_message_size: Some(config.max_message_size),
+                max_frame_size: Some(config.max_frame_size),
+                ..Default::default()
+            };
             let stream = WebSocketStream::from_raw_socket(
                 tokio_io,
                 tokio_tungstenite::tungstenite::protocol::Role::Client,
-                None,
+                Some(ws_config),
             )
             .await;
 

crates/openwire/src/bridge.rs

@@ -42,6 +42,8 @@ impl Interceptor for BridgeInterceptor {
         let transparent_compression = normalization.as_ref().copied().unwrap_or(false);
         Box::pin(async move {
             normalization?;
+            #[cfg(feature = "compression")]
+            let ctx = exchange.context().clone();
             let response = next.run(exchange).await?;
             #[cfg(feature = "compression")]
             {
@@ -50,6 +52,7 @@ impl Interceptor for BridgeInterceptor {
                         response,
                         &request_method,
                         max_decompressed_body_bytes,
+                        ctx,
                     ));
                 }
             }

crates/openwire/src/compression.rs

@@ -12,7 +12,7 @@ use http::header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH, RANGE};
 use http::{HeaderMap, HeaderValue, Method, Response, StatusCode};
 use http_body::{Body, Frame, SizeHint};
 use http_body_util::BodyExt;
-use openwire_core::{RequestBody, ResponseBody, WireError};
+use openwire_core::{CallContext, RequestBody, ResponseBody, WireError};
 use pin_project_lite::pin_project;
 
 const ACCEPTED_ENCODINGS: HeaderValue = HeaderValue::from_static("br, gzip, deflate, zstd");
@@ -42,6 +42,7 @@ pub(crate) fn decode_response(
     response: Response<ResponseBody>,
     request_method: &Method,
     max_decompressed_body_bytes: usize,
+    ctx: CallContext,
 ) -> Response<ResponseBody> {
     if !response_can_have_body(request_method, response.status()) {
         return response;
@@ -62,7 +63,7 @@ pub(crate) fn decode_response(
         .map(|encoding| encoding.as_str())
         .collect::<Vec<_>>()
         .join(", ");
-    let body = DecodedResponseBody::new(body, encodings, label, max_decompressed_body_bytes);
+    let body = DecodedResponseBody::new(body, encodings, label, max_decompressed_body_bytes, ctx);
     Response::from_parts(parts, ResponseBody::new(body.boxed()))
 }
 
@@ -150,6 +151,7 @@ pin_project! {
         label: String,
         max_decompressed_body_bytes: usize,
         decoded_bytes: usize,
+        ctx: CallContext,
     }
 }
 
@@ -159,6 +161,7 @@ impl DecodedResponseBody {
         encodings: Vec<ResponseEncoding>,
         label: String,
         max_decompressed_body_bytes: usize,
+        ctx: CallContext,
     ) -> Self {
         let stream = body.into_data_stream().map_err(wire_error_to_io);
         let reader = stream.into_async_read();
@@ -173,6 +176,7 @@ impl DecodedResponseBody {
             label,
             max_decompressed_body_bytes,
             decoded_bytes: 0,
+            ctx,
         }
     }
 }
@@ -191,6 +195,7 @@ impl Body for DecodedResponseBody {
             Poll::Ready(Ok(0)) => Poll::Ready(None),
             Poll::Ready(Ok(read)) => {
                 let Some(total) = this.decoded_bytes.checked_add(read) else {
+                    this.ctx.mark_body_force_discard();
                     return Poll::Ready(Some(Err(WireError::body(
                         format!(
                             "decompressed {} response exceeded size limit {}",
@@ -200,6 +205,7 @@ impl Body for DecodedResponseBody {
                     ))));
                 };
                 if total > *this.max_decompressed_body_bytes {
+                    this.ctx.mark_body_force_discard();
                     return Poll::Ready(Some(Err(WireError::body(
                         format!(
                             "decompressed {} response exceeded size limit {}",
@@ -214,6 +220,7 @@ impl Body for DecodedResponseBody {
                 )))))
             }
             Poll::Ready(Err(error)) => {
+                this.ctx.mark_body_force_discard();
                 Poll::Ready(Some(Err(io_error_to_wire(error, this.label.as_str()))))
             }
             Poll::Pending => Poll::Pending,
@@ -479,7 +486,16 @@ mod tests {
             .body(ResponseBody::empty())
             .expect("response");
 
-        let response = decode_response(response, &Method::GET, DEFAULT_MAX_DECOMPRESSED_BODY_BYTES);
+        let response = decode_response(
+            response,
+            &Method::GET,
+            DEFAULT_MAX_DECOMPRESSED_BODY_BYTES,
+            openwire_core::CallContext::new(
+                std::sync::Arc::new(openwire_core::NoopEventListener)
+                    as openwire_core::SharedEventListener,
+                None,
+            ),
+        );
 
         assert!(response.headers().get(CONTENT_ENCODING).is_none());
         assert!(response.headers().get(CONTENT_LENGTH).is_none());
@@ -493,7 +509,16 @@ mod tests {
             .body(ResponseBody::empty())
             .expect("response");
 
-        let response = decode_response(response, &Method::GET, DEFAULT_MAX_DECOMPRESSED_BODY_BYTES);
+        let response = decode_response(
+            response,
+            &Method::GET,
+            DEFAULT_MAX_DECOMPRESSED_BODY_BYTES,
+            openwire_core::CallContext::new(
+                std::sync::Arc::new(openwire_core::NoopEventListener)
+                    as openwire_core::SharedEventListener,
+                None,
+            ),
+        );
 
         assert_eq!(response.headers().get(CONTENT_ENCODING).unwrap(), "made-up");
         assert_eq!(response.headers().get(CONTENT_LENGTH).unwrap(), "20");

crates/openwire/src/connection/fast_fallback.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use futures_channel::mpsc;
-use futures_util::future::{AbortHandle, Abortable};
+use futures_util::future::{select, AbortHandle, Abortable, Either};
 use futures_util::stream::StreamExt;
 use hyper::rt::Timer;
 use hyper::Uri;
@@ -376,6 +376,8 @@ impl FastFallbackDialer {
         route_plan: RoutePlan,
         deps: DirectDialDeps,
     ) -> Result<(BoxConnection, FastFallbackOutcome), WireError> {
+        let connect_timeout = deps.connect_timeout;
+        let timer = deps.runtime.timer.clone();
         self.dial_route_plan(
             ctx,
             uri,
@@ -398,7 +400,18 @@ impl FastFallbackDialer {
             },
             move |ctx, uri, _route, stream| {
                 let tls_connector = deps.tls_connector.clone();
-                async move { finalize_direct_connection(ctx, uri, stream, tls_connector).await }
+                let timer = timer.clone();
+                async move {
+                    finalize_direct_connection(
+                        ctx,
+                        uri,
+                        stream,
+                        tls_connector,
+                        timer,
+                        connect_timeout,
+                    )
+                    .await
+                }
             },
         )
         .await
@@ -410,6 +423,8 @@ async fn finalize_direct_connection(
     uri: Uri,
     stream: BoxConnection,
     tls_connector: Option<Arc<dyn TlsConnector>>,
+    timer: SharedTimer,
+    connect_timeout: Option<Duration>,
 ) -> Result<BoxConnection, WireError> {
     if !uri
         .scheme_str()
@@ -425,7 +440,42 @@ async fn finalize_direct_connection(
         ));
     };
 
-    tls_connector.connect(ctx, uri, stream).await
+    with_connect_timeout(
+        timer,
+        connect_timeout,
+        tls_connector.connect(ctx, uri, stream),
+        "TLS handshake",
+    )
+    .await
+}
+
+/// Applies `connect_timeout` to a connect-stage future (TLS or protocol bind).
+pub(crate) async fn with_connect_timeout<T, F>(
+    timer: SharedTimer,
+    connect_timeout: Option<Duration>,
+    future: F,
+    stage: &'static str,
+) -> Result<T, WireError>
+where
+    F: Future<Output = Result<T, WireError>>,
+{
+    let Some(timeout) = connect_timeout else {
+        return future.await;
+    };
+    if timeout.is_zero() {
+        return Err(WireError::connect_timeout(format!(
+            "connect timed out before {stage}"
+        )));
+    }
+
+    let future = Box::pin(future);
+    let sleep = timer.sleep(timeout);
+    match select(future, sleep).await {
+        Either::Left((result, _sleep)) => result,
+        Either::Right((_ready, _future)) => Err(WireError::connect_timeout(format!(
+            "connect timed out after {timeout:?} during {stage}"
+        ))),
+    }
 }
 
 fn failure_stage(error: &WireError) -> ConnectFailureStage {

crates/openwire/src/connection/mod.rs

@@ -12,7 +12,8 @@ pub(crate) use exchange_finder::{
     ResolvedAddress,
 };
 pub(crate) use fast_fallback::{
-    DirectDialDeps, FastFallbackDialer, FastFallbackOutcome, FastFallbackRuntime,
+    with_connect_timeout, DirectDialDeps, FastFallbackDialer, FastFallbackOutcome,
+    FastFallbackRuntime,
 };
 pub(crate) use limits::{
     ConnectionAvailability, ConnectionLimiter, ConnectionPermit, RequestAdmissionLimiter,