mirror of
https://github.com/openai/codex.git
synced 2026-11-03 16:23:51 +00:00
Use copy-on-write storage for history snapshots (#34390)
## Why Cloning a `ContextManager` deep-copies every `ResponseItem`, even when callers only need a read-only history snapshot. ## What changed - Store history items in an `Arc<Vec<ResponseItem>>` so cloned snapshots share their items until one is mutated. - Reuse an owned vector when consuming an unshared snapshot and clone it only when other snapshots still reference it. - Estimate session token counts from a snapshot after releasing the session-state lock. ## Testing Added a regression test that verifies cloned histories share their backing items before mutation and diverge without affecting each other afterward. GitOrigin-RevId: 24605888209ae9ff78cadb04d9c36463f7e2dbe0
This commit is contained in:
1 parent
6bf4845b60
commit
45ac251e17
4 files changed
+52
-20
No files matched your search
@@ -32,13 +32,15 @@ use codex_utils_output_truncation::truncate_function_output_items_with_policy;
|
||||
use codex_utils_output_truncation::truncate_text;
|
||||
use std::num::NonZeroUsize;
|
||||
use std::ops::Deref;
|
||||
use std::sync::Arc;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
/// Transcript of thread history
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub(crate) struct ContextManager {
|
||||
/// The oldest items are at the beginning of the vector.
|
||||
items: Vec<ResponseItem>,
|
||||
/// The oldest items are at the beginning of the vector. Snapshots share the vector until a
|
||||
/// caller needs to mutate it, avoiding deep copies for read-only history consumers.
|
||||
items: Arc<Vec<ResponseItem>>,
|
||||
/// Bumped whenever history is rewritten, such as compaction or rollback.
|
||||
history_version: u64,
|
||||
token_info: Option<TokenUsageInfo>,
|
||||
@@ -60,7 +62,7 @@ pub(crate) struct ContextManager {
|
||||
impl ContextManager {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
items: Vec::new(),
|
||||
items: Arc::new(Vec::new()),
|
||||
history_version: 0,
|
||||
token_info: TokenUsageInfo::new_or_append(
|
||||
&None, &None, /*model_context_window*/ None,
|
||||
@@ -130,8 +132,8 @@ impl ContextManager {
|
||||
continue;
|
||||
}
|
||||
|
||||
let processed = self.process_item(item_ref, policy);
|
||||
self.items.push(processed);
|
||||
let processed = Self::process_item(item_ref, policy);
|
||||
Arc::make_mut(&mut self.items).push(processed);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -140,7 +142,7 @@ impl ContextManager {
|
||||
/// is stripped from messages and tool outputs according to `input_modalities`.
|
||||
pub(crate) fn for_prompt(mut self, input_modalities: &[InputModality]) -> Vec<ResponseItem> {
|
||||
self.normalize_history(input_modalities);
|
||||
self.items
|
||||
Arc::unwrap_or_clone(self.items)
|
||||
}
|
||||
|
||||
/// Returns raw items in the history.
|
||||
@@ -150,7 +152,7 @@ impl ContextManager {
|
||||
|
||||
/// Returns raw items in the history and consumes the snapshot.
|
||||
pub(crate) fn into_raw_items(self) -> Vec<ResponseItem> {
|
||||
self.items
|
||||
Arc::unwrap_or_clone(self.items)
|
||||
}
|
||||
|
||||
pub(crate) fn history_version(&self) -> u64 {
|
||||
@@ -188,17 +190,18 @@ impl ContextManager {
|
||||
if !self.items.is_empty() {
|
||||
// Remove the oldest item (front of the list). Items are ordered from
|
||||
// oldest → newest, so index 0 is the first entry recorded.
|
||||
let removed = self.items.remove(0);
|
||||
let items = Arc::make_mut(&mut self.items);
|
||||
let removed = items.remove(0);
|
||||
// If the removed item participates in a call/output pair, also remove
|
||||
// its corresponding counterpart to keep the invariants intact without
|
||||
// running a full normalization pass.
|
||||
normalize::remove_corresponding_for(&mut self.items, &removed);
|
||||
normalize::remove_corresponding_for(items, &removed);
|
||||
self.world_state_baseline = None;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn replace(&mut self, items: Vec<ResponseItem>) {
|
||||
self.items = items;
|
||||
self.items = Arc::new(items);
|
||||
self.history_version = self.history_version.saturating_add(1);
|
||||
self.world_state_baseline = None;
|
||||
}
|
||||
@@ -227,7 +230,7 @@ impl ContextManager {
|
||||
let snapshot = self.items.clone();
|
||||
let user_positions = user_message_positions(&snapshot);
|
||||
let Some(&first_instruction_turn_idx) = user_positions.first() else {
|
||||
self.replace(snapshot);
|
||||
self.replace(Arc::unwrap_or_clone(snapshot));
|
||||
return;
|
||||
};
|
||||
|
||||
@@ -323,20 +326,22 @@ impl ContextManager {
|
||||
/// 2. every output has a corresponding call entry
|
||||
/// 3. unsupported image and audio content is stripped from messages and tool outputs
|
||||
fn normalize_history(&mut self, input_modalities: &[InputModality]) {
|
||||
let items = Arc::make_mut(&mut self.items);
|
||||
|
||||
// all function/tool calls must have a corresponding output
|
||||
normalize::ensure_call_outputs_present(&mut self.items);
|
||||
normalize::ensure_call_outputs_present(items);
|
||||
|
||||
// all outputs must have a corresponding function/tool call
|
||||
normalize::remove_orphan_outputs(&mut self.items);
|
||||
normalize::remove_orphan_outputs(items);
|
||||
|
||||
// strip images when model does not support them
|
||||
normalize::strip_images_when_unsupported(input_modalities, &mut self.items);
|
||||
normalize::strip_images_when_unsupported(input_modalities, items);
|
||||
|
||||
// strip audio when model does not support it
|
||||
normalize::strip_audio_when_unsupported(input_modalities, &mut self.items);
|
||||
normalize::strip_audio_when_unsupported(input_modalities, items);
|
||||
}
|
||||
|
||||
fn process_item(&self, item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem {
|
||||
fn process_item(item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem {
|
||||
let policy_with_serialization_budget = policy * 1.2;
|
||||
match item {
|
||||
ResponseItem::FunctionCallOutput {
|
||||
|
||||
@@ -444,6 +444,31 @@ fn for_prompt_preserves_inter_agent_assistant_messages() {
|
||||
assert_eq!(history.for_prompt(&default_input_modalities()), vec![item]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cloned_history_shares_items_until_mutated() {
|
||||
let first = assistant_msg(&"first ".repeat(1_024));
|
||||
let second = assistant_msg("second");
|
||||
let history = create_history_with_items(vec![first.clone()]);
|
||||
let mut snapshot = history.clone();
|
||||
|
||||
assert!(std::ptr::eq(
|
||||
history.raw_items().as_ptr(),
|
||||
snapshot.raw_items().as_ptr()
|
||||
));
|
||||
|
||||
snapshot.record_items(
|
||||
std::slice::from_ref(&second),
|
||||
TruncationPolicy::Tokens(10_000),
|
||||
);
|
||||
|
||||
assert!(!std::ptr::eq(
|
||||
history.raw_items().as_ptr(),
|
||||
snapshot.raw_items().as_ptr()
|
||||
));
|
||||
assert_eq!(history.raw_items(), std::slice::from_ref(&first));
|
||||
assert_eq!(snapshot.raw_items(), &[first, second]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_last_n_user_turns_treats_inter_agent_assistant_messages_as_instruction_turns() {
|
||||
let first_turn = user_input_text_msg("first");
|
||||
|
||||
@@ -62,8 +62,10 @@ pub(crate) async fn build_realtime_startup_context(
|
||||
) -> Option<String> {
|
||||
let config = sess.get_config().await;
|
||||
let cwd = config.cwd.clone();
|
||||
let history = sess.clone_history().await;
|
||||
let current_thread_section = build_current_thread_section(history.raw_items());
|
||||
let current_thread_section = {
|
||||
let history = sess.clone_history().await;
|
||||
build_current_thread_section(history.raw_items())
|
||||
};
|
||||
let recent_threads = load_recent_threads(sess).await;
|
||||
let recent_work_section = build_recent_work_section(&cwd, &recent_threads).await;
|
||||
let workspace_section = build_workspace_section_with_user_root(&cwd, home_dir()).await;
|
||||
|
||||
@@ -1217,8 +1217,8 @@ impl Session {
|
||||
&self,
|
||||
turn_context: &TurnContext,
|
||||
) -> Option<i64> {
|
||||
let state = self.state.lock().await;
|
||||
state.history.estimate_token_count(turn_context)
|
||||
let history = self.clone_history().await;
|
||||
history.estimate_token_count(turn_context)
|
||||
}
|
||||
|
||||
pub(crate) async fn get_base_instructions(&self) -> BaseInstructions {
|
||||
|
||||
Reference in new issue
Block a user