use crate::{
core::{
self,
agent_loop::{self, AgentLoop},
cache_warmer::CacheWarmer,
compaction::{self, CompactionReason},
extensions::{
BoundaryOutcome, ExecuteToolFn, Extension, ExtensionHook, ExtensionTool,
ExtensionUiRequest, GrammarSampling, InjectedChildTool, InjectedToolHandler,
MakeSubAgentFn, NestedCallLog, SubAgentControls, SubAgentRunner, SubAgentSpec,
ToolExecCtx, ToolExposure, ToolLoadout, UiNotifyLevel, is_deferred_tool_activated,
},
model_resolver::{self, ModelEntry},
model_scope::ScopedModel,
provider::{
AgentMessage, ContentBlock, GRAMMAR_SCHEMA_KEY, ModelConfig, StreamEvent, Usage,
usage::combine_usage,
},
session_manager::{self, Session, SessionCreateOptions},
session_v4, settings_manager,
skills::Skill,
system_prompt::{SystemPromptOptions, build_system_prompt},
tool_names,
tools::{self, ToolError, ToolExecutionMode, ToolResult, ToolResultAttachment},
virtual_models::{self, FailedRoute, ModelRouteReason, ModelRouteRequest},
},
error::{Error, Result},
extensions::{codemode, util::truncate_chars},
utils::{
image::{ImageResizeLimits, normalize_base64_attachment},
time::now_ms,
},
};
use futures_util::future::BoxFuture;
use serde_json::{Value, json};
use std::{
collections::{HashMap, HashSet, VecDeque},
sync::{
Arc, Mutex, OnceLock,
atomic::{AtomicBool, AtomicU32, Ordering},
},
time::{Duration, Instant},
};
use tokio::task::JoinHandle;
use tools::BashOutcome;
pub type JsonSink = Box<dyn FnMut(Value) + Send>;
pub type FinishTurn = Box<dyn FnMut(&TurnContext) -> Option<FinishTurnAction> + Send>;
pub type PrepareNextTurn = Box<dyn FnMut(&TurnContext) -> Option<NextTurnSnapshot> + Send>;
const SUBAGENT_GRACE_TURNS: u32 = 5;
const MAX_NESTED_CALLS: usize = 32;
const MAX_NESTED_TOOL_DEPTH: u32 = 4;
const NESTED_CALL_TEXT_CAP: usize = 4096;
const SUBAGENT_WRAP_UP: &str =
"Wrap up immediately — provide your final answer now. Do not start new work.";
const SUBAGENT_TOOL_NAMES: [&str; 4] = [
"Agent",
"SubagentWorkflow",
"get_subagent_result",
"steer_subagent",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FinishTurnAction {
End,
Continue,
}
#[derive(Debug, Clone)]
pub struct TurnContext {
pub message: AgentMessage,
pub tool_results: Vec<AgentMessage>,
pub context: Vec<AgentMessage>,
pub new_messages: Vec<AgentMessage>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ToolSelection {
Default(Vec<String>),
Allowlist(Vec<String>),
NoTools,
NoBuiltinTools,
}
impl ToolSelection {
pub fn tools(names: Vec<String>) -> Self {
ToolSelection::Default(names)
}
fn is_allowed(&self, name: &str, exclude: &[String]) -> bool {
if tool_names::matches_any(exclude, name) {
return false;
}
match self {
ToolSelection::Allowlist(entries) => {
tool_names::matches_any(entries, name) || self.is_mcp_reserved(name)
}
ToolSelection::NoTools => false,
ToolSelection::Default(_) | ToolSelection::NoBuiltinTools => true,
}
}
fn is_mcp_reserved(&self, name: &str) -> bool {
match self {
ToolSelection::Allowlist(entries) => {
!tool_names::matches_any(entries, name)
&& !tool_names::allowlist_filters_mcp(entries)
&& tool_names::is_mcp_tool_name(name)
}
_ => false,
}
}
fn declares_to_model(&self, name: &str, exposure: ToolExposure) -> bool {
!self.is_mcp_reserved(name) || exposure != ToolExposure::Direct
}
}
impl Default for ToolSelection {
fn default() -> Self {
ToolSelection::Default(Vec::new())
}
}
#[derive(Clone, Default)]
pub struct ToolRebuildCtx {
pub selection: ToolSelection,
pub exclude_tools: Vec<String>,
pub system_prompt: Option<String>,
pub append_system_prompt: Option<String>,
pub context_files: Vec<(String, String)>,
pub injected_tools: Arc<Vec<InjectedChildTool>>,
pub agent_id: Option<String>,
pub depth: u32,
}
fn compose_tools(
cwd: &str,
skills: &[Skill],
rebuild_ctx: &ToolRebuildCtx,
injected: &[InjectedChildTool],
) -> (String, Vec<(String, String, Value)>, Vec<String>) {
let selection = &rebuild_ctx.selection;
let exclude_tools = &rebuild_ctx.exclude_tools;
let mut selected: Vec<String> = match selection {
ToolSelection::Default(names) => names.clone(),
ToolSelection::Allowlist(_) => tools::all_tool_names()
.into_iter()
.filter(|n| selection.is_allowed(n, exclude_tools))
.collect(),
ToolSelection::NoTools | ToolSelection::NoBuiltinTools => Vec::new(),
};
for ext in core::extensions::registered() {
selected = ext.filter_tools(selected);
}
if !exclude_tools.is_empty() {
selected.retain(|t| !tool_names::matches_any(exclude_tools, t));
}
let tools_defs = tools::tool_defs(&selected);
let mut snippets = HashMap::new();
let mut prompt_guidelines: Vec<String> = Vec::new();
for t in &tools_defs {
if !t.snippet.trim().is_empty() {
snippets.insert(t.name.to_string(), t.snippet.clone());
}
prompt_guidelines.extend(t.prompt_guidelines.clone());
}
for ext in core::extensions::registered() {
let own = visible_extension_tools(ext.as_ref());
for t in own {
if !selection.is_allowed(&t.name, exclude_tools) {
continue;
}
if !selection.declares_to_model(&t.name, t.exposure) {
continue;
}
if !selected.contains(&t.name) {
selected.push(t.name.clone());
}
if !t.snippet.trim().is_empty() {
snippets.insert(t.name.clone(), t.snippet.clone());
}
prompt_guidelines.extend(t.prompt_guidelines.clone());
}
}
for t in injected {
snippets.insert(t.meta.name.clone(), t.meta.snippet.clone());
prompt_guidelines.extend(t.meta.prompt_guidelines.clone());
if !selected.iter().any(|n| n == &t.meta.name) {
selected.push(t.meta.name.clone());
}
}
let mut tools: Vec<(String, String, Value)> = tools_defs
.into_iter()
.map(|t| (t.name.to_string(), t.description, t.parameters))
.collect();
for ext in core::extensions::registered() {
let own = visible_extension_tools(ext.as_ref());
for t in own {
if !selection.is_allowed(&t.name, exclude_tools) {
continue;
}
if !selection.declares_to_model(&t.name, t.exposure) {
continue;
}
let parameters =
parameters_with_grammar_sampling(&t.parameters, t.grammar_sampling.as_ref());
if let Some(slot) = tools.iter_mut().find(|(n, _, _)| *n == t.name) {
*slot = (t.name.clone(), t.description.clone(), parameters);
} else {
tools.push((t.name.clone(), t.description.clone(), parameters));
}
}
}
for t in injected {
let entry = (
t.meta.name.clone(),
t.meta.description.clone(),
t.meta.parameters.clone(),
);
if let Some(slot) = tools.iter_mut().find(|(n, _, _)| *n == t.meta.name) {
*slot = entry;
} else {
tools.push(entry);
}
}
let hidden = apply_prepare_loadout(rebuild_ctx, &mut tools, injected);
let sp = build_system_prompt(&SystemPromptOptions {
cwd: cwd.to_string(),
selected_tools: Some(selected.clone()),
hidden_tools: hidden.clone(),
tool_snippets: snippets,
prompt_guidelines,
append_system_prompt: rebuild_ctx.append_system_prompt.clone(),
context_files: rebuild_ctx.context_files.clone(),
skills: skills.to_vec(),
custom_prompt: rebuild_ctx.system_prompt.clone(),
});
(sp, tools, hidden)
}
fn builtin_tool_descriptor(name: &str, description: &str, parameters: &Value) -> ExtensionTool {
let def = tools::tool_def(name);
let snippet = def.as_ref().map(|d| d.snippet.clone()).unwrap_or_default();
let mut tool = ExtensionTool::simple(name, description, parameters.clone(), &snippet);
tool.output_schema = def.and_then(|d| d.output_schema);
tool
}
fn parameters_with_grammar_sampling(schema: &Value, sampling: Option<&GrammarSampling>) -> Value {
let Some(sampling) = sampling else {
return schema.clone();
};
let Some(input) = core::provider::grammar_input_property(schema) else {
return schema.clone();
};
let mut grammar = serde_json::Map::new();
if let Some(lark) = &sampling.openai_lark {
grammar.insert("openai_lark".to_string(), json!(lark));
}
if let Some(regex) = &sampling.openai_regex {
grammar.insert("openai_regex".to_string(), json!(regex));
}
grammar.insert("input".to_string(), json!(input));
let mut out = schema.clone();
if let Some(obj) = out.as_object_mut() {
obj.insert(GRAMMAR_SCHEMA_KEY.to_string(), Value::Object(grammar));
}
out
}
fn script_tools_for(
active: &[(String, String, Value)],
injected: &[InjectedChildTool],
selection: &ToolSelection,
exclude: &[String],
) -> Vec<ExtensionTool> {
let names: Vec<&str> = active.iter().map(|(n, _, _)| n.as_str()).collect();
let builtin_names = tools::all_tool_names();
let mut out: Vec<ExtensionTool> = active
.iter()
.filter(|(n, _, _)| builtin_names.iter().any(|b| b == n))
.map(|(n, d, p)| builtin_tool_descriptor(n, d, p))
.collect();
for ext in core::extensions::registered() {
for t in ext.filter_extension_tools(ext.tools()) {
if t.name == codemode::TOOL_NAME {
continue;
}
let callable = if t.exposure == ToolExposure::Direct {
names.contains(&t.name.as_str())
|| (selection.is_mcp_reserved(&t.name)
&& !tool_names::matches_any(exclude, &t.name))
} else {
t.exposure
.is_script_callable(is_deferred_tool_activated(&t.name))
};
if callable {
out.push(t);
}
}
}
for t in injected {
out.push(t.meta.clone());
}
out
}
fn apply_prepare_loadout(
rebuild_ctx: &ToolRebuildCtx,
tools: &mut [(String, String, Value)],
injected: &[InjectedChildTool],
) -> Vec<String> {
let builtin_names = tools::all_tool_names();
let mut declared: Vec<ExtensionTool> = tools
.iter()
.filter(|(n, _, _)| builtin_names.iter().any(|b| b == n))
.map(|(n, d, p)| builtin_tool_descriptor(n, d, p))
.collect();
for ext in core::extensions::registered() {
for t in visible_extension_tools(ext.as_ref()) {
if rebuild_ctx
.selection
.is_allowed(&t.name, &rebuild_ctx.exclude_tools)
&& rebuild_ctx.selection.declares_to_model(&t.name, t.exposure)
{
declared.push(t);
}
}
}
let loadout = ToolLoadout {
declared,
callable: script_tools_for(
tools,
injected,
&rebuild_ctx.selection,
&rebuild_ctx.exclude_tools,
),
};
let mut hidden: Vec<String> = Vec::new();
for ext in core::extensions::registered() {
let changes = ext.prepare_loadout(&loadout);
for (name, description) in changes.descriptions {
if let Some(slot) = tools.iter_mut().find(|(n, _, _)| *n == name) {
slot.1 = description;
}
}
for name in changes.hidden_declarations {
if !hidden.contains(&name) {
hidden.push(name);
}
}
}
hidden
}
fn visible_extension_tools(ext: &dyn Extension) -> Vec<ExtensionTool> {
ext.filter_extension_tools(ext.tools())
.into_iter()
.filter(|t| {
t.exposure
.is_model_visible(is_deferred_tool_activated(&t.name))
})
.collect()
}
#[derive(Clone)]
struct ToolIndex<'a> {
name: &'a str,
injected: Option<&'a InjectedChildTool>,
extension: Option<ExtensionTool>,
builtin: bool,
}
impl<'a> ToolIndex<'a> {
fn resolve(injected: &'a [InjectedChildTool], name: &'a str) -> Self {
let inj = injected.iter().find(|t| t.meta.name == name);
let builtin = Agent::is_builtin_tool(name);
let extension = if inj.is_none() && !builtin {
find_extension_tool(name)
} else {
None
};
ToolIndex {
name,
injected: inj,
extension,
builtin,
}
}
fn exists(&self) -> bool {
self.injected.is_some() || self.builtin || self.extension.is_some()
}
fn is_builtin(&self) -> bool {
self.builtin
}
fn injected_handler(&self) -> Option<&'a InjectedToolHandler> {
self.injected.map(|t| &t.handler)
}
fn prepare_arguments(&self, args: &Value) -> Value {
if let Some(t) = self.injected
&& let Some(f) = t.meta.prepare_arguments
{
return f(args);
}
if self.builtin
&& let Some(d) = tools::tool_def(self.name)
&& let Some(f) = d.prepare_arguments
{
return f(args);
}
if let Some(f) = self.extension.as_ref().and_then(|t| t.prepare_arguments) {
return f(args);
}
args.clone()
}
fn parameters(&self) -> Option<Value> {
if let Some(t) = self.injected {
return Some(t.meta.parameters.clone());
}
if self.builtin
&& let Some(d) = tools::tool_def(self.name)
{
return Some(d.parameters.clone());
}
self.extension.as_ref().map(|t| t.parameters.clone())
}
fn execution_mode(&self) -> Option<ToolExecutionMode> {
if let Some(t) = self.injected
&& let Some(m) = t.meta.execution_mode
{
return Some(m);
}
if self.builtin
&& let Some(m) = tools::tool_execution_mode(self.name)
{
return Some(m);
}
self.extension.as_ref().and_then(|t| t.execution_mode)
}
}
fn find_extension_tool(name: &str) -> Option<ExtensionTool> {
core::extensions::registered()
.iter()
.find_map(|ext| ext.tools().into_iter().find(|t| t.name == name))
}
fn note_warning_once(key: &str) -> bool {
static SEEN: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
SEEN.get_or_init(Default::default)
.lock()
.unwrap()
.insert(key.to_string())
}
fn same_tool_names(recorded: Option<&[String]>, current: &[String]) -> bool {
let mut recorded: Vec<&str> = recorded
.unwrap_or_default()
.iter()
.map(String::as_str)
.collect();
let mut current: Vec<&str> = current.iter().map(String::as_str).collect();
recorded.sort_unstable();
current.sort_unstable();
recorded == current
}
#[derive(Debug, Clone)]
pub struct NextTurnSnapshot {
pub model: Option<ModelConfig>,
pub thinking_level: Option<String>,
pub context: Option<Vec<AgentMessage>>,
}
pub struct Agent {
pub model: ModelConfig,
pub routed_model: Option<ModelConfig>,
pub routed_thinking_level: Option<String>,
pub(crate) pending_failed_route: Option<FailedRoute>,
pub api_key_override: Option<String>,
pub skills: Vec<Skill>,
pub cwd: String,
pub messages: Vec<AgentMessage>,
pub system_prompt: String,
pub tools: Vec<(String, String, Value)>,
pub hidden_declarations: Vec<String>,
pub rebuild_ctx: ToolRebuildCtx,
settings_default_tools: Option<Vec<String>>,
settings_default_tool_modifiers: Vec<String>,
pub thinking_level: Option<String>,
pub last_usage: Option<Usage>,
pub session: Option<Session>,
pub json_sink: Option<Arc<Mutex<JsonSink>>>,
#[allow(clippy::type_complexity)]
pub on_response: Option<Arc<Mutex<Box<dyn FnMut(&Value) + Send>>>>,
#[allow(clippy::type_complexity)]
pub on_payload: Option<Arc<Mutex<Box<dyn FnMut(&mut Value) + Send>>>>,
pub run_abort: Arc<AtomicBool>,
pub steer_mode: String,
pub follow_up_mode: String,
pub auto_compaction: bool,
pub auto_retry: bool,
pub is_compacting: bool,
pub overflow_recovery_attempted: bool,
pub(crate) last_user_batch: Vec<AgentMessage>,
pub runtime_steer_inbox: Arc<Mutex<VecDeque<AgentMessage>>>,
pub is_subagent: bool,
pub offline: bool,
pub model_cycle: Vec<ScopedModel>,
pub model_is_fallback: bool,
pub tool_execution: ToolExecutionMode,
pub finish_turn: Option<Mutex<FinishTurn>>,
pub prepare_next_turn: Option<Mutex<PrepareNextTurn>>,
pub agent_loop: Option<AgentLoop>,
pub cache_warmer: CacheWarmer,
}
struct FinalizedToolCall {
call: ContentBlock,
name: String,
args: Value,
text: String,
details: Option<Value>,
usage: Option<Usage>,
terminate: bool,
is_error: bool,
attachments: Vec<ToolResultAttachment>,
content_override: Option<Vec<ContentBlock>>,
duration_ms: Option<u64>,
}
impl FinalizedToolCall {
fn call_id(&self) -> String {
match &self.call {
ContentBlock::ToolCall { id, .. } => id.clone(),
_ => String::new(),
}
}
fn error_result(
call: &ContentBlock,
name: &str,
args: &Value,
text: &str,
terminate: bool,
) -> Self {
FinalizedToolCall {
call: call.clone(),
name: name.to_string(),
args: args.clone(),
text: text.to_string(),
details: None,
usage: None,
terminate,
is_error: true,
attachments: Vec::new(),
content_override: None,
duration_ms: None,
}
}
fn from_tool_result(call: &ContentBlock, name: &str, args: &Value, t: ToolResult) -> Self {
let details = match t.structured_content {
Some(sc) => {
let mut map = match t.details {
Some(Value::Object(o)) => o,
_ => serde_json::Map::new(),
};
map.insert("structuredContent".to_string(), sc);
Some(Value::Object(map))
}
None => t.details,
};
FinalizedToolCall {
call: call.clone(),
name: name.to_string(),
args: args.clone(),
text: t.text,
details,
usage: t.usage,
terminate: t.terminate,
is_error: t.is_error,
attachments: t.attachments,
content_override: None,
duration_ms: None,
}
}
fn to_result_json(&self) -> Value {
let mut m = serde_json::Map::new();
m.insert(
"content".to_string(),
json!([{ "type": "text", "text": self.text }]),
);
m.insert(
"details".to_string(),
self.details.clone().unwrap_or_else(|| json!({})),
);
if let Some(u) = &self.usage {
m.insert(
"usage".to_string(),
serde_json::to_value(u).unwrap_or(Value::Null),
);
}
if self.terminate {
m.insert("terminate".to_string(), json!(true));
}
Value::Object(m)
}
}
struct AgentTemplate {
model: ModelConfig,
api_key_override: Option<String>,
skills: Vec<Skill>,
cwd: String,
system_prompt: String,
tools: Vec<(String, String, Value)>,
rebuild_ctx: ToolRebuildCtx,
thinking_level: Option<String>,
offline: bool,
auto_retry: bool,
tool_execution: ToolExecutionMode,
parent_messages: Arc<Vec<AgentMessage>>,
parent_session_id: Option<String>,
}
impl Agent {
const MAX_AGENT_RETRY_DELAY_MS: u64 = 60_000;
fn snapshot_template(&self) -> AgentTemplate {
let system_override = self
.rebuild_ctx
.system_prompt
.clone()
.or_else(|| Some(self.system_prompt.clone()))
.unwrap_or_default();
AgentTemplate {
model: self.model.clone(),
api_key_override: self.api_key_override.clone(),
skills: self.skills.clone(),
cwd: self.cwd.clone(),
system_prompt: system_override,
tools: self.tools.clone(),
rebuild_ctx: self.rebuild_ctx.clone(),
thinking_level: self.thinking_level.clone(),
offline: self.offline,
auto_retry: self.auto_retry,
tool_execution: self.tool_execution,
parent_messages: Arc::new(self.messages.clone()),
parent_session_id: self.session.as_ref().map(|s| s.session_id.clone()),
}
}
pub fn exec_ctx(&self) -> ToolExecCtx {
let child_tpl = Arc::new(self.snapshot_template());
let entries = self.session_branch_entries();
make_exec_ctx(
&self.cwd,
&child_tpl,
&self.run_abort,
&self.json_sink,
None,
entries,
)
}
fn session_branch_entries(&self) -> Option<Arc<Vec<Value>>> {
let session = self.session.as_ref()?;
let entries = match session.get_leaf_id() {
Some(leaf) => session.get_branch(leaf),
None => Vec::new(),
};
Some(Arc::new(entries))
}
fn child_from_template(tpl: &AgentTemplate) -> Agent {
Agent {
model: tpl.model.clone(),
routed_model: None,
routed_thinking_level: None,
pending_failed_route: None,
api_key_override: tpl.api_key_override.clone(),
skills: tpl.skills.clone(),
cwd: tpl.cwd.clone(),
messages: Vec::new(),
system_prompt: tpl.system_prompt.clone(),
tools: tpl.tools.clone(),
hidden_declarations: Vec::new(),
rebuild_ctx: tpl.rebuild_ctx.clone(),
settings_default_tools: None,
settings_default_tool_modifiers: Vec::new(),
thinking_level: tpl.thinking_level.clone(),
last_usage: None,
session: None,
json_sink: None,
on_response: None,
on_payload: None,
run_abort: Arc::new(AtomicBool::new(false)),
steer_mode: String::from("one-at-a-time"),
follow_up_mode: String::from("one-at-a-time"),
auto_compaction: false,
auto_retry: tpl.auto_retry,
is_compacting: false,
overflow_recovery_attempted: false,
last_user_batch: Vec::new(),
runtime_steer_inbox: Arc::new(Mutex::new(VecDeque::new())),
is_subagent: true,
offline: tpl.offline,
model_cycle: Vec::new(),
model_is_fallback: false,
tool_execution: tpl.tool_execution,
finish_turn: None,
prepare_next_turn: None,
agent_loop: None,
cache_warmer: CacheWarmer::new(),
}
}
#[allow(clippy::too_many_arguments)]
pub fn new(
model_entry: ModelEntry,
cwd: String,
selection: ToolSelection,
exclude_tools: Vec<String>,
thinking_level: Option<String>,
system_prompt: Option<String>,
append_system_prompt: Option<String>,
context_files: Vec<(String, String)>,
session: Option<Session>,
api_key_override: Option<String>,
offline: bool,
skills: Vec<Skill>,
) -> Result<Agent> {
let model = model_resolver::model_config_from_entry(
&model_entry,
api_key_override.clone(),
session.as_ref().map(|s| s.session_id.clone()),
);
let thinking_level = thinking_level.map(|l| model.clamp_thinking_level(&l));
let rebuild_ctx = ToolRebuildCtx {
selection: selection.clone(),
exclude_tools: exclude_tools.clone(),
system_prompt: system_prompt.clone(),
append_system_prompt: append_system_prompt.clone(),
context_files: context_files.clone(),
injected_tools: Arc::new(Vec::new()),
agent_id: None,
depth: 0,
};
let mut history = Vec::new();
if let Some(s) = &session {
history = s.build_context_messages();
core::extensions::restore_deferred_activations(&history);
}
let (sp, tools, hidden) = compose_tools(&cwd, &skills, &rebuild_ctx, &[]);
Ok(Agent {
model,
routed_model: None,
routed_thinking_level: None,
pending_failed_route: None,
api_key_override,
skills,
cwd,
messages: history,
system_prompt: sp,
tools,
hidden_declarations: hidden,
rebuild_ctx,
settings_default_tools: None,
settings_default_tool_modifiers: Vec::new(),
thinking_level,
last_usage: None,
session,
json_sink: None,
on_response: None,
on_payload: None,
run_abort: Arc::new(AtomicBool::new(false)),
steer_mode: settings_manager::read_settings_steering_mode(),
follow_up_mode: settings_manager::read_settings_follow_up_mode(),
auto_compaction: settings_manager::read_settings_auto_compact(),
auto_retry: true,
is_compacting: false,
overflow_recovery_attempted: false,
last_user_batch: Vec::new(),
runtime_steer_inbox: Arc::new(Mutex::new(VecDeque::new())),
is_subagent: false,
offline,
model_cycle: Vec::new(),
model_is_fallback: false,
tool_execution: ToolExecutionMode::Parallel,
finish_turn: None,
prepare_next_turn: None,
agent_loop: None,
cache_warmer: CacheWarmer::new(),
})
}
pub fn set_settings_default_tools(&mut self, names: Option<Vec<String>>) {
self.settings_default_tools = names;
}
pub fn set_default_tool_modifiers(&mut self, modifiers: Vec<String>) {
self.settings_default_tool_modifiers = modifiers;
}
pub fn reload_default_tools(&mut self) {
let Some(previous) = self.settings_default_tools.take() else {
return;
};
let raw = settings_manager::read_settings_default_tools().unwrap_or_else(|| {
settings_manager::DEFAULT_TOOL_NAMES
.iter()
.map(|s| s.to_string())
.collect()
});
let next =
settings_manager::apply_tool_modifiers(&raw, &self.settings_default_tool_modifiers);
if let ToolSelection::Default(current) = &mut self.rebuild_ctx.selection {
for name in next.iter().filter(|n| !previous.contains(*n)).cloned() {
if !current.contains(&name) {
current.push(name);
}
}
}
self.settings_default_tools = Some(next);
}
pub fn rebuild_tools(&mut self) {
let (sp, tools, hidden) = compose_tools(
&self.cwd,
&self.skills,
&self.rebuild_ctx,
&self.rebuild_ctx.injected_tools,
);
self.system_prompt = sp;
self.tools = tools;
self.hidden_declarations = hidden;
if let Some(s) = self.session.as_mut() {
let names: Vec<String> = self.tools.iter().map(|t| t.0.clone()).collect();
if !same_tool_names(s.last_active_tools().as_deref(), &names) {
s.append_active_tools_change(&names);
}
}
}
pub fn get_active_tools(&self) -> Vec<String> {
self.tools.iter().map(|t| t.0.clone()).collect()
}
pub fn get_all_tools(&self) -> Vec<String> {
if matches!(self.rebuild_ctx.selection, ToolSelection::NoTools) {
return Vec::new();
}
let mut names = tools::all_tool_names();
for ext in core::extensions::registered() {
let own = ext.filter_extension_tools(ext.tools());
for t in own {
if !names.contains(&t.name) {
names.push(t.name);
}
}
}
names.retain(|n| {
self.rebuild_ctx
.selection
.is_allowed(n, &self.rebuild_ctx.exclude_tools)
});
names
}
pub fn set_active_tools(&mut self, names: Vec<String>) {
self.rebuild_ctx.selection = ToolSelection::Default(names);
self.rebuild_tools();
}
pub fn clear_all_queues(&mut self) {
self.runtime_steer_inbox.lock().unwrap().clear();
}
pub fn reset(&mut self) -> Result<()> {
self.messages.clear();
self.last_user_batch.clear();
self.clear_all_queues();
Ok(())
}
pub(crate) fn emit(&mut self, event: Value) {
if let Some(sink) = self.json_sink.as_ref()
&& let Ok(mut g) = sink.lock()
{
let f: &mut dyn FnMut(Value) = g.as_mut();
f(event);
}
}
pub fn set_on_response(&mut self, cb: Box<dyn FnMut(&Value) + Send>) {
self.on_response = Some(Arc::new(Mutex::new(cb)));
}
pub fn set_on_payload(&mut self, cb: Box<dyn FnMut(&mut Value) + Send>) {
self.on_payload = Some(Arc::new(Mutex::new(cb)));
}
pub fn attach_json_sink(&mut self, sink: JsonSink) {
self.json_sink = Some(Arc::new(Mutex::new(sink)));
}
pub fn request_abort(&self) {
self.run_abort.store(true, Ordering::Relaxed);
}
pub fn abort_and_detach(&mut self) -> Arc<AtomicBool> {
self.request_abort();
std::mem::replace(&mut self.run_abort, Arc::new(AtomicBool::new(false)))
}
pub fn switch_model(&mut self, provider: &str, model_id: &str, persist: bool) -> Result<()> {
let entry = model_resolver::find_model(provider, model_id)?;
self.model = model_resolver::model_config_from_entry(
&entry,
self.api_key_override.clone(),
self.session.as_ref().map(|s| s.session_id.clone()),
);
if let Some(s) = self.session.as_mut() {
s.append_model_change(provider, model_id);
}
if persist {
_ = settings_manager::write_default_model_and_provider(&entry.provider, &entry.id);
}
self.model_is_fallback = false;
let current = self
.thinking_level
.clone()
.unwrap_or_else(|| "off".to_string());
self.set_thinking_level(¤t);
Ok(())
}
pub fn get_available_thinking_levels(&self) -> Vec<&'static str> {
self.model.supported_thinking_levels()
}
pub fn selected_is_virtual(&self) -> bool {
virtual_models::is_virtual_api(&self.model.api)
}
pub fn effective_model(&self) -> &ModelConfig {
self.routed_model.as_ref().unwrap_or(&self.model)
}
pub fn effective_thinking_level(&self) -> Option<String> {
self.routed_thinking_level
.clone()
.or_else(|| self.thinking_level.clone())
}
pub fn clear_routed_model(&mut self) {
self.routed_model = None;
self.routed_thinking_level = None;
}
pub fn take_pending_failed_route(&mut self) -> Option<FailedRoute> {
self.pending_failed_route.take()
}
pub async fn direct_model(&self) -> Result<ModelConfig> {
if !self.selected_is_virtual() {
return Ok(self.model.clone());
}
let request = ModelRouteRequest {
selected: &self.model,
thinking_level: self.thinking_level.as_deref(),
reason: ModelRouteReason::Direct,
previous: None,
failed: None,
state: None,
messages: &self.messages,
abort: Some(self.run_abort.clone()),
};
let route = virtual_models::resolve_route(
request,
self.api_key_override.clone(),
self.session.as_ref().map(|s| s.session_id.clone()),
)
.await?;
Ok(route.model)
}
pub async fn summarization_model(&self) -> Result<ModelConfig> {
match &self.routed_model {
Some(m) => Ok(m.clone()),
None => self.direct_model().await,
}
}
pub fn set_session_name(&mut self, name: &str) -> bool {
let Some(s) = self.session.as_mut() else {
return false;
};
s.name = Some(name.to_string());
s.append_session_info(name);
self.emit(json!({
"type": "session_info_changed",
"name": name,
}));
true
}
pub fn set_thinking_level(&mut self, level: &str) -> String {
let (effective, changed) = self.apply_thinking_level(level);
if changed {
self.emit(json!({
"type": "thinking_level_changed",
"level": effective,
}));
}
effective
}
pub fn set_thinking_level_silent(&mut self, level: &str, persist: bool) -> String {
let effective = self.apply_thinking_level(level).0;
if persist {
_ = settings_manager::write_settings_default_thinking_level(level);
}
effective
}
fn apply_thinking_level(&mut self, level: &str) -> (String, bool) {
let effective = self.model.clamp_thinking_level(level);
let changed = self.thinking_level.as_deref() != Some(effective.as_str());
if changed {
self.thinking_level = Some(effective.clone());
if let Some(s) = self.session.as_mut() {
s.append_thinking_level_change(&effective);
}
}
(effective, changed)
}
pub async fn maybe_compact(&mut self, force: bool) -> Result<bool> {
let reason = if force {
compaction::CompactionReason::Manual
} else {
compaction::CompactionReason::Threshold
};
Ok(self.maybe_compact_with(reason, None).await?.is_some())
}
pub async fn maybe_compact_with(
&mut self,
reason: compaction::CompactionReason,
custom_instructions: Option<&str>,
) -> Result<Option<Value>> {
self.is_compacting = true;
let result = self.maybe_compact_inner(reason, custom_instructions).await;
self.is_compacting = false;
result
}
async fn maybe_compact_inner(
&mut self,
reason: compaction::CompactionReason,
custom_instructions: Option<&str>,
) -> Result<Option<Value>> {
let limit_model = self
.summarization_model()
.await
.unwrap_or_else(|_| self.model.clone());
let settings = compaction::CompactionSettings::from_settings()
.for_model(&limit_model.provider, &limit_model.model_id);
let estimate = compaction::estimate_context_tokens(&self.messages);
let context_tokens = estimate.tokens;
let should = match reason {
compaction::CompactionReason::Manual | compaction::CompactionReason::Overflow => true,
compaction::CompactionReason::Threshold => {
let latest_compaction_ts = self.session.as_ref().and_then(|s| {
s.state.get_entries().iter().rev().find_map(|e| match e {
session_v4::Entry::Compaction { base, .. } => Some(base.timestamp),
_ => None,
})
});
let stale = compaction::usage_anchor_is_stale(
&estimate,
&self.messages,
latest_compaction_ts,
);
!stale
&& compaction::should_compact(
context_tokens,
limit_model.context_window,
&settings,
)
}
};
if !should {
return Ok(None);
}
self.emit(json!({
"type": "compaction_start",
"reason": reason.as_str(),
}));
let mut previous_summary: Option<String> = None;
let mut summarize_start = 0usize;
for (i, m) in self.messages.iter().enumerate() {
if m.role == "compactionSummary" {
previous_summary = Some(m.text());
summarize_start = i + 1;
break;
}
}
let keep_recent = settings.keep_recent_tokens;
let cut_index =
compaction::find_compaction_cut_index(&self.messages, summarize_start, keep_recent);
if cut_index <= summarize_start {
return Ok(None);
}
let messages_to_summarize = self.messages[summarize_start..cut_index].to_vec();
let kept_messages = self.messages[cut_index..].to_vec();
if messages_to_summarize.is_empty() {
return Ok(None);
}
let conversation_text = compaction::serialize_conversation(&messages_to_summarize);
let summary_prompt =
compaction::build_summary_prompt(previous_summary.as_deref(), custom_instructions);
let user_prompt = conversation_text.to_string();
for ext in core::extensions::registered() {
ext.on_before_compact(context_tokens, reason.as_str());
}
let mut custom_summary: Option<String> = None;
for ext in core::extensions::registered() {
if let Some(s) = ext.custom_summarize(
compaction::SUMMARIZATION_SYSTEM_PROMPT,
&user_prompt,
&summary_prompt,
) {
custom_summary = Some(s);
break;
}
}
let summary_result = match custom_summary {
Some(s) => Ok((s, Usage::default())),
None => {
let reserve = settings.reserve_tokens.max(1);
let budget = ((reserve as f64) * 0.8).floor() as u32;
let max_tokens = limit_model.max_tokens.unwrap_or(u32::MAX).min(budget);
core::provider::simple_completion(
&limit_model,
compaction::SUMMARIZATION_SYSTEM_PROMPT,
&format!("{}\n\n{}", user_prompt, summary_prompt),
Some(max_tokens),
)
.await
}
};
let (summary, summary_usage) = match summary_result {
Ok(v) => v,
Err(e) => {
for ext in core::extensions::registered() {
ext.on_compact_failure(reason.as_str());
}
return Err(e);
}
};
if summary.trim().is_empty() {
for ext in core::extensions::registered() {
ext.on_compact_failure(reason.as_str());
}
return Err(Error::msg(
"compaction summary generation failed: model returned empty content",
));
}
let mut file_ops = compaction::FileOperations::default();
if let Some(s) = self.session.as_ref()
&& let Some(prev) = s
.get_entries()
.iter()
.rev()
.find(|e| {
e.get("type").and_then(|v| v.as_str()) == Some("compaction")
&& e.get("fromHook").and_then(|v| v.as_bool()) != Some(true)
})
.and_then(|e| e.get("details").cloned())
{
if let Some(prev_read) = prev.get("readFiles").and_then(|v| v.as_array()) {
for f in prev_read {
if let Some(fs) = f.as_str() {
file_ops.read.insert(fs.to_string());
}
}
}
if let Some(prev_mod) = prev.get("modifiedFiles").and_then(|v| v.as_array()) {
for f in prev_mod {
if let Some(fs) = f.as_str() {
file_ops.edited.insert(fs.to_string());
}
}
}
}
for msg in &messages_to_summarize {
compaction::extract_file_ops_from_message(msg, &mut file_ops);
}
let (read_files, modified_files) = compaction::compute_file_lists(&file_ops);
let mut summary_with_files = summary;
summary_with_files.push_str(&compaction::format_file_operations(
&read_files,
&modified_files,
));
let first_kept_id = self
.session
.as_ref()
.and_then(|s| s.first_kept_entry_id_for_cut(&self.messages, cut_index))
.or_else(|| kept_messages.iter().find_map(|m| m.entry_id.clone()))
.or_else(|| {
self.session
.as_ref()
.and_then(|s| s.get_leaf_id().map(|x| x.to_string()))
});
let tokens_before = context_tokens;
let details = json!({ "readFiles": read_files, "modifiedFiles": modified_files });
if let Some(s) = self.session.as_mut() {
s.append_compaction(
&summary_with_files,
u64::from(tokens_before),
first_kept_id.as_deref(),
Some(details.clone()),
Some(&summary_usage),
);
}
let projected = self
.session
.as_ref()
.map(|s| s.build_context_messages())
.filter(|msgs| !msgs.is_empty());
match projected {
Some(msgs) => self.messages = msgs,
None => {
let mut summary_msg = AgentMessage::user_text("");
summary_msg.role = "compactionSummary".into();
summary_msg.content = vec![ContentBlock::Text {
text: summary_with_files.clone(),
text_signature: None,
}];
let mut new_messages = vec![summary_msg];
new_messages.extend(kept_messages);
self.messages = new_messages;
}
}
let estimated_tokens_after = compaction::estimate_context_tokens(&self.messages).tokens;
self.emit(json!({
"type": "compaction_end",
"reason": reason.as_str(),
"result": "completed",
"tokensBefore": tokens_before,
"estimatedTokensAfter": estimated_tokens_after,
}));
let usage_json = serde_json::to_value(&summary_usage).unwrap_or(Value::Null);
Ok(Some(json!({
"summary": summary_with_files,
"firstKeptEntryId": first_kept_id,
"tokensBefore": tokens_before,
"estimatedTokensAfter": estimated_tokens_after,
"reason": reason.as_str(),
"usage": usage_json,
"details": details,
})))
}
pub async fn navigate_tree(
&mut self,
target_id: &str,
summarize: bool,
) -> Result<(bool, Option<String>)> {
let summary_model = if summarize {
Some(self.summarization_model().await?)
} else {
None
};
let Some(session) = self.session.as_mut() else {
return Err(Error::msg("no current session"));
};
let Some(old_leaf) = session.get_leaf_id() else {
return Err(Error::msg("session has no leaf"));
};
if old_leaf == target_id {
return Ok((false, None));
}
let mut editor_text: Option<String> = None;
let mut effective_target = target_id.to_string();
if let Some(entry) = session
.get_entries()
.iter()
.find(|e| e.get("id").and_then(|v| v.as_str()) == Some(target_id))
{
let ty = entry.get("type").and_then(|v| v.as_str()).unwrap_or("");
if ty == "message" {
let role = entry
.get("message")
.and_then(|m| m.get("role").and_then(|r| r.as_str()))
.unwrap_or("");
if role == "user" {
editor_text = entry
.get("message")
.and_then(|m| m.get("content").and_then(|c| c.as_array()))
.map(|blocks| {
blocks
.iter()
.filter_map(|b| b.get("text").and_then(|t| t.as_str()))
.collect::<Vec<_>>()
.join("\n")
})
.filter(|s| !s.is_empty());
if let Some(parent) = entry.get("parentId").and_then(|v| v.as_str()) {
effective_target = parent.to_string();
}
}
}
}
let target_id = effective_target;
if !session
.get_entries()
.iter()
.any(|e| e.get("id").and_then(|v| v.as_str()) == Some(target_id.as_str()))
{
return Err(Error::EntryNotFound(target_id.clone()));
}
let old_branch = session.get_branch(old_leaf);
let target_branch = session.get_branch(&target_id);
let old_ids: HashSet<String> = old_branch
.iter()
.filter_map(|e| e.get("id").and_then(|v| v.as_str()).map(|s| s.to_string()))
.collect();
let common_ancestor = target_branch.iter().rev().find(|e| {
e.get("id")
.and_then(|v| v.as_str())
.is_some_and(|id| old_ids.contains(id))
});
let common_id = common_ancestor
.and_then(|e| e.get("id").and_then(|v| v.as_str()))
.unwrap_or("");
let abandoned: Vec<Value> = old_branch
.iter()
.skip_while(|e| {
e.get("id")
.and_then(|v| v.as_str())
.is_some_and(|id| common_id.is_empty() || id != common_id)
})
.skip(1)
.cloned()
.collect();
if summarize && !abandoned.is_empty() {
let mut messages: Vec<AgentMessage> = Vec::new();
let mut file_ops = compaction::FileOperations::default();
for entry in &abandoned {
let ty = entry.get("type").and_then(|v| v.as_str()).unwrap_or("");
match ty {
"message" => {
let Some(msg) = entry.get("message") else {
continue;
};
let Ok(mut msg) = serde_json::from_value::<AgentMessage>(msg.clone())
else {
continue;
};
if msg.role == "toolResult" {
continue;
}
msg.entry_id = entry
.get("id")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
compaction::extract_file_ops_from_message(&msg, &mut file_ops);
messages.push(msg);
}
"branch_summary" => {
if let Some(summary) = entry.get("summary").and_then(|v| v.as_str()) {
let mut m = AgentMessage::user_text(summary);
m.role = "branchSummary".to_string();
messages.push(m);
}
if let Some(details) = entry.get("details") {
if let Some(read) = details.get("readFiles").and_then(|v| v.as_array())
{
for f in read {
if let Some(fs) = f.as_str() {
file_ops.read.insert(fs.to_string());
}
}
}
if let Some(modified) =
details.get("modifiedFiles").and_then(|v| v.as_array())
{
for f in modified {
if let Some(fs) = f.as_str() {
file_ops.edited.insert(fs.to_string());
}
}
}
}
}
"compaction" => {
if let Some(summary) = entry.get("summary").and_then(|v| v.as_str()) {
let mut m = AgentMessage::user_text(summary);
m.role = "compactionSummary".to_string();
messages.push(m);
}
}
_ => {}
}
}
if messages.is_empty() {
session.set_leaf(Some(&target_id))?;
} else {
let conversation = compaction::serialize_conversation(&messages);
let prompt = format!(
"<conversation>\n{}\n</conversation>\n\n{}",
conversation,
compaction::BRANCH_SUMMARY_PROMPT
);
let summary_model = summary_model
.as_ref()
.expect("summary model resolved before borrowing session when summarize=true");
let (mut summary, summary_usage) = core::provider::simple_completion(
summary_model,
compaction::SUMMARIZATION_SYSTEM_PROMPT,
&prompt,
Some(branch_summary_max_tokens(summary_model)),
)
.await?;
summary = format!("{}{}", compaction::BRANCH_SUMMARY_PREAMBLE, summary);
let (read_files, modified_files) = compaction::compute_file_lists(&file_ops);
summary.push_str(&compaction::format_file_operations(
&read_files,
&modified_files,
));
let details = json!({ "readFiles": read_files, "modifiedFiles": modified_files });
session.append_branch_summary(
&target_id,
&summary,
Some(details),
Some(&summary_usage),
)?;
}
} else {
session.set_leaf(Some(&target_id))?;
}
let messages = session.build_context_messages();
self.messages = messages;
Ok((true, editor_text))
}
pub fn rewind_last_user_input(&mut self) -> Result<Option<String>> {
let Some(session) = self.session.as_mut() else {
return Ok(None);
};
let removed = session.rewind_last_user_input()?;
if removed.is_some() {
self.messages = session.build_context_messages();
}
Ok(removed)
}
pub async fn prompt(&mut self, user_message: &str) -> Result<String> {
let msg = AgentMessage::user_text(user_message);
self.prompt_message(msg).await
}
pub async fn prompt_message(&mut self, msg: AgentMessage) -> Result<String> {
self.prompt_messages(vec![msg]).await
}
pub async fn prompt_messages(&mut self, msgs: Vec<AgentMessage>) -> Result<String> {
self.reset_run_state();
self.inject_prompt_messages(msgs);
self.compact_before_run().await;
let mut result = self.run_with_recovery().await;
let mut continuation_index = 0usize;
while let Some(outcome) = self.settle_boundary(continuation_index) {
self.apply_boundary_drafts(&outcome);
if !outcome.append.is_empty() {
self.inject_prompt_messages(outcome.append);
}
if !outcome.r#continue {
break;
}
self.overflow_recovery_attempted = false;
result = self.run_with_recovery().await;
continuation_index += 1;
}
self.emit_agent_settled();
result
}
async fn run_with_recovery(&mut self) -> Result<String> {
let result = self.run_loop().await;
if self.should_recover_overflow() {
self.recover_from_overflow(result).await
} else {
self.retry_transient_errors(result).await
}
}
fn settle_boundary(&mut self, continuation_index: usize) -> Option<BoundaryOutcome> {
if !core::extensions::has_boundary_handlers() {
return None;
}
if self.run_abort.load(Ordering::Relaxed) {
return None;
}
if self.messages.is_empty() {
return None;
}
let context_messages: Vec<Value> = self
.messages
.iter()
.map(|m| serde_json::to_value(m).unwrap_or(Value::Null))
.collect();
let has_queued_messages = !self.runtime_steer_inbox.lock().unwrap().is_empty();
let event = json!({
"type": "agent_before_settle",
"outcome": self.activity_outcome(),
"contextMessages": context_messages,
"hasQueuedMessages": has_queued_messages,
"continuationIndex": continuation_index,
"agentScope": self.agent_scope(),
});
let outcome = core::extensions::dispatch_boundary(&event);
if !outcome.r#continue && !outcome.has_drafts() {
return None;
}
Some(outcome)
}
pub(crate) fn apply_boundary_drafts(&mut self, outcome: &BoundaryOutcome) -> bool {
if !outcome.has_drafts() {
return false;
}
let estimated_tokens = if outcome
.compactions
.iter()
.any(|c| c.tokens_before.is_none())
{
u64::from(compaction::estimate_context_tokens(&self.messages).tokens)
} else {
0
};
let Some(session) = self.session.as_mut() else {
return false;
};
let mut changed = false;
for draft in &outcome.edits {
let replacement = draft
.replacement
.clone()
.map(|content| session_v4::ContextEditReplacement { content });
if session
.append_context_edit(&draft.target_id, replacement)
.is_ok()
{
changed = true;
}
}
for compaction in &outcome.compactions {
let tokens_before = compaction.tokens_before.unwrap_or(estimated_tokens);
session.append_retain_none_compaction(
&compaction.summary,
tokens_before,
compaction.details.clone(),
None,
);
changed = true;
}
if changed {
self.refresh_context_from_session();
}
changed
}
fn refresh_context_from_session(&mut self) {
let Some(session) = self.session.as_ref() else {
return;
};
let rebuilt = session.build_context_messages();
if !rebuilt.is_empty() {
self.messages = rebuilt;
}
}
fn agent_scope(&self) -> &'static str {
if self.is_subagent { "subagent" } else { "main" }
}
fn activity_outcome(&self) -> &'static str {
match self
.messages
.iter()
.rev()
.find(|m| m.role == "assistant")
.and_then(|m| m.stop_reason.as_deref())
{
Some("aborted") => "aborted",
Some("error") => "error",
_ => "completed",
}
}
fn reset_run_state(&mut self) {
self.run_abort.store(false, Ordering::Relaxed); self.overflow_recovery_attempted = false; self.routed_model = None; self.routed_thinking_level = None;
self.pending_failed_route = None; }
fn inject_prompt_messages(&mut self, msgs: Vec<AgentMessage>) {
let mut injected = Vec::with_capacity(msgs.len());
for mut msg in msgs {
if let Some(s) = self.session.as_mut() {
let entry_id = s.append_message(&msg);
msg.entry_id = Some(entry_id);
}
self.messages.push(msg.clone());
injected.push(msg);
}
self.last_user_batch = injected;
}
async fn compact_before_run(&mut self) {
if self.auto_compaction && !self.selected_is_virtual() {
_ = self
.maybe_compact_with(CompactionReason::Threshold, None)
.await;
}
}
fn last_message_is_overflow(&self) -> bool {
let model = self.effective_model();
self.messages.last().is_some_and(|m| {
m.role == "assistant"
&& Self::message_is_context_overflow(m, model.context_window, &model.provider)
})
}
fn should_recover_overflow(&self) -> bool {
if !self.auto_retry || !self.last_message_is_overflow() {
return false;
}
let model_id = self.effective_model().model_id.clone();
self.messages
.last()
.and_then(|m| m.model.clone())
.map(|m| m == model_id)
.unwrap_or(true)
}
fn last_assistant_error(&self) -> Option<String> {
self.messages.last().and_then(|m| {
(m.role == "assistant" && m.stop_reason.as_deref() == Some("error"))
.then(|| m.error_message.clone())
.flatten()
})
}
fn omit_entry_from_context(&mut self, entry_id: Option<&str>) -> bool {
let Some(entry_id) = entry_id else {
return false;
};
let Some(session) = self.session.as_mut() else {
return false;
};
session.append_context_edit(entry_id, None).is_ok()
}
fn capture_failed_route(&mut self) {
if !self.selected_is_virtual() {
return;
}
let failed = self.messages.last().and_then(|m| {
if m.role != "assistant" || m.stop_reason.as_deref() != Some("error") {
return None;
}
if m.api.as_deref().is_some_and(virtual_models::is_virtual_api) {
return None;
}
let (provider, model_id) = (m.provider.clone()?, m.model.clone()?);
Some(virtual_models::FailedRoute {
provider,
model_id,
thinking_level: m.thinking_level.clone(),
error_message: m.error_message.clone(),
stop_reason: m.stop_reason.clone(),
})
});
if failed.is_some() {
self.pending_failed_route = failed;
}
}
fn emit_overflow_recovery_failed(&mut self) {
self.emit(json!({
"type": "compaction_end",
"reason": "overflow",
"result": "failed",
"willRetry": false,
"errorMessage":
"Context overflow recovery failed after one compact-and-retry attempt.",
}));
}
async fn recover_from_overflow(&mut self, result: Result<String>) -> Result<String> {
if self.overflow_recovery_attempted {
self.emit_overflow_recovery_failed();
return result;
}
self.overflow_recovery_attempted = true;
self.capture_failed_route();
if self.messages.last().is_some_and(|m| m.role == "assistant") {
let omitted_id = self.messages.last().and_then(|m| m.entry_id.clone());
self.omit_entry_from_context(omitted_id.as_deref());
self.messages.pop();
}
let ok = matches!(
self.maybe_compact_with(CompactionReason::Overflow, None)
.await,
Ok(Some(_))
);
if !ok {
return result;
}
let retried = self.run_loop().await;
if self.last_message_is_overflow() {
self.emit_overflow_recovery_failed();
}
retried
}
async fn retry_transient_errors(&mut self, result: Result<String>) -> Result<String> {
if !self.auto_retry || !settings_manager::read_settings_retry_enabled() {
return result;
}
let max_retries = settings_manager::read_settings_retry_max_retries();
let base_delay = settings_manager::read_settings_retry_base_delay_ms();
let mut attempt = 0u32;
let mut result = result;
while attempt < max_retries {
let Some(err_msg) = self.last_assistant_error() else {
break;
};
if !Self::is_retryable_assistant_error(&err_msg) {
break;
}
attempt += 1;
let delay_ms =
(base_delay * (1u64 << (attempt - 1))).min(Self::MAX_AGENT_RETRY_DELAY_MS);
let omitted_id = self.messages.last().and_then(|m| m.entry_id.clone());
self.capture_failed_route();
self.messages.pop();
self.omit_entry_from_context(omitted_id.as_deref());
self.emit(json!({
"type": "auto_retry_start",
"attempt": attempt,
"maxAttempts": max_retries,
"delayMs": delay_ms,
"errorMessage": err_msg,
}));
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
result = self.run_loop().await;
self.emit(json!({
"type": "auto_retry_end",
"attempt": attempt,
"success": result.is_ok(),
"finalError": result.as_ref().err().map(|e| e.to_string()),
}));
}
result
}
fn is_retryable_assistant_error(msg: &str) -> bool {
let m = msg.to_lowercase();
const NON_RETRYABLE: [&str; 10] = [
"gousagelimiterror",
"freeusagelimiterror",
"monthly usage limit reached",
"insufficient_quota",
"out of budget",
"quota exceeded",
"billing",
"rate limit reached",
"daily limit",
"subscription_sharing_usage_limit_exceeded",
];
if NON_RETRYABLE.iter().any(|p| m.contains(p)) {
return false;
}
const RETRYABLE: &[&str] = &[
"overloaded",
"server_busy",
"servers are currently busy",
"currently experiencing high demand",
"model is at capacity",
"rate limit",
"too many requests",
"429",
"500",
"502",
"503",
"504",
"520",
"524",
"service unavailable",
"server error",
"internal error",
"provider returned error",
"exceeded request buffer limit while retrying upstream",
"network",
"connection",
"connection refused",
"connection reset",
"connection lost",
"other side closed",
"error sending request",
"error trying to connect",
"request failed",
"fetch failed",
"getaddrinfo",
"enotfound",
"eai_again",
"upstream connect",
"reset before headers",
"socket hang up",
"timed out",
"timeout",
"terminated",
"websocket closed",
"websocket error",
"ended without",
"stream ended before message_stop",
"stream ended before a terminal response event",
"http2 request did not get a response",
"http2",
"retry delay",
"you can retry your request",
"try your request again",
"please retry your request",
"resourceexhausted",
];
RETRYABLE.iter().any(|p| m.contains(p))
}
async fn run_loop(&mut self) -> Result<String> {
match self.agent_loop.clone() {
Some(arc) => {
let f: &dyn Fn(&mut Agent) -> BoxFuture<'_, Result<String>> = arc.as_ref();
f(self).await
}
None => agent_loop::default_loop(self).await,
}
}
pub fn set_agent_loop(&mut self, f: AgentLoop) {
self.agent_loop = Some(f);
}
pub fn reset_agent_loop(&mut self) {
self.agent_loop = None;
}
pub(crate) fn emit_agent_end(&mut self, messages: &[AgentMessage]) {
let list: Vec<Value> = messages
.iter()
.map(|m| serde_json::to_value(m).unwrap())
.collect();
self.emit(json!({ "type": "agent_end", "messages": list }));
}
pub(crate) fn emit_agent_settled(&mut self) {
self.cache_warmer.on_agent_settled();
let aborted = self.run_abort.load(Ordering::Relaxed);
self.emit(json!({ "type": "agent_settled", "aborted": aborted }));
}
pub(crate) fn emit_abort_failure(&mut self, new_messages: &mut Vec<AgentMessage>) {
let (model_id, provider, api) = {
let m = self.effective_model();
(m.model_id.clone(), m.provider.clone(), m.api.clone())
};
let mut failure = AgentMessage {
role: "assistant".to_string(),
thinking_level: None,
content: vec![ContentBlock::Text {
text: String::new(),
text_signature: None,
}],
tool_call_id: None,
tool_name: None,
is_error: false,
stop_reason: Some("aborted".to_string()),
error_message: Some("Operation aborted".to_string()),
model: Some(model_id),
provider: Some(provider),
api: Some(api),
response_model: None,
response_id: None,
diagnostics: None,
usage: None,
deferred: None,
raw_stop_reason: None,
end_turn: None,
timestamp: now_ms(),
duration_ms: None,
details: None,
citations: None,
entry_id: None,
};
let v = serde_json::to_value(&failure).unwrap();
self.emit(json!({ "type": "message_start", "message": v.clone() }));
self.emit(json!({ "type": "message_end", "message": v.clone() }));
self.emit(json!({ "type": "turn_end", "message": v, "toolResults": [] }));
if let Some(s) = self.session.as_mut() {
let eid = s.append_message(&failure);
failure.entry_id = Some(eid);
}
new_messages.push(failure.clone());
self.messages.push(failure);
}
pub(crate) fn take_steering_batch(&mut self) -> Vec<AgentMessage> {
let mut out: Vec<AgentMessage> = Vec::new();
{
let mut q = self.runtime_steer_inbox.lock().unwrap();
if self.steer_mode == "all" {
out.extend(q.drain(..));
} else if let Some(m) = q.pop_front() {
out.push(m);
}
}
out
}
pub(crate) async fn execute_tool_calls(
&mut self,
tool_calls: &[ContentBlock],
run_id: &str,
assistant_entry_id: Option<&str>,
assistant_msg: &AgentMessage,
context: &[AgentMessage],
) -> (Vec<(ContentBlock, AgentMessage)>, bool) {
let prepared = Self::prepare_tool_calls(
tool_calls,
assistant_msg,
context,
&self.rebuild_ctx.injected_tools,
);
let sequential = self.tool_execution == ToolExecutionMode::Sequential
|| prepared.iter().any(|p| {
ToolIndex::resolve(&self.rebuild_ctx.injected_tools, &p.name).execution_mode()
== Some(ToolExecutionMode::Sequential)
});
let sink = self.json_sink.clone();
let abort = self.run_abort.clone();
let cwd = self.cwd.clone();
let child_tpl = Arc::new(self.snapshot_template());
core::extensions::dispatch_exec_ctx(&make_exec_ctx(
&cwd,
&child_tpl,
&abort,
&sink,
None,
self.session_branch_entries(),
));
let mut tool_results: Vec<(ContentBlock, AgentMessage)> = Vec::new();
let terminates = if sequential {
self.execute_tool_batch_sequential(
&prepared,
run_id,
assistant_entry_id,
sink,
abort,
child_tpl,
&mut tool_results,
)
.await
} else {
self.execute_tool_batch_parallel(
&prepared,
run_id,
assistant_entry_id,
sink,
abort,
cwd,
child_tpl,
&mut tool_results,
)
.await
};
let terminate = !terminates.is_empty() && terminates.iter().all(|b| *b);
(tool_results, terminate)
}
fn prepare_tool_calls(
tool_calls: &[ContentBlock],
assistant_msg: &AgentMessage,
context: &[AgentMessage],
injected: &[InjectedChildTool],
) -> Vec<Prepared> {
let mut prepared: Vec<Prepared> = Vec::new();
for tc in tool_calls {
let (id, name, arguments) = match tc {
ContentBlock::ToolCall {
id,
name,
arguments,
..
} => (id.clone(), name.clone(), arguments.clone()),
_ => continue,
};
let mut args = arguments.clone();
let mut immediate: Option<String> = None;
let mut immediate_terminate = false;
let tools_index = ToolIndex::resolve(injected, &name);
if !tools_index.exists() {
immediate = Some(format!("Tool {} not found", name));
} else {
args = tools_index.prepare_arguments(&args);
if let Some(parameters) = tools_index.parameters() {
match tools::validate_tool_arguments(&name, ¶meters, &args) {
Ok(validated) => args = validated,
Err(msg) => immediate = Some(msg),
}
}
if immediate.is_none() {
for ext in core::extensions::registered() {
if ext.hooks().contains(&ExtensionHook::BeforeToolCall) {
match ext.before_tool_call_ext(
&name,
&args,
Some(assistant_msg),
context,
&arguments,
) {
Ok(out) => {
if out.block {
let reason = out.reason.unwrap_or_else(|| {
"Tool execution was blocked".to_string()
});
immediate = Some(reason);
immediate_terminate = out.terminate;
break;
}
if let Some(new_args) = out.args {
args = new_args;
}
}
Err(e) => {
immediate = Some(e.0);
break;
}
}
}
}
}
}
prepared.push(Prepared {
call: ContentBlock::ToolCall {
id,
name: name.clone(),
arguments,
thought_signature: None,
namespace: None,
},
name,
args,
immediate,
immediate_terminate,
});
}
prepared
}
#[allow(clippy::too_many_arguments)]
async fn execute_tool_batch_sequential(
&mut self,
prepared: &[Prepared],
run_id: &str,
assistant_entry_id: Option<&str>,
sink: Option<Arc<Mutex<JsonSink>>>,
abort: Arc<AtomicBool>,
child_tpl: Arc<AgentTemplate>,
tool_results: &mut Vec<(ContentBlock, AgentMessage)>,
) -> Vec<bool> {
let mut terminates: Vec<bool> = Vec::new();
for (tool_seq, p) in prepared.iter().enumerate() {
if abort.load(Ordering::Relaxed) {
break; }
if tool_seq == usize::MAX {
}
emit_start(&sink, p);
let tool_start = Instant::now();
let (f, dur) = if let Some(msg) = &p.immediate {
(
FinalizedToolCall::error_result(
&p.call,
&p.name,
&p.args,
msg,
p.immediate_terminate,
),
None,
)
} else {
let on_chunk = Self::bash_chunk_callback(
self.json_sink.clone(),
&p.tool_call_id(),
&p.name,
&p.args,
);
let idx = ToolIndex::resolve(&self.rebuild_ctx.injected_tools, &p.name);
let mut exec_ctx: Option<ToolExecCtx> = None;
let result = if let Some(handler) = idx.injected_handler() {
handler(p.args.clone()) } else if idx.is_builtin() {
let supports_images = self.model.input.iter().any(|i| i == "image");
let bash_env = self.current_bash_env();
Self::execute_builtin_tool(
&self.cwd,
&p.name,
&p.args,
on_chunk,
supports_images,
bash_env,
self.model.image_resize,
)
.await
} else {
let entries = tool_hosts_scripts(&p.name)
.then(|| self.session_branch_entries())
.flatten();
let ctx = make_exec_ctx(
&self.cwd,
&child_tpl,
&abort,
&self.json_sink,
Some(&p.tool_call_id()),
entries,
);
exec_ctx = Some(ctx.clone());
self.execute_extension_tool_async(&p.name, &p.args, ctx)
.await
};
let mut f = match result {
Ok(t) => FinalizedToolCall::from_tool_result(&p.call, &p.name, &p.args, t),
Err(e) => {
FinalizedToolCall::error_result(&p.call, &p.name, &p.args, &e.0, false)
}
};
let dur = Some(tool_start.elapsed().as_millis() as u64);
f.duration_ms = dur;
attach_nested_calls(&mut f, exec_ctx.as_ref());
apply_after_tool_call(&mut f, &p.name, &p.args);
(f, dur)
};
emit_tool_execution_end(&sink, &f);
self.push_tool_result_message(&f, dur, tool_results);
self.record_tool_started(
run_id,
assistant_entry_id,
tool_seq,
&f,
tool_results.last(),
);
terminates.push(f.terminate);
}
terminates
}
#[allow(clippy::too_many_arguments)]
async fn execute_tool_batch_parallel(
&mut self,
prepared: &[Prepared],
run_id: &str,
assistant_entry_id: Option<&str>,
sink: Option<Arc<Mutex<JsonSink>>>,
abort: Arc<AtomicBool>,
cwd: String,
child_tpl: Arc<AgentTemplate>,
tool_results: &mut Vec<(ContentBlock, AgentMessage)>,
) -> Vec<bool> {
let mut locks: HashMap<String, Arc<tokio::sync::Mutex<()>>> = HashMap::new();
let mut slots: Vec<Option<(FinalizedToolCall, Option<u64>)>> = Vec::new();
let mut futs: Vec<BoxFuture<'static, (FinalizedToolCall, Option<u64>)>> = Vec::new();
let mut fut_slots: Vec<usize> = Vec::new();
for p in prepared {
if abort.load(Ordering::Relaxed) {
break;
}
emit_start(&sink, p);
if let Some(msg) = &p.immediate {
let f = FinalizedToolCall::error_result(
&p.call,
&p.name,
&p.args,
msg,
p.immediate_terminate,
);
emit_tool_execution_end(&sink, &f);
slots.push(Some((f, None)));
} else {
let lock: Option<Arc<tokio::sync::Mutex<()>>> = if (p.name == "edit"
|| p.name == "write")
&& let Some(path) = p.args.get("path").and_then(|v| v.as_str())
{
Some(
locks
.entry(Self::mutation_key(&cwd, path))
.or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
.clone(),
)
} else {
None
};
slots.push(None);
let entries = tool_hosts_scripts(&p.name)
.then(|| self.session_branch_entries())
.flatten();
let mut fut = execute_one(p, cwd.clone(), &sink, &abort, &child_tpl, entries);
if let Some(l) = lock {
fut = Box::pin(async move {
let _guard = l.lock().await;
fut.await
});
}
fut_slots.push(slots.len() - 1);
futs.push(fut);
}
}
let executed = futures_util::future::join_all(futs).await;
for (pos, f) in fut_slots.into_iter().zip(executed) {
slots[pos] = Some(f);
}
let mut terminates: Vec<bool> = Vec::new();
for (seq_idx, (f, dur)) in slots.into_iter().flatten().enumerate() {
self.push_tool_result_message(&f, dur, tool_results);
self.record_tool_started(run_id, assistant_entry_id, seq_idx, &f, tool_results.last());
terminates.push(f.terminate);
}
terminates
}
fn current_bash_env(&self) -> Option<tools::BashSessionEnv> {
let env = tools::BashSessionEnv {
session_id: self.session.as_ref().map(|s| s.session_id.clone()),
session_file: self
.session
.as_ref()
.and_then(|s| s.get_session_file())
.map(|p| p.to_string_lossy().to_string()),
provider: Some(self.model.provider.clone()),
model: Some(self.model.model_id.clone()),
reasoning_level: self.thinking_level.clone(),
expose: settings_manager::read_settings_expose_session_environment(),
};
Some(env)
}
fn record_tool_started(
&mut self,
run_id: &str,
assistant_entry_id: Option<&str>,
tool_index: usize,
fin: &FinalizedToolCall,
last_result: Option<&(ContentBlock, AgentMessage)>,
) {
let Some(s) = self.session.as_mut() else {
return;
};
let result_entry_id = last_result.and_then(|(_, m)| m.entry_id.clone());
s.append_lane_record(&session_v4::LaneRecord::ToolStarted {
base: session_v4::RecordBase {
id: session_manager::new_record_id("tool"),
seq: 0,
lane: String::new(),
timestamp: now_ms() as i64,
},
run_id: run_id.to_string(),
assistant_entry_id: assistant_entry_id.unwrap_or("").to_string(),
tool_index,
tool_call_id: fin.call_id(),
tool_name: fin.name.clone(),
effective_args: fin.args.clone(),
result_entry_id: result_entry_id.unwrap_or_default(),
replay: "never".to_string(),
});
}
pub(crate) fn record_step_attempt(
&mut self,
run_id: &str,
attempt: &mut u32,
result_entry_id: Option<&str>,
) {
let Some(result_entry_id) = result_entry_id else {
return;
};
if let Some(s) = self.session.as_mut() {
*attempt += 1;
s.append_lane_record(&session_v4::LaneRecord::StepAttempt {
base: session_v4::RecordBase {
id: session_manager::new_record_id("step"),
seq: 0,
lane: String::new(),
timestamp: now_ms() as i64,
},
run_id: run_id.to_string(),
step: "assistant".to_string(),
attempt: *attempt,
result_entry_id: result_entry_id.to_string(),
compaction_reason: None,
});
}
}
fn push_tool_result_message(
&mut self,
fin: &FinalizedToolCall,
duration_ms: Option<u64>,
tool_results: &mut Vec<(ContentBlock, AgentMessage)>,
) {
let content: Vec<ContentBlock> = if let Some(override_blocks) = &fin.content_override {
override_blocks.clone()
} else if fin.attachments.is_empty() {
vec![ContentBlock::Text {
text: fin.text.clone(),
text_signature: None,
}]
} else {
let limits = self.model.image_resize;
let mut blocks: Vec<ContentBlock> =
if settings_manager::read_settings_images_block_images() {
vec![ContentBlock::Text {
text:
"[Image omitted: images are blocked by settings (images.blockImages).]"
.to_string(),
text_signature: None,
}]
} else {
fin.attachments
.iter()
.map(|a| {
let (data, mime_type) =
normalize_base64_attachment(&a.data_base64, &a.mime_type, limits);
ContentBlock::Image { data, mime_type }
})
.collect()
};
blocks.push(ContentBlock::Text {
text: fin.text.clone(),
text_signature: None,
});
blocks
};
let result_msg = AgentMessage {
role: "toolResult".to_string(),
thinking_level: None,
content,
tool_call_id: Some(fin.call_id()),
tool_name: Some(fin.name.clone()),
is_error: fin.is_error,
stop_reason: None,
error_message: None,
model: None,
provider: None,
api: None,
response_model: None,
response_id: None,
diagnostics: None,
usage: fin.usage.clone(),
deferred: None,
raw_stop_reason: None,
end_turn: None,
timestamp: now_ms(),
duration_ms,
details: fin.details.clone(),
citations: None,
entry_id: None,
};
let result_msg = if let Some(s) = self.session.as_mut() {
let entry_id = s.append_message(&result_msg);
let mut rm = result_msg;
rm.entry_id = Some(entry_id);
rm
} else {
result_msg
};
let v = serde_json::to_value(&result_msg).unwrap();
self.emit(json!({ "type": "message_start", "message": v.clone() }));
self.emit(json!({ "type": "message_end", "message": v }));
tool_results.push((fin.call.clone(), result_msg));
}
fn mutation_key(cwd: &str, path: &str) -> String {
let p = tools::resolve_tool_path(path, cwd);
std::fs::canonicalize(&p)
.unwrap_or(p)
.to_string_lossy()
.into_owned()
}
fn bash_chunk_callback(
sink: Option<Arc<Mutex<JsonSink>>>,
tool_call_id: &str,
tool_name: &str,
args: &Value,
) -> Option<tools::OnChunkCallback> {
let sink = sink?;
let tool_call_id = tool_call_id.to_string();
let tool_name = tool_name.to_string();
let args = args.clone();
let last_emit = Arc::new(Mutex::new(Instant::now()));
Some(Arc::new(tokio::sync::Mutex::new(
Box::new(move |chunk: &str| {
let now = Instant::now();
let mut last = match last_emit.lock() {
Ok(g) => g,
Err(_) => return,
};
if now.duration_since(*last) < Duration::from_millis(100) {
return;
}
*last = now;
drop(last);
if let Ok(mut g) = sink.lock() {
let f: &mut dyn FnMut(Value) = g.as_mut();
f(json!({
"type": "tool_execution_update",
"toolCallId": tool_call_id,
"toolName": tool_name,
"args": args,
"partialResult": { "content": [ { "type": "text", "text": chunk } ] }
}));
}
}) as Box<dyn FnMut(&str) + Send>,
)))
}
fn is_overflow_error_str(msg: &str) -> bool {
let m = msg.to_lowercase();
const NON_OVERFLOW: [&str; 4] = [
"throttling",
"service unavailable",
"rate limit",
"too many requests",
];
if NON_OVERFLOW.iter().any(|p| m.contains(p)) {
return false;
}
const OVERFLOW: &[&str] = &[
"prompt is too long",
"prompt too long",
"prompt exceeds max length",
"request_too_large",
"input is too long for requested model",
"exceeds the context window",
"exceeds the maximum context length",
"maximum context length",
"exceeds the maximum allowed input length",
"input token count exceeds",
"exceeds the maximum number of tokens",
"maximum prompt length",
"reduce the length of the messages",
"maximum context length is",
"exceeds the limit of",
"exceeds the available context size",
"greater than the context length",
"context window exceeds limit",
"exceeded model token limit",
"too large for model",
"is longer than the model",
"configured context size",
"model_context_window_exceeded",
"range of input length",
"context length exceeded",
"context_length_exceeded",
"context_length",
"too many tokens",
"token limit exceeded",
"tokens exceeded",
];
OVERFLOW.iter().any(|p| m.contains(p))
}
fn is_bodyless_status_error(msg: &str) -> bool {
let lower = msg.to_lowercase();
let Some(rest) = lower.strip_prefix("provider returned ") else {
return false;
};
if !(rest.starts_with("400") || rest.starts_with("413")) {
return false;
}
msg.split_once(": ")
.is_some_and(|(_, body)| body.trim().is_empty())
}
fn message_is_context_overflow(
msg: &AgentMessage,
context_window: u32,
provider: &str,
) -> bool {
let stop = msg.stop_reason.as_deref();
let usage = msg.usage.as_ref();
let input_tokens = usage
.map(|u| u.input as u64 + u.cache_read as u64)
.unwrap_or(0);
if stop == Some("error") {
let Some(err) = msg.error_message.as_deref() else {
return false;
};
return Self::is_overflow_error_str(err)
|| (provider == "cerebras" && Self::is_bodyless_status_error(err));
}
if context_window == 0 {
return false;
}
if stop == Some("stop") && input_tokens > context_window as u64 {
return true;
}
if stop == Some("length")
&& usage.is_some_and(|u| u.output == 0)
&& input_tokens * 100 >= context_window as u64 * 99
{
return true;
}
false
}
pub fn child_agent(&self, custom_system_prompt: Option<String>) -> Agent {
let system_override = custom_system_prompt
.or_else(|| self.rebuild_ctx.system_prompt.clone())
.or_else(|| Some(self.system_prompt.clone()));
let mut tpl = self.snapshot_template();
if let Some(sp) = system_override {
tpl.system_prompt = sp;
}
let mut child = Agent::child_from_template(&tpl);
child.agent_loop = self.agent_loop.clone();
child
}
pub fn run_sub_agent<'a>(&'a self, user_message: &'a str) -> BoxFuture<'a, Result<String>> {
let mut child = self.child_agent(None);
Box::pin(async move { child.prompt(user_message).await })
}
pub fn spawn_sub_agent(&self, user_message: String) -> JoinHandle<Result<String>> {
let mut child = self.child_agent(None);
tokio::spawn(async move { child.prompt(&user_message).await })
}
fn is_builtin_tool(name: &str) -> bool {
tools::is_builtin_tool_name(name)
}
fn is_shell_tool(name: &str) -> bool {
matches!(name, "bash" | "powershell")
}
fn shell_output_details(out: &BashOutcome) -> Value {
json!({
"truncation": if out.truncated {
json!({
"truncated": true,
"reason": "length",
"fullOutputPath": out.full_output_path,
})
} else {
Value::Null
},
"fullOutputPath": out.full_output_path,
})
}
async fn execute_builtin_tool(
cwd: &str,
name: &str,
arguments: &Value,
on_bash_chunk: Option<tools::OnChunkCallback>,
model_supports_images: bool,
bash_env: Option<tools::BashSessionEnv>,
image_resize: ImageResizeLimits,
) -> std::result::Result<tools::ToolResult, tools::ToolError> {
match name {
"read" => {
let path = arguments
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("read tool is missing path argument".to_string())
})?;
let offset = arguments.get("offset").and_then(|v| v.as_i64());
let limit = arguments.get("limit").and_then(|v| v.as_i64());
let image_opts = tools::ReadImageOptions {
model_supports_images,
auto_resize: settings_manager::read_settings_images_auto_resize(),
block_images: settings_manager::read_settings_images_block_images(),
resize_limits: image_resize,
};
tools::execute_read_with_options(path, offset, limit, cwd, image_opts).await
}
"write" => {
let path = arguments
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("write tool is missing path argument".to_string())
})?;
let content = arguments
.get("content")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("write tool is missing content argument".to_string())
})?;
tools::execute_write(path, content, cwd).await
}
"edit" => {
let path = arguments
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("edit tool is missing path argument".to_string())
})?;
let edits = arguments.get("edits").ok_or_else(|| {
tools::ToolError("edit tool is missing edits argument".to_string())
})?;
let (result, _patch, _first_line) = tools::execute_edit(path, edits, cwd).await?;
Ok(result)
}
#[cfg(windows)]
"powershell" => {
let command = arguments
.get("command")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("powershell tool is missing command argument".to_string())
})?;
let timeout = arguments.get("timeout").and_then(|v| v.as_f64());
let out = tools::execute_powershell(command, timeout, cwd, on_bash_chunk).await;
if out.cancelled || out.timed_out {
Err(tools::ToolError(out.text))
} else {
let details = Self::shell_output_details(&out);
let mut result = if out.exit_code.is_some_and(|c| c != 0) {
tools::ToolResult::error(out.text)
} else {
tools::ToolResult::text(out.text)
};
result.details = Some(details);
result.structured_content = out.structured;
Ok(result)
}
}
"bash" => {
let command = arguments
.get("command")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("bash tool is missing command argument".to_string())
})?;
let timeout = arguments.get("timeout").and_then(|v| v.as_f64());
let id = session_manager::new_record_id("tool");
let out = match bash_env {
Some(env) => {
tools::execute_bash_with_env(
command,
timeout,
cwd,
Some(&id),
on_bash_chunk,
&env,
)
.await
}
None => {
tools::execute_bash_with(command, timeout, cwd, Some(&id), on_bash_chunk)
.await
}
};
if out.cancelled || out.timed_out {
Err(tools::ToolError(out.text))
} else {
let details = Self::shell_output_details(&out);
let mut result = if out.exit_code.is_some_and(|c| c != 0) {
tools::ToolResult::error(out.text)
} else {
tools::ToolResult::text(out.text)
};
result.details = Some(details);
result.structured_content = out.structured;
Ok(result)
}
}
"grep" => {
let pattern = arguments
.get("pattern")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("grep tool is missing pattern argument".to_string())
})?;
tools::execute_grep(
pattern,
arguments.get("path").and_then(|v| v.as_str()),
arguments.get("glob").and_then(|v| v.as_str()),
arguments
.get("ignoreCase")
.and_then(|v| v.as_bool())
.unwrap_or(false),
arguments
.get("literal")
.and_then(|v| v.as_bool())
.unwrap_or(false),
arguments.get("context").and_then(|v| v.as_i64()),
arguments.get("limit").and_then(|v| v.as_i64()),
cwd,
)
.await
}
"find" => {
let pattern = arguments
.get("pattern")
.and_then(|v| v.as_str())
.ok_or_else(|| {
tools::ToolError("find tool is missing pattern argument".to_string())
})?;
tools::execute_find(
pattern,
arguments.get("path").and_then(|v| v.as_str()),
arguments.get("limit").and_then(|v| v.as_i64()),
cwd,
)
.await
}
"ls" => {
tools::execute_ls(
arguments.get("path").and_then(|v| v.as_str()),
arguments.get("limit").and_then(|v| v.as_i64()),
cwd,
)
.await
}
_ => Err(tools::ToolError(format!("Unknown tool: {}", name))),
}
}
async fn execute_extension_tool_async(
&self,
name: &str,
args: &Value,
ctx: core::extensions::ToolExecCtx,
) -> std::result::Result<tools::ToolResult, tools::ToolError> {
for ext in core::extensions::registered() {
if ext.tools().iter().any(|t| t.name == name) {
return ext
.execute_tool_async(name.to_string(), args.clone(), ctx.clone())
.await;
}
}
Err(tools::ToolError(format!("Tool {} not found", name)))
}
}
struct Prepared {
call: ContentBlock,
name: String,
args: Value,
immediate: Option<String>, immediate_terminate: bool,
}
impl Prepared {
fn tool_call_id(&self) -> String {
match &self.call {
ContentBlock::ToolCall { id, .. } => id.clone(),
_ => String::new(),
}
}
}
fn emit_start(sink: &Option<Arc<Mutex<JsonSink>>>, p: &Prepared) {
if let Some(sink) = sink
&& let Ok(mut g) = sink.lock()
{
let f: &mut dyn FnMut(Value) = g.as_mut();
let args = serde_json::to_value(&p.args).unwrap_or(Value::Null);
f(json!({
"type": "tool_execution_start",
"toolCallId": p.tool_call_id(),
"toolName": p.name,
"args": args,
"startedAtEpochMs": now_ms()
}));
}
}
fn emit_tool_execution_end(sink: &Option<Arc<Mutex<JsonSink>>>, f: &FinalizedToolCall) {
if let Some(sink) = sink
&& let Ok(mut g) = sink.lock()
{
let fptr: &mut dyn FnMut(Value) = g.as_mut();
let mut payload = json!({
"type": "tool_execution_end",
"toolCallId": f.call_id(),
"toolName": f.name,
"result": f.to_result_json(),
"isError": f.is_error
});
if let Some(ms) = f.duration_ms
&& let Some(obj) = payload.as_object_mut()
{
obj.insert("durationMs".to_string(), json!(ms));
}
fptr(payload);
}
}
fn apply_after_tool_call(fin: &mut FinalizedToolCall, name: &str, args: &Value) {
for ext in core::extensions::registered() {
if ext.hooks().contains(&ExtensionHook::AfterToolCall) {
match ext.after_tool_call_ext(name, args, &fin.text, fin.is_error) {
Ok(out) => {
fin.text = out.text;
if let Some(v) = out.is_error {
fin.is_error = v;
}
if let Some(v) = out.usage {
fin.usage = Some(v);
}
if let Some(v) = out.terminate {
fin.terminate = v;
}
if let Some(content) = out.content {
fin.content_override = Some(content);
}
if let Some(details) = out.details {
fin.details = Some(details);
}
}
Err(e) => {
fin.text = e.0;
fin.is_error = true;
fin.usage = None;
fin.terminate = false;
fin.details = None;
fin.content_override = None;
}
}
}
}
}
struct MaterializedSubAgent {
agent: Agent,
abort: Arc<AtomicBool>,
session_path: Option<String>,
}
impl SubAgentRunner for MaterializedSubAgent {
fn controls(&self) -> SubAgentControls {
let inbox = self.agent.runtime_steer_inbox.clone();
SubAgentControls {
steer: Arc::new(move |text: String| {
if let Ok(mut q) = inbox.lock() {
q.push_back(AgentMessage::user_text(&text));
}
}),
abort: self.abort.clone(),
session_path: self.session_path.clone(),
}
}
fn run(&mut self, prompt: String) -> BoxFuture<'_, std::result::Result<String, String>> {
Box::pin(async move { self.agent.prompt(&prompt).await.map_err(|e| e.to_string()) })
}
}
fn materialize_sub_agent(
tpl: &AgentTemplate,
spec: SubAgentSpec,
) -> std::result::Result<MaterializedSubAgent, String> {
let mut child = Agent::child_from_template(tpl);
let abort = child.run_abort.clone();
if let Some(cwd) = spec.cwd.clone() {
child.cwd = cwd;
}
if spec.inherit_context {
child.messages = tpl.parent_messages.as_ref().clone();
}
if let Some(msgs) = spec.context_messages.clone() {
child.messages = msgs;
}
let session_path = if let Some(path) = spec.resume_session_path.as_deref() {
let session = Session::open(path)
.map_err(|e| format!("failed to resume sub-agent session {path}: {e}"))?;
child.messages = session.build_context_messages();
let path = session
.get_session_file()
.map(|p| p.to_string_lossy().to_string())
.or_else(|| Some(path.to_string()));
child.session = Some(session);
path
} else if spec.persist_session {
let parent_session_id = spec
.session_parent_id
.clone()
.or_else(|| tpl.parent_session_id.clone());
let session = Session::create_with(SessionCreateOptions {
cwd: child.cwd.clone(),
session_dir: spec.session_dir.clone(),
persist: true,
parent_session_id,
..Default::default()
})
.map_err(|e| format!("failed to create sub-agent session: {e}"))?;
let path = session
.session_file
.as_ref()
.map(|p| p.to_string_lossy().to_string());
child.session = Some(session);
path
} else {
None
};
if let Some(model) = spec.model.as_deref() {
let (provider, model_id) = model
.split_once('/')
.ok_or_else(|| format!("invalid model {model:?}: expected \"provider/modelId\""))?;
child
.switch_model(provider, model_id, false)
.map_err(|e| format!("model {model:?} is not available: {e}"))?;
}
if let Some(level) = spec.thinking.as_deref() {
child.set_thinking_level_silent(level, false);
}
if let Some(body) = spec.system_prompt.clone() {
child.rebuild_ctx.system_prompt = Some(body);
}
if let Some(body) = spec.append_system_prompt.clone() {
child.rebuild_ctx.append_system_prompt = Some(body);
}
if spec.clear_context_files {
child.rebuild_ctx.context_files.clear();
}
if spec.clear_skills {
child.skills.clear();
}
let universe = child.get_all_tools();
let mut selected: Vec<String> = match spec.tools.as_ref() {
None => universe,
Some(list) => {
if list.is_empty() || list.iter().any(|t| t == "none") {
Vec::new()
} else if list.iter().any(|t| t == "*" || t == "all") {
universe
} else {
universe
.into_iter()
.filter(|n| list.iter().any(|t| t == n))
.collect()
}
}
};
if spec.isolated || spec.extensions.is_some() || !spec.exclude_extensions.is_empty() {
let mut owner: HashMap<String, String> = HashMap::new();
for ext in core::extensions::registered() {
for t in ext.tools() {
owner.insert(t.name, ext.name().to_string());
}
}
if spec.isolated {
selected.retain(|n| !owner.contains_key(n));
}
if let Some(allow) = spec.extensions.as_ref() {
selected.retain(|n| match owner.get(n) {
None => true,
Some(ext) => allow.iter().any(|a| a == ext),
});
}
if !spec.exclude_extensions.is_empty() {
selected.retain(|n| match owner.get(n) {
None => true,
Some(ext) => !spec.exclude_extensions.iter().any(|d| d == ext),
});
}
}
if !spec.disallowed_tools.is_empty() {
selected.retain(|n| !spec.disallowed_tools.iter().any(|d| d == n));
}
let mut selected: Vec<String> = selected
.into_iter()
.filter(|n| {
!SUBAGENT_TOOL_NAMES.contains(&n.as_str())
|| spec.tools.as_ref().is_some_and(|l| l.contains(n))
})
.collect();
selected.dedup();
if spec.tools.as_ref().is_some_and(|l| !l.is_empty()) && selected.is_empty() {
return Err(format!(
"agent type {:?}: none of the configured tools are available",
spec.agent_type
));
}
child.rebuild_ctx.agent_id = spec.agent_id.clone();
child.rebuild_ctx.depth = spec.depth;
child.rebuild_ctx.selection = ToolSelection::Allowlist(selected);
for t in &spec.injected_tools {
let shadowed =
Agent::is_builtin_tool(&t.meta.name) || find_extension_tool(&t.meta.name).is_some();
if shadowed && note_warning_once(&format!("injected-tool-shadow|{}", t.meta.name)) {
core::extensions::request_ui(ExtensionUiRequest::Notify {
text: format!(
"subagent: injected tool {:?} shadows an existing tool inside this sub-agent",
t.meta.name
),
level: UiNotifyLevel::Warning,
});
}
}
child.rebuild_ctx.injected_tools = Arc::new(spec.injected_tools.clone());
child.rebuild_tools();
child.json_sink = spec.events.clone();
if let Some(max_turns) = spec.max_turns.filter(|m| *m > 0) {
let grace = spec.grace_turns.unwrap_or(SUBAGENT_GRACE_TURNS);
let counter = Arc::new(AtomicU32::new(0));
let inbox = child.runtime_steer_inbox.clone();
child.finish_turn = Some(Mutex::new(Box::new(move |ctx: &TurnContext| {
if matches!(
ctx.message.stop_reason.as_deref(),
Some("error") | Some("aborted")
) {
return None;
}
let turn = counter.fetch_add(1, Ordering::SeqCst) + 1;
if turn == max_turns
&& let Ok(mut q) = inbox.lock()
{
q.push_back(AgentMessage::user_text(SUBAGENT_WRAP_UP));
}
(turn >= max_turns.saturating_add(grace)).then_some(FinishTurnAction::End)
})));
}
Ok(MaterializedSubAgent {
agent: child,
abort,
session_path,
})
}
fn make_exec_ctx(
cwd: &str,
tpl: &Arc<AgentTemplate>,
parent_abort: &Arc<AtomicBool>,
sink: &Option<Arc<Mutex<JsonSink>>>,
parent_tool_call_id: Option<&str>,
branch_entries: Option<Arc<Vec<Value>>>,
) -> ToolExecCtx {
build_exec_ctx(
cwd,
tpl,
parent_abort,
sink,
parent_tool_call_id.map(str::to_string),
Arc::new(Mutex::new(NestedCallLog::default())),
0,
false,
branch_entries,
)
}
#[allow(clippy::too_many_arguments)]
fn build_exec_ctx(
cwd: &str,
tpl: &Arc<AgentTemplate>,
parent_abort: &Arc<AtomicBool>,
sink: &Option<Arc<Mutex<JsonSink>>>,
parent_tool_call_id: Option<String>,
log: Arc<Mutex<NestedCallLog>>,
nested_depth: u32,
script_call: bool,
branch_entries: Option<Arc<Vec<Value>>>,
) -> ToolExecCtx {
let make_sub_agent = {
let tpl = tpl.clone();
Arc::new(move |spec: SubAgentSpec| {
materialize_sub_agent(&tpl, spec).map(|r| Box::new(r) as Box<dyn SubAgentRunner>)
}) as Arc<MakeSubAgentFn>
};
let execute_tool: ExecuteToolFn = {
let cwd = cwd.to_string();
let tpl = tpl.clone();
let abort = parent_abort.clone();
let sink = sink.clone();
let log = log.clone();
let parent = parent_tool_call_id.clone();
let entries = branch_entries.clone();
Arc::new(move |name: String, args: Value| {
let cwd = cwd.clone();
let tpl = tpl.clone();
let abort = abort.clone();
let sink = sink.clone();
let log = log.clone();
let parent = parent.clone();
let entries = entries.clone();
Box::pin(async move {
run_nested_tool(
&cwd,
&tpl,
&abort,
&sink,
&parent,
&name,
args,
nested_depth,
&log,
entries,
)
.await
}) as BoxFuture<'static, std::result::Result<ToolResult, ToolError>>
})
};
ToolExecCtx {
cwd: cwd.to_string(),
make_sub_agent,
parent_abort: parent_abort.clone(),
agent_id: tpl.rebuild_ctx.agent_id.clone(),
depth: tpl.rebuild_ctx.depth,
parent_model: Some(format!("{}/{}", tpl.model.provider, tpl.model.model_id)),
execute_tool,
parent_tool_call_id,
nested_calls: log,
session_branch_entries: branch_entries,
script_tools: Arc::new(script_tools_for(
&tpl.tools,
&tpl.rebuild_ctx.injected_tools,
&tpl.rebuild_ctx.selection,
&tpl.rebuild_ctx.exclude_tools,
)),
script_call,
}
}
fn tool_hosts_scripts(name: &str) -> bool {
core::extensions::registered()
.iter()
.any(|ext| ext.hosts_scripts() && ext.tools().iter().any(|t| t.name == name))
}
#[allow(clippy::too_many_arguments)]
async fn run_nested_tool(
cwd: &str,
tpl: &Arc<AgentTemplate>,
abort: &Arc<AtomicBool>,
sink: &Option<Arc<Mutex<JsonSink>>>,
parent_call_id: &Option<String>,
name: &str,
args: Value,
nested_depth: u32,
log: &Arc<Mutex<NestedCallLog>>,
branch_entries: Option<Arc<Vec<Value>>>,
) -> std::result::Result<ToolResult, ToolError> {
let refusal: Option<String> = if nested_depth >= MAX_NESTED_TOOL_DEPTH {
Some(format!(
"nested tool call depth limit reached ({MAX_NESTED_TOOL_DEPTH}); refusing to call {name}"
))
} else if abort.load(Ordering::Relaxed) {
Some("Operation aborted".to_string())
} else {
let idx = ToolIndex::resolve(&tpl.rebuild_ctx.injected_tools, name);
if !idx.exists() {
Some(format!("Tool {name} not found"))
} else {
match find_extension_tool(name) {
Some(tool)
if !tool
.exposure
.is_nested_callable(is_deferred_tool_activated(name)) =>
{
Some(format!(
"Tool {name} cannot be called from another tool (exposure: {})",
tool.exposure.as_str()
))
}
_ => None,
}
}
};
let idx = ToolIndex::resolve(&tpl.rebuild_ctx.injected_tools, name);
let args = match refusal {
Some(_) => args,
None => match idx.parameters() {
Some(schema) => match tools::validate_tool_arguments(name, &schema, &args) {
Ok(valid) => valid,
Err(msg) => {
return refuse_nested_tool(sink, parent_call_id, name, args, msg, log);
}
},
None => args,
},
};
if let Some(reason) = refusal {
return refuse_nested_tool(sink, parent_call_id, name, args, reason, log);
}
let call_id = new_nested_call_id(parent_call_id);
let started = Instant::now();
emit_nested_start(sink, &call_id, parent_call_id.as_deref(), name, &args);
let child_log = Arc::new(Mutex::new(NestedCallLog::default()));
let result = if let Some(handler) = idx.injected_handler() {
handler(args.clone())
} else if idx.is_builtin() {
let supports_images = tpl.model.input.iter().any(|i| i == "image");
Agent::execute_builtin_tool(
cwd,
name,
&args,
None,
supports_images,
None,
tpl.model.image_resize,
)
.await
} else {
let child_ctx = build_exec_ctx(
cwd,
tpl,
abort,
sink,
Some(call_id.clone()),
child_log.clone(),
nested_depth + 1,
tool_hosts_scripts(name),
branch_entries.clone(),
);
let mut found = None;
for ext in core::extensions::registered() {
if ext.tools().iter().any(|t| t.name == name) {
found = Some(
ext.execute_tool_async(name.to_string(), args.clone(), child_ctx.clone())
.await,
);
break;
}
}
found.unwrap_or_else(|| Err(ToolError(format!("Tool {name} not found"))))
};
let duration_ms = started.elapsed().as_millis() as u64;
let (is_error, text, structured, usage) = match &result {
Ok(t) => (
t.is_error,
t.text.clone(),
t.structured_content.clone(),
t.usage.clone(),
),
Err(e) => (true, e.0.clone(), None, None),
};
let children = take_log(&child_log);
let mut record = serde_json::Map::new();
record.insert("toolCallId".into(), json!(call_id));
record.insert("toolName".into(), json!(name));
record.insert("args".into(), nested_args_value(&args));
record.insert("isError".into(), json!(is_error));
record.insert(
"text".into(),
json!(truncate_chars(&text, NESTED_CALL_TEXT_CAP)),
);
record.insert("durationMs".into(), json!(duration_ms));
if let Some(s) = &structured {
record.insert("structuredContent".into(), s.clone());
}
if let Some(u) = &usage {
record.insert(
"usage".into(),
serde_json::to_value(u).unwrap_or(Value::Null),
);
}
if !children.calls.is_empty() {
record.insert("nestedCalls".into(), Value::Array(children.calls.clone()));
}
if children.dropped > 0 {
record.insert("nestedCallsDropped".into(), json!(children.dropped));
}
emit_nested_end(sink, &call_id, parent_call_id.as_deref(), &record);
let mut log = log.lock().unwrap();
if log.calls.len() < MAX_NESTED_CALLS {
log.calls.push(Value::Object(record));
} else {
log.dropped = log.dropped.saturating_add(1);
}
if let Some(own) = usage {
log.usage = Some(match log.usage.take() {
Some(prev) => combine_usage(&prev, &own),
None => own,
});
}
if let Some(child) = children.usage {
log.usage = Some(match log.usage.take() {
Some(prev) => combine_usage(&prev, &child),
None => child,
});
}
result
}
fn new_nested_call_id(parent_call_id: &Option<String>) -> String {
format!(
"nested_{}_{}",
parent_call_id.as_deref().unwrap_or("root"),
nested_call_seq()
)
}
fn nested_args_value(args: &Value) -> Value {
let json = args.to_string();
if json.chars().count() <= NESTED_CALL_TEXT_CAP {
args.clone()
} else {
json!({
"_truncated": truncate_chars(&json, NESTED_CALL_TEXT_CAP)
})
}
}
fn refuse_nested_tool(
sink: &Option<Arc<Mutex<JsonSink>>>,
parent_call_id: &Option<String>,
name: &str,
args: Value,
reason: String,
log: &Arc<Mutex<NestedCallLog>>,
) -> std::result::Result<ToolResult, ToolError> {
let call_id = new_nested_call_id(parent_call_id);
emit_nested_start(sink, &call_id, parent_call_id.as_deref(), name, &args);
let record = json!({
"toolCallId": call_id,
"toolName": name,
"args": nested_args_value(&args),
"isError": true,
"text": reason,
"durationMs": 0,
});
if let Some(obj) = record.as_object() {
emit_nested_end(sink, &call_id, parent_call_id.as_deref(), obj);
}
let mut log = log.lock().unwrap();
if log.calls.len() < MAX_NESTED_CALLS {
log.calls.push(record);
} else {
log.dropped = log.dropped.saturating_add(1);
}
Err(ToolError(reason))
}
fn nested_call_seq() -> u64 {
static SEQ: AtomicU32 = AtomicU32::new(0);
SEQ.fetch_add(1, Ordering::Relaxed) as u64
}
fn take_log(log: &Arc<Mutex<NestedCallLog>>) -> NestedCallLog {
std::mem::take(&mut *log.lock().unwrap())
}
fn emit_nested_start(
sink: &Option<Arc<Mutex<JsonSink>>>,
call_id: &str,
parent: Option<&str>,
name: &str,
args: &Value,
) {
if let Some(sink) = sink
&& let Ok(mut g) = sink.lock()
{
let f: &mut dyn FnMut(Value) = g.as_mut();
f(json!({
"type": "tool_execution_start",
"toolCallId": call_id,
"toolName": name,
"parentToolCallId": parent,
"args": args,
"startedAtEpochMs": now_ms()
}));
}
}
fn emit_nested_end(
sink: &Option<Arc<Mutex<JsonSink>>>,
call_id: &str,
parent: Option<&str>,
record: &serde_json::Map<String, Value>,
) {
if let Some(sink) = sink
&& let Ok(mut g) = sink.lock()
{
let f: &mut dyn FnMut(Value) = g.as_mut();
f(json!({
"type": "tool_execution_end",
"toolCallId": call_id,
"toolName": record.get("toolName"),
"parentToolCallId": parent,
"isError": record.get("isError").and_then(|v| v.as_bool()).unwrap_or(false),
"result": {
"content": [{ "type": "text", "text": record.get("text").cloned().unwrap_or(Value::Null) }],
"details": { "nestedCall": Value::Object(record.clone()) },
},
}));
}
}
fn attach_nested_calls(fin: &mut FinalizedToolCall, ctx: Option<&ToolExecCtx>) {
let Some(ctx) = ctx else {
return;
};
let log = ctx.take_nested_calls();
if let Some(u) = log.usage {
fin.usage = Some(match fin.usage.take() {
Some(prev) => combine_usage(&prev, &u),
None => u,
});
}
if log.calls.is_empty() && log.dropped == 0 {
return;
}
let mut details = match fin.details.take() {
Some(Value::Object(o)) => o,
_ => serde_json::Map::new(),
};
details.insert("nestedCalls".into(), Value::Array(log.calls));
if log.dropped > 0 {
details.insert("nestedCallsDropped".into(), json!(log.dropped));
}
fin.details = Some(Value::Object(details));
}
fn execute_one(
prepared: &Prepared,
cwd: String,
sink: &Option<Arc<Mutex<JsonSink>>>,
abort: &Arc<AtomicBool>,
tpl: &Arc<AgentTemplate>,
branch_entries: Option<Arc<Vec<Value>>>,
) -> BoxFuture<'static, (FinalizedToolCall, Option<u64>)> {
let call = prepared.call.clone();
let name = prepared.name.clone();
let args = prepared.args.clone();
let abort = abort.clone();
let sink = sink.clone();
let tpl = tpl.clone();
let on_chunk = Agent::bash_chunk_callback(sink.clone(), &prepared.tool_call_id(), &name, &args);
let parent_tool_call_id = prepared.tool_call_id();
Box::pin(async move {
if abort.load(Ordering::Relaxed) {
let f =
FinalizedToolCall::error_result(&call, &name, &args, "Operation aborted", false);
emit_tool_execution_end(&sink, &f);
return (f, None);
}
let tool_start = Instant::now();
let idx = ToolIndex::resolve(&tpl.rebuild_ctx.injected_tools, &name);
let mut exec_ctx: Option<ToolExecCtx> = None;
let result = if let Some(handler) = idx.injected_handler() {
handler(args.clone()) } else if idx.is_builtin() {
let model_supports_images = tpl.model.input.iter().any(|i| i == "image");
Agent::execute_builtin_tool(
&cwd,
&name,
&args,
on_chunk,
model_supports_images,
None,
tpl.model.image_resize,
)
.await
} else {
let ctx = make_exec_ctx(
&cwd,
&tpl,
&abort,
&sink,
Some(&parent_tool_call_id),
branch_entries.clone(),
);
exec_ctx = Some(ctx.clone());
let mut found = None;
for ext in core::extensions::registered() {
if ext.tools().iter().any(|t| t.name == name) {
found = Some(
ext.execute_tool_async(name.clone(), args.clone(), ctx)
.await,
);
break;
}
}
found.unwrap_or_else(|| Err(tools::ToolError(format!("Tool {} not found", name))))
};
let mut f = match result {
Ok(t) => FinalizedToolCall::from_tool_result(&call, &name, &args, t),
Err(e) => FinalizedToolCall::error_result(&call, &name, &args, &e.0, false),
};
let call_id = match &call {
ContentBlock::ToolCall { id, .. } => id.clone(),
_ => String::new(),
};
if Agent::is_shell_tool(&name)
&& !f.is_error
&& let Some(details) = f.details.as_ref()
&& let Some(g) = &sink
&& let Ok(mut g) = g.lock()
{
let cb: &mut dyn FnMut(Value) = g.as_mut();
cb(json!({
"type": "tool_execution_update",
"toolCallId": call_id,
"toolName": name,
"args": args,
"partialResult": { "content": [ { "type": "text", "text": f.text } ] },
"details": details,
}));
}
let dur = Some(tool_start.elapsed().as_millis() as u64);
f.duration_ms = dur;
attach_nested_calls(&mut f, exec_ctx.as_ref());
apply_after_tool_call(&mut f, &name, &args);
emit_tool_execution_end(&sink, &f);
(f, dur)
})
}
pub(crate) fn stream_event_to_json(ev: StreamEvent, partial: Option<&AgentMessage>) -> Value {
match ev {
StreamEvent::Start { .. } => json!({ "type": "start" }),
StreamEvent::TextStart { content_index, .. } => {
json!({ "type": "text_start", "contentIndex": content_index })
}
StreamEvent::TextDelta {
content_index,
delta,
..
} => {
json!({ "type": "text_delta", "contentIndex": content_index, "delta": delta })
}
StreamEvent::TextEnd {
content_index,
content,
..
} => {
json!({ "type": "text_end", "contentIndex": content_index, "content": content })
}
StreamEvent::ThinkingStart { content_index, .. } => {
json!({ "type": "thinking_start", "contentIndex": content_index })
}
StreamEvent::ThinkingDelta {
content_index,
delta,
..
} => {
json!({ "type": "thinking_delta", "contentIndex": content_index, "delta": delta })
}
StreamEvent::ThinkingEnd {
content_index,
content,
..
} => {
json!({ "type": "thinking_end", "contentIndex": content_index, "content": content })
}
StreamEvent::ToolCallStart { content_index, .. } => {
let (call_id, tool_name) = partial
.and_then(|p| p.content.get(content_index))
.and_then(|b| match b {
ContentBlock::ToolCall { id, name, .. } => Some((id.clone(), name.clone())),
_ => None,
})
.unwrap_or_default();
let mut v = serde_json::Map::new();
v.insert("type".into(), json!("toolcall_start"));
v.insert("contentIndex".into(), json!(content_index));
v.insert("id".into(), json!(call_id));
v.insert("toolName".into(), json!(tool_name));
Value::Object(v)
}
StreamEvent::ToolCallDelta {
content_index,
delta,
..
} => {
json!({ "type": "toolcall_delta", "contentIndex": content_index, "delta": delta })
}
StreamEvent::ToolCallEnd {
content_index,
tool_call,
..
} => {
let tool_call = serde_json::to_value(tool_call).unwrap_or(Value::Null);
json!({ "type": "toolcall_end", "contentIndex": content_index, "toolCall": tool_call })
}
}
}
fn branch_summary_max_tokens(model: &ModelConfig) -> u32 {
model
.max_tokens
.filter(|m| *m > 0)
.unwrap_or(u32::MAX)
.min(4096)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::extensions::ToolExposure;
#[tokio::test]
async fn retryable_error_classification_matches_pi() {
assert!(Agent::is_retryable_assistant_error(
"503 service unavailable, retrying..."
));
assert!(Agent::is_retryable_assistant_error(
"connection reset by peer (network error)"
));
assert!(Agent::is_retryable_assistant_error(
"stream ended without response"
));
assert!(Agent::is_retryable_assistant_error(
"overloaded 429 too many requests"
));
assert!(Agent::is_retryable_assistant_error(
"provider returned 429 Too Many Requests: Selected model is at capacity. Please try again later."
));
assert!(Agent::is_retryable_assistant_error(
"provider returned error: server_busy"
));
assert!(Agent::is_retryable_assistant_error(
"The servers are currently busy, please retry"
));
assert!(Agent::is_retryable_assistant_error(
"provider returned 520: 520 status code (no body)"
));
assert!(Agent::is_retryable_assistant_error(
"The system is currently experiencing high demand and cannot process your request."
));
assert!(Agent::is_retryable_assistant_error(
"Error: request failed: error sending request for url (https://opencode.ai/zen/go/v1/chat/completions)"
));
assert!(Agent::is_retryable_assistant_error(
"request failed: error sending request for url (https://x): error trying to connect: tcp connect error: Connection refused (os error 111)"
));
assert!(!Agent::is_retryable_assistant_error(
"Monthly usage limit reached for free tier"
));
assert!(!Agent::is_retryable_assistant_error(
"insufficient_quota: out of budget"
));
assert!(!Agent::is_retryable_assistant_error(
"GoUsageLimitError: quota exceeded for billing"
));
assert!(!Agent::is_retryable_assistant_error(
"provider returned 429 Too Many Requests: {\"error\":{\"code\":\"subscription_sharing_usage_limit_exceeded\"}}"
));
assert!(!Agent::is_retryable_assistant_error("random model error"));
}
use crate::core::provider::Citation;
use serde_json::{Value as J, json};
use tokio::{
io::{AsyncReadExt as _, AsyncWriteExt as _},
net::TcpListener,
};
fn subagent_test_lock() -> (SubagentAuthGuard, crate::extensions::subagent::TestLock) {
let auth = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let lock = crate::extensions::subagent::test_lock();
(SubagentAuthGuard(auth), lock)
}
struct SubagentAuthGuard(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>);
async fn read_request(socket: &mut tokio::net::TcpStream) {
let mut buf = Vec::new();
let mut tmp = [0u8; 8192];
let header_end = loop {
let n = socket.read(&mut tmp).await.unwrap();
if n == 0 {
break buf.len();
}
buf.extend_from_slice(&tmp[..n]);
if let Some(pos) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
break pos + 4;
}
};
let Some(len) = content_length(&buf[..header_end.min(buf.len())]) else {
return;
};
while buf.len() < header_end + len {
let n = socket.read(&mut tmp).await.unwrap();
if n == 0 {
break;
}
buf.extend_from_slice(&tmp[..n]);
}
}
fn content_length(headers: &[u8]) -> Option<usize> {
let text = String::from_utf8_lossy(headers).to_ascii_lowercase();
let rest = &text[text.find("content-length:")? + "content-length:".len()..];
let end = rest.find(['\r', '\n']).unwrap_or(rest.len());
rest[..end].trim().parse().ok()
}
async fn mock_server_multi(
responses: Vec<Vec<&'static str>>,
) -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let handle = tokio::spawn(async move {
for frames in responses {
let (mut socket, _) = match tokio::time::timeout(
std::time::Duration::from_millis(5000),
listener.accept(),
)
.await
{
Ok(Ok(pair)) => pair,
_ => break,
};
read_request(&mut socket).await;
socket
.write_all(
b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n",
)
.await
.unwrap();
for f in frames {
let payload = format!("data: {}\n\n", f);
socket.write_all(payload.as_bytes()).await.unwrap();
}
socket.shutdown().await.unwrap();
}
});
(format!("http://127.0.0.1:{}", addr.port()), handle)
}
async fn mock_json_once(body: String) -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let handle = tokio::spawn(async move {
let Ok(Ok((mut socket, _))) =
tokio::time::timeout(std::time::Duration::from_millis(5_000), listener.accept())
.await
else {
return;
};
read_request(&mut socket).await;
let response = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
);
socket.write_all(response.as_bytes()).await.unwrap();
socket.shutdown().await.unwrap();
});
(format!("http://127.0.0.1:{}", addr.port()), handle)
}
enum MockResp {
Err(&'static str, &'static str),
Sse(Vec<&'static str>),
}
async fn mock_server_sequence(
responses: Vec<MockResp>,
) -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let handle = tokio::spawn(async move {
for resp in responses {
let (mut socket, _) = match tokio::time::timeout(
std::time::Duration::from_millis(5000),
listener.accept(),
)
.await
{
Ok(Ok(pair)) => pair,
_ => break,
};
read_request(&mut socket).await;
match resp {
MockResp::Err(status, body) => {
let out = format!(
"{status}\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n{body}"
);
socket.write_all(out.as_bytes()).await.unwrap();
}
MockResp::Sse(frames) => {
socket
.write_all(
b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n",
)
.await
.unwrap();
for f in frames {
let payload = format!("data: {}\n\n", f);
socket.write_all(payload.as_bytes()).await.unwrap();
}
}
}
socket.shutdown().await.unwrap();
}
});
(format!("http://127.0.0.1:{}", addr.port()), handle)
}
fn frame(s: String) -> &'static str {
Box::leak(s.into_boxed_str())
}
fn text_frame(text: &str, _id: &str) -> Vec<&'static str> {
vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}"#,
frame(format!(
r#"{{"id":"cmpl_t","choices":[{{"index":0,"delta":{{"content":"{}"}}}}]}}"#,
text
)),
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}"#,
r#"[DONE]"#,
]
}
fn tool_call_frame(tool_name: &str, args: &str, id: &str) -> Vec<&'static str> {
vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
frame(format!(
r#"{{"id":"cmpl_t","choices":[{{"index":0,"delta":{{"tool_calls":[{{"index":0,"id":"{}","type":"function","function":{{"name":"{}","arguments":"{}"}}}}]}}}}]}}"#,
id, tool_name, args
)),
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
]
}
#[test]
fn parameters_with_grammar_sampling_annotates_schema() {
use crate::core::extensions::Extension as _;
let tool = crate::extensions::codemode::Codemode
.tools()
.into_iter()
.find(|t| t.name == "codemode")
.expect("codemode 工具");
let annotated =
parameters_with_grammar_sampling(&tool.parameters, tool.grammar_sampling.as_ref());
assert_eq!(annotated["x-prux-grammar"]["input"], "code");
assert!(
annotated["x-prux-grammar"]["openai_lark"]
.as_str()
.unwrap_or_default()
.contains("options_source")
);
let spec = crate::core::provider::grammar_tool(&annotated).expect("语法约束可解析");
assert_eq!(spec.syntax, "lark");
assert_eq!(spec.input, "code");
assert!(
parameters_with_grammar_sampling(&tool.parameters, None)
.get("x-prux-grammar")
.is_none()
);
}
#[test]
fn branch_summary_max_tokens_caps_at_4096_and_model() {
let base = ModelConfig {
model_type: Default::default(),
image_resize: crate::utils::image::ImageResizeLimits::default(),
allowed_fallback_models: Vec::new(),
provider: "test".into(),
model_id: "m".into(),
base_url: "https://example.com/v1".into(),
api_key: String::new(),
api: "openai-completions".into(),
output: Vec::new(),
input: vec!["text".into()],
reasoning: true,
max_tokens: None,
temperature: None,
context_window: 200000,
thinking_format: String::new(),
supports_reasoning_effort: false,
thinking_level_map: None,
auth_header: false,
supports_developer_role: false,
requires_reasoning_content_on_assistant_messages: false,
supports_usage_in_streaming: true,
supports_store: true,
supports_finish_reason: true,
requires_assistant_after_tool_result: false,
max_tokens_field: "max_tokens".into(),
supports_strict_mode: false,
send_session_affinity_headers: None,
session_affinity_format: None,
cost: None,
session_id: None,
max_retry_delay_ms: None,
sampling_params: None,
sampling_params_by_thinking_level: None,
provider_routing: None,
thinking_budgets: None,
};
assert_eq!(branch_summary_max_tokens(&base), 4096);
let mut small = base.clone();
small.max_tokens = Some(1024);
assert_eq!(branch_summary_max_tokens(&small), 1024);
let mut zero = base.clone();
zero.max_tokens = Some(0);
assert_eq!(branch_summary_max_tokens(&zero), 4096);
let mut big = base.clone();
big.max_tokens = Some(60000);
assert_eq!(branch_summary_max_tokens(&big), 4096);
}
#[test]
fn deferred_handle_serializes_roundtrip() {
let h = crate::core::provider::DeferredHandle {
provider: "test".into(),
model_id: "m".into(),
api: "openai-responses".into(),
id: "resp_1|call_2".into(),
..Default::default()
};
let v = serde_json::to_value(&h).unwrap();
assert_eq!(v["id"], "resp_1|call_2");
assert_eq!(v["modelId"], "m");
}
#[test]
fn on_response_callback_receives_llm_result_summary() {
let rt = test_rt();
let frames1 = vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"content":"hello world"}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":null}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}"#,
r#"[DONE]"#,
];
let (base, handle) = rt.block_on(mock_server_multi(vec![frames1]));
let mut agent = test_agent(&base, ".");
let got: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
let got2 = got.clone();
agent.set_on_response(Box::new(move |summary: &J| {
*got2.lock().unwrap() = Some(summary["text"].as_str().unwrap_or("").to_string());
}));
let _ = rt.block_on(agent.prompt("hi"));
let _ = rt.block_on(handle);
let text = (*got.lock().unwrap()).clone().unwrap_or_default();
assert!(
text.contains("hello world"),
"on_response summary: {text:?}"
);
}
#[test]
fn citations_serialize_roundtrip() {
let c = Citation {
kind: "citation".into(),
start_index: Some(1),
end_index: Some(5),
url: Some("https://example.com".into()),
title: Some("Doc".into()),
text: Some("cited".into()),
};
let mut m = AgentMessage::user_text("");
m.role = "assistant".into();
m.citations = Some(vec![c]);
let v = serde_json::to_value(&m).unwrap();
assert_eq!(v["citations"][0]["url"], "https://example.com");
assert_eq!(v["citations"][0]["title"], "Doc");
let m2 = AgentMessage::user_text("x");
let v2 = serde_json::to_value(&m2).unwrap();
assert!(v2.get("citations").is_none(), "{v2}");
}
#[test]
fn request_abort_is_cooperative_and_emits_aborted_failure() {
let mut agent = test_agent("http://127.0.0.1:1", ".");
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
agent.request_abort();
assert!(agent.run_abort.load(Ordering::Relaxed));
let _ = collected;
}
#[test]
fn agent_settled_reports_whether_the_run_was_aborted() {
let mut agent = test_agent("http://127.0.0.1:1", ".");
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
agent.emit_agent_settled();
agent.request_abort();
agent.emit_agent_settled();
let evs = collected.lock().unwrap();
assert_eq!(evs.len(), 2);
assert_eq!(evs[0]["type"], "agent_settled");
assert_eq!(evs[0]["aborted"], false, "未取消的整轮: {}", evs[0]);
assert_eq!(evs[1]["aborted"], true, "已取消的整轮: {}", evs[1]);
}
#[test]
fn reset_run_state_resets_without_reallocating_abort_signal() {
let mut agent = test_agent("http://127.0.0.1:1", ".");
let captured = agent.run_abort.clone();
agent.run_abort.store(true, Ordering::Relaxed);
agent.reset_run_state();
assert!(!captured.load(Ordering::Relaxed), "reset 必须复位同一信号");
assert!(
Arc::ptr_eq(&captured, &agent.run_abort),
"Arc 身份必须保持,否则旧句柄永远停不下新一轮"
);
}
#[test]
fn abort_and_detach_signals_old_signal_and_installs_fresh_one() {
let mut agent = test_agent("http://127.0.0.1:1", ".");
let round = agent.run_abort.clone();
let old = agent.abort_and_detach();
assert!(round.load(Ordering::Relaxed), "本轮信号必须置位");
assert!(
old.load(Ordering::Relaxed),
"返回的旧信号即已被置位的那一个"
);
assert!(
!agent.run_abort.load(Ordering::Relaxed),
"换上的新信号必须是未置位"
);
assert!(
!agent.exec_ctx().parent_abort.load(Ordering::Relaxed),
"刷新后的执行上下文不得再带已取消标记"
);
}
#[test]
fn compaction_cut_point_never_splits_assistant_turn() {
let mut msgs: Vec<AgentMessage> = Vec::new();
msgs.push(AgentMessage::user_text("q1"));
let mut a1 = AgentMessage::user_text("");
a1.role = "assistant".into();
a1.content = vec![ContentBlock::Text {
text: "r1".into(),
text_signature: None,
}];
msgs.push(a1);
msgs.push(AgentMessage::user_text("q2"));
let mut a2 = AgentMessage::user_text("");
a2.role = "assistant".into();
a2.content = vec![ContentBlock::Text {
text: "r2".into(),
text_signature: None,
}];
msgs.push(a2);
let mut agent = test_agent("http://127.0.0.1:1", ".");
agent.messages = msgs;
agent.auto_compaction = false;
assert_eq!(agent.messages.len(), 4);
}
#[test]
fn threshold_compaction_skipped_when_usage_anchor_predates_compaction() {
let dir = tempfile::tempdir().unwrap();
let mut agent = test_agent("http://127.0.0.1:1", ".");
let mut sess = Session::create(".", Some(dir.path().to_path_buf()), false).unwrap();
let mut old = AgentMessage::user_text("");
old.role = "assistant".into();
old.timestamp = 1_000;
old.usage = Some(Usage {
total_tokens: 5_000_000,
..Usage::default()
});
sess.append_message(&old);
sess.append_compaction("## Goal\nsum", 10, None, None, None);
agent.session = Some(sess);
agent.messages = vec![old];
agent.auto_compaction = true;
let rt = test_rt();
let r =
rt.block_on(agent.maybe_compact_with(compaction::CompactionReason::Threshold, None));
assert!(
matches!(r, Ok(None)),
"陈旧 usage 锚点应跳过阈值压缩: {:?}",
r
);
}
#[test]
fn compaction_end_reports_shrunk_context_for_footer() {
let rt = test_rt();
let events: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let mut summarized = AgentMessage::user_text(&"A".repeat(300_000));
summarized.timestamp = 1;
let mut summarized_reply = AgentMessage::user_text("head answer");
summarized_reply.role = "assistant".into();
summarized_reply.timestamp = 2;
let mut kept_user = AgentMessage::user_text(&"B".repeat(200_000));
kept_user.timestamp = 3;
let mut kept_assistant = AgentMessage::user_text("tail answer");
kept_assistant.role = "assistant".into();
kept_assistant.timestamp = 4;
let messages = vec![summarized, summarized_reply, kept_user, kept_assistant];
let tokens_before = compaction::estimate_context_tokens(&messages).tokens;
let (base, handle) = rt.block_on(mock_server_sequence(vec![MockResp::Sse(text_frame(
"summary text",
"s",
))]));
let mut agent = test_agent(&base, ".");
let sink_events = events.clone();
agent.attach_json_sink(Box::new(move |e: Value| {
sink_events.lock().unwrap().push(e);
}));
agent.messages = messages;
let r = rt.block_on(agent.maybe_compact_with(compaction::CompactionReason::Manual, None));
let _ = rt.block_on(handle);
assert!(matches!(r, Ok(Some(_))), "应完成压缩: {:?}", r);
assert_eq!(agent.messages.len(), 3);
assert_eq!(agent.messages[0].role, "compactionSummary");
assert!(
agent.messages[1].text().starts_with("BBBB"),
"切点之后的消息原样保留"
);
assert_eq!(agent.messages[2].text(), "tail answer");
let evs = events.lock().unwrap();
let ev = evs
.iter()
.find(|e| e["type"] == "compaction_end" && e["result"] == "completed")
.expect("compaction_end(completed) 事件");
assert_eq!(
ev["tokensBefore"].as_u64(),
Some(tokens_before as u64),
"压缩前 token 数"
);
let after = ev["estimatedTokensAfter"].as_u64().unwrap();
assert!(
after < tokens_before as u64,
"压缩后估算应下降(footer 占用随之更新): {tokens_before} → {after}"
);
let pure: u64 = agent
.messages
.iter()
.map(compaction::estimate_tokens)
.sum::<u32>() as u64;
assert_eq!(after, pure, "压缩后估算 = 重建上下文纯估算");
}
#[test]
fn overflow_error_patterns_cover_providers() {
assert!(Agent::is_overflow_error_str(
"prompt is too long: 213462 tokens > 200000 maximum"
));
assert!(Agent::is_overflow_error_str("Prompt exceeds max length"));
assert!(Agent::is_overflow_error_str(
"request_too_large: Request exceeds the maximum size"
));
assert!(Agent::is_overflow_error_str(
"Your input exceeds the context window of this model"
));
assert!(Agent::is_overflow_error_str(
"exceeds the model's maximum context length of 131072 tokens"
));
assert!(Agent::is_overflow_error_str(
"This model's maximum prompt length is 131072 but the request contains 537812 tokens"
));
assert!(Agent::is_overflow_error_str(
"Please reduce the length of the messages or completion"
));
assert!(Agent::is_overflow_error_str(
"The input token count (1196265) exceeds the maximum number of tokens allowed (1048575)"
));
assert!(Agent::is_overflow_error_str(
"prompt token count of 1000 exceeds the limit of 999"
));
assert!(Agent::is_overflow_error_str(
"Invalid params: context window exceeds limit"
));
assert!(Agent::is_overflow_error_str(
"Prompt contains 90001 tokens ... too large for model with 200000 maximum context length"
));
assert!(Agent::is_overflow_error_str(
"Prompt has 90001 tokens, but the configured context size is 200000 tokens"
));
assert!(Agent::is_overflow_error_str(
"The input (90001 tokens) is longer than the model's context length (20000 tokens)"
));
assert!(Agent::is_overflow_error_str(
"Range of input length should be [1, 90000]"
));
assert!(!Agent::is_overflow_error_str(
"ThrottlingException: Too many tokens, please wait before trying again"
));
assert!(!Agent::is_overflow_error_str(
"Rate limit reached for anthropic"
));
assert!(!Agent::is_overflow_error_str("Service unavailable"));
}
#[test]
fn context_overflow_covers_error_silent_and_length_stop() {
let mut m = AgentMessage::user_text("");
m.role = "assistant".into();
m.stop_reason = Some("error".into());
m.error_message = Some("prompt is too long".into());
assert!(Agent::message_is_context_overflow(&m, 200_000, "anthropic"));
m.error_message = Some("provider returned 400 Bad Request: ".into());
assert!(Agent::message_is_context_overflow(&m, 200_000, "cerebras"));
assert!(!Agent::message_is_context_overflow(&m, 200_000, "openai"));
m.stop_reason = Some("stop".into());
m.error_message = None;
m.usage = Some(Usage {
input: 210_000,
..Usage::default()
});
assert!(Agent::message_is_context_overflow(&m, 200_000, "zai"));
m.usage = Some(Usage {
input: 100_000,
..Usage::default()
});
assert!(!Agent::message_is_context_overflow(&m, 200_000, "zai"));
m.stop_reason = Some("length".into());
m.usage = Some(Usage {
input: 199_000,
output: 0,
..Usage::default()
});
assert!(Agent::message_is_context_overflow(&m, 200_000, "xiaomi"));
m.usage = Some(Usage {
input: 199_000,
output: 5,
..Usage::default()
});
assert!(!Agent::message_is_context_overflow(&m, 200_000, "xiaomi"));
}
#[test]
fn overflow_removes_failed_message_compacts_and_retries_once() {
let rt = test_rt();
let events: std::sync::Arc<std::sync::Mutex<Vec<Value>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let overflow_body =
r#"{"error":{"message":"prompt is too long: 500000 tokens > 200000 maximum"}}"#;
let (base, handle) = rt.block_on(mock_server_sequence(vec![
MockResp::Err("HTTP/1.1 400 Bad Request", overflow_body),
MockResp::Sse(text_frame("summary text", "s")),
MockResp::Sse(text_frame("retried reply", "r")),
]));
let mut agent = test_agent(&base, ".");
agent.auto_retry = true;
let sink_events = events.clone();
agent.attach_json_sink(Box::new(move |e: Value| {
sink_events.lock().unwrap().push(e);
}));
let result = rt.block_on(agent.prompt_messages(vec![
AgentMessage::user_text("hello"),
AgentMessage::user_text(&"x".repeat(400_000)), ]));
let _ = rt.block_on(handle);
let r = result.unwrap();
assert_eq!(r, "retried reply");
assert!(
!agent
.messages
.iter()
.any(|m| m.stop_reason.as_deref() == Some("error")),
"重试后 messages 不应残留 error assistant"
);
assert_eq!(
agent.messages.first().map(|m| m.role.as_str()),
Some("compactionSummary")
);
let evs = events.lock().unwrap();
assert!(
evs.iter()
.any(|e| e["type"] == "compaction_start" && e["reason"] == "overflow")
);
assert!(
evs.iter()
.any(|e| e["type"] == "compaction_end" && e["reason"] == "overflow")
);
assert!(
evs.iter().filter(|e| e["type"] == "agent_end").count() > 1,
"溢出恢复应产生多次 agent_end(每次尝试一次): {evs:?}"
);
assert_eq!(
evs.iter().filter(|e| e["type"] == "agent_settled").count(),
1,
"整轮应只发一次 agent_settled: {evs:?}"
);
assert_eq!(
evs.last().map(|e| e["type"].as_str()),
Some(Some("agent_settled")),
"agent_settled 应为最后一个事件"
);
assert_eq!(
evs.last().map(|e| e["aborted"].as_bool()),
Some(Some(false)),
"正常完成的整轮 aborted=false"
);
}
#[test]
fn overflow_retry_guard_compacts_only_once() {
let rt = test_rt();
let events: std::sync::Arc<std::sync::Mutex<Vec<Value>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let overflow_body =
r#"{"error":{"message":"prompt is too long: 500000 tokens > 200000 maximum"}}"#;
let (base, handle) = rt.block_on(mock_server_sequence(vec![
MockResp::Err("HTTP/1.1 400 Bad Request", overflow_body),
MockResp::Sse(text_frame("summary text", "s")),
MockResp::Err("HTTP/1.1 400 Bad Request", overflow_body),
]));
let mut agent = test_agent(&base, ".");
agent.auto_retry = true;
let sink_events = events.clone();
agent.attach_json_sink(Box::new(move |e: Value| {
sink_events.lock().unwrap().push(e);
}));
let result = rt.block_on(agent.prompt_messages(vec![
AgentMessage::user_text("hello"),
AgentMessage::user_text(&"x".repeat(400_000)), ]));
let _ = rt.block_on(handle);
assert!(result.is_ok(), "error 回合不抛 Err: {:?}", result.err());
let evs = events.lock().unwrap();
let starts = evs
.iter()
.filter(|e| e["type"] == "compaction_start" && e["reason"] == "overflow")
.count();
assert_eq!(starts, 1, "overflow 压缩必须只发生一次");
assert!(evs.iter().any(|e| {
e["type"] == "compaction_end" && e["reason"] == "overflow" && e["willRetry"] == false
}));
}
#[test]
fn retry_omits_abandoned_attempt_from_rebuilt_context() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let dir = crate::test_support::agent_dir_override().expect("temp agent dir");
std::fs::write(dir.join("settings.json"), r#"{"retry":{"baseDelayMs":1}}"#).unwrap();
let rt = test_rt();
let (base, handle) = rt.block_on(mock_server_sequence(vec![
MockResp::Err(
"HTTP/1.1 503 Service Unavailable",
r#"{"error":{"message":"overloaded"}}"#,
),
MockResp::Sse(text_frame("recovered", "r")),
]));
let mut agent = session_agent(&base, ".");
agent.auto_retry = true;
let text = rt.block_on(agent.prompt("go")).unwrap();
let _ = rt.block_on(handle);
assert_eq!(text, "recovered");
let session = agent.session.as_ref().unwrap();
let entries = session.get_entries();
let edit = entries
.iter()
.find(|e| e["type"] == "context_edit")
.expect("被放弃的 attempt 应落库 context_edit");
assert!(edit["replacement"].is_null(), "省略形态 replacement=null");
let target = edit["targetId"].as_str().unwrap().to_string();
let target_entry = entries
.iter()
.find(|e| e["id"] == target.as_str())
.expect("目标条目仍在原始 transcript 里");
assert_eq!(target_entry["message"]["stopReason"], "error");
let path = session
.get_session_file()
.unwrap()
.to_string_lossy()
.to_string();
let reopened = crate::core::session_manager::Session::open(&path).unwrap();
let msgs = reopened.build_context_messages();
assert!(
!msgs
.iter()
.any(|m| m.entry_id.as_deref() == Some(target.as_str())),
"resume 后不应包含被放弃的 attempt: {msgs:?}"
);
assert!(
reopened
.get_entries()
.iter()
.any(|e| e["id"] == target.as_str()),
"原始 transcript 仍完整保留该条目"
);
assert!(msgs.iter().any(|m| m.text() == "recovered"));
}
#[test]
fn overflow_recovery_omits_failed_attempt_persistently() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = test_rt();
let overflow_body =
r#"{"error":{"message":"prompt is too long: 500000 tokens > 200000 maximum"}}"#;
let (base, handle) = rt.block_on(mock_server_sequence(vec![
MockResp::Err("HTTP/1.1 400 Bad Request", overflow_body),
MockResp::Sse(text_frame("summary text", "s")),
MockResp::Sse(text_frame("retried reply", "r")),
]));
let mut agent = session_agent(&base, ".");
agent.auto_retry = true;
let result = rt.block_on(agent.prompt_messages(vec![
AgentMessage::user_text("hello"),
AgentMessage::user_text(&"x".repeat(400_000)),
]));
let _ = rt.block_on(handle);
assert_eq!(result.unwrap(), "retried reply");
let session = agent.session.as_ref().unwrap();
let entries = session.get_entries();
let edit = entries
.iter()
.find(|e| e["type"] == "context_edit")
.expect("overflow 恢复也应落库 context_edit");
assert!(edit["replacement"].is_null());
let target = edit["targetId"].as_str().unwrap().to_string();
let path = session
.get_session_file()
.unwrap()
.to_string_lossy()
.to_string();
let reopened = crate::core::session_manager::Session::open(&path).unwrap();
let msgs = reopened.build_context_messages();
assert!(
!msgs
.iter()
.any(|m| m.entry_id.as_deref() == Some(target.as_str())),
"resume 后被放弃的溢出 attempt 不应回到上下文"
);
}
#[test]
fn usage_reasoning_fields_present_in_provider_results() {
let u = Usage {
input: 1,
output: 2,
reasoning: Some(1),
..Default::default()
};
let v = serde_json::to_value(&u).unwrap();
assert_eq!(v["reasoning"], 1);
assert_eq!(v["reasoning"], v["reasoning"]);
}
fn test_agent(base: &str, cwd: &str) -> Agent {
test_agent_with_input(base, cwd, vec!["text".into()])
}
fn test_agent_with_input(base: &str, cwd: &str, input: Vec<String>) -> Agent {
let entry = model_resolver::ModelEntry {
model_type: Default::default(),
id: "test-model".into(),
name: "Test".into(),
api: "openai-completions".into(),
base_url: base.into(),
provider: "test".into(),
input_limits: crate::utils::image::InputLimits::default(),
allowed_fallback_models: Vec::new(),
sampling_params: None,
sampling_params_by_thinking_level: None,
provider_routing: None,
supports_usage_in_streaming: true,
supports_store: true,
supports_finish_reason: true,
requires_assistant_after_tool_result: false,
supports_strict_mode: false,
send_session_affinity_headers: None,
session_affinity_format: None,
output: Vec::new(),
input,
reasoning: false,
context_window: 200_000,
max_tokens: 4096,
max_tokens_field: "max_tokens".into(),
cost: None,
supports_developer_role: false,
requires_reasoning_content_on_assistant_messages: false,
thinking_format: String::new(),
supports_reasoning_effort: false,
thinking_level_map: None,
auth_header: false,
supports_openai_grammar_tools: false,
};
Agent::new(
entry,
cwd.to_string(),
ToolSelection::tools(vec!["read".into(), "write".into()]),
Vec::new(),
None,
None,
None,
Vec::new(),
None,
Some("sk-test".to_string()),
true,
Vec::new(),
)
.unwrap()
}
fn test_agent_selection(selection: ToolSelection, exclude: Vec<String>) -> Agent {
let entry = model_resolver::ModelEntry {
model_type: Default::default(),
id: "test-model".into(),
name: "Test".into(),
api: "openai-completions".into(),
base_url: "http://127.0.0.1:1".into(),
provider: "test".into(),
input_limits: crate::utils::image::InputLimits::default(),
allowed_fallback_models: Vec::new(),
sampling_params: None,
sampling_params_by_thinking_level: None,
provider_routing: None,
supports_usage_in_streaming: true,
supports_store: true,
supports_finish_reason: true,
requires_assistant_after_tool_result: false,
supports_strict_mode: false,
send_session_affinity_headers: None,
session_affinity_format: None,
output: Vec::new(),
input: vec!["text".into()],
reasoning: false,
context_window: 200_000,
max_tokens: 4096,
max_tokens_field: "max_tokens".into(),
cost: None,
supports_developer_role: false,
requires_reasoning_content_on_assistant_messages: false,
thinking_format: String::new(),
supports_reasoning_effort: false,
thinking_level_map: None,
auth_header: false,
supports_openai_grammar_tools: false,
};
Agent::new(
entry,
"/tmp".to_string(),
selection,
exclude,
None,
None,
None,
Vec::new(),
None,
Some("sk-test".to_string()),
true,
Vec::new(),
)
.unwrap()
}
#[test]
fn every_declared_builtin_tool_resolves() {
let injected: Vec<InjectedChildTool> = Vec::new();
for name in tools::all_tool_names() {
assert!(Agent::is_builtin_tool(&name), "{name} 未被视为内置");
assert!(
ToolIndex::resolve(&injected, &name).exists(),
"{name} 已声明给模型却解析不到"
);
}
assert_eq!(Agent::is_builtin_tool("powershell"), cfg!(windows));
}
#[test]
fn deferred_tools_hidden_until_activated() {
use crate::core::extensions::{
Extension, ExtensionTool, ToolExposure, activate_deferred_tools,
clear_deferred_activations, register_extension, unregister_extension,
};
let _g = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let _ad = crate::test_support::AgentDirGuard::temp();
crate::core::settings_manager::write_disabled_extensions(&[]).ok();
struct DeferredExt;
impl Extension for DeferredExt {
fn name(&self) -> &str {
"deferred-demo"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![
ExtensionTool::simple(
"deferred_weather",
"Look up the weather",
serde_json::json!({ "type": "object" }),
"weather lookup",
)
.with_exposure(ToolExposure::Deferred),
ExtensionTool::simple(
"eager_ping",
"Ping a host",
serde_json::json!({ "type": "object" }),
"ping",
),
]
}
}
register_extension(DeferredExt);
clear_deferred_activations();
let selection = || ToolSelection::Default(vec!["read".to_string()]);
let agent = test_agent_selection(selection(), vec![]);
let active = agent.get_active_tools();
assert!(
active.iter().any(|n| n == "eager_ping"),
"非延迟的扩展工具照常可见: {active:?}"
);
assert!(
!active.iter().any(|n| n == "deferred_weather"),
"延迟工具未激活前不得进工具表: {active:?}"
);
assert!(
!agent.system_prompt.contains("deferred_weather"),
"延迟工具不得进系统提示"
);
activate_deferred_tools(&["deferred_weather".to_string()]);
let agent2 = test_agent_selection(selection(), vec![]);
let active2 = agent2.get_active_tools();
assert!(
active2.iter().any(|n| n == "deferred_weather"),
"激活后应进工具表: {active2:?}"
);
clear_deferred_activations();
let agent3 = test_agent_selection(selection(), vec![]);
assert!(
!agent3
.get_active_tools()
.iter()
.any(|n| n == "deferred_weather"),
"清空激活集合后重新隐藏"
);
assert!(unregister_extension("deferred-demo"));
}
#[test]
fn extension_tools_join_available_tools_section() {
use crate::core::extensions::{
Extension, ExtensionTool, register_extension, unregister_extension,
};
let _g = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let _ad = crate::test_support::AgentDirGuard::temp();
crate::core::settings_manager::write_disabled_extensions(&[]).ok();
struct SnippetProbe;
impl Extension for SnippetProbe {
fn name(&self) -> &str {
"snippet-probe"
}
fn tools(&self) -> Vec<ExtensionTool> {
let mut listed = ExtensionTool::simple(
"probe_listed",
"A tool with a snippet",
serde_json::json!({ "type": "object" }),
"probe one-liner",
);
listed.prompt_guidelines = vec!["probe guideline".to_string()];
let silent = ExtensionTool::simple(
"probe_silent",
"A tool without a snippet",
serde_json::json!({ "type": "object" }),
"",
);
vec![listed, silent]
}
}
register_extension(SnippetProbe);
let agent = test_agent_selection(ToolSelection::Default(vec!["read".to_string()]), vec![]);
let active = agent.get_active_tools();
assert!(active.iter().any(|n| n == "probe_listed"), "{active:?}");
assert!(
active.iter().any(|n| n == "probe_silent"),
"无 snippet 的工具仍要进工具表:{active:?}"
);
assert!(
agent
.system_prompt
.contains("- probe_listed: probe one-liner"),
"扩展工具的 snippet 必须进 Available tools:{}",
agent.system_prompt
);
assert!(
agent.system_prompt.contains("- probe guideline"),
"扩展工具的 guideline 必须进 Guidelines"
);
assert!(
!agent.system_prompt.contains("probe_silent"),
"无 snippet 的扩展工具不进 Available tools 段:{}",
agent.system_prompt
);
assert!(unregister_extension("snippet-probe"));
}
#[test]
fn tool_selection_no_tools_disables_all() {
let agent = test_agent_selection(ToolSelection::NoTools, vec![]);
assert!(
agent.get_active_tools().is_empty(),
"active: {:?}",
agent.get_active_tools()
);
assert!(
agent.get_all_tools().is_empty(),
"all: {:?}",
agent.get_all_tools()
);
}
#[test]
fn tool_selection_no_builtin_keeps_extensions() {
let agent = test_agent_selection(ToolSelection::NoBuiltinTools, vec![]);
let active = agent.get_active_tools();
assert!(
active.iter().all(|n| !tools::all_tool_names().contains(n)),
"active 不应含内置: {:?}",
active
);
let all = agent.get_all_tools();
assert!(
tools::all_tool_names().iter().all(|b| all.contains(b)),
"all: {:?}",
all
);
}
#[test]
fn tool_selection_allowlist_applies_to_all() {
let list = vec!["read".to_string(), "edit".to_string()];
let agent = test_agent_selection(ToolSelection::Allowlist(list.clone()), vec![]);
let active = agent.get_active_tools();
assert!(
active.iter().all(|n| list.contains(n)),
"active 超出名单: {:?}",
active
);
assert!(active.contains(&"read".to_string()));
assert!(!active.contains(&"bash".to_string()));
assert!(!active.contains(&"write".to_string()));
let all = agent.get_all_tools();
assert!(all.iter().all(|n| list.contains(n)), "all: {:?}", all);
}
#[test]
fn tool_selection_default_and_exclude() {
let agent = test_agent_selection(
ToolSelection::tools(vec!["read".into(), "write".into()]),
vec!["write".to_string()],
);
let active = agent.get_active_tools();
assert!(active.contains(&"read".to_string()));
assert!(!active.contains(&"write".to_string()));
assert!(!active.contains(&"bash".to_string()));
let all = agent.get_all_tools();
assert!(
!all.contains(&"write".to_string()),
"exclude 应作用于 getAllTools: {:?}",
all
);
}
#[test]
fn tool_selection_supports_wildcard_entries() {
let agent = test_agent_selection(ToolSelection::Allowlist(vec!["re*".into()]), vec![]);
let active = agent.get_active_tools();
assert!(
active.contains(&"read".to_string()),
"re* 应命中 read: {active:?}"
);
assert!(!active.contains(&"bash".to_string()));
assert!(!active.contains(&"grep".to_string()));
let agent = test_agent_selection(
ToolSelection::tools(vec!["read".into(), "bash".into()]),
vec!["b*".into()],
);
let active = agent.get_active_tools();
assert!(active.contains(&"read".to_string()));
assert!(
!active.contains(&"bash".to_string()),
"--exclude-tools 的通配条目应生效: {active:?}"
);
}
#[test]
fn allowlist_keeps_mcp_tools_unless_pointed_at() {
use crate::core::extensions::{
Extension, ExtensionTool, register_extension, unregister_extension,
};
struct McpGatewayProbe;
impl Extension for McpGatewayProbe {
fn name(&self) -> &str {
"mcp-gateway-probe"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![ExtensionTool::simple(
tool_names::MCP_GATEWAY_TOOL_NAME,
"MCP gateway",
serde_json::json!({ "type": "object" }),
"mcp gateway",
)]
}
}
let _g = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let _ad = crate::test_support::AgentDirGuard::temp();
crate::core::settings_manager::write_disabled_extensions(&[]).ok();
register_extension(McpGatewayProbe);
let agent = test_agent_selection(ToolSelection::Allowlist(vec!["read".into()]), vec![]);
let active = agent.get_active_tools();
assert!(active.contains(&"read".to_string()));
assert!(
!active.contains(&"mcp".to_string()),
"未点名的 MCP 工具不得对模型声明: {active:?}"
);
assert!(
agent.get_all_tools().contains(&"mcp".to_string()),
"未点名的 MCP 工具仍应保留注册"
);
assert!(
!agent.system_prompt.contains("MCP gateway"),
"未声明的 MCP 工具不得进系统提示"
);
let script = script_tools_for(
&[],
&[],
&ToolSelection::Allowlist(vec!["read".into()]),
&[],
);
assert!(
script.iter().any(|t| t.name == "mcp"),
"保留的 MCP 工具应对脚本可调: {:?}",
script.iter().map(|t| &t.name).collect::<Vec<_>>()
);
let excluded = test_agent_selection(
ToolSelection::Allowlist(vec!["read".into()]),
vec!["mcp".to_string()],
);
assert!(!excluded.get_all_tools().contains(&"mcp".to_string()));
let filtered = test_agent_selection(
ToolSelection::Allowlist(vec!["read".into(), "mcp__docs__*".into()]),
vec![],
);
assert!(!filtered.get_all_tools().contains(&"mcp".to_string()));
let pointed = test_agent_selection(ToolSelection::Allowlist(vec!["mcp".into()]), vec![]);
assert!(pointed.get_active_tools().contains(&"mcp".to_string()));
assert!(unregister_extension("mcp-gateway-probe"));
}
#[test]
fn reload_reapplies_default_tool_modifiers() {
let _g = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let _ad = crate::test_support::AgentDirGuard::temp();
let path = crate::core::settings_manager::agent_dir().join("settings.json");
std::fs::write(&path, r#"{"defaultTools":["read"]}"#).unwrap();
let base = vec!["read".to_string()];
let modifiers = vec!["+bash".to_string()];
let mut agent = test_agent_selection(
ToolSelection::tools(settings_manager::apply_tool_modifiers(&base, &modifiers)),
vec![],
);
agent.set_settings_default_tools(Some(settings_manager::apply_tool_modifiers(
&base, &modifiers,
)));
agent.set_default_tool_modifiers(modifiers);
std::fs::write(&path, r#"{"defaultTools":["read","edit"]}"#).unwrap();
agent.reload_default_tools();
let ToolSelection::Default(names) = &agent.rebuild_ctx.selection else {
panic!("Default 选择不应被改写成其它变体");
};
assert_eq!(
names,
&vec!["read".to_string(), "bash".to_string(), "edit".to_string()]
);
}
#[test]
fn reload_default_tools_only_adds_new_names() {
let _g = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let _ad = crate::test_support::AgentDirGuard::temp();
let path = crate::core::settings_manager::agent_dir().join("settings.json");
std::fs::write(&path, r#"{"defaultTools":["read"]}"#).unwrap();
let mut agent = test_agent_selection(ToolSelection::tools(Vec::new()), vec![]);
agent.set_settings_default_tools(Some(vec!["read".to_string()]));
std::fs::write(&path, r#"{"defaultTools":["read","bash"]}"#).unwrap();
agent.reload_default_tools();
let ToolSelection::Default(names) = &agent.rebuild_ctx.selection else {
panic!("Default 选择不应被改写成其它变体");
};
assert_eq!(names, &vec!["bash".to_string()]);
std::fs::write(&path, r#"{"defaultTools":["read"]}"#).unwrap();
agent.reload_default_tools();
let ToolSelection::Default(names) = &agent.rebuild_ctx.selection else {
panic!("Default 选择不应被改写成其它变体");
};
assert_eq!(names, &vec!["bash".to_string()]);
}
#[test]
fn reload_default_tools_skips_cli_override() {
let _g = crate::test_support::AUTH_TEST_LOCK
.lock()
.unwrap_or_else(|e| e.into_inner());
let _ad = crate::test_support::AgentDirGuard::temp();
let path = crate::core::settings_manager::agent_dir().join("settings.json");
std::fs::write(&path, r#"{"defaultTools":["read","bash"]}"#).unwrap();
let mut agent =
test_agent_selection(ToolSelection::Allowlist(vec!["read".to_string()]), vec![]);
agent.reload_default_tools();
let ToolSelection::Allowlist(names) = &agent.rebuild_ctx.selection else {
panic!("--tools 的 allowlist 不应被 /reload 改写");
};
assert_eq!(names, &vec!["read".to_string()]);
}
fn ev_types(events: &[J]) -> Vec<String> {
events
.iter()
.map(|e| {
e.get("type")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string()
})
.collect()
}
fn test_rt() -> tokio::runtime::Runtime {
tokio::runtime::Runtime::new().unwrap()
}
fn events_of<'a>(events: &'a [J], ty: &'a str) -> Vec<&'a J> {
events
.iter()
.filter(|e| e.get("type").and_then(|v| v.as_str()) == Some(ty))
.collect()
}
#[test]
fn run_loop_text_reply_pi_event_sequence() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir();
let cwd = cwd.to_string_lossy().to_string();
let frames = text_frame("Hello out there", "cmpl_t");
let (base, handle) = mock_server_multi(vec![frames]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("hi").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Hello out there");
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(types.first().map(String::as_str), Some("agent_start"));
assert_eq!(types[1], "turn_start");
assert!(types.contains(&"message_start".to_string()));
assert!(types.contains(&"message_update".to_string()));
assert!(types.contains(&"turn_end".to_string()));
assert!(types.contains(&"agent_end".to_string()));
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
1
);
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
assert_eq!(msgs.len(), 2);
assert_eq!(msgs[0]["role"], "user");
assert_eq!(msgs[0]["content"][0]["text"], "hi");
assert_eq!(msgs[1]["role"], "assistant");
assert_eq!(msgs[1]["content"][0]["text"], "Hello out there");
});
}
#[test]
fn virtual_model_routes_request_and_records_physical_model() {
let _ad = crate::test_support::AgentDirGuard::temp();
let rt = test_rt();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) =
mock_server_multi(vec![text_frame("routed reply", "cmpl_v")]).await;
std::fs::write(
crate::core::settings_manager::agent_dir().join("models.json"),
json!({"providers":{"mockprov":{
"apiKey":"sk-mock",
"baseUrl":base,
"api":"openai-completions",
"models":[{"id":"physical","name":"Physical"}]
}}})
.to_string(),
)
.unwrap();
struct FixedRouter;
impl virtual_models::VirtualModelRouter for FixedRouter {
fn route<'a>(
&'a self,
_r: virtual_models::ModelRouteRequest<'a>,
) -> BoxFuture<'a, Result<virtual_models::ModelRoute>> {
Box::pin(async {
Ok(virtual_models::ModelRoute {
provider: "mockprov".into(),
model_id: "physical".into(),
thinking_level: None,
state: None,
})
})
}
}
virtual_models::register(
virtual_models::VirtualModelDefinition::new(
"mockprov",
"auto",
"Auto",
Arc::new(FixedRouter),
),
virtual_models::STANDALONE_OWNER,
);
let listed = model_resolver::list_models("mockprov");
assert!(listed.iter().any(|(id, _)| id == "auto"), "{listed:?}");
let entry = model_resolver::find_model("mockprov", "auto").unwrap();
assert_eq!(entry.api, virtual_models::VIRTUAL_MODEL_API);
let mut agent = Agent::new(
entry,
cwd,
ToolSelection::tools(vec!["read".into()]),
Vec::new(),
None,
None,
None,
Vec::new(),
None,
None,
false,
Vec::new(),
)
.unwrap();
assert!(agent.selected_is_virtual());
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("hi").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "routed reply");
let last = agent
.messages
.iter()
.rev()
.find(|m| m.role == "assistant")
.unwrap();
assert_eq!(last.provider.as_deref(), Some("mockprov"));
assert_eq!(last.model.as_deref(), Some("physical"));
let events = collected.lock().unwrap();
let routed = events_of(&events, "model_routed");
assert_eq!(routed.len(), 1, "{events:?}");
assert_eq!(routed[0]["model"], "physical");
virtual_models::unregister("mockprov", "auto");
});
}
#[test]
fn virtual_model_routing_failure_ends_run_with_error_message() {
let _ad = crate::test_support::AgentDirGuard::temp();
let rt = test_rt();
rt.block_on(async {
struct FailingRouter;
impl virtual_models::VirtualModelRouter for FailingRouter {
fn route<'a>(
&'a self,
_r: virtual_models::ModelRouteRequest<'a>,
) -> BoxFuture<'a, Result<virtual_models::ModelRoute>> {
Box::pin(async { Err(Error::msg("router exploded")) })
}
}
virtual_models::register(
virtual_models::VirtualModelDefinition::new(
"prux-vm-fail",
"auto",
"Auto",
Arc::new(FailingRouter),
),
virtual_models::STANDALONE_OWNER,
);
let entry = model_resolver::find_model("prux-vm-fail", "auto").unwrap();
let mut agent = Agent::new(
entry,
std::env::temp_dir().to_string_lossy().to_string(),
ToolSelection::tools(vec!["read".into()]),
Vec::new(),
None,
None,
None,
Vec::new(),
None,
None,
false,
Vec::new(),
)
.unwrap();
_ = agent.prompt("hi").await;
let last = agent.messages.last().unwrap();
assert_eq!(last.role, "assistant");
assert_eq!(last.stop_reason.as_deref(), Some("error"));
assert!(
last.error_message
.as_deref()
.unwrap_or_default()
.contains("router exploded"),
"{:?}",
last.error_message
);
assert_eq!(last.api.as_deref(), Some(virtual_models::VIRTUAL_MODEL_API));
virtual_models::unregister("prux-vm-fail", "auto");
});
}
#[test]
fn run_loop_tool_then_final_reply() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let tmp =
std::env::temp_dir().join(format!("prux-test-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
std::fs::write(tmp.join("a.rs"), "fn main() {}").unwrap();
let cwd = tmp.to_string_lossy().to_string();
let frames1 = tool_call_frame("read", r#"{\"path\":\"a.rs\"}"#, "call_1");
let frames2 = text_frame("Read done", "cmpl_2");
let (base, handle) = mock_server_multi(vec![frames1, frames2]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("read it").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Read done");
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(types.iter().filter(|t| t.as_str() == "turn_end").count(), 2);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
2
);
assert_eq!(
types
.iter()
.filter(|t| t.as_str() == "tool_execution_start")
.count(),
1
);
assert_eq!(
types
.iter()
.filter(|t| t.as_str() == "tool_execution_end")
.count(),
1
);
let tr_start = events
.iter()
.filter(|e| e["type"] == "message_start" && e["message"]["role"] == "toolResult")
.count();
assert_eq!(tr_start, 1);
let turn_end = events_of(&events, "turn_end")[0];
let tool_results = turn_end["toolResults"].as_array().unwrap();
assert_eq!(tool_results.len(), 1);
assert_eq!(tool_results[0]["role"], "toolResult");
assert_eq!(tool_results[0]["isError"], false);
assert!(
tool_results[0]["content"][0]["text"]
.as_str()
.unwrap()
.contains("fn main() {}")
);
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let roles: Vec<&str> = msgs.iter().map(|m| m["role"].as_str().unwrap()).collect();
assert_eq!(roles, vec!["user", "assistant", "toolResult", "assistant"]);
assert_eq!(msgs[1]["content"][0]["type"], "toolCall");
});
}
#[test]
fn run_loop_length_stop_fails_truncated_tool_calls_then_continues() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let frames1 = tool_call_frame("write", r#"{\"path\":\"x.txt\",\"content\":\"trunc"#, "call_1");
let frames1 = {
let mut v = frames1;
let idx = v.len() - 2;
v[idx] = r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"length"}]}"#;
v
};
let frames2 = text_frame("Recovered", "cmpl_2");
let (base, handle) = mock_server_multi(vec![frames1, frames2]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("write x").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Recovered");
let events = collected.lock().unwrap().clone();
assert_eq!(ev_types(&events).iter().filter(|t| t.as_str() == "tool_execution_start").count(), 1);
assert_eq!(ev_types(&events).iter().filter(|t| t.as_str() == "tool_execution_end").count(), 1);
let tr = events
.iter()
.filter(|e| e["type"] == "message_start" && e["message"]["role"] == "toolResult")
.count();
assert_eq!(tr, 1);
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let tr_msg = msgs.iter().find(|m| m["role"] == "toolResult").unwrap();
assert_eq!(tr_msg["isError"], true);
assert!(tr_msg["content"][0]["text"]
.as_str()
.unwrap()
.contains("output token limit"));
assert_eq!(msgs.last().unwrap()["content"][0]["text"], "Recovered");
});
}
#[test]
fn run_loop_error_stop_terminates_cleanly() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let frames = vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":"partial"}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"error"}]}"#,
r#"[DONE]"#,
];
let (base, handle) = mock_server_multi(vec![frames]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let _text = agent.prompt("boom").await.unwrap();
handle.await.unwrap();
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert!(types.contains(&"turn_end".to_string()));
assert!(types.contains(&"agent_end".to_string()));
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
assert_eq!(msgs.last().unwrap()["role"], "assistant");
assert_eq!(msgs.last().unwrap()["content"][0]["text"], "partial");
});
}
#[test]
fn steering_is_injected_via_runtime_inbox() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("First", "a")]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
agent
.runtime_steer_inbox
.lock()
.unwrap()
.push_back(AgentMessage::user_text("steer-msg"));
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "First");
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
1
);
let user_msgs: Vec<&str> = events
.iter()
.filter(|e| e["type"] == "message_start" && e["message"]["role"] == "user")
.map(|e| e["message"]["content"][0]["text"].as_str().unwrap())
.collect();
assert!(user_msgs.contains(&"go"));
assert!(user_msgs.contains(&"steer-msg"));
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let roles: Vec<&str> = msgs.iter().map(|m| m["role"].as_str().unwrap()).collect();
assert_eq!(roles, vec!["user", "user", "assistant"]);
assert_eq!(msgs[1]["content"][0]["text"], "steer-msg");
assert_eq!(msgs[2]["content"][0]["text"], "First");
});
}
#[test]
fn agent_end_carries_only_new_messages_not_history() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("Only new", "a")]).await;
let mut agent = test_agent(&base, &cwd);
agent
.messages
.push(AgentMessage::user_text("old history 1"));
agent
.messages
.push(AgentMessage::user_text("old history 2"));
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("now").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Only new");
let collected_events = collected.lock().unwrap().clone();
let agent_end = events_of(&collected_events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let roles: Vec<&str> = msgs.iter().map(|m| m["role"].as_str().unwrap()).collect();
assert_eq!(roles, vec!["user", "assistant"]);
assert_eq!(msgs[0]["content"][0]["text"], "now");
assert_eq!(msgs[1]["content"][0]["text"], "Only new");
});
}
#[test]
fn finish_turn_end_terminates_before_next_turn() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![
tool_call_frame("read", r#"{\"path\":\"a.rs\"}"#, "call_1"),
text_frame("Should not run", "b"),
])
.await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
agent.finish_turn = Some(std::sync::Mutex::new(Box::new(|_| {
Some(FinishTurnAction::End)
})));
let result = agent.prompt("go").await;
handle.await.unwrap();
let _ = result;
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
1
);
assert_eq!(types.iter().filter(|t| t.as_str() == "turn_end").count(), 1);
assert!(types.contains(&"agent_end".to_string()));
assert!(
!events
.iter()
.any(|e| e["message"]["content"][0]["text"] == "Should not run")
);
});
}
#[test]
fn finish_turn_continue_forces_extra_request() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) =
mock_server_multi(vec![text_frame("First", "a"), text_frame("Second", "b")]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let calls = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
let counter = calls.clone();
agent.finish_turn = Some(std::sync::Mutex::new(Box::new(move |_| {
if counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
Some(FinishTurnAction::Continue)
} else {
None
}
})));
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
2,
"Continue 应触发第二次请求"
);
assert_eq!(text, "Second");
});
}
#[test]
fn finish_turn_runs_on_error_and_decision_ignored() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("boom", "a")]).await;
let mut agent = test_agent(&base, &cwd);
agent.model.base_url = "http://127.0.0.1:1".to_string();
agent.auto_retry = false;
let seen: std::sync::Arc<std::sync::Mutex<Vec<Option<String>>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let rec = seen.clone();
agent.finish_turn = Some(std::sync::Mutex::new(Box::new(move |ctx| {
rec.lock().unwrap().push(ctx.message.stop_reason.clone());
Some(FinishTurnAction::End)
})));
let _ = agent.prompt("go").await;
handle.await.unwrap();
let seen = seen.lock().unwrap().clone();
assert_eq!(seen.len(), 1, "硬退出也应调用一次 finishTurn");
assert_eq!(seen[0].as_deref(), Some("error"));
});
}
struct BoundaryProbe;
#[derive(Default)]
struct ProbeState {
turn: u8,
settle: u8,
seen: Vec<J>,
}
thread_local! {
static PROBE: std::cell::RefCell<ProbeState> = std::cell::RefCell::new(ProbeState::default());
}
impl core::extensions::Extension for BoundaryProbe {
fn name(&self) -> &str {
"boundary-probe-test"
}
fn tools(&self) -> Vec<core::extensions::ExtensionTool> {
Vec::new()
}
fn hooks(&self) -> Vec<core::extensions::ExtensionHook> {
vec![core::extensions::ExtensionHook::Boundary]
}
fn on_boundary(&self, event: &J) -> Option<core::extensions::BoundaryOutcome> {
let ty = event["type"].as_str().unwrap_or("");
PROBE.with(|p| {
let mut p = p.borrow_mut();
p.seen.push(event.clone());
let slot = match ty {
"turn_end" => &mut p.turn,
"agent_before_settle" => &mut p.settle,
_ => return None,
};
match *slot {
1 => {
*slot = 0; Some(core::extensions::BoundaryOutcome::continue_once())
}
2 => Some(core::extensions::BoundaryOutcome::end()),
3 => {
*slot = 0;
Some(core::extensions::BoundaryOutcome::continue_with(vec![
crate::core::provider::AgentMessage::user_text("probe-nudge"),
]))
}
4 => {
*slot = 0;
let target = event["messageEntryId"].as_str().unwrap_or("").to_string();
Some(core::extensions::BoundaryOutcome::default().with_omit(target))
}
5 => {
*slot = 0;
Some(
core::extensions::BoundaryOutcome::default()
.with_retain_none_compaction("probe-self-retain"),
)
}
_ => None,
}
})
}
}
#[test]
fn turn_end_boundary_context_edit_omits_entry() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("done", "a")]).await;
let mut agent = session_agent(&base, &cwd);
probe_set("turn_end", 4);
let text = agent.prompt("go").await.unwrap();
assert_eq!(text, "done");
handle.await.unwrap();
let seen = probe_seen("turn_end");
let target = seen[0]["messageEntryId"]
.as_str()
.expect("turn_end 带 messageEntryId")
.to_string();
let entries = agent.session.as_ref().unwrap().get_entries();
let edit = entries
.iter()
.find(|e| e["type"] == "context_edit")
.expect("边界草稿应落库为 context_edit");
assert_eq!(edit["targetId"], target.as_str());
assert!(edit["replacement"].is_null());
assert!(entries.iter().any(|e| e["id"] == target.as_str()));
let rebuilt = agent.session.as_ref().unwrap().build_context_messages();
assert!(
!rebuilt
.iter()
.any(|m| m.entry_id.as_deref() == Some(target.as_str())),
"被省略的条目不应出现在重建后的上下文: {rebuilt:?}"
);
});
}
#[test]
fn turn_end_boundary_retain_none_compaction_self_retains() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("done", "a")]).await;
let mut agent = session_agent(&base, &cwd);
probe_set("turn_end", 5);
let text = agent.prompt("go").await.unwrap();
assert_eq!(text, "done");
handle.await.unwrap();
let entries = agent.session.as_ref().unwrap().get_entries();
let comp = entries
.iter()
.find(|e| e["type"] == "compaction")
.expect("边界草稿应落库为 compaction");
assert_eq!(
comp["firstKeptEntryId"].as_str(),
comp["id"].as_str(),
"retain-none:firstKeptEntryId 自引用(pi `firstKeptEntryId ?? id`)"
);
assert!(
comp.get("retainedTail").is_none(),
"新格式不再内联保留尾部: {comp}"
);
assert_eq!(comp["summary"], "probe-self-retain");
let rebuilt = agent.session.as_ref().unwrap().build_context_messages();
assert_eq!(rebuilt.len(), 1);
assert_eq!(rebuilt[0].role, "compactionSummary");
assert_eq!(agent.messages.len(), 1);
assert_eq!(agent.messages[0].role, "compactionSummary");
});
}
struct ContextProbe;
thread_local! {
static CTX_PROBE: std::cell::RefCell<(Option<String>, Vec<J>)> =
const { std::cell::RefCell::new((None, Vec::new())) };
}
impl core::extensions::Extension for ContextProbe {
fn name(&self) -> &str {
"context-with-system-probe-test"
}
fn tools(&self) -> Vec<core::extensions::ExtensionTool> {
Vec::new()
}
fn hooks(&self) -> Vec<core::extensions::ExtensionHook> {
vec![core::extensions::ExtensionHook::ContextWithSystem]
}
fn transform_context_with_system(
&self,
messages: &mut Vec<AgentMessage>,
) -> crate::error::Result<()> {
CTX_PROBE.with(|p| {
let mut p = p.borrow_mut();
p.1.push(json!({
"roles": messages.iter().map(|m| m.role.clone()).collect::<Vec<_>>(),
}));
if let Some(text) = p.0.clone()
&& let Some(first) = messages.first_mut()
{
first.content = vec![crate::core::provider::ContentBlock::Text {
text,
text_signature: None,
}];
}
});
Ok(())
}
}
fn ctx_probe_reset(rewrite: Option<String>) {
static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
ONCE.get_or_init(|| core::extensions::register_extension(ContextProbe));
CTX_PROBE.with(|p| *p.borrow_mut() = (rewrite, Vec::new()));
}
fn ctx_probe_seen() -> Vec<J> {
CTX_PROBE.with(|p| p.borrow().1.clone())
}
#[test]
fn context_with_system_sees_and_rewrites_prompt() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("ok", "a")]).await;
ctx_probe_reset(Some("rewritten prompt".to_string()));
let mut agent = session_agent(&base, &cwd);
agent.system_prompt = "original prompt".to_string();
let sent: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let cap = sent.clone();
agent.set_on_payload(Box::new(move |p: &mut J| {
cap.lock().unwrap().push(p.clone());
}));
let text = agent.prompt("go").await.unwrap();
assert_eq!(text, "ok");
let seen = ctx_probe_seen();
assert_eq!(seen.len(), 1, "每次 LLM 调用投递一次");
assert_eq!(seen[0]["roles"][0], "system", "system 消息应在 [0]");
assert_eq!(seen[0]["roles"][1], "user");
assert_eq!(
agent.system_prompt, "original prompt",
"会话自身的系统提示词不被改写(只影响本次请求)"
);
let payloads = sent.lock().unwrap().clone();
let msgs = payloads[0]["messages"].as_array().unwrap().clone();
assert_eq!(msgs[0]["role"], "system");
assert_eq!(
msgs[0]["content"], "rewritten prompt",
"请求体应使用边界改写后的系统提示词"
);
handle.await.unwrap();
let (base, handle) = mock_server_multi(vec![text_frame("ok2", "b")]).await;
ctx_probe_reset(None);
let mut agent = session_agent(&base, &cwd);
agent.system_prompt = "keep me".to_string();
let text = agent.prompt("go").await.unwrap();
assert_eq!(text, "ok2");
let seen = ctx_probe_seen();
assert_eq!(seen[0]["roles"][0], "system");
handle.await.unwrap();
});
}
fn probe_reset() {
static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
ONCE.get_or_init(|| core::extensions::register_extension(BoundaryProbe));
PROBE.with(|p| *p.borrow_mut() = ProbeState::default());
}
fn probe_set(kind: &str, mode: u8) {
PROBE.with(|p| {
let mut p = p.borrow_mut();
match kind {
"turn_end" => p.turn = mode,
"agent_before_settle" => p.settle = mode,
_ => unreachable!(),
}
});
}
fn probe_seen(kind: &str) -> Vec<J> {
PROBE.with(|p| {
p.borrow()
.seen
.iter()
.filter(|e| e["type"].as_str() == Some(kind))
.cloned()
.collect()
})
}
fn session_agent(base: &str, cwd: &str) -> Agent {
let mut agent = test_agent(base, cwd);
let dir =
std::env::temp_dir().join(format!("prux-boundary-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&dir).unwrap();
agent.session = Some(
crate::core::session_manager::Session::create_with(
crate::core::session_manager::SessionCreateOptions {
cwd: cwd.to_string(),
session_dir: Some(dir),
persist: true,
..Default::default()
},
)
.unwrap(),
);
agent
}
#[test]
fn rewind_last_user_input_rebuilds_agent_context() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let mut agent = session_agent("http://127.0.0.1:1", &cwd);
let session = agent.session.as_mut().unwrap();
session.append_message(&AgentMessage::user_text("first"));
let mut a1 = AgentMessage::user_text("a1");
a1.role = "assistant".into();
session.append_message(&a1);
session.append_message(&AgentMessage::user_text("second"));
let mut a2 = AgentMessage::user_text("a2");
a2.role = "assistant".into();
session.append_message(&a2);
agent.messages = agent.session.as_ref().unwrap().build_context_messages();
assert_eq!(agent.messages.len(), 4);
let removed = agent.rewind_last_user_input().unwrap();
assert_eq!(removed.as_deref(), Some("second"));
let texts: Vec<String> = agent.messages.iter().map(|m| m.text()).collect();
assert_eq!(texts, vec!["first", "a1"]);
let session = agent.session.as_mut().unwrap();
session.rewind_last_user_input().unwrap(); agent.messages = agent.session.as_ref().unwrap().build_context_messages();
assert!(agent.messages.is_empty());
assert_eq!(agent.rewind_last_user_input().unwrap(), None);
assert!(agent.messages.is_empty());
}
#[test]
fn exec_ctx_branch_entries_follow_current_branch() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let mut agent = session_agent("http://127.0.0.1:1", &cwd);
let (id1, id2, id3) = {
let session = agent.session.as_mut().unwrap();
let id1 = session.append_message(&AgentMessage::user_text("first"));
session.append_custom_entry("codemode-store", Some(json!({ "set": { "a": 1 } })));
let id2 = session.append_message(&AgentMessage::user_text("second"));
session.set_leaf(Some(&id1)).unwrap();
let id3 =
session.append_custom_entry("codemode-store", Some(json!({ "set": { "b": 2 } })));
(id1, id2, id3)
};
let ids = |entries: &[Value]| -> Vec<String> {
entries
.iter()
.filter_map(|e| e.get("id").and_then(|v| v.as_str()).map(str::to_string))
.collect()
};
let ctx = agent.exec_ctx();
let entries = ctx.session_branch_entries.expect("有会话时应带分支快照");
let branch = ids(&entries);
assert!(
branch.contains(&id1) && branch.contains(&id3),
"含本分支的条目: {branch:?}"
);
assert!(
!branch.contains(&id2),
"离开活动分支的条目不得出现: {branch:?}"
);
assert_eq!(
session_branch_store_payload(&entries),
Some(json!({ "b": 2 })),
"只看到本分支上的 store 条目(a=1 写在另一条分支上)"
);
{
let session = agent.session.as_mut().unwrap();
session.set_leaf(Some(&id2)).unwrap();
}
let ctx = agent.exec_ctx();
let entries = ctx.session_branch_entries.unwrap();
let branch = ids(&entries);
assert!(
branch.contains(&id1) && branch.contains(&id2),
"含本分支的条目: {branch:?}"
);
assert!(!branch.contains(&id3), "切回后不应看到另一分支: {branch:?}");
assert_eq!(
session_branch_store_payload(&entries),
Some(json!({ "a": 1 })),
"切回后看到的是本分支的 store 写入"
);
agent.session = None;
assert!(agent.exec_ctx().session_branch_entries.is_none());
}
fn session_branch_store_payload(entries: &[Value]) -> Option<Value> {
let mut merged = serde_json::Map::new();
for entry in entries {
if entry.get("customType").and_then(|v| v.as_str()) != Some("codemode-store") {
continue;
}
if let Some(set) = entry
.get("data")
.and_then(|d| d.get("set"))
.and_then(|v| v.as_object())
{
for (k, v) in set {
merged.insert(k.clone(), v.clone());
}
}
}
(!merged.is_empty()).then_some(Value::Object(merged))
}
#[test]
fn turn_end_boundary_fields_and_continue() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) =
mock_server_multi(vec![text_frame("First", "a"), text_frame("Second", "b")]).await;
let mut agent = session_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
probe_set("turn_end", 1); let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Second");
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
2
);
assert_eq!(types.iter().filter(|t| t.as_str() == "turn_end").count(), 2);
let seen = probe_seen("turn_end");
assert_eq!(seen.len(), 2, "每轮都应投递 turn_end 边界");
assert_eq!(seen[0]["turnIndex"], 0);
assert_eq!(seen[1]["turnIndex"], 1);
assert_eq!(seen[0]["outcome"], "completed");
assert!(
seen[0]["messageEntryId"]
.as_str()
.is_some_and(|s| !s.is_empty()),
"messageEntryId 应为落库条目 id: {}",
seen[0]
);
assert_eq!(seen[0]["toolResultEntryIds"].as_array().unwrap().len(), 0);
let emitted = events_of(&events, "turn_end")[0];
assert_eq!(emitted["turnIndex"], 0);
assert!(
emitted["messageEntryId"]
.as_str()
.is_some_and(|s| !s.is_empty())
);
assert_eq!(emitted["outcome"], "completed");
});
}
#[test]
fn turn_end_boundary_reports_tool_result_entry_ids() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![
tool_call_frame("read", r#"{\"path\":\"a.rs\"}"#, "call_1"),
text_frame("done", "b"),
])
.await;
let mut agent = session_agent(&base, &cwd);
let _ = agent.prompt("go").await.unwrap();
handle.await.unwrap();
let seen = probe_seen("turn_end");
assert_eq!(seen.len(), 2);
let ids = seen[0]["toolResultEntryIds"].as_array().unwrap();
assert_eq!(
ids.len(),
1,
"工具轮应有 1 个 toolResult entryId: {}",
seen[0]
);
assert!(ids[0].as_str().is_some_and(|s| !s.is_empty()));
assert_eq!(seen[0]["outcome"], "completed");
assert_eq!(seen[1]["toolResultEntryIds"].as_array().unwrap().len(), 0);
});
}
#[test]
fn turn_end_boundary_end_stops_run() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![
tool_call_frame("read", r#"{\"path\":\"a.rs\"}"#, "call_1"),
text_frame("Should not run", "b"),
])
.await;
let mut agent = session_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
probe_set("turn_end", 2); let _ = agent.prompt("go").await.unwrap();
handle.await.unwrap();
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
1
);
assert!(
!events
.iter()
.any(|e| e["message"]["content"][0]["text"] == "Should not run")
);
});
}
#[test]
fn turn_end_boundary_ignored_on_hard_exit() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("boom", "a")]).await;
let mut agent = session_agent(&base, &cwd);
agent.model.base_url = "http://127.0.0.1:1".to_string(); agent.auto_retry = false;
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
probe_set("turn_end", 1); let _ = agent.prompt("go").await;
handle.await.unwrap();
let seen = probe_seen("turn_end");
assert_eq!(seen.len(), 1, "硬退出也应投递一次 turn_end 边界");
assert_eq!(seen[0]["outcome"], "error");
assert!(
seen[0]["messageEntryId"]
.as_str()
.is_some_and(|s| !s.is_empty()),
"硬退出轮也应先落库再投边界: {}",
seen[0]
);
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
1,
"硬退出不得因边界决策续跑"
);
});
}
#[test]
fn settle_boundary_continue_runs_another_round() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) =
mock_server_multi(vec![text_frame("First", "a"), text_frame("Second", "b")]).await;
let mut agent = session_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
probe_set("agent_before_settle", 1);
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Second", "结算边界续跑后应取新一轮回复");
let events = collected.lock().unwrap().clone();
let types = ev_types(&events);
assert_eq!(
types.iter().filter(|t| t.as_str() == "turn_start").count(),
2,
"continue 应触发第二次 run_loop"
);
assert_eq!(
types.iter().filter(|t| t.as_str() == "agent_end").count(),
2
);
assert_eq!(
types
.iter()
.filter(|t| t.as_str() == "agent_settled")
.count(),
1
);
assert_eq!(types.last().map(String::as_str), Some("agent_settled"));
let seen = probe_seen("agent_before_settle");
assert_eq!(seen.len(), 2, "每次结算都应投递边界");
assert_eq!(seen[0]["outcome"], "completed");
assert_eq!(seen[0]["continuationIndex"], 0);
assert_eq!(seen[1]["continuationIndex"], 1, "续跑序号递增");
assert_eq!(seen[0]["agentScope"], "main");
assert_eq!(seen[0]["hasQueuedMessages"], false);
assert!(
!seen[0]["contextMessages"].as_array().unwrap().is_empty(),
"边界应携带当前上下文"
);
});
}
#[test]
fn settle_boundary_appends_messages_before_continuation() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
probe_reset();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) =
mock_server_multi(vec![text_frame("First", "a"), text_frame("Second", "b")]).await;
let mut agent = session_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
probe_set("agent_before_settle", 3);
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Second", "追加消息后应再跑一轮");
assert!(
agent
.messages
.iter()
.any(|m| m.role == "user" && m.text().contains("probe-nudge")),
"边界追加的消息应进入上下文"
);
let events = collected.lock().unwrap().clone();
let nudged = events.iter().any(|e| {
e["type"] == "message_end"
&& e["message"]["role"] == "user"
&& e["message"]["content"]
.as_array()
.is_some_and(|c| c.iter().any(|b| b["text"] == "probe-nudge"))
});
assert!(nudged, "追加的 user 消息应发出 message_end 事件");
});
}
#[test]
fn parallel_execution_runs_multiple_tools() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let tmp = std::env::temp_dir().join(format!("prux-par-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
std::fs::write(tmp.join("a.rs"), "AAA").unwrap();
std::fs::write(tmp.join("b.rs"), "BBB").unwrap();
let cwd = tmp.to_string_lossy().to_string();
let frames1 = vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"read","arguments":"{\"path\":\"a.rs\"}"}}]}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"c2","type":"function","function":{"name":"read","arguments":"{\"path\":\"b.rs\"}"}}]}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let frames2 = text_frame("Both read", "c");
let (base, handle) = mock_server_multi(vec![frames1, frames2]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("read both").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Both read");
let events = collected.lock().unwrap().clone();
assert_eq!(ev_types(&events).iter().filter(|t| t.as_str() == "tool_execution_start").count(), 2);
assert_eq!(ev_types(&events).iter().filter(|t| t.as_str() == "tool_execution_end").count(), 2);
assert_eq!(
events.iter()
.filter(|e| e["type"] == "message_start" && e["message"]["role"] == "toolResult")
.count(),
2
);
});
}
#[test]
fn parallel_path_read_keeps_image_attachment() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let tmp =
std::env::temp_dir().join(format!("prux-img-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
let png = tmp.join("pic.png");
let mut buf = Vec::new();
image::RgbaImage::from_pixel(4, 4, image::Rgba([1, 2, 3, 255]))
.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)
.unwrap();
std::fs::write(&png, buf).unwrap();
let cwd = tmp.to_string_lossy().to_string();
let frames1 = vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"read","arguments":"{\"path\":\"pic.png\"}"}}]}}]}"#,
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let (base, handle) = mock_server_multi(vec![frames1, text_frame("done", "c")]).await;
let mut agent =
test_agent_with_input(&base, &cwd, vec!["text".into(), "image".into()]);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
agent.prompt("look").await.unwrap();
handle.await.unwrap();
let events = collected.lock().unwrap().clone();
let tr = events
.iter()
.find(|e| {
e["type"] == "message_start" && e["message"]["role"] == "toolResult"
})
.expect("toolResult message_start");
let content = tr["message"]["content"].as_array().unwrap();
let img = content
.iter()
.find(|b| b["type"] == "image")
.unwrap_or_else(|| panic!("read 应回传 image 块,实际: {content:?}"));
assert_eq!(img["mimeType"], "image/png");
assert!(
!content.iter().any(|b| b["type"] == "text"
&& b["text"]
.as_str()
.unwrap_or("")
.contains("not supported by this model")),
"不应出现 model does not support images 占位"
);
});
}
#[test]
fn block_images_drops_tool_result_attachments() {
let _ad = crate::test_support::AgentDirGuard::temp();
let tmp =
std::env::temp_dir().join(format!("prux-block-img-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
let mut agent = test_agent_with_input(
"http://127.0.0.1:1",
&tmp.to_string_lossy(),
vec!["text".into(), "image".into()],
);
let fin = FinalizedToolCall {
call: ContentBlock::ToolCall {
id: "c1".into(),
name: "codemode".into(),
arguments: json!({}),
thought_signature: None,
namespace: None,
},
name: "codemode".into(),
args: json!({}),
text: "generated an image".into(),
details: None,
usage: None,
terminate: false,
is_error: false,
attachments: vec![crate::core::tools::ToolResultAttachment {
data_base64: "AAAA".into(),
mime_type: "image/png".into(),
original_size: None,
converted_from: None,
}],
content_override: None,
duration_ms: None,
};
let mut out = Vec::new();
agent.push_tool_result_message(&fin, Some(1), &mut out);
let content = &out[0].1.content;
assert!(
content
.iter()
.any(|b| matches!(b, ContentBlock::Image { .. })),
"默认应回传图片块: {content:?}"
);
settings_manager::write_settings_images_block_images(true).unwrap();
let mut out = Vec::new();
agent.push_tool_result_message(&fin, Some(1), &mut out);
settings_manager::write_settings_images_block_images(false).unwrap();
let content = &out[0].1.content;
assert!(
!content
.iter()
.any(|b| matches!(b, ContentBlock::Image { .. })),
"blockImages 打开时不应有图片块: {content:?}"
);
assert!(
content.iter().any(|b| matches!(
b,
ContentBlock::Text { text, .. } if text.contains("images.blockImages")
)),
"应给出被阻止的文本占位: {content:?}"
);
}
#[test]
fn child_agent_is_state_isolated() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, _handle) = mock_server_multi(vec![text_frame("x", "x")]).await;
let mut agent = test_agent(&base, &cwd);
agent.messages.push(AgentMessage::user_text("parent msg"));
let mut child = agent.child_agent(None);
assert!(child.messages.is_empty());
assert!(child.session.is_none());
assert_eq!(child.cwd, agent.cwd);
assert_eq!(child.model.base_url, agent.model.base_url);
agent
.runtime_steer_inbox
.lock()
.unwrap()
.push_back(AgentMessage::user_text("parent steer"));
assert!(child.runtime_steer_inbox.lock().unwrap().is_empty());
child.messages.push(AgentMessage::user_text("child msg"));
assert_eq!(agent.messages.len(), 1);
});
}
#[test]
fn run_sub_agent_returns_isolated_result() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("sub done", "s")]).await;
let agent = test_agent(&base, &cwd);
let text = agent.run_sub_agent("do sub task").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "sub done");
assert!(agent.messages.is_empty());
});
}
#[test]
fn spawn_sub_agents_run_concurrently() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) =
mock_server_multi(vec![text_frame("one", "1"), text_frame("two", "2")]).await;
let agent = test_agent(&base, &cwd);
let h1 = agent.spawn_sub_agent("task one".to_string());
let h2 = agent.spawn_sub_agent("task two".to_string());
let r1 = h1.await.unwrap().unwrap();
let r2 = h2.await.unwrap().unwrap();
handle.await.unwrap();
let mut got = vec![r1, r2];
got.sort();
assert_eq!(got, vec!["one".to_string(), "two".to_string()]);
});
}
#[test]
fn tool_execution_event_payloads_match_pi() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let tmp =
std::env::temp_dir().join(format!("prux-test-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
std::fs::write(tmp.join("a.rs"), "fn main() {}").unwrap();
let cwd = tmp.to_string_lossy().to_string();
let frames1 = tool_call_frame("read", r#"{\"path\":\"a.rs\"}"#, "call_1");
let frames2 = text_frame("Read done", "cmpl_2");
let (base, handle) = mock_server_multi(vec![frames1, frames2]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let _ = agent.prompt("read it").await.unwrap();
handle.await.unwrap();
let events = collected.lock().unwrap().clone();
let start = events
.iter()
.find(|e| e["type"] == "tool_execution_start")
.unwrap();
assert_eq!(start["toolCallId"], "call_1");
assert_eq!(start["toolName"], "read");
assert_eq!(start["args"], json!({ "path": "a.rs" }));
let end = events
.iter()
.find(|e| e["type"] == "tool_execution_end")
.unwrap();
assert_eq!(end["toolCallId"], "call_1");
assert_eq!(end["isError"], false);
assert!(
end["durationMs"].as_u64().is_some(),
"执行过的工具应带 durationMs: {end}"
);
let result = end["result"].as_object().unwrap();
assert!(result.contains_key("content"), "result 必须是对象");
assert!(result.contains_key("details"), "result 必须含 details");
let text = result["content"][0]["text"].as_str().unwrap();
assert!(text.contains("fn main() {}"), "got: {text}");
let ends = events.iter().filter(|e| e["type"] == "agent_end").count();
assert_eq!(ends, 1, "agent_end 只发一次");
});
}
#[test]
fn unknown_tool_reports_pi_not_found() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let frames1 = tool_call_frame("no_such_tool", r#"{}"#, "call_1");
let frames2 = text_frame("Recovered", "cmpl_2");
let (base, handle) = mock_server_multi(vec![frames1, frames2]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("hi").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "Recovered");
let events = collected.lock().unwrap().clone();
let turn_end = events_of(&events, "turn_end");
let first = turn_end[0]["toolResults"].as_array().unwrap();
assert_eq!(first.len(), 1);
assert_eq!(first[0]["isError"], true);
assert_eq!(first[0]["toolName"], "no_such_tool");
let err_text = first[0]["content"][0]["text"].as_str().unwrap();
assert_eq!(err_text, "Tool no_such_tool not found");
assert!(
first[0]["durationMs"].is_null(),
"immediate 调用不应带 durationMs: {:?}",
first[0]
);
assert_eq!(
events
.iter()
.filter(|e| e["type"] == "tool_execution_start")
.count(),
1
);
assert_eq!(
events
.iter()
.filter(|e| e["type"] == "tool_execution_end")
.count(),
1
);
assert_eq!(
events.iter().filter(|e| e["type"] == "agent_end").count(),
1
);
});
}
#[test]
fn before_tool_block_terminate_stops_batch_when_all_terminate() {
struct BlockReads;
impl crate::core::extensions::Extension for BlockReads {
fn name(&self) -> &str {
"block-reads-test"
}
fn tools(&self) -> Vec<crate::core::extensions::ExtensionTool> {
vec![crate::core::extensions::ExtensionTool::simple(
"blk_read",
"test",
json!({ "type": "object", "properties": {} }),
"",
)]
}
fn hooks(&self) -> Vec<crate::core::extensions::ExtensionHook> {
vec![crate::core::extensions::ExtensionHook::BeforeToolCall]
}
fn execute_tool(
&self,
name: &str,
_args: &Value,
) -> std::result::Result<tools::ToolResult, tools::ToolError> {
if name == "blk_read" {
Ok(tools::ToolResult::text("executed anyway"))
} else {
Err(tools::ToolError("unknown".to_string()))
}
}
fn before_tool_call_ext(
&self,
name: &str,
_args: &Value,
assistant_message: Option<&crate::core::provider::AgentMessage>,
context: &[crate::core::provider::AgentMessage],
_all_args: &Value,
) -> std::result::Result<crate::core::extensions::BeforeToolCallOutcome, tools::ToolError>
{
assert!(
assistant_message.is_some(),
"beforeToolCall 应收到 assistantMessage"
);
assert!(!context.is_empty(), "beforeToolCall 应收到当前上下文");
if name == "blk_read" {
Ok(crate::core::extensions::BeforeToolCallOutcome {
block: true,
reason: Some("test block".to_string()),
terminate: true,
args: None,
})
} else {
Ok(crate::core::extensions::BeforeToolCallOutcome::allow())
}
}
}
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(BlockReads);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let frames1 = tool_call_frame("blk_read", r#"{}"#, "call_1");
let (base, handle) = mock_server_multi(vec![frames1]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let _ = agent.prompt("hi").await.unwrap();
handle.await.unwrap();
let events = collected.lock().unwrap().clone();
assert_eq!(events.iter().filter(|e| e["type"] == "turn_end").count(), 1);
let turn_end = events_of(&events, "turn_end")[0];
let tr = turn_end["toolResults"].as_array().unwrap();
assert_eq!(tr.len(), 1);
assert_eq!(tr[0]["isError"], true);
assert_eq!(tr[0]["content"][0]["text"], "test block");
assert_eq!(
events.iter().filter(|e| e["type"] == "agent_end").count(),
1
);
assert!(crate::core::extensions::unregister_extension(
"block-reads-test"
));
});
}
#[test]
fn tool_duration_excludes_after_tool_call_hook() {
use crate::core::extensions::{AfterToolCallOutcome, ExtensionHook, ExtensionTool};
struct SlowHook;
impl crate::core::extensions::Extension for SlowHook {
fn name(&self) -> &str {
"slow-hook-test"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![ExtensionTool::simple(
"slow_probe",
"test",
json!({ "type": "object", "properties": {} }),
"",
)]
}
fn hooks(&self) -> Vec<ExtensionHook> {
vec![ExtensionHook::AfterToolCall]
}
fn execute_tool(
&self,
name: &str,
_args: &Value,
) -> std::result::Result<tools::ToolResult, tools::ToolError> {
if name == "slow_probe" {
Ok(tools::ToolResult::text("ok"))
} else {
Err(tools::ToolError("unknown".to_string()))
}
}
fn after_tool_call_ext(
&self,
name: &str,
_args: &Value,
result_text: &str,
_is_error: bool,
) -> std::result::Result<AfterToolCallOutcome, tools::ToolError> {
if name != "slow_probe" {
return Ok(AfterToolCallOutcome {
text: result_text.to_string(),
content: None,
details: None,
is_error: None,
usage: None,
terminate: None,
});
}
std::thread::sleep(std::time::Duration::from_millis(400));
Ok(AfterToolCallOutcome {
text: format!("{result_text} (hooked)"),
content: None,
details: None,
is_error: None,
usage: None,
terminate: None,
})
}
}
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(SlowHook);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![
tool_call_frame("slow_probe", "{}", "call_1"),
text_frame("done", "cmpl_d"),
])
.await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let _ = agent.prompt("hi").await.unwrap();
handle.await.unwrap();
let events = collected.lock().unwrap().clone();
let tr = events_of(&events, "turn_end")[0]["toolResults"]
.as_array()
.unwrap()
.clone();
assert_eq!(tr.len(), 1);
assert_eq!(tr[0]["content"][0]["text"], "ok (hooked)");
let dur = tr[0]["durationMs"].as_u64().unwrap();
assert!(
dur < 300,
"Took 不得把 after_tool_call 钩子的 400ms 算进去,实际 {dur}ms"
);
assert!(crate::core::extensions::unregister_extension(
"slow-hook-test"
));
});
}
fn materialize_expect_err(tpl: &AgentTemplate, spec: core::extensions::SubAgentSpec) -> String {
match materialize_sub_agent(tpl, spec) {
Ok(_) => panic!("expected materialization to fail"),
Err(e) => e,
}
}
#[test]
fn subagent_materialize_narrows_tools_and_blocks_recursion() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let agent = test_agent("http://127.0.0.1:1", "/tmp");
let tpl = agent.snapshot_template();
let full =
materialize_sub_agent(&tpl, core::extensions::SubAgentSpec::new("x")).unwrap();
let names = full.agent.get_active_tools();
for sub in [
"Agent",
"SubagentWorkflow",
"get_subagent_result",
"steer_subagent",
] {
assert!(
!names.iter().any(|n| n == sub),
"递归防护失效:child 拿到 {sub}"
);
}
assert!(
agent.get_all_tools().iter().any(|n| n == "Agent"),
"父级应能看到 Agent 工具"
);
assert!(names.iter().any(|n| n == "read"));
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.tools = Some(vec!["read".to_string(), "grep".to_string()]);
let narrowed = materialize_sub_agent(&tpl, spec).unwrap();
let got = narrowed.agent.get_active_tools();
assert_eq!(got, vec!["read".to_string(), "grep".to_string()]);
let mut bad = core::extensions::SubAgentSpec::new("x");
bad.tools = Some(vec!["definitely_not_a_tool".to_string()]);
let err = materialize_expect_err(&tpl, bad);
assert!(err.contains("none of the configured tools"), "{err}");
let mut nested = core::extensions::SubAgentSpec::new("x");
nested.tools = Some(vec!["read".to_string(), "Agent".to_string()]);
let allowed = materialize_sub_agent(&tpl, nested).unwrap();
let got = allowed.agent.get_active_tools();
assert!(got.iter().any(|n| n == "Agent"), "{got:?}");
assert!(!got.iter().any(|n| n == "AgentWorkflow"), "{got:?}");
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn materialized_child_abort_handle_stops_the_run() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let tmp = std::env::temp_dir().join(format!(
"prux-child-abort-{}",
crate::utils::time::now_ms()
));
std::fs::create_dir_all(&tmp).unwrap();
std::fs::write(tmp.join("a.rs"), "fn main() {}").unwrap();
let cwd = tmp.to_string_lossy().to_string();
let slot: Arc<std::sync::Mutex<Option<Arc<AtomicBool>>>> =
Arc::new(std::sync::Mutex::new(None));
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let count = requests.clone();
let slot_srv = slot.clone();
let server = tokio::spawn(async move {
let rounds = [
tool_call_frame("read", r#"{\"path\":\"a.rs\"}"#, "call_1"),
text_frame("done", "cmpl_2"),
];
for frames in rounds {
let Ok(Ok((mut socket, _))) = tokio::time::timeout(
std::time::Duration::from_millis(2_000),
listener.accept(),
)
.await
else {
break;
};
if count.fetch_add(1, Ordering::SeqCst) == 0
&& let Some(flag) = slot_srv.lock().unwrap().clone()
{
flag.store(true, Ordering::Relaxed);
}
read_request(&mut socket).await;
socket
.write_all(
b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n",
)
.await
.unwrap();
for f in frames {
socket
.write_all(format!("data: {f}\n\n").as_bytes())
.await
.unwrap();
}
socket.shutdown().await.unwrap();
}
});
let agent = test_agent(&base, &cwd);
let tpl = agent.snapshot_template();
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.tools = Some(vec!["read".to_string()]);
let mut runner = materialize_sub_agent(&tpl, spec).unwrap();
*slot.lock().unwrap() = Some(runner.controls().abort.clone());
let res = tokio::time::timeout(
std::time::Duration::from_secs(3),
tokio::spawn(async move { runner.run("go".into()).await }),
)
.await
.expect("置位中止后 child 应在下一轮界停下")
.expect("run task 不应 panic");
assert!(res.is_ok(), "协作式取消走 Ok 路径: {res:?}");
assert_eq!(
requests.load(Ordering::SeqCst),
1,
"中止后不得再发起下一次 LLM 调用(说明句柄没作用在运行中的 child 上)"
);
server.abort();
});
}
fn test_injected_tool(
name: &str,
parameters: Value,
) -> (core::extensions::InjectedChildTool, Arc<Mutex<Vec<Value>>>) {
let calls: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let sink = calls.clone();
let tool = core::extensions::InjectedChildTool {
meta: core::extensions::ExtensionTool {
exposure: crate::core::extensions::ToolExposure::Direct,
namespace: None,
annotations: crate::core::extensions::ToolAnnotations::default(),
output_schema: None,
name: name.to_string(),
description: format!("{name} injected description"),
label: Some("Injected".to_string()),
parameters,
snippet: format!("{name}: injected snippet"),
prompt_guidelines: vec![format!("{name} injected guideline")],
constrained_sampling: false,
render_shell: None,
execution_mode: None,
prepare_arguments: None,
grammar_sampling: None,
},
handler: Arc::new(move |args: Value| {
sink.lock().unwrap().push(args);
Ok(ToolResult::text("Recorded."))
}),
};
(tool, calls)
}
#[test]
fn tool_index_prefers_injected_over_builtin() {
let (tool, _) = test_injected_tool(
"read",
json!({ "type": "object", "properties": { "only_injected": { "type": "string" } } }),
);
let injected = vec![tool];
let idx = ToolIndex::resolve(&injected, "read");
assert!(idx.exists());
assert!(idx.is_builtin(), "名字仍是内置的,但元数据归注入");
assert!(idx.injected_handler().is_some(), "命中注入就该走 handler");
let params = idx.parameters().expect("parameters");
assert!(
params["properties"].get("only_injected").is_some(),
"注入的 schema 必须胜过内置的:{params}"
);
let base = ToolIndex::resolve(&injected, "read");
assert!(base.is_builtin());
assert!(ToolIndex::resolve(&injected, "bash").exists());
assert!(!ToolIndex::resolve(&injected, "definitely_not_a_tool").exists());
assert!(
ToolIndex::resolve(&injected, "definitely_not_a_tool")
.parameters()
.is_none()
);
}
#[test]
fn tool_index_carries_injected_prepare_arguments_and_mode() {
let (mut tool, _) = test_injected_tool("read", json!({ "type": "object" }));
tool.meta.execution_mode = Some(ToolExecutionMode::Sequential);
tool.meta.prepare_arguments = Some(|args: &Value| {
if args.get("only_injected_prepared").is_some() {
json!({ "prepared": true })
} else {
args.clone()
}
});
let injected = vec![tool];
let idx = ToolIndex::resolve(&injected, "read");
assert_eq!(idx.execution_mode(), Some(ToolExecutionMode::Sequential));
assert_eq!(
idx.prepare_arguments(&json!({ "only_injected_prepared": 1 })),
json!({ "prepared": true })
);
}
#[test]
fn injected_tools_join_child_table_and_survive_scoping() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let agent = test_agent("http://127.0.0.1:1", "/tmp");
let tpl = agent.snapshot_template();
let (tool, _) = test_injected_tool("StructuredOutput", json!({ "type": "object" }));
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.tools = Some(vec!["read".to_string()]);
spec.disallowed_tools = vec!["StructuredOutput".to_string()];
spec.isolated = true;
spec.extensions = Some(Vec::new());
spec.injected_tools = vec![tool];
let child = materialize_sub_agent(&tpl, spec).unwrap();
let names = child.agent.get_active_tools();
assert!(names.iter().any(|n| n == "read"));
assert!(
names.iter().any(|n| n == "StructuredOutput"),
"注入工具必须无条件并入:{names:?}"
);
assert!(
child
.agent
.system_prompt
.contains("StructuredOutput: injected snippet"),
"注入工具的 snippet 必须进 Available tools"
);
assert!(
child
.agent
.system_prompt
.contains("StructuredOutput injected guideline"),
"注入工具的 guideline 必须进 Guidelines"
);
let universe = child.agent.get_all_tools();
assert!(universe.iter().any(|n| n == "read"), "{universe:?}");
assert!(
!universe.iter().any(|n| n == "StructuredOutput"),
"{universe:?}"
);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn subagent_materialize_honours_isolated_and_disallowed_tools() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let agent = test_agent("http://127.0.0.1:1", "/tmp");
let tpl = agent.snapshot_template();
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.isolated = true;
let isolated = materialize_sub_agent(&tpl, spec).unwrap();
let names = isolated.agent.get_active_tools();
for ext_tool in [
"Agent",
"SubagentWorkflow",
"get_subagent_result",
"steer_subagent",
] {
assert!(
!names.iter().any(|n| n == ext_tool),
"isolated 下不应有扩展工具 {ext_tool}"
);
}
assert!(
names.iter().any(|n| n == "read"),
"内置工具应保留: {names:?}"
);
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.disallowed_tools = vec!["read".to_string(), "grep".to_string()];
let filtered = materialize_sub_agent(&tpl, spec).unwrap();
let names = filtered.agent.get_active_tools();
assert!(!names.iter().any(|n| n == "read"), "read 应被剔除");
assert!(!names.iter().any(|n| n == "grep"), "grep 应被剔除");
assert!(names.iter().any(|n| n == "bash"));
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.isolated = true;
spec.tools = Some(vec!["read".to_string(), "Agent".to_string()]);
let both = materialize_sub_agent(&tpl, spec).unwrap();
assert_eq!(both.agent.get_active_tools(), vec!["read".to_string()]);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn subagent_materialize_applies_extension_scope_and_skill_clearing() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let agent = test_agent("http://127.0.0.1:1", "/tmp");
let tpl = agent.snapshot_template();
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.extensions = Some(vec!["some-other-ext".to_string()]);
let scoped = materialize_sub_agent(&tpl, spec).unwrap();
let names = scoped.agent.get_active_tools();
assert!(!names.iter().any(|n| n == "Agent"), "{names:?}");
assert!(
names.iter().any(|n| n == "read"),
"内置不受扩展作用域影响: {names:?}"
);
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.exclude_extensions = vec!["subagent".to_string()];
let scoped = materialize_sub_agent(&tpl, spec).unwrap();
let names = scoped.agent.get_active_tools();
assert!(!names.iter().any(|n| n == "Agent"), "{names:?}");
assert!(!names.iter().any(|n| n == "steer_subagent"), "{names:?}");
let mut parent = test_agent("http://127.0.0.1:1", "/tmp");
parent.skills = vec![crate::core::skills::Skill {
name: "demo".to_string(),
description: "demo skill".to_string(),
path: std::path::PathBuf::from("/tmp/demo/SKILL.md"),
base_dir: std::path::PathBuf::from("/tmp/demo"),
disable_model_invocation: false,
source: "/tmp".to_string(),
}];
let tpl2 = parent.snapshot_template();
let kept =
materialize_sub_agent(&tpl2, core::extensions::SubAgentSpec::new("x")).unwrap();
assert_eq!(kept.agent.skills.len(), 1, "默认继承父级技能索引");
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.clear_skills = true;
let cleared = materialize_sub_agent(&tpl2, spec).unwrap();
assert!(
cleared.agent.skills.is_empty(),
"skills:false 应清空继承的技能索引"
);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn subagent_materialize_switches_prompt_mode() {
let agent = test_agent("http://127.0.0.1:1", "/tmp");
let mut parent = agent;
parent.rebuild_ctx.context_files =
vec![("AGENTS.md".to_string(), "project rules".to_string())];
let tpl = parent.snapshot_template();
let mut spec = core::extensions::SubAgentSpec::new("Explore");
spec.system_prompt = Some("ROLE ONLY".to_string());
spec.clear_context_files = true;
let child = materialize_sub_agent(&tpl, spec).unwrap();
assert!(child.agent.system_prompt.starts_with("ROLE ONLY"));
assert!(!child.agent.system_prompt.contains("project rules"));
let mut keep = core::extensions::SubAgentSpec::new("Explore");
keep.system_prompt = Some("ROLE ONLY".to_string());
let child2 = materialize_sub_agent(&tpl, keep).unwrap();
assert!(child2.agent.system_prompt.contains("project rules"));
let mut app = core::extensions::SubAgentSpec::new("general-purpose");
app.append_system_prompt = Some("BRIDGE".to_string());
let child3 = materialize_sub_agent(&tpl, app).unwrap();
assert!(child3.agent.system_prompt.contains("BRIDGE"));
assert!(child3.agent.system_prompt.contains("project rules"));
assert!(!child3.agent.system_prompt.starts_with("BRIDGE"));
}
#[test]
fn subagent_materialize_inherits_context_only_when_asked() {
let mut agent = test_agent("http://127.0.0.1:1", "/tmp");
agent.messages.push(AgentMessage::user_text("parent turn"));
let tpl = agent.snapshot_template();
let child = materialize_sub_agent(&tpl, core::extensions::SubAgentSpec::new("x")).unwrap();
assert!(child.agent.messages.is_empty());
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.inherit_context = true;
let forked = materialize_sub_agent(&tpl, spec).unwrap();
assert_eq!(forked.agent.messages.len(), 1);
assert_eq!(forked.agent.messages[0].text(), "parent turn");
}
#[test]
fn subagent_materialize_rejects_unresolvable_model() {
let agent = test_agent("http://127.0.0.1:1", "/tmp");
let tpl = agent.snapshot_template();
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.model = Some("no-such-provider/no-such-model".to_string());
let err = materialize_expect_err(&tpl, spec);
assert!(err.contains("not available"), "{err}");
assert!(err.contains("no-such-provider/no-such-model"), "{err}");
let mut bad = core::extensions::SubAgentSpec::new("x");
bad.model = Some("just-a-model".to_string());
let err = materialize_expect_err(&tpl, bad);
assert!(err.contains("provider/modelId"), "{err}");
}
#[test]
fn subagent_materialize_wires_graceful_max_turns() {
let agent = test_agent("http://127.0.0.1:1", "/tmp");
let tpl = agent.snapshot_template();
let plain = materialize_sub_agent(&tpl, core::extensions::SubAgentSpec::new("x")).unwrap();
assert!(plain.agent.finish_turn.is_none());
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.max_turns = Some(2);
let mut child = materialize_sub_agent(&tpl, spec).unwrap();
let ctx = TurnContext {
message: AgentMessage::user_text(""),
tool_results: Vec::new(),
context: Vec::new(),
new_messages: Vec::new(),
};
let stop = |a: &mut Agent| {
a.finish_turn
.as_ref()
.and_then(|m| {
let mut g = m.lock().unwrap();
let f: &mut dyn FnMut(&TurnContext) -> Option<FinishTurnAction> = g.as_mut();
f(&ctx)
})
.is_some_and(|d| d == FinishTurnAction::End)
};
assert!(!stop(&mut child.agent), "第 1 轮不应停机");
assert!(
child.agent.runtime_steer_inbox.lock().unwrap().is_empty(),
"第 1 轮不应注入 wrap-up"
);
assert!(!stop(&mut child.agent), "到限轮仍应给宽限");
let injected: Vec<String> = child
.agent
.runtime_steer_inbox
.lock()
.unwrap()
.iter()
.map(|m| m.text())
.collect();
assert_eq!(injected.len(), 1, "到限轮应注入一条 wrap-up");
assert!(injected[0].contains("Wrap up"), "{injected:?}");
for turn in 3..(2 + SUBAGENT_GRACE_TURNS) {
assert!(!stop(&mut child.agent), "第 {turn} 轮仍在宽限期内");
}
assert!(
stop(&mut child.agent),
"第 {} 轮应停机",
2 + SUBAGENT_GRACE_TURNS
);
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.max_turns = Some(2);
let child = materialize_sub_agent(&tpl, spec).unwrap();
let err_ctx = TurnContext {
message: AgentMessage {
stop_reason: Some("error".to_string()),
..AgentMessage::user_text("")
},
tool_results: Vec::new(),
context: Vec::new(),
new_messages: Vec::new(),
};
let decision = child.agent.finish_turn.as_ref().and_then(|m| {
let mut g = m.lock().unwrap();
let f: &mut dyn FnMut(&TurnContext) -> Option<FinishTurnAction> = g.as_mut();
f(&err_ctx)
});
assert_eq!(decision, None, "硬退出决策被忽略");
assert!(
child.agent.runtime_steer_inbox.lock().unwrap().is_empty(),
"硬退出不应注入 wrap-up"
);
}
#[test]
fn subagent_materialize_resumes_existing_session() {
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let (base, handle) = mock_server_multi(vec![text_frame("first answer", "c1")]).await;
let mut parent = test_agent(&base, "/tmp");
let dir = std::env::temp_dir().join(format!(
"prux-subagent-resume-{}",
crate::utils::time::now_ms()
));
std::fs::create_dir_all(&dir).unwrap();
parent.session = Some(
crate::core::session_manager::Session::create_with(
crate::core::session_manager::SessionCreateOptions {
cwd: "/tmp".to_string(),
session_dir: Some(dir.clone()),
persist: true,
..Default::default()
},
)
.unwrap(),
);
parent.prompt("remember this").await.unwrap();
handle.await.unwrap();
let path = parent
.session
.as_ref()
.and_then(|s| s.get_session_file())
.map(|p| p.to_string_lossy().to_string())
.expect("session file");
let tpl = parent.snapshot_template();
let mut spec = core::extensions::SubAgentSpec::new("x");
spec.resume_session_path = Some(path.clone());
spec.inherit_context = false;
let resumed = materialize_sub_agent(&tpl, spec).unwrap();
let texts: Vec<String> = resumed.agent.messages.iter().map(|m| m.text()).collect();
assert!(
texts.iter().any(|t| t.contains("remember this")),
"未回灌上下文: {texts:?}"
);
assert!(
texts.iter().any(|t| t.contains("first answer")),
"未回灌回答: {texts:?}"
);
assert_eq!(
resumed
.agent
.session
.as_ref()
.and_then(|s| s.get_session_file())
.map(|p| p.to_string_lossy().to_string())
.as_deref(),
Some(path.as_str())
);
});
}
#[test]
fn subagent_background_spawn_notifies_and_persists_child_session() {
use crate::core::extensions::Extension as _;
let _lock = subagent_test_lock();
let agent_home = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
while crate::core::extensions::take_pending_ui().is_some() {}
let (base, handle) = mock_server_multi(vec![text_frame("child work done", "c1")]).await;
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let mut parent = test_agent(&base, &cwd);
let sess_dir =
session_manager::default_session_dir(&cwd, &settings_manager::agent_dir());
let parent_session = Session::create_with(session_manager::SessionCreateOptions {
cwd: cwd.clone(),
session_dir: Some(sess_dir.clone()),
persist: true,
..Default::default()
})
.unwrap();
let parent_id = parent_session.session_id.clone();
parent.session = Some(parent_session);
let tpl = std::sync::Arc::new(parent.snapshot_template());
let ctx = make_exec_ctx(
&cwd,
&tpl,
&std::sync::Arc::new(AtomicBool::new(false)),
&None,
None,
None,
);
let ext = crate::extensions::subagent::Subagent;
let out = ext
.execute_tool_async(
"Agent".to_string(),
json!({
"prompt": "do the thing",
"description": "do thing",
"subagent_type": "general-purpose",
"run_in_background": true
}),
ctx.clone(),
)
.await
.expect("Agent tool");
assert!(out.text.contains("launched"), "{}", out.text);
assert!(
out.text.contains("do not sleep or poll"),
"后台结果应明确禁止轮询: {}",
out.text
);
let id = out.details.as_ref().unwrap()["id"]
.as_str()
.unwrap()
.to_string();
assert_eq!(out.details.as_ref().unwrap()["background"], true);
let mut notified: Option<String> = None;
for _ in 0..500 {
while let Some(r) = crate::core::extensions::take_pending_ui() {
if let core::extensions::ExtensionUiRequest::Continuation { message } = r {
notified = Some(message.text());
}
}
if notified.is_some() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let text = notified.expect("后台完成必须投递一条 Continuation");
assert!(text.contains("<subagent_result"), "{text}");
assert!(text.contains("status=\"completed\""), "{text}");
assert!(text.contains(&id), "{text}");
assert!(text.contains("child work done"), "{text}");
let res = ext
.execute_tool_async(
"get_subagent_result".to_string(),
json!({ "agent_id": id, "wait": true }),
ctx.clone(),
)
.await
.expect("get_subagent_result");
assert!(res.text.contains("completed"), "{}", res.text);
assert!(res.text.contains("child work done"), "{}", res.text);
let sessions = session_manager::list_sessions(&sess_dir);
let child = sessions
.iter()
.filter_map(|p| Session::open(&p.to_string_lossy()).ok())
.find(|s| s.parent_session_id() == Some(parent_id.as_str()))
.expect("子会话应已落盘并带 parentSessionId");
assert_eq!(child.cwd, cwd);
let _ = agent_home;
handle.await.unwrap();
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn subagent_extension_runs_agent_tool() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let f1 = tool_call_frame(
"Agent",
r#"{\"prompt\":\"do x\",\"description\":\"do x\",\"subagent_type\":\"general-purpose\",\"run_in_background\":false}"#,
"call_1",
);
let f2 = text_frame("child done", "cmpl_c");
let f3 = text_frame("parent final", "cmpl_p");
let (base, handle) = mock_server_multi(vec![f1, f2, f3]).await;
let mut agent = test_agent(&base, &cwd);
let collected: std::sync::Arc<std::sync::Mutex<Vec<J>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("hi").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "parent final");
let events = collected.lock().unwrap().clone();
let turn_end = events_of(&events, "turn_end");
let tr = turn_end[0]["toolResults"].as_array().unwrap();
assert_eq!(tr.len(), 1);
assert_eq!(tr[0]["isError"], false);
let t = tr[0]["content"][0]["text"].as_str().unwrap();
assert!(t.contains("child done"), "got: {t}");
assert!(t.contains("[subagent "), "状态头缺失: {t}");
assert!(t.contains("completed"), "状态缺失: {t}");
assert_eq!(tr[0]["details"]["type"], "general-purpose");
assert_eq!(tr[0]["details"]["status"], "completed");
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
fn json_string_escape(s: &str) -> String {
let quoted = serde_json::to_string(s).unwrap();
quoted[1..quoted.len() - 1].to_string()
}
#[test]
fn autoresearch_job_loops_until_the_completion_marker() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let f1 = text_frame("first change [continue]", "cmpl_1");
let f2 = text_frame("verified [goal-complete]", "cmpl_2");
let (base, handle) = mock_server_multi(vec![f1, f2]).await;
let agent = test_agent(&base, &cwd);
let id = crate::extensions::subagent::autoresearch::start_with(
agent.exec_ctx(),
tokio::runtime::Handle::current(),
"make the renderer faster",
None,
);
let snapshot = await_autoresearch_terminal(id).await;
handle.await.unwrap();
assert_eq!(
snapshot.status,
crate::extensions::subagent::autoresearch::JobStatus::Completed,
"{snapshot:?}"
);
assert_eq!(snapshot.iteration, 2, "{snapshot:?}");
assert!(
snapshot
.latest_result
.as_deref()
.is_some_and(|r| r.contains("[goal-complete]")),
"{snapshot:?}"
);
let agent_id = snapshot.agent_id.clone().expect("作业应留下子代理记录 id");
let path = crate::extensions::subagent::agent_session_path(&agent_id)
.expect("子代理记录应带子会话路径(persist_session=true)");
let text = std::fs::read_to_string(&path).expect("子会话文件应落盘");
assert!(
text.contains("first change") && text.contains("verified"),
"同一个子会话里应同时有两轮的回答({path}): {text}"
);
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
let mut notes: Vec<(String, crate::core::extensions::UiNotifyLevel)> = Vec::new();
let mut continuation = false;
loop {
while let Some(req) = crate::core::extensions::take_pending_ui() {
match req {
crate::core::extensions::ExtensionUiRequest::Notify { text, level } => {
notes.push((text, level));
}
crate::core::extensions::ExtensionUiRequest::NotifyRich {
spans,
level,
} => {
notes.push((
spans.into_iter().map(|span| span.text).collect::<String>(),
level,
));
}
crate::core::extensions::ExtensionUiRequest::Continuation { .. } => {
continuation = true;
}
_ => {}
}
}
let settled = notes
.iter()
.any(|(text, _)| text.contains("make the renderer faster"));
if settled || std::time::Instant::now() >= deadline {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert!(!continuation, "作业完成不该往主会话注入消息");
let ours: Vec<_> = notes
.iter()
.filter(|(text, _)| text.contains("make the renderer faster"))
.collect();
assert_eq!(ours.len(), 1, "{notes:?}");
assert!(ours[0].0.contains("completed"), "{:?}", ours[0]);
assert_eq!(ours[0].1, crate::core::extensions::UiNotifyLevel::Success);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
async fn await_autoresearch_terminal(
id: crate::extensions::subagent::autoresearch::JobId,
) -> crate::extensions::subagent::autoresearch::JobSnapshot {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
loop {
let snapshot = crate::extensions::subagent::autoresearch::snapshot(id)
.expect("job should be registered");
if snapshot.status.is_finished() {
return snapshot;
}
assert!(
std::time::Instant::now() < deadline,
"autoresearch job {id} did not settle in time: {snapshot:?}"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
#[test]
fn workflow_schema_agent_gets_a_structured_object_end_to_end() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let script = "export const meta = { name: 'x', description: 'd' }\n\
const r = await agent('x', { schema: { type: 'object', required: \
['answer'], properties: { answer: { type: 'string' } } } }); return r";
let f1 = tool_call_frame(
"StructuredOutput",
&json_string_escape(&json!({ "answer": "42" }).to_string()),
"call_2",
);
let f2 = text_frame("child final", "cmpl_c");
let (base, handle) = mock_server_multi(vec![f1, f2]).await;
let agent = test_agent(&base, &cwd);
let out = crate::core::extensions::Extension::execute_tool_async(
&crate::extensions::subagent::Subagent,
"SubagentWorkflow".to_string(),
json!({ "script": script }),
agent.exec_ctx(),
)
.await
.expect("workflow tool");
let run_id = out.details.as_ref().unwrap()["taskId"]
.as_str()
.expect("details.taskId")
.to_string();
let run = await_run_terminal(&run_id).await;
handle.await.unwrap();
assert_eq!(
run.value,
Some(json!({ "answer": "42" })),
"脚本拿到的应是 StructuredOutput 交回的**对象**"
);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn workflow_agent_resume_continues_the_same_child_session() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let f1 = text_frame("first answer", "cmpl_a");
let f2 = text_frame("second answer", "cmpl_b");
let (base, handle) = mock_server_multi(vec![f1, f2]).await;
let agent = test_agent(&base, &cwd);
let script = "export const meta = { name: 'res', description: 'd' }\n\
const a = await agent('one', { label: 'loop' })\n\
const b = await agent('more', { resume: 'loop' })\n\
return { a, b }";
let out = crate::core::extensions::Extension::execute_tool_async(
&crate::extensions::subagent::Subagent,
"SubagentWorkflow".to_string(),
json!({ "script": script }),
agent.exec_ctx(),
)
.await
.expect("workflow tool");
let run_id = out.details.as_ref().unwrap()["taskId"]
.as_str()
.expect("details.taskId")
.to_string();
let run = await_run_terminal(&run_id).await;
handle.await.unwrap();
let value = run.value.clone().unwrap();
assert_eq!(value["a"], "first answer", "{:?}", run.error);
assert_eq!(
value["b"], "second answer",
"resume 应续跑并拿到新答案: {value}"
);
let ids: Vec<String> = run
.log()
.iter()
.filter(|e| {
e["type"] == "workflow_agent"
&& e["state"] == "done"
&& e["recordId"].is_string()
})
.filter_map(|e| e["recordId"].as_str().map(str::to_string))
.collect();
assert_eq!(ids.len(), 2, "{ids:?}");
assert_eq!(ids[0], ids[1], "resume 复用被续跑 child 的记录: {ids:?}");
let log = run.log();
let mid_run: Vec<&serde_json::Value> = log
.iter()
.filter(|e| {
e["type"] == "workflow_agent"
&& e["recordId"].is_string()
&& e["state"] != "done"
})
.collect();
assert!(
!mid_run.is_empty()
&& mid_run
.iter()
.all(|e| e["recordId"].as_str() == Some(ids[0].as_str())),
"运行中条目应带同一个 recordId: {mid_run:?}"
);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn nested_subagents_can_delegate_one_more_level() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let main_call = tool_call_frame(
"Agent",
r#"{\"prompt\":\"child task\",\"description\":\"child\",\"subagent_type\":\"general-purpose\",\"run_in_background\":false}"#,
"call_child",
);
let child_call = tool_call_frame(
"Agent",
r#"{\"prompt\":\"grand task\",\"description\":\"grand\",\"subagent_type\":\"general-purpose\",\"run_in_background\":false}"#,
"call_grand",
);
let (base, handle) = mock_server_multi(vec![
main_call,
child_call,
text_frame("grand answer", "cmpl_g"),
text_frame("child saw grand answer", "cmpl_c"),
text_frame("done", "cmpl_m"),
])
.await;
let mut agent = test_agent(&base, &cwd);
agent.set_active_tools(vec!["read".to_string(), "Agent".to_string()]);
let answer = agent.prompt("go").await.expect("main turn");
handle.await.unwrap();
assert!(answer.contains("done"), "主的最终回答: {answer}");
let tree = crate::extensions::subagent::record_tree_for_test();
let (child_id, _, _) = tree
.iter()
.find(|(_, _, depth)| *depth == 1)
.expect("应有一个 depth 1 的子代理")
.clone();
let (grand_id, grand_parent, _) = tree
.iter()
.find(|(_, _, depth)| *depth == 2)
.expect("子代理应派出一个 depth 2 的孙代理")
.clone();
assert_eq!(
grand_parent.as_deref(),
Some(child_id.as_str()),
"孙的父是这个子代理: {tree:?}"
);
let (agent_type, result) =
crate::extensions::subagent::record_for_test(&grand_id).expect("孙记录");
assert_eq!(agent_type, "general-purpose");
assert!(
result.as_deref().unwrap_or_default().contains("grand answer"),
"孙的结果: {result:?}"
);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
async fn await_run_terminal(
run_id: &str,
) -> crate::extensions::subagent::workflow::task::WorkflowRun {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
loop {
let run = crate::extensions::subagent::workflow::task::find(run_id)
.expect("run should be registered");
if run.status.is_terminal() {
return run;
}
assert!(
std::time::Instant::now() < deadline,
"run {run_id} did not settle in time"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
#[test]
fn workflow_schema_agent_without_the_tool_call_yields_null() {
let _lock = subagent_test_lock();
let _agent_dir = crate::test_support::AgentDirGuard::temp();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::subagent::Subagent);
crate::core::extensions::set_extension_enabled("subagent", true);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let script = "export const meta = { name: 'x', description: 'd' }\n\
const r = await agent('x', { schema: { type: 'object', required: \
['answer'], properties: { answer: { type: 'string' } } } }); \
return { got: r }";
let f1 = text_frame(&json_string_escape("{\"answer\": \"42\"}"), "cmpl_c");
let (base, handle) = mock_server_multi(vec![f1]).await;
let agent = test_agent(&base, &cwd);
let out = crate::core::extensions::Extension::execute_tool_async(
&crate::extensions::subagent::Subagent,
"SubagentWorkflow".to_string(),
json!({ "script": script }),
agent.exec_ctx(),
)
.await
.expect("workflow tool");
let run_id = out.details.as_ref().unwrap()["taskId"]
.as_str()
.expect("details.taskId")
.to_string();
let run = await_run_terminal(&run_id).await;
handle.await.unwrap();
assert_eq!(
run.value,
Some(json!({ "got": null })),
"没从工具交回就是 null,正文里的 JSON 不算"
);
crate::core::extensions::set_extension_enabled("subagent", false);
_ = crate::core::extensions::unregister_extension("subagent");
});
}
#[test]
fn sequential_extension_batch_assigns_source_order_tool_indexes() {
struct SeqProbe;
impl crate::core::extensions::Extension for SeqProbe {
fn name(&self) -> &str {
"seq-probe-test"
}
fn tools(&self) -> Vec<crate::core::extensions::ExtensionTool> {
let mut probe = crate::core::extensions::ExtensionTool::simple(
"seq_probe",
"test tool that forces sequential batch execution",
json!({
"type": "object",
"properties": { "n": { "type": "integer" } },
"required": ["n"]
}),
"",
);
probe.execution_mode = Some(ToolExecutionMode::Sequential);
vec![probe]
}
fn execute_tool(
&self,
name: &str,
args: &Value,
) -> std::result::Result<tools::ToolResult, tools::ToolError> {
if name == "seq_probe" {
Ok(tools::ToolResult::text(format!(
"probe {}",
args.get("n").and_then(|v| v.as_i64()).unwrap_or(-1)
)))
} else {
Err(tools::ToolError("unknown".to_string()))
}
}
}
let _lock = subagent_test_lock();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
crate::core::extensions::register_extension(SeqProbe);
let tmp = std::env::temp_dir()
.join(format!("prux-v4-seq-probe-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
let cwd = tmp.to_string_lossy().to_string();
let sess_dir = tmp.join("sessions");
let session = crate::core::session_manager::Session::create(
&cwd,
Some(sess_dir.clone()),
true,
)
.unwrap();
let f1: Vec<&'static str> = vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
frame(
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"seq_probe","arguments":"{\"n\":1}"}}]}}]}"#
.to_string(),
),
frame(
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"call_2","type":"function","function":{"name":"seq_probe","arguments":"{\"n\":2}"}}]}}]}"#
.to_string(),
),
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let f2 = text_frame("parent final", "cmpl_p");
let (base, handle) = mock_server_multi(vec![f1, f2]).await;
let mut agent = test_agent(&base, &cwd);
agent.session = Some(session);
let r = agent.prompt("go").await;
handle.await.unwrap();
assert!(r.is_ok(), "prompt failed: {:?}", r.err());
drop(agent);
let path = crate::core::session_manager::list_sessions(&sess_dir)
.pop()
.expect("session file");
let opened =
crate::core::session_manager::Session::open(&path.to_string_lossy()).unwrap();
let entries = opened.get_entries();
let ts: Vec<&J> = entries
.iter()
.filter(|e| e.get("type").and_then(|v| v.as_str()) == Some("tool_started"))
.collect();
let calls: Vec<(String, i64)> = ts
.iter()
.map(|e| {
(
e["toolCallId"].as_str().unwrap().to_string(),
e["toolIndex"].as_i64().unwrap(),
)
})
.collect();
assert!(calls.contains(&("call_1".into(), 0)), "{calls:?}");
assert!(calls.contains(&("call_2".into(), 1)), "{calls:?}");
let red = opened.reduce_lane();
assert!(red.is_ok(), "reduce_lane 应通过: {:?}", red.err());
_ = crate::core::extensions::unregister_extension("seq-probe-test");
});
}
struct NestedProbe;
impl crate::core::extensions::Extension for NestedProbe {
fn name(&self) -> &str {
"nested-probe-test"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![
ExtensionTool::simple(
"nested_probe",
"compose other tools",
json!({ "type": "object", "properties": {} }),
"",
),
ExtensionTool::simple(
"hidden_probe",
"internal tool, not visible to the model",
json!({ "type": "object", "properties": {} }),
"",
)
.with_exposure(ToolExposure::Hidden),
ExtensionTool::simple(
"model_only_probe",
"visible to the model only",
json!({ "type": "object", "properties": {} }),
"",
)
.with_exposure(ToolExposure::ModelOnly),
]
}
fn execute_tool_async(
&self,
name: String,
_args: Value,
ctx: ToolExecCtx,
) -> BoxFuture<'static, std::result::Result<ToolResult, tools::ToolError>> {
Box::pin(async move {
match name.as_str() {
"nested_probe" => {
let bash = ctx
.execute_tool(
"bash".to_string(),
json!({ "command": "echo nested-out" }),
)
.await?;
let hidden = ctx
.execute_tool("hidden_probe".to_string(), json!({}))
.await?;
let denied = ctx
.execute_tool("model_only_probe".to_string(), json!({}))
.await;
let missing = ctx.execute_tool("nope_probe".to_string(), json!({})).await;
Ok(ToolResult {
text: format!(
"{}|{}|denied={}|missing={}",
bash.text.trim(),
hidden.text,
denied.is_err(),
missing.is_err()
),
structured_content: Some(json!({ "composed": 2 })),
..Default::default()
})
}
"hidden_probe" => Ok(ToolResult {
text: "hidden ran".to_string(),
usage: Some(Usage {
input: 10,
output: 5,
total_tokens: 15,
..Default::default()
}),
..Default::default()
}),
"model_only_probe" => Ok(ToolResult::text("should never run nested")),
other => Err(tools::ToolError(format!("unknown tool: {other}"))),
}
})
}
}
#[test]
fn nested_tool_calls_are_recorded_counted_and_evented() {
let _lock = subagent_test_lock();
let rt = test_rt();
rt.block_on(async {
crate::core::extensions::register_extension(NestedProbe);
let tmp = std::env::temp_dir()
.join(format!("prux-nested-probe-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
let cwd = tmp.to_string_lossy().to_string();
let f1: Vec<&'static str> = vec![
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
frame(
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_nested","type":"function","function":{"name":"nested_probe","arguments":"{}"}}]}}]}"#
.to_string(),
),
r#"{"id":"cmpl_t","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let (base, handle) = mock_server_multi(vec![f1, text_frame("done", "cmpl_p")]).await;
let mut agent = test_agent(&base, &cwd);
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "done");
_ = crate::core::extensions::unregister_extension("nested-probe-test");
let events = collected.lock().unwrap().clone();
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let tr = msgs
.iter()
.find(|m| m["role"] == "toolResult")
.expect("toolResult 消息");
let tr_text = tr["content"][0]["text"].as_str().unwrap();
assert!(tr_text.contains("nested-out"), "{tr_text}");
assert!(tr_text.contains("hidden ran"), "{tr_text}");
assert!(tr_text.contains("denied=true"), "{tr_text}");
assert!(tr_text.contains("missing=true"), "{tr_text}");
assert_eq!(tr["details"]["structuredContent"]["composed"], 2);
assert_eq!(tr["usage"]["totalTokens"], 15, "{tr:?}");
let nested = tr["details"]["nestedCalls"].as_array().unwrap();
let names: Vec<&str> = nested
.iter()
.map(|c| c["toolName"].as_str().unwrap())
.collect();
assert_eq!(
names,
vec!["bash", "hidden_probe", "model_only_probe", "nope_probe"],
"{nested:?}"
);
assert!(nested[0]["text"].as_str().unwrap().contains("nested-out"));
assert_eq!(nested[2]["isError"], true);
assert!(
nested[2]["text"]
.as_str()
.unwrap()
.contains("exposure: model-only"),
"{nested:?}"
);
assert!(
nested[3]["text"].as_str().unwrap().contains("not found"),
"{nested:?}"
);
assert_eq!(nested[1]["usage"]["totalTokens"], 15, "{nested:?}");
let nested_ends: Vec<&J> = events
.iter()
.filter(|e| {
e["type"] == "tool_execution_end"
&& e["parentToolCallId"].as_str() == Some("call_nested")
})
.collect();
assert_eq!(nested_ends.len(), 4, "每个嵌套调用一条结束事件");
let top_ends = events
.iter()
.filter(|e| {
e["type"] == "tool_execution_end"
&& !crate::core::extensions::is_nested_call_event(e)
})
.count();
assert_eq!(top_ends, 1, "只有顶层 nested_probe 一条不带 parent");
});
}
#[test]
fn nested_tool_calls_are_bounded_and_truncated() {
struct EchoProbe;
impl crate::core::extensions::Extension for EchoProbe {
fn name(&self) -> &str {
"nested-echo-test"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![ExtensionTool::simple(
"echo_probe",
"echo text back",
json!({ "type": "object", "properties": { "text": { "type": "string" } } }),
"",
)]
}
fn execute_tool(
&self,
name: &str,
args: &Value,
) -> std::result::Result<ToolResult, tools::ToolError> {
if name != "echo_probe" {
return Err(tools::ToolError("unknown".to_string()));
}
Ok(ToolResult::text(
args.get("text").and_then(|v| v.as_str()).unwrap_or(""),
))
}
}
let _lock = subagent_test_lock();
let rt = test_rt();
rt.block_on(async {
let _ad = crate::test_support::AgentDirGuard::temp();
crate::core::extensions::register_extension(EchoProbe);
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, _handle) = mock_server_multi(vec![]).await;
let agent = test_agent(&base, &cwd);
let tpl = Arc::new(agent.snapshot_template());
let abort = Arc::new(AtomicBool::new(false));
let log = Arc::new(Mutex::new(NestedCallLog::default()));
let big = "x".repeat(NESTED_CALL_TEXT_CAP + 100);
for i in 0..(MAX_NESTED_CALLS + 3) {
let out = run_nested_tool(
&cwd,
&tpl,
&abort,
&None,
&Some("call_parent".to_string()),
"echo_probe",
json!({ "text": big }),
0,
&log,
None,
)
.await;
assert!(
out.is_ok(),
"第 {i} 次嵌套调用: {:?}",
out.err().map(|e| e.0)
);
}
_ = crate::core::extensions::unregister_extension("nested-echo-test");
let taken = take_log(&log);
assert_eq!(taken.calls.len(), MAX_NESTED_CALLS);
assert_eq!(taken.dropped, 3);
let first = &taken.calls[0];
assert_eq!(first["toolName"], "echo_probe");
assert!(
first["text"].as_str().unwrap().chars().count() <= NESTED_CALL_TEXT_CAP + 1,
"文本应被截断"
);
assert!(first["args"]["_truncated"].is_string(), "{first:?}");
assert!(take_log(&log).calls.is_empty());
});
}
#[test]
fn nested_tool_calls_respect_depth_limit_and_abort() {
let _lock = subagent_test_lock();
let rt = test_rt();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, _handle) = mock_server_multi(vec![]).await;
let agent = test_agent(&base, &cwd);
let tpl = Arc::new(agent.snapshot_template());
let log = Arc::new(Mutex::new(NestedCallLog::default()));
let abort = Arc::new(AtomicBool::new(false));
let too_deep = run_nested_tool(
&cwd,
&tpl,
&abort,
&None,
&None,
"bash",
json!({ "command": "echo x" }),
MAX_NESTED_TOOL_DEPTH,
&log,
None,
)
.await
.unwrap_err();
assert!(too_deep.0.contains("depth limit"), "{}", too_deep.0);
let abort = Arc::new(AtomicBool::new(true));
let aborted = run_nested_tool(
&cwd,
&tpl,
&abort,
&None,
&None,
"bash",
json!({ "command": "echo x" }),
0,
&log,
None,
)
.await
.unwrap_err();
assert!(aborted.0.contains("aborted"), "{}", aborted.0);
let taken = take_log(&log);
assert_eq!(taken.calls.len(), 2, "{:?}", taken.calls);
assert_eq!(taken.calls[0]["isError"], true);
assert!(
taken.calls[0]["text"]
.as_str()
.unwrap()
.contains("depth limit")
);
assert!(taken.calls[1]["text"].as_str().unwrap().contains("aborted"));
assert!(take_log(&log).calls.is_empty(), "账本取走后清空");
});
}
#[test]
fn exposure_controls_model_visible_tool_surface() {
struct ExposureProbe;
impl crate::core::extensions::Extension for ExposureProbe {
fn name(&self) -> &str {
"exposure-probe-test"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![
ExtensionTool::simple("expo_direct", "visible direct tool", json!({}), ""),
ExtensionTool::simple(
"expo_model_only",
"visible model-only tool",
json!({}),
"",
)
.with_exposure(ToolExposure::ModelOnly),
ExtensionTool::simple("expo_hidden", "hidden internal tool", json!({}), "")
.with_exposure(ToolExposure::Hidden),
ExtensionTool::simple("expo_codemode", "codemode-only tool", json!({}), "")
.with_exposure(ToolExposure::Codemode),
ExtensionTool::simple("expo_deferred", "deferred tool", json!({}), "")
.with_exposure(ToolExposure::Deferred),
]
}
}
let _lock = subagent_test_lock();
let _ad = crate::test_support::AgentDirGuard::temp();
crate::core::extensions::register_extension(ExposureProbe);
crate::core::extensions::clear_deferred_activations();
let ctx = ToolRebuildCtx::default();
let names = |tools: &[(String, String, Value)]| -> Vec<String> {
tools.iter().map(|(n, _, _)| n.clone()).collect()
};
let (prompt, tools, _hidden) = compose_tools(".", &[], &ctx, &[]);
let visible = names(&tools);
assert!(visible.contains(&"expo_direct".to_string()), "{visible:?}");
assert!(
visible.contains(&"expo_model_only".to_string()),
"{visible:?}"
);
assert!(!visible.contains(&"expo_hidden".to_string()), "{visible:?}");
assert!(
!visible.contains(&"expo_codemode".to_string()),
"{visible:?}"
);
assert!(
!visible.contains(&"expo_deferred".to_string()),
"{visible:?}"
);
assert!(
!prompt.contains("hidden internal tool"),
"提示里不应出现 hidden 工具"
);
let deferred: Vec<String> = crate::core::extensions::unactivated_deferred_tools()
.into_iter()
.map(|t| t.name)
.collect();
assert!(
deferred.contains(&"expo_deferred".to_string()),
"{deferred:?}"
);
crate::core::extensions::activate_deferred_tools(&["expo_deferred".to_string()]);
let (_, tools, _) = compose_tools(".", &[], &ctx, &[]);
assert!(names(&tools).contains(&"expo_deferred".to_string()));
crate::core::extensions::clear_deferred_activations();
_ = crate::core::extensions::unregister_extension("exposure-probe-test");
}
#[test]
fn codemode_prepare_loadout_inlines_callable_tools() {
struct CodemodeProbe;
impl crate::core::extensions::Extension for CodemodeProbe {
fn name(&self) -> &str {
"codemode-probe-test"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![
ExtensionTool::simple(
"cm_probe_tool",
"probe description",
json!({"type": "object", "properties": {"v": {"type": "string"}}}),
"",
),
ExtensionTool::simple(
"cm_internal_tool",
"codemode-only probe",
json!({"type": "object", "properties": {}}),
"",
)
.with_exposure(ToolExposure::Codemode),
]
}
}
let _lock = subagent_test_lock();
let _ad = crate::test_support::AgentDirGuard::temp();
crate::core::extensions::register_extension(CodemodeProbe);
crate::core::extensions::register_extension(crate::extensions::codemode::Codemode);
crate::core::extensions::set_extension_enabled("codemode", true);
let ctx = ToolRebuildCtx::default();
let (_, tools, hidden) = compose_tools(".", &[], &ctx, &[]);
assert!(hidden.is_empty(), "{hidden:?}");
let codemode = tools
.iter()
.find(|(n, _, _)| n == crate::extensions::codemode::TOOL_NAME)
.expect("codemode 应在工具表里");
assert!(
codemode.1.contains("### `cm_internal_tool`"),
"codemode 描述应内联非 direct 的可调工具:{}",
codemode.1
);
assert!(
!codemode.1.contains("### `cm_probe_tool`"),
"direct 工具不进 codemode 描述:{}",
codemode.1
);
assert!(tools.iter().all(|(n, _, _)| n != "cm_internal_tool"));
let probe = tools
.iter()
.find(|(n, _, _)| n == "cm_probe_tool")
.expect("探针工具应可见");
assert!(
probe.1.contains("codemode tool declaration:"),
"direct 工具的描述应附上脚本调用样例:{}",
probe.1
);
crate::core::extensions::set_extension_enabled("codemode", false);
let (_, tools, _) = compose_tools(".", &[], &ctx, &[]);
assert!(tools.iter().all(|(n, _, _)| n != "codemode"));
_ = crate::core::extensions::unregister_extension("codemode");
_ = crate::core::extensions::unregister_extension("codemode-probe-test");
}
#[test]
fn codemode_only_mode_hides_direct_declarations() {
struct OnlyModeProbe;
impl crate::core::extensions::Extension for OnlyModeProbe {
fn name(&self) -> &str {
"codemode-only-probe-test"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![ExtensionTool::simple(
"cm_only_tool",
"only-mode probe",
json!({"type": "object", "properties": {}}),
"",
)]
}
}
let _lock = subagent_test_lock();
let _ad = crate::test_support::AgentDirGuard::temp();
let cfg_dir = crate::core::settings_manager::agent_dir().join("extensions");
std::fs::create_dir_all(&cfg_dir).unwrap();
std::fs::write(
cfg_dir.join("codemode.json"),
r#"{"mode": "only", "inlineBudget": 100000}"#,
)
.unwrap();
crate::core::extensions::register_extension(OnlyModeProbe);
crate::core::extensions::register_extension(crate::extensions::codemode::Codemode);
crate::core::extensions::set_extension_enabled("codemode", true);
let ctx = ToolRebuildCtx {
selection: ToolSelection::tools(vec!["read".into(), "bash".into()]),
..ToolRebuildCtx::default()
};
let (prompt, tools, hidden) = compose_tools(".", &[], &ctx, &[]);
assert!(
hidden.contains(&"cm_only_tool".to_string()),
"only 模式应隐藏 direct 工具的声明:{hidden:?}"
);
assert!(
hidden.contains(&"read".to_string()) && hidden.contains(&"bash".to_string()),
"内置工具的声明同样被摘掉:{hidden:?}"
);
assert!(
!prompt.contains("- read:") && !prompt.contains("- bash:"),
"only 模式下内置工具不应出现在系统提示里:{prompt}"
);
let codemode = tools
.iter()
.find(|(n, _, _)| n == crate::extensions::codemode::TOOL_NAME)
.expect("codemode 应在工具表里");
assert!(
codemode.1.contains("### `cm_only_tool`"),
"only 模式下可调工具应全部列进 codemode 描述:{}",
codemode.1
);
_ = crate::core::extensions::unregister_extension("codemode");
_ = crate::core::extensions::unregister_extension("codemode-only-probe-test");
}
#[test]
fn codemode_script_orchestrates_nested_tools() {
let _lock = subagent_test_lock();
let rt = test_rt();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::codemode::Codemode);
crate::core::extensions::set_extension_enabled("codemode", true);
let tmp = std::env::temp_dir().join(format!(
"prux-codemode-e2e-{}",
crate::utils::time::now_ms()
));
std::fs::create_dir_all(&tmp).unwrap();
let file = tmp.join("note.txt");
std::fs::write(&file, "codemode-e2e-content").unwrap();
let cwd = tmp.to_string_lossy().to_string();
let code = format!(
"const out = await tools.read({{path: {}}});\nreturn out.trim();",
serde_json::to_string(&file.to_string_lossy().to_string()).unwrap()
);
let delta = serde_json::json!({
"id": "cmpl_c",
"choices": [{"index": 0, "delta": {"tool_calls": [{
"index": 0,
"id": "call_codemode",
"type": "function",
"function": {
"name": crate::extensions::codemode::TOOL_NAME,
"arguments": serde_json::json!({"code": code}).to_string(),
}
}]}}]
})
.to_string();
let f1: Vec<&'static str> = vec![
r#"{"id":"cmpl_c","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
frame(delta),
r#"{"id":"cmpl_c","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let (base, handle) = mock_server_multi(vec![f1, text_frame("done", "cmpl_p")]).await;
let mut agent = test_agent(&base, &cwd);
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "done");
_ = crate::core::extensions::unregister_extension("codemode");
_ = std::fs::remove_dir_all(&tmp);
let events = collected.lock().unwrap().clone();
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let tr = msgs
.iter()
.find(|m| m["role"] == "toolResult")
.expect("toolResult 消息");
let tr_text = tr["content"][0]["text"].as_str().unwrap();
assert!(tr_text.contains("Script completed"), "{tr_text}");
assert!(tr_text.contains("codemode-e2e-content"), "{tr_text}");
let nested = tr["details"]["nestedCalls"].as_array().unwrap();
assert_eq!(nested.len(), 1, "{nested:?}");
assert_eq!(nested[0]["toolName"], "read");
assert_eq!(nested[0]["isError"], false);
assert!(
nested[0]["text"].as_str().unwrap().contains("codemode-e2e-content"),
"{nested:?}"
);
assert_eq!(
msgs.iter()
.filter(|m| m["role"] == "toolResult")
.count(),
1,
"嵌套结果不单独成为会话消息"
);
});
}
struct StructuredErrorProbe;
impl crate::core::extensions::Extension for StructuredErrorProbe {
fn name(&self) -> &str {
"structured-error-probe-test"
}
fn tools(&self) -> Vec<ExtensionTool> {
vec![
ExtensionTool::simple(
"err_with_schema",
"fails but carries structured output",
json!({ "type": "object", "properties": {} }),
"",
)
.with_exposure(ToolExposure::Codemode)
.with_output_schema(
json!({ "type": "object", "properties": { "ok": { "type": "boolean" } } }),
),
ExtensionTool::simple(
"err_plain",
"fails without structured output",
json!({ "type": "object", "properties": {} }),
"",
)
.with_exposure(ToolExposure::Codemode),
]
}
fn execute_tool_async(
&self,
name: String,
_args: Value,
_ctx: ToolExecCtx,
) -> BoxFuture<'static, std::result::Result<ToolResult, tools::ToolError>> {
Box::pin(async move {
match name.as_str() {
"err_with_schema" => Ok(ToolResult::error("probe failed")
.with_structured_content(json!({ "ok": false }))),
"err_plain" => Ok(ToolResult::error("plain failure")),
other => Err(tools::ToolError(format!("unknown tool: {other}"))),
}
})
}
}
#[test]
fn codemode_script_reads_structured_error_result() {
let _lock = subagent_test_lock();
let rt = test_rt();
rt.block_on(async {
crate::core::extensions::register_extension(StructuredErrorProbe);
crate::core::extensions::register_extension(crate::extensions::codemode::Codemode);
crate::core::extensions::set_extension_enabled("codemode", true);
let tmp = std::env::temp_dir().join(format!(
"prux-codemode-error-{}",
crate::utils::time::now_ms()
));
std::fs::create_dir_all(&tmp).unwrap();
let cwd = tmp.to_string_lossy().to_string();
let code = "let a;\ntry { a = await tools.err_with_schema({}); } catch (e) { a = 'threw:' + e.message; }\nlet b;\ntry { b = await tools.err_plain({}); } catch (e) { b = 'threw:' + e.message; }\nreturn JSON.stringify({a, b});";
let delta = serde_json::json!({
"id": "cmpl_e",
"choices": [{"index": 0, "delta": {"tool_calls": [{
"index": 0,
"id": "call_codemode_err",
"type": "function",
"function": {
"name": crate::extensions::codemode::TOOL_NAME,
"arguments": serde_json::json!({"code": code}).to_string(),
}
}]}}]
})
.to_string();
let f1: Vec<&'static str> = vec![
r#"{"id":"cmpl_e","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
frame(delta),
r#"{"id":"cmpl_e","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let (base, handle) = mock_server_multi(vec![f1, text_frame("done", "cmpl_p")]).await;
let mut agent = test_agent(&base, &cwd);
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "done");
_ = crate::core::extensions::unregister_extension("structured-error-probe-test");
_ = crate::core::extensions::unregister_extension("codemode");
_ = std::fs::remove_dir_all(&tmp);
let events = collected.lock().unwrap().clone();
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let tr = msgs
.iter()
.find(|m| m["role"] == "toolResult")
.expect("toolResult 消息");
let tr_text = tr["content"][0]["text"].as_str().unwrap();
assert!(tr_text.contains("\"ok\":false"), "{tr_text}");
assert!(tr_text.contains("threw:plain failure"), "{tr_text}");
let nested = tr["details"]["nestedCalls"].as_array().unwrap();
let with_schema = nested
.iter()
.find(|c| c["toolName"] == "err_with_schema")
.expect("err_with_schema 记录");
assert_eq!(with_schema["isError"], true);
assert_eq!(with_schema["structuredContent"]["ok"], false);
let plain = nested
.iter()
.find(|c| c["toolName"] == "err_plain")
.expect("err_plain 记录");
assert_eq!(plain["isError"], true);
assert!(plain.get("structuredContent").is_none(), "{plain:?}");
});
}
#[test]
fn codemode_script_receives_full_bash_output() {
let _lock = subagent_test_lock();
let rt = test_rt();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::codemode::Codemode);
crate::core::extensions::set_extension_enabled("codemode", true);
let tmp = std::env::temp_dir().join(format!(
"prux-codemode-bash-{}",
crate::utils::time::now_ms()
));
std::fs::create_dir_all(&tmp).unwrap();
let cwd = tmp.to_string_lossy().to_string();
let code = "const r = await tools.bash({command: \"yes x | head -n 700000\"});\nreturn JSON.stringify({truncated: r.truncated, marker: r.output.includes('bytes omitted'), code: r.exit_code});";
let delta = serde_json::json!({
"id": "cmpl_b",
"choices": [{"index": 0, "delta": {"tool_calls": [{
"index": 0,
"id": "call_codemode_bash",
"type": "function",
"function": {
"name": crate::extensions::codemode::TOOL_NAME,
"arguments": serde_json::json!({"code": code}).to_string(),
}
}]}}]
})
.to_string();
let f1: Vec<&'static str> = vec![
r#"{"id":"cmpl_b","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
frame(delta),
r#"{"id":"cmpl_b","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let (base, handle) = mock_server_multi(vec![f1, text_frame("done", "cmpl_p")]).await;
let mut agent = test_agent(&base, &cwd);
agent.set_active_tools(vec!["read".into(), "write".into(), "bash".into()]);
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
assert_eq!(text, "done");
_ = crate::core::extensions::unregister_extension("codemode");
_ = std::fs::remove_dir_all(&tmp);
let events = collected.lock().unwrap().clone();
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let tr = msgs
.iter()
.find(|m| m["role"] == "toolResult")
.expect("toolResult 消息");
let tr_text = tr["content"][0]["text"].as_str().unwrap();
assert!(tr_text.contains("\"truncated\":true"), "{tr_text}");
assert!(tr_text.contains("\"marker\":true"), "{tr_text}");
assert!(tr_text.contains("\"code\":0"), "{tr_text}");
});
}
#[test]
fn codemode_script_runs_classifier_and_counts_usage() {
let _lock = subagent_test_lock();
crate::test_support::pin_test_agent_dir();
let models_path = crate::core::settings_manager::agent_dir().join("models.json");
let rt = test_rt();
rt.block_on(async {
crate::core::extensions::register_extension(crate::extensions::codemode::Codemode);
crate::core::extensions::set_extension_enabled("codemode", true);
let (classifier_base, classifier_handle) = mock_json_once(
json!({
"answers": { "sentiment": { "type": "noul", "noul": 0.87 } },
"usage": { "input_tokens": 120, "output_tokens": 8 }
})
.to_string(),
)
.await;
std::fs::write(
&models_path,
json!({
"providers": {
"openrouter": { "baseUrl": classifier_base, "apiKey": "sk-test", "models": [] }
}
})
.to_string(),
)
.unwrap();
let tmp = std::env::temp_dir().join(format!(
"prux-codemode-model-{}",
crate::utils::time::now_ms()
));
std::fs::create_dir_all(&tmp).unwrap();
let cwd = tmp.to_string_lossy().to_string();
let code = "const m = await models.getModelOfType(\"classifier\", \"openrouter\", \"~typesafe/jev-latest\");\nconst r = await models.classify(m, {state: {text: \"great\"}, questions: {sentiment: {type: \"bool\", instructions: \"Is it positive?\", criteria: {true: \"positive\", false: \"negative\"}}}});\nreturn JSON.stringify({p: r.answers.sentiment.probability, stop: r.stopReason, selected: m && m.id});";
let delta = serde_json::json!({
"id": "cmpl_m",
"choices": [{"index": 0, "delta": {"tool_calls": [{
"index": 0,
"id": "call_codemode_model",
"type": "function",
"function": {
"name": crate::extensions::codemode::TOOL_NAME,
"arguments": serde_json::json!({"code": code}).to_string(),
}
}]}}]
})
.to_string();
let f1: Vec<&'static str> = vec![
r#"{"id":"cmpl_m","choices":[{"index":0,"delta":{"role":"assistant","content":null}}]}"#,
frame(delta),
r#"{"id":"cmpl_m","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}"#,
r#"[DONE]"#,
];
let (base, handle) = mock_server_multi(vec![f1, text_frame("done", "cmpl_p")]).await;
let mut agent = test_agent(&base, &cwd);
let collected: Arc<Mutex<Vec<J>>> = Arc::new(Mutex::new(Vec::new()));
let sink = collected.clone();
agent.attach_json_sink(Box::new(move |e: J| sink.lock().unwrap().push(e)));
let text = agent.prompt("go").await.unwrap();
handle.await.unwrap();
let _ = tokio::time::timeout(std::time::Duration::from_secs(5), classifier_handle).await;
assert_eq!(text, "done");
_ = crate::core::extensions::unregister_extension("codemode");
_ = std::fs::remove_dir_all(&tmp);
let _ = std::fs::remove_file(&models_path);
let events = collected.lock().unwrap().clone();
let agent_end = events_of(&events, "agent_end")[0];
let msgs = agent_end["messages"].as_array().unwrap();
let tr = msgs
.iter()
.find(|m| m["role"] == "toolResult")
.expect("toolResult 消息");
let tr_text = tr["content"][0]["text"].as_str().unwrap();
assert!(tr_text.contains("\"p\":0.87"), "{tr_text}");
assert!(tr_text.contains("\"stop\":\"stop\""), "{tr_text}");
assert!(
tr_text.contains("\"selected\":\"~typesafe/jev-latest\""),
"{tr_text}"
);
assert_eq!(tr["usage"]["totalTokens"], 128, "{tr:?}");
let calls = tr["details"]["modelCalls"].as_array().unwrap();
assert_eq!(calls.len(), 1, "{calls:?}");
assert_eq!(calls[0]["name"], "models.classify");
assert_eq!(calls[0]["status"], "ok");
assert_eq!(calls[0]["args"], "openrouter/~typesafe/jev-latest");
assert!(calls[0]["durationMs"].is_number());
});
}
#[test]
fn structured_content_lands_in_details() {
let call = ContentBlock::ToolCall {
id: "c1".to_string(),
name: "mcp".to_string(),
arguments: json!({}),
thought_signature: None,
namespace: None,
};
let fin = FinalizedToolCall::from_tool_result(
&call,
"mcp",
&json!({}),
ToolResult {
text: "ok".to_string(),
details: Some(json!({ "action": "call" })),
structured_content: Some(json!({ "rows": [1, 2] })),
..Default::default()
},
);
let details = fin.details.as_ref().unwrap();
assert_eq!(details["action"], "call");
assert_eq!(details["structuredContent"]["rows"][1], 2);
assert_eq!(
fin.to_result_json()["details"]["structuredContent"]["rows"][0],
1
);
let fin = FinalizedToolCall::from_tool_result(
&call,
"x",
&json!({}),
ToolResult {
text: "ok".to_string(),
structured_content: Some(json!({ "a": 1 })),
..Default::default()
},
);
assert_eq!(fin.details.as_ref().unwrap()["structuredContent"]["a"], 1);
}
#[test]
fn is_nested_call_event_requires_parent_id() {
assert!(!crate::core::extensions::is_nested_call_event(
&json!({ "type": "tool_execution_end", "toolName": "bash" })
));
assert!(!crate::core::extensions::is_nested_call_event(
&json!({ "type": "tool_execution_end", "parentToolCallId": J::Null })
));
assert!(crate::core::extensions::is_nested_call_event(
&json!({ "type": "tool_execution_end", "parentToolCallId": "call_1" })
));
}
#[test]
fn tool_management_api_get_and_set_active() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let cwd = std::env::temp_dir().to_string_lossy().to_string();
let (base, handle) = mock_server_multi(vec![text_frame("ok", "c1")]).await;
let mut agent = test_agent(&base, &cwd);
let active = agent.get_active_tools();
assert!(active.contains(&"read".to_string()));
assert!(active.contains(&"write".to_string()));
let all = agent.get_all_tools();
assert!(all.contains(&"read".to_string()));
assert!(all.contains(&"bash".to_string()));
agent.set_active_tools(vec!["read".to_string()]);
let active2 = agent.get_active_tools();
assert!(active2.contains(&"read".to_string()));
assert!(
!active2.contains(&"write".to_string()),
"write 应从活动列表移除, got: {active2:?}"
);
handle.await.unwrap();
});
}
#[test]
fn run_loop_writes_v4_durable_records() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let tmp =
std::env::temp_dir().join(format!("prux-v4-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
std::fs::write(tmp.join("a.rs"), "fn main() {}").unwrap();
let cwd = tmp.to_string_lossy().to_string();
let sess_dir = tmp.join("sessions");
let session =
crate::core::session_manager::Session::create(&cwd, Some(sess_dir.clone()), true)
.unwrap();
let frames1 = tool_call_frame("read", r#"{\"path\":\"a.rs\"}"#, "call_1");
let frames2 = text_frame("Read done", "cmpl_2");
let (base, handle) = mock_server_multi(vec![frames1, frames2]).await;
let entry = model_resolver::ModelEntry {
model_type: Default::default(),
id: "test-model".into(),
name: "Test".into(),
api: "openai-completions".into(),
base_url: base.clone(),
provider: "test".into(),
input_limits: crate::utils::image::InputLimits::default(),
allowed_fallback_models: Vec::new(),
sampling_params: None,
sampling_params_by_thinking_level: None,
provider_routing: None,
supports_usage_in_streaming: true,
supports_store: true,
supports_finish_reason: true,
requires_assistant_after_tool_result: false,
supports_strict_mode: false,
send_session_affinity_headers: None,
session_affinity_format: None,
output: Vec::new(),
input: vec!["text".into()],
reasoning: false,
context_window: 200_000,
max_tokens: 4096,
max_tokens_field: "max_tokens".into(),
cost: None,
supports_developer_role: false,
requires_reasoning_content_on_assistant_messages: false,
thinking_format: String::new(),
supports_reasoning_effort: false,
thinking_level_map: None,
auth_header: false,
supports_openai_grammar_tools: false,
};
let mut agent = Agent::new(
entry,
cwd.clone(),
ToolSelection::tools(vec!["read".into(), "write".into()]),
Vec::new(),
None,
None,
None,
Vec::new(),
Some(session),
Some("sk-test".to_string()),
true,
Vec::new(),
)
.unwrap();
let _ = agent.prompt("read it").await.unwrap();
handle.await.unwrap();
drop(agent);
let path = crate::core::session_manager::list_sessions(&sess_dir)
.pop()
.expect("session file");
let opened =
crate::core::session_manager::Session::open(&path.to_string_lossy()).unwrap();
let entries = opened.get_entries();
let types: Vec<&str> = entries
.iter()
.filter_map(|e| e.get("type").and_then(|v| v.as_str()))
.collect();
assert!(types.contains(&"operation_started"), "{types:?}");
assert!(types.contains(&"step_attempt"), "{types:?}");
assert!(types.contains(&"tool_started"), "{types:?}");
assert!(types.contains(&"operation_finished"), "{types:?}");
let ts = entries
.iter()
.find(|e| e.get("type").and_then(|v| v.as_str()) == Some("tool_started"))
.unwrap();
assert!(ts.get("assistantEntryId").is_some());
assert!(ts.get("toolIndex").is_some());
assert_eq!(ts["toolCallId"], "call_1");
assert_eq!(ts["toolName"], "read");
assert!(ts.get("resultEntryId").is_some());
assert_eq!(ts["replay"], "never");
let os = entries
.iter()
.find(|e| e.get("type").and_then(|v| v.as_str()) == Some("operation_started"))
.unwrap();
let of = entries
.iter()
.find(|e| e.get("type").and_then(|v| v.as_str()) == Some("operation_finished"))
.unwrap();
assert_eq!(os["id"], of["runId"]);
assert_eq!(of["outcome"], "completed");
let red = opened.reduce_lane();
assert!(red.is_ok(), "reduce_lane 应通过: {:?}", red.err());
});
}
#[test]
fn run_loop_writes_failed_finish_on_provider_error() {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let tmp =
std::env::temp_dir().join(format!("prux-v4-err-{}", crate::utils::time::now_ms()));
std::fs::create_dir_all(&tmp).unwrap();
std::fs::write(tmp.join("a.rs"), "fn main() {}").unwrap();
let cwd = tmp.to_string_lossy().to_string();
let sess_dir = tmp.join("sessions");
let session =
crate::core::session_manager::Session::create(&cwd, Some(sess_dir.clone()), true)
.unwrap();
let (base, handle) = mock_server_sequence(vec![MockResp::Err(
"HTTP/1.1 500 Internal Server Error",
r#"{"error":{"message":"boom"}}"#,
)])
.await;
let entry = model_resolver::ModelEntry {
model_type: Default::default(),
id: "test-model".into(),
name: "Test".into(),
api: "openai-completions".into(),
base_url: base.clone(),
provider: "test".into(),
input_limits: crate::utils::image::InputLimits::default(),
allowed_fallback_models: Vec::new(),
sampling_params: None,
sampling_params_by_thinking_level: None,
provider_routing: None,
supports_usage_in_streaming: true,
supports_store: true,
supports_finish_reason: true,
requires_assistant_after_tool_result: false,
supports_strict_mode: false,
send_session_affinity_headers: None,
session_affinity_format: None,
output: Vec::new(),
input: vec!["text".into()],
reasoning: false,
context_window: 200_000,
max_tokens: 4096,
max_tokens_field: "max_tokens".into(),
cost: None,
supports_developer_role: false,
requires_reasoning_content_on_assistant_messages: false,
thinking_format: String::new(),
supports_reasoning_effort: false,
thinking_level_map: None,
auth_header: false,
supports_openai_grammar_tools: false,
};
let mut agent = Agent::new(
entry,
cwd.clone(),
ToolSelection::tools(vec!["read".into(), "write".into()]),
Vec::new(),
None,
None,
None,
Vec::new(),
Some(session),
Some("sk-test".to_string()),
true,
Vec::new(),
)
.unwrap();
let _ = agent.prompt("read it").await;
handle.await.unwrap();
drop(agent);
let path = crate::core::session_manager::list_sessions(&sess_dir)
.pop()
.expect("session file");
let opened =
crate::core::session_manager::Session::open(&path.to_string_lossy()).unwrap();
let entries = opened.get_entries();
let of = entries
.iter()
.find(|e| e.get("type").and_then(|v| v.as_str()) == Some("operation_finished"))
.unwrap_or_else(|| panic!("必须落 operation_finished: {entries:#?}"));
assert_eq!(of["outcome"], "failed", "error 回合 outcome=failed");
assert_eq!(opened.open_operations_count(), 0);
assert!(opened.reduce_lane().is_ok());
});
}
}