Port state api for restate-sdk-gen

packages/libs/restate-sdk-gen/e2e/state.e2e.test.ts

@@ -9,34 +9,38 @@
  * https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
  */
 
-// State e2e: a virtual-object counter exercises state() against a real
-// Restate runtime — get/set are journaled, clear and clearAll affect
-// persistent VO state, keys() returns the live key set.
+// State e2e: a virtual-object counter exercises the new per-key accessor API
+// against a real Restate runtime.
 //
-// Two flavors of the API are demonstrated:
-//
-//   - Typed: state<CounterState>() gives keyof-checked names and
-//     per-key value types. get("count") infers Future<number | null>;
-//     state.get("nope") would be a type error.
-//   - Shared: read-only handlers use sharedState<CounterState>(),
-//     which drops set/clear/clearAll from the type — calling them is
-//     a compile error, mirroring what the runtime would do anyway
-//     (ObjectSharedContext has no write methods).
+// Demonstrates:
+//   - state(config) defined at module level (lazy — resolves ops on method call).
+//     count has default 0 → get() returns Future<number> (never null).
+//     secondary uses typed<string>() → get() returns Future<string | null>.
+//   - state(config) with factory default (closure) for mutable defaults.
+//   - getAllStateKeys(): list live keys.
+//   - clear() on per-key accessors.
 //
 // Both runtime modes: default + alwaysReplay.
 
 import { afterAll, beforeAll, describe, expect, test } from "vitest";
 import * as restate from "@restatedev/restate-sdk";
 import * as clients from "@restatedev/restate-sdk-clients";
 import { RestateTestEnvironment } from "@restatedev/restate-sdk-testcontainers";
-import { gen, execute, state, sharedState } from "@restatedev/restate-sdk-gen";
-
-// Typed-state shape for the counter object. Keys and value types here
-// flow through to state<CounterState>() at every call site below.
-type CounterState = {
-  count: number;
-  secondary: string;
-};
+import {
+  gen,
+  execute,
+  state,
+  typed,
+  getAllStateKeys,
+} from "@restatedev/restate-sdk-gen";
+
+// Defined once at module level — lazy, resolves ops only when methods are called.
+// count has a static default 0 (get returns Future<number>).
+// secondary uses typed<string>() — typed but no default (get returns Future<string | null>).
+const counterState = state({
+  count: { default: 0 },
+  secondary: typed<string>(),
+});
 
 const counterObj = restate.object({
   name: "counter",
@@ -45,28 +49,26 @@ const counterObj = restate.object({
       execute(
         ctx,
         gen(function* () {
-          const s = state<CounterState>();
-          const current = (yield* s.get("count")) ?? 0;
+          const current = yield* counterState.count.get(); // Future<number>
           const next = current + n;
-          s.set("count", next);
+          counterState.count.set(next);
           return next;
         })
       ),
 
-    // Read-only handler — uses sharedState() so writes wouldn't compile.
     current: async (ctx: restate.ObjectSharedContext): Promise<number> =>
       execute(
         ctx,
         gen(function* () {
-          return (yield* sharedState<CounterState>().get("count")) ?? 0;
+          return yield* counterState.count.get(); // still Future<number> — default 0
         })
       ),
 
     keys: async (ctx: restate.ObjectSharedContext): Promise<string[]> =>
       execute(
         ctx,
         gen(function* () {
-          return yield* sharedState<CounterState>().keys();
+          return yield* getAllStateKeys();
         })
       ),
 
@@ -77,9 +79,8 @@ const counterObj = restate.object({
       execute(
         ctx,
         gen(function* () {
-          const s = state<CounterState>();
-          s.set("secondary", value);
-          const count = (yield* s.get("count")) ?? 0;
+          counterState.secondary.set(value);
+          const count: number = yield* counterState.count.get();
           return { count, secondary: value };
         })
       ),
@@ -88,15 +89,43 @@ const counterObj = restate.object({
       execute(
         ctx,
         gen(function* () {
-          state<CounterState>().clear("secondary");
+          counterState.secondary.clear();
         })
       ),
 
     reset: async (ctx: restate.ObjectContext): Promise<void> =>
       execute(
         ctx,
         gen(function* () {
-          state<CounterState>().clearAll();
+          counterState.count.clear();
+          counterState.secondary.clear();
+        })
+      ),
+  },
+});
+
+const listState = state({
+  items: { default: () => [] as string[] },
+});
+
+// Object for testing factory defaults.
+const listObj = restate.object({
+  name: "list",
+  handlers: {
+    get: async (ctx: restate.ObjectSharedContext) =>
+      execute(
+        ctx,
+        gen(function* () {
+          return { items: yield* listState.items.get() };
+        })
+      ),
+
+    append: async (ctx: restate.ObjectContext, item: string): Promise<void> =>
+      execute(
+        ctx,
+        gen(function* () {
+          const current = yield* listState.items.get();
+          listState.items.set([...current, item]);
         })
       ),
   },
@@ -113,7 +142,7 @@ describe.each(modes)("state — $name mode", ({ alwaysReplay }) => {
 
   beforeAll(async () => {
     env = await RestateTestEnvironment.start({
-      services: [counterObj],
+      services: [counterObj, listObj],
       alwaysReplay,
     });
     ingress = clients.connect({ url: env.baseUrl() });
@@ -132,15 +161,18 @@ describe.each(modes)("state — $name mode", ({ alwaysReplay }) => {
     expect(await client.current()).toBe(10);
   });
 
+  test("static default: count returns 0 for a fresh object", async () => {
+    const key = `fresh-${alwaysReplay ? "replay" : "default"}`;
+    expect(await ingress.objectClient(counterObj, key).current()).toBe(0);
+  });
+
   test("isolation: each VO key has independent state", async () => {
     const k1 = `iso1-${alwaysReplay ? "replay" : "default"}`;
     const k2 = `iso2-${alwaysReplay ? "replay" : "default"}`;
-    const c1 = ingress.objectClient(counterObj, k1);
-    const c2 = ingress.objectClient(counterObj, k2);
-    await c1.add(5);
-    await c2.add(11);
-    expect(await c1.current()).toBe(5);
-    expect(await c2.current()).toBe(11);
+    await ingress.objectClient(counterObj, k1).add(5);
+    await ingress.objectClient(counterObj, k2).add(11);
+    expect(await ingress.objectClient(counterObj, k1).current()).toBe(5);
+    expect(await ingress.objectClient(counterObj, k2).current()).toBe(11);
   });
 
   test("keys(): lists all set keys, drops cleared ones", async () => {
@@ -153,7 +185,7 @@ describe.each(modes)("state — $name mode", ({ alwaysReplay }) => {
     expect(await client.keys()).toEqual(["count"]);
   });
 
-  test("clearAll: wipes all state for this object", async () => {
+  test("reset: clears all keys", async () => {
     const key = `clear-${alwaysReplay ? "replay" : "default"}`;
     const client = ingress.objectClient(counterObj, key);
     await client.add(42);
@@ -163,4 +195,20 @@ describe.each(modes)("state — $name mode", ({ alwaysReplay }) => {
     expect(await client.keys()).toEqual([]);
     expect(await client.current()).toBe(0);
   });
+
+  test("factory default: fresh VO key starts with an empty array", async () => {
+    const key = `factory-${alwaysReplay ? "replay" : "default"}`;
+    const client = ingress.objectClient(listObj, key);
+    expect((await client.get()).items).toEqual([]);
+    await client.append("a");
+    await client.append("b");
+    expect((await client.get()).items).toEqual(["a", "b"]);
+  });
+
+  test("factory default: each VO key gets its own independent array", async () => {
+    const k1 = `factory-iso1-${alwaysReplay ? "replay" : "default"}`;
+    const k2 = `factory-iso2-${alwaysReplay ? "replay" : "default"}`;
+    await ingress.objectClient(listObj, k1).append("x");
+    expect((await ingress.objectClient(listObj, k2).get()).items).toEqual([]);
+  });
 });

packages/libs/restate-sdk-gen/examples/tutorial/src/07-state.ts

@@ -11,60 +11,107 @@
 
 // Tier 7: virtual-object state.
 //
-// Maps to guide.md §"Working with state". A virtual object owns a
-// per-key state slot the SDK persists between invocations. `state()`
-// gives you the read-write store; `sharedState()` is the read-only
-// view available from `ObjectSharedContext` handlers (concurrent reads
-// while no writer holds the key).
+// Two APIs:
 //
-// Typed state: pass a shape (`{ counter: number }`) to get keyof-checked
-// names and per-key value types. Without it, names are `string` and
-// values are inferred per call (untyped, like the SDK's default).
+//   state(config) — per-key typed accessors. Each key in the config gets a
+//     .get() / .set() / .clear() accessor. Keys with a `default` return
+//     Future<T> (never null); others return Future<T | null>.
+//
+//   state<TShape>() — per-key accessors without config. All keys return
+//     Future<T | null>.
+//
+//   getState / setState / clearState / clearAllState / getAllStateKeys —
+//     flat untyped functions for dynamic key names or simple cases.
 
 import * as restate from "@restatedev/restate-sdk";
-import { gen, execute, state, sharedState } from "@restatedev/restate-sdk-gen";
+import {
+  gen,
+  execute,
+  state,
+  getState,
+  setState,
+  clearState,
+  getAllStateKeys,
+} from "@restatedev/restate-sdk-gen";
 
-type CounterState = {
-  counter: number;
-};
+// Description of the state fields used by counter
+const counterState = state({
+  counter: { default: 0 },
+});
 
 export const counter = restate.object({
   name: "counter",
   handlers: {
-    // Read-only handler: takes ObjectSharedContext, uses sharedState().
-    // Multiple `get` invocations can run concurrently for the same key.
+    // Read-only handler: shared context, uses state() (same API — write
+    // methods throw at runtime in a shared context).
     get: async (ctx: restate.ObjectSharedContext): Promise<number> =>
       execute(
         ctx,
         gen(function* () {
-          return (yield* sharedState<CounterState>().get("counter")) ?? 0;
+          // Default is 0 — no null-coalescing needed.
+          return yield* counterState.counter.get();
         })
       ),
 
-    // Read-write handler: takes ObjectContext, uses state(). Exclusive
-    // access to the key for the duration of the invocation.
+    // Read-write handler: default applied, so no ?? needed.
     add: async (
       ctx: restate.ObjectContext,
       addend: number
     ): Promise<{ oldValue: number; newValue: number }> =>
       execute(
         ctx,
         gen(function* () {
-          const s = state<CounterState>();
-          const oldValue = (yield* s.get("counter")) ?? 0;
+          const oldValue: number = yield* counterState.counter.get(); // Future<number>
           const newValue = oldValue + addend;
-          s.set("counter", newValue);
+          counterState.counter.set(newValue);
           return { oldValue, newValue };
         })
       ),
 
-    // Clear the counter back to zero (well, deletes the entry; `get`
-    // returns 0 by falling through the `?? 0`).
     reset: async (ctx: restate.ObjectContext): Promise<void> =>
       execute(
         ctx,
         gen(function* () {
-          state<CounterState>().clear("counter");
+          counterState.counter.clear();
+        })
+      ),
+
+    // Flat untyped API — useful when key names are dynamic.
+    setRaw: async (
+      ctx: restate.ObjectContext,
+      payload: { key: string; value: string }
+    ): Promise<void> =>
+      execute(
+        ctx,
+        gen(function* () {
+          setState(payload.key, payload.value);
+        })
+      ),
+
+    getRaw: async (
+      ctx: restate.ObjectContext,
+      key: string
+    ): Promise<string | null> =>
+      execute(
+        ctx,
+        gen(function* () {
+          return yield* getState<string>(key);
+        })
+      ),
+
+    clearRaw: async (ctx: restate.ObjectContext, key: string): Promise<void> =>
+      execute(
+        ctx,
+        gen(function* () {
+          clearState(key);
+        })
+      ),
+
+    keys: async (ctx: restate.ObjectSharedContext): Promise<string[]> =>
+      execute(
+        ctx,
+        gen(function* () {
+          return yield* getAllStateKeys();
         })
       ),
   },

packages/libs/restate-sdk-gen/examples/tutorial/src/08-clients.ts

@@ -36,7 +36,6 @@ import {
   serviceSendClient,
   objectClient,
   state,
-  sharedState,
 } from "@restatedev/restate-sdk-gen";
 import type { counter } from "./07-state.js";
 
@@ -70,7 +69,7 @@ export const greeter = restate.service({
 // id often surfaces through an external system (webhook, queue) and
 // the holder VO is just a convenient stash.
 
-type HolderState = { id: string };
+const holderState = state<{ id: string }>();
 
 export const awakeableHolder = restate.object({
   name: "awakeableHolder",
@@ -79,7 +78,7 @@ export const awakeableHolder = restate.object({
       execute(
         ctx,
         gen(function* () {
-          state<HolderState>().set("id", id);
+          holderState.id.set(id);
         })
       ),
 
@@ -90,12 +89,12 @@ export const awakeableHolder = restate.object({
       execute(
         ctx,
         gen(function* () {
-          const id = yield* state<HolderState>().get("id");
+          const id = yield* holderState.id.get();
           if (!id) {
             throw new restate.TerminalError("no awakeable registered yet");
           }
           resolveAwakeable(id, payload);
-          state<HolderState>().clear("id");
+          holderState.id.clear();
         })
       ),
 
@@ -105,7 +104,7 @@ export const awakeableHolder = restate.object({
       execute(
         ctx,
         gen(function* () {
-          return (yield* sharedState<HolderState>().get("id")) ?? null;
+          return yield* holderState.id.get();
         })
       ),
   },