diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index afb3cbbed9..f1d78b7479 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2740,6 +2740,7 @@ dependencies = [ "ctor 0.6.3", "futures", "pretty_assertions", + "prost 0.14.3", "reqwest", "serde", "serde_json", diff --git a/codex-rs/code-mode/src/runtime/callbacks.rs b/codex-rs/code-mode/src/runtime/callbacks.rs index a9755f6eb0..c3a648ae32 100644 --- a/codex-rs/code-mode/src/runtime/callbacks.rs +++ b/codex-rs/code-mode/src/runtime/callbacks.rs @@ -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::() 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::() 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()); diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 200a47c989..f989312da1 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -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, } @@ -126,6 +128,7 @@ pub(crate) enum RuntimeEvent { ToolCall { id: String, name: ToolName, + kind: CodeModeToolKind, input: Option, }, Notify { diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 7326c834e2..3ab8913f19 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -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 diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 80a72a88ec..8eb1697fc6 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -8,17 +8,20 @@ use codex_tools::ToolSpec; use super::ExecContext; use super::PUBLIC_TOOL_NAME; -use super::build_enabled_tools; use super::handle_runtime_response; use super::is_exec_tool_name; pub struct CodeModeExecuteHandler { spec: ToolSpec, + nested_tool_specs: Vec, } impl CodeModeExecuteHandler { - pub(crate) fn new(spec: ToolSpec) -> Self { - Self { spec } + pub(crate) fn new(spec: ToolSpec, nested_tool_specs: Vec) -> Self { + Self { + spec, + nested_tool_specs, + } } async fn execute( @@ -31,7 +34,8 @@ impl CodeModeExecuteHandler { let args = codex_code_mode::parse_exec_source(&code).map_err(FunctionCallError::RespondToModel)?; let exec = ExecContext { session, turn }; - let enabled_tools = build_enabled_tools(&exec).await; + let enabled_tools = + codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); let stored_values = exec .session .services diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 9ee8e0352a..ff9f8c8893 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -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; @@ -28,13 +29,9 @@ use crate::tools::context::ToolPayload; use crate::tools::parallel::ToolCallRuntime; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; -use crate::tools::router::ToolRouterParams; -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; use codex_utils_output_truncation::truncate_function_output_items_with_policy; @@ -260,36 +257,6 @@ fn truncate_code_mode_result( truncate_function_output_items_with_policy(&items, policy) } -pub(super) async fn build_enabled_tools( - exec: &ExecContext, -) -> Vec { - let router = build_nested_router(exec).await; - let specs = router.specs(); - collect_code_mode_tool_definitions(&specs) -} - -#[expect( - clippy::await_holding_invalid_type, - reason = "nested tool router construction reads through the session-owned manager guard" -)] -async fn build_nested_router(exec: &ExecContext) -> ToolRouter { - let nested_tools_config = exec.turn.tools_config.for_code_mode_nested_tools(); - let mcp_connection_manager = exec.session.services.mcp_connection_manager.read().await; - let listed_mcp_tools = mcp_connection_manager.list_all_tools().await; - - ToolRouter::from_config( - &nested_tools_config, - ToolRouterParams { - deferred_mcp_tools: None, - mcp_tools: Some(listed_mcp_tools), - unavailable_called_tools: Vec::new(), - discoverable_tools: None, - extension_tool_bundles: extension_tool_bundles(exec.session.as_ref()), - dynamic_tools: exec.turn.dynamic_tools.as_slice(), - }, - ) -} - async fn call_nested_tool( _exec: ExecContext, tool_runtime: ToolCallRuntime, @@ -300,6 +267,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 +276,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 +299,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, - tool_name: &ToolName, -) -> Result { - 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, + tool_kind: CodeModeToolKind, tool_name: &ToolName, input: Option, ) -> Result { - 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 +341,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:?}"), + } + } +} diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index 7cc9e61e30..9ecdf2340d 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -9,10 +9,13 @@ use crate::tools::registry::ToolHandler; use crate::turn_timing::now_unix_timestamp_ms; use codex_protocol::dynamic_tools::DynamicToolCallRequest; use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::protocol::DynamicToolCallResponseEvent; use codex_protocol::protocol::EventMsg; use codex_tools::ToolName; +use codex_tools::ToolSpec; +use codex_tools::dynamic_tool_to_loadable_tool_spec; use serde_json::Value; use std::time::Instant; use tokio::sync::oneshot; @@ -20,11 +23,17 @@ use tracing::warn; pub struct DynamicToolHandler { tool_name: ToolName, + spec: Option, } impl DynamicToolHandler { - pub fn new(tool_name: ToolName) -> Self { - Self { tool_name } + pub fn new(tool: &DynamicToolSpec) -> Option { + let tool_name = ToolName::new(tool.namespace.clone(), tool.name.clone()); + let spec = dynamic_tool_to_loadable_tool_spec(tool).ok()?.into(); + Some(Self { + tool_name, + spec: Some(spec), + }) } } @@ -35,6 +44,10 @@ impl ToolHandler for DynamicToolHandler { self.tool_name.clone() } + fn spec(&self) -> Option { + self.spec.clone() + } + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 8bb8337337..89666c80aa 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -14,7 +14,11 @@ use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolTelemetryTags; use codex_mcp::ToolInfo; +use codex_tools::ResponsesApiNamespace; +use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolName; +use codex_tools::ToolSpec; +use codex_tools::mcp_tool_to_responses_api_tool; use serde_json::Map; use serde_json::Value; @@ -35,6 +39,26 @@ impl ToolHandler for McpHandler { self.tool_info.canonical_tool_name() } + fn spec(&self) -> Option { + let tool_name = self.tool_name(); + let namespace_name = tool_name.namespace.as_ref()?; + let tool = mcp_tool_to_responses_api_tool(&tool_name, &self.tool_info.tool).ok()?; + let description = self + .tool_info + .namespace_description + .as_deref() + .map(str::trim) + .filter(|description| !description.is_empty()) + .map(str::to_string) + .unwrap_or_default(); + + Some(ToolSpec::Namespace(ResponsesApiNamespace { + name: namespace_name.clone(), + description, + tools: vec![ResponsesApiNamespaceTool::Function(tool)], + })) + } + fn supports_parallel_tool_calls(&self) -> bool { self.tool_info.supports_parallel_tool_calls } diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index d7fe22e4f4..2c62e6c15d 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -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 { - self.router.find_spec(tool_name) - } - pub(crate) fn create_diff_consumer( &self, tool_name: &codex_tools::ToolName, diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 01b3caa34c..95e7d17ce5 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -168,7 +168,11 @@ pub(crate) struct PostToolUsePayload { pub(crate) tool_response: Value, } -trait AnyToolHandler: Send + Sync { +pub(crate) trait AnyToolHandler: Send + Sync { + fn tool_name(&self) -> ToolName; + + fn spec(&self) -> Option; + fn supports_parallel_tool_calls(&self) -> bool; fn matches_kind(&self, payload: &ToolPayload) -> bool; @@ -197,6 +201,14 @@ impl AnyToolHandler for T where T: ToolHandler, { + fn tool_name(&self) -> ToolName { + ToolHandler::tool_name(self) + } + + fn spec(&self) -> Option { + ToolHandler::spec(self) + } + fn supports_parallel_tool_calls(&self) -> bool { ToolHandler::supports_parallel_tool_calls(self) } @@ -531,24 +543,17 @@ impl ToolRegistry { pub struct ToolRegistryBuilder { handlers: HashMap>, specs: Vec, - code_mode_enabled: bool, } impl ToolRegistryBuilder { - pub fn new(code_mode_enabled: bool) -> Self { + pub fn new() -> Self { Self { handlers: HashMap::new(), specs: Vec::new(), - code_mode_enabled, } } pub(crate) fn push_spec(&mut self, spec: ToolSpec) { - let spec = if self.code_mode_enabled { - codex_tools::augment_tool_spec_for_code_mode(spec) - } else { - spec - }; self.specs.push(spec); } @@ -556,17 +561,32 @@ impl ToolRegistryBuilder { where H: ToolHandler + 'static, { + self.register_any_handler(handler); + } + + pub(crate) fn register_any_handler(&mut self, handler: Arc) { + self.register_any_handler_internal(handler, /*include_spec*/ true); + } + + pub(crate) fn register_any_handler_without_spec(&mut self, handler: Arc) { + self.register_any_handler_internal(handler, /*include_spec*/ false); + } + + fn register_any_handler_internal( + &mut self, + handler: Arc, + include_spec: bool, + ) { let name = handler.tool_name(); if self.handlers.contains_key(&name) { error_or_panic(format!("handler for tool {name} already registered")); return; } - if let Some(spec) = handler.spec() { + if include_spec && let Some(spec) = handler.spec() { self.push_spec(spec); } - let handler: Arc = handler; self.handlers.insert(name, handler); } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index d6acb80fb4..dc744321d1 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -63,16 +63,13 @@ fn handler_looks_up_namespaced_aliases_explicitly() { } #[test] -fn register_handler_adds_handler_and_augments_specs_for_code_mode() { - let mut builder = ToolRegistryBuilder::new(/*code_mode_enabled*/ true); +fn register_handler_adds_handler_and_spec() { + let mut builder = ToolRegistryBuilder::new(); builder.register_handler(Arc::new(GetGoalHandler)); let (specs, registry) = builder.build(); assert_eq!(specs.len(), 1); - assert_eq!( - specs[0], - codex_tools::augment_tool_spec_for_code_mode(create_get_goal_tool()) - ); + assert_eq!(specs[0], create_get_goal_tool()); assert!(registry.has_handler(&codex_tools::ToolName::plain(GET_GOAL_TOOL_NAME))); } diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index ed3f1a720b..4bc366bf91 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -37,7 +37,6 @@ pub struct ToolCall { pub struct ToolRouter { registry: ToolRegistry, - specs: Vec, model_visible_specs: Vec, } @@ -90,44 +89,14 @@ impl ToolRouter { Self { registry, - specs, model_visible_specs, } } - pub fn specs(&self) -> Vec { - self.specs.clone() - } - pub fn model_visible_specs(&self) -> Vec { self.model_visible_specs.clone() } - pub fn find_spec(&self, tool_name: &ToolName) -> Option { - 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, diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index e78bd9402f..4d6119cd5c 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -232,7 +232,7 @@ async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> } #[tokio::test] -async fn model_visible_specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { +async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; let hidden_tool = "hidden_dynamic_tool"; let visible_tool = "visible_dynamic_tool"; @@ -273,15 +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()] - ); assert_eq!( namespace_function_names(&router.model_visible_specs(), "codex_app"), vec![visible_tool.to_string()] @@ -339,12 +330,6 @@ async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow:: }, ); - assert!( - router - .find_spec(&ToolName::plain("extension_echo")) - .is_some(), - "expected extension-provided tool spec to be registered" - ); assert!( router .model_visible_specs() diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index fc21ff3375..8c54c21479 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -70,7 +70,6 @@ pub(crate) fn build_specs_with_discoverable_tools( use crate::tools::tool_search_entry::build_tool_search_entries_for_config; let mcp_tool_plan_inputs = mcp_tools.as_deref().map(map_mcp_tools_for_plan); - let deferred_mcp_tool_sources = deferred_mcp_tools.as_deref(); let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let min_wait_timeout_ms = if config.multi_agent_v2 { @@ -85,12 +84,12 @@ pub(crate) fn build_specs_with_discoverable_tools( DEFAULT_WAIT_TIMEOUT_MS.clamp(min_wait_timeout_ms, MAX_WAIT_TIMEOUT_MS); let deferred_dynamic_tools = dynamic_tools .iter() - .filter(|tool| tool.defer_loading && (config.namespace_tools || tool.namespace.is_none())) + .filter(|tool| tool.defer_loading) .cloned() .collect::>(); let tool_search_entries = build_tool_search_entries_for_config( config, - deferred_mcp_tool_sources, + deferred_mcp_tools.as_deref(), &deferred_dynamic_tools, ); let mut builder = build_tool_registry_builder( @@ -99,7 +98,7 @@ pub(crate) fn build_specs_with_discoverable_tools( mcp_tools: mcp_tool_plan_inputs .as_ref() .map(|inputs| inputs.mcp_tools.as_slice()), - deferred_mcp_tools: deferred_mcp_tool_sources, + deferred_mcp_tools: deferred_mcp_tools.as_deref(), tool_namespaces: mcp_tool_plan_inputs .as_ref() .map(|inputs| &inputs.tool_namespaces), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 69456a36ef..7232dddd85 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -27,6 +27,7 @@ use crate::tools::handlers::ViewImageHandler; use crate::tools::handlers::WriteStdinHandler; use crate::tools::handlers::agent_jobs::ReportAgentJobResultHandler; use crate::tools::handlers::agent_jobs::SpawnAgentsOnCsvHandler; +use crate::tools::handlers::extension_tools::extension_tool_spec; use crate::tools::handlers::multi_agents::CloseAgentHandler; use crate::tools::handlers::multi_agents::ResumeAgentHandler; use crate::tools::handlers::multi_agents::SendInputHandler; @@ -44,12 +45,11 @@ use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::hosted_spec::WebSearchToolOptions; use crate::tools::hosted_spec::create_image_generation_tool; use crate::tools::hosted_spec::create_web_search_tool; +use crate::tools::registry::AnyToolHandler; use crate::tools::registry::ToolRegistryBuilder; use crate::tools::spec_plan_types::ToolRegistryBuildParams; use crate::tools::spec_plan_types::agent_type_description; -use codex_mcp::ToolInfo; use codex_protocol::openai_models::ConfigShellToolType; -use codex_tools::ResponsesApiNamespace; use codex_tools::ResponsesApiNamespaceTool; use codex_tools::ToolEnvironmentMode; use codex_tools::ToolName; @@ -57,12 +57,9 @@ use codex_tools::ToolSearchSource; use codex_tools::ToolSearchSourceInfo; use codex_tools::ToolSpec; use codex_tools::ToolsConfig; -use codex_tools::coalesce_loadable_tool_specs; use codex_tools::collect_code_mode_exec_prompt_tool_definitions; use codex_tools::collect_tool_search_source_infos; use codex_tools::default_namespace_description; -use codex_tools::dynamic_tool_to_loadable_tool_spec; -use codex_tools::mcp_tool_to_responses_api_tool; use std::collections::BTreeMap; use std::collections::HashSet; use std::sync::Arc; @@ -71,8 +68,21 @@ pub fn build_tool_registry_builder( config: &ToolsConfig, params: ToolRegistryBuildParams<'_>, ) -> ToolRegistryBuilder { - let mut builder = ToolRegistryBuilder::new(config.code_mode_enabled); - let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled; + let mut builder = ToolRegistryBuilder::new(); + let all_deferred_tools = params + .deferred_mcp_tools + .into_iter() + .flatten() + .map(codex_mcp::ToolInfo::canonical_tool_name) + .chain( + params + .dynamic_tools + .iter() + .filter(|tool| tool.defer_loading) + .map(|tool| ToolName::new(tool.namespace.clone(), tool.name.clone())), + ) + .collect::>(); + let handlers = collect_handler_tools(config, params); if config.code_mode_enabled { let namespace_descriptions = params @@ -89,16 +99,18 @@ pub fn build_tool_registry_builder( ) }) .collect::>(); - let nested_config = config.for_code_mode_nested_tools(); - let nested_builder = build_tool_registry_builder( - &nested_config, - ToolRegistryBuildParams { - discoverable_tools: None, - ..params - }, + let mut code_mode_nested_tool_specs = handlers + .iter() + .filter_map(|handler| handler.spec()) + .collect::>(); + code_mode_nested_tool_specs.extend( + params + .extension_tool_bundles + .iter() + .filter_map(|bundle| extension_tool_spec(bundle.spec()).ok()), ); let mut enabled_tools = - collect_code_mode_exec_prompt_tool_definitions(nested_builder.specs().iter()); + collect_code_mode_exec_prompt_tool_definitions(code_mode_nested_tool_specs.iter()); enabled_tools .sort_by(|left, right| compare_code_mode_tools(left, right, &namespace_descriptions)); builder.register_handler(Arc::new(CodeModeExecuteHandler::new( @@ -106,40 +118,162 @@ pub fn build_tool_registry_builder( &enabled_tools, &namespace_descriptions, config.code_mode_only_enabled, - config.search_tool - && params - .deferred_mcp_tools - .is_some_and(|tools| !tools.is_empty()), + config.search_tool && !all_deferred_tools.is_empty(), ), + code_mode_nested_tool_specs, ))); builder.register_handler(Arc::new(CodeModeWaitHandler)); } + let mut non_deferred_specs = Vec::new(); + for handler in &handlers { + let tool_name = handler.tool_name(); + if !all_deferred_tools.contains(&tool_name) + && let Some(spec) = handler.spec() + { + non_deferred_specs.push(spec); + } + } + + if let Some(web_search_tool) = create_web_search_tool(WebSearchToolOptions { + web_search_mode: config.web_search_mode, + web_search_config: config.web_search_config.as_ref(), + web_search_tool_type: config.web_search_tool_type, + }) { + non_deferred_specs.push(web_search_tool); + } + if config.image_gen_tool { + non_deferred_specs.push(create_image_generation_tool("png")); + } + + for spec in merge_into_namespaces(non_deferred_specs) { + if !config.namespace_tools && matches!(spec, ToolSpec::Namespace(_)) { + continue; + } + let spec = if config.code_mode_enabled { + codex_tools::augment_tool_spec_for_code_mode(spec) + } else { + spec + }; + builder.push_spec(spec); + } + + for handler in handlers { + builder.register_any_handler_without_spec(handler); + } + + if config.search_tool && config.namespace_tools && !all_deferred_tools.is_empty() { + let mut search_source_infos = params + .deferred_mcp_tools + .map(|mcp_tools| { + collect_tool_search_source_infos(mcp_tools.iter().map(|tool| ToolSearchSource { + server_name: tool.server_name.as_str(), + connector_name: tool.connector_name.as_deref(), + description: tool.namespace_description.as_deref(), + })) + }) + .unwrap_or_default(); + + if params.dynamic_tools.iter().any(|tool| { + all_deferred_tools.contains(&ToolName::new(tool.namespace.clone(), tool.name.clone())) + }) { + search_source_infos.push(ToolSearchSourceInfo { + name: "Dynamic tools".to_string(), + description: Some("Tools provided by the current Codex thread.".to_string()), + }); + } + + builder.register_handler(Arc::new(ToolSearchHandler::new( + params.tool_search_entries.to_vec(), + search_source_infos, + ))); + } + + for bundle in params.extension_tool_bundles.iter().cloned() { + builder.register_tool_bundle(bundle); + } + + builder +} + +fn merge_into_namespaces(specs: Vec) -> Vec { + let mut merged_specs = Vec::with_capacity(specs.len()); + let mut namespace_indices = BTreeMap::::new(); + for spec in specs { + match spec { + ToolSpec::Namespace(mut namespace) => { + if let Some(index) = namespace_indices.get(&namespace.name).copied() { + let ToolSpec::Namespace(existing_namespace) = &mut merged_specs[index] else { + unreachable!("namespace index must point to a namespace spec"); + }; + if existing_namespace.description.trim().is_empty() + && !namespace.description.trim().is_empty() + { + existing_namespace.description = namespace.description; + } + existing_namespace.tools.append(&mut namespace.tools); + continue; + } + + namespace_indices.insert(namespace.name.clone(), merged_specs.len()); + merged_specs.push(ToolSpec::Namespace(namespace)); + } + spec => merged_specs.push(spec), + } + } + + for spec in &mut merged_specs { + let ToolSpec::Namespace(namespace) = spec else { + continue; + }; + + namespace.tools.sort_by(|left, right| match (left, right) { + ( + ResponsesApiNamespaceTool::Function(left), + ResponsesApiNamespaceTool::Function(right), + ) => left.name.cmp(&right.name), + }); + + if namespace.description.trim().is_empty() { + namespace.description = default_namespace_description(&namespace.name); + } + } + + merged_specs +} + +fn collect_handler_tools( + config: &ToolsConfig, + params: ToolRegistryBuildParams<'_>, +) -> Vec> { + let exec_permission_approvals_enabled = config.exec_permission_approvals_enabled; + let mut handlers = Vec::>::new(); + if config.environment_mode.has_environment() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); match &config.shell_type { ConfigShellToolType::Default => { - builder.register_handler(Arc::new(ShellHandler::new(ShellToolOptions { + handlers.push(Arc::new(ShellHandler::new(ShellToolOptions { exec_permission_approvals_enabled, }))); } ConfigShellToolType::Local => { - builder.register_handler(Arc::new(LocalShellHandler::new())); + handlers.push(Arc::new(LocalShellHandler::new())); } ConfigShellToolType::UnifiedExec => { - builder.register_handler(Arc::new(ExecCommandHandler::new( + handlers.push(Arc::new(ExecCommandHandler::new( ExecCommandHandlerOptions { allow_login_shell: config.allow_login_shell, exec_permission_approvals_enabled, include_environment_id, }, ))); - builder.register_handler(Arc::new(WriteStdinHandler)); + handlers.push(Arc::new(WriteStdinHandler)); } ConfigShellToolType::Disabled => {} ConfigShellToolType::ShellCommand => { - builder.register_handler(Arc::new(ShellCommandHandler::new( + handlers.push(Arc::new(ShellCommandHandler::new( ShellCommandHandlerOptions { backend_config: config.shell_command_backend, allow_login_shell: config.allow_login_shell, @@ -155,101 +289,62 @@ pub fn build_tool_registry_builder( { match &config.shell_type { ConfigShellToolType::Default => { - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(LocalShellHandler::default())); - builder.register_handler(Arc::new(ShellCommandHandler::from( + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(LocalShellHandler::default())); + handlers.push(Arc::new(ShellCommandHandler::from( config.shell_command_backend, ))); } ConfigShellToolType::Local => { - builder.register_handler(Arc::new(ShellHandler::default())); - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(ShellCommandHandler::from( + handlers.push(Arc::new(ShellHandler::default())); + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(ShellCommandHandler::from( config.shell_command_backend, ))); } ConfigShellToolType::UnifiedExec => { - builder.register_handler(Arc::new(ShellHandler::default())); - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(LocalShellHandler::default())); - builder.register_handler(Arc::new(ShellCommandHandler::from( + handlers.push(Arc::new(ShellHandler::default())); + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(LocalShellHandler::default())); + handlers.push(Arc::new(ShellCommandHandler::from( config.shell_command_backend, ))); } ConfigShellToolType::ShellCommand => { - builder.register_handler(Arc::new(ShellHandler::default())); - builder.register_handler(Arc::new(ContainerExecHandler)); - builder.register_handler(Arc::new(LocalShellHandler::default())); + handlers.push(Arc::new(ShellHandler::default())); + handlers.push(Arc::new(ContainerExecHandler)); + handlers.push(Arc::new(LocalShellHandler::default())); } ConfigShellToolType::Disabled => {} } } if params.mcp_tools.is_some() { - builder.register_handler(Arc::new(ListMcpResourcesHandler)); - builder.register_handler(Arc::new(ListMcpResourceTemplatesHandler)); - builder.register_handler(Arc::new(ReadMcpResourceHandler)); + handlers.push(Arc::new(ListMcpResourcesHandler)); + handlers.push(Arc::new(ListMcpResourceTemplatesHandler)); + handlers.push(Arc::new(ReadMcpResourceHandler)); } - builder.register_handler(Arc::new(PlanHandler)); + handlers.push(Arc::new(PlanHandler)); if config.goal_tools { - builder.register_handler(Arc::new(GetGoalHandler)); - builder.register_handler(Arc::new(CreateGoalHandler)); - builder.register_handler(Arc::new(UpdateGoalHandler)); + handlers.push(Arc::new(GetGoalHandler)); + handlers.push(Arc::new(CreateGoalHandler)); + handlers.push(Arc::new(UpdateGoalHandler)); } - builder.register_handler(Arc::new(RequestUserInputHandler { + handlers.push(Arc::new(RequestUserInputHandler { available_modes: config.request_user_input_available_modes.clone(), })); if config.request_permissions_tool_enabled { - builder.register_handler(Arc::new(RequestPermissionsHandler)); - } - - let deferred_dynamic_tools = params - .dynamic_tools - .iter() - .filter(|tool| tool.defer_loading && (config.namespace_tools || tool.namespace.is_none())) - .collect::>(); - let deferred_mcp_tools_for_search = if config.namespace_tools { - params.deferred_mcp_tools - } else { - None - }; - - if config.search_tool - && (deferred_mcp_tools_for_search.is_some() || !deferred_dynamic_tools.is_empty()) - { - let mut search_source_infos = deferred_mcp_tools_for_search - .map(|deferred_mcp_tools| { - collect_tool_search_source_infos(deferred_mcp_tools.iter().map(|tool| { - ToolSearchSource { - server_name: tool.server_name.as_str(), - connector_name: tool.connector_name.as_deref(), - description: tool.namespace_description.as_deref(), - } - })) - }) - .unwrap_or_default(); - - if !deferred_dynamic_tools.is_empty() { - search_source_infos.push(ToolSearchSourceInfo { - name: "Dynamic tools".to_string(), - description: Some("Tools provided by the current Codex thread.".to_string()), - }); - } - - builder.register_handler(Arc::new(ToolSearchHandler::new( - params.tool_search_entries.to_vec(), - search_source_infos, - ))); + handlers.push(Arc::new(RequestPermissionsHandler)); } if config.tool_suggest && let Some(discoverable_tools) = params.discoverable_tools.filter(|tools| !tools.is_empty()) { - builder.register_handler(Arc::new(RequestPluginInstallHandler::new( + handlers.push(Arc::new(RequestPluginInstallHandler::new( discoverable_tools, ))); } @@ -257,7 +352,7 @@ pub fn build_tool_registry_builder( if config.environment_mode.has_environment() && config.apply_patch_tool_type.is_some() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); - builder.register_handler(Arc::new(ApplyPatchHandler::new(include_environment_id))); + handlers.push(Arc::new(ApplyPatchHandler::new(include_environment_id))); } if config @@ -265,25 +360,13 @@ pub fn build_tool_registry_builder( .iter() .any(|tool| tool == "test_sync_tool") { - builder.register_handler(Arc::new(TestSyncHandler)); - } - - if let Some(web_search_tool) = create_web_search_tool(WebSearchToolOptions { - web_search_mode: config.web_search_mode, - web_search_config: config.web_search_config.as_ref(), - web_search_tool_type: config.web_search_tool_type, - }) { - builder.push_spec(web_search_tool); - } - - if config.image_gen_tool { - builder.push_spec(create_image_generation_tool("png")); + handlers.push(Arc::new(TestSyncHandler)); } if config.environment_mode.has_environment() { let include_environment_id = matches!(config.environment_mode, ToolEnvironmentMode::Multiple); - builder.register_handler(Arc::new(ViewImageHandler::new(ViewImageToolOptions { + handlers.push(Arc::new(ViewImageHandler::new(ViewImageToolOptions { can_request_original_image_detail: config.can_request_original_image_detail, include_environment_id, }))); @@ -293,7 +376,7 @@ pub fn build_tool_registry_builder( if config.multi_agent_v2 { let agent_type_description = agent_type_description(config, params.default_agent_type_description); - builder.register_handler(Arc::new(SpawnAgentHandlerV2::new(SpawnAgentToolOptions { + handlers.push(Arc::new(SpawnAgentHandlerV2::new(SpawnAgentToolOptions { available_models: config.available_models.clone(), agent_type_description, hide_agent_type_model_reasoning: config.hide_spawn_agent_metadata, @@ -301,17 +384,17 @@ pub fn build_tool_registry_builder( usage_hint_text: config.spawn_agent_usage_hint_text.clone(), max_concurrent_threads_per_session: config.max_concurrent_threads_per_session, }))); - builder.register_handler(Arc::new(SendMessageHandlerV2)); - builder.register_handler(Arc::new(FollowupTaskHandlerV2)); - builder.register_handler(Arc::new(WaitAgentHandlerV2::new( + handlers.push(Arc::new(SendMessageHandlerV2)); + handlers.push(Arc::new(FollowupTaskHandlerV2)); + handlers.push(Arc::new(WaitAgentHandlerV2::new( params.wait_agent_timeouts, ))); - builder.register_handler(Arc::new(CloseAgentHandlerV2)); - builder.register_handler(Arc::new(ListAgentsHandlerV2)); + handlers.push(Arc::new(CloseAgentHandlerV2)); + handlers.push(Arc::new(ListAgentsHandlerV2)); } else { let agent_type_description = agent_type_description(config, params.default_agent_type_description); - builder.register_handler(Arc::new(SpawnAgentHandler::new(SpawnAgentToolOptions { + handlers.push(Arc::new(SpawnAgentHandler::new(SpawnAgentToolOptions { available_models: config.available_models.clone(), agent_type_description, hide_agent_type_model_reasoning: config.hide_spawn_agent_metadata, @@ -319,122 +402,45 @@ pub fn build_tool_registry_builder( usage_hint_text: config.spawn_agent_usage_hint_text.clone(), max_concurrent_threads_per_session: config.max_concurrent_threads_per_session, }))); - builder.register_handler(Arc::new(SendInputHandler)); - builder.register_handler(Arc::new(ResumeAgentHandler)); - builder.register_handler(Arc::new(WaitAgentHandler::new(params.wait_agent_timeouts))); - builder.register_handler(Arc::new(CloseAgentHandler)); + handlers.push(Arc::new(SendInputHandler)); + handlers.push(Arc::new(ResumeAgentHandler)); + handlers.push(Arc::new(WaitAgentHandler::new(params.wait_agent_timeouts))); + handlers.push(Arc::new(CloseAgentHandler)); } } if config.agent_jobs_tools { - builder.register_handler(Arc::new(SpawnAgentsOnCsvHandler)); + handlers.push(Arc::new(SpawnAgentsOnCsvHandler)); if config.agent_jobs_worker_tools { - builder.register_handler(Arc::new(ReportAgentJobResultHandler)); + handlers.push(Arc::new(ReportAgentJobResultHandler)); } } if let Some(mcp_tools) = params.mcp_tools { - let mut entries = mcp_tools - .iter() - .map(|tool| (tool.canonical_tool_name(), tool)) - .collect::>(); - entries.sort_by(|(left, _), (right, _)| left.cmp(right)); - let mut namespace_entries = BTreeMap::new(); - - for (tool_name, tool) in entries { - let Some(namespace) = tool_name.namespace.as_ref() else { - tracing::error!("Skipping MCP tool `{tool_name}`: MCP tools must be namespaced"); - continue; - }; - namespace_entries - .entry(namespace.clone()) - .or_insert_with(Vec::new) - .push((tool_name, tool)); - } - - for (namespace, mut entries) in namespace_entries { - entries.sort_by_key(|(tool_name, _)| tool_name.name.clone()); - let tool_namespace = params - .tool_namespaces - .and_then(|namespaces| namespaces.get(&namespace)); - let description = tool_namespace - .and_then(|namespace| namespace.description.as_deref()) - .map(str::trim) - .filter(|description| !description.is_empty()) - .map(str::to_string) - .unwrap_or_else(|| { - let namespace_name = tool_namespace - .map(|namespace| namespace.name.as_str()) - .unwrap_or(namespace.as_str()); - default_namespace_description(namespace_name) - }); - let mut tools = Vec::new(); - for (tool_name, tool) in entries { - match mcp_tool_to_responses_api_tool(&tool_name, &tool.tool) { - Ok(converted_tool) => { - tools.push(ResponsesApiNamespaceTool::Function(converted_tool)); - builder.register_handler(Arc::new(McpHandler::new(tool.clone()))); - } - Err(error) => { - tracing::error!( - "Failed to convert `{tool_name}` MCP tool to OpenAI tool: {error:?}" - ); - } - } - } - - if config.namespace_tools && !tools.is_empty() { - builder.push_spec(ToolSpec::Namespace(ResponsesApiNamespace { - name: namespace, - description, - tools, - })); - } - } - } - - let mut dynamic_tool_specs = Vec::new(); - for tool in params.dynamic_tools { - match dynamic_tool_to_loadable_tool_spec(tool) { - Ok(loadable_tool) => { - let handler_name = ToolName::new(tool.namespace.clone(), tool.name.clone()); - dynamic_tool_specs.push(loadable_tool); - builder.register_handler(Arc::new(DynamicToolHandler::new(handler_name))); - } - Err(error) => { - tracing::error!( - "Failed to convert dynamic tool {:?} to OpenAI tool: {error:?}", - tool.name - ); - } - } - } - for spec in coalesce_loadable_tool_specs(dynamic_tool_specs) { - let spec = spec.into(); - if config.namespace_tools || !matches!(spec, ToolSpec::Namespace(_)) { - builder.push_spec(spec); + for tool in mcp_tools { + handlers.push(Arc::new(McpHandler::new(tool.clone()))); } } if let Some(deferred_mcp_tools) = params.deferred_mcp_tools { - let directly_registered_mcp_tools = params - .mcp_tools - .into_iter() - .flatten() - .map(ToolInfo::canonical_tool_name) - .collect::>(); for tool in deferred_mcp_tools { - if !directly_registered_mcp_tools.contains(&tool.canonical_tool_name()) { - builder.register_handler(Arc::new(McpHandler::new(tool.clone()))); - } + handlers.push(Arc::new(McpHandler::new(tool.clone()))); } } - for bundle in params.extension_tool_bundles.iter().cloned() { - builder.register_tool_bundle(bundle); + for tool in params.dynamic_tools { + let Some(handler) = DynamicToolHandler::new(tool).map(Arc::new) else { + tracing::error!( + "Failed to convert dynamic tool {:?} to OpenAI tool", + tool.name + ); + continue; + }; + + handlers.push(handler); } - builder + handlers } fn compare_code_mode_tools( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index a4132391a2..7c876d3738 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -1579,20 +1579,7 @@ fn search_tool_description_lists_each_mcp_source_once() { let (tools, registry) = build_specs( &tools_config, - Some(HashMap::from([ - ( - ToolName::namespaced("mcp__codex_apps__calendar", "_create_event"), - mcp_tool( - "calendar_create_event", - "Create calendar event", - serde_json::json!({"type": "object"}), - ), - ), - ( - ToolName::namespaced("mcp__rmcp__", "echo"), - mcp_tool("echo", "Echo", serde_json::json!({"type": "object"})), - ), - ])), + /*mcp_tools*/ None, Some(vec![ deferred_mcp_tool( "_create_event", @@ -1810,32 +1797,15 @@ fn search_tool_registers_for_deferred_dynamic_tools() { panic!("expected tool_search tool"); }; assert!(description.contains("- Dynamic tools: Tools provided by the current Codex thread.")); - assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME, "codex_app"]); - assert_eq!( - tools - .iter() - .filter(|tool| tool.name() == "codex_app") - .count(), - 1 - ); - assert_eq!( - namespace_function_names(&tools, "codex_app"), - vec![ - "automation_update".to_string(), - "automation_list".to_string() - ] - ); - for tool_name in ["automation_update", "automation_list"] { - let dynamic_tool = find_namespace_function_tool(&tools, "codex_app", tool_name); - assert_eq!(dynamic_tool.defer_loading, Some(true)); - } + assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME]); + assert_lacks_tool_name(&tools, "codex_app"); assert!(registry.has_handler(&ToolName::plain(TOOL_SEARCH_TOOL_NAME))); assert!(registry.has_handler(&ToolName::namespaced("codex_app", "automation_update"))); assert!(registry.has_handler(&ToolName::namespaced("codex_app", "automation_list"))); } #[test] -fn search_tool_keeps_plain_deferred_dynamic_tools_when_namespace_tools_are_disabled() { +fn search_tool_is_hidden_for_deferred_dynamic_tools_when_namespace_tools_are_disabled() { let model_info = search_capable_model_info(); let mut features = Features::with_defaults(); features.enable(Feature::ToolSearch); @@ -1875,9 +1845,12 @@ fn search_tool_keeps_plain_deferred_dynamic_tools_when_namespace_tools_are_disab &dynamic_tools, ); - assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME, "plain_dynamic"]); + assert_lacks_tool_name(&tools, TOOL_SEARCH_TOOL_NAME); assert_lacks_tool_name(&tools, "codex_app"); - assert!(registry.has_handler(&ToolName::plain(TOOL_SEARCH_TOOL_NAME))); + assert_lacks_tool_name(&tools, "plain_dynamic"); + assert!(!registry.has_handler(&ToolName::plain(TOOL_SEARCH_TOOL_NAME))); + assert!(registry.has_handler(&ToolName::namespaced("codex_app", "automation_update"))); + assert!(registry.has_handler(&ToolName::plain("plain_dynamic"))); } #[test] diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index 0086542859..52f92b281d 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -365,7 +365,7 @@ async fn assert_model_tools( .iter() .map(ToolSpec::name) .collect::>(); - assert_eq!(&tool_names, &expected_tools,); + assert_eq!(&tool_names, &expected_tools); } async fn assert_default_model_tools( @@ -396,14 +396,14 @@ async fn test_build_specs_gpt5_codex_default() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -421,14 +421,14 @@ async fn test_build_specs_gpt51_codex_default() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -448,14 +448,14 @@ async fn test_build_specs_gpt5_codex_unified_exec_web_search() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -475,14 +475,14 @@ async fn test_build_specs_gpt51_codex_unified_exec_web_search() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -500,14 +500,14 @@ async fn test_gpt_5_1_codex_max_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -525,14 +525,14 @@ async fn test_codex_5_1_mini_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -550,14 +550,14 @@ async fn test_gpt_5_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -575,14 +575,14 @@ async fn test_gpt_5_1_defaults() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -602,14 +602,14 @@ async fn test_gpt_5_1_codex_max_unified_exec_web_search() { "update_plan", "request_user_input", "apply_patch", - "web_search", - "image_generation", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", + "image_generation", ], ) .await; @@ -841,7 +841,7 @@ async fn request_plugin_install_requires_apps_and_plugins_features() { } #[tokio::test] -async fn search_tool_description_handles_no_enabled_mcp_tools() { +async fn search_tool_is_hidden_without_deferred_tools() { let model_info = search_capable_model_info().await; let mut features = Features::with_defaults(); features.enable(Feature::Apps); @@ -865,13 +865,11 @@ async fn search_tool_description_handles_no_enabled_mcp_tools() { &[], ) .build(); - let search_tool = find_tool(&tools, TOOL_SEARCH_TOOL_NAME); - let ToolSpec::ToolSearch { description, .. } = search_tool else { - panic!("expected tool_search tool"); - }; - - assert!(description.contains("None currently enabled.")); - assert!(!description.contains("{{source_descriptions}}")); + assert!( + !tools + .iter() + .any(|tool| tool.name() == TOOL_SEARCH_TOOL_NAME) + ); } #[tokio::test] @@ -1003,7 +1001,7 @@ async fn search_tool_registers_namespaced_mcp_tool_aliases() { } #[tokio::test] -async fn tool_search_entries_skip_namespace_outputs_when_namespace_tools_are_disabled() { +async fn tool_search_entries_skip_mcp_outputs_when_namespace_tools_are_disabled() { let model_info = search_capable_model_info().await; let mut features = Features::with_defaults(); features.enable(Feature::ToolSearch); @@ -1048,11 +1046,15 @@ async fn tool_search_entries_skip_namespace_outputs_when_namespace_tools_are_dis .map(|entry| entry.output) .collect::>(); - assert_eq!(outputs.len(), 1); - match &outputs[0] { - LoadableToolSpec::Function(tool) => assert_eq!(tool.name, "plain_dynamic"), - LoadableToolSpec::Namespace(_) => panic!("namespace tool_search output should be hidden"), - } + assert_eq!(outputs.len(), 2); + assert!(outputs.iter().any(|output| match output { + LoadableToolSpec::Namespace(namespace) => namespace.name == "codex_app", + LoadableToolSpec::Function(_) => false, + })); + assert!(outputs.iter().any(|output| match output { + LoadableToolSpec::Function(tool) => tool.name == "plain_dynamic", + LoadableToolSpec::Namespace(_) => false, + })); } #[tokio::test] diff --git a/codex-rs/core/src/tools/tool_search_entry.rs b/codex-rs/core/src/tools/tool_search_entry.rs index 6a8ccd6402..cbaa4d87c9 100644 --- a/codex-rs/core/src/tools/tool_search_entry.rs +++ b/codex-rs/core/src/tools/tool_search_entry.rs @@ -63,11 +63,7 @@ pub(crate) fn build_tool_search_entries_for_config( } else { None }; - let dynamic_tools = dynamic_tools - .iter() - .filter(|tool| config.namespace_tools || tool.namespace.is_none()) - .cloned() - .collect::>(); + let dynamic_tools = dynamic_tools.to_vec(); build_tool_search_entries(mcp_tools, &dynamic_tools) } diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index b81bb06bb9..5b861ab219 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/codex-rs/core/tests/suite/prompt_caching.rs @@ -179,13 +179,13 @@ async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> { "update_plan", "request_user_input", "apply_patch", - "web_search", "view_image", "spawn_agent", "send_input", "resume_agent", "wait_agent", "close_agent", + "web_search", ]); let body0 = req1.single_request().body_json(); diff --git a/codex-rs/exec-server/BUILD.bazel b/codex-rs/exec-server/BUILD.bazel index 224536da8e..e94a5c0043 100644 --- a/codex-rs/exec-server/BUILD.bazel +++ b/codex-rs/exec-server/BUILD.bazel @@ -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", ], diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 936fa412f1..09a9a71ea0 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -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 } diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 81664eaca0..1eaf6e69eb 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -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 diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 8ca1eb0280..4bdc09a80e 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -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); diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index d8c147127c..872f16ce32 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -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; diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto new file mode 100644 index 0000000000..46527d80cc --- /dev/null +++ b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto @@ -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 {} diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs new file mode 100644 index 0000000000..072a003dac --- /dev/null +++ b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs @@ -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, +} +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, +} +#[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 {} diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs new file mode 100644 index 0000000000..bce787cfc2 --- /dev/null +++ b/codex-rs/exec-server/src/relay.rs @@ -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) -> 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 { + 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 { + 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 { + 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 { + frame.encode_to_vec() +} + +fn decode_relay_message_frame(payload: &[u8]) -> Result { + RelayMessageFrame::decode(payload) + .map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}"))) +} + +fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> { + serde_json::to_vec(message).map_err(ExecServerError::Json) +} + +pub(crate) fn harness_connection_from_websocket( + stream: WebSocketStream, + 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( + stream: WebSocketStream, + 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::>(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 = 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, + disconnected_tx: watch::Sender, +} + +impl VirtualStream { + async fn disconnect(self, reason: Option) { + 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>, +) -> 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, + } +} diff --git a/codex-rs/exec-server/src/relay_proto.rs b/codex-rs/exec-server/src/relay_proto.rs new file mode 100644 index 0000000000..b8a938b8c7 --- /dev/null +++ b/codex-rs/exec-server/src/relay_proto.rs @@ -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; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index bb22105c19..32c0d5bc8e 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -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}"); diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs new file mode 100644 index 0000000000..d228db9ee5 --- /dev/null +++ b/codex-rs/exec-server/tests/relay.rs @@ -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(®istry) + .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( + websocket: &mut WebSocketStream, + 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( + websocket: &mut WebSocketStream, + 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( + websocket: &mut WebSocketStream, + 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( + websocket: &mut WebSocketStream, + count: usize, +) -> Result> +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( + websocket: &mut WebSocketStream, +) -> 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 { + 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 { + 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(()) +} diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 955ae44c00..3be2404a9a 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -231,7 +231,8 @@ pub enum Feature { TuiAppServer, /// Prevent idle system sleep while a turn is actively running. PreventIdleSleep, - /// Enable workspace-specific owner nudge copy and prompts in the TUI. + /// Removed compatibility flag retained as a no-op now that workspace owner + /// usage nudges are always enabled. WorkspaceOwnerUsageNudge, /// Legacy rollout flag for Responses API WebSocket transport experiments. ResponsesWebsockets, @@ -1153,7 +1154,7 @@ pub const FEATURES: &[FeatureSpec] = &[ FeatureSpec { id: Feature::WorkspaceOwnerUsageNudge, key: "workspace_owner_usage_nudge", - stage: Stage::UnderDevelopment, + stage: Stage::Removed, default_enabled: false, }, FeatureSpec { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 9efc995313..38b6181221 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -2890,18 +2890,7 @@ impl ChatWidget { self.maybe_send_next_queued_input(); } - fn workspace_owner_usage_nudge_enabled(&self) -> bool { - self.config - .features - .enabled(Feature::WorkspaceOwnerUsageNudge) - } - fn on_rate_limit_error(&mut self, error_kind: RateLimitErrorKind, message: String) { - if !self.workspace_owner_usage_nudge_enabled() { - self.on_error(message); - return; - } - let rate_limit_reached_type = self.codex_rate_limit_reached_type.map(|kind| { if matches!(error_kind, RateLimitErrorKind::UsageLimit) { match kind { @@ -6919,9 +6908,7 @@ impl ChatWidget { } fn open_workspace_owner_nudge_prompt(&mut self, credit_type: AddCreditsNudgeCreditType) { - if !self.workspace_owner_usage_nudge_enabled() - || self.add_credits_nudge_email_in_flight.is_some() - { + if self.add_credits_nudge_email_in_flight.is_some() { return; } @@ -6969,10 +6956,6 @@ impl ChatWidget { &mut self, credit_type: AddCreditsNudgeCreditType, ) -> bool { - if !self.workspace_owner_usage_nudge_enabled() { - return false; - } - self.add_credits_nudge_email_in_flight = Some(credit_type); true } @@ -6985,9 +6968,6 @@ impl ChatWidget { .add_credits_nudge_email_in_flight .take() .unwrap_or(AddCreditsNudgeCreditType::Credits); - if !self.workspace_owner_usage_nudge_enabled() { - return; - } let message = match (credit_type, result) { (AddCreditsNudgeCreditType::Credits, Ok(AddCreditsNudgeEmailStatus::Sent)) => { "Workspace owner notified." diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index 6cc9abc704..4cc01d7591 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -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::(); 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(); @@ -796,57 +797,9 @@ async fn rate_limit_switch_prompt_popup_snapshot() { assert_chatwidget_snapshot!("rate_limit_switch_prompt_popup", popup); } -#[tokio::test] -async fn workspace_owner_usage_nudge_flag_disabled_keeps_generic_rate_limit_error() { - { - let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - let mut limits = snapshot(/*percent*/ 100.0); - limits.rate_limit_reached_type = - Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); - chat.on_rate_limit_snapshot(Some(limits)); - - chat.on_rate_limit_error( - RateLimitErrorKind::UsageLimit, - "Usage limit reached.".to_string(), - ); - let rendered = drain_insert_history(&mut rx) - .into_iter() - .map(|lines| lines_to_single_string(&lines)) - .collect::(); - assert!( - rendered.contains("Usage limit reached."), - "rendered: {rendered}" - ); - } - - { - let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - let mut limits = snapshot(/*percent*/ 100.0); - limits.rate_limit_reached_type = - Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); - chat.on_rate_limit_snapshot(Some(limits)); - - chat.on_rate_limit_error( - RateLimitErrorKind::UsageLimit, - "Usage limit reached.".to_string(), - ); - let popup = render_bottom_popup(&chat, /*width*/ 100); - assert!( - !popup.contains("Request a limit increase from your owner"), - "popup: {popup}" - ); - assert_no_owner_nudge_or_rate_limit_refresh(&mut rx); - } -} - -fn enable_workspace_owner_usage_nudge(chat: &mut ChatWidget) { - chat.set_feature_enabled(Feature::WorkspaceOwnerUsageNudge, /*enabled*/ true); -} - #[tokio::test] async fn workspace_member_credits_depleted_prompts_and_sends_credits() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberCreditsDepleted); chat.on_rate_limit_snapshot(Some(limits)); @@ -866,7 +819,6 @@ async fn workspace_member_credits_depleted_prompts_and_sends_credits() { #[tokio::test] async fn workspace_member_usage_limit_prompts_and_sends_usage_limit() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); chat.on_rate_limit_snapshot(Some(limits)); @@ -886,7 +838,6 @@ async fn workspace_member_usage_limit_prompts_and_sends_usage_limit() { #[tokio::test] async fn header_rate_limit_snapshot_preserves_member_limit_type_for_error_prompt() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut usage_limits = snapshot(/*percent*/ 100.0); usage_limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); @@ -916,7 +867,6 @@ async fn header_rate_limit_snapshot_preserves_member_limit_type_for_error_prompt #[tokio::test] async fn usage_limit_error_remaps_stale_member_credits_state_to_usage_limit_prompt() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberCreditsDepleted); chat.on_rate_limit_snapshot(Some(limits)); @@ -953,7 +903,6 @@ async fn workspace_owner_limit_states_do_not_prompt_for_owner_nudge() { ), ] { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(limit_type); chat.on_rate_limit_snapshot(Some(limits)); @@ -983,7 +932,6 @@ async fn workspace_owner_limit_states_render_state_specific_messages() { let mut rendered_cases = Vec::new(); for (limit_type, error_kind, expected) in cases { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(limit_type); chat.on_rate_limit_snapshot(Some(limits)); @@ -1006,7 +954,6 @@ async fn workspace_owner_limit_states_render_state_specific_messages() { #[tokio::test] async fn missing_rate_limit_reached_type_does_not_prompt_or_refresh() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); chat.on_rate_limit_snapshot(Some(snapshot(/*percent*/ 100.0))); chat.on_rate_limit_error( @@ -1021,7 +968,6 @@ async fn missing_rate_limit_reached_type_does_not_prompt_or_refresh() { #[tokio::test] async fn workspace_owner_nudge_default_no_dismisses_without_sending() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberCreditsDepleted); chat.on_rate_limit_snapshot(Some(limits)); @@ -1038,7 +984,6 @@ async fn workspace_owner_nudge_default_no_dismisses_without_sending() { #[tokio::test] async fn workspace_owner_nudge_reappears_after_dismissing_no() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); let mut limits = snapshot(/*percent*/ 100.0); limits.rate_limit_reached_type = Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached); chat.on_rate_limit_snapshot(Some(limits)); @@ -1081,7 +1026,6 @@ async fn workspace_owner_credits_nudge_completion_renders_feedback() { let mut rendered_cases = Vec::new(); for (result, expected) in cases { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); chat.start_add_credits_nudge_email_request(AddCreditsNudgeCreditType::Credits); chat.finish_add_credits_nudge_email_request(result); let rendered = drain_insert_history(&mut rx) @@ -1118,7 +1062,6 @@ async fn workspace_owner_usage_limit_nudge_completion_renders_feedback() { let mut rendered_cases = Vec::new(); for (result, expected) in cases { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - enable_workspace_owner_usage_nudge(&mut chat); chat.start_add_credits_nudge_email_request(AddCreditsNudgeCreditType::UsageLimit); chat.finish_add_credits_nudge_email_request(result); let rendered = drain_insert_history(&mut rx)