use async_trait::async_trait;
use crate::agentic::compress::is_compaction_text;
use crate::metrics::TurnMetrics;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemMessage {
pub role: Role,
pub content: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Role {
System,
User,
Assistant,
Tool,
}
impl Role {
pub fn as_str(&self) -> &'static str {
match self {
Self::System => "system",
Self::User => "user",
Self::Assistant => "assistant",
Self::Tool => "tool",
}
}
}
impl MemMessage {
pub fn system(content: impl Into<String>) -> Self {
Self {
role: Role::System,
content: content.into(),
}
}
pub fn user(content: impl Into<String>) -> Self {
Self {
role: Role::User,
content: content.into(),
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: Role::Assistant,
content: content.into(),
}
}
}
#[derive(Debug, Clone)]
pub struct SessionContext {
pub workspace: String,
pub session_id: String,
}
#[async_trait]
pub trait MemoryProvider: Send + Sync {
fn name(&self) -> &str;
async fn initialize(&mut self, _ctx: &SessionContext) -> anyhow::Result<()> {
Ok(())
}
fn system_prompt_block(&self) -> Option<String> {
None
}
async fn prefetch(&self, _query: &str) -> anyhow::Result<String> {
Ok(String::new())
}
fn build_messages(&self, system_prompt: &str, new_task: &str) -> Vec<MemMessage>;
async fn sync_turn(&mut self, user: &str, assistant: &str, metrics: &TurnMetrics);
fn reset(&mut self) {}
fn restore_turns(&mut self, _turns: &[crate::ConversationTurn]) {}
fn take_compaction_record(&mut self) -> Option<String> {
None
}
async fn on_pre_compress(&self, _messages: &[MemMessage]) -> String {
String::new()
}
async fn on_session_end(&mut self, _messages: &[MemMessage]) {}
fn usage(&self) -> Option<(String, usize, usize)> {
None
}
fn add_note(&mut self, _fact: &str) -> anyhow::Result<()> {
Err(NotesUnsupported.into())
}
fn replace_note(&mut self, _old_substring: &str, _new_text: &str) -> anyhow::Result<()> {
Err(NotesUnsupported.into())
}
fn remove_note(&mut self, _substring: &str) -> anyhow::Result<()> {
Err(NotesUnsupported.into())
}
}
#[derive(Debug, thiserror::Error)]
#[error("this memory provider does not support persistent notes")]
pub struct NotesUnsupported;
pub struct MemoryManager {
providers: Vec<Box<dyn MemoryProvider>>,
}
impl MemoryManager {
pub fn new() -> Self {
Self {
providers: Vec::new(),
}
}
pub fn add_provider(&mut self, p: impl MemoryProvider + 'static) {
self.providers.push(Box::new(p));
}
pub async fn initialize_all(&mut self, ctx: &SessionContext) {
for p in &mut self.providers {
if let Err(e) = p.initialize(ctx).await {
tracing::warn!(provider = p.name(), error = %e, "memory provider init failed");
}
}
}
pub fn build_system_prompt_additions(&self) -> String {
self.providers
.iter()
.filter_map(|p| p.system_prompt_block())
.collect::<Vec<_>>()
.join("\n\n")
}
pub async fn prefetch_all(&self, query: &str) -> String {
let mut parts = Vec::new();
for p in &self.providers {
match p.prefetch(query).await {
Ok(s) if !s.is_empty() => parts.push(s),
Err(e) => {
tracing::warn!(provider = p.name(), error = %e, "prefetch failed");
}
_ => {}
}
}
parts.join("\n\n")
}
pub fn build_messages(&self, system_prompt: &str, new_task: &str) -> Vec<MemMessage> {
for p in &self.providers {
let msgs = p.build_messages(system_prompt, new_task);
if !msgs.is_empty() {
return msgs;
}
}
vec![
MemMessage::system(system_prompt),
MemMessage::user(new_task),
]
}
pub async fn sync_all(&mut self, user: &str, assistant: &str, metrics: &TurnMetrics) {
for p in &mut self.providers {
p.sync_turn(user, assistant, metrics).await;
}
}
pub fn reset_all(&mut self) {
for p in &mut self.providers {
p.reset();
}
}
pub fn restore_turns(&mut self, turns: &[crate::ConversationTurn]) {
for p in &mut self.providers {
p.restore_turns(turns);
}
}
pub fn take_compaction_record(&mut self) -> Option<String> {
self.providers
.iter_mut()
.find_map(|p| p.take_compaction_record())
}
pub async fn on_pre_compress(&self, messages: &[MemMessage]) -> String {
let mut parts = Vec::new();
for p in &self.providers {
let s = p.on_pre_compress(messages).await;
if !s.is_empty() {
parts.push(s);
}
}
parts.join("\n\n")
}
pub async fn on_session_end(&mut self, messages: &[MemMessage]) {
for p in &mut self.providers {
p.on_session_end(messages).await;
}
}
pub fn usage(&self) -> Vec<(String, usize, usize)> {
self.providers.iter().filter_map(|p| p.usage()).collect()
}
pub fn add_note(&mut self, fact: &str) -> anyhow::Result<()> {
self.route_note_op(|p| p.add_note(fact))
}
pub fn replace_note(&mut self, old_substring: &str, new_text: &str) -> anyhow::Result<()> {
self.route_note_op(|p| p.replace_note(old_substring, new_text))
}
pub fn remove_note(&mut self, substring: &str) -> anyhow::Result<()> {
self.route_note_op(|p| p.remove_note(substring))
}
fn route_note_op(
&mut self,
mut op: impl FnMut(&mut dyn MemoryProvider) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
let mut first_err: Option<anyhow::Error> = None;
for p in &mut self.providers {
match op(p.as_mut()) {
Ok(()) => return Ok(()),
Err(e) if e.is::<NotesUnsupported>() => continue,
Err(e) => {
first_err.get_or_insert(e);
}
}
}
match first_err {
Some(e) => Err(e),
None => anyhow::bail!(
"no note-capable memory provider registered — add [memory] provider = \"note_store\" to newt.toml"
),
}
}
}
impl Default for MemoryManager {
fn default() -> Self {
Self::new()
}
}
pub struct RollingWindow {
max_turns: usize,
history: Vec<(String, String)>,
}
impl RollingWindow {
pub fn new(max_turns: usize) -> Self {
Self {
max_turns: max_turns.max(1),
history: Vec::new(),
}
}
pub fn from_config() -> Self {
let window = newt_core_memory_window();
Self::new(window)
}
pub fn turn_count(&self) -> usize {
self.history.len()
}
}
fn newt_core_memory_window() -> usize {
crate::Config::resolve()
.ok()
.and_then(|c| c.memory)
.map(|m| m.window)
.unwrap_or(20)
}
#[async_trait]
impl MemoryProvider for RollingWindow {
fn name(&self) -> &str {
"rolling_window"
}
fn build_messages(&self, system_prompt: &str, new_task: &str) -> Vec<MemMessage> {
let mut msgs = vec![MemMessage::system(system_prompt)];
let start = self.history.len().saturating_sub(self.max_turns);
for (user, asst) in &self.history[start..] {
msgs.push(MemMessage::user(user));
if !asst.is_empty() {
msgs.push(MemMessage::assistant(asst));
}
}
msgs.push(MemMessage::user(new_task));
msgs
}
async fn sync_turn(&mut self, user: &str, assistant: &str, _metrics: &TurnMetrics) {
self.history.push((user.to_string(), assistant.to_string()));
if self.history.len() > self.max_turns * 2 {
let drain_to = self.history.len() - self.max_turns;
self.history.drain(..drain_to);
}
}
fn reset(&mut self) {
self.history.clear();
}
fn restore_turns(&mut self, turns: &[crate::ConversationTurn]) {
self.history = turns
.iter()
.map(|t| (t.user.clone(), t.assistant.clone()))
.collect();
}
fn usage(&self) -> Option<(String, usize, usize)> {
Some((
"turns".into(),
self.history.len().min(self.max_turns),
self.max_turns,
))
}
}
pub const DEFAULT_CONTEXT_TOKENS: u32 = 8_192;
#[derive(Debug, Clone)]
struct TurnRecord {
user: String,
assistant: String,
est_tokens: u32,
}
fn turn_content_estimate(user: &str, assistant: &str, est: crate::tokens::TokenEstimation) -> u32 {
(est.tokens_for_chars(user.len()) + est.tokens_for_chars(assistant.len())) as u32
}
fn restored_token_anchor(
turns: &[crate::ConversationTurn],
est: crate::tokens::TokenEstimation,
) -> (Option<u32>, i64) {
let Some(pos) = turns.iter().rposition(|t| t.tokens_in.is_some()) else {
return (None, 0);
};
let mut delta = est.tokens_for_chars(turns[pos].assistant.len()) as i64;
for t in &turns[pos + 1..] {
delta += i64::from(turn_content_estimate(&t.user, &t.assistant, est));
}
(turns[pos].tokens_in, delta)
}
pub struct TokenBudget {
max_tokens: u32,
threshold_pct: f32,
history: Vec<TurnRecord>,
pruned_count: usize,
last_prompt_tokens: Option<u32>,
delta_since_prompt: i64,
est: crate::tokens::TokenEstimation,
}
impl TokenBudget {
pub fn new(max_tokens: u32, threshold_pct: f32) -> Self {
Self {
max_tokens: max_tokens.max(512),
threshold_pct: threshold_pct.clamp(0.1, 0.99),
history: Vec::new(),
pruned_count: 0,
last_prompt_tokens: None,
delta_since_prompt: 0,
est: crate::tokens::TokenEstimation::default(),
}
}
#[must_use]
pub fn with_estimation(mut self, est: crate::tokens::TokenEstimation) -> Self {
self.est = est;
self
}
pub fn with_budget(mut self, tokens: u32) -> Self {
self.max_tokens = tokens.max(512);
self
}
fn budget_tokens(&self) -> u32 {
(self.max_tokens as f32 * self.threshold_pct) as u32
}
fn used_tokens(&self) -> u32 {
match self.last_prompt_tokens {
Some(p) => (i64::from(p) + self.delta_since_prompt).max(0) as u32,
None => self.history.iter().map(|r| r.est_tokens).sum(),
}
}
fn prune_to_budget(&mut self) -> usize {
let budget = self.budget_tokens();
let mut dropped = 0;
while self.used_tokens() > budget && !self.history.is_empty() {
let removed = self.history.remove(0);
self.delta_since_prompt -= i64::from(removed.est_tokens);
dropped += 1;
}
dropped
}
}
#[async_trait]
impl MemoryProvider for TokenBudget {
fn name(&self) -> &str {
"token_budget"
}
fn build_messages(&self, system_prompt: &str, new_task: &str) -> Vec<MemMessage> {
let mut msgs = vec![MemMessage::system(system_prompt)];
for r in &self.history {
msgs.push(MemMessage::user(&r.user));
if !r.assistant.is_empty() {
msgs.push(MemMessage::assistant(&r.assistant));
}
}
msgs.push(MemMessage::user(new_task));
msgs
}
async fn sync_turn(&mut self, user: &str, assistant: &str, metrics: &TurnMetrics) {
let est = turn_content_estimate(user, assistant, self.est);
self.history.push(TurnRecord {
user: user.to_string(),
assistant: assistant.to_string(),
est_tokens: est,
});
match metrics.usage {
Some(u) => {
self.last_prompt_tokens = Some(u.input_tokens);
self.delta_since_prompt = self.est.tokens_for_chars(assistant.len()) as i64;
}
None => self.delta_since_prompt += i64::from(est),
}
let dropped = self.prune_to_budget();
self.pruned_count += dropped;
if dropped > 0 {
tracing::info!(
dropped,
budget = self.budget_tokens(),
used = self.used_tokens(),
"TokenBudget pruned old turns"
);
}
}
fn reset(&mut self) {
self.history.clear();
self.pruned_count = 0;
self.last_prompt_tokens = None;
self.delta_since_prompt = 0;
}
fn restore_turns(&mut self, turns: &[crate::ConversationTurn]) {
self.history = turns
.iter()
.map(|t| TurnRecord {
user: t.user.clone(),
assistant: t.assistant.clone(),
est_tokens: turn_content_estimate(&t.user, &t.assistant, self.est),
})
.collect();
let (anchor, delta) = restored_token_anchor(turns, self.est);
self.last_prompt_tokens = anchor;
self.delta_since_prompt = delta;
self.pruned_count = self.prune_to_budget();
}
fn usage(&self) -> Option<(String, usize, usize)> {
Some((
"tokens".into(),
self.used_tokens() as usize,
self.budget_tokens() as usize,
))
}
}
pub use crate::notes::NoteStore;
pub const MEMORY_INDEX_BUDGET: usize = 12;
pub struct MemoryIndex {
rows: Vec<(usize, String)>,
truncated: bool,
notes_path: std::path::PathBuf,
}
impl MemoryIndex {
pub fn new(notes_path: impl Into<std::path::PathBuf>) -> Self {
Self {
rows: Vec::new(),
truncated: false,
notes_path: notes_path.into(),
}
}
pub fn default_path() -> Self {
let path = crate::Config::user_config_path()
.map(|p| p.with_file_name("NOTES.md"))
.unwrap_or_else(|| std::path::PathBuf::from("NOTES.md"));
Self::new(path)
}
pub fn rows(&self) -> &[(usize, String)] {
&self.rows
}
fn capture_from(&mut self, notes: &NoteStore) {
let all = notes.index_entries();
self.truncated = all.len() > MEMORY_INDEX_BUDGET;
self.rows = all
.into_iter()
.take(MEMORY_INDEX_BUDGET)
.map(|(id, title)| (id, title.to_string()))
.collect();
}
}
#[async_trait]
impl MemoryProvider for MemoryIndex {
fn name(&self) -> &str {
"memory_index"
}
async fn initialize(&mut self, ctx: &SessionContext) -> anyhow::Result<()> {
let mut notes = NoteStore::new(self.notes_path.clone(), NoteStore::DEFAULT_CHAR_LIMIT);
notes.initialize(ctx).await?;
self.capture_from(¬es);
Ok(())
}
fn system_prompt_block(&self) -> Option<String> {
if self.rows.is_empty() {
return None;
}
let mut block =
String::from("## Memory index (call `memory_fetch` with an id to read a body)\n");
for (id, title) in &self.rows {
let title = if title.is_empty() {
"(untitled)"
} else {
title
};
block.push_str(&format!("- note:{id} {title}\n"));
}
if self.truncated {
block
.push_str("(more notes exist than are listed — use `recall` to search the rest)\n");
}
Some(block.trim_end().to_string())
}
fn build_messages(&self, _system_prompt: &str, _new_task: &str) -> Vec<MemMessage> {
Vec::new()
}
async fn sync_turn(&mut self, _user: &str, _assistant: &str, _metrics: &TurnMetrics) {}
}
#[derive(Clone)]
struct SumTurn {
user: String,
assistant: String,
est_tokens: u32,
}
impl SumTurn {
fn new(
user: impl Into<String>,
assistant: impl Into<String>,
est: crate::tokens::TokenEstimation,
) -> Self {
let user = user.into();
let assistant = assistant.into();
let est_tokens = turn_content_estimate(&user, &assistant, est);
Self {
user,
assistant,
est_tokens,
}
}
fn to_wire(&self) -> Vec<serde_json::Value> {
let mut v = Vec::with_capacity(2);
if !self.user.is_empty() {
v.push(serde_json::json!({"role": "user", "content": self.user}));
}
if !self.assistant.is_empty() {
v.push(serde_json::json!({"role": "assistant", "content": self.assistant}));
}
v
}
}
fn wire_to_history(
messages: &[serde_json::Value],
est: crate::tokens::TokenEstimation,
) -> Vec<SumTurn> {
let mut out: Vec<SumTurn> = Vec::new();
for m in messages {
let content = m["content"].as_str().unwrap_or_default();
match m["role"].as_str() {
Some("user") => out.push(SumTurn::new(content, "", est)),
Some("assistant") => {
match out.last_mut() {
Some(last)
if last.assistant.is_empty()
&& !last.user.is_empty()
&& !is_compaction_text(&last.user) =>
{
last.assistant = content.to_string();
last.est_tokens = turn_content_estimate(&last.user, &last.assistant, est);
}
_ => out.push(SumTurn::new("", content, est)),
}
}
_ => {}
}
}
out
}
pub struct Summarizing {
max_tokens: u32,
threshold_pct: f32,
history: Vec<SumTurn>,
prev_summary: String,
compress_count: usize,
state: crate::agentic::CompressState,
summarizer: Option<crate::agentic::Summarizer>,
last_prompt_tokens: Option<u32>,
delta_since_prompt: i64,
pending_record: Option<String>,
est: crate::tokens::TokenEstimation,
summary_input_cap_floor_chars: usize,
}
impl Summarizing {
pub fn new(max_tokens: u32) -> Self {
Self {
max_tokens: max_tokens.max(1),
threshold_pct: 0.80,
history: Vec::new(),
prev_summary: String::new(),
compress_count: 0,
state: crate::agentic::CompressState::new(),
summarizer: None,
last_prompt_tokens: None,
delta_since_prompt: 0,
pending_record: None,
est: crate::tokens::TokenEstimation::default(),
summary_input_cap_floor_chars: 8_192,
}
}
#[must_use]
pub fn with_estimation(
mut self,
est: crate::tokens::TokenEstimation,
summary_input_cap_floor_chars: usize,
) -> Self {
self.est = est;
self.summary_input_cap_floor_chars = summary_input_cap_floor_chars;
self
}
pub fn with_budget(mut self, tokens: u32) -> Self {
self.max_tokens = tokens.max(1);
self
}
pub fn with_summarizer(
mut self,
f: impl Fn(String) -> crate::agentic::SummarizeFuture + Send + Sync + 'static,
) -> Self {
self.summarizer = Some(Box::new(f));
self
}
fn budget(&self) -> u32 {
(self.max_tokens as f32 * self.threshold_pct) as u32
}
fn used_tokens(&self) -> u32 {
match self.last_prompt_tokens {
Some(p) => (i64::from(p) + self.delta_since_prompt).max(0) as u32,
None => self.history.iter().map(|t| t.est_tokens).sum(),
}
}
async fn compress_via_pipeline(&mut self) {
use crate::agentic::compress::{compress, CompressAction, CompressRequest};
if self.history.is_empty() {
return;
}
let messages: Vec<serde_json::Value> =
self.history.iter().flat_map(SumTurn::to_wire).collect();
let task = self
.history
.iter()
.find(|t| !t.user.is_empty() && !is_compaction_text(&t.user))
.map(|t| t.user.clone())
.unwrap_or_default();
let outcome = compress(
CompressRequest {
messages: &messages,
budget: self.budget() as usize,
max_messages: None,
task: &task,
hard_budget: true,
authoritative: true,
focus: None,
est: self.est,
summary_input_cap_floor_chars: self.summary_input_cap_floor_chars,
compaction_store: None,
},
self.summarizer.as_deref(),
&mut self.state,
)
.await;
if !outcome.fired {
return;
}
self.history = wire_to_history(&outcome.messages, self.est);
self.delta_since_prompt += outcome.tokens_after as i64 - outcome.tokens_before as i64;
if matches!(
outcome.action,
CompressAction::Summarized | CompressAction::StaticFallback
) {
if let Some(c) = self.history.iter().find(|t| is_compaction_text(&t.user)) {
self.prev_summary = c.user.clone();
self.pending_record = Some(c.user.clone());
self.compress_count += 1;
}
}
tracing::info!(
compress_count = self.compress_count,
action = outcome.action.describe(),
tokens_before = outcome.tokens_before,
tokens_after = outcome.tokens_after,
"Summarizing: compressed context via shared pipeline"
);
}
}
#[async_trait]
impl MemoryProvider for Summarizing {
fn name(&self) -> &str {
"summarizing"
}
fn build_messages(&self, system_prompt: &str, new_task: &str) -> Vec<MemMessage> {
let mut msgs = vec![MemMessage::system(system_prompt)];
for t in &self.history {
if !t.user.is_empty() {
msgs.push(MemMessage::user(&t.user));
}
if !t.assistant.is_empty() {
msgs.push(MemMessage::assistant(&t.assistant));
}
}
msgs.push(MemMessage::user(new_task));
msgs
}
async fn sync_turn(&mut self, user: &str, assistant: &str, metrics: &TurnMetrics) {
let est = turn_content_estimate(user, assistant, self.est);
self.history.push(SumTurn::new(user, assistant, self.est));
match metrics.usage {
Some(u) => {
self.last_prompt_tokens = Some(u.input_tokens);
self.delta_since_prompt = self.est.tokens_for_chars(assistant.len()) as i64;
}
None => self.delta_since_prompt += i64::from(est),
}
if self.used_tokens() > self.budget() && !self.state.is_disabled() {
self.compress_via_pipeline().await;
}
}
fn reset(&mut self) {
self.history.clear();
self.prev_summary.clear();
self.compress_count = 0;
self.state.reset();
self.last_prompt_tokens = None;
self.delta_since_prompt = 0;
self.pending_record = None;
}
fn restore_turns(&mut self, turns: &[crate::ConversationTurn]) {
let cut = turns
.iter()
.rposition(|t| is_compaction_text(&t.user) && t.assistant.is_empty());
let live = match cut {
Some(k) => &turns[k + 1..],
None => turns,
};
self.history.clear();
self.prev_summary.clear();
if let Some(k) = cut {
self.history
.push(SumTurn::new(turns[k].user.clone(), "", self.est));
self.prev_summary = turns[k].user.clone();
}
self.history.extend(
live.iter()
.map(|t| SumTurn::new(&*t.user, &*t.assistant, self.est)),
);
self.compress_count = 0;
self.pending_record = None;
self.state.reset();
let (anchor, delta) = restored_token_anchor(live, self.est);
self.last_prompt_tokens = anchor;
self.delta_since_prompt = delta;
}
fn take_compaction_record(&mut self) -> Option<String> {
self.pending_record.take()
}
async fn on_pre_compress(&self, _messages: &[MemMessage]) -> String {
if self.prev_summary.is_empty() {
String::new()
} else {
format!("Previous compression summary:\n{}", self.prev_summary)
}
}
fn usage(&self) -> Option<(String, usize, usize)> {
Some((
"tokens".into(),
self.used_tokens() as usize,
self.budget() as usize,
))
}
}
pub const DEFAULT_SOUL: &str = "\
You are newt, a small, fast, local-first agentic coder. \
Be concise and direct. \
You have tools: run_command, read_file, write_file, edit_file, list_dir, find, use_skill, web_fetch. \
Use them to actually complete tasks rather than describing what to do.\n\
\n\
## How to work\n\
\n\
**One change at a time.** Read only the files you need for the immediate next step. \
Make the change. Commit it. Then move to the next step. \
Never accumulate multiple uncommitted edits — a committed partial result survives a crash; \
an uncommitted complete result does not.\n\
\n\
**Read minimum, act fast.** Resist reading the entire codebase before acting. \
Read the specific file or function you are about to change, make the change, \
commit, then read the next thing. The session has a finite context window — \
every token spent reading is a token not spent writing.\n\
\n\
**Prefer edit_file over write_file.** For any existing file, use edit_file \
to replace a specific string — you only generate the change, not the whole file. \
Only use write_file when creating a new file or when you have generated the \
complete contents in full. write_file will refuse if the new content is \
significantly shorter than the original.\n\
\n\
**Stop when blocked.** If the same tool call fails twice in a row with the \
same error, stop immediately and tell the user what blocked you and why. \
Do not try alternative installation methods, do not loop, do not pivot to \
answering a different question. Two identical failures are a signal to report, \
not to retry. One sentence explaining the block is worth more than ten more \
failed tool calls.\n\
\n\
**Seek ground truth.** After every action, verify what actually happened — \
not what you intended. Do not proceed on assumptions about your own actions; \
confirm them. After writing a file, the tool reports the new line count — \
check it matches what you expected. After editing code, the tool reports \
whether it compiled — if it did not, fix the error before committing. \
Before committing, confirm you are on the right branch. \
A belief that something worked is worthless; a tool result that confirms it \
is ground truth.\n\
\n\
**Never describe a code change — make it.** Do not paste code into the chat. \
If the task requires a code change, call edit_file or write_file immediately. \
A markdown code block in the conversation is invisible to the filesystem — \
it does not modify any file. Write the code once, into the file, via the tool. \
Showing code in text is NOT completing a task; calling the tool IS.\n\
\n\
**Exploration budget.** Treat read-only rounds (list_dir, read_file) as expensive. \
Spend at most three consecutive rounds on exploration before making a write. \
Once you have read the file you need, stop reading and call edit_file or write_file. \
Continued reading without writing means you are lost — make your best attempt \
at the change based on what you have already read, then verify.\n\
\n\
**Working code first, then the three Cs.** Make it work, then make it right. \
Shipping a working result that hardcodes a list or a constant to get there is \
fine; functional results come first. Then RETURN to the three Cs: lift hardcoded \
knowledge (keyword lists, magic values, language or domain rules) into pure DATA \
that is Composed, Configured, and Convention-driven, so a new case is config, \
not code. Don't let this block shipping; do circle back and de-hardcode once it \
works.";
pub struct SoulProvider {
soul: String,
pub source: SoulSource,
override_path: Option<std::path::PathBuf>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SoulSource {
Default,
Global,
Workspace,
Explicit(std::path::PathBuf),
}
impl std::fmt::Display for SoulSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Default => write!(f, "built-in default"),
Self::Global => write!(f, "~/.newt/soul.md"),
Self::Workspace => write!(f, ".newt/soul.md"),
Self::Explicit(p) => write!(f, "{}", p.display()),
}
}
}
impl SoulProvider {
pub fn new(override_path: Option<std::path::PathBuf>) -> Self {
Self {
soul: DEFAULT_SOUL.to_string(),
source: SoulSource::Default,
override_path,
}
}
pub fn from_config() -> Self {
let override_path = crate::Config::resolve()
.ok()
.and_then(|c| c.memory)
.and_then(|m| m.soul_file)
.map(std::path::PathBuf::from);
Self::new(override_path)
}
fn try_load(path: &std::path::Path) -> Option<String> {
let text = std::fs::read_to_string(path).ok()?;
let trimmed = text.trim().to_string();
if trimmed.is_empty() {
None
} else {
Some(trimmed)
}
}
pub fn load(&mut self, workspace: &str) {
if let Some(ref p) = self.override_path {
if let Some(text) = Self::try_load(p) {
self.soul = text;
self.source = SoulSource::Explicit(p.clone());
return;
}
}
let ws_soul = std::path::Path::new(workspace)
.join(".newt")
.join("soul.md");
if let Some(text) = Self::try_load(&ws_soul) {
self.soul = text;
self.source = SoulSource::Workspace;
return;
}
if let Some(global) = crate::Config::user_config_path().map(|p| p.with_file_name("soul.md"))
{
if let Some(text) = Self::try_load(&global) {
self.soul = text;
self.source = SoulSource::Global;
}
}
}
}
#[async_trait]
impl MemoryProvider for SoulProvider {
fn name(&self) -> &str {
"soul"
}
async fn initialize(&mut self, ctx: &SessionContext) -> anyhow::Result<()> {
self.load(&ctx.workspace);
tracing::info!(source = %self.source, "soul loaded");
Ok(())
}
fn system_prompt_block(&self) -> Option<String> {
Some(self.soul.clone())
}
fn build_messages(&self, _system_prompt: &str, _new_task: &str) -> Vec<MemMessage> {
Vec::new()
}
async fn sync_turn(&mut self, _user: &str, _assistant: &str, _metrics: &TurnMetrics) {}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metrics::TokenUsage;
fn dummy_metrics() -> TurnMetrics {
TurnMetrics {
elapsed_ms: 100,
usage: Some(TokenUsage {
input_tokens: 10,
output_tokens: 5,
}),
cost_usd: Some(0.0),
model_id: "test".into(),
endpoint: "http://localhost".into(),
..Default::default()
}
}
#[tokio::test]
async fn rolling_window_empty_produces_two_messages() {
let rw = RollingWindow::new(5);
let msgs = rw.build_messages("sys", "hello");
assert_eq!(msgs.len(), 2);
assert_eq!(msgs[0].role, Role::System);
assert_eq!(msgs[1].role, Role::User);
assert_eq!(msgs[1].content, "hello");
}
#[tokio::test]
async fn rolling_window_includes_history() {
let mut rw = RollingWindow::new(5);
rw.sync_turn("q1", "a1", &dummy_metrics()).await;
rw.sync_turn("q2", "a2", &dummy_metrics()).await;
let msgs = rw.build_messages("sys", "q3");
assert_eq!(msgs.len(), 6);
assert_eq!(msgs[1].content, "q1");
assert_eq!(msgs[2].content, "a1");
assert_eq!(msgs[5].content, "q3");
}
#[tokio::test]
async fn rolling_window_caps_at_max_turns() {
let mut rw = RollingWindow::new(2);
for i in 0..5u32 {
rw.sync_turn(&format!("q{i}"), &format!("a{i}"), &dummy_metrics())
.await;
}
let msgs = rw.build_messages("sys", "q5");
assert_eq!(msgs.len(), 6);
assert_eq!(msgs[1].content, "q3");
assert_eq!(msgs[3].content, "q4");
assert_eq!(msgs[5].content, "q5");
}
#[tokio::test]
async fn rolling_window_usage_reports_correctly() {
let mut rw = RollingWindow::new(10);
rw.sync_turn("q", "a", &dummy_metrics()).await;
rw.sync_turn("q", "a", &dummy_metrics()).await;
let (label, cur, max) = rw.usage().unwrap();
assert_eq!(label, "turns");
assert_eq!(cur, 2);
assert_eq!(max, 10);
}
#[tokio::test]
async fn memory_manager_routes_to_provider() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
let msgs = mgr.build_messages("sys", "hello");
assert_eq!(msgs[0].role, Role::System);
assert_eq!(msgs.last().unwrap().content, "hello");
}
fn metrics_with_input(input_tokens: u32) -> TurnMetrics {
let mut m = dummy_metrics();
m.usage = Some(TokenUsage {
input_tokens,
output_tokens: 20,
});
m
}
#[tokio::test]
async fn token_budget_prunes_oldest_when_over_budget() {
let mut tb = TokenBudget::new(512, 0.99);
let budget = 506;
let big = "x".repeat(200); tb.sync_turn(&big, &big, &metrics_with_input(200)).await;
assert_eq!(tb.history.len(), 1);
tb.sync_turn(&big, &big, &metrics_with_input(520)).await;
assert!(tb.used_tokens() <= budget, "used = {}", tb.used_tokens());
assert_eq!(tb.history.len(), 1, "the oldest turn must have been pruned");
}
#[tokio::test]
async fn token_budget_uses_metrics_when_available() {
let mut tb = TokenBudget::new(1000, 1.0);
let mut m = dummy_metrics();
m.usage = Some(crate::metrics::TokenUsage {
input_tokens: 30,
output_tokens: 20,
});
tb.sync_turn("q", "a", &m).await;
assert_eq!(tb.used_tokens(), 31); }
#[tokio::test]
async fn token_budget_no_runaway_drift_b3_regression() {
let mut tb = TokenBudget::new(8_192, 0.80); let mut old_running_sum: u64 = 0;
let turns = 20u32;
for i in 0..turns {
let input = 2_582 + (4_748 - 2_582) * i / (turns - 1);
old_running_sum += u64::from(input) + 20;
tb.sync_turn("reply ok", "ok", &metrics_with_input(input))
.await;
}
let used = u64::from(tb.used_tokens());
assert!(
(4_748..4_800).contains(&used),
"used must track the last real prompt, got {used}"
);
assert_eq!(tb.history.len(), turns as usize, "no spurious pruning");
assert_eq!(tb.pruned_count, 0);
assert!(
used * 5 < old_running_sum,
"anchored used ({used}) must be at least 5× below the old \
running sum ({old_running_sum})"
);
}
fn stub_summarizer(
reply: &'static str,
) -> impl Fn(String) -> crate::agentic::SummarizeFuture + Send + Sync {
move |_req: String| -> crate::agentic::SummarizeFuture {
Box::pin(async move { Ok(reply.to_string()) })
}
}
fn capturing_summarizer(
reply: &'static str,
calls: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
) -> impl Fn(String) -> crate::agentic::SummarizeFuture + Send + Sync {
move |req: String| -> crate::agentic::SummarizeFuture {
calls.lock().unwrap().push(req);
Box::pin(async move { Ok(reply.to_string()) })
}
}
fn turn_with_tokens(
user: &str,
assistant: &str,
tokens_in: Option<u32>,
tokens_out: Option<u32>,
) -> crate::ConversationTurn {
let mut t = crate::ConversationTurn::new(user, assistant);
t.tokens_in = tokens_in;
t.tokens_out = tokens_out;
t
}
#[tokio::test]
async fn summarizing_no_runaway_compression_b3_regression() {
let mut s = Summarizing::new(8_192) .with_summarizer(stub_summarizer("SUMMARY"));
for i in 0..20u32 {
let input = 2_582 + (4_748 - 2_582) * i / 19;
s.sync_turn("reply ok", "ok", &metrics_with_input(input))
.await;
}
assert_eq!(
s.compress_count, 0,
"prompts never exceeded the budget — compression must not fire"
);
assert!(s.prev_summary.is_empty());
assert_eq!(s.history.len(), 20);
}
#[test]
fn token_budget_construction_injects_budget() {
let tb = TokenBudget::new(24_000, 0.80);
let (label, _used, budget) = tb.usage().unwrap();
assert_eq!(label, "tokens");
assert_eq!(budget, 19_200); }
#[test]
fn token_budget_with_budget_overrides_constructor_value() {
let tb = TokenBudget::new(512, 0.80).with_budget(24_000);
assert_eq!(tb.usage().unwrap().2, 19_200);
let clamped = TokenBudget::new(4_096, 1.0).with_budget(0);
assert_eq!(clamped.max_tokens, 512);
}
#[test]
fn summarizing_construction_injects_budget() {
let s = Summarizing::new(32_768);
let (label, _used, budget) = s.usage().unwrap();
assert_eq!(label, "tokens");
assert_eq!(budget, 26_214); }
#[test]
fn summarizing_with_budget_overrides_constructor_value() {
let s = Summarizing::new(100).with_budget(32_768);
assert_eq!(s.usage().unwrap().2, 26_214);
let clamped = Summarizing::new(100).with_budget(0);
assert_eq!(clamped.max_tokens, 1);
}
#[test]
fn token_budget_does_not_sit_at_static_default_with_capability_data() {
let capability_derived = 24_000;
let tb = TokenBudget::new(capability_derived, 0.80);
let static_default_budget = (DEFAULT_CONTEXT_TOKENS as f32 * 0.80) as usize;
assert_ne!(
tb.usage().unwrap().2,
static_default_budget,
"provider budget must reflect injected capability data, \
not the static default"
);
assert_eq!(tb.usage().unwrap().2, 19_200);
}
#[tokio::test]
async fn soul_provider_uses_default_when_no_file() {
let mut sp = SoulProvider::new(None);
let ctx = SessionContext {
workspace: "/nonexistent".into(),
session_id: "s".into(),
};
sp.initialize(&ctx).await.unwrap();
assert_eq!(sp.source, SoulSource::Default);
let block = sp.system_prompt_block().unwrap();
assert!(block.contains("newt"), "default soul should mention newt");
}
#[test]
fn default_soul_lists_all_current_tools() {
for tool in [
"run_command",
"read_file",
"write_file",
"edit_file",
"list_dir",
"find",
"use_skill",
"web_fetch",
] {
assert!(
DEFAULT_SOUL.contains(tool),
"DEFAULT_SOUL must advertise `{tool}`"
);
}
}
#[tokio::test]
async fn soul_provider_loads_workspace_soul() {
let dir = tempfile::tempdir().unwrap();
let newt_dir = dir.path().join(".newt");
std::fs::create_dir_all(&newt_dir).unwrap();
std::fs::write(newt_dir.join("soul.md"), "You are a Django expert.").unwrap();
let mut sp = SoulProvider::new(None);
let ctx = SessionContext {
workspace: dir.path().to_string_lossy().into(),
session_id: "s".into(),
};
sp.initialize(&ctx).await.unwrap();
assert_eq!(sp.source, SoulSource::Workspace);
let block = sp.system_prompt_block().unwrap();
assert!(block.contains("Django"), "should use workspace soul");
}
#[tokio::test]
async fn soul_provider_explicit_path_wins() {
let dir = tempfile::tempdir().unwrap();
let soul_file = dir.path().join("custom_soul.md");
std::fs::write(&soul_file, "You are a security auditor.").unwrap();
let ws_dir = tempfile::tempdir().unwrap();
let newt_dir = ws_dir.path().join(".newt");
std::fs::create_dir_all(&newt_dir).unwrap();
std::fs::write(newt_dir.join("soul.md"), "You are a Django expert.").unwrap();
let mut sp = SoulProvider::new(Some(soul_file.clone()));
let ctx = SessionContext {
workspace: ws_dir.path().to_string_lossy().into(),
session_id: "s".into(),
};
sp.initialize(&ctx).await.unwrap();
assert_eq!(sp.source, SoulSource::Explicit(soul_file));
let block = sp.system_prompt_block().unwrap();
assert!(
block.contains("security auditor"),
"explicit path should win"
);
}
#[tokio::test]
async fn soul_provider_empty_workspace_soul_falls_through() {
let dir = tempfile::tempdir().unwrap();
let newt_dir = dir.path().join(".newt");
std::fs::create_dir_all(&newt_dir).unwrap();
std::fs::write(newt_dir.join("soul.md"), " ").unwrap();
let mut sp = SoulProvider::new(None);
let ctx = SessionContext {
workspace: dir.path().to_string_lossy().into(),
session_id: "s".into(),
};
sp.initialize(&ctx).await.unwrap();
assert_eq!(sp.source, SoulSource::Default);
}
#[tokio::test]
async fn summarizing_compresses_when_over_budget() {
let mut s = Summarizing::new(100) .with_summarizer(stub_summarizer("SUMMARY"));
let big = "x".repeat(200);
for i in 0..5u32 {
s.sync_turn(&big, &big, &metrics_with_input(10 + i)).await;
}
assert_eq!(s.compress_count, 0);
s.sync_turn(&big, &big, &metrics_with_input(120)).await;
assert!(s.compress_count >= 1, "compress_count={}", s.compress_count);
assert!(
s.prev_summary.starts_with(crate::agentic::SUMMARY_PREFIX),
"prev_summary must be the marked compaction message"
);
assert!(s.prev_summary.contains("SUMMARY"));
assert!(s.prev_summary.contains(crate::agentic::SUMMARY_END_MARKER));
assert!(
s.history
.iter()
.any(|t| t.user.starts_with(crate::agentic::SUMMARY_PREFIX)
&& t.assistant.is_empty()),
"the compaction message must live in history as a lone user entry"
);
}
#[tokio::test]
async fn summarizing_compresses_repeatedly() {
let mut s = Summarizing::new(100).with_summarizer(stub_summarizer("SUMMARY"));
let text = "x".repeat(120);
for i in 0..20u32 {
s.sync_turn(&text, &text, &metrics_with_input(50 * (i + 1)))
.await;
}
assert!(s.compress_count >= 1);
}
#[tokio::test]
async fn summarizing_delegates_to_shared_pipeline() {
let calls = std::sync::Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
let mut s =
Summarizing::new(100).with_summarizer(capturing_summarizer("PIPE-SUM", calls.clone()));
let secret = "sk-aaaaaaaaaaaaaaaaaaaaaaaa";
let big = "x".repeat(200);
s.sync_turn(
"the original task",
&format!("noted {secret}"),
&metrics_with_input(10),
)
.await;
for i in 0..4u32 {
s.sync_turn(&big, &big, &metrics_with_input(11 + i)).await;
}
assert!(calls.lock().unwrap().is_empty(), "under budget — no calls");
s.sync_turn(&big, &big, &metrics_with_input(120)).await;
let reqs = calls.lock().unwrap();
assert_eq!(reqs.len(), 1, "exactly one summarizer call per compression");
let req = &reqs[0];
assert!(
req.contains("## Conversation middle to summarise"),
"must be the shared pipeline's request template"
);
assert!(req.contains("## Original Task"));
assert!(req.contains("the original task"), "task anchored verbatim");
assert!(
!req.contains(secret),
"secret must not reach the summarizer"
);
assert!(req.contains("[REDACTED]"), "shared redaction pass applied");
drop(reqs);
assert!(s.prev_summary.starts_with(crate::agentic::SUMMARY_PREFIX));
assert!(s.prev_summary.contains("PIPE-SUM"));
assert!(s.prev_summary.contains(crate::agentic::SUMMARY_END_MARKER));
}
#[tokio::test]
async fn summarizing_compaction_record_is_minted_once_and_drained() {
let mut s = Summarizing::new(100).with_summarizer(stub_summarizer("SUMMARY"));
assert!(s.take_compaction_record().is_none(), "nothing minted yet");
let big = "x".repeat(200);
for i in 0..5u32 {
s.sync_turn(&big, &big, &metrics_with_input(10 + i)).await;
}
s.sync_turn(&big, &big, &metrics_with_input(120)).await;
let record = s
.take_compaction_record()
.expect("compression must mint a record");
assert!(record.starts_with(crate::agentic::SUMMARY_PREFIX));
assert_eq!(record, s.prev_summary, "the record IS the chain head");
assert!(
s.take_compaction_record().is_none(),
"the record is drained on take — never persisted twice"
);
}
#[tokio::test]
async fn memory_manager_routes_take_compaction_record() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(50)); mgr.add_provider(Summarizing::new(100).with_summarizer(stub_summarizer("SUMMARY")));
assert!(mgr.take_compaction_record().is_none());
let big = "x".repeat(200);
for i in 0..5u32 {
mgr.sync_all(&big, &big, &metrics_with_input(10 + i)).await;
}
mgr.sync_all(&big, &big, &metrics_with_input(120)).await;
let record = mgr
.take_compaction_record()
.expect("manager must surface the Summarizing provider's record");
assert!(record.starts_with(crate::agentic::SUMMARY_PREFIX));
assert!(mgr.take_compaction_record().is_none());
}
#[tokio::test]
async fn summarizing_restore_rehydrates_compaction_summary() {
let marked = format!(
"{}\nsummary of earlier work\n{}",
crate::agentic::SUMMARY_PREFIX,
crate::agentic::SUMMARY_END_MARKER
);
let turns = vec![
crate::ConversationTurn::new("old task", "old reply"), crate::ConversationTurn::new(marked.clone(), ""), crate::ConversationTurn::new("recent task", "recent reply"),
];
let mut s = Summarizing::new(8_192);
s.restore_turns(&turns);
let pre = s.on_pre_compress(&[]).await;
assert!(pre.contains("summary of earlier work"), "got: {pre:?}");
let msgs = s.build_messages("sys", "next");
assert!(msgs
.iter()
.any(|m| m.content.starts_with(crate::agentic::SUMMARY_PREFIX)));
assert!(!msgs.iter().any(|m| m.content == "old task"));
assert!(msgs.iter().any(|m| m.content == "recent task"));
assert!(msgs.iter().any(|m| m.content == "recent reply"));
assert!(!msgs.iter().any(|m| m.content.is_empty()));
}
#[tokio::test]
async fn summarizing_restore_never_resummarizes() {
let calls = std::sync::Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
let mut s =
Summarizing::new(10).with_summarizer(capturing_summarizer("SUMMARY", calls.clone()));
let big = "x".repeat(400);
let turns: Vec<crate::ConversationTurn> = (0..6)
.map(|i| crate::ConversationTurn::new(format!("q{i} {big}"), format!("a{i} {big}")))
.collect();
s.restore_turns(&turns); assert!(
calls.lock().unwrap().is_empty(),
"restore must never call the summarizer"
);
assert_eq!(s.history.len(), 6, "restored history intact");
}
#[tokio::test]
async fn restored_compaction_chains_into_next_compression() {
let calls = std::sync::Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
let mut s =
Summarizing::new(100).with_summarizer(capturing_summarizer("NEW-SUM", calls.clone()));
let marked = format!(
"{}\nCHAIN-ME: facts from the first compaction\n{}",
crate::agentic::SUMMARY_PREFIX,
crate::agentic::SUMMARY_END_MARKER
);
let big = "y".repeat(100);
let mut turns = vec![crate::ConversationTurn::new(marked.clone(), "")];
for i in 0..4 {
turns.push(crate::ConversationTurn::new(
format!("task {i} {big}"),
format!("reply {i} {big}"),
));
}
s.restore_turns(&turns);
s.sync_turn(&big, &big, &metrics_with_input(200)).await;
let reqs = calls.lock().unwrap();
assert_eq!(reqs.len(), 1, "one call through the shared path");
assert!(
reqs[0].contains("CHAIN-ME"),
"the previous summary must be summarizer INPUT (the chain), got: {}",
&reqs[0]
);
let anchored = reqs[0]
.split("## Original Task")
.nth(1)
.and_then(|rest| rest.lines().nth(1).map(str::to_string))
.unwrap_or_default();
assert!(
anchored.starts_with("task 0"),
"task anchor must skip the compaction message, got: {anchored:?}"
);
drop(reqs);
let compactions: Vec<&SumTurn> = s
.history
.iter()
.filter(|t| t.user.starts_with(crate::agentic::SUMMARY_PREFIX))
.collect();
assert_eq!(compactions.len(), 1, "exactly one compaction in history");
assert!(compactions[0].user.contains("NEW-SUM"));
assert!(s.prev_summary.contains("NEW-SUM"));
}
#[test]
fn token_budget_restore_anchors_on_column_tokens() {
let mut tb = TokenBudget::new(100_000, 0.80);
let turns = vec![
turn_with_tokens("u1", "a1", Some(900), Some(10)),
turn_with_tokens("u2", "aa", Some(1_000), Some(12)),
];
tb.restore_turns(&turns);
assert_eq!(tb.last_prompt_tokens, Some(1_000));
assert_eq!(tb.used_tokens(), 1_001);
}
#[test]
fn token_budget_restore_null_columns_fall_back_to_estimate() {
let mut tb = TokenBudget::new(100_000, 0.80);
let text = "x".repeat(40); let turns = vec![
crate::ConversationTurn::new(text.clone(), text.clone()),
crate::ConversationTurn::new(text.clone(), text.clone()),
];
tb.restore_turns(&turns);
assert_eq!(
tb.last_prompt_tokens, None,
"no measurement in the store → no anchor, only estimates"
);
assert_eq!(tb.used_tokens(), 40, "2 turns × (10 + 10) est tokens");
}
#[test]
fn token_budget_restore_unmeasured_tail_extends_delta() {
let mut tb = TokenBudget::new(100_000, 0.80);
let turns = vec![
turn_with_tokens("u1", "aaaa", Some(1_000), Some(8)), crate::ConversationTurn::new("xxxx", "yyyy"), ];
tb.restore_turns(&turns);
assert_eq!(tb.last_prompt_tokens, Some(1_000));
assert_eq!(tb.used_tokens(), 1_003);
}
#[test]
fn summarizing_restore_anchors_on_column_tokens() {
let mut s = Summarizing::new(100_000);
let turns = vec![
turn_with_tokens("u1", "a1", Some(700), Some(9)),
turn_with_tokens("u2", "aaaa", Some(2_000), Some(11)),
];
s.restore_turns(&turns);
assert_eq!(s.last_prompt_tokens, Some(2_000));
assert_eq!(s.used_tokens(), 2_001); }
#[test]
fn summarizing_restore_ignores_measurements_before_the_cut() {
let marked = format!(
"{}\nolder work compressed\n{}",
crate::agentic::SUMMARY_PREFIX,
crate::agentic::SUMMARY_END_MARKER
);
let mut s = Summarizing::new(100_000);
let turns = vec![
turn_with_tokens("old", "old", Some(50_000), Some(10)),
crate::ConversationTurn::new(marked, ""),
crate::ConversationTurn::new("after", "compaction"), ];
s.restore_turns(&turns);
assert_eq!(
s.last_prompt_tokens, None,
"a pre-compression measurement must not anchor the cut working set"
);
}
#[test]
fn rolling_window_restores_compaction_record_without_empty_assistant() {
let marked = format!(
"{}\nearlier work\n{}",
crate::agentic::SUMMARY_PREFIX,
crate::agentic::SUMMARY_END_MARKER
);
let mut rw = RollingWindow::new(10);
rw.restore_turns(&[
crate::ConversationTurn::new(marked.clone(), ""),
crate::ConversationTurn::new("next task", "next reply"),
]);
let msgs = rw.build_messages("sys", "go");
assert!(msgs.iter().any(|m| m.content == marked));
assert!(!msgs.iter().any(|m| m.content.is_empty()));
}
#[tokio::test]
async fn memory_manager_on_pre_compress() {
let mgr = MemoryManager::new();
let result = mgr.on_pre_compress(&[]).await;
assert!(result.is_empty());
}
#[tokio::test]
async fn memory_manager_on_session_end() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
mgr.on_session_end(&[]).await; }
#[tokio::test]
async fn memory_manager_prefetch_all_empty() {
let mgr = MemoryManager::new();
let result = mgr.prefetch_all("query").await;
assert!(result.is_empty());
}
#[tokio::test]
async fn memory_manager_build_system_prompt_additions_from_note_store() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
std::fs::write(&path, "fact one\nfact two").unwrap();
let mut ns = NoteStore::new(path, 2200);
let ctx = SessionContext {
workspace: "/ws".into(),
session_id: "s".into(),
};
ns.initialize(&ctx).await.unwrap();
let mut mgr = MemoryManager::new();
mgr.add_provider(ns);
let additions = mgr.build_system_prompt_additions();
assert!(additions.contains("fact one"));
}
#[tokio::test]
async fn memory_manager_add_note_fails_with_no_note_store() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
let err = mgr.add_note("fact").unwrap_err().to_string();
assert!(
err.contains("no note-capable memory provider"),
"guidance error expected: {err}"
);
}
#[tokio::test]
async fn token_budget_usage_reporting() {
let tb = TokenBudget::new(1000, 0.80);
let (label, cur, max) = tb.usage().unwrap();
assert_eq!(label, "tokens");
assert_eq!(cur, 0);
assert_eq!(max, 800); }
#[tokio::test]
async fn token_budget_does_not_prune_within_budget() {
let mut tb = TokenBudget::new(200, 1.0); let mut m = dummy_metrics();
m.usage = Some(crate::metrics::TokenUsage {
input_tokens: 50,
output_tokens: 50,
});
tb.sync_turn("q", "a", &m).await; assert_eq!(tb.history.len(), 1);
}
#[tokio::test]
async fn token_budget_estimate_fallback_counts_each_turn_once() {
let mut tb = TokenBudget::new(1000, 1.0);
let mut m = dummy_metrics();
m.usage = None; let text = "x".repeat(40); tb.sync_turn(&text, &text, &m).await;
tb.sync_turn(&text, &text, &m).await;
assert_eq!(tb.used_tokens(), 40, "2 turns × (10 + 10) est tokens");
}
#[tokio::test]
async fn summarizing_fallback_placeholder_when_no_summarizer() {
let mut s = Summarizing::new(10); for i in 0..6u32 {
s.sync_turn(
&format!("question {i}"),
&format!("answer {i}"),
&metrics_with_input(6 + 4 * i),
)
.await;
}
let compaction = s
.history
.iter()
.find(|t| t.user.starts_with(crate::agentic::SUMMARY_PREFIX))
.expect("static fallback marker should be inserted");
assert!(
compaction
.user
.contains("Summary generation was unavailable"),
"got: {}",
compaction.user
);
}
#[tokio::test]
async fn summarizing_usage_reporting() {
let s = Summarizing::new(1000);
let (label, cur, max) = s.usage().unwrap();
assert_eq!(label, "tokens");
assert_eq!(cur, 0);
assert_eq!(max, 800); }
#[tokio::test]
async fn summarizing_on_pre_compress_returns_prev_summary() {
let mut s = Summarizing::new(10).with_summarizer(stub_summarizer("PRIOR SUMMARY"));
for _ in 0..5u32 {
s.sync_turn("question text", "answer text", &metrics_with_input(2))
.await;
}
s.sync_turn("question text", "answer text", &metrics_with_input(50))
.await;
let pre = s.on_pre_compress(&[]).await;
assert!(
pre.contains("PRIOR SUMMARY"),
"compression must have run and set prev_summary, got: {pre:?}"
);
}
#[tokio::test]
async fn rolling_window_on_session_end_noop() {
let mut rw = RollingWindow::new(5);
rw.on_session_end(&[]).await; }
#[tokio::test]
async fn rolling_window_add_note_returns_notes_unsupported() {
let mut rw = RollingWindow::new(5);
let err = rw.add_note("fact").unwrap_err();
assert!(err.is::<NotesUnsupported>());
}
#[tokio::test]
async fn rolling_window_replace_and_remove_note_return_notes_unsupported() {
let mut rw = RollingWindow::new(5);
assert!(rw
.replace_note("old", "new")
.unwrap_err()
.is::<NotesUnsupported>());
assert!(rw.remove_note("old").unwrap_err().is::<NotesUnsupported>());
}
#[tokio::test]
async fn memory_manager_replace_and_remove_route_to_note_store() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
mgr.add_provider(NoteStore::new(path.clone(), 2200));
let ctx = SessionContext {
workspace: "/ws".into(),
session_id: "s".into(),
};
mgr.initialize_all(&ctx).await;
mgr.add_note("model alpha is the fast tier").unwrap();
mgr.add_note("workspace uses just check").unwrap();
mgr.replace_note("alpha", "model beta is the fast tier")
.unwrap();
let raw = std::fs::read_to_string(&path).unwrap();
assert!(raw.contains("model beta"), "{raw}");
assert!(!raw.contains("model alpha"), "{raw}");
mgr.remove_note("just check").unwrap();
let raw = std::fs::read_to_string(&path).unwrap();
assert!(!raw.contains("just check"), "{raw}");
assert!(raw.contains("model beta"), "other entry untouched: {raw}");
let err = mgr.remove_note("nonexistent").unwrap_err().to_string();
assert!(err.contains("no entry contains"), "{err}");
}
#[tokio::test]
async fn memory_manager_replace_and_remove_fail_with_no_note_store() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
let err = mgr.replace_note("a", "b").unwrap_err().to_string();
assert!(err.contains("no note-capable memory provider"), "{err}");
let err = mgr.remove_note("a").unwrap_err().to_string();
assert!(err.contains("no note-capable memory provider"), "{err}");
}
#[tokio::test]
async fn memory_manager_add_note_routes_to_note_store() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
mgr.add_provider(NoteStore::new(path.clone(), 2200));
let ctx = SessionContext {
workspace: "/ws".into(),
session_id: "s".into(),
};
mgr.initialize_all(&ctx).await;
mgr.add_note("the answer is 42").unwrap();
let raw = std::fs::read_to_string(&path).unwrap();
assert!(raw.contains("the answer is 42"));
}
struct CustomNotes {
notes: Vec<String>,
}
#[async_trait]
impl MemoryProvider for CustomNotes {
fn name(&self) -> &str {
"custom_notes"
}
fn build_messages(&self, _system_prompt: &str, _new_task: &str) -> Vec<MemMessage> {
Vec::new()
}
async fn sync_turn(&mut self, _user: &str, _assistant: &str, _metrics: &TurnMetrics) {}
fn add_note(&mut self, fact: &str) -> anyhow::Result<()> {
self.notes.push(fact.to_string());
Ok(())
}
}
#[tokio::test]
async fn memory_manager_add_note_first_ok_wins_regardless_of_name() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5)); mgr.add_provider(CustomNotes { notes: Vec::new() });
mgr.add_note("routed by capability, not by name").unwrap();
}
#[tokio::test]
async fn memory_manager_add_note_surfaces_curator_error() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
mgr.add_provider(NoteStore::new(path, 40));
let ctx = SessionContext {
workspace: "/ws".into(),
session_id: "s".into(),
};
mgr.initialize_all(&ctx).await;
mgr.add_note("an entry that fits").unwrap();
let err = mgr.add_note(&"x".repeat(80)).unwrap_err().to_string();
assert!(
err.contains("Replace or remove existing entries first"),
"curator error must propagate: {err}"
);
assert!(err.contains("an entry that fits"), "{err}");
}
#[tokio::test]
async fn memory_manager_sync_all() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
mgr.sync_all("q", "a", &dummy_metrics()).await;
let usage = mgr.usage();
assert_eq!(usage[0].1, 1); }
#[tokio::test]
async fn memory_manager_reset_all_clears_conversation_history() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
mgr.sync_all("old task", "old reply", &dummy_metrics())
.await;
let before = mgr.build_messages("system", "new task");
assert!(before.iter().any(|m| m.content == "old task"));
assert!(before.iter().any(|m| m.content == "old reply"));
mgr.reset_all();
let after = mgr.build_messages("system", "new task");
assert!(!after.iter().any(|m| m.content == "old task"));
assert!(!after.iter().any(|m| m.content == "old reply"));
assert!(after.iter().any(|m| m.content == "new task"));
}
#[tokio::test]
async fn memory_manager_restore_turns_replaces_conversation_history() {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(5));
mgr.sync_all("old task", "old reply", &dummy_metrics())
.await;
mgr.restore_turns(&[
crate::ConversationTurn::new("restored task", "restored reply"),
crate::ConversationTurn::new("follow up", "followed up"),
]);
let messages = mgr.build_messages("system", "new task");
assert!(!messages.iter().any(|m| m.content == "old task"));
assert!(!messages.iter().any(|m| m.content == "old reply"));
assert!(messages.iter().any(|m| m.content == "restored task"));
assert!(messages.iter().any(|m| m.content == "restored reply"));
assert!(messages.iter().any(|m| m.content == "follow up"));
assert!(messages.iter().any(|m| m.content == "followed up"));
assert!(messages.iter().any(|m| m.content == "new task"));
}
#[tokio::test]
async fn memory_manager_fallback_with_no_providers() {
let mgr = MemoryManager::new();
let msgs = mgr.build_messages("sys", "task");
assert_eq!(msgs.len(), 2);
}
fn index_ctx(dir: &std::path::Path) -> SessionContext {
SessionContext {
workspace: dir.to_string_lossy().into(),
session_id: "s".into(),
}
}
async fn seed_notes(path: &std::path::Path, n: usize) {
let mut ns = NoteStore::new(path, NoteStore::DEFAULT_CHAR_LIMIT);
ns.initialize(&index_ctx(path.parent().unwrap()))
.await
.unwrap();
for i in 0..n {
ns.add(&format!("note number {i}\nbody line for {i}"))
.unwrap();
}
}
#[tokio::test]
async fn memory_index_stays_under_pinned_budget() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
seed_notes(&path, MEMORY_INDEX_BUDGET + 25).await;
let mut idx = MemoryIndex::new(&path);
idx.initialize(&index_ctx(dir.path())).await.unwrap();
assert!(
idx.rows().len() <= MEMORY_INDEX_BUDGET,
"index surface ({}) exceeds the pinned budget ({MEMORY_INDEX_BUDGET})",
idx.rows().len()
);
assert_eq!(idx.rows().len(), MEMORY_INDEX_BUDGET, "fills to the cap");
let block = idx.system_prompt_block().unwrap();
assert!(block.contains("use `recall`"), "overflow hint: {block}");
}
#[tokio::test]
async fn memory_index_lists_ids_and_titles_not_bodies() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
seed_notes(&path, 2).await;
let mut idx = MemoryIndex::new(&path);
idx.initialize(&index_ctx(dir.path())).await.unwrap();
let block = idx.system_prompt_block().unwrap();
assert!(block.contains("note:1 note number 0"), "got: {block}");
assert!(block.contains("note:2 note number 1"), "got: {block}");
assert!(!block.contains("body line for 0"), "body leaked: {block}");
assert!(
block.contains("call `memory_fetch`"),
"names the fetch tool: {block}"
);
}
#[tokio::test]
async fn memory_index_is_system_prompt_only() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
seed_notes(&path, 1).await;
let mut idx = MemoryIndex::new(&path);
idx.initialize(&index_ctx(dir.path())).await.unwrap();
assert!(idx.build_messages("sys", "task").is_empty());
}
#[tokio::test]
async fn memory_index_empty_notes_contributes_no_block() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
let mut idx = MemoryIndex::new(&path); idx.initialize(&index_ctx(dir.path())).await.unwrap();
assert!(idx.system_prompt_block().is_none());
}
#[tokio::test]
async fn disclosure_frozen_default_is_bit_for_bit_unchanged() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("NOTES.md");
seed_notes(&path, 3).await;
async fn build(
path: &std::path::Path,
ws: &std::path::Path,
with_index: bool,
) -> (String, Vec<MemMessage>) {
let mut mgr = MemoryManager::new();
mgr.add_provider(RollingWindow::new(20));
mgr.add_provider(NoteStore::new(path, NoteStore::DEFAULT_CHAR_LIMIT));
if with_index {
mgr.add_provider(MemoryIndex::new(path));
}
mgr.initialize_all(&SessionContext {
workspace: ws.to_string_lossy().into(),
session_id: "s".into(),
})
.await;
(
mgr.build_system_prompt_additions(),
mgr.build_messages("sys", "task"),
)
}
let (frozen_sys, frozen_msgs) = build(&path, dir.path(), false).await;
let (frozen_sys2, frozen_msgs2) = build(&path, dir.path(), false).await;
assert_eq!(frozen_sys, frozen_sys2);
assert_eq!(frozen_msgs, frozen_msgs2);
assert!(
!frozen_sys.contains("Memory index"),
"frozen mode must NOT add the block: {frozen_sys}"
);
let (index_sys, _index_msgs) = build(&path, dir.path(), true).await;
assert!(
index_sys.contains("Memory index"),
"index mode adds the block: {index_sys}"
);
}
}