use super::*;
#[test]
fn agent_prompt_cache_key_session_requests_reuse_bounded_hash() {
let temp = tempfile::TempDir::new().unwrap();
std::fs::write(temp.path().join("file.txt"), "file text").unwrap();
let manager = crate::sessions::SessionManager::new(temp.path().join("sessions"));
let session = manager.create().unwrap();
let provider = ScriptedProvider::new(vec![read_done("cache_read"), text_done("done")]);
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
agent
.run_print_with_tools_streaming_output(
&provider,
"read file",
Some(&tools),
Some(&session),
temp.path(),
None,
)
.unwrap();
let requests = provider.requests();
let expected = expected_prompt_cache_key_for_session_id(session.id());
assert_eq!(requests.len(), 2);
assert_eq!(expected.len(), "magi-code-session-".len() + 32);
assert!(expected.starts_with("magi-code-session-"));
assert!(!expected.contains(session.id()));
assert_eq!(requests[0].prompt_cache_key(), Some(expected.as_str()));
assert_eq!(requests[1].prompt_cache_key(), Some(expected.as_str()));
}
#[test]
fn preflight_projection_preserves_first_request_and_budget() {
for (provider_id, model, reasoning_updates) in [
("local-ai", "model", false),
("openai-codex", "gpt-4o", false),
("openai-codex", "gpt-4o", true),
] {
let temp = tempfile::TempDir::new().unwrap();
let session = crate::sessions::SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
append_preflight_history(&session, temp.path(), "earlier request");
let tools = ToolRuntime::new(temp.path()).unwrap();
tools.set_tool_enabled("read", false).unwrap();
tools.set_tool_enabled("bash", false).unwrap();
let agent = AgentSession::from_system_prompt(
"system instructions".to_string(),
AgentSessionConfig::for_model(model),
)
.with_provider_id(provider_id)
.with_reasoning_updates(reasoning_updates);
let preflight = agent
.project_prompt_for_preflight("next", &session, Some(&tools))
.unwrap();
let projection = preflight.projection;
let identity = expected_prompt_cache_key_for_session_id(session.id());
let expected = preflight
.request
.to_owned_request()
.with_prompt_cache_key(&identity)
.with_conversation_id(&identity)
.with_semantic_progress_timeout(Duration::from_secs(5));
let agent = agent.with_context_budget(ContextBudget {
max_tokens: projection.tokens,
reserve_tokens: 0,
..ContextBudget::default()
});
let provider = ScriptedProvider::new(vec![text_done("answer")]);
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output_inner(
&provider,
AgentRunRequest {
session: Some(&session),
tools: Some(&tools),
output_sink: Some(&mut sink),
semantic_progress_timeout: Some(Duration::from_secs(5)),
..run_request("next", temp.path())
},
None,
false,
Arc::new(AtomicU64::new(0)),
Some(preflight),
)
.unwrap();
let requests = provider.requests();
assert_eq!(requests.len(), 1);
assert_eq!(requests[0].to_owned_request(), expected);
assert!(
requests[0]
.messages()
.iter()
.any(|m| m.content == "earlier request")
);
assert!(sink.outputs.iter().any(|event| matches!(
event,
OutputEvent::ContextUsage { current_tokens, source, request_sequence: 0, .. }
if *current_tokens == projection.tokens && *source == projection.source
)));
assert_eq!(
session
.read_events()
.unwrap()
.iter()
.filter(|event| {
event.event_type == "user_input" && event.payload["text"] == "next"
})
.count(),
1
);
}
}
#[test]
fn preflight_projection_rejects_changed_context_before_provider() {
for change in [
"append",
"external rewrite",
"rotation",
"prompt",
"system",
"model",
"provider",
"tools",
] {
let temp = tempfile::TempDir::new().unwrap();
let session = crate::sessions::SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
let old_history = "aaaaaaaaaaaaaaaa ".repeat(4_096);
let larger_context = "a ".repeat(old_history.len() / 2);
append_preflight_history(&session, temp.path(), &old_history);
let tools = ToolRuntime::new(temp.path()).unwrap();
tools.set_tool_enabled("read", false).unwrap();
let mut agent = AgentSession::from_system_prompt(
"system".to_string(),
AgentSessionConfig::for_model("gpt-4o"),
)
.with_provider_id("openai-codex");
let preflight = agent
.project_prompt_for_preflight("next", &session, Some(&tools))
.unwrap();
let threshold = preflight.projection.tokens;
let mut prompt = "next";
match change {
"append" => append_preflight_history(&session, temp.path(), &larger_context),
"external rewrite" => {
let contents = std::fs::read_to_string(session.path()).unwrap();
let edited = contents.replace(&old_history, &larger_context);
assert_eq!(edited.len(), contents.len());
std::fs::write(session.path(), edited).unwrap();
}
"rotation" => crate::sessions::record_session_compaction(
&session,
temp.path(),
&larger_context,
"openai-codex",
"gpt-4o",
session.read_events().unwrap().len(),
)
.unwrap(),
"prompt" => prompt = &larger_context,
"system" => agent = agent.with_appended_system_prompt(&larger_context),
"model" => agent = agent.with_model("unknown-model"),
"provider" => agent = agent.with_provider_id("local-ai"),
"tools" => tools.set_tool_enabled("read", true).unwrap(),
_ => unreachable!(),
}
let current_tokens = agent
.clone()
.project_prompt_input_tokens(prompt, &session, Some(&tools))
.unwrap();
assert!(
current_tokens > threshold,
"{change}: {current_tokens} <= {threshold}"
);
let agent = agent.with_context_budget(ContextBudget {
max_tokens: threshold,
reserve_tokens: 0,
..ContextBudget::default()
});
let before_execution = std::fs::read(session.path()).unwrap();
let provider = ScriptedProvider::new(vec![text_done("must not run")]);
let error = agent
.run_print_with_tools_streaming_output_inner(
&provider,
AgentRunRequest {
session: Some(&session),
tools: Some(&tools),
..run_request(prompt, temp.path())
},
None,
false,
Arc::new(AtomicU64::new(0)),
Some(preflight),
)
.unwrap_err();
let budget = error.downcast_ref::<ContextBudgetError>().unwrap();
assert_eq!(budget.phase(), ContextBudgetPhase::Initial, "{change}");
assert_eq!(budget.estimated_tokens(), current_tokens, "{change}");
assert!(provider.requests().is_empty(), "{change}");
assert_eq!(
std::fs::read(session.path()).unwrap(),
before_execution,
"{change}"
);
}
}
#[test]
fn preflight_projection_keeps_prompt_snapshot_and_fresh_replay_diagnostics() {
let temp = tempfile::TempDir::new().unwrap();
let session = crate::sessions::SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
append_preflight_history(&session, temp.path(), "earlier request");
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
let effective_prompt = AgentSession::effective_prompt_for_projection(
"inspect #git-status",
Some(&tools),
InvocationMode::Print,
);
assert!(effective_prompt.contains("success=\"false\""));
let preflight = agent
.project_prompt_for_preflight(&effective_prompt, &session, Some(&tools))
.unwrap();
assert!(
std::process::Command::new("git")
.args(["init", "--quiet"])
.current_dir(temp.path())
.status()
.unwrap()
.success()
);
let history = std::fs::read_to_string(session.path()).unwrap();
std::fs::write(session.path(), format!("{history}{{not-json}}\n")).unwrap();
let provider = ScriptedProvider::new(vec![text_done("answer")]);
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output_inner(
&provider,
AgentRunRequest {
effective_prompt: Some(&effective_prompt),
session: Some(&session),
tools: Some(&tools),
output_sink: Some(&mut sink),
..run_request("inspect #git-status", temp.path())
},
None,
false,
Arc::new(AtomicU64::new(0)),
Some(preflight),
)
.unwrap();
let requests = provider.requests();
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].messages().last().unwrap().content,
effective_prompt
);
assert_eq!(effective_prompt.matches("<context_injections>").count(), 1);
assert!(sink.outputs.iter().any(|event| matches!(
event, OutputEvent::Diagnostic { level, message }
if level == "warning" && message.contains("session replay")
)));
assert_eq!(
session
.read_events_tolerant()
.unwrap()
.events
.iter()
.filter(|event| {
event.event_type == "user_input" && event.payload["text"] == effective_prompt
})
.count(),
1
);
}
fn append_preflight_history(session: &Session, cwd: &Path, text: &str) {
for (kind, payload) in [
(SessionEventKind::UserInput, json!({"text": text})),
(
SessionEventKind::AssistantOutput,
json!({"text": "earlier answer"}),
),
(SessionEventKind::TurnStatus, json!({"status": "completed"})),
] {
session
.append(&SessionEvent::new_kind(
kind,
session.id().to_string(),
cwd.to_path_buf(),
payload,
))
.unwrap();
}
}
#[test]
fn parent_projection_replay_is_reused_by_execution() {
projection_replay_is_reused_by_execution(crate::output::InvocationMode::Print);
}
#[test]
fn child_projection_replay_is_reused_by_execution() {
projection_replay_is_reused_by_execution(crate::output::InvocationMode::Subagent);
}
fn projection_replay_is_reused_by_execution(invocation_mode: crate::output::InvocationMode) {
let temp = tempfile::TempDir::new().unwrap();
let session = crate::sessions::SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
for (kind, payload) in [
(SessionEventKind::UserInput, json!({"text": "earlier"})),
(SessionEventKind::AssistantOutput, json!({"text": "answer"})),
(SessionEventKind::TurnStatus, json!({"status": "completed"})),
] {
session
.append(&SessionEvent::new_kind(
kind,
session.id().to_string(),
temp.path().to_path_buf(),
payload,
))
.unwrap();
}
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
agent
.project_prompt_input_tokens("next", &session, None)
.unwrap();
let provider = ScriptedProvider::new(vec![vec![done()]]);
agent
.run_print_with_tools_streaming_output_cancellable(
&provider,
AgentRunRequest {
session: Some(&session),
invocation_mode,
..run_request("next", temp.path())
},
)
.unwrap();
let metrics = agent.replay_cache_metrics();
assert_eq!(metrics.full_scans, 1);
assert_eq!(metrics.replay_builds, 1);
assert_eq!(metrics.cache_hits, 1);
}
#[test]
fn subagent_prompt_cache_key_uses_identity_material_not_session_id() {
let temp = tempfile::TempDir::new().unwrap();
let manager = crate::sessions::SessionManager::new(temp.path().join("sessions"));
let session_a = manager.create().unwrap();
let session_b = manager.create().unwrap();
let provider = ScriptedProvider::new(vec![vec![done()], vec![done()], vec![done()]]);
let agent = AgentSession::from_system_prompt(
"system profile prompt".to_string(),
AgentSessionConfig::for_model("model").with_context_budget(ContextBudget {
enabled: false,
..ContextBudget::default()
}),
);
for (session, agent_id) in [
(&session_a, "identity-a"),
(&session_b, "identity-a"),
(&session_b, "identity-b"),
] {
agent
.run_print_with_tools_streaming_output_cancellable(
&provider,
AgentRunRequest {
session: Some(session),
invocation_mode: crate::output::InvocationMode::Subagent,
agent_id: Some(agent_id.to_string()),
..run_request("child prompt", temp.path())
},
)
.unwrap();
}
let requests = provider.requests();
let expected_a =
expected_prompt_cache_key_for_subagent(&agent.provider_id, &agent.model, "identity-a");
let expected_b =
expected_prompt_cache_key_for_subagent(&agent.provider_id, &agent.model, "identity-b");
assert_eq!(expected_a.len(), "magi-code-session-".len() + 32);
assert!(expected_a.starts_with("magi-code-session-"));
assert_eq!(requests[0].prompt_cache_key(), Some(expected_a.as_str()));
assert_eq!(requests[1].prompt_cache_key(), Some(expected_a.as_str()));
assert_eq!(requests[2].prompt_cache_key(), Some(expected_b.as_str()));
assert_eq!(
requests[0].conversation_id(),
Some(expected_prompt_cache_key_for_session_id(session_a.id()).as_str())
);
assert_eq!(
requests[1].conversation_id(),
Some(expected_prompt_cache_key_for_session_id(session_b.id()).as_str())
);
assert_ne!(requests[0].conversation_id(), requests[1].conversation_id());
assert_ne!(expected_a, expected_b);
assert_ne!(
expected_a,
expected_prompt_cache_key_for_session_id(session_a.id())
);
}
#[test]
fn agent_prompt_cache_key_sessionless_requests_omit_key() {
let temp = tempfile::TempDir::new().unwrap();
let provider = ScriptedProvider::new(vec![text_done("done")]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
agent
.run_print_with_tools_streaming_output(&provider, "prompt", None, None, temp.path(), None)
.unwrap();
let requests = provider.requests();
assert_eq!(requests.len(), 1);
assert_eq!(requests[0].prompt_cache_key(), None);
}
#[test]
fn top_level_prompt_context_injections_reach_provider_and_session() {
let temp = tempfile::TempDir::new().unwrap();
let manager = crate::sessions::SessionManager::new(temp.path().join("sessions"));
let session = manager.create().unwrap();
let provider = ScriptedProvider::new(vec![vec![done()]]);
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::from_system_prompt(
String::new(),
AgentSessionConfig::for_model("model").with_context_budget(ContextBudget {
enabled: false,
..ContextBudget::default()
}),
);
agent
.run_print_with_tools_streaming_output_cancellable(
&provider,
AgentRunRequest {
tools: Some(&tools),
session: Some(&session),
..run_request("inspect #git-status", temp.path())
},
)
.unwrap();
let requests = provider.requests();
let messages = requests[0].messages();
let prompt = &messages.last().unwrap().content;
assert!(prompt.starts_with("inspect #git-status\n\n<context_injections>"));
assert!(prompt.contains("name=\"git-status\""));
assert!(prompt.contains("command=\"git status\""));
assert!(prompt.contains("success=\"false\""));
let events = session.read_events().unwrap();
assert_eq!(events[0].event_type, "user_input");
assert_eq!(events[0].payload["text"], *prompt);
}
#[test]
fn diff_tag_without_tools_reaches_provider_and_persistence_with_explicit_error() {
let temp = tempfile::TempDir::new().unwrap();
let session = crate::sessions::SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
let provider = ScriptedProvider::new(vec![vec![done()]]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output_cancellable(
&provider,
AgentRunRequest {
session: Some(&session),
output_sink: Some(&mut sink),
..run_request("inspect #diff-changes", temp.path())
},
)
.unwrap();
let requests = provider.requests();
let messages = requests[0].messages();
let prompt = &messages.last().unwrap().content;
assert!(prompt.contains("Diff comments unavailable: Not a Git worktree"));
assert!(!prompt.contains("tools are disabled"));
assert_eq!(session.read_events().unwrap()[0].payload["text"], *prompt);
assert_eq!(
sink.outputs[0],
OutputEvent::UserPrompt {
text: prompt.clone()
}
);
}
#[test]
fn saved_diff_comments_match_live_transcript_provider_and_session() {
use crate::diff_review::{CommentChange, ReviewSide};
use crate::output::UserPromptOrigin;
let temp = tempfile::TempDir::new().unwrap();
let base = std::fs::canonicalize(temp.path()).unwrap();
let env = crate::test_support::env::env_lock();
let _home = env.save("MC_HOME");
env.set_var("MC_HOME", base.join("home"));
let root = base.join("repo");
std::fs::create_dir(&root).unwrap();
let result = std::process::Command::new("git")
.args(["init", "-q"])
.current_dir(&root)
.output()
.unwrap();
assert!(result.status.success(), "{:?}", result);
std::fs::write(root.join("guard.rs"), "guard();\n").unwrap();
let snapshot = crate::diff_review::load_snapshot(
&root,
&AgentCancellation::default(),
&mut crate::diff_review::SnapshotCache::default(),
)
.unwrap();
crate::diff_review::save_comment(
&root,
CommentChange::Save {
id: None,
path: "guard.rs".into(),
side: ReviewSide::Changed,
line: 1,
text: "Check the guard\nbefore continuing".into(),
anchor: crate::diff_review::ReviewAnchor::capture(
snapshot
.files
.iter()
.find(|file| file.path == "guard.rs")
.unwrap(),
ReviewSide::Changed,
1,
),
},
&AgentCancellation::default(),
)
.unwrap();
for (origin, mid_turn) in [
(UserPromptOrigin::User, false),
(UserPromptOrigin::Steering, false),
(UserPromptOrigin::Steering, true),
] {
let session = crate::sessions::SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
let provider = ScriptedProvider::new(vec![text_done("first"), text_done("second")]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
let mut sink = CapturingOutputSink::default();
let steering = AgentSteering::new();
if mid_turn {
steering
.try_enqueue("inspect #diff-changes".into())
.unwrap();
}
agent
.run_print_with_tools_streaming_output_cancellable_with_steering(
&provider,
AgentRunRequest {
session: Some(&session),
output_sink: Some(&mut sink),
prompt_origin: if mid_turn {
UserPromptOrigin::User
} else {
origin
},
..run_request(
if mid_turn {
"start"
} else {
"inspect #diff-changes"
},
&root,
)
},
steering,
)
.unwrap();
let events = session.read_events().unwrap();
let persisted = events
.iter()
.filter(|event| event.event_type == "user_input")
.nth(usize::from(mid_turn))
.unwrap();
let expanded = persisted.payload["text"].as_str().unwrap();
assert!(expanded.contains("guard.rs Changed:1"), "{expanded}");
assert!(
expanded.contains("Check the guard\nbefore continuing"),
"{expanded}"
);
assert_eq!(expanded.matches("Check the guard").count(), 1);
let expected_event = match origin {
UserPromptOrigin::User => OutputEvent::UserPrompt {
text: expanded.into(),
},
_ => OutputEvent::SteeringInjected {
text: expanded.into(),
},
};
assert!(sink.outputs.contains(&expected_event));
assert!(
provider
.requests()
.last()
.unwrap()
.messages()
.iter()
.any(|message| message.content == expanded)
);
}
}
#[test]
fn subagent_invocation_does_not_expand_hash_context_tags() {
let temp = tempfile::TempDir::new().unwrap();
let provider = ScriptedProvider::new(vec![vec![done()]]);
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::from_system_prompt(
String::new(),
AgentSessionConfig::for_model("model").with_context_budget(ContextBudget {
enabled: false,
..ContextBudget::default()
}),
);
agent
.run_print_with_tools_streaming_output_cancellable(
&provider,
AgentRunRequest {
tools: Some(&tools),
invocation_mode: crate::output::InvocationMode::Subagent,
..run_request("child sees #git-status #diff-changes", temp.path())
},
)
.unwrap();
let requests = provider.requests();
assert_eq!(
requests[0].messages().last().unwrap().content,
"child sees #git-status #diff-changes"
);
}
#[test]
fn context_budget_checks_effective_injected_prompt() {
let temp = tempfile::TempDir::new().unwrap();
let provider = ScriptedProvider::new(vec![vec![done()]]);
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::from_system_prompt(
String::new(),
AgentSessionConfig::for_model("model").with_context_budget(ContextBudget {
enabled: true,
max_tokens: 10,
reserve_tokens: 0,
keep_recent_tokens: 0,
..ContextBudget::default()
}),
);
let error = agent
.run_print_with_tools_streaming_output_cancellable(
&provider,
AgentRunRequest {
tools: Some(&tools),
invocation_mode: crate::output::InvocationMode::MissionControl,
..run_request("#git-status", temp.path())
},
)
.unwrap_err()
.to_string();
assert!(error.contains("exceeding threshold 10 tokens"), "{error}");
assert!(provider.requests().is_empty());
}
#[test]
fn cancellation_during_provider_callback_stops_active_stream_promptly() {
let cancel = Arc::new(AtomicBool::new(false));
let provider = ActiveCancelingProvider::new(Arc::clone(&cancel), 100);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
let error = agent
.run_print_with_tools_streaming_output_cancellable(
&provider,
AgentRunRequest {
cancellation: AgentCancellation::new(cancel),
..run_request("stream forever", Path::new("."))
},
)
.unwrap_err();
assert!(is_run_canceled(&error));
assert!(
provider.sent_events() < 100,
"provider should stop when cancellation is observed instead of draining the script"
);
}
#[test]
fn context_usage_estimate_emits_before_first_provider_response_with_budget_max() {
let provider = ScriptedProvider::new(vec![text_done("hello")]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default()).with_context_budget(
crate::context::ContextBudget {
max_tokens: 16_000,
reserve_tokens: 0,
..crate::context::ContextBudget::default()
},
);
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"say hi",
None,
None,
Path::new("."),
Some(&mut sink),
)
.unwrap();
assert!(matches!(sink.outputs[0], OutputEvent::UserPrompt { .. }));
assert!(matches!(
sink.outputs[1],
OutputEvent::ContextUsage {
current_tokens,
max_tokens: 16_000,
..
} if current_tokens > 0
));
assert!(matches!(
sink.outputs[2],
OutputEvent::AssistantDelta { .. }
));
}
#[test]
fn streaming_context_usage_projection_updates_during_assistant_delta() {
let long_delta = "tokenized streaming output ".repeat(80);
let provider = ScriptedProvider::new(vec![vec![text(long_delta.clone()), done()]]);
let agent = AgentSession::new("gpt-5.5", &[], &SkillDiscovery::default()).with_context_budget(
crate::context::ContextBudget {
max_tokens: 16_000,
reserve_tokens: 0,
..crate::context::ContextBudget::default()
},
);
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"say hi",
None,
None,
Path::new("."),
Some(&mut sink),
)
.unwrap();
let delta_position = sink
.outputs
.iter()
.position(
|event| matches!(event, OutputEvent::AssistantDelta { text } if text == &long_delta),
)
.unwrap();
let projection_position = sink
.outputs
.iter()
.position(|event| {
matches!(
event,
OutputEvent::ContextUsage {
source: crate::output::ContextUsageSource::TokenizerProjection,
request_sequence: 0,
..
}
)
})
.expect("streaming projection event");
assert!(delta_position < projection_position);
}
#[test]
fn streaming_context_usage_projection_emits_after_256_byte_delta() {
let first_delta = "x".repeat(256);
let second_delta = "second delta".to_string();
let provider = ScriptedProvider::new(vec![vec![
text(first_delta.clone()),
text(second_delta.clone()),
done(),
]]);
let agent = AgentSession::new("gpt-5.5", &[], &SkillDiscovery::default()).with_context_budget(
crate::context::ContextBudget {
max_tokens: 16_000,
reserve_tokens: 0,
..crate::context::ContextBudget::default()
},
);
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"say hi",
None,
None,
Path::new("."),
Some(&mut sink),
)
.unwrap();
let first_position = sink
.outputs
.iter()
.position(
|event| matches!(event, OutputEvent::AssistantDelta { text } if text == &first_delta),
)
.expect("first assistant delta");
let second_position = sink
.outputs
.iter()
.position(
|event| matches!(event, OutputEvent::AssistantDelta { text } if text == &second_delta),
)
.expect("second assistant delta");
let projection_position = sink
.outputs
.iter()
.enumerate()
.find_map(|(index, event)| {
(index > first_position
&& index < second_position
&& matches!(
event,
OutputEvent::ContextUsage {
source: crate::output::ContextUsageSource::TokenizerProjection,
request_sequence: 0,
..
}
))
.then_some(index)
})
.expect("projection between assistant deltas");
assert!(first_position < projection_position);
assert!(projection_position < second_position);
let initial_tokens = sink
.outputs
.iter()
.find_map(|event| match event {
OutputEvent::ContextUsage {
current_tokens,
source: crate::output::ContextUsageSource::TokenizerEstimate,
request_sequence: 0,
..
} => Some(*current_tokens),
_ => None,
})
.expect("initial tokenizer estimate");
let projected_tokens = match &sink.outputs[projection_position] {
OutputEvent::ContextUsage { current_tokens, .. } => *current_tokens,
_ => unreachable!("projection position is context usage"),
};
assert!(
projected_tokens.saturating_sub(initial_tokens) >= 16,
"initial={initial_tokens}, projected={projected_tokens}"
);
}
#[test]
fn streaming_context_usage_projection_counts_pending_delta_batch() {
let provider = ScriptedProvider::new(vec![
(0..6)
.map(|_| text("x".repeat(100)))
.chain(std::iter::once(done()))
.collect(),
]);
let agent =
AgentSession::new("model", &[], &SkillDiscovery::default()).with_provider_id("local-ai");
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"say hi",
None,
None,
Path::new("."),
Some(&mut sink),
)
.unwrap();
let projected_tokens = sink
.outputs
.iter()
.filter_map(|event| match event {
OutputEvent::ContextUsage {
current_tokens,
source: crate::output::ContextUsageSource::FallbackProjection,
..
} => Some(*current_tokens),
_ => None,
})
.max()
.expect("fallback projection event");
assert!(
projected_tokens >= 150,
"projected_tokens={projected_tokens}"
);
}
#[test]
fn streaming_context_usage_projection_is_coalesced_for_small_deltas() {
let provider = ScriptedProvider::new(vec![
(0..10)
.map(|_| text("a"))
.chain(std::iter::once(done()))
.collect(),
]);
let agent = AgentSession::new("gpt-5.5", &[], &SkillDiscovery::default());
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"say hi",
None,
None,
Path::new("."),
Some(&mut sink),
)
.unwrap();
let projections = sink
.outputs
.iter()
.filter(|event| {
matches!(
event,
OutputEvent::ContextUsage {
source: crate::output::ContextUsageSource::TokenizerProjection,
..
}
)
})
.count();
assert_eq!(projections, 1);
}
#[test]
fn context_usage_emits_again_before_tool_continuation_request() {
let temp = tempfile::TempDir::new().unwrap();
std::fs::write(temp.path().join("file.txt"), "hello").unwrap();
let provider = ScriptedProvider::new(vec![read_done("call_1"), text_done("done")]);
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::new("model", &[], &SkillDiscovery::default()).with_context_budget(
crate::context::ContextBudget {
max_tokens: 16_000,
reserve_tokens: 0,
..crate::context::ContextBudget::default()
},
);
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"read file",
Some(&tools),
None,
temp.path(),
Some(&mut sink),
)
.unwrap();
let usage_positions = sink
.outputs
.iter()
.enumerate()
.filter_map(|(index, event)| match event {
OutputEvent::ContextUsage {
max_tokens,
source:
crate::output::ContextUsageSource::FallbackEstimate
| crate::output::ContextUsageSource::TokenizerEstimate,
..
} => {
assert_eq!(*max_tokens, 16_000);
Some(index)
}
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(usage_positions.len(), 2);
let tool_result_position = sink
.outputs
.iter()
.position(|event| matches!(event, OutputEvent::ToolResult { .. }))
.unwrap();
let second_delta_position = sink
.outputs
.iter()
.position(|event| matches!(event, OutputEvent::AssistantDelta { text } if text == "done"))
.unwrap();
assert!(tool_result_position < usage_positions[1]);
assert!(usage_positions[1] < second_delta_position);
}
#[test]
fn context_usage_refreshes_from_provider_usage_input_tokens() {
let provider = ScriptedProvider::new(vec![vec![
ProviderEvent::Usage(Usage {
input: 42,
output: 9000,
total: 9042,
..Usage::default()
}),
done(),
]]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default()).with_context_budget(
crate::context::ContextBudget {
max_tokens: 16_000,
reserve_tokens: 0,
..crate::context::ContextBudget::default()
},
);
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"usage",
None,
None,
Path::new("."),
Some(&mut sink),
)
.unwrap();
assert!(sink.outputs.iter().any(|event| matches!(
event,
OutputEvent::ContextUsage {
current_tokens: 42,
max_tokens: 16_000,
..
}
)));
}
#[test]
fn anthropic_context_projection_calibrates_next_request_from_provider_usage() {
let temp = tempfile::TempDir::new().unwrap();
std::fs::write(temp.path().join("file.txt"), "hello").unwrap();
let provider = ScriptedProvider::new(vec![
vec![
ProviderEvent::Usage(Usage {
input: 100,
output: 1,
cache_read: 40,
cache_write: 20,
total: 161,
..Usage::default()
}),
read_call("call_1"),
done(),
],
vec![done()],
]);
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::new("claude-test", &[], &SkillDiscovery::default())
.with_provider_id(crate::providers::ANTHROPIC_PROVIDER)
.with_context_budget(crate::context::ContextBudget {
max_tokens: 16_000,
reserve_tokens: 0,
..crate::context::ContextBudget::default()
});
let mut sink = CapturingOutputSink::default();
agent
.run_print_with_tools_streaming_output(
&provider,
"read file",
Some(&tools),
None,
temp.path(),
Some(&mut sink),
)
.unwrap();
let request_usage = sink
.outputs
.iter()
.filter_map(|event| match event {
OutputEvent::ContextUsage {
current_tokens,
source,
request_sequence,
..
} => Some((*current_tokens, *source, *request_sequence)),
_ => None,
})
.collect::<Vec<_>>();
assert!(
request_usage
.iter()
.any(|usage| { *usage == (160, crate::output::ContextUsageSource::ProviderExact, 0,) })
);
assert!(request_usage.iter().any(|(tokens, source, sequence)| {
*sequence == 1
&& *tokens > 160
&& *source == crate::output::ContextUsageSource::LastProviderUsage
}));
}
#[test]
fn agent_continues_past_former_provider_tool_budget() {
let temp = tempfile::TempDir::new().unwrap();
std::fs::write(temp.path().join("file.txt"), "hello").unwrap();
let session = crate::sessions::SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
let responses = (0..65)
.map(|index| vec![read_call(&format!("call_{index}")), done()])
.chain(std::iter::once(text_done("done")))
.collect::<Vec<_>>();
let provider = ScriptedProvider::new(responses);
let tools = ToolRuntime::new(temp.path()).unwrap();
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
let output = agent
.run_print_with_tools(&provider, "loop", Some(&tools), Some(&session), temp.path())
.unwrap();
assert_eq!(output.text, "done");
assert_eq!(output.tool_results.len(), 65);
assert!(output.tool_results.iter().all(|result| result.success));
let requests = provider.requests();
assert_eq!(requests.len(), 66);
assert_eq!(
requests[65].tool_results().last().unwrap().call_id,
"call_64"
);
assert!(requests[65].response_items().iter().any(|item| {
item.get("type").and_then(serde_json::Value::as_str) == Some("function_call")
&& item.get("call_id").and_then(serde_json::Value::as_str) == Some("call_64")
}));
assert!(requests[65].response_items().iter().any(|item| {
item.get("type").and_then(serde_json::Value::as_str) == Some("function_call_output")
&& item.get("call_id").and_then(serde_json::Value::as_str) == Some("call_64")
}));
}
#[test]
fn resumed_session_replays_prior_history_before_current_prompt() {
let temp = tempfile::TempDir::new().unwrap();
let manager = crate::sessions::SessionManager::new(temp.path().join("sessions"));
let session = manager.create().unwrap();
for (event_type, payload) in [
("user_input", json!({"text":"prior request"})),
("assistant_output", json!({"text":"prior answer"})),
] {
session
.append(&SessionEvent::new(
event_type,
session.id().to_string(),
temp.path().to_path_buf(),
payload,
))
.unwrap();
}
let provider = ScriptedProvider::new(vec![vec![done()]]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
agent
.run_print_with_tools(
&provider,
"current request",
None,
Some(&session),
temp.path(),
)
.unwrap();
let items = provider.requests()[0].conversation_items().to_vec();
assert!(matches!(
&items[0],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::System
));
assert!(matches!(
&items[1],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User && message.content == "prior request"
));
assert!(matches!(
&items[2],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::Assistant && message.content == "prior answer"
));
assert!(matches!(
items.last().unwrap(),
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User && message.content == "current request"
));
}
#[test]
fn prompt_context_replays_full_structured_session_events() {
let temp = tempfile::TempDir::new().unwrap();
let manager = crate::sessions::SessionManager::new(temp.path().join("sessions"));
let session = manager.create().unwrap();
let old_type = format!("{}_{}", "context", "compaction");
session
.append(&SessionEvent::new(
old_type.clone(),
session.id().to_string(),
temp.path().to_path_buf(),
json!({"summary":"old summary line"}),
))
.unwrap();
session
.append(&SessionEvent::new(
"user_input",
session.id().to_string(),
temp.path().to_path_buf(),
json!({"text":"recent line"}),
))
.unwrap();
let provider = ScriptedProvider::new(vec![vec![done()]]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default());
agent
.run_print_with_tools(&provider, "next", None, Some(&session), temp.path())
.unwrap();
let requests = provider.requests();
assert!(matches!(
&requests[0].conversation_items()[0],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::System
&& !message.content.contains("Session context for continuity")
));
assert!(requests[0].conversation_items().iter().any(|item| matches!(
item,
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User
&& message.content == "recent line"
)));
assert!(matches!(
requests[0].conversation_items().last().unwrap(),
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User && message.content == "next"
));
}
#[test]
fn extremely_small_context_budget_fails_before_provider() {
let temp = tempfile::TempDir::new().unwrap();
let manager = crate::sessions::SessionManager::new(temp.path().join("sessions"));
let session = manager.create().unwrap();
session
.append(&SessionEvent::new(
"user_input",
session.id().to_string(),
temp.path().to_path_buf(),
json!({"text":"tiny"}),
))
.unwrap();
let provider = ScriptedProvider::new(vec![vec![done()]]);
let agent = AgentSession::new("model", &[], &SkillDiscovery::default()).with_context_budget(
crate::context::ContextBudget {
max_tokens: 1,
reserve_tokens: 0,
keep_recent_tokens: 1,
..crate::context::ContextBudget::default()
},
);
let error = agent
.run_print_with_tools(&provider, "next", None, Some(&session), temp.path())
.unwrap_err()
.to_string();
assert!(error.contains("full session history request is estimated"));
assert!(error.contains("exceeding threshold"));
assert!(provider.requests().is_empty());
}