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::InconsistentToolHistory => {
1089            LoopError::Stream(StreamError::InconsistentToolHistory(message))
1090        }
1091        StreamErrorKind::Empty => LoopError::Stream(StreamError::Empty),
1092        StreamErrorKind::Aborted => LoopError::Aborted,
1093        StreamErrorKind::ContextOverflow => {
1094            LoopError::Stream(StreamError::ContextOverflow(message))
1095        }
1096    }
1097}
1098
1099fn now_ms() -> u64 {
1100    SystemTime::now()
1101        .duration_since(UNIX_EPOCH)
1102        .map(|d| d.as_millis() as u64)
1103        .unwrap_or(0)
1104}
1105
1106/// Walk back through `messages` and return the most recent provider
1107/// usage block reported on an assistant turn, if any. `None` on the
1108/// very first turn or when the active provider doesn't surface usage.
1109fn last_provider_usage(messages: &[AgentMessage]) -> Option<Usage> {
1110    messages.iter().rev().find_map(|message| match message {
1111        AgentMessage::Assistant {
1112            usage: Some(usage), ..
1113        } => Some(usage.clone()),
1114        _ => None,
1115    })
1116}
1117
1118fn build_tool_schemas(
1119    config: &LoopConfig,
1120    allowlist: Option<&std::collections::HashSet<String>>,
1121) -> Vec<ToolSchema> {
1122    config
1123        .tools
1124        .iter()
1125        .filter(|tool| match allowlist {
1126            Some(set) => set.contains(tool.name()),
1127            None => true,
1128        })
1129        .map(|tool| ToolSchema {
1130            name: tool.name().to_string(),
1131            description: tool.description().to_string(),
1132            parameters: tool.parameters_schema(),
1133        })
1134        .collect()
1135}
1136
1137/// Poll every registered `ToolGate` plugin and intersect their
1138/// allowlists. Returns `None` when no plugin returned an allowlist
1139/// (the common case — no narrowing). Returns `Some(set)` when at
1140/// least one plugin is gating; multiple gates compose by intersection
1141/// unless their non-empty allowlists conflict to an empty set, in which
1142/// case the highest-priority gate wins and a typed conflict event is
1143/// emitted.
1144/// Resolve the per-turn tool allowlist by composing every registered
1145/// `ToolGate` plugin (intersection) and emit one
1146/// [`AgentEvent::ToolGateApplied`] per gate so observers can attribute
1147/// the final allowlist to specific plugins.
1148async fn collect_tool_allowlist_with_events(
1149    config: &LoopConfig,
1150    iteration: usize,
1151    messages: &[AgentMessage],
1152) -> Option<std::collections::HashSet<String>> {
1153    if config.plugins.tool_gate.is_empty() {
1154        return None;
1155    }
1156    let conversation_id = config.conversation_id.as_deref();
1157    let available_tool_names: Vec<&str> = config.tools.iter().map(|t| t.name()).collect();
1158    let mut decisions: Vec<GateAllowDecision> = Vec::new();
1159    for gate in &config.plugins.tool_gate {
1160        let ctx = crate::plugin::ToolGateContext {
1161            iteration,
1162            messages,
1163            conversation_id,
1164            available_tool_names: &available_tool_names,
1165        };
1166        let decision = gate.next_turn_tool_allowlist(ctx).await;
1167        emit(
1168            config,
1169            AgentEvent::ToolGateApplied {
1170                iteration,
1171                plugin: gate.name(),
1172                allow: decision.as_ref().map(|set| {
1173                    let mut sorted: Vec<String> = set.iter().cloned().collect();
1174                    sorted.sort();
1175                    sorted
1176                }),
1177            },
1178        )
1179        .await;
1180        if let Some(set) = decision {
1181            let suppresses_advisory =
1182                gate.suppresses_advisory_gates(crate::plugin::ToolGateContext {
1183                    iteration,
1184                    messages,
1185                    conversation_id,
1186                    available_tool_names: &available_tool_names,
1187                });
1188            decisions.push(GateAllowDecision {
1189                plugin: gate.name(),
1190                priority: gate.conflict_priority(),
1191                class: gate.tool_gate_class(),
1192                suppresses_advisory,
1193                allow: set,
1194            });
1195        }
1196    }
1197    let suppression_priority = decisions
1198        .iter()
1199        .filter(|decision| decision.suppresses_advisory)
1200        .map(|decision| decision.priority)
1201        .max();
1202    let active_decisions = decisions
1203        .iter()
1204        .filter(|decision| {
1205            !matches!(
1206                suppression_priority,
1207                Some(priority)
1208                    if decision.class == crate::plugin::ToolGateClass::Advisory
1209                        && decision.priority < priority
1210            )
1211        })
1212        .collect::<Vec<_>>();
1213    let mut combined: Option<std::collections::HashSet<String>> = None;
1214    for decision in &active_decisions {
1215        combined = Some(match combined {
1216            Some(prev) => prev.intersection(&decision.allow).cloned().collect(),
1217            None => decision.allow.clone(),
1218        });
1219    }
1220    if combined.as_ref().is_some_and(|allow| allow.is_empty()) {
1221        let non_empty_decisions = active_decisions
1222            .iter()
1223            .filter(|decision| !decision.allow.is_empty())
1224            .map(|decision| (decision.plugin, decision.priority, decision.allow.clone()))
1225            .collect::<Vec<_>>();
1226        let resolved = resolve_empty_tool_gate_intersection(&non_empty_decisions);
1227        let (chosen_plugin, allow, reason) = match resolved {
1228            Some((plugin, allow, reason)) => (Some(plugin.to_string()), allow, reason),
1229            None => (
1230                None,
1231                std::collections::HashSet::new(),
1232                "all gating plugins returned empty allowlists".to_string(),
1233            ),
1234        };
1235        let sorted_allow = sorted_tool_names(&allow);
1236        emit(
1237            config,
1238            AgentEvent::ToolGateConflictResolved {
1239                iteration,
1240                plugins: active_decisions
1241                    .iter()
1242                    .map(|decision| decision.plugin.to_string())
1243                    .collect(),
1244                chosen_plugin,
1245                allow: sorted_allow,
1246                reason,
1247            },
1248        )
1249        .await;
1250        return if allow.is_empty() { None } else { Some(allow) };
1251    }
1252    combined
1253}
1254
1255struct GateAllowDecision {
1256    plugin: &'static str,
1257    priority: i32,
1258    class: crate::plugin::ToolGateClass,
1259    suppresses_advisory: bool,
1260    allow: std::collections::HashSet<String>,
1261}
1262
1263fn resolve_empty_tool_gate_intersection(
1264    decisions: &[(&'static str, i32, std::collections::HashSet<String>)],
1265) -> Option<(&'static str, std::collections::HashSet<String>, String)> {
1266    decisions
1267        .iter()
1268        .max_by(|(left_plugin, left_priority, left), (right_plugin, right_priority, right)| {
1269            left_priority
1270                .cmp(right_priority)
1271                .then_with(|| right.len().cmp(&left.len()))
1272                .then_with(|| right_plugin.cmp(left_plugin))
1273        })
1274        .map(|(plugin, priority, allow)| {
1275            (
1276                *plugin,
1277                allow.clone(),
1278                format!(
1279                    "empty intersection repaired by highest-priority owner `{plugin}` (priority {priority})"
1280                ),
1281            )
1282        })
1283}
1284
1285fn sorted_tool_names(set: &std::collections::HashSet<String>) -> Vec<String> {
1286    let mut sorted: Vec<String> = set.iter().cloned().collect();
1287    sorted.sort();
1288    sorted
1289}
1290
1291#[cfg(test)]
1292mod tests {
1293    use super::*;
1294    use crate::config::AgentBuilder;
1295    use crate::plugin::{
1296        FollowUpSource, Plugin, PluginCapabilities, ToolGate, ToolGateClass, ToolGateContext,
1297    };
1298    use crate::stream::{ReasoningEffort, StreamFn};
1299    use crate::types::{AssistantBlock, UserContent};
1300    use futures::stream::{self, BoxStream};
1301    use std::sync::{
1302        atomic::{AtomicUsize, Ordering},
1303        Arc, Mutex,
1304    };
1305
1306    #[test]
1307    fn inconsistent_tool_history_stream_kind_stays_typed_at_loop_boundary() {
1308        let error = loop_error_from_stream_kind(
1309            StreamErrorKind::InconsistentToolHistory,
1310            "interleaved tool result batch".into(),
1311        );
1312
1313        assert!(matches!(
1314            error,
1315            LoopError::Stream(StreamError::InconsistentToolHistory(message))
1316                if message == "interleaved tool result batch"
1317        ));
1318    }
1319
1320    fn empty_assistant_message() -> AgentMessage {
1321        AgentMessage::Assistant {
1322            content: AssistantContent { blocks: Vec::new() },
1323            stop_reason: StopReason::Other,
1324            error_message: None,
1325            timestamp: None,
1326            usage: None,
1327        }
1328    }
1329
1330    fn text_assistant_message(text: impl Into<String>) -> AgentMessage {
1331        AgentMessage::Assistant {
1332            content: AssistantContent::text(text),
1333            stop_reason: StopReason::EndTurn,
1334            error_message: None,
1335            timestamp: None,
1336            usage: None,
1337        }
1338    }
1339
1340    fn tool_call_assistant_message(name: impl Into<String>, id: impl Into<String>) -> AgentMessage {
1341        AgentMessage::Assistant {
1342            content: AssistantContent::with_tool_calls(
1343                None,
1344                vec![crate::tool::ToolCall {
1345                    id: id.into(),
1346                    name: name.into(),
1347                    arguments: serde_json::json!({}),
1348                }],
1349            ),
1350            stop_reason: StopReason::ToolUse,
1351            error_message: None,
1352            timestamp: None,
1353            usage: None,
1354        }
1355    }
1356
1357    #[derive(Default)]
1358    struct EmptyThenTextStream {
1359        calls: AtomicUsize,
1360    }
1361
1362    #[derive(Default)]
1363    struct ZeroOutputThenTextStream {
1364        calls: AtomicUsize,
1365        requests: Mutex<Vec<StreamRequest>>,
1366    }
1367
1368    impl ZeroOutputThenTextStream {
1369        fn requests(&self) -> Vec<StreamRequest> {
1370            self.requests.lock().unwrap().clone()
1371        }
1372    }
1373
1374    #[derive(Default)]
1375    struct RepeatedTextStream {
1376        calls: AtomicUsize,
1377    }
1378
1379    #[derive(Default)]
1380    struct EmptyStopsAroundProgressStream {
1381        calls: AtomicUsize,
1382    }
1383
1384    struct CountingFollowUp {
1385        remaining: AtomicUsize,
1386    }
1387
1388    struct TerminalOnlyGate;
1389    struct TerminalWithStatusGate;
1390
1391    /// A product protocol policy that declares several delivery/status
1392    /// tools (beyond the configured fallback tool) as terminal, so an
1393    /// allowlist narrowed to `{message_info, message_result}` still
1394    /// classifies as terminal-only. The core ships none of these names;
1395    /// they live behind the policy.
1396    struct TestTerminalPolicy;
1397    impl crate::protocol::ProtocolPolicy for TestTerminalPolicy {
1398        fn terminal_tool_names(&self) -> std::collections::HashSet<String> {
1399            [
1400                "message_info",
1401                "message_ask",
1402                "message_result",
1403                "terminator",
1404            ]
1405            .iter()
1406            .map(|s| s.to_string())
1407            .collect()
1408        }
1409    }
1410    struct StaticAllowGate {
1411        name: &'static str,
1412        tools: &'static [&'static str],
1413        priority: i32,
1414        class: ToolGateClass,
1415        suppresses_advisory: bool,
1416    }
1417
1418    impl Plugin for TerminalOnlyGate {
1419        fn name(&self) -> &'static str {
1420            "terminal_only_gate"
1421        }
1422
1423        fn capabilities(&self) -> PluginCapabilities {
1424            PluginCapabilities::tool_gate()
1425        }
1426    }
1427
1428    #[async_trait::async_trait]
1429    impl ToolGate for TerminalOnlyGate {
1430        async fn next_turn_tool_allowlist(
1431            &self,
1432            _ctx: ToolGateContext<'_>,
1433        ) -> Option<std::collections::HashSet<String>> {
1434            Some(["message_result".to_string()].into_iter().collect())
1435        }
1436    }
1437
1438    impl Plugin for TerminalWithStatusGate {
1439        fn name(&self) -> &'static str {
1440            "terminal_with_status_gate"
1441        }
1442
1443        fn capabilities(&self) -> PluginCapabilities {
1444            PluginCapabilities::tool_gate()
1445        }
1446    }
1447
1448    #[async_trait::async_trait]
1449    impl ToolGate for TerminalWithStatusGate {
1450        async fn next_turn_tool_allowlist(
1451            &self,
1452            _ctx: ToolGateContext<'_>,
1453        ) -> Option<std::collections::HashSet<String>> {
1454            Some(
1455                ["message_info".to_string(), "message_result".to_string()]
1456                    .into_iter()
1457                    .collect(),
1458            )
1459        }
1460    }
1461
1462    impl Plugin for StaticAllowGate {
1463        fn name(&self) -> &'static str {
1464            self.name
1465        }
1466
1467        fn capabilities(&self) -> PluginCapabilities {
1468            PluginCapabilities::tool_gate()
1469        }
1470    }
1471
1472    #[async_trait::async_trait]
1473    impl ToolGate for StaticAllowGate {
1474        fn conflict_priority(&self) -> i32 {
1475            self.priority
1476        }
1477
1478        fn tool_gate_class(&self) -> ToolGateClass {
1479            self.class
1480        }
1481
1482        fn suppresses_advisory_gates(&self, _ctx: ToolGateContext<'_>) -> bool {
1483            self.suppresses_advisory
1484        }
1485
1486        async fn next_turn_tool_allowlist(
1487            &self,
1488            _ctx: ToolGateContext<'_>,
1489        ) -> Option<std::collections::HashSet<String>> {
1490            Some(self.tools.iter().map(|name| (*name).to_string()).collect())
1491        }
1492    }
1493
1494    impl CountingFollowUp {
1495        fn new(remaining: usize) -> Self {
1496            Self {
1497                remaining: AtomicUsize::new(remaining),
1498            }
1499        }
1500    }
1501
1502    impl Plugin for CountingFollowUp {
1503        fn name(&self) -> &'static str {
1504            "counting_follow_up"
1505        }
1506
1507        fn capabilities(&self) -> PluginCapabilities {
1508            PluginCapabilities::follow_up()
1509        }
1510    }
1511
1512    #[async_trait::async_trait]
1513    impl FollowUpSource for CountingFollowUp {
1514        async fn next_follow_up_messages(&self) -> Vec<AgentMessage> {
1515            let used = self
1516                .remaining
1517                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
1518                    remaining.checked_sub(1)
1519                })
1520                .unwrap_or(0);
1521            if used == 0 {
1522                return Vec::new();
1523            }
1524            vec![AgentMessage::System {
1525                content: "retry after no-tool stop".into(),
1526                timestamp: None,
1527            }]
1528        }
1529    }
1530
1531    #[async_trait::async_trait]
1532    impl StreamFn for EmptyThenTextStream {
1533        async fn stream(
1534            &self,
1535            _request: StreamRequest,
1536            _signal: CancellationToken,
1537        ) -> BoxStream<'static, StreamEvent> {
1538            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1539            let partial = empty_assistant_message();
1540            if call == 0 {
1541                return Box::pin(stream::iter(vec![
1542                    StreamEvent::Start {
1543                        partial: partial.clone(),
1544                    },
1545                    StreamEvent::Error {
1546                        partial,
1547                        kind: StreamErrorKind::Empty,
1548                        message: "empty provider response".to_string(),
1549                    },
1550                ]));
1551            }
1552            Box::pin(stream::iter(vec![
1553                StreamEvent::Start { partial },
1554                StreamEvent::Done {
1555                    message: text_assistant_message("recovered"),
1556                },
1557            ]))
1558        }
1559    }
1560
1561    #[async_trait::async_trait]
1562    impl StreamFn for RepeatedTextStream {
1563        async fn stream(
1564            &self,
1565            _request: StreamRequest,
1566            _signal: CancellationToken,
1567        ) -> BoxStream<'static, StreamEvent> {
1568            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1569            let partial = empty_assistant_message();
1570            Box::pin(stream::iter(vec![
1571                StreamEvent::Start { partial },
1572                StreamEvent::Done {
1573                    message: text_assistant_message(format!("plain stop {call}")),
1574                },
1575            ]))
1576        }
1577    }
1578
1579    #[async_trait::async_trait]
1580    impl StreamFn for EmptyStopsAroundProgressStream {
1581        async fn stream(
1582            &self,
1583            _request: StreamRequest,
1584            _signal: CancellationToken,
1585        ) -> BoxStream<'static, StreamEvent> {
1586            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1587            let partial = empty_assistant_message();
1588            let message = match call {
1589                0 | 2 | 4 => text_assistant_message(format!("plain stop {call}")),
1590                1 | 3 => tool_call_assistant_message("progress", format!("tc-progress-{call}")),
1591                5 => tool_call_assistant_message("terminator", "tc-terminator"),
1592                other => panic!("unexpected stream call after terminal turn: {other}"),
1593            };
1594            Box::pin(stream::iter(vec![
1595                StreamEvent::Start { partial },
1596                StreamEvent::Done { message },
1597            ]))
1598        }
1599    }
1600
1601    #[async_trait::async_trait]
1602    impl StreamFn for ZeroOutputThenTextStream {
1603        async fn stream(
1604            &self,
1605            request: StreamRequest,
1606            _signal: CancellationToken,
1607        ) -> BoxStream<'static, StreamEvent> {
1608            self.requests.lock().unwrap().push(request);
1609            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1610            let partial = empty_assistant_message();
1611            if call == 0 {
1612                return Box::pin(stream::iter(vec![
1613                    StreamEvent::Start {
1614                        partial: partial.clone(),
1615                    },
1616                    StreamEvent::Error {
1617                        partial,
1618                        kind: StreamErrorKind::ZeroOutputTransport,
1619                        message: "response body decode failed before output".to_string(),
1620                    },
1621                ]));
1622            }
1623            Box::pin(stream::iter(vec![
1624                StreamEvent::Start { partial },
1625                StreamEvent::Done {
1626                    message: text_assistant_message("recovered from transport"),
1627                },
1628            ]))
1629        }
1630    }
1631
1632    /// Overflows on the first call, then (once the history is shrunk)
1633    /// returns text. Records each request so a test can assert what the
1634    /// retried call actually sent.
1635    #[derive(Default)]
1636    struct OverflowThenTextStream {
1637        calls: AtomicUsize,
1638        requests: Mutex<Vec<StreamRequest>>,
1639    }
1640
1641    #[async_trait::async_trait]
1642    impl StreamFn for OverflowThenTextStream {
1643        async fn stream(
1644            &self,
1645            request: StreamRequest,
1646            _signal: CancellationToken,
1647        ) -> BoxStream<'static, StreamEvent> {
1648            self.requests.lock().unwrap().push(request);
1649            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1650            let partial = empty_assistant_message();
1651            if call == 0 {
1652                return Box::pin(stream::iter(vec![
1653                    StreamEvent::Start {
1654                        partial: partial.clone(),
1655                    },
1656                    StreamEvent::Error {
1657                        partial,
1658                        kind: StreamErrorKind::ContextOverflow,
1659                        message: "maximum context length exceeded".to_string(),
1660                    },
1661                ]));
1662            }
1663            Box::pin(stream::iter(vec![
1664                StreamEvent::Start { partial },
1665                StreamEvent::Done {
1666                    message: text_assistant_message("recovered after shrink"),
1667                },
1668            ]))
1669        }
1670    }
1671
1672    /// Recovery that drops every message except the last — enough to
1673    /// prove the loop persists the shrink and retries.
1674    struct KeepLastRecovery {
1675        calls: Arc<AtomicUsize>,
1676    }
1677
1678    #[async_trait::async_trait]
1679    impl crate::plugin::ContextOverflowRecovery for KeepLastRecovery {
1680        async fn recover(
1681            &self,
1682            mut messages: Vec<AgentMessage>,
1683            _cx: &TransformContext<'_>,
1684        ) -> Vec<AgentMessage> {
1685            self.calls.fetch_add(1, Ordering::SeqCst);
1686            if let Some(last) = messages.pop() {
1687                vec![last]
1688            } else {
1689                messages
1690            }
1691        }
1692        fn max_attempts(&self) -> u8 {
1693            2
1694        }
1695        fn name(&self) -> &'static str {
1696            "keep_last_recovery"
1697        }
1698    }
1699
1700    #[tokio::test]
1701    async fn context_overflow_is_recovered_by_shrinking_and_retrying() {
1702        let stream = Arc::new(OverflowThenTextStream::default());
1703        let recovery_calls = Arc::new(AtomicUsize::new(0));
1704        let config = AgentBuilder::new()
1705            .stream(stream.clone())
1706            .model_id("test-model")
1707            .overflow_recovery(KeepLastRecovery {
1708                calls: recovery_calls.clone(),
1709            })
1710            .build()
1711            .expect("config builds");
1712        let mut context = AgentContext::new("system").with_messages(vec![
1713            AgentMessage::User {
1714                content: UserContent::Text("first".to_string()),
1715                timestamp: None,
1716            },
1717            AgentMessage::User {
1718                content: UserContent::Text("keep me".to_string()),
1719                timestamp: None,
1720            },
1721        ]);
1722
1723        let (assistant, _allowlist) =
1724            stream_with_overflow_recovery(&mut context, &config, &CancellationToken::new(), 0)
1725                .await
1726                .expect("overflow recovery should retry");
1727
1728        let AgentMessage::Assistant { content, .. } = assistant else {
1729            panic!("expected assistant response");
1730        };
1731        assert_eq!(content.plain_text(), "recovered after shrink");
1732        assert_eq!(stream.calls.load(Ordering::SeqCst), 2, "one retry");
1733        assert_eq!(recovery_calls.load(Ordering::SeqCst), 1);
1734        // The shrink is persisted into the live transcript…
1735        assert_eq!(context.messages.len(), 1);
1736        // …and the retried request sent only the shrunk history.
1737        let retried = &stream.requests.lock().unwrap()[1];
1738        assert_eq!(retried.messages.len(), 1);
1739    }
1740
1741    #[tokio::test]
1742    async fn context_overflow_without_recovery_propagates() {
1743        let stream = Arc::new(OverflowThenTextStream::default());
1744        let config = AgentBuilder::new()
1745            .stream(stream.clone())
1746            .model_id("test-model")
1747            .build()
1748            .expect("config builds");
1749        let mut context = AgentContext::new("system").with_messages(vec![AgentMessage::User {
1750            content: UserContent::Text("hi".to_string()),
1751            timestamp: None,
1752        }]);
1753
1754        let result =
1755            stream_with_overflow_recovery(&mut context, &config, &CancellationToken::new(), 0)
1756                .await;
1757        assert!(matches!(
1758            result,
1759            Err(LoopError::Stream(StreamError::ContextOverflow(_)))
1760        ));
1761        assert_eq!(stream.calls.load(Ordering::SeqCst), 1, "no retry");
1762    }
1763
1764    #[test]
1765    fn wrapped_up_is_complete() {
1766        assert!(LoopOutcome::Done.is_complete());
1767        assert!(LoopOutcome::WrappedUp.is_complete());
1768        assert!(!LoopOutcome::HitMaxIterations.is_complete());
1769    }
1770
1771    #[tokio::test]
1772    async fn empty_stream_response_is_retried_before_returning() {
1773        let stream = Arc::new(EmptyThenTextStream::default());
1774        let config = AgentBuilder::new()
1775            .stream(stream.clone())
1776            .model_id("test-model")
1777            .build()
1778            .expect("config builds");
1779        let context = AgentContext::new("system").with_messages(vec![AgentMessage::User {
1780            content: UserContent::Text("continue".to_string()),
1781            timestamp: None,
1782        }]);
1783
1784        let (assistant, _allowlist) =
1785            stream_with_max_tokens_recovery(&context, &config, &CancellationToken::new(), 0)
1786                .await
1787                .expect("second stream attempt should recover");
1788
1789        let AgentMessage::Assistant { content, .. } = assistant else {
1790            panic!("expected assistant response");
1791        };
1792        assert_eq!(content.plain_text(), "recovered");
1793        assert_eq!(stream.calls.load(Ordering::SeqCst), 2);
1794    }
1795
1796    #[tokio::test]
1797    async fn zero_output_transport_error_is_retried_before_returning() {
1798        let stream = Arc::new(ZeroOutputThenTextStream::default());
1799        let config = AgentBuilder::new()
1800            .stream(stream.clone())
1801            .model_id("test-model")
1802            .reasoning(ReasoningEffort::High)
1803            .build()
1804            .expect("config builds");
1805        let context = AgentContext::new("system").with_messages(vec![AgentMessage::User {
1806            content: UserContent::Text("continue".to_string()),
1807            timestamp: None,
1808        }]);
1809
1810        let (assistant, _allowlist) =
1811            stream_with_max_tokens_recovery(&context, &config, &CancellationToken::new(), 0)
1812                .await
1813                .expect("second zero-output transport attempt should recover");
1814
1815        let AgentMessage::Assistant { content, .. } = assistant else {
1816            panic!("expected assistant response");
1817        };
1818        assert_eq!(content.plain_text(), "recovered from transport");
1819        assert_eq!(stream.calls.load(Ordering::SeqCst), 2);
1820
1821        let requests = stream.requests();
1822        assert_eq!(requests.len(), 2);
1823        assert_eq!(requests[0].reasoning, ReasoningEffort::High);
1824        assert_eq!(
1825            requests[1].reasoning,
1826            ReasoningEffort::Minimal,
1827            "zero-output replay should lower high reasoning so reasoning-heavy private-only spins can produce a tool call"
1828        );
1829        assert!(
1830            requests[1].messages.iter().any(|message| matches!(
1831                message,
1832                AgentMessage::System { content, .. }
1833                    if content.contains("transport recovery")
1834                        && content.contains("no visible assistant text")
1835                        && content.contains("no usable tool call")
1836                        && content.contains("unusable burst of partial tool calls")
1837                        && content.contains("exactly one next structured tool call")
1838                        && content.contains("next structured tool call")
1839            )),
1840            "zero-output replay must carry explicit recovery context"
1841        );
1842    }
1843
1844    /// `StreamFn` that emits one assistant turn with a single
1845    /// `terminator` tool call, then panics on subsequent invocations
1846    /// — the test asserts the loop never re-enters the LLM.
1847    struct TerminatorOnlyStream {
1848        calls: AtomicUsize,
1849    }
1850
1851    impl Default for TerminatorOnlyStream {
1852        fn default() -> Self {
1853            Self {
1854                calls: AtomicUsize::new(0),
1855            }
1856        }
1857    }
1858
1859    #[async_trait::async_trait]
1860    impl StreamFn for TerminatorOnlyStream {
1861        async fn stream(
1862            &self,
1863            _request: StreamRequest,
1864            _signal: CancellationToken,
1865        ) -> BoxStream<'static, StreamEvent> {
1866            let call = self.calls.fetch_add(1, Ordering::SeqCst);
1867            assert_eq!(
1868                call, 0,
1869                "terminate-on-turn-1 test must NOT re-enter the LLM after a successful terminator"
1870            );
1871            let partial = empty_assistant_message();
1872            let assistant = AgentMessage::Assistant {
1873                content: AssistantContent {
1874                    blocks: vec![AssistantBlock::ToolCall(crate::tool::ToolCall {
1875                        id: "tc-terminator-1".into(),
1876                        name: "terminator".into(),
1877                        arguments: serde_json::json!({}),
1878                    })],
1879                },
1880                stop_reason: StopReason::ToolUse,
1881                error_message: None,
1882                timestamp: None,
1883                usage: None,
1884            };
1885            Box::pin(stream::iter(vec![
1886                StreamEvent::Start { partial },
1887                StreamEvent::Done { message: assistant },
1888            ]))
1889        }
1890    }
1891
1892    /// Tool that always votes `terminate=true`. Mirrors the contract a
1893    /// downstream terminal/delivery tool upholds.
1894    struct TerminatorTool;
1895
1896    #[async_trait::async_trait]
1897    impl crate::tool::AgentTool for TerminatorTool {
1898        fn name(&self) -> &str {
1899            "terminator"
1900        }
1901
1902        fn description(&self) -> &str {
1903            "test terminator"
1904        }
1905
1906        fn parameters_schema(&self) -> serde_json::Value {
1907            serde_json::json!({"type": "object"})
1908        }
1909
1910        async fn execute(
1911            &self,
1912            _call_id: &str,
1913            _args: serde_json::Value,
1914            _signal: CancellationToken,
1915            _update: tokio::sync::mpsc::UnboundedSender<crate::tool::ToolResult>,
1916        ) -> Result<crate::tool::ToolResult, crate::error::ToolError> {
1917            Ok(crate::tool::ToolResult {
1918                content: vec![crate::types::ToolResultBlock::Text(
1919                    crate::types::TextContent {
1920                        text: "delivered".into(),
1921                    },
1922                )],
1923                is_error: false,
1924                details: serde_json::Value::Null,
1925                terminate: true,
1926                narration: None,
1927            })
1928        }
1929    }
1930
1931    struct ProgressTool;
1932
1933    #[async_trait::async_trait]
1934    impl crate::tool::AgentTool for ProgressTool {
1935        fn name(&self) -> &str {
1936            "progress"
1937        }
1938
1939        fn description(&self) -> &str {
1940            "test progress tool"
1941        }
1942
1943        fn parameters_schema(&self) -> serde_json::Value {
1944            serde_json::json!({"type": "object"})
1945        }
1946
1947        async fn execute(
1948            &self,
1949            _call_id: &str,
1950            _args: serde_json::Value,
1951            _signal: CancellationToken,
1952            _update: tokio::sync::mpsc::UnboundedSender<crate::tool::ToolResult>,
1953        ) -> Result<crate::tool::ToolResult, crate::error::ToolError> {
1954            Ok(crate::tool::ToolResult::text("made progress"))
1955        }
1956    }
1957
1958    /// `SteeringSource` that always returns one wrap-up message. Used
1959    /// to prove the loop does NOT poll steering after a terminator
1960    /// vote (otherwise this would re-enter the LLM and trip the
1961    /// `assert_eq!(call, 0)` in `TerminatorOnlyStream`).
1962    struct AlwaysSteer {
1963        polls: Arc<AtomicUsize>,
1964    }
1965
1966    impl Plugin for AlwaysSteer {
1967        fn name(&self) -> &'static str {
1968            "always_steer"
1969        }
1970
1971        fn capabilities(&self) -> PluginCapabilities {
1972            PluginCapabilities {
1973                steering: true,
1974                ..PluginCapabilities::default()
1975            }
1976        }
1977    }
1978
1979    #[async_trait::async_trait]
1980    impl crate::plugin::SteeringSource for AlwaysSteer {
1981        async fn next_steering_messages(&self) -> Vec<AgentMessage> {
1982            self.polls.fetch_add(1, Ordering::SeqCst);
1983            vec![AgentMessage::System {
1984                content: "wrap up now".into(),
1985                timestamp: None,
1986            }]
1987        }
1988    }
1989
1990    #[tokio::test]
1991    async fn terminator_vote_skips_post_batch_steering_collection() {
1992        // Regression: a `SteeringSource` whose firing condition lines
1993        // up with the same turn the model delivers (e.g.
1994        // `graceful_turn_limit` reaching its soft limit on the delivery
1995        // turn) used to re-enter the loop and prompt the model for
1996        // ANOTHER turn after a clean terminator. The model's drift on
1997        // that extra turn corrupted the user-visible answer in
1998        // production. With the fix, a unanimous terminator vote is a
1999        // hard exit — steering sources are not polled once the run has
2000        // decided it's done.
2001        let stream = Arc::new(TerminatorOnlyStream::default());
2002        let polls = Arc::new(AtomicUsize::new(0));
2003        let mut tool_registry = crate::tool::ToolRegistry::new();
2004        tool_registry = tool_registry.with(Arc::new(TerminatorTool));
2005        let config = AgentBuilder::new()
2006            .stream(stream.clone())
2007            .model_id("test-model")
2008            .tools(tool_registry)
2009            .steering(AlwaysSteer {
2010                polls: polls.clone(),
2011            })
2012            .build()
2013            .expect("config builds");
2014        let context = AgentContext::new("system");
2015        let prompts = vec![AgentMessage::User {
2016            content: UserContent::Text("deliver".to_string()),
2017            timestamp: None,
2018        }];
2019
2020        let result = run(prompts, context, &config, CancellationToken::new())
2021            .await
2022            .expect("run completes after one terminator turn");
2023
2024        // Exactly one LLM call — the terminator turn.
2025        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2026        // Outcome is a clean Done, not WrappedUp (no graceful flag) and
2027        // not HitMaxIterations.
2028        assert_eq!(result.outcome, LoopOutcome::Done);
2029        // Steering source is consulted exactly once — the pre-loop
2030        // priming poll at the top of `inner_run`. After the terminator
2031        // batch, `collect_steering` MUST NOT fire again.
2032        assert_eq!(
2033            polls.load(Ordering::SeqCst),
2034            1,
2035            "steering source polled more than once — terminator vote did not gate post-batch re-entry"
2036        );
2037    }
2038
2039    /// `FollowUpSource` that always emits one nudge. Counts polls so
2040    /// the test can prove `collect_follow_up` is NOT invoked after a
2041    /// terminator batch.
2042    struct AlwaysFollowUp {
2043        polls: Arc<AtomicUsize>,
2044    }
2045
2046    impl Plugin for AlwaysFollowUp {
2047        fn name(&self) -> &'static str {
2048            "always_follow_up"
2049        }
2050
2051        fn capabilities(&self) -> PluginCapabilities {
2052            PluginCapabilities::follow_up()
2053        }
2054    }
2055
2056    #[async_trait::async_trait]
2057    impl FollowUpSource for AlwaysFollowUp {
2058        async fn next_follow_up_messages(&self) -> Vec<AgentMessage> {
2059            self.polls.fetch_add(1, Ordering::SeqCst);
2060            vec![AgentMessage::System {
2061                content: "deliver something".into(),
2062                timestamp: None,
2063            }]
2064        }
2065    }
2066
2067    #[tokio::test]
2068    async fn terminator_vote_skips_post_batch_follow_up_collection() {
2069        // Mirror of the steering test for the follow-up source path.
2070        // `FollowUpSource` exists to nudge the model toward a
2071        // terminator when it failed to emit one — not to overrule a
2072        // terminator the model already cast. After a clean delivery,
2073        // follow-up must be silent.
2074        let stream = Arc::new(TerminatorOnlyStream::default());
2075        let polls = Arc::new(AtomicUsize::new(0));
2076        let mut tool_registry = crate::tool::ToolRegistry::new();
2077        tool_registry = tool_registry.with(Arc::new(TerminatorTool));
2078        let config = AgentBuilder::new()
2079            .stream(stream.clone())
2080            .model_id("test-model")
2081            .tools(tool_registry)
2082            .follow_up(AlwaysFollowUp {
2083                polls: polls.clone(),
2084            })
2085            .build()
2086            .expect("config builds");
2087        let context = AgentContext::new("system");
2088        let prompts = vec![AgentMessage::User {
2089            content: UserContent::Text("deliver".to_string()),
2090            timestamp: None,
2091        }];
2092
2093        let result = run(prompts, context, &config, CancellationToken::new())
2094            .await
2095            .expect("run completes after one terminator turn");
2096
2097        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2098        assert_eq!(result.outcome, LoopOutcome::Done);
2099        assert_eq!(
2100            polls.load(Ordering::SeqCst),
2101            0,
2102            "follow-up source polled after a terminator vote — terminator did not gate post-batch re-entry"
2103        );
2104    }
2105
2106    #[tokio::test]
2107    async fn exhausted_empty_outcome_budget_returns_typed_loop_error() {
2108        let stream = Arc::new(RepeatedTextStream::default());
2109        let config = AgentBuilder::new()
2110            .stream(stream.clone())
2111            .model_id("test-model")
2112            .empty_outcome_retry_budget(1)
2113            .follow_up(CountingFollowUp::new(1))
2114            .build()
2115            .expect("config builds");
2116        let context = AgentContext::new("system");
2117        let prompts = vec![AgentMessage::User {
2118            content: UserContent::Text("continue".to_string()),
2119            timestamp: None,
2120        }];
2121
2122        let err = run(prompts, context, &config, CancellationToken::new())
2123            .await
2124            .expect_err("second no-tool stop should exhaust the budget");
2125
2126        assert!(
2127            matches!(
2128                err,
2129                LoopError::EmptyOutcomeBudgetExhausted {
2130                    budget: 1,
2131                    observed: 2,
2132                }
2133            ),
2134            "unexpected error: {err:?}"
2135        );
2136        assert_eq!(stream.calls.load(Ordering::SeqCst), 2);
2137    }
2138
2139    #[tokio::test]
2140    async fn empty_tool_gate_intersection_prefers_delivery_repair_owner() {
2141        let (sink, mut rx) = crate::event::ChannelSink::new();
2142        let config = AgentBuilder::new()
2143            .stream(Arc::new(RepeatedTextStream::default()))
2144            .event_sink(Arc::new(sink))
2145            .tool_gate_arc(Arc::new(StaticAllowGate {
2146                name: "delivery_repair_gate",
2147                tools: &["browser_interact"],
2148                priority: 100,
2149                class: ToolGateClass::Required,
2150                suppresses_advisory: false,
2151            }))
2152            .tool_gate_arc(Arc::new(StaticAllowGate {
2153                name: "terminal_message_guard",
2154                tools: &["message_result"],
2155                priority: 10,
2156                class: ToolGateClass::Required,
2157                suppresses_advisory: false,
2158            }))
2159            .build()
2160            .expect("config builds");
2161
2162        let allow = collect_tool_allowlist_with_events(&config, 3, &[])
2163            .await
2164            .expect("conflict repair should keep a non-empty allowlist");
2165
2166        assert_eq!(
2167            allow,
2168            ["browser_interact".to_string()].into_iter().collect()
2169        );
2170
2171        let mut saw_conflict = false;
2172        while let Ok(event) = rx.try_recv() {
2173            if let AgentEvent::ToolGateConflictResolved {
2174                chosen_plugin,
2175                allow,
2176                ..
2177            } = event
2178            {
2179                saw_conflict = true;
2180                assert_eq!(chosen_plugin.as_deref(), Some("delivery_repair_gate"));
2181                assert_eq!(allow, vec!["browser_interact".to_string()]);
2182            }
2183        }
2184        assert!(saw_conflict, "tool-gate deadlock should be diagnosable");
2185    }
2186
2187    #[tokio::test]
2188    async fn repair_owner_suppresses_advisory_gate_before_plan_only_intersection() {
2189        let config = AgentBuilder::new()
2190            .stream(Arc::new(RepeatedTextStream::default()))
2191            .tool_gate_arc(Arc::new(StaticAllowGate {
2192                name: "delivery_repair_gate",
2193                tools: &["plan", "file_write"],
2194                priority: 100,
2195                class: ToolGateClass::Required,
2196                suppresses_advisory: true,
2197            }))
2198            .tool_gate_arc(Arc::new(StaticAllowGate {
2199                name: "wrap_up_gate",
2200                tools: &["plan", "message_result", "message_ask"],
2201                priority: 0,
2202                class: ToolGateClass::Advisory,
2203                suppresses_advisory: false,
2204            }))
2205            .build()
2206            .expect("config builds");
2207
2208        let allow = collect_tool_allowlist_with_events(&config, 3, &[])
2209            .await
2210            .expect("repair owner should keep its own allowlist");
2211
2212        assert_eq!(
2213            allow,
2214            ["plan".to_string(), "file_write".to_string()]
2215                .into_iter()
2216                .collect()
2217        );
2218    }
2219
2220    #[tokio::test]
2221    async fn productive_tool_batch_resets_empty_outcome_budget() {
2222        let stream = Arc::new(EmptyStopsAroundProgressStream::default());
2223        let mut tool_registry = crate::tool::ToolRegistry::new();
2224        tool_registry = tool_registry
2225            .with(Arc::new(ProgressTool))
2226            .with(Arc::new(TerminatorTool));
2227        let config = AgentBuilder::new()
2228            .stream(stream.clone())
2229            .model_id("test-model")
2230            .tools(tool_registry)
2231            .empty_outcome_retry_budget(1)
2232            .follow_up(CountingFollowUp::new(3))
2233            .build()
2234            .expect("config builds");
2235        let context = AgentContext::new("system");
2236        let prompts = vec![AgentMessage::User {
2237            content: UserContent::Text("continue".to_string()),
2238            timestamp: None,
2239        }];
2240
2241        let result = run(prompts, context, &config, CancellationToken::new())
2242            .await
2243            .expect("productive tool batches should reset the empty-outcome budget");
2244
2245        assert_eq!(result.outcome, LoopOutcome::Done);
2246        assert_eq!(stream.calls.load(Ordering::SeqCst), 6);
2247    }
2248
2249    #[tokio::test]
2250    async fn terminal_only_plain_text_fallback_synthesizes_terminal_result() {
2251        let stream = Arc::new(RepeatedTextStream::default());
2252        let mut tool_registry = crate::tool::ToolRegistry::new();
2253        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2254        let config = AgentBuilder::new()
2255            .stream(stream.clone())
2256            .model_id("auto-tool-provider")
2257            .tools(tool_registry)
2258            .tool_gate_arc(Arc::new(TerminalOnlyGate))
2259            .plain_text_terminal_fallback_tool("message_result")
2260            .empty_outcome_retry_budget(0)
2261            .build()
2262            .expect("config builds");
2263        let context = AgentContext::new("system");
2264        let prompts = vec![AgentMessage::User {
2265            content: UserContent::Text("answer directly".to_string()),
2266            timestamp: None,
2267        }];
2268
2269        let result = run(prompts, context, &config, CancellationToken::new())
2270            .await
2271            .expect("plain text should be converted on terminal-only turn");
2272
2273        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2274        assert_eq!(result.outcome, LoopOutcome::Done);
2275        assert!(result.messages.iter().any(|message| matches!(
2276            message,
2277            AgentMessage::ToolResult {
2278                tool_name,
2279                content,
2280                is_error: false,
2281                ..
2282            } if tool_name == "message_result"
2283                && content.plain_text() == "plain stop 0"
2284        )));
2285    }
2286
2287    #[tokio::test]
2288    async fn eager_plain_text_fallback_fires_without_terminal_only_allowlist() {
2289        // Providers in the "auto-when-forced" class can never be
2290        // wire-forced into a tool call, so prose IS their failure mode.
2291        // The eager flag lifts the "allowlist must already be narrowed"
2292        // precondition so the fallback fires on the FIRST plain-text
2293        // stop instead of after a narrowing gate has burned 2-3 nudge
2294        // turns.
2295        //
2296        // No `tool_gate_arc` is installed in this test, so the catalog
2297        // stays at the full registry — exactly the situation where the
2298        // non-eager path would refuse to convert and the run would die
2299        // on the empty-outcome budget.
2300        let stream = Arc::new(RepeatedTextStream::default());
2301        let mut tool_registry = crate::tool::ToolRegistry::new();
2302        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2303        let config = AgentBuilder::new()
2304            .stream(stream.clone())
2305            .model_id("auto-tool-provider-eager")
2306            .tools(tool_registry)
2307            .plain_text_terminal_fallback_tool("message_result")
2308            .plain_text_terminal_fallback_eager(true)
2309            .empty_outcome_retry_budget(0)
2310            .build()
2311            .expect("config builds");
2312        let context = AgentContext::new("system");
2313        let prompts = vec![AgentMessage::User {
2314            content: UserContent::Text("answer directly".to_string()),
2315            timestamp: None,
2316        }];
2317
2318        let result = run(prompts, context, &config, CancellationToken::new())
2319            .await
2320            .expect("eager fallback should convert plain text on first stop");
2321
2322        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2323        assert_eq!(result.outcome, LoopOutcome::Done);
2324        assert!(result.messages.iter().any(|message| matches!(
2325            message,
2326            AgentMessage::ToolResult {
2327                tool_name,
2328                content,
2329                is_error: false,
2330                ..
2331            } if tool_name == "message_result"
2332                && content.plain_text() == "plain stop 0"
2333        )));
2334    }
2335
2336    #[tokio::test]
2337    async fn eager_nudge_mode_injects_protocol_recovery_before_synthesizing() {
2338        // With `plain_text_terminal_fallback_eager_nudge(true)` the eager
2339        // path nudges the model with a protocol-recovery system message
2340        // on each consecutive plain-text stop, up to
2341        // `MAX_PLAIN_TEXT_NUDGE_RETRIES`. After the cap a synthesizer
2342        // fires as a last resort so the run still terminates with the
2343        // model's prose as the delivered text — never silently, never
2344        // forever. Verifies the recovery path is observable in the
2345        // emitted message stream (the model sees the nudges in context)
2346        // and that the synthesizer ultimately delivers the first
2347        // substantive plain-text answer, not later recovery drift.
2348        let stream = Arc::new(RepeatedTextStream::default());
2349        let mut tool_registry = crate::tool::ToolRegistry::new();
2350        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2351        let config = AgentBuilder::new()
2352            .stream(stream.clone())
2353            .model_id("auto-tool-provider-eager-nudge")
2354            .tools(tool_registry)
2355            .plain_text_terminal_fallback_tool("message_result")
2356            .plain_text_terminal_fallback_eager(true)
2357            .plain_text_terminal_fallback_eager_nudge(true)
2358            .build()
2359            .expect("config builds");
2360        let context = AgentContext::new("system");
2361        let prompts = vec![AgentMessage::User {
2362            content: UserContent::Text("answer directly".to_string()),
2363            timestamp: None,
2364        }];
2365
2366        let result = run(prompts, context, &config, CancellationToken::new())
2367            .await
2368            .expect("nudge mode should eventually synthesize after retries");
2369
2370        // MAX_PLAIN_TEXT_NUDGE_RETRIES = 2 → two nudges fire, then on the
2371        // third empty stop the synthesizer takes over. Total LLM calls = 3.
2372        assert_eq!(stream.calls.load(Ordering::SeqCst), 3);
2373        assert_eq!(result.outcome, LoopOutcome::Done);
2374
2375        let nudge_count = result
2376            .messages
2377            .iter()
2378            .filter(|m| matches!(m, AgentMessage::System { content, .. } if content == crate::protocol::DEFAULT_PLAIN_TEXT_RECOVERY_PROMPT))
2379            .count();
2380        assert_eq!(
2381            nudge_count, 2,
2382            "expected two protocol-recovery system messages in the run output, got {nudge_count}",
2383        );
2384
2385        let synthesized_text = result
2386            .messages
2387            .iter()
2388            .find_map(|message| match message {
2389                AgentMessage::ToolResult {
2390                    tool_name,
2391                    content,
2392                    is_error: false,
2393                    ..
2394                } if tool_name == "message_result" => Some(content.plain_text()),
2395                _ => None,
2396            })
2397            .expect("a terminal tool result should be synthesized as last resort");
2398        assert_eq!(
2399            synthesized_text, "plain stop 0",
2400            "synthesizer should deliver the first preserved plain text, not later recovery drift",
2401        );
2402    }
2403
2404    #[test]
2405    fn plain_text_fallback_candidate_skips_obvious_clarifying_questions() {
2406        assert!(!should_preserve_plain_text_terminal_candidate(
2407            "Continue what, exactly? What's your next move?"
2408        ));
2409        assert!(!should_preserve_plain_text_terminal_candidate(
2410            "Would you like me to proceed?"
2411        ));
2412        assert!(should_preserve_plain_text_terminal_candidate(
2413            "# Machine Learning\n\nMachine learning is the branch of artificial intelligence that studies systems which improve from data."
2414        ));
2415    }
2416
2417    #[tokio::test]
2418    async fn non_eager_plain_text_fallback_still_requires_narrowed_allowlist() {
2419        // Default behaviour preserved: when eager is NOT set and the
2420        // turn allowlist is the full catalog, plain text is NOT
2421        // converted — the run dies on the empty-outcome budget,
2422        // matching the pre-eager contract for every other model.
2423        let stream = Arc::new(RepeatedTextStream::default());
2424        let mut tool_registry = crate::tool::ToolRegistry::new();
2425        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2426        let config = AgentBuilder::new()
2427            .stream(stream.clone())
2428            .model_id("non-eager-provider")
2429            .tools(tool_registry)
2430            .plain_text_terminal_fallback_tool("message_result")
2431            // Eager NOT set → defaults to false.
2432            .empty_outcome_retry_budget(0)
2433            .build()
2434            .expect("config builds");
2435        let context = AgentContext::new("system");
2436        let prompts = vec![AgentMessage::User {
2437            content: UserContent::Text("answer directly".to_string()),
2438            timestamp: None,
2439        }];
2440
2441        let err = run(prompts, context, &config, CancellationToken::new())
2442            .await
2443            .expect_err("non-eager fallback must not convert without narrowed allowlist");
2444
2445        assert!(
2446            matches!(err, LoopError::EmptyOutcomeBudgetExhausted { .. }),
2447            "unexpected error: {err:?}"
2448        );
2449    }
2450
2451    #[tokio::test]
2452    async fn terminal_plain_text_fallback_allows_status_delivery_gate() {
2453        let stream = Arc::new(RepeatedTextStream::default());
2454        let mut tool_registry = crate::tool::ToolRegistry::new();
2455        tool_registry = tool_registry.with(Arc::new(TerminalNamedTool("message_result")));
2456        let config = AgentBuilder::new()
2457            .stream(stream.clone())
2458            .model_id("auto-tool-provider")
2459            .tools(tool_registry)
2460            .protocol_policy(Arc::new(TestTerminalPolicy))
2461            .tool_gate_arc(Arc::new(TerminalWithStatusGate))
2462            .plain_text_terminal_fallback_tool("message_result")
2463            .empty_outcome_retry_budget(0)
2464            .build()
2465            .expect("config builds");
2466        let context = AgentContext::new("system");
2467        let prompts = vec![AgentMessage::User {
2468            content: UserContent::Text("answer directly".to_string()),
2469            timestamp: None,
2470        }];
2471
2472        let result = run(prompts, context, &config, CancellationToken::new())
2473            .await
2474            .expect(
2475                "plain text should be converted when only status and terminal tools are allowed",
2476            );
2477
2478        assert_eq!(stream.calls.load(Ordering::SeqCst), 1);
2479        assert_eq!(result.outcome, LoopOutcome::Done);
2480        assert!(result.messages.iter().any(|message| matches!(
2481            message,
2482            AgentMessage::ToolResult {
2483                tool_name,
2484                content,
2485                is_error: false,
2486                ..
2487            } if tool_name == "message_result"
2488                && content.plain_text() == "plain stop 0"
2489        )));
2490    }
2491
2492    struct TerminalNamedTool(&'static str);
2493
2494    #[async_trait::async_trait]
2495    impl crate::tool::AgentTool for TerminalNamedTool {
2496        fn name(&self) -> &str {
2497            self.0
2498        }
2499
2500        fn description(&self) -> &str {
2501            "test terminal tool"
2502        }
2503
2504        fn parameters_schema(&self) -> serde_json::Value {
2505            serde_json::json!({"type": "object"})
2506        }
2507
2508        async fn execute(
2509            &self,
2510            _call_id: &str,
2511            _args: serde_json::Value,
2512            _signal: CancellationToken,
2513            _update: tokio::sync::mpsc::UnboundedSender<crate::tool::ToolResult>,
2514        ) -> Result<crate::tool::ToolResult, crate::error::ToolError> {
2515            Ok(crate::tool::ToolResult {
2516                content: vec![crate::types::ToolResultBlock::Text(
2517                    crate::types::TextContent {
2518                        text: "not used".into(),
2519                    },
2520                )],
2521                is_error: false,
2522                details: serde_json::Value::Null,
2523                terminate: true,
2524                narration: None,
2525            })
2526        }
2527    }
2528}