From 62839fec5df9dfd9442ffd7226a789c63d7451c7 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 3 Aug 2026 11:54:15 +0000 Subject: [PATCH] Store the extension registry directly in its builder (#36744) ## What changed Make `ExtensionRegistryBuilder` wrap an `ExtensionRegistry` and register contributors directly on it. Return the wrapped registry from `build()` instead of copying each field into a new value. GitOrigin-RevId: ae04de99c17ea951156f9892151b4196c70026bf --- codex-rs/ext/extension-api/src/registry.rs | 95 +++++++++------------- 1 file changed, 37 insertions(+), 58 deletions(-) diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index 842155ea29..f34c88cc3b 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -20,37 +20,27 @@ use crate::TurnLifecycleContributor; /// Mutable registry used while hosts register typed runtime contributions. pub struct ExtensionRegistryBuilder { - event_sink: Arc, - thread_lifecycle_contributors: Vec>>, - turn_lifecycle_contributors: Vec>, - config_contributors: Vec>>, - token_usage_contributors: Vec>, - skill_invocation_contributors: Vec>, - context_contributors: Vec>, - mcp_server_contributors: Vec>>, - turn_input_contributors: Vec>, - tool_contributors: Vec>, - tool_lifecycle_contributors: Vec>, - turn_item_contributors: Vec>, - approval_review_contributors: Vec>, + registry: ExtensionRegistry, } impl Default for ExtensionRegistryBuilder { fn default() -> Self { Self { - event_sink: Arc::new(NoopExtensionEventSink), - thread_lifecycle_contributors: Vec::new(), - turn_lifecycle_contributors: Vec::new(), - config_contributors: Vec::new(), - token_usage_contributors: Vec::new(), - skill_invocation_contributors: Vec::new(), - approval_review_contributors: Vec::new(), - context_contributors: Vec::new(), - mcp_server_contributors: Vec::new(), - turn_input_contributors: Vec::new(), - tool_contributors: Vec::new(), - tool_lifecycle_contributors: Vec::new(), - turn_item_contributors: Vec::new(), + registry: ExtensionRegistry { + event_sink: Arc::new(NoopExtensionEventSink), + thread_lifecycle_contributors: Vec::new(), + turn_lifecycle_contributors: Vec::new(), + config_contributors: Vec::new(), + token_usage_contributors: Vec::new(), + skill_invocation_contributors: Vec::new(), + approval_review_contributors: Vec::new(), + context_contributors: Vec::new(), + mcp_server_contributors: Vec::new(), + turn_input_contributors: Vec::new(), + tool_contributors: Vec::new(), + tool_lifecycle_contributors: Vec::new(), + turn_item_contributors: Vec::new(), + }, } } } @@ -63,20 +53,19 @@ impl ExtensionRegistryBuilder { /// Creates an empty registry builder with a host-provided event sink. pub fn with_event_sink(event_sink: Arc) -> Self { - Self { - event_sink, - ..Self::default() - } + let mut builder = Self::default(); + builder.registry.event_sink = event_sink; + builder } /// Returns the host event sink to pass into extension constructors. pub fn event_sink(&self) -> Arc { - Arc::clone(&self.event_sink) + Arc::clone(&self.registry.event_sink) } /// Registers one approval-review contributor. pub fn approval_review_contributor(&mut self, contributor: Arc) { - self.approval_review_contributors.push(contributor); + self.registry.approval_review_contributors.push(contributor); } /// Registers one thread-lifecycle contributor. @@ -84,22 +73,24 @@ impl ExtensionRegistryBuilder { &mut self, contributor: Arc>, ) { - self.thread_lifecycle_contributors.push(contributor); + self.registry + .thread_lifecycle_contributors + .push(contributor); } /// Registers one turn-lifecycle contributor. pub fn turn_lifecycle_contributor(&mut self, contributor: Arc) { - self.turn_lifecycle_contributors.push(contributor); + self.registry.turn_lifecycle_contributors.push(contributor); } /// Registers one config contributor. pub fn config_contributor(&mut self, contributor: Arc>) { - self.config_contributors.push(contributor); + self.registry.config_contributors.push(contributor); } /// Registers one token-usage contributor. pub fn token_usage_contributor(&mut self, contributor: Arc) { - self.token_usage_contributors.push(contributor); + self.registry.token_usage_contributors.push(contributor); } /// Registers one skill-invocation contributor. @@ -107,56 +98,44 @@ impl ExtensionRegistryBuilder { &mut self, contributor: Arc, ) { - self.skill_invocation_contributors.push(contributor); + self.registry + .skill_invocation_contributors + .push(contributor); } /// Registers one prompt contributor. pub fn prompt_contributor(&mut self, contributor: Arc) { - self.context_contributors.push(contributor); + self.registry.context_contributors.push(contributor); } /// Registers one runtime MCP server contributor. pub fn mcp_server_contributor(&mut self, contributor: Arc>) { - self.mcp_server_contributors.push(contributor); + self.registry.mcp_server_contributors.push(contributor); } /// Registers one turn-input contributor. pub fn turn_input_contributor(&mut self, contributor: Arc) { - self.turn_input_contributors.push(contributor); + self.registry.turn_input_contributors.push(contributor); } /// Registers one native tool contributor. pub fn tool_contributor(&mut self, contributor: Arc) { - self.tool_contributors.push(contributor); + self.registry.tool_contributors.push(contributor); } /// Registers one tool-lifecycle contributor. pub fn tool_lifecycle_contributor(&mut self, contributor: Arc) { - self.tool_lifecycle_contributors.push(contributor); + self.registry.tool_lifecycle_contributors.push(contributor); } /// Registers one ordered turn-item contributor. pub fn turn_item_contributor(&mut self, contributor: Arc) { - self.turn_item_contributors.push(contributor); + self.registry.turn_item_contributors.push(contributor); } /// Finishes construction and returns the immutable registry. pub fn build(self) -> ExtensionRegistry { - ExtensionRegistry { - event_sink: self.event_sink, - thread_lifecycle_contributors: self.thread_lifecycle_contributors, - turn_lifecycle_contributors: self.turn_lifecycle_contributors, - config_contributors: self.config_contributors, - token_usage_contributors: self.token_usage_contributors, - skill_invocation_contributors: self.skill_invocation_contributors, - approval_review_contributors: self.approval_review_contributors, - context_contributors: self.context_contributors, - mcp_server_contributors: self.mcp_server_contributors, - turn_input_contributors: self.turn_input_contributors, - tool_contributors: self.tool_contributors, - tool_lifecycle_contributors: self.tool_lifecycle_contributors, - turn_item_contributors: self.turn_item_contributors, - } + self.registry } }