code-mode: clean up terminal cell dispatch gates

This commit is contained in:
Channing Conger
2026-06-21 07:22:16 +00:00
parent 18fb5edda3
commit 54d30e5fa1
12 changed files with 204 additions and 11 deletions
+18 -10
View File
@@ -291,7 +291,7 @@ async fn run_cell<H: CellHost>(
content_items: std::mem::take(&mut content_items),
error_text: Some("exec runtime ended unexpectedly".to_string()),
};
let rejected_event = match host
let (completion_committed, rejected_event) = match host
.commit_completion(
HashMap::new(),
event,
@@ -300,14 +300,18 @@ async fn run_cell<H: CellHost>(
)
.await
{
CompletionCommit::Committed => None,
CompletionCommit::Rejected(event) => Some(event),
CompletionCommit::Committed => (true, None),
CompletionCommit::Rejected(event) => (false, Some(event)),
};
match cell_state.deliver_completion(
let completion_delivery = cell_state.deliver_completion(
observer
.take()
.map(|observer| (observer.mode, observer.response_tx)),
) {
);
if completion_committed {
host.terminal();
}
match completion_delivery {
CompletionDelivery::Delivered => break,
CompletionDelivery::Buffered => {}
CompletionDelivery::Rejected(response_tx) => {
@@ -438,7 +442,7 @@ async fn run_cell<H: CellHost>(
content_items: std::mem::take(&mut content_items),
error_text,
};
let rejected_event = match host
let (completion_committed, rejected_event) = match host
.commit_completion(
stored_value_writes,
event,
@@ -447,14 +451,18 @@ async fn run_cell<H: CellHost>(
)
.await
{
CompletionCommit::Committed => None,
CompletionCommit::Rejected(event) => Some(event),
CompletionCommit::Committed => (true, None),
CompletionCommit::Rejected(event) => (false, Some(event)),
};
match cell_state.deliver_completion(
let completion_delivery = cell_state.deliver_completion(
observer
.take()
.map(|observer| (observer.mode, observer.response_tx)),
) {
);
if completion_committed {
host.terminal();
}
match completion_delivery {
CompletionDelivery::Delivered => break,
CompletionDelivery::Buffered => {}
CompletionDelivery::Rejected(response_tx) => {
@@ -52,6 +52,8 @@ impl CellHost for TestHost {
cell_state.commit_completion(event, pending_initial_yield_items, || {})
}
fn terminal(&self) {}
async fn closed(&self) {}
}
@@ -86,6 +88,8 @@ impl CellHost for RecordingHost {
})
}
fn terminal(&self) {}
async fn closed(&self) {}
}
@@ -60,6 +60,9 @@ pub(crate) trait CellHost: Send + Sync + 'static {
cell_state: Arc<CellState>,
) -> impl Future<Output = CompletionCommit> + Send;
/// Reports that execution reached a terminal state, even when its result remains buffered.
fn terminal(&self);
fn closed(&self) -> impl Future<Output = ()> + Send;
}
+6
View File
@@ -49,6 +49,8 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
#[derive(Default)]
@@ -251,6 +253,10 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
)
.await
}
fn cell_closed(&self, cell_id: &runtime::CellId) {
self.delegate.cell_closed(&protocol_cell_id(cell_id));
}
}
fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest {
@@ -19,6 +19,11 @@ enum DelegateEvent {
NotificationStarted,
NotificationFinished,
ToolStarted,
CellClosed(CellId),
}
fn record_cell_closed(events_tx: &mpsc::UnboundedSender<DelegateEvent>, cell_id: &CellId) {
let _ = events_tx.send(DelegateEvent::CellClosed(cell_id.clone()));
}
struct BlockingDelegate {
@@ -100,6 +105,10 @@ impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate {
std::future::pending().await
})
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
@@ -134,6 +143,10 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
}
})
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
@@ -157,6 +170,10 @@ impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
impl BlockingDelegate {
@@ -211,6 +228,10 @@ impl CodeModeSessionDelegate for BlockingDelegate {
Err("cancelled".to_string())
})
}
fn cell_closed(&self, cell_id: &CellId) {
record_cell_closed(&self.events_tx, cell_id);
}
}
fn cell_id(value: &str) -> CellId {
@@ -372,6 +393,39 @@ await tools.block({});
);
}
#[tokio::test]
async fn background_completion_notifies_the_delegate_without_another_observation() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate);
let created_cell_id = service
.create_cell(execute_request(
r#"await new Promise(resolve => setTimeout(resolve, 100)); text("done");"#,
))
.await
.unwrap();
assert_eq!(
tokio::time::timeout(
Duration::from_secs(2),
service.observe(ObserveRequest {
cell_id: created_cell_id.clone(),
yield_time_ms: 1,
}),
)
.await
.expect("initial observation should yield while the cell is still running")
.unwrap(),
CellOutcome::LiveCell(RuntimeResponse::Yielded {
cell_id: created_cell_id.clone(),
content_items: Vec::new(),
})
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(created_cell_id)
);
}
#[tokio::test]
async fn returns_and_resumes_from_the_pending_frontier() {
let (delegate, mut events_rx) = BlockingDelegate::new();
+18
View File
@@ -68,6 +68,8 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate {
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
#[derive(Debug, PartialEq)]
@@ -82,6 +84,9 @@ enum RecordedDelegateCall {
text: String,
cancellation_requested: bool,
},
CellClosed {
cell_id: CellId,
},
}
struct RecordingDelegate {
@@ -150,6 +155,15 @@ impl CodeModeSessionDelegate for RecordingDelegate {
.expect("test must provide one result per notification")
})
}
fn cell_closed(&self, cell_id: &CellId) {
self.calls
.lock()
.unwrap()
.push(RecordedDelegateCall::CellClosed {
cell_id: cell_id.clone(),
});
}
}
fn execute_request(source: &str) -> CreateCellRequest {
@@ -1204,6 +1218,7 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel
.await,
Err("notification failed".to_string())
);
runtime::SessionRuntimeDelegate::cell_closed(&adapter, &runtime::CellId::new("cell-d4"));
assert_eq!(
delegate.take_calls(),
@@ -1234,6 +1249,9 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel
text: "progress".to_string(),
cancellation_requested: true,
},
RecordedDelegateCall::CellClosed {
cell_id: cell_id("cell-d4"),
},
]
);
}
@@ -4,6 +4,7 @@ use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
@@ -177,6 +178,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
let host = Arc::new(RuntimeCellHost {
cell_id: cell_id.clone(),
inner: Arc::clone(&self.inner),
terminal_notified: AtomicBool::new(false),
});
let mut cells = self.inner.cells.lock().await;
if self.inner.shutdown_token.is_cancelled() {
@@ -220,6 +222,7 @@ impl PendingEvent {
struct RuntimeCellHost<D: SessionRuntimeDelegate> {
cell_id: CellId,
inner: Arc<Inner<D>>,
terminal_notified: AtomicBool,
}
impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
@@ -275,8 +278,15 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
})
}
fn terminal(&self) {
if !self.terminal_notified.swap(true, Ordering::AcqRel) {
self.inner.delegate.cell_closed(&self.cell_id);
}
}
async fn closed(&self) {
self.inner.cells.lock().await.remove(&self.cell_id);
self.terminal();
}
}
@@ -1,6 +1,7 @@
use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::task::Context;
use std::task::Poll;
use std::task::Waker;
@@ -33,6 +34,8 @@ impl SessionRuntimeDelegate for RecordingDelegate {
) -> Result<(), String> {
Ok(())
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
#[tokio::test]
@@ -42,6 +45,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
let host = RuntimeCellHost {
cell_id: CellId::new("terminating-writer"),
inner: Arc::clone(&runtime.inner),
terminal_notified: AtomicBool::new(false),
};
let completion = CellEvent::Completed {
content_items: vec![OutputItem::Text {
@@ -118,7 +118,8 @@ pub struct NestedToolCall {
/// Host callbacks used by cells owned by a [`super::SessionRuntime`].
///
/// Implementations should forward callback cancellation tokens to downstream
/// work. The runtime stops awaiting callbacks once cancellation begins.
/// work. The runtime stops awaiting callbacks once cancellation begins. Terminal
/// notifications must be non-blocking and are not acknowledged or retried.
pub trait SessionRuntimeDelegate: Send + Sync + 'static {
fn invoke_tool(
&self,
@@ -133,6 +134,11 @@ pub trait SessionRuntimeDelegate: Send + Sync + 'static {
text: String,
cancellation_token: CancellationToken,
) -> impl Future<Output = Result<(), String>> + Send;
/// Reports that execution is terminal and will issue no more callbacks.
///
/// The terminal result may remain buffered for a later observation.
fn cell_closed(&self, cell_id: &CellId);
}
/// A failure reported by a session runtime operation.