feat: zod catchall

packages/cojson/src/coValues/coMap.ts

@@ -52,6 +52,8 @@ export class RawCoMapView<
   };
   /** @internal */
   knownTransactions: CoValueKnownState["sessions"];
+  /** @internal */
+  totalProcessedTransactions: number;
 
   /** @internal */
   ignorePrivateTransactions: boolean;
@@ -75,6 +77,7 @@ export class RawCoMapView<
     this.ops = {};
     this.latest = {};
     this.knownTransactions = {};
+    this.totalProcessedTransactions = 0;
 
     this.processNewTransactions();
   }
@@ -136,6 +139,8 @@ export class RawCoMapView<
     for (const [key, entries] of changedEntries.entries()) {
       this.latest[key] = entries[entries.length - 1];
     }
+
+    this.totalProcessedTransactions += nextValidTransactions.length;
   }
 
   isTimeTravelEntity() {

packages/jazz-react-core/src/index.ts

@@ -1,3 +1,4 @@
 export * from "./hooks.js";
 export * from "./provider.js";
 export * from "./auth/index.js";
+export * from "./useCoState2.js";

packages/jazz-react-core/src/tests/useCoState2.test.ts

@@ -0,0 +1,349 @@
+// @vitest-environment happy-dom
+
+import { cojsonInternals } from "cojson";
+import { CoValue, Group, ID, SchemaV2 } from "jazz-tools";
+import { beforeEach, describe, expect, expectTypeOf, it } from "vitest";
+import { createJazzTestAccount, setupJazzTestSync } from "../testing.js";
+import { useCoState2 } from "../useCoState2.js";
+import { act, renderHook, waitFor } from "./testUtils.js";
+
+beforeEach(async () => {
+  await setupJazzTestSync();
+
+  await createJazzTestAccount({
+    isCurrentActiveAccount: true,
+  });
+});
+
+beforeEach(() => {
+  cojsonInternals.CO_VALUE_LOADING_CONFIG.MAX_RETRIES = 1;
+  cojsonInternals.CO_VALUE_LOADING_CONFIG.TIMEOUT = 1;
+});
+
+const { co, z } = SchemaV2;
+
+describe("useCoState2", () => {
+  it("should return the correct value", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+    });
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const map = TestMap.create({
+      value: "123",
+    });
+
+    const { result } = renderHook(
+      () => useCoState2(TestMap, map.$jazz.id, {}),
+      {
+        account,
+      },
+    );
+
+    expect(result.current?.value).toBe("123");
+  });
+
+  it("should update the value when the coValue changes", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+    });
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const map = TestMap.create({
+      value: "123",
+    });
+
+    const { result } = renderHook(
+      () => useCoState2(TestMap, map.$jazz.id, {}),
+      {
+        account,
+      },
+    );
+
+    expect(result.current?.value).toBe("123");
+
+    act(() => {
+      map.$jazz.set("value", "456");
+    });
+
+    expect(result.current?.value).toBe("456");
+  });
+
+  it("should load nested values if requested", async () => {
+    const TestNestedMap = co.map({
+      value: z.string(),
+    });
+
+    const TestMap = co.map({
+      value: z.string(),
+      nested: TestNestedMap,
+    });
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const map = TestMap.create({
+      value: "123",
+      nested: TestNestedMap.create({
+        value: "456",
+      }),
+    });
+
+    const { result } = renderHook(
+      () =>
+        useCoState2(TestMap, map.$jazz.id, {
+          resolve: {
+            nested: true,
+          },
+        }),
+      {
+        account,
+      },
+    );
+
+    expect(result.current?.value).toBe("123");
+    expect(result.current?.nested.value).toBe("456");
+  });
+
+  it("should load nested values when $requested", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+      nested: co.map({
+        value: z.string(),
+      }),
+    });
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const map = TestMap.create({
+      value: "123",
+      nested: {
+        value: "456",
+      },
+    });
+
+    expect(map.nested).toEqual({
+      value: "456",
+    });
+
+    const { result } = renderHook(
+      () => useCoState2(TestMap, map.$jazz.id, {}),
+      {
+        account,
+      },
+    );
+
+    result.current?.$jazz.request({
+      resolve: {
+        nested: true,
+      },
+    });
+
+    expect(result.current?.value).toBe("123");
+
+    await waitFor(() => {
+      expect(result.current?.nested).toEqual({
+        value: "456",
+      });
+    });
+  });
+
+  it.skip("should return null if the coValue is not found", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+    });
+
+    const map = TestMap.create({
+      value: "123",
+    });
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const { result } = renderHook(
+      () => useCoState2(TestMap, (map.$jazz.id + "123") as any),
+      {
+        account,
+      },
+    );
+
+    expect(result.current).toBeUndefined();
+
+    await waitFor(() => {
+      expect(result.current).toBeNull();
+    });
+  });
+
+  it.skip("should return null if the coValue is not accessible", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+    });
+
+    const someoneElse = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const map = TestMap.create(
+      {
+        value: "123",
+      },
+      someoneElse,
+    );
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const { result } = renderHook(() => useCoState2(TestMap, map.$jazz.id), {
+      account,
+    });
+
+    expect(result.current).toBeUndefined();
+
+    await waitFor(() => {
+      expect(result.current).toBeNull();
+    });
+  });
+
+  it("should not return null if the coValue is shared with everyone", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+    });
+
+    const someoneElse = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const group = Group.create(someoneElse);
+    group.addMember("everyone", "reader");
+
+    const map = TestMap.create(
+      {
+        value: "123",
+      },
+      group,
+    );
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const { result } = renderHook(() => useCoState2(TestMap, map.$jazz.id), {
+      account,
+    });
+
+    expect(result.current).toBeUndefined();
+
+    await waitFor(() => {
+      expect(result.current?.value).toBe("123");
+    });
+  });
+
+  it.skip("should return a value when the coValue becomes accessible", async () => {
+    const TestMap = co.map({
+      value: z.string(),
+    });
+
+    const someoneElse = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const group = Group.create(someoneElse);
+
+    const map = TestMap.create(
+      {
+        value: "123",
+      },
+      group,
+    );
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const { result } = renderHook(() => useCoState2(TestMap, map.$jazz.id), {
+      account,
+    });
+
+    expect(result.current).toBeUndefined();
+
+    await waitFor(() => {
+      expect(result.current).toBeNull();
+    });
+
+    group.addMember("everyone", "reader");
+
+    await waitFor(() => {
+      expect(result.current).not.toBeNull();
+      expect(result.current?.value).toBe("123");
+    });
+  });
+
+  it.skip("should update when an inner coValue is updated", async () => {
+    const TestNestedMap = co.map({
+      value: z.string(),
+    });
+
+    const TestMap = co.map({
+      value: z.string(),
+      nested: TestNestedMap,
+    });
+
+    const someoneElse = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const everyone = Group.create(someoneElse);
+    everyone.addMember("everyone", "reader");
+    const group = Group.create(someoneElse);
+
+    const map = TestMap.create(
+      {
+        value: "123",
+        nested: {
+          value: "456",
+        },
+      },
+      everyone,
+    );
+
+    const account = await createJazzTestAccount({
+      isCurrentActiveAccount: true,
+    });
+
+    const { result } = renderHook(
+      () =>
+        useCoState2(TestMap, map.$jazz.id, {
+          resolve: {
+            nested: true,
+          },
+        }),
+      {
+        account,
+      },
+    );
+
+    expect(result.current).toBeUndefined();
+
+    await waitFor(() => {
+      expect(result.current).not.toBeUndefined();
+    });
+
+    expect(result.current?.nested).toBeUndefined();
+    group.addMember("everyone", "reader");
+
+    await waitFor(() => {
+      expect(result.current?.nested?.value).toBe("456");
+    });
+  });
+});

packages/jazz-react-core/src/useCoState2.ts

@@ -0,0 +1,139 @@
+import { AnonymousJazzAgent, ID, SchemaV2 } from "jazz-tools";
+import type { Account } from "jazz-tools";
+import type {
+  CoMapSchema,
+  Loaded,
+  RelationsToResolve,
+  RelationsToResolveStrict,
+} from "jazz-tools/dist/schema/schema.js";
+import { useCallback, useRef, useState, useSyncExternalStore } from "react";
+import { useJazzContextManager } from "./hooks.js";
+import {
+  getCurrentAccountFromContextManager,
+  subscribeToContextManager,
+} from "./utils.js";
+
+export function createCoValueObservable<
+  V extends CoMapSchema<any>,
+  const R extends RelationsToResolve<V>,
+>() {
+  let currentValue: Loaded<V, R> | undefined | null = undefined;
+  let subscriberCount = 0;
+
+  function subscribe(
+    cls: V,
+    id: ID<V>,
+    options: {
+      loadAs: Account | AnonymousJazzAgent;
+      resolve?: RelationsToResolveStrict<V, R>;
+      onUnavailable?: () => void;
+      onUnauthorized?: () => void;
+      syncResolution?: boolean;
+    },
+    listener: () => void,
+  ) {
+    subscriberCount++;
+
+    const unsubscribe = SchemaV2.subscribeToCoValue(
+      cls,
+      id,
+      {
+        loadAs: options.loadAs,
+        resolve: options.resolve,
+        onUnavailable: () => {
+          currentValue = null;
+          options.onUnavailable?.();
+        },
+        onUnauthorized: () => {
+          currentValue = null;
+          options.onUnauthorized?.();
+        },
+      },
+      (value) => {
+        currentValue = value;
+        listener();
+      },
+    );
+
+    return () => {
+      unsubscribe();
+      subscriberCount--;
+      if (subscriberCount === 0) {
+        currentValue = undefined;
+      }
+    };
+  }
+
+  const observable = {
+    getCurrentValue: () => currentValue,
+    subscribe,
+  };
+
+  return observable;
+}
+
+function useCoValueObservable<
+  V extends CoMapSchema<any>,
+  const R extends RelationsToResolve<V>,
+>() {
+  const [initialValue] = useState(() => createCoValueObservable<V, R>());
+  const ref = useRef(initialValue);
+
+  return {
+    getCurrentValue() {
+      return ref.current.getCurrentValue();
+    },
+    getCurrentObservable() {
+      return ref.current;
+    },
+    reset() {
+      ref.current = createCoValueObservable<V, R>();
+    },
+  };
+}
+
+export function useCoState2<
+  V extends CoMapSchema<any>,
+  const R extends RelationsToResolve<V> = true,
+>(
+  Schema: V,
+  id: ID<V> | undefined,
+  options?: { resolve?: RelationsToResolveStrict<V, R> },
+): Loaded<V, R> | undefined | null {
+  const contextManager = useJazzContextManager();
+
+  const observable = useCoValueObservable<V, R>();
+
+  const value = useSyncExternalStore<Loaded<V, R> | undefined | null>(
+    useCallback(
+      (callback) => {
+        if (!id) return () => {};
+
+        // We subscribe to the context manager to react to the account updates
+        // faster than the useSyncExternalStore callback update to keep the isAuthenticated state
+        // up to date with the data when logging in and out.
+        return subscribeToContextManager(contextManager, () => {
+          const agent = getCurrentAccountFromContextManager(contextManager);
+          observable.reset();
+
+          return observable.getCurrentObservable().subscribe(
+            Schema,
+            id,
+            {
+              loadAs: agent,
+              resolve: options?.resolve,
+              onUnauthorized: callback,
+              onUnavailable: callback,
+            },
+            callback,
+          );
+        });
+      },
+      [Schema, id, contextManager],
+    ),
+    () => observable.getCurrentValue(),
+    () => observable.getCurrentValue(),
+  );
+
+  return value;
+}

packages/jazz-react/src/hooks.tsx

@@ -58,4 +58,5 @@ export {
   useJazzContext,
   useAccount,
   useAccountOrGuest,
+  useCoState2,
 } from "jazz-react-core";

packages/jazz-react/src/index.ts

@@ -8,6 +8,7 @@ export {
   experimental_useInboxSender,
   useJazzContext,
   useAuthSecretStorage,
+  useCoState2,
 } from "./hooks.js";
 
 export { createInviteLink, parseInviteLink } from "jazz-browser";

packages/jazz-tools/package.json

@@ -21,7 +21,8 @@
   "dependencies": {
     "@scure/bip39": "^1.3.0",
     "cojson": "workspace:*",
-    "fast-myers-diff": "^3.2.0"
+    "fast-myers-diff": "^3.2.0",
+    "zod": "^3.24.2"
   },
   "scripts": {
     "format-and-lint": "biome check .",

packages/jazz-tools/src/index.ts

@@ -1,3 +1,4 @@
 export * from "./exports.js";
 
 export { MAX_RECOMMENDED_TX_SIZE, cojsonInternals } from "cojson";
+export { SchemaV2 } from "./schema/index.js";

packages/jazz-tools/src/lib/cache.ts

@@ -13,4 +13,7 @@ export const coValuesCache = {
     weakMap.set(raw, computed);
     return computed;
   },
+  set: (raw: RawCoValue, value: CoValue) => {
+    weakMap.set(raw, value);
+  },
 };

packages/jazz-tools/src/schema/coMap/instance.ts

@@ -0,0 +1,392 @@
+import { CoValueUniqueness, JsonValue, RawAccount, RawCoMap } from "cojson";
+import { ZodType, ZodTypeAny } from "zod";
+import type { Account } from "../../coValues/account.js";
+import type { Group } from "../../coValues/group.js";
+import { RegisteredSchemas } from "../../coValues/registeredSchemas.js";
+import { AnonymousJazzAgent, ID } from "../../internal.js";
+import { coValuesCache } from "../../lib/cache.js";
+import { isOptional } from "../coValue/optional.js";
+import { SelfReference, isSelfReference } from "../coValue/self.js";
+import {
+  Loaded,
+  RelationsToResolve,
+  RelationsToResolveStrict,
+} from "../coValue/types.js";
+import { CoValueResolutionNode, ensureCoValueLoaded } from "../subscribe.js";
+import {
+  CoMapInit,
+  CoMapSchema,
+  CoMapSchemaKey,
+  CoValueSchema,
+} from "./schema.js";
+
+type Relations<D extends CoValueSchema<any>> = D extends CoMapSchema<infer S>
+  ? {
+      [K in keyof S]: S[K] extends CoMapSchema<any>
+        ? S[K]
+        : S[K] extends SelfReference
+          ? D
+          : never;
+    }
+  : never;
+
+type RelationsKeys<D extends CoValueSchema<any>> = keyof Relations<D> &
+  (string | number);
+
+type ChildMap<D extends CoMapSchema<any>> = Map<
+  RelationsKeys<D>,
+  Loaded<any, any> | undefined
+>;
+
+type PropertyType<
+  D extends CoMapSchema<any>,
+  K extends CoMapSchemaKey<D>,
+> = CoMapInit<D>[K];
+
+export type CoMap<
+  D extends CoMapSchema<any>,
+  R extends RelationsToResolve<D> = true,
+> = {
+  $jazz: CoMapJazzApi<D, R>;
+};
+
+export class CoMapJazzApi<
+  D extends CoMapSchema<any>,
+  R extends RelationsToResolve<D> = true,
+> {
+  raw: RawCoMap;
+  schema: D;
+  id: ID<D>;
+  _resolutionNode: CoValueResolutionNode<D, R> | undefined;
+  refs: ChildMap<D> = new Map();
+  protected lastUpdateTx: number;
+  declare _instance: Loaded<D, R>;
+
+  constructor(
+    schema: D,
+    raw: RawCoMap,
+    resolutionNode?: CoValueResolutionNode<D, R>,
+  ) {
+    this.schema = schema;
+    this.raw = raw;
+    this.lastUpdateTx = raw.totalProcessedTransactions;
+    this.id = raw.id as unknown as ID<D>;
+    this._resolutionNode = resolutionNode;
+  }
+
+  _setInstance(instance: CoMap<D, R>) {
+    this._instance = instance as unknown as Loaded<D, R>;
+  }
+
+  _fillRef<K extends RelationsKeys<D>>(key: K, value: Loaded<any, any>) {
+    const descriptor = this.schema.get(key);
+
+    if (descriptor && isRelationRef(descriptor)) {
+      this.refs.set(key, value);
+      Object.defineProperty(this._instance, key, {
+        value,
+        writable: false,
+        enumerable: true,
+        configurable: true,
+      });
+    } else {
+      throw new Error(`Field ${key} is not a reference`);
+    }
+  }
+
+  set<K extends CoMapSchemaKey<D>>(key: K, value: PropertyType<D, K>) {
+    const descriptor = this.schema.get(key);
+
+    if (descriptor && isRelationRef(descriptor)) {
+      if (!value) {
+        this.refs.delete(key as RelationsKeys<D>);
+      } else {
+        if (!isCoValue(value)) {
+          // To support inline CoMap creation on set
+          value = getSchemaFromDescriptor(this.schema, key).create(
+            value as CoMapInit<any>,
+            this.owner,
+          ) as PropertyType<D, K>;
+        }
+
+        this.refs.set(key as RelationsKeys<D>, value as Loaded<any, any>);
+      }
+    }
+
+    setValue(this.raw, this.schema, key, value as JsonValue);
+
+    return this.updated();
+  }
+
+  updated(refs?: ChildMap<D>): Loaded<D, R> {
+    if (this.lastUpdateTx === this.raw.totalProcessedTransactions && !refs) {
+      return this._instance;
+    }
+
+    return createCoMapFromRaw<D, R>(
+      this.schema,
+      this.raw,
+      refs ?? this.refs,
+      this._resolutionNode,
+    );
+  }
+
+  /**
+   * Given an already loaded `CoMapSchema`, ensure that the specified fields are loaded to the specified depth.
+   *
+   * Works like `CoMapSchema.load()`, but you don't need to pass the ID or the account to load as again.
+   *
+   * @category Subscription & Loading
+   */
+  ensureLoaded<O extends RelationsToResolve<D>>(options: {
+    resolve: RelationsToResolveStrict<D, O>;
+  }): Promise<Loaded<D, O>> {
+    return ensureCoValueLoaded<D, R, O>(this._instance, {
+      resolve: options.resolve,
+    });
+  }
+
+  /**
+   * Wait for the `CoMapSchema` to be uploaded to the other peers.
+   *
+   * @category Subscription & Loading
+   */
+  waitForSync(options?: { timeout?: number }) {
+    return this.raw.core.waitForSync(options);
+  }
+
+  get _loadedAs(): Account | AnonymousJazzAgent {
+    const rawAccount = this.raw.core.node.account;
+
+    if (rawAccount instanceof RawAccount) {
+      return coValuesCache.get(rawAccount, () =>
+        RegisteredSchemas["Account"].fromRaw(rawAccount),
+      );
+    }
+
+    return new AnonymousJazzAgent(this.raw.core.node);
+  }
+
+  get owner(): Account | Group {
+    return coValuesCache.get(this.raw.group, () =>
+      this.raw.group instanceof RawAccount
+        ? RegisteredSchemas["Account"].fromRaw(this.raw.group)
+        : RegisteredSchemas["Group"].fromRaw(this.raw.group),
+    );
+  }
+}
+
+export function createCoMap<D extends CoMapSchema<any>>(
+  schema: D,
+  init: CoMapInit<D>,
+  owner: Account | Group,
+  uniqueness?: CoValueUniqueness,
+) {
+  const { raw, refs } = createCoMapFromInit(init, owner, schema, uniqueness);
+
+  return createCoMapFromRaw<D, true>(schema, raw, refs);
+}
+
+export function createCoMapFromRaw<
+  D extends CoMapSchema<any>,
+  R extends RelationsToResolve<D>,
+>(
+  schema: D,
+  raw: RawCoMap,
+  refs?: ChildMap<D>,
+  resolutionNode?: CoValueResolutionNode<D, R>,
+) {
+  const instance = Object.create({
+    $jazz: new CoMapJazzApi(schema, raw, resolutionNode),
+  }) as CoMap<D, R>;
+  instance.$jazz._setInstance(instance);
+
+  const isRecord = false;
+  const fields = isRecord ? raw.keys() : schema.keys();
+
+  for (const key of fields) {
+    Object.defineProperty(instance, key, {
+      value: getValue(raw, schema, key as CoMapSchemaKey<D>),
+      writable: false,
+      enumerable: true,
+      configurable: true,
+    });
+  }
+
+  if (refs) {
+    for (const [key, value] of refs.entries()) {
+      if (value) {
+        instance.$jazz._fillRef(key as any, value);
+      }
+    }
+  }
+
+  return instance as unknown as Loaded<D, R>;
+}
+
+function getValue<D extends CoMapSchema<any>>(
+  raw: RawCoMap,
+  schema: D,
+  key: CoMapSchemaKey<D>,
+) {
+  const descriptor = schema.get(key);
+
+  if (descriptor && typeof key === "string") {
+    const value = raw.get(key);
+
+    if (descriptor instanceof CoMapSchema || isSelfReference(descriptor)) {
+      if (value === undefined) {
+        return undefined;
+      } else {
+        return null;
+      }
+    } else {
+      try {
+        return descriptor.parse(value);
+      } catch (error) {
+        throw new Error(
+          `Failed to parse field ${key}: ${JSON.stringify(error)}`,
+        );
+      }
+    }
+  } else {
+    return undefined;
+  }
+}
+
+function setValue<D extends CoMapSchema<any>>(
+  raw: RawCoMap,
+  schema: D,
+  key: CoMapSchemaKey<D>,
+  value: JsonValue,
+) {
+  const descriptor = schema.get(key);
+
+  if (descriptor && typeof key === "string") {
+    if (isRelationRef(descriptor)) {
+      if (value === null || value === undefined) {
+        if (isOptional(descriptor)) {
+          raw.set(key, undefined);
+        } else {
+          throw new Error(`Field ${key} is required`);
+        }
+      } else {
+        if (value && typeof value === "object" && "$jazz" in value) {
+          raw.set(
+            key,
+            (value as unknown as Loaded<CoMapSchema<{}>, true>).$jazz.id,
+          );
+        } else {
+          throw new Error(`The value assigned to ${key} is not a reference`);
+        }
+      }
+    } else {
+      // TODO: Provide better parse errors with the field information
+      try {
+        raw.set(key, descriptor.parse(value));
+      } catch (error) {
+        throw new Error(
+          `Failed to parse field ${key}: ${JSON.stringify(error)}`,
+        );
+      }
+    }
+
+    return true;
+  }
+}
+
+function createCoMapFromInit<D extends CoMapSchema<any>>(
+  init: CoMapInit<D> | undefined,
+  owner: Account | Group,
+  schema: D,
+  uniqueness?: CoValueUniqueness,
+) {
+  const rawOwner = owner._raw;
+
+  const rawInit = {} as {
+    [key in CoMapSchemaKey<D>]: JsonValue | undefined;
+  };
+
+  const refs = new Map<string, Loaded<any, any>>();
+
+  if (init) {
+    const fields = schema.keys() as (CoMapSchemaKey<D> & string)[];
+
+    for (const key of fields) {
+      const initValue = init[key] as
+        | Loaded<CoValueSchema<{}>>
+        | CoMapInit<any>
+        | undefined
+        | null;
+
+      const descriptor = schema.get(key);
+
+      if (!descriptor) {
+        continue;
+      }
+
+      if (isRelationRef(descriptor)) {
+        if (initValue === null || initValue === undefined) {
+          if (isOptional(descriptor)) {
+            rawInit[key] = undefined;
+          } else {
+            throw new Error(`Field ${key} is required`);
+          }
+        } else {
+          let instance: Loaded<CoValueSchema<{}>>;
+
+          if ("$jazz" in initValue) {
+            instance = initValue as Loaded<CoValueSchema<{}>>;
+          } else {
+            instance = getSchemaFromDescriptor(schema, key).create(
+              initValue,
+              owner,
+            ) as Loaded<CoValueSchema<{}>>;
+          }
+
+          rawInit[key] = instance.$jazz.id;
+          refs.set(key as string, instance);
+        }
+      } else {
+        try {
+          rawInit[key] = descriptor.parse(initValue);
+        } catch (error) {
+          throw new Error(
+            `Failed to parse field ${key}: ${JSON.stringify(error)}`,
+          );
+        }
+      }
+    }
+  }
+
+  const raw = rawOwner.createMap(rawInit, null, "private", uniqueness);
+
+  return { raw, refs };
+}
+
+function getSchemaFromDescriptor<
+  S extends CoMapSchema<any>,
+  K extends CoMapSchemaKey<S>,
+>(schema: S, key: K) {
+  const descriptor = schema.get(key);
+
+  if (descriptor && isRelationRef(descriptor)) {
+    if (isSelfReference(descriptor)) {
+      return schema;
+    } else {
+      return descriptor;
+    }
+  } else {
+    throw new Error(`Field ${String(key)} is not a reference`);
+  }
+}
+
+export function isRelationRef(
+  descriptor: CoMapSchema<any> | ZodTypeAny | SelfReference,
+): descriptor is CoMapSchema<any> | SelfReference {
+  return descriptor instanceof CoMapSchema || isSelfReference(descriptor);
+}
+
+export function isCoValue(value: unknown): value is CoMap<any, any> {
+  return typeof value === "object" && value !== null && "$jazz" in value;
+}

packages/jazz-tools/src/schema/coMap/record.ts

@@ -0,0 +1,18 @@
+import { ZodString } from "zod";
+import { CoMapSchema, CoMapSchemaField } from "./schema.js";
+
+export const RecordSymbol = "$$record$$";
+export type RecordSymbol = typeof RecordSymbol;
+
+export type RecordDefinition<
+  K extends ZodString,
+  V extends CoMapSchemaField,
+> = {
+  key: K;
+  value: V;
+};
+
+export type IsRecord<T extends CoMapSchema<any>> =
+  T["shape"][typeof RecordSymbol] extends RecordDefinition<any, any>
+    ? true
+    : false;

packages/jazz-tools/src/schema/coMap/schema.ts

@@ -0,0 +1,191 @@
+import { CoValueUniqueness } from "cojson";
+import { TypeOf, ZodString, ZodTypeAny, z } from "zod";
+import { Account } from "../../coValues/account.js";
+import { Group } from "../../coValues/group.js";
+import { parseCoValueCreateOptions } from "../../internal.js";
+import { addOptional, carryOptional, optional } from "../coValue/optional.js";
+import { SelfReference, markSelfReferenceAsOptional } from "../coValue/self.js";
+import {
+  IsDepthLimit,
+  addQuestionMarks,
+  flatten,
+} from "../coValue/typeUtils.js";
+import { Loaded, LoadedCoMap, ValidateResolve } from "../coValue/types.js";
+import { createCoMap } from "./instance.js";
+import { RecordDefinition, RecordSymbol } from "./record.js";
+
+export type CoMapSchemaField = CoMapSchema<any> | ZodTypeAny | SelfReference;
+
+export type CoMapSchemaShape = {
+  [RecordSymbol]?: RecordDefinition<any, CoMapSchemaField>;
+} & {
+  [key: string]: CoMapSchemaField;
+};
+
+export type RecordKeyType<S extends CoMapSchema<any>> =
+  S["shape"][RecordSymbol] extends RecordDefinition<any, any>
+    ? TypeOf<S["shape"][RecordSymbol]["key"]>
+    : never;
+
+export type RecordValueType<S extends CoMapSchema<any>> =
+  S["shape"][RecordSymbol] extends RecordDefinition<any, any>
+    ? S["shape"][RecordSymbol]["value"]
+    : never;
+
+export type CoMapSchemaKey<S extends CoMapSchema<any>> = Exclude<
+  keyof S["shape"],
+  RecordSymbol
+>;
+
+export type CoMapFieldType<
+  S extends CoMapSchema<any>,
+  K extends CoMapSchemaKey<S>,
+> = S["shape"][K] extends CoMapSchemaField ? S["shape"][K] : never;
+
+export type CoValueSchema<S extends CoMapSchemaShape> = CoMapSchema<S>;
+
+export type UnwrapReference<
+  D extends CoMapSchema<any>,
+  K extends CoMapSchemaKey<D>,
+> = CoMapFieldType<D, K> extends CoValueSchema<any>
+  ? CoMapFieldType<D, K>
+  : CoMapFieldType<D, K> extends SelfReference
+    ? D
+    : never;
+
+export type CoMapInit<
+  D extends CoMapSchema<any>,
+  CurrentDepth extends number[] = [],
+> = IsDepthLimit<CurrentDepth> extends true
+  ? {}
+  : addQuestionMarks<{
+      [K in CoMapSchemaKey<D>]: CoMapFieldType<D, K> extends ZodTypeAny
+        ? TypeOf<CoMapFieldType<D, K>>
+        : UnwrapReference<D, K> extends CoMapSchema<any>
+          ?
+              | CoMapInit<UnwrapReference<D, K>, [0, ...CurrentDepth]>
+              | LoadedCoMap<UnwrapReference<D, K>, any>
+              | addOptional<UnwrapReference<D, K>>
+              | markSelfReferenceAsOptional<CoMapFieldType<D, K>> // Self references are always optional
+          : never;
+    }>;
+
+export type CoMapInitStrict<
+  D extends CoMapSchema<any>,
+  I,
+> = I extends CoMapInit<D> ? CoMapInit<D> : I;
+
+/**
+ * This is a simplified version of CoMapInit that only includes the keys that are defined in the schema.
+ * It is used to build the resolve type for the create method without paying the compelxity cost of the full CoMapInit.
+ */
+type CoMapSimpleInit<
+  D extends CoMapSchema<any>,
+  CurrentDepth extends number[] = [],
+> = IsDepthLimit<CurrentDepth> extends true
+  ? {}
+  : {
+      [K in CoMapSchemaKey<D>]?: CoMapFieldType<D, K> extends ZodTypeAny
+        ? TypeOf<CoMapFieldType<D, K>>
+        : any;
+    };
+
+export type CoMapInitToRelationsToResolve<
+  S extends CoMapSchemaShape,
+  I extends CoMapSimpleInit<CoMapSchema<S>>,
+  CurrentDepth extends number[] = [],
+> = IsDepthLimit<CurrentDepth> extends true
+  ? true
+  : ValidateResolve<
+      CoMapSchema<S>,
+      {
+        [K in CoMapSchemaKey<CoMapSchema<S>>]: UnwrapReference<
+          CoMapSchema<S>,
+          K
+        > extends CoMapSchema<infer ChildSchema>
+          ? I[K] extends LoadedCoMap<CoMapSchema<ChildSchema>, infer R>
+            ? R
+            : I[K] extends CoMapSimpleInit<CoMapSchema<ChildSchema>>
+              ? CoMapInitToRelationsToResolve<
+                  ChildSchema,
+                  I[K],
+                  [0, ...CurrentDepth]
+                >
+              : never
+          : never;
+      },
+      true
+    >;
+
+export class CoMapSchema<S extends CoMapSchemaShape> {
+  shape: S;
+
+  constructor(schema: S) {
+    this.shape = schema;
+  }
+
+  optional() {
+    return optional(this);
+  }
+
+  get(key: CoMapSchemaKey<CoMapSchema<S>>) {
+    const descriptor = this.shape[key];
+
+    if (descriptor) {
+      return descriptor;
+    }
+
+    return this.getRecordDescriptor(key);
+  }
+
+  keys() {
+    return Object.keys(this.shape) as (keyof S & string)[];
+  }
+
+  getRecordDescriptor(key: CoMapSchemaKey<CoMapSchema<S>>) {
+    if (!this.shape[RecordSymbol]) {
+      return undefined;
+    }
+
+    const { key: keySchema, value } = this.shape[RecordSymbol];
+
+    if (!keySchema.safeParse(key).success) {
+      return undefined;
+    }
+
+    return value;
+  }
+
+  isRecord() {
+    return this.shape[RecordSymbol] !== undefined;
+  }
+
+  catchall<V extends CoMapSchemaField>(descriptor: V) {
+    const newShape = {
+      [RecordSymbol]: { key: z.string(), value: descriptor },
+      ...this.shape,
+    };
+
+    const newSchema = new CoMapSchema(newShape);
+
+    return carryOptional(this, newSchema);
+  }
+
+  create<I extends CoMapInit<CoMapSchema<S>>>(
+    init: CoMapInitStrict<CoMapSchema<S>, I>,
+    options?:
+      | {
+          owner: Account | Group;
+          unique?: CoValueUniqueness["uniqueness"];
+        }
+      | Account
+      | Group,
+  ): Loaded<
+    CoMapSchema<S>,
+    CoMapInitToRelationsToResolve<S, I>,
+    "non-nullable" // We want the loaded type to reflect the init input as we know for sure if values are available or not
+  > {
+    const { owner, uniqueness } = parseCoValueCreateOptions(options);
+    return createCoMap<CoMapSchema<S>>(this, init, owner, uniqueness) as any;
+  }
+}

packages/jazz-tools/src/schema/coMapWithZod.load.test.ts

@@ -0,0 +1,814 @@
+import {
+  assert,
+  beforeEach,
+  describe,
+  expect,
+  expectTypeOf,
+  it,
+  vi,
+} from "vitest";
+import { Group } from "../exports.js";
+import { createJazzTestAccount, setupJazzTestSync } from "../testing.js";
+import { waitFor } from "../tests/utils.js";
+import { Loaded, co, z } from "./schema.js";
+import { loadCoValue, subscribeToCoValue } from "./subscribe.js";
+
+beforeEach(async () => {
+  await setupJazzTestSync();
+
+  await createJazzTestAccount({
+    isCurrentActiveAccount: true,
+  });
+});
+
+describe("CoMap with Zod", () => {
+  describe("load", () => {
+    it("should load a CoMap without nested values", async () => {
+      const anotherAccount = await createJazzTestAccount();
+
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const group = Group.create(anotherAccount);
+      group.addMember("everyone", "reader");
+
+      const john = Person.create(
+        {
+          name: "John",
+          age: 30,
+          address: { street: "123 Main St" },
+        },
+        group,
+      );
+
+      const loaded = await loadCoValue(Person, john.$jazz.id, {
+        resolve: true,
+      });
+
+      assert(loaded);
+
+      expect(loaded.name).toBe("John");
+      expect(loaded.age).toBe(30);
+      expect(loaded.address).toBe(null);
+
+      expectTypeOf(loaded.address).toEqualTypeOf<null>();
+    });
+
+    it("should load a CoMap with nested values", async () => {
+      const anotherAccount = await createJazzTestAccount();
+
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const group = Group.create(anotherAccount);
+      group.addMember("everyone", "reader");
+
+      const john = Person.create(
+        {
+          name: "John",
+          age: 30,
+          address: { street: "123 Main St" },
+        },
+        group,
+      );
+
+      const loaded = await loadCoValue(Person, john.$jazz.id, {
+        resolve: { address: true },
+      });
+
+      assert(loaded);
+
+      expect(loaded.name).toBe("John");
+      expect(loaded.age).toBe(30);
+      expect(loaded.address).toEqual({ street: "123 Main St" });
+
+      expectTypeOf(loaded.address).toEqualTypeOf<
+        Loaded<typeof Person.shape.address, true>
+      >();
+    });
+
+    it("should load a CoMap with self references", async () => {
+      const anotherAccount = await createJazzTestAccount();
+
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        friend: co.self(),
+      });
+
+      const group = Group.create(anotherAccount);
+      group.addMember("everyone", "reader");
+
+      const john = Person.create(
+        {
+          name: "John",
+          age: 30,
+          friend: {
+            name: "Jane",
+            age: 20,
+          },
+        },
+        group,
+      );
+
+      const loaded = await loadCoValue(Person, john.$jazz.id, {
+        resolve: { friend: true },
+      });
+
+      assert(loaded);
+
+      expect(loaded.name).toBe("John");
+      expect(loaded.age).toBe(30);
+      expect(loaded.friend).toEqual({
+        name: "Jane",
+        age: 20,
+      });
+
+      expectTypeOf(loaded.friend).toEqualTypeOf<
+        Loaded<typeof Person, true> | undefined | null
+      >();
+    });
+
+    it("should load a CoMap with nested values if an optional nested value is missing", async () => {
+      const anotherAccount = await createJazzTestAccount();
+
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co
+          .map({
+            street: z.string(),
+          })
+          .optional(),
+      });
+
+      const group = Group.create(anotherAccount);
+      group.addMember("everyone", "reader");
+
+      const john = Person.create(
+        {
+          name: "John",
+          age: 30,
+        },
+        group,
+      );
+
+      const loaded = await loadCoValue(Person, john.$jazz.id, {
+        resolve: { address: true },
+      });
+
+      assert(loaded);
+
+      expect(loaded.name).toBe("John");
+      expect(loaded.age).toBe(30);
+      expect(loaded.address).toBeUndefined();
+
+      expectTypeOf(loaded.address).toEqualTypeOf<
+        Loaded<typeof Person.shape.address, true> | undefined | null
+      >();
+    });
+
+    it.todo(
+      "should return undefined if the value is not available",
+      async () => {
+        const Person = co.map({
+          name: z.string(),
+          age: z.number(),
+          address: co.map({
+            street: z.string(),
+          }),
+        });
+
+        const john = Person.create({
+          name: "John",
+          age: 30,
+          address: { street: "123 Main St" },
+        });
+
+        const loaded = await loadCoValue(Person, (john.$jazz.id + "1") as any, {
+          resolve: { address: true },
+        });
+
+        expect(loaded).toBeUndefined();
+      },
+    );
+
+    it.todo(
+      "should return undefined if one of the nested values is not available",
+      async () => {
+        const anotherAccount = await createJazzTestAccount();
+        const Person = co.map({
+          name: z.string(),
+          age: z.number(),
+          address: co.map({
+            street: z.string(),
+          }),
+        });
+
+        const group = Group.create(anotherAccount);
+        group.addMember("everyone", "reader");
+
+        const john = Person.create(
+          {
+            name: "John",
+            age: 30,
+            address: { street: "123 Main St" },
+          },
+          group,
+        );
+
+        john.$jazz.raw.set("address", "co_z1");
+
+        await john.$jazz.waitForSync();
+
+        const loaded = await loadCoValue(Person, john.$jazz.id, {
+          resolve: { address: true },
+        });
+
+        expect(loaded).toBeUndefined();
+      },
+    );
+
+    it.todo(
+      "should return undefined if the value is not accessible",
+      async () => {
+        const anotherAccount = await createJazzTestAccount();
+
+        const Person = co.map({
+          name: z.string(),
+          age: z.number(),
+          address: co.map({
+            street: z.string(),
+          }),
+        });
+
+        const john = Person.create(
+          {
+            name: "John",
+            age: 30,
+            address: { street: "123 Main St" },
+          },
+          anotherAccount,
+        );
+
+        await john.$jazz.waitForSync();
+
+        const loaded = await loadCoValue(Person, john.$jazz.id as any, {
+          resolve: { address: true },
+        });
+
+        expect(loaded).toBeUndefined();
+      },
+    );
+
+    it.todo(
+      "should return undefined if one of the nested values is not accessible",
+      async () => {
+        const anotherAccount = await createJazzTestAccount();
+        const Person = co.map({
+          name: z.string(),
+          age: z.number(),
+          address: co.map({
+            street: z.string(),
+          }),
+        });
+
+        const group = Group.create(anotherAccount);
+        group.addMember("everyone", "reader");
+
+        const john = Person.create(
+          {
+            name: "John",
+            age: 30,
+            address: Person.shape.address.create({ street: "123 Main St" }),
+          },
+          group,
+        );
+
+        await john.$jazz.waitForSync();
+
+        const loaded = await loadCoValue(Person, john.$jazz.id, {
+          resolve: { address: true },
+        });
+
+        expect(loaded).toBeUndefined();
+      },
+    );
+  });
+
+  describe("ensureLoaded", () => {
+    it("should load the nested values", async () => {
+      const anotherAccount = await createJazzTestAccount();
+
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const group = Group.create(anotherAccount);
+      group.addMember("everyone", "reader");
+
+      const john = Person.create(
+        {
+          name: "John",
+          age: 30,
+          address: { street: "123 Main St" },
+        },
+        group,
+      );
+
+      const loaded = await loadCoValue(Person, john.$jazz.id, {
+        resolve: true,
+      });
+
+      assert(loaded);
+
+      const { address } = await loaded.$jazz.ensureLoaded({
+        resolve: { address: true },
+      });
+
+      expect(address).toEqual(john.address);
+      expectTypeOf(address).toEqualTypeOf<
+        Loaded<typeof Person.shape.address, true>
+      >();
+    });
+
+    it.todo(
+      "should throw if one of the nested values is not available",
+      async () => {
+        const anotherAccount = await createJazzTestAccount();
+        const Person = co.map({
+          name: z.string(),
+          age: z.number(),
+          address: co.map({
+            street: z.string(),
+          }),
+        });
+
+        const group = Group.create(anotherAccount);
+        group.addMember("everyone", "reader");
+
+        const john = Person.create(
+          {
+            name: "John",
+            age: 30,
+            address: { street: "123 Main St" },
+          },
+          group,
+        );
+
+        john.$jazz.raw.set("address", "co_z1");
+
+        await john.$jazz.waitForSync();
+
+        const loaded = await loadCoValue(Person, john.$jazz.id, {
+          resolve: true,
+        });
+
+        assert(loaded);
+
+        await expect(
+          loaded.$jazz.ensureLoaded({
+            resolve: { address: true },
+          }),
+        ).rejects.toThrow();
+      },
+    );
+
+    it.todo(
+      "should throw if one of the nested values is not accessible",
+      async () => {
+        const anotherAccount = await createJazzTestAccount();
+        const Person = co.map({
+          name: z.string(),
+          age: z.number(),
+          address: co.map({
+            street: z.string(),
+          }),
+        });
+
+        const group = Group.create(anotherAccount);
+        group.addMember("everyone", "reader");
+
+        const john = Person.create(
+          {
+            name: "John",
+            age: 30,
+            address: Person.shape.address.create({ street: "123 Main St" }),
+          },
+          group,
+        );
+
+        await john.$jazz.waitForSync();
+
+        const loaded = await loadCoValue(Person, john.$jazz.id, {
+          resolve: true,
+        });
+
+        assert(loaded);
+
+        await expect(
+          loaded.$jazz.ensureLoaded({
+            resolve: { address: true },
+          }),
+        ).rejects.toThrow();
+      },
+    );
+  });
+
+  describe("subscribe", () => {
+    it("should syncronously load a locally available CoMap without nested values", async () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      let result: any;
+
+      subscribeToCoValue(Person, john.$jazz.id, { resolve: true }, (value) => {
+        result = value;
+      });
+
+      const johnAfterSubscribe = result as Loaded<typeof Person, true>;
+
+      expect(johnAfterSubscribe.name).toBe("John");
+      expect(johnAfterSubscribe.age).toBe(30);
+      expect(johnAfterSubscribe.address).toBe(null);
+    });
+
+    it("should syncronously load a locally available CoMap with nested values", async () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      let result: any;
+
+      subscribeToCoValue(
+        Person,
+        john.$jazz.id,
+        { resolve: { address: true } },
+        (value) => {
+          result = value;
+        },
+      );
+
+      expect(result).toEqual({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+    });
+  });
+
+  it("should emit on updates", async () => {
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co.map({
+        street: z.string(),
+      }),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+      address: { street: "123 Main St" },
+    });
+
+    let result: any;
+
+    subscribeToCoValue(
+      Person,
+      john.$jazz.id,
+      { resolve: { address: true } },
+      (value) => {
+        result = value;
+      },
+    );
+
+    const resultBeforeSet = result;
+
+    john.$jazz.set("name", "Jane");
+
+    expect(result).toEqual({
+      name: "Jane",
+      age: 30,
+      address: { street: "123 Main St" },
+    });
+
+    expect(resultBeforeSet).not.toBe(result);
+  });
+
+  it("should emit on nested values updates", async () => {
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co.map({
+        street: z.string(),
+      }),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+      address: { street: "123 Main St" },
+    });
+
+    let result: any;
+
+    subscribeToCoValue(
+      Person,
+      john.$jazz.id,
+      { resolve: { address: true } },
+      (value) => {
+        result = value;
+      },
+    );
+
+    const resultBeforeSet = result;
+
+    john.$jazz.set("address", { street: "456 Main St" });
+
+    expect(result).toEqual({
+      name: "John",
+      age: 30,
+      address: { street: "456 Main St" },
+    });
+
+    expect(resultBeforeSet).not.toBe(result);
+  });
+
+  it("should not emit updates if a stale nested value is updated", async () => {
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co.map({
+        street: z.string(),
+      }),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+      address: { street: "123 Main St" },
+    });
+
+    let result: any;
+
+    const spy = vi.fn();
+
+    subscribeToCoValue(
+      Person,
+      john.$jazz.id,
+      { resolve: { address: true } },
+      (value) => {
+        result = value;
+        spy(value);
+      },
+    );
+
+    john.$jazz.set("address", { street: "456 Main St" });
+
+    spy.mockReset();
+
+    john.address.$jazz.set("street", "updated");
+
+    expect(result).toEqual({
+      name: "John",
+      age: 30,
+      address: { street: "456 Main St" },
+    });
+
+    expect(spy).not.toHaveBeenCalled();
+  });
+
+  it("should emit when an optional nested value becomes missing", async () => {
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co
+        .map({
+          street: z.string(),
+        })
+        .optional(),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+      address: { street: "123 Main St" },
+    });
+
+    let result: any;
+
+    subscribeToCoValue(
+      Person,
+      john.$jazz.id,
+      { resolve: { address: true } },
+      (value) => {
+        result = value;
+      },
+    );
+
+    const resultBeforeSet = result;
+
+    john.$jazz.set("address", undefined);
+
+    expect(result).toEqual({
+      name: "John",
+      age: 30,
+    });
+
+    expect(resultBeforeSet).not.toBe(result);
+  });
+
+  it("should set the property as null if we have the value but it is not requested", async () => {
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co
+        .map({
+          street: z.string(),
+        })
+        .optional(),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+      address: { street: "123 Main St" },
+    });
+
+    let result: any;
+
+    subscribeToCoValue(Person, john.$jazz.id, { resolve: true }, (value) => {
+      result = value;
+    });
+
+    expect(result.address).toBe(null);
+  });
+
+  it("should set the property as undefined if we have the nested coValue is optional and missing", async () => {
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co
+        .map({
+          street: z.string(),
+        })
+        .optional(),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+    });
+
+    let result: any;
+
+    subscribeToCoValue(Person, john.$jazz.id, { resolve: true }, (value) => {
+      result = value;
+    });
+
+    expect(result.address).toBe(undefined);
+  });
+
+  it("should emit the update only when the nested coValue is loaded", async () => {
+    const anotherAccount = await createJazzTestAccount();
+
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co.map({
+        street: z.string(),
+      }),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+      address: { street: "123 Main St" },
+    });
+
+    let result: any;
+
+    subscribeToCoValue(
+      Person,
+      john.$jazz.id,
+      { resolve: { address: true } },
+      (value) => {
+        result = value;
+      },
+    );
+
+    const firstResult = result as Loaded<typeof Person, { address: true }>;
+
+    const group = Group.create(anotherAccount);
+    group.addMember("everyone", "reader");
+
+    const newAddress = Person.shape.address.create(
+      { street: "456 Main St" },
+      group,
+    );
+
+    john.$jazz.set("address", newAddress);
+
+    expect(result).toBe(firstResult);
+
+    await waitFor(() => {
+      expect(result).not.toBe(firstResult);
+    });
+
+    expect(result).toEqual({
+      name: "John",
+      age: 30,
+      address: { street: "456 Main St" },
+    });
+  });
+
+  it("should preserve references of unchanged nested values", async () => {
+    const Person = co.map({
+      name: z.string(),
+      age: z.number(),
+      address: co.map({
+        street: z.string(),
+      }),
+      friend: co.self(),
+    });
+
+    const john = Person.create({
+      name: "John",
+      age: 30,
+      address: { street: "123 Main St" },
+      friend: {
+        name: "Jane",
+        age: 20,
+        address: { street: "456 Main St" },
+      },
+    });
+
+    let result: any;
+
+    subscribeToCoValue(
+      Person,
+      john.$jazz.id,
+      { resolve: { address: true, friend: { address: true } } },
+      (value) => {
+        result = value;
+      },
+    );
+
+    expect(result).toEqual(john);
+
+    const resultBeforeSet = result;
+
+    john.address.$jazz.set("street", "456 Main St");
+
+    expect(result).toEqual({
+      name: "John",
+      age: 30,
+      address: { street: "456 Main St" },
+      friend: {
+        name: "Jane",
+        age: 20,
+        address: { street: "456 Main St" },
+      },
+    });
+
+    expect(resultBeforeSet).not.toBe(result);
+    expect(resultBeforeSet.friend).toBe(result.friend);
+    expect(resultBeforeSet.address).not.toBe(result.address);
+  });
+});

packages/jazz-tools/src/schema/coMapWithZod.test.ts

@@ -0,0 +1,414 @@
+import { beforeEach, describe, expect, expectTypeOf, it } from "vitest";
+import { createJazzTestAccount } from "../testing.js";
+import { CoMapInit } from "./coMap/schema.js";
+import { Loaded, co, z } from "./schema.js";
+
+beforeEach(async () => {
+  await createJazzTestAccount({
+    isCurrentActiveAccount: true,
+  });
+});
+
+describe("CoMap - with zod based schema", () => {
+  describe("init", () => {
+    it("should create a CoMap with basic property access", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      const john = Person.create({ name: "John", age: 30 });
+
+      expect(john.name).toBe("John");
+      expect(john.age).toBe(30);
+    });
+
+    it("should generate the right CoMapInit type", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      expectTypeOf<CoMapInit<typeof Person>>().toMatchTypeOf<{
+        name: string;
+        age: number;
+      }>();
+    });
+
+    it("should throw an error if a required field is missing", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      // @ts-expect-error - age is required
+      expect(() => Person.create({ name: "John" })).toThrow(
+        /^Failed to parse field age/,
+      );
+    });
+
+    it("should not throw an error if a optional field is missing", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number().optional(),
+      });
+
+      const john = Person.create({ name: "John" });
+
+      expect(john.age).toBeUndefined();
+      expect(john.name).toBe("John");
+    });
+
+    it("should create a CoMap with nested values", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      expect(john.name).toBe("John");
+      expect(john.age).toBe(30);
+      expect(john.address.street).toBe("123 Main St");
+      expect(john.address.$jazz.owner).toBe(john.$jazz.owner);
+    });
+
+    it("should retrurn a loaded type when a reference is passed on create", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co
+          .map({
+            street: z.string(),
+          })
+          .optional(),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      expectTypeOf<typeof john.address>().toMatchTypeOf<
+        Loaded<typeof Person.shape.address>
+      >();
+    });
+
+    it("should be possible to reference a nested map schema to split group creation", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: Person.shape.address.create({ street: "123 Main St" }),
+      });
+
+      expect(john.name).toBe("John");
+      expect(john.age).toBe(30);
+      expect(john.address.street).toBe("123 Main St");
+      expect(john.address.$jazz.owner).not.toBe(john.$jazz.owner);
+    });
+
+    it("should throw an error if a required ref is missing", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      // @ts-expect-error - address is required
+      expect(() => Person.create({ name: "John", age: 30 })).toThrow(
+        /^Field address is required/,
+      );
+    });
+
+    it("should not throw an error if a required ref is missing", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co
+          .map({
+            street: z.string(),
+          })
+          .optional(),
+      });
+
+      const john = Person.create({ name: "John", age: 30 });
+
+      expect(john.address).toBeUndefined();
+    });
+
+    it("should create a CoMap with self references", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        friend: co.self(),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        friend: {
+          name: "Jane",
+          age: 20,
+          friend: { name: "Bob", age: 20 },
+        },
+      });
+
+      expect(john.friend.friend.name).toBe("Bob");
+      expect(john.friend.name).toBe("Jane");
+    });
+
+    it("should return a loaded type when a self reference is passed on create", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        friend: co.self(),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        friend: {
+          name: "Jane",
+          age: 20,
+          friend: { name: "Bob", age: 20 },
+        },
+      });
+
+      expectTypeOf<typeof john.friend.friend.name>().toEqualTypeOf<string>();
+      expectTypeOf<typeof john.friend.friend.friend>().toEqualTypeOf<never>();
+    });
+
+    it("should not throw an error if a self reference is missing", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        friend: co.self(),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+      });
+
+      expect(john.friend).toBeUndefined();
+    });
+
+    it("should support extra properties with catchall", () => {
+      const Person = co
+        .map({
+          name: z.string(),
+          age: z.number(),
+        })
+        .catchall(z.string());
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        extra: "extra",
+      });
+
+      expect(john.extra).toBe("extra");
+    });
+  });
+
+  describe("json format", () => {
+    it("should properly serialize to JSON", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      const john = Person.create({ name: "John", age: 30 });
+
+      expect(JSON.stringify(john)).toMatchInlineSnapshot(
+        `"{"name":"John","age":30}"`,
+      );
+    });
+
+    it("should properly serialize nested values to JSON", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      expect(JSON.stringify(john)).toMatchInlineSnapshot(
+        `"{"name":"John","age":30,"address":{"street":"123 Main St"}}"`,
+      );
+    });
+
+    it("should work with Object.entries", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      const john = Person.create({ name: "John", age: 30 });
+
+      expect(Object.entries(john)).toEqual([
+        ["name", "John"],
+        ["age", 30],
+      ]);
+    });
+
+    it("should work on equality checks", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      expect(john).toEqual({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+    });
+  });
+
+  describe("updates", () => {
+    it("should not change original property after calling $jazz.set", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      const john = Person.create({ name: "John", age: 30 });
+
+      const jane = john.$jazz.set("name", "Jane");
+
+      expect(john.name).toBe("John");
+      expect(jane.name).toBe("Jane");
+    });
+
+    it("should update nested values", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      const johnAfterMoving = john.$jazz.set(
+        "address",
+        Person.shape.address.create({
+          street: "456 Main St",
+        }),
+      );
+
+      expect(john.address.street).toBe("123 Main St");
+      expect(johnAfterMoving.address.street).toBe("456 Main St");
+      expect(johnAfterMoving.address.$jazz.owner).not.toBe(john.$jazz.owner);
+      expect(john.$jazz.updated().address).toBe(johnAfterMoving.address);
+    });
+
+    it("should update nested values with JSON data", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        address: co.map({
+          street: z.string(),
+        }),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        address: { street: "123 Main St" },
+      });
+
+      const johnAfterMoving = john.$jazz.set("address", {
+        street: "456 Main St",
+      });
+
+      expect(john.address.street).toBe("123 Main St");
+      expect(johnAfterMoving.address.street).toBe("456 Main St");
+      expect(john.$jazz.updated().address).toBe(johnAfterMoving.address);
+      expect(johnAfterMoving.address.$jazz.owner).toBe(john.$jazz.owner);
+    });
+
+    it("should update nested values with self references", () => {
+      const Person = co.map({
+        name: z.string(),
+        age: z.number(),
+        friend: co.self(),
+      });
+
+      const john = Person.create({
+        name: "John",
+        age: 30,
+        friend: {
+          name: "Jane",
+          age: 20,
+        },
+      });
+
+      const johnWithANewFriend = john.$jazz.set("friend", {
+        name: "Bob",
+        age: 20,
+      });
+
+      expect(john.friend.name).toBe("Jane");
+      expect(johnWithANewFriend.friend?.name).toBe("Bob");
+      expect(john.$jazz.updated().friend).toBe(johnWithANewFriend.friend);
+
+      const name = johnWithANewFriend.friend?.name;
+
+      // TODO: It would be interesting to keep the Loaded type as non-nullable when returning updated values
+      // because based on the set value we know for sure that the value is or is not undefined
+      expectTypeOf<typeof name>().toMatchTypeOf<string | undefined>();
+    });
+
+    it("should return the same instance on $updated if there are no changes", () => {
+      const MyCoMap = co.map({
+        name: z.string(),
+        age: z.number(),
+      });
+
+      const myCoMap = MyCoMap.create({ name: "John", age: 30 });
+
+      expect(myCoMap.$jazz.updated()).toBe(myCoMap);
+    });
+  });
+});

packages/jazz-tools/src/schema/coValue/optional.ts

@@ -0,0 +1,51 @@
+import { CoValueSchema } from "../coMap/schema.js";
+import { SelfReference } from "./self.js";
+
+export const OptionalSymbol = "Optional" as const;
+
+export type Optional<T extends CoValueSchema<any> | SelfReference> = T & {
+  [OptionalSymbol]: true;
+};
+
+export type isOptional<T> = T extends {
+  [OptionalSymbol]: true;
+}
+  ? true
+  : false;
+
+export type addOptional<T> = isOptional<T> extends true ? undefined : never;
+
+export function carryOptional<
+  S extends CoValueSchema<any>,
+  V extends Optional<CoValueSchema<any>> | CoValueSchema<any>,
+>(
+  value: V,
+  newSchema: S,
+): V extends Optional<CoValueSchema<any>> ? Optional<S> : S {
+  if (isOptional(value)) {
+    return optional(newSchema);
+  } else {
+    return newSchema as V extends Optional<CoValueSchema<any>>
+      ? Optional<S>
+      : S;
+  }
+}
+
+export function optional<T extends CoValueSchema<any> | SelfReference>(
+  value: T,
+): Optional<T> {
+  return Object.create(value, {
+    [OptionalSymbol]: {
+      value: true,
+      writable: true,
+      enumerable: false,
+      configurable: false,
+    },
+  });
+}
+
+export function isOptional<T extends CoValueSchema<any> | SelfReference>(
+  value: T,
+): value is Optional<T> {
+  return OptionalSymbol in value;
+}

packages/jazz-tools/src/schema/coValue/self.ts

@@ -0,0 +1,29 @@
+import { CoValueSchema } from "../coMap/schema.js";
+import { optional } from "./optional.js";
+
+export const SelfReferenceSymbol = "SelfReference" as const;
+
+export type SelfReference = {
+  [SelfReferenceSymbol]: true;
+};
+
+export type markSelfReferenceAsOptional<T> = T extends SelfReference
+  ? undefined
+  : never;
+
+export function self() {
+  // Self references are always optional
+  const selfRef = optional({
+    [SelfReferenceSymbol]: true,
+  });
+
+  selfRef[SelfReferenceSymbol] = true;
+
+  return selfRef as SelfReference;
+}
+
+export function isSelfReference(value: unknown): value is SelfReference {
+  return (
+    typeof value === "object" && value !== null && SelfReferenceSymbol in value
+  );
+}

packages/jazz-tools/src/schema/coValue/typeUtils.ts

@@ -0,0 +1,21 @@
+export type optionalKeys<T extends object> = {
+  [k in keyof T]: undefined extends T[k] ? k : never;
+}[keyof T];
+
+export type requiredKeys<T extends object> = {
+  [k in keyof T]: undefined extends T[k] ? never : k;
+}[keyof T];
+
+export type addQuestionMarks<T extends object, _O = any> = {
+  [K in requiredKeys<T>]: T[K];
+} & {
+  [K in optionalKeys<T>]?: T[K];
+} & { [k in keyof T]?: unknown };
+
+export type identity<T> = T;
+export type flatten<T> = identity<{ [k in keyof T]: T[k] }>;
+
+type DEPTH_LIMIT = 5;
+
+export type IsDepthLimit<CurrentDepth extends number[]> =
+  DEPTH_LIMIT extends CurrentDepth["length"] ? true : false;

packages/jazz-tools/src/schema/coValue/types.ts

@@ -0,0 +1,107 @@
+import { TypeOf, ZodTypeAny } from "zod";
+import { CoMap } from "../coMap/instance.js";
+import {
+  CoMapFieldType,
+  CoMapSchema,
+  CoMapSchemaKey,
+  CoValueSchema,
+  UnwrapReference,
+} from "../coMap/schema.js";
+import { Optional } from "./optional.js";
+import { SelfReference } from "./self.js";
+import { IsDepthLimit, flatten } from "./typeUtils.js";
+
+export type RelationsToResolveStrict<
+  T extends CoValueSchema<any>,
+  V,
+> = V extends RelationsToResolve<T> ? RelationsToResolve<T> : V;
+
+export type RelationsToResolve<
+  S extends CoValueSchema<any>,
+  CurrentDepth extends number[] = [],
+> =
+  | true
+  | (IsDepthLimit<CurrentDepth> extends true
+      ? true
+      : S extends CoMapSchema<any>
+        ?
+            | {
+                [K in CoMapSchemaKey<S>]?: UnwrapReference<
+                  S,
+                  K
+                > extends CoValueSchema<any>
+                  ? RelationsToResolve<
+                      UnwrapReference<S, K>,
+                      [0, ...CurrentDepth]
+                    >
+                  : never;
+              }
+            | true
+        : true);
+
+export type isResolveLeaf<R> = R extends boolean | undefined
+  ? true
+  : keyof R extends never // R = {}
+    ? true
+    : false;
+
+export type Loaded<
+  S extends CoValueSchema<any>,
+  R extends RelationsToResolve<S> = true,
+  Options extends "nullable" | "non-nullable" = "nullable",
+  CurrentDepth extends number[] = [],
+> = R extends never
+  ? never
+  : S extends CoMapSchema<any>
+    ? LoadedCoMap<S, R, Options, CurrentDepth>
+    : never;
+
+export type LoadedCoMap<
+  S extends CoMapSchema<any>,
+  R extends RelationsToResolve<S>,
+  Options extends "nullable" | "non-nullable" = "non-nullable",
+  CurrentDepth extends number[] = [],
+> = flatten<
+  (S extends CoMapSchema<any>
+    ? {
+        [K in CoMapSchemaKey<S>]: CoMapFieldType<S, K> extends ZodTypeAny
+          ? TypeOf<CoMapFieldType<S, K>>
+          : UnwrapReference<S, K> extends CoValueSchema<any>
+            ? K extends keyof R
+              ? R[K] extends RelationsToResolve<UnwrapReference<S, K>>
+                ? IsDepthLimit<CurrentDepth> & isResolveLeaf<R> extends false
+                  ?
+                      | Loaded<
+                          UnwrapReference<S, K>,
+                          R[K],
+                          Options,
+                          [0, ...CurrentDepth]
+                        >
+                      | addNullable<Options, CoMapFieldType<S, K>>
+                  : null
+                : null
+              : null
+            : UnwrapZodType<CoMapFieldType<S, K>, never>;
+      }
+    : never) &
+    CoMap<S, R>
+>;
+
+export type UnwrapZodType<T, O> = T extends ZodTypeAny ? TypeOf<T> : O;
+
+export type ValidateResolve<
+  D extends CoValueSchema<any>,
+  I,
+  E,
+> = I extends RelationsToResolve<D> ? I : E;
+
+export type addNullable<
+  O extends "nullable" | "non-nullable",
+  T,
+> = O extends "nullable"
+  ? T extends Optional<infer U>
+    ? null | undefined
+    : T extends SelfReference
+      ? undefined | null
+      : never
+  : never;

packages/jazz-tools/src/schema/index.ts

@@ -0,0 +1,16 @@
+import { co, z } from "./schema.js";
+import {
+  ensureCoValueLoaded,
+  loadCoValue,
+  subscribeToCoValue,
+} from "./subscribe.js";
+
+export const SchemaV2 = {
+  co,
+  z,
+  subscribeToCoValue: subscribeToCoValue,
+  ensureCoValueLoaded: ensureCoValueLoaded,
+  loadCoValue: loadCoValue,
+} as const;
+
+export type SchemaV2 = typeof SchemaV2;

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

@@ -0,0 +1,24 @@
+import { z } from "zod";
+import { CoMapSchema } from "./coMap/schema.js";
+import { CoMapSchemaShape } from "./coMap/schema.js";
+import { optional } from "./coValue/optional.js";
+import { self } from "./coValue/self.js";
+
+function map<S extends CoMapSchemaShape>(schema: S) {
+  return new CoMapSchema(schema);
+}
+
+export const co = {
+  map,
+  self,
+  optional,
+};
+
+export { z };
+export type { CoMap } from "./coMap/instance.js";
+export type { CoMapSchema } from "./coMap/schema.js";
+export type {
+  RelationsToResolve,
+  RelationsToResolveStrict,
+  Loaded,
+} from "./coValue/types.js";

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

@@ -0,0 +1,332 @@
+import { LocalNode, RawCoMap } from "cojson";
+import { Account } from "../exports.js";
+import { activeAccountContext } from "../implementation/activeAccountContext.js";
+import { AnonymousJazzAgent, ID } from "../internal.js";
+import { CoMap, createCoMapFromRaw, isRelationRef } from "./coMap/instance.js";
+import { CoMapSchema, CoValueSchema } from "./coMap/schema.js";
+import { isSelfReference } from "./coValue/self.js";
+import {
+  Loaded,
+  LoadedCoMap,
+  RelationsToResolve,
+  RelationsToResolveStrict,
+} from "./coValue/types.js";
+
+type SubscribeListener<
+  D extends CoValueSchema<any>,
+  R extends RelationsToResolve<D>,
+> = (value: Loaded<D, R>, unsubscribe: () => void) => void;
+
+function createResolvablePromise<T>() {
+  let resolve!: (value: T) => void;
+  let reject!: (reason?: any) => void;
+
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+
+  return { promise, resolve, reject };
+}
+
+type ResolvablePromise<T> = ReturnType<typeof createResolvablePromise<T>>;
+
+class Subscription {
+  _unsubscribe: () => void = () => {};
+  unsubscribed = false;
+
+  value: RawCoMap | undefined;
+  status: "unknown" | "loading" | "loaded" | "unauthorized" | "unavailable" =
+    "unknown";
+
+  constructor(
+    public node: LocalNode,
+    public id: ID<CoValueSchema<any>>,
+    public listener: (value: RawCoMap) => void,
+  ) {
+    const value = this.node.coValuesStore.get(this.id as any);
+
+    if (value.state.type === "available") {
+      this.status = "loaded";
+      this.subscribe(value.state.coValue.getCurrentContent() as RawCoMap);
+    } else {
+      this.status = "loading";
+      this.node.load(this.id as any).then((value) => {
+        if (this.unsubscribed) return;
+        // TODO handle the error states which should be transitive
+        if (value !== "unavailable") {
+          this.status = "loaded";
+          this.subscribe(value as RawCoMap);
+        } else {
+          this.status = "unavailable";
+        }
+      });
+    }
+  }
+
+  subscribe(value: RawCoMap) {
+    if (this.unsubscribed) return;
+
+    this._unsubscribe = value.subscribe((value) => {
+      this.listener(value);
+    });
+    this.listener(value);
+  }
+
+  unsubscribe() {
+    if (this.unsubscribed) return;
+    this.unsubscribed = true;
+    this._unsubscribe();
+  }
+}
+
+export class CoValueResolutionNode<
+  D extends CoValueSchema<any>,
+  R extends RelationsToResolve<D>,
+> {
+  childNodes = new Map<
+    string,
+    CoValueResolutionNode<CoValueSchema<any>, any>
+  >();
+  childValues = new Map<string, Loaded<any, any> | undefined>();
+  value: Loaded<D, R> | undefined;
+  promise: ResolvablePromise<void> | undefined;
+  subscription: Subscription;
+  listener: ((value: Loaded<D, R>) => void) | undefined;
+
+  constructor(
+    public node: LocalNode,
+    public resolve: RelationsToResolve<D>,
+    public id: ID<D>,
+    public schema: D,
+  ) {
+    this.subscription = new Subscription(node, id, (value) => {
+      this.handleUpdate(value);
+    });
+  }
+
+  handleUpdate(value: RawCoMap) {
+    if (!this.value) {
+      this.value = createCoMapFromRaw<D, R>(
+        this.schema,
+        value,
+        this.childValues,
+        this,
+      );
+      this.loadChildren();
+      if (this.isLoaded()) {
+        this.listener?.(this.value);
+      }
+    } else if (this.isLoaded()) {
+      const changesOnChildren = this.loadChildren();
+
+      if (this.isLoaded()) {
+        const value = this.value.$jazz.updated(
+          changesOnChildren ? this.childValues : undefined,
+        ) as Loaded<D, R>;
+
+        if (value !== this.value) {
+          this.value = value;
+          this.listener?.(this.value);
+        }
+      }
+    }
+  }
+
+  handleChildUpdate = (key: string, value: Loaded<any, any>) => {
+    this.childValues.set(key, value);
+
+    if (this.value && this.isLoaded()) {
+      this.value = this.value.$jazz.updated(this.childValues) as Loaded<D, R>;
+      this.listener?.(this.value);
+    }
+  };
+
+  isLoaded() {
+    if (!this.value) return false;
+
+    for (const value of this.childValues.values()) {
+      if (value === undefined) return false;
+    }
+
+    return true;
+  }
+
+  setListener(listener: (value: Loaded<D, R>) => void) {
+    this.listener = listener;
+    if (this.value && this.isLoaded()) {
+      this.listener(this.value);
+    }
+  }
+
+  loadChildren() {
+    const { node, resolve, schema } = this;
+
+    const raw = this.value?.$jazz.raw;
+
+    if (raw === undefined) {
+      throw new Error("RefNode is not initialized");
+    }
+
+    if (typeof resolve !== "object" || resolve === null) {
+      return false;
+    }
+
+    let hasChanged = false;
+
+    for (const key of Object.keys(resolve)) {
+      const value = raw.get(key);
+      const descriptor = schema.get(key);
+
+      if (descriptor === undefined) {
+        continue;
+      }
+
+      if (this.childNodes.has(key)) {
+        const child = this.childNodes.get(key);
+
+        if (child) {
+          if (child.id !== value) {
+            hasChanged = true;
+            const childNode = this.childNodes.get(key);
+
+            if (childNode) {
+              childNode.destroy();
+            }
+
+            this.childNodes.delete(key);
+            this.childValues.delete(key);
+          } else {
+            continue;
+          }
+        }
+      }
+
+      if (value && isRelationRef(descriptor) && resolve[key]) {
+        hasChanged = true;
+        let childSchema = descriptor as CoValueSchema<any>;
+
+        if (isSelfReference(childSchema)) {
+          childSchema = this.schema;
+        }
+
+        this.childValues.set(key, undefined);
+        const child = new CoValueResolutionNode(
+          node,
+          resolve[key] as RelationsToResolve<any>,
+          raw.get(key) as ID<any>,
+          childSchema,
+        );
+        child.setListener((value) => this.handleChildUpdate(key, value));
+        this.childNodes.set(key, child);
+      }
+    }
+
+    return hasChanged;
+  }
+
+  destroy() {
+    this.subscription.unsubscribe();
+    this.childNodes.forEach((child) => child.destroy());
+  }
+}
+
+export function subscribeToCoValue<
+  D extends CoValueSchema<any>,
+  R extends RelationsToResolve<D>,
+>(
+  schema: D,
+  id: ID<D>,
+  options: {
+    resolve?: RelationsToResolveStrict<D, R>;
+    loadAs?: Account | AnonymousJazzAgent;
+    onUnavailable?: () => void;
+    onUnauthorized?: () => void;
+  },
+  listener: SubscribeListener<D, R>,
+) {
+  const loadAs = options.loadAs ?? activeAccountContext.get();
+  const node = "node" in loadAs ? loadAs.node : loadAs._raw.core.node;
+
+  const resolve = options.resolve ?? true;
+
+  let unsubscribed = false;
+
+  const rootNode = new CoValueResolutionNode<D, R>(
+    node,
+    resolve,
+    id as ID<D>,
+    schema,
+  );
+  rootNode.setListener(handleUpdate);
+
+  function unsubscribe() {
+    unsubscribed = true;
+    rootNode.subscription.unsubscribe();
+  }
+
+  function handleUpdate(value: Loaded<D, R>) {
+    if (unsubscribed) return;
+
+    listener(value, unsubscribe);
+  }
+
+  return unsubscribe;
+}
+
+export function loadCoValue<
+  D extends CoValueSchema<any>,
+  R extends RelationsToResolve<D>,
+>(
+  schema: D,
+  id: ID<D>,
+  options?: {
+    resolve?: RelationsToResolveStrict<D, R>;
+    loadAs?: Account | AnonymousJazzAgent;
+  },
+) {
+  return new Promise<Loaded<D, R> | undefined>((resolve) => {
+    subscribeToCoValue<D, R>(
+      schema,
+      id,
+      {
+        resolve: options?.resolve,
+        loadAs: options?.loadAs,
+        onUnavailable: () => {
+          resolve(undefined);
+        },
+        onUnauthorized: () => {
+          resolve(undefined);
+        },
+      },
+      (value, unsubscribe) => {
+        resolve(value);
+        unsubscribe();
+      },
+    );
+  });
+}
+
+export async function ensureCoValueLoaded<
+  D extends CoValueSchema<any>,
+  I extends RelationsToResolve<D>,
+  R extends RelationsToResolve<D>,
+>(
+  existing: Loaded<D, I>,
+  options?: { resolve?: RelationsToResolveStrict<D, R> } | undefined,
+) {
+  const response = await loadCoValue<D, R>(
+    existing.$jazz.schema as D,
+    existing.$jazz.id as ID<D>,
+    {
+      loadAs: existing.$jazz._loadedAs,
+      resolve: options?.resolve,
+    },
+  );
+
+  if (!response) {
+    throw new Error("Failed to deeply load CoValue " + existing.$jazz.id);
+  }
+
+  return response;
+}

pnpm-lock.yaml

@@ -2082,6 +2082,9 @@ importers:
       fast-myers-diff:
         specifier: ^3.2.0
         version: 3.2.0
+      zod:
+        specifier: ^3.24.2
+        version: 3.24.2
     devDependencies:
       tsup:
         specifier: 8.3.5
@@ -11998,6 +12001,9 @@ packages:
   zod@3.24.1:
     resolution: {integrity: sha512-muH7gBL9sI1nciMZV67X5fTKKBLtwpZ5VBp1vsOQzj1MhrBZ4wlVCm3gedKZWLp0Oyel8sIGfeiz54Su+OVT+A==}
 
+  zod@3.24.2:
+    resolution: {integrity: sha512-lY7CDW43ECgW9u1TcT3IoXHflywfVqDYze4waEz812jR/bZ8FHDsl7pFQoSZTz5N+2NqRXs8GBwnAwo3ZNxqhQ==}
+
 snapshots:
 
   '@0no-co/graphql.web@1.0.12(graphql@16.10.0)':
@@ -22980,3 +22986,5 @@ snapshots:
   zod@3.22.3: {}
 
   zod@3.24.1: {}
+
+  zod@3.24.2: {}

tests/e2e/src/pages/TestInput.tsx

@@ -1,8 +1,8 @@
-import { useAccount, useCoState } from "jazz-react";
-import { CoMap, Group, ID, co } from "jazz-tools";
+import { useAccount, useCoState2 as useCoState } from "jazz-react";
+import { Group, ID, SchemaV2, co } from "jazz-tools";
 import { useEffect, useState } from "react";
 
-export class InputTestCoMap extends CoMap {
+export class InputTestCoMap extends SchemaV2.CoMap {
   title = co.string;
 }
 
@@ -18,7 +18,7 @@ export function TestInput() {
 
     group.addMember("everyone", "writer");
 
-    setId(InputTestCoMap.create({ title: "" }, { owner: group }).id);
+    setId(InputTestCoMap.create({ title: "" }, { owner: group }).$id);
   }, [me]);
 
   if (!coMap) return null;
@@ -28,7 +28,9 @@ export function TestInput() {
       value={coMap?.title ?? ""}
       onChange={(e) => {
         if (!coMap) return;
-        coMap.title = e.target.value;
+
+        // @ts-expect-error
+        coMap.$set("title", e.target.value);
       }}
     />
   );