GenarrativeAI/Genarrative
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: workerize external generation

Browse Source
This commit is contained in:
kdletters
2026-06-05 17:29:08 +08:00
parent 5150925947
commit 8d54ea3374
60 changed files with 5285 additions and 700 deletions
Download Patch File Download Diff File Expand all files Collapse all files

766
server-rs/crates/spacetime-module/src/external_generation.rs Normal file
View File

@@ -0,0 +1,766 @@
use crate::*;
const EXTERNAL_GENERATION_STATUS_PENDING: &str = "pending";
const EXTERNAL_GENERATION_STATUS_RUNNING: &str = "running";
const EXTERNAL_GENERATION_STATUS_COMPLETED: &str = "completed";
const EXTERNAL_GENERATION_STATUS_FAILED: &str = "failed";
const EXTERNAL_GENERATION_STATUS_CANCELLED: &str = "cancelled";
#[spacetimedb::table(
accessor = external_generation_job,
index(
accessor = by_external_generation_job_status_available,
btree(columns = [status, available_at])
),
index(
accessor = by_external_generation_job_worker_id,
btree(columns = [worker_id])
),
index(
accessor = by_external_generation_job_source,
btree(columns = [source_module, source_entity_id])
),
index(
accessor = by_external_generation_job_owner_user_id,
btree(columns = [owner_user_id])
)
)]
#[derive(Clone)]
pub struct ExternalGenerationJob {
#[primary_key]
pub(crate) job_id: String,
#[unique]
pub(crate) dedupe_key: String,
pub(crate) job_kind: String,
pub(crate) owner_user_id: String,
pub(crate) source_module: String,
pub(crate) source_entity_id: String,
pub(crate) request_label: String,
pub(crate) request_payload_json: String,
pub(crate) status: String,
pub(crate) attempt: u32,
pub(crate) max_attempts: u32,
pub(crate) last_error_message: Option<String>,
pub(crate) worker_id: Option<String>,
pub(crate) lease_expires_at: Option<Timestamp>,
pub(crate) available_at: Timestamp,
pub(crate) result_payload_json: Option<String>,
pub(crate) created_at: Timestamp,
pub(crate) started_at: Option<Timestamp>,
pub(crate) completed_at: Option<Timestamp>,
pub(crate) updated_at: Timestamp,
#[default(None::<String>)]
pub(crate) lease_token: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct ExternalGenerationJobEnqueueInput {
pub job_id: String,
pub dedupe_key: String,
pub job_kind: String,
pub owner_user_id: String,
pub source_module: String,
pub source_entity_id: String,
pub request_label: String,
pub request_payload_json: String,
pub max_attempts: u32,
pub available_at_micros: i64,
pub created_at_micros: i64,
}
#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct ExternalGenerationJobClaimInput {
pub worker_id: String,
pub limit: u32,
pub lease_expires_at_micros: i64,
pub claimed_at_micros: i64,
}
#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct ExternalGenerationJobRenewLeaseInput {
pub job_id: String,
pub worker_id: String,
pub lease_token: String,
pub lease_expires_at_micros: i64,
pub renewed_at_micros: i64,
}
#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct ExternalGenerationJobCompleteInput {
pub job_id: String,
pub worker_id: String,
pub lease_token: String,
pub result_payload_json: Option<String>,
pub completed_at_micros: i64,
}
#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct ExternalGenerationJobFailInput {
pub job_id: String,
pub worker_id: String,
pub lease_token: String,
pub error_message: String,
pub retry_after_micros: i64,
pub failed_at_micros: i64,
}
#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct ExternalGenerationJobSnapshot {
pub job_id: String,
pub dedupe_key: String,
pub job_kind: String,
pub owner_user_id: String,
pub source_module: String,
pub source_entity_id: String,
pub request_label: String,
pub request_payload_json: String,
pub status: String,
pub attempt: u32,
pub max_attempts: u32,
pub last_error_message: Option<String>,
pub worker_id: Option<String>,
pub lease_expires_at_micros: Option<i64>,
pub available_at_micros: i64,
pub result_payload_json: Option<String>,
pub created_at_micros: i64,
pub started_at_micros: Option<i64>,
pub completed_at_micros: Option<i64>,
pub updated_at_micros: i64,
pub lease_token: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, SpacetimeType)]
pub struct ExternalGenerationJobProcedureResult {
pub ok: bool,
pub job: Option<ExternalGenerationJobSnapshot>,
pub jobs: Vec<ExternalGenerationJobSnapshot>,
pub error_message: Option<String>,
}
#[spacetimedb::procedure]
pub fn enqueue_external_generation_job_and_return(
ctx: &mut ProcedureContext,
input: ExternalGenerationJobEnqueueInput,
) -> ExternalGenerationJobProcedureResult {
match ctx.try_with_tx(|tx| enqueue_external_generation_job_tx(tx, input.clone())) {
Ok(job) => single_external_generation_job_result(job),
Err(message) => failed_external_generation_job_result(message),
}
}
#[spacetimedb::procedure]
pub fn claim_external_generation_jobs_and_return(
ctx: &mut ProcedureContext,
input: ExternalGenerationJobClaimInput,
) -> ExternalGenerationJobProcedureResult {
match ctx.try_with_tx(|tx| claim_external_generation_jobs_tx(tx, input.clone())) {
Ok(jobs) => ExternalGenerationJobProcedureResult {
ok: true,
job: None,
jobs,
error_message: None,
},
Err(message) => failed_external_generation_job_result(message),
}
}
#[spacetimedb::procedure]
pub fn complete_external_generation_job_and_return(
ctx: &mut ProcedureContext,
input: ExternalGenerationJobCompleteInput,
) -> ExternalGenerationJobProcedureResult {
match ctx.try_with_tx(|tx| complete_external_generation_job_tx(tx, input.clone())) {
Ok(job) => single_external_generation_job_result(job),
Err(message) => failed_external_generation_job_result(message),
}
}
#[spacetimedb::procedure]
pub fn renew_external_generation_job_lease_and_return(
ctx: &mut ProcedureContext,
input: ExternalGenerationJobRenewLeaseInput,
) -> ExternalGenerationJobProcedureResult {
match ctx.try_with_tx(|tx| renew_external_generation_job_lease_tx(tx, input.clone())) {
Ok(job) => single_external_generation_job_result(job),
Err(message) => failed_external_generation_job_result(message),
}
}
#[spacetimedb::procedure]
pub fn fail_external_generation_job_and_return(
ctx: &mut ProcedureContext,
input: ExternalGenerationJobFailInput,
) -> ExternalGenerationJobProcedureResult {
match ctx.try_with_tx(|tx| fail_external_generation_job_tx(tx, input.clone())) {
Ok(job) => single_external_generation_job_result(job),
Err(message) => failed_external_generation_job_result(message),
}
}
fn enqueue_external_generation_job_tx(
ctx: &ReducerContext,
input: ExternalGenerationJobEnqueueInput,
) -> Result<ExternalGenerationJobSnapshot, String> {
validate_required("external_generation_job.job_id", &input.job_id)?;
validate_required("external_generation_job.dedupe_key", &input.dedupe_key)?;
validate_required("external_generation_job.job_kind", &input.job_kind)?;
validate_required(
"external_generation_job.owner_user_id",
&input.owner_user_id,
)?;
validate_required(
"external_generation_job.source_module",
&input.source_module,
)?;
validate_required(
"external_generation_job.source_entity_id",
&input.source_entity_id,
)?;
validate_required(
"external_generation_job.request_label",
&input.request_label,
)?;
validate_required(
"external_generation_job.request_payload_json",
&input.request_payload_json,
)?;
if let Some(row) = ctx
.db
.external_generation_job()
.dedupe_key()
.find(&input.dedupe_key)
{
return Ok(map_external_generation_job_row(row));
}
if ctx
.db
.external_generation_job()
.job_id()
.find(&input.job_id)
.is_some()
{
return Err("external_generation_job.job_id 已存在".to_string());
}
let now = Timestamp::from_micros_since_unix_epoch(input.created_at_micros);
let available_at = Timestamp::from_micros_since_unix_epoch(input.available_at_micros);
let row = ExternalGenerationJob {
job_id: input.job_id.trim().to_string(),
dedupe_key: input.dedupe_key.trim().to_string(),
job_kind: input.job_kind.trim().to_string(),
owner_user_id: input.owner_user_id.trim().to_string(),
source_module: input.source_module.trim().to_string(),
source_entity_id: input.source_entity_id.trim().to_string(),
request_label: input.request_label.trim().to_string(),
request_payload_json: input.request_payload_json.trim().to_string(),
status: EXTERNAL_GENERATION_STATUS_PENDING.to_string(),
attempt: 0,
max_attempts: input.max_attempts.max(1),
last_error_message: None,
worker_id: None,
lease_expires_at: None,
available_at,
result_payload_json: None,
created_at: now,
started_at: None,
completed_at: None,
updated_at: now,
lease_token: None,
};
ctx.db.external_generation_job().insert(row.clone());
Ok(map_external_generation_job_row(row))
}
fn claim_external_generation_jobs_tx(
ctx: &ReducerContext,
input: ExternalGenerationJobClaimInput,
) -> Result<Vec<ExternalGenerationJobSnapshot>, String> {
validate_required("external_generation_job.worker_id", &input.worker_id)?;
if input.limit == 0 {
return Ok(Vec::new());
}
let claim_time = ctx.timestamp;
let lease_duration_micros = duration_between_micros(
input.lease_expires_at_micros,
input.claimed_at_micros,
"external_generation_job.lease_duration",
)?;
let lease_expires_at = timestamp_after_micros(claim_time, lease_duration_micros);
let worker_id = input.worker_id.trim().to_string();
let limit = input.limit.min(64) as usize;
let mut candidates = Vec::new();
candidates.extend(
ctx.db
.external_generation_job()
.by_external_generation_job_status_available()
.filter(&EXTERNAL_GENERATION_STATUS_PENDING.to_string())
.filter(|row| is_external_generation_job_claimable(row, claim_time)),
);
candidates.extend(
ctx.db
.external_generation_job()
.by_external_generation_job_status_available()
.filter(&EXTERNAL_GENERATION_STATUS_RUNNING.to_string())
.filter(|row| is_external_generation_job_claimable(row, claim_time)),
);
candidates.sort_by(|left, right| {
left.available_at
.to_micros_since_unix_epoch()
.cmp(&right.available_at.to_micros_since_unix_epoch())
.then_with(|| {
left.created_at
.to_micros_since_unix_epoch()
.cmp(&right.created_at.to_micros_since_unix_epoch())
})
.then_with(|| left.job_id.cmp(&right.job_id))
});
let mut claimed = Vec::new();
for mut row in candidates.into_iter().take(limit) {
let next_attempt = row.attempt.saturating_add(1);
let lease_token = build_external_generation_lease_token(
&row.job_id,
&worker_id,
next_attempt,
claim_time,
);
row.status = EXTERNAL_GENERATION_STATUS_RUNNING.to_string();
row.worker_id = Some(worker_id.clone());
row.lease_expires_at = Some(lease_expires_at);
row.lease_token = Some(lease_token);
row.attempt = next_attempt;
if row.started_at.is_none() {
row.started_at = Some(claim_time);
}
row.updated_at = claim_time;
persist_external_generation_job_row(ctx, row.clone());
claimed.push(map_external_generation_job_row(row));
}
Ok(claimed)
}
fn complete_external_generation_job_tx(
ctx: &ReducerContext,
input: ExternalGenerationJobCompleteInput,
) -> Result<ExternalGenerationJobSnapshot, String> {
let mut row = get_worker_owned_external_generation_job(
ctx,
&input.job_id,
&input.worker_id,
&input.lease_token,
)?;
let completed_at = ctx.timestamp;
row.status = EXTERNAL_GENERATION_STATUS_COMPLETED.to_string();
row.result_payload_json = input
.result_payload_json
.and_then(|value| normalize_optional_text(value.as_str()));
row.lease_expires_at = None;
row.completed_at = Some(completed_at);
row.updated_at = completed_at;
persist_external_generation_job_row(ctx, row.clone());
Ok(map_external_generation_job_row(row))
}
fn renew_external_generation_job_lease_tx(
ctx: &ReducerContext,
input: ExternalGenerationJobRenewLeaseInput,
) -> Result<ExternalGenerationJobSnapshot, String> {
let mut row = get_worker_owned_external_generation_job(
ctx,
&input.job_id,
&input.worker_id,
&input.lease_token,
)?;
let renewed_at = ctx.timestamp;
let lease_duration_micros = duration_between_micros(
input.lease_expires_at_micros,
input.renewed_at_micros,
"external_generation_job.lease_duration",
)?;
row.lease_expires_at = Some(timestamp_after_micros(renewed_at, lease_duration_micros));
row.updated_at = renewed_at;
persist_external_generation_job_row(ctx, row.clone());
Ok(map_external_generation_job_row(row))
}
fn fail_external_generation_job_tx(
ctx: &ReducerContext,
input: ExternalGenerationJobFailInput,
) -> Result<ExternalGenerationJobSnapshot, String> {
let error_message = input.error_message.trim();
if error_message.is_empty() {
return Err("external_generation_job.error_message 不能为空".to_string());
}
let mut row = get_worker_owned_external_generation_job(
ctx,
&input.job_id,
&input.worker_id,
&input.lease_token,
)?;
let failed_at = ctx.timestamp;
let retry_delay_micros = duration_between_micros(
input.retry_after_micros,
input.failed_at_micros,
"external_generation_job.retry_delay",
)?;
row.last_error_message = Some(error_message.to_string());
row.lease_expires_at = None;
row.worker_id = None;
row.lease_token = None;
row.updated_at = failed_at;
if row.attempt < row.max_attempts {
row.status = EXTERNAL_GENERATION_STATUS_PENDING.to_string();
row.available_at = timestamp_after_micros(failed_at, retry_delay_micros);
} else {
row.status = EXTERNAL_GENERATION_STATUS_FAILED.to_string();
row.completed_at = Some(failed_at);
}
persist_external_generation_job_row(ctx, row.clone());
Ok(map_external_generation_job_row(row))
}
pub(crate) fn validate_external_generation_job_lease_for_tx(
ctx: &ReducerContext,
job_id: &str,
worker_id: &str,
lease_token: &str,
expected_job_kinds: &[&str],
expected_owner_user_id: &str,
expected_source_module: &str,
expected_source_entity_ids: &[String],
) -> Result<(), String> {
let row = get_worker_owned_external_generation_job(ctx, job_id, worker_id, lease_token)?;
if !expected_job_kinds.is_empty()
&& !expected_job_kinds
.iter()
.any(|expected| row.job_kind.trim() == expected.trim())
{
return Err("external_generation_job job_kind 与业务写回不匹配".to_string());
}
if row.owner_user_id.trim() != expected_owner_user_id.trim() {
return Err("external_generation_job owner_user_id 与业务写回不匹配".to_string());
}
if row.source_module.trim() != expected_source_module.trim() {
return Err("external_generation_job source_module 与业务写回不匹配".to_string());
}
if !expected_source_entity_ids
.iter()
.any(|expected| row.source_entity_id.trim() == expected.trim())
{
return Err("external_generation_job source_entity_id 与业务写回不匹配".to_string());
}
Ok(())
}
fn get_worker_owned_external_generation_job(
ctx: &ReducerContext,
job_id: &str,
worker_id: &str,
lease_token: &str,
) -> Result<ExternalGenerationJob, String> {
validate_required("external_generation_job.job_id", job_id)?;
validate_required("external_generation_job.worker_id", worker_id)?;
validate_required("external_generation_job.lease_token", lease_token)?;
let row = ctx
.db
.external_generation_job()
.job_id()
.find(&job_id.trim().to_string())
.ok_or_else(|| "external_generation_job 不存在".to_string())?;
if row.status != EXTERNAL_GENERATION_STATUS_RUNNING {
return Err("external_generation_job 当前不是 running 状态".to_string());
}
if !is_external_generation_job_owned_by_worker(&row, worker_id) {
return Err("external_generation_job worker lease 不匹配".to_string());
}
if !is_external_generation_job_owned_by_lease_token(&row, lease_token) {
return Err("external_generation_job lease token 不匹配".to_string());
}
if !is_external_generation_job_lease_active(&row, ctx.timestamp) {
return Err("external_generation_job lease 已过期".to_string());
}
Ok(row)
}
fn is_external_generation_job_owned_by_worker(
row: &ExternalGenerationJob,
worker_id: &str,
) -> bool {
row.worker_id.as_deref() == Some(worker_id.trim())
}
fn is_external_generation_job_owned_by_lease_token(
row: &ExternalGenerationJob,
lease_token: &str,
) -> bool {
row.lease_token.as_deref() == Some(lease_token.trim())
}
fn is_external_generation_job_lease_active(row: &ExternalGenerationJob, now: Timestamp) -> bool {
row.lease_expires_at
.map(|lease_expires_at| lease_expires_at > now)
.unwrap_or(false)
}
fn is_external_generation_job_claimable(row: &ExternalGenerationJob, now: Timestamp) -> bool {
match row.status.as_str() {
EXTERNAL_GENERATION_STATUS_PENDING => row.available_at <= now,
EXTERNAL_GENERATION_STATUS_RUNNING => row
.lease_expires_at
.map(|lease_expires_at| lease_expires_at <= now)
.unwrap_or(true),
EXTERNAL_GENERATION_STATUS_COMPLETED
| EXTERNAL_GENERATION_STATUS_FAILED
| EXTERNAL_GENERATION_STATUS_CANCELLED => false,
_ => false,
}
}
fn persist_external_generation_job_row(ctx: &ReducerContext, row: ExternalGenerationJob) {
ctx.db
.external_generation_job()
.job_id()
.delete(&row.job_id);
ctx.db.external_generation_job().insert(row);
}
fn map_external_generation_job_row(row: ExternalGenerationJob) -> ExternalGenerationJobSnapshot {
ExternalGenerationJobSnapshot {
job_id: row.job_id,
dedupe_key: row.dedupe_key,
job_kind: row.job_kind,
owner_user_id: row.owner_user_id,
source_module: row.source_module,
source_entity_id: row.source_entity_id,
request_label: row.request_label,
request_payload_json: row.request_payload_json,
status: row.status,
attempt: row.attempt,
max_attempts: row.max_attempts,
last_error_message: row.last_error_message,
worker_id: row.worker_id,
lease_expires_at_micros: row
.lease_expires_at
.map(|value| value.to_micros_since_unix_epoch()),
available_at_micros: row.available_at.to_micros_since_unix_epoch(),
result_payload_json: row.result_payload_json,
created_at_micros: row.created_at.to_micros_since_unix_epoch(),
started_at_micros: row
.started_at
.map(|value| value.to_micros_since_unix_epoch()),
completed_at_micros: row
.completed_at
.map(|value| value.to_micros_since_unix_epoch()),
updated_at_micros: row.updated_at.to_micros_since_unix_epoch(),
lease_token: row.lease_token,
}
}
fn single_external_generation_job_result(
job: ExternalGenerationJobSnapshot,
) -> ExternalGenerationJobProcedureResult {
ExternalGenerationJobProcedureResult {
ok: true,
job: Some(job),
jobs: Vec::new(),
error_message: None,
}
}
fn failed_external_generation_job_result(message: String) -> ExternalGenerationJobProcedureResult {
ExternalGenerationJobProcedureResult {
ok: false,
job: None,
jobs: Vec::new(),
error_message: Some(message),
}
}
fn validate_required(field: &str, value: &str) -> Result<(), String> {
if value.trim().is_empty() {
return Err(format!("{field} 不能为空"));
}
Ok(())
}
fn duration_between_micros(later: i64, earlier: i64, field: &str) -> Result<i64, String> {
let duration = later.saturating_sub(earlier);
if duration <= 0 {
return Err(format!("{field} 必须大于 0"));
}
Ok(duration)
}
fn timestamp_after_micros(timestamp: Timestamp, duration_micros: i64) -> Timestamp {
Timestamp::from_micros_since_unix_epoch(
timestamp
.to_micros_since_unix_epoch()
.saturating_add(duration_micros.max(0)),
)
}
fn build_external_generation_lease_token(
job_id: &str,
worker_id: &str,
attempt: u32,
claimed_at: Timestamp,
) -> String {
format!(
"{}:{}:{}:{}",
job_id.trim(),
worker_id.trim(),
attempt,
claimed_at.to_micros_since_unix_epoch()
)
}
fn normalize_optional_text(value: &str) -> Option<String> {
let normalized = value.trim();
if normalized.is_empty() {
None
} else {
Some(normalized.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn external_generation_job_result_failure_is_structured() {
let result = failed_external_generation_job_result("失败".to_string());
assert!(!result.ok);
assert_eq!(result.error_message.as_deref(), Some("失败"));
assert!(result.jobs.is_empty());
}
#[test]
fn pending_job_is_claimable_only_after_available_time() {
let mut row = external_generation_job_fixture(EXTERNAL_GENERATION_STATUS_PENDING);
row.available_at = micros(1_000);
assert!(!is_external_generation_job_claimable(&row, micros(999)));
assert!(is_external_generation_job_claimable(&row, micros(1_000)));
}
#[test]
fn running_job_is_claimable_only_after_lease_expires() {
let mut row = external_generation_job_fixture(EXTERNAL_GENERATION_STATUS_RUNNING);
row.lease_expires_at = Some(micros(2_000));
assert!(!is_external_generation_job_claimable(&row, micros(1_999)));
assert!(is_external_generation_job_claimable(&row, micros(2_000)));
}
#[test]
fn terminal_job_is_not_claimable() {
for status in [
EXTERNAL_GENERATION_STATUS_COMPLETED,
EXTERNAL_GENERATION_STATUS_FAILED,
EXTERNAL_GENERATION_STATUS_CANCELLED,
] {
let row = external_generation_job_fixture(status);
assert!(!is_external_generation_job_claimable(&row, micros(10_000)));
}
}
#[test]
fn worker_ownership_requires_matching_trimmed_worker_id() {
let mut row = external_generation_job_fixture(EXTERNAL_GENERATION_STATUS_RUNNING);
row.worker_id = Some("worker-a".to_string());
assert!(is_external_generation_job_owned_by_worker(
&row,
" worker-a "
));
assert!(!is_external_generation_job_owned_by_worker(
&row, "worker-b"
));
}
#[test]
fn worker_ownership_requires_matching_trimmed_lease_token() {
let mut row = external_generation_job_fixture(EXTERNAL_GENERATION_STATUS_RUNNING);
row.lease_token = Some("job-1:worker-a:1:1000".to_string());
assert!(is_external_generation_job_owned_by_lease_token(
&row,
" job-1:worker-a:1:1000 "
));
assert!(!is_external_generation_job_owned_by_lease_token(
&row,
"job-1:worker-a:2:2000"
));
}
#[test]
fn worker_lease_is_active_only_before_expiry() {
let mut row = external_generation_job_fixture(EXTERNAL_GENERATION_STATUS_RUNNING);
row.lease_expires_at = Some(micros(2_000));
assert!(is_external_generation_job_lease_active(&row, micros(1_999)));
assert!(!is_external_generation_job_lease_active(
&row,
micros(2_000)
));
}
#[test]
fn lease_token_changes_with_claim_attempt() {
let first =
build_external_generation_lease_token("extgen-test", "worker-a", 1, micros(1_000));
let second =
build_external_generation_lease_token("extgen-test", "worker-a", 2, micros(2_000));
assert_ne!(first, second);
}
#[test]
fn positive_duration_between_client_times_is_preserved() {
assert_eq!(
duration_between_micros(3_500, 1_000, "external_generation_job.lease_duration"),
Ok(2_500),
);
assert!(duration_between_micros(1_000, 1_000, "duration").is_err());
}
fn external_generation_job_fixture(status: &str) -> ExternalGenerationJob {
ExternalGenerationJob {
job_id: "extgen-test".to_string(),
dedupe_key: "puzzle:compile:test".to_string(),
job_kind: "puzzle_compile_draft".to_string(),
owner_user_id: "user-1".to_string(),
source_module: "puzzle".to_string(),
source_entity_id: "session-1".to_string(),
request_label: "拼图首关草稿生成".to_string(),
request_payload_json: r#"{"sessionId":"session-1"}"#.to_string(),
status: status.to_string(),
attempt: 0,
max_attempts: 1,
last_error_message: None,
worker_id: None,
lease_expires_at: None,
available_at: micros(0),
result_payload_json: None,
created_at: micros(0),
started_at: None,
completed_at: None,
updated_at: micros(0),
lease_token: None,
}
}
fn micros(value: i64) -> Timestamp {
Timestamp::from_micros_since_unix_epoch(value)
}
}

2
server-rs/crates/spacetime-module/src/lib.rs
View File

@@ -30,6 +30,7 @@ mod big_fish;
mod custom_world;
mod domain_types;
mod entry;
mod external_generation;
mod gameplay;
mod jump_hop;
mod match3d;
@@ -49,6 +50,7 @@ pub use big_fish::*;
pub use custom_world::*;
pub use domain_types::*;
pub use entry::*;
pub use external_generation::*;
pub use gameplay::*;
pub use jump_hop::*;
pub use match3d::*;

1
server-rs/crates/spacetime-module/src/migration.rs
View File

@@ -178,6 +178,7 @@ macro_rules! migration_tables {
ai_text_chunk,
ai_result_reference,
ai_task_event,
external_generation_job,
runtime_snapshot,
runtime_setting,
creation_entry_config,

241
server-rs/crates/spacetime-module/src/puzzle.rs
View File

@@ -12,16 +12,17 @@ use module_puzzle::{
PuzzleAgentSessionGetInput, PuzzleAgentSessionProcedureResult, PuzzleAgentSessionSnapshot,
PuzzleAgentStage, PuzzleAnchorPack, PuzzleDraftCompileFailureInput, PuzzleDraftCompileInput,
PuzzleFormDraftSaveInput, PuzzleGeneratedImageCandidate, PuzzleGeneratedImagesSaveInput,
PuzzleLeaderboardEntry, PuzzleLeaderboardSubmitInput, PuzzlePublicationStatus,
PuzzlePublishInput, PuzzleRecommendedNextWork, PuzzleResultDraft, PuzzleRunDragInput,
PuzzleRunGetInput, PuzzleRunNextLevelInput, PuzzleRunPauseInput, PuzzleRunProcedureResult,
PuzzleRunPropInput, PuzzleRunSnapshot, PuzzleRunStartInput, PuzzleRunSwapInput,
PuzzleRuntimeLevelStatus, PuzzleSelectCoverImageInput, PuzzleUiBackgroundSaveInput,
PuzzleWorkDeleteInput, PuzzleWorkGetInput, PuzzleWorkLikeRecordInput as PuzzleWorkLikeInput,
PuzzleWorkPointIncentiveClaimInput, PuzzleWorkProcedureResult, PuzzleWorkProfile,
PuzzleWorkRemixInput, PuzzleWorkUpsertInput, PuzzleWorksListInput, PuzzleWorksProcedureResult,
apply_publish_overrides_to_draft, apply_selected_candidate, build_form_draft_from_seed,
build_result_preview, compile_result_draft_from_seed, create_work_profile, infer_anchor_pack,
PuzzleLeaderboardEntry, PuzzleLeaderboardSubmitInput, PuzzleLevelGenerationFailureInput,
PuzzlePublicationStatus, PuzzlePublishInput, PuzzleRecommendedNextWork, PuzzleResultDraft,
PuzzleRunDragInput, PuzzleRunGetInput, PuzzleRunNextLevelInput, PuzzleRunPauseInput,
PuzzleRunProcedureResult, PuzzleRunPropInput, PuzzleRunSnapshot, PuzzleRunStartInput,
PuzzleRunSwapInput, PuzzleRuntimeLevelStatus, PuzzleSelectCoverImageInput,
PuzzleUiBackgroundSaveInput, PuzzleWorkDeleteInput, PuzzleWorkGetInput,
PuzzleWorkLikeRecordInput as PuzzleWorkLikeInput, PuzzleWorkPointIncentiveClaimInput,
PuzzleWorkProcedureResult, PuzzleWorkProfile, PuzzleWorkRemixInput, PuzzleWorkUpsertInput,
PuzzleWorksListInput, PuzzleWorksProcedureResult, apply_publish_overrides_to_draft,
apply_selected_candidate, build_form_draft_from_seed, build_result_preview,
compile_result_draft_from_seed, create_work_profile, infer_anchor_pack,
mark_failed_puzzle_result_draft_generation, normalize_puzzle_draft, normalize_puzzle_levels,
normalize_theme_tags, publish_work_profile, replace_puzzle_level, select_next_profiles,
selected_profile_level_after_runtime_level, selected_puzzle_level, tag_similarity_score,
@@ -36,9 +37,14 @@ use spacetimedb::{
};
use crate::auth::user_account;
use crate::validate_external_generation_job_lease_for_tx;
const PUZZLE_POINT_INCENTIVE_DEFAULT_U64: u64 = 0;
const WORK_VISIBLE_DEFAULT: bool = true;
const PUZZLE_EXTERNAL_GENERATION_SOURCE_MODULE: &str = "puzzle";
const PUZZLE_COMPILE_DRAFT_JOB_KIND: &str = "puzzle_compile_draft";
const PUZZLE_GENERATE_IMAGES_JOB_KIND: &str = "puzzle_generate_images";
const PUZZLE_GENERATE_UI_BACKGROUND_JOB_KIND: &str = "puzzle_generate_ui_background";
/// 拼图 Agent session 真相表。
/// 当前只保存结构化字段与 JSON 草稿，不提前拆出更多编辑态子表。
@@ -388,6 +394,25 @@ pub fn mark_puzzle_draft_generation_failed(
}
}
#[spacetimedb::procedure]
pub fn mark_puzzle_level_generation_failed(
ctx: &mut ProcedureContext,
input: PuzzleLevelGenerationFailureInput,
) -> PuzzleAgentSessionProcedureResult {
match ctx.try_with_tx(|tx| mark_puzzle_level_generation_failed_tx(tx, input.clone())) {
Ok(session) => PuzzleAgentSessionProcedureResult {
ok: true,
session: Some(session),
error_message: None,
},
Err(message) => PuzzleAgentSessionProcedureResult {
ok: false,
session: None,
error_message: Some(message),
},
}
}
/// 保存拼图入口表单草稿。
/// 中文注释：该 procedure 只更新 session 与创作中心草稿卡，不触发图片生成或发布校验。
#[spacetimedb::procedure]
@@ -978,6 +1003,26 @@ fn compile_puzzle_agent_draft_tx(
ctx: &TxContext,
input: PuzzleDraftCompileInput,
) -> Result<PuzzleAgentSessionSnapshot, String> {
match (
input.external_generation_job_id.as_deref(),
input.external_generation_worker_id.as_deref(),
input.external_generation_lease_token.as_deref(),
) {
(Some(job_id), Some(worker_id), Some(lease_token)) => {
validate_puzzle_external_generation_write_guard(
ctx,
job_id,
worker_id,
lease_token,
&[PUZZLE_COMPILE_DRAFT_JOB_KIND],
&input.session_id,
&input.owner_user_id,
None,
)?;
}
(None, None, None) => {}
_ => return Err("拼图草稿编译外部生成 guard 不完整".to_string()),
}
let row = get_owned_session_row(ctx, &input.session_id, &input.owner_user_id)?;
if row.seed_text.trim().is_empty() {
return Err("请先填写拼图作品信息".to_string());
@@ -1028,6 +1073,16 @@ fn mark_puzzle_draft_generation_failed_tx(
ctx: &TxContext,
input: PuzzleDraftCompileFailureInput,
) -> Result<PuzzleAgentSessionSnapshot, String> {
validate_puzzle_external_generation_write_guard(
ctx,
&input.external_generation_job_id,
&input.external_generation_worker_id,
&input.external_generation_lease_token,
&[PUZZLE_COMPILE_DRAFT_JOB_KIND],
&input.session_id,
&input.owner_user_id,
None,
)?;
let row = get_owned_session_row(ctx, &input.session_id, &input.owner_user_id)?;
let updated_at = Timestamp::from_micros_since_unix_epoch(input.failed_at_micros);
let draft = match deserialize_optional_draft(&row.draft_json)? {
@@ -1079,6 +1134,88 @@ fn mark_puzzle_draft_generation_failed_tx(
)
}
fn mark_puzzle_level_generation_failed_tx(
ctx: &TxContext,
input: PuzzleLevelGenerationFailureInput,
) -> Result<PuzzleAgentSessionSnapshot, String> {
validate_puzzle_external_generation_write_guard(
ctx,
&input.external_generation_job_id,
&input.external_generation_worker_id,
&input.external_generation_lease_token,
&[
PUZZLE_GENERATE_IMAGES_JOB_KIND,
PUZZLE_GENERATE_UI_BACKGROUND_JOB_KIND,
],
&input.session_id,
&input.owner_user_id,
input.level_id.as_deref(),
)?;
let row = get_owned_session_row(ctx, &input.session_id, &input.owner_user_id)?;
let updated_at = Timestamp::from_micros_since_unix_epoch(input.failed_at_micros);
let mut draft = match deserialize_optional_draft(&row.draft_json)? {
Some(draft) => draft,
None => {
let anchor_pack = deserialize_anchor_pack(&row.anchor_pack_json)?;
let messages = list_session_messages(ctx, &row.session_id);
compile_result_draft_from_seed(&anchor_pack, &messages, Some(&row.seed_text))
}
};
if let Some(levels) = deserialize_optional_levels_input(input.levels_json.as_deref())? {
// 中文注释：新增关卡可能还没完成自动保存，失败回写必须以本次 action 快照作为目标集合。
draft.levels = levels;
}
draft = mark_puzzle_level_generation_failed_draft(draft, input.level_id.as_deref())?;
let next_stage = resolve_failed_puzzle_agent_stage(row.stage, &draft);
upsert_puzzle_draft_work_profile(
ctx,
&row.session_id,
&row.owner_user_id,
&draft,
input.failed_at_micros,
)?;
replace_puzzle_agent_session(
ctx,
&row,
PuzzleAgentSessionRow {
session_id: row.session_id.clone(),
owner_user_id: row.owner_user_id.clone(),
seed_text: row.seed_text.clone(),
current_turn: row.current_turn,
progress_percent: row.progress_percent.max(94),
stage: next_stage,
anchor_pack_json: row.anchor_pack_json.clone(),
draft_json: Some(serialize_json(&draft)),
last_assistant_reply: Some(input.error_message),
published_profile_id: row.published_profile_id.clone(),
created_at: row.created_at,
updated_at,
},
);
get_puzzle_agent_session_tx(
ctx,
PuzzleAgentSessionGetInput {
session_id: input.session_id,
owner_user_id: input.owner_user_id,
},
)
}
fn mark_puzzle_level_generation_failed_draft(
draft: PuzzleResultDraft,
level_id: Option<&str>,
) -> Result<PuzzleResultDraft, String> {
let target_level =
selected_puzzle_level(&draft, level_id).ok_or_else(|| "拼图关卡不存在".to_string())?;
let mut next_level = target_level;
next_level.generation_status = "failed".to_string();
let mut draft = replace_puzzle_level(&draft, next_level).map_err(|error| error.to_string())?;
module_puzzle::sync_primary_level_fields(&mut draft);
Ok(draft)
}
fn resolve_failed_puzzle_agent_stage(
current_stage: PuzzleAgentStage,
draft: &PuzzleResultDraft,
@@ -1094,6 +1231,34 @@ fn resolve_failed_puzzle_agent_stage(
}
}
fn validate_puzzle_external_generation_write_guard(
ctx: &TxContext,
job_id: &str,
worker_id: &str,
lease_token: &str,
expected_job_kinds: &[&str],
session_id: &str,
owner_user_id: &str,
level_id: Option<&str>,
) -> Result<(), String> {
let session_entity_id = session_id.trim().to_string();
let mut source_entity_ids = vec![session_entity_id.clone()];
if let Some(level_id) = level_id.map(str::trim).filter(|value| !value.is_empty()) {
source_entity_ids.push(format!("{session_entity_id}:{level_id}"));
}
validate_external_generation_job_lease_for_tx(
ctx.as_ref(),
job_id,
worker_id,
lease_token,
expected_job_kinds,
owner_user_id,
PUZZLE_EXTERNAL_GENERATION_SOURCE_MODULE,
&source_entity_ids,
)
}
fn save_puzzle_form_draft_tx(
ctx: &TxContext,
input: PuzzleFormDraftSaveInput,
@@ -1151,6 +1316,19 @@ fn save_puzzle_generated_images_tx(
ctx: &TxContext,
input: PuzzleGeneratedImagesSaveInput,
) -> Result<PuzzleAgentSessionSnapshot, String> {
validate_puzzle_external_generation_write_guard(
ctx,
&input.external_generation_job_id,
&input.external_generation_worker_id,
&input.external_generation_lease_token,
&[
PUZZLE_COMPILE_DRAFT_JOB_KIND,
PUZZLE_GENERATE_IMAGES_JOB_KIND,
],
&input.session_id,
&input.owner_user_id,
input.level_id.as_deref(),
)?;
let row = get_owned_session_row(ctx, &input.session_id, &input.owner_user_id)?;
let mut draft = deserialize_draft_required(&row.draft_json)?;
let previous_primary_level_name = draft.level_name.clone();
@@ -1235,6 +1413,16 @@ fn save_puzzle_ui_background_tx(
ctx: &TxContext,
input: PuzzleUiBackgroundSaveInput,
) -> Result<PuzzleAgentSessionSnapshot, String> {
validate_puzzle_external_generation_write_guard(
ctx,
&input.external_generation_job_id,
&input.external_generation_worker_id,
&input.external_generation_lease_token,
&[PUZZLE_GENERATE_UI_BACKGROUND_JOB_KIND],
&input.session_id,
&input.owner_user_id,
input.level_id.as_deref(),
)?;
let row = get_owned_session_row(ctx, &input.session_id, &input.owner_user_id)?;
let mut draft = deserialize_draft_required(&row.draft_json)?;
if let Some(levels) = deserialize_optional_levels_input(input.levels_json.as_deref())? {
@@ -3897,6 +4085,39 @@ mod tests {
);
}
#[test]
fn level_generation_failure_only_marks_target_level_failed() {
let anchor_pack = infer_anchor_pack("画面描述：一只猫在雨夜灯牌下回头。", None);
let mut draft = compile_result_draft_from_seed(
&anchor_pack,
&[],
Some("画面描述：一只猫在雨夜灯牌下回头。"),
);
draft.levels[0].generation_status = "ready".to_string();
draft.levels[0].cover_image_src = Some("/generated-puzzle-assets/first.png".to_string());
let mut second_level = draft.levels[0].clone();
second_level.level_id = "puzzle-level-2".to_string();
second_level.level_name = "第二关".to_string();
second_level.picture_description = "第二关画面".to_string();
second_level.cover_image_src = None;
second_level.cover_asset_id = None;
second_level.candidates = Vec::new();
second_level.selected_candidate_id = None;
second_level.generation_status = "generating".to_string();
draft.levels.push(second_level);
let failed = mark_puzzle_level_generation_failed_draft(draft, Some("puzzle-level-2"))
.expect("target level should be marked failed");
assert_eq!(failed.levels[0].generation_status, "ready");
assert_eq!(
failed.levels[0].cover_image_src.as_deref(),
Some("/generated-puzzle-assets/first.png")
);
assert_eq!(failed.levels[1].generation_status, "failed");
assert_eq!(failed.generation_status, "ready");
}
#[test]
fn puzzle_recommendation_score_prefers_same_author_weight() {
let left = PuzzleWorkProfile {