code-mode: tighten lifecycle tests and fix CI

This commit is contained in:
Channing Conger
2026-06-19 18:51:09 +00:00
parent dfc1727e0d
commit f0b82c0ea1
8 changed files with 96 additions and 190 deletions
+8 -2
View File
@@ -321,7 +321,7 @@ async fn run_cell<H: CellHost>(
send_termination_events(
&accepting_requests,
observer.take(),
None,
/*termination*/ None,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -455,7 +455,7 @@ async fn run_cell<H: CellHost>(
send_termination_events(
&accepting_requests,
observer.take(),
None,
/*termination*/ None,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -548,6 +548,12 @@ fn send_or_buffer_completion(
completed_event: &mut Option<CellEvent>,
accepting_requests: &AtomicBool,
) -> bool {
if observer
.as_ref()
.is_some_and(|observer| observer.response_tx.is_closed())
{
*observer = None;
}
if observer.is_some() {
accepting_requests.store(false, Ordering::Release);
send_observer_event(observer.take(), event);
+27 -83
View File
@@ -267,26 +267,6 @@ async fn yield_timer_preempts_buffered_runtime_output() {
harness.task.await.unwrap();
}
#[tokio::test]
async fn queued_termination_preempts_unobserved_runtime_completion() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(60)));
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::new(),
error_text: None,
})
.unwrap();
let termination = harness.handle.terminate();
let terminated = Ok(CellEvent::Terminated {
content_items: Vec::new(),
});
assert_eq!(termination.await, terminated.clone());
assert_eq!(harness.initial_event_rx.await.unwrap(), terminated);
harness.task.await.unwrap();
}
#[tokio::test]
async fn termination_preempts_result_without_committing_stored_values() {
let host = Arc::new(RecordingHost::default());
@@ -302,6 +282,10 @@ async fn termination_preempts_result_without_committing_stored_values() {
})
.unwrap();
let termination = harness.handle.terminate();
assert_eq!(
harness.handle.terminate().await,
Err(CellError::AlreadyTerminating)
);
let terminated = Ok(CellEvent::Terminated {
content_items: Vec::new(),
@@ -394,53 +378,6 @@ async fn session_shutdown_during_result_callback_drain_prevents_stored_value_com
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test]
async fn termination_during_result_callback_drain_does_not_wait_for_the_notification() {
let host = Arc::new(PendingNotificationHost::default());
let harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
Arc::clone(&host),
);
harness
.event_tx
.send(RuntimeEvent::Notify {
call_id: "call-1".to_string(),
text: "pending".to_string(),
})
.unwrap();
tokio::task::yield_now().await;
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::from([("key".to_string(), JsonValue::Bool(true))]),
error_text: None,
})
.unwrap();
tokio::task::yield_now().await;
let termination = harness.handle.terminate();
let terminated = tokio::time::timeout(Duration::from_millis(/*millis*/ 100), termination)
.await
.expect("termination should not await a non-cooperative notification");
assert_eq!(
terminated,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
assert_eq!(
harness.initial_event_rx.await.unwrap(),
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
drop(harness.event_tx);
harness.task.await.unwrap();
assert!(!host.committed.load(Ordering::Acquire));
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test]
async fn session_shutdown_during_stored_value_commit_does_not_wait_for_the_host() {
let (host, mut commit_started_rx) = BlockingCommitHost::new();
@@ -559,9 +496,11 @@ async fn observer_receives_a_buffered_completion_after_commit_finishes() {
#[tokio::test]
async fn dropped_observer_does_not_block_the_next_observer() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(
/*secs*/ 60,
)));
let (host, mut commit_started_rx) = BlockingCommitHost::new();
let harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
Arc::clone(&host),
);
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
assert_eq!(
harness.initial_event_rx.await.unwrap(),
@@ -572,26 +511,31 @@ async fn dropped_observer_does_not_block_the_next_observer() {
let abandoned_observer = harness
.handle
.observe(ObserveMode::YieldAfter(Duration::ZERO));
.observe(ObserveMode::YieldAfter(Duration::from_secs(
/*secs*/ 60,
)));
drop(abandoned_observer);
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::new(),
error_text: None,
})
.unwrap();
assert_eq!(commit_started_rx.recv().await, Some(()));
let next_observer = harness
.handle
.observe(ObserveMode::YieldAfter(Duration::ZERO));
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
.observe(ObserveMode::YieldAfter(Duration::from_secs(
/*secs*/ 60,
)));
host.commit_release.add_permits(/*n*/ 1);
assert_eq!(
next_observer.await,
Ok(CellEvent::Yielded {
content_items: Vec::new(),
})
);
let termination = harness.handle.terminate();
drop(harness.event_tx);
assert_eq!(
termination.await,
Ok(CellEvent::Terminated {
Ok(CellEvent::Completed {
content_items: Vec::new(),
error_text: None,
})
);
harness.task.await.unwrap();
@@ -234,14 +234,15 @@ async fn next_event(events_rx: &mut mpsc::UnboundedReceiver<DelegateEvent>) -> D
#[tokio::test]
async fn yields_and_resumes() {
let service = CodeModeService::new();
let (delegate, mut events_rx) = ReleasableNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = service
.execute(ExecuteRequest {
source: r#"
text("before");
yield_control();
store("finished", true);
text("after");
notify("resumed");
"#
.to_string(),
yield_time_ms: Some(60_000),
@@ -259,30 +260,11 @@ text("after");
}],
}
);
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), async {
for cell_number in 2.. {
let response = service
.execute(ExecuteRequest {
cell_id: cell_id(&cell_number.to_string()),
..execute_request(r#"text(String(load("finished")));"#)
})
.await
.unwrap();
let RuntimeResponse::Result { content_items, .. } = response else {
panic!("expected stored-value probe to complete");
};
if content_items
== vec![FunctionCallOutputContentItem::InputText {
text: "true".to_string(),
}]
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
delegate.release_notification();
assert_eq!(
service
.wait(WaitRequest {
@@ -299,20 +281,6 @@ text("after");
error_text: None,
})
);
assert_eq!(
service
.wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
})
.await
.unwrap(),
WaitOutcome::MissingCell(RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: Vec::new(),
error_text: Some("exec cell 1 not found".to_string()),
})
);
}
#[tokio::test]
@@ -364,7 +332,7 @@ text("after");
}
#[tokio::test]
async fn observed_natural_completion_wins_over_termination() {
async fn buffered_natural_completion_wins_over_termination() {
let service = CodeModeService::new();
let cell = service
.execute(execute_request(
@@ -448,7 +416,7 @@ async fn termination_discards_pending_callbacks_before_responding() {
}
#[tokio::test]
async fn shutdown_does_not_await_notifications_during_natural_completion() {
async fn shutdown_cancels_a_non_cooperative_notification() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate));
service
@@ -471,7 +439,7 @@ async fn shutdown_does_not_await_notifications_during_natural_completion() {
}
#[tokio::test]
async fn termination_does_not_await_notifications_during_natural_completion() {
async fn termination_cancels_a_non_cooperative_notification() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate);
service
@@ -568,25 +536,6 @@ async fn termination_discards_pending_tools_before_responding() {
);
}
#[tokio::test]
async fn shutdown_discards_pending_tools_before_returning() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
service
.execute(ExecuteRequest {
enabled_tools: vec![blocking_tool()],
source: r#"await tools.block({});"#.to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
let shutdown_service = Arc::clone(&service);
assert_eq!(shutdown_service.shutdown().await, Ok(()));
}
#[tokio::test]
async fn shutdown_completes_when_a_nested_tool_ignores_cancellation() {
let (delegate, mut events_rx) = NeverResolvingToolDelegate::new();
+25 -21
View File
@@ -88,33 +88,37 @@ fn execute_request_with_id(cell_id: &str, source: &str) -> ExecuteRequest {
async fn terminal_cells_are_unrouted_before_they_are_unregistered() {
let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate)));
let cell_id = CellId::new("1");
let event = runtime
.execute(
execute_request(r#"text("done");"#),
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
)
.await
.unwrap();
assert_eq!(
event,
CellEvent::Completed {
content_items: vec![OutputItem::Text {
text: "done".to_string(),
}],
error_text: None,
}
);
assert_eq!(
runtime
.observe(&cell_id, ObserveMode::PendingFrontier)
.execute(
execute_request("while (true) {}"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await,
Err(Error::MissingCell(cell_id.clone()))
Ok(CellEvent::Yielded {
content_items: Vec::new(),
})
);
let mut cell_count = runtime.inner.cell_count_tx.subscribe();
let cells = runtime.inner.cells.lock().await;
let handle = cells.get(&cell_id).unwrap().clone();
assert_eq!(
runtime.terminate(&cell_id).await,
Err(Error::MissingCell(cell_id.clone()))
handle.terminate().await,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
assert!(cells.contains_key(&cell_id));
assert_eq!(
handle.observe(ObserveMode::PendingFrontier).await,
Err(CellError::Closed)
);
assert_eq!(handle.terminate().await, Err(CellError::Closed));
drop(cells);
cell_count.changed().await.unwrap();
assert_eq!(*cell_count.borrow_and_update(), 0);
assert_eq!(runtime.shutdown().await, Ok(()));
}