use super::{ContextState, Msg, Summary};
use crate::state::now_ms;
use serde_json::{Value, json};
#[derive(Debug, Clone)]
pub struct CompactionRequest {
pub fold: usize,
pub system: String,
pub input: String,
pub output_schema: Value,
pub version: u64,
}
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.";
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 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;
}
}
while fold > 0 && ctx.messages.get(fold).is_some_and(|m| !m.is_user()) {
fold -= 1;
}
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())
)
}
}
}
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,
}
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![]));
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"
);
let big = ctx_with(20);
let req = plan_compaction(&big, 10, Some(1)).unwrap();
assert_eq!(req.fold, 18);
}
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
}
#[test]
fn fold_never_leaves_an_assistant_or_a_tool_result_first() {
let c = mixed_ctx();
let req = plan_compaction(&c, 1, None).unwrap();
assert_eq!(req.fold, 6);
assert!(c.messages[req.fold].is_user());
assert_eq!(plan_compaction(&c, 5, None).unwrap().fold, 6);
let req = plan_compaction(&c, 7, None).unwrap();
assert_eq!(req.fold, 2);
assert!(c.messages[req.fold].is_user());
let req = plan_compaction(&c, 2, Some(1)).unwrap();
assert!(c.messages[req.fold].is_user());
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());
}
#[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; };
assert!(
first.is_user(),
"keep_last {keep_last} target {target:?} → fold {} left {first:?} first",
req.fold
);
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);
let req2 = plan_compaction(&ctx_with(10), 3, None).unwrap();
assert!(apply_compaction(&mut c, &req2, &json!({})).is_err());
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")
);
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 ("))
);
}
}