jazz-tools 0.13.31 → 0.14.0

package/src/implementation/zodSchema/zodCo.ts

@@ -0,0 +1,409 @@
+import { CoValueUniqueness } from "cojson";
+import z from "zod/v4";
+import {
+  Account,
+  AccountCreationProps,
+  AccountInstance,
+  AccountSchema,
+  AnonymousJazzAgent,
+  AnyCoMapSchema,
+  CoFeed,
+  CoFeedSchema,
+  CoListSchema,
+  CoMapSchema,
+  CoPlainText,
+  CoProfileSchema,
+  CoRecordSchema,
+  CoRichText,
+  DefaultProfileShape,
+  FileStream,
+  FileStreamSchema,
+  Group,
+  ImageDefinition,
+  PlainTextSchema,
+  Simplify,
+  zodSchemaToCoSchema,
+} from "../../internal.js";
+import { RichTextSchema } from "./schemaTypes/RichTextSchema.js";
+
+export const coMapDefiner = <Shape extends z.core.$ZodLooseShape>(
+  shape: Shape,
+): CoMapSchema<Shape> => {
+  const objectSchema = z.object(shape).meta({
+    collaborative: true,
+  });
+
+  type CleanedType = Pick<
+    typeof objectSchema,
+    "_zod" | "def" | "~standard" | "catchall"
+  >;
+
+  const coMapSchema = objectSchema as unknown as CleanedType & {
+    collaborative: true;
+    create: CoMapSchema<Shape>["create"];
+    load: CoMapSchema<Shape>["load"];
+    subscribe: CoMapSchema<Shape>["subscribe"];
+    findUnique: CoMapSchema<Shape>["findUnique"];
+    catchall: CoMapSchema<Shape>["catchall"];
+    withHelpers: CoMapSchema<Shape>["withHelpers"];
+  };
+
+  coMapSchema.collaborative = true;
+
+  coMapSchema.create = function (this: CoMapSchema<Shape>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).create(...args);
+  } as CoMapSchema<Shape>["create"];
+
+  coMapSchema.load = function (this: CoMapSchema<Shape>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).load(...args);
+  } as CoMapSchema<Shape>["load"];
+
+  coMapSchema.subscribe = function (this: CoMapSchema<Shape>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).subscribe(...args);
+  } as CoMapSchema<Shape>["subscribe"];
+
+  coMapSchema.findUnique = function (
+    this: CoMapSchema<Shape>,
+    unique: CoValueUniqueness["uniqueness"],
+    ownerID: string,
+    as?: Account | Group | AnonymousJazzAgent,
+  ) {
+    return (zodSchemaToCoSchema(this) as any).findUnique(unique, ownerID, as);
+  } as CoMapSchema<Shape>["findUnique"];
+
+  const oldCatchall = coMapSchema.catchall;
+  coMapSchema.catchall = function (
+    this: CoMapSchema<Shape>,
+    schema: z.core.$ZodType,
+  ) {
+    return { ...this, ...oldCatchall(schema) } as any;
+  } as CoMapSchema<Shape>["catchall"] as any;
+
+  coMapSchema.withHelpers = function (
+    this: CoMapSchema<Shape>,
+    helpers: (Self: CoMapSchema<Shape>) => object,
+  ) {
+    return { ...this, ...helpers(this) };
+  } as CoMapSchema<Shape>["withHelpers"];
+
+  return coMapSchema as unknown as CoMapSchema<Shape>;
+};
+
+const coAccountDefiner = <
+  Shape extends {
+    profile: AnyCoMapSchema<{
+      name: z.core.$ZodString<string>;
+      inbox?: z.core.$ZodOptional<z.core.$ZodString>;
+      inboxInvite?: z.core.$ZodOptional<z.core.$ZodString>;
+    }>;
+    root: AnyCoMapSchema;
+  },
+>(
+  shape: Shape = {
+    profile: co.map({
+      name: z.string(),
+      inbox: z.optional(z.string()),
+      inboxInvite: z.optional(z.string()),
+    }),
+    root: co.map({}),
+  } as unknown as Shape,
+): AccountSchema<Shape> => {
+  const objectSchema = z.object(shape).meta({
+    collaborative: true,
+  });
+
+  type CleanedType = Pick<
+    typeof objectSchema,
+    "_zod" | "def" | "~standard" | "catchall"
+  >;
+
+  const accountSchema = objectSchema as unknown as CleanedType & {
+    collaborative: true;
+    builtin: "Account";
+    migration?: (
+      account: AccountInstance<Shape>,
+      creationProps?: AccountCreationProps,
+    ) => void;
+
+    create: AccountSchema<Shape>["create"];
+    createAs: AccountSchema<Shape>["createAs"];
+    getMe: AccountSchema<Shape>["getMe"];
+    load: AccountSchema<Shape>["load"];
+    subscribe: AccountSchema<Shape>["subscribe"];
+    withHelpers: AccountSchema<Shape>["withHelpers"];
+    withMigration: AccountSchema<Shape>["withMigration"];
+  };
+
+  accountSchema.collaborative = true;
+  accountSchema.builtin = "Account";
+
+  accountSchema.create = function (this: AccountSchema<Shape>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).create(...args);
+  } as AccountSchema<Shape>["create"];
+
+  accountSchema.createAs = function (
+    this: AccountSchema<Shape>,
+    ...args: any[]
+  ) {
+    return (zodSchemaToCoSchema(this) as any).createAs(...args);
+  } as AccountSchema<Shape>["createAs"];
+
+  accountSchema.getMe = function (this: AccountSchema<Shape>) {
+    return (zodSchemaToCoSchema(this) as any).getMe();
+  } as AccountSchema<Shape>["getMe"];
+
+  accountSchema.load = function (this: AccountSchema<Shape>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).load(...args);
+  } as AccountSchema<Shape>["load"];
+
+  accountSchema.subscribe = function (
+    this: AccountSchema<Shape>,
+    ...args: any[]
+  ) {
+    return (zodSchemaToCoSchema(this) as any).subscribe(...args);
+  } as AccountSchema<Shape>["subscribe"];
+
+  accountSchema.withHelpers = function (
+    this: CoMapSchema<Shape>,
+    helpers: object,
+  ) {
+    return { ...this, ...helpers };
+  } as CoMapSchema<Shape>["withHelpers"];
+
+  accountSchema.withMigration = function (
+    this: AccountSchema<Shape>,
+    migration: (
+      account: AccountInstance<Shape>,
+      creationProps?: AccountCreationProps,
+    ) => void,
+  ) {
+    return { ...this, migration };
+  } as AccountSchema<Shape>["withMigration"];
+
+  return accountSchema as unknown as AccountSchema<Shape>;
+};
+
+const coRecordDefiner = <
+  K extends z.core.$ZodString<string>,
+  V extends z.core.$ZodType,
+>(
+  _keyType: K,
+  valueType: V,
+): CoRecordSchema<K, V> => {
+  return coMapDefiner({}).catchall(valueType) as unknown as CoRecordSchema<
+    K,
+    V
+  >;
+};
+
+const coListDefiner = <T extends z.core.$ZodType>(
+  element: T,
+): CoListSchema<T> => {
+  const arraySchema = z.array(element).meta({
+    collaborative: true,
+  });
+
+  type CleanedType = Pick<typeof arraySchema, "_zod" | "def" | "~standard">;
+
+  const coListSchema = arraySchema as unknown as CleanedType & {
+    collaborative: true;
+    create: CoListSchema<T>["create"];
+    load: CoListSchema<T>["load"];
+    subscribe: CoListSchema<T>["subscribe"];
+    withHelpers: CoListSchema<T>["withHelpers"];
+  };
+
+  coListSchema.collaborative = true;
+
+  coListSchema.create = function (this: CoListSchema<T>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).create(...args);
+  } as CoListSchema<T>["create"];
+
+  coListSchema.load = function (this: CoListSchema<T>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).load(...args);
+  } as CoListSchema<T>["load"];
+
+  coListSchema.subscribe = function (this: CoListSchema<T>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).subscribe(...args);
+  } as CoListSchema<T>["subscribe"];
+
+  coListSchema.withHelpers = function (
+    this: CoListSchema<T>,
+    helpers: (Self: CoListSchema<T>) => object,
+  ) {
+    return { ...this, ...helpers(this) };
+  } as CoListSchema<T>["withHelpers"];
+
+  return coListSchema;
+};
+
+const coProfileDefiner = <
+  Shape extends z.core.$ZodLooseShape = Simplify<DefaultProfileShape>,
+>(
+  shape: Shape & {
+    name?: z.core.$ZodString<string>;
+    inbox?: z.core.$ZodOptional<z.core.$ZodString>;
+    inboxInvite?: z.core.$ZodOptional<z.core.$ZodString>;
+  } = {} as any,
+): CoProfileSchema<Shape> => {
+  const base = coMapDefiner({
+    ...(shape ?? {}),
+    name: z.string(),
+    inbox: z.optional(z.string()),
+    inboxInvite: z.optional(z.string()),
+  });
+  return {
+    ...base,
+    // enforce that the owner is a group
+    create: ((init: any, options: { owner: Group } | Group) => {
+      return base.create(init, options);
+    }) as CoProfileSchema<Shape>["create"],
+  } as CoProfileSchema<Shape>;
+};
+
+const coFeedDefiner = <T extends z.core.$ZodType>(
+  element: T,
+): CoFeedSchema<T> => {
+  const placeholderSchema = z.instanceof(CoFeed);
+
+  const coFeedSchema = placeholderSchema as unknown as Pick<
+    typeof placeholderSchema,
+    "_zod" | "def" | "~standard"
+  > & {
+    collaborative: true;
+    builtin: "CoFeed";
+    element: T;
+    create: CoFeedSchema<T>["create"];
+    load: CoFeedSchema<T>["load"];
+    subscribe: CoFeedSchema<T>["subscribe"];
+  };
+
+  coFeedSchema.collaborative = true;
+  coFeedSchema.builtin = "CoFeed";
+  coFeedSchema.element = element;
+
+  coFeedSchema.create = function (this: CoFeedSchema<T>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).create(...args);
+  } as CoFeedSchema<T>["create"];
+
+  coFeedSchema.load = function (this: CoFeedSchema<T>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).load(...args);
+  } as CoFeedSchema<T>["load"];
+
+  coFeedSchema.subscribe = function (this: CoFeedSchema<T>, ...args: any[]) {
+    return (zodSchemaToCoSchema(this) as any).subscribe(...args);
+  } as CoFeedSchema<T>["subscribe"];
+
+  return coFeedSchema as unknown as CoFeedSchema<T>;
+};
+
+export const coFileStreamDefiner = (): FileStreamSchema => {
+  const placeholderSchema = z.instanceof(FileStream);
+
+  const fileStreamSchema = placeholderSchema as unknown as Pick<
+    typeof placeholderSchema,
+    "_zod" | "def" | "~standard"
+  > & {
+    collaborative: true;
+    builtin: "FileStream";
+    create: (typeof FileStream)["create"];
+    createFromBlob: (typeof FileStream)["createFromBlob"];
+  };
+
+  fileStreamSchema.collaborative = true;
+  fileStreamSchema.builtin = "FileStream";
+
+  fileStreamSchema.create = function (options: any) {
+    return FileStream.create(options);
+  } as (typeof FileStream)["create"];
+
+  fileStreamSchema.createFromBlob = function (blob: Blob, options: any) {
+    return FileStream.createFromBlob(blob, options);
+  } as (typeof FileStream)["createFromBlob"];
+
+  return fileStreamSchema;
+};
+
+export const coPlainTextDefiner = (): PlainTextSchema => {
+  const placeholderSchema = z.instanceof(CoPlainText);
+
+  const plainTextSchema = placeholderSchema as unknown as Pick<
+    typeof placeholderSchema,
+    "_zod" | "def" | "~standard"
+  > & {
+    collaborative: true;
+    builtin: "CoPlainText";
+    create: PlainTextSchema["create"];
+    load: PlainTextSchema["load"];
+    subscribe: PlainTextSchema["subscribe"];
+    fromRaw: PlainTextSchema["fromRaw"];
+  };
+
+  plainTextSchema.collaborative = true;
+  plainTextSchema.builtin = "CoPlainText";
+
+  plainTextSchema.create = function (...args: any[]) {
+    return (CoPlainText as any).create(...args);
+  } as PlainTextSchema["create"];
+
+  plainTextSchema.load = function (...args: any[]) {
+    return (CoPlainText as any).load(...args);
+  } as PlainTextSchema["load"];
+
+  plainTextSchema.subscribe = function (...args: any[]) {
+    return (CoPlainText as any).subscribe(...args);
+  } as PlainTextSchema["subscribe"];
+
+  plainTextSchema.fromRaw = function (...args: any[]) {
+    return (CoPlainText as any).fromRaw(...args);
+  } as PlainTextSchema["fromRaw"];
+
+  return plainTextSchema;
+};
+
+export const coRichTextDefiner = (): RichTextSchema => {
+  const placeholderSchema = z.instanceof(CoRichText);
+
+  const richTextSchema = placeholderSchema as unknown as Pick<
+    typeof placeholderSchema,
+    "_zod" | "def" | "~standard"
+  > & {
+    collaborative: true;
+    builtin: "CoRichText";
+    create: RichTextSchema["create"];
+    load: RichTextSchema["load"];
+    subscribe: RichTextSchema["subscribe"];
+  };
+
+  richTextSchema.collaborative = true;
+  richTextSchema.builtin = "CoRichText";
+
+  richTextSchema.create = function (...args: any[]) {
+    return (CoRichText as any).create(...args);
+  } as RichTextSchema["create"];
+
+  richTextSchema.load = function (...args: any[]) {
+    return (CoRichText as any).load(...args);
+  } as RichTextSchema["load"];
+
+  richTextSchema.subscribe = function (...args: any[]) {
+    return (CoRichText as any).subscribe(...args);
+  } as RichTextSchema["subscribe"];
+
+  return richTextSchema;
+};
+
+export const co = {
+  map: coMapDefiner,
+  record: coRecordDefiner,
+  list: coListDefiner,
+  feed: coFeedDefiner,
+  plainText: coPlainTextDefiner,
+  richText: coRichTextDefiner,
+  fileStream: coFileStreamDefiner,
+  image: (): typeof ImageDefinition => {
+    return ImageDefinition;
+  },
+  account: coAccountDefiner,
+  profile: coProfileDefiner,
+};

package/src/implementation/zodSchema/zodSchema.ts

@@ -0,0 +1,116 @@
+import { LocalNode, RawAccount } from "cojson";
+import z from "zod/v4";
+import {
+  Account,
+  AccountClass,
+  CoValue,
+  CoValueClass,
+  CoValueFromRaw,
+  RefsToResolve,
+  RefsToResolveStrict,
+  Resolved,
+  Simplify,
+} from "../../internal.js";
+import { AnyAccountSchema } from "./schemaTypes/AccountSchema.js";
+import { AnyCoFeedSchema } from "./schemaTypes/CoFeedSchema.js";
+import { AnyCoListSchema } from "./schemaTypes/CoListSchema.js";
+import { AnyCoMapSchema, CoMapInitZod } from "./schemaTypes/CoMapSchema.js";
+import { AnyCoRecordSchema } from "./schemaTypes/CoRecordSchema.js";
+import { FileStreamSchema } from "./schemaTypes/FileStreamSchema.js";
+import { PlainTextSchema } from "./schemaTypes/PlainTextSchema.js";
+import { RichTextSchema } from "./schemaTypes/RichTextSchema.js";
+import { InstanceOfSchema } from "./typeConverters/InstanceOfSchema.js";
+import { InstanceOfSchemaCoValuesNullable } from "./typeConverters/InstanceOfSchemaCoValuesNullable.js";
+import { InstanceOrPrimitiveOfSchema } from "./typeConverters/InstanceOrPrimitiveOfSchema.js";
+import { InstanceOrPrimitiveOfSchemaCoValuesNullable } from "./typeConverters/InstanceOrPrimitiveOfSchemaCoValuesNullable.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 FullyOrPartiallyLoaded<S extends z.core.$ZodType | CoValueClass> =
+  InstanceOrPrimitiveOfSchema<S> extends CoValue
+    ? NonNullable<InstanceOrPrimitiveOfSchemaCoValuesNullable<S>>
+    : InstanceOrPrimitiveOfSchema<S>;
+
+export type ZodPrimitiveSchema =
+  | z.core.$ZodString
+  | z.core.$ZodNumber
+  | z.core.$ZodBoolean
+  | z.core.$ZodNull
+  | 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.core.$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;
+}
+
+export type CoValueOrZodSchema = CoValueClass | AnyCoSchema;
+
+export type CoValueClassFromZodSchema<S extends z.core.$ZodType> = CoValueClass<
+  InstanceOfSchema<S>
+> &
+  CoValueFromRaw<InstanceOfSchema<S>> &
+  (S extends AnyAccountSchema ? AccountClassEssentials : {});
+
+type AccountClassEssentials = {
+  fromRaw: <A extends Account>(this: AccountClass<A>, raw: RawAccount) => A;
+  fromNode: <A extends Account>(this: AccountClass<A>, node: LocalNode) => A;
+};
+
+type AnyCoSchema =
+  | AnyCoMapSchema
+  | AnyAccountSchema
+  | AnyCoRecordSchema
+  | AnyCoListSchema
+  | AnyCoFeedSchema
+  | AnyCoUnionSchema
+  | PlainTextSchema
+  | RichTextSchema
+  | FileStreamSchema;
+
+export type Loaded<
+  T extends CoValueClass | AnyCoSchema,
+  R extends ResolveQuery<T> = true,
+> = Resolved<NonNullable<InstanceOfSchemaCoValuesNullable<T>>, R>;
+
+export type ResolveQuery<T extends CoValueClass | AnyCoSchema> = RefsToResolve<
+  NonNullable<InstanceOfSchemaCoValuesNullable<T>>
+>;
+
+export type ResolveQueryStrict<
+  T extends CoValueClass | AnyCoSchema,
+  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;

package/src/internal.ts

@@ -1,13 +1,51 @@
 export * from "./coValues/interfaces.js";
+export * from "./coValues/CoValueBase.js";
 export * from "./implementation/inspect.js";
 export * from "./implementation/symbols.js";
 
+export * from "./coValues/registeredSchemas.js";
+
+export * from "./coValues/coMap.js";
+export * from "./coValues/coList.js";
+export * from "./coValues/coFeed.js";
+export * from "./coValues/account.js";
+export * from "./coValues/group.js";
+export * from "./coValues/profile.js";
+export * from "./coValues/inbox.js";
+export * from "./coValues/coPlainText.js";
+export * from "./coValues/coRichText.js";
+export * from "./coValues/schemaUnion.js";
+
 export * from "./coValues/deepLoading.js";
 export * from "./implementation/anonymousJazzAgent.js";
+export * from "./implementation/activeAccountContext.js";
 export * from "./implementation/errors.js";
 export * from "./implementation/refs.js";
 export * from "./implementation/schema.js";
+export * from "./subscribe/SubscriptionScope.js";
+export * from "./subscribe/types.js";
 export * from "./subscribe/index.js";
+export * from "./lib/cache.js";
 export * from "./implementation/createContext.js";
 
+export * from "./types.js";
+export * from "./implementation/zodSchema/zodSchema.js";
+export * from "./implementation/zodSchema/zodCo.js";
+export * from "./implementation/zodSchema/schemaTypes/CoMapSchema.js";
+export * from "./implementation/zodSchema/schemaTypes/CoRecordSchema.js";
+export * from "./implementation/zodSchema/schemaTypes/CoListSchema.js";
+export * from "./implementation/zodSchema/schemaTypes/CoFeedSchema.js";
+export * from "./implementation/zodSchema/schemaTypes/AccountSchema.js";
+export * from "./implementation/zodSchema/schemaTypes/FileStreamSchema.js";
+export * from "./implementation/zodSchema/schemaTypes/PlainTextSchema.js";
+export * from "./implementation/zodSchema/typeConverters/InstanceOrPrimitiveOfSchema.js";
+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 "./coValues/extensions/imageDef.js";
+
+export * from "./implementation/ContextManager.js";
+
 import "./implementation/devtoolsFormatters.js";

package/src/subscribe/SubscriptionScope.ts

@@ -1,6 +1,8 @@
 import type { LocalNode, RawCoValue } from "cojson";
-import type { CoFeed, CoList, CoMap } from "../exports.js";
 import {
+  CoFeed,
+  CoList,
+  CoMap,
   type CoValue,
   type ID,
   type RefEncoded,

package/src/subscribe/utils.ts

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

package/src/testing.ts

@@ -4,17 +4,14 @@ import { PureJSCrypto } from "cojson/dist/crypto/PureJSCrypto";
 import {
   Account,
   AccountClass,
+  type AnonymousJazzAgent,
   AuthCredentials,
-  JazzContextManagerAuthProps,
-} from "./exports.js";
-import {
+  type CoValueClass,
+  CoValueFromRaw,
   JazzContextManager,
+  JazzContextManagerAuthProps,
   JazzContextManagerBaseProps,
-} from "./implementation/ContextManager.js";
-import { activeAccountContext } from "./implementation/activeAccountContext.js";
-import {
-  type AnonymousJazzAgent,
-  type CoValueClass,
+  activeAccountContext,
   createAnonymousJazzContext,
   createJazzContext,
   randomSessionProvider,
@@ -76,13 +73,13 @@ export function getPeerConnectedToTestSyncServer() {
 const SecretSeedMap = new Map<string, Uint8Array>();
 let isMigrationActive = false;
 
-export async function createJazzTestAccount<Acc extends Account>(options?: {
+export async function createJazzTestAccount<A extends Account>(options?: {
   isCurrentActiveAccount?: boolean;
-  AccountSchema?: CoValueClass<Acc>;
+  AccountSchema?: CoValueClass<A>;
   creationProps?: Record<string, unknown>;
-}): Promise<Acc> {
+}): Promise<A> {
   const AccountSchema = (options?.AccountSchema ??
-    Account) as unknown as TestAccountSchema<Acc>;
+    Account) as unknown as TestAccountSchema<Account>;
   const peers = [];
   if (syncServer.current) {
     peers.push(getPeerConnectedToTestSyncServer());
@@ -134,7 +131,7 @@ export async function createJazzTestAccount<Acc extends Account>(options?: {
     activeAccountContext.set(account);
   }
 
-  return account;
+  return account as A;
 }
 
 export function setActiveAccount(account: Account) {
@@ -155,7 +152,7 @@ export async function createJazzTestGuest() {
 export type TestJazzContextManagerProps<Acc extends Account> =
   JazzContextManagerBaseProps<Acc> & {
     defaultProfileName?: string;
-    AccountSchema?: AccountClass<Acc>;
+    AccountSchema?: AccountClass<Acc> & CoValueFromRaw<Acc>;
     isAuthenticated?: boolean;
   };
 
@@ -194,7 +191,8 @@ export class TestJazzContextManager<
 
     context.updateContext(
       {
-        AccountSchema: account.constructor as AccountClass<Acc>,
+        AccountSchema: account.constructor as AccountClass<Acc> &
+          CoValueFromRaw<Acc>,
         ...props,
       },
       {
@@ -247,7 +245,7 @@ export class TestJazzContextManager<
       );
     }
 
-    const context = await createJazzContext<Acc>({
+    const context = await createJazzContext({
       credentials: authProps?.credentials,
       defaultProfileName: props.defaultProfileName,
       newAccountProps: authProps?.newAccountProps,