openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-07 15:40:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve context baselines across nested agent forks (#41424)

Browse Source 
## What changed

- Treat a surviving full world-state snapshot as a context baseline when a fork removes the associated user message.
- Restore the previous turn settings and reference context from that baseline, without treating the segment as a user turn for rollback.
- Ignore partial snapshots and full snapshots superseded by compaction.

## Testing

- Cover resumed rollout reconstruction with removed task messages, partial snapshots, and compaction.
- Verify nested agents inherit developer instructions and environment context exactly once across history modes and compacted histories.

GitOrigin-RevId: 9f9f6992d7ea8cff9ac82a7fcdf0c8fb764db81c
This commit is contained in:
jif
2026-08-28 21:10:24 +00:00
committed by copyberry
parent 0918cd2c08
commit f9cdc90c2c
3 changed files with 333 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
codex-rs/core/src/session/rollout_reconstruction.rs
View File

@@ -78,6 +78,16 @@ fn finalize_active_segment<'a>(
return;
}
// Full world-state snapshots are persisted after installing initial context. They still
// establish a baseline when a child fork removes the parent turn's agent message. Do not
// count these context-only segments as user turns for rollback, or use a snapshot from
// before the segment's latest compaction.
let has_context_baseline = active_segment.counts_as_user_turn
|| active_segment
.world_state_replay
.iter()
.take_while(|item| !matches!(item, RolloutItem::Compacted(_)))
.any(|item| matches!(item, RolloutItem::WorldState(state) if state.full));
world_state_replay.extend(active_segment.world_state_replay);
// A surviving replacement-history checkpoint is a complete history base. Once we
@@ -92,15 +102,15 @@ fn finalize_active_segment<'a>(
*window = active_segment.window;
}
// `previous_turn_settings` come from the newest surviving user turn that established them.
if previous_turn_settings.is_none() && active_segment.counts_as_user_turn {
// Restore settings from the newest surviving context baseline.
if previous_turn_settings.is_none() && has_context_baseline {
*previous_turn_settings = active_segment.previous_turn_settings;
}
// `reference_context_item` comes from the newest surviving user turn baseline, or
// `reference_context_item` comes from the newest surviving context baseline, or
// from a surviving compaction that explicitly cleared that baseline.
if matches!(reference_context_item, TurnReferenceContextItem::NeverSet)
&& (active_segment.counts_as_user_turn
&& (has_context_baseline
|| matches!(
active_segment.reference_context_item,
TurnReferenceContextItem::Cleared

80
codex-rs/core/src/session/rollout_reconstruction_tests.rs
View File

@@ -24,6 +24,7 @@ use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;
use std::path::PathBuf;
use test_case::test_case;
use uuid::Uuid;
macro_rules! object {
@@ -181,8 +182,16 @@ async fn record_initial_history_ignores_security_risk_scores() {
);
}
#[derive(Clone, Copy)]
enum BaselineTurnInput {
UserMessage,
RemovedByFork,
}
#[test_case(BaselineTurnInput::UserMessage; "user turn")]
#[test_case(BaselineTurnInput::RemovedByFork; "fork removed the task message")]
#[tokio::test]
async fn record_initial_history_restores_world_state_baseline() {
async fn record_initial_history_restores_world_state_baseline(input: BaselineTurnInput) {
let (session, turn_context) = make_session_and_context().await;
let turn_context = Arc::new(turn_context);
let world_state = build_world_state_from_turn_context(&session, &turn_context).await;
@@ -200,8 +209,19 @@ async fn record_initial_history_restores_world_state_baseline() {
world_state_items.push(RolloutItem::WorldState(WorldStateItem::full(
world_state.snapshot().into_object(),
)));
let context_item = turn_context.to_turn_context_item();
let rollout_items = match input {
BaselineTurnInput::UserMessage => {
completed_user_turn_rollout(context_item.clone(), world_state_items)
}
BaselineTurnInput::RemovedByFork => {
world_state_items.push(RolloutItem::TurnContext(context_item.clone()));
world_state_items
}
};
// Exercise the persisted representation, not just an in-memory fork.
let rollout_items =
completed_user_turn_rollout(turn_context.to_turn_context_item(), world_state_items);
serde_json::from_value(serde_json::to_value(rollout_items).unwrap()).unwrap();
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
@@ -210,6 +230,20 @@ async fn record_initial_history_restores_world_state_baseline() {
rollout_path: Some(PathBuf::from("/tmp/resume.jsonl")),
}))
.await;
assert_eq!(
(
session.previous_turn_settings().await,
serde_json::to_value(session.reference_context_item().await).unwrap(),
),
(
Some(PreviousTurnSettings {
model: context_item.model.clone(),
comp_hash: context_item.comp_hash.clone(),
realtime_active: context_item.realtime_active,
}),
serde_json::to_value(Some(context_item)).unwrap(),
)
);
let step_context = StepContext::for_test(Arc::clone(&turn_context));
session
.record_context_updates_and_set_reference_context_item(&step_context)
@@ -674,6 +708,11 @@ async fn reconstruct_history_rollback_skips_non_user_turns_for_history_and_metad
},
)),
RolloutItem::ResponseItem(standalone_assistant.into()),
RolloutItem::WorldState(WorldStateItem::full(object!({}))),
RolloutItem::TurnContext(TurnContextItem {
turn_id: Some(standalone_turn_id.clone()),
..first_context_item.clone()
}),
RolloutItem::EventMsg(EventMsg::TurnComplete(
codex_protocol::protocol::TurnCompleteEvent {
turn_id: standalone_turn_id,
@@ -1053,11 +1092,42 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp
);
}
#[derive(Clone, Copy)]
enum MissingContextBaseline {
BareTurnContext,
WorldStatePatch,
CompactedSnapshot,
}
#[test_case(MissingContextBaseline::BareTurnContext; "bare turn context")]
#[test_case(MissingContextBaseline::WorldStatePatch; "patch without a full snapshot")]
#[test_case(MissingContextBaseline::CompactedSnapshot; "full snapshot before compaction")]
#[tokio::test]
async fn record_initial_history_resumed_bare_turn_context_does_not_seed_reference_context_item() {
async fn record_initial_history_requires_surviving_full_snapshot_without_user_turn(
baseline: MissingContextBaseline,
) {
let (session, turn_context) = make_session_and_context().await;
let previous_context_item = turn_context.to_turn_context_item();
let rollout_items = vec![RolloutItem::TurnContext(previous_context_item.clone())];
let mut rollout_items = match baseline {
MissingContextBaseline::BareTurnContext => Vec::new(),
MissingContextBaseline::WorldStatePatch => {
vec![RolloutItem::WorldState(WorldStateItem::patch(object!({})))]
}
MissingContextBaseline::CompactedSnapshot => vec![
RolloutItem::WorldState(WorldStateItem::full(object!({}))),
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
mcp_resource_origins: None,
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
],
};
rollout_items.push(RolloutItem::TurnContext(
turn_context.to_turn_context_item(),
));
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {

244
codex-rs/core/tests/suite/subagent_notifications.rs
View File

@@ -33,6 +33,8 @@ use core_test_support::responses::assert_parent_turn;
use core_test_support::responses::assert_root_turn;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_completed_with_tokens;
use core_test_support::responses::ev_function_call;
use core_test_support::responses::ev_function_call_with_namespace;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::ev_tool_search_call;
@@ -1122,6 +1124,248 @@ async fn spawned_child_receives_forked_parent_context(
Ok(())
}
#[derive(Clone, Copy, Debug)]
enum GrandchildParentContext {
FullHistory,
LastTurn,
NoHistory,
Compacted,
}
#[test_case(GrandchildParentContext::FullHistory, ThreadHistoryMode::Legacy; "legacy full history")]
#[test_case(GrandchildParentContext::LastTurn, ThreadHistoryMode::Legacy; "legacy last turn")]
#[test_case(GrandchildParentContext::NoHistory, ThreadHistoryMode::Legacy; "legacy no history")]
#[test_case(GrandchildParentContext::FullHistory, ThreadHistoryMode::Paginated; "paginated full history")]
#[test_case(GrandchildParentContext::LastTurn, ThreadHistoryMode::Paginated; "paginated last turn")]
#[test_case(GrandchildParentContext::NoHistory, ThreadHistoryMode::Paginated; "paginated no history")]
#[test_case(GrandchildParentContext::Compacted, ThreadHistoryMode::Legacy; "legacy full history after compaction")]
#[test_case(GrandchildParentContext::Compacted, ThreadHistoryMode::Paginated; "paginated full history after compaction")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn grandchild_full_fork_preserves_context_baseline(
parent_context: GrandchildParentContext,
history_mode: ThreadHistoryMode,
) -> Result<()> {
skip_if_no_network!(Ok(()));
const ROOT_PROMPT: &str = "root: delegate the context check";
const CHILD_TASK: &str = "child: delegate the context check";
const GRANDCHILD_TASK: &str = "grandchild: check inherited context";
const ROOT_CALL: &str = "root-context-baseline-spawn";
const CHILD_CALL: &str = "child-context-baseline-spawn";
const INSTRUCTIONS: &str = "UNIQUE_CONTEXT_BASELINE_DEVELOPER_INSTRUCTIONS";
const COMPACT_PROMPT: &str = "CONTEXT_BASELINE_COMPACTION_PROMPT";
const COMPACT_SUMMARY: &str = "CONTEXT_BASELINE_COMPACTION_SUMMARY";
const PRELUDE_CALL: &str = "context-baseline-prelude-call";
let server = start_mock_server().await;
let (parent_fork_turns, compact_parent) = match parent_context {
GrandchildParentContext::FullHistory => ("all", false),
GrandchildParentContext::LastTurn => ("1", false),
GrandchildParentContext::NoHistory => ("none", false),
GrandchildParentContext::Compacted => ("all", true),
};
let root_spawn_args = serde_json::to_string(&json!({
"task_name": "child",
"message": CHILD_TASK,
"fork_turns": parent_fork_turns,
}))?;
let root_log = mount_sse_once_match(
&server,
|req: &wiremock::Request| {
body_contains(req, ROOT_PROMPT)
&& !body_contains(req, CHILD_TASK)
&& !body_contains(req, GRANDCHILD_TASK)
&& !body_contains(req, ROOT_CALL)
},
sse(vec![
ev_response_created("baseline-root"),
ev_function_call_with_namespace(
ROOT_CALL,
MULTI_AGENT_V2_NAMESPACE,
"spawn_agent",
&root_spawn_args,
),
ev_completed("baseline-root"),
]),
)
.await;
let child_spawn_args = serde_json::to_string(&json!({
"task_name": "grandchild",
"message": GRANDCHILD_TASK,
"fork_turns": "all",
}))?;
if compact_parent {
mount_sse_once_match(
&server,
|req: &wiremock::Request| {
body_contains(req, CHILD_TASK)
&& !body_contains(req, ROOT_CALL)
&& !body_contains(req, PRELUDE_CALL)
&& !body_contains(req, COMPACT_SUMMARY)
},
sse(vec![
ev_response_created("baseline-prelude"),
ev_function_call(
PRELUDE_CALL,
"update_plan",
r#"{"plan":[{"step":"Check inherited context","status":"in_progress"}]}"#,
),
ev_completed_with_tokens("baseline-prelude", /*total_tokens*/ 250_000),
]),
)
.await;
mount_sse_once_match(
&server,
|req: &wiremock::Request| body_contains(req, COMPACT_PROMPT),
sse(vec![
ev_response_created("baseline-compaction"),
ev_assistant_message("baseline-summary", COMPACT_SUMMARY),
ev_completed("baseline-compaction"),
]),
)
.await;
}
let child_log = mount_sse_once_match(
&server,
move |req: &wiremock::Request| {
body_contains(
req,
if compact_parent {
COMPACT_SUMMARY
} else {
CHILD_TASK
},
) && !body_contains(req, GRANDCHILD_TASK)
&& !body_contains(req, ROOT_CALL)
&& !body_contains(req, CHILD_CALL)
&& !body_contains(req, COMPACT_PROMPT)
},
sse(vec![
ev_response_created("baseline-child"),
ev_function_call_with_namespace(
CHILD_CALL,
MULTI_AGENT_V2_NAMESPACE,
"spawn_agent",
&child_spawn_args,
),
ev_completed("baseline-child"),
]),
)
.await;
let grandchild_log = mount_sse_once_match(
&server,
|req: &wiremock::Request| {
body_contains(req, GRANDCHILD_TASK) && !body_contains(req, CHILD_CALL)
},
sse(vec![
ev_response_created("baseline-grandchild"),
ev_assistant_message("baseline-grandchild-answer", "done"),
ev_completed("baseline-grandchild"),
]),
)
.await;
let _parent_followups = mount_sse_sequence(
&server,
vec![
sse(vec![ev_completed("baseline-parent-finished-1")]),
sse(vec![ev_completed("baseline-parent-finished-2")]),
],
)
.await;
let test = test_codex()
.with_history_mode(history_mode)
.with_config(move |config| {
config
.features
.enable(Feature::Collab)
.expect("test config should allow feature update");
config
.features
.enable(Feature::MultiAgentV2)
.expect("test config should allow feature update");
config.model = Some(V2_DEFAULT_MODEL.to_string());
config.agent_default_subagent_model = Some(V2_DEFAULT_MODEL.to_string());
config.developer_instructions = Some(INSTRUCTIONS.to_string());
if compact_parent {
// Use local compaction so the test controls the replacement history.
config.model_provider.name = "test-provider".to_string();
config
.features
.disable(Feature::RemoteCompactionV2)
.expect("test config should allow feature update");
config.compact_prompt = Some(COMPACT_PROMPT.to_string());
config.model_auto_compact_token_limit = Some(200_000);
config.model_context_window = Some(1_000_000);
}
})
.build_with_auto_env(&server)
.await?;
test.submit_turn(ROOT_PROMPT).await?;
let root_request = root_log.single_request();
let mut descendant_requests = Vec::new();
for (mock, agent_name) in [
(&child_log, "/root/child"),
(&grandchild_log, "/root/child/grandchild"),
] {
let request = timeout(Duration::from_secs(/*secs*/ 10), async {
loop {
let request = mock.requests().into_iter().find(|request| {
request.body_json()["client_metadata"]["x-codex-turn-metadata"]
.as_str()
.and_then(|text| serde_json::from_str::<Value>(text).ok())
.is_some_and(|metadata| metadata["agent_name"] == agent_name)
&& (!compact_parent || request.body_contains_text(COMPACT_SUMMARY))
});
if let Some(request) = request {
break request;
}
sleep(Duration::from_millis(/*millis*/ 10)).await;
}
})
.await?;
let thread_id = ThreadId::from_string(
request.body_json()["client_metadata"]["thread_id"]
.as_str()
.expect("descendant thread id"),
)?;
let thread = test.thread_manager.get_thread(thread_id).await?;
timeout(Duration::from_secs(/*secs*/ 10), async {
while !matches!(thread.agent_status().await, AgentStatus::Completed(_)) {
sleep(Duration::from_millis(/*millis*/ 10)).await;
}
})
.await?;
descendant_requests.push(request);
}
let context_counts = [
&root_request,
&descendant_requests[0],
&descendant_requests[1],
]
.map(|request| {
(
request
.message_input_texts("developer")
.iter()
.filter(|text| text.contains(INSTRUCTIONS))
.count(),
request
.message_input_texts("user")
.iter()
.filter(|text| text.contains("<environment_context>"))
.count(),
)
});
assert_eq!(
context_counts,
[(1, 1); 3],
"Initial context should appear once per agent: {parent_context:?}, {history_mode:?}"
);
assert!(!descendant_requests[1].body_contains_text(CHILD_TASK));
Ok(())
}
#[derive(Clone, Copy)]
enum FullHistoryV2ModelSelection {
ConfiguredDefault,