diff --git a/codex-rs/app-server-protocol/src/protocol/v2/agent.rs b/codex-rs/app-server-protocol/src/protocol/v2/agent.rs index 360a87a18d..c26ef231ba 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/agent.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/agent.rs @@ -91,6 +91,11 @@ pub enum AgentLifecycleEffect { #[ts(export_to = "v2/")] pub struct AgentRun { pub agent_id: String, + // Stable identity of the admitted idempotency key, not the key itself: + // admission persists only a one-way SHA-256 digest so a caller-supplied + // key can never be recovered from local state. Kept under the original + // wire name to avoid an unversioned protocol break; use it for + // correlation, never as an echo of the submitted key. #[ts(type = "string | null")] pub idempotency_key: Option, #[ts(type = "string | null")] 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 d3118d8d7d..ae061bd732 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; @@ -50,12 +51,18 @@ use codex_app_server_protocol::WorktreeReconcileResponse; use codex_app_server_protocol::WorktreeReleaseParams; use codex_app_server_protocol::WorktreeReleaseResponse; use codex_app_server_protocol::WorktreeSessionMode; +use codex_backend_client::Client as BackendClient; 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::BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT; use codex_background_agent::BackgroundAgentDesiredState; use codex_background_agent::BackgroundAgentEvent; use codex_background_agent::BackgroundAgentExecutionHandleParams; +use codex_background_agent::BackgroundAgentExecutionSnapshot; use codex_background_agent::BackgroundAgentExecutionSnapshotParams; use codex_background_agent::BackgroundAgentPendingInteraction; use codex_background_agent::BackgroundAgentPendingInteractionCreateParams; @@ -94,6 +101,7 @@ use codex_git_utils::worktree_has_commits_after; use codex_protocol::approvals::ElicitationAction; use codex_protocol::config_types::SandboxMode; use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::Op; @@ -140,6 +148,9 @@ const BACKGROUND_AGENT_WORKER_STDERR_DIR: &str = "workers"; const MANAGED_WORKTREE_CLEANUP_BATCH_LIMIT: u32 = 20; const MANAGED_WORKTREE_CLEANUP_RETRY_DELAY: ChronoDuration = ChronoDuration::minutes(15); const BACKGROUND_AGENT_USAGE_PROFILE_WAIT_PREFIX: &str = "usage_profile_wait_until:"; +const BACKGROUND_AGENT_WORKER_GENERATION_ENV: &str = "CODEWITH_BACKGROUND_AGENT_WORKER_GENERATION"; +const BACKGROUND_AGENT_WORKER_SUPERVISOR_ID_ENV: &str = + "CODEWITH_BACKGROUND_AGENT_WORKER_SUPERVISOR_ID"; struct AgentStartManagedWorktree { worktree_id: String, @@ -264,7 +275,9 @@ impl ThreadRequestProcessor { let Some(run) = context.state_db.get_run(run_id.as_str()).await? else { anyhow::bail!("background-agent run not found: {run_id}"); }; - if !should_start_background_run(&run) { + let preclaimed_generation = + background_agent_worker_preclaimed_generation(&run, context.supervisor_id.as_str())?; + if preclaimed_generation.is_none() && !should_start_background_run(&run) { debug!(run_id = %run.id, status = ?run.status, "background-agent worker had nothing to start"); return Ok(()); } @@ -273,7 +286,9 @@ impl ThreadRequestProcessor { let mut active = context.active_workers.lock().await; active.insert(run.id.clone(), token.clone()); } - let result = run_background_agent_worker(context.clone(), run.clone(), token).await; + let result = + run_background_agent_worker(context.clone(), run.clone(), token, preclaimed_generation) + .await; context.active_workers.lock().await.remove(run.id.as_str()); match &result { Err(err) if is_background_agent_ownership_lost(err) => { @@ -329,6 +344,7 @@ impl ThreadRequestProcessor { &self, mut params: AgentStartParams, ) -> Result, JSONRPCErrorError> { + self.validate_agent_start_admission(&mut params)?; let managed_worktree = self .trusted_agent_start_managed_worktree( params.cwd.as_deref(), @@ -364,7 +380,12 @@ impl ThreadRequestProcessor { } let response = self .background_agent_state_processor() - .agent_start_inner(params) + .agent_start_inner( + params, + managed_worktree + .as_ref() + .map(|worktree| worktree.worktree_id.as_str()), + ) .await?; if let Some(worktree) = managed_worktree.as_ref() { let snapshot_cwd_matches = response @@ -1645,6 +1666,43 @@ impl ThreadRequestProcessor { .get(run_id) .cloned(); if let Some(handle) = process_handle { + if let Some(state_db) = self.state_db.as_ref() { + match state_db.get_background_agent_run(run_id).await { + Ok(Some(run)) => { + if let Err(err) = stop_tracked_background_agent_worker_process( + state_db, + &self.background_agent_worker_processes, + &run, + &handle, + ) + .await + { + warn!( + run_id, + "failed to stop and finalize background-agent worker process: {err}" + ); + } + return; + } + Ok(None) => {} + Err(load_err) => { + warn!( + run_id, + "failed to load background-agent run before stopping worker process: {load_err}" + ); + if let Err(stop_err) = stop_background_agent_worker_process(&handle).await { + warn!( + run_id, + "failed to stop background-agent worker process after state read failure: {stop_err}" + ); + } + // Retain the tracked handle so periodic pruning can + // retry durable finalization after the state read + // recovers, even when the OS process is already gone. + return; + } + } + } match stop_background_agent_worker_process(&handle).await { Ok(_) => { self.background_agent_worker_processes @@ -1666,6 +1724,38 @@ impl ThreadRequestProcessor { BackgroundAgentRequestProcessor::new(self.state_db.clone()) } + fn validate_agent_start_admission( + &self, + params: &mut AgentStartParams, + ) -> Result<(), JSONRPCErrorError> { + 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!({ + "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, + })); + 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 { @@ -1833,6 +1923,17 @@ 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, + BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + ) + .await? + { + continue; + } let token = CancellationToken::new(); { let mut active = context.active_workers.lock().await; @@ -1844,7 +1945,13 @@ async fn reconcile_background_agents( let worker_context = context.clone(); context.task_tracker.spawn(async move { let run_id = run.id.clone(); - if let Err(err) = run_background_agent_worker(worker_context.clone(), run, token).await + if let Err(err) = run_background_agent_worker( + worker_context.clone(), + run, + token, + /*preclaimed_generation*/ None, + ) + .await { error!(run_id, "background agent worker failed: {err}"); if let Err(mark_err) = @@ -1878,7 +1985,7 @@ async fn reconcile_background_agent_worker_processes( reconcile_managed_worktree_cleanup(&context.state_db).await?; rehydrate_background_agent_worker_processes(&context).await?; let cleanup_retried_run_ids = - retry_pending_background_agent_worker_process_cleanup(&context).await; + retry_pending_background_agent_worker_process_cleanup(&context).await?; prune_finished_background_agent_worker_processes(&context).await?; let runs = match only_run_id { Some(run_id) => context @@ -1918,6 +2025,17 @@ 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, + BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + ) + .await? + { + continue; + } if context .active_worker_processes .lock() @@ -1926,14 +2044,54 @@ async fn reconcile_background_agent_worker_processes( { continue; } + let process_lease_id = format!( + "{}:{}:{}", + context.supervisor_id, + run.id, + run.generation.saturating_add(1) + ); + let Some(generation) = 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, + BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + ) + .await? + else { + continue; + }; let stderr_log_path = background_agent_worker_stderr_log_path(&context, run.id.as_str()); let command = WorkerProcessCommand::new(&context.codex_bin, &stderr_log_path) .arg(OsString::from("app-server")) .arg(OsString::from("--listen")) .arg(OsString::from("off")) .arg(OsString::from("--background-agent-worker")) - .arg(OsString::from(run.id.clone())); - let handle = WorkerProcessController::default().spawn(command).await?; + .arg(OsString::from(run.id.clone())) + .env( + BACKGROUND_AGENT_WORKER_SUPERVISOR_ID_ENV, + context.supervisor_id.as_str(), + ) + .env( + BACKGROUND_AGENT_WORKER_GENERATION_ENV, + generation.to_string(), + ); + let handle = match WorkerProcessController::default().spawn(command).await { + Ok(handle) => handle, + Err(err) => { + fail_claimed_background_agent_worker_process( + &context, + run.id.as_str(), + generation, + "failed to spawn background-agent worker process", + &json!({"reason": "worker_process_spawn_failed"}), + ) + .await?; + return Err(err); + } + }; context .active_worker_processes .lock() @@ -1944,19 +2102,44 @@ async fn reconcile_background_agent_worker_processes( if let Some(hook) = context.worker_process_spawn_hook.as_ref() { hook(&handle)?; } + let stderr_log_path = stderr_log_path.display().to_string(); + let job_id = format!("worker-process:{}", run.id); + if !context + .state_db + .record_background_agent_execution_handle(BackgroundAgentExecutionHandleParams { + run_id: run.id.as_str(), + supervisor_id: context.supervisor_id.as_str(), + generation, + pid: Some(i64::from(handle.pid)), + pgid: handle.pgid.map(i64::from), + job_id: Some(job_id.as_str()), + start_token: handle.start_token.as_deref(), + stderr_log_path: Some(stderr_log_path.as_str()), + }) + .await? + { + anyhow::bail!("background-agent worker process lost ownership after spawn"); + } context .state_db - .append_background_agent_event( + .append_background_agent_event_for_supervisor( run.id.as_str(), + context.supervisor_id.as_str(), + generation, "agent.workerProcessSpawned", &json!({ "supervisorId": context.supervisor_id, + "generation": generation, "pid": handle.pid, "pgid": handle.pgid, - "stderrLogPath": stderr_log_path.display().to_string(), + "stderrLogPath": stderr_log_path, }), + /*allow_terminal_current*/ false, ) - .await + .await? + .ok_or_else(|| { + anyhow::anyhow!("background-agent worker process lost ownership after spawn") + }) } .await; if let Err(err) = event_result { @@ -1977,6 +2160,14 @@ async fn reconcile_background_agent_worker_processes( .lock() .await .remove(run.id.as_str()); + fail_claimed_background_agent_worker_process( + &context, + run.id.as_str(), + generation, + "background-agent worker process bookkeeping failed", + &json!({"reason": "worker_process_bookkeeping_failed"}), + ) + .await?; return Err(err); } Err(stop_err) => { @@ -1990,6 +2181,66 @@ async fn reconcile_background_agent_worker_processes( Ok(()) } +async fn fail_claimed_background_agent_worker_process( + context: &BackgroundAgentProcessSupervisorContext, + run_id: &str, + generation: i64, + status_reason: &str, + event_payload_json: &Value, +) -> anyhow::Result<()> { + if let Some(run) = context.state_db.get_background_agent_run(run_id).await? + && run.supervisor_id.as_deref() == Some(context.supervisor_id.as_str()) + && run.generation == generation + && run.desired_state != BackgroundAgentDesiredState::Running + { + return finalize_stopped_background_agent_process_for_run( + &context.state_db, + &run, + json!({ + "reason": "worker_process_stopped_after_desired_state_change", + "failure": event_payload_json, + }), + ) + .await; + } + let status_payload_json = json!({ + "phase": "failed", + "reason": event_payload_json, + }); + let status_event = context + .state_db + .append_background_agent_status_event_for_supervisor( + BackgroundAgentStatusEventForSupervisorParams { + run_id, + supervisor_id: context.supervisor_id.as_str(), + generation, + status: BackgroundAgentRunStatus::Failed, + status_reason: Some(status_reason), + event_type: "agent.failed", + event_payload_json, + summary: Some("Failed"), + pending_interaction_count: 0, + status_payload_json: &status_payload_json, + }, + ) + .await?; + if status_event.is_none() { + return Err(background_agent_ownership_lost(run_id, generation)); + } + context + .state_db + .finish_background_agent_process_lease( + run_id, + context.supervisor_id.as_str(), + generation, + /*exit_code*/ Some(1), + /*exit_signal*/ None, + Some(status_reason), + ) + .await?; + Ok(()) +} + async fn stop_tracked_background_agent_worker_process( state_db: &StateDbHandle, active_worker_processes: &Arc>>, @@ -2045,7 +2296,7 @@ async fn stop_background_agent_worker_process_with_context( async fn retry_pending_background_agent_worker_process_cleanup( context: &BackgroundAgentProcessSupervisorContext, -) -> HashSet { +) -> anyhow::Result> { let run_ids = context .worker_process_cleanup_pending .lock() @@ -2081,6 +2332,27 @@ async fn retry_pending_background_agent_worker_process_cleanup( .lock() .await .remove(run_id.as_str()); + if let Some(run) = context.state_db.get_run(run_id.as_str()).await? { + finalize_stopped_background_agent_process_for_run( + &context.state_db, + &run, + json!({"reason": "worker_process_cleanup_retry_succeeded"}), + ) + .await?; + if run.desired_state == BackgroundAgentDesiredState::Running + && run.supervisor_id.as_deref() == Some(context.supervisor_id.as_str()) + && run.status == BackgroundAgentRunStatus::Starting + { + fail_claimed_background_agent_worker_process( + context, + run_id.as_str(), + run.generation, + "background-agent worker process bookkeeping failed", + &json!({"reason": "worker_process_bookkeeping_failed"}), + ) + .await?; + } + } retried.insert(run_id); } Err(err) => { @@ -2091,7 +2363,7 @@ async fn retry_pending_background_agent_worker_process_cleanup( } } } - retried + Ok(retried) } async fn finalize_stopped_background_agent_process_for_run( @@ -2727,6 +2999,30 @@ async fn prune_finished_background_agent_worker_processes( }), ) .await?; + if run.desired_state == BackgroundAgentDesiredState::Running + && run.supervisor_id.as_deref() == Some(context.supervisor_id.as_str()) + && matches!( + run.status, + BackgroundAgentRunStatus::Starting + | BackgroundAgentRunStatus::Running + | BackgroundAgentRunStatus::WaitingOnApproval + | BackgroundAgentRunStatus::WaitingOnUser + ) + { + fail_claimed_background_agent_worker_process( + context, + run.id.as_str(), + run.generation, + "worker process exited after durable claim", + &json!({ + "reason": "worker_process_exited_after_claim", + "pid": handle.pid, + "pgid": handle.pgid, + "stderrLogPath": handle.stderr_log_path.display().to_string(), + }), + ) + .await?; + } } context .state_db @@ -2806,6 +3102,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 @@ -2821,6 +3118,63 @@ fn should_start_background_run(run: &BackgroundAgentRun) -> bool { true } +fn background_agent_worker_preclaimed_generation( + run: &BackgroundAgentRun, + supervisor_id: &str, +) -> anyhow::Result> { + let Some(value) = std::env::var_os(BACKGROUND_AGENT_WORKER_GENERATION_ENV) else { + return Ok(None); + }; + let generation = value + .to_string_lossy() + .parse::() + .context("invalid preclaimed background-agent worker generation")?; + if run.supervisor_id.as_deref() != Some(supervisor_id) + || run.generation != generation + || run.status != BackgroundAgentRunStatus::Starting + { + anyhow::bail!( + "preclaimed background-agent worker ownership does not match durable run state" + ); + } + Ok(Some(generation)) +} + +async fn wait_for_background_agent_worker_spawn_receipt( + context: &BackgroundAgentWorkerContext, + run_id: &str, + generation: i64, + last_event_seq: i64, +) -> anyhow::Result<()> { + let after_seq = last_event_seq.saturating_sub(20); + let wait = async { + loop { + let events = context + .state_db + .list_events_after(run_id, Some(after_seq), Some(50)) + .await?; + if events.iter().any(|event| { + event.event_type == "agent.workerProcessSpawned" + && event + .payload_json + .get("supervisorId") + .and_then(Value::as_str) + == Some(context.supervisor_id.as_str()) + && event.payload_json.get("generation").and_then(Value::as_i64) + == Some(generation) + }) { + return Ok(()); + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }; + tokio::time::timeout(Duration::from_secs(10), wait) + .await + .with_context(|| { + format!("timed out waiting for durable worker-process spawn receipt for run {run_id}") + })? +} + fn is_terminal_background_agent_status(status: BackgroundAgentRunStatus) -> bool { matches!( status, @@ -2841,25 +3195,46 @@ async fn run_background_agent_worker( context: BackgroundAgentWorkerContext, run: BackgroundAgentRun, cancel_token: CancellationToken, + preclaimed_generation: Option, ) -> anyhow::Result<()> { - let process_lease_id = format!( - "{}:{}:{}", - context.supervisor_id, - run.id, - run.generation.saturating_add(1) - ); - let Some(generation) = retry_transient_sqlite_busy("claim background agent supervisor", || { - context.state_db.claim_background_agent_supervisor( + let generation = match preclaimed_generation { + Some(generation) => generation, + None => { + let process_lease_id = format!( + "{}:{}:{}", + context.supervisor_id, + run.id, + run.generation.saturating_add(1) + ); + let Some(generation) = + retry_transient_sqlite_busy("claim background agent supervisor", || { + 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, + BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + ) + }) + .await? + else { + debug!(run_id = %run.id, "background agent run was not claimable"); + return Ok(()); + }; + generation + } + }; + if context.worker_process && preclaimed_generation.is_some() { + wait_for_background_agent_worker_spawn_receipt( + &context, run.id.as_str(), - context.supervisor_id.as_str(), - process_lease_id.as_str(), + generation, + run.last_event_seq, ) - }) - .await? - else { - debug!(run_id = %run.id, "background agent run was not claimable"); - return Ok(()); - }; + .await?; + } let (config, initial_execution_payload) = match resolve_background_agent_config(&context, &run).await? { @@ -2933,12 +3308,11 @@ async fn run_background_agent_worker( ) .await?; - let should_submit_prompt = run.thread_id.is_none() && run.rollout_path.is_none(); - let prompt = if should_submit_prompt { - Some(load_background_agent_prompt(&context.state_db, run.id.as_str()).await?) - } else { - None - }; + let prompt = load_background_agent_prompt(&context.state_db, run.id.as_str()).await?; + let initial_prompt_submission_id = background_agent_initial_prompt_submission_id(&run.id); + let initial_prompt_sha256 = + codex_state::StateRuntime::background_agent_identity_sha256(prompt.as_bytes()); + let resuming_durable_thread = run.rollout_path.is_some(); let NewThread { thread_id, thread, @@ -2954,10 +3328,21 @@ async fn run_background_agent_worker( }), ) .await?; - let rollout_path = session_configured + let durable_rollout_path = session_configured .rollout_path - .as_ref() - .map(|path| path.to_string_lossy().to_string()); + .as_deref() + .ok_or_else(|| anyhow::anyhow!("background-agent thread has no durable rollout path"))?; + let rollout_path = Some(durable_rollout_path.to_string_lossy().to_string()); + let initial_prompt_recorded = if resuming_durable_thread { + background_agent_initial_prompt_is_recorded( + durable_rollout_path, + initial_prompt_submission_id.as_str(), + prompt.as_str(), + ) + .await? + } else { + false + }; let thread_id_string = thread_id.to_string(); let session_id_string = session_configured.session_id.to_string(); let binding_params = BackgroundAgentThreadBindingParams { @@ -2988,26 +3373,34 @@ 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(), ) .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_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( + &execution_snapshot_params, + 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(), @@ -3019,9 +3412,16 @@ async fn run_background_agent_worker( ) .await?; - if should_submit_prompt { - let prompt = - prompt.ok_or_else(|| anyhow::anyhow!("background agent prompt missing for new run"))?; + if initial_prompt_recorded { + record_background_agent_initial_prompt_receipt( + &context.state_db, + run.id.as_str(), + thread_id, + initial_prompt_submission_id.as_str(), + initial_prompt_sha256.as_str(), + ) + .await?; + } else { with_startup_heartbeat( &context, run.id.as_str(), @@ -3029,17 +3429,21 @@ async fn run_background_agent_worker( "submit initial prompt", async { thread - .submit(Op::UserInput { - items: vec![UserInput::Text { - text: prompt, - text_elements: Vec::new(), - }], - final_output_json_schema: None, - responsesapi_client_metadata: None, - additional_context: Default::default(), - environments: None, - thread_settings: Default::default(), - }) + .submit_user_input_with_client_user_message_id( + Op::UserInput { + items: vec![UserInput::Text { + text: prompt, + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + environments: None, + thread_settings: Default::default(), + }, + /*trace*/ None, + Some(initial_prompt_submission_id.clone()), + ) .await?; Ok(()) }, @@ -3065,6 +3469,21 @@ async fn run_background_agent_worker( } event = thread.next_event() => { let event = event?; + if matches!( + &event.msg, + EventMsg::UserMessage(message) + if message.client_id.as_deref() + == Some(initial_prompt_submission_id.as_str()) + ) { + record_background_agent_initial_prompt_receipt( + &context.state_db, + run.id.as_str(), + thread_id, + initial_prompt_submission_id.as_str(), + initial_prompt_sha256.as_str(), + ) + .await?; + } if handle_background_agent_event( &context, run.id.as_str(), @@ -3072,7 +3491,10 @@ async fn run_background_agent_worker( thread.clone(), event.msg, &cancel_token, - &mut turn_error, + BackgroundAgentTurnState { + auth_profile_ref: run.auth_profile_ref.as_deref(), + turn_error: &mut turn_error, + }, ).await? { return Ok(()); } @@ -3131,16 +3553,31 @@ 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()) - .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()); + .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" + ) + })?; + validate_background_agent_initial_execution_snapshot(run, &snapshot)?; + let initial_execution_payload = Some(snapshot.payload_json.clone()); + let payload = snapshot.payload_json.as_object(); + let package_fingerprint = payload + .and_then(|payload| payload.get("packageFingerprint")) + .and_then(Value::as_str); + if package_fingerprint != Some(BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT) { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: persisted background agent runtime package is incompatible" + ); + } let cwd = payload .and_then(|payload| payload.get("cwd")) .and_then(Value::as_str) @@ -3230,7 +3667,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, @@ -3245,91 +3682,178 @@ 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)?; - - 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" + 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" ); - 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) + if let Some(retry_at) = + background_agent_exact_profile_retry_at(context, &config, run.auth_profile_ref.as_deref()) + .await { - 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, }) } -async fn insert_initial_goal_for_background_thread( - state_db: &codex_state::StateRuntime, - run_id: &str, - thread_id: codex_protocol::ThreadId, - initial_execution_payload: Option<&Value>, +fn validate_background_agent_initial_execution_snapshot( + run: &BackgroundAgentRun, + snapshot: &BackgroundAgentExecutionSnapshot, ) -> anyhow::Result<()> { - let Some(objective) = - initial_execution_payload.and_then(initial_goal_objective_from_execution_payload) - else { - return Ok(()); - }; - let Some(goal) = retry_transient_sqlite_busy("insert background agent initial goal", || { - state_db.thread_goals().insert_thread_goal( - thread_id, - objective.as_str(), - codex_state::ThreadGoalStatus::Active, - /*token_budget*/ None, + let payload = snapshot.payload_json.as_object().ok_or_else(|| { + anyhow::anyhow!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: initial execution snapshot must be an object" ) - }) - .await? - else { - return Ok(()); + })?; + let required = |key: &str| { + payload.get(key).ok_or_else(|| { + anyhow::anyhow!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: initial execution snapshot is missing {key}" + ) + }) }; - let event_payload = json!({ - "threadId": thread_id.to_string(), - "goalId": goal.goal_id, - "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) - }) - .await?; - Ok(()) -} - -fn initial_goal_objective_from_execution_payload(payload: &Value) -> Option { - payload + let required_string = |key: &str| -> anyhow::Result<&str> { + required(key)?.as_str().filter(|value| !value.is_empty()).ok_or_else(|| { + anyhow::anyhow!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: initial execution snapshot field {key} must be a non-empty string" + ) + }) + }; + let optional_string = |key: &str| -> anyhow::Result> { + let value = required(key)?; + if value.is_null() { + Ok(None) + } else { + value.as_str().map(Some).ok_or_else(|| { + anyhow::anyhow!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: initial execution snapshot field {key} must be a string or null" + ) + }) + } + }; + + required_string("cwd")?; + for key in ["model", "provider", "serviceTier"] { + optional_string(key)?; + } + let roots = required("workspaceRoots")?; + if !roots.is_null() + && !roots.as_array().is_some_and(|roots| { + roots + .iter() + .all(|root| root.as_str().is_some_and(|root| !root.is_empty())) + }) + { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: initial execution snapshot field workspaceRoots must be an array of non-empty strings or null" + ); + } + let permission_profile = required("permissionProfile")?; + if !permission_profile.is_null() && !permission_profile.is_object() { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: initial execution snapshot field permissionProfile must be an object or null" + ); + } + let expected_auth_profile_identity = run.auth_profile_ref.as_deref().map(|profile| { + codex_state::StateRuntime::background_agent_identity_sha256(profile.as_bytes()) + }); + if optional_string("authProfileIdentitySha256")? != expected_auth_profile_identity.as_deref() { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH}: persisted execution snapshot auth profile does not match the admitted run" + ); + } + if optional_string("managedWorktreeId")?.is_some_and(str::is_empty) { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: managed worktree id must not be empty" + ); + } + if required_string("configFingerprint")? + != run.config_fingerprint.as_deref().unwrap_or_default() + || snapshot.config_fingerprint.as_deref() != run.config_fingerprint.as_deref() + { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: persisted execution snapshot configuration fingerprint does not match the admitted run" + ); + } + if required_string("versionFingerprint")? != BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION + || run.version_fingerprint.as_deref() != Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION) + { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: persisted execution snapshot admission schema is incompatible" + ); + } + if required_string("packageFingerprint")? != BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT + { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: persisted execution snapshot runtime package is incompatible" + ); + } + let recovery_policy = required_string("recoveryPolicy")?; + if recovery_policy != snapshot.recovery_policy.as_str() { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_SCHEMA_MISMATCH}: persisted execution snapshot recovery policy does not match its envelope" + ); + } + Ok(()) +} + +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<()> { + let Some(objective) = + initial_execution_payload.and_then(initial_goal_objective_from_execution_payload) + else { + return Ok(()); + }; + let Some(goal) = retry_transient_sqlite_busy("insert background agent initial goal", || { + state_db.thread_goals().insert_thread_goal( + thread_id, + objective.as_str(), + codex_state::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + }) + .await? + else { + return Ok(()); + }; + let event_payload = json!({ + "threadId": thread_id.to_string(), + "goalId": goal.goal_id, + "objective": goal.objective, + }); + retry_transient_sqlite_busy("append background agent initial goal event", || { + state_db.append_background_agent_event_for_supervisor( + run_id, + supervisor_id, + generation, + "agent.initialGoalCreated", + &event_payload, + /*allow_terminal_current*/ false, + ) + }) + .await? + .ok_or_else(|| background_agent_ownership_lost(run_id, generation))?; + Ok(()) +} + +fn initial_goal_objective_from_execution_payload(payload: &Value) -> Option { + payload .get("initialGoalObjective") .and_then(Value::as_str) .map(str::trim) @@ -3370,14 +3894,86 @@ async fn defer_background_agent_for_usage_profile_wait( Ok(()) } -fn background_agent_broker_retry_at( +async fn background_agent_exact_profile_retry_at( + context: &BackgroundAgentWorkerContext, config: &codex_core::config::Config, - retry_at: i64, + auth_profile_ref: Option<&str>, ) -> 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) + let scoped_auth_manager = context + .auth_manager + .shared_scoped_auth_profile(auth_profile_ref.map(str::to_string)) + .await; + let auth = scoped_auth_manager.auth().await?; + if !auth.uses_codex_backend() { + return None; + } + let client = BackendClient::from_auth(config.chatgpt_base_url.clone(), &auth).ok()?; + let snapshots = client.get_rate_limits_many().await.ok()?; + let now = Utc::now().timestamp(); + let mut exhausted = false; + let mut retry_at = None; + for snapshot in &snapshots { + let snapshot = codex_core::usage_profile_health::UsageProfileRateLimitSnapshot { + limit_id: snapshot.limit_id.as_deref(), + limit_name: snapshot.limit_name.as_deref(), + primary: snapshot.primary.as_ref().map(|window| { + codex_core::usage_profile_health::UsageProfileRateLimitWindow { + used_percent: window.used_percent, + window_minutes: window.window_minutes, + resets_at: window.resets_at, + } + }), + secondary: snapshot.secondary.as_ref().map(|window| { + codex_core::usage_profile_health::UsageProfileRateLimitWindow { + used_percent: window.used_percent, + window_minutes: window.window_minutes, + resets_at: window.resets_at, + } + }), + }; + exhausted |= + codex_core::usage_profile_health::exhausted_auto_switch_window_for_snapshot(&snapshot) + .is_some(); + if let Some(snapshot_retry_at) = + codex_core::usage_profile_health::earliest_exhausted_reset_at(&snapshot, now) + { + retry_at = Some(retry_at.map_or(snapshot_retry_at, |current: i64| { + current.min(snapshot_retry_at) + })); + } + } + exhausted.then(|| { + background_agent_usage_retry_at( + retry_at, + config.usage_self_heal.reset_retry_buffer_secs, + config.usage_self_heal.initial_backoff_secs, + now, + ) + }) +} + +fn background_agent_usage_retry_at( + reset_at: Option, + reset_retry_buffer_secs: u64, + initial_backoff_secs: u64, + now: i64, +) -> DateTime { + let retry_at = reset_at + .filter(|reset_at| *reset_at > now) + .map(|reset_at| { + reset_at.saturating_add(i64::try_from(reset_retry_buffer_secs).unwrap_or(i64::MAX)) + }) + .unwrap_or_else(|| { + now.saturating_add(i64::try_from(initial_backoff_secs.max(1)).unwrap_or(i64::MAX)) + }); + DateTime::::from_timestamp(retry_at, /*nsecs*/ 0).unwrap_or_else(Utc::now) +} + +fn background_agent_error_is_usage_limit(event: &codex_protocol::protocol::ErrorEvent) -> bool { + matches!( + event.codex_error_info, + Some(CodexErrorInfo::UsageLimitExceeded) + ) } fn background_agent_usage_profile_wait_reason(retry_at: DateTime) -> String { @@ -3417,6 +4013,75 @@ fn background_agent_snapshot_sandbox_mode(value: &Value) -> anyhow::Result String { + let identity = codex_state::StateRuntime::background_agent_identity_sha256(run_id.as_bytes()); + format!("background-agent-initial-prompt:{identity}") +} + +fn background_agent_initial_prompt_is_in_history( + history: &InitialHistory, + submission_id: &str, + prompt: &str, +) -> bool { + history.scan_rollout_items(|item| { + matches!( + item, + codex_protocol::protocol::RolloutItem::EventMsg(EventMsg::UserMessage(message)) + if message.client_id.as_deref() == Some(submission_id) + || (message.client_id.is_none() && message.message == prompt) + ) + }) +} + +async fn background_agent_initial_prompt_is_recorded( + rollout_path: &Path, + submission_id: &str, + prompt: &str, +) -> anyhow::Result { + let history = codex_rollout::RolloutRecorder::get_rollout_history(rollout_path) + .await + .with_context(|| { + format!( + "failed to inspect background-agent rollout {} for initial prompt", + rollout_path.display() + ) + })?; + Ok(background_agent_initial_prompt_is_in_history( + &history, + submission_id, + prompt, + )) +} + +async fn record_background_agent_initial_prompt_receipt( + state_db: &codex_state::StateRuntime, + run_id: &str, + thread_id: codex_protocol::ThreadId, + submission_id: &str, + prompt_sha256: &str, +) -> anyhow::Result<()> { + let receipt_identity = + codex_state::StateRuntime::background_agent_identity_sha256(submission_id.as_bytes()); + let receipt_key = format!("initial-prompt:{receipt_identity}"); + let diagnostics = json!({ + "threadId": thread_id.to_string(), + "submissionId": submission_id, + "promptSha256": prompt_sha256, + }); + retry_transient_sqlite_busy("record background agent initial prompt receipt", || { + state_db.append_background_agent_lifecycle_receipt( + run_id, + "agent.initialPromptRecorded", + receipt_key.as_str(), + /*generation*/ 0, + /*attempt*/ None, + &diagnostics, + ) + }) + .await?; + Ok(()) +} + async fn load_background_agent_prompt( state_db: &StateDbHandle, run_id: &str, @@ -3426,7 +4091,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 @@ -3437,6 +4111,11 @@ async fn load_background_agent_prompt( .ok_or_else(|| anyhow::anyhow!("background agent prompt missing from start event")) } +struct BackgroundAgentTurnState<'a> { + auth_profile_ref: Option<&'a str>, + turn_error: &'a mut Option, +} + async fn handle_background_agent_event( context: &BackgroundAgentWorkerContext, run_id: &str, @@ -3444,11 +4123,11 @@ async fn handle_background_agent_event( thread: Arc, msg: EventMsg, cancel_token: &CancellationToken, - turn_error: &mut Option, + turn_state: BackgroundAgentTurnState<'_>, ) -> anyhow::Result { match msg { EventMsg::TurnStarted(event) => { - *turn_error = None; + *turn_state.turn_error = None; tolerate_transient_busy( append_status( context, @@ -3722,11 +4401,33 @@ async fn handle_background_agent_event( .await?; } EventMsg::Error(event) => { + if background_agent_error_is_usage_limit(&event) { + let config = thread.config().await; + let retry_at = background_agent_exact_profile_retry_at( + context, + &config, + turn_state.auth_profile_ref, + ) + .await + .unwrap_or_else(|| { + background_agent_usage_retry_at( + /*reset_at*/ None, + config.usage_self_heal.reset_retry_buffer_secs, + config.usage_self_heal.initial_backoff_secs, + Utc::now().timestamp(), + ) + }); + defer_background_agent_for_usage_profile_wait( + context, run_id, generation, retry_at, + ) + .await?; + return Ok(true); + } // Codex-core emits `EventMsg::Error` immediately before ending a turn // (e.g. pre-sampling/auto compaction failed with // `context_length_exceeded`). Record it so the trailing `TurnComplete` // can be reclassified as a failure instead of a silent success. - *turn_error = Some(event.message.clone()); + *turn_state.turn_error = Some(event.message.clone()); tolerate_transient_busy( append_background_agent_event_with_retry( context, @@ -3749,9 +4450,10 @@ async fn handle_background_agent_event( // re-growing) the run. A no-message turn without an observed error is // still classified as a failure by `background_agent_turn_complete_status` // below, but with the generic reason. - if let Some(reason) = - turn_completion_failure_reason(event.last_agent_message.as_deref(), turn_error) - { + if let Some(reason) = turn_completion_failure_reason( + event.last_agent_message.as_deref(), + turn_state.turn_error, + ) { let status_reason = format!("turn failed: {reason}"); append_status( context, @@ -4258,36 +4960,40 @@ 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_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( - 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(), - ), - }) + .upsert_background_agent_status_snapshot_for_supervisor( + &status_snapshot_params, + context.supervisor_id.as_str(), + generation, + ) }) - .await?; + .await? + .ok_or_else(|| background_agent_ownership_lost(run_id, generation))?; Ok(()) } @@ -4312,14 +5018,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( @@ -4509,7 +5221,13 @@ fn status_snapshot_payload( } } if let Some(execution_payload) = execution_payload { - for key in ["model", "provider", "serviceTier", "cwd", "authProfileRef"] { + for key in [ + "model", + "provider", + "serviceTier", + "cwd", + "authProfileIdentitySha256", + ] { if let Some(value) = execution_payload.get(key).cloned() { object.insert(key.to_string(), value); } @@ -4728,6 +5446,9 @@ fn tolerate_transient_busy( } pub(super) fn new_background_agent_supervisor_id() -> String { + if let Some(supervisor_id) = std::env::var_os(BACKGROUND_AGENT_WORKER_SUPERVISOR_ID_ENV) { + return supervisor_id.to_string_lossy().into_owned(); + } format!( "{}:{}:{}", BACKGROUND_AGENT_SUPERVISOR_PREFIX, @@ -5168,8 +5889,7 @@ done } #[tokio::test] - async fn post_spawn_event_failure_retries_tracked_worker_cleanup_after_stop_failure() - -> anyhow::Result<()> { + async fn post_spawn_event_failure_retry_finalizes_concurrent_stop() -> anyhow::Result<()> { let temp = TempDir::new()?; let state_db = codex_state::StateRuntime::init(temp.path().to_path_buf(), "test-provider".to_string()) @@ -5183,7 +5903,7 @@ done let fail_stop = Arc::new(AtomicBool::new(true)); let fail_stop_for_hook = Arc::clone(&fail_stop); let context = BackgroundAgentProcessSupervisorContext { - state_db, + state_db: Arc::clone(&state_db), supervisor_id: "process-supervisor-test".to_string(), active_worker_processes: Arc::clone(&active_worker_processes), worker_process_cleanup_pending: Arc::new(Mutex::new(HashSet::new())), @@ -5230,6 +5950,19 @@ done ); assert!(worker_process_exists(descendant_pid)?); + state_db + .set_background_agent_desired_state( + "forced-stop-run", + BackgroundAgentDesiredState::Stopped, + ) + .await?; + state_db + .update_background_agent_run_status( + "forced-stop-run", + BackgroundAgentRunStatus::Stopping, + Some("stop requested"), + ) + .await?; fail_stop.store(false, Ordering::SeqCst); reconcile_background_agent_worker_processes(context, Some("forced-stop-run".to_string())) .await?; @@ -5244,6 +5977,15 @@ done tokio::time::sleep(Duration::from_millis(25)).await; } assert!(active_worker_processes.lock().await.is_empty()); + let run = state_db + .get_background_agent_run("forced-stop-run") + .await? + .expect("run should exist"); + assert_eq!(run.status, BackgroundAgentRunStatus::Cancelled); + assert_eq!( + run.status_reason.as_deref(), + Some("worker process stopped after stop request") + ); drop(fixture); assert!(!temp_path.exists()); Ok(()) @@ -5302,7 +6044,13 @@ done .collect::>(); assert_eq!( event_types, - vec!["agent.started", "agent.workerProcessSpawned"] + vec![ + "agent.admitted", + "agent.started", + "agent.claimed", + "agent.heartbeat", + "agent.workerProcessSpawned", + ] ); let spawn_event = events .last() @@ -5321,6 +6069,38 @@ done spawn_event.get("pgid").and_then(Value::as_u64), Some(u64::from(handle.pid)) ); + let run = state_db + .get_background_agent_run("run/with path") + .await? + .expect("run should remain durable"); + assert_eq!(run.status, BackgroundAgentRunStatus::Starting); + assert_eq!( + run.supervisor_id.as_deref(), + Some("process-supervisor-test") + ); + + let duplicate_spawned = Arc::new(AtomicBool::new(false)); + let duplicate_spawned_for_hook = Arc::clone(&duplicate_spawned); + let competing_context = BackgroundAgentProcessSupervisorContext { + state_db: Arc::clone(&state_db), + supervisor_id: "competing-process-supervisor".to_string(), + active_worker_processes: Arc::new(Mutex::new(HashMap::new())), + worker_process_cleanup_pending: Arc::new(Mutex::new(HashSet::new())), + codex_home: temp.path().to_path_buf(), + codex_bin: fixture.program.clone(), + worker_process_spawn_hook: Some(Arc::new(move |_| { + duplicate_spawned_for_hook.store(true, Ordering::SeqCst); + Ok(()) + })), + worker_process_stop_hook: None, + }; + reconcile_background_agent_worker_processes( + competing_context, + Some("run/with path".to_string()), + ) + .await?; + assert!(!duplicate_spawned.load(Ordering::SeqCst)); + fixture.release()?; wait_for_worker_process_exit(&WorkerProcessController::default(), &handle).await?; @@ -5337,6 +6117,40 @@ done assert_eq!(underscore_component, "run_a"); } + #[test] + fn usage_limit_error_defers_exact_profile_until_reset_or_bounded_retry() { + let usage_error = codex_protocol::protocol::ErrorEvent { + message: "usage exhausted".to_string(), + codex_error_info: Some(CodexErrorInfo::UsageLimitExceeded), + }; + let ordinary_error = codex_protocol::protocol::ErrorEvent { + message: "ordinary failure".to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + }; + assert!(background_agent_error_is_usage_limit(&usage_error)); + assert!(!background_agent_error_is_usage_limit(&ordinary_error)); + + let now = 1_000; + assert_eq!( + background_agent_usage_retry_at( + Some(2_000), + /*reset_retry_buffer_secs*/ 30, + /*initial_backoff_secs*/ 15, + now, + ) + .timestamp(), + 2_030 + ); + assert_eq!( + background_agent_usage_retry_at( + /*reset_at*/ None, /*reset_retry_buffer_secs*/ 30, + /*initial_backoff_secs*/ 15, now, + ) + .timestamp(), + 1_015 + ); + } + #[tokio::test] async fn usage_profile_wait_status_reason_defers_queued_run_until_reset() -> anyhow::Result<()> { @@ -5399,6 +6213,97 @@ done ); } + #[test] + fn initial_prompt_recovery_uses_deterministic_rollout_identity() { + let run_id = "initial-prompt-run"; + let submission_id = background_agent_initial_prompt_submission_id(run_id); + let history = + InitialHistory::Forked(vec![codex_protocol::protocol::RolloutItem::EventMsg( + EventMsg::UserMessage(codex_protocol::protocol::UserMessageEvent { + client_id: Some(submission_id.clone()), + message: "durable prompt".to_string(), + ..Default::default() + }), + )]); + let unrelated = + InitialHistory::Forked(vec![codex_protocol::protocol::RolloutItem::EventMsg( + EventMsg::UserMessage(codex_protocol::protocol::UserMessageEvent { + client_id: Some("another-submission".to_string()), + message: "durable prompt".to_string(), + ..Default::default() + }), + )]); + let legacy = InitialHistory::Forked(vec![codex_protocol::protocol::RolloutItem::EventMsg( + EventMsg::UserMessage(codex_protocol::protocol::UserMessageEvent { + client_id: None, + message: "durable prompt".to_string(), + ..Default::default() + }), + )]); + + assert!(background_agent_initial_prompt_is_in_history( + &history, + submission_id.as_str(), + "durable prompt" + )); + assert!(!background_agent_initial_prompt_is_in_history( + &unrelated, + submission_id.as_str(), + "durable prompt" + )); + assert!(background_agent_initial_prompt_is_in_history( + &legacy, + submission_id.as_str(), + "durable prompt" + )); + } + + #[tokio::test] + async fn initial_prompt_receipt_replays_after_rollout_recovery() -> anyhow::Result<()> { + let temp = TempDir::new()?; + let state_db = + codex_state::StateRuntime::init(temp.path().to_path_buf(), "test-provider".to_string()) + .await?; + seed_queued_run(state_db.as_ref(), "initial-prompt-receipt").await?; + let thread_id = codex_protocol::ThreadId::new(); + let submission_id = background_agent_initial_prompt_submission_id("initial-prompt-receipt"); + let prompt_sha256 = + codex_state::StateRuntime::background_agent_identity_sha256(b"durable prompt"); + + record_background_agent_initial_prompt_receipt( + state_db.as_ref(), + "initial-prompt-receipt", + thread_id, + submission_id.as_str(), + prompt_sha256.as_str(), + ) + .await?; + record_background_agent_initial_prompt_receipt( + state_db.as_ref(), + "initial-prompt-receipt", + thread_id, + submission_id.as_str(), + prompt_sha256.as_str(), + ) + .await?; + + let events = state_db + .list_background_agent_events_after( + "initial-prompt-receipt", + /*after_seq*/ None, + None, + ) + .await?; + assert_eq!( + events + .iter() + .filter(|event| event.event_type == "agent.initialPromptRecorded") + .count(), + 1 + ); + Ok(()) + } + #[tokio::test] async fn insert_initial_goal_for_background_thread_creates_active_goal_once() -> anyhow::Result<()> { @@ -5407,6 +6312,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 ", @@ -5415,6 +6324,8 @@ done insert_initial_goal_for_background_thread( state_db.as_ref(), "goal-run", + "goal-supervisor", + generation, thread_id, Some(&initial_execution_payload), ) @@ -5455,6 +6366,8 @@ done insert_initial_goal_for_background_thread( state_db.as_ref(), "goal-run", + "goal-supervisor", + generation, thread_id, Some(&initial_execution_payload), ) @@ -5532,7 +6445,7 @@ done assert_eq!(run.status, BackgroundAgentRunStatus::Failed); assert_eq!( run.status_reason.as_deref(), - Some("worker process exited before claiming run") + Some("worker process exited after durable claim") ); let snapshot = context .state_db @@ -5763,6 +6676,85 @@ done Ok(()) } + #[tokio::test] + async fn claimed_spawn_failure_finalizes_concurrent_stop() -> anyhow::Result<()> { + let temp = TempDir::new()?; + let state_db = + codex_state::StateRuntime::init(temp.path().to_path_buf(), "test-provider".to_string()) + .await?; + seed_queued_run(state_db.as_ref(), "stopped-before-spawn").await?; + let generation = state_db + .claim_background_agent_supervisor( + "stopped-before-spawn", + "process-supervisor-test", + "lease-1", + ) + .await? + .expect("run should be claimed"); + state_db + .set_background_agent_desired_state( + "stopped-before-spawn", + BackgroundAgentDesiredState::Stopped, + ) + .await?; + state_db + .update_background_agent_run_status( + "stopped-before-spawn", + BackgroundAgentRunStatus::Stopping, + Some("stop requested"), + ) + .await?; + let context = BackgroundAgentProcessSupervisorContext { + state_db: Arc::clone(&state_db), + supervisor_id: "process-supervisor-test".to_string(), + active_worker_processes: Arc::new(Mutex::new(HashMap::new())), + worker_process_cleanup_pending: Arc::new(Mutex::new(HashSet::new())), + codex_home: temp.path().to_path_buf(), + codex_bin: PathBuf::from("/bin/true"), + worker_process_spawn_hook: None, + worker_process_stop_hook: None, + }; + + fail_claimed_background_agent_worker_process( + &context, + "stopped-before-spawn", + generation, + "failed to spawn background-agent worker process", + &json!({"reason": "worker_process_spawn_failed"}), + ) + .await?; + + let run = state_db + .get_background_agent_run("stopped-before-spawn") + .await? + .expect("run should exist"); + assert_eq!(run.status, BackgroundAgentRunStatus::Cancelled); + assert_eq!( + run.status_reason.as_deref(), + Some("worker process stopped after stop request") + ); + let event_types = state_db + .list_background_agent_events_after( + "stopped-before-spawn", + /*after_seq*/ None, + /*limit*/ None, + ) + .await? + .into_iter() + .map(|event| event.event_type) + .collect::>(); + assert_eq!( + event_types, + vec![ + "agent.admitted", + "agent.started", + "agent.claimed", + "agent.cancelled", + ] + ); + Ok(()) + } + #[tokio::test] async fn process_supervisor_finalizes_missing_stopped_handle_during_prune() -> anyhow::Result<()> { @@ -5894,36 +6886,50 @@ 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, + "packageFingerprint": BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + "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("version-test".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, + codex_background_agent::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 85cacebb10..54e24a08a8 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,12 @@ 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_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::BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT; use codex_background_agent::BackgroundAgentDesiredState; use codex_background_agent::BackgroundAgentEvent; use codex_background_agent::BackgroundAgentExecutionSnapshot; @@ -73,6 +78,7 @@ 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; @@ -87,17 +93,18 @@ use codex_rollout::StateDbHandle; use codex_state::ManagedWorktreeAssignmentTarget; 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; 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 { @@ -111,9 +118,11 @@ impl BackgroundAgentRequestProcessor { pub(super) async fn agent_start_inner( &self, - params: AgentStartParams, + mut params: AgentStartParams, + required_managed_worktree_id: Option<&str>, ) -> Result { let state_db = self.state_db()?; + normalize_agent_start_schema(&mut params)?; let AgentStartParams { prompt, initial_goal_objective, @@ -138,50 +147,38 @@ 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}" + )) + })?; + // Runs admitted by an older codewith build can never be claimed by this + // binary, so release the admission capacity they still hold instead of + // failing every future admission with a capacity error. + retry_transient_sqlite_busy("release incompatible background agent admissions", || { + state_db.terminalize_incompatible_background_agent_runs( + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + ) + }) + .await + .map_err(|err| { + internal_error(format!( + "failed to release incompatible background agent admissions: {err}" + )) + })?; 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( @@ -194,152 +191,78 @@ 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 + 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 create background agent: {err}")) + internal_error(format!( + "failed to fingerprint background agent configuration: {err}" + )) })?, + ), }; - let created_new_run = run.id == agent_id; - let execution_payload = initial_execution_snapshot_payload( - &run, - InitialExecutionSnapshotPayloadParams { + 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(), - }, - ); - 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}" - )) - })?, - } + auth_profile_ref: auth_profile_ref.as_deref(), + required_managed_worktree_id, + config_fingerprint: config_fingerprint.as_deref(), + version_fingerprint: version_fingerprint.as_deref(), + }); + 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: config_fingerprint.clone(), + version_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 mut events = state_db - .list_events_after(run.id.as_str(), /*after_seq*/ None, Some(1)) - .await - .map_err(|err| { - internal_error(format!("failed to list background agent events: {err}")) - })?; - match events.pop() { - 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", - }), - ) - .await - .map_err(|err| { - internal_error(format!("failed to append background agent event: {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 snapshot = match state_db - .get_status_snapshot(run.id.as_str()) + 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 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), @@ -572,69 +495,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()) @@ -650,80 +578,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? @@ -839,7 +699,7 @@ 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 +1048,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, )) } @@ -1197,10 +1057,13 @@ struct InitialExecutionSnapshotPayloadParams<'a> { initial_goal_objective: Option<&'a str>, execution_context: Option<&'a AgentExecutionContextParams>, recovery_policy: &'a str, + auth_profile_ref: Option<&'a str>, + required_managed_worktree_id: Option<&'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!({ @@ -1213,7 +1076,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()), @@ -1251,13 +1113,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(), + "authProfileIdentitySha256": params + .auth_profile_ref + .map(|profile| format!("{:x}", Sha256::digest(profile.as_bytes()))), + "managedWorktreeId": params.required_managed_worktree_id, + "configFingerprint": params.config_fingerprint, + "versionFingerprint": params.version_fingerprint, + "packageFingerprint": BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, "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 { @@ -1277,6 +1162,43 @@ fn validate_agent_prompt(prompt: String) -> Result { Ok(prompt) } +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) { + 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) + } +} + +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> { @@ -1319,18 +1241,6 @@ fn map_background_agent_event_replay_error(err: anyhow::Error) -> JSONRPCErrorEr } } -async fn append_background_agent_event_with_retry( - state_db: &codex_state::StateRuntime, - run_id: &str, - event_type: &str, - payload_json: &Value, -) -> anyhow::Result { - retry_transient_sqlite_busy("append background agent event", || { - state_db.append_event(run_id, event_type, payload_json) - }) - .await -} - fn decode_offset_cursor(cursor: Option<&str>) -> Result { let Some(cursor) = cursor else { return Ok(0); 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..3fc5b06dbd 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,11 @@ 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; +use codex_background_agent::BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT; use codex_protocol::ThreadId; use codex_protocol::models::PermissionProfile; use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS; @@ -78,6 +83,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; @@ -87,6 +94,11 @@ use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); const INVALID_REQUEST_ERROR_CODE: i64 = -32600; +enum SeededAgentExecution { + Inert, + Claimable { cwd: String }, +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result<()> { let codex_home = TempDir::new()?; @@ -153,7 +165,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 +211,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( @@ -210,13 +222,10 @@ 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.workerStarting" - ); + 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() @@ -232,7 +241,7 @@ async fn agent_start_list_read_and_events_survive_app_server_restart() -> Result let stale_cursor_error = agent_events_error( &mut restarted, &agent_id, - Some("event:1".to_string()), + Some("event:0".to_string()), Some(1), ) .await?; @@ -250,7 +259,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())?; @@ -261,6 +270,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?; @@ -274,7 +284,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", @@ -283,8 +293,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")) @@ -330,6 +350,7 @@ async fn agent_list_pages_beyond_state_default_cap() -> Result<()> { agent_id.as_str(), /*idempotency_key*/ None, "paged run", + SeededAgentExecution::Inert, ) .await?; state_db @@ -390,7 +411,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() @@ -413,8 +433,12 @@ 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"), None); assert_eq!( - start.execution_snapshot.payload.get("authProfileRef"), + start + .execution_snapshot + .payload + .get("authProfileIdentitySha256"), Some(&JsonValue::Null) ); assert_eq!( @@ -457,6 +481,75 @@ 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)) + ); + 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 compatible_agent_id = compatible.agent.agent_id; + + 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_eq!( + agent_list(&mut mcp) + .await? + .data + .into_iter() + .map(|agent| agent.agent_id) + .collect::>(), + vec![compatible_agent_id] + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn agent_start_uses_validated_managed_worktree_cwd() -> Result<()> { let codex_home = TempDir::new()?; @@ -749,36 +842,11 @@ async fn supervisor_periodically_starts_durable_queued_runs() -> Result<()> { agent_id.as_str(), /*idempotency_key*/ None, "picked up by periodic supervisor", + SeededAgentExecution::Claimable { + cwd: codex_home.path().display().to_string(), + }, ) .await?; - state_db - .create_background_agent_execution_snapshot(&BackgroundAgentExecutionSnapshotParams { - run_id: agent_id.clone(), - snapshot_kind: "initial_execution_context".to_string(), - payload_json: json!({ - "snapshotSource": "state-seeded-test", - "cwd": codex_home.path().display().to_string(), - "workspaceRoots": [codex_home.path().display().to_string()], - "model": "mock-model", - "provider": "mock_provider", - "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()), - }) - .await?; - state_db - .upsert_background_agent_status_snapshot(&BackgroundAgentStatusSnapshotParams { - run_id: agent_id.clone(), - seq: 1, - status: StateBackgroundAgentRunStatus::Queued, - desired_state: StateBackgroundAgentDesiredState::Running, - summary: Some("Queued".to_string()), - pending_interaction_count: 0, - last_event_seq: 1, - payload_json: json!({"phase": "queued"}), - }) - .await?; let completed = wait_for_agent_status(&mut mcp, agent_id.as_str(), AgentRunStatus::Completed).await?; @@ -786,7 +854,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() @@ -813,6 +881,7 @@ async fn agent_lifecycle_and_pending_interaction_flow() -> Result<()> { agent_id.as_str(), /*idempotency_key*/ None, "wait for approval", + SeededAgentExecution::Inert, ) .await?; state_db @@ -1059,6 +1128,7 @@ async fn agent_stop_preserves_delete_requested_desired_state() -> Result<()> { agent_id.as_str(), /*idempotency_key*/ None, "delete then stop", + SeededAgentExecution::Inert, ) .await?; drop(state_db); @@ -1115,6 +1185,7 @@ async fn agent_pending_interaction_respond_rejects_invalid_responded_payload() - agent_id.as_str(), /*idempotency_key*/ None, "wait for approval", + SeededAgentExecution::Inert, ) .await?; state_db @@ -1223,7 +1294,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( @@ -1240,8 +1311,21 @@ async fn agent_diagnostics_reports_quota_and_overloaded_admission() -> Result<() &format!("quota-run-{index}"), Some(format!("quota-idempotency-{index}")), &format!("quota run {index}"), + SeededAgentExecution::Inert, ) .await?; + // Park the filler runs in `stopping`: they still consume an admission + // slot, but unlike queued/orphaned rows they are neither claimable nor + // reclaimable, so this exercises capacity backpressure rather than the + // incompatible-runtime reaper (covered by + // `agent_start_reclaims_capacity_stranded_by_a_runtime_upgrade`). + state_db + .update_background_agent_run_status( + &format!("quota-run-{index}"), + StateBackgroundAgentRunStatus::Stopping, + Some("parked by quota test"), + ) + .await?; } let first_agent_id = "quota-run-0".to_string(); @@ -1252,7 +1336,6 @@ async fn agent_diagnostics_reports_quota_and_overloaded_admission() -> Result<() let full = initial; assert_eq!(full.active_run_count, full.max_active_runs_per_user); - assert_eq!(full.queued_run_count, full.max_active_runs_per_user); assert_eq!(full.available_active_run_slots, 0); assert!(!full.admission_allowed); assert_eq!( @@ -1260,7 +1343,7 @@ async fn agent_diagnostics_reports_quota_and_overloaded_admission() -> Result<() vec!["active_run_limit".to_string()] ); assert_eq!( - run_status_count(&full, AgentRunStatus::Queued), + run_status_count(&full, AgentRunStatus::Stopping), full.max_active_runs_per_user ); @@ -1274,13 +1357,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 @@ -1300,15 +1391,13 @@ async fn agent_diagnostics_reports_quota_and_overloaded_admission() -> Result<() assert!(available.admission_allowed); assert!(available.backpressure_reasons.is_empty()); - let accepted = start_agent( - &mut mcp, - start_params( - "new run after a slot opens", - Some("quota-idempotency-after-slot".to_string()), - codex_home.path(), - ), - ) - .await?; + let accepted_params = start_params( + "new run after a slot opens", + Some("quota-idempotency-after-slot".to_string()), + codex_home.path(), + ); + let retry_params = accepted_params.clone(); + let accepted = start_agent(&mut mcp, accepted_params).await?; assert_ne!(accepted.agent.agent_id, first_agent_id); wait_for_agent_status( &mut mcp, @@ -1317,16 +1406,82 @@ async fn agent_diagnostics_reports_quota_and_overloaded_admission() -> Result<() ) .await?; - let retry = start_agent( + let retry = start_agent(&mut mcp, retry_params).await?; + assert_eq!(retry.agent.agent_id, accepted.agent.agent_id); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn agent_start_reclaims_capacity_stranded_by_a_runtime_upgrade() -> Result<()> { + let codex_home = TempDir::new()?; + let server = + create_mock_responses_server_sequence(vec![create_final_assistant_message_sse_response( + "accepted after reclaim", + )?]) + .await; + write_config(codex_home.path(), server.uri().as_str())?; + + let mut mcp = init_mcp(codex_home.path()).await?; + let state_db = init_state_db(codex_home.path()).await?; + // Seeded with a foreign `packageFingerprint`, exactly like runs admitted by + // a previous codewith build: queued, still counted against the admission + // quota, and unclaimable by this binary forever. + for index in 0..8 { + seed_queued_agent_run( + state_db.as_ref(), + &format!("stranded-run-{index}"), + Some(format!("stranded-idempotency-{index}")), + &format!("stranded run {index}"), + SeededAgentExecution::Inert, + ) + .await?; + } + + let full = agent_daemon_diagnostics(&mut mcp).await?; + assert_eq!(full.active_run_count, full.max_active_runs_per_user); + assert_eq!(full.available_active_run_slots, 0); + assert!(!full.admission_allowed); + + // Admission must self-heal instead of failing closed forever. + let accepted = start_agent( &mut mcp, start_params( - "idempotent retry is not new pressure", - Some("quota-idempotency-0".to_string()), + "new run after a runtime upgrade", + Some("reclaimed-idempotency".to_string()), codex_home.path(), ), ) .await?; - assert_eq!(retry.agent.agent_id, first_agent_id); + wait_for_agent_status( + &mut mcp, + accepted.agent.agent_id.as_str(), + AgentRunStatus::Completed, + ) + .await?; + + for index in 0..8 { + let stranded = state_db + .get_background_agent_run(&format!("stranded-run-{index}")) + .await? + .expect("stranded run should still be readable"); + assert_eq!(stranded.status, StateBackgroundAgentRunStatus::Failed); + assert!( + stranded + .status_reason + .as_deref() + .is_some_and(|reason| reason.contains("runtime package is incompatible")), + "unexpected status reason: {:?}", + stranded.status_reason + ); + } + let reclaimed = agent_daemon_diagnostics(&mut mcp).await?; + assert_eq!(reclaimed.active_run_count, 0); + assert_eq!( + reclaimed.available_active_run_slots, + reclaimed.max_active_runs_per_user + ); + assert!(reclaimed.admission_allowed); Ok(()) } @@ -1929,6 +2084,7 @@ async fn worktree_cleanup_retains_nonterminal_owner_agent_worktree() -> Result<( "agent-run-cleanup-guard", /*idempotency_key*/ None, "guard nonterminal owner cleanup", + SeededAgentExecution::Inert, ) .await?; state_db @@ -2654,6 +2810,7 @@ async fn worktree_attach_assigns_thread_and_rejects_ambiguous_targets() -> Resul "agent-run-attach", /*idempotency_key*/ None, "attach this agent to a worktree", + SeededAgentExecution::Inert, ) .await?; state_db @@ -3011,6 +3168,7 @@ async fn create_background_agent_git_worktree_lease( agent_id, /*idempotency_key*/ None, "release a leased background-agent worktree", + SeededAgentExecution::Inert, ) .await?; let branch = format!("codewith/{agent_id}"); @@ -3312,37 +3470,87 @@ async fn seed_queued_agent_run( agent_id: &str, idempotency_key: Option, prompt: &str, + execution: SeededAgentExecution, ) -> 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 mut execution_payload = 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, + "authProfileIdentitySha256": null, + "managedWorktreeId": null, + "configFingerprint": "cfg-test", + "versionFingerprint": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + "packageFingerprint": "incompatible-test-runtime", + "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", + "midTurnCrashSemantics": "abort_mid_turn_resume_at_safe_boundary", + }); + if let SeededAgentExecution::Claimable { cwd } = execution { + let Some(payload) = execution_payload.as_object_mut() else { + anyhow::bail!("seeded execution payload should be an object"); + }; + payload.insert("workspaceRoots".to_string(), json!([cwd.as_str()])); + payload.insert("cwd".to_string(), json!(cwd)); + payload.insert("model".to_string(), json!("mock-model")); + payload.insert("provider".to_string(), json!("mock_provider")); + payload.insert("serviceTier".to_string(), json!("default")); + payload.insert("approvalPolicy".to_string(), json!("never")); + payload.insert( + "packageFingerprint".to_string(), + json!(BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT), + ); + } + let execution_snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: agent_id.to_string(), + snapshot_kind: "initial_execution_context".to_string(), + payload_json: execution_payload, + 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("version-test".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(()) @@ -3371,7 +3579,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..4d88f0ada5 100644 --- a/codex-rs/background-agent/ARCHITECTURE.md +++ b/codex-rs/background-agent/ARCHITECTURE.md @@ -23,6 +23,81 @@ 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. That check stays +fail-closed but recoverable: `agent start` stops and replaces a mismatched +daemon instead of dead-ending, while read-only daemon status reports the +mismatch together with the exact remediation command. For the same reason, +reconciliation fails closed any queued or orphaned run whose persisted admission +schema or execution-snapshot package fingerprint no longer matches the installed +binary. Such a run can never be claimed again, so releasing its admission slot +(with an explicit lifecycle receipt) is what keeps a routine upgrade from +permanently consuming capacity. 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. + +### Persisted-secret Invariant + +**No persisted background-agent column may hold a recoverable plaintext +secret.** The invariant is enforced structurally rather than column by column: +every caller-supplied string that admission writes to `background_agent_runs` +is produced by `RedactedBackgroundAgentRunColumns::from_params`, and every +execution-snapshot column by +`RedactedBackgroundAgentExecutionSnapshotColumns::from_params`. Both the write +path and the idempotent-replay comparison read those same projections, so a +column cannot be protected on write while being compared raw on replay. Each +column is therefore exactly one of: + +- **Digested** — `idempotency_key` is caller-controlled and opaque, so it is + stored only as a one-way SHA-256 digest; dedupe, replay, and + immutable-identity comparisons all run against that digest, and the plaintext + key is never written to or read back from local state. Lifecycle receipt keys + and `admission_identity_sha256` are likewise digests. Execution snapshots + persist only the auth-profile alias digest, never the alias. +- **Redacted** — `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`, `status_reason`, + `config_fingerprint`, `version_fingerprint`, and the execution-snapshot + `snapshot_kind` / `payload_json` / `recovery_policy` / `config_fingerprint`. + These all have to survive readable — the worker loads the same auth profile, + and the CLI/TUI/app-server surface the refs — so they go through state + redaction, which leaves non-secret values byte-identical and replaces only + credential-shaped material. Event, receipt, and status-snapshot payloads are + redacted as JSON. The thread-binding update redacts the same thread columns. +- **Provably non-secret** — `id` (runtime-generated identifier and the + primary/foreign key joining every background-agent table), the closed enums + (`desired_state`, `status`, `retention_state`), supervisor-side process + bookkeeping (`supervisor_id`, `generation`, `pid`, `pgid`, `job_id`, + `start_token`, `stderr_log_path`, lease ids), and integer timestamps and + counters. None of these can carry caller input. + +Credential or account payloads do not belong in any reference projection. +`background_agent_admission_persists_no_plaintext_secret_in_any_column` guards +the invariant generically: it admits a run whose every caller-supplied string +carries the same secret-shaped token, then scans every cell of every +`background_agent*` table for that token, so a newly added raw binding fails +the suite without anyone having to remember to extend a per-column assertion. + ## Run And Thread Relationship A background-agent run owns background execution. A thread owns transcript and @@ -71,16 +146,45 @@ delete requested, or deleted. ## Ownership And Liveness Supervisors claim runs by writing `supervisor_id`, incrementing `generation`, and -creating a process lease. Worker handles are recorded as `pid`, `pgid`, or -`job_id` when the execution backend owns an OS process. Heartbeats update the -claimed generation. Reconciliation may orphan stale non-terminal runs after the -configured heartbeat timeout, then re-claim only runs whose `desired_state` is -`running`. +creating a process lease. A process supervisor must claim before spawning and +pass that exact supervisor generation to the child; competing hosts therefore +lose the durable compare-and-swap before creating an OS process. Worker handles +are recorded as `pid`, `pgid`, or `job_id` when the execution backend owns an OS +process. Heartbeats update the claimed generation. Reconciliation may orphan +stale non-terminal runs after the configured heartbeat timeout, then re-claim +only runs whose `desired_state` is `running`. 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. +Terminal receipts therefore replay with the same run projection and attempt +binding, while receipt insertion and cursor advancement commit in one state +transaction. +The receipt `attempt` is optional operation metadata and is never synthesized +from the supervisor fencing generation. The authoritative cross-system attempt +remains in immutable admitted `spawn_linkage_json` reached through `runId`. + +Initial prompt delivery uses a deterministic client message id derived from the +run. Recovery scans the append-only rollout for that id (and the exact legacy +prompt shape), records `agent.initialPromptRecorded` only after durable rollout +evidence exists, and submits only when that evidence is absent. A crash after +thread binding or submission can therefore neither drop nor duplicate the +initial prompt. + +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..b55909fc1b 100644 --- a/codex-rs/background-agent/src/daemon.rs +++ b/codex-rs/background-agent/src/daemon.rs @@ -14,6 +14,9 @@ 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; @@ -31,6 +34,15 @@ 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); +/// Exact operator remediation for an incompatible-but-running daemon. +pub const DAEMON_INCOMPATIBLE_REMEDIATION: &str = "run `codewith agent daemon stop` to stop the mismatched daemon, or run \ + `codewith agent start`, which stops and replaces it automatically"; +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) @@ -98,16 +110,41 @@ impl BackgroundAgentDaemon { ) })?; let _lock = acquire_lock(&self.paths.lock_file()).await?; + let mut replaced_stop_report = None; if let Some(record) = read_pid_record(&self.paths.pid_file()).await? { match self.controller.status(&record.handle).await? { WorkerProcessStatus::Running => { - return self - .output( - BackgroundAgentDaemonStatus::AlreadyRunning, - Some(record), - /*stop_report*/ None, - ) - .await; + match ensure_daemon_record_compatible(&record) { + Ok(()) => { + return self + .output( + BackgroundAgentDaemonStatus::AlreadyRunning, + Some(record), + /*stop_report*/ None, + ) + .await; + } + Err(incompatible) => { + // Stay fail-closed but recoverable. Reusing a daemon + // built from a different package, protocol, admission + // schema, or capability set is never allowed, yet a + // routine binary upgrade leaves exactly that daemon + // running and would otherwise dead-end every + // `codewith agent start`. Stop and replace it here; + // durable runs are re-claimed by the replacement. + let stop_report = self.controller.stop(&record.handle).await.map_err( + |stop_error| { + incompatible.context(format!( + "failed to stop the incompatible background-agent daemon \ + automatically; stop it with \ + `codewith agent daemon stop` and retry: {stop_error:#}" + )) + }, + )?; + remove_pid_file(&self.paths.pid_file()).await?; + replaced_stop_report = Some(stop_report); + } + } } WorkerProcessStatus::Missing | WorkerProcessStatus::StalePidRecord => { remove_pid_file(&self.paths.pid_file()).await?; @@ -126,6 +163,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 { @@ -135,12 +175,13 @@ impl BackgroundAgentDaemon { ))), }; } - self.output( - BackgroundAgentDaemonStatus::Started, - Some(record), - /*stop_report*/ None, - ) - .await + let status = if replaced_stop_report.is_some() { + BackgroundAgentDaemonStatus::ReplacedIncompatible + } else { + BackgroundAgentDaemonStatus::Started + }; + self.output(status, Some(record), replaced_stop_report) + .await } pub async fn status(&self) -> Result { @@ -155,7 +196,14 @@ impl BackgroundAgentDaemon { .await; }; let status = match self.controller.status(&record.handle).await? { - WorkerProcessStatus::Running => BackgroundAgentDaemonStatus::Running, + // `status` is read-only and grants no capability, so it reports the + // incompatibility (with the exact remediation command) instead of + // erroring out. Every path that would actually *use* the daemon + // still refuses to reuse an incompatible one. + WorkerProcessStatus::Running => match ensure_daemon_record_compatible(&record) { + Ok(()) => BackgroundAgentDaemonStatus::Running, + Err(_) => BackgroundAgentDaemonStatus::IncompatibleRunning, + }, WorkerProcessStatus::Missing | WorkerProcessStatus::StalePidRecord => { BackgroundAgentDaemonStatus::StalePidRecord } @@ -214,16 +262,31 @@ 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(); + let remediation = matches!(status, BackgroundAgentDaemonStatus::IncompatibleRunning) + .then(|| DAEMON_INCOMPATIBLE_REMEDIATION.to_string()); 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(), stderr_tail, stop_report, + remediation, }) } } @@ -234,6 +297,12 @@ pub enum BackgroundAgentDaemonStatus { Started, AlreadyRunning, Running, + /// A daemon is alive but was built from a package, protocol, admission + /// schema, or capability set this client refuses to reuse. `start` replaces + /// it automatically; read-only callers get [`DAEMON_INCOMPATIBLE_REMEDIATION`]. + IncompatibleRunning, + /// An incompatible daemon was stopped and replaced by a matching one. + ReplacedIncompatible, NotRunning, Stopped, StalePidRecord, @@ -247,11 +316,17 @@ 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, pub stderr_tail: Option, pub stop_report: Option, + /// Actionable operator remediation, populated whenever the reported status + /// is not by itself recoverable. + pub remediation: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -259,6 +334,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 +533,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 +552,108 @@ mod tests { Ok(()) } + #[tokio::test] + async fn daemon_status_reports_incompatible_running_daemon_with_remediation() -> 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 output = daemon.status().await?; + + assert_eq!( + output.status, + BackgroundAgentDaemonStatus::IncompatibleRunning + ); + assert_eq!( + output.remediation.as_deref(), + Some(DAEMON_INCOMPATIBLE_REMEDIATION) + ); + // Read-only status must never mutate the daemon it reports on. + assert_eq!( + controller.status(&existing).await?, + WorkerProcessStatus::Running + ); + let _ = controller.stop(&existing).await; + Ok(()) + } + + #[tokio::test] + async fn daemon_start_replaces_incompatible_running_daemon() -> 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 replacement_path = temp_dir.path().join("replacement.sh"); + fs::write(&replacement_path, "#!/bin/sh\nsleep 60\n").await?; + let mut permissions = std::fs::metadata(&replacement_path)?.permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&replacement_path, permissions)?; + let daemon = BackgroundAgentDaemon::with_controller( + BackgroundAgentDaemonPaths::new(&replacement_path, temp_dir.path()), + controller, + ); + write_pid_record(&daemon.paths.pid_file(), &record).await?; + + let output = daemon.start().await?; + + // An upgraded binary must not dead-end: the mismatched daemon is + // stopped and replaced instead of failing `agent start` outright. + assert_eq!( + output.status, + BackgroundAgentDaemonStatus::ReplacedIncompatible + ); + assert!(output.stop_report.is_some()); + assert_ne!(output.pid, Some(existing.pid)); + assert_ne!( + controller.status(&existing).await?, + WorkerProcessStatus::Running + ); + let after_replacement = daemon.status().await?; + assert_eq!( + after_replacement.status, + BackgroundAgentDaemonStatus::Running + ); + assert_eq!(after_replacement.remediation, None); + let _ = daemon.stop().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 +809,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..46aa79805e 100644 --- a/codex-rs/background-agent/src/lib.rs +++ b/codex-rs/background-agent/src/lib.rs @@ -29,15 +29,36 @@ 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 BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT: &str = concat!( + "codewith.background-agent.runtime.v1:", + env!("CARGO_PKG_VERSION") +); +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 /// ownership, heartbeat, and status independently from loaded app-server /// threads or connected clients. pub trait AgentRunStore { - fn create_run( + fn admit_run( &self, params: BackgroundAgentRunCreateParams, + start_event_payload_json: &serde_json::Value, + execution_snapshot_params: &BackgroundAgentExecutionSnapshotParams, + max_active_runs: i64, ) -> impl Future> + Send; fn get_run( @@ -66,6 +87,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, @@ -75,16 +101,32 @@ 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, ) -> impl Future> + Send; + /// Fails closed unclaimed runs that this binary can never claim because + /// their persisted admission schema or runtime package fingerprint predates + /// the installed build, releasing the admission capacity they hold. + fn terminalize_incompatible_runs(&self) -> impl Future> + Send; + fn claim_supervisor( &self, run_id: &str, supervisor_id: &str, process_lease_id: &str, + required_version_fingerprint: &str, + required_package_fingerprint: &str, ) -> impl Future>> + Send; fn record_execution_handle( @@ -101,11 +143,21 @@ pub trait AgentRunStore { } impl AgentRunStore for codex_state::StateRuntime { - async fn create_run( + async fn admit_run( &self, params: BackgroundAgentRunCreateParams, + start_event_payload_json: &serde_json::Value, + execution_snapshot_params: &BackgroundAgentExecutionSnapshotParams, + max_active_runs: i64, ) -> anyhow::Result { - self.create_background_agent_run(¶ms).await + self.admit_background_agent_run( + ¶ms, + start_event_payload_json, + execution_snapshot_params, + max_active_runs, + ) + .await + .map(|(run, _created, _event, _execution_snapshot, _status_snapshot)| run) } async fn get_run(&self, run_id: &str) -> anyhow::Result> { @@ -138,6 +190,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, @@ -151,19 +211,53 @@ 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 } + async fn terminalize_incompatible_runs(&self) -> anyhow::Result { + self.terminalize_incompatible_background_agent_runs( + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + ) + .await + } + async fn claim_supervisor( &self, run_id: &str, supervisor_id: &str, process_lease_id: &str, + required_version_fingerprint: &str, + required_package_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, + required_package_fingerprint, + ) + .await } async fn record_execution_handle( diff --git a/codex-rs/background-agent/src/process_lifecycle.rs b/codex-rs/background-agent/src/process_lifecycle.rs index abca282ef3..2131da5d1d 100644 --- a/codex-rs/background-agent/src/process_lifecycle.rs +++ b/codex-rs/background-agent/src/process_lifecycle.rs @@ -35,6 +35,7 @@ const DEFAULT_STDERR_TAIL_BYTES: u64 = 4096; pub struct WorkerProcessCommand { pub program: PathBuf, pub args: Vec, + pub env: Vec<(OsString, OsString)>, pub cwd: Option, pub stderr_log_path: PathBuf, } @@ -44,6 +45,7 @@ impl WorkerProcessCommand { Self { program: program.into(), args: Vec::new(), + env: Vec::new(), cwd: None, stderr_log_path: stderr_log_path.into(), } @@ -59,6 +61,11 @@ impl WorkerProcessCommand { self } + pub fn env(mut self, key: impl Into, value: impl Into) -> Self { + self.env.push((key.into(), value.into())); + self + } + pub fn cwd(mut self, cwd: impl Into) -> Self { self.cwd = Some(cwd.into()); self @@ -183,6 +190,7 @@ impl WorkerProcessController { let mut command = Command::new(&request.program); command .args(&request.args) + .envs(request.env) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::from(stderr_log.into_std().await)); diff --git a/codex-rs/background-agent/src/supervisor.rs b/codex-rs/background-agent/src/supervisor.rs index 4d68bc7325..ea12c91bdf 100644 --- a/codex-rs/background-agent/src/supervisor.rs +++ b/codex-rs/background-agent/src/supervisor.rs @@ -14,7 +14,6 @@ use crate::BackgroundAgentExecutionHandleParams; use crate::BackgroundAgentPendingInteractionStatus; use crate::BackgroundAgentRun; use crate::BackgroundAgentRunStatus; -use crate::BackgroundAgentStatusSnapshotParams; use crate::PendingInteractionLedger; use crate::SupervisorReconcileReport; @@ -83,6 +82,10 @@ where .store() .orphan_stale_runs(self.config.heartbeat_timeout) .await?; + // A codewith upgrade makes previously admitted-but-unclaimed runs + // permanently unclaimable; reap them here so their admission capacity + // is released instead of leaking until the user hits the cap. + self.store().terminalize_incompatible_runs().await?; let runs = self .store() .list_runs(Some(self.config.reconcile_limit)) @@ -185,43 +188,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 { @@ -238,6 +235,8 @@ where run_id.as_str(), self.config.supervisor_id.as_str(), process_lease_id.as_str(), + crate::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + crate::BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, ) .await? else { @@ -268,59 +267,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) @@ -418,6 +420,7 @@ mod tests { use super::*; use crate::AgentExecutionHandle; + use crate::BackgroundAgentExecutionSnapshotParams; #[derive(Debug, Clone, Default)] struct RecordingExecution { @@ -550,7 +553,12 @@ mod tests { .into_iter() .map(|event| event.event_type) .collect::>(), - vec!["agent.workerStartFailed".to_string()] + vec![ + "agent.admitted".to_string(), + "agent.started".to_string(), + "agent.claimed".to_string(), + "agent.workerStartFailed".to_string() + ] ); Ok(()) } @@ -653,7 +661,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 @@ -665,6 +673,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() ] @@ -761,6 +771,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() @@ -786,25 +798,45 @@ 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, - }) + .admit_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", + "packageFingerprint": + crate::BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + }), + recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), + config_fingerprint: None, + }, + /*max_active_runs*/ 8, + ) .await } } diff --git a/codex-rs/cli/src/agent_cmd.rs b/codex-rs/cli/src/agent_cmd.rs index cc1dd14515..f1f03fba75 100644 --- a/codex-rs/cli/src/agent_cmd.rs +++ b/codex-rs/cli/src/agent_cmd.rs @@ -4,22 +4,25 @@ 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_PROFILE_MISMATCH; +use codex_background_agent::BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION; +use codex_background_agent::BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT; +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; 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; 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; @@ -356,37 +359,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 @@ -982,6 +958,48 @@ async fn attach_agent(state_db: &StateRuntime, cmd: AgentLogsCommand) -> anyhow: })) } +/// Resolves the auth-profile alias that will be admitted with the run. +/// +/// The admitted alias is always the one config resolution honored +/// (`Config::selected_auth_profile`), never the raw `--auth-profile` string. +/// The previous `.or(auth_profile)` fallback could admit a raw flag value that +/// resolution had *not* selected, so the run carried an alias no layer had +/// vetted; this fails closed instead whenever the two disagree. +/// +/// Scope, deliberately: this checks agreement, not existence. Config resolution +/// validates an alias's *syntax* and trims it, and nothing in codewith — the +/// app-server admission path included — resolves an alias against the +/// credential store before admitting. Adding a store lookup only here would +/// make the CLI stricter than the app-server contract this run is admitted +/// under, so alias existence stays the worker's failure to report. Comparison +/// is against the trimmed request so a value resolution accepted verbatim is +/// not rejected for whitespace alone. +fn resolve_agent_start_auth_profile( + runtime_context: Option<&AgentStartRuntimeContext>, + auth_profile: Option<&str>, +) -> anyhow::Result> { + let Some(context) = runtime_context else { + if let Some(requested) = auth_profile { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH}: cannot resolve auth profile \ + `{requested}` without a loaded configuration" + ); + } + return Ok(None); + }; + if let Some(requested) = auth_profile + && context.auth_profile_ref.as_deref() != Some(requested.trim()) + { + let resolved = context.auth_profile_ref.as_deref().unwrap_or(""); + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH}: requested auth profile `{requested}` \ + is not the profile configuration resolved (`{resolved}`); refusing to admit an alias \ + no configuration layer selected" + ); + } + Ok(context.auth_profile_ref.clone()) +} + async fn start_agent( state_db: &StateRuntime, cmd: AgentStartCommand, @@ -993,19 +1011,9 @@ 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()?; + let daemon_output = background_agent_daemon()?.start().await?; let agent_id = new_agent_id(); let explicit_cwd = cmd.cwd.is_some(); @@ -1016,52 +1024,33 @@ async fn start_agent( let workspace_roots = agent_start_snapshot_workspace_roots(runtime_context, &cwd, explicit_cwd); let permission_profile = agent_start_snapshot_permission_profile(runtime_context, &cwd, explicit_cwd)?; - let auth_profile_ref = runtime_context - .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"); - // 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, - request_id: None, - source: "cli".to_string(), - prompt_snapshot_ref: prompt_snapshot_ref.clone(), - 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: 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()), - }; - let run = retry_on_busy("create background agent run", || { - state_db.create_background_agent_run(&create_params) - }) - .await - .context("failed to create background agent")?; + let auth_profile_ref = resolve_agent_start_auth_profile(runtime_context, auth_profile)?; + 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, - "promptSnapshotRef": prompt_snapshot_ref, + "promptSha256": StateRuntime::background_agent_identity_sha256(prompt.as_bytes()), + "promptSnapshotRef": prompt_snapshot_ref.as_str(), }); - let event = retry_on_busy("append background agent start event", || { - state_db.append_background_agent_event( - agent_id.as_str(), - "agent.started", - &start_event_payload, - ) - }) - .await - .context("failed to append background agent start event")?; let snapshot_params = BackgroundAgentExecutionSnapshotParams { run_id: agent_id.clone(), snapshot_kind: "initial_execution_context".to_string(), @@ -1069,7 +1058,6 @@ async fn start_agent( "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, @@ -1077,34 +1065,68 @@ async fn start_agent( "provider": runtime_context.and_then(|context| context.provider.as_deref()), "serviceTier": runtime_context .and_then(|context| context.service_tier.as_deref()), + "authProfileIdentitySha256": auth_profile_ref.as_deref().map(|profile| { + StateRuntime::background_agent_identity_sha256(profile.as_bytes()) + }), + "managedWorktreeId": null, + "configFingerprint": config_fingerprint.as_str(), + "versionFingerprint": BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + "packageFingerprint": BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, "recoveryPolicy": "abort_mid_turn_resume_at_safe_boundary", }), recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), - config_fingerprint: None, + 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, + request_id: None, + source: "cli".to_string(), + prompt_snapshot_ref: prompt_snapshot_ref.clone(), + 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: auth_profile_ref.clone(), + status_reason: Some("queued by codewith agent start".to_string()), + config_fingerprint: Some(config_fingerprint), + version_fingerprint: Some(BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION.to_string()), }; - retry_on_busy("create background agent execution snapshot", || { - state_db.create_background_agent_execution_snapshot(&snapshot_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 create background agent execution snapshot")?; - let status_snapshot_params = BackgroundAgentStatusSnapshotParams { - run_id: agent_id, - 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) + .context("failed to reconcile stale background agents before admission")?; + // Runs admitted by an older codewith build can never be claimed by this + // binary, so release the admission capacity they still hold instead of + // failing every future start with a capacity error. + retry_on_busy("release incompatible background agent admissions", || { + state_db.terminalize_incompatible_background_agent_runs( + BACKGROUND_AGENT_ADMISSION_SCHEMA_VERSION, + BACKGROUND_AGENT_RUNTIME_COMPATIBILITY_FINGERPRINT, + ) }) .await - .context("failed to create background agent status snapshot")?; - let daemon = background_agent_daemon()?; - let daemon_output = daemon.start().await?; - Ok(json!({ "agent": run_json(run), "created": true, "daemon": daemon_output })) + .context("failed to release incompatible background agent admissions")?; + 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, + ) + }) + .await + .context("failed to admit background agent")?; + Ok(json!({ "agent": run_json(run), "created": created, "daemon": daemon_output })) } fn agent_start_snapshot_workspace_roots( @@ -1156,34 +1178,53 @@ 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 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 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 +1241,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 +1253,7 @@ async fn diagnostics_json(state_db: &StateRuntime) -> anyhow::Result { | BackgroundAgentRunStatus::WaitingOnApproval | BackgroundAgentRunStatus::WaitingOnUser | BackgroundAgentRunStatus::Stopping + | BackgroundAgentRunStatus::Orphaned ) }) .map(|(_, count)| *count) @@ -1283,7 +1325,8 @@ fn background_agent_status_is_terminal(status: BackgroundAgentRunStatus) -> bool fn run_json(run: BackgroundAgentRun) -> Value { json!({ "agentId": run.id, - "idempotencyKey": run.idempotency_key, + // Persisted as a one-way digest; the original key is never stored. + "idempotencyKeySha256": run.idempotency_key, "source": run.source, "promptSnapshotRef": run.prompt_snapshot_ref, "threadId": run.thread_id, @@ -1441,6 +1484,63 @@ mod tests { ); } + #[test] + fn agent_start_admits_only_the_auth_profile_configuration_resolved() { + let context = |auth_profile_ref: Option<&str>| AgentStartRuntimeContext { + cwd: PathBuf::from("/launcher"), + workspace_roots: vec![PathBuf::from("/launcher")], + auth_profile_ref: auth_profile_ref.map(str::to_string), + approval_policy: None, + permission_profile: PermissionProfile::read_only(), + model: None, + provider: None, + service_tier: None, + }; + + // The resolved alias is what gets admitted, with or without the flag. + assert_eq!( + resolve_agent_start_auth_profile(Some(&context(Some("work"))), Some("work")) + .expect("agreeing profile"), + Some("work".to_string()) + ); + assert_eq!( + resolve_agent_start_auth_profile(Some(&context(Some("work"))), None) + .expect("env-resolved profile"), + Some("work".to_string()) + ); + // Config resolution trims the alias, so whitespace alone is not a mismatch. + assert_eq!( + resolve_agent_start_auth_profile(Some(&context(Some("work"))), Some(" work ")) + .expect("trimmed profile"), + Some("work".to_string()) + ); + assert_eq!( + resolve_agent_start_auth_profile(Some(&context(None)), None).expect("no profile"), + None + ); + + // A raw flag that resolution did not select must never reach admission, + // which is exactly what the previous `.or(auth_profile)` fallback did. + for (resolved, requested) in [(Some("work"), "personal"), (None, "personal")] { + let error = resolve_agent_start_auth_profile(Some(&context(resolved)), Some(requested)) + .expect_err("diverging profile must fail closed"); + assert!( + error + .to_string() + .contains(BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH), + "unexpected error: {error}" + ); + } + let error = resolve_agent_start_auth_profile(/*runtime_context*/ None, Some("work")) + .expect_err("an unloaded configuration cannot resolve a profile"); + assert!( + error + .to_string() + .contains(BACKGROUND_AGENT_ADMISSION_PROFILE_MISMATCH), + "unexpected error: {error}" + ); + } + #[test] fn explicit_cwd_snapshot_preserves_read_only_permission_profile() { let runtime_context = AgentStartRuntimeContext { 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..fc8c2a5a48 --- /dev/null +++ b/codex-rs/state/migrations/0061_background_agent_lifecycle_receipts.sql @@ -0,0 +1,84 @@ +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; + +-- Idempotency keys are persisted going forward as one-way SHA-256 digests and +-- auth-profile aliases keep flowing through state redaction, so neither column +-- can hold recoverable secret material. Legacy rows are deliberately left +-- untouched: SQLite cannot compute SHA-256 during migration, the rows below are +-- failed closed anyway, and rewriting them would only re-encode whatever the +-- secrets doctor has already redacted. A legacy key therefore no longer matches +-- a digest lookup, which at worst costs idempotent replay against terminal +-- pre-upgrade runs. + +-- Runs admitted before lifecycle receipts existed cannot prove a complete, +-- replayable admission. Fail them closed during the upgrade instead of +-- leaving them permanently active but unclaimable. +UPDATE background_agent_runs +SET + desired_state = 'stopped', + status = 'failed', + status_reason = 'background agent admission predates durable lifecycle receipts', + updated_at = unixepoch(), + completed_at = COALESCE(completed_at, unixepoch()) +WHERE + admission_ready_at IS NULL + AND status IN ( + 'queued', + 'starting', + 'running', + 'waiting_on_approval', + 'waiting_on_user', + 'stopping', + 'orphaned' + ); + +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/model/background_agent/run.rs b/codex-rs/state/src/model/background_agent/run.rs index 4314a92f1f..916905c733 100644 --- a/codex-rs/state/src/model/background_agent/run.rs +++ b/codex-rs/state/src/model/background_agent/run.rs @@ -273,6 +273,9 @@ impl TryFrom for BackgroundAgentRun { fn try_from(value: BackgroundAgentRunRow) -> Result { Ok(Self { id: value.id, + // `idempotency_key` is persisted as a one-way digest and + // `auth_profile_ref` is persisted through state redaction; neither + // is ever reconstructed back into its original plaintext here. idempotency_key: value.idempotency_key, request_id: value.request_id, source: value.source, diff --git a/codex-rs/state/src/runtime/background_agents/events.rs b/codex-rs/state/src/runtime/background_agents/events.rs index 12588c739c..59a4339d20 100644 --- a/codex-rs/state/src/runtime/background_agents/events.rs +++ b/codex-rs/state/src/runtime/background_agents/events.rs @@ -1,5 +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>, @@ -55,6 +63,282 @@ 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, + event_type: &str, + receipt_key: &str, + generation: i64, + attempt: Option, + diagnostics_json: &serde_json::Value, + now: i64, +) -> anyhow::Result { + 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(), + ) + .await? + { + return Ok(event); + } + + let orphaned_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 receipt_key = ? + "#, + ) + .bind(run_id) + .bind(receipt_key) + .fetch_optional(&mut **tx) + .await?; + if let Some(orphaned_event) = orphaned_event { + let payload_json: serde_json::Value = + serde_json::from_str(orphaned_event.payload_json.as_str())?; + let identity_matches = orphaned_event.event_type == event_type + && payload_json + .get("receiptKey") + .and_then(serde_json::Value::as_str) + == Some(receipt_key) + && payload_json + .get("generation") + .and_then(serde_json::Value::as_i64) + == Some(generation) + && payload_json + .get("attempt") + .and_then(serde_json::Value::as_i64) + == attempt + && payload_json.get("diagnostics") == Some(&diagnostics_json); + if !identity_matches { + anyhow::bail!( + "{BACKGROUND_AGENT_RECEIPT_IDENTITY_MISMATCH}: \ + orphaned receipt event is bound to a different lifecycle operation" + ); + } + 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(orphaned_event.id) + .bind(orphaned_event.seq) + .bind(event_type) + .bind(generation) + .bind(attempt) + .bind(operation_identity_sha256) + .bind(orphaned_event.payload_json.as_str()) + .bind(orphaned_event.created_at) + .execute(&mut **tx) + .await?; + return BackgroundAgentEvent::try_from(orphaned_event); + } + + let mut payload = diagnostics_json + .get("eventPayload") + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + if let Some(event_payload) = diagnostics_json.get("eventPayload") + && !event_payload.is_object() + { + payload.insert("eventPayload".to_string(), event_payload.clone()); + } + payload.insert( + "receiptKey".to_string(), + serde_json::Value::String(receipt_key.to_string()), + ); + payload.insert( + "runId".to_string(), + serde_json::Value::String(run_id.to_string()), + ); + payload.insert("generation".to_string(), generation.into()); + payload.insert( + "attempt".to_string(), + attempt.map_or(serde_json::Value::Null, Into::into), + ); + payload.insert("occurredAt".to_string(), now.into()); + payload.insert("diagnostics".to_string(), diagnostics_json.clone()); + let payload_json = crate::redacted_local_state_json(&serde_json::Value::Object(payload)); + 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.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#" +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, + }) +} + +#[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, @@ -71,6 +355,81 @@ 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, + 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, @@ -110,7 +469,10 @@ LIMIT ? let result = sqlx::query( r#" DELETE FROM background_agent_events -WHERE run_id = ? AND seq < ? +WHERE + run_id = ? + AND seq < ? + AND event_type NOT IN ('agent.started', 'agent.startRecovered') "#, ) .bind(run_id) @@ -158,3 +520,50 @@ GROUP BY r.id Ok(()) } } + +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); + 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/mod.rs b/codex-rs/state/src/runtime/background_agents/mod.rs index 9cbd5f3c94..7c01307343 100644 --- a/codex-rs/state/src/runtime/background_agents/mod.rs +++ b/codex-rs/state/src/runtime/background_agents/mod.rs @@ -8,5 +8,13 @@ mod worktrees; mod tests; pub(in crate::runtime) use events::append_background_agent_event_in_tx; +pub(in crate::runtime) use runs::ExistingBackgroundAgentAdmissionIdentity; +pub(in crate::runtime) use runs::RedactedBackgroundAgentExecutionSnapshotColumns; +pub(in crate::runtime) use runs::background_agent_admission_identity_sha256; +pub(in crate::runtime) use runs::background_agent_idempotency_key_digest; +pub(in crate::runtime) use runs::count_live_or_recoverable_background_agent_runs_in_tx; +pub(in crate::runtime) use runs::insert_background_agent_run_in_tx; +pub(in crate::runtime) use runs::recover_or_validate_background_agent_initial_state_in_tx; +pub(in crate::runtime) use runs::validate_existing_background_agent_admission_in_tx; use super::*; diff --git a/codex-rs/state/src/runtime/background_agents/runs.rs b/codex-rs/state/src/runtime/background_agents/runs.rs index d8d8831fc6..08c70fe154 100644 --- a/codex-rs/state/src/runtime/background_agents/runs.rs +++ b/codex-rs/state/src/runtime/background_agents/runs.rs @@ -1,98 +1,249 @@ use super::*; +use sha2::Digest; +use sha2::Sha256; + +const BACKGROUND_AGENT_ADMISSION_CAPACITY_EXCEEDED: &str = + "background_agent_admission_capacity_exceeded"; +const BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH: &str = + "background_agent_admission_identity_mismatch"; +const BACKGROUND_AGENT_RUNTIME_INCOMPATIBLE_REASON: &str = + "background agent runtime package is incompatible with the installed build"; + +/// One-way storage form for a caller-supplied idempotency key. +/// +/// Idempotency keys are opaque, caller-controlled strings that can carry +/// credential-shaped material, so they are never persisted in a form that can +/// be reversed back into the original value. Dedupe and replay lookups compare +/// digests, which preserves byte-exact identity matching without keeping the +/// plaintext in the local state database. +pub(in crate::runtime) fn background_agent_idempotency_key_digest(value: &str) -> String { + StateRuntime::background_agent_identity_sha256(value.as_bytes()) +} + +#[derive(sqlx::FromRow)] +struct BackgroundAgentSupervisorClaimState { + generation: i64, + desired_state: String, + status: String, + retention_state: String, + admission_identity_sha256: Option, + admission_ready_at: Option, +} + +#[derive(sqlx::FromRow)] +struct BackgroundAgentDeleteState { + supervisor_id: Option, + generation: i64, + status: String, + retention_state: String, + status_reason: Option, +} + +pub(in crate::runtime) enum ExistingBackgroundAgentAdmissionIdentity<'a> { + RunFields, + AdmissionDigest(&'a str), +} 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 { - let idempotency_key = params.idempotency_key.as_deref().map(redact_state_string); - if let Some(idempotency_key) = idempotency_key.as_deref() - && let Some(existing) = self - .get_background_agent_run_by_idempotency_key(idempotency_key) - .await? + self.create_background_agent_run_row(params) + .await + .map(|(run, _created)| run) + } + + 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, + 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, + ExistingBackgroundAgentAdmissionIdentity::AdmissionDigest( + admission_identity_sha256.as_str(), + ), + ) + .await? { - return Ok(existing); + let (event, execution_snapshot_id) = + recover_or_validate_background_agent_initial_state_in_tx( + &mut tx, + existing_id.as_str(), + params, + admission_identity_sha256.as_str(), + start_event_payload_json, + &execution_snapshot_params, + max_active_runs, + ) + .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(); - let spawn_linkage_json = params - .spawn_linkage_json - .as_ref() - .map(redact_state_json_string) - .transpose()?; - let insert_result = sqlx::query( + insert_background_agent_run_in_tx(&mut tx, params, now).await?; + append_background_agent_admission_receipt_in_tx(&mut tx, params.id.as_str(), params, 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#" -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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +UPDATE background_agent_runs +SET admission_identity_sha256 = ?, admission_ready_at = ?, updated_at = ? +WHERE id = ? "#, ) - .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(admission_identity_sha256) .bind(now) .bind(now) - .execute(self.pool.as_ref()) - .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); - } - return Err(err.into()); + .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_background_agent_run_row( + &self, + params: &BackgroundAgentRunCreateParams, + ) -> anyhow::Result<(BackgroundAgentRun, bool)> { + 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 + && let Some(existing_id) = validate_existing_background_agent_admission_in_tx( + &mut tx, + idempotency_key, + params, + ExistingBackgroundAgentAdmissionIdentity::RunFields, + ) + .await? + { + 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)); } - self.get_background_agent_run(params.id.as_str()) + let now = Utc::now().timestamp(); + 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()) + .await? + .ok_or_else(|| anyhow::anyhow!("failed to load background agent run {}", params.id))?; + 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 run {}", params.id)) + .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( @@ -231,7 +382,30 @@ OFFSET ? r#" SELECT status, COUNT(*) as count FROM background_agent_runs -WHERE retention_state = 'active' +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 + ) + ) GROUP BY status "#, ) @@ -260,7 +434,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()) @@ -269,6 +448,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) @@ -327,7 +507,46 @@ 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? + { + tx.commit().await?; + return Ok(Some(event)); + } let result = sqlx::query( r#" UPDATE background_agent_runs @@ -367,11 +586,14 @@ WHERE return Ok(None); } - let event = super::events::append_background_agent_event_in_tx( + let event = super::events::append_background_agent_lifecycle_receipt_in_tx( &mut tx, params.run_id, params.event_type, - params.event_payload_json, + receipt_key.as_str(), + params.generation, + /*attempt*/ None, + &operation_diagnostics_json, now, ) .await?; @@ -400,114 +622,488 @@ WHERE Ok(Some(event)) } - pub async fn bind_background_agent_thread( + pub async fn request_background_agent_stop_for_generation( &self, - params: &BackgroundAgentThreadBindingParams, + 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 operation_diagnostics_json = serde_json::json!({ + "statusReason": status_reason, + "diagnostics": diagnostics_json, + }); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; let result = sqlx::query( r#" UPDATE background_agent_runs SET - thread_id = ?, - thread_store_kind = ?, - thread_store_id = ?, - rollout_path = ?, - updated_at = ? + 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 supervisor_id = ? - AND generation = ? - AND status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') + 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(params.thread_id.as_str()) - .bind(params.thread_store_kind.as_str()) - .bind(params.thread_store_id.as_deref()) - .bind(params.rollout_path.as_deref()) + .bind(BackgroundAgentDesiredState::Stopped.as_str()) + .bind(BackgroundAgentRunStatus::Cancelled.as_str()) + .bind(BackgroundAgentRunStatus::Stopping.as_str()) + .bind(redact_state_string(status_reason)) .bind(now) - .bind(params.run_id.as_str()) - .bind(params.supervisor_id.as_str()) - .bind(params.generation) - .execute(self.pool.as_ref()) + .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?; - Ok(result.rows_affected() > 0) - } - pub async fn set_background_agent_desired_state( - &self, - run_id: &str, - desired_state: BackgroundAgentDesiredState, - ) -> anyhow::Result { - let now = Utc::now().timestamp(); - let result = sqlx::query( + let status = if result.rows_affected() == 0 { + let idempotent_status: Option = sqlx::query_scalar( + r#" +SELECT status +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?; + let Some(status) = idempotent_status else { + tx.commit().await?; + return Ok(false); + }; + status + } else { + 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}"); + super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + run_id, + "agent.stopRequested", + receipt_key.as_str(), + expected_generation, + /*attempt*/ None, + &operation_diagnostics_json, + now, + ) + .await?; + if status == BackgroundAgentRunStatus::Cancelled { + super::interactions::terminalize_active_background_agent_pending_interactions_in_tx( + &mut tx, + run_id, + BackgroundAgentPendingInteractionStatus::Cancelled, + &operation_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#" -UPDATE background_agent_runs -SET desired_state = ?, updated_at = ? -WHERE id = ? +SELECT COUNT(*) +FROM background_agent_pending_interactions +WHERE run_id = ? AND status IN (?, ?) "#, ) - .bind(desired_state.as_str()) - .bind(now) .bind(run_id) - .execute(self.pool.as_ref()) + .bind(BackgroundAgentPendingInteractionStatus::Pending.as_str()) + .bind(BackgroundAgentPendingInteractionStatus::Delivered.as_str()) + .fetch_one(&mut *tx) .await?; - Ok(result.rows_affected() > 0) + 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": operation_diagnostics_json, + }), + }, + now, + ) + .await?; + tx.commit().await?; + Ok(true) } - pub async fn request_background_agent_delete(&self, run_id: &str) -> anyhow::Result { + pub async fn bind_background_agent_thread( + &self, + params: &BackgroundAgentThreadBindingParams, + ) -> anyhow::Result { let now = Utc::now().timestamp(); 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, - delete_after = COALESCE(delete_after, ?), + thread_id = ?, + thread_store_kind = ?, + thread_store_id = ?, + rollout_path = ?, updated_at = ? -WHERE id = ? AND retention_state != ? +WHERE + id = ? + AND supervisor_id = ? + AND generation = ? + AND status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') "#, ) - .bind(BackgroundAgentDesiredState::Deleted.as_str()) - .bind(crate::BackgroundAgentRetentionState::DeleteRequested.as_str()) - .bind(BackgroundAgentRunStatus::Stopping.as_str()) - .bind("delete requested") + // Same invariant as admission: these columns are caller-visible strings, + // so they are written through state redaction rather than raw. + .bind(redact_state_string(params.thread_id.as_str())) + .bind(redact_state_string(params.thread_store_kind.as_str())) + .bind(params.thread_store_id.as_deref().map(redact_state_string)) + .bind(params.rollout_path.as_deref().map(redact_state_string)) .bind(now) - .bind(now) - .bind(run_id) - .bind(crate::BackgroundAgentRetentionState::Deleted.as_str()) + .bind(params.run_id.as_str()) + .bind(params.supervisor_id.as_str()) + .bind(params.generation) .execute(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } - pub async fn orphan_stale_background_agent_runs( + pub async fn create_background_agent_execution_snapshot_for_supervisor( &self, - heartbeat_timeout: Duration, - ) -> anyhow::Result { + params: &BackgroundAgentExecutionSnapshotParams, + supervisor_id: &str, + generation: i64, + ) -> anyhow::Result> { let now = Utc::now().timestamp(); - let timeout_secs = i64::try_from(heartbeat_timeout.as_secs()).unwrap_or(i64::MAX); - let stale_before = now.saturating_sub(timeout_secs); - let orphan_candidates = sqlx::query_as::<_, (String, String, i64)>( + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let current: Option = sqlx::query_scalar( r#" -SELECT id, supervisor_id, generation +SELECT 1 FROM background_agent_runs WHERE - desired_state = ? - AND retention_state = ? - AND supervisor_id IS NOT NULL + 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, + desired_state: BackgroundAgentDesiredState, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let result = sqlx::query( + r#" +UPDATE background_agent_runs +SET desired_state = ?, updated_at = ? +WHERE id = ? + "#, + ) + .bind(desired_state.as_str()) + .bind(now) + .bind(run_id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + 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 = sqlx::query_as::<_, BackgroundAgentDeleteState>( + r#" +SELECT supervisor_id, generation, status, retention_state, status_reason +FROM background_agent_runs +WHERE id = ? + "#, + ) + .bind(run_id) + .fetch_optional(&mut *tx) + .await?; + let Some(BackgroundAgentDeleteState { + supervisor_id, + generation, + status, + retention_state, + status_reason: stored_status_reason, + }) = current + else { + tx.commit().await?; + return Ok(false); + }; + let current_status = BackgroundAgentRunStatus::parse(status.as_str())?; + let replaying = matches!(retention_state.as_str(), "delete_requested" | "deleted"); + let (next_status, status_reason, terminalize_immediately) = if replaying { + let status_reason = stored_status_reason.ok_or_else(|| { + anyhow::anyhow!( + "background agent delete receipt replay is missing its status reason" + ) + })?; + (current_status, status_reason, false) + } else { + 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" + }; + ( + next_status, + status_reason.to_string(), + terminalize_immediately, + ) + }; + if !replaying { + let result = sqlx::query( + r#" +UPDATE background_agent_runs +SET + desired_state = ?, + retention_state = ?, + status = ?, + status_reason = ?, + delete_after = COALESCE(delete_after, ?), + 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.as_str()) + .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", + "statusReason": status_reason.as_str(), + }); + let receipt_key = format!("delete:{generation}"); + 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 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: next_status, + desired_state: BackgroundAgentDesiredState::Deleted, + summary: Some(status_reason.clone()), + pending_interaction_count, + last_event_seq, + payload_json: serde_json::json!({ + "phase": next_status.as_str(), + "reason": "delete_requested", + "statusReason": status_reason.as_str(), + }), + }, + now, + ) + .await?; + tx.commit().await?; + Ok(true) + } + + pub async fn orphan_stale_background_agent_runs( + &self, + heartbeat_timeout: Duration, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let timeout_secs = i64::try_from(heartbeat_timeout.as_secs()).unwrap_or(i64::MAX); + let stale_before = now.saturating_sub(timeout_secs); + let orphan_candidates = sqlx::query_as::<_, (String, String, i64)>( + r#" +SELECT id, supervisor_id, generation +FROM background_agent_runs +WHERE + desired_state = ? + AND retention_state = ? + AND supervisor_id IS NOT NULL AND status IN ('starting', 'running', 'waiting_on_approval', 'waiting_on_user') AND (heartbeat_at IS NULL OR heartbeat_at <= ?) "#, @@ -858,12 +1454,148 @@ WHERE run_id: &str, supervisor_id: &str, process_lease_id: &str, + ) -> anyhow::Result> { + self.claim_background_agent_supervisor_inner( + run_id, + supervisor_id, + process_lease_id, + /*required_version_fingerprint*/ None, + /*required_package_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, + required_package_fingerprint: &str, + ) -> anyhow::Result> { + self.claim_background_agent_supervisor_inner( + run_id, + supervisor_id, + process_lease_id, + Some(required_version_fingerprint), + Some(required_package_fingerprint), + ) + .await + } + + pub async fn background_agent_admission_is_ready( + &self, + run_id: &str, + required_version_fingerprint: &str, + required_package_fingerprint: &str, + ) -> anyhow::Result { + let ready: Option = sqlx::query_scalar( + r#" +SELECT 1 +FROM background_agent_runs +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_lifecycle_receipts r + WHERE + r.run_id = background_agent_runs.id + AND r.event_type = 'agent.admitted' + ) + 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 json_extract(s.payload_json, '$.packageFingerprint') = ? + AND ( + json_extract(s.payload_json, '$.managedWorktreeId') IS NULL + OR EXISTS ( + SELECT 1 + FROM managed_worktree_assignments a + WHERE + a.worktree_id = json_extract( + s.payload_json, + '$.managedWorktreeId' + ) + AND a.agent_run_id = background_agent_runs.id + AND a.detached_at_ms IS NULL + ) + ) + ) + AND EXISTS ( + SELECT 1 + FROM background_agent_status_snapshots s + WHERE s.run_id = background_agent_runs.id + ) + "#, + ) + .bind(run_id) + .bind(required_version_fingerprint) + .bind(required_package_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?; + 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>, + required_package_fingerprint: Option<&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( + let current = sqlx::query_as::<_, BackgroundAgentSupervisorClaimState>( r#" -SELECT generation, desired_state, status, retention_state +SELECT + generation, + desired_state, + status, + retention_state, + admission_identity_sha256, + admission_ready_at FROM background_agent_runs WHERE id = ? "#, @@ -871,13 +1603,24 @@ WHERE id = ? .bind(run_id) .fetch_optional(&mut *tx) .await?; - let Some((current_generation, desired_state, status, retention_state)) = current else { + let Some(BackgroundAgentSupervisorClaimState { + generation: 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); @@ -901,6 +1644,55 @@ 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_lifecycle_receipts r + WHERE + r.run_id = background_agent_runs.id + AND r.event_type = 'agent.admitted' + ) + 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 json_extract(s.payload_json, '$.packageFingerprint') = ? + AND ( + json_extract(s.payload_json, '$.managedWorktreeId') IS NULL + OR EXISTS ( + SELECT 1 + FROM managed_worktree_assignments a + WHERE + a.worktree_id = json_extract( + s.payload_json, + '$.managedWorktreeId' + ) + AND a.agent_run_id = background_agent_runs.id + AND a.detached_at_ms IS NULL + ) + ) + ) + AND EXISTS ( + SELECT 1 + FROM background_agent_status_snapshots s + WHERE s.run_id = background_agent_runs.id + ) + ) + ) "#, ) .bind(supervisor_id) @@ -914,6 +1706,9 @@ WHERE .bind(current_generation) .bind(BackgroundAgentDesiredState::Running.as_str()) .bind(crate::BackgroundAgentRetentionState::Active.as_str()) + .bind(required_version_fingerprint) + .bind(required_version_fingerprint) + .bind(required_package_fingerprint) .execute(&mut *tx) .await?; @@ -948,6 +1743,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, + /*attempt*/ None, + &serde_json::json!({ + "supervisorId": supervisor_id, + "processLeaseId": process_lease_id, + }), + now, + ) + .await?; tx.commit().await?; Ok(Some(generation)) } @@ -1014,6 +1824,25 @@ WHERE run_id = ? AND supervisor_id = ? AND generation = ? .bind(params.generation) .execute(&mut *tx) .await?; + let receipt_key = format!("heartbeat:{}", params.generation); + super::events::append_background_agent_lifecycle_receipt_in_tx( + &mut tx, + params.run_id, + "agent.heartbeat", + receipt_key.as_str(), + params.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, + ) + .await?; } tx.commit().await?; @@ -1148,24 +1977,169 @@ WHERE run_id = ? AND supervisor_id = ? AND generation = ? Ok(result.rows_affected() > 0) } - pub async fn get_background_agent_run_by_idempotency_key( + /// Fails closed any unclaimed run that the currently installed runtime can + /// never claim, so that an upgrade cannot permanently strand admission + /// capacity. + /// + /// A claim requires the persisted admission schema *and* the execution + /// snapshot `packageFingerprint` to match the running binary. Queued and + /// orphaned rows keep consuming a live-or-recoverable capacity slot, so + /// after a version bump those rows would otherwise be both unclaimable and + /// undeletable by any reconciliation pass. Terminalizing them releases the + /// slot and leaves an explicit lifecycle receipt describing why. + pub async fn terminalize_incompatible_background_agent_runs( &self, - idempotency_key: &str, - ) -> anyhow::Result> { - let row = sqlx::query_as::<_, BackgroundAgentRunRow>( + required_version_fingerprint: &str, + required_package_fingerprint: &str, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + // The `NOT EXISTS` / fingerprint predicate below is the exact negation + // of the compatibility clause enforced by + // `claim_background_agent_supervisor_compatible`; keep the two in sync. + let candidates = sqlx::query_scalar::<_, String>( r#" -SELECT - 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, +SELECT id +FROM background_agent_runs +WHERE + desired_state = ? + AND retention_state = ? + AND status IN ('queued', 'orphaned') + AND ( + COALESCE(version_fingerprint, '') != ? + OR NOT 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 json_extract(s.payload_json, '$.packageFingerprint') = ? + ) + ) + "#, + ) + .bind(BackgroundAgentDesiredState::Running.as_str()) + .bind(crate::BackgroundAgentRetentionState::Active.as_str()) + .bind(required_version_fingerprint) + .bind(required_package_fingerprint) + .fetch_all(self.pool.as_ref()) + .await?; + + let mut finalized = 0; + for run_id in candidates { + let mut tx = self.pool.begin().await?; + let result = sqlx::query( + r#" +UPDATE background_agent_runs +SET + desired_state = ?, + status = ?, + status_reason = ?, + crash_reason = COALESCE(crash_reason, ?), + updated_at = ?, + completed_at = COALESCE(completed_at, ?) +WHERE + id = ? + AND desired_state = ? + AND retention_state = ? + AND status IN ('queued', 'orphaned') + AND ( + COALESCE(version_fingerprint, '') != ? + OR NOT 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 json_extract(s.payload_json, '$.packageFingerprint') = ? + ) + ) + "#, + ) + .bind(BackgroundAgentDesiredState::Stopped.as_str()) + .bind(BackgroundAgentRunStatus::Failed.as_str()) + .bind(BACKGROUND_AGENT_RUNTIME_INCOMPATIBLE_REASON) + .bind(BACKGROUND_AGENT_RUNTIME_INCOMPATIBLE_REASON) + .bind(now) + .bind(now) + .bind(run_id.as_str()) + .bind(BackgroundAgentDesiredState::Running.as_str()) + .bind(crate::BackgroundAgentRetentionState::Active.as_str()) + .bind(required_version_fingerprint) + .bind(required_package_fingerprint) + .execute(&mut *tx) + .await?; + + if result.rows_affected() == 0 { + tx.commit().await?; + continue; + } + + sqlx::query( + r#" +UPDATE background_agent_process_leases +SET + status = 'stopped', + exit_reason = COALESCE(exit_reason, ?), + updated_at = ?, + stopped_at = COALESCE(stopped_at, ?) +WHERE run_id = ? AND status != 'stopped' + "#, + ) + .bind(BACKGROUND_AGENT_RUNTIME_INCOMPATIBLE_REASON) + .bind(now) + .bind(now) + .bind(run_id.as_str()) + .execute(&mut *tx) + .await?; + + let payload_json = serde_json::json!({ + "reason": "runtime_package_incompatible", + "requiredVersionFingerprint": required_version_fingerprint, + "requiredPackageFingerprint": required_package_fingerprint, + }); + super::interactions::terminalize_active_background_agent_pending_interactions_in_tx( + &mut tx, + run_id.as_str(), + BackgroundAgentPendingInteractionStatus::Cancelled, + &payload_json, + now, + ) + .await?; + append_terminal_stale_background_agent_status_in_tx( + &mut tx, + run_id.as_str(), + BackgroundAgentRunStatus::Failed, + BACKGROUND_AGENT_RUNTIME_INCOMPATIBLE_REASON, + "agent.failed", + &payload_json, + now, + ) + .await?; + tx.commit().await?; + finalized += 1; + } + + Ok(finalized) + } + + pub async fn get_background_agent_run_by_idempotency_key( + &self, + idempotency_key: &str, + ) -> anyhow::Result> { + let row = sqlx::query_as::<_, BackgroundAgentRunRow>( + r#" +SELECT + 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, worktree_lease_id, @@ -1199,13 +2173,665 @@ FROM background_agent_runs WHERE idempotency_key = ? "#, ) - .bind(idempotency_key) + .bind(background_agent_idempotency_key_digest(idempotency_key)) .fetch_optional(self.pool.as_ref()) .await?; row.map(BackgroundAgentRun::try_from).transpose() } } +async fn ensure_background_agent_capacity_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + max_active_runs: i64, +) -> anyhow::Result<()> { + let active_run_count = count_live_or_recoverable_background_agent_runs_in_tx(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(()) +} + +pub(in crate::runtime) async fn count_live_or_recoverable_background_agent_runs_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, +) -> anyhow::Result { + 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 + .map_err(anyhow::Error::from) +} + +/// Secret-safe projection of every string column that admission writes to +/// `background_agent_runs`. +/// +/// State invariant: **no persisted background-agent column may hold a +/// recoverable plaintext secret.** Every caller-supplied string reaches SQLite +/// through this one struct, which is the enforcement point for that invariant: +/// [`insert_background_agent_run_in_tx`] binds only values produced here, and +/// [`validate_existing_background_agent_admission_in_tx`] compares stored rows +/// against the same projection, so a column cannot be protected on write while +/// being compared raw on replay (or vice versa). +/// +/// Per-column protection rationale: +/// * `idempotency_key` — one-way SHA-256 digest; the raw key is only ever +/// compared against a digest of the caller-supplied key and is never read +/// back, so it must never be recoverable. +/// * every other column — [`redact_state_string`] (or `redact_state_json_string` +/// for `spawn_linkage_json`). These are surfaced on read (TUI/CLI/app-server +/// responses) and are matched verbatim on idempotent replay, so a digest +/// would break both display and replay; pattern redaction keeps non-secret +/// values byte-identical while never persisting secret-shaped material. +/// * `id`, `desired_state`, `status`, `retention_state`, `created_at`, +/// `updated_at` are not projected here: `id` is a runtime-generated +/// identifier and the primary/foreign key joining every background-agent +/// table, the state columns are closed enums owned by this module, and the +/// timestamps are integers. None can carry caller input. +struct RedactedBackgroundAgentRunColumns { + idempotency_key: Option, + request_id: Option, + source: String, + prompt_snapshot_ref: String, + input_snapshot_ref: Option, + thread_id: Option, + thread_store_kind: String, + thread_store_id: Option, + rollout_path: Option, + parent_thread_id: Option, + parent_agent_run_id: Option, + spawn_linkage_json: Option, + auth_profile_ref: Option, + status_reason: Option, + config_fingerprint: Option, + version_fingerprint: Option, +} + +impl RedactedBackgroundAgentRunColumns { + fn from_params(params: &BackgroundAgentRunCreateParams) -> anyhow::Result { + Ok(Self { + idempotency_key: params + .idempotency_key + .as_deref() + .map(background_agent_idempotency_key_digest), + request_id: params.request_id.as_deref().map(redact_state_string), + source: redact_state_string(params.source.as_str()), + prompt_snapshot_ref: redact_state_string(params.prompt_snapshot_ref.as_str()), + input_snapshot_ref: params + .input_snapshot_ref + .as_deref() + .map(redact_state_string), + thread_id: params.thread_id.as_deref().map(redact_state_string), + thread_store_kind: redact_state_string(params.thread_store_kind.as_str()), + thread_store_id: params.thread_store_id.as_deref().map(redact_state_string), + rollout_path: params.rollout_path.as_deref().map(redact_state_string), + parent_thread_id: params.parent_thread_id.as_deref().map(redact_state_string), + parent_agent_run_id: params + .parent_agent_run_id + .as_deref() + .map(redact_state_string), + spawn_linkage_json: params + .spawn_linkage_json + .as_ref() + .map(redact_state_json_string) + .transpose()?, + auth_profile_ref: params.auth_profile_ref.as_deref().map(redact_state_string), + status_reason: params.status_reason.as_deref().map(redact_state_string), + config_fingerprint: params + .config_fingerprint + .as_deref() + .map(redact_state_string), + version_fingerprint: params + .version_fingerprint + .as_deref() + .map(redact_state_string), + }) + } +} + +/// Secret-safe projection of the execution-snapshot columns, for the same +/// reason as [`RedactedBackgroundAgentRunColumns`]: the initial snapshot is +/// written during admission and re-compared byte-for-byte on replay, so write +/// and compare must share one definition. +pub(in crate::runtime) struct RedactedBackgroundAgentExecutionSnapshotColumns { + pub(in crate::runtime) snapshot_kind: String, + pub(in crate::runtime) payload_json: String, + pub(in crate::runtime) recovery_policy: String, + pub(in crate::runtime) config_fingerprint: Option, +} + +impl RedactedBackgroundAgentExecutionSnapshotColumns { + pub(in crate::runtime) fn from_params( + params: &BackgroundAgentExecutionSnapshotParams, + ) -> anyhow::Result { + Ok(Self { + snapshot_kind: redact_state_string(params.snapshot_kind.as_str()), + payload_json: redact_state_json_string(¶ms.payload_json)?, + recovery_policy: redact_state_string(params.recovery_policy.as_str()), + config_fingerprint: params + .config_fingerprint + .as_deref() + .map(redact_state_string), + }) + } +} + +pub(in crate::runtime) async fn insert_background_agent_run_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + params: &BackgroundAgentRunCreateParams, + now: i64, +) -> anyhow::Result<()> { + let columns = RedactedBackgroundAgentRunColumns::from_params(params)?; + 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(columns.idempotency_key.as_deref()) + .bind(columns.request_id.as_deref()) + .bind(columns.source.as_str()) + .bind(columns.prompt_snapshot_ref.as_str()) + .bind(columns.input_snapshot_ref.as_deref()) + .bind(columns.thread_id.as_deref()) + .bind(columns.thread_store_kind.as_str()) + .bind(columns.thread_store_id.as_deref()) + .bind(columns.rollout_path.as_deref()) + .bind(columns.parent_thread_id.as_deref()) + .bind(columns.parent_agent_run_id.as_deref()) + .bind(columns.spawn_linkage_json.as_deref()) + .bind(columns.auth_profile_ref.as_deref()) + .bind(BackgroundAgentDesiredState::Running.as_str()) + .bind(BackgroundAgentRunStatus::Queued.as_str()) + .bind(columns.status_reason.as_deref()) + .bind(columns.config_fingerprint.as_deref()) + .bind(columns.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(()) +} + +pub(in crate::runtime) 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)?) + )) +} + +async fn append_background_agent_admission_receipt_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + run_id: &str, + params: &BackgroundAgentRunCreateParams, + now: i64, +) -> anyhow::Result { + let receipt_identity = params.idempotency_key.as_deref().unwrap_or(run_id); + let receipt_key = format!( + "admission:{:x}", + Sha256::digest(receipt_identity.as_bytes()) + ); + super::events::append_background_agent_lifecycle_receipt_in_tx( + tx, + run_id, + "agent.admitted", + receipt_key.as_str(), + /*generation*/ 0, + /*attempt*/ None, + &serde_json::json!({ + "source": params.source, + "requestId": params.request_id, + "versionFingerprint": params.version_fingerprint, + }), + now, + ) + .await +} + +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 columns = RedactedBackgroundAgentExecutionSnapshotColumns::from_params(params)?; + 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(columns.snapshot_kind.as_str()) + .bind(columns.payload_json.as_str()) + .bind(columns.recovery_policy.as_str()) + .bind(columns.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) +} + +pub(in crate::runtime) async fn recover_or_validate_background_agent_initial_state_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + run_id: &str, + params: &BackgroundAgentRunCreateParams, + admission_identity_sha256: &str, + start_event_payload_json: &serde_json::Value, + execution_snapshot_params: &BackgroundAgentExecutionSnapshotParams, + max_active_runs: i64, +) -> 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, admission_ready_at): (Option, Option) = sqlx::query_as( + "SELECT admission_identity_sha256, admission_ready_at \ + 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" + ); + } + if admission_ready_at.is_none() { + ensure_background_agent_capacity_in_tx(tx, max_active_runs).await?; + } + + 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(); + append_background_agent_admission_receipt_in_tx(tx, run_id, params, now).await?; + 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 if stored_identity.is_none() => { + super::events::append_background_agent_event_in_tx( + tx, + run_id, + "agent.started", + &proposed_event_payload, + now, + ) + .await? + } + None => { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + admitted background agent is missing its authoritative start event" + ); + } + }; + + 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?; + // Same rule as the run columns: replay compares against the persisted, + // secret-safe projection, never against the raw request. + let proposed_snapshot_columns = + RedactedBackgroundAgentExecutionSnapshotColumns::from_params(&execution_snapshot_params)?; + let proposed_snapshot_payload: serde_json::Value = + serde_json::from_str(proposed_snapshot_columns.payload_json.as_str())?; + 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 payload_json != proposed_snapshot_payload + || recovery_policy != proposed_snapshot_columns.recovery_policy + || config_fingerprint != proposed_snapshot_columns.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)) +} + +pub(in crate::runtime) async fn validate_existing_background_agent_admission_in_tx( + tx: &mut sqlx::Transaction<'_, Sqlite>, + idempotency_key: &str, + params: &BackgroundAgentRunCreateParams, + identity: ExistingBackgroundAgentAdmissionIdentity<'_>, +) -> anyhow::Result> { + let existing = sqlx::query_as::< + _, + ( + String, + Option, + String, + String, + Option, + Option, + String, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + ), + >( + r#" +SELECT + id, + 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, + config_fingerprint, + version_fingerprint, + admission_identity_sha256 +FROM background_agent_runs +WHERE idempotency_key = ? + "#, + ) + .bind(background_agent_idempotency_key_digest(idempotency_key)) + .fetch_optional(&mut **tx) + .await?; + let Some(( + existing_id, + 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, + config_fingerprint, + version_fingerprint, + admission_identity_sha256, + )) = existing + else { + return Ok(None); + }; + if let ( + ExistingBackgroundAgentAdmissionIdentity::AdmissionDigest(requested_identity), + Some(stored_identity), + ) = (identity, admission_identity_sha256.as_deref()) + { + if stored_identity != requested_identity { + anyhow::bail!( + "{BACKGROUND_AGENT_ADMISSION_IDENTITY_MISMATCH}: \ + idempotency key is already bound to different prompt or execution context" + ); + } + return Ok(Some(existing_id)); + } + // Compare against the same secret-safe projection that admission persisted; + // reading a column raw here would silently reject every replay of a + // redacted value (and invites the write side to drift back to plaintext). + let requested = RedactedBackgroundAgentRunColumns::from_params(params)?; + let identity_matches = request_id == requested.request_id + && source == requested.source + && prompt_snapshot_ref == requested.prompt_snapshot_ref + && input_snapshot_ref == requested.input_snapshot_ref + && thread_id == requested.thread_id + && thread_store_kind == requested.thread_store_kind + && thread_store_id == requested.thread_store_id + && rollout_path == requested.rollout_path + && parent_thread_id == requested.parent_thread_id + && parent_agent_run_id == requested.parent_agent_run_id + && spawn_linkage_json == requested.spawn_linkage_json + && auth_profile_ref == requested.auth_profile_ref + && config_fingerprint == requested.config_fingerprint + && version_fingerprint == requested.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 +2841,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, + /*attempt*/ 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 +2879,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/snapshots.rs b/codex-rs/state/src/runtime/background_agents/snapshots.rs index d06aa65e13..34cd9a4aeb 100644 --- a/codex-rs/state/src/runtime/background_agents/snapshots.rs +++ b/codex-rs/state/src/runtime/background_agents/snapshots.rs @@ -86,7 +86,7 @@ WHERE run_id = ? params: &BackgroundAgentExecutionSnapshotParams, ) -> anyhow::Result { let now = Utc::now().timestamp(); - let payload_json = redact_state_json_string(¶ms.payload_json)?; + let columns = RedactedBackgroundAgentExecutionSnapshotColumns::from_params(params)?; let mut tx = self.pool.begin().await?; let seq: i64 = sqlx::query_scalar( "SELECT COALESCE(MAX(seq), 0) + 1 FROM background_agent_execution_snapshots WHERE run_id = ?", @@ -109,10 +109,10 @@ INSERT INTO background_agent_execution_snapshots ( ) .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(columns.snapshot_kind.as_str()) + .bind(columns.payload_json.as_str()) + .bind(columns.recovery_policy.as_str()) + .bind(columns.config_fingerprint.as_deref()) .bind(now) .execute(&mut *tx) .await? @@ -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> { diff --git a/codex-rs/state/src/runtime/background_agents/tests.rs b/codex-rs/state/src/runtime/background_agents/tests.rs index 20a5ff5f55..88f4df6184 100644 --- a/codex-rs/state/src/runtime/background_agents/tests.rs +++ b/codex-rs/state/src/runtime/background_agents/tests.rs @@ -80,37 +80,1315 @@ 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:{idempotency_key}: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()), + } +} + +async fn admit_run( + runtime: &StateRuntime, + params: &BackgroundAgentRunCreateParams, + max_active_runs: i64, +) -> anyhow::Result<(BackgroundAgentRun, bool)> { + let snapshot_params = admission_snapshot_params(params); + let (run, created, _, _, _) = runtime + .admit_background_agent_run( + params, + &admission_start_payload(params), + &snapshot_params, + max_active_runs, + ) + .await?; + Ok((run, created)) +} + +fn admission_start_payload(params: &BackgroundAgentRunCreateParams) -> serde_json::Value { + let prompt = format!( + "prompt for {}", + params + .idempotency_key + .as_deref() + .unwrap_or(params.id.as_str()) + ); + json!({ + "cwd": "/tmp/admission-test", + "prompt": prompt, + "promptSnapshotRef": params.prompt_snapshot_ref, + "initialGoalObjective": "test admission", + }) +} + +fn admission_snapshot_params( + params: &BackgroundAgentRunCreateParams, +) -> BackgroundAgentExecutionSnapshotParams { + 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", + "authProfileIdentitySha256": params.auth_profile_ref.as_deref().map(|profile| { + StateRuntime::background_agent_identity_sha256(profile.as_bytes()) + }), + "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(), + } +} + +#[tokio::test] +async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let first = create_run(runtime.as_ref()).await?; + let second = runtime + .create_background_agent_run(&BackgroundAgentRunCreateParams { + id: "run-duplicate".to_string(), + idempotency_key: Some("idem-1".to_string()), + request_id: Some("req-1".to_string()), + source: "cli".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(), + 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: Some(json!({"agentPath": ["reviewer"]})), + auth_profile_ref: Some("profile:default".to_string()), + status_reason: None, + config_fingerprint: Some("cfg-1".to_string()), + version_fingerprint: Some("version-1".to_string()), + }) + .await?; + + assert_eq!(second, first); + assert_eq!( + runtime.list_background_agent_runs(/*limit*/ None).await?, + vec![first] + ); + 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) = + 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) = + 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(), 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!( + 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_lifecycle_receipts \ + WHERE run_id = ? AND event_type = 'agent.admitted'", + ) + .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", + "codex-state:test", + ) + .await? + ); + assert!( + runtime + .claim_background_agent_supervisor_compatible( + first.id.as_str(), + "supervisor-without-admission-receipt", + "lease-without-admission-receipt", + "codewith.background-agent.admission.v1", + "codex-state:test", + ) + .await? + .is_none() + ); + let (_, recovered) = admit_run(runtime.as_ref(), &retry_params, /*max_active_runs*/ 2).await?; + assert!(!recovered); + assert!( + runtime + .background_agent_admission_is_ready( + first.id.as_str(), + "codewith.background-agent.admission.v1", + "codex-state:test", + ) + .await? + ); + 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", + "codex-state:test", + ) + .await? + ); + assert!( + runtime + .claim_background_agent_supervisor_compatible( + first.id.as_str(), + "supervisor-after-corruption", + "lease-after-corruption", + "codewith.background-agent.admission.v1", + "codex-state:test", + ) + .await? + .is_none() + ); + 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?; + admit_run( + runtime.as_ref(), + &admission_params("admitted-1", "admission-key", "profile-a"), + 2, + ) + .await?; + + 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 + .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_retry_uses_immutable_identity_after_thread_binding() +-> 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) = + admit_run(runtime.as_ref(), &first_params, /*max_active_runs*/ 2).await?; + assert!(created); + let generation = runtime + .claim_background_agent_supervisor_compatible( + first.id.as_str(), + "binding-supervisor", + "binding-process-lease", + "codewith.background-agent.admission.v1", + "codex-state:test", + ) + .await? + .expect("admitted run should be claimable"); + assert!( + runtime + .bind_background_agent_thread(&BackgroundAgentThreadBindingParams { + run_id: first.id.clone(), + supervisor_id: "binding-supervisor".to_string(), + generation, + thread_id: "bound-thread".to_string(), + thread_store_kind: "thread-store".to_string(), + thread_store_id: Some("bound-thread".to_string()), + rollout_path: Some("/tmp/bound-rollout.jsonl".to_string()), + }) + .await? + ); + + let retry_params = admission_params("admitted-retry", "admission-key", "profile-a"); + let (retry, created) = + admit_run(runtime.as_ref(), &retry_params, /*max_active_runs*/ 2).await?; + assert!(!created); + assert_eq!(retry.id, first.id); + + let mut conflicting_params = retry_params; + conflicting_params.prompt_snapshot_ref = "inline:different:prompt".to_string(); + let error = admit_run( + runtime.as_ref(), + &conflicting_params, + /*max_active_runs*/ 2, + ) + .await + .expect_err("changed immutable admission identity must be rejected"); + assert!( + error + .to_string() + .contains("background_agent_admission_identity_mismatch") + ); + Ok(()) +} + +#[tokio::test] +async fn background_agent_admission_never_persists_recoverable_identity_plaintext() +-> anyhow::Result<()> { + let sqlite_home = unique_temp_dir(); + let runtime = StateRuntime::init(sqlite_home.clone(), "test-provider".to_string()).await?; + let idempotency_key = format!("{}{}", "sk-proj-", "a".repeat(32)); + let auth_profile_ref = format!("{}{}", "sk-proj-", "b".repeat(32)); + assert!(crate::local_state_string_contains_secret(&idempotency_key)); + assert!(crate::local_state_string_contains_secret(&auth_profile_ref)); + let mut params = admission_params( + "opaque-admission", + idempotency_key.as_str(), + auth_profile_ref.as_str(), + ); + params.request_id = Some("opaque-admission-request".to_string()); + params.prompt_snapshot_ref = "inline:opaque-admission:prompt".to_string(); + params.thread_id = Some("thread-opaque-admission".to_string()); + let (run, created) = admit_run(runtime.as_ref(), ¶ms, /*max_active_runs*/ 2).await?; + + assert!(created); + // The idempotency key is caller-controlled and opaque, so it is persisted + // only as a one-way digest and is never reconstructed on read. + let expected_digest = + StateRuntime::background_agent_identity_sha256(idempotency_key.as_bytes()); + assert_eq!( + run.idempotency_key.as_deref(), + Some(expected_digest.as_str()) + ); + assert_ne!( + run.idempotency_key.as_deref(), + Some(idempotency_key.as_str()) + ); + // The auth-profile alias stays usable for config loading, but it is written + // through state redaction so a secret-shaped alias never lands in plaintext. + let expected_redacted_profile = crate::redact_local_state_string(&auth_profile_ref); + assert_ne!(expected_redacted_profile, auth_profile_ref); + assert_eq!( + run.auth_profile_ref.as_deref(), + Some(expected_redacted_profile.as_str()) + ); + let stored = sqlx::query_as::<_, (String, String)>( + "SELECT idempotency_key, auth_profile_ref FROM background_agent_runs WHERE id = ?", + ) + .bind(run.id.as_str()) + .fetch_one(runtime.pool.as_ref()) + .await?; + assert_eq!(stored.0, expected_digest); + assert_eq!(stored.1, expected_redacted_profile); + assert!(!crate::local_state_string_contains_secret(&stored.0)); + assert!(!crate::local_state_string_contains_secret(&stored.1)); + let execution_snapshot = runtime + .get_background_agent_initial_execution_snapshot(run.id.as_str()) + .await? + .expect("admission snapshot should exist"); + assert_eq!( + execution_snapshot + .payload_json + .get("authProfileIdentitySha256"), + Some(&json!(StateRuntime::background_agent_identity_sha256( + auth_profile_ref.as_bytes() + ))) + ); + assert!(!serde_json::to_string(&execution_snapshot.payload_json)?.contains(&auth_profile_ref)); + // The secrets doctor must never find plaintext to redact here, because + // admission already refused to write any. + let doctor_report = + crate::run_local_state_secrets_doctor(crate::LocalStateSecretsDoctorOptions { + codex_home: sqlite_home.clone(), + sqlite_home, + repair: true, + }) + .await?; + assert_eq!(doctor_report.redacted_sqlite_cells, 0); + assert!(!doctor_report.has_findings()); + let after_doctor = runtime + .get_background_agent_run(run.id.as_str()) + .await? + .expect("run should survive secrets repair"); + assert_eq!( + after_doctor.idempotency_key.as_deref(), + Some(expected_digest.as_str()) + ); + assert_eq!( + after_doctor.auth_profile_ref.as_deref(), + Some(expected_redacted_profile.as_str()) + ); + // Byte-exact idempotent replay still works off the digest. + assert_eq!( + runtime + .get_background_agent_run_by_idempotency_key(idempotency_key.as_str()) + .await? + .map(|run| run.id), + Some(run.id.clone()) + ); + assert_eq!( + runtime + .get_background_agent_run_by_idempotency_key( + format!("{}{}", "sk-proj-", "c".repeat(32)).as_str() + ) + .await? + .map(|run| run.id), + None + ); + Ok(()) +} + +/// Generic guard for the background-agent state invariant: *no* persisted +/// column may hold a recoverable plaintext secret. +/// +/// Per-column assertions have twice passed while a sibling column still +/// leaked, so this test asserts the invariant structurally instead: every +/// caller-supplied string on the admission request carries the same +/// secret-shaped token, and afterwards every cell of every `background_agent*` +/// table is scanned for that token. A newly added column therefore fails here +/// automatically unless it is digested, redacted, or provably non-secret. +#[tokio::test] +async fn background_agent_admission_persists_no_plaintext_secret_in_any_column() +-> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let secret = format!("{}{}", "sk-proj-", "W5wN".repeat(12)); + assert!(crate::local_state_string_contains_secret(&secret)); + let tainted = |field: &str| format!("inline:{secret}:{field}"); + + let mut params = admission_params( + "plaintext-sweep-run", + tainted("idempotency-key").as_str(), + tainted("auth-profile").as_str(), + ); + // `id` is the one string the caller does not choose: it is generated by the + // runtime and is the primary/foreign key joining every background-agent + // table, so it must stay verbatim. Everything else is caller-controlled and + // is therefore tainted here. + params.request_id = Some(tainted("request-id")); + params.source = tainted("source"); + params.prompt_snapshot_ref = tainted("prompt-snapshot-ref"); + params.input_snapshot_ref = Some(tainted("input-snapshot-ref")); + params.thread_id = Some(tainted("thread-id")); + params.thread_store_kind = tainted("thread-store-kind"); + params.thread_store_id = Some(tainted("thread-store-id")); + params.rollout_path = Some(tainted("rollout-path")); + params.parent_thread_id = Some(tainted("parent-thread-id")); + params.parent_agent_run_id = Some(tainted("parent-agent-run-id")); + params.spawn_linkage_json = Some(json!({"agentPath": [tainted("spawn-linkage")]})); + params.status_reason = Some(tainted("status-reason")); + params.config_fingerprint = Some(tainted("config-fingerprint")); + params.version_fingerprint = Some(tainted("version-fingerprint")); + + let mut snapshot_params = admission_snapshot_params(¶ms); + snapshot_params.recovery_policy = tainted("recovery-policy"); + let start_payload = admission_start_payload(¶ms); + let (run, created, _, _, _) = runtime + .admit_background_agent_run( + ¶ms, + &start_payload, + &snapshot_params, + /*max_active_runs*/ 2, + ) + .await?; + assert!(created); + + let tables = sqlx::query_scalar::<_, String>( + "SELECT name FROM sqlite_master \ + WHERE type = 'table' AND name LIKE 'background\\_agent%' ESCAPE '\\' \ + ORDER BY name", + ) + .fetch_all(runtime.pool.as_ref()) + .await?; + assert!( + tables.len() >= 5, + "expected the background-agent schema to be present, found {tables:?}" + ); + + let mut leaks: Vec = Vec::new(); + let mut scanned_columns = 0usize; + let mut conn = runtime.pool.acquire().await?; + for table in &tables { + let columns = + sqlx::query_scalar::<_, String>("SELECT name FROM pragma_table_info(?) ORDER BY cid") + .bind(table.clone()) + .fetch_all(&mut *conn) + .await?; + for column in &columns { + scanned_columns += 1; + // The sqlite driver requires `'static` SQL, and the table/column + // names come from the schema (never from input), so leaking the + // handful of generated scan statements is safe inside this test. + let scan_sql: &'static str = Box::leak( + format!( + "SELECT COUNT(*) FROM \"{table}\" \ + WHERE instr(CAST(\"{column}\" AS TEXT), ?) > 0" + ) + .into_boxed_str(), + ); + let hits = sqlx::query_scalar::<_, i64>(scan_sql) + .bind(secret.clone()) + .fetch_one(&mut *conn) + .await?; + if hits > 0 { + leaks.push(format!("{table}.{column} ({hits} row(s))")); + } + } + } + drop(conn); + assert!(scanned_columns > 0, "no background-agent columns scanned"); + assert!( + leaks.is_empty(), + "recoverable plaintext secret persisted in: {}", + leaks.join(", ") + ); + + // The invariant must hold without disabling the feature: the run is still + // readable and its refs still round-trip (redacted, not dropped). + let stored = runtime + .get_background_agent_run(run.id.as_str()) + .await? + .expect("admitted run should be readable"); + assert_eq!( + stored.prompt_snapshot_ref, + crate::redact_local_state_string(tainted("prompt-snapshot-ref")) + ); + assert!(!crate::local_state_string_contains_secret( + &stored.prompt_snapshot_ref + )); + // Idempotent replay still resolves the same run from the raw key. + let (replayed, replayed_created, _, _, _) = runtime + .admit_background_agent_run( + ¶ms, + &start_payload, + &snapshot_params, + /*max_active_runs*/ 2, + ) + .await?; + assert!(!replayed_created); + assert_eq!(replayed.id, run.id); + Ok(()) +} + +#[tokio::test] +async fn background_agent_runtime_upgrade_reclaims_stranded_admission_capacity() +-> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let params = admission_params("stranded-run", "stranded-key", "profile:default"); + let (run, created) = admit_run(runtime.as_ref(), ¶ms, /*max_active_runs*/ 1).await?; + assert!(created); + + // The daemon dies mid-run, so the row goes to `orphaned` while still + // holding the single admission slot. + runtime + .claim_background_agent_supervisor(run.id.as_str(), "supervisor-old", "lease-old") + .await? + .expect("queued run should be claimable by the admitting runtime"); + assert_eq!( + runtime + .orphan_stale_background_agent_runs(Duration::ZERO) + .await?, + 1 + ); + assert_eq!( + runtime + .get_background_agent_run(run.id.as_str()) + .await? + .map(|run| run.status), + Some(BackgroundAgentRunStatus::Orphaned) + ); + + // Operator upgrades codewith: the persisted package fingerprint no longer + // matches the running binary, so the row can never be claimed again. + let upgraded_package_fingerprint = "codex-state:test-upgraded"; + let upgraded_version_fingerprint = params + .version_fingerprint + .as_deref() + .expect("admission fixtures set a version fingerprint"); + assert_eq!( + runtime + .claim_background_agent_supervisor_compatible( + run.id.as_str(), + "supervisor-new", + "lease-new", + upgraded_version_fingerprint, + upgraded_package_fingerprint, + ) + .await?, + None + ); + let capacity_error = admit_run( + runtime.as_ref(), + &admission_params("post-upgrade-run", "post-upgrade-key", "profile:default"), + /*max_active_runs*/ 1, + ) + .await + .expect_err("the stranded run must still be holding the only admission slot"); + assert!( + capacity_error + .to_string() + .contains("background_agent_admission_capacity_exceeded"), + "unexpected error: {capacity_error}" + ); + + // Reconciliation must release the slot instead of leaking it forever. + assert_eq!( + runtime + .terminalize_incompatible_background_agent_runs( + upgraded_version_fingerprint, + upgraded_package_fingerprint, + ) + .await?, + 1 + ); + let stranded = runtime + .get_background_agent_run(run.id.as_str()) + .await? + .expect("stranded run should still be readable"); + assert_eq!(stranded.status, BackgroundAgentRunStatus::Failed); + assert_eq!(stranded.desired_state, BackgroundAgentDesiredState::Stopped); + assert!(stranded.completed_at.is_some()); + assert!( + stranded + .status_reason + .as_deref() + .is_some_and(|reason| reason.contains("runtime package is incompatible")), + "unexpected status reason: {:?}", + stranded.status_reason + ); + let receipts: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM background_agent_lifecycle_receipts \ + WHERE run_id = ? AND event_type = 'agent.failed'", + ) + .bind(run.id.as_str()) + .fetch_one(runtime.pool.as_ref()) + .await?; + assert_eq!(receipts, 1); + + // A repeat pass is a no-op, and new admissions succeed again. + assert_eq!( + runtime + .terminalize_incompatible_background_agent_runs( + upgraded_version_fingerprint, + upgraded_package_fingerprint, + ) + .await?, + 0 + ); + let (recovered, created) = admit_run( + runtime.as_ref(), + &admission_params("post-upgrade-run", "post-upgrade-key", "profile:default"), + /*max_active_runs*/ 1, + ) + .await?; + assert!(created); + assert_eq!(recovered.status, BackgroundAgentRunStatus::Queued); + 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?; + 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(), + Some("supervisor-1"), + generation, + "capacity test stop", + &json!({"reason": "capacity_test"}), + ) + .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 + .update_background_agent_run_status( + "admitted-2", + BackgroundAgentRunStatus::Orphaned, + Some("recoverable orphan"), + ) + .await?; + 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() + .contains("background_agent_admission_capacity_exceeded") + ); + Ok(()) +} + +#[tokio::test] +async fn partial_admission_recovery_cannot_bypass_full_capacity() -> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let partial_params = admission_params("partial-run", "partial-key", "profile-a"); + runtime.create_background_agent_run(&partial_params).await?; + let (active_run, created) = admit_run( + runtime.as_ref(), + &admission_params("active-run", "active-key", "profile-a"), + /*max_active_runs*/ 1, + ) + .await?; + assert!(created); + + let error = admit_run( + runtime.as_ref(), + &partial_params, + /*max_active_runs*/ 1, + ) + .await + .expect_err("recovering a partial admission must consume capacity atomically"); + assert!( + error + .to_string() + .contains("background_agent_admission_capacity_exceeded") + ); + assert!( + runtime + .list_background_agent_events_after( + "partial-run", + /*after_seq*/ None, + /*limit*/ None, + ) + .await? + .is_empty() + ); + assert!( + !runtime + .background_agent_admission_is_ready( + "partial-run", + "codewith.background-agent.admission.v1", + "codex-state:test", + ) + .await? + ); + + assert!( + runtime + .request_background_agent_stop_for_generation( + active_run.id.as_str(), + /*expected_supervisor_id*/ None, + active_run.generation, + "capacity released for partial recovery", + &json!({"reason": "capacity_released"}), + ) + .await? + ); + let (recovered, created) = admit_run( + runtime.as_ref(), + &partial_params, + /*max_active_runs*/ 1, + ) + .await?; + assert!(!created); + assert_eq!(recovered.id, "partial-run"); + assert_eq!( + runtime + .list_background_agent_events_after( + "partial-run", + /*after_seq*/ None, + /*limit*/ None, + ) + .await? + .into_iter() + .map(|event| event.event_type) + .collect::>(), + vec!["agent.admitted".to_string(), "agent.started".to_string()] + ); + assert!( + runtime + .background_agent_admission_is_ready( + "partial-run", + "codewith.background-agent.admission.v1", + "codex-state:test", + ) + .await? + ); + Ok(()) +} + +#[tokio::test] +async fn partial_admission_recovery_requires_exact_persisted_execution_identity() +-> anyhow::Result<()> { + let field_cases = [ + ("permission missing", "permissionProfile", None), + ( + "permission different", + "permissionProfile", + Some(json!({"type": "danger-full-access"})), + ), + ("model missing", "model", None), + ("model different", "model", Some(json!("other-model"))), + ("workspace missing", "workspaceRoots", None), + ( + "workspace different", + "workspaceRoots", + Some(json!(["/tmp/other-workspace"])), + ), + ("provider missing", "provider", None), + ( + "provider different", + "provider", + Some(json!("other-provider")), + ), + ("service tier missing", "serviceTier", None), + ( + "service tier different", + "serviceTier", + Some(json!("priority")), + ), + ( + "auth profile identity missing", + "authProfileIdentitySha256", + None, + ), + ( + "auth profile identity different", + "authProfileIdentitySha256", + Some(json!(StateRuntime::background_agent_identity_sha256( + b"profile-b" + ))), + ), + ("package missing", "packageFingerprint", None), + ( + "package different", + "packageFingerprint", + Some(json!("codex-state:other")), + ), + ("recovery policy missing", "recoveryPolicy", None), + ( + "recovery policy different", + "recoveryPolicy", + Some(json!("restart_from_beginning")), + ), + ]; + for (index, (case, field, stored_value)) in field_cases.into_iter().enumerate() { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let params = admission_params( + format!("partial-identity-{index}").as_str(), + format!("partial-identity-key-{index}").as_str(), + "profile-a", + ); + runtime.create_background_agent_run(¶ms).await?; + let mut stored_snapshot = admission_snapshot_params(¶ms); + let stored_payload = stored_snapshot + .payload_json + .as_object_mut() + .expect("test execution snapshot must be an object"); + match stored_value { + Some(value) => { + stored_payload.insert(field.to_string(), value); + } + None => { + stored_payload.remove(field); + } + } + runtime + .create_background_agent_execution_snapshot(&stored_snapshot) + .await?; + + let error = admit_run(runtime.as_ref(), ¶ms, /*max_active_runs*/ 1) + .await + .expect_err(case); + assert!( + error + .to_string() + .contains("background_agent_admission_identity_mismatch"), + "{case}: {error}" + ); + } + + for (index, (case, recovery_policy, config_fingerprint)) in [ + ( + "recovery policy column different", + "restart_from_beginning", + Some("config-v1"), + ), + ( + "config fingerprint column missing", + "abort_mid_turn_resume_at_safe_boundary", + None, + ), + ( + "config fingerprint column different", + "abort_mid_turn_resume_at_safe_boundary", + Some("config-v2"), + ), + ] + .into_iter() + .enumerate() + { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let params = admission_params( + format!("partial-column-{index}").as_str(), + format!("partial-column-key-{index}").as_str(), + "profile-a", + ); + runtime.create_background_agent_run(¶ms).await?; + let mut stored_snapshot = admission_snapshot_params(¶ms); + stored_snapshot.recovery_policy = recovery_policy.to_string(); + stored_snapshot.config_fingerprint = config_fingerprint.map(str::to_string); + runtime + .create_background_agent_execution_snapshot(&stored_snapshot) + .await?; + + let error = admit_run(runtime.as_ref(), ¶ms, /*max_active_runs*/ 1) + .await + .expect_err(case); + assert!( + error + .to_string() + .contains("background_agent_admission_identity_mismatch"), + "{case}: {error}" + ); + } + Ok(()) +} + #[tokio::test] -async fn background_agent_run_create_is_idempotent() -> anyhow::Result<()> { +async fn compatible_claim_rejects_persisted_runtime_package_skew() -> anyhow::Result<()> { let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; - let first = create_run(runtime.as_ref()).await?; - let second = runtime - .create_background_agent_run(&BackgroundAgentRunCreateParams { - id: "run-duplicate".to_string(), - idempotency_key: Some("idem-1".to_string()), - request_id: Some("req-duplicate".to_string()), - source: "cli".to_string(), - prompt_snapshot_ref: "prompt://duplicate".to_string(), - 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: None, - config_fingerprint: None, - version_fingerprint: None, - }) + let (run, created) = admit_run( + runtime.as_ref(), + &admission_params("package-run", "package-key", "profile-a"), + /*max_active_runs*/ 1, + ) + .await?; + assert!(created); + assert!( + !runtime + .background_agent_admission_is_ready( + run.id.as_str(), + "codewith.background-agent.admission.v1", + "codex-state:next", + ) + .await? + ); + assert!( + runtime + .claim_background_agent_supervisor_compatible( + run.id.as_str(), + "newer-supervisor", + "newer-lease", + "codewith.background-agent.admission.v1", + "codex-state:next", + ) + .await? + .is_none() + ); + assert!( + runtime + .claim_background_agent_supervisor_compatible( + run.id.as_str(), + "compatible-supervisor", + "compatible-lease", + "codewith.background-agent.admission.v1", + "codex-state:test", + ) + .await? + .is_some() + ); + 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 conflict = runtime + .append_background_agent_lifecycle_receipt( + run.id.as_str(), + "agent.testReceipt", + "test-receipt", + 1, + 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!(second, first); + 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!( - runtime.list_background_agent_runs(/*limit*/ None).await?, - vec![first] + retry + .payload_json + .pointer("/diagnostics/truncated") + .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 stop_and_delete_retries_replay_exact_lifecycle_operation_identity() -> anyhow::Result<()> { + let runtime = StateRuntime::init(unique_temp_dir(), "test-provider".to_string()).await?; + let run = create_run_with_id(runtime.as_ref(), "stop-retry").await?; + let stop_diagnostics = json!({"reason": "operator_requested"}); + assert!( + runtime + .request_background_agent_stop_for_generation( + run.id.as_str(), + /*expected_supervisor_id*/ None, + run.generation, + "operator requested stop", + &stop_diagnostics, + ) + .await? + ); + let first_stop_events = runtime + .list_background_agent_events_after(run.id.as_str(), None, None) + .await?; + assert!( + runtime + .request_background_agent_stop_for_generation( + run.id.as_str(), + /*expected_supervisor_id*/ None, + run.generation, + "operator requested stop", + &stop_diagnostics, + ) + .await? + ); + assert_eq!( + runtime + .list_background_agent_events_after(run.id.as_str(), None, None) + .await?, + first_stop_events + ); + for (status_reason, diagnostics) in [ + ("different stop reason", stop_diagnostics.clone()), + ( + "operator requested stop", + json!({"reason": "different_request"}), + ), + ] { + let error = runtime + .request_background_agent_stop_for_generation( + run.id.as_str(), + /*expected_supervisor_id*/ None, + run.generation, + status_reason, + &diagnostics, + ) + .await + .expect_err("conflicting stop receipt retry must fail closed"); + assert!( + error + .to_string() + .contains("background agent lifecycle receipt identity mismatch") + ); + } + + let delete_run = create_run_with_id(runtime.as_ref(), "delete-retry").await?; + assert!( + runtime + .request_background_agent_delete(delete_run.id.as_str()) + .await? + ); + let first_delete_events = runtime + .list_background_agent_events_after(delete_run.id.as_str(), None, None) + .await?; + assert!( + runtime + .request_background_agent_delete(delete_run.id.as_str()) + .await? + ); + assert_eq!( + runtime + .list_background_agent_events_after(delete_run.id.as_str(), None, None) + .await?, + first_delete_events + ); + sqlx::query("UPDATE background_agent_runs SET status_reason = ? WHERE id = ?") + .bind("conflicting delete reason") + .bind(delete_run.id.as_str()) + .execute(runtime.pool.as_ref()) + .await?; + let error = runtime + .request_background_agent_delete(delete_run.id.as_str()) + .await + .expect_err("conflicting persisted delete reason must fail receipt replay"); + assert!( + error + .to_string() + .contains("background agent lifecycle receipt identity mismatch") + ); + 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(()) } @@ -754,7 +2032,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(()) } @@ -819,7 +2102,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( @@ -830,6 +2113,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() @@ -866,6 +2151,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) @@ -908,6 +2219,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") @@ -924,6 +2302,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?; @@ -1016,7 +2451,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? @@ -1369,7 +2804,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, 1); + 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(()) } diff --git a/codex-rs/state/src/runtime/workflow_orchestrator.rs b/codex-rs/state/src/runtime/workflow_orchestrator.rs index d5043132e7..1b89360601 100644 --- a/codex-rs/state/src/runtime/workflow_orchestrator.rs +++ b/codex-rs/state/src/runtime/workflow_orchestrator.rs @@ -1,5 +1,12 @@ use super::*; +use crate::runtime::background_agents::ExistingBackgroundAgentAdmissionIdentity; use crate::runtime::background_agents::append_background_agent_event_in_tx; +use crate::runtime::background_agents::background_agent_admission_identity_sha256; +use crate::runtime::background_agents::background_agent_idempotency_key_digest; +use crate::runtime::background_agents::count_live_or_recoverable_background_agent_runs_in_tx; +use crate::runtime::background_agents::insert_background_agent_run_in_tx; +use crate::runtime::background_agents::recover_or_validate_background_agent_initial_state_in_tx; +use crate::runtime::background_agents::validate_existing_background_agent_admission_in_tx; use crate::runtime::workflow_automation::arm_workflow_timers_for_succeeded_step_in_tx; use crate::runtime::workflows::WorkflowRunEventAppend; use crate::runtime::workflows::append_workflow_run_event_in_tx; @@ -471,7 +478,8 @@ async fn admit_ready_workflow_branches_in_tx( .saturating_sub(active_counts.branch_count) .min(limits.max_agents.saturating_sub(active_counts.branch_count)); if let Some(max_active_background_agent_runs) = params.max_active_background_agent_runs { - let active_background_runs = active_background_agent_run_count_in_tx(tx).await?; + let active_background_runs = + count_live_or_recoverable_background_agent_runs_in_tx(tx).await?; capacity = capacity.min(max_active_background_agent_runs.saturating_sub(active_background_runs)); } @@ -917,22 +925,6 @@ WHERE run_id = ? }) } -async fn active_background_agent_run_count_in_tx( - tx: &mut sqlx::Transaction<'_, Sqlite>, -) -> anyhow::Result { - sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM background_agent_runs -WHERE retention_state = 'active' - AND status NOT IN ('completed', 'failed', 'cancelled') - "#, - ) - .fetch_one(&mut **tx) - .await - .map_err(anyhow::Error::from) -} - async fn ready_branch_candidates_in_tx( tx: &mut sqlx::Transaction<'_, Sqlite>, run_id: &str, @@ -1098,7 +1090,7 @@ FROM background_agent_runs WHERE idempotency_key = ? "#, ) - .bind(idempotency_key) + .bind(background_agent_idempotency_key_digest(idempotency_key)) .fetch_optional(&mut **tx) .await .map_err(anyhow::Error::from) @@ -1119,20 +1111,6 @@ async fn create_background_branch_run_if_missing_in_tx( now_ms, } = branch; let now = now_ms.div_euclid(1000); - let existing: Option = sqlx::query_scalar( - r#" -SELECT id -FROM background_agent_runs -WHERE id = ? - "#, - ) - .bind(background_agent_run_id) - .fetch_optional(&mut **tx) - .await?; - if existing.is_some() { - return Ok(false); - } - let prompt = render_workflow_branch_prompt(WorkflowBranchPrompt { run_id: run.run_id.as_str(), step_id: candidate.step_id.as_str(), @@ -1141,8 +1119,7 @@ WHERE id = ? parallel_group: candidate.parallel_group.as_deref(), }); let prompt_snapshot_ref = format!("workflow:{}:step:{}:prompt", run.run_id, candidate.step_id); - let source = WORKFLOW_BRANCH_SOURCE.to_string(); - let spawn_linkage_json = redact_state_json_string(&json!({ + let spawn_linkage_json = json!({ "schemaVersion": "workflow.branch_spawn/v0", "redactionVersion": 1, "workflowRunId": run.run_id.as_str(), @@ -1152,69 +1129,73 @@ WHERE id = ? "stepRunId": candidate.step_run_id.as_str(), "agentId": candidate.agent_id.as_str(), "parallelGroup": candidate.parallel_group.as_deref(), - }))?; - 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 (?, ?, NULL, ?, ?, NULL, NULL, ?, NULL, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(background_agent_run_id) - .bind(redact_state_string(idempotency_key)) - .bind(source) - .bind(redact_state_string(prompt_snapshot_ref.as_str())) - .bind(WORKFLOW_BRANCH_THREAD_STORE_KIND) - .bind( - run.source_thread_id + }); + let run_params = BackgroundAgentRunCreateParams { + id: background_agent_run_id.to_string(), + idempotency_key: Some(idempotency_key.to_string()), + request_id: None, + source: WORKFLOW_BRANCH_SOURCE.to_string(), + prompt_snapshot_ref: prompt_snapshot_ref.clone(), + input_snapshot_ref: None, + thread_id: None, + thread_store_kind: WORKFLOW_BRANCH_THREAD_STORE_KIND.to_string(), + thread_store_id: None, + rollout_path: None, + parent_thread_id: run + .source_thread_id .as_ref() .map(std::string::ToString::to_string), + parent_agent_run_id: params.parent_agent_run_id.clone(), + spawn_linkage_json: Some(spawn_linkage_json), + auth_profile_ref: params.auth_profile_ref.clone(), + status_reason: Some("queued by workflow branch admission".to_string()), + config_fingerprint: params.config_fingerprint.clone(), + version_fingerprint: params.version_fingerprint.clone(), + }; + let start_event_payload = json!({ + "cwd": Value::Null, + "prompt": prompt, + "promptSnapshotRef": prompt_snapshot_ref, + }); + let execution_snapshot_params = BackgroundAgentExecutionSnapshotParams { + run_id: background_agent_run_id.to_string(), + snapshot_kind: "initial_execution_context".to_string(), + payload_json: branch_execution_payload( + run, + candidate, + model_route_json, + workspace_json, + params, + ), + recovery_policy: "abort_mid_turn_resume_at_safe_boundary".to_string(), + config_fingerprint: params.config_fingerprint.clone(), + }; + let admission_identity_sha256 = background_agent_admission_identity_sha256( + &run_params, + &start_event_payload, + &execution_snapshot_params, + )?; + let created = validate_existing_background_agent_admission_in_tx( + tx, + idempotency_key, + &run_params, + ExistingBackgroundAgentAdmissionIdentity::AdmissionDigest( + admission_identity_sha256.as_str(), + ), ) - .bind(params.parent_agent_run_id.as_deref()) - .bind(spawn_linkage_json.as_str()) - .bind(params.auth_profile_ref.as_deref().map(redact_state_string)) - .bind(BackgroundAgentDesiredState::Running.as_str()) - .bind(BackgroundAgentRunStatus::Queued.as_str()) - .bind(redact_state_string("queued by workflow branch admission")) - .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?; - - let event = append_background_agent_event_in_tx( + .await? + .is_none(); + if created { + insert_background_agent_run_in_tx(tx, &run_params, now).await?; + } + let (event, _execution_snapshot_id) = recover_or_validate_background_agent_initial_state_in_tx( tx, background_agent_run_id, - "agent.started", - &json!({ - "cwd": Value::Null, - "prompt": prompt, - "promptSnapshotRef": prompt_snapshot_ref, - }), - now, + &run_params, + admission_identity_sha256.as_str(), + &start_event_payload, + &execution_snapshot_params, + params.max_active_background_agent_runs.unwrap_or(i64::MAX), ) .await?; let status_payload = json!({ @@ -1235,17 +1216,7 @@ INSERT INTO background_agent_runs ( }, ) .await?; - insert_background_agent_execution_snapshot_in_tx( - tx, - background_agent_run_id, - "initial_execution_context", - branch_execution_payload(run, candidate, model_route_json, workspace_json, params), - "abort_mid_turn_resume_at_safe_boundary", - params.config_fingerprint.as_deref(), - now, - ) - .await?; - Ok(true) + Ok(created) } async fn upsert_background_agent_status_snapshot_in_tx( @@ -1298,58 +1269,6 @@ ON CONFLICT(run_id) DO UPDATE SET Ok(()) } -async fn insert_background_agent_execution_snapshot_in_tx( - tx: &mut sqlx::Transaction<'_, Sqlite>, - run_id: &str, - snapshot_kind: &str, - payload_json: Value, - recovery_policy: &str, - config_fingerprint: Option<&str>, - now: i64, -) -> anyhow::Result<()> { - let seq: i64 = sqlx::query_scalar( - "SELECT COALESCE(MAX(seq), 0) + 1 FROM background_agent_execution_snapshots WHERE run_id = ?", - ) - .bind(run_id) - .fetch_one(&mut **tx) - .await?; - sqlx::query( - r#" -INSERT INTO background_agent_execution_snapshots ( - run_id, - seq, - snapshot_kind, - payload_json, - recovery_policy, - config_fingerprint, - created_at -) VALUES (?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(run_id) - .bind(seq) - .bind(snapshot_kind) - .bind(redact_state_json_string(&payload_json)?) - .bind(recovery_policy) - .bind(config_fingerprint) - .bind(now) - .execute(&mut **tx) - .await?; - sqlx::query( - r#" -UPDATE background_agent_runs -SET last_snapshot_seq = ?, updated_at = ? -WHERE id = ? - "#, - ) - .bind(seq) - .bind(now) - .bind(run_id) - .execute(&mut **tx) - .await?; - Ok(()) -} - fn branch_execution_payload( run: &crate::WorkflowRun, candidate: &ReadyBranchCandidate, @@ -1372,7 +1291,10 @@ fn branch_execution_payload( "serviceTier": model_route_json.get("service_tier"), "approvalPolicy": model_route_json.get("approval_policy"), "permissionProfile": model_route_json.get("permission_profile"), - "authProfileRef": params.auth_profile_ref.as_deref(), + "authProfileIdentitySha256": params + .auth_profile_ref + .as_deref() + .map(|profile| StateRuntime::background_agent_identity_sha256(profile.as_bytes())), "workspace": workspace_json, "envSnapshotPolicy": "inherit-minimal", "maxRuntimeSeconds": workflow_state_data(&run.limits_json).get("max_step_runtime_seconds"), @@ -3320,6 +3242,21 @@ WHERE plan_id = ? AND key = ? .expect("admission should succeed") .expect("run should be owned"); assert_eq!(1, admitted.admitted.len()); + let admitted_run_id = admitted.admitted[0].background_agent_run_id.clone(); + assert_eq!( + runtime + .list_background_agent_events_after( + admitted_run_id.as_str(), + /*after_seq*/ None, + /*limit*/ None, + ) + .await + .expect("admission events should load") + .into_iter() + .map(|event| event.event_type) + .collect::>(), + vec!["agent.admitted".to_string(), "agent.started".to_string()] + ); let retry = admit_test_workflow_run_branches( &runtime, diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index bd170e5841..72eb6fc9aa 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -265,6 +265,9 @@ 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 +1654,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),