agentd/context/

compact.rs

// SPDX-License-Identifier: Apache-2.0
//! **Compaction** (RFC 0026 §5.2): when a context's token estimate crosses
//! `context.compact_at × model_window` (or `context.compact` is called), the
//! older messages are summarized by a structured `think` into the summary
//! block, the last `keep_last` messages stay verbatim, the plan stays
//! verbatim, skill bodies not referenced in the kept window are evicted (the
//! names stay), the version bumps and the record is checkpointed.
//!
//! The runtime never calls the model itself, so compaction is two halves: a
//! pure **plan** ([`plan_compaction`] — which messages to fold + the prompt +
//! the output schema for the summarizer) and a pure **apply**
//! ([`apply_compaction`] — fold the summarizer's verdict into the context).
//! The turn worker runs the `think` in between.

use super::{ContextState, Msg, Summary};
use crate::state::now_ms;
use serde_json::{Value, json};

/// A prepared compaction: what to summarize and how.
#[derive(Debug, Clone)]
pub struct CompactionRequest {
    /// The number of leading messages that will be folded.
    pub fold: usize,
    /// The summarizer prompt (system).
    pub system: String,
    /// The summarizer input (user).
    pub input: String,
    /// The summarizer's output schema.
    pub output_schema: Value,
    /// The context version this plan was made against (apply refuses drift).
    pub version: u64,
}

/// The summarizer's output schema (RFC 0026 §5.1 summary block).
pub fn summary_schema() -> Value {
    json!({
        "type": "object",
        "properties": {
            "goals": {"type": "array", "items": {"type": "string"}},
            "decisions": {"type": "array", "items": {"type": "string"}},
            "open": {"type": "array", "items": {"type": "string"}},
            "facts": {"type": "array", "items": {"type": "string"}},
            "narrative": {"type": "string"}
        },
        "required": ["goals", "decisions", "open", "facts"],
        "additionalProperties": false
    })
}

const SUMMARIZER_SYSTEM: &str = "You compact an agent's conversation memory. Read the transcript excerpt and \
produce a faithful structured summary: goals (what is being pursued), decisions (what was decided and why), \
open (unresolved questions, pending work, promises made), facts (concrete facts, values, identifiers, results \
worth remembering). Keep entries short and specific; never invent; keep identifiers, numbers and names verbatim. \
Reply with ONLY one JSON object matching the schema.";

/// Decide what to fold. `keep_last` messages stay verbatim; a fold always
/// ends before the kept window, never splits an assistant tool-call from its
/// tool results, and always leaves a `User` message first in the kept window.
/// Returns `None` when there is nothing worth folding (fewer than
/// `keep_last + 2` messages, or no user message to fold up to).
pub fn plan_compaction(
    ctx: &ContextState,
    keep_last: usize,
    target_tokens: Option<u64>,
) -> Option<CompactionRequest> {
    let n = ctx.messages.len();
    if n < keep_last + 2 {
        return None;
    }
    let mut fold = n - keep_last;
    // If a target is given, fold more aggressively until the estimate of the
    // kept tail is under it (but always keep at least 2 messages).
    if let Some(target) = target_tokens {
        let mut kept: u64 = ctx.messages[fold..].iter().map(Msg::est_tokens).sum();
        while kept > target && n - fold > 2 {
            kept -= ctx.messages[fold].est_tokens();
            fold += 1;
        }
    }
    // One boundary correction, and it subsumes the tool-round rule: walk back
    // until the first kept message is a `User`. Two things go wrong otherwise,
    // and because the context is durable a single bad fold poisons every later
    // turn and survives a restart:
    //   * an `Assistant` first sends `messages[0].role == "assistant"` in the
    //     anthropic dialect, which the API rejects — the summary block is no
    //     shield, system messages hoist into the top-level `system` field
    //     (`intel/anthropic.rs`), so the assistant really is first on the wire;
    //   * a `Tool` first is a `tool_result` block whose `tool_use` was folded
    //     away — a dangling id, the split the old tool-round rule guarded
    //     against from the other end.
    // Walking back only ever keeps *more*, so both the "at least 2 kept" floor
    // and the target-token loop above stay satisfied. `get` (not an index) is
    // load-bearing: `keep_last` is caller-supplied and may be 0, and folding
    // everything is fine — the next turn appends its user message first.
    while fold > 0 && ctx.messages.get(fold).is_some_and(|m| !m.is_user()) {
        fold -= 1;
    }
    // No user message at or before the boundary: decline rather than fold to a
    // shape the provider refuses. An oversized context is recoverable; a
    // checkpointed context that cannot be sent is not.
    if fold == 0 {
        return None;
    }
    let mut input = String::new();
    if !ctx.summary.is_empty() {
        input.push_str("Previous summary (already compacted; extend it, do not lose it):\n");
        input.push_str(&ctx.summary.render());
        input.push('\n');
    }
    input.push_str("Transcript excerpt to compact:\n");
    for m in &ctx.messages[..fold] {
        input.push_str(&render_for_summary(m));
        input.push('\n');
    }
    Some(CompactionRequest {
        fold,
        system: SUMMARIZER_SYSTEM.to_string(),
        input,
        output_schema: summary_schema(),
        version: ctx.version,
    })
}

fn render_for_summary(m: &Msg) -> String {
    const CAP: usize = 2000;
    let clip = |s: &str| {
        if s.chars().count() > CAP {
            format!("{}…", s.chars().take(CAP).collect::<String>())
        } else {
            s.to_string()
        }
    };
    match m {
        Msg::System { text, .. } => format!("[system] {}", clip(text)),
        Msg::Note { text, .. } => format!("[note] {}", clip(text)),
        Msg::User {
            text, principal, ..
        } => format!(
            "[user{}] {}",
            principal
                .as_deref()
                .map(|p| format!(" {p}"))
                .unwrap_or_default(),
            clip(text)
        ),
        Msg::Assistant {
            text, tool_calls, ..
        } => {
            let calls: Vec<String> = tool_calls
                .iter()
                .map(|c| format!("{}({})", c.name, clip(&c.arguments.to_string())))
                .collect();
            format!(
                "[assistant] {}{}",
                clip(text.as_deref().unwrap_or("")),
                if calls.is_empty() {
                    String::new()
                } else {
                    format!(" calls: {}", calls.join(", "))
                }
            )
        }
        Msg::Tool {
            name,
            content,
            is_error,
            ..
        } => {
            format!(
                "[tool {name}{}] {}",
                if *is_error { " error" } else { "" },
                clip(&content.to_string())
            )
        }
    }
}

/// Fold the summarizer's verdict into the context: absorb the summary, drop
/// the folded messages, bump the version, evict unreferenced skill bodies
/// (names stay — the caller drops the bodies from its cache), recount.
/// Refuses when the context changed since the plan (`version` drift).
pub fn apply_compaction(
    ctx: &mut ContextState,
    req: &CompactionRequest,
    verdict: &Value,
) -> Result<CompactionOutcome, String> {
    if ctx.version != req.version {
        return Err(format!(
            "context version moved from {} to {} during compaction",
            req.version, ctx.version
        ));
    }
    if req.fold > ctx.messages.len() {
        return Err("compaction fold exceeds the message count".into());
    }
    let mut newer: Summary = match verdict {
        Value::Object(_) => serde_json::from_value(verdict.clone())
            .map_err(|e| format!("summary does not match the schema: {e}"))?,
        Value::String(s) => Summary {
            narrative: Some(s.clone()),
            ..Default::default()
        },
        _ => return Err("summary verdict must be an object".into()),
    };
    newer.covers_messages = req.fold as u64;
    newer.updated = now_ms();
    let before_tokens = ctx.est_tokens;
    ctx.summary.absorb(newer);
    ctx.messages.drain(..req.fold);
    ctx.version += 1;
    ctx.recount();
    ctx.touch();
    Ok(CompactionOutcome {
        folded: req.fold,
        version: ctx.version,
        before_tokens,
        after_tokens: ctx.est_tokens,
    })
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CompactionOutcome {
    pub folded: usize,
    pub version: u64,
    pub before_tokens: u64,
    pub after_tokens: u64,
}

/// A degraded compaction for when the summarizer is unavailable: fold the
/// older messages into a plain narrative built from their rendered lines
/// (truncated). Never loses the plan or the skill names.
pub fn apply_fallback(
    ctx: &mut ContextState,
    req: &CompactionRequest,
) -> Result<CompactionOutcome, String> {
    let mut lines: Vec<String> = ctx.messages[..req.fold.min(ctx.messages.len())]
        .iter()
        .map(render_for_summary)
        .collect();
    let mut narrative = lines.join("\n");
    while narrative.len() > 8_000 && lines.len() > 1 {
        lines.remove(0);
        narrative = format!("(earlier messages elided)\n{}", lines.join("\n"));
    }
    apply_compaction(ctx, req, &Value::String(narrative))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::ContextKind;
    use crate::wire::intel::ToolCall;

    fn ctx_with(n: usize) -> ContextState {
        let mut c = ContextState::new(ContextKind::Conversation, 1000);
        for i in 0..n {
            c.append(Msg::user(
                format!("message number {i} with some words in it"),
                None,
            ));
        }
        c
    }

    #[test]
    fn plan_keeps_the_tail_and_does_not_split_tool_rounds() {
        let mut c = ctx_with(6);
        c.append(Msg::assistant(
            None,
            vec![ToolCall {
                id: "c1".into(),
                name: "memory.get".into(),
                arguments: json!({"key": "k"}),
            }],
        ));
        c.append(Msg::tool(
            "c1",
            "memory.get",
            json!({"found": false}),
            false,
        ));
        c.append(Msg::assistant(Some("done".into()), vec![]));
        // 9 messages; keep_last 2 → fold 7 → messages[7] is the tool result, [6]
        // its assistant call, [5] the user that opened the round → fold 5, the
        // last boundary that leaves a user message first (see the walk-back).
        let req = plan_compaction(&c, 2, None).unwrap();
        assert_eq!(req.fold, 5);
        assert!(c.messages[req.fold].is_user());
        assert!(req.input.contains("[user] message number 0"));
        assert!(req.input.contains("[user] message number 4"));
        assert!(
            !req.input.contains("[user] message number 5"),
            "the user that opened the tool round is kept, not folded"
        );
        assert!(
            !req.input.contains("memory.get"),
            "the tool round stays verbatim"
        );
        assert!(
            plan_compaction(&ctx_with(3), 2, None).is_none(),
            "too short to fold"
        );
        // A target folds more aggressively (never below 2 kept).
        let big = ctx_with(20);
        let req = plan_compaction(&big, 10, Some(1)).unwrap();
        assert_eq!(req.fold, 18);
    }

    /// A mixed transcript: `user, assistant, user, assistant(call), tool,
    /// assistant, user, assistant(call), tool, tool, assistant`.
    fn mixed_ctx() -> ContextState {
        let mut c = ContextState::new(ContextKind::Conversation, 1000);
        let call = |id: &str| ToolCall {
            id: id.into(),
            name: "memory.get".into(),
            arguments: json!({"key": id}),
        };
        c.append(Msg::user("first ask with a few words", None));
        c.append(Msg::assistant(Some("first answer".into()), vec![]));
        c.append(Msg::user("second ask with a few words", None));
        c.append(Msg::assistant(None, vec![call("c1")]));
        c.append(Msg::tool(
            "c1",
            "memory.get",
            json!({"found": false}),
            false,
        ));
        c.append(Msg::assistant(Some("second answer".into()), vec![]));
        c.append(Msg::user("third ask with a few words", None));
        c.append(Msg::assistant(None, vec![call("c2"), call("c3")]));
        c.append(Msg::tool("c2", "memory.get", json!({"found": true}), false));
        c.append(Msg::tool("c3", "memory.get", json!({"found": true}), false));
        c.append(Msg::assistant(Some("third answer".into()), vec![]));
        c
    }

    /// The kept window must open on a `User`. The anthropic dialect hoists
    /// system messages (summary + plan included) into the top-level `system`
    /// field, so an assistant left first by a fold really is `messages[0]` on
    /// the wire and the API rejects it — and the context is durable, so that
    /// fold poisons every later turn until someone edits the record by hand.
    #[test]
    fn fold_never_leaves_an_assistant_or_a_tool_result_first() {
        let c = mixed_ctx();
        // keep_last 1 → fold 10 (the trailing assistant); 9 and 8 are tool
        // results, 7 their assistant call → the user at 6 is the boundary.
        let req = plan_compaction(&c, 1, None).unwrap();
        assert_eq!(req.fold, 6);
        assert!(c.messages[req.fold].is_user());
        // keep_last 5 → fold 6 already lands on the user; nothing to correct.
        assert_eq!(plan_compaction(&c, 5, None).unwrap().fold, 6);
        // keep_last 7 → fold 4 is a tool result → back past its call at 3 to
        // the user at 2 (the old rule stopped at 3, an assistant).
        let req = plan_compaction(&c, 7, None).unwrap();
        assert_eq!(req.fold, 2);
        assert!(c.messages[req.fold].is_user());
        // An aggressive target folds forward first, then walks back the same way.
        let req = plan_compaction(&c, 2, Some(1)).unwrap();
        assert!(c.messages[req.fold].is_user());
        // A transcript with no user message at all is refused rather than
        // folded into a shape the provider will not accept.
        let mut none = ContextState::new(ContextKind::Conversation, 1000);
        for i in 0..6 {
            none.append(Msg::assistant(Some(format!("thought {i}")), vec![]));
        }
        assert!(plan_compaction(&none, 2, None).is_none());
    }

    /// Property-style: over every fold point the two callers can ask for, the
    /// kept window opens on a `User` and carries no orphaned tool result.
    #[test]
    fn every_fold_point_keeps_a_user_first_and_no_orphan_tool_result() {
        let c = mixed_ctx();
        let n = c.messages.len();
        for keep_last in 0..=n {
            for target in [None, Some(0), Some(1), Some(60), Some(10_000)] {
                let Some(req) = plan_compaction(&c, keep_last, target) else {
                    continue;
                };
                let kept = &c.messages[req.fold..];
                let Some(first) = kept.first() else {
                    continue; // keep_last 0: everything folds, nothing to lead
                };
                assert!(
                    first.is_user(),
                    "keep_last {keep_last} target {target:?} → fold {} left {first:?} first",
                    req.fold
                );
                // No kept tool result may reference a call that was folded away.
                let calls: Vec<&str> = kept
                    .iter()
                    .flat_map(|m| match m {
                        Msg::Assistant { tool_calls, .. } => tool_calls.as_slice(),
                        _ => &[],
                    })
                    .map(|tc| tc.id.as_str())
                    .collect();
                for m in kept {
                    if let Msg::Tool { id, .. } = m {
                        assert!(
                            calls.contains(&id.as_str()),
                            "keep_last {keep_last} target {target:?} → fold {} orphaned {id}",
                            req.fold
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn apply_absorbs_the_summary_bumps_version_and_recounts() {
        let mut c = ctx_with(10);
        c.plan = Some(super::super::plan::Plan::create("goal", &[json!("a")], 32).unwrap());
        c.load_skill("review", "h", 8).unwrap();
        let before = c.est_tokens;
        let req = plan_compaction(&c, 3, None).unwrap();
        let out = apply_compaction(
            &mut c,
            &req,
            &json!({"goals": ["finish"], "decisions": [], "open": ["q1"], "facts": ["n=7"]}),
        )
        .unwrap();
        assert_eq!(out.folded, 7);
        assert_eq!(out.version, 2);
        assert_eq!(c.messages.len(), 3);
        assert_eq!(c.summary.goals, vec!["finish".to_string()]);
        assert_eq!(c.summary.covers_messages, 7);
        assert!(c.est_tokens < before);
        assert!(c.plan.is_some(), "plan kept verbatim");
        assert_eq!(c.skills.len(), 1, "skill names kept");
        assert!(c.dirty);
        // Version drift is refused.
        let req2 = plan_compaction(&ctx_with(10), 3, None).unwrap();
        assert!(apply_compaction(&mut c, &req2, &json!({})).is_err());
        // Fallback path.
        let mut c2 = ctx_with(10);
        let req = plan_compaction(&c2, 3, None).unwrap();
        let out = apply_fallback(&mut c2, &req).unwrap();
        assert_eq!(out.folded, 7);
        assert!(
            c2.summary
                .narrative
                .as_deref()
                .unwrap()
                .contains("message number 0")
        );
        // A wire slice carries the summary + plan first.
        let wire = c.to_wire();
        assert!(
            matches!(&wire[0], crate::wire::intel::Message::System(s) if s.starts_with("Summary of earlier"))
        );
        assert!(
            matches!(&wire[1], crate::wire::intel::Message::System(s) if s.starts_with("Plan ("))
        );
    }
}