Add turn-scoped context contributions (#28911)

## Summary
- keep context injection on a single ContextContributor trait
- split context injection into thread-scoped and turn-scoped
contribution methods
- wire turn-scoped fragments into initial context assembly so extensions
can contribute context from turn-local state
This commit is contained in:
jif
2026-06-18 19:40:28 +02:00
committed by GitHub
parent 790213ded0
commit 9684ec25be
14 changed files with 307 additions and 33 deletions
+41 -3
View File
@@ -12,12 +12,14 @@ use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::PromptFragment;
use codex_extension_api::PromptSlot;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::TokenUsageContributor;
use codex_extension_api::ToolCall;
use codex_extension_api::ToolContributor;
use codex_extension_api::ToolExecutor;
use codex_extension_api::ToolLifecycleContributor;
use codex_extension_api::TurnContextContributionInput;
use codex_extension_api::TurnInputContext;
use codex_extension_api::TurnInputContributor;
use codex_extension_api::TurnItemContributor;
@@ -34,7 +36,7 @@ use pretty_assertions::assert_eq;
struct AllContributors;
impl ContextContributor for AllContributors {
fn contribute<'a>(
fn contribute_thread_context<'a>(
&'a self,
_session_store: &'a ExtensionData,
_thread_store: &'a ExtensionData,
@@ -147,7 +149,7 @@ async fn build_round_trips_every_contributor_category() {
struct NamedContextContributor(&'static str);
impl ContextContributor for NamedContextContributor {
fn contribute<'a>(
fn contribute_thread_context<'a>(
&'a self,
_session_store: &'a ExtensionData,
_thread_store: &'a ExtensionData,
@@ -158,6 +160,20 @@ impl ContextContributor for NamedContextContributor {
}
}
struct NamedTurnContextContributor(&'static str);
impl ContextContributor for NamedTurnContextContributor {
fn contribute_turn_context<'a>(
&'a self,
_input: TurnContextContributionInput<'a>,
) -> ExtensionFuture<'a, Vec<PromptFragment>> {
Box::pin(std::future::ready(vec![PromptFragment::new(
PromptSlot::ContextualUser,
self.0,
)]))
}
}
struct RecordingTurnItemContributor {
name: &'static str,
calls: Arc<Mutex<Vec<&'static str>>>,
@@ -186,6 +202,8 @@ async fn contributors_preserve_registration_order() {
let mut builder = ExtensionRegistryBuilder::<()>::new();
builder.prompt_contributor(Arc::new(NamedContextContributor("first")));
builder.prompt_contributor(Arc::new(NamedContextContributor("second")));
builder.prompt_contributor(Arc::new(NamedTurnContextContributor("turn-first")));
builder.prompt_contributor(Arc::new(NamedTurnContextContributor("turn-second")));
for name in ["first", "second"] {
builder.turn_item_contributor(Arc::new(RecordingTurnItemContributor {
name,
@@ -199,7 +217,25 @@ async fn contributors_preserve_registration_order() {
let mut fragments = Vec::new();
for contributor in registry.context_contributors() {
fragments.extend(contributor.contribute(&session_store, &thread_store).await);
fragments.extend(
contributor
.contribute_thread_context(&session_store, &thread_store)
.await,
);
}
for contributor in registry.context_contributors() {
fragments.extend(
contributor
.contribute_turn_context(TurnContextContributionInput {
thread_id: codex_protocol::ThreadId::default(),
turn_id: turn_store.level_id(),
session_store: &session_store,
thread_store: &thread_store,
turn_store: &turn_store,
model_context_window: Some(123),
})
.await,
);
}
let mut item = TurnItem::HookPrompt(HookPromptItem {
id: "item".to_string(),
@@ -217,6 +253,8 @@ async fn contributors_preserve_registration_order() {
vec![
PromptFragment::developer_policy("first"),
PromptFragment::developer_policy("second"),
PromptFragment::new(PromptSlot::ContextualUser, "turn-first"),
PromptFragment::new(PromptSlot::ContextualUser, "turn-second"),
]
);
assert_eq!(