use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use futures::StreamExt as _;
use serde::Deserialize;
use serde_json::json;
use tokio::sync::mpsc;
use wcode_harness::agent::{Agent, AgentConfig};
use wcode_harness::compaction::{CompactOutcome, CompactionPolicy};
use wcode_harness::event::{AgentEvent, LlmStreamEvent};
use wcode_harness::hooks::HooksSet;
use wcode_harness::message::{AgentMessage, StopReason};
use wcode_harness::session::{Session, SessionEntry};
use wcode_harness::streamfn::{
LlmEndpoint, LlmOpts, LlmProfile, LlmProvider, LlmStream, ModelProfiles, StreamFn,
};
use wcode_harness::tool::{ToolContext, ToolOutput, TypedTool, erased};
#[derive(Debug, Clone)]
struct StreamCall {
ctx: Vec<AgentMessage>,
system: String,
model: String,
endpoint: LlmEndpoint,
base_url: Option<String>,
api_key: Option<String>,
}
#[derive(Clone, Default)]
struct Recorder {
script: Arc<Mutex<VecDeque<Vec<LlmStreamEvent>>>>,
calls: Arc<Mutex<Vec<StreamCall>>>,
}
impl Recorder {
fn push(&self, events: Vec<LlmStreamEvent>) {
self.script.lock().unwrap().push_back(events);
}
fn calls(&self) -> Vec<StreamCall> {
self.calls.lock().unwrap().clone()
}
}
fn fake_stream_fn(rec: &Recorder) -> StreamFn {
let rec = rec.clone();
Arc::new(
move |ctx: &[AgentMessage],
system: &str,
_tools: &[rig::completion::ToolDefinition],
_opts: &LlmOpts| {
rec.calls.lock().unwrap().push(StreamCall {
ctx: ctx.to_vec(),
system: system.to_string(),
model: _opts.model.clone(),
endpoint: _opts.endpoint,
base_url: _opts.base_url.clone(),
api_key: _opts.api_key.clone(),
});
let events = rec.script.lock().unwrap().pop_front().unwrap_or_default();
Box::pin(futures::stream::iter(events)) as LlmStream
},
)
}
#[derive(Deserialize, schemars::JsonSchema)]
struct GateArgs {
text: String,
}
struct GateTool {
entered: mpsc::UnboundedSender<()>,
release: Arc<tokio::sync::Notify>,
}
#[async_trait::async_trait]
impl TypedTool for GateTool {
type Args = GateArgs;
fn name(&self) -> &str {
"gate"
}
fn description(&self) -> &str {
"blocks until released"
}
async fn execute(&self, args: GateArgs, _ctx: &ToolContext) -> ToolOutput {
self.entered.send(()).ok();
self.release.notified().await;
ToolOutput {
output: format!("ran:{}", args.text),
is_error: false,
diff: None,
path: None,
}
}
}
fn agent_config(
stream_fn: StreamFn,
tools: Vec<wcode_harness::tool::Tool>,
session: Option<Session>,
) -> AgentConfig {
agent_config_with_llm(
stream_fn,
tools,
session,
LlmOpts {
model: "m1".into(),
..LlmOpts::default()
},
)
}
fn agent_config_with_llm(
stream_fn: StreamFn,
tools: Vec<wcode_harness::tool::Tool>,
session: Option<Session>,
llm: LlmOpts,
) -> AgentConfig {
AgentConfig {
system: "sys".into(),
tools,
llm,
stream_fn,
hooks: HooksSet::default(),
session,
context: Vec::new(),
working_dir: std::path::PathBuf::new(),
max_turns: wcode_harness::loop_::DEFAULT_MAX_TURNS,
parallel_tools: true,
compaction: wcode_harness::compaction::CompactionPolicy::default(),
plan_mode: wcode_harness::hooks::PlanModeHandle::new(),
}
}
#[tokio::test]
async fn steer_during_blocked_tool_reaches_next_turn_stream() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("after steer".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let slot: Arc<OnceLock<mpsc::UnboundedSender<AgentMessage>>> = Arc::new(OnceLock::new());
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let mut agent = Agent::new(agent_config(fake_stream_fn(&rec), vec![gate], None));
slot.set(agent.steer_sender()).unwrap();
let steer = slot.get().unwrap().clone();
tokio::spawn(async move {
entered_rx.recv().await.unwrap();
steer
.send(AgentMessage::user_text("mid-run steer"))
.unwrap();
release.notify_one();
});
let (tx, _rx) = mpsc::unbounded_channel();
let res = agent.run("hi", tx).await;
assert_eq!(res.unwrap(), StopReason::Stop);
let calls = rec.calls();
assert_eq!(calls.len(), 2);
let ctx2 = &calls[1].ctx;
assert_eq!(ctx2.len(), 4);
assert!(
matches!(&ctx2[3], AgentMessage::User { .. } if ctx2[3].as_text() == "mid-run steer"),
"steer must appear before the next turn's stream call, got: {:?}",
ctx2
);
}
#[tokio::test]
async fn cancel_returns_aborted() {
let stream_fn: StreamFn = Arc::new(move |_ctx, _sys, _tools, _opts| {
let delta = futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())]);
Box::pin(delta.chain(futures::stream::pending())) as LlmStream
});
let mut agent = Agent::new(agent_config(stream_fn, vec![], None));
let token = agent.cancel_token();
let (tx, mut rx) = mpsc::unbounded_channel();
tokio::spawn(async move {
while let Some(ev) = rx.recv().await {
if matches!(ev, AgentEvent::MessageUpdate { .. }) {
token.cancel();
break;
}
}
});
let res = agent.run("hi", tx).await;
assert_eq!(res.unwrap(), StopReason::Aborted);
assert_eq!(agent.messages().len(), 2);
assert!(matches!(
&agent.messages()[1],
AgentMessage::Assistant { stop_reason: StopReason::Aborted, .. }
if agent.messages()[1].as_text() == "part"
));
}
#[tokio::test]
async fn error_before_any_delta_persists_only_user_message() {
let stream_fn: StreamFn = Arc::new(|_ctx, _sys, _tools, _opts| {
Box::pin(futures::stream::iter(vec![LlmStreamEvent::Error {
message: "boom".into(),
fatal: true,
}])) as LlmStream
});
let dir = tempfile::tempdir().unwrap();
let session = Session::create(dir.path()).unwrap();
let path = session.path().unwrap().to_path_buf();
let mut agent = Agent::new(agent_config(stream_fn, vec![], Some(session)));
let (tx, _rx) = mpsc::unbounded_channel();
let res = agent.run("hi", tx).await;
assert_eq!(res.unwrap(), StopReason::Error);
assert_eq!(
agent.messages().len(),
1,
"only the user message, no empty assistant: {:?}",
agent.messages()
);
let reopened = Session::open(&path).unwrap();
let msgs = reopened.messages();
assert_eq!(
msgs.len(),
1,
"session must not contain an empty assistant: {msgs:?}"
);
assert!(matches!(&msgs[0], AgentMessage::User { .. } if msgs[0].as_text() == "hi"));
}
#[tokio::test]
async fn session_file_contains_all_messages_after_run() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("hello".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let dir = tempfile::tempdir().unwrap();
let session = Session::create(dir.path()).unwrap();
let path = session.path().unwrap().to_path_buf();
let mut agent = Agent::new(agent_config(fake_stream_fn(&rec), vec![], Some(session)));
let (tx, _rx) = mpsc::unbounded_channel();
let res = agent.run("hi", tx).await;
assert_eq!(res.unwrap(), StopReason::Stop);
assert_eq!(agent.session_path(), Some(path.as_path()));
assert_eq!(agent.messages().len(), 2);
assert_eq!(rec.calls()[0].system, "sys");
let reopened = Session::open(&path).unwrap();
let msgs = reopened.messages();
assert_eq!(msgs.len(), 2);
assert!(matches!(&msgs[0], AgentMessage::User { .. } if msgs[0].as_text() == "hi"));
assert!(matches!(&msgs[1], AgentMessage::Assistant { .. } if msgs[1].as_text() == "hello"));
}
#[tokio::test]
async fn session_persists_messages_before_run_finishes() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::ToolCall {
id: "c1".into(),
name: "gate".into(),
arguments: json!({ "text": "hi" }),
},
LlmStreamEvent::Done {
stop_reason: StopReason::ToolUse,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("done".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let (entered_tx, mut entered_rx) = mpsc::unbounded_channel::<()>();
let release = Arc::new(tokio::sync::Notify::new());
let gate = erased(GateTool {
entered: entered_tx,
release: release.clone(),
});
let dir = tempfile::tempdir().unwrap();
let session = Session::create(dir.path()).unwrap();
let path = session.path().unwrap().to_path_buf();
let mut agent = Agent::new(agent_config(
fake_stream_fn(&rec),
vec![gate],
Some(session),
));
let (tx, _rx) = mpsc::unbounded_channel();
let handle = tokio::spawn(async move { agent.run("hi", tx).await.unwrap() });
entered_rx.recv().await.unwrap();
let mid = Session::open(&path).unwrap();
let msgs = mid.messages();
assert_eq!(
msgs.len(),
2,
"messages must be on disk pre-finish: {msgs:?}"
);
assert!(matches!(&msgs[0], AgentMessage::User { .. } if msgs[0].as_text() == "hi"));
release.notify_one();
assert_eq!(handle.await.unwrap(), StopReason::Stop);
let after = Session::open(&path).unwrap();
let msgs = after.messages();
assert_eq!(
msgs.len(),
4,
"user + tool-call assistant + ToolResult + final assistant"
);
assert!(matches!(&msgs[3], AgentMessage::Assistant { .. } if msgs[3].as_text() == "done"));
}
#[tokio::test]
async fn steer_between_runs_is_queued_for_next_run() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("one".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("two".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let mut agent = Agent::new(agent_config(fake_stream_fn(&rec), vec![], None));
let (tx1, _rx1) = mpsc::unbounded_channel();
agent.run("hi", tx1).await.unwrap();
assert_eq!(rec.calls().len(), 1);
agent.steer(AgentMessage::user_text("queued"));
let (tx2, _rx2) = mpsc::unbounded_channel();
let res = agent.run("again", tx2).await;
assert_eq!(res.unwrap(), StopReason::Stop);
let calls = rec.calls();
assert_eq!(calls.len(), 2);
let ctx2 = &calls[1].ctx;
assert_eq!(ctx2.len(), 4);
assert!(
matches!(&ctx2[3], AgentMessage::User { .. } if ctx2[3].as_text() == "queued"),
"steer queued between runs must reach the next run, got: {:?}",
ctx2
);
assert_eq!(agent.messages().len(), 5);
}
#[tokio::test]
async fn steer_sender_clone_survives_runs() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("one".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("two".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let mut agent = Agent::new(agent_config(fake_stream_fn(&rec), vec![], None));
let steer = agent.steer_sender();
let (tx1, _rx1) = mpsc::unbounded_channel();
agent.run("hi", tx1).await.unwrap();
assert_eq!(rec.calls().len(), 1);
let _ = steer.send(AgentMessage::user_text("straggler"));
let (tx2, _rx2) = mpsc::unbounded_channel();
let res = agent.run("again", tx2).await;
assert_eq!(res.unwrap(), StopReason::Stop);
let calls = rec.calls();
assert_eq!(calls.len(), 2);
let ctx2 = &calls[1].ctx;
assert!(
matches!(&ctx2[3], AgentMessage::User { .. } if ctx2[3].as_text() == "straggler"),
"pre-run steer clone must reach the next run, got: {:?}",
ctx2
);
}
#[tokio::test]
async fn cancel_token_refreshed_between_runs() {
let call = Arc::new(AtomicUsize::new(0));
let stream_fn: StreamFn = {
let call = call.clone();
Arc::new(move |_ctx, _sys, _tools, _opts| {
if call.fetch_add(1, Ordering::SeqCst) == 0 {
Box::pin(
futures::stream::iter(vec![LlmStreamEvent::TextDelta("part".into())])
.chain(futures::stream::pending()),
) as LlmStream
} else {
Box::pin(futures::stream::iter(vec![
LlmStreamEvent::TextDelta("after".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
])) as LlmStream
}
})
};
let mut agent = Agent::new(agent_config(stream_fn, vec![], None));
let token = agent.cancel_token();
let (tx1, mut rx1) = mpsc::unbounded_channel();
tokio::spawn(async move {
while let Some(ev) = rx1.recv().await {
if matches!(ev, AgentEvent::MessageUpdate { .. }) {
token.cancel();
break;
}
}
});
let res1 = agent.run("hi", tx1).await;
assert_eq!(res1.unwrap(), StopReason::Aborted);
let (tx2, _rx2) = mpsc::unbounded_channel();
let res2 = agent.run("again", tx2).await;
assert_eq!(res2.unwrap(), StopReason::Stop);
assert_eq!(call.load(Ordering::SeqCst), 2);
assert!(
matches!(
agent.messages().last().unwrap(),
AgentMessage::Assistant { stop_reason: StopReason::Stop, .. }
if agent.messages().last().unwrap().as_text() == "after"
),
"run 2 must stream normally, got: {:?}",
agent.messages()
);
}
#[tokio::test]
async fn seeded_context_reaches_first_stream_call() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("ok".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let mut cfg = agent_config(fake_stream_fn(&rec), vec![], None);
cfg.context = vec![
AgentMessage::user_text("old question"),
AgentMessage::Assistant {
content: vec![wcode_harness::message::ContentBlock::Text {
text: "old answer".into(),
}],
stop_reason: StopReason::Stop,
usage: None,
model: None,
},
];
let mut agent = Agent::new(cfg);
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("new question", tx).await.unwrap();
let ctx = &rec.calls()[0].ctx;
assert_eq!(ctx.len(), 3);
assert_eq!(ctx[0].as_text(), "old question");
assert_eq!(ctx[1].as_text(), "old answer");
assert_eq!(ctx[2].as_text(), "new question");
}
#[tokio::test]
async fn set_model_swaps_llm_and_logs_session_change() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("a".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
rec.push(vec![
LlmStreamEvent::TextDelta("b".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let dir = tempfile::tempdir().unwrap();
let session = Session::create(dir.path()).unwrap();
let path = session.path().unwrap().to_path_buf();
let mut agent = Agent::new(agent_config(fake_stream_fn(&rec), vec![], Some(session)));
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("hi", tx).await.unwrap();
assert_eq!(rec.calls()[0].model, "m1");
agent.set_model("m2".into()).unwrap();
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("again", tx).await.unwrap();
assert_eq!(rec.calls()[1].model, "m2");
let reopened = Session::open(&path).unwrap();
assert_eq!(reopened.model().as_deref(), Some("m2"));
}
#[tokio::test]
async fn set_model_settles_the_provider_and_round_trips() {
let rec = Recorder::default();
for _ in 0..3 {
rec.push(vec![
LlmStreamEvent::TextDelta("x".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
}
let base = LlmProvider {
endpoint: LlmEndpoint::Chat,
base_url: Some("http://base".into()),
api_key: Some("base-key".into()),
};
let mut profiles = ModelProfiles::new(base.clone());
profiles.insert(
"m2",
LlmProfile {
endpoint: Some(LlmEndpoint::Responses),
base_url: Some("http://profile".into()),
api_key: Some("profile-key".into()),
},
);
let llm = LlmOpts {
model: "m1".into(),
endpoint: base.endpoint,
base_url: base.base_url.clone(),
api_key: base.api_key.clone(),
model_profiles: profiles,
..LlmOpts::default()
};
let dir = tempfile::tempdir().unwrap();
let session = Session::create(dir.path()).unwrap();
let mut agent = Agent::new(agent_config_with_llm(
fake_stream_fn(&rec),
vec![],
Some(session),
llm,
));
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("one", tx).await.unwrap();
agent.set_model("m2".into()).unwrap();
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("two", tx).await.unwrap();
agent.set_model("m1".into()).unwrap();
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("three", tx).await.unwrap();
let provider = |i: usize| {
let c = &rec.calls()[i];
(
c.model.clone(),
c.endpoint,
c.base_url.clone(),
c.api_key.clone(),
)
};
assert_eq!(
[provider(0), provider(1), provider(2)],
[
(
"m1".to_string(),
LlmEndpoint::Chat,
Some("http://base".to_string()),
Some("base-key".to_string()),
),
(
"m2".to_string(),
LlmEndpoint::Responses,
Some("http://profile".to_string()),
Some("profile-key".to_string()),
),
(
"m1".to_string(),
LlmEndpoint::Chat,
Some("http://base".to_string()),
Some("base-key".to_string()),
),
],
"a switch must settle on the launch provider, not leak the previous model's"
);
}
#[tokio::test]
async fn set_model_on_unmapped_id_settles_to_the_base_provider() {
let rec = Recorder::default();
for _ in 0..2 {
rec.push(vec![
LlmStreamEvent::TextDelta("x".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
}
let base = LlmProvider {
endpoint: LlmEndpoint::Chat,
base_url: Some("http://base".into()),
api_key: Some("base-key".into()),
};
let mut profiles = ModelProfiles::new(base.clone());
profiles.insert(
"other",
LlmProfile {
endpoint: Some(LlmEndpoint::Responses),
base_url: Some("http://other".into()),
api_key: None,
},
);
let llm = LlmOpts {
model: "m1".into(),
endpoint: base.endpoint,
base_url: base.base_url.clone(),
api_key: base.api_key.clone(),
model_profiles: profiles,
..LlmOpts::default()
};
let dir = tempfile::tempdir().unwrap();
let session = Session::create(dir.path()).unwrap();
let mut agent = Agent::new(agent_config_with_llm(
fake_stream_fn(&rec),
vec![],
Some(session),
llm,
));
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("one", tx).await.unwrap();
agent.set_model("m2".into()).unwrap();
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("two", tx).await.unwrap();
let c = &rec.calls()[1];
assert_eq!(c.model, "m2");
assert_eq!(c.endpoint, LlmEndpoint::Chat, "unmapped id === the base");
assert_eq!(c.base_url.as_deref(), Some("http://base"));
assert_eq!(c.api_key.as_deref(), Some("base-key"));
}
#[test]
fn all_none_profile_leaves_the_base_provider() {
let base = LlmProvider {
endpoint: LlmEndpoint::Chat,
base_url: Some("http://base/v1".into()),
api_key: Some("base-key".into()),
};
let mut profiles = ModelProfiles::new(base.clone());
profiles.insert("m", LlmProfile::default());
let mut llm = LlmOpts {
model: "m".into(),
endpoint: LlmEndpoint::Chat,
base_url: Some("http://base/v1".into()),
api_key: Some("base-key".into()),
model_profiles: profiles,
..LlmOpts::default()
};
llm.settle_provider();
assert_eq!(llm.endpoint, base.endpoint);
assert_eq!(llm.base_url, base.base_url);
assert_eq!(llm.api_key, base.api_key);
}
#[test]
fn profile_with_no_endpoint_keeps_the_base_endpoint() {
let mut profiles = ModelProfiles::new(LlmProvider {
endpoint: LlmEndpoint::Responses,
base_url: Some("http://base/v1".into()),
api_key: None,
});
profiles.insert(
"m",
LlmProfile {
endpoint: None,
base_url: Some("http://over/v1".into()),
api_key: None,
},
);
let mut llm = LlmOpts {
model: "m".into(),
endpoint: LlmEndpoint::Chat,
model_profiles: profiles,
..LlmOpts::default()
};
llm.settle_provider();
assert_eq!(
llm.endpoint,
LlmEndpoint::Responses,
"a None endpoint inherits the base endpoint, not Chat"
);
assert_eq!(llm.base_url.as_deref(), Some("http://over/v1"));
}
#[test]
fn default_model_profiles_settle_is_a_noop() {
let mut llm = LlmOpts {
model: "m".into(),
endpoint: LlmEndpoint::Responses,
base_url: Some("http://keep".into()),
api_key: Some("keep".into()),
..LlmOpts::default()
};
llm.settle_provider();
assert_eq!(llm.endpoint, LlmEndpoint::Responses);
assert_eq!(llm.base_url.as_deref(), Some("http://keep"));
assert_eq!(llm.api_key.as_deref(), Some("keep"));
}
#[tokio::test]
async fn set_effort_swaps_llm_and_logs_session_change() {
let rec = Recorder::default();
for _ in 0..3 {
rec.push(vec![
LlmStreamEvent::TextDelta("x".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
}
let rec2 = rec.clone();
let effort_seen = Arc::new(Mutex::new(Vec::new()));
let stream_fn = {
let seen = effort_seen.clone();
let inner = fake_stream_fn(&rec2);
Arc::new(
move |ctx: &[AgentMessage],
system: &str,
tools: &[rig::completion::ToolDefinition],
opts: &LlmOpts| {
seen.lock().unwrap().push(opts.effort.clone());
inner(ctx, system, tools, opts)
},
) as StreamFn
};
let dir = tempfile::tempdir().unwrap();
let session = Session::create(dir.path()).unwrap();
let path = session.path().unwrap().to_path_buf();
let mut agent = Agent::new(agent_config(stream_fn, vec![], Some(session)));
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("hi", tx).await.unwrap();
agent.set_effort(Some("high".into())).unwrap();
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("again", tx).await.unwrap();
agent.set_effort(None).unwrap();
let (tx, _rx) = mpsc::unbounded_channel();
agent.run("third", tx).await.unwrap();
assert_eq!(
*effort_seen.lock().unwrap(),
vec![None, Some("high".to_string()), None]
);
let reopened = Session::open(&path).unwrap();
assert_eq!(reopened.effort(), Some(None));
}
#[tokio::test]
async fn compact_replaces_prefix_with_a_summary() {
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("Context: condensed history".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let mut cfg = agent_config(fake_stream_fn(&rec), vec![], None);
cfg.context = (0..20)
.map(|i| AgentMessage::user_text(format!("message {i} {}", "x".repeat(400))))
.collect();
cfg.compaction = CompactionPolicy {
keep_recent_tokens: 200,
keep_recent_turns: 0,
..Default::default()
};
let mut agent = Agent::new(cfg);
let before = agent.messages().len();
match agent.compact(None).await.unwrap() {
CompactOutcome::Done {
summarized, kept, ..
} => {
assert_eq!(summarized + kept, before, "accounting covers all of ctx");
assert_eq!(kept, 2, "kept the newest messages within the budget");
assert_eq!(agent.messages().len(), kept + 1, "summary + kept");
assert!(
agent.messages()[0].as_text().contains("condensed history"),
"the summary replaces the summarized prefix"
);
}
other => panic!("expected Done, got {other:?}"),
}
}
#[tokio::test]
async fn compact_is_a_noop_when_history_fits_the_budget() {
let rec = Recorder::default();
let mut cfg = agent_config(fake_stream_fn(&rec), vec![], None);
cfg.context = vec![AgentMessage::user_text("short")];
let mut agent = Agent::new(cfg);
let out = agent.compact(None).await.unwrap();
assert!(matches!(out, CompactOutcome::NothingToDo));
assert_eq!(agent.messages().len(), 1, "nothing dropped");
assert!(rec.calls().is_empty(), "the summarizer was never called");
}
#[tokio::test]
async fn compact_persists_and_resumes_to_summary_plus_tail() {
let dir = tempfile::tempdir().unwrap();
let rec = Recorder::default();
rec.push(vec![
LlmStreamEvent::TextDelta("summary body".into()),
LlmStreamEvent::Done {
stop_reason: StopReason::Stop,
usage: None,
},
]);
let mut session = Session::create(dir.path()).unwrap();
for i in 0..10 {
session
.append(SessionEntry::Message {
id: format!("id{i}"),
parent_id: None,
message: AgentMessage::user_text(format!("m{i} {}", "x".repeat(400))),
})
.unwrap();
}
let path = session.path().unwrap().to_path_buf();
let seeded = session.messages();
let mut cfg = agent_config(fake_stream_fn(&rec), vec![], Some(session));
cfg.context = seeded;
cfg.compaction = CompactionPolicy {
keep_recent_tokens: 200, keep_recent_turns: 0,
..Default::default()
};
let mut agent = Agent::new(cfg);
agent.compact(None).await.unwrap();
let ctx = agent.messages();
assert_eq!(ctx.len(), 3, "summary + two kept");
assert!(ctx[0].as_text().contains("summary body"));
assert!(ctx[1].as_text().starts_with("m8 "));
assert!(ctx[2].as_text().starts_with("m9 "));
let reopened = Session::open(&path).unwrap();
let resumed: Vec<String> = reopened.messages().iter().map(|m| m.as_text()).collect();
assert_eq!(resumed.len(), 3, "summary + two kept after resume");
assert!(resumed[0].contains("summary body"), "summary leads: {resumed:?}");
assert!(resumed[1].starts_with("m8 "));
assert!(resumed[2].starts_with("m9 "));
assert_eq!(
reopened.message_count(),
10,
"all messages stay on disk; only the view is compacted"
);
}