Skip to main content

harness/
agent_loop.rs

1use async_trait::async_trait;
2use futures::{stream::StreamExt, FutureExt};
3use serde_json::{json, Value};
4use std::sync::Arc;
5use std::time::Duration;
6use tokio::sync::mpsc;
7use tokio_util::sync::CancellationToken;
8
9use crate::compaction::{
10    estimate_messages_tokens, CompactionContext, CompactionStrategy, SummarizeCompactionStrategy,
11};
12use crate::event::{HarnessInternalEvent, HarnessUsage, NativeHarnessError, NativeTurnInput};
13use crate::model::{
14    AssistantThinking, CapabilitySupport, ChatMessage, HostedCapability, HostedTool, ModelChunk,
15    ModelClient, ModelClientError, ModelTurnInput,
16};
17use crate::runner::NativeHarness;
18use crate::tools::{
19    bounded::BoundedToolRuntime, ToolFailure, ToolFailureKind, ToolInvocation, ToolOutcome,
20    ToolRuntime, ToolRuntimeError, ToolSpec,
21};
22
23/// Optional compaction wiring: strategy + the model client used to run
24/// the summarize request + the resolved context-window cap. All three
25/// must travel together; without any of them the loop can't make a
26/// useful compaction decision. `AgentLoopHarness::with_compaction`
27/// installs it once and the per-turn loop checks it between steps.
28#[derive(Clone)]
29pub struct CompactionPolicy {
30    pub strategy: Arc<dyn CompactionStrategy>,
31    /// Model client used by `strategy.compact` to run the summarize
32    /// request. Usually the same provider as the main turn model so the
33    /// Anthropic cache prefix stays hot; tests may swap in a fake.
34    pub model_client: Arc<dyn ModelClient>,
35    pub context_window_tokens: u64,
36}
37
38impl CompactionPolicy {
39    /// Build a compaction policy from a custom strategy.
40    ///
41    /// The strategy receives the full message history plus a
42    /// [`CompactionContext`] containing `model_client`, `context_window_tokens`,
43    /// and the current tool specs. Custom strategies may ignore the model
44    /// client entirely, or use it to produce their own summaries.
45    pub fn new(
46        strategy: Arc<dyn CompactionStrategy>,
47        model_client: Arc<dyn ModelClient>,
48        context_window_tokens: u64,
49    ) -> Self {
50        Self {
51            strategy,
52            model_client,
53            context_window_tokens,
54        }
55    }
56
57    /// Build the default summarizing compaction policy used by the harness.
58    ///
59    /// This preserves the existing compaction behavior: old oversized tool
60    /// outputs are pruned first, and when pruning is insufficient the history is
61    /// folded with [`SummarizeCompactionStrategy::default`].
62    pub fn summarizing(model_client: Arc<dyn ModelClient>, context_window_tokens: u64) -> Self {
63        Self::new(
64            Arc::new(SummarizeCompactionStrategy::default()),
65            model_client,
66            context_window_tokens,
67        )
68    }
69}
70
71/// Default mid-stream idle timeout: how long `consume_step_stream` waits
72/// for the *next* model chunk before declaring the connection stalled.
73/// Generous enough to cover extended-thinking pauses (a model can legitimately
74/// go silent for tens of seconds while reasoning) yet bounded so a silently
75/// wedged upstream (TCP open, no FIN/RST, no bytes) can't park the turn forever.
76const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(90);
77
78/// Default stream-layer reconnect budget — how many times we re-establish the
79/// SSE stream after a stall / mid-stream transport drop *before any output has
80/// reached the user*. Separate from `MAX_RETRIES` (the request-establish
81/// budget); stream re-establishment tends to succeed on retry since the
82/// failure is usually a transient gateway / long-lived-connection hiccup, so
83/// this is set higher (6).
84const DEFAULT_STREAM_MAX_ATTEMPTS: u32 = 6;
85
86/// Selects how web search is exposed for an agent turn.
87#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
88#[serde(rename_all = "snake_case")]
89pub enum WebSearchMode {
90    /// Do not expose either provider-hosted or harness-managed web search.
91    #[default]
92    Off,
93    /// Prefer proven provider-hosted search, otherwise use a managed
94    /// `web_search` function tool when the runtime supplies one.
95    Auto,
96    /// Require provider-hosted search. Custom clients reporting `Unknown` are
97    /// allowed to attempt it; explicit `Unsupported` support fails fast.
98    Native,
99    /// Require a harness-managed `web_search` function tool.
100    Managed,
101}
102
103#[derive(Clone)]
104pub struct AgentLoopHarness<M, R> {
105    model: M,
106    /// Every runtime is wrapped so repair + validation + tracing + the
107    /// safety-net output cap apply uniformly, regardless of which concrete
108    /// runtime the caller passed to [`AgentLoopHarness::new`].
109    tools: BoundedToolRuntime<R>,
110    max_steps: usize,
111    compaction: Option<CompactionPolicy>,
112    tool_choice: crate::model::ToolChoice,
113    web_search: WebSearchMode,
114    parallel_tool_calls: Option<bool>,
115    stream_idle_timeout: Duration,
116    stream_max_attempts: u32,
117    turn_end_grace: Duration,
118}
119
120impl<M, R: ToolRuntime> AgentLoopHarness<M, R> {
121    pub fn new(model: M, tools: R) -> Self {
122        Self {
123            model,
124            tools: BoundedToolRuntime::new(tools),
125            max_steps: 8,
126            compaction: None,
127            tool_choice: crate::model::ToolChoice::Auto,
128            web_search: WebSearchMode::Off,
129            parallel_tool_calls: None,
130            stream_idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT,
131            stream_max_attempts: DEFAULT_STREAM_MAX_ATTEMPTS,
132            turn_end_grace: DEFAULT_TURN_END_GRACE,
133        }
134    }
135
136    /// Cap the number of LLM steps per turn. `0` means unlimited — the
137    /// loop only ends when the model stops calling tools (or on
138    /// cancel/error), so callers passing `0` should keep their own
139    /// liveness backstop (idle/wall-clock) around the turn.
140    pub fn with_max_steps(mut self, max_steps: usize) -> Self {
141        self.max_steps = max_steps;
142        self
143    }
144
145    /// Attach a compaction policy. The loop will call
146    /// `policy.strategy.should_compact` before every step and run
147    /// `policy.strategy.compact` when it fires. Without a policy
148    /// installed the loop never compacts — fine for short / test
149    /// conversations, fatal for long production sessions.
150    pub fn with_compaction(mut self, policy: CompactionPolicy) -> Self {
151        self.compaction = Some(policy);
152        self
153    }
154
155    /// Constrain how the model selects tools this turn.
156    /// Defaults to `Auto`. See `ToolChoice` for variants.
157    pub fn with_tool_choice(mut self, choice: crate::model::ToolChoice) -> Self {
158        self.tool_choice = choice;
159        self
160    }
161
162    /// Configure web-search routing. Search is `Off` by default because it may
163    /// incur network egress and provider charges.
164    pub fn with_web_search(mut self, mode: WebSearchMode) -> Self {
165        self.web_search = mode;
166        self
167    }
168
169    /// OpenAI-only: whether the model may emit multiple `tool_use`
170    /// blocks in one response. `None` ⇒ provider default (true on
171    /// OpenAI). Ignored by Anthropic (multi tool_use is implicit).
172    pub fn with_parallel_tool_calls(mut self, parallel: Option<bool>) -> Self {
173        self.parallel_tool_calls = parallel;
174        self
175    }
176
177    /// Override mid-stream resilience knobs. `idle_timeout` is how long a step
178    /// waits for the next model chunk before declaring a stall;
179    /// `max_attempts` is the stream-layer reconnect budget (total stream
180    /// attempts, so `max_attempts = 1` disables reconnection). Primarily for
181    /// tests, which inject a sub-second timeout so a stall surfaces fast
182    /// instead of after the 90s production default.
183    pub fn with_stream_resilience(mut self, idle_timeout: Duration, max_attempts: u32) -> Self {
184        self.stream_idle_timeout = idle_timeout;
185        self.stream_max_attempts = max_attempts.max(1);
186        self
187    }
188
189    /// Harness-level default for how long the loop waits for in-flight tool
190    /// futures to resolve after cancellation before publishing
191    /// `TurnEnd{interrupt}`. A per-turn `NativeTurnInput.turn_end_grace`
192    /// overrides this. Defaults to [`DEFAULT_TURN_END_GRACE`] (1s, the
193    /// historical hard-coded window) so existing behavior is unchanged.
194    ///
195    /// Hosts driving remote tool runtimes (a sandbox exec whose remote
196    /// process needs a TERM-grace followed by SIGKILL, for example) should
197    /// raise this so the turn-end event only lands once sandbox-side
198    /// writers have actually stopped. Otherwise a next turn may start
199    /// while the previous turn's process is still writing to the shared
200    /// workspace. Whatever is still unresolved when the window expires is
201    /// reported in `TurnEnd.pending_tools` so the host can fence
202    /// accordingly.
203    pub fn with_turn_end_grace(mut self, grace: Duration) -> Self {
204        self.turn_end_grace = grace;
205        self
206    }
207}
208
209#[async_trait]
210impl<M, R> NativeHarness for AgentLoopHarness<M, R>
211where
212    M: ModelClient + Clone + Send + Sync + 'static,
213    R: ToolRuntime + Clone + Send + Sync + 'static,
214{
215    async fn run_turn(
216        &self,
217        input: NativeTurnInput,
218    ) -> Result<mpsc::Receiver<Result<HarnessInternalEvent, NativeHarnessError>>, NativeHarnessError>
219    {
220        let (tx, rx) = mpsc::channel(16);
221        let model = self.model.clone();
222        let tools = self.tools.clone();
223        let max_steps = self.max_steps;
224        let compaction = self.compaction.clone();
225        let tool_choice = self.tool_choice.clone();
226        let web_search = self.web_search;
227        let parallel_tool_calls = self.parallel_tool_calls;
228        let stream_idle_timeout = self.stream_idle_timeout;
229        let stream_max_attempts = self.stream_max_attempts;
230        let turn_end_grace = self.turn_end_grace;
231
232        tokio::spawn(async move {
233            let tx_for_panic = tx.clone();
234            let result = std::panic::AssertUnwindSafe(run_loop(
235                model,
236                tools,
237                RunLoopConfig {
238                    max_steps,
239                    compaction,
240                    tool_choice,
241                    web_search,
242                    parallel_tool_calls,
243                    stream_idle_timeout,
244                    stream_max_attempts,
245                    turn_end_grace,
246                },
247                input,
248                tx,
249            ))
250            .catch_unwind()
251            .await;
252            if let Err(payload) = result {
253                let detail = panic_payload_to_string(payload.as_ref());
254                tracing::error!(
255                    target: "harness::agent_loop",
256                    panic = %detail,
257                    "native agent loop panicked"
258                );
259                let _ = tx_for_panic
260                    .send(Err(NativeHarnessError::Failed(format!(
261                        "agent loop panicked: {detail}"
262                    ))))
263                    .await;
264            }
265        });
266
267        Ok(rx)
268    }
269}
270
271/// Test whether the cancel token (if any) has been signalled.
272fn cancel_fired(token: Option<&CancellationToken>) -> bool {
273    token.is_some_and(|t| t.is_cancelled())
274}
275
276fn is_silent_stop(text: &str, stop_reason: &str) -> bool {
277    text.trim().is_empty() && matches!(stop_reason, "end_turn" | "max_tokens")
278}
279
280struct RunLoopConfig {
281    max_steps: usize,
282    compaction: Option<CompactionPolicy>,
283    tool_choice: crate::model::ToolChoice,
284    web_search: WebSearchMode,
285    parallel_tool_calls: Option<bool>,
286    stream_idle_timeout: Duration,
287    stream_max_attempts: u32,
288    /// Harness-level default for how long to wait for in-flight tool
289    /// futures after cancel before publishing `TurnEnd{interrupt}`.
290    /// A per-turn `NativeTurnInput.turn_end_grace` overrides this.
291    turn_end_grace: Duration,
292}
293
294/// Historical default: after the cancel token fires, give cancellation-aware
295/// runtimes this long to resolve in-flight tool futures (e.g. finish remote
296/// cleanup) before `TurnEnd{interrupt}` is published and the remaining
297/// futures are dropped. Tool futures ignored past this window are reported
298/// via `TurnEnd.pending_tools`. Override per harness
299/// ([`AgentLoopHarness::with_turn_end_grace`]) or per turn
300/// (`NativeTurnInput.turn_end_grace`) for runtimes that need longer.
301pub const DEFAULT_TURN_END_GRACE: Duration = Duration::from_secs(1);
302
303fn resolve_web_search_tools(
304    mode: WebSearchMode,
305    native_support: CapabilitySupport,
306    mut tools: Vec<ToolSpec>,
307) -> Result<(Vec<ToolSpec>, Vec<HostedTool>), String> {
308    let has_managed = tools.iter().any(|tool| tool.name == "web_search");
309    let remove_managed = |tools: &mut Vec<ToolSpec>| {
310        tools.retain(|tool| tool.name != "web_search");
311    };
312
313    match mode {
314        WebSearchMode::Off => {
315            remove_managed(&mut tools);
316            Ok((tools, vec![]))
317        }
318        WebSearchMode::Auto if native_support == CapabilitySupport::Supported => {
319            remove_managed(&mut tools);
320            Ok((tools, vec![HostedTool::WebSearch]))
321        }
322        WebSearchMode::Auto if has_managed => Ok((tools, vec![])),
323        WebSearchMode::Auto => Ok((tools, vec![])),
324        WebSearchMode::Native if native_support == CapabilitySupport::Unsupported => Err(
325            "web search mode is Native, but the model client does not support hosted web search"
326                .into(),
327        ),
328        WebSearchMode::Native => {
329            remove_managed(&mut tools);
330            Ok((tools, vec![HostedTool::WebSearch]))
331        }
332        WebSearchMode::Managed if !has_managed => Err(
333            "web search mode is Managed, but the tool runtime does not provide `web_search`".into(),
334        ),
335        WebSearchMode::Managed => Ok((tools, vec![])),
336    }
337}
338
339async fn run_loop<M, R>(
340    model: M,
341    tools: R,
342    config: RunLoopConfig,
343    input: NativeTurnInput,
344    tx: mpsc::Sender<Result<HarnessInternalEvent, NativeHarnessError>>,
345) where
346    M: ModelClient + Send + Sync,
347    R: ToolRuntime + Clone + Send + Sync + 'static,
348{
349    let system_prompt = input.system_prompt.clone();
350    let cancel_token = input.cancel_token.clone();
351    let context_path = input.context_path.clone();
352    // Snapshot and route web-search tools once per turn. This prevents native
353    // and managed `web_search` from being advertised together and keeps the
354    // provider prompt prefix stable across the inner tool loop.
355    let (tools_snapshot, hosted_tools) = match resolve_web_search_tools(
356        config.web_search,
357        model.hosted_capability(HostedCapability::WebSearch),
358        tools.specs(),
359    ) {
360        Ok(selection) => selection,
361        Err(message) => {
362            let _ = tx
363                .send(Err(NativeHarnessError::ModelBadRequest(message)))
364                .await;
365            return;
366        }
367    };
368    // Seed history: load from context JSONL when a path is provided
369    // (persistent mode), otherwise use the in-memory prior_messages.
370    let mut messages: Vec<ChatMessage> = if let Some(ref path) = context_path {
371        crate::context::jsonl::load_context(path).await
372    } else {
373        input.prior_messages
374    };
375    messages.push(ChatMessage::User {
376        content: input.prompt_text,
377        attachments: input.attachments,
378    });
379    // Cursor: how many messages have been flushed to the context JSONL.
380    // Set to messages.len() after the initial User flush (Some path),
381    // or 0 when running in-memory (None — ctx_written is never read).
382    let mut ctx_written: usize = match context_path.as_deref() {
383        None => 0,
384        Some(path) => {
385            let start = messages.len() - 1;
386            crate::context::jsonl::append_context(path, &messages[start..]).await;
387            messages.len()
388        }
389    };
390    // Per-turn accumulated token usage. Each model call may report a
391    // fresh `HarnessUsage` (provider reports per-call counts, not deltas);
392    // we sum them so `TurnEnd.usage` reflects what the whole turn cost.
393    let mut total_usage = HarnessUsage::default();
394    let mut saw_any_usage = false;
395
396    // Fired by `cancel_token.cancel()` from RD on InterruptDispatch.
397    // Emits a single TurnEnd{interrupt} and returns. We check at three
398    // load-bearing points: before each step, before tool dispatch, and
399    // (cheapest of all) inside `consume_step_stream`'s select! on every
400    // chunk await.
401    macro_rules! check_cancel {
402        () => {
403            if cancel_fired(cancel_token.as_ref()) {
404                let _ = tx
405                    .send(Ok(HarnessInternalEvent::TurnEnd {
406                        stop_reason: "interrupt".into(),
407                        usage: saw_any_usage.then(|| total_usage.clone()),
408                        pending_tools: vec![],
409                        final_messages: if context_path.is_none() {
410                            messages.clone()
411                        } else {
412                            vec![]
413                        },
414                    }))
415                    .await;
416                return;
417            }
418        };
419    }
420
421    for step in 0.. {
422        // max_steps == 0 ⇒ unlimited: only the model finishing (or
423        // cancel/error) ends the turn. Otherwise break to the trailing
424        // TurnEnd{max_turns} once the cap is hit.
425        if config.max_steps != 0 && step >= config.max_steps {
426            break;
427        }
428        check_cancel!();
429        // Compaction check — purely additive, never fails the turn. If
430        // the strategy errors out (e.g. provider returned empty
431        // summary), we leave `messages` untouched and let the next
432        // step / turn try again. This keeps "context overflow" as the
433        // worst case: HR sees a model error and decides how to react.
434        if let Some(policy) = &config.compaction {
435            if policy
436                .strategy
437                .should_compact(&messages, policy.context_window_tokens)
438            {
439                let original_count = messages.len();
440                let original_tokens = estimate_messages_tokens(&messages);
441                let cctx = CompactionContext {
442                    system_prompt: system_prompt.clone(),
443                    model_client: policy.model_client.clone(),
444                    context_window_tokens: policy.context_window_tokens,
445                    tools: tools_snapshot.clone(),
446                };
447                match policy.strategy.compact(messages.clone(), &cctx).await {
448                    Ok(outcome) => {
449                        let compacted_count = outcome.messages.len();
450                        let compacted_tokens = estimate_messages_tokens(&outcome.messages);
451                        messages = outcome.messages;
452                        // Compaction summarize-call usage attributed two
453                        // places: into the turn-level total (so HR sees
454                        // the full cost) AND into compaction_*_tokens
455                        // sub-buckets (so HR can isolate what compaction
456                        // alone cost).
457                        if let Some(u) = outcome.usage.as_ref() {
458                            saw_any_usage = true;
459                            total_usage.input_tokens += u.input_tokens;
460                            total_usage.output_tokens += u.output_tokens;
461                            total_usage.cache_read_input_tokens += u.cache_read_input_tokens;
462                            total_usage.cache_creation_input_tokens +=
463                                u.cache_creation_input_tokens;
464                            total_usage.compaction_input_tokens += u.input_tokens;
465                            total_usage.compaction_output_tokens += u.output_tokens;
466                        }
467                        // Structured tracing for operators / dashboards.
468                        // Token counts are estimator output (4 chars/token),
469                        // not provider-reported — labelled in field name.
470                        tracing::info!(
471                            target: "harness::compaction",
472                            step,
473                            original_message_count = original_count,
474                            compacted_message_count = compacted_count,
475                            original_estimated_tokens = original_tokens,
476                            compacted_estimated_tokens = compacted_tokens,
477                            context_window_tokens = policy.context_window_tokens,
478                            "compaction applied"
479                        );
480                        // Rewrite the context JSONL with the compacted history.
481                        if let Some(ref path) = context_path {
482                            crate::context::jsonl::rewrite_context(path, &messages).await;
483                            ctx_written = messages.len();
484                        }
485                        if tx
486                            .send(Ok(HarnessInternalEvent::CompactionApplied {
487                                original_message_count: original_count,
488                                compacted_message_count: compacted_count,
489                                original_tokens,
490                                compacted_tokens,
491                            }))
492                            .await
493                            .is_err()
494                        {
495                            return;
496                        }
497                    }
498                    Err(e) => {
499                        tracing::warn!(
500                            target: "harness::compaction",
501                            step,
502                            error = %e,
503                            "compaction skipped; history retained as-is, model call may now fail with context overflow"
504                        );
505                    }
506                }
507            }
508        }
509
510        // ── Model call with retry for transient errors ────────────────────────
511        // Non-retryable errors (bad config, auth, context overflow) surface
512        // immediately. Retryable errors (rate-limit, network, 5xx) back off
513        // exponentially up to MAX_RETRIES before giving up.
514        const MAX_RETRIES: u32 = 3;
515        const BASE_BACKOFF_MS: u64 = 1_000;
516        const MAX_BACKOFF_MS: u64 = 16_000;
517
518        let model_input = ModelTurnInput {
519            system_prompt: system_prompt.clone(),
520            messages: messages.clone(),
521            tools: tools_snapshot.clone(),
522            hosted_tools: hosted_tools.clone(),
523            tool_choice: config.tool_choice.clone(),
524            parallel_tool_calls: config.parallel_tool_calls,
525        };
526
527        // Per-step stream lifecycle with two independent retry budgets:
528        //   * establish — `model.stream()` erroring before any stream exists.
529        //     Retried up to MAX_RETRIES (request-layer transient faults).
530        //   * consume — a stall / drop *mid-stream*. Retried up to
531        //     `stream_max_attempts`, but ONLY while `had_progress == false`:
532        //     once output has reached the user, re-issuing the request would
533        //     duplicate it, so a mid-stream failure becomes terminal.
534        // The two are nested: each reconnect re-runs establishment (with its
535        // own request-layer retry) before consuming again.
536        let mut stream_attempt = 0u32;
537        let outcome = 'stream: loop {
538            let stream = {
539                let mut attempt = 0u32;
540                loop {
541                    match model.stream(model_input.clone()).await {
542                        Ok(s) => break s,
543                        Err(e) => {
544                            if e.retryable() && attempt < MAX_RETRIES {
545                                let delay_ms =
546                                    (BASE_BACKOFF_MS * (1 << attempt)).min(MAX_BACKOFF_MS);
547                                tracing::warn!(
548                                    attempt,
549                                    delay_ms,
550                                    error = %e,
551                                    "model call failed (retryable) — backing off"
552                                );
553                                if !backoff_sleep(delay_ms, cancel_token.as_ref()).await {
554                                    let _ = tx
555                                        .send(Err(NativeHarnessError::ModelOther(
556                                            "interrupted during retry backoff".into(),
557                                        )))
558                                        .await;
559                                    return;
560                                }
561                                attempt += 1;
562                            } else {
563                                // Non-retryable (config error, auth, etc.) or retries exhausted.
564                                // Surface the error immediately so the user can act on it.
565                                tracing::error!(
566                                    attempt,
567                                    error = %e,
568                                    retryable = e.retryable(),
569                                    "model call failed — terminating turn"
570                                );
571                                let _ = tx.send(Err(model_error_to_native(e))).await;
572                                return;
573                            }
574                        }
575                    }
576                }
577            };
578
579            // Consume the per-step stream: forward TextDelta chunks live
580            // (token-level emit) and accumulate the tool-call state so we
581            // can either dispatch a tool or finalise a message at the end.
582            // The idle watchdog inside fires if the stream goes silent.
583            match consume_step_stream(
584                stream,
585                &tx,
586                step,
587                cancel_token.as_ref(),
588                config.stream_idle_timeout,
589            )
590            .await
591            {
592                Ok(StepDrain::Complete(o)) => break 'stream o,
593                Ok(StepDrain::Cancelled) => {
594                    let _ = tx
595                        .send(Ok(HarnessInternalEvent::TurnEnd {
596                            stop_reason: "interrupt".into(),
597                            usage: saw_any_usage.then(|| total_usage.clone()),
598                            pending_tools: vec![],
599                            final_messages: if context_path.is_none() {
600                                messages.clone()
601                            } else {
602                                vec![]
603                            },
604                        }))
605                        .await;
606                    return;
607                }
608                Err(StepFailure::Model { err, had_progress }) => {
609                    // Reconnect only when nothing has reached the user yet, the
610                    // fault is transient, and the stream budget isn't spent.
611                    if !had_progress
612                        && err.retryable()
613                        && stream_attempt + 1 < config.stream_max_attempts
614                    {
615                        let delay_ms =
616                            (BASE_BACKOFF_MS * (1 << stream_attempt)).min(MAX_BACKOFF_MS);
617                        tracing::warn!(
618                            step,
619                            stream_attempt,
620                            delay_ms,
621                            error = %err,
622                            "model stream failed before any output — reconnecting"
623                        );
624                        if !backoff_sleep(delay_ms, cancel_token.as_ref()).await {
625                            let _ = tx
626                                .send(Err(NativeHarnessError::ModelOther(
627                                    "interrupted during stream reconnect backoff".into(),
628                                )))
629                                .await;
630                            return;
631                        }
632                        stream_attempt += 1;
633                        continue 'stream;
634                    }
635                    // Terminal: output already emitted, non-retryable, or budget
636                    // exhausted. Surface so the user / HR can act on it.
637                    tracing::error!(
638                        step,
639                        stream_attempt,
640                        error = %err,
641                        had_progress,
642                        retryable = err.retryable(),
643                        "model stream failed — terminating turn"
644                    );
645                    let _ = tx.send(Err(model_error_to_native(err))).await;
646                    return;
647                }
648                Err(StepFailure::ChannelClosed) => return,
649                Err(StepFailure::Fatal(e)) => {
650                    let _ = tx.send(Err(e)).await;
651                    return;
652                }
653            }
654        };
655
656        if let Some(u) = outcome.usage.as_ref() {
657            saw_any_usage = true;
658            total_usage.input_tokens += u.input_tokens;
659            total_usage.output_tokens += u.output_tokens;
660            total_usage.cache_read_input_tokens += u.cache_read_input_tokens;
661            total_usage.cache_creation_input_tokens += u.cache_creation_input_tokens;
662        }
663
664        match outcome.next {
665            StepNext::Message { text, stop_reason } => {
666                if is_silent_stop(&text, &stop_reason) {
667                    let _ = tx
668                        .send(Err(NativeHarnessError::ModelOther(format!(
669                            "silent_stop: model returned stop_reason={stop_reason} \
670                             with empty text and no tool calls"
671                        ))))
672                        .await;
673                    return;
674                }
675                let assistant_text = (!text.trim().is_empty()).then_some(text);
676                messages.push(ChatMessage::Assistant {
677                    text: assistant_text,
678                    tool_calls: vec![],
679                    thinking: outcome.thinking.clone(),
680                    usage: outcome.usage.clone(),
681                });
682                // Persist the final Assistant message to context JSONL.
683                if let Some(ref path) = context_path {
684                    crate::context::jsonl::append_context(path, &messages[ctx_written..]).await;
685                }
686                // Note: AssistantTextChunk events were already emitted
687                // mid-stream, so there's nothing more to send here.
688                let final_msgs = if context_path.is_none() {
689                    messages.clone()
690                } else {
691                    vec![]
692                };
693                let _ = tx
694                    .send(Ok(HarnessInternalEvent::TurnEnd {
695                        stop_reason,
696                        usage: saw_any_usage.then(|| total_usage.clone()),
697                        pending_tools: vec![],
698                        final_messages: final_msgs,
699                    }))
700                    .await;
701                return;
702            }
703            StepNext::ToolCalls {
704                preface,
705                mut invocations,
706            } => {
707                check_cancel!();
708                // Schema-guided input repair at dispatch time: fix common
709                // shape mistakes from weak models before the tool sees
710                // them. Runs BEFORE the history
711                // push and the ToolCall events so history, wire, and the
712                // actual execution all agree on the (repaired) arguments.
713                // No matching spec (e.g. model hallucinated a tool name) →
714                // leave the input alone; dispatch will fail it as unknown.
715                for inv in &mut invocations {
716                    // Repair lives in the runtime wrapper (single source of
717                    // truth). Running it here — before the history push and
718                    // ToolCall events below — keeps history, wire, and the
719                    // wrapper's (idempotent) dispatch-time repair in agreement.
720                    if let Some(repairs) = tools.repair_invocation(inv) {
721                        inv.raw_emitted_args = None;
722                        tracing::warn!(
723                            target: "harness::tool_repair",
724                            tool = %inv.name,
725                            id = %inv.id,
726                            repairs = ?repairs,
727                            "schema-guided tool input repair applied"
728                        );
729                    }
730                }
731                let preface_text = preface.filter(|s| !s.is_empty());
732                // Record the assistant turn in history BEFORE executing
733                // the tools. Two reasons:
734                //   * the tool_use blocks live in the assistant message
735                //     per the OpenAI / Anthropic protocols;
736                //   * if the tool errors and the loop bails, history
737                //     still reflects "model called X/Y/Z" — useful for
738                //     debugging and possible retry strategies.
739                messages.push(ChatMessage::Assistant {
740                    text: preface_text,
741                    tool_calls: invocations.clone(),
742                    thinking: outcome.thinking.clone(),
743                    usage: outcome.usage.clone(),
744                });
745
746                // Preface AssistantTextChunk was already emitted mid-stream.
747
748                // Emit ToolCall events in declared order so the wire
749                // sees them in a stable sequence (matters for HR's
750                // ordinal assignment in run_dispatch).
751                for inv in &invocations {
752                    if tx
753                        .send(Ok(HarnessInternalEvent::ToolCall {
754                            id: inv.id.clone(),
755                            name: inv.name.clone(),
756                            input: inv.input.clone(),
757                        }))
758                        .await
759                        .is_err()
760                    {
761                        return;
762                    }
763                }
764
765                // Dispatch all invocations concurrently in one aggregate future.
766                // Keeping them out of detached Tokio tasks is important: if the
767                // turn is cancelled, dropping `join` below drops every tool
768                // future, so a runtime that ignores the cancellation token cannot
769                // keep mutating state after TurnEnd.
770                let calls = invocations.iter().cloned().map(|inv| {
771                    let tools = tools.clone();
772                    let cancel_for_task = cancel_token.clone();
773                    let invocation_for_task = inv.clone();
774                    async move {
775                        let call =
776                            tools.invoke_cancellable(invocation_for_task, cancel_for_task.as_ref());
777                        let outcome = match std::panic::AssertUnwindSafe(call).catch_unwind().await
778                        {
779                            Ok(outcome) => outcome,
780                            Err(payload) => Err(ToolRuntimeError::Runtime(format!(
781                                "tool task panicked: {}",
782                                panic_payload_to_string(payload.as_ref())
783                            ))),
784                        };
785                        (inv, outcome)
786                    }
787                });
788                let join = futures::future::join_all(calls);
789                tokio::pin!(join);
790
791                // Jointly decide how the tool batch ended AND which ids (if
792                // any) are still unresolved. `join` is polled at most once:
793                // `join_all` must not be polled again after returning Ready
794                // (second poll would `mem::take` an empty vec).
795                let (pairs_opt, pending_ids) = if let Some(token) = cancel_token.as_ref() {
796                    let grace = input.turn_end_grace.unwrap_or(config.turn_end_grace);
797                    tokio::select! {
798                        biased;
799                        _ = token.cancelled() => {
800                            // Cancel won. Give cancellation-aware runtimes a
801                            // window to perform remote cleanup (for example
802                            // E2B SendSignal, sandbox TERM→KILL tail). On
803                            // expiry, dropping `join` below cancels all
804                            // remaining in-process futures; none are detached.
805                            // Ids still in flight are reported on TurnEnd so
806                            // the host can keep its terminal fence up until
807                            // the sandbox-side writer is really gone.
808                            match tokio::time::timeout(grace, &mut join).await {
809                                // Everything resolved inside the grace
810                                // window; results are discarded because the
811                                // turn is interrupted anyway, but no writer
812                                // can outlive the turn.
813                                Ok(_) => (None, Vec::new()),
814                                Err(_elapsed) => {
815                                    let pending = invocations.iter().map(|inv| inv.id.clone()).collect();
816                                    (None, pending)
817                                }
818                            }
819                        },
820                        results = &mut join => (Some(results), Vec::new()),
821                    }
822                } else {
823                    (Some((&mut join).await), Vec::new())
824                };
825                let pairs = match pairs_opt {
826                    Some(o) => o,
827                    None => {
828                        // Cancel won the select. Cancellation-aware runtimes had
829                        // an opportunity to clean up, and all remaining futures
830                        // are dropped when this scope returns.
831                        if !pending_ids.is_empty() {
832                            tracing::warn!(
833                                target: "harness::agent_loop",
834                                pending_tools = ?pending_ids,
835                                grace_secs = input
836                                    .turn_end_grace
837                                    .unwrap_or(config.turn_end_grace)
838                                    .as_secs_f32(),
839                                "turn cancelled with tool futures still unresolved after grace window"
840                            );
841                        }
842                        let _ = tx
843                            .send(Ok(HarnessInternalEvent::TurnEnd {
844                                stop_reason: "interrupt".into(),
845                                usage: saw_any_usage.then(|| total_usage.clone()),
846                                pending_tools: pending_ids,
847                                final_messages: if context_path.is_none() {
848                                    messages.clone()
849                                } else {
850                                    vec![]
851                                },
852                            }))
853                            .await;
854                        return;
855                    }
856                };
857
858                // Walk invocations + outcomes pairwise to keep ordering
859                // stable. Tool failures are model-observable: the model can
860                // retry with better input or choose a different approach, and
861                // the harness must not disappear behind a missing terminal.
862                for (inv, outcome) in pairs {
863                    let id = inv.id.clone();
864                    let outcome = match outcome {
865                        Ok(o) => {
866                            tracing::info!(
867                                target: "harness::tool",
868                                step,
869                                tool = %inv.name,
870                                id = %id,
871                                success = o.output.is_ok(),
872                                attachments = o.attachments.len(),
873                                "tool invocation completed"
874                            );
875                            o
876                        }
877                        Err(e) => {
878                            let failure = tool_runtime_error_to_failure(
879                                &inv,
880                                e,
881                                tools_snapshot
882                                    .iter()
883                                    .find(|s| s.name == inv.name)
884                                    .map(|s| &s.input_schema),
885                            );
886                            tracing::warn!(
887                                target: "harness::tool",
888                                step,
889                                tool = %inv.name,
890                                id = %id,
891                                failure_kind = ?failure.kind,
892                                failure_message = %failure.message,
893                                "tool invocation failed; returning model-visible ToolResult"
894                            );
895                            ToolOutcome {
896                                output: Err(failure),
897                                attachments: vec![],
898                            }
899                        }
900                    };
901                    let tool_attachments = outcome.attachments;
902                    let output = outcome.output.map_err(|failure| failure.to_string());
903
904                    // Append the tool result to history so the next model
905                    // step sees it. OpenAI's `tool` role expects content as
906                    // a string; we serialize successes verbatim and wrap
907                    // failures into a small JSON object so the model can
908                    // tell the two apart structurally.
909                    let (tool_content, is_error) = match &output {
910                        Ok(value) => (value.to_string(), false),
911                        Err(err) => (json!({ "error": err }).to_string(), true),
912                    };
913                    messages.push(ChatMessage::Tool {
914                        tool_call_id: id.clone(),
915                        content: tool_content,
916                        is_error,
917                        attachments: tool_attachments,
918                    });
919
920                    if tx
921                        .send(Ok(HarnessInternalEvent::ToolResult { id, output }))
922                        .await
923                        .is_err()
924                    {
925                        return;
926                    }
927                }
928                // Flush the Assistant + all Tool messages for this step.
929                if let Some(ref path) = context_path {
930                    crate::context::jsonl::append_context(path, &messages[ctx_written..]).await;
931                    ctx_written = messages.len();
932                }
933                // Continue the loop — next step will see the tool
934                // results in `messages` and decide what to do.
935            }
936        }
937    }
938
939    // max_turns reached — also flush any unflushed messages.
940    if let Some(ref path) = context_path {
941        crate::context::jsonl::append_context(path, &messages[ctx_written..]).await;
942    }
943    let final_msgs = if context_path.is_none() {
944        messages
945    } else {
946        vec![]
947    };
948    let _ = tx
949        .send(Ok(HarnessInternalEvent::TurnEnd {
950            stop_reason: "max_turns".into(),
951            usage: saw_any_usage.then(|| total_usage.clone()),
952            pending_tools: vec![],
953            final_messages: final_msgs,
954        }))
955        .await;
956}
957
958/// 1:1 lift from `ModelClientError` to `NativeHarnessError`. Two enums
959/// because `ModelClient` is provider-facing (the test fixture
960/// `ScriptedModelClient` exists in the same world) and shouldn't have to
961/// know about the harness-runtime variants (`Encode` / `ChannelClosed`
962/// don't apply to it).
963/// Sleep for `delay_ms`, waking early if the cancel token fires. Returns
964/// `true` if the full backoff elapsed, `false` if interrupted by cancel —
965/// callers treat `false` as "abort the turn". Shared by the request-establish
966/// retry and the stream-reconnect retry so both honour InterruptDispatch
967/// mid-backoff.
968async fn backoff_sleep(delay_ms: u64, cancel_token: Option<&CancellationToken>) -> bool {
969    let sleep = tokio::time::sleep(Duration::from_millis(delay_ms));
970    tokio::pin!(sleep);
971    let cancelled = async {
972        if let Some(t) = cancel_token {
973            t.cancelled().await
974        } else {
975            std::future::pending().await
976        }
977    };
978    tokio::select! {
979        _ = &mut sleep => true,
980        _ = cancelled => false,
981    }
982}
983
984fn model_error_to_native(err: ModelClientError) -> NativeHarnessError {
985    match err {
986        ModelClientError::RateLimit(s) => NativeHarnessError::ModelRateLimit(s),
987        ModelClientError::Auth(s) => NativeHarnessError::ModelAuth(s),
988        ModelClientError::ContextOverflow(s) => NativeHarnessError::ModelContextOverflow(s),
989        ModelClientError::BadRequest(s) => NativeHarnessError::ModelBadRequest(s),
990        ModelClientError::ServerError(s) => NativeHarnessError::ModelServerError(s),
991        ModelClientError::Network(s) => NativeHarnessError::ModelNetwork(s),
992        ModelClientError::Other(s) => NativeHarnessError::ModelOther(s),
993    }
994}
995
996fn tool_runtime_error_to_failure(
997    inv: &ToolInvocation,
998    err: ToolRuntimeError,
999    schema: Option<&serde_json::Value>,
1000) -> ToolFailure {
1001    match err {
1002        ToolRuntimeError::Timeout(message) => ToolFailure::new(ToolFailureKind::Timeout, message),
1003        ToolRuntimeError::InvalidInput { tool, message } => {
1004            crate::tools::invalid_input_failure(&tool, message, &inv.input, schema)
1005        }
1006        ToolRuntimeError::UnknownTool(tool) => ToolFailure::new(
1007            ToolFailureKind::InvalidInput,
1008            format!("unknown tool {tool}; choose one of the advertised tools"),
1009        ),
1010        ToolRuntimeError::Runtime(message) => ToolFailure::new(ToolFailureKind::Runtime, message),
1011    }
1012}
1013
1014fn panic_payload_to_string(payload: &(dyn std::any::Any + Send)) -> String {
1015    if let Some(s) = payload.downcast_ref::<&str>() {
1016        (*s).to_string()
1017    } else if let Some(s) = payload.downcast_ref::<String>() {
1018        s.clone()
1019    } else {
1020        "non-string panic payload".into()
1021    }
1022}
1023
1024/// Per-step accumulated state extracted while draining a `ModelChunk`
1025/// stream. `next` carries the "what to do next" decision (final
1026/// message vs tool dispatch); `usage` rides separately because it must
1027/// fold into the turn-level total regardless of the branch above; and
1028/// `thinking` carries the (text + signature) of any extended-thinking
1029/// block produced this step so the next turn's assistant message can
1030/// echo it back verbatim (Anthropic rejects modified thinking blocks).
1031struct StepOutcome {
1032    next: StepNext,
1033    usage: Option<HarnessUsage>,
1034    thinking: Option<AssistantThinking>,
1035}
1036
1037/// Outcome of draining a single step's chunk stream. `Cancelled` is
1038/// distinct from `Complete` so the agent loop can emit a clean
1039/// `TurnEnd { interrupt }` rather than papering over the half-finished
1040/// state as "Message with empty text".
1041enum StepDrain {
1042    Complete(StepOutcome),
1043    Cancelled,
1044}
1045
1046/// Failure modes of draining a single step's chunk stream. Split out so
1047/// `run_loop` can decide between reconnecting (re-establishing the stream)
1048/// and terminating the turn.
1049enum StepFailure {
1050    /// Stream / transport failure (chunk error, premature close, or idle
1051    /// stall). `err` keeps the original `ModelClientError` so the caller can
1052    /// consult `retryable()`; `had_progress` records whether any model output
1053    /// already reached the user this step. A reconnect is only safe when
1054    /// `!had_progress` — re-issuing the request after partial output would
1055    /// duplicate what the user has already seen.
1056    Model {
1057        err: ModelClientError,
1058        had_progress: bool,
1059    },
1060    /// Downstream event channel closed — RD dropped the receiver. Nothing left
1061    /// to send to; never retryable.
1062    ChannelClosed,
1063    /// Non-retryable processing error (e.g. tool-argument JSON decode failure).
1064    /// Surfaced to the user as-is.
1065    Fatal(NativeHarnessError),
1066}
1067
1068enum StepNext {
1069    Message {
1070        text: String,
1071        stop_reason: String,
1072    },
1073    /// Model returned one or more `tool_use` blocks. Multi-element
1074    /// arrays come from providers that ship `parallel_tool_calls`
1075    /// (OpenAI default) or models that emit multiple tool_use
1076    /// blocks in a single Anthropic message. agent_loop dispatches
1077    /// them concurrently via `join_all`.
1078    ToolCalls {
1079        preface: Option<String>,
1080        invocations: Vec<ToolInvocation>,
1081    },
1082}
1083
1084/// Drain one model step's chunk stream. Forwards `TextDelta` chunks to
1085/// the harness output channel live (token-by-token), accumulates the
1086/// tool call (if any), and returns once `ModelChunk::Done` lands. All
1087/// emitted `AssistantTextChunk` events share `msg_id = "msg_native_<step>"`
1088/// so `native_adapter::TextAccumulator` collapses them into a single
1089/// `AdapterEvent::AgentMessage` on the wire.
1090/// Build the `StepFailure` for an idle-watchdog timeout. Classified as
1091/// `ModelClientError::Network` so `retryable()` is true (a stall is a
1092/// transport-level fault, like a dropped connection); whether it actually
1093/// gets retried is gated by `had_progress` in `run_loop`.
1094fn stall_failure(idle_timeout: Duration, had_progress: bool) -> StepFailure {
1095    StepFailure::Model {
1096        err: ModelClientError::Network(format!(
1097            "model stream stalled: no output for {}s (connection open but idle)",
1098            idle_timeout.as_secs()
1099        )),
1100        had_progress,
1101    }
1102}
1103
1104async fn consume_step_stream(
1105    mut stream: futures::stream::BoxStream<'static, Result<ModelChunk, ModelClientError>>,
1106    tx: &mpsc::Sender<Result<HarnessInternalEvent, NativeHarnessError>>,
1107    step: usize,
1108    cancel_token: Option<&CancellationToken>,
1109    idle_timeout: Duration,
1110) -> Result<StepDrain, StepFailure> {
1111    let emit_msg_id = format!("msg_native_{step}");
1112    let emit_thinking_id = format!("thinking_native_{step}");
1113    let mut text_buf = String::new();
1114    let mut text_stream_started = false;
1115    let mut thinking_buf = String::new();
1116    let mut thinking_signature: Option<String> = None;
1117    let mut saw_thinking = false;
1118    let mut tool_states: Vec<ToolBuf> = Vec::new();
1119    let mut stop_reason = "end_turn".to_string();
1120    let mut usage: Option<HarnessUsage> = None;
1121    // Whether any real model output has reached the user this step. Gates
1122    // whether a later stall / drop is safe to retry (see `StepFailure::Model`).
1123    let mut had_progress = false;
1124
1125    loop {
1126        // Mid-stream idle watchdog: a freshly-armed timer each iteration means
1127        // it measures the gap since the *last* chunk, i.e. it resets on every
1128        // chunk we receive. Keepalive / ping frames are dropped by the SSE
1129        // layer before they ever become a `ModelChunk`, so "received a chunk"
1130        // is exactly "the model made progress" — the timer only survives a
1131        // genuine silence, never a heartbeat-only lull.
1132        let idle = tokio::time::sleep(idle_timeout);
1133        tokio::pin!(idle);
1134
1135        // select! arms: cancellation (priority via `biased`), the idle
1136        // watchdog, and the next stream chunk. Without `biased`, tokio's
1137        // randomised polling can starve cancel checks under heavy
1138        // chunk throughput. With it, an InterruptDispatch fires
1139        // exactly one stream poll later — typically <100 µs.
1140        let item = if let Some(token) = cancel_token {
1141            tokio::select! {
1142                biased;
1143                _ = token.cancelled() => {
1144                    return Ok(StepDrain::Cancelled);
1145                }
1146                _ = &mut idle => return Err(stall_failure(idle_timeout, had_progress)),
1147                next = stream.next() => next,
1148            }
1149        } else {
1150            tokio::select! {
1151                _ = &mut idle => return Err(stall_failure(idle_timeout, had_progress)),
1152                next = stream.next() => next,
1153            }
1154        };
1155        let Some(item) = item else { break };
1156        let chunk = match item {
1157            Ok(c) => c,
1158            Err(e) => {
1159                return Err(StepFailure::Model {
1160                    err: e,
1161                    had_progress,
1162                })
1163            }
1164        };
1165        match chunk {
1166            ModelChunk::TextDelta { msg_id: _, delta } => {
1167                if delta.is_empty() {
1168                    continue;
1169                }
1170                text_buf.push_str(&delta);
1171                let mut flush_delta = delta;
1172                if !text_stream_started {
1173                    if text_buf.trim().is_empty() {
1174                        continue;
1175                    }
1176                    text_stream_started = true;
1177                    // First visible chunk: flush any leading whitespace we
1178                    // held back while deciding whether the step is silent.
1179                    flush_delta = text_buf.clone();
1180                }
1181                // A non-empty text delta is model output the user is about to
1182                // see — past this point a stall is no longer safe to retry.
1183                had_progress = true;
1184                // Forward live to harness output. We rewrite msg_id to
1185                // the per-step canonical form so native_adapter groups
1186                // every chunk of this step into one AdapterEvent.
1187                if tx
1188                    .send(Ok(HarnessInternalEvent::AssistantTextChunk {
1189                        msg_id: emit_msg_id.clone(),
1190                        delta: flush_delta,
1191                    }))
1192                    .await
1193                    .is_err()
1194                {
1195                    return Err(StepFailure::ChannelClosed);
1196                }
1197            }
1198            ModelChunk::ThinkingDelta {
1199                thinking_id: _,
1200                delta,
1201                signature,
1202            } => {
1203                // Signature chunks usually arrive without text and vice
1204                // versa; we accept both shapes and latch whichever the
1205                // provider sends. The text part feeds the live
1206                // AssistantThinkingChunk emit; the signature rides on
1207                // the final ChatMessage::Assistant.thinking so the next
1208                // turn can re-send the block verbatim.
1209                if let Some(sig) = signature {
1210                    if !sig.is_empty() {
1211                        thinking_signature = Some(sig);
1212                    }
1213                }
1214                if !delta.is_empty() {
1215                    saw_thinking = true;
1216                    had_progress = true;
1217                    thinking_buf.push_str(&delta);
1218                    if tx
1219                        .send(Ok(HarnessInternalEvent::AssistantThinkingChunk {
1220                            msg_id: emit_thinking_id.clone(),
1221                            delta,
1222                        }))
1223                        .await
1224                        .is_err()
1225                    {
1226                        return Err(StepFailure::ChannelClosed);
1227                    }
1228                }
1229            }
1230            ModelChunk::ToolCallStart { id, name } => {
1231                // A tool call is committed model output. Even though we buffer
1232                // tool args rather than forwarding them live, treat any
1233                // tool-call activity as progress: re-issuing the request after
1234                // the model has started emitting a tool_use risks a divergent
1235                // / duplicated call.
1236                had_progress = true;
1237                tool_states.push(ToolBuf {
1238                    id,
1239                    name,
1240                    args_buf: String::new(),
1241                    early_input: None,
1242                });
1243            }
1244            ModelChunk::ToolCallInputDelta { id, delta } => {
1245                if let Some(s) = tool_states.iter_mut().find(|s| s.id == id) {
1246                    s.args_buf.push_str(&delta);
1247                }
1248            }
1249            ModelChunk::ToolCallEnd { id, input } => {
1250                if let Some(s) = tool_states.iter_mut().find(|s| s.id == id) {
1251                    s.early_input = input;
1252                }
1253            }
1254            ModelChunk::Done {
1255                stop_reason: sr,
1256                usage: u,
1257            } => {
1258                stop_reason = sr;
1259                usage = u;
1260            }
1261        }
1262    }
1263
1264    // Finalised thinking block — `saw_thinking` covers the (rare) case
1265    // where the provider sent only signature + empty text. We only build
1266    // the AssistantThinking if at least one of the two parts landed.
1267    let thinking = if saw_thinking || thinking_signature.is_some() {
1268        Some(AssistantThinking {
1269            text: thinking_buf,
1270            signature: thinking_signature,
1271        })
1272    } else {
1273        None
1274    };
1275
1276    // Tool call takes precedence — see collect_model_response in model.rs
1277    // for the same rule; the model deferred the final answer until the
1278    // tool runs, so we dispatch the tool(s) instead of emitting TurnEnd.
1279    // Multiple tool_use blocks land here when the provider runs
1280    // parallel_tool_calls — we forward all of them to run_loop.
1281    if !tool_states.is_empty() {
1282        let mut invocations = Vec::with_capacity(tool_states.len());
1283        for state in tool_states {
1284            let parsed_input = match state.early_input {
1285                Some(v) => v,
1286                None => {
1287                    let trimmed = state.args_buf.trim();
1288                    if trimmed.is_empty() {
1289                        // Some providers ship the final tool_call with
1290                        // no args delta (e.g. zero-arg tools); treat
1291                        // empty buffer as an empty object.
1292                        Value::Object(serde_json::Map::new())
1293                    } else {
1294                        match serde_json::from_str(trimmed) {
1295                            Ok(v) => v,
1296                            Err(e) => {
1297                                // Weak models truncate / malform streamed
1298                                // arguments; run the repair chain before
1299                                // failing the turn. A rescued (possibly
1300                                // partial) input the tool can reject is
1301                                // strictly better than a dead turn.
1302                                let res = crate::tool_repair::repair_truncated_json(trimmed);
1303                                match serde_json::from_str(&res.repaired) {
1304                                    Ok(v) if res.changed => {
1305                                        tracing::warn!(
1306                                            target: "harness::tool_repair",
1307                                            tool = %state.name,
1308                                            id = %state.id,
1309                                            notes = ?res.notes,
1310                                            "repaired malformed tool arguments"
1311                                        );
1312                                        v
1313                                    }
1314                                    _ => {
1315                                        return Err(StepFailure::Fatal(
1316                                            NativeHarnessError::ModelOther(format!(
1317                                                "decode tool arguments for {id}: {e}",
1318                                                id = state.id
1319                                            )),
1320                                        ))
1321                                    }
1322                                }
1323                            }
1324                        }
1325                    }
1326                }
1327            };
1328            let raw_emitted_args = raw_args_for_input(&state.args_buf, &parsed_input);
1329            invocations.push(ToolInvocation {
1330                id: state.id,
1331                name: state.name,
1332                input: parsed_input,
1333                raw_emitted_args,
1334            });
1335        }
1336        return Ok(StepDrain::Complete(StepOutcome {
1337            next: StepNext::ToolCalls {
1338                preface: (!text_buf.is_empty()).then_some(text_buf),
1339                invocations,
1340            },
1341            usage,
1342            thinking,
1343        }));
1344    }
1345
1346    Ok(StepDrain::Complete(StepOutcome {
1347        next: StepNext::Message {
1348            text: text_buf,
1349            stop_reason,
1350        },
1351        usage,
1352        thinking,
1353    }))
1354}
1355
1356struct ToolBuf {
1357    id: String,
1358    name: String,
1359    args_buf: String,
1360    early_input: Option<Value>,
1361}
1362
1363fn raw_args_for_input(raw: &str, input: &Value) -> Option<String> {
1364    let trimmed = raw.trim();
1365    if trimmed.is_empty() {
1366        return None;
1367    }
1368    match serde_json::from_str::<Value>(trimmed) {
1369        Ok(parsed) if parsed == *input => Some(trimmed.to_string()),
1370        _ => None,
1371    }
1372}
1373
1374#[cfg(test)]
1375mod tests {
1376    use super::*;
1377    use crate::compaction::{CompactionContext, CompactionError, CompactionStrategy};
1378    use crate::model::{ModelChunk, ModelClient, ModelClientError, ModelResponse};
1379    use crate::tools::{ToolInvocation, ToolOutcome};
1380    use crate::{HarnessInternalEvent, MockToolRuntime, ScriptedModelClient};
1381    use async_trait::async_trait;
1382    use futures::stream::{BoxStream, StreamExt};
1383
1384    fn search_spec() -> crate::tools::ToolSpec {
1385        crate::tools::web_search::web_search_spec()
1386    }
1387
1388    fn ordinary_spec() -> crate::tools::ToolSpec {
1389        crate::tools::ToolSpec {
1390            name: "read".into(),
1391            description: "read".into(),
1392            input_schema: serde_json::json!({"type": "object"}),
1393        }
1394    }
1395
1396    #[test]
1397    fn web_search_mode_routes_exactly_one_search_surface() {
1398        let both = || vec![ordinary_spec(), search_spec()];
1399
1400        let (tools, hosted) =
1401            resolve_web_search_tools(WebSearchMode::Off, CapabilitySupport::Supported, both())
1402                .unwrap();
1403        assert!(tools.iter().all(|tool| tool.name != "web_search"));
1404        assert!(hosted.is_empty());
1405
1406        let (tools, hosted) =
1407            resolve_web_search_tools(WebSearchMode::Auto, CapabilitySupport::Supported, both())
1408                .unwrap();
1409        assert!(tools.iter().all(|tool| tool.name != "web_search"));
1410        assert_eq!(hosted, vec![HostedTool::WebSearch]);
1411
1412        for support in [CapabilitySupport::Unsupported, CapabilitySupport::Unknown] {
1413            let (tools, hosted) =
1414                resolve_web_search_tools(WebSearchMode::Auto, support, both()).unwrap();
1415            assert!(tools.iter().any(|tool| tool.name == "web_search"));
1416            assert!(hosted.is_empty());
1417        }
1418
1419        let (tools, hosted) =
1420            resolve_web_search_tools(WebSearchMode::Native, CapabilitySupport::Unknown, both())
1421                .unwrap();
1422        assert!(tools.iter().all(|tool| tool.name != "web_search"));
1423        assert_eq!(hosted, vec![HostedTool::WebSearch]);
1424
1425        assert!(resolve_web_search_tools(
1426            WebSearchMode::Native,
1427            CapabilitySupport::Unsupported,
1428            both(),
1429        )
1430        .is_err());
1431
1432        let (tools, hosted) =
1433            resolve_web_search_tools(WebSearchMode::Managed, CapabilitySupport::Supported, both())
1434                .unwrap();
1435        assert!(tools.iter().any(|tool| tool.name == "web_search"));
1436        assert!(hosted.is_empty());
1437
1438        assert!(resolve_web_search_tools(
1439            WebSearchMode::Managed,
1440            CapabilitySupport::Supported,
1441            vec![ordinary_spec()],
1442        )
1443        .is_err());
1444    }
1445
1446    #[test]
1447    fn auto_without_any_search_surface_degrades_to_no_search() {
1448        let (tools, hosted) = resolve_web_search_tools(
1449            WebSearchMode::Auto,
1450            CapabilitySupport::Unknown,
1451            vec![ordinary_spec()],
1452        )
1453        .unwrap();
1454        assert_eq!(tools.len(), 1);
1455        assert!(hosted.is_empty());
1456    }
1457    use std::sync::atomic::{AtomicUsize, Ordering};
1458    use std::sync::{Arc, Mutex};
1459
1460    /// Test-only model client that returns a scripted sequence of responses.
1461    /// Each `next()` pops the front of the queue. Used to assert how
1462    /// `AgentLoopHarness` folds per-call usage into the turn total.
1463    #[derive(Clone)]
1464    struct QueueModelClient {
1465        queue: Arc<Mutex<Vec<ModelResponse>>>,
1466    }
1467
1468    impl QueueModelClient {
1469        fn new(responses: Vec<ModelResponse>) -> Self {
1470            Self {
1471                queue: Arc::new(Mutex::new(responses)),
1472            }
1473        }
1474    }
1475
1476    #[async_trait]
1477    impl ModelClient for QueueModelClient {
1478        fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
1479            CapabilitySupport::Unsupported
1480        }
1481
1482        async fn stream(
1483            &self,
1484            _input: ModelTurnInput,
1485        ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
1486        {
1487            let mut q = self.queue.lock().unwrap();
1488            if q.is_empty() {
1489                return Err(ModelClientError::Other("queue exhausted".into()));
1490            }
1491            let response = q.remove(0);
1492            let chunks = response_to_chunks(response);
1493            Ok(futures::stream::iter(chunks.into_iter().map(Ok)).boxed())
1494        }
1495    }
1496
1497    /// Render a synthetic `ModelResponse` as the `ModelChunk` sequence the
1498    /// streaming impl would have emitted. Lets QueueModelClient assert
1499    /// agent-loop behaviour without doing real SSE in tests.
1500    fn response_to_chunks(response: ModelResponse) -> Vec<ModelChunk> {
1501        match response {
1502            ModelResponse::Message {
1503                text,
1504                stop_reason,
1505                usage,
1506            } => {
1507                let mut out = Vec::new();
1508                if !text.is_empty() {
1509                    out.push(ModelChunk::TextDelta {
1510                        msg_id: "queue_msg".into(),
1511                        delta: text,
1512                    });
1513                }
1514                out.push(ModelChunk::Done { stop_reason, usage });
1515                out
1516            }
1517            ModelResponse::ToolCall {
1518                preface,
1519                invocation,
1520                usage,
1521            } => {
1522                let mut out = Vec::new();
1523                if let Some(p) = preface {
1524                    if !p.is_empty() {
1525                        out.push(ModelChunk::TextDelta {
1526                            msg_id: "queue_msg".into(),
1527                            delta: p,
1528                        });
1529                    }
1530                }
1531                out.push(ModelChunk::ToolCallStart {
1532                    id: invocation.id.clone(),
1533                    name: invocation.name.clone(),
1534                });
1535                out.push(ModelChunk::ToolCallEnd {
1536                    id: invocation.id.clone(),
1537                    input: Some(invocation.input.clone()),
1538                });
1539                out.push(ModelChunk::Done {
1540                    stop_reason: "end_turn".into(),
1541                    usage,
1542                });
1543                out
1544            }
1545        }
1546    }
1547
1548    fn usage(input: u64, output: u64, cache_read: u64) -> HarnessUsage {
1549        HarnessUsage {
1550            input_tokens: input,
1551            output_tokens: output,
1552            cache_read_input_tokens: cache_read,
1553            cache_creation_input_tokens: 0,
1554            compaction_input_tokens: 0,
1555            compaction_output_tokens: 0,
1556        }
1557    }
1558
1559    #[tokio::test]
1560    async fn agent_loop_accumulates_usage_across_steps() {
1561        // 2 steps: tool call (10/5 tokens) then final message (20/15 tokens).
1562        let model = QueueModelClient::new(vec![
1563            ModelResponse::ToolCall {
1564                preface: None,
1565                invocation: ToolInvocation {
1566                    id: "tc_1".into(),
1567                    name: "bash".into(),
1568                    input: serde_json::json!({"command": "pwd"}),
1569                    raw_emitted_args: None,
1570                },
1571                usage: Some(usage(10, 5, 0)),
1572            },
1573            ModelResponse::Message {
1574                text: "done".into(),
1575                stop_reason: "end_turn".into(),
1576                usage: Some(usage(20, 15, 4)),
1577            },
1578        ]);
1579        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1580        let mut rx = harness
1581            .run_turn(NativeTurnInput {
1582                prompt_text: "pwd".into(),
1583                system_prompt: None,
1584                attachments: vec![],
1585                cancel_token: None,
1586                prior_messages: vec![],
1587                context_path: None,
1588                turn_end_grace: None,
1589            })
1590            .await
1591            .unwrap();
1592
1593        // Drain until TurnEnd and inspect usage.
1594        let mut final_usage = None;
1595        while let Some(item) = rx.recv().await {
1596            if let HarnessInternalEvent::TurnEnd { usage: u, .. } = item.unwrap() {
1597                final_usage = u;
1598                break;
1599            }
1600        }
1601        let u = final_usage.expect("TurnEnd carried usage");
1602        assert_eq!(u.input_tokens, 30);
1603        assert_eq!(u.output_tokens, 20);
1604        assert_eq!(u.cache_read_input_tokens, 4);
1605    }
1606
1607    #[tokio::test]
1608    async fn agent_loop_errors_on_silent_stop() {
1609        let model = QueueModelClient::new(vec![ModelResponse::Message {
1610            text: "   \n".into(),
1611            stop_reason: "end_turn".into(),
1612            usage: None,
1613        }]);
1614        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1615        let mut rx = harness
1616            .run_turn(NativeTurnInput {
1617                prompt_text: "say something".into(),
1618                system_prompt: None,
1619                attachments: vec![],
1620                cancel_token: None,
1621                prior_messages: vec![],
1622                context_path: None,
1623                turn_end_grace: None,
1624            })
1625            .await
1626            .unwrap();
1627
1628        match rx.recv().await.unwrap() {
1629            Err(NativeHarnessError::ModelOther(msg)) => {
1630                assert!(msg.contains("silent_stop"));
1631                assert!(msg.contains("stop_reason=end_turn"));
1632            }
1633            other => panic!("expected silent_stop model error, got {other:?}"),
1634        }
1635    }
1636
1637    #[tokio::test]
1638    async fn agent_loop_errors_on_empty_max_tokens_stop() {
1639        let model = QueueModelClient::new(vec![ModelResponse::Message {
1640            text: "".into(),
1641            stop_reason: "max_tokens".into(),
1642            usage: None,
1643        }]);
1644        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1645        let mut rx = harness
1646            .run_turn(NativeTurnInput {
1647                prompt_text: "think".into(),
1648                system_prompt: None,
1649                attachments: vec![],
1650                cancel_token: None,
1651                prior_messages: vec![],
1652                context_path: None,
1653                turn_end_grace: None,
1654            })
1655            .await
1656            .unwrap();
1657
1658        match rx.recv().await.unwrap() {
1659            Err(NativeHarnessError::ModelOther(msg)) => {
1660                assert!(msg.contains("silent_stop"));
1661                assert!(msg.contains("stop_reason=max_tokens"));
1662            }
1663            other => panic!("expected silent_stop model error, got {other:?}"),
1664        }
1665    }
1666
1667    #[tokio::test]
1668    async fn agent_loop_turn_end_usage_is_none_when_no_step_reported() {
1669        // Provider reports no usage on either step.
1670        let model = QueueModelClient::new(vec![ModelResponse::Message {
1671            text: "ok".into(),
1672            stop_reason: "end_turn".into(),
1673            usage: None,
1674        }]);
1675        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1676        let mut rx = harness
1677            .run_turn(NativeTurnInput {
1678                prompt_text: "noop".into(),
1679                system_prompt: None,
1680                attachments: vec![],
1681                cancel_token: None,
1682                prior_messages: vec![],
1683                context_path: None,
1684                turn_end_grace: None,
1685            })
1686            .await
1687            .unwrap();
1688        let mut saw_usage = None;
1689        while let Some(item) = rx.recv().await {
1690            if let HarnessInternalEvent::TurnEnd { usage, .. } = item.unwrap() {
1691                saw_usage = Some(usage);
1692                break;
1693            }
1694        }
1695        assert_eq!(saw_usage.unwrap(), None);
1696    }
1697
1698    /// Streaming-aware fake client. Emits a pre-computed `ModelChunk`
1699    /// sequence per call — distinct from `QueueModelClient` which uses
1700    /// the `ModelResponse → chunks` translation. Tests that need
1701    /// token-level chunking go through this one.
1702    #[derive(Clone)]
1703    struct StreamingFakeClient {
1704        chunks_per_call: Arc<Mutex<Vec<Vec<ModelChunk>>>>,
1705    }
1706
1707    impl StreamingFakeClient {
1708        fn new(per_call: Vec<Vec<ModelChunk>>) -> Self {
1709            Self {
1710                chunks_per_call: Arc::new(Mutex::new(per_call)),
1711            }
1712        }
1713    }
1714
1715    #[async_trait]
1716    impl ModelClient for StreamingFakeClient {
1717        fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
1718            CapabilitySupport::Unsupported
1719        }
1720
1721        async fn stream(
1722            &self,
1723            _input: ModelTurnInput,
1724        ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
1725        {
1726            let mut bucket = self.chunks_per_call.lock().unwrap();
1727            if bucket.is_empty() {
1728                return Err(ModelClientError::Other("queue exhausted".into()));
1729            }
1730            let chunks = bucket.remove(0);
1731            Ok(futures::stream::iter(chunks.into_iter().map(Ok)).boxed())
1732        }
1733    }
1734
1735    #[tokio::test]
1736    async fn agent_loop_forwards_token_chunks_to_harness_output() {
1737        let model = StreamingFakeClient::new(vec![vec![
1738            ModelChunk::TextDelta {
1739                msg_id: "remote_msg".into(),
1740                delta: "Hel".into(),
1741            },
1742            ModelChunk::TextDelta {
1743                msg_id: "remote_msg".into(),
1744                delta: "lo ".into(),
1745            },
1746            ModelChunk::TextDelta {
1747                msg_id: "remote_msg".into(),
1748                delta: "world".into(),
1749            },
1750            ModelChunk::Done {
1751                stop_reason: "end_turn".into(),
1752                usage: None,
1753            },
1754        ]]);
1755        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1756        let mut rx = harness
1757            .run_turn(NativeTurnInput {
1758                prompt_text: "hi".into(),
1759                system_prompt: None,
1760                attachments: vec![],
1761                cancel_token: None,
1762                prior_messages: vec![],
1763                context_path: None,
1764                turn_end_grace: None,
1765            })
1766            .await
1767            .unwrap();
1768
1769        let mut deltas: Vec<String> = Vec::new();
1770        let mut saw_end = false;
1771        while let Some(item) = rx.recv().await {
1772            match item.unwrap() {
1773                HarnessInternalEvent::AssistantTextChunk { msg_id, delta } => {
1774                    // The harness rewrites msg_id to the step-local form
1775                    // so native_adapter accumulates everything from one
1776                    // step into a single AgentMessage frame.
1777                    assert_eq!(msg_id, "msg_native_0");
1778                    deltas.push(delta);
1779                }
1780                HarnessInternalEvent::TurnEnd { stop_reason, .. } => {
1781                    assert_eq!(stop_reason, "end_turn");
1782                    saw_end = true;
1783                    break;
1784                }
1785                other => panic!("unexpected event: {other:?}"),
1786            }
1787        }
1788        assert_eq!(deltas, vec!["Hel", "lo ", "world"]);
1789        assert!(saw_end);
1790    }
1791
1792    #[tokio::test]
1793    async fn agent_loop_streaming_tool_call_then_summary() {
1794        // Two scripted streams: first dispatches a tool with streamed
1795        // arguments; second returns a final message after the tool
1796        // result. Tests that the agent loop:
1797        //  * accumulates streamed JSON arguments correctly
1798        //  * runs the tool with the parsed value
1799        //  * feeds the tool result back into the next stream's input
1800        let model = StreamingFakeClient::new(vec![
1801            vec![
1802                ModelChunk::TextDelta {
1803                    msg_id: "r1".into(),
1804                    delta: "running ".into(),
1805                },
1806                ModelChunk::ToolCallStart {
1807                    id: "tc_1".into(),
1808                    name: "bash".into(),
1809                },
1810                ModelChunk::ToolCallInputDelta {
1811                    id: "tc_1".into(),
1812                    delta: "{\"command\":".into(),
1813                },
1814                ModelChunk::ToolCallInputDelta {
1815                    id: "tc_1".into(),
1816                    delta: "\"pwd\"}".into(),
1817                },
1818                ModelChunk::ToolCallEnd {
1819                    id: "tc_1".into(),
1820                    input: None,
1821                },
1822                ModelChunk::Done {
1823                    stop_reason: "end_turn".into(),
1824                    usage: None,
1825                },
1826            ],
1827            vec![
1828                ModelChunk::TextDelta {
1829                    msg_id: "r2".into(),
1830                    delta: "done".into(),
1831                },
1832                ModelChunk::Done {
1833                    stop_reason: "end_turn".into(),
1834                    usage: None,
1835                },
1836            ],
1837        ]);
1838        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1839        let mut rx = harness
1840            .run_turn(NativeTurnInput {
1841                prompt_text: "pwd".into(),
1842                system_prompt: None,
1843                attachments: vec![],
1844                cancel_token: None,
1845                prior_messages: vec![],
1846                context_path: None,
1847                turn_end_grace: None,
1848            })
1849            .await
1850            .unwrap();
1851
1852        // Expected event sequence:
1853        //   AssistantTextChunk("running ")
1854        //   ToolCall{ name=bash, input={"command":"pwd"} }
1855        //   ToolResult{ ok }
1856        //   AssistantTextChunk("done")
1857        //   TurnEnd
1858        let ev = rx.recv().await.unwrap().unwrap();
1859        assert!(matches!(
1860            ev,
1861            HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "running "
1862        ));
1863        let ev = rx.recv().await.unwrap().unwrap();
1864        let HarnessInternalEvent::ToolCall { name, input, .. } = ev else {
1865            panic!("expected ToolCall");
1866        };
1867        assert_eq!(name, "bash");
1868        assert_eq!(input["command"], "pwd");
1869        let ev = rx.recv().await.unwrap().unwrap();
1870        assert!(matches!(ev, HarnessInternalEvent::ToolResult { .. }));
1871        let ev = rx.recv().await.unwrap().unwrap();
1872        assert!(matches!(
1873            ev,
1874            HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "done"
1875        ));
1876        let ev = rx.recv().await.unwrap().unwrap();
1877        assert!(matches!(ev, HarnessInternalEvent::TurnEnd { .. }));
1878    }
1879
1880    #[tokio::test]
1881    async fn agent_loop_repairs_truncated_tool_arguments() {
1882        // OpenAI-style streamed arguments cut off mid-object (missing the
1883        // closing brace). Without the repair chain this was a fatal
1884        // ModelOther; with it the args close cleanly and the tool runs.
1885        let model = StreamingFakeClient::new(vec![
1886            vec![
1887                ModelChunk::ToolCallStart {
1888                    id: "tc_trunc".into(),
1889                    name: "bash".into(),
1890                },
1891                ModelChunk::ToolCallInputDelta {
1892                    id: "tc_trunc".into(),
1893                    delta: r#"{"command":"pwd""#.into(), // truncated
1894                },
1895                ModelChunk::ToolCallEnd {
1896                    id: "tc_trunc".into(),
1897                    input: None,
1898                },
1899                ModelChunk::Done {
1900                    stop_reason: "tool_use".into(),
1901                    usage: None,
1902                },
1903            ],
1904            vec![
1905                ModelChunk::TextDelta {
1906                    msg_id: "r2".into(),
1907                    delta: "done".into(),
1908                },
1909                ModelChunk::Done {
1910                    stop_reason: "end_turn".into(),
1911                    usage: None,
1912                },
1913            ],
1914        ]);
1915        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1916        let mut rx = harness
1917            .run_turn(NativeTurnInput {
1918                prompt_text: "pwd".into(),
1919                system_prompt: None,
1920                attachments: vec![],
1921                cancel_token: None,
1922                prior_messages: vec![],
1923                context_path: None,
1924                turn_end_grace: None,
1925            })
1926            .await
1927            .unwrap();
1928
1929        let mut saw_tool_call = false;
1930        let mut saw_turn_end = false;
1931        while let Some(item) = rx.recv().await {
1932            match item.expect("turn must not fail on truncated args") {
1933                HarnessInternalEvent::ToolCall { name, input, .. } => {
1934                    assert_eq!(name, "bash");
1935                    assert_eq!(input["command"], "pwd", "repaired args reach the wire");
1936                    saw_tool_call = true;
1937                }
1938                HarnessInternalEvent::TurnEnd { .. } => {
1939                    saw_turn_end = true;
1940                    break;
1941                }
1942                _ => {}
1943            }
1944        }
1945        assert!(saw_tool_call, "expected ToolCall with repaired input");
1946        assert!(saw_turn_end);
1947    }
1948
1949    /// Records the invocation input the runtime actually received, so a
1950    /// test can assert dispatch saw the schema-repaired arguments.
1951    #[derive(Clone)]
1952    struct ProbeToolRuntime {
1953        seen_input: Arc<Mutex<Option<Value>>>,
1954    }
1955
1956    #[async_trait]
1957    impl ToolRuntime for ProbeToolRuntime {
1958        fn specs(&self) -> Vec<crate::tools::ToolSpec> {
1959            vec![crate::tools::ToolSpec {
1960                name: "probe".into(),
1961                description: "records its input".into(),
1962                input_schema: serde_json::json!({
1963                    "type": "object",
1964                    "properties": {
1965                        "pattern": {"type": "string"},
1966                        "literal": {"type": "boolean"},
1967                        "limit": {"type": "integer"}
1968                    },
1969                    "required": ["pattern"]
1970                }),
1971            }]
1972        }
1973
1974        async fn invoke(
1975            &self,
1976            invocation: ToolInvocation,
1977        ) -> Result<ToolOutcome, ToolRuntimeError> {
1978            *self.seen_input.lock().unwrap() = Some(invocation.input);
1979            Ok(ToolOutcome {
1980                output: Ok(r#"{"ok":true}"#.into()),
1981                attachments: vec![],
1982            })
1983        }
1984    }
1985
1986    #[tokio::test]
1987    async fn agent_loop_applies_schema_repair_before_dispatch() {
1988        // Weak-model shape mistakes — "true" for a boolean, "30" for an
1989        // integer — are coerced against the tool's input_schema before the
1990        // runtime sees them.
1991        let model = StreamingFakeClient::new(vec![
1992            vec![
1993                ModelChunk::ToolCallStart {
1994                    id: "tc_shape".into(),
1995                    name: "probe".into(),
1996                },
1997                ModelChunk::ToolCallEnd {
1998                    id: "tc_shape".into(),
1999                    input: Some(json!({"pattern": "x", "literal": "true", "limit": "30"})),
2000                },
2001                ModelChunk::Done {
2002                    stop_reason: "tool_use".into(),
2003                    usage: None,
2004                },
2005            ],
2006            vec![
2007                ModelChunk::TextDelta {
2008                    msg_id: "r2".into(),
2009                    delta: "done".into(),
2010                },
2011                ModelChunk::Done {
2012                    stop_reason: "end_turn".into(),
2013                    usage: None,
2014                },
2015            ],
2016        ]);
2017        let seen_input = Arc::new(Mutex::new(None));
2018        let tools = ProbeToolRuntime {
2019            seen_input: seen_input.clone(),
2020        };
2021        let harness = AgentLoopHarness::new(model, tools);
2022        let mut rx = harness
2023            .run_turn(NativeTurnInput {
2024                prompt_text: "go".into(),
2025                system_prompt: None,
2026                attachments: vec![],
2027                cancel_token: None,
2028                prior_messages: vec![],
2029                context_path: None,
2030                turn_end_grace: None,
2031            })
2032            .await
2033            .unwrap();
2034
2035        let mut wire_input: Option<Value> = None;
2036        let mut history: Option<Vec<ChatMessage>> = None;
2037        while let Some(item) = rx.recv().await {
2038            match item.unwrap() {
2039                HarnessInternalEvent::ToolCall { input, .. } => wire_input = Some(input),
2040                HarnessInternalEvent::TurnEnd { final_messages, .. } => {
2041                    history = Some(final_messages);
2042                    break;
2043                }
2044                _ => {}
2045            }
2046        }
2047        let repaired = json!({"pattern": "x", "literal": true, "limit": 30});
2048        // Runtime, wire event, and history all agree on the repaired input.
2049        assert_eq!(seen_input.lock().unwrap().clone().unwrap(), repaired);
2050        assert_eq!(wire_input.unwrap(), repaired);
2051        let history = history.unwrap();
2052        let assistant_tool_calls = history
2053            .iter()
2054            .find_map(|m| match m {
2055                ChatMessage::Assistant { tool_calls, .. } if !tool_calls.is_empty() => {
2056                    Some(tool_calls.clone())
2057                }
2058                _ => None,
2059            })
2060            .expect("assistant message with tool_calls in history");
2061        assert_eq!(assistant_tool_calls[0].input, repaired);
2062    }
2063
2064    #[derive(Clone)]
2065    struct TimeoutToolRuntime;
2066
2067    #[async_trait]
2068    impl ToolRuntime for TimeoutToolRuntime {
2069        fn specs(&self) -> Vec<crate::tools::ToolSpec> {
2070            vec![crate::tools::ToolSpec {
2071                name: "slow".into(),
2072                description: "always times out".into(),
2073                input_schema: serde_json::json!({"type": "object"}),
2074            }]
2075        }
2076
2077        async fn invoke(
2078            &self,
2079            _invocation: ToolInvocation,
2080        ) -> Result<ToolOutcome, ToolRuntimeError> {
2081            Err(ToolRuntimeError::Timeout("tool timed out after 1s".into()))
2082        }
2083    }
2084
2085    #[tokio::test]
2086    async fn agent_loop_tool_timeout_is_model_observable_result() {
2087        let model = StreamingFakeClient::new(vec![
2088            vec![
2089                ModelChunk::ToolCallStart {
2090                    id: "tc_timeout".into(),
2091                    name: "slow".into(),
2092                },
2093                ModelChunk::ToolCallEnd {
2094                    id: "tc_timeout".into(),
2095                    input: Some(json!({})),
2096                },
2097                ModelChunk::Done {
2098                    stop_reason: "tool_use".into(),
2099                    usage: None,
2100                },
2101            ],
2102            vec![
2103                ModelChunk::TextDelta {
2104                    msg_id: "r2".into(),
2105                    delta: "recovered".into(),
2106                },
2107                ModelChunk::Done {
2108                    stop_reason: "end_turn".into(),
2109                    usage: None,
2110                },
2111            ],
2112        ]);
2113        let harness = AgentLoopHarness::new(model, TimeoutToolRuntime);
2114        let mut rx = harness
2115            .run_turn(NativeTurnInput {
2116                prompt_text: "run slow".into(),
2117                system_prompt: None,
2118                attachments: vec![],
2119                cancel_token: None,
2120                prior_messages: vec![],
2121                context_path: None,
2122                turn_end_grace: None,
2123            })
2124            .await
2125            .unwrap();
2126
2127        assert!(matches!(
2128            rx.recv().await.unwrap().unwrap(),
2129            HarnessInternalEvent::ToolCall { .. }
2130        ));
2131        match rx.recv().await.unwrap().unwrap() {
2132            HarnessInternalEvent::ToolResult { output, .. } => {
2133                let err = output.unwrap_err();
2134                assert!(err.contains("Timeout"));
2135                assert!(err.contains("tool timed out"));
2136            }
2137            other => panic!("expected timeout ToolResult, got {other:?}"),
2138        }
2139        assert!(matches!(
2140            rx.recv().await.unwrap().unwrap(),
2141            HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "recovered"
2142        ));
2143        assert!(matches!(
2144            rx.recv().await.unwrap().unwrap(),
2145            HarnessInternalEvent::TurnEnd { ref stop_reason, .. } if stop_reason == "end_turn"
2146        ));
2147    }
2148
2149    #[derive(Clone)]
2150    struct RuntimeErrorToolRuntime;
2151
2152    #[async_trait]
2153    impl ToolRuntime for RuntimeErrorToolRuntime {
2154        fn specs(&self) -> Vec<crate::tools::ToolSpec> {
2155            vec![crate::tools::ToolSpec {
2156                name: "flaky".into(),
2157                description: "always returns a runtime failure".into(),
2158                input_schema: serde_json::json!({"type": "object"}),
2159            }]
2160        }
2161
2162        async fn invoke(
2163            &self,
2164            _invocation: ToolInvocation,
2165        ) -> Result<ToolOutcome, ToolRuntimeError> {
2166            Err(ToolRuntimeError::Runtime(
2167                "sandbox exec stream closed".into(),
2168            ))
2169        }
2170    }
2171
2172    #[tokio::test]
2173    async fn agent_loop_tool_runtime_error_is_model_observable_result() {
2174        let model = StreamingFakeClient::new(vec![
2175            vec![
2176                ModelChunk::ToolCallStart {
2177                    id: "tc_runtime".into(),
2178                    name: "flaky".into(),
2179                },
2180                ModelChunk::ToolCallEnd {
2181                    id: "tc_runtime".into(),
2182                    input: Some(json!({})),
2183                },
2184                ModelChunk::Done {
2185                    stop_reason: "tool_use".into(),
2186                    usage: None,
2187                },
2188            ],
2189            vec![
2190                ModelChunk::TextDelta {
2191                    msg_id: "r2".into(),
2192                    delta: "recovered".into(),
2193                },
2194                ModelChunk::Done {
2195                    stop_reason: "end_turn".into(),
2196                    usage: None,
2197                },
2198            ],
2199        ]);
2200        let harness = AgentLoopHarness::new(model, RuntimeErrorToolRuntime);
2201        let mut rx = harness
2202            .run_turn(NativeTurnInput {
2203                prompt_text: "run flaky".into(),
2204                system_prompt: None,
2205                attachments: vec![],
2206                cancel_token: None,
2207                prior_messages: vec![],
2208                context_path: None,
2209                turn_end_grace: None,
2210            })
2211            .await
2212            .unwrap();
2213
2214        assert!(matches!(
2215            rx.recv().await.unwrap().unwrap(),
2216            HarnessInternalEvent::ToolCall { .. }
2217        ));
2218        match rx.recv().await.unwrap().unwrap() {
2219            HarnessInternalEvent::ToolResult { id, output } => {
2220                assert_eq!(id, "tc_runtime");
2221                let err = output.unwrap_err();
2222                assert!(err.contains("Runtime"));
2223                assert!(err.contains("sandbox exec stream closed"));
2224            }
2225            other => panic!("expected runtime-error ToolResult, got {other:?}"),
2226        }
2227        assert!(matches!(
2228            rx.recv().await.unwrap().unwrap(),
2229            HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "recovered"
2230        ));
2231        assert!(matches!(
2232            rx.recv().await.unwrap().unwrap(),
2233            HarnessInternalEvent::TurnEnd { ref stop_reason, .. } if stop_reason == "end_turn"
2234        ));
2235    }
2236
2237    #[tokio::test]
2238    async fn agent_loop_invalid_tool_input_is_model_observable_and_bounded() {
2239        let huge_content = "x".repeat(20_000);
2240        let model = StreamingFakeClient::new(vec![
2241            vec![
2242                ModelChunk::ToolCallStart {
2243                    id: "tc_bad_write".into(),
2244                    name: "write".into(),
2245                },
2246                ModelChunk::ToolCallEnd {
2247                    id: "tc_bad_write".into(),
2248                    input: Some(json!({"content": huge_content})),
2249                },
2250                ModelChunk::Done {
2251                    stop_reason: "tool_use".into(),
2252                    usage: None,
2253                },
2254            ],
2255            vec![
2256                ModelChunk::TextDelta {
2257                    msg_id: "r2".into(),
2258                    delta: "recovered".into(),
2259                },
2260                ModelChunk::Done {
2261                    stop_reason: "end_turn".into(),
2262                    usage: None,
2263                },
2264            ],
2265        ]);
2266        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2267        let mut rx = harness
2268            .run_turn(NativeTurnInput {
2269                prompt_text: "write file".into(),
2270                system_prompt: None,
2271                attachments: vec![],
2272                cancel_token: None,
2273                prior_messages: vec![],
2274                context_path: None,
2275                turn_end_grace: None,
2276            })
2277            .await
2278            .unwrap();
2279
2280        assert!(matches!(
2281            rx.recv().await.unwrap().unwrap(),
2282            HarnessInternalEvent::ToolCall { .. }
2283        ));
2284        match rx.recv().await.unwrap().unwrap() {
2285            HarnessInternalEvent::ToolResult { output, .. } => {
2286                let err = output.unwrap_err();
2287                // Now caught by the wrapper's schema validation BEFORE the
2288                // inner runtime, with a teaching example appended.
2289                assert!(err.contains("The write tool was called with invalid arguments"));
2290                assert!(err.contains("missing required field `path`"), "{err}");
2291                assert!(err.contains("Received fields: content"));
2292                assert!(err.contains("string(20000 chars"));
2293                assert!(
2294                    err.contains("Expected shape"),
2295                    "teaching example missing: {err}"
2296                );
2297                assert!(
2298                    !err.contains(&"x".repeat(2000)),
2299                    "error should not echo full content"
2300                );
2301            }
2302            other => panic!("expected invalid-input ToolResult, got {other:?}"),
2303        }
2304        assert!(matches!(
2305            rx.recv().await.unwrap().unwrap(),
2306            HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "recovered"
2307        ));
2308    }
2309
2310    /// Spy strategy that always fires and records the call count. Lets
2311    /// us verify agent_loop actually consults the compaction policy
2312    /// between steps without depending on a real summarize round trip.
2313    struct CountingCompactionStrategy {
2314        calls: Arc<AtomicUsize>,
2315    }
2316
2317    #[async_trait]
2318    impl CompactionStrategy for CountingCompactionStrategy {
2319        fn should_compact(&self, _messages: &[ChatMessage], _context_window_tokens: u64) -> bool {
2320            true
2321        }
2322
2323        async fn compact(
2324            &self,
2325            _messages: Vec<ChatMessage>,
2326            _ctx: &CompactionContext,
2327        ) -> Result<crate::compaction::CompactionOutcome, CompactionError> {
2328            self.calls.fetch_add(1, Ordering::SeqCst);
2329            // Replace history with a single synthetic user message — the
2330            // test asserts on the call count, not the content shape.
2331            Ok(crate::compaction::CompactionOutcome {
2332                messages: vec![ChatMessage::User {
2333                    content: "<conversation-summary>FOLDED</conversation-summary>".into(),
2334                    attachments: vec![],
2335                }],
2336                usage: None,
2337            })
2338        }
2339    }
2340
2341    /// Spy strategy that reports a fixed `HarnessUsage` from its compact
2342    /// call. Lets us assert that agent_loop forwards compaction usage
2343    /// into the turn-level total + the `compaction_*` sub-buckets.
2344    struct UsageReportingCompactionStrategy {
2345        invoked: Arc<AtomicUsize>,
2346        per_call_usage: HarnessUsage,
2347    }
2348
2349    #[async_trait]
2350    impl CompactionStrategy for UsageReportingCompactionStrategy {
2351        fn should_compact(&self, _: &[ChatMessage], _: u64) -> bool {
2352            // Fire once per step. Since the model fixture below ends the
2353            // turn after one step, this triggers exactly once per
2354            // run_turn.
2355            self.invoked.load(Ordering::SeqCst) == 0
2356        }
2357        async fn compact(
2358            &self,
2359            messages: Vec<ChatMessage>,
2360            _ctx: &CompactionContext,
2361        ) -> Result<crate::compaction::CompactionOutcome, CompactionError> {
2362            self.invoked.fetch_add(1, Ordering::SeqCst);
2363            Ok(crate::compaction::CompactionOutcome {
2364                messages,
2365                usage: Some(self.per_call_usage.clone()),
2366            })
2367        }
2368    }
2369
2370    #[tokio::test]
2371    async fn agent_loop_attributes_compaction_usage_to_subbucket_and_total() {
2372        // Compaction reports 50/20 tokens; main step reports 100/30.
2373        // TurnEnd.usage should sum into 150/50, with compaction_*
2374        // sub-buckets showing the 50/20 isolated.
2375        let model = StreamingFakeClient::new(vec![vec![
2376            ModelChunk::TextDelta {
2377                msg_id: "m".into(),
2378                delta: "done".into(),
2379            },
2380            ModelChunk::Done {
2381                stop_reason: "end_turn".into(),
2382                usage: Some(usage(100, 30, 0)),
2383            },
2384        ]]);
2385        let invoked = Arc::new(AtomicUsize::new(0));
2386        let strategy = UsageReportingCompactionStrategy {
2387            invoked: invoked.clone(),
2388            per_call_usage: usage(50, 20, 0),
2389        };
2390        let policy = CompactionPolicy::new(
2391            Arc::new(strategy),
2392            Arc::new(ScriptedModelClient),
2393            1, // forces should_compact's true branch
2394        );
2395        let harness = AgentLoopHarness::new(model, MockToolRuntime::new()).with_compaction(policy);
2396        let mut rx = harness
2397            .run_turn(NativeTurnInput {
2398                prompt_text: "hi".into(),
2399                system_prompt: None,
2400                attachments: vec![],
2401                cancel_token: None,
2402                prior_messages: vec![],
2403                context_path: None,
2404                turn_end_grace: None,
2405            })
2406            .await
2407            .unwrap();
2408        let mut final_usage = None;
2409        while let Some(item) = rx.recv().await {
2410            if let HarnessInternalEvent::TurnEnd { usage, .. } = item.unwrap() {
2411                final_usage = usage;
2412                break;
2413            }
2414        }
2415        assert_eq!(invoked.load(Ordering::SeqCst), 1);
2416        let u = final_usage.expect("TurnEnd carried usage");
2417        // Main step (100, 30) + compaction (50, 20) = total (150, 50).
2418        assert_eq!(u.input_tokens, 150);
2419        assert_eq!(u.output_tokens, 50);
2420        // Compaction sub-bucket isolates the (50, 20) portion.
2421        assert_eq!(u.compaction_input_tokens, 50);
2422        assert_eq!(u.compaction_output_tokens, 20);
2423    }
2424
2425    /// Stub client that records every `ModelTurnInput.messages` it was
2426    /// asked to stream. Lets the test assert that the compaction-replaced
2427    /// messages are what reaches the model on the next step.
2428    #[derive(Clone)]
2429    struct RecordingFakeClient {
2430        last_messages: Arc<Mutex<Option<Vec<ChatMessage>>>>,
2431        chunks: Vec<ModelChunk>,
2432    }
2433
2434    #[async_trait]
2435    impl ModelClient for RecordingFakeClient {
2436        fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
2437            CapabilitySupport::Unsupported
2438        }
2439
2440        async fn stream(
2441            &self,
2442            input: ModelTurnInput,
2443        ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
2444        {
2445            *self.last_messages.lock().unwrap() = Some(input.messages);
2446            Ok(futures::stream::iter(self.chunks.clone().into_iter().map(Ok)).boxed())
2447        }
2448    }
2449
2450    #[tokio::test]
2451    async fn agent_loop_invokes_compaction_between_steps() {
2452        let calls = Arc::new(AtomicUsize::new(0));
2453        let last_messages = Arc::new(Mutex::new(None::<Vec<ChatMessage>>));
2454        let model = RecordingFakeClient {
2455            last_messages: last_messages.clone(),
2456            chunks: vec![
2457                ModelChunk::TextDelta {
2458                    msg_id: "m".into(),
2459                    delta: "done".into(),
2460                },
2461                ModelChunk::Done {
2462                    stop_reason: "end_turn".into(),
2463                    usage: None,
2464                },
2465            ],
2466        };
2467        // Summary client used by the spy strategy's ctx — not actually
2468        // called because our strategy short-circuits, but we satisfy
2469        // the policy contract.
2470        let summary_client: Arc<dyn ModelClient> = Arc::new(ScriptedModelClient);
2471        let policy = CompactionPolicy::new(
2472            Arc::new(CountingCompactionStrategy {
2473                calls: calls.clone(),
2474            }),
2475            summary_client,
2476            1, // forces should_compact to fire
2477        );
2478
2479        let harness = AgentLoopHarness::new(model, MockToolRuntime::new()).with_compaction(policy);
2480        let mut rx = harness
2481            .run_turn(NativeTurnInput {
2482                prompt_text: "hello".into(),
2483                system_prompt: None,
2484                attachments: vec![],
2485                cancel_token: None,
2486                prior_messages: vec![],
2487                context_path: None,
2488                turn_end_grace: None,
2489            })
2490            .await
2491            .unwrap();
2492        let mut compaction_event: Option<(usize, usize)> = None;
2493        while let Some(item) = rx.recv().await {
2494            match item.unwrap() {
2495                HarnessInternalEvent::CompactionApplied {
2496                    original_message_count,
2497                    compacted_message_count,
2498                    ..
2499                } => {
2500                    compaction_event = Some((original_message_count, compacted_message_count));
2501                }
2502                HarnessInternalEvent::TurnEnd { .. } => break,
2503                _ => {}
2504            }
2505        }
2506        // Compaction ran exactly once before the single model step.
2507        assert_eq!(calls.load(Ordering::SeqCst), 1);
2508        // CompactionApplied event surfaced with sensible counts.
2509        let (orig, comp) = compaction_event.expect("CompactionApplied event emitted");
2510        assert_eq!(orig, 1, "started with 1 message ([User \"hello\"])");
2511        assert_eq!(comp, 1, "spy strategy folded to single User message");
2512        // The model saw the FOLDED messages, not the original
2513        // [User "hello"] prefix.
2514        let observed = last_messages.lock().unwrap().clone().expect("model called");
2515        assert_eq!(observed.len(), 1);
2516        match &observed[0] {
2517            ChatMessage::User { content, .. } => {
2518                assert!(content.contains("FOLDED"), "got {content:?}");
2519            }
2520            other => panic!("expected User, got {other:?}"),
2521        }
2522    }
2523
2524    #[test]
2525    fn compaction_policy_summarizing_uses_default_strategy() {
2526        let policy = CompactionPolicy::summarizing(Arc::new(ScriptedModelClient), 100_000);
2527
2528        assert_eq!(policy.context_window_tokens, 100_000);
2529        assert!(!policy.strategy.should_compact(&[], 100_000));
2530    }
2531
2532    /// Model client whose stream blocks indefinitely until cancelled.
2533    /// Lets the test prove that cancel_token.cancelled() races
2534    /// stream.next() and wins.
2535    #[derive(Clone)]
2536    struct HangingModelClient {
2537        started: Arc<tokio::sync::Notify>,
2538    }
2539
2540    #[async_trait]
2541    impl ModelClient for HangingModelClient {
2542        fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
2543            CapabilitySupport::Unsupported
2544        }
2545
2546        async fn stream(
2547            &self,
2548            _input: ModelTurnInput,
2549        ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
2550        {
2551            // Channel-backed stream whose sender never sends and never
2552            // drops — `rx.next().await` parks forever. Mirrors a real
2553            // LLM that opened the SSE response but hasn't shipped a
2554            // chunk yet (slow first-token time).
2555            let (tx, rx) = mpsc::channel::<Result<ModelChunk, ModelClientError>>(1);
2556            let started = self.started.clone();
2557            tokio::spawn(async move {
2558                // Hold the sender alive for the test's lifetime. Notify
2559                // the test that the stream is "started" so it knows
2560                // when to fire cancel — proves the cancel races a
2561                // pending stream.next(), not the pre-step check.
2562                started.notify_one();
2563                let _retain = tx; // suppress drop warning
2564                let () = std::future::pending().await;
2565            });
2566            Ok(tokio_stream::wrappers::ReceiverStream::new(rx).boxed())
2567        }
2568    }
2569
2570    #[tokio::test]
2571    async fn agent_loop_cancellation_interrupts_in_flight_stream() {
2572        // Without cancel: agent_loop would hang forever waiting for the
2573        // first chunk. With cancel fired *after* the stream began, the
2574        // select! arm in consume_step_stream wins and we get TurnEnd
2575        // with stop_reason "interrupt" within milliseconds.
2576        let started = Arc::new(tokio::sync::Notify::new());
2577        let model = HangingModelClient {
2578            started: started.clone(),
2579        };
2580        let cancel = CancellationToken::new();
2581        let cancel_for_outside = cancel.clone();
2582
2583        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2584        let mut rx = harness
2585            .run_turn(NativeTurnInput {
2586                prompt_text: "hi".into(),
2587                system_prompt: None,
2588                attachments: vec![],
2589                cancel_token: Some(cancel),
2590                prior_messages: vec![],
2591                context_path: None,
2592                turn_end_grace: None,
2593            })
2594            .await
2595            .unwrap();
2596
2597        // Wait for the model to actually start streaming, then cancel.
2598        // (Cancelling before the stream begins would short-circuit at
2599        // the pre-step check_cancel! macro — also correct, but a
2600        // different code path. We want to exercise the select!.)
2601        started.notified().await;
2602        cancel_for_outside.cancel();
2603
2604        // Within a small window we should observe a TurnEnd{interrupt}.
2605        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
2606        let mut saw_interrupt = false;
2607        while tokio::time::Instant::now() < deadline {
2608            tokio::select! {
2609                item = rx.recv() => {
2610                    match item {
2611                        Some(Ok(HarnessInternalEvent::TurnEnd { stop_reason, .. })) => {
2612                            assert_eq!(stop_reason, "interrupt");
2613                            saw_interrupt = true;
2614                            break;
2615                        }
2616                        Some(_) => continue,
2617                        None => break,
2618                    }
2619                }
2620                _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {}
2621            }
2622        }
2623        assert!(saw_interrupt, "expected TurnEnd{{interrupt}} after cancel");
2624    }
2625
2626    /// Per-`stream()`-call scripted client for the mid-stream idle-timeout
2627    /// tests. Each behavior either streams chunks to completion (stream
2628    /// closes), or emits an optional prefix then parks forever without
2629    /// closing — simulating a silently wedged upstream (TCP open, no FIN/RST,
2630    /// no further bytes). `calls` counts establishments so tests can assert
2631    /// whether a reconnect happened.
2632    enum StallBehavior {
2633        /// Stream these chunks, then close (clean end).
2634        Complete(Vec<ModelChunk>),
2635        /// Emit these chunks (possibly none), then hang forever.
2636        EmitThenHang(Vec<ModelChunk>),
2637    }
2638
2639    #[derive(Clone)]
2640    struct StallingModelClient {
2641        behaviors: Arc<Mutex<Vec<StallBehavior>>>,
2642        calls: Arc<AtomicUsize>,
2643    }
2644
2645    impl StallingModelClient {
2646        fn new(behaviors: Vec<StallBehavior>) -> Self {
2647            Self {
2648                behaviors: Arc::new(Mutex::new(behaviors)),
2649                calls: Arc::new(AtomicUsize::new(0)),
2650            }
2651        }
2652    }
2653
2654    #[async_trait]
2655    impl ModelClient for StallingModelClient {
2656        fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
2657            CapabilitySupport::Unsupported
2658        }
2659
2660        async fn stream(
2661            &self,
2662            _input: ModelTurnInput,
2663        ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
2664        {
2665            self.calls.fetch_add(1, Ordering::SeqCst);
2666            // Pop the next scripted behavior; once the script is exhausted,
2667            // default to hanging (covers "every attempt stalls" tests).
2668            let behavior = {
2669                let mut b = self.behaviors.lock().unwrap();
2670                if b.is_empty() {
2671                    StallBehavior::EmitThenHang(vec![])
2672                } else {
2673                    b.remove(0)
2674                }
2675            };
2676            let (tx, rx) = mpsc::channel::<Result<ModelChunk, ModelClientError>>(8);
2677            tokio::spawn(async move {
2678                match behavior {
2679                    StallBehavior::Complete(chunks) => {
2680                        for c in chunks {
2681                            if tx.send(Ok(c)).await.is_err() {
2682                                return;
2683                            }
2684                        }
2685                        // tx dropped here → stream ends cleanly.
2686                    }
2687                    StallBehavior::EmitThenHang(chunks) => {
2688                        for c in chunks {
2689                            if tx.send(Ok(c)).await.is_err() {
2690                                return;
2691                            }
2692                        }
2693                        let _retain = tx; // hold sender open so rx parks
2694                        let () = std::future::pending().await;
2695                    }
2696                }
2697            });
2698            Ok(tokio_stream::wrappers::ReceiverStream::new(rx).boxed())
2699        }
2700    }
2701
2702    #[tokio::test(start_paused = true)]
2703    async fn agent_loop_reconnects_after_stall_before_any_output() {
2704        // First establishment opens the stream then goes silent → idle
2705        // watchdog fires → no output yet, so it's safe to reconnect. Second
2706        // establishment streams a full response. We should see the text
2707        // exactly once and a clean end_turn, with two establishments total.
2708        let model = StallingModelClient::new(vec![
2709            StallBehavior::EmitThenHang(vec![]),
2710            StallBehavior::Complete(vec![
2711                ModelChunk::TextDelta {
2712                    msg_id: "m".into(),
2713                    delta: "ok".into(),
2714                },
2715                ModelChunk::Done {
2716                    stop_reason: "end_turn".into(),
2717                    usage: None,
2718                },
2719            ]),
2720        ]);
2721        let calls = model.calls.clone();
2722        let harness = AgentLoopHarness::new(model, MockToolRuntime::new())
2723            .with_stream_resilience(Duration::from_millis(50), 3);
2724        let mut rx = harness
2725            .run_turn(NativeTurnInput {
2726                prompt_text: "hi".into(),
2727                system_prompt: None,
2728                attachments: vec![],
2729                cancel_token: None,
2730                prior_messages: vec![],
2731                context_path: None,
2732                turn_end_grace: None,
2733            })
2734            .await
2735            .unwrap();
2736
2737        let mut text = String::new();
2738        let mut stop = None;
2739        while let Some(item) = rx.recv().await {
2740            match item.expect("no error expected") {
2741                HarnessInternalEvent::AssistantTextChunk { delta, .. } => text.push_str(&delta),
2742                HarnessInternalEvent::TurnEnd { stop_reason, .. } => {
2743                    stop = Some(stop_reason);
2744                    break;
2745                }
2746                _ => {}
2747            }
2748        }
2749        assert_eq!(stop.as_deref(), Some("end_turn"));
2750        assert_eq!(text, "ok", "text delivered exactly once, no duplication");
2751        assert_eq!(
2752            calls.load(Ordering::SeqCst),
2753            2,
2754            "stream established twice (one reconnect)"
2755        );
2756    }
2757
2758    #[tokio::test(start_paused = true)]
2759    async fn agent_loop_surfaces_error_when_reconnect_budget_exhausted() {
2760        // Every establishment stalls. With max_attempts = 2 we get one
2761        // reconnect, then the second stall is terminal → ModelNetwork error.
2762        let model = StallingModelClient::new(vec![]); // all default to hang
2763        let calls = model.calls.clone();
2764        let harness = AgentLoopHarness::new(model, MockToolRuntime::new())
2765            .with_stream_resilience(Duration::from_millis(50), 2);
2766        let mut rx = harness
2767            .run_turn(NativeTurnInput {
2768                prompt_text: "hi".into(),
2769                system_prompt: None,
2770                attachments: vec![],
2771                cancel_token: None,
2772                prior_messages: vec![],
2773                context_path: None,
2774                turn_end_grace: None,
2775            })
2776            .await
2777            .unwrap();
2778
2779        let mut saw_error = false;
2780        while let Some(item) = rx.recv().await {
2781            match item {
2782                Err(NativeHarnessError::ModelNetwork(msg)) => {
2783                    assert!(msg.contains("stalled"), "got {msg:?}");
2784                    saw_error = true;
2785                    break;
2786                }
2787                Err(other) => panic!("unexpected error variant: {other:?}"),
2788                Ok(_) => {}
2789            }
2790        }
2791        assert!(
2792            saw_error,
2793            "expected ModelNetwork stall error after budget exhausted"
2794        );
2795        assert_eq!(
2796            calls.load(Ordering::SeqCst),
2797            2,
2798            "two establishments (initial + one reconnect)"
2799        );
2800    }
2801
2802    #[tokio::test(start_paused = true)]
2803    async fn agent_loop_does_not_reconnect_after_stall_with_partial_output() {
2804        // Stream emits text (the user now sees it) then stalls. Even though
2805        // the reconnect budget is generous, a stall *after* output is
2806        // terminal — reconnecting would re-issue the request and duplicate
2807        // what was already shown. Expect: text once, then a ModelNetwork
2808        // error, and exactly one establishment (no reconnect).
2809        let model = StallingModelClient::new(vec![StallBehavior::EmitThenHang(vec![
2810            ModelChunk::TextDelta {
2811                msg_id: "m".into(),
2812                delta: "partial".into(),
2813            },
2814        ])]);
2815        let calls = model.calls.clone();
2816        let harness = AgentLoopHarness::new(model, MockToolRuntime::new())
2817            .with_stream_resilience(Duration::from_millis(50), 5);
2818        let mut rx = harness
2819            .run_turn(NativeTurnInput {
2820                prompt_text: "hi".into(),
2821                system_prompt: None,
2822                attachments: vec![],
2823                cancel_token: None,
2824                prior_messages: vec![],
2825                context_path: None,
2826                turn_end_grace: None,
2827            })
2828            .await
2829            .unwrap();
2830
2831        let mut text = String::new();
2832        let mut saw_error = false;
2833        while let Some(item) = rx.recv().await {
2834            match item {
2835                Ok(HarnessInternalEvent::AssistantTextChunk { delta, .. }) => text.push_str(&delta),
2836                Err(NativeHarnessError::ModelNetwork(_)) => {
2837                    saw_error = true;
2838                    break;
2839                }
2840                Err(other) => panic!("unexpected error variant: {other:?}"),
2841                Ok(_) => {}
2842            }
2843        }
2844        assert!(saw_error, "expected terminal ModelNetwork error");
2845        assert_eq!(
2846            text, "partial",
2847            "partial output delivered once, not replayed"
2848        );
2849        assert_eq!(
2850            calls.load(Ordering::SeqCst),
2851            1,
2852            "no reconnect once output has reached the user"
2853        );
2854    }
2855
2856    #[tokio::test]
2857    async fn agent_loop_accumulates_thinking_chunks_and_signature() {
2858        // Anthropic-style step: thinking deltas + signature, then text,
2859        // then Done. Asserts that:
2860        //   * each ThinkingDelta with non-empty text emits an
2861        //     AssistantThinkingChunk;
2862        //   * the signature latches and ends up on
2863        //     ChatMessage::Assistant.thinking;
2864        //   * an empty-text ThinkingDelta carrying a signature does NOT
2865        //     emit a chunk (signature-only chunks are silent).
2866        let model = StreamingFakeClient::new(vec![vec![
2867            ModelChunk::ThinkingDelta {
2868                thinking_id: "th_1".into(),
2869                delta: "let me think...".into(),
2870                signature: None,
2871            },
2872            ModelChunk::ThinkingDelta {
2873                thinking_id: "th_1".into(),
2874                delta: "".into(),
2875                signature: Some("sig_abc".into()),
2876            },
2877            ModelChunk::TextDelta {
2878                msg_id: "m1".into(),
2879                delta: "ok".into(),
2880            },
2881            ModelChunk::Done {
2882                stop_reason: "end_turn".into(),
2883                usage: None,
2884            },
2885        ]]);
2886        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2887        let mut rx = harness
2888            .run_turn(NativeTurnInput {
2889                prompt_text: "hi".into(),
2890                system_prompt: None,
2891                attachments: vec![],
2892                cancel_token: None,
2893                prior_messages: vec![],
2894                context_path: None,
2895                turn_end_grace: None,
2896            })
2897            .await
2898            .unwrap();
2899
2900        let mut thinking_chunks: Vec<String> = Vec::new();
2901        let mut text_chunks: Vec<String> = Vec::new();
2902        let mut saw_end = false;
2903        while let Some(item) = rx.recv().await {
2904            match item.unwrap() {
2905                HarnessInternalEvent::AssistantThinkingChunk { msg_id, delta } => {
2906                    assert_eq!(msg_id, "thinking_native_0");
2907                    thinking_chunks.push(delta);
2908                }
2909                HarnessInternalEvent::AssistantTextChunk { msg_id, delta } => {
2910                    assert_eq!(msg_id, "msg_native_0");
2911                    text_chunks.push(delta);
2912                }
2913                HarnessInternalEvent::TurnEnd { .. } => {
2914                    saw_end = true;
2915                    break;
2916                }
2917                other => panic!("unexpected event: {other:?}"),
2918            }
2919        }
2920        // Only the non-empty thinking delta emits a chunk; signature-only
2921        // chunk is silent.
2922        assert_eq!(thinking_chunks, vec!["let me think..."]);
2923        assert_eq!(text_chunks, vec!["ok"]);
2924        assert!(saw_end);
2925    }
2926
2927    #[tokio::test]
2928    async fn agent_loop_runs_tool_then_final_message() {
2929        let harness = AgentLoopHarness::new(
2930            ScriptedModelClient,
2931            MockToolRuntime::new().with_file("README.md", "hello"),
2932        );
2933        let mut rx = harness
2934            .run_turn(NativeTurnInput {
2935                prompt_text: "read README.md".into(),
2936                system_prompt: None,
2937                attachments: vec![],
2938                cancel_token: None,
2939                prior_messages: vec![],
2940                context_path: None,
2941                turn_end_grace: None,
2942            })
2943            .await
2944            .unwrap();
2945
2946        assert!(matches!(
2947            rx.recv().await.unwrap().unwrap(),
2948            HarnessInternalEvent::AssistantTextChunk { .. }
2949        ));
2950        assert!(matches!(
2951            rx.recv().await.unwrap().unwrap(),
2952            HarnessInternalEvent::ToolCall { ref name, .. } if name == "read"
2953        ));
2954        assert!(matches!(
2955            rx.recv().await.unwrap().unwrap(),
2956            HarnessInternalEvent::ToolResult { .. }
2957        ));
2958        assert!(matches!(
2959            rx.recv().await.unwrap().unwrap(),
2960            HarnessInternalEvent::AssistantTextChunk { .. }
2961        ));
2962        assert!(matches!(
2963            rx.recv().await.unwrap().unwrap(),
2964            HarnessInternalEvent::TurnEnd { .. }
2965        ));
2966        assert!(rx.recv().await.is_none());
2967    }
2968
2969    /// `TurnEnd.final_messages` must reflect the whole conversation:
2970    /// every `prior_messages` entry RD seeded the turn with, plus the
2971    /// new user prompt, plus the assistant's reply. This is the
2972    /// contract RD's `native_history` slot depends on for multi-turn
2973    /// replay — if it ever shrinks (e.g. we accidentally clone before
2974    /// the final push), same-process multi-turn loses history.
2975    #[tokio::test]
2976    async fn agent_loop_turn_end_carries_full_message_history() {
2977        let model = QueueModelClient::new(vec![ModelResponse::Message {
2978            text: "second reply".into(),
2979            stop_reason: "end_turn".into(),
2980            usage: None,
2981        }]);
2982        let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2983        // Simulate "RD captured this from a previous turn".
2984        let prior = vec![
2985            ChatMessage::User {
2986                content: "first prompt".into(),
2987                attachments: vec![],
2988            },
2989            ChatMessage::Assistant {
2990                text: Some("first reply".into()),
2991                tool_calls: vec![],
2992                thinking: None,
2993                usage: None,
2994            },
2995        ];
2996        let mut rx = harness
2997            .run_turn(NativeTurnInput {
2998                prompt_text: "second prompt".into(),
2999                system_prompt: None,
3000                attachments: vec![],
3001                cancel_token: None,
3002                prior_messages: prior,
3003                context_path: None,
3004                turn_end_grace: None,
3005            })
3006            .await
3007            .unwrap();
3008        let mut final_messages: Option<Vec<ChatMessage>> = None;
3009        while let Some(item) = rx.recv().await {
3010            if let HarnessInternalEvent::TurnEnd {
3011                final_messages: m, ..
3012            } = item.unwrap()
3013            {
3014                final_messages = Some(m);
3015                break;
3016            }
3017        }
3018        let msgs = final_messages.expect("TurnEnd carried final_messages");
3019        // [user-1, assistant-1, user-2, assistant-2] — 4 entries.
3020        assert_eq!(msgs.len(), 4, "got {msgs:?}");
3021        match &msgs[0] {
3022            ChatMessage::User { content, .. } => assert_eq!(content, "first prompt"),
3023            other => panic!("msgs[0] not user-1: {other:?}"),
3024        }
3025        match &msgs[1] {
3026            ChatMessage::Assistant { text, .. } => {
3027                assert_eq!(text.as_deref(), Some("first reply"));
3028            }
3029            other => panic!("msgs[1] not assistant-1: {other:?}"),
3030        }
3031        match &msgs[2] {
3032            ChatMessage::User { content, .. } => assert_eq!(content, "second prompt"),
3033            other => panic!("msgs[2] not user-2: {other:?}"),
3034        }
3035        match &msgs[3] {
3036            ChatMessage::Assistant { text, .. } => {
3037                assert_eq!(text.as_deref(), Some("second reply"));
3038            }
3039            other => panic!("msgs[3] not assistant-2: {other:?}"),
3040        }
3041    }
3042
3043    /// Tool runtime that sleeps for a configurable duration before
3044    /// returning. Records the actual concurrency observed (max number
3045    /// of in-flight invocations at any point) so we can assert the
3046    /// agent loop is truly running them in parallel, not interleaving.
3047    #[derive(Clone)]
3048    struct ConcurrencyProbeRuntime {
3049        sleep_for: std::time::Duration,
3050        in_flight: Arc<AtomicUsize>,
3051        max_concurrency: Arc<AtomicUsize>,
3052        call_order: Arc<Mutex<Vec<String>>>,
3053        cancelled: Arc<AtomicUsize>,
3054    }
3055
3056    impl ConcurrencyProbeRuntime {
3057        fn new(sleep_for: std::time::Duration) -> Self {
3058            Self {
3059                sleep_for,
3060                in_flight: Arc::new(AtomicUsize::new(0)),
3061                max_concurrency: Arc::new(AtomicUsize::new(0)),
3062                call_order: Arc::new(Mutex::new(Vec::new())),
3063                cancelled: Arc::new(AtomicUsize::new(0)),
3064            }
3065        }
3066    }
3067
3068    #[async_trait]
3069    impl ToolRuntime for ConcurrencyProbeRuntime {
3070        fn specs(&self) -> Vec<crate::tools::ToolSpec> {
3071            vec![crate::tools::ToolSpec {
3072                name: "slow".into(),
3073                description: "sleeps".into(),
3074                input_schema: serde_json::json!({"type": "object"}),
3075            }]
3076        }
3077
3078        async fn invoke(
3079            &self,
3080            invocation: ToolInvocation,
3081        ) -> Result<ToolOutcome, ToolRuntimeError> {
3082            self.call_order.lock().unwrap().push(invocation.id.clone());
3083            let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
3084            let mut prev = self.max_concurrency.load(Ordering::SeqCst);
3085            while now > prev {
3086                match self.max_concurrency.compare_exchange(
3087                    prev,
3088                    now,
3089                    Ordering::SeqCst,
3090                    Ordering::SeqCst,
3091                ) {
3092                    Ok(_) => break,
3093                    Err(actual) => prev = actual,
3094                }
3095            }
3096            tokio::time::sleep(self.sleep_for).await;
3097            self.in_flight.fetch_sub(1, Ordering::SeqCst);
3098            Ok(ToolOutcome {
3099                output: Ok(serde_json::json!({"slept": true, "id": invocation.id})),
3100                attachments: vec![],
3101            })
3102        }
3103
3104        async fn invoke_cancellable(
3105            &self,
3106            invocation: ToolInvocation,
3107            cancel: Option<&CancellationToken>,
3108        ) -> Result<ToolOutcome, ToolRuntimeError> {
3109            self.call_order.lock().unwrap().push(invocation.id.clone());
3110            let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
3111            let mut prev = self.max_concurrency.load(Ordering::SeqCst);
3112            while now > prev {
3113                match self.max_concurrency.compare_exchange(
3114                    prev,
3115                    now,
3116                    Ordering::SeqCst,
3117                    Ordering::SeqCst,
3118                ) {
3119                    Ok(_) => break,
3120                    Err(actual) => prev = actual,
3121                }
3122            }
3123            if let Some(token) = cancel {
3124                tokio::select! {
3125                    _ = token.cancelled() => {
3126                        self.cancelled.fetch_add(1, Ordering::SeqCst);
3127                        self.in_flight.fetch_sub(1, Ordering::SeqCst);
3128                        Err(ToolRuntimeError::Runtime("cancelled".into()))
3129                    }
3130                    _ = tokio::time::sleep(self.sleep_for) => {
3131                        self.in_flight.fetch_sub(1, Ordering::SeqCst);
3132                        Ok(ToolOutcome {
3133                            output: Ok(serde_json::json!({"slept": true, "id": invocation.id})),
3134                            attachments: vec![],
3135                        })
3136                    }
3137                }
3138            } else {
3139                tokio::time::sleep(self.sleep_for).await;
3140                self.in_flight.fetch_sub(1, Ordering::SeqCst);
3141                Ok(ToolOutcome {
3142                    output: Ok(serde_json::json!({"slept": true, "id": invocation.id})),
3143                    attachments: vec![],
3144                })
3145            }
3146        }
3147    }
3148
3149    /// When the model returns multiple `tool_use` blocks in a single
3150    /// step (parallel_tool_calls on OpenAI / multi tool_use on
3151    /// Anthropic), the agent loop MUST dispatch them concurrently —
3152    /// not sequentially. Before the F3 fix, only the first one was
3153    /// invoked and the rest were silently dropped.
3154    #[tokio::test]
3155    async fn agent_loop_runs_multi_tool_calls_concurrently() {
3156        // One step that emits 3 tool_use blocks back-to-back, then a
3157        // second step that returns a final message.
3158        let model = StreamingFakeClient::new(vec![
3159            vec![
3160                ModelChunk::ToolCallStart {
3161                    id: "tc_a".into(),
3162                    name: "slow".into(),
3163                },
3164                ModelChunk::ToolCallEnd {
3165                    id: "tc_a".into(),
3166                    input: Some(json!({})),
3167                },
3168                ModelChunk::ToolCallStart {
3169                    id: "tc_b".into(),
3170                    name: "slow".into(),
3171                },
3172                ModelChunk::ToolCallEnd {
3173                    id: "tc_b".into(),
3174                    input: Some(json!({})),
3175                },
3176                ModelChunk::ToolCallStart {
3177                    id: "tc_c".into(),
3178                    name: "slow".into(),
3179                },
3180                ModelChunk::ToolCallEnd {
3181                    id: "tc_c".into(),
3182                    input: Some(json!({})),
3183                },
3184                ModelChunk::Done {
3185                    stop_reason: "tool_use".into(),
3186                    usage: None,
3187                },
3188            ],
3189            vec![
3190                ModelChunk::TextDelta {
3191                    msg_id: "remote".into(),
3192                    delta: "done".into(),
3193                },
3194                ModelChunk::Done {
3195                    stop_reason: "end_turn".into(),
3196                    usage: None,
3197                },
3198            ],
3199        ]);
3200        let probe = ConcurrencyProbeRuntime::new(std::time::Duration::from_millis(80));
3201        let max_concurrency = probe.max_concurrency.clone();
3202        let harness = AgentLoopHarness::new(model, probe);
3203
3204        let start = std::time::Instant::now();
3205        let mut rx = harness
3206            .run_turn(NativeTurnInput {
3207                prompt_text: "go".into(),
3208                system_prompt: None,
3209                attachments: vec![],
3210                cancel_token: None,
3211                prior_messages: vec![],
3212                context_path: None,
3213                turn_end_grace: None,
3214            })
3215            .await
3216            .unwrap();
3217        let mut tool_results = 0;
3218        while let Some(item) = rx.recv().await {
3219            match item.unwrap() {
3220                HarnessInternalEvent::ToolResult { .. } => tool_results += 1,
3221                HarnessInternalEvent::TurnEnd { .. } => break,
3222                _ => {}
3223            }
3224        }
3225        let elapsed = start.elapsed();
3226        // All 3 tools surfaced results — none were silently dropped.
3227        assert_eq!(
3228            tool_results, 3,
3229            "expected 3 tool results, got {tool_results}"
3230        );
3231        // Concurrency probe saw all 3 in flight simultaneously.
3232        assert_eq!(
3233            max_concurrency.load(Ordering::SeqCst),
3234            3,
3235            "expected max concurrency 3 (parallel dispatch), got {}",
3236            max_concurrency.load(Ordering::SeqCst)
3237        );
3238        // Wall clock < 3× sleep duration confirms parallelism (3 × 80ms
3239        // = 240ms sequential; parallel should be ~80ms + scheduler
3240        // overhead, allow up to 200ms for slow CI).
3241        assert!(
3242            elapsed < std::time::Duration::from_millis(200),
3243            "elapsed {elapsed:?} suggests sequential execution"
3244        );
3245    }
3246
3247    /// When the cancel token fires while tool invocations are in
3248    /// flight, the agent loop must emit a clean `TurnEnd { interrupt }`
3249    /// and stop — not wait for the tools to drain naturally.
3250    #[tokio::test]
3251    async fn agent_loop_cancels_in_flight_tool_calls() {
3252        let model = StreamingFakeClient::new(vec![vec![
3253            ModelChunk::ToolCallStart {
3254                id: "tc_slow".into(),
3255                name: "slow".into(),
3256            },
3257            ModelChunk::ToolCallEnd {
3258                id: "tc_slow".into(),
3259                input: Some(json!({})),
3260            },
3261            ModelChunk::Done {
3262                stop_reason: "tool_use".into(),
3263                usage: None,
3264            },
3265        ]]);
3266        // 5-second sleep — if cancel didn't propagate, the test would
3267        // take 5s. We assert it returns in < 200ms.
3268        let probe = ConcurrencyProbeRuntime::new(std::time::Duration::from_secs(5));
3269        let cancelled_count = probe.cancelled.clone();
3270        let harness = AgentLoopHarness::new(model, probe);
3271
3272        let cancel = CancellationToken::new();
3273        let cancel_for_input = cancel.clone();
3274        let mut rx = harness
3275            .run_turn(NativeTurnInput {
3276                prompt_text: "go".into(),
3277                system_prompt: None,
3278                attachments: vec![],
3279                cancel_token: Some(cancel_for_input),
3280                prior_messages: vec![],
3281                context_path: None,
3282                turn_end_grace: None,
3283            })
3284            .await
3285            .unwrap();
3286
3287        // Let the tool spin up briefly, then cancel.
3288        tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3289        cancel.cancel();
3290
3291        let start = std::time::Instant::now();
3292        let mut saw_interrupt = false;
3293        while let Some(item) = rx.recv().await {
3294            if let HarnessInternalEvent::TurnEnd { stop_reason, .. } = item.unwrap() {
3295                assert_eq!(stop_reason, "interrupt");
3296                saw_interrupt = true;
3297                break;
3298            }
3299        }
3300        let elapsed = start.elapsed();
3301        assert!(saw_interrupt, "must see interrupt TurnEnd");
3302        assert!(
3303            elapsed < std::time::Duration::from_millis(200),
3304            "cancel propagation took too long: {elapsed:?}"
3305        );
3306        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
3307        while cancelled_count.load(Ordering::SeqCst) == 0 && tokio::time::Instant::now() < deadline
3308        {
3309            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
3310        }
3311        assert_eq!(
3312            cancelled_count.load(Ordering::SeqCst),
3313            1,
3314            "tool runtime must observe the cancellation token"
3315        );
3316    }
3317
3318    #[derive(Clone)]
3319    struct DefaultCancellationProbe {
3320        started: Arc<tokio::sync::Notify>,
3321        completed: Arc<std::sync::atomic::AtomicBool>,
3322    }
3323
3324    #[async_trait]
3325    impl ToolRuntime for DefaultCancellationProbe {
3326        fn specs(&self) -> Vec<crate::tools::ToolSpec> {
3327            vec![crate::tools::ToolSpec {
3328                name: "slow".into(),
3329                description: "records a delayed side effect".into(),
3330                input_schema: serde_json::json!({"type": "object"}),
3331            }]
3332        }
3333
3334        async fn invoke(
3335            &self,
3336            _invocation: ToolInvocation,
3337        ) -> Result<ToolOutcome, ToolRuntimeError> {
3338            self.started.notify_one();
3339            tokio::time::sleep(std::time::Duration::from_millis(200)).await;
3340            self.completed.store(true, Ordering::SeqCst);
3341            Ok(ToolOutcome {
3342                output: Ok(serde_json::json!({"completed": true})),
3343                attachments: vec![],
3344            })
3345        }
3346    }
3347
3348    /// The default `ToolRuntime::invoke_cancellable` must drop an invocation
3349    /// that does not implement custom cancellation. Before the fix, the agent
3350    /// loop detached its task and the delayed side effect happened after
3351    /// TurnEnd{interrupt}.
3352    #[tokio::test]
3353    async fn agent_loop_cancel_drops_default_tool_future() {
3354        let model = StreamingFakeClient::new(vec![vec![
3355            ModelChunk::ToolCallStart {
3356                id: "tc_slow".into(),
3357                name: "slow".into(),
3358            },
3359            ModelChunk::ToolCallEnd {
3360                id: "tc_slow".into(),
3361                input: Some(json!({})),
3362            },
3363            ModelChunk::Done {
3364                stop_reason: "tool_use".into(),
3365                usage: None,
3366            },
3367        ]]);
3368        let started = Arc::new(tokio::sync::Notify::new());
3369        let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3370        let runtime = DefaultCancellationProbe {
3371            started: started.clone(),
3372            completed: completed.clone(),
3373        };
3374        let cancel = CancellationToken::new();
3375        let harness = AgentLoopHarness::new(model, runtime);
3376        let mut rx = harness
3377            .run_turn(NativeTurnInput {
3378                prompt_text: "go".into(),
3379                system_prompt: None,
3380                attachments: vec![],
3381                cancel_token: Some(cancel.clone()),
3382                prior_messages: vec![],
3383                context_path: None,
3384                turn_end_grace: None,
3385            })
3386            .await
3387            .unwrap();
3388
3389        started.notified().await;
3390        cancel.cancel();
3391
3392        let mut saw_interrupt = false;
3393        while let Some(item) = rx.recv().await {
3394            if let HarnessInternalEvent::TurnEnd { stop_reason, .. } = item.unwrap() {
3395                assert_eq!(stop_reason, "interrupt");
3396                saw_interrupt = true;
3397                break;
3398            }
3399        }
3400        assert!(saw_interrupt);
3401        tokio::time::sleep(std::time::Duration::from_millis(300)).await;
3402        assert!(
3403            !completed.load(Ordering::SeqCst),
3404            "cancelled tool produced a side effect after TurnEnd"
3405        );
3406    }
3407
3408    /// Cancellation while the tool is in flight must put the harness
3409    /// into a clean state: TurnEnd.final_messages carries the
3410    /// assistant tool_use blocks but no synthetic tool_result rows
3411    /// (since the tool never finished).
3412    #[tokio::test]
3413    async fn agent_loop_cancel_during_tools_yields_clean_history() {
3414        let model = StreamingFakeClient::new(vec![vec![
3415            ModelChunk::ToolCallStart {
3416                id: "tc_a".into(),
3417                name: "slow".into(),
3418            },
3419            ModelChunk::ToolCallEnd {
3420                id: "tc_a".into(),
3421                input: Some(json!({})),
3422            },
3423            ModelChunk::ToolCallStart {
3424                id: "tc_b".into(),
3425                name: "slow".into(),
3426            },
3427            ModelChunk::ToolCallEnd {
3428                id: "tc_b".into(),
3429                input: Some(json!({})),
3430            },
3431            ModelChunk::Done {
3432                stop_reason: "tool_use".into(),
3433                usage: None,
3434            },
3435        ]]);
3436        let probe = ConcurrencyProbeRuntime::new(std::time::Duration::from_secs(3));
3437        let harness = AgentLoopHarness::new(model, probe);
3438
3439        let cancel = CancellationToken::new();
3440        let cancel_for_input = cancel.clone();
3441        let mut rx = harness
3442            .run_turn(NativeTurnInput {
3443                prompt_text: "go".into(),
3444                system_prompt: None,
3445                attachments: vec![],
3446                cancel_token: Some(cancel_for_input),
3447                prior_messages: vec![],
3448                context_path: None,
3449                turn_end_grace: None,
3450            })
3451            .await
3452            .unwrap();
3453
3454        tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3455        cancel.cancel();
3456
3457        let mut final_msgs = None;
3458        while let Some(item) = rx.recv().await {
3459            if let HarnessInternalEvent::TurnEnd { final_messages, .. } = item.unwrap() {
3460                final_msgs = Some(final_messages);
3461                break;
3462            }
3463        }
3464        let msgs = final_msgs.expect("interrupt TurnEnd");
3465        // History: [user, assistant(tool_use a + b)]
3466        // — no tool_result rows because the tools never finished.
3467        assert_eq!(msgs.len(), 2, "expected 2 messages, got {msgs:?}");
3468        match &msgs[1] {
3469            ChatMessage::Assistant { tool_calls, .. } => {
3470                assert_eq!(tool_calls.len(), 2);
3471                assert_eq!(tool_calls[0].id, "tc_a");
3472                assert_eq!(tool_calls[1].id, "tc_b");
3473            }
3474            other => panic!("msgs[1] not assistant: {other:?}"),
3475        }
3476    }
3477
3478    /// A runtime whose `invoke_cancellable` deliberately ignores the cancel
3479    /// token and resolves only when the remote work is done — the shape of
3480    /// a cancellation-aware sandbox runtime that waits out the remote
3481    /// TERM→SIGKILL tail before confirming the writer is dead (issue #18:
3482    /// that tail regularly exceeds the historical 1s window).
3483    #[derive(Clone)]
3484    struct WriterProbeRuntime {
3485        sleep_for: std::time::Duration,
3486        completed: Arc<std::sync::atomic::AtomicBool>,
3487    }
3488
3489    #[async_trait]
3490    impl ToolRuntime for WriterProbeRuntime {
3491        fn specs(&self) -> Vec<crate::tools::ToolSpec> {
3492            vec![crate::tools::ToolSpec {
3493                name: "slow".into(),
3494                description: "writer that ignores cancel".into(),
3495                input_schema: serde_json::json!({"type": "object"}),
3496            }]
3497        }
3498
3499        async fn invoke(
3500            &self,
3501            _invocation: ToolInvocation,
3502        ) -> Result<ToolOutcome, ToolRuntimeError> {
3503            tokio::time::sleep(self.sleep_for).await;
3504            self.completed.store(true, Ordering::SeqCst);
3505            Ok(ToolOutcome {
3506                output: Ok(json!({"done": true})),
3507                attachments: vec![],
3508            })
3509        }
3510
3511        async fn invoke_cancellable(
3512            &self,
3513            invocation: ToolInvocation,
3514            _cancel: Option<&CancellationToken>,
3515        ) -> Result<ToolOutcome, ToolRuntimeError> {
3516            self.invoke(invocation).await
3517        }
3518    }
3519
3520    fn writer_tool_turn() -> StreamingFakeClient {
3521        StreamingFakeClient::new(vec![vec![
3522            ModelChunk::ToolCallStart {
3523                id: "tc_slow".into(),
3524                name: "slow".into(),
3525            },
3526            ModelChunk::ToolCallEnd {
3527                id: "tc_slow".into(),
3528                input: Some(json!({})),
3529            },
3530            ModelChunk::Done {
3531                stop_reason: "tool_use".into(),
3532                usage: None,
3533            },
3534        ]])
3535    }
3536
3537    /// Issue #18 acceptance 1: a tool future that resolves 400ms after
3538    /// cancel must be drained when the per-turn grace window is large
3539    /// enough — `TurnEnd{interrupt}` carries an empty `pending_tools` and
3540    /// only lands after the writer actually finished.
3541    #[tokio::test]
3542    async fn turn_end_grace_drains_unresolved_tool_futures() {
3543        let model = writer_tool_turn();
3544        let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3545        let runtime = WriterProbeRuntime {
3546            sleep_for: std::time::Duration::from_millis(400),
3547            completed: completed.clone(),
3548        };
3549        let harness = AgentLoopHarness::new(model, runtime);
3550        let cancel = CancellationToken::new();
3551        let mut rx = harness
3552            .run_turn(NativeTurnInput {
3553                prompt_text: "go".into(),
3554                system_prompt: None,
3555                attachments: vec![],
3556                cancel_token: Some(cancel.clone()),
3557                prior_messages: vec![],
3558                context_path: None,
3559                turn_end_grace: Some(std::time::Duration::from_secs(3)),
3560            })
3561            .await
3562            .unwrap();
3563
3564        tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3565        cancel.cancel();
3566
3567        let start = std::time::Instant::now();
3568        let (saw_interrupt, pending) = loop {
3569            match rx.recv().await {
3570                Some(Ok(HarnessInternalEvent::TurnEnd {
3571                    stop_reason,
3572                    pending_tools,
3573                    ..
3574                })) => break (stop_reason == "interrupt", pending_tools),
3575                Some(_) => continue,
3576                None => panic!("channel closed without TurnEnd"),
3577            }
3578        };
3579        let elapsed = start.elapsed();
3580        assert!(saw_interrupt);
3581        assert!(
3582            elapsed >= std::time::Duration::from_millis(350),
3583            "TurnEnd landed at {elapsed:?} — writer future was not drained"
3584        );
3585        assert!(
3586            elapsed < std::time::Duration::from_secs(2),
3587            "TurnEnd landed at {elapsed:?} — grace window not honoured"
3588        );
3589        assert!(
3590            completed.load(Ordering::SeqCst),
3591            "writer must have finished"
3592        );
3593        assert!(
3594            pending.is_empty(),
3595            "drained turn must report no pending tools, got {pending:?}"
3596        );
3597    }
3598
3599    /// Issue #18 acceptance 2: a writer still running when the grace
3600    /// window expires must be reported loudly via `TurnEnd.pending_tools`
3601    /// — not silently passed over — and its future is dropped unf i nished.
3602    #[tokio::test]
3603    async fn turn_end_grace_expiry_reports_pending_tools() {
3604        let model = writer_tool_turn();
3605        let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3606        let runtime = WriterProbeRuntime {
3607            sleep_for: std::time::Duration::from_secs(5),
3608            completed: completed.clone(),
3609        };
3610        let harness = AgentLoopHarness::new(model, runtime);
3611        let cancel = CancellationToken::new();
3612        let mut rx = harness
3613            .run_turn(NativeTurnInput {
3614                prompt_text: "go".into(),
3615                system_prompt: None,
3616                attachments: vec![],
3617                cancel_token: Some(cancel.clone()),
3618                prior_messages: vec![],
3619                context_path: None,
3620                turn_end_grace: Some(std::time::Duration::from_millis(150)),
3621            })
3622            .await
3623            .unwrap();
3624
3625        tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3626        cancel.cancel();
3627
3628        let start = std::time::Instant::now();
3629        let (saw_interrupt, pending) = loop {
3630            match rx.recv().await {
3631                Some(Ok(HarnessInternalEvent::TurnEnd {
3632                    stop_reason,
3633                    pending_tools,
3634                    ..
3635                })) => break (stop_reason == "interrupt", pending_tools),
3636                Some(_) => continue,
3637                None => panic!("channel closed without TurnEnd"),
3638            }
3639        };
3640        let elapsed = start.elapsed();
3641        assert!(saw_interrupt);
3642        assert!(
3643            elapsed >= std::time::Duration::from_millis(120),
3644            "TurnEnd landed at {elapsed:?} — grace window was skipped"
3645        );
3646        assert!(
3647            elapsed < std::time::Duration::from_secs(1),
3648            "TurnEnd landed at {elapsed:?} — should not wait for the writer"
3649        );
3650        assert_eq!(pending, vec!["tc_slow".to_string()]);
3651        // The future was dropped unf i nished: no side effect after TurnEnd.
3652        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
3653        assert!(!completed.load(Ordering::SeqCst));
3654    }
3655
3656    /// The harness-level builder grace applies when the per-turn override
3657    /// is absent, and a per-turn value wins over the builder value.
3658    #[tokio::test]
3659    async fn turn_end_grace_builder_default_overridden_by_per_turn() {
3660        let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3661        // Builder grace 150ms is too short for the 400ms writer...
3662        let harness = AgentLoopHarness::new(
3663            writer_tool_turn(),
3664            WriterProbeRuntime {
3665                sleep_for: std::time::Duration::from_millis(400),
3666                completed: completed.clone(),
3667            },
3668        )
3669        .with_turn_end_grace(std::time::Duration::from_millis(150));
3670        let cancel = CancellationToken::new();
3671        let mut rx = harness
3672            .run_turn(NativeTurnInput {
3673                prompt_text: "go".into(),
3674                system_prompt: None,
3675                attachments: vec![],
3676                cancel_token: Some(cancel.clone()),
3677                prior_messages: vec![],
3678                context_path: None,
3679                turn_end_grace: None,
3680            })
3681            .await
3682            .unwrap();
3683        tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3684        cancel.cancel();
3685        let (_, pending) = loop {
3686            match rx.recv().await {
3687                Some(Ok(HarnessInternalEvent::TurnEnd {
3688                    stop_reason,
3689                    pending_tools,
3690                    ..
3691                })) => {
3692                    assert_eq!(stop_reason, "interrupt");
3693                    break (true, pending_tools);
3694                }
3695                Some(_) => continue,
3696                None => panic!("channel closed without TurnEnd"),
3697            }
3698        };
3699        assert_eq!(pending, vec!["tc_slow".to_string()]);
3700
3701        // ...but the per-turn override (2s) lets the same writer finish.
3702        let completed2 = Arc::new(std::sync::atomic::AtomicBool::new(false));
3703        let harness = AgentLoopHarness::new(
3704            writer_tool_turn(),
3705            WriterProbeRuntime {
3706                sleep_for: std::time::Duration::from_millis(400),
3707                completed: completed2.clone(),
3708            },
3709        )
3710        .with_turn_end_grace(std::time::Duration::from_millis(150));
3711        let cancel = CancellationToken::new();
3712        let mut rx = harness
3713            .run_turn(NativeTurnInput {
3714                prompt_text: "go".into(),
3715                system_prompt: None,
3716                attachments: vec![],
3717                cancel_token: Some(cancel.clone()),
3718                prior_messages: vec![],
3719                context_path: None,
3720                turn_end_grace: Some(std::time::Duration::from_secs(2)),
3721            })
3722            .await
3723            .unwrap();
3724        tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3725        cancel.cancel();
3726        let (_, pending) = loop {
3727            match rx.recv().await {
3728                Some(Ok(HarnessInternalEvent::TurnEnd { pending_tools, .. })) => {
3729                    break (true, pending_tools)
3730                }
3731                Some(_) => continue,
3732                None => panic!("channel closed without TurnEnd"),
3733            }
3734        };
3735        assert!(pending.is_empty());
3736        assert!(completed2.load(Ordering::SeqCst));
3737    }
3738
3739    #[test]
3740    fn default_turn_end_grace_is_unchanged() {
3741        // The historical hard-coded window stays the default: existing
3742        // hosts must not see their cancellation latency change silently.
3743        assert_eq!(DEFAULT_TURN_END_GRACE, Duration::from_secs(1));
3744    }
3745}