diff --git a/server-rs/crates/spacetime-client/src/active/mapper/external_generation.rs b/server-rs/crates/spacetime-client/src/active/mapper/external_generation.rs index d4ba85760..2919a7a75 100644 --- a/server-rs/crates/spacetime-client/src/active/mapper/external_generation.rs +++ b/server-rs/crates/spacetime-client/src/active/mapper/external_generation.rs @@ -65,6 +65,20 @@ impl From for ExternalGenerationJob } } +impl From + for ExternalGenerationJobProviderCheckpointInput +{ + fn from(input: ExternalGenerationJobProviderCheckpointRecordInput) -> Self { + Self { + job_id: input.job_id, + worker_id: input.worker_id, + lease_token: input.lease_token, + provider_kind: input.provider_kind, + provider_task_id: input.provider_task_id, + } + } +} + impl From for ExternalGenerationJobFailInput { fn from(input: ExternalGenerationJobFailRecordInput) -> Self { Self { @@ -300,6 +314,8 @@ pub(crate) fn map_external_generation_job_snapshot( .map(format_timestamp_micros), notification_acknowledged_at_micros: snapshot.notification_acknowledged_at_micros, phase: snapshot.phase, + provider_kind: snapshot.provider_kind, + provider_task_id: snapshot.provider_task_id, } } @@ -414,6 +430,15 @@ pub struct ExternalGenerationJobAcknowledgeRecordInput { pub acknowledged_at_micros: i64, } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalGenerationJobProviderCheckpointRecordInput { + pub job_id: String, + pub worker_id: String, + pub lease_token: String, + pub provider_kind: String, + pub provider_task_id: String, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct ExternalGenerationJobRecord { pub job_id: String, @@ -443,6 +468,9 @@ pub struct ExternalGenerationJobRecord { pub notification_acknowledged_at: Option, pub notification_acknowledged_at_micros: Option, pub phase: Option, + /// provider 任务 checkpoint：有值时只允许续跑查询，不允许再次 submit。 + pub provider_kind: Option, + pub provider_task_id: Option, } #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/server-rs/crates/spacetime-client/src/external_generation.rs b/server-rs/crates/spacetime-client/src/external_generation.rs index 0b2a2abfc..4f96e7817 100644 --- a/server-rs/crates/spacetime-client/src/external_generation.rs +++ b/server-rs/crates/spacetime-client/src/external_generation.rs @@ -360,6 +360,36 @@ impl SpacetimeClient { .await } + /// 写入 provider 任务 checkpoint。 + /// + /// 只在 submit 成功、拿到 provider task id 之后调用一次；procedure 侧对已有 + /// checkpoint 的 job 直接拒绝，调用方不得用重试来覆盖，这是 at-most-once submit + /// 的数据层保证。写入失败意味着本次 attempt 无法安全续跑，按失败处理而不是重新 submit。 + pub async fn set_external_generation_job_provider_checkpoint( + &self, + input: ExternalGenerationJobProviderCheckpointRecordInput, + ) -> Result { + let procedure_input = input.into(); + + self.call_after_connect( + "set_external_generation_job_provider_checkpoint_and_return", + move |connection, sender| { + connection + .procedures() + .set_external_generation_job_provider_checkpoint_and_return_then( + procedure_input, + move |_, result| { + let mapped = result + .map_err(SpacetimeClientError::from_sdk_error) + .and_then(map_external_generation_job_procedure_result); + send_once(&sender, mapped); + }, + ); + }, + ) + .await + } + pub async fn get_external_generation_job( &self, input: ExternalGenerationJobGetRecordInput, diff --git a/server-rs/crates/spacetime-client/src/module_bindings.rs b/server-rs/crates/spacetime-client/src/module_bindings.rs index 00999e3d6..e199c48ca 100644 --- a/server-rs/crates/spacetime-client/src/module_bindings.rs +++ b/server-rs/crates/spacetime-client/src/module_bindings.rs @@ -428,6 +428,7 @@ pub mod external_generation_job_phase_update_failure_kind_type; pub mod external_generation_job_phase_update_input_type; pub mod external_generation_job_phase_update_procedure_result_type; pub mod external_generation_job_procedure_result_type; +pub mod external_generation_job_provider_checkpoint_input_type; pub mod external_generation_job_renew_lease_input_type; pub mod external_generation_job_result_procedure_result_type; pub mod external_generation_job_result_snapshot_type; @@ -849,6 +850,7 @@ pub mod save_editor_project_layout_v_2_ack_procedure; pub mod save_editor_project_layout_v_2_and_return_procedure; pub mod seed_analytics_date_dimensions_reducer; pub mod set_editor_showcase_asset_like_for_viewer_and_return_procedure; +pub mod set_external_generation_job_provider_checkpoint_and_return_procedure; pub mod settle_llm_router_quota_and_return_procedure; pub mod square_hole_agent_message_row_type; pub mod square_hole_agent_message_table; @@ -1337,6 +1339,7 @@ pub use external_generation_job_phase_update_failure_kind_type::ExternalGenerati pub use external_generation_job_phase_update_input_type::ExternalGenerationJobPhaseUpdateInput; pub use external_generation_job_phase_update_procedure_result_type::ExternalGenerationJobPhaseUpdateProcedureResult; pub use external_generation_job_procedure_result_type::ExternalGenerationJobProcedureResult; +pub use external_generation_job_provider_checkpoint_input_type::ExternalGenerationJobProviderCheckpointInput; pub use external_generation_job_renew_lease_input_type::ExternalGenerationJobRenewLeaseInput; pub use external_generation_job_result_procedure_result_type::ExternalGenerationJobResultProcedureResult; pub use external_generation_job_result_snapshot_type::ExternalGenerationJobResultSnapshot; @@ -1758,6 +1761,7 @@ pub use save_editor_project_layout_v_2_ack_procedure::save_editor_project_layout pub use save_editor_project_layout_v_2_and_return_procedure::save_editor_project_layout_v_2_and_return; pub use seed_analytics_date_dimensions_reducer::seed_analytics_date_dimensions; pub use set_editor_showcase_asset_like_for_viewer_and_return_procedure::set_editor_showcase_asset_like_for_viewer_and_return; +pub use set_external_generation_job_provider_checkpoint_and_return_procedure::set_external_generation_job_provider_checkpoint_and_return; pub use settle_llm_router_quota_and_return_procedure::settle_llm_router_quota_and_return; pub use square_hole_agent_message_row_type::SquareHoleAgentMessageRow; pub use square_hole_agent_message_table::*; diff --git a/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_provider_checkpoint_input_type.rs b/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_provider_checkpoint_input_type.rs new file mode 100644 index 000000000..17ce6b40e --- /dev/null +++ b/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_provider_checkpoint_input_type.rs @@ -0,0 +1,19 @@ +// THIS FILE IS AUTOMATICALLY GENERATED BY SPACETIMEDB. EDITS TO THIS FILE +// WILL NOT BE SAVED. MODIFY TABLES IN YOUR MODULE SOURCE CODE INSTEAD. + +#![allow(unused, clippy::all)] +use spacetimedb_sdk::__codegen::{self as __sdk, __lib, __sats, __ws}; + +#[derive(__lib::ser::Serialize, __lib::de::Deserialize, Clone, PartialEq, Debug)] +#[sats(crate = __lib)] +pub struct ExternalGenerationJobProviderCheckpointInput { + pub job_id: String, + pub worker_id: String, + pub lease_token: String, + pub provider_kind: String, + pub provider_task_id: String, +} + +impl __sdk::InModule for ExternalGenerationJobProviderCheckpointInput { + type Module = super::RemoteModule; +} diff --git a/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_snapshot_type.rs b/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_snapshot_type.rs index 8778d141c..5d34d2bf7 100644 --- a/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_snapshot_type.rs +++ b/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_snapshot_type.rs @@ -32,6 +32,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, } impl __sdk::InModule for ExternalGenerationJobSnapshot { diff --git a/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_type.rs b/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_type.rs index f63d494a4..bf47075f3 100644 --- a/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_type.rs +++ b/server-rs/crates/spacetime-client/src/module_bindings/external_generation_job_type.rs @@ -32,6 +32,8 @@ pub struct ExternalGenerationJob { pub refund_ledger_id: Option, pub notification_acknowledged_at: Option<__sdk::Timestamp>, pub phase: Option, + pub provider_kind: Option, + pub provider_task_id: Option, } impl __sdk::InModule for ExternalGenerationJob { @@ -69,6 +71,8 @@ pub struct ExternalGenerationJobCols { pub notification_acknowledged_at: __sdk::__query_builder::Col>, pub phase: __sdk::__query_builder::Col>, + pub provider_kind: __sdk::__query_builder::Col>, + pub provider_task_id: __sdk::__query_builder::Col>, } impl __sdk::__query_builder::HasCols for ExternalGenerationJob { @@ -109,6 +113,8 @@ impl __sdk::__query_builder::HasCols for ExternalGenerationJob { "notification_acknowledged_at", ), phase: __sdk::__query_builder::Col::new(table_name, "phase"), + provider_kind: __sdk::__query_builder::Col::new(table_name, "provider_kind"), + provider_task_id: __sdk::__query_builder::Col::new(table_name, "provider_task_id"), } } } diff --git a/server-rs/crates/spacetime-client/src/module_bindings/set_external_generation_job_provider_checkpoint_and_return_procedure.rs b/server-rs/crates/spacetime-client/src/module_bindings/set_external_generation_job_provider_checkpoint_and_return_procedure.rs new file mode 100644 index 000000000..47bef2d83 --- /dev/null +++ b/server-rs/crates/spacetime-client/src/module_bindings/set_external_generation_job_provider_checkpoint_and_return_procedure.rs @@ -0,0 +1,62 @@ +// THIS FILE IS AUTOMATICALLY GENERATED BY SPACETIMEDB. EDITS TO THIS FILE +// WILL NOT BE SAVED. MODIFY TABLES IN YOUR MODULE SOURCE CODE INSTEAD. + +#![allow(unused, clippy::all)] +use spacetimedb_sdk::__codegen::{self as __sdk, __lib, __sats, __ws}; + +use super::external_generation_job_procedure_result_type::ExternalGenerationJobProcedureResult; +use super::external_generation_job_provider_checkpoint_input_type::ExternalGenerationJobProviderCheckpointInput; + +#[derive(__lib::ser::Serialize, __lib::de::Deserialize, Clone, PartialEq, Debug)] +#[sats(crate = __lib)] +struct SetExternalGenerationJobProviderCheckpointAndReturnArgs { + pub input: ExternalGenerationJobProviderCheckpointInput, +} + +impl __sdk::InModule for SetExternalGenerationJobProviderCheckpointAndReturnArgs { + type Module = super::RemoteModule; +} + +#[allow(non_camel_case_types)] +/// Extension trait for access to the procedure `set_external_generation_job_provider_checkpoint_and_return`. +/// +/// Implemented for [`super::RemoteProcedures`]. +pub trait set_external_generation_job_provider_checkpoint_and_return { + fn set_external_generation_job_provider_checkpoint_and_return( + &self, + input: ExternalGenerationJobProviderCheckpointInput, + ) { + self.set_external_generation_job_provider_checkpoint_and_return_then(input, |_, _| {}); + } + + fn set_external_generation_job_provider_checkpoint_and_return_then( + &self, + input: ExternalGenerationJobProviderCheckpointInput, + + __callback: impl FnOnce( + &super::ProcedureEventContext, + Result, + ) + Send + + 'static, + ); +} + +impl set_external_generation_job_provider_checkpoint_and_return for super::RemoteProcedures { + fn set_external_generation_job_provider_checkpoint_and_return_then( + &self, + input: ExternalGenerationJobProviderCheckpointInput, + + __callback: impl FnOnce( + &super::ProcedureEventContext, + Result, + ) + Send + + 'static, + ) { + self.imp + .invoke_procedure_with_callback::<_, ExternalGenerationJobProcedureResult>( + "set_external_generation_job_provider_checkpoint_and_return", + SetExternalGenerationJobProviderCheckpointAndReturnArgs { input }, + __callback, + ); + } +}