diff --git a/examples/ai-agent-vercel/README.md b/examples/ai-agent-vercel/README.md index 846de305bf..e11694c409 100644 --- a/examples/ai-agent-vercel/README.md +++ b/examples/ai-agent-vercel/README.md @@ -20,7 +20,7 @@ npm run dev ## Features - Actor-per-agent pattern with a coordinating manager Rivet Actor -- Queue-based intake using `c.queue.next` inside the run loop +- Queue-based intake using `for await` over `c.queue.iter(...)` inside the run loop - Streaming AI responses sent to the UI as they arrive - Persistent history stored in Rivet Actor state - Live status updates via events and polling diff --git a/examples/ai-agent-vercel/frontend/App.tsx b/examples/ai-agent-vercel/frontend/App.tsx index 900945b812..a8fe34fcb2 100644 --- a/examples/ai-agent-vercel/frontend/App.tsx +++ b/examples/ai-agent-vercel/frontend/App.tsx @@ -79,7 +79,7 @@ function AgentPanel({ info }: { info: AgentInfo }) { return; } - await agent.connection.queue.message.send({ + await agent.connection.send("message", { text: trimmed, sender: "Operator", }); diff --git a/examples/ai-agent-vercel/src/actors.ts b/examples/ai-agent-vercel/src/actors.ts index 81bd367d5b..b1ebd29171 100644 --- a/examples/ai-agent-vercel/src/actors.ts +++ b/examples/ai-agent-vercel/src/actors.ts @@ -1,6 +1,6 @@ import { openai } from "@ai-sdk/openai"; import { streamText, type CoreMessage } from "ai"; -import { actor, setup } from "rivetkit"; +import { actor, event, queue, setup } from "rivetkit"; export type AgentMessage = { id: string; @@ -27,6 +27,14 @@ export type AgentQueueMessage = { sender?: string; }; +export type AgentResponseEvent = { + messageId: string; + delta: string; + content: string; + done: boolean; + error?: string; +}; + const SYSTEM_PROMPT = "You are a focused AI assistant. Keep responses concise, actionable, and ready for handoff."; @@ -52,13 +60,19 @@ export const agent = actor({ updatedAt: Date.now(), } as AgentStatus, }, + queues: { + message: queue(), + }, + events: { + messageAdded: event(), + status: event(), + response: event(), + }, // The run hook keeps the agent listening for queued messages. run: async (c) => { - while (true) { - const queued = await c.queue.next("message"); - - const body = queued.body as AgentQueueMessage; + for await (const queued of c.queue.iter()) { + const { body } = queued; if (!body?.text || typeof body.text !== "string") { continue; } @@ -102,7 +116,7 @@ export const agent = actor({ let content = ""; for await (const delta of result.textStream) { - if (c.abortSignal.aborted) { + if (c.aborted) { break; } diff --git a/examples/ai-agent/README.md b/examples/ai-agent/README.md index da98be1861..5bddde551e 100644 --- a/examples/ai-agent/README.md +++ b/examples/ai-agent/README.md @@ -15,7 +15,7 @@ npm run dev ## Features - Actor-per-agent pattern with a coordinating manager Rivet Actor -- Queue-based intake using `c.queue.next` inside the run loop +- Queue-based intake using `for await` over `c.queue.iter(...)` inside the run loop - Streaming AI responses sent to the UI as they arrive - Persistent history stored in Rivet Actor state - Live status updates via events and polling diff --git a/examples/ai-agent/frontend/App.tsx b/examples/ai-agent/frontend/App.tsx index 900945b812..a8fe34fcb2 100644 --- a/examples/ai-agent/frontend/App.tsx +++ b/examples/ai-agent/frontend/App.tsx @@ -79,7 +79,7 @@ function AgentPanel({ info }: { info: AgentInfo }) { return; } - await agent.connection.queue.message.send({ + await agent.connection.send("message", { text: trimmed, sender: "Operator", }); diff --git a/examples/ai-agent/src/actors.ts b/examples/ai-agent/src/actors.ts index 81bd367d5b..b1ebd29171 100644 --- a/examples/ai-agent/src/actors.ts +++ b/examples/ai-agent/src/actors.ts @@ -1,6 +1,6 @@ import { openai } from "@ai-sdk/openai"; import { streamText, type CoreMessage } from "ai"; -import { actor, setup } from "rivetkit"; +import { actor, event, queue, setup } from "rivetkit"; export type AgentMessage = { id: string; @@ -27,6 +27,14 @@ export type AgentQueueMessage = { sender?: string; }; +export type AgentResponseEvent = { + messageId: string; + delta: string; + content: string; + done: boolean; + error?: string; +}; + const SYSTEM_PROMPT = "You are a focused AI assistant. Keep responses concise, actionable, and ready for handoff."; @@ -52,13 +60,19 @@ export const agent = actor({ updatedAt: Date.now(), } as AgentStatus, }, + queues: { + message: queue(), + }, + events: { + messageAdded: event(), + status: event(), + response: event(), + }, // The run hook keeps the agent listening for queued messages. run: async (c) => { - while (true) { - const queued = await c.queue.next("message"); - - const body = queued.body as AgentQueueMessage; + for await (const queued of c.queue.iter()) { + const { body } = queued; if (!body?.text || typeof body.text !== "string") { continue; } @@ -102,7 +116,7 @@ export const agent = actor({ let content = ""; for await (const delta of result.textStream) { - if (c.abortSignal.aborted) { + if (c.aborted) { break; } diff --git a/examples/ai-and-user-generated-actors-freestyle/template/src/registry.ts b/examples/ai-and-user-generated-actors-freestyle/template/src/registry.ts index 41dfcd357a..6e23e6b9a6 100644 --- a/examples/ai-and-user-generated-actors-freestyle/template/src/registry.ts +++ b/examples/ai-and-user-generated-actors-freestyle/template/src/registry.ts @@ -1,9 +1,12 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ state: { count: 0, }, + events: { + countChanged: event(), + }, actions: { increment: (c) => { diff --git a/examples/chat-room-vercel/src/actors.ts b/examples/chat-room-vercel/src/actors.ts index 34e333eb1a..8261daa4af 100644 --- a/examples/chat-room-vercel/src/actors.ts +++ b/examples/chat-room-vercel/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type Message = { sender: string; text: string; timestamp: number }; @@ -7,6 +7,9 @@ export const chatRoom = actor({ state: { messages: [] as Message[], }, + events: { + newMessage: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/chat-room/src/actors.ts b/examples/chat-room/src/actors.ts index 34e333eb1a..8261daa4af 100644 --- a/examples/chat-room/src/actors.ts +++ b/examples/chat-room/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type Message = { sender: string; text: string; timestamp: number }; @@ -7,6 +7,9 @@ export const chatRoom = actor({ state: { messages: [] as Message[], }, + events: { + newMessage: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/cloudflare-workers-inline-client/src/registry.ts b/examples/cloudflare-workers-inline-client/src/registry.ts index 4afe732a3c..5939fc3512 100644 --- a/examples/cloudflare-workers-inline-client/src/registry.ts +++ b/examples/cloudflare-workers-inline-client/src/registry.ts @@ -1,7 +1,10 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/cloudflare-workers/src/registry.ts b/examples/cloudflare-workers/src/registry.ts index 4afe732a3c..5939fc3512 100644 --- a/examples/cloudflare-workers/src/registry.ts +++ b/examples/cloudflare-workers/src/registry.ts @@ -1,7 +1,10 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/collaborative-document-vercel/src/actors.ts b/examples/collaborative-document-vercel/src/actors.ts index 3b973c9eb4..a76d83aa31 100644 --- a/examples/collaborative-document-vercel/src/actors.ts +++ b/examples/collaborative-document-vercel/src/actors.ts @@ -1,5 +1,10 @@ -import { actor, setup } from "rivetkit"; -import { applyAwarenessUpdate, Awareness, encodeAwarenessUpdate } from "y-protocols/awareness"; +import { actor, setup, event } from "rivetkit"; +import { + applyAwarenessUpdate, + Awareness, + encodeAwarenessUpdate, + removeAwarenessStates, +} from "y-protocols/awareness"; import { randomUUID } from "node:crypto"; import * as Y from "yjs"; @@ -14,6 +19,7 @@ export type SyncEvent = { update: number[] }; export type AwarenessEvent = { update: number[] }; type DocumentInput = { title: string; createdAt: number }; +type DocumentState = { title: string; createdAt: number; updatedAt: number }; type UpdateKind = "sync" | "awareness"; @@ -25,9 +31,13 @@ export const document = actor({ connState: { clientIds: [] as number[], }, + events: { + sync: event(), + awareness: event(), + }, // Persistent metadata that survives restarts: https://rivet.dev/docs/actors/state - createState: (_c, input: DocumentInput) => ({ + createState: (_c, input: DocumentInput): DocumentState => ({ title: input.title, createdAt: input.createdAt, updatedAt: input.createdAt, @@ -49,7 +59,7 @@ export const document = actor({ if (clientIds.length === 0) { return; } - c.vars.awareness.removeStates(clientIds, "disconnect"); + removeAwarenessStates(c.vars.awareness, clientIds, "disconnect"); const update = encodeAwarenessUpdate(c.vars.awareness, clientIds); c.broadcast("awareness", { update: toNumbers(update) }); conn.state.clientIds = []; diff --git a/examples/collaborative-document/src/actors.ts b/examples/collaborative-document/src/actors.ts index 3b973c9eb4..a76d83aa31 100644 --- a/examples/collaborative-document/src/actors.ts +++ b/examples/collaborative-document/src/actors.ts @@ -1,5 +1,10 @@ -import { actor, setup } from "rivetkit"; -import { applyAwarenessUpdate, Awareness, encodeAwarenessUpdate } from "y-protocols/awareness"; +import { actor, setup, event } from "rivetkit"; +import { + applyAwarenessUpdate, + Awareness, + encodeAwarenessUpdate, + removeAwarenessStates, +} from "y-protocols/awareness"; import { randomUUID } from "node:crypto"; import * as Y from "yjs"; @@ -14,6 +19,7 @@ export type SyncEvent = { update: number[] }; export type AwarenessEvent = { update: number[] }; type DocumentInput = { title: string; createdAt: number }; +type DocumentState = { title: string; createdAt: number; updatedAt: number }; type UpdateKind = "sync" | "awareness"; @@ -25,9 +31,13 @@ export const document = actor({ connState: { clientIds: [] as number[], }, + events: { + sync: event(), + awareness: event(), + }, // Persistent metadata that survives restarts: https://rivet.dev/docs/actors/state - createState: (_c, input: DocumentInput) => ({ + createState: (_c, input: DocumentInput): DocumentState => ({ title: input.title, createdAt: input.createdAt, updatedAt: input.createdAt, @@ -49,7 +59,7 @@ export const document = actor({ if (clientIds.length === 0) { return; } - c.vars.awareness.removeStates(clientIds, "disconnect"); + removeAwarenessStates(c.vars.awareness, clientIds, "disconnect"); const update = encodeAwarenessUpdate(c.vars.awareness, clientIds); c.broadcast("awareness", { update: toNumbers(update) }); conn.state.clientIds = []; diff --git a/examples/cursors-vercel/src/actors.ts b/examples/cursors-vercel/src/actors.ts index 9936f5c704..fd5faabfed 100644 --- a/examples/cursors-vercel/src/actors.ts +++ b/examples/cursors-vercel/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export interface CursorPosition { userId: string; @@ -20,6 +20,11 @@ export const cursorRoom = actor({ state: { textLabels: [] as TextLabel[], }, + events: { + cursorMoved: event(), + textUpdated: event(), + textRemoved: event(), + }, connState: { cursor: null as CursorPosition | null, diff --git a/examples/cursors/src/actors.ts b/examples/cursors/src/actors.ts index 9936f5c704..fd5faabfed 100644 --- a/examples/cursors/src/actors.ts +++ b/examples/cursors/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export interface CursorPosition { userId: string; @@ -20,6 +20,11 @@ export const cursorRoom = actor({ state: { textLabels: [] as TextLabel[], }, + events: { + cursorMoved: event(), + textUpdated: event(), + textRemoved: event(), + }, connState: { cursor: null as CursorPosition | null, diff --git a/examples/custom-serverless-vercel/src/actors.ts b/examples/custom-serverless-vercel/src/actors.ts index 44707e2fef..ade7deb547 100644 --- a/examples/custom-serverless-vercel/src/actors.ts +++ b/examples/custom-serverless-vercel/src/actors.ts @@ -1,9 +1,12 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; const counter = actor({ state: { count: 0, }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/custom-serverless/src/actors.ts b/examples/custom-serverless/src/actors.ts index 44707e2fef..ade7deb547 100644 --- a/examples/custom-serverless/src/actors.ts +++ b/examples/custom-serverless/src/actors.ts @@ -1,9 +1,12 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; const counter = actor({ state: { count: 0, }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/experimental-durable-streams-ai-agent-vercel/src/actors.ts b/examples/experimental-durable-streams-ai-agent-vercel/src/actors.ts index b81f018fa3..90e67fcef5 100644 --- a/examples/experimental-durable-streams-ai-agent-vercel/src/actors.ts +++ b/examples/experimental-durable-streams-ai-agent-vercel/src/actors.ts @@ -1,7 +1,7 @@ import { anthropic } from "@ai-sdk/anthropic"; import type { DurableStream } from "@durable-streams/client"; import { streamText } from "ai"; -import { type ActorContextOf, actor, setup } from "rivetkit"; +import { type ActorContextOf, actor, setup, event } from "rivetkit"; import { getStreams } from "./shared/streams.ts"; import type { PromptMessage, ResponseChunk } from "./shared/types.ts"; @@ -17,6 +17,10 @@ export const aiAgent = actor({ // undefined means start from the beginning of the stream promptStreamOffset: undefined as string | undefined, }), + events: { + responseError: event<{ promptId: string; error: string }>(), + responseComplete: event<{ promptId: string; fullResponse: string }>(), + }, onWake: (c) => { consumeStream(c); @@ -69,7 +73,7 @@ async function consumeStream(c: ActorContextOf) { offset: chunk.offset, }); - if (c.abortSignal.aborted) break; + if (c.aborted) break; c.state.promptStreamOffset = chunk.offset; @@ -129,9 +133,9 @@ async function consumeStream(c: ActorContextOf) { c.log.error({ msg: "error in consumeStream", error, - aborted: c.abortSignal.aborted, + aborted: c.aborted, }); - if (!c.abortSignal.aborted) { + if (!c.aborted) { c.log.error({ msg: "error consuming prompts", error }); } } diff --git a/examples/experimental-durable-streams-ai-agent/src/actors.ts b/examples/experimental-durable-streams-ai-agent/src/actors.ts index b81f018fa3..90e67fcef5 100644 --- a/examples/experimental-durable-streams-ai-agent/src/actors.ts +++ b/examples/experimental-durable-streams-ai-agent/src/actors.ts @@ -1,7 +1,7 @@ import { anthropic } from "@ai-sdk/anthropic"; import type { DurableStream } from "@durable-streams/client"; import { streamText } from "ai"; -import { type ActorContextOf, actor, setup } from "rivetkit"; +import { type ActorContextOf, actor, setup, event } from "rivetkit"; import { getStreams } from "./shared/streams.ts"; import type { PromptMessage, ResponseChunk } from "./shared/types.ts"; @@ -17,6 +17,10 @@ export const aiAgent = actor({ // undefined means start from the beginning of the stream promptStreamOffset: undefined as string | undefined, }), + events: { + responseError: event<{ promptId: string; error: string }>(), + responseComplete: event<{ promptId: string; fullResponse: string }>(), + }, onWake: (c) => { consumeStream(c); @@ -69,7 +73,7 @@ async function consumeStream(c: ActorContextOf) { offset: chunk.offset, }); - if (c.abortSignal.aborted) break; + if (c.aborted) break; c.state.promptStreamOffset = chunk.offset; @@ -129,9 +133,9 @@ async function consumeStream(c: ActorContextOf) { c.log.error({ msg: "error in consumeStream", error, - aborted: c.abortSignal.aborted, + aborted: c.aborted, }); - if (!c.abortSignal.aborted) { + if (!c.aborted) { c.log.error({ msg: "error consuming prompts", error }); } } diff --git a/examples/hello-world-vercel/src/actors.ts b/examples/hello-world-vercel/src/actors.ts index 858c9d1262..a17372a51c 100644 --- a/examples/hello-world-vercel/src/actors.ts +++ b/examples/hello-world-vercel/src/actors.ts @@ -1,10 +1,13 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ // Persistent state that survives restarts: https://rivet.dev/docs/actors/state state: { count: 0, }, + events: { + newCount: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/hello-world/src/actors.ts b/examples/hello-world/src/actors.ts index 858c9d1262..a17372a51c 100644 --- a/examples/hello-world/src/actors.ts +++ b/examples/hello-world/src/actors.ts @@ -1,10 +1,13 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ // Persistent state that survives restarts: https://rivet.dev/docs/actors/state state: { count: 0, }, + events: { + newCount: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/hono-react-vercel/src/actors.ts b/examples/hono-react-vercel/src/actors.ts index 4afe732a3c..5939fc3512 100644 --- a/examples/hono-react-vercel/src/actors.ts +++ b/examples/hono-react-vercel/src/actors.ts @@ -1,7 +1,10 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/hono-react/src/actors.ts b/examples/hono-react/src/actors.ts index 4afe732a3c..5939fc3512 100644 --- a/examples/hono-react/src/actors.ts +++ b/examples/hono-react/src/actors.ts @@ -1,7 +1,10 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/kitchen-sink-vercel/src/actors/demo.ts b/examples/kitchen-sink-vercel/src/actors/demo.ts index c8d16676d8..89628f3dad 100644 --- a/examples/kitchen-sink-vercel/src/actors/demo.ts +++ b/examples/kitchen-sink-vercel/src/actors/demo.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { handleHttpRequest, httpActions } from "./http.ts"; import { handleWebSocket, websocketActions } from "./websocket.ts"; @@ -14,6 +14,11 @@ export const demo = actor({ connState: { connectionTime: 0, }, + events: { + countChanged: event<{ count: number; amount: number }>(), + messageChanged: event<{ message: string }>(), + alarmTriggered: event<{ id: string; time: number; data?: any }>(), + }, onWake: (c) => { c.state.startCount += 1; c.log.info({ @@ -138,7 +143,7 @@ export const demo = actor({ // Events broadcastCustomEvent: (c, eventName: string, data: any) => { - c.broadcast(eventName, data); + c.broadcast(eventName as never, data); return { eventName, data, timestamp: Date.now() }; }, diff --git a/examples/kitchen-sink/src/actors/demo.ts b/examples/kitchen-sink/src/actors/demo.ts index c8d16676d8..89628f3dad 100644 --- a/examples/kitchen-sink/src/actors/demo.ts +++ b/examples/kitchen-sink/src/actors/demo.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { handleHttpRequest, httpActions } from "./http.ts"; import { handleWebSocket, websocketActions } from "./websocket.ts"; @@ -14,6 +14,11 @@ export const demo = actor({ connState: { connectionTime: 0, }, + events: { + countChanged: event<{ count: number; amount: number }>(), + messageChanged: event<{ message: string }>(), + alarmTriggered: event<{ id: string; time: number; data?: any }>(), + }, onWake: (c) => { c.state.startCount += 1; c.log.info({ @@ -138,7 +143,7 @@ export const demo = actor({ // Events broadcastCustomEvent: (c, eventName: string, data: any) => { - c.broadcast(eventName, data); + c.broadcast(eventName as never, data); return { eventName, data, timestamp: Date.now() }; }, diff --git a/examples/multi-region-vercel/src/actors.ts b/examples/multi-region-vercel/src/actors.ts index 3286648d47..bbeeae03ac 100644 --- a/examples/multi-region-vercel/src/actors.ts +++ b/examples/multi-region-vercel/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; import type { Player } from "./types.ts"; export type { Player }; @@ -19,6 +19,11 @@ const gameRoom = actor({ connState: { playerId: null as string | null, }, + events: { + playerJoined: event<{ playerId: string; player: Player }>(), + playerLeft: event<{ playerId: string }>(), + playerMoved: event<{ playerId: string; x: number; y: number }>(), + }, // Handle client connections onConnect: (c, conn) => { diff --git a/examples/multiplayer-game-vercel/frontend/App.tsx b/examples/multiplayer-game-vercel/frontend/App.tsx index 9c5e269ed7..ef861ac9f5 100644 --- a/examples/multiplayer-game-vercel/frontend/App.tsx +++ b/examples/multiplayer-game-vercel/frontend/App.tsx @@ -46,18 +46,19 @@ export function App() { }); useEffect(() => { - if (!matchmaker.connection || roomId) return; + const connection = matchmaker.connection; + if (!connection || roomId) return; - matchmaker.connection - .findGame() - .then((nextRoomId: string) => { + void (async () => { + try { + const nextRoomId = await (await connection.findGame()); setRoomId(nextRoomId); setStatus(`Connected to ${nextRoomId}`); - }) - .catch((err: unknown) => { + } catch (err) { console.error("Failed to find game:", err); setErrorMessage("Matchmaking failed. Try again."); - }); + } + })(); }, [matchmaker.connection, roomId]); useEffect(() => { diff --git a/examples/multiplayer-game-vercel/src/actors.ts b/examples/multiplayer-game-vercel/src/actors.ts index 1eebecda63..f691fa8825 100644 --- a/examples/multiplayer-game-vercel/src/actors.ts +++ b/examples/multiplayer-game-vercel/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; import type { GameState, JoinResult, Player, RoomStats } from "./types.ts"; const MAX_PLAYERS = 10; @@ -43,6 +43,12 @@ const buildGameState = (roomId: string, maxPlayers: number, players: Record; +}; + const matchmaker = actor({ // Coordinator actor that tracks active game rooms. https://rivet.dev/docs/actors/design-patterns state: { @@ -97,14 +103,30 @@ const matchmaker = actor({ const gameRoom = actor({ // Data actor that owns room state and broadcasts changes. https://rivet.dev/docs/actors/state - createState: (_c, input: { roomId: string; maxPlayers: number }) => ({ + createState: ( + _c, + input: { roomId: string; maxPlayers: number }, + ): GameRoomState => ({ roomId: input.roomId, maxPlayers: input.maxPlayers, - players: {} as Record, + players: {}, }), connState: { playerId: null as string | null, }, + events: { + playerLeft: event<{ playerId: string }>(), + gameState: event(), + playerJoined: event<{ playerId: string; player: Player }>(), + playerEaten: event<{ eaterId: string; eatenId: string; eaterMass: number }>(), + playerMoved: event<{ + playerId: string; + x: number; + y: number; + mass: number; + radius: number; + }>(), + }, onDisconnect: (c, conn) => { const playerId = conn.state.playerId; if (!playerId) return; @@ -234,7 +256,10 @@ const resolveCollisions = ( return eaten; }; -const reportRoomCount = async (c: { state: { roomId: string; maxPlayers: number; players: Record }; client: () => any }) => { +const reportRoomCount = async (c: { + state: GameRoomState; + client: () => any; +}) => { const client = c.client(); await client.matchmaker .getOrCreate(["main"]) diff --git a/examples/multiplayer-game/frontend/App.tsx b/examples/multiplayer-game/frontend/App.tsx index 9c5e269ed7..ef861ac9f5 100644 --- a/examples/multiplayer-game/frontend/App.tsx +++ b/examples/multiplayer-game/frontend/App.tsx @@ -46,18 +46,19 @@ export function App() { }); useEffect(() => { - if (!matchmaker.connection || roomId) return; + const connection = matchmaker.connection; + if (!connection || roomId) return; - matchmaker.connection - .findGame() - .then((nextRoomId: string) => { + void (async () => { + try { + const nextRoomId = await (await connection.findGame()); setRoomId(nextRoomId); setStatus(`Connected to ${nextRoomId}`); - }) - .catch((err: unknown) => { + } catch (err) { console.error("Failed to find game:", err); setErrorMessage("Matchmaking failed. Try again."); - }); + } + })(); }, [matchmaker.connection, roomId]); useEffect(() => { diff --git a/examples/multiplayer-game/src/actors.ts b/examples/multiplayer-game/src/actors.ts index 1eebecda63..f691fa8825 100644 --- a/examples/multiplayer-game/src/actors.ts +++ b/examples/multiplayer-game/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; import type { GameState, JoinResult, Player, RoomStats } from "./types.ts"; const MAX_PLAYERS = 10; @@ -43,6 +43,12 @@ const buildGameState = (roomId: string, maxPlayers: number, players: Record; +}; + const matchmaker = actor({ // Coordinator actor that tracks active game rooms. https://rivet.dev/docs/actors/design-patterns state: { @@ -97,14 +103,30 @@ const matchmaker = actor({ const gameRoom = actor({ // Data actor that owns room state and broadcasts changes. https://rivet.dev/docs/actors/state - createState: (_c, input: { roomId: string; maxPlayers: number }) => ({ + createState: ( + _c, + input: { roomId: string; maxPlayers: number }, + ): GameRoomState => ({ roomId: input.roomId, maxPlayers: input.maxPlayers, - players: {} as Record, + players: {}, }), connState: { playerId: null as string | null, }, + events: { + playerLeft: event<{ playerId: string }>(), + gameState: event(), + playerJoined: event<{ playerId: string; player: Player }>(), + playerEaten: event<{ eaterId: string; eatenId: string; eaterMass: number }>(), + playerMoved: event<{ + playerId: string; + x: number; + y: number; + mass: number; + radius: number; + }>(), + }, onDisconnect: (c, conn) => { const playerId = conn.state.playerId; if (!playerId) return; @@ -234,7 +256,10 @@ const resolveCollisions = ( return eaten; }; -const reportRoomCount = async (c: { state: { roomId: string; maxPlayers: number; players: Record }; client: () => any }) => { +const reportRoomCount = async (c: { + state: GameRoomState; + client: () => any; +}) => { const client = c.client(); await client.matchmaker .getOrCreate(["main"]) diff --git a/examples/next-js/src/rivet/registry.ts b/examples/next-js/src/rivet/registry.ts index 1e6cdc4fca..e1a8269066 100644 --- a/examples/next-js/src/rivet/registry.ts +++ b/examples/next-js/src/rivet/registry.ts @@ -1,9 +1,12 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; const counter = actor({ state: { count: 0, }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/per-tenant-database-vercel/src/actors.ts b/examples/per-tenant-database-vercel/src/actors.ts index 6950c86025..74bf865379 100644 --- a/examples/per-tenant-database-vercel/src/actors.ts +++ b/examples/per-tenant-database-vercel/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type Employee = { id: string; @@ -54,6 +54,10 @@ export const companyDatabase = actor({ updated_at: now, }; }, + events: { + employeeAdded: event(), + projectAdded: event(), + }, // Callable functions from clients. https://rivet.dev/docs/actors/actions actions: { diff --git a/examples/per-tenant-database/src/actors.ts b/examples/per-tenant-database/src/actors.ts index 6950c86025..74bf865379 100644 --- a/examples/per-tenant-database/src/actors.ts +++ b/examples/per-tenant-database/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type Employee = { id: string; @@ -54,6 +54,10 @@ export const companyDatabase = actor({ updated_at: now, }; }, + events: { + employeeAdded: event(), + projectAdded: event(), + }, // Callable functions from clients. https://rivet.dev/docs/actors/actions actions: { diff --git a/examples/queue-sandbox-vercel/README.md b/examples/queue-sandbox-vercel/README.md index 5bbc5880ae..df335e1060 100644 --- a/examples/queue-sandbox-vercel/README.md +++ b/examples/queue-sandbox-vercel/README.md @@ -36,7 +36,7 @@ See [`src/actors/sender.ts`](https://github.com/rivet-dev/rivet/tree/main/exampl ### Multi-Queue -Listen to multiple named queues (high, normal, low priority) simultaneously using `c.queue.next(names, { count })`. +Listen to multiple named queues (high, normal, low priority) simultaneously using `for await (const m of c.queue.iter({ names: [...] }))`. See [`src/actors/multi-queue.ts`](https://github.com/rivet-dev/rivet/tree/main/examples/queue-sandbox/src/actors/multi-queue.ts). @@ -48,7 +48,7 @@ See [`src/actors/timeout.ts`](https://github.com/rivet-dev/rivet/tree/main/examp ### Worker -Use the `run` handler to continuously consume queue messages in a loop. The worker polls for jobs and processes them automatically. +Use the `run` handler to continuously consume queue messages in a `for await` loop. See [`src/actors/worker.ts`](https://github.com/rivet-dev/rivet/tree/main/examples/queue-sandbox/src/actors/worker.ts). diff --git a/examples/queue-sandbox/README.md b/examples/queue-sandbox/README.md index 18d032f10e..7bbf5b9826 100644 --- a/examples/queue-sandbox/README.md +++ b/examples/queue-sandbox/README.md @@ -31,7 +31,7 @@ See [`src/actors/sender.ts`](https://github.com/rivet-dev/rivet/tree/main/exampl ### Multi-Queue -Listen to multiple named queues (high, normal, low priority) simultaneously using `c.queue.next(names, { count })`. +Listen to multiple named queues (high, normal, low priority) simultaneously using `for await (const m of c.queue.iter({ names: [...] }))`. See [`src/actors/multi-queue.ts`](https://github.com/rivet-dev/rivet/tree/main/examples/queue-sandbox/src/actors/multi-queue.ts). @@ -43,7 +43,7 @@ See [`src/actors/timeout.ts`](https://github.com/rivet-dev/rivet/tree/main/examp ### Worker -Use the `run` handler to continuously consume queue messages in a loop. The worker polls for jobs and processes them automatically. +Use the `run` handler to continuously consume queue messages in a `for await` loop. See [`src/actors/worker.ts`](https://github.com/rivet-dev/rivet/tree/main/examples/queue-sandbox/src/actors/worker.ts). diff --git a/examples/react-vercel/src/actors.ts b/examples/react-vercel/src/actors.ts index 4afe732a3c..5939fc3512 100644 --- a/examples/react-vercel/src/actors.ts +++ b/examples/react-vercel/src/actors.ts @@ -1,7 +1,10 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/react/src/actors.ts b/examples/react/src/actors.ts index 4afe732a3c..5939fc3512 100644 --- a/examples/react/src/actors.ts +++ b/examples/react/src/actors.ts @@ -1,7 +1,10 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/sandbox-coding-agent-vercel/README.md b/examples/sandbox-coding-agent-vercel/README.md index 0328791a09..54875a4266 100644 --- a/examples/sandbox-coding-agent-vercel/README.md +++ b/examples/sandbox-coding-agent-vercel/README.md @@ -20,7 +20,7 @@ npm run dev ## Features - Actor-per-agent pattern with a coordinating manager Rivet Actor -- Queue-based intake using `c.queue.next` inside the run loop +- Queue-based intake using `for await` over `c.queue.iter(...)` inside the run loop - Sandbox Agent SDK sessions per agent with streamed output - Persistent history stored in Rivet Actor state - Live status updates via events and polling diff --git a/examples/sandbox-coding-agent-vercel/frontend/App.tsx b/examples/sandbox-coding-agent-vercel/frontend/App.tsx index 4b8bcf2107..e991b81d9d 100644 --- a/examples/sandbox-coding-agent-vercel/frontend/App.tsx +++ b/examples/sandbox-coding-agent-vercel/frontend/App.tsx @@ -79,7 +79,7 @@ function AgentPanel({ info }: { info: AgentInfo }) { return; } - await agent.connection.queue.message.send({ + await agent.connection.send("message", { text: trimmed, sender: "Operator", }); diff --git a/examples/sandbox-coding-agent-vercel/src/actors.ts b/examples/sandbox-coding-agent-vercel/src/actors.ts index 208d90e1cf..2788c65b76 100644 --- a/examples/sandbox-coding-agent-vercel/src/actors.ts +++ b/examples/sandbox-coding-agent-vercel/src/actors.ts @@ -1,5 +1,5 @@ import { SandboxAgent } from "sandbox-agent"; -import { actor, setup } from "rivetkit"; +import { actor, event, queue, setup } from "rivetkit"; export type AgentMessage = { id: string; @@ -26,6 +26,14 @@ export type AgentQueueMessage = { sender?: string; }; +export type AgentResponseEvent = { + messageId: string; + delta: string; + content: string; + done: boolean; + error?: string; +}; + type ItemContentPart = { type?: string; text?: string; @@ -121,13 +129,19 @@ export const agent = actor({ } as AgentStatus, sessionId: null as string | null, }, + queues: { + message: queue(), + }, + events: { + messageAdded: event(), + status: event(), + response: event(), + }, // The run hook keeps the agent listening for queued messages. run: async (c) => { - while (true) { - const queued = await c.queue.next("message"); - - const body = queued.body as AgentQueueMessage; + for await (const queued of c.queue.iter()) { + const { body } = queued; if (!body?.text || typeof body.text !== "string") { continue; } @@ -178,17 +192,16 @@ export const agent = actor({ c.state.sessionId = sessionId; } - const eventStream = await client.streamTurn(sessionId, { - message: userMessage.content, - sync: false, - }); + const eventStream = await client.streamTurn(sessionId, { + message: userMessage.content, + }); let content = ""; let assistantItemId: string | null = null; for await (const rawEvent of eventStream) { const event = rawEvent as StreamEvent; - if (c.abortSignal.aborted) { + if (c.aborted) { break; } diff --git a/examples/sandbox-coding-agent/README.md b/examples/sandbox-coding-agent/README.md index d52525c9c1..ed8ff06015 100644 --- a/examples/sandbox-coding-agent/README.md +++ b/examples/sandbox-coding-agent/README.md @@ -15,7 +15,7 @@ npm run dev ## Features - Actor-per-agent pattern with a coordinating manager Rivet Actor -- Queue-based intake using `c.queue.next` inside the run loop +- Queue-based intake using `for await` over `c.queue.iter(...)` inside the run loop - Sandbox Agent SDK sessions per agent with streamed output - Persistent history stored in Rivet Actor state - Live status updates via events and polling diff --git a/examples/sandbox-coding-agent/frontend/App.tsx b/examples/sandbox-coding-agent/frontend/App.tsx index 4b8bcf2107..e991b81d9d 100644 --- a/examples/sandbox-coding-agent/frontend/App.tsx +++ b/examples/sandbox-coding-agent/frontend/App.tsx @@ -79,7 +79,7 @@ function AgentPanel({ info }: { info: AgentInfo }) { return; } - await agent.connection.queue.message.send({ + await agent.connection.send("message", { text: trimmed, sender: "Operator", }); diff --git a/examples/sandbox-coding-agent/src/actors.ts b/examples/sandbox-coding-agent/src/actors.ts index 208d90e1cf..2788c65b76 100644 --- a/examples/sandbox-coding-agent/src/actors.ts +++ b/examples/sandbox-coding-agent/src/actors.ts @@ -1,5 +1,5 @@ import { SandboxAgent } from "sandbox-agent"; -import { actor, setup } from "rivetkit"; +import { actor, event, queue, setup } from "rivetkit"; export type AgentMessage = { id: string; @@ -26,6 +26,14 @@ export type AgentQueueMessage = { sender?: string; }; +export type AgentResponseEvent = { + messageId: string; + delta: string; + content: string; + done: boolean; + error?: string; +}; + type ItemContentPart = { type?: string; text?: string; @@ -121,13 +129,19 @@ export const agent = actor({ } as AgentStatus, sessionId: null as string | null, }, + queues: { + message: queue(), + }, + events: { + messageAdded: event(), + status: event(), + response: event(), + }, // The run hook keeps the agent listening for queued messages. run: async (c) => { - while (true) { - const queued = await c.queue.next("message"); - - const body = queued.body as AgentQueueMessage; + for await (const queued of c.queue.iter()) { + const { body } = queued; if (!body?.text || typeof body.text !== "string") { continue; } @@ -178,17 +192,16 @@ export const agent = actor({ c.state.sessionId = sessionId; } - const eventStream = await client.streamTurn(sessionId, { - message: userMessage.content, - sync: false, - }); + const eventStream = await client.streamTurn(sessionId, { + message: userMessage.content, + }); let content = ""; let assistantItemId: string | null = null; for await (const rawEvent of eventStream) { const event = rawEvent as StreamEvent; - if (c.abortSignal.aborted) { + if (c.aborted) { break; } diff --git a/examples/sandbox-vercel/frontend/App.tsx b/examples/sandbox-vercel/frontend/App.tsx index 44b7231e27..9366a6c15b 100644 --- a/examples/sandbox-vercel/frontend/App.tsx +++ b/examples/sandbox-vercel/frontend/App.tsx @@ -23,8 +23,6 @@ import { type PageConfig, } from "./page-data.ts"; -type ActorName = (typeof registry)["config"]["use"] extends Record ? K & string : never; - const GROUP_ICONS: Record> = { compass: Compass, code: Code, @@ -75,6 +73,28 @@ const { useActor } = createRivetKit( `${location.origin}/api/rivet`, ); +type ActorPanelActor = { + connStatus: string | null; + error: unknown; + handle: { + action: (request: { name: string; args: unknown[] }) => Promise; + fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise; + webSocket: () => Promise; + } | null; + connection: { + on: (event: string, callback: (...args: unknown[]) => void) => () => void; + } | null; +}; + +type LooseActorHook = (options: { + name: string; + key: string | string[]; + params?: Record; + createWithInput?: unknown; +}) => ActorPanelActor; + +const useActorLoose = useActor as unknown as LooseActorHook; + type JsonResult = { ok: true; value: T } | { ok: false; error: string }; function parseJson(value: string): JsonResult { @@ -305,8 +325,8 @@ function ActorView({ actorName, page }: { actorName: string; page: PageConfig }) ? parsedInput.value : undefined; - const actor = useActor({ - name: actorName as ActorName, + const actor = useActorLoose({ + name: actorName, key: parsedKey.ok ? parsedKey.value : "demo", params: resolvedParams, createWithInput: resolvedInput, @@ -401,7 +421,7 @@ function StatePanel({ stateAction, refreshTrigger, }: { - actor: ReturnType; + actor: ActorPanelActor; stateAction: string; refreshTrigger: number; }) { @@ -464,7 +484,7 @@ function StatePanel({ // ── Events Panel ────────────────────────────────── -function EventsPanel({ actor }: { actor: ReturnType }) { +function EventsPanel({ actor }: { actor: ActorPanelActor }) { const [eventName, setEventName] = useState(""); const [events, setEvents] = useState>([]); @@ -554,7 +574,7 @@ function ActionRunner({ templates, onActionComplete, }: { - actor: ReturnType; + actor: ActorPanelActor; templates: ActionTemplate[]; onActionComplete?: () => void; }) { @@ -743,8 +763,8 @@ function RawHttpPanel({ page }: { page: PageConfig }) { setBody(page.rawHttpDefaults?.body ?? ""); }, [page.id, page.actors, page.rawHttpDefaults]); - const actor = useActor({ - name: selectedActor as ActorName, + const actor = useActorLoose({ + name: selectedActor, key: ["demo"], }); @@ -859,8 +879,8 @@ function RawWebSocketPanel({ page }: { page: PageConfig }) { setSelectedActor(page.actors[0] ?? ""); }, [page.id, page.actors]); - const actor = useActor({ - name: selectedActor as ActorName, + const actor = useActorLoose({ + name: selectedActor, key: ["demo"], }); diff --git a/examples/sandbox-vercel/frontend/page-data.ts b/examples/sandbox-vercel/frontend/page-data.ts index 0e3c738026..a15f0ea8db 100644 --- a/examples/sandbox-vercel/frontend/page-data.ts +++ b/examples/sandbox-vercel/frontend/page-data.ts @@ -65,7 +65,7 @@ console.log(metadata.tags, metadata.region);`, });`, kv: `await actor.putText("greeting", "hello"); const value = await actor.getText("greeting");`, - queue: `await actor.queue.work.send({ id: "task-1" }); + queue: `await actor.send("work", { id: "task-1" }); const message = await actor.receiveOne("work");`, workflow: `const workflow = client.order.getOrCreate([orderId]); await workflow.getOrder();`, @@ -247,15 +247,7 @@ export const ACTION_TEMPLATES: Record = { { label: "Ping", action: "ping", args: [] }, { label: "Actor Counts", action: "getActorCounts", args: [] }, ], - queueActor: [ - { label: "Receive One", action: "receiveOne", args: ["work"] }, - ], - sender: [{ label: "Get Messages", action: "getMessages", args: [] }], - multiQueue: [{ label: "Get Messages", action: "getMessages", args: [] }], - timeout: [{ label: "Wait", action: "waitForMessage", args: [2000] }], worker: [{ label: "Get State", action: "getState", args: [] }], - selfSender: [{ label: "Get State", action: "getState", args: [] }], - keepAwake: [{ label: "Get State", action: "getState", args: [] }], order: [{ label: "Get Order", action: "getOrder", args: [] }], timer: [{ label: "Get Timer", action: "getTimer", args: [] }], batch: [{ label: "Get Job", action: "getJob", args: [] }], @@ -808,40 +800,33 @@ export const PAGE_GROUPS: PageGroup[] = [ title: "Queues", icon: "list", pages: [ - { - id: "queue-basics", - title: "Queue Basics", - description: "Send and receive queue messages from actors.", + { + id: "queue-basics", + title: "Queue Basics", + description: "Send and receive queue messages from actors.", docs: [ { label: "Queue", href: "https://rivet.dev/docs/actors/queue", }, ], - actors: ["queueActor", "queueLimitedActor"], - snippet: SNIPPETS.queue, - }, - { - id: "queue-patterns", - title: "Queue Patterns", - description: - "Explore sending, timeouts, workers, and keep-awake patterns.", + actors: ["worker"], + snippet: SNIPPETS.queue, + }, + { + id: "queue-patterns", + title: "Queue Patterns", + description: + "Run a worker loop that consumes queued jobs.", docs: [ { label: "Queue", href: "https://rivet.dev/docs/actors/queue", }, ], - actors: [ - "sender", - "multiQueue", - "timeout", - "worker", - "selfSender", - "keepAwake", - ], - snippet: SNIPPETS.queue, - }, + actors: ["worker"], + snippet: SNIPPETS.queue, + }, { id: "queue-run-loop", title: "Queue in Run Loop", diff --git a/examples/sandbox-vercel/src/actors.ts b/examples/sandbox-vercel/src/actors.ts index 47bcfe5d6c..efc9aff6c9 100644 --- a/examples/sandbox-vercel/src/actors.ts +++ b/examples/sandbox-vercel/src/actors.ts @@ -78,13 +78,7 @@ import { } from "./actors/lifecycle/destroy.ts"; import { hibernationActor } from "./actors/lifecycle/hibernation.ts"; // Queues -import { queueActor, queueLimitedActor } from "./actors/queue/queue.ts"; -import { sender } from "./actors/queue/sender.ts"; -import { multiQueue } from "./actors/queue/multi-queue.ts"; -import { timeout } from "./actors/queue/timeout.ts"; import { worker } from "./actors/queue/worker.ts"; -import { selfSender } from "./actors/queue/self-sender.ts"; -import { keepAwake } from "./actors/queue/keep-awake.ts"; // Workflows import { workflowCounterActor, @@ -179,14 +173,7 @@ export const registry = setup({ destroyObserver, hibernationActor, // Queues - queueActor, - queueLimitedActor, - sender, - multiQueue, - timeout, worker, - selfSender, - keepAwake, // Workflows timer, order, diff --git a/examples/sandbox-vercel/src/actors/ai/ai-agent.ts b/examples/sandbox-vercel/src/actors/ai/ai-agent.ts index 94c957119b..1615f2ec7d 100644 --- a/examples/sandbox-vercel/src/actors/ai/ai-agent.ts +++ b/examples/sandbox-vercel/src/actors/ai/ai-agent.ts @@ -1,6 +1,6 @@ import { openai } from "@ai-sdk/openai"; import { generateText, tool } from "ai"; -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { z } from "zod"; import { getWeather } from "./my-tools.ts"; import type { Message } from "./types.ts"; @@ -10,6 +10,9 @@ export const aiAgent = actor({ state: { messages: [] as Message[], }, + events: { + messageReceived: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/sandbox-vercel/src/actors/connections/conn-state.ts b/examples/sandbox-vercel/src/actors/connections/conn-state.ts index 8312f5aa3e..c18244cdb4 100644 --- a/examples/sandbox-vercel/src/actors/connections/conn-state.ts +++ b/examples/sandbox-vercel/src/actors/connections/conn-state.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export type ConnState = { username: string; @@ -13,6 +13,11 @@ export const connStateActor = actor({ sharedCounter: 0, disconnectionCount: 0, }, + events: { + userConnected: event<{ id: string; username: string; role: string }>(), + userDisconnected: event<{ id: string }>(), + directMessage: event<{ from: string; message: string }>(), + }, // Define connection state createConnState: ( c, diff --git a/examples/sandbox-vercel/src/actors/counter/conn-params.ts b/examples/sandbox-vercel/src/actors/counter/conn-params.ts index 4116a4432c..f0e05b614f 100644 --- a/examples/sandbox-vercel/src/actors/counter/conn-params.ts +++ b/examples/sandbox-vercel/src/actors/counter/conn-params.ts @@ -1,7 +1,10 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const counterWithParams = actor({ state: { count: 0, initializers: [] as string[] }, + events: { + newCount: event<{ count: number; by: string }>(), + }, createConnState: (c, params: { name?: string }) => { return { name: params.name || "anonymous", diff --git a/examples/sandbox-vercel/src/actors/counter/counter-conn.ts b/examples/sandbox-vercel/src/actors/counter/counter-conn.ts index 93d4a6c0af..f06b1b119a 100644 --- a/examples/sandbox-vercel/src/actors/counter/counter-conn.ts +++ b/examples/sandbox-vercel/src/actors/counter/counter-conn.ts @@ -1,9 +1,12 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const counterConn = actor({ state: { connectionCount: 0, }, + events: { + newCount: event(), + }, connState: { count: 0 }, onConnect: (c, conn) => { c.state.connectionCount += 1; diff --git a/examples/sandbox-vercel/src/actors/counter/counter.ts b/examples/sandbox-vercel/src/actors/counter/counter.ts index 8f14bea1c2..9fd224e818 100644 --- a/examples/sandbox-vercel/src/actors/counter/counter.ts +++ b/examples/sandbox-vercel/src/actors/counter/counter.ts @@ -1,7 +1,10 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/sandbox-vercel/src/actors/lifecycle/run.ts b/examples/sandbox-vercel/src/actors/lifecycle/run.ts index 42dce0de96..a6858e5349 100644 --- a/examples/sandbox-vercel/src/actors/lifecycle/run.ts +++ b/examples/sandbox-vercel/src/actors/lifecycle/run.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event, queue } from "rivetkit"; import type { registry } from "../../actors.ts"; export const RUN_SLEEP_TIMEOUT = 500; @@ -15,7 +15,7 @@ export const runWithTicks = actor({ c.state.runStarted = true; c.log.info("run handler started"); - while (!c.abortSignal.aborted) { + while (!c.aborted) { c.state.tickCount += 1; c.state.lastTickAt = Date.now(); c.log.info({ msg: "tick", tickCount: c.state.tickCount }); @@ -57,19 +57,19 @@ export const runWithQueueConsumer = actor({ messagesReceived: [] as Array<{ name: string; body: unknown }>, runStarted: false, }, + queues: { + messages: queue(), + }, run: async (c) => { c.state.runStarted = true; c.log.info("run handler started, waiting for messages"); - while (!c.abortSignal.aborted) { - const message = await c.queue.next("messages", { timeout: 100 }); - if (message) { - c.log.info({ msg: "received message", body: message.body }); - c.state.messagesReceived.push({ - name: message.name, - body: message.body, - }); - } + for await (const message of c.queue.iter()) { + c.log.info({ msg: "received message", body: message.body }); + c.state.messagesReceived.push({ + name: message.name, + body: message.body, + }); } c.log.info("run handler exiting gracefully"); @@ -82,7 +82,7 @@ export const runWithQueueConsumer = actor({ sendMessage: async (c, body: unknown) => { const client = c.client(); const handle = client.runWithQueueConsumer.getForId(c.actorId); - await handle.queue.messages.send(body); + await handle.send("messages", body); return true; }, }, diff --git a/examples/sandbox-vercel/src/actors/lifecycle/scheduled.ts b/examples/sandbox-vercel/src/actors/lifecycle/scheduled.ts index 7bac35bac8..b11b802bd0 100644 --- a/examples/sandbox-vercel/src/actors/lifecycle/scheduled.ts +++ b/examples/sandbox-vercel/src/actors/lifecycle/scheduled.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const scheduled = actor({ state: { @@ -6,6 +6,10 @@ export const scheduled = actor({ scheduledCount: 0, taskHistory: [] as string[], }, + events: { + scheduled: event<{ time: number; count: number }>(), + scheduledWithId: event<{ taskId: string; time: number; count: number }>(), + }, actions: { // Schedule using 'at' with specific timestamp scheduleTaskAt: (c, timestamp: number) => { diff --git a/examples/sandbox-vercel/src/actors/lifecycle/sleep.ts b/examples/sandbox-vercel/src/actors/lifecycle/sleep.ts index 5717a83f07..7414678e71 100644 --- a/examples/sandbox-vercel/src/actors/lifecycle/sleep.ts +++ b/examples/sandbox-vercel/src/actors/lifecycle/sleep.ts @@ -1,4 +1,4 @@ -import { actor, type UniversalWebSocket } from "rivetkit"; +import { actor, event, type UniversalWebSocket } from "rivetkit"; import { promiseWithResolvers } from "rivetkit/utils"; export const SLEEP_TIMEOUT = 1000; @@ -37,6 +37,9 @@ export const sleepWithLongRpc = actor({ state: { startCount: 0, sleepCount: 0 }, createVars: () => ({}) as { longRunningResolve: PromiseWithResolvers }, + events: { + waiting: event<[]>(), + }, onWake: (c) => { c.state.startCount += 1; }, diff --git a/examples/sandbox-vercel/src/actors/queue/keep-awake.ts b/examples/sandbox-vercel/src/actors/queue/keep-awake.ts deleted file mode 100644 index dcf9a0a28e..0000000000 --- a/examples/sandbox-vercel/src/actors/queue/keep-awake.ts +++ /dev/null @@ -1,70 +0,0 @@ -import { actor } from "rivetkit"; - -export interface CurrentTask { - id: string; - startedAt: number; - durationMs: number; -} - -export interface CompletedTask { - id: string; - completedAt: number; -} - -export interface KeepAwakeState { - currentTask: CurrentTask | null; - completedTasks: CompletedTask[]; -} - -export const keepAwake = actor({ - state: { - currentTask: null as CurrentTask | null, - completedTasks: [] as CompletedTask[], - }, - async run(c) { - while (!c.abortSignal.aborted) { - const job = await c.queue.next("tasks", { timeout: 1000 }); - if (job) { - const taskId = crypto.randomUUID(); - const { durationMs } = job.body as { durationMs: number }; - - c.state.currentTask = { - id: taskId, - startedAt: Date.now(), - durationMs, - }; - c.broadcast("taskStarted", c.state.currentTask); - - // Wrap long-running work in keepAwake so actor doesn't sleep - await c.keepAwake( - new Promise((resolve) => setTimeout(resolve, durationMs)), - ); - - c.state.completedTasks.push({ id: taskId, completedAt: Date.now() }); - c.state.currentTask = null; - c.broadcast("taskCompleted", { - taskId, - completedTasks: c.state.completedTasks, - }); - } - } - }, - actions: { - getState(c): KeepAwakeState { - return { - currentTask: c.state.currentTask, - completedTasks: c.state.completedTasks, - }; - }, - clearTasks(c) { - c.state.completedTasks = []; - c.broadcast("taskCompleted", { - taskId: null, - completedTasks: [], - }); - }, - }, - options: { - sleepTimeout: 2000, - }, -}); diff --git a/examples/sandbox-vercel/src/actors/queue/multi-queue.ts b/examples/sandbox-vercel/src/actors/queue/multi-queue.ts deleted file mode 100644 index 57d57e7ad8..0000000000 --- a/examples/sandbox-vercel/src/actors/queue/multi-queue.ts +++ /dev/null @@ -1,31 +0,0 @@ -import { actor } from "rivetkit"; - -export interface QueueMessage { - name: string; - body: unknown; -} - -export const multiQueue = actor({ - state: { - messages: [] as QueueMessage[], - }, - actions: { - async receiveFromQueues(c, names: string[], count: number) { - const msgs = await c.queue.next(names, { count, timeout: 100 }); - if (msgs && msgs.length > 0) { - for (const msg of msgs) { - c.state.messages.push({ name: msg.name, body: msg.body }); - } - c.broadcast("messagesReceived", c.state.messages); - } - return msgs ?? []; - }, - getMessages(c): QueueMessage[] { - return c.state.messages; - }, - clearMessages(c) { - c.state.messages = []; - c.broadcast("messagesReceived", c.state.messages); - }, - }, -}); diff --git a/examples/sandbox-vercel/src/actors/queue/queue.ts b/examples/sandbox-vercel/src/actors/queue/queue.ts deleted file mode 100644 index d77e9a77ff..0000000000 --- a/examples/sandbox-vercel/src/actors/queue/queue.ts +++ /dev/null @@ -1,123 +0,0 @@ -import { actor } from "rivetkit"; -import type { registry } from "../../actors.ts"; - -export const queueActor = actor({ - state: {}, - actions: { - receiveOne: async ( - c, - name: string, - opts?: { count?: number; timeout?: number }, - ) => { - const message = await c.queue.next(name, opts); - if (!message) { - return null; - } - return { name: message.name, body: message.body }; - }, - receiveMany: async ( - c, - names: string[], - opts?: { count?: number; timeout?: number }, - ) => { - const messages = await c.queue.next(names, opts); - return (messages ?? []).map( - (message: { name: string; body: unknown }) => ({ - name: message.name, - body: message.body, - }), - ); - }, - receiveRequest: async ( - c, - request: { - name: string | string[]; - count?: number; - timeout?: number; - }, - ) => { - const messages = await c.queue.next(request); - return (messages ?? []).map( - (message: { name: string; body: unknown }) => ({ - name: message.name, - body: message.body, - }), - ); - }, - sendToSelf: async (c, name: string, body: unknown) => { - const client = c.client(); - const handle = client.queueActor.getForId(c.actorId); - await handle.queue[name].send(body); - return true; - }, - waitForAbort: async (c) => { - setTimeout(() => { - c.destroy(); - }, 10); - await c.queue.next("abort", { timeout: 10_000 }); - return true; - }, - receiveAndComplete: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; - } - await message.complete({ echo: message.body }); - return { name: message.name, body: message.body }; - }, - receiveAndCompleteTwice: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; - } - await message.complete({ ok: true }); - try { - await message.complete({ ok: false }); - return { ok: false }; - } catch (error) { - const actorError = error as { group?: string; code?: string }; - return { group: actorError.group, code: actorError.code }; - } - }, - receiveWithoutWaitComplete: async (c, name: string) => { - const message = await c.queue.next(name); - if (!message) { - return null; - } - try { - await message.complete(); - return { ok: false }; - } catch (error) { - const actorError = error as { group?: string; code?: string }; - return { group: actorError.group, code: actorError.code }; - } - }, - receiveWhilePending: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; - } - let errorPayload: { group?: string; code?: string } | undefined; - try { - await c.queue.next(name); - } catch (error) { - const actorError = error as { group?: string; code?: string }; - errorPayload = { - group: actorError.group, - code: actorError.code, - }; - } - await message.complete({ ok: true }); - return errorPayload ?? { ok: false }; - }, - }, -}); - -export const queueLimitedActor = actor({ - state: {}, - actions: {}, - options: { - maxQueueSize: 1, - maxQueueMessageSize: 64, - }, -}); diff --git a/examples/sandbox-vercel/src/actors/queue/self-sender.ts b/examples/sandbox-vercel/src/actors/queue/self-sender.ts deleted file mode 100644 index bfd1b3cd1a..0000000000 --- a/examples/sandbox-vercel/src/actors/queue/self-sender.ts +++ /dev/null @@ -1,48 +0,0 @@ -import { actor } from "rivetkit"; - -export interface SelfSenderState { - sentCount: number; - receivedCount: number; - messages: unknown[]; -} - -export const selfSender = actor({ - state: { - sentCount: 0, - receivedCount: 0, - messages: [] as unknown[], - }, - actions: { - async receiveFromSelf(c) { - const msg = await c.queue.next("self", { timeout: 100 }); - if (msg) { - c.state.receivedCount += 1; - c.state.messages.push(msg.body); - c.broadcast("received", { - receivedCount: c.state.receivedCount, - message: msg.body, - }); - return msg.body; - } - return null; - }, - getState(c): SelfSenderState { - return { - sentCount: c.state.sentCount, - receivedCount: c.state.receivedCount, - messages: c.state.messages, - }; - }, - clearMessages(c) { - c.state.sentCount = 0; - c.state.receivedCount = 0; - c.state.messages = []; - c.broadcast("sent", { sentCount: 0 }); - c.broadcast("received", { receivedCount: 0, message: null }); - }, - incrementSentCount(c) { - c.state.sentCount += 1; - c.broadcast("sent", { sentCount: c.state.sentCount }); - }, - }, -}); diff --git a/examples/sandbox-vercel/src/actors/queue/sender.ts b/examples/sandbox-vercel/src/actors/queue/sender.ts deleted file mode 100644 index 0db2c49a0b..0000000000 --- a/examples/sandbox-vercel/src/actors/queue/sender.ts +++ /dev/null @@ -1,36 +0,0 @@ -import { actor } from "rivetkit"; - -export interface ReceivedMessage { - name: string; - body: unknown; - receivedAt: number; -} - -export const sender = actor({ - state: { - messages: [] as ReceivedMessage[], - }, - actions: { - getMessages(c): ReceivedMessage[] { - return c.state.messages; - }, - async receiveOne(c) { - const msg = await c.queue.next("task", { timeout: 100 }); - if (msg) { - const received: ReceivedMessage = { - name: msg.name, - body: msg.body, - receivedAt: Date.now(), - }; - c.state.messages.push(received); - c.broadcast("messageReceived", c.state.messages); - return received; - } - return null; - }, - clearMessages(c) { - c.state.messages = []; - c.broadcast("messageReceived", c.state.messages); - }, - }, -}); diff --git a/examples/sandbox-vercel/src/actors/queue/timeout.ts b/examples/sandbox-vercel/src/actors/queue/timeout.ts deleted file mode 100644 index be0ff8fea7..0000000000 --- a/examples/sandbox-vercel/src/actors/queue/timeout.ts +++ /dev/null @@ -1,44 +0,0 @@ -import { actor } from "rivetkit"; - -export interface TimeoutResult { - timedOut: boolean; - message?: unknown; - waitedMs: number; -} - -export interface TimeoutState { - lastResult: TimeoutResult | null; - waitStartedAt: number | null; -} - -export const timeout = actor({ - state: { - lastResult: null as TimeoutResult | null, - waitStartedAt: null as number | null, - }, - actions: { - async waitForMessage(c, timeoutMs: number): Promise { - const startedAt = Date.now(); - c.state.waitStartedAt = startedAt; - c.broadcast("waitStarted", { startedAt, timeoutMs }); - - const msg = await c.queue.next("work", { timeout: timeoutMs }); - - const waitedMs = Date.now() - startedAt; - const result: TimeoutResult = msg - ? { timedOut: false, message: msg.body, waitedMs } - : { timedOut: true, waitedMs }; - - c.state.lastResult = result; - c.state.waitStartedAt = null; - c.broadcast("waitCompleted", result); - return result; - }, - getState(c): TimeoutState { - return { - lastResult: c.state.lastResult, - waitStartedAt: c.state.waitStartedAt, - }; - }, - }, -}); diff --git a/examples/sandbox-vercel/src/actors/queue/worker.ts b/examples/sandbox-vercel/src/actors/queue/worker.ts index e7b710f5dc..41ef9e388e 100644 --- a/examples/sandbox-vercel/src/actors/queue/worker.ts +++ b/examples/sandbox-vercel/src/actors/queue/worker.ts @@ -1,16 +1,28 @@ -import { actor } from "rivetkit"; +import { actor, event, queue } from "rivetkit"; + +export interface WorkerJob { + id: string; + payload: string; +} export interface WorkerState { status: "idle" | "running"; processed: number; - lastJob: unknown; + lastJob: WorkerJob | null; } export const worker = actor({ state: { status: "idle" as "idle" | "running", processed: 0, - lastJob: null as unknown, + lastJob: null as WorkerJob | null, + }, + events: { + statusChanged: event<{ status: "idle" | "running"; processed: number }>(), + jobProcessed: event<{ processed: number; job: WorkerJob }>(), + }, + queues: { + jobs: queue(), }, async run(c) { c.state.status = "running"; @@ -19,16 +31,13 @@ export const worker = actor({ processed: c.state.processed, }); - while (!c.abortSignal.aborted) { - const job = await c.queue.next("jobs", { timeout: 1000 }); - if (job) { - c.state.processed += 1; - c.state.lastJob = job.body; - c.broadcast("jobProcessed", { - processed: c.state.processed, - job: job.body, - }); - } + for await (const job of c.queue.iter()) { + c.state.processed += 1; + c.state.lastJob = job.body; + c.broadcast("jobProcessed", { + processed: c.state.processed, + job: job.body, + }); } c.state.status = "idle"; diff --git a/examples/sandbox-vercel/src/actors/state/sqlite-drizzle/drizzle/migrations.d.ts b/examples/sandbox-vercel/src/actors/state/sqlite-drizzle/drizzle/migrations.d.ts new file mode 100644 index 0000000000..00714979b2 --- /dev/null +++ b/examples/sandbox-vercel/src/actors/state/sqlite-drizzle/drizzle/migrations.d.ts @@ -0,0 +1,6 @@ +declare const migrations: { + journal: unknown; + migrations: Record; +}; + +export default migrations; diff --git a/examples/sandbox-vercel/src/actors/state/vars.ts b/examples/sandbox-vercel/src/actors/state/vars.ts index 7a62319824..cec9116fff 100644 --- a/examples/sandbox-vercel/src/actors/state/vars.ts +++ b/examples/sandbox-vercel/src/actors/state/vars.ts @@ -52,7 +52,7 @@ export const dynamicVarActor = actor({ }; }, actions: { - getVars: (c) => { + getVars: (c: any) => { return c.vars; }, }, @@ -68,7 +68,7 @@ export const uniqueVarActor = actor({ }; }, actions: { - getVars: (c) => { + getVars: (c: any) => { return c.vars; }, }, @@ -84,7 +84,7 @@ export const driverCtxActor = actor({ }; }, actions: { - getVars: (c) => { + getVars: (c: any) => { return c.vars; }, }, diff --git a/examples/sandbox-vercel/src/actors/workflow/approval.ts b/examples/sandbox-vercel/src/actors/workflow/approval.ts index e39af4cf5e..86070ec7eb 100644 --- a/examples/sandbox-vercel/src/actors/workflow/approval.ts +++ b/examples/sandbox-vercel/src/actors/workflow/approval.ts @@ -1,9 +1,9 @@ -// APPROVAL REQUEST (Listen Demo) -// Demonstrates: Message listening with timeout for approval workflows +// APPROVAL REQUEST (Queue Wait Demo) +// Demonstrates: Queue waits with timeout for approval workflows // One actor per approval request - actor key is the request ID -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, event, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; export type RequestStatus = "pending" | "approved" | "rejected" | "timeout"; @@ -21,7 +21,7 @@ export type ApprovalRequest = { type State = ApprovalRequest; -const QUEUE_DECISION = workflowQueueName("decision"); +const QUEUE_DECISION = "decision" as const; const APPROVAL_TIMEOUT_MS = 30000; @@ -30,6 +30,11 @@ export type ApprovalRequestInput = { description?: string; }; +type ApprovalDecision = { + approved: boolean; + approver: string; +}; + export const approval = actor({ createState: (c, input?: ApprovalRequestInput): ApprovalRequest => ({ id: c.key[0] as string, @@ -38,6 +43,13 @@ export const approval = actor({ status: "pending", createdAt: Date.now(), }), + queues: { + decision: queue(), + }, + events: { + requestUpdated: event(), + requestCreated: event(), + }, actions: { getRequest: (c): ApprovalRequest => c.state, @@ -72,10 +84,14 @@ export const approval = actor({ c.broadcast("requestCreated", c.state); }); - const decision = await loopCtx.listenWithTimeout<{ - approved: boolean; - approver: string; - }>("wait-decision", "decision", APPROVAL_TIMEOUT_MS); + const [decisionMessage] = await loopCtx.queue.next( + "wait-decision", + { + names: [QUEUE_DECISION], + timeout: APPROVAL_TIMEOUT_MS, + }, + ); + const decision = decisionMessage?.body ?? null; await loopCtx.step("update-status", async () => { c.state.deciding = false; diff --git a/examples/sandbox-vercel/src/actors/workflow/batch.ts b/examples/sandbox-vercel/src/actors/workflow/batch.ts index 006e894e11..567fb23313 100644 --- a/examples/sandbox-vercel/src/actors/workflow/batch.ts +++ b/examples/sandbox-vercel/src/actors/workflow/batch.ts @@ -2,7 +2,7 @@ // Demonstrates: Loop with persistent state (cursor) for batch processing // One actor per batch job - actor key is the job ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -59,6 +59,11 @@ export const batch = actor({ batches: [], startedAt: Date.now(), }), + events: { + batchProcessed: event(), + stateChanged: event(), + processingComplete: event<{ totalBatches: number; totalItems: number }>(), + }, actions: { getJob: (c): BatchJob => c.state, diff --git a/examples/sandbox-vercel/src/actors/workflow/dashboard.ts b/examples/sandbox-vercel/src/actors/workflow/dashboard.ts index 86887d466e..449e38244f 100644 --- a/examples/sandbox-vercel/src/actors/workflow/dashboard.ts +++ b/examples/sandbox-vercel/src/actors/workflow/dashboard.ts @@ -1,8 +1,8 @@ // DASHBOARD (Join Demo) // Demonstrates: Parallel data fetching with join (wait-all) -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, event, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; export type UserStats = { @@ -45,7 +45,8 @@ export type DashboardState = { type State = DashboardState; -const QUEUE_REFRESH = workflowQueueName("refresh"); +const QUEUE_REFRESH = "refresh"; +type RefreshMessage = Record; async function fetchUserStats(): Promise { await new Promise((r) => setTimeout(r, 800 + Math.random() * 1200)); @@ -87,6 +88,13 @@ export const dashboard = actor({ }, lastRefresh: null as number | null, }, + queues: { + [QUEUE_REFRESH]: queue(), + }, + events: { + stateChanged: event(), + refreshComplete: event(), + }, actions: { refresh: async (c) => { @@ -111,7 +119,9 @@ export const dashboard = actor({ run: async (loopCtx) => { const c = actorCtx(loopCtx); - await loopCtx.listen("wait-refresh", "refresh"); + await loopCtx.queue.next("wait-refresh", { + names: [QUEUE_REFRESH], + }); ctx.log.info({ msg: "starting dashboard refresh" }); diff --git a/examples/sandbox-vercel/src/actors/workflow/history-examples.ts b/examples/sandbox-vercel/src/actors/workflow/history-examples.ts index dd0b0e5fe5..b70a1f1f19 100644 --- a/examples/sandbox-vercel/src/actors/workflow/history-examples.ts +++ b/examples/sandbox-vercel/src/actors/workflow/history-examples.ts @@ -1,5 +1,5 @@ -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, event, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; function delay(ms: number): Promise { @@ -250,19 +250,38 @@ export type WorkflowHistoryFullState = { completedAt?: number; }; -type MessageSeed = { name: string; payload: unknown }; +const QUEUE_ORDER_CREATED = "order:created"; +const QUEUE_ORDER_UPDATED = "order:updated"; +const QUEUE_ORDER_ITEM = "order:item"; +const QUEUE_ORDER_ARTIFACT = "order:artifact"; +const QUEUE_ORDER_READY = "order:ready"; +const QUEUE_ORDER_OPTIONAL = "order:optional"; + +type OrderCreatedMessage = { id: string }; +type OrderUpdatedMessage = { id: string; status: string }; +type OrderItemMessage = { sku: string; qty: number }; +type OrderArtifactMessage = { artifactId: string }; +type OrderReadyMessage = { batch: number }; +type OrderOptionalMessage = { note?: string }; + +type MessageSeed = + | { name: typeof QUEUE_ORDER_CREATED; payload: OrderCreatedMessage } + | { name: typeof QUEUE_ORDER_UPDATED; payload: OrderUpdatedMessage } + | { name: typeof QUEUE_ORDER_ITEM; payload: OrderItemMessage } + | { name: typeof QUEUE_ORDER_ARTIFACT; payload: OrderArtifactMessage } + | { name: typeof QUEUE_ORDER_READY; payload: OrderReadyMessage }; const FULL_WORKFLOW_MESSAGE_SEEDS: MessageSeed[] = [ - { name: "order:created", payload: { id: "order-1" } }, - { name: "order:updated", payload: { id: "order-1", status: "paid" } }, - { name: "order:item", payload: { sku: "sku-0", qty: 1 } }, - { name: "order:item", payload: { sku: "sku-4", qty: 1 } }, - { name: "order:artifact", payload: { artifactId: "artifact-0" } }, - { name: "order:artifact", payload: { artifactId: "artifact-1" } }, - { name: "order:artifact", payload: { artifactId: "artifact-2" } }, - { name: "order:ready", payload: { batch: 3 } }, - { name: "order:ready", payload: { batch: 0 } }, - { name: "order:ready", payload: { batch: 2 } }, + { name: QUEUE_ORDER_CREATED, payload: { id: "order-1" } }, + { name: QUEUE_ORDER_UPDATED, payload: { id: "order-1", status: "paid" } }, + { name: QUEUE_ORDER_ITEM, payload: { sku: "sku-0", qty: 1 } }, + { name: QUEUE_ORDER_ITEM, payload: { sku: "sku-4", qty: 1 } }, + { name: QUEUE_ORDER_ARTIFACT, payload: { artifactId: "artifact-0" } }, + { name: QUEUE_ORDER_ARTIFACT, payload: { artifactId: "artifact-1" } }, + { name: QUEUE_ORDER_ARTIFACT, payload: { artifactId: "artifact-2" } }, + { name: QUEUE_ORDER_READY, payload: { batch: 3 } }, + { name: QUEUE_ORDER_READY, payload: { batch: 0 } }, + { name: QUEUE_ORDER_READY, payload: { batch: 2 } }, ]; const FULL_WORKFLOW_ITEMS = [ @@ -278,12 +297,20 @@ export const workflowHistoryFull = actor({ status: "pending", seededMessages: false, }), + queues: { + [QUEUE_ORDER_CREATED]: queue(), + [QUEUE_ORDER_UPDATED]: queue(), + [QUEUE_ORDER_ITEM]: queue(), + [QUEUE_ORDER_ARTIFACT]: queue(), + [QUEUE_ORDER_READY]: queue(), + [QUEUE_ORDER_OPTIONAL]: queue(), + }, actions: { getState: (c): WorkflowHistoryFullState => c.state, seedMessages: async (c) => { if (c.state.seededMessages) return; for (const seed of FULL_WORKFLOW_MESSAGE_SEEDS) { - await c.queue.send(workflowQueueName(seed.name), seed.payload); + await c.queue.send(seed.name, seed.payload); } c.state.seededMessages = true; }, @@ -381,31 +408,36 @@ export const workflowHistoryFull = actor({ return { readyBy, readyBatchBy }; }); - await ctx.listenN("listen-order-created", "order:created", 1); - await ctx.listenWithTimeout( - "listen-order-updated-timeout", - "order:updated", - 250, - ); - await ctx.listenN("listen-batch-two", "order:item", 2); - await ctx.listenNWithTimeout( - "listen-artifacts-timeout", - "order:artifact", - 3, - 300, - ); - await ctx.listenWithTimeout("listen-optional", "order:optional", 200); - await ctx.listenUntil( - "listen-until", - "order:ready", - Date.now() + 300, - ); - await ctx.listenNUntil( - "listen-batch-until", - "order:ready", - 2, - Date.now() + 400, - ); + await ctx.queue.next("listen-order-created", { + names: [QUEUE_ORDER_CREATED], + count: 1, + }); + await ctx.queue.next("listen-order-updated-timeout", { + names: [QUEUE_ORDER_UPDATED], + timeout: 250, + }); + await ctx.queue.next("listen-batch-two", { + names: [QUEUE_ORDER_ITEM], + count: 2, + }); + await ctx.queue.next("listen-artifacts-timeout", { + names: [QUEUE_ORDER_ARTIFACT], + count: 3, + timeout: 300, + }); + await ctx.queue.next("listen-optional", { + names: [QUEUE_ORDER_OPTIONAL], + timeout: 200, + }); + await ctx.queue.next("listen-until", { + names: [QUEUE_ORDER_READY], + timeout: 300, + }); + await ctx.queue.next("listen-batch-until", { + names: [QUEUE_ORDER_READY], + count: 2, + timeout: 400, + }); await ctx.join("join-dependencies", { inventory: { diff --git a/examples/sandbox-vercel/src/actors/workflow/order.ts b/examples/sandbox-vercel/src/actors/workflow/order.ts index 12b5344744..ad0c0db237 100644 --- a/examples/sandbox-vercel/src/actors/workflow/order.ts +++ b/examples/sandbox-vercel/src/actors/workflow/order.ts @@ -2,7 +2,7 @@ // Demonstrates: Sequential workflow steps with automatic retries // One actor per order - actor key is the order ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -41,6 +41,9 @@ export const order = actor({ step: 0, createdAt: Date.now(), }), + events: { + orderUpdated: event(), + }, actions: { getOrder: (c): Order => c.state, diff --git a/examples/sandbox-vercel/src/actors/workflow/payment.ts b/examples/sandbox-vercel/src/actors/workflow/payment.ts index f227b6945d..13997b5e7b 100644 --- a/examples/sandbox-vercel/src/actors/workflow/payment.ts +++ b/examples/sandbox-vercel/src/actors/workflow/payment.ts @@ -2,7 +2,7 @@ // Demonstrates: Rollback checkpoints with compensating actions // One actor per transaction - actor key is the transaction ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -51,6 +51,11 @@ export const payment = actor({ ], startedAt: Date.now(), }), + events: { + transactionStarted: event(), + transactionUpdated: event(), + transactionCompleted: event(), + }, actions: { getTransaction: (c): Transaction => c.state, diff --git a/examples/sandbox-vercel/src/actors/workflow/race.ts b/examples/sandbox-vercel/src/actors/workflow/race.ts index 0e7ad1ed3c..922c526e1e 100644 --- a/examples/sandbox-vercel/src/actors/workflow/race.ts +++ b/examples/sandbox-vercel/src/actors/workflow/race.ts @@ -2,7 +2,7 @@ // Demonstrates: Race (parallel first-wins) for timeout patterns // One actor per race task - actor key is the task ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -32,6 +32,10 @@ export const race = actor({ status: "running", startedAt: Date.now(), }), + events: { + raceStarted: event(), + raceCompleted: event(), + }, actions: { getTask: (c): RaceTask => c.state, diff --git a/examples/sandbox-vercel/src/actors/workflow/timer.ts b/examples/sandbox-vercel/src/actors/workflow/timer.ts index dec43ffbec..22c0c80636 100644 --- a/examples/sandbox-vercel/src/actors/workflow/timer.ts +++ b/examples/sandbox-vercel/src/actors/workflow/timer.ts @@ -2,7 +2,7 @@ // Demonstrates: Durable sleep that survives restarts // One actor per timer - actor key is the timer ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -28,6 +28,10 @@ export const timer = actor({ durationMs: input?.durationMs ?? 10000, startedAt: Date.now(), }), + events: { + timerStarted: event(), + timerCompleted: event(), + }, actions: { getTimer: (c): Timer => c.state, diff --git a/examples/sandbox-vercel/src/actors/workflow/workflow-fixtures.ts b/examples/sandbox-vercel/src/actors/workflow/workflow-fixtures.ts index 8d7e06e951..3deb92001f 100644 --- a/examples/sandbox-vercel/src/actors/workflow/workflow-fixtures.ts +++ b/examples/sandbox-vercel/src/actors/workflow/workflow-fixtures.ts @@ -1,5 +1,5 @@ -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; const WORKFLOW_GUARD_KV_KEY = "__rivet_actor_workflow_guard_triggered"; @@ -50,18 +50,25 @@ export const workflowQueueActor = actor({ state: { received: [] as unknown[], }, + queues: { + [WORKFLOW_QUEUE_NAME]: queue(), + }, run: workflow(async (ctx) => { await ctx.loop({ name: "queue", run: async (loopCtx) => { const actorLoopCtx = loopCtx as any; - const message = await loopCtx.listen( - "queue-wait", - WORKFLOW_QUEUE_NAME, - ); + const [message] = await loopCtx.queue.next("queue-wait", { + names: [WORKFLOW_QUEUE_NAME], + completable: true, + }); + if (!message || !message.complete) { + return Loop.continue(undefined); + } + const complete = message.complete; await loopCtx.step("store-message", async () => { actorLoopCtx.state.received.push(message.body); - await message.complete({ echo: message.body }); + await complete({ echo: message.body }); }); return Loop.continue(undefined); }, @@ -97,4 +104,4 @@ export const workflowSleepActor = actor({ }, }); -export { WORKFLOW_QUEUE_NAME, workflowQueueName }; +export { WORKFLOW_QUEUE_NAME }; diff --git a/examples/sandbox/frontend/App.tsx b/examples/sandbox/frontend/App.tsx index 44b7231e27..9366a6c15b 100644 --- a/examples/sandbox/frontend/App.tsx +++ b/examples/sandbox/frontend/App.tsx @@ -23,8 +23,6 @@ import { type PageConfig, } from "./page-data.ts"; -type ActorName = (typeof registry)["config"]["use"] extends Record ? K & string : never; - const GROUP_ICONS: Record> = { compass: Compass, code: Code, @@ -75,6 +73,28 @@ const { useActor } = createRivetKit( `${location.origin}/api/rivet`, ); +type ActorPanelActor = { + connStatus: string | null; + error: unknown; + handle: { + action: (request: { name: string; args: unknown[] }) => Promise; + fetch: (input: RequestInfo | URL, init?: RequestInit) => Promise; + webSocket: () => Promise; + } | null; + connection: { + on: (event: string, callback: (...args: unknown[]) => void) => () => void; + } | null; +}; + +type LooseActorHook = (options: { + name: string; + key: string | string[]; + params?: Record; + createWithInput?: unknown; +}) => ActorPanelActor; + +const useActorLoose = useActor as unknown as LooseActorHook; + type JsonResult = { ok: true; value: T } | { ok: false; error: string }; function parseJson(value: string): JsonResult { @@ -305,8 +325,8 @@ function ActorView({ actorName, page }: { actorName: string; page: PageConfig }) ? parsedInput.value : undefined; - const actor = useActor({ - name: actorName as ActorName, + const actor = useActorLoose({ + name: actorName, key: parsedKey.ok ? parsedKey.value : "demo", params: resolvedParams, createWithInput: resolvedInput, @@ -401,7 +421,7 @@ function StatePanel({ stateAction, refreshTrigger, }: { - actor: ReturnType; + actor: ActorPanelActor; stateAction: string; refreshTrigger: number; }) { @@ -464,7 +484,7 @@ function StatePanel({ // ── Events Panel ────────────────────────────────── -function EventsPanel({ actor }: { actor: ReturnType }) { +function EventsPanel({ actor }: { actor: ActorPanelActor }) { const [eventName, setEventName] = useState(""); const [events, setEvents] = useState>([]); @@ -554,7 +574,7 @@ function ActionRunner({ templates, onActionComplete, }: { - actor: ReturnType; + actor: ActorPanelActor; templates: ActionTemplate[]; onActionComplete?: () => void; }) { @@ -743,8 +763,8 @@ function RawHttpPanel({ page }: { page: PageConfig }) { setBody(page.rawHttpDefaults?.body ?? ""); }, [page.id, page.actors, page.rawHttpDefaults]); - const actor = useActor({ - name: selectedActor as ActorName, + const actor = useActorLoose({ + name: selectedActor, key: ["demo"], }); @@ -859,8 +879,8 @@ function RawWebSocketPanel({ page }: { page: PageConfig }) { setSelectedActor(page.actors[0] ?? ""); }, [page.id, page.actors]); - const actor = useActor({ - name: selectedActor as ActorName, + const actor = useActorLoose({ + name: selectedActor, key: ["demo"], }); diff --git a/examples/sandbox/frontend/page-data.ts b/examples/sandbox/frontend/page-data.ts index 0e3c738026..a15f0ea8db 100644 --- a/examples/sandbox/frontend/page-data.ts +++ b/examples/sandbox/frontend/page-data.ts @@ -65,7 +65,7 @@ console.log(metadata.tags, metadata.region);`, });`, kv: `await actor.putText("greeting", "hello"); const value = await actor.getText("greeting");`, - queue: `await actor.queue.work.send({ id: "task-1" }); + queue: `await actor.send("work", { id: "task-1" }); const message = await actor.receiveOne("work");`, workflow: `const workflow = client.order.getOrCreate([orderId]); await workflow.getOrder();`, @@ -247,15 +247,7 @@ export const ACTION_TEMPLATES: Record = { { label: "Ping", action: "ping", args: [] }, { label: "Actor Counts", action: "getActorCounts", args: [] }, ], - queueActor: [ - { label: "Receive One", action: "receiveOne", args: ["work"] }, - ], - sender: [{ label: "Get Messages", action: "getMessages", args: [] }], - multiQueue: [{ label: "Get Messages", action: "getMessages", args: [] }], - timeout: [{ label: "Wait", action: "waitForMessage", args: [2000] }], worker: [{ label: "Get State", action: "getState", args: [] }], - selfSender: [{ label: "Get State", action: "getState", args: [] }], - keepAwake: [{ label: "Get State", action: "getState", args: [] }], order: [{ label: "Get Order", action: "getOrder", args: [] }], timer: [{ label: "Get Timer", action: "getTimer", args: [] }], batch: [{ label: "Get Job", action: "getJob", args: [] }], @@ -808,40 +800,33 @@ export const PAGE_GROUPS: PageGroup[] = [ title: "Queues", icon: "list", pages: [ - { - id: "queue-basics", - title: "Queue Basics", - description: "Send and receive queue messages from actors.", + { + id: "queue-basics", + title: "Queue Basics", + description: "Send and receive queue messages from actors.", docs: [ { label: "Queue", href: "https://rivet.dev/docs/actors/queue", }, ], - actors: ["queueActor", "queueLimitedActor"], - snippet: SNIPPETS.queue, - }, - { - id: "queue-patterns", - title: "Queue Patterns", - description: - "Explore sending, timeouts, workers, and keep-awake patterns.", + actors: ["worker"], + snippet: SNIPPETS.queue, + }, + { + id: "queue-patterns", + title: "Queue Patterns", + description: + "Run a worker loop that consumes queued jobs.", docs: [ { label: "Queue", href: "https://rivet.dev/docs/actors/queue", }, ], - actors: [ - "sender", - "multiQueue", - "timeout", - "worker", - "selfSender", - "keepAwake", - ], - snippet: SNIPPETS.queue, - }, + actors: ["worker"], + snippet: SNIPPETS.queue, + }, { id: "queue-run-loop", title: "Queue in Run Loop", diff --git a/examples/sandbox/src/actors.ts b/examples/sandbox/src/actors.ts index 47bcfe5d6c..efc9aff6c9 100644 --- a/examples/sandbox/src/actors.ts +++ b/examples/sandbox/src/actors.ts @@ -78,13 +78,7 @@ import { } from "./actors/lifecycle/destroy.ts"; import { hibernationActor } from "./actors/lifecycle/hibernation.ts"; // Queues -import { queueActor, queueLimitedActor } from "./actors/queue/queue.ts"; -import { sender } from "./actors/queue/sender.ts"; -import { multiQueue } from "./actors/queue/multi-queue.ts"; -import { timeout } from "./actors/queue/timeout.ts"; import { worker } from "./actors/queue/worker.ts"; -import { selfSender } from "./actors/queue/self-sender.ts"; -import { keepAwake } from "./actors/queue/keep-awake.ts"; // Workflows import { workflowCounterActor, @@ -179,14 +173,7 @@ export const registry = setup({ destroyObserver, hibernationActor, // Queues - queueActor, - queueLimitedActor, - sender, - multiQueue, - timeout, worker, - selfSender, - keepAwake, // Workflows timer, order, diff --git a/examples/sandbox/src/actors/ai/ai-agent.ts b/examples/sandbox/src/actors/ai/ai-agent.ts index 94c957119b..1615f2ec7d 100644 --- a/examples/sandbox/src/actors/ai/ai-agent.ts +++ b/examples/sandbox/src/actors/ai/ai-agent.ts @@ -1,6 +1,6 @@ import { openai } from "@ai-sdk/openai"; import { generateText, tool } from "ai"; -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { z } from "zod"; import { getWeather } from "./my-tools.ts"; import type { Message } from "./types.ts"; @@ -10,6 +10,9 @@ export const aiAgent = actor({ state: { messages: [] as Message[], }, + events: { + messageReceived: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/sandbox/src/actors/connections/conn-state.ts b/examples/sandbox/src/actors/connections/conn-state.ts index 8312f5aa3e..c18244cdb4 100644 --- a/examples/sandbox/src/actors/connections/conn-state.ts +++ b/examples/sandbox/src/actors/connections/conn-state.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export type ConnState = { username: string; @@ -13,6 +13,11 @@ export const connStateActor = actor({ sharedCounter: 0, disconnectionCount: 0, }, + events: { + userConnected: event<{ id: string; username: string; role: string }>(), + userDisconnected: event<{ id: string }>(), + directMessage: event<{ from: string; message: string }>(), + }, // Define connection state createConnState: ( c, diff --git a/examples/sandbox/src/actors/counter/conn-params.ts b/examples/sandbox/src/actors/counter/conn-params.ts index 4116a4432c..f0e05b614f 100644 --- a/examples/sandbox/src/actors/counter/conn-params.ts +++ b/examples/sandbox/src/actors/counter/conn-params.ts @@ -1,7 +1,10 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const counterWithParams = actor({ state: { count: 0, initializers: [] as string[] }, + events: { + newCount: event<{ count: number; by: string }>(), + }, createConnState: (c, params: { name?: string }) => { return { name: params.name || "anonymous", diff --git a/examples/sandbox/src/actors/counter/counter-conn.ts b/examples/sandbox/src/actors/counter/counter-conn.ts index 93d4a6c0af..f06b1b119a 100644 --- a/examples/sandbox/src/actors/counter/counter-conn.ts +++ b/examples/sandbox/src/actors/counter/counter-conn.ts @@ -1,9 +1,12 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const counterConn = actor({ state: { connectionCount: 0, }, + events: { + newCount: event(), + }, connState: { count: 0 }, onConnect: (c, conn) => { c.state.connectionCount += 1; diff --git a/examples/sandbox/src/actors/counter/counter.ts b/examples/sandbox/src/actors/counter/counter.ts index 8f14bea1c2..9fd224e818 100644 --- a/examples/sandbox/src/actors/counter/counter.ts +++ b/examples/sandbox/src/actors/counter/counter.ts @@ -1,7 +1,10 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const counter = actor({ state: { count: 0 }, + events: { + newCount: event(), + }, actions: { increment: (c, x: number) => { c.state.count += x; diff --git a/examples/sandbox/src/actors/lifecycle/run.ts b/examples/sandbox/src/actors/lifecycle/run.ts index 42dce0de96..a6858e5349 100644 --- a/examples/sandbox/src/actors/lifecycle/run.ts +++ b/examples/sandbox/src/actors/lifecycle/run.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event, queue } from "rivetkit"; import type { registry } from "../../actors.ts"; export const RUN_SLEEP_TIMEOUT = 500; @@ -15,7 +15,7 @@ export const runWithTicks = actor({ c.state.runStarted = true; c.log.info("run handler started"); - while (!c.abortSignal.aborted) { + while (!c.aborted) { c.state.tickCount += 1; c.state.lastTickAt = Date.now(); c.log.info({ msg: "tick", tickCount: c.state.tickCount }); @@ -57,19 +57,19 @@ export const runWithQueueConsumer = actor({ messagesReceived: [] as Array<{ name: string; body: unknown }>, runStarted: false, }, + queues: { + messages: queue(), + }, run: async (c) => { c.state.runStarted = true; c.log.info("run handler started, waiting for messages"); - while (!c.abortSignal.aborted) { - const message = await c.queue.next("messages", { timeout: 100 }); - if (message) { - c.log.info({ msg: "received message", body: message.body }); - c.state.messagesReceived.push({ - name: message.name, - body: message.body, - }); - } + for await (const message of c.queue.iter()) { + c.log.info({ msg: "received message", body: message.body }); + c.state.messagesReceived.push({ + name: message.name, + body: message.body, + }); } c.log.info("run handler exiting gracefully"); @@ -82,7 +82,7 @@ export const runWithQueueConsumer = actor({ sendMessage: async (c, body: unknown) => { const client = c.client(); const handle = client.runWithQueueConsumer.getForId(c.actorId); - await handle.queue.messages.send(body); + await handle.send("messages", body); return true; }, }, diff --git a/examples/sandbox/src/actors/lifecycle/scheduled.ts b/examples/sandbox/src/actors/lifecycle/scheduled.ts index 7bac35bac8..b11b802bd0 100644 --- a/examples/sandbox/src/actors/lifecycle/scheduled.ts +++ b/examples/sandbox/src/actors/lifecycle/scheduled.ts @@ -1,4 +1,4 @@ -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; export const scheduled = actor({ state: { @@ -6,6 +6,10 @@ export const scheduled = actor({ scheduledCount: 0, taskHistory: [] as string[], }, + events: { + scheduled: event<{ time: number; count: number }>(), + scheduledWithId: event<{ taskId: string; time: number; count: number }>(), + }, actions: { // Schedule using 'at' with specific timestamp scheduleTaskAt: (c, timestamp: number) => { diff --git a/examples/sandbox/src/actors/lifecycle/sleep.ts b/examples/sandbox/src/actors/lifecycle/sleep.ts index 5717a83f07..7414678e71 100644 --- a/examples/sandbox/src/actors/lifecycle/sleep.ts +++ b/examples/sandbox/src/actors/lifecycle/sleep.ts @@ -1,4 +1,4 @@ -import { actor, type UniversalWebSocket } from "rivetkit"; +import { actor, event, type UniversalWebSocket } from "rivetkit"; import { promiseWithResolvers } from "rivetkit/utils"; export const SLEEP_TIMEOUT = 1000; @@ -37,6 +37,9 @@ export const sleepWithLongRpc = actor({ state: { startCount: 0, sleepCount: 0 }, createVars: () => ({}) as { longRunningResolve: PromiseWithResolvers }, + events: { + waiting: event<[]>(), + }, onWake: (c) => { c.state.startCount += 1; }, diff --git a/examples/sandbox/src/actors/queue/keep-awake.ts b/examples/sandbox/src/actors/queue/keep-awake.ts deleted file mode 100644 index dcf9a0a28e..0000000000 --- a/examples/sandbox/src/actors/queue/keep-awake.ts +++ /dev/null @@ -1,70 +0,0 @@ -import { actor } from "rivetkit"; - -export interface CurrentTask { - id: string; - startedAt: number; - durationMs: number; -} - -export interface CompletedTask { - id: string; - completedAt: number; -} - -export interface KeepAwakeState { - currentTask: CurrentTask | null; - completedTasks: CompletedTask[]; -} - -export const keepAwake = actor({ - state: { - currentTask: null as CurrentTask | null, - completedTasks: [] as CompletedTask[], - }, - async run(c) { - while (!c.abortSignal.aborted) { - const job = await c.queue.next("tasks", { timeout: 1000 }); - if (job) { - const taskId = crypto.randomUUID(); - const { durationMs } = job.body as { durationMs: number }; - - c.state.currentTask = { - id: taskId, - startedAt: Date.now(), - durationMs, - }; - c.broadcast("taskStarted", c.state.currentTask); - - // Wrap long-running work in keepAwake so actor doesn't sleep - await c.keepAwake( - new Promise((resolve) => setTimeout(resolve, durationMs)), - ); - - c.state.completedTasks.push({ id: taskId, completedAt: Date.now() }); - c.state.currentTask = null; - c.broadcast("taskCompleted", { - taskId, - completedTasks: c.state.completedTasks, - }); - } - } - }, - actions: { - getState(c): KeepAwakeState { - return { - currentTask: c.state.currentTask, - completedTasks: c.state.completedTasks, - }; - }, - clearTasks(c) { - c.state.completedTasks = []; - c.broadcast("taskCompleted", { - taskId: null, - completedTasks: [], - }); - }, - }, - options: { - sleepTimeout: 2000, - }, -}); diff --git a/examples/sandbox/src/actors/queue/multi-queue.ts b/examples/sandbox/src/actors/queue/multi-queue.ts deleted file mode 100644 index 57d57e7ad8..0000000000 --- a/examples/sandbox/src/actors/queue/multi-queue.ts +++ /dev/null @@ -1,31 +0,0 @@ -import { actor } from "rivetkit"; - -export interface QueueMessage { - name: string; - body: unknown; -} - -export const multiQueue = actor({ - state: { - messages: [] as QueueMessage[], - }, - actions: { - async receiveFromQueues(c, names: string[], count: number) { - const msgs = await c.queue.next(names, { count, timeout: 100 }); - if (msgs && msgs.length > 0) { - for (const msg of msgs) { - c.state.messages.push({ name: msg.name, body: msg.body }); - } - c.broadcast("messagesReceived", c.state.messages); - } - return msgs ?? []; - }, - getMessages(c): QueueMessage[] { - return c.state.messages; - }, - clearMessages(c) { - c.state.messages = []; - c.broadcast("messagesReceived", c.state.messages); - }, - }, -}); diff --git a/examples/sandbox/src/actors/queue/queue.ts b/examples/sandbox/src/actors/queue/queue.ts deleted file mode 100644 index d77e9a77ff..0000000000 --- a/examples/sandbox/src/actors/queue/queue.ts +++ /dev/null @@ -1,123 +0,0 @@ -import { actor } from "rivetkit"; -import type { registry } from "../../actors.ts"; - -export const queueActor = actor({ - state: {}, - actions: { - receiveOne: async ( - c, - name: string, - opts?: { count?: number; timeout?: number }, - ) => { - const message = await c.queue.next(name, opts); - if (!message) { - return null; - } - return { name: message.name, body: message.body }; - }, - receiveMany: async ( - c, - names: string[], - opts?: { count?: number; timeout?: number }, - ) => { - const messages = await c.queue.next(names, opts); - return (messages ?? []).map( - (message: { name: string; body: unknown }) => ({ - name: message.name, - body: message.body, - }), - ); - }, - receiveRequest: async ( - c, - request: { - name: string | string[]; - count?: number; - timeout?: number; - }, - ) => { - const messages = await c.queue.next(request); - return (messages ?? []).map( - (message: { name: string; body: unknown }) => ({ - name: message.name, - body: message.body, - }), - ); - }, - sendToSelf: async (c, name: string, body: unknown) => { - const client = c.client(); - const handle = client.queueActor.getForId(c.actorId); - await handle.queue[name].send(body); - return true; - }, - waitForAbort: async (c) => { - setTimeout(() => { - c.destroy(); - }, 10); - await c.queue.next("abort", { timeout: 10_000 }); - return true; - }, - receiveAndComplete: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; - } - await message.complete({ echo: message.body }); - return { name: message.name, body: message.body }; - }, - receiveAndCompleteTwice: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; - } - await message.complete({ ok: true }); - try { - await message.complete({ ok: false }); - return { ok: false }; - } catch (error) { - const actorError = error as { group?: string; code?: string }; - return { group: actorError.group, code: actorError.code }; - } - }, - receiveWithoutWaitComplete: async (c, name: string) => { - const message = await c.queue.next(name); - if (!message) { - return null; - } - try { - await message.complete(); - return { ok: false }; - } catch (error) { - const actorError = error as { group?: string; code?: string }; - return { group: actorError.group, code: actorError.code }; - } - }, - receiveWhilePending: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; - } - let errorPayload: { group?: string; code?: string } | undefined; - try { - await c.queue.next(name); - } catch (error) { - const actorError = error as { group?: string; code?: string }; - errorPayload = { - group: actorError.group, - code: actorError.code, - }; - } - await message.complete({ ok: true }); - return errorPayload ?? { ok: false }; - }, - }, -}); - -export const queueLimitedActor = actor({ - state: {}, - actions: {}, - options: { - maxQueueSize: 1, - maxQueueMessageSize: 64, - }, -}); diff --git a/examples/sandbox/src/actors/queue/self-sender.ts b/examples/sandbox/src/actors/queue/self-sender.ts deleted file mode 100644 index bfd1b3cd1a..0000000000 --- a/examples/sandbox/src/actors/queue/self-sender.ts +++ /dev/null @@ -1,48 +0,0 @@ -import { actor } from "rivetkit"; - -export interface SelfSenderState { - sentCount: number; - receivedCount: number; - messages: unknown[]; -} - -export const selfSender = actor({ - state: { - sentCount: 0, - receivedCount: 0, - messages: [] as unknown[], - }, - actions: { - async receiveFromSelf(c) { - const msg = await c.queue.next("self", { timeout: 100 }); - if (msg) { - c.state.receivedCount += 1; - c.state.messages.push(msg.body); - c.broadcast("received", { - receivedCount: c.state.receivedCount, - message: msg.body, - }); - return msg.body; - } - return null; - }, - getState(c): SelfSenderState { - return { - sentCount: c.state.sentCount, - receivedCount: c.state.receivedCount, - messages: c.state.messages, - }; - }, - clearMessages(c) { - c.state.sentCount = 0; - c.state.receivedCount = 0; - c.state.messages = []; - c.broadcast("sent", { sentCount: 0 }); - c.broadcast("received", { receivedCount: 0, message: null }); - }, - incrementSentCount(c) { - c.state.sentCount += 1; - c.broadcast("sent", { sentCount: c.state.sentCount }); - }, - }, -}); diff --git a/examples/sandbox/src/actors/queue/sender.ts b/examples/sandbox/src/actors/queue/sender.ts deleted file mode 100644 index 0db2c49a0b..0000000000 --- a/examples/sandbox/src/actors/queue/sender.ts +++ /dev/null @@ -1,36 +0,0 @@ -import { actor } from "rivetkit"; - -export interface ReceivedMessage { - name: string; - body: unknown; - receivedAt: number; -} - -export const sender = actor({ - state: { - messages: [] as ReceivedMessage[], - }, - actions: { - getMessages(c): ReceivedMessage[] { - return c.state.messages; - }, - async receiveOne(c) { - const msg = await c.queue.next("task", { timeout: 100 }); - if (msg) { - const received: ReceivedMessage = { - name: msg.name, - body: msg.body, - receivedAt: Date.now(), - }; - c.state.messages.push(received); - c.broadcast("messageReceived", c.state.messages); - return received; - } - return null; - }, - clearMessages(c) { - c.state.messages = []; - c.broadcast("messageReceived", c.state.messages); - }, - }, -}); diff --git a/examples/sandbox/src/actors/queue/timeout.ts b/examples/sandbox/src/actors/queue/timeout.ts deleted file mode 100644 index be0ff8fea7..0000000000 --- a/examples/sandbox/src/actors/queue/timeout.ts +++ /dev/null @@ -1,44 +0,0 @@ -import { actor } from "rivetkit"; - -export interface TimeoutResult { - timedOut: boolean; - message?: unknown; - waitedMs: number; -} - -export interface TimeoutState { - lastResult: TimeoutResult | null; - waitStartedAt: number | null; -} - -export const timeout = actor({ - state: { - lastResult: null as TimeoutResult | null, - waitStartedAt: null as number | null, - }, - actions: { - async waitForMessage(c, timeoutMs: number): Promise { - const startedAt = Date.now(); - c.state.waitStartedAt = startedAt; - c.broadcast("waitStarted", { startedAt, timeoutMs }); - - const msg = await c.queue.next("work", { timeout: timeoutMs }); - - const waitedMs = Date.now() - startedAt; - const result: TimeoutResult = msg - ? { timedOut: false, message: msg.body, waitedMs } - : { timedOut: true, waitedMs }; - - c.state.lastResult = result; - c.state.waitStartedAt = null; - c.broadcast("waitCompleted", result); - return result; - }, - getState(c): TimeoutState { - return { - lastResult: c.state.lastResult, - waitStartedAt: c.state.waitStartedAt, - }; - }, - }, -}); diff --git a/examples/sandbox/src/actors/queue/worker.ts b/examples/sandbox/src/actors/queue/worker.ts index e7b710f5dc..41ef9e388e 100644 --- a/examples/sandbox/src/actors/queue/worker.ts +++ b/examples/sandbox/src/actors/queue/worker.ts @@ -1,16 +1,28 @@ -import { actor } from "rivetkit"; +import { actor, event, queue } from "rivetkit"; + +export interface WorkerJob { + id: string; + payload: string; +} export interface WorkerState { status: "idle" | "running"; processed: number; - lastJob: unknown; + lastJob: WorkerJob | null; } export const worker = actor({ state: { status: "idle" as "idle" | "running", processed: 0, - lastJob: null as unknown, + lastJob: null as WorkerJob | null, + }, + events: { + statusChanged: event<{ status: "idle" | "running"; processed: number }>(), + jobProcessed: event<{ processed: number; job: WorkerJob }>(), + }, + queues: { + jobs: queue(), }, async run(c) { c.state.status = "running"; @@ -19,16 +31,13 @@ export const worker = actor({ processed: c.state.processed, }); - while (!c.abortSignal.aborted) { - const job = await c.queue.next("jobs", { timeout: 1000 }); - if (job) { - c.state.processed += 1; - c.state.lastJob = job.body; - c.broadcast("jobProcessed", { - processed: c.state.processed, - job: job.body, - }); - } + for await (const job of c.queue.iter()) { + c.state.processed += 1; + c.state.lastJob = job.body; + c.broadcast("jobProcessed", { + processed: c.state.processed, + job: job.body, + }); } c.state.status = "idle"; diff --git a/examples/sandbox/src/actors/state/sqlite-drizzle/drizzle/migrations.d.ts b/examples/sandbox/src/actors/state/sqlite-drizzle/drizzle/migrations.d.ts new file mode 100644 index 0000000000..00714979b2 --- /dev/null +++ b/examples/sandbox/src/actors/state/sqlite-drizzle/drizzle/migrations.d.ts @@ -0,0 +1,6 @@ +declare const migrations: { + journal: unknown; + migrations: Record; +}; + +export default migrations; diff --git a/examples/sandbox/src/actors/state/vars.ts b/examples/sandbox/src/actors/state/vars.ts index 7a62319824..cec9116fff 100644 --- a/examples/sandbox/src/actors/state/vars.ts +++ b/examples/sandbox/src/actors/state/vars.ts @@ -52,7 +52,7 @@ export const dynamicVarActor = actor({ }; }, actions: { - getVars: (c) => { + getVars: (c: any) => { return c.vars; }, }, @@ -68,7 +68,7 @@ export const uniqueVarActor = actor({ }; }, actions: { - getVars: (c) => { + getVars: (c: any) => { return c.vars; }, }, @@ -84,7 +84,7 @@ export const driverCtxActor = actor({ }; }, actions: { - getVars: (c) => { + getVars: (c: any) => { return c.vars; }, }, diff --git a/examples/sandbox/src/actors/workflow/approval.ts b/examples/sandbox/src/actors/workflow/approval.ts index e39af4cf5e..86070ec7eb 100644 --- a/examples/sandbox/src/actors/workflow/approval.ts +++ b/examples/sandbox/src/actors/workflow/approval.ts @@ -1,9 +1,9 @@ -// APPROVAL REQUEST (Listen Demo) -// Demonstrates: Message listening with timeout for approval workflows +// APPROVAL REQUEST (Queue Wait Demo) +// Demonstrates: Queue waits with timeout for approval workflows // One actor per approval request - actor key is the request ID -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, event, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; export type RequestStatus = "pending" | "approved" | "rejected" | "timeout"; @@ -21,7 +21,7 @@ export type ApprovalRequest = { type State = ApprovalRequest; -const QUEUE_DECISION = workflowQueueName("decision"); +const QUEUE_DECISION = "decision" as const; const APPROVAL_TIMEOUT_MS = 30000; @@ -30,6 +30,11 @@ export type ApprovalRequestInput = { description?: string; }; +type ApprovalDecision = { + approved: boolean; + approver: string; +}; + export const approval = actor({ createState: (c, input?: ApprovalRequestInput): ApprovalRequest => ({ id: c.key[0] as string, @@ -38,6 +43,13 @@ export const approval = actor({ status: "pending", createdAt: Date.now(), }), + queues: { + decision: queue(), + }, + events: { + requestUpdated: event(), + requestCreated: event(), + }, actions: { getRequest: (c): ApprovalRequest => c.state, @@ -72,10 +84,14 @@ export const approval = actor({ c.broadcast("requestCreated", c.state); }); - const decision = await loopCtx.listenWithTimeout<{ - approved: boolean; - approver: string; - }>("wait-decision", "decision", APPROVAL_TIMEOUT_MS); + const [decisionMessage] = await loopCtx.queue.next( + "wait-decision", + { + names: [QUEUE_DECISION], + timeout: APPROVAL_TIMEOUT_MS, + }, + ); + const decision = decisionMessage?.body ?? null; await loopCtx.step("update-status", async () => { c.state.deciding = false; diff --git a/examples/sandbox/src/actors/workflow/batch.ts b/examples/sandbox/src/actors/workflow/batch.ts index 006e894e11..567fb23313 100644 --- a/examples/sandbox/src/actors/workflow/batch.ts +++ b/examples/sandbox/src/actors/workflow/batch.ts @@ -2,7 +2,7 @@ // Demonstrates: Loop with persistent state (cursor) for batch processing // One actor per batch job - actor key is the job ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -59,6 +59,11 @@ export const batch = actor({ batches: [], startedAt: Date.now(), }), + events: { + batchProcessed: event(), + stateChanged: event(), + processingComplete: event<{ totalBatches: number; totalItems: number }>(), + }, actions: { getJob: (c): BatchJob => c.state, diff --git a/examples/sandbox/src/actors/workflow/dashboard.ts b/examples/sandbox/src/actors/workflow/dashboard.ts index 86887d466e..449e38244f 100644 --- a/examples/sandbox/src/actors/workflow/dashboard.ts +++ b/examples/sandbox/src/actors/workflow/dashboard.ts @@ -1,8 +1,8 @@ // DASHBOARD (Join Demo) // Demonstrates: Parallel data fetching with join (wait-all) -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, event, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; export type UserStats = { @@ -45,7 +45,8 @@ export type DashboardState = { type State = DashboardState; -const QUEUE_REFRESH = workflowQueueName("refresh"); +const QUEUE_REFRESH = "refresh"; +type RefreshMessage = Record; async function fetchUserStats(): Promise { await new Promise((r) => setTimeout(r, 800 + Math.random() * 1200)); @@ -87,6 +88,13 @@ export const dashboard = actor({ }, lastRefresh: null as number | null, }, + queues: { + [QUEUE_REFRESH]: queue(), + }, + events: { + stateChanged: event(), + refreshComplete: event(), + }, actions: { refresh: async (c) => { @@ -111,7 +119,9 @@ export const dashboard = actor({ run: async (loopCtx) => { const c = actorCtx(loopCtx); - await loopCtx.listen("wait-refresh", "refresh"); + await loopCtx.queue.next("wait-refresh", { + names: [QUEUE_REFRESH], + }); ctx.log.info({ msg: "starting dashboard refresh" }); diff --git a/examples/sandbox/src/actors/workflow/history-examples.ts b/examples/sandbox/src/actors/workflow/history-examples.ts index dd0b0e5fe5..b70a1f1f19 100644 --- a/examples/sandbox/src/actors/workflow/history-examples.ts +++ b/examples/sandbox/src/actors/workflow/history-examples.ts @@ -1,5 +1,5 @@ -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, event, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; function delay(ms: number): Promise { @@ -250,19 +250,38 @@ export type WorkflowHistoryFullState = { completedAt?: number; }; -type MessageSeed = { name: string; payload: unknown }; +const QUEUE_ORDER_CREATED = "order:created"; +const QUEUE_ORDER_UPDATED = "order:updated"; +const QUEUE_ORDER_ITEM = "order:item"; +const QUEUE_ORDER_ARTIFACT = "order:artifact"; +const QUEUE_ORDER_READY = "order:ready"; +const QUEUE_ORDER_OPTIONAL = "order:optional"; + +type OrderCreatedMessage = { id: string }; +type OrderUpdatedMessage = { id: string; status: string }; +type OrderItemMessage = { sku: string; qty: number }; +type OrderArtifactMessage = { artifactId: string }; +type OrderReadyMessage = { batch: number }; +type OrderOptionalMessage = { note?: string }; + +type MessageSeed = + | { name: typeof QUEUE_ORDER_CREATED; payload: OrderCreatedMessage } + | { name: typeof QUEUE_ORDER_UPDATED; payload: OrderUpdatedMessage } + | { name: typeof QUEUE_ORDER_ITEM; payload: OrderItemMessage } + | { name: typeof QUEUE_ORDER_ARTIFACT; payload: OrderArtifactMessage } + | { name: typeof QUEUE_ORDER_READY; payload: OrderReadyMessage }; const FULL_WORKFLOW_MESSAGE_SEEDS: MessageSeed[] = [ - { name: "order:created", payload: { id: "order-1" } }, - { name: "order:updated", payload: { id: "order-1", status: "paid" } }, - { name: "order:item", payload: { sku: "sku-0", qty: 1 } }, - { name: "order:item", payload: { sku: "sku-4", qty: 1 } }, - { name: "order:artifact", payload: { artifactId: "artifact-0" } }, - { name: "order:artifact", payload: { artifactId: "artifact-1" } }, - { name: "order:artifact", payload: { artifactId: "artifact-2" } }, - { name: "order:ready", payload: { batch: 3 } }, - { name: "order:ready", payload: { batch: 0 } }, - { name: "order:ready", payload: { batch: 2 } }, + { name: QUEUE_ORDER_CREATED, payload: { id: "order-1" } }, + { name: QUEUE_ORDER_UPDATED, payload: { id: "order-1", status: "paid" } }, + { name: QUEUE_ORDER_ITEM, payload: { sku: "sku-0", qty: 1 } }, + { name: QUEUE_ORDER_ITEM, payload: { sku: "sku-4", qty: 1 } }, + { name: QUEUE_ORDER_ARTIFACT, payload: { artifactId: "artifact-0" } }, + { name: QUEUE_ORDER_ARTIFACT, payload: { artifactId: "artifact-1" } }, + { name: QUEUE_ORDER_ARTIFACT, payload: { artifactId: "artifact-2" } }, + { name: QUEUE_ORDER_READY, payload: { batch: 3 } }, + { name: QUEUE_ORDER_READY, payload: { batch: 0 } }, + { name: QUEUE_ORDER_READY, payload: { batch: 2 } }, ]; const FULL_WORKFLOW_ITEMS = [ @@ -278,12 +297,20 @@ export const workflowHistoryFull = actor({ status: "pending", seededMessages: false, }), + queues: { + [QUEUE_ORDER_CREATED]: queue(), + [QUEUE_ORDER_UPDATED]: queue(), + [QUEUE_ORDER_ITEM]: queue(), + [QUEUE_ORDER_ARTIFACT]: queue(), + [QUEUE_ORDER_READY]: queue(), + [QUEUE_ORDER_OPTIONAL]: queue(), + }, actions: { getState: (c): WorkflowHistoryFullState => c.state, seedMessages: async (c) => { if (c.state.seededMessages) return; for (const seed of FULL_WORKFLOW_MESSAGE_SEEDS) { - await c.queue.send(workflowQueueName(seed.name), seed.payload); + await c.queue.send(seed.name, seed.payload); } c.state.seededMessages = true; }, @@ -381,31 +408,36 @@ export const workflowHistoryFull = actor({ return { readyBy, readyBatchBy }; }); - await ctx.listenN("listen-order-created", "order:created", 1); - await ctx.listenWithTimeout( - "listen-order-updated-timeout", - "order:updated", - 250, - ); - await ctx.listenN("listen-batch-two", "order:item", 2); - await ctx.listenNWithTimeout( - "listen-artifacts-timeout", - "order:artifact", - 3, - 300, - ); - await ctx.listenWithTimeout("listen-optional", "order:optional", 200); - await ctx.listenUntil( - "listen-until", - "order:ready", - Date.now() + 300, - ); - await ctx.listenNUntil( - "listen-batch-until", - "order:ready", - 2, - Date.now() + 400, - ); + await ctx.queue.next("listen-order-created", { + names: [QUEUE_ORDER_CREATED], + count: 1, + }); + await ctx.queue.next("listen-order-updated-timeout", { + names: [QUEUE_ORDER_UPDATED], + timeout: 250, + }); + await ctx.queue.next("listen-batch-two", { + names: [QUEUE_ORDER_ITEM], + count: 2, + }); + await ctx.queue.next("listen-artifacts-timeout", { + names: [QUEUE_ORDER_ARTIFACT], + count: 3, + timeout: 300, + }); + await ctx.queue.next("listen-optional", { + names: [QUEUE_ORDER_OPTIONAL], + timeout: 200, + }); + await ctx.queue.next("listen-until", { + names: [QUEUE_ORDER_READY], + timeout: 300, + }); + await ctx.queue.next("listen-batch-until", { + names: [QUEUE_ORDER_READY], + count: 2, + timeout: 400, + }); await ctx.join("join-dependencies", { inventory: { diff --git a/examples/sandbox/src/actors/workflow/order.ts b/examples/sandbox/src/actors/workflow/order.ts index 12b5344744..ad0c0db237 100644 --- a/examples/sandbox/src/actors/workflow/order.ts +++ b/examples/sandbox/src/actors/workflow/order.ts @@ -2,7 +2,7 @@ // Demonstrates: Sequential workflow steps with automatic retries // One actor per order - actor key is the order ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -41,6 +41,9 @@ export const order = actor({ step: 0, createdAt: Date.now(), }), + events: { + orderUpdated: event(), + }, actions: { getOrder: (c): Order => c.state, diff --git a/examples/sandbox/src/actors/workflow/payment.ts b/examples/sandbox/src/actors/workflow/payment.ts index f227b6945d..13997b5e7b 100644 --- a/examples/sandbox/src/actors/workflow/payment.ts +++ b/examples/sandbox/src/actors/workflow/payment.ts @@ -2,7 +2,7 @@ // Demonstrates: Rollback checkpoints with compensating actions // One actor per transaction - actor key is the transaction ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -51,6 +51,11 @@ export const payment = actor({ ], startedAt: Date.now(), }), + events: { + transactionStarted: event(), + transactionUpdated: event(), + transactionCompleted: event(), + }, actions: { getTransaction: (c): Transaction => c.state, diff --git a/examples/sandbox/src/actors/workflow/race.ts b/examples/sandbox/src/actors/workflow/race.ts index 0e7ad1ed3c..922c526e1e 100644 --- a/examples/sandbox/src/actors/workflow/race.ts +++ b/examples/sandbox/src/actors/workflow/race.ts @@ -2,7 +2,7 @@ // Demonstrates: Race (parallel first-wins) for timeout patterns // One actor per race task - actor key is the task ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -32,6 +32,10 @@ export const race = actor({ status: "running", startedAt: Date.now(), }), + events: { + raceStarted: event(), + raceCompleted: event(), + }, actions: { getTask: (c): RaceTask => c.state, diff --git a/examples/sandbox/src/actors/workflow/timer.ts b/examples/sandbox/src/actors/workflow/timer.ts index dec43ffbec..22c0c80636 100644 --- a/examples/sandbox/src/actors/workflow/timer.ts +++ b/examples/sandbox/src/actors/workflow/timer.ts @@ -2,7 +2,7 @@ // Demonstrates: Durable sleep that survives restarts // One actor per timer - actor key is the timer ID -import { actor } from "rivetkit"; +import { actor, event } from "rivetkit"; import { Loop, workflow } from "rivetkit/workflow"; import { actorCtx } from "./_helpers.ts"; @@ -28,6 +28,10 @@ export const timer = actor({ durationMs: input?.durationMs ?? 10000, startedAt: Date.now(), }), + events: { + timerStarted: event(), + timerCompleted: event(), + }, actions: { getTimer: (c): Timer => c.state, diff --git a/examples/sandbox/src/actors/workflow/workflow-fixtures.ts b/examples/sandbox/src/actors/workflow/workflow-fixtures.ts index 8d7e06e951..3deb92001f 100644 --- a/examples/sandbox/src/actors/workflow/workflow-fixtures.ts +++ b/examples/sandbox/src/actors/workflow/workflow-fixtures.ts @@ -1,5 +1,5 @@ -import { actor } from "rivetkit"; -import { Loop, workflow, workflowQueueName } from "rivetkit/workflow"; +import { actor, queue } from "rivetkit"; +import { Loop, workflow } from "rivetkit/workflow"; const WORKFLOW_GUARD_KV_KEY = "__rivet_actor_workflow_guard_triggered"; @@ -50,18 +50,25 @@ export const workflowQueueActor = actor({ state: { received: [] as unknown[], }, + queues: { + [WORKFLOW_QUEUE_NAME]: queue(), + }, run: workflow(async (ctx) => { await ctx.loop({ name: "queue", run: async (loopCtx) => { const actorLoopCtx = loopCtx as any; - const message = await loopCtx.listen( - "queue-wait", - WORKFLOW_QUEUE_NAME, - ); + const [message] = await loopCtx.queue.next("queue-wait", { + names: [WORKFLOW_QUEUE_NAME], + completable: true, + }); + if (!message || !message.complete) { + return Loop.continue(undefined); + } + const complete = message.complete; await loopCtx.step("store-message", async () => { actorLoopCtx.state.received.push(message.body); - await message.complete({ echo: message.body }); + await complete({ echo: message.body }); }); return Loop.continue(undefined); }, @@ -97,4 +104,4 @@ export const workflowSleepActor = actor({ }, }); -export { WORKFLOW_QUEUE_NAME, workflowQueueName }; +export { WORKFLOW_QUEUE_NAME }; diff --git a/examples/scheduling-vercel/src/actors.ts b/examples/scheduling-vercel/src/actors.ts index 7a1784dc2b..72d3f86c96 100644 --- a/examples/scheduling-vercel/src/actors.ts +++ b/examples/scheduling-vercel/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; interface Reminder { id: string; @@ -17,6 +17,9 @@ const reminderActor = actor({ reminders: [] as Reminder[], completedCount: 0, } satisfies ReminderActorState as ReminderActorState, + events: { + reminderTriggered: event(), + }, actions: { // Schedule a reminder with a delay in milliseconds diff --git a/examples/scheduling/src/actors.ts b/examples/scheduling/src/actors.ts index 7a1784dc2b..72d3f86c96 100644 --- a/examples/scheduling/src/actors.ts +++ b/examples/scheduling/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; interface Reminder { id: string; @@ -17,6 +17,9 @@ const reminderActor = actor({ reminders: [] as Reminder[], completedCount: 0, } satisfies ReminderActorState as ReminderActorState, + events: { + reminderTriggered: event(), + }, actions: { // Schedule a reminder with a delay in milliseconds diff --git a/examples/state-vercel/src/actors.ts b/examples/state-vercel/src/actors.ts index 94f251584e..c0a8c9329e 100644 --- a/examples/state-vercel/src/actors.ts +++ b/examples/state-vercel/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type Message = { id: string; @@ -12,6 +12,10 @@ export const chatRoom = actor({ state: { messages: [] as Message[], }, + events: { + newMessage: event(), + messagesCleared: event<[]>(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/state/src/actors.ts b/examples/state/src/actors.ts index 94f251584e..c0a8c9329e 100644 --- a/examples/state/src/actors.ts +++ b/examples/state/src/actors.ts @@ -1,4 +1,4 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type Message = { id: string; @@ -12,6 +12,10 @@ export const chatRoom = actor({ state: { messages: [] as Message[], }, + events: { + newMessage: event(), + messagesCleared: event<[]>(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions diff --git a/examples/stream-vercel/src/actors.ts b/examples/stream-vercel/src/actors.ts index 6f0d160b8c..9d82a0124a 100644 --- a/examples/stream-vercel/src/actors.ts +++ b/examples/stream-vercel/src/actors.ts @@ -1,15 +1,24 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type StreamState = { topValues: number[]; }; +export type StreamUpdate = { + topValues: number[]; + totalCount: number; + highestValue: number | null; +}; + const streamProcessor = actor({ // Persistent state that survives restarts: https://rivet.dev/docs/actors/state state: { topValues: [] as number[], totalValues: 0, }, + events: { + updated: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions @@ -44,7 +53,7 @@ const streamProcessor = actor({ // Sort descending to ensure correct order c.state.topValues.sort((a, b) => b - a); - const result = { + const result: StreamUpdate = { topValues: c.state.topValues, totalCount: c.state.totalValues, highestValue: @@ -61,7 +70,7 @@ const streamProcessor = actor({ c.state.topValues = []; c.state.totalValues = 0; - const result = { + const result: StreamUpdate = { topValues: c.state.topValues, totalCount: c.state.totalValues, highestValue: null, diff --git a/examples/stream/src/actors.ts b/examples/stream/src/actors.ts index 6f0d160b8c..9d82a0124a 100644 --- a/examples/stream/src/actors.ts +++ b/examples/stream/src/actors.ts @@ -1,15 +1,24 @@ -import { actor, setup } from "rivetkit"; +import { actor, setup, event } from "rivetkit"; export type StreamState = { topValues: number[]; }; +export type StreamUpdate = { + topValues: number[]; + totalCount: number; + highestValue: number | null; +}; + const streamProcessor = actor({ // Persistent state that survives restarts: https://rivet.dev/docs/actors/state state: { topValues: [] as number[], totalValues: 0, }, + events: { + updated: event(), + }, actions: { // Callable functions from clients: https://rivet.dev/docs/actors/actions @@ -44,7 +53,7 @@ const streamProcessor = actor({ // Sort descending to ensure correct order c.state.topValues.sort((a, b) => b - a); - const result = { + const result: StreamUpdate = { topValues: c.state.topValues, totalCount: c.state.totalValues, highestValue: @@ -61,7 +70,7 @@ const streamProcessor = actor({ c.state.topValues = []; c.state.totalValues = 0; - const result = { + const result: StreamUpdate = { topValues: c.state.topValues, totalCount: c.state.totalValues, highestValue: null, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 3192d6c6ff..19d086d746 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -4281,6 +4281,9 @@ importers: '@rivetkit/workflow-engine': specifier: workspace:* version: link:../workflow-engine + '@standard-schema/spec': + specifier: ^1.0.0 + version: 1.0.0 cbor-x: specifier: ^1.6.0 version: 1.6.0 diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-drizzle.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-drizzle.ts index 1cd7f297e8..bf74dffeef 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-drizzle.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actor-db-drizzle.ts @@ -16,38 +16,38 @@ export const dbActorDrizzle = actor({ await c.db.execute( `INSERT INTO test_data (value, payload, created_at) VALUES ('${value}', '', ${Date.now()})`, ); - const results = (await c.db.execute( + const results = await c.db.execute<{ id: number }>( `SELECT last_insert_rowid() as id`, - )) as Array<{ id: number }>; + ); return { id: results[0].id }; }, getValues: async (c) => { - const results = (await c.db.execute( - `SELECT * FROM test_data ORDER BY id`, - )) as Array<{ + const results = await c.db.execute<{ id: number; value: string; payload: string; created_at: number; - }>; + }>( + `SELECT * FROM test_data ORDER BY id`, + ); return results; }, getValue: async (c, id: number) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ value: string }>( `SELECT value FROM test_data WHERE id = ${id}`, - )) as Array<{ value: string }>; + ); return results[0]?.value ?? null; }, getCount: async (c) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ count: number }>( `SELECT COUNT(*) as count FROM test_data`, - )) as Array<{ count: number }>; + ); return results[0].count; }, rawSelectCount: async (c) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ count: number }>( `SELECT COUNT(*) as count FROM test_data`, - )) as Array<{ count: number }>; + ); return results[0]?.count ?? 0; }, insertMany: async (c, count: number) => { @@ -88,15 +88,15 @@ export const dbActorDrizzle = actor({ await c.db.execute( `INSERT INTO test_data (value, payload, created_at) VALUES ('payload', '${payload}', ${Date.now()})`, ); - const results = (await c.db.execute( + const results = await c.db.execute<{ id: number }>( `SELECT last_insert_rowid() as id`, - )) as Array<{ id: number }>; + ); return { id: results[0].id, size }; }, getPayloadSize: async (c, id: number) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ size: number }>( `SELECT length(payload) as size FROM test_data WHERE id = ${id}`, - )) as Array<{ size: number }>; + ); return results[0]?.size ?? 0; }, repeatUpdate: async (c, id: number, count: number) => { @@ -119,9 +119,9 @@ export const dbActorDrizzle = actor({ await c.db.execute( `BEGIN; INSERT INTO test_data (value, payload, created_at) VALUES ('${value}', '', ${Date.now()}); UPDATE test_data SET value = '${value}-updated' WHERE id = last_insert_rowid(); COMMIT;`, ); - const results = (await c.db.execute( + const results = await c.db.execute<{ value: string }>( `SELECT value FROM test_data ORDER BY id DESC LIMIT 1`, - )) as Array<{ value: string }>; + ); return results[0]?.value ?? null; }, triggerSleep: (c) => { 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 89ed0d02c9..1eac26184e 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 @@ -22,38 +22,38 @@ export const dbActorRaw = actor({ await c.db.execute( `INSERT INTO test_data (value, payload, created_at) VALUES ('${value}', '', ${Date.now()})`, ); - const results = (await c.db.execute( + const results = await c.db.execute<{ id: number }>( `SELECT last_insert_rowid() as id`, - )) as Array<{ id: number }>; + ); return { id: results[0].id }; }, getValues: async (c) => { - const results = (await c.db.execute( - `SELECT * FROM test_data ORDER BY id`, - )) as Array<{ + const results = await c.db.execute<{ id: number; value: string; payload: string; created_at: number; - }>; + }>( + `SELECT * FROM test_data ORDER BY id`, + ); return results; }, getValue: async (c, id: number) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ value: string }>( `SELECT value FROM test_data WHERE id = ${id}`, - )) as Array<{ value: string }>; + ); return results[0]?.value ?? null; }, getCount: async (c) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ count: number }>( `SELECT COUNT(*) as count FROM test_data`, - )) as Array<{ count: number }>; + ); return results[0].count; }, rawSelectCount: async (c) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ count: number }>( `SELECT COUNT(*) as count FROM test_data`, - )) as Array<{ count: number }>; + ); return results[0].count; }, insertMany: async (c, count: number) => { @@ -94,15 +94,15 @@ export const dbActorRaw = actor({ await c.db.execute( `INSERT INTO test_data (value, payload, created_at) VALUES ('payload', '${payload}', ${Date.now()})`, ); - const results = (await c.db.execute( + const results = await c.db.execute<{ id: number }>( `SELECT last_insert_rowid() as id`, - )) as Array<{ id: number }>; + ); return { id: results[0].id, size }; }, getPayloadSize: async (c, id: number) => { - const results = (await c.db.execute( + const results = await c.db.execute<{ size: number }>( `SELECT length(payload) as size FROM test_data WHERE id = ${id}`, - )) as Array<{ size: number }>; + ); return results[0]?.size ?? 0; }, repeatUpdate: async (c, id: number, count: number) => { @@ -125,9 +125,9 @@ export const dbActorRaw = actor({ await c.db.execute( `BEGIN; INSERT INTO test_data (value, payload, created_at) VALUES ('${value}', '', ${Date.now()}); UPDATE test_data SET value = '${value}-updated' WHERE id = last_insert_rowid(); COMMIT;`, ); - const results = (await c.db.execute( + const results = await c.db.execute<{ value: string }>( `SELECT value FROM test_data ORDER BY id DESC LIMIT 1`, - )) as Array<{ value: string }>; + ); return results[0]?.value ?? null; }, triggerSleep: (c) => { diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/queue.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/queue.ts index cb5eb893b8..c648bf8308 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/queue.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/queue.ts @@ -1,15 +1,39 @@ -import { actor } from "rivetkit"; +import { actor, queue } from "rivetkit"; import type { registry } from "./registry"; +const queueSchemas = { + greeting: queue<{ hello: string }>(), + self: queue<{ value: number }>(), + a: queue(), + b: queue(), + c: queue(), + one: queue(), + two: queue(), + missing: queue(), + abort: queue(), + tasks: queue<{ value: number }, { echo: { value: number } }>(), + timeout: queue<{ value: number }, { ok: true }>(), + nowait: queue<{ value: string }>(), + twice: queue<{ value: string }, { ok: true }>(), +} as const; + +type QueueName = keyof typeof queueSchemas; + export const queueActor = actor({ state: {}, + queues: queueSchemas, actions: { receiveOne: async ( c, - name: string, + name: QueueName, opts?: { count?: number; timeout?: number }, ) => { - const message = await c.queue.next(name, opts); + const messages = await c.queue.next({ + names: [name], + count: opts?.count, + timeout: opts?.timeout, + }); + const message = messages[0]; if (!message) { return null; } @@ -17,98 +41,171 @@ export const queueActor = actor({ }, receiveMany: async ( c, - names: string[], + names: QueueName[], opts?: { count?: number; timeout?: number }, ) => { - const messages = await c.queue.next(names, opts); - return (messages ?? []).map( - (message: { name: string; body: unknown }) => ({ - name: message.name, - body: message.body, - }), - ); + const messages = await c.queue.next({ + names, + count: opts?.count, + timeout: opts?.timeout, + }); + return messages.map((message) => ({ + name: message.name, + body: message.body, + })); }, receiveRequest: async ( c, request: { - name: string | string[]; + names?: QueueName[]; count?: number; timeout?: number; }, ) => { const messages = await c.queue.next(request); - return (messages ?? []).map( - (message: { name: string; body: unknown }) => ({ - name: message.name, - body: message.body, - }), - ); - }, - sendToSelf: async (c, name: string, body: unknown) => { + return messages.map((message) => ({ + name: message.name, + body: message.body, + })); + }, + tryReceiveMany: async ( + c, + request: { + names?: QueueName[]; + count?: number; + }, + ) => { + const messages = await c.queue.tryNext(request); + return messages.map((message) => ({ + name: message.name, + body: message.body, + })); + }, + receiveWithIterator: async (c, name: QueueName) => { + for await (const message of c.queue.iter({ names: [name] })) { + return { name: message.name, body: message.body }; + } + return null; + }, + receiveWithAsyncIterator: async (c) => { + for await (const message of c.queue.iter()) { + return { name: message.name, body: message.body }; + } + return null; + }, + sendToSelf: async (c, name: QueueName, body: unknown) => { const client = c.client(); const handle = client.queueActor.getForId(c.actorId); - await handle.queue[name].send(body); + await handle.send(name, body); return true; }, waitForAbort: async (c) => { setTimeout(() => { c.destroy(); }, 10); - await c.queue.next("abort", { timeout: 10_000 }); + await c.queue.next({ names: ["abort"], timeout: 10_000 }); return true; }, - receiveAndComplete: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; + waitForSignalAbort: async (c) => { + const controller = new AbortController(); + controller.abort(); + try { + await c.queue.next({ + names: ["abort"], + timeout: 10_000, + signal: controller.signal, + }); + return { ok: false }; + } catch (error) { + const actorError = error as { group?: string; code?: string }; + return { group: actorError.group, code: actorError.code }; } - await message.complete({ echo: message.body }); - return { name: message.name, body: message.body }; }, - receiveAndCompleteTwice: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); - if (!message) { - return null; - } - await message.complete({ ok: true }); + waitForActorAbortWithSignal: async (c) => { + const controller = new AbortController(); + setTimeout(() => { + c.destroy(); + }, 10); try { - await message.complete({ ok: false }); + await c.queue.next({ + names: ["abort"], + timeout: 10_000, + signal: controller.signal, + }); return { ok: false }; } catch (error) { const actorError = error as { group?: string; code?: string }; return { group: actorError.group, code: actorError.code }; } }, - receiveWithoutWaitComplete: async (c, name: string) => { - const message = await c.queue.next(name); + iterWithSignalAbort: async (c) => { + const controller = new AbortController(); + controller.abort(); + for await (const _message of c.queue.iter({ + names: ["abort"], + signal: controller.signal, + })) { + return { ok: false }; + } + return { ok: true }; + }, + receiveAndComplete: async (c, name: "tasks") => { + const messages = await c.queue.next({ names: [name], completable: true }); + const message = messages[0]; if (!message) { return null; } + await message.complete({ echo: message.body }); + return { name: message.name, body: message.body }; + }, + receiveWithoutComplete: async (c, name: "tasks") => { + const messages = await c.queue.next({ names: [name], completable: true }); + const message = messages[0]; + if (!message) { + return null; + } + return { name: message.name, body: message.body }; + }, + receiveManualThenNextWithoutComplete: async (c, name: "tasks") => { + const messages = await c.queue.next({ names: [name], completable: true }); + const message = messages[0]; + if (!message) { + return { ok: false, reason: "no_message" }; + } + try { - await message.complete(); - return { ok: false }; + await c.queue.next({ names: [name], timeout: 0 }); + c.destroy(); + return { ok: false, reason: "next_succeeded" }; } catch (error) { + c.destroy(); const actorError = error as { group?: string; code?: string }; return { group: actorError.group, code: actorError.code }; } }, - receiveWhilePending: async (c, name: string) => { - const message = await c.queue.next(name, { wait: true }); + receiveAndCompleteTwice: async (c, name: "twice") => { + const messages = await c.queue.next({ names: [name], completable: true }); + const message = messages[0]; if (!message) { return null; } - let errorPayload: { group?: string; code?: string } | undefined; + await message.complete({ ok: true }); try { - await c.queue.next(name); + await message.complete({ ok: true }); + return { ok: false }; } catch (error) { const actorError = error as { group?: string; code?: string }; - errorPayload = { - group: actorError.group, - code: actorError.code, - }; + return { group: actorError.group, code: actorError.code }; } - await message.complete({ ok: true }); - return errorPayload ?? { ok: false }; + }, + receiveWithoutCompleteMethod: async (c, name: "nowait") => { + const messages = await c.queue.next({ names: [name], completable: true }); + const message = messages[0]; + return { + hasComplete: + message !== undefined && + typeof message.complete === "function", + }; }, }, }); diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/run.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/run.ts index 7d4b87b454..7bfe655b5e 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/run.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/run.ts @@ -15,7 +15,7 @@ export const runWithTicks = actor({ c.state.runStarted = true; c.log.info("run handler started"); - while (!c.abortSignal.aborted) { + while (!c.aborted) { c.state.tickCount += 1; c.state.lastTickAt = Date.now(); c.log.info({ msg: "tick", tickCount: c.state.tickCount }); @@ -56,13 +56,18 @@ export const runWithQueueConsumer = actor({ state: { messagesReceived: [] as Array<{ name: string; body: unknown }>, runStarted: false, + wakeCount: 0, + }, + onWake: (c) => { + c.state.wakeCount += 1; }, run: async (c) => { c.state.runStarted = true; c.log.info("run handler started, waiting for messages"); - while (!c.abortSignal.aborted) { - const message = await c.queue.next("messages", { timeout: 100 }); + while (!c.aborted) { + const messages = await c.queue.next({ names: ["messages"] }); + const message = messages[0]; if (message) { c.log.info({ msg: "received message", body: message.body }); c.state.messagesReceived.push({ @@ -78,11 +83,12 @@ export const runWithQueueConsumer = actor({ getState: (c) => ({ messagesReceived: c.state.messagesReceived, runStarted: c.state.runStarted, + wakeCount: c.state.wakeCount, }), sendMessage: async (c, body: unknown) => { const client = c.client(); const handle = client.runWithQueueConsumer.getForId(c.actorId); - await handle.queue.messages.send(body); + await handle.send("messages", body); return true; }, }, diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/workflow.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/workflow.ts index 3038124b4b..220652333d 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/workflow.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/workflow.ts @@ -1,8 +1,8 @@ import { Loop } from "@rivetkit/workflow-engine"; -import { actor } from "@/actor/mod"; +import { actor, queue } from "@/actor/mod"; import { db } from "@/db/mod"; import { WORKFLOW_GUARD_KV_KEY } from "@/workflow/constants"; -import { workflow, workflowQueueName } from "@/workflow/mod"; +import { workflow } from "@/workflow/mod"; import type { registry } from "./registry"; const WORKFLOW_QUEUE_NAME = "workflow-default"; @@ -52,18 +52,25 @@ export const workflowQueueActor = actor({ state: { received: [] as unknown[], }, + queues: { + [WORKFLOW_QUEUE_NAME]: queue(), + }, run: workflow(async (ctx) => { await ctx.loop({ name: "queue", run: async (loopCtx) => { const actorLoopCtx = loopCtx as any; - const message = await loopCtx.listen( - "queue-wait", - WORKFLOW_QUEUE_NAME, - ); + const [message] = await loopCtx.queue.next("queue-wait", { + names: [WORKFLOW_QUEUE_NAME], + completable: true, + }); + if (!message || !message.complete) { + return Loop.continue(undefined); + } + const complete = message.complete; await loopCtx.step("store-message", async () => { actorLoopCtx.state.received.push(message.body); - await message.complete({ echo: message.body }); + await complete({ echo: message.body }); }); return Loop.continue(undefined); }, @@ -74,7 +81,8 @@ export const workflowQueueActor = actor({ sendAndWait: async (c, payload: unknown) => { const client = c.client(); const handle = client.workflowQueueActor.getForId(c.actorId); - return await handle.queue[workflowQueueName(WORKFLOW_QUEUE_NAME)].send( + return await handle.send( + WORKFLOW_QUEUE_NAME, payload, { wait: true, timeout: 1_000 }, ); @@ -117,25 +125,25 @@ export const workflowAccessActor = actor({ } try { - actorLoopCtx.client(); + actorLoopCtx.client(); } catch (error) { outsideClientError = error instanceof Error ? error.message : String(error); } await loopCtx.step("access-step", async () => { - await actorLoopCtx.db.execute( + await loopCtx.db.execute( `INSERT INTO workflow_access_log (created_at) VALUES (${Date.now()})`, ); - const counts = (await actorLoopCtx.db.execute( + const counts = await loopCtx.db.execute<{ count: number }>( `SELECT COUNT(*) as count FROM workflow_access_log`, - )) as Array<{ count: number }>; - const client = actorLoopCtx.client(); + ); + const client = loopCtx.client(); - actorLoopCtx.state.outsideDbError = outsideDbError; - actorLoopCtx.state.outsideClientError = outsideClientError; - actorLoopCtx.state.insideDbCount = counts[0]?.count ?? 0; - actorLoopCtx.state.insideClientAvailable = + loopCtx.state.outsideDbError = outsideDbError; + loopCtx.state.outsideClientError = outsideClientError; + loopCtx.state.insideDbCount = counts[0]?.count ?? 0; + loopCtx.state.insideClientAvailable = typeof client.workflowQueueActor.getForId === "function"; }); @@ -174,4 +182,4 @@ export const workflowSleepActor = actor({ }, }); -export { WORKFLOW_QUEUE_NAME, workflowQueueName }; +export { WORKFLOW_QUEUE_NAME }; diff --git a/rivetkit-typescript/packages/rivetkit/package.json b/rivetkit-typescript/packages/rivetkit/package.json index b0a7195044..dccaf71807 100644 --- a/rivetkit-typescript/packages/rivetkit/package.json +++ b/rivetkit-typescript/packages/rivetkit/package.json @@ -217,6 +217,7 @@ "@rivetkit/sqlite-vfs": "workspace:*", "@rivetkit/traces": "workspace:*", "@rivetkit/virtual-websocket": "workspace:*", + "@standard-schema/spec": "^1.0.0", "cbor-x": "^1.6.0", "get-port": "^7.1.0", "hono": "^4.7.0", diff --git a/rivetkit-typescript/packages/rivetkit/scripts/bench-sqlite.ts b/rivetkit-typescript/packages/rivetkit/scripts/bench-sqlite.ts index c70c75c6f9..158df4b290 100755 --- a/rivetkit-typescript/packages/rivetkit/scripts/bench-sqlite.ts +++ b/rivetkit-typescript/packages/rivetkit/scripts/bench-sqlite.ts @@ -3,13 +3,14 @@ /** * SQLite Benchmark Script * - * Compares batch vs non-batch performance: - * 1. Filesystem + Native SQLite (better-sqlite3) - * 2. Filesystem + KV Filesystem (wa-sqlite with file-backed KV) + * Compares batch vs non-batch performance across driver configurations. */ import Table from "cli-table3"; -import { createFileSystemDriver } from "@/drivers/file-system/mod"; +import { + createFileSystemDriver, + createMemoryDriver, +} from "@/drivers/file-system/mod"; import { registry } from "../fixtures/driver-test-suite/registry"; interface TimingResult { @@ -52,13 +53,13 @@ async function runBenchmark(client: Client, name: string): Promise { const results: BenchmarkResult[] = []; - // 1. Native SQLite - console.log("1. Native SQLite..."); + // 1. File System + console.log("1. File System..."); try { const { client } = await registry.start({ - driver: createFileSystemDriver({ useNativeSqlite: true }), + driver: createFileSystemDriver({ + path: `/tmp/rivetkit-bench-${crypto.randomUUID()}`, + }), defaultServerPort: 6430, }); - results.push(await runBenchmark(client, "Native SQLite")); + results.push(await runBenchmark(client, "File System")); console.log(" Done"); } catch (err) { console.log(` Skipped: ${err}`); } - // 2. KV Filesystem - console.log("2. KV Filesystem..."); + // 2. Memory + console.log("2. Memory..."); try { const { client } = await registry.start({ - driver: createFileSystemDriver({ useNativeSqlite: false }), + driver: createMemoryDriver(), defaultServerPort: 6431, }); - results.push(await runBenchmark(client, "KV Filesystem")); + results.push(await runBenchmark(client, "Memory")); console.log(" Done"); } catch (err) { console.log(` Skipped: ${err}`); diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/config.ts b/rivetkit-typescript/packages/rivetkit/src/actor/config.ts index 6096342af5..152d2c563e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/config.ts @@ -20,6 +20,7 @@ import type { WebSocketContext, } from "./contexts"; import type { AnyDatabaseProvider } from "./database"; +import type { EventSchemaConfig, QueueSchemaConfig } from "./schema"; export interface ActorTypes< TState, @@ -57,9 +58,10 @@ export interface RunInspectorConfig { const WorkflowInspectorConfigSchema = z.object({ getHistory: zFunction["getHistory"]>(), - onHistoryUpdated: zFunction< - NonNullable["onHistoryUpdated"]> - >().optional(), + onHistoryUpdated: + zFunction< + NonNullable["onHistoryUpdated"]> + >().optional(), }); const RunInspectorConfigSchema = z @@ -79,14 +81,62 @@ export const RunConfigSchema = z.object({ /** Inspector integration for long-running run handlers. */ inspector: RunInspectorConfigSchema.optional(), }); -export type RunConfig = z.infer; +type RunConfigRuntime = z.infer; +export type RunConfig< + TState = unknown, + TConnParams = unknown, + TConnState = unknown, + TVars = unknown, + TInput = unknown, + TDatabase extends AnyDatabaseProvider = AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, +> = Omit & { + run: ( + c: RunContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, + ) => void | Promise; +}; + +type AnyRunConfig = RunConfig< + any, + any, + any, + any, + any, + AnyDatabaseProvider, + any, + any +>; + +export const RUN_FUNCTION_CONFIG_SYMBOL = Symbol.for( + "rivetkit.run_function_config", +); + +interface RunFunctionConfig { + name?: string; + icon?: string; + inspector?: RunInspectorConfig; +} + +type RunFunctionWithConfig = ((...args: any[]) => any) & { + [RUN_FUNCTION_CONFIG_SYMBOL]?: RunFunctionConfig; +}; // Run can be either a function or an object with name/icon/run const zRunHandler = z.union([zFunction(), RunConfigSchema]).optional(); /** Extract the run function from either a function or RunConfig object. */ export function getRunFunction( - run: ((...args: any[]) => any) | RunConfig | undefined, + run: ((...args: any[]) => any) | AnyRunConfig | undefined, ): ((...args: any[]) => any) | undefined { if (!run) return undefined; if (typeof run === "function") return run; @@ -95,17 +145,28 @@ export function getRunFunction( /** Extract run metadata (name/icon) from RunConfig if provided. */ export function getRunMetadata( - run: ((...args: any[]) => any) | RunConfig | undefined, + run: ((...args: any[]) => any) | AnyRunConfig | undefined, ): { name?: string; icon?: string } { - if (!run || typeof run === "function") return {}; + if (!run) return {}; + if (typeof run === "function") { + const config = (run as RunFunctionWithConfig)[ + RUN_FUNCTION_CONFIG_SYMBOL + ]; + if (!config) return {}; + return { name: config.name, icon: config.icon }; + } return { name: run.name, icon: run.icon }; } /** Extract run inspector configuration if provided. */ export function getRunInspectorConfig( - run: ((...args: any[]) => any) | RunConfig | undefined, + run: ((...args: any[]) => any) | AnyRunConfig | undefined, ): RunInspectorConfig | undefined { - if (!run || typeof run === "function") return undefined; + if (!run) return undefined; + if (typeof run === "function") { + return (run as RunFunctionWithConfig)[RUN_FUNCTION_CONFIG_SYMBOL] + ?.inspector; + } return run.inspector; } @@ -129,6 +190,8 @@ export const ActorConfigSchema = z onRequest: zFunction().optional(), onWebSocket: zFunction().optional(), actions: z.record(z.string(), zFunction()).default(() => ({})), + events: z.record(z.string(), z.any()).optional(), + queues: z.record(z.string(), z.any()).optional(), state: z.any().optional(), createState: zFunction().optional(), connState: z.any().optional(), @@ -219,11 +282,13 @@ type CreateState< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig, + TQueues extends QueueSchemaConfig, > = | { state: TState } | { createState: ( - c: CreateContext, + c: CreateContext, input: TInput, ) => TState | Promise; } @@ -241,11 +306,20 @@ type CreateConnState< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig, + TQueues extends QueueSchemaConfig, > = | { connState: TConnState } | { createConnState: ( - c: CreateConnStateContext, + c: CreateConnStateContext< + TState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, params: TConnParams, ) => TConnState | Promise; } @@ -264,6 +338,8 @@ type CreateVars< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig, + TQueues extends QueueSchemaConfig, > = | { /** @@ -276,7 +352,13 @@ type CreateVars< * @experimental */ createVars: ( - c: CreateVarsContext, + c: CreateVarsContext< + TState, + TInput, + TDatabase, + TEvents, + TQueues + >, driverCtx: any, ) => TVars | Promise; } @@ -289,6 +371,8 @@ export interface Actions< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > { [Action: string]: ( c: ActionContext< @@ -297,7 +381,9 @@ export interface Actions< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ...args: any[] ) => any; @@ -320,13 +406,17 @@ interface BaseActorConfig< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig, + TQueues extends QueueSchemaConfig, TActions extends Actions< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, > { /** @@ -336,7 +426,7 @@ interface BaseActorConfig< * This is called before any other lifecycle hooks. */ onCreate?: ( - c: CreateContext, + c: CreateContext, input: TInput, ) => void | Promise; @@ -350,7 +440,9 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ) => void | Promise; @@ -369,7 +461,9 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ) => void | Promise; @@ -390,7 +484,9 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ) => void | Promise; @@ -425,10 +521,21 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ) => void | Promise) - | RunConfig; + | RunConfig< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >; /** * Called when the actor's state changes. @@ -448,7 +555,9 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, newState: TState, ) => void; @@ -464,7 +573,14 @@ interface BaseActorConfig< * @throws Throw an error to reject the connection */ onBeforeConnect?: ( - c: BeforeConnectContext, + c: BeforeConnectContext< + TState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, params: TConnParams, ) => void | Promise; @@ -484,9 +600,20 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues + >, + conn: Conn< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues >, - conn: Conn, ) => void | Promise; /** @@ -505,9 +632,20 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues + >, + conn: Conn< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues >, - conn: Conn, ) => void | Promise; /** @@ -529,7 +667,9 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, name: string, args: unknown[], @@ -554,7 +694,9 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, request: Request, ) => Response | Promise; @@ -576,12 +718,24 @@ interface BaseActorConfig< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, websocket: UniversalWebSocket, ) => void | Promise; actions?: TActions; + + /** + * Schema map for events broadcasted by this actor. + */ + events?: TEvents; + + /** + * Schema map for queue payloads sent by this actor. + */ + queues?: TQueues; } type ActorDatabaseConfig = @@ -603,9 +757,13 @@ export type ActorConfig< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > = Omit< z.infer, | "actions" + | "events" + | "queues" | "onCreate" | "onDestroy" | "onWake" @@ -633,11 +791,49 @@ export type ActorConfig< TVars, TInput, TDatabase, - Actions + TEvents, + TQueues, + Actions< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > + > & + CreateState< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > & + CreateConnState< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > & + CreateVars< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues > & - CreateState & - CreateConnState & - CreateVars & ActorDatabaseConfig; // See description on `ActorConfig` @@ -648,13 +844,17 @@ export type ActorConfigInput< TVars = undefined, TInput = undefined, TDatabase extends AnyDatabaseProvider = undefined, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, TActions extends Actions< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > = Record, > = { types?: ActorTypes< @@ -668,6 +868,8 @@ export type ActorConfigInput< } & Omit< z.input, | "actions" + | "events" + | "queues" | "onCreate" | "onDestroy" | "onWake" @@ -695,11 +897,40 @@ export type ActorConfigInput< TVars, TInput, TDatabase, + TEvents, + TQueues, TActions > & - CreateState & - CreateConnState & - CreateVars & + CreateState< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > & + CreateConnState< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > & + CreateVars< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > & ActorDatabaseConfig; // For testing type definitions: @@ -710,13 +941,17 @@ export function test< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig, + TQueues extends QueueSchemaConfig, TActions extends Actions< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, >( input: ActorConfigInput< @@ -726,16 +961,29 @@ export function test< TVars, TInput, TDatabase, + TEvents, + TQueues, TActions >, -): ActorConfig { +): ActorConfig< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues +> { const config = ActorConfigSchema.parse(input) as ActorConfig< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >; return config; } @@ -959,6 +1207,14 @@ export const DocActorConfigSchema = z .describe( "Map of action name to handler function. Defaults to an empty object.", ), + events: z + .record(z.string(), z.unknown()) + .optional() + .describe("Map of event names to schemas."), + queues: z + .record(z.string(), z.unknown()) + .optional() + .describe("Map of queue names to schemas."), options: DocActorOptionsSchema.optional(), }) .describe("Actor configuration passed to the actor() function."); diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts index 7c7c7dcb9b..4c3e20f234 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts @@ -10,15 +10,22 @@ import { } from "@/schemas/client-protocol-zod/mod"; import { bufferToArrayBuffer } from "@/utils"; import type { AnyDatabaseProvider } from "../database"; -import { InternalError } from "../errors"; +import { EventPayloadInvalid, InternalError } from "../errors"; import type { ActorInstance } from "../instance/mod"; import { CachedSerializer } from "../protocol/serde"; +import { + type EventSchemaConfig, + type InferEventArgs, + type InferSchemaMap, + type QueueSchemaConfig, + validateSchemaSync, +} from "../schema"; import type { ConnDriver } from "./driver"; import { type ConnDataInput, StateManager } from "./state-manager"; export type ConnId = string; -export type AnyConn = Conn; +export type AnyConn = Conn; export const CONN_CONNECTED_SYMBOL = Symbol("connected"); export const CONN_SPEAKS_RIVETKIT_SYMBOL = Symbol("speaksRivetKit"); @@ -34,10 +41,19 @@ export const CONN_SEND_MESSAGE_SYMBOL = Symbol("sendMessage"); * * @see {@link https://rivet.dev/docs/connections|Connection Documentation} */ -export class Conn { - #actor: ActorInstance; +export class Conn< + S, + CP, + CS, + V, + I, + DB extends AnyDatabaseProvider, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, +> { + #actor: ActorInstance; - get [CONN_ACTOR_SYMBOL](): ActorInstance { + get [CONN_ACTOR_SYMBOL](): ActorInstance { return this.#actor; } @@ -122,7 +138,7 @@ export class Conn { * @protected */ constructor( - actor: ActorInstance, + actor: ActorInstance, data: ConnDataInput, ) { this.#actor = actor; @@ -159,6 +175,14 @@ export class Conn { * @param args - The arguments for the event. * @see {@link https://rivet.dev/docs/events|Events Documentation} */ + send( + eventName: K, + ...args: InferEventArgs[K]> + ): void; + send( + eventName: keyof E extends never ? string : never, + ...args: unknown[] + ): void; send(eventName: string, ...args: unknown[]) { this.#assertConnected(); if (!this[CONN_SPEAKS_RIVETKIT_SYMBOL]) { @@ -168,11 +192,27 @@ export class Conn { connType: this[CONN_DRIVER_SYMBOL]?.type, }); } + + const payload = args.length === 1 ? args[0] : args; + const result = validateSchemaSync( + this.#actor.config.events, + eventName as keyof E & string, + payload, + ); + if (!result.success) { + throw new EventPayloadInvalid(eventName, result.issues); + } + const eventArgs = + args.length === 1 + ? [result.data] + : Array.isArray(result.data) + ? (result.data as unknown[]) + : args; this.#actor.emitTraceEvent("message.send", { "rivet.event.name": eventName, "rivet.conn.id": this.id, }); - const eventData = { name: eventName, args }; + const eventData = { name: eventName, args: eventArgs }; this[CONN_SEND_MESSAGE_SYMBOL]( new CachedSerializer( eventData, diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts index 63156343c3..753cabd59f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts @@ -39,7 +39,7 @@ export type ConnData = * Handles automatic state change detection for connection-specific state. */ export class StateManager { - #conn: Conn; + #conn: Conn; /** * Data representing this connection. @@ -50,7 +50,7 @@ export class StateManager { #data!: ConnData; constructor( - conn: Conn, + conn: Conn, data: ConnDataInput, ) { this.#conn = conn; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/action.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/action.ts index 5c1cccc022..6a49a2c17f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/action.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/action.ts @@ -2,6 +2,7 @@ import type { Conn } from "../conn/mod"; import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; import type { ActorInstance } from "../instance/mod"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ConnContext } from "./base/conn"; /** @@ -14,26 +15,33 @@ export class ActionContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ConnContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} /** * Extracts the ActionContext type from an ActorDefinition. */ -export type ActionContextOf = AD extends ActorDefinition< - infer S, - infer CP, - infer CS, - infer V, - infer I, - infer DB extends AnyDatabaseProvider, - any -> - ? ActionContext - : never; +export type ActionContextOf = + AD extends ActorDefinition< + infer S, + infer CP, + infer CS, + infer V, + infer I, + infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, + any + > + ? ActionContext + : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts index 4aba6dcc5a..04543d7d6d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts @@ -5,14 +5,22 @@ import type { Registry } from "@/registry"; import type { Conn, ConnId } from "../../conn/mod"; import type { AnyDatabaseProvider, InferDatabaseClient } from "../../database"; import type { ActorDefinition, AnyActorDefinition } from "../../definition"; +import * as errors from "../../errors"; +import { ActorKv } from "../../instance/kv"; import type { ActorInstance, AnyActorInstance, SaveStateOptions, } from "../../instance/mod"; -import { ActorKv } from "../../instance/kv"; import { ActorQueue } from "../../instance/queue"; import type { Schedule } from "../../schedule"; +import { + type EventSchemaConfig, + type InferEventArgs, + type InferSchemaMap, + type QueueSchemaConfig, + validateSchemaSync, +} from "../../schema"; export const ACTOR_CONTEXT_INTERNAL_SYMBOL = Symbol.for( "rivetkit.actorContextInternal", @@ -28,6 +36,8 @@ export class ActorContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > { [ACTOR_CONTEXT_INTERNAL_SYMBOL]!: AnyActorInstance; #actor: ActorInstance< @@ -36,11 +46,22 @@ export class ActorContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >; #kv: ActorKv | undefined; #queue: - | ActorQueue + | ActorQueue< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > | undefined; constructor( @@ -50,7 +71,9 @@ export class ActorContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ) { this.#actor = actor; @@ -93,9 +116,36 @@ export class ActorContext< * @param name - The name of the event. * @param args - The arguments to send with the event. */ - broadcast>(name: string, ...args: Args): void { + broadcast( + name: K, + ...args: InferEventArgs[K]> + ): void; + broadcast( + name: keyof TEvents extends never ? string : never, + ...args: Array + ): void; + broadcast(name: string, ...args: Array): void { + const payload = args.length === 1 ? args[0] : args; + const result = validateSchemaSync( + this.#actor.config.events, + name as keyof TEvents & string, + payload, + ); + if (!result.success) { + throw new errors.EventPayloadInvalid(name, result.issues); + } + if (args.length === 1) { + this.#actor.eventManager.broadcast(name, result.data); + return; + } + if (Array.isArray(result.data)) { + this.#actor.eventManager.broadcast( + name, + ...(result.data as unknown[]), + ); + return; + } this.#actor.eventManager.broadcast(name, ...args); - return; } /** @@ -114,7 +164,9 @@ export class ActorContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > { if (!this.#queue) { this.#queue = new ActorQueue( @@ -165,7 +217,16 @@ export class ActorContext< */ get conns(): Map< ConnId, - Conn + Conn< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > > { return this.#actor.conns; } @@ -227,6 +288,15 @@ export class ActorContext< return this.#actor.abortSignal; } + /** + * True when the actor is stopping. + * + * Alias for `c.abortSignal.aborted`. + */ + get aborted(): boolean { + return this.#actor.abortSignal.aborted; + } + /** * Forces the actor to sleep. * @@ -258,7 +328,9 @@ export type ActorContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? ActorContext + ? ActorContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn-init.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn-init.ts index a5f5b002df..227ff471c9 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn-init.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn-init.ts @@ -1,9 +1,7 @@ import type { AnyDatabaseProvider } from "../../database"; -import type { - ActorDefinition, - AnyActorDefinition, -} from "../../definition"; +import type { ActorDefinition, AnyActorDefinition } from "../../definition"; import type { ActorInstance } from "../../instance/mod"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../../schema"; import { ActorContext } from "./actor"; /** @@ -15,7 +13,18 @@ export abstract class ConnInitContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, -> extends ActorContext { + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, +> extends ActorContext< + TState, + never, + never, + TVars, + TInput, + TDatabase, + TEvents, + TQueues +> { /** * The incoming request that initiated the connection. * May be undefined for connections initiated without a direct HTTP request. @@ -26,7 +35,16 @@ export abstract class ConnInitContext< * @internal */ constructor( - actor: ActorInstance, + actor: ActorInstance< + TState, + any, + any, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, request: Request | undefined, ) { super(actor as any); @@ -42,7 +60,9 @@ export type ConnInitContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? ConnInitContext + ? ConnInitContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn.ts index 06be2482c3..6c6c5f3506 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn.ts @@ -1,10 +1,8 @@ import type { Conn } from "../../conn/mod"; import type { AnyDatabaseProvider } from "../../database"; -import type { - ActorDefinition, - AnyActorDefinition, -} from "../../definition"; +import type { ActorDefinition, AnyActorDefinition } from "../../definition"; import type { ActorInstance } from "../../instance/mod"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../../schema"; import { ActorContext } from "./actor"; /** @@ -18,13 +16,17 @@ export abstract class ConnContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > { /** * @internal @@ -36,7 +38,9 @@ export abstract class ConnContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, public readonly conn: Conn< TState, @@ -44,21 +48,26 @@ export abstract class ConnContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ) { super(actor); } } -export type ConnContextOf = AD extends ActorDefinition< - infer S, - infer CP, - infer CS, - infer V, - infer I, - infer DB extends AnyDatabaseProvider, - any -> - ? ConnContext - : never; +export type ConnContextOf = + AD extends ActorDefinition< + infer S, + infer CP, + infer CS, + infer V, + infer I, + infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, + any + > + ? ConnContext + : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-action-response.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-action-response.ts index 309ba14bdb..51918bfe3c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-action-response.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-action-response.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -12,13 +13,17 @@ export class BeforeActionResponseContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} export type BeforeActionResponseContextOf = @@ -29,7 +34,9 @@ export type BeforeActionResponseContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? BeforeActionResponseContext + ? BeforeActionResponseContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-connect.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-connect.ts index 4a54b761f2..e640ad4f78 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-connect.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-connect.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ConnInitContext } from "./base/conn-init"; /** @@ -10,7 +11,9 @@ export class BeforeConnectContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, -> extends ConnInitContext {} + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, +> extends ConnInitContext {} export type BeforeConnectContextOf = AD extends ActorDefinition< @@ -20,7 +23,9 @@ export type BeforeConnectContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? BeforeConnectContext + ? BeforeConnectContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/connect.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/connect.ts index 2c1472ac92..1e7d7fbcb3 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/connect.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/connect.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ConnContext } from "./base/conn"; /** @@ -12,13 +13,17 @@ export class ConnectContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ConnContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} export type ConnectContextOf = @@ -29,7 +34,9 @@ export type ConnectContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? ConnectContext + ? ConnectContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-conn-state.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-conn-state.ts index 9953be1470..86bc78496c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-conn-state.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-conn-state.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ConnInitContext } from "./base/conn-init"; /** @@ -11,7 +12,9 @@ export class CreateConnStateContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, -> extends ConnInitContext {} + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, +> extends ConnInitContext {} export type CreateConnStateContextOf = AD extends ActorDefinition< @@ -21,7 +24,9 @@ export type CreateConnStateContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? CreateConnStateContext + ? CreateConnStateContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-vars.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-vars.ts index ef611cefb6..8282acc892 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-vars.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create-vars.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -9,8 +10,18 @@ export class CreateVarsContext< TState, TInput, TDatabase extends AnyDatabaseProvider, -> extends ActorContext {} - + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, +> extends ActorContext< + TState, + never, + never, + never, + TInput, + TDatabase, + TEvents, + TQueues +> {} export type CreateVarsContextOf = AD extends ActorDefinition< @@ -20,7 +31,9 @@ export type CreateVarsContextOf = any, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? CreateVarsContext + ? CreateVarsContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create.ts index 17f860781d..16d9cfc081 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/create.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -9,8 +10,18 @@ export class CreateContext< TState, TInput, TDatabase extends AnyDatabaseProvider, -> extends ActorContext {} - + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, +> extends ActorContext< + TState, + never, + never, + never, + TInput, + TDatabase, + TEvents, + TQueues +> {} export type CreateContextOf = AD extends ActorDefinition< @@ -20,7 +31,9 @@ export type CreateContextOf = any, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? CreateContext + ? CreateContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/destroy.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/destroy.ts index c6304e1a84..85f9ddee3a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/destroy.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/destroy.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -12,16 +13,19 @@ export class DestroyContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} - export type DestroyContextOf = AD extends ActorDefinition< infer S, @@ -30,7 +34,9 @@ export type DestroyContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? DestroyContext + ? DestroyContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/disconnect.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/disconnect.ts index 42743546f6..2a38c5e4f1 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/disconnect.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/disconnect.ts @@ -1,6 +1,7 @@ import type { Conn } from "../conn/mod"; -import type { ActorDefinition, AnyActorDefinition } from "../definition"; import type { AnyDatabaseProvider } from "../database"; +import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -13,16 +14,19 @@ export class DisconnectContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} - export type DisconnectContextOf = AD extends ActorDefinition< infer S, @@ -31,7 +35,9 @@ export type DisconnectContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? DisconnectContext + ? DisconnectContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/request.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/request.ts index 632da2407c..0b53436f64 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/request.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/request.ts @@ -1,7 +1,8 @@ import type { Conn } from "../conn/mod"; -import type { ActorDefinition, AnyActorDefinition } from "../definition"; import type { AnyDatabaseProvider } from "../database"; +import type { ActorDefinition, AnyActorDefinition } from "../definition"; import type { ActorInstance } from "../instance/mod"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ConnContext } from "./base/conn"; /** @@ -14,13 +15,17 @@ export class RequestContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ConnContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > { /** * The incoming HTTP request. @@ -38,9 +43,20 @@ export class RequestContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues + >, + conn: Conn< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues >, - conn: Conn, request?: Request, ) { super(actor, conn); @@ -48,7 +64,6 @@ export class RequestContext< } } - export type RequestContextOf = AD extends ActorDefinition< infer S, @@ -57,7 +72,9 @@ export type RequestContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? RequestContext + ? RequestContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/run.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/run.ts index 63aa2a3ad1..a061db261e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/run.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/run.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -8,7 +9,7 @@ import { ActorContext } from "./base/actor"; * This context is passed to the `run` handler which executes after the actor * starts. It does not block actor startup and is intended for background tasks. * - * Use `c.abortSignal` to detect when the actor is stopping and gracefully exit. + * Use `c.aborted` to detect when the actor is stopping and gracefully exit. */ export class RunContext< TState, @@ -17,13 +18,17 @@ export class RunContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} export type RunContextOf = @@ -34,7 +39,9 @@ export type RunContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? RunContext + ? RunContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/sleep.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/sleep.ts index 1614d1927c..92d1f50c0a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/sleep.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/sleep.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -12,16 +13,19 @@ export class SleepContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} - export type SleepContextOf = AD extends ActorDefinition< infer S, @@ -30,7 +34,9 @@ export type SleepContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? SleepContext + ? SleepContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/state-change.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/state-change.ts index d323c25839..e297414443 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/state-change.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/state-change.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -12,16 +13,19 @@ export class StateChangeContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} - export type StateChangeContextOf = AD extends ActorDefinition< infer S, @@ -30,7 +34,9 @@ export type StateChangeContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? StateChangeContext + ? StateChangeContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/wake.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/wake.ts index b2c4402a2e..138b3e9d9d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/wake.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/wake.ts @@ -1,5 +1,6 @@ import type { AnyDatabaseProvider } from "../database"; import type { ActorDefinition, AnyActorDefinition } from "../definition"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ActorContext } from "./base/actor"; /** @@ -12,16 +13,19 @@ export class WakeContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ActorContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > {} - export type WakeContextOf = AD extends ActorDefinition< infer S, @@ -30,7 +34,9 @@ export type WakeContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? WakeContext + ? WakeContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/websocket.ts b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/websocket.ts index b274953cd9..17e882d841 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/contexts/websocket.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/contexts/websocket.ts @@ -1,7 +1,8 @@ import type { Conn } from "../conn/mod"; -import type { ActorDefinition, AnyActorDefinition } from "../definition"; import type { AnyDatabaseProvider } from "../database"; +import type { ActorDefinition, AnyActorDefinition } from "../definition"; import type { ActorInstance } from "../instance/mod"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { ConnContext } from "./base/conn"; /** @@ -14,13 +15,17 @@ export class WebSocketContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > extends ConnContext< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > { /** * The incoming HTTP request that initiated the WebSocket upgrade. @@ -38,9 +43,20 @@ export class WebSocketContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues + >, + conn: Conn< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues >, - conn: Conn, request?: Request, ) { super(actor, conn); @@ -48,7 +64,6 @@ export class WebSocketContext< } } - export type WebSocketContextOf = AD extends ActorDefinition< infer S, @@ -57,7 +72,9 @@ export type WebSocketContextOf = infer V, infer I, infer DB extends AnyDatabaseProvider, + infer E extends EventSchemaConfig, + infer Q extends QueueSchemaConfig, any > - ? WebSocketContext + ? WebSocketContext : never; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/definition.ts b/rivetkit-typescript/packages/rivetkit/src/actor/definition.ts index 53a2fc3111..9394bef2dd 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/definition.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/definition.ts @@ -1,7 +1,10 @@ import type { RegistryConfig } from "@/registry/config"; +import { DeepMutable } from "@/utils"; import type { Actions, ActorConfig } from "./config"; +import type { ActionContextOf, ActorContext } from "./contexts"; import type { AnyDatabaseProvider } from "./database"; import { ActorInstance } from "./instance/mod"; +import type { EventSchemaConfig, QueueSchemaConfig } from "./schema"; export type AnyActorDefinition = ActorDefinition< any, @@ -10,6 +13,8 @@ export type AnyActorDefinition = ActorDefinition< any, any, any, + any, + any, any >; @@ -20,19 +25,30 @@ export class ActorDefinition< V, I, DB extends AnyDatabaseProvider, - R extends Actions, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, + R extends Actions = Actions< + S, + CP, + CS, + V, + I, + DB, + E, + Q + >, > { - #config: ActorConfig; + #config: ActorConfig; - constructor(config: ActorConfig) { + constructor(config: ActorConfig) { this.#config = config; } - get config(): ActorConfig { + get config(): ActorConfig { return this.#config; } - instantiate(): ActorInstance { + instantiate(): ActorInstance { return new ActorInstance(this.#config); } } diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts b/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts index 6291566e24..d31f3cc339 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/errors.ts @@ -221,39 +221,72 @@ export class QueueMessageInvalid extends ActorError { } } -export class QueueCompleteNotAllowed extends ActorError { - constructor() { +export class EventPayloadInvalid extends ActorError { + constructor(name: string, issues?: unknown[]) { + super( + "event", + "invalid_payload", + `Event payload failed validation for '${name}'.`, + { public: true, metadata: { name, issues } }, + ); + } +} + +export class QueuePayloadInvalid extends ActorError { + constructor(name: string, issues?: unknown[]) { super( "queue", - "complete_not_allowed", - "Queue message completion is only allowed when wait is enabled.", - { public: true }, + "invalid_payload", + `Queue payload failed validation for '${name}'.`, + { public: true, metadata: { name, issues } }, ); } } -export class QueueMessagePending extends ActorError { - constructor() { +export class QueueCompletionPayloadInvalid extends ActorError { + constructor(name: string, issues?: unknown[]) { super( "queue", - "message_pending", - "Queue message is already pending completion.", - { public: true }, + "invalid_completion_payload", + `Queue completion payload failed validation for '${name}'.`, + { public: true, metadata: { name, issues } }, ); } } export class QueueAlreadyCompleted extends ActorError { + constructor() { + super("queue", "already_completed", "Queue message was already completed.", { + public: true, + }); + } +} + +export class QueuePreviousMessageNotCompleted extends ActorError { constructor() { super( "queue", - "already_completed", - "Queue message has already been completed.", + "previous_message_not_completed", + "Previous completable queue message is not completed. Call `message.complete(...)` before receiving the next message.", { public: true }, ); } } +export class QueueCompleteNotConfigured extends ActorError { + constructor(name: string) { + super( + "queue", + "complete_not_configured", + `Queue '${name}' does not support completion responses.`, + { + public: true, + metadata: { name }, + }, + ); + } +} + export class ActorAborted extends ActorError { constructor() { super("actor", "aborted", "Actor aborted.", { public: true }); diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts index fb0aa16af7..11ec88fb3d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/connection-manager.ts @@ -30,6 +30,7 @@ import { } from "../contexts"; import type { AnyDatabaseProvider } from "../database"; import { CachedSerializer } from "../protocol/serde"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { deadline } from "../utils"; import { makeConnKey } from "./keys"; import type { ActorInstance } from "./mod"; @@ -44,22 +45,24 @@ export class ConnectionManager< V, I, DB extends AnyDatabaseProvider, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, > { - #actor: ActorInstance; - #connections = new Map>(); + #actor: ActorInstance; + #connections = new Map>(); /** Connections that have had their state changed and need to be persisted. */ #connsWithPersistChanged = new Set(); - constructor(actor: ActorInstance) { + constructor(actor: ActorInstance) { this.#actor = actor; } - get connections(): Map> { + get connections(): Map> { return this.#connections; } - getConnForId(id: string): Conn | undefined { + getConnForId(id: string): Conn | undefined { return this.#connections.get(id); } @@ -71,7 +74,7 @@ export class ConnectionManager< this.#connsWithPersistChanged.clear(); } - markConnWithPersistChanged(conn: Conn) { + markConnWithPersistChanged(conn: Conn) { invariant( conn.isHibernatable, "cannot mark non-hibernatable conn for persist", @@ -100,7 +103,7 @@ export class ConnectionManager< requestHeaders: Record | undefined, isHibernatable: boolean, isRestoringHibernatable: boolean, - ): Promise> { + ): Promise> { this.#actor.assertReady(); // TODO: Add back @@ -169,7 +172,7 @@ export class ConnectionManager< } // Create connection instance - const conn = new Conn(this.#actor, connData); + const conn = new Conn(this.#actor, connData); conn[CONN_DRIVER_SYMBOL] = driver; return conn; @@ -183,7 +186,7 @@ export class ConnectionManager< * be messed up and cause race conditions that can drop WebSocket messages. * So all async work in prepareConn. */ - connectConn(conn: Conn) { + connectConn(conn: Conn) { invariant(!this.#connections.has(conn.id), "conn already connected"); this.#connections.set(conn.id, conn); @@ -236,7 +239,9 @@ export class ConnectionManager< } } - #reconnectHibernatableConn(driver: ConnDriver): Conn { + #reconnectHibernatableConn( + driver: ConnDriver, + ): Conn { invariant(driver.hibernatable, "missing requestIdBuf"); const existingConn = this.findHibernatableConn( driver.hibernatable.gatewayId, @@ -271,7 +276,7 @@ export class ConnectionManager< return existingConn; } - #disconnectExistingDriver(conn: Conn) { + #disconnectExistingDriver(conn: Conn) { const driver = conn[CONN_DRIVER_SYMBOL]; if (driver?.disconnect) { driver.disconnect( @@ -287,7 +292,7 @@ export class ConnectionManager< * * This is called by `Conn.disconnect`. This should not call `Conn.disconnect.` */ - async connDisconnected(conn: Conn) { + async connDisconnected(conn: Conn) { // Remove from tracking this.#connections.delete(conn.id); @@ -400,7 +405,7 @@ export class ConnectionManager< request: Request | undefined, requestPath: string | undefined, requestHeaders: Record | undefined, - ): Promise> { + ): Promise> { const conn = await this.prepareConn( driver, params, @@ -422,7 +427,7 @@ export class ConnectionManager< restoreConnections(connections: PersistedConn[]) { for (const connPersist of connections) { // Create connection instance - const conn = new Conn(this.#actor, { + const conn = new Conn(this.#actor, { hibernatable: connPersist, }); this.#connections.set(conn.id, conn); @@ -448,7 +453,7 @@ export class ConnectionManager< findHibernatableConn( gatewayIdBuf: ArrayBuffer, requestIdBuf: ArrayBuffer, - ): Conn | undefined { + ): Conn | undefined { return Array.from(this.#connections.values()).find((conn) => { const connStateManager = conn[CONN_STATE_MANAGER_SYMBOL]; const h = connStateManager.hibernatableDataRaw; @@ -471,8 +476,7 @@ export class ConnectionManager< "actor.createConnState", undefined, () => { - const dataOrPromise = - createConnState!(ctx, params); + const dataOrPromise = createConnState!(ctx, params); if (dataOrPromise instanceof Promise) { return deadline( dataOrPromise, @@ -491,7 +495,7 @@ export class ConnectionManager< ); } - #callOnConnect(conn: Conn) { + #callOnConnect(conn: Conn) { const attributes = { "rivet.conn.id": conn.id, "rivet.conn.type": conn[CONN_DRIVER_SYMBOL]?.type, diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/event-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/event-manager.ts index 079fd6eace..7e6cf21bca 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/event-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/event-manager.ts @@ -18,17 +18,30 @@ import { import type { AnyDatabaseProvider } from "../database"; import * as errors from "../errors"; import { CachedSerializer } from "../protocol/serde"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import type { ActorInstance } from "./mod"; /** * Manages event subscriptions and broadcasting for actor instances. * Handles subscription tracking and efficient message distribution to connected clients. */ -export class EventManager { - #actor: ActorInstance; - #subscriptionIndex = new Map>>(); +export class EventManager< + S, + CP, + CS, + V, + I, + DB extends AnyDatabaseProvider, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, +> { + #actor: ActorInstance; + #subscriptionIndex = new Map< + string, + Set> + >(); - constructor(actor: ActorInstance) { + constructor(actor: ActorInstance) { this.#actor = actor; } @@ -43,7 +56,7 @@ export class EventManager { */ addSubscription( eventName: string, - connection: Conn, + connection: Conn, fromPersist: boolean, ) { // Check if already subscribed @@ -94,7 +107,7 @@ export class EventManager { */ removeSubscription( eventName: string, - connection: Conn, + connection: Conn, fromRemoveConn: boolean, ) { // Check if subscription exists @@ -241,7 +254,7 @@ export class EventManager { */ getSubscribers( eventName: string, - ): Set> | undefined { + ): Set> | undefined { return this.#subscriptionIndex.get(eventName); } @@ -264,7 +277,7 @@ export class EventManager { * * @param connection - The connection to clear subscriptions for */ - clearConnectionSubscriptions(connection: Conn) { + clearConnectionSubscriptions(connection: Conn) { for (const eventName of [...connection.subscriptions.values()]) { this.removeSubscription(eventName, connection, true); } diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts index 83fe4e0da7..78a6fa97d4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts @@ -1,11 +1,11 @@ -import invariant from "invariant"; +import type { OtlpExportTraceServiceRequestJson } from "@rivetkit/traces"; import { createTraces, type SpanHandle, type SpanStatusInput, type Traces, } from "@rivetkit/traces"; -import type { OtlpExportTraceServiceRequestJson } from "@rivetkit/traces"; +import invariant from "invariant"; import type { ActorKey } from "@/actor/mod"; import type { Client } from "@/client/client"; import { getBaseLogger, getIncludeTarget, type Logger } from "@/common/log"; @@ -18,7 +18,7 @@ import { CONN_VERSIONED, } from "@/schemas/actor-persist/versioned"; import { EXTRA_ERROR_LOG } from "@/utils"; -import { getRunFunction, type ActorConfig } from "../config"; +import { type ActorConfig, getRunFunction } from "../config"; import type { ConnDriver } from "../conn/driver"; import { createHttpDriver } from "../conn/drivers/http"; import { @@ -44,6 +44,7 @@ import * as errors from "../errors"; import { serializeActorKey } from "../keys"; import { processMessage } from "../protocol/old"; import { Schedule } from "../schedule"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { assertUnreachable, DeadlineError, @@ -53,7 +54,6 @@ import { import { ConnectionManager } from "./connection-manager"; import { EventManager } from "./event-manager"; import { KEYS } from "./keys"; -import { ActorTracesDriver } from "./traces-driver"; import { convertActorFromBarePersisted, type PersistedActor, @@ -61,6 +61,7 @@ import { import { QueueManager } from "./queue-manager"; import { ScheduleManager } from "./schedule-manager"; import { type SaveStateOptions, StateManager } from "./state-manager"; +import { ActorTracesDriver } from "./traces-driver"; export type { SaveStateOptions }; @@ -71,31 +72,50 @@ enum CanSleep { ActiveConns, ActiveHonoHttpRequests, ActiveKeepAwake, + ActiveRun, } /** Actor type alias with all `any` types. Used for `extends` in classes referencing this actor. */ -export type AnyActorInstance = ActorInstance; +export type AnyActorInstance = ActorInstance< + any, + any, + any, + any, + any, + any, + any, + any +>; export type ExtractActorState = - A extends ActorInstance + A extends ActorInstance ? State : never; export type ExtractActorConnParams = - A extends ActorInstance + A extends ActorInstance ? ConnParams : never; export type ExtractActorConnState = - A extends ActorInstance + A extends ActorInstance ? ConnState : never; // MARK: - Main ActorInstance Class -export class ActorInstance { +export class ActorInstance< + S, + CP, + CS, + V, + I, + DB extends AnyDatabaseProvider, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, +> { // MARK: - Core Properties - actorContext: ActorContext; - #config: ActorConfig; + actorContext: ActorContext; + #config: ActorConfig; driver!: ActorDriver; #inlineClient!: Client>; #actorId!: string; @@ -105,15 +125,15 @@ export class ActorInstance { #region!: string; // MARK: - Managers - connectionManager!: ConnectionManager; + connectionManager!: ConnectionManager; - stateManager!: StateManager; + stateManager!: StateManager; - eventManager!: EventManager; + eventManager!: EventManager; - #scheduleManager!: ScheduleManager; + #scheduleManager!: ScheduleManager; - queueManager!: QueueManager; + queueManager!: QueueManager; // MARK: - Logging #log!: Logger; @@ -145,6 +165,8 @@ export class ActorInstance { // MARK: - Background Tasks #backgroundPromises: Promise[] = []; #runPromise?: Promise; + #runHandlerActive = false; + #activeQueueWaitCount = 0; // MARK: - HTTP/WebSocket Tracking #activeHonoHttpRequests = 0; @@ -155,15 +177,14 @@ export class ActorInstance { // MARK: - Inspector #inspectorToken?: string; - #inspector!: ActorInspector; + #inspector: ActorInspector; // MARK: - Tracing #traces!: Traces; // MARK: - Constructor - constructor(config: ActorConfig) { + constructor(config: ActorConfig) { this.#config = config; - this.#inspector = new ActorInspector(this); this.actorContext = new ActorContext(this); this.#inspector = new ActorInspector(this); } @@ -230,14 +251,8 @@ export class ActorInstance { }); } - endTraceSpan( - handle: SpanHandle, - status?: SpanStatusInput, - ): void { - this.#traces.endSpan( - handle, - status ? { status } : undefined, - ); + endTraceSpan(handle: SpanHandle, status?: SpanStatusInput): void { + this.#traces.endSpan(handle, status ? { status } : undefined); } async runInTraceSpan( @@ -248,8 +263,7 @@ export class ActorInstance { const span = this.startTraceSpan(name, attributes); try { const result = this.#traces.withSpan(span, fn); - const resolved = - result instanceof Promise ? await result : result; + const resolved = result instanceof Promise ? await result : result; this.#traces.endSpan(span, { status: { code: "OK" }, }); @@ -280,7 +294,7 @@ export class ActorInstance { }); } - get conns(): Map> { + get conns(): Map> { return this.connectionManager.connections; } @@ -296,7 +310,7 @@ export class ActorInstance { return Object.keys(this.#config.actions ?? {}); } - get config(): ActorConfig { + get config(): ActorConfig { return this.#config; } @@ -580,7 +594,7 @@ export class ActorInstance { val: { eventName: string; subscribe: boolean }; }; }, - conn: Conn, + conn: Conn, ) { await processMessage(message, this, conn, { onExecuteAction: async (ctx, name, args) => { @@ -597,7 +611,7 @@ export class ActorInstance { // MARK: - Action Execution async executeAction( - ctx: ActionContext, + ctx: ActionContext, actionName: string, args: unknown[], ): Promise { @@ -619,15 +633,9 @@ export class ActorInstance { throw new errors.ActionNotFound(actionName); } - this.#activeKeepAwakeCount++; - this.resetSleepTimer(); - - const actionSpan = this.startTraceSpan( - `actor.action.${actionName}`, - { - "rivet.action.name": actionName, - }, - ); + const actionSpan = this.startTraceSpan(`actor.action.${actionName}`, { + "rivet.action.name": actionName, + }); let spanEnded = false; try { @@ -645,12 +653,9 @@ export class ActorInstance { ); let output: unknown; - const maybeThenable = outputOrPromise as { - then?: (onfulfilled?: unknown, onrejected?: unknown) => unknown; - }; - if (maybeThenable && typeof maybeThenable.then === "function") { + if (outputOrPromise instanceof Promise) { output = await deadline( - Promise.resolve(outputOrPromise), + outputOrPromise, this.#config.options.actionTimeout, ); } else { @@ -710,22 +715,13 @@ export class ActorInstance { status: { code: "OK" }, }); } - this.#activeKeepAwakeCount--; - if (this.#activeKeepAwakeCount < 0) { - this.#activeKeepAwakeCount = 0; - this.#rLog.warn({ - msg: "active keep awake count went below 0, this is a RivetKit bug", - ...EXTRA_ERROR_LOG, - }); - } - this.resetSleepTimer(); this.stateManager.savePersistThrottled(); } } // MARK: - HTTP/WebSocket Handlers async handleRawRequest( - conn: Conn, + conn: Conn, request: Request, ): Promise { this.assertReady(); @@ -742,10 +738,10 @@ export class ActorInstance { "http.url": request.url, "rivet.conn.id": conn.id, }, - async () => { - try { - const ctx = new RequestContext(this, conn, request); - const response = await onRequest(ctx, request); + async () => { + try { + const ctx = new RequestContext(this, conn, request); + const response = await onRequest(ctx, request); if (!response) { throw new errors.InvalidRequestHandlerResponse(); } @@ -764,7 +760,7 @@ export class ActorInstance { } handleRawWebSocket( - conn: Conn, + conn: Conn, websocket: UniversalWebSocket, request?: Request, ) { @@ -903,6 +899,24 @@ export class ActorInstance { } } + beginQueueWait() { + this.assertReady(true); + this.#activeQueueWaitCount++; + this.resetSleepTimer(); + } + + endQueueWait() { + this.#activeQueueWaitCount--; + if (this.#activeQueueWaitCount < 0) { + this.#activeQueueWaitCount = 0; + this.#rLog.warn({ + msg: "active queue wait count went below 0, this is a RivetKit bug", + ...EXTRA_ERROR_LOG, + }); + } + this.resetSleepTimer(); + } + // MARK: - Private Helper Methods #initializeTraces() { this.#traces = createTraces({ @@ -950,7 +964,10 @@ export class ActorInstance { const [value] = args; if (typeof value === "string") { message = value; - } else if (typeof value === "number" || typeof value === "boolean") { + } else if ( + typeof value === "number" || + typeof value === "boolean" + ) { message = String(value); } else if (value && typeof value === "object") { const maybeMsg = (value as { msg?: unknown }).msg; @@ -1100,19 +1117,23 @@ export class ActorInstance { let vars: V | undefined; if ("createVars" in this.#config) { const createVars = this.#config.createVars; - vars = await this.runInTraceSpan("actor.createVars", undefined, () => { - const dataOrPromise = createVars!( - this.actorContext as any, - this.driver.getContext(this.#actorId), - ); - if (dataOrPromise instanceof Promise) { - return deadline( - dataOrPromise, - this.#config.options.createVarsTimeout, + vars = await this.runInTraceSpan( + "actor.createVars", + undefined, + () => { + const dataOrPromise = createVars!( + this.actorContext as any, + this.driver.getContext(this.#actorId), ); - } - return dataOrPromise; - }); + if (dataOrPromise instanceof Promise) { + return deadline( + dataOrPromise, + this.#config.options.createVarsTimeout, + ); + } + return dataOrPromise; + }, + ); } else if ("vars" in this.#config) { vars = structuredClone(this.#config.vars); } else { @@ -1138,15 +1159,19 @@ export class ActorInstance { const onSleep = this.#config.onSleep; try { this.#rLog.debug({ msg: "calling onSleep" }); - await this.runInTraceSpan("actor.onSleep", undefined, async () => { - const result = onSleep(this.actorContext); - if (result instanceof Promise) { - await deadline( - result, - this.#config.options.onSleepTimeout, - ); - } - }); + await this.runInTraceSpan( + "actor.onSleep", + undefined, + async () => { + const result = onSleep(this.actorContext); + if (result instanceof Promise) { + await deadline( + result, + this.#config.options.onSleepTimeout, + ); + } + }, + ); this.#rLog.debug({ msg: "onSleep completed" }); } catch (error) { if (error instanceof DeadlineError) { @@ -1162,34 +1187,23 @@ export class ActorInstance { } async #callOnDestroy() { - // Clean up database first - if ("db" in this.#config && this.#config.db && this.#db) { - try { - this.#rLog.debug({ msg: "cleaning up database" }); - await this.#config.db.onDestroy?.(this.#db); - this.#rLog.debug({ msg: "database cleanup completed" }); - } catch (error) { - this.#rLog.error({ - msg: "error cleaning up database", - error: stringifyError(error), - }); - } - } - - // Then call user's onDestroy if (this.#config.onDestroy) { const onDestroy = this.#config.onDestroy; try { this.#rLog.debug({ msg: "calling onDestroy" }); - await this.runInTraceSpan("actor.onDestroy", undefined, async () => { - const result = onDestroy(this.actorContext); - if (result instanceof Promise) { - await deadline( - result, - this.#config.options.onDestroyTimeout, - ); - } - }); + await this.runInTraceSpan( + "actor.onDestroy", + undefined, + async () => { + const result = onDestroy(this.actorContext); + if (result instanceof Promise) { + await deadline( + result, + this.#config.options.onDestroyTimeout, + ); + } + }, + ); this.#rLog.debug({ msg: "onDestroy completed" }); } catch (error) { if (error instanceof DeadlineError) { @@ -1209,6 +1223,8 @@ export class ActorInstance { if (!runFn) return; this.#rLog.debug({ msg: "starting run handler" }); + this.#runHandlerActive = true; + this.resetSleepTimer(); const runSpan = this.startTraceSpan("actor.run"); const runResult = this.#traces.withSpan(runSpan, () => @@ -1252,9 +1268,15 @@ export class ActorInstance { error: stringifyError(error), }); this.startDestroy(); + }) + .finally(() => { + this.#runHandlerActive = false; + this.resetSleepTimer(); }); } else if (runSpan.isActive()) { this.endTraceSpan(runSpan, { code: "OK" }); + this.#runHandlerActive = false; + this.resetSleepTimer(); } } @@ -1284,48 +1306,31 @@ export class ActorInstance { async #setupDatabase() { if ("db" in this.#config && this.#config.db) { - try { - const client = await this.#config.db.createClient({ - actorId: this.#actorId, - overrideRawDatabaseClient: this.driver.overrideRawDatabaseClient - ? () => this.driver.overrideRawDatabaseClient!(this.#actorId) - : undefined, - overrideDrizzleDatabaseClient: this.driver - .overrideDrizzleDatabaseClient - ? () => this.driver.overrideDrizzleDatabaseClient!(this.#actorId) + const client = await this.#config.db.createClient({ + actorId: this.#actorId, + overrideRawDatabaseClient: this.driver.overrideRawDatabaseClient + ? () => this.driver.overrideRawDatabaseClient!(this.#actorId) + : undefined, + overrideDrizzleDatabaseClient: + this.driver.overrideDrizzleDatabaseClient + ? () => + this.driver.overrideDrizzleDatabaseClient!( + this.#actorId, + ) : undefined, - kv: { - batchPut: (entries) => - this.driver.kvBatchPut(this.#actorId, entries), - batchGet: (keys) => this.driver.kvBatchGet(this.#actorId, keys), - batchDelete: (keys) => - this.driver.kvBatchDelete(this.#actorId, keys), - }, - sqliteVfs: this.driver.sqliteVfs, - }); - this.#rLog.info({ msg: "database migration starting" }); - await this.#config.db.onMigrate?.(client); - this.#rLog.info({ msg: "database migration complete" }); - this.#db = client; - } catch (error) { - // Ensure error is properly formatted - if (error instanceof Error) { - this.#rLog.error({ - msg: "database setup failed", - error: stringifyError(error), - }); - throw error; - } - const wrappedError = new Error( - `Database setup failed: ${String(error)}`, - ); - this.#rLog.error({ - msg: "database setup failed with non-Error object", - error: String(error), - errorType: typeof error, - }); - throw wrappedError; - } + kv: { + batchPut: (entries) => + this.driver.kvBatchPut(this.#actorId, entries), + batchGet: (keys) => this.driver.kvBatchGet(this.#actorId, keys), + batchDelete: (keys) => + this.driver.kvBatchDelete(this.#actorId, keys), + }, + sqliteVfs: this.driver.sqliteVfs, + }); + this.#rLog.info({ msg: "database migration starting" }); + await this.#config.db.onMigrate?.(client); + this.#rLog.info({ msg: "database migration complete" }); + this.#db = client; } } @@ -1428,6 +1433,9 @@ export class ActorInstance { if (this.#activeHonoHttpRequests > 0) return CanSleep.ActiveHonoHttpRequests; if (this.#activeKeepAwakeCount > 0) return CanSleep.ActiveKeepAwake; + if (this.#runHandlerActive && this.#activeQueueWaitCount === 0) { + return CanSleep.ActiveRun; + } for (const _conn of this.connectionManager.connections.values()) { // TODO: Add back diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue-manager.ts index 73847f8115..5a2fc5a777 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue-manager.ts @@ -6,10 +6,10 @@ import { QUEUE_METADATA_VERSIONED, } from "@/schemas/actor-persist/versioned"; import { promiseWithResolvers } from "@/utils"; -import { loggerWithoutContext } from "@/actor/log"; import type { AnyDatabaseProvider } from "../database"; import type { ActorDriver } from "../driver"; import * as errors from "../errors"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { decodeQueueMessageKey, KEYS, makeQueueMessageKey } from "./keys"; import type { ActorInstance } from "./mod"; @@ -18,15 +18,6 @@ export interface QueueMessage { name: string; body: unknown; createdAt: number; - failureCount: number; - availableAt: number; - inFlight: boolean; - inFlightAt?: number; -} - -export interface QueueCompletionResult { - status: "completed" | "timedOut"; - response?: unknown; } interface QueueMetadata { @@ -34,40 +25,17 @@ interface QueueMetadata { size: number; } -interface EnqueueOptions { - deferWaiters?: boolean; -} - interface QueueWaiter { id: string; - nameSet: Set; + nameSet?: Set; count: number; - wait: boolean; + completable: boolean; resolve: (messages: QueueMessage[]) => void; reject: (error: Error) => void; - signal?: AbortSignal; - timeoutHandle?: ReturnType; -} - -interface QueueNameWaiter { - id: string; - nameSet: Set; - resolve: () => void; - reject: (error: Error) => void; - signal?: AbortSignal; - abortHandler?: () => void; -} - -interface QueueCompletionWaiter { - id: string; - messageId: bigint; - resolve: (result: QueueCompletionResult) => void; - reject: (error: Error) => void; - timeoutHandle?: ReturnType; } interface MessageListener { - nameSet: Set; + nameSet?: Set; resolve: () => void; reject: (error: Error) => void; actorAbortCleanup?: () => void; @@ -80,25 +48,33 @@ const DEFAULT_METADATA: QueueMetadata = { size: 0, }; -const PENDING_WARNING_MS = 30_000; -const BACKOFF_INITIAL_MS = 1_000; -const BACKOFF_MAX_MS = 5 * 60_000; +interface PendingCompletion { + resolve: (result: { + status: "completed" | "timedOut"; + response?: unknown; + }) => void; + timeoutHandle?: ReturnType; +} -export class QueueManager { - #actor: ActorInstance; +export class QueueManager< + S, + CP, + CS, + V, + I, + DB extends AnyDatabaseProvider, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, +> { + #actor: ActorInstance; #driver: ActorDriver; #waiters = new Map(); - #nameWaiters = new Map(); - #completionWaiters = new Map(); #metadata: QueueMetadata = { ...DEFAULT_METADATA }; - #pendingMessageId: bigint | undefined; - #pendingWarningHandle: ReturnType | undefined; - #redeliveryTimeout: ReturnType | undefined; - #redeliveryAt: number | undefined; #messageListeners = new Set(); + #pendingCompletions = new Map(); constructor( - actor: ActorInstance, + actor: ActorInstance, driver: ActorDriver, ) { this.#actor = actor; @@ -137,16 +113,10 @@ export class QueueManager { await this.#rebuildMetadata(); } this.#actor.inspector.updateQueueSize(this.#metadata.size); - - await this.#recoverInFlightMessages(); } /** Adds a message to the queue with the given name and body. */ - async enqueue( - name: string, - body: unknown, - options: EnqueueOptions = {}, - ): Promise { + async enqueue(name: string, body: unknown): Promise { this.#actor.assertReady(); const sizeLimit = this.#actor.config.options.maxQueueSize; @@ -165,16 +135,15 @@ export class QueueManager { const createdAt = Date.now(); const bodyCborBuffer = cbor.encode(body); - const availableAt = createdAt; const encodedMessage = QUEUE_MESSAGE_VERSIONED.serializeWithEmbeddedVersion( { name, body: new Uint8Array(bodyCborBuffer).buffer as ArrayBuffer, createdAt: BigInt(createdAt), - failureCount: 0, - availableAt: BigInt(availableAt), - inFlight: false, + failureCount: null, + availableAt: null, + inFlight: null, inFlightAt: null, }, ACTOR_PERSIST_CURRENT_VERSION, @@ -208,91 +177,149 @@ export class QueueManager { name, body, createdAt, - failureCount: 0, - availableAt, - inFlight: false, - inFlightAt: undefined, }; this.#actor.resetSleepTimer(); - if (!options.deferWaiters) { - await this.#maybeResolveWaiters(); - } + await this.#maybeResolveWaiters(); this.#notifyMessageListeners(name); return message; } + /** + * Adds a message and waits for completion. + */ async enqueueAndWait( name: string, body: unknown, timeout?: number, - ): Promise { - const message = await this.enqueue(name, body, { - deferWaiters: true, - }); - const completionPromise = this.waitForCompletion(message.id, timeout); - await this.#maybeResolveWaiters(); - return await completionPromise; + ): Promise<{ status: "completed" | "timedOut"; response?: unknown }> { + if (timeout !== undefined && timeout <= 0) { + return { status: "timedOut" }; + } + + const message = await this.enqueue(name, body); + const messageId = message.id.toString(); + const { promise, resolve } = promiseWithResolvers<{ + status: "completed" | "timedOut"; + response?: unknown; + }>(() => {}); + + const pending: PendingCompletion = { resolve }; + if (timeout !== undefined) { + pending.timeoutHandle = setTimeout(() => { + this.#pendingCompletions.delete(messageId); + resolve({ status: "timedOut" }); + }, timeout); + } + this.#pendingCompletions.set(messageId, pending); + + return await promise; + } + + async completeMessage( + message: QueueMessage, + response?: unknown, + ): Promise { + await this.completeMessageById(message.id, response); + } + + async completeMessageById( + messageId: bigint, + response?: unknown, + ): Promise { + const messageIdString = messageId.toString(); + const pending = this.#pendingCompletions.get(messageIdString); + if (pending) { + if (pending.timeoutHandle) { + clearTimeout(pending.timeoutHandle); + } + this.#pendingCompletions.delete(messageIdString); + pending.resolve({ status: "completed", response }); + } + + await this.deleteMessagesById([messageId]); } /** Receives messages from the queue matching the given names. Waits until messages are available or timeout is reached. */ async receive( - names: string[], + names: string[] | undefined, count: number, timeout?: number, abortSignal?: AbortSignal, - wait: boolean = false, - ): Promise { + completable = false, + ): Promise { this.#actor.assertReady(); - if (this.#pendingMessageId !== undefined) { - throw new errors.QueueMessagePending(); - } const limitedCount = Math.max(1, count); - const nameSet = new Set(names); + const nameSet = + names && names.length > 0 ? new Set(names) : undefined; const immediate = await this.#drainMessages( nameSet, limitedCount, - wait, + completable, ); - if (immediate.length > 0 || timeout === 0) { - return timeout === 0 && immediate.length === 0 ? [] : immediate; + if (immediate.length > 0) { + return immediate; + } + if (timeout === 0) { + return []; } - const { promise, resolve, reject } = - promiseWithResolvers((reason) => loggerWithoutContext().warn({ msg: "unhandled queue message waiter rejection", reason })); + const { promise, resolve, reject } = promiseWithResolvers< + QueueMessage[] + >(() => {}); const waiterId = crypto.randomUUID(); + let timeoutHandle: ReturnType | undefined; + let cleanedUp = false; + let actorAbortCleanup: (() => void) | undefined; + let signalAbortCleanup: (() => void) | undefined; + + const cleanup = () => { + if (cleanedUp) { + return; + } + cleanedUp = true; + this.#waiters.delete(waiterId); + if (timeoutHandle) { + clearTimeout(timeoutHandle); + timeoutHandle = undefined; + } + actorAbortCleanup?.(); + signalAbortCleanup?.(); + this.#actor.endQueueWait(); + }; + const resolveWaiter = (messages: QueueMessage[]) => { + cleanup(); + resolve(messages); + }; + const rejectWaiter = (error: Error) => { + cleanup(); + reject(error); + }; + const waiter: QueueWaiter = { id: waiterId, nameSet, count: limitedCount, - wait, - resolve, - reject, - signal: abortSignal, + completable, + resolve: resolveWaiter, + reject: rejectWaiter, }; + this.#actor.beginQueueWait(); + if (timeout !== undefined) { - waiter.timeoutHandle = setTimeout(() => { - this.#waiters.delete(waiterId); - resolve([]); + timeoutHandle = setTimeout(() => { + resolveWaiter([]); }, timeout); } const onAbort = () => { - this.#waiters.delete(waiterId); - if (waiter.timeoutHandle) { - clearTimeout(waiter.timeoutHandle); - } - reject(new errors.ActorAborted()); + rejectWaiter(new errors.ActorAborted()); }; const onStop = () => { - this.#waiters.delete(waiterId); - if (waiter.timeoutHandle) { - clearTimeout(waiter.timeoutHandle); - } - reject(new errors.ActorAborted()); + rejectWaiter(new errors.ActorAborted()); }; const actorAbortSignal = this.#actor.abortSignal; if (actorAbortSignal.aborted) { @@ -300,6 +327,8 @@ export class QueueManager { return promise; } actorAbortSignal.addEventListener("abort", onStop, { once: true }); + actorAbortCleanup = () => + actorAbortSignal.removeEventListener("abort", onStop); if (abortSignal) { if (abortSignal.aborted) { @@ -307,79 +336,41 @@ export class QueueManager { return promise; } abortSignal.addEventListener("abort", onAbort, { once: true }); + signalAbortCleanup = () => + abortSignal.removeEventListener("abort", onAbort); } this.#waiters.set(waiterId, waiter); return promise; } - /** Waits for a specific queue message to complete. */ - async waitForCompletion( - messageId: bigint, - timeout?: number, - ): Promise { - const { promise, resolve, reject } = - promiseWithResolvers((reason) => loggerWithoutContext().warn({ msg: "unhandled queue completion waiter rejection", reason })); - const waiterId = crypto.randomUUID(); - - const waiter: QueueCompletionWaiter = { - id: waiterId, - messageId, - resolve, - reject, - }; - - if (timeout !== undefined) { - waiter.timeoutHandle = setTimeout(() => { - this.#completionWaiters.delete(messageId); - resolve({ status: "timedOut" }); - }, timeout); - } - - this.#completionWaiters.set(messageId, waiter); - return promise; - } - - /** Completes a pending message and optionally responds to any waiter. */ - async complete(message: QueueMessage, response?: unknown): Promise { - if (this.#pendingMessageId !== message.id) { - throw new errors.QueueAlreadyCompleted(); - } - this.#pendingMessageId = undefined; - if (this.#pendingWarningHandle) { - clearTimeout(this.#pendingWarningHandle); - this.#pendingWarningHandle = undefined; - } - - await this.#removeMessages([message], { resolveWaiters: false }); - this.#resolveCompletionWaiter(message.id, { - status: "completed", - response, - }); - - await this.#maybeResolveWaiters(); - } - - /** Waits for messages with any of the specified names to appear in the queue. */ async waitForNames( - names: string[], + names: readonly string[] | undefined, abortSignal?: AbortSignal, ): Promise { - const nameSet = new Set(names); + const nameSet = + names && names.length > 0 ? new Set(names) : undefined; const existing = await this.#loadQueueMessages(); - if (existing.some((message) => nameSet.has(message.name))) { + if (nameSet) { + if (existing.some((message) => nameSet.has(message.name))) { + return; + } + } else if (existing.length > 0) { return; } return await new Promise((resolve, reject) => { + this.#actor.beginQueueWait(); const listener: MessageListener = { nameSet, resolve: () => { this.#removeMessageListener(listener); + this.#actor.endQueueWait(); resolve(); }, reject: (error) => { this.#removeMessageListener(listener); + this.#actor.endQueueWait(); reject(error); }, }; @@ -419,10 +410,7 @@ export class QueueManager { } /** Deletes messages matching the provided IDs. Returns the IDs that were removed. */ - async deleteMessagesById( - ids: bigint[], - options: { resolveWaiters?: boolean } = {}, - ): Promise { + async deleteMessagesById(ids: bigint[]): Promise { if (ids.length === 0) { return []; } @@ -434,52 +422,31 @@ export class QueueManager { if (toRemove.length === 0) { return []; } - await this.#removeMessages(toRemove, { - resolveWaiters: options.resolveWaiters ?? true, - }); + await this.#removeMessages(toRemove); return toRemove.map((entry) => entry.id); } - /** Completes a previously removed message by resolving its waiter, if one exists. */ - async completeById(messageId: bigint, response?: unknown): Promise { - this.#resolveCompletionWaiter(messageId, { - status: "completed", - response, - }); - } - async #drainMessages( - nameSet: Set, + nameSet: Set | undefined, count: number, - wait: boolean, + completable: boolean, ): Promise { if (this.#metadata.size === 0) { return []; } - const now = Date.now(); const entries = await this.#loadQueueMessages(); - const matched = entries.filter( - (entry) => nameSet.has(entry.name) && !entry.inFlight, - ); + const matched = nameSet + ? entries.filter((entry) => nameSet.has(entry.name)) + : entries; if (matched.length === 0) { return []; } - const eligible = matched.filter((entry) => entry.availableAt <= now); - if (eligible.length === 0) { - this.#scheduleRedelivery(matched); - return []; - } - - const selected = eligible.slice(0, wait ? 1 : count); - if (wait) { - await this.#markMessageInFlight(selected[0], now); - return [selected[0]]; + const selected = matched.slice(0, count); + if (!completable) { + await this.#removeMessages(selected); } - - await this.#removeMessages(selected, { resolveWaiters: true }); - - // Emit trace events for received messages + const now = Date.now(); for (const message of selected) { this.#actor.emitTraceEvent("queue.message.receive", { "rivet.queue.name": message.name, @@ -488,7 +455,6 @@ export class QueueManager { "rivet.queue.latency_ms": now - message.createdAt, }); } - return selected; } @@ -506,31 +472,11 @@ export class QueueManager { value, ); const body = cbor.decode(new Uint8Array(decodedPayload.body)); - const failureCount = - decodedPayload.failureCount !== undefined && - decodedPayload.failureCount !== null - ? Number(decodedPayload.failureCount) - : 0; - const availableAt = - decodedPayload.availableAt !== undefined && - decodedPayload.availableAt !== null - ? Number(decodedPayload.availableAt) - : Number(decodedPayload.createdAt); - const inFlight = decodedPayload.inFlight ?? false; - const inFlightAt = - decodedPayload.inFlightAt !== undefined && - decodedPayload.inFlightAt !== null - ? Number(decodedPayload.inFlightAt) - : undefined; decoded.push({ id: messageId, name: decodedPayload.name, body, createdAt: Number(decodedPayload.createdAt), - failureCount, - availableAt, - inFlight, - inFlightAt, }); } catch (error) { this.#actor.rLog.error({ @@ -559,7 +505,7 @@ export class QueueManager { return; } for (const listener of [...this.#messageListeners]) { - if (!listener.nameSet.has(name)) { + if (listener.nameSet && !listener.nameSet.has(name)) { continue; } this.#removeMessageListener(listener); @@ -567,10 +513,7 @@ export class QueueManager { } } - async #removeMessages( - messages: QueueMessage[], - options: { resolveWaiters: boolean }, - ): Promise { + async #removeMessages(messages: QueueMessage[]): Promise { if (messages.length === 0) { return; } @@ -590,91 +533,24 @@ export class QueueManager { ]); this.#actor.inspector.updateQueueSize(this.#metadata.size); - - if (options.resolveWaiters) { - for (const message of messages) { - this.#resolveCompletionWaiter(message.id, { - status: "completed", - response: undefined, - }); - } - } } async #maybeResolveWaiters() { - if (this.#pendingMessageId !== undefined) { - return; - } - if (this.#redeliveryTimeout) { - clearTimeout(this.#redeliveryTimeout); - this.#redeliveryTimeout = undefined; - this.#redeliveryAt = undefined; - } - const hasReceiveWaiters = this.#waiters.size > 0; - const hasNameWaiters = this.#nameWaiters.size > 0; - if (!hasReceiveWaiters && !hasNameWaiters) { - return; - } - - if (hasNameWaiters) { - const entries = await this.#loadQueueMessages(); - const now = Date.now(); - const nameWaiters = [...this.#nameWaiters.values()]; - for (const waiter of nameWaiters) { - if (waiter.signal?.aborted) { - this.#nameWaiters.delete(waiter.id); - waiter.reject(new errors.ActorAborted()); - continue; - } - - const hasMatch = entries.some( - (message) => - waiter.nameSet.has(message.name) && - !message.inFlight && - message.availableAt <= now, - ); - if (!hasMatch) { - continue; - } - - this.#nameWaiters.delete(waiter.id); - if (waiter.abortHandler) { - waiter.signal?.removeEventListener( - "abort", - waiter.abortHandler, - ); - } - waiter.resolve(); - } - } - - if (!hasReceiveWaiters) { + if (this.#waiters.size === 0) { return; } const pending = [...this.#waiters.values()]; for (const waiter of pending) { - if (waiter.signal?.aborted) { - this.#waiters.delete(waiter.id); - waiter.reject(new errors.ActorAborted()); - continue; - } - const messages = await this.#drainMessages( waiter.nameSet, waiter.count, - waiter.wait, + waiter.completable, ); if (messages.length === 0) { continue; } this.#waiters.delete(waiter.id); - if (waiter.timeoutHandle) { - clearTimeout(waiter.timeoutHandle); - } waiter.resolve(messages); - if (waiter.wait) { - break; - } } } @@ -706,167 +582,6 @@ export class QueueManager { this.#actor.inspector.updateQueueSize(this.#metadata.size); } - async #markMessageInFlight( - message: QueueMessage, - now: number, - ): Promise { - if (message.inFlight) { - throw new errors.QueueMessagePending(); - } - - message.inFlight = true; - message.inFlightAt = now; - - await this.#persistMessage(message); - - this.#pendingMessageId = message.id; - this.#pendingWarningHandle = setTimeout(() => { - if (this.#pendingMessageId === message.id) { - this.#actor.rLog.warn({ - msg: "queue message pending for over 30s", - messageId: message.id.toString(), - name: message.name, - }); - } - }, PENDING_WARNING_MS); - } - - async #persistMessage(message: QueueMessage): Promise { - const bodyCborBuffer = cbor.encode(message.body); - const encodedMessage = - QUEUE_MESSAGE_VERSIONED.serializeWithEmbeddedVersion( - { - name: message.name, - body: new Uint8Array(bodyCborBuffer).buffer as ArrayBuffer, - createdAt: BigInt(message.createdAt), - failureCount: message.failureCount, - availableAt: BigInt(message.availableAt), - inFlight: message.inFlight, - inFlightAt: - message.inFlightAt !== undefined - ? BigInt(message.inFlightAt) - : null, - }, - ACTOR_PERSIST_CURRENT_VERSION, - ); - - await this.#driver.kvBatchPut(this.#actor.id, [ - [makeQueueMessageKey(message.id), encodedMessage], - ]); - } - - async #recoverInFlightMessages(): Promise { - const entries = await this.#driver.kvListPrefix( - this.#actor.id, - KEYS.QUEUE_PREFIX, - ); - - const updates: [Uint8Array, Uint8Array][] = []; - const now = Date.now(); - - for (const [key, value] of entries) { - try { - const messageId = decodeQueueMessageKey(key); - const decodedPayload = - QUEUE_MESSAGE_VERSIONED.deserializeWithEmbeddedVersion( - value, - ); - const inFlight = decodedPayload.inFlight ?? false; - if (!inFlight) { - continue; - } - - const failureCount = - (decodedPayload.failureCount !== undefined && - decodedPayload.failureCount !== null - ? Number(decodedPayload.failureCount) - : 0) + 1; - const availableAt = now + this.#computeBackoffMs(failureCount); - - const updatedMessage = - QUEUE_MESSAGE_VERSIONED.serializeWithEmbeddedVersion( - { - name: decodedPayload.name, - body: decodedPayload.body, - createdAt: decodedPayload.createdAt, - failureCount, - availableAt: BigInt(availableAt), - inFlight: false, - inFlightAt: null, - }, - ACTOR_PERSIST_CURRENT_VERSION, - ); - - updates.push([key, updatedMessage]); - - this.#actor.rLog.warn({ - msg: "recovering in-flight queue message", - messageId: messageId.toString(), - failureCount, - availableAt, - }); - } catch (error) { - this.#actor.rLog.error({ - msg: "failed to recover in-flight queue message", - error, - }); - } - } - - if (updates.length > 0) { - await this.#driver.kvBatchPut(this.#actor.id, updates); - } - } - - #scheduleRedelivery(messages: QueueMessage[]): void { - if (messages.length === 0) { - return; - } - const nextAvailableAt = messages.reduce((min, message) => { - return message.availableAt < min ? message.availableAt : min; - }, messages[0].availableAt); - - if ( - this.#redeliveryAt !== undefined && - this.#redeliveryAt <= nextAvailableAt - ) { - return; - } - - if (this.#redeliveryTimeout) { - clearTimeout(this.#redeliveryTimeout); - } - - const delay = Math.max(0, nextAvailableAt - Date.now()); - this.#redeliveryAt = nextAvailableAt; - this.#redeliveryTimeout = setTimeout(() => { - this.#redeliveryTimeout = undefined; - this.#redeliveryAt = undefined; - void this.#maybeResolveWaiters(); - }, delay); - } - - #resolveCompletionWaiter( - messageId: bigint, - result: QueueCompletionResult, - ): void { - const waiter = this.#completionWaiters.get(messageId); - if (!waiter) { - return; - } - this.#completionWaiters.delete(messageId); - if (waiter.timeoutHandle) { - clearTimeout(waiter.timeoutHandle); - } - waiter.resolve(result); - } - - #computeBackoffMs(failureCount: number): number { - const exp = Math.max(0, failureCount - 1); - const delay = Math.min(BACKOFF_MAX_MS, BACKOFF_INITIAL_MS * 2 ** exp); - return delay; - } - #serializeMetadata(): Uint8Array { return QUEUE_METADATA_VERSIONED.serializeWithEmbeddedVersion( { @@ -876,4 +591,5 @@ export class QueueManager { ACTOR_PERSIST_CURRENT_VERSION, ); } + } diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue.ts index 638d18846c..65eaace13f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/queue.ts @@ -1,112 +1,287 @@ -import type { AnyDatabaseProvider } from "../database"; import * as errors from "../errors"; -import type { QueueManager, QueueMessage as QueueMessageRecord } from "./queue-manager"; +import type { AnyDatabaseProvider } from "../database"; +import type { + EventSchemaConfig, + InferQueueCompleteMap, + InferSchemaMap, + QueueSchemaConfig, +} from "../schema"; +import { joinAbortSignals } from "../utils"; +import type { QueueManager, QueueMessage } from "./queue-manager"; + +export type QueueMessageOf = Omit< + QueueMessage, + "name" | "body" +> & { + name: Name; + body: Body; +}; + +export type QueueName = keyof TQueues & string; +export type QueueFilterName = + keyof TQueues extends never ? string : QueueName; + +type QueueMessageForName< + TQueues extends QueueSchemaConfig, + TName extends QueueFilterName, +> = keyof TQueues extends never + ? QueueMessage + : TName extends QueueName + ? QueueMessageOf[TName]> + : never; + +type QueueCompleteArgs = undefined extends T + ? [response?: T] + : [response: T]; + +type QueueCompleteArgsForName< + TQueues extends QueueSchemaConfig, + TName extends QueueFilterName, +> = keyof TQueues extends never + ? [response?: unknown] + : TName extends QueueName + ? [InferQueueCompleteMap[TName]] extends [never] + ? [response?: unknown] + : QueueCompleteArgs[TName]> + : [response?: unknown]; + +type QueueCompletableMessageForName< + TQueues extends QueueSchemaConfig, + TName extends QueueFilterName, +> = QueueMessageForName & { + complete( + ...args: QueueCompleteArgsForName + ): Promise; +}; + +export type QueueResultMessageForName< + TQueues extends QueueSchemaConfig, + TName extends QueueFilterName, + TCompletable extends boolean, +> = TCompletable extends true + ? QueueCompletableMessageForName + : QueueMessageForName; -/** Options for receiving messages from the queue. */ -export interface QueueReceiveOptions { +/** Options for receiving queue messages. */ +export interface QueueNextOptions< + TName extends string = string, + TCompletable extends boolean = boolean, +> { + /** Queue names to receive from. If omitted, reads from all queue names. */ + names?: readonly TName[]; /** Maximum number of messages to receive. Defaults to 1. */ count?: number; - /** Timeout in milliseconds to wait for messages. Waits indefinitely if not specified. */ + /** Timeout in milliseconds. Omit to wait indefinitely. */ timeout?: number; - /** When true, message must be manually completed. */ - wait?: boolean; + /** Optional abort signal for this receive call. */ + signal?: AbortSignal; + /** Whether to return completable messages. */ + completable?: TCompletable; } -/** Request object for receiving messages from the queue. */ -export interface QueueReceiveRequest extends QueueReceiveOptions { - /** Queue name or names to receive from. */ - name: string | string[]; +/** Options for non-blocking queue reads. */ +export interface QueueTryNextOptions< + TName extends string = string, + TCompletable extends boolean = boolean, +> { + /** Queue names to receive from. If omitted, reads from all queue names. */ + names?: readonly TName[]; + /** Maximum number of messages to receive. Defaults to 1. */ + count?: number; + /** Whether to return completable messages. */ + completable?: TCompletable; +} + +/** Options for queue async iteration. */ +export interface QueueIterOptions< + TName extends string = string, + TCompletable extends boolean = boolean, +> { + /** Queue names to receive from. If omitted, reads from all queue names. */ + names?: readonly TName[]; + /** Optional abort signal for this iterator. */ + signal?: AbortSignal; + /** Whether to return completable messages. */ + completable?: TCompletable; } /** User-facing queue interface exposed on ActorContext. */ -export class ActorQueue { - #queueManager: QueueManager; +export class ActorQueue< + S, + CP, + CS, + V, + I, + DB extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, +> { + #queueManager: QueueManager; #abortSignal: AbortSignal; + #pendingCompletableMessageIds = new Set(); constructor( - queueManager: QueueManager, + queueManager: QueueManager, abortSignal: AbortSignal, ) { this.#queueManager = queueManager; this.#abortSignal = abortSignal; } - /** Receives the next message from a single queue. Returns undefined if no message available. */ - next( - name: string, - opts?: QueueReceiveOptions, - ): Promise; - /** Receives messages from multiple queues. Returns messages matching any of the queue names. */ - next( - name: string[], - opts?: QueueReceiveOptions, - ): Promise; - /** Receives messages using a request object for full control over options. */ - next(request: QueueReceiveRequest): Promise; - async next( - nameOrRequest: string | string[] | QueueReceiveRequest, - opts: QueueReceiveOptions = {}, - ): Promise { - const request = - typeof nameOrRequest === "object" && !Array.isArray(nameOrRequest) - ? nameOrRequest - : { name: nameOrRequest }; - const mergedOptions = request === nameOrRequest ? request : opts; - const names = Array.isArray(request.name) - ? request.name - : [request.name]; - const count = mergedOptions.count ?? 1; - - const messages = await this.#queueManager.receive( - names, - count, - mergedOptions.timeout, - this.#abortSignal, - mergedOptions.wait ?? false, - ); + async next< + const TName extends QueueFilterName, + const TCompletable extends boolean = false, + >( + opts?: QueueNextOptions, + ): Promise>> { + const resolvedOpts = (opts ?? {}) as QueueNextOptions< + TName, + TCompletable + >; + const completable = resolvedOpts.completable === true; - if (Array.isArray(request.name)) { - return messages?.map((message) => - this.#toQueueMessage(message, mergedOptions.wait ?? false), - ); + if (this.#pendingCompletableMessageIds.size > 0) { + throw new errors.QueuePreviousMessageNotCompleted(); } - if (!messages || messages.length === 0) { - return undefined; + const names = this.#normalizeNames(resolvedOpts.names); + const count = Math.max(1, resolvedOpts.count ?? 1); + const { signal, cleanup } = joinAbortSignals( + this.#abortSignal, + resolvedOpts.signal, + ); + const messages = await this.#queueManager + .receive( + names, + count, + resolvedOpts.timeout, + signal, + completable, + ) + .finally(cleanup); + if (!completable) { + return messages as Array< + QueueResultMessageForName + >; } + return messages.map((message) => this.#makeCompletableMessage(message)) as unknown as Array< + QueueResultMessageForName + >; + } - return this.#toQueueMessage(messages[0], mergedOptions.wait ?? false); + async tryNext< + const TName extends QueueFilterName, + const TCompletable extends boolean = false, + >( + opts?: QueueTryNextOptions, + ): Promise>> { + const resolvedOpts = (opts ?? {}) as QueueTryNextOptions< + TName, + TCompletable + >; + if (resolvedOpts.completable === true) { + return (await this.next({ + names: resolvedOpts.names, + count: resolvedOpts.count, + timeout: 0, + completable: true, + })) as Array>; + } + return (await this.next({ + names: resolvedOpts.names, + count: resolvedOpts.count, + timeout: 0, + })) as Array>; } - #toQueueMessage( - message: QueueMessageRecord, - wait: boolean, - ): QueueMessage { - const base: QueueMessage = { - id: message.id.toString(), - name: message.name, - body: message.body, - complete: async (data?: unknown) => { - if (!wait) { - throw new errors.QueueCompleteNotAllowed(); + async *iter< + const TName extends QueueFilterName, + const TCompletable extends boolean = false, + >( + opts?: QueueIterOptions, + ): AsyncIterableIterator< + QueueResultMessageForName + > { + const resolvedOpts = (opts ?? {}) as QueueIterOptions< + TName, + TCompletable + >; + while (!this.#abortSignal.aborted) { + try { + const messages = resolvedOpts.completable === true + ? await this.next({ + names: resolvedOpts.names, + count: 1, + signal: resolvedOpts.signal, + completable: true, + }) + : await this.next({ + names: resolvedOpts.names, + count: 1, + signal: resolvedOpts.signal, + }); + if (messages.length === 0) { + continue; } - await this.#queueManager.complete(message, data); - }, - }; - - return base; + yield messages[0] as QueueResultMessageForName< + TQueues, + TName, + TCompletable + >; + } catch (error) { + if (error instanceof errors.ActorAborted) { + return; + } + throw error; + } + } } /** Sends a message to the specified queue. */ + send( + name: K, + body: InferSchemaMap[K], + ): Promise; + send( + name: keyof TQueues extends never ? string : never, + body: unknown, + ): Promise; async send(name: string, body: unknown): Promise { - const message = await this.#queueManager.enqueue(name, body); - return this.#toQueueMessage(message, false); + return await this.#queueManager.enqueue(name, body); + } + + #normalizeNames(names: readonly string[] | undefined): string[] | undefined { + if (!names || names.length === 0) { + return undefined; + } + return [...new Set(names)]; } -} -export interface QueueMessage { - name: string; - body: T; - id: string; - complete(data?: unknown): Promise; + #makeCompletableMessage( + message: QueueMessage, + ): QueueMessage & { + complete: (response?: unknown) => Promise; + } { + const messageId = message.id.toString(); + this.#pendingCompletableMessageIds.add(messageId); + + let completed = false; + const completableMessage = { + ...message, + complete: async (response?: unknown) => { + if (completed) { + throw new errors.QueueAlreadyCompleted(); + } + completed = true; + try { + await this.#queueManager.completeMessage(message, response); + this.#pendingCompletableMessageIds.delete(messageId); + } catch (error) { + completed = false; + throw error; + } + }, + }; + return completableMessage; + } } diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/schedule-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/schedule-manager.ts index 638827fcac..666b044e5d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/schedule-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/schedule-manager.ts @@ -6,6 +6,7 @@ import { } from "@/utils"; import type { AnyDatabaseProvider } from "../database"; import type { ActorDriver } from "../driver"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import type { ActorInstance } from "./mod"; import type { PersistedScheduleEvent } from "./persisted"; @@ -13,15 +14,24 @@ import type { PersistedScheduleEvent } from "./persisted"; * Manages scheduled events and alarms for actor instances. * Handles event scheduling, alarm triggers, and automatic event execution. */ -export class ScheduleManager { - #actor: ActorInstance; +export class ScheduleManager< + S, + CP, + CS, + V, + I, + DB extends AnyDatabaseProvider, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, +> { + #actor: ActorInstance; #actorDriver: ActorDriver; #alarmWriteQueue = new SinglePromiseQueue(); #config: any; // ActorConfig type #persist: any; // Reference to PersistedActor constructor( - actor: ActorInstance, + actor: ActorInstance, actorDriver: ActorDriver, config: any, ) { diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts b/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts index 38bd6fb2b9..ab7fdb025a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/instance/state-manager.ts @@ -12,6 +12,7 @@ import { type AnyConn, CONN_STATE_MANAGER_SYMBOL } from "../conn/mod"; import { convertConnToBarePersistedConn } from "../conn/persisted"; import type { ActorDriver } from "../driver"; import * as errors from "../errors"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; import { isConnStatePath, isStatePath } from "../utils"; import { KEYS, makeConnKey } from "./keys"; import type { ActorInstance } from "./mod"; @@ -36,8 +37,15 @@ export interface SaveStateOptions { * Manages actor state persistence, proxying, and synchronization. * Handles automatic state change detection and throttled persistence to KV storage. */ -export class StateManager { - #actor: ActorInstance; +export class StateManager< + S, + CP, + CS, + I, + E extends EventSchemaConfig = Record, + Q extends QueueSchemaConfig = Record, +> { + #actor: ActorInstance; #actorDriver: ActorDriver; // State tracking @@ -58,7 +66,7 @@ export class StateManager { #stateSaveInterval: number; constructor( - actor: ActorInstance, + actor: ActorInstance, actorDriver: ActorDriver, config: any, ) { diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts b/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts index d9906fbc9a..88ba4c8a27 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/mod.ts @@ -7,6 +7,8 @@ import { } from "./config"; import type { AnyDatabaseProvider } from "./database"; import { ActorDefinition } from "./definition"; +import { event as schemaEvent, queue as schemaQueue } from "./schema"; +import type { EventSchemaConfig, QueueSchemaConfig } from "./schema"; export function actor< TState, @@ -15,13 +17,26 @@ export function actor< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, TActions extends Actions< TState, TConnParams, TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues + > = Actions< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues >, >( input: ActorConfigInput< @@ -31,6 +46,8 @@ export function actor< TVars, TInput, TDatabase, + TEvents, + TQueues, TActions >, ): ActorDefinition< @@ -40,6 +57,8 @@ export function actor< TVars, TInput, TDatabase, + TEvents, + TQueues, TActions > { const config = ActorConfigSchema.parse(input) as ActorConfig< @@ -48,7 +67,9 @@ export function actor< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >; return new ActorDefinition(config); } @@ -84,3 +105,6 @@ export { createActorRouter, } from "./router"; export { routeWebSocket } from "./router-websocket-endpoints"; +export type { Type } from "./schema"; +export const event = schemaEvent; +export const queue = schemaQueue; diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/protocol/old.ts b/rivetkit-typescript/packages/rivetkit/src/actor/protocol/old.ts index 5a6654f684..85c19fcb39 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/protocol/old.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/protocol/old.ts @@ -29,6 +29,7 @@ import { import { CONN_SEND_MESSAGE_SYMBOL, type Conn } from "../conn/mod"; import { ActionContext } from "../contexts"; import type { ActorInstance } from "../instance/mod"; +import type { EventSchemaConfig, QueueSchemaConfig } from "../schema"; interface MessageEventOpts { encoding: Encoding; @@ -139,19 +140,21 @@ export interface ProcessMessageHandler< V, I, DB extends AnyDatabaseProvider, + E extends EventSchemaConfig, + Q extends QueueSchemaConfig, > { onExecuteAction?: ( - ctx: ActionContext, + ctx: ActionContext, name: string, args: unknown[], ) => Promise; onSubscribe?: ( eventName: string, - conn: Conn, + conn: Conn, ) => Promise; onUnsubscribe?: ( eventName: string, - conn: Conn, + conn: Conn, ) => Promise; } @@ -162,6 +165,8 @@ export async function processMessage< V, I, DB extends AnyDatabaseProvider, + E extends EventSchemaConfig, + Q extends QueueSchemaConfig, >( message: { body: @@ -174,9 +179,9 @@ export async function processMessage< val: { eventName: string; subscribe: boolean }; }; }, - actor: ActorInstance, - conn: Conn, - handler: ProcessMessageHandler, + actor: ActorInstance, + conn: Conn, + handler: ProcessMessageHandler, ) { let actionId: bigint | undefined; let actionName: string | undefined; @@ -199,7 +204,10 @@ export async function processMessage< actionName: name, }); - const ctx = new ActionContext(actor, conn); + const ctx = new ActionContext( + actor, + conn, + ); // Process the action request and wait for the result // This will wait for async actions to complete diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/schema.ts b/rivetkit-typescript/packages/rivetkit/src/actor/schema.ts new file mode 100644 index 0000000000..db29ca3cd9 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/actor/schema.ts @@ -0,0 +1,177 @@ +import type { StandardSchemaV1 } from "@standard-schema/spec"; +import { Unsupported } from "./errors"; + +export interface EventTypeToken { + readonly _eventType?: T; +} + +export interface QueueTypeToken { + readonly _queueMessage?: TMessage; + readonly _queueComplete?: TComplete; +} + +/** @deprecated Use `event()`. */ +export type Type = EventTypeToken; + +export function event(..._args: unknown[]): EventTypeToken { + return {} as EventTypeToken; +} + +export function queue( + ..._args: unknown[] +): QueueTypeToken< + TMessage, + TComplete +> { + return {} as QueueTypeToken; +} + +export type PrimitiveSchema = StandardSchemaV1 | EventTypeToken; + +export interface QueueSchemaDefinition { + message: PrimitiveSchema; + complete?: PrimitiveSchema; +} + +export type EventSchema = PrimitiveSchema; +export type QueueSchema = + | PrimitiveSchema + | QueueSchemaDefinition + | QueueTypeToken; +export type EventSchemaConfig = Record; +export type QueueSchemaConfig = Record; +export type AnySchemaConfig = EventSchemaConfig | QueueSchemaConfig; + +/** @deprecated Use `EventSchema` or `QueueSchema`. */ +export type Schema = QueueSchema; +/** @deprecated Use `EventSchemaConfig` or `QueueSchemaConfig`. */ +export type SchemaConfig = QueueSchemaConfig; + +export type InferSchema = + T extends QueueSchemaDefinition + ? InferSchema + : T extends QueueTypeToken + ? M + : T extends StandardSchemaV1 + ? O + : T extends EventTypeToken + ? R + : never; + +export type InferSchemaMap> = { + [K in keyof T]: InferSchema; +}; + +export type InferQueueComplete = + T extends QueueTypeToken + ? [C] extends [never] + ? never + : C + : T extends QueueSchemaDefinition + ? T["complete"] extends PrimitiveSchema + ? InferSchema + : never + : never; + +export type InferQueueCompleteMap = { + [K in keyof T]: InferQueueComplete; +}; + +export type InferEventArgs = T extends readonly unknown[] + ? number extends T["length"] + ? [T] + : T + : [T]; + +export type ValidationResult = + | { success: true; data: T } + | { success: false; issues: unknown[] }; + +export function isStandardSchema(value: unknown): value is StandardSchemaV1 { + return typeof value === "object" && value !== null && "~standard" in value; +} + +export function isQueueSchemaDefinition( + value: unknown, +): value is QueueSchemaDefinition { + return ( + typeof value === "object" && + value !== null && + "message" in value && + (value as { message?: unknown }).message !== undefined + ); +} + +function getValidationSchema( + schema: QueueSchema | EventSchema | undefined, +): QueueSchema | EventSchema | undefined { + if (!schema) { + return undefined; + } + if (isQueueSchemaDefinition(schema)) { + return schema.message; + } + return schema; +} + +function isPromiseLike(value: unknown): value is PromiseLike { + return ( + typeof value === "object" && + value !== null && + "then" in value && + typeof (value as { then?: unknown }).then === "function" + ); +} + +export async function validateSchema( + schemas: T | undefined, + key: keyof T & string, + data: unknown, +): Promise[typeof key]>> { + const schema = getValidationSchema(schemas?.[key]); + + if (!schema) { + return { success: true, data: data as InferSchemaMap[typeof key] }; + } + + if (isStandardSchema(schema)) { + const result = await schema["~standard"].validate(data); + if (result.issues) { + return { success: false, issues: [...result.issues] }; + } + return { + success: true, + data: result.value as InferSchemaMap[typeof key], + }; + } + + return { success: true, data: data as InferSchemaMap[typeof key] }; +} + +export function validateSchemaSync( + schemas: T | undefined, + key: keyof T & string, + data: unknown, +): ValidationResult[typeof key]> { + const schema = getValidationSchema(schemas?.[key]); + + if (!schema) { + return { success: true, data: data as InferSchemaMap[typeof key] }; + } + + if (isStandardSchema(schema)) { + const result = schema["~standard"].validate(data); + if (isPromiseLike(result)) { + throw new Unsupported("async schema validation"); + } + if (result.issues) { + return { success: false, issues: [...result.issues] }; + } + return { + success: true, + data: result.value as InferSchemaMap[typeof key], + }; + } + + return { success: true, data: data as InferSchemaMap[typeof key] }; +} diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/utils.ts b/rivetkit-typescript/packages/rivetkit/src/actor/utils.ts index 4b47c1c2d2..cf9d6c1765 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/utils.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/utils.ts @@ -84,6 +84,46 @@ export class Lock { } } +export interface JoinedAbortSignal { + signal?: AbortSignal; + cleanup: () => void; +} + +export function joinAbortSignals( + ...signals: Array +): JoinedAbortSignal { + const activeSignals = signals.filter( + (signal): signal is AbortSignal => signal !== undefined, + ); + if (activeSignals.length === 0) { + return { signal: undefined, cleanup: () => {} }; + } + if (activeSignals.length === 1) { + return { signal: activeSignals[0], cleanup: () => {} }; + } + + const controller = new AbortController(); + if (activeSignals.some((signal) => signal.aborted)) { + controller.abort(); + return { signal: controller.signal, cleanup: () => {} }; + } + + const cleanup = () => { + for (const signal of activeSignals) { + signal.removeEventListener("abort", onAbort); + } + }; + const onAbort = () => { + controller.abort(); + cleanup(); + }; + for (const signal of activeSignals) { + signal.addEventListener("abort", onAbort, { once: true }); + } + + return { signal: controller.signal, cleanup }; +} + export function generateSecureToken(length = 32) { const array = new Uint8Array(length); crypto.getRandomValues(array); diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts index 64284a7c0f..3a71d4cab2 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-common.ts @@ -21,7 +21,7 @@ export type ActorActionFunction< */ export type ActorDefinitionActions = // biome-ignore lint/suspicious/noExplicitAny: safe to use any here - AD extends ActorDefinition + AD extends ActorDefinition ? { [K in keyof R]: R[K] extends ( ...args: infer Args diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts index 04b3e8dbd5..73eae2fb94 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-conn.ts @@ -30,7 +30,13 @@ import { checkForSchedulingError, queryActor } from "./actor-query"; import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client"; import * as errors from "./errors"; import { logger } from "./log"; -import { createQueueProxy, createQueueSender } from "./queue"; +import { + createQueueSender, + type QueueSendNoWaitOptions, + type QueueSendOptions, + type QueueSendResult, + type QueueSendWaitOptions, +} from "./queue"; import { type WebSocketMessage as ConnMessage, messageLength, @@ -197,8 +203,22 @@ export class ActorConnRaw { this.#keepNodeAliveInterval = setInterval(() => 60_000); } - get queue() { - return createQueueProxy(this.#queueSender); + send( + name: string, + body: unknown, + options: QueueSendWaitOptions, + ): Promise; + send( + name: string, + body: unknown, + options?: QueueSendNoWaitOptions, + ): Promise; + send( + name: string, + body: unknown, + options?: QueueSendOptions, + ): Promise { + return this.#queueSender.send(name, body, options as any); } /** diff --git a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts index 9c9367bdec..7d29e8ed99 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/actor-handle.ts @@ -29,7 +29,13 @@ import { checkForSchedulingError, queryActor } from "./actor-query"; import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client"; import { ActorError, isSchedulingError } from "./errors"; import { logger } from "./log"; -import { createQueueProxy, createQueueSender } from "./queue"; +import { + createQueueSender, + type QueueSendNoWaitOptions, + type QueueSendOptions, + type QueueSendResult, + type QueueSendWaitOptions, +} from "./queue"; import { rawHttpFetch, rawWebSocket } from "./raw-utils"; import { sendHttpRequest } from "./utils"; @@ -80,8 +86,22 @@ export class ActorHandleRaw { }); } - get queue() { - return createQueueProxy(this.#queueSender); + send( + name: string, + body: unknown, + options: QueueSendWaitOptions, + ): Promise; + send( + name: string, + body: unknown, + options?: QueueSendNoWaitOptions, + ): Promise; + send( + name: string, + body: unknown, + options?: QueueSendOptions, + ): Promise { + return this.#queueSender.send(name, body, options as any); } /** diff --git a/rivetkit-typescript/packages/rivetkit/src/client/client.ts b/rivetkit-typescript/packages/rivetkit/src/client/client.ts index df298bc484..f4efbe1f46 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/client.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/client.ts @@ -497,14 +497,10 @@ function createActorProxy( return value; } - if (prop === "queue") { - return Reflect.get(target, prop, receiver); - } - // Create action function that preserves 'this' context if (typeof prop === "string") { // If JS is attempting to calling this as a promise, ignore it - if (prop === "then" || prop === "queue") return undefined; + if (prop === "then") return undefined; let method = methodCache.get(prop); if (!method) { @@ -518,10 +514,8 @@ function createActorProxy( // Support for 'in' operator has(target: ActorHandleRaw, prop: string | symbol) { - // All string properties are potentially action functions - if (typeof prop === "string") { - return prop !== "queue"; - } + // All string properties are potentially action functions. + if (typeof prop === "string") return true; // For symbols, defer to the target's own has behavior return Reflect.has(target, prop); }, @@ -549,9 +543,6 @@ function createActorProxy( return targetDescriptor; } if (typeof prop === "string") { - if (prop === "queue") { - return undefined; - } // Make action methods appear non-enumerable return { configurable: true, diff --git a/rivetkit-typescript/packages/rivetkit/src/client/queue.ts b/rivetkit-typescript/packages/rivetkit/src/client/queue.ts index 2fbbde9eec..b848228c5f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/client/queue.ts +++ b/rivetkit-typescript/packages/rivetkit/src/client/queue.ts @@ -20,27 +20,29 @@ export interface QueueSender { send( name: string, body: unknown, - options?: QueueSendOptions | AbortSignal, - ): Promise; -} - -export interface QueueNameSender { + options: QueueSendWaitOptions, + ): Promise; send( + name: string, body: unknown, - options?: QueueSendOptions | AbortSignal, - ): Promise; + options?: QueueSendNoWaitOptions, + ): Promise; } -export type QueueProxy = QueueSender & { - [key: string]: QueueNameSender; -}; - -export interface QueueSendOptions { - wait?: boolean; +export interface QueueSendWaitOptions { + wait: true; timeout?: number; signal?: AbortSignal; } +export interface QueueSendNoWaitOptions { + wait?: false; + timeout?: never; + signal?: AbortSignal; +} + +export type QueueSendOptions = QueueSendWaitOptions | QueueSendNoWaitOptions; + export interface QueueSendResult { status: "completed" | "timedOut"; response?: unknown; @@ -53,123 +55,92 @@ interface QueueSenderOptions { } export function createQueueSender(senderOptions: QueueSenderOptions): QueueSender { - return { - async send( - name: string, - body: unknown, - options?: QueueSendOptions | AbortSignal, - ): Promise { - const normalizedOptions = - options instanceof AbortSignal ? { signal: options } : options; - const wait = normalizedOptions?.wait ?? false; - const timeout = normalizedOptions?.timeout; - - const result = await sendHttpRequest< - protocol.HttpQueueSendRequest, - protocol.HttpQueueSendResponse, - HttpQueueSendRequestJson, - HttpQueueSendResponseJson, - { body: unknown; wait?: boolean; timeout?: number; name?: string }, - QueueSendResult - >({ - url: `http://actor/queue/${encodeURIComponent(name)}`, - method: "POST", - headers: { - [HEADER_ENCODING]: senderOptions.encoding, - ...(senderOptions.params !== undefined - ? { - [HEADER_CONN_PARAMS]: JSON.stringify( - senderOptions.params, - ), - } - : {}), - }, - body: { body, wait, timeout }, - encoding: senderOptions.encoding, - customFetch: senderOptions.customFetch, - signal: normalizedOptions?.signal, - requestVersion: CLIENT_PROTOCOL_CURRENT_VERSION, - requestVersionedDataHandler: HTTP_QUEUE_SEND_REQUEST_VERSIONED, - responseVersion: CLIENT_PROTOCOL_CURRENT_VERSION, - responseVersionedDataHandler: - HTTP_QUEUE_SEND_RESPONSE_VERSIONED, - requestZodSchema: HttpQueueSendRequestSchema, - responseZodSchema: HttpQueueSendResponseSchema, - requestToJson: (value): HttpQueueSendRequestJson => ({ - ...value, - name, - }), - requestToBare: (value): protocol.HttpQueueSendRequest => ({ - name: value.name ?? name, - body: bufferToArrayBuffer(cbor.encode(value.body)), - wait: value.wait ?? false, - timeout: value.timeout !== undefined ? BigInt(value.timeout) : null, - }), - responseFromJson: (json): QueueSendResult => { - if (json.response === undefined) { - return { status: json.status as "completed" | "timedOut" }; - } - return { - status: json.status as "completed" | "timedOut", - response: json.response, - }; - }, - responseFromBare: (bare): QueueSendResult => { - if (bare.response === null || bare.response === undefined) { - return { status: bare.status as "completed" | "timedOut" }; - } - return { - status: bare.status as "completed" | "timedOut", - response: cbor.decode(new Uint8Array(bare.response)), - }; - }, - }); - - if (wait) { - return result; - } - return; - }, - }; -} - -export function createQueueProxy(sender: QueueSender): QueueProxy { - const methodCache = new Map(); - return new Proxy(sender, { - get(target, prop: string | symbol, receiver: unknown) { - if (typeof prop === "symbol") { - return Reflect.get(target, prop, receiver); - } + async function send( + name: string, + body: unknown, + options: QueueSendWaitOptions, + ): Promise; + async function send( + name: string, + body: unknown, + options?: QueueSendNoWaitOptions, + ): Promise; + async function send( + name: string, + body: unknown, + options?: QueueSendOptions, + ): Promise { + const wait = options?.wait ?? false; + const timeout = options?.timeout; - if (prop in target) { - const value = Reflect.get(target, prop, target); - if (typeof value === "function") { - return value.bind(target); + const result = await sendHttpRequest< + protocol.HttpQueueSendRequest, + protocol.HttpQueueSendResponse, + HttpQueueSendRequestJson, + HttpQueueSendResponseJson, + { body: unknown; wait?: boolean; timeout?: number; name?: string }, + QueueSendResult + >({ + url: `http://actor/queue/${encodeURIComponent(name)}`, + method: "POST", + headers: { + [HEADER_ENCODING]: senderOptions.encoding, + ...(senderOptions.params !== undefined + ? { + [HEADER_CONN_PARAMS]: JSON.stringify( + senderOptions.params, + ), + } + : {}), + }, + body: { body, wait, timeout }, + encoding: senderOptions.encoding, + customFetch: senderOptions.customFetch, + signal: options?.signal, + requestVersion: CLIENT_PROTOCOL_CURRENT_VERSION, + requestVersionedDataHandler: HTTP_QUEUE_SEND_REQUEST_VERSIONED, + responseVersion: CLIENT_PROTOCOL_CURRENT_VERSION, + responseVersionedDataHandler: + HTTP_QUEUE_SEND_RESPONSE_VERSIONED, + requestZodSchema: HttpQueueSendRequestSchema, + responseZodSchema: HttpQueueSendResponseSchema, + requestToJson: (value): HttpQueueSendRequestJson => ({ + ...value, + name, + }), + requestToBare: (value): protocol.HttpQueueSendRequest => ({ + name: value.name ?? name, + body: bufferToArrayBuffer(cbor.encode(value.body)), + wait: value.wait ?? false, + timeout: value.timeout !== undefined ? BigInt(value.timeout) : null, + }), + responseFromJson: (json): QueueSendResult => { + if (json.response === undefined) { + return { status: json.status as "completed" | "timedOut" }; + } + return { + status: json.status as "completed" | "timedOut", + response: json.response, + }; + }, + responseFromBare: (bare): QueueSendResult => { + if (bare.response === null || bare.response === undefined) { + return { status: bare.status as "completed" | "timedOut" }; } - return value; - } + return { + status: bare.status as "completed" | "timedOut", + response: cbor.decode(new Uint8Array(bare.response)), + }; + }, + }); - if (prop === "then") return undefined; + if (wait) { + return result; + } + return; + } - if (typeof prop === "string") { - let method = methodCache.get(prop); - if (!method) { - method = { - send: ( - body: unknown, - options?: QueueSendOptions | AbortSignal, - ) => target.send(prop, body, options), - }; - methodCache.set(prop, method); - } - return method; - } - }, - has(target, prop: string | symbol) { - if (typeof prop === "string") { - return true; - } - return Reflect.has(target, prop); - }, - }) as QueueProxy; + return { + send, + }; } diff --git a/rivetkit-typescript/packages/rivetkit/src/db/config.ts b/rivetkit-typescript/packages/rivetkit/src/db/config.ts index e94c9b4dbe..50d1e7682e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/db/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/db/config.ts @@ -66,7 +66,10 @@ export type DatabaseProvider = { * Raw database client with basic exec method */ export interface RawDatabaseClient { - exec: (query: string, ...args: unknown[]) => Promise | unknown[]; + exec: = Record>( + query: string, + ...args: unknown[] + ) => Promise | TRow[]; } /** @@ -77,10 +80,12 @@ export interface DrizzleDatabaseClient { // For now, just a marker interface } -type ExecuteFunction = ( +type ExecuteFunction = < + TRow extends Record = Record, +>( query: string, ...args: unknown[] -) => Promise; +) => Promise; export type RawAccess = { /** diff --git a/rivetkit-typescript/packages/rivetkit/src/db/drizzle/mod.ts b/rivetkit-typescript/packages/rivetkit/src/db/drizzle/mod.ts index 2b3ec95b5c..a7875071f1 100644 --- a/rivetkit-typescript/packages/rivetkit/src/db/drizzle/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/db/drizzle/mod.ts @@ -163,9 +163,36 @@ export function db< const client = proxyDrizzle(callback, config); return Object.assign(client, { - execute: async (query: string, ...args: unknown[]) => { - const result = await callback(query, args, "all"); - return result.rows; + execute: async < + TRow extends Record = Record, + >( + query: string, + ...args: unknown[] + ): Promise => { + if (args.length > 0) { + const { rows, columns } = await waDb.query(query, args); + return rows.map((row: unknown[]) => { + const rowObj: Record = {}; + for (let i = 0; i < row.length; i++) { + rowObj[columns[i]] = row[i]; + } + return rowObj; + }) as TRow[]; + } + + const results: Record[] = []; + let columnNames: string[] | null = null; + await waDb.exec(query, (row: unknown[], columns: string[]) => { + if (!columnNames) { + columnNames = columns; + } + const rowObj: Record = {}; + for (let i = 0; i < row.length; i++) { + rowObj[columnNames[i]] = row[i]; + } + results.push(rowObj); + }); + return results as TRow[]; }, close: async () => { await waDb.close(); diff --git a/rivetkit-typescript/packages/rivetkit/src/db/mod.ts b/rivetkit-typescript/packages/rivetkit/src/db/mod.ts index 2f586b5b0c..1cd28de725 100644 --- a/rivetkit-typescript/packages/rivetkit/src/db/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/db/mod.ts @@ -46,8 +46,13 @@ export function db({ if (override) { // Use the override return { - execute: async (query, ...args) => { - return override.exec(query, ...args); + execute: async < + TRow extends Record = Record, + >( + query: string, + ...args: unknown[] + ): Promise => { + return await override.exec(query, ...args); }, close: async () => { // Override clients don't need cleanup @@ -64,7 +69,12 @@ export function db({ const db = await ctx.sqliteVfs.open(ctx.actorId, kvStore); return { - execute: async (query, ...args) => { + execute: async < + TRow extends Record = Record, + >( + query: string, + ...args: unknown[] + ): Promise => { if (args.length > 0) { // Use parameterized query when args are provided const { rows, columns } = await db.query(query, args); @@ -74,7 +84,7 @@ export function db({ rowObj[columns[i]] = row[i]; } return rowObj; - }); + }) as TRow[]; } // Use exec for non-parameterized queries @@ -90,7 +100,7 @@ export function db({ } results.push(rowObj); }); - return results; + return results as TRow[]; }, close: async () => { await db.close(); diff --git a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-kv.ts b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-kv.ts index 4c05cb55b7..31b97a2055 100644 --- a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-kv.ts +++ b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-kv.ts @@ -3,10 +3,25 @@ import { setupDriverTest } from "../utils"; import { describe, expect, test, type TestContext } from "vitest"; export function runActorKvTests(driverTestConfig: DriverTestConfig) { - describe("Actor KV Tests", () => { - test("supports text encoding and decoding", async (c: TestContext) => { - const { client } = await setupDriverTest(c, driverTestConfig); - const kvHandle = client.kvActor.getOrCreate(["kv-text"]); + type KvTextHandle = { + putText: (key: string, value: string) => Promise; + getText: (key: string) => Promise; + listText: (prefix: string) => Promise>; + }; + + type KvArrayBufferHandle = { + roundtripArrayBuffer: (key: string, bytes: number[]) => Promise; + }; + + describe("Actor KV Tests", () => { + test("supports text encoding and decoding", async (c: TestContext) => { + const { client: rawClient } = await setupDriverTest( + c, + driverTestConfig, + ); + const client = rawClient as any; + const kvHandle = + client.kvActor.getOrCreate(["kv-text"]) as unknown as KvTextHandle; await kvHandle.putText("greeting", "hello"); const value = await kvHandle.getText("greeting"); @@ -26,19 +41,25 @@ export function runActorKvTests(driverTestConfig: DriverTestConfig) { test( "supports arrayBuffer encoding and decoding", async (c: TestContext) => { - const { client } = await setupDriverTest(c, driverTestConfig); - const kvHandle = client.kvActor.getOrCreate(["kv-array-buffer"]); - - const values = await kvHandle.roundtripArrayBuffer("bytes", [ - 4, - 8, - 15, - 16, - 23, - 42, - ]); - expect(values).toEqual([4, 8, 15, 16, 23, 42]); - }, - ); - }); -} + const { client: rawClient } = await setupDriverTest( + c, + driverTestConfig, + ); + const client = rawClient as any; + const kvHandle = client.kvActor.getOrCreate([ + "kv-array-buffer", + ]) as unknown as KvArrayBufferHandle; + + const values = await kvHandle.roundtripArrayBuffer("bytes", [ + 4, + 8, + 15, + 16, + 23, + 42, + ]); + expect(values).toEqual([4, 8, 15, 16, 23, 42]); + }, + ); + }); + } diff --git a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-queue.ts b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-queue.ts index 8bce3bbd3d..4db8bdafc9 100644 --- a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-queue.ts +++ b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-queue.ts @@ -9,7 +9,7 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate(["client-send"]); - await handle.queue.greeting.send({ hello: "world" }); + await handle.send("greeting", { hello: "world" }); const message = await handle.receiveOne("greeting"); expect(message).toEqual({ @@ -32,9 +32,9 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate(["receive-array"]); - await handle.queue.a.send(1); - await handle.queue.b.send(2); - await handle.queue.c.send(3); + await handle.send("a", 1); + await handle.send("b", 2); + await handle.send("c", 3); const messages = await handle.receiveMany(["a", "b"], { count: 2 }); expect(messages).toEqual([ @@ -47,11 +47,11 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate(["receive-request"]); - await handle.queue.one.send("first"); - await handle.queue.two.send("second"); + await handle.send("one", "first"); + await handle.send("two", "second"); const messages = await handle.receiveRequest({ - name: ["one", "two"], + names: ["one", "two"], count: 2, }); expect(messages).toEqual([ @@ -60,6 +60,22 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { ]); }); + test("next defaults to all names when names is omitted", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate([ + "receive-request-all", + ]); + + await handle.send("one", "first"); + await handle.send("two", "second"); + + const messages = await handle.receiveRequest({ count: 2 }); + expect(messages).toEqual([ + { name: "one", body: "first" }, + { name: "two", body: "second" }, + ]); + }); + test("next timeout returns empty array", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate(["receive-timeout"]); @@ -70,6 +86,17 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { expect(messages).toEqual([]); }); + test("tryNext does not wait and returns empty array", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate(["try-next-empty"]); + + const messages = await handle.tryReceiveMany({ + names: ["missing"], + count: 1, + }); + expect(messages).toEqual([]); + }); + test("abort throws ActorAborted", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate(["abort-test"]); @@ -83,17 +110,49 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { } }); + test("next supports signal abort", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate(["signal-abort-next"]); + + const result = await handle.waitForSignalAbort(); + expect(result).toEqual({ + group: "actor", + code: "aborted", + }); + }); + + test("next supports actor abort when signal is provided", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate([ + "actor-abort-with-signal-next", + ]); + + const result = await handle.waitForActorAbortWithSignal(); + expect(result).toEqual({ + group: "actor", + code: "aborted", + }); + }); + + test("iter supports signal abort", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate(["signal-abort-iter"]); + + const result = await handle.iterWithSignalAbort(); + expect(result).toEqual({ ok: true }); + }); + test("enforces queue size limit", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); const key = `size-limit-${Date.now()}-${Math.random().toString(16).slice(2)}`; const handle = client.queueLimitedActor.getOrCreate([key]); - await handle.queue.message.send(1); + await handle.send("message", 1); await waitFor(driverTestConfig, 10); try { - await handle.queue.message.send(2); + await handle.send("message", 2); expect.fail("expected queue full error"); } catch (error) { expect(error).toBeInstanceOf(Error); @@ -115,7 +174,7 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { const largePayload = "a".repeat(200); try { - await handle.queue.oversize.send(largePayload); + await handle.send("oversize", largePayload); expect.fail("expected message_too_large error"); } catch (error) { expect((error as ActorError).group).toBe("queue"); @@ -128,7 +187,7 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { const handle = client.queueActor.getOrCreate(["wait-complete"]); const actionPromise = handle.receiveAndComplete("tasks"); - const result = await handle.queue.tasks.send( + const result = await handle.send("tasks", { value: 123 }, { wait: true, timeout: 1_000 }, ); @@ -144,7 +203,7 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate(["wait-timeout"]); - const resultPromise = handle.queue.timeout.send( + const resultPromise = handle.send("timeout", { value: 456 }, { wait: true, timeout: 50 }, ); @@ -152,31 +211,74 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { await waitFor(driverTestConfig, 60); const result = await resultPromise; - expect(result?.status).toBe("timedOut"); + expect(result.status).toBe("timedOut"); }); - test("complete throws when wait is false", async (c) => { + test("manual receive retries message when not completed", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate([ - "complete-not-allowed", + "manual-retry-uncompleted", ]); - await handle.queue.nowait.send({ value: "test" }); - const result = await handle.receiveWithoutWaitComplete("nowait"); + await handle.send("tasks", { value: 789 }); + const first = await handle.receiveWithoutComplete("tasks"); + expect(first).toEqual({ name: "tasks", body: { value: 789 } }); + + const retried = await handle.receiveOne("tasks", { timeout: 1_000 }); + expect(retried).toEqual({ name: "tasks", body: { value: 789 } }); + }); + + test("next throws when previous manual message is not completed", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate([ + "manual-next-requires-complete", + ]); + await handle.send("tasks", { value: 111 }); + const result = await handle.receiveManualThenNextWithoutComplete( + "tasks", + ); expect(result).toEqual({ group: "queue", - code: "complete_not_allowed", + code: "previous_message_not_completed", }); }); + test("manual receive includes complete even without completion schema", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate([ + "complete-not-allowed", + ]); + + await handle.send("nowait", { value: "test" }); + const result = await handle.receiveWithoutCompleteMethod("nowait"); + + expect(result).toEqual({ + hasComplete: true, + }); + }); + + test("manual receive retries queues without completion schema until completed", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate([ + "complete-not-allowed-consume", + ]); + + await handle.send("nowait", { value: "test" }); + const result = await handle.receiveWithoutCompleteMethod("nowait"); + expect(result).toEqual({ hasComplete: true }); + + const next = await handle.receiveOne("nowait", { timeout: 1_000 }); + expect(next).toEqual({ name: "nowait", body: { value: "test" } }); + }); + test("complete throws when called twice", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); const handle = client.queueActor.getOrCreate([ "complete-twice", ]); - await handle.queue.twice.send({ value: "test" }); + await handle.send("twice", { value: "test" }); const result = await handle.receiveAndCompleteTwice("twice"); expect(result).toEqual({ @@ -185,17 +287,38 @@ export function runActorQueueTests(driverTestConfig: DriverTestConfig) { }); }); - test("next throws when message pending", async (c) => { + test("wait send no longer requires queue completion schema", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); - const handle = client.queueActor.getOrCreate(["pending-next"]); + const handle = client.queueActor.getOrCreate([ + "missing-completion-schema", + ]); - await handle.queue.pending.send({ value: "test" }); - const result = await handle.receiveWhilePending("pending"); + const result = await handle.send( + "nowait", + { value: "test" }, + { wait: true, timeout: 50 }, + ); + expect(result).toEqual({ status: "timedOut" }); + }); - expect(result).toEqual({ - group: "queue", - code: "message_pending", - }); + test("iter can consume queued messages", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate(["iter-consume"]); + + await handle.send("one", "first"); + const message = await handle.receiveWithIterator("one"); + expect(message).toEqual({ name: "one", body: "first" }); + }); + + test("queue async iterator can consume queued messages", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + const handle = client.queueActor.getOrCreate([ + "async-iter-consume", + ]); + + await handle.send("two", "second"); + const message = await handle.receiveWithAsyncIterator(); + expect(message).toEqual({ name: "two", body: "second" }); }); }); } diff --git a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-run.ts b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-run.ts index d634a6137e..5665106d05 100644 --- a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-run.ts +++ b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-run.ts @@ -38,26 +38,28 @@ export function runActorRunTests(driverTestConfig: DriverTestConfig) { expect(state2.tickCount).toBeGreaterThan(count1); }); - test("run handler exits gracefully on actor stop", async (c) => { - const { client } = await setupDriverTest(c, driverTestConfig); + test("active run handler keeps actor awake past sleep timeout", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); - const actor = client.runWithTicks.getOrCreate([ - "run-graceful-exit", - ]); + const actor = client.runWithTicks.getOrCreate([ + "run-stays-awake", + ]); - // Wait for run to start - await waitFor(driverTestConfig, 100); + // Wait for run to start + await waitFor(driverTestConfig, 100); - const state1 = await actor.getState(); - expect(state1.runStarted).toBe(true); + const state1 = await actor.getState(); + expect(state1.runStarted).toBe(true); + const tickCount1 = state1.tickCount; - // Wait for sleep timeout to trigger sleep - await waitFor(driverTestConfig, RUN_SLEEP_TIMEOUT + 300); + // Active run loops should keep the actor awake. + await waitFor(driverTestConfig, RUN_SLEEP_TIMEOUT + 300); - // Wake actor again and check state persisted - const state2 = await actor.getState(); - expect(state2.runExited).toBe(true); - }); + const state2 = await actor.getState(); + expect(state2.runStarted).toBe(true); + expect(state2.runExited).toBe(false); + expect(state2.tickCount).toBeGreaterThan(tickCount1); + }); test("actor without run handler works normally", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); @@ -76,8 +78,8 @@ export function runActorRunTests(driverTestConfig: DriverTestConfig) { expect(state2.wakeCount).toBe(2); }); - test("run handler can consume from queue", async (c) => { - const { client } = await setupDriverTest(c, driverTestConfig); + test("run handler can consume from queue", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); const actor = client.runWithQueueConsumer.getOrCreate([ "queue-consumer", @@ -105,11 +107,28 @@ export function runActorRunTests(driverTestConfig: DriverTestConfig) { type: "test", value: 2, }); - expect(state.messagesReceived[2].body).toEqual({ - type: "test", - value: 3, + expect(state.messagesReceived[2].body).toEqual({ + type: "test", + value: 3, + }); + }); + + test("queue-waiting run handler can sleep and resume", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + + const actor = client.runWithQueueConsumer.getOrCreate([ + "queue-consumer-sleep", + ]); + + await waitFor(driverTestConfig, 100); + const state1 = await actor.getState(); + expect(state1.runStarted).toBe(true); + + await waitFor(driverTestConfig, RUN_SLEEP_TIMEOUT + 500); + const state2 = await actor.getState(); + + expect(state2.wakeCount).toBeGreaterThan(state1.wakeCount); }); - }); test("run handler that exits early triggers destroy", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); diff --git a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-workflow.ts b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-workflow.ts index a572b1db11..37a8472e6e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-workflow.ts +++ b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-workflow.ts @@ -1,7 +1,6 @@ import { describe, expect, test } from "vitest"; import { WORKFLOW_QUEUE_NAME, - workflowQueueName, } from "../../../fixtures/driver-test-suite/workflow"; import type { DriverTestConfig } from "../mod"; import { setupDriverTest, waitFor } from "../utils"; @@ -21,19 +20,20 @@ export function runActorWorkflowTests(driverTestConfig: DriverTestConfig) { expect(state.guardTriggered).toBe(true); }); - test("consumes queue messages via workflow listen", async (c) => { + test("consumes queue messages via workflow queue.next", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); const actor = client.workflowQueueActor.getOrCreate(["workflow-queue"]); - const queueHandle = actor.queue[workflowQueueName(WORKFLOW_QUEUE_NAME)]; - await queueHandle.send({ hello: "world" }); + await actor.send(WORKFLOW_QUEUE_NAME, { + hello: "world", + }); await waitFor(driverTestConfig, 200); const messages = await actor.getMessages(); expect(messages).toEqual([{ hello: "world" }]); }); - test("workflow listen supports completing wait sends", async (c) => { + test("workflow queue.next supports completing wait sends", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig); const actor = client.workflowQueueActor.getOrCreate([ "workflow-queue-wait", diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/default.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/default.ts index 574f091bdf..c12b294678 100644 --- a/rivetkit-typescript/packages/rivetkit/src/drivers/default.ts +++ b/rivetkit-typescript/packages/rivetkit/src/drivers/default.ts @@ -28,6 +28,11 @@ export function chooseDefaultDriver( return createEngineDriver(); } - loggerWithoutContext().debug({ msg: "using default file system driver" }); - return createFileSystemOrMemoryDriver(true); + loggerWithoutContext().debug({ + msg: "using default file system driver", + storagePath: config.storagePath, + }); + return createFileSystemOrMemoryDriver(true, { + path: config.storagePath, + }); } diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts index f010a5f6ac..667d60d6a1 100644 --- a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts +++ b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts @@ -14,8 +14,6 @@ import { CURRENT_VERSION as FILE_SYSTEM_DRIVER_CURRENT_VERSION, } from "@/schemas/file-system-driver/versioned"; import { - arrayBuffersEqual, - bufferToArrayBuffer, type LongTimeoutHandle, promiseWithResolvers, setLongTimeout, @@ -33,6 +31,13 @@ import { ensureDirectoryExistsSync, getStoragePath, } from "./utils"; +import { + computePrefixUpperBound, + ensureUint8Array, + loadSqliteRuntime, + type SqliteRuntime, + type SqliteRuntimeDatabase, +} from "./sqlite-runtime"; // Actor handler to track running instances @@ -78,6 +83,8 @@ export interface FileSystemDriverOptions { persist?: boolean; /** Custom path for storage */ customPath?: string; + /** Deprecated option retained for explicit migration to sqlite-only KV. */ + useNativeSqlite?: boolean; } /** @@ -90,6 +97,8 @@ export class FileSystemGlobalState { #alarmsDir: string; #persist: boolean; + #sqliteRuntime: SqliteRuntime; + #actorKvDatabases = new Map(); /** SQLite VFS instance for this driver. */ readonly sqliteVfs = new SqliteVfs(); @@ -120,8 +129,14 @@ export class FileSystemGlobalState { } constructor(options: FileSystemDriverOptions = {}) { - const { persist = true, customPath } = options; + const { persist = true, customPath, useNativeSqlite = true } = options; + if (!useNativeSqlite) { + throw new Error( + "File-system driver no longer supports non-SQLite KV storage.", + ); + } this.#persist = persist; + this.#sqliteRuntime = loadSqliteRuntime(); this.#storagePath = persist ? (customPath ?? getStoragePath()) : "/tmp"; const path = getNodePath(); this.#stateDir = path.join(this.#storagePath, "state"); @@ -146,6 +161,7 @@ export class FileSystemGlobalState { msg: "file system driver ready", dir: this.#storagePath, actorCount: this.#actorCountOnStartup, + sqliteRuntime: this.#sqliteRuntime.kind, }); // Cleanup stale temp files on startup @@ -157,8 +173,21 @@ export class FileSystemGlobalState { error: err, }); } + + try { + this.#migrateLegacyKvToSqliteOnStartupSync(); + } catch (error) { + logger().error({ + msg: "failed legacy kv startup migration", + error, + }); + throw error; + } } else { - logger().debug({ msg: "memory driver ready" }); + logger().debug({ + msg: "memory driver ready", + sqliteRuntime: this.#sqliteRuntime.kind, + }); } } @@ -174,6 +203,154 @@ export class FileSystemGlobalState { return getNodePath().join(this.#alarmsDir, actorId); } + #getActorKvDatabasePath(actorId: string): string { + if (this.#persist) { + return this.getActorDbPath(actorId); + } + return ":memory:"; + } + + #ensureActorKvTables(db: SqliteRuntimeDatabase): void { + db.exec(` + CREATE TABLE IF NOT EXISTS kv ( + key BLOB PRIMARY KEY NOT NULL, + value BLOB NOT NULL + ) + `); + } + + #getOrCreateActorKvDatabase(actorId: string): SqliteRuntimeDatabase { + const existing = this.#actorKvDatabases.get(actorId); + if (existing) { + return existing; + } + + const dbPath = this.#getActorKvDatabasePath(actorId); + if (this.#persist) { + const path = getNodePath(); + ensureDirectoryExistsSync(path.dirname(dbPath)); + } + + let db: SqliteRuntimeDatabase; + try { + db = this.#sqliteRuntime.open(dbPath); + } catch (error) { + throw new Error( + `failed to open actor kv database for actor ${actorId} at ${dbPath}: ${error}`, + ); + } + + this.#ensureActorKvTables(db); + this.#actorKvDatabases.set(actorId, db); + return db; + } + + #closeActorKvDatabase(actorId: string): void { + const db = this.#actorKvDatabases.get(actorId); + if (!db) { + return; + } + + try { + db.close(); + } finally { + this.#actorKvDatabases.delete(actorId); + } + } + + #putKvEntriesInDb( + db: SqliteRuntimeDatabase, + entries: [Uint8Array, Uint8Array][], + ): void { + if (entries.length === 0) { + return; + } + + db.exec("BEGIN"); + try { + for (const [key, value] of entries) { + db.run("INSERT OR REPLACE INTO kv (key, value) VALUES (?, ?)", [ + key, + value, + ]); + } + db.exec("COMMIT"); + } catch (error) { + try { + db.exec("ROLLBACK"); + } catch { + // Ignore rollback errors, original error is more actionable. + } + throw error; + } + } + + #isKvDbPopulated(db: SqliteRuntimeDatabase): boolean { + const row = db.get<{ count: number | bigint }>( + "SELECT COUNT(*) AS count FROM kv", + ); + const count = row ? Number(row.count) : 0; + return count > 0; + } + + #migrateLegacyKvToSqliteOnStartupSync(): void { + const fsSync = getNodeFsSync(); + if (!fsSync.existsSync(this.#stateDir)) { + return; + } + + const actorIds = fsSync + .readdirSync(this.#stateDir) + .filter((id) => !id.includes(".tmp.")); + + for (const actorId of actorIds) { + const statePath = this.getActorStatePath(actorId); + let state: schema.ActorState; + try { + const stateBytes = fsSync.readFileSync(statePath); + state = ACTOR_STATE_VERSIONED.deserializeWithEmbeddedVersion( + new Uint8Array(stateBytes), + ); + } catch (error) { + logger().warn({ + msg: "failed to parse actor state during startup migration", + actorId, + error, + }); + continue; + } + + if (!state.kvStorage || state.kvStorage.length === 0) { + continue; + } + + const dbPath = this.getActorDbPath(actorId); + const path = getNodePath(); + ensureDirectoryExistsSync(path.dirname(dbPath)); + const db = this.#sqliteRuntime.open(dbPath); + try { + this.#ensureActorKvTables(db); + if (this.#isKvDbPopulated(db)) { + continue; + } + + const legacyEntries = state.kvStorage.map((entry) => [ + new Uint8Array(entry.key), + new Uint8Array(entry.value), + ]) as [Uint8Array, Uint8Array][]; + this.#putKvEntriesInDb(db, legacyEntries); + + logger().info({ + msg: "migrated legacy actor kv storage to sqlite", + actorId, + entryCount: legacyEntries.length, + }); + } finally { + db.close(); + } + } + } + async *getActorsIterator(params: { cursor?: string; }): AsyncGenerator { @@ -260,15 +437,8 @@ export class FileSystemGlobalState { entry.generation = crypto.randomUUID(); } - // Initialize storage - const kvStorage: schema.ActorKvEntry[] = []; + // Initialize storage (runtime KV is stored in SQLite; state.kvStorage is legacy-only) const initialKvState = getInitialActorKvState(input); - for (const [key, value] of initialKvState) { - kvStorage.push({ - key: bufferToArrayBuffer(key), - value: bufferToArrayBuffer(value), - }); - } // Initialize metadata await this.#withActorWrite(actorId, async (lockedEntry) => { @@ -277,7 +447,7 @@ export class FileSystemGlobalState { name, key, createdAt: BigInt(Date.now()), - kvStorage, + kvStorage: [], startTs: null, connectableTs: null, sleepTs: null, @@ -291,6 +461,10 @@ export class FileSystemGlobalState { lockedEntry.state, ); } + if (initialKvState.length > 0) { + const db = this.#getOrCreateActorKvDatabase(actorId); + this.#putKvEntriesInDb(db, initialKvState); + } }); return entry; @@ -336,10 +510,16 @@ export class FileSystemGlobalState { const fs = getNodeFs(); const stateData = await fs.readFile(stateFilePath); - // Cache the loaded state in handler - entry.state = ACTOR_STATE_VERSIONED.deserializeWithEmbeddedVersion( - new Uint8Array(stateData), - ); + const loadedState = + ACTOR_STATE_VERSIONED.deserializeWithEmbeddedVersion( + new Uint8Array(stateData), + ); + + // Runtime reads/writes are SQLite-only; legacy kvStorage is for one-time startup migration. + entry.state = { + ...loadedState, + kvStorage: [], + }; return entry; } catch (innerError: any) { @@ -381,39 +561,36 @@ export class FileSystemGlobalState { entry.generation = crypto.randomUUID(); } - // Initialize kvStorage with the initial persist data - const kvStorage: schema.ActorKvEntry[] = []; - const initialKvState = getInitialActorKvState(input); - for (const [key, value] of initialKvState) { - kvStorage.push({ - key: bufferToArrayBuffer(key), - value: bufferToArrayBuffer(value), - }); - } + // Initialize storage (runtime KV is stored in SQLite; state.kvStorage is legacy-only) + const initialKvState = getInitialActorKvState(input); - await this.#withActorWrite(actorId, async (lockedEntry) => { - lockedEntry.state = { - actorId, - name, - key: key as readonly string[], - createdAt: BigInt(Date.now()), - kvStorage, - startTs: null, - connectableTs: null, - sleepTs: null, - destroyTs: null, - }; - if (this.#persist) { - await this.#performWrite( + await this.#withActorWrite(actorId, async (lockedEntry) => { + lockedEntry.state = { actorId, - lockedEntry.generation, - lockedEntry.state, - ); - } - }); + name, + key: key as readonly string[], + createdAt: BigInt(Date.now()), + kvStorage: [], + startTs: null, + connectableTs: null, + sleepTs: null, + destroyTs: null, + }; + if (this.#persist) { + await this.#performWrite( + actorId, + lockedEntry.generation, + lockedEntry.state, + ); + } + if (initialKvState.length > 0) { + const db = this.#getOrCreateActorKvDatabase(actorId); + this.#putKvEntriesInDb(db, initialKvState); + } + }); + } + return entry; } - return entry; - } async sleepActor(actorId: string) { invariant( @@ -461,11 +638,12 @@ export class FileSystemGlobalState { // Stop actor invariant(actor.actor, "actor should be loaded"); await actor.actor.onStop("sleep"); - } finally { - // Ensure any pending KV writes finish before removing the entry. - await this.#withActorWrite(actorId, async () => {}); - actor.stopPromise?.resolve(); - actor.stopPromise = undefined; + } finally { + // Ensure any pending KV writes finish before removing the entry. + await this.#withActorWrite(actorId, async () => {}); + this.#closeActorKvDatabase(actorId); + actor.stopPromise?.resolve(); + actor.stopPromise = undefined; // Remove from map after stop is complete this.#actors.delete(actorId); @@ -515,8 +693,9 @@ export class FileSystemGlobalState { await actor.actor.onStop("destroy"); } - // Ensure any pending KV writes finish before deleting files. - await this.#withActorWrite(actorId, async () => {}); + // Ensure any pending KV writes finish before deleting files. + await this.#withActorWrite(actorId, async () => {}); + this.#closeActorKvDatabase(actorId); // Clear alarm timeout if exists if (actor.alarmTimeout) { @@ -924,25 +1103,25 @@ export class FileSystemGlobalState { connectableTs: now, sleepTs: null, // Clear sleep timestamp when actor wakes up }; - if (this.#persist) { - await this.#performWrite( - actorId, - lockedEntry.generation, - lockedEntry.state, - ); - } - }); + if (this.#persist) { + await this.#performWrite( + actorId, + lockedEntry.generation, + lockedEntry.state, + ); + } + }); // Finish entry.startPromise.resolve(); entry.startPromise = undefined; return entry.actor; - } catch (innerError) { - const error = new Error( - `Failed to start actor ${actorId}: ${innerError}`, - { cause: innerError }, - ); + } catch (innerError) { + const error = new Error( + `Failed to start actor ${actorId}: ${innerError}`, + { cause: innerError }, + ); entry.startPromise?.reject(error); entry.startPromise = undefined; throw error; @@ -1147,46 +1326,8 @@ export class FileSystemGlobalState { } throw new Error(`Actor ${actorId} state not loaded`); } - - // Create a mutable copy of kvStorage - const newKvStorage = [...entry.state.kvStorage]; - - // Update kvStorage with new entries - for (const [key, value] of entries) { - // Find existing entry with the same key - const existingIndex = newKvStorage.findIndex((e) => - arrayBuffersEqual(e.key, bufferToArrayBuffer(key)), - ); - - if (existingIndex >= 0) { - // Replace existing entry with new one - newKvStorage[existingIndex] = { - key: bufferToArrayBuffer(key), - value: bufferToArrayBuffer(value), - }; - } else { - // Add new entry - newKvStorage.push({ - key: bufferToArrayBuffer(key), - value: bufferToArrayBuffer(value), - }); - } - } - - // Update state with new kvStorage - entry.state = { - ...entry.state, - kvStorage: newKvStorage, - }; - - // Save state to disk - if (this.#persist) { - await this.#performWrite( - actorId, - entry.generation, - entry.state, - ); - } + const db = this.#getOrCreateActorKvDatabase(actorId); + this.#putKvEntriesInDb(db, entries); }); } @@ -1207,18 +1348,18 @@ export class FileSystemGlobalState { } } + const db = this.#getOrCreateActorKvDatabase(actorId); const results: (Uint8Array | null)[] = []; for (const key of keys) { - // Find entry with the same key - const foundEntry = entry.state.kvStorage.find((e) => - arrayBuffersEqual(e.key, bufferToArrayBuffer(key)), + const row = db.get<{ value: Uint8Array | ArrayBuffer }>( + "SELECT value FROM kv WHERE key = ?", + [key], ); - - if (foundEntry) { - results.push(new Uint8Array(foundEntry.value)); - } else { + if (!row) { results.push(null); + continue; } + results.push(ensureUint8Array(row.value, "value")); } return results; } @@ -1235,34 +1376,24 @@ export class FileSystemGlobalState { } throw new Error(`Actor ${actorId} state not loaded`); } - - // Create a mutable copy of kvStorage - const newKvStorage = [...entry.state.kvStorage]; - - // Delete entries from kvStorage - for (const key of keys) { - const indexToDelete = newKvStorage.findIndex((e) => - arrayBuffersEqual(e.key, bufferToArrayBuffer(key)), - ); - - if (indexToDelete >= 0) { - newKvStorage.splice(indexToDelete, 1); - } + if (keys.length === 0) { + return; } - // Update state with new kvStorage - entry.state = { - ...entry.state, - kvStorage: newKvStorage, - }; - - // Save state to disk - if (this.#persist) { - await this.#performWrite( - actorId, - entry.generation, - entry.state, - ); + const db = this.#getOrCreateActorKvDatabase(actorId); + db.exec("BEGIN"); + try { + for (const key of keys) { + db.run("DELETE FROM kv WHERE key = ?", [key]); + } + db.exec("COMMIT"); + } catch (error) { + try { + db.exec("ROLLBACK"); + } catch { + // Ignore rollback errors, original error is more actionable. + } + throw error; } }); } @@ -1284,23 +1415,21 @@ export class FileSystemGlobalState { } } - const results: [Uint8Array, Uint8Array][] = []; - for (const kvEntry of entry.state.kvStorage) { - const keyBytes = new Uint8Array(kvEntry.key); - // Check if key starts with prefix - if (keyBytes.length >= prefix.length) { - let hasPrefix = true; - for (let i = 0; i < prefix.length; i++) { - if (keyBytes[i] !== prefix[i]) { - hasPrefix = false; - break; - } - } - if (hasPrefix) { - results.push([keyBytes, new Uint8Array(kvEntry.value)]); - } - } - } - return results; + const db = this.#getOrCreateActorKvDatabase(actorId); + const upperBound = computePrefixUpperBound(prefix); + const rows = upperBound + ? db.all<{ key: Uint8Array | ArrayBuffer; value: Uint8Array | ArrayBuffer }>( + "SELECT key, value FROM kv WHERE key >= ? AND key < ? ORDER BY key ASC", + [prefix, upperBound], + ) + : db.all<{ key: Uint8Array | ArrayBuffer; value: Uint8Array | ArrayBuffer }>( + "SELECT key, value FROM kv WHERE key >= ? ORDER BY key ASC", + [prefix], + ); + + return rows.map((row) => [ + ensureUint8Array(row.key, "key"), + ensureUint8Array(row.value, "value"), + ]); } } diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/mod.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/mod.ts index 56c1c6e74d..81ce39e784 100644 --- a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/mod.ts @@ -16,6 +16,8 @@ export { getStoragePath } from "./utils"; const CreateFileSystemDriverOptionsSchema = z.object({ /** Custom path for storage. */ path: z.string().optional(), + /** Deprecated: file-system driver KV is now always SQLite-backed. */ + useNativeSqlite: z.boolean().optional(), }); type CreateFileSystemDriverOptionsInput = z.input< @@ -28,9 +30,16 @@ export function createFileSystemOrMemoryDriver( ): DriverConfig { importNodeDependencies(); + if (options?.useNativeSqlite === false) { + throw new Error( + "File-system driver no longer supports non-SQLite KV storage. Remove useNativeSqlite: false.", + ); + } + const stateOptions: FileSystemDriverOptions = { persist, customPath: options?.path, + useNativeSqlite: true, }; const state = new FileSystemGlobalState(stateOptions); const driverConfig: DriverConfig = { diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/sqlite-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/sqlite-runtime.ts new file mode 100644 index 0000000000..323b7ebd2a --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/sqlite-runtime.ts @@ -0,0 +1,206 @@ +import { getRequireFn } from "@/utils/node"; + +type SqliteRuntimeKind = "bun" | "node" | "better-sqlite3"; + +interface SqliteStatement { + run(...params: unknown[]): unknown; + get(...params: unknown[]): unknown; + all(...params: unknown[]): unknown[]; +} + +interface SqliteRawDatabase { + exec(sql: string): unknown; + close(): unknown; + prepare?(sql: string): SqliteStatement; + query?(sql: string): SqliteStatement; +} + +export interface SqliteRuntimeDatabase { + exec(sql: string): void; + run(sql: string, params?: readonly unknown[]): void; + get>( + sql: string, + params?: readonly unknown[], + ): T | undefined; + all>( + sql: string, + params?: readonly unknown[], + ): T[]; + close(): void; +} + +export interface SqliteRuntime { + kind: SqliteRuntimeKind; + open(path: string): SqliteRuntimeDatabase; +} + +function normalizeParams(params: readonly unknown[] | undefined): unknown[] { + if (!params || params.length === 0) { + return []; + } + + return params.map((value) => { + if (value instanceof Uint8Array) { + return Buffer.from(value); + } + return value; + }); +} + +function createPreparedDatabaseAdapter( + rawDb: SqliteRawDatabase, + prepare: (sql: string) => SqliteStatement, +): SqliteRuntimeDatabase { + return { + exec: (sql) => { + rawDb.exec(sql); + }, + run: (sql, params) => { + const stmt = prepare(sql); + stmt.run(...normalizeParams(params)); + }, + get: >( + sql: string, + params?: readonly unknown[], + ) => { + const stmt = prepare(sql); + return stmt.get(...normalizeParams(params)) as T | undefined; + }, + all: >( + sql: string, + params?: readonly unknown[], + ) => { + const stmt = prepare(sql); + return stmt.all(...normalizeParams(params)) as T[]; + }, + close: () => { + rawDb.close(); + }, + }; +} + +export function loadSqliteRuntime(): SqliteRuntime { + const requireFn = getRequireFn(); + const loadErrors: string[] = []; + + try { + const bunSqlite = requireFn(/* webpackIgnore: true */ "bun:sqlite") as { + Database: new (path: string) => SqliteRawDatabase; + }; + if (typeof bunSqlite.Database === "function") { + return { + kind: "bun", + open: (path) => { + const rawDb = new bunSqlite.Database(path); + const query = (sql: string) => { + if (!rawDb.query) { + throw new Error("bun:sqlite database missing query method"); + } + return rawDb.query(sql); + }; + return createPreparedDatabaseAdapter(rawDb, query); + }, + }; + } + } catch (error) { + loadErrors.push(`bun:sqlite unavailable: ${String(error)}`); + } + + try { + const nodeSqlite = requireFn(/* webpackIgnore: true */ "node:sqlite") as { + DatabaseSync: new (path: string) => SqliteRawDatabase; + }; + if (typeof nodeSqlite.DatabaseSync === "function") { + return { + kind: "node", + open: (path) => { + const rawDb = new nodeSqlite.DatabaseSync(path); + const prepare = (sql: string) => { + if (!rawDb.prepare) { + throw new Error( + "node:sqlite DatabaseSync missing prepare method", + ); + } + return rawDb.prepare(sql); + }; + return createPreparedDatabaseAdapter(rawDb, prepare); + }, + }; + } + } catch (error) { + loadErrors.push(`node:sqlite unavailable: ${String(error)}`); + } + + try { + const betterSqlite3Module = requireFn( + /* webpackIgnore: true */ "better-sqlite3", + ) as { + default?: new (path: string) => SqliteRawDatabase; + } | (new (path: string) => SqliteRawDatabase); + const BetterSqlite3 = + typeof betterSqlite3Module === "function" + ? betterSqlite3Module + : betterSqlite3Module.default; + + if (typeof BetterSqlite3 === "function") { + return { + kind: "better-sqlite3", + open: (path) => { + const rawDb = new BetterSqlite3(path); + const prepare = (sql: string) => { + if (!rawDb.prepare) { + throw new Error( + "better-sqlite3 database missing prepare method", + ); + } + return rawDb.prepare(sql); + }; + return createPreparedDatabaseAdapter(rawDb, prepare); + }, + }; + } + } catch (error) { + loadErrors.push(`better-sqlite3 unavailable: ${String(error)}`); + throw new Error( + `No SQLite runtime available. Tried bun:sqlite, node:sqlite, and better-sqlite3. Install better-sqlite3 (e.g. "pnpm add better-sqlite3") if native runtimes are unavailable.\n${loadErrors.join("\n")}`, + ); + } + + throw new Error( + `No SQLite runtime available. Tried bun:sqlite, node:sqlite, and better-sqlite3.\n${loadErrors.join("\n")}`, + ); +} + +export function computePrefixUpperBound( + prefix: Uint8Array, +): Uint8Array | undefined { + if (prefix.length === 0) { + return undefined; + } + + const upperBound = new Uint8Array(prefix); + for (let i = upperBound.length - 1; i >= 0; i--) { + if (upperBound[i] !== 0xff) { + upperBound[i] += 1; + return upperBound.slice(0, i + 1); + } + } + return undefined; +} + +export function ensureUint8Array( + value: unknown, + fieldName: string, +): Uint8Array { + if (value instanceof Uint8Array) { + return value; + } + if (value instanceof ArrayBuffer) { + return new Uint8Array(value); + } + if (ArrayBuffer.isView(value)) { + const view = value as ArrayBufferView; + return new Uint8Array(view.buffer, view.byteOffset, view.byteLength); + } + throw new Error(`SQLite row field "${fieldName}" is not binary data`); +} diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/config/index.ts b/rivetkit-typescript/packages/rivetkit/src/registry/config/index.ts index 40c1b055d3..a4c7e71ac6 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/config/index.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/config/index.ts @@ -1,13 +1,15 @@ -import { z } from "zod/v4"; -import type { ActorDefinition, AnyActorDefinition } from "@/actor/definition"; +import { z } from "zod"; import { getRunMetadata } from "@/actor/config"; +import type { ActorDefinition, AnyActorDefinition } from "@/actor/definition"; import { type Logger, LogLevelSchema } from "@/common/log"; import { ENGINE_ENDPOINT } from "@/engine-process/constants"; import { InspectorConfigSchema } from "@/inspector/config"; +import { DeepReadonly } from "@/utils"; import { tryParseEndpoint } from "@/utils/endpoint-parser"; import { getRivetEndpoint, getRivetEngine, + getRivetkitStoragePath, getRivetNamespace, getRivetToken, isDev, @@ -15,13 +17,12 @@ import { import { type DriverConfig, DriverConfigSchema } from "./driver"; import { RunnerConfigSchema } from "./runner"; import { ServerlessConfigSchema } from "./serverless"; -import { DeepReadonly } from "@/utils"; export { DriverConfigSchema, type DriverConfig }; export const ActorsSchema = z.record( z.string(), - z.custom>(), + z.custom>(), ); export type RegistryActors = z.infer; @@ -45,6 +46,17 @@ export const RegistryConfigSchema = z // MARK: Driver driver: DriverConfigSchema.optional(), + /** + * Storage path for RivetKit file-system state when using the default driver. + * + * If not set, RivetKit uses the platform default data location. + * + * Can also be set via RIVETKIT_STORAGE_PATH. + */ + storagePath: z + .string() + .optional() + .transform((val) => val ?? getRivetkitStoragePath()), // MARK: Networking /** @experimental */ @@ -290,9 +302,22 @@ export function buildActorNames( export const DocInspectorConfigSchema = z .object({ - enabled: z.boolean().optional().describe("Whether to enable the Rivet Inspector. Defaults to true in development mode."), - token: z.string().optional().describe("Token used to access the Inspector."), - defaultEndpoint: z.string().optional().describe("Default RivetKit server endpoint for Rivet Inspector to connect to."), + enabled: z + .boolean() + .optional() + .describe( + "Whether to enable the Rivet Inspector. Defaults to true in development mode.", + ), + token: z + .string() + .optional() + .describe("Token used to access the Inspector."), + defaultEndpoint: z + .string() + .optional() + .describe( + "Default RivetKit server endpoint for Rivet Inspector to connect to.", + ), }) .optional() .describe("Inspector configuration for debugging and development."); @@ -300,53 +325,184 @@ export const DocInspectorConfigSchema = z export const DocConfigureRunnerPoolSchema = z .object({ name: z.string().optional().describe("Name of the runner pool."), - url: z.string().describe("URL of the serverless platform to configure runners."), - headers: z.record(z.string(), z.string()).optional().describe("Headers to include in requests to the serverless platform."), - maxRunners: z.number().optional().describe("Maximum number of runners in the pool."), - minRunners: z.number().optional().describe("Minimum number of runners to keep warm."), - requestLifespan: z.number().optional().describe("Maximum lifespan of a request in milliseconds."), - runnersMargin: z.number().optional().describe("Buffer margin for scaling runners."), - slotsPerRunner: z.number().optional().describe("Number of actor slots per runner."), - metadata: z.record(z.string(), z.unknown()).optional().describe("Additional metadata to pass to the serverless platform."), - metadataPollInterval: z.number().optional().describe("Interval in milliseconds between metadata polls from the engine. Defaults to 10000 milliseconds (10 seconds)."), + url: z + .string() + .describe("URL of the serverless platform to configure runners."), + headers: z + .record(z.string(), z.string()) + .optional() + .describe( + "Headers to include in requests to the serverless platform.", + ), + maxRunners: z + .number() + .optional() + .describe("Maximum number of runners in the pool."), + minRunners: z + .number() + .optional() + .describe("Minimum number of runners to keep warm."), + requestLifespan: z + .number() + .optional() + .describe("Maximum lifespan of a request in milliseconds."), + runnersMargin: z + .number() + .optional() + .describe("Buffer margin for scaling runners."), + slotsPerRunner: z + .number() + .optional() + .describe("Number of actor slots per runner."), + metadata: z + .record(z.string(), z.unknown()) + .optional() + .describe( + "Additional metadata to pass to the serverless platform.", + ), + metadataPollInterval: z + .number() + .optional() + .describe( + "Interval in milliseconds between metadata polls from the engine. Defaults to 10000 milliseconds (10 seconds).", + ), }) .optional(); -export const DocServerlessConfigSchema = z.object({ - spawnEngine: z.boolean().optional().describe("Downloads and starts the full Rust engine process. Auto-enabled in development mode when no endpoint is provided. Default: false"), - engineVersion: z.string().optional().describe("Version of the engine to download. Defaults to the current RivetKit version."), - configureRunnerPool: DocConfigureRunnerPoolSchema.describe("Automatically configure serverless runners in the engine."), - basePath: z.string().optional().describe("Base path for serverless API routes. Default: '/api/rivet'"), - publicEndpoint: z.string().optional().describe("The endpoint that clients should connect to. Supports URL auth syntax: https://namespace:token@api.rivet.dev"), - publicToken: z.string().optional().describe("Token that clients should use when connecting via the public endpoint."), -}).describe("Configuration for serverless deployment mode."); - -export const DocRunnerConfigSchema = z.object({ - totalSlots: z.number().optional().describe("Total number of actor slots available. Default: 100000"), - runnerName: z.string().optional().describe("Name of this runner. Default: 'default'"), - runnerKey: z.string().optional().describe("Authentication key for the runner."), - version: z.number().optional().describe("Version number of this runner. Default: 1"), -}).describe("Configuration for runner mode."); +export const DocServerlessConfigSchema = z + .object({ + spawnEngine: z + .boolean() + .optional() + .describe( + "Downloads and starts the full Rust engine process. Auto-enabled in development mode when no endpoint is provided. Default: false", + ), + engineVersion: z + .string() + .optional() + .describe( + "Version of the engine to download. Defaults to the current RivetKit version.", + ), + configureRunnerPool: DocConfigureRunnerPoolSchema.describe( + "Automatically configure serverless runners in the engine.", + ), + basePath: z + .string() + .optional() + .describe( + "Base path for serverless API routes. Default: '/api/rivet'", + ), + publicEndpoint: z + .string() + .optional() + .describe( + "The endpoint that clients should connect to. Supports URL auth syntax: https://namespace:token@api.rivet.dev", + ), + publicToken: z + .string() + .optional() + .describe( + "Token that clients should use when connecting via the public endpoint.", + ), + }) + .describe("Configuration for serverless deployment mode."); + +export const DocRunnerConfigSchema = z + .object({ + totalSlots: z + .number() + .optional() + .describe("Total number of actor slots available. Default: 100000"), + runnerName: z + .string() + .optional() + .describe("Name of this runner. Default: 'default'"), + runnerKey: z + .string() + .optional() + .describe("Authentication key for the runner."), + version: z + .number() + .optional() + .describe("Version number of this runner. Default: 1"), + }) + .describe("Configuration for runner mode."); export const DocRegistryConfigSchema = z .object({ - use: z.record(z.string(), z.unknown()).describe("Actor definitions. Keys are actor names, values are actor definitions."), - maxIncomingMessageSize: z.number().optional().describe("Maximum size of incoming WebSocket messages in bytes. Default: 65536"), - maxOutgoingMessageSize: z.number().optional().describe("Maximum size of outgoing WebSocket messages in bytes. Default: 1048576"), - noWelcome: z.boolean().optional().describe("Disable the welcome message on startup. Default: false"), + use: z + .record(z.string(), z.unknown()) + .describe( + "Actor definitions. Keys are actor names, values are actor definitions.", + ), + storagePath: z + .string() + .optional() + .describe( + "Storage path for RivetKit file-system state when using the default driver. Can also be set via RIVETKIT_STORAGE_PATH.", + ), + maxIncomingMessageSize: z + .number() + .optional() + .describe( + "Maximum size of incoming WebSocket messages in bytes. Default: 65536", + ), + maxOutgoingMessageSize: z + .number() + .optional() + .describe( + "Maximum size of outgoing WebSocket messages in bytes. Default: 1048576", + ), + noWelcome: z + .boolean() + .optional() + .describe("Disable the welcome message on startup. Default: false"), logging: z .object({ - level: LogLevelSchema.optional().describe("Log level for RivetKit. Default: 'warn'"), + level: LogLevelSchema.optional().describe( + "Log level for RivetKit. Default: 'warn'", + ), }) .optional() .describe("Logging configuration."), - endpoint: z.string().optional().describe("Endpoint URL to connect to Rivet Engine. Supports URL auth syntax: https://namespace:token@api.rivet.dev. Can also be set via RIVET_ENDPOINT environment variable."), - token: z.string().optional().describe("Authentication token for Rivet Engine. Can also be set via RIVET_TOKEN environment variable."), - namespace: z.string().optional().describe("Namespace to use. Default: 'default'. Can also be set via RIVET_NAMESPACE environment variable."), - headers: z.record(z.string(), z.string()).optional().describe("Additional headers to include in requests to Rivet Engine."), - serveManager: z.boolean().optional().describe("Whether to start the local manager server. Auto-determined based on endpoint and NODE_ENV if not specified."), - managerBasePath: z.string().optional().describe("Base path for the manager API. Default: '/'"), - managerPort: z.number().optional().describe("Port to run the manager on. Default: 6420"), + endpoint: z + .string() + .optional() + .describe( + "Endpoint URL to connect to Rivet Engine. Supports URL auth syntax: https://namespace:token@api.rivet.dev. Can also be set via RIVET_ENDPOINT environment variable.", + ), + token: z + .string() + .optional() + .describe( + "Authentication token for Rivet Engine. Can also be set via RIVET_TOKEN environment variable.", + ), + namespace: z + .string() + .optional() + .describe( + "Namespace to use. Default: 'default'. Can also be set via RIVET_NAMESPACE environment variable.", + ), + headers: z + .record(z.string(), z.string()) + .optional() + .describe( + "Additional headers to include in requests to Rivet Engine.", + ), + serveManager: z + .boolean() + .optional() + .describe( + "Whether to start the local manager server. Auto-determined based on endpoint and NODE_ENV if not specified.", + ), + managerBasePath: z + .string() + .optional() + .describe("Base path for the manager API. Default: '/'"), + managerPort: z + .number() + .optional() + .describe("Port to run the manager on. Default: 6420"), inspector: DocInspectorConfigSchema, serverless: DocServerlessConfigSchema.optional(), runner: DocRunnerConfigSchema.optional(), diff --git a/rivetkit-typescript/packages/rivetkit/src/utils/env-vars.ts b/rivetkit-typescript/packages/rivetkit/src/utils/env-vars.ts index 8544d31374..ea2ad5cef8 100644 --- a/rivetkit-typescript/packages/rivetkit/src/utils/env-vars.ts +++ b/rivetkit-typescript/packages/rivetkit/src/utils/env-vars.ts @@ -44,6 +44,8 @@ export const getRivetkitInspectorToken = (): string | undefined => getEnvUniversal("RIVET_INSPECTOR_TOKEN"); export const getRivetkitInspectorDisable = (): boolean => getEnvUniversal("RIVET_INSPECTOR_DISABLE") === "1"; +export const getRivetkitStoragePath = (): string | undefined => + getEnvUniversal("RIVETKIT_STORAGE_PATH"); // Logging configuration // DEPRECATED: LOG_LEVEL will be removed in a future version diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts index af39f3cc1b..89d67bc30b 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts @@ -2,6 +2,17 @@ import type { RunContext } from "@/actor/contexts/run"; import type { Client } from "@/client/client"; import type { Registry } from "@/registry"; import type { AnyDatabaseProvider, InferDatabaseClient } from "@/actor/database"; +import type { + QueueFilterName, + QueueNextOptions, + QueueResultMessageForName, +} from "@/actor/instance/queue"; +import type { + EventSchemaConfig, + InferEventArgs, + InferSchemaMap, + QueueSchemaConfig, +} from "@/actor/schema"; import type { WorkflowContextInterface } from "@rivetkit/workflow-engine"; import type { BranchConfig, @@ -10,10 +21,72 @@ import type { LoopConfig, LoopResult, StepConfig, - WorkflowListenMessage, + WorkflowQueueMessage, } from "@rivetkit/workflow-engine"; import { WORKFLOW_GUARD_KV_KEY } from "./constants"; +type WorkflowActorQueueNextOptions< + TName extends string, + TCompletable extends boolean, +> = Omit, "signal">; + +type WorkflowActorQueueNextOptionsFallback = Omit< + QueueNextOptions, + "signal" +>; + +type ActorWorkflowLoopConfig< + S, + T, + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig, + TQueues extends QueueSchemaConfig, +> = Omit, "run"> & { + run: ( + ctx: ActorWorkflowContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, + state: S, + ) => Promise>; +}; + +type ActorWorkflowBranchConfig< + TOutput, + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig, + TQueues extends QueueSchemaConfig, +> = { + run: ( + ctx: ActorWorkflowContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, + ) => Promise; +}; + export class ActorWorkflowContext< TState, TConnParams, @@ -21,17 +94,37 @@ export class ActorWorkflowContext< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, > implements WorkflowContextInterface { #inner: WorkflowContextInterface; - #runCtx: RunContext; + #runCtx: RunContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >; #actorAccessDepth = 0; #allowActorAccess = false; #guardViolation = false; constructor( inner: WorkflowContextInterface, - runCtx: RunContext, + runCtx: RunContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, ) { this.#inner = inner; this.#runCtx = runCtx; @@ -45,39 +138,136 @@ export class ActorWorkflowContext< return this.#inner.abortSignal; } - async step( - nameOrConfig: string | Parameters[0], - run?: () => Promise, - ): Promise { + get queue() { + const self = this; + function next< + const TName extends QueueFilterName, + const TCompletable extends boolean = false, + >( + name: string, + opts?: WorkflowActorQueueNextOptions, + ): Promise>>; + function next( + name: string, + opts?: WorkflowActorQueueNextOptionsFallback, + ): Promise< + Array< + QueueResultMessageForName< + TQueues, + QueueFilterName, + TCompletable + > + > + >; + async function next( + name: string, + opts?: WorkflowActorQueueNextOptions, + ): Promise>> { + const messages = await self.#inner.queue.next(name, opts); + return messages.map((message) => self.#toActorQueueMessage(message)); + } + + function send( + name: K, + body: InferSchemaMap[K], + ): Promise; + function send( + name: keyof TQueues extends never ? string : never, + body: unknown, + ): Promise; + async function send(name: string, body: unknown): Promise { + await self.#runCtx.queue.send(name as never, body as never); + } + + return { + next, + send, + }; + } + + async step( + nameOrConfig: string | Parameters[0], + run?: () => Promise, + ): Promise { if (typeof nameOrConfig === "string") { if (!run) { throw new Error("Step run function missing"); } - return await this.#wrapActive(() => - this.#inner.step(nameOrConfig, () => - this.#withActorAccess(run), - ), - ); - } - const stepConfig = nameOrConfig as StepConfig; - const config: StepConfig = { - ...stepConfig, - run: () => this.#withActorAccess(stepConfig.run), - }; - return await this.#wrapActive(() => this.#inner.step(config)); + return await this.#wrapActive(() => + this.#inner.step(nameOrConfig, () => this.#withActorAccess(run)), + ); } + const stepConfig = nameOrConfig as StepConfig; + const config: StepConfig = { + ...stepConfig, + run: () => this.#withActorAccess(stepConfig.run), + }; + return await this.#wrapActive(() => this.#inner.step(config)); + } + async loop( + name: string, + run: ( + ctx: ActorWorkflowContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, + ) => Promise>, + ): Promise; async loop( name: string, run: ( ctx: WorkflowContextInterface, ) => Promise>, ): Promise; + async loop( + config: ActorWorkflowLoopConfig< + S, + T, + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, + ): Promise; async loop(config: LoopConfig): Promise; async loop( - nameOrConfig: string | LoopConfig, + nameOrConfig: + | string + | LoopConfig + | ActorWorkflowLoopConfig< + unknown, + unknown, + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, run?: ( - ctx: WorkflowContextInterface, + ctx: ActorWorkflowContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, ) => Promise>, ): Promise { if (typeof nameOrConfig === "string") { @@ -106,68 +296,34 @@ export class ActorWorkflowContext< return this.#inner.sleepUntil(name, timestampMs); } - listen( - name: string, - messageName: string | string[], - ): Promise> { - return this.#inner.listen(name, messageName); - } - - listenN( - name: string, - messageName: string, - limit: number, - ): Promise { - return this.#inner.listenN(name, messageName, limit); - } - - listenWithTimeout( - name: string, - messageName: string, - timeoutMs: number, - ): Promise { - return this.#inner.listenWithTimeout(name, messageName, timeoutMs); - } - - listenUntil( - name: string, - messageName: string, - timestampMs: number, - ): Promise { - return this.#inner.listenUntil(name, messageName, timestampMs); - } - - listenNWithTimeout( - name: string, - messageName: string, - limit: number, - timeoutMs: number, - ): Promise { - return this.#inner.listenNWithTimeout( - name, - messageName, - limit, - timeoutMs, - ); - } - - listenNUntil( - name: string, - messageName: string, - limit: number, - timestampMs: number, - ): Promise { - return this.#inner.listenNUntil(name, messageName, limit, timestampMs); - } - async rollbackCheckpoint(name: string): Promise { await this.#wrapActive(() => this.#inner.rollbackCheckpoint(name)); } + async join< + T extends Record< + string, + ActorWorkflowBranchConfig< + unknown, + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + > + >, + >( + name: string, + branches: T, + ): Promise<{ [K in keyof T]: Awaited> }>; async join>>( name: string, branches: T, - ): Promise<{ [K in keyof T]: BranchOutput }> { + ): Promise<{ [K in keyof T]: BranchOutput }>; + async join(name: string, branches: Record>) { const wrappedBranches = Object.fromEntries( Object.entries(branches).map(([key, branch]) => [ key, @@ -176,13 +332,30 @@ export class ActorWorkflowContext< branch.run(this.#createChildContext(ctx)), }, ]), - ) as T; - - return (await this.#wrapActive(() => + ) as Record>; + return await this.#wrapActive(() => this.#inner.join(name, wrappedBranches), - )) as { [K in keyof T]: BranchOutput }; + ); } + async race( + name: string, + branches: Array<{ + name: string; + run: ( + ctx: ActorWorkflowContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, + ) => Promise; + }>, + ): Promise<{ winner: string; value: T }>; async race( name: string, branches: Array<{ @@ -246,8 +419,37 @@ export class ActorWorkflowContext< return this.#runCtx.actorId; } - broadcast>(name: string, ...args: Args): void { - this.#runCtx.broadcast(name, ...args); + broadcast( + name: K, + ...args: InferEventArgs[K]> + ): void; + broadcast( + name: keyof TEvents extends never ? string : never, + ...args: Array + ): void; + broadcast(name: string, ...args: Array): void { + this.#runCtx.broadcast( + name as never, + ...((args as unknown[]) as never[]), + ); + } + + #toActorQueueMessage( + message: WorkflowQueueMessage, + ): WorkflowQueueMessage & { id: bigint } { + let id: bigint; + try { + id = BigInt(message.id); + } catch { + throw new Error(`Invalid queue message id "${message.id}"`); + } + return { + id, + name: message.name, + body: message.body, + createdAt: message.createdAt, + ...(message.complete ? { complete: message.complete } : {}), + }; } async #wrapActive(run: () => Promise): Promise { @@ -321,7 +523,9 @@ export class ActorWorkflowContext< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues > { return new ActorWorkflowContext(ctx, this.#runCtx); } diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts index 7997b588cc..3213597adb 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts @@ -9,77 +9,66 @@ import type { WorkflowMessageDriver, } from "@rivetkit/workflow-engine"; -const WORKFLOW_QUEUE_PREFIX = "__workflow:"; - -export function workflowQueueName(name: string): string { - return `${WORKFLOW_QUEUE_PREFIX}${name}`; -} - -function stripWorkflowQueueName(name: string): string | null { - if (!name.startsWith(WORKFLOW_QUEUE_PREFIX)) { - return null; - } - return name.slice(WORKFLOW_QUEUE_PREFIX.length); -} - function stripWorkflowKey(prefixed: Uint8Array): Uint8Array { return prefixed.slice(KEYS.WORKFLOW_PREFIX.length); } class ActorWorkflowMessageDriver implements WorkflowMessageDriver { #actor: AnyActorInstance; - #runCtx: RunContext; - #completionHandles = new Map Promise>(); + #runCtx: RunContext; constructor( actor: AnyActorInstance, - runCtx: RunContext, + runCtx: RunContext, ) { this.#actor = actor; this.#runCtx = runCtx; } async loadMessages(): Promise { - const queueMessages = await this.#runCtx.keepAwake( - this.#actor.queueManager.getMessages(), - ); - const now = Date.now(); - - const workflowMessages: Message[] = []; - for (const queueMessage of queueMessages) { - if (queueMessage.inFlight || queueMessage.availableAt > now) { - continue; - } - - const workflowName = stripWorkflowQueueName(queueMessage.name); - if (!workflowName) continue; - const id = queueMessage.id.toString(); - this.#completionHandles.set(id, async (response?: unknown) => { - await this.#runCtx.keepAwake( - this.#actor.queueManager.completeById(queueMessage.id, response), - ); - }); - workflowMessages.push({ - id, - name: workflowName, - data: queueMessage.body, - sentAt: queueMessage.createdAt, - complete: async (response?: unknown) => { - await this.completeMessage(id, response); - }, - }); - } - - return workflowMessages; + // Actor-backed workflows use receiveMessages() directly and do not + // mirror queue messages into workflow-engine storage. + return []; } async addMessage(message: Message): Promise { await this.#runCtx.keepAwake( - this.#actor.queueManager.enqueue( - workflowQueueName(message.name), - message.data, + this.#actor.queueManager.enqueue(message.name, message.data), + ); + } + + async receiveMessages(opts: { + names?: readonly string[]; + count: number; + completable: boolean; + }): Promise { + const messages = await this.#runCtx.keepAwake( + this.#actor.queueManager.receive( + opts.names && opts.names.length > 0 ? [...opts.names] : undefined, + opts.count, + 0, + undefined, + opts.completable, ), ); + return messages.map((message) => ({ + id: message.id.toString(), + name: message.name, + data: message.body, + sentAt: message.createdAt, + ...(opts.completable + ? { + complete: async (response?: unknown) => { + await this.#runCtx.keepAwake( + this.#actor.queueManager.completeMessage( + message, + response, + ), + ); + }, + } + : {}), + })); } async deleteMessages(messageIds: string[]): Promise { @@ -103,34 +92,13 @@ class ActorWorkflowMessageDriver implements WorkflowMessageDriver { } const deleted = await this.#runCtx.keepAwake( - this.#actor.queueManager.deleteMessagesById(validIds, { - resolveWaiters: false, - }), + this.#actor.queueManager.deleteMessagesById(validIds), ); - for (const id of deleted) { - const idString = id.toString(); - if (this.#completionHandles.has(idString)) { - continue; - } - this.#completionHandles.set(idString, async (response?: unknown) => { - await this.#runCtx.keepAwake( - this.#actor.queueManager.completeById(id, response), - ); - }); - } - return deleted.map((id) => id.toString()); } async completeMessage(messageId: string, response?: unknown): Promise { - const complete = this.#completionHandles.get(messageId); - if (complete) { - await complete(response); - this.#completionHandles.delete(messageId); - return; - } - let parsedId: bigint; try { parsedId = BigInt(messageId); @@ -139,7 +107,7 @@ class ActorWorkflowMessageDriver implements WorkflowMessageDriver { } await this.#runCtx.keepAwake( - this.#actor.queueManager.completeById(parsedId, response), + this.#actor.queueManager.completeMessageById(parsedId, response), ); } } @@ -148,21 +116,17 @@ export class ActorWorkflowDriver implements EngineDriver { readonly workerPollInterval = 100; readonly messageDriver: WorkflowMessageDriver; #actor: AnyActorInstance; - #runCtx: RunContext; + #runCtx: RunContext; constructor( actor: AnyActorInstance, - runCtx: RunContext, + runCtx: RunContext, ) { this.#actor = actor; this.#runCtx = runCtx; this.messageDriver = new ActorWorkflowMessageDriver(actor, runCtx); } - #log(msg: string, data?: Record) { - this.#runCtx.log.info({ msg: `[workflow-driver] ${msg}`, ...data }); - } - async get(key: Uint8Array): Promise { const [value] = await this.#runCtx.keepAwake( this.#actor.driver.kvBatchGet(this.#actor.id, [ @@ -250,7 +214,9 @@ export class ActorWorkflowDriver implements EngineDriver { messageNames: string[], abortSignal: AbortSignal, ): Promise { - const queueNames = messageNames.map((name) => workflowQueueName(name)); - return this.#actor.queueManager.waitForNames(queueNames, abortSignal); + return this.#actor.queueManager.waitForNames( + messageNames.length > 0 ? messageNames : undefined, + abortSignal, + ); } } diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts index 75bb637e28..23ef812b8c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts @@ -2,16 +2,16 @@ import { ACTOR_CONTEXT_INTERNAL_SYMBOL } from "@/actor/contexts/base/actor"; import type { RunContext } from "@/actor/contexts/run"; import type { AnyDatabaseProvider } from "@/actor/database"; import type { AnyActorInstance } from "@/actor/instance/mod"; -import type { RunConfig } from "@/actor/config"; +import type { EventSchemaConfig, QueueSchemaConfig } from "@/actor/schema"; +import { RUN_FUNCTION_CONFIG_SYMBOL } from "@/actor/config"; import { stringifyError } from "@/utils"; import { runWorkflow } from "@rivetkit/workflow-engine"; import invariant from "invariant"; import { ActorWorkflowContext } from "./context"; -import { ActorWorkflowDriver, workflowQueueName } from "./driver"; +import { ActorWorkflowDriver } from "./driver"; import { createWorkflowInspectorAdapter } from "./inspector"; export { Loop } from "@rivetkit/workflow-engine"; -export { workflowQueueName } from "./driver"; export { ActorWorkflowContext } from "./context"; export function workflow< @@ -21,6 +21,8 @@ export function workflow< TVars, TInput, TDatabase extends AnyDatabaseProvider, + TEvents extends EventSchemaConfig = Record, + TQueues extends QueueSchemaConfig = Record, >( fn: ( ctx: ActorWorkflowContext< @@ -29,22 +31,37 @@ export function workflow< TConnState, TVars, TInput, - TDatabase + TDatabase, + TEvents, + TQueues >, ) => Promise, -): RunConfig { +): ( + c: RunContext< + TState, + TConnParams, + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, +) => Promise { const workflowInspector = createWorkflowInspectorAdapter(); async function run( runCtx: RunContext< TState, TConnParams, - TConnState, - TVars, - TInput, - TDatabase - >, - ): Promise { + TConnState, + TVars, + TInput, + TDatabase, + TEvents, + TQueues + >, + ): Promise { const actor = ( runCtx as unknown as { [ACTOR_CONTEXT_INTERNAL_SYMBOL]?: AnyActorInstance; @@ -92,9 +109,16 @@ export function workflow< }); } - return { + const runWithConfig = run as typeof run & { + [RUN_FUNCTION_CONFIG_SYMBOL]?: { + icon?: string; + inspector?: { workflow: typeof workflowInspector.adapter }; + }; + }; + runWithConfig[RUN_FUNCTION_CONFIG_SYMBOL] = { icon: "diagram-project", - run: run as RunConfig["run"], inspector: { workflow: workflowInspector.adapter }, }; + + return runWithConfig; } diff --git a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts index f469e77e48..2105dcdbc6 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/actor-types.test.ts @@ -1,9 +1,25 @@ import { describe, expectTypeOf, it } from "vitest"; +import { actor, event, queue } from "@/actor/mod"; import type { ActorContext, ActorContextOf } from "@/actor/contexts"; import type { ActorDefinition } from "@/actor/definition"; import type { DatabaseProviderContext } from "@/db/config"; +import { db } from "@/db/mod"; +import { workflow } from "@/workflow/mod"; describe("ActorDefinition", () => { + describe("schema config types", () => { + it("events do not accept queue-style schemas", () => { + actor({ + state: {}, + events: { + // @ts-expect-error events must use primitive schemas, not queue definitions. + invalid: queue<{ foo: string }, { ok: true }>(), + }, + actions: {}, + }); + }); + }); + describe("ActorContextOf type utility", () => { it("should correctly extract the context type from an ActorDefinition", () => { // Define some simple types for testing @@ -42,6 +58,8 @@ describe("ActorDefinition", () => { TestVars, TestInput, TestDatabase, + Record, + Record, TestActions >; @@ -55,7 +73,9 @@ describe("ActorDefinition", () => { TestConnState, TestVars, TestInput, - TestDatabase + TestDatabase, + Record, + Record > >(); @@ -73,9 +93,168 @@ describe("ActorDefinition", () => { TestConnState, TestVars, TestInput, - TestDatabase + TestDatabase, + Record, + Record > >(); }); }); + + describe("queue type inference", () => { + const queueTypeActor = actor({ + state: {}, + queues: { + foo: queue<{ fooBody: string }>(), + bar: queue<{ barBody: number }>(), + completable: queue<{ input: string }, { output: string }>(), + }, + actions: {}, + }); + + type QueueTypeContext = ActorContextOf; + + async function receiveFooBar(c: QueueTypeContext) { + return await c.queue.next({ + names: ["foo", "bar"] as const, + }); + } + + async function receiveCompletableManual(c: QueueTypeContext) { + return await c.queue.next({ + names: ["completable"] as const, + completable: true, + }); + } + + async function receiveFromAllQueues(c: QueueTypeContext) { + for await (const message of c.queue.iter()) { + return message; + } + + throw new Error("queue iteration terminated unexpectedly"); + } + + it("narrows message body by queue name", () => { + type ReceivedFooBar = Awaited>[number]; + type FooBody = Extract["body"]; + type BarBody = Extract["body"]; + + expectTypeOf().toEqualTypeOf<{ fooBody: string }>(); + expectTypeOf().toEqualTypeOf<{ barBody: number }>(); + }); + + it("completable queue messages expose correctly typed complete()", () => { + type ManualMessage = Awaited< + ReturnType + >[number]; + type CompleteArgs = ManualMessage extends { + complete: (...args: infer TArgs) => Promise; + } + ? TArgs + : never; + + expectTypeOf().toEqualTypeOf< + [response: { output: string }] + >(); + }); + + it("infers queue body types when iterating c.queue.iter()", () => { + type Received = Awaited>; + type FooBody = Extract["body"]; + type BarBody = Extract["body"]; + type CompletableBody = Extract< + Received, + { name: "completable" } + >["body"]; + + expectTypeOf().toEqualTypeOf<{ fooBody: string }>(); + expectTypeOf().toEqualTypeOf<{ barBody: number }>(); + expectTypeOf().toEqualTypeOf<{ input: string }>(); + }); + }); + + describe("workflow context type inference", () => { + it("infers queue and event types for workflow ctx", () => { + actor({ + state: {}, + queues: { + foo: queue<{ fooBody: string }>(), + bar: queue<{ barBody: number }>(), + }, + events: { + updated: event<{ count: number }>(), + pair: event<[number, string]>(), + }, + run: workflow(async (ctx) => { + const [single] = await ctx.queue.next("wait-single", { + names: ["foo"] as const, + }); + if (single && single.name === "foo") { + expectTypeOf(single.body).toEqualTypeOf<{ + fooBody: string; + }>(); + } + + const [union] = await ctx.queue.next("wait-union", { + names: ["foo", "bar"], + }); + if (union?.name === "foo") { + expectTypeOf(union.body).toEqualTypeOf<{ + fooBody: string; + }>(); + } + if (union?.name === "bar") { + expectTypeOf(union.body).toEqualTypeOf<{ + barBody: number; + }>(); + } + + ctx.broadcast("updated", { count: 1 }); + ctx.broadcast("pair", 1, "ok"); + // @ts-expect-error wrong payload shape + ctx.broadcast("updated", { count: "no" }); + // @ts-expect-error unknown event name + ctx.broadcast("missing", { count: 1 }); + }), + actions: {}, + }); + }); + + it("does not require explicit queue.next body generic for single-queue actors", () => { + type Decision = { approved: boolean; approver: string }; + actor({ + state: {}, + queues: { + decision: queue(), + }, + run: workflow(async (ctx) => { + const [message] = await ctx.queue.next("wait-decision", { + names: ["decision"], + }); + if (message) { + expectTypeOf(message.body).toEqualTypeOf(); + } + }), + actions: {}, + }); + }); + }); + + describe("database type inference", () => { + it("supports typed rows for c.db.execute", () => { + actor({ + state: {}, + db: db(), + actions: { + readFoo: async (c) => { + const rows = await c.db.execute<{ foo: string }>( + "SELECT foo FROM bar", + ); + expectTypeOf(rows).toEqualTypeOf>(); + }, + }, + }); + }); + }); }); diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver-file-system-native-sqlite.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver-file-system-native-sqlite.test.ts index ff665f20d2..38ba057752 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver-file-system-native-sqlite.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/driver-file-system-native-sqlite.test.ts @@ -14,13 +14,12 @@ runDriverTests({ join(__dirname, "../fixtures/driver-test-suite/registry.ts"), async () => { return { - driver: createFileSystemOrMemoryDriver( - true, - { - path: `/tmp/test-${crypto.randomUUID()}`, - useNativeSqlite: true, - }, - ), + driver: createFileSystemOrMemoryDriver( + true, + { + path: `/tmp/test-${crypto.randomUUID()}`, + }, + ), }; }, ); diff --git a/rivetkit-typescript/packages/rivetkit/tests/file-system-kv-migration.test.ts b/rivetkit-typescript/packages/rivetkit/tests/file-system-kv-migration.test.ts new file mode 100644 index 0000000000..7783d019c5 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/file-system-kv-migration.test.ts @@ -0,0 +1,128 @@ +import { + copyFileSync, + mkdirSync, + mkdtempSync, + readFileSync, + readdirSync, + rmSync, +} from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; +import { describe, expect, it } from "vitest"; +import { importNodeDependencies } from "@/utils/node"; +import { FileSystemGlobalState } from "@/drivers/file-system/global-state"; +import { loadSqliteRuntime } from "@/drivers/file-system/sqlite-runtime"; + +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); +const fixtureStateDir = join(__dirname, "fixtures", "legacy-kv", "state"); + +function makeStorageFromFixtures(): string { + const storageRoot = mkdtempSync(join(tmpdir(), "rivetkit-kv-migration-")); + + const stateDir = join(storageRoot, "state"); + const dbDir = join(storageRoot, "databases"); + const alarmsDir = join(storageRoot, "alarms"); + mkdirSync(stateDir, { recursive: true }); + mkdirSync(dbDir, { recursive: true }); + mkdirSync(alarmsDir, { recursive: true }); + + for (const fileName of readdirSync(fixtureStateDir)) { + copyFileSync(join(fixtureStateDir, fileName), join(stateDir, fileName)); + } + + return storageRoot; +} + +describe("file-system driver legacy KV startup migration", () => { + it("migrates legacy actor kvStorage into sqlite databases on startup", async () => { + importNodeDependencies(); + const storageRoot = makeStorageFromFixtures(); + try { + const actorOneStatePath = join(storageRoot, "state", "legacy-actor-one"); + const actorOneStateBefore = readFileSync(actorOneStatePath); + + const state = new FileSystemGlobalState({ + persist: true, + customPath: storageRoot, + useNativeSqlite: true, + }); + + const alpha = await state.kvBatchGet("legacy-actor-one", [ + encoder.encode("alpha"), + ]); + expect(alpha[0]).not.toBeNull(); + expect(decoder.decode(alpha[0] ?? new Uint8Array())).toBe("one"); + + const prefixed = await state.kvListPrefix( + "legacy-actor-one", + encoder.encode("prefix:"), + ); + expect(prefixed).toHaveLength(2); + + const sqliteRuntime = loadSqliteRuntime(); + const actorTwoDb = sqliteRuntime.open( + join(storageRoot, "databases", "legacy-actor-two.db"), + ); + const actorTwoRow = actorTwoDb.get<{ value: Uint8Array | ArrayBuffer }>( + "SELECT value FROM kv WHERE key = ?", + [encoder.encode("beta")], + ); + expect(actorTwoRow).toBeDefined(); + expect( + decoder.decode( + (actorTwoRow?.value as Uint8Array | ArrayBuffer) ?? new Uint8Array(), + ), + ).toBe("two"); + actorTwoDb.close(); + + // Migration must not mutate legacy state files. + const actorOneStateAfter = readFileSync(actorOneStatePath); + expect(Buffer.compare(actorOneStateBefore, actorOneStateAfter)).toBe(0); + } finally { + rmSync(storageRoot, { recursive: true, force: true }); + } + }); + + it("does not overwrite sqlite data when database is already populated", async () => { + importNodeDependencies(); + const storageRoot = makeStorageFromFixtures(); + try { + const sqliteRuntime = loadSqliteRuntime(); + const actorDbPath = join(storageRoot, "databases", "legacy-actor-one.db"); + const db = sqliteRuntime.open(actorDbPath); + db.exec(` + CREATE TABLE IF NOT EXISTS kv ( + key BLOB PRIMARY KEY NOT NULL, + value BLOB NOT NULL + ) + `); + db.run("INSERT OR REPLACE INTO kv (key, value) VALUES (?, ?)", [ + encoder.encode("alpha"), + encoder.encode("existing"), + ]); + db.close(); + + const state = new FileSystemGlobalState({ + persist: true, + customPath: storageRoot, + useNativeSqlite: true, + }); + void state; + const checkDb = sqliteRuntime.open(actorDbPath); + const alpha = checkDb.get<{ value: Uint8Array | ArrayBuffer }>( + "SELECT value FROM kv WHERE key = ?", + [encoder.encode("alpha")], + ); + expect(alpha).toBeDefined(); + expect( + decoder.decode( + (alpha?.value as Uint8Array | ArrayBuffer) ?? new Uint8Array(), + ), + ).toBe("existing"); + checkDb.close(); + } finally { + rmSync(storageRoot, { recursive: true, force: true }); + } + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/legacy-kv/state/legacy-actor-one b/rivetkit-typescript/packages/rivetkit/tests/fixtures/legacy-kv/state/legacy-actor-one new file mode 100644 index 0000000000..caac5e87ea Binary files /dev/null and b/rivetkit-typescript/packages/rivetkit/tests/fixtures/legacy-kv/state/legacy-actor-one differ diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/legacy-kv/state/legacy-actor-two b/rivetkit-typescript/packages/rivetkit/tests/fixtures/legacy-kv/state/legacy-actor-two new file mode 100644 index 0000000000..799a59984a Binary files /dev/null and b/rivetkit-typescript/packages/rivetkit/tests/fixtures/legacy-kv/state/legacy-actor-two differ diff --git a/rivetkit-typescript/packages/rivetkit/tests/registry-config-storage-path.test.ts b/rivetkit-typescript/packages/rivetkit/tests/registry-config-storage-path.test.ts new file mode 100644 index 0000000000..4d269f2bb2 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/registry-config-storage-path.test.ts @@ -0,0 +1,41 @@ +import { RegistryConfigSchema } from "@/registry/config"; +import { describe, expect, test } from "vitest"; + +describe.sequential("registry config storagePath", () => { + test("reads storagePath from RIVETKIT_STORAGE_PATH when unset in config", () => { + const previous = process.env.RIVETKIT_STORAGE_PATH; + try { + process.env.RIVETKIT_STORAGE_PATH = "/tmp/rivetkit-storage-env"; + const parsed = RegistryConfigSchema.parse({ + use: {}, + }); + + expect(parsed.storagePath).toBe("/tmp/rivetkit-storage-env"); + } finally { + if (previous === undefined) { + delete process.env.RIVETKIT_STORAGE_PATH; + } else { + process.env.RIVETKIT_STORAGE_PATH = previous; + } + } + }); + + test("config storagePath overrides RIVETKIT_STORAGE_PATH", () => { + const previous = process.env.RIVETKIT_STORAGE_PATH; + try { + process.env.RIVETKIT_STORAGE_PATH = "/tmp/rivetkit-storage-env"; + const parsed = RegistryConfigSchema.parse({ + use: {}, + storagePath: "/tmp/rivetkit-storage-config", + }); + + expect(parsed.storagePath).toBe("/tmp/rivetkit-storage-config"); + } finally { + if (previous === undefined) { + delete process.env.RIVETKIT_STORAGE_PATH; + } else { + process.env.RIVETKIT_STORAGE_PATH = previous; + } + } + }); +}); diff --git a/rivetkit-typescript/packages/sqlite-vfs/src/vfs.ts b/rivetkit-typescript/packages/sqlite-vfs/src/vfs.ts index 6bbca6e8e3..c811c54574 100644 --- a/rivetkit-typescript/packages/sqlite-vfs/src/vfs.ts +++ b/rivetkit-typescript/packages/sqlite-vfs/src/vfs.ts @@ -214,7 +214,7 @@ export class SqliteVfs { // Initialize wa-sqlite module - each instance gets its own module const module = await SQLiteESMFactory({ wasmBinary }); - this.#sqlite3 = Factory(module) as SQLite3Api; + this.#sqlite3 = Factory(module) as unknown as SQLite3Api; // Create and register VFS with unique name this.#sqliteSystem = new SqliteSystem(this.#sqlite3, `kv-vfs-${this.#instanceId}`); diff --git a/rivetkit-typescript/packages/workflow-engine/QUICKSTART.md b/rivetkit-typescript/packages/workflow-engine/QUICKSTART.md index 5875d2f870..f0ad053745 100644 --- a/rivetkit-typescript/packages/workflow-engine/QUICKSTART.md +++ b/rivetkit-typescript/packages/workflow-engine/QUICKSTART.md @@ -28,7 +28,9 @@ async function orderWorkflow(ctx: WorkflowContextInterface, orderId: string) { }); // Wait for shipping confirmation (external message) - const tracking = await ctx.listen("wait-shipping", "shipment-confirmed"); + const [tracking] = await ctx.queue.next("wait-shipping", { + names: ["shipment-confirmed"], + }); // Send notification await ctx.step("notify-customer", async () => { @@ -213,46 +215,34 @@ await ctx.sleepUntil("wait-midnight", midnightTimestamp); Short sleeps (< `driver.workerPollInterval`) wait in memory. Longer sleeps yield to the scheduler and set an alarm for wake-up. -### Messages (Listen) +### Queue Messages Wait for external events delivered via `handle.message()`. ```typescript // Wait for a single message -const data = await ctx.listen("payment", "payment-completed"); +const [data] = await ctx.queue.next("payment", { + names: ["payment-completed"], +}); // Wait for N messages -const items = await ctx.listenN("batch", "item-added", 10); - -// Wait with timeout (returns null on timeout) -const result = await ctx.listenWithTimeout( - "api-response", - "response-received", - 30000 -); +const items = await ctx.queue.next("batch", { + names: ["item-added"], + count: 10, +}); -// Wait until timestamp (returns null on timeout) -const result = await ctx.listenUntil( - "api-response", - "response-received", - deadline -); +// Wait with timeout (returns [] on timeout) +const result = await ctx.queue.next("api-response", { + names: ["response-received"], + timeout: 30000, +}); // Wait for up to N messages with timeout -const items = await ctx.listenNWithTimeout( - "batch", - "item-added", - 10, // max items - 60000 // timeout ms -); - -// Wait for up to N messages until timestamp -const items = await ctx.listenNUntil( - "batch", - "item-added", - 10, // max items - deadline // timestamp -); +const timedBatch = await ctx.queue.next("batch-timeout", { + names: ["item-added"], + count: 10, + timeout: 60000, +}); ``` **Message delivery:** Messages are loaded once at workflow start. If a message is sent during execution, the workflow yields and picks it up on the next run. @@ -429,7 +419,7 @@ type WorkflowState = ## Best Practices -1. **Unique step names** - Each step/loop/sleep/listen within a scope must have a unique name +1. **Unique step names** - Each step/loop/sleep/queue.next within a scope must have a unique name 2. **Deterministic code** - Workflow code outside of steps must be deterministic. Don't use `Math.random()`, `Date.now()`, or read external state outside steps. diff --git a/rivetkit-typescript/packages/workflow-engine/architecture.md b/rivetkit-typescript/packages/workflow-engine/architecture.md index 59e2f6f9a6..93e60c2fc3 100644 --- a/rivetkit-typescript/packages/workflow-engine/architecture.md +++ b/rivetkit-typescript/packages/workflow-engine/architecture.md @@ -143,7 +143,7 @@ This optimization reduces storage size when the same names appear many times. ┌─────────────────────────────────────────────────────────────┐ │ WorkflowContextImpl │ │ Implements WorkflowContext interface │ -│ - step(), loop(), sleep(), listen(), join(), race() │ +│ - step(), loop(), sleep(), queue.next(), join(), race() │ │ - Manages current location │ │ - Creates branch contexts for parallel execution │ └─────────────────────────────────────────────────────────────┘ @@ -202,7 +202,7 @@ This optimization reduces storage size when the same names appear many times. ### Sleep/Message Yielding ``` -1. ctx.sleep() or ctx.listen() called +1. ctx.sleep() or ctx.queue.next() called 2. Check if deadline passed or message available (in memory) 3. If not ready: a. Throw SleepError or MessageWaitError diff --git a/rivetkit-typescript/packages/workflow-engine/docs/control-flow.md b/rivetkit-typescript/packages/workflow-engine/docs/control-flow.md index 66625d5012..d51e1bab16 100644 --- a/rivetkit-typescript/packages/workflow-engine/docs/control-flow.md +++ b/rivetkit-typescript/packages/workflow-engine/docs/control-flow.md @@ -74,17 +74,19 @@ const { winner, value } = await ctx.race("timeout", [ ## Messages in Control Flow -`ctx.listen()` and its variants pause the workflow until a message arrives. Message names are part of history, so keep them stable and unique. +`ctx.queue.next()` pauses the workflow until matching messages arrive. Queue wait names are part of history, so keep them stable and unique. ```ts -const approval = await ctx.listen("approval", "approval-granted"); +const [approval] = await ctx.queue.next("approval", { + names: ["approval-granted"], +}); ``` Messages are loaded at workflow start. If a message arrives during execution, the workflow yields and picks it up on the next run. ## Best Practices -- Use stable names for steps, loops, joins, races, and listens. +- Use stable names for steps, loops, joins, races, and queue waits. - Keep all nondeterministic work inside steps. - Use loop state to avoid native `while`/`for` loops. - Handle cancellation via `ctx.abortSignal` in long-running branches. @@ -92,4 +94,4 @@ Messages are loaded at workflow start. If a message arrives during execution, th ## Related - `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:155` for loop usage. -- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:207` for message waits. +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:207` for queue waits. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md b/rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md index f52ee553a6..b0bce652df 100644 --- a/rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md +++ b/rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md @@ -4,11 +4,13 @@ Long-running workflows can pause, sleep, and resume across process restarts. Thi ## Yielding Execution -Use sleep and listen helpers to yield control while waiting: +Use sleep and queue helpers to yield control while waiting: ```ts await ctx.sleep("wait-5-min", 5 * 60 * 1000); -const message = await ctx.listen("wait-approval", "approval"); +const [message] = await ctx.queue.next("wait-approval", { + names: ["approval"], +}); ``` When a workflow yields, `runWorkflow` returns a `WorkflowResult` with `state: "sleeping"`. The driver alarm or a message wake-up triggers the next run. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md b/rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md index 4e8b751c5c..edfa100266 100644 --- a/rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md +++ b/rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md @@ -8,43 +8,49 @@ Workflows can pause until external events arrive. This enables human approvals, - Messages are loaded at workflow start. - If a message arrives while the workflow is running, the workflow yields and picks it up on the next run. -In live mode (`runWorkflow(..., { mode: "live" })`), incoming messages can also wake a workflow waiting on `ctx.listen()`. +In live mode (`runWorkflow(..., { mode: "live" })`), incoming messages can also wake a workflow waiting on `ctx.queue.next()`. -## Listening for Messages +## Waiting for Queue Messages ```ts -const approval = await ctx.listen("wait-approval", "approval-granted"); +const [approval] = await ctx.queue.next("wait-approval", { + names: ["approval-granted"], +}); ``` -Use the `listen*` variants to wait for multiple messages or apply timeouts: +Use `count` and `timeout` to wait for batches or apply deadlines: ```ts -const items = await ctx.listenN("batch", "item-added", 10); -const result = await ctx.listenWithTimeout("approval", "approval-granted", 60000); +const items = await ctx.queue.next("batch", { + names: ["item-added"], + count: 10, +}); +const result = await ctx.queue.next("approval", { + names: ["approval-granted"], + timeout: 60000, +}); ``` ## Deadlines and Timeouts -Use `listenUntil` or `listenWithTimeout` to model approval windows: +Use `timeout` to model approval windows: ```ts -const approval = await ctx.listenUntil( - "approval-window", - "approval-granted", - Date.now() + 24 * 60 * 60 * 1000, -); +const [approval] = await ctx.queue.next("approval-window", { + names: ["approval-granted"], + timeout: 24 * 60 * 60 * 1000, +}); ``` -If the deadline passes, the method returns `null` instead of throwing. +If the deadline passes, `ctx.queue.next(...)` returns `[]`. ## Human-in-the-Loop Example ```ts -const approval = await ctx.listenWithTimeout( - "manual-approval", - "approval-granted", - 30 * 60 * 1000, -); +const [approval] = await ctx.queue.next("manual-approval", { + names: ["approval-granted"], + timeout: 30 * 60 * 1000, +}); if (!approval) { await ctx.step("notify-timeout", () => sendTimeoutNotice()); @@ -62,5 +68,5 @@ await ctx.step("proceed", () => runApprovedWork()); ## Related -- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:207` for listen helpers. +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:207` for queue waits. - `rivetkit-typescript/packages/workflow-engine/architecture.md:218` for message delivery details. diff --git a/rivetkit-typescript/packages/workflow-engine/src/context.ts b/rivetkit-typescript/packages/workflow-engine/src/context.ts index d6e8d8a47b..3bb29e0fc1 100644 --- a/rivetkit-typescript/packages/workflow-engine/src/context.ts +++ b/rivetkit-typescript/packages/workflow-engine/src/context.ts @@ -24,14 +24,13 @@ import { registerName, } from "./location.js"; import { - consumeMessage, consumeMessages, createEntry, deleteEntriesWithPrefix, flush, - getEntry, getOrCreateMetadata, loadMetadata, + peekMessages, setEntry, } from "./storage.js"; import type { @@ -49,7 +48,9 @@ import type { StepConfig, Storage, WorkflowContextInterface, - WorkflowListenMessage, + WorkflowQueue, + WorkflowQueueMessage, + WorkflowQueueNextOptions, WorkflowMessageDriver, } from "./types.js"; import { sleep } from "./utils.js"; @@ -66,7 +67,7 @@ export const DEFAULT_LOOP_HISTORY_EVERY = 20; export const DEFAULT_LOOP_HISTORY_KEEP = 20; export const DEFAULT_STEP_TIMEOUT = 30000; // 30 seconds -const LISTEN_HISTORY_MESSAGE_MARKER = "__rivetWorkflowListenMessage"; +const QUEUE_HISTORY_MESSAGE_MARKER = "__rivetWorkflowQueueMessage"; /** * Calculate backoff delay with exponential backoff. @@ -113,6 +114,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface { private rollbackCheckpointSet: boolean; /** Track names used in current execution to detect duplicates */ private usedNamesInExecution = new Set(); + private pendingCompletableMessageIds = new Set(); private historyNotifier?: () => void; private logger?: Logger; @@ -142,6 +144,13 @@ export class WorkflowContextImpl implements WorkflowContextInterface { return this.abortController.signal; } + get queue(): WorkflowQueue { + return { + next: async (name, opts) => await this.queueNext(name, opts), + send: async (name, body) => await this.queueSend(name, body), + }; + } + isEvicted(): boolean { return this.abortSignal.aborted; } @@ -206,12 +215,12 @@ export class WorkflowContextImpl implements WorkflowContextInterface { private checkDuplicateName(name: string): void { const fullKey = locationToKey(this.storage, this.currentLocation) + "/" + name; - if (this.usedNamesInExecution.has(fullKey)) { - throw new HistoryDivergedError( - `Duplicate entry name "${name}" at location "${locationToKey(this.storage, this.currentLocation)}". ` + - `Each step/loop/sleep/listen/join/race must have a unique name within its scope.`, - ); - } + if (this.usedNamesInExecution.has(fullKey)) { + throw new HistoryDivergedError( + `Duplicate entry name "${name}" at location "${locationToKey(this.storage, this.currentLocation)}". ` + + `Each step/loop/sleep/queue.next/join/race must have a unique name within its scope.`, + ); + } this.usedNamesInExecution.add(fullKey); } @@ -953,66 +962,54 @@ export class WorkflowContextImpl implements WorkflowContextInterface { this.rollbackCheckpointSet = true; } - // === Listen === - // - // IMPORTANT: Messages are loaded once at workflow start (in loadStorage). - // If a message is sent via handle.message() DURING workflow execution, - // it won't be visible until the next execution. The workflow will yield - // (SleepError/MessageWaitError), then on the next run, loadStorage() will - // pick up the new message. This is intentional - no polling during execution. - - async listen( - name: string, - messageName: string | string[], - ): Promise> { - this.assertNotInProgress(); - this.checkEvicted(); + // === Queue === - this.entryInProgress = true; - try { - const messages = await this.executeListenN(name, messageName, 1); - const message = messages[0]; - if (!message) { - throw new HistoryDivergedError("Expected message for listen()"); - } - return this.toListenMessage(message); - } finally { - this.entryInProgress = false; + private async queueSend(name: string, body: unknown): Promise { + const message: Message = { + id: crypto.randomUUID(), + name, + data: body, + sentAt: Date.now(), + }; + if (!this.messageDriver.receiveMessages) { + this.storage.messages.push(message); } + await this.messageDriver.addMessage(message); } - async listenN( + private async queueNext( name: string, - messageName: string, - limit: number, - ): Promise { + opts?: WorkflowQueueNextOptions, + ): Promise>> { this.assertNotInProgress(); this.checkEvicted(); this.entryInProgress = true; try { - const messages = await this.executeListenN(name, messageName, limit); - await Promise.all( - messages.map((message) => this.completeConsumedMessage(message)), - ); - return messages.map((message) => message.data as T); + return await this.executeQueueNext(name, opts); } finally { this.entryInProgress = false; } } - private async executeListenN( + private async executeQueueNext( name: string, - messageName: string | string[], - limit: number, - ): Promise { - const messageNames = this.normalizeMessageNames(messageName); + opts?: WorkflowQueueNextOptions, + ): Promise>> { + if (this.pendingCompletableMessageIds.size > 0) { + throw new Error( + "Previous completable queue message is not completed. Call `message.complete(...)` before receiving the next message.", + ); + } + + const resolvedOpts = opts ?? {}; + const messageNames = this.normalizeQueueNames(resolvedOpts.names); const messageNameLabel = this.messageNamesLabel(messageNames); + const count = Math.max(1, resolvedOpts.count ?? 1); + const completable = resolvedOpts.completable === true; - // Check for duplicate name in current execution this.checkDuplicateName(name); - // Check for replay: first check if we have a count entry const countLocation = appendName( this.storage, this.currentLocation, @@ -1020,99 +1017,103 @@ export class WorkflowContextImpl implements WorkflowContextInterface { ); const countKey = locationToKey(this.storage, countLocation); const existingCount = this.storage.history.entries.get(countKey); - - // Mark count entry as visited this.markVisited(countKey); - this.stopRollbackIfMissing(existingCount); - if (existingCount && existingCount.kind.type === "message") { - // Replay: read all recorded messages - const count = existingCount.kind.data.data as number; - const results: Message[] = []; - - for (let i = 0; i < count; i++) { - const messageLocation = appendName( - this.storage, - this.currentLocation, - `${name}:${i}`, - ); - const messageKey = locationToKey(this.storage, messageLocation); + let deadline: number | undefined; + let deadlineEntry: Entry | undefined; + if (resolvedOpts.timeout !== undefined) { + const deadlineLocation = appendName( + this.storage, + this.currentLocation, + `${name}:deadline`, + ); + const deadlineKey = locationToKey(this.storage, deadlineLocation); + deadlineEntry = this.storage.history.entries.get(deadlineKey); + this.markVisited(deadlineKey); + this.stopRollbackIfMissing(deadlineEntry); + + if (deadlineEntry && deadlineEntry.kind.type === "sleep") { + deadline = deadlineEntry.kind.data.deadline; + } else { + deadline = Date.now() + resolvedOpts.timeout; + const created = createEntry(deadlineLocation, { + type: "sleep", + data: { deadline, state: "pending" }, + }); + setEntry(this.storage, deadlineLocation, created); + created.dirty = true; + await this.flushStorage(); + deadlineEntry = created; + } + } - // Mark each message entry as visited - this.markVisited(messageKey); + if (existingCount && existingCount.kind.type === "message") { + const replayCount = existingCount.kind.data.data as number; + return await this.readReplayQueueMessages( + name, + replayCount, + completable, + ); + } - const existingMessage = - this.storage.history.entries.get(messageKey); - if ( - existingMessage && - existingMessage.kind.type === "message" - ) { - results.push( - this.fromHistoryListenMessage( - existingMessage.kind.data.name, - existingMessage.kind.data.data, - ), - ); - } + const now = Date.now(); + if (deadline !== undefined && now >= deadline) { + if (deadlineEntry && deadlineEntry.kind.type === "sleep") { + deadlineEntry.kind.data.state = "completed"; + deadlineEntry.dirty = true; } - - return results; + await this.recordQueueCountEntry( + countLocation, + `${messageNameLabel}:count`, + 0, + ); + return []; } - // Try to consume messages immediately - const messages = await consumeMessages( - this.storage, - this.messageDriver, + const received = await this.receiveMessagesNow( messageNames, - limit, + count, + completable, ); - - if (messages.length > 0) { - // Record each message in history with indexed names - for (let i = 0; i < messages.length; i++) { - const messageLocation = appendName( - this.storage, - this.currentLocation, - `${name}:${i}`, - ); - const messageEntry = createEntry(messageLocation, { - type: "message", - data: { - name: messages[i].name, - data: this.toHistoryListenMessage(messages[i]), - }, - }); - setEntry(this.storage, messageLocation, messageEntry); - - // Mark as visited - this.markVisited(locationToKey(this.storage, messageLocation)); + if (received.length > 0) { + const historyMessages = received.map((message) => + this.toWorkflowQueueMessage(message), + ); + if (deadlineEntry && deadlineEntry.kind.type === "sleep") { + deadlineEntry.kind.data.state = "interrupted"; + deadlineEntry.dirty = true; } - - // Record the count for replay - const countEntry = createEntry(countLocation, { - type: "message", - data: { - name: `${messageNameLabel}:count`, - data: messages.length, - }, - }); - setEntry(this.storage, countLocation, countEntry); - - await this.flushStorage(); - - return messages; + await this.recordQueueMessages( + name, + countLocation, + messageNames, + historyMessages, + ); + const queueMessages = received.map((message, index) => + this.createQueueMessage(message, completable, { + historyLocation: appendName( + this.storage, + this.currentLocation, + `${name}:${index}`, + ), + }), + ); + return queueMessages; } - // No messages found, throw to yield to scheduler - throw new MessageWaitError(messageNames); + if (deadline === undefined) { + throw new MessageWaitError(messageNames); + } + throw new SleepError(deadline, messageNames); } - private normalizeMessageNames(messageName: string | string[]): string[] { - const names = Array.isArray(messageName) ? messageName : [messageName]; + private normalizeQueueNames(names?: readonly string[]): string[] { + if (!names || names.length === 0) { + return []; + } const deduped: string[] = []; const seen = new Set(); - for (const name of names) { if (seen.has(name)) { continue; @@ -1120,466 +1121,280 @@ export class WorkflowContextImpl implements WorkflowContextInterface { seen.add(name); deduped.push(name); } - - if (deduped.length === 0) { - throw new Error("listen() requires at least one message name"); - } - return deduped; } private messageNamesLabel(messageNames: string[]): string { + if (messageNames.length === 0) { + return "*"; + } return messageNames.length === 1 ? messageNames[0] : messageNames.join("|"); } - private toListenMessage(message: Message): WorkflowListenMessage { - return { - id: message.id, - name: message.name, - body: message.data as T, - complete: async (response?: unknown) => { - if (message.complete) { - await message.complete(response); - return; - } - if (this.messageDriver.completeMessage) { - await this.messageDriver.completeMessage(message.id, response); - } - }, - }; - } - - private async completeConsumedMessage(message: Message): Promise { - if (message.complete) { - await message.complete(); - return; - } - if (message.id && this.messageDriver.completeMessage) { - await this.messageDriver.completeMessage(message.id); - } - } - - private toHistoryListenMessage(message: Message): unknown { - return { - [LISTEN_HISTORY_MESSAGE_MARKER]: 1, - id: message.id, - name: message.name, - body: message.data, - }; - } - - private fromHistoryListenMessage(name: string, value: unknown): Message { - if ( - typeof value === "object" && - value !== null && - (value as Record)[LISTEN_HISTORY_MESSAGE_MARKER] === 1 - ) { - const serialized = value as Record; - const id = - typeof serialized.id === "string" ? serialized.id : ""; - const serializedName = - typeof serialized.name === "string" ? serialized.name : name; - const complete = async (response?: unknown) => { - if (!id || !this.messageDriver.completeMessage) { - return; - } - await this.messageDriver.completeMessage(id, response); - }; - - return { - id, - name: serializedName, - data: serialized.body, - sentAt: 0, - complete, - }; - } - - return { - id: "", - name, - data: value, - sentAt: 0, - }; - } - - async listenWithTimeout( - name: string, - messageName: string, - timeoutMs: number, - ): Promise { - const deadline = Date.now() + timeoutMs; - return this.listenUntil(name, messageName, deadline); - } - - async listenUntil( - name: string, - messageName: string, - timestampMs: number, - ): Promise { - this.assertNotInProgress(); - this.checkEvicted(); - - this.entryInProgress = true; - try { - return await this.executeListenUntil( - name, - messageName, - timestampMs, - ); - } finally { - this.entryInProgress = false; - } - } - - private async executeListenUntil( - name: string, - messageName: string, - deadline: number, - ): Promise { - // Check for duplicate name in current execution - this.checkDuplicateName(name); - - const sleepLocation = appendName( - this.storage, - this.currentLocation, - name, - ); - const messageLocation = appendName( - this.storage, - this.currentLocation, - `${name}:message`, - ); - const sleepKey = locationToKey(this.storage, sleepLocation); - const messageKey = locationToKey(this.storage, messageLocation); - - // Mark entries as visited for validateComplete - this.markVisited(sleepKey); - this.markVisited(messageKey); - - const existingSleep = this.storage.history.entries.get(sleepKey); - - this.stopRollbackIfMissing(existingSleep); - - // Check for replay - if (existingSleep && existingSleep.kind.type === "sleep") { - const sleepData = existingSleep.kind.data; - if (sleepData.state === "completed") { - return null; - } - - if (sleepData.state === "interrupted") { - const existingMessage = this.storage.history.entries.get(messageKey); - if ( - existingMessage && - existingMessage.kind.type === "message" - ) { - const replayedMessage = this.fromHistoryListenMessage( - existingMessage.kind.data.name, - existingMessage.kind.data.data, - ); - await this.completeConsumedMessage(replayedMessage); - return replayedMessage.data as T; - } - throw new HistoryDivergedError( - "Expected message entry after interrupted sleep", - ); - } - - this.stopRollbackIfIncomplete(true); - - deadline = sleepData.deadline; - } else { - this.stopRollbackIfIncomplete(true); - - // Create sleep entry - const sleepEntry = createEntry(sleepLocation, { - type: "sleep", - data: { deadline, state: "pending" }, + private async receiveMessagesNow( + messageNames: string[], + count: number, + completable: boolean, + ): Promise { + if (this.messageDriver.receiveMessages) { + return await this.messageDriver.receiveMessages({ + names: messageNames.length > 0 ? messageNames : undefined, + count, + completable, }); - setEntry(this.storage, sleepLocation, sleepEntry); - sleepEntry.dirty = true; - await this.flushStorage(); } - - const now = Date.now(); - const remaining = deadline - now; - - // Deadline passed, check for message one more time - if (remaining <= 0) { - const message = await consumeMessage( + if (completable) { + return peekMessages( this.storage, - this.messageDriver, - messageName, + messageNames.length > 0 ? messageNames : [], + count, ); - const sleepEntry = getEntry(this.storage, sleepLocation)!; - - if (message) { - if (sleepEntry.kind.type === "sleep") { - sleepEntry.kind.data.state = "interrupted"; - } - sleepEntry.dirty = true; - - const messageEntry = createEntry(messageLocation, { - type: "message", - data: { - name: message.name, - data: this.toHistoryListenMessage(message), - }, - }); - setEntry(this.storage, messageLocation, messageEntry); - await this.flushStorage(); - await this.completeConsumedMessage(message); - - return message.data as T; - } - - if (sleepEntry.kind.type === "sleep") { - sleepEntry.kind.data.state = "completed"; - } - sleepEntry.dirty = true; - await this.flushStorage(); - return null; } - - // Check for message (messages are loaded at workflow start, no polling needed) - const message = await consumeMessage( + return await consumeMessages( this.storage, this.messageDriver, - messageName, + messageNames.length > 0 ? messageNames : [], + count, ); - if (message) { - const sleepEntry = getEntry(this.storage, sleepLocation)!; - if (sleepEntry.kind.type === "sleep") { - sleepEntry.kind.data.state = "interrupted"; - } - sleepEntry.dirty = true; + } + private async recordQueueMessages( + name: string, + countLocation: Location, + messageNames: string[], + messages: Array>, + ): Promise { + for (let i = 0; i < messages.length; i++) { + const messageLocation = appendName( + this.storage, + this.currentLocation, + `${name}:${i}`, + ); const messageEntry = createEntry(messageLocation, { type: "message", data: { - name: message.name, - data: this.toHistoryListenMessage(message), + name: messages[i].name, + data: this.toHistoryQueueMessage(messages[i]), }, }); setEntry(this.storage, messageLocation, messageEntry); - await this.flushStorage(); - await this.completeConsumedMessage(message); - - return message.data as T; + this.markVisited(locationToKey(this.storage, messageLocation)); } + await this.recordQueueCountEntry( + countLocation, + `${this.messageNamesLabel(messageNames)}:count`, + messages.length, + ); + } - // Message not available, yield to scheduler until deadline or message - throw new SleepError(deadline, [messageName]); + private async recordQueueCountEntry( + countLocation: Location, + countLabel: string, + count: number, + ): Promise { + const countEntry = createEntry(countLocation, { + type: "message", + data: { + name: countLabel, + data: count, + }, + }); + setEntry(this.storage, countLocation, countEntry); + await this.flushStorage(); } - async listenNWithTimeout( + private async readReplayQueueMessages( name: string, - messageName: string, - limit: number, - timeoutMs: number, - ): Promise { - this.assertNotInProgress(); - this.checkEvicted(); - - this.entryInProgress = true; - try { - return await this.executeListenNWithTimeout( - name, - messageName, - limit, - timeoutMs, + count: number, + completable: boolean, + ): Promise>> { + const results: Array> = []; + for (let i = 0; i < count; i++) { + const messageLocation = appendName( + this.storage, + this.currentLocation, + `${name}:${i}`, + ); + const messageKey = locationToKey(this.storage, messageLocation); + this.markVisited(messageKey); + const existingMessage = this.storage.history.entries.get(messageKey); + if (!existingMessage || existingMessage.kind.type !== "message") { + throw new HistoryDivergedError( + `Expected queue message "${name}:${i}" in history`, + ); + } + const parsed = this.fromHistoryQueueMessage( + existingMessage.kind.data.name, + existingMessage.kind.data.data, + ); + results.push( + this.createQueueMessage(parsed.message, completable, { + historyLocation: messageLocation, + completed: parsed.completed, + replay: true, + }), ); - } finally { - this.entryInProgress = false; } + return results; } - private async executeListenNWithTimeout( - name: string, - messageName: string, - limit: number, - timeoutMs: number, - ): Promise { - // Check for duplicate name in current execution - this.checkDuplicateName(name); - - // Use a sleep entry to store the deadline for replay - const sleepLocation = appendName( - this.storage, - this.currentLocation, - `${name}:deadline`, - ); - const sleepKey = locationToKey(this.storage, sleepLocation); - const existingSleep = this.storage.history.entries.get(sleepKey); + private toWorkflowQueueMessage(message: Message): WorkflowQueueMessage { + return { + id: message.id, + name: message.name, + body: message.data as T, + createdAt: message.sentAt, + }; + } - this.markVisited(sleepKey); + private createQueueMessage( + message: Message, + completable: boolean, + opts?: { + historyLocation?: Location; + completed?: boolean; + replay?: boolean; + }, + ): WorkflowQueueMessage { + const queueMessage = this.toWorkflowQueueMessage(message); + if (!completable) { + return queueMessage; + } + + if (opts?.replay && opts.completed) { + return { + ...queueMessage, + complete: async () => { + // No-op: this message was already completed in a prior run. + }, + }; + } - this.stopRollbackIfMissing(existingSleep); + const messageId = message.id; + this.pendingCompletableMessageIds.add(messageId); + let completed = false; - let deadline: number; + return { + ...queueMessage, + complete: async (response?: unknown) => { + if (completed) { + throw new Error("Queue message already completed"); + } + completed = true; + try { + await this.completeMessage(message, response); + await this.markQueueMessageCompleted(opts?.historyLocation); + this.pendingCompletableMessageIds.delete(messageId); + } catch (error) { + completed = false; + throw error; + } + }, + }; + } - if (existingSleep && existingSleep.kind.type === "sleep") { - // Replay: use stored deadline - deadline = existingSleep.kind.data.deadline; - } else { - // New execution: calculate and store deadline - deadline = Date.now() + timeoutMs; - const sleepEntry = createEntry(sleepLocation, { - type: "sleep", - data: { deadline, state: "pending" }, - }); - setEntry(this.storage, sleepLocation, sleepEntry); - sleepEntry.dirty = true; - // Flush immediately to persist deadline before potential SleepError - await this.flushStorage(); + private async markQueueMessageCompleted( + historyLocation: Location | undefined, + ): Promise { + if (!historyLocation) { + return; } - - return this.executeListenNUntilImpl( - name, - messageName, - limit, - deadline, + const key = locationToKey(this.storage, historyLocation); + const entry = this.storage.history.entries.get(key); + if (!entry || entry.kind.type !== "message") { + return; + } + const parsed = this.fromHistoryQueueMessage( + entry.kind.data.name, + entry.kind.data.data, + ); + entry.kind.data.data = this.toHistoryQueueMessage( + this.toWorkflowQueueMessage(parsed.message), + true, ); + entry.dirty = true; + await this.flushStorage(); } - async listenNUntil( - name: string, - messageName: string, - limit: number, - timestampMs: number, - ): Promise { - this.assertNotInProgress(); - this.checkEvicted(); - - // Check for duplicate name in current execution - this.checkDuplicateName(name); - - this.entryInProgress = true; - try { - return await this.executeListenNUntilImpl( - name, - messageName, - limit, - timestampMs, - ); - } finally { - this.entryInProgress = false; + private async completeMessage( + message: Message, + response?: unknown, + ): Promise { + if (message.complete) { + await message.complete(response); + return; } - } - - /** - * Internal implementation for listenNUntil with proper replay support. - * Stores the count and individual messages for deterministic replay. - */ - private async executeListenNUntilImpl( - name: string, - messageName: string, - limit: number, - deadline: number, - ): Promise { - // Check for replay: look for count entry - const countLocation = appendName( - this.storage, - this.currentLocation, - `${name}:count`, + if (this.messageDriver.completeMessage) { + await this.messageDriver.completeMessage(message.id, response); + return; + } + const deleted = await this.messageDriver.deleteMessages([message.id]); + if (!deleted.includes(message.id)) { + return; + } + const idx = this.storage.messages.findIndex((entry) => + entry.id === message.id ); - const countKey = locationToKey(this.storage, countLocation); - const existingCount = this.storage.history.entries.get(countKey); - - this.markVisited(countKey); - - this.stopRollbackIfMissing(existingCount); - - if (existingCount && existingCount.kind.type === "message") { - // Replay: read all recorded messages - const count = existingCount.kind.data.data as number; - const results: T[] = []; - - for (let i = 0; i < count; i++) { - const messageLocation = appendName( - this.storage, - this.currentLocation, - `${name}:${i}`, - ); - const messageKey = locationToKey(this.storage, messageLocation); - - this.markVisited(messageKey); - - const existingMessage = this.storage.history.entries.get(messageKey); - if ( - existingMessage && - existingMessage.kind.type === "message" - ) { - const replayedMessage = this.fromHistoryListenMessage( - existingMessage.kind.data.name, - existingMessage.kind.data.data, - ); - await this.completeConsumedMessage(replayedMessage); - results.push(replayedMessage.data as T); - } - } - - return results; + if (idx !== -1) { + this.storage.messages.splice(idx, 1); } + } - // New execution: collect messages until timeout or limit reached - const results: T[] = []; - - for (let i = 0; i < limit; i++) { - const now = Date.now(); - if (now >= deadline) { - break; - } - - // Try to consume a message - const message = await consumeMessage( - this.storage, - this.messageDriver, - messageName, - ); - if (!message) { - // No message available - check if we should wait - if (results.length === 0) { - // No messages yet - yield to scheduler until deadline or message - throw new SleepError(deadline, [messageName]); - } - // We have some messages - return what we have - break; - } + private toHistoryQueueMessage( + message: WorkflowQueueMessage, + completed = false, + ): unknown { + return { + [QUEUE_HISTORY_MESSAGE_MARKER]: 1, + id: message.id, + name: message.name, + body: message.body, + createdAt: message.createdAt, + completed, + }; + } - // Record the message - const messageLocation = appendName( - this.storage, - this.currentLocation, - `${name}:${i}`, - ); - const messageEntry = createEntry(messageLocation, { - type: "message", - data: { - name: message.name, - data: this.toHistoryListenMessage(message), + private fromHistoryQueueMessage( + name: string, + value: unknown, + ): { message: Message; completed: boolean } { + if ( + typeof value === "object" && + value !== null && + (value as Record)[QUEUE_HISTORY_MESSAGE_MARKER] === 1 + ) { + const serialized = value as Record; + const id = + typeof serialized.id === "string" ? serialized.id : ""; + const serializedName = + typeof serialized.name === "string" ? serialized.name : name; + const createdAt = + typeof serialized.createdAt === "number" ? serialized.createdAt : 0; + const completed = + typeof serialized.completed === "boolean" + ? serialized.completed + : false; + return { + message: { + id, + name: serializedName, + data: serialized.body, + sentAt: createdAt, }, - }); - setEntry(this.storage, messageLocation, messageEntry); - this.markVisited(locationToKey(this.storage, messageLocation)); - await this.completeConsumedMessage(message); - - results.push(message.data as T); + completed, + }; } - - // Record the count for replay - const countEntry = createEntry(countLocation, { - type: "message", - data: { name: `${messageName}:count`, data: results.length }, - }); - setEntry(this.storage, countLocation, countEntry); - - await this.flushStorage(); - - return results; + return { + message: { + id: "", + name, + data: value, + sentAt: 0, + }, + completed: false, + }; } // === Join === diff --git a/rivetkit-typescript/packages/workflow-engine/src/index.ts b/rivetkit-typescript/packages/workflow-engine/src/index.ts index e9f231a3c8..511a6a1109 100644 --- a/rivetkit-typescript/packages/workflow-engine/src/index.ts +++ b/rivetkit-typescript/packages/workflow-engine/src/index.ts @@ -100,7 +100,9 @@ export type { WorkflowContextInterface, WorkflowFunction, WorkflowHandle, - WorkflowListenMessage, + WorkflowQueue, + WorkflowQueueMessage, + WorkflowQueueNextOptions, WorkflowMessageDriver, WorkflowResult, WorkflowRunMode, @@ -200,9 +202,13 @@ function notifyMessage(runtime: LiveRuntime, name: string): void { for (let i = 0; i < runtime.messageWaiters.length; i++) { const waiter = runtime.messageWaiters[i]; - const matchIndex = runtime.pendingMessageNames.findIndex((pending) => - waiter.names.includes(pending), - ); + const matchIndex = waiter.names.length === 0 + ? runtime.pendingMessageNames.length > 0 + ? 0 + : -1 + : runtime.pendingMessageNames.findIndex((pending) => + waiter.names.includes(pending) + ); if (matchIndex !== -1) { runtime.pendingMessageNames.splice(matchIndex, 1); runtime.messageWaiters.splice(i, 1); @@ -221,9 +227,13 @@ async function waitForMessage( throw new EvictedError(); } - const matchIndex = runtime.pendingMessageNames.findIndex((pending) => - names.includes(pending), - ); + const matchIndex = names.length === 0 + ? runtime.pendingMessageNames.length > 0 + ? 0 + : -1 + : runtime.pendingMessageNames.findIndex((pending) => + names.includes(pending) + ); if (matchIndex !== -1) { runtime.pendingMessageNames.splice(matchIndex, 1); return; @@ -525,12 +535,11 @@ async function executeLiveWorkflow( return result; } - const hasMessages = - result.waitingForMessages && result.waitingForMessages.length > 0; - const hasDeadline = result.sleepUntil !== undefined; + const hasMessages = result.waitingForMessages !== undefined; + const hasDeadline = result.sleepUntil !== undefined; if (hasMessages && hasDeadline) { - // Wait for EITHER a message OR the deadline (for listenWithTimeout) + // Wait for EITHER a message OR the deadline (for queue.next timeout) try { const messagePromise = driver.waitForMessages ? awaitWithEviction( diff --git a/rivetkit-typescript/packages/workflow-engine/src/storage.ts b/rivetkit-typescript/packages/workflow-engine/src/storage.ts index 8c6fff7965..62c3671f19 100644 --- a/rivetkit-typescript/packages/workflow-engine/src/storage.ts +++ b/rivetkit-typescript/packages/workflow-engine/src/storage.ts @@ -398,6 +398,31 @@ export async function consumeMessage( return messages[0] ?? null; } +/** + * Peek up to N messages from the queue without consuming them. + */ +export function peekMessages( + storage: Storage, + messageName: string | string[], + limit: number, +): Message[] { + const messageNameSet = new Set( + Array.isArray(messageName) ? messageName : [messageName], + ); + const includeAll = messageNameSet.size === 0; + const results: Message[] = []; + for (const message of storage.messages) { + if (!includeAll && !messageNameSet.has(message.name)) { + continue; + } + results.push(message); + if (results.length >= limit) { + break; + } + } + return results; +} + /** * Consume up to N messages from the queue. * @@ -415,13 +440,14 @@ export async function consumeMessages( const messageNameSet = new Set( Array.isArray(messageName) ? messageName : [messageName], ); + const includeAll = messageNameSet.size === 0; // Find all matching messages up to limit (don't modify memory yet) const toConsume: { message: Message; index: number }[] = []; let count = 0; for (let i = 0; i < storage.messages.length && count < limit; i++) { - if (messageNameSet.has(storage.messages[i].name)) { + if (includeAll || messageNameSet.has(storage.messages[i].name)) { toConsume.push({ message: storage.messages[i], index: i }); count++; } diff --git a/rivetkit-typescript/packages/workflow-engine/src/types.ts b/rivetkit-typescript/packages/workflow-engine/src/types.ts index 9b2d918b35..335d249d68 100644 --- a/rivetkit-typescript/packages/workflow-engine/src/types.ts +++ b/rivetkit-typescript/packages/workflow-engine/src/types.ts @@ -204,13 +204,45 @@ export interface Message { } /** - * Message handle returned by listen(). + * Options for receiving queue messages in workflows. */ -export interface WorkflowListenMessage { - id: string; +export interface WorkflowQueueNextOptions { + /** + * Queue names to receive from. + * If omitted, receives from all queue names. + */ + names?: readonly string[]; + /** Maximum number of messages to receive. Defaults to 1. */ + count?: number; + /** + * Timeout in milliseconds. + * Omit to wait indefinitely. + */ + timeout?: number; + /** Whether returned messages must be manually completed. */ + completable?: boolean; +} + +/** + * Message returned by workflow queue operations. + */ +export interface WorkflowQueueMessage { + id: string | bigint; name: string; - body: T; - complete(response?: unknown): Promise; + body: TBody; + createdAt: number; + complete?(response?: unknown): Promise; +} + +/** + * Workflow queue interface. + */ +export interface WorkflowQueue { + next( + name: string, + opts?: WorkflowQueueNextOptions, + ): Promise>>; + send(name: string, body: unknown): Promise; } /** @@ -276,6 +308,17 @@ export interface Storage { export interface WorkflowMessageDriver { loadMessages(): Promise; addMessage(message: Message): Promise; + /** + * Optionally receive messages directly from the host queue implementation. + * This is used by actor-backed workflows to reuse native queue behavior. + * + * The operation must be non-blocking and return immediately. + */ + receiveMessages?(opts: { + names?: readonly string[]; + count: number; + completable: boolean; + }): Promise; /** * Delete the specified messages and return the IDs that were successfully removed. */ @@ -353,6 +396,7 @@ export type BranchOutput = T extends BranchConfig ? O : never; export interface WorkflowContextInterface { readonly workflowId: string; readonly abortSignal: AbortSignal; + readonly queue: WorkflowQueue; step(name: string, run: () => Promise): Promise; step(config: StepConfig): Promise; @@ -368,34 +412,6 @@ export interface WorkflowContextInterface { sleep(name: string, durationMs: number): Promise; sleepUntil(name: string, timestampMs: number): Promise; - listen( - name: string, - messageName: string | string[], - ): Promise>; - listenN(name: string, messageName: string, limit: number): Promise; - listenWithTimeout( - name: string, - messageName: string, - timeoutMs: number, - ): Promise; - listenUntil( - name: string, - messageName: string, - timestampMs: number, - ): Promise; - listenNWithTimeout( - name: string, - messageName: string, - limit: number, - timeoutMs: number, - ): Promise; - listenNUntil( - name: string, - messageName: string, - limit: number, - timestampMs: number, - ): Promise; - rollbackCheckpoint(name: string): Promise; join>>( diff --git a/rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts index 149da0baa0..7a166fe6c9 100644 --- a/rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts +++ b/rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts @@ -18,7 +18,12 @@ for (const mode of modes) { it("should send messages via handle", async () => { const workflow = async (ctx: WorkflowContextInterface) => { - const message = await ctx.listen("wait", "message-name"); + const [message] = await ctx.queue.next("wait", { + names: ["message-name"], + }); + if (!message) { + throw new Error("Expected message"); + } return message.body; }; diff --git a/rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts index 27da0e0467..b67fefbcf4 100644 --- a/rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts +++ b/rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts @@ -31,10 +31,12 @@ for (const mode of modes) { it("should wait for messages", async () => { const workflow = async (ctx: WorkflowContextInterface) => { - const message = await ctx.listen( - "wait-message", - "my-message", - ); + const [message] = await ctx.queue.next("wait-message", { + names: ["my-message"], + }); + if (!message) { + throw new Error("Expected message"); + } return message.body; }; @@ -55,12 +57,14 @@ for (const mode of modes) { expect(result.output).toBe("payload"); }); - it("should listen for any message in a name set", async () => { + it("should wait for any message in a name set", async () => { const workflow = async (ctx: WorkflowContextInterface) => { - const message = await ctx.listen("wait-many", [ - "first", - "second", - ]); + const [message] = await ctx.queue.next("wait-many", { + names: ["first", "second"], + }); + if (!message) { + throw new Error("Expected message"); + } return { name: message.name, body: message.body }; }; @@ -107,13 +111,15 @@ for (const mode of modes) { ), ); - const workflow = async (ctx: WorkflowContextInterface) => { - const message = await ctx.listen( - "wait-message", - "my-message", - ); - return message.body; - }; + const workflow = async (ctx: WorkflowContextInterface) => { + const [message] = await ctx.queue.next("wait-message", { + names: ["my-message"], + }); + if (!message) { + throw new Error("Expected message"); + } + return message.body; + }; const result = await runWorkflow( "wf-1", @@ -129,7 +135,7 @@ for (const mode of modes) { expect(result.output).toBe("hello"); }); - it("listen should return a completable message handle", async () => { + it("queue.next should return completable messages", async () => { const completions: Array<{ id: string; response?: unknown }> = []; const pending = [ buildMessagePayload("my-message", "hello", "msg-1") as { @@ -165,11 +171,20 @@ for (const mode of modes) { }; driver.messageDriver = messageDriver; - const workflow = async (ctx: WorkflowContextInterface) => { - const message = await ctx.listen("wait-message", "my-message"); - await message.complete({ ok: true }); - return message.body; - }; + const workflow = async (ctx: WorkflowContextInterface) => { + const [message] = await ctx.queue.next("wait-message", { + names: ["my-message"], + completable: true, + }); + if (!message) { + throw new Error("Expected message"); + } + if (!message.complete) { + throw new Error("Expected completable message"); + } + await message.complete({ ok: true }); + return message.body; + }; const result = await runWorkflow( "wf-1", @@ -186,7 +201,142 @@ for (const mode of modes) { expect(completions).toEqual([{ id: "msg-1", response: { ok: true } }]); }); - it("should collect multiple messages with listenN", async () => { + it("replay should not block the next completable queue.next", async () => { + if (mode !== "yield") { + return; + } + + await driver.set( + buildMessageKey("msg-1"), + serializeMessage(buildMessagePayload("my-message", "one", "msg-1")), + ); + await driver.set( + buildMessageKey("msg-2"), + serializeMessage(buildMessagePayload("my-message", "two", "msg-2")), + ); + + const workflow = async (ctx: WorkflowContextInterface) => { + const [first] = await ctx.queue.next("wait-first", { + names: ["my-message"], + completable: true, + }); + if (!first || !first.complete) { + throw new Error("Expected first completable message"); + } + const completeFirst = first.complete; + await ctx.step("complete-first", async () => { + await completeFirst({ ok: "first" }); + return first.body; + }); + + await ctx.sleep("between", 120); + + const [second] = await ctx.queue.next("wait-second", { + names: ["my-message"], + completable: true, + }); + if (!second || !second.complete) { + throw new Error("Expected second completable message"); + } + const completeSecond = second.complete; + await ctx.step("complete-second", async () => { + await completeSecond({ ok: "second" }); + return second.body; + }); + + return [first.body, second.body] as const; + }; + + const firstRun = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(firstRun.state).toBe("sleeping"); + + await new Promise((resolve) => setTimeout(resolve, 140)); + + const secondRun = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(secondRun.state).toBe("completed"); + expect(secondRun.output).toEqual(["one", "two"]); + }); + + it("replay should keep blocking if completable message was not completed", async () => { + if (mode !== "yield") { + return; + } + + await driver.set( + buildMessageKey("msg-1"), + serializeMessage( + buildMessagePayload("my-message", "one", "msg-1"), + ), + ); + await driver.set( + buildMessageKey("msg-2"), + serializeMessage( + buildMessagePayload("my-message", "two", "msg-2"), + ), + ); + + const workflow = async (ctx: WorkflowContextInterface) => { + const [first] = await ctx.queue.next("wait-first", { + names: ["my-message"], + completable: true, + }); + if (!first || !first.complete) { + throw new Error("Expected first completable message"); + } + + // Intentionally do not complete the message. + await ctx.sleep("between", 120); + + await ctx.queue.next("wait-second", { + names: ["my-message"], + completable: true, + }); + + return first.body; + }; + + const firstRun = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(firstRun.state).toBe("sleeping"); + + await new Promise((resolve) => setTimeout(resolve, 140)); + + const secondRunHandle = runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ); + await expect(secondRunHandle.result).rejects.toThrow( + "Previous completable queue message is not completed.", + ); + + const queued = await driver.messageDriver.loadMessages(); + expect(queued.map((message) => message.id).sort()).toEqual([ + "msg-1", + "msg-2", + ]); + }); + + it("should collect multiple messages with queue.next count", async () => { await driver.set( buildMessageKey("1"), serializeMessage(buildMessagePayload("batch", "a", "1")), @@ -196,9 +346,13 @@ for (const mode of modes) { serializeMessage(buildMessagePayload("batch", "b", "2")), ); - const workflow = async (ctx: WorkflowContextInterface) => { - return await ctx.listenN("batch-wait", "batch", 2); - }; + const workflow = async (ctx: WorkflowContextInterface) => { + const messages = await ctx.queue.next("batch-wait", { + names: ["batch"], + count: 2, + }); + return messages.map((message) => message.body); + }; const result = await runWorkflow( "wf-1", @@ -214,14 +368,14 @@ for (const mode of modes) { expect(result.output).toEqual(["a", "b"]); }); - it("should time out listenWithTimeout", async () => { - const workflow = async (ctx: WorkflowContextInterface) => { - return await ctx.listenWithTimeout( - "timeout", - "missing", - 50, - ); - }; + it("should time out queue.next", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + const messages = await ctx.queue.next("timeout", { + names: ["missing"], + timeout: 50, + }); + return messages[0]?.body ?? null; + }; if (mode === "yield") { const result1 = await runWorkflow( @@ -258,7 +412,7 @@ for (const mode of modes) { expect(result.output).toBeNull(); }); - it("should return a message before listenUntil deadline", async () => { + it("should return a message before queue.next timeout", async () => { const messageId = generateId(); await driver.set( buildMessageKey(messageId), @@ -267,13 +421,13 @@ for (const mode of modes) { ), ); - const workflow = async (ctx: WorkflowContextInterface) => { - return await ctx.listenUntil( - "deadline", - "deadline", - Date.now() + 1000, - ); - }; + const workflow = async (ctx: WorkflowContextInterface) => { + const messages = await ctx.queue.next("deadline", { + names: ["deadline"], + timeout: 1000, + }); + return messages[0]?.body ?? null; + }; const result = await runWorkflow( "wf-1", @@ -289,15 +443,15 @@ for (const mode of modes) { expect(result.output).toBe("data"); }); - it("should wait for listenNWithTimeout messages", async () => { - const workflow = async (ctx: WorkflowContextInterface) => { - return await ctx.listenNWithTimeout( - "batch", - "batch", - 2, - 5000, - ); - }; + it("should wait for queue.next timeout messages", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + const messages = await ctx.queue.next("batch", { + names: ["batch"], + count: 2, + timeout: 5000, + }); + return messages.map((message) => message.body); + }; const handle = runWorkflow("wf-1", workflow, undefined, driver, { mode, @@ -327,7 +481,7 @@ for (const mode of modes) { expect(result.output).toEqual(["first"]); }); - it("should respect limits and FIFO ordering in listenNUntil", async () => { + it("should respect limits and FIFO ordering in queue.next", async () => { await driver.set( buildMessageKey("1"), serializeMessage(buildMessagePayload("fifo", "first", "1")), @@ -341,14 +495,14 @@ for (const mode of modes) { serializeMessage(buildMessagePayload("fifo", "third", "3")), ); - const workflow = async (ctx: WorkflowContextInterface) => { - return await ctx.listenNUntil( - "fifo", - "fifo", - 2, - Date.now() + 1000, - ); - }; + const workflow = async (ctx: WorkflowContextInterface) => { + const messages = await ctx.queue.next("fifo", { + names: ["fifo"], + count: 2, + timeout: 1000, + }); + return messages.map((message) => message.body); + }; const result = await runWorkflow( "wf-1", @@ -381,9 +535,14 @@ for (const mode of modes) { return "ready"; }); - const message = await ctx.listen("wait", "mid"); - return message.body; - }; + const [message] = await ctx.queue.next("wait", { + names: ["mid"], + }); + if (!message) { + throw new Error("Expected message"); + } + return message.body; + }; const handle = runWorkflow("wf-1", workflow, undefined, driver, { mode, diff --git a/rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts index a6ec0165c7..356828ab57 100644 --- a/rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts +++ b/rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts @@ -47,9 +47,11 @@ for (const mode of modes) { return Loop.continue({ count: state.count + 1 }); }, }); - await ctx.sleep("old-sleep", 0); - await ctx.listen("old-listen", "old-message"); - await ctx.join("old-join", { + await ctx.sleep("old-sleep", 0); + await ctx.queue.next("old-listen", { + names: ["old-message"], + }); + await ctx.join("old-join", { branch: { run: async () => "ok", }, diff --git a/website/src/content/docs/actors/lifecycle.mdx b/website/src/content/docs/actors/lifecycle.mdx index 2a53894b3a..e99c265225 100644 --- a/website/src/content/docs/actors/lifecycle.mdx +++ b/website/src/content/docs/actors/lifecycle.mdx @@ -267,7 +267,7 @@ const tickActor = actor({ run: async (c) => { c.log.info("Background loop started"); - while (!c.abortSignal.aborted) { + while (!c.aborted) { c.state.tickCount++; c.log.info({ msg: "tick", count: c.state.tickCount }); @@ -300,9 +300,10 @@ const queueConsumer = actor({ run: async (c) => { c.log.info("Queue consumer started"); - while (!c.abortSignal.aborted) { - // Wait for next message with timeout - const message = await c.queue.next("tasks", { timeout: 1000 }); + while (!c.aborted) { + // Wait for next message with timeout. + const messages = await c.queue.next({ names: ["tasks"], timeout: 1000 }); + const message = messages[0]; if (message) { c.log.info({ msg: "processing message", body: message.body }); @@ -858,7 +859,7 @@ const counter = actor({ // Background task (does not block startup) run: async (c) => { - while (!c.abortSignal.aborted) { + while (!c.aborted) { // Example: periodic logging console.log(`Counter "${c.state.name}" is at ${c.state.count}`); await new Promise((resolve) => { diff --git a/website/src/content/docs/actors/queue.mdx b/website/src/content/docs/actors/queue.mdx index 949a6beeb8..cd805754e4 100644 --- a/website/src/content/docs/actors/queue.mdx +++ b/website/src/content/docs/actors/queue.mdx @@ -1,29 +1,28 @@ --- title: "Queue Messages" -description: "Send durable messages to Rivet Actors and optionally wait for completion." +description: "Send durable queue messages to Rivet Actors and consume them from run loops." skill: true --- -Rivet Actors include a pull-based queue for durable message processing. Clients enqueue messages, and actors pull them with `c.queue.next()`. - -Queues are durable: messages persist until the actor acknowledges them. When `wait: true`, the actor must call `msg.complete()` to remove the message. +Rivet Actors include a pull-based queue for durable background processing. ## Send Messages (Client) -Fire-and-forget sends a message and returns immediately: +Use `handle.send(name, body)` for fire-and-forget: -```typescript @nocheck -const handle = client.myActor.getOrCreate(["orders"]); +```typescript +const handle = client.worker.getOrCreate(["main"]); -await handle.queue.tasks.send({ id: "order-123" }); +await handle.send("jobs", { id: "job-1" }); ``` -To wait for the actor to finish processing: +Use `wait: true` for request/response: -```typescript @nocheck -const result = await handle.queue.tasks.send( - { id: "order-123" }, - { wait: true, timeout: 30_000 } +```typescript +const result = await handle.send( + "jobs", + { id: "job-1" }, + { wait: true, timeout: 30_000 }, ); if (result.status === "completed") { @@ -33,109 +32,85 @@ if (result.status === "completed") { } ``` -## Receive Messages (Actor) +## Queue Schema -Actors pull messages with `c.queue.next()`: +Define queue message types under `queues`. Use `complete` when a queue supports manual completion responses. -```typescript @nocheck -const msg = await c.queue.next("tasks"); -if (!msg) return; +```typescript +import { actor, queue } from "rivetkit"; -// Process the message... +const worker = actor({ + state: {}, + queues: { + jobs: queue<{ id: string }, { ok: true }>(), + logs: queue<{ line: string }>(), + }, + actions: {}, +}); ``` -Each message includes a stable `id` string you can log or correlate across systems. - -Use `wait: true` to hold the message until you explicitly complete it: - -```typescript @nocheck -const msg = await c.queue.next("tasks", { wait: true }); -if (!msg) return; +## Receive Messages (Actor) -// Process the message... -await msg.complete(); -``` +### `next` -You can also send data back to the waiting client: +`next` returns an array and can block until messages are available. -```typescript @nocheck -await msg.complete({ result: "ok" }); +```typescript +const messages = await c.queue.next({ + names: ["jobs"], + count: 10, + timeout: 1000, + signal: abortController.signal, +}); ``` -## Manual Completion - -When `wait: true`, the message is marked in-flight and remains persisted until completion. Calling `msg.complete()` removes the message and resolves any waiting client. - -If `wait: false`, calling `msg.complete()` throws `QueueCompleteNotAllowed`. - -If a client sends with `wait: true` but the actor receives without `wait: true`, the message auto-completes and the waiting client receives `{ status: \"completed\", response: undefined }`. - -## Persistence & Redelivery - -- Messages persist until `msg.complete()` is called. -- In-flight state persists across restarts. -- If the actor crashes before completion, the message is redelivered with exponential backoff. -- Completion responses are not persisted and are only delivered to waiting clients. - -Backoff schedule: - -- 1s, 2s, 4s, 8s, ... (capped at 5 minutes) - -Messages are retried indefinitely. There is no forced timeout on the actor side. +If no messages arrive before timeout, `next` returns `[]`. -## Error Handling +### `tryNext` -- `QueueCompleteNotAllowed`: `msg.complete()` called when `wait: false`. -- `QueueMessagePending`: `c.queue.next()` called while a previous `wait: true` message is still pending. -- `QueueAlreadyCompleted`: `msg.complete()` called more than once. +`tryNext` is non-blocking and immediately returns `[]` when empty. -If a message remains pending for more than 30 seconds, the actor logs a warning but does not time out automatically. +```typescript +const messages = await c.queue.tryNext({ names: ["jobs"], count: 10 }); +``` -## HTTP API +### `iter` -Use the vanilla HTTP API to enqueue messages: +`iter` returns an async iterator yielding one message at a time. -``` -POST /queue/:name -Body: { body: { ... } } -Response: { status: "completed" } +```typescript +for await (const message of c.queue.iter({ + names: ["jobs"], + signal: abortController.signal, +})) { + // process message +} ``` -Wait for completion: +### Iterate All Queue Names -``` -POST /queue/:name -Body: { body: { ... }, wait: true, timeout: 30000 } -Response: { status: "completed", response: { ... } } -``` - -Timeouts still return HTTP 200: +Use `iter()` without `names` to consume across all queue names. -``` -{ status: "timedOut" } +```typescript +for await (const message of c.queue.iter()) { + // process message +} ``` -## Examples +## Completable Messages -### Fire-and-forget +Use `completable: true` to receive messages that expose `message.complete(...)`. -```typescript @nocheck -await handle.queue.emails.send({ to: "user@example.com" }); +```typescript +for await (const message of c.queue.iter({ names: ["jobs"], completable: true })) { + await message.complete({ ok: true }); +} ``` -### Request-response +The message stays in the queue until `message.complete(...)` is called. -```typescript @nocheck -// Client -const result = await handle.queue.work.send( - { input: "value" }, - { wait: true, timeout: 10_000 } -); +## Abort Behavior -// Actor -const msg = await c.queue.next("work", { wait: true }); -if (msg) { - const output = doWork(msg.body); - await msg.complete({ output }); -} -``` +Use `c.aborted` for loop exit conditions when needed. + +Never wrap `c.queue.next(...)` in `try/catch` for normal shutdown handling. Queue receive calls throw special abort errors during actor shutdown so the run handler can stop cleanly. diff --git a/website/src/content/docs/general/environment-variables.mdx b/website/src/content/docs/general/environment-variables.mdx index 1818a11a26..4a522b6cb3 100644 --- a/website/src/content/docs/general/environment-variables.mdx +++ b/website/src/content/docs/general/environment-variables.mdx @@ -45,6 +45,12 @@ These variables configure how clients connect to your actors. | `RIVET_INSPECTOR_TOKEN` | Token for accessing the Rivet Inspector | | `RIVET_INSPECTOR_DISABLE` | Set to `1` to disable the inspector | +## Storage + +| Environment Variable | Description | +|---------------------|-------------| +| `RIVETKIT_STORAGE_PATH` | Overrides the default file-system storage path used by RivetKit when using the default driver. | + ## Logging | Environment Variable | Description | @@ -55,4 +61,3 @@ These variables configure how clients connect to your actors. | `RIVET_LOG_MESSAGE` | Set to `1` to include message formatting | | `RIVET_LOG_ERROR_STACK` | Set to `1` to include error stack traces | | `RIVET_LOG_HEADERS` | Set to `1` to log request headers | - diff --git a/website/src/metadata/skill-base-rivetkit.md b/website/src/metadata/skill-base-rivetkit.md index a38c2cfd69..3d53554e9c 100644 --- a/website/src/metadata/skill-base-rivetkit.md +++ b/website/src/metadata/skill-base-rivetkit.md @@ -92,6 +92,10 @@ The RivetKit OpenAPI specification is available in the skill directory at `opena +## Code Style + +- Use object property syntax (arrow functions) for lifecycle hooks and `run`, not method shorthand. For example, use `onWake: async (c) => { ... }` instead of `async onWake(c) { ... }`. This applies to `onCreate`, `onDestroy`, `onWake`, `onSleep`, `onStateChange`, `onBeforeConnect`, `onConnect`, `onDisconnect`, `onBeforeActionResponse`, `onRequest`, `onWebSocket`, and `run`. + ## Misc Notes - The Rivet domain is rivet.dev, not rivet.gg