diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml index 4b5c97fe72..e65bae65c7 100644 --- a/.github/workflows/publish.yaml +++ b/.github/workflows/publish.yaml @@ -399,6 +399,8 @@ jobs: node-version: "22" registry-url: "https://registry.npmjs.org" cache: pnpm + - name: Install OIDC-capable npm + run: npm install --global npm@11.16.0 - run: pnpm install --frozen-lockfile - uses: ./.github/actions/docker-setup with: @@ -568,7 +570,6 @@ jobs: - name: Publish npm packages env: - NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }} SKIP_WASM_BUILD: "1" run: | pnpm --filter=publish exec tsx src/ci/bin.ts publish-npm \ diff --git a/rivetkit-rust/engine/artifacts/errors/queue.message_identity_mismatch.json b/rivetkit-rust/engine/artifacts/errors/queue.message_identity_mismatch.json new file mode 100644 index 0000000000..8024dad60e --- /dev/null +++ b/rivetkit-rust/engine/artifacts/errors/queue.message_identity_mismatch.json @@ -0,0 +1,5 @@ +{ + "code": "message_identity_mismatch", + "group": "queue", + "message": "Queue message identity does not match" +} \ No newline at end of file diff --git a/rivetkit-rust/packages/actor-persist/src/versioned.rs b/rivetkit-rust/packages/actor-persist/src/versioned.rs index 8b59de1a4e..d64bd2d9e0 100644 --- a/rivetkit-rust/packages/actor-persist/src/versioned.rs +++ b/rivetkit-rust/packages/actor-persist/src/versioned.rs @@ -491,3 +491,47 @@ impl OwnedVersionedData for LastPushedAlarm { Vec:: Result>::new() } } + +/// Logical run-handler deadline stored separately from the shared physical +/// alarm. This uses its own versioned type so its persisted meaning cannot be +/// conflated with `LastPushedAlarm` even though both currently encode an +/// optional millisecond timestamp. +pub enum RunWakeAt { + V1(Option), +} + +impl OwnedVersionedData for RunWakeAt { + type Latest = Option; + + fn wrap_latest(latest: Self::Latest) -> Self { + Self::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + Self::V1(data) => Ok(data), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + 1 => Ok(Self::V1(serde_bare::from_slice(payload)?)), + _ => bail!("invalid run wake deadline version: {version}"), + } + } + + fn serialize_version(self, version: u16) -> Result> { + match (self, version) { + (Self::V1(data), 1) => serde_bare::to_vec(&data).map_err(Into::into), + (_, version) => bail!("unexpected run wake deadline version: {version}"), + } + } + + fn deserialize_converters() -> Vec Result> { + Vec:: Result>::new() + } + + fn serialize_converters() -> Vec Result> { + Vec:: Result>::new() + } +} diff --git a/rivetkit-rust/packages/actor-persist/tests/versioned.rs b/rivetkit-rust/packages/actor-persist/tests/versioned.rs index c8750590a1..f9a0e4bee4 100644 --- a/rivetkit-rust/packages/actor-persist/tests/versioned.rs +++ b/rivetkit-rust/packages/actor-persist/tests/versioned.rs @@ -27,3 +27,15 @@ fn actor_decodes_legacy_raw_v4_when_current_v4_accepts_bytes() { assert_eq!(decoded.scheduled_events[0].args, Some(vec![0x01, 0x99])); } + +#[test] +fn run_wake_deadline_round_trips_with_embedded_version() { + let encoded = versioned::RunWakeAt::wrap_latest(Some(1_725_000_000_123)) + .serialize_with_embedded_version(1) + .expect("encode run wake deadline"); + assert_eq!(&encoded[..2], &[1, 0]); + + let decoded = versioned::RunWakeAt::deserialize_with_embedded_version(&encoded) + .expect("decode run wake deadline"); + assert_eq!(decoded, Some(1_725_000_000_123)); +} diff --git a/rivetkit-rust/packages/rivetkit-core/schemas/testing/workflow-fixture/v1.bare b/rivetkit-rust/packages/rivetkit-core/schemas/testing/workflow-fixture/v1.bare new file mode 100644 index 0000000000..9bf52c0f5c --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/schemas/testing/workflow-fixture/v1.bare @@ -0,0 +1,83 @@ +# Test-only logical actor SQLite fixture. +# +# The Rust definitions in `src/testing/workflow_fixture.rs` intentionally mirror +# this schema. The outer payload is encoded with vbare's embedded version header. + +type FixtureMetadata struct { + fixtureName: str + sourceRivetkitVersion: str + sourceWorkflowVersion: str + sourceRevision: str + actorId: str + registryKey: str + internalSchemaVersion: i64 + fakeClockSeed: u64 + generatedIdSeed: u64 +} + +type RuntimeRow struct { + lastPushedAlarm: optional + inspectorToken: optional + queueNextId: i64 +} + +type MetaRow struct { + key: str + value: data +} + +type ActorRow struct { + hasInitialized: i64 + input: optional +} + +type WorkflowRow struct { + key: data + value: data +} + +type QueueRow struct { + id: i64 + name: str + body: data + createdAt: i64 +} + +type ScheduleEventRow struct { + eventId: str + triggerAt: i64 + action: str + args: optional + kind: i64 + cronExpression: optional + timezone: optional + intervalMs: optional + lastStartedAt: optional + maxHistory: i64 +} + +type ScheduleHistoryRow struct { + id: i64 + scheduleId: str + action: str + scheduledAt: i64 + firedAt: i64 + finishedAt: optional + result: i64 + errorGroup: optional + errorCode: optional + errorMessage: optional + errorMetadata: optional +} + +type WorkflowFixture struct { + metadata: FixtureMetadata + metaRows: list + runtime: optional + actor: optional + actorState: optional + workflowRows: list + queueRows: list + scheduleEvents: list + scheduleHistory: list +} diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index fbe40d9740..3c575129ad 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -75,6 +75,7 @@ pub(crate) struct ActorContextInner { pub(super) current_state: RwLock>, pub(super) persisted: RwLock, pub(super) last_pushed_alarm: RwLock>, + pub(super) run_wake_at: RwLock>, pub(super) state_save_interval: Duration, pub(super) state_dirty: AtomicBool, pub(super) state_revision: AtomicU64, @@ -88,7 +89,7 @@ pub(crate) struct ActorContextInner { pub(super) last_save_at: Mutex>, pub(super) pending_save: Mutex>, pub(super) tracked_persist: Mutex>>, - pub(super) save_guard: AsyncMutex<()>, + pub(super) save_guard: Arc>, pub(super) in_flight_state_writes: AtomicUsize, pub(super) state_write_completion: Notify, pub(super) on_state_change_in_flight: AtomicUsize, @@ -112,6 +113,7 @@ pub(crate) struct ActorContextInner { // being moved out of the lock. pub(super) schedule_pending_alarm_writes: Mutex>>, pub(super) schedule_local_alarm_epoch: AtomicU64, + pub(super) schedule_alarm_push_epoch: AtomicU64, pub(super) schedule_alarm_dispatch_enabled: AtomicBool, pub(super) schedule_dirty_since_push: AtomicBool, pub(super) schedule_mutation_lock: AsyncMutex<()>, @@ -297,6 +299,7 @@ impl ActorContext { current_state: RwLock::new(Vec::new()), persisted: RwLock::new(PersistedActor::default()), last_pushed_alarm: RwLock::new(None), + run_wake_at: RwLock::new(None), state_save_interval, state_dirty: AtomicBool::new(false), state_revision: AtomicU64::new(0), @@ -309,7 +312,7 @@ impl ActorContext { last_save_at: Mutex::new(None), pending_save: Mutex::new(None), tracked_persist: Mutex::new(None), - save_guard: AsyncMutex::new(()), + save_guard: Arc::new(AsyncMutex::new(())), in_flight_state_writes: AtomicUsize::new(0), state_write_completion: Notify::new(), on_state_change_in_flight: AtomicUsize::new(0), @@ -326,6 +329,7 @@ impl ActorContext { schedule_local_alarm_task: Mutex::new(None), schedule_pending_alarm_writes: Mutex::new(Vec::new()), schedule_local_alarm_epoch: AtomicU64::new(0), + schedule_alarm_push_epoch: AtomicU64::new(0), schedule_alarm_dispatch_enabled: AtomicBool::new(true), // A fresh actor context has no in-process record of a successful // envoy alarm push yet, so the first sync must always push. @@ -1123,7 +1127,7 @@ impl ActorContext { *self.0.hibernated_connection_liveness_override.write() = Some(pairs.into_iter().collect()); } - fn prepare_state_deltas( + pub(super) fn prepare_state_deltas( &self, deltas: Vec, ) -> Result<(Vec, PendingHibernationChanges)> { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs index 52f8955911..15927e5bcb 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/mod.rs @@ -2,12 +2,16 @@ use std::collections::{BTreeSet, HashMap}; use std::io::Cursor; use anyhow::{Context, Result, bail}; +use rivetkit_actor_persist::versioned as persist_versioned; use crate::actor::connection::{ PersistedConnection, PersistedSubscription, encode_persisted_connection, }; use crate::actor::keys::make_workflow_key; use crate::actor::messages::WorkflowKvWrite; +use crate::actor::persist::{ + decode_latest_with_embedded_version, encode_latest_with_embedded_version, +}; use crate::actor::queue::{PersistedQueueMessage, QueueMetadata}; use crate::actor::state::PersistedActor; use crate::error::KvRuntimeError; @@ -28,6 +32,8 @@ pub(crate) const USER_KV_BATCH_GET_MAX_KEYS: usize = 128; const QUEUE_METADATA_PAGE_ROWS: usize = 128; const QUEUE_MESSAGE_IDS_PER_QUERY: usize = 128; const WORKFLOW_KV_VALUE_LIMIT: usize = 256 * 1024; +pub(crate) const RUN_WAKE_AT_META_KEY: &str = "run_wake_at"; +const RUN_WAKE_AT_VERSION: u16 = 1; /// Depot rejects SQLite commits that dirty more than `MAX_COMMIT_RAW_DIRTY_BYTES` /// (320 pages * 4 KiB = 1.3 MiB) in `engine/packages/depot/src/conveyer/constants.rs`. @@ -72,6 +78,7 @@ where pub(crate) struct InternalActorSnapshot { pub actor: PersistedActor, pub last_pushed_alarm: Option, + pub run_wake_at: Option, } pub(crate) async fn load_actor_snapshot(db: &SqliteDb) -> Result> { @@ -87,6 +94,7 @@ pub(crate) async fn load_actor_snapshot(db: &SqliteDb) -> Result Result, connections: &[PersistedConnection], removed_connections: &[String], @@ -477,6 +486,21 @@ pub(crate) async fn load_queue_messages(db: &SqliteDb) -> Result Result> { + let id = i64::try_from(id).context("queue message id exceeds sqlite integer range")?; + let result = db + .query( + LOAD_QUEUE_MESSAGE_NAME_SQL, + Some(vec![BindParam::Integer(id)]), + ) + .await + .context("load internal queue message name")?; + let Some(row) = result.rows.first() else { + return Ok(None); + }; + Ok(Some(read_text(row, 0, "queue message name")?)) +} + pub(crate) async fn load_queue_messages_matching( db: &SqliteDb, names: Option<&BTreeSet>, @@ -855,14 +879,17 @@ fn build_workflow_kv_statements(writes: &[WorkflowKvWrite]) -> Result Result<()> { let row_count = statements.len(); - let payload_bytes = statements - .iter() - .flat_map(|statement| statement.params.iter().flatten()) - .map(bind_param_payload_len) - .fold(0usize, usize::saturating_add); + let payload_bytes = statement_bind_payload_len(statements); + validate_atomic_state_transaction_budget(row_count, payload_bytes) +} + +pub(crate) fn validate_atomic_state_transaction_budget( + row_count: usize, + payload_bytes: usize, +) -> Result<()> { if row_count > KV_TX_MAX_ROWS || payload_bytes > KV_TX_MAX_PAYLOAD_BYTES { bail!( - "atomic actor state and workflow flush exceeds sqlite transaction budget: {row_count} rows and {payload_bytes} bytes (limits: {} rows and {} bytes)", + "atomic SQLite and actor state transaction exceeds transaction budget: {row_count} rows and {payload_bytes} bytes (limits: {} rows and {} bytes)", KV_TX_MAX_ROWS, KV_TX_MAX_PAYLOAD_BYTES ); @@ -870,7 +897,15 @@ fn validate_atomic_workflow_flush(statements: &[SqliteBatchStatement]) -> Result Ok(()) } -fn bind_param_payload_len(param: &BindParam) -> usize { +pub(crate) fn statement_bind_payload_len(statements: &[SqliteBatchStatement]) -> usize { + statements + .iter() + .flat_map(|statement| statement.params.iter().flatten()) + .map(bind_param_payload_len) + .fold(0usize, usize::saturating_add) +} + +pub(crate) fn bind_param_payload_len(param: &BindParam) -> usize { match param { BindParam::Null => 0, BindParam::Integer(_) | BindParam::Float(_) => 8, @@ -995,6 +1030,42 @@ pub(crate) async fn persist_last_pushed_alarm(db: &SqliteDb, alarm_ts: Option Result> { + let result = db + .query( + LOAD_RUN_WAKE_AT_SQL, + Some(vec![BindParam::Text(RUN_WAKE_AT_META_KEY.to_owned())]), + ) + .await + .context("load internal run wake deadline")?; + let Some(row) = result.rows.first() else { + return Ok(None); + }; + let payload = read_blob(row, 0, "run wake deadline")?; + decode_latest_with_embedded_version::( + &payload, + "run wake deadline", + ) +} + +pub(crate) async fn persist_run_wake_at(db: &SqliteDb, wake_at: Option) -> Result<()> { + let payload = encode_latest_with_embedded_version::( + wake_at, + RUN_WAKE_AT_VERSION, + "run wake deadline", + )?; + db.execute( + UPSERT_RUN_WAKE_AT_SQL, + Some(vec![ + BindParam::Text(RUN_WAKE_AT_META_KEY.to_owned()), + BindParam::Blob(payload), + ]), + ) + .await + .context("persist internal run wake deadline")?; + Ok(()) +} + pub(crate) async fn load_inspector_token(db: &SqliteDb) -> Result> { let result = db .query(LOAD_INSPECTOR_TOKEN_SQL, None) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs index 703fdb5dbe..dad1417c00 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/queries.rs @@ -65,6 +65,7 @@ pub(crate) const HAS_QUEUE_MESSAGES_FOR_NAME_SQL: &str = "SELECT 1 FROM _rivet_queue INDEXED BY _rivet_queue_name_id WHERE name = ? LIMIT 1"; pub(crate) const LOAD_QUEUE_MESSAGE_METADATA_PAGE_SQL: &str = "SELECT id, name FROM _rivet_queue WHERE id > ? ORDER BY id LIMIT ?"; +pub(crate) const LOAD_QUEUE_MESSAGE_NAME_SQL: &str = "SELECT name FROM _rivet_queue WHERE id = ?"; pub(crate) fn load_queue_messages_by_ids_sql(id_count: usize) -> String { let placeholders = std::iter::repeat_n("?", id_count) .collect::>() @@ -86,10 +87,12 @@ pub(crate) const UPSERT_WORKFLOW_KV_SQL: &str = "INSERT INTO _rivet_wf_kv (key, pub(crate) const LOAD_LAST_PUSHED_ALARM_SQL: &str = "SELECT last_pushed_alarm FROM _rivet_runtime WHERE id = 1"; +pub(crate) const LOAD_RUN_WAKE_AT_SQL: &str = "SELECT value FROM _rivet_meta WHERE key = ?"; pub(crate) const LOAD_INSPECTOR_TOKEN_SQL: &str = "SELECT inspector_token FROM _rivet_runtime WHERE id = 1"; pub(crate) const UPSERT_QUEUE_NEXT_ID_SQL: &str = "INSERT INTO _rivet_runtime (id, last_pushed_alarm, inspector_token, queue_next_id) VALUES (1, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET queue_next_id = excluded.queue_next_id"; pub(crate) const UPSERT_LAST_PUSHED_ALARM_SQL: &str = "INSERT INTO _rivet_runtime (id, last_pushed_alarm, inspector_token, queue_next_id) VALUES (1, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET last_pushed_alarm = excluded.last_pushed_alarm"; +pub(crate) const UPSERT_RUN_WAKE_AT_SQL: &str = "INSERT INTO _rivet_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value"; pub(crate) const UPSERT_INSPECTOR_TOKEN_SQL: &str = "INSERT INTO _rivet_runtime (id, last_pushed_alarm, inspector_token, queue_next_id) VALUES (1, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET inspector_token = excluded.inspector_token"; pub(crate) const LOAD_META_TEXT_SQL: &str = "SELECT value FROM _rivet_meta WHERE key = ?"; pub(crate) const UPSERT_META_TEXT_SQL: &str = "INSERT INTO _rivet_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value"; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs index 5f9b1c3372..41c23e44b2 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/internal_storage/schema.rs @@ -10,9 +10,10 @@ const SCHEMA_VERSION_KEY: &str = "schema_version"; // `_rivet_meta` is the bootstrap root created before the numbered migrations. // `schema_version` cannot live in a table created by those migrations, and // `kv_import_state` must survive clearing partially imported runtime tables so -// interrupted imports can be detected and retried. This is not a general- -// purpose runtime KV store; its text accessors are migration bookkeeping only. -// W[bootstrap + import bookkeeping only | point upsert | <100 B | 1-page map] +// interrupted imports can be detected and retried. Fixed core-owned logical +// metadata may also live here when adding a column would break older runtimes' +// ability to open the database. This is not a general-purpose runtime KV store. +// W[bootstrap + core metadata only | point upsert | <100 B | 1-page map] pub(crate) const CREATE_META_TABLE: &str = r#" CREATE TABLE IF NOT EXISTS _rivet_meta ( key TEXT PRIMARY KEY, @@ -135,6 +136,9 @@ CREATE TABLE _rivet_queue ( CREATE INDEX _rivet_queue_name_id ON _rivet_queue (name, id) "#, + // RivetKit core owns this table's creation and migrations. External workflow + // packages are format-compatible clients and must never create or migrate it; + // any format change requires bidirectional old/new compatibility fixtures. // W[per workflow step flush | keyed upsert + range delete | values <=256 KiB | verbatim fdb-tuple keys in one clustered tree] r#" CREATE TABLE _rivet_wf_kv ( @@ -148,7 +152,7 @@ CREATE TABLE _rivet_user_kv ( key BLOB PRIMARY KEY, value BLOB NOT NULL ) STRICT, WITHOUT ROWID -"#, + "#, ]]; pub(crate) async fn ensure_internal_schema(db: &SqliteDb) -> Result<()> { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs index 644f3df3cc..7e36f9d09d 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs @@ -348,6 +348,11 @@ pub enum ActorEvent { entry_id: Option, reply: Reply>>, }, + /// The persisted logical run deadline became due. Foreign runtimes restart + /// their run callback without restarting the actor runtime itself. + RunWake { + reply: Reply<()>, + }, } impl ActorEvent { @@ -378,6 +383,7 @@ impl ActorEvent { Self::Destroy { .. } => "destroy", Self::WorkflowHistoryRequested { .. } => "workflow_history_requested", Self::WorkflowReplayRequested { .. } => "workflow_replay_requested", + Self::RunWake { .. } => "run_wake", } } } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/mod.rs index 6f0ae6dbdc..837629ac4c 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/mod.rs @@ -48,7 +48,7 @@ pub use sqlite::{ BindParam, ColumnValue, ExecResult, ExecuteResult, QueryResult, SqliteBackend, SqliteBatchStatement, SqliteDb, SqliteTransaction, }; -pub use state::RequestSaveOpts; +pub use state::{ActorStateTransaction, RequestSaveOpts}; pub use task::{ ActionDispatchResult, ActorTask, DispatchCommand, HttpDispatchResult, LifecycleCommand, LifecycleEvent, LifecycleState, diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs index fd6203c48f..128237a8a4 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/queue.rs @@ -195,6 +195,19 @@ struct QueueCompleteNotConfigured { name: String, } +#[derive(RivetError, Serialize, Deserialize)] +#[error( + "queue", + "message_identity_mismatch", + "Queue message identity does not match", + "Queue message {message_id} is named '{actual_name}', not '{expected_name}'." +)] +struct QueueMessageIdentityMismatch { + message_id: u64, + expected_name: String, + actual_name: String, +} + #[derive(RivetError)] #[error("actor", "aborted", "Actor aborted")] struct QueueActorAborted; @@ -487,6 +500,49 @@ impl ActorContext { } } + /// Completes a persisted message by durable identity. Missing or already + /// completed IDs are idempotent; an expected-name mismatch leaves the row + /// untouched. + pub async fn complete_persisted_message( + &self, + message_id: u64, + expected_name: &str, + response: Option>, + ) -> Result { + self.ensure_initialized().await?; + let _receive_guard = self.0.queue_receive_lock.lock().await; + let Some(_) = self + .verify_persisted_message_identity_unlocked(message_id, expected_name) + .await? + else { + return Ok(false); + }; + self.complete_message_by_id_unlocked(message_id, response) + .await?; + Ok(true) + } + + async fn verify_persisted_message_identity_unlocked( + &self, + message_id: u64, + expected_name: &str, + ) -> Result> { + let Some(actual_name) = + internal_storage::load_queue_message_name(self.sql(), message_id).await? + else { + return Ok(None); + }; + if actual_name != expected_name { + return Err(QueueMessageIdentityMismatch { + message_id, + expected_name: expected_name.to_owned(), + actual_name: actual_name.clone(), + } + .build()); + } + Ok(Some(actual_name)) + } + pub fn try_next(&self, opts: QueueTryNextOpts) -> Result> { let mut messages = self.try_next_batch(QueueTryNextBatchOpts { names: opts.names, @@ -675,6 +731,16 @@ impl ActorContext { &self, message_id: u64, response: Option>, + ) -> Result<()> { + let _receive_guard = self.0.queue_receive_lock.lock().await; + self.complete_message_by_id_unlocked(message_id, response) + .await + } + + async fn complete_message_by_id_unlocked( + &self, + message_id: u64, + response: Option>, ) -> Result<()> { self.remove_messages(vec![message_id]).await?; if let Some(waiter) = self.remove_completion_waiter(message_id).await { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs index d55eaf1a3c..39d7b0329e 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/schedule.rs @@ -38,6 +38,14 @@ pub(crate) const GLOBAL_HISTORY_PRUNE_INTERVAL: usize = 100; pub(crate) const GLOBAL_HISTORY_RETAINED_ROWS: i64 = MAX_ACTOR_HISTORY - GLOBAL_HISTORY_PRUNE_INTERVAL as i64; +fn min_deadline(left: Option, right: Option) -> Option { + match (left, right) { + (Some(left), Some(right)) => Some(left.min(right)), + (Some(value), None) | (None, Some(value)) => Some(value), + (None, None) => None, + } +} + pub(super) type InternalKeepAwakeCallback = Arc>) -> BoxFuture<'static, Result<()>> + Send + Sync>; pub(super) type LocalAlarmCallback = Arc BoxFuture<'static, ()> + Send + Sync>; @@ -807,6 +815,29 @@ impl ActorContext { .store(true, Ordering::SeqCst); } + /// Sets the durable logical deadline that should restart the foreign run + /// handler. This deadline shares one physical alarm with scheduled actions, + /// but remains independently addressable and clearable. + pub async fn set_run_wake_at(&self, wake_at: Option) -> Result<()> { + let _mutation = self.0.schedule_mutation_lock.lock().await; + self.persist_run_wake_at(wake_at).await?; + self.mark_schedule_dirty(); + self.sync_alarm().await + } + + pub(crate) async fn consume_due_run_wake(&self) -> Result> { + let _mutation = self.0.schedule_mutation_lock.lock().await; + let Some(wake_at) = self.run_wake_at() else { + return Ok(None); + }; + if wake_at > self.schedule_now_timestamp_ms() { + return Ok(None); + } + self.persist_run_wake_at(None).await?; + self.mark_schedule_dirty(); + Ok(Some(wake_at)) + } + async fn next_schedule_timestamp(&self, future_only: bool) -> Result> { let (sql, params) = if future_only { ( @@ -840,12 +871,18 @@ impl ActorContext { { anyhow::bail!("injected schedule alarm sync failure"); } - let next_alarm = self.next_schedule_timestamp(false).await?; + let next_alarm = min_deadline( + self.next_schedule_timestamp(false).await?, + self.run_wake_at(), + ); self.sync_alarm_timestamp(next_alarm) } async fn sync_future_alarm(&self) -> Result<()> { - let next_alarm = self.next_schedule_timestamp(true).await?; + let next_alarm = min_deadline( + self.next_schedule_timestamp(true).await?, + self.run_wake_at(), + ); self.sync_alarm_timestamp(next_alarm) } @@ -967,6 +1004,11 @@ impl ActorContext { timestamp_ms: Option, generation: Option, ) { + let push_epoch = self + .0 + .schedule_alarm_push_epoch + .fetch_add(1, Ordering::SeqCst) + .wrapping_add(1); let (ack_tx, ack_rx) = oneshot::channel(); envoy_handle.set_alarm_with_ack( self.actor_id().to_owned(), @@ -981,7 +1023,11 @@ impl ActorContext { handle.spawn( async move { let _ = ack_rx.await; - if let Err(error) = state_ctx.persist_last_pushed_alarm(timestamp_ms).await { + let is_current = state_ctx.0.schedule_alarm_push_epoch.load(Ordering::SeqCst) + == push_epoch && state_ctx.last_pushed_alarm() == timestamp_ms; + if is_current + && let Err(error) = state_ctx.persist_last_pushed_alarm(timestamp_ms).await + { tracing::error!( ?error, ?timestamp_ms, diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs index 332a0c3a7f..f12270b2ce 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs @@ -226,7 +226,8 @@ impl ActorContext { } } - pub(crate) fn run_handler_active(&self) -> bool { + #[doc(hidden)] + pub fn run_handler_active(&self) -> bool { self.0.sleep.run_handler_active_count.load(Ordering::SeqCst) > 0 } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/state.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/state.rs index ceeb4c77e6..0bccf063dc 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/state.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/state.rs @@ -9,7 +9,7 @@ use anyhow::{Context, Result}; use rivetkit_actor_persist::{generated::v4 as persist_v4, versioned as persist_versioned}; #[cfg(not(feature = "wasm-runtime"))] use tokio::runtime::Handle; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{Mutex as AsyncMutex, OwnedMutexGuard, mpsc, oneshot}; use tokio::task::JoinHandle; #[cfg(test)] use tokio::time::timeout; @@ -28,6 +28,7 @@ use crate::actor::task_types::StateMutationReason; use crate::error::ActorRuntime; #[cfg(feature = "wasm-runtime")] use crate::runtime::RuntimeSpawner; +use crate::sqlite::{BindParam, ExecuteResult, QueryResult, SqliteTransaction}; use crate::types::SaveStateOpts; #[cfg(test)] @@ -84,6 +85,136 @@ pub struct OnStateChangeGuard { ctx: Option, } +/// A SQLite transaction that owns the actor state-save exclusion until it is +/// committed or rolled back. +#[derive(Clone)] +pub struct ActorStateTransaction { + owner: Arc>, +} + +struct ActorStateTransactionOwner { + ctx: ActorContext, + transaction: SqliteTransaction, + save_guard: Option>, + write_guard: Option, + finalized: bool, + committed: bool, + save_request_revision: u64, + statement_count: usize, + bind_payload_bytes: usize, +} + +impl Drop for ActorStateTransactionOwner { + fn drop(&mut self) { + if !self.finalized { + let transaction = self.transaction.clone(); + #[cfg(not(feature = "wasm-runtime"))] + if let Ok(handle) = Handle::try_current() { + handle.spawn(async move { + if let Err(error) = transaction.rollback().await { + tracing::debug!(?error, "dropped actor state transaction rollback failed"); + } + }); + } + #[cfg(feature = "wasm-runtime")] + RuntimeSpawner::spawn(async move { + if let Err(error) = transaction.rollback().await { + tracing::debug!(?error, "dropped actor state transaction rollback failed"); + } + }); + } + if !self.committed { + self.ctx.schedule_save(None); + } + } +} + +impl ActorStateTransaction { + pub async fn exec(&self, _sql: impl Into) -> Result { + let owner = self.owner.lock().await; + if owner.finalized { + return Err(anyhow::anyhow!( + "actor state transaction is already finalized" + )); + } + Err(anyhow::anyhow!( + "actor state transactions only support single-statement execute calls" + )) + } + + pub async fn execute( + &self, + sql: impl Into, + params: Option>, + ) -> Result { + let mut owner = self.owner.lock().await; + if owner.finalized { + return Err(anyhow::anyhow!( + "actor state transaction is already finalized" + )); + } + let sql = sql.into(); + let payload_bytes = match params.as_deref() { + Some(params) => params + .iter() + .map(internal_storage::bind_param_payload_len) + .fold(0usize, usize::saturating_add), + None => sql.len(), + }; + let result = owner.transaction.execute(sql, params).await?; + owner.statement_count = owner.statement_count.saturating_add(1); + owner.bind_payload_bytes = owner.bind_payload_bytes.saturating_add(payload_bytes); + Ok(result) + } + + pub async fn commit(&self, deltas: Vec) -> Result<()> { + let mut owner = self.owner.lock().await; + if owner.finalized { + return Err(anyhow::anyhow!( + "actor state transaction is already finalized" + )); + } + + let save_request_revision = owner.save_request_revision; + let transaction = owner.transaction.clone(); + let statement_count = owner.statement_count; + let bind_payload_bytes = owner.bind_payload_bytes; + let result = owner + .ctx + .commit_state_transaction( + &transaction, + deltas, + save_request_revision, + statement_count, + bind_payload_bytes, + ) + .await; + owner.finalized = true; + owner.committed = result.is_ok(); + owner.write_guard.take(); + owner.save_guard.take(); + if result.is_err() { + owner.ctx.schedule_save(None); + } + result + } + + pub async fn rollback(&self) -> Result<()> { + let mut owner = self.owner.lock().await; + if owner.finalized { + return Err(anyhow::anyhow!( + "actor state transaction is already finalized" + )); + } + let result = owner.transaction.rollback().await; + owner.finalized = true; + owner.write_guard.take(); + owner.save_guard.take(); + owner.ctx.schedule_save(None); + result + } +} + impl OnStateChangeGuard { fn new(ctx: ActorContext) -> Self { ctx.on_state_change_started(); @@ -100,6 +231,38 @@ impl Drop for OnStateChangeGuard { } impl ActorContext { + pub async fn begin_state_transaction( + &self, + timeout: Option, + ) -> Result { + self.clear_pending_save(); + let save_guard = Arc::clone(&self.0.save_guard).lock_owned().await; + self.wait_for_in_flight_writes().await; + let transaction = match self.sql().begin_transaction(timeout).await { + Ok(transaction) => transaction, + Err(error) => { + drop(save_guard); + self.schedule_save(None); + return Err(error); + } + }; + let write_guard = self.begin_write(); + let save_request_revision = self.save_request_revision(); + Ok(ActorStateTransaction { + owner: Arc::new(AsyncMutex::new(ActorStateTransactionOwner { + ctx: self.clone(), + transaction, + save_guard: Some(save_guard), + write_guard: Some(write_guard), + finalized: false, + committed: false, + save_request_revision, + statement_count: 0, + bind_payload_bytes: 0, + })), + }) + } + pub fn state(&self) -> Vec { self.0.current_state.read().clone() } @@ -479,6 +642,99 @@ impl ActorContext { Ok(()) } + async fn commit_state_transaction( + &self, + transaction: &SqliteTransaction, + deltas: Vec, + save_request_revision: u64, + statement_count: usize, + bind_payload_bytes: usize, + ) -> Result<()> { + let (deltas, pending_hibernation_changes) = match self.prepare_state_deltas(deltas) { + Ok(prepared) => prepared, + Err(error) => { + let _ = transaction.rollback().await; + return Err(error); + } + }; + let commit_result = async { + let revision = self.0.state_revision.load(Ordering::SeqCst); + let mut persisted = self.persisted(); + let mut next_state = None; + let mut actor_to_persist = None; + let mut connections_to_persist: Vec = Vec::new(); + let mut connections_to_delete = Vec::new(); + + for delta in deltas { + match delta { + StateDelta::ActorState(bytes) => { + next_state = Some(bytes.clone()); + persisted.state = bytes; + } + StateDelta::ConnHibernation { conn: _, bytes } => { + connections_to_persist.push( + decode_persisted_connection(&bytes) + .context("decode hibernatable connection state delta")?, + ); + } + StateDelta::ConnHibernationRemoved(conn) => { + connections_to_delete.push(conn); + } + } + } + + if next_state.is_some() { + actor_to_persist = Some(persisted); + } + let statements = internal_storage::build_actor_core_and_connection_statements( + actor_to_persist.as_ref(), + &connections_to_persist, + &connections_to_delete, + )?; + internal_storage::validate_atomic_state_transaction_budget( + statement_count.saturating_add(statements.len()), + bind_payload_bytes + .saturating_add(internal_storage::statement_bind_payload_len(&statements)), + )?; + for statement in statements { + transaction + .execute(statement.sql, statement.params) + .await + .context("persist actor state inside sqlite transaction")?; + } + transaction + .commit() + .await + .context("commit sqlite transaction with actor state")?; + + if let Some(state) = next_state { + self.0.persisted.write().state = state.clone(); + *self.0.current_state.write() = state; + } + for connection in &connections_to_persist { + if let Some(handle) = self.connection(&connection.id) { + handle.set_state_initial(connection.state.clone()); + } + } + *self.0.last_save_at.lock() = Some(StdInstant::now()); + if self.0.state_revision.load(Ordering::SeqCst) == revision { + self.0.state_dirty.store(false, Ordering::SeqCst); + } + self.mark_save_request_completed(save_request_revision); + self.finish_save_request(save_request_revision); + self.record_state_updated(); + Ok(()) + } + .await; + + if let Err(error) = commit_result { + self.restore_pending_hibernation_changes(pending_hibernation_changes); + let _ = transaction.rollback().await; + return Err(error); + } + Ok(()) + } + pub(crate) async fn wait_for_pending_writes(&self) { loop { if let Some(handle) = self.take_tracked_persist() { @@ -598,6 +854,22 @@ impl ActorContext { Ok(()) } + pub(crate) fn load_run_wake_at(&self, wake_at: Option) { + *self.0.run_wake_at.write() = wake_at; + } + + pub fn run_wake_at(&self) -> Option { + *self.0.run_wake_at.read() + } + + pub(crate) async fn persist_run_wake_at(&self, wake_at: Option) -> Result<()> { + internal_storage::persist_run_wake_at(self.sql(), wake_at) + .await + .context("persist run wake deadline to sqlite")?; + self.load_run_wake_at(wake_at); + Ok(()) + } + pub(crate) fn set_initial_state(&self, state: Vec) { *self.0.current_state.write() = state.clone(); self.0.persisted.write().state = state; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index cbb8737b63..274f924e16 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -317,6 +317,7 @@ struct SleepGraceState { struct PersistedStartup { actor: PersistedActor, last_pushed_alarm: Option, + run_wake_at: Option, } struct PendingLifecycleReply { @@ -1178,6 +1179,7 @@ impl ActorTask { let core_init_result: Result<()> = async { self.ctx.load_persisted_actor(persisted.actor); self.ctx.load_last_pushed_alarm(persisted.last_pushed_alarm); + self.ctx.load_run_wake_at(persisted.run_wake_at); // New manual-startup runtimes must not persist initialization until the // runtime startup_ready handshake completes. The runtime preamble owns // initial state creation. @@ -1277,6 +1279,7 @@ impl ActorTask { return Ok(PersistedStartup { actor: snapshot.actor, last_pushed_alarm: snapshot.last_pushed_alarm, + run_wake_at: snapshot.run_wake_at, }); } Ok(PersistedStartup { @@ -1285,6 +1288,7 @@ impl ActorTask { ..PersistedActor::default() }, last_pushed_alarm: None, + run_wake_at: None, }) } @@ -1410,7 +1414,50 @@ impl ActorTask { return Ok(()); } - self.ctx.drain_overdue_scheduled_events().await + let due_run_wake = self.ctx.consume_due_run_wake().await?; + if let Err(error) = self.ctx.drain_overdue_scheduled_events().await { + if self.lifecycle != LifecycleState::DestroyGrace + && let Some(wake_at) = due_run_wake + && let Err(restore_error) = self.ctx.set_run_wake_at(Some(wake_at)).await + { + tracing::error!( + ?restore_error, + wake_at, + "failed to restore run wake after schedule alarm dispatch failed", + ); + } + return Err(error); + } + // Destroy is terminal. Consume the logical deadline so it cannot keep a + // past physical alarm armed, but never restart the foreign run handler + // after its destroy cleanup has started. + if self.lifecycle == LifecycleState::DestroyGrace { + return Ok(()); + } + if let Some(wake_at) = due_run_wake { + let (reply_tx, reply_rx) = oneshot::channel(); + if let Err(error) = self.ctx.try_send_actor_event( + ActorEvent::RunWake { + reply: Reply::from(reply_tx), + }, + "run_wake", + ) { + self.ctx.set_run_wake_at(Some(wake_at)).await?; + return Err(error).context("dispatch due run wake"); + } + let restart_result = match reply_rx.await { + Ok(result) => result, + Err(error) => { + self.ctx.set_run_wake_at(Some(wake_at)).await?; + return Err(error).context("receive due run wake restart reply"); + } + }; + if let Err(error) = restart_result { + self.ctx.set_run_wake_at(Some(wake_at)).await?; + return Err(error).context("restart run handler for due wake"); + } + } + Ok(()) } fn handle_run_handle_outcome( diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 8212027a1b..dbe4b7c71f 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -139,17 +139,17 @@ pub use actor::sqlite::{ BindParam, ColumnValue, ExecResult, ExecuteResult, QueryResult, SqliteBackend, SqliteBatchStatement, SqliteDb, SqliteTransaction, }; -pub use actor::state::RequestSaveOpts; +pub use actor::state::{ActorStateTransaction, RequestSaveOpts}; pub use actor::task::{ ActionDispatchResult, ActorTask, DispatchCommand, HttpDispatchResult, LifecycleCommand, LifecycleEvent, LifecycleState, }; -pub use rivet_envoy_client::config::ResponseChunk; pub use actor::task_types::ShutdownKind; pub use actor::work_registry::{ActorWorkKind, ActorWorkPolicy}; pub use error::ActorLifecycle; pub use inspector::{Inspector, InspectorSnapshot}; pub use registry::{CoreRegistry, EngineSpawnMode, ServeConfig}; +pub use rivet_envoy_client::config::ResponseChunk; pub use runtime::{RuntimeBoxFuture, RuntimeSpawner, boxed_runtime_future}; pub use serverless::{CoreServerlessRuntime, ServerlessRequest, ServerlessResponse}; pub use types::{ diff --git a/rivetkit-rust/packages/rivetkit-core/src/testing.rs b/rivetkit-rust/packages/rivetkit-core/src/testing.rs index 437ddbc9ca..f6102d8c84 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/testing.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/testing.rs @@ -29,6 +29,8 @@ use crate::actor::kv::LegacyActorKv; use crate::sqlite::SqliteDb; use crate::types::ActorKey; +pub mod workflow_fixture; + /// Reusable in-memory SQLite store for constructing fully configured contexts. /// Contexts created from the same harness observe the same database. #[derive(Clone)] diff --git a/rivetkit-rust/packages/rivetkit-core/src/testing/workflow_fixture.rs b/rivetkit-rust/packages/rivetkit-core/src/testing/workflow_fixture.rs new file mode 100644 index 0000000000..14a2d3b2d1 --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/testing/workflow_fixture.rs @@ -0,0 +1,524 @@ +//! Versioned logical SQLite fixtures for workflow upgrade tests. +//! +//! This module is compiled only for tests or with `test-support`. It deliberately +//! exposes no production actor-database import path: callers can dump or restore +//! only the fixed Rivet-owned tables represented by [`WorkflowFixture`]. + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +use vbare::OwnedVersionedData; + +use crate::ActorContext; +use crate::actor::keys::WORKFLOW_STORAGE_PREFIX; +use crate::sqlite::{BindParam, ColumnValue, SqliteBatchStatement}; + +const FIXTURE_VERSION: u16 = 1; +const SCHEMA_VERSION_META_KEY: &str = "schema_version"; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureMetadata { + pub fixture_name: String, + pub source_rivetkit_version: String, + pub source_workflow_version: String, + pub source_revision: String, + pub actor_id: String, + pub registry_key: String, + pub internal_schema_version: i64, + pub fake_clock_seed: u64, + pub generated_id_seed: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureRuntimeRow { + pub last_pushed_alarm: Option, + pub inspector_token: Option, + pub queue_next_id: i64, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureMetaRow { + pub key: String, + pub value: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureActorRow { + pub has_initialized: i64, + pub input: Option>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureWorkflowRow { + pub key: Vec, + pub value: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureQueueRow { + pub id: i64, + pub name: String, + pub body: Vec, + pub created_at: i64, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureScheduleEventRow { + pub event_id: String, + pub trigger_at: i64, + pub action: String, + pub args: Option>, + pub kind: i64, + pub cron_expression: Option, + pub timezone: Option, + pub interval_ms: Option, + pub last_started_at: Option, + pub max_history: i64, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixtureScheduleHistoryRow { + pub id: i64, + pub schedule_id: String, + pub action: String, + pub scheduled_at: i64, + pub fired_at: i64, + pub finished_at: Option, + pub result: i64, + pub error_group: Option, + pub error_code: Option, + pub error_message: Option, + pub error_metadata: Option>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkflowFixture { + pub metadata: WorkflowFixtureMetadata, + pub meta_rows: Vec, + pub runtime: Option, + pub actor: Option, + pub actor_state: Option>, + pub workflow_rows: Vec, + pub queue_rows: Vec, + pub schedule_events: Vec, + pub schedule_history: Vec, +} + +enum VersionedWorkflowFixture { + V1(WorkflowFixture), +} + +impl OwnedVersionedData for VersionedWorkflowFixture { + type Latest = WorkflowFixture; + + fn wrap_latest(latest: Self::Latest) -> Self { + Self::V1(latest) + } + + fn unwrap_latest(self) -> Result { + match self { + Self::V1(fixture) => Ok(fixture), + } + } + + fn deserialize_version(payload: &[u8], version: u16) -> Result { + match version { + FIXTURE_VERSION => Ok(Self::V1(serde_bare::from_slice(payload)?)), + _ => bail!("unsupported workflow fixture version {version}"), + } + } + + fn serialize_version(self, version: u16) -> Result> { + match (self, version) { + (Self::V1(fixture), FIXTURE_VERSION) => { + serde_bare::to_vec(&fixture).map_err(Into::into) + } + (_, version) => bail!("unsupported workflow fixture version {version}"), + } + } +} + +impl WorkflowFixture { + pub fn encode(&self) -> Result> { + VersionedWorkflowFixture::wrap_latest(self.clone()) + .serialize_with_embedded_version(FIXTURE_VERSION) + } + + pub fn decode(bytes: &[u8]) -> Result { + VersionedWorkflowFixture::deserialize_with_embedded_version(bytes) + } +} + +/// Dumps the fixed set of Rivet-owned logical rows required to resume a +/// workflow. Every query has an explicit ordering so equivalent databases +/// produce byte-identical fixtures regardless of SQLite page layout. +pub async fn dump_workflow_fixture( + ctx: &ActorContext, + metadata: WorkflowFixtureMetadata, +) -> Result { + let db = ctx.sql(); + let meta_rows = db + .query("SELECT key, value FROM _rivet_meta ORDER BY key", None) + .await + .context("dump workflow fixture metadata rows")? + .rows + .iter() + .map(|row| { + Ok(WorkflowFixtureMetaRow { + key: text(row, 0, "metadata key")?, + value: blob(row, 1, "metadata value")?, + }) + }) + .collect::>>()?; + let runtime = db + .query( + "SELECT last_pushed_alarm, inspector_token, queue_next_id FROM _rivet_runtime WHERE id = 1", + None, + ) + .await + .context("dump workflow fixture runtime")? + .rows + .first() + .map(|row| { + Ok::<_, anyhow::Error>(WorkflowFixtureRuntimeRow { + last_pushed_alarm: optional_integer(row, 0, "last_pushed_alarm")?, + inspector_token: optional_text(row, 1, "inspector_token")?, + queue_next_id: integer(row, 2, "queue_next_id")?, + }) + }) + .transpose()?; + let actor = db + .query( + "SELECT has_initialized, input FROM _rivet_actor WHERE id = 1", + None, + ) + .await + .context("dump workflow fixture actor")? + .rows + .first() + .map(|row| { + Ok::<_, anyhow::Error>(WorkflowFixtureActorRow { + has_initialized: integer(row, 0, "has_initialized")?, + input: optional_blob(row, 1, "input")?, + }) + }) + .transpose()?; + let actor_state = db + .query("SELECT state FROM _rivet_actor_state WHERE id = 1", None) + .await + .context("dump workflow fixture actor state")? + .rows + .first() + .map(|row| blob(row, 0, "state")) + .transpose()?; + + let workflow_rows = db + .query("SELECT key, value FROM _rivet_wf_kv ORDER BY key", None) + .await + .context("dump workflow fixture history")? + .rows + .iter() + .map(|row| { + Ok(WorkflowFixtureWorkflowRow { + key: blob(row, 0, "workflow key")?, + value: blob(row, 1, "workflow value")?, + }) + }) + .collect::>>()?; + let queue_rows = db + .query( + "SELECT id, name, body, created_at FROM _rivet_queue ORDER BY id", + None, + ) + .await + .context("dump workflow fixture queue")? + .rows + .iter() + .map(|row| { + Ok(WorkflowFixtureQueueRow { + id: integer(row, 0, "queue id")?, + name: text(row, 1, "queue name")?, + body: blob(row, 2, "queue body")?, + created_at: integer(row, 3, "queue created_at")?, + }) + }) + .collect::>>()?; + let schedule_events = db + .query( + "SELECT event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history FROM _rivet_schedule_events ORDER BY event_id", + None, + ) + .await + .context("dump workflow fixture schedule events")? + .rows + .iter() + .map(|row| { + Ok(WorkflowFixtureScheduleEventRow { + event_id: text(row, 0, "schedule event_id")?, + trigger_at: integer(row, 1, "schedule trigger_at")?, + action: text(row, 2, "schedule action")?, + args: optional_blob(row, 3, "schedule args")?, + kind: integer(row, 4, "schedule kind")?, + cron_expression: optional_text(row, 5, "schedule cron_expression")?, + timezone: optional_text(row, 6, "schedule timezone")?, + interval_ms: optional_integer(row, 7, "schedule interval_ms")?, + last_started_at: optional_integer(row, 8, "schedule last_started_at")?, + max_history: integer(row, 9, "schedule max_history")?, + }) + }) + .collect::>>()?; + let schedule_history = db + .query( + "SELECT id, schedule_id, action, scheduled_at, fired_at, finished_at, result, error_group, error_code, error_message, error_metadata FROM _rivet_schedule_history ORDER BY id", + None, + ) + .await + .context("dump workflow fixture schedule history")? + .rows + .iter() + .map(|row| { + Ok(WorkflowFixtureScheduleHistoryRow { + id: integer(row, 0, "schedule history id")?, + schedule_id: text(row, 1, "schedule history schedule_id")?, + action: text(row, 2, "schedule history action")?, + scheduled_at: integer(row, 3, "schedule history scheduled_at")?, + fired_at: integer(row, 4, "schedule history fired_at")?, + finished_at: optional_integer(row, 5, "schedule history finished_at")?, + result: integer(row, 6, "schedule history result")?, + error_group: optional_text(row, 7, "schedule history error_group")?, + error_code: optional_text(row, 8, "schedule history error_code")?, + error_message: optional_text(row, 9, "schedule history error_message")?, + error_metadata: optional_blob(row, 10, "schedule history error_metadata")?, + }) + }) + .collect::>>()?; + + Ok(WorkflowFixture { + metadata, + meta_rows, + runtime, + actor, + actor_state, + workflow_rows, + queue_rows, + schedule_events, + schedule_history, + }) +} + +/// Restores a decoded fixture into an empty test actor database. This function +/// is unavailable without `test-support` and accepts no caller-provided SQL. +pub async fn restore_workflow_fixture(ctx: &ActorContext, fixture: &WorkflowFixture) -> Result<()> { + validate_fixture_for_restore(fixture)?; + let mut statements = vec![ + statement("DELETE FROM _rivet_schedule_history", None), + statement("DELETE FROM _rivet_schedule_events", None), + statement("DELETE FROM _rivet_queue", None), + statement("DELETE FROM _rivet_wf_kv", None), + statement("DELETE FROM _rivet_actor_state", None), + statement("DELETE FROM _rivet_actor", None), + statement("DELETE FROM _rivet_runtime", None), + statement("DELETE FROM _rivet_meta", None), + ]; + for row in &fixture.meta_rows { + statements.push(statement( + "INSERT INTO _rivet_meta (key, value) VALUES (?, ?)", + Some(vec![ + BindParam::Text(row.key.clone()), + BindParam::Blob(row.value.clone()), + ]), + )); + } + + if let Some(row) = &fixture.runtime { + statements.push(statement( + "INSERT INTO _rivet_runtime (id, last_pushed_alarm, inspector_token, queue_next_id) VALUES (1, ?, ?, ?)", + Some(vec![ + optional_integer_param(row.last_pushed_alarm), + optional_text_param(row.inspector_token.clone()), + BindParam::Integer(row.queue_next_id), + ]), + )); + } + if let Some(row) = &fixture.actor { + statements.push(statement( + "INSERT INTO _rivet_actor (id, has_initialized, input) VALUES (1, ?, ?)", + Some(vec![ + BindParam::Integer(row.has_initialized), + optional_blob_param(row.input.clone()), + ]), + )); + } + if let Some(state) = &fixture.actor_state { + statements.push(statement( + "INSERT INTO _rivet_actor_state (id, state) VALUES (1, ?)", + Some(vec![BindParam::Blob(state.clone())]), + )); + } + for row in &fixture.workflow_rows { + statements.push(statement( + "INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?)", + Some(vec![ + BindParam::Blob(row.key.clone()), + BindParam::Blob(row.value.clone()), + ]), + )); + } + for row in &fixture.queue_rows { + statements.push(statement( + "INSERT INTO _rivet_queue (id, name, body, created_at) VALUES (?, ?, ?, ?)", + Some(vec![ + BindParam::Integer(row.id), + BindParam::Text(row.name.clone()), + BindParam::Blob(row.body.clone()), + BindParam::Integer(row.created_at), + ]), + )); + } + for row in &fixture.schedule_events { + statements.push(statement( + "INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + Some(vec![ + BindParam::Text(row.event_id.clone()), + BindParam::Integer(row.trigger_at), + BindParam::Text(row.action.clone()), + optional_blob_param(row.args.clone()), + BindParam::Integer(row.kind), + optional_text_param(row.cron_expression.clone()), + optional_text_param(row.timezone.clone()), + optional_integer_param(row.interval_ms), + optional_integer_param(row.last_started_at), + BindParam::Integer(row.max_history), + ]), + )); + } + for row in &fixture.schedule_history { + statements.push(statement( + "INSERT INTO _rivet_schedule_history (id, schedule_id, action, scheduled_at, fired_at, finished_at, result, error_group, error_code, error_message, error_metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + Some(vec![ + BindParam::Integer(row.id), + BindParam::Text(row.schedule_id.clone()), + BindParam::Text(row.action.clone()), + BindParam::Integer(row.scheduled_at), + BindParam::Integer(row.fired_at), + optional_integer_param(row.finished_at), + BindParam::Integer(row.result), + optional_text_param(row.error_group.clone()), + optional_text_param(row.error_code.clone()), + optional_text_param(row.error_message.clone()), + optional_blob_param(row.error_metadata.clone()), + ]), + )); + } + + ctx.sql() + .execute_batch(statements) + .await + .context("restore logical workflow fixture")?; + Ok(()) +} + +fn validate_fixture_for_restore(fixture: &WorkflowFixture) -> Result<()> { + let schema_row = fixture + .meta_rows + .iter() + .find(|row| row.key == SCHEMA_VERSION_META_KEY) + .context("workflow fixture is missing its schema_version metadata row")?; + let schema_bytes: [u8; 8] = schema_row + .value + .as_slice() + .try_into() + .context("workflow fixture schema_version must be an i64 little-endian blob")?; + let stored_schema_version = i64::from_le_bytes(schema_bytes); + if stored_schema_version != fixture.metadata.internal_schema_version { + bail!( + "workflow fixture schema metadata mismatch: row is {stored_schema_version}, fixture declares {}", + fixture.metadata.internal_schema_version, + ); + } + if stored_schema_version != crate::actor::internal_storage::schema::INTERNAL_SCHEMA_VERSION { + bail!( + "workflow fixture schema {stored_schema_version} cannot be restored into schema {}", + crate::actor::internal_storage::schema::INTERNAL_SCHEMA_VERSION, + ); + } + if let Some(row) = fixture + .workflow_rows + .iter() + .find(|row| !row.key.starts_with(&WORKFLOW_STORAGE_PREFIX)) + { + bail!( + "workflow fixture row escaped the {:?} namespace: {:?}", + WORKFLOW_STORAGE_PREFIX, + row.key, + ); + } + Ok(()) +} + +fn statement(sql: &str, params: Option>) -> SqliteBatchStatement { + SqliteBatchStatement { + sql: sql.to_owned(), + params, + } +} + +fn integer(row: &[ColumnValue], index: usize, label: &str) -> Result { + match row.get(index) { + Some(ColumnValue::Integer(value)) => Ok(*value), + value => bail!("invalid {label}: expected INTEGER, found {value:?}"), + } +} + +fn optional_integer(row: &[ColumnValue], index: usize, label: &str) -> Result> { + match row.get(index) { + Some(ColumnValue::Null) => Ok(None), + Some(ColumnValue::Integer(value)) => Ok(Some(*value)), + value => bail!("invalid {label}: expected NULL or INTEGER, found {value:?}"), + } +} + +fn text(row: &[ColumnValue], index: usize, label: &str) -> Result { + match row.get(index) { + Some(ColumnValue::Text(value)) => Ok(value.clone()), + value => bail!("invalid {label}: expected TEXT, found {value:?}"), + } +} + +fn optional_text(row: &[ColumnValue], index: usize, label: &str) -> Result> { + match row.get(index) { + Some(ColumnValue::Null) => Ok(None), + Some(ColumnValue::Text(value)) => Ok(Some(value.clone())), + value => bail!("invalid {label}: expected NULL or TEXT, found {value:?}"), + } +} + +fn blob(row: &[ColumnValue], index: usize, label: &str) -> Result> { + match row.get(index) { + Some(ColumnValue::Blob(value)) => Ok(value.clone()), + value => bail!("invalid {label}: expected BLOB, found {value:?}"), + } +} + +fn optional_blob(row: &[ColumnValue], index: usize, label: &str) -> Result>> { + match row.get(index) { + Some(ColumnValue::Null) => Ok(None), + Some(ColumnValue::Blob(value)) => Ok(Some(value.clone())), + value => bail!("invalid {label}: expected NULL or BLOB, found {value:?}"), + } +} + +fn optional_integer_param(value: Option) -> BindParam { + value.map_or(BindParam::Null, BindParam::Integer) +} + +fn optional_text_param(value: Option) -> BindParam { + value.map_or(BindParam::Null, BindParam::Text) +} + +fn optional_blob_param(value: Option>) -> BindParam { + value.map_or(BindParam::Null, BindParam::Blob) +} + +#[cfg(test)] +#[path = "../../tests/workflow_fixture.rs"] +mod tests; diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs index 28613a9342..a2c5b1247e 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/counter.rs @@ -136,6 +136,7 @@ fn counter_factory() -> ActorFactory { ActorEvent::WorkflowReplayRequested { entry_id: _, reply } => { reply.send(Ok(None)); } + ActorEvent::RunWake { .. } => {} } } diff --git a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs index f07d92eab2..c8b9c2c5e6 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/integration/sqlite_corruption_fuzz.rs @@ -629,6 +629,7 @@ fn sqlite_fuzz_factory() -> ActorFactory { ActorEvent::WorkflowReplayRequested { entry_id: _, reply } => { reply.send(Ok(None)); } + ActorEvent::RunWake { .. } => {} } } diff --git a/rivetkit-rust/packages/rivetkit-core/tests/internal_schema.rs b/rivetkit-rust/packages/rivetkit-core/tests/internal_schema.rs index bd340e1f18..aadf4bc08e 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/internal_schema.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/internal_schema.rs @@ -58,3 +58,58 @@ fn unpublished_schema_has_explicit_values_and_minimal_constraints() { ); } } + +#[test] +fn logical_run_wake_metadata_keeps_the_v1_schema_openable() { + use rivetkit_actor_persist::versioned::RunWakeAt; + use vbare::OwnedVersionedData; + + assert_eq!(INTERNAL_SCHEMA_VERSION, 1); + let conn = rusqlite::Connection::open_in_memory().expect("open v1 fixture database"); + initialize_test_schema(&conn).expect("initialize v1 actor schema"); + let logical_wake = RunWakeAt::wrap_latest(Some(1_723_456_789_000)) + .serialize_with_embedded_version(1) + .expect("encode logical run wake"); + conn.execute( + "INSERT INTO _rivet_meta (key, value) VALUES (?1, ?2)", + rusqlite::params![ + crate::actor::internal_storage::RUN_WAKE_AT_META_KEY, + logical_wake.clone() + ], + ) + .expect("persist reserved metadata row"); + conn.execute( + "INSERT INTO _rivet_runtime (id, last_pushed_alarm, inspector_token, queue_next_id) VALUES (1, ?1, NULL, 2)", + rusqlite::params![1_723_456_789_500_i64], + ) + .expect("persist v1 runtime row"); + + let stored_schema: Vec = conn + .query_row( + LOAD_META_TEXT_SQL, + rusqlite::params![SCHEMA_VERSION_KEY], + |row| row.get(0), + ) + .expect("read schema version as an old runtime would"); + assert_eq!(decode_schema_version(&stored_schema).unwrap(), 1); + let legacy_runtime: (Option, Option, i64) = conn + .query_row( + "SELECT last_pushed_alarm, inspector_token, queue_next_id FROM _rivet_runtime WHERE id = 1", + [], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), + ) + .expect("open runtime through the v1 projection"); + assert_eq!(legacy_runtime, (Some(1_723_456_789_500), None, 2)); + let stored_wake: Vec = conn + .query_row( + LOAD_META_TEXT_SQL, + rusqlite::params![crate::actor::internal_storage::RUN_WAKE_AT_META_KEY], + |row| row.get(0), + ) + .expect("preserve unknown metadata row"); + assert_eq!(stored_wake, logical_wake); + assert_eq!( + RunWakeAt::deserialize_with_embedded_version(&stored_wake).unwrap(), + Some(1_723_456_789_000), + ); +} diff --git a/rivetkit-rust/packages/rivetkit-core/tests/migrate_kv_to_sqlite.rs b/rivetkit-rust/packages/rivetkit-core/tests/migrate_kv_to_sqlite.rs index b85c7a1a0a..bc7d8c6070 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/migrate_kv_to_sqlite.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/migrate_kv_to_sqlite.rs @@ -698,6 +698,7 @@ async fn imports_legacy_kv_snapshot_to_sqlite_once() -> Result<()> { ..actor.clone() }, last_pushed_alarm: Some(5678), + run_wake_at: None, }) ); let schedule_rows = ctx diff --git a/rivetkit-rust/packages/rivetkit-core/tests/queue.rs b/rivetkit-rust/packages/rivetkit-core/tests/queue.rs index dc59742181..1060672ad1 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/queue.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/queue.rs @@ -71,6 +71,63 @@ mod moved_tests { ); } + #[tokio::test] + async fn wait_for_available_does_not_consume_or_reorder_messages() { + let queue = test_queue(); + crate::actor::internal_storage::schema::ensure_internal_schema(queue.sql()) + .await + .expect("initialize queue storage"); + queue.send("first", b"one").await.expect("send first"); + queue.send("target", b"two").await.expect("send target"); + + queue + .wait_for_names_available(vec!["target".to_owned()], QueueWaitOpts::default()) + .await + .expect("wait for matching queue message"); + + let messages = queue.inspect_messages().await.expect("inspect queue"); + assert_eq!( + messages + .iter() + .map(|message| message.name.as_str()) + .collect::>(), + vec!["first", "target"], + ); + } + + #[tokio::test] + async fn durable_completion_verifies_persisted_name_and_is_idempotent() { + let queue = test_queue(); + crate::actor::internal_storage::schema::ensure_internal_schema(queue.sql()) + .await + .expect("initialize queue storage"); + let message = queue + .send("expected", b"body") + .await + .expect("send queue message"); + let error = queue + .complete_persisted_message(message.id, "wrong", None) + .await + .expect_err("wrong name must fail while completing"); + let error = rivet_error::RivetError::extract(&error); + assert_eq!(error.group(), "queue"); + assert_eq!(error.code(), "message_identity_mismatch"); + assert_eq!(queue.inspect_messages().await.expect("inspect").len(), 1); + + assert!( + queue + .complete_persisted_message(message.id, "expected", None) + .await + .expect("complete matching message") + ); + assert!( + !queue + .complete_persisted_message(message.id, "expected", None) + .await + .expect("repeat completion is an idempotent miss") + ); + } + #[tokio::test] async fn next_batch_supports_large_name_filters_without_sql_bind_expansion() { let queue = test_queue(); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/schedule.rs b/rivetkit-rust/packages/rivetkit-core/tests/schedule.rs index c3e98280f3..0dbee10c20 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/schedule.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/schedule.rs @@ -768,4 +768,96 @@ mod moved_tests { advance_schedule_time(&ctx, BASE_TIME + 5_000, Duration::from_millis(1)).await; assert_eq!(fired.load(Ordering::SeqCst), 1); } + + #[tokio::test] + async fn run_wake_and_schedule_share_the_earliest_alarm_without_overwriting() { + let ctx = context("actor-run-wake-multiplex"); + ctx.set_run_wake_at(Some(BASE_TIME + 10_000)).await.unwrap(); + ctx.wait_for_pending_alarm_writes().await; + assert_eq!( + crate::actor::internal_storage::load_last_pushed_alarm(ctx.sql()) + .await + .unwrap(), + Some(BASE_TIME + 10_000), + ); + + ctx.at(BASE_TIME + 5_000, "tick", &[]).await.unwrap(); + ctx.wait_for_pending_alarm_writes().await; + assert_eq!( + crate::actor::internal_storage::load_last_pushed_alarm(ctx.sql()) + .await + .unwrap(), + Some(BASE_TIME + 5_000), + ); + + ctx.set_schedule_time_for_tests(BASE_TIME + 5_000); + let schedules = ctx.take_due_schedule_dispatches().await.unwrap(); + assert_eq!(schedules.len(), 1); + assert_eq!(ctx.run_wake_at(), Some(BASE_TIME + 10_000)); + ctx.wait_for_pending_alarm_writes().await; + assert_eq!( + crate::actor::internal_storage::load_last_pushed_alarm(ctx.sql()) + .await + .unwrap(), + Some(BASE_TIME + 10_000), + ); + + ctx.set_schedule_time_for_tests(BASE_TIME + 10_000); + assert_eq!( + ctx.consume_due_run_wake().await.unwrap(), + Some(BASE_TIME + 10_000) + ); + assert_eq!(ctx.run_wake_at(), None); + } + + #[tokio::test] + async fn clearing_run_wake_preserves_a_later_schedule_alarm() { + let ctx = context("actor-run-wake-clear"); + ctx.at(BASE_TIME + 10_000, "tick", &[]).await.unwrap(); + ctx.set_run_wake_at(Some(BASE_TIME + 5_000)).await.unwrap(); + ctx.wait_for_pending_alarm_writes().await; + assert_eq!( + crate::actor::internal_storage::load_last_pushed_alarm(ctx.sql()) + .await + .unwrap(), + Some(BASE_TIME + 5_000), + ); + + ctx.set_run_wake_at(None).await.unwrap(); + ctx.wait_for_pending_alarm_writes().await; + assert_eq!( + crate::actor::internal_storage::load_last_pushed_alarm(ctx.sql()) + .await + .unwrap(), + Some(BASE_TIME + 10_000), + ); + } + + #[tokio::test] + async fn run_wake_is_persisted_independently_from_transport_alarm() { + let harness = ActorContextHarness::new(); + let ctx = harness.context("actor-run-wake-persist", "actor", Vec::new(), "local"); + ctx.set_schedule_time_for_tests(BASE_TIME); + ctx.set_run_wake_at(Some(BASE_TIME + 7_000)).await.unwrap(); + ctx.wait_for_pending_alarm_writes().await; + crate::actor::internal_storage::persist_last_pushed_alarm( + ctx.sql(), + Some(BASE_TIME + 9_000), + ) + .await + .unwrap(); + + assert_eq!( + crate::actor::internal_storage::load_run_wake_at(ctx.sql()) + .await + .unwrap(), + Some(BASE_TIME + 7_000), + ); + assert_eq!( + crate::actor::internal_storage::load_last_pushed_alarm(ctx.sql()) + .await + .unwrap(), + Some(BASE_TIME + 9_000), + ); + } } diff --git a/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs b/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs index c4f54b77cc..eedba335de 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/sql_efficiency.rs @@ -116,9 +116,6 @@ fn fixture(row_count: usize) -> Connection { let mut user_kv_insert = tx .prepare("INSERT INTO _rivet_user_kv (key, value) VALUES (?, x'01')") .expect("prepare user kv seed"); - let mut workflow_kv_insert = tx - .prepare("INSERT INTO _rivet_wf_kv (key, value) VALUES (?, x'01')") - .expect("prepare workflow kv seed"); let mut schedule_insert = tx .prepare("INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES (?, ?, 'run', x'01', ?, NULL, NULL, NULL, NULL, 100)") .expect("prepare schedule seed"); @@ -137,9 +134,6 @@ fn fixture(row_count: usize) -> Connection { user_kv_insert .execute([key.as_bytes()]) .expect("seed user kv"); - workflow_kv_insert - .execute([key.as_bytes()]) - .expect("seed workflow kv"); let event_id = if index % 3 == 0 { format!("at:{key}") } else { @@ -163,7 +157,6 @@ fn fixture(row_count: usize) -> Connection { drop(conn_state_insert); drop(queue_insert); drop(user_kv_insert); - drop(workflow_kv_insert); drop(schedule_insert); drop(history_insert); tx.commit().expect("commit fixture seed"); @@ -355,6 +348,12 @@ fn query_catalog() -> Vec { params: vec![0_i64.into(), 128_i64.into()], expectation: indexed(None, &["_rivet_queue"]), }, + QueryCase { + id: "queue.load_name", + sql: internal_storage::LOAD_QUEUE_MESSAGE_NAME_SQL.into(), + params: vec![1_i64.into()], + expectation: indexed(None, &["_rivet_queue"]), + }, QueryCase { id: "queue.load_ids", sql: internal_storage::load_queue_messages_by_ids_sql(3), @@ -465,6 +464,12 @@ fn query_catalog() -> Vec { params: vec![], expectation: indexed(None, &["_rivet_runtime"]), }, + QueryCase { + id: "runtime.run_wake", + sql: internal_storage::LOAD_RUN_WAKE_AT_SQL.into(), + params: vec![text(internal_storage::RUN_WAKE_AT_META_KEY)], + expectation: indexed(None, &["_rivet_meta"]), + }, QueryCase { id: "runtime.inspector_token", sql: internal_storage::LOAD_INSPECTOR_TOKEN_SQL.into(), diff --git a/rivetkit-rust/packages/rivetkit-core/tests/state.rs b/rivetkit-rust/packages/rivetkit-core/tests/state.rs index 3156a64c15..3e8c260eb4 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/state.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/state.rs @@ -24,6 +24,7 @@ mod moved_tests { use crate::actor::messages::StateDelta; use crate::actor::task::LifecycleEvent; use crate::kv::tests::new_in_memory; + use crate::sqlite::BindParam; use crate::{ActorContext, RequestSaveOpts}; use super::{ @@ -323,6 +324,600 @@ mod moved_tests { ); } + #[tokio::test] + async fn state_transaction_commits_user_sql_actor_state_and_dirty_connection_together() { + let ctx = new_with_kv( + "actor-state-tx", + "state-tx", + Vec::new(), + "local", + new_in_memory(), + ); + ctx.sql() + .execute( + "CREATE TABLE user_values (id INTEGER PRIMARY KEY, value TEXT NOT NULL)", + None, + ) + .await + .expect("create user table"); + + let conn = ConnHandle::new("conn-state-tx", Vec::new(), vec![1], true); + conn.configure_hibernation(Some(HibernatableConnectionMetadata { + gateway_id: *b"gate", + request_id: *b"tx01", + server_message_index: 1, + client_message_index: 2, + request_path: "/ws".to_owned(), + request_headers: Default::default(), + })); + ctx.add_conn(conn.clone()); + conn.set_state_initial(vec![7, 8, 9]); + ctx.request_hibernation_transport_save(conn.id()); + let removed_conn = ConnHandle::new("conn-state-tx-removed", Vec::new(), vec![3], true); + removed_conn.configure_hibernation(Some(HibernatableConnectionMetadata { + gateway_id: *b"gate", + request_id: *b"tx03", + server_message_index: 3, + client_message_index: 4, + request_path: "/removed".to_owned(), + request_headers: Default::default(), + })); + ctx.add_conn(removed_conn.clone()); + ctx.save_state(vec![StateDelta::ConnHibernation { + conn: removed_conn.id().to_owned(), + bytes: vec![3], + }]) + .await + .expect("seed connection that will be removed"); + ctx.request_hibernation_transport_removal(removed_conn.id().to_owned()); + + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + transaction + .execute( + "INSERT INTO user_values (id, value) VALUES (1, 'committed')", + None, + ) + .await + .expect("insert user row"); + transaction + .commit(vec![StateDelta::ActorState(vec![4, 5, 6])]) + .await + .expect("commit state transaction"); + + let rows = ctx + .sql() + .query("SELECT value FROM user_values WHERE id = 1", None) + .await + .expect("query user row"); + assert_eq!( + rows.rows, + vec![vec![crate::sqlite::ColumnValue::Text("committed".into())]], + ); + let actor = internal_storage::load_actor_snapshot(ctx.sql()) + .await + .expect("load actor snapshot") + .expect("actor snapshot should exist") + .actor; + assert_eq!(actor.state, vec![4, 5, 6]); + let connections = internal_storage::load_connections(ctx.sql()) + .await + .expect("load connection snapshots"); + assert_eq!(connections.len(), 1); + assert_eq!(connections[0].id, conn.id()); + assert_eq!(connections[0].state, vec![7, 8, 9]); + assert!(!ctx.has_pending_hibernation_changes()); + } + + #[tokio::test] + async fn state_transaction_begin_failure_restores_scheduled_state_save() { + let ctx = new_with_kv( + "actor-state-tx-begin-failure", + "state-tx-begin-failure", + Vec::new(), + "local", + new_in_memory(), + ); + ctx.set_input(Some(vec![1])); + assert!(ctx.0.pending_save.lock().is_some()); + + let result = ctx.begin_state_transaction(Some(Duration::ZERO)).await; + assert!( + result.is_err(), + "zero timeout must reject transaction begin" + ); + + assert!( + ctx.0.pending_save.lock().is_some(), + "begin failure must reschedule the save cleared before acquiring exclusion", + ); + } + + #[tokio::test] + async fn state_transaction_callback_rollback_reschedules_state_and_reverts_sql() { + let ctx = new_with_kv( + "actor-state-tx-rollback", + "state-tx-rollback", + Vec::new(), + "local", + new_in_memory(), + ); + ctx.sql() + .execute( + "CREATE TABLE user_values (id INTEGER PRIMARY KEY, value TEXT NOT NULL)", + None, + ) + .await + .expect("create user table"); + ctx.set_input(Some(vec![1])); + + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + transaction + .execute( + "INSERT INTO user_values (id, value) VALUES (1, 'rolled-back')", + None, + ) + .await + .expect("insert user row"); + transaction + .rollback() + .await + .expect("roll back state transaction"); + + let rows = ctx + .sql() + .query("SELECT value FROM user_values", None) + .await + .expect("query user rows"); + assert!(rows.rows.is_empty()); + assert!(ctx.0.pending_save.lock().is_some()); + } + + #[tokio::test] + async fn state_transaction_finalization_is_one_shot_across_clones() { + let ctx = new_with_kv( + "actor-state-tx-finalize", + "state-tx-finalize", + Vec::new(), + "local", + new_in_memory(), + ); + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + let cloned = transaction.clone(); + transaction + .commit(Vec::new()) + .await + .expect("first finalization should commit"); + + assert!(cloned.commit(Vec::new()).await.is_err()); + assert!(cloned.rollback().await.is_err()); + ctx.sql() + .execute("SELECT 1", None) + .await + .expect("coordinator should be released after one commit"); + } + + #[tokio::test] + async fn state_transaction_combined_statement_budget_accepts_exact_boundary() { + let ctx = new_with_kv( + "actor-state-tx-budget-boundary", + "state-tx-budget-boundary", + Vec::new(), + "local", + new_in_memory(), + ); + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + // Actor state adds two internal upserts, reaching the combined 128-row + // limit exactly. + for _ in 0..126 { + transaction + .execute("SELECT 1", None) + .await + .expect("execute boundary statement"); + } + transaction + .commit(vec![StateDelta::ActorState(vec![1])]) + .await + .expect("exact transaction budget boundary should commit"); + assert_eq!( + internal_storage::load_actor_snapshot(ctx.sql()) + .await + .expect("load actor snapshot") + .expect("actor snapshot should exist") + .actor + .state, + vec![1], + ); + } + + #[tokio::test] + async fn state_transaction_combined_payload_budget_accepts_exact_boundary() { + let ctx = new_with_kv( + "actor-state-tx-payload-boundary", + "state-tx-payload-boundary", + Vec::new(), + "local", + new_in_memory(), + ); + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + // The actor upsert contributes one eight-byte integer and the state + // upsert contributes the one-byte state below. + transaction + .execute( + "SELECT ?", + Some(vec![BindParam::Blob(vec![ + 0; + internal_storage::KV_TX_MAX_PAYLOAD_BYTES + - 9 + ])]), + ) + .await + .expect("execute exact payload boundary statement"); + transaction + .commit(vec![StateDelta::ActorState(vec![1])]) + .await + .expect("exact payload budget boundary should commit"); + } + + #[tokio::test] + async fn state_transaction_rejects_multi_statement_exec() { + let ctx = new_with_kv( + "actor-state-tx-exec", + "state-tx-exec", + Vec::new(), + "local", + new_in_memory(), + ); + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + let error = transaction + .exec("SELECT 1; SELECT 2") + .await + .expect_err("state-aware transactions must reject exec"); + assert!(format!("{error:#}").contains("single-statement execute")); + transaction + .rollback() + .await + .expect("rejected exec should leave the transaction rollback-safe"); + } + + #[tokio::test] + async fn state_transaction_counts_unbound_sql_text_toward_payload_budget() { + let ctx = new_with_kv( + "actor-state-tx-inline-payload", + "state-tx-inline-payload", + Vec::new(), + "local", + new_in_memory(), + ); + ctx.sql() + .execute("CREATE TABLE user_values (value TEXT NOT NULL)", None) + .await + .expect("create user table"); + + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + transaction + .execute("INSERT INTO user_values VALUES ('before')", None) + .await + .expect("insert user row"); + transaction + .execute( + format!( + "SELECT '{}'", + "x".repeat(internal_storage::KV_TX_MAX_PAYLOAD_BYTES) + ), + None, + ) + .await + .expect("execute oversized inline literal before commit validation"); + let error = transaction + .commit(Vec::new()) + .await + .expect_err("inline SQL text must count toward the payload budget"); + assert!(format!("{error:#}").contains("exceeds transaction budget")); + + let rows = ctx + .sql() + .query("SELECT value FROM user_values", None) + .await + .expect("query rolled-back user rows"); + assert!(rows.rows.is_empty()); + } + + #[tokio::test] + async fn state_transaction_combined_statement_budget_overflow_rolls_back_user_sql() { + let ctx = new_with_kv( + "actor-state-tx-budget-overflow", + "state-tx-budget-overflow", + Vec::new(), + "local", + new_in_memory(), + ); + ctx.sql() + .execute( + "CREATE TABLE user_values (id INTEGER PRIMARY KEY, value TEXT NOT NULL)", + None, + ) + .await + .expect("create user table"); + ctx.sql() + .execute( + "INSERT INTO user_values (id, value) VALUES (1, 'before')", + None, + ) + .await + .expect("seed user row"); + + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + transaction + .execute("UPDATE user_values SET value = 'after' WHERE id = 1", None) + .await + .expect("update user row"); + for _ in 0..126 { + transaction + .execute("SELECT 1", None) + .await + .expect("execute counted statement"); + } + let error = transaction + .commit(vec![StateDelta::ActorState(vec![2])]) + .await + .expect_err("combined state statements must overflow row budget"); + assert!(format!("{error:#}").contains("exceeds transaction budget")); + + let rows = ctx + .sql() + .query("SELECT value FROM user_values WHERE id = 1", None) + .await + .expect("query rolled-back user row"); + assert_eq!( + rows.rows, + vec![vec![crate::sqlite::ColumnValue::Text("before".into())]], + ); + assert!( + internal_storage::load_actor_snapshot(ctx.sql()) + .await + .expect("load actor snapshot") + .is_none(), + ); + } + + #[tokio::test] + async fn actor_state_payload_overflow_rolls_back_user_sql() { + let ctx = new_with_kv( + "actor-state-tx-state-overflow", + "state-tx-state-overflow", + Vec::new(), + "local", + new_in_memory(), + ); + ctx.sql() + .execute( + "CREATE TABLE user_values (id INTEGER PRIMARY KEY, value TEXT NOT NULL)", + None, + ) + .await + .expect("create user table"); + ctx.sql() + .execute( + "INSERT INTO user_values (id, value) VALUES (1, 'before')", + None, + ) + .await + .expect("seed user row"); + + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + transaction + .execute("UPDATE user_values SET value = 'after' WHERE id = 1", None) + .await + .expect("update user row"); + transaction + .execute( + "SELECT ?", + Some(vec![BindParam::Blob(vec![ + 0; + internal_storage::KV_TX_MAX_PAYLOAD_BYTES + - 8 + ])]), + ) + .await + .expect("execute user payload at pre-state limit"); + transaction + .commit(vec![StateDelta::ActorState(vec![2])]) + .await + .expect_err("actor state byte must overflow combined payload budget"); + + let rows = ctx + .sql() + .query("SELECT value FROM user_values WHERE id = 1", None) + .await + .expect("query rolled-back user row"); + assert_eq!( + rows.rows, + vec![vec![crate::sqlite::ColumnValue::Text("before".into())]], + ); + assert!( + internal_storage::load_actor_snapshot(ctx.sql()) + .await + .expect("load actor snapshot") + .is_none(), + ); + } + + #[tokio::test] + async fn dropping_state_transaction_releases_state_save_exclusion() { + let ctx = new_with_kv( + "actor-state-tx-drop", + "state-tx-drop", + Vec::new(), + "local", + new_in_memory(), + ); + let save_guard = Arc::clone(&ctx.0.save_guard); + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + assert!(save_guard.try_lock().is_err()); + + drop(transaction); + let acquired_save_guard = + tokio::time::timeout(Duration::from_millis(100), save_guard.lock()) + .await + .expect("dropping transaction must release save exclusion"); + drop(acquired_save_guard); + tokio::time::timeout(Duration::from_secs(1), ctx.sql().execute("SELECT 1", None)) + .await + .expect("dropping transaction must release sqlite coordinator") + .expect("regular sqlite work should resume after dropped transaction"); + } + + #[tokio::test] + async fn state_transaction_failure_rolls_back_sql_and_restores_hibernation_changes() { + let ctx = new_with_kv( + "actor-state-tx-failure", + "state-tx-failure", + Vec::new(), + "local", + new_in_memory(), + ); + ctx.sql() + .execute( + "CREATE TABLE user_values (id INTEGER PRIMARY KEY, value TEXT NOT NULL)", + None, + ) + .await + .expect("create user table"); + ctx.sql() + .execute( + "INSERT INTO user_values (id, value) VALUES (1, 'before')", + None, + ) + .await + .expect("seed user row"); + + let conn = ConnHandle::new("conn-state-tx-failure", Vec::new(), vec![1], true); + conn.configure_hibernation(Some(HibernatableConnectionMetadata { + gateway_id: *b"gate", + request_id: *b"tx02", + server_message_index: 1, + client_message_index: 2, + request_path: "/ws".to_owned(), + request_headers: Default::default(), + })); + ctx.add_conn(conn.clone()); + conn.set_state_initial(vec![2]); + ctx.request_hibernation_transport_save(conn.id()); + + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + transaction + .execute("UPDATE user_values SET value = 'after' WHERE id = 1", None) + .await + .expect("update user row"); + transaction + .commit(vec![ + StateDelta::ActorState(vec![9]), + StateDelta::ConnHibernation { + conn: "missing-connection".to_owned(), + bytes: vec![3], + }, + ]) + .await + .expect_err("invalid hibernation delta must fail the atomic commit"); + + let rows = ctx + .sql() + .query("SELECT value FROM user_values WHERE id = 1", None) + .await + .expect("query user row"); + assert_eq!( + rows.rows, + vec![vec![crate::sqlite::ColumnValue::Text("before".into())]], + ); + assert!( + internal_storage::load_actor_snapshot(ctx.sql()) + .await + .expect("load actor snapshot") + .is_none(), + ); + assert!(ctx.has_pending_hibernation_changes()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn state_mutation_after_transaction_snapshot_remains_dirty() { + let (entered_tx, mut entered_rx) = mpsc::unbounded_channel(); + let release = Arc::new(Semaphore::new(0)); + let ctx = Arc::new(new_with_kv_and_write_gate( + "actor-state-tx-revision", + "state-tx-revision", + Vec::new(), + "local", + new_in_memory(), + TestSqliteWriteGate { + sql_prefix: "INSERT INTO _rivet_actor (", + entered_tx, + release: release.clone(), + }, + )); + ctx.set_input(Some(vec![1])); + let transaction = ctx + .begin_state_transaction(None) + .await + .expect("begin state transaction"); + let commit = tokio::spawn({ + let transaction = transaction.clone(); + async move { + transaction + .commit(vec![StateDelta::ActorState(vec![1])]) + .await + } + }); + entered_rx + .recv() + .await + .expect("commit should reach actor snapshot write"); + + ctx.set_input(Some(vec![2])); + release.add_permits(1); + commit + .await + .expect("commit task should not panic") + .expect("state transaction should commit"); + + assert!( + ctx.0.state_dirty.load(std::sync::atomic::Ordering::SeqCst), + "mutation after the serialized revision must remain dirty", + ); + } + #[tokio::test] async fn save_state_applies_actor_upsert_and_hibernation_delete_in_one_batch() { let kv = new_in_memory(); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index afa7091149..d71fd2711d 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -2704,6 +2704,194 @@ pub(crate) mod moved_tests { ); } + #[tokio::test] + async fn startup_restores_persisted_run_wake_before_runtime_preamble() { + let ctx = new_with_kv( + "actor-startup-run-wake", + "task-startup-run-wake", + Vec::new(), + "local", + new_in_memory(), + ); + crate::actor::internal_storage::schema::ensure_internal_schema(ctx.sql()) + .await + .expect("initialize schema"); + crate::actor::internal_storage::persist_actor_core_and_connections( + ctx.sql(), + Some(&PersistedActor { + has_initialized: true, + ..PersistedActor::default() + }), + &[], + &[], + ) + .await + .expect("persist actor row"); + crate::actor::internal_storage::persist_run_wake_at(ctx.sql(), Some(4_000)) + .await + .expect("persist logical run wake"); + let (handle, mut rx) = test_envoy_handle(); + ctx.configure_envoy(handle, Some(14)); + ctx.set_schedule_time_for_tests(1_000); + let mut task = new_task(ctx.clone()); + let (start_tx, start_rx) = oneshot::channel(); + + task.handle_lifecycle(LifecycleCommand::Start { reply: start_tx }) + .await; + start_rx + .await + .expect("start reply should send") + .expect("start should succeed"); + + assert_eq!(ctx.run_wake_at(), Some(4_000)); + assert_eq!( + recv_alarm_now(&mut rx, "actor-startup-run-wake", Some(14)), + Some(4_000), + ); + } + + #[tokio::test] + async fn startup_delivers_a_persisted_run_wake_that_became_overdue() { + let ctx = new_with_kv( + "actor-startup-overdue-run-wake", + "task-startup-overdue-run-wake", + Vec::new(), + "local", + new_in_memory(), + ); + crate::actor::internal_storage::schema::ensure_internal_schema(ctx.sql()) + .await + .expect("initialize schema"); + crate::actor::internal_storage::persist_actor_core_and_connections( + ctx.sql(), + Some(&PersistedActor { + has_initialized: true, + ..PersistedActor::default() + }), + &[], + &[], + ) + .await + .expect("persist actor row"); + crate::actor::internal_storage::persist_run_wake_at(ctx.sql(), Some(500)) + .await + .expect("persist overdue logical run wake"); + let (handle, mut rx) = test_envoy_handle(); + ctx.configure_envoy(handle, Some(15)); + ctx.set_schedule_time_for_tests(1_000); + + let wake_count = Arc::new(AtomicUsize::new(0)); + let factory = Arc::new(ActorFactory::new(Default::default(), { + let wake_count = wake_count.clone(); + move |start| { + let wake_count = wake_count.clone(); + Box::pin(async move { + let mut events = start.events; + while let Some(event) = events.recv().await { + match event { + ActorEvent::RunWake { reply } => { + wake_count.fetch_add(1, Ordering::SeqCst); + reply.send(Ok(())); + } + ActorEvent::BeginSleep => {} + ActorEvent::FinalizeSleep { reply } | ActorEvent::Destroy { reply } => { + reply.send(Ok(())); + break; + } + _ => {} + } + } + Ok(()) + }) + } + })); + let mut task = new_task_with_factory(ctx.clone(), factory); + let (start_tx, start_rx) = oneshot::channel(); + + task.handle_lifecycle(LifecycleCommand::Start { reply: start_tx }) + .await; + start_rx + .await + .expect("start reply should send") + .expect("start should succeed"); + + assert_eq!( + recv_alarm_now(&mut rx, "actor-startup-overdue-run-wake", Some(15)), + Some(500), + ); + let (alarm_tx, alarm_rx) = oneshot::channel(); + task.handle_lifecycle(LifecycleCommand::FireAlarm { reply: alarm_tx }) + .await; + alarm_rx + .await + .expect("alarm reply should send") + .expect("overdue alarm should fire"); + assert_eq!( + recv_alarm_now(&mut rx, "actor-startup-overdue-run-wake", Some(15)), + None, + ); + wait_for_count(&wake_count, 1).await; + assert_eq!(ctx.run_wake_at(), None); + assert_eq!( + crate::actor::internal_storage::load_run_wake_at(ctx.sql()) + .await + .expect("load persisted run wake"), + None, + ); + } + + #[tokio::test] + async fn stale_alarm_ack_cannot_overwrite_the_latest_transport_deadline() { + let ctx = new_with_kv( + "actor-stale-alarm-ack", + "task-stale-alarm-ack", + Vec::new(), + "local", + new_in_memory(), + ); + let (handle, mut rx) = test_envoy_handle(); + ctx.configure_envoy(handle, Some(13)); + ctx.set_schedule_time_for_tests(1_000); + + ctx.set_run_wake_at(Some(3_000)).await.expect("first alarm"); + let first_ack = match rx.recv().await.expect("first alarm message") { + ToEnvoyMessage::SetAlarm { + alarm_ts, + ack_tx: Some(ack_tx), + .. + } => { + assert_eq!(alarm_ts, Some(3_000)); + ack_tx + } + _ => panic!("expected first alarm message"), + }; + + ctx.set_run_wake_at(Some(2_000)) + .await + .expect("second alarm"); + let second_ack = match rx.recv().await.expect("second alarm message") { + ToEnvoyMessage::SetAlarm { + alarm_ts, + ack_tx: Some(ack_tx), + .. + } => { + assert_eq!(alarm_ts, Some(2_000)); + ack_tx + } + _ => panic!("expected second alarm message"), + }; + + second_ack.send(()).expect("ack latest alarm"); + first_ack.send(()).expect("ack stale alarm"); + ctx.wait_for_pending_alarm_writes().await; + assert_eq!( + crate::actor::internal_storage::load_last_pushed_alarm(ctx.sql()) + .await + .expect("load persisted transport alarm"), + Some(2_000), + ); + } + #[tokio::test] async fn fire_due_alarms_dispatches_overdue_work_during_sleep_grace() { let ctx = new_with_kv( @@ -2737,6 +2925,211 @@ pub(crate) mod moved_tests { assert!(ctx.list_scheduled_events().await.unwrap().is_empty()); } + #[tokio::test] + async fn fire_due_alarms_dispatches_run_wake_once_alongside_schedules() { + let ctx = new_with_kv( + "actor-run-wake", + "task-run-wake", + Vec::new(), + "local", + new_in_memory(), + ); + let (events_tx, mut events_rx) = mpsc::unbounded_channel(); + ctx.configure_actor_events(Some(events_tx)); + ctx.set_schedule_time_for_tests(100); + ctx.set_run_wake_at(Some(100)) + .await + .expect("persist run wake"); + ctx.at(100, "tick", &[]).await.expect("persist schedule"); + + let mut task = new_task(ctx.clone()); + task.lifecycle = LifecycleState::Started; + let fire = tokio::spawn(async move { + let result = task.fire_due_alarms().await; + (task, result) + }); + + let mut saw_action = false; + let mut saw_run_wake = false; + for _ in 0..2 { + match events_rx.recv().await.expect("due actor event") { + ActorEvent::Action { reply, .. } => { + saw_action = true; + reply.send(Ok(Vec::new())); + } + ActorEvent::RunWake { reply } => { + saw_run_wake = true; + reply.send(Ok(())); + } + other => panic!("unexpected due event {}", other.kind()), + } + } + assert!(saw_action); + assert!(saw_run_wake); + let (mut task, fire_result) = fire.await.expect("alarm task should not panic"); + fire_result.expect("dispatch due alarms"); + assert_eq!(ctx.run_wake_at(), None); + + task.fire_due_alarms() + .await + .expect("repeat alarm delivery should be a no-op"); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + + #[tokio::test] + async fn failed_run_wake_restart_restores_the_logical_deadline() { + let ctx = new_with_kv( + "actor-run-wake-restart-retry", + "task-run-wake-restart-retry", + Vec::new(), + "local", + new_in_memory(), + ); + let (events_tx, mut events_rx) = mpsc::unbounded_channel(); + ctx.configure_actor_events(Some(events_tx)); + ctx.set_schedule_time_for_tests(100); + ctx.set_run_wake_at(Some(100)) + .await + .expect("persist run wake"); + + let mut task = new_task(ctx.clone()); + task.lifecycle = LifecycleState::Started; + let fire = tokio::spawn(async move { task.fire_due_alarms().await }); + match events_rx.recv().await.expect("due run wake event") { + ActorEvent::RunWake { reply } => { + reply.send(Err(anyhow::anyhow!("runtime restart unavailable"))); + } + other => panic!("expected run wake, got {}", other.kind()), + } + + let error = fire + .await + .expect("alarm task should not panic") + .expect_err("failed restart should fail alarm dispatch"); + assert!(format!("{error:#}").contains("restart run handler for due wake")); + assert_eq!(ctx.run_wake_at(), Some(100)); + assert_eq!( + crate::actor::internal_storage::load_run_wake_at(ctx.sql()) + .await + .expect("load restored run wake"), + Some(100), + ); + } + + #[tokio::test] + async fn failed_run_wake_delivery_restores_the_logical_deadline() { + let ctx = new_with_kv( + "actor-run-wake-retry", + "task-run-wake-retry", + Vec::new(), + "local", + new_in_memory(), + ); + let (events_tx, events_rx) = mpsc::unbounded_channel(); + drop(events_rx); + ctx.configure_actor_events(Some(events_tx)); + ctx.set_schedule_time_for_tests(100); + ctx.set_run_wake_at(Some(100)) + .await + .expect("persist run wake"); + + let mut task = new_task(ctx.clone()); + task.lifecycle = LifecycleState::Started; + let error = task + .fire_due_alarms() + .await + .expect_err("closed runtime inbox should fail delivery"); + + assert!(format!("{error:#}").contains("dispatch due run wake")); + assert_eq!(ctx.run_wake_at(), Some(100)); + assert_eq!( + crate::actor::internal_storage::load_run_wake_at(ctx.sql()) + .await + .expect("load restored run wake"), + Some(100), + ); + } + + #[tokio::test] + async fn failed_schedule_drain_restores_the_consumed_run_wake() { + let ctx = new_with_kv( + "actor-run-wake-schedule-failure", + "task-run-wake-schedule-failure", + Vec::new(), + "local", + new_in_memory(), + ); + let (events_tx, mut events_rx) = mpsc::unbounded_channel(); + ctx.configure_actor_events(Some(events_tx)); + ctx.set_schedule_time_for_tests(100); + ctx.set_run_wake_at(Some(100)) + .await + .expect("persist run wake"); + ctx.sql() + .execute( + "INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES ('invalid', 100, 'tick', NULL, 99, NULL, NULL, NULL, NULL, 0)", + None, + ) + .await + .expect("seed malformed due schedule"); + + let mut task = new_task(ctx.clone()); + task.lifecycle = LifecycleState::Started; + task.fire_due_alarms() + .await + .expect_err("malformed schedule should fail alarm dispatch"); + + assert_eq!(ctx.run_wake_at(), Some(100)); + assert_eq!( + crate::actor::internal_storage::load_run_wake_at(ctx.sql()) + .await + .expect("load restored run wake"), + Some(100), + ); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + + #[tokio::test] + async fn destroy_grace_does_not_restart_the_run_handler() { + let ctx = new_with_kv( + "actor-run-wake-destroy", + "task-run-wake-destroy", + Vec::new(), + "local", + new_in_memory(), + ); + let (events_tx, mut events_rx) = mpsc::unbounded_channel(); + ctx.configure_actor_events(Some(events_tx)); + ctx.set_schedule_time_for_tests(100); + ctx.set_run_wake_at(Some(100)) + .await + .expect("persist run wake"); + + let mut task = new_task(ctx.clone()); + task.lifecycle = LifecycleState::DestroyGrace; + task.fire_due_alarms() + .await + .expect("destroy grace alarm should be ignored"); + + assert_eq!(ctx.run_wake_at(), None); + assert_eq!( + crate::actor::internal_storage::load_run_wake_at(ctx.sql()) + .await + .expect("load consumed destroy wake"), + None, + ); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + #[tokio::test(start_paused = true)] async fn sleep_shutdown_preserves_driver_alarm_after_cleanup() { let ctx = new_with_kv( diff --git a/rivetkit-rust/packages/rivetkit-core/tests/workflow_fixture.rs b/rivetkit-rust/packages/rivetkit-core/tests/workflow_fixture.rs new file mode 100644 index 0000000000..1f0f70187c --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/tests/workflow_fixture.rs @@ -0,0 +1,231 @@ +use super::*; +use crate::ActorKey; +use crate::testing::ActorContextHarness; +use rivetkit_actor_persist::versioned::RunWakeAt; +use vbare::OwnedVersionedData; + +fn metadata() -> WorkflowFixtureMetadata { + WorkflowFixtureMetadata { + fixture_name: "typed-roundtrip".to_owned(), + source_rivetkit_version: "2.3.7".to_owned(), + source_workflow_version: "2.3.7".to_owned(), + source_revision: "legacy-revision".to_owned(), + actor_id: "workflow-fixture".to_owned(), + registry_key: "workflowFixture".to_owned(), + internal_schema_version: 1, + fake_clock_seed: 1_723_456_789_000, + generated_id_seed: 42, + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn logical_fixture_roundtrips_every_persisted_workflow_row_type() { + let source_harness = ActorContextHarness::new(); + let source = source_harness.context( + "workflow-fixture", + "workflowFixture", + ActorKey::default(), + "local", + ); + let statements = vec![ + statement( + "INSERT INTO _rivet_runtime (id, last_pushed_alarm, inspector_token, queue_next_id) VALUES (1, ?, ?, ?)", + Some(vec![ + BindParam::Integer(1_723_456_790_000), + BindParam::Null, + BindParam::Integer(10), + ]), + ), + statement( + "INSERT INTO _rivet_meta (key, value) VALUES (?, ?)", + Some(vec![ + BindParam::Text("run_wake_at".to_owned()), + BindParam::Blob( + RunWakeAt::wrap_latest(Some(1_723_456_789_500)) + .serialize_with_embedded_version(1) + .expect("encode logical run wake"), + ), + ]), + ), + statement( + "INSERT INTO _rivet_actor (id, has_initialized, input) VALUES (1, ?, ?)", + Some(vec![BindParam::Integer(1), BindParam::Null]), + ), + statement( + "INSERT INTO _rivet_actor_state (id, state) VALUES (1, ?)", + Some(vec![BindParam::Blob(vec![0, 0xff, 7, 0])]), + ), + statement( + "INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?)", + Some(vec![ + BindParam::Blob(vec![6, 1, 0xff]), + BindParam::Blob(vec![0, 1, 0xff]), + ]), + ), + statement( + "INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?)", + Some(vec![ + BindParam::Blob(vec![6, 1, 0]), + BindParam::Blob(Vec::new()), + ]), + ), + statement( + "INSERT INTO _rivet_queue (id, name, body, created_at) VALUES (?, ?, ?, ?)", + Some(vec![ + BindParam::Integer(9), + BindParam::Text("approval".to_owned()), + BindParam::Blob(vec![0xd9, 0x01, 0x02]), + BindParam::Integer(1_723_456_789_100), + ]), + ), + statement( + "INSERT INTO _rivet_schedule_events (event_id, trigger_at, action, args, kind, cron_expression, timezone, interval_ms, last_started_at, max_history) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + Some(vec![ + BindParam::Text("scheduled-action".to_owned()), + BindParam::Integer(1_723_456_790_000), + BindParam::Text("tick".to_owned()), + BindParam::Blob(vec![0x81, 0x01]), + BindParam::Integer(0), + BindParam::Null, + BindParam::Text("UTC".to_owned()), + BindParam::Null, + BindParam::Integer(1_723_456_788_000), + BindParam::Integer(3), + ]), + ), + statement( + "INSERT INTO _rivet_schedule_history (id, schedule_id, action, scheduled_at, fired_at, finished_at, result, error_group, error_code, error_message, error_metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + Some(vec![ + BindParam::Integer(1), + BindParam::Text("scheduled-action".to_owned()), + BindParam::Text("tick".to_owned()), + BindParam::Integer(1_723_456_788_000), + BindParam::Integer(1_723_456_788_010), + BindParam::Null, + BindParam::Integer(0), + BindParam::Text("schedule".to_owned()), + BindParam::Text("still_running".to_owned()), + BindParam::Null, + BindParam::Blob(vec![0xa0]), + ]), + ), + ]; + source + .sql() + .execute_batch(statements) + .await + .expect("seed source actor database"); + + let fixture = dump_workflow_fixture(&source, metadata()) + .await + .expect("dump source fixture"); + assert_eq!( + fixture + .workflow_rows + .iter() + .map(|row| row.key.clone()) + .collect::>(), + vec![vec![6, 1, 0], vec![6, 1, 0xff]], + "workflow rows must use bytewise canonical ordering and retain the hidden namespace", + ); + let bytes = fixture.encode().expect("encode fixture"); + assert_eq!(&bytes[..2], &[1, 0], "fixture needs a vbare version header"); + let decoded = WorkflowFixture::decode(&bytes).expect("decode fixture"); + assert_eq!(decoded, fixture); + + let target_harness = ActorContextHarness::new(); + let target = target_harness.context( + "workflow-fixture", + "workflowFixture", + ActorKey::default(), + "local", + ); + restore_workflow_fixture(&target, &decoded) + .await + .expect("restore fixture"); + let roundtrip = dump_workflow_fixture(&target, metadata()) + .await + .expect("dump restored fixture"); + assert_eq!(roundtrip, fixture); + assert_eq!(roundtrip.actor_state, Some(vec![0, 0xff, 7, 0])); + assert_eq!(roundtrip.runtime.as_ref().unwrap().inspector_token, None); + let run_wake_row = roundtrip + .meta_rows + .iter() + .find(|row| row.key == "run_wake_at") + .expect("logical run wake metadata row"); + assert_eq!( + RunWakeAt::deserialize_with_embedded_version(&run_wake_row.value) + .expect("decode logical run wake"), + Some(1_723_456_789_500), + ); + assert_eq!(roundtrip.schedule_events[0].args, Some(vec![0x81, 0x01])); + assert_eq!(roundtrip.schedule_events[0].interval_ms, None); + assert_eq!(roundtrip.schedule_history[0].finished_at, None); + assert_eq!(roundtrip.schedule_history[0].error_message, None); +} + +#[test] +fn fixture_rejects_unknown_embedded_version() { + let fixture = WorkflowFixture { + metadata: metadata(), + meta_rows: Vec::new(), + runtime: None, + actor: None, + actor_state: None, + workflow_rows: Vec::new(), + queue_rows: Vec::new(), + schedule_events: Vec::new(), + schedule_history: Vec::new(), + }; + let mut bytes = fixture.encode().expect("encode fixture"); + bytes[..2].copy_from_slice(&2u16.to_le_bytes()); + assert!( + WorkflowFixture::decode(&bytes) + .unwrap_err() + .to_string() + .contains("unsupported workflow fixture version 2") + ); +} + +#[test] +fn fixture_restore_validation_rejects_schema_and_namespace_ambiguity() { + let base = WorkflowFixture { + metadata: metadata(), + meta_rows: vec![WorkflowFixtureMetaRow { + key: "schema_version".to_owned(), + value: 1_i64.to_le_bytes().to_vec(), + }], + runtime: None, + actor: None, + actor_state: None, + workflow_rows: Vec::new(), + queue_rows: Vec::new(), + schedule_events: Vec::new(), + schedule_history: Vec::new(), + }; + validate_fixture_for_restore(&base).expect("valid v1 fixture"); + + let mut wrong_schema = base.clone(); + wrong_schema.meta_rows[0].value = 2_i64.to_le_bytes().to_vec(); + assert!( + validate_fixture_for_restore(&wrong_schema) + .unwrap_err() + .to_string() + .contains("schema metadata mismatch") + ); + + let mut wrong_namespace = base; + wrong_namespace + .workflow_rows + .push(WorkflowFixtureWorkflowRow { + key: vec![6, 2, 0], + value: vec![1], + }); + assert!( + validate_fixture_for_restore(&wrong_namespace) + .unwrap_err() + .to_string() + .contains("escaped the [6, 1] namespace") + ); +} diff --git a/rivetkit-rust/packages/rivetkit/src/event.rs b/rivetkit-rust/packages/rivetkit/src/event.rs index e16520a26d..84f1511062 100644 --- a/rivetkit-rust/packages/rivetkit/src/event.rs +++ b/rivetkit-rust/packages/rivetkit/src/event.rs @@ -180,9 +180,10 @@ impl RuntimeEvent { unreachable!("DisconnectConn is handled by foreign-runtime adapters") } ActorEvent::WorkflowHistoryRequested { .. } - | ActorEvent::WorkflowReplayRequested { .. } => { + | ActorEvent::WorkflowReplayRequested { .. } + | ActorEvent::RunWake { .. } => { unreachable!( - "workflow events are handled by the TypeScript runtime; Rust actors never host workflows" + "workflow/run-wake events are handled by the TypeScript runtime; Rust actors never host workflows" ) } } diff --git a/rivetkit-rust/packages/rivetkit/src/start.rs b/rivetkit-rust/packages/rivetkit/src/start.rs index 32e6c56d57..d454cf932b 100644 --- a/rivetkit-rust/packages/rivetkit/src/start.rs +++ b/rivetkit-rust/packages/rivetkit/src/start.rs @@ -437,6 +437,7 @@ async fn handle_actor_event( ActorEvent::WorkflowReplayRequested { reply, .. } => { reply.send(Err(not_configured("workflow replay"))); } + ActorEvent::RunWake { .. } => {} } Ok(false) diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index c4b07c9249..721468aa6f 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -299,6 +299,7 @@ export declare class ActorContext { takePendingHibernationChanges(): Array dirtyHibernatableConns(): Array saveState(payload: StateDeltaPayload): Promise + beginStateTransaction(timeoutMs?: number | undefined | null): Promise saveStateAndWorkflowBatch(writes: Array): Promise actorId(): string name(): string @@ -313,6 +314,7 @@ export declare class ActorContext { aborted(): boolean runHandlerActive(): boolean restartRunHandler(): void + setRunWakeAt(timestampMs?: number | undefined | null): Promise beginKeepAwake(): number endKeepAwake(regionId: number): void keepAwake(promise: Promise): void @@ -366,6 +368,12 @@ export declare class JsSqliteTransaction { commit(): Promise rollback(): Promise } +export declare class JsActorStateTransaction { + execute(sql: string, params?: Array | undefined | null): Promise + exec(sql: string): Promise + commit(payload: StateDeltaPayload): Promise + rollback(): Promise +} export declare class Kv { get(key: Buffer): Promise put(key: Buffer, value: Buffer): Promise @@ -383,6 +391,7 @@ export declare class Queue { nextBatch(options?: JsQueueNextBatchOptions | undefined | null, signal?: CancellationToken | undefined | null): Promise> waitForNames(names: Array, options?: JsQueueWaitOptions | undefined | null, signal?: CancellationToken | undefined | null): Promise waitForNamesAvailable(names: Array, options?: JsQueueWaitOptions | undefined | null, signal?: CancellationToken | undefined | null): Promise + completePersisted(messageId: bigint, expectedName: string, response?: Buffer | undefined | null): Promise enqueueAndWait(name: string, body: Buffer, options?: JsQueueEnqueueAndWaitOptions | undefined | null, signal?: CancellationToken | undefined | null): Promise tryNext(options?: JsQueueTryNextOptions | undefined | null): QueueMessage | null tryNextBatch(options?: JsQueueTryNextBatchOptions | undefined | null): Array diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.js b/rivetkit-typescript/packages/rivetkit-napi/index.js index 18e5a172db..c8d90d43e8 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.js +++ b/rivetkit-typescript/packages/rivetkit-napi/index.js @@ -310,7 +310,7 @@ if (!nativeBinding) { throw new Error(`Failed to load native binding`) } -const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, NapiActorFactory, CancellationToken, ConnHandle, JsNativeDatabase, JsSqliteTransaction, Kv, Queue, QueueMessage, CoreRegistry, Schedule, WebSocket } = nativeBinding +const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, NapiActorFactory, CancellationToken, ConnHandle, JsNativeDatabase, JsSqliteTransaction, JsActorStateTransaction, Kv, Queue, QueueMessage, CoreRegistry, Schedule, WebSocket } = nativeBinding module.exports.ActorContext = ActorContext module.exports.decodeInspectorRequest = decodeInspectorRequest @@ -320,6 +320,7 @@ module.exports.CancellationToken = CancellationToken module.exports.ConnHandle = ConnHandle module.exports.JsNativeDatabase = JsNativeDatabase module.exports.JsSqliteTransaction = JsSqliteTransaction +module.exports.JsActorStateTransaction = JsActorStateTransaction module.exports.Kv = Kv module.exports.Queue = Queue module.exports.QueueMessage = QueueMessage diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index 437e2e8b6f..c4a91f0f1a 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -27,7 +27,7 @@ use tokio_util::sync::CancellationToken as CoreCancellationToken; use crate::actor_factory::BridgeRivetErrorContext; use crate::connection::ConnHandle; -use crate::database::JsNativeDatabase; +use crate::database::{JsActorStateTransaction, JsNativeDatabase, transaction_timeout}; use crate::kv::Kv; use crate::queue::Queue; use crate::schedule::Schedule; @@ -413,6 +413,19 @@ impl ActorContext { .map_err(napi_anyhow_error) } + #[napi] + pub async fn begin_state_transaction( + &self, + timeout_ms: Option, + ) -> napi::Result { + let timeout = timeout_ms.map(transaction_timeout).transpose()?; + self.inner + .begin_state_transaction(timeout) + .await + .map(JsActorStateTransaction::new) + .map_err(napi_anyhow_error) + } + #[napi] pub async fn save_state_and_workflow_batch( &self, @@ -518,6 +531,14 @@ impl ActorContext { self.shared.run_restart().map_err(napi_anyhow_error) } + #[napi] + pub async fn set_run_wake_at(&self, timestamp_ms: Option) -> napi::Result<()> { + self.inner + .set_run_wake_at(timestamp_ms) + .await + .map_err(napi_anyhow_error) + } + #[napi] pub fn begin_keep_awake(&self) -> u32 { self.shared.begin_keep_awake(self.inner.keep_awake_region()) diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/database.rs b/rivetkit-typescript/packages/rivetkit-napi/src/database.rs index 792105ee29..fb160c1bfa 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/database.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/database.rs @@ -2,12 +2,14 @@ use std::time::Duration; use napi::bindgen_prelude::Buffer; use napi_derive::napi; +use rivetkit_core::ActorStateTransaction as CoreActorStateTransaction; use rivetkit_core::sqlite::{ BindParam, ColumnValue, ExecuteResult as CoreExecuteResult, QueryResult as CoreQueryResult, SqliteBatchStatement as CoreSqliteBatchStatement, SqliteDb as CoreSqliteDb, SqliteTransaction as CoreSqliteTransaction, }; +use crate::actor_context::{StateDeltaPayload, state_deltas_from_payload}; use crate::{NapiInvalidArgument, napi_anyhow_error}; #[napi] #[derive(Clone)] @@ -22,6 +24,18 @@ pub struct JsSqliteTransaction { transaction: CoreSqliteTransaction, } +#[napi] +#[derive(Clone)] +pub struct JsActorStateTransaction { + transaction: CoreActorStateTransaction, +} + +impl JsActorStateTransaction { + pub(crate) fn new(transaction: CoreActorStateTransaction) -> Self { + Self { transaction } + } +} + impl JsNativeDatabase { pub(crate) fn new(db: CoreSqliteDb, actor_id: Option) -> Self { tracing::debug!( @@ -245,7 +259,49 @@ impl JsSqliteTransaction { } } -fn transaction_timeout(timeout_ms: f64) -> napi::Result { +#[napi] +impl JsActorStateTransaction { + #[napi] + pub async fn execute( + &self, + sql: String, + params: Option>, + ) -> napi::Result { + let params = params.map(js_bind_params_to_core).transpose()?; + self.transaction + .execute(sql, params) + .await + .map(core_execute_result_to_js) + .map_err(crate::napi_anyhow_error) + } + + #[napi] + pub async fn exec(&self, sql: String) -> napi::Result { + self.transaction + .exec(sql) + .await + .map(core_query_result_to_js) + .map_err(crate::napi_anyhow_error) + } + + #[napi] + pub async fn commit(&self, payload: StateDeltaPayload) -> napi::Result<()> { + self.transaction + .commit(state_deltas_from_payload(payload)) + .await + .map_err(crate::napi_anyhow_error) + } + + #[napi] + pub async fn rollback(&self) -> napi::Result<()> { + self.transaction + .rollback() + .await + .map_err(crate::napi_anyhow_error) + } +} + +pub(crate) fn transaction_timeout(timeout_ms: f64) -> napi::Result { if !timeout_ms.is_finite() || timeout_ms <= 0.0 { return Err(napi_anyhow_error( NapiInvalidArgument { diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs index 9da9f0847e..378e1bb19b 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -730,6 +730,12 @@ pub(crate) async fn dispatch_event( call_workflow_replay(&callback, &ctx, entry_id).await }); } + ActorEvent::RunWake { reply } => { + reply.send( + ctx.restart_run_handler() + .map_err(|error| anyhow::anyhow!(error.to_string())), + ); + } } } diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/queue.rs b/rivetkit-typescript/packages/rivetkit-napi/src/queue.rs index 2d53c83e9f..0cd5331ae0 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/queue.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/queue.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use napi::bindgen_prelude::Buffer; +use napi::bindgen_prelude::{BigInt, Buffer}; use napi_derive::napi; use parking_lot::Mutex; use rivetkit_core::{ @@ -170,6 +170,33 @@ impl Queue { .map_err(napi_anyhow_error) } + #[napi] + pub async fn complete_persisted( + &self, + message_id: BigInt, + expected_name: String, + response: Option, + ) -> napi::Result { + let (negative, message_id, lossless) = message_id.get_u64(); + if negative || !lossless { + return Err(napi_anyhow_error( + NapiInvalidArgument { + argument: "messageId".to_owned(), + reason: "must be a non-negative 64-bit bigint".to_owned(), + } + .build(), + )); + } + self.inner + .complete_persisted_message( + message_id, + &expected_name, + response.map(|value| value.to_vec()), + ) + .await + .map_err(napi_anyhow_error) + } + #[napi] pub async fn enqueue_and_wait( &self, diff --git a/rivetkit-typescript/packages/rivetkit-wasm/index.d.ts b/rivetkit-typescript/packages/rivetkit-wasm/index.d.ts index f95b9d10e8..36ba556061 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-wasm/index.d.ts @@ -16,6 +16,7 @@ export class ActorContext { connectConn(params: Uint8Array, request: any): Promise; requestSave(opts: any): void; registerTask(promise: Promise): void; + beginStateTransaction(timeout_ms?: number | null): Promise; runtimeState(): any; endKeepAwake(region_id: number): void; beginKeepAwake(): number; @@ -121,6 +122,11 @@ export class Queue { ): Promise; inspectMessages(): Promise>; waitForNamesAvailable(names: any, options: any): Promise; + completePersisted( + message_id: bigint, + expected_name: string, + response?: Uint8Array | null, + ): Promise; send(name: string, body: Uint8Array): Promise; maxSize(): number; reset(): Promise; @@ -187,6 +193,15 @@ export class SqliteTransaction { rollback(): Promise; } +export class ActorStateTransaction { + private constructor(); + free(): void; + exec(sql: string): Promise; + execute(sql: string, params: any): Promise; + commit(payload: any): Promise; + rollback(): Promise; +} + export class WebSocketHandle { private constructor(); free(): void; diff --git a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs index ca6147e4ea..781cad21eb 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs @@ -574,7 +574,9 @@ async fn run_actor_adapter(callbacks: WasmCallbacks, start: ActorStart) -> Resul ); } preamble?; - start_run_handler(&callbacks, &ctx); + if callbacks.run.is_some() { + start_run_handler(&callbacks, &ctx)?; + } while let Some(event) = events.recv().await { dispatch_event(&callbacks, &ctx, event).await; @@ -583,10 +585,13 @@ async fn run_actor_adapter(callbacks: WasmCallbacks, start: ActorStart) -> Resul Ok(()) } -fn start_run_handler(callbacks: &WasmCallbacks, ctx: &WasmActorContext) { +fn start_run_handler(callbacks: &WasmCallbacks, ctx: &WasmActorContext) -> Result<()> { let Some(callback) = callbacks.run.clone() else { - return; + return Err(anyhow!("wasm run handler is not configured")); }; + if ctx.inner.run_handler_active() { + return Err(anyhow!("wasm run handler is already active")); + } let ctx = ctx.clone(); ctx.inner.begin_run_handler(); spawn_local(async move { @@ -602,6 +607,7 @@ fn start_run_handler(callbacks: &WasmCallbacks, ctx: &WasmActorContext) { } ctx.inner.end_run_handler(); }); + Ok(()) } async fn run_preamble( @@ -801,6 +807,9 @@ async fn dispatch_event(callbacks: &WasmCallbacks, ctx: &WasmActorContext, event } reply.send(result); } + ActorEvent::RunWake { reply } => { + reply.send(start_run_handler(callbacks, ctx)); + } ActorEvent::HttpRequest { request, reply } => { let callback = callbacks.on_request.clone(); let ctx = ctx.clone(); @@ -1292,6 +1301,19 @@ impl WasmActorContext { .map_err(anyhow_to_js_error) } + #[wasm_bindgen(js_name = beginStateTransaction)] + pub async fn begin_state_transaction( + &self, + timeout_ms: Option, + ) -> Result { + let timeout = timeout_ms.map(transaction_timeout).transpose()?; + self.inner + .begin_state_transaction(timeout) + .await + .map(|inner| WasmActorStateTransaction { inner }) + .map_err(anyhow_to_js_error) + } + #[wasm_bindgen(js_name = saveStateAndWorkflowBatch)] pub async fn save_state_and_workflow_batch(&self, writes: JsValue) -> Result<(), JsValue> { let writes: Vec = serde_wasm_bindgen::from_value(writes)?; @@ -1405,6 +1427,17 @@ impl WasmActorContext { .map_err(anyhow_to_js_error) } + #[wasm_bindgen(js_name = setRunWakeAt)] + pub async fn set_run_wake_at(&self, timestamp_ms: Option) -> Result<(), JsValue> { + let timestamp_ms = timestamp_ms + .filter(|value| value.is_finite()) + .map(|value| value.trunc() as i64); + self.inner + .set_run_wake_at(timestamp_ms) + .await + .map_err(anyhow_to_js_error) + } + #[wasm_bindgen] pub fn sleep(&self) -> Result<(), JsValue> { self.inner.sleep().map_err(anyhow_to_js_error) @@ -1492,8 +1525,8 @@ impl WasmActorContext { } #[wasm_bindgen(js_name = restartRunHandler)] - pub fn restart_run_handler(&self) { - start_run_handler(&self.callbacks, self); + pub fn restart_run_handler(&self) -> Result<(), JsValue> { + start_run_handler(&self.callbacks, self).map_err(anyhow_to_js_error) } #[wasm_bindgen(js_name = beginKeepAwake)] @@ -1985,6 +2018,19 @@ impl WasmQueue { Ok(()) } + #[wasm_bindgen(js_name = completePersisted)] + pub async fn complete_persisted( + &self, + message_id: u64, + expected_name: String, + response: Option>, + ) -> Result { + self.inner + .complete_persisted_message(message_id, &expected_name, response) + .await + .map_err(anyhow_to_js_error) + } + #[wasm_bindgen(js_name = enqueueAndWait)] pub async fn enqueue_and_wait( &self, @@ -2236,6 +2282,43 @@ impl WasmSqliteTransaction { } } +#[wasm_bindgen(js_name = ActorStateTransaction)] +pub struct WasmActorStateTransaction { + inner: rivetkit_core::ActorStateTransaction, +} + +#[wasm_bindgen(js_class = ActorStateTransaction)] +impl WasmActorStateTransaction { + #[wasm_bindgen] + pub async fn exec(&self, sql: String) -> Result { + self.inner + .exec(sql) + .await + .map(query_result_to_js) + .map_err(anyhow_to_js_error) + } + + #[wasm_bindgen] + pub async fn execute(&self, sql: String, params: JsValue) -> Result { + self.inner + .execute(sql, bind_params_from_js(params)?) + .await + .map(execute_result_to_js) + .map_err(anyhow_to_js_error) + } + + #[wasm_bindgen] + pub async fn commit(&self, payload: JsValue) -> Result<(), JsValue> { + let deltas = state_delta_payload_from_js(payload).map_err(anyhow_to_js_error)?; + self.inner.commit(deltas).await.map_err(anyhow_to_js_error) + } + + #[wasm_bindgen] + pub async fn rollback(&self) -> Result<(), JsValue> { + self.inner.rollback().await.map_err(anyhow_to_js_error) + } +} + fn transaction_timeout(timeout_ms: f64) -> Result { if !timeout_ms.is_finite() || timeout_ms <= 0.0 { return Err(js_error( diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-raw.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-raw.ts index 891e58879b..d5bdb6e3c6 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-raw.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-raw.ts @@ -722,6 +722,7 @@ export const actorRuntimeSocketWithoutDb = actor({ getActorRuntimeSocketPath: async (c) => { return (await c.actorRuntimeSocket()).path; }, + destroy: (c) => c.destroy(), }, options: { enableActorRuntimeSocket: true, diff --git a/rivetkit-typescript/packages/rivetkit/package.json b/rivetkit-typescript/packages/rivetkit/package.json index 4f52492cf0..eb2e4c2c0a 100644 --- a/rivetkit-typescript/packages/rivetkit/package.json +++ b/rivetkit-typescript/packages/rivetkit/package.json @@ -136,6 +136,16 @@ "default": "./dist/tsup/inspector/mod.cjs" } }, + "./experimental/inspector/workflow": { + "import": { + "types": "./dist/tsup/inspector/workflow.d.ts", + "default": "./dist/tsup/inspector/workflow.js" + }, + "require": { + "types": "./dist/tsup/inspector/workflow.d.cts", + "default": "./dist/tsup/inspector/workflow.cjs" + } + }, "./inspector-tab": { "import": { "types": "./dist/tsup/inspector-tab/mod.d.ts", @@ -181,7 +191,7 @@ "./dist/tsup/chunk-*.cjs" ], "scripts": { - "build": "tsup src/mod.ts src/client/mod.ts src/common/log.ts src/common/websocket.ts src/actor/errors.ts src/utils.ts src/workflow/mod.ts src/test/mod.ts src/inspector/mod.ts src/inspector-tab/mod.ts src/db/mod.ts src/db/drizzle.ts src/dynamic/mod.ts src/unstable/migrations.ts && tsup src/agent-os/index.ts --no-clean --out-dir dist/tsup/agent-os", + "build": "tsup src/mod.ts src/client/mod.ts src/common/log.ts src/common/websocket.ts src/actor/errors.ts src/utils.ts src/workflow/mod.ts src/test/mod.ts src/inspector/mod.ts src/inspector/workflow.ts src/inspector-tab/mod.ts src/db/mod.ts src/db/drizzle.ts src/dynamic/mod.ts src/unstable/migrations.ts && tsup src/agent-os/index.ts --no-clean --out-dir dist/tsup/agent-os && node scripts/check-built-commonjs.mjs", "build:browser": "tsup --config tsup.browser.config.ts", "check-types": "tsc --noEmit", "lint": "biome check . && pnpm run check:test-skips && pnpm run check:wait-for-comments", diff --git a/rivetkit-typescript/packages/rivetkit/scripts/check-built-commonjs.mjs b/rivetkit-typescript/packages/rivetkit/scripts/check-built-commonjs.mjs new file mode 100644 index 0000000000..8c0bd53d93 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/scripts/check-built-commonjs.mjs @@ -0,0 +1,21 @@ +import { readFileSync, readdirSync } from "node:fs"; +import { createRequire } from "node:module"; + +const outputDir = new URL("../dist/tsup/", import.meta.url); +for (const relativePath of readdirSync(outputDir, { recursive: true })) { + if (!relativePath.endsWith(".cjs")) continue; + const source = readFileSync(new URL(relativePath, outputDir), "utf8"); + if ( + source.includes("import.meta.url") || + /\bimport_meta\d*\.url\b/.test(source) + ) { + throw new Error(`${relativePath} contains an invalid CommonJS import.meta URL`); + } +} + +const require = createRequire(import.meta.url); +const rivetkit = require("../dist/tsup/mod.cjs"); + +if (typeof rivetkit.actor !== "function") { + throw new Error("CommonJS build does not export actor()"); +} diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/config.test.ts b/rivetkit-typescript/packages/rivetkit/src/actor/config.test.ts index 22c7f61c10..f7fc4ad62e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/config.test.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/config.test.ts @@ -1,5 +1,12 @@ import { describe, expect, test } from "vitest"; -import { ActorOptionsSchema } from "./config"; +import { + ActorOptionsSchema, + type DefineRunHandlerOptions, + defineRunHandler, + getRunFunction, + getRunInspectorKind, + getRunMetadata, +} from "./config"; describe("ActorOptionsSchema", () => { test("keeps the Actor Runtime Socket opt-in", () => { @@ -12,3 +19,60 @@ describe("ActorOptionsSchema", () => { ).toBe(true); }); }); + +describe("defineRunHandler", () => { + test("preserves the callable and exposes static metadata without creating an inspector", () => { + let inspectorCreates = 0; + const run = async (value: number): Promise => String(value); + const defined = defineRunHandler(run, { + name: "Durable import", + icon: "diagram-project", + inspectorKind: "workflow", + createInspector: () => { + inspectorCreates += 1; + return { + inspector: { + workflow: { + getHistory: () => null, + getState: async () => null, + onHistoryUpdated: () => () => {}, + replayFromStep: async () => null, + }, + }, + }; + }, + }); + + const exactType: (value: number) => Promise = defined; + expect(exactType).toBe(run); + expect(getRunFunction(defined)).toBe(run); + expect(getRunMetadata(defined)).toEqual({ + name: "Durable import", + icon: "diagram-project", + }); + expect(getRunInspectorKind(defined)).toBe("workflow"); + expect(inspectorCreates).toBe(0); + }); + + test("requires static inspector metadata and its factory together", () => { + expect(() => + defineRunHandler(async () => {}, { + inspectorKind: "workflow", + } as unknown as DefineRunHandlerOptions), + ).toThrow("requires inspectorKind and createInspector together"); + expect(() => + defineRunHandler(async () => {}, { + createInspector: () => ({ + inspector: { + workflow: { + getHistory: () => null, + getState: async () => null, + onHistoryUpdated: () => () => {}, + replayFromStep: async () => null, + }, + }, + }), + } as unknown as DefineRunHandlerOptions), + ).toThrow("requires inspectorKind and createInspector together"); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/config.ts b/rivetkit-typescript/packages/rivetkit/src/actor/config.ts index eae0b73744..6df2030d9a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/config.ts @@ -6,6 +6,7 @@ import type { } from "@/common/database/config"; import type { UniversalWebSocket } from "@/common/websocket-interface"; import type { Registry } from "@/registry"; +import type { WorkflowState } from "@/inspector/workflow"; import { flattenActionHandlers } from "./actions"; import type { BaseActorDefinition } from "./definition"; import type { @@ -64,7 +65,7 @@ type ActorKvListOptions< type ActorClientFor = T extends Registry ? Client : T; /** - * @deprecated Actor KV is deprecated. Use embedded SQLite (`c.db` / `c.sql`) + * @deprecated Actor KV is deprecated. Use embedded SQLite (`c.db`) * or actor state instead. */ export interface ActorKv { @@ -332,6 +333,16 @@ export interface ActorQueue< names: readonly TName[], opts?: QueueWaitOptions, ): Promise; + /** @experimental */ + waitForAvailable>( + names?: readonly TName[], + opts?: Omit, "completable">, + ): Promise; + /** @experimental */ + complete>( + message: { id: bigint; name: TName }, + ...args: QueueCompleteArgsForName + ): Promise; enqueueAndWait>( name: TName, body: QueueMessageForName["body"], @@ -389,7 +400,7 @@ export interface ActorContext< state: TState; vars: TVars; /** - * @deprecated Actor KV is deprecated. Use embedded SQLite (`db` / `sql`) + * @deprecated Actor KV is deprecated. Use embedded SQLite (`c.db`) * or actor state instead. */ readonly kv: ActorKv; @@ -397,6 +408,8 @@ export interface ActorContext< readonly schedule: ActorSchedule; readonly cron: ActorCron; readonly queue: ActorQueue; + /** @experimental */ + readonly run: ActorRun; readonly actorId: string; readonly name: string; readonly key: string[]; @@ -434,6 +447,11 @@ export interface ActorContext< [key: string]: any; } +export interface ActorRun { + /** @experimental Sets or clears the durable deadline that restarts this actor's run handler. */ + setWakeAt(timestamp: number | null): Promise; +} + export type ActionContext< TState, TConnParams, @@ -821,20 +839,27 @@ const zActionTree = z export type InspectorUnsubscribe = () => void; +/** @experimental */ export interface WorkflowInspectorConfig { getHistory: () => THistory | null; + getState?: () => Promise; onHistoryUpdated?: ( listener: (history: THistory) => void, ) => InspectorUnsubscribe; replayFromStep?: (entryId?: string) => Promise; } +/** @experimental */ export interface RunInspectorConfig { workflow?: WorkflowInspectorConfig; } const WorkflowInspectorConfigSchema = z.object({ getHistory: zFunction["getHistory"]>(), + getState: + zFunction< + NonNullable["getState"]> + >().optional(), onHistoryUpdated: zFunction< NonNullable["onHistoryUpdated"]> @@ -981,9 +1006,70 @@ type AnyRunConfig = RunConfig< any >; -export const RUN_FUNCTION_CONFIG_SYMBOL = Symbol.for( - "rivetkit.run_function_config", -); +const RUN_FUNCTION_CONFIG_SYMBOL = Symbol.for("rivetkit.run_function_config"); + +export type RunInspectorKind = "workflow"; + +export interface RunWithInactiveOptions { + /** Restart the run handler before releasing the exclusive gate on success. */ + restartOnSuccess?: boolean; +} + +/** @experimental */ +export interface RunControl { + run: { + /** + * Executes only while no run handler is active or waiting to start. + * + * This is a non-blocking acquisition: concurrent replay/control requests + * fail instead of waiting behind an active run. + */ + withInactive( + options: RunWithInactiveOptions, + callback: () => T | Promise, + ): Promise; + }; +} + +/** @experimental */ +export interface RunInspectorFactoryContext { + actorId: string; + control: RunControl; +} + +/** @experimental */ +export interface RunInspectorFactoryResult { + inspector: { + workflow: Required> & { + getState: () => Promise; + }; + }; + /** Releases actor-specific listeners and encoded history. */ + dispose?: () => void; +} + +interface RunHandlerDisplayOptions { + name?: string; + icon?: string; +} + +/** @experimental */ +export type DefineRunHandlerOptions = + RunHandlerDisplayOptions & + ( + | { + inspectorKind?: never; + createInspector?: never; + } + | { + /** Static metadata used to install Inspector callbacks at registry build time. */ + inspectorKind: RunInspectorKind; + /** Creates one Inspector adapter for each live actor. */ + createInspector: ( + context: RunInspectorFactoryContext, + ) => RunInspectorFactoryResult; + } + ); interface RunFunctionConfig { name?: string; @@ -992,12 +1078,69 @@ interface RunFunctionConfig { inspectorFactory?: (actor: unknown) => RunInspectorConfig | undefined; /** Release any per-actor inspector state held for this actor id. */ disposeInspector?: (actorId: string) => void; + inspectorKind?: RunInspectorKind; + createInspector?: ( + context: RunInspectorFactoryContext, + ) => RunInspectorFactoryResult; } type RunFunctionWithConfig = ((...args: any[]) => any) & { [RUN_FUNCTION_CONFIG_SYMBOL]?: RunFunctionConfig; }; +/** + * @experimental Adds supported RivetKit run metadata while preserving the function's exact + * callable type. + */ +export function defineRunHandler< + TRun extends (...args: any[]) => any, + THistory = unknown, +>(run: TRun, options: DefineRunHandlerOptions): TRun { + if ( + (options.inspectorKind === undefined) !== + (options.createInspector === undefined) + ) { + throw new TypeError( + "defineRunHandler requires inspectorKind and createInspector together", + ); + } + + Object.defineProperty(run, RUN_FUNCTION_CONFIG_SYMBOL, { + configurable: false, + enumerable: false, + writable: false, + value: { + name: options.name, + icon: options.icon, + inspectorKind: options.inspectorKind, + createInspector: + options.createInspector as RunFunctionConfig["createInspector"], + } satisfies RunFunctionConfig, + }); + + return run; +} + +/** @internal */ +export function getRunInspectorKind( + run: ((...args: any[]) => any) | AnyRunConfig | undefined, +): RunInspectorKind | undefined { + if (!run || typeof run !== "function") return undefined; + return (run as RunFunctionWithConfig)[RUN_FUNCTION_CONFIG_SYMBOL] + ?.inspectorKind; +} + +/** @internal */ +export function createRunInspector( + run: ((...args: any[]) => any) | AnyRunConfig | undefined, + context: RunInspectorFactoryContext, +): RunInspectorFactoryResult | undefined { + if (!run || typeof run !== "function") return undefined; + return (run as RunFunctionWithConfig)[ + RUN_FUNCTION_CONFIG_SYMBOL + ]?.createInspector?.(context); +} + // Run can be either a function or an object with name/icon/run const zRunHandler = z.union([zFunction(), RunConfigSchema]).optional(); @@ -1042,6 +1185,21 @@ export function getRunInspectorConfig( return run.inspector; } +/** @internal */ +export function hasRunInspectorConfig( + run: ((...args: any[]) => any) | AnyRunConfig | undefined, +): boolean { + if (!run) return false; + if (typeof run !== "function") return run.inspector !== undefined; + const config = (run as RunFunctionWithConfig)[RUN_FUNCTION_CONFIG_SYMBOL]; + return ( + config?.inspectorKind !== undefined || + config?.createInspector !== undefined || + config?.inspector !== undefined || + config?.inspectorFactory !== undefined + ); +} + /** Release per-actor inspector state for a destroyed actor, if the run handler registered a disposer. */ export function disposeRunInspector( run: ((...args: any[]) => any) | AnyRunConfig | undefined, diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts b/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts index d373f76e2c..7ab0df3e68 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts @@ -14,6 +14,7 @@ export type { UniversalWebSocket, } from "@/common/websocket-interface"; export type * from "./config"; +export { defineRunHandler } from "./config"; export type { ActionContextOf, BeforeActionResponseContextOf, @@ -57,4 +58,12 @@ export { UserError, type UserErrorOptions, } from "./errors"; -export { event, queue, type Type } from "./schema"; +export { + event, + type EventSchemaConfig, + type InferEventArgs, + type InferSchemaMap, + queue, + type QueueSchemaConfig, + type Type, +} from "./schema"; diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts index 250ce081a4..1f5f7befed 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts @@ -48,20 +48,11 @@ type ActorActionMap = { : never; }; -type ActionsOf = - AD extends BaseActorDefinition< - any, - any, - any, - any, - any, - any, - any, - any, - infer R - > - ? R - : never; +type ActionsOf = AD["config"] extends { + actions?: infer R; +} + ? R + : never; export interface ActorGatewayOptions { skipReadyWait?: boolean; diff --git a/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts b/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts index c94e44228f..73a9b954e2 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts @@ -9,7 +9,7 @@ export interface ActorMetricsLike { export type InferDatabaseClient = DBProvider extends DatabaseProvider ? Awaited> - : never; + : RawAccess; export type SqliteBindings = unknown[] | Record; @@ -165,7 +165,17 @@ export type RawAccess = { /** Runs a callback in an isolated SQLite transaction. */ transaction: ( callback: (tx: RawAccess) => Promise | T, - options?: { timeout?: number }, + options?: { + timeout?: number; + /** @experimental */ + experimental?: { + /** + * Atomically includes actor and hibernatable connection state. + * Only single-statement `execute` calls are supported in the transaction. + */ + includeState?: boolean; + }; + }, ) => Promise; /** * Returns native SQLite metrics when the active runtime supports them. diff --git a/rivetkit-typescript/packages/rivetkit/src/common/database/mod.test.ts b/rivetkit-typescript/packages/rivetkit/src/common/database/mod.test.ts index 1c322bcd4e..b6cef144c5 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/database/mod.test.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/database/mod.test.ts @@ -8,7 +8,7 @@ import type { SqliteExecuteResult, SqliteTransactionDatabase, } from "./config"; -import { db } from "./mod"; +import { db, registerNativeStateTransactionOpener } from "./mod"; let logLines: string[]; @@ -16,6 +16,7 @@ class FakeSqliteDatabase implements SqliteDatabase { failSql = new Map(); executeCalls: { sql: string; params?: SqliteBindings }[] = []; transactionTimeouts: Array = []; + stateTransactionTimeouts: Array = []; async exec(): Promise {} @@ -43,6 +44,22 @@ class FakeSqliteDatabase implements SqliteDatabase { }; } + async beginStateTransaction( + timeoutMs?: number, + ): Promise { + this.stateTransactionTimeouts.push(timeoutMs); + this.record("BEGIN_STATE"); + return { + exec: async () => {}, + execute: async (sql, params) => { + this.record(sql, params); + return emptyResult(); + }, + commit: async () => this.record("COMMIT"), + rollback: async () => this.record("ROLLBACK"), + }; + } + async executeBatch( statements: Array<{ sql: string; params?: SqliteBindings }>, ): Promise { @@ -81,7 +98,8 @@ function emptyResult(): SqliteExecuteResult { } function testProviderContext( - database: SqliteDatabase, + database: FakeSqliteDatabase, + includeStateTransactions = false, ): DatabaseProviderContext { return { actorId: "actor-a", @@ -91,7 +109,13 @@ function testProviderContext( batchDelete: async () => {}, deleteRange: async () => {}, }, - nativeDatabaseProvider: { open: async () => database }, + nativeDatabaseProvider: includeStateTransactions + ? registerNativeStateTransactionOpener( + { open: async () => database }, + async (timeoutMs?: number) => + await database.beginStateTransaction(timeoutMs), + ) + : { open: async () => database }, }; } @@ -193,6 +217,81 @@ describe("db", () => { ]); }); + test("uses the state-aware transaction bridge when explicitly requested", async () => { + const nativeDb = new FakeSqliteDatabase(); + const client = await db().createClient( + testProviderContext(nativeDb, true), + ); + + await client.transaction( + async (tx) => { + await tx.execute( + "INSERT INTO items(value) VALUES (?)", + "inside", + ); + }, + { + timeout: 120_000, + experimental: { includeState: true }, + }, + ); + + expect(nativeDb.transactionTimeouts).toEqual([]); + expect(nativeDb.stateTransactionTimeouts).toEqual([120_000]); + expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([ + "BEGIN_STATE", + "INSERT INTO items(value) VALUES (?)", + "COMMIT", + ]); + }); + + test("rolls back state-aware transactions when the callback throws", async () => { + const nativeDb = new FakeSqliteDatabase(); + const client = await db().createClient( + testProviderContext(nativeDb, true), + ); + + await expect( + client.transaction( + async (tx) => { + await tx.execute( + "INSERT INTO items(value) VALUES (?)", + "inside", + ); + throw new Error("callback failed"); + }, + { experimental: { includeState: true } }, + ), + ).rejects.toThrow("callback failed"); + expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([ + "BEGIN_STATE", + "INSERT INTO items(value) VALUES (?)", + "ROLLBACK", + ]); + }); + + test("rejects nested state-aware transactions", async () => { + const nativeDb = new FakeSqliteDatabase(); + const client = await db().createClient( + testProviderContext(nativeDb, true), + ); + + await expect( + client.transaction( + async (tx) => { + await tx.transaction(async () => {}, { + experimental: { includeState: true }, + }); + }, + { experimental: { includeState: true } }, + ), + ).rejects.toThrow("not supported for nested transactions"); + expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([ + "BEGIN_STATE", + "ROLLBACK", + ]); + }); + test("validates transaction timeouts", async () => { const client = await db().createClient( testProviderContext(new FakeSqliteDatabase()), diff --git a/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts b/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts index f4a9584811..ac6adc31e6 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts @@ -1,6 +1,7 @@ import { getLogger } from "@/common/log"; import type { DatabaseProvider, + NativeDatabaseProvider, RawAccess, SqliteDatabase, SqliteTransactionDatabase, @@ -20,6 +21,28 @@ interface DatabaseFactoryConfig { warnOnManualTransactions?: boolean; } +const builtInDatabaseProviders = new WeakSet(); +const nativeStateTransactionOpeners = new WeakMap< + NativeDatabaseProvider, + (timeoutMs?: number) => Promise +>(); + +/** @internal */ +export function registerNativeStateTransactionOpener< + T extends NativeDatabaseProvider, +>( + provider: T, + opener: (timeoutMs?: number) => Promise, +): T { + nativeStateTransactionOpeners.set(provider, opener); + return provider; +} + +/** @internal */ +export function isBuiltInDatabaseProvider(provider: object): boolean { + return builtInDatabaseProviders.has(provider); +} + function hasMultipleStatements(query: string): boolean { const trimmed = query.trim().replace(/;+$/, "").trimEnd(); return trimmed.includes(";"); @@ -29,7 +52,7 @@ export function db({ onMigrate, warnOnManualTransactions = true, }: DatabaseFactoryConfig = {}): DatabaseProvider { - return { + const provider: DatabaseProvider = { createClient: async (ctx) => { const nativeDatabaseProvider = ctx.nativeDatabaseProvider; if (!nativeDatabaseProvider) { @@ -128,12 +151,32 @@ export function db({ }, transaction: async ( callback: (tx: RawAccess) => Promise | T, - options?: { timeout?: number }, + options?: { + timeout?: number; + experimental?: { includeState?: boolean }; + }, ): Promise => { validateTransactionTimeout(options?.timeout); - const transaction = await db.beginTransaction( - options?.timeout, - ); + if (transactionScoped && options?.experimental?.includeState) { + throw new Error( + "experimental.includeState is not supported for nested transactions", + ); + } + const transaction = options?.experimental?.includeState + ? await (() => { + const beginStateTransaction = + ctx.nativeDatabaseProvider && + nativeStateTransactionOpeners.get( + ctx.nativeDatabaseProvider, + ); + if (!beginStateTransaction) { + throw new Error( + "experimental.includeState is only supported by RivetKit's embedded database provider", + ); + } + return beginStateTransaction(options?.timeout); + })() + : await db.beginTransaction(options?.timeout); const tx = createClient(transaction, true); try { const result = await callback(tx); @@ -167,6 +210,8 @@ export function db({ } }, }; + builtInDatabaseProviders.add(provider); + return provider; } function rowToObject>( diff --git a/rivetkit-typescript/packages/rivetkit/src/common/database/native-database.ts b/rivetkit-typescript/packages/rivetkit/src/common/database/native-database.ts index c3965535ad..174976b493 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/database/native-database.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/database/native-database.ts @@ -82,6 +82,9 @@ export interface JsNativeDatabaseLike { statements: NativeBatchStatement[], ): Promise; beginTransaction(timeoutMs?: number): Promise; + beginStateTransaction?( + timeoutMs?: number, + ): Promise; query( sql: string, params?: NativeBindParam[] | null, @@ -95,6 +98,12 @@ export interface JsNativeDatabaseLike { close(): Promise; } +export type StateAwareSqliteDatabase = SqliteDatabase & { + beginStateTransaction( + timeoutMs?: number, + ): Promise; +}; + export interface JsNativeTransactionLike { exec(sql: string): Promise; execute( @@ -312,7 +321,7 @@ class NativeCloseGate { export function wrapJsNativeDatabase( database: JsNativeDatabaseLike, -): SqliteDatabase { +): StateAwareSqliteDatabase { const gate = new NativeCloseGate(); let closePromise: Promise | undefined; let lastInsertRowId: number | null = null; @@ -435,6 +444,27 @@ export function wrapJsNativeDatabase( } }); }, + async beginStateTransaction( + timeoutMs?: number, + ): Promise { + if (!database.beginStateTransaction) { + throw new Error("actor state transactions are not configured"); + } + const release = gate.enter(); + let transaction: JsNativeTransactionLike; + try { + transaction = await database.beginStateTransaction(timeoutMs); + } catch (error) { + enrichNativeDatabaseError(database, error); + } finally { + release(); + } + return wrapTransaction(database, transaction, gate, (result) => { + if (result.lastInsertRowId !== undefined) { + lastInsertRowId = result.lastInsertRowId; + } + }); + }, async run(sql: string, params?: SqliteBindings): Promise { await executeNative(sql, params); }, diff --git a/rivetkit-typescript/packages/rivetkit/src/db/mod.ts b/rivetkit-typescript/packages/rivetkit/src/db/mod.ts index c5daf6abd5..eac8c3fb16 100644 --- a/rivetkit-typescript/packages/rivetkit/src/db/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/db/mod.ts @@ -2,6 +2,7 @@ export type { AnyDatabaseProvider, DatabaseProvider, DatabaseProviderContext, + InferDatabaseClient, NativeDatabaseProvider, RawAccess, RawDatabaseClient, diff --git a/rivetkit-typescript/packages/rivetkit/src/inspector/workflow.ts b/rivetkit-typescript/packages/rivetkit/src/inspector/workflow.ts new file mode 100644 index 0000000000..b51fac61f6 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/inspector/workflow.ts @@ -0,0 +1,57 @@ +import type { JsonCompatValue } from "@/common/encoding"; +import { encodeCborCompat } from "@/serde"; +import { bufferToArrayBuffer } from "@/utils"; + +export type { + WorkflowBranchStatus, + WorkflowCbor, + WorkflowEntry, + WorkflowEntryKind, + WorkflowEntryMetadata, + WorkflowHistory, + WorkflowJoinEntry, + WorkflowLocation, + WorkflowLoopEntry, + WorkflowLoopIterationMarker, + WorkflowMessageEntry, + WorkflowNameIndex, + WorkflowPathSegment, + WorkflowRaceEntry, + WorkflowRemovedEntry, + WorkflowRollbackCheckpointEntry, + WorkflowSleepEntry, + WorkflowStepEntry, + WorkflowVersionCheckEntry, +} from "@/common/bare/transport/v1"; +export { + WorkflowBranchStatusType, + WorkflowEntryStatus, + WorkflowSleepState, +} from "@/common/bare/transport/v1"; +export { + decodeWorkflowHistoryTransport, + encodeWorkflowHistoryTransport, +} from "@/common/inspector-transport"; + +/** @experimental State exposed by a durable workflow run handler to the Inspector. */ +export type WorkflowState = + | "pending" + | "running" + | "sleeping" + | "failed" + | "completed" + | "cancelled" + | "rolling_back"; + +/** @experimental The raw workflow Inspector adapter consumed by RivetKit's transport. */ +export interface WorkflowInspectorAdapter { + getHistory: () => ArrayBuffer | null; + getState: () => Promise; + onHistoryUpdated: (listener: (history: ArrayBuffer) => void) => () => void; + replayFromStep: (entryId?: string) => Promise; +} + +/** @experimental Encodes a workflow Inspector value with RivetKit's CBOR-compatible codec. */ +export function encodeWorkflowInspectorValue(value: unknown): ArrayBuffer { + return bufferToArrayBuffer(encodeCborCompat(value as JsonCompatValue)); +} diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index fc930fe660..87180b34cc 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -8,6 +8,7 @@ import type { } from "@rivetkit/rivetkit-napi"; import type { ActorContextHandle, + ActorStateTransactionHandle, ActorFactoryHandle, CancellationTokenHandle, ConnHandle, @@ -59,6 +60,9 @@ type NapiSqlBatchStatement = Parameters< type NapiSqlTransaction = Awaited< ReturnType >; +type NapiActorStateTransaction = Awaited< + ReturnType +>; function asNativeRegistry(handle: RegistryHandle): NativeCoreRegistry { return handle as unknown as NativeCoreRegistry; @@ -86,6 +90,12 @@ function asNativeSqlTransaction( return handle as unknown as NapiSqlTransaction; } +function asNativeActorStateTransaction( + handle: ActorStateTransactionHandle, +): NapiActorStateTransaction { + return handle as unknown as NapiActorStateTransaction; +} + function asNativeCancellationToken( handle: CancellationTokenHandle, ): NativeCancellationToken { @@ -420,6 +430,13 @@ export class NapiCoreRuntime implements CoreRuntime { asNativeActorContext(ctx).setAlarm(timestampMs); } + async actorSetRunWakeAt( + ctx: ActorContextHandle, + timestampMs?: number | undefined | null, + ): Promise { + await asNativeActorContext(ctx).setRunWakeAt(timestampMs); + } + actorRequestSave( ctx: ActorContextHandle, opts?: RuntimeRequestSaveOpts | undefined | null, @@ -762,6 +779,49 @@ export class NapiCoreRuntime implements CoreRuntime { await asNativeSqlTransaction(transaction).rollback(); } + async actorBeginStateTransaction( + ctx: ActorContextHandle, + timeoutMs?: number, + ): Promise { + return (await asNativeActorContext(ctx).beginStateTransaction( + timeoutMs, + )) as unknown as ActorStateTransactionHandle; + } + + async actorStateTransactionExec( + transaction: ActorStateTransactionHandle, + sql: string, + ): Promise { + return await asNativeActorStateTransaction(transaction).exec(sql); + } + + async actorStateTransactionExecute( + transaction: ActorStateTransactionHandle, + sql: string, + params?: RuntimeSqlBindParams, + ): Promise { + const result = await asNativeActorStateTransaction(transaction).execute( + sql, + toNapiSqlBindParams(params), + ); + return normalizeRuntimeSqlExecuteResult(result); + } + + async actorStateTransactionCommit( + transaction: ActorStateTransactionHandle, + payload: RuntimeStateDeltaPayload, + ): Promise { + await asNativeActorStateTransaction(transaction).commit( + toNapiStateDeltaPayload(payload), + ); + } + + async actorStateTransactionRollback( + transaction: ActorStateTransactionHandle, + ): Promise { + await asNativeActorStateTransaction(transaction).rollback(); + } + async actorSqlQuery( ctx: ActorContextHandle, sql: string, @@ -862,6 +922,21 @@ export class NapiCoreRuntime implements CoreRuntime { ); } + async actorQueueCompletePersisted( + ctx: ActorContextHandle, + messageId: bigint, + expectedName: string, + response?: RuntimeBytes | undefined | null, + ): Promise { + return await asNativeActorContext(ctx) + .queue() + .completePersisted( + messageId, + expectedName, + response == null ? response : toNapiBuffer(response), + ); + } + async actorQueueEnqueueAndWait( ctx: ActorContextHandle, name: string, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index c9fe937f74..909ce2c431 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -15,6 +15,7 @@ import { disposeRunInspector, getRunFunction, getRunInspectorConfig, + hasRunInspectorConfig, RAW_STATE_SYMBOL, type ScheduledEventInfo, type WorkflowInspectorConfig, @@ -45,6 +46,11 @@ import { import { convertRegistryConfigToClientConfig } from "@/client/config"; import { HEADER_CONN_PARAMS } from "@/common/actor-router-consts"; import type { AnyDatabaseProvider } from "@/common/database/config"; +import { + db as defaultDb, + isBuiltInDatabaseProvider, + registerNativeStateTransactionOpener, +} from "@/common/database/mod"; import { wrapJsNativeDatabase } from "@/common/database/native-database"; import { assertJsonCompatValue, type JsonCompatValue } from "@/common/encoding"; import { isResponseLike, type ResponseLike } from "@/common/fetch-like"; @@ -80,6 +86,7 @@ import { validateQueueBody, validateQueueComplete, } from "./native-validation"; +import { RunHandlerCoordinator } from "./run-handler-coordinator"; import type { ActorContextHandle, ActorFactoryHandle, @@ -379,6 +386,30 @@ function databaseClientNotReadyError(): RivetError { ); } +function guardCustomDatabaseStateTransactions(client: T): T { + const transaction = Reflect.get(client, "transaction"); + if (typeof transaction !== "function") { + return client; + } + const guardedTransaction = ( + callback: unknown, + options?: { experimental?: { includeState?: boolean } }, + ) => { + if (options?.experimental?.includeState) { + throw new Error( + "experimental.includeState is only supported by RivetKit's embedded database provider", + ); + } + return Reflect.apply(transaction, client, [callback, options]); + }; + return new Proxy(client, { + get(target, property, receiver) { + if (property === "transaction") return guardedTransaction; + return Reflect.get(target, property, receiver); + }, + }); +} + function stateNotEnabledError(): RivetError { return new RivetError( "actor", @@ -518,6 +549,7 @@ async function closeNativeDatabaseClient( function getOrCreateNativeSqlDatabase( runtime: CoreRuntime, ctx: ActorContextHandle, + serializeState: () => RuntimeStateDeltaPayload, ): ReturnType { const runtimeState = getNativeRuntimeState(runtime, ctx); const cachedDatabase = runtimeState.sql; @@ -549,6 +581,29 @@ function getOrCreateNativeSqlDatabase( runtime.actorSqlTransactionRollback(transaction), }; }, + beginStateTransaction: async (timeoutMs) => { + const transaction = await runtime.actorBeginStateTransaction( + ctx, + timeoutMs, + ); + return { + exec: (sql) => + runtime.actorStateTransactionExec(transaction, sql), + execute: (sql, params) => + runtime.actorStateTransactionExecute( + transaction, + sql, + params, + ), + commit: () => + runtime.actorStateTransactionCommit( + transaction, + serializeState(), + ), + rollback: () => + runtime.actorStateTransactionRollback(transaction), + }; + }, query: (sql, params) => runtime.actorSqlQuery(ctx, sql, params), run: (sql, params) => runtime.actorSqlRun(ctx, sql, params), metrics: () => runtime.actorSqlMetrics(ctx), @@ -1874,6 +1929,32 @@ class NativeQueueAdapter { } } + async waitForAvailable( + names?: readonly string[], + options?: { timeout?: number; signal?: AbortSignal }, + ): Promise { + await this.waitForNamesAvailable(names ?? [], options); + } + + async complete( + message: { id: bigint; name: string }, + response?: unknown, + ): Promise { + const validatedResponse = validateQueueComplete( + this.#schemas, + message.name, + response, + ); + await callNative(() => + this.#runtime.actorQueueCompletePersisted( + this.#ctx, + message.id, + message.name, + encodeValue(validatedResponse), + ), + ); + } + async enqueueAndWait( name: string, body: unknown, @@ -2563,7 +2644,7 @@ export class ActorContextHandleAdapter { #queue?: NativeQueueAdapter; #request?: Request; #schedule?: NativeScheduleAdapter; - #sql?: ReturnType; + #run?: { setWakeAt(timestamp: number | null): Promise }; #runHandlerActiveProvider?: () => boolean; #onStateChange?: NativeOnStateChangeHandler; #stateEnabled: boolean; @@ -2608,13 +2689,6 @@ export class ActorContextHandleAdapter { return this.#kv; } - get sql() { - if (!this.#sql) { - this.#sql = getOrCreateNativeSqlDatabase(this.#runtime, this.#ctx); - } - return this.#sql; - } - async actorRuntimeSocket(): Promise<{ path: string }> { return await callNative(() => this.#runtime.actorRuntimeSocketProvision(this.#ctx), @@ -2720,6 +2794,27 @@ export class ActorContextHandleAdapter { return this.#queue; } + get run() { + if (!this.#run) { + this.#run = { + setWakeAt: async (timestamp: number | null) => { + if ( + timestamp !== null && + (!Number.isSafeInteger(timestamp) || timestamp < 0) + ) { + throw new TypeError( + "Run wake timestamp must be a non-negative safe integer or null", + ); + } + await callNative(() => + this.#runtime.actorSetRunWakeAt(this.#ctx, timestamp), + ); + }, + }; + } + return this.#run; + } + get schedule(): NativeScheduleAdapter { if (!this.#schedule) { this.#schedule = new NativeScheduleAdapter( @@ -2846,7 +2941,7 @@ export class ActorContextHandleAdapter { } const actorId = this.actorId; - const client = await this.#databaseProvider.createClient({ + const createdClient = await this.#databaseProvider.createClient({ actorId, kv: { batchPut: async (entries) => { @@ -2867,16 +2962,28 @@ export class ActorContextHandleAdapter { log: { debug: (obj) => logger().debug(obj), }, - nativeDatabaseProvider: { - open: async (requestedActorId) => { - void requestedActorId; - return getOrCreateNativeSqlDatabase( + nativeDatabaseProvider: registerNativeStateTransactionOpener( + { + open: async (requestedActorId) => { + void requestedActorId; + return getOrCreateNativeSqlDatabase( + this.#runtime, + this.#ctx, + () => this.serializeForTick("save"), + ); + }, + }, + async (timeoutMs) => + await getOrCreateNativeSqlDatabase( this.#runtime, this.#ctx, - ); - }, - }, + () => this.serializeForTick("save"), + ).beginStateTransaction(timeoutMs), + ), }); + const client = isBuiltInDatabaseProvider(this.#databaseProvider) + ? createdClient + : guardCustomDatabaseStateTransactions(createdClient as object); runtimeState.databaseClient = { client, }; @@ -2904,7 +3011,6 @@ export class ActorContextHandleAdapter { async closeDatabase(): Promise { this.#db = undefined; - this.#sql = undefined; await closeNativeDatabaseClient(this.#runtime, this.#ctx); await closeNativeSqlDatabase(this.#runtime, this.#ctx); } @@ -3126,7 +3232,6 @@ export class ActorContextHandleAdapter { // down so the request-save and onStateChange always run. this.#flushStateChange(); this.#abortSignalCleanup?.(); - this.#sql = undefined; } #createActorAbortSignal(): AbortSignal { @@ -3486,7 +3591,7 @@ function buildActorConfig( return { name: options.name as string | undefined, icon: options.icon as string | undefined, - hasDatabase: config.db !== undefined || usesRemoteSqlite, + hasDatabase: true, remoteSqlite: usesRemoteSqlite, enableActorRuntimeSocket: options.enableActorRuntimeSocket === true, hasState: @@ -3632,7 +3737,10 @@ export function buildNativeFactory( definition: AnyActorDefinition, ): ActorFactoryHandle { const config = definition.config as Record; - const databaseProvider = config.db as AnyDatabaseProvider; + const databaseProvider = (config.db ?? defaultDb()) as Exclude< + AnyDatabaseProvider, + undefined + >; const actionHandlers = flattenActionHandlers(config.actions); const schemaConfig: NativeValidationConfig = { actionInputSchemas: flattenActionInputSchemas( @@ -3651,15 +3759,23 @@ export function buildNativeFactory( { encoding: "bare" }, ); const nativeRunHandlerActiveByActorId = new Map(); + const runHandlerCoordinator = hasRunInspectorConfig(config.run) + ? new RunHandlerCoordinator(config.run) + : undefined; const isNativeRunHandlerActive = (ctx: ActorContextHandle) => nativeRunHandlerActiveByActorId.get( callNativeSync(() => runtime.actorId(ctx)), ) ?? false; - const getNativeWorkflowInspector = (ctx: ActorContextHandle) => - getRunInspectorConfig( - config.run, - callNativeSync(() => runtime.actorId(ctx)), - )?.workflow as NativeWorkflowInspectorConfig | undefined; + const getNativeWorkflowInspector = (ctx: ActorContextHandle) => { + const actorId = callNativeSync(() => runtime.actorId(ctx)); + const restart = () => + callNativeSync(() => runtime.actorRestartRunHandler(ctx)); + return (runHandlerCoordinator?.getInspector(actorId, restart) + ?.workflow ?? + getRunInspectorConfig(config.run, actorId)?.workflow) as + | NativeWorkflowInspectorConfig + | undefined; + }; const onStateChange = typeof config.onStateChange === "function" ? (actorCtx: ActorContextHandleAdapter, nextState: unknown) => { @@ -3781,6 +3897,9 @@ export function buildNativeFactory( const workflowState = async () => (await getNativeWorkflowInspector(ctx)?.getState?.()) ?? null; const actorCtx = makeActorCtx(ctx, jsRequest); + const sql = getOrCreateNativeSqlDatabase(runtime, ctx, () => + actorCtx.serializeForTick("save"), + ); try { if ( url.pathname === "/inspector/state" && @@ -3952,7 +4071,7 @@ export function buildNativeFactory( url.pathname === "/inspector/database/schema" && jsRequest.method === "GET" ) { - const db = actorCtx.sql; + const db = sql; const tables = queryRows( await db.query( "SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' AND name NOT LIKE '__drizzle_%' ORDER BY name", @@ -4006,7 +4125,7 @@ export function buildNativeFactory( ); const quoted = `"${table.replace(/"/g, '""')}"`; const rows = queryRows( - await actorCtx.sql.query( + await sql.query( `SELECT * FROM ${quoted} LIMIT ? OFFSET ?`, [ Math.max(0, Math.min(limit, 500)), @@ -4050,13 +4169,13 @@ export function buildNativeFactory( body.properties as Record, ); const rows = queryRows( - await actorCtx.sql.query(body.sql, bindings), + await sql.query(body.sql, bindings), ); return jsonResponse({ rows: jsonSafe(rows) }); } const args = Array.isArray(body.args) ? body.args : []; const rows = queryRows( - await actorCtx.sql.query(body.sql, args), + await sql.query(body.sql, args), ); return jsonResponse({ rows: jsonSafe(rows) }); } @@ -4086,7 +4205,7 @@ export function buildNativeFactory( rpcs: Object.keys(actionHandlers).sort(), queueSize: inspectorSnapshot.queueSize, isStateEnabled: stateEnabled, - isDatabaseEnabled: databaseProvider !== undefined, + isDatabaseEnabled: true, isWorkflowEnabled: getNativeWorkflowInspector(ctx) !== undefined, workflowState: await workflowState(), @@ -4342,6 +4461,7 @@ export function buildNativeFactory( async (error: unknown, payload: { ctx: ActorContextHandle }) => { const { ctx } = unwrapTsfnPayload(error, payload); const actorCtx = makeActorCtx(ctx); + const actorId = callNativeSync(() => runtime.actorId(ctx)); // TODO: Move this save hook into cleanupNativeSleepRuntimeState // so immediate and deferred sleep cleanup share one save-state // path instead of passing a callback through cleanup. @@ -4372,6 +4492,9 @@ export function buildNativeFactory( saveActorState, ); } finally { + nativeRunHandlerActiveByActorId.delete(actorId); + runHandlerCoordinator?.destroy(actorId); + disposeRunInspector(config.run, actorId); await actorCtx.dispose(); } } @@ -4381,16 +4504,17 @@ export function buildNativeFactory( async (error: unknown, payload: { ctx: ActorContextHandle }) => { const { ctx } = unwrapTsfnPayload(error, payload); const actorCtx = makeActorCtx(ctx); + const actorId = callNativeSync(() => runtime.actorId(ctx)); + // Close run control before user cleanup so replay cannot race actor + // destruction. Recreating this actor id receives a fresh controller. + nativeRunHandlerActiveByActorId.delete(actorId); + runHandlerCoordinator?.destroy(actorId); + disposeRunInspector(config.run, actorId); try { if (typeof config.onDestroy === "function") { await config.onDestroy(actorCtx); } } finally { - const actorId = callNativeSync(() => runtime.actorId(ctx)); - // Release actorId-keyed state so it does not accumulate per - // destroyed actor. - nativeRunHandlerActiveByActorId.delete(actorId); - disposeRunInspector(config.run, actorId); resolveNativeDestroy(runtime, ctx); await actorCtx.closeDatabase(); clearNativeRuntimeState(runtime, ctx); @@ -4789,67 +4913,74 @@ export function buildNativeFactory( ) => { const { ctx } = unwrapTsfnPayload(error, payload); const actorId = callNativeSync(() => runtime.actorId(ctx)); - const actorCtx = makeActorCtx(ctx); - nativeRunHandlerActiveByActorId.set(actorId, true); - try { - await run(actorCtx); - } finally { - // Delete rather than set(false): an absent entry already - // reads as inactive, and deleting keeps this map bounded - // to currently-running handlers instead of accumulating an - // entry per actor id forever. - nativeRunHandlerActiveByActorId.delete(actorId); - await actorCtx.dispose(); + const executeRun = async () => { + const actorCtx = makeActorCtx(ctx); + nativeRunHandlerActiveByActorId.set(actorId, true); + try { + await run(actorCtx); + } finally { + nativeRunHandlerActiveByActorId.delete(actorId); + await actorCtx.dispose(); + } + }; + if (runHandlerCoordinator) { + await runHandlerCoordinator.run( + actorId, + () => + callNativeSync(() => + runtime.actorRestartRunHandler(ctx), + ), + executeRun, + ); + } else { + await executeRun(); } }, ); })(), - getWorkflowHistory: - getRunInspectorConfig(config.run) !== undefined - ? wrapNativeCallback( - async ( - error: unknown, - payload: { ctx: ActorContextHandle }, - ) => { - const { ctx } = unwrapTsfnPayload(error, payload); - const history = - getNativeWorkflowInspector(ctx)?.getHistory(); - return history == null - ? undefined - : encodeValue(history); + getWorkflowHistory: hasRunInspectorConfig(config.run) + ? wrapNativeCallback( + async ( + error: unknown, + payload: { ctx: ActorContextHandle }, + ) => { + const { ctx } = unwrapTsfnPayload(error, payload); + const history = + getNativeWorkflowInspector(ctx)?.getHistory(); + return history == null + ? undefined + : encodeValue(history); + }, + ) + : undefined, + replayWorkflow: hasRunInspectorConfig(config.run) + ? wrapNativeCallback( + async ( + error: unknown, + payload: { + ctx: ActorContextHandle; + entryId?: string; }, - ) - : undefined, - replayWorkflow: - getRunInspectorConfig(config.run) !== undefined - ? wrapNativeCallback( - async ( - error: unknown, - payload: { - ctx: ActorContextHandle; - entryId?: string; - }, - ) => { - const { ctx, entryId } = unwrapTsfnPayload( - error, - payload, - ); - const workflowInspector = - getNativeWorkflowInspector(ctx); - if (!workflowInspector?.replayFromStep) { - return undefined; - } + ) => { + const { ctx, entryId } = unwrapTsfnPayload( + error, + payload, + ); + const workflowInspector = + getNativeWorkflowInspector(ctx); + if (!workflowInspector?.replayFromStep) { + return undefined; + } - const history = - (await workflowInspector.replayFromStep( - entryId, - )) ?? null; - return history == null - ? undefined - : encodeValue(history); - }, - ) - : undefined, + const history = + (await workflowInspector.replayFromStep(entryId)) ?? + null; + return history == null + ? undefined + : encodeValue(history); + }, + ) + : undefined, actions: Object.fromEntries( Object.entries(actionHandlers).map(([name, handler]) => [ name, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.test.ts b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.test.ts new file mode 100644 index 0000000000..edbb603e8e --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.test.ts @@ -0,0 +1,232 @@ +import { describe, expect, test, vi } from "vitest"; +import { defineRunHandler, type RunControl } from "@/actor/config"; +import { RunHandlerCoordinator } from "./run-handler-coordinator"; + +class Deferred { + readonly promise: Promise; + resolve!: (value: T | PromiseLike) => void; + reject!: (reason?: unknown) => void; + + constructor() { + this.promise = new Promise((resolve, reject) => { + this.resolve = resolve; + this.reject = reject; + }); + } +} + +function createCoordinator() { + let control: RunControl | undefined; + const dispose = vi.fn(); + const createInspector = vi.fn((context: { control: RunControl }) => { + control = context.control; + return { + inspector: { + workflow: { + getHistory: () => null, + getState: async () => null, + onHistoryUpdated: () => () => {}, + replayFromStep: async () => null, + }, + }, + dispose, + }; + }); + const run = defineRunHandler(async () => {}, { + inspectorKind: "workflow", + createInspector, + }); + const coordinator = new RunHandlerCoordinator(run); + const restart = vi.fn(); + const inspector = coordinator.getInspector("actor-1", restart); + + return { + control: control as unknown as RunControl, + coordinator, + createInspector, + dispose, + inspector, + restart, + }; +} + +describe("RunHandlerCoordinator", () => { + test("fails explicitly when a JavaScript factory omits required workflow controls", async () => { + const coordinator = new RunHandlerCoordinator( + defineRunHandler(async () => {}, { + inspectorKind: "workflow", + createInspector: (() => ({ + inspector: { workflow: { getHistory: () => null } }, + })) as never, + }), + ); + + await expect( + coordinator.run( + "actor-1", + () => {}, + async () => {}, + ), + ).rejects.toThrow( + "createInspector returned an invalid workflow adapter for actor actor-1", + ); + }); + + test("initializes and disposes Inspector state even when it is never queried", async () => { + const dispose = vi.fn(); + const createInspector = vi.fn(() => ({ + inspector: { + workflow: { + getHistory: () => null, + getState: async () => null, + onHistoryUpdated: () => () => {}, + replayFromStep: async () => null, + }, + }, + dispose, + })); + const coordinator = new RunHandlerCoordinator( + defineRunHandler(async () => {}, { + inspectorKind: "workflow", + createInspector, + }), + ); + + await coordinator.run( + "actor-1", + () => {}, + async () => {}, + ); + expect(createInspector).toHaveBeenCalledOnce(); + coordinator.destroy("actor-1"); + expect(dispose).toHaveBeenCalledOnce(); + }); + + test("creates one inspector per live actor and disposes it exactly once", () => { + const subject = createCoordinator(); + expect( + subject.coordinator.getInspector("actor-1", subject.restart), + ).toBe(subject.inspector); + expect(subject.createInspector).toHaveBeenCalledOnce(); + + subject.coordinator.destroy("actor-1"); + subject.coordinator.destroy("actor-1"); + expect(subject.dispose).toHaveBeenCalledOnce(); + + const recreated = subject.coordinator.getInspector( + "actor-1", + subject.restart, + ); + expect(recreated).not.toBe(subject.inspector); + expect(subject.createInspector).toHaveBeenCalledTimes(2); + }); + + test("rejects replay while a run is active", async () => { + const subject = createCoordinator(); + const active = new Deferred(); + const started = new Deferred(); + const running = subject.coordinator.run( + "actor-1", + subject.restart, + async () => { + started.resolve(); + await active.promise; + }, + ); + await started.promise; + + await expect( + subject.control.run.withInactive( + { restartOnSuccess: true }, + async () => {}, + ), + ).rejects.toMatchObject({ + group: "actor", + code: "run_handler_unavailable", + }); + + active.resolve(); + await running; + }); + + test("blocks starts during replay and restarts once before releasing them", async () => { + const subject = createCoordinator(); + const replay = new Deferred(); + const replayStarted = new Deferred(); + const events: string[] = []; + subject.restart.mockImplementation(() => { + events.push("restart"); + }); + + const exclusive = subject.control.run.withInactive( + { restartOnSuccess: true }, + async () => { + events.push("replay"); + replayStarted.resolve(); + await replay.promise; + }, + ); + await replayStarted.promise; + + const queued = subject.coordinator.run( + "actor-1", + subject.restart, + async () => { + events.push("run"); + }, + ); + await Promise.resolve(); + expect(events).toEqual(["replay"]); + + replay.resolve(); + await exclusive; + await queued; + expect(events).toEqual(["replay", "restart", "run"]); + expect(subject.restart).toHaveBeenCalledOnce(); + }); + + test("drops starts queued behind a failed replay and leaves the run inactive", async () => { + const subject = createCoordinator(); + const replay = new Deferred(); + const replayStarted = new Deferred(); + const run = vi.fn(); + const expected = new Error("rewrite failed"); + + const exclusive = subject.control.run.withInactive( + { restartOnSuccess: true }, + async () => { + replayStarted.resolve(); + await replay.promise; + throw expected; + }, + ); + await replayStarted.promise; + const queued = subject.coordinator.run("actor-1", subject.restart, run); + + replay.resolve(); + await expect(exclusive).rejects.toBe(expected); + await expect(queued).resolves.toBeUndefined(); + expect(run).not.toHaveBeenCalled(); + expect(subject.restart).not.toHaveBeenCalled(); + }); + + test("allows only one concurrent replay and closes old controls on destroy", async () => { + const subject = createCoordinator(); + const replay = new Deferred(); + const first = subject.control.run.withInactive({}, async () => { + await replay.promise; + }); + + await expect( + subject.control.run.withInactive({}, async () => {}), + ).rejects.toMatchObject({ code: "run_handler_unavailable" }); + subject.coordinator.destroy("actor-1"); + replay.resolve(); + await expect(first).rejects.toMatchObject({ + code: "run_handler_unavailable", + }); + await expect( + subject.control.run.withInactive({}, async () => {}), + ).rejects.toMatchObject({ code: "run_handler_unavailable" }); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts new file mode 100644 index 0000000000..f27b420812 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts @@ -0,0 +1,214 @@ +import { + createRunInspector, + getRunInspectorKind, + type RunControl, + type RunInspectorConfig, + type RunInspectorFactoryResult, + type RunInspectorKind, +} from "@/actor/config"; +import { RivetError } from "@/actor/errors"; + +type RunHandler = ((...args: any[]) => any) | { run: (...args: any[]) => any }; + +interface ActorRunState { + active: boolean; + closed: boolean; + exclusive: boolean; + exclusiveGeneration: number; + exclusiveOutcome?: "success" | "failure"; + queued: number; + restart?: () => void | Promise; + inspectorInitialized: boolean; + inspector?: RunInspectorFactoryResult; + waiters: Set<() => void>; +} + +function runUnavailable(actorId: string, reason: string): RivetError { + return new RivetError( + "actor", + "run_handler_unavailable", + `Run handler control is unavailable for actor ${actorId}: ${reason}.`, + { + public: true, + statusCode: 409, + metadata: { actorId, reason }, + }, + ); +} + +function notifyStateChanged(state: ActorRunState): void { + const waiters = [...state.waiters]; + state.waiters.clear(); + for (const waiter of waiters) { + waiter(); + } +} + +function waitForStateChange(state: ActorRunState): Promise { + return new Promise((resolve) => { + state.waiters.add(resolve); + }); +} + +/** Coordinates a run handler and its Inspector controls within one registry. */ +export class RunHandlerCoordinator { + readonly #run: RunHandler | undefined; + readonly #states = new Map(); + + constructor(run: RunHandler | undefined) { + this.#run = run; + } + + get inspectorKind(): RunInspectorKind | undefined { + return getRunInspectorKind(this.#run); + } + + async run( + actorId: string, + restart: () => void | Promise, + callback: () => T | Promise, + ): Promise { + const state = this.#getOrCreate(actorId); + state.restart = restart; + this.#initializeInspector(actorId, state); + state.queued += 1; + const blockedGeneration = state.exclusive + ? state.exclusiveGeneration + : undefined; + + try { + while (!state.closed && (state.exclusive || state.active)) { + await waitForStateChange(state); + } + } finally { + state.queued -= 1; + } + + if ( + state.closed || + (blockedGeneration !== undefined && + state.exclusiveGeneration === blockedGeneration && + state.exclusiveOutcome === "failure") + ) { + return undefined; + } + + state.active = true; + try { + return await callback(); + } finally { + state.active = false; + notifyStateChanged(state); + } + } + + getInspector( + actorId: string, + restart: () => void | Promise, + ): RunInspectorConfig | undefined { + const state = this.#getOrCreate(actorId); + state.restart = restart; + this.#initializeInspector(actorId, state); + return state.inspector?.inspector; + } + + destroy(actorId: string): void { + const state = this.#states.get(actorId); + if (!state) return; + + state.closed = true; + state.inspector?.dispose?.(); + state.inspector = undefined; + notifyStateChanged(state); + this.#states.delete(actorId); + } + + #control(actorId: string, state: ActorRunState): RunControl { + return { + run: { + withInactive: async (options, callback) => { + if (state.closed) { + throw runUnavailable(actorId, "actor is destroyed"); + } + if (state.active || state.queued > 0 || state.exclusive) { + throw runUnavailable( + actorId, + "the run handler is active or waiting to start", + ); + } + + state.exclusive = true; + state.exclusiveGeneration += 1; + state.exclusiveOutcome = undefined; + try { + const result = await callback(); + if (state.closed) { + throw runUnavailable(actorId, "actor is destroyed"); + } + state.exclusiveOutcome = "success"; + if (options.restartOnSuccess) { + if (!state.restart) { + throw runUnavailable( + actorId, + "the runtime restart callback is not available", + ); + } + await state.restart(); + } + return result; + } catch (error) { + state.exclusiveOutcome = "failure"; + throw error; + } finally { + state.exclusive = false; + notifyStateChanged(state); + } + }, + }, + }; + } + + #getOrCreate(actorId: string): ActorRunState { + let state = this.#states.get(actorId); + if (!state) { + state = { + active: false, + closed: false, + exclusive: false, + exclusiveGeneration: 0, + inspectorInitialized: false, + queued: 0, + waiters: new Set(), + }; + this.#states.set(actorId, state); + } + return state; + } + + #initializeInspector(actorId: string, state: ActorRunState): void { + if (state.inspectorInitialized) return; + const inspector = createRunInspector(this.#run, { + actorId, + control: this.#control(actorId, state), + }); + if (this.inspectorKind === "workflow") { + const workflow = + inspector === undefined + ? undefined + : inspector.inspector.workflow; + if ( + !workflow || + typeof workflow.getHistory !== "function" || + typeof workflow.getState !== "function" || + typeof workflow.onHistoryUpdated !== "function" || + typeof workflow.replayFromStep !== "function" + ) { + throw new TypeError( + `defineRunHandler createInspector returned an invalid workflow adapter for actor ${actorId}`, + ); + } + } + state.inspector = inspector; + state.inspectorInitialized = true; + } +} diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index 6827db8a3c..a358d176bb 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -16,6 +16,7 @@ export type ConnHandle = OpaqueHandle<"conn">; export type WebSocketHandle = OpaqueHandle<"webSocket">; export type CancellationTokenHandle = OpaqueHandle<"cancellationToken">; export type SqliteTransactionHandle = OpaqueHandle<"sqliteTransaction">; +export type ActorStateTransactionHandle = OpaqueHandle<"actorStateTransaction">; export type RuntimeBytes = Uint8Array; @@ -479,6 +480,10 @@ export interface CoreRuntime { ctx: ActorContextHandle, timestampMs?: number | undefined | null, ): void; + actorSetRunWakeAt( + ctx: ActorContextHandle, + timestampMs?: number | undefined | null, + ): Promise; actorRequestSave( ctx: ActorContextHandle, opts?: RuntimeRequestSaveOpts | undefined | null, @@ -617,6 +622,26 @@ export interface CoreRuntime { actorSqlTransactionRollback( transaction: SqliteTransactionHandle, ): Promise; + actorBeginStateTransaction( + ctx: ActorContextHandle, + timeoutMs?: number, + ): Promise; + actorStateTransactionExec( + transaction: ActorStateTransactionHandle, + sql: string, + ): Promise; + actorStateTransactionExecute( + transaction: ActorStateTransactionHandle, + sql: string, + params?: RuntimeSqlBindParams, + ): Promise; + actorStateTransactionCommit( + transaction: ActorStateTransactionHandle, + payload: RuntimeStateDeltaPayload, + ): Promise; + actorStateTransactionRollback( + transaction: ActorStateTransactionHandle, + ): Promise; actorSqlQuery( ctx: ActorContextHandle, sql: string, @@ -656,6 +681,12 @@ export interface CoreRuntime { options?: RuntimeQueueWaitOptions | undefined | null, signal?: CancellationTokenHandle | undefined | null, ): Promise; + actorQueueCompletePersisted( + ctx: ActorContextHandle, + messageId: bigint, + expectedName: string, + response?: RuntimeBytes | undefined | null, + ): Promise; actorQueueEnqueueAndWait( ctx: ActorContextHandle, name: string, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index 5051edffb8..c3eacf6bdc 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -14,6 +14,7 @@ import type { } from "./config"; import type { ActorContextHandle, + ActorStateTransactionHandle, ActorFactoryHandle, CancellationTokenHandle, ConnHandle, @@ -399,6 +400,17 @@ export class WasmCoreRuntime implements CoreRuntime { ); } + async actorSetRunWakeAt( + ctx: ActorContextHandle, + timestampMs?: number | undefined | null, + ): Promise { + await callHandleAsync( + asWasmActorContext(ctx), + "setRunWakeAt", + optionalWasmNumber(timestampMs), + ); + } + actorRequestSave( ctx: ActorContextHandle, opts?: RuntimeRequestSaveOpts | undefined | null, @@ -803,6 +815,75 @@ export class WasmCoreRuntime implements CoreRuntime { ); } + async actorBeginStateTransaction( + ctx: ActorContextHandle, + timeoutMs?: number, + ): Promise { + return (await callWasm(() => + ( + asWasmActorContext(ctx) as unknown as { + beginStateTransaction( + timeoutMs?: number, + ): Promise; + } + ).beginStateTransaction(timeoutMs), + )) as ActorStateTransactionHandle; + } + + async actorStateTransactionExec( + transaction: ActorStateTransactionHandle, + sql: string, + ): Promise { + return await callWasm(() => + ( + transaction as unknown as { + exec(sql: string): Promise; + } + ).exec(sql), + ); + } + + async actorStateTransactionExecute( + transaction: ActorStateTransactionHandle, + sql: string, + params?: RuntimeSqlBindParams, + ): Promise { + const result = await callWasm(() => + ( + transaction as unknown as { + execute( + sql: string, + params?: RuntimeSqlBindParams, + ): Promise; + } + ).execute(sql, params), + ); + return normalizeRuntimeSqlExecuteResult(result); + } + + async actorStateTransactionCommit( + transaction: ActorStateTransactionHandle, + payload: RuntimeStateDeltaPayload, + ): Promise { + await callWasm(() => + ( + transaction as unknown as { + commit(payload: RuntimeStateDeltaPayload): Promise; + } + ).commit(payload), + ); + } + + async actorStateTransactionRollback( + transaction: ActorStateTransactionHandle, + ): Promise { + await callWasm(() => + ( + transaction as unknown as { rollback(): Promise } + ).rollback(), + ); + } + async actorSqlQuery( ctx: ActorContextHandle, sql: string, @@ -912,6 +993,22 @@ export class WasmCoreRuntime implements CoreRuntime { ); } + async actorQueueCompletePersisted( + ctx: ActorContextHandle, + messageId: bigint, + expectedName: string, + response?: RuntimeBytes | undefined | null, + ): Promise { + const queue = childHandle(asWasmActorContext(ctx), "queue"); + return await callHandleAsync( + queue, + "completePersisted", + messageId, + expectedName, + response, + ); + } + async actorQueueEnqueueAndWait( ctx: ActorContextHandle, name: string, diff --git a/rivetkit-typescript/packages/rivetkit/src/utils/node.ts b/rivetkit-typescript/packages/rivetkit/src/utils/node.ts index 2e70c1560b..a4af2453f7 100644 --- a/rivetkit-typescript/packages/rivetkit/src/utils/node.ts +++ b/rivetkit-typescript/packages/rivetkit/src/utils/node.ts @@ -24,15 +24,10 @@ let hasImportedDependencies = false; // We use require() instead of await import() because registry.start() cannot // be async and needs immediate access to Node.js modules during setup. export function getRequireFn() { - // TODO: This causes issues in tsup - // CommonJS context - use global require - // if (typeof require !== "undefined") { - // console.log("existing require"); - // return require; - // } - - // ESM context - use createRequire with import.meta.url - return createRequire(import.meta.url); + // This loader only resolves Node built-ins, so its resolution base is + // intentionally process-local. Avoid import.meta here: tsup preserves it in + // CommonJS output, which makes the published require entrypoint invalid. + return createRequire(`${process.cwd()}/package.json`); } /** diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts index c1ba096378..bce8c52032 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts @@ -8,7 +8,7 @@ import type { import type { RunContext } from "@/actor/config"; import type { AnyStaticActorInstance } from "@/actor/definition"; import { makeWorkflowKey, workflowStoragePrefix } from "@/actor/keys"; -import type { SqliteDatabase } from "@/common/database/config"; +import type { RawAccess } from "@/common/database/config"; const WORKFLOW_STORAGE_PREFIX = workflowStoragePrefix(); // Keep workflow flushes below depot's 320-dirty-page commit ceiling. The @@ -68,52 +68,40 @@ function normalizeSqlBlob(value: unknown): Uint8Array { throw new Error("workflow sqlite value was not a blob"); } -function runtimeSqlFromContext( +function runtimeDbFromContext( runCtx?: RunContext, -): SqliteDatabase | undefined { - const sql = (runCtx as unknown as { sql?: unknown } | undefined)?.sql; +): RawAccess | undefined { + const db = (runCtx as unknown as { db?: unknown } | undefined)?.db; if ( - sql && - typeof sql === "object" && - "query" in sql && - "execute" in sql && - "executeBatch" in sql && - "run" in sql + db && + typeof db === "object" && + "execute" in db && + "transaction" in db ) { - return sql as SqliteDatabase; + return db as RawAccess; } return undefined; } -type WorkflowSqliteDatabase = SqliteDatabase & - Required>; - class WorkflowStorage { - #sql: WorkflowSqliteDatabase; + #db: RawAccess; - constructor(sql?: SqliteDatabase) { - if (!sql) { + constructor(db?: RawAccess) { + if (!db) { throw new Error( "workflow storage requires embedded SQLite; actor KV workflow storage is no longer supported", ); } - if (!sql.executeBatch) { - throw new Error( - "workflow storage requires a SQLite database with executeBatch support", - ); - } - this.#sql = sql as WorkflowSqliteDatabase; + this.#db = db; } async get(key: Uint8Array): Promise { const prefixed = makeWorkflowKey(key); - const result = await this.#sql.query( + const rows = await this.#db.execute<{ value: unknown }>( "SELECT value FROM _rivet_wf_kv WHERE key = ?", - [prefixed], + prefixed, ); - return result.rows[0]?.[0] == null - ? null - : normalizeSqlBlob(result.rows[0][0]); + return rows[0]?.value == null ? null : normalizeSqlBlob(rows[0].value); } async set(key: Uint8Array, value: Uint8Array): Promise { @@ -122,18 +110,20 @@ class WorkflowStorage { async delete(key: Uint8Array): Promise { const prefixed = makeWorkflowKey(key); - await this.#sql.run("DELETE FROM _rivet_wf_kv WHERE key = ?", [ + await this.#db.execute( + "DELETE FROM _rivet_wf_kv WHERE key = ?", prefixed, - ]); + ); } async deletePrefix(prefix: Uint8Array): Promise { const start = makeWorkflowKey(prefix); const end = computeUpperBound(start); if (end) { - await this.#sql.run( + await this.#db.execute( "DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?", - [start, end], + start, + end, ); return; } @@ -148,9 +138,10 @@ class WorkflowStorage { async deleteRange(start: Uint8Array, end: Uint8Array): Promise { const prefixedStart = makeWorkflowKey(start); const prefixedEnd = makeWorkflowKey(end); - await this.#sql.run( + await this.#db.execute( "DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?", - [prefixedStart, prefixedEnd], + prefixedStart, + prefixedEnd, ); } @@ -168,29 +159,33 @@ class WorkflowStorage { async batch(writes: KVWrite[]): Promise { if (writes.length === 0) return; for (const chunk of chunkWorkflowWrites(writes)) { - await this.#sql.executeBatch( - chunk.map(({ key, value }) => ({ - sql: "INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", - params: [makeWorkflowKey(key), value], - })), - ); + await this.#db.transaction(async (tx) => { + for (const { key, value } of chunk) { + await tx.execute( + "INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", + makeWorkflowKey(key), + value, + ); + } + }); } } async listRaw(prefix: Uint8Array): Promise { const end = computeUpperBound(prefix); - const result = end - ? await this.#sql.query( + const rows = end + ? await this.#db.execute<{ key: unknown; value: unknown }>( "SELECT key, value FROM _rivet_wf_kv WHERE key >= ? AND key < ? ORDER BY key ASC", - [prefix, end], + prefix, + end, ) - : await this.#sql.query( + : await this.#db.execute<{ key: unknown; value: unknown }>( "SELECT key, value FROM _rivet_wf_kv WHERE key >= ? ORDER BY key ASC", - [prefix], + prefix, ); - return result.rows.map((row) => [ - normalizeSqlBlob(row[0]), - normalizeSqlBlob(row[1]), + return rows.map((row) => [ + normalizeSqlBlob(row.key), + normalizeSqlBlob(row.value), ]); } @@ -200,14 +195,17 @@ class WorkflowStorage { start < keys.length; start += WORKFLOW_SQLITE_MAX_BATCH_ROWS ) { - await this.#sql.executeBatch( - keys - .slice(start, start + WORKFLOW_SQLITE_MAX_BATCH_ROWS) - .map((key) => ({ - sql: "DELETE FROM _rivet_wf_kv WHERE key = ?", - params: [key], - })), - ); + await this.#db.transaction(async (tx) => { + for (const key of keys.slice( + start, + start + WORKFLOW_SQLITE_MAX_BATCH_ROWS, + )) { + await tx.execute( + "DELETE FROM _rivet_wf_kv WHERE key = ?", + key, + ); + } + }); } } } @@ -331,7 +329,7 @@ export class ActorWorkflowDriver implements EngineDriver { this.#actor = actor; this.#runCtx = runCtx; this.messageDriver = new ActorWorkflowMessageDriver(actor, runCtx); - this.#storage = new WorkflowStorage(runtimeSqlFromContext(runCtx)); + this.#storage = new WorkflowStorage(runtimeDbFromContext(runCtx)); } async get(key: Uint8Array): Promise { @@ -425,7 +423,7 @@ export class ActorWorkflowControlDriver implements EngineDriver { runCtx?: RunContext, ) { this.#actor = actor; - this.#storage = new WorkflowStorage(runtimeSqlFromContext(runCtx)); + this.#storage = new WorkflowStorage(runtimeDbFromContext(runCtx)); } async get(key: Uint8Array): Promise { diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/inspector.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/inspector.ts index 155a65154f..02822ce894 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/inspector.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/inspector.ts @@ -10,14 +10,15 @@ import type { WorkflowHistorySnapshot, WorkflowState, } from "@rivetkit/workflow-engine"; -import type * as inspectorSchema from "@/common/bare/generated/inspector/v4"; -import * as transport from "@/common/bare/transport/v1"; -import type { JsonCompatValue } from "@/common/encoding"; -import { encodeWorkflowHistoryTransport } from "@/common/inspector-transport"; -import { encodeCborCompat } from "@/serde"; -import { assertUnreachable, bufferToArrayBuffer } from "@/utils"; +import * as transport from "@/inspector/workflow"; +import { + encodeWorkflowHistoryTransport, + encodeWorkflowInspectorValue, + type WorkflowInspectorAdapter as PublicWorkflowInspectorAdapter, +} from "@/inspector/workflow"; +import { assertUnreachable } from "@/utils"; -type HistoryListener = (history: inspectorSchema.WorkflowHistory) => void; +type HistoryListener = (history: ArrayBuffer) => void; function createHistoryEmitter() { const listeners = new Set(); @@ -27,7 +28,7 @@ function createHistoryEmitter() { listeners.add(listener); return () => listeners.delete(listener); }, - emit: (history: inspectorSchema.WorkflowHistory) => { + emit: (history: ArrayBuffer) => { for (const listener of listeners) { listener(history); } @@ -35,35 +36,23 @@ function createHistoryEmitter() { }; } -export interface WorkflowInspectorAdapter { - getHistory: () => inspectorSchema.WorkflowHistory | null; - getState: () => Promise; - onHistoryUpdated: ( - listener: (history: inspectorSchema.WorkflowHistory) => void, - ) => () => void; - replayFromStep: ( - entryId?: string, - ) => Promise; -} +export type WorkflowInspectorAdapter = PublicWorkflowInspectorAdapter; export function createWorkflowInspectorAdapter(): { adapter: WorkflowInspectorAdapter; update: (snapshot: WorkflowHistorySnapshot) => void; setGetState: (fn: () => Promise) => void; setReplayFromStep: ( - fn: ( - entryId?: string, - ) => Promise, + fn: (entryId?: string) => Promise, ) => void; } { const emitter = createHistoryEmitter(); - let history: inspectorSchema.WorkflowHistory | null = null; + let history: ArrayBuffer | null = null; let getState: () => Promise = async () => null; - let replayFromStep: ( - entryId?: string, - ) => Promise = async () => { - throw new Error("Workflow replay controls are not initialized"); - }; + let replayFromStep: (entryId?: string) => Promise = + async () => { + throw new Error("Workflow replay controls are not initialized"); + }; const adapter: WorkflowInspectorAdapter = { getHistory: () => history, @@ -92,7 +81,7 @@ export function createWorkflowInspectorAdapter(): { } function encodeCbor(value: unknown): ArrayBuffer { - return bufferToArrayBuffer(encodeCborCompat(value as JsonCompatValue)); + return encodeWorkflowInspectorValue(value); } function encodeOptionalCbor(value: unknown): ArrayBuffer | null { diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts index 8c50a13da8..7a41686ea3 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts @@ -13,10 +13,11 @@ import { type WorkflowErrorEvent, } from "@rivetkit/workflow-engine"; import invariant from "invariant"; -import type { RunContext } from "@/actor/config"; import { ACTOR_CONTEXT_INTERNAL_SYMBOL, - RUN_FUNCTION_CONFIG_SYMBOL, + defineRunHandler, + type RunContext, + type RunControl, } from "@/actor/config"; import type { AnyStaticActorInstance } from "@/actor/definition"; import { isActorAbortedError, RivetError } from "@/actor/errors"; @@ -96,6 +97,14 @@ function isWorkflowReplayBlockedByRunningEntry(error: unknown): boolean { ); } +function isRunHandlerUnavailable(error: unknown): boolean { + return ( + error instanceof RivetError && + error.group === "actor" && + error.code === "run_handler_unavailable" + ); +} + export interface WorkflowOptions< TState, TConnParams, @@ -166,10 +175,10 @@ export function workflow< >, ) => Promise { const onError = options.onError; - const workflowInspectors = new Map< - string, - ReturnType - >(); + type WorkflowInspectorRegistration = ReturnType< + typeof createWorkflowInspectorAdapter + > & { control?: RunControl }; + const workflowInspectors = new Map(); function getWorkflowInspector(actorId: string) { let workflowInspector = workflowInspectors.get(actorId); @@ -204,31 +213,35 @@ export function workflow< const controlDriver = new ActorWorkflowControlDriver(actor, runCtx); workflowInspector.setReplayFromStep(async (entryId) => { const workflowState = await workflowInspector.adapter.getState(); - if ( - actor.isRunHandlerActive() || - workflowState === "pending" || - workflowState === "running" - ) { + if (workflowState === "pending" || workflowState === "running") { throw workflowReplayInFlightError(); } - let snapshot: Awaited>; + const control = workflowInspector.control; + invariant(control, "workflow Inspector control is not initialized"); try { - snapshot = await replayWorkflowFromStep( - actor.id, - controlDriver, - entryId, - { scheduleAlarm: false }, + return await control.run.withInactive( + { restartOnSuccess: true }, + async () => { + const snapshot = await replayWorkflowFromStep( + actor.id, + controlDriver, + entryId, + { scheduleAlarm: false }, + ); + workflowInspector.update(snapshot); + return workflowInspector.adapter.getHistory(); + }, ); } catch (error) { - if (isWorkflowReplayBlockedByRunningEntry(error)) { + if ( + isWorkflowReplayBlockedByRunningEntry(error) || + isRunHandlerUnavailable(error) + ) { throw workflowReplayInFlightError(); } throw error; } - workflowInspector.update(snapshot); - await actor.restartRunHandler(); - return workflowInspector.adapter.getHistory(); }); const handle = runWorkflow( @@ -286,57 +299,22 @@ export function workflow< } } - const runWithConfig = run as typeof run & { - [RUN_FUNCTION_CONFIG_SYMBOL]?: { - icon?: string; - inspectorFactory?: (actor: unknown) => unknown; - disposeInspector?: (actorId: string) => void; - }; - }; - runWithConfig[RUN_FUNCTION_CONFIG_SYMBOL] = { + return defineRunHandler(run, { icon: "diagram-project", - // Drop the per-actor inspector when the actor is destroyed so this map - // does not retain one inspector (and its encoded history) per actor id - // for the process lifetime. - disposeInspector: (actorId) => { - workflowInspectors.delete(actorId); - }, - inspectorFactory: (actor) => { - const actorId = resolveWorkflowInspectorActorId(actor); + inspectorKind: "workflow", + createInspector: ({ actorId, control }) => { + const workflowInspector = getWorkflowInspector(actorId); + workflowInspector.control = control; return { - workflow: actorId - ? getWorkflowInspector(actorId).adapter - : { - getHistory: () => null, - onHistoryUpdated: () => () => {}, - replayFromStep: async () => null, - }, + inspector: { workflow: workflowInspector.adapter }, + dispose: () => { + // Do not let a stale disposer remove a newly-created adapter for + // the same actor id. + if (workflowInspectors.get(actorId) === workflowInspector) { + workflowInspectors.delete(actorId); + } + }, }; }, - }; - - return runWithConfig; -} - -function resolveWorkflowInspectorActorId(actor: unknown): string | undefined { - if (typeof actor === "string" && actor.length > 0) { - return actor; - } - - if (!actor || typeof actor !== "object") { - return undefined; - } - - const candidate = actor as { - id?: unknown; - actorId?: unknown; - }; - if (typeof candidate.id === "string" && candidate.id.length > 0) { - return candidate.id; - } - if (typeof candidate.actorId === "string" && candidate.actorId.length > 0) { - return candidate.actorId; - } - - return undefined; + }); } diff --git a/rivetkit-typescript/packages/rivetkit/tests/database-state-transaction.test.ts b/rivetkit-typescript/packages/rivetkit/tests/database-state-transaction.test.ts new file mode 100644 index 0000000000..a51f7f5e53 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/database-state-transaction.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, test } from "vitest"; +import type { DatabaseProvider, RawAccess } from "@/common/database/config"; +import { ActorContextHandleAdapter } from "@/registry/native"; + +describe("experimental database state transactions", () => { + test("rejects includeState for custom database providers", async () => { + const customClient: RawAccess = { + execute: async () => [], + transaction: async (callback) => await callback(customClient), + close: async () => {}, + }; + const provider: DatabaseProvider = { + createClient: async () => customClient, + onMigrate: async () => {}, + }; + const runtimeState = {}; + const context = new ActorContextHandleAdapter( + { + actorId: () => "actor-a", + actorRuntimeState: () => runtimeState, + } as never, + {} as never, + undefined, + {}, + provider, + ); + await context.prepare(); + + await expect(context.db.transaction(async () => 42)).resolves.toBe(42); + expect(() => + context.db.transaction(async () => {}, { + experimental: { includeState: true }, + }), + ).toThrow( + "experimental.includeState is only supported by RivetKit's embedded database provider", + ); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-db.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-db.test.ts index 30fc001f7c..dbbf2d0863 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-db.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-db.test.ts @@ -231,7 +231,7 @@ describeDriverMatrix( driverTestConfig.runtime !== "native" || driverTestConfig.sqliteBackend !== "local", )( - "rejects Actor Runtime Socket provisioning when opt-in or SQLite is missing", + "requires opt-in and provisions the default embedded database", async (c) => { const { client } = await setupDriverTest( c, @@ -252,12 +252,11 @@ describeDriverMatrix( client.actorRuntimeSocketWithoutDb.getOrCreate([ `runtime-socket-without-db-${crypto.randomUUID()}`, ]); - await expect( - withoutDb.getActorRuntimeSocketPath(), - ).rejects.toMatchObject({ - group: "actor_runtime_socket", - code: "database_unavailable", - }); + const path = await withoutDb.getActorRuntimeSocketPath(); + expect(path).toBeTruthy(); + expect(existsSync(path)).toBe(true); + await withoutDb.destroy(); + await vi.waitFor(() => expect(existsSync(path)).toBe(false)); }, dbTestTimeout, ); diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/engine-restart-serverless-runtime.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/engine-restart-serverless-runtime.ts index 4e7dd5df35..791a1456e1 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/engine-restart-serverless-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/engine-restart-serverless-runtime.ts @@ -42,10 +42,12 @@ interface HeartbeatVars { } const rawSqlDatabaseProvider = { - createClient: async () => ({ - execute: async () => [], - close: async () => {}, - }), + createClient: async (ctx: any) => { + if (!ctx.nativeDatabaseProvider) { + throw new Error("native SQLite is required"); + } + return await ctx.nativeDatabaseProvider.open(ctx.actorId); + }, onMigrate: async () => {}, }; @@ -203,7 +205,7 @@ const sqliteCounter = actor({ return; } - const database = ctx.sql as SqliteDatabase; + const database = ctx.db as SqliteDatabase; vars.heartbeatSeq = 0; logRuntimeEvent("heartbeat_on_wake", { actorId: ctx.actorId, @@ -300,7 +302,7 @@ const sqliteCounter = actor({ }, actions: { tick: async (ctx, payloadBytes = 4096) => { - const database = ctx.sql as SqliteDatabase; + const database = ctx.db as SqliteDatabase; const payload = "x".repeat(Math.max(0, Math.trunc(payloadBytes))); const now = Date.now(); @@ -352,7 +354,7 @@ const sqliteCounter = actor({ } }, getCount: async (ctx) => { - const database = ctx.sql as SqliteDatabase; + const database = ctx.db as SqliteDatabase; await ensureTables(database); const rows = await database.query( "SELECT count FROM restart_counter WHERE id = ?", @@ -368,7 +370,7 @@ const sqliteCounter = actor({ payloadBytes?: number; }, ) => { - const database = ctx.sql as SqliteDatabase; + const database = ctx.db as SqliteDatabase; const payload = "x".repeat( Math.max(0, Math.trunc(input.payloadBytes ?? 8192)), ); diff --git a/rivetkit-typescript/packages/rivetkit/tests/inspector-workflow-surface.test.ts b/rivetkit-typescript/packages/rivetkit/tests/inspector-workflow-surface.test.ts new file mode 100644 index 0000000000..4cb912053f --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/inspector-workflow-surface.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, test } from "vitest"; +import type { WorkflowHistory } from "@/common/bare/transport/v1"; +import { + decodeWorkflowHistoryTransport, + encodeWorkflowHistoryTransport, +} from "@/common/inspector-transport"; +import { + decodeWorkflowHistoryTransport as decodePublicWorkflowHistory, + encodeWorkflowHistoryTransport as encodePublicWorkflowHistory, + encodeWorkflowInspectorValue, + WorkflowEntryStatus, + type WorkflowInspectorAdapter, + WorkflowSleepState, + type WorkflowState, +} from "@/inspector/workflow"; +import { encodeCborCompat } from "@/serde"; + +function bytes(value: ArrayBuffer): Uint8Array { + return new Uint8Array(value); +} + +describe("rivetkit/experimental/inspector/workflow", () => { + test("preserves the existing raw BARE history bytes", () => { + const history: WorkflowHistory = { + nameRegistry: ["root", "delay"], + entries: [ + { + id: "sleep-1", + location: [ + { tag: "WorkflowNameIndex", val: 0 }, + { + tag: "WorkflowLoopIterationMarker", + val: { loop: 1, iteration: 2 }, + }, + ], + kind: { + tag: "WorkflowSleepEntry", + val: { + deadline: 1_234n, + state: WorkflowSleepState.PENDING, + }, + }, + }, + ], + entryMetadata: new Map([ + [ + "sleep-1", + { + status: WorkflowEntryStatus.RUNNING, + error: null, + attempts: 2, + lastAttemptAt: 1_200n, + createdAt: 1_000n, + completedAt: null, + rollbackCompletedAt: null, + rollbackError: null, + }, + ], + ]), + }; + + const publicBytes = encodePublicWorkflowHistory(history); + const existingBytes = encodeWorkflowHistoryTransport(history); + expect(bytes(publicBytes)).toEqual(bytes(existingBytes)); + expect(decodePublicWorkflowHistory(publicBytes)).toEqual(history); + expect(decodeWorkflowHistoryTransport(publicBytes)).toEqual(history); + }); + + test("preserves the existing CBOR-compatible value bytes", () => { + const value = { count: 3n, nested: [null, "ok"] }; + expect(bytes(encodeWorkflowInspectorValue(value))).toEqual( + new Uint8Array(encodeCborCompat(value)), + ); + }); + + test("requires state on the public workflow adapter contract", async () => { + const state: WorkflowState = "sleeping"; + const adapter: WorkflowInspectorAdapter = { + getHistory: () => null, + getState: async () => state, + onHistoryUpdated: () => () => {}, + replayFromStep: async () => null, + }; + + await expect(adapter.getState()).resolves.toBe("sleeping"); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/tests/nested-actions.test.ts b/rivetkit-typescript/packages/rivetkit/tests/nested-actions.test.ts index fce6306034..95aff2d60e 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/nested-actions.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/nested-actions.test.ts @@ -4,10 +4,18 @@ import { flattenActionHandlers, flattenActionInputSchemas, } from "../src/actor/actions"; -import { actor } from "../src/actor/definition"; +import { + type ActionContext, + hasRunInspectorConfig, +} from "../src/actor/config"; +import { + actor, + type BaseActorDefinition, +} from "../src/actor/definition"; import type { ActorDefinitionActions } from "../src/client/actor-common"; import type { ActorHandleRaw } from "../src/client/actor-handle"; import { createActorProxy } from "../src/client/client"; +import type { RawAccess } from "../src/common/database/config"; describe("nested actions", () => { test("preserves nested handler context and client action types", () => { @@ -32,6 +40,63 @@ describe("nested actions", () => { >().returns.toEqualTypeOf>(); }); + test("preserves default database context and actions on base definitions", () => { + const definition = actor({ + actions: { + query: async (c, value: number) => { + expectTypeOf(c.db).toEqualTypeOf(); + return value.toString(); + }, + }, + }); + type DefinitionActions = ActorDefinitionActions; + expectTypeOf() + .parameter(0) + .toEqualTypeOf(); + + type DefaultContext = ActionContext< + undefined, + undefined, + undefined, + undefined, + undefined, + undefined + >; + type DefaultActions = { + query: (c: DefaultContext, value: number) => string; + }; + type WorkflowStyleDefinition = BaseActorDefinition< + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + Record, + Record, + DefaultActions + >; + type WorkflowStyleActions = + ActorDefinitionActions; + expectTypeOf() + .parameter(0) + .toEqualTypeOf(); + expectTypeOf< + WorkflowStyleActions["query"] + >().returns.toEqualTypeOf>(); + }); + + test("detects legacy run inspector metadata without invoking its factory", () => { + const inspectorFactory = vi.fn(() => undefined); + const run = () => {}; + Object.defineProperty(run, Symbol.for("rivetkit.run_function_config"), { + value: { inspectorFactory }, + }); + + expect(hasRunInspectorConfig(run)).toBe(true); + expect(inspectorFactory).not.toHaveBeenCalled(); + }); + test("dispatches nested proxy calls with dotted names", async () => { const action = vi.fn().mockResolvedValue("created"); const handle = createActorProxy({ diff --git a/rivetkit-typescript/packages/rivetkit/tests/package-surface.test.ts b/rivetkit-typescript/packages/rivetkit/tests/package-surface.test.ts index f76029ee17..f6580b980c 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/package-surface.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/package-surface.test.ts @@ -34,6 +34,11 @@ import { TO_SERVER_VERSIONED, type TransportWorkflowHistory, } from "rivetkit/inspector/client"; +import { + decodeWorkflowHistoryTransport as decodePublicWorkflowHistory, + encodeWorkflowHistoryTransport as encodePublicWorkflowHistory, + WorkflowEntryStatus, +} from "rivetkit/experimental/inspector/workflow"; import { setupTest } from "rivetkit/test"; import { jsonParseCompat, jsonStringifyCompat } from "rivetkit/utils"; import { describe, expect, test } from "vitest"; @@ -48,7 +53,10 @@ const contextTypeSmokeActor = rivetkit.actor({ userId: params.userId, }), actions: { - increment: (ctx, amount: number) => (ctx.state.count += amount), + increment: (ctx, amount: number) => { + ctx.state.count += amount; + return ctx.state.count; + }, }, run: async (ctx) => { ctx.state.count += 1; @@ -88,6 +96,9 @@ describe("package surface", () => { test("restores supported package entrypoints", () => { expect(packageJson.exports).toHaveProperty("./test"); expect(packageJson.exports).toHaveProperty("./inspector"); + expect(packageJson.exports).toHaveProperty( + "./experimental/inspector/workflow", + ); expect(packageJson.exports).toHaveProperty("./inspector/client"); expect(packageJson.exports).toHaveProperty("./db"); expect(packageJson.exports).toHaveProperty("./db/drizzle"); @@ -96,6 +107,9 @@ describe("package surface", () => { test("restored package entrypoints resolve", () => { expect(setupTest).toBeTypeOf("function"); expect(decodeWorkflowHistoryTransport).toBeTypeOf("function"); + expect(decodePublicWorkflowHistory).toBeTypeOf("function"); + expect(encodePublicWorkflowHistory).toBeTypeOf("function"); + expect(WorkflowEntryStatus.COMPLETED).toBe("COMPLETED"); expect(rawDb).toBeTypeOf("function"); expect(drizzleDb).toBeTypeOf("function"); expect(defineConfig).toBeTypeOf("function"); diff --git a/rivetkit-typescript/packages/rivetkit/tests/platforms/CLAUDE.md b/rivetkit-typescript/packages/rivetkit/tests/platforms/CLAUDE.md index 6a20962e6c..da2a4956a3 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/platforms/CLAUDE.md +++ b/rivetkit-typescript/packages/rivetkit/tests/platforms/CLAUDE.md @@ -7,7 +7,7 @@ - Cloudflare Workers, Supabase Functions, and Deno fixtures should share the same docs-shaped SQLite counter actor source with only platform bootstrapping differences. - Use `buildPlatformSqliteCounterActorSource()` for the shared actor in package-based fixtures (Cloudflare/Supabase); use `buildPlatformSqliteCounterRegistrySource(...)` for the raw-`setup()` Deno fixture. - Do not use lower-level registry builders, private generated wasm paths, or repo-local `pkg*` imports in platform app code. -- Raw `ctx.sql` platform fixtures still need a `db` provider so runtime SQLite is enabled. +- Raw SQLite platform fixtures use an explicit `db` provider that returns the runtime-native client through `ctx.db`. - Cloudflare Workers need a fetch-upgrade `WebSocket` shim for wasm envoy connections; it lives inside `@rivetkit/cloudflare-workers` (installed on `globalThis` by the package), so fixtures and user code must not hand-roll one. - Deno fixtures need `--allow-sys` because public `rivetkit` imports `pino`, which reads `os.hostname()`. - Deno fixtures should load wasm bytes from the public `@rivetkit/rivetkit-wasm/rivetkit_wasm_bg.wasm` export with `import.meta.resolve` plus `Deno.readFile`. diff --git a/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-platform-harness.ts b/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-platform-harness.ts index 9e147266d4..3d21d49b62 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-platform-harness.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-platform-harness.ts @@ -107,10 +107,10 @@ const PLATFORM_SQLITE_COUNTER_ACTOR_BODY = `interface SqliteDatabase { const COUNTER_ID = 1; const rawSqlDatabaseProvider = { -\tcreateClient: async () => ({ -\t\texecute: async () => [], -\t\tclose: async () => {}, -\t}), +\tcreateClient: async (ctx: any) => { +\t\tif (!ctx.nativeDatabaseProvider) throw new Error("native SQLite is required"); +\t\treturn await ctx.nativeDatabaseProvider.open(ctx.actorId); +\t}, \tonMigrate: async () => {}, }; @@ -164,14 +164,14 @@ async function readLifecycleCounts(db: SqliteDatabase): Promise<{ export const sqliteCounter = actor({ \tdb: rawSqlDatabaseProvider, \tonWake: async (ctx) => { -\t\tawait recordLifecycleEvent(ctx.sql as SqliteDatabase, "wake"); +\t\tawait recordLifecycleEvent(ctx.db as SqliteDatabase, "wake"); \t}, \tonSleep: async (ctx) => { -\t\tawait recordLifecycleEvent(ctx.sql as SqliteDatabase, "sleep"); +\t\tawait recordLifecycleEvent(ctx.db as SqliteDatabase, "sleep"); \t}, \tactions: { \t\tincrement: async (ctx, amount = 1) => { -\t\t\tconst db = ctx.sql as SqliteDatabase; +\t\t\tconst db = ctx.db as SqliteDatabase; \t\t\tawait ensureCounterTable(db); \t\t\tawait db.run( \t\t\t\t"INSERT INTO platform_counter (id, count) VALUES (?, ?) ON CONFLICT(id) DO UPDATE SET count = count + excluded.count", @@ -181,13 +181,13 @@ export const sqliteCounter = actor({ \t\t\treturn await readCounter(db); \t\t}, \t\tgetCount: async (ctx) => { -\t\t\tconst db = ctx.sql as SqliteDatabase; +\t\t\tconst db = ctx.db as SqliteDatabase; \t\t\tawait ensureCounterTable(db); \t\t\treturn await readCounter(db); \t\t}, \t\tgetLifecycleCounts: async (ctx) => { -\t\t\treturn await readLifecycleCounts(ctx.sql as SqliteDatabase); +\t\t\treturn await readLifecycleCounts(ctx.db as SqliteDatabase); \t\t}, \t\ttriggerSleep: (ctx) => { \t\t\tctx.sleep(); diff --git a/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-registry.ts b/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-registry.ts index bb805cdf7f..586afa08f4 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-registry.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/platforms/shared-registry.ts @@ -19,10 +19,12 @@ interface SqliteDatabase { const COUNTER_ID = 1; const rawSqlDatabaseProvider = { - createClient: async () => ({ - execute: async () => [], - close: async () => {}, - }), + createClient: async (ctx: any) => { + if (!ctx.nativeDatabaseProvider) { + throw new Error("native SQLite is required"); + } + return await ctx.nativeDatabaseProvider.open(ctx.actorId); + }, onMigrate: async () => {}, }; @@ -86,14 +88,14 @@ async function readLifecycleCounts(db: SqliteDatabase): Promise<{ export const sqliteCounterActor = actor({ db: rawSqlDatabaseProvider, onWake: async (ctx) => { - await recordLifecycleEvent(ctx.sql as SqliteDatabase, "wake"); + await recordLifecycleEvent(ctx.db as SqliteDatabase, "wake"); }, onSleep: async (ctx) => { - await recordLifecycleEvent(ctx.sql as SqliteDatabase, "sleep"); + await recordLifecycleEvent(ctx.db as SqliteDatabase, "sleep"); }, actions: { increment: async (ctx, amount = 1) => { - const db = ctx.sql as SqliteDatabase; + const db = ctx.db as SqliteDatabase; await ensureCounterTable(db); await db.run( ` @@ -107,13 +109,13 @@ export const sqliteCounterActor = actor({ return await readCounter(db); }, getCount: async (ctx) => { - const db = ctx.sql as SqliteDatabase; + const db = ctx.db as SqliteDatabase; await ensureCounterTable(db); return await readCounter(db); }, getLifecycleCounts: async (ctx) => { - return await readLifecycleCounts(ctx.sql as SqliteDatabase); + return await readLifecycleCounts(ctx.db as SqliteDatabase); }, triggerSleep: (ctx) => { ctx.sleep(); diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index f0ed09a7cf..e1bb24aea3 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -4,7 +4,7 @@ import { decodeBridgeRivetError, type RivetError, } from "@/actor/errors"; -import { actor } from "@/actor/mod"; +import { actor, defineRunHandler, type RunControl } from "@/actor/mod"; import { type RegistryConfig, RegistryConfigSchema } from "@/registry/config"; import { NapiCoreRuntime } from "@/registry/napi-runtime"; import { buildNativeFactory } from "@/registry/native"; @@ -35,6 +35,10 @@ type NativeCallbacks = { input?: Uint8Array; }, ) => Promise; + run?: ( + error: unknown, + payload: { ctx: ActorContextHandle }, + ) => Promise; actions: Record< string, ( @@ -94,6 +98,7 @@ class ParityScenario { readonly registerTask = new Gate(); readonly saves: unknown[] = []; registerTaskCompleted = false; + runRestarts = 0; } class FakeActorContext { @@ -166,6 +171,10 @@ class FakeActorContext { abortSignal(): AbortSignal { return this.abortController.signal; } + + restartRunHandler(): void { + this.scenario.runRestarts += 1; + } } class FakeCancellationToken { @@ -420,6 +429,68 @@ async function invokePromotedStatus( } describe("CoreRuntime NAPI and wasm parity", () => { + test.each([ + "napi", + "wasm", + ] as const)("%s shares the public run-control gate with runtime starts", async (kind) => { + const runtimeCase = createRuntimeCase(kind); + const started = new Gate(); + let control: RunControl | undefined; + const definition = actor({ + run: defineRunHandler( + async () => { + started.markStarted(); + await started.released; + }, + { + inspectorKind: "workflow", + createInspector: (context) => { + control = context.control; + return { + inspector: { + workflow: { + getHistory: () => null, + getState: async () => null, + onHistoryUpdated: () => () => {}, + replayFromStep: async () => null, + }, + }, + }; + }, + }, + ), + }); + const factory = buildNativeFactory( + runtimeCase.runtime, + registryConfig(definition), + definition, + ) as unknown as FakeActorFactory; + const callbacks = factory.callbacks as NativeCallbacks & { + getWorkflowHistory?: NativeCallbacks["run"]; + }; + const ctx = new FakeActorContext(runtimeCase.scenario); + + await callbacks.getWorkflowHistory?.(null, { ctx }); + expect(control).toBeDefined(); + const running = callbacks.run?.(null, { ctx }); + await started.started; + + await expect( + control?.run.withInactive({}, async () => {}), + ).rejects.toMatchObject({ + group: "actor", + code: "run_handler_unavailable", + }); + + started.release(); + await running; + await control?.run.withInactive( + { restartOnSuccess: true }, + async () => {}, + ); + expect(runtimeCase.scenario.runRestarts).toBe(1); + }); + test("scheduled fire metadata is appended after validated action args", async () => { const runtimeCase = createRuntimeCase("napi"); let received: unknown[] = []; diff --git a/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts b/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts index e8163417da..1020028149 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts @@ -348,6 +348,7 @@ function fakeWasmBindings( input: encodeValue({ host: host.kind }), }); ctx.stateBytes = Buffer.from(initialState); + await factory.callbacks.onMigrate(null, { ctx, isNew: true }); let actionSettled = false; const actionPromise = factory.callbacks.actions.smoke(null, { @@ -366,11 +367,23 @@ function fakeWasmBindings( }, ); - await scenario.actionReconnect.started; + await Promise.race([ + scenario.actionReconnect.started, + actionPromise.then(() => { + throw new Error("action settled before the reconnect gate"); + }), + ]); host.reconnect(config, "during-action"); scenario.actionReconnect.release(); - await scenario.remoteWriteReconnect.started; + await Promise.race([ + scenario.remoteWriteReconnect.started, + actionPromise.then(() => { + throw new Error( + "action settled before the remote-write reconnect gate", + ); + }), + ]); host.reconnect(config, "during-remote-write-sql"); scenario.remoteWriteReconnect.release(); @@ -443,13 +456,6 @@ async function runHostSmoke(kind: HostKind): Promise { const registry = runtime.createRegistry(); const definition = actor({ state: { count: 0 }, - db: { - createClient: async () => ({ - execute: async () => [], - close: async () => {}, - }), - onMigrate: async () => {}, - }, actions: { smoke: async (c, label: string) => { c.state.count += 1; @@ -459,21 +465,22 @@ async function runHostSmoke(kind: HostKind): Promise { scenario.actionReconnect.markStarted(); await scenario.actionReconnect.released; - await c.sql.execute( + await c.db.execute( "INSERT INTO smoke_events (host) VALUES (?)", - [label], + label, ); scenario.remoteWriteReconnect.markStarted(); await scenario.remoteWriteReconnect.released; - await c.sql.execute( + await c.db.execute( "UPDATE smoke_events SET host = ? WHERE id = ?", - [label, 1], + label, + 1, ); - const rows = await c.sql.query( + const rows = await c.db.execute( "SELECT host FROM smoke_events WHERE host = ?", - [label], + label, ); await c.saveState({ immediate: true }); void ( @@ -489,7 +496,7 @@ async function runHostSmoke(kind: HostKind): Promise { return { stateCount: c.state.count, kvValue, - sqlRows: rows.rows.length, + sqlRows: rows.length, }; }, }, diff --git a/rivetkit-typescript/packages/rivetkit/tsconfig.json b/rivetkit-typescript/packages/rivetkit/tsconfig.json index 4ba9ab3b75..64b9531cce 100644 --- a/rivetkit-typescript/packages/rivetkit/tsconfig.json +++ b/rivetkit-typescript/packages/rivetkit/tsconfig.json @@ -11,6 +11,7 @@ "rivetkit/db/drizzle": ["./src/db/drizzle.ts"], "rivetkit/dynamic": ["./src/dynamic/mod.ts"], "rivetkit/errors": ["./src/actor/errors.ts"], + "rivetkit/experimental/inspector/workflow": ["./src/inspector/workflow.ts"], "rivetkit/utils": ["./src/utils.ts"], "rivetkit/agent-os": ["./src/agent-os/index.ts"] }