From 561114823c4f95e62f96f0394ed5d0c9c5d486be Mon Sep 17 00:00:00 2001 From: Penguin <50603637+PenguinDOOM@users.noreply.github.com> Date: Fri, 31 Jul 2026 06:49:33 +0900 Subject: [PATCH 1/3] Add active upstream websocket ping --- src/ws_pump.rs | 79 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 75 insertions(+), 4 deletions(-) diff --git a/src/ws_pump.rs b/src/ws_pump.rs index ed96024..43b484e 100644 --- a/src/ws_pump.rs +++ b/src/ws_pump.rs @@ -6,6 +6,7 @@ use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::sync::{Mutex, mpsc}; use tokio::task::JoinHandle; +use tokio::time::{Duration, Instant, MissedTickBehavior}; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; use tracing::debug; @@ -36,13 +37,24 @@ enum OutboundCommand { } const OUTBOUND_CHANNEL_CAPACITY: usize = 32; +const UPSTREAM_PING_INTERVAL: Duration = Duration::from_secs(30); impl LiveUpstreamWebSocket { - pub fn from_stream(_stream: WebSocketStream) -> Self + pub fn from_stream(stream: WebSocketStream) -> Self where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { - let (mut writer, mut reader) = _stream.split(); + Self::from_stream_with_ping_interval(stream, UPSTREAM_PING_INTERVAL) + } + + fn from_stream_with_ping_interval( + stream: WebSocketStream, + ping_interval: Duration, + ) -> Self + where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + { + let (mut writer, mut reader) = stream.split(); let (outbound_tx, mut outbound_rx) = mpsc::channel(OUTBOUND_CHANNEL_CAPACITY); let (inbound_tx, inbound_rx) = mpsc::unbounded_channel(); let close_metadata = Arc::new(Mutex::new(None)); @@ -51,12 +63,26 @@ impl LiveUpstreamWebSocket { let task_is_closed = Arc::clone(&is_closed); let task = tokio::spawn(async move { + let mut ping_timer = tokio::time::interval_at( + Instant::now() + ping_interval, + ping_interval, + ); + ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay); + debug!( outbound_capacity = OUTBOUND_CHANNEL_CAPACITY, + ping_interval_secs = ping_interval.as_secs_f64(), "ws_pump_started" ); loop { tokio::select! { + _ = ping_timer.tick() => { + if let Err(error) = writer.send(Message::Ping(Vec::new())).await { + record_error(&task_close_metadata, error.to_string()).await; + break; + } + debug!("ws_pump_ping_sent"); + } outbound = outbound_rx.recv() => match outbound { Some(OutboundCommand::Text(text)) => { if let Err(error) = writer.send(Message::Text(text)).await { @@ -89,7 +115,9 @@ impl LiveUpstreamWebSocket { } debug!(payload_len, "ws_pump_pong_sent"); } - Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Pong(payload))) => { + debug!(payload_len = payload.len(), "ws_pump_pong_received"); + } Some(Ok(Message::Close(frame))) => { let metadata = UpstreamCloseMetadata { code: frame.as_ref().map(|frame| u16::from(frame.code)), @@ -181,8 +209,8 @@ fn outbound_channel_closed_metadata() -> UpstreamCloseMetadata { #[cfg(test)] mod tests { use super::*; - use std::time::Duration; use tokio::net::TcpListener; + use tokio::sync::oneshot; use tokio::time::timeout; use tokio_tungstenite::accept_async; use tokio_tungstenite::connect_async; @@ -208,6 +236,49 @@ mod tests { pump } + #[tokio::test] + async fn websocket_pump_sends_active_ping_after_interval() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind listener"); + let address = listener.local_addr().expect("local addr"); + let (ping_seen_tx, ping_seen_rx) = oneshot::channel(); + + let accept_task = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept client"); + let websocket = accept_async(stream).await.expect("accept websocket"); + let (_writer, mut reader) = websocket.split(); + + while let Some(message) = reader.next().await { + match message.expect("read websocket message") { + Message::Ping(payload) => { + assert!(payload.is_empty()); + let _ = ping_seen_tx.send(()); + break; + } + Message::Close(_) => break, + _ => {} + } + } + }); + + let (stream, _) = connect_async(format!("ws://{address}")) + .await + .expect("connect websocket"); + let pump = LiveUpstreamWebSocket::from_stream_with_ping_interval( + stream, + Duration::from_millis(20), + ); + + timeout(Duration::from_secs(2), ping_seen_rx) + .await + .expect("active ping should be sent") + .expect("server should report active ping"); + + drop(pump); + accept_task.await.expect("accept task"); + } + #[tokio::test] async fn websocket_pump_records_metadata_when_outbound_channel_closes() { let mut pump = connect_test_pump().await; From bae452dd2646d9cd4ced7a5dccae4a76795b886a Mon Sep 17 00:00:00 2001 From: Penguin <50603637+PenguinDOOM@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:39:38 +0900 Subject: [PATCH 2/3] Strip unsupported prompt cache options --- src/responses/upstream/original.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/responses/upstream/original.rs b/src/responses/upstream/original.rs index 0ab93e6..5e8ce08 100644 --- a/src/responses/upstream/original.rs +++ b/src/responses/upstream/original.rs @@ -12,6 +12,7 @@ const UNSUPPORTED_RESPONSE_FIELDS: &[&str] = &[ "max_output_tokens", "max_tokens", "max_completion_tokens", + "prompt_cache_options", "truncation", ]; @@ -188,6 +189,9 @@ mod tests { "max_output_tokens": 32, "max_tokens": 64, "max_completion_tokens": 96, + "prompt_cache_options": { + "retention": "in_memory" + }, "truncation": "auto" })) .expect("response.create payload"); @@ -198,6 +202,7 @@ mod tests { assert!(payload.get("max_output_tokens").is_none()); assert!(payload.get("max_tokens").is_none()); assert!(payload.get("max_completion_tokens").is_none()); + assert!(payload.get("prompt_cache_options").is_none()); assert!(payload.get("truncation").is_none()); } @@ -409,6 +414,9 @@ mod tests { json!({ "model": "gpt-test", "instructions": "keep", + "prompt_cache_options": { + "retention": "in_memory" + }, "truncation": "auto" }), "resp_intermediate", @@ -429,6 +437,7 @@ mod tests { assert_eq!(payload["input"][0]["type"], "function_call_output"); assert_eq!(payload["input"][0]["call_id"], "call_123"); assert_eq!(payload["input"][0]["output"], "done"); + assert!(payload.get("prompt_cache_options").is_none()); assert!(payload.get("truncation").is_none()); } } From f4a662a65f884053a33edcf61e7a5d8fad19356e Mon Sep 17 00:00:00 2001 From: PenguinDOOM Date: Fri, 31 Jul 2026 14:46:29 +0900 Subject: [PATCH 3/3] style: reformat websocket ping timer initialization - Format the ping timer initialization consistently --- src/ws_pump.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/ws_pump.rs b/src/ws_pump.rs index 43b484e..8c5cdf7 100644 --- a/src/ws_pump.rs +++ b/src/ws_pump.rs @@ -63,10 +63,8 @@ impl LiveUpstreamWebSocket { let task_is_closed = Arc::clone(&is_closed); let task = tokio::spawn(async move { - let mut ping_timer = tokio::time::interval_at( - Instant::now() + ping_interval, - ping_interval, - ); + let mut ping_timer = + tokio::time::interval_at(Instant::now() + ping_interval, ping_interval); ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay); debug!(