11import type { ActorQueue , ActorRun , RunContext } from "rivetkit" ;
2- import {
3- WORKFLOW_STORAGE_V1 ,
4- type WorkflowStorageHandle ,
5- } from "rivetkit/storage" ;
2+ import type { RawAccess } from "rivetkit/db" ;
63import type {
74 EngineDriver ,
85 KVEntry ,
@@ -12,6 +9,14 @@ import type {
129 WorkflowMessageIdentity ,
1310} from "../index.js" ;
1411
12+ const WORKFLOW_STORAGE_PREFIX = new Uint8Array ( [ 6 , 1 ] ) ;
13+ const WORKFLOW_UPSERT_SQL =
14+ "INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value" ;
15+
16+ const WORKFLOW_SQLITE_MAX_VALUE_BYTES = 256 * 1024 ;
17+ const WORKFLOW_SQLITE_MAX_BATCH_ROWS = 128 ;
18+ const WORKFLOW_SQLITE_MAX_BATCH_BYTES = 512 * 1024 ;
19+
1520function track < T > (
1621 runCtx : RunContext < any , any , any , any , any , any , any , any > ,
1722 promise : Promise < T > ,
@@ -25,6 +30,169 @@ function track<T>(
2530 return promise ;
2631}
2732
33+ function prefixWorkflowKey ( key : Uint8Array ) : Uint8Array {
34+ const prefixed = new Uint8Array (
35+ WORKFLOW_STORAGE_PREFIX . byteLength + key . byteLength ,
36+ ) ;
37+ prefixed . set ( WORKFLOW_STORAGE_PREFIX ) ;
38+ prefixed . set ( key , WORKFLOW_STORAGE_PREFIX . byteLength ) ;
39+ return prefixed ;
40+ }
41+
42+ function stripWorkflowKey ( key : Uint8Array ) : Uint8Array {
43+ if (
44+ key . byteLength < WORKFLOW_STORAGE_PREFIX . byteLength ||
45+ ! WORKFLOW_STORAGE_PREFIX . every ( ( byte , index ) => key [ index ] === byte )
46+ ) {
47+ throw new Error ( "workflow SQLite key escaped the [6, 1] namespace" ) ;
48+ }
49+ return key . slice ( WORKFLOW_STORAGE_PREFIX . byteLength ) ;
50+ }
51+
52+ function computeUpperBound ( prefix : Uint8Array ) : Uint8Array {
53+ const upperBound = prefix . slice ( ) ;
54+ for ( let index = upperBound . length - 1 ; index >= 0 ; index -- ) {
55+ if ( upperBound [ index ] !== 0xff ) {
56+ upperBound [ index ] ++ ;
57+ return upperBound . slice ( 0 , index + 1 ) ;
58+ }
59+ }
60+
61+ // Every workflow key begins with 6, so a finite upper bound always exists.
62+ throw new Error ( "workflow storage prefix has no upper bound" ) ;
63+ }
64+
65+ function normalizeSqlBlob ( value : unknown ) : Uint8Array {
66+ if ( value instanceof Uint8Array ) {
67+ return value ;
68+ }
69+ if ( value instanceof ArrayBuffer ) {
70+ return new Uint8Array ( value ) ;
71+ }
72+ if ( ArrayBuffer . isView ( value ) ) {
73+ return new Uint8Array ( value . buffer , value . byteOffset , value . byteLength ) ;
74+ }
75+ if ( Array . isArray ( value ) ) {
76+ const bytes = new Uint8Array ( value . length ) ;
77+ for ( const [ index , byte ] of value . entries ( ) ) {
78+ if ( ! Number . isInteger ( byte ) || byte < 0 || byte > 255 ) {
79+ throw new Error ( "workflow SQLite value was not a byte array" ) ;
80+ }
81+ bytes [ index ] = byte ;
82+ }
83+ return bytes ;
84+ }
85+ throw new Error ( "workflow SQLite value was not a blob" ) ;
86+ }
87+
88+ function validateWrites ( writes : KVWrite [ ] ) : void {
89+ if ( writes . length > WORKFLOW_SQLITE_MAX_BATCH_ROWS ) {
90+ throw new Error (
91+ `Workflow batch contains ${ writes . length } ${ WORKFLOW_SQLITE_MAX_BATCH_ROWS } ,
92+ ) ;
93+ }
94+
95+ let batchBytes = 0 ;
96+ for ( const write of writes ) {
97+ if ( write . value . byteLength > WORKFLOW_SQLITE_MAX_VALUE_BYTES ) {
98+ throw new Error (
99+ `Workflow value is ${ write . value . byteLength } ${ WORKFLOW_SQLITE_MAX_VALUE_BYTES } ,
100+ ) ;
101+ }
102+ batchBytes +=
103+ WORKFLOW_STORAGE_PREFIX . byteLength +
104+ write . key . byteLength +
105+ write . value . byteLength ;
106+ }
107+
108+ if ( batchBytes > WORKFLOW_SQLITE_MAX_BATCH_BYTES ) {
109+ throw new Error (
110+ `Workflow batch is ${ batchBytes } ${ WORKFLOW_SQLITE_MAX_BATCH_BYTES } ,
111+ ) ;
112+ }
113+ }
114+
115+ class WorkflowStorage {
116+ #db: RawAccess ;
117+
118+ constructor ( db : RawAccess ) {
119+ this . #db = db ;
120+ }
121+
122+ async get ( key : Uint8Array ) : Promise < Uint8Array | null > {
123+ const rows = await this . #db. execute < { value : unknown } > (
124+ "SELECT value FROM _rivet_wf_kv WHERE key = ?" ,
125+ prefixWorkflowKey ( key ) ,
126+ ) ;
127+ const value = rows [ 0 ] ?. value ;
128+ return value == null ? null : normalizeSqlBlob ( value ) ;
129+ }
130+
131+ async set ( key : Uint8Array , value : Uint8Array ) : Promise < void > {
132+ await this . batch ( [ { key, value } ] , false ) ;
133+ }
134+
135+ async delete ( key : Uint8Array ) : Promise < void > {
136+ await this . #db. execute (
137+ "DELETE FROM _rivet_wf_kv WHERE key = ?" ,
138+ prefixWorkflowKey ( key ) ,
139+ ) ;
140+ }
141+
142+ async deletePrefix ( prefix : Uint8Array ) : Promise < void > {
143+ const start = prefixWorkflowKey ( prefix ) ;
144+ await this . #db. execute (
145+ "DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?" ,
146+ start ,
147+ computeUpperBound ( start ) ,
148+ ) ;
149+ }
150+
151+ async deleteRange ( start : Uint8Array , end : Uint8Array ) : Promise < void > {
152+ await this . #db. execute (
153+ "DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?" ,
154+ prefixWorkflowKey ( start ) ,
155+ prefixWorkflowKey ( end ) ,
156+ ) ;
157+ }
158+
159+ async list ( prefix : Uint8Array ) : Promise < KVEntry [ ] > {
160+ const start = prefixWorkflowKey ( prefix ) ;
161+ const rows = await this . #db. execute < { key : unknown ; value : unknown } > (
162+ "SELECT key, value FROM _rivet_wf_kv WHERE key >= ? AND key < ? ORDER BY key ASC" ,
163+ start ,
164+ computeUpperBound ( start ) ,
165+ ) ;
166+ return rows . map ( ( row ) => ( {
167+ key : stripWorkflowKey ( normalizeSqlBlob ( row . key ) ) ,
168+ value : normalizeSqlBlob ( row . value ) ,
169+ } ) ) ;
170+ }
171+
172+ async batch ( writes : KVWrite [ ] , includeState : boolean ) : Promise < void > {
173+ if ( writes . length === 0 ) return ;
174+ validateWrites ( writes ) ;
175+
176+ const commit = async ( tx : RawAccess ) => {
177+ for ( const write of writes ) {
178+ await tx . execute (
179+ WORKFLOW_UPSERT_SQL ,
180+ prefixWorkflowKey ( write . key ) ,
181+ write . value ,
182+ ) ;
183+ }
184+ } ;
185+
186+ if ( includeState ) {
187+ await this . #db. transaction ( commit , {
188+ experimental : { includeState : true } ,
189+ } ) ;
190+ } else {
191+ await this . #db. transaction ( commit ) ;
192+ }
193+ }
194+ }
195+
28196class ActorWorkflowMessageDriver implements WorkflowMessageDriver {
29197 #runCtx: RunContext < any , any , any , any , any , any , any , any > ;
30198 #queue: ActorQueue ;
@@ -95,15 +263,15 @@ export class ActorWorkflowDriver implements EngineDriver {
95263 readonly workerPollInterval = 100 ;
96264 readonly messageDriver : WorkflowMessageDriver ;
97265 #runCtx: RunContext < any , any , any , any , any , any , any , any > ;
98- #storage: WorkflowStorageHandle ;
266+ #storage: WorkflowStorage ;
99267 #queue: ActorQueue ;
100268 #run: ActorRun ;
101269
102270 constructor ( runCtx : RunContext < any , any , any , any , any , any , any , any > ) {
103271 this . #runCtx = runCtx ;
104272 this . messageDriver = new ActorWorkflowMessageDriver ( runCtx ) ;
105273 this . #queue = runCtx . queue ;
106- this . #storage = runCtx . storage . open ( WORKFLOW_STORAGE_V1 ) ;
274+ this . #storage = new WorkflowStorage ( runCtx . db ) ;
107275 this . #run = runCtx . run ;
108276 }
109277
@@ -132,9 +300,7 @@ export class ActorWorkflowDriver implements EngineDriver {
132300 }
133301
134302 async batch ( writes : KVWrite [ ] ) : Promise < void > {
135- if ( writes . length === 0 ) return ;
136-
137- await track ( this . #runCtx, this . #storage. flushWithState ( writes ) ) ;
303+ await track ( this . #runCtx, this . #storage. batch ( writes , true ) ) ;
138304 }
139305
140306 async setAlarm ( _workflowId : string , wakeAt : number ) : Promise < void > {
@@ -184,11 +350,11 @@ export class ActorWorkflowControlDriver implements EngineDriver {
184350 readonly workerPollInterval = 100 ;
185351 readonly messageDriver : WorkflowMessageDriver =
186352 new NoopWorkflowMessageDriver ( ) ;
187- #storage: WorkflowStorageHandle ;
353+ #storage: WorkflowStorage ;
188354 #run: ActorRun ;
189355
190356 constructor ( runCtx : RunContext < any , any , any , any , any , any , any , any > ) {
191- this . #storage = runCtx . storage . open ( WORKFLOW_STORAGE_V1 ) ;
357+ this . #storage = new WorkflowStorage ( runCtx . db ) ;
192358 this . #run = runCtx . run ;
193359 }
194360
@@ -217,11 +383,7 @@ export class ActorWorkflowControlDriver implements EngineDriver {
217383 }
218384
219385 async batch ( writes : KVWrite [ ] ) : Promise < void > {
220- if ( writes . length === 0 ) {
221- return ;
222- }
223-
224- await this . #storage. batch ( writes ) ;
386+ await this . #storage. batch ( writes , false ) ;
225387 }
226388
227389 async setAlarm ( _workflowId : string , wakeAt : number ) : Promise < void > {
0 commit comments