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 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 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", temp.path())
},
)
.unwrap();
let requests = provider.requests();
assert_eq!(
requests[0].messages().last().unwrap().content,
"child sees #git-status"
);
}
#[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());
}