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

fix(core): decouple websocket cache invalidation from window lineage

Browse Source
This commit is contained in:
Ningyi Xie
2026-06-01 12:21:57 -07:00
parent 1cd1ee3268
commit 8d59fc9c52
4 changed files with 50 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
codex-rs/core/src/client.rs
View File

@@ -172,8 +172,6 @@ pub(crate) struct CompactConversationRequestSettings {
struct ModelClientState {
session_id: SessionId,
thread_id: ThreadId,
// Monotonic context epoch used to avoid reusing remotely cached state after history rewrites.
// Forks may inherit a nonzero epoch when their initial history already contains rewrites.
window_generation: AtomicU64,
installation_id: String,
provider: SharedModelProvider,
@@ -187,7 +185,7 @@ struct ModelClientState {
include_attestation: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
disable_websockets: AtomicBool,
cached_websocket_session: StdMutex<WebsocketSession>,
cached_websocket_session: StdMutex<CachedWebsocketSession>,
}
/// Resolved API client setup for a single request attempt.
@@ -244,7 +242,7 @@ pub struct ModelClient {
pub struct ModelClientSession {
client: ModelClient,
websocket_session: WebsocketSession,
window_generation: u64,
websocket_cache_generation: u64,
/// Turn state for sticky routing.
///
/// This is an `OnceLock` that stores the turn state value received from the server
@@ -273,6 +271,12 @@ struct WebsocketSession {
connection_reused: StdMutex<bool>,
}
#[derive(Debug, Default)]
struct CachedWebsocketSession {
generation: u64,
session: WebsocketSession,
}
impl WebsocketSession {
fn set_connection_reused(&self, connection_reused: bool) {
*self
@@ -361,7 +365,7 @@ impl ModelClient {
include_attestation,
attestation_provider,
disable_websockets: AtomicBool::new(false),
cached_websocket_session: StdMutex::new(WebsocketSession::default()),
cached_websocket_session: StdMutex::new(CachedWebsocketSession::default()),
}),
prompt_cache_key_override: None,
}
@@ -386,10 +390,11 @@ impl ModelClient {
/// This constructor does not perform network I/O itself; the session opens a websocket lazily
/// when the first stream request is issued.
pub fn new_session(&self) -> ModelClientSession {
let (websocket_cache_generation, websocket_session) = self.take_cached_websocket_session();
ModelClientSession {
client: self.clone(),
websocket_session: self.take_cached_websocket_session(),
window_generation: self.state.window_generation.load(Ordering::Relaxed),
websocket_session,
websocket_cache_generation,
turn_state: Arc::new(OnceLock::new()),
}
}
@@ -402,12 +407,12 @@ impl ModelClient {
self.state
.window_generation
.store(window_generation, Ordering::Relaxed);
self.store_cached_websocket_session(WebsocketSession::default());
self.invalidate_cached_websocket_session();
}
pub(crate) fn advance_window_generation(&self) {
self.state.window_generation.fetch_add(1, Ordering::Relaxed);
self.store_cached_websocket_session(WebsocketSession::default());
self.invalidate_cached_websocket_session();
}
pub(crate) fn current_window_id(&self) -> String {
@@ -416,21 +421,36 @@ impl ModelClient {
format!("{thread_id}:{window_generation}")
}
fn take_cached_websocket_session(&self) -> WebsocketSession {
pub(crate) fn invalidate_cached_websocket_session(&self) {
let mut cached_websocket_session = self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
std::mem::take(&mut *cached_websocket_session)
cached_websocket_session.generation = cached_websocket_session.generation.saturating_add(1);
cached_websocket_session.session = WebsocketSession::default();
}
fn store_cached_websocket_session(&self, websocket_session: WebsocketSession) {
*self
fn take_cached_websocket_session(&self) -> (u64, WebsocketSession) {
let mut cached_websocket_session = self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session;
.unwrap_or_else(std::sync::PoisonError::into_inner);
let generation = cached_websocket_session.generation;
let session = std::mem::take(&mut cached_websocket_session.session);
(generation, session)
}
fn store_cached_websocket_session(&self, generation: u64, websocket_session: WebsocketSession) {
let mut cached_websocket_session = self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if cached_websocket_session.generation == generation {
cached_websocket_session.session = websocket_session;
}
}
pub(crate) fn force_http_fallback(
@@ -450,7 +470,7 @@ impl ModelClient {
);
}
self.store_cached_websocket_session(WebsocketSession::default());
self.invalidate_cached_websocket_session();
activated
}
@@ -979,13 +999,9 @@ impl ModelClient {
impl Drop for ModelClientSession {
fn drop(&mut self) {
// A history invalidation may advance the generation before an older turn session drops.
if self.window_generation != self.client.state.window_generation.load(Ordering::Relaxed) {
return;
}
let websocket_session = std::mem::take(&mut self.websocket_session);
self.client
.store_cached_websocket_session(websocket_session);
.store_cached_websocket_session(self.websocket_cache_generation, websocket_session);
}
}

4
codex-rs/core/src/session/handlers.rs
View File

@@ -568,7 +568,9 @@ pub async fn thread_rollback(sess: &Arc<Session>, sub_id: String, num_turns: u32
sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice())
.await;
sess.recompute_token_usage(turn_context.as_ref()).await;
sess.services.model_client.advance_window_generation();
sess.services
.model_client
.invalidate_cached_websocket_session();
sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())])
.await;

35
codex-rs/core/src/session/session.rs
View File

@@ -16,25 +16,6 @@ use codex_protocol::protocol::TurnEnvironmentSelections;
use std::sync::OnceLock;
use tokio::sync::Semaphore;
// This is a request-cache invalidation epoch, not a hash of logical context. A rollback must use a
// new generation even when it reconstructs a previously seen prefix, because the remote window
// with the old ID may already include requests from the rolled-back suffix.
fn window_generation_from_history_invalidations(items: &[RolloutItem]) -> u64 {
u64::try_from(
items
.iter()
.filter(|item| {
matches!(
item,
RolloutItem::Compacted(_)
| RolloutItem::EventMsg(EventMsg::ThreadRolledBack(_))
)
})
.count(),
)
.unwrap_or(u64::MAX)
}
/// Context for an initialized model agent
///
/// A session has at most 1 running task at a time, and can be interrupted by user input.
@@ -537,13 +518,15 @@ impl Session {
InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id,
};
let window_generation = match &initial_history {
InitialHistory::Resumed(resumed_history) => {
window_generation_from_history_invalidations(&resumed_history.history)
}
InitialHistory::Forked(history) => {
window_generation_from_history_invalidations(history)
}
InitialHistory::New | InitialHistory::Cleared => 0,
InitialHistory::Resumed(resumed_history) => u64::try_from(
resumed_history
.history
.iter()
.filter(|item| matches!(item, RolloutItem::Compacted(_)))
.count(),
)
.unwrap_or(u64::MAX),
InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => 0,
};
// Kick off independent async setup tasks in parallel to reduce startup latency.
//

88
codex-rs/core/tests/suite/window_headers.rs
View File

@@ -21,8 +21,7 @@ use pretty_assertions::assert_eq;
use std::sync::Arc;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_empty_fork() -> Result<()>
{
async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_fork() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
@@ -103,7 +102,7 @@ async fn window_id_advances_after_compact_persists_on_resume_and_resets_on_empty
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn window_id_advances_after_rollback_and_persists_on_resume() -> Result<()> {
async fn window_id_stays_stable_after_rollback_and_resume() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
@@ -157,91 +156,8 @@ async fn window_id_advances_after_rollback_and_persists_on_resume() -> Result<()
vec![
(initial_thread_id.clone(), 0),
(initial_thread_id.clone(), 0),
(initial_thread_id.clone(), 1),
(initial_thread_id, 1),
]
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn forked_compacted_history_keeps_window_generation_on_resume() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let request_log = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_assistant_message("msg-1", "first reply"),
ev_completed("resp-1"),
]),
sse(vec![
ev_assistant_message("msg-2", "summary"),
ev_completed("resp-2"),
]),
sse(vec![ev_completed("resp-3")]),
sse(vec![ev_completed("resp-4")]),
],
)
.await;
let mut builder = test_codex().with_config(|config| {
config.model_provider.name = "Non-OpenAI Model provider".to_string();
config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string());
});
let initial = builder.build(&server).await?;
let initial_thread = Arc::clone(&initial.codex);
let rollout_path = initial
.session_configured
.rollout_path
.clone()
.expect("rollout path");
submit_user_turn(&initial_thread, "before compact").await?;
submit_compact_turn(&initial_thread).await?;
shutdown_thread(&initial_thread).await?;
let forked = initial
.thread_manager
.fork_thread(
/*snapshot*/ usize::MAX,
initial.config.clone(),
rollout_path,
/*thread_source*/ None,
/*persist_extended_history*/ false,
/*parent_trace*/ None,
)
.await?;
let fork_rollout_path = forked
.session_configured
.rollout_path
.clone()
.expect("fork rollout path");
submit_user_turn(&forked.thread, "after fork").await?;
shutdown_thread(&forked.thread).await?;
let resumed = builder
.resume(&server, initial.home.clone(), fork_rollout_path)
.await?;
submit_user_turn(&resumed.codex, "after fork resume").await?;
shutdown_thread(&resumed.codex).await?;
let requests = request_log.requests();
assert_eq!(requests.len(), 4, "expected four model requests");
let window_ids = requests.iter().map(window_id_parts).collect::<Vec<_>>();
let initial_thread_id = window_ids[0].0.clone();
let forked_thread_id = window_ids[2].0.clone();
assert_ne!(forked_thread_id, initial_thread_id);
assert_eq!(
window_ids,
vec![
(initial_thread_id.clone(), 0),
(initial_thread_id, 0),
(forked_thread_id.clone(), 1),
(forked_thread_id, 1),
]
);