Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion codex-rs/core/src/codex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -485,7 +485,7 @@ pub(crate) struct SessionSettingsUpdate {

impl Session {
/// Don't expand the number of mutated arguments on config. We are in the process of getting rid of it.
fn build_per_turn_config(session_configuration: &SessionConfiguration) -> Config {
pub(crate) fn build_per_turn_config(session_configuration: &SessionConfiguration) -> Config {
// todo(aibrahim): store this state somewhere else so we don't need to mut config
let config = session_configuration.original_config_do_not_use.clone();
let mut per_turn_config = (*config).clone();
Expand Down
8 changes: 8 additions & 0 deletions codex-rs/core/src/features.rs
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,8 @@ pub enum Feature {
PowershellUtf8,
/// Compress request bodies (zstd) when sending streaming requests to codex-backend.
EnableRequestCompression,
/// Enable collab tools.
Collab,
}

impl Feature {
Expand Down Expand Up @@ -398,6 +400,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::Experimental,
default_enabled: false,
},
FeatureSpec {
id: Feature::Collab,
key: "collab",
stage: Stage::Experimental,
default_enabled: false,
},
FeatureSpec {
id: Feature::Tui2,
key: "tui2",
Expand Down
194 changes: 194 additions & 0 deletions codex-rs/core/src/tools/handlers/collab.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,194 @@
use crate::codex::TurnContext;
use crate::config::Config;
use crate::error::CodexErr;
use crate::function_tool::FunctionCallError;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::parse_arguments;
use crate::tools::registry::ToolHandler;
use crate::tools::registry::ToolKind;
use async_trait::async_trait;
use codex_protocol::ThreadId;
use serde::Deserialize;

pub struct CollabHandler;

pub(crate) const DEFAULT_WAIT_TIMEOUT_MS: i64 = 30_000;
pub(crate) const MAX_WAIT_TIMEOUT_MS: i64 = 300_000;

#[derive(Debug, Deserialize)]
struct SpawnAgentArgs {
message: String,
}

#[derive(Debug, Deserialize)]
struct SendInputArgs {
id: String,
message: String,
}

#[derive(Debug, Deserialize)]
struct WaitArgs {
id: String,
timeout_ms: Option<i64>,
}

#[derive(Debug, Deserialize)]
struct CloseAgentArgs {
id: String,
}

#[async_trait]
impl ToolHandler for CollabHandler {
fn kind(&self) -> ToolKind {
ToolKind::Function
}

fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(payload, ToolPayload::Function { .. })
}

async fn handle(&self, invocation: ToolInvocation) -> Result<ToolOutput, FunctionCallError> {
let ToolInvocation {
session,
turn,
tool_name,
payload,
..
} = invocation;

let arguments = match payload {
ToolPayload::Function { arguments } => arguments,
_ => {
return Err(FunctionCallError::RespondToModel(
"collab handler received unsupported payload".to_string(),
));
}
};

match tool_name.as_str() {
"spawn_agent" => handle_spawn_agent(session, turn, arguments).await,
"send_input" => handle_send_input(session, arguments).await,
"wait" => handle_wait(arguments).await,
"close_agent" => handle_close_agent(arguments).await,
other => Err(FunctionCallError::RespondToModel(format!(
"unsupported collab tool {other}"
))),
}
}
}

async fn handle_spawn_agent(
session: std::sync::Arc<crate::codex::Session>,
turn: std::sync::Arc<TurnContext>,
arguments: String,
) -> Result<ToolOutput, FunctionCallError> {
let args: SpawnAgentArgs = parse_arguments(&arguments)?;
if args.message.trim().is_empty() {
return Err(FunctionCallError::RespondToModel(
"Empty message can't be send to an agent".to_string(),
));
}
let config = build_agent_spawn_config(turn.as_ref())?;
let result = session
.services
.agent_control
.spawn_agent(config, args.message, true)
.await
.map_err(|err| FunctionCallError::Fatal(err.to_string()))?;

Ok(ToolOutput::Function {
content: format!("agent_id: {result}"),
success: Some(true),
content_items: None,
})
}

async fn handle_send_input(
session: std::sync::Arc<crate::codex::Session>,
arguments: String,
) -> Result<ToolOutput, FunctionCallError> {
let args: SendInputArgs = parse_arguments(&arguments)?;
let agent_id = agent_id(&args.id)?;
if args.message.trim().is_empty() {
return Err(FunctionCallError::RespondToModel(
"Empty message can't be send to an agent".to_string(),
));
}
let content = session
.services
.agent_control
.send_prompt(agent_id, args.message)
.await
.map_err(|err| match err {
CodexErr::ThreadNotFound(id) => {
FunctionCallError::RespondToModel(format!("agent with id {id} not found"))
}
err => FunctionCallError::Fatal(err.to_string()),
})?;

Ok(ToolOutput::Function {
content,
success: Some(true),
content_items: None,
})
}

async fn handle_wait(arguments: String) -> Result<ToolOutput, FunctionCallError> {
let args: WaitArgs = parse_arguments(&arguments)?;
let _agent_id = agent_id(&args.id)?;

let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS);
if timeout_ms <= 0 {
return Err(FunctionCallError::RespondToModel(
"timeout_ms must be greater than zero".to_string(),
));
}
let _timeout_ms = timeout_ms.min(MAX_WAIT_TIMEOUT_MS);
// TODO(jif): implement agent wait once lifecycle tracking is wired up.
Err(FunctionCallError::Fatal("wait not implemented".to_string()))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid fatal errors for wait tool

With Feature::Collab enabled, build_specs registers the wait tool, so the model is allowed to call it; this handler always returns FunctionCallError::Fatal, which ToolRouter::dispatch_tool_call treats as a hard failure rather than a tool output. That means any wait call will abort the entire turn instead of returning a normal tool failure, making the collab flow unusable whenever the model tries to wait. Consider hiding the tool until implemented or returning a non-fatal failure response.

Useful? React with 👍 / 👎.

}

async fn handle_close_agent(arguments: String) -> Result<ToolOutput, FunctionCallError> {
let args: CloseAgentArgs = parse_arguments(&arguments)?;
let _agent_id = agent_id(&args.id)?;
// TODO(jif): implement agent shutdown and return the final status.
Err(FunctionCallError::Fatal(
"close_agent not implemented".to_string(),
Comment on lines +157 to +158
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid fatal errors for close_agent tool

The close_agent tool is also registered when collab is enabled, but its handler always returns FunctionCallError::Fatal. Since fatal tool errors short-circuit dispatch instead of producing a tool failure output, any close_agent call will fail the entire turn. If this tool is exposed to the model before being implemented, it will cause abrupt failures rather than a recoverable tool error.

Useful? React with 👍 / 👎.

))
}

fn agent_id(id: &str) -> Result<ThreadId, FunctionCallError> {
ThreadId::from_string(id)
.map_err(|e| FunctionCallError::RespondToModel(format!("invalid agent id {id}: {e:?}")))
}

fn build_agent_spawn_config(turn: &TurnContext) -> Result<Config, FunctionCallError> {
let base_config = turn.client.config();
let mut config = (*base_config).clone();
config.model = Some(turn.client.get_model());
config.model_provider = turn.client.get_provider();
config.model_reasoning_effort = turn.client.get_reasoning_effort();
config.model_reasoning_summary = turn.client.get_reasoning_summary();
config.developer_instructions = turn.developer_instructions.clone();
config.base_instructions = turn.base_instructions.clone();
config.compact_prompt = turn.compact_prompt.clone();
config.user_instructions = turn.user_instructions.clone();
config.shell_environment_policy = turn.shell_environment_policy.clone();
config.codex_linux_sandbox_exe = turn.codex_linux_sandbox_exe.clone();
config.cwd = turn.cwd.clone();
config
.approval_policy
.set(turn.approval_policy)
.map_err(|err| {
FunctionCallError::RespondToModel(format!("approval_policy is invalid: {err}"))
})?;
config
.sandbox_policy
.set(turn.sandbox_policy.clone())
.map_err(|err| {
FunctionCallError::RespondToModel(format!("sandbox_policy is invalid: {err}"))
})?;
Ok(config)
}
2 changes: 2 additions & 0 deletions codex-rs/core/src/tools/handlers/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
pub mod apply_patch;
pub(crate) mod collab;
mod grep_files;
mod list_dir;
mod mcp;
Expand All @@ -15,6 +16,7 @@ use serde::Deserialize;

use crate::function_tool::FunctionCallError;
pub use apply_patch::ApplyPatchHandler;
pub use collab::CollabHandler;
pub use grep_files::GrepFilesHandler;
pub use list_dir::ListDirHandler;
pub use mcp::McpHandler;
Expand Down
Loading
Loading