Skip to main content

clark_agent/
run.rs

1//! The canonical agent loop.
2//!
3//! One free function each for run/start and continue — no god-class.
4//!
5//! Shape:
6//!
7//! ```text
8//! agent_start
9//!  └ loop:                          ← outer (follow-up) loop
10//!     turn_start
11//!     [pending steering messages]   ← injected before LLM call
12//!     stream assistant response     ← StreamFn → AssistantMessage
13//!     execute tool batch (if any)   ← parallel/sequential dispatch
14//!     turn_end
15//!     ↻ until no more tool calls AND no steering ready
16//!     check follow-up               ← post-stop injection
17//!  agent_end
18//! ```
19//!
20//! Termination is unanimous-tool-vote: a batch ends the run only when
21//! every finalized tool result sets `terminate = true`. One tool wanting
22//! to stop does not stop the batch.
23
24use futures::stream::StreamExt;
25use std::time::{SystemTime, UNIX_EPOCH};
26use tokio_util::sync::CancellationToken;
27
28use crate::config::LoopConfig;
29use crate::error::{LoopError, StreamError};
30use crate::event::AgentEvent;
31use crate::exec::{execute_tool_batch, ExecutedBatch};
32use crate::plugin::TransformContext;
33use crate::stream::{ReasoningEffort, StreamErrorKind, StreamEvent, StreamRequest, ToolSchema};
34use crate::types::{
35    AgentContext, AgentMessage, AssistantContent, StopReason, ToolResultContent, Usage,
36};
37
38const EMPTY_STREAM_MAX_ATTEMPTS: u8 = 3;
39const EMPTY_STREAM_RETRY_INITIAL_DELAY: std::time::Duration = std::time::Duration::from_millis(250);
40const ZERO_OUTPUT_TRANSPORT_MAX_ATTEMPTS: u8 = 2;
41const ZERO_OUTPUT_TRANSPORT_RETRY_INITIAL_DELAY: std::time::Duration =
42    std::time::Duration::from_millis(500);
43const ZERO_OUTPUT_TRANSPORT_RECOVERY_CONTEXT: &str = "\
44[runtime context — transport recovery, not user instruction]\n\
45The previous provider attempt produced no actionable output: no visible assistant text and no usable tool call reached the runtime. \
46It may have produced private-only reasoning or an unusable burst of partial tool calls. \
47Do not continue with private reasoning only. Re-read the latest observation and immediately choose exactly one next structured tool call; \
48if the answer is ready, use the final response tool.";
49
50/// Hard cap on consecutive plain-text-fallback nudges before the loop
51/// falls back to synthesizing a terminal tool result as a last resort.
52/// Two nudges plus one synthesize keeps the recovery window bounded
53/// without leaning on a caller-configured `empty_outcome_retry_budget`.
54const MAX_PLAIN_TEXT_NUDGE_RETRIES: usize = 2;
55
56/// Outcome label for a completed run.
57///
58/// Distinguishes natural termination from budget-pressure terminations so
59/// callers (notably parent agents reading a subagent's tool result) can
60/// reason about whether the answer is complete or partial. All variants
61/// are non-error — a hard error becomes [`LoopError`].
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum LoopOutcome {
64    /// Model emitted a final assistant turn with no tool calls and no
65    /// pending steering. The natural happy path.
66    Done,
67    /// The graceful turn-limit plugin injected a wrap-up steering message
68    /// and the model produced a clean final turn within the grace window.
69    /// Result text reflects the model's deliberate close-out, not a
70    /// truncated transcript.
71    WrappedUp,
72    /// `max_iterations` was reached before the model wrapped up. The run
73    /// stopped at the cap, but earlier turns are still in the transcript.
74    /// The most recent assistant turn may have had pending tool calls.
75    HitMaxIterations,
76}
77
78impl LoopOutcome {
79    /// Whether this outcome implies a clean, non-partial final answer.
80    pub fn is_complete(self) -> bool {
81        matches!(self, LoopOutcome::Done | LoopOutcome::WrappedUp)
82    }
83
84    /// Short stable label suitable for logs and tool-result prefixes.
85    pub fn label(self) -> &'static str {
86        match self {
87            LoopOutcome::Done => "done",
88            LoopOutcome::WrappedUp => "wrapped_up",
89            LoopOutcome::HitMaxIterations => "hit_max_iterations",
90        }
91    }
92}
93
94/// Result of a completed run: emitted messages plus a typed outcome label.
95///
96/// Returned by [`run`] and [`run_continue`]. `messages` is the slice of
97/// messages produced **during this run** (not the full transcript).
98/// `outcome` lets callers distinguish a natural close from a budget-driven
99/// one without inspecting message content.
100#[derive(Debug, Clone)]
101pub struct RunResult {
102    pub messages: Vec<AgentMessage>,
103    pub outcome: LoopOutcome,
104}
105
106/// Run the loop with one or more starting prompts.
107///
108/// The prompts are appended to the context's existing message list, then
109/// the loop runs until natural stop (no more tool calls, no follow-up).
110/// Returns the messages produced **during this run** plus a typed outcome
111/// label — not the full transcript. Callers that want the full transcript
112/// should fold prior messages into their own state, or read from the
113/// event sink.
114pub async fn run(
115    prompts: Vec<AgentMessage>,
116    context: AgentContext,
117    config: &LoopConfig,
118    signal: CancellationToken,
119) -> Result<RunResult, LoopError> {
120    let mut current = context;
121    let mut new_messages = prompts.clone();
122
123    current.messages.extend(prompts.iter().cloned());
124
125    emit(config, AgentEvent::AgentStart).await;
126    if let Some(identity) = current.identity.clone() {
127        emit(config, AgentEvent::RunIdentified { identity }).await;
128    }
129    emit(config, AgentEvent::TurnStart).await;
130    for prompt in &prompts {
131        emit(
132            config,
133            AgentEvent::MessageStart {
134                message: prompt.clone(),
135            },
136        )
137        .await;
138        emit(
139            config,
140            AgentEvent::MessageEnd {
141                message: prompt.clone(),
142            },
143        )
144        .await;
145    }
146
147    let outcome = inner_run(&mut current, &mut new_messages, config, &signal).await?;
148
149    Ok(RunResult {
150        messages: new_messages,
151        outcome,
152    })
153}
154
155/// Continue an existing context without adding a new prompt.
156///
157/// Used when the trailing message is already a `User` (e.g., steering
158/// queued externally) or `ToolResult` (e.g., an out-of-band tool result
159/// was injected). Errors if the trailing message is `Assistant` — the
160/// model would not respond to its own message.
161pub async fn run_continue(
162    context: AgentContext,
163    config: &LoopConfig,
164    signal: CancellationToken,
165) -> Result<RunResult, LoopError> {
166    let last = context
167        .messages
168        .last()
169        .ok_or_else(|| LoopError::InvalidContinuation("no messages in context".into()))?;
170    if matches!(last, AgentMessage::Assistant { .. }) {
171        return Err(LoopError::InvalidContinuation(
172            "trailing message is assistant".into(),
173        ));
174    }
175
176    let mut current = context;
177    let mut new_messages = Vec::new();
178
179    emit(config, AgentEvent::AgentStart).await;
180    if let Some(identity) = current.identity.clone() {
181        emit(config, AgentEvent::RunIdentified { identity }).await;
182    }
183    emit(config, AgentEvent::TurnStart).await;
184
185    let outcome = inner_run(&mut current, &mut new_messages, config, &signal).await?;
186
187    Ok(RunResult {
188        messages: new_messages,
189        outcome,
190    })
191}
192
193// ─── Internals ─────────────────────────────────────────────────────
194
195async fn emit(config: &LoopConfig, event: AgentEvent) {
196    config.event_sink.emit(event.clone()).await;
197    for observer in &config.plugins.event_observer {
198        observer.on_event(&event).await;
199    }
200}
201
202async fn inner_run(
203    current: &mut AgentContext,
204    new_messages: &mut Vec<AgentMessage>,
205    config: &LoopConfig,
206    signal: &CancellationToken,
207) -> Result<LoopOutcome, LoopError> {
208    let mut first_turn = true;
209    let mut iterations: usize = 0;
210    let mut empty_outcomes_seen: usize = 0;
211    let mut last_turn_stopped_without_tool = false;
212    let mut plain_text_terminal_fallback_candidate: Option<AgentMessage> = None;
213
214    // Steering messages may already be queued (caller produced them
215    // before calling `run`).
216    let mut pending = collect_steering(config).await;
217
218    'outer: loop {
219        let mut has_more_tool_calls = true;
220        // Did the most recent tool batch vote terminate? Reset per
221        // outer iteration so a follow-up-driven re-entry starts clean.
222        //
223        // When the batch produces a unanimous terminator (every
224        // finalized result votes `terminate = true`), the run is over —
225        // `SteeringSource` and `FollowUpSource` plugins must NOT
226        // re-prompt the model with another LLM call. Without this
227        // guard a steering source whose firing condition lined up
228        // with the same turn (e.g. `graceful_turn_limit` reaching
229        // its soft limit on the same turn the model delivered)
230        // would inject a wrap-up message and the loop would burn
231        // another turn after a clean delivery — observed in production,
232        // where a model drifted into hallucinated content on the
233        // wrap-up re-entry after the prior batch had already produced
234        // the correct terminal delivery.
235        let mut last_batch_terminated = false;
236
237        while has_more_tool_calls || !pending.is_empty() {
238            if signal.is_cancelled() {
239                return Err(LoopError::Aborted);
240            }
241            if let Some(max) = config.max_iterations {
242                if iterations >= max {
243                    // Hit the iteration cap. Break out of the inner
244                    // loop so the follow-up sources get one last
245                    // chance to inject a terminator nudge before the
246                    // run ends. The outer loop's own cap-check (added
247                    // below) ensures we don't loop forever.
248                    break;
249                }
250            }
251            iterations += 1;
252
253            if !first_turn {
254                emit(config, AgentEvent::TurnStart).await;
255            } else {
256                first_turn = false;
257            }
258
259            // Inject any pending steering messages before the next LLM call.
260            if !pending.is_empty() {
261                for msg in pending.drain(..) {
262                    emit(
263                        config,
264                        AgentEvent::MessageStart {
265                            message: msg.clone(),
266                        },
267                    )
268                    .await;
269                    emit(
270                        config,
271                        AgentEvent::MessageEnd {
272                            message: msg.clone(),
273                        },
274                    )
275                    .await;
276                    current.messages.push(msg.clone());
277                    new_messages.push(msg);
278                }
279            }
280
281            // Stream one assistant response, applying the configured
282            // max-tokens recovery ladder if a turn comes back truncated,
283            // and the context-overflow recovery hook if the request is
284            // rejected for exceeding the model's window. `iteration` is
285            // 0-indexed and counts LLM calls within this run — `iterations`
286            // was already incremented above for cap-checking, so the
287            // 0-indexed turn number is `iterations - 1`.
288            let (assistant, turn_allowlist) =
289                stream_with_overflow_recovery(current, config, signal, iterations - 1).await?;
290            // The assistant message must land in *both* the live conversation
291            // (so the next turn's request body includes it — providers reject
292            // tool messages that don't follow a matching assistant tool_call)
293            // and the run's emitted-messages tail.
294            current.messages.push(assistant.clone());
295            new_messages.push(assistant.clone());
296
297            // Stop on stream-level error/abort. Well-behaved
298            // transports surface these as `StreamEvent::Error`, which
299            // `stream_assistant_response` converts to `LoopError`
300            // before returning. Keep this branch as a guard for
301            // transports that incorrectly finalize a `Done` message
302            // with an error stop reason.
303            let stop_reason = match &assistant {
304                AgentMessage::Assistant { stop_reason, .. } => *stop_reason,
305                _ => StopReason::Other,
306            };
307            if matches!(stop_reason, StopReason::Error | StopReason::Aborted) {
308                let loop_error = match &assistant {
309                    AgentMessage::Assistant {
310                        stop_reason: StopReason::Aborted,
311                        ..
312                    } => LoopError::Aborted,
313                    AgentMessage::Assistant { error_message, .. } => LoopError::Stream(
314                        StreamError::Transient(error_message.clone().unwrap_or_else(|| {
315                            "assistant stream ended with error stop reason".into()
316                        })),
317                    ),
318                    _ => LoopError::Stream(StreamError::Transient(
319                        "assistant stream ended with error stop reason".into(),
320                    )),
321                };
322                emit(
323                    config,
324                    AgentEvent::TurnEnd {
325                        message: assistant,
326                        tool_results: Vec::new(),
327                    },
328                )
329                .await;
330                emit(
331                    config,
332                    AgentEvent::AgentEnd {
333                        messages: new_messages.clone(),
334                    },
335                )
336                .await;
337                return Err(loop_error);
338            }
339
340            // Extract tool calls.
341            let tool_calls: Vec<_> = match &assistant {
342                AgentMessage::Assistant { content, .. } => {
343                    content.tool_calls().into_iter().cloned().collect()
344                }
345                _ => Vec::new(),
346            };
347            last_turn_stopped_without_tool = tool_calls.is_empty();
348            if last_turn_stopped_without_tool {
349                empty_outcomes_seen = empty_outcomes_seen.saturating_add(1);
350            }
351
352            let mut tool_result_messages = Vec::new();
353            has_more_tool_calls = false;
354
355            if tool_calls.is_empty() {
356                if let Some(tool_name) = config.plain_text_terminal_fallback_tool.as_deref() {
357                    let eager = config.plain_text_terminal_fallback_eager;
358                    let terminal_tool_names = config.protocol.terminal_tool_names();
359                    let narrowed_to_terminators = is_terminal_only_allowlist(
360                        turn_allowlist.as_ref(),
361                        tool_name,
362                        &terminal_tool_names,
363                    );
364                    let preserve_plain_text_candidate = plain_assistant_text(&assistant)
365                        .is_some_and(|text| should_preserve_plain_text_terminal_candidate(&text));
366                    if plain_text_terminal_fallback_candidate.is_none()
367                        && preserve_plain_text_candidate
368                    {
369                        plain_text_terminal_fallback_candidate = Some(assistant.clone());
370                    }
371                    let nudge_mode = config.plain_text_terminal_fallback_eager_nudge
372                        && eager
373                        && !narrowed_to_terminators
374                        && empty_outcomes_seen <= MAX_PLAIN_TEXT_NUDGE_RETRIES;
375                    if nudge_mode {
376                        // Catalog still contains real work tools (e.g. `plan`)
377                        // but the model emitted prose. Inject an explicit
378                        // protocol-recovery system message and force the
379                        // inner loop to re-stream rather than laundering
380                        // the prose into a synthetic `message_result`.
381                        // After MAX_PLAIN_TEXT_NUDGE_RETRIES the synthesizer
382                        // below fires as a last resort, preferring the first
383                        // preserved non-clarifying answer so retry drift does
384                        // not replace a good response with recovery chatter.
385                        //
386                        // Push directly into `current.messages` (mirrors the
387                        // synthesize path) rather than `pending`, which is
388                        // overwritten by `collect_steering` at end-of-iter.
389                        // Set `has_more_tool_calls = true` to satisfy the
390                        // inner while-loop's continuation predicate.
391                        //
392                        // The recovery prose comes from the active
393                        // `ProtocolPolicy` (which may name the product's
394                        // delivery / ask tools); the core falls back to a
395                        // generic, vocabulary-free nudge.
396                        let available_tool_names: Vec<&str> =
397                            config.tools.iter().map(|t| t.name()).collect();
398                        let nudge_text = config
399                            .protocol
400                            .plain_text_recovery_prompt(crate::protocol::PlainTextRecoveryContext {
401                                messages: &current.messages,
402                                iteration: iterations - 1,
403                                available_tool_names: &available_tool_names,
404                                terminal_fallback_tool: Some(tool_name),
405                            })
406                            .unwrap_or_else(|| {
407                                crate::protocol::DEFAULT_PLAIN_TEXT_RECOVERY_PROMPT.to_string()
408                            });
409                        let nudge = AgentMessage::System {
410                            content: nudge_text,
411                            timestamp: Some(now_ms()),
412                        };
413                        current.messages.push(nudge.clone());
414                        new_messages.push(nudge);
415                        has_more_tool_calls = true;
416                    } else if let Some(result_msg) = synthesize_plain_text_terminal_result(
417                        plain_text_terminal_fallback_candidate
418                            .as_ref()
419                            .unwrap_or(&assistant),
420                        turn_allowlist.as_ref(),
421                        tool_name,
422                        eager,
423                        &terminal_tool_names,
424                    ) {
425                        plain_text_terminal_fallback_candidate = None;
426                        last_turn_stopped_without_tool = false;
427                        empty_outcomes_seen = 0;
428                        last_batch_terminated = true;
429                        current.messages.push(result_msg.clone());
430                        new_messages.push(result_msg.clone());
431                        tool_result_messages.push(result_msg);
432                    }
433                }
434            } else {
435                let ExecutedBatch {
436                    messages,
437                    terminate,
438                } = execute_tool_batch(
439                    &assistant,
440                    tool_calls,
441                    current,
442                    config,
443                    signal,
444                    turn_allowlist.as_ref(),
445                )
446                .await?;
447
448                // A real tool batch is forward progress; the empty-outcome
449                // budget tracks being stuck, not lifetime empty stops.
450                empty_outcomes_seen = 0;
451                plain_text_terminal_fallback_candidate = None;
452                tool_result_messages = messages;
453                has_more_tool_calls = !terminate;
454                last_batch_terminated = terminate;
455
456                for result_msg in &tool_result_messages {
457                    current.messages.push(result_msg.clone());
458                    new_messages.push(result_msg.clone());
459                }
460            }
461
462            emit(
463                config,
464                AgentEvent::TurnEnd {
465                    message: assistant,
466                    tool_results: tool_result_messages,
467                },
468            )
469            .await;
470
471            // Drain any new steering messages that arrived during the
472            // turn — except when the batch just emitted a unanimous
473            // terminator. A clean terminator vote is the model's
474            // "we're done" signal; further steering would re-prompt
475            // past the delivery and let the model drift.
476            pending = if last_batch_terminated {
477                Vec::new()
478            } else {
479                collect_steering(config).await
480            };
481        }
482
483        // Inner loop exhausted: either (a) the model produced no tool
484        // calls AND no steering is queued, or (b) we hit the iteration
485        // cap. In either case, give the follow-up sources one last
486        // chance to inject a terminator nudge before declaring the
487        // run done. To prevent infinite looping when a follow-up
488        // re-arms but we're already past the cap, exit unconditionally
489        // if the cap was hit.
490        let cap_hit = config.max_iterations.is_some_and(|max| iterations >= max);
491        // Skip the follow-up source pass when the last batch
492        // terminated for the same reason steering is skipped above:
493        // a clean terminator vote means the run is done; follow-up
494        // sources exist to nudge the model toward a terminator when
495        // it failed to emit one, not to overrule one it already cast.
496        let follow_up = if last_batch_terminated {
497            Vec::new()
498        } else {
499            collect_follow_up(config).await
500        };
501        if last_turn_stopped_without_tool {
502            if let Some(budget) = config.empty_outcome_retry_budget {
503                if empty_outcomes_seen > budget {
504                    emit(
505                        config,
506                        AgentEvent::AgentEnd {
507                            messages: new_messages.clone(),
508                        },
509                    )
510                    .await;
511                    return Err(LoopError::EmptyOutcomeBudgetExhausted {
512                        budget,
513                        observed: empty_outcomes_seen,
514                    });
515                }
516            }
517        }
518        if !follow_up.is_empty() && !cap_hit {
519            pending = follow_up;
520            continue 'outer;
521        }
522        // If the cap was hit but a follow-up was produced, append it
523        // to the transcript so listeners see the final nudge — but do
524        // NOT re-enter the LLM loop. The user-facing run still ends
525        // with this message as the last appended turn.
526        if cap_hit {
527            for msg in follow_up {
528                emit(
529                    config,
530                    AgentEvent::MessageStart {
531                        message: msg.clone(),
532                    },
533                )
534                .await;
535                emit(
536                    config,
537                    AgentEvent::MessageEnd {
538                        message: msg.clone(),
539                    },
540                )
541                .await;
542                current.messages.push(msg.clone());
543                new_messages.push(msg);
544            }
545        }
546
547        break;
548    }
549
550    emit(
551        config,
552        AgentEvent::AgentEnd {
553            messages: new_messages.clone(),
554        },
555    )
556    .await;
557
558    // Classify outcome.
559    // - HitMaxIterations: hard cap was reached before the model stopped
560    //   tool-calling. The transcript may end on a turn that wanted to do
561    //   more.
562    // - WrappedUp: the graceful-turn-limit plugin fired its one-shot
563    //   wrap-up steer AND we exited naturally (cap not hit). The model
564    //   responded to the warning and produced a clean close.
565    // - Done: natural termination with no budget pressure.
566    let cap_hit_final = config.max_iterations.is_some_and(|max| iterations >= max);
567    let wrap_up_fired = config
568        .grace_signal
569        .as_ref()
570        .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed));
571    let outcome = if cap_hit_final {
572        LoopOutcome::HitMaxIterations
573    } else if wrap_up_fired {
574        LoopOutcome::WrappedUp
575    } else {
576        LoopOutcome::Done
577    };
578    Ok(outcome)
579}
580
581async fn collect_steering(config: &LoopConfig) -> Vec<AgentMessage> {
582    let mut out = Vec::new();
583    for source in &config.plugins.steering {
584        out.extend(source.next_steering_messages().await);
585    }
586    out
587}
588
589async fn collect_follow_up(config: &LoopConfig) -> Vec<AgentMessage> {
590    let mut out = Vec::new();
591    for source in &config.plugins.follow_up {
592        out.extend(source.next_follow_up_messages().await);
593    }
594    out
595}
596
597fn synthesize_plain_text_terminal_result(
598    assistant: &AgentMessage,
599    turn_allowlist: Option<&std::collections::HashSet<String>>,
600    tool_name: &str,
601    eager: bool,
602    terminal_tool_names: &std::collections::HashSet<String>,
603) -> Option<AgentMessage> {
604    // The default contract is "only convert plain text once the runtime
605    // has narrowed the catalog to terminators" — preserves strict
606    // delivery shape for everyone else. When `eager` is set the gate is
607    // lifted: the host has signalled this provider can never honor
608    // forced tool choice, so prose IS the failure mode and the nudge
609    // cycle that normally narrows the allowlist would just burn turns.
610    if !eager && !is_terminal_only_allowlist(turn_allowlist, tool_name, terminal_tool_names) {
611        return None;
612    }
613    let text = plain_assistant_text(assistant)?;
614    Some(AgentMessage::ToolResult {
615        tool_call_id: format!("plain_text_terminal_fallback_{}", now_ms()),
616        tool_name: tool_name.to_string(),
617        content: ToolResultContent::text(text),
618        is_error: false,
619        narration: Some(
620            "Converted plain assistant text into terminal delivery for an auto-tool-choice provider."
621                .to_string(),
622        ),
623        details: None,
624        timestamp: Some(now_ms()),
625    })
626}
627
628fn plain_assistant_text(assistant: &AgentMessage) -> Option<String> {
629    let AgentMessage::Assistant { content, .. } = assistant else {
630        return None;
631    };
632    let text = crate::strip_thinking_tags(&content.plain_text())
633        .trim()
634        .to_string();
635    (!text.is_empty()).then_some(text)
636}
637
638fn should_preserve_plain_text_terminal_candidate(text: &str) -> bool {
639    !looks_like_permission_or_clarification_question(text)
640}
641
642fn looks_like_permission_or_clarification_question(text: &str) -> bool {
643    let trimmed = text.trim();
644    if !trimmed.contains('?') {
645        return false;
646    }
647    let lower = trimmed.to_ascii_lowercase();
648    let starts_with_prompt = [
649        "would you like",
650        "shall i",
651        "should i",
652        "do you want",
653        "what would you like",
654        "what do you need",
655        "what's your next move",
656        "what is your next move",
657        "continue what",
658    ]
659    .iter()
660    .any(|prefix| lower.starts_with(prefix));
661    starts_with_prompt
662        || (trimmed.len() <= 500
663            && lower.contains("what")
664            && (lower.contains("next") || lower.contains("continue")))
665}
666
667/// Whether a turn's allowlist has narrowed to "terminal only" — it
668/// contains the configured fallback terminal tool and nothing but
669/// terminal/delivery tools. The set of *other* names that count as
670/// terminal comes from the active [`crate::protocol::ProtocolPolicy`]
671/// ([`crate::protocol::ProtocolPolicy::terminal_tool_names`]); the core
672/// hardcodes no product tool names. With the default policy (empty extra
673/// set) an allowlist is terminal-only exactly when it contains only the
674/// fallback tool itself.
675fn is_terminal_only_allowlist(
676    turn_allowlist: Option<&std::collections::HashSet<String>>,
677    terminal_tool: &str,
678    terminal_tool_names: &std::collections::HashSet<String>,
679) -> bool {
680    let Some(allowlist) = turn_allowlist else {
681        return false;
682    };
683    !allowlist.is_empty()
684        && allowlist.contains(terminal_tool)
685        && allowlist
686            .iter()
687            .all(|tool| tool == terminal_tool || terminal_tool_names.contains(tool))
688}
689
690// ─── Stream one assistant response ─────────────────────────────────
691
692/// Wrap [`stream_assistant_response`] with the configured max-output-
693/// tokens recovery ladder. When recovery is disabled (the default), this
694/// reduces to a single call. When enabled, a `StopReason::MaxTokens`
695/// turn is discarded and the next attempt re-streams with a larger
696/// cap until the ladder runs out or the model produces a non-truncated
697/// turn.
698///
699/// Discarded turns *do* fire `MessageStart`/`MessageEnd` from the
700/// inner streamer — listeners that care must correlate via the
701/// `OutputTokensEscalation` event that this wrapper emits before each
702/// retry. Persistence layers should treat the message that immediately
703/// precedes an `OutputTokensEscalation` as overridden by the next
704/// `MessageEnd`.
705/// Wraps [`stream_with_max_tokens_recovery`] with context-overflow
706/// recovery: when the provider rejects the request for exceeding its
707/// window ([`StreamError::ContextOverflow`]) and an overflow-recovery
708/// hook is installed, shrink `current.messages` in place (persisting it
709/// so later turns don't re-expand), emit the diff event, and retry the
710/// same LLM call — bounded by the hook's `max_attempts`. With no hook,
711/// or once attempts are exhausted, the overflow propagates unchanged
712/// (today's behavior). A recovery that fails to shrink also stops the
713/// loop rather than spinning.
714async fn stream_with_overflow_recovery(
715    current: &mut AgentContext,
716    config: &LoopConfig,
717    signal: &CancellationToken,
718    iteration: usize,
719) -> Result<(AgentMessage, Option<std::collections::HashSet<String>>), LoopError> {
720    let mut attempts: u8 = 0;
721    loop {
722        match stream_with_max_tokens_recovery(current, config, signal, iteration).await {
723            Err(LoopError::Stream(StreamError::ContextOverflow(message))) => {
724                let Some(recovery) = config.overflow_recovery.clone() else {
725                    return Err(LoopError::Stream(StreamError::ContextOverflow(message)));
726                };
727                if attempts >= recovery.max_attempts() || signal.is_cancelled() {
728                    return Err(LoopError::Stream(StreamError::ContextOverflow(message)));
729                }
730                attempts = attempts.saturating_add(1);
731
732                // Compute the observables before taking the history, so the
733                // borrow doesn't collide with the `mem::take` below.
734                let usage = last_provider_usage(&current.messages);
735                let cx = TransformContext {
736                    signal,
737                    model_id: config.model_id.as_deref().unwrap_or(""),
738                    iteration,
739                    last_provider_usage: usage.as_ref(),
740                    estimator: &*config.token_estimator,
741                };
742                let before = std::mem::take(&mut current.messages);
743                let before_size = cx.estimator.estimate_messages(&before);
744                let after = recovery.recover(before.clone(), &cx).await;
745
746                // No-progress guard: a recovery that didn't actually SHRINK the
747                // history (measured in estimated tokens, not message count —
748                // compaction can trade many messages for a summary + tail
749                // without reducing the count) would just overflow again.
750                // Surface the overflow instead of retrying forever.
751                if cx.estimator.estimate_messages(&after) >= before_size {
752                    current.messages = before;
753                    return Err(LoopError::Stream(StreamError::ContextOverflow(message)));
754                }
755                emit(
756                    config,
757                    AgentEvent::ContextTransformApplied {
758                        iteration,
759                        plugin: recovery.name(),
760                        before,
761                        after: after.clone(),
762                    },
763                )
764                .await;
765                current.messages = after;
766                // Retry the same LLM call against the shrunk history.
767            }
768            other => return other,
769        }
770    }
771}
772
773async fn stream_with_max_tokens_recovery(
774    context: &AgentContext,
775    config: &LoopConfig,
776    signal: &CancellationToken,
777    iteration: usize,
778) -> Result<(AgentMessage, Option<std::collections::HashSet<String>>), LoopError> {
779    let mut current_cap = config.max_output_tokens;
780    let mut max_tokens_attempt: u8 = 0;
781    let mut empty_stream_attempts: u8 = 0;
782    let mut zero_output_transport_attempts: u8 = 0;
783    let mut zero_output_recovery_context: Option<AgentContext> = None;
784    let mut reasoning = config.reasoning;
785
786    loop {
787        let attempt_context = zero_output_recovery_context.as_ref().unwrap_or(context);
788        let (assistant, allowlist) = match stream_assistant_response(
789            attempt_context,
790            config,
791            signal,
792            iteration,
793            current_cap,
794            reasoning,
795        )
796        .await
797        {
798            Ok(pair) => pair,
799            Err(LoopError::Stream(StreamError::Empty))
800                if empty_stream_attempts + 1 < EMPTY_STREAM_MAX_ATTEMPTS =>
801            {
802                empty_stream_attempts = empty_stream_attempts.saturating_add(1);
803                let delay = EMPTY_STREAM_RETRY_INITIAL_DELAY * u32::from(empty_stream_attempts);
804                tokio::select! {
805                    _ = signal.cancelled() => return Err(LoopError::Aborted),
806                    _ = tokio::time::sleep(delay) => {}
807                }
808                continue;
809            }
810            Err(LoopError::Stream(StreamError::ZeroOutputTransport(_)))
811                if zero_output_transport_attempts + 1 < ZERO_OUTPUT_TRANSPORT_MAX_ATTEMPTS =>
812            {
813                zero_output_transport_attempts = zero_output_transport_attempts.saturating_add(1);
814                zero_output_recovery_context =
815                    Some(context_with_zero_output_transport_recovery(context));
816                reasoning = zero_output_transport_retry_reasoning(config.reasoning);
817                let delay = ZERO_OUTPUT_TRANSPORT_RETRY_INITIAL_DELAY
818                    * u32::from(zero_output_transport_attempts);
819                tokio::select! {
820                    _ = signal.cancelled() => return Err(LoopError::Aborted),
821                    _ = tokio::time::sleep(delay) => {}
822                }
823                continue;
824            }
825            Err(err) => return Err(err),
826        };
827
828        let stop_reason = match &assistant {
829            AgentMessage::Assistant { stop_reason, .. } => *stop_reason,
830            _ => StopReason::Other,
831        };
832        if stop_reason != StopReason::MaxTokens {
833            return Ok((assistant, allowlist));
834        }
835        let Some(recovery) = config.max_output_tokens_recovery.as_ref() else {
836            return Ok((assistant, allowlist));
837        };
838        if max_tokens_attempt >= recovery.max_attempts {
839            return Ok((assistant, allowlist));
840        }
841        // No starting cap means there's no number to scale from. Refuse
842        // recovery rather than guess — the deployment hadn't pinned a
843        // cap, so the truncation came from a provider-side limit we
844        // don't know how to raise.
845        let Some(prev_cap) = current_cap else {
846            return Ok((assistant, allowlist));
847        };
848        let Some(new_cap) = recovery.next_cap(prev_cap, max_tokens_attempt) else {
849            return Ok((assistant, allowlist));
850        };
851
852        max_tokens_attempt = max_tokens_attempt.saturating_add(1);
853        emit(
854            config,
855            AgentEvent::OutputTokensEscalation {
856                attempt: max_tokens_attempt,
857                prev_cap,
858                new_cap,
859            },
860        )
861        .await;
862        current_cap = Some(new_cap);
863        // Discard the truncated `assistant` by simply not pushing it
864        // into the caller's transcript. The MessageStart/MessageEnd
865        // events for it already fired from the inner streamer; the
866        // OutputTokensEscalation event above is the listener's signal
867        // to roll the previous pair back from any projection.
868    }
869}
870
871async fn stream_assistant_response(
872    context: &AgentContext,
873    config: &LoopConfig,
874    signal: &CancellationToken,
875    iteration: usize,
876    max_output_tokens: Option<u32>,
877    reasoning: ReasoningEffort,
878) -> Result<(AgentMessage, Option<std::collections::HashSet<String>>), LoopError> {
879    // Apply context transforms in registration order. The
880    // `TransformContext` carries the cancellation signal plus a few
881    // cheap observables (model id, iteration, last-turn provider
882    // usage, token estimator) so each transform can decide locally
883    // without the loop widening the trait per-knob.
884    let last_provider_usage = last_provider_usage(&context.messages);
885    let cx = TransformContext {
886        signal,
887        model_id: config.model_id.as_deref().unwrap_or(""),
888        iteration,
889        last_provider_usage: last_provider_usage.as_ref(),
890        estimator: &*config.token_estimator,
891    };
892    let mut messages = context.messages.clone();
893    // Each transform's diff is observable so post-mortems can attribute
894    // a specific compaction (shrinker, microcompactor, history-repair,
895    // …) to the missing slice the model went on to misuse. Cloning is
896    // cheap relative to the actual transform work, and the eval-side
897    // observer is the one consumer that wants this much detail; other
898    // sinks ignore the variant.
899    for transform in &config.plugins.context_transform {
900        // Cheap pre-check: plugins that can locally decide they have
901        // nothing to do (no browser snapshots, history under budget, …)
902        // skip the clone + diff-event entirely. Default impl returns
903        // `true`, so plugins that haven't opted in still run on every
904        // round.
905        if !transform.should_run(&messages, &cx) {
906            continue;
907        }
908        let before = messages.clone();
909        messages = transform.transform(messages, &cx).await;
910        emit(
911            config,
912            AgentEvent::ContextTransformApplied {
913                iteration,
914                plugin: transform.name(),
915                before,
916                after: messages.clone(),
917            },
918        )
919        .await;
920    }
921
922    // Consult any registered ToolGate plugins for a per-turn allowlist.
923    // Each plugin returns `Some(set)` to narrow the advertised tools for
924    // exactly this LLM call. Multiple plugins compose by intersection;
925    // `None` plugins do not constrain. See `ToolGate` docs for rationale.
926    let allowlist = collect_tool_allowlist_with_events(config, iteration, &messages).await;
927
928    let tools = build_tool_schemas(config, allowlist.as_ref());
929    // Final snapshot of what the loop is about to send, captured after
930    // every transform/gate. Observers (eval per-turn dump, debugger,
931    // replay) take this as the source of truth for "what did the
932    // model see this turn?".
933    emit(
934        config,
935        AgentEvent::ProviderRequestPrepared {
936            iteration,
937            model_id: config.model_id.clone(),
938            system_prompt: context.system_prompt.clone(),
939            messages: messages.clone(),
940            tools: tools.clone(),
941            temperature: config.temperature,
942            max_output_tokens,
943        },
944    )
945    .await;
946    let request = StreamRequest {
947        system_prompt: context.system_prompt.clone(),
948        messages,
949        tools,
950        temperature: config.temperature,
951        max_output_tokens,
952        reasoning,
953        provider_extras: config
954            .provider_extras
955            .clone()
956            .unwrap_or(serde_json::Value::Null),
957        // `tool_choice: "required"` on every turn. The LLM-in-charge
958        // contract is "context → LLM → tool call → append result →
959        // repeat" — the model's job is to pick a tool, not emit
960        // narration. This assumes the catalog includes a terminal
961        // text-delivery tool, so required-on-every-turn doesn't trap the
962        // model: when the work is done it calls that delivery tool to
963        // return the answer. If the model loops on verification instead,
964        // the bug is in the catalog or prompt — not in the requirement.
965        force_tool_call: true,
966    };
967
968    let mut stream = config.stream.stream(request, signal.clone()).await;
969
970    let mut last_partial: Option<AgentMessage> = None;
971
972    while let Some(event) = stream.next().await {
973        match event {
974            StreamEvent::Start { partial } => {
975                emit(
976                    config,
977                    AgentEvent::MessageStart {
978                        message: partial.clone(),
979                    },
980                )
981                .await;
982                last_partial = Some(partial);
983            }
984            StreamEvent::Chunk(chunk) => {
985                if let Some(ref partial) = last_partial {
986                    emit(
987                        config,
988                        AgentEvent::MessageUpdate {
989                            partial: partial.clone(),
990                            chunk,
991                        },
992                    )
993                    .await;
994                }
995            }
996            StreamEvent::Done { message } => {
997                emit(
998                    config,
999                    AgentEvent::MessageEnd {
1000                        message: message.clone(),
1001                    },
1002                )
1003                .await;
1004                return Ok((message, allowlist));
1005            }
1006            StreamEvent::Error {
1007                partial,
1008                kind,
1009                message,
1010            } => {
1011                let stop_reason = match kind {
1012                    StreamErrorKind::Aborted => StopReason::Aborted,
1013                    _ => StopReason::Error,
1014                };
1015                let error_message = AgentMessage::Assistant {
1016                    content: match &partial {
1017                        AgentMessage::Assistant { content, .. } => content.clone(),
1018                        _ => AssistantContent { blocks: Vec::new() },
1019                    },
1020                    stop_reason,
1021                    error_message: Some(message.clone()),
1022                    timestamp: Some(now_ms()),
1023                    usage: None,
1024                };
1025                emit(
1026                    config,
1027                    AgentEvent::MessageEnd {
1028                        message: error_message.clone(),
1029                    },
1030                )
1031                .await;
1032                return Err(loop_error_from_stream_kind(kind, message));
1033            }
1034        }
1035    }
1036
1037    // Stream ended without `Done` or `Error`. Synthesize an empty
1038    // assistant message so the loop can recover.
1039    let empty = AgentMessage::Assistant {
1040        content: AssistantContent { blocks: Vec::new() },
1041        stop_reason: StopReason::Error,
1042        error_message: Some("stream ended without terminal event".into()),
1043        timestamp: Some(now_ms()),
1044        usage: None,
1045    };
1046    emit(
1047        config,
1048        AgentEvent::MessageEnd {
1049            message: empty.clone(),
1050        },
1051    )
1052    .await;
1053    Err(LoopError::Stream(StreamError::Empty))
1054}
1055
1056fn context_with_zero_output_transport_recovery(context: &AgentContext) -> AgentContext {
1057    let mut recovered = context.clone();
1058    recovered.messages.push(AgentMessage::System {
1059        content: ZERO_OUTPUT_TRANSPORT_RECOVERY_CONTEXT.to_string(),
1060        timestamp: Some(now_ms()),
1061    });
1062    recovered
1063}
1064
1065fn zero_output_transport_retry_reasoning(reasoning: ReasoningEffort) -> ReasoningEffort {
1066    match reasoning {
1067        ReasoningEffort::Medium | ReasoningEffort::High | ReasoningEffort::XHigh => {
1068            ReasoningEffort::Minimal
1069        }
1070        ReasoningEffort::None | ReasoningEffort::Minimal | ReasoningEffort::Low => reasoning,
1071    }
1072}
1073
1074fn loop_error_from_stream_kind(kind: StreamErrorKind, message: String) -> LoopError {
1075    // StreamFn implementations own transport retries. Once an error
1076    // reaches the loop, it is the terminal outcome of that provider
1077    // attempt and must not be reclassified as a successful assistant
1078    // turn.
1079    match kind {
1080        StreamErrorKind::Transient => LoopError::Stream(StreamError::Transient(message)),
1081        StreamErrorKind::ProviderRateLimited => {
1082            LoopError::Stream(StreamError::ProviderRateLimited(message))
1083        }
1084        StreamErrorKind::ZeroOutputTransport => {
1085            LoopError::Stream(StreamError::ZeroOutputTransport(message))
1086        }
1087        StreamErrorKind::Fatal => LoopError::Stream(StreamError::Fatal(message)),
1088        StreamErrorKind::Empty => LoopError::Stream(StreamError::Empty),
1089        StreamErrorKind::Aborted => LoopError::Aborted,
1090        StreamErrorKind::ContextOverflow => {
1091            LoopError::Stream(StreamError::ContextOverflow(message))
1092        }
1093    }
1094}
1095
1096fn now_ms() -> u64 {
1097    SystemTime::now()
1098        .duration_since(UNIX_EPOCH)
1099        .map(|d| d.as_millis() as u64)
1100        .unwrap_or(0)
1101}
1102
1103/// Walk back through `messages` and return the most recent provider
1104/// usage block reported on an assistant turn, if any. `None` on the
1105/// very first turn or when the active provider doesn't surface usage.
1106fn last_provider_usage(messages: &[AgentMessage]) -> Option<Usage> {
1107    messages.iter().rev().find_map(|message| match message {
1108        AgentMessage::Assistant {
1109            usage: Some(usage), ..
1110        } => Some(usage.clone()),
1111        _ => None,
1112    })
1113}
1114
1115fn build_tool_schemas(
1116    config: &LoopConfig,
1117    allowlist: Option<&std::collections::HashSet<String>>,
1118) -> Vec<ToolSchema> {
1119    config
1120        .tools
1121        .iter()
1122        .filter(|tool| match allowlist {
1123            Some(set) => set.contains(tool.name()),
1124            None => true,
1125        })
1126        .map(|tool| ToolSchema {
1127            name: tool.name().to_string(),
1128            description: tool.description().to_string(),
1129            parameters: tool.parameters_schema(),
1130        })
1131        .collect()
1132}
1133
1134/// Poll every registered `ToolGate` plugin and intersect their
1135/// allowlists. Returns `None` when no plugin returned an allowlist
1136/// (the common case — no narrowing). Returns `Some(set)` when at
1137/// least one plugin is gating; multiple gates compose by intersection
1138/// unless their non-empty allowlists conflict to an empty set, in which
1139/// case the highest-priority gate wins and a typed conflict event is
1140/// emitted.
1141/// Resolve the per-turn tool allowlist by composing every registered
1142/// `ToolGate` plugin (intersection) and emit one
1143/// [`AgentEvent::ToolGateApplied`] per gate so observers can attribute
1144/// the final allowlist to specific plugins.
1145async fn collect_tool_allowlist_with_events(
1146    config: &LoopConfig,
1147    iteration: usize,
1148    messages: &[AgentMessage],
1149) -> Option<std::collections::HashSet<String>> {
1150    if config.plugins.tool_gate.is_empty() {
1151        return None;
1152    }
1153    let conversation_id = config.conversation_id.as_deref();
1154    let available_tool_names: Vec<&str> = config.tools.iter().map(|t| t.name()).collect();
1155    let mut decisions: Vec<GateAllowDecision> = Vec::new();
1156    for gate in &config.plugins.tool_gate {
1157        let ctx = crate::plugin::ToolGateContext {
1158            iteration,
1159            messages,
1160            conversation_id,
1161            available_tool_names: &available_tool_names,
1162        };
1163        let decision = gate.next_turn_tool_allowlist(ctx).await;
1164        emit(
1165            config,
1166            AgentEvent::ToolGateApplied {
1167                iteration,
1168                plugin: gate.name(),
1169                allow: decision.as_ref().map(|set| {
1170                    let mut sorted: Vec<String> = set.iter().cloned().collect();
1171                    sorted.sort();
1172                    sorted
1173                }),
1174            },
1175        )
1176        .await;
1177        if let Some(set) = decision {
1178            let suppresses_advisory =
1179                gate.suppresses_advisory_gates(crate::plugin::ToolGateContext {
1180                    iteration,
1181                    messages,
1182                    conversation_id,
1183                    available_tool_names: &available_tool_names,
1184                });
1185            decisions.push(GateAllowDecision {
1186                plugin: gate.name(),
1187                priority: gate.conflict_priority(),
1188                class: gate.tool_gate_class(),
1189                suppresses_advisory,
1190                allow: set,
1191            });
1192        }
1193    }
1194    let suppression_priority = decisions
1195        .iter()
1196        .filter(|decision| decision.suppresses_advisory)
1197        .map(|decision| decision.priority)
1198        .max();
1199    let active_decisions = decisions
1200        .iter()
1201        .filter(|decision| {
1202            !matches!(
1203                suppression_priority,
1204                Some(priority)
1205                    if decision.class == crate::plugin::ToolGateClass::Advisory
1206                        && decision.priority < priority
1207            )
1208        })
1209        .collect::<Vec<_>>();
1210    let mut combined: Option<std::collections::HashSet<String>> = None;
1211    for decision in &active_decisions {
1212        combined = Some(match combined {
1213            Some(prev) => prev.intersection(&decision.allow).cloned().collect(),
1214            None => decision.allow.clone(),
1215        });
1216    }
1217    if combined.as_ref().is_some_and(|allow| allow.is_empty()) {
1218        let non_empty_decisions = active_decisions
1219            .iter()
1220            .filter(|decision| !decision.allow.is_empty())
1221            .map(|decision| (decision.plugin, decision.priority, decision.allow.clone()))
1222            .collect::<Vec<_>>();
1223        let resolved = resolve_empty_tool_gate_intersection(&non_empty_decisions);
1224        let (chosen_plugin, allow, reason) = match resolved {
1225            Some((plugin, allow, reason)) => (Some(plugin.to_string()), allow, reason),
1226            None => (
1227                None,
1228                std::collections::HashSet::new(),
1229                "all gating plugins returned empty allowlists".to_string(),
1230            ),
1231        };
1232        let sorted_allow = sorted_tool_names(&allow);
1233        emit(
1234            config,
1235            AgentEvent::ToolGateConflictResolved {
1236                iteration,
1237                plugins: active_decisions
1238                    .iter()
1239                    .map(|decision| decision.plugin.to_string())
1240                    .collect(),
1241                chosen_plugin,
1242                allow: sorted_allow,
1243                reason,
1244            },
1245        )
1246        .await;
1247        return if allow.is_empty() { None } else { Some(allow) };
1248    }
1249    combined
1250}
1251
1252struct GateAllowDecision {
1253    plugin: &'static str,
1254    priority: i32,
1255    class: crate::plugin::ToolGateClass,
1256    suppresses_advisory: bool,
1257    allow: std::collections::HashSet<String>,
1258}
1259
1260fn resolve_empty_tool_gate_intersection(
1261    decisions: &[(&'static str, i32, std::collections::HashSet<String>)],
1262) -> Option<(&'static str, std::collections::HashSet<String>, String)> {
1263    decisions
1264        .iter()
1265        .max_by(|(left_plugin, left_priority, left), (right_plugin, right_priority, right)| {
1266            left_priority
1267                .cmp(right_priority)
1268                .then_with(|| right.len().cmp(&left.len()))
1269                .then_with(|| right_plugin.cmp(left_plugin))
1270        })
1271        .map(|(plugin, priority, allow)| {
1272            (
1273                *plugin,
1274                allow.clone(),
1275                format!(
1276                    "empty intersection repaired by highest-priority owner `{plugin}` (priority {priority})"
1277                ),
1278            )
1279        })
1280}
1281
1282fn sorted_tool_names(set: &std::collections::HashSet<String>) -> Vec<String> {
1283    let mut sorted: Vec<String> = set.iter().cloned().collect();
1284    sorted.sort();
1285    sorted
1286}
1287
1288#[cfg(test)]
1289mod tests {
1290    use super::*;
1291    use crate::config::AgentBuilder;
1292    use crate::plugin::{
1293        FollowUpSource, Plugin, PluginCapabilities, ToolGate, ToolGateClass, ToolGateContext,
1294    };
1295    use crate::stream::{ReasoningEffort, StreamFn};
1296    use crate::types::{AssistantBlock, UserContent};
1297    use futures::stream::{self, BoxStream};
1298    use std::sync::{
1299        atomic::{AtomicUsize, Ordering},
1300        Arc, Mutex,
1301    };
1302
1303    fn empty_assistant_message() -> AgentMessage {
1304        AgentMessage::Assistant {
1305            content: AssistantContent { blocks: Vec::new() },
1306            stop_reason: StopReason::Other,
1307            error_message: None,
1308            timestamp: None,
1309            usage: None,
1310        }
1311    }
1312
1313    fn text_assistant_message(text: impl Into<String>) -> AgentMessage {
1314        AgentMessage::Assistant {
1315            content: AssistantContent::text(text),
1316            stop_reason: StopReason::EndTurn,
1317            error_message: None,
1318            timestamp: None,
1319            usage: None,
1320        }
1321    }
1322
1323    fn tool_call_assistant_message(name: impl Into<String>, id: impl Into<String>) -> AgentMessage {
1324        AgentMessage::Assistant {
1325            content: AssistantContent::with_tool_calls(
1326                None,
1327                vec![crate::tool::ToolCall {
1328                    id: id.into(),
1329                    name: name.into(),
1330                    arguments: serde_json::json!({}),
1331                }],
1332            ),
1333            stop_reason: StopReason::ToolUse,
1334            error_message: None,
1335            timestamp: None,
1336            usage: None,
1337        }
1338    }
1339
1340    #[derive(Default)]
1341    struct EmptyThenTextStream {
1342        calls: AtomicUsize,
1343    }
1344
1345    #[derive(Default)]
1346    struct ZeroOutputThenTextStream {
1347        calls: AtomicUsize,
1348        requests: Mutex<Vec<StreamRequest>>,
1349    }
1350
1351    impl ZeroOutputThenTextStream {
1352        fn requests(&self) -> Vec<StreamRequest> {
1353            self.requests.lock().unwrap().clone()
1354        }
1355    }
1356
1357    #[derive(Default)]
1358    struct RepeatedTextStream {
1359        calls: AtomicUsize,
1360    }
1361
1362    #[derive(Default)]
1363    struct EmptyStopsAroundProgressStream {
1364        calls: AtomicUsize,
1365    }
1366
1367    struct CountingFollowUp {
1368        remaining: AtomicUsize,
1369    }
1370
1371    struct TerminalOnlyGate;
1372    struct TerminalWithStatusGate;
1373
1374    /// A product protocol policy that declares several delivery/status
1375    /// tools (beyond the configured fallback tool) as terminal, so an
1376    /// allowlist narrowed to `{message_info, message_result}` still
1377    /// classifies as terminal-only. The core ships none of these names;
1378    /// they live behind the policy.
1379    struct TestTerminalPolicy;
1380    impl crate::protocol::ProtocolPolicy for TestTerminalPolicy {
1381        fn terminal_tool_names(&self) -> std::collections::HashSet<String> {
1382            [
1383                "message_info",
1384                "message_ask",
1385                "message_result",
1386                "terminator",
1387            ]
1388            .iter()
1389            .map(|s| s.to_string())
1390            .collect()
1391        }
1392    }
1393    struct StaticAllowGate {
1394        name: &'static str,
1395        tools: &'static [&'static str],
1396        priority: i32,
1397        class: ToolGateClass,
1398        suppresses_advisory: bool,
1399    }
1400
1401    impl Plugin for TerminalOnlyGate {
1402        fn name(&self) -> &'static str {
1403            "terminal_only_gate"
1404        }
1405
1406        fn capabilities(&self) -> PluginCapabilities {
1407            PluginCapabilities::tool_gate()
1408        }
1409    }
1410
1411    #[async_trait::async_trait]
1412    impl ToolGate for TerminalOnlyGate {
1413        async fn next_turn_tool_allowlist(
1414            &self,
1415            _ctx: ToolGateContext<'_>,
1416        ) -> Option<std::collections::HashSet<String>> {
1417            Some(["message_result".to_string()].into_iter().collect())
1418        }
1419    }
1420
1421    impl Plugin for TerminalWithStatusGate {
1422        fn name(&self) -> &'static str {
1423            "terminal_with_status_gate"
1424        }
1425
1426        fn capabilities(&self) -> PluginCapabilities {
1427            PluginCapabilities::tool_gate()
1428        }
1429    }
1430
1431    #[async_trait::async_trait]
1432    impl ToolGate for TerminalWithStatusGate {
1433        async fn next_turn_tool_allowlist(
1434            &self,
1435            _ctx: ToolGateContext<'_>,
1436        ) -> Option<std::collections::HashSet<String>> {
1437            Some(
1438                ["message_info".to_string(), "message_result".to_string()]
1439                    .into_iter()
1440                    .collect(),
1441            )
1442        }
1443    }
1444
1445    impl Plugin for StaticAllowGate {
1446        fn name(&self) -> &'static str {
1447            self.name
1448        }
1449
1450        fn capabilities(&self) -> PluginCapabilities {
1451            PluginCapabilities::tool_gate()
1452        }
1453    }
1454
1455    #[async_trait::async_trait]
1456    impl ToolGate for StaticAllowGate {
1457        fn conflict_priority(&self) -> i32 {
1458            self.priority
1459        }
1460
1461        fn tool_gate_class(&self) -> ToolGateClass {
1462            self.class
1463        }
1464
1465        fn suppresses_advisory_gates(&self, _ctx: ToolGateContext<'_>) -> bool {
1466            self.suppresses_advisory
1467        }
1468
1469        async fn next_turn_tool_allowlist(
1470            &self,
1471            _ctx: ToolGateContext<'_>,
1472        ) -> Option<std::collections::HashSet<String>> {
1473            Some(self.tools.iter().map(|name| (*name).to_string()).collect())
1474        }
1475    }
1476
1477    impl CountingFollowUp {
1478        fn new(remaining: usize) -> Self {
1479            Self {
1480                remaining: AtomicUsize::new(remaining),
1481            }
1482        }
1483    }
1484
1485    impl Plugin for CountingFollowUp {
1486        fn name(&self) -> &'static str {
1487            "counting_follow_up"
1488        }
1489
1490        fn capabilities(&self) -> PluginCapabilities {
1491            PluginCapabilities::follow_up()
1492        }
1493    }
1494
1495    #[async_trait::async_trait]
1496    impl FollowUpSource for CountingFollowUp {
1497        async fn next_follow_up_messages(&self) -> Vec<AgentMessage> {
1498            let used = self
1499                .remaining
1500                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
1501                    remaining.checked_sub(1)
1502                })
1503                .unwrap_or(0);
1504            if used == 0 {
1505                return Vec::new();
1506            }
1507            vec![AgentMessage::System {
1508                content: "retry after no-tool stop".into(),
1509                timestamp: None,
1510            }]
1511        }
1512    }
1513
1514    #[async_trait::async_trait]
1515    impl StreamFn for EmptyThenTextStream {
1516        async fn stream(
1517            &self,
1518            _request: StreamRequest,
1519            _signal: CancellationToken,
1520        ) -> BoxStream<'static, StreamEvent> {
1521            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1522            let partial = empty_assistant_message();
1523            if call == 0 {
1524                return Box::pin(stream::iter(vec![
1525                    StreamEvent::Start {
1526                        partial: partial.clone(),
1527                    },
1528                    StreamEvent::Error {
1529                        partial,
1530                        kind: StreamErrorKind::Empty,
1531                        message: "empty provider response".to_string(),
1532                    },
1533                ]));
1534            }
1535            Box::pin(stream::iter(vec![
1536                StreamEvent::Start { partial },
1537                StreamEvent::Done {
1538                    message: text_assistant_message("recovered"),
1539                },
1540            ]))
1541        }
1542    }
1543
1544    #[async_trait::async_trait]
1545    impl StreamFn for RepeatedTextStream {
1546        async fn stream(
1547            &self,
1548            _request: StreamRequest,
1549            _signal: CancellationToken,
1550        ) -> BoxStream<'static, StreamEvent> {
1551            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1552            let partial = empty_assistant_message();
1553            Box::pin(stream::iter(vec![
1554                StreamEvent::Start { partial },
1555                StreamEvent::Done {
1556                    message: text_assistant_message(format!("plain stop {call}")),
1557                },
1558            ]))
1559        }
1560    }
1561
1562    #[async_trait::async_trait]
1563    impl StreamFn for EmptyStopsAroundProgressStream {
1564        async fn stream(
1565            &self,
1566            _request: StreamRequest,
1567            _signal: CancellationToken,
1568        ) -> BoxStream<'static, StreamEvent> {
1569            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1570            let partial = empty_assistant_message();
1571            let message = match call {
1572                0 | 2 | 4 => text_assistant_message(format!("plain stop {call}")),
1573                1 | 3 => tool_call_assistant_message("progress", format!("tc-progress-{call}")),
1574                5 => tool_call_assistant_message("terminator", "tc-terminator"),
1575                other => panic!("unexpected stream call after terminal turn: {other}"),
1576            };
1577            Box::pin(stream::iter(vec![
1578                StreamEvent::Start { partial },
1579                StreamEvent::Done { message },
1580            ]))
1581        }
1582    }
1583
1584    #[async_trait::async_trait]
1585    impl StreamFn for ZeroOutputThenTextStream {
1586        async fn stream(
1587            &self,
1588            request: StreamRequest,
1589            _signal: CancellationToken,
1590        ) -> BoxStream<'static, StreamEvent> {
1591            self.requests.lock().unwrap().push(request);
1592            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1593            let partial = empty_assistant_message();
1594            if call == 0 {
1595                return Box::pin(stream::iter(vec![
1596                    StreamEvent::Start {
1597                        partial: partial.clone(),
1598                    },
1599                    StreamEvent::Error {
1600                        partial,
1601                        kind: StreamErrorKind::ZeroOutputTransport,
1602                        message: "response body decode failed before output".to_string(),
1603                    },
1604                ]));
1605            }
1606            Box::pin(stream::iter(vec![
1607                StreamEvent::Start { partial },
1608                StreamEvent::Done {
1609                    message: text_assistant_message("recovered from transport"),
1610                },
1611            ]))
1612        }
1613    }
1614
1615    /// Overflows on the first call, then (once the history is shrunk)
1616    /// returns text. Records each request so a test can assert what the
1617    /// retried call actually sent.
1618    #[derive(Default)]
1619    struct OverflowThenTextStream {
1620        calls: AtomicUsize,
1621        requests: Mutex<Vec<StreamRequest>>,
1622    }
1623
1624    #[async_trait::async_trait]
1625    impl StreamFn for OverflowThenTextStream {
1626        async fn stream(
1627            &self,
1628            request: StreamRequest,
1629            _signal: CancellationToken,
1630        ) -> BoxStream<'static, StreamEvent> {
1631            self.requests.lock().unwrap().push(request);
1632            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1633            let partial = empty_assistant_message();
1634            if call == 0 {
1635                return Box::pin(stream::iter(vec![
1636                    StreamEvent::Start {
1637                        partial: partial.clone(),
1638                    },
1639                    StreamEvent::Error {
1640                        partial,
1641                        kind: StreamErrorKind::ContextOverflow,
1642                        message: "maximum context length exceeded".to_string(),
1643                    },
1644                ]));
1645            }
1646            Box::pin(stream::iter(vec![
1647                StreamEvent::Start { partial },
1648                StreamEvent::Done {
1649                    message: text_assistant_message("recovered after shrink"),
1650                },
1651            ]))
1652        }
1653    }
1654
1655    /// Recovery that drops every message except the last — enough to
1656    /// prove the loop persists the shrink and retries.
1657    struct KeepLastRecovery {
1658        calls: Arc<AtomicUsize>,
1659    }
1660
1661    #[async_trait::async_trait]
1662    impl crate::plugin::ContextOverflowRecovery for KeepLastRecovery {
1663        async fn recover(
1664            &self,
1665            mut messages: Vec<AgentMessage>,
1666            _cx: &TransformContext<'_>,
1667        ) -> Vec<AgentMessage> {
1668            self.calls.fetch_add(1, Ordering::SeqCst);
1669            if let Some(last) = messages.pop() {
1670                vec![last]
1671            } else {
1672                messages
1673            }
1674        }
1675        fn max_attempts(&self) -> u8 {
1676            2
1677        }
1678        fn name(&self) -> &'static str {
1679            "keep_last_recovery"
1680        }
1681    }
1682
1683    #[tokio::test]
1684    async fn context_overflow_is_recovered_by_shrinking_and_retrying() {
1685        let stream = Arc::new(OverflowThenTextStream::default());
1686        let recovery_calls = Arc::new(AtomicUsize::new(0));
1687        let config = AgentBuilder::new()
1688            .stream(stream.clone())
1689            .model_id("test-model")
1690            .overflow_recovery(KeepLastRecovery {
1691                calls: recovery_calls.clone(),
1692            })
1693            .build()
1694            .expect("config builds");
1695        let mut context = AgentContext::new("system").with_messages(vec![
1696            AgentMessage::User {
1697                content: UserContent::Text("first".to_string()),
1698                timestamp: None,
1699            },
1700            AgentMessage::User {
1701                content: UserContent::Text("keep me".to_string()),
1702                timestamp: None,
1703            },
1704        ]);
1705
1706        let (assistant, _allowlist) =
1707            stream_with_overflow_recovery(&mut context, &config, &CancellationToken::new(), 0)
1708                .await
1709                .expect("overflow recovery should retry");
1710
1711        let AgentMessage::Assistant { content, .. } = assistant else {
1712            panic!("expected assistant response");
1713        };
1714        assert_eq!(content.plain_text(), "recovered after shrink");
1715        assert_eq!(stream.calls.load(Ordering::SeqCst), 2, "one retry");
1716        assert_eq!(recovery_calls.load(Ordering::SeqCst), 1);
1717        // The shrink is persisted into the live transcript…
1718        assert_eq!(context.messages.len(), 1);
1719        // …and the retried request sent only the shrunk history.
1720        let retried = &stream.requests.lock().unwrap()[1];
1721        assert_eq!(retried.messages.len(), 1);
1722    }
1723
1724    #[tokio::test]
1725    async fn context_overflow_without_recovery_propagates() {
1726        let stream = Arc::new(OverflowThenTextStream::default());
1727        let config = AgentBuilder::new()
1728            .stream(stream.clone())
1729            .model_id("test-model")
1730            .build()
1731            .expect("config builds");
1732        let mut context = AgentContext::new("system").with_messages(vec![AgentMessage::User {
1733            content: UserContent::Text("hi".to_string()),
1734            timestamp: None,
1735        }]);
1736
1737        let result =
1738            stream_with_overflow_recovery(&mut context, &config, &CancellationToken::new(), 0)
1739                .await;
1740        assert!(matches!(
1741            result,
1742            Err(LoopError::Stream(StreamError::ContextOverflow(_)))
1743        ));
1744        assert_eq!(stream.calls.load(Ordering::SeqCst), 1, "no retry");
1745    }
1746
1747    #[test]
1748    fn wrapped_up_is_complete() {
1749        assert!(LoopOutcome::Done.is_complete());
1750        assert!(LoopOutcome::WrappedUp.is_complete());
1751        assert!(!LoopOutcome::HitMaxIterations.is_complete());
1752    }
1753
1754    #[tokio::test]
1755    async fn empty_stream_response_is_retried_before_returning() {
1756        let stream = Arc::new(EmptyThenTextStream::default());
1757        let config = AgentBuilder::new()
1758            .stream(stream.clone())
1759            .model_id("test-model")
1760            .build()
1761            .expect("config builds");
1762        let context = AgentContext::new("system").with_messages(vec![AgentMessage::User {
1763            content: UserContent::Text("continue".to_string()),
1764            timestamp: None,
1765        }]);
1766
1767        let (assistant, _allowlist) =
1768            stream_with_max_tokens_recovery(&context, &config, &CancellationToken::new(), 0)
1769                .await
1770                .expect("second stream attempt should recover");
1771
1772        let AgentMessage::Assistant { content, .. } = assistant else {
1773            panic!("expected assistant response");
1774        };
1775        assert_eq!(content.plain_text(), "recovered");
1776        assert_eq!(stream.calls.load(Ordering::SeqCst), 2);
1777    }
1778
1779    #[tokio::test]
1780    async fn zero_output_transport_error_is_retried_before_returning() {
1781        let stream = Arc::new(ZeroOutputThenTextStream::default());
1782        let config = AgentBuilder::new()
1783            .stream(stream.clone())
1784            .model_id("test-model")
1785            .reasoning(ReasoningEffort::High)
1786            .build()
1787            .expect("config builds");
1788        let context = AgentContext::new("system").with_messages(vec![AgentMessage::User {
1789            content: UserContent::Text("continue".to_string()),
1790            timestamp: None,
1791        }]);
1792
1793        let (assistant, _allowlist) =
1794            stream_with_max_tokens_recovery(&context, &config, &CancellationToken::new(), 0)
1795                .await
1796                .expect("second zero-output transport attempt should recover");
1797
1798        let AgentMessage::Assistant { content, .. } = assistant else {
1799            panic!("expected assistant response");
1800        };
1801        assert_eq!(content.plain_text(), "recovered from transport");
1802        assert_eq!(stream.calls.load(Ordering::SeqCst), 2);
1803
1804        let requests = stream.requests();
1805        assert_eq!(requests.len(), 2);
1806        assert_eq!(requests[0].reasoning, ReasoningEffort::High);
1807        assert_eq!(
1808            requests[1].reasoning,
1809            ReasoningEffort::Minimal,
1810            "zero-output replay should lower high reasoning so reasoning-heavy private-only spins can produce a tool call"
1811        );
1812        assert!(
1813            requests[1].messages.iter().any(|message| matches!(
1814                message,
1815                AgentMessage::System { content, .. }
1816                    if content.contains("transport recovery")
1817                        && content.contains("no visible assistant text")
1818                        && content.contains("no usable tool call")
1819                        && content.contains("unusable burst of partial tool calls")
1820                        && content.contains("exactly one next structured tool call")
1821                        && content.contains("next structured tool call")
1822            )),
1823            "zero-output replay must carry explicit recovery context"
1824        );
1825    }
1826
1827    /// `StreamFn` that emits one assistant turn with a single
1828    /// `terminator` tool call, then panics on subsequent invocations
1829    /// — the test asserts the loop never re-enters the LLM.
1830    struct TerminatorOnlyStream {
1831        calls: AtomicUsize,
1832    }
1833
1834    impl Default for TerminatorOnlyStream {
1835        fn default() -> Self {
1836            Self {
1837                calls: AtomicUsize::new(0),
1838            }
1839        }
1840    }
1841
1842    #[async_trait::async_trait]
1843    impl StreamFn for TerminatorOnlyStream {
1844        async fn stream(
1845            &self,
1846            _request: StreamRequest,
1847            _signal: CancellationToken,
1848        ) -> BoxStream<'static, StreamEvent> {
1849            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1850            assert_eq!(
1851                call, 0,
1852                "terminate-on-turn-1 test must NOT re-enter the LLM after a successful terminator"
1853            );
1854            let partial = empty_assistant_message();
1855            let assistant = AgentMessage::Assistant {
1856                content: AssistantContent {
1857                    blocks: vec![AssistantBlock::ToolCall(crate::tool::ToolCall {
1858                        id: "tc-terminator-1".into(),
1859                        name: "terminator".into(),
1860                        arguments: serde_json::json!({}),
1861                    })],
1862                },
1863                stop_reason: StopReason::ToolUse,
1864                error_message: None,
1865                timestamp: None,
1866                usage: None,
1867            };
1868            Box::pin(stream::iter(vec![
1869                StreamEvent::Start { partial },
1870                StreamEvent::Done { message: assistant },
1871            ]))
1872        }
1873    }
1874
1875    /// Tool that always votes `terminate=true`. Mirrors the contract a
1876    /// downstream terminal/delivery tool upholds.
1877    struct TerminatorTool;
1878
1879    #[async_trait::async_trait]
1880    impl crate::tool::AgentTool for TerminatorTool {
1881        fn name(&self) -> &str {
1882            "terminator"
1883        }
1884
1885        fn description(&self) -> &str {
1886            "test terminator"
1887        }
1888
1889        fn parameters_schema(&self) -> serde_json::Value {
1890            serde_json::json!({"type": "object"})
1891        }
1892
1893        async fn execute(
1894            &self,
1895            _call_id: &str,
1896            _args: serde_json::Value,
1897            _signal: CancellationToken,
1898            _update: tokio::sync::mpsc::UnboundedSender<crate::tool::ToolResult>,
1899        ) -> Result<crate::tool::ToolResult, crate::error::ToolError> {
1900            Ok(crate::tool::ToolResult {
1901                content: vec![crate::types::ToolResultBlock::Text(
1902                    crate::types::TextContent {
1903                        text: "delivered".into(),
1904                    },
1905                )],
1906                is_error: false,
1907                details: serde_json::Value::Null,
1908                terminate: true,
1909                narration: None,
1910            })
1911        }
1912    }
1913
1914    struct ProgressTool;
1915
1916    #[async_trait::async_trait]
1917    impl crate::tool::AgentTool for ProgressTool {
1918        fn name(&self) -> &str {
1919            "progress"
1920        }
1921
1922        fn description(&self) -> &str {
1923            "test progress tool"
1924        }
1925
1926        fn parameters_schema(&self) -> serde_json::Value {
1927            serde_json::json!({"type": "object"})
1928        }
1929
1930        async fn execute(
1931            &self,
1932            _call_id: &str,
1933            _args: serde_json::Value,
1934            _signal: CancellationToken,
1935            _update: tokio::sync::mpsc::UnboundedSender<crate::tool::ToolResult>,
1936        ) -> Result<crate::tool::ToolResult, crate::error::ToolError> {
1937            Ok(crate::tool::ToolResult::text("made progress"))
1938        }
1939    }
1940
1941    /// `SteeringSource` that always returns one wrap-up message. Used
1942    /// to prove the loop does NOT poll steering after a terminator
1943    /// vote (otherwise this would re-enter the LLM and trip the
1944    /// `assert_eq!(call, 0)` in `TerminatorOnlyStream`).
1945    struct AlwaysSteer {
1946        polls: Arc<AtomicUsize>,
1947    }
1948
1949    impl Plugin for AlwaysSteer {
1950        fn name(&self) -> &'static str {
1951            "always_steer"
1952        }
1953
1954        fn capabilities(&self) -> PluginCapabilities {
1955            PluginCapabilities {
1956                steering: true,
1957                ..PluginCapabilities::default()
1958            }
1959        }
1960    }
1961
1962    #[async_trait::async_trait]
1963    impl crate::plugin::SteeringSource for AlwaysSteer {
1964        async fn next_steering_messages(&self) -> Vec<AgentMessage> {
1965            self.polls.fetch_add(1, Ordering::SeqCst);
1966            vec![AgentMessage::System {
1967                content: "wrap up now".into(),
1968                timestamp: None,
1969            }]
1970        }
1971    }
1972
1973    #[tokio::test]
1974    async fn terminator_vote_skips_post_batch_steering_collection() {
1975        // Regression: a `SteeringSource` whose firing condition lines
1976        // up with the same turn the model delivers (e.g.
1977        // `graceful_turn_limit` reaching its soft limit on the delivery
1978        // turn) used to re-enter the loop and prompt the model for
1979        // ANOTHER turn after a clean terminator. The model's drift on
1980        // that extra turn corrupted the user-visible answer in
1981        // production. With the fix, a unanimous terminator vote is a
1982        // hard exit — steering sources are not polled once the run has
1983        // decided it's done.
1984        let stream = Arc::new(TerminatorOnlyStream::default());
1985        let polls = Arc::new(AtomicUsize::new(0));
1986        let mut tool_registry = crate::tool::ToolRegistry::new();
1987        tool_registry = tool_registry.with(Arc::new(TerminatorTool));
1988        let config = AgentBuilder::new()
1989            .stream(stream.clone())
1990            .model_id("test-model")
1991            .tools(tool_registry)
1992            .steering(AlwaysSteer {
1993                polls: polls.clone(),
1994            })
1995            .build()
1996            .expect("config builds");
1997        let context = AgentContext::new("system");
1998        let prompts = vec![AgentMessage::User {
1999            content: UserContent::Text("deliver".to_string()),
2000            timestamp: None,
2001        }];
2002
2003        let result = run(prompts, context, &config, CancellationToken::new())
2004            .await
2005            .expect("run completes after one terminator turn");
2006
2007        // Exactly one LLM call — the terminator turn.
2008        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2009        // Outcome is a clean Done, not WrappedUp (no graceful flag) and
2010        // not HitMaxIterations.
2011        assert_eq!(result.outcome, LoopOutcome::Done);
2012        // Steering source is consulted exactly once — the pre-loop
2013        // priming poll at the top of `inner_run`. After the terminator
2014        // batch, `collect_steering` MUST NOT fire again.
2015        assert_eq!(
2016            polls.load(Ordering::SeqCst),
2017            1,
2018            "steering source polled more than once — terminator vote did not gate post-batch re-entry"
2019        );
2020    }
2021
2022    /// `FollowUpSource` that always emits one nudge. Counts polls so
2023    /// the test can prove `collect_follow_up` is NOT invoked after a
2024    /// terminator batch.
2025    struct AlwaysFollowUp {
2026        polls: Arc<AtomicUsize>,
2027    }
2028
2029    impl Plugin for AlwaysFollowUp {
2030        fn name(&self) -> &'static str {
2031            "always_follow_up"
2032        }
2033
2034        fn capabilities(&self) -> PluginCapabilities {
2035            PluginCapabilities::follow_up()
2036        }
2037    }
2038
2039    #[async_trait::async_trait]
2040    impl FollowUpSource for AlwaysFollowUp {
2041        async fn next_follow_up_messages(&self) -> Vec<AgentMessage> {
2042            self.polls.fetch_add(1, Ordering::SeqCst);
2043            vec![AgentMessage::System {
2044                content: "deliver something".into(),
2045                timestamp: None,
2046            }]
2047        }
2048    }
2049
2050    #[tokio::test]
2051    async fn terminator_vote_skips_post_batch_follow_up_collection() {
2052        // Mirror of the steering test for the follow-up source path.
2053        // `FollowUpSource` exists to nudge the model toward a
2054        // terminator when it failed to emit one — not to overrule a
2055        // terminator the model already cast. After a clean delivery,
2056        // follow-up must be silent.
2057        let stream = Arc::new(TerminatorOnlyStream::default());
2058        let polls = Arc::new(AtomicUsize::new(0));
2059        let mut tool_registry = crate::tool::ToolRegistry::new();
2060        tool_registry = tool_registry.with(Arc::new(TerminatorTool));
2061        let config = AgentBuilder::new()
2062            .stream(stream.clone())
2063            .model_id("test-model")
2064            .tools(tool_registry)
2065            .follow_up(AlwaysFollowUp {
2066                polls: polls.clone(),
2067            })
2068            .build()
2069            .expect("config builds");
2070        let context = AgentContext::new("system");
2071        let prompts = vec![AgentMessage::User {
2072            content: UserContent::Text("deliver".to_string()),
2073            timestamp: None,
2074        }];
2075
2076        let result = run(prompts, context, &config, CancellationToken::new())
2077            .await
2078            .expect("run completes after one terminator turn");
2079
2080        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2081        assert_eq!(result.outcome, LoopOutcome::Done);
2082        assert_eq!(
2083            polls.load(Ordering::SeqCst),
2084            0,
2085            "follow-up source polled after a terminator vote — terminator did not gate post-batch re-entry"
2086        );
2087    }
2088
2089    #[tokio::test]
2090    async fn exhausted_empty_outcome_budget_returns_typed_loop_error() {
2091        let stream = Arc::new(RepeatedTextStream::default());
2092        let config = AgentBuilder::new()
2093            .stream(stream.clone())
2094            .model_id("test-model")
2095            .empty_outcome_retry_budget(1)
2096            .follow_up(CountingFollowUp::new(1))
2097            .build()
2098            .expect("config builds");
2099        let context = AgentContext::new("system");
2100        let prompts = vec![AgentMessage::User {
2101            content: UserContent::Text("continue".to_string()),
2102            timestamp: None,
2103        }];
2104
2105        let err = run(prompts, context, &config, CancellationToken::new())
2106            .await
2107            .expect_err("second no-tool stop should exhaust the budget");
2108
2109        assert!(
2110            matches!(
2111                err,
2112                LoopError::EmptyOutcomeBudgetExhausted {
2113                    budget: 1,
2114                    observed: 2,
2115                }
2116            ),
2117            "unexpected error: {err:?}"
2118        );
2119        assert_eq!(stream.calls.load(Ordering::SeqCst), 2);
2120    }
2121
2122    #[tokio::test]
2123    async fn empty_tool_gate_intersection_prefers_delivery_repair_owner() {
2124        let (sink, mut rx) = crate::event::ChannelSink::new();
2125        let config = AgentBuilder::new()
2126            .stream(Arc::new(RepeatedTextStream::default()))
2127            .event_sink(Arc::new(sink))
2128            .tool_gate_arc(Arc::new(StaticAllowGate {
2129                name: "delivery_repair_gate",
2130                tools: &["browser_interact"],
2131                priority: 100,
2132                class: ToolGateClass::Required,
2133                suppresses_advisory: false,
2134            }))
2135            .tool_gate_arc(Arc::new(StaticAllowGate {
2136                name: "terminal_message_guard",
2137                tools: &["message_result"],
2138                priority: 10,
2139                class: ToolGateClass::Required,
2140                suppresses_advisory: false,
2141            }))
2142            .build()
2143            .expect("config builds");
2144
2145        let allow = collect_tool_allowlist_with_events(&config, 3, &[])
2146            .await
2147            .expect("conflict repair should keep a non-empty allowlist");
2148
2149        assert_eq!(
2150            allow,
2151            ["browser_interact".to_string()].into_iter().collect()
2152        );
2153
2154        let mut saw_conflict = false;
2155        while let Ok(event) = rx.try_recv() {
2156            if let AgentEvent::ToolGateConflictResolved {
2157                chosen_plugin,
2158                allow,
2159                ..
2160            } = event
2161            {
2162                saw_conflict = true;
2163                assert_eq!(chosen_plugin.as_deref(), Some("delivery_repair_gate"));
2164                assert_eq!(allow, vec!["browser_interact".to_string()]);
2165            }
2166        }
2167        assert!(saw_conflict, "tool-gate deadlock should be diagnosable");
2168    }
2169
2170    #[tokio::test]
2171    async fn repair_owner_suppresses_advisory_gate_before_plan_only_intersection() {
2172        let config = AgentBuilder::new()
2173            .stream(Arc::new(RepeatedTextStream::default()))
2174            .tool_gate_arc(Arc::new(StaticAllowGate {
2175                name: "delivery_repair_gate",
2176                tools: &["plan", "file_write"],
2177                priority: 100,
2178                class: ToolGateClass::Required,
2179                suppresses_advisory: true,
2180            }))
2181            .tool_gate_arc(Arc::new(StaticAllowGate {
2182                name: "wrap_up_gate",
2183                tools: &["plan", "message_result", "message_ask"],
2184                priority: 0,
2185                class: ToolGateClass::Advisory,
2186                suppresses_advisory: false,
2187            }))
2188            .build()
2189            .expect("config builds");
2190
2191        let allow = collect_tool_allowlist_with_events(&config, 3, &[])
2192            .await
2193            .expect("repair owner should keep its own allowlist");
2194
2195        assert_eq!(
2196            allow,
2197            ["plan".to_string(), "file_write".to_string()]
2198                .into_iter()
2199                .collect()
2200        );
2201    }
2202
2203    #[tokio::test]
2204    async fn productive_tool_batch_resets_empty_outcome_budget() {
2205        let stream = Arc::new(EmptyStopsAroundProgressStream::default());
2206        let mut tool_registry = crate::tool::ToolRegistry::new();
2207        tool_registry = tool_registry
2208            .with(Arc::new(ProgressTool))
2209            .with(Arc::new(TerminatorTool));
2210        let config = AgentBuilder::new()
2211            .stream(stream.clone())
2212            .model_id("test-model")
2213            .tools(tool_registry)
2214            .empty_outcome_retry_budget(1)
2215            .follow_up(CountingFollowUp::new(3))
2216            .build()
2217            .expect("config builds");
2218        let context = AgentContext::new("system");
2219        let prompts = vec![AgentMessage::User {
2220            content: UserContent::Text("continue".to_string()),
2221            timestamp: None,
2222        }];
2223
2224        let result = run(prompts, context, &config, CancellationToken::new())
2225            .await
2226            .expect("productive tool batches should reset the empty-outcome budget");
2227
2228        assert_eq!(result.outcome, LoopOutcome::Done);
2229        assert_eq!(stream.calls.load(Ordering::SeqCst), 6);
2230    }
2231
2232    #[tokio::test]
2233    async fn terminal_only_plain_text_fallback_synthesizes_terminal_result() {
2234        let stream = Arc::new(RepeatedTextStream::default());
2235        let mut tool_registry = crate::tool::ToolRegistry::new();
2236        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2237        let config = AgentBuilder::new()
2238            .stream(stream.clone())
2239            .model_id("auto-tool-provider")
2240            .tools(tool_registry)
2241            .tool_gate_arc(Arc::new(TerminalOnlyGate))
2242            .plain_text_terminal_fallback_tool("message_result")
2243            .empty_outcome_retry_budget(0)
2244            .build()
2245            .expect("config builds");
2246        let context = AgentContext::new("system");
2247        let prompts = vec![AgentMessage::User {
2248            content: UserContent::Text("answer directly".to_string()),
2249            timestamp: None,
2250        }];
2251
2252        let result = run(prompts, context, &config, CancellationToken::new())
2253            .await
2254            .expect("plain text should be converted on terminal-only turn");
2255
2256        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2257        assert_eq!(result.outcome, LoopOutcome::Done);
2258        assert!(result.messages.iter().any(|message| matches!(
2259            message,
2260            AgentMessage::ToolResult {
2261                tool_name,
2262                content,
2263                is_error: false,
2264                ..
2265            } if tool_name == "message_result"
2266                && content.plain_text() == "plain stop 0"
2267        )));
2268    }
2269
2270    #[tokio::test]
2271    async fn eager_plain_text_fallback_fires_without_terminal_only_allowlist() {
2272        // Providers in the "auto-when-forced" class can never be
2273        // wire-forced into a tool call, so prose IS their failure mode.
2274        // The eager flag lifts the "allowlist must already be narrowed"
2275        // precondition so the fallback fires on the FIRST plain-text
2276        // stop instead of after a narrowing gate has burned 2-3 nudge
2277        // turns.
2278        //
2279        // No `tool_gate_arc` is installed in this test, so the catalog
2280        // stays at the full registry — exactly the situation where the
2281        // non-eager path would refuse to convert and the run would die
2282        // on the empty-outcome budget.
2283        let stream = Arc::new(RepeatedTextStream::default());
2284        let mut tool_registry = crate::tool::ToolRegistry::new();
2285        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2286        let config = AgentBuilder::new()
2287            .stream(stream.clone())
2288            .model_id("auto-tool-provider-eager")
2289            .tools(tool_registry)
2290            .plain_text_terminal_fallback_tool("message_result")
2291            .plain_text_terminal_fallback_eager(true)
2292            .empty_outcome_retry_budget(0)
2293            .build()
2294            .expect("config builds");
2295        let context = AgentContext::new("system");
2296        let prompts = vec![AgentMessage::User {
2297            content: UserContent::Text("answer directly".to_string()),
2298            timestamp: None,
2299        }];
2300
2301        let result = run(prompts, context, &config, CancellationToken::new())
2302            .await
2303            .expect("eager fallback should convert plain text on first stop");
2304
2305        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2306        assert_eq!(result.outcome, LoopOutcome::Done);
2307        assert!(result.messages.iter().any(|message| matches!(
2308            message,
2309            AgentMessage::ToolResult {
2310                tool_name,
2311                content,
2312                is_error: false,
2313                ..
2314            } if tool_name == "message_result"
2315                && content.plain_text() == "plain stop 0"
2316        )));
2317    }
2318
2319    #[tokio::test]
2320    async fn eager_nudge_mode_injects_protocol_recovery_before_synthesizing() {
2321        // With `plain_text_terminal_fallback_eager_nudge(true)` the eager
2322        // path nudges the model with a protocol-recovery system message
2323        // on each consecutive plain-text stop, up to
2324        // `MAX_PLAIN_TEXT_NUDGE_RETRIES`. After the cap a synthesizer
2325        // fires as a last resort so the run still terminates with the
2326        // model's prose as the delivered text — never silently, never
2327        // forever. Verifies the recovery path is observable in the
2328        // emitted message stream (the model sees the nudges in context)
2329        // and that the synthesizer ultimately delivers the first
2330        // substantive plain-text answer, not later recovery drift.
2331        let stream = Arc::new(RepeatedTextStream::default());
2332        let mut tool_registry = crate::tool::ToolRegistry::new();
2333        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2334        let config = AgentBuilder::new()
2335            .stream(stream.clone())
2336            .model_id("auto-tool-provider-eager-nudge")
2337            .tools(tool_registry)
2338            .plain_text_terminal_fallback_tool("message_result")
2339            .plain_text_terminal_fallback_eager(true)
2340            .plain_text_terminal_fallback_eager_nudge(true)
2341            .build()
2342            .expect("config builds");
2343        let context = AgentContext::new("system");
2344        let prompts = vec![AgentMessage::User {
2345            content: UserContent::Text("answer directly".to_string()),
2346            timestamp: None,
2347        }];
2348
2349        let result = run(prompts, context, &config, CancellationToken::new())
2350            .await
2351            .expect("nudge mode should eventually synthesize after retries");
2352
2353        // MAX_PLAIN_TEXT_NUDGE_RETRIES = 2 → two nudges fire, then on the
2354        // third empty stop the synthesizer takes over. Total LLM calls = 3.
2355        assert_eq!(stream.calls.load(Ordering::SeqCst), 3);
2356        assert_eq!(result.outcome, LoopOutcome::Done);
2357
2358        let nudge_count = result
2359            .messages
2360            .iter()
2361            .filter(|m| matches!(m, AgentMessage::System { content, .. } if content == crate::protocol::DEFAULT_PLAIN_TEXT_RECOVERY_PROMPT))
2362            .count();
2363        assert_eq!(
2364            nudge_count, 2,
2365            "expected two protocol-recovery system messages in the run output, got {nudge_count}",
2366        );
2367
2368        let synthesized_text = result
2369            .messages
2370            .iter()
2371            .find_map(|message| match message {
2372                AgentMessage::ToolResult {
2373                    tool_name,
2374                    content,
2375                    is_error: false,
2376                    ..
2377                } if tool_name == "message_result" => Some(content.plain_text()),
2378                _ => None,
2379            })
2380            .expect("a terminal tool result should be synthesized as last resort");
2381        assert_eq!(
2382            synthesized_text, "plain stop 0",
2383            "synthesizer should deliver the first preserved plain text, not later recovery drift",
2384        );
2385    }
2386
2387    #[test]
2388    fn plain_text_fallback_candidate_skips_obvious_clarifying_questions() {
2389        assert!(!should_preserve_plain_text_terminal_candidate(
2390            "Continue what, exactly? What's your next move?"
2391        ));
2392        assert!(!should_preserve_plain_text_terminal_candidate(
2393            "Would you like me to proceed?"
2394        ));
2395        assert!(should_preserve_plain_text_terminal_candidate(
2396            "# Machine Learning\n\nMachine learning is the branch of artificial intelligence that studies systems which improve from data."
2397        ));
2398    }
2399
2400    #[tokio::test]
2401    async fn non_eager_plain_text_fallback_still_requires_narrowed_allowlist() {
2402        // Default behaviour preserved: when eager is NOT set and the
2403        // turn allowlist is the full catalog, plain text is NOT
2404        // converted — the run dies on the empty-outcome budget,
2405        // matching the pre-eager contract for every other model.
2406        let stream = Arc::new(RepeatedTextStream::default());
2407        let mut tool_registry = crate::tool::ToolRegistry::new();
2408        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2409        let config = AgentBuilder::new()
2410            .stream(stream.clone())
2411            .model_id("non-eager-provider")
2412            .tools(tool_registry)
2413            .plain_text_terminal_fallback_tool("message_result")
2414            // Eager NOT set → defaults to false.
2415            .empty_outcome_retry_budget(0)
2416            .build()
2417            .expect("config builds");
2418        let context = AgentContext::new("system");
2419        let prompts = vec![AgentMessage::User {
2420            content: UserContent::Text("answer directly".to_string()),
2421            timestamp: None,
2422        }];
2423
2424        let err = run(prompts, context, &config, CancellationToken::new())
2425            .await
2426            .expect_err("non-eager fallback must not convert without narrowed allowlist");
2427
2428        assert!(
2429            matches!(err, LoopError::EmptyOutcomeBudgetExhausted { .. }),
2430            "unexpected error: {err:?}"
2431        );
2432    }
2433
2434    #[tokio::test]
2435    async fn terminal_plain_text_fallback_allows_status_delivery_gate() {
2436        let stream = Arc::new(RepeatedTextStream::default());
2437        let mut tool_registry = crate::tool::ToolRegistry::new();
2438        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2439        let config = AgentBuilder::new()
2440            .stream(stream.clone())
2441            .model_id("auto-tool-provider")
2442            .tools(tool_registry)
2443            .protocol_policy(Arc::new(TestTerminalPolicy))
2444            .tool_gate_arc(Arc::new(TerminalWithStatusGate))
2445            .plain_text_terminal_fallback_tool("message_result")
2446            .empty_outcome_retry_budget(0)
2447            .build()
2448            .expect("config builds");
2449        let context = AgentContext::new("system");
2450        let prompts = vec![AgentMessage::User {
2451            content: UserContent::Text("answer directly".to_string()),
2452            timestamp: None,
2453        }];
2454
2455        let result = run(prompts, context, &config, CancellationToken::new())
2456            .await
2457            .expect(
2458                "plain text should be converted when only status and terminal tools are allowed",
2459            );
2460
2461        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2462        assert_eq!(result.outcome, LoopOutcome::Done);
2463        assert!(result.messages.iter().any(|message| matches!(
2464            message,
2465            AgentMessage::ToolResult {
2466                tool_name,
2467                content,
2468                is_error: false,
2469                ..
2470            } if tool_name == "message_result"
2471                && content.plain_text() == "plain stop 0"
2472        )));
2473    }
2474
2475    struct TerminalNamedTool(&'static str);
2476
2477    #[async_trait::async_trait]
2478    impl crate::tool::AgentTool for TerminalNamedTool {
2479        fn name(&self) -> &str {
2480            self.0
2481        }
2482
2483        fn description(&self) -> &str {
2484            "test terminal tool"
2485        }
2486
2487        fn parameters_schema(&self) -> serde_json::Value {
2488            serde_json::json!({"type": "object"})
2489        }
2490
2491        async fn execute(
2492            &self,
2493            _call_id: &str,
2494            _args: serde_json::Value,
2495            _signal: CancellationToken,
2496            _update: tokio::sync::mpsc::UnboundedSender<crate::tool::ToolResult>,
2497        ) -> Result<crate::tool::ToolResult, crate::error::ToolError> {
2498            Ok(crate::tool::ToolResult {
2499                content: vec![crate::types::ToolResultBlock::Text(
2500                    crate::types::TextContent {
2501                        text: "not used".into(),
2502                    },
2503                )],
2504                is_error: false,
2505                details: serde_json::Value::Null,
2506                terminate: true,
2507                narration: None,
2508            })
2509        }
2510    }
2511}