mirror of
https://github.com/openai/codex.git
synced 2026-11-11 14:32:25 +00:00
Restore V2 agent identities on root thread resume (#32837)
## Why After a cold resume, a Multi-Agent V2 root thread must still be able to address its persisted descendants even though their runtimes are not reopened eagerly. ## What changed - Restore metadata for open descendant agents when resuming a V2 root thread. - Preserve each descendant's path, nickname, and role while leaving its runtime unloaded until a message targets it. - Reuse the spawn metadata preparation path so restored identities retain the same reservation and naming checks as newly spawned agents. ## Testing Added an end-to-end cold-resume test that follows up with a persisted agent by name and verifies that the original agent thread is loaded on demand. GitOrigin-RevId: a1663c1ae60eefadb948cd39f933cd3bbb49a492
This commit is contained in:
1 parent
19a000259d
commit
088239294a
5 files changed
+363
-19
No files matched your search
@@ -540,6 +540,32 @@ impl AgentControl {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn prepare_agent_metadata(
|
||||||
|
&self,
|
||||||
|
reservation: &mut crate::agent::registry::SpawnReservation,
|
||||||
|
config: &Config,
|
||||||
|
agent_path: Option<AgentPath>,
|
||||||
|
agent_role: Option<String>,
|
||||||
|
preferred_agent_nickname: Option<String>,
|
||||||
|
) -> CodexResult<AgentMetadata> {
|
||||||
|
if let Some(agent_path) = agent_path.as_ref() {
|
||||||
|
reservation.reserve_agent_path(agent_path)?;
|
||||||
|
}
|
||||||
|
let candidate_names = spawn::agent_nickname_candidates(config, agent_role.as_deref());
|
||||||
|
let candidate_name_refs: Vec<&str> = candidate_names.iter().map(String::as_str).collect();
|
||||||
|
let agent_nickname = Some(reservation.reserve_agent_nickname_with_preference(
|
||||||
|
&candidate_name_refs,
|
||||||
|
preferred_agent_nickname.as_deref(),
|
||||||
|
)?);
|
||||||
|
Ok(AgentMetadata {
|
||||||
|
agent_id: None,
|
||||||
|
agent_path,
|
||||||
|
agent_nickname,
|
||||||
|
agent_role,
|
||||||
|
last_task_message: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
fn prepare_thread_spawn(
|
fn prepare_thread_spawn(
|
||||||
&self,
|
&self,
|
||||||
@@ -554,29 +580,20 @@ impl AgentControl {
|
|||||||
if depth == 1 {
|
if depth == 1 {
|
||||||
self.state.register_root_thread(parent_thread_id);
|
self.state.register_root_thread(parent_thread_id);
|
||||||
}
|
}
|
||||||
if let Some(agent_path) = agent_path.as_ref() {
|
let agent_metadata = self.prepare_agent_metadata(
|
||||||
reservation.reserve_agent_path(agent_path)?;
|
reservation,
|
||||||
}
|
config,
|
||||||
let candidate_names = spawn::agent_nickname_candidates(config, agent_role.as_deref());
|
agent_path,
|
||||||
let candidate_name_refs: Vec<&str> = candidate_names.iter().map(String::as_str).collect();
|
agent_role,
|
||||||
let agent_nickname = Some(reservation.reserve_agent_nickname_with_preference(
|
preferred_agent_nickname,
|
||||||
&candidate_name_refs,
|
)?;
|
||||||
preferred_agent_nickname.as_deref(),
|
|
||||||
)?);
|
|
||||||
let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
||||||
parent_thread_id,
|
parent_thread_id,
|
||||||
depth,
|
depth,
|
||||||
agent_path: agent_path.clone(),
|
agent_path: agent_metadata.agent_path.clone(),
|
||||||
agent_nickname: agent_nickname.clone(),
|
agent_nickname: agent_metadata.agent_nickname.clone(),
|
||||||
agent_role: agent_role.clone(),
|
agent_role: agent_metadata.agent_role.clone(),
|
||||||
});
|
});
|
||||||
let agent_metadata = AgentMetadata {
|
|
||||||
agent_id: None,
|
|
||||||
agent_path,
|
|
||||||
agent_nickname,
|
|
||||||
agent_role,
|
|
||||||
last_task_message: None,
|
|
||||||
};
|
|
||||||
Ok((session_source, agent_metadata))
|
Ok((session_source, agent_metadata))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -95,6 +95,77 @@ fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: &
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl AgentControl {
|
impl AgentControl {
|
||||||
|
/// Restore persisted V2 agent identities without reopening their runtimes.
|
||||||
|
pub(crate) async fn restore_v2_agent_metadata(
|
||||||
|
&self,
|
||||||
|
config: &Config,
|
||||||
|
root_thread_id: ThreadId,
|
||||||
|
) {
|
||||||
|
self.state.register_root_thread(root_thread_id);
|
||||||
|
|
||||||
|
let Ok(state) = self.upgrade() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let Some(agent_graph_store) = state.agent_graph_store() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let descendant_ids = match agent_graph_store
|
||||||
|
.list_thread_spawn_descendants(
|
||||||
|
root_thread_id,
|
||||||
|
Some(codex_agent_graph_store::ThreadSpawnEdgeStatus::Open),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(descendant_ids) => descendant_ids,
|
||||||
|
Err(err) => {
|
||||||
|
warn!("failed to restore persisted V2 agent metadata for {root_thread_id}: {err}");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
for thread_id in descendant_ids {
|
||||||
|
if self.state.agent_metadata_for_thread(thread_id).is_some() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let restore_result = async {
|
||||||
|
let stored_thread = state
|
||||||
|
.read_stored_thread(ReadThreadParams {
|
||||||
|
thread_id,
|
||||||
|
include_archived: true,
|
||||||
|
include_history: false,
|
||||||
|
})
|
||||||
|
.await?;
|
||||||
|
let stored_agent_path = stored_thread
|
||||||
|
.agent_path
|
||||||
|
.as_deref()
|
||||||
|
.map(AgentPath::try_from)
|
||||||
|
.transpose()
|
||||||
|
.map_err(|err| {
|
||||||
|
CodexErr::InvalidRequest(format!("invalid stored agent path: {err}"))
|
||||||
|
})?;
|
||||||
|
let mut reservation = self.state.reserve_spawn_slot(/*max_threads*/ None)?;
|
||||||
|
let mut metadata = self.prepare_agent_metadata(
|
||||||
|
&mut reservation,
|
||||||
|
config,
|
||||||
|
stored_agent_path.or_else(|| stored_thread.source.get_agent_path()),
|
||||||
|
stored_thread
|
||||||
|
.agent_role
|
||||||
|
.or_else(|| stored_thread.source.get_agent_role()),
|
||||||
|
stored_thread
|
||||||
|
.agent_nickname
|
||||||
|
.or_else(|| stored_thread.source.get_nickname()),
|
||||||
|
)?;
|
||||||
|
metadata.agent_id = Some(thread_id);
|
||||||
|
reservation.commit(metadata);
|
||||||
|
Ok::<(), CodexErr>(())
|
||||||
|
}
|
||||||
|
.await;
|
||||||
|
if let Err(err) = restore_result {
|
||||||
|
warn!("failed to restore V2 agent metadata for {thread_id}: {err}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Spawn a new agent thread and submit the initial prompt.
|
/// Spawn a new agent thread and submit the initial prompt.
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub(crate) async fn spawn_agent(
|
pub(crate) async fn spawn_agent(
|
||||||
|
|||||||
@@ -793,6 +793,14 @@ impl ThreadManager {
|
|||||||
let (session_source, thread_source) = initial_history
|
let (session_source, thread_source) = initial_history
|
||||||
.get_resumed_session_sources()
|
.get_resumed_session_sources()
|
||||||
.unwrap_or_else(|| (self.state.session_source.clone(), None));
|
.unwrap_or_else(|| (self.state.session_source.clone(), None));
|
||||||
|
if let InitialHistory::Resumed(resumed) = &initial_history
|
||||||
|
&& initial_history.get_multi_agent_version() == Some(MultiAgentVersion::V2)
|
||||||
|
&& !session_source.is_non_root_agent()
|
||||||
|
{
|
||||||
|
agent_control
|
||||||
|
.restore_v2_agent_metadata(&config, resumed.conversation_id)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
Box::pin(self.state.spawn_thread_with_source(
|
Box::pin(self.state.spawn_thread_with_source(
|
||||||
config,
|
config,
|
||||||
initial_history,
|
initial_history,
|
||||||
|
|||||||
@@ -80,6 +80,7 @@ mod model_visible_layout;
|
|||||||
mod models_cache_ttl;
|
mod models_cache_ttl;
|
||||||
mod models_etag_responses;
|
mod models_etag_responses;
|
||||||
mod multi_agent_mode;
|
mod multi_agent_mode;
|
||||||
|
mod multi_agent_resume;
|
||||||
mod network_approval;
|
mod network_approval;
|
||||||
mod openai_file_mcp;
|
mod openai_file_mcp;
|
||||||
mod otel;
|
mod otel;
|
||||||
|
|||||||
@@ -0,0 +1,247 @@
|
|||||||
|
use anyhow::Result;
|
||||||
|
use codex_features::Feature;
|
||||||
|
use codex_protocol::protocol::AgentStatus;
|
||||||
|
use core_test_support::responses::ev_assistant_message;
|
||||||
|
use core_test_support::responses::ev_completed;
|
||||||
|
use core_test_support::responses::ev_function_call_with_namespace;
|
||||||
|
use core_test_support::responses::ev_response_created;
|
||||||
|
use core_test_support::responses::mount_sse_once_match;
|
||||||
|
use core_test_support::responses::sse;
|
||||||
|
use core_test_support::responses::start_mock_server;
|
||||||
|
use core_test_support::test_codex::test_codex;
|
||||||
|
use pretty_assertions::assert_eq;
|
||||||
|
use serde_json::Value;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::time::Instant;
|
||||||
|
use tokio::time::sleep;
|
||||||
|
|
||||||
|
const COLLABORATION_NAMESPACE: &str = "collaboration";
|
||||||
|
const SPAWN_CALL_ID: &str = "spawn-worker";
|
||||||
|
const FOLLOWUP_CALL_ID: &str = "followup-worker";
|
||||||
|
const INITIAL_PROMPT: &str = "spawn a durable worker";
|
||||||
|
const INITIAL_TASK: &str = "inspect the repository";
|
||||||
|
const FOLLOWUP_PROMPT: &str = "continue the durable worker";
|
||||||
|
const FOLLOWUP_TASK: &str = "inspect the tests too";
|
||||||
|
|
||||||
|
fn decoded_body(request: &wiremock::Request) -> Option<Vec<u8>> {
|
||||||
|
let is_zstd = request
|
||||||
|
.headers
|
||||||
|
.get("content-encoding")
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.is_some_and(|value| {
|
||||||
|
value
|
||||||
|
.split(',')
|
||||||
|
.any(|entry| entry.trim().eq_ignore_ascii_case("zstd"))
|
||||||
|
});
|
||||||
|
if is_zstd {
|
||||||
|
zstd::stream::decode_all(std::io::Cursor::new(&request.body)).ok()
|
||||||
|
} else {
|
||||||
|
Some(request.body.clone())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn body_contains(request: &wiremock::Request, text: &str) -> bool {
|
||||||
|
decoded_body(request)
|
||||||
|
.and_then(|body| String::from_utf8(body).ok())
|
||||||
|
.is_some_and(|body| body.contains(text))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn request_has_input_type(request: &wiremock::Request, input_type: &str) -> bool {
|
||||||
|
decoded_body(request)
|
||||||
|
.and_then(|body| serde_json::from_slice::<Value>(&body).ok())
|
||||||
|
.and_then(|body| body.get("input").and_then(Value::as_array).cloned())
|
||||||
|
.is_some_and(|items| {
|
||||||
|
items
|
||||||
|
.iter()
|
||||||
|
.any(|item| item.get("type").and_then(Value::as_str) == Some(input_type))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn configure_multi_agent_v2(config: &mut codex_core::config::Config) {
|
||||||
|
config
|
||||||
|
.features
|
||||||
|
.enable(Feature::Collab)
|
||||||
|
.expect("test config should allow feature update");
|
||||||
|
config
|
||||||
|
.features
|
||||||
|
.enable(Feature::MultiAgentV2)
|
||||||
|
.expect("test config should allow feature update");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||||
|
async fn cold_root_resume_restores_agent_identity_and_reloads_target_on_followup() -> Result<()> {
|
||||||
|
let server = start_mock_server().await;
|
||||||
|
let spawn_args = serde_json::to_string(&json!({
|
||||||
|
"message": INITIAL_TASK,
|
||||||
|
"task_name": "worker",
|
||||||
|
}))?;
|
||||||
|
mount_sse_once_match(
|
||||||
|
&server,
|
||||||
|
|request: &wiremock::Request| body_contains(request, INITIAL_PROMPT),
|
||||||
|
sse(vec![
|
||||||
|
ev_response_created("resp-spawn-1"),
|
||||||
|
ev_function_call_with_namespace(
|
||||||
|
SPAWN_CALL_ID,
|
||||||
|
COLLABORATION_NAMESPACE,
|
||||||
|
"spawn_agent",
|
||||||
|
&spawn_args,
|
||||||
|
),
|
||||||
|
ev_completed("resp-spawn-1"),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let initial_child_request = mount_sse_once_match(
|
||||||
|
&server,
|
||||||
|
|request: &wiremock::Request| {
|
||||||
|
request_has_input_type(request, "agent_message") && body_contains(request, INITIAL_TASK)
|
||||||
|
},
|
||||||
|
sse(vec![
|
||||||
|
ev_response_created("resp-worker-1"),
|
||||||
|
ev_assistant_message("msg-worker-1", "initial task complete"),
|
||||||
|
ev_completed("resp-worker-1"),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
mount_sse_once_match(
|
||||||
|
&server,
|
||||||
|
|request: &wiremock::Request| {
|
||||||
|
body_contains(request, SPAWN_CALL_ID)
|
||||||
|
&& !request_has_input_type(request, "agent_message")
|
||||||
|
},
|
||||||
|
sse(vec![
|
||||||
|
ev_response_created("resp-spawn-2"),
|
||||||
|
ev_assistant_message("msg-spawn-2", "worker spawned"),
|
||||||
|
ev_completed("resp-spawn-2"),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let mut initial_builder = test_codex().with_config(configure_multi_agent_v2);
|
||||||
|
let initial = initial_builder.build(&server).await?;
|
||||||
|
let root_thread_id = initial.session_configured.thread_id;
|
||||||
|
let home = initial.home.clone();
|
||||||
|
let rollout_path = initial
|
||||||
|
.codex
|
||||||
|
.rollout_path()
|
||||||
|
.expect("root rollout path")
|
||||||
|
.to_path_buf();
|
||||||
|
initial.submit_turn(INITIAL_PROMPT).await?;
|
||||||
|
|
||||||
|
let deadline = Instant::now() + Duration::from_secs(2);
|
||||||
|
let worker_thread_id = loop {
|
||||||
|
if let Some(thread_id) = initial
|
||||||
|
.thread_manager
|
||||||
|
.list_thread_ids()
|
||||||
|
.await
|
||||||
|
.into_iter()
|
||||||
|
.find(|thread_id| *thread_id != root_thread_id)
|
||||||
|
{
|
||||||
|
break thread_id;
|
||||||
|
}
|
||||||
|
if Instant::now() >= deadline {
|
||||||
|
anyhow::bail!("timed out waiting for spawned worker");
|
||||||
|
}
|
||||||
|
sleep(Duration::from_millis(10)).await;
|
||||||
|
};
|
||||||
|
let worker_thread = initial.thread_manager.get_thread(worker_thread_id).await?;
|
||||||
|
let deadline = Instant::now() + Duration::from_secs(2);
|
||||||
|
loop {
|
||||||
|
if matches!(
|
||||||
|
worker_thread.agent_status().await,
|
||||||
|
AgentStatus::Completed(_)
|
||||||
|
) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if Instant::now() >= deadline {
|
||||||
|
anyhow::bail!("timed out waiting for worker completion");
|
||||||
|
}
|
||||||
|
sleep(Duration::from_millis(10)).await;
|
||||||
|
}
|
||||||
|
assert!(
|
||||||
|
initial_child_request
|
||||||
|
.requests()
|
||||||
|
.iter()
|
||||||
|
.any(|request| request.body_contains_text(INITIAL_TASK))
|
||||||
|
);
|
||||||
|
worker_thread.flush_rollout().await?;
|
||||||
|
initial.codex.flush_rollout().await?;
|
||||||
|
drop(worker_thread);
|
||||||
|
drop(initial);
|
||||||
|
|
||||||
|
let followup_args = serde_json::to_string(&json!({
|
||||||
|
"target": "worker",
|
||||||
|
"message": FOLLOWUP_TASK,
|
||||||
|
}))?;
|
||||||
|
mount_sse_once_match(
|
||||||
|
&server,
|
||||||
|
|request: &wiremock::Request| body_contains(request, FOLLOWUP_PROMPT),
|
||||||
|
sse(vec![
|
||||||
|
ev_response_created("resp-followup-1"),
|
||||||
|
ev_function_call_with_namespace(
|
||||||
|
FOLLOWUP_CALL_ID,
|
||||||
|
COLLABORATION_NAMESPACE,
|
||||||
|
"followup_task",
|
||||||
|
&followup_args,
|
||||||
|
),
|
||||||
|
ev_completed("resp-followup-1"),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let followup_child_request = mount_sse_once_match(
|
||||||
|
&server,
|
||||||
|
|request: &wiremock::Request| {
|
||||||
|
request_has_input_type(request, "agent_message")
|
||||||
|
&& body_contains(request, FOLLOWUP_TASK)
|
||||||
|
},
|
||||||
|
sse(vec![
|
||||||
|
ev_response_created("resp-worker-2"),
|
||||||
|
ev_assistant_message("msg-worker-2", "follow-up complete"),
|
||||||
|
ev_completed("resp-worker-2"),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
mount_sse_once_match(
|
||||||
|
&server,
|
||||||
|
|request: &wiremock::Request| {
|
||||||
|
body_contains(request, FOLLOWUP_CALL_ID)
|
||||||
|
&& !request_has_input_type(request, "agent_message")
|
||||||
|
},
|
||||||
|
sse(vec![
|
||||||
|
ev_response_created("resp-followup-2"),
|
||||||
|
ev_assistant_message("msg-followup-2", "follow-up sent"),
|
||||||
|
ev_completed("resp-followup-2"),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let mut resume_builder = test_codex().with_config(configure_multi_agent_v2);
|
||||||
|
let resumed = resume_builder.resume(&server, home, rollout_path).await?;
|
||||||
|
assert_eq!(
|
||||||
|
resumed.thread_manager.list_thread_ids().await,
|
||||||
|
vec![root_thread_id]
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
resumed
|
||||||
|
.thread_manager
|
||||||
|
.get_thread(worker_thread_id)
|
||||||
|
.await
|
||||||
|
.is_err()
|
||||||
|
);
|
||||||
|
|
||||||
|
resumed.submit_turn(FOLLOWUP_PROMPT).await?;
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
followup_child_request
|
||||||
|
.requests()
|
||||||
|
.iter()
|
||||||
|
.any(|request| request.body_contains_text(FOLLOWUP_TASK))
|
||||||
|
);
|
||||||
|
resumed
|
||||||
|
.thread_manager
|
||||||
|
.get_thread(worker_thread_id)
|
||||||
|
.await
|
||||||
|
.expect("follow-up should lazily reload the original worker");
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
Reference in new issue
Block a user