jazz-tools 0.15.9 → 0.15.11

package/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);
+  }
 }

package/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";

package/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,9 +43,16 @@ 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) {
   if (isCoValueClass(schema)) {
     return coField.ref(schema);
+  } else if (isCoValueSchema(schema)) {
+    if (isAnyCoOptionalSchema(schema)) {
+      return coField.ref(schema.getCoValueClass(), {
+        optional: true,
+      });
+    }
+    return coField.ref(schema.getCoValueClass());
   } else {
     if ("_zod" in schema) {
       if (schema._zod.def.type === "optional") {
@@ -45,7 +62,7 @@ export function zodFieldToCoFieldDef(schema: FieldSchema) {
         if (isCoValueClass(inner)) {
           return coField.ref(inner, { optional: true });
         } else {
-          return zodFieldToCoFieldDef(inner);
+          return schemaFieldToCoFieldDef(inner);
         }
       } else if (schema._zod.def.type === "string") {
         return coField.string;
@@ -58,8 +75,8 @@ 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,
         );
       } else if (schema._zod.def.type === "date") {
         return coField.optional.Date;
@@ -67,8 +84,8 @@ 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,
         );
       } else if (
         schema._zod.def.type === "default" ||
@@ -78,9 +95,9 @@ 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,
         );
       } else if (schema._zod.def.type === "literal") {
         if (
@@ -112,7 +129,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);
         } else {
           throw new Error(`Unsupported custom zod type`);
         }

package/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;

package/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> & {

package/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]>>;

package/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;