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

Version Packages

examples/chat-svelte/CHANGELOG.md

@@ -1,5 +1,27 @@
 # passkey-svelte
 
+## 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.102",
   "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/jazz-paper-scissors/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,

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-side-validation/vercel.json

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

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,23 @@
 # cojson-storage-indexeddb
 
+## 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.13",
   "main": "dist/index.js",
   "type": "module",
   "types": "dist/index.d.ts",

packages/cojson-storage-sqlite/CHANGELOG.md

@@ -1,5 +1,23 @@
 # cojson-storage-sqlite
 
+## 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.13",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "license": "MIT",

packages/cojson-transport-ws/CHANGELOG.md

@@ -1,5 +1,23 @@
 # cojson-transport-nodejs-ws
 
+## 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.13",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "license": "MIT",

packages/cojson/CHANGELOG.md

@@ -1,5 +1,11 @@
 # cojson
 
+## 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.13",
   "devDependencies": {
     "@opentelemetry/sdk-metrics": "^2.0.0",
     "libsql": "^0.5.13",

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/tests/sync.load.test.ts

@@ -54,6 +54,38 @@ describe("loading coValues from server", () => {
     `);
   });
 
+  test("coValue load throws on invalid id", async () => {
+    const { node } = setupTestNode({
+      connected: true,
+    });
+
+    expect(async () => await node.load("test" as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id test",
+    );
+    expect(async () => await node.load(null as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id null",
+    );
+    expect(async () => await node.load(undefined as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id undefined",
+    );
+    expect(async () => await node.load(1 as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id 1",
+    );
+    expect(async () => await node.load({} as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id [object Object]",
+    );
+    expect(async () => await node.load([] as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id []",
+    );
+    expect(async () => await node.load(["test"] as any)).rejects.toThrow(
+      'Trying to load CoValue with invalid id ["test"]',
+    );
+    expect(async () => await node.load((() => {}) as any)).rejects.toThrow(
+      "Trying to load CoValue with invalid id () => {\n    }",
+    );
+    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();

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

@@ -991,7 +991,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/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>
 
@@ -11911,4 +11911,4 @@ export function cn(...inputs: ClassValue[]) {
 ```ts
 /// <reference types="vite/client" />
 
-```
+```

packages/jazz-auth-betterauth/CHANGELOG.md

@@ -1,5 +1,33 @@
 # jazz-auth-betterauth
 
+## 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.13",
 	"type": "module",
 	"main": "dist/index.js",
 	"types": "src/index.ts",

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

@@ -1,5 +1,23 @@
 # jazz-betterauth-client-plugin
 
+## 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.13",
   "type": "module",
   "main": "dist/index.js",
   "types": "src/index.ts",

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

@@ -1,5 +1,30 @@
 # jazz-betterauth-server-plugin
 
+## 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.13",
   "type": "module",
   "main": "dist/index.js",
   "types": "src/index.ts",

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

@@ -1,5 +1,36 @@
 # jazz-react-auth-betterauth
 
+## 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.13",
 	"type": "module",
 	"main": "dist/index.js",
 	"types": "src/index.tsx",

packages/jazz-run/CHANGELOG.md

@@ -1,5 +1,36 @@
 # jazz-run
 
+## 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.13",
   "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.13",
+    "cojson-storage-sqlite": "workspace:0.15.13",
+    "cojson-transport-ws": "workspace:0.15.13",
     "effect": "^3.6.5",
-    "jazz-tools": "workspace:0.15.10",
+    "jazz-tools": "workspace:0.15.13",
     "ws": "^8.14.2"
   },
   "devDependencies": {

packages/jazz-tools/CHANGELOG.md

@@ -1,5 +1,34 @@
 # jazz-tools
 
+## 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.13",
   "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/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/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,

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

@@ -6,9 +6,12 @@ import {
 } from "../../../internal.js";
 import { z } from "../zodReExport.js";
 
-export type FileStreamSchema = z.core.$ZodCustom<FileStream, unknown> & {
+export type AnyFileStreamSchema = z.core.$ZodCustom<FileStream, unknown> & {
   collaborative: true;
   builtin: "FileStream";
+};
+
+export type FileStreamSchema = AnyFileStreamSchema & {
   create(options?: { owner?: Account | Group } | Account | Group): FileStream;
   createFromBlob(
     blob: Blob | File,
@@ -40,5 +43,32 @@ export type FileStreamSchema = z.core.$ZodCustom<FileStream, unknown> & {
     id: string,
     listener: (value: FileStream, unsubscribe: () => void) => void,
   ): () => void;
-  getCoSchema: () => typeof FileStream;
+  getCoValueClass: () => typeof FileStream;
 };
+
+export function enrichFileStreamSchema(
+  schema: AnyFileStreamSchema,
+  coValueClass: typeof FileStream,
+): FileStreamSchema {
+  return Object.assign(schema, {
+    create: (...args: any[]) => {
+      return coValueClass.create(...args);
+    },
+    createFromBlob: (...args: [any, ...any[]]) => {
+      return coValueClass.createFromBlob(...args);
+    },
+    load: (...args: [any, ...any[]]) => {
+      return coValueClass.load(...args);
+    },
+    loadAsBlob: (...args: [any, ...any[]]) => {
+      return coValueClass.loadAsBlob(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as FileStreamSchema;
+}

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

@@ -3,9 +3,12 @@ import { Account, CoPlainText, Group } from "../../../internal.js";
 import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
 import { z } from "../zodReExport.js";
 
-export type PlainTextSchema = z.core.$ZodCustom<CoPlainText, unknown> & {
+export type AnyPlainTextSchema = z.core.$ZodCustom<CoPlainText, unknown> & {
   collaborative: true;
   builtin: "CoPlainText";
+};
+
+export type PlainTextSchema = AnyPlainTextSchema & {
   create(
     text: string,
     options?: { owner: Account | Group } | Account | Group,
@@ -24,5 +27,30 @@ export type PlainTextSchema = z.core.$ZodCustom<CoPlainText, unknown> & {
     listener: (value: CoPlainText, unsubscribe: () => void) => void,
   ): () => void;
   fromRaw(raw: RawCoPlainText): CoPlainText;
-  getCoSchema: () => typeof CoPlainText;
+  getCoValueClass: () => typeof CoPlainText;
 };
+
+export function enrichPlainTextSchema(
+  schema: AnyPlainTextSchema,
+  coValueClass: typeof CoPlainText,
+): PlainTextSchema {
+  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);
+    },
+    fromRaw: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.fromRaw(...args);
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as PlainTextSchema;
+}

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

@@ -2,9 +2,12 @@ import { Account, CoRichText, Group } from "../../../internal.js";
 import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
 import { z } from "../zodReExport.js";
 
-export type RichTextSchema = z.core.$ZodCustom<CoRichText, unknown> & {
+export type AnyRichTextSchema = z.core.$ZodCustom<CoRichText, unknown> & {
   collaborative: true;
   builtin: "CoRichText";
+};
+
+export type RichTextSchema = AnyRichTextSchema & {
   create(
     text: string,
     options?: { owner: Account | Group } | Account | Group,
@@ -22,5 +25,26 @@ export type RichTextSchema = z.core.$ZodCustom<CoRichText, unknown> & {
     id: string,
     listener: (value: CoRichText, unsubscribe: () => void) => void,
   ): () => void;
-  getCoSchema: () => typeof CoRichText;
+  getCoValueClass: () => typeof CoRichText;
 };
+
+export function enrichRichTextSchema(
+  schema: AnyRichTextSchema,
+  coValueClass: typeof CoRichText,
+): RichTextSchema {
+  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);
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as RichTextSchema;
+}

packages/jazz-tools/src/tools/implementation/zodSchema/typeConverters/InstanceOfSchema.ts

@@ -1,5 +1,7 @@
 import {
   Account,
+  AnyAccountSchema,
+  AnyCoRecordSchema,
   CoFeed,
   CoList,
   CoMap,
@@ -11,24 +13,19 @@ import {
 import { AnyCoFeedSchema } from "../schemaTypes/CoFeedSchema.js";
 import { AnyCoListSchema } from "../schemaTypes/CoListSchema.js";
 import { AnyCoMapSchema } from "../schemaTypes/CoMapSchema.js";
-import { FileStreamSchema } from "../schemaTypes/FileStreamSchema.js";
-import { PlainTextSchema } from "../schemaTypes/PlainTextSchema.js";
-import { RichTextSchema } from "../schemaTypes/RichTextSchema.js";
+import { AnyFileStreamSchema } from "../schemaTypes/FileStreamSchema.js";
+import { AnyPlainTextSchema } from "../schemaTypes/PlainTextSchema.js";
+import { AnyRichTextSchema } from "../schemaTypes/RichTextSchema.js";
 import { z } from "../zodReExport.js";
 import { InstanceOrPrimitiveOfSchema } from "./InstanceOrPrimitiveOfSchema.js";
 
 export type InstanceOfSchema<S extends CoValueClass | z.core.$ZodType> =
   S extends z.core.$ZodType
-    ? S extends z.core.$ZodObject<infer Shape> & {
-        collaborative: true;
-        builtin: "Account";
-      }
+    ? S extends AnyAccountSchema<infer Shape>
       ? {
           [key in keyof Shape]: InstanceOrPrimitiveOfSchema<Shape[key]>;
         } & Account
-      : S extends z.core.$ZodRecord<infer K, infer V> & {
-            collaborative: true;
-          }
+      : S extends AnyCoRecordSchema<infer K, infer V>
         ? {
             [key in z.output<K> & string]: InstanceOrPrimitiveOfSchema<V>;
           } & CoMap
@@ -45,11 +42,11 @@ export type InstanceOfSchema<S extends CoValueClass | z.core.$ZodType> =
             ? CoList<InstanceOrPrimitiveOfSchema<T>>
             : S extends AnyCoFeedSchema<infer T>
               ? CoFeed<InstanceOrPrimitiveOfSchema<T>>
-              : S extends PlainTextSchema
+              : S extends AnyPlainTextSchema
                 ? CoPlainText
-                : S extends RichTextSchema
+                : S extends AnyRichTextSchema
                   ? CoRichText
-                  : S extends FileStreamSchema
+                  : S extends AnyFileStreamSchema
                     ? FileStream
                     : S extends z.core.$ZodOptional<infer Inner>
                       ? InstanceOrPrimitiveOfSchema<Inner>

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

@@ -1,5 +1,12 @@
 import { RawAccount, RawCoList, RawCoMap } from "cojson";
-import { zodSchemaToCoSchema } from "./runtimeConverters/zodSchemaToCoSchema.js";
+import { CoValueClass, CoValueFromRaw } from "../../internal.js";
+import {
+  isAnyCoValueSchema,
+  zodSchemaToCoSchema,
+} from "./runtimeConverters/zodSchemaToCoSchema.js";
+import { AccountSchema } from "./schemaTypes/AccountSchema.js";
+import { CoListSchema } from "./schemaTypes/CoListSchema.js";
+import { CoMapSchema } from "./schemaTypes/CoMapSchema.js";
 import { z } from "./zodReExport.js";
 
 export function schemaUnionDiscriminatorFor(
@@ -8,7 +15,7 @@ export function schemaUnionDiscriminatorFor(
   if (isUnionOfCoMapsDeeply(schema)) {
     if (!schema._zod.disc || schema._zod.disc.size == 0) {
       throw new Error(
-        "z.union() of collaborative types is not supported, use z.discriminatedUnion() instead",
+        "z.union() of collaborative types is not supported, use co.discriminatedUnion() instead",
       );
     }
 
@@ -17,14 +24,14 @@ export function schemaUnionDiscriminatorFor(
 
     if (!field) {
       throw new Error(
-        "z.discriminatedUnion() of collaborative types with non-existent discriminator key is not supported",
+        "co.discriminatedUnion() of collaborative types with non-existent discriminator key is not supported",
       );
     }
 
     for (const value of field.values) {
       if (typeof value !== "string" && typeof value !== "number") {
         throw new Error(
-          "z.discriminatedUnion() of collaborative types with non-string or non-number discriminator value is not supported",
+          "co.discriminatedUnion() of collaborative types with non-string or non-number discriminator value is not supported",
         );
       }
     }
@@ -42,7 +49,7 @@ export function schemaUnionDiscriminatorFor(
         }
       } else {
         throw new Error(
-          "Unsupported zod type in z.discriminatedUnion() of collaborative types",
+          "Unsupported zod type in co.discriminatedUnion() of collaborative types",
         );
       }
     }
@@ -50,7 +57,7 @@ export function schemaUnionDiscriminatorFor(
     const determineSchema = (_raw: RawCoMap | RawAccount | RawCoList) => {
       if (_raw instanceof RawCoList) {
         throw new Error(
-          "z.discriminatedUnion() of collaborative types is not supported for CoLists",
+          "co.discriminatedUnion() of collaborative types is not supported for CoLists",
         );
       }
 
@@ -90,12 +97,17 @@ export function schemaUnionDiscriminatorFor(
         }
 
         if (match) {
-          return zodSchemaToCoSchema(option);
+          const coValueSchema = zodSchemaToCoSchema(option) as
+            | CoMapSchema<any>
+            | AccountSchema
+            | CoListSchema<any>;
+          return coValueSchema.getCoValueClass() as CoValueClass<any> &
+            CoValueFromRaw<any>;
         }
       }
 
       throw new Error(
-        "z.discriminatedUnion() of collaborative types with no matching discriminator value found",
+        "co.discriminatedUnion() of collaborative types with no matching discriminator value found",
       );
     };
 
@@ -120,11 +132,7 @@ export function isUnionOfCoMapsDeeply(
 function isCoMapOrUnionOfCoMapsDeeply(
   schema: z.core.$ZodType,
 ): schema is z.core.$ZodDiscriminatedUnion {
-  if (
-    schema instanceof z.core.$ZodObject &&
-    "collaborative" in schema &&
-    schema.collaborative
-  ) {
+  if (schema instanceof z.core.$ZodObject && isAnyCoValueSchema(schema)) {
     return true;
   } else if (schema instanceof z.core.$ZodUnion) {
     return schema._zod.def.options.every(isCoMapOrUnionOfCoMapsDeeply);
@@ -137,6 +145,6 @@ export function isUnionOfPrimitivesDeeply(schema: z.core.$ZodType) {
   if (schema instanceof z.core.$ZodUnion) {
     return schema._zod.def.options.every(isUnionOfPrimitivesDeeply);
   } else {
-    return !("collaborative" in schema);
+    return !isAnyCoValueSchema(schema);
   }
 }

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

@@ -1,8 +1,11 @@
 import {
-  type Account,
-  type AccountCreationProps,
   type AccountSchema,
-  type AnyCoMapSchema,
+  AnyAccountSchema,
+  AnyCoFeedSchema,
+  AnyCoListSchema,
+  AnyCoSchema,
+  AnyFileStreamSchema,
+  AnyPlainTextSchema,
   BaseAccountShape,
   CoFeed,
   type CoFeedSchema,
@@ -20,121 +23,31 @@ import {
   type Simplify,
   zodSchemaToCoSchema,
 } from "../../internal.js";
-import { RichTextSchema } from "./schemaTypes/RichTextSchema.js";
+import {
+  AnyDiscriminableCoSchema,
+  CoDiscriminatedUnionSchema,
+} from "./schemaTypes/CoDiscriminatedUnionSchema.js";
+import {
+  CoOptionalSchema,
+  createCoOptionalSchema,
+} from "./schemaTypes/CoOptionalSchema.js";
+import {
+  AnyRichTextSchema,
+  RichTextSchema,
+} from "./schemaTypes/RichTextSchema.js";
 import { z } from "./zodReExport.js";
 
-function enrichCoMapSchema<Shape extends z.core.$ZodLooseShape>(
-  schema: z.ZodObject<
-    { -readonly [P in keyof Shape]: Shape[P] },
-    z.core.$strip
-  >,
-) {
-  const baseCatchall = schema.catchall;
-
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    findUnique: (...args: any[]) => {
-      return coSchema.findUnique(...args);
-    },
-    upsertUnique: (...args: any[]) => {
-      return coSchema.upsertUnique(...args);
-    },
-    loadUnique: (...args: any[]) => {
-      return coSchema.loadUnique(...args);
-    },
-    catchall: (index: z.core.$ZodType) => {
-      return enrichCoMapSchema(baseCatchall(index));
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    withMigration: (migration: (value: any) => undefined) => {
-      coSchema.prototype.migrate = migration;
-
-      return enrichedSchema;
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as CoMapSchema<Shape>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
 export const coMapDefiner = <Shape extends z.core.$ZodLooseShape>(
   shape: Shape,
 ): CoMapSchema<Shape> => {
   const objectSchema = z.object(shape).meta({
     collaborative: true,
   });
-
-  return enrichCoMapSchema(objectSchema);
-};
-
-function enrichAccountSchema<Shape extends BaseAccountShape>(
-  schema: z.ZodObject<Shape, z.core.$strip>,
-) {
-  const enrichedSchema = Object.assign(schema, {
+  const enrichedSchema = Object.assign(objectSchema, {
     collaborative: true,
-    builtin: "Account",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    createAs: (...args: any[]) => {
-      return coSchema.createAs(...args);
-    },
-    getMe: (...args: any[]) => {
-      return coSchema.getMe(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    fromRaw: (...args: any[]) => {
-      return coSchema.fromRaw(...args);
-    },
-    withMigration: (
-      migration: (
-        value: any,
-        creationProps?: AccountCreationProps,
-      ) => void | Promise<void>,
-    ) => {
-      (coSchema.prototype as Account).migrate = async function (
-        this,
-        creationProps,
-      ) {
-        await migration(this, creationProps);
-      };
-
-      return enrichedSchema;
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as AccountSchema<Shape>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
+  });
+  return zodSchemaToCoSchema(enrichedSchema);
+};
 
 /**
  * Defines a collaborative account schema for Jazz applications.
@@ -183,11 +96,14 @@ export const coAccountDefiner = <Shape extends BaseAccountShape>(
     root: coMapDefiner({}),
   } as unknown as Shape,
 ): AccountSchema<Shape> => {
-  const objectSchema = z.object(shape).meta({
+  const schema = z.object(shape).meta({
     collaborative: true,
   });
-
-  return enrichAccountSchema(objectSchema) as unknown as AccountSchema<Shape>;
+  const enrichedSchema = Object.assign(schema, {
+    collaborative: true,
+    builtin: "Account",
+  }) as AnyAccountSchema<Shape>;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
 export const coRecordDefiner = <
@@ -203,190 +119,86 @@ export const coRecordDefiner = <
   >;
 };
 
-function enrichCoListSchema<T extends z.core.$ZodType>(schema: z.ZodArray<T>) {
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as CoListSchema<T>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
 export const coListDefiner = <T extends z.core.$ZodType>(
   element: T,
 ): CoListSchema<T> => {
-  const arraySchema = z.array(element).meta({
+  const schema = z.array(element).meta({
     collaborative: true,
   });
-
-  return enrichCoListSchema(arraySchema);
+  const enrichedSchema = Object.assign(schema, {
+    collaborative: true,
+  }) as AnyCoListSchema<T>;
+  return zodSchemaToCoSchema(enrichedSchema) as unknown as CoListSchema<T>;
 };
 
 export const coProfileDefiner = <
   Shape extends z.core.$ZodLooseShape = Simplify<DefaultProfileShape>,
 >(
-  shape: Shape & {
-    name?: z.core.$ZodString<string>;
-    inbox?: z.core.$ZodOptional<z.core.$ZodString>;
-    inboxInvite?: z.core.$ZodOptional<z.core.$ZodString>;
-  } = {} as any,
+  shape: Shape & Partial<DefaultProfileShape> = {} as any,
 ): CoProfileSchema<Shape> => {
-  const ehnancedShape = Object.assign(shape ?? {}, {
+  const ehnancedShape = Object.assign(shape, {
     name: z.string(),
     inbox: z.optional(z.string()),
     inboxInvite: z.optional(z.string()),
   });
-
   return coMapDefiner(ehnancedShape) as CoProfileSchema<Shape>;
 };
 
-function enrichCoFeedSchema<T extends z.core.$ZodType>(
-  schema: z.ZodCustom<CoFeed<unknown>, unknown>,
-  element: T,
-) {
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    builtin: "CoFeed",
-    element,
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as CoFeedSchema<T>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
 export const coFeedDefiner = <T extends z.core.$ZodType>(
   element: T,
 ): CoFeedSchema<T> => {
-  return enrichCoFeedSchema(z.instanceof(CoFeed), element);
+  const schema = z.instanceof(CoFeed);
+  const enrichedSchema = Object.assign(schema, {
+    collaborative: true,
+    builtin: "CoFeed",
+    element,
+  }) as AnyCoFeedSchema<T>;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
-function enrichFileStreamSchema(schema: z.ZodCustom<FileStream, unknown>) {
+export const coFileStreamDefiner = (): FileStreamSchema => {
+  const schema = z.instanceof(FileStream);
   const enrichedSchema = Object.assign(schema, {
     collaborative: true,
     builtin: "FileStream",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    createFromBlob: (...args: any[]) => {
-      return coSchema.createFromBlob(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    loadAsBlob: (...args: any[]) => {
-      return coSchema.loadAsBlob(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as FileStreamSchema;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
-export const coFileStreamDefiner = (): FileStreamSchema => {
-  return enrichFileStreamSchema(z.instanceof(FileStream));
+  }) as AnyFileStreamSchema;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
-function enrichPlainTextSchema(schema: z.ZodCustom<CoPlainText, unknown>) {
+export const coPlainTextDefiner = (): PlainTextSchema => {
+  const schema = z.instanceof(CoPlainText);
   const enrichedSchema = Object.assign(schema, {
     collaborative: true,
     builtin: "CoPlainText",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    fromRaw: (...args: any[]) => {
-      return coSchema.fromRaw(...args);
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as PlainTextSchema;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
-export const coPlainTextDefiner = (): PlainTextSchema => {
-  return enrichPlainTextSchema(z.instanceof(CoPlainText));
+  }) as AnyPlainTextSchema;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
-function enrichRichTextSchema(schema: z.ZodCustom<CoRichText, unknown>) {
+export const coRichTextDefiner = (): RichTextSchema => {
+  const schema = z.instanceof(CoRichText);
   const enrichedSchema = Object.assign(schema, {
     collaborative: true,
     builtin: "CoRichText",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as RichTextSchema;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
-export const coRichTextDefiner = (): RichTextSchema => {
-  return enrichRichTextSchema(z.instanceof(CoRichText));
+  }) as AnyRichTextSchema;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
 export const coImageDefiner = (): typeof ImageDefinition => {
   return ImageDefinition;
 };
+
+export const coOptionalDefiner = <T extends AnyCoSchema>(
+  schema: T,
+): CoOptionalSchema<T> => {
+  return createCoOptionalSchema(schema);
+};
+
+export const coDiscriminatedUnionDefiner = <
+  T extends readonly [AnyDiscriminableCoSchema, ...AnyDiscriminableCoSchema[]],
+>(
+  discriminator: string,
+  schemas: T,
+): CoDiscriminatedUnionSchema<T> => {
+  const schema = z.discriminatedUnion(discriminator, schemas as any);
+  return zodSchemaToCoSchema(schema);
+};

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

@@ -1,13 +1,19 @@
+import {
+  ZodArray,
+  ZodTuple,
+  ZodUnion,
+  core,
+  array as zodArray,
+  tuple as zodTuple,
+  union as zodUnion,
+} from "zod/v4";
 export {
   string,
   number,
   boolean,
-  union,
   object,
-  array,
   templateLiteral,
   json,
-  tuple,
   date,
   emoji,
   base64,
@@ -29,8 +35,8 @@ export {
   int32,
   strictObject,
   discriminatedUnion,
-  //   intersection,
-  //   record,
+  // record,
+  // intersection,
   int,
   optional,
   type ZodOptional,
@@ -42,3 +48,26 @@ export {
   type ZodDiscriminatedUnion,
   z,
 } from "zod/v4";
+
+type NonCoZodType = core.$ZodType & { collaborative?: false };
+
+export function union<const T extends readonly NonCoZodType[]>(
+  options: T,
+  params?: string | core.$ZodUnionParams,
+): ZodUnion<T> {
+  return zodUnion(options, params);
+}
+
+export function array<T extends NonCoZodType>(
+  element: T,
+  params?: string | core.$ZodArrayParams,
+): ZodArray<T> {
+  return zodArray(element, params);
+}
+
+export function tuple<T extends readonly [NonCoZodType, ...NonCoZodType[]]>(
+  options: T,
+  params?: string | core.$ZodTupleParams,
+): ZodTuple<T> {
+  return zodTuple(options, params);
+}

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

@@ -2,23 +2,47 @@ import { LocalNode, RawAccount } from "cojson";
 import {
   Account,
   AccountClass,
-  CoValue,
   CoValueClass,
   CoValueFromRaw,
+  InstanceOfSchema,
   RefsToResolve,
   RefsToResolveStrict,
   Resolved,
   Simplify,
 } from "../../internal.js";
-import { AnyAccountSchema } from "./schemaTypes/AccountSchema.js";
-import { AnyCoFeedSchema } from "./schemaTypes/CoFeedSchema.js";
-import { AnyCoListSchema } from "./schemaTypes/CoListSchema.js";
-import { AnyCoMapSchema, CoMapInitZod } from "./schemaTypes/CoMapSchema.js";
-import { AnyCoRecordSchema } from "./schemaTypes/CoRecordSchema.js";
-import { FileStreamSchema } from "./schemaTypes/FileStreamSchema.js";
-import { PlainTextSchema } from "./schemaTypes/PlainTextSchema.js";
-import { RichTextSchema } from "./schemaTypes/RichTextSchema.js";
-import { InstanceOfSchema } from "./typeConverters/InstanceOfSchema.js";
+import {
+  AccountSchema,
+  AnyAccountSchema,
+  BaseAccountShape,
+} from "./schemaTypes/AccountSchema.js";
+import {
+  AnyDiscriminableCoSchema,
+  CoDiscriminatedUnionSchema,
+} from "./schemaTypes/CoDiscriminatedUnionSchema.js";
+import { AnyCoFeedSchema, CoFeedSchema } from "./schemaTypes/CoFeedSchema.js";
+import { AnyCoListSchema, CoListSchema } from "./schemaTypes/CoListSchema.js";
+import {
+  AnyCoMapSchema,
+  CoMapInitZod,
+  CoMapSchema,
+} from "./schemaTypes/CoMapSchema.js";
+import { AnyCoOptionalSchema } from "./schemaTypes/CoOptionalSchema.js";
+import {
+  AnyCoRecordSchema,
+  CoRecordSchema,
+} from "./schemaTypes/CoRecordSchema.js";
+import {
+  AnyFileStreamSchema,
+  FileStreamSchema,
+} from "./schemaTypes/FileStreamSchema.js";
+import {
+  AnyPlainTextSchema,
+  PlainTextSchema,
+} from "./schemaTypes/PlainTextSchema.js";
+import {
+  AnyRichTextSchema,
+  RichTextSchema,
+} from "./schemaTypes/RichTextSchema.js";
 import { InstanceOfSchemaCoValuesNullable } from "./typeConverters/InstanceOfSchemaCoValuesNullable.js";
 import { z } from "./zodReExport.js";
 
@@ -47,7 +71,7 @@ export type AnyCoUnionSchema = z.core.$ZodDiscriminatedUnion<
 // this is a series of hacks to work around z4 removing _zod at runtime from z.core.$ZodType
 export function isZodObject(
   schema: z.core.$ZodType,
-): schema is z.core.$ZodObject<any, any> {
+): schema is z.ZodObject<any, any> {
   return (schema as any).def?.type === "object";
 }
 
@@ -67,19 +91,53 @@ export function getDef<S extends z.core.$ZodType>(schema: S): S["_zod"]["def"] {
   return (schema as any).def;
 }
 
+// TODO rename. This represents a CoValue class or a CoValue schema
 export type CoValueOrZodSchema = CoValueClass | AnyCoSchema;
 
-export type CoValueClassFromZodSchema<S extends z.core.$ZodType> = CoValueClass<
-  InstanceOfSchema<S>
-> &
-  CoValueFromRaw<InstanceOfSchema<S>> &
-  (S extends AnyAccountSchema ? AccountClassEssentials : {});
+// TODO rename to CoValueSchemaFromCoProtoSchema
+export type CoValueSchemaFromZodSchema<S extends z.core.$ZodType> =
+  S extends z.core.$ZodType
+    ? S extends AnyAccountSchema<infer Shape extends BaseAccountShape>
+      ? AccountSchema<Shape>
+      : S extends AnyCoRecordSchema<infer K, infer V>
+        ? CoRecordSchema<K, V>
+        : S extends AnyCoMapSchema<infer Shape, infer Config>
+          ? CoMapSchema<Shape, Config>
+          : S extends AnyCoListSchema<infer T>
+            ? CoListSchema<T>
+            : S extends AnyCoFeedSchema<infer T>
+              ? CoFeedSchema<T>
+              : S extends AnyPlainTextSchema
+                ? PlainTextSchema
+                : S extends AnyRichTextSchema
+                  ? RichTextSchema
+                  : S extends AnyFileStreamSchema
+                    ? FileStreamSchema
+                    : S extends z.core.$ZodOptional<infer Inner>
+                      ? CoValueSchemaFromZodSchema<Inner>
+                      : S extends z.core.$ZodUnion<
+                            infer Members extends readonly [
+                              AnyDiscriminableCoSchema,
+                              ...AnyDiscriminableCoSchema[],
+                            ]
+                          >
+                        ? CoDiscriminatedUnionSchema<Members>
+                        : never
+    : never;
+
+export type CoValueClassFromAnySchema<S extends CoValueOrZodSchema> =
+  S extends CoValueClass<any>
+    ? S
+    : CoValueClass<InstanceOfSchema<S>> &
+        CoValueFromRaw<InstanceOfSchema<S>> &
+        (S extends AnyAccountSchema ? AccountClassEssentials : {});
 
 type AccountClassEssentials = {
   fromRaw: <A extends Account>(this: AccountClass<A>, raw: RawAccount) => A;
   fromNode: <A extends Account>(this: AccountClass<A>, node: LocalNode) => A;
 };
 
+// TODO rename to ProtoCoSchema?
 export type AnyCoSchema =
   | AnyCoMapSchema
   | AnyAccountSchema
@@ -87,9 +145,10 @@ export type AnyCoSchema =
   | AnyCoListSchema
   | AnyCoFeedSchema
   | AnyCoUnionSchema
-  | PlainTextSchema
-  | RichTextSchema
-  | FileStreamSchema;
+  | AnyCoOptionalSchema
+  | AnyPlainTextSchema
+  | AnyRichTextSchema
+  | AnyFileStreamSchema;
 
 export type Loaded<
   T extends CoValueClass | AnyCoSchema,

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

@@ -33,6 +33,7 @@ export * from "./implementation/zodSchema/zodSchema.js";
 export * from "./implementation/zodSchema/zodCo.js";
 export * as co from "./implementation/zodSchema/coExport.js";
 export * from "./implementation/zodSchema/schemaTypes/CoMapSchema.js";
+export * from "./implementation/zodSchema/schemaTypes/CoDiscriminatedUnionSchema.js";
 export * from "./implementation/zodSchema/schemaTypes/CoRecordSchema.js";
 export * from "./implementation/zodSchema/schemaTypes/CoListSchema.js";
 export * from "./implementation/zodSchema/schemaTypes/CoFeedSchema.js";

packages/jazz-tools/src/tools/subscribe/CoValueCoreSubscription.ts

@@ -28,6 +28,10 @@ export class CoValueCoreSubscription {
             this.subscribeToState();
             this.listener("unavailable");
           }
+        })
+        .catch((error) => {
+          console.error("Unexpected error loading CoValue: ", error);
+          this.listener("unavailable");
         });
     }
   }

packages/jazz-tools/src/tools/tests/ContextManager.test.ts

@@ -1,16 +1,15 @@
+import { StorageAPI } from "cojson";
 import { WasmCrypto } from "cojson/crypto/WasmCrypto";
 import { beforeEach, describe, expect, test, vi } from "vitest";
 import {
   Account,
   AccountClass,
   AuthSecretStorage,
-  CoMap,
   Group,
   InMemoryKVStore,
   JazzAuthContext,
   KvStoreContext,
   co,
-  coField,
   z,
 } from "../exports";
 import {
@@ -26,6 +25,7 @@ import {
   CoValueFromRaw,
   InstanceOfSchema,
   Loaded,
+  anySchemaToCoSchema,
   zodSchemaToCoSchema,
 } from "../internal";
 import {
@@ -33,6 +33,7 @@ import {
   getPeerConnectedToTestSyncServer,
   setupJazzTestSync,
 } from "../testing";
+import { createAsyncStorage, getDbPath } from "./testStorage";
 
 const Crypto = await WasmCrypto.create();
 
@@ -41,12 +42,14 @@ class TestJazzContextManager<Acc extends Account> extends JazzContextManager<
   JazzContextManagerBaseProps<Acc> & {
     defaultProfileName?: string;
     AccountSchema?: AccountClass<Acc>;
+    storage?: string;
   }
 > {
   async getNewContext(
     props: JazzContextManagerBaseProps<Acc> & {
       defaultProfileName?: string;
       AccountSchema?: AccountClass<Acc> & CoValueFromRaw<Acc>;
+      storage?: string;
     },
     authProps?: JazzContextManagerAuthProps,
   ) {
@@ -59,6 +62,7 @@ class TestJazzContextManager<Acc extends Account> extends JazzContextManager<
       sessionProvider: randomSessionProvider,
       authSecretStorage: this.getAuthSecretStorage(),
       AccountSchema: props.AccountSchema,
+      storage: await createAsyncStorage({ filename: props.storage }),
     });
 
     return {
@@ -77,12 +81,14 @@ class TestJazzContextManager<Acc extends Account> extends JazzContextManager<
 describe("ContextManager", () => {
   let manager: TestJazzContextManager<Account>;
   let authSecretStorage: AuthSecretStorage;
+  let storage: StorageAPI;
 
   function getCurrentValue() {
     return manager.getCurrentValue() as JazzAuthContext<Account>;
   }
 
   beforeEach(async () => {
+    storage = await createAsyncStorage({});
     KvStoreContext.getInstance().initialize(new InMemoryKVStore());
     authSecretStorage = new AuthSecretStorage();
     await authSecretStorage.clear();
@@ -233,6 +239,78 @@ describe("ContextManager", () => {
     expect(onAnonymousAccountDiscarded).not.toHaveBeenCalled();
   });
 
+  test("onAnonymousAccountDiscarded should not block the authentication when storage is active", async () => {
+    const dbFilename = getDbPath();
+
+    const AccountRoot = co.map({
+      value: z.string(),
+      get transferredRoot(): z.ZodOptional<typeof AccountRoot> {
+        return co.optional(AccountRoot);
+      },
+    });
+
+    let lastRootId: string | undefined;
+
+    const CustomAccount = co
+      .account({
+        root: AccountRoot,
+        profile: co.profile(),
+      })
+      .withMigration(async (account) => {
+        account.root = AccountRoot.create(
+          {
+            value: "Hello",
+          },
+          Group.create(this).makePublic(),
+        );
+      });
+
+    const customManager = new TestJazzContextManager<
+      InstanceOfSchema<typeof CustomAccount>
+    >();
+
+    await customManager.createContext({
+      AccountSchema: anySchemaToCoSchema(CustomAccount),
+      storage: dbFilename,
+      onAnonymousAccountDiscarded: async (anonymousAccount) => {
+        const anonymousAccountWithRoot = await anonymousAccount.ensureLoaded({
+          resolve: { root: true },
+        });
+
+        const me = await CustomAccount.getMe().ensureLoaded({
+          resolve: { root: true },
+        });
+
+        me.root.transferredRoot = anonymousAccountWithRoot.root;
+      },
+    });
+
+    const prevContextNode = customManager.getCurrentValue()!.node;
+
+    expect(prevContextNode.storage).toBeDefined();
+
+    const account = (
+      customManager.getCurrentValue() as JazzAuthContext<
+        InstanceOfSchema<typeof CustomAccount>
+      >
+    ).me;
+
+    await customManager.authenticate({
+      accountID: account.id,
+      accountSecret: account._raw.core.node.getCurrentAgent().agentSecret,
+      provider: "test",
+    });
+
+    // The storage should be closed and set to undefined
+    expect(prevContextNode.storage).toBeUndefined();
+
+    const me = await CustomAccount.getMe().ensureLoaded({
+      resolve: { root: { transferredRoot: true } },
+    });
+
+    expect(me.root.transferredRoot?.value).toBe("Hello");
+  });
+
   test("the migration should be applied correctly on existing accounts", async () => {
     const AccountRoot = co.map({
       value: z.string(),
@@ -258,7 +336,7 @@ describe("ContextManager", () => {
 
     // Create initial anonymous context
     await customManager.createContext({
-      AccountSchema: zodSchemaToCoSchema(CustomAccount),
+      AccountSchema: anySchemaToCoSchema(CustomAccount),
     });
 
     const account = (
@@ -267,8 +345,6 @@ describe("ContextManager", () => {
       >
     ).me;
 
-    console.log("before", account._refs.root?.id);
-
     await customManager.authenticate({
       accountID: account.id,
       accountSecret: account._raw.core.node.getCurrentAgent().agentSecret,
@@ -279,8 +355,6 @@ describe("ContextManager", () => {
       resolve: { root: true },
     });
 
-    console.log("after", me._refs.root?.id);
-
     expect(me.root.id).toBe(lastRootId);
   });
 
@@ -313,7 +387,7 @@ describe("ContextManager", () => {
 
     // Create initial anonymous context
     await customManager.createContext({
-      AccountSchema: zodSchemaToCoSchema(CustomAccount),
+      AccountSchema: anySchemaToCoSchema(CustomAccount),
     });
 
     const account = (
@@ -339,7 +413,7 @@ describe("ContextManager", () => {
     const AccountRoot = co.map({
       value: z.string(),
       get transferredRoot(): z.ZodOptional<typeof AccountRoot> {
-        return z.optional(AccountRoot);
+        return co.optional(AccountRoot);
       },
     });
 
@@ -385,7 +459,7 @@ describe("ContextManager", () => {
     // Create initial anonymous context
     await customManager.createContext({
       onAnonymousAccountDiscarded,
-      AccountSchema: zodSchemaToCoSchema(CustomAccount),
+      AccountSchema: anySchemaToCoSchema(CustomAccount),
     });
 
     const account = await createJazzTestAccount({

packages/jazz-tools/src/tools/tests/coDiscriminatedUnion.test-d.ts

@@ -0,0 +1,57 @@
+import { beforeEach, describe, test } from "vitest";
+import { Loaded, co, z } from "../exports.js";
+import { createJazzTestAccount, setupJazzTestSync } from "../testing.js";
+
+describe("co.discriminatedUnion", () => {
+  beforeEach(async () => {
+    await setupJazzTestSync();
+
+    await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+      creationProps: { name: "Hermes Puggington" },
+    });
+  });
+
+  test("can use co.discriminatedUnion with CoValue schemas as values", () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+    const Cat = co.map({
+      type: z.literal("cat"),
+    });
+    const Person = co.map({
+      pet: co.discriminatedUnion("type", [Dog, Cat]),
+    });
+
+    const person = Person.create({
+      pet: Dog.create({
+        type: "dog",
+      }),
+    });
+
+    person.pet = Cat.create({
+      type: "cat",
+    });
+
+    type ExpectedType = {
+      pet: Loaded<typeof Dog> | Loaded<typeof Cat>;
+    };
+
+    function matches(value: ExpectedType) {
+      return value;
+    }
+
+    matches(person);
+  });
+
+  test("cannot use co.discriminatedUnion with zod schemas as values", () => {
+    const Person = co.map({
+      pet: co.discriminatedUnion("type", [
+        // @ts-expect-error: cannot use co.discriminatedUnion with a Zod schema
+        z.object({
+          type: z.string(),
+        }),
+      ]),
+    });
+  });
+});

packages/jazz-tools/src/tools/tests/coDiscriminatedUnion.test.ts

@@ -0,0 +1,93 @@
+import { beforeEach, describe, expect, test, vi } from "vitest";
+import { Loaded, co, z } from "../exports.js";
+import { createJazzTestAccount, setupJazzTestSync } from "../testing.js";
+import { waitFor } from "./utils.js";
+
+describe("co.discriminatedUnion", () => {
+  beforeEach(async () => {
+    await setupJazzTestSync();
+
+    await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+      creationProps: { name: "Hermes Puggington" },
+    });
+  });
+
+  test("use co.discriminatedUnion with CoValue schemas as values", () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+    const Cat = co.map({
+      type: z.literal("cat"),
+    });
+    const Person = co.map({
+      pet: co.discriminatedUnion("type", [Dog, Cat]),
+    });
+
+    const person = Person.create({
+      pet: Dog.create({
+        type: "dog",
+      }),
+    });
+
+    expect(person.pet.type).toEqual("dog");
+
+    person.pet = Cat.create({
+      type: "cat",
+    });
+
+    expect(person.pet.type).toEqual("cat");
+  });
+
+  test("load CoValue instances using the DiscriminatedUnion schema", async () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+    const Cat = co.map({
+      type: z.literal("cat"),
+    });
+    const Pet = co.discriminatedUnion("type", [Dog, Cat]);
+
+    const dog = Dog.create({ type: "dog" });
+    const loadedPet = await Pet.load(dog.id);
+    expect(loadedPet?.type).toEqual("dog");
+  });
+
+  test("subscribe to CoValue instances using the DiscriminatedUnion schema", async () => {
+    const Person = co.map({
+      name: z.string(),
+    });
+    const Dog = co.map({
+      type: z.literal("dog"),
+      name: z.string(),
+      owner: Person,
+    });
+    const Cat = co.map({
+      type: z.literal("cat"),
+      name: z.string(),
+    });
+    const Pet = co.discriminatedUnion("type", [Dog, Cat]);
+
+    const dog = Dog.create({
+      type: "dog",
+      name: "Rex",
+      owner: Person.create({ name: "John Doe" }),
+    });
+
+    const updates: Loaded<typeof Pet>[] = [];
+    const spy = vi.fn((pet) => updates.push(pet));
+
+    Pet.subscribe(dog.id, {}, (pet) => {
+      expect(pet.type).toEqual("dog");
+      spy(pet);
+    });
+
+    expect(spy).not.toHaveBeenCalled();
+
+    await waitFor(() => expect(spy).toHaveBeenCalled());
+
+    expect(spy).toHaveBeenCalledTimes(1);
+
+    expect(updates[0]?.name).toEqual("Rex");
+  });
+});

packages/jazz-tools/src/tools/tests/coList.test.ts

@@ -1,15 +1,7 @@
 import { WasmCrypto } from "cojson/crypto/WasmCrypto";
 import { beforeEach, describe, expect, test, vi } from "vitest";
-import {
-  Account,
-  CoList,
-  CoMap,
-  Group,
-  coField,
-  subscribeToCoValue,
-  z,
-} from "../index.js";
-import { Loaded, co, zodSchemaToCoSchema } from "../internal.js";
+import { Account, Group, subscribeToCoValue, z } from "../index.js";
+import { Loaded, anySchemaToCoSchema, co } from "../internal.js";
 import { createJazzTestAccount, setupJazzTestSync } from "../testing.js";
 import { waitFor } from "./utils.js";
 
@@ -127,7 +119,7 @@ describe("Simple CoList operations", async () => {
         name: z.string(),
       });
 
-      const Recipe = co.list(z.optional(Ingredient));
+      const Recipe = co.list(co.optional(Ingredient));
 
       const recipe = Recipe.create(
         [
@@ -588,7 +580,7 @@ describe("CoList subscription", async () => {
     const spy = vi.fn((list) => updates.push(list));
 
     subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         syncResolution: true,

packages/jazz-tools/src/tools/tests/coMap.record.test.ts

@@ -281,7 +281,7 @@ describe("CoMap.Record", async () => {
     });
 
     test("Is ok to omit optional fields", () => {
-      const TestRecord = co.record(z.string(), z.optional(NestedRecord));
+      const TestRecord = co.record(z.string(), co.optional(NestedRecord));
 
       expectTypeOf<typeof TestRecord.create>().toBeCallableWith(
         {
@@ -315,7 +315,7 @@ describe("CoMap.Record", async () => {
 
     const PersonRecord = co.record(
       z.string(),
-      z.discriminatedUnion("type", [Base, IssueRepro]),
+      co.discriminatedUnion("type", [Base, IssueRepro]),
     );
 
     const person = IssueRepro.create({
@@ -350,7 +350,7 @@ describe("CoMap.Record", async () => {
 
     const PersonRecord = co.record(
       z.string(),
-      z.discriminatedUnion("type", [Base, IssueRepro]),
+      co.discriminatedUnion("type", [Base, IssueRepro]),
     );
 
     const person = IssueRepro.create({
@@ -385,7 +385,7 @@ describe("CoMap.Record", async () => {
 
     const PersonRecord = co.record(
       z.string(),
-      z.discriminatedUnion("type", [Base, IssueRepro]),
+      co.discriminatedUnion("type", [Base, IssueRepro]),
     );
 
     const person = IssueRepro.create({

packages/jazz-tools/src/tools/tests/coMap.test-d.ts

@@ -159,7 +159,7 @@ describe("CoMap", async () => {
     test("Comap with recursive optional reference", () => {
       const Recursive = co.map({
         get child(): z.ZodOptional<typeof Recursive> {
-          return z.optional(Recursive);
+          return co.optional(Recursive);
         },
       });
 
@@ -185,7 +185,7 @@ describe("CoMap", async () => {
         age: z.number(),
         // TODO: would be nice if this didn't need a type annotation
         get friend(): z.ZodOptional<typeof Person> {
-          return z.optional(Person);
+          return co.optional(Person);
         },
       });
 
@@ -280,7 +280,7 @@ describe("CoMap", async () => {
 
     const MapWithEnumOfMaps = co.map({
       name: z.string(),
-      child: z.discriminatedUnion("type", [ChildA, ChildB]),
+      child: co.discriminatedUnion("type", [ChildA, ChildB]),
     });
 
     const mapWithEnum = MapWithEnumOfMaps.create({

packages/jazz-tools/src/tools/tests/coMap.test.ts

@@ -12,7 +12,7 @@ import {
 } from "vitest";
 import { Group, co, subscribeToCoValue, z } from "../exports.js";
 import { Account } from "../index.js";
-import { ID, Loaded, zodSchemaToCoSchema } from "../internal.js";
+import { Loaded, anySchemaToCoSchema } from "../internal.js";
 import {
   createJazzTestAccount,
   getPeerConnectedToTestSyncServer,
@@ -145,7 +145,7 @@ describe("CoMap", async () => {
         age: z.number(),
         // TODO: would be nice if this didn't need a type annotation
         get friend(): z.ZodOptional<typeof Person> {
-          return z.optional(Person);
+          return co.optional(Person);
         },
       });
 
@@ -182,7 +182,7 @@ describe("CoMap", async () => {
         name: z.string(),
         age: z.number(),
         get friend(): z.ZodOptional<typeof Person> {
-          return z.optional(Person);
+          return co.optional(Person);
         },
       });
 
@@ -211,7 +211,7 @@ describe("CoMap", async () => {
         name: z.string(),
         age: z.number(),
         get friend(): z.ZodOptional<typeof Person> {
-          return z.optional(Person);
+          return co.optional(Person);
         },
       });
 
@@ -400,7 +400,7 @@ describe("CoMap", async () => {
 
     const MapWithEnumOfMaps = co.map({
       name: z.string(),
-      child: z.discriminatedUnion("type", [ChildA, ChildB]),
+      child: co.discriminatedUnion("type", [ChildA, ChildB]),
     });
 
     const mapWithEnum = MapWithEnumOfMaps.create({
@@ -801,7 +801,7 @@ describe("CoMap resolution", async () => {
     const spy = vi.fn((person) => updates.push(person));
 
     subscribeToCoValue(
-      zodSchemaToCoSchema(Person), // TODO: we should get rid of the conversion in the future
+      anySchemaToCoSchema(Person), // TODO: we should get rid of the conversion in the future
       person.id,
       {
         syncResolution: true,
@@ -1003,7 +1003,7 @@ describe("CoMap applyDiff", async () => {
     birthday: z.date(),
     nested: NestedMap,
     optionalField: z.string().optional(),
-    optionalNested: z.optional(NestedMap),
+    optionalNested: co.optional(NestedMap),
   });
 
   test("Basic applyDiff", () => {
@@ -1698,7 +1698,7 @@ describe("Creating and finding unique CoMaps", async () => {
       dateValue: z.date(),
     });
 
-    const AttributeValue = z.discriminatedUnion("type", [
+    const AttributeValue = co.discriminatedUnion("type", [
       StringAttributeValue,
       NumberAttributeValue,
       DateAttributeValue,
@@ -1714,7 +1714,7 @@ describe("Creating and finding unique CoMaps", async () => {
       attributeValue: AttributeValue,
     });
 
-    const Tag = z.discriminatedUnion("type", [
+    const Tag = co.discriminatedUnion("type", [
       AttributeTag,
       StringTag,
       DateTag,
@@ -1764,7 +1764,7 @@ describe("Creating and finding unique CoMaps", async () => {
       error: HttpError,
     });
 
-    const ErrorResponse = z.discriminatedUnion("type", [
+    const ErrorResponse = co.discriminatedUnion("type", [
       ClientError,
       ServerError,
       NetworkError,

packages/jazz-tools/src/tools/tests/coOptional.test.ts

@@ -0,0 +1,39 @@
+import { beforeEach, describe, test } from "vitest";
+import { CoPlainText, co, z } from "../exports.js";
+import { createJazzTestAccount, setupJazzTestSync } from "../testing.js";
+
+describe("co.optional", () => {
+  beforeEach(async () => {
+    await setupJazzTestSync();
+
+    await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+      creationProps: { name: "Hermes Puggington" },
+    });
+  });
+
+  test("can use co.optional with CoValue schemas as values", () => {
+    const Person = co.map({
+      preferredName: co.optional(co.plainText()),
+    });
+
+    const person = Person.create({});
+
+    type ExpectedType = {
+      preferredName?: CoPlainText;
+    };
+
+    function matches(value: ExpectedType) {
+      return value;
+    }
+
+    matches(person);
+  });
+
+  test("cannot use co.optional with zod schemas as values", () => {
+    const Person = co.map({
+      // @ts-expect-error: cannot use co.optional with a Zod schema
+      preferredName: co.optional(z.string()),
+    });
+  });
+});

packages/jazz-tools/src/tools/tests/createContext.test.ts

@@ -1,4 +1,4 @@
-import { AgentSecret, RawCoMap } from "cojson";
+import { AgentSecret } from "cojson";
 import { WasmCrypto } from "cojson/crypto/WasmCrypto";
 import { beforeEach, describe, expect, test, vi } from "vitest";
 import {
@@ -9,6 +9,7 @@ import {
   ID,
   InMemoryKVStore,
   KvStoreContext,
+  anySchemaToCoSchema,
   co,
   createAnonymousJazzContext,
   createJazzContext,
@@ -90,7 +91,7 @@ describe("createContext methods", () => {
       });
 
       expect(context.account).toBeInstanceOf(
-        zodSchemaToCoSchema(CustomAccount),
+        anySchemaToCoSchema(CustomAccount),
       );
     });
 
@@ -199,7 +200,7 @@ describe("createContext methods", () => {
       });
 
       expect(context.account).toBeInstanceOf(
-        zodSchemaToCoSchema(CustomAccount),
+        anySchemaToCoSchema(CustomAccount),
       );
     });
 
@@ -350,7 +351,7 @@ describe("createContext methods", () => {
       });
 
       expect(context.account).toBeInstanceOf(
-        zodSchemaToCoSchema(CustomAccount),
+        anySchemaToCoSchema(CustomAccount),
       );
     });
   });

packages/jazz-tools/src/tools/tests/deepLoading.test.ts

@@ -36,7 +36,7 @@ const TestList = co.list(InnerMap);
 
 const TestMap = co.map({
   list: TestList,
-  optionalRef: z.optional(InnermostMap),
+  optionalRef: co.optional(InnermostMap),
 });
 
 describe("Deep loading with depth arg", async () => {
@@ -599,7 +599,7 @@ describe("Deep loading with unauthorized account", async () => {
     });
 
     const Lv2 = co.map({
-      lv3: z.optional(Lv3),
+      lv3: co.optional(Lv3),
     });
 
     const Lv1 = co.map({
@@ -737,7 +737,7 @@ describe("Deep loading with unauthorized account", async () => {
     const Person = co.map({
       name: z.string(),
       get friends(): z.ZodOptional<typeof Friends> {
-        return z.optional(Friends);
+        return co.optional(Friends);
       },
     });
     const Friends: CoListSchema<typeof Person> = co.list(Person); // TODO: annoying that we have to annotate

packages/jazz-tools/src/tools/tests/inbox.test.ts

@@ -1,14 +1,6 @@
 import { describe, expect, it, vi } from "vitest";
-import {
-  Account,
-  CoMap,
-  Group,
-  Inbox,
-  InboxSender,
-  Profile,
-  z,
-} from "../exports";
-import { Loaded, co, coField, zodSchemaToCoSchema } from "../internal";
+import { Group, Inbox, InboxSender, z } from "../exports";
+import { Loaded, anySchemaToCoSchema, co } from "../internal";
 import { setupTwoNodes, waitFor } from "./utils";
 
 const Message = co.map({
@@ -26,7 +18,7 @@ describe("Inbox", () => {
 
       const { clientAccount: sender, serverAccount: receiver } =
         await setupTwoNodes({
-          ServerAccountSchema: zodSchemaToCoSchema(WorkerAccount),
+          ServerAccountSchema: anySchemaToCoSchema(WorkerAccount),
         });
 
       await expect(() => InboxSender.load(receiver.id, sender)).rejects.toThrow(

packages/jazz-tools/src/tools/tests/load.test.ts

@@ -33,6 +33,129 @@ test("load a value", async () => {
   expect(john?.name).toBe("John");
 });
 
+test("return null if id is invalid", async () => {
+  const Person = co.map({
+    name: z.string(),
+  });
+
+  const john = await Person.load("test");
+  expect(john).toBeNull();
+});
+
+test("load a missing optional value (co.optional)", async () => {
+  const Dog = co.map({
+    name: z.string(),
+  });
+
+  const Person = co.map({
+    name: z.string(),
+    dog: co.optional(Dog),
+  });
+
+  const group = Group.create();
+  const map = Person.create({ name: "John" }, group);
+  group.addMember("everyone", "reader");
+
+  const alice = await createJazzTestAccount();
+
+  const john = await Person.load(map.id, {
+    loadAs: alice,
+    resolve: { dog: true },
+  });
+
+  assert(john);
+
+  expect(john.name).toBe("John");
+  expect(john.dog).toBeUndefined();
+});
+
+test("load a missing optional value (z.optional)", async () => {
+  const Dog = co.map({
+    name: z.string(),
+  });
+
+  const Person = co.map({
+    name: z.string(),
+    dog: z.optional(Dog),
+  });
+
+  const group = Group.create();
+  const map = Person.create({ name: "John" }, group);
+  group.addMember("everyone", "reader");
+
+  const alice = await createJazzTestAccount();
+
+  const john = await Person.load(map.id, {
+    loadAs: alice,
+    resolve: { dog: true },
+  });
+
+  assert(john);
+
+  expect(john.name).toBe("John");
+  expect(john.dog).toBeUndefined();
+});
+
+test("load a missing optional value (Schema.optional)", async () => {
+  const Dog = co.map({
+    name: z.string(),
+  });
+
+  const Person = co.map({
+    name: z.string(),
+    dog: Dog.optional(),
+  });
+
+  const group = Group.create();
+  const map = Person.create({ name: "John" }, group);
+  group.addMember("everyone", "reader");
+
+  const alice = await createJazzTestAccount();
+
+  const john = await Person.load(map.id, {
+    loadAs: alice,
+    resolve: { dog: true },
+  });
+
+  assert(john);
+
+  expect(john.name).toBe("John");
+  expect(john.dog).toBeUndefined();
+});
+
+test("load a missing optional value (optional discrminatedUnion)", async () => {
+  const Dog = co.map({
+    type: z.literal("dog"),
+    name: z.string(),
+  });
+
+  const Cat = co.map({
+    type: z.literal("cat"),
+    name: z.string(),
+  });
+
+  const Person = co.map({
+    name: z.string(),
+    pet: co.discriminatedUnion("type", [Dog, Cat]).optional(),
+  });
+
+  const group = Group.create();
+  const map = Person.create({ name: "John" }, group);
+  group.addMember("everyone", "reader");
+
+  const alice = await createJazzTestAccount();
+
+  const john = await Person.load(map.id, {
+    loadAs: alice,
+    resolve: { pet: true },
+  });
+
+  assert(john);
+
+  expect(john.name).toBe("John");
+  expect(john.pet).toBeUndefined();
+});
+
 test("retry an unavailable value", async () => {
   const Person = co.map({
     name: z.string(),

packages/jazz-tools/src/tools/tests/patterns/requestToJoin.test.ts

@@ -1,5 +1,5 @@
 import { assert, describe, expect, test } from "vitest";
-import { Account, CoList, CoMap, Group, ID, co, z } from "../../exports";
+import { Account, Group, co, z } from "../../exports";
 import { createJazzTestAccount, linkAccounts } from "../../testing";
 
 const RequestToJoin = co.map({

packages/jazz-tools/src/tools/tests/schemaUnion.test.ts

@@ -6,7 +6,6 @@ import {
   Loaded,
   anySchemaToCoSchema,
   co,
-  loadCoValue,
   subscribeToCoValue,
   z,
 } from "../exports.js";
@@ -24,7 +23,7 @@ const BlueButtonWidget = co.map({
   blueness: z.number(),
 });
 
-const ButtonWidget = z.discriminatedUnion("type", [
+const ButtonWidget = co.discriminatedUnion("type", [
   RedButtonWidget,
   BlueButtonWidget,
 ]);
@@ -40,7 +39,7 @@ const CheckboxWidget = co.map({
   checked: z.boolean(),
 });
 
-const WidgetUnion = z.discriminatedUnion("type", [
+const WidgetUnion = co.discriminatedUnion("type", [
   ButtonWidget,
   SliderWidget,
   CheckboxWidget,
@@ -72,25 +71,15 @@ describe("SchemaUnion", () => {
       { owner: me },
     );
 
-    const loadedButtonWidget = await loadCoValue(
-      anySchemaToCoSchema(WidgetUnion),
-      buttonWidget.id,
-      {
-        loadAs: me,
-      },
-    );
-    const loadedSliderWidget = await loadCoValue(
-      anySchemaToCoSchema(WidgetUnion),
-      sliderWidget.id,
-      {
-        loadAs: me,
-      },
-    );
-    const loadedCheckboxWidget = await loadCoValue(
-      anySchemaToCoSchema(WidgetUnion),
-      checkboxWidget.id,
-      { loadAs: me },
-    );
+    const loadedButtonWidget = await WidgetUnion.load(buttonWidget.id, {
+      loadAs: me,
+    });
+    const loadedSliderWidget = await WidgetUnion.load(sliderWidget.id, {
+      loadAs: me,
+    });
+    const loadedCheckboxWidget = await WidgetUnion.load(checkboxWidget.id, {
+      loadAs: me,
+    });
 
     expect(loadedButtonWidget?.type).toBe("button");
     expect(loadedSliderWidget?.type).toBe("slider");

packages/jazz-tools/src/tools/tests/subscribe.test.ts

@@ -7,27 +7,13 @@ import {
   onTestFinished,
   vi,
 } from "vitest";
+import { Account, Group, cojsonInternals, z } from "../index.js";
 import {
-  Account,
-  CoFeed,
-  CoList,
-  CoMap,
-  FileStream,
-  Group,
-  coField,
-  cojsonInternals,
-  z,
-} from "../index.js";
-import {
-  CoMapInstance,
-  ID,
-  InstanceOrPrimitiveOfSchema,
   Loaded,
-  Resolved,
+  anySchemaToCoSchema,
   co,
   createCoValueObservable,
   subscribeToCoValue,
-  zodSchemaToCoSchema,
 } from "../internal.js";
 import {
   createJazzTestAccount,
@@ -43,7 +29,7 @@ const ReactionsFeed = co.feed(z.string());
 const Message = co.map({
   text: z.string(),
   reactions: ReactionsFeed,
-  attachment: z.optional(co.fileStream()),
+  attachment: co.optional(co.fileStream()),
 });
 
 const ChatRoom = co.map({
@@ -84,7 +70,7 @@ describe("subscribeToCoValue", () => {
     let result = null as Loaded<typeof ChatRoom, true> | null;
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(ChatRoom),
+      anySchemaToCoSchema(ChatRoom),
       chatRoom.id,
       { loadAs: meOnSecondPeer },
       (value) => {
@@ -131,7 +117,7 @@ describe("subscribeToCoValue", () => {
     let result = null as Loaded<typeof ChatRoom, {}> | null;
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(ChatRoom),
+      anySchemaToCoSchema(ChatRoom),
       chatRoom.id,
       {
         loadAs: meOnSecondPeer,
@@ -172,7 +158,7 @@ describe("subscribeToCoValue", () => {
     messages.push(createMessage(me, "Hello"));
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(ChatRoom),
+      anySchemaToCoSchema(ChatRoom),
       chatRoom.id,
       {
         loadAs: meOnSecondPeer,
@@ -215,7 +201,7 @@ describe("subscribeToCoValue", () => {
     const updateFn = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(ChatRoom),
+      anySchemaToCoSchema(ChatRoom),
       chatRoom.id,
       {
         loadAs: meOnSecondPeer,
@@ -275,7 +261,7 @@ describe("subscribeToCoValue", () => {
     >[];
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(ChatRoom),
+      anySchemaToCoSchema(ChatRoom),
       chatRoom.id,
       {
         loadAs: meOnSecondPeer,
@@ -346,7 +332,7 @@ describe("subscribeToCoValue", () => {
     const updateFn = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(ChatRoom),
+      anySchemaToCoSchema(ChatRoom),
       chatRoom.id,
       {
         loadAs: meOnSecondPeer,
@@ -410,7 +396,7 @@ describe("subscribeToCoValue", () => {
     const updateFn = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: account,
@@ -471,7 +457,7 @@ describe("subscribeToCoValue", () => {
     const onUnauthorized = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: reader,
@@ -542,7 +528,7 @@ describe("subscribeToCoValue", () => {
     const onUnavailable = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: reader,
@@ -617,7 +603,7 @@ describe("subscribeToCoValue", () => {
     const onUnavailable = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: reader,
@@ -654,7 +640,7 @@ describe("subscribeToCoValue", () => {
       value: z.string(),
     });
 
-    const TestList = co.list(z.optional(TestMap));
+    const TestList = co.list(co.optional(TestMap));
 
     const reader = await createJazzTestAccount({
       isCurrentActiveAccount: true,
@@ -687,7 +673,7 @@ describe("subscribeToCoValue", () => {
     const onUnavailable = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: reader,
@@ -718,7 +704,7 @@ describe("subscribeToCoValue", () => {
       value: z.string(),
     });
 
-    const TestList = co.list(z.optional(TestMap));
+    const TestList = co.list(co.optional(TestMap));
 
     const creator = await createJazzTestAccount({
       isCurrentActiveAccount: true,
@@ -739,7 +725,7 @@ describe("subscribeToCoValue", () => {
     });
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: creator,
@@ -806,7 +792,7 @@ describe("subscribeToCoValue", () => {
     });
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: creator,
@@ -895,7 +881,7 @@ describe("subscribeToCoValue", () => {
     });
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(Person),
+      anySchemaToCoSchema(Person),
       person.id,
       {
         loadAs: reader,
@@ -982,7 +968,7 @@ describe("subscribeToCoValue", () => {
     const onUnavailable = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(TestList),
+      anySchemaToCoSchema(TestList),
       list.id,
       {
         loadAs: reader,
@@ -1069,7 +1055,7 @@ describe("subscribeToCoValue", () => {
     const onUnavailable = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(PersonList),
+      anySchemaToCoSchema(PersonList),
       list.id,
       {
         loadAs: reader,
@@ -1165,7 +1151,7 @@ describe("subscribeToCoValue", () => {
     const onUnavailable = vi.fn();
 
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(PersonList),
+      anySchemaToCoSchema(PersonList),
       list.id,
       {
         loadAs: reader,
@@ -1248,7 +1234,7 @@ describe("subscribeToCoValue", () => {
 
     // Test subscribing to the large coValue
     const unsubscribe = subscribeToCoValue(
-      zodSchemaToCoSchema(LargeDataset),
+      anySchemaToCoSchema(LargeDataset),
       largeMap.id,
       {
         loadAs: alice,
@@ -1314,7 +1300,7 @@ describe("createCoValueObservable", () => {
     const mockListener = vi.fn();
 
     const unsubscribe = observable.subscribe(
-      zodSchemaToCoSchema(TestMap),
+      anySchemaToCoSchema(TestMap),
       testMap.id,
       {
         loadAs: meOnSecondPeer,
@@ -1343,7 +1329,7 @@ describe("createCoValueObservable", () => {
     const mockListener = vi.fn();
 
     const unsubscribe = observable.subscribe(
-      zodSchemaToCoSchema(TestMap),
+      anySchemaToCoSchema(TestMap),
       testMap.id,
       {
         loadAs: meOnSecondPeer,

packages/jazz-tools/src/tools/tests/testStorage.ts

@@ -0,0 +1,70 @@
+import { randomUUID } from "node:crypto";
+import { unlinkSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { SQLiteDatabaseDriverAsync, getSqliteStorageAsync } from "cojson";
+import Database, { type Database as DatabaseT } from "libsql";
+import { onTestFinished } from "vitest";
+
+class LibSQLSqliteAsyncDriver implements SQLiteDatabaseDriverAsync {
+  private readonly db: DatabaseT;
+
+  constructor(filename: string) {
+    this.db = new Database(filename, {});
+  }
+
+  async initialize() {
+    await this.db.pragma("journal_mode = WAL");
+  }
+
+  async run(sql: string, params: unknown[]) {
+    this.db.prepare(sql).run(params);
+  }
+
+  async query<T>(sql: string, params: unknown[]): Promise<T[]> {
+    return this.db.prepare(sql).all(params) as T[];
+  }
+
+  async get<T>(sql: string, params: unknown[]): Promise<T | undefined> {
+    return this.db.prepare(sql).get(params) as T | undefined;
+  }
+
+  async transaction(callback: () => unknown) {
+    await this.run("BEGIN TRANSACTION", []);
+
+    try {
+      await callback();
+      await this.run("COMMIT", []);
+    } catch (error) {
+      await this.run("ROLLBACK", []);
+    }
+  }
+
+  async closeDb() {
+    this.db.close();
+  }
+}
+
+export async function createAsyncStorage({ filename }: { filename?: string }) {
+  const storage = await getSqliteStorageAsync(
+    new LibSQLSqliteAsyncDriver(getDbPath(filename)),
+  );
+
+  onTestFinished(() => {
+    storage.close();
+  });
+
+  return storage;
+}
+
+export function getDbPath(defaultDbPath?: string) {
+  const dbPath = defaultDbPath ?? join(tmpdir(), `test-${randomUUID()}.db`);
+
+  if (!defaultDbPath) {
+    onTestFinished(() => {
+      unlinkSync(dbPath);
+    });
+  }
+
+  return dbPath;
+}

packages/jazz-tools/src/tools/tests/zod.test-d.ts

@@ -0,0 +1,70 @@
+import { describe, test } from "vitest";
+import { z } from "../exports";
+import { co } from "../internal";
+
+describe("CoValue and Zod schema compatibility", () => {
+  test("cannot use z.record with CoValue schemas as values", () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+
+    const Person = co.map({
+      // @ts-expect-error: cannot use z.record with a CoValue schema
+      // (z.record is not exported by jazz-tools)
+      pets: z.record(z.string(), Dog),
+    });
+  });
+
+  test("cannot use z.union with CoValue schemas as values", () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+
+    const Cat = co.map({
+      type: z.literal("cat"),
+    });
+
+    const Person = co.map({
+      // @ts-expect-error: cannot use z.union with a CoValue schema
+      pets: z.union([Dog, Cat]),
+    });
+  });
+
+  test("cannot use z.intersection with CoValue schemas as values", () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+
+    const Cat = co.map({
+      type: z.literal("cat"),
+    });
+
+    const Person = co.map({
+      // @ts-expect-error: cannot use z.intersection with a CoValue schema
+      // (z.intersection is not exported by jazz-tools)
+      pets: z.intersection(Dog, Cat),
+    });
+  });
+
+  test("cannot use z.array with CoValue schemas as values", () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+
+    const Person = co.map({
+      // @ts-expect-error: cannot use z.array with a CoValue schema
+      pets: z.array(Dog),
+    });
+  });
+
+  test("cannot use z.tuple with CoValue schemas as values", () => {
+    const Dog = co.map({
+      type: z.literal("dog"),
+    });
+
+    const Person = co.map({
+      // @ts-expect-error: cannot use z.tuple with a CoValue schema
+      pets: z.tuple([Dog]),
+    });
+  });
+});

packages/jazz-tools/src/tools/tests/zod.test.ts

@@ -1,10 +1,18 @@
-import { describe, expect, it, vi } from "vitest";
+import { beforeAll, describe, expect, it, vi } from "vitest";
 import { z } from "../exports.js";
 import { co } from "../internal.js";
-import { createJazzTestAccount } from "../testing.js";
+import { createJazzTestAccount, setupJazzTestSync } from "../testing.js";
 
 describe("co.map and Zod schema compatibility", () => {
-  // Helper function to create a test account
+  beforeAll(async () => {
+    await setupJazzTestSync();
+
+    await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+      creationProps: { name: "Hermes Puggington" },
+    });
+  });
+
   describe("Primitive types", () => {
     it("should handle string fields", async () => {
       const schema = co.map({
@@ -348,12 +356,63 @@ describe("co.map and Zod schema compatibility", () => {
     //     ]),
     //   });
     //   const account = await createJazzTestAccount();
-    //   const successMap = schema.create({ result: { status: "success", data: "data" } }, account);
-    //   const failedMap = schema.create({ result: { status: "failed", error: "error" } }, account);
+    //   const successMap = schema.create(
+    //     { result: { status: "success", data: "data" } },
+    //     account,
+    //   );
+    //   const failedMap = schema.create(
+    //     { result: { status: "failed", error: "error" } },
+    //     account,
+    //   );
     //   expect(successMap.result).toEqual({ status: "success", data: "data" });
     //   expect(failedMap.result).toEqual({ status: "failed", error: "error" });
     // });
 
+    it("should handle discriminated unions of CoValues", () => {
+      const Dog = co.map({
+        type: z.literal("dog"),
+        breed: z.string(),
+      });
+
+      const Cat = co.map({
+        type: z.literal("cat"),
+        name: z.string(),
+      });
+
+      const Person = co.map({
+        pet: z.discriminatedUnion("type", [Dog, Cat]),
+      });
+
+      const person = Person.create({
+        pet: Dog.create({ type: "dog", breed: "Labrador" }),
+      });
+
+      expect(person.pet).toEqual({ type: "dog", breed: "Labrador" });
+
+      person.pet = Cat.create({ type: "cat", name: "Whiskers" });
+
+      expect(person.pet).toEqual({ type: "cat", name: "Whiskers" });
+    });
+
+    it("should handle optional CoValues", () => {
+      const Dog = co.map({
+        name: z.string(),
+        breed: z.string(),
+      });
+
+      const Person = co.map({
+        pet: z.optional(Dog),
+      });
+
+      const person = Person.create({});
+
+      expect(person.pet).toBeUndefined();
+
+      person.pet = Dog.create({ name: "Rex", breed: "Labrador" });
+
+      expect(person.pet).toEqual({ name: "Rex", breed: "Labrador" });
+    });
+
     // it("should handle intersections", async () => {
     //   const schema = co.map({
     //     value: z.intersection(

pnpm-lock.yaml

@@ -1772,19 +1772,19 @@ importers:
         specifier: ^0.25.5
         version: 0.25.8(effect@3.11.9)
       cojson:
-        specifier: workspace:0.15.10
+        specifier: workspace:0.15.13
         version: link:../cojson
       cojson-storage-sqlite:
-        specifier: workspace:0.15.10
+        specifier: workspace:0.15.13
         version: link:../cojson-storage-sqlite
       cojson-transport-ws:
-        specifier: workspace:0.15.10
+        specifier: workspace:0.15.13
         version: link:../cojson-transport-ws
       effect:
         specifier: ^3.6.5
         version: 3.11.9
       jazz-tools:
-        specifier: workspace:0.15.10
+        specifier: workspace:0.15.13
         version: link:../jazz-tools
       ws:
         specifier: ^8.14.2

starters/react-passkey-auth/CHANGELOG.md

@@ -1,5 +1,27 @@
 # jazz-react-tailwind-starter
 
+## 0.0.133
+
+### Patch Changes
+
+- Updated dependencies [6c76ff8]
+  - jazz-tools@0.15.13
+
+## 0.0.132
+
+### Patch Changes
+
+- Updated dependencies [d1c1b0c]
+- Updated dependencies [cf4ad72]
+  - jazz-tools@0.15.12
+
+## 0.0.131
+
+### Patch Changes
+
+- Updated dependencies [bdc9aee]
+  - jazz-tools@0.15.11
+
 ## 0.0.130
 
 ### Patch Changes

starters/react-passkey-auth/package.json

@@ -1,7 +1,7 @@
 {
   "name": "jazz-react-passkey-auth-starter",
   "private": true,
-  "version": "0.0.130",
+  "version": "0.0.133",
   "type": "module",
   "scripts": {
     "dev": "vite",

starters/svelte-passkey-auth/CHANGELOG.md

@@ -1,5 +1,27 @@
 # svelte-passkey-auth
 
+## 0.0.107
+
+### Patch Changes
+
+- Updated dependencies [6c76ff8]
+  - jazz-tools@0.15.13
+
+## 0.0.106
+
+### Patch Changes
+
+- Updated dependencies [d1c1b0c]
+- Updated dependencies [cf4ad72]
+  - jazz-tools@0.15.12
+
+## 0.0.105
+
+### Patch Changes
+
+- Updated dependencies [bdc9aee]
+  - jazz-tools@0.15.11
+
 ## 0.0.104
 
 ### Patch Changes

starters/svelte-passkey-auth/package.json

@@ -1,6 +1,6 @@
 {
   "name": "svelte-passkey-auth",
-  "version": "0.0.104",
+  "version": "0.0.107",
   "type": "module",
   "private": true,
   "scripts": {

tests/e2e/src/pages/FileStream/UploaderPeer.tsx

@@ -1,4 +1,4 @@
-import { zodSchemaToCoSchema } from "jazz-tools";
+import { anySchemaToCoSchema } from "jazz-tools";
 import { useAccount, useCoState } from "jazz-tools/react";
 import { useEffect, useState } from "react";
 import { createCredentiallessIframe } from "../../lib/createCredentiallessIframe";
@@ -44,7 +44,7 @@ export function UploaderPeer() {
     // The downloader peer will set the syncCompleted to true when the download is complete.
     // We use this to measure the sync duration.
     await waitForCoValue(
-      zodSchemaToCoSchema(UploadedFile),
+      anySchemaToCoSchema(UploadedFile),
       file.id,
       (value) => value.syncCompleted,
       { loadAs: account.me },