mod approve;
mod config;
mod queue;
mod run;
mod session;
mod sink;
mod terminal;
pub use approve::{AllowAll, ApprovalRequest, Approver, Decision};
pub use config::EngineConfig;
pub use queue::{InputQueue, MID_RUN_PREAMBLE};
pub use session::Session;
pub use sink::{EventSink, FnSink, NullSink};
pub use tokio_util::sync::CancellationToken;
#[cfg(test)]
mod tests {
#![allow(clippy::unnecessary_literal_bound)]
use super::*;
use async_trait::async_trait;
use locode_protocol::{
ContentBlock, Conversation, Event, Message, ReasoningFormat, Role, Status, Usage,
reconstruct_conversation,
};
use locode_provider::{
Completion, ConversationRequest, MockProvider, Provider, ProviderError, StopReason,
};
use locode_tools::{Registry, Tool, ToolCtx, ToolError, ToolKind, ToolOutput};
use serde::Serialize;
use serde_json::{Value, json};
use std::sync::{Arc, Mutex};
use std::time::Duration;
#[derive(Serialize)]
struct EchoOut {
echoed: String,
}
impl ToolOutput for EchoOut {
fn to_prompt_text(&self) -> String {
self.echoed.clone()
}
}
struct Echo;
#[async_trait]
impl Tool for Echo {
type Args = Value;
type Output = EchoOut;
fn kind(&self) -> ToolKind {
ToolKind::Shell
}
fn description(&self) -> &str {
"echo"
}
async fn run(&self, _ctx: &ToolCtx, args: Value) -> Result<EchoOut, ToolError> {
Ok(EchoOut {
echoed: args.to_string(),
})
}
}
struct Boom;
#[async_trait]
impl Tool for Boom {
type Args = Value;
type Output = EchoOut;
fn kind(&self) -> ToolKind {
ToolKind::Shell
}
fn description(&self) -> &str {
"boom"
}
async fn run(&self, _ctx: &ToolCtx, _args: Value) -> Result<EchoOut, ToolError> {
Err(ToolError::Fatal("boom aborted the turn".into()))
}
}
fn text_turn(text: &str) -> Completion {
Completion {
content: vec![ContentBlock::Text { text: text.into() }],
usage: Usage::default(),
stop: StopReason::EndTurn,
}
}
fn tool_turn(id: &str, name: &str) -> Completion {
Completion {
content: vec![ContentBlock::ToolUse {
id: id.into(),
name: name.into(),
input: json!({}),
}],
usage: Usage::default(),
stop: StopReason::ToolUse,
}
}
fn config() -> EngineConfig {
EngineConfig {
session_id: "sess-1".into(),
harness: "grok".into(),
api_schema: "mock".into(),
model: "mock-1".into(),
max_turns: None,
resample_retries: 2,
resample_backoff: Duration::ZERO, instructions: locode_instructions::InstructionsConfig {
enabled: false,
..Default::default()
},
..EngineConfig::default()
}
}
fn session_with(
script: Vec<Result<Completion, ProviderError>>,
registry: Registry,
cfg: EngineConfig,
) -> (Session, Arc<Mutex<Vec<Event>>>) {
let events = Arc::new(Mutex::new(Vec::new()));
let sink_events = Arc::clone(&events);
let sink = Box::new(FnSink(move |event| {
sink_events.lock().unwrap().push(event);
}));
let provider = Arc::new(MockProvider::with_results(script));
let session = Session::new(provider, registry, vec![], cfg, sink);
(session, events)
}
fn echo_registry() -> Registry {
let mut reg = Registry::new();
reg.register("echo", Echo);
reg
}
fn dump(events: &Arc<Mutex<Vec<Event>>>) -> Vec<Event> {
events.lock().unwrap().clone()
}
#[tokio::test]
async fn completed_with_no_tools() {
let (mut s, events) =
session_with(vec![Ok(text_turn("all done"))], Registry::new(), config());
let report = s.run_text("hi").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(report.final_message.as_deref(), Some("all done"));
assert_eq!(report.turns, 1);
assert!(report.tool_calls.is_empty());
assert_eq!(report.api_schema, "mock");
let evs = dump(&events);
assert!(matches!(evs.first(), Some(Event::Init { .. })));
assert!(matches!(evs.last(), Some(Event::Result { .. })));
}
#[tokio::test]
async fn streaming_run_emits_text_deltas_and_the_whole_message() {
let mut cfg = config();
cfg.streaming = true;
let (mut s, events) = session_with(
vec![Ok(text_turn("hello streamed world"))],
Registry::new(),
cfg,
);
let report = s.run_text("hi").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(
report.final_message.as_deref(),
Some("hello streamed world")
);
let evs = dump(&events);
let delta_text: String = evs
.iter()
.filter_map(|e| match e {
Event::MessageDelta { text } => Some(text.as_str()),
_ => None,
})
.collect();
assert_eq!(delta_text, "hello streamed world", "{evs:?}");
let n_deltas = evs
.iter()
.filter(|e| matches!(e, Event::MessageDelta { .. }))
.count();
assert!(n_deltas > 1, "expected multiple deltas, got {n_deltas}");
assert!(
evs.iter().any(|e| matches!(
e,
Event::Message { message } if message.role == Role::Assistant
)),
"whole assistant Message still emitted: {evs:?}"
);
let first_delta = evs
.iter()
.position(|e| matches!(e, Event::MessageDelta { .. }))
.expect("a delta");
let asst_msg = evs
.iter()
.position(
|e| matches!(e, Event::Message { message } if message.role == Role::Assistant),
)
.expect("assistant message");
assert!(
first_delta < asst_msg,
"deltas come before the whole message"
);
}
#[tokio::test]
async fn non_streaming_run_emits_no_deltas() {
let (mut s, events) =
session_with(vec![Ok(text_turn("no stream"))], Registry::new(), config());
let _ = s.run_text("hi").await;
let evs = dump(&events);
assert!(
!evs.iter().any(|e| matches!(e, Event::MessageDelta { .. })),
"default (non-streaming) run must not emit deltas: {evs:?}"
);
}
#[tokio::test]
async fn streaming_and_non_streaming_reports_match() {
let (mut a, _ea) = session_with(
vec![Ok(text_turn("same result"))],
Registry::new(),
config(),
);
let mut cfg = config();
cfg.streaming = true;
let (mut b, _eb) = session_with(vec![Ok(text_turn("same result"))], Registry::new(), cfg);
let ra = a.run_text("go").await;
let rb = b.run_text("go").await;
assert_eq!(ra.status, rb.status);
assert_eq!(ra.final_message, rb.final_message);
assert_eq!(ra.turns, rb.turns);
assert_eq!(ra.tool_calls.len(), rb.tool_calls.len());
}
#[tokio::test]
async fn tool_call_then_complete() {
let (mut s, _e) = session_with(
vec![Ok(tool_turn("c1", "echo")), Ok(text_turn("done"))],
echo_registry(),
config(),
);
let report = s.run_text("go").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(report.turns, 2);
assert_eq!(report.tool_calls.len(), 1);
assert!(report.tool_calls[0].ok);
assert_eq!(report.tool_calls[0].name, "echo");
}
#[tokio::test]
async fn hits_max_turns_after_dispatch() {
let mut cfg = config();
cfg.max_turns = Some(2);
let (mut s, _e) = session_with(
vec![
Ok(tool_turn("c1", "echo")),
Ok(tool_turn("c2", "echo")),
Ok(tool_turn("c3", "echo")),
],
echo_registry(),
cfg,
);
let report = s.run_text("go").await;
assert_eq!(report.status, Status::MaxTurns);
assert_eq!(report.turns, 2);
assert_eq!(report.tool_calls.len(), 2);
}
#[tokio::test]
async fn model_error_after_bounded_retry() {
let script = vec![
Err(ProviderError::Transport("reset".into())),
Err(ProviderError::Transport("reset".into())),
Err(ProviderError::Transport("reset".into())),
];
let (mut s, events) = session_with(script, Registry::new(), config());
let report = s.run_text("go").await;
assert_eq!(report.status, Status::ModelError);
assert!(report.error.is_some());
assert_eq!(report.turns, 0);
let retries = dump(&events)
.iter()
.filter(|e| matches!(e, Event::Error { .. }))
.count();
assert_eq!(retries, 2);
}
#[tokio::test]
async fn model_error_non_retryable_is_immediate() {
let (mut s, events) = session_with(
vec![Err(ProviderError::ContextOverflow)],
Registry::new(),
config(),
);
let report = s.run_text("go").await;
assert_eq!(report.status, Status::ModelError);
let retries = dump(&events)
.iter()
.filter(|e| matches!(e, Event::Error { .. }))
.count();
assert_eq!(retries, 0, "a non-retryable error must not resample");
}
#[tokio::test]
async fn fatal_tool_error_ends_the_run() {
let mut reg = Registry::new();
reg.register("boom", Boom);
let (mut s, _e) = session_with(vec![Ok(tool_turn("c1", "boom"))], reg, config());
let report = s.run_text("go").await;
assert_eq!(report.status, Status::Error);
assert!(report.error.is_some());
assert_eq!(report.tool_calls.len(), 1);
assert!(!report.tool_calls[0].ok);
}
#[tokio::test]
async fn empty_completion_resamples_then_succeeds() {
let empty = Completion {
content: vec![ContentBlock::Reasoning {
format: ReasoningFormat::Anthropic,
text: "thinking only".into(),
signature: Some("sig".into()),
payload: None,
}],
usage: Usage::default(),
stop: StopReason::MaxTokens,
};
let (mut session, _events) = session_with(
vec![Ok(empty), Ok(text_turn("recovered"))],
echo_registry(),
config(),
);
let report = session.run_text("go").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(report.final_message.as_deref(), Some("recovered"));
assert_eq!(report.stop_reason.as_deref(), Some("end_turn"));
}
#[tokio::test]
async fn queued_input_rides_the_tool_result_batch_after_the_results() {
let (mut session, events) = session_with(
vec![Ok(tool_turn("c1", "echo")), Ok(text_turn("done"))],
echo_registry(),
config(),
);
let queue = session.input_queue();
queue.push("actually, use tabs");
let report = session.run_text("go").await;
assert_eq!(report.status, Status::Completed);
let batch = dump(&events)
.into_iter()
.find_map(|e| match e {
Event::Message { message }
if message.role == Role::User
&& message
.content
.iter()
.any(|b| matches!(b, ContentBlock::ToolResult { .. })) =>
{
Some(message)
}
_ => None,
})
.expect("a tool-result batch was appended");
let kinds: Vec<&str> = batch
.content
.iter()
.map(|b| match b {
ContentBlock::ToolResult { .. } => "result",
ContentBlock::Text { .. } => "text",
_ => "other",
})
.collect();
assert_eq!(
kinds,
vec!["result", "text"],
"the queued text must follow the results, never precede them"
);
let text = batch
.content
.iter()
.find_map(|b| match b {
ContentBlock::Text { text } => Some(text.clone()),
_ => None,
})
.expect("the queued text landed");
assert!(
text.starts_with(crate::MID_RUN_PREAMBLE),
"the mid-run path is marked: {text}"
);
assert!(text.contains("actually, use tabs"));
assert!(queue.is_empty(), "draining consumes");
}
#[tokio::test]
async fn queued_input_with_no_tool_calls_stays_for_the_fallback() {
let (mut session, _events) = session_with(
vec![Ok(text_turn("nothing to do"))],
echo_registry(),
config(),
);
let queue = session.input_queue();
queue.push("one more thing");
let report = session.run_text("go").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(
queue.pending(),
vec!["one more thing".to_string()],
"no carrier this run — the frontend submits it as an ordinary prompt"
);
}
#[tokio::test]
async fn truncated_tool_call_is_not_executed_and_names_the_cause() {
let truncated = Completion {
content: vec![ContentBlock::ToolUse {
id: "c1".into(),
name: "echo".into(),
input: json!({}), }],
usage: Usage::default(),
stop: StopReason::MaxTokens,
};
let (mut session, events) = session_with(
vec![Ok(truncated), Ok(text_turn("smaller this time"))],
echo_registry(),
config(),
);
let report = session.run_text("write a huge file").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(report.final_message.as_deref(), Some("smaller this time"));
assert!(
report.tool_calls.is_empty(),
"a call that never ran is not recorded, matching the cancel path"
);
let explained = dump(&events).iter().any(|e| match e {
Event::Message { message } => message.content.iter().any(|b| {
matches!(
b,
ContentBlock::ToolResult { tool_use_id, is_error: true, content, .. }
if tool_use_id == "c1"
&& content.iter().any(|c| matches!(
c,
locode_protocol::ResultChunk::Text { text }
if text.contains("output-token limit")
&& text.contains("max_tokens")
&& text.contains("Do not repeat the call unchanged")
))
)
}),
_ => false,
});
assert!(
explained,
"the model must see the truncation, not a 'missing field' decode error"
);
}
#[tokio::test]
async fn truncation_after_a_finished_tool_call_still_dispatches() {
let cut_after_call = Completion {
content: vec![
ContentBlock::ToolUse {
id: "c1".into(),
name: "echo".into(),
input: json!({"complete": true}),
},
ContentBlock::Text {
text: "and then I was cut off mid-sent".into(),
},
],
usage: Usage::default(),
stop: StopReason::MaxTokens,
};
let (mut session, _events) = session_with(
vec![Ok(cut_after_call), Ok(text_turn("done"))],
echo_registry(),
config(),
);
let report = session.run_text("go").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(report.tool_calls.len(), 1, "the finished call still ran");
assert!(report.tool_calls[0].ok);
}
#[tokio::test]
async fn persistent_empty_completions_are_model_error() {
let empty = || Completion {
content: vec![],
usage: Usage::default(),
stop: StopReason::MaxTokens,
};
let (mut session, _events) = session_with(
vec![Ok(empty()), Ok(empty()), Ok(empty())],
echo_registry(),
config(),
);
let report = session.run_text("go").await;
assert_eq!(report.status, Status::ModelError);
assert!(
report
.error
.as_deref()
.unwrap_or("")
.contains("empty completion"),
"error names the cause: {:?}",
report.error
);
assert_eq!(report.stop_reason, None, "no completion was accepted");
}
#[tokio::test]
async fn mid_batch_abort_synthesizes_results() {
let mut reg = Registry::new();
reg.register("boom", Boom);
reg.register("echo", Echo);
let completion = Completion {
content: vec![
ContentBlock::ToolUse {
id: "c_boom".into(),
name: "boom".into(),
input: json!({}),
},
ContentBlock::ToolUse {
id: "c_echo".into(),
name: "echo".into(),
input: json!({}),
},
],
usage: Usage::default(),
stop: StopReason::ToolUse,
};
let (mut s, events) = session_with(vec![Ok(completion)], reg, config());
let report = s.run_text("go").await;
assert_eq!(report.status, Status::Error);
let evs = dump(&events);
let answered: Vec<String> = evs
.iter()
.filter_map(|e| match e {
Event::Message { message } if message.role == Role::User => Some(&message.content),
_ => None,
})
.flatten()
.filter_map(|b| match b {
ContentBlock::ToolResult { tool_use_id, .. } => Some(tool_use_id.clone()),
_ => None,
})
.collect();
assert!(answered.iter().any(|id| id == "c_boom"));
assert!(
answered.iter().any(|id| id == "c_echo"),
"the un-run echo must be paired"
);
assert_eq!(report.tool_calls.len(), 1);
}
#[tokio::test]
async fn thinking_block_is_appended_verbatim() {
let completion = Completion {
content: vec![
ContentBlock::Reasoning {
format: ReasoningFormat::Anthropic,
text: "reasoning".into(),
signature: Some("sig-xyz".into()),
payload: None,
},
ContentBlock::Text {
text: "answer".into(),
},
],
usage: Usage::default(),
stop: StopReason::EndTurn,
};
let (mut s, events) = session_with(vec![Ok(completion)], Registry::new(), config());
let report = s.run_text("think").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(report.final_message.as_deref(), Some("answer"));
let has_thinking = dump(&events).iter().any(|e| match e {
Event::Message { message } if message.role == Role::Assistant => {
message.content.iter().any(|b| {
matches!(
b,
ContentBlock::Reasoning { signature: Some(sig), .. } if sig == "sig-xyz"
)
})
}
_ => false,
});
assert!(
has_thinking,
"thinking + signature must survive into history"
);
}
#[tokio::test]
async fn events_reconstruct_the_history() {
let (mut s, events) = session_with(
vec![Ok(tool_turn("c1", "echo")), Ok(text_turn("done"))],
echo_registry(),
config(),
);
let _ = s.run_text("go").await;
let rebuilt: Conversation = reconstruct_conversation(&dump(&events));
let roles: Vec<Role> = rebuilt.messages.iter().map(|m| m.role).collect();
assert_eq!(
roles,
vec![Role::User, Role::Assistant, Role::User, Role::Assistant]
);
}
use std::sync::atomic::{AtomicUsize, Ordering};
struct Counting(Arc<AtomicUsize>);
#[async_trait]
impl Tool for Counting {
type Args = Value;
type Output = EchoOut;
fn kind(&self) -> ToolKind {
ToolKind::Shell
}
fn description(&self) -> &str {
"counting"
}
async fn run(&self, _ctx: &ToolCtx, _args: Value) -> Result<EchoOut, ToolError> {
self.0.fetch_add(1, Ordering::SeqCst);
Ok(EchoOut {
echoed: "ran".into(),
})
}
}
type SeenKinds = Arc<Mutex<Vec<(String, Option<ToolKind>)>>>;
struct DenyNamed {
deny: Vec<&'static str>,
seen_kinds: SeenKinds,
}
#[async_trait]
impl Approver for DenyNamed {
async fn decide(&self, request: &ApprovalRequest<'_>) -> Decision {
self.seen_kinds
.lock()
.unwrap()
.push((request.tool_name.to_owned(), request.kind));
if self.deny.contains(&request.tool_name) {
Decision::Deny {
reason: format!("{} is not allowed here", request.tool_name),
}
} else {
Decision::Allow
}
}
}
fn approvals(events: &Arc<Mutex<Vec<Event>>>) -> Vec<(String, String, String)> {
dump(events)
.iter()
.filter_map(|e| match e {
Event::Approval {
tool_use_id,
tool_name,
decision,
..
} => Some((tool_use_id.clone(), tool_name.clone(), decision.clone())),
_ => None,
})
.collect()
}
#[tokio::test]
async fn deny_is_a_soft_paired_error_and_the_run_continues() {
let ran = Arc::new(AtomicUsize::new(0));
let mut reg = Registry::new();
reg.register("counting", Counting(Arc::clone(&ran)));
let (s, events) = session_with(
vec![Ok(tool_turn("c1", "counting")), Ok(text_turn("done"))],
reg,
config(),
);
let seen = Arc::new(Mutex::new(Vec::new()));
let mut s = s.with_approver(Arc::new(DenyNamed {
deny: vec!["counting"],
seen_kinds: Arc::clone(&seen),
}));
let report = s.run_text("go").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(ran.load(Ordering::SeqCst), 0, "denied tool must not run");
assert_eq!(report.tool_calls.len(), 1);
let record = &report.tool_calls[0];
assert!(!record.ok);
assert_eq!(
record.denial_reason.as_deref(),
Some("counting is not allowed here")
);
assert_eq!(record.kind, "shell", "kind still recorded on denial");
let denied_result = dump(&events).iter().any(|e| match e {
Event::Message { message } => message.content.iter().any(|b| {
matches!(
b,
ContentBlock::ToolResult { tool_use_id, is_error: true, content, .. }
if tool_use_id == "c1"
&& content.iter().any(|c| matches!(
c,
locode_protocol::ResultChunk::Text { text }
if text == "tool call denied: counting is not allowed here"
))
)
}),
_ => false,
});
assert!(denied_result, "the model sees the denial reason, paired");
assert_eq!(
approvals(&events),
vec![("c1".into(), "counting".into(), "deny".into())]
);
}
#[tokio::test]
async fn deny_then_allow_within_one_batch_keeps_order_and_pairing() {
let ran = Arc::new(AtomicUsize::new(0));
let mut reg = Registry::new();
reg.register("blocked", Counting(Arc::clone(&ran)));
reg.register("echo", Echo);
let batch = Completion {
content: vec![
ContentBlock::ToolUse {
id: "c1".into(),
name: "blocked".into(),
input: json!({}),
},
ContentBlock::ToolUse {
id: "c2".into(),
name: "echo".into(),
input: json!({}),
},
],
usage: Usage::default(),
stop: StopReason::ToolUse,
};
let (s, events) = session_with(vec![Ok(batch), Ok(text_turn("done"))], reg, config());
let mut s = s.with_approver(Arc::new(DenyNamed {
deny: vec!["blocked"],
seen_kinds: Arc::new(Mutex::new(Vec::new())),
}));
let report = s.run_text("go").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(ran.load(Ordering::SeqCst), 0);
let pairs: Vec<(String, bool)> = dump(&events)
.iter()
.filter_map(|e| match e {
Event::Message { message } if message.role == Role::User => Some(&message.content),
_ => None,
})
.flatten()
.filter_map(|b| match b {
ContentBlock::ToolResult {
tool_use_id,
is_error,
..
} => Some((tool_use_id.clone(), *is_error)),
_ => None,
})
.collect();
assert_eq!(pairs, vec![("c1".into(), true), ("c2".into(), false)]);
assert_eq!(report.tool_calls.len(), 2);
assert!(report.tool_calls[0].denial_reason.is_some());
assert_eq!(report.tool_calls[0].kind, "shell");
assert!(report.tool_calls[1].ok);
assert_eq!(report.tool_calls[1].denial_reason, None);
assert_eq!(
approvals(&events),
vec![
("c1".into(), "blocked".into(), "deny".into()),
("c2".into(), "echo".into(), "allow".into()),
]
);
}
#[tokio::test]
async fn approval_request_carries_the_registry_kind() {
let seen = Arc::new(Mutex::new(Vec::new()));
let (s, _e) = session_with(
vec![Ok(tool_turn("c1", "echo")), Ok(text_turn("done"))],
echo_registry(),
config(),
);
let mut s = s.with_approver(Arc::new(DenyNamed {
deny: vec![],
seen_kinds: Arc::clone(&seen),
}));
let _ = s.run_text("go").await;
let seen = seen.lock().unwrap();
assert_eq!(seen.len(), 1);
assert_eq!(seen[0].0, "echo");
assert_eq!(
seen[0].1,
Some(ToolKind::Shell),
"kind resolves from the registry pre-dispatch"
);
}
#[tokio::test]
async fn async_approver_suspends_the_call_until_resolved() {
struct OneshotApprover(Mutex<Option<tokio::sync::oneshot::Receiver<Decision>>>);
#[async_trait]
impl Approver for OneshotApprover {
async fn decide(&self, _request: &ApprovalRequest<'_>) -> Decision {
let rx = self.0.lock().unwrap().take().expect("one decision");
rx.await.expect("decider dropped")
}
}
let (tx, rx) = tokio::sync::oneshot::channel();
let (s, _e) = session_with(
vec![Ok(tool_turn("c1", "echo")), Ok(text_turn("done"))],
echo_registry(),
config(),
);
let mut s = s.with_approver(Arc::new(OneshotApprover(Mutex::new(Some(rx)))));
let ui = tokio::spawn(async move {
tokio::task::yield_now().await;
let _ = tx.send(Decision::Allow);
});
let report = s.run_text("go").await;
ui.await.expect("ui task");
assert_eq!(report.status, Status::Completed);
assert_eq!(report.tool_calls.len(), 1);
assert!(report.tool_calls[0].ok);
}
#[tokio::test]
async fn allowed_calls_emit_approval_events_by_default() {
let (mut s, events) = session_with(
vec![Ok(tool_turn("c1", "echo")), Ok(text_turn("done"))],
echo_registry(),
config(),
);
let report = s.run_text("go").await;
assert_eq!(report.status, Status::Completed);
assert_eq!(
approvals(&events),
vec![("c1".into(), "echo".into(), "allow".into())]
);
assert_eq!(report.tool_calls[0].denial_reason, None);
}
struct HalfStreamsThenFails {
attempts: std::sync::atomic::AtomicU32,
}
#[async_trait]
impl Provider for HalfStreamsThenFails {
#[allow(clippy::unnecessary_literal_bound)]
fn api_schema(&self) -> &str {
"mock"
}
async fn complete(
&self,
_request: &ConversationRequest,
) -> Result<Completion, ProviderError> {
unreachable!("this test runs streaming")
}
async fn stream(
&self,
_request: &ConversationRequest,
on_delta: &mut (dyn FnMut(locode_provider::CompletionDelta) + Send),
) -> Result<Completion, ProviderError> {
let n = self
.attempts
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if n == 0 {
on_delta(locode_provider::CompletionDelta::Text("Hel".into()));
on_delta(locode_provider::CompletionDelta::Text("lo wor".into()));
return Err(ProviderError::Transport("lossy stream".into()));
}
on_delta(locode_provider::CompletionDelta::Text("Hello world".into()));
Ok(Completion {
content: vec![ContentBlock::Text {
text: "Hello world".into(),
}],
usage: Usage::default(),
stop: locode_provider::StopReason::EndTurn,
})
}
}
#[tokio::test]
async fn a_partial_stream_that_resamples_annuls_its_deltas() {
let mut cfg = config();
cfg.streaming = true;
let provider = std::sync::Arc::new(HalfStreamsThenFails {
attempts: std::sync::atomic::AtomicU32::new(0),
});
let events = Arc::new(Mutex::new(Vec::new()));
let sink_events = Arc::clone(&events);
let sink = Box::new(FnSink(move |event| {
sink_events.lock().unwrap().push(event);
}));
let mut session = Session::new(provider, Registry::new(), vec![], cfg, sink);
let report = session.run_text("go").await;
assert_eq!(report.status, Status::Completed, "the retry succeeded");
let evs = dump(&events);
let resets = evs
.iter()
.filter(|e| matches!(e, Event::MessageDeltaReset { .. }))
.count();
assert_eq!(resets, 1, "exactly one annulment, for the failed attempt");
let reset_at = evs
.iter()
.position(|e| matches!(e, Event::MessageDeltaReset { .. }))
.expect("reset emitted");
let last_delta = evs
.iter()
.rposition(|e| matches!(e, Event::MessageDelta { .. }))
.expect("the retry streamed");
assert!(
reset_at < last_delta,
"reset must come before the re-stream"
);
}
struct HangingProvider;
#[async_trait]
impl Provider for HangingProvider {
#[allow(clippy::unnecessary_literal_bound)]
fn api_schema(&self) -> &str {
"mock"
}
async fn complete(
&self,
_request: &ConversationRequest,
) -> Result<Completion, ProviderError> {
tokio::time::sleep(Duration::from_hours(1)).await;
Err(ProviderError::Transport("unreachable".into()))
}
}
struct WaitsForCancel;
#[async_trait]
impl Tool for WaitsForCancel {
type Args = Value;
type Output = EchoOut;
fn kind(&self) -> ToolKind {
ToolKind::Shell
}
fn description(&self) -> &str {
"waits"
}
async fn run(&self, ctx: &ToolCtx, _args: Value) -> Result<EchoOut, ToolError> {
ctx.cancel.cancelled().await;
Ok(EchoOut {
echoed: "stopped cooperatively".into(),
})
}
}
#[tokio::test]
async fn cancel_mid_sample_yields_cancelled_report() {
let events = Arc::new(Mutex::new(Vec::new()));
let sink_events = Arc::clone(&events);
let sink = Box::new(FnSink(move |event| {
sink_events.lock().unwrap().push(event);
}));
let mut s = Session::new(
Arc::new(HangingProvider),
Registry::new(),
vec![],
config(),
sink,
);
let handle = s.cancel_handle();
let canceller = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(20)).await;
handle.cancel();
handle.cancel(); });
let report = s.run_text("go").await;
canceller.await.expect("canceller");
assert_eq!(report.status, Status::Cancelled);
assert_eq!(report.error, None, "cancelled is a stop, not a fault");
assert_eq!(report.final_message, None, "no assistant text this run");
assert_eq!(report.turns, 0, "no completion was accepted");
let roles: Vec<Role> = s.history().iter().map(|m| m.role).collect();
assert_eq!(roles, vec![Role::User]);
let evs = dump(&events);
assert!(
matches!(evs.last(), Some(Event::Result { report }) if report.status == Status::Cancelled)
);
}
#[tokio::test]
async fn cancel_mid_batch_pairs_the_rest_synthetically() {
let mut reg = Registry::new();
reg.register("waits", WaitsForCancel);
reg.register("echo", Echo);
let batch = Completion {
content: vec![
ContentBlock::ToolUse {
id: "c_wait".into(),
name: "waits".into(),
input: json!({}),
},
ContentBlock::ToolUse {
id: "c_echo".into(),
name: "echo".into(),
input: json!({}),
},
],
usage: Usage::default(),
stop: StopReason::ToolUse,
};
let (s, events) = session_with(vec![Ok(batch)], reg, config());
let mut s = s; let handle = s.cancel_handle();
let canceller = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(20)).await;
handle.cancel();
});
let report = s.run_text("go").await;
canceller.await.expect("canceller");
assert_eq!(report.status, Status::Cancelled);
assert_eq!(report.tool_calls.len(), 1);
assert_eq!(report.tool_calls[0].id, "c_wait");
assert!(report.tool_calls[0].ok);
assert_eq!(report.tool_calls[0].denial_reason, None);
let pairs: Vec<(String, bool)> = dump(&events)
.iter()
.filter_map(|e| match e {
Event::Message { message } if message.role == Role::User => Some(&message.content),
_ => None,
})
.flatten()
.filter_map(|b| match b {
ContentBlock::ToolResult {
tool_use_id,
is_error,
..
} => Some((tool_use_id.clone(), *is_error)),
_ => None,
})
.collect();
assert_eq!(
pairs,
vec![("c_wait".into(), false), ("c_echo".into(), true)]
);
assert_eq!(
approvals(&events),
vec![("c_wait".into(), "waits".into(), "allow".into())]
);
}
#[tokio::test]
async fn cancelled_session_continues_on_the_next_run_with_a_fresh_token() {
let mut reg = Registry::new();
reg.register("waits", WaitsForCancel);
let (s, _e) = session_with(
vec![Ok(tool_turn("c1", "waits")), Ok(text_turn("second run"))],
reg,
config(),
);
let mut s = s;
let handle1 = s.cancel_handle();
let canceller = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(20)).await;
handle1.cancel();
});
let r1 = s.run_text("q1").await;
canceller.await.expect("canceller");
assert_eq!(r1.status, Status::Cancelled);
assert!(!s.cancel_handle().is_cancelled());
let r2 = s.run_text("q2").await;
assert_eq!(r2.status, Status::Completed);
assert_eq!(r2.final_message.as_deref(), Some("second run"));
assert!(s.history().len() >= 4, "history: {:?}", s.history().len());
}
struct CapturingProvider {
inner: MockProvider,
requests: Arc<Mutex<Vec<Vec<Message>>>>,
}
#[async_trait]
impl Provider for CapturingProvider {
#[allow(clippy::unnecessary_literal_bound)]
fn api_schema(&self) -> &str {
"mock"
}
async fn complete(
&self,
request: &ConversationRequest,
) -> Result<Completion, ProviderError> {
self.requests.lock().unwrap().push(request.messages.clone());
self.inner.complete(request).await
}
}
#[allow(clippy::type_complexity)]
fn capturing_session_with(
script: Vec<Result<Completion, ProviderError>>,
registry: Registry,
) -> (
Session,
Arc<Mutex<Vec<Vec<Message>>>>,
Arc<Mutex<Vec<Event>>>,
) {
let requests = Arc::new(Mutex::new(Vec::new()));
let events = Arc::new(Mutex::new(Vec::new()));
let sink_events = Arc::clone(&events);
let sink = Box::new(FnSink(move |event| {
sink_events.lock().unwrap().push(event);
}));
let provider = Arc::new(CapturingProvider {
inner: MockProvider::with_results(script),
requests: Arc::clone(&requests),
});
let session = Session::new(provider, registry, vec![], config(), sink);
(session, requests, events)
}
fn user_text(message: &Message) -> Option<&str> {
match (message.role, message.content.as_slice()) {
(Role::User, [ContentBlock::Text { text }]) => Some(text.as_str()),
_ => None,
}
}
#[tokio::test]
async fn second_run_continues_the_conversation() {
let (mut s, requests, _e) = capturing_session_with(
vec![
Ok(text_turn("first answer")),
Ok(text_turn("second answer")),
],
Registry::new(),
);
let r1 = s.run_text("q1").await;
let r2 = s.run_text("q2").await;
assert_eq!(r1.status, Status::Completed);
assert_eq!(r2.status, Status::Completed);
assert_eq!(r2.final_message.as_deref(), Some("second answer"));
let reqs = requests.lock().unwrap();
assert_eq!(reqs.len(), 2);
let run2 = &reqs[1];
assert_eq!(run2.len(), 3, "user q1, assistant, user q2: {run2:?}");
assert_eq!(user_text(&run2[0]), Some("q1"));
assert_eq!(run2[1].role, Role::Assistant);
assert_eq!(user_text(&run2[2]), Some("q2"));
let roles: Vec<Role> = s.history().iter().map(|m| m.role).collect();
assert_eq!(
roles,
vec![Role::User, Role::Assistant, Role::User, Role::Assistant]
);
}
fn capturing_with_cfg(
script: Vec<Result<Completion, ProviderError>>,
cfg: EngineConfig,
) -> (Session, Arc<Mutex<Vec<Vec<Message>>>>) {
let requests = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(CapturingProvider {
inner: MockProvider::with_results(script),
requests: Arc::clone(&requests),
});
let session = Session::new(provider, Registry::new(), vec![], cfg, Box::new(NullSink));
(session, requests)
}
fn instr_config(cwd: std::path::PathBuf) -> EngineConfig {
EngineConfig {
cwd,
instructions: locode_instructions::InstructionsConfig {
global_file: false,
..Default::default()
},
..config()
}
}
fn reminder_text(msgs: &[Message]) -> Option<String> {
msgs.iter()
.find_map(|m| match (m.role, m.content.as_slice()) {
(Role::User, [ContentBlock::Text { text }])
if text.starts_with("<system-reminder>") =>
{
Some(text.clone())
}
_ => None,
})
}
fn reminder_count(msgs: &[Message]) -> usize {
msgs.iter()
.filter(|m| {
matches!(
(m.role, m.content.as_slice()),
(Role::User, [ContentBlock::Text { text }]) if text.starts_with("<system-reminder>")
)
})
.count()
}
#[tokio::test]
async fn project_instructions_injected_once_before_prompt() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
std::fs::write(root.join("AGENTS.md"), "be terse").unwrap();
let (mut s, requests) = capturing_with_cfg(
vec![Ok(text_turn("ok1")), Ok(text_turn("ok2"))],
instr_config(root),
);
s.run_text("q1").await;
s.run_text("q2").await;
let reqs = requests.lock().unwrap();
let run1 = &reqs[0];
let rem = reminder_text(run1).expect("instructions injected on run 1");
assert!(rem.contains("## From:"), "labeled: {rem}");
assert!(rem.contains("be terse"), "content present: {rem}");
let rem_idx = run1
.iter()
.position(|m| reminder_text(std::slice::from_ref(m)).is_some());
let q1_idx = run1.iter().position(|m| user_text(m) == Some("q1"));
assert!(rem_idx < q1_idx, "reminder comes before the prompt");
assert_eq!(reminder_count(&reqs[1]), 1, "not re-injected on run 2");
}
fn resumed_with_cfg(
script: Vec<Result<Completion, ProviderError>>,
cfg: EngineConfig,
replayed: Vec<Message>,
) -> (Session, Arc<Mutex<Vec<Vec<Message>>>>) {
let requests = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(CapturingProvider {
inner: MockProvider::with_results(script),
requests: Arc::clone(&requests),
});
let session = Session::new(provider, Registry::new(), replayed, cfg, Box::new(NullSink));
(session, requests)
}
#[tokio::test]
async fn resuming_does_not_re_inject_unchanged_instructions() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
std::fs::write(root.join("AGENTS.md"), "be terse").unwrap();
let (mut first, requests) =
capturing_with_cfg(vec![Ok(text_turn("ok"))], instr_config(root.clone()));
first.run_text("q1").await;
let replayed = requests.lock().unwrap()[0].clone();
assert_eq!(reminder_count(&replayed), 1, "precondition: injected once");
let (mut resumed, resumed_requests) =
resumed_with_cfg(vec![Ok(text_turn("ok2"))], instr_config(root), replayed);
resumed.run_text("q2").await;
let reqs = resumed_requests.lock().unwrap();
assert_eq!(
reminder_count(&reqs[0]),
1,
"still the one from before the resume, not a second copy: {:#?}",
reqs[0]
);
}
#[tokio::test]
async fn resuming_re_injects_instructions_that_changed_while_away() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
std::fs::write(root.join("AGENTS.md"), "be terse").unwrap();
let (mut first, requests) =
capturing_with_cfg(vec![Ok(text_turn("ok"))], instr_config(root.clone()));
first.run_text("q1").await;
let replayed = requests.lock().unwrap()[0].clone();
std::fs::write(root.join("AGENTS.md"), "be verbose").unwrap();
let (mut resumed, resumed_requests) =
resumed_with_cfg(vec![Ok(text_turn("ok2"))], instr_config(root), replayed);
resumed.run_text("q2").await;
let reqs = resumed_requests.lock().unwrap();
assert_eq!(reminder_count(&reqs[0]), 2, "the new body joins the old");
let latest = reqs[0]
.iter()
.rev()
.find_map(|m| reminder_text(std::slice::from_ref(m)))
.expect("a reminder");
assert!(latest.contains("be verbose"), "{latest}");
assert!(
latest.contains("replace all previously provided"),
"banner present: {latest}"
);
}
#[tokio::test]
async fn instructions_dropped_from_the_transcript_are_re_injected() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
std::fs::write(root.join("AGENTS.md"), "be terse").unwrap();
let (mut first, requests) =
capturing_with_cfg(vec![Ok(text_turn("ok"))], instr_config(root.clone()));
first.run_text("q1").await;
let compacted: Vec<Message> = requests.lock().unwrap()[0]
.iter()
.filter(|m| reminder_text(std::slice::from_ref(m)).is_none())
.cloned()
.collect();
let (mut after, after_requests) =
resumed_with_cfg(vec![Ok(text_turn("ok2"))], instr_config(root), compacted);
after.run_text("q2").await;
assert_eq!(
reminder_count(&after_requests.lock().unwrap()[0]),
1,
"re-injected after being compacted away"
);
}
fn skills_config(cwd: std::path::PathBuf) -> EngineConfig {
EngineConfig {
cwd: cwd.clone(),
skills: locode_skills::SkillsConfig::enabled(),
..instr_config(cwd)
}
}
fn write_skill(root: &std::path::Path, name: &str, description: &str) {
let dir = root.join(".agents/skills").join(name);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
dir.join("SKILL.md"),
format!("---\nname: {name}\ndescription: {description}\n---\n# {name}\n"),
)
.unwrap();
}
#[tokio::test]
async fn skills_listing_injected_once_then_quiet() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
write_skill(&root, "commit", "Make a commit");
let (mut s, requests) = capturing_with_cfg(
vec![Ok(text_turn("a")), Ok(text_turn("b"))],
skills_config(root),
);
s.run_text("q1").await;
s.run_text("q2").await;
let reqs = requests.lock().unwrap();
let listing = reqs[0]
.iter()
.filter_map(|m| reminder_text(std::slice::from_ref(m)))
.find(|t| t.contains("skills are available"))
.expect("listing injected");
assert!(listing.contains(r#"<skill name="commit""#), "{listing}");
assert!(listing.contains("Make a commit"), "{listing}");
assert!(
listing.contains("SKILL.md\">"),
"the path attribute: {listing}"
);
let count = |msgs: &[Message]| {
msgs.iter()
.filter(|m| {
reminder_text(std::slice::from_ref(m))
.is_some_and(|t| t.contains("skills are available"))
})
.count()
};
assert_eq!(count(&reqs[1]), 1, "unchanged ⇒ not re-sent");
}
#[tokio::test]
async fn adding_a_skill_re_sends_the_entire_listing() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
write_skill(&root, "commit", "Make a commit");
let (mut s, requests) = capturing_with_cfg(
vec![Ok(text_turn("a")), Ok(text_turn("b")), Ok(text_turn("c"))],
skills_config(root.clone()),
);
s.run_text("q1").await;
write_skill(&root, "review", "Review a diff"); s.run_text("q2").await; s.run_text("q3").await;
let reqs = requests.lock().unwrap();
let listing = |msgs: &[Message]| {
msgs.iter()
.filter_map(|m| reminder_text(std::slice::from_ref(m)))
.rfind(|t| t.contains("skills are available"))
};
assert!(
!listing(&reqs[1]).unwrap().contains(r#"name="review""#),
"not yet — the scan that would see it runs at the end of this run"
);
let third = listing(&reqs[2]).expect("re-sent");
assert!(
third.contains(r#"name="commit""#),
"old skill included: {third}"
);
assert!(
third.contains(r#"name="review""#),
"new skill included: {third}"
);
}
#[tokio::test]
async fn removing_the_last_skill_announces_it() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
write_skill(&root, "commit", "Make a commit");
let (mut s, requests) = capturing_with_cfg(
vec![Ok(text_turn("a")), Ok(text_turn("b")), Ok(text_turn("c"))],
skills_config(root.clone()),
);
s.run_text("q1").await;
std::fs::remove_dir_all(root.join(".agents/skills/commit")).unwrap();
s.run_text("q2").await; s.run_text("q3").await;
let reqs = requests.lock().unwrap();
let last = reqs[2]
.iter()
.filter_map(|m| reminder_text(std::slice::from_ref(m)))
.next_back()
.expect("a reminder");
assert!(last.contains("No skills are currently available"), "{last}");
}
#[tokio::test]
async fn no_skills_ever_means_no_message_at_all() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
let (mut s, requests) = capturing_with_cfg(vec![Ok(text_turn("a"))], skills_config(root));
s.run_text("q1").await;
let reqs = requests.lock().unwrap();
assert!(
!reqs[0]
.iter()
.any(|m| reminder_text(std::slice::from_ref(m))
.is_some_and(|t| t.contains("skills"))),
"silence, not a denial"
);
}
#[tokio::test]
async fn project_instructions_absent_when_disabled() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
std::fs::write(root.join("AGENTS.md"), "be terse").unwrap();
let mut cfg = instr_config(root);
cfg.instructions.enabled = false;
let (mut s, requests) = capturing_with_cfg(vec![Ok(text_turn("ok"))], cfg);
s.run_text("q").await;
assert!(reminder_text(&requests.lock().unwrap()[0]).is_none());
}
#[tokio::test]
async fn project_instructions_absent_when_no_agents_md() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
let (mut s, requests) = capturing_with_cfg(vec![Ok(text_turn("ok"))], instr_config(root));
s.run_text("q").await;
assert!(reminder_text(&requests.lock().unwrap()[0]).is_none());
}
#[tokio::test]
async fn project_instructions_replace_banner_on_edit() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
let agents = root.join("AGENTS.md");
std::fs::write(&agents, "v1 rules").unwrap();
let (mut s, requests) = capturing_with_cfg(
vec![Ok(text_turn("ok1")), Ok(text_turn("ok2"))],
instr_config(root),
);
s.run_text("q1").await;
std::fs::write(&agents, "v2 rules").unwrap(); s.run_text("q2").await;
let reqs = requests.lock().unwrap();
let run2 = &reqs[1];
let banner = run2
.iter()
.find_map(|m| match (m.role, m.content.as_slice()) {
(Role::User, [ContentBlock::Text { text }])
if text.contains("replace all previously provided") =>
{
Some(text.clone())
}
_ => None,
})
.expect("replace banner on edit");
assert!(banner.contains("v2 rules"), "new content: {banner}");
assert!(!banner.contains("v1 rules"), "not the old content");
assert_eq!(reminder_count(run2), 2);
}
#[tokio::test]
async fn project_instructions_removal_banner_on_delete() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
let agents = root.join("AGENTS.md");
std::fs::write(&agents, "rules").unwrap();
let (mut s, requests) = capturing_with_cfg(
vec![Ok(text_turn("ok1")), Ok(text_turn("ok2"))],
instr_config(root),
);
s.run_text("q1").await;
std::fs::remove_file(&agents).unwrap(); s.run_text("q2").await;
let reqs = requests.lock().unwrap();
assert!(
reqs[1].iter().any(|m| matches!(
(m.role, m.content.as_slice()),
(Role::User, [ContentBlock::Text { text }]) if text.contains("no longer apply")
)),
"removal notice on delete"
);
}
#[tokio::test]
async fn project_instructions_not_reinjected_when_unchanged() {
let dir = tempfile::tempdir().unwrap();
let root = std::fs::canonicalize(dir.path()).unwrap();
std::fs::create_dir(root.join(".git")).unwrap();
std::fs::write(root.join("AGENTS.md"), "stable").unwrap();
let (mut s, requests) = capturing_with_cfg(
vec![
Ok(text_turn("ok1")),
Ok(text_turn("ok2")),
Ok(text_turn("ok3")),
],
instr_config(root),
);
s.run_text("q1").await;
s.run_text("q2").await;
s.run_text("q3").await;
assert_eq!(reminder_count(&requests.lock().unwrap()[2]), 1);
}
#[tokio::test]
async fn init_emitted_once_across_runs_with_one_result_each() {
let (mut s, events) = session_with(
vec![Ok(text_turn("one")), Ok(text_turn("two"))],
Registry::new(),
config(),
);
let _ = s.run_text("q1").await;
let _ = s.run_text("q2").await;
let evs = dump(&events);
let inits = evs
.iter()
.filter(|e| matches!(e, Event::Init { .. }))
.count();
let results = evs
.iter()
.filter(|e| matches!(e, Event::Result { .. }))
.count();
assert_eq!(inits, 1, "Init is once per session, not per run");
assert_eq!(results, 2, "one Result per run");
assert!(
matches!(evs.first(), Some(Event::Init { .. })),
"Init still opens the stream"
);
}
#[tokio::test]
async fn report_counts_are_per_run_not_cumulative() {
let mut t1 = tool_turn("c1", "echo");
t1.usage = Usage {
input_tokens: 10,
output_tokens: 5,
..Usage::default()
};
let t2 = text_turn("done one");
let mut t3 = text_turn("done two");
t3.usage = Usage {
input_tokens: 20,
output_tokens: 7,
..Usage::default()
};
let (mut s, _e) = session_with(vec![Ok(t1), Ok(t2), Ok(t3)], echo_registry(), config());
let r1 = s.run_text("q1").await;
let r2 = s.run_text("q2").await;
assert_eq!(r1.turns, 2);
assert_eq!(r1.tool_calls.len(), 1);
assert_eq!(r2.turns, 1, "run 2 counts its own turns only");
assert!(r2.tool_calls.is_empty());
assert_eq!(r2.usage.input_tokens, 20, "usage is per-run");
assert_eq!(r2.usage.output_tokens, 7);
}
#[tokio::test]
async fn two_run_stream_reconstructs_the_full_conversation() {
let (mut s, events) = session_with(
vec![
Ok(tool_turn("c1", "echo")),
Ok(text_turn("done one")),
Ok(text_turn("done two")),
],
echo_registry(),
config(),
);
let _ = s.run_text("q1").await;
let _ = s.run_text("q2").await;
let rebuilt: Conversation = reconstruct_conversation(&dump(&events));
let roles: Vec<Role> = rebuilt.messages.iter().map(|m| m.role).collect();
assert_eq!(
roles,
vec![
Role::User,
Role::Assistant,
Role::User,
Role::Assistant,
Role::User,
Role::Assistant,
]
);
assert_eq!(rebuilt.messages.as_slice(), s.history());
}
#[tokio::test]
async fn continues_after_model_error() {
let (mut s, requests, _e) = capturing_session_with(
vec![Err(ProviderError::ContextOverflow), Ok(text_turn("ok now"))],
Registry::new(),
);
let r1 = s.run_text("q1").await;
let r2 = s.run_text("q2").await;
assert_eq!(r1.status, Status::ModelError);
assert_eq!(r2.status, Status::Completed);
let reqs = requests.lock().unwrap();
let run2 = &reqs[1];
assert_eq!(run2.len(), 2, "user q1 + user q2: {run2:?}");
assert_eq!(user_text(&run2[0]), Some("q1"));
assert_eq!(user_text(&run2[1]), Some("q2"));
}
#[tokio::test]
async fn continues_after_fatal_tool_error_with_valid_pairing() {
let mut reg = Registry::new();
reg.register("boom", Boom);
let (mut s, requests, _e) = capturing_session_with(
vec![Ok(tool_turn("c1", "boom")), Ok(text_turn("recovered"))],
reg,
);
let r1 = s.run_text("q1").await;
let r2 = s.run_text("q2").await;
assert_eq!(r1.status, Status::Error);
assert_eq!(r2.status, Status::Completed);
let reqs = requests.lock().unwrap();
let run2 = &reqs[1];
assert_eq!(run2.len(), 4, "q1, assistant, tool_result, q2: {run2:?}");
assert!(
run2[1]
.content
.iter()
.any(|b| matches!(b, ContentBlock::ToolUse { id, .. } if id == "c1"))
);
assert!(run2[2].content.iter().any(|b| matches!(
b,
ContentBlock::ToolResult { tool_use_id, is_error: true, .. } if tool_use_id == "c1"
)));
assert_eq!(user_text(&run2[3]), Some("q2"));
}
#[tokio::test]
async fn usage_is_summed_across_turns() {
let mut first = tool_turn("c1", "echo");
first.usage = Usage {
input_tokens: 10,
output_tokens: 5,
..Usage::default()
};
let mut second = text_turn("done");
second.usage = Usage {
input_tokens: 20,
output_tokens: 7,
..Usage::default()
};
let (mut s, _e) = session_with(vec![Ok(first), Ok(second)], echo_registry(), config());
let report = s.run_text("go").await;
assert_eq!(report.usage.input_tokens, 30);
assert_eq!(report.usage.output_tokens, 12);
}
#[tokio::test]
async fn setting_the_model_announces_it_without_touching_the_preamble() {
let preamble = vec![Message {
role: Role::System,
content: vec![ContentBlock::Text {
text: "You are powered by the model old-1.".into(),
}],
}];
let provider = Arc::new(MockProvider::with_results(vec![Ok(text_turn("ok"))]));
let mut s = Session::new(
provider.clone(),
Registry::new(),
preamble.clone(),
config(),
Box::new(NullSink),
);
let notice = s.set_model(provider, "new-2");
s.announce(notice);
assert_eq!(
s.history()[0],
preamble[0],
"the preamble is untouched — the trace already recorded it"
);
let last = s.history().last().expect("announcement appended");
assert_eq!(last.role, Role::User);
let ContentBlock::Text { text } = &last.content[0] else {
panic!("text block")
};
assert!(text.starts_with("<system-reminder>"), "{text}");
assert!(text.contains("is now new-2"), "{text}");
assert!(
text.contains("out of date"),
"corrects the stale line: {text}"
);
}
#[tokio::test]
async fn context_usage_is_the_final_turn_not_the_sum() {
let mut first = tool_turn("c1", "echo");
first.usage = Usage {
input_tokens: 10,
output_tokens: 5,
..Usage::default()
};
let mut second = text_turn("done");
second.usage = Usage {
input_tokens: 20,
output_tokens: 7,
cache_read_tokens: Some(4),
cache_creation_tokens: Some(3),
..Usage::default()
};
let (mut s, _e) = session_with(vec![Ok(first), Ok(second)], echo_registry(), config());
let report = s.run_text("go").await;
assert_eq!(report.context_usage.input_tokens, 20, "the last turn only");
assert_eq!(report.context_usage.output_tokens, 7);
assert_eq!(
report.context_usage.context_tokens(),
20 + 4 + 3 + 7,
"both cache counters are prompt tokens"
);
assert_eq!(report.usage.input_tokens, 30, "the sum is still the sum");
}
}