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 18:40:28 +01:00
committed by GitHub
Unverified
parent 790213ded0
commit 9684ec25be
14 changed files with 307 additions and 33 deletions
+1 -1
View File
@@ -107,7 +107,7 @@ impl codex_extension_api::ThreadLifecycleContributor<Config> for GuardianMemoryC
}
impl codex_extension_api::ContextContributor for GuardianMemoryContextProbe {
fn contribute<'a>(
fn contribute_thread_context<'a>(
&'a self,
_session_store: &'a codex_extension_api::ExtensionData,
thread_store: &'a codex_extension_api::ExtensionData,
+113 -16
View File
@@ -58,7 +58,9 @@ use codex_exec_server::Environment;
use codex_exec_server::EnvironmentManager;
use codex_extension_api::ExtensionDataInit;
use codex_extension_api::LoadedUserInstructions;
use codex_extension_api::PromptFragment;
use codex_extension_api::PromptSlot;
use codex_extension_api::TurnContextContributionInput;
use codex_features::FEATURES;
use codex_features::Feature;
use codex_features::unstable_features_warning_event;
@@ -911,6 +913,25 @@ async fn thread_title_from_thread_store(
(!title.is_empty() && thread.preview.trim() != title).then(|| title.to_string())
}
fn push_prompt_fragment(
fragment: PromptFragment,
developer_sections: &mut Vec<String>,
contextual_user_sections: &mut Vec<String>,
separate_developer_sections: &mut Vec<String>,
) {
match fragment.slot() {
PromptSlot::DeveloperPolicy | PromptSlot::DeveloperCapabilities => {
developer_sections.push(fragment.text().to_string());
}
PromptSlot::ContextualUser => {
contextual_user_sections.push(fragment.text().to_string());
}
PromptSlot::SeparateDeveloper => {
separate_developer_sections.push(fragment.text().to_string());
}
}
}
impl Session {
pub(crate) async fn app_server_client_metadata(&self) -> AppServerClientMetadata {
let state = self.state.lock().await;
@@ -2860,6 +2881,57 @@ impl Session {
}
}
async fn build_turn_context_contribution_items(
&self,
turn_context: &TurnContext,
) -> Vec<ResponseItem> {
let mut developer_sections = Vec::new();
let mut contextual_user_sections = Vec::new();
let mut separate_developer_sections = Vec::new();
let context_contributors = self.services.extensions.context_contributors().to_vec();
for contributor in &context_contributors {
for fragment in contributor
.contribute_turn_context(TurnContextContributionInput {
thread_id: self.thread_id(),
turn_id: turn_context.sub_id.as_str(),
session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data,
turn_store: turn_context.extension_data.as_ref(),
model_context_window: turn_context.model_context_window(),
})
.await
{
push_prompt_fragment(
fragment,
&mut developer_sections,
&mut contextual_user_sections,
&mut separate_developer_sections,
);
}
}
let mut items = Vec::with_capacity(3);
if let Some(developer_message) =
crate::context_manager::updates::build_developer_update_item(developer_sections)
{
items.push(developer_message);
}
for section in separate_developer_sections {
if let Some(developer_message) =
crate::context_manager::updates::build_developer_update_item(vec![section])
{
items.push(developer_message);
}
}
if let Some(contextual_user_message) =
crate::context_manager::updates::build_contextual_user_message(contextual_user_sections)
{
items.push(contextual_user_message);
}
items
}
pub(crate) async fn build_initial_context(
&self,
turn_context: &TurnContext,
@@ -3026,25 +3098,40 @@ impl Session {
developer_sections.push(plugin_instructions.render());
}
let context_contributors = self.services.extensions.context_contributors().to_vec();
for contributor in context_contributors {
for contributor in &context_contributors {
for fragment in contributor
.contribute(
.contribute_thread_context(
&self.services.session_extension_data,
&self.services.thread_extension_data,
)
.await
{
match fragment.slot() {
PromptSlot::DeveloperPolicy | PromptSlot::DeveloperCapabilities => {
developer_sections.push(fragment.text().to_string());
}
PromptSlot::ContextualUser => {
contextual_user_sections.push(fragment.text().to_string());
}
PromptSlot::SeparateDeveloper => {
separate_developer_sections.push(fragment.text().to_string());
}
}
push_prompt_fragment(
fragment,
&mut developer_sections,
&mut contextual_user_sections,
&mut separate_developer_sections,
);
}
}
for contributor in &context_contributors {
for fragment in contributor
.contribute_turn_context(TurnContextContributionInput {
thread_id: self.thread_id(),
turn_id: turn_context.sub_id.as_str(),
session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data,
turn_store: turn_context.extension_data.as_ref(),
model_context_window: turn_context.model_context_window(),
})
.await
{
push_prompt_fragment(
fragment,
&mut developer_sections,
&mut contextual_user_sections,
&mut separate_developer_sections,
);
}
}
if let Some(user_instructions) = turn_context.user_instructions.as_deref() {
@@ -3211,15 +3298,25 @@ impl Session {
let state = self.state.lock().await;
state.reference_context_item()
};
let turn_context_item = turn_context.to_turn_context_item();
if reference_context_item.as_ref() == Some(&turn_context_item) {
return;
}
let should_inject_full_context = reference_context_item.is_none();
let context_items = if should_inject_full_context {
let mut context_items = if should_inject_full_context {
self.build_initial_context(turn_context).await
} else {
// Steady-state path: append only context diffs to minimize token overhead.
// Steady-state path: append only built-in context diffs here; turn-scoped extension
// context is added below.
self.build_settings_update_items(reference_context_item.as_ref(), turn_context)
.await
};
let turn_context_item = turn_context.to_turn_context_item();
if !should_inject_full_context {
context_items.extend(
self.build_turn_context_contribution_items(turn_context)
.await,
);
}
if !context_items.is_empty() {
self.record_conversation_items(turn_context, &context_items)
.await;
+91 -1
View File
@@ -7540,9 +7540,13 @@ async fn make_multi_agent_v2_usage_hint_test_session(
struct PromptExtensionTestContributor;
struct PromptExtensionTestState;
struct TurnContextExtensionTestContributor;
struct TurnContextExtensionTestState {
expected_model_context_window: Option<i64>,
}
impl codex_extension_api::ContextContributor for PromptExtensionTestContributor {
fn contribute<'a>(
fn contribute_thread_context<'a>(
&'a self,
_session_store: &'a codex_extension_api::ExtensionData,
thread_store: &'a codex_extension_api::ExtensionData,
@@ -7571,6 +7575,31 @@ fn prompt_extension_test_registry()
Arc::new(builder.build())
}
impl codex_extension_api::ContextContributor for TurnContextExtensionTestContributor {
fn contribute_turn_context<'a>(
&'a self,
input: codex_extension_api::TurnContextContributionInput<'a>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Vec<codex_extension_api::PromptFragment>> + Send + 'a>,
> {
Box::pin(async move {
let Some(state) = input.turn_store.get::<TurnContextExtensionTestState>() else {
return Vec::new();
};
(input.model_context_window == state.expected_model_context_window
&& input.model_context_window.is_some()
&& !input.turn_id.is_empty())
.then(|| {
codex_extension_api::PromptFragment::developer_policy(
"turn context extension enabled",
)
})
.into_iter()
.collect()
})
}
}
#[tokio::test]
async fn build_initial_context_includes_prompt_fragments_from_extensions() {
let (mut session, turn_context) = make_session_and_context().await;
@@ -7592,6 +7621,67 @@ async fn build_initial_context_includes_prompt_fragments_from_extensions() {
);
}
#[tokio::test]
async fn build_initial_context_includes_turn_context_fragments_from_extensions() {
let (mut session, mut turn_context) = make_session_and_context().await;
let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
builder.prompt_contributor(Arc::new(TurnContextExtensionTestContributor));
session.services.extensions = Arc::new(builder.build());
turn_context.model_info.context_window = Some(100);
turn_context.model_info.effective_context_window_percent = 50;
turn_context
.extension_data
.insert(TurnContextExtensionTestState {
expected_model_context_window: Some(50),
});
let initial_context = session.build_initial_context(&turn_context).await;
let developer_messages = developer_message_texts(&initial_context);
assert!(
developer_messages
.iter()
.flatten()
.any(|text| *text == "turn context extension enabled"),
"expected turn context extension developer text, got {developer_messages:?}"
);
}
#[tokio::test]
async fn record_context_updates_includes_turn_context_fragments_on_steady_state_turns() {
let (mut session, mut turn_context) = make_session_and_context().await;
let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
builder.prompt_contributor(Arc::new(TurnContextExtensionTestContributor));
session.services.extensions = Arc::new(builder.build());
turn_context.model_info.context_window = Some(200);
turn_context.model_info.effective_context_window_percent = 25;
turn_context
.extension_data
.insert(TurnContextExtensionTestState {
expected_model_context_window: Some(50),
});
let mut previous_context_item = turn_context.to_turn_context_item();
previous_context_item.turn_id = Some("previous-turn-id".to_string());
{
let mut state = session.state.lock().await;
state.set_reference_context_item(Some(previous_context_item));
}
session
.record_context_updates_and_set_reference_context_item(&turn_context)
.await;
let history = session.clone_history().await;
let developer_messages = developer_message_texts(history.raw_items());
assert!(
developer_messages
.iter()
.flatten()
.any(|text| *text == "turn context extension enabled"),
"expected steady-state turn context extension developer text, got {developer_messages:?}"
);
}
#[tokio::test]
async fn build_initial_context_omits_prompt_fragments_without_extension_state() {
let (mut session, turn_context) = make_session_and_context().await;