tower_llm/auto_compaction/

mod.rs

//! Auto-compaction layer for managing conversation context length
//!
//! This module provides automatic conversation compaction when context limits are reached.
//! It can work reactively (on error) or proactively (when approaching limits).
//!
//! # Features
//! - Reactive compaction on context length errors
//! - Proactive compaction based on token thresholds
//! - Configurable compaction strategies
//! - Custom compaction prompts
//! - Integration with existing provider infrastructure
//!
//! # Example
//!
//! ```no_run
//! use tower_llm::{Agent, auto_compaction::{CompactionPolicy, CompactionStrategy, ProactiveThreshold}};
//! use std::sync::Arc;
//! use async_openai::{Client, config::OpenAIConfig};
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
//! let client = Arc::new(Client::<OpenAIConfig>::new());
//!
//! // Configure compaction to trigger at 1000 tokens
//! let policy = CompactionPolicy {
//!     compaction_model: "gpt-4o-mini".to_string(),
//!     proactive_threshold: Some(ProactiveThreshold {
//!         token_threshold: 1000,
//!         percentage_threshold: None,
//!     }),
//!     compaction_strategy: CompactionStrategy::PreserveSystemAndRecent { recent_count: 5 },
//!     ..Default::default()
//! };
//!
//! let agent = Agent::builder(client)
//!     .model("gpt-4o")
//!     .auto_compaction(policy)
//!     .build();
//! # Ok(())
//! # }
//! ```

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use async_openai::types::{
    ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs,
    ChatCompletionRequestToolMessageArgs, ChatCompletionRequestUserMessageArgs,
    CreateChatCompletionRequest, CreateChatCompletionRequestArgs,
};
use tower::{BoxError, Layer, Service, ServiceExt};

use crate::core::StepOutcome;
use crate::provider::ModelService;

// ===== Core Types =====

/// Strategy for handling orphaned tool calls during compaction
#[derive(Clone, Debug)]
pub enum OrphanedToolCallStrategy {
    /// Drop orphaned tool calls before compaction and re-append after (default)
    DropAndReappend,
    /// Exclude messages with orphaned tool calls from compaction range
    ExcludeFromCompaction,
    /// Add placeholder tool responses for orphaned calls
    AddPlaceholderResponses,
    /// Fail compaction if orphaned tool calls are detected
    FailOnOrphaned,
}

impl Default for OrphanedToolCallStrategy {
    fn default() -> Self {
        Self::DropAndReappend
    }
}

/// Configuration for auto-compaction behavior
#[derive(Clone)]
pub struct CompactionPolicy {
    /// Model to use for compaction (e.g., "gpt-4o-mini")
    pub compaction_model: String,

    /// Optional: Proactive threshold for triggering compaction
    pub proactive_threshold: Option<ProactiveThreshold>,

    /// Strategy for selecting messages to compact
    pub compaction_strategy: CompactionStrategy,

    /// Prompt to use for compaction
    pub compaction_prompt: CompactionPrompt,

    /// Maximum retries for compaction attempts
    pub max_compaction_attempts: usize,

    /// Strategy for handling orphaned tool calls
    pub orphaned_tool_call_strategy: OrphanedToolCallStrategy,
}

impl Default for CompactionPolicy {
    fn default() -> Self {
        Self {
            compaction_model: "gpt-4o-mini".to_string(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::PreserveSystemAndRecent { recent_count: 10 },
            compaction_prompt: CompactionPrompt::Default,
            max_compaction_attempts: 2,
            orphaned_tool_call_strategy: OrphanedToolCallStrategy::default(),
        }
    }
}

/// Threshold configuration for proactive compaction
#[derive(Clone, Debug)]
pub struct ProactiveThreshold {
    /// Token count at which to trigger compaction
    pub token_threshold: usize,

    /// Alternative: percentage of model's context window (0.0-1.0)
    pub percentage_threshold: Option<f32>,
}

/// Strategy for selecting which messages to compact
#[derive(Clone)]
pub enum CompactionStrategy {
    /// Compact all messages except the last N
    CompactAllButLast(usize),

    /// Compact messages older than N turns (user+assistant pairs)
    CompactOlderThan(usize),

    /// Keep system prompt and last N messages, compact the middle
    PreserveSystemAndRecent { recent_count: usize },

    /// Custom function to select which messages to compact
    Custom(CompactionRangeFn),
}

/// Range of messages to compact
#[derive(Clone, Debug)]
pub struct CompactionRange {
    /// Start index (inclusive)
    pub start: usize,
    /// End index (exclusive)
    pub end: usize,
}

/// Type alias for custom compaction range function
pub type CompactionRangeFn =
    Arc<dyn Fn(&[ChatCompletionRequestMessage]) -> CompactionRange + Send + Sync>;

/// Type alias for dynamic prompt generation function
pub type PromptGeneratorFn = Arc<dyn Fn(&[ChatCompletionRequestMessage]) -> String + Send + Sync>;

/// Compaction prompt configuration
#[derive(Clone)]
pub enum CompactionPrompt {
    /// Use default prompt optimized for conversation summarization
    Default,

    /// Custom static prompt
    Custom(String),

    /// Dynamic prompt based on messages being compacted
    Dynamic(PromptGeneratorFn),
}

impl CompactionPrompt {
    fn generate(&self, messages: &[ChatCompletionRequestMessage]) -> String {
        match self {
            CompactionPrompt::Default => {
                "Please provide a concise summary of the conversation above, preserving:\n\
                1. The user's original intent and requirements\n\
                2. Key decisions, conclusions, and action items\n\
                3. Important context, constraints, and technical details\n\
                4. Any errors or issues encountered and their resolutions\n\
                5. Current state of any ongoing tasks\n\n\
                Format the summary as a clear narrative that maintains conversation continuity."
                    .to_string()
            }
            CompactionPrompt::Custom(prompt) => prompt.clone(),
            CompactionPrompt::Dynamic(f) => f(messages),
        }
    }
}

/// Token counter trait for estimating message tokens
pub trait TokenCounter: Send + Sync {
    fn count_messages(&self, messages: &[ChatCompletionRequestMessage]) -> usize;
}

/// Simple token counter that estimates based on character count
#[derive(Clone)]
pub struct SimpleTokenCounter {
    chars_per_token: f32,
}

impl SimpleTokenCounter {
    pub fn new() -> Self {
        Self {
            chars_per_token: 4.0, // Rough estimate
        }
    }
}

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

impl TokenCounter for SimpleTokenCounter {
    fn count_messages(&self, messages: &[ChatCompletionRequestMessage]) -> usize {
        let total_chars: usize = messages
            .iter()
            .map(|msg| {
                // Estimate character count based on message type
                match msg {
                    ChatCompletionRequestMessage::System(m) => {
                        match &m.content {
                            async_openai::types::ChatCompletionRequestSystemMessageContent::Text(t) => t.len(),
                            async_openai::types::ChatCompletionRequestSystemMessageContent::Array(_) => 100, // Rough estimate
                        }
                    }
                    ChatCompletionRequestMessage::User(m) => {
                        match &m.content {
                            async_openai::types::ChatCompletionRequestUserMessageContent::Text(t) => t.len(),
                            async_openai::types::ChatCompletionRequestUserMessageContent::Array(_) => 100,
                        }
                    }
                    ChatCompletionRequestMessage::Assistant(m) => {
                        // Assistant messages have optional content
                        m.content.as_ref().map(|c| match c {
                            async_openai::types::ChatCompletionRequestAssistantMessageContent::Text(t) => t.len(),
                            async_openai::types::ChatCompletionRequestAssistantMessageContent::Array(_) => 100,
                        }).unwrap_or(0)
                    }
                    ChatCompletionRequestMessage::Tool(m) => {
                        // Tool messages have content enum
                        match &m.content {
                            async_openai::types::ChatCompletionRequestToolMessageContent::Text(t) => t.len(),
                            async_openai::types::ChatCompletionRequestToolMessageContent::Array(_) => 100,
                        }
                    }
                    ChatCompletionRequestMessage::Function(m) => {
                        m.content.as_ref().map(|c| c.len()).unwrap_or(0)
                    }
                    ChatCompletionRequestMessage::Developer(m) => {
                        match &m.content {
                            async_openai::types::ChatCompletionRequestDeveloperMessageContent::Text(t) => t.len(),
                            async_openai::types::ChatCompletionRequestDeveloperMessageContent::Array(_) => 100,
                        }
                    }
                }
            })
            .sum();

        (total_chars as f32 / self.chars_per_token) as usize
    }
}

// ===== Layer Implementation =====

/// Layer that adds auto-compaction to a Step service
pub struct AutoCompactionLayer<P, C> {
    policy: CompactionPolicy,
    provider: P,
    token_counter: Arc<C>,
}

impl<P, C> AutoCompactionLayer<P, C> {
    pub fn new(policy: CompactionPolicy, provider: P, token_counter: C) -> Self {
        Self {
            policy,
            provider,
            token_counter: Arc::new(token_counter),
        }
    }
}

/// Service wrapper that performs auto-compaction
pub struct AutoCompaction<S, P, C> {
    pub(crate) inner: Arc<tokio::sync::Mutex<S>>,
    pub(crate) policy: CompactionPolicy,
    pub(crate) provider: Arc<tokio::sync::Mutex<P>>,
    pub(crate) token_counter: Arc<C>,
}

impl<S, P, C> Layer<S> for AutoCompactionLayer<P, C>
where
    P: Clone,
    C: Clone,
{
    type Service = AutoCompaction<S, P, C>;

    fn layer(&self, inner: S) -> Self::Service {
        AutoCompaction {
            inner: Arc::new(tokio::sync::Mutex::new(inner)),
            policy: self.policy.clone(),
            provider: Arc::new(tokio::sync::Mutex::new(self.provider.clone())),
            token_counter: self.token_counter.clone(),
        }
    }
}

impl<S, P, C> AutoCompaction<S, P, C>
where
    P: ModelService + Send + 'static,
    P::Future: Send + 'static,
    C: TokenCounter + 'static,
{
    /// Find all orphaned tool calls in the conversation
    fn find_orphaned_tool_calls(messages: &[ChatCompletionRequestMessage]) -> Vec<usize> {
        let mut orphaned_indices = Vec::new();

        for (i, msg) in messages.iter().enumerate() {
            if let ChatCompletionRequestMessage::Assistant(asst) = msg {
                if let Some(tool_calls) = &asst.tool_calls {
                    if !tool_calls.is_empty() {
                        // Check if all tool calls have corresponding responses
                        let mut all_calls_have_responses = true;
                        for tool_call in tool_calls {
                            let mut found_response = false;
                            // Look for a tool response with matching ID after this message
                            for subsequent_msg in messages.iter().skip(i + 1) {
                                if let ChatCompletionRequestMessage::Tool(tool_msg) = subsequent_msg
                                {
                                    if tool_msg.tool_call_id == tool_call.id {
                                        found_response = true;
                                        break;
                                    }
                                }
                                // Stop looking if we hit another assistant message (conversation moved on)
                                if matches!(
                                    subsequent_msg,
                                    ChatCompletionRequestMessage::Assistant(_)
                                ) {
                                    break;
                                }
                            }
                            if !found_response {
                                all_calls_have_responses = false;
                                break;
                            }
                        }

                        if !all_calls_have_responses {
                            orphaned_indices.push(i);
                        }
                    }
                }
            }
        }

        orphaned_indices
    }

    /// Check if messages have orphaned tool calls (tool calls without responses)
    fn has_orphaned_tool_calls(messages: &[ChatCompletionRequestMessage]) -> bool {
        !Self::find_orphaned_tool_calls(messages).is_empty()
    }

    /// Remove all orphaned tool call messages and return them separately
    fn extract_all_orphaned_tool_calls(
        messages: Vec<ChatCompletionRequestMessage>,
    ) -> (
        Vec<ChatCompletionRequestMessage>,
        Vec<(usize, ChatCompletionRequestMessage)>,
    ) {
        let orphaned_indices = Self::find_orphaned_tool_calls(&messages);

        if orphaned_indices.is_empty() {
            return (messages, Vec::new());
        }

        let mut cleaned = Vec::new();
        let mut orphaned = Vec::new();

        for (i, msg) in messages.into_iter().enumerate() {
            if orphaned_indices.contains(&i) {
                orphaned.push((i, msg));
            } else {
                cleaned.push(msg);
            }
        }

        (cleaned, orphaned)
    }

    /// Compact messages based on the configured strategy
    pub(crate) async fn compact_messages(
        &self,
        messages: Vec<ChatCompletionRequestMessage>,
    ) -> Result<Vec<ChatCompletionRequestMessage>, BoxError> {
        // Handle orphaned tool calls based on strategy
        let (messages_to_compact, orphaned_tool_calls) = if Self::has_orphaned_tool_calls(&messages)
        {
            match self.policy.orphaned_tool_call_strategy {
                OrphanedToolCallStrategy::DropAndReappend => {
                    let orphaned_count = Self::find_orphaned_tool_calls(&messages).len();
                    tracing::debug!("Detected {} orphaned tool call message(s), dropping and will re-append after compaction", orphaned_count);
                    Self::extract_all_orphaned_tool_calls(messages)
                }
                OrphanedToolCallStrategy::ExcludeFromCompaction => {
                    tracing::debug!("Detected orphaned tool calls, excluding from compaction");
                    // For this strategy, we keep orphaned messages but exclude them from compaction
                    // This is handled by keeping all messages and adjusting the range
                    (messages, Vec::new())
                }
                OrphanedToolCallStrategy::AddPlaceholderResponses => {
                    tracing::debug!("Detected orphaned tool calls, adding placeholder responses");
                    // Add placeholder tool responses for all orphaned tool calls
                    let mut fixed_messages = messages.clone();
                    let orphaned_indices = Self::find_orphaned_tool_calls(&messages);

                    // Process in reverse order to maintain indices
                    for &idx in orphaned_indices.iter().rev() {
                        if let ChatCompletionRequestMessage::Assistant(asst) = &messages[idx] {
                            if let Some(tool_calls) = &asst.tool_calls {
                                // Insert placeholder responses right after the assistant message
                                let mut placeholders = Vec::new();
                                for tool_call in tool_calls {
                                    let placeholder =
                                        ChatCompletionRequestToolMessageArgs::default()
                                            .content("[Pending tool execution]")
                                            .tool_call_id(&tool_call.id)
                                            .build()?;
                                    placeholders.push(placeholder.into());
                                }
                                // Insert all placeholders after the assistant message
                                for (i, placeholder) in placeholders.into_iter().enumerate() {
                                    fixed_messages.insert(idx + 1 + i, placeholder);
                                }
                            }
                        }
                    }
                    (fixed_messages, Vec::new())
                }
                OrphanedToolCallStrategy::FailOnOrphaned => {
                    let orphaned_count = Self::find_orphaned_tool_calls(&messages).len();
                    return Err(format!(
                        "Cannot compact: conversation has {} orphaned tool call message(s)",
                        orphaned_count
                    )
                    .into());
                }
            }
        } else {
            (messages, Vec::new())
        };

        // Determine which messages to compact based on strategy
        let range = match &self.policy.compaction_strategy {
            CompactionStrategy::CompactAllButLast(n) => {
                let start = 0;
                let end = messages_to_compact.len().saturating_sub(*n);
                CompactionRange { start, end }
            }
            CompactionStrategy::CompactOlderThan(turns) => {
                // Count back N user+assistant pairs
                let mut turn_count = 0;
                let mut cutoff = messages_to_compact.len();
                for (i, msg) in messages_to_compact.iter().enumerate().rev() {
                    if matches!(msg, ChatCompletionRequestMessage::User(_)) {
                        turn_count += 1;
                        if turn_count >= *turns {
                            cutoff = i;
                            break;
                        }
                    }
                }
                CompactionRange {
                    start: 0,
                    end: cutoff,
                }
            }
            CompactionStrategy::PreserveSystemAndRecent { recent_count } => {
                // Find first non-system message
                let start = messages_to_compact
                    .iter()
                    .position(|m| !matches!(m, ChatCompletionRequestMessage::System(_)))
                    .unwrap_or(0);
                let end = messages_to_compact.len().saturating_sub(*recent_count);
                CompactionRange {
                    start,
                    end: end.max(start),
                }
            }
            CompactionStrategy::Custom(f) => f(&messages_to_compact),
        };

        // If nothing to compact, return original with orphaned tool calls if any
        if range.start >= range.end {
            let mut result = messages_to_compact;
            if !orphaned_tool_calls.is_empty() {
                // Re-append orphaned messages
                for (_, orphaned_msg) in orphaned_tool_calls {
                    result.push(orphaned_msg);
                }
            }
            return Ok(result);
        }

        // Extract messages to compact, extending end to include trailing tool responses
        let mut adjusted_end = range.end;
        if adjusted_end > range.start {
            if let ChatCompletionRequestMessage::Assistant(asst) =
                &messages_to_compact[adjusted_end - 1]
            {
                if let Some(tool_calls) = &asst.tool_calls {
                    let ids: std::collections::HashSet<&str> =
                        tool_calls.iter().map(|tc| tc.id.as_str()).collect();
                    while adjusted_end < messages_to_compact.len() {
                        if let ChatCompletionRequestMessage::Tool(t) =
                            &messages_to_compact[adjusted_end]
                        {
                            if ids.contains(t.tool_call_id.as_str()) {
                                adjusted_end += 1;
                                continue;
                            }
                        }
                        break;
                    }
                }
            }
        }
        let to_compact = &messages_to_compact[range.start..adjusted_end];
        let prompt = self.policy.compaction_prompt.generate(to_compact);

        // Build compaction request
        let mut builder = CreateChatCompletionRequestArgs::default();
        builder.model(&self.policy.compaction_model);

        // Start with a system prompt
        let mut compact_messages = vec![ChatCompletionRequestSystemMessageArgs::default()
            .content(prompt)
            .build()?
            .into()];

        // Identify the last assistant-with-tool_calls inside the slice
        let mut last_asst_rel_idx: Option<usize> = None;
        for (i, msg) in to_compact.iter().enumerate() {
            if let ChatCompletionRequestMessage::Assistant(asst) = msg {
                if asst
                    .tool_calls
                    .as_ref()
                    .map(|v| !v.is_empty())
                    .unwrap_or(false)
                {
                    last_asst_rel_idx = Some(i);
                }
            }
        }

        if let Some(rel_idx) = last_asst_rel_idx {
            // Gather tool ids from the last assistant
            let abs_last = range.start + rel_idx;
            let mut remaining_ids: std::collections::HashSet<String> =
                if let ChatCompletionRequestMessage::Assistant(asst) =
                    &messages_to_compact[abs_last]
                {
                    asst.tool_calls
                        .as_ref()
                        .map(|tc| tc.iter().map(|t| t.id.clone()).collect())
                        .unwrap_or_else(std::collections::HashSet::new)
                } else {
                    std::collections::HashSet::new()
                };

            // Phase 1: push messages up to and including the assistant (skip leading system)
            let mut pushed_non_system = false;
            for msg in &to_compact[..=rel_idx] {
                if !pushed_non_system {
                    if matches!(msg, ChatCompletionRequestMessage::System(_)) {
                        continue;
                    } else {
                        pushed_non_system = true;
                    }
                }
                compact_messages.push(msg.clone());
            }

            // Phase 2: scan forward beyond the slice for matching tool responses, skipping non-tool messages
            let mut extra_tools: Vec<ChatCompletionRequestMessage> = Vec::new();
            if !remaining_ids.is_empty() {
                let mut idx = abs_last + 1;
                while idx < messages_to_compact.len() && !remaining_ids.is_empty() {
                    match &messages_to_compact[idx] {
                        ChatCompletionRequestMessage::Tool(t) => {
                            if remaining_ids.contains(&t.tool_call_id) {
                                extra_tools.push(messages_to_compact[idx].clone());
                                remaining_ids.remove(&t.tool_call_id);
                            }
                        }
                        ChatCompletionRequestMessage::Assistant(_) => {
                            break; // stop at next assistant
                        }
                        _ => {
                            // skip non-tool messages
                        }
                    }
                    idx += 1;
                }
            }

            // Append extra tool responses (contiguously after the assistant)
            compact_messages.extend(extra_tools);

            // Add instruction to summarize if needed
            let last_is_user = matches!(
                compact_messages.last(),
                Some(ChatCompletionRequestMessage::User(_))
            );
            if !last_is_user {
                compact_messages.push(
                    ChatCompletionRequestUserMessageArgs::default()
                        .content("Please provide a summary of the above conversation following the instructions.")
                        .build()?
                        .into(),
                );
            }

            builder.messages(compact_messages);
        } else {
            // Fallback: no assistant with tool_calls inside the slice; include the slice as-is (skipping leading system)
            let mut pushed_non_system = false;
            for msg in to_compact {
                if !pushed_non_system {
                    if matches!(msg, ChatCompletionRequestMessage::System(_)) {
                        continue;
                    } else {
                        pushed_non_system = true;
                    }
                }
                compact_messages.push(msg.clone());
            }

            let last_is_user = matches!(
                compact_messages.last(),
                Some(ChatCompletionRequestMessage::User(_))
            );
            if !last_is_user {
                compact_messages.push(
                    ChatCompletionRequestUserMessageArgs::default()
                        .content("Please provide a summary of the above conversation following the instructions.")
                        .build()?
                        .into(),
                );
            }

            builder.messages(compact_messages);
        }
        let compact_req = builder.build()?;

        // Call compaction model
        let mut provider = self.provider.lock().await;
        let response = ServiceExt::ready(&mut *provider)
            .await?
            .call(compact_req)
            .await?;
        drop(provider);

        // Build new message list with compacted summary
        let mut result = Vec::new();

        // Keep messages before the compacted range
        for msg in &messages_to_compact[..range.start] {
            result.push(msg.clone());
        }

        // Add the compacted summary as a user message to maintain role validity
        if let Some(summary) = response.assistant.content {
            result.push(
                ChatCompletionRequestUserMessageArgs::default()
                    .content(format!("[Previous conversation summary]: {}", summary))
                    .build()?
                    .into(),
            );
        }

        // Keep messages after the compacted range
        for msg in &messages_to_compact[adjusted_end..] {
            result.push(msg.clone());
        }

        // Re-append orphaned tool calls if we dropped them earlier
        if !orphaned_tool_calls.is_empty() {
            tracing::debug!(
                "Re-appending {} orphaned tool call message(s) after compaction",
                orphaned_tool_calls.len()
            );

            // Sort by original index to maintain relative order
            let mut sorted_orphaned = orphaned_tool_calls;
            sorted_orphaned.sort_by_key(|(idx, _)| *idx);

            // For simplicity, append all orphaned messages at the end
            // This maintains the conversation flow while ensuring all tool calls are preserved
            for (_, orphaned_msg) in sorted_orphaned {
                result.push(orphaned_msg);
            }
        }

        // Normalize: drop empty assistant messages and coalesce adjacent assistants
        let mut normalized: Vec<ChatCompletionRequestMessage> = Vec::with_capacity(result.len());
        for msg in result.into_iter() {
            if let Some(ChatCompletionRequestMessage::Assistant(prev)) = normalized.last_mut() {
                if let ChatCompletionRequestMessage::Assistant(curr) = &msg {
                    let curr_empty = matches!(&curr.content, Some(async_openai::types::ChatCompletionRequestAssistantMessageContent::Text(t)) if t.is_empty());
                    let no_calls = prev
                        .tool_calls
                        .as_ref()
                        .map(|v| v.is_empty())
                        .unwrap_or(true)
                        && curr
                            .tool_calls
                            .as_ref()
                            .map(|v| v.is_empty())
                            .unwrap_or(true);
                    if curr_empty && no_calls {
                        // skip pushing empty assistant
                        continue;
                    }
                    if no_calls {
                        let prev_text = match &prev.content {
                            Some(async_openai::types::ChatCompletionRequestAssistantMessageContent::Text(t)) => Some(t.clone()),
                            _ => None,
                        };
                        let curr_text = match &curr.content {
                            Some(async_openai::types::ChatCompletionRequestAssistantMessageContent::Text(t)) => Some(t.clone()),
                            _ => None,
                        };
                        if let (Some(p), Some(c)) = (prev_text, curr_text) {
                            let combined = if p.is_empty() {
                                c
                            } else if c.is_empty() {
                                p
                            } else {
                                format!("{}\n\n{}", p, c)
                            };
                            prev.content = Some(async_openai::types::ChatCompletionRequestAssistantMessageContent::Text(combined));
                            continue;
                        }
                    }
                }
            }

            // default push (dropping empty assistants)
            if let ChatCompletionRequestMessage::Assistant(a) = &msg {
                let is_empty = matches!(&a.content, Some(async_openai::types::ChatCompletionRequestAssistantMessageContent::Text(t)) if t.is_empty())
                    && a.tool_calls.as_ref().map(|v| v.is_empty()).unwrap_or(true);
                if is_empty {
                    continue;
                }
            }
            normalized.push(msg);
        }

        Ok(normalized)
    }

    /// Check if the error is a context length error
    fn is_context_length_error(error: &BoxError) -> bool {
        let error_str = error.to_string().to_lowercase();
        error_str.contains("context_length_exceeded")
            || error_str.contains("context length")
            || error_str.contains("maximum context")
            || error_str.contains("token limit")
    }
}

impl<S, P, C> Service<CreateChatCompletionRequest> for AutoCompaction<S, P, C>
where
    S: Service<CreateChatCompletionRequest, Response = StepOutcome, Error = BoxError>
        + Send
        + 'static,
    S::Future: Send + 'static,
    P: ModelService + Send + 'static,
    P::Future: Send + 'static,
    C: TokenCounter + 'static,
{
    type Response = StepOutcome;
    type Error = BoxError;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(
        &mut self,
        _cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        std::task::Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: CreateChatCompletionRequest) -> Self::Future {
        let inner = self.inner.clone();
        let policy = self.policy.clone();
        let provider = self.provider.clone();
        let token_counter = self.token_counter.clone();

        Box::pin(async move {
            let mut current_messages = req.messages.clone();
            let mut attempts = 0;

            // Check for proactive compaction
            if let Some(threshold) = &policy.proactive_threshold {
                let token_count = token_counter.count_messages(&current_messages);

                if token_count > threshold.token_threshold {
                    // Proactively compact
                    let compactor = AutoCompaction {
                        inner: inner.clone(),
                        policy: policy.clone(),
                        provider: provider.clone(),
                        token_counter: token_counter.clone(),
                    };

                    match compactor.compact_messages(current_messages.clone()).await {
                        Ok(compacted) => {
                            current_messages = compacted;
                        }
                        Err(e) => {
                            // Log but don't fail - continue with original messages
                            tracing::warn!("Proactive compaction failed: {}", e);
                        }
                    }
                }
            }

            loop {
                // Build request with current messages
                let mut builder = CreateChatCompletionRequestArgs::default();
                builder.model(&req.model);
                builder.messages(current_messages.clone());
                if let Some(t) = req.temperature {
                    builder.temperature(t);
                }
                #[allow(deprecated)]
                if let Some(mt) = req.max_tokens {
                    builder.max_tokens(mt);
                }
                if let Some(mct) = req.max_completion_tokens {
                    builder.max_completion_tokens(mct);
                }
                if let Some(tools) = req.tools.clone() {
                    builder.tools(tools);
                }
                let current_req = builder.build()?;

                // Try to call inner service
                let mut guard = inner.lock().await;
                let result = ServiceExt::ready(&mut *guard)
                    .await?
                    .call(current_req)
                    .await;
                drop(guard);

                match result {
                    Ok(outcome) => return Ok(outcome),
                    Err(e) => {
                        // Check if it's a context length error
                        if Self::is_context_length_error(&e)
                            && attempts < policy.max_compaction_attempts
                        {
                            attempts += 1;

                            // Attempt reactive compaction
                            let compactor = AutoCompaction {
                                inner: inner.clone(),
                                policy: policy.clone(),
                                provider: provider.clone(),
                                token_counter: token_counter.clone(),
                            };

                            match compactor.compact_messages(current_messages.clone()).await {
                                Ok(compacted) => {
                                    tracing::info!(
                                        "Reactive compaction successful, retrying request"
                                    );
                                    current_messages = compacted;
                                    // Loop will retry with compacted messages
                                }
                                Err(compact_err) => {
                                    return Err(format!(
                                        "Context length error and compaction failed: original={}, compact={}",
                                        e, compact_err
                                    ).into());
                                }
                            }
                        } else {
                            return Err(e);
                        }
                    }
                }
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::validation::gen;
    use crate::validation::{validate_conversation, ValidationPolicy};
    use async_openai::types::ChatCompletionRequestAssistantMessageArgs;
    use proptest::prelude::*;
    use std::sync::Arc as StdArc;
    use tokio::sync::Mutex as AsyncMutex;
    use tower::Service;

    #[test]
    fn test_simple_token_counter() {
        let counter = SimpleTokenCounter::new();
        let messages = vec![
            ChatCompletionRequestSystemMessageArgs::default()
                .content("You are a helpful assistant.")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestUserMessageArgs::default()
                .content("Hello, how are you?")
                .build()
                .unwrap()
                .into(),
        ];

        let count = counter.count_messages(&messages);
        assert!(count > 0);
        // Rough estimate: ~50 chars / 4 = ~12 tokens
        assert!((10..=20).contains(&count));
    }

    // =============================================================================================
    // Boundary and slice-invariance tests for compaction preflight requests
    // =============================================================================================

    /// Build a minimal AutoCompaction with a capturing provider used only for compact_messages
    fn make_compactor_with_policy(
        policy: CompactionPolicy,
        provider: CapturingProvider,
    ) -> AutoCompaction<(), CapturingProvider, SimpleTokenCounter> {
        AutoCompaction {
            inner: Arc::new(tokio::sync::Mutex::new(())),
            policy,
            provider: Arc::new(tokio::sync::Mutex::new(provider)),
            token_counter: Arc::new(SimpleTokenCounter::new()),
        }
    }

    fn asst_with_tool_calls(ids: &[&str]) -> ChatCompletionRequestMessage {
        use async_openai::types::{
            ChatCompletionMessageToolCall, ChatCompletionToolType, FunctionCall,
        };
        let calls: Vec<ChatCompletionMessageToolCall> = ids
            .iter()
            .map(|id| ChatCompletionMessageToolCall {
                id: (*id).to_string(),
                r#type: ChatCompletionToolType::Function,
                function: FunctionCall {
                    name: "x".to_string(),
                    arguments: "{}".to_string(),
                },
            })
            .collect();
        ChatCompletionRequestAssistantMessageArgs::default()
            .content("")
            .tool_calls(calls)
            .build()
            .unwrap()
            .into()
    }

    fn tool_msg(id: &str, content: &str) -> ChatCompletionRequestMessage {
        ChatCompletionRequestToolMessageArgs::default()
            .tool_call_id(id)
            .content(content)
            .build()
            .unwrap()
            .into()
    }

    fn user_msg(text: &str) -> ChatCompletionRequestMessage {
        ChatCompletionRequestUserMessageArgs::default()
            .content(text)
            .build()
            .unwrap()
            .into()
    }

    fn system_msg(text: &str) -> ChatCompletionRequestMessage {
        ChatCompletionRequestSystemMessageArgs::default()
            .content(text)
            .build()
            .unwrap()
            .into()
    }

    #[test]
    fn compaction_preflight_excludes_tools_at_slice_end_custom_end() {
        // TODO: tests_compaction_slice_end_missing_tools
        // Construct a transcript that is acceptable under default validation (contiguity not enforced),
        // but where we select a slice that ends immediately after an assistant with tool_calls
        // (thus excluding tool responses), provoking a preflight request that violates tool invariants.

        // Messages: [sys, user1, assistant(tc c1), user_interrupt, tool(c1)]
        let msgs = vec![
            system_msg("S"),
            user_msg("U1"),
            asst_with_tool_calls(&["c1"]),
            user_msg("interrupt"),
            tool_msg("c1", "{}"),
        ];

        // Policy: Custom compaction range to include up to (but not including) the interrupt user
        // so the slice ends with the assistant with tool_calls.
        let range_fn: PromptGeneratorFn = Arc::new(|_m| "".to_string());
        let range_select: CompactionRangeFn = Arc::new(|m| {
            // Find assistant with tool_calls index (expect index 2), set end = idx + 1
            let idx = m
                .iter()
                .position(|mm| matches!(mm, ChatCompletionRequestMessage::Assistant(_)))
                .unwrap();
            CompactionRange {
                start: 1,
                end: idx + 1,
            }
        });
        let policy = CompactionPolicy {
            compaction_model: "gpt-4o-mini".into(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::Custom(range_select),
            compaction_prompt: CompactionPrompt::Dynamic(range_fn),
            max_compaction_attempts: 1,
            orphaned_tool_call_strategy: OrphanedToolCallStrategy::DropAndReappend,
        };

        // Capturing provider to inspect the preflight request sent to the model
        let provider = CapturingProvider::new();
        let provider_clone = provider.clone();
        let ac = make_compactor_with_policy(policy, provider_clone);

        let rt = tokio::runtime::Runtime::new().unwrap();
        let out = rt
            .block_on(async move { ac.compact_messages(msgs.clone()).await })
            .unwrap();

        // Preflight request should violate tool invariants: assistant tool_calls without tool responses
        let req = rt
            .block_on(async move { provider.get().await })
            .expect("captured");
        let strict = ValidationPolicy {
            require_user_first: false,
            allow_repeated_roles: true, // allow tool blocks
            enforce_contiguous_tool_responses: true,
            ..Default::default()
        };
        assert!(
            validate_conversation(&req.messages, &strict).is_none(),
            "expected preflight to be valid when slice ends after assistant (slice-aware)"
        );

        // Post-compaction output exists (we're not asserting validity here)
        assert!(!out.is_empty());
    }

    #[test]
    fn compaction_slice_starts_with_tool_message_custom_start() {
        // TODO: tests_compaction_slice_starts_with_tool
        // Construct a valid transcript, then select a slice that begins at a tool message.
        let msgs = vec![
            system_msg("S"),
            user_msg("U1"),
            asst_with_tool_calls(&["c1"]),
            tool_msg("c1", "{}"),
            user_msg("U2"),
        ];

        // Custom range: start at the first Tool message index
        let range_fn: PromptGeneratorFn = Arc::new(|_m| "".to_string());
        let range_select: CompactionRangeFn = Arc::new(|m| {
            let start = m
                .iter()
                .position(|mm| matches!(mm, ChatCompletionRequestMessage::Tool(_)))
                .unwrap();
            CompactionRange {
                start,
                end: start + 1,
            }
        });
        let policy = CompactionPolicy {
            compaction_model: "gpt-4o-mini".into(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::Custom(range_select),
            compaction_prompt: CompactionPrompt::Dynamic(range_fn),
            max_compaction_attempts: 1,
            orphaned_tool_call_strategy: OrphanedToolCallStrategy::DropAndReappend,
        };

        let provider = CapturingProvider::new();
        let ac = make_compactor_with_policy(policy, provider.clone());
        let rt = tokio::runtime::Runtime::new().unwrap();
        let _ = rt
            .block_on(async move { ac.compact_messages(msgs).await })
            .unwrap();
        let req = rt
            .block_on(async move { provider.get().await })
            .expect("captured");

        // Starting with a Tool message should be invalid (tool before assistant)
        let strict = ValidationPolicy {
            require_user_first: false,
            ..Default::default()
        };
        let violations = validate_conversation(&req.messages, &strict).expect("violations");
        assert!(
            violations.iter().any(|v| matches!(
                v.code,
                crate::validation::ViolationCode::ToolBeforeAssistant { .. }
            )),
            "expected ToolBeforeAssistant violation in preflight request"
        );
    }

    #[test]
    fn compaction_appends_user_after_assistant_with_missing_tools() {
        // TODO: tests_compaction_appended_user_after_tools_missing
        // Similar to the first test, but assert the appended user instruction appears
        // immediately after an assistant with tool_calls when no tools are in-slice.
        let msgs = vec![
            system_msg("S"),
            user_msg("U1"),
            asst_with_tool_calls(&["c1"]),
            user_msg("interrupt"),
            tool_msg("c1", "{}"),
        ];

        // Select range that ends right after assistant (excludes tool response and the interrupt)
        let range_fn: PromptGeneratorFn = Arc::new(|_m| "".to_string());
        let range_select: CompactionRangeFn = Arc::new(|m| {
            let idx = m
                .iter()
                .position(|mm| matches!(mm, ChatCompletionRequestMessage::Assistant(_)))
                .unwrap();
            CompactionRange {
                start: 1,
                end: idx + 1,
            }
        });
        let policy = CompactionPolicy {
            compaction_model: "gpt-4o-mini".into(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::Custom(range_select),
            compaction_prompt: CompactionPrompt::Dynamic(range_fn),
            max_compaction_attempts: 1,
            orphaned_tool_call_strategy: OrphanedToolCallStrategy::DropAndReappend,
        };

        let provider = CapturingProvider::new();
        let ac = make_compactor_with_policy(policy, provider.clone());
        let rt = tokio::runtime::Runtime::new().unwrap();
        let _ = rt
            .block_on(async move { ac.compact_messages(msgs).await })
            .unwrap();
        let req = rt
            .block_on(async move { provider.get().await })
            .expect("captured");

        // The compactor appends a user instruction if the slice doesn't end with a user
        // We ended with assistant, so a user instruction should have been appended.
        assert!(
            matches!(
                req.messages.last(),
                Some(ChatCompletionRequestMessage::User(_))
            ),
            "expected appended user instruction at end of preflight"
        );
    }

    #[test]
    fn compaction_boundary_extends_to_include_all_tool_responses_for_last_assistant() {
        // TODO: tests_compaction_multicall_boundary_complete
        // Assistant with two tool_calls followed by two tool responses; end is set at assistant,
        // compactor should extend to include both contiguous tool responses.
        let msgs = vec![
            system_msg("S"),
            user_msg("U1"),
            asst_with_tool_calls(&["c1", "c2"]),
            tool_msg("c1", "r1"),
            tool_msg("c2", "r2"),
            user_msg("U2"),
        ];

        let range_fn: PromptGeneratorFn = Arc::new(|_m| "".to_string());
        let range_select: CompactionRangeFn = Arc::new(|m| {
            let idx = m
                .iter()
                .position(|mm| matches!(mm, ChatCompletionRequestMessage::Assistant(_)))
                .unwrap();
            CompactionRange {
                start: 1,
                end: idx + 1,
            }
        });
        let policy = CompactionPolicy {
            compaction_model: "gpt-4o-mini".into(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::Custom(range_select),
            compaction_prompt: CompactionPrompt::Dynamic(range_fn),
            max_compaction_attempts: 1,
            orphaned_tool_call_strategy: OrphanedToolCallStrategy::DropAndReappend,
        };

        let provider = CapturingProvider::new();
        let ac = make_compactor_with_policy(policy, provider.clone());
        let rt = tokio::runtime::Runtime::new().unwrap();
        let _ = rt
            .block_on(async move { ac.compact_messages(msgs).await })
            .unwrap();
        let req = rt
            .block_on(async move { provider.get().await })
            .expect("captured");

        // Expect both tool responses present (no violations under strict contiguity policy)
        let strict = ValidationPolicy {
            require_user_first: false,
            allow_repeated_roles: true,
            enforce_contiguous_tool_responses: true,
            ..Default::default()
        };
        if let Some(violations) = validate_conversation(&req.messages, &strict) {
            eprintln!(
                "strict policy: allow_repeated_roles={}, enforce_contiguous_tool_responses={}",
                strict.allow_repeated_roles, strict.enforce_contiguous_tool_responses
            );
            eprintln!("violations: {:?}", violations);
            panic!("expected preflight to be valid when contiguous tools exist");
        }
    }

    #[test]
    fn compaction_global_ok_but_slice_orphaned_missing_tools() {
        // TODO: tests_compaction_global_vs_slice_orphan
        // Full transcript is valid (assistant tool_calls followed by tool),
        // but the chosen slice excludes the tool response, making the preflight invalid.
        let full = vec![
            system_msg("S"),
            user_msg("U1"),
            asst_with_tool_calls(&["c1"]),
            tool_msg("c1", "{}"),
            user_msg("U2"),
        ];

        let base_policy = ValidationPolicy {
            require_user_first: false,
            ..Default::default()
        };
        assert!(validate_conversation(&full, &base_policy).is_none());

        // Custom range excludes the tool response (end at assistant)
        let range_fn: PromptGeneratorFn = Arc::new(|_m| "".to_string());
        let range_select: CompactionRangeFn = Arc::new(|m| {
            let idx = m
                .iter()
                .position(|mm| matches!(mm, ChatCompletionRequestMessage::Assistant(_)))
                .unwrap();
            CompactionRange {
                start: 1,
                end: idx + 1,
            }
        });
        let policy = CompactionPolicy {
            compaction_model: "gpt-4o-mini".into(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::Custom(range_select),
            compaction_prompt: CompactionPrompt::Dynamic(range_fn),
            max_compaction_attempts: 1,
            orphaned_tool_call_strategy: OrphanedToolCallStrategy::DropAndReappend,
        };

        let provider = CapturingProvider::new();
        let ac = make_compactor_with_policy(policy, provider.clone());
        let rt = tokio::runtime::Runtime::new().unwrap();
        let _ = rt
            .block_on(async move { ac.compact_messages(full).await })
            .unwrap();
        let req = rt
            .block_on(async move { provider.get().await })
            .expect("captured");

        let strict = ValidationPolicy {
            require_user_first: false,
            allow_repeated_roles: true,
            enforce_contiguous_tool_responses: true,
            ..Default::default()
        };
        assert!(validate_conversation(&req.messages, &strict).is_none());
    }

    #[test]
    fn compaction_strategy_comparison_drop_vs_placeholder() {
        // TODO: tests_compaction_property_strategy_comparison
        // Same input/slice, compare DropAndReappend vs AddPlaceholderResponses
        let msgs = vec![
            system_msg("S"),
            user_msg("U1"),
            asst_with_tool_calls(&["c1"]),
            user_msg("interrupt"),
            tool_msg("c1", "{}"),
        ];

        let range_fn: PromptGeneratorFn = Arc::new(|_m| "".to_string());
        let range_select: CompactionRangeFn = Arc::new(|m| {
            let idx = m
                .iter()
                .position(|mm| matches!(mm, ChatCompletionRequestMessage::Assistant(_)))
                .unwrap();
            CompactionRange {
                start: 1,
                end: idx + 1,
            }
        });

        let mk_pol = |strat| CompactionPolicy {
            compaction_model: "gpt-4o-mini".into(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::Custom(range_select.clone()),
            compaction_prompt: CompactionPrompt::Dynamic(range_fn.clone()),
            max_compaction_attempts: 1,
            orphaned_tool_call_strategy: strat,
        };

        // DropAndReappend
        let prov1 = CapturingProvider::new();
        let ac1 = make_compactor_with_policy(
            mk_pol(OrphanedToolCallStrategy::DropAndReappend),
            prov1.clone(),
        );
        let rt = tokio::runtime::Runtime::new().unwrap();
        let _ = rt
            .block_on(async { ac1.compact_messages(msgs.clone()).await })
            .unwrap();
        let req1 = rt.block_on(async move { prov1.get().await }).unwrap();

        // AddPlaceholderResponses
        let prov2 = CapturingProvider::new();
        let ac2 = make_compactor_with_policy(
            mk_pol(OrphanedToolCallStrategy::AddPlaceholderResponses),
            prov2.clone(),
        );
        let _ = rt
            .block_on(async { ac2.compact_messages(msgs.clone()).await })
            .unwrap();
        let req2 = rt.block_on(async move { prov2.get().await }).unwrap();

        let strict = ValidationPolicy {
            require_user_first: false,
            allow_repeated_roles: true,
            enforce_contiguous_tool_responses: true,
            ..Default::default()
        };
        if validate_conversation(&req2.messages, &strict).is_some() {
            eprintln!(
                "placeholder violations: {:?}",
                validate_conversation(&req2.messages, &strict).unwrap()
            );
        }
        // With slice-aware boundary repair, both strategies should yield valid preflight
        assert!(
            validate_conversation(&req1.messages, &strict).is_none(),
            "DropAndReappend should be valid with slice-aware boundary repair"
        );
        assert!(
            validate_conversation(&req2.messages, &strict).is_none(),
            "AddPlaceholderResponses should be valid with slice-aware boundary repair"
        );
    }

    #[test]
    fn compaction_on_merged_history_like_handoff_preflight_validity() {
        // TODO: tests_compaction_handoff_merged_preflight
        // Simulate merged history (original + accumulated) and ensure preflight validity can fail
        // with the strict policy when slice cuts across tool boundaries.
        let merged = vec![
            system_msg("S"),
            user_msg("U1"),
            asst_with_tool_calls(&["c1"]),
            user_msg("interrupt"),
            tool_msg("c1", "{}"),
            user_msg("U2"),
        ];

        let range_fn: PromptGeneratorFn = Arc::new(|_m| "".to_string());
        let range_select: CompactionRangeFn = Arc::new(|m| {
            // Cut right after assistant (exclude interrupt/user and tools)
            let idx = m
                .iter()
                .position(|mm| matches!(mm, ChatCompletionRequestMessage::Assistant(_)))
                .unwrap();
            CompactionRange {
                start: 1,
                end: idx + 1,
            }
        });
        let policy = CompactionPolicy {
            compaction_model: "gpt-4o-mini".into(),
            proactive_threshold: None,
            compaction_strategy: CompactionStrategy::Custom(range_select),
            compaction_prompt: CompactionPrompt::Dynamic(range_fn),
            max_compaction_attempts: 1,
            orphaned_tool_call_strategy: OrphanedToolCallStrategy::DropAndReappend,
        };

        let provider = CapturingProvider::new();
        let ac = make_compactor_with_policy(policy, provider.clone());
        let rt = tokio::runtime::Runtime::new().unwrap();
        let _ = rt
            .block_on(async move { ac.compact_messages(merged).await })
            .unwrap();
        let req = rt
            .block_on(async move { provider.get().await })
            .expect("captured");

        let strict = ValidationPolicy {
            require_user_first: false,
            enforce_contiguous_tool_responses: true,
            ..Default::default()
        };
        // TDD red: expect valid preflight with slice-aware compaction
        assert!(validate_conversation(&req.messages, &strict).is_none());
    }

    #[test]
    fn test_compaction_range_preserve_system() {
        let strategy = CompactionStrategy::PreserveSystemAndRecent { recent_count: 2 };
        let messages = vec![
            ChatCompletionRequestSystemMessageArgs::default()
                .content("System")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestUserMessageArgs::default()
                .content("User1")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("Assistant1")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestUserMessageArgs::default()
                .content("User2")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("Assistant2")
                .build()
                .unwrap()
                .into(),
        ];

        let range = match strategy {
            CompactionStrategy::PreserveSystemAndRecent { recent_count } => {
                let start = messages
                    .iter()
                    .position(|m| !matches!(m, ChatCompletionRequestMessage::System(_)))
                    .unwrap_or(0);
                let end = messages.len().saturating_sub(recent_count);
                CompactionRange {
                    start,
                    end: end.max(start),
                }
            }
            _ => panic!("Wrong strategy"),
        };

        assert_eq!(range.start, 1); // After system message
        assert_eq!(range.end, 3); // Keep last 2 messages
    }

    #[test]
    fn test_compaction_strategy_compact_all_but_last() {
        let strategy = CompactionStrategy::CompactAllButLast(3);
        let messages: Vec<ChatCompletionRequestMessage> = vec![
            ChatCompletionRequestUserMessageArgs::default()
                .content("1")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("2")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestUserMessageArgs::default()
                .content("3")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("4")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestUserMessageArgs::default()
                .content("5")
                .build()
                .unwrap()
                .into(),
        ];

        let range = match strategy {
            CompactionStrategy::CompactAllButLast(n) => {
                let start = 0;
                let end = messages.len().saturating_sub(n);
                CompactionRange { start, end }
            }
            _ => panic!("Wrong strategy"),
        };

        assert_eq!(range.start, 0);
        assert_eq!(range.end, 2); // Keep last 3 messages
    }

    #[test]
    fn test_compaction_prompt_default() {
        let prompt = CompactionPrompt::Default;
        let messages = vec![];
        let generated = prompt.generate(&messages);
        assert!(generated.contains("concise summary"));
        assert!(generated.contains("conversation continuity"));
    }

    #[test]
    fn test_compaction_prompt_custom() {
        let prompt = CompactionPrompt::Custom("Custom prompt".to_string());
        let messages = vec![];
        let generated = prompt.generate(&messages);
        assert_eq!(generated, "Custom prompt");
    }

    #[test]
    fn test_compaction_prompt_dynamic() {
        let prompt = CompactionPrompt::Dynamic(Arc::new(|msgs| {
            format!("Summarize {} messages", msgs.len())
        }));
        let messages = vec![ChatCompletionRequestUserMessageArgs::default()
            .content("test")
            .build()
            .unwrap()
            .into()];
        let generated = prompt.generate(&messages);
        assert_eq!(generated, "Summarize 1 messages");
    }

    #[test]
    fn test_is_context_length_error() {
        let test_cases = vec![
            ("context_length_exceeded", true),
            ("Error: context length exceeded", true),
            ("Maximum context reached", true),
            ("Token limit exceeded", true),
            ("Some other error", false),
            ("Network timeout", false),
        ];

        for (error_msg, expected) in test_cases {
            let error: BoxError = error_msg.into();
            assert_eq!(
                AutoCompaction::<DummyService, DummyProvider, SimpleTokenCounter>::is_context_length_error(&error),
                expected,
                "Failed for: {}",
                error_msg
            );
        }
    }

    #[test]
    fn test_proactive_threshold() {
        let threshold = ProactiveThreshold {
            token_threshold: 1000,
            percentage_threshold: Some(0.8),
        };

        assert_eq!(threshold.token_threshold, 1000);
        assert_eq!(threshold.percentage_threshold, Some(0.8));
    }

    #[test]
    fn test_compaction_policy_default() {
        let policy = CompactionPolicy::default();
        assert_eq!(policy.compaction_model, "gpt-4o-mini");
        assert!(policy.proactive_threshold.is_none());
        assert_eq!(policy.max_compaction_attempts, 2);
        matches!(
            policy.compaction_strategy,
            CompactionStrategy::PreserveSystemAndRecent { .. }
        );
        matches!(
            policy.orphaned_tool_call_strategy,
            OrphanedToolCallStrategy::DropAndReappend
        );
    }

    #[test]
    fn test_orphaned_tool_call_detection() {
        use async_openai::types::{
            ChatCompletionMessageToolCall, ChatCompletionRequestToolMessageArgs,
            ChatCompletionToolType, FunctionCall,
        };

        // Create messages with an orphaned tool call
        let messages = vec![
            ChatCompletionRequestUserMessageArgs::default()
                .content("Please help me")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("")
                .tool_calls(vec![ChatCompletionMessageToolCall {
                    id: "call_123".to_string(),
                    r#type: ChatCompletionToolType::Function,
                    function: FunctionCall {
                        name: "some_tool".to_string(),
                        arguments: "{}".to_string(),
                    },
                }])
                .build()
                .unwrap()
                .into(),
            // No tool response message - this is orphaned!
        ];

        assert!(has_orphaned_tool_calls(&messages));

        // Now add the tool response - no longer orphaned
        let mut messages_with_response = messages.clone();
        messages_with_response.push(
            ChatCompletionRequestToolMessageArgs::default()
                .content("Tool result")
                .tool_call_id("call_123")
                .build()
                .unwrap()
                .into(),
        );

        assert!(!has_orphaned_tool_calls(&messages_with_response));
    }

    #[derive(Clone)]
    struct CapturingProvider {
        captured: StdArc<AsyncMutex<Option<CreateChatCompletionRequest>>>,
    }
    impl CapturingProvider {
        fn new() -> Self {
            Self {
                captured: StdArc::new(AsyncMutex::new(None)),
            }
        }
        async fn get(&self) -> Option<CreateChatCompletionRequest> {
            self.captured.lock().await.clone()
        }
    }
    impl Service<CreateChatCompletionRequest> for CapturingProvider {
        type Response = crate::provider::ProviderResponse;
        type Error = BoxError;
        type Future = std::pin::Pin<
            Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
        >;
        fn poll_ready(
            &mut self,
            _cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Result<(), Self::Error>> {
            std::task::Poll::Ready(Ok(()))
        }
        fn call(&mut self, req: CreateChatCompletionRequest) -> Self::Future {
            let captured = self.captured.clone();
            Box::pin(async move {
                *captured.lock().await = Some(req);
                #[allow(deprecated)]
                let assistant = async_openai::types::ChatCompletionResponseMessage {
                    content: Some("summary".into()),
                    role: async_openai::types::Role::Assistant,
                    tool_calls: None,
                    refusal: None,
                    audio: None,
                    function_call: None,
                };
                Ok(crate::provider::ProviderResponse {
                    assistant,
                    prompt_tokens: 1,
                    completion_tokens: 1,
                })
            })
        }
    }

    proptest! {
        #[test]
        fn compaction_preflight_and_post_valid_for_valid_inputs(msgs in gen::valid_conversation(gen::GeneratorConfig::default())) {
            let provider = CapturingProvider::new();
            let provider_clone = provider.clone();
            let policy = CompactionPolicy {
                compaction_model: "gpt-4o-mini".into(),
                compaction_strategy: CompactionStrategy::CompactAllButLast(1),
                orphaned_tool_call_strategy: OrphanedToolCallStrategy::AddPlaceholderResponses,
                ..Default::default()
            };
            let counter = SimpleTokenCounter::new();
            let dummy_inner = (); // unused by compact_messages
            let ac = AutoCompaction {
                inner: Arc::new(tokio::sync::Mutex::new(dummy_inner)),
                policy,
                provider: Arc::new(tokio::sync::Mutex::new(provider_clone)),
                token_counter: Arc::new(counter),
            };
            let rt = tokio::runtime::Runtime::new().unwrap();
            let out = rt.block_on(async move { ac.compact_messages(msgs.clone()).await.unwrap() });
            // Preflight validity: captured request sent to compactor must be valid
            let req = rt.block_on(async move { provider.get().await }).expect("captured");
            prop_assert!(validate_conversation(&req.messages, &ValidationPolicy::default()).is_none());
            // Post-compaction validity: output must be valid
            prop_assert!(validate_conversation(&out, &ValidationPolicy::default()).is_none());
        }
    }

    #[test]
    fn test_multiple_orphaned_tool_calls_in_middle() {
        use async_openai::types::{
            ChatCompletionMessageToolCall, ChatCompletionToolType, FunctionCall,
        };

        // Test case similar to the production error - orphaned tool calls in the middle
        let messages = vec![
            ChatCompletionRequestUserMessageArgs::default()
                .content("First user message")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("First response")
                .build()
                .unwrap()
                .into(),
            // First orphaned tool call (in the middle)
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("")
                .tool_calls(vec![ChatCompletionMessageToolCall {
                    id: "call_middle".to_string(),
                    r#type: ChatCompletionToolType::Function,
                    function: FunctionCall {
                        name: "middle_tool".to_string(),
                        arguments: "{}".to_string(),
                    },
                }])
                .build()
                .unwrap()
                .into(),
            // No tool response - orphaned!
            // Conversation continues...
            ChatCompletionRequestUserMessageArgs::default()
                .content("Continue conversation")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("Continuing...")
                .build()
                .unwrap()
                .into(),
            // Another orphaned tool call at the end
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("")
                .tool_calls(vec![ChatCompletionMessageToolCall {
                    id: "call_end".to_string(),
                    r#type: ChatCompletionToolType::Function,
                    function: FunctionCall {
                        name: "end_tool".to_string(),
                        arguments: "{}".to_string(),
                    },
                }])
                .build()
                .unwrap()
                .into(),
        ];

        // Should detect both orphaned tool calls
        assert!(has_orphaned_tool_calls(&messages));
        let orphaned_indices = find_orphaned_tool_calls_test(&messages);
        assert_eq!(orphaned_indices.len(), 2);
        assert_eq!(orphaned_indices[0], 2); // Middle orphaned message
        assert_eq!(orphaned_indices[1], 5); // End orphaned message

        // Extract orphaned messages
        let (cleaned, orphaned) = extract_orphaned_tool_calls(messages.clone());
        assert_eq!(cleaned.len(), 4); // 6 original - 2 orphaned = 4
        assert_eq!(orphaned.len(), 2);

        // Verify the correct messages were extracted
        assert_eq!(orphaned[0].0, 2);
        assert_eq!(orphaned[1].0, 5);
    }

    #[test]
    fn test_mixed_orphaned_and_valid_tool_calls() {
        use async_openai::types::{
            ChatCompletionMessageToolCall, ChatCompletionRequestToolMessageArgs,
            ChatCompletionToolType, FunctionCall,
        };

        let messages = vec![
            ChatCompletionRequestUserMessageArgs::default()
                .content("User message")
                .build()
                .unwrap()
                .into(),
            // Valid tool call with response
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("")
                .tool_calls(vec![ChatCompletionMessageToolCall {
                    id: "call_valid".to_string(),
                    r#type: ChatCompletionToolType::Function,
                    function: FunctionCall {
                        name: "valid_tool".to_string(),
                        arguments: "{}".to_string(),
                    },
                }])
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestToolMessageArgs::default()
                .content("Valid tool response")
                .tool_call_id("call_valid")
                .build()
                .unwrap()
                .into(),
            // Orphaned tool call
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("")
                .tool_calls(vec![ChatCompletionMessageToolCall {
                    id: "call_orphaned".to_string(),
                    r#type: ChatCompletionToolType::Function,
                    function: FunctionCall {
                        name: "orphaned_tool".to_string(),
                        arguments: "{}".to_string(),
                    },
                }])
                .build()
                .unwrap()
                .into(),
            // No response for orphaned call
            ChatCompletionRequestUserMessageArgs::default()
                .content("Another message")
                .build()
                .unwrap()
                .into(),
        ];

        assert!(has_orphaned_tool_calls(&messages));
        let orphaned_indices = find_orphaned_tool_calls_test(&messages);
        assert_eq!(orphaned_indices.len(), 1);
        assert_eq!(orphaned_indices[0], 3); // Only the orphaned one

        let (cleaned, orphaned) = extract_orphaned_tool_calls(messages.clone());
        assert_eq!(cleaned.len(), 4); // Valid tool call and response remain
        assert_eq!(orphaned.len(), 1);
        assert_eq!(orphaned[0].0, 3);
    }

    #[test]
    fn test_multiple_tool_calls_partial_orphaned() {
        use async_openai::types::{
            ChatCompletionMessageToolCall, ChatCompletionRequestToolMessageArgs,
            ChatCompletionToolType, FunctionCall,
        };

        // Assistant message with multiple tool calls, only some have responses
        let messages = vec![
            ChatCompletionRequestUserMessageArgs::default()
                .content("User message")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("")
                .tool_calls(vec![
                    ChatCompletionMessageToolCall {
                        id: "call_1".to_string(),
                        r#type: ChatCompletionToolType::Function,
                        function: FunctionCall {
                            name: "tool_1".to_string(),
                            arguments: "{}".to_string(),
                        },
                    },
                    ChatCompletionMessageToolCall {
                        id: "call_2".to_string(),
                        r#type: ChatCompletionToolType::Function,
                        function: FunctionCall {
                            name: "tool_2".to_string(),
                            arguments: "{}".to_string(),
                        },
                    },
                ])
                .build()
                .unwrap()
                .into(),
            // Only response for call_1, not call_2
            ChatCompletionRequestToolMessageArgs::default()
                .content("Response for call_1")
                .tool_call_id("call_1")
                .build()
                .unwrap()
                .into(),
            // call_2 is orphaned!
        ];

        assert!(has_orphaned_tool_calls(&messages));
        let orphaned_indices = find_orphaned_tool_calls_test(&messages);
        assert_eq!(orphaned_indices.len(), 1);
        assert_eq!(orphaned_indices[0], 1); // The assistant message with partial orphaned calls
    }

    #[test]
    fn test_orphaned_tool_call_strategy_drop_and_reappend() {
        use async_openai::types::{
            ChatCompletionMessageToolCall, ChatCompletionToolType, FunctionCall,
        };

        let messages = vec![
            ChatCompletionRequestUserMessageArgs::default()
                .content("User message")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("Assistant response")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestUserMessageArgs::default()
                .content("Another user message")
                .build()
                .unwrap()
                .into(),
            ChatCompletionRequestAssistantMessageArgs::default()
                .content("")
                .tool_calls(vec![ChatCompletionMessageToolCall {
                    id: "orphaned_call".to_string(),
                    r#type: ChatCompletionToolType::Function,
                    function: FunctionCall {
                        name: "orphaned_tool".to_string(),
                        arguments: "{}".to_string(),
                    },
                }])
                .build()
                .unwrap()
                .into(),
        ];

        let (cleaned, orphaned) = extract_orphaned_tool_calls(messages.clone());

        // Should have removed the last message with tool calls
        assert_eq!(cleaned.len(), 3);
        assert_eq!(orphaned.len(), 1);

        // The orphaned message should be the one with tool calls
        let (original_idx, orphaned_msg) = &orphaned[0];
        assert_eq!(*original_idx, 3); // The 4th message (index 3) was orphaned
        if let ChatCompletionRequestMessage::Assistant(asst) = orphaned_msg {
            assert!(asst.tool_calls.is_some());
        } else {
            panic!("Expected assistant message");
        }
    }

    // Helper functions for testing
    fn has_orphaned_tool_calls(messages: &[ChatCompletionRequestMessage]) -> bool {
        // Use the actual implementation's logic
        !find_orphaned_tool_calls_test(messages).is_empty()
    }

    fn find_orphaned_tool_calls_test(messages: &[ChatCompletionRequestMessage]) -> Vec<usize> {
        let mut orphaned_indices = Vec::new();

        for (i, msg) in messages.iter().enumerate() {
            if let ChatCompletionRequestMessage::Assistant(asst) = msg {
                if let Some(tool_calls) = &asst.tool_calls {
                    if !tool_calls.is_empty() {
                        // Check if all tool calls have corresponding responses
                        let mut all_calls_have_responses = true;
                        for tool_call in tool_calls {
                            let mut found_response = false;
                            // Look for a tool response with matching ID after this message
                            for subsequent_msg in messages.iter().skip(i + 1) {
                                if let ChatCompletionRequestMessage::Tool(tool_msg) = subsequent_msg
                                {
                                    if tool_msg.tool_call_id == tool_call.id {
                                        found_response = true;
                                        break;
                                    }
                                }
                                // Stop looking if we hit another assistant message (conversation moved on)
                                if matches!(
                                    subsequent_msg,
                                    ChatCompletionRequestMessage::Assistant(_)
                                ) {
                                    break;
                                }
                            }
                            if !found_response {
                                all_calls_have_responses = false;
                                break;
                            }
                        }

                        if !all_calls_have_responses {
                            orphaned_indices.push(i);
                        }
                    }
                }
            }
        }

        orphaned_indices
    }

    fn extract_orphaned_tool_calls(
        messages: Vec<ChatCompletionRequestMessage>,
    ) -> (
        Vec<ChatCompletionRequestMessage>,
        Vec<(usize, ChatCompletionRequestMessage)>,
    ) {
        let orphaned_indices = find_orphaned_tool_calls_test(&messages);

        if orphaned_indices.is_empty() {
            return (messages, Vec::new());
        }

        let mut cleaned = Vec::new();
        let mut orphaned = Vec::new();

        for (i, msg) in messages.into_iter().enumerate() {
            if orphaned_indices.contains(&i) {
                orphaned.push((i, msg));
            } else {
                cleaned.push(msg);
            }
        }

        (cleaned, orphaned)
    }

    // Dummy types for testing
    struct DummyService;
    struct DummyProvider;

    impl Service<CreateChatCompletionRequest> for DummyService {
        type Response = crate::core::StepOutcome;
        type Error = BoxError;
        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

        fn poll_ready(
            &mut self,
            _: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Result<(), Self::Error>> {
            std::task::Poll::Ready(Ok(()))
        }

        fn call(&mut self, _: CreateChatCompletionRequest) -> Self::Future {
            Box::pin(async {
                Ok(crate::core::StepOutcome::Done {
                    messages: vec![],
                    aux: Default::default(),
                })
            })
        }
    }

    impl Service<CreateChatCompletionRequest> for DummyProvider {
        type Response = crate::provider::ProviderResponse;
        type Error = BoxError;
        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

        fn poll_ready(
            &mut self,
            _: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Result<(), Self::Error>> {
            std::task::Poll::Ready(Ok(()))
        }

        fn call(&mut self, _: CreateChatCompletionRequest) -> Self::Future {
            Box::pin(async {
                #[allow(deprecated)]
                let assistant = async_openai::types::ChatCompletionResponseMessage {
                    content: Some("summary".to_string()),
                    role: async_openai::types::Role::Assistant,
                    tool_calls: None,
                    refusal: None,
                    audio: None,
                    function_call: None,
                };
                Ok(crate::provider::ProviderResponse {
                    assistant,
                    prompt_tokens: 10,
                    completion_tokens: 10,
                })
            })
        }
    }
}