Skip to main content

codex_extension_api/
contributors.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use codex_context_fragments::ContextualUserFragment;
6use codex_protocol::items::TurnItem;
7use codex_protocol::protocol::ReviewDecision;
8use codex_protocol::protocol::TokenUsageInfo;
9use codex_tools::ToolCall;
10use codex_tools::ToolExecutor;
11
12use crate::ExtensionData;
13
14mod context;
15mod mcp;
16mod prompt;
17mod skill_invocation;
18mod thread_lifecycle;
19mod tool_lifecycle;
20mod turn_input;
21mod turn_lifecycle;
22mod world_state;
23
24pub use context::TurnContextContributionInput;
25pub use mcp::McpServerContribution;
26pub use mcp::McpServerContributionContext;
27pub use prompt::PromptFragment;
28pub use prompt::PromptSlot;
29pub use skill_invocation::SkillInvocationInput;
30pub use skill_invocation::SkillInvocationKind;
31pub use thread_lifecycle::ThreadIdleInput;
32pub use thread_lifecycle::ThreadOriginator;
33pub use thread_lifecycle::ThreadResumeInput;
34pub use thread_lifecycle::ThreadStartInput;
35pub use thread_lifecycle::ThreadStopInput;
36pub use tool_lifecycle::ToolCallOutcome;
37pub use tool_lifecycle::ToolCallSource;
38pub use tool_lifecycle::ToolFinishInput;
39pub use tool_lifecycle::ToolLifecycleFuture;
40pub use tool_lifecycle::ToolStartInput;
41pub use turn_input::TurnInputContext;
42pub use turn_input::TurnInputEnvironment;
43pub use turn_lifecycle::TurnAbortInput;
44pub use turn_lifecycle::TurnErrorInput;
45pub use turn_lifecycle::TurnStartInput;
46pub use turn_lifecycle::TurnStopInput;
47pub use world_state::PreviousWorldStateSection;
48pub use world_state::RenderedWorldStateFragment;
49pub use world_state::WorldStateContributionInput;
50pub use world_state::WorldStateSectionContribution;
51
52/// Boxed, sendable future returned by asynchronous extension contributors.
53pub type ExtensionFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
54
55/// Extension contribution that resolves runtime MCP servers from host config.
56///
57/// Contributors run in registration order. Later contributions for the same
58/// name replace earlier ones. Implementations must contribute only names they
59/// own and must apply any source-specific policy before returning a server.
60/// Thread-scoped resolution exposes the host-seeded thread inputs; global
61/// resolution exposes none and must not imply a local fallback. Thread inputs
62/// are frozen for the runtime and do not include lifecycle-contributor state.
63/// Auto-discovered plugin servers are resolved by the plugin manager. A
64/// thread-selected plugin contribution must carry its own package provenance.
65pub trait McpServerContributor<C: Sync>: Send + Sync {
66    /// Stable identity used for registration provenance and conflict diagnostics.
67    fn id(&self) -> &'static str;
68
69    fn contribute<'a>(
70        &'a self,
71        context: McpServerContributionContext<'a, C>,
72    ) -> ExtensionFuture<'a, Vec<McpServerContribution>>;
73}
74
75/// Extension contribution that adds prompt fragments during prompt assembly.
76///
77/// Implementations should use the method matching the scope needed by the
78/// fragment: thread/session context for stable inputs, and turn context for
79/// fragments that depend on turn-local host state.
80pub trait ContextContributor: Send + Sync {
81    /// Returns thread-scoped context using capabilities from the current sampling step.
82    fn contribute_thread_context<'a>(
83        &'a self,
84        session_store: &'a ExtensionData,
85        thread_store: &'a ExtensionData,
86        step_store: &'a ExtensionData,
87    ) -> ExtensionFuture<'a, Vec<PromptFragment>> {
88        Box::pin(async move {
89            let _self = self;
90            let _session_store = session_store;
91            let _thread_store = thread_store;
92            let _step_store = step_store;
93            Vec::new()
94        })
95    }
96
97    fn contribute_turn_context<'a>(
98        &'a self,
99        input: TurnContextContributionInput<'a>,
100    ) -> ExtensionFuture<'a, Vec<PromptFragment>> {
101        Box::pin(async move {
102            let _self = self;
103            let _input = input;
104            Vec::new()
105        })
106    }
107
108    fn contribute_world_state<'a>(
109        &'a self,
110        input: WorldStateContributionInput<'a>,
111    ) -> ExtensionFuture<'a, Vec<WorldStateSectionContribution>> {
112        Box::pin(async move {
113            let _self = self;
114            let _input = input;
115            Vec::new()
116        })
117    }
118}
119
120/// Contributor for host-owned thread lifecycle gates.
121///
122/// Implementations should use these callbacks to seed, rehydrate, or flush
123/// extension-private thread state. Heavy dependencies belong on the extension
124/// value created by the host, not in these inputs.
125pub trait ThreadLifecycleContributor<C: Sync>: Send + Sync {
126    /// Called after host startup has initialized the thread-scoped store.
127    fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> {
128        Box::pin(async move {
129            let _self = self;
130            let _input = input;
131        })
132    }
133
134    /// Called after the host constructs a runtime from persisted history.
135    fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> {
136        Box::pin(async move {
137            let _self = self;
138            let _input = input;
139        })
140    }
141
142    /// Called after the host has drained immediately pending thread work.
143    ///
144    /// Implementations may use host capabilities captured by the extension to
145    /// submit follow-up input. The host remains responsible for deciding
146    /// whether that input starts a turn, is queued, or is ignored.
147    fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> {
148        Box::pin(async move {
149            let _self = self;
150            let _input = input;
151        })
152    }
153
154    /// Called before the host drops the thread runtime and thread-scoped store.
155    fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> {
156        Box::pin(async move {
157            let _self = self;
158            let _input = input;
159        })
160    }
161}
162
163/// Contributor for host-owned turn lifecycle gates.
164///
165/// Implementations should use these callbacks to seed, observe, or clear
166/// extension-private turn state. The host exposes stable identifiers and
167/// extension stores instead of core runtime objects.
168pub trait TurnLifecycleContributor: Send + Sync {
169    /// Called after turn-scoped extension stores are created, before the task
170    /// for the turn starts running.
171    fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> {
172        Box::pin(async move {
173            let _self = self;
174            let _input = input;
175        })
176    }
177
178    /// Called before the host drops the completed turn runtime and turn store.
179    fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> {
180        Box::pin(async move {
181            let _self = self;
182            let _input = input;
183        })
184    }
185
186    /// Called after the host aborts a running turn.
187    fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> {
188        Box::pin(async move {
189            let _self = self;
190            let _input = input;
191        })
192    }
193
194    /// Called when the host observes an error for a running turn.
195    fn on_turn_error<'a>(&'a self, input: TurnErrorInput<'a>) -> ExtensionFuture<'a, ()> {
196        Box::pin(async move {
197            let _self = self;
198            let _input = input;
199        })
200    }
201}
202
203/// Extension contribution that can add turn-local model input.
204///
205/// Implementations should resolve only the model-visible input they own and
206/// must preserve authority boundaries for external resources. Expensive or
207/// host-specific dependencies belong on the extension value installed by the
208/// host, not in this input.
209pub trait TurnInputContributor: Send + Sync {
210    /// Returns additional contextual fragments for one submitted turn.
211    ///
212    /// `step_store` contains host capabilities bound to the sampling step that
213    /// will consume these fragments.
214    fn contribute<'a>(
215        &'a self,
216        input: TurnInputContext,
217        session_store: &'a ExtensionData,
218        thread_store: &'a ExtensionData,
219        turn_store: &'a ExtensionData,
220        step_store: &'a ExtensionData,
221    ) -> ExtensionFuture<'a, Vec<Box<dyn ContextualUserFragment + Send>>>;
222}
223
224/// Contributor for host-owned configuration changes.
225///
226/// Implementations should treat the supplied values as immutable before/after
227/// snapshots of the effective thread configuration.
228pub trait ConfigContributor<C>: Send + Sync {
229    /// Called after the host commits a changed thread configuration.
230    fn on_config_changed(
231        &self,
232        _session_store: &ExtensionData,
233        _thread_store: &ExtensionData,
234        _previous_config: &C,
235        _new_config: &C,
236    ) {
237    }
238}
239
240/// Contributor for token usage checkpoints reported by the model provider.
241///
242/// Implementations should keep this callback cheap. The host calls it after
243/// updating cached token usage and before emitting the corresponding client
244/// token-count notification.
245pub trait TokenUsageContributor: Send + Sync {
246    /// Called each time the host records token usage from a model response.
247    fn on_token_usage<'a>(
248        &'a self,
249        _session_store: &'a ExtensionData,
250        _thread_store: &'a ExtensionData,
251        _turn_store: &'a ExtensionData,
252        _token_usage: &'a TokenUsageInfo,
253    ) -> ExtensionFuture<'a, ()> {
254        Box::pin(async move {
255            let _self = self;
256            let _inputs = (_session_store, _thread_store, _turn_store, _token_usage);
257        })
258    }
259}
260
261/// Contributor for skill invocations observed by the host or an owning extension.
262///
263/// Implementations should treat the skill resource as an opaque identity and keep this callback
264/// cheap because it runs inline with skill loading or command dispatch.
265pub trait SkillInvocationContributor: Send + Sync {
266    /// Called after one explicit skill load or deduplicated implicit skill invocation is observed.
267    fn on_skill_invocation<'a>(
268        &'a self,
269        _input: SkillInvocationInput<'a>,
270    ) -> ExtensionFuture<'a, ()> {
271        Box::pin(async move {
272            let _self = self;
273            let _input = _input;
274        })
275    }
276}
277
278/// Extension contribution that exposes native tools owned by a feature.
279pub trait ToolContributor: Send + Sync {
280    /// Returns native tools bound to the supplied sampling-step capabilities.
281    fn tools(
282        &self,
283        session_store: &ExtensionData,
284        thread_store: &ExtensionData,
285        step_store: &ExtensionData,
286    ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>>;
287}
288
289/// Contributor for host-owned tool lifecycle gates.
290///
291/// Implementations should use these callbacks to observe tool execution without
292/// inspecting or rewriting tool input/output. Use `ToolContributor` for owning a
293/// tool implementation and hooks for policy that needs tool payloads.
294pub trait ToolLifecycleContributor: Send + Sync {
295    /// Called once the host has accepted a tool call for execution.
296    fn on_tool_start<'a>(&'a self, _input: ToolStartInput<'a>) -> ToolLifecycleFuture<'a> {
297        Box::pin(std::future::ready(()))
298    }
299
300    /// Called after the tool call returns, is blocked, fails, or is cancelled.
301    fn on_tool_finish<'a>(&'a self, _input: ToolFinishInput<'a>) -> ToolLifecycleFuture<'a> {
302        Box::pin(std::future::ready(()))
303    }
304}
305
306/// Extension contribution that can claim rendered approval-review prompts.
307pub trait ApprovalReviewContributor: Send + Sync {
308    fn contribute<'a>(
309        &'a self,
310        session_store: &'a ExtensionData,
311        thread_store: &'a ExtensionData,
312        prompt: &'a str,
313    ) -> ExtensionFuture<'a, Option<ReviewDecision>>;
314}
315
316/// Ordered post-processing contribution for one parsed turn item.
317///
318/// Implementations may mutate the item before it is emitted and may use the
319/// explicitly exposed thread- and turn-lifetime stores when they need durable
320/// extension-private state.
321pub trait TurnItemContributor: Send + Sync {
322    fn contribute<'a>(
323        &'a self,
324        thread_store: &'a ExtensionData,
325        turn_store: &'a ExtensionData,
326        item: &'a mut TurnItem,
327    ) -> ExtensionFuture<'a, Result<(), String>>;
328}