agy-bridge 0.10.0

Async Rust bridge and native runtime for the Google Antigravity SDK
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
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
//! Step streaming and forwarding handlers.
//!
//! ## Triple-channel fan-out
//!
//! Each streaming event (text, thought, tool-call) is forwarded to **three**
//! channels simultaneously, serving different consumer patterns:
//!
//! | Channel     | Purpose                                       |
//! |-------------|-----------------------------------------------|
//! | `event_tx`  | Timeline replay: all events in arrival order  |
//! | `chunk_tx`  | Unified stream: typed chunks for consumers    |
//! | `text_tx` / `thought_tx` / `tool_call_tx` | Type-specific streams |
//!
//! This fan-out is intentional: consumers may attach to whichever channel
//! suits their use-case (e.g. a CLI only needs `text_tx`, while a UI might
//! consume the full `event_tx` timeline).
//!
//! ## Subscription gating (deadlock safety)
//!
//! Each channel is only written when a consumer has *subscribed* to it (by
//! calling the matching handle accessor). A channel nobody drains is skipped
//! entirely, so its bounded buffer can never fill and block the writer. This
//! makes every consumption pattern deadlock-free while still delivering every
//! item to the channels that *are* consumed. See `StreamSubscriptions` and
//! `ChatResponseWriter::fan_out`.

use pyo3::prelude::*;

// ── Stream error strategy ─────────────────────────────────────────────────
//
// agy-bridge wraps the Python SDK and does NOT classify individual error
// steps.  The Go backend (localharness) handles retries and decides what's
// fatal vs recoverable.  The Python SDK raises exceptions for fatal errors
// (e.g. `AntigravityExecutionError` after exhausted retries) and yields
// error steps for intermediate failures (which the backend may retry).
//
// We simply:
// 1. Forward every error step to timeline consumers (step channel).
// 2. Track the last error and whether output arrived *after* it.
// 3. At stream end: if no output came after the last error → propagate it
//    to error_tx so handle.text() returns Err(StreamError).
//    If the backend recovered (produced text after the error) → skip it.
//
// This avoids fragile string-matching on backend-specific messages like
// "model output", "terminated", or "Retryable error".
use super::AgentId;

/// What a forwarded step contained — used by the stream loop to track
/// whether the stream produced useful output.
enum StepContent {
    /// The step had text content and/or tool calls.
    Output,
    /// The step was an error (no content produced, error forwarded). Carries
    /// the harness-reported HTTP status code (`0` when unknown).
    Error { message: String, http_code: u16 },
    /// The step had no content and no error (e.g. metadata-only).
    Empty,
}

/// Maximum number of *consecutive* model-quality error steps tolerated within a
/// single stream before we stop pulling the SDK iterator and fail the turn.
///
/// The Python SDK retries a "thinking-only / invalid output" turn internally by
/// re-issuing the model call, and with some backends each re-issue is a fresh
/// RPC — a brand-new subprocess and network connection. A context that
/// *deterministically* yields invalid output would otherwise retry up to the
/// SDK's own (large) ceiling, spawning a burst of RPCs per turn that can exhaust
/// connection/socket resources. Cutting the stream after a few attempts hands control to the
/// orchestrator's turn-level recovery ladder (drop-the-bad-turn + corrective
/// nudge, then respawn), which breaks the loop far more cheaply — and more
/// effectively — than blind re-generation on the same poisoned context.
///
/// Transient API errors (e.g. HTTP 503) are deliberately *not* counted here: the
/// SDK's backoff/retry for those is genuinely useful and left intact.
const DEFAULT_MAX_CONSECUTIVE_MODEL_ERRORS: u32 = 3;

/// Maximum number of *consecutive* thinking-only / empty steps tolerated within a
/// single generation before we abort the stream.
///
/// Some models can enter a rumination loop: they stream thinking deltas
/// indefinitely without ever emitting text or a tool call, and the generation
/// only terminates with an "invalid output" error much later. While that single
/// stream runs (often minutes), the underlying subprocess keeps opening sockets,
/// so a *single* runaway generation — not just a retry loop — can balloon
/// fd/socket usage into the thousands. Aborting after a generous ceiling of
/// purely-non-productive steps caps the subprocess's lifetime and hands the turn
/// to the orchestrator's recovery ladder.
///
/// The ceiling is deliberately high so that legitimately long chains of thought
/// that *do* eventually produce output (which resets the counter) are never cut.
const DEFAULT_MAX_CONSECUTIVE_EMPTY_STEPS: u32 = 500;

/// Configurable thresholds for stream abort safety valves.
///
/// These limits prevent runaway SDK retry loops and rumination streams from
/// exhausting resources. A value of `0` disables the corresponding limit
/// entirely, giving pure SDK pass-through behavior.
#[derive(Debug, Clone, Copy)]
pub(crate) struct StreamLimits {
    /// Max consecutive model-quality errors before aborting. 0 = unlimited.
    pub max_model_errors: u32,
    /// Max consecutive thinking-only/empty steps before aborting. 0 = unlimited.
    pub max_empty_steps: u32,
    /// Buffer size for streaming response channels.
    pub channel_buffer: usize,
}

impl Default for StreamLimits {
    fn default() -> Self {
        Self {
            max_model_errors: DEFAULT_MAX_CONSECUTIVE_MODEL_ERRORS,
            max_empty_steps: DEFAULT_MAX_CONSECUTIVE_EMPTY_STEPS,
            channel_buffer: crate::streaming::DEFAULT_CHANNEL_BUFFER,
        }
    }
}

impl StreamLimits {
    /// Build from `RuntimeConfig` overrides, falling back to defaults.
    pub fn from_config(config: &super::config::RuntimeConfig) -> Self {
        Self {
            max_model_errors: config
                .max_consecutive_model_errors
                .unwrap_or(DEFAULT_MAX_CONSECUTIVE_MODEL_ERRORS),
            max_empty_steps: config
                .max_consecutive_empty_steps
                .unwrap_or(DEFAULT_MAX_CONSECUTIVE_EMPTY_STEPS),
            channel_buffer: config
                .streaming_channel_buffer
                .unwrap_or(crate::streaming::DEFAULT_CHANNEL_BUFFER),
        }
    }
}

/// Synthetic error recorded when a generation is aborted for producing only
/// thinking/empty steps. Contains the "model output" marker so it is classified
/// as a model-quality failure by [`is_model_quality_error`] and routed through
/// the same recovery path as a backend-reported invalid-output error.
const RUNAWAY_THINKING_ERROR: &str = "aborted: model output contained only thinking with no text or tool calls \
     after too many consecutive steps (runaway rumination)";

/// Whether an error message denotes a model-*quality* failure (empty /
/// thought-only / invalid tool-call output) as opposed to a transient transport
/// or API error. Kept as a single predicate so the streaming loop and the
/// per-step logging agree on the classification.
fn is_model_quality_error(message: &str) -> bool {
    message.contains("model output")
}

/// Running state of the end-of-stream error state machine.
///
/// Tracks the last error seen, whether useful output arrived *after* it (which
/// makes the error stale), how many model-quality errors have occurred back to
/// back (used to break out of an SDK re-generation loop early), and how many
/// thinking-only steps have occurred back to back (used to abort a single
/// runaway rumination stream).
#[derive(Default)]
struct StreamErrorState {
    last_error: Option<String>,
    last_http_code: u16,
    output_after_error: bool,
    consecutive_model_errors: u32,
    consecutive_empty_steps: u32,
    limits: StreamLimits,
}

impl StreamErrorState {
    fn new(limits: StreamLimits) -> Self {
        Self {
            limits,
            ..Self::default()
        }
    }

    /// Fold one step's classified [`StepContent`] into the state.
    ///
    /// Returns `true` when the caller should stop pulling the iterator, either
    /// because too many consecutive model-quality errors have occurred (the SDK
    /// is re-generating a deterministically-bad turn) or because a single
    /// generation has produced too many consecutive thinking-only steps (a
    /// runaway rumination stream).
    fn observe(&mut self, content: &StepContent) -> bool {
        match content {
            StepContent::Error { message, http_code } => {
                // An error breaks any thinking-only streak.
                self.consecutive_empty_steps = 0;
                if is_model_quality_error(message) {
                    self.consecutive_model_errors += 1;
                } else {
                    // A transport/API error breaks the model-quality streak;
                    // the SDK's own retry policy governs those.
                    self.consecutive_model_errors = 0;
                }
                self.last_error = Some(message.clone());
                self.last_http_code = *http_code;
                self.output_after_error = false;
                self.limits.max_model_errors > 0
                    && self.consecutive_model_errors >= self.limits.max_model_errors
            }
            StepContent::Output => {
                // Any usable output means the turn is making progress: forget
                // both streaks and mark the last error (if any) stale.
                self.consecutive_model_errors = 0;
                self.consecutive_empty_steps = 0;
                if self.last_error.is_some() {
                    self.output_after_error = true;
                }
                false
            }
            StepContent::Empty => {
                self.consecutive_empty_steps += 1;
                if self.limits.max_empty_steps > 0
                    && self.consecutive_empty_steps >= self.limits.max_empty_steps
                {
                    // Record a synthetic model-quality error so the end-of-stream
                    // logic propagates it and the orchestrator recovers.
                    self.last_error = Some(RUNAWAY_THINKING_ERROR.to_string());
                    // Synthetic model-quality failure — no HTTP status applies.
                    self.last_http_code = crate::error::HTTP_CODE_UNKNOWN;
                    self.output_after_error = false;
                    true
                } else {
                    false
                }
            }
        }
    }
}

async fn forward_step_to_writer(
    writer: &crate::streaming::ChatResponseWriter,
    mut step: crate::types::Step,
    agent_id: AgentId,
    streamed_text: &mut String,
) -> StepContent {
    // ── Error detection ─────────────────────────────────────────────────
    // The Python SDK sets `step.status = ERROR` and/or populates
    // `step.error` for API-level failures.  Actual Python exceptions are
    // caught earlier by `classify_py_step_error`.  We only need to check
    // the structured fields — never string-match step content.
    let has_error_status = step.status == crate::types::StepStatus::Error;
    let has_error_field = !step.error.is_empty();

    if has_error_status || has_error_field {
        // Log for observability, forward to step channel for timeline
        // consumers, but do NOT break the stream.  The SDK decides when
        // the iterator is done (StopAsyncIteration or exception).
        let error_msg = format_error_message(&step);
        let http_code = step.http_code;
        let is_model_quality = is_model_quality_error(&error_msg);
        tracing::warn!(
            agent_id = ?agent_id,
            http_code,
            error = %error_msg,
            "{}",
            if is_model_quality {
                "Model produced invalid output. Stream continues (backend will retry)"
            } else {
                "Error step received. Stream continues (backend controls iteration)"
            }
        );
        // Do NOT eagerly send to error_tx here: the backend may retry after
        // a recoverable error and produce valid text in a subsequent step.
        // The end-of-stream logic in `stream_steps_to_writer` tracks
        // whether output arrived *after* the last error and only sends
        // to error_tx if no output followed (i.e. the error was fatal).
        // Forward the error step so timeline consumers see it.
        crate::streaming::ChatResponseWriter::fan_out(
            &writer.subs.step,
            &writer.step_tx,
            std::mem::take(&mut step),
            "step",
        )
        .await;
        return StepContent::Error {
            message: error_msg,
            http_code,
        };
    }

    // ── Extract summary info before forwarding consumes the data ────────
    let step_idx = step.step_index;
    let tool_names: Vec<String> = step.tool_calls.iter().map(|tc| tc.name.clone()).collect();
    let usage_summary = step.usage_metadata.as_ref().map(|u| {
        format!(
            "{}p/{}o/{}t",
            // NOLINT: zero is the correct default for missing token counts
            u.prompt_token_count.unwrap_or(0),
            // NOLINT: zero is the correct default for missing token counts
            u.candidates_token_count.unwrap_or(0),
            // NOLINT: zero is the correct default for missing token counts
            u.thoughts_token_count.unwrap_or(0),
        )
    });
    let text_len = step.content.len() + step.content_delta.len();
    let thinking_len = step.thinking.len() + step.thinking_delta.len();
    let has_tool_calls = !step.tool_calls.is_empty();
    // A "complete response" step marks the end of a model message. The SDK may
    // emit it *in addition to* the incremental delta steps, repeating the full
    // text; once handled, reset the per-message dedup accumulator so the next
    // message measures duplicates afresh.
    let is_complete_response = step.is_complete_response == Some(true);

    // ── Normal content forwarding ───────────────────────────────────────
    // Every fan-out is subscription-gated and non-fatal: a view nobody drains
    // is skipped (so it can never fill its buffer and block the writer), and a
    // dropped receiver never aborts the stream.
    forward_text(writer, &mut step, streamed_text).await;
    if is_complete_response {
        streamed_text.clear();
    }
    forward_thoughts(writer, &mut step).await;
    forward_tool_calls(writer, &mut step, agent_id).await;
    apply_step_metadata(writer, &mut step);

    crate::streaming::ChatResponseWriter::fan_out(&writer.subs.step, &writer.step_tx, step, "step")
        .await;

    // ── Structured step summary ─────────────────────────────────────────
    if !tool_names.is_empty() {
        tracing::info!(
            agent_id = ?agent_id,
            step = step_idx,
            tools = ?tool_names,
            usage = ?usage_summary,
            "tool_call"
        );
    } else if text_len > 0 || thinking_len > 0 {
        tracing::debug!(
            agent_id = ?agent_id,
            text_len,
            thinking_len,
            usage = ?usage_summary,
            "model_output"
        );
    }

    if text_len > 0 || has_tool_calls {
        StepContent::Output
    } else {
        StepContent::Empty
    }
}

/// Extract a human-readable error message from a step's error fields.
fn format_error_message(step: &crate::types::Step) -> String {
    if !step.error.is_empty() {
        return step.error.clone();
    }
    let content = if step.content.is_empty() {
        &step.content_delta
    } else {
        &step.content
    };
    format!("Step error (status={:?}): {content}", step.status)
}

/// Extract text content from the step and fan it out to the (subscribed)
/// `event_tx` (timeline), `chunk_tx` (unified), and `text_tx` (type-specific)
/// channels. Unsubscribed views are skipped; no send can block or fail fatally.
async fn forward_text(
    writer: &crate::streaming::ChatResponseWriter,
    step: &mut crate::types::Step,
    streamed_text: &mut String,
) {
    // Mirror the SDK's `receive_chunks`: text is part of the primary stream
    // only for the model's user-facing output (source == MODEL && target ==
    // USER). Text from other sources (user echoes, system) or targeted at
    // subagents/the environment is not yielded as a text chunk by the SDK, so
    // the bridge must not leak it into the text/chunk/event streams either.
    let is_model = step.source == crate::types::StepSource::Model;
    let is_target_user = step.target == crate::types::StepTarget::User;
    if !(is_model && is_target_user) {
        return;
    }

    let (raw, is_delta) = if step.content_delta.is_empty() {
        (std::mem::take(&mut step.content), false)
    } else {
        (std::mem::take(&mut step.content_delta), true)
    };
    if raw.is_empty() {
        return;
    }

    // De-duplicate a consolidated "complete response" step that repeats text
    // already streamed via deltas within the same message. The SDK yields
    // *both* the incremental deltas and a final full-content step per turn;
    // forwarding both would double the response text. Unlike the SDK's
    // delta-only `receive_chunks`, the bridge also surfaces the final full
    // content when a turn produced no deltas, so `.text()` is never empty for
    // non-streaming turns.
    let Some(text) = dedup_model_text(raw, is_delta, streamed_text) else {
        return;
    };

    crate::streaming::ChatResponseWriter::fan_out(
        &writer.subs.event,
        &writer.event_tx,
        crate::streaming::ResponseEvent::TextChunk(text.clone()),
        "event",
    )
    .await;
    crate::streaming::ChatResponseWriter::fan_out(
        &writer.subs.chunk,
        &writer.chunk_tx,
        crate::streaming::StreamChunk::Text(text.clone()),
        "chunk",
    )
    .await;
    crate::streaming::ChatResponseWriter::fan_out(&writer.subs.text, &writer.text_tx, text, "text")
        .await;
}

/// Compute the *new* model text to stream, de-duplicating a consolidated
/// full-content step against text already streamed via deltas within the same
/// model message. Returns `None` when the step carries no new text.
///
/// The Antigravity SDK emits, per turn, incremental delta steps *and* a final
/// "complete response" step whose `content` repeats the full message. Emitting
/// both doubles the response; this reconciles them:
/// - a delta is always new text;
/// - a full-content step is skipped if it merely repeats the streamed deltas,
///   or trimmed to its new tail if it grew beyond them;
/// - with no prior deltas (non-streaming turn) the content is emitted as-is.
fn dedup_model_text(raw: String, is_delta: bool, streamed: &mut String) -> Option<String> {
    if is_delta {
        // Incremental token: always new text.
        streamed.push_str(&raw);
        return Some(raw);
    }
    if streamed.is_empty() {
        // Non-streaming turn (content only, no prior deltas): emit as-is.
        streamed.push_str(&raw);
        return Some(raw);
    }
    if raw == *streamed {
        // Exact consolidation of the streamed deltas — nothing new.
        return None;
    }
    if let Some(suffix) = raw.strip_prefix(streamed.as_str()) {
        // Content grew beyond what we streamed: emit only the new tail.
        let suffix = suffix.to_owned();
        streamed.push_str(&suffix);
        return Some(suffix);
    }
    // Unrelated full content (e.g. a fresh message without deltas): emit it and
    // reset the baseline so any later consolidation is measured against it.
    streamed.clear();
    streamed.push_str(&raw);
    Some(raw)
}

/// Extract thinking content from the step and fan it out to the (subscribed)
/// `event_tx` (timeline), `chunk_tx` (unified), and `thought_tx` (type-specific)
/// channels. Unsubscribed views are skipped; no send can block or fail fatally.
async fn forward_thoughts(
    writer: &crate::streaming::ChatResponseWriter,
    step: &mut crate::types::Step,
) {
    // Mirror the SDK's `receive_chunks`: thoughts are streamed only for the
    // model's user-facing output (source == MODEL && target == USER). Thoughts
    // from other sources or targeted at subagents/the environment are not
    // yielded by the SDK, so the bridge must not leak them either.
    let is_model = step.source == crate::types::StepSource::Model;
    let is_target_user = step.target == crate::types::StepTarget::User;
    if !(is_model && is_target_user) {
        return;
    }

    let thinking = if step.thinking_delta.is_empty() {
        std::mem::take(&mut step.thinking)
    } else {
        std::mem::take(&mut step.thinking_delta)
    };
    if thinking.is_empty() {
        return;
    }
    crate::streaming::ChatResponseWriter::fan_out(
        &writer.subs.event,
        &writer.event_tx,
        crate::streaming::ResponseEvent::ThoughtChunk(thinking.clone()),
        "event",
    )
    .await;
    crate::streaming::ChatResponseWriter::fan_out(
        &writer.subs.chunk,
        &writer.chunk_tx,
        crate::streaming::StreamChunk::Thought(thinking.clone()),
        "chunk",
    )
    .await;
    crate::streaming::ChatResponseWriter::fan_out(
        &writer.subs.thought,
        &writer.thought_tx,
        thinking,
        "thought",
    )
    .await;
}

/// Extract tool calls from the step and fan each out to the (subscribed)
/// `event_tx` (timeline), `chunk_tx` (unified), and `tool_call_tx`
/// (type-specific) channels. Unsubscribed views are skipped; no send can block
/// or fail fatally.
async fn forward_tool_calls(
    writer: &crate::streaming::ChatResponseWriter,
    step: &mut crate::types::Step,
    agent_id: AgentId,
) {
    for tc in std::mem::take(&mut step.tool_calls) {
        tracing::debug!(
            agent_id = ?agent_id,
            tool = %tc.name,
            "Streaming tool call event"
        );
        let event = crate::streaming::ToolCallEvent {
            name: tc.name,
            args: tc.args,
            id: tc.id,
            canonical_path: tc.canonical_path,
        };
        crate::streaming::ChatResponseWriter::fan_out(
            &writer.subs.event,
            &writer.event_tx,
            crate::streaming::ResponseEvent::ToolCall(event.clone()),
            "event",
        )
        .await;
        crate::streaming::ChatResponseWriter::fan_out(
            &writer.subs.chunk,
            &writer.chunk_tx,
            crate::streaming::StreamChunk::ToolCall(event.clone()),
            "chunk",
        )
        .await;
        crate::streaming::ChatResponseWriter::fan_out(
            &writer.subs.tool_call,
            &writer.tool_call_tx,
            event,
            "tool_call",
        )
        .await;
    }
}

/// Transfer usage and structured-output metadata from the step to the writer's
/// shared state so the [`ChatResponseHandle`] can read them after completion.
fn apply_step_metadata(
    writer: &crate::streaming::ChatResponseWriter,
    step: &mut crate::types::Step,
) {
    if let Some(usage) = step.usage_metadata.take() {
        writer.set_usage(usage);
    }
    if let Some(out) = step.structured_output.take() {
        writer.set_structured_output(out);
    }
}

enum StepIterationResult {
    Step(Box<crate::types::Step>),
    Stop,
    Error(String),
}

/// Classify a Python step-iteration error.
///
/// Returns `Stop` for `StopAsyncIteration` (normal end of stream) or
/// `Error` for any other exception.
fn classify_py_step_error(err: &pyo3::PyErr, agent_id: AgentId) -> StepIterationResult {
    let is_stop =
        Python::attach(|py| err.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py));
    if is_stop {
        tracing::debug!(agent_id = ?agent_id, "Step stream ended (StopAsyncIteration)");
        return StepIterationResult::Stop;
    }
    let err_msg = Python::attach(|py| crate::error::classify_py_error(py, err).to_string());
    tracing::error!(agent_id = ?agent_id, error = %err_msg, "Python step iteration failed");
    StepIterationResult::Error(err_msg)
}

async fn process_next_step_iteration(
    aiter_py: &Py<PyAny>,
    agent_id: AgentId,
) -> StepIterationResult {
    let next_fut = Python::attach(|py| -> PyResult<_> {
        let aiter_bound = aiter_py.bind(py);
        let coro = aiter_bound.call_method0("__anext__")?;
        pyo3_async_runtimes::tokio::into_future(coro)
    });

    let next_fut = match next_fut {
        Ok(fut) => fut,
        Err(e) => return classify_py_step_error(&e, agent_id),
    };

    let step_py = match next_fut.await {
        Ok(obj) => obj,
        Err(e) => return classify_py_step_error(&e, agent_id),
    };

    Python::attach(|py| {
        let step_bound = step_py.bind(py);
        if step_bound.is_none() {
            return StepIterationResult::Stop;
        }
        match super::py_scripts::to_dict_py(step_bound)
            .and_then(|d| d.extract::<crate::types::Step>())
        {
            Ok(step) => StepIterationResult::Step(Box::new(step)),
            Err(e) => {
                let err_msg = format!("Failed to extract Step from Python object: {e}");
                tracing::error!(agent_id = ?agent_id, "{err_msg}");
                StepIterationResult::Error(err_msg)
            }
        }
    })
}

pub async fn stream_steps_to_writer(
    writer: &crate::streaming::ChatResponseWriter,
    agent_id: AgentId,
    aiter_py: &Py<PyAny>,
    limits: StreamLimits,
) {
    tracing::debug!(agent_id = ?agent_id, ?limits, "Starting step streaming");

    // Track the last error and whether useful output arrived *after* it.
    // If the SDK recovered (produced text/tool-calls after an error step),
    // the error is stale and should not be propagated.
    let mut state = StreamErrorState::new(limits);
    // Accumulates model text already streamed within the current message so a
    // consolidated "complete response" step is not re-emitted on top of its
    // deltas (which would double the response text).
    let mut streamed_text = String::new();

    loop {
        match process_next_step_iteration(aiter_py, agent_id).await {
            StepIterationResult::Step(step) => {
                let content =
                    forward_step_to_writer(writer, *step, agent_id, &mut streamed_text).await;
                if state.observe(&content) {
                    // Either the SDK is re-generating a deterministically-bad
                    // turn in a tight loop (each attempt a fresh RPC /
                    // subprocess with some backends), or a single generation
                    // is ruminating without ever producing output. Both hold a
                    // subprocess open and churn sockets. Stop pulling the
                    // iterator and let the orchestrator's turn-level recovery
                    // ladder (drop-the-bad-turn + corrective nudge, then respawn)
                    // break the loop instead of streaming/re-generating forever.
                    tracing::warn!(
                        agent_id = ?agent_id,
                        consecutive_model_errors = state.consecutive_model_errors,
                        consecutive_empty_steps = state.consecutive_empty_steps,
                        "Stopping stream (repeated invalid output or runaway \
                         thinking-only rumination) — handing off to orchestrator recovery"
                    );
                    break;
                }
            }
            StepIterationResult::Stop => break,
            StepIterationResult::Error(err_msg) => {
                // Python exception (not a step) — always fatal. No structured
                // HTTP status is available, so classification falls back to the
                // message text.
                send_stream_error(writer, err_msg, crate::error::HTTP_CODE_UNKNOWN);
                return;
            }
        }
    }

    // ── Stream ended (iterator exhausted, or cut short after repeated
    // invalid model output) ─────────────────────────────────────────
    // Propagate the last error only if no useful output followed it. If the
    // SDK recovered (produced text/tool-calls after the error), the error is
    // stale and the stream effectively succeeded.
    if let Some(error_msg) = state.last_error {
        if state.output_after_error {
            tracing::info!(
                agent_id = ?agent_id,
                error = %error_msg,
                "Stream recovered after error — not propagating"
            );
        } else {
            tracing::warn!(
                agent_id = ?agent_id,
                error = %error_msg,
                "Stream ended with unrecovered error — propagating"
            );
            send_stream_error(writer, error_msg, state.last_http_code);
        }
    }
}

fn send_stream_error(
    writer: &crate::streaming::ChatResponseWriter,
    message: String,
    http_code: u16,
) {
    // Use try_send to avoid deadlock: the error channel has capacity 1.
    // The writer must be dropped to close the text channel (which
    // handle.text() is waiting on). If error_tx is already full, the
    // first error wins; subsequent errors are logged but not queued.
    if let Err(e) = writer
        .error_tx
        .try_send(crate::streaming::StreamError::with_http_code(
            message, http_code,
        ))
    {
        tracing::debug!("Error channel full or closed (first error wins): {e}");
    }
}

#[cfg(test)]
#[path = "streaming_tests.rs"]
mod tests;