mirror of
https://github.com/codewhale-hq/Codewhale.git
synced 2026-11-08 08:04:25 -10:00
fix: preserve canceled tool completion receipts through one event primitive
Move the available-capacity receipt policy into the existing reserve_event_capacity authority, deleting the separate Engine caller fast path. Migrate all 12 reservation consumers: stream/start/admission/terminal handoff remain strict; ordinary usage, activity, parallel tool completion and owned MCP notice receipts retain available capacity after cancellation and release Full waits. Keep terminal leases and existing typed outcomes unchanged. Correct the owned collapsible_if without lint suppression. Add actual parallel run_turn tool completion witnesses for available and full queues plus exact activity outcome/span witnesses; formatting and whitespace passed. Prior frozen 3ae checkpoint combined 33 passed / 0 failed and scoped Clippy failed on that one nested conditional. New coherent source batch qualification pending. Signed-off-by: CodeWhale Bot <[email protected]>
This commit is contained in:
1 parent
3c49ab172c
commit
f8774da858
5 files changed
+359
-47
No files matched your search
@@ -2162,13 +2162,21 @@ impl Engine {
|
||||
approval_mode: ApprovalMode,
|
||||
) {
|
||||
let turn_control = self.begin_turn_control();
|
||||
let Ok(terminal_permit) =
|
||||
streaming::reserve_event_capacity(&self.tx_event, Some(&self.cancel_token)).await
|
||||
let Ok(terminal_permit) = streaming::reserve_event_capacity(
|
||||
&self.tx_event,
|
||||
Some(&self.cancel_token),
|
||||
streaming::EventReservationPolicy::Strict,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Ok(start_permit) =
|
||||
streaming::reserve_event_capacity(&self.tx_event, Some(&self.cancel_token)).await
|
||||
let Ok(start_permit) = streaming::reserve_event_capacity(
|
||||
&self.tx_event,
|
||||
Some(&self.cancel_token),
|
||||
streaming::EventReservationPolicy::Strict,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
@@ -5470,10 +5478,18 @@ impl Engine {
|
||||
// create a completion with no start. The production queue has 256
|
||||
// slots; these local permits do not create a second event authority.
|
||||
let admission = async {
|
||||
let terminal =
|
||||
streaming::reserve_event_capacity(&self.tx_event, Some(&self.cancel_token)).await?;
|
||||
let started =
|
||||
streaming::reserve_event_capacity(&self.tx_event, Some(&self.cancel_token)).await?;
|
||||
let terminal = streaming::reserve_event_capacity(
|
||||
&self.tx_event,
|
||||
Some(&self.cancel_token),
|
||||
streaming::EventReservationPolicy::Strict,
|
||||
)
|
||||
.await?;
|
||||
let started = streaming::reserve_event_capacity(
|
||||
&self.tx_event,
|
||||
Some(&self.cancel_token),
|
||||
streaming::EventReservationPolicy::Strict,
|
||||
)
|
||||
.await?;
|
||||
Ok::<_, streaming::EventSendError>((terminal, started))
|
||||
};
|
||||
let (terminal_permit, start_permit) = match admission.await {
|
||||
|
||||
@@ -15,14 +15,32 @@ pub(super) enum EventSendError {
|
||||
Closed,
|
||||
}
|
||||
|
||||
/// Reserve from the one event queue. A terminal reservation is held by its
|
||||
/// turn; ordinary sends consume theirs immediately. A live turn retains
|
||||
/// lossless backpressure, cancellation releases its wait, and idle sends have
|
||||
/// no turn token to cancel them.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum EventReservationPolicy {
|
||||
/// Cancellation wins before admitting work or a new stream observation.
|
||||
Strict,
|
||||
/// Preserve a completed observation when capacity is already available.
|
||||
Receipt,
|
||||
}
|
||||
|
||||
/// Reserve from the one event queue. Lifecycle and terminal handoff permits
|
||||
/// are held locally; ordinary receipts consume theirs immediately. Receipts
|
||||
/// may enter available capacity after cancellation, but cancellation always
|
||||
/// releases a wait on a full queue. Idle sends have no turn token to cancel.
|
||||
pub(super) async fn reserve_event_capacity(
|
||||
tx: &tokio::sync::mpsc::Sender<super::Event>,
|
||||
cancel: Option<&tokio_util::sync::CancellationToken>,
|
||||
policy: EventReservationPolicy,
|
||||
) -> Result<tokio::sync::mpsc::OwnedPermit<super::Event>, EventSendError> {
|
||||
if policy == EventReservationPolicy::Receipt {
|
||||
match tx.clone().try_reserve_owned() {
|
||||
Ok(permit) => return Ok(permit),
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
|
||||
return Err(EventSendError::Closed);
|
||||
}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {}
|
||||
}
|
||||
}
|
||||
let reserve = tx.clone().reserve_owned();
|
||||
match cancel {
|
||||
Some(cancel) => tokio::select! {
|
||||
@@ -38,7 +56,13 @@ impl super::Engine {
|
||||
/// Stream observations always belong to the decoder's current turn,
|
||||
/// including direct test/embedding calls that do not enqueue an Op.
|
||||
pub(super) async fn send_stream_event(&self, event: super::Event) -> bool {
|
||||
match reserve_event_capacity(&self.tx_event, Some(&self.cancel_token)).await {
|
||||
match reserve_event_capacity(
|
||||
&self.tx_event,
|
||||
Some(&self.cancel_token),
|
||||
EventReservationPolicy::Strict,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => {
|
||||
permit.send(event);
|
||||
true
|
||||
@@ -60,15 +84,12 @@ impl super::Engine {
|
||||
.active
|
||||
.as_ref()
|
||||
.map(|control| control.cancel.clone());
|
||||
let permit = match self.tx_event.clone().try_reserve_owned() {
|
||||
Ok(permit) => permit,
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
|
||||
return Err(EventSendError::Closed);
|
||||
}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
|
||||
reserve_event_capacity(&self.tx_event, cancel.as_ref()).await?
|
||||
}
|
||||
};
|
||||
let permit = reserve_event_capacity(
|
||||
&self.tx_event,
|
||||
cancel.as_ref(),
|
||||
EventReservationPolicy::Receipt,
|
||||
)
|
||||
.await?;
|
||||
permit.send(event);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -611,6 +611,181 @@ async fn event_capacity_cancelled_repl_child_keeps_unknown_cost_and_discards_ker
|
||||
);
|
||||
drop(turn_guard);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn event_capacity_cancelled_parallel_tool_keeps_completed_span_and_call_when_available() {
|
||||
use crate::llm_client::mock::{MockLlmClient, canned};
|
||||
use crate::tools::spec::{ToolCapability, ToolSpec};
|
||||
use codewhale_protocol::engine_owner::OwnerOperationOutcome;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
struct CompleteThenCancel {
|
||||
cancel: tokio_util::sync::CancellationToken,
|
||||
tx: mpsc::Sender<Event>,
|
||||
fill_queue: bool,
|
||||
executed: Arc<AtomicUsize>,
|
||||
}
|
||||
#[async_trait::async_trait]
|
||||
impl ToolSpec for CompleteThenCancel {
|
||||
fn name(&self) -> &str {
|
||||
"fixture_complete_then_cancel"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"Finish an observed operation before firing its turn cancellation token."
|
||||
}
|
||||
fn input_schema(&self) -> Value {
|
||||
json!({"type": "object"})
|
||||
}
|
||||
fn capabilities(&self) -> Vec<ToolCapability> {
|
||||
vec![ToolCapability::ReadOnly]
|
||||
}
|
||||
fn supports_parallel(&self) -> bool {
|
||||
true
|
||||
}
|
||||
async fn execute(&self, _: Value, _: &ToolContext) -> Result<ToolResult, ToolError> {
|
||||
self.executed.fetch_add(1, Ordering::SeqCst);
|
||||
if self.fill_queue {
|
||||
while self
|
||||
.tx
|
||||
.try_send(Event::status("full at tool completion"))
|
||||
.is_ok()
|
||||
{}
|
||||
}
|
||||
self.cancel.cancel();
|
||||
Ok(ToolResult::success("completed before cancellation"))
|
||||
}
|
||||
}
|
||||
|
||||
for fill_queue in [false, true] {
|
||||
let workspace = tempdir().unwrap();
|
||||
let mock = Arc::new(MockLlmClient::new(vec![
|
||||
tool_batch_turn(&[
|
||||
("one", "fixture_complete_then_cancel", "{}"),
|
||||
("two", "fixture_complete_then_cancel", "{}"),
|
||||
]),
|
||||
canned::simple_text_turn("must not run after cancellation"),
|
||||
]));
|
||||
let (mut engine, handle) = Engine::new_with_model_client(
|
||||
deterministic_engine_config(workspace.path()),
|
||||
&Config::default(),
|
||||
mock.clone(),
|
||||
);
|
||||
let turn_guard = engine.begin_turn_control();
|
||||
let executed = Arc::new(AtomicUsize::new(0));
|
||||
let mut registry = crate::tools::ToolRegistry::new(ToolContext::new(workspace.path()));
|
||||
registry.register(Arc::new(CompleteThenCancel {
|
||||
cancel: engine.cancel_token.clone(),
|
||||
tx: engine.tx_event.clone(),
|
||||
fill_queue,
|
||||
executed: executed.clone(),
|
||||
}));
|
||||
let tools = Some(registry.to_api_tools_with_cache(true));
|
||||
let policy = test_tool_surface(&engine, registry, tools, AppMode::Agent);
|
||||
let mut turn = TurnContext::new(4);
|
||||
let (status, error) = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
engine.run_turn(&mut turn, policy, None, None),
|
||||
)
|
||||
.await
|
||||
.expect("actual parallel tool completion cannot park on a full cancelled queue");
|
||||
assert_eq!(status, TurnOutcomeStatus::Interrupted, "{error:?}");
|
||||
assert!(error.is_none());
|
||||
assert_eq!(
|
||||
executed.load(Ordering::SeqCst),
|
||||
1,
|
||||
"the peer never executes"
|
||||
);
|
||||
assert_eq!(
|
||||
mock.call_count(),
|
||||
1,
|
||||
"no provider continuation after cancellation"
|
||||
);
|
||||
let mut rx = handle.rx_event.write().await;
|
||||
let events: Vec<_> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
|
||||
let starts: Vec<_> = events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
Event::OperationActivityStarted {
|
||||
span_id,
|
||||
activity_kind,
|
||||
} => Some((span_id, activity_kind)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(starts.len(), 1, "one actual tool activity was admitted");
|
||||
let completed: Vec<_> = events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
Event::OperationActivityCompleted {
|
||||
span_id,
|
||||
activity_kind,
|
||||
outcome,
|
||||
} => Some((span_id, activity_kind, outcome)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
let calls: Vec<_> = events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
Event::ToolCallComplete {
|
||||
id,
|
||||
model_call,
|
||||
result,
|
||||
..
|
||||
} => Some((id, model_call, result)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
if fill_queue {
|
||||
assert!(
|
||||
completed.is_empty() && calls.is_empty(),
|
||||
"cancelled full waits release without inventing delivery"
|
||||
);
|
||||
} else {
|
||||
assert_eq!(
|
||||
completed.len(),
|
||||
1,
|
||||
"preserve the completed activity after cancellation"
|
||||
);
|
||||
assert_eq!(completed[0].0, starts[0].0, "retain the span relationship");
|
||||
assert_eq!(completed[0].1, starts[0].1);
|
||||
assert_eq!(*completed[0].2, OwnerOperationOutcome::Succeeded);
|
||||
assert_eq!(
|
||||
calls.len(),
|
||||
2,
|
||||
"completed call and cancelled peer both settle exactly once"
|
||||
);
|
||||
let succeeded: Vec<_> = calls
|
||||
.iter()
|
||||
.filter(|(_, _, result)| result.as_ref().is_ok_and(|result| result.success))
|
||||
.collect();
|
||||
assert_eq!(succeeded.len(), 1);
|
||||
assert_eq!(
|
||||
succeeded[0].2.as_ref().unwrap().content,
|
||||
"completed before cancellation"
|
||||
);
|
||||
assert!(
|
||||
starts[0].0.starts_with(&format!("{}#", succeeded[0].0)),
|
||||
"span retains its completed execution identity"
|
||||
);
|
||||
let provider_ids: HashSet<_> = calls
|
||||
.iter()
|
||||
.map(|(_, model_call, _)| model_call.as_ref().unwrap().provider_id.as_str())
|
||||
.collect();
|
||||
assert_eq!(provider_ids, HashSet::from(["one", "two"]));
|
||||
let call_starts: HashSet<_> = events
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
Event::ToolCallStarted { id, .. } => Some(id),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert!(calls.iter().all(|(id, _, _)| call_starts.contains(id)));
|
||||
}
|
||||
drop(rx);
|
||||
drop(turn_guard);
|
||||
}
|
||||
}
|
||||
|
||||
const REPRESENTATIVE_FIXTURE_ID: &str = "representative-v1";
|
||||
const REPRESENTATIVE_PROJECT_AUTHORITY: &str = "REPRESENTATIVE_PROJECT_AUTHORITY";
|
||||
const REPRESENTATIVE_PROJECT_AUTHORITY_BODY: &str = concat!(
|
||||
|
||||
@@ -24,15 +24,15 @@ fn inherited_interactive_shell_refusal(tool_name: &str, interactive: bool) -> Op
|
||||
.map(|message| ToolError::execution_failed(message.to_string()))
|
||||
}
|
||||
|
||||
/// Pairs one `OperationActivityStarted` with exactly one
|
||||
/// `OperationActivityCompleted`.
|
||||
/// Pairs an observed `OperationActivityStarted` with at most one completion.
|
||||
///
|
||||
/// The turn loop drops an in-flight tool future when the user cancels
|
||||
/// (`tokio::select!` on the cancel token, or `drop(tool_tasks)` for a parallel
|
||||
/// batch), so a Completed sent inline after the await would never be sent.
|
||||
/// Dropping an unfinished span sends `Completed { Cancelled }` with
|
||||
/// `try_send`: best effort, like the other guards here, because `Drop` cannot
|
||||
/// await a full channel.
|
||||
/// await a full channel. Cancellation may release a completion parked on a
|
||||
/// full queue; the turn's reserved terminal observation settles its lifecycle.
|
||||
pub(super) struct OperationSpanGuard {
|
||||
tx: mpsc::Sender<Event>,
|
||||
span: Option<(String, codewhale_protocol::engine_owner::OwnerActivityKind)>,
|
||||
@@ -56,9 +56,15 @@ impl OperationSpanGuard {
|
||||
cancel: Option<CancellationToken>,
|
||||
) -> Self {
|
||||
let span_id = Self::span_id(call_id);
|
||||
// `Sender::send` is cancel safe: if this future is dropped the event
|
||||
// Reservation is cancel safe: if this future is dropped the event
|
||||
// was either sent or not, and the guard is armed only once it was.
|
||||
let sent = match super::streaming::reserve_event_capacity(&tx, cancel.as_ref()).await {
|
||||
let sent = match super::streaming::reserve_event_capacity(
|
||||
&tx,
|
||||
cancel.as_ref(),
|
||||
super::streaming::EventReservationPolicy::Strict,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => {
|
||||
permit.send(Event::OperationActivityStarted {
|
||||
span_id: span_id.clone(),
|
||||
@@ -76,16 +82,19 @@ impl OperationSpanGuard {
|
||||
}
|
||||
|
||||
async fn complete(mut self, outcome: codewhale_protocol::engine_owner::OwnerOperationOutcome) {
|
||||
if let Some((span_id, activity_kind)) = self.span.take() {
|
||||
if let Ok(permit) =
|
||||
super::streaming::reserve_event_capacity(&self.tx, self.cancel.as_ref()).await
|
||||
{
|
||||
permit.send(Event::OperationActivityCompleted {
|
||||
span_id,
|
||||
activity_kind,
|
||||
outcome,
|
||||
});
|
||||
}
|
||||
if let Some((span_id, activity_kind)) = self.span.take()
|
||||
&& let Ok(permit) = super::streaming::reserve_event_capacity(
|
||||
&self.tx,
|
||||
self.cancel.as_ref(),
|
||||
super::streaming::EventReservationPolicy::Receipt,
|
||||
)
|
||||
.await
|
||||
{
|
||||
permit.send(Event::OperationActivityCompleted {
|
||||
span_id,
|
||||
activity_kind,
|
||||
outcome,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -186,13 +195,21 @@ impl InteractiveTerminalGuard {
|
||||
if !interactive {
|
||||
return Ok(Self { resume: None });
|
||||
}
|
||||
let resume = super::streaming::reserve_event_capacity(&tx, cancel.as_ref())
|
||||
.await
|
||||
.map_err(terminal_handoff_send_error)?;
|
||||
let resume = super::streaming::reserve_event_capacity(
|
||||
&tx,
|
||||
cancel.as_ref(),
|
||||
super::streaming::EventReservationPolicy::Strict,
|
||||
)
|
||||
.await
|
||||
.map_err(terminal_handoff_send_error)?;
|
||||
let ack = Arc::new(tokio::sync::Notify::new());
|
||||
let pause = super::streaming::reserve_event_capacity(&tx, cancel.as_ref())
|
||||
.await
|
||||
.map_err(terminal_handoff_send_error)?;
|
||||
let pause = super::streaming::reserve_event_capacity(
|
||||
&tx,
|
||||
cancel.as_ref(),
|
||||
super::streaming::EventReservationPolicy::Strict,
|
||||
)
|
||||
.await
|
||||
.map_err(terminal_handoff_send_error)?;
|
||||
// No await separates the pause send and guard installation. A
|
||||
// cancelled reservation above never paused and needs no resume.
|
||||
pause.send(Event::PauseEvents {
|
||||
@@ -311,7 +328,13 @@ impl Engine {
|
||||
let url = url.to_string();
|
||||
let tx = tx_event.clone();
|
||||
notice_task = Some(super::turn_heartbeat::AbortOnDrop(tokio::spawn(async move {
|
||||
if let Ok(permit) = super::streaming::reserve_event_capacity(&tx, None).await {
|
||||
if let Ok(permit) = super::streaming::reserve_event_capacity(
|
||||
&tx,
|
||||
None,
|
||||
super::streaming::EventReservationPolicy::Receipt,
|
||||
)
|
||||
.await
|
||||
{
|
||||
permit.send(Event::status(format!(
|
||||
"◆ auth required: sign in to MCP server '{server}' in your browser — {url}"
|
||||
)));
|
||||
@@ -684,6 +707,80 @@ mod tests {
|
||||
drop(guard);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn event_capacity_cancelled_activity_preserves_exact_outcomes_and_releases_full_queue() {
|
||||
use codewhale_protocol::engine_owner::{OwnerActivityKind, OwnerOperationOutcome};
|
||||
for outcome in [
|
||||
OwnerOperationOutcome::Succeeded,
|
||||
OwnerOperationOutcome::Failed,
|
||||
OwnerOperationOutcome::Denied,
|
||||
OwnerOperationOutcome::Cancelled,
|
||||
] {
|
||||
let (tx, mut rx) = mpsc::channel(2);
|
||||
let cancel = CancellationToken::new();
|
||||
let span = OperationSpanGuard::start(
|
||||
tx,
|
||||
"completed-call",
|
||||
OwnerActivityKind::Tool,
|
||||
Some(cancel.clone()),
|
||||
)
|
||||
.await;
|
||||
cancel.cancel();
|
||||
tokio::time::timeout(Duration::from_millis(100), span.complete(outcome))
|
||||
.await
|
||||
.expect("completed activity preserves available capacity after cancellation");
|
||||
let Event::OperationActivityStarted {
|
||||
span_id,
|
||||
activity_kind,
|
||||
} = rx.try_recv().unwrap()
|
||||
else {
|
||||
panic!("activity start");
|
||||
};
|
||||
let Event::OperationActivityCompleted {
|
||||
span_id: completed_span,
|
||||
activity_kind: completed_kind,
|
||||
outcome: completed_outcome,
|
||||
} = rx
|
||||
.try_recv()
|
||||
.expect("cancellation must not discard the completion")
|
||||
else {
|
||||
panic!("activity completion");
|
||||
};
|
||||
assert_eq!(completed_span, span_id);
|
||||
assert_eq!(completed_kind, activity_kind);
|
||||
assert_eq!(
|
||||
completed_outcome, outcome,
|
||||
"retain the observed outcome exactly"
|
||||
);
|
||||
assert!(rx.try_recv().is_err(), "exactly one completion");
|
||||
}
|
||||
|
||||
let (tx, mut rx) = mpsc::channel(1);
|
||||
let cancel = CancellationToken::new();
|
||||
let span = OperationSpanGuard::start(
|
||||
tx,
|
||||
"full-call",
|
||||
OwnerActivityKind::Tool,
|
||||
Some(cancel.clone()),
|
||||
)
|
||||
.await;
|
||||
cancel.cancel();
|
||||
tokio::time::timeout(
|
||||
Duration::from_millis(100),
|
||||
span.complete(OwnerOperationOutcome::Succeeded),
|
||||
)
|
||||
.await
|
||||
.expect("a full queue cannot retain a cancelled activity sender");
|
||||
assert!(matches!(
|
||||
rx.try_recv().unwrap(),
|
||||
Event::OperationActivityStarted { .. }
|
||||
));
|
||||
assert!(
|
||||
rx.try_recv().is_err(),
|
||||
"no sender survives to publish after draining"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_heartbeat_is_delayed_for_fast_tool() {
|
||||
let (tx, mut rx) = mpsc::channel(4);
|
||||
|
||||
@@ -4521,9 +4521,12 @@ impl Engine {
|
||||
.map(|result| result.content_blocks.clone())
|
||||
.unwrap_or_default();
|
||||
let legacy_result = result.map(RichToolResult::into_result);
|
||||
if let Ok(permit) =
|
||||
super::streaming::reserve_event_capacity(&tx_event, Some(&cancel_token))
|
||||
.await
|
||||
if let Ok(permit) = super::streaming::reserve_event_capacity(
|
||||
&tx_event,
|
||||
Some(&cancel_token),
|
||||
super::streaming::EventReservationPolicy::Receipt,
|
||||
)
|
||||
.await
|
||||
{
|
||||
permit.send(Event::ToolCallComplete {
|
||||
model_call: plan.model_call.clone(),
|
||||
|
||||
Reference in new issue
Block a user