mirror of
https://github.com/openai/codex.git
synced 2026-11-14 21:40:39 +00:00
Add early yielding for code-mode observations (#48123)
## What changed Add an optional `preempt` signal to `CodeModeSession::execute` and `CodeModeSession::wait` so callers can end an observation early while the cell continues running and remains available for later waits. Support the signal in the in-process runtime and across both remote transports: `operation/yield` with the negotiated `yield-observation` capability for stdio, and `YieldObservation` for gRPC. Preserve signals received before gRPC observation registration, and retain normal timeout behavior for older hosts that lack support. ## Testing Add runtime and transport tests covering early execute and wait yields followed by successful cell completion, yields before observation registration, and suppression of yield frames for older stdio hosts. GitOrigin-RevId: 3f8049d8f0582be6b8f0832cc03f2d974a0b8f18
This commit is contained in:
@@ -197,7 +197,7 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
|
||||
..execute_request("await tools.block({}); await tools.block({});")
|
||||
};
|
||||
let a = service
|
||||
.execute(request.clone(), delegate_a.clone())
|
||||
.execute(request.clone(), delegate_a.clone(), /*preempt*/ None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(next_event(&mut events_a).await, DelegateEvent::ToolStarted);
|
||||
@@ -206,7 +206,10 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
|
||||
RuntimeResponse::Yielded { .. }
|
||||
));
|
||||
|
||||
let b = service.execute(request, delegate_b.clone()).await.unwrap();
|
||||
let b = service
|
||||
.execute(request, delegate_b.clone(), /*preempt*/ None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(next_event(&mut events_b).await, DelegateEvent::ToolStarted);
|
||||
assert!(matches!(
|
||||
b.initial_response().await.unwrap(),
|
||||
@@ -219,10 +222,13 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
|
||||
delegate_a.release_tool();
|
||||
assert_eq!(
|
||||
service
|
||||
.wait(WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000
|
||||
})
|
||||
.wait(
|
||||
WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
},
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
WaitOutcome::LiveCell(RuntimeResponse::Result {
|
||||
@@ -267,6 +273,7 @@ async fn yields_and_resumes() {
|
||||
..execute_request("")
|
||||
},
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -283,10 +290,13 @@ async fn yields_and_resumes() {
|
||||
);
|
||||
assert_eq!(
|
||||
service
|
||||
.wait(WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
})
|
||||
.wait(
|
||||
WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
},
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
WaitOutcome::LiveCell(RuntimeResponse::Result {
|
||||
@@ -359,6 +369,7 @@ async fn observed_natural_completion_wins_over_termination() {
|
||||
.execute(
|
||||
execute_request(r#"yield_control(); store("finished", true); text("done");"#),
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -380,6 +391,7 @@ async fn observed_natural_completion_wins_over_termination() {
|
||||
..execute_request(r#"text(String(load("finished")));"#)
|
||||
},
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
@@ -422,6 +434,7 @@ async fn termination_cancels_pending_callbacks_before_responding() {
|
||||
.execute(
|
||||
execute_request(r#"notify("pending"); await new Promise(() => {});"#),
|
||||
delegate.clone(),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -462,7 +475,11 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
|
||||
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
|
||||
let service = Arc::new(InProcessCodeModeSession::new());
|
||||
service
|
||||
.execute(execute_request(r#"notify("pending");"#), delegate.clone())
|
||||
.execute(
|
||||
execute_request(r#"notify("pending");"#),
|
||||
delegate.clone(),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -495,6 +512,7 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
|
||||
.execute(
|
||||
execute_request(r#"notify("pending"); await new Promise(() => {});"#),
|
||||
delegate.clone(),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -548,6 +566,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
|
||||
.execute(
|
||||
execute_request("await new Promise(() => {});"),
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -562,17 +581,23 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
|
||||
);
|
||||
|
||||
let first_observer = service
|
||||
.begin_wait(WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
})
|
||||
.begin_wait(
|
||||
WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
},
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
service
|
||||
.wait(WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
})
|
||||
.wait(
|
||||
WaitRequest {
|
||||
cell_id: cell_id("1"),
|
||||
yield_time_ms: 60_000,
|
||||
},
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap_err(),
|
||||
"exec cell 1 already has an active observer"
|
||||
@@ -606,6 +631,7 @@ async fn natural_completion_cleans_up_callbacks_before_responding() {
|
||||
..execute_request("")
|
||||
},
|
||||
delegate.clone(),
|
||||
/*preempt*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
Reference in New Issue
Block a user