mirror of
https://github.com/openai/codex.git
synced 2026-11-01 19:33:35 +00:00
Add automation dispatch claims
This commit is contained in:
Andrei Eternal
@@ -117,6 +117,9 @@ pub const MEMORIES_DB_FILENAME: &str = "memories_1.sqlite";
|
||||
pub const AUTOMATIONS_DB_FILENAME: &str = "automations_1.sqlite";
|
||||
pub const STATE_DB_FILENAME: &str = "state_5.sqlite";
|
||||
pub const DEFAULT_AUTOMATION_RRULE: &str = "FREQ=HOURLY;INTERVAL=24;BYMINUTE=0";
|
||||
pub const AUTOMATION_CLAIM_LEASE_SECS: i64 = 120;
|
||||
pub const AUTOMATION_RETRY_BACKOFF_SECS: i64 = 10;
|
||||
pub const AUTOMATION_RETRY_BUDGET: i64 = 3;
|
||||
pub const AUTOMATION_RUN_JITTER_WINDOW_SECS: i64 = 120;
|
||||
|
||||
/// Errors encountered during DB operations. Tags: [stage]
|
||||
|
||||
@@ -1,8 +1,15 @@
|
||||
use super::automation_schedule::AutomationSchedule;
|
||||
use super::automation_schedule::compute_next_run_at;
|
||||
use super::*;
|
||||
use crate::AUTOMATION_CLAIM_LEASE_SECS;
|
||||
use crate::AUTOMATION_RETRY_BACKOFF_SECS;
|
||||
use crate::AUTOMATION_RETRY_BUDGET;
|
||||
use crate::Automation;
|
||||
use crate::AutomationCreateParams;
|
||||
use crate::AutomationDispatchClaim;
|
||||
use crate::AutomationDispatchMode;
|
||||
use crate::AutomationDispatchOutcome;
|
||||
use crate::AutomationDispatchRetryOutcome;
|
||||
use crate::AutomationDispatchSettings;
|
||||
use crate::AutomationStatus;
|
||||
use crate::AutomationTarget;
|
||||
@@ -266,6 +273,358 @@ WHERE id = ?
|
||||
let result = query.execute(self.automations_pool.as_ref()).await?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
pub async fn claim_due_automation_dispatch(
|
||||
&self,
|
||||
claimed_by: &str,
|
||||
) -> anyhow::Result<Option<AutomationDispatchClaim>> {
|
||||
let now = Utc::now();
|
||||
let now_ts = now.timestamp();
|
||||
let lease_until = now_ts.saturating_add(AUTOMATION_CLAIM_LEASE_SECS.max(0));
|
||||
let mut tx = self.automations_pool.begin_with("BEGIN IMMEDIATE").await?;
|
||||
let row = sqlx::query_as::<_, AutomationRow>(
|
||||
r#"
|
||||
SELECT *
|
||||
FROM automations
|
||||
WHERE (
|
||||
(status = 'ACTIVE' AND in_flight_run_at IS NULL AND next_run_at IS NOT NULL AND next_run_at <= ?)
|
||||
OR
|
||||
(in_flight_run_at IS NOT NULL
|
||||
AND (status = 'ACTIVE' OR in_flight_dispatch_mode = 'manual')
|
||||
AND (retry_at IS NULL OR retry_at <= ?))
|
||||
)
|
||||
AND (claimed_by IS NULL OR lease_until IS NULL OR lease_until <= ?)
|
||||
ORDER BY COALESCE(retry_at, in_flight_run_at, next_run_at) ASC, id ASC
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(now_ts)
|
||||
.bind(now_ts)
|
||||
.bind(now_ts)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
let claim = match row {
|
||||
Some(row) => claim_automation_row(&mut tx, row, claimed_by, now, lease_until).await?,
|
||||
None => None,
|
||||
};
|
||||
tx.commit().await?;
|
||||
Ok(claim)
|
||||
}
|
||||
|
||||
pub async fn claim_automation_run_now(
|
||||
&self,
|
||||
id: &str,
|
||||
claimed_by: &str,
|
||||
) -> anyhow::Result<AutomationDispatchOutcome> {
|
||||
self.claim_automation_run_now_inner(id, /*expected_owner_thread_id*/ None, claimed_by)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn claim_automation_run_now_if_owner(
|
||||
&self,
|
||||
id: &str,
|
||||
expected_owner_thread_id: ThreadId,
|
||||
claimed_by: &str,
|
||||
) -> anyhow::Result<AutomationDispatchOutcome> {
|
||||
self.claim_automation_run_now_inner(id, Some(expected_owner_thread_id), claimed_by)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn claim_automation_run_now_inner(
|
||||
&self,
|
||||
id: &str,
|
||||
expected_owner_thread_id: Option<ThreadId>,
|
||||
claimed_by: &str,
|
||||
) -> anyhow::Result<AutomationDispatchOutcome> {
|
||||
let now = Utc::now();
|
||||
let now_ts = now.timestamp();
|
||||
let lease_until = now_ts.saturating_add(AUTOMATION_CLAIM_LEASE_SECS.max(0));
|
||||
let mut tx = self.automations_pool.begin_with("BEGIN IMMEDIATE").await?;
|
||||
let Some(row) = load_automation_row_tx(&mut tx, id).await? else {
|
||||
tx.commit().await?;
|
||||
return Ok(AutomationDispatchOutcome::NotFound);
|
||||
};
|
||||
if let Some(expected_owner_thread_id) = expected_owner_thread_id
|
||||
&& row.owner_thread_id != expected_owner_thread_id.to_string()
|
||||
{
|
||||
tx.commit().await?;
|
||||
return Ok(AutomationDispatchOutcome::NotFound);
|
||||
}
|
||||
if claim_is_active(&row, now_ts) {
|
||||
tx.commit().await?;
|
||||
return Ok(AutomationDispatchOutcome::AlreadyClaimed);
|
||||
}
|
||||
|
||||
let claim = if row.in_flight_run_at.is_some() {
|
||||
claim_automation_row(&mut tx, row, claimed_by, now, lease_until).await?
|
||||
} else {
|
||||
claim_automation_row_manual(&mut tx, row, claimed_by, now, lease_until).await?
|
||||
};
|
||||
tx.commit().await?;
|
||||
|
||||
Ok(match claim {
|
||||
Some(claim) => AutomationDispatchOutcome::Claimed(Box::new(claim)),
|
||||
None => AutomationDispatchOutcome::AlreadyClaimed,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn mark_automation_dispatch_started(
|
||||
&self,
|
||||
automation_id: &str,
|
||||
ownership_token: &str,
|
||||
) -> anyhow::Result<bool> {
|
||||
let now = Utc::now().timestamp();
|
||||
let lease_until = now.saturating_add(AUTOMATION_CLAIM_LEASE_SECS.max(0));
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
last_dispatch_started_at = ?,
|
||||
lease_until = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND ownership_token = ?
|
||||
AND claimed_by IS NOT NULL
|
||||
"#,
|
||||
)
|
||||
.bind(now)
|
||||
.bind(lease_until)
|
||||
.bind(now)
|
||||
.bind(automation_id)
|
||||
.bind(ownership_token)
|
||||
.execute(self.automations_pool.as_ref())
|
||||
.await?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
pub async fn renew_automation_dispatch_lease(
|
||||
&self,
|
||||
automation_id: &str,
|
||||
ownership_token: &str,
|
||||
) -> anyhow::Result<bool> {
|
||||
let now = Utc::now().timestamp();
|
||||
let lease_until = now.saturating_add(AUTOMATION_CLAIM_LEASE_SECS.max(0));
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
lease_until = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND ownership_token = ?
|
||||
AND claimed_by IS NOT NULL
|
||||
"#,
|
||||
)
|
||||
.bind(lease_until)
|
||||
.bind(now)
|
||||
.bind(automation_id)
|
||||
.bind(ownership_token)
|
||||
.execute(self.automations_pool.as_ref())
|
||||
.await?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
pub async fn checkpoint_automation_dispatch_progress(
|
||||
&self,
|
||||
automation_id: &str,
|
||||
ownership_token: &str,
|
||||
next_cwd_index: usize,
|
||||
last_error: Option<&str>,
|
||||
) -> anyhow::Result<bool> {
|
||||
let now = Utc::now().timestamp();
|
||||
let lease_until = now.saturating_add(AUTOMATION_CLAIM_LEASE_SECS.max(0));
|
||||
let next_cwd_index =
|
||||
i64::try_from(next_cwd_index).map_err(|_| anyhow::anyhow!("cwd index overflow"))?;
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
dispatch_cwd_index = ?,
|
||||
last_error = ?,
|
||||
lease_until = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND ownership_token = ?
|
||||
AND claimed_by IS NOT NULL
|
||||
"#,
|
||||
)
|
||||
.bind(next_cwd_index)
|
||||
.bind(last_error)
|
||||
.bind(lease_until)
|
||||
.bind(now)
|
||||
.bind(automation_id)
|
||||
.bind(ownership_token)
|
||||
.execute(self.automations_pool.as_ref())
|
||||
.await?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
pub async fn release_automation_dispatch_after_retryable_failure(
|
||||
&self,
|
||||
automation_id: &str,
|
||||
ownership_token: &str,
|
||||
error_message: &str,
|
||||
) -> anyhow::Result<AutomationDispatchRetryOutcome> {
|
||||
let Some(row) = load_automation_row(self.automations_pool.as_ref(), automation_id).await?
|
||||
else {
|
||||
return Ok(AutomationDispatchRetryOutcome::LostClaim);
|
||||
};
|
||||
if row.ownership_token.as_deref() != Some(ownership_token) || row.claimed_by.is_none() {
|
||||
return Ok(AutomationDispatchRetryOutcome::LostClaim);
|
||||
}
|
||||
if row.attempt_count >= AUTOMATION_RETRY_BUDGET {
|
||||
let marked_terminal = self
|
||||
.mark_automation_dispatch_failed_terminal(
|
||||
automation_id,
|
||||
ownership_token,
|
||||
error_message,
|
||||
)
|
||||
.await?;
|
||||
return Ok(if marked_terminal {
|
||||
AutomationDispatchRetryOutcome::MarkedTerminal
|
||||
} else {
|
||||
AutomationDispatchRetryOutcome::LostClaim
|
||||
});
|
||||
}
|
||||
|
||||
let now = Utc::now().timestamp();
|
||||
let retry_at = now.saturating_add(AUTOMATION_RETRY_BACKOFF_SECS.max(0));
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
claimed_by = NULL,
|
||||
ownership_token = NULL,
|
||||
lease_until = NULL,
|
||||
retry_at = ?,
|
||||
last_error = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND ownership_token = ?
|
||||
"#,
|
||||
)
|
||||
.bind(retry_at)
|
||||
.bind(error_message)
|
||||
.bind(now)
|
||||
.bind(automation_id)
|
||||
.bind(ownership_token)
|
||||
.execute(self.automations_pool.as_ref())
|
||||
.await?;
|
||||
Ok(if result.rows_affected() == 1 {
|
||||
AutomationDispatchRetryOutcome::ReleasedForRetry
|
||||
} else {
|
||||
AutomationDispatchRetryOutcome::LostClaim
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn mark_automation_dispatch_failed_terminal(
|
||||
&self,
|
||||
automation_id: &str,
|
||||
ownership_token: &str,
|
||||
error_message: &str,
|
||||
) -> anyhow::Result<bool> {
|
||||
let now = Utc::now().timestamp();
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
claimed_by = NULL,
|
||||
ownership_token = NULL,
|
||||
lease_until = NULL,
|
||||
in_flight_run_at = NULL,
|
||||
in_flight_dispatch_mode = NULL,
|
||||
dispatch_cwd_index = 0,
|
||||
retry_at = NULL,
|
||||
attempt_count = 0,
|
||||
status = CASE
|
||||
WHEN in_flight_dispatch_mode = ? THEN ?
|
||||
ELSE status
|
||||
END,
|
||||
next_run_at = CASE
|
||||
WHEN in_flight_dispatch_mode = ? THEN NULL
|
||||
ELSE next_run_at
|
||||
END,
|
||||
last_error = ?,
|
||||
last_dispatch_succeeded = 0,
|
||||
last_dispatch_completed_at = ?,
|
||||
last_dispatch_failed_at = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND ownership_token = ?
|
||||
"#,
|
||||
)
|
||||
.bind(AutomationDispatchMode::Scheduled.as_str())
|
||||
.bind(AutomationStatus::Paused.as_str())
|
||||
.bind(AutomationDispatchMode::Scheduled.as_str())
|
||||
.bind(error_message)
|
||||
.bind(now)
|
||||
.bind(now)
|
||||
.bind(now)
|
||||
.bind(automation_id)
|
||||
.bind(ownership_token)
|
||||
.execute(self.automations_pool.as_ref())
|
||||
.await?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
pub async fn mark_automation_dispatch_completed(
|
||||
&self,
|
||||
claim: &AutomationDispatchClaim,
|
||||
last_error: Option<&str>,
|
||||
) -> anyhow::Result<bool> {
|
||||
let mut tx = self.automations_pool.begin_with("BEGIN IMMEDIATE").await?;
|
||||
let Some(row) = load_automation_row_tx(&mut tx, claim.automation.id.as_str()).await? else {
|
||||
tx.commit().await?;
|
||||
return Ok(false);
|
||||
};
|
||||
if row.ownership_token.as_deref() != Some(claim.ownership_token.as_str()) {
|
||||
tx.commit().await?;
|
||||
return Ok(false);
|
||||
}
|
||||
let automation = Automation::try_from(row)?;
|
||||
let now = Utc::now().timestamp();
|
||||
let success = last_error.is_none();
|
||||
let next_run_at = match automation.status {
|
||||
AutomationStatus::Active => automation.next_run_at,
|
||||
AutomationStatus::Paused => None,
|
||||
};
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
claimed_by = NULL,
|
||||
ownership_token = NULL,
|
||||
lease_until = NULL,
|
||||
in_flight_run_at = NULL,
|
||||
in_flight_dispatch_mode = NULL,
|
||||
dispatch_cwd_index = 0,
|
||||
retry_at = NULL,
|
||||
attempt_count = 0,
|
||||
last_error = ?,
|
||||
last_run_at = ?,
|
||||
next_run_at = ?,
|
||||
last_dispatch_succeeded = ?,
|
||||
last_dispatch_completed_at = ?,
|
||||
last_dispatch_failed_at = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND ownership_token = ?
|
||||
"#,
|
||||
)
|
||||
.bind(last_error)
|
||||
.bind(now)
|
||||
.bind(next_run_at.map(|value| value.timestamp()))
|
||||
.bind(success)
|
||||
.bind(now)
|
||||
.bind((!success).then_some(now))
|
||||
.bind(now)
|
||||
.bind(automation.id.as_str())
|
||||
.bind(claim.ownership_token.as_str())
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
tx.commit().await?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_automation_row(
|
||||
@@ -436,6 +795,206 @@ fn target_thread_id(target: &AutomationTarget) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn claim_automation_row(
|
||||
conn: &mut sqlx::SqliteConnection,
|
||||
row: AutomationRow,
|
||||
claimed_by: &str,
|
||||
now: DateTime<Utc>,
|
||||
lease_until: i64,
|
||||
) -> anyhow::Result<Option<AutomationDispatchClaim>> {
|
||||
let now_ts = now.timestamp();
|
||||
let automation = Automation::try_from(row.clone())?;
|
||||
let dispatch_mode = match row.in_flight_dispatch_mode.as_deref() {
|
||||
Some(value) => AutomationDispatchMode::parse(value)?,
|
||||
None => AutomationDispatchMode::Scheduled,
|
||||
};
|
||||
let in_flight_run_at = match row.in_flight_run_at {
|
||||
Some(value) => epoch_seconds_to_datetime(value)?,
|
||||
None => now,
|
||||
};
|
||||
let next_run_at_after_claim =
|
||||
if row.in_flight_run_at.is_some() && dispatch_mode == AutomationDispatchMode::Scheduled {
|
||||
automation.next_run_at
|
||||
} else {
|
||||
compute_next_run_at_for_claim(&automation, dispatch_mode, now)?
|
||||
};
|
||||
let ownership_token = Uuid::new_v4().to_string();
|
||||
|
||||
let result = if row.in_flight_run_at.is_some() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
claimed_by = ?,
|
||||
ownership_token = ?,
|
||||
lease_until = ?,
|
||||
retry_at = NULL,
|
||||
attempt_count = attempt_count + 1,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND in_flight_run_at IS NOT NULL
|
||||
AND (claimed_by IS NULL OR lease_until IS NULL OR lease_until <= ?)
|
||||
AND (retry_at IS NULL OR retry_at <= ?)
|
||||
"#,
|
||||
)
|
||||
.bind(claimed_by)
|
||||
.bind(ownership_token.as_str())
|
||||
.bind(lease_until)
|
||||
.bind(now_ts)
|
||||
.bind(row.id.as_str())
|
||||
.bind(now_ts)
|
||||
.bind(now_ts)
|
||||
.execute(&mut *conn)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
claimed_by = ?,
|
||||
ownership_token = ?,
|
||||
lease_until = ?,
|
||||
in_flight_run_at = ?,
|
||||
in_flight_dispatch_mode = ?,
|
||||
dispatch_cwd_index = 0,
|
||||
retry_at = NULL,
|
||||
attempt_count = 1,
|
||||
last_error = NULL,
|
||||
next_run_at = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND in_flight_run_at IS NULL
|
||||
AND status = 'ACTIVE'
|
||||
AND next_run_at IS NOT NULL
|
||||
AND next_run_at <= ?
|
||||
AND (claimed_by IS NULL OR lease_until IS NULL OR lease_until <= ?)
|
||||
"#,
|
||||
)
|
||||
.bind(claimed_by)
|
||||
.bind(ownership_token.as_str())
|
||||
.bind(lease_until)
|
||||
.bind(in_flight_run_at.timestamp())
|
||||
.bind(dispatch_mode.as_str())
|
||||
.bind(next_run_at_after_claim.map(|value| value.timestamp()))
|
||||
.bind(now_ts)
|
||||
.bind(row.id.as_str())
|
||||
.bind(now_ts)
|
||||
.bind(now_ts)
|
||||
.execute(&mut *conn)
|
||||
.await?
|
||||
};
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(AutomationDispatchClaim {
|
||||
automation,
|
||||
ownership_token,
|
||||
dispatch_mode,
|
||||
in_flight_run_at,
|
||||
dispatch_cwd_index: usize::try_from(row.dispatch_cwd_index)
|
||||
.map_err(|_| anyhow::anyhow!("invalid persisted cwd index"))?,
|
||||
attempt_count: row.attempt_count + 1,
|
||||
next_run_at_after_claim,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn claim_automation_row_manual(
|
||||
conn: &mut sqlx::SqliteConnection,
|
||||
row: AutomationRow,
|
||||
claimed_by: &str,
|
||||
now: DateTime<Utc>,
|
||||
lease_until: i64,
|
||||
) -> anyhow::Result<Option<AutomationDispatchClaim>> {
|
||||
let now_ts = now.timestamp();
|
||||
let automation = Automation::try_from(row.clone())?;
|
||||
let ownership_token = Uuid::new_v4().to_string();
|
||||
let next_run_at_after_claim =
|
||||
compute_next_run_at_for_claim(&automation, AutomationDispatchMode::Manual, now)?;
|
||||
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE automations
|
||||
SET
|
||||
claimed_by = ?,
|
||||
ownership_token = ?,
|
||||
lease_until = ?,
|
||||
in_flight_run_at = ?,
|
||||
in_flight_dispatch_mode = ?,
|
||||
dispatch_cwd_index = 0,
|
||||
retry_at = NULL,
|
||||
attempt_count = 1,
|
||||
last_error = NULL,
|
||||
next_run_at = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
AND in_flight_run_at IS NULL
|
||||
AND (claimed_by IS NULL OR lease_until IS NULL OR lease_until <= ?)
|
||||
"#,
|
||||
)
|
||||
.bind(claimed_by)
|
||||
.bind(ownership_token.as_str())
|
||||
.bind(lease_until)
|
||||
.bind(now_ts)
|
||||
.bind(AutomationDispatchMode::Manual.as_str())
|
||||
.bind(next_run_at_after_claim.map(|value| value.timestamp()))
|
||||
.bind(now_ts)
|
||||
.bind(row.id.as_str())
|
||||
.bind(now_ts)
|
||||
.execute(&mut *conn)
|
||||
.await?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(AutomationDispatchClaim {
|
||||
automation,
|
||||
ownership_token,
|
||||
dispatch_mode: AutomationDispatchMode::Manual,
|
||||
in_flight_run_at: now,
|
||||
dispatch_cwd_index: 0,
|
||||
attempt_count: 1,
|
||||
next_run_at_after_claim,
|
||||
}))
|
||||
}
|
||||
|
||||
fn claim_is_active(row: &AutomationRow, now_ts: i64) -> bool {
|
||||
row.claimed_by.is_some()
|
||||
&& row
|
||||
.lease_until
|
||||
.is_some_and(|lease_until| lease_until > now_ts)
|
||||
}
|
||||
|
||||
fn compute_next_run_at_for_claim(
|
||||
automation: &Automation,
|
||||
dispatch_mode: AutomationDispatchMode,
|
||||
now: DateTime<Utc>,
|
||||
) -> anyhow::Result<Option<DateTime<Utc>>> {
|
||||
if automation.status != AutomationStatus::Active {
|
||||
return Ok(None);
|
||||
}
|
||||
if dispatch_mode == AutomationDispatchMode::Manual
|
||||
&& automation
|
||||
.next_run_at
|
||||
.is_some_and(|next_run_at| next_run_at > now)
|
||||
{
|
||||
return Ok(automation.next_run_at);
|
||||
}
|
||||
let schedule = AutomationSchedule::parse(automation.rrule.as_str())?;
|
||||
compute_next_run_at(
|
||||
automation.target.kind(),
|
||||
automation.id.as_str(),
|
||||
&schedule,
|
||||
Some(now),
|
||||
now,
|
||||
)
|
||||
}
|
||||
|
||||
fn epoch_seconds_to_datetime(secs: i64) -> anyhow::Result<DateTime<Utc>> {
|
||||
DateTime::<Utc>::from_timestamp(secs, 0)
|
||||
.ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {secs}"))
|
||||
}
|
||||
|
||||
struct DispatchSettingsJson {
|
||||
workspace_roots: Option<String>,
|
||||
approval_policy: Option<String>,
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
use super::*;
|
||||
use crate::AUTOMATION_RETRY_BUDGET;
|
||||
use crate::AutomationCreateParams;
|
||||
use crate::AutomationDispatchMode;
|
||||
use crate::AutomationDispatchOutcome;
|
||||
use crate::AutomationDispatchRetryOutcome;
|
||||
use crate::AutomationStatus;
|
||||
use crate::AutomationTarget;
|
||||
use crate::AutomationUpdateParams;
|
||||
use crate::runtime::test_support::unique_temp_dir;
|
||||
use chrono::Duration;
|
||||
use chrono::Utc;
|
||||
use codex_protocol::ThreadId;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::PathBuf;
|
||||
@@ -130,3 +136,197 @@ async fn active_heartbeat_targets_are_unique() {
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(codex_home).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn due_claim_advances_schedule_and_completes() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string())
|
||||
.await
|
||||
.expect("state runtime should initialize");
|
||||
let automation = runtime
|
||||
.create_automation(&cron_create_params(
|
||||
thread_id(),
|
||||
codex_home.join("workspace"),
|
||||
))
|
||||
.await
|
||||
.expect("create automation");
|
||||
sqlx::query("UPDATE automations SET next_run_at = ? WHERE id = ?")
|
||||
.bind((Utc::now() - Duration::seconds(1)).timestamp())
|
||||
.bind(automation.id.as_str())
|
||||
.execute(runtime.automations_pool.as_ref())
|
||||
.await
|
||||
.expect("force due automation");
|
||||
|
||||
let claim = runtime
|
||||
.claim_due_automation_dispatch("worker-a")
|
||||
.await
|
||||
.expect("claim due automation")
|
||||
.expect("automation should be due");
|
||||
|
||||
assert_eq!(claim.automation.id, automation.id);
|
||||
assert_eq!(claim.dispatch_mode, AutomationDispatchMode::Scheduled);
|
||||
assert_eq!(claim.dispatch_cwd_index, 0);
|
||||
assert_eq!(claim.attempt_count, 1);
|
||||
assert!(claim.next_run_at_after_claim.is_some());
|
||||
assert_eq!(
|
||||
runtime
|
||||
.claim_due_automation_dispatch("worker-b")
|
||||
.await
|
||||
.expect("second claim should not error"),
|
||||
None
|
||||
);
|
||||
assert!(
|
||||
runtime
|
||||
.mark_automation_dispatch_started(
|
||||
claim.automation.id.as_str(),
|
||||
claim.ownership_token.as_str(),
|
||||
)
|
||||
.await
|
||||
.expect("mark started")
|
||||
);
|
||||
assert!(
|
||||
runtime
|
||||
.checkpoint_automation_dispatch_progress(
|
||||
claim.automation.id.as_str(),
|
||||
claim.ownership_token.as_str(),
|
||||
/*next_cwd_index*/ 1,
|
||||
/*last_error*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("checkpoint")
|
||||
);
|
||||
assert!(
|
||||
runtime
|
||||
.mark_automation_dispatch_completed(&claim, /*last_error*/ None)
|
||||
.await
|
||||
.expect("mark completed")
|
||||
);
|
||||
let reloaded = runtime
|
||||
.get_automation(automation.id.as_str())
|
||||
.await
|
||||
.expect("load automation")
|
||||
.expect("automation should exist");
|
||||
assert_eq!(reloaded.last_run_at.is_some(), true);
|
||||
assert_eq!(reloaded.next_run_at, claim.next_run_at_after_claim);
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(codex_home).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retryable_failure_requeues_until_retry_budget_is_exhausted() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string())
|
||||
.await
|
||||
.expect("state runtime should initialize");
|
||||
let automation = runtime
|
||||
.create_automation(&cron_create_params(
|
||||
thread_id(),
|
||||
codex_home.join("workspace"),
|
||||
))
|
||||
.await
|
||||
.expect("create automation");
|
||||
sqlx::query("UPDATE automations SET next_run_at = ? WHERE id = ?")
|
||||
.bind((Utc::now() - Duration::seconds(1)).timestamp())
|
||||
.bind(automation.id.as_str())
|
||||
.execute(runtime.automations_pool.as_ref())
|
||||
.await
|
||||
.expect("force due automation");
|
||||
|
||||
let claim = runtime
|
||||
.claim_due_automation_dispatch("worker-a")
|
||||
.await
|
||||
.expect("claim due automation")
|
||||
.expect("automation should be due");
|
||||
assert_eq!(
|
||||
runtime
|
||||
.release_automation_dispatch_after_retryable_failure(
|
||||
claim.automation.id.as_str(),
|
||||
claim.ownership_token.as_str(),
|
||||
"temporary failure",
|
||||
)
|
||||
.await
|
||||
.expect("release retry"),
|
||||
AutomationDispatchRetryOutcome::ReleasedForRetry
|
||||
);
|
||||
assert_eq!(
|
||||
runtime
|
||||
.claim_due_automation_dispatch("worker-b")
|
||||
.await
|
||||
.expect("retry should wait for retry_at"),
|
||||
None
|
||||
);
|
||||
sqlx::query("UPDATE automations SET retry_at = ?, attempt_count = ? WHERE id = ?")
|
||||
.bind((Utc::now() - Duration::seconds(1)).timestamp())
|
||||
.bind(AUTOMATION_RETRY_BUDGET)
|
||||
.bind(automation.id.as_str())
|
||||
.execute(runtime.automations_pool.as_ref())
|
||||
.await
|
||||
.expect("make retry terminal");
|
||||
|
||||
let terminal_claim = runtime
|
||||
.claim_due_automation_dispatch("worker-c")
|
||||
.await
|
||||
.expect("claim retry")
|
||||
.expect("retry should be due");
|
||||
assert_eq!(
|
||||
runtime
|
||||
.release_automation_dispatch_after_retryable_failure(
|
||||
terminal_claim.automation.id.as_str(),
|
||||
terminal_claim.ownership_token.as_str(),
|
||||
"still failing",
|
||||
)
|
||||
.await
|
||||
.expect("terminal retry"),
|
||||
AutomationDispatchRetryOutcome::MarkedTerminal
|
||||
);
|
||||
let reloaded = runtime
|
||||
.get_automation(automation.id.as_str())
|
||||
.await
|
||||
.expect("load automation")
|
||||
.expect("automation should exist");
|
||||
assert_eq!(reloaded.status, AutomationStatus::Paused);
|
||||
assert_eq!(reloaded.next_run_at, None);
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(codex_home).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_now_dispatches_paused_automation_once() {
|
||||
let codex_home = unique_temp_dir();
|
||||
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string())
|
||||
.await
|
||||
.expect("state runtime should initialize");
|
||||
let mut params = cron_create_params(thread_id(), codex_home.join("workspace"));
|
||||
params.status = AutomationStatus::Paused;
|
||||
let automation = runtime
|
||||
.create_automation(¶ms)
|
||||
.await
|
||||
.expect("create automation");
|
||||
|
||||
let outcome = runtime
|
||||
.claim_automation_run_now(automation.id.as_str(), "worker-a")
|
||||
.await
|
||||
.expect("claim run now");
|
||||
let AutomationDispatchOutcome::Claimed(claim) = outcome else {
|
||||
panic!("expected paused automation to be claimed");
|
||||
};
|
||||
assert_eq!(claim.dispatch_mode, AutomationDispatchMode::Manual);
|
||||
assert_eq!(claim.next_run_at_after_claim, None);
|
||||
assert!(
|
||||
runtime
|
||||
.mark_automation_dispatch_completed(&claim, /*last_error*/ None)
|
||||
.await
|
||||
.expect("mark completed")
|
||||
);
|
||||
|
||||
let reloaded = runtime
|
||||
.get_automation(automation.id.as_str())
|
||||
.await
|
||||
.expect("load automation")
|
||||
.expect("automation should exist");
|
||||
assert_eq!(reloaded.status, AutomationStatus::Paused);
|
||||
assert_eq!(reloaded.last_run_at.is_some(), true);
|
||||
assert_eq!(reloaded.next_run_at, None);
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(codex_home).await;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user