diff --git a/codex-rs/app-server/src/request_processors/thread_workflow_processor.rs b/codex-rs/app-server/src/request_processors/thread_workflow_processor.rs index bbc084144..7dcbc53c7 100644 --- a/codex-rs/app-server/src/request_processors/thread_workflow_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_workflow_processor.rs @@ -378,15 +378,23 @@ impl ThreadWorkflowRequestProcessor { if !workflow_run_belongs_to_thread(&state_db, thread_id, params.run_id.as_str()).await? { return Ok(ThreadWorkflowRunCancelResponse { run: None }); } + let run_id = params.run_id; let cancel_params = codex_state::WorkflowRunCancelParams { - run_id: params.run_id, + run_id: run_id.clone(), reason: params .reason .and_then(normalize_optional_string) .unwrap_or_else(|| "user requested cancellation".to_string()), }; - let run = state_db - .request_workflow_run_cancel(cancel_params) + let activation_service = self + .activation_service + .as_ref() + .ok_or_else(|| internal_error("workflow activation service is not initialized"))?; + let run = activation_service + .cancel_workflow_run( + cancel_params, + crate::extensions::workflow_activation_config(&self.config), + ) .await .map_err(|err| internal_error(format!("failed to cancel thread workflow run: {err}")))? .map(api_thread_workflow_run_snapshot_from_state); diff --git a/codex-rs/app-server/tests/suite/v2/workflow.rs b/codex-rs/app-server/tests/suite/v2/workflow.rs index 3ec0495dc..25b99d40c 100644 --- a/codex-rs/app-server/tests/suite/v2/workflow.rs +++ b/codex-rs/app-server/tests/suite/v2/workflow.rs @@ -472,9 +472,38 @@ async fn workflow_run_lifecycle_projects_tasks_and_returns_sanitized_state() -> assert_does_not_leak(&serde_json::to_string(&cancel_resp.result)?)?; let cancelled = to_response::(cancel_resp)?; assert_eq!( - Some(ThreadWorkflowRunStatus::CancelRequested), + Some(ThreadWorkflowRunStatus::Cancelled), cancelled.run.as_ref().map(|run| run.run.status) ); + let runtime = open_state_runtime(codex_home.path()).await?; + let terminal = timeout(DEFAULT_TIMEOUT, async { + loop { + let snapshot = runtime + .workflows() + .get_workflow_run_snapshot(started.run.run.run_id.as_str()) + .await? + .ok_or_else(|| anyhow::anyhow!("cancelled workflow run disappeared"))?; + if snapshot.run.status.is_terminal() { + return Result::<_, anyhow::Error>::Ok(snapshot); + } + sleep(std::time::Duration::from_millis(50)).await; + } + }) + .await??; + assert_eq!( + codex_state::WorkflowRunStatus::Cancelled, + terminal.run.status + ); + let projected_plan = runtime + .thread_goals() + .list_thread_goal_plans(parse_thread_id(thread_id.as_str())?) + .await? + .pop() + .ok_or_else(|| anyhow::anyhow!("cancelled projected goal plan is missing"))?; + assert_eq!( + codex_state::ThreadGoalPlanStatus::Cancelled, + projected_plan.plan.status + ); assert!(!marker.exists()); Ok(()) diff --git a/codex-rs/ext/workflows/src/activation.rs b/codex-rs/ext/workflows/src/activation.rs index 9e9b25023..9158d03a4 100644 --- a/codex-rs/ext/workflows/src/activation.rs +++ b/codex-rs/ext/workflows/src/activation.rs @@ -30,6 +30,7 @@ use codex_state::WorkflowProviderCreditReservationStatus; use codex_state::WorkflowRunAdvanceParams; use codex_state::WorkflowRunBranchAdmissionParams; use codex_state::WorkflowRunBranchReconcileParams; +use codex_state::WorkflowRunCancelParams; use codex_state::WorkflowRunClaimParams; use codex_state::WorkflowRunCreateParams; use codex_state::WorkflowRunFenceParams; @@ -51,6 +52,7 @@ use codex_workflows::WorkflowRouteReceipt; use codex_workflows::WorkflowRouteRuntime; use codex_workflows::WorkflowSpec; use codex_workflows::WorkflowVerifier; +#[cfg(test)] use codex_workflows::WorkflowWorkspaceMode; use codex_workflows::admit_workflow_model_route_for_runtime; use codex_workflows::parse_workflow_yaml; @@ -66,7 +68,6 @@ use tokio_util::sync::CancellationToken; use crate::provider_credit::WorkflowProviderCreditAuthority; use crate::provider_credit::unsupported_workflow_provider_credit_authority; -use crate::provider_credit::workflow_finite_budget; const WORKFLOW_SUPERVISOR_POLL_INTERVAL: Duration = Duration::from_millis(250); const WORKFLOW_HEARTBEAT_INTERVAL: Duration = Duration::from_millis(500); @@ -263,6 +264,37 @@ pub struct WorkflowStartOutcome { pub goal_plan: Option, } +#[derive(Debug)] +pub(crate) struct WorkflowStartError { + pub code: &'static str, + pub stage: &'static str, + pub message: &'static str, + pub recovered: bool, + pub snapshot: Option, + pub goal_plan: Option, + source: anyhow::Error, +} + +struct WorkflowStartFailureContext<'a> { + thread_id: ThreadId, + idempotency_key: Option<&'a str>, + code: &'static str, + stage: &'static str, + message: &'static str, +} + +impl std::fmt::Display for WorkflowStartError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "{}: {}", self.code, self.message) + } +} + +impl std::error::Error for WorkflowStartError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.source.source() + } +} + #[derive(Clone)] pub struct WorkflowActivationService { state_db: Arc, @@ -302,10 +334,6 @@ impl WorkflowActivationService { .await? .ok_or_else(|| anyhow::anyhow!("workflow spec record not found"))?; let spec = parse_workflow_yaml(spec_record.source_yaml.as_str())?; - validate_workflow_routes_before_reservation( - &spec, - &request.activation_config.route_runtime, - )?; let create_params = WorkflowRunCreateParams { workflow_record_id: request.workflow_record_id, source_thread_id: Some(request.source_thread_id), @@ -324,27 +352,96 @@ impl WorkflowActivationService { .create_workflow_run(create_params) .await? }; + let projection_params = WorkflowGoalPlanProjectionParams { + workflow_run_id: snapshot.run.run_id.clone(), + thread_id: request.source_thread_id, + idempotency_key: request.idempotency_key.clone(), + }; + let goal_plan = match retry_workflow_state(WorkflowStateOperation::ProjectGoalPlan, || { + self.state_db + .project_workflow_run_to_goal_plan(projection_params.clone()) + }) + .await + { + Ok(goal_plan) => goal_plan, + Err(error) => { + return Err(self + .recover_start_failure( + &snapshot, + WorkflowStartFailureContext { + thread_id: request.source_thread_id, + idempotency_key: request.idempotency_key.as_deref(), + code: "workflow_start_projection_failed", + stage: "goal_plan_projection", + message: "workflow run creation succeeded but goal-plan projection failed", + }, + error, + ) + .await); + } + }; + if snapshot.run.status.is_terminal() { + return Ok(WorkflowStartOutcome { + snapshot, + goal_plan, + }); + } let mut activation_config = request.activation_config; let claim_params = WorkflowRunClaimParams { run_id: snapshot.run.run_id.clone(), owner_id: self.owner_instance_id.to_string(), lease_duration_ms: None, }; - let initial_claim = retry_workflow_state(WorkflowStateOperation::ClaimRun, || { + let initial_claim = match retry_workflow_state(WorkflowStateOperation::ClaimRun, || { self.state_db.claim_workflow_run(claim_params.clone()) }) - .await?; + .await + { + Ok(initial_claim) => initial_claim, + Err(error) => { + return Err(self + .recover_start_failure( + &snapshot, + WorkflowStartFailureContext { + thread_id: request.source_thread_id, + idempotency_key: request.idempotency_key.as_deref(), + code: "workflow_start_claim_failed", + stage: "run_claim", + message: "workflow run was created and projected but could not be claimed", + }, + error, + ) + .await); + } + }; if let Some(claim) = initial_claim.as_ref() { let fence = WorkflowRunFenceParams { run_id: snapshot.run.run_id.clone(), owner_id: self.owner_instance_id.to_string(), generation: claim.generation, }; - activation_config.route_runtime.credit_control = self + activation_config.route_runtime.credit_control = match self .provider_credit_authority .reserve_or_restore(&fence, &spec, &activation_config.route_runtime) - .await?; - validate_workflow_routes_before_effects(&spec, &activation_config.route_runtime)?; + .await + { + Ok(credit_control) => credit_control, + Err(error) => { + return Err(self + .recover_start_failure( + &snapshot, + WorkflowStartFailureContext { + thread_id: request.source_thread_id, + idempotency_key: request.idempotency_key.as_deref(), + code: "workflow_start_provider_credit_failed", + stage: "provider_credit", + message: "provider credit reservation failed after durable workflow creation", + }, + error, + ) + .await); + } + }; } let snapshot = retry_workflow_state(WorkflowStateOperation::LoadSnapshot, || { self.state_db @@ -353,16 +450,6 @@ impl WorkflowActivationService { }) .await? .ok_or_else(|| anyhow::anyhow!("workflow run disappeared after credit reservation"))?; - let projection_params = WorkflowGoalPlanProjectionParams { - workflow_run_id: snapshot.run.run_id.clone(), - thread_id: request.source_thread_id, - idempotency_key: request.idempotency_key, - }; - let goal_plan = retry_workflow_state(WorkflowStateOperation::ProjectGoalPlan, || { - self.state_db - .project_workflow_run_to_goal_plan(projection_params.clone()) - }) - .await?; self.activate_with_claim( snapshot.run.run_id.clone(), activation_config, @@ -375,6 +462,103 @@ impl WorkflowActivationService { }) } + async fn recover_start_failure( + &self, + snapshot: &WorkflowRunSnapshot, + context: WorkflowStartFailureContext<'_>, + source: anyhow::Error, + ) -> anyhow::Error { + let run_id = snapshot.run.run_id.as_str(); + let recovered_snapshot = match self + .state_db + .request_workflow_run_cancel(WorkflowRunCancelParams { + run_id: run_id.to_string(), + reason: "workflow start failed".to_string(), + }) + .await + { + Ok(Some(cancelled)) if cancelled.run.status == WorkflowRunStatus::CancelRequested => { + let generation = match self + .state_db + .claim_workflow_run(WorkflowRunClaimParams { + run_id: run_id.to_string(), + owner_id: self.owner_instance_id.to_string(), + lease_duration_ms: None, + }) + .await + { + Ok(claim) => claim.map(|claim| claim.generation), + Err(error) => { + tracing::warn!( + workflow_run_id = %run_id, + "workflow start recovery could not claim the cancelled run: {error}" + ); + None + } + }; + if let Some(generation) = generation + && let Err(error) = self + .state_db + .advance_workflow_run(WorkflowRunAdvanceParams { + run_id: run_id.to_string(), + owner_id: self.owner_instance_id.to_string(), + generation, + }) + .await + { + tracing::warn!( + workflow_run_id = %run_id, + "workflow start recovery could not finalize cancellation: {error}" + ); + } + self.state_db + .workflows() + .get_workflow_run_snapshot(run_id) + .await + .ok() + .flatten() + } + Ok(snapshot) => snapshot, + Err(error) => { + tracing::warn!( + workflow_run_id = %run_id, + "workflow start recovery could not request cancellation: {error}" + ); + self.state_db + .workflows() + .get_workflow_run_snapshot(run_id) + .await + .ok() + .flatten() + } + }; + let goal_plan = self + .state_db + .project_workflow_run_to_goal_plan(WorkflowGoalPlanProjectionParams { + workflow_run_id: run_id.to_string(), + thread_id: context.thread_id, + idempotency_key: context.idempotency_key.map(str::to_string), + }) + .await + .ok() + .flatten(); + let recovered = recovered_snapshot + .as_ref() + .is_some_and(|snapshot| snapshot.run.status == WorkflowRunStatus::Cancelled) + && goal_plan.as_ref().is_none_or(|projection| { + projection.snapshot.plan.status == codex_state::ThreadGoalPlanStatus::Cancelled + }); + anyhow::Error::new(WorkflowStartError { + code: context.code, + stage: context.stage, + message: context.message, + recovered, + snapshot: recovered_snapshot, + goal_plan, + source, + }) + } + pub async fn activate_thread_runs( &self, thread_id: ThreadId, @@ -415,6 +599,42 @@ impl WorkflowActivationService { .await; } + pub async fn cancel_workflow_run( + &self, + params: WorkflowRunCancelParams, + config: WorkflowActivationConfig, + ) -> anyhow::Result> { + let run_id = params.run_id.clone(); + let snapshot = self.state_db.request_workflow_run_cancel(params).await?; + if snapshot + .as_ref() + .is_some_and(|snapshot| snapshot.run.status == WorkflowRunStatus::CancelRequested) + { + self.activate(run_id.clone(), config).await; + for _ in 0..20 { + let current = retry_workflow_state(WorkflowStateOperation::LoadSnapshot, || { + self.state_db + .workflows() + .get_workflow_run_snapshot(run_id.as_str()) + }) + .await?; + if current + .as_ref() + .is_none_or(|snapshot| snapshot.run.status.is_terminal()) + { + return Ok(current); + } + tokio::time::sleep(WORKFLOW_SUPERVISOR_POLL_INTERVAL).await; + } + } + retry_workflow_state(WorkflowStateOperation::LoadSnapshot, || { + self.state_db + .workflows() + .get_workflow_run_snapshot(run_id.as_str()) + }) + .await + } + async fn activate_with_claim( &self, run_id: String, @@ -531,7 +751,6 @@ impl WorkflowActivationService { .provider_credit_authority .reserve_or_restore(&fence, &spec, &config.route_runtime) .await?; - validate_workflow_routes_before_effects(&spec, &config.route_runtime)?; self.drive_generation(run_id, claim.generation, &spec, config.clone()) .await?; @@ -635,7 +854,6 @@ impl WorkflowActivationService { .provider_credit_authority .reserve_or_restore(&fence, spec, &config.route_runtime) .await?; - validate_workflow_routes_before_effects(spec, &config.route_runtime)?; let fingerprint = activation_config_fingerprint(&config)?; let admission_params = WorkflowRunBranchAdmissionParams { run_id: run_id.to_string(), @@ -1599,7 +1817,9 @@ fn validate_verifier_route( ), None => "shared_repository", }; - let current = admit_workflow_model_route_for_runtime(&requested, runtime, worktree_mode)?; + let route_runtime = workflow_route_runtime_for_requested(&requested, runtime); + let current = + admit_workflow_model_route_for_runtime(&requested, &route_runtime, worktree_mode)?; if current != admitted { anyhow::bail!( "workflow_route_receipt_mismatch: verifier route differs from branch admission" @@ -1609,6 +1829,41 @@ fn validate_verifier_route( Ok(admitted) } +fn workflow_route_runtime_for_requested( + requested: &WorkflowModelRoute, + runtime: &WorkflowRouteRuntime, +) -> WorkflowRouteRuntime { + let explicit_auth_profile = requested.routing.as_ref().and_then(|routing| { + routing + .decision + .as_ref() + .and_then(|decision| decision.auth_profile.clone()) + .or_else(|| routing.request.context.auth_profile.clone()) + }); + WorkflowRouteRuntime { + model_gateway: Some(requested.model_gateway.clone()), + provider: Some(requested.provider.clone()), + model: Some(requested.model.clone()), + reasoning: Some(requested.reasoning.clone()), + service_tier: requested.service_tier.clone(), + auth_profile: explicit_auth_profile.or_else(|| { + (runtime.provider.as_deref() == Some(requested.provider.as_str())) + .then(|| runtime.auth_profile.clone()) + .flatten() + }), + approval_policy: requested + .approval_policy + .clone() + .or_else(|| runtime.approval_policy.clone()), + permission_profile: requested + .permission_profile + .clone() + .or_else(|| runtime.permission_profile.clone()), + context_window_tokens: runtime.context_window_tokens, + credit_control: runtime.credit_control.clone(), + } +} + #[derive(Serialize)] #[serde(rename_all = "camelCase")] struct WorkflowActivationFingerprint<'a> { @@ -1626,6 +1881,7 @@ fn activation_config_fingerprint(config: &WorkflowActivationConfig) -> anyhow::R Ok(format!("{:x}", Sha256::digest(serde_json::to_vec(&value)?))) } +#[cfg(test)] fn validate_workflow_routes_before_effects( spec: &codex_workflows::WorkflowSpec, runtime: &WorkflowRouteRuntime, @@ -1648,24 +1904,6 @@ fn validate_workflow_routes_before_effects( Ok(()) } -fn validate_workflow_routes_before_reservation( - spec: &codex_workflows::WorkflowSpec, - runtime: &WorkflowRouteRuntime, -) -> anyhow::Result<()> { - let mut preflight_runtime = runtime.clone(); - preflight_runtime.credit_control = match workflow_finite_budget(spec)? { - Some(ceiling_usd) => WorkflowProviderCreditControl::Reserved { - reservation_id: "preflight-provider-credit-reservation".to_string(), - ceiling_usd: ceiling_usd.clone(), - spent_usd: "0".to_string(), - remaining_usd: ceiling_usd, - exhausted: false, - }, - None => WorkflowProviderCreditControl::NotRequested, - }; - validate_workflow_routes_before_effects(spec, &preflight_runtime) -} - fn workflow_credit_terminal_status( run_status: &WorkflowRunStatus, ) -> WorkflowProviderCreditReservationStatus { diff --git a/codex-rs/ext/workflows/src/manager_tool.rs b/codex-rs/ext/workflows/src/manager_tool.rs index fee3a9612..42b622c17 100644 --- a/codex-rs/ext/workflows/src/manager_tool.rs +++ b/codex-rs/ext/workflows/src/manager_tool.rs @@ -22,6 +22,7 @@ use serde_json::json; use crate::WorkflowActivationConfig; use crate::WorkflowActivationService; use crate::WorkflowStartRequest; +use crate::activation::WorkflowStartError; use crate::manager_output::goal_plan_projection_json; use crate::manager_output::run_snapshot_json; use crate::manager_output::run_summary_json; @@ -386,7 +387,7 @@ impl ManageWorkflowTool { "error": "workflow not found for current thread", })); } - let outcome = activation_service + let outcome = match activation_service .start_workflow_run(WorkflowStartRequest { workflow_record_id, source_thread_id: thread_id, @@ -394,7 +395,35 @@ impl ManageWorkflowTool { activation_config: activation_config.clone(), }) .await - .map_err(|_| respond("failed to start workflow run"))?; + { + Ok(outcome) => outcome, + Err(error) => { + if let Some(error) = error.downcast_ref::() { + return Ok(json!({ + "action": "start", + "run": error.snapshot.as_ref().map(run_snapshot_json), + "goalPlan": error.goal_plan.clone().map(goal_plan_projection_json), + "error": { + "code": error.code, + "stage": error.stage, + "message": error.message, + "recovered": error.recovered, + }, + })); + } + return Ok(json!({ + "action": "start", + "run": null, + "goalPlan": null, + "error": { + "code": "workflow_start_failed", + "stage": "activation", + "message": "workflow activation failed before durable recovery was available", + "recovered": false, + }, + })); + } + }; let run = run_snapshot_json(&outcome.snapshot); let goal_plan = outcome.goal_plan.map(goal_plan_projection_json); Ok(json!({ @@ -475,17 +504,20 @@ impl ManageWorkflowTool { } async fn cancel_run(&self, args: ManageWorkflowArgs) -> Result { - let (state_db, _) = self.runtime()?; + let (_, _, activation_service, activation_config) = self.activation_runtime()?; let run_id = required_field(args.run_id, "run_id", "cancel")?; if self.thread_run_snapshot(run_id.as_str()).await?.is_none() { return Ok(not_found_run_response("cancel")); } - let snapshot = state_db - .request_workflow_run_cancel(codex_state::WorkflowRunCancelParams { - run_id, - reason: normalize_optional_string(args.reason) - .unwrap_or_else(|| "model requested cancellation".to_string()), - }) + let snapshot = activation_service + .cancel_workflow_run( + codex_state::WorkflowRunCancelParams { + run_id, + reason: normalize_optional_string(args.reason) + .unwrap_or_else(|| "model requested cancellation".to_string()), + }, + activation_config.clone(), + ) .await .map_err(|_| respond("failed to cancel workflow run"))?; Ok(json!({ @@ -635,6 +667,7 @@ mod tests { use std::sync::Arc; use std::sync::atomic::AtomicBool; + use async_trait::async_trait; use codex_extension_api::ConversationHistory; use codex_extension_api::FunctionCallError; use codex_extension_api::NoopTurnItemEmitter; @@ -651,6 +684,9 @@ mod tests { use super::MANAGE_WORKFLOW_TOOL_NAME; use super::ManageWorkflowTool; + use crate::WorkflowActivationConfig; + use crate::WorkflowActivationService; + use crate::provider_credit::WorkflowProviderCreditAuthority; const MANAGE_WORKFLOW_TEST_YAML: &str = r#" schema_version: "workflow.codex.codewith/v0" @@ -758,6 +794,87 @@ cleanup: on_complete: [] "#; + fn heterogeneous_manage_workflow_yaml() -> String { + MANAGE_WORKFLOW_TEST_YAML + .replacen( + r#" - id: "adversarial_reviewer" + display_name: "Reviewer-Hypatia" + role: "Independently review the workflow lifecycle candidate." + model: + model_gateway: "hasna" + provider: "openai" + model: "gpt-5.4" + reasoning: "high""#, + r#" - id: "adversarial_reviewer" + display_name: "Reviewer-Hypatia" + role: "Independently review the workflow lifecycle candidate." + model: + model_gateway: "hasna" + provider: "anthropic" + model: "claude-fable-5" + reasoning: "high""#, + 1, + ) + .replacen( + r#" - id: "initial_adversarial_review" + title: "Run the initial adversarial review" + agent: "adversarial_reviewer" + model: + model_gateway: "hasna" + provider: "openai" + model: "gpt-5.4" + reasoning: "high""#, + r#" - id: "initial_adversarial_review" + title: "Run the initial adversarial review" + agent: "adversarial_reviewer" + model: + model_gateway: "hasna" + provider: "anthropic" + model: "claude-fable-5" + reasoning: "high""#, + 1, + ) + } + + fn workflow_activation_config() -> WorkflowActivationConfig { + WorkflowActivationConfig { + route_runtime: codex_workflows::WorkflowRouteRuntime { + model_gateway: Some("hasna".to_string()), + provider: Some("openai".to_string()), + model: Some("gpt-5.4".to_string()), + reasoning: Some("high".to_string()), + approval_policy: Some("never".to_string()), + permission_profile: Some("read-only".to_string()), + context_window_tokens: Some(256_000), + ..Default::default() + }, + ..Default::default() + } + } + + #[derive(Debug)] + struct FailingProviderCreditAuthority; + + #[async_trait] + impl WorkflowProviderCreditAuthority for FailingProviderCreditAuthority { + async fn reserve_or_restore( + &self, + _fence: &codex_state::WorkflowRunFenceParams, + _spec: &codex_workflows::WorkflowSpec, + _runtime: &codex_workflows::WorkflowRouteRuntime, + ) -> anyhow::Result { + anyhow::bail!("provider credit reservation unavailable") + } + + async fn reconcile_terminal( + &self, + _fence: &codex_state::WorkflowRunFenceParams, + _status: codex_state::WorkflowProviderCreditReservationStatus, + ) -> anyhow::Result<()> { + Ok(()) + } + } + #[tokio::test] async fn manage_workflow_lifecycle_returns_sanitized_state() { let tempdir = tempfile::tempdir().expect("tempdir"); @@ -908,12 +1025,194 @@ cleanup: }), ) .await; - assert_eq!(cancel["run"]["run"]["status"], "cancel_requested"); + assert_eq!(cancel["run"]["run"]["status"], "cancelled"); assert_eq!( cancel["run"]["run"]["statusReason"], - "user requested workflow cancellation" + "workflow run cancelled" ); assert!(!cancel.to_string().contains("should-not-leak")); + let cancelled_goal_plan = state_db + .thread_goals() + .list_thread_goal_plans(thread_id) + .await + .expect("cancelled goal plan should list") + .pop() + .expect("cancelled goal plan should exist"); + assert_eq!( + codex_state::ThreadGoalPlanStatus::Cancelled, + cancelled_goal_plan.plan.status + ); + assert!( + cancelled_goal_plan + .nodes + .iter() + .all(|node| node.status == codex_state::ThreadGoalPlanNodeStatus::Cancelled) + ); + let repeated_cancel = call_tool( + &tool, + json!({ + "action": "cancel", + "run_id": cancel["run"]["run"]["runId"].as_str().expect("run id"), + }), + ) + .await; + assert_eq!( + cancel["run"]["run"]["runId"], + repeated_cancel["run"]["run"]["runId"] + ); + assert_eq!(repeated_cancel["run"]["run"]["status"], "cancelled"); + } + + #[tokio::test] + async fn manage_workflow_heterogeneous_start_is_idempotent() { + let tempdir = tempfile::tempdir().expect("tempdir"); + let state_db = codex_state::StateRuntime::init( + tempdir.path().to_path_buf(), + "test-provider".to_string(), + ) + .await + .expect("state runtime should initialize"); + let thread_id = codex_protocol::ThreadId::new(); + let mut thread = codex_state::ThreadMetadataBuilder::new( + thread_id, + state_db.codex_home().join("rollout.jsonl"), + chrono::Utc::now(), + codex_protocol::protocol::SessionSource::Cli, + ); + thread.cwd = tempdir.path().to_path_buf(); + state_db + .upsert_thread(&thread.build("test-provider")) + .await + .expect("thread metadata should insert"); + let tool = + ManageWorkflowTool::new(Arc::new(AtomicBool::new(true)), state_db.clone(), thread_id); + + let create = call_tool( + &tool, + json!({ + "action": "create", + "yaml": heterogeneous_manage_workflow_yaml(), + }), + ) + .await; + let workflow_record_id = create["workflow"]["workflowRecordId"] + .as_str() + .expect("workflow id") + .to_string(); + + let first = call_tool( + &tool, + json!({ + "action": "start", + "workflow_record_id": workflow_record_id, + "idempotency_key": "heterogeneous-run", + }), + ) + .await; + let replay = call_tool( + &tool, + json!({ + "action": "start", + "workflow_record_id": create["workflow"]["workflowRecordId"], + "idempotency_key": "heterogeneous-run", + }), + ) + .await; + + assert_eq!(first["run"]["run"]["runId"], replay["run"]["run"]["runId"]); + assert_eq!(first["goalPlan"]["planId"], replay["goalPlan"]["planId"]); + let runs = state_db + .workflows() + .list_thread_workflow_runs_page(thread_id, /*cursor*/ None, /*limit*/ 10) + .await + .expect("workflow runs should list"); + assert_eq!(1, runs.data.len()); + let plans = state_db + .thread_goals() + .list_thread_goal_plans(thread_id) + .await + .expect("workflow goal plans should list"); + assert_eq!(1, plans.len()); + } + + #[tokio::test] + async fn manage_workflow_start_failure_returns_structured_recovered_state() { + let tempdir = tempfile::tempdir().expect("tempdir"); + let state_db = codex_state::StateRuntime::init( + tempdir.path().to_path_buf(), + "test-provider".to_string(), + ) + .await + .expect("state runtime should initialize"); + let thread_id = codex_protocol::ThreadId::new(); + let mut thread = codex_state::ThreadMetadataBuilder::new( + thread_id, + state_db.codex_home().join("rollout.jsonl"), + chrono::Utc::now(), + codex_protocol::protocol::SessionSource::Cli, + ); + thread.cwd = tempdir.path().to_path_buf(); + state_db + .upsert_thread(&thread.build("test-provider")) + .await + .expect("thread metadata should insert"); + let activation_service = Arc::new( + WorkflowActivationService::new_with_provider_credit_authority( + state_db.clone(), + Arc::new(FailingProviderCreditAuthority), + ), + ); + let tool = ManageWorkflowTool::new_with_activation( + Arc::new(AtomicBool::new(true)), + state_db.clone(), + thread_id, + activation_service, + workflow_activation_config(), + ); + let create = call_tool( + &tool, + json!({ + "action": "create", + "yaml": MANAGE_WORKFLOW_TEST_YAML.replace(" max_tokens: 1000", ""), + }), + ) + .await; + + let start = call_tool( + &tool, + json!({ + "action": "start", + "workflow_record_id": create["workflow"]["workflowRecordId"], + "idempotency_key": "recoverable-start-failure", + }), + ) + .await; + + assert_eq!(start["action"], "start"); + assert_eq!( + start["error"]["code"], + "workflow_start_provider_credit_failed" + ); + assert_eq!(start["error"]["stage"], "provider_credit"); + assert_eq!(start["error"]["recovered"], true); + assert_eq!(start["run"]["run"]["status"], "cancelled"); + assert_eq!(start["goalPlan"]["status"], "cancelled"); + let runs = state_db + .workflows() + .list_thread_workflow_runs_page(thread_id, /*cursor*/ None, /*limit*/ 10) + .await + .expect("workflow runs should list"); + assert_eq!(1, runs.data.len()); + let plans = state_db + .thread_goals() + .list_thread_goal_plans(thread_id) + .await + .expect("workflow goal plans should list"); + assert_eq!(1, plans.len()); + assert_eq!( + codex_state::ThreadGoalPlanStatus::Cancelled, + plans[0].plan.status + ); } #[tokio::test] diff --git a/codex-rs/state/src/runtime/workflow_goal_plan_projections.rs b/codex-rs/state/src/runtime/workflow_goal_plan_projections.rs index f2d0da4e7..f5d1b7bbc 100644 --- a/codex-rs/state/src/runtime/workflow_goal_plan_projections.rs +++ b/codex-rs/state/src/runtime/workflow_goal_plan_projections.rs @@ -62,13 +62,6 @@ impl StateRuntime { else { return Ok(None); }; - if !is_workflow_run_projectable(&run) { - anyhow::bail!( - "workflow run {} cannot be projected after reaching terminal status {}", - run.run.run_id, - run.run.status.as_str() - ); - } if run.run.source_thread_id != Some(params.thread_id) { anyhow::bail!( "workflow run {} does not belong to thread {}", @@ -103,6 +96,13 @@ impl GoalStore { tx.commit().await?; return workflow_goal_plan_projection_from_row(row, /*created*/ false, snapshot); } + if !is_workflow_run_projectable(params.run) { + anyhow::bail!( + "workflow run {} cannot be projected after reaching terminal status {}", + params.run.run.run_id, + params.run.run.status.as_str() + ); + } let plan_params = thread_goal_plan_from_workflow_run(params.run, params.thread_id)?; let plan = insert_thread_goal_plan_in_tx(&mut tx, plan_params, params.now_ms).await?; @@ -171,6 +171,66 @@ WHERE plan_id = ? Ok(vec![snapshot]) } + pub(crate) async fn cancel_workflow_goal_plan_projection( + &self, + run_id: &str, + ) -> anyhow::Result> { + let Some(plan_id) = + workflow_goal_plan_projection_plan_id_in_pool(&self.pool, run_id).await? + else { + return Ok(Vec::new()); + }; + let now_ms = datetime_to_epoch_millis(Utc::now()); + let mut tx = self.pool.begin().await?; + sqlx::query( + r#" +UPDATE thread_goals +SET status = ?, updated_at_ms = ? +WHERE status NOT IN ('complete', 'cancelled') + AND EXISTS ( + SELECT 1 + FROM thread_goal_plan_nodes node + WHERE node.projected_goal_id = thread_goals.goal_id + AND node.plan_id = ? + ) + "#, + ) + .bind(crate::ThreadGoalStatus::Cancelled.as_str()) + .bind(now_ms) + .bind(plan_id.as_str()) + .execute(&mut *tx) + .await?; + sqlx::query( + r#" +UPDATE thread_goal_plan_nodes +SET status = ?, updated_at_ms = ? +WHERE plan_id = ? + AND status NOT IN ('complete', 'cancelled') + "#, + ) + .bind(crate::ThreadGoalPlanNodeStatus::Cancelled.as_str()) + .bind(now_ms) + .bind(plan_id.as_str()) + .execute(&mut *tx) + .await?; + sqlx::query( + r#" +UPDATE thread_goal_plans +SET status = ?, updated_at_ms = ? +WHERE plan_id = ? + AND status NOT IN ('complete', 'cancelled') + "#, + ) + .bind(crate::ThreadGoalPlanStatus::Cancelled.as_str()) + .bind(now_ms) + .bind(plan_id.as_str()) + .execute(&mut *tx) + .await?; + let snapshot = snapshot_thread_goal_plan_in_tx(&mut tx, plan_id.as_str()).await?; + tx.commit().await?; + Ok(vec![snapshot]) + } + pub(crate) async fn pause_workflow_goal_plan_projection( &self, run_id: &str, diff --git a/codex-rs/state/src/runtime/workflow_orchestrator.rs b/codex-rs/state/src/runtime/workflow_orchestrator.rs index fe75a21f6..f48b34224 100644 --- a/codex-rs/state/src/runtime/workflow_orchestrator.rs +++ b/codex-rs/state/src/runtime/workflow_orchestrator.rs @@ -408,7 +408,7 @@ WHERE run_id = ? tx.commit().await?; if snapshot.run.status == crate::WorkflowRunStatus::Cancelled { self.thread_goals - .block_workflow_goal_plan_projection(params.run_id.as_str()) + .cancel_workflow_goal_plan_projection(params.run_id.as_str()) .await?; } Ok(Some(WorkflowRunAdvanceOutcome { snapshot, changed })) @@ -580,7 +580,7 @@ WHERE run_id = ? let outcome = self.workflows.request_workflow_run_cancel(params).await?; if outcome.as_ref().is_some_and(|outcome| outcome.changed) { self.thread_goals - .block_workflow_goal_plan_projection(run_id.as_str()) + .cancel_workflow_goal_plan_projection(run_id.as_str()) .await?; } Ok(outcome.map(|outcome| outcome.snapshot)) @@ -1552,7 +1552,7 @@ async fn create_background_branch_run_if_missing_in_tx( .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(), + auth_profile_ref: route_receipt.effective.auth_profile.clone(), status_reason: Some("queued by workflow branch admission".to_string()), config_fingerprint: params.config_fingerprint.clone(), version_fingerprint: params.version_fingerprint.clone(), @@ -2005,8 +2005,9 @@ fn branch_execution_payload( .permission_profile .as_deref() .map(|id| json!({"id": id, "extends": null})), - "authProfileIdentitySha256": params - .auth_profile_ref + "authProfileIdentitySha256": route_receipt + .effective + .auth_profile .as_deref() .map(|profile| StateRuntime::background_agent_identity_sha256(profile.as_bytes())), "workspace": branch.workspace_json, @@ -2241,7 +2242,43 @@ fn workflow_branch_route_receipt( WorkflowWorkspaceMode::IsolatedWorktree => "isolated_worktree", WorkflowWorkspaceMode::SharedRepository => "shared_repository", }; - admit_workflow_model_route_for_runtime(&requested, runtime, worktree_mode) + let effective_runtime = workflow_route_runtime_for_requested(&requested, runtime); + admit_workflow_model_route_for_runtime(&requested, &effective_runtime, worktree_mode) +} + +fn workflow_route_runtime_for_requested( + requested: &WorkflowModelRoute, + runtime: &WorkflowRouteRuntime, +) -> WorkflowRouteRuntime { + let explicit_auth_profile = requested.routing.as_ref().and_then(|routing| { + routing + .decision + .as_ref() + .and_then(|decision| decision.auth_profile.clone()) + .or_else(|| routing.request.context.auth_profile.clone()) + }); + WorkflowRouteRuntime { + model_gateway: Some(requested.model_gateway.clone()), + provider: Some(requested.provider.clone()), + model: Some(requested.model.clone()), + reasoning: Some(requested.reasoning.clone()), + service_tier: requested.service_tier.clone(), + auth_profile: explicit_auth_profile.or_else(|| { + (runtime.provider.as_deref() == Some(requested.provider.as_str())) + .then(|| runtime.auth_profile.clone()) + .flatten() + }), + approval_policy: requested + .approval_policy + .clone() + .or_else(|| runtime.approval_policy.clone()), + permission_profile: requested + .permission_profile + .clone() + .or_else(|| runtime.permission_profile.clone()), + context_window_tokens: runtime.context_window_tokens, + credit_control: runtime.credit_control.clone(), + } } fn validate_workflow_permission_profile_json( @@ -2873,7 +2910,7 @@ async fn cancel_workflow_run_in_tx( UPDATE workflow_runs SET status = ?, - status_reason = COALESCE(status_reason, ?), + status_reason = ?, reason_code = ?, owner_id = NULL, owner_instance_id = NULL, @@ -3670,6 +3707,14 @@ cleanup: model: "gpt-5.4" reasoning: "high" permission_profile: "read-only" +"# + .to_string() + } else if index == 0 && title_suffix == "fable-child" { + r#" model: + model_gateway: "hasna" + provider: "anthropic" + model: "claude-fable-5" + reasoning: "high" "# .to_string() } else if index == 0 && title_suffix == "bootstrap-context" { @@ -4589,6 +4634,86 @@ WHERE worktree_id = ? ); } + #[tokio::test] + async fn workflow_branch_admission_uses_requested_fable_route_from_openai_coordinator() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(); + upsert_test_thread(&runtime, thread_id).await; + let run = create_unprojected_run( + &runtime, + thread_id, + "wf_branch_fable_child", + parallel_branch_workflow_yaml( + "wf_branch_fable_child", + /*step_count*/ 1, + /*max_parallel_steps*/ 1, + /*max_agents*/ 2, + /*max_worktrees*/ 1, + "fable-child", + ), + ) + .await; + let owner_id = "fable-child-owner"; + let generation = claim_and_advance(&runtime, run.run.run_id.as_str(), owner_id).await; + + let admitted = admit_test_workflow_run_branches( + &runtime, + WorkflowRunBranchAdmissionParams { + run_id: run.run.run_id, + owner_id: owner_id.to_string(), + generation, + auth_profile_ref: Some("openai-coordinator".to_string()), + config_fingerprint: Some("cfg-fable-child".to_string()), + version_fingerprint: Some("version-fable-child".to_string()), + runtime_package_fingerprint: Some("package-fable-child".to_string()), + permission_profile_json: read_only_permission_profile_json(), + route_runtime: workflow_route_runtime(), + parent_agent_run_id: None, + max_active_background_agent_runs: Some(10), + }, + ) + .await + .expect("Fable branch admission should succeed") + .expect("workflow run should remain owned"); + + assert_eq!(1, admitted.admitted.len()); + let branch = &admitted.admitted[0]; + assert_eq!("anthropic", branch.route_receipt.effective.provider); + assert_eq!("claude-fable-5", branch.route_receipt.effective.model); + assert_eq!(None, branch.route_receipt.effective.auth_profile); + let background_run = runtime + .get_background_agent_run(branch.background_agent_run_id.as_str()) + .await + .expect("background run should load") + .expect("background run should exist"); + assert_eq!(None, background_run.auth_profile_ref); + let execution_snapshot = runtime + .get_latest_background_agent_execution_snapshot(branch.background_agent_run_id.as_str()) + .await + .expect("execution snapshot should load") + .expect("execution snapshot should exist"); + assert_eq!( + Some("anthropic"), + execution_snapshot + .payload_json + .get("provider") + .and_then(Value::as_str) + ); + assert_eq!( + Some("claude-fable-5"), + execution_snapshot + .payload_json + .get("model") + .and_then(Value::as_str) + ); + assert!( + execution_snapshot + .payload_json + .get("authProfileIdentitySha256") + .is_none_or(Value::is_null) + ); + } + #[tokio::test] async fn workflow_branch_admission_delivers_bounded_redacted_bootstrap_privately() { let runtime = test_runtime().await; @@ -6218,7 +6343,7 @@ WHERE run_id = ? } #[tokio::test] - async fn workflow_cancel_blocks_projected_goal_plan_and_sanitizes_reason() { + async fn workflow_cancel_terminalizes_projected_goal_plan_and_sanitizes_reason() { let runtime = test_runtime().await; let thread_id = test_thread_id(); upsert_test_thread(&runtime, thread_id).await; @@ -6266,9 +6391,9 @@ WHERE run_id = ? .pop() .expect("projection plan should exist"); assert_eq!(projection.plan_id, plan.plan.plan_id); - assert_eq!(crate::ThreadGoalPlanStatus::Blocked, plan.plan.status); + assert_eq!(crate::ThreadGoalPlanStatus::Cancelled, plan.plan.status); assert_eq!( - vec![crate::ThreadGoalPlanNodeStatus::Blocked], + vec![crate::ThreadGoalPlanNodeStatus::Cancelled], plan.nodes .iter() .map(|node| node.status) @@ -6539,6 +6664,19 @@ WHERE run_id = ? .iter() .all(|verifier| verifier.status == crate::WorkflowRunStepVerifierStatus::Skipped) ); + let plan = runtime + .thread_goals() + .list_thread_goal_plans(thread_id) + .await + .expect("goal plans should list") + .pop() + .expect("projection plan should exist"); + assert_eq!(crate::ThreadGoalPlanStatus::Cancelled, plan.plan.status); + assert!( + plan.nodes + .iter() + .all(|node| node.status == crate::ThreadGoalPlanNodeStatus::Cancelled) + ); } fn gated_approval_workflow_yaml() -> String {