Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src-rust/crates/api/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1400,6 +1400,7 @@ impl StreamAccumulator {
}
}
StreamEvent::MessageStop => AnthropicStreamEvent::MessageStop,
StreamEvent::Ping => AnthropicStreamEvent::Ping,
StreamEvent::Error {
error_type,
message,
Expand Down
3 changes: 3 additions & 0 deletions src-rust/crates/api/src/provider_types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

// ---------------------------------------------------------------------------
Expand Down
10 changes: 9 additions & 1 deletion src-rust/crates/api/src/providers/anthropic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,7 @@ impl AnthropicProvider {
error_type,
message,
}),
AnthropicStreamEvent::Ping => None,
AnthropicStreamEvent::Ping => Some(StreamEvent::Ping),
}
}

Expand Down Expand Up @@ -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);
Expand Down
1 change: 1 addition & 0 deletions src-rust/crates/query/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,5 @@ dirs = { workspace = true }
tokio-util = { workspace = true }

[dev-dependencies]
tokio = { workspace = true, features = ["test-util"] }
tempfile = { workspace = true }
46 changes: 46 additions & 0 deletions src-rust/crates/query/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down
Loading