roder_core/

compaction_runtime.rs

use std::collections::HashMap;

use futures::StreamExt;
use roder_api::events::{ThreadId, TurnId};
use roder_api::inference::{
    AgentInferenceRequest, InferenceEvent, InferenceTurnContext, InstructionBundle, MessageDelta,
    ModelSelection, OutputConfig, ReasoningConfig, RuntimeHints, RuntimeProfile,
};
use roder_api::tools::ToolChoice;
use roder_api::transcript::{TranscriptItem, UserMessage};
use std::sync::Mutex;

use crate::compaction::{
    CompactionOptions, accept_llm_compaction_summary, build_compaction_summary_prompt,
    build_compaction_verify_prompt, estimate_prompt_tokens,
};
use crate::runtime::Runtime;

#[derive(Default)]
pub(crate) struct CompactionState {
    tokens: u32,
    generation: u64,
}

impl Runtime {
    pub(crate) fn record_compaction_hysteresis(&self, thread_id: &ThreadId, trigger_tokens: u32) {
        if let Ok(mut state) = self.compaction_hysteresis.lock() {
            let state = state.entry(thread_id.clone()).or_default();
            state.tokens = trigger_tokens;
            state.generation = state.generation.wrapping_add(1);
        }
    }

    pub(crate) fn compaction_hysteresis_baseline(&self, thread_id: &ThreadId) -> Option<u32> {
        self.compaction_hysteresis
            .lock()
            .ok()
            .and_then(|state| state.get(thread_id).map(|state| state.tokens))
    }

    pub(crate) fn compaction_generation(&self, thread_id: &ThreadId) -> u64 {
        self.compaction_hysteresis
            .lock()
            .ok()
            .and_then(|state| state.get(thread_id).map(|state| state.generation))
            .unwrap_or(0)
    }

    pub(crate) fn compaction_options_for_turn(
        &self,
        thread_id: &ThreadId,
        allow_repeat: bool,
    ) -> CompactionOptions {
        CompactionOptions {
            allow_repeat,
            force: false,
            hysteresis_baseline: self.compaction_hysteresis_baseline(thread_id),
            preserve_hint: None,
        }
    }

    pub async fn force_compact_thread(
        &self,
        thread_id: &ThreadId,
        turn_id: &TurnId,
        preserve_hint: Option<String>,
    ) -> anyhow::Result<ForceCompactOutcome> {
        let _thread_admission = self.thread_admission(thread_id).await;
        let cfg = self.status().await;
        let selection = self
            .parent_model_selection_for_subagents(thread_id, turn_id)
            .await;
        let provider = selection
            .as_ref()
            .map(|selection| selection.provider.clone())
            .unwrap_or(cfg.default_provider.clone());
        let model = selection
            .as_ref()
            .map(|selection| selection.model.clone())
            .unwrap_or(cfg.default_model.clone());
        let transcript = self.transcript_for_force_compact(thread_id).await?;
        if transcript.is_empty() {
            return Ok(ForceCompactOutcome {
                compacted: false,
                reason: Some("empty_transcript".to_string()),
                estimated_tokens_before: 0,
                estimated_tokens_after: 0,
            });
        }
        let estimated_before = estimate_prompt_tokens(&transcript);
        // Manual compaction runs only on an idle task. A second sampling
        // request against a snapshot taken during tool execution can omit
        // effects that complete before the boundary is committed.
        if self.active_turn_for_thread(thread_id).await.is_some() {
            return Ok(ForceCompactOutcome {
                compacted: false,
                reason: Some("turn_active".into()),
                estimated_tokens_before: estimated_before,
                estimated_tokens_after: estimated_before,
            });
        }

        let compacted = self
            .compact_transcript_if_needed(
                thread_id,
                turn_id,
                &provider,
                &model,
                transcript,
                CompactionOptions {
                    allow_repeat: true,
                    force: true,
                    hysteresis_baseline: None,
                    preserve_hint: preserve_hint.filter(|text| !text.trim().is_empty()),
                },
            )
            .await?;
        let estimated_after = estimate_prompt_tokens(&compacted);
        Ok(ForceCompactOutcome {
            compacted: estimated_after < estimated_before
                || compacted
                    .iter()
                    .any(crate::compaction::is_compaction_boundary),
            reason: None,
            estimated_tokens_before: estimated_before,
            estimated_tokens_after: estimated_after,
        })
    }

    async fn transcript_for_force_compact(
        &self,
        thread_id: &ThreadId,
    ) -> anyhow::Result<Vec<TranscriptItem>> {
        let Some(store) = &self.thread_store else {
            return Ok(Vec::new());
        };
        let Some(snapshot) = store.load_thread(thread_id).await? else {
            return Ok(Vec::new());
        };
        let mut out = Vec::new();
        for turn in snapshot.turns {
            out.extend(turn.items);
        }
        Ok(crate::compaction::trim_to_last_compaction_boundary(out))
    }

    pub(crate) async fn summarize_compaction_head(
        &self,
        thread_id: &ThreadId,
        turn_id: &TurnId,
        provider: &str,
        model: &str,
        head: &[TranscriptItem],
        preserve_hint: Option<&str>,
    ) -> anyhow::Result<Option<String>> {
        if head.is_empty() {
            return Ok(None);
        }
        // Prefer a full-head LLM snapshot; if the summary request itself hits the
        // provider prompt limit, shrink the head and retry once, then fall back
        // to deterministic summarization.
        let mut summary_head = head.to_vec();
        let draft = loop {
            match self
                .run_compaction_summary_inference(
                    thread_id,
                    turn_id,
                    provider,
                    model,
                    build_compaction_summary_prompt(&summary_head, preserve_hint),
                )
                .await
            {
                Ok(draft) => break draft,
                Err(err)
                    if crate::compaction::is_context_limit_failure_message(&err.to_string()) =>
                {
                    if summary_head.len() <= 1 {
                        return Ok(None);
                    }
                    // Drop the oldest half so the summary prompt fits.
                    let keep_from = summary_head.len() / 2;
                    summary_head = summary_head.split_off(keep_from);
                    continue;
                }
                Err(_) => {
                    // Non-context failures already fall through to deterministic
                    // summaries via Ok(None) below when the provider returns empty.
                    return Ok(None);
                }
            }
        };
        let Some(draft) = draft else {
            return Ok(None);
        };
        if !accept_llm_compaction_summary(head, &draft) {
            return Ok(None);
        }
        let verified = match self
            .run_compaction_summary_inference(
                thread_id,
                turn_id,
                provider,
                model,
                build_compaction_verify_prompt(&draft),
            )
            .await
        {
            Ok(Some(text)) => text,
            Ok(None) | Err(_) => draft.clone(),
        };
        if accept_llm_compaction_summary(head, &verified) {
            Ok(Some(verified))
        } else if accept_llm_compaction_summary(head, &draft) {
            Ok(Some(draft))
        } else {
            Ok(None)
        }
    }

    async fn run_compaction_summary_inference(
        &self,
        thread_id: &ThreadId,
        turn_id: &TurnId,
        provider: &str,
        model: &str,
        prompt: String,
    ) -> anyhow::Result<Option<String>> {
        let engine = self.engine_for(provider)?;
        let request = AgentInferenceRequest {
            model: ModelSelection {
                provider: provider.to_string(),
                model: model.to_string(),
            },
            instructions: InstructionBundle {
                system: Some(
                    "You compress conversation history into durable state snapshots.".to_string(),
                ),
                developer: None,
                developer_context: None,
            },
            transcript: vec![TranscriptItem::UserMessage(UserMessage::text(prompt))],
            tools: Vec::new(),
            tool_choice: ToolChoice::None,
            reasoning: ReasoningConfig::default(),
            output: OutputConfig::default(),
            runtime: RuntimeHints {
                profile: RuntimeProfile::Interactive,
                ..RuntimeHints::default()
            },
            metadata: serde_json::json!({ "roderCompactionSummary": true }),
        };
        let ctx = InferenceTurnContext {
            thread_id,
            turn_id,
            tool_executor: None,
        };
        let mut stream = engine.stream_turn(ctx, request).await?;
        let mut text = String::new();
        let mut completed = false;
        while let Some(event) = stream.next().await {
            match event? {
                InferenceEvent::MessageDelta(MessageDelta { text: delta, .. }) => {
                    text.push_str(&delta)
                }
                InferenceEvent::Failed(failure) => {
                    if crate::compaction::is_context_limit_failure_message(&failure.message) {
                        anyhow::bail!("{}", failure.message);
                    }
                    // Non-context provider failures should not abort the parent
                    // turn compaction path — fall back to deterministic summary.
                    return Ok(None);
                }
                InferenceEvent::Completed(_) => {
                    completed = true;
                    break;
                }
                _ => {}
            }
        }
        if !completed || text.trim().is_empty() {
            Ok(None)
        } else {
            Ok(Some(text.trim().to_string()))
        }
    }
}

#[derive(Debug, Clone)]
pub struct ForceCompactOutcome {
    pub compacted: bool,
    pub reason: Option<String>,
    pub estimated_tokens_before: u32,
    pub estimated_tokens_after: u32,
}

pub(crate) fn compaction_hysteresis_state() -> Mutex<HashMap<ThreadId, CompactionState>> {
    Mutex::new(HashMap::new())
}