openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-26 18:38:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Persist agent task state and stabilize CI

Browse Source
This commit is contained in:
adrian
2026-04-14 10:13:16 -07:00
parent c66f085580
commit 1176180367
24 changed files with 238 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control.rs
+1
View File
@@ -116,6 +116,7 @@ fn keep_forked_rollout_item(item: &RolloutItem) -> bool {
| ResponseItem::Compaction { .. }
| ResponseItem::Other,
) => false,
RolloutItem::SessionState(_) => false,
RolloutItem::Compacted(_)
| RolloutItem::EventMsg(_)
| RolloutItem::SessionMeta(_)
codex-rs/core/src/agent/control_tests.rs
+1 -1
View File
@@ -186,7 +186,7 @@ async fn wait_for_subagent_notification(parent_thread: &Arc<CodexThread>) -> boo
sleep(Duration::from_millis(25)).await;
}
};
timeout(Duration::from_secs(2), wait).await.is_ok()
timeout(Duration::from_secs(5), wait).await.is_ok()
}
async fn persist_thread_for_tree_resume(thread: &Arc<CodexThread>, message: &str) {
codex-rs/core/src/agent_identity.rs
+10 -3
View File
@@ -153,14 +153,16 @@ impl AgentIdentityManager {
Ok(stored_identity)
}
pub(crate) async fn task_matches_current_binding(&self, task: &RegisteredAgentTask) -> bool {
pub(crate) async fn task_matches_current_identity(&self, task: &RegisteredAgentTask) -> bool {
if !self.feature_enabled {
return false;
}
self.current_auth_binding()
self.current_stored_identity()
.await
.is_some_and(|(_, binding)| task.matches_binding(&binding))
.is_some_and(|stored_identity| {
stored_identity.agent_runtime_id == task.agent_runtime_id
})
}
async fn current_auth_binding(&self) -> Option<(CodexAuth, AgentIdentityBinding)> {
@@ -177,6 +179,11 @@ impl AgentIdentityManager {
binding.map(|binding| (auth, binding))
}
async fn current_stored_identity(&self) -> Option<StoredAgentIdentity> {
let (auth, binding) = self.current_auth_binding().await?;
self.load_stored_identity(&auth, &binding).ok().flatten()
}
async fn register_agent_identity(
&self,
binding: &AgentIdentityBinding,
codex-rs/core/src/agent_identity/task_registration.rs
+15 -27
View File
@@ -2,6 +2,7 @@ use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use codex_protocol::protocol::SessionAgentTask;
use crypto_box::SecretKey as Curve25519SecretKey;
use ed25519_dalek::Signer as _;
use serde::Deserialize;
@@ -16,9 +17,6 @@ const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct RegisteredAgentTask {
pub(crate) binding_id: String,
pub(crate) chatgpt_account_id: String,
pub(crate) chatgpt_user_id: Option<String>,
pub(crate) agent_runtime_id: String,
pub(crate) task_id: String,
pub(crate) registered_at: String,
@@ -82,9 +80,6 @@ impl AgentIdentityManager {
.await
.with_context(|| format!("failed to parse agent task response from {url}"))?;
let registered_task = RegisteredAgentTask {
binding_id: stored_identity.binding_id.clone(),
chatgpt_account_id: stored_identity.chatgpt_account_id.clone(),
chatgpt_user_id: stored_identity.chatgpt_user_id.clone(),
agent_runtime_id: stored_identity.agent_runtime_id.clone(),
task_id: decrypt_task_id_response(
&stored_identity,
@@ -107,18 +102,20 @@ impl AgentIdentityManager {
}
impl RegisteredAgentTask {
pub(super) fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool {
binding.matches_parts(
&self.binding_id,
&self.chatgpt_account_id,
self.chatgpt_user_id.as_deref(),
)
pub(crate) fn to_session_agent_task(&self) -> SessionAgentTask {
SessionAgentTask {
agent_runtime_id: self.agent_runtime_id.clone(),
task_id: self.task_id.clone(),
registered_at: self.registered_at.clone(),
}
}
pub(crate) fn has_same_binding(&self, other: &Self) -> bool {
self.binding_id == other.binding_id
&& self.chatgpt_account_id == other.chatgpt_account_id
&& self.chatgpt_user_id == other.chatgpt_user_id
pub(crate) fn from_session_agent_task(task: SessionAgentTask) -> Self {
Self {
agent_runtime_id: task.agent_runtime_id,
task_id: task.task_id,
registered_at: task.registered_at,
}
}
}
@@ -242,9 +239,6 @@ mod tests {
assert_eq!(
task,
RegisteredAgentTask {
binding_id: "chatgpt-account-account-123".to_string(),
chatgpt_account_id: "account-123".to_string(),
chatgpt_user_id: Some("user-123".to_string()),
agent_runtime_id: "agent-123".to_string(),
task_id: "task_123".to_string(),
registered_at: task.registered_at.clone(),
@@ -331,9 +325,6 @@ mod tests {
assert_eq!(
task,
RegisteredAgentTask {
binding_id: "chatgpt-account-account-123".to_string(),
chatgpt_account_id: "account-123".to_string(),
chatgpt_user_id: Some("user-123".to_string()),
agent_runtime_id: "agent-123".to_string(),
task_id: "task_123".to_string(),
registered_at: task.registered_at.clone(),
@@ -342,7 +333,7 @@ mod tests {
}
#[tokio::test]
async fn task_matches_current_binding_rejects_stale_auth_binding() {
async fn task_matches_current_identity_rejects_stale_registered_identity() {
let auth_manager =
AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456")));
let manager = AgentIdentityManager::new_for_tests(
@@ -352,15 +343,12 @@ mod tests {
SessionSource::Cli,
);
let task = RegisteredAgentTask {
binding_id: "chatgpt-account-account-123".to_string(),
chatgpt_account_id: "account-123".to_string(),
chatgpt_user_id: Some("user-123".to_string()),
agent_runtime_id: "agent-123".to_string(),
task_id: "task_123".to_string(),
registered_at: "2026-03-23T12:00:00Z".to_string(),
};
assert!(!manager.task_matches_current_binding(&task).await);
assert!(!manager.task_matches_current_identity(&task).await);
}
async fn mount_human_biscuit(
codex-rs/core/src/codex.rs
+78 -24
View File
@@ -121,6 +121,7 @@ use codex_protocol::protocol::RawResponseItemEvent;
use codex_protocol::protocol::ReviewRequest;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SessionStateUpdate;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnContextItem;
@@ -845,6 +846,7 @@ pub(crate) struct Session {
pub(crate) services: SessionServices,
js_repl: Arc<JsReplHandle>,
next_internal_sub_id: AtomicU64,
agent_task_registration_lock: Mutex<()>,
}
#[derive(Clone, Debug)]
@@ -1529,7 +1531,68 @@ impl Session {
handlers::shutdown(self, self.next_internal_sub_id()).await;
}
async fn cached_agent_task_for_current_binding(&self) -> Option<RegisteredAgentTask> {
fn latest_persisted_agent_task(
rollout_items: &[RolloutItem],
) -> Option<Option<RegisteredAgentTask>> {
rollout_items.iter().rev().find_map(|item| match item {
RolloutItem::SessionState(update) => Some(
update
.agent_task
.clone()
.map(RegisteredAgentTask::from_session_agent_task),
),
_ => None,
})
}
async fn restore_persisted_agent_task(&self, rollout_items: &[RolloutItem]) {
let Some(agent_task) = Self::latest_persisted_agent_task(rollout_items).flatten() else {
return;
};
let mut state = self.state.lock().await;
state.set_agent_task(agent_task);
}
async fn persist_agent_task_update(&self, agent_task: Option<&RegisteredAgentTask>) {
self.persist_rollout_items(&[RolloutItem::SessionState(SessionStateUpdate {
agent_task: agent_task.map(RegisteredAgentTask::to_session_agent_task),
})])
.await;
}
async fn clear_cached_agent_task(&self, agent_task: &RegisteredAgentTask) {
let cleared = {
let mut state = self.state.lock().await;
if state.agent_task().as_ref() == Some(agent_task) {
state.clear_agent_task();
true
} else {
false
}
};
if cleared {
self.persist_agent_task_update(/*agent_task*/ None).await;
}
}
async fn cache_agent_task(&self, agent_task: RegisteredAgentTask) -> RegisteredAgentTask {
let changed = {
let mut state = self.state.lock().await;
if state.agent_task().as_ref() == Some(&agent_task) {
false
} else {
state.set_agent_task(agent_task.clone());
true
}
};
if changed {
self.persist_agent_task_update(Some(&agent_task)).await;
}
agent_task
}
async fn cached_agent_task_for_current_identity(&self) -> Option<RegisteredAgentTask> {
let agent_task = {
let state = self.state.lock().await;
state.agent_task()
@@ -1538,7 +1601,7 @@ impl Session {
if self
.services
.agent_identity_manager
.task_matches_current_binding(&agent_task)
.task_matches_current_identity(&agent_task)
.await
{
debug!(
@@ -1552,17 +1615,19 @@ impl Session {
debug!(
agent_runtime_id = %agent_task.agent_runtime_id,
task_id = %agent_task.task_id,
"discarding cached agent task because auth binding changed"
"discarding cached agent task because the registered agent identity changed"
);
let mut state = self.state.lock().await;
if state.agent_task().as_ref() == Some(&agent_task) {
state.clear_agent_task();
}
self.clear_cached_agent_task(&agent_task).await;
None
}
async fn ensure_agent_task_registered(&self) -> anyhow::Result<Option<RegisteredAgentTask>> {
if let Some(agent_task) = self.cached_agent_task_for_current_binding().await {
if let Some(agent_task) = self.cached_agent_task_for_current_identity().await {
return Ok(Some(agent_task));
}
let _guard = self.agent_task_registration_lock.lock().await;
if let Some(agent_task) = self.cached_agent_task_for_current_identity().await {
return Ok(Some(agent_task));
}
@@ -1575,31 +1640,18 @@ impl Session {
if !self
.services
.agent_identity_manager
.task_matches_current_binding(&agent_task)
.task_matches_current_identity(&agent_task)
.await
{
debug!(
agent_runtime_id = %agent_task.agent_runtime_id,
task_id = %agent_task.task_id,
"discarding newly registered agent task because auth binding changed"
"discarding newly registered agent task because the registered agent identity changed"
);
continue;
}
{
let mut state = self.state.lock().await;
if let Some(existing_agent_task) = state.agent_task() {
if existing_agent_task.has_same_binding(&agent_task) {
return Ok(Some(existing_agent_task));
}
debug!(
agent_runtime_id = %existing_agent_task.agent_runtime_id,
task_id = %existing_agent_task.task_id,
"replacing cached agent task because auth binding changed"
);
}
state.set_agent_task(agent_task.clone());
}
let agent_task = self.cache_agent_task(agent_task).await;
info!(
thread_id = %self.conversation_id,
@@ -2221,6 +2273,7 @@ impl Session {
services,
js_repl,
next_internal_sub_id: AtomicU64::new(0),
agent_task_registration_lock: Mutex::new(()),
});
if let Some(network_policy_decider_session) = network_policy_decider_session {
let mut guard = network_policy_decider_session.write().await;
@@ -2490,6 +2543,7 @@ impl Session {
}
InitialHistory::Resumed(resumed_history) => {
let rollout_items = resumed_history.history;
self.restore_persisted_agent_task(&rollout_items).await;
let previous_turn_settings = self
.apply_rollout_reconstruction(&turn_context, &rollout_items)
.await;
codex-rs/core/src/codex/rollout_reconstruction.rs
+4 -1
View File
@@ -207,7 +207,9 @@ impl Session {
active_segment.get_or_insert_with(ActiveReplaySegment::default);
active_segment.counts_as_user_turn |= is_user_turn_boundary(response_item);
}
RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {}
RolloutItem::EventMsg(_)
| RolloutItem::SessionMeta(_)
| RolloutItem::SessionState(_) => {}
}
if base_replacement_history.is_some()
@@ -275,6 +277,7 @@ impl Session {
history.drop_last_n_user_turns(rollback.num_turns);
}
RolloutItem::EventMsg(_)
| RolloutItem::SessionState(_)
| RolloutItem::TurnContext(_)
| RolloutItem::SessionMeta(_) => {}
}
codex-rs/core/src/codex_tests.rs
+63
View File
@@ -1077,6 +1077,67 @@ async fn record_initial_history_reconstructs_resumed_transcript() {
assert_eq!(expected, history.raw_items());
}
#[tokio::test]
async fn record_initial_history_restores_latest_persisted_agent_task() {
let (session, _turn_context) = make_session_and_context().await;
let expected = RegisteredAgentTask {
agent_runtime_id: "agent-123".to_string(),
task_id: "task-123".to_string(),
registered_at: "2026-03-23T12:00:00Z".to_string(),
};
let rollout_items = vec![
RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate {
agent_task: Some(expected.to_session_agent_task()),
}),
RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate {
agent_task: None,
}),
RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate {
agent_task: Some(expected.to_session_agent_task()),
}),
];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: rollout_items,
rollout_path: PathBuf::from("/tmp/resume.jsonl"),
}))
.await;
assert_eq!(session.state.lock().await.agent_task(), Some(expected));
}
#[tokio::test]
async fn record_initial_history_honors_cleared_persisted_agent_task() {
let (session, _turn_context) = make_session_and_context().await;
let rollout_items = vec![
RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate {
agent_task: Some(
RegisteredAgentTask {
agent_runtime_id: "agent-123".to_string(),
task_id: "task-123".to_string(),
registered_at: "2026-03-23T12:00:00Z".to_string(),
}
.to_session_agent_task(),
),
}),
RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate {
agent_task: None,
}),
];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: rollout_items,
rollout_path: PathBuf::from("/tmp/resume.jsonl"),
}))
.await;
assert_eq!(session.state.lock().await.agent_task(), None);
}
#[tokio::test]
async fn record_initial_history_new_defers_initial_context_until_first_turn() {
let (session, _turn_context) = make_session_and_context().await;
@@ -2954,6 +3015,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
services,
js_repl,
next_internal_sub_id: AtomicU64::new(0),
agent_task_registration_lock: Mutex::new(()),
};
(session, turn_context)
@@ -3804,6 +3866,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
services,
js_repl,
next_internal_sub_id: AtomicU64::new(0),
agent_task_registration_lock: Mutex::new(()),
});
(session, turn_context, rx_event)
codex-rs/core/src/state/session_tests.rs
+1 -7
View File
@@ -39,10 +39,7 @@ async fn set_agent_task_persists_plaintext_task_for_session_reuse() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
let agent_task = RegisteredAgentTask {
binding_id: "chatgpt-account-account-123".to_string(),
chatgpt_account_id: "account-123".to_string(),
chatgpt_user_id: Some("user-123".to_string()),
agent_runtime_id: "agent_123".to_string(),
agent_runtime_id: "agent-123".to_string(),
task_id: "task_123".to_string(),
registered_at: "2026-03-23T12:00:00Z".to_string(),
};
@@ -57,9 +54,6 @@ async fn clear_agent_task_removes_cached_task() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
let agent_task = RegisteredAgentTask {
binding_id: "chatgpt-account-account-123".to_string(),
chatgpt_account_id: "account-123".to_string(),
chatgpt_user_id: Some("user-123".to_string()),
agent_runtime_id: "agent_123".to_string(),
task_id: "task_123".to_string(),
registered_at: "2026-03-23T12:00:00Z".to_string(),