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:
cassirer-openai
2026-10-28 01:29:34 +00:00
committed by copyberry
parent c0e88f3c04
commit 68b67f0ff4
19 changed files with 367 additions and 100 deletions
+25 -7
View File
@@ -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)?;
+35 -12
View File
@@ -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(),
}],