Skip to main content

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