mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +00:00
Preserve rollout tracing across sub-agent reloads (#57615)
## Why Reloading an evicted sub-agent disabled its rollout tracing, leaving later turns out of the trace. Replacement Code Mode sessions could also reuse cell IDs from earlier sessions. ## What changed - Reattach reloaded children to the root rollout trace and emit `ThreadResumed` without duplicating `ThreadStarted`. - Reopen existing threads during replay, preserving their identity and history while clearing stale runtime cell and terminal process mappings. - Allocate Code Mode cell generations across sessions for both gRPC and stdio hosts to prevent ID reuse after reload. - Record an aborted trace when session startup fails before the session loop begins. ## Testing Add regression coverage for child eviction and reload, including a failed reload, continued inference tracing, distinct cell IDs, and preserved mail and follow-up context. Extend replay tests for history preservation and reused runtime handles, and host tests for cell IDs across session replacement and reconnection. GitOrigin-RevId: 6c146c04b5cda9295afb7d973c0d5474b363a822
This commit is contained in:
committed by
copyberry
parent
c0e88f3c04
commit
68b67f0ff4
@@ -976,7 +976,7 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() ->
|
||||
.await?;
|
||||
assert_eq!(
|
||||
actual,
|
||||
text_response("1", "fast-before", actual.code_mode_host_duration())
|
||||
text_response("g2:1", "fast-before", actual.code_mode_host_duration())
|
||||
);
|
||||
|
||||
delegate.release.add_permits(/*n*/ 1);
|
||||
@@ -991,7 +991,7 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() ->
|
||||
let actual = fast_response?;
|
||||
assert_eq!(
|
||||
actual,
|
||||
text_response("2", "fast-during", actual.code_mode_host_duration())
|
||||
text_response("g2:2", "fast-during", actual.code_mode_host_duration())
|
||||
);
|
||||
let actual = slow_response
|
||||
.context("large unary tool response did not complete")?
|
||||
@@ -1128,7 +1128,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
|
||||
.await?;
|
||||
assert_eq!(
|
||||
actual,
|
||||
text_response("1", "isolated", actual.code_mode_host_duration())
|
||||
text_response("g2:1", "isolated", actual.code_mode_host_duration())
|
||||
);
|
||||
let actual = limited
|
||||
.terminate(cell_id("1").into())
|
||||
@@ -1221,6 +1221,24 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
|
||||
let actual = execute(
|
||||
&session,
|
||||
request(r#"text("before reload");"#),
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(
|
||||
actual,
|
||||
text_response("1", "before reload", actual.code_mode_host_duration())
|
||||
);
|
||||
session.shutdown().await.map_err(anyhow::Error::msg)?;
|
||||
// Reloading a thread creates a new session through the same shared provider.
|
||||
let session = provider
|
||||
.clone()
|
||||
.create_session()
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
|
||||
let saved = execute(
|
||||
&session,
|
||||
request(r#"store("saved", "before restart"); text(load("saved"));"#),
|
||||
@@ -1229,7 +1247,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
|
||||
.await?;
|
||||
assert_eq!(
|
||||
saved,
|
||||
text_response("1", "before restart", saved.code_mode_host_duration())
|
||||
text_response("g2:1", "before restart", saved.code_mode_host_duration())
|
||||
);
|
||||
|
||||
let mut pending = request("await tools.echo({generation: 1}); await new Promise(() => {});");
|
||||
@@ -1240,7 +1258,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
let old_cell_id = started.cell_id.clone();
|
||||
assert_eq!(old_cell_id, cell_id("2"));
|
||||
assert_eq!(old_cell_id, cell_id("g2:2"));
|
||||
assert!(matches!(
|
||||
started.initial_response().await,
|
||||
Ok(RuntimeResponse::Yielded { .. })
|
||||
@@ -1372,7 +1390,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
|
||||
);
|
||||
let mut replacement_cell_ids = [callback_cell_id.as_str(), concurrent_cell_id.as_str()];
|
||||
replacement_cell_ids.sort_unstable();
|
||||
assert_eq!(replacement_cell_ids, ["g2:1", "g2:2"]);
|
||||
assert_eq!(replacement_cell_ids, ["g3:1", "g3:2"]);
|
||||
assert_eq!(
|
||||
delegate
|
||||
.invocations
|
||||
@@ -1405,7 +1423,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
let replacement_cell_id = started.cell_id.clone();
|
||||
assert_eq!(replacement_cell_id, cell_id("g2:3"));
|
||||
assert_eq!(replacement_cell_id, cell_id("g3:3"));
|
||||
let actual = started
|
||||
.initial_response()
|
||||
.await
|
||||
|
||||
@@ -41,7 +41,7 @@ try { image(value); } catch { text(calls); }
|
||||
.await?;
|
||||
assert_eq!(
|
||||
other_response,
|
||||
text_response("1", "yes", other_response.code_mode_host_duration())
|
||||
text_response("g2:1", "yes", other_response.code_mode_host_duration())
|
||||
);
|
||||
|
||||
let delegate = Arc::new(RecordingDelegate::default());
|
||||
@@ -90,7 +90,7 @@ try { image(value); } catch { text(calls); }
|
||||
.await?;
|
||||
assert_eq!(
|
||||
other_response,
|
||||
text_response("2", "yes", other_response.code_mode_host_duration())
|
||||
text_response("g2:2", "yes", other_response.code_mode_host_duration())
|
||||
);
|
||||
session.shutdown().await.map_err(anyhow::Error::msg)?;
|
||||
other_session.shutdown().await.map_err(anyhow::Error::msg)?;
|
||||
|
||||
@@ -92,7 +92,7 @@ async fn account_change_revokes_active_rpc_and_renews_the_cached_client() -> Res
|
||||
.await?;
|
||||
assert_eq!(
|
||||
response,
|
||||
text_response("1", "replacement", response.code_mode_host_duration())
|
||||
text_response("g2:1", "replacement", response.code_mode_host_duration())
|
||||
);
|
||||
session.shutdown().await.map_err(anyhow::Error::msg)?;
|
||||
replacement.shutdown().await.map_err(anyhow::Error::msg)?;
|
||||
|
||||
@@ -393,6 +393,29 @@ async fn remote_session_persists_values_forwards_delegates_and_controls_cells()
|
||||
.await
|
||||
.expect("create remote session");
|
||||
|
||||
let actual = execute(
|
||||
&session,
|
||||
execute_request(""),
|
||||
Arc::new(RecordingDelegate::default()),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
actual,
|
||||
RuntimeResponse::Result {
|
||||
code_mode_host_duration: actual.code_mode_host_duration(),
|
||||
cell_id: cell_id("1"),
|
||||
content_items: Vec::new(),
|
||||
result_buffered_for: Duration::ZERO,
|
||||
error_text: None,
|
||||
}
|
||||
);
|
||||
session.shutdown().await.expect("shutdown before reload");
|
||||
// Reloading a thread must not reuse cell IDs from its previous session.
|
||||
let session = provider
|
||||
.create_session()
|
||||
.await
|
||||
.expect("create replacement session");
|
||||
|
||||
let actual = execute(
|
||||
&session,
|
||||
execute_request(r#"store("key", "persisted");"#),
|
||||
@@ -403,7 +426,7 @@ async fn remote_session_persists_values_forwards_delegates_and_controls_cells()
|
||||
actual,
|
||||
RuntimeResponse::Result {
|
||||
code_mode_host_duration: actual.code_mode_host_duration(),
|
||||
cell_id: cell_id("1"),
|
||||
cell_id: cell_id("g2:1"),
|
||||
content_items: Vec::new(),
|
||||
result_buffered_for: Duration::ZERO,
|
||||
error_text: None,
|
||||
@@ -432,7 +455,7 @@ text(result.value);
|
||||
actual,
|
||||
RuntimeResponse::Result {
|
||||
code_mode_host_duration: actual.code_mode_host_duration(),
|
||||
cell_id: cell_id("2"),
|
||||
cell_id: cell_id("g2:2"),
|
||||
content_items: vec![FunctionCallOutputContentItem::InputText {
|
||||
text: "output".to_string(),
|
||||
}],
|
||||
@@ -443,7 +466,7 @@ text(result.value);
|
||||
assert_eq!(
|
||||
*delegate.invocations.lock().expect("invocations lock"),
|
||||
vec![CodeModeNestedToolCall {
|
||||
cell_id: cell_id("2"),
|
||||
cell_id: cell_id("g2:2"),
|
||||
runtime_tool_call_id: "tool-1".to_string(),
|
||||
tool_name: ToolName::plain("echo"),
|
||||
tool_kind: CodeModeToolKind::Function,
|
||||
@@ -452,7 +475,7 @@ text(result.value);
|
||||
);
|
||||
assert_eq!(
|
||||
*delegate.notifications.lock().expect("notifications lock"),
|
||||
vec![("call-2".to_string(), cell_id("2"), "notice".to_string())]
|
||||
vec![("call-2".to_string(), cell_id("g2:2"), "notice".to_string())]
|
||||
);
|
||||
|
||||
let mut pending_request = execute_request("await new Promise(() => {});");
|
||||
@@ -463,14 +486,14 @@ text(result.value);
|
||||
actual,
|
||||
RuntimeResponse::Yielded {
|
||||
code_mode_host_duration: actual.code_mode_host_duration(),
|
||||
cell_id: cell_id("3"),
|
||||
cell_id: cell_id("g2:3"),
|
||||
content_items: Vec::new(),
|
||||
}
|
||||
);
|
||||
let actual = session
|
||||
.wait(
|
||||
WaitRequest {
|
||||
cell_id: cell_id("3"),
|
||||
cell_id: cell_id("g2:3"),
|
||||
yield_time_ms: 1,
|
||||
max_output_tokens: None,
|
||||
},
|
||||
@@ -482,19 +505,19 @@ text(result.value);
|
||||
actual,
|
||||
WaitOutcome::LiveCell(RuntimeResponse::Yielded {
|
||||
code_mode_host_duration: actual.code_mode_host_duration(),
|
||||
cell_id: cell_id("3"),
|
||||
cell_id: cell_id("g2:3"),
|
||||
content_items: Vec::new(),
|
||||
})
|
||||
);
|
||||
let actual = session
|
||||
.terminate(cell_id("3").into())
|
||||
.terminate(cell_id("g2:3").into())
|
||||
.await
|
||||
.expect("terminate cell");
|
||||
assert_eq!(
|
||||
actual,
|
||||
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
|
||||
code_mode_host_duration: actual.code_mode_host_duration(),
|
||||
cell_id: cell_id("3"),
|
||||
cell_id: cell_id("g2:3"),
|
||||
content_items: Vec::new(),
|
||||
})
|
||||
);
|
||||
@@ -502,14 +525,14 @@ text(result.value);
|
||||
session.shutdown().await.expect("shutdown remote session");
|
||||
assert_eq!(
|
||||
*delegate.closed_cells.lock().expect("closed cells lock"),
|
||||
vec![cell_id("2"), cell_id("3")]
|
||||
vec![cell_id("g2:2"), cell_id("g2:3")]
|
||||
);
|
||||
assert_eq!(
|
||||
*first_delegate
|
||||
.closed_cells
|
||||
.lock()
|
||||
.expect("closed cells lock"),
|
||||
vec![cell_id("1")]
|
||||
vec![cell_id("g2:1")]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1089,7 +1112,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
|
||||
actual,
|
||||
RuntimeResponse::Result {
|
||||
code_mode_host_duration: actual.code_mode_host_duration(),
|
||||
cell_id: cell_id("g2:1"),
|
||||
cell_id: cell_id("g3:1"),
|
||||
content_items: vec![FunctionCallOutputContentItem::InputText {
|
||||
text: "replacement".to_string(),
|
||||
}],
|
||||
|
||||
Reference in New Issue
Block a user