polyc-llm 2026.8.0

Provider-agnostic LLM trait + wire types for polychrome.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
//! Turn helpers: a [`StubProvider`] for wiring/tests and [`collect_turn`],
//! which folds a provider's [`Chunk`] stream into a single [`TurnOutput`].
//!
//! `collect_turn` is the output half of the bridge between this crate's
//! streaming vocabulary and the message-granular wire types: the harness drains
//! a provider stream into a `TurnOutput`, then maps that to wire `Message`s.

use async_trait::async_trait;
use futures::{Stream, StreamExt, stream};

use crate::{
    Chunk, CompletionRequest, LlmProvider, StopReason, Usage, error::DummyError, request::ToolCall,
};

/// An incremental event observed while folding a turn, for live streaming.
///
/// Surfaces like Slack `chat.appendStream` or a streaming CLI consume these;
/// the buffered [`TurnOutput`] is still returned in full — this is a side
/// channel, not a replacement.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TurnStreamEvent {
    /// A freshly-generated piece of answer text (concatenate to reconstruct).
    TextDelta(String),
    /// A freshly-generated piece of model reasoning ("thinking") text, distinct
    /// from the answer. Observers may render it as a collapsed thought; it is
    /// never concatenated into the answer text.
    ReasoningDelta(String),
    /// The model has begun a tool call (`id` + `name` known up front).
    ToolStarted {
        /// Provider-assigned call id.
        id: String,
        /// Name of the tool being called.
        name: String,
    },
}

/// The fully-assembled result of one turn, folded from a [`Chunk`] stream.
#[derive(Debug, Default, Clone)]
pub struct TurnOutput {
    /// Concatenated text deltas.
    pub text: String,
    /// Concatenated reasoning ("thinking") deltas, kept separate from `text`.
    /// Empty for providers/models that don't expose reasoning.
    pub reasoning: String,
    /// Completed tool calls, in arrival order.
    pub tool_calls: Vec<ToolCall>,
    /// Final token accounting (last [`Chunk::Usage`] seen).
    pub usage: Usage,
    /// Whether the provider's native web-search-grounding primitive actually
    /// fired this turn — folded from [`Chunk::Grounded`], which a provider
    /// emits only on response-side proof of use, never merely because
    /// grounding was allowed on the request. `false` for every provider that
    /// doesn't support native grounding, correctly: it structurally cannot
    /// have fired there.
    pub grounded: bool,
    /// Why the turn ended, if the stream reported it.
    pub stop: Option<StopReason>,
}

/// Drain a provider stream into a [`TurnOutput`].
///
/// Text deltas concatenate; a tool call accretes from its
/// `ToolCallStart`/`ToolCallArgsDelta`/`ToolCallEnd` run (matched by `id`);
/// usage and stop reason are taken from their chunks.
///
/// # Errors
///
/// Propagates the first `Err` item from the stream.
pub async fn collect_turn<S, E>(stream: S) -> Result<TurnOutput, E>
where
    S: Stream<Item = Result<Chunk, E>> + Unpin,
{
    collect_turn_observed(stream, async |_| {}).await
}

/// Like [`collect_turn`], but observes each streamable event as it arrives.
///
/// Invokes `on_event` for each text delta / tool start while still folding and
/// returning the complete [`TurnOutput`]. `on_event` is `async` and is
/// `.await`ed in place before the next stream item is polled: a caller
/// forwarding onto a bounded channel (`Sender::send`) genuinely applies
/// backpressure here — a slow consumer on the other end stalls this fold
/// (and, transitively, the provider stream poll loop) rather than letting
/// events buffer without limit. Pass a caller-owned clone of a bounded sender
/// (cloned ONCE outside this call, not per event — `futures::mpsc` grants
/// every live sender its own reserved slot, so a fresh clone per event would
/// silently defeat the bound).
///
/// # Errors
///
/// Propagates the first `Err` item from the stream.
pub async fn collect_turn_observed<S, E, F>(mut stream: S, mut on_event: F) -> Result<TurnOutput, E>
where
    S: Stream<Item = Result<Chunk, E>> + Unpin,
    F: AsyncFnMut(TurnStreamEvent),
{
    let mut out = TurnOutput::default();
    // In-progress tool calls, kept in start order and matched by id. A provider
    // may interleave several calls (OpenAI's `parallel_tool_calls` defaults to
    // true) and/or defer all their `ToolCallEnd`s to the end of the stream, so a
    // single `Option` would let a second `ToolCallStart` clobber the first and
    // an `ToolCallEnd` close the wrong call. Matching by id throughout keeps
    // every parallel call intact regardless of emission order.
    let mut pending: Vec<ToolCall> = Vec::new();
    while let Some(item) = stream.next().await {
        match item? {
            Chunk::TextDelta(s) => {
                on_event(TurnStreamEvent::TextDelta(s.clone())).await;
                out.text.push_str(&s);
            }
            Chunk::ReasoningDelta(s) => {
                on_event(TurnStreamEvent::ReasoningDelta(s.clone())).await;
                out.reasoning.push_str(&s);
            }
            Chunk::ToolCallStart {
                id,
                name,
                signature,
            } => {
                on_event(TurnStreamEvent::ToolStarted {
                    id: id.clone(),
                    name: name.clone(),
                })
                .await;
                pending.push(ToolCall {
                    id,
                    name,
                    args_json: String::new(),
                    signature,
                });
            }
            Chunk::ToolCallArgsDelta {
                id,
                args_json_delta,
            } => {
                if let Some(tc) = pending.iter_mut().find(|tc| tc.id == id) {
                    tc.args_json.push_str(&args_json_delta);
                }
            }
            Chunk::ToolCallEnd { id } => {
                // Move the matching call to the output in completion order. An
                // unmatched id is ignored (defensive); calls still open at EOF
                // are flushed after the loop so none are silently dropped.
                if let Some(pos) = pending.iter().position(|tc| tc.id == id) {
                    out.tool_calls.push(pending.remove(pos));
                }
            }
            Chunk::Usage(u) => out.usage = u,
            Chunk::Grounded => out.grounded = true,
            // A `ToolUse` stop is sticky against a *later* `EndTurn`. Some
            // providers stream the tool call in one event and then a separate
            // trailing terminator event carrying an end-of-turn finish reason;
            // letting that later `EndTurn` overwrite the `ToolUse` stop would
            // make the agent loop skip executing the tool and end the turn with
            // no output.
            //
            // A *hard* stop (MaxTokens / Refusal / StopSequence) is the
            // opposite: it means the turn was truncated or refused, so it must
            // win over an earlier `ToolUse` — the tool call may be incomplete
            // and must not be executed.
            Chunk::Stop(r) => {
                let keep_tool_use =
                    out.stop == Some(StopReason::ToolUse) && matches!(r, StopReason::EndTurn);
                if !keep_tool_use {
                    out.stop = Some(r);
                }
            }
        }
    }
    // Flush any call that started (and may have accreted args) but whose
    // `ToolCallEnd` never arrived — a provider that omits the terminator must
    // not lose the call.
    out.tool_calls.append(&mut pending);
    Ok(out)
}

/// Env var: emit a synthetic tool call for `<name>` on the stub provider.
///
/// First `complete()` of a turn emits a synthetic tool call for the named
/// tool, subsequent calls (once a `tool_result` has landed in the
/// transcript) fall back to canned `EndTurn` text. Empty / unset keeps the
/// canned-text behaviour. Used by the HITL resume loopback verification to
/// drive the data path without a real provider backend.
pub const STUB_TOOL_CALL_ENV: &str = "POLYCHROME_STUB_TOOL_CALL";

fn stub_tool_name() -> Option<String> {
    std::env::var(STUB_TOOL_CALL_ENV)
        .ok()
        .filter(|s| !s.is_empty())
}

/// The stub's tool-call sequence: [`STUB_TOOL_CALL_ENV`] split on commas.
///
/// One name is the common case; a comma-separated list drives a multi-step
/// turn (e.g. `read_tool,post_tool`), emitting the Nth tool after N tool
/// results have landed. Whitespace around each name is trimmed.
fn stub_tool_sequence() -> Vec<String> {
    stub_tool_name()
        .into_iter()
        .flat_map(|s| {
            s.split(',')
                .map(str::trim)
                .filter(|s| !s.is_empty())
                .map(str::to_owned)
                .collect::<Vec<_>>()
        })
        .collect()
}

/// The args-delta for the sequence tool `name` at step `idx`.
///
/// Contract, in precedence order:
/// - [`STUB_TOOL_ARGS_ENV`] unset/empty → `"{}"`.
/// - Single-tool sequence (`seq_len == 1`) → the whole env value verbatim
///   (the original single-tool behavior).
/// - Multi-tool sequence → the env value MUST be a JSON object keyed by tool
///   name; return the entry for `name` serialized, or `"{}"` if absent.
///
/// Unvalidated scaffolding: whatever is returned reaches the tool as its args
/// delta verbatim so the tool's own schema validation reports any mismatch.
fn stub_tool_args_for(name: &str, seq_len: usize) -> String {
    let Some(raw) = std::env::var(STUB_TOOL_ARGS_ENV)
        .ok()
        .filter(|s| !s.is_empty())
    else {
        return "{}".to_owned();
    };
    if seq_len <= 1 {
        return raw;
    }
    match serde_json::from_str::<serde_json::Value>(&raw) {
        Ok(serde_json::Value::Object(map)) => map
            .get(name)
            .map_or_else(|| "{}".to_owned(), ToString::to_string),
        _ => "{}".to_owned(),
    }
}

/// Env var: the JSON-object-literal args delta for the [`STUB_TOOL_CALL_ENV`]
/// synthetic tool call.
///
/// This is wiring/test scaffolding, not a validated input: the value must be
/// a JSON object literal matching the target tool's input schema. Invalid
/// JSON is passed through verbatim as the args delta so the tool's own
/// schema validation reports it downstream — this crate does no validation
/// of its own. Empty / unset defaults to `"{}"`, matching a tool with no
/// required arguments.
pub const STUB_TOOL_ARGS_ENV: &str = "POLYCHROME_STUB_TOOL_ARGS";

/// A canned [`LlmProvider`] for wiring and tests.
///
/// Emits two text deltas, a usage tally, and an end-of-turn stop. No
/// network, no credentials.
///
/// When [`STUB_TOOL_CALL_ENV`] is set, the first `complete()` of a turn
/// emits a synthetic tool call (id `stub-call-1`) for that tool name and
/// the caller's function-calling loop drives the rest. Subsequent calls
/// in the same turn fall back to the `EndTurn` text path. Used by the
/// HITL resume loopback verification.
///
/// The call's args delta is `"{}"` unless [`STUB_TOOL_ARGS_ENV`] overrides
/// it. This is wiring/test scaffolding: the override must be a JSON object
/// literal matching the target tool's schema, and is passed through
/// unvalidated.
#[derive(Clone, Copy, Default)]
pub struct StubProvider;

#[async_trait]
impl LlmProvider for StubProvider {
    type Error = DummyError;

    async fn complete(
        &self,
        req: CompletionRequest,
    ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error> {
        // If POLYCHROME_STUB_TOOL_CALL is set and we haven't yet seen a
        // matching tool_result in the transcript, emit the synthetic tool
        // call. Otherwise fall through to canned text.
        let sequence = stub_tool_sequence();
        if !sequence.is_empty() {
            // Emit the Nth tool once N tool-results have landed, so a
            // comma-separated sequence drives a multi-step turn deterministically
            // (e.g. an open-world read that taints, then a gated egress). After
            // the last tool's result, fall through to the canned end-turn text.
            let results_seen = req
                .messages
                .iter()
                .flat_map(|m| m.content.iter())
                .filter(|c| matches!(c, crate::Content::ToolResult(_)))
                .count();
            if let Some(tool_name) = sequence.get(results_seen) {
                let id = format!("stub-call-{}", results_seen + 1);
                let chunks = vec![
                    Ok(Chunk::tool_call_start(&id, tool_name)),
                    Ok(Chunk::tool_call_args_delta(
                        &id,
                        stub_tool_args_for(tool_name, sequence.len()),
                    )),
                    Ok(Chunk::tool_call_end(&id)),
                    Ok(Chunk::Stop(StopReason::ToolUse)),
                ];
                return Ok(stream::iter(chunks).boxed());
            }
        }
        let chunks = vec![
            Ok(Chunk::text_delta("Hello from the ")),
            Ok(Chunk::text_delta("stub provider.")),
            Ok(Chunk::Usage(Usage {
                input_tokens: 5,
                output_tokens: 4,
                ..Default::default()
            })),
            Ok(Chunk::Stop(StopReason::EndTurn)),
        ];
        Ok(stream::iter(chunks).boxed())
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use super::*;

    #[tokio::test]
    async fn stub_provider_collects_into_text() {
        let stream = StubProvider
            .complete(CompletionRequest::new("stub"))
            .await
            .expect("stream opens");
        let out = collect_turn(stream).await.expect("collect");
        assert_eq!(out.text, "Hello from the stub provider.");
        assert!(out.tool_calls.is_empty());
        assert_eq!(out.usage.output_tokens, 4);
        assert_eq!(out.stop, Some(StopReason::EndTurn));
    }

    // With STUB_TOOL_ARGS_ENV set, the synthetic tool call's args delta is
    // exactly that value — useless-against-a-strict-schema "{}" is only the
    // fallback, not the only option.
    #[tokio::test]
    async fn stub_provider_emits_the_configured_tool_args() {
        temp_env::async_with_vars(
            [
                (STUB_TOOL_CALL_ENV, Some("search")),
                (STUB_TOOL_ARGS_ENV, Some(r#"{"q":"rust"}"#)),
            ],
            async {
                let stream = StubProvider
                    .complete(CompletionRequest::new("stub"))
                    .await
                    .expect("stream opens");
                let out = collect_turn(stream).await.expect("collect");
                assert_eq!(out.tool_calls.len(), 1);
                assert_eq!(out.tool_calls[0].args_json, r#"{"q":"rust"}"#);
            },
        )
        .await;
    }

    // With STUB_TOOL_ARGS_ENV unset, the synthetic tool call's args delta
    // falls back to "{}" — today's behavior, preserved.
    #[tokio::test]
    async fn stub_provider_defaults_tool_args_to_empty_object() {
        temp_env::async_with_vars(
            [
                (STUB_TOOL_CALL_ENV, Some("search")),
                (STUB_TOOL_ARGS_ENV, None),
            ],
            async {
                let stream = StubProvider
                    .complete(CompletionRequest::new("stub"))
                    .await
                    .expect("stream opens");
                let out = collect_turn(stream).await.expect("collect");
                assert_eq!(out.tool_calls.len(), 1);
                assert_eq!(out.tool_calls[0].args_json, "{}");
            },
        )
        .await;
    }

    // With STUB_TOOL_CALL_ENV unset, the canned-text path is unaffected by
    // the new args knob.
    #[tokio::test]
    async fn stub_provider_canned_text_path_is_unchanged_by_the_args_knob() {
        temp_env::async_with_vars(
            [
                (STUB_TOOL_CALL_ENV, None),
                (STUB_TOOL_ARGS_ENV, Some(r#"{"q":"rust"}"#)),
            ],
            async {
                let stream = StubProvider
                    .complete(CompletionRequest::new("stub"))
                    .await
                    .expect("stream opens");
                let out = collect_turn(stream).await.expect("collect");
                assert_eq!(out.text, "Hello from the stub provider.");
                assert!(out.tool_calls.is_empty());
            },
        )
        .await;
    }

    #[tokio::test]
    async fn collect_assembles_tool_call_from_deltas() {
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::text_delta("calling ")),
            Ok(Chunk::tool_call_start("c1", "search")),
            Ok(Chunk::tool_call_args_delta("c1", r#"{"q":"#)),
            Ok(Chunk::tool_call_args_delta("c1", r#""rust"}"#)),
            Ok(Chunk::tool_call_end("c1")),
            Ok(Chunk::Stop(StopReason::ToolUse)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert_eq!(out.text, "calling ");
        assert_eq!(out.tool_calls.len(), 1);
        assert_eq!(out.tool_calls[0].name, "search");
        assert_eq!(out.tool_calls[0].args_json, r#"{"q":"rust"}"#);
        assert_eq!(out.stop, Some(StopReason::ToolUse));
    }

    #[tokio::test]
    async fn collect_keeps_parallel_tool_calls_with_deferred_ends() {
        // Two interleaved calls whose `ToolCallEnd`s are both deferred to the
        // end of the stream (the OpenAI-compatible provider's shape). A single
        // `Option` would drop call 0 and close the survivor with the wrong end;
        // id-matching must preserve both, in completion order.
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::tool_call_start("c0", "search")),
            Ok(Chunk::tool_call_args_delta("c0", r#"{"q":"a"}"#)),
            Ok(Chunk::tool_call_start("c1", "fetch")),
            Ok(Chunk::tool_call_args_delta("c1", r#"{"u":"b"}"#)),
            Ok(Chunk::tool_call_end("c0")),
            Ok(Chunk::tool_call_end("c1")),
            Ok(Chunk::Stop(StopReason::ToolUse)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert_eq!(out.tool_calls.len(), 2, "both parallel calls preserved");
        assert_eq!(out.tool_calls[0].id, "c0");
        assert_eq!(out.tool_calls[0].name, "search");
        assert_eq!(out.tool_calls[0].args_json, r#"{"q":"a"}"#);
        assert_eq!(out.tool_calls[1].id, "c1");
        assert_eq!(out.tool_calls[1].name, "fetch");
        assert_eq!(out.tool_calls[1].args_json, r#"{"u":"b"}"#);
        assert_eq!(out.stop, Some(StopReason::ToolUse));
    }

    #[tokio::test]
    async fn collect_flushes_a_call_left_open_at_eof() {
        // A provider that omits the terminal `ToolCallEnd` must not lose the
        // call — it is flushed when the stream ends.
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::tool_call_start("c0", "search")),
            Ok(Chunk::tool_call_args_delta("c0", r#"{"q":"a"}"#)),
            Ok(Chunk::Stop(StopReason::ToolUse)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert_eq!(out.tool_calls.len(), 1);
        assert_eq!(out.tool_calls[0].args_json, r#"{"q":"a"}"#);
    }

    #[tokio::test]
    async fn tool_use_stop_is_sticky_against_later_end_turn() {
        // Provider streams the tool call (ToolUse) then a trailing terminator
        // event (EndTurn). The terminator must NOT clobber ToolUse, else the
        // agent loop skips the tool.
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::tool_call_start("c1", "search")),
            Ok(Chunk::tool_call_end("c1")),
            Ok(Chunk::Stop(StopReason::ToolUse)),
            Ok(Chunk::Stop(StopReason::EndTurn)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert_eq!(out.stop, Some(StopReason::ToolUse));
    }

    #[tokio::test]
    async fn hard_stop_wins_over_earlier_tool_use() {
        // A later MaxTokens (truncation) MUST override an earlier ToolUse so
        // the agent doesn't execute a tool call with truncated arguments.
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::tool_call_start("c1", "search")),
            Ok(Chunk::tool_call_end("c1")),
            Ok(Chunk::Stop(StopReason::ToolUse)),
            Ok(Chunk::Stop(StopReason::MaxTokens)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert_eq!(out.stop, Some(StopReason::MaxTokens));
    }

    #[tokio::test]
    async fn collect_folds_reasoning_separately_from_text() {
        // Reasoning deltas accumulate into `reasoning`, never into `text`.
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::reasoning_delta("first ")),
            Ok(Chunk::reasoning_delta("thought")),
            Ok(Chunk::text_delta("the answer")),
            Ok(Chunk::Stop(StopReason::EndTurn)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert_eq!(out.reasoning, "first thought");
        assert_eq!(out.text, "the answer");
    }

    #[tokio::test]
    async fn observed_reasoning_deltas_are_emitted() {
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::reasoning_delta("hmm")),
            Ok(Chunk::text_delta("ok")),
            Ok(Chunk::Stop(StopReason::EndTurn)),
        ];
        let mut events = Vec::new();
        let out = collect_turn_observed(stream::iter(chunks), async |e| events.push(e))
            .await
            .expect("collect");
        assert_eq!(out.reasoning, "hmm");
        assert!(events.contains(&TurnStreamEvent::ReasoningDelta("hmm".to_owned())));
        assert!(events.contains(&TurnStreamEvent::TextDelta("ok".to_owned())));
    }

    #[tokio::test]
    async fn collect_folds_grounded_evidence() {
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::text_delta("per recent sources, ")),
            Ok(Chunk::grounded()),
            Ok(Chunk::text_delta("it's sunny.")),
            Ok(Chunk::Stop(StopReason::EndTurn)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert!(
            out.grounded,
            "a Grounded chunk anywhere in the stream must fold to true"
        );
    }

    #[tokio::test]
    async fn collect_defaults_grounded_to_false() {
        // No Chunk::Grounded anywhere — e.g. a provider that doesn't support
        // native grounding at all, or a response that didn't ground.
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::text_delta("the answer")),
            Ok(Chunk::Stop(StopReason::EndTurn)),
        ];
        let out = collect_turn(stream::iter(chunks)).await.expect("collect");
        assert!(!out.grounded);
    }

    #[tokio::test]
    async fn collect_propagates_error() {
        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::text_delta("partial")),
            Err(DummyError::Other("mid-stream fault".to_owned())),
        ];
        let res = collect_turn(stream::iter(chunks)).await;
        assert!(res.is_err());
    }

    /// `#251`: `collect_turn_observed` folding onto a bounded channel must
    /// genuinely stall when the channel is full and undrained — the whole
    /// point of switching `on_event` to `AsyncFnMut` is that a caller's
    /// `Sender::send(..).await` blocks the fold (and, transitively, stops
    /// polling the underlying provider stream) instead of buffering without
    /// limit. Proven two ways: the wrapped stream's poll count plateaus while
    /// the channel is full, and the fold only completes after the channel
    /// drains.
    #[tokio::test]
    async fn collect_turn_observed_backpressures_on_a_full_bounded_channel() {
        use std::pin::Pin;
        use std::sync::Arc;
        use std::sync::atomic::{AtomicUsize, Ordering};
        use std::task::{Context, Poll};

        use futures::SinkExt;

        /// Counts every `poll_next` call on the wrapped stream, so the test
        /// can observe that the fold has stopped driving the stream forward
        /// (not merely that the spawned task hasn't been scheduled yet).
        struct CountingStream<S> {
            inner: S,
            polls: Arc<AtomicUsize>,
        }

        impl<S: Stream + Unpin> Stream for CountingStream<S> {
            type Item = S::Item;
            fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
                self.polls.fetch_add(1, Ordering::SeqCst);
                let this = self.get_mut();
                Pin::new(&mut this.inner).poll_next(cx)
            }
        }

        let chunks: Vec<Result<Chunk, DummyError>> = vec![
            Ok(Chunk::text_delta("a")),
            Ok(Chunk::text_delta("b")),
            Ok(Chunk::text_delta("c")),
            Ok(Chunk::Stop(StopReason::EndTurn)),
        ];
        let polls = Arc::new(AtomicUsize::new(0));
        let stream = CountingStream {
            inner: stream::iter(chunks),
            polls: polls.clone(),
        };

        // Capacity 1 with a single, never-cloned sender: `futures::mpsc`
        // grants every live `Sender` a guaranteed slot on top of the shared
        // buffer, so with exactly one sender the channel absorbs 2 events
        // before a 3rd send blocks. The plan for this change is explicit that
        // cloning the sender per event (instead of once, reused) would give
        // each clone its own slot and silently defeat the bound — this test's
        // closure captures `tx` by move and reuses the same instance.
        let (tx, mut rx) = futures::channel::mpsc::channel::<TurnStreamEvent>(1);
        let handle = tokio::spawn(async move {
            let mut tx = tx;
            collect_turn_observed(stream, async move |ev| {
                let _ = tx.send(ev).await;
            })
            .await
        });

        // Let the fold run until it genuinely stalls on the full channel.
        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        assert!(
            !handle.is_finished(),
            "fold must not complete while the channel is full and undrained"
        );
        let stalled_at = polls.load(Ordering::SeqCst);
        assert!(
            stalled_at < 4,
            "stream must not have been fully drained while the channel is full"
        );
        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        assert_eq!(
            polls.load(Ordering::SeqCst),
            stalled_at,
            "poll count must plateau while the channel is full — proof the stall is real"
        );

        // Draining unblocks the stalled send, letting the fold resume and finish.
        let mut texts = Vec::new();
        while texts.len() < 3 {
            match rx.next().await {
                Some(TurnStreamEvent::TextDelta(s)) => texts.push(s),
                Some(_) => {}
                None => break,
            }
        }
        let out = tokio::time::timeout(std::time::Duration::from_secs(5), handle)
            .await
            .expect("fold must complete once the channel drains")
            .expect("task join")
            .expect("collect");
        assert_eq!(out.text, "abc");
        assert_eq!(texts, vec!["a".to_owned(), "b".to_owned(), "c".to_owned()]);
    }
}