From 775628c52cebf53032a39a1b9279b6fbc11d6402 Mon Sep 17 00:00:00 2001
From: Jack Kleeman <jackkleeman@gmail.com>
Date: Thu, 5 Sep 2024 14:47:37 +0100
Subject: [PATCH] Add state types

---
 README.md      |  6 +++---
 src/lib.ts     | 43 ++++++++++++++++++++-----------------------
 src/promise.ts |  8 ++++----
 3 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/README.md b/README.md
index 026a7d1..3ae785d 100644
--- a/README.md
+++ b/README.md
@@ -15,9 +15,9 @@ avoid blocking the event loop.
 The service is set up and managed automatically by interpreting the state
 machine definition, and can be deployed as a Lambda or as a long-lived service.
 
-In `app.ts` you will see an example of an XState machine that uses cross-machine
-communication, delays, and Promise actors, all running in Restate. However,
-most XState machines should work out of the box; this is still experimental, so
+In [`example/app.ts`](./example/app.ts) you will see an example of an XState machine
+that uses cross-machine communication, delays, and Promise actors, all running in Restate.
+Most XState machines should work out of the box, but this is still experimental, so
 we haven't tested everything yet!
 
 To try out this example:
diff --git a/src/lib.ts b/src/lib.ts
index 826534c..99eff1a 100644
--- a/src/lib.ts
+++ b/src/lib.ts
@@ -44,14 +44,11 @@ export interface RestateActorSystem<T extends ActorSystemInfo>
     event: AnyEventObject
   ) => void;
   api: XStateApi<ActorLogicFrom<T>>;
-  ctx: restate.ObjectContext;
+  ctx: restate.ObjectContext<State>;
   systemName: string;
 }
 
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export type AnyRestateActorSystem = RestateActorSystem<any>;
-
-export type SerialisableActorRef = {
+type SerialisableActorRef = {
   id: string;
   sessionId: string;
   _parent?: SerialisableActorRef;
@@ -80,16 +77,19 @@ type SerialisableScheduledEvent = {
   uuid: string;
 };
 
+type State = {
+  events: { [key: string]: SerialisableScheduledEvent };
+  children: { [key: string]: SerialisableActorRef };
+  snapshot: Snapshot<unknown>;
+};
+
 async function createSystem<T extends ActorSystemInfo>(
-  ctx: restate.ObjectContext,
+  ctx: restate.ObjectContext<State>,
   api: XStateApi<ActorLogicFrom<T>>,
   systemName: string
 ): Promise<RestateActorSystem<T>> {
-  const events =
-    (await ctx.get<{ [key: string]: SerialisableScheduledEvent }>("events")) ??
-    {};
-  const childrenByID =
-    (await ctx.get<{ [key: string]: SerialisableActorRef }>("children")) ?? {};
+  const events = (await ctx.get("events")) ?? {};
+  const childrenByID = (await ctx.get("children")) ?? {};
 
   const children = new Map<string, ActorRefEventSender>();
   const keyedActors = new Map<keyof T["actors"], AnyActorRef | undefined>();
@@ -288,15 +288,15 @@ export interface ActorRefEventSender extends AnyActorRef {
   _send: (event: AnyEventObject) => void;
 }
 
-export async function createActor<TLogic extends AnyStateMachine>(
-  ctx: restate.ObjectContext,
+async function createActor<TLogic extends AnyStateMachine>(
+  ctx: restate.ObjectContext<State>,
   api: XStateApi<TLogic>,
   systemName: string,
   logic: TLogic,
   options?: ActorOptions<TLogic>
 ): Promise<ActorEventSender<TLogic>> {
   const system = await createSystem(ctx, api, systemName);
-  const snapshot = (await ctx.get<Snapshot<unknown>>("snapshot")) ?? undefined;
+  const snapshot = (await ctx.get("snapshot")) ?? undefined;
 
   const parent: ActorRefEventSender = {
     id: "fakeRoot",
@@ -361,7 +361,7 @@ const actorObject = <TLogic extends AnyStateMachine>(
     name: path,
     handlers: {
       create: async (
-        ctx: restate.ObjectContext,
+        ctx: restate.ObjectContext<State>,
         request?: { input?: InputFrom<TLogic> }
       ): Promise<Snapshot<unknown>> => {
         const systemName = ctx.key;
@@ -381,7 +381,7 @@ const actorObject = <TLogic extends AnyStateMachine>(
         return root.getPersistedSnapshot();
       },
       send: async (
-        ctx: restate.ObjectContext,
+        ctx: restate.ObjectContext<State>,
         request?: {
           scheduledEvent?: SerialisableScheduledEvent;
           source?: SerialisableActorRef;
@@ -396,10 +396,7 @@ const actorObject = <TLogic extends AnyStateMachine>(
         }
 
         if (request.scheduledEvent) {
-          const events =
-            (await ctx.get<{ [key: string]: SerialisableScheduledEvent }>(
-              "events"
-            )) ?? {};
+          const events = (await ctx.get("events")) ?? {};
           const scheduledEventId = createScheduledEventId(
             request.scheduledEvent.source,
             request.scheduledEvent.id
@@ -430,7 +427,7 @@ const actorObject = <TLogic extends AnyStateMachine>(
 
         let actor;
         if (request.target) {
-          actor = (root.system as AnyRestateActorSystem).actor(
+          actor = (root.system as RestateActorSystem<ActorSystemInfo>).actor(
             request.target.sessionId
           );
           if (!actor) {
@@ -442,7 +439,7 @@ const actorObject = <TLogic extends AnyStateMachine>(
           actor = root;
         }
 
-        (root.system as AnyRestateActorSystem)._relay(
+        (root.system as RestateActorSystem<ActorSystemInfo>)._relay(
           request.source,
           actor,
           request.event
@@ -454,7 +451,7 @@ const actorObject = <TLogic extends AnyStateMachine>(
         return nextSnapshot;
       },
       snapshot: async (
-        ctx: restate.ObjectContext,
+        ctx: restate.ObjectContext<State>,
         systemName: string
       ): Promise<Snapshot<unknown>> => {
         const root = await createActor(ctx, api, systemName, logic);
diff --git a/src/promise.ts b/src/promise.ts
index a915922..b8eef2c 100644
--- a/src/promise.ts
+++ b/src/promise.ts
@@ -12,10 +12,10 @@ import type {
 } from "xstate";
 import {
   serialiseActorRef,
-  type AnyRestateActorSystem,
   type ActorRefEventSender,
+  type RestateActorSystem,
 } from "./lib.js";
-import type { Context } from "@restatedev/restate-sdk";
+import type { ObjectSharedContext } from "@restatedev/restate-sdk";
 
 export type PromiseSnapshot<TOutput, TInput> = Snapshot<TOutput> & {
   input: TInput | undefined;
@@ -32,7 +32,7 @@ export type PromiseCreator<TOutput, TInput extends NonReducibleUnknown> = ({
   ctx,
 }: {
   input: TInput;
-  ctx: Context;
+  ctx: ObjectSharedContext;
 }) => PromiseLike<TOutput>;
 
 export type PromiseActorLogic<TOutput, TInput = unknown> = ActorLogic<
@@ -102,7 +102,7 @@ export function fromPromise<TOutput, TInput extends NonReducibleUnknown>(
         return;
       }
 
-      const rs = system as AnyRestateActorSystem;
+      const rs = system as RestateActorSystem<ActorSystemInfo>;
 
       rs.ctx.objectSendClient(rs.api, rs.systemName).invokePromise({
         self: serialiseActorRef(self),