openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge 1f12d07fd1 into sapling-pr-archive-bolinfest

Browse Source
This commit is contained in:
Michael Bolin
2026-05-12 16:51:01 -07:00
committed by GitHub
parent c9a38fad25 1f12d07fd1
commit 3d9af7a5b5
45 changed files with 1166 additions and 239 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1 -1
View File
@@ -2516,7 +2516,6 @@ dependencies = [
"codex-utils-path",
"codex-utils-plugins",
"codex-utils-pty",
"codex-utils-readiness",
"codex-utils-stream-parser",
"codex-utils-string",
"codex-utils-template",
@@ -2742,6 +2741,7 @@ dependencies = [
"ctor 0.6.3",
"futures",
"pretty_assertions",
"prost 0.14.3",
"reqwest",
"serde",
"serde_json",
codex-rs/Cargo.toml
-1
View File
@@ -228,7 +228,6 @@ codex-utils-output-truncation = { path = "utils/output-truncation" }
codex-utils-path = { path = "utils/path-utils" }
codex-utils-plugins = { path = "utils/plugins" }
codex-utils-pty = { path = "utils/pty" }
codex-utils-readiness = { path = "utils/readiness" }
codex-utils-rustls-provider = { path = "utils/rustls-provider" }
codex-utils-sandbox-summary = { path = "utils/sandbox-summary" }
codex-utils-sleep-inhibitor = { path = "utils/sleep-inhibitor" }
codex-rs/code-mode/src/runtime/callbacks.rs
+4 -7
View File
@@ -42,20 +42,16 @@ pub(super) fn tool_callback(
let promise = resolver.get_promise(scope);
let resolver = v8::Global::new(scope, resolver);
let tool_name = {
let (tool_name, tool_kind) = {
let Some(state) = scope.get_slot::<RuntimeState>() else {
throw_type_error(scope, "runtime state unavailable");
return;
};
let Some(tool_name) = state
.enabled_tools
.get(tool_index)
.map(|tool| tool.tool_name.clone())
else {
let Some(tool) = state.enabled_tools.get(tool_index) else {
throw_type_error(scope, "tool callback data is out of range");
return;
};
tool_name
(tool.tool_name.clone(), tool.kind)
};
let Some(state) = scope.get_slot_mut::<RuntimeState>() else {
@@ -69,6 +65,7 @@ pub(super) fn tool_callback(
let _ = event_tx.send(RuntimeEvent::ToolCall {
id,
name: tool_name,
kind: tool_kind,
input,
});
retval.set(promise.into());
codex-rs/code-mode/src/runtime/mod.rs
+3
View File
@@ -14,6 +14,7 @@ use serde::Serialize;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
use crate::description::CodeModeToolKind;
use crate::description::EnabledToolMetadata;
use crate::description::ToolDefinition;
use crate::description::enabled_tool_metadata;
@@ -97,6 +98,7 @@ pub struct CodeModeNestedToolCall {
pub cell_id: String,
pub runtime_tool_call_id: String,
pub tool_name: ToolName,
pub tool_kind: CodeModeToolKind,
pub input: Option<JsonValue>,
}
@@ -126,6 +128,7 @@ pub(crate) enum RuntimeEvent {
ToolCall {
id: String,
name: ToolName,
kind: CodeModeToolKind,
input: Option<JsonValue>,
},
Notify {
codex-rs/code-mode/src/service.rs
+7 -1
View File
@@ -385,11 +385,17 @@ async fn run_session_control(
text,
}).await;
}
RuntimeEvent::ToolCall { id, name, input } => {
RuntimeEvent::ToolCall {
id,
name,
kind,
input,
} => {
let tool_call = CodeModeNestedToolCall {
cell_id: cell_id.clone(),
runtime_tool_call_id: id,
tool_name: name,
tool_kind: kind,
input,
};
let _ = inner
codex-rs/config/src/config_toml.rs
+3
View File
@@ -162,6 +162,9 @@ pub struct ConfigToml {
/// Whether to inject the `<apps_instructions>` developer block.
pub include_apps_instructions: Option<bool>,
/// Whether to inject the `<collaboration_mode>` developer block.
pub include_collaboration_mode_instructions: Option<bool>,
/// Whether to inject the `<environment_context>` user block.
pub include_environment_context: Option<bool>,
codex-rs/config/src/profile_toml.rs
+1
View File
@@ -57,6 +57,7 @@ pub struct ConfigProfile {
pub include_apply_patch_tool: Option<bool>,
pub include_permissions_instructions: Option<bool>,
pub include_apps_instructions: Option<bool>,
pub include_collaboration_mode_instructions: Option<bool>,
pub include_environment_context: Option<bool>,
pub experimental_use_unified_exec_tool: Option<bool>,
pub experimental_use_freeform_apply_patch: Option<bool>,
codex-rs/core/Cargo.toml
-1
View File
@@ -70,7 +70,6 @@ codex-utils-output-truncation = { workspace = true }
codex-utils-path = { workspace = true }
codex-utils-plugins = { workspace = true }
codex-utils-pty = { workspace = true }
codex-utils-readiness = { workspace = true }
codex-utils-string = { workspace = true }
codex-utils-stream-parser = { workspace = true }
codex-utils-template = { workspace = true }
codex-rs/core/config.schema.json
+7
View File
@@ -622,6 +622,9 @@
"include_apps_instructions": {
"type": "boolean"
},
"include_collaboration_mode_instructions": {
"type": "boolean"
},
"include_environment_context": {
"type": "boolean"
},
@@ -4358,6 +4361,10 @@
"description": "Whether to inject the `<apps_instructions>` developer block.",
"type": "boolean"
},
"include_collaboration_mode_instructions": {
"description": "Whether to inject the `<collaboration_mode>` developer block.",
"type": "boolean"
},
"include_environment_context": {
"description": "Whether to inject the `<environment_context>` user block.",
"type": "boolean"
codex-rs/core/src/config/config_tests.rs
+7
View File
@@ -7369,6 +7369,7 @@ async fn test_precedence_fixture_with_o3_profile() -> std::io::Result<()> {
guardian_policy_config: None,
include_permissions_instructions: true,
include_apps_instructions: true,
include_collaboration_mode_instructions: true,
include_skill_instructions: true,
include_environment_context: true,
compact_prompt: None,
@@ -7816,6 +7817,7 @@ async fn test_precedence_fixture_with_gpt3_profile() -> std::io::Result<()> {
guardian_policy_config: None,
include_permissions_instructions: true,
include_apps_instructions: true,
include_collaboration_mode_instructions: true,
include_skill_instructions: true,
include_environment_context: true,
compact_prompt: None,
@@ -7977,6 +7979,7 @@ async fn test_precedence_fixture_with_zdr_profile() -> std::io::Result<()> {
guardian_policy_config: None,
include_permissions_instructions: true,
include_apps_instructions: true,
include_collaboration_mode_instructions: true,
include_skill_instructions: true,
include_environment_context: true,
compact_prompt: None,
@@ -8123,6 +8126,7 @@ async fn test_precedence_fixture_with_gpt5_profile() -> std::io::Result<()> {
guardian_policy_config: None,
include_permissions_instructions: true,
include_apps_instructions: true,
include_collaboration_mode_instructions: true,
include_skill_instructions: true,
include_environment_context: true,
compact_prompt: None,
@@ -9360,6 +9364,7 @@ async fn prompt_instruction_blocks_can_be_disabled_from_config_and_profiles() ->
codex_home.path().join(CONFIG_TOML_FILE),
r#"include_permissions_instructions = false
include_apps_instructions = false
include_collaboration_mode_instructions = false
include_environment_context = false
profile = "chatty"
@@ -9368,6 +9373,7 @@ include_instructions = false
[profiles.chatty]
include_permissions_instructions = true
include_collaboration_mode_instructions = true
include_environment_context = true
"#,
)?;
@@ -9380,6 +9386,7 @@ include_environment_context = true
assert!(config.include_permissions_instructions);
assert!(!config.include_apps_instructions);
assert!(config.include_collaboration_mode_instructions);
assert!(!config.include_skill_instructions);
assert!(config.include_environment_context);
Ok(())
codex-rs/core/src/config/mod.rs
+8
View File
@@ -532,6 +532,9 @@ pub struct Config {
/// Whether to inject the `<apps_instructions>` developer block.
pub include_apps_instructions: bool,
/// Whether to inject the `<collaboration_mode>` developer block.
pub include_collaboration_mode_instructions: bool,
/// Whether to inject the `<skills_instructions>` developer block.
pub include_skill_instructions: bool,
@@ -2877,6 +2880,10 @@ impl Config {
.include_apps_instructions
.or(cfg.include_apps_instructions)
.unwrap_or(true);
let include_collaboration_mode_instructions = config_profile
.include_collaboration_mode_instructions
.or(cfg.include_collaboration_mode_instructions)
.unwrap_or(true);
let include_skill_instructions = cfg
.skills
.as_ref()
@@ -3093,6 +3100,7 @@ impl Config {
commit_attribution,
include_permissions_instructions,
include_apps_instructions,
include_collaboration_mode_instructions,
include_skill_instructions,
include_environment_context,
// The config.toml omits "_mode" because it's a config file. However, "_mode"
codex-rs/core/src/context_manager/updates.rs
+4
View File
@@ -80,6 +80,10 @@ fn build_collaboration_mode_update_item(
previous: Option<&TurnContextItem>,
next: &TurnContext,
) -> Option<String> {
if !next.config.include_collaboration_mode_instructions {
return None;
}
let prev = previous?;
if prev.collaboration_mode.as_ref() != Some(&next.collaboration_mode) {
// If the next mode has empty developer instructions, this returns None and we emit no
codex-rs/core/src/session/config_lock.rs
+2
View File
@@ -133,6 +133,8 @@ fn save_config_resolved_fields(
lock_config.model_verbosity = config.model_verbosity;
lock_config.include_permissions_instructions = Some(config.include_permissions_instructions);
lock_config.include_apps_instructions = Some(config.include_apps_instructions);
lock_config.include_collaboration_mode_instructions =
Some(config.include_collaboration_mode_instructions);
lock_config.include_environment_context = Some(config.include_environment_context);
lock_config.background_terminal_max_timeout = Some(config.background_terminal_max_timeout);
codex-rs/core/src/session/mod.rs
+3 -3
View File
@@ -358,7 +358,6 @@ use codex_tools::ToolEnvironmentMode;
use codex_tools::ToolsConfig;
use codex_tools::ToolsConfigParams;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_readiness::ReadinessFlag;
#[cfg(test)]
use codex_utils_stream_parser::ProposedPlanSegment;
@@ -2623,8 +2622,9 @@ impl Session {
developer_sections.push(memory_prompt);
}
// Add developer instructions from collaboration_mode if they exist and are non-empty
if let Some(collab_instructions) =
CollaborationModeInstructions::from_collaboration_mode(&collaboration_mode)
if turn_context.config.include_collaboration_mode_instructions
&& let Some(collab_instructions) =
CollaborationModeInstructions::from_collaboration_mode(&collaboration_mode)
{
developer_sections.push(collab_instructions.render());
}
codex-rs/core/src/session/review.rs
-1
View File
@@ -147,7 +147,6 @@ pub(super) async fn spawn_review_thread(
final_output_json_schema: None,
codex_self_exe: parent_turn_context.codex_self_exe.clone(),
codex_linux_sandbox_exe: parent_turn_context.codex_linux_sandbox_exe.clone(),
tool_call_gate: Arc::new(ReadinessFlag::new()),
dynamic_tools: parent_turn_context.dynamic_tools.clone(),
truncation_policy: model_info.truncation_policy.into(),
turn_metadata_state,
codex-rs/core/src/session/tests.rs
-4
View File
@@ -2843,10 +2843,6 @@ async fn turn_context_with_model_updates_model_fields() {
updated.truncation_policy,
expected_model_info.truncation_policy.into()
);
assert!(!Arc::ptr_eq(
&updated.tool_call_gate,
&turn_context.tool_call_gate
));
}
#[test]
codex-rs/core/src/session/turn_context.rs
-3
View File
@@ -90,7 +90,6 @@ pub struct TurnContext {
pub(crate) final_output_json_schema: Option<Value>,
pub(crate) codex_self_exe: Option<PathBuf>,
pub(crate) codex_linux_sandbox_exe: Option<PathBuf>,
pub(crate) tool_call_gate: Arc<ReadinessFlag>,
pub(crate) truncation_policy: TruncationPolicy,
pub(crate) dynamic_tools: Vec<DynamicToolSpec>,
pub(crate) turn_metadata_state: Arc<TurnMetadataState>,
@@ -277,7 +276,6 @@ impl TurnContext {
final_output_json_schema: self.final_output_json_schema.clone(),
codex_self_exe: self.codex_self_exe.clone(),
codex_linux_sandbox_exe: self.codex_linux_sandbox_exe.clone(),
tool_call_gate: Arc::new(ReadinessFlag::new()),
truncation_policy,
dynamic_tools: self.dynamic_tools.clone(),
turn_metadata_state: self.turn_metadata_state.clone(),
@@ -585,7 +583,6 @@ impl Session {
final_output_json_schema: None,
codex_self_exe: per_turn_config.codex_self_exe.clone(),
codex_linux_sandbox_exe: per_turn_config.codex_linux_sandbox_exe.clone(),
tool_call_gate: Arc::new(ReadinessFlag::new()),
truncation_policy: model_info.truncation_policy.into(),
dynamic_tools: session_configuration.dynamic_tools.clone(),
turn_metadata_state,
codex-rs/core/src/tools/code_mode/mod.rs
+53 -32
View File
@@ -8,6 +8,7 @@ use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeToolKind;
use codex_code_mode::CodeModeTurnHost;
use codex_code_mode::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem;
@@ -33,7 +34,6 @@ use crate::tools::router::extension_tool_bundles;
use crate::unified_exec::resolve_max_tokens;
use codex_features::Feature;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_tools::collect_code_mode_tool_definitions;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy;
@@ -300,6 +300,7 @@ async fn call_nested_tool(
cell_id,
runtime_tool_call_id,
tool_name,
tool_kind,
input,
} = invocation;
if is_exec_tool_name(&tool_name) {
@@ -308,11 +309,10 @@ async fn call_nested_tool(
)));
}
let payload =
match build_nested_tool_payload(tool_runtime.find_spec(&tool_name), &tool_name, input) {
Ok(payload) => payload,
Err(error) => return Err(FunctionCallError::RespondToModel(error)),
};
let payload = match build_nested_tool_payload(tool_kind, &tool_name, input) {
Ok(payload) => payload,
Err(error) => return Err(FunctionCallError::RespondToModel(error)),
};
let call = ToolCall {
tool_name,
@@ -332,36 +332,14 @@ async fn call_nested_tool(
Ok(result.code_mode_result())
}
fn tool_kind_for_spec(spec: &ToolSpec) -> codex_code_mode::CodeModeToolKind {
if matches!(spec, ToolSpec::Freeform(_)) {
codex_code_mode::CodeModeToolKind::Freeform
} else {
codex_code_mode::CodeModeToolKind::Function
}
}
fn tool_kind_for_name(
spec: Option<ToolSpec>,
tool_name: &ToolName,
) -> Result<codex_code_mode::CodeModeToolKind, String> {
spec.as_ref()
.map(tool_kind_for_spec)
.ok_or_else(|| format!("tool `{tool_name}` is not enabled in {PUBLIC_TOOL_NAME}"))
}
fn build_nested_tool_payload(
spec: Option<ToolSpec>,
tool_kind: CodeModeToolKind,
tool_name: &ToolName,
input: Option<JsonValue>,
) -> Result<ToolPayload, String> {
let actual_kind = tool_kind_for_name(spec, tool_name)?;
match actual_kind {
codex_code_mode::CodeModeToolKind::Function => {
build_function_tool_payload(tool_name, input)
}
codex_code_mode::CodeModeToolKind::Freeform => {
build_freeform_tool_payload(tool_name, input)
}
match tool_kind {
CodeModeToolKind::Function => build_function_tool_payload(tool_name, input),
CodeModeToolKind::Freeform => build_freeform_tool_payload(tool_name, input),
}
}
@@ -396,3 +374,46 @@ fn build_freeform_tool_payload(
_ => Err(format!("tool `{tool_name}` expects a string input")),
}
}
#[cfg(test)]
mod tests {
use super::build_nested_tool_payload;
use crate::tools::context::ToolPayload;
use codex_code_mode::CodeModeToolKind;
use codex_tools::ToolName;
use serde_json::json;
#[test]
fn build_nested_tool_payload_uses_function_kind() {
let payload = build_nested_tool_payload(
CodeModeToolKind::Function,
&ToolName::plain("example"),
Some(json!({ "value": 1 })),
)
.expect("function payload should serialize");
match payload {
ToolPayload::Function { arguments } => {
assert_eq!(arguments, r#"{"value":1}"#.to_string());
}
other => panic!("expected function payload, got {other:?}"),
}
}
#[test]
fn build_nested_tool_payload_uses_freeform_kind() {
let payload = build_nested_tool_payload(
CodeModeToolKind::Freeform,
&ToolName::plain("example"),
Some(json!("hello")),
)
.expect("freeform payload should preserve string input");
match payload {
ToolPayload::Custom { input } => {
assert_eq!(input, "hello".to_string());
}
other => panic!("expected freeform payload, got {other:?}"),
}
}
}
codex-rs/core/src/tools/handlers/apply_patch.rs
-4
View File
@@ -311,10 +311,6 @@ impl ToolHandler for ApplyPatchHandler {
matches!(payload, ToolPayload::Custom { .. })
}
async fn is_mutating(&self, _invocation: &ToolInvocation) -> bool {
true
}
fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>> {
Some(Box::<ApplyPatchArgumentDiffConsumer>::default())
}
codex-rs/core/src/tools/handlers/dynamic.rs
-4
View File
@@ -35,10 +35,6 @@ impl ToolHandler for DynamicToolHandler {
self.tool_name.clone()
}
async fn is_mutating(&self, _invocation: &ToolInvocation) -> bool {
true
}
async fn handle(&self, invocation: ToolInvocation) -> Result<Self::Output, FunctionCallError> {
let ToolInvocation {
session,
codex-rs/core/src/tools/handlers/extension_tools.rs
+1 -6
View File
@@ -83,10 +83,6 @@ impl ToolHandler for BundledToolHandler {
self.arguments_from_payload(payload).is_some()
}
async fn is_mutating(&self, _invocation: &ToolInvocation) -> bool {
true
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
let arguments = self.arguments_from_payload(&invocation.payload)?;
Some(PreToolUsePayload {
@@ -189,7 +185,7 @@ mod tests {
}
#[tokio::test]
async fn exposes_generic_hook_payloads_and_is_conservatively_mutating() {
async fn exposes_generic_hook_payloads() {
let bundle = codex_tool_api::ToolBundle::new(
codex_tool_api::FunctionToolSpec {
name: "extension_echo".to_string(),
@@ -225,7 +221,6 @@ mod tests {
value: json!({ "ok": true }),
};
assert!(ToolHandler::is_mutating(&handler, &invocation).await);
assert_eq!(
ToolHandler::pre_tool_use_payload(&handler, &invocation),
Some(PreToolUsePayload {
codex-rs/core/src/tools/handlers/shell/container_exec.rs
-11
View File
@@ -1,5 +1,4 @@
use codex_protocol::models::ShellToolCallParams;
use codex_shell_command::is_safe_command::is_known_safe_command;
use codex_tools::ToolName;
use crate::function_tool::FunctionCallError;
@@ -33,16 +32,6 @@ impl ToolHandler for ContainerExecHandler {
matches!(payload, ToolPayload::Function { .. })
}
async fn is_mutating(&self, invocation: &ToolInvocation) -> bool {
let ToolPayload::Function { arguments } = &invocation.payload else {
return true;
};
serde_json::from_str::<ShellToolCallParams>(arguments)
.map(|params| !is_known_safe_command(&params.command))
.unwrap_or(true)
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
shell_function_pre_tool_use_payload(invocation)
}
codex-rs/core/src/tools/handlers/shell/local_shell.rs
+2 -11
View File
@@ -1,4 +1,3 @@
use codex_shell_command::is_safe_command::is_known_safe_command;
use codex_tools::ToolName;
use crate::function_tool::FunctionCallError;
@@ -42,20 +41,12 @@ impl ToolHandler for LocalShellHandler {
self.include_spec.then(create_local_shell_tool)
}
fn supports_parallel_tool_calls(&self) -> bool {
self.include_spec
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(payload, ToolPayload::LocalShell { .. })
}
async fn is_mutating(&self, invocation: &ToolInvocation) -> bool {
let ToolPayload::LocalShell { params } = &invocation.payload else {
return true;
};
!is_known_safe_command(&params.command)
fn supports_parallel_tool_calls(&self) -> bool {
self.include_spec
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
codex-rs/core/src/tools/handlers/shell/shell_command.rs
+2 -24
View File
@@ -1,6 +1,5 @@
use codex_protocol::ThreadId;
use codex_protocol::models::ShellCommandToolCallParams;
use codex_shell_command::is_safe_command::is_known_safe_command;
use codex_tools::ShellCommandBackendConfig;
use codex_tools::ToolName;
@@ -141,33 +140,12 @@ impl ToolHandler for ShellCommandHandler {
})
}
fn supports_parallel_tool_calls(&self) -> bool {
self.options.is_some()
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(payload, ToolPayload::Function { .. })
}
async fn is_mutating(&self, invocation: &ToolInvocation) -> bool {
let ToolPayload::Function { arguments } = &invocation.payload else {
return true;
};
serde_json::from_str::<ShellCommandToolCallParams>(arguments)
.map(|params| {
let use_login_shell = match Self::resolve_use_login_shell(
params.login,
invocation.turn.tools_config.allow_login_shell,
) {
Ok(use_login_shell) => use_login_shell,
Err(_) => return true,
};
let shell = invocation.session.user_shell();
let command = Self::base_command(shell.as_ref(), &params.command, use_login_shell);
!is_known_safe_command(&command)
})
.unwrap_or(true)
fn supports_parallel_tool_calls(&self) -> bool {
self.options.is_some()
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
codex-rs/core/src/tools/handlers/shell/shell_handler.rs
+2 -13
View File
@@ -1,6 +1,5 @@
use codex_protocol::ThreadId;
use codex_protocol::models::ShellToolCallParams;
use codex_shell_command::is_safe_command::is_known_safe_command;
use codex_tools::ToolName;
use crate::exec::ExecCapturePolicy;
@@ -74,22 +73,12 @@ impl ToolHandler for ShellHandler {
self.options.map(create_shell_tool)
}
fn supports_parallel_tool_calls(&self) -> bool {
self.options.is_some()
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(payload, ToolPayload::Function { .. })
}
async fn is_mutating(&self, invocation: &ToolInvocation) -> bool {
let ToolPayload::Function { arguments } = &invocation.payload else {
return true;
};
serde_json::from_str::<ShellToolCallParams>(arguments)
.map(|params| !is_known_safe_command(&params.command))
.unwrap_or(true)
fn supports_parallel_tool_calls(&self) -> bool {
self.options.is_some()
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
+2 -27
View File
@@ -26,7 +26,6 @@ use crate::unified_exec::generate_chunk_id;
use codex_features::Feature;
use codex_otel::SessionTelemetry;
use codex_otel::TOOL_CALL_UNIFIED_EXEC_METRIC;
use codex_shell_command::is_safe_command::is_known_safe_command;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_utils_output_truncation::approx_token_count;
@@ -85,36 +84,12 @@ impl ToolHandler for ExecCommandHandler {
))
}
fn supports_parallel_tool_calls(&self) -> bool {
true
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(payload, ToolPayload::Function { .. })
}
async fn is_mutating(&self, invocation: &ToolInvocation) -> bool {
let ToolPayload::Function { arguments } = &invocation.payload else {
tracing::error!(
"This should never happen, invocation payload is wrong: {:?}",
invocation.payload
);
return true;
};
let Ok(params) = parse_arguments::<ExecCommandArgs>(arguments) else {
return true;
};
let command = match get_command(
&params,
invocation.session.user_shell(),
&invocation.turn.tools_config.unified_exec_shell_mode,
invocation.turn.tools_config.allow_login_shell,
) {
Ok(command) => command,
Err(_) => return true,
};
!is_known_safe_command(&command)
fn supports_parallel_tool_calls(&self) -> bool {
true
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs
-4
View File
@@ -45,10 +45,6 @@ impl ToolHandler for WriteStdinHandler {
matches!(payload, ToolPayload::Function { .. })
}
async fn is_mutating(&self, _invocation: &ToolInvocation) -> bool {
true
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
codex-rs/core/src/tools/parallel.rs
-5
View File
@@ -22,7 +22,6 @@ use crate::tools::router::ToolCallSource;
use crate::tools::router::ToolRouter;
use codex_protocol::error::CodexErr;
use codex_protocol::models::ResponseInputItem;
use codex_tools::ToolSpec;
#[derive(Clone)]
pub(crate) struct ToolCallRuntime {
@@ -49,10 +48,6 @@ impl ToolCallRuntime {
}
}
pub(crate) fn find_spec(&self, tool_name: &codex_tools::ToolName) -> Option<ToolSpec> {
self.router.find_spec(tool_name)
}
pub(crate) fn create_diff_consumer(
&self,
tool_name: &codex_tools::ToolName,
codex-rs/core/src/tools/registry.rs
+2 -25
View File
@@ -27,7 +27,6 @@ use codex_protocol::protocol::EventMsg;
use codex_tool_api::ToolBundle as ExtensionToolBundle;
use codex_tools::ToolName;
use codex_tools::ToolSpec;
use codex_utils_readiness::Readiness;
use futures::future::BoxFuture;
use serde_json::Value;
use tracing::warn;
@@ -62,17 +61,6 @@ pub trait ToolHandler: Send + Sync {
async { Vec::new() }
}
/// Returns `true` if the [ToolInvocation] *might* mutate the environment of the
/// user (through file system, OS operations, ...).
/// This function must remains defensive and return `true` if a doubt exist on the
/// exact effect of a ToolInvocation.
fn is_mutating(
&self,
_invocation: &ToolInvocation,
) -> impl std::future::Future<Output = bool> + Send {
async { false }
}
fn post_tool_use_payload(
&self,
_invocation: &ToolInvocation,
@@ -185,8 +173,6 @@ trait AnyToolHandler: Send + Sync {
fn matches_kind(&self, payload: &ToolPayload) -> bool;
fn is_mutating<'a>(&'a self, invocation: &'a ToolInvocation) -> BoxFuture<'a, bool>;
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload>;
fn with_updated_hook_input(
@@ -219,10 +205,6 @@ where
ToolHandler::matches_kind(self, payload)
}
fn is_mutating<'a>(&'a self, invocation: &'a ToolInvocation) -> BoxFuture<'a, bool> {
Box::pin(ToolHandler::is_mutating(self, invocation))
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
ToolHandler::pre_tool_use_payload(self, invocation)
}
@@ -305,7 +287,8 @@ impl ToolRegistry {
}
pub(crate) fn supports_parallel_tool_calls(&self, name: &ToolName) -> Option<bool> {
Some(self.handler(name)?.supports_parallel_tool_calls())
let handler = self.handler(name)?;
Some(handler.supports_parallel_tool_calls())
}
#[expect(
@@ -424,7 +407,6 @@ impl ToolRegistry {
}
}
let is_mutating = handler.is_mutating(&invocation).await;
let response_cell = tokio::sync::Mutex::new(None);
let invocation_for_tool = invocation.clone();
let log_payload = invocation.payload.log_payload();
@@ -440,11 +422,6 @@ impl ToolRegistry {
let handler = handler.clone();
let response_cell = &response_cell;
async move {
if is_mutating {
tracing::trace!("waiting for tool gate");
invocation_for_tool.turn.tool_call_gate.wait_ready().await;
tracing::trace!("tool gate released");
}
match handler.handle_any(invocation_for_tool).await {
Ok(result) => {
let preview = result.result.log_preview();
codex-rs/core/src/tools/router.rs
-25
View File
@@ -103,31 +103,6 @@ impl ToolRouter {
self.model_visible_specs.clone()
}
pub fn find_spec(&self, tool_name: &ToolName) -> Option<ToolSpec> {
self.specs.iter().find_map(|spec| match spec {
ToolSpec::Function(tool)
if tool_name.namespace.is_none() && tool.name == tool_name.name =>
{
Some(spec.clone())
}
ToolSpec::Freeform(tool)
if tool_name.namespace.is_none() && tool.name == tool_name.name =>
{
Some(spec.clone())
}
ToolSpec::Namespace(namespace) => namespace.tools.iter().find_map(|tool| match tool {
ResponsesApiNamespaceTool::Function(tool)
if tool_name.namespace.as_deref() == Some(namespace.name.as_str())
&& tool.name == tool_name.name =>
{
Some(ToolSpec::Function(tool.clone()))
}
_ => None,
}),
_ => None,
})
}
pub(crate) fn create_diff_consumer(
&self,
tool_name: &ToolName,
codex-rs/core/src/tools/router_tests.rs
+3 -8
View File
@@ -273,11 +273,6 @@ async fn model_visible_specs_filter_deferred_dynamic_tools() -> anyhow::Result<(
},
);
assert!(
router
.find_spec(&ToolName::namespaced("codex_app", hidden_tool))
.is_some()
);
assert_eq!(
namespace_function_names(&router.specs(), "codex_app"),
vec![hidden_tool.to_string(), visible_tool.to_string()]
@@ -341,8 +336,9 @@ async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow::
assert!(
router
.find_spec(&ToolName::plain("extension_echo"))
.is_some(),
.specs()
.iter()
.any(|spec| spec.name() == "extension_echo"),
"expected extension-provided tool spec to be registered"
);
assert!(
@@ -361,7 +357,6 @@ async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow::
call_id: "call-extension".to_string(),
})?
.expect("function_call should produce a tool call");
let result = router
.dispatch_tool_call_with_code_mode_result(
Arc::new(session),
codex-rs/core/tests/suite/collaboration_instructions.rs
+54
View File
@@ -210,6 +210,60 @@ async fn collaboration_instructions_added_on_user_turn() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn collaboration_instructions_omitted_when_disabled() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let req = mount_sse_once(
&server,
sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
)
.await;
let mut builder = test_codex().with_config(|config| {
config.include_collaboration_mode_instructions = false;
});
let test = builder.build(&server).await?;
let collaboration_mode = collab_mode_with_instructions(Some("turn instructions"));
test.codex
.submit(Op::UserTurn {
environments: None,
items: vec![UserInput::Text {
text: "hello".into(),
text_elements: Vec::new(),
}],
cwd: test.config.cwd.to_path_buf(),
approval_policy: test.config.permissions.approval_policy.value(),
approvals_reviewer: None,
sandbox_policy: test.config.legacy_sandbox_policy(),
permission_profile: None,
model: test.session_configured.model.clone(),
effort: None,
summary: Some(
test.config
.model_reasoning_summary
.unwrap_or(codex_protocol::config_types::ReasoningSummary::Auto),
),
service_tier: None,
collaboration_mode: Some(collaboration_mode),
final_output_json_schema: None,
personality: None,
})
.await?;
wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let input = req.single_request().input();
let dev_texts = developer_texts(&input);
assert_eq!(
count_messages_containing(&dev_texts, COLLABORATION_MODE_OPEN_TAG),
0
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn override_then_next_turn_uses_updated_collaboration_instructions() -> Result<()> {
skip_if_no_network!(Ok(()));
codex-rs/exec-server/BUILD.bazel
+3
View File
@@ -10,6 +10,9 @@ codex_rust_crate(
# they install process-global test-binary dispatch state, and the remote
# exec-server cases already rely on serialization around the full CLI path.
integration_test_args = ["--test-threads=1"],
integration_compile_data_extra = [
"src/proto/codex.exec_server.relay.v1.rs",
],
extra_binaries = [
"//codex-rs/bwrap:bwrap",
],
codex-rs/exec-server/Cargo.toml
+1
View File
@@ -26,6 +26,7 @@ codex-utils-pty = { workspace = true }
codex-utils-rustls-provider = { workspace = true }
futures = { workspace = true }
reqwest = { workspace = true, features = ["json", "rustls-tls", "stream"] }
prost = "0.14.3"
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
thiserror = { workspace = true }
codex-rs/exec-server/README.md
+59 -1
View File
@@ -30,7 +30,65 @@ It requires a bearer token in `CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN`.
Wire framing:
- websocket: one JSON-RPC message per websocket text frame
- local websocket: one JSON-RPC message per websocket frame
- remote websocket: binary protobuf relay frames carrying JSON-RPC payloads
## Remote Relay Message Format
In remote mode, the harness and executor communicate through rendezvous using
`codex.exec_server.relay.v1.RelayMessageFrame`; the checked-in schema is in
`src/proto/codex.exec_server.relay.v1.proto`. The relay frame carries stream
identity plus endpoint-owned reliability metadata:
```text
version
stream_id
body // data | ack_frame | resume | reset | heartbeat
ack // highest contiguous peer segment seq received
ack_bits // bitset for peer segment seqs after ack
seq // data only: segment sequence number
segment_index // data only: 0-based index within message
segment_count // data only: number of segments in message
payload // data only: JSON-RPC message bytes or segment bytes
next_seq // resume only: next sender seq
reason // reset only: reset reason
```
`stream_id` identifies one virtual harness/executor JSON-RPC session on the
executor websocket. The harness generates a UUIDv4 `stream_id`; the executor
demuxes frames by `stream_id` and runs an independent `ConnectionProcessor` per
stream.
Use segment-level sequence numbers for reliability:
```text
seq = 0, 1, 2, 3, ...
```
Use contiguous segment sequence ranges to identify and stitch a segmented
application message:
```text
message_start_seq = seq - segment_index
segment_index = 0
segment_count = 1
```
`message_start_seq` is derived by the receiver, not sent on the wire. For
unsplit messages, `message_start_seq == seq`, `segment_index == 0`, and
`segment_count == 1`.
Use cumulative `ack` plus fixed-size `ack_bits` instead of variable ack ranges:
```text
ack = highest contiguous received segment seq
bit i in ack_bits acknowledges seq = ack + 1 + i
```
Send `ack` and `ack_bits` redundantly on every outbound frame. Acks are not
themselves acked. Acks, retries, duplicate suppression, segmentation, and
reassembly are endpoint responsibilities; rendezvous only routes relay frames
by `stream_id`.
## Lifecycle
codex-rs/exec-server/src/client_transport.rs
+18 -8
View File
@@ -15,6 +15,7 @@ use crate::client_api::RemoteExecServerConnectArgs;
use crate::client_api::StdioExecServerCommand;
use crate::client_api::StdioExecServerConnectArgs;
use crate::connection::JsonRpcConnection;
use crate::relay::harness_connection_from_websocket;
const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment";
@@ -69,14 +70,13 @@ impl ExecServerClient {
source,
})?;
Self::connect(
JsonRpcConnection::from_websocket(
stream,
format!("exec-server websocket {websocket_url}"),
),
args.into(),
)
.await
let connection_label = format!("exec-server websocket {websocket_url}");
let connection = if is_rendezvous_harness_url(&websocket_url) {
harness_connection_from_websocket(stream, connection_label)
} else {
JsonRpcConnection::from_websocket(stream, connection_label)
};
Self::connect(connection, args.into()).await
}
pub(crate) async fn connect_stdio_command(
@@ -120,6 +120,16 @@ impl ExecServerClient {
}
}
fn is_rendezvous_harness_url(websocket_url: &str) -> bool {
let Some((_path, query)) = websocket_url.split_once('?') else {
return false;
};
query
.split('&')
.filter_map(|pair| pair.split_once('='))
.any(|(key, value)| key == "role" && value == "harness")
}
fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command {
let mut command = Command::new(&stdio_command.program);
command.args(&stdio_command.args);
codex-rs/exec-server/src/lib.rs
+2
View File
@@ -13,6 +13,8 @@ mod local_process;
mod process;
mod process_id;
mod protocol;
mod relay;
mod relay_proto;
mod remote;
mod remote_file_system;
mod remote_process;
codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto
+37
View File
@@ -0,0 +1,37 @@
syntax = "proto3";
package codex.exec_server.relay.v1;
message RelayMessageFrame {
uint32 version = 1;
string stream_id = 2;
uint32 ack = 3;
uint32 ack_bits = 4;
oneof body {
RelayData data = 5;
RelayAck ack_frame = 6;
RelayResume resume = 7;
RelayReset reset = 8;
RelayHeartbeat heartbeat = 9;
}
}
message RelayData {
uint32 seq = 1;
uint32 segment_index = 2;
uint32 segment_count = 3;
bytes payload = 4;
}
message RelayAck {}
message RelayResume {
uint32 next_seq = 1;
}
message RelayReset {
string reason = 1;
}
message RelayHeartbeat {}
codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs
+54
View File
@@ -0,0 +1,54 @@
// This file is @generated by prost-build.
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct RelayMessageFrame {
#[prost(uint32, tag = "1")]
pub version: u32,
#[prost(string, tag = "2")]
pub stream_id: ::prost::alloc::string::String,
#[prost(uint32, tag = "3")]
pub ack: u32,
#[prost(uint32, tag = "4")]
pub ack_bits: u32,
#[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9")]
pub body: ::core::option::Option<relay_message_frame::Body>,
}
pub mod relay_message_frame {
#[derive(Clone, PartialEq, ::prost::Oneof)]
pub enum Body {
#[prost(message, tag = "5")]
Data(super::RelayData),
#[prost(message, tag = "6")]
AckFrame(super::RelayAck),
#[prost(message, tag = "7")]
Resume(super::RelayResume),
#[prost(message, tag = "8")]
Reset(super::RelayReset),
#[prost(message, tag = "9")]
Heartbeat(super::RelayHeartbeat),
}
}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct RelayData {
#[prost(uint32, tag = "1")]
pub seq: u32,
#[prost(uint32, tag = "2")]
pub segment_index: u32,
#[prost(uint32, tag = "3")]
pub segment_count: u32,
#[prost(bytes = "vec", tag = "4")]
pub payload: ::prost::alloc::vec::Vec<u8>,
}
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
pub struct RelayAck {}
#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
pub struct RelayResume {
#[prost(uint32, tag = "1")]
pub next_seq: u32,
}
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
pub struct RelayReset {
#[prost(string, tag = "1")]
pub reason: ::prost::alloc::string::String,
}
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
pub struct RelayHeartbeat {}
codex-rs/exec-server/src/relay.rs
+455
View File
@@ -0,0 +1,455 @@
use std::collections::HashMap;
use codex_app_server_protocol::JSONRPCMessage;
use futures::SinkExt;
use futures::StreamExt;
use prost::Message as ProstMessage;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::Message;
use tracing::debug;
use tracing::warn;
use uuid::Uuid;
use crate::ExecServerError;
use crate::connection::CHANNEL_CAPACITY;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::connection::JsonRpcTransport;
use crate::relay_proto::RelayData;
use crate::relay_proto::RelayMessageFrame;
use crate::relay_proto::RelayResume;
use crate::relay_proto::relay_message_frame;
use crate::server::ConnectionProcessor;
const RELAY_MESSAGE_FRAME_VERSION: u32 = 1;
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
enum RelayFrameBodyKind {
Data,
Ack,
Resume,
Reset,
Heartbeat,
}
impl RelayMessageFrame {
fn data(stream_id: String, seq: u32, payload: Vec<u8>) -> Self {
Self {
version: RELAY_MESSAGE_FRAME_VERSION,
stream_id,
ack: 0,
ack_bits: 0,
body: Some(relay_message_frame::Body::Data(RelayData {
seq,
segment_index: 0,
segment_count: 1,
payload,
})),
}
}
fn resume(stream_id: String) -> Self {
Self {
version: RELAY_MESSAGE_FRAME_VERSION,
stream_id,
ack: 0,
ack_bits: 0,
body: Some(relay_message_frame::Body::Resume(RelayResume {
next_seq: 0,
})),
}
}
fn validate(&self) -> Result<RelayFrameBodyKind, ExecServerError> {
if self.version != RELAY_MESSAGE_FRAME_VERSION {
return Err(ExecServerError::Protocol(format!(
"unsupported relay message frame version {}",
self.version
)));
}
if self.stream_id.trim().is_empty() {
return Err(ExecServerError::Protocol(
"relay message frame is missing stream_id".to_string(),
));
}
match self.body.as_ref() {
Some(relay_message_frame::Body::Data(data)) => {
if data.segment_index != 0 || data.segment_count != 1 || data.payload.is_empty() {
return Err(ExecServerError::Protocol(
"relay data message frame is missing required fields".to_string(),
));
}
Ok(RelayFrameBodyKind::Data)
}
Some(relay_message_frame::Body::AckFrame(_)) => Ok(RelayFrameBodyKind::Ack),
Some(relay_message_frame::Body::Resume(_)) => Ok(RelayFrameBodyKind::Resume),
Some(relay_message_frame::Body::Reset(reset)) => {
if reset.reason.is_empty() {
return Err(ExecServerError::Protocol(
"relay reset message frame is missing reason".to_string(),
));
}
Ok(RelayFrameBodyKind::Reset)
}
Some(relay_message_frame::Body::Heartbeat(_)) => Ok(RelayFrameBodyKind::Heartbeat),
None => Err(ExecServerError::Protocol(
"relay message frame is missing body".to_string(),
)),
}
}
fn into_jsonrpc_message(self) -> Result<JSONRPCMessage, ExecServerError> {
let kind = self.validate()?;
if kind != RelayFrameBodyKind::Data {
return Err(ExecServerError::Protocol(
"expected relay data message frame".to_string(),
));
}
let payload = match self.body {
Some(relay_message_frame::Body::Data(data)) => data.payload,
_ => Vec::new(),
};
serde_json::from_slice(&payload).map_err(ExecServerError::Json)
}
fn into_reset_reason(self) -> Option<String> {
match self.body {
Some(relay_message_frame::Body::Reset(reset)) if !reset.reason.is_empty() => {
Some(reset.reason)
}
_ => None,
}
}
}
fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec<u8> {
frame.encode_to_vec()
}
fn decode_relay_message_frame(payload: &[u8]) -> Result<RelayMessageFrame, ExecServerError> {
RelayMessageFrame::decode(payload)
.map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}")))
}
fn jsonrpc_payload(message: &JSONRPCMessage) -> Result<Vec<u8>, ExecServerError> {
serde_json::to_vec(message).map_err(ExecServerError::Json)
}
pub(crate) fn harness_connection_from_websocket<S>(
stream: WebSocketStream<S>,
connection_label: String,
) -> JsonRpcConnection
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
let stream_id = Uuid::new_v4().to_string();
let (mut websocket_writer, mut websocket_reader) = stream.split();
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (disconnected_tx, disconnected_rx) = watch::channel(false);
let reader_label = connection_label;
let reader_stream_id = stream_id.clone();
let incoming_tx_for_reader = incoming_tx;
let disconnected_tx_for_reader = disconnected_tx.clone();
let reader_task = tokio::spawn(async move {
loop {
match websocket_reader.next().await {
Some(Ok(Message::Binary(payload))) => {
let frame = match decode_relay_message_frame(payload.as_ref()) {
Ok(frame) => frame,
Err(err) => {
let _ = incoming_tx_for_reader
.send(JsonRpcConnectionEvent::MalformedMessage {
reason: format!(
"failed to parse relay message frame from {reader_label}: {err}"
),
})
.await;
continue;
}
};
if frame.stream_id != reader_stream_id {
continue;
}
let kind = match frame.validate() {
Ok(kind) => kind,
Err(err) => {
let _ = incoming_tx_for_reader
.send(JsonRpcConnectionEvent::MalformedMessage {
reason: err.to_string(),
})
.await;
continue;
}
};
match kind {
RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() {
Ok(message) => {
if incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Message(message))
.await
.is_err()
{
break;
}
}
Err(err) => {
let _ = incoming_tx_for_reader
.send(JsonRpcConnectionEvent::MalformedMessage {
reason: err.to_string(),
})
.await;
}
},
RelayFrameBodyKind::Reset => {
let _ = disconnected_tx_for_reader.send(true);
let _ = incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Disconnected {
reason: frame.into_reset_reason(),
})
.await;
break;
}
RelayFrameBodyKind::Ack
| RelayFrameBodyKind::Resume
| RelayFrameBodyKind::Heartbeat => {}
}
}
Some(Ok(Message::Close(_))) | None => {
let _ = disconnected_tx_for_reader.send(true);
let _ = incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Disconnected { reason: None })
.await;
break;
}
Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {}
Some(Ok(Message::Text(_))) => {
let _ = incoming_tx_for_reader
.send(JsonRpcConnectionEvent::MalformedMessage {
reason: "relay exec-server transport expects binary protobuf frames"
.to_string(),
})
.await;
}
Some(Err(err)) => {
let _ = disconnected_tx_for_reader.send(true);
let _ = incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Disconnected {
reason: Some(format!(
"failed to read relay websocket frame from {reader_label}: {err}"
)),
})
.await;
break;
}
}
}
});
let writer_task = tokio::spawn(async move {
let resume = RelayMessageFrame::resume(stream_id.clone());
if websocket_writer
.send(Message::Binary(encode_relay_message_frame(&resume).into()))
.await
.is_err()
{
let _ = disconnected_tx.send(true);
return;
}
let mut next_seq = 0u32;
while let Some(message) = outgoing_rx.recv().await {
let payload = match jsonrpc_payload(&message) {
Ok(payload) => payload,
Err(err) => {
warn!("failed to serialize JSON-RPC payload for relay transport: {err}");
break;
}
};
let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload);
next_seq = next_seq.wrapping_add(1);
if websocket_writer
.send(Message::Binary(encode_relay_message_frame(&frame).into()))
.await
.is_err()
{
let _ = disconnected_tx.send(true);
break;
}
}
});
JsonRpcConnection {
outgoing_tx,
incoming_rx,
disconnected_rx,
task_handles: vec![reader_task, writer_task],
transport: JsonRpcTransport::Plain,
}
}
pub(crate) async fn run_multiplexed_executor<S>(
stream: WebSocketStream<S>,
processor: ConnectionProcessor,
) where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
let (mut websocket_writer, mut websocket_reader) = stream.split();
let (physical_outgoing_tx, mut physical_outgoing_rx) =
mpsc::channel::<Vec<u8>>(CHANNEL_CAPACITY);
let writer_task = tokio::spawn(async move {
while let Some(encoded) = physical_outgoing_rx.recv().await {
if websocket_writer
.send(Message::Binary(encoded.into()))
.await
.is_err()
{
break;
}
}
});
let mut streams: HashMap<String, VirtualStream> = HashMap::new();
loop {
let frame = match websocket_reader.next().await {
Some(Ok(Message::Binary(payload))) => {
match decode_relay_message_frame(payload.as_ref()) {
Ok(frame) => frame,
Err(err) => {
warn!("dropping malformed relay message frame from harness: {err}");
continue;
}
}
}
Some(Ok(Message::Close(_))) | None => break,
Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue,
Some(Ok(Message::Text(_))) => {
warn!("dropping non-binary relay message frame from harness");
continue;
}
Some(Err(err)) => {
debug!("multiplexed executor websocket read failed: {err}");
break;
}
};
let kind = match frame.validate() {
Ok(kind) => kind,
Err(err) => {
warn!("dropping invalid relay message frame: {err}");
continue;
}
};
match kind {
RelayFrameBodyKind::Data => {
let stream_id = frame.stream_id.clone();
let message = match frame.into_jsonrpc_message() {
Ok(message) => message,
Err(err) => {
warn!("dropping malformed relay data message frame: {err}");
continue;
}
};
let stream = streams.entry(stream_id.clone()).or_insert_with(|| {
spawn_virtual_stream(
stream_id.clone(),
processor.clone(),
physical_outgoing_tx.clone(),
)
});
if stream
.incoming_tx
.send(JsonRpcConnectionEvent::Message(message))
.await
.is_err()
{
streams.remove(&stream_id);
}
}
RelayFrameBodyKind::Reset => {
if let Some(stream) = streams.remove(&frame.stream_id) {
stream.disconnect(frame.into_reset_reason()).await;
}
}
RelayFrameBodyKind::Ack
| RelayFrameBodyKind::Resume
| RelayFrameBodyKind::Heartbeat => {}
}
}
for (_stream_id, stream) in streams {
stream.disconnect(/*reason*/ None).await;
}
drop(physical_outgoing_tx);
let _ = writer_task.await;
}
struct VirtualStream {
incoming_tx: mpsc::Sender<JsonRpcConnectionEvent>,
disconnected_tx: watch::Sender<bool>,
}
impl VirtualStream {
async fn disconnect(self, reason: Option<String>) {
let _ = self.disconnected_tx.send(true);
let _ = self
.incoming_tx
.send(JsonRpcConnectionEvent::Disconnected { reason })
.await;
}
}
fn spawn_virtual_stream(
stream_id: String,
processor: ConnectionProcessor,
physical_outgoing_tx: mpsc::Sender<Vec<u8>>,
) -> VirtualStream {
let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (disconnected_tx, disconnected_rx) = watch::channel(false);
let writer_stream_id = stream_id;
let writer_task = tokio::spawn(async move {
let mut next_seq = 0u32;
while let Some(message) = json_outgoing_rx.recv().await {
let payload = match jsonrpc_payload(&message) {
Ok(payload) => payload,
Err(err) => {
warn!("failed to serialize virtual stream JSON-RPC payload: {err}");
break;
}
};
let frame = RelayMessageFrame::data(writer_stream_id.clone(), next_seq, payload);
next_seq = next_seq.wrapping_add(1);
if physical_outgoing_tx
.send(encode_relay_message_frame(&frame))
.await
.is_err()
{
break;
}
}
});
let connection = JsonRpcConnection {
outgoing_tx: json_outgoing_tx,
incoming_rx,
disconnected_rx,
task_handles: vec![writer_task],
transport: JsonRpcTransport::Plain,
};
tokio::spawn(async move {
processor.run_connection(connection).await;
});
VirtualStream {
incoming_tx,
disconnected_tx,
}
}
codex-rs/exec-server/src/relay_proto.rs
+7
View File
@@ -0,0 +1,7 @@
#[path = "proto/codex.exec_server.relay.v1.rs"]
mod generated;
pub(crate) use generated::RelayData;
pub(crate) use generated::RelayMessageFrame;
pub(crate) use generated::RelayResume;
pub(crate) use generated::relay_message_frame;
codex-rs/exec-server/src/remote.rs
+3 -8
View File
@@ -11,7 +11,7 @@ use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use crate::ExecServerError;
use crate::ExecServerRuntimePaths;
use crate::connection::JsonRpcConnection;
use crate::relay::run_multiplexed_executor;
use crate::server::ConnectionProcessor;
pub const CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR: &str =
@@ -113,7 +113,7 @@ impl RemoteExecutorConfig {
Self::with_bearer_token(base_url, executor_id, read_remote_bearer_token_from_env()?)
}
fn with_bearer_token(
pub fn with_bearer_token(
base_url: String,
executor_id: String,
bearer_token: String,
@@ -150,12 +150,7 @@ pub async fn run_remote_executor(
match connect_async(response.url.as_str()).await {
Ok((websocket, _)) => {
backoff = Duration::from_secs(1);
processor
.run_connection(JsonRpcConnection::from_websocket(
websocket,
"remote exec-server websocket".to_string(),
))
.await;
run_multiplexed_executor(websocket, processor.clone()).await;
}
Err(err) => {
warn!("failed to connect remote exec-server websocket: {err}");
codex-rs/exec-server/tests/relay.rs
+353
View File
@@ -0,0 +1,353 @@
mod common;
#[path = "../src/proto/codex.exec_server.relay.v1.rs"]
mod relay_proto;
use std::collections::HashMap;
use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use anyhow::bail;
use codex_app_server_protocol::JSONRPCError;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_exec_server::ExecServerRuntimePaths;
use codex_exec_server::InitializeParams;
use codex_exec_server::InitializeResponse;
use codex_exec_server::RemoteExecutorConfig;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use prost::Message as ProstMessage;
use relay_proto::RelayData;
use relay_proto::RelayMessageFrame;
use relay_proto::RelayReset;
use relay_proto::relay_message_frame;
use tokio::net::TcpListener;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use uuid::Uuid;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
const EXECUTOR_ID: &str = "exec-mux-test";
const REGISTRY_TOKEN: &str = "registry-token";
const RELAY_MESSAGE_FRAME_VERSION: u32 = 1;
const TEST_TIMEOUT: Duration = Duration::from_secs(5);
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn multiplexed_remote_executor_routes_independent_virtual_streams() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let rendezvous_url = format!("ws://{}", listener.local_addr()?);
let registry = MockServer::start().await;
Mock::given(method("POST"))
.and(path(format!("/cloud/executor/{EXECUTOR_ID}/register")))
.and(header("authorization", format!("Bearer {REGISTRY_TOKEN}")))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"executor_id": EXECUTOR_ID,
"url": rendezvous_url,
})))
.mount(&registry)
.await;
let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?;
let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?;
let config = RemoteExecutorConfig::with_bearer_token(
registry.uri(),
EXECUTOR_ID.to_string(),
REGISTRY_TOKEN.to_string(),
)?;
let remote_executor = tokio::spawn(codex_exec_server::run_remote_executor(
config,
runtime_paths,
));
let (socket, _peer_addr) = timeout(TEST_TIMEOUT, listener.accept())
.await
.context("remote executor should connect to fake rendezvous")??;
let mut websocket = timeout(TEST_TIMEOUT, accept_async(socket))
.await
.context("fake rendezvous should accept executor websocket")??;
let stream_a = "stream-a";
let stream_b = "stream-b";
send_relay_message(
&mut websocket,
stream_a,
/*seq*/ 0,
initialize_request(/*id*/ 1, "relay-test-a")?,
)
.await?;
send_relay_message(
&mut websocket,
stream_b,
/*seq*/ 0,
initialize_request(/*id*/ 1, "relay-test-b")?,
)
.await?;
let initialize_responses = read_relay_messages_by_stream(&mut websocket, /*count*/ 2).await?;
let session_a =
assert_initialize_response(initialize_responses.get(stream_a), stream_a, /*id*/ 1)?;
let session_b =
assert_initialize_response(initialize_responses.get(stream_b), stream_b, /*id*/ 1)?;
assert_ne!(session_a, session_b);
send_relay_message(
&mut websocket,
stream_a,
/*seq*/ 1,
notification("initialized", serde_json::json!({})),
)
.await?;
send_relay_message(
&mut websocket,
stream_b,
/*seq*/ 1,
notification("initialized", serde_json::json!({})),
)
.await?;
send_relay_message(
&mut websocket,
stream_a,
/*seq*/ 2,
request(/*id*/ 2, "test/unknown-a", serde_json::json!({})),
)
.await?;
send_relay_message(
&mut websocket,
stream_b,
/*seq*/ 2,
request(/*id*/ 2, "test/unknown-b", serde_json::json!({})),
)
.await?;
let unknown_method_responses =
read_relay_messages_by_stream(&mut websocket, /*count*/ 2).await?;
assert_error_response(
unknown_method_responses.get(stream_a),
stream_a,
/*id*/ 2,
"test/unknown-a",
)?;
assert_error_response(
unknown_method_responses.get(stream_b),
stream_b,
/*id*/ 2,
"test/unknown-b",
)?;
send_relay_reset(&mut websocket, stream_a, "test_reset").await?;
send_relay_message(
&mut websocket,
stream_b,
/*seq*/ 3,
request(
/*id*/ 3,
"test/unknown-b-after-reset",
serde_json::json!({}),
),
)
.await?;
let (stream_id, message) = read_relay_message(&mut websocket).await?;
assert_eq!(stream_id, stream_b);
assert_error_response(
Some(&message),
stream_b,
/*id*/ 3,
"test/unknown-b-after-reset",
)?;
websocket.close(None).await?;
remote_executor.abort();
let _ = remote_executor.await;
Ok(())
}
async fn send_relay_message<S>(
websocket: &mut WebSocketStream<S>,
stream_id: &str,
seq: u32,
message: JSONRPCMessage,
) -> Result<()>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
let payload = serde_json::to_vec(&message)?;
let frame = RelayMessageFrame {
version: RELAY_MESSAGE_FRAME_VERSION,
stream_id: stream_id.to_string(),
ack: 0,
ack_bits: 0,
body: Some(relay_message_frame::Body::Data(RelayData {
seq,
segment_index: 0,
segment_count: 1,
payload,
})),
};
send_relay_frame(websocket, frame).await
}
async fn send_relay_reset<S>(
websocket: &mut WebSocketStream<S>,
stream_id: &str,
reason: &str,
) -> Result<()>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
send_relay_frame(
websocket,
RelayMessageFrame {
version: RELAY_MESSAGE_FRAME_VERSION,
stream_id: stream_id.to_string(),
ack: 0,
ack_bits: 0,
body: Some(relay_message_frame::Body::Reset(RelayReset {
reason: reason.to_string(),
})),
},
)
.await
}
async fn send_relay_frame<S>(
websocket: &mut WebSocketStream<S>,
frame: RelayMessageFrame,
) -> Result<()>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
websocket
.send(Message::Binary(frame.encode_to_vec().into()))
.await?;
Ok(())
}
async fn read_relay_messages_by_stream<S>(
websocket: &mut WebSocketStream<S>,
count: usize,
) -> Result<HashMap<String, JSONRPCMessage>>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
let mut messages = HashMap::new();
for _ in 0..count {
let (stream_id, message) = read_relay_message(websocket).await?;
if messages.insert(stream_id.clone(), message).is_some() {
bail!("received duplicate response for stream {stream_id}");
}
}
Ok(messages)
}
async fn read_relay_message<S>(
websocket: &mut WebSocketStream<S>,
) -> Result<(String, JSONRPCMessage)>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
loop {
let frame = timeout(TEST_TIMEOUT, websocket.next())
.await
.context("timed out waiting for relay frame")?
.ok_or_else(|| anyhow!("executor websocket closed"))??;
match frame {
Message::Binary(bytes) => {
let frame = RelayMessageFrame::decode(bytes.as_ref())?;
let stream_id = frame.stream_id;
let Some(relay_message_frame::Body::Data(data)) = frame.body else {
continue;
};
let message = serde_json::from_slice(&data.payload)?;
return Ok((stream_id, message));
}
Message::Ping(_) | Message::Pong(_) => {}
Message::Close(_) => bail!("executor websocket closed"),
Message::Text(_) => bail!("executor sent text frame on relay websocket"),
Message::Frame(_) => {}
}
}
}
fn initialize_request(id: i64, client_name: &str) -> Result<JSONRPCMessage> {
Ok(request(
id,
"initialize",
serde_json::to_value(InitializeParams {
client_name: client_name.to_string(),
resume_session_id: None,
})?,
))
}
fn request(id: i64, method: &str, params: serde_json::Value) -> JSONRPCMessage {
JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(id),
method: method.to_string(),
params: Some(params),
trace: None,
})
}
fn notification(method: &str, params: serde_json::Value) -> JSONRPCMessage {
JSONRPCMessage::Notification(JSONRPCNotification {
method: method.to_string(),
params: Some(params),
})
}
fn assert_initialize_response(
message: Option<&JSONRPCMessage>,
stream_id: &str,
id: i64,
) -> Result<Uuid> {
let message = message.ok_or_else(|| anyhow!("missing initialize response for {stream_id}"))?;
let JSONRPCMessage::Response(JSONRPCResponse {
id: response_id,
result,
}) = message
else {
bail!("expected initialize response for {stream_id}, got {message:?}");
};
assert_eq!(response_id, &RequestId::Integer(id));
let response: InitializeResponse = serde_json::from_value(result.clone())?;
Ok(Uuid::parse_str(&response.session_id)?)
}
fn assert_error_response(
message: Option<&JSONRPCMessage>,
stream_id: &str,
id: i64,
expected_method: &str,
) -> Result<()> {
let message = message.ok_or_else(|| anyhow!("missing error response for {stream_id}"))?;
let JSONRPCMessage::Error(JSONRPCError {
id: response_id,
error,
}) = message
else {
bail!("expected error response for {stream_id}, got {message:?}");
};
assert_eq!(response_id, &RequestId::Integer(id));
assert!(
error.message.contains(expected_method),
"expected error for {stream_id} to mention {expected_method}, got {}",
error.message
);
Ok(())
}
codex-rs/thread-manager-sample/src/main.rs
+1
View File
@@ -184,6 +184,7 @@ fn new_config(model: Option<String>, arg0_paths: Arg0DispatchPaths) -> anyhow::R
guardian_policy_config: None,
include_permissions_instructions: false,
include_apps_instructions: false,
include_collaboration_mode_instructions: false,
include_skill_instructions: false,
include_environment_context: false,
compact_prompt: None,
codex-rs/tui/src/chatwidget/tests/status_and_layout.rs
+2 -1
View File
@@ -401,6 +401,7 @@ async fn completed_plan_table_tail_skips_provisional_history_insert() {
}
#[tokio::test]
#[cfg_attr(target_os = "windows", ignore = "disabled on windows")]
async fn configured_pet_load_is_deferred_until_after_construction() {
let (tx_raw, mut rx) = unbounded_channel::<AppEvent>();
let tx = AppEventSender::new(tx_raw);
@@ -434,7 +435,7 @@ async fn configured_pet_load_is_deferred_until_after_construction() {
let chat = ChatWidget::new_with_app_event(init);
assert!(!chat.ambient_pet_image_enabled());
let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
let event = tokio::time::timeout(std::time::Duration::from_secs(/*secs*/ 30), rx.recv())
.await
.unwrap()
.unwrap();