mirror of
https://github.com/openai/codex.git
synced 2026-05-25 13:34:51 +00:00
9531e932ef
## Why Extension lifecycle hooks sit on the host/extension boundary, but the current trait surface only allows synchronous callbacks. That forces extensions that need to seed, rehydrate, observe, or flush extension-owned state during thread and turn transitions to either block inside the callback or move async work into separate host plumbing. This PR makes those lifecycle callbacks awaitable so extension implementations can perform async work directly at the lifecycle point where the host already has the relevant session, thread, or turn stores available. ## What changed - Makes `ThreadLifecycleContributor` and `TurnLifecycleContributor` async in `codex-extension-api`. - Awaits thread start/resume/stop and turn start/stop/abort lifecycle callbacks from `codex-core`. - Updates the guardian and memories extensions to implement the async lifecycle trait surface. - Updates the existing lifecycle tests to use async contributor implementations. - Adds `async-trait` to the crates that now expose or implement these async object-safe lifecycle traits. ## Testing - Existing `codex-core` lifecycle tests were updated to cover async implementations for thread stop and turn abort ordering.
72 lines
2.2 KiB
Rust
72 lines
2.2 KiB
Rust
use std::sync::Arc;
|
|
|
|
use codex_core::config::Config;
|
|
use codex_extension_api::AgentSpawnFuture;
|
|
use codex_extension_api::AgentSpawner;
|
|
use codex_extension_api::ExtensionRegistryBuilder;
|
|
use codex_extension_api::ThreadLifecycleContributor;
|
|
use codex_extension_api::ThreadStartInput;
|
|
use codex_protocol::ThreadId;
|
|
|
|
/// Guardian extension dependencies supplied by the host at construction time.
|
|
#[derive(Clone, Debug)]
|
|
pub struct GuardianExtension<S> {
|
|
agent_spawner: S,
|
|
}
|
|
|
|
impl<S> GuardianExtension<S> {
|
|
/// Creates a guardian extension with its host-provided agent spawn helper.
|
|
pub fn new(agent_spawner: S) -> Self {
|
|
Self { agent_spawner }
|
|
}
|
|
|
|
/// Delegates one guardian-owned subagent spawn request to the host helper.
|
|
pub fn spawn_subagent<'a, R>(
|
|
&'a self,
|
|
forked_from_thread_id: ThreadId,
|
|
request: R,
|
|
) -> AgentSpawnFuture<'a, <S as AgentSpawner<R>>::Spawned, <S as AgentSpawner<R>>::Error>
|
|
where
|
|
S: AgentSpawner<R>,
|
|
{
|
|
self.agent_spawner
|
|
.spawn_subagent(forked_from_thread_id, request)
|
|
}
|
|
}
|
|
|
|
/// Thread-local guardian state captured when the host starts a thread.
|
|
#[derive(Clone, Copy, Debug)]
|
|
pub struct GuardianThreadContext {
|
|
forked_from_thread_id: ThreadId,
|
|
}
|
|
|
|
impl GuardianThreadContext {
|
|
/// Returns the thread that future guardian subagents should fork from by default.
|
|
pub fn forked_from_thread_id(&self) -> ThreadId {
|
|
self.forked_from_thread_id
|
|
}
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl<S> ThreadLifecycleContributor<Config> for GuardianExtension<S>
|
|
where
|
|
S: Send + Sync,
|
|
{
|
|
async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
|
|
let Ok(forked_from_thread_id) = ThreadId::from_string(input.thread_store.level_id()) else {
|
|
return;
|
|
};
|
|
input.thread_store.insert(GuardianThreadContext {
|
|
forked_from_thread_id,
|
|
});
|
|
}
|
|
}
|
|
|
|
/// Installs the guardian contributors into the extension registry.
|
|
pub fn install<S>(registry: &mut ExtensionRegistryBuilder<Config>, agent_spawner: S)
|
|
where
|
|
S: Send + Sync + 'static,
|
|
{
|
|
registry.thread_lifecycle_contributor(Arc::new(GuardianExtension::new(agent_spawner)));
|
|
}
|