| 1 | //! Admission gates between a provider response and tool execution: a failed |
| 2 | //! stream (C02-05), the step-budget final report (C02-10), the stream content |
| 3 | //! cap over tool arguments (C02-13), and bounded tool errors (C02-12). Every |
| 4 | //! fixture counts handler entries; none runs a shell or reaches a provider. |
| 5 | |
| 6 | use super::*; |
| 7 | use crate::llm_client::mock::{MockLlmClient, canned}; |
| 8 | use crate::tools::spec::{ApprovalRequirement, ToolCapability, ToolSpec}; |
| 9 | use std::sync::atomic::{AtomicUsize, Ordering}; |
| 10 | |
| 11 | const WRITE_TOOL: &str = "fixture_write"; |
| 12 | |
| 13 | /// A mutating tool that needs no approval, so a call that is admitted runs. |
| 14 | struct CountingWriteTool(Arc<AtomicUsize>); |
| 15 | |
| 16 | #[async_trait::async_trait] |
| 17 | impl ToolSpec for CountingWriteTool { |
| 18 | fn name(&self) -> &str { |
| 19 | WRITE_TOOL |
| 20 | } |
| 21 | fn description(&self) -> &str { |
| 22 | "Count every execution of a mutating fixture." |
| 23 | } |
| 24 | fn input_schema(&self) -> Value { |
| 25 | json!({"type": "object"}) |
| 26 | } |
| 27 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 28 | vec![ToolCapability::WritesFiles] |
| 29 | } |
| 30 | fn approval_requirement(&self) -> ApprovalRequirement { |
| 31 | ApprovalRequirement::Auto |
| 32 | } |
| 33 | async fn execute(&self, _: Value, _: &ToolContext) -> Result<ToolResult, ToolError> { |
| 34 | self.0.fetch_add(1, Ordering::SeqCst); |
| 35 | Ok(ToolResult::success("wrote")) |
| 36 | } |
| 37 | } |
| 38 | |
| 39 | fn write_surface( |
| 40 | engine: &Engine, |
| 41 | workspace: &Path, |
| 42 | executions: Arc<AtomicUsize>, |
| 43 | ) -> ToolSurfacePolicy { |
| 44 | let mut registry = crate::tools::ToolRegistry::new(ToolContext::new(workspace)); |
| 45 | registry.register(Arc::new(CountingWriteTool(executions))); |
| 46 | let tools = Some(registry.to_api_tools_with_cache(true)); |
| 47 | test_tool_surface(engine, registry, tools, AppMode::Agent) |
| 48 | } |
| 49 | |
| 50 | /// Streams a scripted response per call, `Err` items included, which the |
| 51 | /// queue-driven mock cannot express. |
| 52 | struct ScriptedStreamClient { |
| 53 | calls: AtomicUsize, |
| 54 | script: fn(usize) -> Vec<anyhow::Result<StreamEvent>>, |
| 55 | } |
| 56 | |
| 57 | #[async_trait::async_trait] |
| 58 | impl crate::core::model_client::ModelClient for ScriptedStreamClient { |
| 59 | fn provider_name(&self) -> &str { |
| 60 | "scripted" |
| 61 | } |
| 62 | fn model(&self) -> &str { |
| 63 | "scripted-model" |
| 64 | } |
| 65 | async fn create_message( |
| 66 | &self, |
| 67 | _request: codewhale_models::MessageRequest, |
| 68 | ) -> anyhow::Result<codewhale_models::MessageResponse> { |
| 69 | anyhow::bail!("scripted fixture only streams") |
| 70 | } |
| 71 | async fn create_message_stream( |
| 72 | &self, |
| 73 | _request: codewhale_models::MessageRequest, |
| 74 | ) -> anyhow::Result<crate::llm_client::StreamEventBox> { |
| 75 | let call = self.calls.fetch_add(1, Ordering::SeqCst); |
| 76 | Ok(Box::pin(futures_util::stream::iter((self.script)(call)))) |
| 77 | } |
| 78 | async fn health_check(&self) -> anyhow::Result<bool> { |
| 79 | Ok(true) |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | async fn drain_events(handle: &EngineHandle) -> Vec<Event> { |
| 84 | let mut rx = handle.rx_event.write().await; |
| 85 | std::iter::from_fn(|| rx.try_recv().ok()).collect() |
| 86 | } |
| 87 | |
| 88 | fn not_started(result: &Result<ToolResult, ToolError>) -> bool { |
| 89 | result.as_ref().is_ok_and(|output| { |
| 90 | !output.success |
| 91 | && output |
| 92 | .metadata |
| 93 | .as_ref() |
| 94 | .is_some_and(|metadata| metadata["side_effect_status"] == "not_started") |
| 95 | }) |
| 96 | } |
| 97 | |
| 98 | #[tokio::test] |
| 99 | async fn terminal_stream_error_after_a_complete_tool_call_never_reaches_the_handler() { |
| 100 | // C02-05: a complete mutating call followed by a provider error frame. |
| 101 | let workspace = tempdir().unwrap(); |
| 102 | let mock = Arc::new(MockLlmClient::new(vec![ |
| 103 | vec![ |
| 104 | canned::message_start("tool-then-error"), |
| 105 | canned::tool_use_block_start(0, "call-write", WRITE_TOOL), |
| 106 | canned::tool_input_delta(0, r#"{"path":"target.txt"}"#), |
| 107 | canned::block_stop(0), |
| 108 | StreamEvent::Error { |
| 109 | error: json!({ "message": "Model not exist." }), |
| 110 | }, |
| 111 | ], |
| 112 | canned::simple_text_turn("this second request must never be issued"), |
| 113 | ])); |
| 114 | let (mut engine, handle) = Engine::new_with_model_client( |
| 115 | deterministic_engine_config(workspace.path()), |
| 116 | &Config::default(), |
| 117 | mock.clone(), |
| 118 | ); |
| 119 | engine.session.auto_approve = true; |
| 120 | let executions = Arc::new(AtomicUsize::new(0)); |
| 121 | let surface = write_surface(&engine, workspace.path(), executions.clone()); |
| 122 | let mut turn = crate::core::turn::TurnContext::new(4); |
| 123 | |
| 124 | let (status, error) = engine.run_turn(&mut turn, surface, None, None).await; |
| 125 | |
| 126 | assert_eq!( |
| 127 | executions.load(Ordering::SeqCst), |
| 128 | 0, |
| 129 | "the handler must never run" |
| 130 | ); |
| 131 | assert_eq!(status, TurnOutcomeStatus::Failed); |
| 132 | assert!( |
| 133 | error |
| 134 | .as_deref() |
| 135 | .is_some_and(|e| e.contains("Model not exist.")), |
| 136 | "{error:?}" |
| 137 | ); |
| 138 | assert_eq!( |
| 139 | mock.call_count(), |
| 140 | 1, |
| 141 | "a failed response authorizes no further request" |
| 142 | ); |
| 143 | assert_eq!( |
| 144 | turn.stop_diagnostics.last_response_tool_calls_suppressed, |
| 145 | Some(1) |
| 146 | ); |
| 147 | let completions: Vec<_> = drain_events(&handle) |
| 148 | .await |
| 149 | .into_iter() |
| 150 | .filter_map(|event| match event { |
| 151 | Event::ToolCallComplete { result, .. } => Some(result), |
| 152 | _ => None, |
| 153 | }) |
| 154 | .collect(); |
| 155 | assert_eq!(completions.len(), 1, "the call is settled exactly once"); |
| 156 | assert!(not_started(&completions[0]), "{:?}", completions[0]); |
| 157 | assert!( |
| 158 | !engine |
| 159 | .session |
| 160 | .messages |
| 161 | .iter() |
| 162 | .flat_map(|message| &message.content) |
| 163 | .any(|block| matches!( |
| 164 | block, |
| 165 | ContentBlock::ToolUse { .. } | ContentBlock::ToolResult { .. } |
| 166 | )), |
| 167 | "no unpaired or fabricated tool history" |
| 168 | ); |
| 169 | } |
| 170 | |
| 171 | #[tokio::test] |
| 172 | async fn exhausted_stream_resume_never_executes_the_failed_batch() { |
| 173 | // C02-05, retry path: every attempt streams a complete call and then |
| 174 | // drops. Each resume discards its batch; the last, unretried attempt is |
| 175 | // refused at admission instead of executing. |
| 176 | fn tool_then_drop(_call: usize) -> Vec<anyhow::Result<StreamEvent>> { |
| 177 | vec![ |
| 178 | Ok(canned::message_start("tool-then-drop")), |
| 179 | Ok(canned::tool_use_block_start(0, "call-write", WRITE_TOOL)), |
| 180 | Ok(canned::tool_input_delta(0, "{}")), |
| 181 | Ok(canned::block_stop(0)), |
| 182 | Err(anyhow::anyhow!( |
| 183 | "Stream read error: error decoding response body" |
| 184 | )), |
| 185 | ] |
| 186 | } |
| 187 | let workspace = tempdir().unwrap(); |
| 188 | let client = Arc::new(ScriptedStreamClient { |
| 189 | calls: AtomicUsize::new(0), |
| 190 | script: tool_then_drop, |
| 191 | }); |
| 192 | let (mut engine, _handle) = Engine::new_with_model_client( |
| 193 | EngineConfig { |
| 194 | terminal_chrome_enabled: false, |
| 195 | ..deterministic_engine_config(workspace.path()) |
| 196 | }, |
| 197 | &Config::default(), |
| 198 | client.clone(), |
| 199 | ); |
| 200 | engine.session.auto_approve = true; |
| 201 | let resumes = engine.config.stream_retry_limits.max_resumes as usize; |
| 202 | let executions = Arc::new(AtomicUsize::new(0)); |
| 203 | let surface = write_surface(&engine, workspace.path(), executions.clone()); |
| 204 | let mut turn = crate::core::turn::TurnContext::new(4); |
| 205 | |
| 206 | let (status, error) = engine.run_turn(&mut turn, surface, None, None).await; |
| 207 | |
| 208 | assert_eq!( |
| 209 | executions.load(Ordering::SeqCst), |
| 210 | 0, |
| 211 | "no attempt's batch may run" |
| 212 | ); |
| 213 | assert_eq!(client.calls.load(Ordering::SeqCst), 1 + resumes); |
| 214 | assert_eq!(status, TurnOutcomeStatus::Failed); |
| 215 | assert!( |
| 216 | error |
| 217 | .as_deref() |
| 218 | .is_some_and(|e| e.contains("error decoding response body")), |
| 219 | "{error:?}" |
| 220 | ); |
| 221 | } |
| 222 | |
| 223 | #[tokio::test] |
| 224 | async fn tool_argument_bytes_count_toward_the_stream_content_cap() { |
| 225 | // C02-13: a response made only of tool-argument JSON must still hit the |
| 226 | // per-step content cap, and the capped call must not run. |
| 227 | let workspace = tempdir().unwrap(); |
| 228 | let huge_args = format!(r#"{{"blob":"{}"}}"#, "x".repeat(4096)); |
| 229 | let mock = Arc::new(MockLlmClient::new(vec![ |
| 230 | canned::tool_call_turn("call-big", WRITE_TOOL, &huge_args), |
| 231 | canned::simple_text_turn("this second request must never be issued"), |
| 232 | ])); |
| 233 | let (mut engine, _handle) = Engine::new_with_model_client( |
| 234 | EngineConfig { |
| 235 | stream_max_content_bytes: 64, |
| 236 | ..deterministic_engine_config(workspace.path()) |
| 237 | }, |
| 238 | &Config::default(), |
| 239 | mock.clone(), |
| 240 | ); |
| 241 | engine.session.auto_approve = true; |
| 242 | let executions = Arc::new(AtomicUsize::new(0)); |
| 243 | let surface = write_surface(&engine, workspace.path(), executions.clone()); |
| 244 | let mut turn = crate::core::turn::TurnContext::new(4); |
| 245 | |
| 246 | let (status, _error) = engine.run_turn(&mut turn, surface, None, None).await; |
| 247 | |
| 248 | assert_eq!(executions.load(Ordering::SeqCst), 0); |
| 249 | assert_eq!(status, TurnOutcomeStatus::Failed); |
| 250 | assert_eq!(mock.call_count(), 1); |
| 251 | } |
| 252 | |
| 253 | #[tokio::test] |
| 254 | async fn text_fallback_tool_calls_share_the_response_tool_limit() { |
| 255 | // C02-14: calls parsed from text markers obey the same per-response |
| 256 | // ceiling as native tool starts. An over-limit batch fails the turn |
| 257 | // before any call is announced, planned or run. |
| 258 | let workspace = tempdir().unwrap(); |
| 259 | let over_limit = super::super::streaming::MAX_TOOL_CALLS_PER_RESPONSE + 1; |
| 260 | let text: String = (0..over_limit) |
| 261 | .map(|_| format!(r#"[TOOL_CALL]{{"tool": "{WRITE_TOOL}", "args": {{}}}}[/TOOL_CALL]"#)) |
| 262 | .collect(); |
| 263 | let mock = Arc::new(MockLlmClient::new(vec![ |
| 264 | canned::simple_text_turn(&text), |
| 265 | canned::simple_text_turn("this second request must never be issued"), |
| 266 | ])); |
| 267 | let (mut engine, handle) = Engine::new_with_model_client( |
| 268 | deterministic_engine_config(workspace.path()), |
| 269 | &Config::default(), |
| 270 | mock.clone(), |
| 271 | ); |
| 272 | engine.session.auto_approve = true; |
| 273 | let executions = Arc::new(AtomicUsize::new(0)); |
| 274 | let surface = write_surface(&engine, workspace.path(), executions.clone()); |
| 275 | let mut turn = crate::core::turn::TurnContext::new(4); |
| 276 | |
| 277 | // Undrained, an admitted over-limit batch cannot even be announced: its |
| 278 | // 257 starts overflow the event queue. Bound the wait so that is a |
| 279 | // failure, not a hang. |
| 280 | let (status, _error) = tokio::time::timeout( |
| 281 | std::time::Duration::from_secs(30), |
| 282 | engine.run_turn(&mut turn, surface, None, None), |
| 283 | ) |
| 284 | .await |
| 285 | .expect("an over-limit text batch must be refused, not admitted"); |
| 286 | |
| 287 | assert_eq!(executions.load(Ordering::SeqCst), 0, "no parsed call runs"); |
| 288 | assert_eq!(status, TurnOutcomeStatus::Failed); |
| 289 | assert_eq!(mock.call_count(), 1); |
| 290 | assert_eq!( |
| 291 | turn.stop_diagnostics.last_response_tool_calls_suppressed, |
| 292 | Some(over_limit) |
| 293 | ); |
| 294 | assert!( |
| 295 | !drain_events(&handle) |
| 296 | .await |
| 297 | .iter() |
| 298 | .any(|event| matches!(event, Event::ToolCallStarted { .. })), |
| 299 | "no over-limit call is announced" |
| 300 | ); |
| 301 | } |
| 302 | |
| 303 | #[tokio::test] |
| 304 | async fn step_budget_final_report_is_report_only() { |
| 305 | // C02-10: the one response granted after the step budget is spent asks |
| 306 | // for no tools, and a call it returns anyway is refused, not executed. |
| 307 | let workspace = tempdir().unwrap(); |
| 308 | let mock = Arc::new(MockLlmClient::new(vec![ |
| 309 | canned::tool_call_turn("call-step-1", WRITE_TOOL, "{}"), |
| 310 | canned::tool_call_turn("call-final-report", WRITE_TOOL, "{}"), |
| 311 | ])); |
| 312 | let (mut engine, handle) = Engine::new_with_model_client( |
| 313 | deterministic_engine_config(workspace.path()), |
| 314 | &Config::default(), |
| 315 | mock.clone(), |
| 316 | ); |
| 317 | engine.session.auto_approve = true; |
| 318 | let executions = Arc::new(AtomicUsize::new(0)); |
| 319 | let surface = write_surface(&engine, workspace.path(), executions.clone()); |
| 320 | let mut turn = crate::core::turn::TurnContext::new(1); |
| 321 | |
| 322 | let (status, error) = engine.run_turn(&mut turn, surface, None, None).await; |
| 323 | |
| 324 | assert_eq!( |
| 325 | executions.load(Ordering::SeqCst), |
| 326 | 1, |
| 327 | "only the in-budget step runs" |
| 328 | ); |
| 329 | assert_eq!(status, TurnOutcomeStatus::Failed); |
| 330 | assert!( |
| 331 | error |
| 332 | .as_deref() |
| 333 | .is_some_and(|e| e.contains("Maximum model steps")), |
| 334 | "{error:?}" |
| 335 | ); |
| 336 | let requests = mock.captured_requests(); |
| 337 | assert_eq!(requests.len(), 2); |
| 338 | assert_ne!(requests[0].tool_choice, Some(json!("none"))); |
| 339 | assert_eq!(requests[1].tool_choice, Some(json!("none"))); |
| 340 | let refused = drain_events(&handle).await.into_iter().any(|event| { |
| 341 | matches!( |
| 342 | event, |
| 343 | Event::ToolCallComplete { |
| 344 | model_call: Some(ref call), |
| 345 | result: Err(ToolError::PermissionDenied { ref message }), |
| 346 | .. |
| 347 | } if call.provider_id == "call-final-report" && message.contains("final report") |
| 348 | ) |
| 349 | }); |
| 350 | assert!(refused, "the final-report call is refused with the reason"); |
| 351 | } |
| 352 | |
| 353 | #[test] |
| 354 | fn oversized_tool_errors_are_bounded_before_fanout() { |
| 355 | // C02-12: an error is published to the event stream and the session like |
| 356 | // a result, so an oversized one gets the same bounded projection. |
| 357 | struct HugeErrorTool { |
| 358 | as_result: bool, |
| 359 | } |
| 360 | #[async_trait::async_trait] |
| 361 | impl ToolSpec for HugeErrorTool { |
| 362 | fn name(&self) -> &str { |
| 363 | "fixture_huge_error" |
| 364 | } |
| 365 | fn description(&self) -> &str { |
| 366 | "Fail with an oversized message." |
| 367 | } |
| 368 | fn input_schema(&self) -> Value { |
| 369 | json!({"type": "object"}) |
| 370 | } |
| 371 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 372 | vec![ToolCapability::ReadOnly] |
| 373 | } |
| 374 | async fn execute(&self, _: Value, _: &ToolContext) -> Result<ToolResult, ToolError> { |
| 375 | let message = format!("FIRST{}LAST", "e".repeat(300_000)); |
| 376 | if self.as_result { |
| 377 | Ok(ToolResult::error(message)) |
| 378 | } else { |
| 379 | Err(ToolError::execution_failed(message)) |
| 380 | } |
| 381 | } |
| 382 | } |
| 383 | |
| 384 | with_artifact_home(|home| { |
| 385 | tokio::runtime::Builder::new_current_thread() |
| 386 | .enable_all() |
| 387 | .build() |
| 388 | .unwrap() |
| 389 | .block_on(async { |
| 390 | for as_result in [false, true] { |
| 391 | let mock = Arc::new(MockLlmClient::new(vec![ |
| 392 | canned::tool_call_turn("call-huge", "fixture_huge_error", "{}"), |
| 393 | canned::simple_text_turn("done"), |
| 394 | ])); |
| 395 | let (mut engine, handle) = Engine::new_with_model_client( |
| 396 | deterministic_engine_config(home), |
| 397 | &Config::default(), |
| 398 | mock.clone(), |
| 399 | ); |
| 400 | let mut registry = crate::tools::ToolRegistry::new(ToolContext::new(home)); |
| 401 | registry.register(Arc::new(HugeErrorTool { as_result })); |
| 402 | let tools = Some(registry.to_api_tools_with_cache(true)); |
| 403 | let surface = test_tool_surface(&engine, registry, tools, AppMode::Agent); |
| 404 | let mut turn = crate::core::turn::TurnContext::new(4); |
| 405 | let (status, error) = engine.run_turn(&mut turn, surface, None, None).await; |
| 406 | assert_eq!(status, TurnOutcomeStatus::Completed, "{error:?}"); |
| 407 | let published = drain_events(&handle) |
| 408 | .await |
| 409 | .into_iter() |
| 410 | .find_map(|event| match event { |
| 411 | Event::ToolCallComplete { result, .. } => Some(match result { |
| 412 | Ok(output) => output.content, |
| 413 | Err(error) => error.to_string(), |
| 414 | }), |
| 415 | _ => None, |
| 416 | }) |
| 417 | .expect("the call completes"); |
| 418 | assert!( |
| 419 | published.len() < crate::tools::truncate::SPILLOVER_THRESHOLD_BYTES, |
| 420 | "as_result={as_result}: {} bytes published", |
| 421 | published.len() |
| 422 | ); |
| 423 | assert!(published.contains("FIRST"), "the head survives"); |
| 424 | assert!(published.contains("LAST"), "the tail survives"); |
| 425 | assert!( |
| 426 | published.contains(crate::tools::truncate::SPILLOVER_RECOVERY_HINT), |
| 427 | "the omitted range stays recoverable" |
| 428 | ); |
| 429 | } |
| 430 | }); |
| 431 | }); |
| 432 | } |
| 433 |