Add new methods to Schedule API

zephraph · zephraph · commit afe83c20eb26 · 2025-03-31T13:08:19.000-04:00
diff --git a/docs/concepts/schedule.mdx b/docs/concepts/schedule.mdx
@@ -21,6 +21,9 @@ Parameters:
 - `fn` (string): The name of the action to be executed.
 - `...args` (unknown[]): Additional arguments to pass to the function.
 
+Returns:
+- `Promise<string>`: A unique identifier for the scheduled event.
+
 ### `c.schedule.at(timestamp, fn, ...args)`
 
 Schedules a function to be executed at a specific timestamp. This function persists across actor restarts, upgrades, or crashes.
@@ -31,6 +34,41 @@ Parameters:
 - `fn` (string): The name of the action to be executed.
 - `...args` (unknown[]): Additional arguments to pass to the function.
 
+Returns:
+- `Promise<string>`: A unique identifier for the scheduled event.
+
+### `c.schedule.list()`
+
+Lists all scheduled events for the actor.
+
+Returns:
+- `Promise<Alarm[]>`: An array of scheduled alarms, where each alarm has the following properties:
+  - `id` (string): The unique identifier of the alarm
+  - `createdAt` (number): The timestamp when the alarm was created
+  - `triggersAt` (number): The timestamp when the alarm will trigger
+  - `fn` (string): The name of the action to be executed
+  - `args` (unknown[]): The arguments to pass to the function
+
+### `c.schedule.get(alarmId)`
+
+Gets details about a specific scheduled event.
+
+Parameters:
+- `alarmId` (string): The unique identifier of the alarm to retrieve.
+
+Returns:
+- `Promise<Alarm | undefined>`: The alarm details if found, undefined otherwise.
+
+### `c.schedule.cancel(alarmId)`
+
+Cancels a scheduled event.
+
+Parameters:
+- `alarmId` (string): The unique identifier of the alarm to cancel.
+
+Returns:
+- `Promise<void>`
+
 ## Scheduling Private Actions
 
 Currently, scheduling can only trigger public actions. If the scheduled action is private, it needs to be secured with something like a token.
@@ -46,7 +84,7 @@ const reminderService = actor({
   },
   
   actions: {
-    setReminder: (c, userId, message, delayMs) => {
+    setReminder: async (c, userId, message, delayMs) => {
       const reminderId = crypto.randomUUID();
       
       // Store the reminder in state
@@ -57,7 +95,89 @@ const reminderService = actor({
       };
       
       // Schedule the sendReminder action to run after the delay
-      c.after(delayMs, "sendReminder", reminderId);
+      // Store the alarmId for potential cancellation
+      const alarmId = await c.schedule.after(delayMs, "sendReminder", reminderId);
+      
+      return { reminderId, alarmId };
+    },
+    
+    cancelReminder: async (c, reminderId) => {
+      const reminder = c.state.reminders[reminderId];
+      if (!reminder) return { success: false };
+      
+      // Cancel the scheduled reminder
+      await c.schedule.cancel(reminder.alarmId);
+      
+      // Clean up the reminder
+      delete c.state.reminders[reminderId];
+      
+      return { success: true };
+    },
+    
+    sendReminder: (c, reminderId) => {
+      const reminder = c.state.reminders[reminderId];
+      if (!reminder) return;
+      
+      // Find the user's connection if they're online
+      const userConn = c.conns.find(
+        conn => conn.state.userId === reminder.userId
+      );
+      
+      if (userConn) {
+        // Send the reminder to the user
+        userConn.send("reminder", {
+          message: reminder.message,
+          scheduledAt: reminder.scheduledFor
+        });
+      } else {
+        // If user is offline, store reminder for later delivery
+        // ...
+      }
+      
+      // Clean up the processed reminder
+      delete c.state.reminders[reminderId];
+    }
+  }
+});
+```
+
+## Testing Schedules
+
+```typescript
+import { actor } from "actor-core";
+
+const reminderService = actor({
+  state: {
+    reminders: {}
+  },
+  
+  actions: {
+    setReminder: async (c, userId, message, delayMs) => {
+      const reminderId = crypto.randomUUID();
+      
+      // Store the reminder in state
+      c.state.reminders[reminderId] = {
+        userId,
+        message,
+        scheduledFor: Date.now() + delayMs
+      };
+      
+      // Schedule the sendReminder action to run after the delay
+      // Store the alarmId for potential cancellation
+      const alarmId = await c.schedule.after(delayMs, "sendReminder", reminderId);
+      
+      return { reminderId, alarmId };
+    },
+    
+    cancelReminder: async (c, reminderId) => {
+      const reminder = c.state.reminders[reminderId];
+      if (!reminder) return { success: false };
+      
+      // Cancel the scheduled reminder
+      await c.schedule.cancel(reminder.alarmId);
+      
+      // Clean up the reminder
+      delete c.state.reminders[reminderId];
       
       return { reminderId };
     },
diff --git a/packages/actor-core/src/actor/schedule.ts b/packages/actor-core/src/actor/schedule.ts
@@ -15,6 +15,15 @@ interface ScheduleIndexEvent {
 
 interface ScheduleEvent {
 	timestamp: number;
+	createdAt: number;
+	fn: string;
+	args: unknown[];
+}
+
+export interface Alarm {
+	id: string;
+	createdAt: number;
+	triggersAt: number;
 	fn: string;
 	args: unknown[];
 }
@@ -28,27 +37,109 @@ export class Schedule {
 		this.#driver = driver;
 	}
 
-	async after(duration: number, fn: string, ...args: unknown[]) {
-		await this.#scheduleEvent(Date.now() + duration, fn, args);
+	async after(duration: number, fn: string, ...args: unknown[]): Promise<string> {
+		return this.#scheduleEvent(Date.now() + duration, fn, args);
+	}
+
+	async at(timestamp: number, fn: string, ...args: unknown[]): Promise<string> {
+		return this.#scheduleEvent(timestamp, fn, args);
+	}
+
+	async get(alarmId: string): Promise<Alarm | null> {
+		const event = await this.#driver.kvGet(
+			this.#actor.id,
+			KEYS.SCHEDULE.event(alarmId)
+		) as ScheduleEvent | undefined;
+
+		if (!event) return null;
+
+		return {
+			id: alarmId,
+			createdAt: event.createdAt,
+			triggersAt: event.timestamp,
+			fn: event.fn,
+			args: event.args
+		};
+	}
+
+	async list(): Promise<readonly Alarm[]> {
+		const schedule: ScheduleState = ((await this.#driver.kvGet(
+			this.#actor.id,
+			KEYS.SCHEDULE.SCHEDULE
+		)) as ScheduleState) ?? { events: [] };
+
+		const alarms: Alarm[] = [];
+		for (const event of schedule.events) {
+			const scheduleEvent = await this.#driver.kvGet(
+				this.#actor.id,
+				KEYS.SCHEDULE.event(event.eventId)
+			) as ScheduleEvent;
+
+			if (scheduleEvent) {
+				alarms.push(Object.freeze({
+					id: event.eventId,
+					createdAt: scheduleEvent.createdAt,
+					triggersAt: scheduleEvent.timestamp,
+					fn: scheduleEvent.fn,
+					args: scheduleEvent.args
+				}));
+			}
+		}
+
+		return Object.freeze(alarms);
 	}
 
-	async at(timestamp: number, fn: string, ...args: unknown[]) {
-		await this.#scheduleEvent(timestamp, fn, args);
+	async cancel(alarmId: string): Promise<void> {
+		// Get the schedule index
+		const schedule: ScheduleState = ((await this.#driver.kvGet(
+			this.#actor.id,
+			KEYS.SCHEDULE.SCHEDULE
+		)) as ScheduleState) ?? { events: [] };
+
+		// Find and remove the event from the index
+		const eventIndex = schedule.events.findIndex(x => x.eventId === alarmId);
+		if (eventIndex === -1) return;
+
+		const [removedEvent] = schedule.events.splice(eventIndex, 1);
+
+		// Delete the event data
+		await this.#driver.kvDelete(
+			this.#actor.id,
+			KEYS.SCHEDULE.event(alarmId)
+		);
+
+		// Update the schedule index
+		await this.#driver.kvPut(
+			this.#actor.id,
+			KEYS.SCHEDULE.SCHEDULE,
+			schedule
+		);
+
+		// If we removed the first event (next to execute), update the alarm
+		if (eventIndex === 0) {
+			if (schedule.events.length > 0) {
+				// Set alarm to next event
+				await this.#driver.setAlarm(this.#actor, schedule.events[0].timestamp);
+			} else {
+				// No more events, delete the alarm
+				await this.#driver.deleteAlarm(this.#actor);
+			}
+		}
 	}
 
 	async #scheduleEvent(
 		timestamp: number,
 		fn: string,
 		args: unknown[],
-	): Promise<void> {
+	): Promise<string> {
 		// Save event
 		const eventId = crypto.randomUUID();
 		await this.#driver.kvPut(
 			this.#actor.id,
-			// biome-ignore lint/suspicious/noExplicitAny: <explanation>
-			KEYS.SCHEDULE.event(eventId) as any,
+			KEYS.SCHEDULE.event(eventId),
 			{
 				timestamp,
+				createdAt: Date.now(),
 				fn,
 				args,
 			},
@@ -58,8 +149,7 @@ export class Schedule {
 		// Read index
 		const schedule: ScheduleState = ((await this.#driver.kvGet(
 			this.#actor.id,
-			// biome-ignore lint/suspicious/noExplicitAny: <explanation>
-			KEYS.SCHEDULE.SCHEDULE as any,
+			KEYS.SCHEDULE.SCHEDULE,
 		)) as ScheduleState) ?? {
 			events: [],
 		};
@@ -84,6 +174,8 @@ export class Schedule {
 		if (insertIndex === 0 || schedule.events.length === 1) {
 			await this.#driver.setAlarm(this.#actor, newEvent.timestamp);
 		}
+
+		return eventId;
 	}
 
 	async __onAlarm() {
