From 9684ec25be97707ea5a0cc4dc0a2d23d466c4b41 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 18 Jun 2026 18:40:28 +0100 Subject: [PATCH] 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 --- codex-rs/core/src/guardian/tests.rs | 2 +- codex-rs/core/src/session/mod.rs | 129 +++++++++++++++--- codex-rs/core/src/session/tests.rs | 92 ++++++++++++- .../examples/enabled_extensions.rs | 6 +- .../shared_state_extension.rs | 4 +- .../ext/extension-api/src/contributors.rs | 28 +++- .../extension-api/src/contributors/context.rs | 20 +++ codex-rs/ext/extension-api/src/lib.rs | 1 + codex-rs/ext/extension-api/tests/registry.rs | 44 +++++- codex-rs/ext/memories/src/extension.rs | 2 +- codex-rs/ext/memories/src/tests.rs | 2 +- codex-rs/ext/skills/src/extension.rs | 2 +- codex-rs/ext/skills/tests/skills_extension.rs | 6 +- codex-rs/protocol/src/protocol.rs | 2 +- 14 files changed, 307 insertions(+), 33 deletions(-) create mode 100644 codex-rs/ext/extension-api/src/contributors/context.rs diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index d6eb93dd1..15024990a 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -107,7 +107,7 @@ impl codex_extension_api::ThreadLifecycleContributor 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, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0ecc70681..84e60a3d2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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, + contextual_user_sections: &mut Vec, + separate_developer_sections: &mut Vec, +) { + 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 { + 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; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 724c35f2e..63e8783cc 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -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, +} 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> + Send + 'a>, + > { + Box::pin(async move { + let Some(state) = input.turn_store.get::() 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; diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs index 9027c2824..45b178bc8 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -74,7 +74,11 @@ async fn contribute_prompt( ) -> Vec { 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, + ); } fragments } diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs index 531f65b99..414a67215 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs @@ -17,7 +17,7 @@ pub fn install(registry: &mut ExtensionRegistryBuilder<()>) { struct StyleContributor; impl ContextContributor for StyleContributor { - fn contribute<'a>( + fn contribute_thread_context<'a>( &'a self, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, @@ -37,7 +37,7 @@ impl ContextContributor for StyleContributor { struct UsageContributor; impl ContextContributor for UsageContributor { - fn contribute<'a>( + fn contribute_thread_context<'a>( &'a self, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 3a0fee49a..7d16efd3d 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -11,6 +11,7 @@ use codex_tools::ToolExecutor; use crate::ExtensionData; +mod context; mod mcp; mod prompt; mod thread_lifecycle; @@ -18,6 +19,7 @@ mod tool_lifecycle; mod turn_input; mod turn_lifecycle; +pub use context::TurnContextContributionInput; pub use mcp::McpServerContribution; pub use mcp::McpServerContributionContext; pub use prompt::PromptFragment; @@ -62,12 +64,34 @@ pub trait McpServerContributor: Send + Sync { } /// Extension contribution that adds prompt fragments during prompt assembly. +/// +/// Implementations should use the method matching the scope needed by the +/// fragment: thread/session context for stable inputs, and turn context for +/// fragments that depend on turn-local host state. pub trait ContextContributor: Send + Sync { - fn contribute<'a>( + fn contribute_thread_context<'a>( &'a self, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, - ) -> ExtensionFuture<'a, Vec>; + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let _self = self; + let _session_store = session_store; + let _thread_store = thread_store; + Vec::new() + }) + } + + fn contribute_turn_context<'a>( + &'a self, + input: TurnContextContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let _self = self; + let _input = input; + Vec::new() + }) + } } /// Contributor for host-owned thread lifecycle gates. diff --git a/codex-rs/ext/extension-api/src/contributors/context.rs b/codex-rs/ext/extension-api/src/contributors/context.rs new file mode 100644 index 000000000..fb6239ec4 --- /dev/null +++ b/codex-rs/ext/extension-api/src/contributors/context.rs @@ -0,0 +1,20 @@ +use codex_protocol::ThreadId; + +use crate::ExtensionData; + +/// Host context available while extensions contribute turn-scoped context fragments. +#[derive(Clone, Copy)] +pub struct TurnContextContributionInput<'a> { + /// Stable host-owned thread identifier. + pub thread_id: ThreadId, + /// Stable host-owned turn identifier. + pub turn_id: &'a str, + /// Store scoped to the host session runtime. + pub session_store: &'a ExtensionData, + /// Store scoped to this thread runtime. + pub thread_store: &'a ExtensionData, + /// Store scoped to this turn. + pub turn_store: &'a ExtensionData, + /// Effective model context window for this turn, when known. + pub model_context_window: Option, +} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index e646a53a5..de0e0c99f 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -54,6 +54,7 @@ pub use contributors::ToolLifecycleContributor; pub use contributors::ToolLifecycleFuture; pub use contributors::ToolStartInput; pub use contributors::TurnAbortInput; +pub use contributors::TurnContextContributionInput; pub use contributors::TurnErrorInput; pub use contributors::TurnInputContext; pub use contributors::TurnInputContributor; diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index 79773a98c..59b4f8c5e 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -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> { + Box::pin(std::future::ready(vec![PromptFragment::new( + PromptSlot::ContextualUser, + self.0, + )])) + } +} + struct RecordingTurnItemContributor { name: &'static str, calls: Arc>>, @@ -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!( diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs index 23c80a877..464ce2457 100644 --- a/codex-rs/ext/memories/src/extension.rs +++ b/codex-rs/ext/memories/src/extension.rs @@ -48,7 +48,7 @@ impl MemoriesExtensionConfig { } impl ContextContributor for MemoriesExtension { - fn contribute<'a>( + fn contribute_thread_context<'a>( &'a self, _session_store: &'a ExtensionData, thread_store: &'a ExtensionData, diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index b78c20c9f..9cc038b80 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -181,7 +181,7 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() { }); let fragments = extension - .contribute(&ExtensionData::new("session"), &thread_store) + .contribute_thread_context(&ExtensionData::new("session"), &thread_store) .await; assert_eq!(fragments.len(), 1); diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index c82ce6642..0a3da0972 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -103,7 +103,7 @@ impl ContextContributor for SkillsExtension where C: Send + Sync + 'static, { - fn contribute<'a>( + fn contribute_thread_context<'a>( &'a self, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index cd6f46d9c..16d9bdd1a 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -183,7 +183,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in .await; let prompt_fragments = registry.context_contributors()[0] - .contribute(&session_store, &thread_store) + .contribute_thread_context(&session_store, &thread_store) .await; assert_eq!(1, prompt_fragments.len()); assert!( @@ -228,7 +228,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in read_request_keys(&read_requests) ); let rebuilt_prompt_fragments = registry.context_contributors()[0] - .contribute(&session_store, &thread_store) + .contribute_thread_context(&session_store, &thread_store) .await; assert_eq!(1, rebuilt_prompt_fragments.len()); assert!(rebuilt_prompt_fragments[0].text().contains("lint-fix")); @@ -294,7 +294,7 @@ async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult { .await; let initial_fragments = registry.context_contributors()[0] - .contribute(&session_store, &thread_store) + .contribute_thread_context(&session_store, &thread_store) .await; assert!(initial_fragments.is_empty()); let EventMsg::Warning(warning) = event_rx.try_recv()?.msg else { diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index d51a30add..a71eb64fe 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2992,7 +2992,7 @@ pub struct TurnContextNetworkItem { /// context updates, and again after mid-turn compaction when replacement /// history re-establishes full context, so resume/fork replay can recover the /// latest durable baseline. -#[derive(Serialize, Deserialize, Clone, Debug, JsonSchema, TS)] +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema, TS)] pub struct TurnContextItem { #[serde(default, skip_serializing_if = "Option::is_none")] pub turn_id: Option,