| @@ -10,7 +10,7 @@ use agent_types::compression::CompressedView; |
| use agent_types::context::prompt::result::PromptBuildResult; | use agent_types::context::prompt::result::PromptBuildResult; |
| use agent_types::events::ToolResultEvent; | use agent_types::events::ToolResultEvent; |
| use agent_types::outcome::{AgentError, AgentOutcome}; | use agent_types::outcome::{AgentError, AgentOutcome}; |
| -use agent_types::tool::{RawToolCall, RawToolOutcome, ToolExecutionResult}; | +use agent_types::tool::{EffectProfile, RawToolCall, RawToolOutcome, ToolExecutionResult}; |
| use agent_types::{ | use agent_types::{ |
| AssistantMessage, ChatMessage, ContentBlock, LlmError, MessageRole, StreamChunk, ToolUseBlock, | AssistantMessage, ChatMessage, ContentBlock, LlmError, MessageRole, StreamChunk, ToolUseBlock, |
| }; | }; |
| @@ -177,8 +177,8 @@ pub async fn run_agent_loop( |
| return Err(error); | return Err(error); |
| } | } |
| update_turn_span_after_llm(&mut ctx).await; | update_turn_span_after_llm(&mut ctx).await; |
| - let suspended_call = match tool_exec(&mut ctx).await { | + let suspended_calls = match tool_exec(&mut ctx).await { |
| - Ok(suspended_call) => suspended_call, | + Ok(suspended_calls) => suspended_calls, |
| Err(error) => { | Err(error) => { |
| end_turn_span( | end_turn_span( |
| &mut ctx, | &mut ctx, |
| @@ -196,7 +196,7 @@ pub async fn run_agent_loop( |
| return Err(error); | return Err(error); |
| } | } |
| }; | }; |
| - if let Some(suspended_call) = suspended_call { | + if !suspended_calls.is_empty() { |
| end_turn_span( | end_turn_span( |
| &mut ctx, | &mut ctx, |
| TraceOutcome::Ok, | TraceOutcome::Ok, |
| @@ -211,7 +211,7 @@ pub async fn run_agent_loop( |
| "suspended", | "suspended", |
| ) | ) |
| .await; | .await; |
| - return Ok(LoopRunResult::Suspended(suspended_call)); | + return Ok(LoopRunResult::Suspended(suspended_calls)); |
| } | } |
| decide(&mut ctx); | decide(&mut ctx); |
| | |
| @@ -705,6 +705,92 @@ fn microcompact(ctx: &mut LoopContext<'_>) { |
| } | } |
| } | } |
| | |
| +fn prune_stale_tool_output(messages: &mut [ChatMessage]) { |
| + const KEEP_RECENT_TOOL_BYTES: usize = 40_000; |
| + const MIN_PRUNABLE_BYTES: usize = 1_000; |
| + const PRUNED_MARKER: &str = |
| + "[older tool output pruned to save context — re-run the tool or read the file if you still need it]"; |
| + let mut kept = 0usize; |
| + let mut pruned = 0usize; |
| + for message in messages.iter_mut().rev() { |
| + for block in message.blocks.iter_mut() { |
| + if let ContentBlock::ToolResult { output, .. } = block { |
| + if output.as_str() == PRUNED_MARKER { |
| + continue; |
| + } |
| + if kept < KEEP_RECENT_TOOL_BYTES { |
| + kept += output.len(); |
| + } else if output.len() > MIN_PRUNABLE_BYTES { |
| + *output = PRUNED_MARKER.to_string(); |
| + pruned += 1; |
| + } |
| + } |
| + } |
| + } |
| + if pruned > 0 { |
| + tracing::debug!(pruned, "pruned stale tool output beyond recent window"); |
| + } |
| +} |
| + |
| + |
| + |
| + |
| +fn live_context_snippets( |
| + ctx: &LoopContext<'_>, |
| +) -> Vec<agent_types::context::prompt::MemorySnippet> { |
| + use agent_types::context::prompt::MemorySnippet; |
| + let mut snippets = Vec::new(); |
| + |
| + let turn = ctx.turn.turn_number; |
| + let max_turns = ctx.snapshot.max_turns; |
| + let tokens_used = ctx.state.token_usage.total_tokens; |
| + let remaining = max_turns.saturating_sub(turn); |
| + if max_turns > 0 && remaining <= 5 { |
| + let horizon = format!( |
| + "- turn: {turn}/{max_turns} ({remaining} remaining)\n- tokens used so far: ~{tokens_used}\n- NEARING THE TURN LIMIT — stop investigating and converge now: apply your best fix, save the files, and finish this turn. A committed partial fix beats an unfinished exploration that gets cut off." |
| + ); |
| + snippets.push(MemorySnippet { |
| + source: "horizon".to_string(), |
| + content: horizon, |
| + relevance_score: 1.0, |
| + }); |
| + } |
| + |
| + let window = ctx.snapshot.token_budget_config.total_budget; |
| + let context_input = ctx.state.token_usage.prompt_tokens; |
| + if window > 0 { |
| + let pct = context_input.saturating_mul(100) / window; |
| + if pct >= 25 { |
| + let mut line = |
| + format!("- context window: ~{pct}% used ({context_input}/{window} input tokens)"); |
| + if pct >= 75 { |
| + line.push_str( |
| + " — running full; converge and finish before the earliest context is compacted away.", |
| + ); |
| + } |
| + snippets.push(MemorySnippet { |
| + source: "budget".to_string(), |
| + content: line, |
| + relevance_score: 0.95, |
| + }); |
| + } |
| + } |
| + |
| + |
| + |
| + if let Some(runtime_view) = ctx.input.runtime_view.as_ref() { |
| + for line in tool::open_todo_lines(runtime_view.as_ref()) { |
| + snippets.push(MemorySnippet { |
| + source: "plan".to_string(), |
| + content: line, |
| + relevance_score: 0.9, |
| + }); |
| + } |
| + } |
| + |
| + snippets |
| +} |
| + |
| async fn build_messages(ctx: &mut LoopContext<'_>) -> Result<(), AgentError> { | async fn build_messages(ctx: &mut LoopContext<'_>) -> Result<(), AgentError> { |
| let skill_summaries = ctx.snapshot.skill_registry.list_skills(); | let skill_summaries = ctx.snapshot.skill_registry.list_skills(); |
| | |
| @@ -733,12 +819,23 @@ async fn build_messages(ctx: &mut LoopContext<'_>) -> Result<(), AgentError> { |
| None | None |
| }; | }; |
| | |
| + let is_final_turn = |
| + ctx.snapshot.max_turns > 0 && ctx.turn.turn_number >= ctx.snapshot.max_turns; |
| + let visible_tools = if is_final_turn { |
| + Vec::new() |
| + } else { |
| + ctx.input.visible_tools.clone() |
| + }; |
| + |
| + let mut projected_messages = ctx.state.messages.read().clone(); |
| + prune_stale_tool_output(&mut projected_messages); |
| + |
| let input = PromptBuildInput { | let input = PromptBuildInput { |
| system_prompt: ctx.snapshot.system_prompt.to_string(), | system_prompt: ctx.snapshot.system_prompt.to_string(), |
| - messages: ctx.state.messages.read().clone(), | + messages: projected_messages, |
| - visible_tools: ctx.input.visible_tools.clone(), | + visible_tools, |
| skill_summaries, | skill_summaries, |
| - memory_snippets: Vec::new(), | + memory_snippets: live_context_snippets(ctx), |
| environment: agent_types::context::prompt::EnvironmentInfo { | environment: agent_types::context::prompt::EnvironmentInfo { |
| model: String::new(), | model: String::new(), |
| cwd: String::new(), | cwd: String::new(), |
| @@ -1018,7 +1115,11 @@ const TRANSIENT_MAX_DELAY_MS: u64 = 60_000; |
| fn is_transient(error: &LlmError) -> bool { | fn is_transient(error: &LlmError) -> bool { |
| matches!( | matches!( |
| error, | error, |
| - LlmError::RateLimited { .. } | LlmError::HttpError(_) | LlmError::Timeout | + LlmError::RateLimited { .. } |
| + | LlmError::HttpError(_) |
| + | LlmError::Timeout |
| + | LlmError::StreamError { .. } |
| + | LlmError::IoError(_) |
| ) | ) |
| } | } |
| | |
| @@ -1110,7 +1211,7 @@ fn stream_assistant_chunk( |
| } | } |
| } | } |
| | |
| -async fn tool_exec(ctx: &mut LoopContext<'_>) -> Result<Option<SuspendedToolCall>, AgentError> { | +async fn tool_exec(ctx: &mut LoopContext<'_>) -> Result<Vec<SuspendedToolCall>, AgentError> { |
| let has_tool_calls = ctx | let has_tool_calls = ctx |
| .turn | .turn |
| .assistant_message | .assistant_message |
| @@ -1118,22 +1219,23 @@ async fn tool_exec(ctx: &mut LoopContext<'_>) -> Result<Option<SuspendedToolCall |
| .map_or(false, |m| m.has_tool_calls()); | .map_or(false, |m| m.has_tool_calls()); |
| | |
| if ctx.turn.assistant_message.is_none() { | if ctx.turn.assistant_message.is_none() { |
| - return Ok(None); | + return Ok(Vec::new()); |
| } | } |
| | |
| if !has_tool_calls || !ctx.snapshot.feature_flags.tool_execution { | if !has_tool_calls || !ctx.snapshot.feature_flags.tool_execution { |
| append_assistant_to_history(ctx); | append_assistant_to_history(ctx); |
| - return Ok(None); | + return Ok(Vec::new()); |
| } | } |
| | |
| if ctx.input.runtime_view.is_none() { | if ctx.input.runtime_view.is_none() { |
| append_assistant_to_history(ctx); | append_assistant_to_history(ctx); |
| - return Ok(None); | + return Ok(Vec::new()); |
| } | } |
| | |
| | |
| if let Some(msg) = ctx.turn.assistant_message.as_mut() { | if let Some(msg) = ctx.turn.assistant_message.as_mut() { |
| synthesize_missing_call_ids(msg, ctx.state.turn_count); | synthesize_missing_call_ids(msg, ctx.state.turn_count); |
| + repair_tool_names(msg, &ctx.input.visible_tools); |
| } | } |
| | |
| let tool_calls: Vec<ToolUseBlock> = ctx | let tool_calls: Vec<ToolUseBlock> = ctx |
| @@ -1145,7 +1247,7 @@ async fn tool_exec(ctx: &mut LoopContext<'_>) -> Result<Option<SuspendedToolCall |
| .clone(); | .clone(); |
| | |
| if ctx.input.agent_id.is_none() { | if ctx.input.agent_id.is_none() { |
| - return Ok(None); | + return Ok(Vec::new()); |
| } | } |
| | |
| | |
| @@ -1209,6 +1311,9 @@ async fn tool_exec(ctx: &mut LoopContext<'_>) -> Result<Option<SuspendedToolCall |
| (valid_calls, invalid_calls) | (valid_calls, invalid_calls) |
| }; | }; |
| | |
| + let mut valid_calls = valid_calls; |
| + valid_calls.sort_by_key(|tc| tc.tool_name == "join_subagent"); |
| + |
| if let Some(msg) = ctx.turn.assistant_message.as_mut() { | if let Some(msg) = ctx.turn.assistant_message.as_mut() { |
| msg.tool_calls = valid_calls.clone(); | msg.tool_calls = valid_calls.clone(); |
| } | } |
| @@ -1250,9 +1355,8 @@ async fn tool_exec(ctx: &mut LoopContext<'_>) -> Result<Option<SuspendedToolCall |
| } | } |
| } | } |
| | |
| - // Execute valid tool calls (original logic). | + // Pass 1 — build every call (borrows ctx for the per-call tool filter). |
| - let runtime_view = ctx.input.runtime_view.as_ref().unwrap(); | + let mut built = Vec::with_capacity(valid_calls.len()); |
| - | |
| for tc in &valid_calls { | for tc in &valid_calls { |
| let raw_tool_call = RawToolCall { | let raw_tool_call = RawToolCall { |
| call_id: tc.call_id.clone(), | call_id: tc.call_id.clone(), |
| @@ -1270,56 +1374,196 @@ async fn tool_exec(ctx: &mut LoopContext<'_>) -> Result<Option<SuspendedToolCall |
| ctx.snapshot.tool_registry.as_ref(), | ctx.snapshot.tool_registry.as_ref(), |
| ); | ); |
| | |
| - let tool_call = match ToolCallBuilderImpl::new() | + match ToolCallBuilderImpl::new() |
| .with_raw_llm_tool_call(raw_tool_call) | .with_raw_llm_tool_call(raw_tool_call) |
| .with_tool_filter(per_call_filter) | .with_tool_filter(per_call_filter) |
| .build() | .build() |
| { | { |
| - Ok(tool_call) => tool_call, | + Ok(tool_call) => built.push(Ok(tool_call)), |
| - Err(error) => { | + Err(error) => built.push(Err(build_framework_failed_tool_result( |
| - let result = build_framework_failed_tool_result( | + fallback_final_call, |
| - fallback_final_call, | + format!("tool call build failed: {error}"), |
| - format!("tool call build failed: {error}"), | + ))), |
| - ); | + } |
| - emit_tool_result_event(ctx, &result); | + } |
| - let tool_result_message = build_tool_result_message(&result); | |
| - ctx.state.messages.write().push(tool_result_message); | |
| - ctx.turn.tool_results.push(result); | |
| - continue; | |
| - } | |
| - }; | |
| | |
| - let result = match tool_call.execute(&**runtime_view).await { | + let serialize_batch = { |
| - Ok(result) => result, | + let profiles: std::collections::HashMap<&str, &EffectProfile> = ctx |
| - Err(error) => { | + .input |
| - let result = | + .visible_tools |
| - build_framework_failed_tool_result(fallback_final_call, error.to_string()); | + .iter() |
| - emit_tool_result_event(ctx, &result); | + .map(|tool| (tool.name().0.as_str(), tool.effect_profile())) |
| - let tool_result_message = build_tool_result_message(&result); | + .collect(); |
| - ctx.state.messages.write().push(tool_result_message); | + !built.iter().filter_map(|b| b.as_ref().ok()).all(|call| { |
| - ctx.turn.tool_results.push(result); | + profiles |
| - continue; | + .get(call.final_call().tool_name.as_str()) |
| + .is_some_and(|profile| is_parallel_safe(profile)) |
| + }) |
| + }; |
| + |
| + |
| + |
| + |
| + let runtime_view = ctx.input.runtime_view.clone().unwrap(); |
| + let exec_outcomes: Vec<_> = if serialize_batch { |
| + let mut outcomes = Vec::new(); |
| + for call in built.iter().filter_map(|b| b.as_ref().ok()) { |
| + outcomes.push(call.execute(&*runtime_view).await); |
| + } |
| + outcomes |
| + } else { |
| + futures_util::future::join_all( |
| + built |
| + .iter() |
| + .filter_map(|b| b.as_ref().ok().map(|call| call.execute(&*runtime_view))), |
| + ) |
| + .await |
| + }; |
| + |
| + |
| + |
| + let mut exec_outcomes = exec_outcomes.into_iter(); |
| + let mut results: Vec<ToolExecutionResult> = Vec::with_capacity(built.len()); |
| + for entry in built { |
| + match entry { |
| + Err(failed_result) => results.push(failed_result), |
| + Ok(tool_call) => { |
| + let result = match exec_outcomes.next().expect("one outcome per executed call") { |
| + Ok(result) => result, |
| + Err(error) => build_framework_failed_tool_result( |
| + tool_call.final_call().clone(), |
| + error.to_string(), |
| + ), |
| + }; |
| + results.push(result); |
| } | } |
| - }; | + } |
| - let should_stop_after_result = should_stop_after_tool_result(ctx, &result); | + } |
| + |
| + |
| + |
| + |
| + let mut streak_note: Option<String> = None; |
| + let mut suspended_calls: Vec<SuspendedToolCall> = Vec::new(); |
| + let mut stop_after_batch = false; |
| + for result in results { |
| + ctx.state.tool_executed = true; |
| + if should_stop_after_tool_result(ctx, &result) { |
| + stop_after_batch = true; |
| + } |
| emit_tool_result_event(ctx, &result); | emit_tool_result_event(ctx, &result); |
| | |
| if let Some(suspended_call) = SuspendedToolCall::from_tool_result(&result) { | if let Some(suspended_call) = SuspendedToolCall::from_tool_result(&result) { |
| + |
| + |
| ctx.turn.tool_results.push(result); | ctx.turn.tool_results.push(result); |
| - return Ok(Some(suspended_call)); | + suspended_calls.push(suspended_call); |
| + continue; |
| } | } |
| | |
| let tool_result_message = build_tool_result_message(&result); | let tool_result_message = build_tool_result_message(&result); |
| ctx.state.messages.write().push(tool_result_message); | ctx.state.messages.write().push(tool_result_message); |
| + |
| + |
| + if let Some(note) = update_tool_failure_streak(ctx, &result) { |
| + streak_note = Some(note); |
| + } |
| ctx.turn.tool_results.push(result); | ctx.turn.tool_results.push(result); |
| + } |
| | |
| - if should_stop_after_result { | + if stop_after_batch && suspended_calls.is_empty() { |
| - ctx.turn.force_return_complete = true; | + ctx.turn.force_return_complete = true; |
| - break; | + } |
| + |
| + |
| + |
| + |
| + if suspended_calls.is_empty() { |
| + if let Some(note) = streak_note { |
| + ctx.state.messages.write().push(ChatMessage::user(note)); |
| } | } |
| } | } |
| | |
| - Ok(None) | + Ok(suspended_calls) |
| +} |
| + |
| +const REPEATED_FAILURE_THRESHOLD: u32 = 3; |
| +const REPEATED_SUCCESS_THRESHOLD: u32 = 3; |
| + |
| +fn update_tool_failure_streak( |
| + ctx: &mut LoopContext<'_>, |
| + result: &ToolExecutionResult, |
| +) -> Option<String> { |
| + let sig = tool_call_signature(result); |
| + if is_failure_result(result) { |
| + ctx.state.last_success_sig = None; |
| + ctx.state.repeated_success_count = 0; |
| + if ctx.state.last_failure_sig == Some(sig) { |
| + ctx.state.repeated_failure_count += 1; |
| + } else { |
| + ctx.state.last_failure_sig = Some(sig); |
| + ctx.state.repeated_failure_count = 1; |
| + } |
| + if ctx.state.repeated_failure_count >= REPEATED_FAILURE_THRESHOLD { |
| + let count = ctx.state.repeated_failure_count; |
| + let tool = result.tool_name().to_string(); |
| + ctx.state.repeated_failure_count = 0; |
| + ctx.state.last_failure_sig = None; |
| + return Some(format!( |
| + "The `{tool}` call has now failed {count} times in a row with identical arguments. \ |
| + Stop retrying it unchanged — change approach: fix the arguments, read the relevant \ |
| + file or state to understand why it fails, or use a different tool to reach the goal." |
| + )); |
| + } |
| + return None; |
| + } |
| + ctx.state.last_failure_sig = None; |
| + ctx.state.repeated_failure_count = 0; |
| + if ctx.state.last_success_sig == Some(sig) { |
| + ctx.state.repeated_success_count += 1; |
| + } else { |
| + ctx.state.last_success_sig = Some(sig); |
| + ctx.state.repeated_success_count = 1; |
| + } |
| + if ctx.state.repeated_success_count >= REPEATED_SUCCESS_THRESHOLD { |
| + let count = ctx.state.repeated_success_count; |
| + let tool = result.tool_name().to_string(); |
| + ctx.state.repeated_success_count = 0; |
| + ctx.state.last_success_sig = None; |
| + return Some(format!( |
| + "The `{tool}` call has now run {count} times in a row with identical arguments and the \ |
| + same result — that output is already in your context above. Stop repeating it: use \ |
| + what you have, or take a different action toward the goal." |
| + )); |
| + } |
| + None |
| +} |
| + |
| +fn is_parallel_safe(profile: &EffectProfile) -> bool { |
| + !profile.writes_filesystem |
| + && !profile.side_effects |
| + && (profile.reads_filesystem || profile.network_access) |
| +} |
| + |
| +fn is_failure_result(result: &ToolExecutionResult) -> bool { |
| + matches!( |
| + result, |
| + ToolExecutionResult::Completed { |
| + raw_outcome: RawToolOutcome::Error { .. }, |
| + .. |
| + } | ToolExecutionResult::Failed { .. } |
| + | ToolExecutionResult::Denied { .. } |
| + ) |
| +} |
| + |
| +fn tool_call_signature(result: &ToolExecutionResult) -> u64 { |
| + use std::hash::{Hash, Hasher}; |
| + let mut hasher = std::collections::hash_map::DefaultHasher::new(); |
| + result.tool_name().hash(&mut hasher); |
| + serde_json::to_string(&result.final_call().input) |
| + .unwrap_or_default() |
| + .hash(&mut hasher); |
| + hasher.finish() |
| } | } |
| | |
| |
| @@ -1333,6 +1577,41 @@ fn synthesize_missing_call_ids(msg: &mut AssistantMessage, turn: u32) { |
| } | } |
| } | } |
| | |
| +fn repair_tool_names( |
| + msg: &mut AssistantMessage, |
| + visible: &[std::sync::Arc<dyn agent_contracts::tool::ToolSpecView>], |
| +) { |
| + if visible.is_empty() { |
| + return; |
| + } |
| + let normalize = |s: &str| -> String { |
| + s.chars() |
| + .filter(|c| c.is_ascii_alphanumeric()) |
| + .map(|c| c.to_ascii_lowercase()) |
| + .collect() |
| + }; |
| + let mut canonical = std::collections::HashSet::new(); |
| + let mut normalized = std::collections::HashMap::new(); |
| + for tool in visible { |
| + let name = tool.name().0.clone(); |
| + normalized |
| + .entry(normalize(&name)) |
| + .or_insert_with(|| name.clone()); |
| + canonical.insert(name); |
| + } |
| + for tc in msg.tool_calls.iter_mut() { |
| + if canonical.contains(&tc.tool_name) { |
| + continue; |
| + } |
| + if let Some(fixed) = normalized.get(&normalize(&tc.tool_name)) { |
| + if *fixed != tc.tool_name { |
| + tracing::debug!(from = %tc.tool_name, to = %fixed, "repaired tool name"); |
| + tc.tool_name = fixed.clone(); |
| + } |
| + } |
| + } |
| +} |
| + |
| fn is_valid_tool_call(tc: &ToolUseBlock) -> bool { | fn is_valid_tool_call(tc: &ToolUseBlock) -> bool { |
| is_valid_tool_call_id(&tc.call_id) && is_valid_tool_name(&tc.tool_name) | is_valid_tool_call_id(&tc.call_id) && is_valid_tool_name(&tc.tool_name) |
| } | } |
| @@ -1588,6 +1867,52 @@ fn decide(ctx: &mut LoopContext<'_>) { |
| } | } |
| } | } |
| | |
| + |
| + |
| + |
| + |
| + if !ctx.state.plan_nudged && ctx.turn.turn_number < ctx.snapshot.max_turns { |
| + let open = ctx |
| + .input |
| + .runtime_view |
| + .as_ref() |
| + .map(|runtime_view| tool::open_todo_lines(runtime_view.as_ref())) |
| + .unwrap_or_default(); |
| + if !open.is_empty() { |
| + ctx.state.plan_nudged = true; |
| + let reminder = format!( |
| + "You are about to stop, but your plan still has {} open item(s):\n{}\n\ |
| + Finish them now, or call todo_write to mark them completed/cancelled if they no longer apply — then stop.", |
| + open.len(), |
| + open.join("\n") |
| + ); |
| + ctx.state.messages.write().push(ChatMessage::user(reminder)); |
| + ctx.turn.decision = Some(LoopDecision::Continue); |
| + return; |
| + } |
| + } |
| + |
| + |
| + |
| + |
| + if !ctx.state.completion_nudged |
| + && ctx.turn.turn_number < ctx.snapshot.max_turns |
| + && ctx.state.tool_executed |
| + { |
| + ctx.state.completion_nudged = true; |
| + let checklist = "You are about to finish. Before you stop, re-read the ORIGINAL task and verify, do not assume:\n\ |
| + 1. Every requirement it states is met — including any exact error message, return value, output, or edge case it names; if it specifies a behavior, you have a check that exercises THAT behavior, not a different one that merely passes.\n\ |
| + 2. Your change is robust to changed inputs — different numbers, empty/None/zero, other files or config — not only the one case you tried, and it does not mutate shared state or leave unintended side effects.\n\ |
| + 3. Review the change once from three angles: as the test engineer who will grade it, as a QA reviewer hunting regressions, and as the user who filed the task.\n\ |
| + If any check fails, fix it now. If all hold, stop again and you are done."; |
| + ctx.state |
| + .messages |
| + .write() |
| + .push(ChatMessage::user(checklist.to_string())); |
| + ctx.turn.decision = Some(LoopDecision::Continue); |
| + return; |
| + } |
| + |
| ctx.turn.decision = Some(LoopDecision::ReturnComplete); | ctx.turn.decision = Some(LoopDecision::ReturnComplete); |
| } | } |
| | |
| @@ -1836,9 +2161,10 @@ mod tests { |
| use std::sync::{Arc, Mutex as StdMutex}; | use std::sync::{Arc, Mutex as StdMutex}; |
| | |
| use agent_contracts::context::budget::TokenBudgetPolicy; | use agent_contracts::context::budget::TokenBudgetPolicy; |
| - use agent_contracts::tool::ToolSpecView; | + use agent_contracts::tool::{ToolExecutor, ToolFilter, ToolRegistry, ToolSpecView}; |
| use agent_contracts::{ | use agent_contracts::{ |
| - CompressionPipeline, LlmProvider, PromptBuilder, ProviderCapabilities, SkillRegistry, | + CompressionPipeline, LlmProvider, PromptBuilder, ProviderCapabilities, RuntimeView, |
| + SkillRegistry, |
| }; | }; |
| use agent_llm::LlmRequestExt; | use agent_llm::LlmRequestExt; |
| use agent_types::common::ids::{AgentId, ToolId, ToolName}; | use agent_types::common::ids::{AgentId, ToolId, ToolName}; |
| @@ -1846,7 +2172,9 @@ mod tests { |
| use agent_types::context::prompt::{PromptBuildError, PromptBuildResult}; | use agent_types::context::prompt::{PromptBuildError, PromptBuildResult}; |
| use agent_types::context::{FeatureFlags, TokenBudgetConfig}; | use agent_types::context::{FeatureFlags, TokenBudgetConfig}; |
| use agent_types::events::LoopEndSummary; | use agent_types::events::LoopEndSummary; |
| + use agent_types::tool::execution_types::{ToolExecutionError, ToolExecutorOutput}; |
| use agent_types::tool::spec_types::{EffectProfile, InputSchemaRef, OutputContract}; | use agent_types::tool::spec_types::{EffectProfile, InputSchemaRef, OutputContract}; |
| + use agent_types::tool::FinalToolCall; |
| use agent_types::{ | use agent_types::{ |
| AssistantMessage, LlmError, LlmRequest, LlmResponse, StopReason, StreamChunk, ToolUseBlock, | AssistantMessage, LlmError, LlmRequest, LlmResponse, StopReason, StreamChunk, ToolUseBlock, |
| Usage, | Usage, |
| @@ -2313,11 +2641,18 @@ mod tests { |
| fn test_runtime_with_max_turns( | fn test_runtime_with_max_turns( |
| provider: Arc<LlmProviderWrapper>, | provider: Arc<LlmProviderWrapper>, |
| max_turns: u32, | max_turns: u32, |
| + ) -> AgentRuntime { |
| + test_runtime_with_registry(provider, max_turns, Arc::new(EmptyToolRegistry::new())) |
| + } |
| + |
| + fn test_runtime_with_registry( |
| + provider: Arc<LlmProviderWrapper>, |
| + max_turns: u32, |
| + tool_registry: Arc<dyn ToolRegistry>, |
| ) -> AgentRuntime { | ) -> AgentRuntime { |
| let prompt_builder: Arc<dyn PromptBuilder> = Arc::new(FixedPromptBuilder); | let prompt_builder: Arc<dyn PromptBuilder> = Arc::new(FixedPromptBuilder); |
| let compression_pipeline: Arc<dyn CompressionPipeline> = | let compression_pipeline: Arc<dyn CompressionPipeline> = |
| Arc::new(compact::PassthroughCompressionPipeline::new()); | Arc::new(compact::PassthroughCompressionPipeline::new()); |
| - let tool_registry = Arc::new(EmptyToolRegistry::new()); | |
| let skill_registry: Arc<dyn SkillRegistry> = Arc::new(EmptySkillRegistry::new()); | let skill_registry: Arc<dyn SkillRegistry> = Arc::new(EmptySkillRegistry::new()); |
| let budget_config = TokenBudgetConfig { | let budget_config = TokenBudgetConfig { |
| total_budget: 4096, | total_budget: 4096, |
| @@ -2628,7 +2963,10 @@ mod tests { |
| outcome, | outcome, |
| LoopRunResult::Complete(AgentOutcome::Complete { .. }) | LoopRunResult::Complete(AgentOutcome::Complete { .. }) |
| )); | )); |
| - assert_eq!(loop_state.turn_count, 2); | + // turn 1: synthesize + run the tool; turn 2: model stops and the |
| + |
| + |
| + assert_eq!(loop_state.turn_count, 3); |
| | |
| let messages = loop_state.messages.read(); | let messages = loop_state.messages.read(); |
| let tool_use = messages.iter().find_map(|m| { | let tool_use = messages.iter().find_map(|m| { |
| @@ -2826,4 +3164,303 @@ mod tests { |
| let secrets3 = extract_secrets_from_messages(&vec![message_other_tool]); | let secrets3 = extract_secrets_from_messages(&vec![message_other_tool]); |
| assert_eq!(secrets3.len(), 0); | assert_eq!(secrets3.len(), 0); |
| } | } |
| + |
| + #[test] |
| + fn is_parallel_safe_allows_pure_readers_only() { |
| + let reader = EffectProfile { |
| + reads_filesystem: true, |
| + writes_filesystem: false, |
| + network_access: false, |
| + side_effects: false, |
| + }; |
| + let network_reader = EffectProfile { |
| + reads_filesystem: false, |
| + writes_filesystem: false, |
| + network_access: true, |
| + side_effects: false, |
| + }; |
| + assert!(is_parallel_safe(&reader)); |
| + assert!(is_parallel_safe(&network_reader)); |
| + } |
| + |
| + #[test] |
| + fn is_parallel_safe_serializes_writers_side_effects_and_interactive() { |
| + let writer = EffectProfile { |
| + reads_filesystem: true, |
| + writes_filesystem: true, |
| + network_access: false, |
| + side_effects: false, |
| + }; |
| + let side_effecting = EffectProfile { |
| + reads_filesystem: true, |
| + writes_filesystem: true, |
| + network_access: false, |
| + side_effects: true, |
| + }; |
| + let stateful = EffectProfile { |
| + reads_filesystem: false, |
| + writes_filesystem: false, |
| + network_access: false, |
| + side_effects: true, |
| + }; |
| + let interactive = EffectProfile::default(); |
| + assert!(!is_parallel_safe(&writer)); |
| + assert!(!is_parallel_safe(&side_effecting)); |
| + assert!(!is_parallel_safe(&stateful)); |
| + assert!( |
| + !is_parallel_safe(&interactive), |
| + "an interactive prompt declares no read/network and must serialize" |
| + ); |
| + } |
| + |
| + #[tokio::test] |
| + async fn completion_nudge_skips_conversational_run_with_visible_tools() { |
| + let provider = Arc::new(LlmProviderWrapper::new( |
| + Arc::new(StreamingTestProvider::new()), |
| + None, |
| + None, |
| + )); |
| + let runtime = test_runtime(provider); |
| + let input = AgentLoopInput::new("explain how this code works") |
| + .with_agent_id(AgentId("test-agent".to_string())) |
| + .with_visible_tools(dummy_visible_tools()) |
| + .with_runtime_view(Arc::new(NoopRuntimeView::new())); |
| + let mut loop_state = LoopState::new(uuid::Uuid::new_v4()); |
| + |
| + let outcome = run_agent_loop(&runtime, &mut loop_state, input) |
| + .await |
| + .expect("conversational loop should complete"); |
| + |
| + assert!(matches!( |
| + outcome, |
| + LoopRunResult::Complete(AgentOutcome::Complete { .. }) |
| + )); |
| + assert!( |
| + !loop_state.tool_executed, |
| + "no tool ran, so the run is conversational" |
| + ); |
| + assert_eq!(loop_state.turn_count, 1); |
| + } |
| + |
| + struct AlwaysSucceedsExecutor { |
| + spec: Arc<VisibleToolSpec>, |
| + } |
| + |
| + #[async_trait] |
| + impl ToolExecutor for AlwaysSucceedsExecutor { |
| + fn spec(&self) -> &dyn ToolSpecView { |
| + self.spec.as_ref() |
| + } |
| + |
| + async fn invoke( |
| + &self, |
| + call: &FinalToolCall, |
| + _runtime: &dyn RuntimeView, |
| + ) -> Result<ToolExecutorOutput, ToolExecutionError> { |
| + Ok(ToolExecutorOutput::Completed { |
| + raw_outcome: RawToolOutcome::Success { |
| + output: format!("ran {}", call.call_id), |
| + }, |
| + }) |
| + } |
| + } |
| + |
| + struct SingleToolRegistry { |
| + spec: Arc<VisibleToolSpec>, |
| + executor: Arc<dyn ToolExecutor>, |
| + } |
| + |
| + impl SingleToolRegistry { |
| + fn new() -> Self { |
| + let spec = Arc::new(VisibleToolSpec { |
| + id: ToolId("tool.peek".to_string()), |
| + name: ToolName("peek".to_string()), |
| + description: "Read-only peek".to_string(), |
| + input_schema: InputSchemaRef { |
| + schema: serde_json::json!({"type": "object"}), |
| + }, |
| + output_contract: OutputContract { |
| + description: "peeked".to_string(), |
| + }, |
| + effect_profile: EffectProfile { |
| + reads_filesystem: true, |
| + writes_filesystem: false, |
| + network_access: false, |
| + side_effects: false, |
| + }, |
| + }); |
| + let executor: Arc<dyn ToolExecutor> = Arc::new(AlwaysSucceedsExecutor { |
| + spec: Arc::clone(&spec), |
| + }); |
| + Self { spec, executor } |
| + } |
| + |
| + fn visible(&self) -> Vec<Arc<dyn ToolSpecView>> { |
| + vec![Arc::clone(&self.spec) as Arc<dyn ToolSpecView>] |
| + } |
| + } |
| + |
| + impl ToolRegistry for SingleToolRegistry { |
| + fn get_executor(&self, id: &ToolId) -> Option<Arc<dyn ToolExecutor>> { |
| + (id == self.spec.id()).then(|| Arc::clone(&self.executor)) |
| + } |
| + |
| + fn get_spec(&self, id: &ToolId) -> Option<&dyn ToolSpecView> { |
| + (id == self.spec.id()).then(|| self.spec.as_ref() as &dyn ToolSpecView) |
| + } |
| + |
| + fn list_specs(&self) -> Vec<&dyn ToolSpecView> { |
| + vec![self.spec.as_ref()] |
| + } |
| + |
| + fn filter_for(&self, _agent_id: &AgentId) -> Box<dyn ToolFilter> { |
| + tool_filter_from_specs(&self.visible(), self) |
| + } |
| + } |
| + |
| + struct TwoToolCallProvider { |
| + capabilities: ProviderCapabilities, |
| + calls: Arc<StdMutex<usize>>, |
| + } |
| + |
| + impl TwoToolCallProvider { |
| + fn new() -> Self { |
| + Self { |
| + capabilities: ProviderCapabilities { |
| + supports_streaming: true, |
| + supports_tool_calls: true, |
| + supports_json_mode: false, |
| + max_context_window: 4096, |
| + model_name: "two-tool-call-test".to_string(), |
| + }, |
| + calls: Arc::new(StdMutex::new(0)), |
| + } |
| + } |
| + } |
| + |
| + #[async_trait] |
| + impl LlmProvider for TwoToolCallProvider { |
| + async fn complete(&self, _request: &LlmRequest) -> Result<LlmResponse, LlmError> { |
| + panic!("streaming path should use complete_stream instead of complete"); |
| + } |
| + |
| + async fn complete_stream( |
| + &self, |
| + _request: &LlmRequest, |
| + _on_chunk: &(dyn Fn(StreamChunk) + Send + Sync), |
| + ) -> Result<LlmResponse, LlmError> { |
| + let call_number = { |
| + let mut calls = self.calls.lock().expect("call counter mutex poisoned"); |
| + *calls += 1; |
| + *calls |
| + }; |
| + |
| + if call_number == 1 { |
| + return Ok(LlmResponse { |
| + message: AssistantMessage { |
| + text: Some("calling tools".to_string()), |
| + reasoning_content: None, |
| + tool_calls: vec![ |
| + ToolUseBlock { |
| + call_id: "call_a".to_string(), |
| + tool_name: "peek".to_string(), |
| + input: serde_json::json!({}), |
| + }, |
| + ToolUseBlock { |
| + call_id: "call_b".to_string(), |
| + tool_name: "peek".to_string(), |
| + input: serde_json::json!({}), |
| + }, |
| + ], |
| + usage: Usage { |
| + prompt_tokens: 5, |
| + completion_tokens: 3, |
| + total_tokens: 8, |
| + }, |
| + stop_reason: StopReason::ToolUse, |
| + }, |
| + kv_cache_chunk_hashes: vec![], |
| + }); |
| + } |
| + |
| + Ok(LlmResponse { |
| + message: AssistantMessage { |
| + text: Some("done".to_string()), |
| + reasoning_content: None, |
| + tool_calls: Vec::new(), |
| + usage: Usage { |
| + prompt_tokens: 5, |
| + completion_tokens: 1, |
| + total_tokens: 6, |
| + }, |
| + stop_reason: StopReason::EndTurn, |
| + }, |
| + kv_cache_chunk_hashes: vec![], |
| + }) |
| + } |
| + |
| + fn capabilities(&self) -> &ProviderCapabilities { |
| + &self.capabilities |
| + } |
| + } |
| + |
| + #[tokio::test] |
| + async fn mid_batch_stop_still_records_every_executed_result() { |
| + let registry = Arc::new(SingleToolRegistry::new()); |
| + let visible = registry.visible(); |
| + let provider = Arc::new(LlmProviderWrapper::new( |
| + Arc::new(TwoToolCallProvider::new()), |
| + None, |
| + None, |
| + )); |
| + let runtime = test_runtime_with_registry(provider, 4, registry); |
| + let input = AgentLoopInput::new("go") |
| + .with_agent_id(AgentId("test-agent".to_string())) |
| + .with_visible_tools(visible) |
| + .with_runtime_view(Arc::new(NoopRuntimeView::new())) |
| + .with_stop_rules([LoopStopRule::AfterSuccessfulTool { |
| + tool_name: "peek".to_string(), |
| + }]); |
| + let mut loop_state = LoopState::new(uuid::Uuid::new_v4()); |
| + |
| + let outcome = run_agent_loop(&runtime, &mut loop_state, input) |
| + .await |
| + .expect("loop should complete via the stop rule"); |
| + |
| + assert!(matches!( |
| + outcome, |
| + LoopRunResult::Complete(AgentOutcome::Complete { .. }) |
| + )); |
| + assert_eq!(loop_state.turn_count, 1); |
| + |
| + let messages = loop_state.messages.read(); |
| + let tool_use_ids: Vec<String> = messages |
| + .iter() |
| + .flat_map(|m| m.blocks.iter()) |
| + .filter_map(|b| match b { |
| + ContentBlock::ToolUse { call_id, .. } => Some(call_id.clone()), |
| + _ => None, |
| + }) |
| + .collect(); |
| + let tool_result_ids: Vec<String> = messages |
| + .iter() |
| + .flat_map(|m| m.blocks.iter()) |
| + .filter_map(|b| match b { |
| + ContentBlock::ToolResult { call_id, .. } => Some(call_id.clone()), |
| + _ => None, |
| + }) |
| + .collect(); |
| + |
| + assert_eq!( |
| + tool_use_ids, |
| + vec!["call_a".to_string(), "call_b".to_string()], |
| + "both tool calls should be in history" |
| + ); |
| + assert_eq!( |
| + tool_result_ids, |
| + vec!["call_a".to_string(), "call_b".to_string()], |
| + "every executed tool_use must keep its paired tool_result even when an \ |
| + earlier call in the batch triggered the stop rule" |
| + ); |
| + } |
| } | } |
| |