Merge pull request #2659 from garden-co/changeset-release/main

Version Packages

biome.json

@@ -14,7 +14,7 @@
       "tests/jazz-svelte/src/**",
       "examples/*svelte*/**",
       "starters/*svelte*/**",
-      "examples/jazz-paper-scissors/src/routeTree.gen.ts",
+      "examples/server-worker-inbox/src/routeTree.gen.ts",
       "homepage/homepage/**",
       "**/package.json"
     ]

examples/chat-svelte/CHANGELOG.md

@@ -1,5 +1,35 @@
 # passkey-svelte
 
+## 0.0.103
+
+### Patch Changes
+
+- Updated dependencies [a584590]
+- Updated dependencies [9acccb5]
+  - jazz-tools@0.15.14
+
+## 0.0.102
+
+### Patch Changes
+
+- Updated dependencies [6c76ff8]
+  - jazz-tools@0.15.13
+
+## 0.0.101
+
+### Patch Changes
+
+- Updated dependencies [d1c1b0c]
+- Updated dependencies [cf4ad72]
+  - jazz-tools@0.15.12
+
+## 0.0.100
+
+### Patch Changes
+
+- Updated dependencies [bdc9aee]
+  - jazz-tools@0.15.11
+
 ## 0.0.99
 
 ### Patch Changes

examples/chat-svelte/package.json

@@ -1,6 +1,6 @@
 {
   "name": "chat-svelte",
-  "version": "0.0.99",
+  "version": "0.0.103",
   "type": "module",
   "private": true,
   "scripts": {

examples/chat/src/schema.ts

@@ -1,8 +1,8 @@
-import { co, z } from "jazz-tools";
+import { co } from "jazz-tools";
 
 export const Message = co.map({
   text: co.plainText(),
-  image: z.optional(co.image()),
+  image: co.optional(co.image()),
 });
 export type Message = co.loaded<typeof Message>;
 

examples/filestream/src/schema.ts

@@ -1,7 +1,7 @@
-import { co, z } from "jazz-tools";
+import { co } from "jazz-tools";
 
 export const JazzProfile = co.profile({
-  file: z.optional(co.fileStream()),
+  file: co.optional(co.fileStream()),
 });
 
 export const JazzAccount = co.account({

examples/form/src/schema.ts

@@ -28,16 +28,16 @@ export const BubbleTeaOrder = co.map({
   addOns: ListOfBubbleTeaAddOns,
   deliveryDate: z.date(),
   withMilk: z.boolean(),
-  instructions: z.optional(co.plainText()),
+  instructions: co.optional(co.plainText()),
 });
 
 export const DraftBubbleTeaOrder = co
   .map({
     baseTea: z.optional(z.literal([...BubbleTeaBaseTeaTypes])),
-    addOns: z.optional(ListOfBubbleTeaAddOns),
+    addOns: co.optional(ListOfBubbleTeaAddOns),
     deliveryDate: z.optional(z.date()),
     withMilk: z.optional(z.boolean()),
-    instructions: z.optional(co.plainText()),
+    instructions: co.optional(co.plainText()),
   })
   .withHelpers((Self) => ({
     hasChanges(order: Loaded<typeof Self> | undefined) {

examples/image-upload/src/schema.ts

@@ -1,7 +1,7 @@
-import { co, z } from "jazz-tools";
+import { co } from "jazz-tools";
 
 export const JazzProfile = co.profile({
-  image: z.optional(co.image()),
+  image: co.optional(co.image()),
 });
 
 export const JazzAccount = co.account({

examples/music-player/src/1_schema.ts

@@ -66,7 +66,7 @@ export const MusicaAccountRoot = co.map({
   // track and playlist
   // You can also add the position in time if you want make it possible
   // to resume the song
-  activeTrack: z.optional(MusicTrack),
+  activeTrack: co.optional(MusicTrack),
   activePlaylist: Playlist,
 
   exampleDataLoaded: z.optional(z.boolean()),

examples/server-worker-http/vercel.json

@@ -0,0 +1,3 @@
+{
+  "ignoreCommand": "echo true"
+}

examples/server-worker-inbox/src/schema.ts

@@ -8,7 +8,7 @@ export type Player = co.loaded<typeof Player>;
 
 export const Game = co.map({
   player1: Player,
-  player2: z.optional(Player),
+  player2: co.optional(Player),
   outcome: z.optional(z.literal(["player1", "player2", "draw"])),
   player1Score: z.number(),
   player2Score: z.number(),
@@ -17,8 +17,8 @@ export type Game = co.loaded<typeof Game>;
 
 export const WaitingRoom = co.map({
   account1: co.account(),
-  account2: z.optional(co.account()),
-  game: z.optional(Game),
+  account2: co.optional(co.account()),
+  game: co.optional(Game),
 });
 export type WaitingRoom = co.loaded<typeof WaitingRoom>;
 
@@ -47,7 +47,7 @@ export const JoinGameRequest = co.map({
 });
 export type JoinGameRequest = co.loaded<typeof JoinGameRequest>;
 
-export const InboxMessage = z.discriminatedUnion("type", [
+export const InboxMessage = co.discriminatedUnion("type", [
   PlayIntent,
   NewGameIntent,
   CreateGameRequest,

homepage/homepage/content/docs/authentication/authentication-states.mdx

@@ -94,7 +94,7 @@ const MusicaAccountRoot = co.map({
 
 const MusicaAccount = co.account({
   root: MusicaAccountRoot,
-  profile: co.profile({}),
+  profile: co.profile(),
 });
 type MusicaAccount = co.loaded<typeof MusicaAccount>
 

homepage/homepage/content/docs/design-patterns/organization.mdx

@@ -67,7 +67,7 @@ export const JazzAccountRoot = co.map({
 export const JazzAccount = co
   .account({
     root: JazzAccountRoot,
-    profile: co.profile({}),
+    profile: co.profile(),
   })
   .withMigration((account) => {
     if (account.root === undefined) {
@@ -130,7 +130,7 @@ const JazzAccountRoot = co.map({
 
 const JazzAccount = co.account({
   root: JazzAccountRoot,
-  profile: co.profile({}),
+  profile: co.profile(),
 });
 
 // ---cut---

homepage/homepage/content/docs/project-setup/server-side.mdx

@@ -11,7 +11,7 @@ The main detail to understand when using Jazz server-side is that Server Workers
 
 This lets you share CoValues with Server Workers, having precise access control by adding the Worker to `Groups` with specific roles just like you would with other users.
 
-[See the full example here.](https://github.com/garden-co/jazz/tree/main/examples/jazz-paper-scissors)
+[See the full example here.](https://github.com/garden-co/jazz/tree/main/examples/server-worker-inbox)
 
 <Alert variant="info" className="mt-4 flex gap-2 items-center">Requires at least Node.js v20.</Alert>
 

homepage/homepage/content/docs/schemas/covalues.mdx

@@ -202,7 +202,7 @@ See the corresponding sections for [creating](/docs/using-covalues/filestreams#c
 
 ### Unions of CoMaps (declaration)
 
-You can declare unions of CoMaps that have discriminating fields, using `z.discriminatedUnion()`.
+You can declare unions of CoMaps that have discriminating fields, using `co.discriminatedUnion()`.
 
 <CodeGroup>
 ```ts twoslash
@@ -220,7 +220,7 @@ const SliderWidget = co.map({
   max: z.number(),
 });
 
-const WidgetUnion = z.discriminatedUnion("type", [ButtonWidget, SliderWidget]);
+const WidgetUnion = co.discriminatedUnion("type", [ButtonWidget, SliderWidget]);
 ```
 
 </CodeGroup>
@@ -321,7 +321,7 @@ const Company = co.map({
 
 #### Optional References
 
-You can make references optional with `z.optional()`:
+You can make references optional with `co.optional()`:
 
 <CodeGroup>
 ```ts twoslash
@@ -331,7 +331,7 @@ const Pet = co.map({
 });
 // ---cut---
 const Person = co.map({
-  pet: z.optional(Pet),
+  pet: co.optional(Pet),
 });
 ```
 </CodeGroup>
@@ -372,7 +372,7 @@ const ListOfPeople = co.list(Person);
 
 </CodeGroup>
 
-Note: similarly, if you use modifiers like `z.optional()` you'll need to help TypeScript along:
+Note: similarly, if you use modifiers like `co.optional()` you'll need to help TypeScript along:
 
 <CodeGroup>
 ```ts twoslash
@@ -381,7 +381,7 @@ import { co, z } from "jazz-tools";
 const Person = co.map({
   name: z.string(),
   get bestFriend(): z.ZodOptional<typeof Person> {
-    return z.optional(Person);
+    return co.optional(Person);
   }
 });
 ```

homepage/homepage/content/docs/using-covalues/comaps.mdx

@@ -24,7 +24,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 export type Project = co.loaded<typeof Project>;
 ```
@@ -54,7 +54,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 const Inventory = co.record(z.string(), z.number());
 // ---cut---
@@ -90,7 +90,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 
 // ---cut---
@@ -134,7 +134,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 const project = Project.create(
   {
@@ -165,7 +165,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 const project = Project.create(
   {
@@ -197,7 +197,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
   get subProject() {
     return Project.optional();
   }
@@ -219,9 +219,9 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
   get subProjects(): z.ZodOptional<CoListSchema<typeof Project>> {
-    return z.optional(co.list(Project));
+    return co.optional(co.list(Project));
   }
 });
 export type Project = co.loaded<typeof Project>;
@@ -265,7 +265,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 const Inventory = co.record(z.string(), z.number());
 const project = Project.create(
@@ -297,7 +297,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 const Inventory = co.record(z.string(), z.number());
 const project = Project.create(
@@ -347,7 +347,7 @@ const Project = co.map({
   name: z.string(),
   startDate: z.date(),
   status: z.literal(["planning", "active", "completed"]),
-  coordinator: z.optional(Member),
+  coordinator: co.optional(Member),
 });
 const Inventory = co.record(z.string(), z.number());
 const project = Project.create(
@@ -445,7 +445,7 @@ const TaskV2 = co.map({
 
 // Export the discriminated union; because some users might
 // not be able to run the migration
-export const Task = z.discriminatedUnion("version", [
+export const Task = co.discriminatedUnion("version", [
   TaskV1,
   TaskV2,
 ]);

packages/cojson-storage-indexeddb/CHANGELOG.md

@@ -1,5 +1,30 @@
 # cojson-storage-indexeddb
 
+## 0.15.14
+
+### Patch Changes
+
+- Updated dependencies [70ce7c5]
+  - cojson@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- cojson@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- cojson@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- cojson@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/cojson-storage-indexeddb/package.json

@@ -1,6 +1,6 @@
 {
   "name": "cojson-storage-indexeddb",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "main": "dist/index.js",
   "type": "module",
   "types": "dist/index.d.ts",

packages/cojson-storage-sqlite/CHANGELOG.md

@@ -1,5 +1,30 @@
 # cojson-storage-sqlite
 
+## 0.15.14
+
+### Patch Changes
+
+- Updated dependencies [70ce7c5]
+  - cojson@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- cojson@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- cojson@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- cojson@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/cojson-storage-sqlite/package.json

@@ -1,7 +1,7 @@
 {
   "name": "cojson-storage-sqlite",
   "type": "module",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "license": "MIT",

packages/cojson-transport-ws/CHANGELOG.md

@@ -1,5 +1,31 @@
 # cojson-transport-nodejs-ws
 
+## 0.15.14
+
+### Patch Changes
+
+- 70ce7c5: Introduce the persistent peers. Used to mark the WebSocket connections to the server as persistent, and wait for reconnection before failing load.
+- Updated dependencies [70ce7c5]
+  - cojson@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- cojson@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- cojson@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- cojson@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/cojson-transport-ws/package.json

@@ -1,7 +1,7 @@
 {
   "name": "cojson-transport-ws",
   "type": "module",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "license": "MIT",

packages/cojson-transport-ws/src/createWebSocketPeer.ts

@@ -165,6 +165,6 @@ export function createWebSocketPeer({
     incoming,
     outgoing,
     role,
-    deletePeerStateOnClose,
+    persistent: !deletePeerStateOnClose,
   };
 }

packages/cojson/CHANGELOG.md

@@ -1,5 +1,17 @@
 # cojson
 
+## 0.15.14
+
+### Patch Changes
+
+- 70ce7c5: Introduce the persistent peers. Used to mark the WebSocket connections to the server as persistent, and wait for reconnection before failing load.
+
+## 0.15.13
+
+## 0.15.12
+
+## 0.15.11
+
 ## 0.15.10
 
 ## 0.15.9

packages/cojson/package.json

@@ -25,7 +25,7 @@
   },
   "type": "module",
   "license": "MIT",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "devDependencies": {
     "@opentelemetry/sdk-metrics": "^2.0.0",
     "libsql": "^0.5.13",

packages/cojson/src/PeerState.ts

@@ -97,6 +97,10 @@ export class PeerState {
     return this.peer.incoming;
   }
 
+  get persistent() {
+    return this.peer.persistent;
+  }
+
   pushOutgoingMessage(msg: SyncMessage) {
     this.peer.outgoing.push(msg);
   }

packages/cojson/src/coValueCore/coValueCore.ts

@@ -1056,37 +1056,33 @@ export class CoValueCore {
       return;
     }
 
-    const peersToActuallyLoadFrom = [] as PeerState[];
-
     for (const peer of peers) {
-      const currentState = this.peers.get(peer.id)?.type;
+      const currentState = this.peers.get(peer.id)?.type ?? "unknown";
 
-      if (
-        !currentState ||
-        currentState === "unknown" ||
-        currentState === "unavailable"
-      ) {
-        peersToActuallyLoadFrom.push(peer);
+      if (currentState === "unknown" || currentState === "unavailable") {
         this.markPending(peer.id);
+        this.internalLoadFromPeer(peer);
       }
     }
-
-    for (const peer of peersToActuallyLoadFrom) {
-      this.internalLoadFromPeer(peer);
-    }
   }
 
   internalLoadFromPeer(peer: PeerState) {
-    if (peer.closed) {
+    if (peer.closed && !peer.persistent) {
       this.markNotFoundInPeer(peer.id);
       return;
     }
 
-    peer.pushOutgoingMessage({
-      action: "load",
-      ...this.knownState(),
-    });
-    peer.trackLoadRequestSent(this.id);
+    /**
+     * On reconnection persistent peers will automatically fire the load request
+     * as part of the reconnection process.
+     */
+    if (!peer.closed) {
+      peer.pushOutgoingMessage({
+        action: "load",
+        ...this.knownState(),
+      });
+      peer.trackLoadRequestSent(this.id);
+    }
 
     return new Promise<void>((resolve) => {
       const markNotFound = () => {
@@ -1100,7 +1096,9 @@ export class CoValueCore {
       };
 
       const timeout = setTimeout(markNotFound, CO_VALUE_LOADING_CONFIG.TIMEOUT);
-      const removeCloseListener = peer.addCloseListener(markNotFound);
+      const removeCloseListener = peer.persistent
+        ? undefined
+        : peer.addCloseListener(markNotFound);
 
       const listener = (state: CoValueCore) => {
         const peerState = state.peers.get(peer.id);
@@ -1111,7 +1109,7 @@ export class CoValueCore {
           peerState?.type === "unavailable"
         ) {
           this.listeners.delete(listener);
-          removeCloseListener();
+          removeCloseListener?.();
           clearTimeout(timeout);
           resolve();
         }

packages/cojson/src/config.ts

@@ -13,6 +13,14 @@ export const CO_VALUE_LOADING_CONFIG = {
   RETRY_DELAY: 3000,
 };
 
+export function setCoValueLoadingMaxRetries(maxRetries: number) {
+  CO_VALUE_LOADING_CONFIG.MAX_RETRIES = maxRetries;
+}
+
+export function setCoValueLoadingTimeout(timeout: number) {
+  CO_VALUE_LOADING_CONFIG.TIMEOUT = timeout;
+}
+
 export function setCoValueLoadingRetryDelay(delay: number) {
   CO_VALUE_LOADING_CONFIG.RETRY_DELAY = delay;
 }

packages/cojson/src/localNode.ts

@@ -82,6 +82,11 @@ export class LocalNode {
     this.storage = storage;
   }
 
+  removeStorage() {
+    this.storage?.close();
+    this.storage = undefined;
+  }
+
   getCoValue(id: RawCoID) {
     let entry = this.coValues.get(id);
 
@@ -348,12 +353,10 @@ export class LocalNode {
     skipLoadingFromPeer?: PeerID,
     skipRetry?: boolean,
   ): Promise<CoValueCore> {
-    if (!id) {
-      throw new Error("Trying to load CoValue with undefined id");
-    }
-
-    if (!id.startsWith("co_z")) {
-      throw new Error(`Trying to load CoValue with invalid id ${id}`);
+    if (typeof id !== "string" || !id.startsWith("co_z")) {
+      throw new TypeError(
+        `Trying to load CoValue with invalid id ${Array.isArray(id) ? JSON.stringify(id) : id}`,
+      );
     }
 
     if (this.crashed) {

packages/cojson/src/streamUtils.ts

@@ -13,9 +13,11 @@ export function connectedPeers(
   {
     peer1role = "client",
     peer2role = "client",
+    persistent = false,
   }: {
     peer1role?: Peer["role"];
     peer2role?: Peer["role"];
+    persistent?: boolean;
   } = {},
 ): [Peer, Peer] {
   const from1to2 = new ConnectedPeerChannel();
@@ -26,6 +28,7 @@ export function connectedPeers(
     incoming: from2to1,
     outgoing: from1to2,
     role: peer2role,
+    persistent,
   };
 
   const peer1AsPeer: Peer = {
@@ -33,6 +36,7 @@ export function connectedPeers(
     incoming: from1to2,
     outgoing: from2to1,
     role: peer1role,
+    persistent,
   };
 
   return [peer1AsPeer, peer2AsPeer];

packages/cojson/src/sync.ts

@@ -88,7 +88,7 @@ export interface Peer {
   outgoing: OutgoingPeerChannel;
   role: "server" | "client";
   priority?: number;
-  deletePeerStateOnClose?: boolean;
+  persistent?: boolean;
 }
 
 export function combinedKnownStates(
@@ -157,8 +157,7 @@ export class SyncManager {
 
   getServerPeers(excludePeerId?: PeerID): PeerState[] {
     return this.getPeers().filter(
-      (peer) =>
-        peer.role === "server" && peer.id !== excludePeerId && !peer.closed,
+      (peer) => peer.role === "server" && peer.id !== excludePeerId,
     );
   }
 
@@ -353,7 +352,7 @@ export class SyncManager {
       unsubscribeFromKnownStatesUpdates();
       this.peersCounter.add(-1, { role: peer.role });
 
-      if (peer.deletePeerStateOnClose && this.peers[peer.id] === peerState) {
+      if (!peer.persistent && this.peers[peer.id] === peerState) {
         delete this.peers[peer.id];
       }
     });
@@ -794,8 +793,8 @@ export class SyncManager {
 
     const peerState = this.peers[peerId];
 
-    // The peer has been closed, so it isn't possible to sync
-    if (!peerState || peerState.closed) {
+    // The peer has been closed and is not persistent, so it isn't possible to sync
+    if (!peerState) {
       return;
     }
 
@@ -830,7 +829,7 @@ export class SyncManager {
     return this.local.storage?.waitForSync(id, this.local.getCoValue(id));
   }
 
-  waitForSync(id: RawCoID, timeout = 30_000) {
+  waitForSync(id: RawCoID, timeout = 60_000) {
     const peers = this.getPeers();
 
     return Promise.all(

packages/cojson/src/tests/SyncStateManager.test.ts

@@ -251,9 +251,11 @@ describe("SyncStateManager", () => {
     ).toEqual({ uploaded: true });
   });
 
-  test("should skip closed peers", async () => {
+  test("should skip non-persistent closed peers", async () => {
     const client = setupTestNode();
-    const { peerState } = client.connectToSyncServer();
+    const { peerState } = client.connectToSyncServer({
+      persistent: false,
+    });
 
     peerState.gracefulShutdown();
 
@@ -263,6 +265,42 @@ describe("SyncStateManager", () => {
     await map.core.waitForSync();
   });
 
+  test("should wait for persistent closed peers to reconnect", async () => {
+    const client = setupTestNode();
+    const { peerState } = client.connectToSyncServer({
+      persistent: true,
+    });
+
+    peerState.gracefulShutdown();
+
+    const group = client.node.createGroup();
+    const map = group.createMap();
+
+    const promise = map.core.waitForSync().then(() => "waitForSync");
+
+    const result = await Promise.race([
+      promise,
+      new Promise((resolve) => {
+        setTimeout(() => resolve("timeout"), 10);
+      }),
+    ]);
+
+    expect(result).toBe("timeout");
+
+    client.connectToSyncServer({
+      persistent: true,
+    });
+
+    const result2 = await Promise.race([
+      promise,
+      new Promise((resolve) => {
+        setTimeout(() => resolve("timeout"), 10);
+      }),
+    ]);
+
+    expect(result2).toBe("waitForSync");
+  });
+
   test("should skip client peers that are not subscribed to the coValue", async () => {
     const server = setupTestNode({ isSyncServer: true });
     const client = setupTestNode();

packages/cojson/src/tests/sync.load.test.ts

@@ -1,6 +1,9 @@
 import { beforeEach, describe, expect, test, vi } from "vitest";
 import { expectMap } from "../coValue";
-import { setCoValueLoadingRetryDelay } from "../config";
+import {
+  CO_VALUE_LOADING_CONFIG,
+  setCoValueLoadingRetryDelay,
+} from "../config";
 import { RawCoMap } from "../exports";
 import {
   SyncMessagesLog,
@@ -54,6 +57,43 @@ describe("loading coValues from server", () => {
     `);
   });
 
+  test("coValue load throws on invalid id", async () => {
+    const { node } = setupTestNode({
+      connected: true,
+    });
+
+    await expect(async () => await node.load("test" as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id test",
+    );
+    await expect(async () => await node.load(null as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id null",
+    );
+    await expect(async () => await node.load(undefined as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id undefined",
+    );
+    await expect(async () => await node.load(1 as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id 1",
+    );
+    await expect(async () => await node.load({} as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id [object Object]",
+    );
+    await expect(async () => await node.load([] as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id []",
+    );
+    await expect(async () => await node.load(["test"] as any)).rejects.toThrow(
+      'Trying to load CoValue with invalid id ["test"]',
+    );
+    await expect(
+      async () => await node.load((() => {}) as any),
+    ).rejects.toMatchInlineSnapshot(`
+      [TypeError: Trying to load CoValue with invalid id () => {
+            }]
+    `);
+    await expect(
+      async () => await node.load(new Date() as any),
+    ).rejects.toThrow();
+  });
+
   test("unavailable coValue retry with skipRetry set to true", async () => {
     const client = setupTestNode();
     const client2 = setupTestNode();
@@ -470,9 +510,12 @@ describe("loading coValues from server", () => {
     `);
   });
 
-  test("should mark the coValue as unavailable if the peer is closed", async () => {
+  test("should wait for a persistent peer to reconnect before marking the coValue as unavailable", async () => {
     const client = setupTestNode();
-    const { peerState } = client.connectToSyncServer();
+    const connection1 = client.connectToSyncServer({
+      persistent: true,
+    });
+    connection1.peerState.gracefulShutdown();
 
     const group = jazzCloud.node.createGroup();
     group.addMember("everyone", "writer");
@@ -480,13 +523,15 @@ describe("loading coValues from server", () => {
     const map = group.createMap({
       test: "value",
     });
-
     const promise = client.node.load(map.id);
 
-    // Close the peer connection
-    peerState.gracefulShutdown();
+    await new Promise((resolve) => setTimeout(resolve, 10));
 
-    expect(await promise).toEqual("unavailable");
+    client.connectToSyncServer();
+
+    const coValue = await promise;
+
+    expect(coValue).not.toBe("unavailable");
 
     expect(
       SyncMessagesLog.getMessages({
@@ -496,6 +541,60 @@ describe("loading coValues from server", () => {
     ).toMatchInlineSnapshot(`
       [
         "client -> server | LOAD Map sessions: empty",
+        "server -> client | CONTENT Group header: true new: After: 0 New: 5",
+        "server -> client | CONTENT Map header: true new: After: 0 New: 1",
+        "client -> server | KNOWN Group sessions: header/5",
+        "client -> server | KNOWN Map sessions: header/1",
+      ]
+    `);
+  });
+
+  test("should handle reconnections in the middle of a load with a persistent peer", async () => {
+    const client = setupTestNode();
+    const connection1 = client.connectToSyncServer({
+      persistent: true,
+    });
+
+    const group = jazzCloud.node.createGroup();
+    group.addMember("everyone", "writer");
+
+    const map = group.createMap({
+      test: "value",
+    });
+
+    blockMessageTypeOnOutgoingPeer(connection1.peerOnServer, "content", {
+      id: map.id,
+      once: true,
+    });
+
+    const promise = client.node.load(map.id);
+
+    await new Promise((resolve) => setTimeout(resolve, 10));
+
+    // Close the peer connection
+    connection1.peerState.gracefulShutdown();
+
+    client.connectToSyncServer();
+
+    const coValue = await promise;
+
+    expect(coValue).not.toBe("unavailable");
+
+    expect(
+      SyncMessagesLog.getMessages({
+        Group: group.core,
+        Map: map.core,
+      }),
+    ).toMatchInlineSnapshot(`
+      [
+        "client -> server | LOAD Map sessions: empty",
+        "server -> client | CONTENT Group header: true new: After: 0 New: 5",
+        "client -> server | KNOWN Group sessions: header/5",
+        "client -> server | LOAD Map sessions: empty",
+        "client -> server | LOAD Group sessions: header/5",
+        "server -> client | KNOWN Group sessions: header/5",
+        "server -> client | CONTENT Map header: true new: After: 0 New: 1",
+        "client -> server | KNOWN Map sessions: header/1",
       ]
     `);
   });
@@ -872,4 +971,62 @@ describe("loading coValues from server", () => {
       ]
     `);
   });
+
+  test("should retry loading from a closed persistent peer after a timeout", async () => {
+    vi.useFakeTimers();
+
+    const client = setupTestNode();
+
+    const connection1 = client.connectToSyncServer({
+      persistent: true,
+    });
+
+    // Close the peer connection
+    connection1.peerState.gracefulShutdown();
+
+    const group = jazzCloud.node.createGroup();
+    group.addMember("everyone", "reader");
+
+    const map = group.createMap();
+    map.set("hello", "world");
+
+    const promise = loadCoValueOrFail(client.node, map.id);
+
+    await vi.advanceTimersByTimeAsync(
+      CO_VALUE_LOADING_CONFIG.TIMEOUT +
+        CO_VALUE_LOADING_CONFIG.RETRY_DELAY +
+        10,
+    );
+
+    client.connectToSyncServer({
+      persistent: true,
+    });
+
+    await vi.advanceTimersByTimeAsync(
+      CO_VALUE_LOADING_CONFIG.TIMEOUT +
+        CO_VALUE_LOADING_CONFIG.RETRY_DELAY +
+        10,
+    );
+
+    const coValue = await promise;
+
+    expect(coValue).not.toBe("unavailable");
+
+    expect(
+      SyncMessagesLog.getMessages({
+        Group: group.core,
+        Map: map.core,
+      }),
+    ).toMatchInlineSnapshot(`
+      [
+        "client -> server | LOAD Map sessions: empty",
+        "server -> client | CONTENT Group header: true new: After: 0 New: 5",
+        "server -> client | CONTENT Map header: true new: After: 0 New: 1",
+        "client -> server | KNOWN Group sessions: header/5",
+        "client -> server | KNOWN Map sessions: header/1",
+      ]
+    `);
+
+    vi.useRealTimers();
+  });
 });

packages/cojson/src/tests/sync.mesh.test.ts

@@ -26,6 +26,7 @@ function setupMesh() {
     ourName: "edge-italy",
     syncServerName: "core",
     syncServer: coreServer.node,
+    persistent: true,
   });
   edgeItaly.addStorage({
     ourName: "edge-italy",
@@ -36,6 +37,7 @@ function setupMesh() {
     ourName: "edge-france",
     syncServerName: "core",
     syncServer: coreServer.node,
+    persistent: true,
   });
   edgeFrance.addStorage({
     ourName: "edge-france",

packages/cojson/src/tests/sync.test.ts

@@ -111,10 +111,12 @@ test("Can sync a coValue with private transactions through a server to another c
   expect(mapOnClient2.get("hello")).toEqual("world");
 });
 
-test("should keep the peer state when the peer closes", async () => {
+test("should keep the peer state when the peer closes if persistent is true", async () => {
   const client = setupTestNode();
 
-  const { peer, peerState, peerOnServer } = client.connectToSyncServer();
+  const { peer, peerState, peerOnServer } = client.connectToSyncServer({
+    persistent: true,
+  });
 
   const group = jazzCloud.node.createGroup();
   const map = group.createMap();
@@ -133,12 +135,12 @@ test("should keep the peer state when the peer closes", async () => {
   expect(syncManager.peers[peer.id]).not.toBeUndefined();
 });
 
-test("should delete the peer state when the peer closes if deletePeerStateOnClose is true", async () => {
+test("should delete the peer state when the peer closes if persistent is false", async () => {
   const client = setupTestNode();
 
-  const { peer, peerState, peerOnServer } = client.connectToSyncServer();
-
-  peer.deletePeerStateOnClose = true;
+  const { peer, peerState, peerOnServer } = client.connectToSyncServer({
+    persistent: false,
+  });
 
   const group = jazzCloud.node.createGroup();
   const map = group.createMap();
@@ -991,7 +993,7 @@ describe("LocalNode.load", () => {
 
     // @ts-expect-error Testing with undefined ID
     await expect(client.node.load(undefined)).rejects.toThrow(
-      "Trying to load CoValue with undefined id",
+      "Trying to load CoValue with invalid id undefined",
     );
   });
 

packages/cojson/src/tests/testUtils.ts

@@ -429,6 +429,7 @@ export function getSyncServerConnectedPeer(opts: {
   ourName?: string;
   syncServer?: LocalNode;
   peerId: string;
+  persistent?: boolean;
 }) {
   const currentSyncServer = opts?.syncServer ?? syncServer.current;
 
@@ -451,6 +452,7 @@ export function getSyncServerConnectedPeer(opts: {
       role: "client",
       name: opts.ourName,
     },
+    persistent: opts?.persistent,
   });
 
   currentSyncServer.syncManager.addPeer(peer2);
@@ -483,6 +485,7 @@ export function setupTestNode(
     syncServerName?: string;
     ourName?: string;
     syncServer?: LocalNode;
+    persistent?: boolean;
   }) {
     const { peer, peerStateOnServer, peerOnServer } =
       getSyncServerConnectedPeer({
@@ -490,6 +493,7 @@ export function setupTestNode(
         syncServerName: opts?.syncServerName,
         ourName: opts?.ourName,
         syncServer: opts?.syncServer,
+        persistent: opts?.persistent,
       });
 
     node.syncManager.addPeer(peer);
@@ -646,12 +650,22 @@ export type SyncTestMessage = {
 export function connectedPeersWithMessagesTracking(opts: {
   peer1: { id: string; role: Peer["role"]; name?: string };
   peer2: { id: string; role: Peer["role"]; name?: string };
+  persistent?: boolean;
 }) {
   const [peer1, peer2] = connectedPeers(opts.peer1.id, opts.peer2.id, {
     peer1role: opts.peer1.role,
     peer2role: opts.peer2.role,
+    persistent: opts.persistent,
   });
 
+  // If the persistent option is not provided, we default to true for the server and false for the client
+  // Trying to mimic the real world behavior of the sync server
+  if (opts.persistent === undefined) {
+    peer1.persistent = opts.peer1.role === "server";
+
+    peer2.persistent = opts.peer2.role === "server";
+  }
+
   const peer1Push = peer1.outgoing.push;
   peer1.outgoing.push = (msg) => {
     if (typeof msg !== "string") {

packages/cursor-docs/.cursor/docs/llms-full.md

@@ -2136,7 +2136,7 @@ See the corresponding sections for [creating](/docs/using-covalues/filestreams#c
 
 ### Unions of CoMaps (declaration)
 
-You can declare unions of CoMaps that have discriminating fields, using `z.discriminatedUnion()`.
+You can declare unions of CoMaps that have discriminating fields, using `co.discriminatedUnion()`.
 
 <CodeGroup>
 ```ts twoslash
@@ -2153,7 +2153,7 @@ const SliderWidget = co.map({
   max: z.number(),
 });
 
-const WidgetUnion = z.discriminatedUnion([ButtonWidget, SliderWidget]);
+const WidgetUnion = co.discriminatedUnion([ButtonWidget, SliderWidget]);
 ```
 </CodeGroup>
 
@@ -7890,7 +7890,7 @@ export const JazzAccountRoot = co.map({
 export const JazzAccount = co
   .account({
     root: JazzAccountRoot,
-    profile: co.profile({}),
+    profile: co.profile(),
   })
   .withMigration((account) => {
     if (account.root === undefined) {
@@ -7950,7 +7950,7 @@ const JazzAccountRoot = co.map({
 
 const JazzAccount = co.account({
   root: JazzAccountRoot,
-  profile: co.profile({}),
+  profile: co.profile(),
 });
 
 // ---cut---
@@ -11911,4 +11911,4 @@ export function cn(...inputs: ClassValue[]) {
 ```ts
 /// <reference types="vite/client" />
 
-```
+```

packages/jazz-auth-betterauth/CHANGELOG.md

@@ -1,5 +1,44 @@
 # jazz-auth-betterauth
 
+## 0.15.14
+
+### Patch Changes
+
+- Updated dependencies [70ce7c5]
+- Updated dependencies [a584590]
+- Updated dependencies [9acccb5]
+  - cojson@0.15.14
+  - jazz-tools@0.15.14
+  - jazz-betterauth-client-plugin@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- Updated dependencies [6c76ff8]
+  - jazz-tools@0.15.13
+  - jazz-betterauth-client-plugin@0.15.13
+  - cojson@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- Updated dependencies [d1c1b0c]
+- Updated dependencies [cf4ad72]
+  - jazz-tools@0.15.12
+  - jazz-betterauth-client-plugin@0.15.12
+  - cojson@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- Updated dependencies [bdc9aee]
+  - jazz-tools@0.15.11
+  - jazz-betterauth-client-plugin@0.15.11
+  - cojson@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/jazz-auth-betterauth/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "jazz-auth-betterauth",
-	"version": "0.15.10",
+	"version": "0.15.14",
 	"type": "module",
 	"main": "dist/index.js",
 	"types": "src/index.ts",

packages/jazz-betterauth-client-plugin/CHANGELOG.md

@@ -1,5 +1,29 @@
 # jazz-betterauth-client-plugin
 
+## 0.15.14
+
+### Patch Changes
+
+- jazz-betterauth-server-plugin@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- jazz-betterauth-server-plugin@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- jazz-betterauth-server-plugin@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- jazz-betterauth-server-plugin@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/jazz-betterauth-client-plugin/package.json

@@ -1,6 +1,6 @@
 {
   "name": "jazz-betterauth-client-plugin",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "type": "module",
   "main": "dist/index.js",
   "types": "src/index.ts",

packages/jazz-betterauth-server-plugin/CHANGELOG.md

@@ -1,5 +1,40 @@
 # jazz-betterauth-server-plugin
 
+## 0.15.14
+
+### Patch Changes
+
+- Updated dependencies [70ce7c5]
+- Updated dependencies [a584590]
+- Updated dependencies [9acccb5]
+  - cojson@0.15.14
+  - jazz-tools@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- Updated dependencies [6c76ff8]
+  - jazz-tools@0.15.13
+  - cojson@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- Updated dependencies [d1c1b0c]
+- Updated dependencies [cf4ad72]
+  - jazz-tools@0.15.12
+  - cojson@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- Updated dependencies [bdc9aee]
+  - jazz-tools@0.15.11
+  - cojson@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/jazz-betterauth-server-plugin/package.json

@@ -1,6 +1,6 @@
 {
   "name": "jazz-betterauth-server-plugin",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "type": "module",
   "main": "dist/index.js",
   "types": "src/index.ts",

packages/jazz-react-auth-betterauth/CHANGELOG.md

@@ -1,5 +1,48 @@
 # jazz-react-auth-betterauth
 
+## 0.15.14
+
+### Patch Changes
+
+- Updated dependencies [70ce7c5]
+- Updated dependencies [a584590]
+- Updated dependencies [9acccb5]
+  - cojson@0.15.14
+  - jazz-tools@0.15.14
+  - jazz-auth-betterauth@0.15.14
+  - jazz-betterauth-client-plugin@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- Updated dependencies [6c76ff8]
+  - jazz-tools@0.15.13
+  - jazz-auth-betterauth@0.15.13
+  - jazz-betterauth-client-plugin@0.15.13
+  - cojson@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- Updated dependencies [d1c1b0c]
+- Updated dependencies [cf4ad72]
+  - jazz-tools@0.15.12
+  - jazz-auth-betterauth@0.15.12
+  - jazz-betterauth-client-plugin@0.15.12
+  - cojson@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- Updated dependencies [bdc9aee]
+  - jazz-tools@0.15.11
+  - jazz-auth-betterauth@0.15.11
+  - jazz-betterauth-client-plugin@0.15.11
+  - cojson@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/jazz-react-auth-betterauth/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "jazz-react-auth-betterauth",
-	"version": "0.15.10",
+	"version": "0.15.14",
 	"type": "module",
 	"main": "dist/index.js",
 	"types": "src/index.tsx",

packages/jazz-run/CHANGELOG.md

@@ -1,5 +1,48 @@
 # jazz-run
 
+## 0.15.14
+
+### Patch Changes
+
+- Updated dependencies [70ce7c5]
+- Updated dependencies [a584590]
+- Updated dependencies [9acccb5]
+  - cojson-transport-ws@0.15.14
+  - cojson@0.15.14
+  - jazz-tools@0.15.14
+  - cojson-storage-sqlite@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- Updated dependencies [6c76ff8]
+  - jazz-tools@0.15.13
+  - cojson@0.15.13
+  - cojson-storage-sqlite@0.15.13
+  - cojson-transport-ws@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- Updated dependencies [d1c1b0c]
+- Updated dependencies [cf4ad72]
+  - jazz-tools@0.15.12
+  - cojson@0.15.12
+  - cojson-storage-sqlite@0.15.12
+  - cojson-transport-ws@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- Updated dependencies [bdc9aee]
+  - jazz-tools@0.15.11
+  - cojson@0.15.11
+  - cojson-storage-sqlite@0.15.11
+  - cojson-transport-ws@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/jazz-run/package.json

@@ -3,7 +3,7 @@
   "bin": "./dist/index.js",
   "type": "module",
   "license": "MIT",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "exports": {
     "./startSyncServer": {
       "import": "./dist/startSyncServer.js",
@@ -28,11 +28,11 @@
     "@effect/printer-ansi": "^0.34.5",
     "@effect/schema": "^0.71.1",
     "@effect/typeclass": "^0.25.5",
-    "cojson": "workspace:0.15.10",
-    "cojson-storage-sqlite": "workspace:0.15.10",
-    "cojson-transport-ws": "workspace:0.15.10",
+    "cojson": "workspace:0.15.14",
+    "cojson-storage-sqlite": "workspace:0.15.14",
+    "cojson-transport-ws": "workspace:0.15.14",
     "effect": "^3.6.5",
-    "jazz-tools": "workspace:0.15.10",
+    "jazz-tools": "workspace:0.15.14",
     "ws": "^8.14.2"
   },
   "devDependencies": {

packages/jazz-run/src/tests/startWorker.test.ts

@@ -230,8 +230,7 @@ describe("startWorker integration", () => {
     await worker2.done();
   });
 
-  // Flaky test, fails randomly on CI
-  test.skip("worker reconnects when sync server is closed and reopened", async () => {
+  test("worker reconnects when sync server is closed and reopened", async () => {
     const worker1 = await setup();
     const worker2 = await setupWorker(worker1.syncServer);
 
@@ -267,10 +266,6 @@ describe("startWorker integration", () => {
       db: "",
     });
 
-    // Wait for reconnection
-    await worker1.waitForConnection();
-    await worker2.waitForConnection();
-
     await worker1.worker.waitForAllCoValuesSync();
 
     // Verify both old and new values are synced

packages/jazz-tools/CHANGELOG.md

@@ -1,5 +1,45 @@
 # jazz-tools
 
+## 0.15.14
+
+### Patch Changes
+
+- a584590: Prevent resolving discriminated union fields
+- 9acccb5: Export `WithHelpers` type used in CoValue schemas
+- Updated dependencies [70ce7c5]
+  - cojson-transport-ws@0.15.14
+  - cojson@0.15.14
+  - cojson-storage-indexeddb@0.15.14
+
+## 0.15.13
+
+### Patch Changes
+
+- 6c76ff8: Fix load failures when loading a missing ref declared with z.optional and Schema.optional
+  - cojson@0.15.13
+  - cojson-storage-indexeddb@0.15.13
+  - cojson-transport-ws@0.15.13
+
+## 0.15.12
+
+### Patch Changes
+
+- d1c1b0c: Fix stuck authentication when using onAnonymousAccountDiscarded with a storage
+- cf4ad72: fix unhandled rejection on CoValue.load
+  - cojson@0.15.12
+  - cojson-storage-indexeddb@0.15.12
+  - cojson-transport-ws@0.15.12
+
+## 0.15.11
+
+### Patch Changes
+
+- bdc9aee: - Add `co.optional` and `co.discriminatedUnion`. You can now `load` and `subcribe` to schemas created with `co.discriminatedUnion`.
+  - Improved type-checking around `z.` schemas to prevent invalid combinations with `co.` schemas.
+  - cojson@0.15.11
+  - cojson-storage-indexeddb@0.15.11
+  - cojson-transport-ws@0.15.11
+
 ## 0.15.10
 
 ### Patch Changes

packages/jazz-tools/package.json

@@ -139,7 +139,7 @@
   },
   "type": "module",
   "license": "MIT",
-  "version": "0.15.10",
+  "version": "0.15.14",
   "dependencies": {
     "@manuscripts/prosemirror-recreate-steps": "^0.1.4",
     "@scure/base": "1.2.1",
@@ -163,6 +163,7 @@
     "zod": "3.25.28"
   },
   "scripts": {
+    "check": "tsc --noEmit",
     "format-and-lint": "biome check .",
     "format-and-lint:fix": "biome check . --write",
     "dev": "tsup --watch --dts",

packages/jazz-tools/src/react-core/tests/useAccount.test.ts

@@ -1,16 +1,6 @@
 // @vitest-environment happy-dom
 
-import {
-  Account,
-  CoMap,
-  Loaded,
-  RefsToResolve,
-  Resolved,
-  co,
-  coField,
-  z,
-  zodSchemaToCoSchema,
-} from "jazz-tools";
+import { RefsToResolve, co, z, zodSchemaToCoSchema } from "jazz-tools";
 import { beforeEach, describe, expect, it } from "vitest";
 import { useAccount, useJazzContextManager } from "../hooks.js";
 import { useIsAuthenticated } from "../index.js";

packages/jazz-tools/src/react-core/tests/useCoState.test.ts

@@ -47,6 +47,26 @@ describe("useCoState", () => {
     expect(result.current?.value).toBe("123");
   });
 
+  it("should return null on invalid id", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+    });
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const { result } = renderHook(() => useCoState(TestMap, "test", {}), {
+      account,
+    });
+
+    expect(result.current).toBeUndefined();
+
+    await waitFor(() => {
+      expect(result.current).toBeNull();
+    });
+  });
+
   it("should update the value when the coValue changes", async () => {
     const TestMap = co.map({
       value: z.string(),

packages/jazz-tools/src/tools/coValues/CoValueBase.ts

@@ -11,10 +11,10 @@ import {
   anySchemaToCoSchema,
   coValuesCache,
   inspect,
+  isCoValueSchema,
 } from "../internal.js";
 import type {
   Account,
-  CoValueClassFromZodSchema,
   Group,
   InstanceOfSchemaCoValuesNullable,
 } from "../internal.js";
@@ -101,7 +101,7 @@ export class CoValueBase implements CoValue {
     : S extends z.core.$ZodType
       ? NonNullable<InstanceOfSchemaCoValuesNullable<S>>
       : never {
-    const cl = "getCoSchema" in schema ? (schema as any).getCoSchema() : schema;
+    const cl = isCoValueSchema(schema) ? schema.getCoValueClass() : schema;
 
     if (this.constructor === cl) {
       return this as any;

packages/jazz-tools/src/tools/coValues/deepLoading.ts

@@ -2,11 +2,23 @@ import { SessionID } from "cojson";
 import { ItemsSym } from "../internal.js";
 import { type Account } from "./account.js";
 import { CoFeedEntry } from "./coFeed.js";
-import { type CoKeys, type CoMap } from "./coMap.js";
+import { type CoKeys } from "./coMap.js";
 import { type CoValue, type ID } from "./interfaces.js";
 
 type NotNull<T> = Exclude<T, null>;
 
+/**
+ * Used to check if T is a union type.
+ *
+ * If T is a union type, the left hand side of the extends becomes a union of function types.
+ * The right hand side is always a single function type.
+ */
+type IsUnion<T, U = T> = (T extends any ? (x: T) => void : never) extends (
+  x: U,
+) => void
+  ? false
+  : true;
+
 export type RefsToResolve<
   V,
   DepthLimit extends number = 10,
@@ -16,56 +28,58 @@ export type RefsToResolve<
   | (DepthLimit extends CurrentDepth["length"]
       ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
         any
-      : // Basically V extends CoList - but if we used that we'd introduce circularity into the definition of CoList itself
-        V extends Array<infer Item>
-        ?
-            | {
-                $each: RefsToResolve<
-                  NotNull<Item>,
-                  DepthLimit,
-                  [0, ...CurrentDepth]
-                >;
-                $onError?: null;
-              }
-            | boolean
-        : // Basically V extends CoMap | Group | Account - but if we used that we'd introduce circularity into the definition of CoMap itself
-          V extends { _type: "CoMap" | "Group" | "Account" }
+      : IsUnion<NonNullable<V>> extends true
+        ? true
+        : // Basically V extends CoList - but if we used that we'd introduce circularity into the definition of CoList itself
+          V extends Array<infer Item>
           ?
-              | ({
-                  [Key in CoKeys<V> as NonNullable<V[Key]> extends CoValue
-                    ? Key
-                    : never]?: RefsToResolve<
-                    NonNullable<V[Key]>,
+              | {
+                  $each: RefsToResolve<
+                    NotNull<Item>,
                     DepthLimit,
                     [0, ...CurrentDepth]
                   >;
-                } & { $onError?: null })
-              | (ItemsSym extends keyof V
-                  ? {
+                  $onError?: null;
+                }
+              | boolean
+          : // Basically V extends CoMap | Group | Account - but if we used that we'd introduce circularity into the definition of CoMap itself
+            V extends { _type: "CoMap" | "Group" | "Account" }
+            ?
+                | ({
+                    [Key in CoKeys<V> as NonNullable<V[Key]> extends CoValue
+                      ? Key
+                      : never]?: RefsToResolve<
+                      NonNullable<V[Key]>,
+                      DepthLimit,
+                      [0, ...CurrentDepth]
+                    >;
+                  } & { $onError?: null })
+                | (ItemsSym extends keyof V
+                    ? {
+                        $each: RefsToResolve<
+                          NonNullable<V[ItemsSym]>,
+                          DepthLimit,
+                          [0, ...CurrentDepth]
+                        >;
+                        $onError?: null;
+                      }
+                    : never)
+                | boolean
+            : V extends {
+                  _type: "CoStream";
+                  byMe: CoFeedEntry<infer Item> | undefined;
+                }
+              ?
+                  | {
                       $each: RefsToResolve<
-                        NonNullable<V[ItemsSym]>,
+                        NotNull<Item>,
                         DepthLimit,
                         [0, ...CurrentDepth]
                       >;
                       $onError?: null;
                     }
-                  : never)
-              | boolean
-          : V extends {
-                _type: "CoStream";
-                byMe: CoFeedEntry<infer Item> | undefined;
-              }
-            ?
-                | {
-                    $each: RefsToResolve<
-                      NotNull<Item>,
-                      DepthLimit,
-                      [0, ...CurrentDepth]
-                    >;
-                    $onError?: null;
-                  }
-                | boolean
-            : boolean);
+                  | boolean
+              : boolean);
 
 export type RefsToResolveStrict<T, V> = V extends RefsToResolve<T>
   ? RefsToResolve<T>

packages/jazz-tools/src/tools/coValues/inbox.ts

@@ -3,14 +3,12 @@ import { CoStreamItem, RawCoStream } from "cojson";
 import {
   type Account,
   CoValue,
-  CoValueClass,
   CoValueOrZodSchema,
   ID,
   InstanceOfSchema,
   activeAccountContext,
   anySchemaToCoSchema,
   loadCoValue,
-  zodSchemaToCoSchema,
 } from "../internal.js";
 
 export type InboxInvite = `${CoID<MessagesStream>}/${InviteSecret}`;

packages/jazz-tools/src/tools/coValues/schemaUnion.ts

@@ -1,8 +1,19 @@
 import {
+  Account,
+  AnonymousJazzAgent,
   CoValue,
   CoValueBase,
   CoValueClass,
   CoValueFromRaw,
+  ID,
+  RefsToResolve,
+  RefsToResolveStrict,
+  Resolved,
+  SubscribeListenerOptions,
+  SubscribeRestArgs,
+  loadCoValueWithoutMe,
+  parseSubscribeRestArgs,
+  subscribeToCoValueWithoutMe,
 } from "../internal.js";
 
 /**
@@ -99,4 +110,55 @@ export abstract class SchemaUnion extends CoValueBase implements CoValue {
   static fromRaw<V extends CoValue>(this: CoValueClass<V>, raw: V["_raw"]): V {
     throw new Error("Not implemented");
   }
+
+  /**
+   * Load a `SchemaUnion` with a given ID, as a given account.
+   *
+   * Note: The `resolve` option is not supported for `SchemaUnion`s due to https://github.com/garden-co/jazz/issues/2639
+   *
+   * @category Subscription & Loading
+   */
+  static load<M extends SchemaUnion>(
+    this: CoValueClass<M>,
+    id: ID<M>,
+    options?: {
+      loadAs?: Account | AnonymousJazzAgent;
+      skipRetry?: boolean;
+    },
+  ): Promise<Resolved<M, true> | null> {
+    return loadCoValueWithoutMe(this, id, options);
+  }
+
+  /**
+   * Load and subscribe to a `CoMap` with a given ID, as a given account.
+   *
+   * Automatically also subscribes to updates to all referenced/nested CoValues as soon as they are accessed in the listener.
+   *
+   * Returns an unsubscribe function that you should call when you no longer need updates.
+   *
+   * Also see the `useCoState` hook to reactively subscribe to a CoValue in a React component.
+   *
+   * Note: The `resolve` option is not supported for `SchemaUnion`s due to https://github.com/garden-co/jazz/issues/2639
+   *
+   * @category Subscription & Loading
+   */
+  static subscribe<M extends SchemaUnion>(
+    this: CoValueClass<M>,
+    id: ID<M>,
+    listener: (value: Resolved<M, true>, unsubscribe: () => void) => void,
+  ): () => void;
+  static subscribe<M extends SchemaUnion>(
+    this: CoValueClass<M>,
+    id: ID<M>,
+    options: SubscribeListenerOptions<M, true>,
+    listener: (value: Resolved<M, true>, unsubscribe: () => void) => void,
+  ): () => void;
+  static subscribe<M extends SchemaUnion>(
+    this: CoValueClass<M>,
+    id: ID<M>,
+    ...args: SubscribeRestArgs<M, true>
+  ): () => void {
+    const { options, listener } = parseSubscribeRestArgs(args);
+    return subscribeToCoValueWithoutMe<M, true>(this, id, options, listener);
+  }
 }

packages/jazz-tools/src/tools/exports.ts

@@ -35,6 +35,7 @@ export type {
   TextPos,
   AccountClass,
   AccountCreationProps,
+  WithHelpers,
 } from "./internal.js";
 
 export {

packages/jazz-tools/src/tools/implementation/ContextManager.ts

@@ -279,8 +279,12 @@ export class JazzContextManager<
           },
         );
 
-      prevContext.node.syncManager.addPeer(currentAccountAsPeer);
+      // Closing storage on the prevContext to avoid conflicting transactions and getting stuck on waitForAllCoValuesSync
+      // The storage is reachable through currentContext using the connectedPeers
+      prevContext.node.removeStorage();
+
       currentContext.node.syncManager.addPeer(prevAccountAsPeer);
+      prevContext.node.syncManager.addPeer(currentAccountAsPeer);
 
       try {
         await this.props.onAnonymousAccountDiscarded?.(prevContext.me);

packages/jazz-tools/src/tools/implementation/zodSchema/coExport.ts

@@ -10,4 +10,6 @@ export {
   coImageDefiner as image,
   coAccountDefiner as account,
   coProfileDefiner as profile,
+  coOptionalDefiner as optional,
+  coDiscriminatedUnionDefiner as discriminatedUnion,
 } from "./zodCo.js";

packages/jazz-tools/src/tools/implementation/zodSchema/runtimeConverters/zodFieldToCoFieldDef.ts

@@ -1,5 +1,6 @@
 import { CoMap, CoValueClass, isCoValueClass } from "../../../internal.js";
 import { coField } from "../../schema.js";
+import { isAnyCoOptionalSchema } from "../schemaTypes/CoOptionalSchema.js";
 import {
   isUnionOfCoMapsDeeply,
   isUnionOfPrimitivesDeeply,
@@ -12,10 +13,19 @@ import {
   ZodReadonly,
   z,
 } from "../zodReExport.js";
-import { ZodPrimitiveSchema } from "../zodSchema.js";
-import { zodSchemaToCoSchemaOrKeepPrimitive } from "./zodSchemaToCoSchema.js";
+import { AnyCoSchema, ZodPrimitiveSchema } from "../zodSchema.js";
+import {
+  isCoValueSchema,
+  zodSchemaToCoSchemaOrKeepPrimitive,
+} from "./zodSchemaToCoSchema.js";
 
-type FieldSchema =
+/**
+ * Types of objects that can be nested inside CoValue schema containers
+ */
+type SchemaField =
+  // Schemas created with co.map(), co.record(), co.list(), etc.
+  | AnyCoSchema
+  // CoValue classes created with class syntax, or framework-provided classes like Group
   | CoValueClass
   | ZodPrimitiveSchema
   | z.core.$ZodOptional<z.core.$ZodType>
@@ -33,20 +43,35 @@ type FieldSchema =
   | z.core.$ZodCatch<z.core.$ZodType>
   | (z.core.$ZodCustom<any, any> & { builtin: any });
 
-export function zodFieldToCoFieldDef(schema: FieldSchema) {
+export function schemaFieldToCoFieldDef(
+  schema: SchemaField,
+  isOptional = false,
+) {
   if (isCoValueClass(schema)) {
-    return coField.ref(schema);
+    if (isOptional) {
+      return coField.ref(schema, { optional: true });
+    } else {
+      return coField.ref(schema);
+    }
+  } else if (isCoValueSchema(schema)) {
+    if (isAnyCoOptionalSchema(schema)) {
+      return coField.ref(schema.getCoValueClass(), {
+        optional: true,
+      });
+    }
+
+    if (isOptional) {
+      return coField.ref(schema.getCoValueClass(), { optional: true });
+    } else {
+      return coField.ref(schema.getCoValueClass());
+    }
   } else {
     if ("_zod" in schema) {
       if (schema._zod.def.type === "optional") {
         const inner = zodSchemaToCoSchemaOrKeepPrimitive(
           schema._zod.def.innerType,
         );
-        if (isCoValueClass(inner)) {
-          return coField.ref(inner, { optional: true });
-        } else {
-          return zodFieldToCoFieldDef(inner);
-        }
+        return schemaFieldToCoFieldDef(inner, true);
       } else if (schema._zod.def.type === "string") {
         return coField.string;
       } else if (schema._zod.def.type === "number") {
@@ -58,8 +83,9 @@ export function zodFieldToCoFieldDef(schema: FieldSchema) {
       } else if (schema._zod.def.type === "enum") {
         return coField.string;
       } else if (schema._zod.def.type === "readonly") {
-        return zodFieldToCoFieldDef(
-          (schema as unknown as ZodReadonly).def.innerType as FieldSchema,
+        return schemaFieldToCoFieldDef(
+          (schema as unknown as ZodReadonly).def.innerType as SchemaField,
+          isOptional,
         );
       } else if (schema._zod.def.type === "date") {
         return coField.optional.Date;
@@ -67,8 +93,9 @@ export function zodFieldToCoFieldDef(schema: FieldSchema) {
         return coField.string;
       } else if (schema._zod.def.type === "lazy") {
         // Mostly to support z.json()
-        return zodFieldToCoFieldDef(
-          (schema as unknown as ZodLazy).unwrap() as FieldSchema,
+        return schemaFieldToCoFieldDef(
+          (schema as unknown as ZodLazy).unwrap() as SchemaField,
+          isOptional,
         );
       } else if (
         schema._zod.def.type === "default" ||
@@ -78,9 +105,10 @@ export function zodFieldToCoFieldDef(schema: FieldSchema) {
           "z.default()/z.catch() are not supported in collaborative schemas. They will be ignored.",
         );
 
-        return zodFieldToCoFieldDef(
+        return schemaFieldToCoFieldDef(
           (schema as unknown as ZodDefault | ZodCatch).def
-            .innerType as FieldSchema,
+            .innerType as SchemaField,
+          isOptional,
         );
       } else if (schema._zod.def.type === "literal") {
         if (
@@ -112,7 +140,7 @@ export function zodFieldToCoFieldDef(schema: FieldSchema) {
         return coField.json();
       } else if (schema._zod.def.type === "custom") {
         if ("builtin" in schema) {
-          return zodFieldToCoFieldDef(schema.builtin);
+          return schemaFieldToCoFieldDef(schema.builtin, isOptional);
         } else {
           throw new Error(`Unsupported custom zod type`);
         }
@@ -120,9 +148,13 @@ export function zodFieldToCoFieldDef(schema: FieldSchema) {
         if (isUnionOfPrimitivesDeeply(schema)) {
           return coField.json();
         } else if (isUnionOfCoMapsDeeply(schema)) {
-          return coField.ref<CoValueClass<CoMap>>(
-            schemaUnionDiscriminatorFor(schema),
-          );
+          const result = schemaUnionDiscriminatorFor(schema);
+
+          if (isOptional) {
+            return coField.ref<CoValueClass<CoMap>>(result, { optional: true });
+          } else {
+            return coField.ref<CoValueClass<CoMap>>(result);
+          }
         } else {
           throw new Error(
             "z.union()/z.discriminatedUnion() of mixed collaborative and non-collaborative types is not supported",

packages/jazz-tools/src/tools/implementation/zodSchema/runtimeConverters/zodSchemaToCoSchema.ts

@@ -9,89 +9,140 @@ import {
   CoValueClass,
   FileStream,
   SchemaUnion,
+  enrichAccountSchema,
+  enrichCoDiscriminatedUnionSchema,
+  enrichCoFeedSchema,
+  enrichCoListSchema,
+  enrichCoMapSchema,
+  enrichFileStreamSchema,
+  enrichPlainTextSchema,
   isCoValueClass,
 } from "../../../internal.js";
 import { coField } from "../../schema.js";
+import { isAnyCoOptionalSchema } from "../schemaTypes/CoOptionalSchema.js";
+import { enrichRichTextSchema } from "../schemaTypes/RichTextSchema.js";
 import {
   isUnionOfCoMapsDeeply,
   schemaUnionDiscriminatorFor,
 } from "../unionUtils.js";
 import { z } from "../zodReExport.js";
 import {
-  CoValueClassFromZodSchema,
+  AnyCoSchema,
+  CoValueClassFromAnySchema,
+  CoValueOrZodSchema,
+  CoValueSchemaFromZodSchema,
   ZodPrimitiveSchema,
   getDef,
   isZodArray,
   isZodCustom,
   isZodObject,
 } from "../zodSchema.js";
-import { zodFieldToCoFieldDef } from "./zodFieldToCoFieldDef.js";
+import { schemaFieldToCoFieldDef } from "./zodFieldToCoFieldDef.js";
 
-let coSchemasForZodSchemas = new Map<z.core.$ZodType, CoValueClass>();
+let coSchemasForZodSchemas = new Map<z.core.$ZodType, AnyCoSchema>();
 
-export function tryZodSchemaToCoSchema<S extends z.core.$ZodType>(
+export function isAnyCoValueSchema(
+  schema: z.core.$ZodType | CoValueClass,
+): schema is AnyCoSchema {
+  return "collaborative" in schema && schema.collaborative === true;
+}
+
+export function isCoValueSchema(
+  schema: z.core.$ZodType | CoValueClass,
+): schema is CoValueSchemaFromZodSchema<AnyCoSchema> {
+  return isAnyCoValueSchema(schema) && "getCoValueClass" in schema;
+}
+
+/**
+ * Convert a Zod schema into a CoValue schema.
+ *
+ * @param schema A Zod schema that may represent a CoValue schema
+ * @returns The CoValue schema matching the provided ProtoCoSchema, or `null` if the Zod schema
+ * does not match a CoValue schema.
+ */
+function tryZodSchemaToCoSchema<S extends z.core.$ZodType>(
   schema: S,
-): CoValueClassFromZodSchema<S> | null {
-  if ("collaborative" in schema && schema.collaborative) {
+): CoValueSchemaFromZodSchema<S> | null {
+  if (isAnyCoValueSchema(schema)) {
     if (coSchemasForZodSchemas.has(schema)) {
-      return coSchemasForZodSchemas.get(schema) as CoValueClassFromZodSchema<S>;
+      return coSchemasForZodSchemas.get(
+        schema,
+      ) as CoValueSchemaFromZodSchema<S>;
     }
 
-    if (isZodObject(schema)) {
+    if (isAnyCoOptionalSchema(schema)) {
+      // Optional schemas are not supported as top-level schemas
+      return null;
+    } else if (isZodObject(schema)) {
       const def = getDef(schema);
 
       const ClassToExtend =
         "builtin" in schema && schema.builtin === "Account" ? Account : CoMap;
 
-      const coSchema = class ZCoMap extends ClassToExtend {
+      const coValueClass = class ZCoMap extends ClassToExtend {
         constructor(options: { fromRaw: RawCoMap } | undefined) {
           super(options);
           for (const [field, fieldType] of Object.entries(
             def.shape as z.core.$ZodShape,
           )) {
-            (this as any)[field] = zodFieldToCoFieldDef(
+            (this as any)[field] = schemaFieldToCoFieldDef(
               zodSchemaToCoSchemaOrKeepPrimitive(fieldType),
             );
           }
           if (def.catchall) {
-            (this as any)[coField.items] = zodFieldToCoFieldDef(
+            (this as any)[coField.items] = schemaFieldToCoFieldDef(
               zodSchemaToCoSchemaOrKeepPrimitive(def.catchall),
             );
           }
         }
       };
 
-      coSchemasForZodSchemas.set(schema, coSchema as unknown as CoValueClass);
-      return coSchema as unknown as CoValueClassFromZodSchema<S>;
+      const coValueSchema =
+        ClassToExtend === Account
+          ? enrichAccountSchema(schema as any, coValueClass as any)
+          : enrichCoMapSchema(schema as any, coValueClass as any);
+
+      coSchemasForZodSchemas.set(schema, coValueSchema);
+      return coValueSchema as unknown as CoValueSchemaFromZodSchema<S>;
     } else if (isZodArray(schema)) {
       const def = getDef(schema);
-      const coSchema = class ZCoList extends CoList {
+      const coValueClass = class ZCoList extends CoList {
         constructor(options: { fromRaw: RawCoList } | undefined) {
           super(options);
-          (this as any)[coField.items] = zodFieldToCoFieldDef(
+          (this as any)[coField.items] = schemaFieldToCoFieldDef(
             zodSchemaToCoSchemaOrKeepPrimitive(def.element),
           );
         }
       };
 
-      coSchemasForZodSchemas.set(schema, coSchema);
-      return coSchema as unknown as CoValueClassFromZodSchema<S>;
+      const coValueSchema = enrichCoListSchema(schema, coValueClass as any);
+
+      coSchemasForZodSchemas.set(schema, coValueSchema);
+      return coValueSchema as unknown as CoValueSchemaFromZodSchema<S>;
     } else if (isZodCustom(schema)) {
       if ("builtin" in schema) {
         if (schema.builtin === "CoFeed" && "element" in schema) {
-          return CoFeed.Of(
-            zodFieldToCoFieldDef(
+          const coValueClass = CoFeed.Of(
+            schemaFieldToCoFieldDef(
               zodSchemaToCoSchemaOrKeepPrimitive(
                 schema.element as z.core.$ZodType,
               ),
             ),
-          ) as unknown as CoValueClassFromZodSchema<S>;
+          );
+          const coValueSchema = enrichCoFeedSchema(schema, coValueClass as any);
+          return coValueSchema as unknown as CoValueSchemaFromZodSchema<S>;
         } else if (schema.builtin === "FileStream") {
-          return FileStream as unknown as CoValueClassFromZodSchema<S>;
+          const coValueClass = FileStream;
+          const coValueSchema = enrichFileStreamSchema(schema, coValueClass);
+          return coValueSchema as unknown as CoValueSchemaFromZodSchema<S>;
         } else if (schema.builtin === "CoPlainText") {
-          return CoPlainText as unknown as CoValueClassFromZodSchema<S>;
+          const coValueClass = CoPlainText;
+          const coValueSchema = enrichPlainTextSchema(schema, coValueClass);
+          return coValueSchema as unknown as CoValueSchemaFromZodSchema<S>;
         } else if (schema.builtin === "CoRichText") {
-          return CoRichText as unknown as CoValueClassFromZodSchema<S>;
+          const coValueClass = CoRichText;
+          const coValueSchema = enrichRichTextSchema(schema, coValueClass);
+          return coValueSchema as unknown as CoValueSchemaFromZodSchema<S>;
         } else {
           throw new Error(`Unsupported builtin type: ${schema.builtin}`);
         }
@@ -105,9 +156,12 @@ export function tryZodSchemaToCoSchema<S extends z.core.$ZodType>(
     }
   } else if (schema instanceof z.core.$ZodDiscriminatedUnion) {
     if (isUnionOfCoMapsDeeply(schema)) {
-      return SchemaUnion.Of(
-        schemaUnionDiscriminatorFor(schema),
-      ) as unknown as CoValueClassFromZodSchema<S>;
+      const coValueClass = SchemaUnion.Of(schemaUnionDiscriminatorFor(schema));
+      const coValueSchema = enrichCoDiscriminatedUnionSchema(
+        schema as any,
+        coValueClass as any,
+      );
+      return coValueSchema as unknown as CoValueSchemaFromZodSchema<S>;
     } else {
       throw new Error(
         "z.discriminatedUnion() of non-collaborative types is not supported as a top-level schema",
@@ -118,19 +172,9 @@ export function tryZodSchemaToCoSchema<S extends z.core.$ZodType>(
   }
 }
 
-export function zodSchemaToCoSchema<
-  S extends
-    | z.core.$ZodType
-    | (z.core.$ZodObject<any, any> & {
-        builtin: "Account";
-        migration?: (account: any, creationProps?: { name: string }) => void;
-      })
-    | (z.core.$ZodCustom<any, any> & { builtin: "FileStream" })
-    | (z.core.$ZodCustom<any, any> & {
-        builtin: "CoFeed";
-        element: z.core.$ZodType;
-      }),
->(schema: S): CoValueClassFromZodSchema<S> {
+export function zodSchemaToCoSchema<S extends z.core.$ZodType | AnyCoSchema>(
+  schema: S,
+): CoValueSchemaFromZodSchema<S> {
   const coSchema = tryZodSchemaToCoSchema(schema);
   if (!coSchema) {
     throw new Error(
@@ -140,38 +184,24 @@ export function zodSchemaToCoSchema<
   return coSchema;
 }
 
-export function anySchemaToCoSchema<
-  S extends
-    | CoValueClass
-    | z.core.$ZodType
-    | (z.core.$ZodObject<any, any> & {
-        builtin: "Account";
-        migration?: (account: any, creationProps?: { name: string }) => void;
-      })
-    | (z.core.$ZodCustom<any, any> & { builtin: "FileStream" })
-    | (z.core.$ZodCustom<any, any> & {
-        builtin: "CoFeed";
-        element: z.core.$ZodType;
-      }),
->(
+// TODO this should be coValueClassOrAnySchemaToCoValueClass
+export function anySchemaToCoSchema<S extends CoValueOrZodSchema>(
   schema: S,
-): S extends CoValueClass
-  ? S
-  : S extends z.core.$ZodType
-    ? CoValueClassFromZodSchema<S>
-    : never {
+): CoValueClassFromAnySchema<S> {
   if (isCoValueClass(schema)) {
     return schema as any;
-  } else if ("getCoSchema" in schema) {
-    return (schema as any).getCoSchema() as any;
+  } else if (isCoValueSchema(schema)) {
+    return schema.getCoValueClass() as any;
   } else if ("def" in schema) {
-    const coSchema = tryZodSchemaToCoSchema(schema as z.core.$ZodType);
+    const coSchema = tryZodSchemaToCoSchema(
+      schema as z.core.$ZodType | AnyCoSchema,
+    );
     if (!coSchema) {
       throw new Error(
         `Unsupported zod type: ${(schema.def as any)?.type || JSON.stringify(schema)}`,
       );
     }
-    return coSchema as any;
+    return coSchema.getCoValueClass() as any;
   }
 
   throw new Error(`Unsupported schema: ${JSON.stringify(schema)}`);
@@ -179,7 +209,7 @@ export function anySchemaToCoSchema<
 
 export function zodSchemaToCoSchemaOrKeepPrimitive<S extends z.core.$ZodType>(
   schema: S,
-): CoValueClassFromZodSchema<S> | ZodPrimitiveSchema {
+): CoValueSchemaFromZodSchema<S> | ZodPrimitiveSchema {
   const coSchema = tryZodSchemaToCoSchema(schema);
   if (!coSchema) {
     return schema as any;

packages/jazz-tools/src/tools/implementation/zodSchema/schemaTypes/AccountSchema.ts

@@ -1,5 +1,10 @@
 import { CryptoProvider } from "cojson";
-import { Account, Group, RefsToResolveStrict } from "../../../internal.js";
+import {
+  Account,
+  AccountCreationProps,
+  Group,
+  RefsToResolveStrict,
+} from "../../../internal.js";
 import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
 import { InstanceOrPrimitiveOfSchema } from "../typeConverters/InstanceOrPrimitiveOfSchema.js";
 import { InstanceOrPrimitiveOfSchemaCoValuesNullable } from "../typeConverters/InstanceOrPrimitiveOfSchemaCoValuesNullable.js";
@@ -57,9 +62,63 @@ export type AccountSchema<
     ) => void,
   ): AccountSchema<Shape>;
 
-  getCoSchema: () => typeof Account;
+  getCoValueClass: () => typeof Account;
 };
 
+export function enrichAccountSchema<Shape extends BaseAccountShape>(
+  schema: AnyAccountSchema<Shape>,
+  coValueClass: typeof Account,
+): AccountSchema<Shape> {
+  const enrichedSchema = Object.assign(schema, {
+    create: (...args: any[]) => {
+      // @ts-expect-error
+      return coValueClass.create(...args);
+    },
+    createAs: (...args: any[]) => {
+      // @ts-expect-error
+      return coValueClass.createAs(...args);
+    },
+    getMe: (...args: any[]) => {
+      // @ts-expect-error
+      return coValueClass.getMe(...args);
+    },
+    load: (...args: any[]) => {
+      // @ts-expect-error
+      return coValueClass.load(...args);
+    },
+    subscribe: (...args: any[]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
+      return Object.assign(schema, helpers(schema));
+    },
+    fromRaw: (...args: any[]) => {
+      // @ts-expect-error
+      return coValueClass.fromRaw(...args);
+    },
+    withMigration: (
+      migration: (
+        value: any,
+        creationProps?: AccountCreationProps,
+      ) => void | Promise<void>,
+    ) => {
+      (coValueClass.prototype as Account).migrate = async function (
+        this,
+        creationProps,
+      ) {
+        await migration(this, creationProps);
+      };
+
+      return enrichedSchema;
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as AccountSchema<Shape>;
+  return enrichedSchema;
+}
+
 export type DefaultProfileShape = {
   name: z.core.$ZodString<string>;
   inbox: z.core.$ZodOptional<z.core.$ZodString>;
@@ -71,7 +130,7 @@ export type CoProfileSchema<
   Config extends z.core.$ZodObjectConfig = z.core.$ZodObjectConfig,
 > = CoMapSchema<Shape & DefaultProfileShape, Config, Group>;
 
-// less precise verion to avoid circularity issues and allow matching against
+// less precise version to avoid circularity issues and allow matching against
 export type AnyAccountSchema<
   Shape extends z.core.$ZodLooseShape = z.core.$ZodLooseShape,
 > = z.core.$ZodObject<Shape> & {

packages/jazz-tools/src/tools/implementation/zodSchema/schemaTypes/CoDiscriminatedUnionSchema.ts

@@ -0,0 +1,90 @@
+import {
+  Account,
+  AnonymousJazzAgent,
+  AnyCoSchema,
+  InstanceOrPrimitiveOfSchemaCoValuesNullable,
+  RefsToResolve,
+  RefsToResolveStrict,
+  Resolved,
+  SchemaUnion,
+  SubscribeListenerOptions,
+} from "../../../internal.js";
+import { z } from "../zodReExport.js";
+
+export type AnyDiscriminableCoSchema = AnyCoSchema &
+  z.core.$ZodTypeDiscriminable;
+
+export type AnyCoDiscriminatedUnionSchema<
+  Types extends readonly [
+    AnyDiscriminableCoSchema,
+    ...AnyDiscriminableCoSchema[],
+  ],
+> = z.ZodDiscriminatedUnion<Types> & {
+  collaborative: true;
+};
+
+export type CoDiscriminatedUnionSchema<
+  Types extends readonly [
+    AnyDiscriminableCoSchema,
+    ...AnyDiscriminableCoSchema[],
+  ],
+> = AnyCoDiscriminatedUnionSchema<Types> & {
+  load(
+    id: string,
+    options?: {
+      loadAs?: Account | AnonymousJazzAgent;
+      skipRetry?: boolean;
+    },
+  ): Promise<Resolved<
+    CoDiscriminatedUnionInstanceCoValuesNullable<Types> & SchemaUnion,
+    true
+  > | null>;
+
+  subscribe(
+    id: string,
+    options: SubscribeListenerOptions<
+      CoDiscriminatedUnionInstanceCoValuesNullable<Types> & SchemaUnion,
+      true
+    >,
+    listener: (
+      value: Resolved<
+        CoDiscriminatedUnionInstanceCoValuesNullable<Types> & SchemaUnion,
+        true
+      >,
+      unsubscribe: () => void,
+    ) => void,
+  ): () => void;
+
+  getCoValueClass: () => typeof SchemaUnion;
+};
+
+export function enrichCoDiscriminatedUnionSchema<
+  Types extends readonly [
+    AnyDiscriminableCoSchema,
+    ...AnyDiscriminableCoSchema[],
+  ],
+>(
+  schema: z.ZodDiscriminatedUnion<Types>,
+  coValueClass: typeof SchemaUnion,
+): CoDiscriminatedUnionSchema<Types> {
+  return Object.assign(schema, {
+    load: (...args: [any, ...any]) => {
+      // @ts-expect-error
+      return coValueClass.load(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as CoDiscriminatedUnionSchema<Types>;
+}
+
+type CoDiscriminatedUnionInstanceCoValuesNullable<
+  Types extends readonly [
+    AnyDiscriminableCoSchema,
+    ...AnyDiscriminableCoSchema[],
+  ],
+> = NonNullable<InstanceOrPrimitiveOfSchemaCoValuesNullable<Types[number]>>;

packages/jazz-tools/src/tools/implementation/zodSchema/schemaTypes/CoFeedSchema.ts

@@ -58,10 +58,35 @@ export type CoFeedSchema<T extends z.core.$ZodType> = z.core.$ZodCustom<
     ) => void,
   ): () => void;
 
-  getCoSchema: () => typeof CoFeed;
+  getCoValueClass: () => typeof CoFeed;
 };
 
-// less precise verion to avoid circularity issues and allow matching against
+export function enrichCoFeedSchema<T extends z.core.$ZodType>(
+  schema: AnyCoFeedSchema<T>,
+  coValueClass: typeof CoFeed,
+): CoFeedSchema<T> {
+  return Object.assign(schema, {
+    create: (...args: [any, ...any[]]) => {
+      return coValueClass.create(...args);
+    },
+    load: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.load(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
+      return Object.assign(schema, helpers(schema));
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as CoFeedSchema<T>;
+}
+
+// less precise version to avoid circularity issues and allow matching against
 export type AnyCoFeedSchema<T extends z.core.$ZodType = z.core.$ZodType> =
   z.core.$ZodCustom<any, unknown> & {
     collaborative: true;

packages/jazz-tools/src/tools/implementation/zodSchema/schemaTypes/CoListSchema.ts

@@ -19,9 +19,7 @@ type CoListInit<T extends z.core.$ZodType> = Array<
     : NonNullable<InstanceOrPrimitiveOfSchemaCoValuesNullable<T>>
 >;
 
-export type CoListSchema<T extends z.core.$ZodType> = z.core.$ZodArray<T> & {
-  collaborative: true;
-
+export type CoListSchema<T extends z.core.$ZodType> = AnyCoListSchema<T> & {
   create: (
     items: CoListInit<T>,
     options?: { owner: Account | Group } | Account | Group,
@@ -52,10 +50,34 @@ export type CoListSchema<T extends z.core.$ZodType> = z.core.$ZodArray<T> & {
     helpers: (Self: S) => T,
   ): WithHelpers<S, T>;
 
-  getCoSchema: () => typeof CoList;
+  getCoValueClass: () => typeof CoList;
 };
 
-// less precise verion to avoid circularity issues and allow matching against
+export function enrichCoListSchema<T extends z.core.$ZodType>(
+  schema: AnyCoListSchema<T>,
+  coValueClass: typeof CoList,
+): CoListSchema<T> {
+  return Object.assign(schema, {
+    create: (...args: [any, ...any[]]) => {
+      return coValueClass.create(...args);
+    },
+    load: (...args: [any, ...any[]]) => {
+      return coValueClass.load(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
+      return Object.assign(schema, helpers(schema));
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as CoListSchema<T>;
+}
+
+// less precise version to avoid circularity issues and allow matching against
 export type AnyCoListSchema<T extends z.core.$ZodType = z.core.$ZodType> =
   z.core.$ZodArray<T> & { collaborative: true };
 

packages/jazz-tools/src/tools/implementation/zodSchema/schemaTypes/CoMapSchema.ts

@@ -8,6 +8,7 @@ import {
   Resolved,
   Simplify,
   SubscribeListenerOptions,
+  zodSchemaToCoSchema,
 } from "../../../internal.js";
 import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
 import { InstanceOrPrimitiveOfSchema } from "../typeConverters/InstanceOrPrimitiveOfSchema.js";
@@ -19,10 +20,8 @@ export type CoMapSchema<
   Shape extends z.core.$ZodLooseShape,
   Config extends z.core.$ZodObjectConfig = z.core.$ZodObjectConfig,
   Owner extends Account | Group = Account | Group,
-> = z.core.$ZodObject<Shape, Config> &
+> = AnyCoMapSchema<Shape, Config> &
   z.$ZodTypeDiscriminable & {
-    collaborative: true;
-
     create: (
       init: Simplify<CoMapInitZod<Shape>>,
       options?:
@@ -145,9 +144,65 @@ export type CoMapSchema<
       ) => undefined,
     ): CoMapSchema<Shape, Config, Owner>;
 
-    getCoSchema: () => typeof CoMap;
+    getCoValueClass: () => typeof CoMap;
   };
 
+export function enrichCoMapSchema<
+  Shape extends z.core.$ZodLooseShape,
+  Config extends z.core.$ZodObjectConfig,
+>(
+  schema: AnyCoMapSchema<Shape, Config>,
+  coValueClass: typeof CoMap,
+): CoMapSchema<Shape, Config> {
+  // @ts-expect-error schema is actually a z.ZodObject, but we need to use z.core.$ZodObject to avoid circularity issues
+  const baseCatchall = schema.catchall;
+  const coValueSchema = Object.assign(schema, {
+    create: (...args: [any, ...any[]]) => {
+      return coValueClass.create(...args);
+    },
+    load: (...args: [any, ...any[]]) => {
+      return coValueClass.load(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    findUnique: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.findUnique(...args);
+    },
+    upsertUnique: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.upsertUnique(...args);
+    },
+    loadUnique: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.loadUnique(...args);
+    },
+    catchall: (index: z.core.$ZodType) => {
+      const newSchema = baseCatchall(index);
+      // TODO avoid repeating this with coMapDefiner
+      const enrichedSchema = Object.assign(newSchema, {
+        collaborative: true,
+      }) as AnyCoMapSchema<Shape, Config>;
+      return zodSchemaToCoSchema(enrichedSchema);
+    },
+    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
+      return Object.assign(schema, helpers(schema));
+    },
+    withMigration: (migration: (value: any) => undefined) => {
+      // @ts-expect-error TODO check
+      coValueClass.prototype.migrate = migration;
+
+      return coValueSchema;
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as CoMapSchema<Shape, Config>;
+  return coValueSchema;
+}
+
 export type optionalKeys<Shape extends z.core.$ZodLooseShape> = {
   [key in keyof Shape]: Shape[key] extends z.core.$ZodOptional<any>
     ? key
@@ -170,7 +225,7 @@ export type CoMapInitZod<Shape extends z.core.$ZodLooseShape> = {
   >;
 } & { [key in keyof Shape]?: unknown };
 
-// less precise verion to avoid circularity issues and allow matching against
+// less precise version to avoid circularity issues and allow matching against
 export type AnyCoMapSchema<
   Shape extends z.core.$ZodLooseShape = z.core.$ZodLooseShape,
   Config extends z.core.$ZodObjectConfig = z.core.$ZodObjectConfig,

packages/jazz-tools/src/tools/implementation/zodSchema/schemaTypes/CoOptionalSchema.ts

@@ -0,0 +1,33 @@
+import { isAnyCoValueSchema } from "../runtimeConverters/zodSchemaToCoSchema.js";
+import { z } from "../zodReExport.js";
+import { AnyCoSchema, CoValueSchemaFromZodSchema } from "../zodSchema.js";
+
+export type AnyCoOptionalSchema<
+  Shape extends z.core.$ZodType = z.core.$ZodType,
+> = z.ZodOptional<Shape> & {
+  collaborative: true;
+};
+
+export type CoOptionalSchema<Shape extends z.core.$ZodType = z.core.$ZodType> =
+  AnyCoOptionalSchema<Shape> & {
+    getCoValueClass: () => CoValueSchemaFromZodSchema<AnyCoSchema>["getCoValueClass"];
+  };
+
+export function createCoOptionalSchema<T extends AnyCoSchema>(
+  schema: T,
+): CoOptionalSchema<T> {
+  return Object.assign(z.optional(schema), {
+    collaborative: true,
+    getCoValueClass: () => {
+      return (
+        schema as CoValueSchemaFromZodSchema<AnyCoSchema>
+      ).getCoValueClass();
+    },
+  }) as unknown as CoOptionalSchema<T>;
+}
+
+export function isAnyCoOptionalSchema(
+  schema: z.core.$ZodType,
+): schema is CoOptionalSchema<z.core.$ZodType> {
+  return isAnyCoValueSchema(schema) && schema._zod.def.type === "optional";
+}

packages/jazz-tools/src/tools/implementation/zodSchema/schemaTypes/CoRecordSchema.ts

@@ -28,9 +28,7 @@ type CoRecordInit<
 export type CoRecordSchema<
   K extends z.core.$ZodString<string>,
   V extends z.core.$ZodType,
-> = z.core.$ZodRecord<K, V> & {
-  collaborative: true;
-
+> = AnyCoRecordSchema<K, V> & {
   create: (
     init: Simplify<CoRecordInit<K, V>>,
     options?:
@@ -83,10 +81,10 @@ export type CoRecordSchema<
     this: S,
     helpers: (Self: S) => T,
   ): WithHelpers<S, T>;
-  getCoSchema: () => typeof CoMap;
+  getCoValueClass: () => typeof CoMap;
 };
 
-// less precise verion to avoid circularity issues and allow matching against
+// less precise version to avoid circularity issues and allow matching against
 export type AnyCoRecordSchema<
   K extends z.core.$ZodString<string> = z.core.$ZodString<string>,
   V extends z.core.$ZodType = z.core.$ZodType,