diff --git a/src-rust/crates/api/src/lib.rs b/src-rust/crates/api/src/lib.rs index 4637643f..ed7c9c8a 100644 --- a/src-rust/crates/api/src/lib.rs +++ b/src-rust/crates/api/src/lib.rs @@ -1400,6 +1400,7 @@ impl StreamAccumulator { } } StreamEvent::MessageStop => AnthropicStreamEvent::MessageStop, + StreamEvent::Ping => AnthropicStreamEvent::Ping, StreamEvent::Error { error_type, message, diff --git a/src-rust/crates/api/src/provider_types.rs b/src-rust/crates/api/src/provider_types.rs index b98e6527..2e7d5b32 100644 --- a/src-rust/crates/api/src/provider_types.rs +++ b/src-rust/crates/api/src/provider_types.rs @@ -163,6 +163,9 @@ pub enum StreamEvent { /// Incremental reasoning / scratchpad delta (alias used by some providers). ReasoningDelta { index: usize, reasoning: String }, + + /// Content-free keepalive that resets the consuming stream stall timer. + Ping, } // --------------------------------------------------------------------------- diff --git a/src-rust/crates/api/src/providers/anthropic.rs b/src-rust/crates/api/src/providers/anthropic.rs index 075653d9..a8b2fbef 100644 --- a/src-rust/crates/api/src/providers/anthropic.rs +++ b/src-rust/crates/api/src/providers/anthropic.rs @@ -215,7 +215,7 @@ impl AnthropicProvider { error_type, message, }), - AnthropicStreamEvent::Ping => None, + AnthropicStreamEvent::Ping => Some(StreamEvent::Ping), } } @@ -367,6 +367,14 @@ mod tests { Box::pin(stream::iter(events.into_iter().map(Ok))) } + #[test] + fn keepalive_ping_reaches_the_stream_consumer() { + assert!( + AnthropicProvider::map_stream_event(AnthropicStreamEvent::Ping).is_some(), + "keepalive pings must reach the consumer's stall timer" + ); + } + #[tokio::test] async fn collected_response_rejects_eof_without_message_stop() { let provider_id = ProviderId::new(ProviderId::ANTHROPIC); diff --git a/src-rust/crates/query/Cargo.toml b/src-rust/crates/query/Cargo.toml index 40bf185a..663439f7 100644 --- a/src-rust/crates/query/Cargo.toml +++ b/src-rust/crates/query/Cargo.toml @@ -21,4 +21,5 @@ dirs = { workspace = true } tokio-util = { workspace = true } [dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } tempfile = { workspace = true } diff --git a/src-rust/crates/query/src/lib.rs b/src-rust/crates/query/src/lib.rs index 589e15bb..89ae7dd4 100644 --- a/src-rust/crates/query/src/lib.rs +++ b/src-rust/crates/query/src/lib.rs @@ -3615,6 +3615,8 @@ fn map_to_anthropic_event( }) } StreamEvent::MessageStop => Some(AnthropicStreamEvent::MessageStop), + // The consumer already reset its stall deadline before mapping this event. + StreamEvent::Ping => None, StreamEvent::Error { error_type, message, @@ -5118,6 +5120,50 @@ mod tests { )); } + #[tokio::test(start_paused = true)] + async fn provider_keepalives_extend_liveness_without_transcript_content() { + let (sender, receiver) = mpsc::channel(8); + let stream = futures::stream::unfold(receiver, |mut receiver| async { + receiver.recv().await.map(|event| (event, receiver)) + }); + let producer = tokio::spawn(async move { + for _ in 0..4 { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + sender.send(Ok(StreamEvent::Ping)).await.unwrap(); + } + sender.send(Ok(StreamEvent::MessageStop)).await.unwrap(); + }); + let cancel = tokio_util::sync::CancellationToken::new(); + let mut accumulator = StreamAccumulator::new(); + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let mut terminal_error = None; + let outcome = consume_provider_stream( + Box::pin(stream), + &mut ProviderStreamContext { + cancel_token: &cancel, + accumulator: &mut accumulator, + event_tx: Some(&event_tx), + provider_id: "anthropic", + model_id: "test", + terminal_stream_error: &mut terminal_error, + }, + std::time::Duration::from_millis(150), + ) + .await; + producer.await.unwrap(); + assert_eq!(outcome, ProviderStreamOutcome::Completed); + assert!(terminal_error.is_none()); + assert!(accumulator.received_message_stop()); + let (message, usage, _) = accumulator.finish(); + assert!(message.content_blocks().is_empty()); + assert_eq!(usage.output_tokens, 0); + assert!(matches!( + event_rx.try_recv(), + Ok(QueryEvent::Stream(AnthropicStreamEvent::MessageStop)) + )); + assert!(event_rx.try_recv().is_err()); + } + #[tokio::test] async fn provider_stream_error_finalizes_partial_response_before_returning_error() { let error = ProviderError::StreamError {