Compare commits

...

3 Commits

Author SHA1 Message Date
Charles Cunningham
3ca93fad63 core: generate context diffs after pre-turn compaction 2026-02-18 17:28:50 -08:00
Charles Cunningham
2783c0c927 core: use TurnContext collaboration mode in TurnContextItem 2026-02-18 17:06:33 -08:00
Charles Cunningham
3dced9e1be core: track previous context item in history for diffs 2026-02-18 17:06:33 -08:00
10 changed files with 147 additions and 86 deletions

View File

@@ -660,10 +660,7 @@ impl TurnContext {
.unwrap_or(compact::SUMMARIZATION_PROMPT)
}
pub(crate) fn to_turn_context_item(
&self,
collaboration_mode: CollaborationMode,
) -> TurnContextItem {
pub(crate) fn to_turn_context_item(&self) -> TurnContextItem {
TurnContextItem {
turn_id: Some(self.sub_id.clone()),
cwd: self.cwd.clone(),
@@ -672,7 +669,7 @@ impl TurnContext {
network: self.turn_context_network_item(),
model: self.model_info.slug.clone(),
personality: self.personality,
collaboration_mode: Some(collaboration_mode),
collaboration_mode: Some(self.collaboration_mode.clone()),
effort: self.reasoning_effort,
summary: self.reasoning_summary,
user_instructions: self.user_instructions.clone(),
@@ -1993,11 +1990,6 @@ impl Session {
.await
}
pub(crate) async fn current_collaboration_mode(&self) -> CollaborationMode {
let state = self.state.lock().await;
state.session_configuration.collaboration_mode.clone()
}
pub(crate) fn is_model_switch_developer_message(item: &ResponseItem) -> bool {
let ResponseItem::Message { role, content, .. } = item else {
return false;
@@ -2012,7 +2004,7 @@ impl Session {
}
fn build_settings_update_items(
&self,
previous_context: Option<&Arc<TurnContext>>,
previous_context: Option<&TurnContextItem>,
resumed_model: Option<&str>,
current_context: &TurnContext,
) -> Vec<ResponseItem> {
@@ -2023,7 +2015,7 @@ impl Session {
let shell = self.user_shell();
let exec_policy = self.services.exec_policy.current();
crate::context_manager::updates::build_settings_update_items(
previous_context.map(Arc::as_ref),
previous_context,
resumed_model,
current_context,
shell.as_ref(),
@@ -2658,6 +2650,16 @@ impl Session {
state.clone_history()
}
pub(crate) async fn previous_context_item(&self) -> Option<TurnContextItem> {
let state = self.state.lock().await;
state.previous_context_item()
}
pub(crate) async fn set_previous_context_item(&self, item: Option<TurnContextItem>) {
let mut state = self.state.lock().await;
state.set_previous_context_item(item);
}
pub(crate) async fn update_token_usage_info(
&self,
turn_context: &TurnContext,
@@ -3121,7 +3123,9 @@ impl Session {
async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiver<Submission>) {
// Seed with context in case there is an OverrideTurnContext first.
let mut previous_context: Option<Arc<TurnContext>> = Some(sess.new_default_turn().await);
let initial_context = sess.new_default_turn().await;
sess.set_previous_context_item(Some(initial_context.to_turn_context_item()))
.await;
// To break out of this loop, send Op::Shutdown.
while let Ok(sub) = rx_sub.recv().await {
@@ -3171,8 +3175,7 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
.await;
}
Op::UserInput { .. } | Op::UserTurn { .. } => {
handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op, &mut previous_context)
.await;
handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op).await;
}
Op::ExecApproval {
id: approval_id,
@@ -3249,13 +3252,7 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
handlers::set_thread_name(&sess, sub.id.clone(), name).await;
}
Op::RunUserShellCommand { command } => {
handlers::run_user_shell_command(
&sess,
sub.id.clone(),
command,
&mut previous_context,
)
.await;
handlers::run_user_shell_command(&sess, sub.id.clone(), command).await;
}
Op::ResolveElicitation {
server_name,
@@ -3283,7 +3280,6 @@ mod handlers {
use crate::codex::Session;
use crate::codex::SessionSettingsUpdate;
use crate::codex::SteerInputError;
use crate::codex::TurnContext;
use crate::codex::spawn_review_thread;
use crate::config::Config;
@@ -3360,12 +3356,7 @@ mod handlers {
}
}
pub async fn user_input_or_turn(
sess: &Arc<Session>,
sub_id: String,
op: Op,
previous_context: &mut Option<Arc<TurnContext>>,
) {
pub async fn user_input_or_turn(sess: &Arc<Session>, sub_id: String, op: Op) {
let (items, updates) = match op {
Op::UserTurn {
cwd,
@@ -3427,32 +3418,16 @@ mod handlers {
// Attempt to inject input into current task.
if let Err(SteerInputError::NoActiveTurn(items)) = sess.steer_input(items, None).await {
sess.seed_initial_context_if_needed(&current_context).await;
let previous_model = sess.previous_model().await;
let update_items = sess.build_settings_update_items(
previous_context.as_ref(),
previous_model.as_deref(),
&current_context,
);
if !update_items.is_empty() {
sess.record_conversation_items(&current_context, &update_items)
.await;
}
sess.refresh_mcp_servers_if_requested(&current_context)
.await;
let regular_task = sess.take_startup_regular_task().await.unwrap_or_default();
sess.spawn_task(Arc::clone(&current_context), items, regular_task)
.await;
*previous_context = Some(current_context);
}
}
pub async fn run_user_shell_command(
sess: &Arc<Session>,
sub_id: String,
command: String,
previous_context: &mut Option<Arc<TurnContext>>,
) {
pub async fn run_user_shell_command(sess: &Arc<Session>, sub_id: String, command: String) {
if let Some((turn_context, cancellation_token)) =
sess.active_turn_context_and_cancellation_token().await
{
@@ -3477,7 +3452,8 @@ mod handlers {
UserShellCommandTask::new(command),
)
.await;
*previous_context = Some(turn_context);
sess.set_previous_context_item(Some(turn_context.to_turn_context_item()))
.await;
}
pub async fn resolve_elicitation(
@@ -4293,6 +4269,19 @@ pub(crate) async fn run_turn(
error!("Failed to run pre-sampling compact");
return None;
}
let previous_model = sess.previous_model().await;
let previous_context_item = sess.previous_context_item().await;
let update_items = sess.build_settings_update_items(
previous_context_item.as_ref(),
previous_model.as_deref(),
turn_context.as_ref(),
);
if !update_items.is_empty() {
sess.record_conversation_items(turn_context.as_ref(), &update_items)
.await;
}
sess.set_previous_context_item(Some(turn_context.to_turn_context_item()))
.await;
let skills_outcome = Some(
sess.services
@@ -5463,9 +5452,7 @@ async fn try_run_sampling_request(
prompt: &Prompt,
cancellation_token: CancellationToken,
) -> CodexResult<SamplingRequestResult> {
let collaboration_mode = sess.current_collaboration_mode().await;
let rollout_item =
RolloutItem::TurnContext(turn_context.to_turn_context_item(collaboration_mode));
let rollout_item = RolloutItem::TurnContext(turn_context.to_turn_context_item());
feedback_tags!(
model = turn_context.model_info.slug.clone(),
@@ -7732,8 +7719,12 @@ mod tests {
.expect("rebuild config layer stack with network requirements");
current_context.config = Arc::new(config);
let update_items =
session.build_settings_update_items(Some(&previous_context), None, &current_context);
let previous_context_item = previous_context.to_turn_context_item();
let update_items = session.build_settings_update_items(
Some(&previous_context_item),
None,
&current_context,
);
let environment_update = update_items
.iter()

View File

@@ -118,13 +118,7 @@ async fn run_compact_task_inner(
// Reuse one client session so turn-scoped state (sticky routing, websocket append tracking)
// survives retries within this compact turn.
// TODO: If we need to guarantee the persisted mode always matches the prompt used for this
// turn, capture it in TurnContext at creation time. Using SessionConfiguration here avoids
// duplicating model settings on TurnContext, but an Op after turn start could update the
// session config before this write occurs.
let collaboration_mode = sess.current_collaboration_mode().await;
let rollout_item =
RolloutItem::TurnContext(turn_context.to_turn_context_item(collaboration_mode));
let rollout_item = RolloutItem::TurnContext(turn_context.to_turn_context_item());
sess.persist_rollout_items(&[rollout_item]).await;
loop {

View File

@@ -18,6 +18,7 @@ use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::InputModality;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::TurnContextItem;
use std::ops::Deref;
/// Transcript of thread history
@@ -26,6 +27,12 @@ pub(crate) struct ContextManager {
/// The oldest items are at the beginning of the vector.
items: Vec<ResponseItem>,
token_info: Option<TokenUsageInfo>,
/// Previous turn context snapshot used for diffing context and producing
/// model-visible settings update items.
///
/// When this is `None`, settings diffing treats the next turn as having no
/// baseline and emits a full reinjection of context state.
previous_context_item: Option<TurnContextItem>,
}
#[derive(Debug, Clone, Copy, Default)]
@@ -41,6 +48,7 @@ impl ContextManager {
Self {
items: Vec::new(),
token_info: TokenUsageInfo::new_or_append(&None, &None, None),
previous_context_item: None,
}
}
@@ -52,6 +60,14 @@ impl ContextManager {
self.token_info = info;
}
pub(crate) fn set_previous_context_item(&mut self, item: Option<TurnContextItem>) {
self.previous_context_item = item;
}
pub(crate) fn previous_context_item(&self) -> Option<TurnContextItem> {
self.previous_context_item.clone()
}
pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
match &mut self.token_info {
Some(info) => info.fill_to_context_window(context_window),

View File

@@ -6,26 +6,27 @@ use codex_protocol::config_types::Personality;
use codex_protocol::models::DeveloperInstructions;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::protocol::TurnContextItem;
fn build_environment_update_item(
previous: Option<&TurnContext>,
previous: Option<&TurnContextItem>,
next: &TurnContext,
shell: &Shell,
) -> Option<ResponseItem> {
let prev = previous?;
let prev_context = EnvironmentContext::from_turn_context(prev, shell);
let prev_context = EnvironmentContext::from_turn_context_item(prev, shell);
let next_context = EnvironmentContext::from_turn_context(next, shell);
if prev_context.equals_except_shell(&next_context) {
return None;
}
Some(ResponseItem::from(EnvironmentContext::diff(
prev, next, shell,
)))
Some(ResponseItem::from(
EnvironmentContext::diff_from_turn_context_item(prev, next, shell),
))
}
fn build_permissions_update_item(
previous: Option<&TurnContext>,
previous: Option<&TurnContextItem>,
next: &TurnContext,
exec_policy: &Policy,
) -> Option<ResponseItem> {
@@ -46,11 +47,11 @@ fn build_permissions_update_item(
}
fn build_collaboration_mode_update_item(
previous: Option<&TurnContext>,
previous: Option<&TurnContextItem>,
next: &TurnContext,
) -> Option<ResponseItem> {
let prev = previous?;
if prev.collaboration_mode != next.collaboration_mode {
if prev.collaboration_mode.as_ref() != Some(&next.collaboration_mode) {
// If the next mode has empty developer instructions, this returns None and we emit no
// update, so prior collaboration instructions remain in the prompt history.
Some(DeveloperInstructions::from_collaboration_mode(&next.collaboration_mode)?.into())
@@ -60,7 +61,7 @@ fn build_collaboration_mode_update_item(
}
fn build_personality_update_item(
previous: Option<&TurnContext>,
previous: Option<&TurnContextItem>,
next: &TurnContext,
personality_feature_enabled: bool,
) -> Option<ResponseItem> {
@@ -68,7 +69,7 @@ fn build_personality_update_item(
return None;
}
let previous = previous?;
if next.model_info.slug != previous.model_info.slug {
if next.model_info.slug != previous.model {
return None;
}
@@ -96,12 +97,11 @@ pub(crate) fn personality_message_for(
}
pub(crate) fn build_model_instructions_update_item(
previous: Option<&TurnContext>,
previous: Option<&TurnContextItem>,
resumed_model: Option<&str>,
next: &TurnContext,
) -> Option<ResponseItem> {
let previous_model =
resumed_model.or_else(|| previous.map(|prev| prev.model_info.slug.as_str()))?;
let previous_model = resumed_model.or_else(|| previous.map(|prev| prev.model.as_str()))?;
if previous_model == next.model_info.slug {
return None;
}
@@ -115,7 +115,7 @@ pub(crate) fn build_model_instructions_update_item(
}
pub(crate) fn build_settings_update_items(
previous: Option<&TurnContext>,
previous: Option<&TurnContextItem>,
resumed_model: Option<&str>,
next: &TurnContext,
shell: &Shell,

View File

@@ -4,6 +4,8 @@ use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::ENVIRONMENT_CONTEXT_CLOSE_TAG;
use codex_protocol::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG;
use codex_protocol::protocol::TurnContextItem;
use codex_protocol::protocol::TurnContextNetworkItem;
use serde::Deserialize;
use serde::Serialize;
use std::path::PathBuf;
@@ -44,8 +46,12 @@ impl EnvironmentContext {
self.cwd == *cwd && self.network == *network
}
pub fn diff(before: &TurnContext, after: &TurnContext, shell: &Shell) -> Self {
let before_network = Self::network_from_turn_context(before);
pub fn diff_from_turn_context_item(
before: &TurnContextItem,
after: &TurnContext,
shell: &Shell,
) -> Self {
let before_network = Self::network_from_turn_context_item(before);
let after_network = Self::network_from_turn_context(after);
let cwd = if before.cwd != after.cwd {
Some(after.cwd.clone())
@@ -68,6 +74,14 @@ impl EnvironmentContext {
)
}
pub fn from_turn_context_item(turn_context_item: &TurnContextItem, shell: &Shell) -> Self {
Self::new(
Some(turn_context_item.cwd.clone()),
shell.clone(),
Self::network_from_turn_context_item(turn_context_item),
)
}
fn network_from_turn_context(turn_context: &TurnContext) -> Option<NetworkContext> {
let network = turn_context
.config
@@ -81,6 +95,19 @@ impl EnvironmentContext {
denied_domains: network.denied_domains.clone().unwrap_or_default(),
})
}
fn network_from_turn_context_item(
turn_context_item: &TurnContextItem,
) -> Option<NetworkContext> {
let TurnContextNetworkItem {
allowed_domains,
denied_domains,
} = turn_context_item.network.as_ref()?;
Some(NetworkContext {
allowed_domains: allowed_domains.clone(),
denied_domains: denied_domains.clone(),
})
}
}
impl EnvironmentContext {

View File

@@ -11,6 +11,7 @@ use crate::protocol::TokenUsage;
use crate::protocol::TokenUsageInfo;
use crate::tasks::RegularTask;
use crate::truncate::TruncationPolicy;
use codex_protocol::protocol::TurnContextItem;
/// Persistent, session-scoped state previously stored directly on `Session`.
pub(crate) struct SessionState {
@@ -80,6 +81,14 @@ impl SessionState {
self.history.set_token_info(info);
}
pub(crate) fn set_previous_context_item(&mut self, item: Option<TurnContextItem>) {
self.history.set_previous_context_item(item);
}
pub(crate) fn previous_context_item(&self) -> Option<TurnContextItem> {
self.history.previous_context_item()
}
// Token/rate limit helpers
pub(crate) fn update_token_info_from_usage(
&mut self,

View File

@@ -3007,7 +3007,7 @@ async fn snapshot_request_shape_pre_turn_compaction_including_incoming_user_mess
insta::assert_snapshot!(
"pre_turn_compaction_including_incoming_shapes",
format_labeled_requests_snapshot(
"Pre-turn auto-compaction with a context override emits the context diff in the compact request while the incoming user message is still excluded.",
"Pre-turn auto-compaction with a context override excludes incoming-turn context diffs from the compact request, then appends those diffs immediately before the incoming user message after compaction.",
&[
("Local Compaction Request", &requests[2]),
("Local Post-Compaction History Layout", &requests[3]),
@@ -3021,11 +3021,21 @@ async fn snapshot_request_shape_pre_turn_compaction_including_incoming_user_mess
.any(|text| text == "USER_THREE"),
"current behavior excludes incoming user message from pre-turn compaction input"
);
let compact_body = requests[2].body_json().to_string();
assert!(
!compact_body.contains(PRETURN_CONTEXT_DIFF_CWD),
"pre-turn compaction request should exclude incoming context diff items"
);
let follow_up_user_texts = requests[3].message_input_texts("user");
assert!(
follow_up_user_texts.iter().any(|text| text == "USER_THREE"),
"expected post-compaction follow-up request to keep incoming user text"
);
let follow_up_body = requests[3].body_json().to_string();
assert!(
follow_up_body.contains(PRETURN_CONTEXT_DIFF_CWD),
"post-compaction follow-up request should include the incoming context diff items"
);
let follow_up_user_images = requests[3].message_input_image_urls("user");
assert!(
follow_up_user_images

View File

@@ -1390,13 +1390,27 @@ async fn snapshot_request_shape_remote_pre_turn_compaction_including_incoming_us
insta::assert_snapshot!(
"remote_pre_turn_compaction_including_incoming_shapes",
format_labeled_requests_snapshot(
"Remote pre-turn auto-compaction with a context override emits the context diff in the compact request while excluding the incoming user message.",
"Remote pre-turn auto-compaction with a context override excludes incoming-turn context diffs from the compact request, then appends those diffs immediately before the incoming user message after compaction.",
&[
("Remote Compaction Request", &compact_request),
("Remote Post-Compaction History Layout", &requests[2]),
]
)
);
assert!(
!compact_request
.body_json()
.to_string()
.contains(PRETURN_CONTEXT_DIFF_CWD),
"pre-turn remote compaction request should exclude incoming context diff items"
);
assert!(
requests[2]
.body_json()
.to_string()
.contains(PRETURN_CONTEXT_DIFF_CWD),
"post-compaction remote request should include incoming context diff items"
);
assert_eq!(
requests[2]
.message_input_texts("user")

View File

@@ -1,8 +1,8 @@
---
source: core/tests/suite/compact.rs
expression: "format_labeled_requests_snapshot(\"Pre-turn auto-compaction with a context override emits the context diff in the compact request while the incoming user message is still excluded.\",\n&[(\"Local Compaction Request\", &requests[2]),\n(\"Local Post-Compaction History Layout\", &requests[3]),])"
expression: "format_labeled_requests_snapshot(\"Pre-turn auto-compaction with a context override excludes incoming-turn context diffs from the compact request, then appends those diffs immediately before the incoming user message after compaction.\",\n&[(\"Local Compaction Request\", &requests[2]),\n(\"Local Post-Compaction History Layout\", &requests[3]),])"
---
Scenario: Pre-turn auto-compaction with a context override emits the context diff in the compact request while the incoming user message is still excluded.
Scenario: Pre-turn auto-compaction with a context override excludes incoming-turn context diffs from the compact request, then appends those diffs immediately before the incoming user message after compaction.
## Local Compaction Request
00:message/developer:<PERMISSIONS_INSTRUCTIONS>
@@ -12,8 +12,7 @@ Scenario: Pre-turn auto-compaction with a context override emits the context dif
04:message/assistant:FIRST_REPLY
05:message/user:USER_TWO
06:message/assistant:SECOND_REPLY
07:message/user:<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD>
08:message/user:<SUMMARIZATION_PROMPT>
07:message/user:<SUMMARIZATION_PROMPT>
## Local Post-Compaction History Layout
00:message/developer:<PERMISSIONS_INSTRUCTIONS>
@@ -22,4 +21,5 @@ Scenario: Pre-turn auto-compaction with a context override emits the context dif
03:message/user:USER_ONE
04:message/user:USER_TWO
05:message/user:<COMPACTION_SUMMARY>\nPRE_TURN_SUMMARY
06:message/user:<image> | <input_image:image_url> | </image> | USER_THREE
06:message/user:<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD>
07:message/user:<image> | <input_image:image_url> | </image> | USER_THREE

View File

@@ -1,8 +1,8 @@
---
source: core/tests/suite/compact_remote.rs
expression: "format_labeled_requests_snapshot(\"Remote pre-turn auto-compaction with a context override emits the context diff in the compact request while excluding the incoming user message.\",\n&[(\"Remote Compaction Request\", &compact_request),\n(\"Remote Post-Compaction History Layout\", &requests[2]),])"
expression: "format_labeled_requests_snapshot(\"Remote pre-turn auto-compaction with a context override excludes incoming-turn context diffs from the compact request, then appends those diffs immediately before the incoming user message after compaction.\",\n&[(\"Remote Compaction Request\", &compact_request),\n(\"Remote Post-Compaction History Layout\", &requests[2]),])"
---
Scenario: Remote pre-turn auto-compaction with a context override emits the context diff in the compact request while excluding the incoming user message.
Scenario: Remote pre-turn auto-compaction with a context override excludes incoming-turn context diffs from the compact request, then appends those diffs immediately before the incoming user message after compaction.
## Remote Compaction Request
00:message/developer:<PERMISSIONS_INSTRUCTIONS>
@@ -12,7 +12,6 @@ Scenario: Remote pre-turn auto-compaction with a context override emits the cont
04:message/assistant:REMOTE_FIRST_REPLY
05:message/user:USER_TWO
06:message/assistant:REMOTE_SECOND_REPLY
07:message/user:<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD>
## Remote Post-Compaction History Layout
00:message/user:USER_ONE
@@ -21,4 +20,5 @@ Scenario: Remote pre-turn auto-compaction with a context override emits the cont
03:message/user:<AGENTS_MD>
04:message/user:<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD>
05:message/user:<COMPACTION_SUMMARY>\nREMOTE_PRE_TURN_SUMMARY
06:message/user:USER_THREE
06:message/user:<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD>
07:message/user:USER_THREE