Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: require terminal stream completion

Browse Source
This commit is contained in:
Hare
2026-09-29 13:43:36 +09:00
parent 1980da2880
commit 36295edc08
3 changed files with 151 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+41 -1
View File
@@ -1529,12 +1529,28 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
self.tool_call_collector.begin_response();
let mut stream = ResponseStreamPump::start(stream);
let mut early_tools: Option<EarlyToolExecutionBatch> = None;
let mut response_completed = false;
let mut event_count: usize = 0;
loop {
@@ -3229,6 +3247,14 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}));
}
}
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 {
crates/agen/tests/parallel_execution_test.rs
+103 -2
View File
@@ -260,6 +260,7 @@ impl Tool for ContextRecordingTool {
struct ControlledStreamClient {
receiver: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<Result<Event, ClientError>>>>>,
completion_support: ToolCallCompletionSupport,
stream_calls: Arc<AtomicUsize>,
}
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<ResponseStream, ClientError> {
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);
crates/worker/src/feature/builtin/flow_transition.rs
+7 -1
View File
@@ -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,
})),
])))
}
}