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()); } } diff --git a/src/ws_pump.rs b/src/ws_pump.rs index ed96024..8c5cdf7 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,24 @@ 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 +113,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 +207,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 +234,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;