From 26f89719027f35bee271361e40110551f7f72d92 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 18:16:56 +0300 Subject: [PATCH 01/15] feat(background-agent): enforce durable admission --- .../background_agent_live.rs | 62 ++- .../background_agent_processor.rs | 276 ++++++----- .../tests/suite/v2/background_agent.rs | 48 +- codex-rs/background-agent/ARCHITECTURE.md | 30 ++ codex-rs/background-agent/src/daemon.rs | 107 +++- codex-rs/background-agent/src/lib.rs | 57 +++ codex-rs/background-agent/src/supervisor.rs | 64 ++- codex-rs/cli/src/agent_cmd.rs | 116 +++-- ...61_background_agent_lifecycle_receipts.sql | 6 + .../src/runtime/background_agents/events.rs | 135 ++++++ .../src/runtime/background_agents/runs.rs | 456 +++++++++++++++++- .../src/runtime/background_agents/tests.rs | 257 +++++++++- codex-rs/tui/src/app_server_session.rs | 8 +- 13 files changed, 1366 insertions(+), 256 deletions(-) create mode 100644 codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql diff --git a/codex-rs/app-server/src/request_processors/background_agent_live.rs b/codex-rs/app-server/src/request_processors/background_agent_live.rs index d3118d8d7d..daeecc95df 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_live.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_live.rs @@ -6,6 +6,7 @@ use super::worktree_paths::path_to_api_string; use super::worktree_paths::paths_equivalent; use crate::error_code::internal_error; use crate::error_code::invalid_params; +use crate::error_code::invalid_request; use anyhow::Context; use chrono::DateTime; use chrono::Duration as ChronoDuration; @@ -53,6 +54,9 @@ use codex_app_server_protocol::WorktreeSessionMode; use codex_background_agent::AgentEventJournal; use codex_background_agent::AgentRunStore; use codex_background_agent::AgentSnapshotStore; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; use codex_background_agent::BackgroundAgentDesiredState; use codex_background_agent::BackgroundAgentEvent; use codex_background_agent::BackgroundAgentExecutionHandleParams; @@ -329,6 +333,7 @@ impl ThreadRequestProcessor { &self, mut params: AgentStartParams, ) -> Result, JSONRPCErrorError> { + self.validate_agent_start_admission(¶ms)?; let managed_worktree = self .trusted_agent_start_managed_worktree( params.cwd.as_deref(), @@ -1666,6 +1671,37 @@ impl ThreadRequestProcessor { BackgroundAgentRequestProcessor::new(self.state_db.clone()) } + fn validate_agent_start_admission( + &self, + params: &AgentStartParams, + ) -> Result<(), JSONRPCErrorError> { + if params.version_fingerprint.as_deref() + != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) + { + let mut error = invalid_request("background agent admission schema is incompatible"); + error.data = Some(json!({ + "errorCode": BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH, + "requestedSchema": params.version_fingerprint.as_deref(), + "supportedSchema": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + })); + return Err(error); + } + if let Some(requested_profile) = params.auth_profile_ref.as_deref() + && self.config.selected_auth_profile.as_deref() != Some(requested_profile) + { + let mut error = invalid_request( + "background agent auth profile does not match the app-server profile", + ); + error.data = Some(json!({ + "errorCode": BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH, + "requestedProfile": requested_profile, + "selectedProfile": self.config.selected_auth_profile.as_deref(), + })); + return Err(error); + } + Ok(()) + } + fn freeze_start_execution_context(&self, params: &mut AgentStartParams) { let context = params.execution_context.get_or_insert_with(|| { Box::new(AgentExecutionContextParams { @@ -3131,6 +3167,11 @@ async fn resolve_background_agent_config( context: &BackgroundAgentWorkerContext, run: &BackgroundAgentRun, ) -> anyhow::Result { + if run.version_fingerprint.as_deref() != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: persisted background agent admission schema is incompatible" + ); + } let snapshot = context .state_db .get_latest_execution_snapshot(run.id.as_str()) @@ -3141,6 +3182,14 @@ async fn resolve_background_agent_config( let payload = snapshot .as_ref() .and_then(|snapshot| snapshot.payload_json.as_object()); + let snapshot_auth_profile_ref = payload + .and_then(|payload| payload.get("authProfileRef")) + .and_then(Value::as_str); + if snapshot_auth_profile_ref != run.auth_profile_ref.as_deref() { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH}: persisted execution snapshot auth profile does not match the admitted run" + ); + } let cwd = payload .and_then(|payload| payload.get("cwd")) .and_then(Value::as_str) @@ -3250,6 +3299,17 @@ async fn resolve_background_agent_config( .load_with_overrides(request_overrides.clone(), config_overrides.clone()) .await .map_err(anyhow::Error::from)?; + if config.selected_auth_profile != run.auth_profile_ref { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH}: loaded worker auth profile does not match the admitted run" + ); + } + if run.auth_profile_ref.is_some() { + return Ok(BackgroundAgentConfigResolution::Ready { + config: Box::new(config), + initial_execution_payload, + }); + } let broker_decision = super::usage_profile_broker::resolve_dispatch_auth_profile( &context.auth_manager, @@ -5912,7 +5972,7 @@ done auth_profile_ref: None, status_reason: Some("queued by process supervisor test".to_string()), config_fingerprint: Some("cfg-test".to_string()), - version_fingerprint: Some("version-test".to_string()), + version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), }) .await?; state_db diff --git a/codex-rs/app-server/src/request_processors/background_agent_processor.rs b/codex-rs/app-server/src/request_processors/background_agent_processor.rs index 85cacebb10..50499ab82f 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_processor.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_processor.rs @@ -60,7 +60,11 @@ use codex_app_server_protocol::WorktreeReadResponse; use codex_background_agent::AgentEventJournal; use codex_background_agent::AgentRunStore; use codex_background_agent::AgentSnapshotStore; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED; use codex_background_agent::BACKGROUND_AGENT_EVENT_CURSOR_COMPACTED; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; use codex_background_agent::BackgroundAgentDesiredState; use codex_background_agent::BackgroundAgentEvent; use codex_background_agent::BackgroundAgentExecutionSnapshot; @@ -76,6 +80,7 @@ use codex_background_agent::BackgroundAgentStatusSnapshotParams; use codex_background_agent::LifecycleAction; use codex_background_agent::LifecycleEffect; use codex_background_agent::PendingInteractionLedger; +use codex_background_agent::DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS; use codex_background_agent::lifecycle_effect_for; use codex_protocol::ThreadId; use codex_protocol::approvals::ElicitationAction; @@ -90,14 +95,14 @@ use codex_state::ManagedWorktreeDetachParams; use serde_json::Value; use serde_json::json; use std::path::Path; +use std::time::Duration; use uuid::Uuid; const DEFAULT_AGENT_LIST_LIMIT: usize = 50; const MAX_AGENT_LIST_LIMIT: usize = 200; -const DEFAULT_MAX_ACTIVE_AGENT_RUNS_PER_USER: i64 = 8; +const AGENT_ADMISSION_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(30); const AGENT_BACKPRESSURE_ACTIVE_RUN_LIMIT: &str = "active_run_limit"; const AGENT_EVENT_CURSOR_PREFIX: &str = "event:"; -static AGENT_START_ADMISSION_LOCK: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(1); #[derive(Clone)] pub(crate) struct BackgroundAgentRequestProcessor { @@ -114,6 +119,7 @@ impl BackgroundAgentRequestProcessor { params: AgentStartParams, ) -> Result { let state_db = self.state_db()?; + validate_agent_start_schema(¶ms)?; let AgentStartParams { prompt, initial_goal_objective, @@ -138,47 +144,15 @@ impl BackgroundAgentRequestProcessor { let execution_context = execution_context.map(|context| *context); let prompt = validate_agent_prompt(prompt)?; let initial_goal_objective = validate_agent_initial_goal_objective(initial_goal_objective)?; - let mut existing_run = match idempotency_key.as_deref() { - Some(idempotency_key) => state_db - .get_run_by_idempotency_key(idempotency_key) - .await - .map_err(|err| { - internal_error(format!( - "failed to load background agent idempotency key: {err}" - )) - })?, - None => None, - }; - let _admission_permit = if existing_run.is_none() { - let permit = AGENT_START_ADMISSION_LOCK.acquire().await.map_err(|err| { - internal_error(format!( - "failed to acquire background agent admission permit: {err}" - )) - })?; - if let Some(idempotency_key) = idempotency_key.as_deref() { - existing_run = state_db - .get_run_by_idempotency_key(idempotency_key) - .await - .map_err(|err| { - internal_error(format!( - "failed to load background agent idempotency key: {err}" - )) - })?; - } - Some(permit) - } else { - None - }; - let new_run_requested = existing_run.is_none(); - if new_run_requested { - let quota = load_agent_quota_snapshot(state_db.as_ref()).await?; - if !quota.admission_allowed() { - return Err(overloaded(format!( - "background agent queue is overloaded: {} active run(s), max {}", - quota.active_run_count, quota.max_active_runs_per_user - ))); - } - } + retry_transient_sqlite_busy("reconcile stale background agents before admission", || { + state_db.orphan_stale_background_agent_runs(AGENT_ADMISSION_HEARTBEAT_TIMEOUT) + }) + .await + .map_err(|err| { + internal_error(format!( + "failed to reconcile stale background agents before admission: {err}" + )) + })?; let agent_id = Uuid::now_v7().to_string(); let prompt_snapshot_ref = prompt_snapshot_ref.unwrap_or_else(|| format!("inline:{agent_id}:prompt")); @@ -194,34 +168,36 @@ impl BackgroundAgentRequestProcessor { .as_ref() .and_then(|context| context.recovery_policy.clone()) .unwrap_or_else(|| "abort_mid_turn_resume_at_safe_boundary".to_string()); - let run = match existing_run { - Some(run) => run, - None => state_db - .create_run(BackgroundAgentRunCreateParams { - id: agent_id.clone(), - idempotency_key, - request_id, - source, - prompt_snapshot_ref, - input_snapshot_ref, - thread_id, - thread_store_kind, - thread_store_id, - rollout_path, - parent_thread_id, - parent_agent_run_id, - spawn_linkage_json: spawn_linkage, - auth_profile_ref, - status_reason: Some("queued for background-agent supervisor".to_string()), - config_fingerprint, - version_fingerprint, - }) - .await - .map_err(|err| { - internal_error(format!("failed to create background agent: {err}")) - })?, + let create_params = BackgroundAgentRunCreateParams { + id: agent_id.clone(), + idempotency_key, + request_id, + source, + prompt_snapshot_ref, + input_snapshot_ref, + thread_id, + thread_store_kind, + thread_store_id, + rollout_path, + parent_thread_id, + parent_agent_run_id, + spawn_linkage_json: spawn_linkage, + auth_profile_ref, + status_reason: Some("queued for background-agent supervisor".to_string()), + config_fingerprint, + version_fingerprint, }; - let created_new_run = run.id == agent_id; + let (run, created_new_run) = retry_transient_sqlite_busy( + "admit background agent", + || { + state_db.admit_run( + create_params.clone(), + DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, + ) + }, + ) + .await + .map_err(map_background_agent_admission_error)?; let execution_payload = initial_execution_snapshot_payload( &run, InitialExecutionSnapshotPayloadParams { @@ -289,13 +265,16 @@ impl BackgroundAgentRequestProcessor { internal_error(format!("failed to append background agent event: {err}")) })? } else { - let mut events = state_db - .list_events_after(run.id.as_str(), /*after_seq*/ None, Some(1)) + let events = state_db + .list_events_after(run.id.as_str(), /*after_seq*/ None, Some(100)) .await .map_err(|err| { internal_error(format!("failed to list background agent events: {err}")) })?; - match events.pop() { + match events + .into_iter() + .find(|event| event.event_type == "agent.started") + { Some(event) => event, None => append_background_agent_event_with_retry( state_db.as_ref(), @@ -572,69 +551,74 @@ impl BackgroundAgentRequestProcessor { agent: Some(api_agent_run_from_state(run)), }); } - state_db - .set_desired_state(run.id.as_str(), BackgroundAgentDesiredState::Stopped) - .await - .map_err(|err| { - internal_error(format!( - "failed to update background agent desired state: {err}" - )) - })?; if !is_terminal_agent_status(run.status) { - let terminalize_immediately = should_terminalize_unclaimed_agent_run(&run); - let status = if terminalize_immediately { - BackgroundAgentRunStatus::Cancelled - } else { - BackgroundAgentRunStatus::Stopping - }; - let status_reason = if terminalize_immediately { - "stop requested before worker claim" - } else { - "stop requested" - }; - state_db - .update_run_status(run.id.as_str(), status, Some(status_reason)) - .await - .map_err(|err| { - internal_error(format!("failed to update background agent status: {err}")) - })?; - append_background_agent_event_with_retry( + let mut observed = run.clone(); + let mut stopped = false; + for _ in 0..2 { + let terminalize_immediately = should_terminalize_unclaimed_agent_run(&observed); + let status_reason = if terminalize_immediately { + "stop requested before worker claim" + } else { + "stop requested" + }; + stopped = state_db + .request_background_agent_stop_for_generation( + observed.id.as_str(), + observed.supervisor_id.as_deref(), + observed.generation, + status_reason, + &json!({ + "reason": "client_requested_stop", + }), + ) + .await + .map_err(|err| { + internal_error(format!( + "failed to request fenced background agent stop: {err}" + )) + })?; + if stopped { + break; + } + let Some(latest) = self + .load_agent_run(state_db.as_ref(), observed.id.as_str()) + .await? + else { + break; + }; + if is_terminal_agent_status(latest.status) { + stopped = true; + break; + } + observed = latest; + } + if !stopped { + return Err(internal_error( + "background agent ownership changed during stop request", + )); + } + cancel_active_pending_interactions_for_run( state_db.as_ref(), run.id.as_str(), - "agent.stopRequested", - &json!({ - "reason": "client_requested_stop", - }), + "client_requested_stop", ) - .await - .map_err(|err| { - internal_error(format!("failed to append background agent event: {err}")) - })?; - cancel_active_pending_interactions_for_run( + .await?; + let latest = self + .load_agent_run(state_db.as_ref(), run.id.as_str()) + .await? + .ok_or_else(|| internal_error("background agent disappeared during stop"))?; + upsert_lifecycle_status_snapshot( state_db.as_ref(), run.id.as_str(), + latest.status, + if latest.status == BackgroundAgentRunStatus::Cancelled { + "Stopped" + } else { + "Stopping" + }, "client_requested_stop", ) .await?; - if terminalize_immediately { - upsert_lifecycle_status_snapshot( - state_db.as_ref(), - run.id.as_str(), - status, - "Stopped", - "client_requested_stop", - ) - .await?; - } else { - upsert_lifecycle_status_snapshot( - state_db.as_ref(), - run.id.as_str(), - status, - "Stopping", - "client_requested_stop", - ) - .await?; - } } let run = self .load_agent_run(state_db.as_ref(), run.id.as_str()) @@ -839,7 +823,8 @@ impl BackgroundAgentRequestProcessor { &self, ) -> Result { let Some(state_db) = self.state_db.clone() else { - let quota = AgentQuotaSnapshot::empty(DEFAULT_MAX_ACTIVE_AGENT_RUNS_PER_USER); + let quota = + AgentQuotaSnapshot::empty(DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS); return Ok(AgentDaemonDiagnosticsResponse { state_store_available: false, active_run_count: quota.active_run_count, @@ -1188,7 +1173,7 @@ async fn load_agent_quota_snapshot( .map_err(|err| internal_error(format!("failed to count background agents: {err}")))?; Ok(AgentQuotaSnapshot::from_status_counts( runs_by_status, - DEFAULT_MAX_ACTIVE_AGENT_RUNS_PER_USER, + DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, )) } @@ -1277,6 +1262,39 @@ fn validate_agent_prompt(prompt: String) -> Result { Ok(prompt) } +fn validate_agent_start_schema(params: &AgentStartParams) -> Result<(), JSONRPCErrorError> { + if params.version_fingerprint.as_deref() == Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) { + return Ok(()); + } + let mut error = invalid_request("background agent admission schema is incompatible"); + error.data = Some(json!({ + "errorCode": BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH, + "requestedSchema": params.version_fingerprint.as_deref(), + "supportedSchema": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + })); + Err(error) +} + +fn map_background_agent_admission_error(err: anyhow::Error) -> JSONRPCErrorError { + let message = err.to_string(); + if message.contains(BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED) { + let mut error = overloaded("background agent admission capacity is exhausted"); + error.data = Some(json!({ + "errorCode": BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED, + "maxActiveRuns": DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, + })); + return error; + } + if message.contains(BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH) { + let mut error = invalid_request("background agent idempotency identity does not match"); + error.data = Some(json!({ + "errorCode": BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH, + })); + return error; + } + internal_error(format!("failed to admit background agent: {message}")) +} + fn validate_agent_initial_goal_objective( objective: Option, ) -> Result, JSONRPCErrorError> { diff --git a/codex-rs/app-server/tests/suite/v2/background_agent.rs b/codex-rs/app-server/tests/suite/v2/background_agent.rs index 5d449a5c05..7588ebc420 100644 --- a/codex-rs/app-server/tests/suite/v2/background_agent.rs +++ b/codex-rs/app-server/tests/suite/v2/background_agent.rs @@ -48,6 +48,9 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; use codex_app_server_protocol::WorktreeAttachResponse; use codex_app_server_protocol::WorktreeCleanupPolicy; use codex_app_server_protocol::WorktreeCleanupResponse; @@ -153,7 +156,7 @@ async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result start.execution_snapshot.payload.get("approvalPolicy"), Some(&json!("never")) ); - assert_eq!(start.event.seq, 1); + assert_eq!(start.event.seq, 2); assert_eq!(start.event.event_type, "agent.started"); assert_eq!( start.event.payload.get("prompt"), @@ -199,7 +202,7 @@ async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result let first_events_page = agent_events_page(&mut restarted, &agent_id, /*cursor*/ None, Some(1)).await?; assert_eq!(first_events_page.data.len(), 1); - assert_eq!(first_events_page.data[0].event_type, "agent.started"); + assert_eq!(first_events_page.data[0].event_type, "agent.admitted"); assert_eq!(first_events_page.next_cursor, Some("event:1".to_string())); let second_events_page = agent_events_page( @@ -212,7 +215,7 @@ async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result assert_eq!(second_events_page.data.len(), 1); assert_eq!( second_events_page.data[0].event_type, - "agent.workerStarting" + "agent.started" ); assert_eq!(second_events_page.next_cursor, Some("event:2".to_string())); let all_events = @@ -390,7 +393,6 @@ async fn agent_start_freezes_authority_from_server_config() -> Result<()> { codex_home.path(), ); params.cwd = Some("/tmp/client-selected-cwd".to_string()); - params.auth_profile_ref = Some("client-selected-auth-profile".to_string()); let context = params .execution_context .as_mut() @@ -457,6 +459,40 @@ async fn agent_start_freezes_authority_from_server_config() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn agent_start_rejects_profile_and_schema_mismatches_before_admission() -> Result<()> { + let codex_home = TempDir::new()?; + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + write_config(codex_home.path(), server.uri().as_str())?; + let mut mcp = init_mcp(codex_home.path()).await?; + + let mut profile_mismatch = start_params( + "reject mismatched profile", + Some("profile-mismatch".to_string()), + codex_home.path(), + ); + profile_mismatch.auth_profile_ref = Some("client-selected-auth-profile".to_string()); + let error = start_agent_error(&mut mcp, profile_mismatch).await?; + assert_eq!( + error.error.data.as_ref().and_then(|data| data.get("errorCode")), + Some(&json!(BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH)) + ); + + let mut schema_mismatch = start_params( + "reject mismatched schema", + Some("schema-mismatch".to_string()), + codex_home.path(), + ); + schema_mismatch.version_fingerprint = Some("older-schema".to_string()); + let error = start_agent_error(&mut mcp, schema_mismatch).await?; + assert_eq!( + error.error.data.as_ref().and_then(|data| data.get("errorCode")), + Some(&json!(BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH)) + ); + assert!(agent_list(&mut mcp).await?.data.is_empty()); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn agent_start_uses_validated_managed_worktree_cwd() -> Result<()> { let codex_home = TempDir::new()?; @@ -3331,7 +3367,7 @@ async fn seed_queued_agent_run( auth_profile_ref: None, status_reason: Some("queued by quota test".to_string()), config_fingerprint: Some("cfg-test".to_string()), - version_fingerprint: Some("version-test".to_string()), + version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), }) .await?; state_db @@ -3371,7 +3407,7 @@ fn start_params( spawn_linkage: None, auth_profile_ref: None, config_fingerprint: Some("cfg-test".to_string()), - version_fingerprint: Some("version-test".to_string()), + version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), execution_context: Some(Box::new(AgentExecutionContextParams { workspace_roots: Some(vec![codex_home.display().to_string()]), approval_policy: Some(AskForApproval::Never), diff --git a/codex-rs/background-agent/ARCHITECTURE.md b/codex-rs/background-agent/ARCHITECTURE.md index 2844570947..4d35bca6ed 100644 --- a/codex-rs/background-agent/ARCHITECTURE.md +++ b/codex-rs/background-agent/ARCHITECTURE.md @@ -23,6 +23,22 @@ Those legacy records can be linked from a background-agent run for history and compatibility, but they are read-only linkage inputs from the background-agent system's perspective. +## Admission Contract + +Admission is one `BEGIN IMMEDIATE` state transaction. The transaction first +looks up an idempotency key, validates that it is bound to the same request, +source, thread linkage, exact auth-profile alias, config fingerprint, and +admission schema, then either adopts that run or counts live/recoverable rows +and inserts exactly one new run. Stopped or terminal rows do not consume +capacity; queued, owned, waiting, stopping, and recoverable orphaned rows do. + +The CLI, TUI, app server, persisted run, execution snapshot, and daemon pid +record use `codewith.background-agent.admission.v1` as the fail-closed schema +contract. A running daemon is reused only when package version, daemon protocol, +admission schema, and required capability set all match. An explicitly admitted +auth-profile alias must match the app-server profile and remains exact during +recovery; it is never silently replaced by profile auto-switching. + ## Run And Thread Relationship A background-agent run owns background execution. A thread owns transcript and @@ -81,6 +97,20 @@ The app-server may host a live worker bridge, but app-server client connection lifetime is not liveness. Dropping a TUI or CLI connection detaches subscribers only; it must not delete the run or imply worker death. +## Lifecycle Receipts And Fencing + +Lifecycle receipts live in `background_agent_events`, alongside progress +events. Each receipt has a unique `(run_id, receipt_key)` identity and records +the run, generation, attempt, timestamp, and bounded redacted diagnostics. +Retries return the existing receipt instead of advancing the event cursor. +Admission, claim/recovery, first heartbeat for a generation, status transitions, +orphaning, stop, and cancellation all use deterministic receipt keys. + +Supervisor-owned heartbeat, status, stop, and process-finalization mutations +compare both `supervisor_id` and `generation`. A stale owner cannot stop or +complete a reclaimed generation. User-requested stop reloads and retries the +current generation once if ownership changes while the request is in flight. + ## Attach, Detach, Stop, Delete Attach returns a durable snapshot: run row, status snapshot, latest execution diff --git a/codex-rs/background-agent/src/daemon.rs b/codex-rs/background-agent/src/daemon.rs index dbee87e35d..53c1f43450 100644 --- a/codex-rs/background-agent/src/daemon.rs +++ b/codex-rs/background-agent/src/daemon.rs @@ -20,6 +20,9 @@ use crate::process_lifecycle::WorkerProcessHandle; use crate::process_lifecycle::WorkerProcessLogTail; use crate::process_lifecycle::WorkerProcessStatus; use crate::process_lifecycle::WorkerProcessStopReport; +use crate::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; +use crate::BACKGROUND_AGENT_DAEMON_INCOMPATIBLE; +use crate::BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION; const DAEMON_STATE_DIR_NAME: &str = "background-agent-daemon"; const DAEMON_PID_FILE_NAME: &str = "daemon.json"; @@ -31,6 +34,12 @@ const LOCK_POLL_INTERVAL: Duration = Duration::from_millis(50); const LOCK_TIMEOUT: Duration = Duration::from_secs(10); const DAEMON_STOP_GRACE_PERIOD: Duration = Duration::from_secs(35); const DAEMON_HARD_KILL_TIMEOUT: Duration = Duration::from_secs(15); +const DAEMON_CAPABILITIES: &[&str] = &[ + "durable-admission", + "exact-auth-profile", + "generation-fencing", + "lifecycle-receipts", +]; pub fn background_agent_daemon_state_dir(codex_home: &Path) -> PathBuf { codex_home.join(DAEMON_STATE_DIR_NAME) @@ -101,6 +110,7 @@ impl BackgroundAgentDaemon { if let Some(record) = read_pid_record(&self.paths.pid_file()).await? { match self.controller.status(&record.handle).await? { WorkerProcessStatus::Running => { + ensure_daemon_record_compatible(&record)?; return self .output( BackgroundAgentDaemonStatus::AlreadyRunning, @@ -126,6 +136,9 @@ impl BackgroundAgentDaemon { let record = BackgroundAgentDaemonPidRecord { handle, version: env!("CARGO_PKG_VERSION").to_string(), + protocol_version: BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION, + admission_schema_version: BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string(), + capabilities: daemon_capabilities(), }; if let Err(publication_error) = write_pid_record(&self.paths.pid_file(), &record).await { return match self.controller.stop(&record.handle).await { @@ -214,11 +227,23 @@ impl BackgroundAgentDaemon { None => None, }; let handle = record.as_ref().map(|record| &record.handle); + let version = record.as_ref().map(|record| record.version.clone()); + let protocol_version = record.as_ref().map(|record| record.protocol_version); + let admission_schema_version = record + .as_ref() + .map(|record| record.admission_schema_version.clone()); + let capabilities = record + .as_ref() + .map(|record| record.capabilities.clone()) + .unwrap_or_default(); Ok(BackgroundAgentDaemonOutput { status, pid: handle.map(|handle| handle.pid), pgid: handle.and_then(|handle| handle.pgid), - version: record.map(|record| record.version), + version, + protocol_version, + admission_schema_version, + capabilities, state_dir: self.paths.state_dir.clone(), pid_file: self.paths.pid_file(), stderr_log_path: self.paths.stderr_log_path(), @@ -247,6 +272,9 @@ pub struct BackgroundAgentDaemonOutput { pub pid: Option, pub pgid: Option, pub version: Option, + pub protocol_version: Option, + pub admission_schema_version: Option, + pub capabilities: Vec, pub state_dir: PathBuf, pub pid_file: PathBuf, pub stderr_log_path: PathBuf, @@ -259,6 +287,34 @@ pub struct BackgroundAgentDaemonOutput { struct BackgroundAgentDaemonPidRecord { handle: WorkerProcessHandle, version: String, + #[serde(default)] + protocol_version: u32, + #[serde(default)] + admission_schema_version: String, + #[serde(default)] + capabilities: Vec, +} + +fn daemon_capabilities() -> Vec { + DAEMON_CAPABILITIES + .iter() + .map(|capability| (*capability).to_string()) + .collect() +} + +fn ensure_daemon_record_compatible(record: &BackgroundAgentDaemonPidRecord) -> Result<()> { + let compatible = record.version == env!("CARGO_PKG_VERSION") + && record.protocol_version == BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION + && record.admission_schema_version == BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION + && DAEMON_CAPABILITIES + .iter() + .all(|capability| record.capabilities.iter().any(|value| value == capability)); + if !compatible { + bail!( + "{BACKGROUND_AGENT_DAEMON_INCOMPATIBLE}: running daemon package/protocol/schema or capabilities do not match this client" + ); + } + Ok(()) } async fn read_pid_record(path: &Path) -> Result> { @@ -430,7 +486,10 @@ mod tests { .await?; let record = BackgroundAgentDaemonPidRecord { handle: existing.clone(), - version: "test".to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + protocol_version: BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION, + admission_schema_version: BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string(), + capabilities: daemon_capabilities(), }; let daemon = BackgroundAgentDaemon::new(BackgroundAgentDaemonPaths::new( "/bin/false", @@ -446,6 +505,47 @@ mod tests { Ok(()) } + #[tokio::test] + async fn daemon_start_rejects_incompatible_running_pid_record() -> Result<()> { + let temp_dir = TempDir::new().expect("temp dir"); + let controller = WorkerProcessController::with_timeouts( + Duration::from_millis(50), + Duration::from_secs(5), + ); + let existing = controller + .spawn( + WorkerProcessCommand::new("/bin/sh", temp_dir.path().join("existing.stderr.log")) + .arg("-c") + .arg("sleep 60"), + ) + .await?; + let record = BackgroundAgentDaemonPidRecord { + handle: existing.clone(), + version: env!("CARGO_PKG_VERSION").to_string(), + protocol_version: BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION, + admission_schema_version: "older-schema".to_string(), + capabilities: daemon_capabilities(), + }; + let daemon = BackgroundAgentDaemon::new(BackgroundAgentDaemonPaths::new( + "/bin/false", + temp_dir.path(), + )); + write_pid_record(&daemon.paths.pid_file(), &record).await?; + + let error = daemon + .start() + .await + .expect_err("incompatible daemon must not be reused"); + + assert!(error.to_string().contains(BACKGROUND_AGENT_DAEMON_INCOMPATIBLE)); + assert_eq!( + controller.status(&existing).await?, + WorkerProcessStatus::Running + ); + let _ = controller.stop(&existing).await; + Ok(()) + } + #[tokio::test] async fn daemon_start_stops_worker_when_pid_record_write_fails() -> Result<()> { let temp_dir = TempDir::new().expect("temp dir"); @@ -601,6 +701,9 @@ mod tests { stderr_log_path: temp_dir.path().join("stale.stderr.log"), }, version: "test".to_string(), + protocol_version: 0, + admission_schema_version: String::new(), + capabilities: Vec::new(), }; write_pid_record(&daemon.paths.pid_file(), &record).await?; diff --git a/codex-rs/background-agent/src/lib.rs b/codex-rs/background-agent/src/lib.rs index 0b02151c8e..94bcd1c4da 100644 --- a/codex-rs/background-agent/src/lib.rs +++ b/codex-rs/background-agent/src/lib.rs @@ -29,6 +29,21 @@ pub use codex_state::BackgroundAgentWorktreeLeaseCreateParams; pub use supervisor::DurableAgentSupervisor; pub use supervisor::DurableAgentSupervisorConfig; +pub const BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION: &str = + "codewith.background-agent.admission.v1"; +pub const BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED: &str = + "background_agent_admission_capacity_exceeded"; +pub const BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH: &str = + "background_agent_admission_identity_mismatch"; +pub const BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH: &str = + "background_agent_admission_profile_mismatch"; +pub const BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH: &str = + "background_agent_admission_schema_mismatch"; +pub const BACKGROUND_AGENT_DAEMON_INCOMPATIBLE: &str = + "background_agent_daemon_incompatible"; +pub const BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION: u32 = 1; +pub const DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS: i64 = 8; + /// Durable run roster used by the background-agent supervisor. /// /// Implementations are expected to persist run identity, desired state, liveness @@ -40,6 +55,12 @@ pub trait AgentRunStore { params: BackgroundAgentRunCreateParams, ) -> impl Future> + Send; + fn admit_run( + &self, + params: BackgroundAgentRunCreateParams, + max_active_runs: i64, + ) -> impl Future> + Send; + fn get_run( &self, run_id: &str, @@ -75,6 +96,15 @@ pub trait AgentRunStore { fn request_delete_run(&self, run_id: &str) -> impl Future> + Send; + fn request_stop_run( + &self, + run_id: &str, + expected_supervisor_id: Option<&str>, + expected_generation: i64, + status_reason: &str, + diagnostics_json: &serde_json::Value, + ) -> impl Future> + Send; + fn orphan_stale_runs( &self, heartbeat_timeout: Duration, @@ -108,6 +138,15 @@ impl AgentRunStore for codex_state::StateRuntime { self.create_background_agent_run(¶ms).await } + async fn admit_run( + &self, + params: BackgroundAgentRunCreateParams, + max_active_runs: i64, + ) -> anyhow::Result<(BackgroundAgentRun, bool)> { + self.admit_background_agent_run(¶ms, max_active_runs) + .await + } + async fn get_run(&self, run_id: &str) -> anyhow::Result> { self.get_background_agent_run(run_id).await } @@ -151,6 +190,24 @@ impl AgentRunStore for codex_state::StateRuntime { self.request_background_agent_delete(run_id).await } + async fn request_stop_run( + &self, + run_id: &str, + expected_supervisor_id: Option<&str>, + expected_generation: i64, + status_reason: &str, + diagnostics_json: &serde_json::Value, + ) -> anyhow::Result { + self.request_background_agent_stop_for_generation( + run_id, + expected_supervisor_id, + expected_generation, + status_reason, + diagnostics_json, + ) + .await + } + async fn orphan_stale_runs(&self, heartbeat_timeout: Duration) -> anyhow::Result { self.orphan_stale_background_agent_runs(heartbeat_timeout) .await diff --git a/codex-rs/background-agent/src/supervisor.rs b/codex-rs/background-agent/src/supervisor.rs index 4d68bc7325..57b98340bd 100644 --- a/codex-rs/background-agent/src/supervisor.rs +++ b/codex-rs/background-agent/src/supervisor.rs @@ -185,43 +185,37 @@ where ) { return Ok(true); } - self.store() - .set_desired_state(run_id, BackgroundAgentDesiredState::Stopped) + if is_terminal_agent_status(run.status) { + return Ok(true); + } + let terminalize_immediately = should_terminalize_unclaimed_agent_run(&run); + let status_reason = if terminalize_immediately { + "stop requested before worker claim" + } else { + "stop requested" + }; + let stopped = self + .store() + .request_stop_run( + run_id, + run.supervisor_id.as_deref(), + run.generation, + status_reason, + &json!({ + "reason": "supervisor_requested_stop", + "supervisorId": self.config.supervisor_id, + }), + ) + .await?; + if stopped && terminalize_immediately { + cancel_active_pending_interactions_for_run( + self.store(), + run_id, + "supervisor_requested_stop", + ) .await?; - if !is_terminal_agent_status(run.status) { - let terminalize_immediately = should_terminalize_unclaimed_agent_run(&run); - let status = if terminalize_immediately { - BackgroundAgentRunStatus::Cancelled - } else { - BackgroundAgentRunStatus::Stopping - }; - let status_reason = if terminalize_immediately { - "stop requested before worker claim" - } else { - "stop requested" - }; - self.store() - .update_run_status(run_id, status, Some(status_reason)) - .await?; - self.store() - .append_event( - run_id, - "agent.stopRequested", - &json!({ - "reason": "supervisor_requested_stop", - }), - ) - .await?; - if terminalize_immediately { - cancel_active_pending_interactions_for_run( - self.store(), - run_id, - "supervisor_requested_stop", - ) - .await?; - } } - Ok(true) + Ok(stopped) } async fn start_run(&self, run: BackgroundAgentRun) -> anyhow::Result { diff --git a/codex-rs/cli/src/agent_cmd.rs b/codex-rs/cli/src/agent_cmd.rs index cc1dd14515..07e547a1ae 100644 --- a/codex-rs/cli/src/agent_cmd.rs +++ b/codex-rs/cli/src/agent_cmd.rs @@ -4,9 +4,12 @@ use serde_json::Value; use serde_json::json; use std::path::Path; use std::path::PathBuf; +use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; +use codex_background_agent::DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS; use codex_background_agent::daemon::BackgroundAgentDaemon; use codex_background_agent::daemon::BackgroundAgentDaemonPaths; use codex_background_agent::daemon::background_agent_daemon_state_dir; @@ -993,17 +996,6 @@ async fn start_agent( if prompt.is_empty() { anyhow::bail!("agent prompt must not be empty"); } - if let Some(idempotency_key) = cmd.idempotency_key.as_deref() - && let Some(run) = retry_on_busy("load background agent idempotency key", || { - state_db.get_background_agent_run_by_idempotency_key(idempotency_key) - }) - .await - .context("failed to load background agent idempotency key")? - { - let daemon = background_agent_daemon()?; - let daemon_output = daemon.start().await?; - return Ok(json!({ "agent": run_json(run), "created": false, "daemon": daemon_output })); - } ensure_background_agent_supported_platform()?; @@ -1041,21 +1033,35 @@ async fn start_agent( auth_profile_ref: auth_profile_ref.clone(), status_reason: Some("queued by codewith agent start".to_string()), config_fingerprint: None, - version_fingerprint: Some(env!("CARGO_PKG_VERSION").to_string()), + version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), }; - let run = retry_on_busy("create background agent run", || { - state_db.create_background_agent_run(&create_params) + retry_on_busy("reconcile stale background agents before admission", || { + state_db.orphan_stale_background_agent_runs(Duration::from_secs(30)) + }) + .await + .context("failed to reconcile stale background agents before admission")?; + let (run, created) = retry_on_busy("admit background agent run", || { + state_db.admit_background_agent_run( + &create_params, + DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, + ) }) .await - .context("failed to create background agent")?; + .context("failed to admit background agent")?; + if !created { + let daemon = background_agent_daemon()?; + let daemon_output = daemon.start().await?; + return Ok(json!({ "agent": run_json(run), "created": false, "daemon": daemon_output })); + } + let admitted_agent_id = run.id.clone(); let start_event_payload = json!({ "cwd": cwd.display().to_string(), "prompt": prompt, - "promptSnapshotRef": prompt_snapshot_ref, + "promptSnapshotRef": run.prompt_snapshot_ref, }); let event = retry_on_busy("append background agent start event", || { state_db.append_background_agent_event( - agent_id.as_str(), + admitted_agent_id.as_str(), "agent.started", &start_event_payload, ) @@ -1063,7 +1069,7 @@ async fn start_agent( .await .context("failed to append background agent start event")?; let snapshot_params = BackgroundAgentExecutionSnapshotParams { - run_id: agent_id.clone(), + run_id: admitted_agent_id.clone(), snapshot_kind: "initial_execution_context".to_string(), payload_json: json!({ "snapshotSource": "codewith agent start", @@ -1088,7 +1094,7 @@ async fn start_agent( .await .context("failed to create background agent execution snapshot")?; let status_snapshot_params = BackgroundAgentStatusSnapshotParams { - run_id: agent_id, + run_id: admitted_agent_id.clone(), seq: event.seq, status: BackgroundAgentRunStatus::Queued, desired_state: BackgroundAgentDesiredState::Running, @@ -1104,6 +1110,11 @@ async fn start_agent( .context("failed to create background agent status snapshot")?; let daemon = background_agent_daemon()?; let daemon_output = daemon.start().await?; + let run = retry_on_busy("reload admitted background agent", || { + state_db.get_background_agent_run(admitted_agent_id.as_str()) + }) + .await? + .unwrap_or(run); Ok(json!({ "agent": run_json(run), "created": true, "daemon": daemon_output })) } @@ -1156,34 +1167,46 @@ async fn stop_agent( else { return Ok(None); }; - state_db - .set_background_agent_desired_state(agent_id, BackgroundAgentDesiredState::Stopped) - .await - .context("failed to update background agent desired state")?; if !background_agent_status_is_terminal(run.status) { - let terminalize_immediately = should_terminalize_unclaimed_agent_run(&run); - let status = if terminalize_immediately { - BackgroundAgentRunStatus::Cancelled - } else { - BackgroundAgentRunStatus::Stopping - }; - let status_reason = if terminalize_immediately { - "stop requested by codewith agent stop before worker claim" - } else { - "stop requested by codewith agent stop" - }; - state_db - .update_background_agent_run_status(agent_id, status, Some(status_reason)) - .await - .context("failed to update background agent status")?; - state_db - .append_background_agent_event( - agent_id, - "agent.stopRequested", - &json!({"reason": "cli_requested_stop"}), - ) + let mut observed = run; + let mut stopped = false; + let stop_diagnostics = json!({"reason": "cli_requested_stop"}); + for _ in 0..2 { + let status_reason = if should_terminalize_unclaimed_agent_run(&observed) { + "stop requested by codewith agent stop before worker claim" + } else { + "stop requested by codewith agent stop" + }; + stopped = retry_on_busy("request fenced background agent stop", || { + state_db.request_background_agent_stop_for_generation( + agent_id, + observed.supervisor_id.as_deref(), + observed.generation, + status_reason, + &stop_diagnostics, + ) + }) .await - .context("failed to append background agent stop event")?; + .context("failed to request background agent stop")?; + if stopped { + break; + } + let Some(latest) = state_db + .get_background_agent_run(agent_id) + .await + .context("failed to reload background agent during stop")? + else { + return Ok(None); + }; + if background_agent_status_is_terminal(latest.status) { + stopped = true; + break; + } + observed = latest; + } + if !stopped { + anyhow::bail!("background agent ownership changed during stop request"); + } } state_db .get_background_agent_run(agent_id) @@ -1200,7 +1223,7 @@ async fn diagnostics_json(state_db: &StateRuntime) -> anyhow::Result { .count_background_agent_pending_interactions(/*status*/ None) .await .context("failed to count background agent pending interactions")?; - let max_active_runs_per_user = 8_i64; + let max_active_runs_per_user = DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS; let active_run_count = counts .iter() .filter(|(status, _)| { @@ -1212,6 +1235,7 @@ async fn diagnostics_json(state_db: &StateRuntime) -> anyhow::Result { | BackgroundAgentRunStatus::WaitingOnApproval | BackgroundAgentRunStatus::WaitingOnUser | BackgroundAgentRunStatus::Stopping + | BackgroundAgentRunStatus::Orphaned ) }) .map(|(_, count)| *count) diff --git a/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql b/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql new file mode 100644 index 0000000000..0434aa111d --- /dev/null +++ b/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql @@ -0,0 +1,6 @@ +ALTER TABLE background_agent_events + ADD COLUMN receipt_key TEXT; + +CREATE UNIQUE INDEX idx_background_agent_events_receipt_key + ON background_agent_events(run_id, receipt_key) + WHERE receipt_key IS NOT NULL; diff --git a/codex-rs/state/src/runtime/background_agents/events.rs b/codex-rs/state/src/runtime/background_agents/events.rs index 12588c739c..199b7ec20b 100644 --- a/codex-rs/state/src/runtime/background_agents/events.rs +++ b/codex-rs/state/src/runtime/background_agents/events.rs @@ -1,6 +1,9 @@ use super::*; use crate::BACKGROUND_AGENT_EVENT_CURSOR_COMPACTED; +const MAX_BACKGROUND_AGENT_RECEIPT_DIAGNOSTICS_BYTES: usize = 4 * 1024; +const MAX_BACKGROUND_AGENT_RECEIPT_DIAGNOSTICS_PREVIEW_CHARS: usize = 1_024; + pub(in crate::runtime) async fn append_background_agent_event_in_tx( tx: &mut sqlx::Transaction<'_, Sqlite>, run_id: &str, @@ -55,6 +58,93 @@ WHERE id = ? }) } +pub(in crate::runtime) async fn append_background_agent_lifecycle_receipt_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + run_id: &str, + event_type: &str, + receipt_key: &str, + generation: i64, + attempt: Option, + diagnostics_json: &serde_json::Value, + now: i64, +) -> anyhow::Result { + if let Some(row) = sqlx::query_as::<_, BackgroundAgentEventRow>( + r#" +SELECT id, run_id, seq, event_type, payload_json, created_at +FROM background_agent_events +WHERE run_id = ? AND receipt_key = ? + "#, + ) + .bind(run_id) + .bind(receipt_key) + .fetch_optional(&mut **tx) + .await? + { + return BackgroundAgentEvent::try_from(row); + } + + let diagnostics_json = bounded_background_agent_receipt_diagnostics(diagnostics_json)?; + let payload_json = crate::redacted_local_state_json(&serde_json::json!({ + "receiptKey": receipt_key, + "runId": run_id, + "generation": generation, + "attempt": attempt, + "occurredAt": now, + "diagnostics": diagnostics_json, + })); + let serialized_payload = serde_json::to_string(&payload_json)?; + let seq: i64 = sqlx::query_scalar( + "SELECT COALESCE(MAX(seq), 0) + 1 FROM background_agent_events WHERE run_id = ?", + ) + .bind(run_id) + .fetch_one(&mut **tx) + .await?; + let id = sqlx::query( + r#" +INSERT INTO background_agent_events ( + run_id, + seq, + event_type, + payload_json, + created_at, + receipt_key +) VALUES (?, ?, ?, ?, ?, ?) + "#, + ) + .bind(run_id) + .bind(seq) + .bind(event_type) + .bind(serialized_payload) + .bind(now) + .bind(receipt_key) + .execute(&mut **tx) + .await? + .last_insert_rowid(); + + sqlx::query( + r#" +UPDATE background_agent_runs +SET last_event_seq = ?, updated_at = ? +WHERE id = ? + "#, + ) + .bind(seq) + .bind(now) + .bind(run_id) + .execute(&mut **tx) + .await?; + let created_at = DateTime::::from_timestamp(now, 0) + .ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {now}"))?; + Ok(BackgroundAgentEvent { + id, + run_id: run_id.to_string(), + seq, + event_type: event_type.to_string(), + payload_json, + created_at, + }) +} + impl StateRuntime { pub async fn append_background_agent_event( &self, @@ -71,6 +161,32 @@ impl StateRuntime { Ok(event) } + pub async fn append_background_agent_lifecycle_receipt( + &self, + run_id: &str, + event_type: &str, + receipt_key: &str, + generation: i64, + attempt: Option, + diagnostics_json: &serde_json::Value, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let event = append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + run_id, + event_type, + receipt_key, + generation, + attempt, + diagnostics_json, + now, + ) + .await?; + tx.commit().await?; + Ok(event) + } + pub async fn list_background_agent_events_after( &self, run_id: &str, @@ -158,3 +274,22 @@ GROUP BY r.id Ok(()) } } + +fn bounded_background_agent_receipt_diagnostics( + diagnostics_json: &serde_json::Value, +) -> anyhow::Result { + let diagnostics_json = crate::redacted_local_state_json(diagnostics_json); + let serialized = serde_json::to_string(&diagnostics_json)?; + if serialized.len() <= MAX_BACKGROUND_AGENT_RECEIPT_DIAGNOSTICS_BYTES { + return Ok(diagnostics_json); + } + let preview = serialized + .chars() + .take(MAX_BACKGROUND_AGENT_RECEIPT_DIAGNOSTICS_PREVIEW_CHARS) + .collect::(); + Ok(serde_json::json!({ + "truncated": true, + "originalBytes": serialized.len(), + "preview": preview, + })) +} diff --git a/codex-rs/state/src/runtime/background_agents/runs.rs b/codex-rs/state/src/runtime/background_agents/runs.rs index d8d8831fc6..5467544ddd 100644 --- a/codex-rs/state/src/runtime/background_agents/runs.rs +++ b/codex-rs/state/src/runtime/background_agents/runs.rs @@ -1,17 +1,83 @@ use super::*; +const BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED: &str = + "background_agent_admission_capacity_exceeded"; +const BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH: &str = + "background_agent_admission_identity_mismatch"; + impl StateRuntime { pub async fn create_background_agent_run( &self, params: &BackgroundAgentRunCreateParams, ) -> anyhow::Result { + self.create_or_admit_background_agent_run(params, /*max_active_runs*/ None) + .await + .map(|(run, _created)| run) + } + + pub async fn admit_background_agent_run( + &self, + params: &BackgroundAgentRunCreateParams, + max_active_runs: i64, + ) -> anyhow::Result<(BackgroundAgentRun, bool)> { + self.create_or_admit_background_agent_run(params, Some(max_active_runs)) + .await + } + + async fn create_or_admit_background_agent_run( + &self, + params: &BackgroundAgentRunCreateParams, + max_active_runs: Option, + ) -> anyhow::Result<(BackgroundAgentRun, bool)> { let idempotency_key = params.idempotency_key.as_deref().map(redact_state_string); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; if let Some(idempotency_key) = idempotency_key.as_deref() - && let Some(existing) = self - .get_background_agent_run_by_idempotency_key(idempotency_key) - .await? + && let Some(existing_id) = validate_existing_background_agent_admission_in_tx( + &mut tx, + idempotency_key, + params, + ) + .await? { - return Ok(existing); + tx.commit().await?; + let existing = self + .get_background_agent_run(existing_id.as_str()) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "background agent {existing_id} disappeared after idempotent admission" + ) + })?; + return Ok((existing, false)); + } + + if let Some(max_active_runs) = max_active_runs { + let active_run_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM background_agent_runs +WHERE + desired_state = 'running' + AND retention_state = 'active' + AND status IN ( + 'queued', + 'starting', + 'running', + 'waiting_on_approval', + 'waiting_on_user', + 'stopping', + 'orphaned' + ) + "#, + ) + .fetch_one(&mut *tx) + .await?; + if active_run_count >= max_active_runs { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED}: \ + {active_run_count} live or recoverable run(s), max {max_active_runs}" + ); + } } let now = Utc::now().timestamp(); @@ -75,24 +141,46 @@ INSERT INTO background_agent_runs ( .bind(crate::BackgroundAgentRetentionState::Active.as_str()) .bind(now) .bind(now) - .execute(self.pool.as_ref()) + .execute(&mut *tx) .await; if let Err(err) = insert_result { - if idempotency_key.is_some() - && is_background_agent_unique_constraint_violation(&err) - && let Some(idempotency_key) = idempotency_key.as_deref() - && let Some(existing) = self - .get_background_agent_run_by_idempotency_key(idempotency_key) - .await? - { - return Ok(existing); + if is_background_agent_unique_constraint_violation(&err) { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + request identity conflicts with an existing background agent" + ); } return Err(err.into()); } - self.get_background_agent_run(params.id.as_str()) + if max_active_runs.is_some() { + let receipt_key = format!( + "admission:{}", + idempotency_key.as_deref().unwrap_or(params.id.as_str()) + ); + super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + params.id.as_str(), + "agent.admitted", + receipt_key.as_str(), + 0, + Some(1), + &serde_json::json!({ + "source": params.source, + "requestId": params.request_id, + "authProfileRef": params.auth_profile_ref, + "versionFingerprint": params.version_fingerprint, + }), + now, + ) + .await?; + } + tx.commit().await?; + let run = self + .get_background_agent_run(params.id.as_str()) .await? - .ok_or_else(|| anyhow::anyhow!("failed to load background agent run {}", params.id)) + .ok_or_else(|| anyhow::anyhow!("failed to load background agent run {}", params.id))?; + Ok((run, true)) } pub async fn get_background_agent_run( @@ -231,7 +319,18 @@ OFFSET ? r#" SELECT status, COUNT(*) as count FROM background_agent_runs -WHERE retention_state = 'active' +WHERE + retention_state = 'active' + AND desired_state = 'running' + AND status IN ( + 'queued', + 'starting', + 'running', + 'waiting_on_approval', + 'waiting_on_user', + 'stopping', + 'orphaned' + ) GROUP BY status "#, ) @@ -367,10 +466,19 @@ WHERE return Ok(None); } - let event = super::events::append_background_agent_event_in_tx( + let receipt_key = format!( + "status:{}:{}:{}", + params.generation, + params.event_type, + params.status.as_str() + ); + let event = super::events::append_background_agent_lifecycle_receipt_in_tx( &mut tx, params.run_id, params.event_type, + receipt_key.as_str(), + params.generation, + None, params.event_payload_json, now, ) @@ -400,6 +508,160 @@ WHERE Ok(Some(event)) } + pub async fn request_background_agent_stop_for_generation( + &self, + run_id: &str, + expected_supervisor_id: Option<&str>, + expected_generation: i64, + status_reason: &str, + diagnostics_json: &serde_json::Value, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let result = sqlx::query( + r#" +UPDATE background_agent_runs +SET + desired_state = ?, + status = CASE + WHEN supervisor_id IS NULL OR status IN ('queued', 'orphaned') THEN ? + ELSE ? + END, + status_reason = ?, + updated_at = ?, + completed_at = CASE + WHEN supervisor_id IS NULL OR status IN ('queued', 'orphaned') + THEN COALESCE(completed_at, ?) + ELSE completed_at + END +WHERE + id = ? + AND generation = ? + AND ( + (supervisor_id IS NULL AND ? IS NULL) + OR supervisor_id = ? + ) + AND retention_state = ? + AND status IN ( + 'queued', + 'starting', + 'running', + 'waiting_on_approval', + 'waiting_on_user', + 'stopping', + 'orphaned' + ) + "#, + ) + .bind(BackgroundAgentDesiredState::Stopped.as_str()) + .bind(BackgroundAgentRunStatus::Cancelled.as_str()) + .bind(BackgroundAgentRunStatus::Stopping.as_str()) + .bind(redact_state_string(status_reason)) + .bind(now) + .bind(now) + .bind(run_id) + .bind(expected_generation) + .bind(expected_supervisor_id) + .bind(expected_supervisor_id) + .bind(crate::BackgroundAgentRetentionState::Active.as_str()) + .execute(&mut *tx) + .await?; + + if result.rows_affected() == 0 { + let idempotent_match: Option = sqlx::query_scalar( + r#" +SELECT 1 +FROM background_agent_runs +WHERE + id = ? + AND generation = ? + AND ( + (supervisor_id IS NULL AND ? IS NULL) + OR supervisor_id = ? + ) + AND desired_state = ? + AND status IN ('stopping', 'cancelled') + "#, + ) + .bind(run_id) + .bind(expected_generation) + .bind(expected_supervisor_id) + .bind(expected_supervisor_id) + .bind(BackgroundAgentDesiredState::Stopped.as_str()) + .fetch_optional(&mut *tx) + .await?; + tx.commit().await?; + return Ok(idempotent_match.is_some()); + } + + let status = sqlx::query_scalar::<_, String>( + "SELECT status FROM background_agent_runs WHERE id = ?", + ) + .bind(run_id) + .fetch_one(&mut *tx) + .await?; + let status = BackgroundAgentRunStatus::parse(status.as_str())?; + let receipt_key = format!("stop:{expected_generation}"); + let event = super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + run_id, + "agent.stopRequested", + receipt_key.as_str(), + expected_generation, + None, + diagnostics_json, + now, + ) + .await?; + if status == BackgroundAgentRunStatus::Cancelled { + super::interactions::terminalize_active_background_agent_pending_interactions_in_tx( + &mut tx, + run_id, + BackgroundAgentPendingInteractionStatus::Cancelled, + diagnostics_json, + now, + ) + .await?; + } + let last_event_seq: i64 = + sqlx::query_scalar("SELECT last_event_seq FROM background_agent_runs WHERE id = ?") + .bind(run_id) + .fetch_one(&mut *tx) + .await?; + let pending_interaction_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM background_agent_pending_interactions +WHERE run_id = ? AND status IN (?, ?) + "#, + ) + .bind(run_id) + .bind(BackgroundAgentPendingInteractionStatus::Pending.as_str()) + .bind(BackgroundAgentPendingInteractionStatus::Delivered.as_str()) + .fetch_one(&mut *tx) + .await?; + super::snapshots::upsert_background_agent_status_snapshot_in_tx( + &mut tx, + &BackgroundAgentStatusSnapshotParams { + run_id: run_id.to_string(), + seq: last_event_seq, + status, + desired_state: BackgroundAgentDesiredState::Stopped, + summary: Some(status_reason.to_string()), + pending_interaction_count, + last_event_seq, + payload_json: serde_json::json!({ + "reason": status_reason, + "event": diagnostics_json, + }), + }, + now, + ) + .await?; + tx.commit().await?; + Ok(true) + } + pub async fn bind_background_agent_thread( &self, params: &BackgroundAgentThreadBindingParams, @@ -948,6 +1210,21 @@ INSERT INTO background_agent_process_leases ( .execute(&mut *tx) .await?; + let receipt_key = format!("claim:{generation}"); + super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + run_id, + "agent.claimed", + receipt_key.as_str(), + generation, + Some(generation), + &serde_json::json!({ + "supervisorId": supervisor_id, + "processLeaseId": process_lease_id, + }), + now, + ) + .await?; tx.commit().await?; Ok(Some(generation)) } @@ -1014,6 +1291,20 @@ WHERE run_id = ? AND supervisor_id = ? AND generation = ? .bind(params.generation) .execute(&mut *tx) .await?; + let receipt_key = format!("heartbeat:{generation}"); + super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + run_id, + "agent.heartbeat", + receipt_key.as_str(), + generation, + None, + &serde_json::json!({ + "supervisorId": supervisor_id, + }), + now, + ) + .await?; } tx.commit().await?; @@ -1206,6 +1497,111 @@ WHERE idempotency_key = ? } } +async fn validate_existing_background_agent_admission_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + idempotency_key: &str, + params: &BackgroundAgentRunCreateParams, +) -> anyhow::Result> { + let existing = sqlx::query_as::< + _, + ( + String, + Option, + String, + Option, + Option, + String, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + ), + >( + r#" +SELECT + id, + request_id, + source, + input_snapshot_ref, + thread_id, + thread_store_kind, + thread_store_id, + rollout_path, + parent_thread_id, + parent_agent_run_id, + spawn_linkage_json, + auth_profile_ref, + config_fingerprint, + version_fingerprint +FROM background_agent_runs +WHERE idempotency_key = ? + "#, + ) + .bind(idempotency_key) + .fetch_optional(&mut **tx) + .await?; + let Some(( + existing_id, + request_id, + source, + input_snapshot_ref, + thread_id, + thread_store_kind, + thread_store_id, + rollout_path, + parent_thread_id, + parent_agent_run_id, + spawn_linkage_json, + auth_profile_ref, + config_fingerprint, + version_fingerprint, + )) = existing + else { + return Ok(None); + }; + let requested_spawn_linkage_json = params + .spawn_linkage_json + .as_ref() + .map(redact_state_json_string) + .transpose()?; + let identity_matches = request_id + == params + .request_id + .as_deref() + .map(redact_state_string) + && source == params.source + && input_snapshot_ref + == params + .input_snapshot_ref + .as_deref() + .map(redact_state_string) + && thread_id == params.thread_id + && thread_store_kind == params.thread_store_kind + && thread_store_id == params.thread_store_id + && rollout_path == params.rollout_path + && parent_thread_id == params.parent_thread_id + && parent_agent_run_id == params.parent_agent_run_id + && spawn_linkage_json == requested_spawn_linkage_json + && auth_profile_ref + == params + .auth_profile_ref + .as_deref() + .map(redact_state_string) + && config_fingerprint == params.config_fingerprint + && version_fingerprint == params.version_fingerprint; + if !identity_matches { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + idempotency key is already bound to a different background agent identity" + ); + } + Ok(Some(existing_id)) +} + async fn append_terminal_stale_background_agent_status_in_tx( tx: &mut sqlx::Transaction<'_, Sqlite>, run_id: &str, @@ -1215,20 +1611,28 @@ async fn append_terminal_stale_background_agent_status_in_tx( event_payload_json: &serde_json::Value, now: i64, ) -> anyhow::Result<()> { - super::events::append_background_agent_event_in_tx( + let generation: i64 = + sqlx::query_scalar("SELECT generation FROM background_agent_runs WHERE id = ?") + .bind(run_id) + .fetch_one(&mut **tx) + .await?; + let receipt_key = format!("lifecycle:{generation}:{event_type}:{}", status.as_str()); + let event = super::events::append_background_agent_lifecycle_receipt_in_tx( tx, run_id, event_type, + receipt_key.as_str(), + generation, + None, event_payload_json, now, ) .await?; - let (last_event_seq, desired_state): (i64, String) = sqlx::query_as( - "SELECT last_event_seq, desired_state FROM background_agent_runs WHERE id = ?", - ) - .bind(run_id) - .fetch_one(&mut **tx) - .await?; + let desired_state: String = + sqlx::query_scalar("SELECT desired_state FROM background_agent_runs WHERE id = ?") + .bind(run_id) + .fetch_one(&mut **tx) + .await?; let pending_interaction_count: i64 = sqlx::query_scalar( r#" SELECT COUNT(*) @@ -1245,12 +1649,12 @@ WHERE run_id = ? AND status IN (?, ?) tx, &BackgroundAgentStatusSnapshotParams { run_id: run_id.to_string(), - seq: last_event_seq, + seq: event.seq, status, desired_state: BackgroundAgentDesiredState::parse(desired_state.as_str())?, summary: Some(status_reason.to_string()), pending_interaction_count, - last_event_seq, + last_event_seq: event.seq, payload_json: serde_json::json!({ "reason": status_reason, "event": event_payload_json, diff --git a/codex-rs/state/src/runtime/background_agents/tests.rs b/codex-rs/state/src/runtime/background_agents/tests.rs index 20a5ff5f55..b9f04363a6 100644 --- a/codex-rs/state/src/runtime/background_agents/tests.rs +++ b/codex-rs/state/src/runtime/background_agents/tests.rs @@ -80,6 +80,32 @@ async fn create_run_with_id( .await } +fn admission_params( + id: &str, + idempotency_key: &str, + auth_profile_ref: &str, +) -> BackgroundAgentRunCreateParams { + BackgroundAgentRunCreateParams { + id: id.to_string(), + idempotency_key: Some(idempotency_key.to_string()), + request_id: Some(format!("request-{idempotency_key}")), + source: "admission-test".to_string(), + prompt_snapshot_ref: format!("inline:{id}:prompt"), + input_snapshot_ref: None, + thread_id: Some(format!("thread-{idempotency_key}")), + thread_store_kind: "background-agent".to_string(), + thread_store_id: Some("state.sqlite".to_string()), + rollout_path: None, + parent_thread_id: Some("parent-thread".to_string()), + parent_agent_run_id: Some("parent-run".to_string()), + spawn_linkage_json: Some(json!({"agentPath": ["worker"]})), + auth_profile_ref: Some(auth_profile_ref.to_string()), + status_reason: Some("queued by admission test".to_string()), + config_fingerprint: Some("config-v1".to_string()), + version_fingerprint: Some("codewith.background-agent.admission.v1".to_string()), + } +} + #[tokio::test] async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; @@ -88,21 +114,21 @@ async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { .create_background_agent_run(&BackgroundAgentRunCreateParams { id: "run-duplicate".to_string(), idempotency_key: Some("idem-1".to_string()), - request_id: Some("req-duplicate".to_string()), + request_id: Some("req-1".to_string()), source: "cli".to_string(), prompt_snapshot_ref: "prompt://duplicate".to_string(), - input_snapshot_ref: None, - thread_id: None, + input_snapshot_ref: Some("input://run-1".to_string()), + thread_id: Some("thread-1".to_string()), thread_store_kind: "local".to_string(), - thread_store_id: None, - rollout_path: None, - parent_thread_id: None, + thread_store_id: Some("state_5.sqlite".to_string()), + rollout_path: Some("/tmp/rollout.jsonl".to_string()), + parent_thread_id: Some("parent-thread".to_string()), parent_agent_run_id: None, - spawn_linkage_json: None, - auth_profile_ref: None, + spawn_linkage_json: Some(json!({"agentPath": ["reviewer"]})), + auth_profile_ref: Some("profile:default".to_string()), status_reason: None, - config_fingerprint: None, - version_fingerprint: None, + config_fingerprint: Some("cfg-1".to_string()), + version_fingerprint: Some("version-1".to_string()), }) .await?; @@ -114,6 +140,160 @@ async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn background_agent_admission_create_or_adopt_is_atomic_and_receipted() +-> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let first_params = admission_params("admitted-1", "admission-key", "profile-a"); + let (first, created) = runtime + .admit_background_agent_run(&first_params, /*max_active_runs*/ 2) + .await?; + assert!(created); + + let retry_params = admission_params("admitted-retry", "admission-key", "profile-a"); + let (retry, created) = runtime + .admit_background_agent_run(&retry_params, /*max_active_runs*/ 2) + .await?; + assert!(!created); + assert_eq!(retry.id, first.id); + assert_eq!(runtime.list_background_agent_runs(None).await?.len(), 1); + let events = runtime + .list_background_agent_events_after(first.id.as_str(), None, None) + .await?; + assert_eq!(events.len(), 1); + assert_eq!(events[0].event_type, "agent.admitted"); + assert_eq!( + events[0].payload_json.get("receiptKey"), + Some(&json!("admission:admission-key")) + ); + Ok(()) +} + +#[tokio::test] +async fn background_agent_admission_rejects_idempotency_identity_mismatch() +-> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + runtime + .admit_background_agent_run( + &admission_params("admitted-1", "admission-key", "profile-a"), + 2, + ) + .await?; + + let error = runtime + .admit_background_agent_run( + &admission_params("admitted-2", "admission-key", "profile-b"), + 2, + ) + .await + .expect_err("profile mismatch must not adopt the existing run"); + + assert!( + error + .to_string() + .contains("background_agent_admission_identity_mismatch") + ); + assert_eq!(runtime.list_background_agent_runs(None).await?.len(), 1); + Ok(()) +} + +#[tokio::test] +async fn background_agent_admission_counts_only_live_or_recoverable_runs() +-> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let (first, _) = runtime + .admit_background_agent_run( + &admission_params("admitted-1", "admission-key-1", "profile-a"), + 1, + ) + .await?; + assert!( + runtime + .request_background_agent_stop_for_generation( + first.id.as_str(), + None, + 0, + "capacity test stop", + &json!({"reason": "capacity_test"}), + ) + .await? + ); + let (_, created) = runtime + .admit_background_agent_run( + &admission_params("admitted-2", "admission-key-2", "profile-a"), + 1, + ) + .await?; + assert!(created); + + runtime + .update_background_agent_run_status( + "admitted-2", + BackgroundAgentRunStatus::Orphaned, + Some("recoverable orphan"), + ) + .await?; + let error = runtime + .admit_background_agent_run( + &admission_params("admitted-3", "admission-key-3", "profile-a"), + 1, + ) + .await + .expect_err("recoverable orphan must consume capacity"); + assert!( + error + .to_string() + .contains("background_agent_admission_capacity_exceeded") + ); + Ok(()) +} + +#[tokio::test] +async fn background_agent_lifecycle_receipts_dedupe_redact_and_bound_diagnostics() +-> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let run = create_run(runtime.as_ref()).await?; + let diagnostics = json!({ + "apiKey": "sk-secret-test-value", + "blob": "x".repeat(8 * 1024), + }); + + let first = runtime + .append_background_agent_lifecycle_receipt( + run.id.as_str(), + "agent.testReceipt", + "test-receipt", + 1, + Some(1), + &diagnostics, + ) + .await?; + let retry = runtime + .append_background_agent_lifecycle_receipt( + run.id.as_str(), + "agent.testReceipt", + "test-receipt", + 1, + Some(2), + &diagnostics, + ) + .await?; + + assert_eq!(retry.id, first.id); + assert_eq!(retry.seq, first.seq); + let serialized = serde_json::to_string(&retry.payload_json)?; + assert!(!serialized.contains("sk-secret-test-value")); + assert!(serialized.len() < 2 * 1024); + assert_eq!( + retry + .payload_json + .pointer("/diagnostics/truncated") + .and_then(serde_json::Value::as_bool), + Some(true) + ); + Ok(()) +} + #[tokio::test] async fn legacy_thread_and_agent_job_rows_do_not_populate_background_agent_roster() -> anyhow::Result<()> { @@ -924,6 +1104,63 @@ async fn stale_generation_cannot_update_status_or_create_interactions_after_recl Ok(()) } +#[tokio::test] +async fn stale_generation_cannot_cancel_reclaimed_run() -> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + create_run(runtime.as_ref()).await?; + let first_generation = runtime + .claim_background_agent_supervisor("run-1", "supervisor-1", "lease-1") + .await? + .expect("run should be claimed"); + assert_eq!( + runtime + .orphan_stale_background_agent_runs(Duration::ZERO) + .await?, + 1 + ); + let second_generation = runtime + .claim_background_agent_supervisor("run-1", "supervisor-2", "lease-2") + .await? + .expect("orphaned run should be reclaimed"); + + assert!( + !runtime + .request_background_agent_stop_for_generation( + "run-1", + Some("supervisor-1"), + first_generation, + "stale stop", + &json!({"reason": "stale_stop"}), + ) + .await? + ); + let running = runtime + .get_background_agent_run("run-1") + .await? + .expect("run should exist"); + assert_eq!(running.desired_state, BackgroundAgentDesiredState::Running); + assert_eq!(running.supervisor_id.as_deref(), Some("supervisor-2")); + + assert!( + runtime + .request_background_agent_stop_for_generation( + "run-1", + Some("supervisor-2"), + second_generation, + "current stop", + &json!({"reason": "current_stop"}), + ) + .await? + ); + let stopped = runtime + .get_background_agent_run("run-1") + .await? + .expect("run should exist"); + assert_eq!(stopped.desired_state, BackgroundAgentDesiredState::Stopped); + assert_eq!(stopped.status, BackgroundAgentRunStatus::Stopping); + Ok(()) +} + #[tokio::test] async fn stale_stopping_run_is_cancelled_and_lease_stopped() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index bd170e5841..6489909413 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -265,6 +265,10 @@ const JSONRPC_INVALID_REQUEST: i64 = -32600; const JSONRPC_METHOD_NOT_FOUND: i64 = -32601; const MISSION_CONTROL_OVERVIEW_LIMIT: u32 = 50; const THREAD_SETTINGS_UPDATE_METHOD: &str = "thread/settings/update"; +// The TUI intentionally depends on the wire protocol rather than the server +// implementation crate; the server validates this explicit contract value. +const BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION: &str = + "codewith.background-agent.admission.v1"; fn bootstrap_request_error(context: &'static str, err: TypedRequestError) -> color_eyre::Report { color_eyre::eyre::eyre!("{context}: {err}") @@ -1651,7 +1655,9 @@ impl AppServerSession { spawn_linkage: None, auth_profile_ref, config_fingerprint: None, - version_fingerprint: None, + version_fingerprint: Some( + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string(), + ), execution_context: workspace_roots.map(|workspace_roots| { Box::new(AgentExecutionContextParams { workspace_roots: Some(workspace_roots), From a4dc676dc0f5f3b31d1eae8281cb8860c385e469 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 18:21:58 +0300 Subject: [PATCH 02/15] style(background-agent): apply repository formatting --- .../background_agent_live.rs | 3 +-- .../background_agent_processor.rs | 22 ++++++++---------- .../tests/suite/v2/background_agent.rs | 23 +++++++++++-------- codex-rs/background-agent/src/daemon.rs | 12 ++++++---- codex-rs/background-agent/src/lib.rs | 3 +-- codex-rs/cli/src/agent_cmd.rs | 6 ++--- .../src/runtime/background_agents/runs.rs | 23 +++++-------------- .../src/runtime/background_agents/tests.rs | 10 ++++---- codex-rs/tui/src/app_server_session.rs | 3 +-- 9 files changed, 46 insertions(+), 59 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/background_agent_live.rs b/codex-rs/app-server/src/request_processors/background_agent_live.rs index daeecc95df..22b809e869 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_live.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_live.rs @@ -1675,8 +1675,7 @@ impl ThreadRequestProcessor { &self, params: &AgentStartParams, ) -> Result<(), JSONRPCErrorError> { - if params.version_fingerprint.as_deref() - != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) + if params.version_fingerprint.as_deref() != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) { let mut error = invalid_request("background agent admission schema is incompatible"); error.data = Some(json!({ diff --git a/codex-rs/app-server/src/request_processors/background_agent_processor.rs b/codex-rs/app-server/src/request_processors/background_agent_processor.rs index 50499ab82f..cdc7412f49 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_processor.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_processor.rs @@ -61,10 +61,10 @@ use codex_background_agent::AgentEventJournal; use codex_background_agent::AgentRunStore; use codex_background_agent::AgentSnapshotStore; use codex_background_agent::BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED; -use codex_background_agent::BACKGROUND_AGENT_EVENT_CURSOR_COMPACTED; use codex_background_agent::BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH; use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH; use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; +use codex_background_agent::BACKGROUND_AGENT_EVENT_CURSOR_COMPACTED; use codex_background_agent::BackgroundAgentDesiredState; use codex_background_agent::BackgroundAgentEvent; use codex_background_agent::BackgroundAgentExecutionSnapshot; @@ -77,10 +77,10 @@ use codex_background_agent::BackgroundAgentRunCreateParams; use codex_background_agent::BackgroundAgentRunStatus; use codex_background_agent::BackgroundAgentStatusSnapshot; use codex_background_agent::BackgroundAgentStatusSnapshotParams; +use codex_background_agent::DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS; use codex_background_agent::LifecycleAction; use codex_background_agent::LifecycleEffect; use codex_background_agent::PendingInteractionLedger; -use codex_background_agent::DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS; use codex_background_agent::lifecycle_effect_for; use codex_protocol::ThreadId; use codex_protocol::approvals::ElicitationAction; @@ -187,15 +187,12 @@ impl BackgroundAgentRequestProcessor { config_fingerprint, version_fingerprint, }; - let (run, created_new_run) = retry_transient_sqlite_busy( - "admit background agent", - || { - state_db.admit_run( - create_params.clone(), - DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, - ) - }, - ) + let (run, created_new_run) = retry_transient_sqlite_busy("admit background agent", || { + state_db.admit_run( + create_params.clone(), + DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, + ) + }) .await .map_err(map_background_agent_admission_error)?; let execution_payload = initial_execution_snapshot_payload( @@ -823,8 +820,7 @@ impl BackgroundAgentRequestProcessor { &self, ) -> Result { let Some(state_db) = self.state_db.clone() else { - let quota = - AgentQuotaSnapshot::empty(DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS); + let quota = AgentQuotaSnapshot::empty(DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS); return Ok(AgentDaemonDiagnosticsResponse { state_store_available: false, active_run_count: quota.active_run_count, diff --git a/codex-rs/app-server/tests/suite/v2/background_agent.rs b/codex-rs/app-server/tests/suite/v2/background_agent.rs index 7588ebc420..82b6834898 100644 --- a/codex-rs/app-server/tests/suite/v2/background_agent.rs +++ b/codex-rs/app-server/tests/suite/v2/background_agent.rs @@ -48,9 +48,6 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; -use codex_background_agent::BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH; -use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH; -use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; use codex_app_server_protocol::WorktreeAttachResponse; use codex_app_server_protocol::WorktreeCleanupPolicy; use codex_app_server_protocol::WorktreeCleanupResponse; @@ -66,6 +63,9 @@ use codex_app_server_protocol::WorktreeMergeCandidateStatus; use codex_app_server_protocol::WorktreeOwnerKind; use codex_app_server_protocol::WorktreeReadResponse; use codex_app_server_protocol::WorktreeReleaseResponse; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; use codex_protocol::ThreadId; use codex_protocol::models::PermissionProfile; use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS; @@ -213,10 +213,7 @@ async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result ) .await?; assert_eq!(second_events_page.data.len(), 1); - assert_eq!( - second_events_page.data[0].event_type, - "agent.started" - ); + assert_eq!(second_events_page.data[0].event_type, "agent.started"); assert_eq!(second_events_page.next_cursor, Some("event:2".to_string())); let all_events = agent_events_page(&mut restarted, &agent_id, /*cursor*/ None, Some(20)).await?; @@ -474,7 +471,11 @@ async fn agent_start_rejects_profile_and_schema_mismatches_before_admission() -> profile_mismatch.auth_profile_ref = Some("client-selected-auth-profile".to_string()); let error = start_agent_error(&mut mcp, profile_mismatch).await?; assert_eq!( - error.error.data.as_ref().and_then(|data| data.get("errorCode")), + error + .error + .data + .as_ref() + .and_then(|data| data.get("errorCode")), Some(&json!(BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH)) ); @@ -486,7 +487,11 @@ async fn agent_start_rejects_profile_and_schema_mismatches_before_admission() -> schema_mismatch.version_fingerprint = Some("older-schema".to_string()); let error = start_agent_error(&mut mcp, schema_mismatch).await?; assert_eq!( - error.error.data.as_ref().and_then(|data| data.get("errorCode")), + error + .error + .data + .as_ref() + .and_then(|data| data.get("errorCode")), Some(&json!(BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH)) ); assert!(agent_list(&mut mcp).await?.data.is_empty()); diff --git a/codex-rs/background-agent/src/daemon.rs b/codex-rs/background-agent/src/daemon.rs index 53c1f43450..dc0ef7eaf2 100644 --- a/codex-rs/background-agent/src/daemon.rs +++ b/codex-rs/background-agent/src/daemon.rs @@ -14,15 +14,15 @@ use tokio::time::Instant; #[cfg(unix)] use tokio::time::sleep; +use crate::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; +use crate::BACKGROUND_AGENT_DAEMON_INCOMPATIBLE; +use crate::BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION; use crate::process_lifecycle::WorkerProcessCommand; use crate::process_lifecycle::WorkerProcessController; use crate::process_lifecycle::WorkerProcessHandle; use crate::process_lifecycle::WorkerProcessLogTail; use crate::process_lifecycle::WorkerProcessStatus; use crate::process_lifecycle::WorkerProcessStopReport; -use crate::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; -use crate::BACKGROUND_AGENT_DAEMON_INCOMPATIBLE; -use crate::BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION; const DAEMON_STATE_DIR_NAME: &str = "background-agent-daemon"; const DAEMON_PID_FILE_NAME: &str = "daemon.json"; @@ -537,7 +537,11 @@ mod tests { .await .expect_err("incompatible daemon must not be reused"); - assert!(error.to_string().contains(BACKGROUND_AGENT_DAEMON_INCOMPATIBLE)); + assert!( + error + .to_string() + .contains(BACKGROUND_AGENT_DAEMON_INCOMPATIBLE) + ); assert_eq!( controller.status(&existing).await?, WorkerProcessStatus::Running diff --git a/codex-rs/background-agent/src/lib.rs b/codex-rs/background-agent/src/lib.rs index 94bcd1c4da..953dd08cdf 100644 --- a/codex-rs/background-agent/src/lib.rs +++ b/codex-rs/background-agent/src/lib.rs @@ -39,8 +39,7 @@ pub const BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH: &str = "background_agent_admission_profile_mismatch"; pub const BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH: &str = "background_agent_admission_schema_mismatch"; -pub const BACKGROUND_AGENT_DAEMON_INCOMPATIBLE: &str = - "background_agent_daemon_incompatible"; +pub const BACKGROUND_AGENT_DAEMON_INCOMPATIBLE: &str = "background_agent_daemon_incompatible"; pub const BACKGROUND_AGENT_DAEMON_PROTOCOL_VERSION: u32 = 1; pub const DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS: i64 = 8; diff --git a/codex-rs/cli/src/agent_cmd.rs b/codex-rs/cli/src/agent_cmd.rs index 07e547a1ae..49f7917da1 100644 --- a/codex-rs/cli/src/agent_cmd.rs +++ b/codex-rs/cli/src/agent_cmd.rs @@ -1041,10 +1041,8 @@ async fn start_agent( .await .context("failed to reconcile stale background agents before admission")?; let (run, created) = retry_on_busy("admit background agent run", || { - state_db.admit_background_agent_run( - &create_params, - DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, - ) + state_db + .admit_background_agent_run(&create_params, DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS) }) .await .context("failed to admit background agent")?; diff --git a/codex-rs/state/src/runtime/background_agents/runs.rs b/codex-rs/state/src/runtime/background_agents/runs.rs index 5467544ddd..556bc9d59a 100644 --- a/codex-rs/state/src/runtime/background_agents/runs.rs +++ b/codex-rs/state/src/runtime/background_agents/runs.rs @@ -32,12 +32,9 @@ impl StateRuntime { let idempotency_key = params.idempotency_key.as_deref().map(redact_state_string); let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; if let Some(idempotency_key) = idempotency_key.as_deref() - && let Some(existing_id) = validate_existing_background_agent_admission_in_tx( - &mut tx, - idempotency_key, - params, - ) - .await? + && let Some(existing_id) = + validate_existing_background_agent_admission_in_tx(&mut tx, idempotency_key, params) + .await? { tx.commit().await?; let existing = self @@ -602,7 +599,7 @@ WHERE .await?; let status = BackgroundAgentRunStatus::parse(status.as_str())?; let receipt_key = format!("stop:{expected_generation}"); - let event = super::events::append_background_agent_lifecycle_receipt_in_tx( + super::events::append_background_agent_lifecycle_receipt_in_tx( &mut tx, run_id, "agent.stopRequested", @@ -1568,11 +1565,7 @@ WHERE idempotency_key = ? .as_ref() .map(redact_state_json_string) .transpose()?; - let identity_matches = request_id - == params - .request_id - .as_deref() - .map(redact_state_string) + let identity_matches = request_id == params.request_id.as_deref().map(redact_state_string) && source == params.source && input_snapshot_ref == params @@ -1586,11 +1579,7 @@ WHERE idempotency_key = ? && parent_thread_id == params.parent_thread_id && parent_agent_run_id == params.parent_agent_run_id && spawn_linkage_json == requested_spawn_linkage_json - && auth_profile_ref - == params - .auth_profile_ref - .as_deref() - .map(redact_state_string) + && auth_profile_ref == params.auth_profile_ref.as_deref().map(redact_state_string) && config_fingerprint == params.config_fingerprint && version_fingerprint == params.version_fingerprint; if !identity_matches { diff --git a/codex-rs/state/src/runtime/background_agents/tests.rs b/codex-rs/state/src/runtime/background_agents/tests.rs index b9f04363a6..96c8fdcfbe 100644 --- a/codex-rs/state/src/runtime/background_agents/tests.rs +++ b/codex-rs/state/src/runtime/background_agents/tests.rs @@ -141,8 +141,8 @@ async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { } #[tokio::test] -async fn background_agent_admission_create_or_adopt_is_atomic_and_receipted() --> anyhow::Result<()> { +async fn background_agent_admission_create_or_adopt_is_atomic_and_receipted() -> anyhow::Result<()> +{ let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; let first_params = admission_params("admitted-1", "admission-key", "profile-a"); let (first, created) = runtime @@ -170,8 +170,7 @@ async fn background_agent_admission_create_or_adopt_is_atomic_and_receipted() } #[tokio::test] -async fn background_agent_admission_rejects_idempotency_identity_mismatch() --> anyhow::Result<()> { +async fn background_agent_admission_rejects_idempotency_identity_mismatch() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; runtime .admit_background_agent_run( @@ -198,8 +197,7 @@ async fn background_agent_admission_rejects_idempotency_identity_mismatch() } #[tokio::test] -async fn background_agent_admission_counts_only_live_or_recoverable_runs() --> anyhow::Result<()> { +async fn background_agent_admission_counts_only_live_or_recoverable_runs() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; let (first, _) = runtime .admit_background_agent_run( diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 6489909413..72eb6fc9aa 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -267,8 +267,7 @@ const MISSION_CONTROL_OVERVIEW_LIMIT: u32 = 50; const THREAD_SETTINGS_UPDATE_METHOD: &str = "thread/settings/update"; // The TUI intentionally depends on the wire protocol rather than the server // implementation crate; the server validates this explicit contract value. -const BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION: &str = - "codewith.background-agent.admission.v1"; +const BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION: &str = "codewith.background-agent.admission.v1"; fn bootstrap_request_error(context: &'static str, err: TypedRequestError) -> color_eyre::Report { color_eyre::eyre::eyre!("{context}: {err}") From 2e5c18242bb36082335731674a15264044019eef Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 18:30:00 +0300 Subject: [PATCH 03/15] fix(background-agent): recover partial CLI admission --- codex-rs/background-agent/ARCHITECTURE.md | 13 +++ codex-rs/cli/src/agent_cmd.rs | 108 ++++++++++++------ .../src/runtime/background_agents/runs.rs | 8 +- 3 files changed, 92 insertions(+), 37 deletions(-) diff --git a/codex-rs/background-agent/ARCHITECTURE.md b/codex-rs/background-agent/ARCHITECTURE.md index 4d35bca6ed..7116561696 100644 --- a/codex-rs/background-agent/ARCHITECTURE.md +++ b/codex-rs/background-agent/ARCHITECTURE.md @@ -39,6 +39,16 @@ admission schema, and required capability set all match. An explicitly admitted auth-profile alias must match the app-server profile and remains exact during recovery; it is never silently replaced by profile auto-switching. +Cross-system execution references remain a projection, not a second task or PR +lifecycle store. Callers place the authoritative Todos root, PR group, leaf, +worker-run, writer-generation, and attempt references in `spawn_linkage_json`; +the Repos-owned lease remains `worktree_lease_id`. Both are reached from every +lifecycle receipt through its foreign-keyed `runId`. The supervisor +`generation` is only a local process-fencing counter and must never be +interpreted as the projected writer generation. Auth-profile references are +opaque validated aliases: state redaction runs before persistence, and +credential or account payloads do not belong in either reference projection. + ## Run And Thread Relationship A background-agent run owns background execution. A thread owns transcript and @@ -105,6 +115,9 @@ the run, generation, attempt, timestamp, and bounded redacted diagnostics. Retries return the existing receipt instead of advancing the event cursor. Admission, claim/recovery, first heartbeat for a generation, status transitions, orphaning, stop, and cancellation all use deterministic receipt keys. +Terminal receipts therefore replay with the same run projection and attempt +binding, while receipt insertion and cursor advancement commit in one state +transaction. Supervisor-owned heartbeat, status, stop, and process-finalization mutations compare both `supervisor_id` and `generation`. A stale owner cannot stop or diff --git a/codex-rs/cli/src/agent_cmd.rs b/codex-rs/cli/src/agent_cmd.rs index 49f7917da1..e35dbcdf29 100644 --- a/codex-rs/cli/src/agent_cmd.rs +++ b/codex-rs/cli/src/agent_cmd.rs @@ -1046,26 +1046,51 @@ async fn start_agent( }) .await .context("failed to admit background agent")?; - if !created { - let daemon = background_agent_daemon()?; - let daemon_output = daemon.start().await?; - return Ok(json!({ "agent": run_json(run), "created": false, "daemon": daemon_output })); - } let admitted_agent_id = run.id.clone(); let start_event_payload = json!({ "cwd": cwd.display().to_string(), "prompt": prompt, "promptSnapshotRef": run.prompt_snapshot_ref, }); - let event = retry_on_busy("append background agent start event", || { - state_db.append_background_agent_event( - admitted_agent_id.as_str(), - "agent.started", - &start_event_payload, - ) - }) - .await - .context("failed to append background agent start event")?; + let existing_start_event = if created { + None + } else { + retry_on_busy("load background agent start event", || { + state_db.list_background_agent_events_after( + admitted_agent_id.as_str(), + /*after_seq*/ None, + Some(100), + ) + }) + .await + .context("failed to load background agent start event")? + .into_iter() + .find(|event| { + matches!( + event.event_type.as_str(), + "agent.started" | "agent.startRecovered" + ) + }) + }; + let event = match existing_start_event { + Some(event) => event, + None => { + let event_type = if created { + "agent.started" + } else { + "agent.startRecovered" + }; + retry_on_busy("append background agent start event", || { + state_db.append_background_agent_event( + admitted_agent_id.as_str(), + event_type, + &start_event_payload, + ) + }) + .await + .context("failed to append background agent start event")? + } + }; let snapshot_params = BackgroundAgentExecutionSnapshotParams { run_id: admitted_agent_id.clone(), snapshot_kind: "initial_execution_context".to_string(), @@ -1086,26 +1111,43 @@ async fn start_agent( recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), config_fingerprint: None, }; - retry_on_busy("create background agent execution snapshot", || { - state_db.create_background_agent_execution_snapshot(&snapshot_params) - }) - .await - .context("failed to create background agent execution snapshot")?; - let status_snapshot_params = BackgroundAgentStatusSnapshotParams { - run_id: admitted_agent_id.clone(), - seq: event.seq, - status: BackgroundAgentRunStatus::Queued, - desired_state: BackgroundAgentDesiredState::Running, - summary: Some("Queued".to_string()), - pending_interaction_count: 0, - last_event_seq: event.seq, - payload_json: json!({"phase": "queued"}), - }; - retry_on_busy("create background agent status snapshot", || { - state_db.upsert_background_agent_status_snapshot(&status_snapshot_params) + let execution_snapshot_exists = + retry_on_busy("load background agent execution snapshot", || { + state_db.get_latest_background_agent_execution_snapshot(admitted_agent_id.as_str()) + }) + .await + .context("failed to load background agent execution snapshot")? + .is_some(); + if !execution_snapshot_exists { + retry_on_busy("create background agent execution snapshot", || { + state_db.create_background_agent_execution_snapshot(&snapshot_params) + }) + .await + .context("failed to create background agent execution snapshot")?; + } + let status_snapshot_exists = retry_on_busy("load background agent status snapshot", || { + state_db.get_background_agent_status_snapshot(admitted_agent_id.as_str()) }) .await - .context("failed to create background agent status snapshot")?; + .context("failed to load background agent status snapshot")? + .is_some(); + if !status_snapshot_exists { + let status_snapshot_params = BackgroundAgentStatusSnapshotParams { + run_id: admitted_agent_id.clone(), + seq: event.seq, + status: BackgroundAgentRunStatus::Queued, + desired_state: BackgroundAgentDesiredState::Running, + summary: Some("Queued".to_string()), + pending_interaction_count: 0, + last_event_seq: event.seq, + payload_json: json!({"phase": "queued"}), + }; + retry_on_busy("create background agent status snapshot", || { + state_db.upsert_background_agent_status_snapshot(&status_snapshot_params) + }) + .await + .context("failed to create background agent status snapshot")?; + } let daemon = background_agent_daemon()?; let daemon_output = daemon.start().await?; let run = retry_on_busy("reload admitted background agent", || { @@ -1113,7 +1155,7 @@ async fn start_agent( }) .await? .unwrap_or(run); - Ok(json!({ "agent": run_json(run), "created": true, "daemon": daemon_output })) + Ok(json!({ "agent": run_json(run), "created": created, "daemon": daemon_output })) } fn agent_start_snapshot_workspace_roots( diff --git a/codex-rs/state/src/runtime/background_agents/runs.rs b/codex-rs/state/src/runtime/background_agents/runs.rs index 556bc9d59a..67231a47c1 100644 --- a/codex-rs/state/src/runtime/background_agents/runs.rs +++ b/codex-rs/state/src/runtime/background_agents/runs.rs @@ -1288,16 +1288,16 @@ WHERE run_id = ? AND supervisor_id = ? AND generation = ? .bind(params.generation) .execute(&mut *tx) .await?; - let receipt_key = format!("heartbeat:{generation}"); + let receipt_key = format!("heartbeat:{}", params.generation); super::events::append_background_agent_lifecycle_receipt_in_tx( &mut tx, - run_id, + params.run_id, "agent.heartbeat", receipt_key.as_str(), - generation, + params.generation, None, &serde_json::json!({ - "supervisorId": supervisor_id, + "supervisorId": params.supervisor_id, }), now, ) From aca308be6a776572cb431029c3194eba0bad98dc Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 18:32:59 +0300 Subject: [PATCH 04/15] fix(background-agent): preserve recovered run status --- codex-rs/cli/src/agent_cmd.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/codex-rs/cli/src/agent_cmd.rs b/codex-rs/cli/src/agent_cmd.rs index e35dbcdf29..fb54fefb7b 100644 --- a/codex-rs/cli/src/agent_cmd.rs +++ b/codex-rs/cli/src/agent_cmd.rs @@ -1132,15 +1132,22 @@ async fn start_agent( .context("failed to load background agent status snapshot")? .is_some(); if !status_snapshot_exists { + let snapshot_run = retry_on_busy("reload background agent for status snapshot", || { + state_db.get_background_agent_run(admitted_agent_id.as_str()) + }) + .await + .context("failed to reload background agent for status snapshot")? + .unwrap_or_else(|| run.clone()); + let last_event_seq = snapshot_run.last_event_seq.max(event.seq); let status_snapshot_params = BackgroundAgentStatusSnapshotParams { run_id: admitted_agent_id.clone(), - seq: event.seq, - status: BackgroundAgentRunStatus::Queued, - desired_state: BackgroundAgentDesiredState::Running, - summary: Some("Queued".to_string()), + seq: last_event_seq, + status: snapshot_run.status, + desired_state: snapshot_run.desired_state, + summary: Some(snapshot_run.status.as_str().to_string()), pending_interaction_count: 0, - last_event_seq: event.seq, - payload_json: json!({"phase": "queued"}), + last_event_seq, + payload_json: json!({"phase": snapshot_run.status.as_str()}), }; retry_on_busy("create background agent status snapshot", || { state_db.upsert_background_agent_status_snapshot(&status_snapshot_params) From 1a55a367b177f7ce304558b501412fc7a110916a Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 18:36:49 +0300 Subject: [PATCH 05/15] test: account for durable lifecycle receipts --- codex-rs/background-agent/src/supervisor.rs | 5 ++++- .../state/src/runtime/background_agents/tests.rs | 15 +++++++++++---- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/codex-rs/background-agent/src/supervisor.rs b/codex-rs/background-agent/src/supervisor.rs index 57b98340bd..b89591a39b 100644 --- a/codex-rs/background-agent/src/supervisor.rs +++ b/codex-rs/background-agent/src/supervisor.rs @@ -544,7 +544,10 @@ mod tests { .into_iter() .map(|event| event.event_type) .collect::>(), - vec!["agent.workerStartFailed".to_string()] + vec![ + "agent.claimed".to_string(), + "agent.workerStartFailed".to_string() + ] ); Ok(()) } diff --git a/codex-rs/state/src/runtime/background_agents/tests.rs b/codex-rs/state/src/runtime/background_agents/tests.rs index 96c8fdcfbe..9c316b2c93 100644 --- a/codex-rs/state/src/runtime/background_agents/tests.rs +++ b/codex-rs/state/src/runtime/background_agents/tests.rs @@ -932,7 +932,12 @@ async fn stale_supervisor_lease_is_orphaned_and_reclaimable() -> anyhow::Result< .into_iter() .map(|event| event.event_type) .collect::>(), - vec!["agent.orphaned".to_string()] + vec![ + "agent.claimed".to_string(), + "agent.heartbeat".to_string(), + "agent.orphaned".to_string(), + "agent.claimed".to_string() + ] ); Ok(()) } @@ -997,7 +1002,7 @@ async fn orphaning_waiting_run_terminalizes_pending_interactions() -> anyhow::Re .expect("status snapshot should exist"); assert_eq!(status_snapshot.status, BackgroundAgentRunStatus::Orphaned); assert_eq!(status_snapshot.pending_interaction_count, 0); - assert_eq!(status_snapshot.last_event_seq, 3); + assert_eq!(status_snapshot.last_event_seq, 5); assert_eq!( runtime .list_background_agent_events_after( @@ -1008,6 +1013,8 @@ async fn orphaning_waiting_run_terminalizes_pending_interactions() -> anyhow::Re .map(|event| event.event_type) .collect::>(), vec![ + "agent.claimed".to_string(), + "agent.heartbeat".to_string(), "interaction.created".to_string(), "interaction.workerNoLongerWaiting".to_string(), "agent.orphaned".to_string() @@ -1251,7 +1258,7 @@ async fn stale_stopping_run_is_cancelled_and_lease_stopped() -> anyhow::Result<( status_snapshot.summary.as_deref(), Some("stop heartbeat stale") ); - assert_eq!(status_snapshot.last_event_seq, 3); + assert_eq!(status_snapshot.last_event_seq, 5); let interaction = runtime .get_background_agent_pending_interaction("pending-1") .await? @@ -1604,7 +1611,7 @@ async fn delete_request_for_claimed_run_becomes_stopping_and_stale_cancelled() - status_snapshot.summary.as_deref(), Some("stop heartbeat stale") ); - assert_eq!(status_snapshot.last_event_seq, 1); + assert_eq!(status_snapshot.last_event_seq, 3); Ok(()) } From b6795dcea11e5e913812683dc03140a5e5be7ffc Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 18:46:15 +0300 Subject: [PATCH 06/15] style(state): label receipt arguments --- codex-rs/state/src/runtime/background_agents/runs.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/codex-rs/state/src/runtime/background_agents/runs.rs b/codex-rs/state/src/runtime/background_agents/runs.rs index 67231a47c1..022471a885 100644 --- a/codex-rs/state/src/runtime/background_agents/runs.rs +++ b/codex-rs/state/src/runtime/background_agents/runs.rs @@ -160,7 +160,7 @@ INSERT INTO background_agent_runs ( params.id.as_str(), "agent.admitted", receipt_key.as_str(), - 0, + /*generation*/ 0, Some(1), &serde_json::json!({ "source": params.source, @@ -475,7 +475,7 @@ WHERE params.event_type, receipt_key.as_str(), params.generation, - None, + /*attempt*/ None, params.event_payload_json, now, ) @@ -605,7 +605,7 @@ WHERE "agent.stopRequested", receipt_key.as_str(), expected_generation, - None, + /*attempt*/ None, diagnostics_json, now, ) @@ -1295,7 +1295,7 @@ WHERE run_id = ? AND supervisor_id = ? AND generation = ? "agent.heartbeat", receipt_key.as_str(), params.generation, - None, + /*attempt*/ None, &serde_json::json!({ "supervisorId": params.supervisor_id, }), @@ -1612,7 +1612,7 @@ async fn append_terminal_stale_background_agent_status_in_tx( event_type, receipt_key.as_str(), generation, - None, + /*attempt*/ None, event_payload_json, now, ) From 09bf176508697e24f4891132d4b7896de0d9032a Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 18:48:50 +0300 Subject: [PATCH 07/15] style(state): scope receipt helper lint --- codex-rs/state/src/runtime/background_agents/events.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/codex-rs/state/src/runtime/background_agents/events.rs b/codex-rs/state/src/runtime/background_agents/events.rs index 199b7ec20b..484c70e2da 100644 --- a/codex-rs/state/src/runtime/background_agents/events.rs +++ b/codex-rs/state/src/runtime/background_agents/events.rs @@ -58,6 +58,7 @@ WHERE id = ? }) } +#[allow(clippy::too_many_arguments)] pub(in crate::runtime) async fn append_background_agent_lifecycle_receipt_in_tx( tx: &mut sqlx::Transaction<'_, Sqlite>, run_id: &str, From ff3ab38c0e10ec37e0be1eddee00a6e56f2e9574 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 19:24:15 +0300 Subject: [PATCH 08/15] test(background-agent): cover receipt-expanded events --- .../tests/suite/v2/background_agent.rs | 25 +++++++++++++------ 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/background_agent.rs b/codex-rs/app-server/tests/suite/v2/background_agent.rs index 82b6834898..f4693da1b2 100644 --- a/codex-rs/app-server/tests/suite/v2/background_agent.rs +++ b/codex-rs/app-server/tests/suite/v2/background_agent.rs @@ -63,6 +63,7 @@ use codex_app_server_protocol::WorktreeMergeCandidateStatus; use codex_app_server_protocol::WorktreeOwnerKind; use codex_app_server_protocol::WorktreeReadResponse; use codex_app_server_protocol::WorktreeReleaseResponse; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED; use codex_background_agent::BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH; use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH; use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; @@ -216,7 +217,7 @@ async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result assert_eq!(second_events_page.data[0].event_type, "agent.started"); assert_eq!(second_events_page.next_cursor, Some("event:2".to_string())); let all_events = - agent_events_page(&mut restarted, &agent_id, /*cursor*/ None, Some(20)).await?; + agent_events_page(&mut restarted, &agent_id, /*cursor*/ None, Some(200)).await?; let event_types = all_events .data .iter() @@ -827,7 +828,7 @@ async fn supervisor_periodically_starts_durable_queued_runs() -> Result<()> { assert_eq!(agent.agent_id, agent_id); assert!(agent.thread_id.is_some()); - let events = agent_events_page(&mut mcp, agent_id.as_str(), /*cursor*/ None, Some(20)).await?; + let events = agent_events_page(&mut mcp, agent_id.as_str(), /*cursor*/ None, Some(200)).await?; let event_types = events .data .iter() @@ -1264,7 +1265,7 @@ async fn worker_fails_when_turn_completes_without_agent_message() -> Result<()> } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn agent_diagnostics_reports_quota_and_overloaded_admission() -> Result<()> { +async fn agent_diagnostics_reports_quota_and_capacity_exhaustion() -> Result<()> { let codex_home = TempDir::new()?; let server = create_mock_responses_server_sequence(vec![create_final_assistant_message_sse_response( @@ -1315,13 +1316,21 @@ async fn agent_diagnostics_reports_quota_and_overloaded_admission() -> Result<() ) .await?; assert_eq!(rejected.error.code, -32001); - assert!( + assert_eq!( rejected .error - .message - .contains("background agent queue is overloaded"), - "unexpected overloaded error: {}", - rejected.error.message + .data + .as_ref() + .and_then(|data| data.get("errorCode")), + Some(&json!(BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED)) + ); + assert_eq!( + rejected + .error + .data + .as_ref() + .and_then(|data| data.get("maxActiveRuns")), + Some(&json!(full.max_active_runs_per_user)) ); state_db From 07465e4edd6fa77e6a39d013f91bbe625228f557 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 19:39:30 +0300 Subject: [PATCH 09/15] fix(cli): remove stale agent state import --- codex-rs/cli/src/agent_cmd.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/codex-rs/cli/src/agent_cmd.rs b/codex-rs/cli/src/agent_cmd.rs index fb54fefb7b..0cb661fdd6 100644 --- a/codex-rs/cli/src/agent_cmd.rs +++ b/codex-rs/cli/src/agent_cmd.rs @@ -16,7 +16,6 @@ use codex_background_agent::daemon::background_agent_daemon_state_dir; use codex_background_agent::daemon::ensure_supported_platform as ensure_background_agent_supported_platform; use codex_core::config::find_codex_home; use codex_protocol::models::PermissionProfile; -use codex_state::BackgroundAgentDesiredState; use codex_state::BackgroundAgentExecutionSnapshotParams; use codex_state::BackgroundAgentPendingInteractionStatus; use codex_state::BackgroundAgentRun; From 3ab85946c9a9764aefb0d4b0c1b946864e07e2c1 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 19:45:58 +0300 Subject: [PATCH 10/15] test(background-agent): replay quota admission identity --- .../tests/suite/v2/background_agent.rs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/background_agent.rs b/codex-rs/app-server/tests/suite/v2/background_agent.rs index f4693da1b2..a458f4aaba 100644 --- a/codex-rs/app-server/tests/suite/v2/background_agent.rs +++ b/codex-rs/app-server/tests/suite/v2/background_agent.rs @@ -1367,15 +1367,16 @@ async fn agent_diagnostics_reports_quota_and_capacity_exhaustion() -> Result<()> ) .await?; - let retry = start_agent( - &mut mcp, - start_params( - "idempotent retry is not new pressure", - Some("quota-idempotency-0".to_string()), - codex_home.path(), - ), - ) - .await?; + let mut retry_params = start_params( + "quota run 0", + Some("quota-idempotency-0".to_string()), + codex_home.path(), + ); + retry_params.source = Some("quota-test".to_string()); + retry_params.prompt_snapshot_ref = Some("inline:quota-run-0:prompt".to_string()); + retry_params.cwd = None; + retry_params.execution_context = None; + let retry = start_agent(&mut mcp, retry_params).await?; assert_eq!(retry.agent.agent_id, first_agent_id); Ok(()) From e553bd36ef78cd76de476b4ace26fa7c0dc7ada3 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 21:10:02 +0300 Subject: [PATCH 11/15] fix(background-agent): harden durable admission --- .../background_agent_live.rs | 396 +++--- .../background_agent_processor.rs | 312 ++--- .../tests/suite/v2/background_agent.rs | 133 +- codex-rs/background-agent/src/daemon.rs | 14 +- codex-rs/background-agent/src/lib.rs | 39 +- codex-rs/background-agent/src/supervisor.rs | 159 ++- codex-rs/cli/src/agent_cmd.rs | 220 +-- ...61_background_agent_lifecycle_receipts.sql | 47 + .../src/runtime/background_agents/events.rs | 213 ++- .../src/runtime/background_agents/runs.rs | 1210 ++++++++++++++--- .../src/runtime/background_agents/tests.rs | 478 ++++++- 11 files changed, 2293 insertions(+), 928 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/background_agent_live.rs b/codex-rs/app-server/src/request_processors/background_agent_live.rs index 22b809e869..dd99a5e3bd 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_live.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_live.rs @@ -333,7 +333,7 @@ impl ThreadRequestProcessor { &self, mut params: AgentStartParams, ) -> Result, JSONRPCErrorError> { - self.validate_agent_start_admission(¶ms)?; + self.validate_agent_start_admission(&mut params)?; let managed_worktree = self .trusted_agent_start_managed_worktree( params.cwd.as_deref(), @@ -1673,9 +1673,13 @@ impl ThreadRequestProcessor { fn validate_agent_start_admission( &self, - params: &AgentStartParams, + params: &mut AgentStartParams, ) -> Result<(), JSONRPCErrorError> { - if params.version_fingerprint.as_deref() != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) + if params.version_fingerprint.is_none() { + params.version_fingerprint = + Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()); + } else if params.version_fingerprint.as_deref() + != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) { let mut error = invalid_request("background agent admission schema is incompatible"); error.data = Some(json!({ @@ -1693,8 +1697,6 @@ impl ThreadRequestProcessor { ); error.data = Some(json!({ "errorCode": BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH, - "requestedProfile": requested_profile, - "selectedProfile": self.config.selected_auth_profile.as_deref(), })); return Err(error); } @@ -1868,6 +1870,16 @@ async fn reconcile_background_agents( if !should_start_background_run(&run) { continue; } + if !context + .state_db + .background_agent_admission_is_ready( + run.id.as_str(), + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + ) + .await? + { + continue; + } let token = CancellationToken::new(); { let mut active = context.active_workers.lock().await; @@ -1953,6 +1965,16 @@ async fn reconcile_background_agent_worker_processes( if !should_start_background_run(&run) { continue; } + if !context + .state_db + .background_agent_admission_is_ready( + run.id.as_str(), + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + ) + .await? + { + continue; + } if context .active_worker_processes .lock() @@ -2841,6 +2863,7 @@ async fn worker_process_start_token(_worker_process: bool) -> anyhow::Result bool { if run.desired_state != BackgroundAgentDesiredState::Running + || run.version_fingerprint.as_deref() != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) || !matches!( run.status, BackgroundAgentRunStatus::Queued | BackgroundAgentRunStatus::Orphaned @@ -2884,11 +2907,14 @@ async fn run_background_agent_worker( run.generation.saturating_add(1) ); let Some(generation) = retry_transient_sqlite_busy("claim background agent supervisor", || { - context.state_db.claim_background_agent_supervisor( - run.id.as_str(), - context.supervisor_id.as_str(), - process_lease_id.as_str(), - ) + context + .state_db + .claim_background_agent_supervisor_compatible( + run.id.as_str(), + context.supervisor_id.as_str(), + process_lease_id.as_str(), + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + ) }) .await? else { @@ -2896,23 +2922,11 @@ async fn run_background_agent_worker( return Ok(()); }; - let (config, initial_execution_payload) = - match resolve_background_agent_config(&context, &run).await? { - BackgroundAgentConfigResolution::Ready { - config, - initial_execution_payload, - } => (*config, initial_execution_payload), - BackgroundAgentConfigResolution::UsageProfileWait { retry_at } => { - defer_background_agent_for_usage_profile_wait( - &context, - run.id.as_str(), - generation, - retry_at, - ) - .await?; - return Ok(()); - } - }; + let BackgroundAgentConfigResolution::Ready { + config, + initial_execution_payload, + } = resolve_background_agent_config(&context, &run).await?; + let config = *config; let pid_value = i64::from(std::process::id()); let pid = Some(pid_value); @@ -3023,6 +3037,8 @@ async fn run_background_agent_worker( insert_initial_goal_for_background_thread( &context.state_db, run.id.as_str(), + context.supervisor_id.as_str(), + generation, thread_id, initial_execution_payload.as_ref(), ) @@ -3030,19 +3046,24 @@ async fn run_background_agent_worker( retry_transient_sqlite_busy("create background agent execution snapshot", || { context .state_db - .create_execution_snapshot(BackgroundAgentExecutionSnapshotParams { - run_id: run.id.clone(), - snapshot_kind: "worker_thread_bound".to_string(), - payload_json: json!({ - "threadId": thread_id.to_string(), - "sessionId": session_id_string, - "rolloutPath": rollout_path, - }), - recovery_policy: "resume_or_orphan".to_string(), - config_fingerprint: run.config_fingerprint.clone(), - }) + .create_background_agent_execution_snapshot_for_supervisor( + &BackgroundAgentExecutionSnapshotParams { + run_id: run.id.clone(), + snapshot_kind: "worker_thread_bound".to_string(), + payload_json: json!({ + "threadId": thread_id.to_string(), + "sessionId": session_id_string, + "rolloutPath": rollout_path, + }), + recovery_policy: "resume_or_orphan".to_string(), + config_fingerprint: run.config_fingerprint.clone(), + }, + context.supervisor_id.as_str(), + generation, + ) }) - .await?; + .await? + .ok_or_else(|| background_agent_ownership_lost(run.id.as_str(), generation))?; append_status( &context, run.id.as_str(), @@ -3157,9 +3178,6 @@ enum BackgroundAgentConfigResolution { config: Box, initial_execution_payload: Option, }, - UsageProfileWait { - retry_at: DateTime, - }, } async fn resolve_background_agent_config( @@ -3173,22 +3191,15 @@ async fn resolve_background_agent_config( } let snapshot = context .state_db - .get_latest_execution_snapshot(run.id.as_str()) - .await?; - let initial_execution_payload = snapshot - .as_ref() - .map(|snapshot| snapshot.payload_json.clone()); - let payload = snapshot - .as_ref() - .and_then(|snapshot| snapshot.payload_json.as_object()); - let snapshot_auth_profile_ref = payload - .and_then(|payload| payload.get("authProfileRef")) - .and_then(Value::as_str); - if snapshot_auth_profile_ref != run.auth_profile_ref.as_deref() { - anyhow::bail!( - "{BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH}: persisted execution snapshot auth profile does not match the admitted run" - ); - } + .get_background_agent_initial_execution_snapshot(run.id.as_str()) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "background agent admission is missing its initial execution context snapshot" + ) + })?; + let initial_execution_payload = Some(snapshot.payload_json.clone()); + let payload = snapshot.payload_json.as_object(); let cwd = payload .and_then(|payload| payload.get("cwd")) .and_then(Value::as_str) @@ -3278,7 +3289,7 @@ async fn resolve_background_agent_config( } let request_overrides = (!request_overrides.is_empty()).then_some(request_overrides); - let mut config_overrides = ConfigOverrides { + let config_overrides = ConfigOverrides { model, model_provider, service_tier, @@ -3293,9 +3304,9 @@ async fn resolve_background_agent_config( main_execve_wrapper_exe: context.arg0_paths.main_execve_wrapper_exe.clone(), ..Default::default() }; - let config = context + let mut config = context .config_manager - .load_with_overrides(request_overrides.clone(), config_overrides.clone()) + .load_with_overrides(request_overrides, config_overrides) .await .map_err(anyhow::Error::from)?; if config.selected_auth_profile != run.auth_profile_ref { @@ -3303,49 +3314,8 @@ async fn resolve_background_agent_config( "{BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH}: loaded worker auth profile does not match the admitted run" ); } - if run.auth_profile_ref.is_some() { - return Ok(BackgroundAgentConfigResolution::Ready { - config: Box::new(config), - initial_execution_payload, - }); - } - - let broker_decision = super::usage_profile_broker::resolve_dispatch_auth_profile( - &context.auth_manager, - &config, - config_overrides.auth_profile.clone(), - ) - .await; - if let Some(profile) = broker_decision.selected_profile.as_ref() { - tracing::debug!( - run_id = %run.id, - auth_profile = %profile, - reason = ?broker_decision.reason, - "usage profile broker selected auth profile for background agent" - ); - config_overrides.auth_profile = Some(Some(profile.clone())); - return context - .config_manager - .load_with_overrides(request_overrides, config_overrides) - .await - .map(Box::new) - .map(|config| BackgroundAgentConfigResolution::Ready { - config, - initial_execution_payload, - }) - .map_err(anyhow::Error::from); - } - if let Some(retry_at) = broker_decision.retry_at - && let Some(retry_at) = background_agent_broker_retry_at(&config, retry_at) - { - tracing::debug!( - run_id = %run.id, - retry_at = %retry_at.to_rfc3339(), - reason = ?broker_decision.reason, - "usage profile broker deferred background agent" - ); - return Ok(BackgroundAgentConfigResolution::UsageProfileWait { retry_at }); - } + config.auth_profile_auto_switch.enabled = false; + config.usage_self_heal.enabled = false; Ok(BackgroundAgentConfigResolution::Ready { config: Box::new(config), initial_execution_payload, @@ -3355,6 +3325,8 @@ async fn resolve_background_agent_config( async fn insert_initial_goal_for_background_thread( state_db: &codex_state::StateRuntime, run_id: &str, + supervisor_id: &str, + generation: i64, thread_id: codex_protocol::ThreadId, initial_execution_payload: Option<&Value>, ) -> anyhow::Result<()> { @@ -3381,9 +3353,17 @@ async fn insert_initial_goal_for_background_thread( "objective": goal.objective, }); retry_transient_sqlite_busy("append background agent initial goal event", || { - state_db.append_background_agent_event(run_id, "agent.initialGoalCreated", &event_payload) + state_db.append_background_agent_event_for_supervisor( + run_id, + supervisor_id, + generation, + "agent.initialGoalCreated", + &event_payload, + /*allow_terminal_current*/ false, + ) }) - .await?; + .await? + .ok_or_else(|| background_agent_ownership_lost(run_id, generation))?; Ok(()) } @@ -3396,49 +3376,6 @@ fn initial_goal_objective_from_execution_payload(payload: &Value) -> Option, -) -> anyhow::Result<()> { - let status_reason = background_agent_usage_profile_wait_reason(retry_at); - append_status( - context, - run_id, - generation, - BackgroundAgentRunStatus::Queued, - status_reason.as_str(), - "agent.usageProfileWait", - json!({ - "retryAt": retry_at.timestamp(), - }), - ) - .await?; - retry_transient_sqlite_busy("finish background agent usage profile wait lease", || { - context.state_db.finish_background_agent_process_lease( - run_id, - context.supervisor_id.as_str(), - generation, - /*exit_code*/ None, - /*exit_signal*/ None, - Some("usage profile wait"), - ) - }) - .await?; - Ok(()) -} - -fn background_agent_broker_retry_at( - config: &codex_core::config::Config, - retry_at: i64, -) -> 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) -} - fn background_agent_usage_profile_wait_reason(retry_at: DateTime) -> String { format!( "{BACKGROUND_AGENT_USAGE_PROFILE_WAIT_PREFIX}{}", @@ -3485,7 +3422,16 @@ async fn load_background_agent_prompt( .await?; events .into_iter() - .find(|event| event.event_type == "agent.started") + .find(|event| { + matches!( + event.event_type.as_str(), + "agent.started" | "agent.startRecovered" + ) && event + .payload_json + .get("prompt") + .and_then(Value::as_str) + .is_some() + }) .and_then(|event| { event .payload_json @@ -4320,33 +4266,38 @@ async fn upsert_interaction_status_snapshot( retry_transient_sqlite_busy("upsert background agent waiting status snapshot", || { context .state_db - .upsert_status_snapshot(BackgroundAgentStatusSnapshotParams { - run_id: run_id.to_string(), - seq: run.last_event_seq, - status, - desired_state: run.desired_state, - summary: Some( - status_summary(status, "waiting for pending interaction").to_string(), - ), - pending_interaction_count: pending_count, - last_event_seq: run.last_event_seq, - payload_json: status_snapshot_payload( + .upsert_background_agent_status_snapshot_for_supervisor( + &BackgroundAgentStatusSnapshotParams { + run_id: run_id.to_string(), + seq: run.last_event_seq, status, - waiting_reason, - generation, - &json!({ - "interactionId": interaction.id, - "workerRequestId": interaction.worker_request_id, - "kind": interaction.kind.as_str(), - "waitingReason": waiting_reason, - "requestPayload": interaction.request_payload_json, - "timeoutAt": interaction.timeout_at.map(|value| value.timestamp()), - }), - execution_payload.as_ref(), - ), - }) + desired_state: run.desired_state, + summary: Some( + status_summary(status, "waiting for pending interaction").to_string(), + ), + pending_interaction_count: pending_count, + last_event_seq: run.last_event_seq, + payload_json: status_snapshot_payload( + status, + waiting_reason, + generation, + &json!({ + "interactionId": interaction.id, + "workerRequestId": interaction.worker_request_id, + "kind": interaction.kind.as_str(), + "waitingReason": waiting_reason, + "requestPayload": interaction.request_payload_json, + "timeoutAt": interaction.timeout_at.map(|value| value.timestamp()), + }), + execution_payload.as_ref(), + ), + }, + context.supervisor_id.as_str(), + generation, + ) }) - .await?; + .await? + .ok_or_else(|| background_agent_ownership_lost(run_id, generation))?; Ok(()) } @@ -4371,14 +4322,20 @@ async fn append_background_agent_event_with_retry( payload_json: &Value, allow_terminal_current: bool, ) -> anyhow::Result { - ensure_background_agent_worker_current(context, run_id, generation, allow_terminal_current) - .await?; - retry_transient_sqlite_busy("append background agent event", || { + let event = retry_transient_sqlite_busy("append background agent event", || { context .state_db - .append_event(run_id, event_type, payload_json) + .append_background_agent_event_for_supervisor( + run_id, + context.supervisor_id.as_str(), + generation, + event_type, + payload_json, + allow_terminal_current, + ) }) - .await + .await?; + event.ok_or_else(|| background_agent_ownership_lost(run_id, generation)) } async fn ensure_background_agent_worker_current( @@ -5466,6 +5423,10 @@ done codex_state::StateRuntime::init(temp.path().to_path_buf(), "test-provider".to_string()) .await?; seed_queued_run(state_db.as_ref(), "goal-run").await?; + let generation = state_db + .claim_background_agent_supervisor("goal-run", "goal-supervisor", "goal-process-lease") + .await? + .expect("seeded run should be claimable"); let thread_id = codex_protocol::ThreadId::new(); let initial_execution_payload = json!({ "initialGoalObjective": " Investigate flaky test ", @@ -5474,6 +5435,8 @@ done insert_initial_goal_for_background_thread( state_db.as_ref(), "goal-run", + "goal-supervisor", + generation, thread_id, Some(&initial_execution_payload), ) @@ -5514,6 +5477,8 @@ done insert_initial_goal_for_background_thread( state_db.as_ref(), "goal-run", + "goal-supervisor", + generation, thread_id, Some(&initial_execution_payload), ) @@ -5953,36 +5918,49 @@ done state_db: &codex_state::StateRuntime, run_id: &str, ) -> anyhow::Result<()> { + let start_event_payload = json!({ + "cwd": null, + "prompt": "process supervisor test", + "promptSnapshotRef": format!("inline:{run_id}:prompt"), + }); + let execution_snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: run_id.to_string(), + snapshot_kind: "initial_execution_context".to_string(), + payload_json: json!({ + "cwd": null, + "configFingerprint": "cfg-test", + "versionFingerprint": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", + }), + recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), + config_fingerprint: Some("cfg-test".to_string()), + }; state_db - .create_background_agent_run(&codex_state::BackgroundAgentRunCreateParams { - id: run_id.to_string(), - idempotency_key: None, - request_id: None, - source: "process-supervisor-test".to_string(), - prompt_snapshot_ref: format!("inline:{run_id}:prompt"), - input_snapshot_ref: None, - thread_id: None, - thread_store_kind: "background-agent".to_string(), - thread_store_id: None, - rollout_path: None, - parent_thread_id: None, - parent_agent_run_id: None, - spawn_linkage_json: None, - auth_profile_ref: None, - status_reason: Some("queued by process supervisor test".to_string()), - config_fingerprint: Some("cfg-test".to_string()), - version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), - }) - .await?; - state_db - .append_background_agent_event( - run_id, - "agent.started", - &json!({ - "cwd": null, - "prompt": "process supervisor test", - "promptSnapshotRef": format!("inline:{run_id}:prompt"), - }), + .admit_background_agent_run( + &codex_state::BackgroundAgentRunCreateParams { + id: run_id.to_string(), + idempotency_key: None, + request_id: None, + source: "process-supervisor-test".to_string(), + prompt_snapshot_ref: format!("inline:{run_id}:prompt"), + input_snapshot_ref: None, + thread_id: None, + thread_store_kind: "background-agent".to_string(), + thread_store_id: None, + rollout_path: None, + parent_thread_id: None, + parent_agent_run_id: None, + spawn_linkage_json: None, + auth_profile_ref: None, + status_reason: Some("queued by process supervisor test".to_string()), + config_fingerprint: Some("cfg-test".to_string()), + version_fingerprint: Some( + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string(), + ), + }, + &start_event_payload, + &execution_snapshot_params, + DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, ) .await?; Ok(()) diff --git a/codex-rs/app-server/src/request_processors/background_agent_processor.rs b/codex-rs/app-server/src/request_processors/background_agent_processor.rs index cdc7412f49..536d1e0c28 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_processor.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_processor.rs @@ -94,6 +94,8 @@ use codex_state::ManagedWorktreeAttachParams; use codex_state::ManagedWorktreeDetachParams; use serde_json::Value; use serde_json::json; +use sha2::Digest; +use sha2::Sha256; use std::path::Path; use std::time::Duration; use uuid::Uuid; @@ -116,10 +118,10 @@ impl BackgroundAgentRequestProcessor { pub(super) async fn agent_start_inner( &self, - params: AgentStartParams, + mut params: AgentStartParams, ) -> Result { let state_db = self.state_db()?; - validate_agent_start_schema(¶ms)?; + normalize_agent_start_schema(&mut params)?; let AgentStartParams { prompt, initial_goal_objective, @@ -154,8 +156,13 @@ impl BackgroundAgentRequestProcessor { )) })?; let agent_id = Uuid::now_v7().to_string(); - let prompt_snapshot_ref = - prompt_snapshot_ref.unwrap_or_else(|| format!("inline:{agent_id}:prompt")); + let prompt_snapshot_ref = prompt_snapshot_ref.unwrap_or_else(|| { + let identity = idempotency_key + .as_deref() + .map(|key| format!("{:x}", Sha256::digest(key.as_bytes()))) + .unwrap_or_else(|| agent_id.clone()); + format!("inline:{identity}:prompt") + }); let source = source.unwrap_or_else(|| "app-server".to_string()); let thread_store_kind = thread_store_kind.unwrap_or_else(|| "background-agent".to_string()); validate_agent_start_rollout_path( @@ -168,6 +175,31 @@ impl BackgroundAgentRequestProcessor { .as_ref() .and_then(|context| context.recovery_policy.clone()) .unwrap_or_else(|| "abort_mid_turn_resume_at_safe_boundary".to_string()); + let config_fingerprint = match config_fingerprint { + Some(config_fingerprint) => Some(config_fingerprint), + None => Some( + background_agent_config_fingerprint( + cwd.as_deref(), + initial_goal_objective.as_deref(), + auth_profile_ref.as_deref(), + execution_context.as_ref(), + ) + .map_err(|err| { + internal_error(format!( + "failed to fingerprint background agent configuration: {err}" + )) + })?, + ), + }; + let execution_payload = + initial_execution_snapshot_payload(InitialExecutionSnapshotPayloadParams { + cwd: cwd.as_deref(), + initial_goal_objective: initial_goal_objective.as_deref(), + execution_context: execution_context.as_ref(), + recovery_policy: recovery_policy.as_str(), + config_fingerprint: config_fingerprint.as_deref(), + version_fingerprint: version_fingerprint.as_deref(), + }); let create_params = BackgroundAgentRunCreateParams { id: agent_id.clone(), idempotency_key, @@ -184,138 +216,35 @@ impl BackgroundAgentRequestProcessor { spawn_linkage_json: spawn_linkage, auth_profile_ref, status_reason: Some("queued for background-agent supervisor".to_string()), - config_fingerprint, + config_fingerprint: config_fingerprint.clone(), version_fingerprint, }; - let (run, created_new_run) = retry_transient_sqlite_busy("admit background agent", || { - state_db.admit_run( - create_params.clone(), - DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, - ) - }) - .await - .map_err(map_background_agent_admission_error)?; - let execution_payload = initial_execution_snapshot_payload( - &run, - InitialExecutionSnapshotPayloadParams { - cwd: cwd.as_deref(), - initial_goal_objective: initial_goal_objective.as_deref(), - execution_context: execution_context.as_ref(), - recovery_policy: recovery_policy.as_str(), - }, - ); - let execution_snapshot = if created_new_run { - state_db - .create_execution_snapshot(BackgroundAgentExecutionSnapshotParams { - run_id: run.id.clone(), - snapshot_kind: "initial_execution_context".to_string(), - payload_json: execution_payload, - recovery_policy: recovery_policy.clone(), - config_fingerprint: run.config_fingerprint.clone(), - }) - .await - .map_err(|err| { - internal_error(format!( - "failed to create background agent execution snapshot: {err}" - )) - })? - } else { - match state_db - .get_latest_execution_snapshot(run.id.as_str()) - .await - .map_err(|err| { - internal_error(format!( - "failed to load background agent execution snapshot: {err}" - )) - })? { - Some(snapshot) => snapshot, - None => state_db - .create_execution_snapshot(BackgroundAgentExecutionSnapshotParams { - run_id: run.id.clone(), - snapshot_kind: "initial_execution_context".to_string(), - payload_json: execution_payload, - recovery_policy: recovery_policy.clone(), - config_fingerprint: run.config_fingerprint.clone(), - }) - .await - .map_err(|err| { - internal_error(format!( - "failed to create background agent execution snapshot: {err}" - )) - })?, - } + let prompt_sha256 = format!("{:x}", Sha256::digest(prompt.as_bytes())); + let start_event_payload = json!({ + "cwd": cwd, + "prompt": prompt, + "promptSha256": prompt_sha256, + "promptSnapshotRef": create_params.prompt_snapshot_ref.as_str(), + "initialGoalObjective": initial_goal_objective.as_deref(), + }); + let execution_snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: agent_id, + snapshot_kind: "initial_execution_context".to_string(), + payload_json: execution_payload, + recovery_policy, + config_fingerprint, }; - let event = if created_new_run { - append_background_agent_event_with_retry( - state_db.as_ref(), - run.id.as_str(), - "agent.started", - &json!({ - "cwd": cwd, - "prompt": prompt, - "promptSnapshotRef": run.prompt_snapshot_ref.as_str(), - "initialGoalObjective": initial_goal_objective.as_deref(), - }), - ) - .await - .map_err(|err| { - internal_error(format!("failed to append background agent event: {err}")) - })? - } else { - let events = state_db - .list_events_after(run.id.as_str(), /*after_seq*/ None, Some(100)) - .await - .map_err(|err| { - internal_error(format!("failed to list background agent events: {err}")) - })?; - match events - .into_iter() - .find(|event| event.event_type == "agent.started") - { - Some(event) => event, - None => append_background_agent_event_with_retry( - state_db.as_ref(), - run.id.as_str(), - "agent.startRecovered", - &json!({ - "reason": "idempotent_start_without_start_event", - }), + let (run, _created_new_run, event, execution_snapshot, snapshot) = + retry_transient_sqlite_busy("admit background agent", || { + state_db.admit_background_agent_run( + &create_params, + &start_event_payload, + &execution_snapshot_params, + DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, ) - .await - .map_err(|err| { - internal_error(format!("failed to append background agent event: {err}")) - })?, - } - }; - let snapshot = match state_db - .get_status_snapshot(run.id.as_str()) + }) .await - .map_err(|err| { - internal_error(format!("failed to load background agent snapshot: {err}")) - })? { - Some(snapshot) => snapshot, - None => state_db - .upsert_status_snapshot(BackgroundAgentStatusSnapshotParams { - run_id: run.id.clone(), - seq: event.seq, - status: run.status, - desired_state: run.desired_state, - summary: Some("Queued".to_string()), - pending_interaction_count: 0, - last_event_seq: event.seq, - payload_json: json!({ - "phase": "queued", - }), - }) - .await - .map_err(|err| { - internal_error(format!("failed to update background agent snapshot: {err}")) - })?, - }; - let run = self - .load_agent_run(state_db.as_ref(), run.id.as_str()) - .await? - .ok_or_else(|| internal_error("background agent disappeared after create"))?; + .map_err(map_background_agent_admission_error)?; Ok(AgentStartResponse { agent: api_agent_run_from_state(run), @@ -631,80 +560,12 @@ impl BackgroundAgentRequestProcessor { params: AgentDeleteParams, ) -> Result { let state_db = self.state_db()?; - let existing_run = self - .load_agent_run(state_db.as_ref(), params.agent_id.as_str()) - .await?; let deleted = state_db .request_delete_run(params.agent_id.as_str()) .await .map_err(|err| { internal_error(format!("failed to request background agent delete: {err}")) })?; - if deleted { - let non_terminal_existing_run = existing_run - .as_ref() - .filter(|run| !is_terminal_agent_status(run.status)); - let terminalized_immediately = - non_terminal_existing_run.is_some_and(should_terminalize_unclaimed_agent_run); - if terminalized_immediately { - state_db - .update_run_status( - params.agent_id.as_str(), - BackgroundAgentRunStatus::Cancelled, - Some("delete requested before worker claim"), - ) - .await - .map_err(|err| { - internal_error(format!("failed to update background agent status: {err}")) - })?; - } - append_background_agent_event_with_retry( - state_db.as_ref(), - params.agent_id.as_str(), - "agent.deleteRequested", - &json!({ - "reason": "client_requested_delete", - }), - ) - .await - .map_err(|err| { - internal_error(format!("failed to append background agent event: {err}")) - })?; - if non_terminal_existing_run.is_some() { - cancel_active_pending_interactions_for_run( - state_db.as_ref(), - params.agent_id.as_str(), - "client_requested_delete", - ) - .await?; - let status = if terminalized_immediately { - BackgroundAgentRunStatus::Cancelled - } else { - BackgroundAgentRunStatus::Stopping - }; - upsert_lifecycle_status_snapshot( - state_db.as_ref(), - params.agent_id.as_str(), - status, - if terminalized_immediately { - "Deleted" - } else { - "Deleting" - }, - "client_requested_delete", - ) - .await?; - } else if let Some(existing_run) = existing_run.as_ref() { - upsert_lifecycle_status_snapshot( - state_db.as_ref(), - params.agent_id.as_str(), - existing_run.status, - "Deleted", - "client_requested_delete", - ) - .await?; - } - } let run = self .load_agent_run(state_db.as_ref(), params.agent_id.as_str()) .await? @@ -1178,10 +1039,11 @@ struct InitialExecutionSnapshotPayloadParams<'a> { initial_goal_objective: Option<&'a str>, execution_context: Option<&'a AgentExecutionContextParams>, recovery_policy: &'a str, + config_fingerprint: Option<&'a str>, + version_fingerprint: Option<&'a str>, } fn initial_execution_snapshot_payload( - run: &BackgroundAgentRun, params: InitialExecutionSnapshotPayloadParams<'_>, ) -> serde_json::Value { json!({ @@ -1194,7 +1056,6 @@ fn initial_execution_snapshot_payload( "approvalPolicy": params .execution_context .and_then(|context| context.approval_policy), - "authProfileRef": run.auth_profile_ref.as_deref(), "permissionProfile": params .execution_context .and_then(|context| context.permission_profile.as_ref()), @@ -1232,13 +1093,36 @@ fn initial_execution_snapshot_payload( "maxTokens": params .execution_context .and_then(|context| context.max_tokens), - "configFingerprint": run.config_fingerprint.as_deref(), - "versionFingerprint": run.version_fingerprint.as_deref(), + "configFingerprint": params.config_fingerprint, + "versionFingerprint": params.version_fingerprint, + "packageFingerprint": format!( + "{}:{}", + env!("CARGO_PKG_NAME"), + env!("CARGO_PKG_VERSION") + ), "recoveryPolicy": params.recovery_policy, "midTurnCrashSemantics": "abort_mid_turn_resume_at_safe_boundary", }) } +fn background_agent_config_fingerprint( + cwd: Option<&str>, + initial_goal_objective: Option<&str>, + auth_profile_ref: Option<&str>, + execution_context: Option<&AgentExecutionContextParams>, +) -> anyhow::Result { + let config_identity = json!({ + "cwd": cwd, + "initialGoalObjective": initial_goal_objective, + "authProfileRef": auth_profile_ref, + "executionContext": execution_context, + }); + Ok(format!( + "{:x}", + Sha256::digest(serde_json::to_vec(&config_identity)?) + )) +} + fn validate_agent_prompt(prompt: String) -> Result { let actual_chars = prompt.chars().count(); if actual_chars > MAX_USER_INPUT_TEXT_CHARS { @@ -1258,17 +1142,21 @@ fn validate_agent_prompt(prompt: String) -> Result { Ok(prompt) } -fn validate_agent_start_schema(params: &AgentStartParams) -> Result<(), JSONRPCErrorError> { +fn normalize_agent_start_schema(params: &mut AgentStartParams) -> Result<(), JSONRPCErrorError> { + if params.version_fingerprint.is_none() { + params.version_fingerprint = Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()); + } if params.version_fingerprint.as_deref() == Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) { - return Ok(()); + Ok(()) + } else { + let mut error = invalid_request("background agent admission schema is incompatible"); + error.data = Some(json!({ + "errorCode": BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH, + "requestedSchema": params.version_fingerprint.as_deref(), + "supportedSchema": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + })); + Err(error) } - let mut error = invalid_request("background agent admission schema is incompatible"); - error.data = Some(json!({ - "errorCode": BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH, - "requestedSchema": params.version_fingerprint.as_deref(), - "supportedSchema": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, - })); - Err(error) } fn map_background_agent_admission_error(err: anyhow::Error) -> JSONRPCErrorError { diff --git a/codex-rs/app-server/tests/suite/v2/background_agent.rs b/codex-rs/app-server/tests/suite/v2/background_agent.rs index a458f4aaba..3ca45b4228 100644 --- a/codex-rs/app-server/tests/suite/v2/background_agent.rs +++ b/codex-rs/app-server/tests/suite/v2/background_agent.rs @@ -82,6 +82,8 @@ use pretty_assertions::assert_eq; use serde::de::DeserializeOwned; use serde_json::Value as JsonValue; use serde_json::json; +use sha2::Digest; +use sha2::Sha256; use std::path::Path; use std::process::Command; use std::sync::Arc; @@ -251,7 +253,7 @@ async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn agent_start_records_initial_goal_objective() -> Result<()> { +async fn agent_start_binds_prompt_and_initial_goal_to_idempotency_identity() -> Result<()> { let codex_home = TempDir::new()?; let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; write_config(codex_home.path(), server.uri().as_str())?; @@ -262,6 +264,7 @@ async fn agent_start_records_initial_goal_objective() -> Result<()> { codex_home.path(), ); params.initial_goal_objective = Some("Investigate flaky regression".to_string()); + let exact_retry_params = params.clone(); let mut mcp = init_mcp(codex_home.path()).await?; let start = start_agent(&mut mcp, params).await?; @@ -275,7 +278,7 @@ async fn agent_start_records_initial_goal_objective() -> Result<()> { Some(&json!("Investigate flaky regression")) ); - let retry = start_agent( + let conflicting_retry = start_agent_error( &mut mcp, start_params( "retry with different params", @@ -284,8 +287,18 @@ async fn agent_start_records_initial_goal_objective() -> Result<()> { ), ) .await?; + assert_eq!( + conflicting_retry + .error + .data + .as_ref() + .and_then(|data| data.get("errorCode")), + Some(&json!("background_agent_admission_identity_mismatch")) + ); + let retry = start_agent(&mut mcp, exact_retry_params).await?; assert_eq!(retry.agent.agent_id, start.agent.agent_id); + assert_eq!(retry.event, start.event); assert_eq!( retry.execution_snapshot.payload.get("initialGoalObjective"), Some(&json!("Investigate flaky regression")) @@ -413,10 +426,7 @@ async fn agent_start_freezes_authority_from_server_config() -> Result<()> { start.execution_snapshot.payload.get("workspaceRoots"), Some(&json!(["/tmp/client-root"])) ); - assert_eq!( - start.execution_snapshot.payload.get("authProfileRef"), - Some(&JsonValue::Null) - ); + assert_eq!(start.execution_snapshot.payload.get("authProfileRef"), None); assert_eq!( start.execution_snapshot.payload.get("approvalPolicy"), Some(&json!("never")) @@ -479,6 +489,24 @@ async fn agent_start_rejects_profile_and_schema_mismatches_before_admission() -> .and_then(|data| data.get("errorCode")), Some(&json!(BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH)) ); + assert_eq!( + error.error.data, + Some(json!({ + "errorCode": BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH, + })) + ); + + let mut omitted_schema = start_params( + "accept omitted schema", + Some("omitted-schema".to_string()), + codex_home.path(), + ); + omitted_schema.version_fingerprint = None; + let compatible = start_agent(&mut mcp, omitted_schema).await?; + assert_eq!( + compatible.agent.version_fingerprint.as_deref(), + Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) + ); let mut schema_mismatch = start_params( "reject mismatched schema", @@ -3364,36 +3392,71 @@ async fn seed_queued_agent_run( idempotency_key: Option, prompt: &str, ) -> Result<()> { + let start_event_payload = json!({ + "cwd": null, + "prompt": prompt, + "promptSha256": format!("{:x}", Sha256::digest(prompt.as_bytes())), + "promptSnapshotRef": format!("inline:{agent_id}:prompt"), + "initialGoalObjective": null, + }); + let execution_snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: agent_id.to_string(), + snapshot_kind: "initial_execution_context".to_string(), + payload_json: json!({ + "snapshotSource": "agent/start", + "cwd": null, + "initialGoalObjective": null, + "workspaceRoots": null, + "approvalPolicy": null, + "permissionProfile": null, + "sandboxPolicy": null, + "networkPolicy": null, + "model": null, + "provider": null, + "serviceTier": null, + "mcpToolAllowlist": null, + "envSnapshotPolicy": "inherit-minimal", + "shellSnapshot": null, + "configSourceHashes": null, + "maxRuntimeSeconds": null, + "maxTokens": null, + "configFingerprint": "cfg-test", + "versionFingerprint": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + "packageFingerprint": format!( + "{}:{}", + env!("CARGO_PKG_NAME"), + env!("CARGO_PKG_VERSION") + ), + "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", + "midTurnCrashSemantics": "abort_mid_turn_resume_at_safe_boundary", + }), + recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), + config_fingerprint: Some("cfg-test".to_string()), + }; state_db - .create_background_agent_run(&BackgroundAgentRunCreateParams { - id: agent_id.to_string(), - idempotency_key, - request_id: None, - source: "quota-test".to_string(), - prompt_snapshot_ref: format!("inline:{agent_id}:prompt"), - input_snapshot_ref: None, - thread_id: None, - thread_store_kind: "background-agent".to_string(), - thread_store_id: None, - rollout_path: None, - parent_thread_id: None, - parent_agent_run_id: None, - spawn_linkage_json: None, - auth_profile_ref: None, - status_reason: Some("queued by quota test".to_string()), - config_fingerprint: Some("cfg-test".to_string()), - version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), - }) - .await?; - state_db - .append_background_agent_event( - agent_id, - "agent.started", - &json!({ - "cwd": null, - "prompt": prompt, - "promptSnapshotRef": format!("inline:{agent_id}:prompt"), - }), + .admit_background_agent_run( + &BackgroundAgentRunCreateParams { + id: agent_id.to_string(), + idempotency_key, + request_id: None, + source: "quota-test".to_string(), + prompt_snapshot_ref: format!("inline:{agent_id}:prompt"), + input_snapshot_ref: None, + thread_id: None, + thread_store_kind: "background-agent".to_string(), + thread_store_id: None, + rollout_path: None, + parent_thread_id: None, + parent_agent_run_id: None, + spawn_linkage_json: None, + auth_profile_ref: None, + status_reason: Some("queued by quota test".to_string()), + config_fingerprint: Some("cfg-test".to_string()), + version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), + }, + &start_event_payload, + &execution_snapshot_params, + /*max_active_runs*/ 1_000, ) .await?; Ok(()) diff --git a/codex-rs/background-agent/src/daemon.rs b/codex-rs/background-agent/src/daemon.rs index dc0ef7eaf2..d4b7694280 100644 --- a/codex-rs/background-agent/src/daemon.rs +++ b/codex-rs/background-agent/src/daemon.rs @@ -168,7 +168,10 @@ impl BackgroundAgentDaemon { .await; }; let status = match self.controller.status(&record.handle).await? { - WorkerProcessStatus::Running => BackgroundAgentDaemonStatus::Running, + WorkerProcessStatus::Running => { + ensure_daemon_record_compatible(&record)?; + BackgroundAgentDaemonStatus::Running + } WorkerProcessStatus::Missing | WorkerProcessStatus::StalePidRecord => { BackgroundAgentDaemonStatus::StalePidRecord } @@ -536,12 +539,21 @@ mod tests { .start() .await .expect_err("incompatible daemon must not be reused"); + let status_error = daemon + .status() + .await + .expect_err("incompatible daemon must not be reported as running"); assert!( error .to_string() .contains(BACKGROUND_AGENT_DAEMON_INCOMPATIBLE) ); + assert!( + status_error + .to_string() + .contains(BACKGROUND_AGENT_DAEMON_INCOMPATIBLE) + ); assert_eq!( controller.status(&existing).await?, WorkerProcessStatus::Running diff --git a/codex-rs/background-agent/src/lib.rs b/codex-rs/background-agent/src/lib.rs index 953dd08cdf..f18c2e1a2e 100644 --- a/codex-rs/background-agent/src/lib.rs +++ b/codex-rs/background-agent/src/lib.rs @@ -54,12 +54,6 @@ pub trait AgentRunStore { params: BackgroundAgentRunCreateParams, ) -> impl Future> + Send; - fn admit_run( - &self, - params: BackgroundAgentRunCreateParams, - max_active_runs: i64, - ) -> impl Future> + Send; - fn get_run( &self, run_id: &str, @@ -86,6 +80,11 @@ pub trait AgentRunStore { status_reason: Option<&str>, ) -> impl Future> + Send; + fn append_status_event_for_supervisor( + &self, + params: BackgroundAgentStatusEventForSupervisorParams<'_>, + ) -> impl Future>> + Send; + fn set_desired_state( &self, run_id: &str, @@ -114,6 +113,7 @@ pub trait AgentRunStore { run_id: &str, supervisor_id: &str, process_lease_id: &str, + required_version_fingerprint: &str, ) -> impl Future>> + Send; fn record_execution_handle( @@ -137,15 +137,6 @@ impl AgentRunStore for codex_state::StateRuntime { self.create_background_agent_run(¶ms).await } - async fn admit_run( - &self, - params: BackgroundAgentRunCreateParams, - max_active_runs: i64, - ) -> anyhow::Result<(BackgroundAgentRun, bool)> { - self.admit_background_agent_run(¶ms, max_active_runs) - .await - } - async fn get_run(&self, run_id: &str) -> anyhow::Result> { self.get_background_agent_run(run_id).await } @@ -176,6 +167,14 @@ impl AgentRunStore for codex_state::StateRuntime { .await } + async fn append_status_event_for_supervisor( + &self, + params: BackgroundAgentStatusEventForSupervisorParams<'_>, + ) -> anyhow::Result> { + self.append_background_agent_status_event_for_supervisor(params) + .await + } + async fn set_desired_state( &self, run_id: &str, @@ -217,9 +216,15 @@ impl AgentRunStore for codex_state::StateRuntime { run_id: &str, supervisor_id: &str, process_lease_id: &str, + required_version_fingerprint: &str, ) -> anyhow::Result> { - self.claim_background_agent_supervisor(run_id, supervisor_id, process_lease_id) - .await + self.claim_background_agent_supervisor_compatible( + run_id, + supervisor_id, + process_lease_id, + required_version_fingerprint, + ) + .await } async fn record_execution_handle( diff --git a/codex-rs/background-agent/src/supervisor.rs b/codex-rs/background-agent/src/supervisor.rs index b89591a39b..c3c74098f7 100644 --- a/codex-rs/background-agent/src/supervisor.rs +++ b/codex-rs/background-agent/src/supervisor.rs @@ -11,10 +11,10 @@ use crate::AgentSnapshotStore; use crate::AgentSupervisor; use crate::BackgroundAgentDesiredState; use crate::BackgroundAgentExecutionHandleParams; +use crate::BackgroundAgentExecutionSnapshotParams; use crate::BackgroundAgentPendingInteractionStatus; use crate::BackgroundAgentRun; use crate::BackgroundAgentRunStatus; -use crate::BackgroundAgentStatusSnapshotParams; use crate::PendingInteractionLedger; use crate::SupervisorReconcileReport; @@ -232,6 +232,7 @@ where run_id.as_str(), self.config.supervisor_id.as_str(), process_lease_id.as_str(), + crate::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, ) .await? else { @@ -262,59 +263,62 @@ where self.execution.stop(handle).await?; return Ok(false); } - self.store() - .update_run_status( - run_id.as_str(), - BackgroundAgentRunStatus::Running, - Some("worker started"), - ) - .await?; - self.store() - .append_event( - run_id.as_str(), - "agent.workerStarted", - &json!({ - "processLeaseId": handle.process_lease_id, - "pid": handle.pid, - "pgid": handle.pgid, - "jobId": handle.job_id, - "generation": generation, - }), + let event_payload = json!({ + "processLeaseId": handle.process_lease_id, + "pid": handle.pid, + "pgid": handle.pgid, + "jobId": handle.job_id, + "generation": generation, + }); + let running = self + .store() + .append_status_event_for_supervisor( + crate::BackgroundAgentStatusEventForSupervisorParams { + run_id: run_id.as_str(), + supervisor_id: self.config.supervisor_id.as_str(), + generation, + status: BackgroundAgentRunStatus::Running, + status_reason: Some("worker started"), + event_type: "agent.workerStarted", + event_payload_json: &event_payload, + summary: Some("Running"), + pending_interaction_count: 0, + status_payload_json: &json!({ + "phase": "running", + }), + }, ) .await?; - self.store() - .upsert_status_snapshot(BackgroundAgentStatusSnapshotParams { - run_id: run_id.clone(), - seq: generation, - status: BackgroundAgentRunStatus::Running, - desired_state: BackgroundAgentDesiredState::Running, - summary: Some("Running".to_string()), - pending_interaction_count: 0, - last_event_seq: 0, - payload_json: json!({ - "phase": "running", - }), - }) - .await?; + if running.is_none() { + self.execution.stop(handle).await?; + return Ok(false); + } Ok(true) } Err(err) => { let reason = format!("worker start failed: {err}"); - self.store() - .update_run_status( - run_id.as_str(), - BackgroundAgentRunStatus::Failed, - Some(reason.as_str()), - ) - .await?; - self.store() - .append_event( - run_id.as_str(), - "agent.workerStartFailed", - &json!({ - "generation": generation, - "error": err.to_string(), - }), + let event_payload = json!({ + "generation": generation, + "error": err.to_string(), + }); + let _failed = self + .store() + .append_status_event_for_supervisor( + crate::BackgroundAgentStatusEventForSupervisorParams { + run_id: run_id.as_str(), + supervisor_id: self.config.supervisor_id.as_str(), + generation, + status: BackgroundAgentRunStatus::Failed, + status_reason: Some(reason.as_str()), + event_type: "agent.workerStartFailed", + event_payload_json: &event_payload, + summary: Some(reason.as_str()), + pending_interaction_count: 0, + status_payload_json: &json!({ + "phase": "failed", + "reason": reason, + }), + }, ) .await?; Ok(false) @@ -782,26 +786,43 @@ mod tests { } async fn create_run(runtime: &StateRuntime, id: &str) -> anyhow::Result { - runtime - .create_background_agent_run(&BackgroundAgentRunCreateParams { - id: id.to_string(), - idempotency_key: None, - request_id: None, - source: "test".to_string(), - prompt_snapshot_ref: format!("prompt://{id}"), - input_snapshot_ref: None, - thread_id: None, - thread_store_kind: "local".to_string(), - thread_store_id: None, - rollout_path: None, - parent_thread_id: None, - parent_agent_run_id: None, - spawn_linkage_json: None, - auth_profile_ref: None, - status_reason: Some("created".to_string()), - config_fingerprint: None, - version_fingerprint: None, - }) - .await + let (run, _, _, _, _) = runtime + .admit_background_agent_run( + &BackgroundAgentRunCreateParams { + id: id.to_string(), + idempotency_key: None, + request_id: None, + source: "test".to_string(), + prompt_snapshot_ref: format!("prompt://{id}"), + input_snapshot_ref: None, + thread_id: None, + thread_store_kind: "local".to_string(), + thread_store_id: None, + rollout_path: None, + parent_thread_id: None, + parent_agent_run_id: None, + spawn_linkage_json: None, + auth_profile_ref: None, + status_reason: Some("created".to_string()), + config_fingerprint: None, + version_fingerprint: Some( + crate::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string(), + ), + }, + &json!({ + "prompt": format!("prompt for {id}"), + "promptSnapshotRef": format!("prompt://{id}"), + }), + &BackgroundAgentExecutionSnapshotParams { + run_id: id.to_string(), + snapshot_kind: "initial_execution_context".to_string(), + payload_json: json!({"cwd": "/tmp"}), + recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), + config_fingerprint: None, + }, + /*max_active_runs*/ 8, + ) + .await?; + Ok(run) } } diff --git a/codex-rs/cli/src/agent_cmd.rs b/codex-rs/cli/src/agent_cmd.rs index 0cb661fdd6..8059bc6c3a 100644 --- a/codex-rs/cli/src/agent_cmd.rs +++ b/codex-rs/cli/src/agent_cmd.rs @@ -21,7 +21,6 @@ use codex_state::BackgroundAgentPendingInteractionStatus; use codex_state::BackgroundAgentRun; use codex_state::BackgroundAgentRunCreateParams; use codex_state::BackgroundAgentRunStatus; -use codex_state::BackgroundAgentStatusSnapshotParams; use codex_state::StateRuntime; use codex_state::busy_retry::retry_on_busy; use codex_utils_absolute_path::AbsolutePathBuf; @@ -358,37 +357,10 @@ pub(crate) async fn run_agent_command( ) } AgentSubcommand::Delete(cmd) => { - let existing_run = state_db - .get_background_agent_run(cmd.agent_id.as_str()) - .await - .context("failed to read background agent before delete")?; let deleted = state_db .request_background_agent_delete(cmd.agent_id.as_str()) .await .context("failed to request background agent delete")?; - if deleted { - if existing_run.as_ref().is_some_and(|run| { - !background_agent_status_is_terminal(run.status) - && should_terminalize_unclaimed_agent_run(run) - }) { - state_db - .update_background_agent_run_status( - cmd.agent_id.as_str(), - BackgroundAgentRunStatus::Cancelled, - Some("delete requested by codewith agent delete before worker claim"), - ) - .await - .context("failed to update background agent status after delete")?; - } - state_db - .append_background_agent_event( - cmd.agent_id.as_str(), - "agent.deleteRequested", - &json!({"reason": "cli_requested_delete"}), - ) - .await - .context("failed to append background agent delete event")?; - } let run = state_db .get_background_agent_run(cmd.agent_id.as_str()) .await @@ -997,6 +969,7 @@ async fn start_agent( } ensure_background_agent_supported_platform()?; + let daemon_output = background_agent_daemon()?.start().await?; let agent_id = new_agent_id(); let explicit_cwd = cmd.cwd.is_some(); @@ -1011,13 +984,64 @@ async fn start_agent( .and_then(|context| context.auth_profile_ref.as_deref()) .or(auth_profile) .map(str::to_string); - let prompt_snapshot_ref = format!("inline:{agent_id}:prompt"); + let idempotency_key = cmd.idempotency_key; + let prompt_snapshot_identity = idempotency_key + .as_deref() + .map(|key| StateRuntime::background_agent_identity_sha256(key.as_bytes())) + .unwrap_or_else(|| agent_id.clone()); + let prompt_snapshot_ref = format!("inline:{prompt_snapshot_identity}:prompt"); + let config_identity = json!({ + "cwd": cwd.display().to_string(), + "workspaceRoots": &workspace_roots, + "authProfileRef": &auth_profile_ref, + "approvalPolicy": runtime_context.and_then(|context| context.approval_policy.as_ref()), + "permissionProfile": &permission_profile, + "model": runtime_context.and_then(|context| context.model.as_deref()), + "provider": runtime_context.and_then(|context| context.provider.as_deref()), + "serviceTier": runtime_context.and_then(|context| context.service_tier.as_deref()), + "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", + }); + let config_fingerprint = StateRuntime::background_agent_identity_sha256( + serde_json::to_vec(&config_identity)?.as_slice(), + ); + let start_event_payload = json!({ + "cwd": cwd.display().to_string(), + "prompt": prompt, + "promptSha256": StateRuntime::background_agent_identity_sha256(prompt.as_bytes()), + "promptSnapshotRef": prompt_snapshot_ref.as_str(), + }); + let snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: agent_id.clone(), + snapshot_kind: "initial_execution_context".to_string(), + payload_json: json!({ + "snapshotSource": "codewith agent start", + "cwd": cwd.display().to_string(), + "workspaceRoots": workspace_roots, + "approvalPolicy": runtime_context + .and_then(|context| context.approval_policy.as_ref()), + "permissionProfile": permission_profile, + "model": runtime_context.and_then(|context| context.model.as_deref()), + "provider": runtime_context.and_then(|context| context.provider.as_deref()), + "serviceTier": runtime_context + .and_then(|context| context.service_tier.as_deref()), + "configFingerprint": config_fingerprint.as_str(), + "versionFingerprint": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + "packageFingerprint": format!( + "{}:{}", + env!("CARGO_PKG_NAME"), + env!("CARGO_PKG_VERSION") + ), + "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", + }), + recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), + config_fingerprint: Some(config_fingerprint.clone()), + }; // The state DB is shared across many concurrent processes; every write // below retries transient SQLITE_BUSY / SQLITE_BUSY_SNAPSHOT contention // with backoff instead of failing the whole `agent start` invocation. let create_params = BackgroundAgentRunCreateParams { id: agent_id.clone(), - idempotency_key: cmd.idempotency_key, + idempotency_key, request_id: None, source: "cli".to_string(), prompt_snapshot_ref: prompt_snapshot_ref.clone(), @@ -1031,7 +1055,7 @@ async fn start_agent( spawn_linkage_json: None, auth_profile_ref: auth_profile_ref.clone(), status_reason: Some("queued by codewith agent start".to_string()), - config_fingerprint: None, + config_fingerprint: Some(config_fingerprint), version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), }; retry_on_busy("reconcile stale background agents before admission", || { @@ -1039,128 +1063,17 @@ async fn start_agent( }) .await .context("failed to reconcile stale background agents before admission")?; - let (run, created) = retry_on_busy("admit background agent run", || { - state_db - .admit_background_agent_run(&create_params, DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS) - }) - .await - .context("failed to admit background agent")?; - let admitted_agent_id = run.id.clone(); - let start_event_payload = json!({ - "cwd": cwd.display().to_string(), - "prompt": prompt, - "promptSnapshotRef": run.prompt_snapshot_ref, - }); - let existing_start_event = if created { - None - } else { - retry_on_busy("load background agent start event", || { - state_db.list_background_agent_events_after( - admitted_agent_id.as_str(), - /*after_seq*/ None, - Some(100), - ) - }) - .await - .context("failed to load background agent start event")? - .into_iter() - .find(|event| { - matches!( - event.event_type.as_str(), - "agent.started" | "agent.startRecovered" + let (run, created, _event, _execution_snapshot, _status_snapshot) = + retry_on_busy("admit background agent run", || { + state_db.admit_background_agent_run( + &create_params, + &start_event_payload, + &snapshot_params, + DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, ) }) - }; - let event = match existing_start_event { - Some(event) => event, - None => { - let event_type = if created { - "agent.started" - } else { - "agent.startRecovered" - }; - retry_on_busy("append background agent start event", || { - state_db.append_background_agent_event( - admitted_agent_id.as_str(), - event_type, - &start_event_payload, - ) - }) - .await - .context("failed to append background agent start event")? - } - }; - let snapshot_params = BackgroundAgentExecutionSnapshotParams { - run_id: admitted_agent_id.clone(), - snapshot_kind: "initial_execution_context".to_string(), - payload_json: json!({ - "snapshotSource": "codewith agent start", - "cwd": cwd.display().to_string(), - "workspaceRoots": workspace_roots, - "authProfileRef": auth_profile_ref, - "approvalPolicy": runtime_context - .and_then(|context| context.approval_policy.as_ref()), - "permissionProfile": permission_profile, - "model": runtime_context.and_then(|context| context.model.as_deref()), - "provider": runtime_context.and_then(|context| context.provider.as_deref()), - "serviceTier": runtime_context - .and_then(|context| context.service_tier.as_deref()), - "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", - }), - recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), - config_fingerprint: None, - }; - let execution_snapshot_exists = - retry_on_busy("load background agent execution snapshot", || { - state_db.get_latest_background_agent_execution_snapshot(admitted_agent_id.as_str()) - }) - .await - .context("failed to load background agent execution snapshot")? - .is_some(); - if !execution_snapshot_exists { - retry_on_busy("create background agent execution snapshot", || { - state_db.create_background_agent_execution_snapshot(&snapshot_params) - }) - .await - .context("failed to create background agent execution snapshot")?; - } - let status_snapshot_exists = retry_on_busy("load background agent status snapshot", || { - state_db.get_background_agent_status_snapshot(admitted_agent_id.as_str()) - }) - .await - .context("failed to load background agent status snapshot")? - .is_some(); - if !status_snapshot_exists { - let snapshot_run = retry_on_busy("reload background agent for status snapshot", || { - state_db.get_background_agent_run(admitted_agent_id.as_str()) - }) - .await - .context("failed to reload background agent for status snapshot")? - .unwrap_or_else(|| run.clone()); - let last_event_seq = snapshot_run.last_event_seq.max(event.seq); - let status_snapshot_params = BackgroundAgentStatusSnapshotParams { - run_id: admitted_agent_id.clone(), - seq: last_event_seq, - status: snapshot_run.status, - desired_state: snapshot_run.desired_state, - summary: Some(snapshot_run.status.as_str().to_string()), - pending_interaction_count: 0, - last_event_seq, - payload_json: json!({"phase": snapshot_run.status.as_str()}), - }; - retry_on_busy("create background agent status snapshot", || { - state_db.upsert_background_agent_status_snapshot(&status_snapshot_params) - }) .await - .context("failed to create background agent status snapshot")?; - } - let daemon = background_agent_daemon()?; - let daemon_output = daemon.start().await?; - let run = retry_on_busy("reload admitted background agent", || { - state_db.get_background_agent_run(admitted_agent_id.as_str()) - }) - .await? - .unwrap_or(run); + .context("failed to admit background agent")?; Ok(json!({ "agent": run_json(run), "created": created, "daemon": daemon_output })) } @@ -1213,6 +1126,13 @@ async fn stop_agent( else { return Ok(None); }; + if matches!( + run.retention_state, + codex_state::BackgroundAgentRetentionState::DeleteRequested + | codex_state::BackgroundAgentRetentionState::Deleted + ) { + return Ok(Some(run)); + } if !background_agent_status_is_terminal(run.status) { let mut observed = run; let mut stopped = false; diff --git a/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql b/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql index 0434aa111d..19a8dbfd75 100644 --- a/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql +++ b/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql @@ -1,6 +1,53 @@ +ALTER TABLE background_agent_runs + ADD COLUMN admission_identity_sha256 TEXT; + +ALTER TABLE background_agent_runs + ADD COLUMN admission_ready_at INTEGER; + ALTER TABLE background_agent_events ADD COLUMN receipt_key TEXT; CREATE UNIQUE INDEX idx_background_agent_events_receipt_key ON background_agent_events(run_id, receipt_key) WHERE receipt_key IS NOT NULL; + +CREATE TABLE background_agent_lifecycle_receipts ( + run_id TEXT NOT NULL, + receipt_key TEXT NOT NULL, + event_id INTEGER NOT NULL, + event_seq INTEGER NOT NULL, + event_type TEXT NOT NULL, + generation INTEGER NOT NULL, + attempt INTEGER, + operation_identity_sha256 TEXT NOT NULL, + payload_json TEXT NOT NULL, + created_at INTEGER NOT NULL, + PRIMARY KEY(run_id, receipt_key), + FOREIGN KEY(run_id) REFERENCES background_agent_runs(id) ON DELETE CASCADE +); + +INSERT INTO background_agent_lifecycle_receipts ( + run_id, + receipt_key, + event_id, + event_seq, + event_type, + generation, + attempt, + operation_identity_sha256, + payload_json, + created_at +) +SELECT + run_id, + receipt_key, + id, + seq, + event_type, + COALESCE(json_extract(payload_json, '$.generation'), 0), + json_extract(payload_json, '$.attempt'), + '', + payload_json, + created_at +FROM background_agent_events +WHERE receipt_key IS NOT NULL; diff --git a/codex-rs/state/src/runtime/background_agents/events.rs b/codex-rs/state/src/runtime/background_agents/events.rs index 484c70e2da..8f03c6342c 100644 --- a/codex-rs/state/src/runtime/background_agents/events.rs +++ b/codex-rs/state/src/runtime/background_agents/events.rs @@ -1,8 +1,13 @@ use super::*; use crate::BACKGROUND_AGENT_EVENT_CURSOR_COMPACTED; +use sha2::Digest; +use sha2::Sha256; const MAX_BACKGROUND_AGENT_RECEIPT_DIAGNOSTICS_BYTES: usize = 4 * 1024; const MAX_BACKGROUND_AGENT_RECEIPT_DIAGNOSTICS_PREVIEW_CHARS: usize = 1_024; +const MAX_BACKGROUND_AGENT_RECEIPT_KEY_BYTES: usize = 256; +const BACKGROUND_AGENT_RECEIPT_IDENTITY_MISMATCH: &str = + "background agent lifecycle receipt identity mismatch"; pub(in crate::runtime) async fn append_background_agent_event_in_tx( tx: &mut sqlx::Transaction<'_, Sqlite>, @@ -69,22 +74,29 @@ pub(in crate::runtime) async fn append_background_agent_lifecycle_receipt_in_tx( diagnostics_json: &serde_json::Value, now: i64, ) -> anyhow::Result { - if let Some(row) = sqlx::query_as::<_, BackgroundAgentEventRow>( - r#" -SELECT id, run_id, seq, event_type, payload_json, created_at -FROM background_agent_events -WHERE run_id = ? AND receipt_key = ? - "#, + validate_background_agent_receipt_key(receipt_key)?; + let operation_identity_sha256 = background_agent_receipt_operation_identity_sha256( + event_type, + generation, + attempt, + diagnostics_json, + )?; + let diagnostics_json = bounded_background_agent_receipt_diagnostics(diagnostics_json)?; + if let Some(event) = get_background_agent_lifecycle_receipt_in_tx( + tx, + run_id, + event_type, + receipt_key, + generation, + attempt, + &diagnostics_json, + operation_identity_sha256.as_str(), ) - .bind(run_id) - .bind(receipt_key) - .fetch_optional(&mut **tx) .await? { - return BackgroundAgentEvent::try_from(row); + return Ok(event); } - let diagnostics_json = bounded_background_agent_receipt_diagnostics(diagnostics_json)?; let payload_json = crate::redacted_local_state_json(&serde_json::json!({ "receiptKey": receipt_key, "runId": run_id, @@ -115,12 +127,40 @@ INSERT INTO background_agent_events ( .bind(run_id) .bind(seq) .bind(event_type) - .bind(serialized_payload) + .bind(serialized_payload.as_str()) .bind(now) .bind(receipt_key) .execute(&mut **tx) .await? .last_insert_rowid(); + sqlx::query( + r#" +INSERT INTO background_agent_lifecycle_receipts ( + run_id, + receipt_key, + event_id, + event_seq, + event_type, + generation, + attempt, + operation_identity_sha256, + payload_json, + created_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + "#, + ) + .bind(run_id) + .bind(receipt_key) + .bind(id) + .bind(seq) + .bind(event_type) + .bind(generation) + .bind(attempt) + .bind(operation_identity_sha256) + .bind(serialized_payload) + .bind(now) + .execute(&mut **tx) + .await?; sqlx::query( r#" @@ -146,6 +186,76 @@ WHERE id = ? }) } +#[allow(clippy::too_many_arguments)] +pub(in crate::runtime) async fn get_background_agent_lifecycle_receipt_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + run_id: &str, + event_type: &str, + receipt_key: &str, + generation: i64, + attempt: Option, + diagnostics_json: &serde_json::Value, + operation_identity_sha256: &str, +) -> anyhow::Result> { + validate_background_agent_receipt_key(receipt_key)?; + let row = sqlx::query_as::<_, (i64, i64, String, i64, Option, String, String, i64)>( + r#" +SELECT + event_id, + event_seq, + event_type, + generation, + attempt, + operation_identity_sha256, + payload_json, + created_at +FROM background_agent_lifecycle_receipts +WHERE run_id = ? AND receipt_key = ? + "#, + ) + .bind(run_id) + .bind(receipt_key) + .fetch_optional(&mut **tx) + .await?; + let Some(( + event_id, + event_seq, + stored_event_type, + stored_generation, + stored_attempt, + stored_operation_identity_sha256, + payload_json, + created_at, + )) = row + else { + return Ok(None); + }; + let payload_json: serde_json::Value = serde_json::from_str(payload_json.as_str())?; + let legacy_identity_matches = stored_operation_identity_sha256.is_empty() + && payload_json.get("diagnostics") == Some(diagnostics_json); + if stored_event_type != event_type + || stored_generation != generation + || stored_attempt != attempt + || (stored_operation_identity_sha256 != operation_identity_sha256 + && !legacy_identity_matches) + { + anyhow::bail!( + "{BACKGROUND_AGENT_RECEIPT_IDENTITY_MISMATCH}: \ + receipt key is already bound to a different lifecycle operation" + ); + } + let created_at = DateTime::::from_timestamp(created_at, 0) + .ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {created_at}"))?; + Ok(Some(BackgroundAgentEvent { + id: event_id, + run_id: run_id.to_string(), + seq: event_seq, + event_type: stored_event_type, + payload_json, + created_at, + })) +} + impl StateRuntime { pub async fn append_background_agent_event( &self, @@ -162,6 +272,55 @@ impl StateRuntime { Ok(event) } + #[allow(clippy::too_many_arguments)] + pub async fn append_background_agent_event_for_supervisor( + &self, + run_id: &str, + supervisor_id: &str, + generation: i64, + event_type: &str, + payload_json: &serde_json::Value, + allow_terminal_current: bool, + ) -> anyhow::Result> { + let now = Utc::now().timestamp(); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let current: Option = sqlx::query_scalar( + r#" +SELECT 1 +FROM background_agent_runs +WHERE + id = ? + AND supervisor_id = ? + AND generation = ? + AND ( + ? = 1 + OR status IN ( + 'starting', + 'running', + 'waiting_on_approval', + 'waiting_on_user', + 'stopping' + ) + ) + "#, + ) + .bind(run_id) + .bind(supervisor_id) + .bind(generation) + .bind(allow_terminal_current) + .fetch_optional(&mut *tx) + .await?; + let Some(_) = current else { + tx.commit().await?; + return Ok(None); + }; + let event = + append_background_agent_event_in_tx(&mut tx, run_id, event_type, payload_json, now) + .await?; + tx.commit().await?; + Ok(Some(event)) + } + pub async fn append_background_agent_lifecycle_receipt( &self, run_id: &str, @@ -276,7 +435,35 @@ GROUP BY r.id } } -fn bounded_background_agent_receipt_diagnostics( +fn validate_background_agent_receipt_key(receipt_key: &str) -> anyhow::Result<()> { + if receipt_key.len() > MAX_BACKGROUND_AGENT_RECEIPT_KEY_BYTES { + anyhow::bail!( + "background agent lifecycle receipt key exceeds \ + {MAX_BACKGROUND_AGENT_RECEIPT_KEY_BYTES} bytes" + ); + } + Ok(()) +} + +pub(in crate::runtime) fn background_agent_receipt_operation_identity_sha256( + event_type: &str, + generation: i64, + attempt: Option, + diagnostics_json: &serde_json::Value, +) -> anyhow::Result { + let identity = serde_json::json!({ + "eventType": event_type, + "generation": generation, + "attempt": attempt, + "diagnostics": diagnostics_json, + }); + Ok(format!( + "{:x}", + Sha256::digest(serde_json::to_vec(&identity)?) + )) +} + +pub(in crate::runtime) fn bounded_background_agent_receipt_diagnostics( diagnostics_json: &serde_json::Value, ) -> anyhow::Result { let diagnostics_json = crate::redacted_local_state_json(diagnostics_json); diff --git a/codex-rs/state/src/runtime/background_agents/runs.rs b/codex-rs/state/src/runtime/background_agents/runs.rs index 022471a885..1ea16e657c 100644 --- a/codex-rs/state/src/runtime/background_agents/runs.rs +++ b/codex-rs/state/src/runtime/background_agents/runs.rs @@ -1,4 +1,6 @@ use super::*; +use sha2::Digest; +use sha2::Sha256; const BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED: &str = "background_agent_admission_capacity_exceeded"; @@ -6,11 +8,16 @@ const BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH: &str = "background_agent_admission_identity_mismatch"; impl StateRuntime { + /// Returns a stable digest for opaque background-agent admission identity. + pub fn background_agent_identity_sha256(bytes: &[u8]) -> String { + format!("{:x}", Sha256::digest(bytes)) + } + pub async fn create_background_agent_run( &self, params: &BackgroundAgentRunCreateParams, ) -> anyhow::Result { - self.create_or_admit_background_agent_run(params, /*max_active_runs*/ None) + self.create_background_agent_run_row(params) .await .map(|(run, _created)| run) } @@ -18,20 +25,136 @@ impl StateRuntime { pub async fn admit_background_agent_run( &self, params: &BackgroundAgentRunCreateParams, + start_event_payload_json: &serde_json::Value, + execution_snapshot_params: &BackgroundAgentExecutionSnapshotParams, max_active_runs: i64, - ) -> anyhow::Result<(BackgroundAgentRun, bool)> { - self.create_or_admit_background_agent_run(params, Some(max_active_runs)) - .await + ) -> anyhow::Result<( + BackgroundAgentRun, + bool, + BackgroundAgentEvent, + BackgroundAgentExecutionSnapshot, + BackgroundAgentStatusSnapshot, + )> { + let idempotency_key = params.idempotency_key.as_deref(); + let mut execution_snapshot_params = execution_snapshot_params.clone(); + execution_snapshot_params.run_id.clone_from(¶ms.id); + let admission_identity_sha256 = background_agent_admission_identity_sha256( + params, + start_event_payload_json, + &execution_snapshot_params, + )?; + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + if let Some(idempotency_key) = idempotency_key + && let Some(existing_id) = + validate_existing_background_agent_admission_in_tx(&mut tx, idempotency_key, params) + .await? + { + let (event, execution_snapshot_id) = + recover_or_validate_background_agent_initial_state_in_tx( + &mut tx, + existing_id.as_str(), + admission_identity_sha256.as_str(), + start_event_payload_json, + &execution_snapshot_params, + ) + .await?; + tx.commit().await?; + return self + .load_background_agent_admission_result( + existing_id.as_str(), + /*created*/ false, + event, + execution_snapshot_id, + ) + .await; + } + + ensure_background_agent_capacity_in_tx(&mut tx, max_active_runs).await?; + let now = Utc::now().timestamp(); + insert_background_agent_run_in_tx(&mut tx, params, now).await?; + let receipt_identity = idempotency_key.unwrap_or(params.id.as_str()); + let receipt_key = format!( + "admission:{:x}", + Sha256::digest(receipt_identity.as_bytes()) + ); + super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + params.id.as_str(), + "agent.admitted", + receipt_key.as_str(), + /*generation*/ 0, + Some(1), + &serde_json::json!({ + "source": params.source, + "requestId": params.request_id, + "versionFingerprint": params.version_fingerprint, + }), + now, + ) + .await?; + let start_event_payload_json = background_agent_start_event_payload( + start_event_payload_json, + admission_identity_sha256.as_str(), + ); + let event = super::events::append_background_agent_event_in_tx( + &mut tx, + params.id.as_str(), + "agent.started", + &start_event_payload_json, + now, + ) + .await?; + let execution_snapshot_id = insert_background_agent_execution_snapshot_in_tx( + &mut tx, + &execution_snapshot_params, + now, + ) + .await?; + super::snapshots::upsert_background_agent_status_snapshot_in_tx( + &mut tx, + &BackgroundAgentStatusSnapshotParams { + run_id: params.id.clone(), + seq: event.seq, + status: BackgroundAgentRunStatus::Queued, + desired_state: BackgroundAgentDesiredState::Running, + summary: Some("Queued".to_string()), + pending_interaction_count: 0, + last_event_seq: event.seq, + payload_json: serde_json::json!({"phase": "queued"}), + }, + now, + ) + .await?; + sqlx::query( + r#" +UPDATE background_agent_runs +SET admission_identity_sha256 = ?, admission_ready_at = ?, updated_at = ? +WHERE id = ? + "#, + ) + .bind(admission_identity_sha256) + .bind(now) + .bind(now) + .bind(params.id.as_str()) + .execute(&mut *tx) + .await?; + tx.commit().await?; + self.load_background_agent_admission_result( + params.id.as_str(), + /*created*/ true, + event, + execution_snapshot_id, + ) + .await } - async fn create_or_admit_background_agent_run( + async fn create_background_agent_run_row( &self, params: &BackgroundAgentRunCreateParams, - max_active_runs: Option, ) -> anyhow::Result<(BackgroundAgentRun, bool)> { - let idempotency_key = params.idempotency_key.as_deref().map(redact_state_string); + let idempotency_key = params.idempotency_key.as_deref(); let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; - if let Some(idempotency_key) = idempotency_key.as_deref() + if let Some(idempotency_key) = idempotency_key && let Some(existing_id) = validate_existing_background_agent_admission_in_tx(&mut tx, idempotency_key, params) .await? @@ -48,130 +171,8 @@ impl StateRuntime { return Ok((existing, false)); } - if let Some(max_active_runs) = max_active_runs { - let active_run_count: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM background_agent_runs -WHERE - desired_state = 'running' - AND retention_state = 'active' - AND status IN ( - 'queued', - 'starting', - 'running', - 'waiting_on_approval', - 'waiting_on_user', - 'stopping', - 'orphaned' - ) - "#, - ) - .fetch_one(&mut *tx) - .await?; - if active_run_count >= max_active_runs { - anyhow::bail!( - "{BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED}: \ - {active_run_count} live or recoverable run(s), max {max_active_runs}" - ); - } - } - let now = Utc::now().timestamp(); - let spawn_linkage_json = params - .spawn_linkage_json - .as_ref() - .map(redact_state_json_string) - .transpose()?; - let insert_result = sqlx::query( - r#" -INSERT INTO background_agent_runs ( - id, - idempotency_key, - request_id, - source, - prompt_snapshot_ref, - input_snapshot_ref, - thread_id, - thread_store_kind, - thread_store_id, - rollout_path, - parent_thread_id, - parent_agent_run_id, - spawn_linkage_json, - auth_profile_ref, - desired_state, - status, - status_reason, - config_fingerprint, - version_fingerprint, - retention_state, - created_at, - updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(params.id.as_str()) - .bind(idempotency_key.as_deref()) - .bind(params.request_id.as_deref().map(redact_state_string)) - .bind(params.source.as_str()) - .bind(redact_state_string(params.prompt_snapshot_ref.as_str())) - .bind( - params - .input_snapshot_ref - .as_deref() - .map(redact_state_string), - ) - .bind(params.thread_id.as_deref()) - .bind(params.thread_store_kind.as_str()) - .bind(params.thread_store_id.as_deref()) - .bind(params.rollout_path.as_deref()) - .bind(params.parent_thread_id.as_deref()) - .bind(params.parent_agent_run_id.as_deref()) - .bind(spawn_linkage_json.as_deref()) - .bind(params.auth_profile_ref.as_deref().map(redact_state_string)) - .bind(BackgroundAgentDesiredState::Running.as_str()) - .bind(BackgroundAgentRunStatus::Queued.as_str()) - .bind(params.status_reason.as_deref().map(redact_state_string)) - .bind(params.config_fingerprint.as_deref()) - .bind(params.version_fingerprint.as_deref()) - .bind(crate::BackgroundAgentRetentionState::Active.as_str()) - .bind(now) - .bind(now) - .execute(&mut *tx) - .await; - if let Err(err) = insert_result { - if is_background_agent_unique_constraint_violation(&err) { - anyhow::bail!( - "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ - request identity conflicts with an existing background agent" - ); - } - return Err(err.into()); - } - - if max_active_runs.is_some() { - let receipt_key = format!( - "admission:{}", - idempotency_key.as_deref().unwrap_or(params.id.as_str()) - ); - super::events::append_background_agent_lifecycle_receipt_in_tx( - &mut tx, - params.id.as_str(), - "agent.admitted", - receipt_key.as_str(), - /*generation*/ 0, - Some(1), - &serde_json::json!({ - "source": params.source, - "requestId": params.request_id, - "authProfileRef": params.auth_profile_ref, - "versionFingerprint": params.version_fingerprint, - }), - now, - ) - .await?; - } + insert_background_agent_run_in_tx(&mut tx, params, now).await?; tx.commit().await?; let run = self .get_background_agent_run(params.id.as_str()) @@ -180,6 +181,40 @@ INSERT INTO background_agent_runs ( Ok((run, true)) } + async fn load_background_agent_admission_result( + &self, + run_id: &str, + created: bool, + event: BackgroundAgentEvent, + execution_snapshot_id: i64, + ) -> anyhow::Result<( + BackgroundAgentRun, + bool, + BackgroundAgentEvent, + BackgroundAgentExecutionSnapshot, + BackgroundAgentStatusSnapshot, + )> { + let run = self + .get_background_agent_run(run_id) + .await? + .ok_or_else(|| anyhow::anyhow!("failed to load background agent run {run_id}"))?; + let execution_snapshot = self + .get_background_agent_execution_snapshot(execution_snapshot_id) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "failed to load background agent execution snapshot {execution_snapshot_id}" + ) + })?; + let status_snapshot = self + .get_background_agent_status_snapshot(run_id) + .await? + .ok_or_else(|| { + anyhow::anyhow!("failed to load background agent status snapshot for run {run_id}") + })?; + Ok((run, created, event, execution_snapshot, status_snapshot)) + } + pub async fn get_background_agent_run( &self, run_id: &str, @@ -317,9 +352,7 @@ OFFSET ? SELECT status, COUNT(*) as count FROM background_agent_runs WHERE - retention_state = 'active' - AND desired_state = 'running' - AND status IN ( + status IN ( 'queued', 'starting', 'running', @@ -328,6 +361,20 @@ WHERE 'stopping', 'orphaned' ) + AND ( + status = 'stopping' + OR ( + status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') + AND supervisor_id IS NOT NULL + ) + OR ( + status IN ('queued', 'orphaned') + AND desired_state = 'running' + AND retention_state = 'active' + AND admission_identity_sha256 IS NOT NULL + AND admission_ready_at IS NOT NULL + ) + ) GROUP BY status "#, ) @@ -356,7 +403,12 @@ SET updated_at = ?, started_at = COALESCE(started_at, ?), completed_at = COALESCE(?, completed_at) -WHERE id = ? +WHERE + id = ? + AND ( + status NOT IN ('completed', 'failed', 'cancelled') + OR status = ? + ) "#, ) .bind(status.as_str()) @@ -365,6 +417,7 @@ WHERE id = ? .bind(started_at) .bind(completed_at) .bind(run_id) + .bind(status.as_str()) .execute(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) @@ -423,7 +476,76 @@ WHERE ) -> anyhow::Result> { let now = Utc::now().timestamp(); let (started_at, completed_at) = background_agent_status_timestamps(params.status, now); - let mut tx = self.pool.begin().await?; + let receipt_key = format!( + "status:{}:{}:{}", + params.generation, + params.event_type, + params.status.as_str() + ); + let operation_diagnostics_json = serde_json::json!({ + "status": params.status.as_str(), + "statusReason": params.status_reason, + "eventPayload": params.event_payload_json, + "summary": params.summary, + "pendingInteractionCount": params.pending_interaction_count, + "statusPayload": params.status_payload_json, + }); + let operation_identity_sha256 = + super::events::background_agent_receipt_operation_identity_sha256( + params.event_type, + params.generation, + /*attempt*/ None, + &operation_diagnostics_json, + )?; + let diagnostics_json = super::events::bounded_background_agent_receipt_diagnostics( + &operation_diagnostics_json, + )?; + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + if let Some(event) = super::events::get_background_agent_lifecycle_receipt_in_tx( + &mut tx, + params.run_id, + params.event_type, + receipt_key.as_str(), + params.generation, + /*attempt*/ None, + &diagnostics_json, + operation_identity_sha256.as_str(), + ) + .await? + { + let terminal_projection: Option<(String, Option, Option)> = + sqlx::query_as( + r#" +SELECT + r.status, + s.status, + s.last_event_seq +FROM background_agent_runs r +LEFT JOIN background_agent_status_snapshots s ON s.run_id = r.id +WHERE + r.id = ? + AND r.supervisor_id = ? + AND r.generation = ? + "#, + ) + .bind(params.run_id) + .bind(params.supervisor_id) + .bind(params.generation) + .fetch_optional(&mut *tx) + .await?; + let projection_matches = terminal_projection.is_some_and( + |(run_status, snapshot_status, snapshot_last_event_seq)| { + run_status == params.status.as_str() + && snapshot_status.as_deref() == Some(params.status.as_str()) + && snapshot_last_event_seq == Some(event.seq) + }, + ); + if !projection_matches { + anyhow::bail!("background agent lifecycle receipt terminal projection mismatch"); + } + tx.commit().await?; + return Ok(Some(event)); + } let result = sqlx::query( r#" UPDATE background_agent_runs @@ -463,12 +585,6 @@ WHERE return Ok(None); } - let receipt_key = format!( - "status:{}:{}:{}", - params.generation, - params.event_type, - params.status.as_str() - ); let event = super::events::append_background_agent_lifecycle_receipt_in_tx( &mut tx, params.run_id, @@ -476,7 +592,7 @@ WHERE receipt_key.as_str(), params.generation, /*attempt*/ None, - params.event_payload_json, + &operation_diagnostics_json, now, ) .await?; @@ -693,6 +809,83 @@ WHERE Ok(result.rows_affected() > 0) } + pub async fn create_background_agent_execution_snapshot_for_supervisor( + &self, + params: &BackgroundAgentExecutionSnapshotParams, + supervisor_id: &str, + generation: i64, + ) -> anyhow::Result> { + let now = Utc::now().timestamp(); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let current: Option = sqlx::query_scalar( + r#" +SELECT 1 +FROM background_agent_runs +WHERE + id = ? + AND supervisor_id = ? + AND generation = ? + AND status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') + "#, + ) + .bind(params.run_id.as_str()) + .bind(supervisor_id) + .bind(generation) + .fetch_optional(&mut *tx) + .await?; + if current.is_none() { + tx.commit().await?; + return Ok(None); + } + let snapshot_id = + insert_background_agent_execution_snapshot_in_tx(&mut tx, params, now).await?; + tx.commit().await?; + self.get_background_agent_execution_snapshot(snapshot_id) + .await + } + + pub async fn upsert_background_agent_status_snapshot_for_supervisor( + &self, + params: &BackgroundAgentStatusSnapshotParams, + supervisor_id: &str, + generation: i64, + ) -> anyhow::Result> { + let now = Utc::now().timestamp(); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let current_last_event_seq: Option = sqlx::query_scalar( + r#" +SELECT last_event_seq +FROM background_agent_runs +WHERE + id = ? + AND supervisor_id = ? + AND generation = ? + AND status = ? + AND desired_state = ? + AND status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') + "#, + ) + .bind(params.run_id.as_str()) + .bind(supervisor_id) + .bind(generation) + .bind(params.status.as_str()) + .bind(params.desired_state.as_str()) + .fetch_optional(&mut *tx) + .await?; + let Some(current_last_event_seq) = current_last_event_seq else { + tx.commit().await?; + return Ok(None); + }; + let mut params = params.clone(); + params.seq = params.seq.max(current_last_event_seq); + params.last_event_seq = current_last_event_seq; + super::snapshots::upsert_background_agent_status_snapshot_in_tx(&mut tx, ¶ms, now) + .await?; + tx.commit().await?; + self.get_background_agent_status_snapshot(params.run_id.as_str()) + .await + } + pub async fn set_background_agent_desired_state( &self, run_id: &str, @@ -716,40 +909,147 @@ WHERE id = ? pub async fn request_background_agent_delete(&self, run_id: &str) -> anyhow::Result { let now = Utc::now().timestamp(); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let current: Option<(Option, i64, String, String)> = sqlx::query_as( + r#" +SELECT supervisor_id, generation, status, retention_state +FROM background_agent_runs +WHERE id = ? + "#, + ) + .bind(run_id) + .fetch_optional(&mut *tx) + .await?; + let Some((supervisor_id, generation, status, retention_state)) = current else { + tx.commit().await?; + return Ok(false); + }; + if matches!(retention_state.as_str(), "delete_requested" | "deleted") { + tx.commit().await?; + return Ok(true); + } + let current_status = BackgroundAgentRunStatus::parse(status.as_str())?; + let already_terminal = matches!( + current_status, + BackgroundAgentRunStatus::Completed + | BackgroundAgentRunStatus::Failed + | BackgroundAgentRunStatus::Cancelled + ); + let terminalize_immediately = !already_terminal + && (supervisor_id.is_none() || matches!(status.as_str(), "queued" | "orphaned")); + let next_status = if already_terminal { + current_status + } else if terminalize_immediately { + BackgroundAgentRunStatus::Cancelled + } else { + BackgroundAgentRunStatus::Stopping + }; + let status_reason = if already_terminal { + "delete requested for terminal run" + } else if terminalize_immediately { + "delete requested before worker claim" + } else { + "delete requested" + }; let result = sqlx::query( r#" UPDATE background_agent_runs SET desired_state = ?, retention_state = ?, - status = CASE - WHEN supervisor_id IS NOT NULL - AND status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') - THEN ? - ELSE status - END, - status_reason = CASE - WHEN supervisor_id IS NOT NULL - AND status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') - THEN ? - ELSE status_reason - END, + status = ?, + status_reason = ?, delete_after = COALESCE(delete_after, ?), - updated_at = ? -WHERE id = ? AND retention_state != ? + updated_at = ?, + completed_at = CASE + WHEN ? = ? THEN COALESCE(completed_at, ?) + ELSE completed_at + END +WHERE + id = ? + AND generation = ? + AND ( + (supervisor_id IS NULL AND ? IS NULL) + OR supervisor_id = ? + ) + AND retention_state = ? + "#, + ) + .bind(BackgroundAgentDesiredState::Deleted.as_str()) + .bind(crate::BackgroundAgentRetentionState::DeleteRequested.as_str()) + .bind(next_status.as_str()) + .bind(status_reason) + .bind(now) + .bind(now) + .bind(next_status.as_str()) + .bind(BackgroundAgentRunStatus::Cancelled.as_str()) + .bind(now) + .bind(run_id) + .bind(generation) + .bind(supervisor_id.as_deref()) + .bind(supervisor_id.as_deref()) + .bind(crate::BackgroundAgentRetentionState::Active.as_str()) + .execute(&mut *tx) + .await?; + if result.rows_affected() == 0 { + tx.commit().await?; + return Ok(false); + } + let diagnostics_json = serde_json::json!({"reason": "delete_requested"}); + let receipt_key = format!("delete:{generation}"); + let event = super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + run_id, + "agent.deleteRequested", + receipt_key.as_str(), + generation, + /*attempt*/ None, + &diagnostics_json, + now, + ) + .await?; + if terminalize_immediately { + super::interactions::terminalize_active_background_agent_pending_interactions_in_tx( + &mut tx, + run_id, + BackgroundAgentPendingInteractionStatus::Cancelled, + &diagnostics_json, + now, + ) + .await?; + } + let pending_interaction_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM background_agent_pending_interactions +WHERE run_id = ? AND status IN (?, ?) "#, ) - .bind(BackgroundAgentDesiredState::Deleted.as_str()) - .bind(crate::BackgroundAgentRetentionState::DeleteRequested.as_str()) - .bind(BackgroundAgentRunStatus::Stopping.as_str()) - .bind("delete requested") - .bind(now) - .bind(now) - .bind(run_id) - .bind(crate::BackgroundAgentRetentionState::Deleted.as_str()) - .execute(self.pool.as_ref()) + .bind(run_id) + .bind(BackgroundAgentPendingInteractionStatus::Pending.as_str()) + .bind(BackgroundAgentPendingInteractionStatus::Delivered.as_str()) + .fetch_one(&mut *tx) + .await?; + super::snapshots::upsert_background_agent_status_snapshot_in_tx( + &mut tx, + &BackgroundAgentStatusSnapshotParams { + run_id: run_id.to_string(), + seq: event.seq, + status: next_status, + desired_state: BackgroundAgentDesiredState::Deleted, + summary: Some(status_reason.to_string()), + pending_interaction_count, + last_event_seq: event.seq, + payload_json: serde_json::json!({ + "phase": next_status.as_str(), + "reason": "delete_requested", + }), + }, + now, + ) .await?; - Ok(result.rows_affected() > 0) + tx.commit().await?; + Ok(true) } pub async fn orphan_stale_background_agent_runs( @@ -1118,25 +1418,145 @@ WHERE supervisor_id: &str, process_lease_id: &str, ) -> anyhow::Result> { - let now = Utc::now().timestamp(); - let mut tx = self.pool.begin().await?; - let current: Option<(i64, String, String, String)> = sqlx::query_as( + self.claim_background_agent_supervisor_inner( + run_id, + supervisor_id, + process_lease_id, + /*required_version_fingerprint*/ None, + ) + .await + } + + pub async fn claim_background_agent_supervisor_compatible( + &self, + run_id: &str, + supervisor_id: &str, + process_lease_id: &str, + required_version_fingerprint: &str, + ) -> anyhow::Result> { + self.claim_background_agent_supervisor_inner( + run_id, + supervisor_id, + process_lease_id, + Some(required_version_fingerprint), + ) + .await + } + + pub async fn background_agent_admission_is_ready( + &self, + run_id: &str, + required_version_fingerprint: &str, + ) -> anyhow::Result { + let ready: Option = sqlx::query_scalar( r#" -SELECT generation, desired_state, status, retention_state +SELECT 1 FROM background_agent_runs -WHERE id = ? +WHERE + id = ? + AND version_fingerprint = ? + AND admission_identity_sha256 IS NOT NULL + AND admission_ready_at IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM background_agent_events e + WHERE + e.run_id = background_agent_runs.id + AND e.event_type IN ('agent.started', 'agent.startRecovered') + ) + AND EXISTS ( + SELECT 1 + FROM background_agent_execution_snapshots s + WHERE + s.run_id = background_agent_runs.id + AND s.snapshot_kind = 'initial_execution_context' + ) + AND EXISTS ( + SELECT 1 + FROM background_agent_status_snapshots s + WHERE s.run_id = background_agent_runs.id + ) "#, ) .bind(run_id) - .fetch_optional(&mut *tx) + .bind(required_version_fingerprint) + .fetch_optional(self.pool.as_ref()) + .await?; + Ok(ready.is_some()) + } + + pub async fn get_background_agent_initial_execution_snapshot( + &self, + run_id: &str, + ) -> anyhow::Result> { + let row = sqlx::query_as::<_, BackgroundAgentExecutionSnapshotRow>( + r#" +SELECT + id, + run_id, + seq, + snapshot_kind, + payload_json, + recovery_policy, + config_fingerprint, + created_at +FROM background_agent_execution_snapshots +WHERE run_id = ? AND snapshot_kind = 'initial_execution_context' +ORDER BY seq ASC +LIMIT 1 + "#, + ) + .bind(run_id) + .fetch_optional(self.pool.as_ref()) .await?; - let Some((current_generation, desired_state, status, retention_state)) = current else { + row.map(BackgroundAgentExecutionSnapshot::try_from) + .transpose() + } + + async fn claim_background_agent_supervisor_inner( + &self, + run_id: &str, + supervisor_id: &str, + process_lease_id: &str, + required_version_fingerprint: Option<&str>, + ) -> anyhow::Result> { + let now = Utc::now().timestamp(); + let mut tx = self.pool.begin().await?; + let current: Option<(i64, String, String, String, Option, Option)> = + sqlx::query_as( + r#" +SELECT + generation, + desired_state, + status, + retention_state, + admission_identity_sha256, + admission_ready_at +FROM background_agent_runs +WHERE id = ? + "#, + ) + .bind(run_id) + .fetch_optional(&mut *tx) + .await?; + let Some(( + current_generation, + desired_state, + status, + retention_state, + admission_identity_sha256, + admission_ready_at, + )) = current + else { tx.rollback().await?; return Ok(None); }; + let admission_ready = required_version_fingerprint.is_none() + || (admission_identity_sha256.is_some() && admission_ready_at.is_some()); let eligible = desired_state == BackgroundAgentDesiredState::Running.as_str() && retention_state == crate::BackgroundAgentRetentionState::Active.as_str() - && matches!(status.as_str(), "queued" | "orphaned"); + && matches!(status.as_str(), "queued" | "orphaned") + && admission_ready; if !eligible { tx.rollback().await?; return Ok(None); @@ -1160,6 +1580,33 @@ WHERE AND desired_state = ? AND retention_state = ? AND status IN ('queued', 'orphaned') + AND ( + ? IS NULL + OR ( + version_fingerprint = ? + AND admission_identity_sha256 IS NOT NULL + AND admission_ready_at IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM background_agent_events e + WHERE + e.run_id = background_agent_runs.id + AND e.event_type IN ('agent.started', 'agent.startRecovered') + ) + AND EXISTS ( + SELECT 1 + FROM background_agent_execution_snapshots s + WHERE + s.run_id = background_agent_runs.id + AND s.snapshot_kind = 'initial_execution_context' + ) + AND EXISTS ( + SELECT 1 + FROM background_agent_status_snapshots s + WHERE s.run_id = background_agent_runs.id + ) + ) + ) "#, ) .bind(supervisor_id) @@ -1173,6 +1620,8 @@ WHERE .bind(current_generation) .bind(BackgroundAgentDesiredState::Running.as_str()) .bind(crate::BackgroundAgentRetentionState::Active.as_str()) + .bind(required_version_fingerprint) + .bind(required_version_fingerprint) .execute(&mut *tx) .await?; @@ -1298,6 +1747,11 @@ WHERE run_id = ? AND supervisor_id = ? AND generation = ? /*attempt*/ None, &serde_json::json!({ "supervisorId": params.supervisor_id, + "pid": params.pid, + "pgid": params.pgid, + "jobId": params.job_id, + "startToken": params.start_token, + "stderrLogPath": params.stderr_log_path, }), now, ) @@ -1494,6 +1948,406 @@ WHERE idempotency_key = ? } } +async fn ensure_background_agent_capacity_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + max_active_runs: i64, +) -> anyhow::Result<()> { + let active_run_count: i64 = sqlx::query_scalar( + r#" +SELECT COUNT(*) +FROM background_agent_runs +WHERE + status IN ( + 'queued', + 'starting', + 'running', + 'waiting_on_approval', + 'waiting_on_user', + 'stopping', + 'orphaned' + ) + AND ( + status = 'stopping' + OR ( + status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') + AND supervisor_id IS NOT NULL + ) + OR ( + status IN ('queued', 'orphaned') + AND desired_state = 'running' + AND retention_state = 'active' + AND admission_identity_sha256 IS NOT NULL + AND admission_ready_at IS NOT NULL + ) + ) + "#, + ) + .fetch_one(&mut **tx) + .await?; + if active_run_count >= max_active_runs { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED}: \ + {active_run_count} live or recoverable run(s), max {max_active_runs}" + ); + } + Ok(()) +} + +async fn insert_background_agent_run_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + params: &BackgroundAgentRunCreateParams, + now: i64, +) -> anyhow::Result<()> { + let spawn_linkage_json = params + .spawn_linkage_json + .as_ref() + .map(redact_state_json_string) + .transpose()?; + let insert_result = sqlx::query( + r#" +INSERT INTO background_agent_runs ( + id, + idempotency_key, + request_id, + source, + prompt_snapshot_ref, + input_snapshot_ref, + thread_id, + thread_store_kind, + thread_store_id, + rollout_path, + parent_thread_id, + parent_agent_run_id, + spawn_linkage_json, + auth_profile_ref, + desired_state, + status, + status_reason, + config_fingerprint, + version_fingerprint, + retention_state, + created_at, + updated_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + "#, + ) + .bind(params.id.as_str()) + .bind(params.idempotency_key.as_deref()) + .bind(params.request_id.as_deref().map(redact_state_string)) + .bind(params.source.as_str()) + .bind(params.prompt_snapshot_ref.as_str()) + .bind( + params + .input_snapshot_ref + .as_deref() + .map(redact_state_string), + ) + .bind(params.thread_id.as_deref()) + .bind(params.thread_store_kind.as_str()) + .bind(params.thread_store_id.as_deref()) + .bind(params.rollout_path.as_deref()) + .bind(params.parent_thread_id.as_deref()) + .bind(params.parent_agent_run_id.as_deref()) + .bind(spawn_linkage_json.as_deref()) + .bind(params.auth_profile_ref.as_deref()) + .bind(BackgroundAgentDesiredState::Running.as_str()) + .bind(BackgroundAgentRunStatus::Queued.as_str()) + .bind(params.status_reason.as_deref().map(redact_state_string)) + .bind(params.config_fingerprint.as_deref()) + .bind(params.version_fingerprint.as_deref()) + .bind(crate::BackgroundAgentRetentionState::Active.as_str()) + .bind(now) + .bind(now) + .execute(&mut **tx) + .await; + if let Err(err) = insert_result { + if is_background_agent_unique_constraint_violation(&err) { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + request identity conflicts with an existing background agent" + ); + } + return Err(err.into()); + } + Ok(()) +} + +fn background_agent_admission_identity_sha256( + params: &BackgroundAgentRunCreateParams, + start_event_payload_json: &serde_json::Value, + execution_snapshot_params: &BackgroundAgentExecutionSnapshotParams, +) -> anyhow::Result { + let identity = serde_json::json!({ + "idempotencyKey": params.idempotency_key, + "requestId": params.request_id, + "source": params.source, + "promptSnapshotRef": params.prompt_snapshot_ref, + "inputSnapshotRef": params.input_snapshot_ref, + "threadId": params.thread_id, + "threadStoreKind": params.thread_store_kind, + "threadStoreId": params.thread_store_id, + "rolloutPath": params.rollout_path, + "parentThreadId": params.parent_thread_id, + "parentAgentRunId": params.parent_agent_run_id, + "spawnLinkage": params.spawn_linkage_json, + "authProfileRef": params.auth_profile_ref, + "configFingerprint": params.config_fingerprint, + "versionFingerprint": params.version_fingerprint, + "startEvent": start_event_payload_json, + "executionSnapshot": { + "snapshotKind": execution_snapshot_params.snapshot_kind, + "payload": execution_snapshot_params.payload_json, + "recoveryPolicy": execution_snapshot_params.recovery_policy, + "configFingerprint": execution_snapshot_params.config_fingerprint, + }, + }); + Ok(format!( + "{:x}", + Sha256::digest(serde_json::to_vec(&identity)?) + )) +} + +fn background_agent_start_event_payload( + payload_json: &serde_json::Value, + admission_identity_sha256: &str, +) -> serde_json::Value { + let mut payload_json = payload_json.clone(); + if let Some(payload) = payload_json.as_object_mut() { + payload.insert( + "admissionIdentitySha256".to_string(), + serde_json::Value::String(admission_identity_sha256.to_string()), + ); + } + payload_json +} + +async fn insert_background_agent_execution_snapshot_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + params: &BackgroundAgentExecutionSnapshotParams, + now: i64, +) -> anyhow::Result { + let payload_json = redact_state_json_string(¶ms.payload_json)?; + let seq: i64 = sqlx::query_scalar( + "SELECT COALESCE(MAX(seq), 0) + 1 FROM background_agent_execution_snapshots WHERE run_id = ?", + ) + .bind(params.run_id.as_str()) + .fetch_one(&mut **tx) + .await?; + let id = sqlx::query( + r#" +INSERT INTO background_agent_execution_snapshots ( + run_id, + seq, + snapshot_kind, + payload_json, + recovery_policy, + config_fingerprint, + created_at +) VALUES (?, ?, ?, ?, ?, ?, ?) + "#, + ) + .bind(params.run_id.as_str()) + .bind(seq) + .bind(params.snapshot_kind.as_str()) + .bind(payload_json) + .bind(params.recovery_policy.as_str()) + .bind(params.config_fingerprint.as_deref()) + .bind(now) + .execute(&mut **tx) + .await? + .last_insert_rowid(); + sqlx::query( + r#" +UPDATE background_agent_runs +SET last_snapshot_seq = ?, updated_at = ? +WHERE id = ? + "#, + ) + .bind(seq) + .bind(now) + .bind(params.run_id.as_str()) + .execute(&mut **tx) + .await?; + Ok(id) +} + +async fn recover_or_validate_background_agent_initial_state_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + run_id: &str, + admission_identity_sha256: &str, + start_event_payload_json: &serde_json::Value, + execution_snapshot_params: &BackgroundAgentExecutionSnapshotParams, +) -> anyhow::Result<(BackgroundAgentEvent, i64)> { + let mut execution_snapshot_params = execution_snapshot_params.clone(); + execution_snapshot_params.run_id = run_id.to_string(); + let stored_identity: Option = sqlx::query_scalar( + "SELECT admission_identity_sha256 FROM background_agent_runs WHERE id = ?", + ) + .bind(run_id) + .fetch_one(&mut **tx) + .await?; + if let Some(stored_identity) = stored_identity.as_deref() + && stored_identity != admission_identity_sha256 + { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + idempotency key is already bound to different prompt or execution context" + ); + } + + let existing_event = sqlx::query_as::<_, BackgroundAgentEventRow>( + r#" +SELECT id, run_id, seq, event_type, payload_json, created_at +FROM background_agent_events +WHERE + run_id = ? + AND event_type IN ('agent.started', 'agent.startRecovered') +ORDER BY + CASE WHEN event_type = 'agent.started' THEN 0 ELSE 1 END, + seq ASC +LIMIT 1 + "#, + ) + .bind(run_id) + .fetch_optional(&mut **tx) + .await? + .map(BackgroundAgentEvent::try_from) + .transpose()?; + let proposed_event_payload = + background_agent_start_event_payload(start_event_payload_json, admission_identity_sha256); + let now = Utc::now().timestamp(); + let event = match existing_event { + Some(event) + if stored_identity.is_some() + || legacy_background_agent_start_event_matches( + &event.payload_json, + start_event_payload_json, + ) => + { + event + } + Some(_) => { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + idempotency key is already bound to a different start event" + ); + } + None => { + super::events::append_background_agent_event_in_tx( + tx, + run_id, + "agent.started", + &proposed_event_payload, + now, + ) + .await? + } + }; + + let existing_snapshot = sqlx::query_as::<_, (i64, String, String, Option)>( + r#" +SELECT id, payload_json, recovery_policy, config_fingerprint +FROM background_agent_execution_snapshots +WHERE run_id = ? AND snapshot_kind = 'initial_execution_context' +ORDER BY seq DESC +LIMIT 1 + "#, + ) + .bind(run_id) + .fetch_optional(&mut **tx) + .await?; + let execution_snapshot_id = match existing_snapshot { + Some((id, payload_json, recovery_policy, config_fingerprint)) => { + let payload_json: serde_json::Value = serde_json::from_str(payload_json.as_str())?; + if stored_identity.is_none() + && (!legacy_background_agent_execution_snapshot_matches( + &payload_json, + &execution_snapshot_params.payload_json, + ) || recovery_policy != execution_snapshot_params.recovery_policy + || (config_fingerprint.is_some() + && config_fingerprint != execution_snapshot_params.config_fingerprint)) + { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + idempotency key is already bound to a different execution snapshot" + ); + } + id + } + None => { + insert_background_agent_execution_snapshot_in_tx(tx, &execution_snapshot_params, now) + .await? + } + }; + + let status_snapshot_exists: Option = + sqlx::query_scalar("SELECT 1 FROM background_agent_status_snapshots WHERE run_id = ?") + .bind(run_id) + .fetch_optional(&mut **tx) + .await?; + if status_snapshot_exists.is_none() { + let run_state = sqlx::query_as::<_, (String, String)>( + "SELECT status, desired_state FROM background_agent_runs WHERE id = ?", + ) + .bind(run_id) + .fetch_one(&mut **tx) + .await?; + let status = BackgroundAgentRunStatus::parse(run_state.0.as_str())?; + super::snapshots::upsert_background_agent_status_snapshot_in_tx( + tx, + &BackgroundAgentStatusSnapshotParams { + run_id: run_id.to_string(), + seq: event.seq, + status, + desired_state: BackgroundAgentDesiredState::parse(run_state.1.as_str())?, + summary: Some(status.as_str().to_string()), + pending_interaction_count: 0, + last_event_seq: event.seq, + payload_json: serde_json::json!({"phase": status.as_str()}), + }, + now, + ) + .await?; + } + sqlx::query( + r#" +UPDATE background_agent_runs +SET admission_identity_sha256 = ?, admission_ready_at = COALESCE(admission_ready_at, ?) +WHERE id = ? + "#, + ) + .bind(admission_identity_sha256) + .bind(now) + .bind(run_id) + .execute(&mut **tx) + .await?; + Ok((event, execution_snapshot_id)) +} + +fn legacy_background_agent_start_event_matches( + stored: &serde_json::Value, + proposed: &serde_json::Value, +) -> bool { + ["prompt", "cwd", "promptSnapshotRef", "initialGoalObjective"] + .into_iter() + .all(|key| stored.get(key) == proposed.get(key)) +} + +fn legacy_background_agent_execution_snapshot_matches( + stored: &serde_json::Value, + proposed: &serde_json::Value, +) -> bool { + let (Some(stored), Some(proposed)) = (stored.as_object(), proposed.as_object()) else { + return stored == proposed; + }; + stored.iter().all(|(key, value)| { + matches!( + key.as_str(), + "authProfileRef" | "packageFingerprint" | "configFingerprint" | "versionFingerprint" + ) || proposed.get(key) == Some(value) + }) +} + async fn validate_existing_background_agent_admission_in_tx( tx: &mut sqlx::Transaction<'_, Sqlite>, idempotency_key: &str, @@ -1505,6 +2359,7 @@ async fn validate_existing_background_agent_admission_in_tx( String, Option, String, + String, Option, Option, String, @@ -1523,6 +2378,7 @@ SELECT id, request_id, source, + prompt_snapshot_ref, input_snapshot_ref, thread_id, thread_store_kind, @@ -1545,6 +2401,7 @@ WHERE idempotency_key = ? existing_id, request_id, source, + prompt_snapshot_ref, input_snapshot_ref, thread_id, thread_store_kind, @@ -1567,6 +2424,7 @@ WHERE idempotency_key = ? .transpose()?; let identity_matches = request_id == params.request_id.as_deref().map(redact_state_string) && source == params.source + && prompt_snapshot_ref == params.prompt_snapshot_ref && input_snapshot_ref == params .input_snapshot_ref @@ -1579,7 +2437,7 @@ WHERE idempotency_key = ? && parent_thread_id == params.parent_thread_id && parent_agent_run_id == params.parent_agent_run_id && spawn_linkage_json == requested_spawn_linkage_json - && auth_profile_ref == params.auth_profile_ref.as_deref().map(redact_state_string) + && auth_profile_ref == params.auth_profile_ref && config_fingerprint == params.config_fingerprint && version_fingerprint == params.version_fingerprint; if !identity_matches { diff --git a/codex-rs/state/src/runtime/background_agents/tests.rs b/codex-rs/state/src/runtime/background_agents/tests.rs index 9c316b2c93..5fbd727993 100644 --- a/codex-rs/state/src/runtime/background_agents/tests.rs +++ b/codex-rs/state/src/runtime/background_agents/tests.rs @@ -90,7 +90,7 @@ fn admission_params( idempotency_key: Some(idempotency_key.to_string()), request_id: Some(format!("request-{idempotency_key}")), source: "admission-test".to_string(), - prompt_snapshot_ref: format!("inline:{id}:prompt"), + prompt_snapshot_ref: format!("inline:{idempotency_key}:prompt"), input_snapshot_ref: None, thread_id: Some(format!("thread-{idempotency_key}")), thread_store_kind: "background-agent".to_string(), @@ -106,6 +106,53 @@ fn admission_params( } } +async fn admit_run( + runtime: &StateRuntime, + params: &BackgroundAgentRunCreateParams, + max_active_runs: i64, +) -> anyhow::Result<(BackgroundAgentRun, bool)> { + let prompt = format!( + "prompt for {}", + params + .idempotency_key + .as_deref() + .unwrap_or(params.id.as_str()) + ); + let snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: params.id.clone(), + snapshot_kind: "initial_execution_context".to_string(), + payload_json: json!({ + "cwd": "/tmp/admission-test", + "workspaceRoots": ["/tmp/admission-test"], + "permissionProfile": {"type": "managed"}, + "networkPolicy": "restricted", + "model": "test-model", + "provider": "test-provider", + "serviceTier": "default", + "configFingerprint": params.config_fingerprint, + "versionFingerprint": params.version_fingerprint, + "packageFingerprint": "codex-state:test", + "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", + }), + recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), + config_fingerprint: params.config_fingerprint.clone(), + }; + let (run, created, _, _, _) = runtime + .admit_background_agent_run( + params, + &json!({ + "cwd": "/tmp/admission-test", + "prompt": prompt, + "promptSnapshotRef": params.prompt_snapshot_ref, + "initialGoalObjective": "test admission", + }), + &snapshot_params, + max_active_runs, + ) + .await?; + Ok((run, created)) +} + #[tokio::test] async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; @@ -116,7 +163,7 @@ async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { idempotency_key: Some("idem-1".to_string()), request_id: Some("req-1".to_string()), source: "cli".to_string(), - prompt_snapshot_ref: "prompt://duplicate".to_string(), + prompt_snapshot_ref: "prompt://run-1".to_string(), input_snapshot_ref: Some("input://run-1".to_string()), thread_id: Some("thread-1".to_string()), thread_store_kind: "local".to_string(), @@ -145,26 +192,83 @@ async fn background_agent_admission_create_or_adopt_is_atomic_and_receipted() -> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; let first_params = admission_params("admitted-1", "admission-key", "profile-a"); - let (first, created) = runtime - .admit_background_agent_run(&first_params, /*max_active_runs*/ 2) - .await?; + let (first, created) = + admit_run(runtime.as_ref(), &first_params, /*max_active_runs*/ 2).await?; assert!(created); let retry_params = admission_params("admitted-retry", "admission-key", "profile-a"); - let (retry, created) = runtime - .admit_background_agent_run(&retry_params, /*max_active_runs*/ 2) - .await?; + let (retry, created) = + admit_run(runtime.as_ref(), &retry_params, /*max_active_runs*/ 2).await?; assert!(!created); assert_eq!(retry.id, first.id); assert_eq!(runtime.list_background_agent_runs(None).await?.len(), 1); let events = runtime .list_background_agent_events_after(first.id.as_str(), None, None) .await?; - assert_eq!(events.len(), 1); + assert_eq!(events.len(), 2); assert_eq!(events[0].event_type, "agent.admitted"); + assert_eq!(events[1].event_type, "agent.started"); + assert_ne!( + events[0] + .payload_json + .get("receiptKey") + .and_then(serde_json::Value::as_str), + Some("admission:admission-key") + ); + assert!( + runtime + .get_latest_background_agent_execution_snapshot(first.id.as_str()) + .await? + .is_some() + ); + assert!( + runtime + .get_background_agent_status_snapshot(first.id.as_str()) + .await? + .is_some() + ); + runtime + .create_background_agent_execution_snapshot(&BackgroundAgentExecutionSnapshotParams { + run_id: first.id.clone(), + snapshot_kind: "worker_thread_bound".to_string(), + payload_json: json!({"threadId": "thread-after-admission"}), + recovery_policy: "resume_or_orphan".to_string(), + config_fingerprint: first.config_fingerprint.clone(), + }) + .await?; assert_eq!( - events[0].payload_json.get("receiptKey"), - Some(&json!("admission:admission-key")) + runtime + .get_background_agent_initial_execution_snapshot(first.id.as_str()) + .await? + .expect("initial execution context must remain authoritative") + .snapshot_kind, + "initial_execution_context" + ); + sqlx::query( + "DELETE FROM background_agent_execution_snapshots \ + WHERE run_id = ? AND snapshot_kind = 'initial_execution_context'", + ) + .bind(first.id.as_str()) + .execute(runtime.pool.as_ref()) + .await?; + assert!( + !runtime + .background_agent_admission_is_ready( + first.id.as_str(), + "codewith.background-agent.admission.v1", + ) + .await? + ); + assert!( + runtime + .claim_background_agent_supervisor_compatible( + first.id.as_str(), + "supervisor-after-corruption", + "lease-after-corruption", + "codewith.background-agent.admission.v1", + ) + .await? + .is_none() ); Ok(()) } @@ -172,20 +276,20 @@ async fn background_agent_admission_create_or_adopt_is_atomic_and_receipted() -> #[tokio::test] async fn background_agent_admission_rejects_idempotency_identity_mismatch() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; - runtime - .admit_background_agent_run( - &admission_params("admitted-1", "admission-key", "profile-a"), - 2, - ) - .await?; + admit_run( + runtime.as_ref(), + &admission_params("admitted-1", "admission-key", "profile-a"), + 2, + ) + .await?; - let error = runtime - .admit_background_agent_run( - &admission_params("admitted-2", "admission-key", "profile-b"), - 2, - ) - .await - .expect_err("profile mismatch must not adopt the existing run"); + let error = admit_run( + runtime.as_ref(), + &admission_params("admitted-2", "admission-key", "profile-b"), + 2, + ) + .await + .expect_err("profile mismatch must not adopt the existing run"); assert!( error @@ -196,32 +300,89 @@ async fn background_agent_admission_rejects_idempotency_identity_mismatch() -> a Ok(()) } +#[tokio::test] +async fn background_agent_admission_preserves_opaque_identity_values() -> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let idempotency_key = "sk-opaque-idempotency-key"; + let auth_profile_ref = "sk-opaque-profile-alias"; + let (run, created) = admit_run( + runtime.as_ref(), + &admission_params("opaque-admission", idempotency_key, auth_profile_ref), + 2, + ) + .await?; + + assert!(created); + assert_eq!(run.idempotency_key.as_deref(), Some(idempotency_key)); + assert_eq!(run.auth_profile_ref.as_deref(), Some(auth_profile_ref)); + assert_eq!( + runtime + .get_background_agent_run_by_idempotency_key(idempotency_key) + .await? + .map(|run| run.id), + Some(run.id) + ); + Ok(()) +} + #[tokio::test] async fn background_agent_admission_counts_only_live_or_recoverable_runs() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; - let (first, _) = runtime - .admit_background_agent_run( - &admission_params("admitted-1", "admission-key-1", "profile-a"), - 1, - ) - .await?; + create_run_with_id(runtime.as_ref(), "legacy-incompatible").await?; + assert_eq!( + runtime.count_background_agent_runs_by_status().await?, + Vec::<(BackgroundAgentRunStatus, i64)>::new() + ); + let (first, _) = admit_run( + runtime.as_ref(), + &admission_params("admitted-1", "admission-key-1", "profile-a"), + 1, + ) + .await?; + let generation = runtime + .claim_background_agent_supervisor(first.id.as_str(), "supervisor-1", "lease-1") + .await? + .expect("admitted run should be claimable"); assert!( runtime .request_background_agent_stop_for_generation( first.id.as_str(), - None, - 0, + Some("supervisor-1"), + generation, "capacity test stop", &json!({"reason": "capacity_test"}), ) .await? ); - let (_, created) = runtime - .admit_background_agent_run( - &admission_params("admitted-2", "admission-key-2", "profile-a"), - 1, - ) - .await?; + let error = admit_run( + runtime.as_ref(), + &admission_params("admitted-2", "admission-key-2", "profile-a"), + 1, + ) + .await + .expect_err("claimed stopping run must consume capacity"); + assert!( + error + .to_string() + .contains("background_agent_admission_capacity_exceeded") + ); + assert!( + runtime + .finalize_stopped_background_agent_process( + first.id.as_str(), + "supervisor-1", + generation, + "capacity test process stopped", + &json!({"reason": "capacity_test_process_stopped"}), + ) + .await? + ); + let (_, created) = admit_run( + runtime.as_ref(), + &admission_params("admitted-2", "admission-key-2", "profile-a"), + 1, + ) + .await?; assert!(created); runtime @@ -231,13 +392,13 @@ async fn background_agent_admission_counts_only_live_or_recoverable_runs() -> an Some("recoverable orphan"), ) .await?; - let error = runtime - .admit_background_agent_run( - &admission_params("admitted-3", "admission-key-3", "profile-a"), - 1, - ) - .await - .expect_err("recoverable orphan must consume capacity"); + let error = admit_run( + runtime.as_ref(), + &admission_params("admitted-3", "admission-key-3", "profile-a"), + 1, + ) + .await + .expect_err("recoverable orphan must consume capacity"); assert!( error .to_string() @@ -266,7 +427,7 @@ async fn background_agent_lifecycle_receipts_dedupe_redact_and_bound_diagnostics &diagnostics, ) .await?; - let retry = runtime + let conflict = runtime .append_background_agent_lifecycle_receipt( run.id.as_str(), "agent.testReceipt", @@ -275,6 +436,41 @@ async fn background_agent_lifecycle_receipts_dedupe_redact_and_bound_diagnostics Some(2), &diagnostics, ) + .await + .expect_err("receipt attempt mismatch must fail"); + assert!( + conflict + .to_string() + .contains("background agent lifecycle receipt identity mismatch") + ); + let redaction_collision = runtime + .append_background_agent_lifecycle_receipt( + run.id.as_str(), + "agent.testReceipt", + "test-receipt", + 1, + Some(1), + &json!({ + "apiKey": "sk-different-secret-value", + "blob": "x".repeat(8 * 1024), + }), + ) + .await + .expect_err("distinct raw diagnostics must not collapse after redaction"); + assert!( + redaction_collision + .to_string() + .contains("background agent lifecycle receipt identity mismatch") + ); + let retry = runtime + .append_background_agent_lifecycle_receipt( + run.id.as_str(), + "agent.testReceipt", + "test-receipt", + 1, + Some(1), + &diagnostics, + ) .await?; assert_eq!(retry.id, first.id); @@ -289,6 +485,103 @@ async fn background_agent_lifecycle_receipts_dedupe_redact_and_bound_diagnostics .and_then(serde_json::Value::as_bool), Some(true) ); + assert_eq!( + runtime + .compact_background_agent_events_before_seq("run-1", first.seq + 1) + .await?, + 1 + ); + let compacted_retry = runtime + .append_background_agent_lifecycle_receipt( + run.id.as_str(), + "agent.testReceipt", + "test-receipt", + 1, + Some(1), + &diagnostics, + ) + .await?; + assert_eq!(compacted_retry, first); + + let oversized_receipt_key = "x".repeat(300); + let error = runtime + .append_background_agent_lifecycle_receipt( + run.id.as_str(), + "agent.oversizedReceipt", + oversized_receipt_key.as_str(), + 1, + Some(1), + &json!({}), + ) + .await + .expect_err("caller-controlled receipt keys must be bounded"); + assert!( + error + .to_string() + .contains("background agent lifecycle receipt key exceeds") + ); + Ok(()) +} + +#[tokio::test] +async fn background_agent_terminal_status_receipt_replays_after_commit_and_compaction() +-> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let run = create_run(runtime.as_ref()).await?; + let generation = runtime + .claim_background_agent_supervisor(run.id.as_str(), "supervisor-1", "lease-1") + .await? + .expect("run should be claimable"); + let event_payload = json!({"outcome": "completed"}); + let status_payload = json!({"phase": "completed"}); + let params = || BackgroundAgentStatusEventForSupervisorParams { + run_id: run.id.as_str(), + supervisor_id: "supervisor-1", + generation, + status: BackgroundAgentRunStatus::Completed, + status_reason: Some("worker completed"), + event_type: "agent.completed", + event_payload_json: &event_payload, + summary: Some("Completed"), + pending_interaction_count: 0, + status_payload_json: &status_payload, + }; + + let first = runtime + .append_background_agent_status_event_for_supervisor(params()) + .await? + .expect("current generation should complete"); + let conflict = runtime + .append_background_agent_status_event_for_supervisor( + BackgroundAgentStatusEventForSupervisorParams { + status_reason: Some("different terminal outcome"), + ..params() + }, + ) + .await + .expect_err("terminal receipt replay must bind the full projected operation"); + assert!( + conflict + .to_string() + .contains("background agent lifecycle receipt identity mismatch") + ); + let retry = runtime + .append_background_agent_status_event_for_supervisor(params()) + .await? + .expect("terminal receipt should replay after an ambiguous acknowledgement"); + assert_eq!(retry, first); + + assert!( + runtime + .compact_background_agent_events_before_seq(run.id.as_str(), first.seq + 1) + .await? + > 0 + ); + let compacted_retry = runtime + .append_background_agent_status_event_for_supervisor(params()) + .await? + .expect("terminal receipt should survive event compaction"); + assert_eq!(compacted_retry, first); Ok(()) } @@ -1051,6 +1344,32 @@ async fn stale_generation_cannot_update_status_or_create_interactions_after_recl stderr_log_path: Some("/tmp/run-1.stderr.log"), }) .await?; + let handle_conflict = runtime + .record_background_agent_execution_handle(BackgroundAgentExecutionHandleParams { + run_id: "run-1", + supervisor_id: "supervisor-1", + generation: first_generation, + pid: Some(200), + pgid: Some(200), + job_id: Some("different-job"), + start_token: Some("different-start"), + stderr_log_path: Some("/tmp/different.stderr.log"), + }) + .await + .expect_err("execution handle receipt must bind the exact operation"); + assert!( + handle_conflict + .to_string() + .contains("background agent lifecycle receipt identity mismatch") + ); + assert_eq!( + runtime + .get_background_agent_run("run-1") + .await? + .expect("run should remain current") + .pid, + Some(100) + ); assert_eq!( runtime .orphan_stale_background_agent_runs(Duration::ZERO) @@ -1093,6 +1412,73 @@ async fn stale_generation_cannot_update_status_or_create_interactions_after_recl .await? .is_none() ); + assert!( + runtime + .append_background_agent_event_for_supervisor( + "run-1", + "supervisor-1", + first_generation, + "agent.staleEvent", + &json!({"generation": first_generation}), + /*allow_terminal_current*/ false, + ) + .await? + .is_none() + ); + assert!( + runtime + .create_background_agent_execution_snapshot_for_supervisor( + &BackgroundAgentExecutionSnapshotParams { + run_id: "run-1".to_string(), + snapshot_kind: "stale_generation".to_string(), + payload_json: json!({"generation": first_generation}), + recovery_policy: "resume_or_orphan".to_string(), + config_fingerprint: None, + }, + "supervisor-1", + first_generation, + ) + .await? + .is_none() + ); + assert!( + runtime + .upsert_background_agent_status_snapshot_for_supervisor( + &BackgroundAgentStatusSnapshotParams { + run_id: "run-1".to_string(), + seq: 99, + status: BackgroundAgentRunStatus::Completed, + desired_state: BackgroundAgentDesiredState::Running, + summary: Some("stale completion".to_string()), + pending_interaction_count: 0, + last_event_seq: 99, + payload_json: json!({"generation": first_generation}), + }, + "supervisor-1", + first_generation, + ) + .await? + .is_none() + ); + assert!( + runtime + .append_background_agent_status_event_for_supervisor( + BackgroundAgentStatusEventForSupervisorParams { + run_id: "run-1", + supervisor_id: "supervisor-1", + generation: first_generation, + status: BackgroundAgentRunStatus::Completed, + status_reason: Some("stale completion"), + event_type: "agent.completed", + event_payload_json: &json!({"generation": first_generation}), + summary: Some("Completed"), + pending_interaction_count: 0, + status_payload_json: &json!({"phase": "completed"}), + }, + ) + .await? + .is_none() + ); let run = runtime .get_background_agent_run("run-1") From 8f96a4d3d40b281ef2f36e86a1f5ed3579562b9c Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 21:22:16 +0300 Subject: [PATCH 12/15] fix(state): share execution snapshot lookup --- codex-rs/state/src/runtime/background_agents/snapshots.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/state/src/runtime/background_agents/snapshots.rs b/codex-rs/state/src/runtime/background_agents/snapshots.rs index d06aa65e13..1e532e179b 100644 --- a/codex-rs/state/src/runtime/background_agents/snapshots.rs +++ b/codex-rs/state/src/runtime/background_agents/snapshots.rs @@ -170,7 +170,7 @@ LIMIT 1 .transpose() } - async fn get_background_agent_execution_snapshot( + pub(super) async fn get_background_agent_execution_snapshot( &self, snapshot_id: i64, ) -> anyhow::Result> { From e71bc91678d90cff8f5cc1e7dceab1df3fe26709 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 21:50:13 +0300 Subject: [PATCH 13/15] fix(agent): close durable admission CI failures --- codex-rs/app-server/Cargo.toml | 2 +- .../background_agent_live.rs | 72 +++++++++---------- codex-rs/background-agent/src/supervisor.rs | 10 ++- .../src/runtime/background_agents/runs.rs | 29 +++++--- 4 files changed, 64 insertions(+), 49 deletions(-) diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index aa42c6ef18..f65d07786c 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -88,6 +88,7 @@ futures = { workspace = true } iana-time-zone = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2 = { workspace = true } tempfile = { workspace = true } thiserror = { workspace = true } time = { workspace = true } @@ -130,7 +131,6 @@ rmcp = { workspace = true, default-features = false, features = [ "transport-streamable-http-server", ] } serial_test = { workspace = true } -sha2 = { workspace = true } shlex = { workspace = true } tar = { workspace = true } tokio-tungstenite = { workspace = true } diff --git a/codex-rs/app-server/src/request_processors/background_agent_live.rs b/codex-rs/app-server/src/request_processors/background_agent_live.rs index dd99a5e3bd..f2e03544df 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_live.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_live.rs @@ -3043,21 +3043,22 @@ async fn run_background_agent_worker( initial_execution_payload.as_ref(), ) .await?; + let execution_snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: run.id.clone(), + snapshot_kind: "worker_thread_bound".to_string(), + payload_json: json!({ + "threadId": thread_id.to_string(), + "sessionId": session_id_string, + "rolloutPath": rollout_path, + }), + recovery_policy: "resume_or_orphan".to_string(), + config_fingerprint: run.config_fingerprint.clone(), + }; retry_transient_sqlite_busy("create background agent execution snapshot", || { context .state_db .create_background_agent_execution_snapshot_for_supervisor( - &BackgroundAgentExecutionSnapshotParams { - run_id: run.id.clone(), - snapshot_kind: "worker_thread_bound".to_string(), - payload_json: json!({ - "threadId": thread_id.to_string(), - "sessionId": session_id_string, - "rolloutPath": rollout_path, - }), - recovery_policy: "resume_or_orphan".to_string(), - config_fingerprint: run.config_fingerprint.clone(), - }, + &execution_snapshot_params, context.supervisor_id.as_str(), generation, ) @@ -4263,35 +4264,34 @@ async fn upsert_interaction_status_snapshot( .get("type") .and_then(Value::as_str) .unwrap_or_else(|| interaction.kind.as_str()); + let status_snapshot_params = BackgroundAgentStatusSnapshotParams { + run_id: run_id.to_string(), + seq: run.last_event_seq, + status, + desired_state: run.desired_state, + summary: Some(status_summary(status, "waiting for pending interaction").to_string()), + pending_interaction_count: pending_count, + last_event_seq: run.last_event_seq, + payload_json: status_snapshot_payload( + status, + waiting_reason, + generation, + &json!({ + "interactionId": interaction.id, + "workerRequestId": interaction.worker_request_id, + "kind": interaction.kind.as_str(), + "waitingReason": waiting_reason, + "requestPayload": interaction.request_payload_json, + "timeoutAt": interaction.timeout_at.map(|value| value.timestamp()), + }), + execution_payload.as_ref(), + ), + }; retry_transient_sqlite_busy("upsert background agent waiting status snapshot", || { context .state_db .upsert_background_agent_status_snapshot_for_supervisor( - &BackgroundAgentStatusSnapshotParams { - run_id: run_id.to_string(), - seq: run.last_event_seq, - status, - desired_state: run.desired_state, - summary: Some( - status_summary(status, "waiting for pending interaction").to_string(), - ), - pending_interaction_count: pending_count, - last_event_seq: run.last_event_seq, - payload_json: status_snapshot_payload( - status, - waiting_reason, - generation, - &json!({ - "interactionId": interaction.id, - "workerRequestId": interaction.worker_request_id, - "kind": interaction.kind.as_str(), - "waitingReason": waiting_reason, - "requestPayload": interaction.request_payload_json, - "timeoutAt": interaction.timeout_at.map(|value| value.timestamp()), - }), - execution_payload.as_ref(), - ), - }, + &status_snapshot_params, context.supervisor_id.as_str(), generation, ) diff --git a/codex-rs/background-agent/src/supervisor.rs b/codex-rs/background-agent/src/supervisor.rs index c3c74098f7..fdbb1d8928 100644 --- a/codex-rs/background-agent/src/supervisor.rs +++ b/codex-rs/background-agent/src/supervisor.rs @@ -11,7 +11,6 @@ use crate::AgentSnapshotStore; use crate::AgentSupervisor; use crate::BackgroundAgentDesiredState; use crate::BackgroundAgentExecutionHandleParams; -use crate::BackgroundAgentExecutionSnapshotParams; use crate::BackgroundAgentPendingInteractionStatus; use crate::BackgroundAgentRun; use crate::BackgroundAgentRunStatus; @@ -416,6 +415,7 @@ mod tests { use super::*; use crate::AgentExecutionHandle; + use crate::BackgroundAgentExecutionSnapshotParams; #[derive(Debug, Clone, Default)] struct RecordingExecution { @@ -549,6 +549,8 @@ mod tests { .map(|event| event.event_type) .collect::>(), vec![ + "agent.admitted".to_string(), + "agent.started".to_string(), "agent.claimed".to_string(), "agent.workerStartFailed".to_string() ] @@ -654,7 +656,7 @@ mod tests { delivered.status, BackgroundAgentPendingInteractionStatus::Delivered ); - assert_eq!(snapshot.run.last_event_seq, 2); + assert_eq!(snapshot.run.last_event_seq, 4); assert_eq!( state .runtime @@ -666,6 +668,8 @@ mod tests { .map(|event| event.event_type) .collect::>(), vec![ + "agent.admitted".to_string(), + "agent.started".to_string(), "interaction.created".to_string(), "interaction.delivered".to_string() ] @@ -762,6 +766,8 @@ mod tests { .map(|event| event.event_type) .collect::>(), vec![ + "agent.admitted".to_string(), + "agent.started".to_string(), "interaction.created".to_string(), "agent.stopRequested".to_string(), "interaction.cancelled".to_string() diff --git a/codex-rs/state/src/runtime/background_agents/runs.rs b/codex-rs/state/src/runtime/background_agents/runs.rs index 1ea16e657c..4d73a07429 100644 --- a/codex-rs/state/src/runtime/background_agents/runs.rs +++ b/codex-rs/state/src/runtime/background_agents/runs.rs @@ -7,6 +7,16 @@ const BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED: &str = const BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH: &str = "background_agent_admission_identity_mismatch"; +#[derive(sqlx::FromRow)] +struct BackgroundAgentSupervisorClaimState { + generation: i64, + desired_state: String, + status: String, + retention_state: String, + admission_identity_sha256: Option, + admission_ready_at: Option, +} + impl StateRuntime { /// Returns a stable digest for opaque background-agent admission identity. pub fn background_agent_identity_sha256(bytes: &[u8]) -> String { @@ -1522,9 +1532,8 @@ LIMIT 1 ) -> anyhow::Result> { let now = Utc::now().timestamp(); let mut tx = self.pool.begin().await?; - let current: Option<(i64, String, String, String, Option, Option)> = - sqlx::query_as( - r#" + let current = sqlx::query_as::<_, BackgroundAgentSupervisorClaimState>( + r#" SELECT generation, desired_state, @@ -1535,18 +1544,18 @@ SELECT FROM background_agent_runs WHERE id = ? "#, - ) - .bind(run_id) - .fetch_optional(&mut *tx) - .await?; - let Some(( - current_generation, + ) + .bind(run_id) + .fetch_optional(&mut *tx) + .await?; + let Some(BackgroundAgentSupervisorClaimState { + generation: current_generation, desired_state, status, retention_state, admission_identity_sha256, admission_ready_at, - )) = current + }) = current else { tx.rollback().await?; return Ok(None); From 5f42ceab29ff62cdecb35a9ed2263ea66d9f4997 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 22:13:34 +0300 Subject: [PATCH 14/15] test(state): assert durable delete lifecycle receipts --- .../src/runtime/background_agents/tests.rs | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/codex-rs/state/src/runtime/background_agents/tests.rs b/codex-rs/state/src/runtime/background_agents/tests.rs index 5fbd727993..87858ce09f 100644 --- a/codex-rs/state/src/runtime/background_agents/tests.rs +++ b/codex-rs/state/src/runtime/background_agents/tests.rs @@ -1997,7 +1997,25 @@ async fn delete_request_for_claimed_run_becomes_stopping_and_stale_cancelled() - status_snapshot.summary.as_deref(), Some("stop heartbeat stale") ); - assert_eq!(status_snapshot.last_event_seq, 3); + let events = runtime + .list_background_agent_events_after("run-1", /*after_seq*/ None, /*limit*/ None) + .await?; + assert_eq!( + events + .iter() + .map(|event| event.event_type.as_str()) + .collect::>(), + vec![ + "agent.claimed", + "agent.heartbeat", + "agent.deleteRequested", + "agent.cancelled", + ] + ); + assert_eq!( + status_snapshot.last_event_seq, + events.last().expect("terminal event should exist").seq + ); Ok(()) } From 6c7c9e6c679f3e6febc3aee13c5ebca625563a59 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 23 Jul 2026 22:18:52 +0300 Subject: [PATCH 15/15] fix(app-server): qualify durable agent quota constant --- .../app-server/src/request_processors/background_agent_live.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/app-server/src/request_processors/background_agent_live.rs b/codex-rs/app-server/src/request_processors/background_agent_live.rs index f2e03544df..7341f76bda 100644 --- a/codex-rs/app-server/src/request_processors/background_agent_live.rs +++ b/codex-rs/app-server/src/request_processors/background_agent_live.rs @@ -5960,7 +5960,7 @@ done }, &start_event_payload, &execution_snapshot_params, - DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, + codex_background_agent::DEFAULT_MAX_ACTIVE_BACKGROUND_AGENT_RUNS, ) .await?; Ok(())