Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion codex-rs/core/src/guardian/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
129 changes: 113 additions & 16 deletions codex-rs/core/src/session/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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,
);
Comment thread
jif-oai marked this conversation as resolved.
}
}

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,
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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,
Comment thread
jif-oai marked this conversation as resolved.
);
}
if !context_items.is_empty() {
self.record_conversation_items(turn_context, &context_items)
.await;
Expand Down
92 changes: 91 additions & 1 deletion codex-rs/core/src/session/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down
6 changes: 5 additions & 1 deletion codex-rs/ext/extension-api/examples/enabled_extensions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,11 @@ async fn contribute_prompt(
) -> Vec<codex_extension_api::PromptFragment> {
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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
28 changes: 26 additions & 2 deletions codex-rs/ext/extension-api/src/contributors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,15 @@ use codex_tools::ToolExecutor;

use crate::ExtensionData;

mod context;
mod mcp;
mod prompt;
mod thread_lifecycle;
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;
Expand Down Expand Up @@ -62,12 +64,34 @@ pub trait McpServerContributor<C: Sync>: 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<PromptFragment>>;
) -> ExtensionFuture<'a, Vec<PromptFragment>> {
Comment thread
jif-oai marked this conversation as resolved.
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<PromptFragment>> {
Comment thread
jif-oai marked this conversation as resolved.
Box::pin(async move {
let _self = self;
let _input = input;
Vec::new()
})
}
}

/// Contributor for host-owned thread lifecycle gates.
Expand Down
Loading
Loading