mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user