From 36295edc08fd6dc6263aebd73c24646541de4db2 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 13:43:36 +0900 Subject: [PATCH] fix: require terminal stream completion --- crates/agen/src/engine.rs | 42 ++++++- crates/agen/tests/parallel_execution_test.rs | 105 +++++++++++++++++- .../src/feature/builtin/flow_transition.rs | 8 +- 3 files changed, 151 insertions(+), 4 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 8dd5825a3..737fac9a1 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1529,12 +1529,28 @@ impl Engine { if pause.is_some() { info!(call_id = %expected_tool_use_id, "Paused while awaiting pre-tool policy"); } + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + None, + )?; + batch.dispatched_any = true; return Ok(EarlyToolAdmission::Pause); } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!(call_id = %expected_tool_use_id, "Cancelled while awaiting pre-tool policy"); } + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + None, + )?; + batch.dispatched_any = true; return Err(EngineError::Cancelled); } }; @@ -1664,6 +1680,7 @@ impl Engine { } PreToolAction::Pause => { self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.dispatched_any = true; return Ok(EarlyToolAdmission::Pause); } } @@ -3192,6 +3209,7 @@ impl Engine { self.tool_call_collector.begin_response(); let mut stream = ResponseStreamPump::start(stream); let mut early_tools: Option = None; + let mut response_completed = false; let mut event_count: usize = 0; loop { @@ -3229,6 +3247,14 @@ impl Engine { json!({}), ); } + if matches!( + &event, + Event::Status(StatusEvent { + status: crate::llm_client::event::ResponseStatus::Completed, + }) + ) { + response_completed = true; + } self.emit_stream_event(turn, llm_call, &event); self.timeline.dispatch(&event); if early_dispatch { @@ -3375,7 +3401,21 @@ impl Engine { })); } } - None => break, + None if response_completed => break, + None => { + // A clean transport EOF is not an authoritative provider + // response boundary. Treat it exactly like an interrupted + // stream so any early side effect is terminalized and never + // followed by automatic continuation/regeneration. + self.timeline.flush_usage(); + return Ok(StreamResponseOutput { + completion: StreamCompletion::Interrupted { + reason: "LLM stream ended before provider completion status" + .to_string(), + }, + early_tools, + }); + } } } terminal = async { diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 744e4111a..3638f48d5 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -260,6 +260,7 @@ impl Tool for ContextRecordingTool { struct ControlledStreamClient { receiver: Arc>>>>, completion_support: ToolCallCompletionSupport, + stream_calls: Arc, } impl ControlledStreamClient { @@ -274,15 +275,21 @@ impl ControlledStreamClient { Self { receiver: Arc::new(Mutex::new(Some(rx))), completion_support, + stream_calls: Arc::new(AtomicUsize::new(0)), }, tx, ) } + + fn stream_count(&self) -> usize { + self.stream_calls.load(Ordering::SeqCst) + } } #[async_trait] impl LlmClient for ControlledStreamClient { async fn stream(&self, _request: Request) -> Result { + self.stream_calls.fetch_add(1, Ordering::SeqCst); let mut rx = self .receiver .lock() @@ -633,9 +640,29 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_pause() { pause.send(()).await.unwrap(); drop(tx); - let (output, _history) = run.await.unwrap(); + let (output, mut history) = run.await.unwrap(); assert!(matches!(output.result, EngineRunExit::Paused)); assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: None, + status: Some(agen::llm_client::ItemStatus::Completed), + .. + } if call_id == "call_blocked_policy" + ))); + + let mut restored = Engine::new(MockLlmClient::new(vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })])); + restored.register_tool(probe.definition()); + let _ = restored.resume(&mut history).await; + assert_eq!(probe.call_count(), 1); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_blocked_policy" + ))); } #[tokio::test] @@ -703,7 +730,7 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_cancel() { cancel.send(()).await.unwrap(); drop(tx); - let (output, _history) = tokio::time::timeout(Duration::from_secs(1), run) + let (output, mut history) = tokio::time::timeout(Duration::from_secs(1), run) .await .expect("cancel does not deadlock the bounded stream pump") .unwrap(); @@ -712,6 +739,26 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_cancel() { EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) )); assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: None, + status: Some(agen::llm_client::ItemStatus::Completed), + .. + } if call_id == "call_cancelled_policy" + ))); + + let mut restored = Engine::new(MockLlmClient::new(vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })])); + restored.register_tool(probe.definition()); + let _ = restored.resume(&mut history).await; + assert_eq!(probe.call_count(), 1); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_cancelled_policy" + ))); } #[tokio::test] @@ -1539,6 +1586,60 @@ async fn interrupted_stream_after_early_start_terminalizes_without_continuation_ ); } +#[tokio::test] +async fn clean_eof_without_provider_completion_after_early_start_stops_without_continuation() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let client_probe = client.clone(); + let tool = BarrierTool::new("barrier_clean_eof"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(tool.definition()); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start( + 0, + "call_clean_eof", + "barrier_clean_eof", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) + .await + .expect("tool starts before clean transport EOF"); + probe.release.notify_one(); + drop(tx); + + let (output, history) = run.await.unwrap(); + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Client( + ClientError::Api { + code: Some(ref code), + .. + } + ))) if code == "early_tool_stream_interrupted" + )); + assert_eq!( + client_probe.stream_count(), + 1, + "must not open a continuation request" + ); + assert_eq!(probe.starts.load(Ordering::SeqCst), 1); + assert_eq!( + history + .items() + .filter(|item| matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_clean_eof")) + .count(), + 1 + ); +} + #[tokio::test] async fn unsupported_provider_visibly_defers_opt_in_dispatch() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::ResponseComplete); diff --git a/crates/worker/src/feature/builtin/flow_transition.rs b/crates/worker/src/feature/builtin/flow_transition.rs index 163ac8928..b96d1919d 100644 --- a/crates/worker/src/feature/builtin/flow_transition.rs +++ b/crates/worker/src/feature/builtin/flow_transition.rs @@ -636,7 +636,7 @@ mod tests { use agen::llm_client::error::ClientError; use agen::llm_client::event::{ BlockDelta, BlockMetadata, BlockStart, BlockStop, BlockType, DeltaContent, - Event as LlmEvent, StopReason, + Event as LlmEvent, ResponseStatus, StatusEvent, StopReason, }; use agen::llm_client::types::Request; use agen::tool::ToolExecutionContext; @@ -906,6 +906,9 @@ mod tests { reasoning: None, stop_reason: Some(StopReason::ToolUse), })), + Ok(LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + })), ]))) } else { Ok(Box::pin(stream::iter(vec![ @@ -924,6 +927,9 @@ mod tests { reasoning: None, stop_reason: Some(StopReason::EndTurn), })), + Ok(LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + })), ]))) } }