yoagent 0.26.5

//! Context window management — smart truncation and token counting.
//!
//! The #1 engineering challenge for agents. This module provides:
//! - Token estimation (fast, no external deps)
//! - Tiered compaction (tool output truncation → turn summarization → full summary)
//! - Execution limits (max turns, tokens, duration)
//!
//! Designed based on Claude Code's approach: clear old tool outputs first,
//! then summarize conversation if needed.

use crate::types::*;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

// ---------------------------------------------------------------------------
// Token estimation
// ---------------------------------------------------------------------------

/// Rough token estimate: ~4 chars per token for English text.
/// Good enough for context budgeting. Use tiktoken-rs for precision.
pub fn estimate_tokens(text: &str) -> usize {
    text.len().div_ceil(4)
}

/// Estimate tokens for a single message
pub fn message_tokens(msg: &AgentMessage) -> usize {
    match msg {
        AgentMessage::Llm(m) => match m {
            Message::User { content, .. } => content_tokens(content) + 4,
            Message::Assistant { content, .. } => content_tokens(content) + 4,
            Message::ToolResult {
                content, tool_name, ..
            } => content_tokens(content) + estimate_tokens(tool_name) + 8,
        },
        AgentMessage::Custom(ext) => estimate_tokens(&ext.data.to_string()) + 4,
    }
}

fn content_tokens(content: &[Content]) -> usize {
    content
        .iter()
        .map(|c| match c {
            Content::Text { text } => estimate_tokens(text),
            Content::Image { data, .. } => {
                // Estimate tokens from base64 data length:
                // base64 len * 3/4 = raw bytes; ~750 bytes per token for images.
                // Floor at 85 (Anthropic minimum), cap at 16000.
                let raw_bytes = data.len() * 3 / 4;
                (raw_bytes / 750).clamp(85, 16_000)
            }
            Content::Thinking {
                thinking, redacted, ..
            } => {
                // Redacted reasoning is replayed and billed as input too.
                // Both kinds are base64-like text (Bedrock's `redactedContent`
                // is base64; Anthropic's opaque `redacted_thinking` data is
                // base64-alphabet ciphertext), so len * 3/4 ≈ encrypted
                // bytes, at ~4 bytes per token. A rough estimate: the real
                // count is the hidden reasoning's, which only usage reports.
                estimate_tokens(thinking) + redacted.as_ref().map_or(0, |r| r.len() * 3 / 4 / 4)
            }
            Content::ToolCall {
                name, arguments, ..
            } => estimate_tokens(name) + estimate_tokens(&arguments.to_string()) + 8,
        })
        .sum()
}

/// Estimate total tokens for a message list
pub fn total_tokens(messages: &[AgentMessage]) -> usize {
    messages.iter().map(message_tokens).sum()
}

// ---------------------------------------------------------------------------
// Context tracking (real usage + estimates)
// ---------------------------------------------------------------------------

/// Tracks context size using real token counts from provider responses
/// combined with estimates for messages added after the last response.
///
/// This gives more accurate context size tracking than pure estimation,
/// since providers report actual token counts in their usage data.
///
/// # Example
///
/// ```rust
/// use yoagent::context::ContextTracker;
/// use yoagent::types::Usage;
///
/// let mut tracker = ContextTracker::new();
/// // After receiving an assistant response with usage data:
/// tracker.record_usage(&Usage { input: 1500, output: 200, ..Default::default() }, 3);
/// ```
pub struct ContextTracker {
    /// Last known total token count from provider usage
    last_usage_tokens: Option<usize>,
    /// Index of the message that had the last usage
    last_usage_index: Option<usize>,
}

impl ContextTracker {
    pub fn new() -> Self {
        Self {
            last_usage_tokens: None,
            last_usage_index: None,
        }
    }

    /// Record usage from an assistant response.
    ///
    /// Call this after each assistant message to update the tracker
    /// with real token counts from the provider.
    pub fn record_usage(&mut self, usage: &Usage, message_index: usize) {
        let total = usage.input + usage.output + usage.cache_read + usage.cache_write;
        if total > 0 {
            self.last_usage_tokens = Some(total as usize);
            self.last_usage_index = Some(message_index);
        }
    }

    /// Estimate current context size.
    ///
    /// Uses real usage from the last assistant response as a baseline,
    /// then adds estimates (chars/4) for any messages added since.
    /// Falls back to pure estimation if no usage data is available.
    pub fn estimate_context_tokens(&self, messages: &[AgentMessage]) -> usize {
        match (self.last_usage_tokens, self.last_usage_index) {
            (Some(usage_tokens), Some(idx)) if idx < messages.len() => {
                let trailing: usize = messages[idx + 1..].iter().map(message_tokens).sum();
                usage_tokens + trailing
            }
            _ => total_tokens(messages),
        }
    }

    /// Reset tracking (e.g. after compaction replaces messages).
    pub fn reset(&mut self) {
        self.last_usage_tokens = None;
        self.last_usage_index = None;
    }
}

impl Default for ContextTracker {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// Context configuration
// ---------------------------------------------------------------------------

/// Configuration for context management.
///
/// Sizes the built-in tiered compaction. [`LlmCompaction`](crate::LlmCompaction),
/// the summarizing alternative, reads `keep_first`, `keep_recent` and the token
/// budget the same way, but sizes its *spliced* result by its own retained-tail
/// setting rather than by `compact_target_ratio` / `compact_headroom_turns` —
/// which still apply on its deterministic fallback path.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextConfig {
    /// Maximum context tokens (leave room for response)
    pub max_context_tokens: usize,
    /// Tokens reserved for the system prompt
    pub system_prompt_tokens: usize,
    /// Minimum recent messages to always keep (full detail)
    pub keep_recent: usize,
    /// Minimum first messages to always keep
    pub keep_first: usize,
    /// Max lines to keep per tool result, shared by its text blocks (see
    /// [`truncate_tool_output`]), for tools with no override below.
    ///
    /// Tuned for command output — build logs, test runs, `grep` — where
    /// head+tail is a good cut: the first error is at the top, the summary at
    /// the bottom, and the middle is repetition.
    pub tool_output_max_lines: usize,
    /// Per-tool overrides for `tool_output_max_lines`, keyed by tool name.
    ///
    /// One global number cannot serve every tool. Head+tail truncation suits
    /// command output but is the *wrong* cut for a file read, where the middle
    /// is usually the part that matters — a paging read tool bounds itself far
    /// better than a blind truncation can. The default therefore exempts the
    /// built-in `read_file` (see [`DEFAULT_READ_MAX_LINES`]).
    ///
    /// `usize::MAX` disables truncation for a tool.
    ///
    /// [`DEFAULT_READ_MAX_LINES`]: crate::tools::DEFAULT_READ_MAX_LINES
    #[serde(default)]
    pub tool_output_max_lines_overrides: HashMap<String, usize>,
    /// Fraction of the budget that lossy compaction (Level 2/3) aims for.
    ///
    /// Compaction still *triggers* at the full budget, but once it has to
    /// summarize or drop history it reduces to `budget * ratio` instead of to
    /// whatever just barely fits. Without this headroom the very next turn
    /// crosses the budget again and history is rewritten every single turn,
    /// which discards the provider's prefix cache each time.
    ///
    /// Clamped to `[0.05, 1.0]` (a non-finite value falls back to `1.0`).
    /// `1.0` restores the old compact-to-just-fit behaviour. Default: `0.7`.
    ///
    /// Acts as a ceiling on retention when `compact_headroom_turns` is also
    /// set — the headroom policy may compact harder, never softer.
    pub compact_target_ratio: f32,
    /// Compact hard enough to buy this many more turns before the next
    /// compaction, at the session's observed growth rate.
    ///
    /// A fixed ratio has no idea how fast a session is growing, so the
    /// headroom it leaves is arbitrary and the interval between compactions
    /// collapses as history accumulates. This targets the interval directly:
    ///
    /// ```text
    /// target = budget - turns × observed_growth_per_turn
    /// ```
    ///
    /// The effective ratio is `min(target / budget, compact_target_ratio)`,
    /// floored at [`MIN_HEADROOM_RATIO`], so this can only make compaction
    /// more aggressive than the ratio alone — never less.
    ///
    /// Growth is measured by the agent loop, so this is only honoured there;
    /// a direct [`compact_messages`] call uses `compact_target_ratio` as-is.
    /// `None` disables the policy. Default: `Some(30)`.
    #[serde(default = "default_headroom_turns")]
    pub compact_headroom_turns: Option<usize>,
    /// Apply `tool_output_max_lines` when a tool result is appended, rather
    /// than only once the session is over budget.
    ///
    /// Retroactive truncation is the single largest source of prefix-cache
    /// loss: a tool result that has already been sent in full is rewritten
    /// later, invalidating the cache from that message onward. Capping on the
    /// way in means the bytes the provider cached are the bytes that stay.
    /// It also slows context growth, so compaction runs less often.
    ///
    /// The cost is that a long tool output is trimmed even when the budget
    /// had room for it — which is why tools that tolerate it badly are exempted
    /// via `tool_output_max_lines_overrides` rather than by turning this off.
    /// The full output is always visible in the
    /// [`AgentEvent`] stream either way.
    ///
    /// On by default. Only honoured by the agent loop when a `ContextConfig`
    /// is set.
    pub truncate_tool_output_on_append: bool,
}

impl Default for ContextConfig {
    fn default() -> Self {
        Self {
            max_context_tokens: 100_000,
            system_prompt_tokens: 4_000,
            keep_recent: 10,
            keep_first: 2,
            tool_output_max_lines: 200,
            tool_output_max_lines_overrides: default_tool_output_overrides(),
            compact_target_ratio: 0.7,
            compact_headroom_turns: default_headroom_turns(),
            truncate_tool_output_on_append: true,
        }
    }
}

/// Lower bound on the ratio the headroom policy may derive.
///
/// A session whose growth per turn approaches the whole budget would otherwise
/// ask compaction to discard essentially everything.
///
/// 0.30 since 0.22 (0.15 before). Measured offline with
/// `examples/compaction_sweep` (see `docs/evals/compaction-defaults.md`): at
/// the default 96K budget, tool-heavy sessions compacted to the 0.15 floor lost
/// the opening task prompt in 91% of compactions and the latest request in
/// 41%; at 0.30, 0% and 2%, for about 7% more input tokens and no measurable
/// change in compaction count or prefix-cache hit rate. It does not prevent
/// collapse at small budgets (around 26K), where the newest turn alone can
/// exceed any target.
pub const MIN_HEADROOM_RATIO: f32 = 0.30;

fn default_headroom_turns() -> Option<usize> {
    Some(30)
}

/// Tools whose output the default head+tail cap would damage.
///
/// `read_file` bounds itself by paging (`DEFAULT_READ_MAX_LINES`), which is
/// both lossless and directed; cutting the middle out of a source file on top
/// of that would remove exactly the part the agent asked for.
fn default_tool_output_overrides() -> HashMap<String, usize> {
    HashMap::from([
        ("read_file".to_string(), usize::MAX),
        // The retrieval side of the truncation stash. Without this the marker
        // is a lie: the model follows `shared_state get "tool-out-…"`, and the
        // full text it asks for is re-truncated by the very cap that stashed
        // it — head+tail of what it already had, plus a fresh stash entry
        // against the cap on every attempt. A directed fetch of a known key is
        // the model asking for exactly this content; capping it defeats the
        // feature.
        ("shared_state".to_string(), usize::MAX),
    ])
}

impl ContextConfig {
    /// Derive a context config from a model's context window size.
    ///
    /// Reserves 20% of the context window for output tokens, uses the rest
    /// as the compaction budget. All other settings use defaults.
    pub fn from_context_window(context_window: u32) -> Self {
        let max_context_tokens = (context_window as usize) * 80 / 100;
        Self {
            max_context_tokens,
            ..Default::default()
        }
    }

    /// The compaction ratio to use given the session's observed growth rate.
    ///
    /// Returns `compact_target_ratio` unchanged when no headroom policy is
    /// set, or when the growth rate is not yet known. Otherwise derives a
    /// ratio that leaves room for `compact_headroom_turns` more turns, and
    /// takes whichever of the two is more aggressive.
    ///
    /// See [`compact_headroom_turns`](Self::compact_headroom_turns).
    pub fn effective_target_ratio(&self, growth_tokens_per_turn: f64) -> f32 {
        let Some(turns) = self.compact_headroom_turns else {
            return self.compact_target_ratio;
        };
        // NaN and non-positive growth both mean "no usable estimate yet".
        if turns == 0 || !growth_tokens_per_turn.is_finite() || growth_tokens_per_turn <= 0.0 {
            return self.compact_target_ratio;
        }
        let budget = self
            .max_context_tokens
            .saturating_sub(self.system_prompt_tokens) as f64;
        if budget <= 0.0 {
            return self.compact_target_ratio;
        }
        let target = budget - (turns as f64) * growth_tokens_per_turn;
        let derived = (target / budget) as f32;
        derived
            .min(self.compact_target_ratio)
            .max(MIN_HEADROOM_RATIO)
    }

    /// The line cap that applies to a given tool's output.
    pub fn max_lines_for(&self, tool_name: &str) -> usize {
        self.tool_output_max_lines_overrides
            .get(tool_name)
            .copied()
            .unwrap_or(self.tool_output_max_lines)
    }

    /// The token count lossy compaction reduces to, given the trigger budget.
    fn compaction_target(&self, budget: usize) -> usize {
        let ratio = if self.compact_target_ratio.is_finite() {
            self.compact_target_ratio.clamp(0.05, 1.0)
        } else {
            1.0
        };
        ((budget as f64) * (ratio as f64)) as usize
    }
}

// ---------------------------------------------------------------------------
// Compaction strategy
// ---------------------------------------------------------------------------

/// Strategy for compacting messages when context exceeds budget.
///
/// Implement this to customize what happens during compaction:
/// - Index discarded content into a memory store before removal
/// - Apply custom preservation rules (e.g., always keep decisions)
/// - Emit metadata about what was compressed
///
/// See the [Custom Compaction](https://yologdev.github.io/yoagent/concepts/agent-loop.html#custom-compaction)
/// docs for examples.
pub trait CompactionStrategy: crate::rt::MaybeSend + crate::rt::MaybeSync {
    /// Compact messages to fit within the token budget defined by `config`.
    ///
    /// Called before each LLM turn when `context_config` is set.
    fn compact(&self, messages: Vec<AgentMessage>, config: &ContextConfig) -> Vec<AgentMessage>;

    /// Re-price any model this strategy calls against the process-wide
    /// price table now. Called by [`Agent::reprice`](crate::Agent::reprice);
    /// [`LlmCompaction`](crate::LlmCompaction) re-prices its summarization
    /// model. The default does nothing: a strategy that makes no model
    /// request has nothing to price.
    fn reprice(&self) {}
}

/// Default 3-level compaction: truncate tool outputs → summarize turns → drop middle.
///
/// This is used automatically when no custom `CompactionStrategy` is set.
/// You can also compose it inside a custom strategy — run your logic first,
/// then delegate to `compact_messages()` for the actual reduction.
///
/// This is deterministic and free, but the lossy tiers discard early decisions
/// outright. [`LlmCompaction`](crate::LlmCompaction) is the alternative: it
/// spends tokens on a background summarization request and splices the result
/// in as prose. It costs money per compaction and does not reduce prefix-cache
/// breaks — see its module docs for the trade-off.
pub struct DefaultCompaction;

impl CompactionStrategy for DefaultCompaction {
    fn compact(&self, messages: Vec<AgentMessage>, config: &ContextConfig) -> Vec<AgentMessage> {
        compact_messages(messages, config)
    }
}

// ---------------------------------------------------------------------------
// Tiered compaction
// ---------------------------------------------------------------------------

/// Compact messages to fit within the token budget using tiered strategy.
///
/// - Level 1: Truncate tool outputs (keep head + tail)
/// - Level 2: Summarize old turns (replace details with one-liner)
/// - Level 3: Drop the middle (keep first + recent, and the latest user
///   message). When those exceed the target: the head, the latest user
///   message and the newest turn, over the target if need be, up to the budget
///
/// Each level is tried in order. Returns as soon as messages fit.
///
/// # Prefix-cache behaviour
///
/// Every rewrite of an already-sent message discards the provider's prefix
/// cache from that point on, so compaction is built to keep history
/// byte-stable:
///
/// - Level 1 is idempotent — re-running it on settled history is a no-op, so
///   staying above the budget does not re-truncate the same outputs.
/// - Level 1 alone only has to reach the budget; the lossy levels aim for
///   `compact_target_ratio` of it, so once history is rewritten it takes many
///   turns before it must be rewritten again.
/// - Level 3 drops the smallest span that reaches the target and its marker
///   text is constant, so a boundary that lands in the same place twice
///   reproduces the same bytes.
pub fn compact_messages(messages: Vec<AgentMessage>, config: &ContextConfig) -> Vec<AgentMessage> {
    let budget = config
        .max_context_tokens
        .saturating_sub(config.system_prompt_tokens);

    // Already fits?
    if total_tokens(&messages) <= budget {
        return messages;
    }

    let target = config.compaction_target(budget);
    let before = total_tokens(&messages);

    // Level 1: Truncate tool outputs. Checked against the full budget, not the
    // target: this level is idempotent and cheap to repeat, so there is nothing
    // to gain by escalating to the lossy levels before we have to.
    let compacted = level1_truncate_tool_outputs(&messages, config);
    let after_l1 = total_tokens(&compacted);
    if after_l1 < before {
        tracing::debug!(
            "compaction level 1: tool outputs truncated, {} -> {} tokens",
            before,
            after_l1
        );
    }
    if after_l1 <= budget {
        return compacted;
    }

    // Level 2: Summarize old turns (keep recent N full, summarize the rest)
    let before_l2 = compacted.len();
    let compacted = level2_summarize_old_turns(&compacted, config.keep_recent);
    if compacted.len() != before_l2 {
        tracing::debug!(
            "compaction level 2: old turns summarized, {} -> {} messages ({} tokens)",
            before_l2,
            compacted.len(),
            total_tokens(&compacted)
        );
    }
    if total_tokens(&compacted) <= target {
        return compacted;
    }

    // Level 3: Drop the smallest middle span that reaches the target
    level3_drop_middle(&compacted, config, target, budget)
}

/// Level 1: Truncate long tool outputs to head + tail.
///
/// This is the cheapest compaction — preserves conversation structure,
/// just removes verbose tool output middles. In practice this saves
/// 50-70% of context in coding sessions.
fn level1_truncate_tool_outputs(
    messages: &[AgentMessage],
    config: &ContextConfig,
) -> Vec<AgentMessage> {
    messages
        .iter()
        .map(|msg| truncate_tool_output(msg.clone(), config))
        .collect()
}

/// All `Content::Text` of a message, joined with newlines.
///
/// The hash input for [`tool_output_key`] — not a stored value. Per-block
/// stashing means no single entry holds this string.
pub fn message_text(msg: &AgentMessage) -> String {
    block_texts(msg)
        .into_iter()
        .map(|(_, s)| s)
        .collect::<Vec<_>>()
        .join("\n")
}

/// Prefix marking a key as machine-generated truncation stash rather than a
/// variable a caller or model chose to store.
///
/// Used to keep these out of the system-prompt summary: that text is
/// prefix-cached, and an entry appearing there on every stash would change the
/// prompt each turn.
pub const TOOL_OUTPUT_KEY_PREFIX: &str = "tool-out-";

/// The **base** key for a truncated tool result.
///
/// Nothing is stored under this key: per-block keys derive from it via
/// [`block_key`]. The hash input is every text block joined, so the base is
/// stable for the whole result while each block gets its own entry.
///
/// Combines the tool call id with a hash of the full output. The id alone is
/// not enough: `google.rs` and `google_vertex.rs` synthesize ids as a
/// *per-response index* (`google-fc-0`), which restarts every turn — so turn 1
/// and turn 5 would collide, and turn 1's frozen marker would resolve to turn
/// 5's content. Both files already treat those ids as throwaway, stripping them
/// before sending back to the API.
///
/// Hashing the content keeps the key deterministic on replay while making a
/// collision mean the contents were identical anyway.
pub fn tool_output_key(tool_call_id: &str, full_output: &str) -> String {
    format!(
        "{TOOL_OUTPUT_KEY_PREFIX}{tool_call_id}-{:016x}",
        fnv1a(full_output.as_bytes())
    )
}

/// FNV-1a, for the same reason the GASP recorder uses it: `DefaultHasher` is
/// documented as unstable across Rust releases, and a key that moves between
/// compiler versions silently stops resolving.
fn fnv1a(bytes: &[u8]) -> u64 {
    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
    for byte in bytes {
        hash ^= u64::from(*byte);
        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
    }
    hash
}

/// The stash key for one content block of a tool result.
///
/// Block-qualified because a result can carry several text blocks, each cut
/// with its own marker (from one budget shared across them). Sharing one key
/// across them made every marker resolve to all the blocks concatenated —
/// never the one the model had just been told to fetch.
pub fn block_key(base: &str, block: usize) -> String {
    format!("{base}-b{block}")
}

/// The text of each `Content::Text` block, with its index.
///
/// Indices are positions in the **full content vector** — the same enumeration
/// [`truncate_tool_output_keyed`] uses — so the indices it reports in `marked`
/// address these entries directly.
///
/// They are deliberately *not* a count of text blocks. A text-only ordinal
/// would read more naturally and would silently disagree with the marker keys
/// whenever non-text content sits between two text blocks: `[Text, Image,
/// Text]` yields indices 0 and 2 here and must yield `-b0` and `-b2` there.
pub fn block_texts(msg: &AgentMessage) -> Vec<(usize, String)> {
    match msg {
        AgentMessage::Llm(Message::ToolResult { content, .. }) => content
            .iter()
            .enumerate()
            .filter_map(|(i, c)| match c {
                Content::Text { text } => Some((i, text.clone())),
                _ => None,
            })
            .collect(),
        _ => Vec::new(),
    }
}

/// Truncate, optionally naming a per-block stash key inside each marker.
///
/// Returns the message and the indices of blocks that actually got a marker.
///
/// The key is threaded into marker *generation* rather than substituted
/// afterwards. Rewriting the rendered text would hit every occurrence of the
/// marker's shape — including one that came from the tool's own output, which
/// is ordinary for a coding agent reading a log or a session transcript — and
/// would silently do nothing in the case where truncation emits no marker at
/// all. Generating it means the key is present exactly when a marker is.
///
/// Each marker names *its own block's* key, so what the model fetches is the
/// block whose marker it followed — not the whole result flattened, which is
/// what a single shared key produced and which silently dropped any non-text
/// content sitting between the blocks.
///
/// Line budgets come from the blocks' line counts alone (one budget per result,
/// shared), so the unkeyed and keyed passes cut the same lines.
///
/// Public so a caller implementing their own stash sink can reproduce exactly
/// what the loop does: truncate with `None` to learn which blocks carry a
/// marker, stash those, then truncate again with the base key.
pub fn truncate_tool_output_keyed(
    msg: AgentMessage,
    config: &ContextConfig,
    base_key: Option<&str>,
) -> (AgentMessage, Vec<usize>) {
    match msg {
        AgentMessage::Llm(Message::ToolResult {
            tool_call_id,
            tool_name,
            content,
            is_error,
            timestamp,
        }) => {
            let budgets = block_budgets(&content, config.max_lines_for(&tool_name));
            let mut marked = Vec::new();
            let truncated_content: Vec<Content> = content
                .into_iter()
                .zip(budgets)
                .enumerate()
                .map(|(i, (c, budget))| match c {
                    Content::Text { text } => {
                        let key = base_key.map(|b| block_key(b, i));
                        let (out, emitted) = truncate_text_head_tail(&text, budget, key.as_deref());
                        if emitted {
                            marked.push(i);
                        }
                        Content::Text { text: out }
                    }
                    other => other,
                })
                .collect();

            (
                AgentMessage::Llm(Message::ToolResult {
                    tool_call_id,
                    tool_name,
                    content: truncated_content,
                    is_error,
                    timestamp,
                }),
                marked,
            )
        }
        other => (other, Vec::new()),
    }
}

/// Fit a single tool result's text into the tool's line budget, keeping head
/// and tail.
///
/// The budget is the result's, shared by its text blocks: blocks no longer
/// than an even share stay whole, and the others split what is left, each cut
/// head + tail with its own marker. With a budget of 5 or more a cut block
/// keeps at least 5 lines (head, marker, tail), so a result with many blocks
/// can go over the budget; below 5, each block keeps the budget on its own and
/// gets no marker.
///
/// The budget comes from [`ContextConfig::max_lines_for`], so a tool that
/// tolerates head+tail badly can be exempted. Non-tool-result messages pass
/// through untouched, and the operation is idempotent, so applying it when the
/// message is first appended means compaction never has to rewrite it later —
/// which is what keeps the provider's prefix cache intact. See
/// [`ContextConfig::truncate_tool_output_on_append`].
pub fn truncate_tool_output(msg: AgentMessage, config: &ContextConfig) -> AgentMessage {
    truncate_tool_output_keyed(msg, config, None).0
}

/// Lines the truncation marker occupies: a blank line, the marker, a blank line.
const TRUNCATION_MARKER_LINES: usize = 3;

/// The fewest lines a cut block keeps: one head line, the marker's three lines,
/// one tail line — so every cut block says it was cut (when the budget is at
/// least this large).
const MIN_BLOCK_LINES: usize = TRUNCATION_MARKER_LINES + 2;

/// Each content block's line budget, from one budget for the whole result.
///
/// The text blocks share `max_lines`: when they fit together, every block is
/// left alone; otherwise blocks no longer than an even share stay whole and
/// the rest split what is left evenly (the remainder to the earliest), so a
/// many-block result (an MCP server's) fits the budget one block would. A
/// block not left whole is given at least [`MIN_BLOCK_LINES`] (so it is cut
/// with a marker, or kept whole if shorter): when the even share is smaller
/// than that, the result goes over `max_lines` by at most `MIN_BLOCK_LINES`
/// per such block, never silently. When `max_lines < MIN_BLOCK_LINES` the
/// floor is `max_lines` itself, so every block keeps the whole budget, as
/// before, and gets no marker.
///
/// Idempotent with [`truncate_text_head_tail`]: a cut block is at most its
/// budget long, so the result either fits afterwards (and a second pass
/// leaves it alone) or the share fell below the floor, leaving every block
/// not left whole at most the floor long; a second pass gives each such block
/// at least the smaller of its length and the floor, so nothing is cut again.
fn block_budgets(content: &[Content], max_lines: usize) -> Vec<usize> {
    let lines: Vec<Option<usize>> = content
        .iter()
        .map(|c| match c {
            Content::Text { text } => Some(text.lines().count()),
            _ => None,
        })
        .collect();
    let mut budgets: Vec<usize> = vec![usize::MAX; content.len()];
    let total = lines
        .iter()
        .flatten()
        .fold(0usize, |a, n| a.saturating_add(*n));
    if total <= max_lines {
        return budgets;
    }
    // Shortest first: a block that fits an even share of what is left keeps
    // all its lines and leaves the rest to the others.
    let mut order: Vec<(usize, usize)> = lines
        .iter()
        .enumerate()
        .filter_map(|(i, n)| n.map(|n| (n, i)))
        .collect();
    order.sort_unstable();
    let mut left = max_lines;
    let mut rest = order.len();
    let mut cut = Vec::new();
    for (n, i) in order {
        if cut.is_empty() && n <= left / rest {
            budgets[i] = n;
            left -= n;
            rest -= 1;
        } else {
            cut.push(i);
        }
    }
    // The cut blocks, in content order, share what is left.
    cut.sort_unstable();
    let share = left / cut.len();
    let extra = left % cut.len();
    let floor = MIN_BLOCK_LINES.min(max_lines);
    for (k, i) in cut.into_iter().enumerate() {
        budgets[i] = (share + usize::from(k < extra)).max(floor);
    }
    budgets
}

/// Truncate text keeping the head and tail, dropping the middle.
///
/// The marker is charged against `max_lines`, so the result is exactly
/// `max_lines` lines and truncating it again returns it unchanged. That
/// idempotence matters: compaction runs before every turn once a session is
/// over budget, and a marker that shrank on each pass (`950 lines truncated`
/// → `3 lines truncated`) both misreported the loss and rewrote settled
/// history, discarding the provider's prefix cache a second time.
fn truncate_text_head_tail(text: &str, max_lines: usize, key: Option<&str>) -> (String, bool) {
    let lines: Vec<&str> = text.lines().collect();
    if lines.len() <= max_lines {
        return (text.to_string(), false);
    }

    // Not enough room for head + marker + tail — keep the head and skip the
    // marker rather than emitting something that re-truncates forever. No
    // marker means no place to name a stash key, which is why the caller is
    // told `false` and must not stash: the entry would be unreachable.
    if max_lines <= TRUNCATION_MARKER_LINES + 1 {
        return (lines[..max_lines].join("\n"), false);
    }

    let keep = max_lines - TRUNCATION_MARKER_LINES;
    let head = keep / 2;
    let tail = keep - head;
    let omitted = lines.len() - head - tail;

    let marker = match key {
        Some(k) => {
            format!("[... {omitted} lines truncated — full output: shared_state get \"{k}\" ...]")
        }
        None => format!("[... {omitted} lines truncated ...]"),
    };

    let mut result = lines[..head].join("\n");
    result.push_str(&format!("\n\n{marker}\n\n"));
    result.push_str(&lines[lines.len() - tail..].join("\n"));
    // True means "a marker was emitted", not "the marker names a key" — the
    // caller runs an unkeyed pass first precisely to learn whether there is
    // somewhere to put one.
    (result, true)
}

/// Level 2: Summarize old assistant turns.
///
/// Keeps the last `keep_recent` messages in full detail.
/// For older messages: assistant messages with tool calls get replaced
/// with a short summary, and their tool results get dropped.
fn level2_summarize_old_turns(messages: &[AgentMessage], keep_recent: usize) -> Vec<AgentMessage> {
    let len = messages.len();
    if len <= keep_recent {
        return messages.to_vec();
    }

    // `len - keep_recent` can land in the middle of a turn, which would
    // summarize away an assistant message while its tool results stay in the
    // kept section — an orphan every provider rejects. Pull the boundary back
    // onto the turn start so the whole turn is kept intact.
    let boundary = safe_turn_start(messages, len - keep_recent);
    if boundary == 0 {
        return messages.to_vec();
    }
    let mut result = Vec::new();

    let mut i = 0;
    while i < boundary {
        let msg = &messages[i];
        match msg {
            AgentMessage::Llm(Message::Assistant {
                content, timestamp, ..
            }) => {
                // Summarize: extract text content, skip tool call details
                let text_parts: Vec<&str> = content
                    .iter()
                    .filter_map(|c| match c {
                        Content::Text { text } => {
                            if text.len() > 200 {
                                None // Too long, will be replaced
                            } else {
                                Some(text.as_str())
                            }
                        }
                        _ => None,
                    })
                    .collect();

                let tool_count = content
                    .iter()
                    .filter(|c| matches!(c, Content::ToolCall { .. }))
                    .count();

                let summary = if !text_parts.is_empty() {
                    text_parts.join(" ")
                } else if tool_count > 0 {
                    format!("[Assistant used {} tool(s)]", tool_count)
                } else {
                    "[Assistant response]".into()
                };

                // Inherit the summarized turn's timestamp rather than stamping
                // wall-clock time: compaction must be a pure function of its
                // input, or the same history compacts to different bytes on
                // every pass.
                result.push(AgentMessage::Llm(Message::User {
                    content: vec![Content::Text {
                        text: format!("{SUMMARY_PREFIX}{summary}"),
                    }],
                    timestamp: *timestamp,
                }));

                // Skip following tool results that belong to this turn
                i += 1;
                while i < boundary {
                    if let AgentMessage::Llm(Message::ToolResult { .. }) = &messages[i] {
                        i += 1;
                    } else {
                        break;
                    }
                }
                continue;
            }
            AgentMessage::Llm(Message::ToolResult { .. }) => {
                // Skip orphaned tool results in old section
                i += 1;
                continue;
            }
            other => {
                // Keep user messages as-is (they provide intent)
                result.push(other.clone());
            }
        }
        i += 1;
    }

    // Append recent messages in full
    result.extend_from_slice(&messages[boundary..]);
    result
}

/// Prefix of the user-role message that replaces a summarized old turn
/// (level-2 compaction).
pub(crate) const SUMMARY_PREFIX: &str = "[Summary] ";

/// Marker left in place of dropped history.
///
/// The text is constant on purpose. It sits near the front of the message list,
/// so embedding a message count here would change the bytes of the cached
/// prefix on every compaction pass — the count goes to the debug log instead.
pub(crate) const COMPACTION_MARKER: &str =
    "[Context compacted: earlier messages removed to fit the context window]";

fn compaction_marker(timestamp: u64) -> AgentMessage {
    AgentMessage::Llm(Message::User {
        content: vec![Content::Text {
            text: COMPACTION_MARKER.into(),
        }],
        timestamp,
    })
}

pub(crate) fn message_timestamp(msg: &AgentMessage) -> u64 {
    match msg {
        AgentMessage::Llm(
            Message::User { timestamp, .. }
            | Message::Assistant { timestamp, .. }
            | Message::ToolResult { timestamp, .. },
        ) => *timestamp,
        AgentMessage::Custom(_) => 0,
    }
}

/// Does this message open tool calls that a later message must answer?
fn opens_tool_calls(msg: &AgentMessage) -> bool {
    matches!(
        msg,
        AgentMessage::Llm(Message::Assistant { content, .. })
            if content.iter().any(|c| matches!(c, Content::ToolCall { .. }))
    )
}

fn is_tool_result(msg: &AgentMessage) -> bool {
    matches!(msg, AgentMessage::Llm(Message::ToolResult { .. }))
}

/// Pull a kept-head boundary back so the head never ends on an assistant
/// message whose tool results are about to be dropped, nor partway through
/// those results — providers reject a `tool_use` with no matching
/// `tool_result`, and a result with no call.
pub(crate) fn safe_head_end(messages: &[AgentMessage], mut end: usize) -> usize {
    while end > 0
        && (opens_tool_calls(&messages[end - 1]) || messages.get(end).is_some_and(is_tool_result))
    {
        end -= 1;
    }
    end
}

/// Push a kept-tail boundary forward so the tail never opens on a tool result
/// whose originating tool call was dropped — the mirror-image rejection.
fn safe_tail_start(messages: &[AgentMessage], mut start: usize) -> usize {
    while start < messages.len() && is_tool_result(&messages[start]) {
        start += 1;
    }
    start
}

/// Pull a boundary back onto the start of the turn it lands in, so the turn's
/// assistant message and its tool results stay on the same side of the split.
pub(crate) fn safe_turn_start(messages: &[AgentMessage], mut start: usize) -> usize {
    while start > 0 && start < messages.len() && is_tool_result(&messages[start]) {
        start -= 1;
    }
    start
}

/// Level 3: Drop the smallest middle span that brings the history to `target`.
///
/// Keeps `keep_first` messages at the front and at least `keep_recent` at the
/// back, and removes only as much of the middle as the target requires — the
/// old fixed `first + recent` shape discarded everything in between even when
/// a fraction of it would have done. The latest user message stays when the
/// cut passes it, after the marker. When the head and that tail alone exceed
/// the target, [`keep_protected`] decides what stays.
fn level3_drop_middle(
    messages: &[AgentMessage],
    config: &ContextConfig,
    target: usize,
    budget: usize,
) -> Vec<AgentMessage> {
    let len = messages.len();
    let head_end = safe_head_end(messages, config.keep_first.min(len));
    // The furthest the cut may go: everything past it is the guaranteed tail
    // — `keep_recent` messages, and never less than the newest turn.
    let max_cut = len
        .saturating_sub(config.keep_recent)
        .min(newest_turn(messages, head_end).unwrap_or(len));

    if head_end >= max_cut {
        return keep_protected(messages, config, target, budget);
    }

    // suffix[i] = tokens of messages[i..]
    let mut suffix = vec![0usize; len + 1];
    for i in (0..len).rev() {
        suffix[i] = suffix[i + 1] + message_tokens(&messages[i]);
    }
    let head_tokens: usize = messages[..head_end].iter().map(message_tokens).sum();
    let marker_tokens = message_tokens(&compaction_marker(0));
    let fixed = head_tokens + marker_tokens;
    // The latest user message is kept even when the cut passes it.
    let ask = latest_ask(messages, head_end, len);
    let size = |cut: usize| {
        let ask = ask
            .filter(|&i| i < cut)
            .map_or(0, |i| message_tokens(&messages[i]));
        fixed + ask + suffix[cut]
    };

    // Smallest cut that reaches the target — the least history destroyed.
    let cut = (head_end..=max_cut)
        .find(|&c| size(c) <= target)
        .unwrap_or(max_cut);
    // Snapping forward only drops more, so the target still holds.
    let cut = safe_tail_start(messages, cut).max(head_end);

    let removed = cut - head_end;
    if removed == 0 {
        return messages.to_vec();
    }
    tracing::debug!(
        "compaction level 3: dropping {} of {} messages (target {} tokens)",
        removed,
        len,
        target
    );

    let mut result = messages[..head_end].to_vec();
    result.push(compaction_marker(message_timestamp(&messages[head_end])));
    if let Some(i) = ask.filter(|&i| i < cut) {
        result.push(messages[i].clone());
    }
    result.extend_from_slice(&messages[cut..]);

    // The head and the guaranteed tail alone miss the target.
    if total_tokens(&result) > target {
        return keep_protected(messages, config, target, budget);
    }

    result
}

/// Where the newest turn starts: the last assistant message at or after
/// `from` (its results and any user message after them follow it); `None`
/// when there is none — level 2 summarizes every assistant message when
/// `keep_recent` is 0.
fn newest_turn(messages: &[AgentMessage], from: usize) -> Option<usize> {
    (from..messages.len())
        .rev()
        .find(|&i| matches!(messages[i], AgentMessage::Llm(Message::Assistant { .. })))
}

/// The latest user message in `messages[from..to]` that a person wrote — not
/// one the loop injected (a summary, the compaction marker, a nudge).
fn latest_ask(messages: &[AgentMessage], from: usize, to: usize) -> Option<usize> {
    (from..to).rev().find(|&i| match &messages[i] {
        AgentMessage::Llm(Message::User { content, .. }) => !content
            .iter()
            .any(|c| matches!(c, Content::Text { text } if crate::types::is_loop_injected(text))),
        _ => false,
    })
}

/// What level 3 keeps when the head and the `keep_recent` tail alone exceed
/// the target — a small budget, or bulky tool results. In order: the head
/// (the task prompt), the latest user message (the current ask), the newest
/// turn (the call the model is waiting on and its results), then as much of
/// the remaining tail as fits the target.
///
/// The target is a headroom aim; only `budget` is a hard limit. So when the
/// head, the ask and the newest turn alone exceed the target, they are kept
/// anyway, up to the budget, rather than cut to fit: keeping the newest
/// messages that fit the target dropped the task prompt first, and when the
/// newest turn alone exceeded the target, everything but the marker (#150).
/// Past the budget the head goes first (an ask inside it is still the ask),
/// then the ask; the newest turn is cut last, from its start, down to the
/// marker and the ask.
fn keep_protected(
    messages: &[AgentMessage],
    config: &ContextConfig,
    target: usize,
    budget: usize,
) -> Vec<AgentMessage> {
    let len = messages.len();
    if len == 0 {
        return Vec::new();
    }
    let head_end = safe_head_end(messages, config.keep_first.min(len));
    // No assistant after the head: no turn to protect, only the head and
    // the ask.
    let turn = newest_turn(messages, head_end).unwrap_or(len);
    let ask = latest_ask(messages, head_end, len);

    let mut suffix = vec![0usize; len + 1];
    for i in (0..len).rev() {
        suffix[i] = suffix[i + 1] + message_tokens(&messages[i]);
    }
    let head_tokens: usize = messages[..head_end].iter().map(message_tokens).sum();
    let marker_tokens = message_tokens(&compaction_marker(0));
    // An ask's tokens, when it sits before the kept tail (and so is placed
    // after the marker); inside the tail it is counted there.
    let ask_tokens = |ask: Option<usize>, start: usize| {
        ask.filter(|&i| i < start)
            .map_or(0, |i| message_tokens(&messages[i]))
    };

    // Head, ask (when not inside the kept tail) and the tail from `start`.
    let build = |keep_head: bool, ask: Option<usize>, start: usize| {
        let mut out = Vec::new();
        if keep_head {
            out.extend_from_slice(&messages[..head_end]);
        }
        let from = if keep_head { head_end } else { 0 };
        if start > from {
            out.push(compaction_marker(message_timestamp(&messages[from])));
        }
        if let Some(i) = ask.filter(|&i| i < start) {
            out.push(messages[i].clone());
        }
        out.extend_from_slice(&messages[start..]);
        out
    };
    let size = |start: usize| head_tokens + marker_tokens + ask_tokens(ask, start) + suffix[start];

    if size(turn) <= budget {
        // The longest tail, back towards the head, that still fits the target;
        // never shorter than the newest turn, never opening on a tool result.
        let start = (head_end..turn)
            .filter(|&s| !is_tool_result(&messages[s]))
            .find(|&s| size(s) <= target)
            .unwrap_or(turn);
        tracing::debug!(
            "compaction level 3: keeping head, ask and {} of {} messages ({} tokens; target {}, budget {})",
            len - start,
            len,
            size(start),
            target,
            budget
        );
        return build(true, ask, start);
    }
    // Over the budget even so: the head goes first (an ask inside it is
    // still the ask), then the ask; the newest turn is cut last, from its
    // start, never opening on a tool result.
    let ask = latest_ask(messages, 0, len);
    let fits = |ask: Option<usize>, start: usize| {
        marker_tokens + ask_tokens(ask, start) + suffix[start] <= budget
    };
    let starts = (turn..=len).filter(|&s| s == len || !is_tool_result(&messages[s]));
    let whole = [(ask, turn), (None, turn)].into_iter();
    let cut = [ask, None]
        .into_iter()
        .flat_map(|ask| starts.clone().map(move |s| (ask, s)));
    match whole.chain(cut).find(|&(ask, start)| fits(ask, start)) {
        Some((ask, start)) => build(false, ask, start),
        None => vec![compaction_marker(message_timestamp(&messages[0]))],
    }
}

// ---------------------------------------------------------------------------
// Execution limits
// ---------------------------------------------------------------------------

/// Execution limits for the agent loop
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ExecutionLimits {
    /// Maximum number of turns (LLM calls)
    pub max_turns: usize,
    /// Maximum tokens in one run (each `prompt` / `continue_loop` starts
    /// again at 0), counted as [`token_accounting`](Self::token_accounting)
    /// says. By default ([`TokenAccounting::Processed`]) every turn adds its
    /// whole prompt (`input + cache_read + cache_write`) plus its output.
    ///
    /// The prompt resends the conversation each turn, so that total is about
    /// turns × context size — quadratic in turns while the conversation is
    /// still growing. At a 60K-token context the 1M default is reached around
    /// turn 16, cached or not. It is checked before each turn, so the last
    /// turn can go over by one turn's tokens. For long agentic runs raise it
    /// (or set `usize::MAX` and rely on `max_turns`), or count the
    /// conversation's growth instead ([`TokenAccounting::Growth`]).
    pub max_total_tokens: usize,
    /// What [`max_total_tokens`](Self::max_total_tokens) counts. Default
    /// [`TokenAccounting::Processed`].
    #[serde(default)]
    pub token_accounting: TokenAccounting,
    /// Maximum wall-clock time of one run (default 10 minutes); checked
    /// between turns. A single long agentic run reaches it; raise it for
    /// those.
    pub max_duration: std::time::Duration,
    /// Consecutive identical tool calls tolerated before intervening.
    ///
    /// The cheapest catastrophic failure mode is a model calling one tool with
    /// the same arguments forever. The other three limits all fire eventually,
    /// but only after burning the full turn, token and wall-clock budget — so
    /// the run costs its maximum to discover it achieved nothing.
    ///
    /// **Consecutive**, and that word is load-bearing. A different signature
    /// resets the streak, so a model alternating `[a, b, a, b, …]` forever is
    /// *not* detected. That is a deliberate trade, not an oversight: an agent
    /// working through a list calls one tool repeatedly and legitimately, and
    /// a detector that fired on interleaved repeats would be worse than none.
    /// See `an_interleaved_different_call_resets_the_counter`. Widening this
    /// to a windowed count is tracked separately.
    ///
    /// `None` disables the check. `Some(0)` also disables it — there is no
    /// streak length below which a call has not been made. `Some(1)` steers on
    /// the very first tool call of a run, which is almost never wanted.
    /// Default `Some(3)`.
    pub max_consecutive_identical_tool_calls: Option<usize>,
}

impl Default for ExecutionLimits {
    fn default() -> Self {
        Self {
            max_turns: 50,
            max_total_tokens: 1_000_000,
            token_accounting: TokenAccounting::Processed,
            max_duration: std::time::Duration::from_secs(600),
            max_consecutive_identical_tool_calls: Some(3),
        }
    }
}

impl ExecutionLimits {
    /// Cap on turns before the run stops.
    pub fn with_max_turns(mut self, turns: usize) -> Self {
        self.max_turns = turns;
        self
    }

    /// Cap on tokens processed in one run — every turn's whole prompt plus
    /// its output, so it grows much faster than the conversation. See
    /// [`max_total_tokens`](Self::max_total_tokens).
    pub fn with_max_total_tokens(mut self, tokens: usize) -> Self {
        self.max_total_tokens = tokens;
        self
    }

    /// What [`max_total_tokens`](Self::max_total_tokens) counts: every
    /// prompt the model processes (the default) or how much the
    /// conversation grows. See [`TokenAccounting`].
    pub fn with_token_accounting(mut self, accounting: TokenAccounting) -> Self {
        self.token_accounting = accounting;
        self
    }

    /// Wall-clock cap on the run.
    pub fn with_max_duration(mut self, duration: std::time::Duration) -> Self {
        self.max_duration = duration;
        self
    }

    /// Consecutive identical tool calls tolerated before steering, then
    /// stopping. `None` disables loop detection entirely.
    pub fn with_max_consecutive_identical_tool_calls(mut self, calls: Option<usize>) -> Self {
        self.max_consecutive_identical_tool_calls = calls;
        self
    }
}

/// What [`ExecutionLimits::max_total_tokens`] counts.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum TokenAccounting {
    /// Every token the model processes: each turn adds its whole prompt
    /// (`input + cache_read + cache_write`) plus its output. The prompt
    /// resends the conversation, so the total is about turns × context size —
    /// a ceiling on the work a run does. The default.
    #[default]
    Processed,
    /// How much the conversation grows during the run: the first turn's
    /// output, then each turn's growth over the turn before (its prompt plus
    /// output, less the previous turn's). The conversation the run starts
    /// with — history, system prompt, tool definitions, the new prompt — is
    /// the baseline and is not counted, so a run late in a long conversation
    /// has the same allowance as the first. Without compaction the total is
    /// about what the run added to the conversation; it grows with the
    /// conversation, not with turns × context, and cached and uncached runs
    /// count the same. A compaction's shrink counts as zero, so the total
    /// never goes down; what the conversation grows by after it counts again.
    /// A turn that reports no usage (zero prompt and output) adds nothing and
    /// leaves the baseline where it was.
    Growth,
}

/// How many steered signatures to remember before dropping the oldest.
const MAX_STEERED_SIGNATURES: usize = 32;

/// Stable hash of a tool call's identity, for the steered-signature set.
///
/// Hashes the serialized `Value`, but only after `serde_json` has normalised
/// it — `to_string` on a `Value` emits object keys in the map's own order, so
/// two calls differing only in key order hash the same, matching how
/// `last_signature` compares `Value`s rather than text.
fn signature_hash(name: &str, args: &serde_json::Value) -> u64 {
    let mut h = fnv1a(name.as_bytes());
    h ^= fnv1a(args.to_string().as_bytes());
    h
}

/// What the loop should do about a repeated tool call.
///
/// Ordered by severity, and that ordering is load-bearing: one batch can
/// contain several repeated signatures, and the loop must act on the worst.
/// Declaration order is the severity order, so `Ord` is derived rather than
/// hand-written — a new escalation slots in by being declared in the right
/// place.
///
/// Deliberately **not** `#[non_exhaustive]`: this is control flow, so a new
/// variant should be a compile error for matchers, not a silent wildcard
/// fallthrough. Matches the `ToolDecision` precedent.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
#[must_use = "a loop verdict that is not acted on silently disables loop detection, \
              while still advancing the tracker's state"]
pub enum LoopVerdict {
    /// Nothing repetitive, or below the threshold.
    Continue,
    /// First trip: steer and keep going. A model that is stuck usually
    /// recovers when told so, and aborting here would regress the legitimate
    /// case of retrying a tool that failed transiently.
    Steer {
        tool_name: String,
        repetitions: usize,
    },
    /// The same signature tripped again after being steered about. Stop.
    Abort {
        tool_name: String,
        repetitions: usize,
    },
}

/// Tracks execution state against limits
pub struct ExecutionTracker {
    pub limits: ExecutionLimits,
    pub turns: usize,
    pub tokens_used: usize,
    pub started_at: web_time::Instant,
    /// Prompt plus output of the latest turn recorded with
    /// [`record_turn_usage`](Self::record_turn_usage) that reported usage,
    /// for [`TokenAccounting::Growth`]. Kept in both modes, so switching
    /// `limits.token_accounting` mid-run counts from the next turn.
    last_turn_size: Option<usize>,
    /// The tool signature seen most recently, and how many times in a row.
    ///
    /// A signature is `(name, arguments)` compared as `serde_json::Value`, not
    /// as serialized text — two calls that differ only in key order are the
    /// same call, and string comparison would miss the loop.
    last_signature: Option<(String, serde_json::Value)>,
    consecutive: usize,
    /// Hashes of signatures already steered about, capped. A second trip on one
    /// of these aborts.
    ///
    /// Hashes rather than the signatures themselves, and bounded rather than
    /// unbounded: this held a full clone of every steered call's arguments —
    /// for a file-write tool, the entire file body — inside the one type whose
    /// job is bounding runaway resource use. FNV for the same reason
    /// [`fnv1a`] is used elsewhere here: `DefaultHasher` is documented as
    /// unstable across Rust releases.
    ///
    /// A collision would abort a run one escalation early. At 64 bits, across
    /// the handful of distinct signatures a single run steers about, that is
    /// far less likely than the memory blowup it replaces.
    ///
    /// rather than steering again, so a model that ignores the nudge does not
    /// loop forever between nudges.
    steered: Vec<u64>,
}

impl ExecutionTracker {
    pub fn new(limits: ExecutionLimits) -> Self {
        Self {
            limits,
            turns: 0,
            tokens_used: 0,
            started_at: web_time::Instant::now(),
            last_turn_size: None,
            last_signature: None,
            consecutive: 0,
            steered: Vec::new(),
        }
    }

    /// Fold one turn's tool calls into the repetition counter.
    ///
    /// Counts within a batch as well as across turns: `ToolExecutionStrategy`
    /// defaults to `Parallel`, so a model can emit the same call three times in
    /// one message and never take a second turn.
    pub fn record_tool_calls(&mut self, calls: &[(String, serde_json::Value)]) -> LoopVerdict {
        let Some(threshold) = self.limits.max_consecutive_identical_tool_calls else {
            return LoopVerdict::Continue;
        };
        if threshold == 0 {
            return LoopVerdict::Continue;
        }

        let mut verdict = LoopVerdict::Continue;
        for (name, args) in calls {
            let sig = (name.clone(), args.clone());
            match &self.last_signature {
                Some(prev) if *prev == sig => self.consecutive += 1,
                _ => {
                    // Any different call breaks the streak — a model working
                    // through distinct steps is not looping.
                    self.last_signature = Some(sig.clone());
                    self.consecutive = 1;
                }
            }

            // No `verdict == Continue` guard: with `Parallel` execution a
            // model emits multi-call batches routinely, and stopping at the
            // first escalation both downgraded a later `Abort` to the earlier
            // `Steer` and skipped recording that later signature in `steered`,
            // giving it a free pass on the next turn too. Every call is
            // assessed; the most severe verdict wins via the fold below.
            if self.consecutive >= threshold {
                let repetitions = self.consecutive;
                let sig_hash = signature_hash(name, args);
                let this_call = if self.steered.contains(&sig_hash) {
                    LoopVerdict::Abort {
                        tool_name: name.clone(),
                        repetitions,
                    }
                } else {
                    // Bounded: a run that steers about many distinct
                    // signatures is already pathological, and the oldest
                    // entries are the least likely to recur.
                    if self.steered.len() >= MAX_STEERED_SIGNATURES {
                        self.steered.remove(0);
                    }
                    self.steered.push(sig_hash);
                    // Reset so the abort needs another full run of repeats,
                    // rather than firing on the very next call.
                    self.consecutive = 0;
                    LoopVerdict::Steer {
                        tool_name: name.clone(),
                        repetitions,
                    }
                };
                // Keep the worst verdict in the batch, not the first.
                if this_call > verdict {
                    verdict = this_call;
                }
            }
        }
        verdict
    }

    /// Count one turn that used `tokens`, added as is whatever the
    /// [`TokenAccounting`] (under [`TokenAccounting::Growth`] the caller
    /// supplies the growth). The loop records a model turn with
    /// [`record_turn_usage`](Self::record_turn_usage).
    pub fn record_turn(&mut self, tokens: usize) {
        self.turns += 1;
        self.tokens_used = self.tokens_used.saturating_add(tokens);
    }

    /// Count one model turn from its usage: `prompt` is the whole prompt
    /// (`input + cache_read + cache_write`), `output` what the model
    /// generated. What it adds to [`tokens_used`](Self::tokens_used) depends
    /// on [`ExecutionLimits::token_accounting`].
    pub fn record_turn_usage(&mut self, prompt: usize, output: usize) {
        self.turns += 1;
        let size = prompt.saturating_add(output);
        if size == 0 {
            // No usage reported (a gateway that omits it, a stream cut before
            // its usage): nothing to add, and a zero baseline would make
            // Growth count the whole conversation again on the next turn.
            return;
        }
        let added = match self.limits.token_accounting {
            TokenAccounting::Processed => size,
            TokenAccounting::Growth => match self.last_turn_size {
                Some(last) => size.saturating_sub(last),
                // The conversation the run starts with is the baseline.
                None => output,
            },
        };
        self.last_turn_size = Some(size);
        self.tokens_used = self.tokens_used.saturating_add(added);
    }

    /// Check if any limit has been exceeded. Returns the reason if so.
    pub fn check_limits(&self) -> Option<String> {
        if self.turns >= self.limits.max_turns {
            return Some(format!(
                "Max turns reached ({}/{})",
                self.turns, self.limits.max_turns
            ));
        }
        if self.tokens_used >= self.limits.max_total_tokens {
            return Some(format!(
                "Max tokens reached ({}/{})",
                self.tokens_used, self.limits.max_total_tokens
            ));
        }
        let elapsed = self.started_at.elapsed();
        if elapsed >= self.limits.max_duration {
            return Some(format!(
                "Max duration reached ({:.0}s/{:.0}s)",
                elapsed.as_secs_f64(),
                self.limits.max_duration.as_secs_f64()
            ));
        }
        None
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_estimate_tokens() {
        assert!(estimate_tokens("hello world") > 0);
        assert!(estimate_tokens("hello world") < 10);
        assert_eq!(estimate_tokens(""), 0);
    }

    #[test]
    fn extension_message_token_estimate_includes_message_overhead() {
        let extension = CustomMessage::new("status", serde_json::json!({"state": "ready"}));
        let expected = estimate_tokens(&extension.data.to_string()) + 4;
        assert_eq!(message_tokens(&AgentMessage::Custom(extension)), expected);
    }

    #[test]
    fn level2_summary_keeps_short_assistant_text() {
        let messages = vec![
            AgentMessage::Llm(Message::user("request")),
            AgentMessage::Llm(
                Message::assistant(
                    vec![Content::Text {
                        text: "the parser is ready".into(),
                    }],
                    StopReason::Stop,
                    "test",
                    "test",
                    Usage::default(),
                )
                .with_timestamp(7),
            ),
            AgentMessage::Llm(Message::user("recent")),
        ];

        let compacted = level2_summarize_old_turns(&messages, 1);
        assert!(matches!(
            compacted.get(1),
            Some(AgentMessage::Llm(Message::User { content, timestamp }))
                if *timestamp == 7
                    && content.iter().any(|content| matches!(
                        content,
                        Content::Text { text } if text == "[Summary] the parser is ready"
                    ))
        ));
        assert!(matches!(
            compacted.last(),
            Some(AgentMessage::Llm(Message::User { content, .. }))
                if content.iter().any(|content| matches!(
                    content,
                    Content::Text { text } if text == "recent"
                ))
        ));
    }

    #[test]
    fn test_context_config_from_context_window() {
        let config = ContextConfig::from_context_window(200_000);
        assert_eq!(config.max_context_tokens, 160_000); // 80% of 200K
        assert_eq!(config.system_prompt_tokens, 4_000); // default
        assert_eq!(config.keep_recent, 10); // default

        let config = ContextConfig::from_context_window(1_000_000);
        assert_eq!(config.max_context_tokens, 800_000); // 80% of 1M

        let config = ContextConfig::from_context_window(128_000);
        assert_eq!(config.max_context_tokens, 102_400); // 80% of 128K
    }

    #[test]
    fn test_truncate_head_tail() {
        let text = (1..=100)
            .map(|i| format!("line {}", i))
            .collect::<Vec<_>>()
            .join("\n");
        let result = truncate_text_head_tail(&text, 10, None).0;
        assert!(result.contains("line 1")); // head
        assert!(result.contains("line 100")); // tail
        assert!(result.contains("truncated"));
        assert!(!result.contains("line 50")); // middle removed

        // The marker is charged against the line budget, so the result fits
        // exactly and a second pass leaves it alone.
        assert_eq!(result.lines().count(), 10);
        assert_eq!(truncate_text_head_tail(&result, 10, None).0, result);
    }

    #[test]
    fn test_truncate_is_idempotent_and_keeps_an_honest_count() {
        let text = (1..=1000)
            .map(|i| format!("line {}", i))
            .collect::<Vec<_>>()
            .join("\n");

        let once = truncate_text_head_tail(&text, 50, None).0;
        assert!(once.contains("[... 953 lines truncated ...]"));
        assert_eq!(once.lines().count(), 50);

        // Re-truncating must not shrink the content further or restate the
        // omitted count — that churn rewrote settled history on every pass.
        let twice = truncate_text_head_tail(&once, 50, None).0;
        assert_eq!(twice, once);
        assert_eq!(truncate_text_head_tail(&twice, 50, None).0, once);
    }

    #[test]
    fn test_truncate_degenerate_line_budget() {
        let text = (1..=100)
            .map(|i| format!("line {}", i))
            .collect::<Vec<_>>()
            .join("\n");
        // Too small for head + marker + tail: keep the head, stay idempotent.
        for max_lines in [1usize, 2, 3, 4] {
            let result = truncate_text_head_tail(&text, max_lines, None).0;
            assert_eq!(result.lines().count(), max_lines);
            assert_eq!(truncate_text_head_tail(&result, max_lines, None).0, result);
        }
    }

    // -- helpers for the compaction-shape tests ---------------------------

    fn tool_turn(i: usize, output_lines: usize) -> Vec<AgentMessage> {
        let output = (0..output_lines)
            .map(|n| format!("turn {} line {} {}", i, n, "x".repeat(40)))
            .collect::<Vec<_>>()
            .join("\n");
        vec![
            AgentMessage::Llm(Message::Assistant {
                content: vec![Content::tool_call(
                    format!("tc-{}", i),
                    "bash",
                    serde_json::json!({ "command": format!("ls {}", i) }),
                )],
                stop_reason: StopReason::ToolUse,
                model: "test".into(),
                provider: "test".into(),
                usage: Usage::default(),
                timestamp: i as u64,
                error_message: None,
            }),
            AgentMessage::Llm(Message::ToolResult {
                tool_call_id: format!("tc-{}", i),
                tool_name: "bash".into(),
                content: vec![Content::Text { text: output }],
                is_error: false,
                timestamp: i as u64,
            }),
        ]
    }

    #[test]
    fn tool_boundaries_never_split_an_open_call_from_its_result() {
        let messages = tool_turn(0, 1);
        assert_eq!(safe_head_end(&messages, 1), 0);
        assert_eq!(safe_head_end(&messages, 2), 2);
        assert_eq!(safe_turn_start(&messages, 1), 0);
    }

    fn tool_session(turns: usize, output_lines: usize) -> Vec<AgentMessage> {
        let mut messages = vec![AgentMessage::Llm(Message::user("start"))];
        for i in 0..turns {
            messages.extend(tool_turn(i, output_lines));
        }
        messages
    }

    #[test]
    fn test_compact_target_ratio_leaves_headroom() {
        let messages = tool_session(40, 120);
        let config = ContextConfig {
            max_context_tokens: 8_000,
            system_prompt_tokens: 0,
            compact_target_ratio: 0.5,
            ..Default::default()
        };
        assert!(total_tokens(&messages) > config.max_context_tokens);

        let result = compact_messages(messages, &config);
        // Compacting to just-fits means the next turn compacts again; the
        // point of the ratio is that it lands well clear of the budget.
        assert!(
            total_tokens(&result) <= 4_000,
            "compacted to {} tokens, expected <= 4000 (50% of budget)",
            total_tokens(&result)
        );
    }

    #[test]
    fn test_headroom_policy_derives_the_ratio_from_growth() {
        let config = ContextConfig {
            max_context_tokens: 100_000,
            system_prompt_tokens: 0,
            compact_target_ratio: 0.7,
            compact_headroom_turns: Some(30),
            ..Default::default()
        };

        // 1000 tokens/turn * 30 turns = 30K headroom -> target 70K -> 0.70,
        // which ties the ratio; nothing changes.
        assert!((config.effective_target_ratio(1000.0) - 0.7).abs() < 1e-6);

        // Faster growth needs a deeper cut than the ratio would give.
        assert!((config.effective_target_ratio(2000.0) - 0.4).abs() < 1e-6);

        // Slow growth would allow keeping more, but the ratio is a ceiling.
        assert!((config.effective_target_ratio(100.0) - 0.7).abs() < 1e-6);

        // Runaway growth is floored so compaction cannot wipe history.
        assert_eq!(
            config.effective_target_ratio(1_000_000.0),
            MIN_HEADROOM_RATIO
        );
    }

    #[test]
    fn test_floor_overrides_a_lower_target_ratio_only_under_headroom() {
        let config = ContextConfig {
            max_context_tokens: 100_000,
            system_prompt_tokens: 0,
            compact_target_ratio: 0.2,
            compact_headroom_turns: Some(30),
            ..Default::default()
        };
        // Under the headroom policy a ratio below the floor is raised to it.
        assert_eq!(config.effective_target_ratio(5_000.0), MIN_HEADROOM_RATIO);
        assert_ne!(config.effective_target_ratio(5_000.0), 0.2);
        // Without the policy the configured ratio is used as-is, so the floor
        // never stops a caller who asks to compact harder.
        let unfloored = ContextConfig {
            compact_headroom_turns: None,
            ..config
        };
        assert_eq!(unfloored.effective_target_ratio(5_000.0), 0.2);
    }

    #[test]
    fn test_headroom_policy_is_inert_without_a_growth_estimate() {
        let config = ContextConfig {
            max_context_tokens: 100_000,
            system_prompt_tokens: 0,
            compact_target_ratio: 0.6,
            compact_headroom_turns: Some(30),
            ..Default::default()
        };
        // No growth measured yet, or a degenerate policy.
        assert_eq!(config.effective_target_ratio(0.0), 0.6);
        assert_eq!(config.effective_target_ratio(-5.0), 0.6);
        assert_eq!(config.effective_target_ratio(f64::NAN), 0.6);
        assert_eq!(
            ContextConfig {
                compact_headroom_turns: Some(0),
                ..config.clone()
            }
            .effective_target_ratio(1000.0),
            0.6
        );
        assert_eq!(
            ContextConfig {
                compact_headroom_turns: None,
                ..config
            }
            .effective_target_ratio(9999.0),
            0.6
        );
    }

    #[test]
    fn test_ratio_of_one_restores_compact_to_just_fit() {
        let config = ContextConfig {
            max_context_tokens: 10_000,
            system_prompt_tokens: 0,
            compact_target_ratio: 1.0,
            ..Default::default()
        };
        assert_eq!(config.compaction_target(10_000), 10_000);
    }

    #[test]
    fn test_compaction_target_rejects_nonsense_ratios() {
        let mut config = ContextConfig::default();
        for (ratio, expected) in [
            (0.0f32, 500usize),      // clamped up to the 0.05 floor
            (-1.0, 500),             // ditto
            (2.0, 10_000),           // clamped down to 1.0
            (f32::NAN, 10_000),      // non-finite falls back to 1.0
            (f32::INFINITY, 10_000), // ditto
        ] {
            config.compact_target_ratio = ratio;
            assert_eq!(
                config.compaction_target(10_000),
                expected,
                "ratio {}",
                ratio
            );
        }
    }

    #[test]
    fn test_level3_drops_only_what_the_target_requires() {
        // Budget comfortably holds most of the session, so a correct Level 3
        // sheds a few turns rather than collapsing to first + recent.
        let messages = tool_session(40, 20);
        let config = ContextConfig {
            max_context_tokens: total_tokens(&messages) * 3 / 4,
            system_prompt_tokens: 0,
            compact_target_ratio: 0.9,
            ..Default::default()
        };

        let result = compact_messages(messages.clone(), &config);
        assert!(total_tokens(&result) <= config.max_context_tokens);
        // The old fixed shape kept keep_first + marker + keep_recent = 13.
        assert!(
            result.len() > 13,
            "kept only {} messages; Level 3 is still collapsing history it did not need to",
            result.len()
        );
    }

    #[test]
    fn test_compaction_never_orphans_tool_calls() {
        // Every provider rejects a tool_result with no matching tool_use, and
        // a tool_use with no result. Compaction must not create either.
        for (max_context_tokens, keep_recent, keep_first) in [
            (400usize, 10usize, 2usize),
            (1_000, 10, 2),
            (4_000, 10, 2),
            (20_000, 10, 2),
            // Sweep the Level 2 / Level 3 boundaries across turn parities: a
            // boundary landing mid-turn is exactly how an orphan is created.
            (4_000, 9, 1),
            (4_000, 11, 3),
            (8_000, 7, 4),
            (8_000, 12, 5),
            (12_000, 13, 0),
        ] {
            let messages = tool_session(30, 60);
            let config = ContextConfig {
                max_context_tokens,
                system_prompt_tokens: 0,
                keep_recent,
                keep_first,
                ..Default::default()
            };
            let result = compact_messages(messages, &config);

            let mut open: Vec<String> = Vec::new();
            for msg in &result {
                match msg {
                    AgentMessage::Llm(Message::Assistant { content, .. }) => {
                        open = content
                            .iter()
                            .filter_map(|c| match c {
                                Content::ToolCall { id, .. } => Some(id.clone()),
                                _ => None,
                            })
                            .collect();
                    }
                    AgentMessage::Llm(Message::ToolResult { tool_call_id, .. }) => {
                        let answered = open.iter().position(|id| id == tool_call_id);
                        assert!(
                            answered.is_some(),
                            "orphaned tool result {} at budget {}",
                            tool_call_id,
                            max_context_tokens
                        );
                        open.remove(answered.unwrap());
                    }
                    _ => {
                        assert!(
                            open.is_empty(),
                            "unanswered tool call {:?} at budget {}",
                            open,
                            max_context_tokens
                        );
                    }
                }
            }
            assert!(
                open.is_empty(),
                "history ends on an unanswered tool call at budget {}",
                max_context_tokens
            );
        }
    }

    #[test]
    fn test_compaction_marker_carries_no_drifting_count() {
        // The marker sits near the front of the request. A message count in it
        // would change the cached prefix on every compaction pass.
        let messages = tool_session(40, 120);
        let config = ContextConfig {
            max_context_tokens: 2_000,
            system_prompt_tokens: 0,
            ..Default::default()
        };
        let result = compact_messages(messages, &config);

        let marker = result
            .iter()
            .find_map(|m| match m {
                AgentMessage::Llm(Message::User { content, .. }) => match content.first() {
                    Some(Content::Text { text }) if text.starts_with("[Context compacted") => {
                        Some(text.clone())
                    }
                    _ => None,
                },
                _ => None,
            })
            .expect("expected a compaction marker");
        assert_eq!(marker, COMPACTION_MARKER);
        assert!(!marker.chars().any(|c| c.is_ascii_digit()));
    }

    fn lines_block(prefix: &str, n: usize) -> Content {
        Content::Text {
            text: (0..n)
                .map(|i| format!("{prefix}{i}"))
                .collect::<Vec<_>>()
                .join("\n"),
        }
    }

    fn result_with(content: Vec<Content>) -> AgentMessage {
        AgentMessage::Llm(Message::ToolResult {
            tool_call_id: "tc-1".into(),
            tool_name: "mcp_tool".into(),
            content,
            is_error: false,
            timestamp: 0,
        })
    }

    fn content_of(msg: &AgentMessage) -> &[Content] {
        match msg {
            AgentMessage::Llm(Message::ToolResult { content, .. }) => content,
            _ => panic!("expected a tool result"),
        }
    }

    fn block_lines(msg: &AgentMessage) -> Vec<usize> {
        block_texts(msg)
            .into_iter()
            .map(|(_, t)| t.lines().count())
            .collect()
    }

    #[test]
    fn the_line_budget_is_the_result_s_not_each_block_s() {
        let config = ContextConfig {
            tool_output_max_lines: 200,
            ..Default::default()
        };
        // Twenty blocks of 200 lines: a per-block budget would keep 4,000.
        let many = result_with(
            (0..20)
                .map(|b| lines_block(&format!("b{b}-"), 200))
                .collect(),
        );
        let (once, marked) = truncate_tool_output_keyed(many, &config, None);
        assert_eq!(block_lines(&once).iter().sum::<usize>(), 200);
        assert_eq!(
            marked,
            (0..20).collect::<Vec<_>>(),
            "every cut block says so"
        );
        assert_eq!(
            truncate_tool_output(once.clone(), &config),
            once,
            "idempotent"
        );

        // Blocks that fit together are left alone, however many.
        let fits = result_with((0..4).map(|b| lines_block(&format!("b{b}-"), 50)).collect());
        assert_eq!(truncate_tool_output(fits.clone(), &config), fits);

        // A short block stays whole; the long ones share what is left, the
        // remainder to the earliest; an image keeps its place and its index.
        let mixed = result_with(vec![
            lines_block("a", 300),
            Content::Image {
                data: "x".into(),
                mime_type: "image/png".into(),
            },
            lines_block("s", 10),
            lines_block("c", 300),
        ]);
        let (once, marked) = truncate_tool_output_keyed(mixed, &config, Some("k"));
        assert_eq!(block_lines(&once), [95, 10, 95]);
        assert_eq!(marked, [0, 3]);
        let texts = block_texts(&once);
        assert_eq!(texts[1].1.lines().count(), 10);
        assert!(
            texts[2].1.contains("shared_state get \"k-b3\""),
            "{}",
            texts[2].1
        );
        assert!(matches!(&content_of(&once)[1], Content::Image { .. }));
        assert_eq!(
            truncate_tool_output(once.clone(), &config),
            once,
            "idempotent"
        );
    }

    #[test]
    fn the_shared_budget_s_remainder_and_freed_slack() {
        let config = ContextConfig {
            tool_output_max_lines: 200,
            ..Default::default()
        };
        // 200 over three long blocks: the remainder goes to the earliest.
        let three = result_with(
            (0..3)
                .map(|b| lines_block(&format!("b{b}-"), 300))
                .collect(),
        );
        assert_eq!(
            block_lines(&truncate_tool_output(three, &config)),
            [67, 67, 66]
        );
        // 60 lines is over an even share of 150 (50) but fits once the
        // 10-line block has left 140 for two (70): it stays whole.
        let config = ContextConfig {
            tool_output_max_lines: 150,
            ..Default::default()
        };
        let slack = result_with(vec![
            lines_block("s", 10),
            lines_block("m", 60),
            lines_block("l", 300),
        ]);
        assert_eq!(
            block_lines(&truncate_tool_output(slack, &config)),
            [10, 60, 80]
        );
        // A tool exempted by an override stays whole, however many blocks.
        let mut config = ContextConfig::default();
        config
            .tool_output_max_lines_overrides
            .insert("mcp_tool".into(), usize::MAX);
        let many = result_with(
            (0..20)
                .map(|b| lines_block(&format!("b{b}-"), 200))
                .collect(),
        );
        assert_eq!(truncate_tool_output(many.clone(), &config), many);
    }

    #[test]
    fn more_long_blocks_than_the_budget_can_mark_keep_a_marker_each() {
        let config = ContextConfig {
            tool_output_max_lines: 20,
            ..Default::default()
        };
        // Ten long blocks and room for two lines each: every one keeps the
        // fewest lines a marker needs rather than being cut silently.
        let msg = result_with(
            (0..10)
                .map(|b| lines_block(&format!("b{b}-"), 100))
                .collect(),
        );
        let (once, marked) = truncate_tool_output_keyed(msg, &config, None);
        // One head line, the marker's three, one tail line.
        assert_eq!(block_lines(&once), vec![5; 10]);
        assert_eq!(marked.len(), 10);
        assert_eq!(
            truncate_tool_output(once.clone(), &config),
            once,
            "idempotent"
        );
        // Short blocks past that point stay whole: sixty 4-line blocks keep
        // 240 lines of a 200-line budget, none cut.
        let config = ContextConfig {
            tool_output_max_lines: 200,
            ..Default::default()
        };
        let short = result_with((0..60).map(|b| lines_block(&format!("b{b}-"), 4)).collect());
        assert_eq!(truncate_tool_output(short.clone(), &config), short);
    }

    #[test]
    fn test_truncate_tool_output_helper() {
        let big = (0..500)
            .map(|i| format!("l{}", i))
            .collect::<Vec<_>>()
            .join("\n");
        let msg = AgentMessage::Llm(Message::ToolResult {
            tool_call_id: "tc-1".into(),
            tool_name: "bash".into(),
            content: vec![Content::Text { text: big }],
            is_error: false,
            timestamp: 7,
        });

        let config = ContextConfig {
            tool_output_max_lines: 50,
            ..Default::default()
        };
        let once = truncate_tool_output(msg, &config);
        let twice = truncate_tool_output(once.clone(), &config);
        assert_eq!(once, twice, "on-append truncation must be idempotent");

        match &once {
            AgentMessage::Llm(Message::ToolResult {
                content, timestamp, ..
            }) => {
                assert_eq!(*timestamp, 7, "timestamp must survive truncation");
                let Content::Text { text } = &content[0] else {
                    panic!("expected text")
                };
                assert_eq!(text.lines().count(), 50);
            }
            _ => panic!("expected a tool result"),
        }

        // Non-tool-result messages pass through untouched.
        let user = AgentMessage::Llm(Message::user("hello"));
        assert_eq!(truncate_tool_output(user.clone(), &config), user);
    }

    #[test]
    fn test_level1_truncation() {
        let big_output = (1..=200)
            .map(|i| format!("output line {}", i))
            .collect::<Vec<_>>()
            .join("\n");
        let messages = vec![
            AgentMessage::Llm(Message::user("do something")),
            AgentMessage::Llm(Message::ToolResult {
                tool_call_id: "tc-1".into(),
                tool_name: "bash".into(),
                content: vec![Content::Text { text: big_output }],
                is_error: false,
                timestamp: 0,
            }),
        ];

        let compacted = level1_truncate_tool_outputs(
            &messages,
            &ContextConfig {
                tool_output_max_lines: 20,
                ..Default::default()
            },
        );
        let tool_msg = &compacted[1];
        if let AgentMessage::Llm(Message::ToolResult { content, .. }) = tool_msg {
            if let Content::Text { text } = &content[0] {
                assert!(text.contains("truncated"));
                assert!(text.contains("output line 1")); // head
                assert!(text.contains("output line 200")); // tail
                assert!(text.lines().count() < 50);
            } else {
                panic!("expected text content");
            }
        } else {
            panic!("expected tool result");
        }
    }

    #[test]
    fn test_compact_within_budget() {
        let messages = vec![
            AgentMessage::Llm(Message::user("Hello")),
            AgentMessage::Llm(Message::user("World")),
        ];
        let config = ContextConfig::default();
        let result = compact_messages(messages.clone(), &config);
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_compact_drops_middle_when_needed() {
        let mut messages = Vec::new();
        for i in 0..100 {
            messages.push(AgentMessage::Llm(Message::user(format!(
                "Message {} {}",
                i,
                "x".repeat(200)
            ))));
        }

        let config = ContextConfig {
            max_context_tokens: 500,
            system_prompt_tokens: 100,
            keep_recent: 5,
            keep_first: 2,
            tool_output_max_lines: 20,
            ..Default::default()
        };

        let result = compact_messages(messages, &config);
        assert!(result.len() < 100);
        assert!(result.len() >= 2);
    }

    #[test]
    fn test_context_tracker_no_usage() {
        let tracker = ContextTracker::new();
        let messages = vec![
            AgentMessage::Llm(Message::user("Hello")),
            AgentMessage::Llm(Message::user("World")),
        ];
        // Without usage data, falls back to estimation
        let tokens = tracker.estimate_context_tokens(&messages);
        assert!(tokens > 0);
        assert_eq!(tokens, total_tokens(&messages));
    }

    #[test]
    fn test_context_tracker_with_usage() {
        let mut tracker = ContextTracker::new();
        let messages = vec![
            AgentMessage::Llm(Message::user("Hello")),
            AgentMessage::Llm(Message::Assistant {
                content: vec![Content::Text {
                    text: "Hi there!".into(),
                }],
                stop_reason: StopReason::Stop,
                model: "test".into(),
                provider: "test".into(),
                usage: Usage {
                    input: 100,
                    output: 50,
                    ..Default::default()
                },
                timestamp: 0,
                error_message: None,
            }),
            AgentMessage::Llm(Message::user("Follow up question here")),
        ];
        // Record usage at index 1 (assistant message)
        tracker.record_usage(
            &Usage {
                input: 100,
                output: 50,
                ..Default::default()
            },
            1,
        );
        let tokens = tracker.estimate_context_tokens(&messages);
        // Should be 150 (real usage) + estimate for the trailing user message
        let trailing_estimate = message_tokens(&messages[2]);
        assert_eq!(tokens, 150 + trailing_estimate);
    }

    #[test]
    fn test_context_tracker_reset() {
        let mut tracker = ContextTracker::new();
        tracker.record_usage(
            &Usage {
                input: 1000,
                output: 500,
                ..Default::default()
            },
            0,
        );
        tracker.reset();
        // Index 0 is inside this history, so only the reset makes the
        // tracker fall back (an index past the end falls back regardless).
        let messages = vec![AgentMessage::Llm(Message::user("test"))];
        // After reset, should fall back to estimation
        assert_eq!(
            tracker.estimate_context_tokens(&messages),
            total_tokens(&messages)
        );
    }

    #[test]
    fn tracker_counts_cached_prompt_tokens() {
        let mut tracker = ContextTracker::new();
        let messages = vec![
            AgentMessage::Llm(Message::user("Hello")),
            AgentMessage::Llm(Message::user("World")),
        ];
        // A cached prompt is mostly `cache_read`; every bucket is context.
        tracker.record_usage(
            &Usage {
                input: 100,
                output: 10,
                cache_read: 1_000,
                cache_write: 200,
                ..Default::default()
            },
            1,
        );
        assert_eq!(tracker.estimate_context_tokens(&messages), 1_310);
    }

    #[test]
    fn tracker_ignores_a_response_without_usage() {
        let mut tracker = ContextTracker::new();
        let messages = vec![
            AgentMessage::Llm(Message::user("Hello")),
            AgentMessage::Llm(Message::user("World")),
        ];
        // A provider that reports no usage must not zero the estimate.
        tracker.record_usage(&Usage::default(), 0);
        assert_eq!(
            tracker.estimate_context_tokens(&messages),
            total_tokens(&messages)
        );
    }

    #[test]
    fn tracker_falls_back_when_its_index_is_past_the_history() {
        let mut tracker = ContextTracker::new();
        tracker.record_usage(
            &Usage {
                input: 1_000,
                ..Default::default()
            },
            3,
        );
        // Compaction shrank the history below the recorded index, and the
        // caller did not reset: estimate, never slice past the end.
        let messages = vec![
            AgentMessage::Llm(Message::user("a")),
            AgentMessage::Llm(Message::user("b")),
            AgentMessage::Llm(Message::user("c")),
        ];
        assert_eq!(
            tracker.estimate_context_tokens(&messages),
            total_tokens(&messages)
        );
    }

    #[test]
    fn token_estimates_of_each_content_kind() {
        let one = |content: Content| {
            message_tokens(&AgentMessage::Llm(Message::User {
                content: vec![content],
                timestamp: 0,
            }))
        };
        // 400,000 base64 chars = 300,000 bytes, at 750 bytes a token.
        let image = Content::Image {
            data: "A".repeat(400_000),
            mime_type: "image/png".into(),
        };
        assert_eq!(one(image), 400 + 4);
        // 1,600 base64 chars = 1,200 bytes, at 4 bytes a token.
        let redacted = Content::thinking_redacted(
            crate::provider::ApiProtocol::AnthropicMessages,
            "A".repeat(1_600),
        );
        assert_eq!(one(redacted), 300 + 4);
        // Name (1) + arguments `{}` (1) + 8 of framing.
        let call = Content::tool_call("tc-1", "read", serde_json::json!({}));
        assert_eq!(one(call), 10 + 4);
        // Text (1) + tool name (1) + 8 of framing.
        let result = AgentMessage::Llm(Message::ToolResult {
            tool_call_id: "tc-1".into(),
            tool_name: "bash".into(),
            content: vec![Content::Text {
                text: "abcd".into(),
            }],
            is_error: false,
            timestamp: 0,
        });
        assert_eq!(message_tokens(&result), 10);
    }

    #[test]
    fn the_headroom_floor_is_inert_without_a_growth_estimate() {
        // Below the floor, an inert policy keeps the configured ratio rather
        // than raising it to the floor.
        let config = ContextConfig {
            max_context_tokens: 100_000,
            system_prompt_tokens: 0,
            compact_target_ratio: 0.2,
            compact_headroom_turns: Some(30),
            ..Default::default()
        };
        assert_eq!(config.effective_target_ratio(f64::NAN), 0.2);
        assert_eq!(
            ContextConfig {
                compact_headroom_turns: Some(0),
                ..config
            }
            .effective_target_ratio(1000.0),
            0.2
        );
    }

    #[test]
    fn message_text_joins_the_text_blocks_of_a_tool_result() {
        let msg = result_with(vec![
            Content::Text { text: "a".into() },
            Content::Image {
                data: "x".into(),
                mime_type: "image/png".into(),
            },
            Content::Text { text: "b".into() },
        ]);
        assert_eq!(message_text(&msg), "a\nb");
        assert_eq!(message_text(&AgentMessage::Llm(Message::user("u"))), "");
    }

    #[test]
    fn tool_output_keys_hash_with_fnv1a() {
        // Published FNV-1a 64 vectors: a stashed key must not move between
        // releases, or the markers already in a history stop resolving.
        assert_eq!(tool_output_key("tc", ""), "tool-out-tc-cbf29ce484222325");
        assert_eq!(tool_output_key("tc", "a"), "tool-out-tc-af63dc4c8601ec8c");
        assert_eq!(
            tool_output_key("tc", "foobar"),
            "tool-out-tc-85944171f73967e8"
        );
    }

    fn assistant_with(content: Vec<Content>) -> AgentMessage {
        AgentMessage::Llm(Message::assistant(
            content,
            StopReason::Stop,
            "test",
            "test",
            Usage::default(),
        ))
    }

    fn user_texts(messages: &[AgentMessage]) -> Vec<String> {
        messages
            .iter()
            .map(|m| match m {
                AgentMessage::Llm(Message::User { content, .. }) => content
                    .iter()
                    .filter_map(|c| match c {
                        Content::Text { text } => Some(text.as_str()),
                        _ => None,
                    })
                    .collect::<String>(),
                other => panic!("expected only user messages, got {other:?}"),
            })
            .collect()
    }

    #[test]
    fn level2_summarizes_each_old_turn_by_what_it_said_or_did() {
        let call = |id: &str| Content::tool_call(id, "bash", serde_json::json!({}));
        let text = |n: usize| Content::Text {
            text: "t".repeat(n),
        };
        let messages = vec![
            AgentMessage::Llm(Message::user("task")),
            assistant_with(vec![
                Content::Text {
                    text: "short".into(),
                },
                call("a"),
            ]),
            result_with(vec![text(1)]),
            // 200 chars is short enough to keep.
            assistant_with(vec![text(200)]),
            // Longer text is dropped: what is left is the tool count...
            assistant_with(vec![text(201), call("b"), call("c")]),
            result_with(vec![text(1)]),
            result_with(vec![text(1)]),
            // ...or, with no tools either, a placeholder.
            assistant_with(vec![text(201)]),
            AgentMessage::Llm(Message::user("recent")),
        ];
        let compacted = level2_summarize_old_turns(&messages, 1);
        assert_eq!(
            user_texts(&compacted),
            [
                "task".to_string(),
                "[Summary] short".into(),
                format!("[Summary] {}", "t".repeat(200)),
                "[Summary] [Assistant used 2 tool(s)]".into(),
                "[Summary] [Assistant response]".into(),
                "recent".into(),
            ]
        );
    }

    #[test]
    fn level2_with_nothing_kept_recent_summarizes_every_turn() {
        let messages = vec![
            AgentMessage::Llm(Message::user("task")),
            assistant_with(vec![Content::tool_call("a", "bash", serde_json::json!({}))]),
            result_with(vec![Content::Text { text: "out".into() }]),
        ];
        // The history ends on the turn's tool result.
        let compacted = level2_summarize_old_turns(&messages, 0);
        assert_eq!(
            user_texts(&compacted),
            ["task", "[Summary] [Assistant used 1 tool(s)]"]
        );
    }

    #[test]
    fn a_history_opening_on_tool_results_stops_at_its_start() {
        // Orphaned results at the front (a history loaded from elsewhere):
        // walking back to the turn start stops at index 0.
        let messages = vec![
            result_with(vec![Content::Text { text: "a".into() }]),
            result_with(vec![Content::Text { text: "b".into() }]),
            AgentMessage::Llm(Message::user("next")),
        ];
        assert_eq!(safe_turn_start(&messages, 1), 0);
    }

    fn users(n: usize, tokens: usize) -> Vec<AgentMessage> {
        (0..n)
            .map(|i| {
                AgentMessage::Llm(
                    Message::user(format!("{i:02}{}", "x".repeat(tokens * 4 - 2)))
                        .with_timestamp(100 + i as u64),
                )
            })
            .collect()
    }

    #[test]
    fn a_history_inside_keep_first_and_keep_recent_still_compacts() {
        // Six messages, fewer than `keep_first + keep_recent`, over budget.
        let config = ContextConfig {
            max_context_tokens: 2_000,
            system_prompt_tokens: 0,
            keep_first: 1,
            keep_recent: 10,
            compact_target_ratio: 0.5,
            ..Default::default()
        };
        let messages = users(6, 500);
        let out = compact_messages(messages.clone(), &config);
        assert!(total_tokens(&out) <= 2_000, "{}", total_tokens(&out));
        // The task, the marker, the latest ask.
        assert_eq!(out.len(), 3);
        assert_eq!(out[0], messages[0]);
        assert_eq!(out[2], messages[5]);
    }

    #[test]
    fn level3_with_no_head_kept_opens_on_the_marker() {
        let config = ContextConfig {
            max_context_tokens: 5_000,
            system_prompt_tokens: 0,
            keep_first: 0,
            keep_recent: 2,
            compact_target_ratio: 0.5,
            ..Default::default()
        };
        let messages = users(20, 500);
        let out = compact_messages(messages.clone(), &config);
        assert!(total_tokens(&out) <= 2_500, "{}", total_tokens(&out));
        // The marker, then the newest messages that fit the target.
        assert_eq!(out.len(), 5);
        assert_eq!(out[1..], messages[16..]);
        // It takes the timestamp of the first message it replaces.
        assert_eq!(out[0], compaction_marker(100));
    }

    #[test]
    fn test_execution_limits() {
        let limits = ExecutionLimits {
            max_turns: 3,
            max_total_tokens: 1000,
            max_duration: std::time::Duration::from_secs(60),
            ..Default::default()
        };

        let mut tracker = ExecutionTracker::new(limits);
        assert!(tracker.check_limits().is_none());

        tracker.record_turn(100);
        tracker.record_turn(100);
        assert!(tracker.check_limits().is_none());

        tracker.record_turn(100);
        assert!(tracker.check_limits().is_some());
    }

    /// Five turns of a conversation that grows by 10K tokens a turn (tool
    /// results in, a 1K answer out), from a 10K prompt.
    fn five_growing_turns(tracker: &mut ExecutionTracker) {
        for turn in 0..5 {
            tracker.record_turn_usage(10_000 + turn * 10_000, 1_000);
        }
    }

    #[test]
    fn processed_accounting_counts_every_prompt() {
        let mut tracker = ExecutionTracker::new(ExecutionLimits::default());
        five_growing_turns(&mut tracker);
        // 10K + 20K + … + 50K of prompts, plus 5 × 1K of output.
        assert_eq!(tracker.tokens_used, 155_000);
        assert_eq!(tracker.turns, 5);
    }

    #[test]
    fn growth_accounting_counts_the_conversation_once() {
        let limits = ExecutionLimits::default().with_token_accounting(TokenAccounting::Growth);
        let mut tracker = ExecutionTracker::new(limits);
        five_growing_turns(&mut tracker);
        // What the run added: from the first 10K prompt to the last 50K
        // prompt plus its 1K answer. Each answer is in the next prompt, so it
        // is not counted twice.
        assert_eq!(tracker.tokens_used, 41_000);
        assert_eq!(tracker.turns, 5);
    }

    #[test]
    fn growth_accounting_does_not_count_the_history_a_run_starts_with() {
        let limits = ExecutionLimits::default().with_token_accounting(TokenAccounting::Growth);
        let mut tracker = ExecutionTracker::new(limits);
        // A run late in a 900K conversation: only its own growth counts.
        tracker.record_turn_usage(900_000, 2_000);
        assert_eq!(tracker.tokens_used, 2_000);
        tracker.record_turn_usage(910_000, 2_000);
        assert_eq!(tracker.tokens_used, 12_000);
    }

    #[test]
    fn a_turn_without_usage_adds_nothing_and_keeps_the_baseline() {
        for (accounting, expected) in [
            (TokenAccounting::Processed, 81_000 + 83_000),
            (TokenAccounting::Growth, 1_000 + 2_000),
        ] {
            let limits = ExecutionLimits::default().with_token_accounting(accounting);
            let mut tracker = ExecutionTracker::new(limits);
            tracker.record_turn_usage(80_000, 1_000);
            tracker.record_turn_usage(0, 0);
            tracker.record_turn_usage(82_000, 1_000);
            assert_eq!(tracker.tokens_used, expected, "{accounting:?}");
            assert_eq!(tracker.turns, 3);
        }
    }

    #[test]
    fn growth_accounting_never_goes_down_and_counts_regrowth() {
        let limits = ExecutionLimits::default().with_token_accounting(TokenAccounting::Growth);
        let mut tracker = ExecutionTracker::new(limits);
        tracker.record_turn_usage(10_000, 1_000);
        tracker.record_turn_usage(80_000, 1_000);
        assert_eq!(tracker.tokens_used, 71_000);
        // Compacted: the next prompt is smaller. Nothing is added, nothing
        // taken back.
        tracker.record_turn_usage(20_000, 1_000);
        assert_eq!(tracker.tokens_used, 71_000);
        // Growth after the compaction counts again.
        tracker.record_turn_usage(30_000, 1_000);
        assert_eq!(tracker.tokens_used, 81_000);
    }

    #[test]
    fn growth_accounting_stops_a_run_at_the_cap() {
        let limits = ExecutionLimits::default()
            .with_max_total_tokens(40_000)
            .with_token_accounting(TokenAccounting::Growth);
        let mut tracker = ExecutionTracker::new(limits);
        tracker.record_turn_usage(10_000, 1_000);
        tracker.record_turn_usage(30_000, 1_000);
        assert!(tracker.check_limits().is_none());
        tracker.record_turn_usage(50_000, 1_000);
        assert_eq!(
            tracker.check_limits().as_deref(),
            Some("Max tokens reached (41000/40000)")
        );
    }

    #[test]
    fn limits_without_token_accounting_deserialize_to_processed() {
        let limits: ExecutionLimits = serde_json::from_value(serde_json::json!({
            "max_turns": 5,
            "max_total_tokens": 100,
            "max_duration": {"secs": 1, "nanos": 0},
            "max_consecutive_identical_tool_calls": 3
        }))
        .unwrap();
        assert_eq!(limits.token_accounting, TokenAccounting::Processed);
        let growth = ExecutionLimits::default().with_token_accounting(TokenAccounting::Growth);
        let stored = serde_json::to_value(&growth).unwrap();
        assert_eq!(stored["token_accounting"], "growth");
        let back: ExecutionLimits = serde_json::from_value(stored).unwrap();
        assert_eq!(back.token_accounting, TokenAccounting::Growth);
    }
}

#[cfg(test)]
mod compaction_retention {
    use super::*;
    use crate::types::{Content, StopReason, Usage};

    fn assistant_call(i: usize) -> AgentMessage {
        AgentMessage::Llm(Message::Assistant {
            content: vec![Content::ToolCall {
                id: format!("tc-{i}"),
                name: "fetch_record".into(),
                arguments: serde_json::json!({"n": i}),
                provider_metadata: None,
            }],
            stop_reason: StopReason::ToolUse,
            model: "m".into(),
            provider: "p".into(),
            usage: Usage::default(),
            timestamp: i as u64,
            error_message: None,
        })
    }

    fn tool_result(i: usize) -> AgentMessage {
        AgentMessage::Llm(Message::ToolResult {
            tool_call_id: format!("tc-{i}"),
            tool_name: "fetch_record".into(),
            content: vec![Content::Text {
                text: "record line, nominal, no action required ".repeat(400),
            }],
            is_error: false,
            timestamp: i as u64,
        })
    }

    /// Compaction must leave a transcript a provider will accept, and must
    /// leave the agent something to work with.
    ///
    /// Written after a live long-horizon run reported a compaction collapsing
    /// 22 messages / 21371 tokens to **1 message / 22 tokens** — far below
    /// `keep_first + keep_recent`. That did not reproduce here (retention is
    /// 15-24 messages on the same shape), so this pins the two properties that
    /// matter, in-crate where the live path cannot be reached: no orphaned tool
    /// call, and never emptied to nothing.
    #[test]
    fn compaction_keeps_a_usable_well_formed_transcript() {
        let cfg = ContextConfig {
            max_context_tokens: 30_000,
            keep_recent: 6,
            keep_first: 2,
            ..Default::default()
        };
        for pairs in [6usize, 11, 12, 20, 40] {
            let mut msgs = vec![AgentMessage::Llm(Message::user("go"))];
            for i in 0..pairs {
                msgs.push(assistant_call(i));
                msgs.push(tool_result(i));
            }
            let before = msgs.len();
            let out = compact_messages(msgs, &cfg);

            assert!(
                !out.is_empty(),
                "compaction of {before} messages emptied the transcript"
            );

            // No `tool_use` may be left unanswered — the shape every provider
            // rejects, and the one thing compaction must never produce.
            let mut pending: Vec<String> = Vec::new();
            for m in &out {
                let AgentMessage::Llm(msg) = m else { continue };
                match msg {
                    Message::Assistant { content, .. } => {
                        pending = content
                            .iter()
                            .filter_map(|c| match c {
                                Content::ToolCall { id, .. } => Some(id.clone()),
                                _ => None,
                            })
                            .collect();
                    }
                    Message::ToolResult { tool_call_id, .. } => {
                        if let Some(at) = pending.iter().position(|p| p == tool_call_id) {
                            pending.remove(at);
                        }
                    }
                    Message::User { .. } => {}
                }
            }
            assert!(
                pending.is_empty(),
                "compaction of {before} messages orphaned tool calls {pending:?}"
            );
        }
    }

    /// Every tool result answers a call of the assistant message before it,
    /// and every call is answered before anything else follows.
    fn assert_well_formed(out: &[AgentMessage], ctx: &str) {
        let mut pending: Vec<String> = Vec::new();
        for m in out {
            let AgentMessage::Llm(msg) = m else { continue };
            match msg {
                Message::Assistant { content, .. } => {
                    assert!(pending.is_empty(), "{ctx}: unanswered calls {pending:?}");
                    pending = content
                        .iter()
                        .filter_map(|c| match c {
                            Content::ToolCall { id, .. } => Some(id.clone()),
                            _ => None,
                        })
                        .collect();
                }
                Message::ToolResult { tool_call_id, .. } => {
                    let at = pending.iter().position(|p| p == tool_call_id);
                    assert!(at.is_some(), "{ctx}: orphaned result {tool_call_id}");
                    pending.remove(at.unwrap());
                }
                Message::User { .. } => {
                    assert!(pending.is_empty(), "{ctx}: unanswered calls {pending:?}");
                }
            }
        }
        assert!(pending.is_empty(), "{ctx}: unanswered calls {pending:?}");
    }

    /// Random sessions — multi-call turns, text-only turns, asks, summary
    /// decoys, results from tiny to bulky — under small budgets and every
    /// `keep_first` / `keep_recent`: compaction stays well formed and within
    /// the budget, never leaves only the marker, and — whenever the task, the
    /// latest ask and the newest turn fit the budget together — keeps all
    /// three (#150).
    #[test]
    fn compaction_is_well_formed_within_budget_and_keeps_what_fits() {
        let mut state = 0x9E37_79B9_7F4A_7C15u64;
        let mut next = move |n: u64| {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            state % n
        };
        for case in 0..3_000 {
            let mut msgs = vec![AgentMessage::Llm(Message::user(format!(
                "TASK: {}",
                "t".repeat(4 * (1 + next(400) as usize))
            )))];
            let mut call = 0;
            for turn in 0..(2 + next(24)) {
                let calls = next(4) as usize;
                let mut content = Vec::new();
                if calls == 0 || next(3) == 0 {
                    content.push(Content::Text {
                        text: "a".repeat(4 * (1 + next(800) as usize)),
                    });
                }
                let ids: Vec<usize> = (0..calls).map(|k| call + k).collect();
                call += calls;
                for &id in &ids {
                    if let AgentMessage::Llm(Message::Assistant { content: c, .. }) =
                        assistant_call(id)
                    {
                        content.extend(c);
                    }
                }
                msgs.push(AgentMessage::Llm(Message::Assistant {
                    content,
                    stop_reason: StopReason::ToolUse,
                    model: "m".into(),
                    provider: "p".into(),
                    usage: Usage::default(),
                    timestamp: turn,
                    error_message: None,
                }));
                for id in ids {
                    let tokens = [5, 100, 1_000, 4_000][next(4) as usize];
                    msgs.push(result_of(id, tokens));
                }
                match next(6) {
                    0 => msgs.push(AgentMessage::Llm(Message::user(format!("ASK {turn}")))),
                    1 => msgs.push(AgentMessage::Llm(Message::user(format!(
                        "{SUMMARY_PREFIX}{}",
                        "s".repeat(4 * 60)
                    )))),
                    _ => {}
                }
            }
            let cfg = ContextConfig {
                max_context_tokens: [800, 2_000, 6_000, 20_000][next(4) as usize],
                system_prompt_tokens: 0,
                keep_first: next(5) as usize,
                keep_recent: next(13) as usize,
                compact_target_ratio: [0.15, 0.3, 0.5, 0.7][next(4) as usize],
                ..Default::default()
            };
            let ctx = format!(
                "case {case} (kf {} kr {} ratio {} budget {})",
                cfg.keep_first, cfg.keep_recent, cfg.compact_target_ratio, cfg.max_context_tokens
            );
            // What must survive when it fits: the task (in the head), the
            // latest ask, and the newest turn — measured on what level 3
            // receives (levels 1 and 2 run first, as in `compact_messages`).
            let l1 = level1_truncate_tool_outputs(&msgs, &cfg);
            let l3_input = if total_tokens(&l1) <= cfg.max_context_tokens {
                l1
            } else {
                level2_summarize_old_turns(&l1, cfg.keep_recent)
            };
            let len = l3_input.len();
            let newest = l3_input
                .iter()
                .rposition(|m| matches!(m, AgentMessage::Llm(Message::Assistant { .. })))
                .filter(|&i| i > 0);
            let ask = latest_ask(&l3_input, 1, len);
            // The `keep_first` head goes or stays as a block.
            let head_end = safe_head_end(&l3_input, cfg.keep_first.min(len));
            let protected = l3_input[..head_end]
                .iter()
                .map(message_tokens)
                .sum::<usize>()
                + ask.map_or(0, |i| message_tokens(&l3_input[i]))
                + message_tokens(&compaction_marker(0))
                + newest.map_or(0, |n| {
                    l3_input[n..].iter().map(message_tokens).sum::<usize>()
                });
            let must_keep = head_end >= 1 && protected <= cfg.max_context_tokens;
            let newest_ids: Vec<String> = l3_input[newest.map_or(len, |n| n + 1)..]
                .iter()
                .filter_map(|m| match m {
                    AgentMessage::Llm(Message::ToolResult { tool_call_id, .. }) => {
                        Some(tool_call_id.clone())
                    }
                    _ => None,
                })
                .collect();
            let ask_text = ask.map(|i| message_text(&l3_input[i]));
            let out = compact_messages(msgs, &cfg);
            assert!(!out.is_empty(), "{ctx}: emptied");
            assert_well_formed(&out, &ctx);
            assert!(
                total_tokens(&out) <= cfg.max_context_tokens,
                "{ctx}: {} tokens",
                total_tokens(&out)
            );
            assert!(out.len() > 1, "{ctx}: only the marker survived");
            if must_keep {
                assert!(has_text(&out, "TASK:"), "{ctx}: the task was dropped");
                if let Some(ask) = &ask_text {
                    assert!(has_text(&out, ask), "{ctx}: the ask was dropped");
                }
                for id in &newest_ids {
                    let kept = out.iter().any(|m| {
                        matches!(m, AgentMessage::Llm(Message::ToolResult { tool_call_id, .. }) if tool_call_id == id)
                    });
                    assert!(kept, "{ctx}: the newest turn's result {id} was dropped");
                }
            }
        }
    }

    fn result_of(i: usize, tokens: usize) -> AgentMessage {
        AgentMessage::Llm(Message::ToolResult {
            tool_call_id: format!("tc-{i}"),
            tool_name: "fetch_record".into(),
            content: vec![Content::Text {
                text: "r".repeat(tokens * 4),
            }],
            is_error: false,
            timestamp: i as u64,
        })
    }

    fn has_text(out: &[AgentMessage], needle: &str) -> bool {
        out.iter().any(|m| match m {
            AgentMessage::Llm(Message::User { content, .. }) => content
                .iter()
                .any(|c| matches!(c, Content::Text { text } if text.contains(needle))),
            _ => false,
        })
    }

    fn has_result(out: &[AgentMessage], i: usize) -> bool {
        let id = format!("tc-{i}");
        out.iter().any(
            |m| matches!(m, AgentMessage::Llm(Message::ToolResult { tool_call_id, .. }) if *tool_call_id == id),
        )
    }

    /// The task prompt, an earlier ask, then `pairs` calls whose results
    /// weigh `tokens` each.
    fn session(pairs: usize, tokens: usize) -> Vec<AgentMessage> {
        let mut msgs = vec![AgentMessage::Llm(Message::user("TASK: audit the records"))];
        for i in 0..pairs {
            msgs.push(assistant_call(i));
            msgs.push(result_of(i, tokens));
            if i == pairs / 2 {
                msgs.push(AgentMessage::Llm(Message::user(
                    "ASK: only the failed ones",
                )));
            }
        }
        msgs
    }

    fn small_budget() -> ContextConfig {
        ContextConfig {
            max_context_tokens: 20_000,
            system_prompt_tokens: 0,
            keep_recent: 10,
            keep_first: 2,
            // A 6K target: the 10-message tail (5 x 1.5K) alone exceeds it.
            compact_target_ratio: 0.3,
            ..Default::default()
        }
    }

    /// When the head and the `keep_recent` tail exceed the target, the task
    /// prompt, the latest ask and the newest turn stay — keeping only the
    /// newest messages that fit the target used to drop the task first (#150).
    #[test]
    fn a_tail_over_the_target_keeps_the_task_the_ask_and_the_newest_turn() {
        let cfg = small_budget();
        let out = compact_messages(session(14, 1_500), &cfg);
        assert!(total_tokens(&out) <= 6_000, "{}", total_tokens(&out));
        assert!(has_text(&out, "TASK:"), "the task prompt was dropped");
        assert!(has_text(&out, "failed ones"), "the latest ask was dropped");
        assert!(has_result(&out, 13), "the newest turn was dropped");
        assert!(has_text(&out, COMPACTION_MARKER));
    }

    /// A newest turn bigger than the target is kept, with the task and the
    /// ask, up to the budget — it used to leave only the marker (#150's
    /// `22 -> 1` collapse).
    #[test]
    fn a_newest_turn_over_the_target_is_kept_up_to_the_budget() {
        let cfg = small_budget();
        let mut msgs = session(14, 1_500);
        msgs.push(assistant_call(14));
        msgs.push(result_of(14, 9_000));
        let out = compact_messages(msgs, &cfg);
        let size = total_tokens(&out);
        assert!(size > 6_000 && size <= 20_000, "{size}");
        assert!(has_text(&out, "TASK:") && has_text(&out, "failed ones"));
        assert!(has_result(&out, 14), "the newest turn was dropped");
        assert!(out.len() > 1, "only the marker survived");
    }

    /// Past the budget the head goes first; the ask and the newest turn stay —
    /// the ask even when it was part of the head (`keep_first: 2`).
    #[test]
    fn over_the_budget_the_head_goes_before_the_ask() {
        let cfg = small_budget();
        let mut msgs = vec![AgentMessage::Llm(Message::user(format!(
            "TASK: {}",
            "t".repeat(4 * 3_000)
        )))];
        msgs.push(AgentMessage::Llm(Message::user(
            "ASK: only the failed ones",
        )));
        for i in 0..4 {
            msgs.push(assistant_call(i));
            msgs.push(result_of(i, 1_000));
        }
        // The bulk is the assistant's own text, which level 1 does not cut.
        let mut newest = assistant_call(4);
        if let AgentMessage::Llm(Message::Assistant { content, .. }) = &mut newest {
            content.insert(
                0,
                Content::Text {
                    text: "a".repeat(4 * 17_000),
                },
            );
        }
        msgs.push(newest);
        msgs.push(result_of(4, 100));
        let out = compact_messages(msgs, &cfg);
        assert!(total_tokens(&out) <= 20_000, "{}", total_tokens(&out));
        assert!(!has_text(&out, "TASK:"));
        assert!(has_text(&out, "failed ones"), "the ask was dropped");
        assert!(has_result(&out, 4), "the newest turn was dropped");
    }

    /// A newest turn bigger than the budget cannot be kept; the marker and the
    /// ask are, so the model still knows what it was asked.
    #[test]
    fn a_newest_turn_over_the_budget_leaves_the_marker_and_the_ask() {
        let cfg = small_budget();
        let mut newest = assistant_call(0);
        if let AgentMessage::Llm(Message::Assistant { content, .. }) = &mut newest {
            content.insert(
                0,
                Content::Text {
                    text: "a".repeat(4 * 25_000),
                },
            );
        }
        let msgs = vec![
            AgentMessage::Llm(Message::user("TASK: audit the records")),
            AgentMessage::Llm(Message::user("ASK: only the failed ones")),
            newest,
            result_of(0, 100),
        ];
        let out = compact_messages(msgs, &cfg);
        assert_eq!(out.len(), 2, "{out:?}");
        assert!(has_text(&out[..1], COMPACTION_MARKER));
        assert!(has_text(&out[1..], "failed ones"));
    }

    /// A user message after the newest assistant turn is the latest ask: an
    /// older one is not brought back after the marker.
    #[test]
    fn an_ask_after_the_newest_turn_supersedes_an_older_one() {
        let cfg = small_budget();
        let mut msgs = session(14, 1_500);
        msgs.push(AgentMessage::Llm(Message::assistant(
            vec![Content::Text {
                text: "Done: 3 records failed.".into(),
            }],
            StopReason::Stop,
            "m",
            "p",
            Usage::default(),
        )));
        msgs.push(AgentMessage::Llm(Message::user("NEW: now fix them")));
        let out = compact_messages(msgs, &cfg);
        assert!(has_text(&out, "TASK:") && has_text(&out, "NEW: now fix them"));
        assert!(
            !has_text(&out, "failed ones"),
            "a superseded ask was brought back"
        );
    }

    /// With `keep_recent: 0` level 2 summarizes every assistant message, so
    /// level 3 sees no turn to protect: it still cuts to the target and keeps
    /// the task and the ask.
    #[test]
    fn with_no_assistant_left_level_3_still_cuts_to_the_target() {
        for keep_recent in [0, 1] {
            let cfg = ContextConfig {
                keep_recent,
                ..small_budget()
            };
            let mut msgs = vec![AgentMessage::Llm(Message::user("TASK: audit the records"))];
            for i in 0..14 {
                msgs.push(assistant_call(i));
                msgs.push(result_of(i, 200));
                msgs.push(AgentMessage::Llm(Message::user(format!(
                    "ASK {i}: {}",
                    "u".repeat(4 * 1_500)
                ))));
            }
            let out = compact_messages(msgs, &cfg);
            assert!(
                total_tokens(&out) <= 6_000,
                "kr {keep_recent}: {}",
                total_tokens(&out)
            );
            assert!(
                has_text(&out, "TASK:"),
                "kr {keep_recent}: the task was dropped"
            );
            assert!(
                has_text(&out, "ASK 13:"),
                "kr {keep_recent}: the ask was dropped"
            );
        }
    }

    /// Compacting again after more turns keeps the bytes of what was kept at
    /// the front — the task and the marker — so the prefix cache survives the
    /// second compaction up to there; the ask is still kept.
    #[test]
    fn a_second_compaction_keeps_the_front_byte_for_byte() {
        let cfg = small_budget();
        let first = compact_messages(session(14, 1_500), &cfg);
        let mut next = first.clone();
        for i in 14..20 {
            next.push(assistant_call(i));
            next.push(result_of(i, 1_500));
        }
        let second = compact_messages(next, &cfg);
        let json = |m: &[AgentMessage]| serde_json::to_string(m).unwrap();
        assert_eq!(json(&second[..2]), json(&first[..2]));
        assert!(has_text(&second, "failed ones"), "the ask was dropped");
        assert!(has_result(&second, 19));
    }

    /// A cut that passes the latest ask keeps it, after the marker — not
    /// one the loop injected (a summary).
    #[test]
    fn a_cut_past_the_latest_ask_keeps_it() {
        let cfg = ContextConfig {
            max_context_tokens: 8_000,
            system_prompt_tokens: 0,
            keep_recent: 4,
            keep_first: 1,
            compact_target_ratio: 0.7,
            ..Default::default()
        };
        let mut msgs = vec![
            AgentMessage::Llm(Message::user("TASK: audit the records")),
            AgentMessage::Llm(Message::user("ASK: only the failed ones")),
        ];
        for _ in 0..20 {
            msgs.push(AgentMessage::Llm(Message::user(format!(
                "{SUMMARY_PREFIX}{}",
                "s".repeat(4 * 500)
            ))));
        }
        for i in 0..2 {
            msgs.push(assistant_call(i));
            msgs.push(result_of(i, 100));
        }
        let out = compact_messages(msgs, &cfg);
        assert!(total_tokens(&out) <= 5_600, "{}", total_tokens(&out));
        assert!(has_text(&out, COMPACTION_MARKER), "nothing was cut");
        assert!(has_text(&out, "TASK:") && has_text(&out, "failed ones"));
        // The ask comes right after the marker, before the kept tail.
        let at = |needle: &str| {
            out.iter()
                .position(|m| has_text(std::slice::from_ref(m), needle))
        };
        assert_eq!(at("failed ones"), at(COMPACTION_MARKER).map(|i| i + 1));
    }
}