polyc_llm/turn.rs
1//! Turn helpers: a [`StubProvider`] for wiring/tests and [`collect_turn`],
2//! which folds a provider's [`Chunk`] stream into a single [`TurnOutput`].
3//!
4//! `collect_turn` is the output half of the bridge between this crate's
5//! streaming vocabulary and the message-granular wire types: the harness drains
6//! a provider stream into a `TurnOutput`, then maps that to wire `Message`s.
7
8use async_trait::async_trait;
9use futures::{Stream, StreamExt, stream};
10
11use crate::{
12 Chunk, CompletionRequest, LlmProvider, StopReason, Usage, error::DummyError, request::ToolCall,
13};
14
15/// An incremental event observed while folding a turn, for live streaming.
16///
17/// Surfaces like Slack `chat.appendStream` or a streaming CLI consume these;
18/// the buffered [`TurnOutput`] is still returned in full — this is a side
19/// channel, not a replacement.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum TurnStreamEvent {
22 /// A freshly-generated piece of answer text (concatenate to reconstruct).
23 TextDelta(String),
24 /// A freshly-generated piece of model reasoning ("thinking") text, distinct
25 /// from the answer. Observers may render it as a collapsed thought; it is
26 /// never concatenated into the answer text.
27 ReasoningDelta(String),
28 /// The model has begun a tool call (`id` + `name` known up front).
29 ToolStarted {
30 /// Provider-assigned call id.
31 id: String,
32 /// Name of the tool being called.
33 name: String,
34 },
35}
36
37/// The fully-assembled result of one turn, folded from a [`Chunk`] stream.
38#[derive(Debug, Default, Clone)]
39pub struct TurnOutput {
40 /// Concatenated text deltas.
41 pub text: String,
42 /// Concatenated reasoning ("thinking") deltas, kept separate from `text`.
43 /// Empty for providers/models that don't expose reasoning.
44 pub reasoning: String,
45 /// Completed tool calls, in arrival order.
46 pub tool_calls: Vec<ToolCall>,
47 /// Final token accounting (last [`Chunk::Usage`] seen).
48 pub usage: Usage,
49 /// Whether the provider's native web-search-grounding primitive actually
50 /// fired this turn — folded from [`Chunk::Grounded`], which a provider
51 /// emits only on response-side proof of use, never merely because
52 /// grounding was allowed on the request. `false` for every provider that
53 /// doesn't support native grounding, correctly: it structurally cannot
54 /// have fired there.
55 pub grounded: bool,
56 /// Why the turn ended, if the stream reported it.
57 pub stop: Option<StopReason>,
58}
59
60/// Drain a provider stream into a [`TurnOutput`].
61///
62/// Text deltas concatenate; a tool call accretes from its
63/// `ToolCallStart`/`ToolCallArgsDelta`/`ToolCallEnd` run (matched by `id`);
64/// usage and stop reason are taken from their chunks.
65///
66/// # Errors
67///
68/// Propagates the first `Err` item from the stream.
69pub async fn collect_turn<S, E>(stream: S) -> Result<TurnOutput, E>
70where
71 S: Stream<Item = Result<Chunk, E>> + Unpin,
72{
73 collect_turn_observed(stream, async |_| {}).await
74}
75
76/// Like [`collect_turn`], but observes each streamable event as it arrives.
77///
78/// Invokes `on_event` for each text delta / tool start while still folding and
79/// returning the complete [`TurnOutput`]. `on_event` is `async` and is
80/// `.await`ed in place before the next stream item is polled: a caller
81/// forwarding onto a bounded channel (`Sender::send`) genuinely applies
82/// backpressure here — a slow consumer on the other end stalls this fold
83/// (and, transitively, the provider stream poll loop) rather than letting
84/// events buffer without limit. Pass a caller-owned clone of a bounded sender
85/// (cloned ONCE outside this call, not per event — `futures::mpsc` grants
86/// every live sender its own reserved slot, so a fresh clone per event would
87/// silently defeat the bound).
88///
89/// # Errors
90///
91/// Propagates the first `Err` item from the stream.
92pub async fn collect_turn_observed<S, E, F>(mut stream: S, mut on_event: F) -> Result<TurnOutput, E>
93where
94 S: Stream<Item = Result<Chunk, E>> + Unpin,
95 F: AsyncFnMut(TurnStreamEvent),
96{
97 let mut out = TurnOutput::default();
98 // In-progress tool calls, kept in start order and matched by id. A provider
99 // may interleave several calls (OpenAI's `parallel_tool_calls` defaults to
100 // true) and/or defer all their `ToolCallEnd`s to the end of the stream, so a
101 // single `Option` would let a second `ToolCallStart` clobber the first and
102 // an `ToolCallEnd` close the wrong call. Matching by id throughout keeps
103 // every parallel call intact regardless of emission order.
104 let mut pending: Vec<ToolCall> = Vec::new();
105 while let Some(item) = stream.next().await {
106 match item? {
107 Chunk::TextDelta(s) => {
108 on_event(TurnStreamEvent::TextDelta(s.clone())).await;
109 out.text.push_str(&s);
110 }
111 Chunk::ReasoningDelta(s) => {
112 on_event(TurnStreamEvent::ReasoningDelta(s.clone())).await;
113 out.reasoning.push_str(&s);
114 }
115 Chunk::ToolCallStart {
116 id,
117 name,
118 signature,
119 approval_turn_id,
120 } => {
121 on_event(TurnStreamEvent::ToolStarted {
122 id: id.clone(),
123 name: name.clone(),
124 })
125 .await;
126 pending.push(ToolCall {
127 id,
128 name,
129 args_json: String::new(),
130 signature,
131 approval_turn_id,
132 });
133 }
134 Chunk::ToolCallArgsDelta {
135 id,
136 args_json_delta,
137 } => {
138 if let Some(tc) = pending.iter_mut().find(|tc| tc.id == id) {
139 tc.args_json.push_str(&args_json_delta);
140 }
141 }
142 Chunk::ToolCallEnd { id } => {
143 // Move the matching call to the output in completion order. An
144 // unmatched id is ignored (defensive); calls still open at EOF
145 // are flushed after the loop so none are silently dropped.
146 if let Some(pos) = pending.iter().position(|tc| tc.id == id) {
147 out.tool_calls.push(pending.remove(pos));
148 }
149 }
150 Chunk::Usage(u) => out.usage = u,
151 Chunk::Grounded => out.grounded = true,
152 // A `ToolUse` stop is sticky against a *later* `EndTurn`. Some
153 // providers stream the tool call in one event and then a separate
154 // trailing terminator event carrying an end-of-turn finish reason;
155 // letting that later `EndTurn` overwrite the `ToolUse` stop would
156 // make the agent loop skip executing the tool and end the turn with
157 // no output.
158 //
159 // A *hard* stop (MaxTokens / Refusal / StopSequence) is the
160 // opposite: it means the turn was truncated or refused, so it must
161 // win over an earlier `ToolUse` — the tool call may be incomplete
162 // and must not be executed.
163 Chunk::Stop(r) => {
164 let keep_tool_use =
165 out.stop == Some(StopReason::ToolUse) && matches!(r, StopReason::EndTurn);
166 if !keep_tool_use {
167 out.stop = Some(r);
168 }
169 }
170 }
171 }
172 // Flush any call that started (and may have accreted args) but whose
173 // `ToolCallEnd` never arrived — a provider that omits the terminator must
174 // not lose the call.
175 out.tool_calls.append(&mut pending);
176 Ok(out)
177}
178
179/// Env var: emit a synthetic tool call for `<name>` on the stub provider.
180///
181/// First `complete()` of a turn emits a synthetic tool call for the named
182/// tool, subsequent calls (once a `tool_result` has landed in the
183/// transcript) fall back to canned `EndTurn` text. Empty / unset keeps the
184/// canned-text behaviour. Used by the HITL resume loopback verification to
185/// drive the data path without a real provider backend.
186pub const STUB_TOOL_CALL_ENV: &str = "POLYCHROME_STUB_TOOL_CALL";
187
188fn stub_tool_name() -> Option<String> {
189 std::env::var(STUB_TOOL_CALL_ENV)
190 .ok()
191 .filter(|s| !s.is_empty())
192}
193
194/// The stub's tool-call sequence: [`STUB_TOOL_CALL_ENV`] split on commas.
195///
196/// One name is the common case; a comma-separated list drives a multi-step
197/// turn (e.g. `read_tool,post_tool`), emitting the Nth tool after N tool
198/// results have landed. Whitespace around each name is trimmed.
199fn stub_tool_sequence() -> Vec<String> {
200 stub_tool_name()
201 .into_iter()
202 .flat_map(|s| {
203 s.split(',')
204 .map(str::trim)
205 .filter(|s| !s.is_empty())
206 .map(str::to_owned)
207 .collect::<Vec<_>>()
208 })
209 .collect()
210}
211
212/// The args-delta for the sequence tool `name` at step `idx`.
213///
214/// Contract, in precedence order:
215/// - [`STUB_TOOL_ARGS_ENV`] unset/empty → `"{}"`.
216/// - Single-tool sequence (`seq_len == 1`) → the whole env value verbatim
217/// (the original single-tool behavior).
218/// - Multi-tool sequence → the env value MUST be a JSON object keyed by tool
219/// name; return the entry for `name` serialized, or `"{}"` if absent.
220///
221/// Unvalidated scaffolding: whatever is returned reaches the tool as its args
222/// delta verbatim so the tool's own schema validation reports any mismatch.
223fn stub_tool_args_for(name: &str, seq_len: usize) -> String {
224 let Some(raw) = std::env::var(STUB_TOOL_ARGS_ENV)
225 .ok()
226 .filter(|s| !s.is_empty())
227 else {
228 return "{}".to_owned();
229 };
230 if seq_len <= 1 {
231 return raw;
232 }
233 match serde_json::from_str::<serde_json::Value>(&raw) {
234 Ok(serde_json::Value::Object(map)) => map
235 .get(name)
236 .map_or_else(|| "{}".to_owned(), ToString::to_string),
237 _ => "{}".to_owned(),
238 }
239}
240
241/// Env var: the JSON-object-literal args delta for the [`STUB_TOOL_CALL_ENV`]
242/// synthetic tool call.
243///
244/// This is wiring/test scaffolding, not a validated input: the value must be
245/// a JSON object literal matching the target tool's input schema. Invalid
246/// JSON is passed through verbatim as the args delta so the tool's own
247/// schema validation reports it downstream — this crate does no validation
248/// of its own. Empty / unset defaults to `"{}"`, matching a tool with no
249/// required arguments.
250pub const STUB_TOOL_ARGS_ENV: &str = "POLYCHROME_STUB_TOOL_ARGS";
251
252/// A canned [`LlmProvider`] for wiring and tests.
253///
254/// Emits two text deltas, a usage tally, and an end-of-turn stop. No
255/// network, no credentials.
256///
257/// When [`STUB_TOOL_CALL_ENV`] is set, the first `complete()` of a turn
258/// emits a synthetic tool call (id `stub-call-1`) for that tool name and
259/// the caller's function-calling loop drives the rest. Subsequent calls
260/// in the same turn fall back to the `EndTurn` text path. Used by the
261/// HITL resume loopback verification.
262///
263/// The call's args delta is `"{}"` unless [`STUB_TOOL_ARGS_ENV`] overrides
264/// it. This is wiring/test scaffolding: the override must be a JSON object
265/// literal matching the target tool's schema, and is passed through
266/// unvalidated.
267#[derive(Clone, Copy, Default)]
268pub struct StubProvider;
269
270#[async_trait]
271impl LlmProvider for StubProvider {
272 type Error = DummyError;
273
274 async fn complete(
275 &self,
276 req: CompletionRequest,
277 ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error> {
278 // If POLYCHROME_STUB_TOOL_CALL is set and we haven't yet seen a
279 // matching tool_result in the transcript, emit the synthetic tool
280 // call. Otherwise fall through to canned text.
281 let sequence = stub_tool_sequence();
282 if !sequence.is_empty() {
283 // Emit the Nth tool once N tool-results have landed, so a
284 // comma-separated sequence drives a multi-step turn deterministically
285 // (e.g. an open-world read that taints, then a gated egress). After
286 // the last tool's result, fall through to the canned end-turn text.
287 let results_seen = req
288 .messages
289 .iter()
290 .flat_map(|m| m.content.iter())
291 .filter(|c| matches!(c, crate::Content::ToolResult(_)))
292 .count();
293 if let Some(tool_name) = sequence.get(results_seen) {
294 let id = format!("stub-call-{}", results_seen + 1);
295 let chunks = vec![
296 Ok(Chunk::tool_call_start(&id, tool_name)),
297 Ok(Chunk::tool_call_args_delta(
298 &id,
299 stub_tool_args_for(tool_name, sequence.len()),
300 )),
301 Ok(Chunk::tool_call_end(&id)),
302 Ok(Chunk::Stop(StopReason::ToolUse)),
303 ];
304 return Ok(stream::iter(chunks).boxed());
305 }
306 }
307 let chunks = vec![
308 Ok(Chunk::text_delta("Hello from the ")),
309 Ok(Chunk::text_delta("stub provider.")),
310 Ok(Chunk::Usage(Usage {
311 input_tokens: 5,
312 output_tokens: 4,
313 ..Default::default()
314 })),
315 Ok(Chunk::Stop(StopReason::EndTurn)),
316 ];
317 Ok(stream::iter(chunks).boxed())
318 }
319}
320
321#[cfg(test)]
322mod tests {
323 #![allow(clippy::pedantic, clippy::nursery, missing_docs)]
324
325 use super::*;
326
327 #[tokio::test]
328 async fn stub_provider_collects_into_text() {
329 let stream = StubProvider
330 .complete(CompletionRequest::new("stub"))
331 .await
332 .expect("stream opens");
333 let out = collect_turn(stream).await.expect("collect");
334 assert_eq!(out.text, "Hello from the stub provider.");
335 assert!(out.tool_calls.is_empty());
336 assert_eq!(out.usage.output_tokens, 4);
337 assert_eq!(out.stop, Some(StopReason::EndTurn));
338 }
339
340 // With STUB_TOOL_ARGS_ENV set, the synthetic tool call's args delta is
341 // exactly that value — useless-against-a-strict-schema "{}" is only the
342 // fallback, not the only option.
343 #[tokio::test]
344 async fn stub_provider_emits_the_configured_tool_args() {
345 temp_env::async_with_vars(
346 [
347 (STUB_TOOL_CALL_ENV, Some("search")),
348 (STUB_TOOL_ARGS_ENV, Some(r#"{"q":"rust"}"#)),
349 ],
350 async {
351 let stream = StubProvider
352 .complete(CompletionRequest::new("stub"))
353 .await
354 .expect("stream opens");
355 let out = collect_turn(stream).await.expect("collect");
356 assert_eq!(out.tool_calls.len(), 1);
357 assert_eq!(out.tool_calls[0].args_json, r#"{"q":"rust"}"#);
358 },
359 )
360 .await;
361 }
362
363 // With STUB_TOOL_ARGS_ENV unset, the synthetic tool call's args delta
364 // falls back to "{}" — today's behavior, preserved.
365 #[tokio::test]
366 async fn stub_provider_defaults_tool_args_to_empty_object() {
367 temp_env::async_with_vars(
368 [
369 (STUB_TOOL_CALL_ENV, Some("search")),
370 (STUB_TOOL_ARGS_ENV, None),
371 ],
372 async {
373 let stream = StubProvider
374 .complete(CompletionRequest::new("stub"))
375 .await
376 .expect("stream opens");
377 let out = collect_turn(stream).await.expect("collect");
378 assert_eq!(out.tool_calls.len(), 1);
379 assert_eq!(out.tool_calls[0].args_json, "{}");
380 },
381 )
382 .await;
383 }
384
385 // With STUB_TOOL_CALL_ENV unset, the canned-text path is unaffected by
386 // the new args knob.
387 #[tokio::test]
388 async fn stub_provider_canned_text_path_is_unchanged_by_the_args_knob() {
389 temp_env::async_with_vars(
390 [
391 (STUB_TOOL_CALL_ENV, None),
392 (STUB_TOOL_ARGS_ENV, Some(r#"{"q":"rust"}"#)),
393 ],
394 async {
395 let stream = StubProvider
396 .complete(CompletionRequest::new("stub"))
397 .await
398 .expect("stream opens");
399 let out = collect_turn(stream).await.expect("collect");
400 assert_eq!(out.text, "Hello from the stub provider.");
401 assert!(out.tool_calls.is_empty());
402 },
403 )
404 .await;
405 }
406
407 #[tokio::test]
408 async fn collect_assembles_tool_call_from_deltas() {
409 let chunks: Vec<Result<Chunk, DummyError>> = vec![
410 Ok(Chunk::text_delta("calling ")),
411 Ok(Chunk::tool_call_start("c1", "search")),
412 Ok(Chunk::tool_call_args_delta("c1", r#"{"q":"#)),
413 Ok(Chunk::tool_call_args_delta("c1", r#""rust"}"#)),
414 Ok(Chunk::tool_call_end("c1")),
415 Ok(Chunk::Stop(StopReason::ToolUse)),
416 ];
417 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
418 assert_eq!(out.text, "calling ");
419 assert_eq!(out.tool_calls.len(), 1);
420 assert_eq!(out.tool_calls[0].name, "search");
421 assert_eq!(out.tool_calls[0].args_json, r#"{"q":"rust"}"#);
422 assert_eq!(out.stop, Some(StopReason::ToolUse));
423 }
424
425 #[tokio::test]
426 async fn collect_keeps_parallel_tool_calls_with_deferred_ends() {
427 // Two interleaved calls whose `ToolCallEnd`s are both deferred to the
428 // end of the stream (the OpenAI-compatible provider's shape). A single
429 // `Option` would drop call 0 and close the survivor with the wrong end;
430 // id-matching must preserve both, in completion order.
431 let chunks: Vec<Result<Chunk, DummyError>> = vec![
432 Ok(Chunk::tool_call_start("c0", "search")),
433 Ok(Chunk::tool_call_args_delta("c0", r#"{"q":"a"}"#)),
434 Ok(Chunk::tool_call_start("c1", "fetch")),
435 Ok(Chunk::tool_call_args_delta("c1", r#"{"u":"b"}"#)),
436 Ok(Chunk::tool_call_end("c0")),
437 Ok(Chunk::tool_call_end("c1")),
438 Ok(Chunk::Stop(StopReason::ToolUse)),
439 ];
440 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
441 assert_eq!(out.tool_calls.len(), 2, "both parallel calls preserved");
442 assert_eq!(out.tool_calls[0].id, "c0");
443 assert_eq!(out.tool_calls[0].name, "search");
444 assert_eq!(out.tool_calls[0].args_json, r#"{"q":"a"}"#);
445 assert_eq!(out.tool_calls[1].id, "c1");
446 assert_eq!(out.tool_calls[1].name, "fetch");
447 assert_eq!(out.tool_calls[1].args_json, r#"{"u":"b"}"#);
448 assert_eq!(out.stop, Some(StopReason::ToolUse));
449 }
450
451 #[tokio::test]
452 async fn collect_flushes_a_call_left_open_at_eof() {
453 // A provider that omits the terminal `ToolCallEnd` must not lose the
454 // call — it is flushed when the stream ends.
455 let chunks: Vec<Result<Chunk, DummyError>> = vec![
456 Ok(Chunk::tool_call_start("c0", "search")),
457 Ok(Chunk::tool_call_args_delta("c0", r#"{"q":"a"}"#)),
458 Ok(Chunk::Stop(StopReason::ToolUse)),
459 ];
460 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
461 assert_eq!(out.tool_calls.len(), 1);
462 assert_eq!(out.tool_calls[0].args_json, r#"{"q":"a"}"#);
463 }
464
465 #[tokio::test]
466 async fn tool_use_stop_is_sticky_against_later_end_turn() {
467 // Provider streams the tool call (ToolUse) then a trailing terminator
468 // event (EndTurn). The terminator must NOT clobber ToolUse, else the
469 // agent loop skips the tool.
470 let chunks: Vec<Result<Chunk, DummyError>> = vec![
471 Ok(Chunk::tool_call_start("c1", "search")),
472 Ok(Chunk::tool_call_end("c1")),
473 Ok(Chunk::Stop(StopReason::ToolUse)),
474 Ok(Chunk::Stop(StopReason::EndTurn)),
475 ];
476 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
477 assert_eq!(out.stop, Some(StopReason::ToolUse));
478 }
479
480 #[tokio::test]
481 async fn hard_stop_wins_over_earlier_tool_use() {
482 // A later MaxTokens (truncation) MUST override an earlier ToolUse so
483 // the agent doesn't execute a tool call with truncated arguments.
484 let chunks: Vec<Result<Chunk, DummyError>> = vec![
485 Ok(Chunk::tool_call_start("c1", "search")),
486 Ok(Chunk::tool_call_end("c1")),
487 Ok(Chunk::Stop(StopReason::ToolUse)),
488 Ok(Chunk::Stop(StopReason::MaxTokens)),
489 ];
490 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
491 assert_eq!(out.stop, Some(StopReason::MaxTokens));
492 }
493
494 #[tokio::test]
495 async fn collect_folds_reasoning_separately_from_text() {
496 // Reasoning deltas accumulate into `reasoning`, never into `text`.
497 let chunks: Vec<Result<Chunk, DummyError>> = vec![
498 Ok(Chunk::reasoning_delta("first ")),
499 Ok(Chunk::reasoning_delta("thought")),
500 Ok(Chunk::text_delta("the answer")),
501 Ok(Chunk::Stop(StopReason::EndTurn)),
502 ];
503 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
504 assert_eq!(out.reasoning, "first thought");
505 assert_eq!(out.text, "the answer");
506 }
507
508 #[tokio::test]
509 async fn observed_reasoning_deltas_are_emitted() {
510 let chunks: Vec<Result<Chunk, DummyError>> = vec![
511 Ok(Chunk::reasoning_delta("hmm")),
512 Ok(Chunk::text_delta("ok")),
513 Ok(Chunk::Stop(StopReason::EndTurn)),
514 ];
515 let mut events = Vec::new();
516 let out = collect_turn_observed(stream::iter(chunks), async |e| events.push(e))
517 .await
518 .expect("collect");
519 assert_eq!(out.reasoning, "hmm");
520 assert!(events.contains(&TurnStreamEvent::ReasoningDelta("hmm".to_owned())));
521 assert!(events.contains(&TurnStreamEvent::TextDelta("ok".to_owned())));
522 }
523
524 #[tokio::test]
525 async fn collect_folds_grounded_evidence() {
526 let chunks: Vec<Result<Chunk, DummyError>> = vec![
527 Ok(Chunk::text_delta("per recent sources, ")),
528 Ok(Chunk::grounded()),
529 Ok(Chunk::text_delta("it's sunny.")),
530 Ok(Chunk::Stop(StopReason::EndTurn)),
531 ];
532 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
533 assert!(
534 out.grounded,
535 "a Grounded chunk anywhere in the stream must fold to true"
536 );
537 }
538
539 #[tokio::test]
540 async fn collect_defaults_grounded_to_false() {
541 // No Chunk::Grounded anywhere — e.g. a provider that doesn't support
542 // native grounding at all, or a response that didn't ground.
543 let chunks: Vec<Result<Chunk, DummyError>> = vec![
544 Ok(Chunk::text_delta("the answer")),
545 Ok(Chunk::Stop(StopReason::EndTurn)),
546 ];
547 let out = collect_turn(stream::iter(chunks)).await.expect("collect");
548 assert!(!out.grounded);
549 }
550
551 #[tokio::test]
552 async fn collect_propagates_error() {
553 let chunks: Vec<Result<Chunk, DummyError>> = vec![
554 Ok(Chunk::text_delta("partial")),
555 Err(DummyError::Other("mid-stream fault".to_owned())),
556 ];
557 let res = collect_turn(stream::iter(chunks)).await;
558 assert!(res.is_err());
559 }
560
561 /// `#251`: `collect_turn_observed` folding onto a bounded channel must
562 /// genuinely stall when the channel is full and undrained — the whole
563 /// point of switching `on_event` to `AsyncFnMut` is that a caller's
564 /// `Sender::send(..).await` blocks the fold (and, transitively, stops
565 /// polling the underlying provider stream) instead of buffering without
566 /// limit. Proven two ways: the wrapped stream's poll count plateaus while
567 /// the channel is full, and the fold only completes after the channel
568 /// drains.
569 #[tokio::test]
570 async fn collect_turn_observed_backpressures_on_a_full_bounded_channel() {
571 use std::pin::Pin;
572 use std::sync::Arc;
573 use std::sync::atomic::{AtomicUsize, Ordering};
574 use std::task::{Context, Poll};
575
576 use futures::SinkExt;
577
578 /// Counts every `poll_next` call on the wrapped stream, so the test
579 /// can observe that the fold has stopped driving the stream forward
580 /// (not merely that the spawned task hasn't been scheduled yet).
581 struct CountingStream<S> {
582 inner: S,
583 polls: Arc<AtomicUsize>,
584 }
585
586 impl<S: Stream + Unpin> Stream for CountingStream<S> {
587 type Item = S::Item;
588 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
589 self.polls.fetch_add(1, Ordering::SeqCst);
590 let this = self.get_mut();
591 Pin::new(&mut this.inner).poll_next(cx)
592 }
593 }
594
595 let chunks: Vec<Result<Chunk, DummyError>> = vec![
596 Ok(Chunk::text_delta("a")),
597 Ok(Chunk::text_delta("b")),
598 Ok(Chunk::text_delta("c")),
599 Ok(Chunk::Stop(StopReason::EndTurn)),
600 ];
601 let polls = Arc::new(AtomicUsize::new(0));
602 let stream = CountingStream {
603 inner: stream::iter(chunks),
604 polls: polls.clone(),
605 };
606
607 // Capacity 1 with a single, never-cloned sender: `futures::mpsc`
608 // grants every live `Sender` a guaranteed slot on top of the shared
609 // buffer, so with exactly one sender the channel absorbs 2 events
610 // before a 3rd send blocks. The plan for this change is explicit that
611 // cloning the sender per event (instead of once, reused) would give
612 // each clone its own slot and silently defeat the bound — this test's
613 // closure captures `tx` by move and reuses the same instance.
614 let (tx, mut rx) = futures::channel::mpsc::channel::<TurnStreamEvent>(1);
615 let handle = tokio::spawn(async move {
616 let mut tx = tx;
617 collect_turn_observed(stream, async move |ev| {
618 let _ = tx.send(ev).await;
619 })
620 .await
621 });
622
623 // Let the fold run until it genuinely stalls on the full channel.
624 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
625 assert!(
626 !handle.is_finished(),
627 "fold must not complete while the channel is full and undrained"
628 );
629 let stalled_at = polls.load(Ordering::SeqCst);
630 assert!(
631 stalled_at < 4,
632 "stream must not have been fully drained while the channel is full"
633 );
634 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
635 assert_eq!(
636 polls.load(Ordering::SeqCst),
637 stalled_at,
638 "poll count must plateau while the channel is full — proof the stall is real"
639 );
640
641 // Draining unblocks the stalled send, letting the fold resume and finish.
642 let mut texts = Vec::new();
643 while texts.len() < 3 {
644 match rx.next().await {
645 Some(TurnStreamEvent::TextDelta(s)) => texts.push(s),
646 Some(_) => {}
647 None => break,
648 }
649 }
650 let out = tokio::time::timeout(std::time::Duration::from_secs(5), handle)
651 .await
652 .expect("fold must complete once the channel drains")
653 .expect("task join")
654 .expect("collect");
655 assert_eq!(out.text, "abc");
656 assert_eq!(texts, vec!["a".to_owned(), "b".to_owned(), "c".to_owned()]);
657 }
658}