feat: reload v2 agents on delivery (#26623)

## Summary

This is the first small step toward making multi-agent v2 agents durable
logical agents whose `ThreadManager` residency is only an implementation
detail.

This PR adds a narrow v2 reload-on-delivery hook:

- If a known v2 agent target is already loaded, delivery is unchanged.
- If the target is still registered but missing from `ThreadManager`,
delivery reloads that exact v2 thread from durable rollout history
before submitting the message.
- If the target is unknown, closed, missing from storage, or not a v2
thread, delivery still fails as not found.

The reload is wired only into existing-agent delivery paths: v2
`send_message` / `followup_task`, and legacy `send_input` when its
target is a known v2 agent.

## Stack

1. **Reload on delivery**: load known unloaded v2 agents before
`followup_task`, `send_message`, or `send_input` delivery. This PR.
2. **Residency LRU**: unload idle resident v2 agents from
`ThreadManager` without making them closed or unreachable.
3. **Execution concurrency**: count active non-root turns, not logical
agents or resident idle threads.
4. **Close semantics**: make v2 close interrupt-only and leave durable
agent identity intact.
5. **Resume cleanup**: remove user-facing v2 resume semantics;
addressing an unloaded durable agent reloads it implicitly.

## Validation

- Ran `just fmt`.
- Left broader tests and clippy to CI.
This commit is contained in:
jif
2026-06-05 18:18:29 +02:00
committed by GitHub
Unverified
parent e5af672d73
commit d5e4f01af4
4 changed files with 206 additions and 6 deletions
+73
View File
@@ -109,6 +109,79 @@ impl AgentControl {
.await
}
pub(crate) async fn ensure_v2_agent_loaded(
&self,
config: Config,
thread_id: ThreadId,
) -> CodexResult<()> {
let state = self.upgrade()?;
if state.get_thread(thread_id).await.is_ok() {
return Ok(());
}
if self.state.agent_metadata_for_thread(thread_id).is_none() {
return Err(CodexErr::ThreadNotFound(thread_id));
}
let stored_thread = state
.read_stored_thread(ReadThreadParams {
thread_id,
include_archived: true,
include_history: true,
})
.await?;
let stored_source = stored_thread.source.clone();
let stored_parent_thread_id = stored_thread.parent_thread_id;
let history = stored_thread
.history
.ok_or(CodexErr::ThreadNotFound(thread_id))?
.items;
let initial_history = InitialHistory::Resumed(ResumedHistory {
conversation_id: thread_id,
history,
rollout_path: stored_thread.rollout_path,
});
if initial_history.get_multi_agent_version() != Some(MultiAgentVersion::V2) {
return Err(CodexErr::ThreadNotFound(thread_id));
}
let (session_source, _) = initial_history
.get_resumed_session_sources()
.unwrap_or((stored_source, None));
let parent_thread_id = initial_history
.get_resumed_parent_thread_id()
.or(stored_parent_thread_id);
let inherited_shell_snapshot = self
.inherited_shell_snapshot_for_source(&state, Some(&session_source))
.await;
let inherited_exec_policy = self
.inherited_exec_policy_for_source(&state, Some(&session_source), &config)
.await;
match state
.resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions {
config,
initial_history,
agent_control: self.clone(),
session_source,
parent_thread_id,
inherited_shell_snapshot,
inherited_exec_policy,
})
.await
{
Ok(reloaded_thread) => {
state.notify_thread_created(reloaded_thread.thread_id);
Ok(())
}
Err(err) => {
if state.get_thread(thread_id).await.is_ok() {
return Ok(());
}
Err(err)
}
}
}
async fn spawn_agent_internal(
&self,
config: Config,
+107
View File
@@ -527,6 +527,113 @@ async fn send_inter_agent_communication_without_turn_queues_message_without_trig
));
}
#[tokio::test]
async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() {
let (home, mut config) = test_config().await;
let _ = config.features.enable(Feature::MultiAgentV2);
let _ = config.features.enable(Feature::Sqlite);
let state_db = init_state_db(&config).await;
let manager = ThreadManager::with_models_provider_home_and_state_for_tests(
CodexAuth::from_api_key("dummy"),
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
state_db.clone(),
);
let control = manager.agent_control();
let harness = AgentControlHarness {
_home: home,
config,
state_db,
manager,
control,
};
let (parent_thread_id, _parent_thread) = harness.start_thread().await;
let agent_path = AgentPath::try_from("/root/worker").expect("agent path");
let spawned_agent = harness
.control
.spawn_agent_with_metadata(
harness.config.clone(),
text_input("hello child"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 1,
agent_path: Some(agent_path.clone()),
agent_nickname: None,
agent_role: None,
})),
SpawnAgentOptions {
parent_thread_id: Some(parent_thread_id),
..Default::default()
},
)
.await
.expect("spawn_agent should succeed");
let child_thread = harness
.manager
.get_thread(spawned_agent.thread_id)
.await
.expect("child thread should exist");
child_thread
.inject_response_items(vec![assistant_message(
"child persisted",
Some(MessagePhase::FinalAnswer),
)])
.await
.expect("child rollout should persist with v2 metadata");
child_thread
.shutdown_and_wait()
.await
.expect("child thread should shut down");
assert!(
harness
.manager
.remove_thread(&spawned_agent.thread_id)
.await
.is_some()
);
match harness.manager.get_thread(spawned_agent.thread_id).await {
Err(CodexErr::ThreadNotFound(id)) => assert_eq!(id, spawned_agent.thread_id),
Err(err) => panic!("expected ThreadNotFound, got {err:?}"),
Ok(_) => panic!("expected thread to be removed"),
}
harness
.control
.ensure_v2_agent_loaded(harness.config.clone(), spawned_agent.thread_id)
.await
.expect("known v2 agent should reload");
let _ = harness
.manager
.get_thread(spawned_agent.thread_id)
.await
.expect("reloaded child thread should exist");
let communication = InterAgentCommunication::new(
AgentPath::root(),
agent_path,
Vec::new(),
"hello after reload".to_string(),
/*trigger_turn*/ false,
);
harness
.control
.send_inter_agent_communication(spawned_agent.thread_id, communication.clone())
.await
.expect("send_inter_agent_communication should succeed after reload");
let expected = (
spawned_agent.thread_id,
Op::InterAgentCommunication { communication },
);
let captured = harness
.manager
.captured_ops()
.into_iter()
.find(|entry| *entry == expected);
assert_eq!(captured, Some(expected));
}
#[tokio::test]
async fn encrypted_inter_agent_communication_clears_existing_last_task_message() {
let harness = AgentControlHarness::new().await;
@@ -42,8 +42,17 @@ impl ToolExecutor<ToolInvocation> for Handler {
let receiver_agent = session
.services
.agent_control
.get_agent_metadata(receiver_thread_id)
.unwrap_or_default();
.get_agent_metadata(receiver_thread_id);
if receiver_agent.is_some() {
let resume_config = build_agent_resume_config(turn.as_ref())?;
session
.services
.agent_control
.ensure_v2_agent_loaded(resume_config, receiver_thread_id)
.await
.map_err(|err| collab_agent_error(receiver_thread_id, err))?;
}
let receiver_agent = receiver_agent.unwrap_or_default();
if args.interrupt {
session
.services
@@ -75,7 +75,11 @@ pub(crate) async fn handle_message_string_tool(
.services
.agent_control
.get_agent_metadata(receiver_thread_id)
.unwrap_or_default();
.ok_or_else(|| {
FunctionCallError::RespondToModel(format!(
"agent with id {receiver_thread_id} not found"
))
})?;
if mode == MessageDeliveryMode::TriggerTurn
&& receiver_agent
.agent_path
@@ -86,6 +90,16 @@ pub(crate) async fn handle_message_string_tool(
"Follow-up tasks can't target the root agent".to_string(),
));
}
let receiver_agent_path = receiver_agent.agent_path.clone().ok_or_else(|| {
FunctionCallError::RespondToModel("target agent is missing an agent_path".to_string())
})?;
let resume_config = build_agent_resume_config(turn.as_ref())?;
session
.services
.agent_control
.ensure_v2_agent_loaded(resume_config, receiver_thread_id)
.await
.map_err(|err| collab_agent_error(receiver_thread_id, err))?;
session
.send_event(
&turn,
@@ -99,9 +113,6 @@ pub(crate) async fn handle_message_string_tool(
.into(),
)
.await;
let receiver_agent_path = receiver_agent.agent_path.clone().ok_or_else(|| {
FunctionCallError::RespondToModel("target agent is missing an agent_path".to_string())
})?;
let author = turn
.session_source
.get_agent_path()