From 7a5fdedf8996cdf49df7e58460f6ee36bd495c81 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 22:40:47 +0300 Subject: [PATCH 1/7] Persist provider profile dispatch leases Agent: quintilianus --- .../app-server/src/bespoke_event_handling.rs | 1 + .../request_processors/thread_lifecycle.rs | 42 +- .../thread_mailbox_dispatcher_runtime.rs | 407 ++++++++++- .../thread_schedule_runtime.rs | 306 ++++++++- .../usage_profile_broker.rs | 642 +++++++++++------- codex-rs/app-server/src/thread_state.rs | 32 + codex-rs/core/src/codex_thread.rs | 20 + .../migrations/0068_usage_profile_leases.sql | 31 + codex-rs/state/src/lib.rs | 6 + codex-rs/state/src/runtime.rs | 7 + .../state/src/runtime/usage_profile_leases.rs | 546 +++++++++++++++ 11 files changed, 1723 insertions(+), 317 deletions(-) create mode 100644 codex-rs/state/migrations/0068_usage_profile_leases.sql create mode 100644 codex-rs/state/src/runtime/usage_profile_leases.rs diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 6b20d2545a..082e6004ce 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2844,6 +2844,7 @@ mod tests { schedule_id: "schedule-1".to_string(), run_id: "run-1".to_string(), lease_id: "lease-1".to_string(), + usage_profile_lease: None, goal_id: None, state_db, }, diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 3ac34b5883..5751917df7 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -325,7 +325,12 @@ pub(super) async fn ensure_listener_task_running( event.msg, EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_) | EventMsg::Error(_) ); - let (raw_events_enabled, tracked_scheduled_run, turn_error) = { + let ( + raw_events_enabled, + tracked_scheduled_run, + tracked_usage_profile_lease, + turn_error, + ) = { let mut thread_state = thread_state.lock().await; thread_state.track_current_turn_event(&event.id, &event.msg); let tracked_scheduled_run = if terminal_event { @@ -333,12 +338,18 @@ pub(super) async fn ensure_listener_task_running( } else { None }; + let tracked_usage_profile_lease = if terminal_event { + thread_state.take_usage_profile_turn_lease(&event.id) + } else { + None + }; let turn_error = terminal_event .then(|| thread_state.turn_summary.last_error.clone()) .flatten(); ( thread_state.experimental_raw_events, tracked_scheduled_run, + tracked_usage_profile_lease, turn_error, ) }; @@ -418,6 +429,9 @@ pub(super) async fn ensure_listener_task_running( ) .await; } + if let Some(usage_profile_lease) = tracked_usage_profile_lease { + finish_usage_profile_turn_lease(usage_profile_lease).await; + } } unloading_watchers_open = unloading_state.wait_for_unloading_trigger() => { if !unloading_watchers_open { @@ -455,15 +469,33 @@ pub(super) async fn ensure_listener_task_running( } } - let mut thread_state = thread_state.lock().await; - if thread_state.listener_generation == listener_generation { - thread_state_manager.unregister_listener_command_tx(conversation_id); - thread_state.clear_listener(); + let abandoned_usage_profile_leases = { + let mut thread_state = thread_state.lock().await; + if thread_state.listener_generation == listener_generation { + thread_state_manager.unregister_listener_command_tx(conversation_id); + let leases = thread_state.drain_usage_profile_turn_leases(); + thread_state.clear_listener(); + leases + } else { + Vec::new() + } + }; + for lease in abandoned_usage_profile_leases { + finish_usage_profile_turn_lease(lease).await; } }); Ok(()) } +async fn finish_usage_profile_turn_lease(tracked: crate::thread_state::UsageProfileTurnLease) { + tracked.heartbeat_cancel.cancel(); + super::usage_profile_broker::release_dispatch_auth_profile_lease( + &tracked.state_db, + &tracked.lease, + ) + .await; +} + async fn heartbeat_local_active_session( listener_task_context: &ListenerTaskContext, conversation_id: ThreadId, diff --git a/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs b/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs index 92f9173068..46bb61b98f 100644 --- a/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs @@ -241,12 +241,15 @@ impl ThreadMailboxDispatcherRuntime { let lease_id = claim.attempt.lease_id.clone(); let result = self.deliver_claim(&claim).await; let mut wake_thread_id = None; + let mut wake_usage_profile_lease = None; let durable_transition_ok = match result { MailboxDispatchResult::Delivered { receipt, wake_thread_id: wake, + usage_profile_lease, } => { wake_thread_id = wake; + wake_usage_profile_lease = usage_profile_lease; self.ack_dispatch_claim(state_db, &claim, receipt).await } MailboxDispatchResult::Retry { error, retry_at } => { @@ -262,6 +265,8 @@ impl ThreadMailboxDispatcherRuntime { } }; if !durable_transition_ok { + release_optional_mailbox_usage_profile_lease(state_db, wake_usage_profile_lease.take()) + .await; warn!( message_id = %message_id, "leaving mailbox target dispatch lease to expire because durable transition did not complete" @@ -273,7 +278,10 @@ impl ThreadMailboxDispatcherRuntime { // crashed before this point, the row would still be claimable and simply // redelivered, without the target ever consuming a duplicate. if let Some(thread_id) = wake_thread_id { - self.spawn_pending_work_wake(thread_id); + self.spawn_pending_work_wake(thread_id, wake_usage_profile_lease.take()); + } else { + release_optional_mailbox_usage_profile_lease(state_db, wake_usage_profile_lease.take()) + .await; } // Refresh is best-effort: a per-peer heartbeat failure elsewhere must not // keep this target's dispatch lease alive. Release the lease once this @@ -452,6 +460,7 @@ impl ThreadMailboxDispatcherRuntime { }; let target_peer_id = claim.message.target_thread_id.to_string(); let mut resumed_target = false; + let mut usage_profile_lease = None; let (registry, target_peer) = match registry.get_active(target_peer_id.as_str(), freshness) { Ok(peer) => (registry, peer), @@ -460,16 +469,32 @@ impl ThreadMailboxDispatcherRuntime { return mailbox_target_not_loaded(err); } MailboxLocalDeliveryPolicy::ResumeAndTrigger => { - if let Err(err) = self - .resume_mailbox_target(claim.message.target_thread_id) + let usage_profile_owner = + super::usage_profile_broker::usage_profile_lease_owner( + "mailbox", + claim.attempt.lease_id.as_str(), + ); + match self + .resume_mailbox_target( + claim.message.target_thread_id, + usage_profile_owner.as_str(), + ) .await { - return err.into_dispatch_result(); + Ok(lease) => usage_profile_lease = lease, + Err(err) => return err.into_dispatch_result(), } resumed_target = true; let registry = match self.active_peer_directory.snapshot(freshness.now).await { Ok(registry) => registry, Err(err) => { + if let Some(state_db) = self.state_db.as_ref() { + release_optional_mailbox_usage_profile_lease( + state_db, + usage_profile_lease.take(), + ) + .await; + } return MailboxDispatchResult::retry(format!( "failed to read active session directory after resume: {err}" )); @@ -478,7 +503,16 @@ impl ThreadMailboxDispatcherRuntime { let target_peer = match registry.get_active(target_peer_id.as_str(), freshness) { Ok(peer) => peer, - Err(err) => return mailbox_target_not_loaded(err), + Err(err) => { + if let Some(state_db) = self.state_db.as_ref() { + release_optional_mailbox_usage_profile_lease( + state_db, + usage_profile_lease.take(), + ) + .await; + } + return mailbox_target_not_loaded(err); + } }; (registry, target_peer) } @@ -509,6 +543,7 @@ impl ThreadMailboxDispatcherRuntime { Ok(ActiveChannelDeliveryOutcome::Delivered { .. }) => { MailboxDispatchResult::Delivered { wake_thread_id: mailbox_dispatch_wake_target(delivery, target_peer.thread_id), + usage_profile_lease, receipt: serde_json::json!({ "delivery": if resumed_target { "resumed" } else { "live" }, "recipientPeerId": target_peer.peer_id, @@ -518,44 +553,181 @@ impl ThreadMailboxDispatcherRuntime { }), } } - Ok(ActiveChannelDeliveryOutcome::NotLoaded { .. }) => MailboxDispatchResult::retry( - "target thread became unloaded during mailbox dispatch", - ), - Ok(ActiveChannelDeliveryOutcome::Unsupported { .. }) => MailboxDispatchResult::retry( - "target peer is active but does not support local mailbox dispatch", - ), - Err(err) => MailboxDispatchResult::retry(format!("mailbox dispatch failed: {err}")), + Ok(ActiveChannelDeliveryOutcome::NotLoaded { .. }) => { + if let Some(state_db) = self.state_db.as_ref() { + release_optional_mailbox_usage_profile_lease( + state_db, + usage_profile_lease.take(), + ) + .await; + } + MailboxDispatchResult::retry( + "target thread became unloaded during mailbox dispatch", + ) + } + Ok(ActiveChannelDeliveryOutcome::Unsupported { .. }) => { + if let Some(state_db) = self.state_db.as_ref() { + release_optional_mailbox_usage_profile_lease( + state_db, + usage_profile_lease.take(), + ) + .await; + } + MailboxDispatchResult::retry( + "target peer is active but does not support local mailbox dispatch", + ) + } + Err(err) => { + if let Some(state_db) = self.state_db.as_ref() { + release_optional_mailbox_usage_profile_lease( + state_db, + usage_profile_lease.take(), + ) + .await; + } + MailboxDispatchResult::retry(format!("mailbox dispatch failed: {err}")) + } } } - fn spawn_pending_work_wake(&self, thread_id: ThreadId) { + fn spawn_pending_work_wake( + &self, + thread_id: ThreadId, + usage_profile_lease: Option, + ) { let thread_manager = Arc::clone(&self.thread_manager); + let thread_state_manager = self.thread_state_manager.clone(); + let auth_manager = Arc::clone(&self.auth_manager); + let config = Arc::clone(&self.config); + let state_db = self.state_db.clone(); let cancel_token = self.cancel_token.clone(); + let tasks = self.tasks.clone(); self.tasks.spawn(async move { + let mut usage_profile_lease = usage_profile_lease; + let mut last_provider_recheck = None; for _ in 0..MAILBOX_PENDING_WORK_WAKE_ATTEMPTS { tokio::select! { - _ = cancel_token.cancelled() => return, + _ = cancel_token.cancelled() => break, _ = tokio::time::sleep(MAILBOX_PENDING_WORK_WAKE_INTERVAL) => {} } match thread_manager.get_thread(thread_id).await { Ok(thread) => { - if thread.maybe_start_turn_for_pending_work().await { + if usage_profile_lease.is_some() + && last_provider_recheck.map_or(true, |checked_at: Instant| { + checked_at.elapsed() >= Duration::from_secs(1) + }) + { + let Some(state_db) = state_db.as_ref() else { + break; + }; + let lease = usage_profile_lease + .take() + .expect("usage profile lease checked as present"); + usage_profile_lease = + match super::usage_profile_broker::recheck_dispatch_auth_profile_lease( + &auth_manager, + &config, + state_db, + &lease, + ) + .await + { + Ok(lease) => Some(lease), + Err(_) => return, + }; + last_provider_recheck = Some(Instant::now()); + } + + let turn_id = Uuid::now_v7().to_string(); + let heartbeat_cancel = CancellationToken::new(); + let mut heartbeat_tracked_lease = None; + if let (Some(lease), Some(state_db)) = + (usage_profile_lease.take(), state_db.as_ref()) + { + let tracked = crate::thread_state::UsageProfileTurnLease { + lease, + state_db: state_db.clone(), + heartbeat_cancel: heartbeat_cancel.clone(), + }; + heartbeat_tracked_lease = Some(tracked.clone()); + let previous = thread_state_manager + .thread_state(thread_id) + .await + .lock() + .await + .track_usage_profile_turn_lease( + turn_id.clone(), + tracked, + ); + if let Some(previous) = previous { + previous.heartbeat_cancel.cancel(); + super::usage_profile_broker::release_dispatch_auth_profile_lease( + &previous.state_db, + &previous.lease, + ) + .await; + } + } + + if thread + .maybe_start_turn_for_pending_work_with_sub_id(turn_id.clone()) + .await + { + if let Some(tracked) = heartbeat_tracked_lease { + let heartbeat_thread = Arc::clone(&thread); + let heartbeat_turn_id = turn_id.clone(); + let heartbeat_config = Arc::clone(&config); + tasks.spawn(async move { + heartbeat_mailbox_usage_profile_turn( + heartbeat_thread, + heartbeat_turn_id, + tracked, + heartbeat_config, + ) + .await; + }); + } return; } + + if let Some(tracked) = thread_state_manager + .thread_state(thread_id) + .await + .lock() + .await + .take_usage_profile_turn_lease(turn_id.as_str()) + { + tracked.heartbeat_cancel.cancel(); + usage_profile_lease = Some(tracked.lease); + } } - Err(CodexErr::ThreadNotFound(_)) => return, + Err(CodexErr::ThreadNotFound(_)) => break, Err(err) => { warn!("failed to wake mailbox pending work for thread {thread_id}: {err}"); - return; + break; } } } + if let (Some(state_db), Some(lease)) = (state_db.as_ref(), usage_profile_lease) { + super::usage_profile_broker::release_dispatch_auth_profile_lease( + state_db, + &lease, + ) + .await; + } }); } - async fn resume_mailbox_target(&self, thread_id: ThreadId) -> Result<(), MailboxResumeError> { + async fn resume_mailbox_target( + &self, + thread_id: ThreadId, + usage_profile_owner: &str, + ) -> Result, MailboxResumeError> { match self.thread_manager.get_thread(thread_id).await { - Ok(thread) => return self.ensure_mailbox_listener(thread_id, thread).await, + Ok(thread) => { + self.ensure_mailbox_listener(thread_id, thread).await?; + return Ok(None); + } Err(CodexErr::ThreadNotFound(_)) => {} Err(err) => { return Err(MailboxResumeError::Failed(format!( @@ -604,10 +776,15 @@ impl ThreadMailboxDispatcherRuntime { let broker_decision = super::usage_profile_broker::resolve_dispatch_auth_profile( &self.auth_manager, &self.config, + state_db, + usage_profile_owner, typesafe_overrides.auth_profile.clone(), ) .await; - if let Some(profile) = broker_decision.selected_profile.as_ref() { + let mut usage_profile_lease = broker_decision.lease; + if let Some(profile) = broker_decision.selected_profile.as_ref() + && usage_profile_lease.is_some() + { tracing::debug!( thread_id = %thread_id, auth_profile = %profile, @@ -615,9 +792,9 @@ impl ThreadMailboxDispatcherRuntime { "usage profile broker selected auth profile for mailbox resume" ); typesafe_overrides.auth_profile = Some(Some(profile.clone())); - } else if let Some(retry_at) = broker_decision.retry_at - && let Some(retry_at) = broker_retry_at_datetime(&self.config, retry_at) - { + } else if self.config.auth_profile_auto_switch.enabled { + let retry_at = + mailbox_usage_profile_retry_at(&self.config, broker_decision.retry_at, Utc::now()); tracing::debug!( thread_id = %thread_id, retry_at = %retry_at.to_rfc3339(), @@ -627,7 +804,7 @@ impl ThreadMailboxDispatcherRuntime { return Err(MailboxResumeError::UsageProfileWait { retry_at, error: format!( - "all eligible auth profiles are exhausted; retrying mailbox resume after {}", + "no healthy unused auth profile is available; retrying mailbox resume after {}", retry_at.to_rfc3339() ), }); @@ -635,12 +812,44 @@ impl ThreadMailboxDispatcherRuntime { let config = self .config_manager .load_for_cwd(request_overrides, typesafe_overrides, history_cwd) - .await - .map_err(|err| { - MailboxResumeError::Failed(format!( + .await; + let config = match config { + Ok(config) => config, + Err(err) => { + release_optional_mailbox_usage_profile_lease(state_db, usage_profile_lease.take()) + .await; + return Err(MailboxResumeError::Failed(format!( "failed to load config for mailbox resume: {err}" - )) - })?; + ))); + } + }; + if let Some(lease) = usage_profile_lease.take() { + usage_profile_lease = + match super::usage_profile_broker::recheck_dispatch_auth_profile_lease( + &self.auth_manager, + &self.config, + state_db, + &lease, + ) + .await + { + Ok(lease) => Some(lease), + Err(decision) => { + let retry_at = mailbox_usage_profile_retry_at( + &self.config, + decision.retry_at, + Utc::now(), + ); + return Err(MailboxResumeError::UsageProfileWait { + retry_at, + error: format!( + "no healthy unused auth profile is available; retrying mailbox resume after {}", + retry_at.to_rfc3339() + ), + }); + } + }; + } let thread = self .thread_manager .resume_thread_with_history( @@ -649,12 +858,23 @@ impl ThreadMailboxDispatcherRuntime { Arc::clone(&self.auth_manager), /*parent_trace*/ None, ) - .await - .map(|new_thread| new_thread.thread) - .map_err(|err| { - MailboxResumeError::Failed(format!("failed to resume mailbox target: {err}")) - })?; - self.ensure_mailbox_listener(thread_id, thread).await + .await; + let thread = match thread { + Ok(new_thread) => new_thread.thread, + Err(err) => { + release_optional_mailbox_usage_profile_lease(state_db, usage_profile_lease.take()) + .await; + return Err(MailboxResumeError::Failed(format!( + "failed to resume mailbox target: {err}" + ))); + } + }; + if let Err(err) = self.ensure_mailbox_listener(thread_id, thread).await { + release_optional_mailbox_usage_profile_lease(state_db, usage_profile_lease.take()) + .await; + return Err(err); + } + Ok(usage_profile_lease) } async fn ensure_mailbox_listener( @@ -687,12 +907,69 @@ impl ThreadMailboxDispatcherRuntime { } } +async fn heartbeat_mailbox_usage_profile_turn( + thread: Arc, + turn_id: String, + tracked: crate::thread_state::UsageProfileTurnLease, + config: Arc, +) { + let interval = Duration::from_secs( + config + .auth_profile_auto_switch + .heartbeat_interval_secs + .max(1), + ); + let mut lease = tracked.lease.clone(); + loop { + tokio::select! { + _ = tracked.heartbeat_cancel.cancelled() => return, + _ = tokio::time::sleep(interval) => {} + } + match super::usage_profile_broker::renew_dispatch_auth_profile_lease( + &tracked.state_db, + &config, + &lease, + ) + .await + { + Ok(Some(renewed)) => lease = renewed, + Ok(None) => { + warn!( + turn_id = %turn_id, + "mailbox turn lost durable usage profile lease ownership" + ); + let _ = thread + .abort_turn_if_active( + turn_id.as_str(), + codex_protocol::protocol::TurnAbortReason::Interrupted, + ) + .await; + return; + } + Err(err) => { + warn!( + turn_id = %turn_id, + "failed to heartbeat mailbox turn usage profile lease: {err}" + ); + let _ = thread + .abort_turn_if_active( + turn_id.as_str(), + codex_protocol::protocol::TurnAbortReason::Interrupted, + ) + .await; + return; + } + } + } +} + enum MailboxDispatchResult { Delivered { receipt: serde_json::Value, /// Thread to wake for pending work once the claim is durably acked, or /// `None` for queue-only deliveries that must not trigger a turn. wake_thread_id: Option, + usage_profile_lease: Option, }, Retry { error: String, @@ -908,11 +1185,39 @@ fn truncate_mailbox_dispatch_error(error: String) -> String { .collect() } -fn broker_retry_at_datetime(config: &Config, retry_at: i64) -> Option> { +fn broker_retry_at_datetime( + config: &Config, + retry_at: i64, + now: DateTime, +) -> Option> { let retry_at = DateTime::::from_timestamp(retry_at, /*nsecs*/ 0)?; let buffer_secs = i64::try_from(config.usage_self_heal.reset_retry_buffer_secs).ok()?; let retry_at = retry_at + ChronoDuration::seconds(buffer_secs); - (retry_at > Utc::now()).then_some(retry_at) + (retry_at > now).then_some(retry_at) +} + +fn mailbox_usage_profile_retry_at( + config: &Config, + retry_at: Option, + now: DateTime, +) -> DateTime { + retry_at + .and_then(|retry_at| broker_retry_at_datetime(config, retry_at, now)) + .unwrap_or_else(|| { + now + ChronoDuration::seconds( + i64::try_from(config.usage_self_heal.reset_retry_buffer_secs.max(30)) + .unwrap_or(i64::MAX), + ) + }) +} + +async fn release_optional_mailbox_usage_profile_lease( + state_db: &StateDbHandle, + lease: Option, +) { + if let Some(lease) = lease { + super::usage_profile_broker::release_dispatch_auth_profile_lease(state_db, &lease).await; + } } async fn apply_persisted_mailbox_resume_metadata( @@ -1058,6 +1363,34 @@ mod tests { } } + #[test] + fn delivered_mailbox_result_carries_usage_profile_lease_to_the_wake_edge() { + let now = DateTime::::from_timestamp(1_700_000_000, 0) + .expect("test timestamp should be valid"); + let usage_profile_lease = codex_state::UsageProfileLease { + lease_id: "usage-lease-mailbox-1".to_string(), + identity_sha256: "a".repeat(64), + owner_id: "mailbox:dispatch-lease-1".to_string(), + profile_name: "profile-a".to_string(), + acquired_at: now, + heartbeat_at: now, + expires_at: now + ChronoDuration::seconds(120), + }; + let result = MailboxDispatchResult::Delivered { + receipt: serde_json::json!({"delivery": "resumed"}), + wake_thread_id: None, + usage_profile_lease: Some(usage_profile_lease.clone()), + }; + + match result { + MailboxDispatchResult::Delivered { + usage_profile_lease: Some(carried), + .. + } => assert_eq!(carried, usage_profile_lease), + _ => panic!("delivered mailbox result should preserve the durable profile lease"), + } + } + #[test] fn parse_persisted_permission_profile_accepts_current_metadata_format() { let profile = PermissionProfile::Disabled; diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 278ca500ce..17afba651b 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -240,15 +240,22 @@ impl ThreadScheduleRuntime { let claim_auth_profile = self .claim_auth_profile(&state_db, thread_id, &claim.schedule) .await; + let usage_profile_owner = super::usage_profile_broker::usage_profile_lease_owner( + "schedule", + claim.run.lease_id.as_str(), + ); let broker_decision = super::usage_profile_broker::resolve_dispatch_auth_profile( &self.auth_manager, &self.config, + &state_db, + usage_profile_owner.as_str(), claim_auth_profile.clone(), ) .await; - let claim_auth_profile = match schedule_auth_profile_after_broker_decision( + let selection = match schedule_auth_profile_after_broker_decision( claim_auth_profile, broker_decision, + self.config.auth_profile_auto_switch.enabled, self.config.usage_self_heal.reset_retry_buffer_secs, Utc::now(), ) { @@ -260,19 +267,31 @@ impl ThreadScheduleRuntime { }); } }; + let claim_auth_profile = selection.auth_profile; + let mut usage_profile_lease = selection.lease; let thread = self .load_or_resume_thread(thread_id, claim_auth_profile.clone()) + .await; + let thread = match thread { + Ok(thread) => thread, + Err(error) => { + release_optional_usage_profile_lease(&state_db, usage_profile_lease.take()).await; + return Err(ScheduleSubmitError { + error, + goal_id: None, + }); + } + }; + if let Err(error) = self + .ensure_schedule_listener(thread_id, thread.clone()) .await - .map_err(|error| ScheduleSubmitError { - error, - goal_id: None, - })?; - self.ensure_schedule_listener(thread_id, thread.clone()) - .await - .map_err(|error| ScheduleSubmitError { + { + release_optional_usage_profile_lease(&state_db, usage_profile_lease.take()).await; + return Err(ScheduleSubmitError { error, goal_id: None, - })?; + }); + } let thread_state = self.thread_state_manager.thread_state(thread_id).await; let listener_command_tx = { let thread_state = thread_state.lock().await; @@ -287,13 +306,18 @@ impl ThreadScheduleRuntime { objective, listener_command_tx.clone(), ) - .await - .map_err(|error| { + .await; + let scheduled_goal_id = match scheduled_goal_id { + Ok(goal_id) => goal_id, + Err(error) => { let goal_id = error .downcast_ref::() .map(|held| held.goal_id.clone()); - ScheduleSubmitError { error, goal_id } - })?; + release_optional_usage_profile_lease(&state_db, usage_profile_lease.take()) + .await; + return Err(ScheduleSubmitError { error, goal_id }); + } + }; ( scheduled_goal_thread_prompt( objective, @@ -319,6 +343,31 @@ impl ThreadScheduleRuntime { claim_auth_profile, ); let turn_id = Uuid::now_v7().to_string(); + let usage_profile_lease = match usage_profile_lease.take() { + Some(lease) => { + match super::usage_profile_broker::recheck_dispatch_auth_profile_lease( + &self.auth_manager, + &self.config, + &state_db, + &lease, + ) + .await + { + Ok(lease) => Some(lease), + Err(decision) => { + return Err(ScheduleSubmitError { + error: anyhow::Error::new(schedule_usage_profile_wait( + decision, + self.config.usage_self_heal.reset_retry_buffer_secs, + Utc::now(), + )), + goal_id: scheduled_goal_id, + }); + } + } + } + None => None, + }; let run_start = state_db .thread_schedules() @@ -335,6 +384,7 @@ impl ThreadScheduleRuntime { let run = match run_start { Ok(Some(run)) => run, Ok(None) => { + release_optional_usage_profile_lease(&state_db, usage_profile_lease).await; return Err(ScheduleSubmitError { error: anyhow::anyhow!( "claimed schedule run {} no longer owns the current unexpired lease", @@ -344,15 +394,26 @@ impl ThreadScheduleRuntime { }); } Err(error) => { + release_optional_usage_profile_lease(&state_db, usage_profile_lease).await; return Err(ScheduleSubmitError { error, goal_id: scheduled_goal_id, }); } }; - let ownership_lost = match self.start_lease_heartbeat(state_db.clone(), &run).await { + let ownership_lost = match self + .start_lease_heartbeat( + state_db.clone(), + &run, + Arc::clone(&thread), + turn_id.clone(), + usage_profile_lease.as_ref(), + ) + .await + { Ok(Some(ownership_lost)) => ownership_lost, Ok(None) => { + release_optional_usage_profile_lease(&state_db, usage_profile_lease).await; return Err(ScheduleSubmitError { error: anyhow::anyhow!( "claimed schedule run {} lost lease ownership before dispatch readiness", @@ -362,6 +423,7 @@ impl ThreadScheduleRuntime { }); } Err(error) => { + release_optional_usage_profile_lease(&state_db, usage_profile_lease).await; return Err(ScheduleSubmitError { error, goal_id: scheduled_goal_id, @@ -376,6 +438,7 @@ impl ThreadScheduleRuntime { schedule_id: run.schedule_id.clone(), run_id: run.run_id.clone(), lease_id: run.lease_id.clone(), + usage_profile_lease: usage_profile_lease.clone(), goal_id: run.goal_id.clone(), state_db: state_db.clone(), }, @@ -405,10 +468,11 @@ impl ThreadScheduleRuntime { { Ok(Some(start_result)) => start_result, Ok(None) => { - thread_state + let scheduled_run = thread_state .lock() .await .take_scheduled_run(turn_id.as_str()); + release_tracked_schedule_usage_profile_lease(scheduled_run).await; return Err(ScheduleSubmitError { error: anyhow::anyhow!( "claimed schedule run {} lost lease ownership before turn submission", @@ -418,10 +482,11 @@ impl ThreadScheduleRuntime { }); } Err(error) => { - thread_state + let scheduled_run = thread_state .lock() .await .take_scheduled_run(turn_id.as_str()); + release_tracked_schedule_usage_profile_lease(scheduled_run).await; return Err(ScheduleSubmitError { error, goal_id: scheduled_goal_id, @@ -429,10 +494,11 @@ impl ThreadScheduleRuntime { } }; if let Err(err) = start_result { - thread_state + let scheduled_run = thread_state .lock() .await .take_scheduled_run(turn_id.as_str()); + release_tracked_schedule_usage_profile_lease(scheduled_run).await; if let Some(deferral) = schedule_deferral_for_idle_rejection(&err, Utc::now()) { return Err(ScheduleSubmitError { error: anyhow::Error::new(deferral), @@ -580,6 +646,9 @@ impl ThreadScheduleRuntime { &self, state_db: StateDbHandle, run: &codex_state::ThreadScheduleRun, + thread: Arc, + turn_id: String, + usage_profile_lease: Option<&codex_state::UsageProfileLease>, ) -> anyhow::Result> { let schedule_id = run.schedule_id.clone(); let run_id = run.run_id.clone(); @@ -599,9 +668,28 @@ impl ThreadScheduleRuntime { ownership_lost.cancel(); return Ok(None); } + if let Some(lease) = usage_profile_lease + && super::usage_profile_broker::renew_dispatch_auth_profile_lease( + &state_db, + &self.config, + lease, + ) + .await? + .is_none() + { + ownership_lost.cancel(); + return Ok(None); + } let cancel_token = self.cancel_token.clone(); let heartbeat_ownership_lost = ownership_lost.clone(); + let usage_profile_lease = usage_profile_lease.cloned(); + let config = Arc::clone(&self.config); + let heartbeat_interval = schedule_lease_heartbeat_interval( + usage_profile_lease + .as_ref() + .map(|_| config.auth_profile_auto_switch.heartbeat_interval_secs), + ); self.tasks.spawn(async move { loop { tokio::select! { @@ -609,7 +697,7 @@ impl ThreadScheduleRuntime { heartbeat_ownership_lost.cancel(); break; } - _ = tokio::time::sleep(SCHEDULE_LEASE_HEARTBEAT_INTERVAL) => {} + _ = tokio::time::sleep(heartbeat_interval) => {} } match state_db .thread_schedules() @@ -636,6 +724,35 @@ impl ThreadScheduleRuntime { break; } } + if let Some(lease) = usage_profile_lease.as_ref() { + match super::usage_profile_broker::renew_dispatch_auth_profile_lease( + &state_db, &config, lease, + ) + .await + { + Ok(Some(_)) => {} + Ok(None) => { + heartbeat_ownership_lost.cancel(); + break; + } + Err(err) => { + warn!( + schedule_id = %schedule_id, + "failed to refresh scheduled usage profile lease: {err}" + ); + heartbeat_ownership_lost.cancel(); + break; + } + } + } + } + if usage_profile_lease.is_some() { + let _ = thread + .abort_turn_if_active( + turn_id.as_str(), + codex_protocol::protocol::TurnAbortReason::Interrupted, + ) + .await; } }); Ok(Some(ownership_lost)) @@ -1177,6 +1294,11 @@ pub(super) async fn finish_scheduled_run_after_turn( "failed to finish scheduled thread run: {err}" ), } + release_optional_usage_profile_lease( + &scheduled_run.state_db, + scheduled_run.usage_profile_lease, + ) + .await; } pub(super) async fn recover_scheduled_run_for_terminal_turn( @@ -1195,6 +1317,7 @@ pub(super) async fn recover_scheduled_run_for_terminal_turn( schedule_id: run.schedule_id, run_id: run.run_id, lease_id: run.lease_id, + usage_profile_lease: None, goal_id: run.goal_id, state_db: state_db.clone(), })) @@ -1232,7 +1355,7 @@ impl std::fmt::Display for ScheduleUsageProfileWait { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "all eligible auth profiles are exhausted; retrying scheduled run after {}", + "no healthy unused auth profile is available; retrying scheduled run after {}", self.retry_at.to_rfc3339() ) } @@ -1314,22 +1437,74 @@ fn schedule_deferral_for_idle_rejection( }) } +struct ScheduleAuthProfileSelection { + auth_profile: Option>, + lease: Option, +} + +async fn release_optional_usage_profile_lease( + state_db: &StateDbHandle, + lease: Option, +) { + if let Some(lease) = lease { + super::usage_profile_broker::release_dispatch_auth_profile_lease(state_db, &lease).await; + } +} + +async fn release_tracked_schedule_usage_profile_lease( + scheduled_run: Option, +) { + if let Some(scheduled_run) = scheduled_run { + release_optional_usage_profile_lease( + &scheduled_run.state_db, + scheduled_run.usage_profile_lease, + ) + .await; + } +} + fn schedule_auth_profile_after_broker_decision( current_auth_profile: Option>, decision: super::usage_profile_broker::UsageProfileBrokerDecision, + lease_required: bool, reset_retry_buffer_secs: u64, now: DateTime, -) -> Result>, ScheduleUsageProfileWait> { - if let Some(profile) = decision.selected_profile { - return Ok(Some(Some(profile))); +) -> Result { + if let (Some(profile), Some(lease)) = (decision.selected_profile, decision.lease) { + return Ok(ScheduleAuthProfileSelection { + auth_profile: Some(Some(profile)), + lease: Some(lease), + }); } - if let Some(retry_at) = decision.retry_at - && let Some(retry_at) = - schedule_broker_retry_at_datetime(reset_retry_buffer_secs, retry_at, now) - { - return Err(ScheduleUsageProfileWait { retry_at }); + if !lease_required { + return Ok(ScheduleAuthProfileSelection { + auth_profile: current_auth_profile, + lease: None, + }); } - Ok(current_auth_profile) + Err(schedule_usage_profile_wait( + decision, + reset_retry_buffer_secs, + now, + )) +} + +fn schedule_usage_profile_wait( + decision: super::usage_profile_broker::UsageProfileBrokerDecision, + reset_retry_buffer_secs: u64, + now: DateTime, +) -> ScheduleUsageProfileWait { + let retry_at = decision + .retry_at + .and_then(|retry_at| { + schedule_broker_retry_at_datetime(reset_retry_buffer_secs, retry_at, now) + }) + .unwrap_or_else(|| { + now + ChronoDuration::seconds( + i64::try_from(reset_retry_buffer_secs.max(30)).unwrap_or(i64::MAX), + ) + }); + ScheduleUsageProfileWait { retry_at } } fn schedule_broker_retry_at_datetime( @@ -1353,6 +1528,14 @@ fn duration_seconds_i64(duration: Duration) -> i64 { i64::try_from(duration.as_secs()).unwrap_or(i64::MAX) } +fn schedule_lease_heartbeat_interval( + usage_profile_heartbeat_interval_secs: Option, +) -> Duration { + usage_profile_heartbeat_interval_secs.map_or(SCHEDULE_LEASE_HEARTBEAT_INTERVAL, |seconds| { + SCHEDULE_LEASE_HEARTBEAT_INTERVAL.min(Duration::from_secs(seconds.max(1))) + }) +} + async fn finish_scheduled_run_state( state_db: &StateDbHandle, schedule_id: &str, @@ -2645,8 +2828,18 @@ mod tests { #[test] fn schedule_auth_profile_uses_broker_selected_profile_for_pinned_loop() { + let now = at(/*seconds*/ 1_700_000_000); let decision = super::usage_profile_broker::UsageProfileBrokerDecision { selected_profile: Some("account003".to_string()), + lease: Some(codex_state::UsageProfileLease { + lease_id: "usage-lease-1".to_string(), + identity_sha256: "a".repeat(64), + owner_id: "schedule:lease-1".to_string(), + profile_name: "account003".to_string(), + acquired_at: now, + heartbeat_at: now, + expires_at: now + chrono::Duration::seconds(120), + }), retry_at: None, reason: super::usage_profile_broker::UsageProfileBrokerDecisionReason::SelectedHealthyProfile, }; @@ -2654,12 +2847,17 @@ mod tests { let resolved = schedule_auth_profile_after_broker_decision( Some(Some("account001".to_string())), decision, + /*lease_required*/ true, /*reset_retry_buffer_secs*/ 30, - at(/*seconds*/ 1_700_000_000), + now, ) .expect("broker-selected profile should be usable"); - assert_eq!(Some(Some("account003".to_string())), resolved); + assert_eq!(Some(Some("account003".to_string())), resolved.auth_profile); + assert_eq!( + Some("usage-lease-1"), + resolved.lease.as_ref().map(|lease| lease.lease_id.as_str()) + ); } #[test] @@ -2667,6 +2865,7 @@ mod tests { let now = at(/*seconds*/ 1_700_000_000); let decision = super::usage_profile_broker::UsageProfileBrokerDecision { selected_profile: None, + lease: None, retry_at: Some(now.timestamp() + 120), reason: super::usage_profile_broker::UsageProfileBrokerDecisionReason::NoAvailableProfiles, @@ -2675,6 +2874,7 @@ mod tests { let wait = schedule_auth_profile_after_broker_decision( Some(Some("account001".to_string())), decision, + /*lease_required*/ true, /*reset_retry_buffer_secs*/ 45, now, ) @@ -2687,11 +2887,55 @@ mod tests { wait ); assert_eq!( - "all eligible auth profiles are exhausted; retrying scheduled run after 2023-11-14T22:16:05+00:00", + "no healthy unused auth profile is available; retrying scheduled run after 2023-11-14T22:16:05+00:00", wait.to_string() ); } + #[test] + fn schedule_auth_profile_fails_closed_when_selected_profile_has_no_lease() { + let now = at(/*seconds*/ 1_700_000_000); + let decision = super::usage_profile_broker::UsageProfileBrokerDecision { + selected_profile: Some("account003".to_string()), + lease: None, + retry_at: None, + reason: + super::usage_profile_broker::UsageProfileBrokerDecisionReason::SelectedHealthyProfile, + }; + + let wait = schedule_auth_profile_after_broker_decision( + Some(Some("account001".to_string())), + decision, + /*lease_required*/ true, + /*reset_retry_buffer_secs*/ 30, + now, + ) + .expect_err("a selected profile without a durable lease must fail closed"); + + assert_eq!( + ScheduleUsageProfileWait { + retry_at: now + chrono::Duration::seconds(30), + }, + wait + ); + } + + #[test] + fn scheduled_profile_lease_heartbeat_uses_the_shorter_profile_cadence() { + assert_eq!( + Duration::from_secs(60), + schedule_lease_heartbeat_interval(Some(60)) + ); + assert_eq!( + SCHEDULE_LEASE_HEARTBEAT_INTERVAL, + schedule_lease_heartbeat_interval(Some(10 * 60)) + ); + assert_eq!( + SCHEDULE_LEASE_HEARTBEAT_INTERVAL, + schedule_lease_heartbeat_interval(None) + ); + } + #[test] fn stringified_usage_profile_wait_does_not_downcast() { let wait = ScheduleUsageProfileWait { diff --git a/codex-rs/app-server/src/request_processors/usage_profile_broker.rs b/codex-rs/app-server/src/request_processors/usage_profile_broker.rs index b3c6d257b6..7e82dd9b0d 100644 --- a/codex-rs/app-server/src/request_processors/usage_profile_broker.rs +++ b/codex-rs/app-server/src/request_processors/usage_profile_broker.rs @@ -19,10 +19,7 @@ use std::time::Instant; use tokio::time::timeout; const PROFILE_BROKER_RATE_LIMIT_FETCH_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10); -const PROFILE_BROKER_PROFILE_LEASE_DURATION: Duration = Duration::from_secs(60); - -static PROFILE_BROKER_PROFILE_LEASES: LazyLock>> = - LazyLock::new(|| StdMutex::new(BTreeMap::new())); +const PROFILE_BROKER_FALLBACK_RETRY_DELAY: Duration = Duration::from_secs(30); static PROFILE_BROKER_EXHAUSTED_PROFILE_COOLDOWNS: LazyLock< StdMutex>, > = LazyLock::new(|| StdMutex::new(BTreeMap::new())); @@ -30,6 +27,7 @@ static PROFILE_BROKER_EXHAUSTED_PROFILE_COOLDOWNS: LazyLock< #[derive(Clone, Debug, PartialEq)] pub(super) struct UsageProfileBrokerDecision { pub(super) selected_profile: Option, + pub(super) lease: Option, pub(super) retry_at: Option, pub(super) reason: UsageProfileBrokerDecisionReason, } @@ -38,6 +36,7 @@ impl UsageProfileBrokerDecision { fn no_switch(reason: UsageProfileBrokerDecisionReason) -> Self { Self { selected_profile: None, + lease: None, retry_at: None, reason, } @@ -46,21 +45,42 @@ impl UsageProfileBrokerDecision { fn selected(profile: String, reason: UsageProfileBrokerDecisionReason) -> Self { Self { selected_profile: Some(profile), + lease: None, + retry_at: None, + reason, + } + } + + fn selected_with_lease( + profile: String, + lease: codex_state::UsageProfileLease, + reason: UsageProfileBrokerDecisionReason, + ) -> Self { + Self { + selected_profile: Some(profile), + lease: Some(lease), retry_at: None, reason, } } + + fn unavailable(reason: UsageProfileBrokerDecisionReason, retry_at: Option) -> Self { + Self { + selected_profile: None, + lease: None, + retry_at, + reason, + } + } } #[derive(Clone, Debug, PartialEq)] pub(super) enum UsageProfileBrokerDecisionReason { AutoSwitchDisabled, CurrentProfileAvailable, - CurrentProfileUnknown, ProfileListUnavailable, NoCandidateProfiles, SelectedHealthyProfile, - SelectedUnknownProfile, NoAvailableProfiles, } @@ -73,6 +93,8 @@ struct FetchedProfileHealth { pub(super) async fn resolve_dispatch_auth_profile( auth_manager: &Arc, config: &Config, + state_db: &StateDbHandle, + lease_owner: &str, requested_auth_profile: Option>, ) -> UsageProfileBrokerDecision { let auto_switch = &config.auth_profile_auto_switch; @@ -82,22 +104,6 @@ pub(super) async fn resolve_dispatch_auth_profile( ); } - let current_profile = effective_current_profile(config, requested_auth_profile.as_ref()); - let current_health = fetch_profile_health(auth_manager, config, current_profile).await; - match current_health.health { - UsageProfileHealth::Healthy(_) => { - return UsageProfileBrokerDecision::no_switch( - UsageProfileBrokerDecisionReason::CurrentProfileAvailable, - ); - } - UsageProfileHealth::Unknown => { - return UsageProfileBrokerDecision::no_switch( - UsageProfileBrokerDecisionReason::CurrentProfileUnknown, - ); - } - UsageProfileHealth::Exhausted { .. } => {} - } - let saved_profiles = match codex_login::list_auth_profiles( &config.codex_home, config.cli_auth_credentials_store_mode, @@ -105,16 +111,17 @@ pub(super) async fn resolve_dispatch_auth_profile( Ok(profiles) => profiles, Err(err) => { warn!("usage profile broker could not list auth profiles: {err}"); - return UsageProfileBrokerDecision::no_switch( + return UsageProfileBrokerDecision::unavailable( UsageProfileBrokerDecisionReason::ProfileListUnavailable, + Some(profile_broker_fallback_retry_at(Utc::now())), ); } }; + let current_profile = effective_current_profile(config, requested_auth_profile.as_ref()); let now = Instant::now(); - let profile_leases = active_profile_lease_expirations(now); let exhausted_cooldowns = active_exhausted_profile_cooldown_expirations(now); - let locked_profiles = locked_profile_names(&profile_leases, &exhausted_cooldowns); + let locked_profiles = exhausted_profile_names(&exhausted_cooldowns); let candidates = auth_profile_auto_switch_candidates( current_profile, auto_switch, @@ -126,7 +133,6 @@ pub(super) async fn resolve_dispatch_auth_profile( current_profile, auto_switch, &saved_profiles, - &profile_leases, &exhausted_cooldowns, now, Utc::now().timestamp(), @@ -146,11 +152,204 @@ pub(super) async fn resolve_dispatch_auth_profile( health_by_profile.insert(profile.clone(), fetched.health); } - let decision = choose_dispatch_auth_profile(auto_switch, &candidates, &health_by_profile); - if let Some(profile) = decision.selected_profile.as_ref() { - lease_profile(profile, Instant::now()); + let mut claim_candidates = candidates; + let mut retry_at = earliest_health_retry_at(&health_by_profile); + while !claim_candidates.is_empty() { + let decision = + choose_dispatch_auth_profile(auto_switch, &claim_candidates, &health_by_profile); + let Some(profile_name) = decision.selected_profile else { + merge_retry_at(&mut retry_at, decision.retry_at); + break; + }; + let Some(profile) = saved_profiles + .iter() + .find(|profile| profile.name.as_str() == profile_name.as_str()) + else { + claim_candidates.retain(|candidate| candidate != &profile_name); + continue; + }; + let Some(identity_sha256) = usage_profile_identity_sha256(profile) else { + claim_candidates.retain(|candidate| candidate != &profile_name); + continue; + }; + match state_db + .claim_usage_profile_lease(codex_state::UsageProfileLeaseClaimParams { + identity_sha256: identity_sha256.clone(), + owner_id: lease_owner.to_string(), + profile_name: profile_name.clone(), + now: Utc::now(), + lease_duration: usage_profile_lease_duration(config), + }) + .await + { + Ok(codex_state::UsageProfileLeaseClaimOutcome::Acquired(lease)) => { + let reason = if current_profile == Some(profile_name.as_str()) { + UsageProfileBrokerDecisionReason::CurrentProfileAvailable + } else { + UsageProfileBrokerDecisionReason::SelectedHealthyProfile + }; + return UsageProfileBrokerDecision::selected_with_lease( + profile_name, + lease, + reason, + ); + } + Ok(codex_state::UsageProfileLeaseClaimOutcome::Occupied { expires_at }) => { + merge_retry_at(&mut retry_at, Some(expires_at.timestamp())); + claim_candidates.retain(|candidate| { + saved_profiles + .iter() + .find(|profile| profile.name.as_str() == candidate.as_str()) + .and_then(usage_profile_identity_sha256) + .as_deref() + != Some(identity_sha256.as_str()) + }); + } + Err(err) => { + warn!("usage profile broker could not claim durable profile lease: {err}"); + return UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + Some(profile_broker_fallback_retry_at(Utc::now())), + ); + } + } + } + UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + retry_at.or_else(|| Some(profile_broker_fallback_retry_at(Utc::now()))), + ) +} + +pub(super) async fn recheck_dispatch_auth_profile_lease( + auth_manager: &Arc, + config: &Config, + state_db: &StateDbHandle, + lease: &codex_state::UsageProfileLease, +) -> Result { + let now = Utc::now(); + let valid = state_db + .validate_usage_profile_lease(codex_state::UsageProfileLeaseValidateParams { + lease_id: lease.lease_id.clone(), + identity_sha256: lease.identity_sha256.clone(), + owner_id: lease.owner_id.clone(), + profile_name: lease.profile_name.clone(), + now, + }) + .await; + match valid { + Ok(Some(_)) => {} + Ok(None) => { + return Err(UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + Some(profile_broker_fallback_retry_at(now)), + )); + } + Err(err) => { + warn!("usage profile broker could not validate durable profile lease: {err}"); + return Err(UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + Some(profile_broker_fallback_retry_at(now)), + )); + } + } + + let saved_profiles = match codex_login::list_auth_profiles( + &config.codex_home, + config.cli_auth_credentials_store_mode, + ) { + Ok(profiles) => profiles, + Err(err) => { + warn!("usage profile broker could not list auth profiles during recheck: {err}"); + release_dispatch_auth_profile_lease(state_db, lease).await; + return Err(UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::ProfileListUnavailable, + Some(profile_broker_fallback_retry_at(now)), + )); + } + }; + let identity_matches = saved_profiles + .iter() + .find(|profile| profile.name.as_str() == lease.profile_name.as_str()) + .and_then(usage_profile_identity_sha256) + .is_some_and(|identity| identity == lease.identity_sha256); + if !identity_matches { + release_dispatch_auth_profile_lease(state_db, lease).await; + return Err(UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + Some(profile_broker_fallback_retry_at(now)), + )); + } + + let fetched = + fetch_profile_health(auth_manager, config, Some(lease.profile_name.as_str())).await; + if !matches!(&fetched.health, UsageProfileHealth::Healthy(_)) { + let retry_at = match &fetched.health { + UsageProfileHealth::Exhausted { retry_at } => *retry_at, + UsageProfileHealth::Healthy(_) | UsageProfileHealth::Unknown => None, + }; + if let Some(cooldown_key) = fetched.exhausted_cooldown { + lease_exhausted_profile( + cooldown_key, + config.usage_self_heal.reset_retry_buffer_secs, + Instant::now(), + ); + } + release_dispatch_auth_profile_lease(state_db, lease).await; + return Err(UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + retry_at.or_else(|| Some(profile_broker_fallback_retry_at(now))), + )); + } + + match renew_dispatch_auth_profile_lease(state_db, config, lease).await { + Ok(Some(lease)) => Ok(lease), + Ok(None) => Err(UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + Some(profile_broker_fallback_retry_at(now)), + )), + Err(err) => { + warn!("usage profile broker could not renew durable profile lease: {err}"); + Err(UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + Some(profile_broker_fallback_retry_at(now)), + )) + } + } +} + +pub(super) async fn renew_dispatch_auth_profile_lease( + state_db: &StateDbHandle, + config: &Config, + lease: &codex_state::UsageProfileLease, +) -> anyhow::Result> { + state_db + .renew_usage_profile_lease(codex_state::UsageProfileLeaseRenewParams { + lease_id: lease.lease_id.clone(), + owner_id: lease.owner_id.clone(), + now: Utc::now(), + lease_duration: usage_profile_lease_duration(config), + }) + .await +} + +pub(super) async fn release_dispatch_auth_profile_lease( + state_db: &StateDbHandle, + lease: &codex_state::UsageProfileLease, +) { + if let Err(err) = state_db + .release_usage_profile_lease(codex_state::UsageProfileLeaseReleaseParams { + lease_id: lease.lease_id.clone(), + owner_id: lease.owner_id.clone(), + now: Utc::now(), + }) + .await + { + warn!("usage profile broker could not release durable profile lease: {err}"); } - decision +} + +pub(super) fn usage_profile_lease_owner(scope: &str, owner_lease_id: &str) -> String { + format!("{scope}:{owner_lease_id}") } fn effective_current_profile<'a>( @@ -221,7 +420,23 @@ fn choose_dispatch_auth_profile( candidates: &[String], health_by_profile: &BTreeMap, ) -> UsageProfileBrokerDecision { - let selection = choose_profile_for_auto_switch(config, candidates, health_by_profile); + let healthy_candidates = candidates + .iter() + .filter(|profile| { + matches!( + health_by_profile.get(profile.as_str()), + Some(UsageProfileHealth::Healthy(_)) + ) + }) + .cloned() + .collect::>(); + if healthy_candidates.is_empty() { + return UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + earliest_health_retry_at(health_by_profile), + ); + } + let selection = choose_profile_for_auto_switch(config, &healthy_candidates, health_by_profile); if let Some(profile) = selection.selected_profile { return UsageProfileBrokerDecision::selected( profile, @@ -230,7 +445,7 @@ fn choose_dispatch_auth_profile( UsageProfileBrokerDecisionReason::SelectedHealthyProfile } codex_core::usage_profile_health::UsageProfileSelectionReason::SelectedUnknownProfile => { - UsageProfileBrokerDecisionReason::SelectedUnknownProfile + UsageProfileBrokerDecisionReason::NoAvailableProfiles } codex_core::usage_profile_health::UsageProfileSelectionReason::NoCandidateProfiles | codex_core::usage_profile_health::UsageProfileSelectionReason::NoAvailableProfiles => { @@ -249,11 +464,10 @@ fn choose_dispatch_auth_profile( codex_core::usage_profile_health::UsageProfileSelectionReason::NoAvailableProfiles | codex_core::usage_profile_health::UsageProfileSelectionReason::SelectedHealthyProfile | codex_core::usage_profile_health::UsageProfileSelectionReason::SelectedUnknownProfile => { - UsageProfileBrokerDecision { - selected_profile: None, - retry_at: selection.retry_at, - reason: UsageProfileBrokerDecisionReason::NoAvailableProfiles, - } + UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + selection.retry_at, + ) } } } @@ -271,14 +485,13 @@ fn auth_profile_auto_switch_candidates( let start = current .and_then(|current| ordered.iter().position(|profile| profile == current)) - .map(|index| index + 1) + .map(|index| index) .unwrap_or(0); ordered .iter() .cycle() .skip(start) .take(ordered.len()) - .filter(|profile| current != Some(profile.as_str())) .filter(|profile| !locked_profiles.contains(profile.as_str())) .cloned() .collect() @@ -293,6 +506,10 @@ fn ordered_auth_profiles_for_auto_switch( .filter(|profile| { profile.subscription_provider == AuthProfileSubscriptionProvider::ChatGpt && profile.auth_mode.is_some() + && profile + .account_id + .as_deref() + .is_some_and(|account_id| !account_id.trim().is_empty()) }) .collect::>(); let saved_names = saved_profiles @@ -314,6 +531,49 @@ fn ordered_auth_profiles_for_auto_switch( dedupe_profile_names(ordered) } +fn usage_profile_identity_sha256(profile: &AuthProfile) -> Option { + profile + .account_id + .as_deref() + .filter(|account_id| !account_id.trim().is_empty()) + .map(|account_id| { + codex_state::StateRuntime::usage_profile_identity_sha256(account_id.as_bytes()) + }) +} + +fn usage_profile_lease_duration(config: &Config) -> Duration { + let heartbeat_interval = config + .auth_profile_auto_switch + .heartbeat_interval_secs + .max(1); + let heartbeat_freshness = config + .auth_profile_auto_switch + .heartbeat_freshness_secs + .max(heartbeat_interval.saturating_mul(2)); + Duration::from_secs(heartbeat_freshness) +} + +fn profile_broker_fallback_retry_at(now: DateTime) -> i64 { + let fallback_seconds = + i64::try_from(PROFILE_BROKER_FALLBACK_RETRY_DELAY.as_secs()).unwrap_or(i64::MAX); + now.timestamp().saturating_add(fallback_seconds) +} + +fn earliest_health_retry_at( + health_by_profile: &BTreeMap, +) -> Option { + let mut retry_at = None; + for health in health_by_profile.values() { + if let UsageProfileHealth::Exhausted { + retry_at: exhausted_retry_at, + } = health + { + merge_retry_at(&mut retry_at, *exhausted_retry_at); + } + } + retry_at +} + impl FetchedProfileHealth { fn unknown() -> Self { Self { @@ -368,13 +628,10 @@ fn dedupe_profile_names(profiles: Vec) -> Vec { /// Decide what to do when the lock-filtered candidate list is empty. /// -/// An empty candidate list has two very different meanings: -/// - the user has no sibling profiles configured at all, in which case the -/// dispatch should proceed on the current profile as before, or -/// - sibling profiles exist but every one of them is lease- or -/// cooldown-locked (e.g. during an all-profiles-exhausted window), in which -/// case proceeding would burn the dispatch (and its failure budget) on a -/// known-exhausted profile. +/// An empty candidate list means either no provider-backed profile is eligible +/// or every eligible profile is locally cooldown-locked. Both cases fail +/// closed: dispatch must never bypass a fresh provider-health check and durable +/// identity lease by silently falling back to the current profile. /// /// For the lock-induced case, report `NoAvailableProfiles` with the earliest /// unlock time so callers defer through their existing usage-wait paths @@ -384,7 +641,6 @@ fn empty_candidate_decision( current_profile: Option<&str>, auto_switch: &AuthProfileAutoSwitchConfig, saved_profiles: &[AuthProfile], - profile_leases: &BTreeMap, exhausted_cooldowns: &BTreeMap, now: Instant, now_epoch: i64, @@ -396,21 +652,25 @@ fn empty_candidate_decision( &HashSet::new(), ); if unlocked_candidates.is_empty() { - return UsageProfileBrokerDecision::no_switch( + return UsageProfileBrokerDecision::unavailable( UsageProfileBrokerDecisionReason::NoCandidateProfiles, + Some(now_epoch.saturating_add( + i64::try_from(PROFILE_BROKER_FALLBACK_RETRY_DELAY.as_secs()).unwrap_or(i64::MAX), + )), ); } - UsageProfileBrokerDecision { - selected_profile: None, - retry_at: locked_candidates_retry_at_epoch( - &unlocked_candidates, - profile_leases, - exhausted_cooldowns, - now, - now_epoch, - ), - reason: UsageProfileBrokerDecisionReason::NoAvailableProfiles, - } + UsageProfileBrokerDecision::unavailable( + UsageProfileBrokerDecisionReason::NoAvailableProfiles, + locked_candidates_retry_at_epoch(&unlocked_candidates, exhausted_cooldowns, now, now_epoch) + .or_else(|| { + Some( + now_epoch.saturating_add( + i64::try_from(PROFILE_BROKER_FALLBACK_RETRY_DELAY.as_secs()) + .unwrap_or(i64::MAX), + ), + ) + }), + ) } /// Earliest epoch at which one of the lock-filtered candidate profiles @@ -418,19 +678,12 @@ fn empty_candidate_decision( /// known) or the remaining lease/cooldown duration. fn locked_candidates_retry_at_epoch( candidates: &[String], - profile_leases: &BTreeMap, exhausted_cooldowns: &BTreeMap, now: Instant, now_epoch: i64, ) -> Option { let mut retry_at = None; for profile in candidates { - if let Some(expires_at) = profile_leases.get(profile) { - merge_retry_at( - &mut retry_at, - Some(instant_expiry_epoch(*expires_at, now, now_epoch)), - ); - } for (key, expires_at) in exhausted_cooldowns { if key.profile.as_deref() != Some(profile.as_str()) { continue; @@ -450,45 +703,15 @@ fn instant_expiry_epoch(expires_at: Instant, now: Instant, now_epoch: i64) -> i6 now_epoch.saturating_add(i64::try_from(remaining.as_secs()).unwrap_or(i64::MAX)) } -fn locked_profile_names( - profile_leases: &BTreeMap, +fn exhausted_profile_names( exhausted_cooldowns: &BTreeMap, ) -> HashSet { - let mut locked_profiles = profile_leases.keys().cloned().collect::>(); - locked_profiles.extend( - exhausted_cooldowns - .keys() - .filter_map(|key| key.profile.clone()), - ); - locked_profiles -} - -fn active_profile_lease_expirations(now: Instant) -> BTreeMap { - let Ok(mut leases) = PROFILE_BROKER_PROFILE_LEASES.lock() else { - return BTreeMap::new(); - }; - leases.retain(|_, expires_at| *expires_at > now); - leases.clone() -} - -#[cfg(test)] -fn active_profile_leases(now: Instant) -> HashSet { - active_profile_lease_expirations(now) + exhausted_cooldowns .keys() - .cloned() + .filter_map(|key| key.profile.clone()) .collect() } -fn lease_profile(profile: &str, now: Instant) { - let Ok(mut leases) = PROFILE_BROKER_PROFILE_LEASES.lock() else { - return; - }; - leases.insert( - profile.to_string(), - now + PROFILE_BROKER_PROFILE_LEASE_DURATION, - ); -} - fn active_exhausted_profile_cooldown_expirations( now: Instant, ) -> BTreeMap { @@ -511,7 +734,7 @@ fn lease_exhausted_profile( key.resets_at, Utc::now().timestamp(), reset_retry_buffer_secs, - PROFILE_BROKER_PROFILE_LEASE_DURATION, + PROFILE_BROKER_FALLBACK_RETRY_DELAY, ); cooldowns.insert(key, now + cooldown); } @@ -525,12 +748,16 @@ mod tests { use codex_protocol::protocol::RateLimitWindow; fn chatgpt_profile(name: &str) -> AuthProfile { + chatgpt_profile_with_account(name, Some(name)) + } + + fn chatgpt_profile_with_account(name: &str, account_id: Option<&str>) -> AuthProfile { AuthProfile { name: name.to_string(), subscription_provider: AuthProfileSubscriptionProvider::ChatGpt, auth_mode: Some(AuthMode::Chatgpt), - email: Some(format!("{name}@example.com")), - account_id: Some(format!("acct-{name}")), + email: None, + account_id: account_id.map(|account_id| format!("synthetic-account-{account_id}")), plan: Some("plus".to_string()), active: false, } @@ -560,7 +787,7 @@ mod tests { } #[test] - fn dispatch_candidates_rotate_after_current_profile() { + fn dispatch_candidates_recheck_current_before_rotating() { let profiles = vec![ chatgpt_profile("work"), chatgpt_profile("second"), @@ -574,7 +801,11 @@ mod tests { &profiles, &HashSet::new(), ), - vec!["second".to_string(), "third".to_string()] + vec![ + "work".to_string(), + "second".to_string(), + "third".to_string(), + ] ); } @@ -594,112 +825,48 @@ mod tests { &profiles, &locked_profiles, ), - vec!["third".to_string()] + vec!["work".to_string(), "third".to_string()] ); } #[test] - fn dispatch_profile_leases_exhaust_candidates_until_lease_expiry() { - PROFILE_BROKER_PROFILE_LEASES - .lock() - .expect("profile leases lock") - .clear(); + fn aliases_share_one_opaque_provider_identity() { + let first = chatgpt_profile_with_account("profile-a", Some("shared")); + let second = chatgpt_profile_with_account("profile-alias-a", Some("shared")); - let mut config = config(); - config.strategy = AuthProfileAutoSwitchStrategy::Ordered; - let profiles = vec![ - chatgpt_profile("work"), - chatgpt_profile("second"), - chatgpt_profile("third"), - ]; - let health_by_profile = BTreeMap::from([ - ("second".to_string(), health(/*remaining_percent*/ 20.0)), - ("third".to_string(), health(/*remaining_percent*/ 80.0)), - ]); - let now = Instant::now(); - - let first_candidates = auth_profile_auto_switch_candidates( - Some("work"), - &config, - &profiles, - &active_profile_leases(now), - ); - assert_eq!( - vec!["second".to_string(), "third".to_string()], - first_candidates - ); - let first = choose_dispatch_auth_profile(&config, &first_candidates, &health_by_profile); assert_eq!( - UsageProfileBrokerDecision::selected( - "second".to_string(), - UsageProfileBrokerDecisionReason::SelectedHealthyProfile, - ), - first + usage_profile_identity_sha256(&first), + usage_profile_identity_sha256(&second) ); - lease_profile( - first - .selected_profile - .as_deref() - .expect("first candidate should be selected"), - now, + assert_ne!( + Some(first.name.as_str()), + usage_profile_identity_sha256(&first).as_deref() ); + } - let second_candidates = auth_profile_auto_switch_candidates( - Some("work"), - &config, - &profiles, - &active_profile_leases(now), - ); - assert_eq!(vec!["third".to_string()], second_candidates); - let second = choose_dispatch_auth_profile(&config, &second_candidates, &health_by_profile); + #[test] + fn lease_owner_is_bound_to_dispatch_scope_and_owner_lease() { assert_eq!( - UsageProfileBrokerDecision::selected( - "third".to_string(), - UsageProfileBrokerDecisionReason::SelectedHealthyProfile, - ), - second + "mailbox:dispatch-lease-a", + usage_profile_lease_owner("mailbox", "dispatch-lease-a") ); - lease_profile( - second - .selected_profile - .as_deref() - .expect("second candidate should be selected"), - now, - ); - - let exhausted_candidates = auth_profile_auto_switch_candidates( - Some("work"), - &config, - &profiles, - &active_profile_leases(now), - ); - assert_eq!(Vec::::new(), exhausted_candidates); assert_eq!( - UsageProfileBrokerDecision::no_switch( - UsageProfileBrokerDecisionReason::NoCandidateProfiles - ), - choose_dispatch_auth_profile(&config, &exhausted_candidates, &health_by_profile) + "schedule:schedule-lease-a", + usage_profile_lease_owner("schedule", "schedule-lease-a") ); + } + + #[test] + fn profiles_without_provider_identity_are_not_dispatch_candidates() { + let profiles = vec![ + chatgpt_profile_with_account("profile-a", None), + chatgpt_profile_with_account("profile-b", Some("profile-b")), + ]; - let after_expiry = now + PROFILE_BROKER_PROFILE_LEASE_DURATION + Duration::from_millis(1); - assert_eq!( - HashSet::::new(), - active_profile_leases(after_expiry) - ); assert_eq!( - vec!["second".to_string(), "third".to_string()], - auth_profile_auto_switch_candidates( - Some("work"), - &config, - &profiles, - &active_profile_leases(after_expiry), - ) + vec!["profile-b".to_string()], + ordered_auth_profiles_for_auto_switch(&[], &profiles) ); - - PROFILE_BROKER_PROFILE_LEASES - .lock() - .expect("profile leases lock") - .clear(); } #[test] @@ -711,8 +878,16 @@ mod tests { ]; let now = Instant::now(); let now_epoch = 1_000; - let profile_leases = BTreeMap::new(); let exhausted_cooldowns = BTreeMap::from([ + ( + UsageProfileCooldownKey { + profile: Some("work".to_string()), + limit_id: "codex".to_string(), + window_label: "5h".to_string(), + resets_at: Some(5_000), + }, + now + Duration::from_secs(4_060), + ), ( UsageProfileCooldownKey { profile: Some("second".to_string()), @@ -733,7 +908,7 @@ mod tests { ), ]); - let locked_profiles = locked_profile_names(&profile_leases, &exhausted_cooldowns); + let locked_profiles = exhausted_profile_names(&exhausted_cooldowns); let candidates = auth_profile_auto_switch_candidates( Some("work"), &config(), @@ -745,6 +920,7 @@ mod tests { assert_eq!( UsageProfileBrokerDecision { selected_profile: None, + lease: None, retry_at: Some(3_000), reason: UsageProfileBrokerDecisionReason::NoAvailableProfiles, }, @@ -752,7 +928,6 @@ mod tests { Some("work"), &config(), &profiles, - &profile_leases, &exhausted_cooldowns, now, now_epoch, @@ -761,53 +936,20 @@ mod tests { } #[test] - fn lease_locked_candidates_defer_until_lease_expiry() { - let profiles = vec![ - chatgpt_profile("work"), - chatgpt_profile("second"), - chatgpt_profile("third"), - ]; - let now = Instant::now(); - let now_epoch = 1_000; - let profile_leases = BTreeMap::from([ - ("second".to_string(), now + Duration::from_secs(60)), - ("third".to_string(), now + Duration::from_secs(45)), - ]); - let exhausted_cooldowns = BTreeMap::new(); + fn empty_profile_inventory_fails_closed_with_retry() { + let profiles = Vec::new(); assert_eq!( UsageProfileBrokerDecision { selected_profile: None, - retry_at: Some(1_045), - reason: UsageProfileBrokerDecisionReason::NoAvailableProfiles, + lease: None, + retry_at: Some(1_030), + reason: UsageProfileBrokerDecisionReason::NoCandidateProfiles, }, empty_candidate_decision( - Some("work"), + None, &config(), &profiles, - &profile_leases, - &exhausted_cooldowns, - now, - now_epoch, - ) - ); - } - - #[test] - fn single_profile_users_keep_proceeding_without_candidates() { - let mut config = config(); - config.profiles = vec!["work".to_string()]; - let profiles = vec![chatgpt_profile("work")]; - - assert_eq!( - UsageProfileBrokerDecision::no_switch( - UsageProfileBrokerDecisionReason::NoCandidateProfiles - ), - empty_candidate_decision( - Some("work"), - &config, - &profiles, - &BTreeMap::new(), &BTreeMap::new(), Instant::now(), /*now_epoch*/ 1_000, @@ -840,6 +982,16 @@ mod tests { // Dispatch #1 observes both siblings exhausted and records their // cooldowns, exactly as resolve_dispatch_auth_profile does. + lease_exhausted_profile( + UsageProfileCooldownKey { + profile: Some("cool-a".to_string()), + limit_id: "codex".to_string(), + window_label: "5h".to_string(), + resets_at: Some(now_epoch + 7_200), + }, + /*reset_retry_buffer_secs*/ 300, + now, + ); lease_exhausted_profile( UsageProfileCooldownKey { profile: Some("cool-b".to_string()), @@ -864,9 +1016,8 @@ mod tests { // Dispatch #2 during the cooldown window: candidates are emptied by // the cooldown locks, but the decision must still defer with a // retry time instead of proceeding on the exhausted profile. - let profile_leases = BTreeMap::new(); let exhausted_cooldowns = active_exhausted_profile_cooldown_expirations(now); - let locked_profiles = locked_profile_names(&profile_leases, &exhausted_cooldowns); + let locked_profiles = exhausted_profile_names(&exhausted_cooldowns); let candidates = auth_profile_auto_switch_candidates( Some("cool-a"), &config, @@ -879,7 +1030,6 @@ mod tests { Some("cool-a"), &config, &profiles, - &profile_leases, &exhausted_cooldowns, now, now_epoch, @@ -887,6 +1037,7 @@ mod tests { assert_eq!( UsageProfileBrokerDecision { selected_profile: None, + lease: None, retry_at: Some(now_epoch + 3_600), reason: UsageProfileBrokerDecisionReason::NoAvailableProfiles, }, @@ -920,7 +1071,7 @@ mod tests { } #[test] - fn ordered_dispatch_skips_exhausted_profile_for_unknown_candidate() { + fn ordered_dispatch_fails_closed_for_unknown_candidate() { let mut config = config(); config.strategy = AuthProfileAutoSwitchStrategy::Ordered; let health_by_profile = BTreeMap::from([( @@ -936,10 +1087,12 @@ mod tests { &["second".to_string(), "third".to_string()], &health_by_profile, ), - UsageProfileBrokerDecision::selected( - "third".to_string(), - UsageProfileBrokerDecisionReason::SelectedUnknownProfile, - ) + UsageProfileBrokerDecision { + selected_profile: None, + lease: None, + retry_at: Some(500), + reason: UsageProfileBrokerDecisionReason::NoAvailableProfiles, + } ); } @@ -968,6 +1121,7 @@ mod tests { ), UsageProfileBrokerDecision { selected_profile: None, + lease: None, retry_at: Some(500), reason: UsageProfileBrokerDecisionReason::NoAvailableProfiles, } diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 45351ea9fa..2a187d5e64 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -26,6 +26,7 @@ use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::watch; +use tokio_util::sync::CancellationToken; use tracing::error; type PendingInterruptQueue = Vec; @@ -54,10 +55,18 @@ pub(crate) struct ScheduledThreadScheduleRun { pub(crate) schedule_id: String, pub(crate) run_id: String, pub(crate) lease_id: String, + pub(crate) usage_profile_lease: Option, pub(crate) goal_id: Option, pub(crate) state_db: StateDbHandle, } +#[derive(Clone)] +pub(crate) struct UsageProfileTurnLease { + pub(crate) lease: codex_state::UsageProfileLease, + pub(crate) state_db: StateDbHandle, + pub(crate) heartbeat_cancel: CancellationToken, +} + // ThreadListenerCommand is used to perform operations in the context of the thread listener, for serialization purposes. pub(crate) enum ThreadListenerCommand { // SendThreadResumeResponse is used to resume an already running thread by sending the thread's history to the client and atomically subscribing for new updates. @@ -118,6 +127,7 @@ pub(crate) struct ThreadState { listener_command_tx: Option>, current_turn_history: ThreadHistoryBuilder, scheduled_runs_by_turn_id: HashMap, + usage_profile_leases_by_turn_id: HashMap, listener_thread: Option>, watch_registration: WatchRegistration, } @@ -193,6 +203,28 @@ impl ThreadState { self.scheduled_runs_by_turn_id.contains_key(turn_id) } + pub(crate) fn track_usage_profile_turn_lease( + &mut self, + turn_id: String, + lease: UsageProfileTurnLease, + ) -> Option { + self.usage_profile_leases_by_turn_id.insert(turn_id, lease) + } + + pub(crate) fn take_usage_profile_turn_lease( + &mut self, + turn_id: &str, + ) -> Option { + self.usage_profile_leases_by_turn_id.remove(turn_id) + } + + pub(crate) fn drain_usage_profile_turn_leases(&mut self) -> Vec { + self.usage_profile_leases_by_turn_id + .drain() + .map(|(_, lease)| lease) + .collect() + } + pub(crate) fn track_current_turn_event(&mut self, event_turn_id: &str, event: &EventMsg) { if let EventMsg::TurnStarted(payload) = event { self.turn_summary.started_at = payload.started_at; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index e12006ff9b..f52f548464 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -35,6 +35,7 @@ use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::user_input::UserInput; @@ -483,6 +484,25 @@ impl CodexThread { self.codex.session.maybe_start_turn_for_pending_work().await } + /// Starts a regular pending-work turn with a caller-owned turn id. + /// + /// Durable dispatchers use the stable id to bind external leases to the + /// exact turn that consumes the queued work. + pub async fn maybe_start_turn_for_pending_work_with_sub_id(&self, sub_id: String) -> bool { + self.codex + .session + .maybe_start_turn_for_pending_work_with_sub_id(sub_id) + .await + } + + /// Aborts the exact active turn when a durable execution lease is lost. + pub async fn abort_turn_if_active(&self, turn_id: &str, reason: TurnAbortReason) -> bool { + self.codex + .session + .abort_turn_if_active(turn_id, reason) + .await + } + pub async fn set_app_server_client_info( &self, app_server_client_name: Option, diff --git a/codex-rs/state/migrations/0068_usage_profile_leases.sql b/codex-rs/state/migrations/0068_usage_profile_leases.sql new file mode 100644 index 0000000000..9911c8fa0a --- /dev/null +++ b/codex-rs/state/migrations/0068_usage_profile_leases.sql @@ -0,0 +1,31 @@ +CREATE TABLE usage_profile_leases ( + lease_id TEXT PRIMARY KEY CHECK(LENGTH(TRIM(lease_id)) > 0), + identity_sha256 TEXT NOT NULL CHECK( + LENGTH(identity_sha256) = 64 + AND identity_sha256 = LOWER(identity_sha256) + AND identity_sha256 NOT GLOB '*[^0-9a-f]*' + ), + owner_id TEXT NOT NULL CHECK(LENGTH(TRIM(owner_id)) > 0), + profile_name TEXT NOT NULL CHECK(LENGTH(TRIM(profile_name)) > 0), + acquired_at_ms INTEGER NOT NULL, + heartbeat_at_ms INTEGER NOT NULL, + expires_at_ms INTEGER NOT NULL, + released_at_ms INTEGER, + release_reason TEXT CHECK( + release_reason IS NULL OR release_reason IN ('released', 'expired') + ), + CHECK(expires_at_ms > acquired_at_ms), + CHECK( + (released_at_ms IS NULL AND release_reason IS NULL) + OR + (released_at_ms IS NOT NULL AND release_reason IS NOT NULL) + ) +); + +CREATE UNIQUE INDEX idx_usage_profile_leases_active_identity + ON usage_profile_leases(identity_sha256) + WHERE released_at_ms IS NULL; + +CREATE INDEX idx_usage_profile_leases_active_expiry + ON usage_profile_leases(expires_at_ms) + WHERE released_at_ms IS NULL; diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 6c407144c0..15abe345ff 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -253,6 +253,12 @@ pub use runtime::ThreadScheduleRunForGoalFinishParams; pub use runtime::ThreadScheduleRunLeaseParams; pub use runtime::ThreadScheduleRunStartParams; pub use runtime::ThreadScheduleUpdate; +pub use runtime::UsageProfileLease; +pub use runtime::UsageProfileLeaseClaimOutcome; +pub use runtime::UsageProfileLeaseClaimParams; +pub use runtime::UsageProfileLeaseReleaseParams; +pub use runtime::UsageProfileLeaseRenewParams; +pub use runtime::UsageProfileLeaseValidateParams; pub use runtime::WEBHOOK_EVENT_DEDUPE_CONFLICT_MESSAGE; pub use runtime::WORKFLOW_STEP_APPROVAL_APPROVED; pub use runtime::WORKFLOW_STEP_APPROVAL_PENDING; diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index caea0489a3..5a70139fc7 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -119,6 +119,7 @@ mod schedules; #[cfg(test)] mod test_support; mod threads; +mod usage_profile_leases; mod webhooks; mod workflow_automation; mod workflow_effects; @@ -257,6 +258,12 @@ pub use schedules::ThreadScheduleRunLeaseParams; pub use schedules::ThreadScheduleRunStartParams; pub use schedules::ThreadScheduleUpdate; pub use threads::ThreadFilterOptions; +pub use usage_profile_leases::UsageProfileLease; +pub use usage_profile_leases::UsageProfileLeaseClaimOutcome; +pub use usage_profile_leases::UsageProfileLeaseClaimParams; +pub use usage_profile_leases::UsageProfileLeaseReleaseParams; +pub use usage_profile_leases::UsageProfileLeaseRenewParams; +pub use usage_profile_leases::UsageProfileLeaseValidateParams; pub use webhooks::DEFAULT_WEBHOOK_EVENT_LIST_LIMIT; pub use webhooks::MAX_WEBHOOK_EVENT_LIST_LIMIT; pub use webhooks::WEBHOOK_EVENT_DEDUPE_CONFLICT_MESSAGE; diff --git a/codex-rs/state/src/runtime/usage_profile_leases.rs b/codex-rs/state/src/runtime/usage_profile_leases.rs new file mode 100644 index 0000000000..74edb82139 --- /dev/null +++ b/codex-rs/state/src/runtime/usage_profile_leases.rs @@ -0,0 +1,546 @@ +use super::*; +use sha2::Digest; +use sha2::Sha256; +use sqlx::Row; +use uuid::Uuid; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageProfileLease { + pub lease_id: String, + pub identity_sha256: String, + pub owner_id: String, + pub profile_name: String, + pub acquired_at: DateTime, + pub heartbeat_at: DateTime, + pub expires_at: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageProfileLeaseClaimParams { + pub identity_sha256: String, + pub owner_id: String, + pub profile_name: String, + pub now: DateTime, + pub lease_duration: std::time::Duration, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum UsageProfileLeaseClaimOutcome { + Acquired(UsageProfileLease), + Occupied { expires_at: DateTime }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageProfileLeaseRenewParams { + pub lease_id: String, + pub owner_id: String, + pub now: DateTime, + pub lease_duration: std::time::Duration, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageProfileLeaseValidateParams { + pub lease_id: String, + pub identity_sha256: String, + pub owner_id: String, + pub profile_name: String, + pub now: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageProfileLeaseReleaseParams { + pub lease_id: String, + pub owner_id: String, + pub now: DateTime, +} + +impl StateRuntime { + pub fn usage_profile_identity_sha256(identity: &[u8]) -> String { + format!("{:x}", Sha256::digest(identity)) + } + + pub async fn claim_usage_profile_lease( + &self, + params: UsageProfileLeaseClaimParams, + ) -> anyhow::Result { + validate_usage_profile_lease_id("identity_sha256", params.identity_sha256.as_str())?; + validate_usage_profile_lease_text("owner_id", params.owner_id.as_str())?; + validate_usage_profile_lease_text("profile_name", params.profile_name.as_str())?; + let now_ms = datetime_to_epoch_millis(params.now); + let expires_at = usage_profile_lease_expiry(params.now, params.lease_duration)?; + let expires_at_ms = datetime_to_epoch_millis(expires_at); + let lease_id = Uuid::new_v4().to_string(); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + + sqlx::query( + r#" +UPDATE usage_profile_leases +SET + released_at_ms = ?, + release_reason = 'expired' +WHERE released_at_ms IS NULL + AND expires_at_ms <= ? + "#, + ) + .bind(now_ms) + .bind(now_ms) + .execute(&mut *tx) + .await?; + + let inserted = sqlx::query( + r#" +INSERT OR IGNORE INTO usage_profile_leases ( + lease_id, + identity_sha256, + owner_id, + profile_name, + acquired_at_ms, + heartbeat_at_ms, + expires_at_ms +) VALUES (?, ?, ?, ?, ?, ?, ?) + "#, + ) + .bind(lease_id.as_str()) + .bind(params.identity_sha256.as_str()) + .bind(params.owner_id.as_str()) + .bind(params.profile_name.as_str()) + .bind(now_ms) + .bind(now_ms) + .bind(expires_at_ms) + .execute(&mut *tx) + .await? + .rows_affected() + == 1; + + let outcome = if inserted { + UsageProfileLeaseClaimOutcome::Acquired( + get_usage_profile_lease_by_id_in_tx(&mut tx, lease_id.as_str()) + .await? + .ok_or_else(|| anyhow::anyhow!("usage profile lease disappeared"))?, + ) + } else { + let occupied_expires_at = sqlx::query_scalar::<_, i64>( + r#" +SELECT expires_at_ms +FROM usage_profile_leases +WHERE identity_sha256 = ? + AND released_at_ms IS NULL + AND expires_at_ms > ? + "#, + ) + .bind(params.identity_sha256.as_str()) + .bind(now_ms) + .fetch_optional(&mut *tx) + .await? + .ok_or_else(|| { + anyhow::anyhow!("usage profile identity contention had no active lease") + })?; + UsageProfileLeaseClaimOutcome::Occupied { + expires_at: epoch_millis_to_datetime(occupied_expires_at)?, + } + }; + tx.commit().await?; + Ok(outcome) + } + + pub async fn renew_usage_profile_lease( + &self, + params: UsageProfileLeaseRenewParams, + ) -> anyhow::Result> { + validate_usage_profile_lease_text("lease_id", params.lease_id.as_str())?; + validate_usage_profile_lease_text("owner_id", params.owner_id.as_str())?; + let now_ms = datetime_to_epoch_millis(params.now); + let expires_at = usage_profile_lease_expiry(params.now, params.lease_duration)?; + let expires_at_ms = datetime_to_epoch_millis(expires_at); + let row = sqlx::query( + r#" +UPDATE usage_profile_leases +SET + heartbeat_at_ms = ?, + expires_at_ms = ? +WHERE lease_id = ? + AND owner_id = ? + AND released_at_ms IS NULL + AND expires_at_ms > ? +RETURNING + lease_id, + identity_sha256, + owner_id, + profile_name, + acquired_at_ms, + heartbeat_at_ms, + expires_at_ms + "#, + ) + .bind(now_ms) + .bind(expires_at_ms) + .bind(params.lease_id.as_str()) + .bind(params.owner_id.as_str()) + .bind(now_ms) + .fetch_optional(&self.pool) + .await?; + row.map(usage_profile_lease_from_row).transpose() + } + + pub async fn validate_usage_profile_lease( + &self, + params: UsageProfileLeaseValidateParams, + ) -> anyhow::Result> { + validate_usage_profile_lease_text("lease_id", params.lease_id.as_str())?; + validate_usage_profile_lease_id("identity_sha256", params.identity_sha256.as_str())?; + validate_usage_profile_lease_text("owner_id", params.owner_id.as_str())?; + validate_usage_profile_lease_text("profile_name", params.profile_name.as_str())?; + let row = sqlx::query( + r#" +SELECT + lease_id, + identity_sha256, + owner_id, + profile_name, + acquired_at_ms, + heartbeat_at_ms, + expires_at_ms +FROM usage_profile_leases +WHERE lease_id = ? + AND identity_sha256 = ? + AND owner_id = ? + AND profile_name = ? + AND released_at_ms IS NULL + AND expires_at_ms > ? + "#, + ) + .bind(params.lease_id.as_str()) + .bind(params.identity_sha256.as_str()) + .bind(params.owner_id.as_str()) + .bind(params.profile_name.as_str()) + .bind(datetime_to_epoch_millis(params.now)) + .fetch_optional(&*self.reader_pool) + .await?; + row.map(usage_profile_lease_from_row).transpose() + } + + pub async fn release_usage_profile_lease( + &self, + params: UsageProfileLeaseReleaseParams, + ) -> anyhow::Result { + validate_usage_profile_lease_text("lease_id", params.lease_id.as_str())?; + validate_usage_profile_lease_text("owner_id", params.owner_id.as_str())?; + Ok(sqlx::query( + r#" +UPDATE usage_profile_leases +SET + released_at_ms = ?, + release_reason = 'released' +WHERE lease_id = ? + AND owner_id = ? + AND released_at_ms IS NULL + "#, + ) + .bind(datetime_to_epoch_millis(params.now)) + .bind(params.lease_id.as_str()) + .bind(params.owner_id.as_str()) + .execute(&self.pool) + .await? + .rows_affected() + == 1) + } +} + +async fn get_usage_profile_lease_by_id_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + lease_id: &str, +) -> anyhow::Result> { + let row = sqlx::query( + r#" +SELECT + lease_id, + identity_sha256, + owner_id, + profile_name, + acquired_at_ms, + heartbeat_at_ms, + expires_at_ms +FROM usage_profile_leases +WHERE lease_id = ? + "#, + ) + .bind(lease_id) + .fetch_optional(&mut **tx) + .await?; + row.map(usage_profile_lease_from_row).transpose() +} + +fn usage_profile_lease_from_row(row: sqlx::sqlite::SqliteRow) -> anyhow::Result { + Ok(UsageProfileLease { + lease_id: row.try_get("lease_id")?, + identity_sha256: row.try_get("identity_sha256")?, + owner_id: row.try_get("owner_id")?, + profile_name: row.try_get("profile_name")?, + acquired_at: epoch_millis_to_datetime(row.try_get("acquired_at_ms")?)?, + heartbeat_at: epoch_millis_to_datetime(row.try_get("heartbeat_at_ms")?)?, + expires_at: epoch_millis_to_datetime(row.try_get("expires_at_ms")?)?, + }) +} + +fn usage_profile_lease_expiry( + now: DateTime, + lease_duration: std::time::Duration, +) -> anyhow::Result> { + if lease_duration.is_zero() { + anyhow::bail!("usage profile lease duration must be positive"); + } + now.checked_add_signed(chrono::Duration::from_std(lease_duration)?) + .ok_or_else(|| anyhow::anyhow!("usage profile lease expiry overflowed")) +} + +fn validate_usage_profile_lease_text(label: &str, value: &str) -> anyhow::Result<()> { + if value.trim().is_empty() { + anyhow::bail!("{label} must not be empty"); + } + Ok(()) +} + +fn validate_usage_profile_lease_id(label: &str, value: &str) -> anyhow::Result<()> { + if value.len() != 64 + || value != value.to_ascii_lowercase() + || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) + { + anyhow::bail!("{label} must be a lowercase SHA-256 digest"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::runtime::test_support::unique_temp_dir; + use chrono::TimeDelta; + use pretty_assertions::assert_eq; + + const LEASE_DURATION: std::time::Duration = std::time::Duration::from_secs(60); + + async fn test_runtime(codex_home: std::path::PathBuf) -> std::sync::Arc { + StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should initialize") + } + + fn claim_params( + owner_id: &str, + profile_name: &str, + identity_sha256: &str, + now: DateTime, + ) -> UsageProfileLeaseClaimParams { + UsageProfileLeaseClaimParams { + owner_id: owner_id.to_string(), + profile_name: profile_name.to_string(), + identity_sha256: identity_sha256.to_string(), + now, + lease_duration: LEASE_DURATION, + } + } + + fn acquired(outcome: UsageProfileLeaseClaimOutcome) -> UsageProfileLease { + match outcome { + UsageProfileLeaseClaimOutcome::Acquired(lease) => lease, + UsageProfileLeaseClaimOutcome::Occupied { .. } => { + panic!("expected usage profile lease acquisition") + } + } + } + + #[tokio::test] + async fn concurrent_alias_claims_allow_only_one_identity_owner() { + let codex_home = unique_temp_dir(); + let first_runtime = test_runtime(codex_home.clone()).await; + let second_runtime = test_runtime(codex_home).await; + let identity_sha256 = + StateRuntime::usage_profile_identity_sha256(b"synthetic-provider-account-a"); + let now = DateTime::::from_timestamp_millis(1_800_000_000_000) + .expect("test timestamp should be valid"); + + let (first, second) = tokio::join!( + first_runtime.claim_usage_profile_lease(claim_params( + "lane-a", + "profile-a", + identity_sha256.as_str(), + now, + )), + second_runtime.claim_usage_profile_lease(claim_params( + "lane-b", + "profile-alias-a", + identity_sha256.as_str(), + now, + )), + ); + let outcomes = [ + first.expect("first claim should complete"), + second.expect("second claim should complete"), + ]; + + assert_eq!( + 1, + outcomes + .iter() + .filter(|outcome| matches!(outcome, UsageProfileLeaseClaimOutcome::Acquired(_))) + .count() + ); + assert_eq!( + 1, + outcomes + .iter() + .filter(|outcome| matches!(outcome, UsageProfileLeaseClaimOutcome::Occupied { .. })) + .count() + ); + } + + #[tokio::test] + async fn lease_heartbeat_and_release_are_owner_bound() { + let runtime = test_runtime(unique_temp_dir()).await; + let identity_sha256 = + StateRuntime::usage_profile_identity_sha256(b"synthetic-provider-account-b"); + let now = DateTime::::from_timestamp_millis(1_800_100_000_000) + .expect("test timestamp should be valid"); + let lease = acquired( + runtime + .claim_usage_profile_lease(claim_params( + "lane-a", + "profile-b", + identity_sha256.as_str(), + now, + )) + .await + .expect("lease should be acquired"), + ); + + assert_eq!( + None, + runtime + .renew_usage_profile_lease(UsageProfileLeaseRenewParams { + lease_id: lease.lease_id.clone(), + owner_id: "lane-b".to_string(), + now: now + TimeDelta::seconds(10), + lease_duration: LEASE_DURATION, + }) + .await + .expect("foreign heartbeat should complete"), + ); + assert!( + !runtime + .release_usage_profile_lease(UsageProfileLeaseReleaseParams { + lease_id: lease.lease_id.clone(), + owner_id: "lane-b".to_string(), + now: now + TimeDelta::seconds(10), + }) + .await + .expect("foreign release should complete") + ); + + let renewed = runtime + .renew_usage_profile_lease(UsageProfileLeaseRenewParams { + lease_id: lease.lease_id.clone(), + owner_id: "lane-a".to_string(), + now: now + TimeDelta::seconds(10), + lease_duration: LEASE_DURATION, + }) + .await + .expect("owner heartbeat should complete") + .expect("owner heartbeat should retain lease"); + assert_eq!(now + TimeDelta::seconds(70), renewed.expires_at); + assert_eq!( + Some(renewed.clone()), + runtime + .validate_usage_profile_lease(UsageProfileLeaseValidateParams { + lease_id: renewed.lease_id.clone(), + identity_sha256: renewed.identity_sha256.clone(), + owner_id: renewed.owner_id.clone(), + profile_name: renewed.profile_name.clone(), + now: now + TimeDelta::seconds(10), + }) + .await + .expect("execution-start validation should complete"), + ); + + assert!( + runtime + .release_usage_profile_lease(UsageProfileLeaseReleaseParams { + lease_id: lease.lease_id, + owner_id: "lane-a".to_string(), + now: now + TimeDelta::seconds(11), + }) + .await + .expect("owner release should complete") + ); + assert!(matches!( + runtime + .claim_usage_profile_lease(claim_params( + "lane-b", + "profile-alias-b", + identity_sha256.as_str(), + now + TimeDelta::seconds(12), + )) + .await + .expect("replacement claim should complete"), + UsageProfileLeaseClaimOutcome::Acquired(_) + )); + } + + #[tokio::test] + async fn stale_crashed_lease_recovers_without_reviving_old_owner() { + let codex_home = unique_temp_dir(); + let crashed_runtime = test_runtime(codex_home.clone()).await; + let recovery_runtime = test_runtime(codex_home).await; + let identity_sha256 = + StateRuntime::usage_profile_identity_sha256(b"synthetic-provider-account-c"); + let now = DateTime::::from_timestamp_millis(1_800_200_000_000) + .expect("test timestamp should be valid"); + let crashed_lease = acquired( + crashed_runtime + .claim_usage_profile_lease(claim_params( + "crashed-lane", + "profile-c", + identity_sha256.as_str(), + now, + )) + .await + .expect("crashed owner should acquire lease"), + ); + + assert!(matches!( + recovery_runtime + .claim_usage_profile_lease(claim_params( + "recovery-lane", + "profile-alias-c", + identity_sha256.as_str(), + now + TimeDelta::seconds(59), + )) + .await + .expect("pre-expiry claim should complete"), + UsageProfileLeaseClaimOutcome::Occupied { .. } + )); + assert!(matches!( + recovery_runtime + .claim_usage_profile_lease(claim_params( + "recovery-lane", + "profile-alias-c", + identity_sha256.as_str(), + now + TimeDelta::seconds(61), + )) + .await + .expect("stale recovery claim should complete"), + UsageProfileLeaseClaimOutcome::Acquired(_) + )); + assert_eq!( + None, + crashed_runtime + .renew_usage_profile_lease(UsageProfileLeaseRenewParams { + lease_id: crashed_lease.lease_id, + owner_id: "crashed-lane".to_string(), + now: now + TimeDelta::seconds(61), + lease_duration: LEASE_DURATION, + }) + .await + .expect("stale heartbeat should complete"), + ); + } +} From d9dda9ec78838c72f8d18ac751916c0356befdfe Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 22:57:40 +0300 Subject: [PATCH 2/7] Fix SQLx lease pool executor Agent: quintilianus --- codex-rs/state/src/runtime/usage_profile_leases.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codex-rs/state/src/runtime/usage_profile_leases.rs b/codex-rs/state/src/runtime/usage_profile_leases.rs index 74edb82139..ab5e77ca2b 100644 --- a/codex-rs/state/src/runtime/usage_profile_leases.rs +++ b/codex-rs/state/src/runtime/usage_profile_leases.rs @@ -177,7 +177,7 @@ RETURNING .bind(params.lease_id.as_str()) .bind(params.owner_id.as_str()) .bind(now_ms) - .fetch_optional(&self.pool) + .fetch_optional(&*self.pool) .await?; row.map(usage_profile_lease_from_row).transpose() } @@ -239,7 +239,7 @@ WHERE lease_id = ? .bind(datetime_to_epoch_millis(params.now)) .bind(params.lease_id.as_str()) .bind(params.owner_id.as_str()) - .execute(&self.pool) + .execute(&*self.pool) .await? .rows_affected() == 1) From 8d0132414feb15317d577fe0fc475c79f359acdc Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 23:13:29 +0300 Subject: [PATCH 3/7] Fix schedule profile selection ownership Agent: quintilianus --- .../src/request_processors/thread_schedule_runtime.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 17afba651b..3077a2d527 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -1437,6 +1437,7 @@ fn schedule_deferral_for_idle_rejection( }) } +#[cfg_attr(test, derive(Debug))] struct ScheduleAuthProfileSelection { auth_profile: Option>, lease: Option, @@ -1470,10 +1471,12 @@ fn schedule_auth_profile_after_broker_decision( reset_retry_buffer_secs: u64, now: DateTime, ) -> Result { - if let (Some(profile), Some(lease)) = (decision.selected_profile, decision.lease) { + if let (Some(profile), Some(lease)) = + (decision.selected_profile.as_ref(), decision.lease.as_ref()) + { return Ok(ScheduleAuthProfileSelection { - auth_profile: Some(Some(profile)), - lease: Some(lease), + auth_profile: Some(Some(profile.clone())), + lease: Some(lease.clone()), }); } if !lease_required { From 3d8ad7878024fad4c618e48344a0ee60c68c4381 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 23:37:08 +0300 Subject: [PATCH 4/7] Add required test argument comments Agent: quintilianus --- .../src/request_processors/thread_schedule_runtime.rs | 2 +- .../app-server/src/request_processors/usage_profile_broker.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 3077a2d527..1d6e72bd29 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -2935,7 +2935,7 @@ mod tests { ); assert_eq!( SCHEDULE_LEASE_HEARTBEAT_INTERVAL, - schedule_lease_heartbeat_interval(None) + schedule_lease_heartbeat_interval(/*usage_profile_heartbeat_interval_secs*/ None,) ); } diff --git a/codex-rs/app-server/src/request_processors/usage_profile_broker.rs b/codex-rs/app-server/src/request_processors/usage_profile_broker.rs index 7e82dd9b0d..e3105b660e 100644 --- a/codex-rs/app-server/src/request_processors/usage_profile_broker.rs +++ b/codex-rs/app-server/src/request_processors/usage_profile_broker.rs @@ -859,7 +859,7 @@ mod tests { #[test] fn profiles_without_provider_identity_are_not_dispatch_candidates() { let profiles = vec![ - chatgpt_profile_with_account("profile-a", None), + chatgpt_profile_with_account("profile-a", /*account_id*/ None), chatgpt_profile_with_account("profile-b", Some("profile-b")), ]; @@ -947,7 +947,7 @@ mod tests { reason: UsageProfileBrokerDecisionReason::NoCandidateProfiles, }, empty_candidate_decision( - None, + /*current_profile*/ None, &config(), &profiles, &BTreeMap::new(), From 4d7fd30bb5f6ea14262cbf412952d6535191c80c Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Mon, 10 Aug 2026 00:16:17 +0300 Subject: [PATCH 5/7] Satisfy provider lease Clippy checks Agent: quintilianus --- .../thread_mailbox_dispatcher_runtime.rs | 8 ++++---- .../src/request_processors/usage_profile_broker.rs | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs b/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs index 46bb61b98f..d33e21aee5 100644 --- a/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs @@ -613,16 +613,16 @@ impl ThreadMailboxDispatcherRuntime { match thread_manager.get_thread(thread_id).await { Ok(thread) => { if usage_profile_lease.is_some() - && last_provider_recheck.map_or(true, |checked_at: Instant| { + && last_provider_recheck.is_none_or(|checked_at: Instant| { checked_at.elapsed() >= Duration::from_secs(1) }) { let Some(state_db) = state_db.as_ref() else { break; }; - let lease = usage_profile_lease - .take() - .expect("usage profile lease checked as present"); + let Some(lease) = usage_profile_lease.take() else { + return; + }; usage_profile_lease = match super::usage_profile_broker::recheck_dispatch_auth_profile_lease( &auth_manager, diff --git a/codex-rs/app-server/src/request_processors/usage_profile_broker.rs b/codex-rs/app-server/src/request_processors/usage_profile_broker.rs index e3105b660e..c72a251a2f 100644 --- a/codex-rs/app-server/src/request_processors/usage_profile_broker.rs +++ b/codex-rs/app-server/src/request_processors/usage_profile_broker.rs @@ -485,7 +485,6 @@ fn auth_profile_auto_switch_candidates( let start = current .and_then(|current| ordered.iter().position(|profile| profile == current)) - .map(|index| index) .unwrap_or(0); ordered .iter() From 86e6858dd1777b3c08342eae3bc303b54adc818b Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Mon, 10 Aug 2026 00:30:25 +0300 Subject: [PATCH 6/7] fix: prevent unleased mailbox restart Agent: Trebius --- .../thread_mailbox_dispatcher_runtime.rs | 4 +- .../thread_schedule_runtime.rs | 2 +- codex-rs/core/src/codex_thread.rs | 15 ++++++++ codex-rs/core/src/session/tests.rs | 38 +++++++++++++++++++ codex-rs/core/src/tasks/mod.rs | 37 +++++++++++++++++- 5 files changed, 92 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs b/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs index d33e21aee5..37e2d001ff 100644 --- a/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_mailbox_dispatcher_runtime.rs @@ -939,7 +939,7 @@ async fn heartbeat_mailbox_usage_profile_turn( "mailbox turn lost durable usage profile lease ownership" ); let _ = thread - .abort_turn_if_active( + .abort_turn_if_active_without_pending_work_restart( turn_id.as_str(), codex_protocol::protocol::TurnAbortReason::Interrupted, ) @@ -952,7 +952,7 @@ async fn heartbeat_mailbox_usage_profile_turn( "failed to heartbeat mailbox turn usage profile lease: {err}" ); let _ = thread - .abort_turn_if_active( + .abort_turn_if_active_without_pending_work_restart( turn_id.as_str(), codex_protocol::protocol::TurnAbortReason::Interrupted, ) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 1d6e72bd29..ac193724ac 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -748,7 +748,7 @@ impl ThreadScheduleRuntime { } if usage_profile_lease.is_some() { let _ = thread - .abort_turn_if_active( + .abort_turn_if_active_without_pending_work_restart( turn_id.as_str(), codex_protocol::protocol::TurnAbortReason::Interrupted, ) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index f52f548464..94b44687e7 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -503,6 +503,21 @@ impl CodexThread { .await } + /// Aborts the exact active turn without starting queued mailbox work. + /// + /// Durable dispatchers use this when execution authority is lost so a new + /// turn cannot start before replacement authority is acquired. + pub async fn abort_turn_if_active_without_pending_work_restart( + &self, + turn_id: &str, + reason: TurnAbortReason, + ) -> bool { + self.codex + .session + .abort_turn_if_active_without_pending_work_restart(turn_id, reason) + .await + } + pub async fn set_app_server_client_info( &self, app_server_client_name: Option, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ff713f5a60..84f63bca52 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -10569,6 +10569,44 @@ async fn task_finish_starts_pending_trigger_turn_mailbox_work() { session.abort_all_tasks(TurnAbortReason::Interrupted).await; } +#[tokio::test] +async fn lease_loss_abort_keeps_pending_mailbox_work_queued() { + let (session, turn_context, _rx) = make_session_and_context_with_rx().await; + let session = Arc::new(session); + session + .spawn_task( + Arc::clone(&turn_context), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }, + ) + .await; + session + .input_queue + .enqueue_mailbox_communication(InterAgentCommunication::new( + AgentPath::root(), + AgentPath::root(), + Vec::new(), + "pending trigger".to_string(), + /*trigger_turn*/ true, + )) + .await + .expect("mailbox queue has room"); + + assert!( + session + .abort_turn_if_active_without_pending_work_restart( + turn_context.sub_id.as_str(), + TurnAbortReason::Interrupted, + ) + .await + ); + assert!(session.active_turn.lock().await.is_none()); + assert!(session.has_wake_worthy_pending_mailbox().await); +} + #[tokio::test] async fn pending_trigger_turn_mailbox_work_takes_over_empty_active_turn() { let (session, _turn_context, rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 2ecd5b8d45..a26cedcc05 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -64,6 +64,12 @@ pub(crate) use user_shell::execute_user_shell_command; const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100; const TASK_COMPACT_METRIC: &str = "codex.task.compact"; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PendingWorkAfterAbort { + Resume, + LeaveQueued, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum InterruptedTurnHistoryMarker { Disabled, @@ -636,6 +642,33 @@ impl Session { self: &Arc, turn_id: &str, reason: TurnAbortReason, + ) -> bool { + self.abort_turn_if_active_with_pending_work( + turn_id, + reason, + PendingWorkAfterAbort::Resume, + ) + .await + } + + pub(crate) async fn abort_turn_if_active_without_pending_work_restart( + self: &Arc, + turn_id: &str, + reason: TurnAbortReason, + ) -> bool { + self.abort_turn_if_active_with_pending_work( + turn_id, + reason, + PendingWorkAfterAbort::LeaveQueued, + ) + .await + } + + async fn abort_turn_if_active_with_pending_work( + self: &Arc, + turn_id: &str, + reason: TurnAbortReason, + pending_work: PendingWorkAfterAbort, ) -> bool { let active_turn = { let mut active = self.active_turn.lock().await; @@ -666,7 +699,9 @@ impl Session { // in-flight approval wait can surface as a model-visible rejection before TurnAborted. self.input_queue.clear_pending(&active_turn).await; - if reason == TurnAbortReason::Interrupted { + if reason == TurnAbortReason::Interrupted + && pending_work == PendingWorkAfterAbort::Resume + { self.maybe_start_turn_for_pending_work().await; } From bae61b1418b9069145b82303c9f0c1d5929266f9 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Mon, 10 Aug 2026 00:33:22 +0300 Subject: [PATCH 7/7] style: apply rustfmt to lease abort path Agent: Trebius --- codex-rs/core/src/tasks/mod.rs | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index a26cedcc05..0d0dd84324 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -643,12 +643,8 @@ impl Session { turn_id: &str, reason: TurnAbortReason, ) -> bool { - self.abort_turn_if_active_with_pending_work( - turn_id, - reason, - PendingWorkAfterAbort::Resume, - ) - .await + self.abort_turn_if_active_with_pending_work(turn_id, reason, PendingWorkAfterAbort::Resume) + .await } pub(crate) async fn abort_turn_if_active_without_pending_work_restart( @@ -699,9 +695,7 @@ impl Session { // in-flight approval wait can surface as a model-visible rejection before TurnAborted. self.input_queue.clear_pending(&active_turn).await; - if reason == TurnAbortReason::Interrupted - && pending_work == PendingWorkAfterAbort::Resume - { + if reason == TurnAbortReason::Interrupted && pending_work == PendingWorkAfterAbort::Resume { self.maybe_start_turn_for_pending_work().await; }