codex_extension_api/
registry.rs

1use std::sync::Arc;
2
3use codex_protocol::protocol::ReviewDecision;
4
5use crate::ApprovalReviewContributor;
6use crate::ConfigContributor;
7use crate::ContextContributor;
8use crate::ExtensionData;
9use crate::ExtensionEventSink;
10use crate::McpServerContributor;
11use crate::NoopExtensionEventSink;
12use crate::SkillInvocationContributor;
13use crate::ThreadLifecycleContributor;
14use crate::TokenUsageContributor;
15use crate::ToolContributor;
16use crate::ToolLifecycleContributor;
17use crate::TurnInputContributor;
18use crate::TurnItemContributor;
19use crate::TurnLifecycleContributor;
20
21/// Mutable registry used while hosts register typed runtime contributions.
22pub struct ExtensionRegistryBuilder<C: Sync> {
23    event_sink: Arc<dyn ExtensionEventSink>,
24    thread_lifecycle_contributors: Vec<Arc<dyn ThreadLifecycleContributor<C>>>,
25    turn_lifecycle_contributors: Vec<Arc<dyn TurnLifecycleContributor>>,
26    config_contributors: Vec<Arc<dyn ConfigContributor<C>>>,
27    token_usage_contributors: Vec<Arc<dyn TokenUsageContributor>>,
28    skill_invocation_contributors: Vec<Arc<dyn SkillInvocationContributor>>,
29    context_contributors: Vec<Arc<dyn ContextContributor>>,
30    mcp_server_contributors: Vec<Arc<dyn McpServerContributor<C>>>,
31    turn_input_contributors: Vec<Arc<dyn TurnInputContributor>>,
32    tool_contributors: Vec<Arc<dyn ToolContributor>>,
33    tool_lifecycle_contributors: Vec<Arc<dyn ToolLifecycleContributor>>,
34    turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
35    approval_review_contributors: Vec<Arc<dyn ApprovalReviewContributor>>,
36}
37
38impl<C: Sync> Default for ExtensionRegistryBuilder<C> {
39    fn default() -> Self {
40        Self {
41            event_sink: Arc::new(NoopExtensionEventSink),
42            thread_lifecycle_contributors: Vec::new(),
43            turn_lifecycle_contributors: Vec::new(),
44            config_contributors: Vec::new(),
45            token_usage_contributors: Vec::new(),
46            skill_invocation_contributors: Vec::new(),
47            approval_review_contributors: Vec::new(),
48            context_contributors: Vec::new(),
49            mcp_server_contributors: Vec::new(),
50            turn_input_contributors: Vec::new(),
51            tool_contributors: Vec::new(),
52            tool_lifecycle_contributors: Vec::new(),
53            turn_item_contributors: Vec::new(),
54        }
55    }
56}
57
58impl<C: Sync> ExtensionRegistryBuilder<C> {
59    /// Creates an empty registry builder.
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    /// Creates an empty registry builder with a host-provided event sink.
65    pub fn with_event_sink(event_sink: Arc<dyn ExtensionEventSink>) -> Self {
66        Self {
67            event_sink,
68            ..Self::default()
69        }
70    }
71
72    /// Returns the host event sink to pass into extension constructors.
73    pub fn event_sink(&self) -> Arc<dyn ExtensionEventSink> {
74        Arc::clone(&self.event_sink)
75    }
76
77    /// Registers one approval-review contributor.
78    pub fn approval_review_contributor(&mut self, contributor: Arc<dyn ApprovalReviewContributor>) {
79        self.approval_review_contributors.push(contributor);
80    }
81
82    /// Registers one thread-lifecycle contributor.
83    pub fn thread_lifecycle_contributor(
84        &mut self,
85        contributor: Arc<dyn ThreadLifecycleContributor<C>>,
86    ) {
87        self.thread_lifecycle_contributors.push(contributor);
88    }
89
90    /// Registers one turn-lifecycle contributor.
91    pub fn turn_lifecycle_contributor(&mut self, contributor: Arc<dyn TurnLifecycleContributor>) {
92        self.turn_lifecycle_contributors.push(contributor);
93    }
94
95    /// Registers one config contributor.
96    pub fn config_contributor(&mut self, contributor: Arc<dyn ConfigContributor<C>>) {
97        self.config_contributors.push(contributor);
98    }
99
100    /// Registers one token-usage contributor.
101    pub fn token_usage_contributor(&mut self, contributor: Arc<dyn TokenUsageContributor>) {
102        self.token_usage_contributors.push(contributor);
103    }
104
105    /// Registers one skill-invocation contributor.
106    pub fn skill_invocation_contributor(
107        &mut self,
108        contributor: Arc<dyn SkillInvocationContributor>,
109    ) {
110        self.skill_invocation_contributors.push(contributor);
111    }
112
113    /// Registers one prompt contributor.
114    pub fn prompt_contributor(&mut self, contributor: Arc<dyn ContextContributor>) {
115        self.context_contributors.push(contributor);
116    }
117
118    /// Registers one runtime MCP server contributor.
119    pub fn mcp_server_contributor(&mut self, contributor: Arc<dyn McpServerContributor<C>>) {
120        self.mcp_server_contributors.push(contributor);
121    }
122
123    /// Registers one turn-input contributor.
124    pub fn turn_input_contributor(&mut self, contributor: Arc<dyn TurnInputContributor>) {
125        self.turn_input_contributors.push(contributor);
126    }
127
128    /// Registers one native tool contributor.
129    pub fn tool_contributor(&mut self, contributor: Arc<dyn ToolContributor>) {
130        self.tool_contributors.push(contributor);
131    }
132
133    /// Registers one tool-lifecycle contributor.
134    pub fn tool_lifecycle_contributor(&mut self, contributor: Arc<dyn ToolLifecycleContributor>) {
135        self.tool_lifecycle_contributors.push(contributor);
136    }
137
138    /// Registers one ordered turn-item contributor.
139    pub fn turn_item_contributor(&mut self, contributor: Arc<dyn TurnItemContributor>) {
140        self.turn_item_contributors.push(contributor);
141    }
142
143    /// Finishes construction and returns the immutable registry.
144    pub fn build(self) -> ExtensionRegistry<C> {
145        ExtensionRegistry {
146            event_sink: self.event_sink,
147            thread_lifecycle_contributors: self.thread_lifecycle_contributors,
148            turn_lifecycle_contributors: self.turn_lifecycle_contributors,
149            config_contributors: self.config_contributors,
150            token_usage_contributors: self.token_usage_contributors,
151            skill_invocation_contributors: self.skill_invocation_contributors,
152            approval_review_contributors: self.approval_review_contributors,
153            context_contributors: self.context_contributors,
154            mcp_server_contributors: self.mcp_server_contributors,
155            turn_input_contributors: self.turn_input_contributors,
156            tool_contributors: self.tool_contributors,
157            tool_lifecycle_contributors: self.tool_lifecycle_contributors,
158            turn_item_contributors: self.turn_item_contributors,
159        }
160    }
161}
162
163/// Immutable typed registry produced after extensions are installed.
164pub struct ExtensionRegistry<C: Sync> {
165    event_sink: Arc<dyn ExtensionEventSink>,
166    thread_lifecycle_contributors: Vec<Arc<dyn ThreadLifecycleContributor<C>>>,
167    turn_lifecycle_contributors: Vec<Arc<dyn TurnLifecycleContributor>>,
168    config_contributors: Vec<Arc<dyn ConfigContributor<C>>>,
169    token_usage_contributors: Vec<Arc<dyn TokenUsageContributor>>,
170    skill_invocation_contributors: Vec<Arc<dyn SkillInvocationContributor>>,
171    context_contributors: Vec<Arc<dyn ContextContributor>>,
172    mcp_server_contributors: Vec<Arc<dyn McpServerContributor<C>>>,
173    turn_input_contributors: Vec<Arc<dyn TurnInputContributor>>,
174    tool_contributors: Vec<Arc<dyn ToolContributor>>,
175    tool_lifecycle_contributors: Vec<Arc<dyn ToolLifecycleContributor>>,
176    turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
177    approval_review_contributors: Vec<Arc<dyn ApprovalReviewContributor>>,
178}
179
180impl<C: Sync> ExtensionRegistry<C> {
181    /// Returns the host event sink retained by this registry.
182    pub fn event_sink(&self) -> Arc<dyn ExtensionEventSink> {
183        Arc::clone(&self.event_sink)
184    }
185
186    /// Returns the registered thread-lifecycle contributors.
187    pub fn thread_lifecycle_contributors(&self) -> &[Arc<dyn ThreadLifecycleContributor<C>>] {
188        &self.thread_lifecycle_contributors
189    }
190
191    /// Returns the registered turn-lifecycle contributors.
192    pub fn turn_lifecycle_contributors(&self) -> &[Arc<dyn TurnLifecycleContributor>] {
193        &self.turn_lifecycle_contributors
194    }
195
196    /// Returns the registered config contributors.
197    pub fn config_contributors(&self) -> &[Arc<dyn ConfigContributor<C>>] {
198        &self.config_contributors
199    }
200
201    /// Returns the registered token-usage contributors.
202    pub fn token_usage_contributors(&self) -> &[Arc<dyn TokenUsageContributor>] {
203        &self.token_usage_contributors
204    }
205
206    /// Returns the registered skill-invocation contributors.
207    pub fn skill_invocation_contributors(&self) -> &[Arc<dyn SkillInvocationContributor>] {
208        &self.skill_invocation_contributors
209    }
210
211    /// Claims the first rendered approval-review prompt accepted by an
212    /// installed contributor.
213    pub async fn approval_review(
214        &self,
215        session_store: &ExtensionData,
216        thread_store: &ExtensionData,
217        prompt: &str,
218    ) -> Option<ReviewDecision> {
219        for contributor in &self.approval_review_contributors {
220            if let Some(decision) = contributor
221                .contribute(session_store, thread_store, prompt)
222                .await
223            {
224                return Some(decision);
225            }
226        }
227
228        None
229    }
230
231    /// Returns the registered prompt contributors.
232    pub fn context_contributors(&self) -> &[Arc<dyn ContextContributor>] {
233        &self.context_contributors
234    }
235
236    /// Returns the registered runtime MCP server contributors.
237    pub fn mcp_server_contributors(&self) -> &[Arc<dyn McpServerContributor<C>>] {
238        &self.mcp_server_contributors
239    }
240
241    /// Returns the registered turn-input contributors.
242    pub fn turn_input_contributors(&self) -> &[Arc<dyn TurnInputContributor>] {
243        &self.turn_input_contributors
244    }
245
246    /// Returns the registered native tool contributors.
247    pub fn tool_contributors(&self) -> &[Arc<dyn ToolContributor>] {
248        &self.tool_contributors
249    }
250
251    /// Returns the registered tool-lifecycle contributors.
252    pub fn tool_lifecycle_contributors(&self) -> &[Arc<dyn ToolLifecycleContributor>] {
253        &self.tool_lifecycle_contributors
254    }
255
256    /// Returns the registered ordered turn-item contributors.
257    pub fn turn_item_contributors(&self) -> &[Arc<dyn TurnItemContributor>] {
258        &self.turn_item_contributors
259    }
260}
261
262/// Creates an empty shared registry for hosts that do not register contributions.
263pub fn empty_extension_registry<C: Sync>() -> Arc<ExtensionRegistry<C>> {
264    Arc::new(ExtensionRegistryBuilder::new().build())
265}
codex_extension_api/registry.rs

codex_extension_api/
registry.rs
