feat: add thread spawn source for collab tools (#9769)

This commit is contained in:
jif-oai
2026-01-24 15:21:34 +01:00
committed by GitHub
Unverified
parent eb7558ba85
commit 515ac2cd19
7 changed files with 75 additions and 26 deletions
+17 -9
View File
@@ -39,12 +39,20 @@ impl AgentControl {
&self,
config: crate::config::Config,
prompt: String,
session_source: Option<codex_protocol::protocol::SessionSource>,
) -> CodexResult<ThreadId> {
let state = self.upgrade()?;
let reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?;
// The same `AgentControl` is sent to spawn the thread.
let new_thread = state.spawn_new_thread(config, self.clone()).await?;
let new_thread = match session_source {
Some(session_source) => {
state
.spawn_new_thread_with_source(config, self.clone(), session_source)
.await?
}
None => state.spawn_new_thread(config, self.clone()).await?,
};
reservation.commit(new_thread.thread_id);
// Notify a new thread has been created. This notification will be processed by clients
@@ -268,7 +276,7 @@ mod tests {
let control = AgentControl::default();
let (_home, config) = test_config().await;
let err = control
.spawn_agent(config, "hello".to_string())
.spawn_agent(config, "hello".to_string(), None)
.await
.expect_err("spawn_agent should fail without a manager");
assert_eq!(
@@ -370,7 +378,7 @@ mod tests {
let harness = AgentControlHarness::new().await;
let thread_id = harness
.control
.spawn_agent(harness.config.clone(), "spawned".to_string())
.spawn_agent(harness.config.clone(), "spawned".to_string(), None)
.await
.expect("spawn_agent should succeed");
let _thread = harness
@@ -417,12 +425,12 @@ mod tests {
.expect("start thread");
let first_agent_id = control
.spawn_agent(config.clone(), "hello".to_string())
.spawn_agent(config.clone(), "hello".to_string(), None)
.await
.expect("spawn_agent should succeed");
let err = control
.spawn_agent(config, "hello again".to_string())
.spawn_agent(config, "hello again".to_string(), None)
.await
.expect_err("spawn_agent should respect max threads");
let CodexErr::AgentLimitReached {
@@ -455,7 +463,7 @@ mod tests {
let control = manager.agent_control();
let first_agent_id = control
.spawn_agent(config.clone(), "hello".to_string())
.spawn_agent(config.clone(), "hello".to_string(), None)
.await
.expect("spawn_agent should succeed");
let _ = control
@@ -464,7 +472,7 @@ mod tests {
.expect("shutdown agent");
let second_agent_id = control
.spawn_agent(config.clone(), "hello again".to_string())
.spawn_agent(config.clone(), "hello again".to_string(), None)
.await
.expect("spawn_agent should succeed after shutdown");
let _ = control
@@ -490,12 +498,12 @@ mod tests {
let cloned = control.clone();
let first_agent_id = cloned
.spawn_agent(config.clone(), "hello".to_string())
.spawn_agent(config.clone(), "hello".to_string(), None)
.await
.expect("spawn_agent should succeed");
let err = control
.spawn_agent(config, "hello again".to_string())
.spawn_agent(config, "hello again".to_string(), None)
.await
.expect_err("spawn_agent should respect shared guard");
let CodexErr::AgentLimitReached { max_threads } = err else {
+5 -7
View File
@@ -226,13 +226,11 @@ impl ModelClient {
let mut extra_headers = ApiHeaderMap::new();
if let SessionSource::SubAgent(sub) = &self.state.session_source {
let subagent = if let crate::protocol::SubAgentSource::Other(label) = sub {
label.clone()
} else {
serde_json::to_value(sub)
.ok()
.and_then(|v| v.as_str().map(std::string::ToString::to_string))
.unwrap_or_else(|| "other".to_string())
let subagent = match sub {
crate::protocol::SubAgentSource::Review => "review".to_string(),
crate::protocol::SubAgentSource::Compact => "compact".to_string(),
crate::protocol::SubAgentSource::ThreadSpawn { .. } => "collab_spawn".to_string(),
crate::protocol::SubAgentSource::Other(label) => label.clone(),
};
if let Ok(val) = HeaderValue::from_str(&subagent) {
extra_headers.insert("x-openai-subagent", val);
+31 -2
View File
@@ -314,11 +314,22 @@ impl ThreadManagerState {
config: Config,
agent_control: AgentControl,
) -> CodexResult<NewThread> {
self.spawn_thread(
self.spawn_new_thread_with_source(config, agent_control, self.session_source.clone())
.await
}
pub(crate) async fn spawn_new_thread_with_source(
&self,
config: Config,
agent_control: AgentControl,
session_source: SessionSource,
) -> CodexResult<NewThread> {
self.spawn_thread_with_source(
config,
InitialHistory::New,
Arc::clone(&self.auth_manager),
agent_control,
session_source,
)
.await
}
@@ -330,6 +341,24 @@ impl ThreadManagerState {
initial_history: InitialHistory,
auth_manager: Arc<AuthManager>,
agent_control: AgentControl,
) -> CodexResult<NewThread> {
self.spawn_thread_with_source(
config,
initial_history,
auth_manager,
agent_control,
self.session_source.clone(),
)
.await
}
pub(crate) async fn spawn_thread_with_source(
&self,
config: Config,
initial_history: InitialHistory,
auth_manager: Arc<AuthManager>,
agent_control: AgentControl,
session_source: SessionSource,
) -> CodexResult<NewThread> {
let CodexSpawnOk {
codex, thread_id, ..
@@ -339,7 +368,7 @@ impl ThreadManagerState {
Arc::clone(&self.models_manager),
Arc::clone(&self.skills_manager),
initial_history,
self.session_source.clone(),
session_source,
agent_control,
)
.await?;
+9 -1
View File
@@ -78,6 +78,8 @@ impl ToolHandler for CollabHandler {
mod spawn {
use super::*;
use crate::agent::AgentRole;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use std::sync::Arc;
#[derive(Debug, Deserialize)]
@@ -125,7 +127,13 @@ mod spawn {
let result = session
.services
.agent_control
.spawn_agent(config, prompt.clone())
.spawn_agent(
config,
prompt.clone(),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: session.conversation_id,
})),
)
.await
.map_err(collab_spawn_error);
let (new_thread_id, status) = match &result {