Move previous turn context tracking into ContextManager history (#12179)

## Summary
- add `previous_context_item: Option<TurnContextItem>` to
`ContextManager`
- expose session/state accessors for reading and updating the stored
previous context item
- switch settings diffing to use `TurnContextItem` instead of
`TurnContext`
- remove submission-loop local `previous_context` and persist the
previous context item in history

## Testing
- `just fmt`
- `just fix -p codex-core`
- `cargo test -p codex-core --test all model_switching::`
- `cargo test -p codex-core --test all collaboration_instructions::`
- `cargo test -p codex-core --test all personality::`
- `cargo test -p codex-core --test all
permissions_messages::permissions_message_not_added_when_no_change`
This commit is contained in:
Charley Cunningham
2026-02-19 09:56:20 -08:00
committed by GitHub
Unverified
parent f6fd4cb3f5
commit f2d5842ed1
6 changed files with 103 additions and 66 deletions
+34 -43
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(),
@@ -1995,11 +1992,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;
@@ -2014,7 +2006,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> {
@@ -2025,7 +2017,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(),
@@ -2660,6 +2652,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,
@@ -3123,7 +3125,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 {
@@ -3173,8 +3177,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,
@@ -3251,13 +3254,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,
@@ -3285,7 +3282,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;
@@ -3362,12 +3358,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,
@@ -3430,8 +3421,9 @@ mod handlers {
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 previous_context_item = sess.previous_context_item().await;
let update_items = sess.build_settings_update_items(
previous_context.as_ref(),
previous_context_item.as_ref(),
previous_model.as_deref(),
&current_context,
);
@@ -3445,16 +3437,12 @@ mod handlers {
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);
sess.set_previous_context_item(Some(current_context.to_turn_context_item()))
.await;
}
}
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
{
@@ -3479,7 +3467,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(
@@ -5465,9 +5454,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(),
@@ -7738,8 +7725,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()
+1 -7
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 {
@@ -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),
+14 -14
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,
+29 -2
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 {
+9
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,