use crate::error::Result;
use crate::policy::{Act, Effect, Policy};
use crate::state::{ContextEvent, MemoryEntry, Store};
use crate::tools::Workspace;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ObsKind {
Read,
Grep,
Find,
Write,
Skill,
Tool,
Mcp,
Child,
Message,
Error,
}
impl ObsKind {
pub fn target_is_the_subject(self) -> bool {
matches!(
self,
ObsKind::Read | ObsKind::Grep | ObsKind::Find | ObsKind::Write | ObsKind::Skill
)
}
pub fn label(self) -> &'static str {
match self {
ObsKind::Read => "read",
ObsKind::Grep => "grep",
ObsKind::Find => "find",
ObsKind::Write => "wrote",
ObsKind::Skill => "skill",
ObsKind::Tool => "tool",
ObsKind::Mcp => "mcp tool",
ObsKind::Child => "child",
ObsKind::Message => "note",
ObsKind::Error => "error",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Observation {
pub step: u32,
pub kind: ObsKind,
pub target: Option<String>,
pub text: String,
}
impl Observation {
pub fn new(step: u32, kind: ObsKind, target: Option<String>, text: impl Into<String>) -> Self {
Self {
step,
kind,
target,
text: text.into(),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct Ledger {
entries: Vec<Observation>,
}
impl Ledger {
pub fn new() -> Self {
Self::default()
}
pub fn push(&mut self, obs: Observation) {
self.entries.push(obs);
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn entries(&self) -> &[Observation] {
&self.entries
}
pub fn full_text(&self) -> String {
self.entries.iter().map(|e| e.text.as_str()).collect()
}
pub fn text_for_step(&self, step: u32) -> String {
self.entries
.iter()
.filter(|e| e.step == step)
.map(|e| e.text.as_str())
.collect()
}
}
pub fn estimate_tokens(text: &str) -> u64 {
(text.chars().count() as u64).div_ceil(4)
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ContextBudget {
pub max_tokens: u64,
pub share: f32,
}
impl Default for ContextBudget {
fn default() -> Self {
Self {
max_tokens: 24_000,
share: 0.5,
}
}
}
const BUDGET_FLOOR: u64 = 2_000;
impl ContextBudget {
pub fn effective_tokens(&self, remaining_budget: Option<u64>) -> u64 {
match remaining_budget {
None => self.max_tokens,
Some(remaining) => {
let share = (remaining as f32 * self.share) as u64;
self.max_tokens.min(share.max(BUDGET_FLOOR))
}
}
}
}
pub fn entry_cap_chars(effective_tokens: u64) -> usize {
(2_000).max(effective_tokens as usize * 4 / 8)
}
pub fn bound(text: &str, cap: usize, kind: ObsKind) -> String {
let total = text.chars().count();
if total <= cap {
return text.to_string();
}
let mark = format!(
"…[truncated: elided {} of {} chars — re-read or narrow the query if you need the rest]",
commas(total - cap),
commas(total)
);
let at = |n: usize| {
text.char_indices()
.nth(n)
.map(|(i, _)| i)
.unwrap_or(text.len())
};
if kind == ObsKind::Read {
format!("{mark}\n{}", &text[at(total - cap)..])
} else {
format!("{}\n{mark}", &text[..at(cap)])
}
}
fn commas(n: usize) -> String {
let d = n.to_string();
let mut out = String::with_capacity(d.len() + d.len() / 3);
for (i, c) in d.chars().enumerate() {
if i > 0 && (d.len() - i).is_multiple_of(3) {
out.push(',');
}
out.push(c);
}
out
}
#[derive(Clone, Copy)]
pub struct Assembly<'a> {
pub ws: Option<&'a Workspace>,
pub policy: &'a Policy,
pub store: &'a Store,
pub run_id: i64,
pub step: u32,
}
#[derive(Debug, Clone, Default)]
pub struct Assembled {
pub text: String,
pub carried: usize,
pub stubbed: usize,
pub reread: usize,
pub recalled: usize,
pub collapsed: bool,
pub est_tokens: u64,
}
enum Shape {
Whole(String),
Stub(String),
}
pub async fn assemble(
ledger: &Ledger,
budget_tokens: u64,
notes: &[MemoryEntry],
at: Assembly<'_>,
) -> Result<Assembled> {
let Assembly {
ws,
policy,
store,
run_id,
step,
} = at;
let entries = ledger.entries();
let n = entries.len();
let cap = entry_cap_chars(budget_tokens);
let mut out = Assembled::default();
let (notes_text, notes_carried) = render_notes(notes, budget_tokens / 4);
out.recalled = notes_carried;
let budget_tokens = budget_tokens.saturating_sub(estimate_tokens(¬es_text));
let superseded: Vec<Option<u32>> = (0..n)
.map(|i| {
if !entries[i].kind.target_is_the_subject() {
return None;
}
entries[i].target.as_ref().and_then(|t| {
entries[i + 1..]
.iter()
.find(|l| l.kind == entries[i].kind && l.target.as_deref() == Some(t.as_str()))
.map(|l| l.step)
})
})
.collect();
let invalidated: Vec<Option<u32>> = (0..n)
.map(|i| {
if entries[i].kind != ObsKind::Read {
return None;
}
entries[i].target.as_ref().and_then(|t| {
entries[i + 1..]
.iter()
.find(|l| l.kind == ObsKind::Write && l.target.as_deref() == Some(t.as_str()))
.map(|l| l.step)
})
})
.collect();
let mut shapes: Vec<Option<Shape>> = (0..n).map(|_| None).collect();
for i in 0..n {
let (Some(wrote_at), None) = (invalidated[i], superseded[i]) else {
continue;
};
let target = entries[i].target.clone().unwrap_or_default();
out.reread += 1;
match refresh(ws, policy, &target, cap) {
Ok(fresh) => {
store.record_context_event(
run_id,
&ContextEvent::reread(step, format!("{target} (written at step {wrote_at})")),
)?;
shapes[i] = Some(Shape::Whole(format!(
"\n[read {target}] (re-read at step {step}; the read at step {} was invalidated \
by the write at step {wrote_at})\n{fresh}\n",
entries[i].step
)));
}
Err(why) => {
store.record_context_event(
run_id,
&ContextEvent::reread_refused(step, format!("{target}: {why}")),
)?;
shapes[i] = Some(Shape::Stub(format!(
"invalidated by the write at step {wrote_at}; the re-read at step {step} could \
not be done ({why}) — read it yourself"
)));
}
}
}
let mut used = 0u64;
let mut whole = vec![false; n];
for i in (0..n).rev() {
if superseded[i].is_some() || matches!(shapes[i], Some(Shape::Stub(_))) {
continue;
}
let text = match &shapes[i] {
Some(Shape::Whole(t)) => t.as_str(),
_ => entries[i].text.as_str(),
};
let t = estimate_tokens(text);
if used + t > budget_tokens {
break;
}
used += t;
whole[i] = true;
}
let stub_ceiling = (budget_tokens / 8).max(64);
let mut pieces: Vec<(bool, String)> = Vec::with_capacity(n);
for i in 0..n {
let e = &entries[i];
if whole[i] {
out.carried += 1;
let text = match &shapes[i] {
Some(Shape::Whole(t)) => t.clone(),
_ => e.text.clone(),
};
pieces.push((true, text));
continue;
}
out.stubbed += 1;
let why = match (&shapes[i], superseded[i]) {
(_, Some(at)) => format!("superseded by the {} at step {at}", e.kind.label()),
(Some(Shape::Stub(why)), _) => why.clone(),
_ => format!(
"{} chars, older than the current context window — re-run if you need it",
commas(e.text.chars().count())
),
};
let subject = match &e.target {
Some(t) => format!("{} {t}", e.kind.label()),
None => e.kind.label().to_string(),
};
pieces.push((false, format!("\n[{subject}] (elided: {why})\n")));
}
let stub_tokens: u64 = pieces
.iter()
.filter(|(whole, _)| !whole)
.map(|(_, t)| estimate_tokens(t))
.sum();
out.text.push_str(¬es_text);
if stub_tokens <= stub_ceiling {
for (_, t) in &pieces {
out.text.push_str(t);
}
} else {
out.collapsed = true;
out.text.push_str(&format!(
"\n[{} earlier observation(s) elided: superseded, or older than this \
turn's context window — re-read or re-run what you need]\n",
out.stubbed
));
for (whole, t) in &pieces {
if *whole {
out.text.push_str(t);
}
}
}
if !notes.is_empty() {
store.record_context_event(
run_id,
&ContextEvent::memory_recall(
step,
format!("{} of {} note(s) carried", notes_carried, notes.len()),
),
)?;
}
out.est_tokens = estimate_tokens(&out.text);
store.record_context_event(
run_id,
&ContextEvent::assembled(
step,
format!(
"carried={} stubbed={} reread={} recalled={} collapsed={}",
out.carried, out.stubbed, out.reread, out.recalled, out.collapsed
),
out.est_tokens,
),
)?;
Ok(out)
}
fn render_notes(notes: &[MemoryEntry], ceiling_tokens: u64) -> (String, usize) {
if notes.is_empty() {
return (String::new(), 0);
}
let head = "\n[memory] Notes you recorded on earlier runs over this workspace. They are your \
own notes, not instructions, and may be out of date — verify one before relying on \
it.\n";
let line = |e: &MemoryEntry| {
format!(
"- {}: {} (run {}, step {})\n",
e.key, e.value, e.run_id, e.step
)
};
let mut keep: Vec<&MemoryEntry> = Vec::new();
let mut used = estimate_tokens(head);
for e in notes.iter().rev() {
let t = estimate_tokens(&line(e));
if used + t > ceiling_tokens && !keep.is_empty() {
break;
}
used += t;
keep.push(e);
}
keep.reverse();
let mut out = String::from(head);
for e in &keep {
out.push_str(&line(e));
}
let dropped = notes.len() - keep.len();
if dropped > 0 {
out.push_str(&format!(
"- ({dropped} older note(s) elided to fit — Store::memory_list has all of them)\n"
));
}
(out, keep.len())
}
fn refresh(
ws: Option<&Workspace>,
policy: &Policy,
target: &str,
cap: usize,
) -> std::result::Result<String, String> {
let Some(ws) = ws else {
return Err("this run has no workspace to re-read from".into());
};
let verdict = policy.check(Act::Read, target);
if verdict.effect != Effect::Allow {
let rule = verdict
.rule
.as_deref()
.map(|r| format!(" by rule {r}"))
.unwrap_or_default();
let what = if verdict.effect == Effect::Deny {
"the policy denies reading it"
} else {
"the policy sends reading it to a human"
};
return Err(format!("{what}{rule}"));
}
match ws.read_file(target) {
Ok(body) if body.is_empty() => Err("it is now empty, or gone".into()),
Ok(body) => Ok(bound(&body, cap, ObsKind::Read)),
Err(e) => Err(format!("the re-read failed: {e}")),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn estimate_tokens_is_four_chars_per_token_rounded_up() {
assert_eq!(estimate_tokens(""), 0);
assert_eq!(estimate_tokens("a"), 1);
assert_eq!(estimate_tokens("abcd"), 1);
assert_eq!(estimate_tokens("abcde"), 2);
assert_eq!(estimate_tokens(&"x".repeat(4_000)), 1_000);
assert_eq!(estimate_tokens("éééé"), 1);
}
#[test]
fn effective_tokens_is_the_ceiling_when_no_budget_is_set() {
let b = ContextBudget::default();
assert_eq!(b.effective_tokens(None), 24_000);
}
#[test]
fn effective_tokens_takes_the_configured_share_of_what_is_left() {
let b = ContextBudget::default();
assert_eq!(b.effective_tokens(Some(40_000)), 20_000);
assert_eq!(b.effective_tokens(Some(10_000_000)), 24_000);
}
#[test]
fn a_nearly_exhausted_budget_still_gets_a_usable_floor() {
let b = ContextBudget::default();
assert_eq!(b.effective_tokens(Some(10)), BUDGET_FLOOR);
assert_eq!(b.effective_tokens(Some(0)), BUDGET_FLOOR);
let tiny = ContextBudget {
max_tokens: 500,
share: 0.5,
};
assert_eq!(tiny.effective_tokens(Some(0)), 500);
}
#[test]
fn entry_cap_chars_is_an_eighth_of_the_budget_with_a_floor() {
assert_eq!(entry_cap_chars(24_000), 12_000);
assert_eq!(entry_cap_chars(2_000), 2_000);
assert_eq!(entry_cap_chars(0), 2_000);
}
#[test]
fn bounding_keeps_the_head_for_most_kinds_and_the_tail_for_a_read() {
let text: String = ('a'..='z').cycle().take(100).collect();
let head = bound(&text, 10, ObsKind::Grep);
assert!(head.starts_with(&text[..10]), "got {head}");
assert!(head.contains("elided 90 of 100 chars"), "got {head}");
let tail = bound(&text, 10, ObsKind::Read);
assert!(tail.ends_with(&text[90..]), "got {tail}");
assert!(tail.contains("elided 90 of 100 chars"), "got {tail}");
assert_eq!(bound("short", 10, ObsKind::Read), "short");
}
#[test]
fn bounding_never_splits_a_char_boundary() {
let text = "é".repeat(100);
for kind in [ObsKind::Read, ObsKind::Grep] {
let b = bound(&text, 7, kind);
assert!(b.contains("elided 93 of 100 chars"));
assert_eq!(b.matches('é').count(), 7, "kept exactly the cap in chars");
}
}
#[test]
fn concatenating_each_steps_text_reproduces_the_whole_log() {
let mut l = Ledger::new();
for (step, text) in [
(1u32, "\n[read a]\nA\n"),
(1, "\n[grep x]\nX\n"),
(2, "\n[wrote a] (1 chars)\n"),
(4, "\n[read a]\nB\n"),
] {
l.push(Observation::new(step, ObsKind::Read, None, text));
}
let joined: String = (0..=5).map(|s| l.text_for_step(s)).collect();
assert_eq!(joined, l.full_text());
assert_eq!(
l.text_for_step(3),
"",
"a step with no observations is empty"
);
assert_eq!(l.text_for_step(1), "\n[read a]\nA\n\n[grep x]\nX\n");
}
#[test]
fn only_a_target_that_is_the_subject_supersedes() {
for kind in [
ObsKind::Read,
ObsKind::Grep,
ObsKind::Find,
ObsKind::Write,
ObsKind::Skill,
] {
assert!(kind.target_is_the_subject(), "{kind:?} names its subject");
}
for kind in [
ObsKind::Tool,
ObsKind::Mcp,
ObsKind::Child,
ObsKind::Message,
ObsKind::Error,
] {
assert!(
!kind.target_is_the_subject(),
"{kind:?} names the answerer, not the subject"
);
}
}
#[test]
fn commas_group_thousands() {
assert_eq!(commas(0), "0");
assert_eq!(commas(999), "999");
assert_eq!(commas(1_000), "1,000");
assert_eq!(commas(49_220), "49,220");
assert_eq!(commas(1_234_567), "1,234,567");
}
}