diff --git a/server-rs/crates/spacetime-module/src/external_generation.rs b/server-rs/crates/spacetime-module/src/external_generation.rs index 5e93b162c..07c2cbab4 100644 --- a/server-rs/crates/spacetime-module/src/external_generation.rs +++ b/server-rs/crates/spacetime-module/src/external_generation.rs @@ -15,6 +15,7 @@ const EXTERNAL_GENERATION_EVENT_LEASE_RENEWED: &str = "lease_renewed"; const EXTERNAL_GENERATION_EVENT_COMPLETED: &str = "completed"; const EXTERNAL_GENERATION_EVENT_FAILED: &str = "failed"; const EXTERNAL_GENERATION_EVENT_ACKNOWLEDGED: &str = "acknowledged"; +const EXTERNAL_GENERATION_EVENT_PROVIDER_CHECKPOINT: &str = "provider_checkpoint"; const EXTERNAL_GENERATION_EDITOR_SOURCE_MODULE: &str = "editor-canvas"; const EXTERNAL_GENERATION_FINAL_ATTEMPT_LEASE_EXPIRED_MESSAGE: &str = "worker 最终执行次数的 lease 已过期，任务已终止"; @@ -90,6 +91,12 @@ pub struct ExternalGenerationJob { pub(crate) notification_acknowledged_at: Option, #[default(None::)] pub(crate) phase: Option, + // 中文注释：provider 任务 checkpoint 晚于主表加入。没有 provider_task_id 才允许 submit， + // 有 checkpoint 的 attempt 只允许续跑查询与落库，避免崩溃重试造成二次提交与二次扣费。 + #[default(None::)] + pub(crate) provider_kind: Option, + #[default(None::)] + pub(crate) provider_task_id: Option, } #[spacetimedb::table( @@ -193,6 +200,15 @@ pub struct ExternalGenerationJobPhaseUpdateInput { pub phase: String, } +#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)] +pub struct ExternalGenerationJobProviderCheckpointInput { + pub job_id: String, + pub worker_id: String, + pub lease_token: String, + pub provider_kind: String, + pub provider_task_id: String, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq, SpacetimeType)] pub enum ExternalGenerationJobPhaseUpdateFailureKind { LeaseFencingRejected, @@ -283,6 +299,8 @@ pub struct ExternalGenerationJobSnapshot { pub refund_ledger_id: Option, pub notification_acknowledged_at_micros: Option, pub phase: Option, + pub provider_kind: Option, + pub provider_task_id: Option, } #[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)] @@ -504,6 +522,23 @@ pub fn renew_external_generation_job_lease_and_return( } } +#[spacetimedb::procedure] +pub fn set_external_generation_job_provider_checkpoint_and_return( + ctx: &mut ProcedureContext, + input: ExternalGenerationJobProviderCheckpointInput, +) -> ExternalGenerationJobProcedureResult { + let caller = ctx.sender(); + match ctx.try_with_tx(|tx| { + crate::editor_project_storage::require_editor_generation_runtime_service_identity( + tx, caller, + )?; + set_external_generation_job_provider_checkpoint_tx(tx, input.clone()) + }) { + Ok(job) => single_external_generation_job_result(job), + Err(message) => failed_external_generation_job_result(message), + } +} + #[spacetimedb::procedure] pub fn update_external_generation_job_phase_and_return( ctx: &mut ProcedureContext, @@ -819,6 +854,8 @@ fn enqueue_external_generation_job_tx( refund_ledger_id: None, notification_acknowledged_at: None, phase: None, + provider_kind: None, + provider_task_id: None, }; persist_external_generation_job_row(ctx, row.clone()); insert_external_generation_job_event( @@ -1578,6 +1615,45 @@ fn update_external_generation_job_phase_tx( Ok(map_external_generation_job_row(row)) } +/// 写入 provider 任务 checkpoint。 +/// +/// 只有持有当前租约的 worker 能写；checkpoint 一旦存在就不再允许覆盖， +/// 这是 at-most-once submit 的数据层保证：崩溃后重新 claim 的 attempt 只能读到既有 +/// provider_task_id 并继续查询，不会再去 submit 一次。 +fn set_external_generation_job_provider_checkpoint_tx( + ctx: &ReducerContext, + input: ExternalGenerationJobProviderCheckpointInput, +) -> Result { + let provider_kind = normalize_optional_text(&input.provider_kind) + .ok_or_else(|| "external_generation_job.provider_kind 不能为空".to_string())?; + let provider_task_id = normalize_optional_text(&input.provider_task_id) + .ok_or_else(|| "external_generation_job.provider_task_id 不能为空".to_string())?; + + let mut row = get_worker_owned_external_generation_job( + ctx, + &input.job_id, + &input.worker_id, + &input.lease_token, + )?; + ensure_external_generation_job_provider_checkpoint_absent(&row)?; + + row.provider_kind = Some(provider_kind); + row.provider_task_id = Some(provider_task_id.clone()); + row.updated_at = ctx.timestamp; + persist_external_generation_job_row(ctx, row.clone()); + insert_external_generation_job_event( + ctx, + &row, + EXTERNAL_GENERATION_EVENT_PROVIDER_CHECKPOINT, + Some(format!( + "worker 已记录 provider 任务 checkpoint {provider_task_id}" + )), + Some(input.worker_id), + ctx.timestamp, + ); + Ok(map_external_generation_job_row(row)) +} + fn fail_external_generation_job_tx( ctx: &ReducerContext, input: ExternalGenerationJobFailInput, @@ -1801,6 +1877,19 @@ fn get_worker_owned_external_generation_job( .map_err(|error| error.message) } +/// at-most-once submit 的数据层不变量：已有 checkpoint 的 job 不允许再写第二次。 +fn ensure_external_generation_job_provider_checkpoint_absent( + row: &ExternalGenerationJob, +) -> Result<(), String> { + if row.provider_task_id.is_some() { + return Err( + "external_generation_job.provider_task_id 已存在，禁止覆盖 provider checkpoint" + .to_string(), + ); + } + Ok(()) +} + #[derive(Debug)] struct ExternalGenerationJobPhaseUpdateError { kind: ExternalGenerationJobPhaseUpdateFailureKind, @@ -2573,6 +2662,8 @@ fn map_external_generation_job_row(row: ExternalGenerationJob) -> ExternalGenera refund_ledger_id: row.refund_ledger_id, notification_acknowledged_at_micros, phase: row.phase, + provider_kind: row.provider_kind, + provider_task_id: row.provider_task_id, } } @@ -2641,6 +2732,9 @@ fn map_external_generation_job_summary_to_compat_snapshot( refund_ledger_id: summary.refund_ledger_id, notification_acknowledged_at_micros: summary.notification_acknowledged_at_micros, phase: summary.phase, + // 中文注释：轻量摘要不投影 provider checkpoint，兼容快照按无 checkpoint 返回。 + provider_kind: None, + provider_task_id: None, } } @@ -3095,6 +3189,22 @@ mod tests { } } + #[test] + fn provider_checkpoint_cannot_be_overwritten_after_first_write() { + let mut row = external_generation_job_fixture(EXTERNAL_GENERATION_STATUS_RUNNING); + assert!( + ensure_external_generation_job_provider_checkpoint_absent(&row).is_ok(), + "首次 checkpoint 必须允许写入" + ); + + row.provider_kind = Some("tripo".to_string()); + row.provider_task_id = Some("task-1".to_string()); + assert!( + ensure_external_generation_job_provider_checkpoint_absent(&row).is_err(), + "已有 checkpoint 必须拒绝覆盖，保证 at-most-once submit" + ); + } + #[test] fn external_generation_job_result_failure_is_structured() { let result = failed_external_generation_job_result("失败".to_string()); @@ -3847,6 +3957,8 @@ mod tests { refund_ledger_id: None, notification_acknowledged_at: None, phase: None, + provider_kind: None, + provider_task_id: None, } } diff --git a/server-rs/crates/spacetime-module/src/migration.rs b/server-rs/crates/spacetime-module/src/migration.rs index 4aad91090..62b5e11f7 100644 --- a/server-rs/crates/spacetime-module/src/migration.rs +++ b/server-rs/crates/spacetime-module/src/migration.rs @@ -1593,6 +1593,17 @@ fn normalize_migration_row(table_name: &str, value: &serde_json::Value) -> serde .or_insert(serde_json::Value::Null); } } + if table_name == "external_generation_job" { + if let Some(object) = next_value.as_object_mut() { + // 中文注释：provider 任务 checkpoint 晚于外部生成主表加入，旧迁移包按尚未提交 provider 兼容。 + object + .entry("provider_kind".to_string()) + .or_insert(serde_json::Value::Null); + object + .entry("provider_task_id".to_string()) + .or_insert(serde_json::Value::Null); + } + } if table_name == "big_fish_creation_session" { if let Some(object) = next_value.as_object_mut() { // 中文注释：旧迁移包没有公开游玩次数字段，导入时按新建作品默认 0 兼容。