jazz-tools 0.15.16 → 0.16.0

package/src/tools/implementation/zodSchema/zodCo.ts

@@ -1,52 +1,43 @@
 import {
   type AccountSchema,
-  AnyAccountSchema,
-  AnyCoFeedSchema,
-  AnyCoListSchema,
-  AnyCoSchema,
-  AnyFileStreamSchema,
-  AnyPlainTextSchema,
+  AnyZodOrCoValueSchema,
   BaseAccountShape,
-  CoFeed,
   type CoFeedSchema,
   type CoListSchema,
   type CoMapSchema,
-  CoPlainText,
   type CoProfileSchema,
   type CoRecordSchema,
-  CoRichText,
   type DefaultProfileShape,
-  FileStream,
   type FileStreamSchema,
   ImageDefinition,
   type PlainTextSchema,
   type Simplify,
-  zodSchemaToCoSchema,
+  createCoreAccountSchema,
+  createCoreCoFeedSchema,
+  createCoreCoListSchema,
+  createCoreCoMapSchema,
+  createCoreCoPlainTextSchema,
+  createCoreFileStreamSchema,
+  hydrateCoreCoValueSchema,
 } from "../../internal.js";
 import {
-  AnyDiscriminableCoSchema,
   CoDiscriminatedUnionSchema,
+  DiscriminableCoValueSchemas,
+  createCoreCoDiscriminatedUnionSchema,
 } from "./schemaTypes/CoDiscriminatedUnionSchema.js";
+import { CoOptionalSchema } from "./schemaTypes/CoOptionalSchema.js";
+import { CoreCoValueSchema } from "./schemaTypes/CoValueSchema.js";
 import {
-  CoOptionalSchema,
-  createCoOptionalSchema,
-} from "./schemaTypes/CoOptionalSchema.js";
-import {
-  AnyRichTextSchema,
   RichTextSchema,
+  createCoreCoRichTextSchema,
 } from "./schemaTypes/RichTextSchema.js";
 import { z } from "./zodReExport.js";
 
 export const coMapDefiner = <Shape extends z.core.$ZodLooseShape>(
   shape: Shape,
 ): CoMapSchema<Shape> => {
-  const objectSchema = z.object(shape).meta({
-    collaborative: true,
-  });
-  const enrichedSchema = Object.assign(objectSchema, {
-    collaborative: true,
-  });
-  return zodSchemaToCoSchema(enrichedSchema);
+  const coreSchema = createCoreCoMapSchema(shape);
+  return hydrateCoreCoValueSchema(coreSchema);
 };
 
 /**
@@ -96,19 +87,13 @@ export const coAccountDefiner = <Shape extends BaseAccountShape>(
     root: coMapDefiner({}),
   } as unknown as Shape,
 ): AccountSchema<Shape> => {
-  const schema = z.object(shape).meta({
-    collaborative: true,
-  });
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    builtin: "Account",
-  }) as AnyAccountSchema<Shape>;
-  return zodSchemaToCoSchema(enrichedSchema);
+  const coreSchema = createCoreAccountSchema(shape);
+  return hydrateCoreCoValueSchema(coreSchema);
 };
 
 export const coRecordDefiner = <
   K extends z.core.$ZodString<string>,
-  V extends z.core.$ZodType,
+  V extends AnyZodOrCoValueSchema,
 >(
   _keyType: K,
   valueType: V,
@@ -119,16 +104,11 @@ export const coRecordDefiner = <
   >;
 };
 
-export const coListDefiner = <T extends z.core.$ZodType>(
+export const coListDefiner = <T extends AnyZodOrCoValueSchema>(
   element: T,
 ): CoListSchema<T> => {
-  const schema = z.array(element).meta({
-    collaborative: true,
-  });
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-  }) as AnyCoListSchema<T>;
-  return zodSchemaToCoSchema(enrichedSchema) as unknown as CoListSchema<T>;
+  const coreSchema = createCoreCoListSchema(element);
+  return hydrateCoreCoValueSchema(coreSchema) as unknown as CoListSchema<T>;
 };
 
 export const coProfileDefiner = <
@@ -144,61 +124,48 @@ export const coProfileDefiner = <
   return coMapDefiner(ehnancedShape) as CoProfileSchema<Shape>;
 };
 
-export const coFeedDefiner = <T extends z.core.$ZodType>(
+export const coFeedDefiner = <T extends AnyZodOrCoValueSchema>(
   element: T,
 ): CoFeedSchema<T> => {
-  const schema = z.instanceof(CoFeed);
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    builtin: "CoFeed",
-    element,
-  }) as AnyCoFeedSchema<T>;
-  return zodSchemaToCoSchema(enrichedSchema);
+  const coreSchema = createCoreCoFeedSchema(element);
+  return hydrateCoreCoValueSchema(coreSchema);
 };
 
 export const coFileStreamDefiner = (): FileStreamSchema => {
-  const schema = z.instanceof(FileStream);
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    builtin: "FileStream",
-  }) as AnyFileStreamSchema;
-  return zodSchemaToCoSchema(enrichedSchema);
+  const coreSchema = createCoreFileStreamSchema();
+  return hydrateCoreCoValueSchema(coreSchema);
 };
 
 export const coPlainTextDefiner = (): PlainTextSchema => {
-  const schema = z.instanceof(CoPlainText);
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    builtin: "CoPlainText",
-  }) as AnyPlainTextSchema;
-  return zodSchemaToCoSchema(enrichedSchema);
+  const coreSchema = createCoreCoPlainTextSchema();
+  return hydrateCoreCoValueSchema(coreSchema);
 };
 
 export const coRichTextDefiner = (): RichTextSchema => {
-  const schema = z.instanceof(CoRichText);
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    builtin: "CoRichText",
-  }) as AnyRichTextSchema;
-  return zodSchemaToCoSchema(enrichedSchema);
+  const coreSchema = createCoreCoRichTextSchema();
+  return hydrateCoreCoValueSchema(coreSchema);
 };
 
-export const coImageDefiner = (): typeof ImageDefinition => {
+export type ImageDefinitionSchema = typeof ImageDefinition;
+export const coImageDefiner = (): ImageDefinitionSchema => {
   return ImageDefinition;
 };
 
-export const coOptionalDefiner = <T extends AnyCoSchema>(
+export const coOptionalDefiner = <T extends CoreCoValueSchema>(
   schema: T,
 ): CoOptionalSchema<T> => {
-  return createCoOptionalSchema(schema);
+  return new CoOptionalSchema(schema);
 };
 
 export const coDiscriminatedUnionDefiner = <
-  T extends readonly [AnyDiscriminableCoSchema, ...AnyDiscriminableCoSchema[]],
+  Options extends DiscriminableCoValueSchemas,
 >(
   discriminator: string,
-  schemas: T,
-): CoDiscriminatedUnionSchema<T> => {
-  const schema = z.discriminatedUnion(discriminator, schemas as any);
-  return zodSchemaToCoSchema(schema);
+  schemas: Options,
+): CoDiscriminatedUnionSchema<Options> => {
+  const coreSchema = createCoreCoDiscriminatedUnionSchema(
+    discriminator,
+    schemas,
+  );
+  return hydrateCoreCoValueSchema(coreSchema);
 };

package/src/tools/implementation/zodSchema/zodReExport.ts

@@ -1,17 +1,14 @@
 import {
-  ZodArray,
-  ZodTuple,
-  ZodUnion,
+  ZodObject,
   core,
-  array as zodArray,
-  tuple as zodTuple,
-  union as zodUnion,
+  object as zodObject,
+  strictObject as zodStrictObject,
 } from "zod/v4";
+import { removeGetters } from "../schemaUtils.js";
 export {
   string,
   number,
   boolean,
-  object,
   templateLiteral,
   json,
   date,
@@ -33,12 +30,14 @@ export {
   cidrv6,
   iso,
   int32,
-  strictObject,
+  union,
   discriminatedUnion,
   // record,
   // intersection,
   int,
   optional,
+  array,
+  tuple,
   type ZodOptional,
   type ZodReadonly,
   type ZodLazy,
@@ -49,25 +48,46 @@ export {
   z,
 } from "zod/v4";
 
-type NonCoZodType = core.$ZodType & { collaborative?: false };
+export function object<
+  T extends core.$ZodLooseShape = Partial<Record<never, core.SomeType>>,
+>(
+  shape?: T,
+  params?: string | core.$ZodObjectParams,
+): ZodObject<T, core.$strip> {
+  rejectCoValueSchemas(
+    shape,
+    "z.object() does not support collaborative types as values. Use co.map() instead",
+  );
+  return zodObject(shape, params);
+}
+
+export function strictObject<T extends core.$ZodLooseShape>(
+  shape: T,
+  params?: string | core.$ZodObjectParams,
+): ZodObject<T, core.$strict> {
+  rejectCoValueSchemas(
+    shape,
+    "z.strictObject() does not support collaborative types as values. Use co.map() instead",
+  );
+  return zodStrictObject(shape, params);
+}
 
-export function union<const T extends readonly NonCoZodType[]>(
-  options: T,
-  params?: string | core.$ZodUnionParams,
-): ZodUnion<T> {
-  return zodUnion(options, params);
+function rejectCoValueSchemas(
+  shape: core.$ZodLooseShape | undefined,
+  errorMessage: string,
+) {
+  if (containsCoValueSchema(shape)) {
+    throw Error(errorMessage);
+  }
 }
 
-export function array<T extends NonCoZodType>(
-  element: T,
-  params?: string | core.$ZodArrayParams,
-): ZodArray<T> {
-  return zodArray(element, params);
+function containsCoValueSchema(shape?: core.$ZodLooseShape): boolean {
+  // Remove getters to avoid circularity issues accessing schemas that may not be defined yet
+  return Object.values(removeGetters(shape ?? {})).some(isAnyCoValueSchema);
 }
 
-export function tuple<T extends readonly [NonCoZodType, ...NonCoZodType[]]>(
-  options: T,
-  params?: string | core.$ZodTupleParams,
-): ZodTuple<T> {
-  return zodTuple(options, params);
+// Note: if you're editing this function, edit the `isAnyCoValueSchema`
+// function in `zodSchemaToCoSchema.ts` as well
+function isAnyCoValueSchema(schema: any): boolean {
+  return "collaborative" in schema && schema.collaborative === true;
 }

package/src/tools/implementation/zodSchema/zodSchema.ts

@@ -2,8 +2,10 @@ import { LocalNode, RawAccount } from "cojson";
 import {
   Account,
   AccountClass,
+  CoRecordSchema,
   CoValueClass,
   CoValueFromRaw,
+  CoreCoRecordSchema,
   InstanceOfSchema,
   RefsToResolve,
   RefsToResolveStrict,
@@ -12,47 +14,40 @@ import {
 } from "../../internal.js";
 import {
   AccountSchema,
-  AnyAccountSchema,
   BaseAccountShape,
+  CoreAccountSchema,
 } from "./schemaTypes/AccountSchema.js";
 import {
-  AnyDiscriminableCoSchema,
   CoDiscriminatedUnionSchema,
+  CoreCoDiscriminatedUnionSchema,
 } from "./schemaTypes/CoDiscriminatedUnionSchema.js";
-import { AnyCoFeedSchema, CoFeedSchema } from "./schemaTypes/CoFeedSchema.js";
-import { AnyCoListSchema, CoListSchema } from "./schemaTypes/CoListSchema.js";
+import { CoFeedSchema, CoreCoFeedSchema } from "./schemaTypes/CoFeedSchema.js";
+import { CoListSchema, CoreCoListSchema } from "./schemaTypes/CoListSchema.js";
 import {
-  AnyCoMapSchema,
   CoMapInitZod,
   CoMapSchema,
+  CoreCoMapSchema,
 } from "./schemaTypes/CoMapSchema.js";
-import { AnyCoOptionalSchema } from "./schemaTypes/CoOptionalSchema.js";
 import {
-  AnyCoRecordSchema,
-  CoRecordSchema,
-} from "./schemaTypes/CoRecordSchema.js";
+  CoOptionalSchema,
+  CoreCoOptionalSchema,
+} from "./schemaTypes/CoOptionalSchema.js";
+import { CoreCoValueSchema } from "./schemaTypes/CoValueSchema.js";
 import {
-  AnyFileStreamSchema,
+  CoreFileStreamSchema,
   FileStreamSchema,
 } from "./schemaTypes/FileStreamSchema.js";
 import {
-  AnyPlainTextSchema,
+  CorePlainTextSchema,
   PlainTextSchema,
 } from "./schemaTypes/PlainTextSchema.js";
 import {
-  AnyRichTextSchema,
+  CoreRichTextSchema,
   RichTextSchema,
 } from "./schemaTypes/RichTextSchema.js";
 import { InstanceOfSchemaCoValuesNullable } from "./typeConverters/InstanceOfSchemaCoValuesNullable.js";
 import { z } from "./zodReExport.js";
 
-// defining an extra type for this, otherwise BaseSchema & {...} often
-// gets expanded into a n inferred type that's too long for typescript to print
-export type WithHelpers<
-  Base extends z.core.$ZodType,
-  Helpers extends object,
-> = Base & Helpers;
-
 export type ZodPrimitiveSchema =
   | z.core.$ZodString
   | z.core.$ZodNumber
@@ -61,108 +56,75 @@ export type ZodPrimitiveSchema =
   | z.core.$ZodDate
   | z.core.$ZodLiteral;
 
-export type AnyCoUnionSchema = z.core.$ZodDiscriminatedUnion<
-  (
-    | (z.core.$ZodType & { collaborative: true })
-    | z.core.$ZodDiscriminatedUnion
-  )[]
->;
-
-// this is a series of hacks to work around z4 removing _zod at runtime from z.core.$ZodType
-export function isZodObject(
-  schema: z.core.$ZodType,
-): schema is z.ZodObject<any, any> {
-  return (schema as any).def?.type === "object";
-}
-
-export function isZodArray(
-  schema: z.core.$ZodType,
-): schema is z.core.$ZodArray<any> {
-  return (schema as any).def?.type === "array";
-}
-
-export function isZodCustom(
-  schema: z.core.$ZodType,
-): schema is z.core.$ZodCustom<any, any> {
-  return (schema as any).def?.type === "custom";
-}
-
-export function getDef<S extends z.core.$ZodType>(schema: S): S["_zod"]["def"] {
-  return (schema as any).def;
-}
-
-// TODO rename. This represents a CoValue class or a CoValue schema
-export type CoValueOrZodSchema = CoValueClass | AnyCoSchema;
-
-// TODO rename to CoValueSchemaFromCoProtoSchema
-export type CoValueSchemaFromZodSchema<S extends z.core.$ZodType> =
-  S extends z.core.$ZodType
-    ? S extends AnyAccountSchema<infer Shape extends BaseAccountShape>
-      ? AccountSchema<Shape>
-      : S extends AnyCoRecordSchema<infer K, infer V>
-        ? CoRecordSchema<K, V>
-        : S extends AnyCoMapSchema<infer Shape, infer Config>
-          ? CoMapSchema<Shape, Config>
-          : S extends AnyCoListSchema<infer T>
-            ? CoListSchema<T>
-            : S extends AnyCoFeedSchema<infer T>
-              ? CoFeedSchema<T>
-              : S extends AnyPlainTextSchema
-                ? PlainTextSchema
-                : S extends AnyRichTextSchema
-                  ? RichTextSchema
-                  : S extends AnyFileStreamSchema
-                    ? FileStreamSchema
-                    : S extends z.core.$ZodOptional<infer Inner>
-                      ? CoValueSchemaFromZodSchema<Inner>
-                      : S extends z.core.$ZodUnion<
-                            infer Members extends readonly [
-                              AnyDiscriminableCoSchema,
-                              ...AnyDiscriminableCoSchema[],
-                            ]
-                          >
-                        ? CoDiscriminatedUnionSchema<Members>
-                        : never
-    : never;
-
-export type CoValueClassFromAnySchema<S extends CoValueOrZodSchema> =
+export type CoValueClassOrSchema = CoValueClass | CoreCoValueSchema;
+
+export type CoValueSchemaFromCoreSchema<S extends CoreCoValueSchema> =
+  S extends CoreAccountSchema<infer Shape extends BaseAccountShape>
+    ? AccountSchema<Shape>
+    : S extends CoreCoRecordSchema<infer K, infer V>
+      ? CoRecordSchema<K, V>
+      : S extends CoreCoMapSchema<infer Shape, infer Config>
+        ? CoMapSchema<Shape, Config>
+        : S extends CoreCoListSchema<infer T>
+          ? CoListSchema<T>
+          : S extends CoreCoFeedSchema<infer T>
+            ? CoFeedSchema<T>
+            : S extends CorePlainTextSchema
+              ? PlainTextSchema
+              : S extends CoreRichTextSchema
+                ? RichTextSchema
+                : S extends CoreFileStreamSchema
+                  ? FileStreamSchema
+                  : S extends CoreCoOptionalSchema<infer Inner>
+                    ? CoOptionalSchema<Inner>
+                    : S extends CoreCoDiscriminatedUnionSchema<infer Members>
+                      ? CoDiscriminatedUnionSchema<Members>
+                      : never;
+
+export type CoValueClassFromAnySchema<S extends CoValueClassOrSchema> =
   S extends CoValueClass<any>
     ? S
     : CoValueClass<InstanceOfSchema<S>> &
         CoValueFromRaw<InstanceOfSchema<S>> &
-        (S extends AnyAccountSchema ? AccountClassEssentials : {});
+        (S extends CoreAccountSchema ? AccountClassEssentials : {});
 
 type AccountClassEssentials = {
   fromRaw: <A extends Account>(this: AccountClass<A>, raw: RawAccount) => A;
   fromNode: <A extends Account>(this: AccountClass<A>, node: LocalNode) => A;
 };
 
-// TODO rename to ProtoCoSchema?
-export type AnyCoSchema =
-  | AnyCoMapSchema
-  | AnyAccountSchema
-  | AnyCoRecordSchema
-  | AnyCoListSchema
-  | AnyCoFeedSchema
-  | AnyCoUnionSchema
-  | AnyCoOptionalSchema
-  | AnyPlainTextSchema
-  | AnyRichTextSchema
-  | AnyFileStreamSchema;
+export type AnyCoreCoValueSchema =
+  | CoreCoMapSchema
+  | CoreAccountSchema
+  | CoreCoRecordSchema
+  | CoreCoListSchema
+  | CoreCoFeedSchema
+  | CoreCoDiscriminatedUnionSchema<any>
+  | CoreCoOptionalSchema
+  | CorePlainTextSchema
+  | CoreRichTextSchema
+  | CoreFileStreamSchema;
+
+type AnyZodSchema = z.core.$ZodType;
+
+export type AnyZodOrCoValueSchema = AnyZodSchema | CoreCoValueSchema;
 
 export type Loaded<
-  T extends CoValueClass | AnyCoSchema,
+  T extends CoValueClassOrSchema,
   R extends ResolveQuery<T> = true,
 > = Resolved<NonNullable<InstanceOfSchemaCoValuesNullable<T>>, R>;
 
-export type ResolveQuery<T extends CoValueClass | AnyCoSchema> = RefsToResolve<
+export type ResolveQuery<T extends CoValueClassOrSchema> = RefsToResolve<
   NonNullable<InstanceOfSchemaCoValuesNullable<T>>
 >;
 
 export type ResolveQueryStrict<
-  T extends CoValueClass | AnyCoSchema,
+  T extends CoValueClassOrSchema,
   R extends ResolveQuery<T>,
 > = RefsToResolveStrict<NonNullable<InstanceOfSchemaCoValuesNullable<T>>, R>;
 
-export type InitFor<T extends CoValueClass | AnyCoSchema> =
-  T extends AnyCoMapSchema<infer Shape> ? Simplify<CoMapInitZod<Shape>> : never;
+export type InitFor<T extends CoValueClassOrSchema> = T extends CoreCoMapSchema<
+  infer Shape
+>
+  ? Simplify<CoMapInitZod<Shape>>
+  : never;

package/src/tools/internal.ts

@@ -26,6 +26,7 @@ export * from "./subscribe/SubscriptionScope.js";
 export * from "./subscribe/types.js";
 export * from "./subscribe/index.js";
 export * from "./lib/cache.js";
+export * from "./lib/utilityTypes.js";
 export * from "./implementation/createContext.js";
 
 export * from "./types.js";
@@ -44,8 +45,8 @@ export * from "./implementation/zodSchema/typeConverters/InstanceOrPrimitiveOfSc
 export * from "./implementation/zodSchema/typeConverters/InstanceOrPrimitiveOfSchemaCoValuesNullable.js";
 export * from "./implementation/zodSchema/typeConverters/InstanceOfSchema.js";
 export * from "./implementation/zodSchema/typeConverters/InstanceOfSchemaCoValuesNullable.js";
-export * from "./implementation/zodSchema/runtimeConverters/zodSchemaToCoSchema.js";
-export * from "./implementation/zodSchema/runtimeConverters/zodFieldToCoFieldDef.js";
+export * from "./implementation/zodSchema/runtimeConverters/coValueSchemaTransformation.js";
+export * from "./implementation/zodSchema/runtimeConverters/schemaFieldToCoFieldDef.js";
 export * from "./coValues/extensions/imageDef.js";
 
 export * from "./implementation/ContextManager.js";

package/src/tools/lib/utilityTypes.ts

@@ -0,0 +1,7 @@
+/**
+ * Useful to flatten the type output to improve type hints shown in editors.
+ * And also to transform an interface into a type to aide with assignability.
+ *
+ * Taken from https://github.com/sindresorhus/type-fest/blob/main/source/simplify.d.ts
+ */
+export type Simplify<T> = { [KeyType in keyof T]: T[KeyType] } & {};

package/src/tools/subscribe/utils.ts

@@ -3,7 +3,7 @@ import { RegisteredSchemas } from "../coValues/registeredSchemas.js";
 import {
   CoValue,
   RefEncoded,
-  anySchemaToCoSchema,
+  coValueClassFromCoValueClassOrSchema,
   instantiateRefEncoded,
 } from "../internal.js";
 import { coValuesCache } from "../lib/cache.js";
@@ -14,7 +14,9 @@ export function getOwnerFromRawValue(raw: RawCoValue) {
 
   return coValuesCache.get(owner, () =>
     owner instanceof RawAccount
-      ? anySchemaToCoSchema(RegisteredSchemas["Account"]).fromRaw(owner)
+      ? coValueClassFromCoValueClassOrSchema(
+          RegisteredSchemas["Account"],
+        ).fromRaw(owner)
       : RegisteredSchemas["Group"].fromRaw(owner as any),
   );
 }

package/src/tools/testing.ts

@@ -5,15 +5,15 @@ import {
   Account,
   AccountClass,
   type AnonymousJazzAgent,
-  AnyAccountSchema,
   AuthCredentials,
   CoValueFromRaw,
+  CoreAccountSchema,
   InstanceOfSchema,
   JazzContextManager,
   JazzContextManagerAuthProps,
   JazzContextManagerBaseProps,
   activeAccountContext,
-  anySchemaToCoSchema,
+  coValueClassFromCoValueClassOrSchema,
   createAnonymousJazzContext,
   createJazzContext,
   randomSessionProvider,
@@ -88,14 +88,14 @@ let isMigrationActive = false;
 export async function createJazzTestAccount<
   S extends
     | (AccountClass<Account> & CoValueFromRaw<Account>)
-    | AnyAccountSchema,
+    | CoreAccountSchema,
 >(options?: {
   isCurrentActiveAccount?: boolean;
   AccountSchema?: S;
   creationProps?: Record<string, unknown>;
 }): Promise<InstanceOfSchema<S>> {
   const AccountClass = options?.AccountSchema
-    ? anySchemaToCoSchema(options.AccountSchema)
+    ? coValueClassFromCoValueClassOrSchema(options.AccountSchema)
     : Account;
   const peers = [];
   if (syncServer.current) {

package/src/tools/tests/ContextManager.test.ts

@@ -25,8 +25,7 @@ import {
   CoValueFromRaw,
   InstanceOfSchema,
   Loaded,
-  anySchemaToCoSchema,
-  zodSchemaToCoSchema,
+  coValueClassFromCoValueClassOrSchema,
 } from "../internal";
 import {
   createJazzTestAccount,
@@ -244,7 +243,7 @@ describe("ContextManager", () => {
 
     const AccountRoot = co.map({
       value: z.string(),
-      get transferredRoot(): z.ZodOptional<typeof AccountRoot> {
+      get transferredRoot(): co.Optional<typeof AccountRoot> {
         return co.optional(AccountRoot);
       },
     });
@@ -270,7 +269,7 @@ describe("ContextManager", () => {
     >();
 
     await customManager.createContext({
-      AccountSchema: anySchemaToCoSchema(CustomAccount),
+      AccountSchema: coValueClassFromCoValueClassOrSchema(CustomAccount),
       storage: dbFilename,
       onAnonymousAccountDiscarded: async (anonymousAccount) => {
         const anonymousAccountWithRoot = await anonymousAccount.ensureLoaded({
@@ -336,7 +335,7 @@ describe("ContextManager", () => {
 
     // Create initial anonymous context
     await customManager.createContext({
-      AccountSchema: anySchemaToCoSchema(CustomAccount),
+      AccountSchema: coValueClassFromCoValueClassOrSchema(CustomAccount),
     });
 
     const account = (
@@ -387,7 +386,7 @@ describe("ContextManager", () => {
 
     // Create initial anonymous context
     await customManager.createContext({
-      AccountSchema: anySchemaToCoSchema(CustomAccount),
+      AccountSchema: coValueClassFromCoValueClassOrSchema(CustomAccount),
     });
 
     const account = (
@@ -412,7 +411,7 @@ describe("ContextManager", () => {
   test("onAnonymousAccountDiscarded should work on transfering data between accounts", async () => {
     const AccountRoot = co.map({
       value: z.string(),
-      get transferredRoot(): z.ZodOptional<typeof AccountRoot> {
+      get transferredRoot(): co.Optional<typeof AccountRoot> {
         return co.optional(AccountRoot);
       },
     });
@@ -459,12 +458,12 @@ describe("ContextManager", () => {
     // Create initial anonymous context
     await customManager.createContext({
       onAnonymousAccountDiscarded,
-      AccountSchema: anySchemaToCoSchema(CustomAccount),
+      AccountSchema: coValueClassFromCoValueClassOrSchema(CustomAccount),
     });
 
     const account = await createJazzTestAccount({
       isCurrentActiveAccount: true,
-      AccountSchema: zodSchemaToCoSchema(CustomAccount),
+      AccountSchema: CustomAccount,
     });
 
     await customManager.authenticate({