index.ts

/// <reference types="vite/client" />

import { AsyncLocalStorage } from "node:async_hooks";
import {
  DataModelFromSchemaDefinition,
  DefaultFunctionArgs,
  DocumentByName,
  FunctionReference,
  FunctionReturnType,
  GenericActionCtx,
  GenericDataModel,
  GenericDocument,
  GenericMutationCtx,
  GenericQueryCtx,
  GenericSchema,
  HttpRouter,
  OptionalRestArgs,
  PublicHttpAction,
  RegisteredAction,
  RegisteredMutation,
  RegisteredQuery,
  SchemaDefinition,
  SystemDataModel,
  UserIdentity,
  actionGeneric,
  getFunctionAddress,
  httpActionGeneric,
  makeFunctionReference,
  mutationGeneric,
  queryGeneric,
} from "convex/server";
import {
  GenericId,
  JSONValue,
  Value,
  compareValues,
  convexToJson,
  getConvexSize,
  getDocumentSize,
  jsonToConvex,
} from "convex/values";
import { HeadroomTracker, TransactionMetrics } from "./headroom.js";

type FilterJson =
  | { $eq: [FilterJson, FilterJson] }
  | { $field: string }
  | JSONValue;

type QueryOperator = { filter: FilterJson } | { limit: number };
type Source =
  | { type: "FullTableScan"; tableName: string; order: "asc" | "desc" | null }
  | {
      type: "IndexRange";
      indexName: string;
      range: ReadonlyArray<SerializedRangeExpression>;
      order: "asc" | "desc" | null;
    }
  | {
      type: "Search";
      indexName: string;
      filters: ReadonlyArray<SerializedSearchFilter>;
    };

type SerializedQuery = {
  source: Source;
  operators: Array<QueryOperator>;
};

type SerializedRangeExpression = {
  type: "Eq" | "Gt" | "Gte" | "Lt" | "Lte";
  fieldPath: string;
  value: JSONValue;
};

type SerializedSearchFilter =
  | {
      type: "Search";
      fieldPath: string;
      value: string;
    }
  | {
      type: "Eq";
      fieldPath: string;
      value: JSONValue;
    };

type ScheduledFunction = DocumentByName<
  SystemDataModel,
  "_scheduled_functions"
>;

type StoredDocument = GenericDocument & {
  _id: DocumentId;
  _creationTime: number;
};

type Index = {
  indexDescriptor: string;
  fields: string[];
};

type VectorIndex = {
  indexDescriptor: string;
  vectorField: string;
  dimensions: number;
  filterFields: string[];
};

type QueryId = number;

type TableName = string;

type DocumentId = GenericId<TableName>;

const ROOT_COMPONENT_PATH = "";

class NestedLock {
  private _lock: Promise<void> | null = null;
  private _release: (() => void) | null = null;

  async acquire() {
    while (this._lock !== null) {
      await this._lock;
    }
    this._lock = new Promise((resolve) => {
      this._release = resolve;
    });
  }

  release() {
    this._lock = null;
    this._release!();
    this._release = null;
  }
}

const nestedTxStorage = new AsyncLocalStorage<NestedLock>();

class DatabaseFake {
  private _componentPath: string;
  private _documents: Record<DocumentId, StoredDocument> = {};
  private _storage: Record<DocumentId, Blob> = {};
  private _nextQueryId: QueryId = 1;
  private _nextDocId: number = 10000;
  private _lastCreationTime: number = 0;
  private _queryResults: Record<QueryId, Array<GenericDocument>> = {};
  // TODO: Make this more robust and cleaner
  jobListener: (jobId: string) => void = () => {};
  // Pending writes for each level of transaction nesting.
  // Whenever a child mutation begins, a new level is created for all
  // components.
  // - When a mutation performs updates, they are staged in the last
  //   level of writes for the mutation's component.
  // - When a child mutation commits, the writes are merged up one level.
  // - When the top-level mutation commits, the writes are applied to the
  //   database.
  // - When a mutation is rolled back, last level of writes is discarded.
  private _writes: Array<Record<DocumentId, StoredDocument | null>> = [];

  private _schema: {
    schemaValidation: boolean;
    tables: Map<
      string,
      {
        indexes: Index[];
        vectorIndexes: VectorIndex[];
        documentType: ValidatorJSON;
      }
    >;
  } | null;

  constructor(
    schema: SchemaDefinition<GenericSchema, boolean> | null,
    componentPath: string,
  ) {
    this._schema =
      schema === null
        ? null
        : {
            schemaValidation: (schema as any).schemaValidation,
            tables: new Map(
              Object.entries(schema.tables).map(([name, tableSchema]) => [
                name,
                (tableSchema as any).export(),
              ]),
            ),
          };
    this._componentPath = componentPath;

    this.validateSchema();
  }

  // Called by TransactionManager.begin.
  startTransaction() {
    this._writes.push({});
  }

  get<Table extends TableName>(
    tableName: Table | undefined,
    id: GenericId<Table>,
  ) {
    this._validateId(tableName, id);

    for (let i = this._writes.length - 1; i >= 0; i--) {
      const write = this._writes[i][id];
      if (write !== undefined) {
        return write;
      }
    }

    return this._documents[id] ?? null;
  }

  // Note that this is not the format the real backend
  // uses for IDs.
  private _generateId<TableName extends string>(
    table: TableName,
  ): GenericId<TableName> {
    const id = this._nextDocId.toString() + ";" + table;
    this._nextDocId += 1;
    return id as GenericId<TableName>;
  }

  storeFile(storageId: GenericId<"_storage">, blob: Blob) {
    this._storage[storageId] = blob;
  }

  getFile(storageId: GenericId<"_storage">) {
    if (this.get("_storage", storageId) === null) {
      return null;
    }
    return this._storage[storageId];
  }

  private _addWrite(id: DocumentId, newValue: StoredDocument | null) {
    if (this._writes.length === 0) {
      throw new Error(`Write outside of transaction ${id}`);
    }
    this._writes[this._writes.length - 1][id] = newValue;
  }

  insert<Table extends TableName>(table: Table, value: any) {
    this._validate(table, value);
    const _id = this._generateId(table);
    const now = Date.now();
    const _creationTime =
      now <= this._lastCreationTime ? this._lastCreationTime + 0.001 : now;
    this._lastCreationTime = _creationTime;
    this._addWrite(_id, { ...value, _id, _creationTime });
    return _id;
  }

  patch<Table extends TableName>(
    tableName: Table | undefined,
    id: GenericId<Table>,
    value: Record<string, any>,
  ) {
    this._validateId(tableName, id);

    if (typeof value !== "object") {
      throw new Error(
        `Invalid argument \`value\` in \`db.patch\`, expected object but got '${typeof value}': ${value as any}`,
      );
    }

    const document = this.get(tableName, id);
    if (document === null) {
      throw new Error(`Patch on non-existent document with ID "${id}"`);
    }
    const { _id, _creationTime, ...fields } = document;
    if (value._id !== undefined && value._id !== _id) {
      throw new Error(
        `Provided \`_id\` field value "${value._id}" ` +
          `does not match the document ID "${_id}"`,
      );
    }
    if (
      value._creationTime !== undefined &&
      value._creationTime !== _creationTime
    ) {
      throw new Error(
        `Provided \`_creationTime\` field value ${value._creationTime} ` +
          `does not match the document's creation time ${_creationTime}`,
      );
    }
    delete value["_id"];
    delete value["_creationTime"];
    const convexValue: any = {};
    for (const [key, v] of Object.entries(value)) {
      convexValue[key] = evaluateValue(v);
    }
    const merged = { ...fields, ...convexValue };
    this._validate(tableNameFromId(_id as string)!, merged);
    this._addWrite(id, { _id, _creationTime, ...merged });
  }

  replace<Table extends TableName>(
    tableName: Table | undefined,
    id: DocumentId,
    value: Record<string, any>,
  ) {
    this._validateId(tableName, id);

    if (typeof value !== "object") {
      throw new Error(
        `Invalid argument \`value\` in \`db.replace\`, expected object but got '${typeof value}': ${value as any}`,
      );
    }

    const document = this.get(tableName, id);
    if (document === null) {
      throw new Error(`Replace on non-existent document with ID "${id}"`);
    }
    if (value._id !== undefined && value._id !== document._id) {
      throw new Error(
        `Provided \`_id\` field value "${value._id}" ` +
          `does not match the document ID "${document._id}"`,
      );
    }
    if (
      value._creationTime !== undefined &&
      value._creationTime !== document._creationTime
    ) {
      throw new Error(
        `Provided \`_creationTime\` field value ${value._creationTime} ` +
          `does not match the document's creation time ${document._creationTime}`,
      );
    }
    delete value["_id"];
    delete value["_creationTime"];
    const convexValue: any = {};
    for (const [key, v] of Object.entries(value)) {
      convexValue[key] = evaluateValue(v);
    }
    this._validate(tableNameFromId(document._id as string)!, convexValue);
    this._addWrite(id, {
      ...convexValue,
      _id: document._id,
      _creationTime: document._creationTime,
    });
  }

  delete<Table extends TableName>(
    tableName: Table | undefined,
    id: GenericId<Table>,
  ) {
    this._validateId(tableName, id);

    const document = this.get(tableName, id);
    if (document === null) {
      throw new Error("Delete on non-existent doc");
    }
    this._addWrite(id, null);
  }

  private _validateId(
    expectedTableName: unknown,
    id: unknown,
  ): asserts id is DocumentId {
    if (typeof id !== "string") {
      throw new Error(
        `Invalid argument \`id\`, expected string but got '${typeof id}': ${id as any}`,
      );
    }

    if (expectedTableName === undefined) {
      return;
    }

    if (typeof expectedTableName !== "string") {
      throw new Error(
        `Invalid argument \`tableName\`, expected string but got '${typeof expectedTableName}': ${expectedTableName as any}`,
      );
    }

    const actualTableName = tableNameFromId(id);
    if (actualTableName === null) {
      throw new Error(
        `Invalid argument \`id\`, expected ID value but got '${id}'`,
      );
    }

    if (actualTableName !== expectedTableName) {
      throw new Error(
        `Invalid argument \`id\`, expected ID in table '${expectedTableName}' but got ID in table '${actualTableName}'`,
      );
    }
  }

  commit() {
    const lastWrites = this._writes.pop();
    if (!lastWrites) {
      throw new Error("Transaction already committed or rolled back");
    }
    for (const [id, write] of Object.entries(lastWrites)) {
      const _id = id as DocumentId;
      if (this._writes.length === 0) {
        if (write === null) {
          delete this._documents[_id];
        } else {
          this._documents[_id] = write;
        }
      } else {
        this._addWrite(_id, write);
      }
    }
  }

  rollbackWrites() {
    if (this._writes.length === 0) {
      throw new Error("Transaction already committed or rolled back");
    }
    this._writes.pop();
  }

  _validate(tableName: string, doc: GenericDocument) {
    if (this._schema === null || !this._schema.schemaValidation) {
      return;
    }
    const validator = this._schema.tables.get(tableName)?.documentType;
    if (validator === undefined) {
      return;
    }
    validateValidator(validator, doc);
  }

  startQuery(query: SerializedQuery) {
    const id = this._nextQueryId;
    const results = this._evaluateQuery(query);
    this._queryResults[id] = results;
    this._nextQueryId += 1;
    return id;
  }

  queryNext(queryId: QueryId) {
    const results = this._queryResults[queryId];
    if (results === undefined) {
      throw new Error("Bad queryId");
    }
    if (results.length === 0) {
      return { value: null, done: true };
    } else {
      return { value: results.shift()!, done: false };
    }
  }

  private _encodeCursor(doc: GenericDocument, fieldPaths: string[]): string {
    const keys = fieldPaths.map((fp) => evaluateFieldPath(fp, doc));
    // null sentinel in the array = undefined field value.
    // Everything else is JSON.stringify(convexToJson(k)) producing a string,
    // so null vs string is unambiguous on decode.
    return JSON.stringify(
      keys.map((k) =>
        k === undefined ? null : JSON.stringify(convexToJson(k)),
      ),
    );
  }

  private _docIsAfterCursor(
    doc: GenericDocument,
    cursorKeys: any[],
    fieldPaths: string[],
    order: "asc" | "desc",
  ): number {
    const orderMultiplier = order === "asc" ? 1 : -1;
    for (let i = 0; i < fieldPaths.length; i++) {
      const docVal = evaluateFieldPath(fieldPaths[i], doc);
      const cursorVal = cursorKeys[i];
      const cmp = compareValues(docVal, cursorVal);
      if (cmp !== 0) return cmp * orderMultiplier;
    }
    return 0;
  }

  paginate({
    query,
    cursor,
    endCursor,
    pageSize,
    maximumRowsRead,
    maximumBytesRead,
  }: {
    query: SerializedQuery;
    cursor: string | null;
    endCursor?: string | null;
    pageSize: number;
    maximumRowsRead?: number | null;
    maximumBytesRead?: number | null;
  }) {
    const { sortedDocs, filterFn, fieldPathsToSortBy, order } =
      this._resolveQuerySource(query);

    // Always include _id in sort key fields for uniqueness
    const cursorFields = fieldPathsToSortBy.includes("_id")
      ? fieldPathsToSortBy
      : [...fieldPathsToSortBy, "_id"];

    const parseCursor = (c: string | null | undefined) => {
      if (!c) return null;
      // null = undefined field; strings are JSON.stringify(convexToJson(v))
      return (JSON.parse(c) as any[]).map((k) =>
        k === null ? undefined : jsonToConvex(JSON.parse(k)),
      );
    };

    const cursorKeys = cursor !== "_end_cursor" ? parseCursor(cursor) : null;
    // _end_cursor as endCursor means "to the end of the range" — no boundary
    const hasEndBoundary = !!endCursor && endCursor !== "_end_cursor";
    const endCursorKeys = hasEndBoundary ? parseCursor(endCursor) : null;

    const page: GenericDocument[] = [];
    // null cursor = start of range
    let isInPage = cursor === null;
    let isDone = false;
    let continueCursor: string = "_end_cursor";
    let splitCursor: string | null = null;
    let pageStatus: "SplitRecommended" | "SplitRequired" | null = null;
    let rowsRead = 0;
    let bytesRead = 0;
    const readDocCursors: string[] = [];

    for (let i = 0; i < sortedDocs.length; i++) {
      const doc = sortedDocs[i];
      const isLastDoc = i === sortedDocs.length - 1;

      if (!isInPage) {
        if (cursor === "_end_cursor") {
          // Past the end — no doc is after the end cursor
          continue;
        } else if (cursorKeys) {
          // Sort-key-based cursor: skip docs at or before cursor position
          if (
            this._docIsAfterCursor(doc, cursorKeys, cursorFields, order) > 0
          ) {
            isInPage = true;
          } else {
            continue;
          }
        }
      }

      // endCursor: stop when we reach or pass this position (inclusive boundary)
      if (hasEndBoundary && endCursorKeys) {
        const cmp = this._docIsAfterCursor(
          doc,
          endCursorKeys,
          cursorFields,
          order,
        );
        if (cmp >= 0) {
          // At or past end cursor — include this doc if exactly at cursor, then stop
          if (cmp === 0) {
            rowsRead += 1;
            bytesRead += getDocumentSize(doc);
            readDocCursors.push(this._encodeCursor(doc, cursorFields));
            if (filterFn(doc)) {
              page.push(doc);
            }
            continueCursor = this._encodeCursor(doc, cursorFields);
          }
          break;
        }
      }

      rowsRead += 1;
      bytesRead += getDocumentSize(doc);
      readDocCursors.push(this._encodeCursor(doc, cursorFields));

      // Check bandwidth limits
      let hitLimit = false;
      if (maximumRowsRead && rowsRead >= maximumRowsRead) {
        hitLimit = true;
      }
      if (maximumBytesRead && bytesRead >= maximumBytesRead) {
        hitLimit = true;
      }

      if (filterFn(doc)) {
        page.push(doc);
      }

      if (hitLimit) {
        pageStatus = "SplitRequired";
        continueCursor = this._encodeCursor(doc, cursorFields);
        break;
      }

      if (!endCursor && page.length >= pageSize) {
        if (isLastDoc) {
          // No more docs — don't set continueCursor so isDone becomes true
        } else {
          continueCursor = this._encodeCursor(doc, cursorFields);
        }
        break;
      }
    }

    if (continueCursor === "_end_cursor") {
      isDone = true;
    }

    // Compute splitCursor at midpoint when limits are hit
    if (pageStatus === "SplitRequired" && readDocCursors.length >= 2) {
      const midIdx = Math.floor((readDocCursors.length - 1) / 2);
      splitCursor = readDocCursors[midIdx];
    } else if (
      pageStatus === null &&
      rowsRead > pageSize + 1 &&
      readDocCursors.length >= 2
    ) {
      // Recommend split when we had to scan significantly more rows than pageSize
      pageStatus = "SplitRecommended";
      const midIdx = Math.floor((readDocCursors.length - 1) / 2);
      splitCursor = readDocCursors[midIdx];
    }

    return {
      page,
      isDone,
      continueCursor,
      splitCursor,
      pageStatus,
    };
  }

  private _iterateDocs(
    tableName: string,
    callback: (doc: GenericDocument) => void,
  ) {
    const isInTable = (id: string) => tableNameFromId(id) === tableName;
    const ids = new Set(Object.keys(this._documents).filter(isInTable));
    for (let i = 0; i < this._writes.length; i++) {
      for (const id of Object.keys(this._writes[i]).filter(isInTable)) {
        ids.add(id);
      }
    }
    for (const id of ids) {
      const document = this.get(tableName, id as DocumentId);
      if (document !== null) {
        callback(document);
      }
    }
  }

  /**
   * Resolves the query source: loads documents matching the source (table scan,
   * index range, or search), sorts them, and returns along with the compiled
   * operator filter function. This separates "rows entering the pipeline"
   * (sortedDocs) from "rows exiting" (after filterFn), which is needed for
   * maximumRowsRead tracking in paginate().
   */
  private _resolveQuerySource(query: SerializedQuery): {
    sortedDocs: GenericDocument[];
    filterFn: (doc: GenericDocument) => boolean;
    limit: number | null;
    fieldPathsToSortBy: string[];
    order: "asc" | "desc";
  } {
    const source = query.source;
    const results: GenericDocument[] = [];
    let fieldPathsToSortBy: string[];
    let order: "asc" | "desc";
    switch (source.type) {
      case "FullTableScan": {
        const tableName = source.tableName;
        this._iterateDocs(tableName, (doc) => {
          results.push(doc);
        });
        order = source.order ?? "asc";
        fieldPathsToSortBy = ["_creationTime"];

        break;
      }
      case "IndexRange": {
        const [tableName, indexName] = source.indexName.split(".");
        order = source.order ?? "asc";
        let fields;
        if (indexName === "by_creation_time") {
          fields = ["_creationTime", "_id"];
        } else if (indexName === "by_id") {
          fields = ["_id"];
        } else {
          const indexes = this._schema?.tables.get(tableName)?.indexes;
          const index = indexes?.find(
            ({ indexDescriptor }: { indexDescriptor: string }) =>
              indexDescriptor === indexName,
          );
          if (index === undefined) {
            throw new Error(
              `Cannot use index "${indexName}" for table "${tableName}" because ` +
                `it is not declared in the schema.`,
            );
          }
          fields = index.fields.concat(["_creationTime", "_id"]);
        }
        fieldPathsToSortBy = fields;

        validateIndexRangeExpression(source, fields);

        this._iterateDocs(tableName, (doc) => {
          if (
            source.range.every((filter) => evaluateRangeFilter(doc, filter))
          ) {
            results.push(doc);
          }
        });
        break;
      }
      case "Search": {
        const [tableName] = source.indexName.split(".");
        this._iterateDocs(tableName, (doc) => {
          if (
            source.filters.every((filter) => evaluateSearchFilter(doc, filter))
          ) {
            results.push(doc);
          }
        });
        fieldPathsToSortBy = [];
        order = "asc";
        break;
      }
    }

    const filters = query.operators
      .filter(
        (operator): operator is { filter: FilterJson } => "filter" in operator,
      )
      .map((operator) => operator.filter);

    const limit =
      query.operators.filter(
        (operator): operator is { limit: number } => "limit" in operator,
      )[0] ?? null;

    results.sort((a, b) => {
      const orderMultiplier = order === "asc" ? 1 : -1;
      let v = 0;
      for (const fp of fieldPathsToSortBy) {
        v = compareValues(evaluateFieldPath(fp, a), evaluateFieldPath(fp, b));
        if (v !== 0) {
          return v * orderMultiplier;
        }
      }
      return v * orderMultiplier;
    });

    const filterFn = (doc: GenericDocument) =>
      filters.every((f) => evaluateFilter(doc, f));

    return {
      sortedDocs: results,
      filterFn,
      limit: limit?.limit ?? null,
      fieldPathsToSortBy,
      order,
    };
  }

  private _evaluateQuery(query: SerializedQuery): Array<GenericDocument> {
    const { sortedDocs, filterFn, limit } = this._resolveQuerySource(query);
    let results = sortedDocs.filter(filterFn);
    if (limit !== null) {
      results = results.slice(0, limit);
    }
    return results;
  }

  jobFinished(jobId: string) {
    this.jobListener(jobId);
  }

  vectorSearch(
    tableAndIndexName: string,
    vector: number[],
    expressions: SerializedRangeExpression[],
    limit: number,
  ) {
    const results: GenericDocument[] = [];
    const [tableName, indexName] = tableAndIndexName.split(".");
    this._iterateDocs(tableName, (doc) => {
      if (expressions.every((filter) => evaluateFilter(doc, filter))) {
        results.push(doc);
      }
    });
    const vectorIndexes = this._schema?.tables.get(tableName)?.vectorIndexes;
    const vectorIndex = vectorIndexes?.find(
      ({ indexDescriptor }: { indexDescriptor: string }) =>
        indexDescriptor === indexName,
    );
    if (vectorIndex === undefined) {
      throw new Error(
        `Cannot use vector index "${indexName}" for table "${tableName}" because ` +
          `it is not declared in the schema.`,
      );
    }
    const { vectorField } = vectorIndex;
    const idsAndScores = results.map((doc) => {
      const score = cosineSimilarity(vector, doc[vectorField] as number[]);
      return { _id: doc._id, _score: score };
    });
    idsAndScores.sort((a, b) => b._score - a._score);
    return idsAndScores.slice(0, limit);
  }

  validateSchema() {
    this._schema?.tables.forEach((table, tableName) => {
      if (!isValidIdentifier(tableName)) {
        throw new Error(
          `Table names must be valid identifiers, got "${tableName}"`,
        );
      }

      validateFieldNames(table.documentType);

      table.indexes.forEach(({ indexDescriptor }) => {
        if (!isValidIdentifier(indexDescriptor)) {
          throw new Error(
            `Index names must be valid identifiers, got "${indexDescriptor}"`,
          );
        }
      });
    });
  }

  get componentPath() {
    return this._componentPath;
  }
}

function tableNameFromId(id: string) {
  const parts = id.split(";");
  if (parts.length !== 2) {
    return null;
  }
  return id.split(";")[1];
}

function isSimpleObject(value: unknown) {
  const isObject = typeof value === "object";
  const prototype = Object.getPrototypeOf(value);
  const isSimple =
    prototype === null ||
    prototype === Object.prototype ||
    // Objects generated from other contexts (e.g. across Node.js `vm` modules) will not satisfy the previous
    // conditions but are still simple objects.
    prototype?.constructor?.name === "Object";
  return isObject && isSimple;
}

function evaluateFieldPath(fieldPath: string, document: any) {
  const pathParts = fieldPath.split(".");
  let result: Value | undefined = document;
  for (const p of pathParts) {
    result =
      result !== undefined && result !== null && isSimpleObject(result)
        ? (result as any)[p]
        : undefined;
  }
  return result;
}

function evaluateFilter(
  document: GenericDocument,
  filter: any,
): Value | undefined {
  if (filter.$eq !== undefined) {
    return (
      compareValues(
        evaluateFilter(document, filter.$eq[0]),
        evaluateFilter(document, filter.$eq[1]),
      ) === 0
    );
  }
  if (filter.$neq !== undefined) {
    return (
      compareValues(
        evaluateFilter(document, filter.$neq[0]),
        evaluateFilter(document, filter.$neq[1]),
      ) !== 0
    );
  }
  if (filter.$and !== undefined) {
    return filter.$and.every((child: any) => evaluateFilter(document, child));
  }
  if (filter.$or !== undefined) {
    return filter.$or.some((child: any) => evaluateFilter(document, child));
  }
  if (filter.$not !== undefined) {
    return !evaluateFilter(document, filter.$not);
  }
  if (filter.$gt !== undefined) {
    return (
      evaluateFilter(document, filter.$gt[0])! >
      evaluateFilter(document, filter.$gt[1])!
    );
  }
  if (filter.$gte !== undefined) {
    return (
      evaluateFilter(document, filter.$gte[0])! >=
      evaluateFilter(document, filter.$gte[1])!
    );
  }
  if (filter.$lt !== undefined) {
    return (
      evaluateFilter(document, filter.$lt[0])! <
      evaluateFilter(document, filter.$lt[1])!
    );
  }
  if (filter.$lte !== undefined) {
    return (
      evaluateFilter(document, filter.$lte[0])! <=
      evaluateFilter(document, filter.$lte[1])!
    );
  }
  if (filter.$add !== undefined) {
    return (
      (evaluateFilter(document, filter.$add[0]) as number) +
      (evaluateFilter(document, filter.$add[1]) as number)
    );
  }
  if (filter.$sub !== undefined) {
    return (
      (evaluateFilter(document, filter.$sub[0]) as number) -
      (evaluateFilter(document, filter.$sub[1]) as number)
    );
  }
  if (filter.$mul !== undefined) {
    return (
      (evaluateFilter(document, filter.$mul[0]) as number) *
      (evaluateFilter(document, filter.$mul[1]) as number)
    );
  }
  if (filter.$div !== undefined) {
    return (
      (evaluateFilter(document, filter.$div[0]) as number) /
      (evaluateFilter(document, filter.$div[1]) as number)
    );
  }
  if (filter.$mod !== undefined) {
    return (
      (evaluateFilter(document, filter.$mod[0]) as number) %
      (evaluateFilter(document, filter.$mod[1]) as number)
    );
  }
  if (filter.$field !== undefined) {
    return evaluateFieldPath(filter.$field, document);
  }
  if (filter.$literal !== undefined) {
    return evaluateValue(filter.$literal);
  }
  throw new Error(`not implemented: ${JSON.stringify(filter)}`);
}

function evaluateRangeFilter(
  document: GenericDocument,
  expr: SerializedRangeExpression,
) {
  const result = evaluateFieldPath(expr.fieldPath, document);
  const value = evaluateValue(expr.value);
  switch (expr.type) {
    case "Eq":
      return compareValues(result, value) === 0;
    case "Gt":
      return compareValues(result, value) > 0;
    case "Gte":
      return compareValues(result, value) >= 0;
    case "Lt":
      return compareValues(result, value) < 0;
    case "Lte":
      return compareValues(result, value) <= 0;
  }
}

function evaluateValue(value: JSONValue) {
  if (typeof value === "object" && value !== null && "$undefined" in value) {
    return undefined;
  }
  return jsonToConvex(value);
}

function evaluateSearchFilter(
  document: GenericDocument,
  filter: SerializedSearchFilter,
) {
  const result = evaluateFieldPath(filter.fieldPath, document);
  switch (filter.type) {
    case "Eq":
      return compareValues(result, evaluateValue(filter.value)) === 0;
    case "Search": {
      const queryTerms = filter.value
        .toLowerCase()
        .split(/\s+/)
        .filter((term) => term.length > 0);
      const documentWords = (result as string)
        .split(/\s+/)
        .map((word) => word.toLowerCase());
      return queryTerms.some((queryTerm) =>
        documentWords.some((word) => word.startsWith(queryTerm)),
      );
    }
  }
}

function cosineSimilarity(vecA: number[], vecB: number[]): number {
  let dotProduct = 0.0;
  let normA = 0.0;
  let normB = 0.0;

  for (let i = 0; i < vecA.length; i++) {
    dotProduct += vecA[i] * vecB[i];
    normA += vecA[i] * vecA[i];
    normB += vecB[i] * vecB[i];
  }

  if (normA === 0 || normB === 0) {
    return 0;
  } else {
    return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
  }
}

// An index range expression is always a chained list of:
//
// 0 or more equality expressions defined with .eq.
// [Optionally] A lower bound expression defined with .gt or .gte.
// [Optionally] An upper bound expression defined with .lt or .lte.
function validateIndexRangeExpression(
  source: Source & { type: "IndexRange" },
  fields: string[],
) {
  let fieldIndex = 0;
  let state: "eq" | "gt" | "lt" | "done" = "eq";
  for (const [filterIndex, filter] of source.range.entries()) {
    if (state === "done") {
      throw new Error(
        `Incorrect operator used in \`withIndex\`, cannot chain ` +
          `more operators after both \`.gt\` and \`.lt\` were already used, ` +
          `got \`${printIndexOperator(filter)}\`.`,
      );
    }

    const filterType: "eq" | "gt" | "lt" =
      filter.type == "Gt" || filter.type == "Gte"
        ? "gt"
        : filter.type == "Lt" || filter.type == "Lte"
          ? "lt"
          : "eq";

    switch (`${state}|${filterType}`) {
      // Allow to operate on the current indexed field
      case "eq|eq":
      case "eq|gt":
      case "eq|lt":
        if (filter.fieldPath === fields[fieldIndex]) {
          fieldIndex += 1;
          state = filterType;
          continue;
        }
        throw new Error(
          `Incorrect field used in \`withIndex\`, ` +
            `expected "${fields[fieldIndex]}", got "${filter.fieldPath}"`,
        );

      // Allow to operate on the previous field (gt and lt must operate on same field)
      case "lt|gt":
      case "gt|lt":
        if (fieldIndex > 0 && filter.fieldPath === fields[fieldIndex - 1]) {
          state = "done";
          continue;
        }
        throw new Error(
          `Incorrect field used in \`withIndex\`, ` +
            `\`.gt\` and \`.lt\` must operate on the same field, ` +
            `expected "${fields[fieldIndex - 1]}", got "${filter.fieldPath}"`,
        );

      default:
        throw new Error(
          `Incorrect operator used in \`withIndex\`, ` +
            `cannot chain \`.${filter.type.toLowerCase()}()\` ` +
            `after \`.${source.range[filterIndex - 1].type.toLowerCase()}()\``,
        );
    }
  }
}

function printIndexOperator(filter: SerializedRangeExpression) {
  return `.${filter.type.toLowerCase()}(${filter.fieldPath}, ${JSON.stringify(filter.value)})`;
}

function validateFieldNames(validator: ValidatorJSON) {
  if (validator.type === "object") {
    Object.keys(validator.value).forEach((fieldName) => {
      if (!isValidIdentifier(fieldName)) {
        throw new Error(
          `Field names must be valid identifiers, got "${fieldName}"`,
        );
      }
    });
  }

  if (validator.type === "union") {
    validator.value.forEach(validateFieldNames);
  }
}

function isValidIdentifier(name: string) {
  return /(^_(id|creationTime)$)|^[a-zA-Z][a-zA-Z0-9_]*$/.test(name);
}

type ObjectFieldType = { fieldType: ValidatorJSON; optional: boolean };

type ValidatorJSON =
  | {
      type: "null";
    }
  | { type: "number" }
  | { type: "bigint" }
  | { type: "boolean" }
  | { type: "string" }
  | { type: "bytes" }
  | { type: "any" }
  | {
      type: "literal";
      value: JSONValue;
    }
  | { type: "id"; tableName: string }
  | { type: "array"; value: ValidatorJSON }
  | { type: "object"; value: Record<string, ObjectFieldType> }
  | { type: "union"; value: ValidatorJSON[] };

function validateValidator(validator: ValidatorJSON, value: any) {
  switch (validator.type) {
    case "null": {
      if (value !== null) {
        throw new Error(`Validator error: Expected \`null\`, got \`${value}\``);
      }
      return;
    }
    case "number": {
      if (typeof value !== "number") {
        throw new Error(
          `Validator error: Expected \`number\`, got \`${value}\``,
        );
      }
      return;
    }
    case "bigint": {
      if (typeof value !== "bigint") {
        throw new Error(
          `Validator error: Expected \`bigint\`, got \`${value}\``,
        );
      }
      return;
    }
    case "boolean": {
      if (typeof value !== "boolean") {
        throw new Error(
          `Validator error: Expected \`boolean\`, got \`${value}\``,
        );
      }
      return;
    }
    case "string": {
      if (typeof value !== "string") {
        throw new Error(
          `Validator error: Expected \`string\`, got \`${value}\``,
        );
      }
      return;
    }
    case "bytes": {
      if (!(value instanceof ArrayBuffer)) {
        throw new Error(
          `Validator error: Expected \`ArrayBuffer\`, got \`${value}\``,
        );
      }
      return;
    }
    case "any": {
      return;
    }
    case "literal": {
      if (value !== validator.value) {
        throw new Error(
          `Validator error: Expected \`${
            validator.value as any
          }\`, got \`${value}\``,
        );
      }
      return;
    }
    case "id": {
      if (typeof value !== "string") {
        throw new Error(
          `Validator error: Expected \`string\`, got \`${value}\``,
        );
      }
      if (tableNameFromId(value) !== validator.tableName) {
        throw new Error(
          `Validator error: Expected ID for table "${validator.tableName}", got \`${value}\``,
        );
      }
      return;
    }
    case "array": {
      if (!Array.isArray(value)) {
        throw new Error(
          `Validator error: Expected \`Array\`, got \`${value}\``,
        );
      }
      for (const v of value) {
        validateValidator(validator.value, v);
      }
      return;
    }
    case "union": {
      let isValid = false;
      for (const v of validator.value) {
        try {
          validateValidator(v, value);
          isValid = true;
          break;
        } catch {
          // Ignore
        }
      }
      if (!isValid) {
        throw new Error(
          `Validator error: Expected one of ${validator.value.map((v) => v.type).join(", ")}, got \`${JSON.stringify(convexToJson(value))}\``,
        );
      }
      return;
    }
    case "object": {
      if (typeof value !== "object") {
        throw new Error(
          `Validator error: Expected \`object\`, got \`${value}\``,
        );
      }
      if (!isSimpleObject(value)) {
        throw new Error(
          `Validator error: Expected a plain old JavaScript \`object\`, got \`${value}\``,
        );
      }
      for (const [k, { fieldType, optional }] of Object.entries(
        validator.value,
      )) {
        if (value[k] === undefined) {
          if (!optional) {
            throw new Error(
              `Validator error: Missing required field \`${k}\` in object`,
            );
          }
        } else {
          validateValidator(fieldType, value[k]);
        }
      }
      for (const k of Object.keys(value)) {
        if (validator.value[k] === undefined) {
          throw new Error(
            `Validator error: Unexpected field \`${k}\` in object`,
          );
        }
      }
      return;
    }
  }
}

function syscallImpl() {
  return (op: string, jsonArgs: string) => {
    const args = JSON.parse(jsonArgs);
    const db = getDb();
    switch (op) {
      case "1.0/queryStream": {
        const { query } = args;
        const tracker = getTransactionManager().getHeadroomTracker();
        tracker?.trackIndexRange();
        const queryId = db.startQuery(query);
        return JSON.stringify({ queryId });
      }
      case "1.0/queryCleanup": {
        return JSON.stringify({});
      }
      case "1.0/db/normalizeId": {
        const idString: string = args.idString;
        const isInTable = idString.endsWith(`;${args.table}`);
        return JSON.stringify({
          id: isInTable ? idString : null,
        });
      }
      default: {
        throw new Error(`\`convexTest\` does not support syscall: "${op}"`);
      }
    }
  };
}

class AuthFake {
  constructor(private _userIdentity: any = null) {}

  getUserIdentity(): Promise<any> {
    return Promise.resolve(this._userIdentity);
  }
}

function asyncSyscallImpl() {
  return async (op: string, jsonArgs: string): Promise<string> => {
    const args = JSON.parse(jsonArgs);
    const db = getDb();
    const tracker = getTransactionManager().getHeadroomTracker();
    function getAndTrack(table: string, id: GenericId<string>) {
      const doc = db.get(table, id);
      if (doc !== null) {
        tracker?.trackRead(getDocumentSize(doc));
      }
      return doc;
    }
    switch (op) {
      case "1.0/get": {
        const { table, id } = args;
        tracker?.trackIndexRange();
        const doc = getAndTrack(table, id);
        return JSON.stringify(convexToJson(doc));
      }
      case "1.0/queryStreamNext": {
        const { value, done } = db.queryNext(args.queryId);
        if (!done && value !== null) {
          tracker?.trackRead(getDocumentSize(value));
        }
        return JSON.stringify(convexToJson({ value, done }));
      }
      case "1.0/queryPage": {
        const {
          query,
          cursor,
          endCursor,
          pageSize,
          maximumRowsRead,
          maximumBytesRead,
        } = args;
        const { page, isDone, continueCursor, splitCursor, pageStatus } =
          db.paginate({
            query,
            cursor,
            endCursor,
            pageSize,
            maximumRowsRead,
            maximumBytesRead,
          });
        if (tracker) {
          tracker.trackIndexRange();
          for (const doc of page) {
            tracker.trackRead(getDocumentSize(doc));
          }
        }
        return JSON.stringify(
          convexToJson({
            page,
            isDone,
            continueCursor,
            splitCursor,
            pageStatus,
          }),
        );
      }
      case "1.0/insert": {
        const _id = db.insert(args.table, jsonToConvex(args.value));
        tracker?.trackWrite(getConvexSize(db.get(args.table, _id)));
        return JSON.stringify({ _id });
      }
      case "1.0/shallowMerge": {
        const { table, id, value } = args;
        getAndTrack(table, id);
        db.patch(table, id, value);
        tracker?.trackWrite(getConvexSize(db.get(table, id)));
        return JSON.stringify({});
      }
      case "1.0/replace": {
        const { table, id, value } = args;
        getAndTrack(table, id);
        db.replace(table, id, value);
        tracker?.trackWrite(getConvexSize(db.get(table, id)));
        return JSON.stringify({});
      }
      case "1.0/remove": {
        const { table, id } = args;
        getAndTrack(table, id);
        db.delete(table, id);
        tracker?.trackWrite(0);
        return JSON.stringify({});
      }
      case "1.0/headroom": {
        if (tracker) {
          return JSON.stringify(convexToJson(tracker.getTransactionHeadroom()));
        }
        throw new Error(
          "getTransactionHeadroom() can only be called from a query or mutation. " +
            "It is not available in actions or outside of a Convex function.",
        );
      }
      case "1.0/actions/query": {
        const { name, args: queryArgs } = args;
        return JSON.stringify(
          convexToJson(
            await withAuth().runQuery(makeFunctionReference(name), queryArgs),
          ),
        );
      }
      case "1.0/actions/mutation": {
        const { name, args: mutationArgs } = args;
        return JSON.stringify(
          convexToJson(
            await withAuth().runMutation(
              makeFunctionReference(name),
              mutationArgs,
            ),
          ),
        );
      }
      case "1.0/actions/action": {
        const { name, args: actionArgs } = args;
        return JSON.stringify(
          convexToJson(
            await withAuth().runAction(makeFunctionReference(name), actionArgs),
          ),
        );
      }
      case "1.0/runUdf": {
        const {
          udfType,
          name,
          reference,
          functionHandle,
          args: udfArgsJson,
        } = args;
        const udfArgs = jsonToConvex(udfArgsJson);
        const functionPath = getFunctionPathFromAddress({
          name,
          reference,
          functionHandle,
        });
        if (udfType === "query") {
          return JSON.stringify(
            convexToJson(
              await withAuth().queryFromPath(
                functionPath,
                /* isNested */ true,
                udfArgs,
              ),
            ),
          );
        }
        if (udfType === "mutation") {
          return JSON.stringify(
            convexToJson(
              await withAuth().mutationFromPath(
                functionPath,
                /* isNested */ true,
                udfArgs,
              ),
            ),
          );
        }
        throw new Error(
          `\`convexTest\` does not support udf type: "${udfType}"`,
        );
      }
      case "1.0/createFunctionHandle": {
        const { name, reference, functionHandle } = args;
        const functionPath = getFunctionPathFromAddress({
          name,
          reference,
          functionHandle,
        });
        const handle = createFunctionHandle(functionPath);
        return JSON.stringify(handle);
      }

      case "1.0/actions/schedule": {
        return await withAuth().run(async () => {
          return await getSyscalls().asyncSyscall("1.0/schedule", jsonArgs);
        });
      }
      case "1.0/schedule": {
        const {
          name,
          reference,
          functionHandle,
          args: fnArgs,
          ts: tsInSecs,
        } = args;
        const functionPath = getFunctionPathFromAddress({
          name,
          reference,
          functionHandle,
        });
        // Convert args to a Convex value before storing them or passing them to the function
        const parsedArgs = jsonToConvex(fnArgs);

        tracker?.trackScheduledFunction(getConvexSize(parsedArgs));

        const jobId = db.insert("_scheduled_functions", {
          args: [parsedArgs],
          name: functionPath.udfPath,
          scheduledTime: tsInSecs * 1000,
          state: { kind: "pending" },
        });
        const componentPath = getCurrentComponentPath();
        setTimeout(
          (async () => {
            const canceled = await withAuth().runInComponent(
              componentPath,
              async () => {
                const job = db.get(
                  "_scheduled_functions",
                  jobId,
                ) as ScheduledFunction;
                if (job.state.kind === "canceled") {
                  return true;
                }
                if (job.state.kind !== "pending") {
                  throw new Error(
                    `\`convexTest\` invariant error: Unexpected scheduled function state when starting it: ${job.state.kind}`,
                  );
                }
                db.patch("_scheduled_functions", jobId, {
                  state: { kind: "inProgress" },
                });
                return false;
              },
            );
            if (canceled) {
              return;
            }
            try {
              await withAuth().fun(functionPath, parsedArgs);
            } catch (error) {
              console.error(
                `Error when running scheduled function ${name}`,
                error,
              );
              await withAuth().runInComponent(componentPath, async () => {
                db.patch("_scheduled_functions", jobId, {
                  state: { kind: "failed" },
                  completedTime: Date.now(),
                });
              });
              db.jobFinished(jobId);
              return;
            }
            await withAuth().runInComponent(componentPath, async () => {
              const job = db.get(
                "_scheduled_functions",
                jobId,
              ) as ScheduledFunction;
              if (job.state.kind !== "inProgress") {
                throw new Error(
                  `\`convexTest\` invariant error: Unexpected scheduled function state after it finished running: ${job.state.kind}`,
                );
              }
              db.patch("_scheduled_functions", jobId, {
                state: { kind: "success" },
              });
            });
            db.jobFinished(jobId);
          }) as () => void,
          tsInSecs * 1000 - Date.now(),
        );
        return JSON.stringify(convexToJson(jobId));
      }
      case "1.0/actions/cancel_job": {
        await withAuth().run(async () => {
          await getSyscalls().asyncSyscall("1.0/cancel_job", jsonArgs);
        });
        return JSON.stringify({});
      }
      case "1.0/actions/vectorSearch": {
        const {
          query: { indexName, limit, vector, expressions },
        } = args;
        const results = db.vectorSearch(
          indexName,
          vector,
          // Probably an unintentional implementation in Convex
          // where expressions is only a single expression
          expressions === null ? [] : [expressions],
          limit,
        );
        return JSON.stringify(convexToJson({ results }));
      }
      case "1.0/cancel_job": {
        const { id } = args;
        db.patch("_scheduled_functions", id, { state: { kind: "canceled" } });
        return JSON.stringify({});
      }
      case "1.0/storageDelete": {
        const { storageId } = args;
        await writeToDatabase(async (db) => {
          db.delete("_storage", storageId);
        });
        return JSON.stringify({});
      }
      case "1.0/storageGetUrl": {
        const { storageId } = args;
        const metadata = db.get("_storage", storageId);
        if (metadata === null) {
          return JSON.stringify(null);
        }
        const { sha256 } = metadata;
        // In the real backend the URL ofc isn't the sha
        const url =
          "https://some-deployment.convex.cloud/api/storage/" +
          (sha256 as string);
        return JSON.stringify(convexToJson(url));
      }
      case "1.0/storageGenerateUploadUrl": {
        // In the real backend the token is cryptographically secure
        const url =
          "https://some-deployment.convex.cloud/api/storage/upload?token=" +
          Math.random();
        return JSON.stringify(convexToJson(url));
      }
      case "1.0/count": {
        const { table } = args;
        const queryId = db.startQuery({
          source: { type: "FullTableScan", tableName: table, order: "asc" },
          operators: [],
        });
        let count = 0;

        while (true) {
          const result = db.queryNext(queryId);
          if (result.done) {
            break;
          }
          count += 1;
        }
        return JSON.stringify(count);
      }
      default: {
        throw new Error(
          `\`convexTest\` does not support async syscall: "${op}"`,
        );
      }
    }
  };
}

function jsSyscallImpl() {
  return async (op: string, args: Record<string, any>): Promise<any> => {
    const db = getDb();
    switch (op) {
      case "storage/storeBlob": {
        const { blob } = args as { blob: Blob };
        const storageId = await writeToDatabase(async (db) => {
          return db.insert("_storage", {
            size: blob.size,
            sha256: await blobSha(blob),
          });
        });
        db.storeFile(storageId, blob);
        return storageId;
      }
      case "storage/getBlob": {
        const { storageId } = args as { storageId: GenericId<"_storage"> };
        return db.getFile(storageId);
      }
      default: {
        throw new Error(`\`convexTest\` does not support js syscall: "${op}"`);
      }
    }
  };
}

// If we're in action, wrap the write in a transaction.
async function writeToDatabase<T>(impl: (db: DatabaseFake) => Promise<T>) {
  const db = getDb();
  if (!getTransactionManager().isInTransaction()) {
    return await withAuth().run(async () => {
      return await impl(db);
    });
  } else {
    return await impl(db);
  }
}

async function blobSha(blob: Blob) {
  const arrayBuffer = await blob.arrayBuffer();
  const hashBuffer = await crypto.subtle.digest("SHA-256", arrayBuffer);
  const hashArray = new Uint8Array(hashBuffer);
  return btoa(String.fromCharCode(...hashArray));
}

async function waitForInProgressScheduledFunctions(): Promise<boolean> {
  let hadScheduledFunctions = false;
  for (const componentPath of Object.keys(getConvexGlobal().components)) {
    const inProgressJobs = (await withAuth().runInComponent(
      componentPath,
      async (ctx) => {
        return (
          await ctx.db.system.query("_scheduled_functions").collect()
        ).filter((job: ScheduledFunction) => job.state.kind === "inProgress");
      },
    )) as ScheduledFunction[];
    let numRemaining = inProgressJobs.length;
    if (numRemaining === 0) {
      continue;
    }
    hadScheduledFunctions = true;

    await new Promise<void>((resolve) => {
      getDbForComponent(componentPath).jobListener = () => {
        numRemaining -= 1;
        if (numRemaining === 0) {
          resolve();
        }
      };
    });
  }
  return hadScheduledFunctions;
}

export type TestConvex<SchemaDef extends SchemaDefinition<any, boolean>> =
  TestConvexForDataModelAndIdentity<DataModelFromSchemaDefinition<SchemaDef>>;

export type TestConvexForDataModel<DataModel extends GenericDataModel> = {
  /**
   * Call a public or internal query, or run an inline query function.
   *
   * @param query A {@link FunctionReference} for the query, or an inline
   *   async function that receives a query context.
   * @param args  An arguments object for the query. If this is omitted,
   *   the arguments will be `{}`. Not used for inline functions.
   * @returns A `Promise` of the query's result.
   */
  query: {
    <Query extends FunctionReference<"query", any>>(
      query: Query,
      ...args: OptionalRestArgs<Query>
    ): Promise<FunctionReturnType<Query>>;
    <Output>(
      func: (ctx: GenericQueryCtx<DataModel>) => Promise<Output>,
    ): Promise<Output>;
  };

  /**
   * Call a public or internal mutation, or run an inline mutation function.
   *
   * @param mutation A {@link FunctionReference} for the mutation, or an inline
   *   async function that receives a mutation context.
   * @param args  An arguments object for the mutation. If this is omitted,
   *   the arguments will be `{}`. Not used for inline functions.
   * @returns A `Promise` of the mutation's result.
   */
  mutation: {
    <Mutation extends FunctionReference<"mutation", any>>(
      mutation: Mutation,
      ...args: OptionalRestArgs<Mutation>
    ): Promise<FunctionReturnType<Mutation>>;
    <Output>(
      func: (ctx: GenericMutationCtx<DataModel>) => Promise<Output>,
    ): Promise<Output>;
  };

  /**
   * Call a public or internal action, or run an inline action function.
   *
   * @param action A {@link FunctionReference} for the action, or an inline
   *   async function that receives an action context.
   * @param args  An arguments object for the action. If this is omitted,
   *   the arguments will be `{}`. Not used for inline functions.
   * @returns A `Promise` of the action's result.
   */
  action: {
    <Action extends FunctionReference<"action", any>>(
      action: Action,
      ...args: OptionalRestArgs<Action>
    ): Promise<FunctionReturnType<Action>>;
    <Output>(
      func: (ctx: GenericActionCtx<DataModel>) => Promise<Output>,
    ): Promise<Output>;
  };

  /**
   * Read from and write to the mock backend.
   *
   * @param func The async function that reads or writes to the mock backend.
   *   It receives a ctx as its first argument that conforms to
   *   {@link GenericMutationCtx},
   *   so it can be passed to functions that expect either context, use file
   *   storage, etc.
   * @returns A `Promise` of the function's result.
   */
  run: <Output>(
    func: (
      ctx: GenericMutationCtx<DataModel> &
        Pick<
          GenericActionCtx<DataModel>,
          "storage"
          // ActionCtx support was added in 0.0.40, removed in 0.0.41, should be added back in the future.
          // | "runAction"
          // | "vectorSearch"
        >,
    ) => Promise<Output>,
  ) => Promise<Output>;

  /**
   * Call an HTTP action.
   *
   * @param path The request URL's path and optional query and fragment.
   * @param init Standard `fetch` options.
   */
  fetch(pathQueryFragment: string, init?: RequestInit): Promise<Response>;

  /**
   * Wait for all scheduled functions currently in the "inProgress" state
   * to either finish successfully or fail.
   *
   * Use in combination with `vi.useFakeTimers()` and `vi.runAllTimers()`
   * to control precisely the execution of scheduled functions.
   *
   * Typically:
   * 1. Use `vi.runAllTimers()` or similar to advance
   *   time such that a function is scheduled.
   * 2. Use `finishInProgressScheduledFunctions()` to wait for the function
   */
  finishInProgressScheduledFunctions: () => Promise<void>;

  /**
   * Wait for all currently scheduled functions and any functions they
   * schedule to either finish successfully or fail.
   *
   * Use in combination with `vi.useFakeTimers()` to test scheduled functions.
   *
   * @param advanceTimers Function that advances timers,
   *   usually `vi.runAllTimers`. This function will be called in a loop
   *   with `finishInProgressScheduledFunctions()`.
   */
  finishAllScheduledFunctions: (advanceTimers: () => void) => Promise<void>;
};

export type TestConvexForDataModelAndIdentity<
  DataModel extends GenericDataModel,
> = {
  /**
   * To test functions which depend on the current authenticated user identity
   * you can create a version of the `t` accessor with given user identity
   * attributes.
   * @param identity A subset of {@link UserIdentity} attributes. If you
   *   don't provide `issuer`, `subject` or `tokenIdentifier` they are
   *   generated automatically.
   */
  withIdentity(
    identity: Partial<UserIdentity>,
  ): TestConvexForDataModel<DataModel>;
  registerComponent: (
    componentPath: string,
    schema: SchemaDefinition<GenericSchema, boolean>,
    glob: Record<string, () => Promise<any>>,
  ) => void;
} & TestConvexForDataModel<DataModel>;

function getComponentInfo(componentPath: string): ComponentInfo {
  const convex = getConvexGlobal();
  if (convex.components[componentPath] === undefined) {
    throw new Error(
      `Component "${componentPath}" is not registered. Call "t.registerComponent".`,
    );
  }
  return convex.components[componentPath];
}

function getDbForComponent(componentPath: string) {
  return getComponentInfo(componentPath).db;
}

function getDb() {
  return getDbForComponent(getCurrentComponentPath());
}

function getTransactionManager() {
  return getConvexGlobal().transactionManager;
}

function getCurrentComponentPath() {
  const functionStack = getTransactionManager().getFunctionStack();
  const currentFunctionPath = functionStack[functionStack.length - 1];
  return currentFunctionPath?.componentPath ?? ROOT_COMPONENT_PATH;
}

function getModules() {
  return getModulesForComponent(getCurrentComponentPath());
}

function getModulesForComponent(componentPath: string) {
  return getComponentInfo(componentPath).modules;
}

function getSyscalls() {
  return (global as any).Convex as {
    syscall: ReturnType<typeof syscallImpl>;
    asyncSyscall: ReturnType<typeof asyncSyscallImpl>;
  };
}

function moduleCache(specifiedModules?: Record<string, () => Promise<any>>) {
  const modules = specifiedModules ?? import.meta.glob("./convex/**/*.*s");
  const prefix = findModulesRoot(
    Object.keys(modules),
    specifiedModules !== undefined,
  );
  const modulesWithoutExtension = Object.fromEntries(
    Object.entries(modules).map(([path, module]) => [
      path.replace(/\.[^.]+$/, ""),
      module,
    ]),
  );
  return async (path: string) => {
    const module = modulesWithoutExtension[prefix + path];
    if (module === undefined) {
      throw new Error(`Could not find module for: "${path}"`);
    }
    return await module();
  };
}

function findModulesRoot(modulesPaths: string[], userProvidedModules: boolean) {
  const generatedFilePath = modulesPaths.find((path) =>
    path.includes("_generated"),
  );
  if (generatedFilePath !== undefined) {
    return generatedFilePath.split("_generated", 2)[0];
  }

  throw new Error(
    'Could not find the "_generated" directory, make sure to run ' +
      "`npx convex dev` or `npx convex codegen`. " +
      (userProvidedModules
        ? "Make sure your `import.meta.glob` includes the files in the " +
          '"_generated" directory'
        : "If your Convex functions aren't defined in a directory " +
          'called "convex" sibling to your node_modules, ' +
          "provide the second argument to `convexTest`"),
  );
}

type ComponentInfo = {
  db: DatabaseFake;
  modules: ReturnType<typeof moduleCache>;
};

type FunctionPath = {
  // e.g. "aggregate"
  componentPath: string;
  // e.g. "messages:list"
  udfPath: string;
};

// Per-execution-context function stack. Actions create their own context
// so parallel actions don't corrupt each other's stacks.
const functionStackStorage = new AsyncLocalStorage<FunctionPath[]>();

type ConvexGlobal = {
  components: Record<string, ComponentInfo>;
  transactionManager: TransactionManager;
  syscall: (op: string, args: any) => any;
  asyncSyscall: (op: string, args: any) => Promise<any>;
  jsSyscall: (op: string, args: any) => Promise<any>;
};

function getConvexGlobal(): ConvexGlobal {
  return (global as unknown as { Convex: ConvexGlobal }).Convex;
}

function setConvexGlobal(convex: ConvexGlobal) {
  if (getConvexGlobal()) {
    if (getTransactionManager().isInTransaction()) {
      throw new Error("test began while previous transaction was still open");
    }
  }
  (global as unknown as { Convex: ConvexGlobal }).Convex = convex;
}

class TransactionManager {
  // The DatabaseFake is used in the Convex global,
  // and so it restricts `convexTest` to run one function
  // at a time.
  // We force sequential execution to make sure actions
  // can run mutations in parallel.
  private _waitOnCurrentFunction: Promise<void> | null = null;
  private _markTransactionDone: (() => void) | null = null;
  // Default function stack used when not inside an action's ALS context.
  private _functionStack: FunctionPath[] = [];
  // Root-level nested lock: serializes the first layer of nested sub-transactions.
  // Deeper layers get their own locks via AsyncLocalStorage.
  public readonly rootNestedLock = new NestedLock();

  // Returns the per-context function stack (ALS stack for actions, shared stack otherwise).
  getFunctionStack(): FunctionPath[] {
    return functionStackStorage.getStore() ?? this._functionStack;
  }

  private _headroomTracker: HeadroomTracker | null = null;
  private _limitsConfig: Partial<TransactionMetrics> | boolean;

  constructor(limitsConfig: Partial<TransactionMetrics> | boolean = false) {
    this._limitsConfig = limitsConfig;
  }

  async begin(functionPath: FunctionPath, isNested: boolean) {
    // Take a lock only for the top-level of each transaction.
    // Nested transactions are not isolated so if you `Promise.all` on multiple
    // `ctx.runMutation` or `ctx.runQuery` calls, they won't be serialized.
    if (!isNested) {
      // Note the loop is important, as the current promise might resolve and a
      // new transaction could start before we get woken up.
      // This is the standard pattern for condition variables:
      // https://en.wikipedia.org/wiki/Monitor_(synchronization)
      while (this._waitOnCurrentFunction !== null) {
        await this._waitOnCurrentFunction;
      }
      this._waitOnCurrentFunction = new Promise((resolve) => {
        this._markTransactionDone = resolve;
      });
      this._headroomTracker = new HeadroomTracker(this._limitsConfig);
    }
    this.getFunctionStack().push(functionPath);
    const convex = getConvexGlobal();
    for (const component of Object.values(convex.components)) {
      component.db.startTransaction();
    }
  }

  beginAction(functionPath: FunctionPath) {
    this.getFunctionStack().push(functionPath);
  }

  finishAction() {
    this.getFunctionStack().pop();
  }

  // Used to distinguish between mutation and action execution
  // environment.
  isInTransaction() {
    return this._waitOnCurrentFunction !== null;
  }

  getHeadroomTracker(): HeadroomTracker | null {
    return this._headroomTracker;
  }

  commit(isNested: boolean) {
    const convex = getConvexGlobal();
    for (const component of Object.values(convex.components)) {
      component.db.commit();
    }
    this._endTransaction(isNested);
  }

  rollback(isNested: boolean) {
    const convex = getConvexGlobal();
    for (const component of Object.values(convex.components)) {
      component.db.rollbackWrites();
    }
    this._endTransaction(isNested);
  }

  _endTransaction(isNested: boolean) {
    if (this._markTransactionDone === null) {
      throw new Error("Transaction not started");
    }
    this.getFunctionStack().pop();
    if (!isNested) {
      this._headroomTracker = null;
      this._waitOnCurrentFunction = null;
      this._markTransactionDone();
      this._markTransactionDone = null;
    }
  }
}

/**
 * Call this function at the start of each of your tests.
 *
 * @param schema The default export from your "schema.ts" file.
 * @param modules If you have a custom `functions` path
 *     in convex.json, provide the function modules map by calling
 *     `import.meta.glob("./**\/*.*s")` from a file at the root of
 *     your Convex functions directory.
 * @returns an object which is by convention stored in the `t` variable
 *   and which provides methods for exercising your Convex functions.
 */
export function convexTest<Schema extends GenericSchema>(
  schema?: SchemaDefinition<Schema, boolean>,
  // For example `import.meta.glob("./**/*.*s")`
  modules?: Record<string, () => Promise<any>>,
): TestConvex<SchemaDefinition<Schema, boolean>>;
/**
 * Call this function at the start of each of your tests.
 *
 * @param options An optional configuration object with:
 *   - `schema`: The default export from your "schema.ts" file.
 *   - `modules`: If you have a custom `functions` path
 *     in convex.json, provide the function modules map by calling
 *     `import.meta.glob("./**\/*.*s")` from a file at the root of
 *     your Convex functions directory.
 *   - `transactionLimits`: Configure per-transaction bandwidth limits.
 *     - `false` (default): limits are not enforced
 *     - `true` or `{}`: enforce default Convex limits
 *     - `{ bytesRead: 1024 }`: override specific limits
 *
 * @returns an object which is by convention stored in the `t` variable
 *   and which provides methods for exercising your Convex functions.
 *
 * @example
 * const t = convexTest({ schema });
 * const t = convexTest({ schema, modules });
 */
export function convexTest<Schema extends GenericSchema>(options: {
  schema?: SchemaDefinition<Schema, boolean>;
  /** For example `import.meta.glob("./**\/*.*s")` */
  modules?: Record<string, () => Promise<any>>;
  transactionLimits?: Partial<TransactionMetrics> | boolean;
}): TestConvex<SchemaDefinition<Schema, boolean>>;
export function convexTest<Schema extends GenericSchema>(
  schemaOrOptions?:
    | SchemaDefinition<Schema, boolean>
    | {
        schema?: SchemaDefinition<Schema, boolean>;
        modules?: Record<string, () => Promise<any>>;
        transactionLimits?: Partial<TransactionMetrics> | boolean;
      },
  modules?: Record<string, () => Promise<any>>,
): TestConvex<SchemaDefinition<Schema, boolean>> {
  let schema: SchemaDefinition<Schema, boolean> | undefined = undefined;
  let limitsConfig: Partial<TransactionMetrics> | boolean = false;

  // Detect legacy API by checking for schema's "tables" property
  function isSchemaDefinition(
    schemaOrOptions: unknown,
  ): schemaOrOptions is SchemaDefinition<Schema, boolean> {
    return (
      schemaOrOptions !== null &&
      typeof schemaOrOptions === "object" &&
      "tables" in schemaOrOptions &&
      typeof (schemaOrOptions as any).tables === "object"
    );
  }

  if (schemaOrOptions === undefined) {
    // convexTest() - no arguments
  } else if (isSchemaDefinition(schemaOrOptions)) {
    // convexTest(schema, modules?)
    schema = schemaOrOptions;
  } else {
    // convexTest({ schema?, modules?, transactionLimits? })
    const opts = schemaOrOptions;
    schema = opts.schema;
    modules = opts.modules;
    limitsConfig = opts.transactionLimits ?? false;
  }

  const rootDb = new DatabaseFake(schema ?? null, ROOT_COMPONENT_PATH);
  setConvexGlobal({
    components: {
      [ROOT_COMPONENT_PATH]: {
        db: rootDb,
        modules: moduleCache(modules),
      },
    },
    transactionManager: new TransactionManager(limitsConfig),
    syscall: syscallImpl(),
    asyncSyscall: asyncSyscallImpl(),
    jsSyscall: jsSyscallImpl(),
  });

  return {
    withIdentity(identity: Partial<UserIdentity>) {
      const subject =
        identity.subject ?? "" + simpleHash(JSON.stringify(identity));
      const issuer = identity.issuer ?? "https://convex.test";
      const tokenIdentifier =
        identity.tokenIdentifier ?? `${issuer}|${subject}`;
      return withAuth(
        new AuthFake({ ...identity, subject, issuer, tokenIdentifier }),
      );
    },
    ...withAuth(),
    registerComponent(
      componentPath: string,
      schema: SchemaDefinition<GenericSchema, boolean>,
      glob: Record<string, () => Promise<any>>,
    ) {
      const componentInfo = {
        db: new DatabaseFake(schema, componentPath),
        modules: moduleCache(glob),
      };
      getConvexGlobal().components[componentPath] = componentInfo;
    },
  } as any;
}

function withAuth(auth: AuthFake = new AuthFake()) {
  const runMutationWithHandler = async <T>(
    handler: (ctx: any, args: any) => T,
    args: any,
    extraCtx: any = {},
    functionPath: FunctionPath,
    isNested: boolean,
  ): Promise<T> => {
    const m = mutationGeneric({
      handler: (ctx: any, a: any) => {
        const testCtx = { ...ctx, auth, ...extraCtx };
        return handler(testCtx, a);
      },
    });
    const transactionManager = getTransactionManager();

    // Serialize nested calls at each nesting level to prevent functionStack
    // and transaction level corruption from parallel component sub-calls.
    // Each level gets its own lock via AsyncLocalStorage so nested-within-nested
    // parallel calls also serialize correctly within their own level.
    const parentLock = isNested
      ? (nestedTxStorage.getStore() ?? transactionManager.rootNestedLock)
      : null;
    if (parentLock) {
      await parentLock.acquire();
    }

    await transactionManager.begin(functionPath, isNested);
    try {
      const childLock = new NestedLock();
      const invokeInContext = () =>
        (
          m as unknown as {
            invokeMutation: (args: string) => Promise<string>;
          }
        ).invokeMutation(JSON.stringify(convexToJson([parseArgs(args)])));

      const rawResult = isNested
        ? await nestedTxStorage.run(childLock, invokeInContext)
        : await invokeInContext();

      transactionManager.commit(isNested);
      return jsonToConvex(JSON.parse(rawResult)) as T;
    } catch (e) {
      transactionManager.rollback(isNested);
      throw e;
    } finally {
      if (parentLock) {
        parentLock.release();
      }
    }
  };

  // Shared helper to run a query with a given handler and transaction context
  // Used by both queryFromPath (for function references) and inline queries
  const runQueryWithHandler = async <T>(
    handler: (ctx: any, args: any) => T,
    args: any,
    functionPath: FunctionPath,
    isNested: boolean,
  ): Promise<T> => {
    const q = queryGeneric({
      handler: (ctx: any, a: any) => {
        const testCtx = { ...ctx, auth };
        return handler(testCtx, a);
      },
    });
    const transactionManager = getTransactionManager();

    // Serialize nested calls at each nesting level (same pattern as runTransaction).
    const parentLock = isNested
      ? (nestedTxStorage.getStore() ?? transactionManager.rootNestedLock)
      : null;
    if (parentLock) {
      await parentLock.acquire();
    }

    await transactionManager.begin(functionPath, isNested);
    try {
      const childLock = new NestedLock();
      const invokeInContext = () =>
        (
          q as unknown as { invokeQuery: (args: string) => Promise<string> }
        ).invokeQuery(JSON.stringify(convexToJson([parseArgs(args)])));

      const rawResult = isNested
        ? await nestedTxStorage.run(childLock, invokeInContext)
        : await invokeInContext();

      return jsonToConvex(JSON.parse(rawResult)) as T;
    } finally {
      transactionManager.rollback(isNested);
      if (parentLock) {
        parentLock.release();
      }
    }
  };

  // Shared helper to run an action with a given handler
  // Used by both actionFromPath (for function references) and inline actions
  const runActionWithHandler = async <T>(
    handler: (ctx: any, args: any) => T,
    args: any,
    functionPath: FunctionPath,
    ctxRunQuery: any,
    ctxRunMutation: any,
    ctxRunAction: any,
  ): Promise<T> => {
    const a = actionGeneric({
      handler: (ctx: any, innerArgs: any) => {
        const testCtx = {
          ...ctx,
          runQuery: ctxRunQuery,
          runMutation: ctxRunMutation,
          runAction: ctxRunAction,
          auth,
        };
        return handler(testCtx, innerArgs);
      },
    });
    // Each action gets its own function stack via ALS so parallel actions
    // don't corrupt each other's stacks.
    const stack: FunctionPath[] = [];
    return await functionStackStorage.run(stack, async () => {
      getTransactionManager().beginAction(functionPath);
      const requestId = "" + Math.random();
      try {
        const rawResult = await (
          a as unknown as {
            invokeAction: (requestId: string, args: string) => Promise<string>;
          }
        ).invokeAction(
          requestId,
          JSON.stringify(convexToJson([parseArgs(args)])),
        );
        return jsonToConvex(JSON.parse(rawResult)) as T;
      } finally {
        getTransactionManager().finishAction();
      }
    });
  };

  const byTypeWithPath = {
    queryFromPath: async (
      functionPath: FunctionPath,
      isNested: boolean,
      args: any,
    ) => {
      const func = await getFunctionFromPath(functionPath, "query");
      validateValidator(JSON.parse((func as any).exportArgs()), args ?? {});

      const result = await runQueryWithHandler(
        (ctx, a) => getHandler(func)(ctx, a),
        args,
        functionPath,
        isNested,
      );
      validateReturnValue(func, result, "query", functionPath);
      return result;
    },

    mutationFromPath: async (
      functionPath: FunctionPath,
      isNested: boolean,
      args: any,
    ): Promise<Value> => {
      const func = await getFunctionFromPath(functionPath, "mutation");
      validateValidator(JSON.parse((func as any).exportArgs()), args ?? {});

      const result = await runMutationWithHandler(
        getHandler(func),
        args,
        {},
        functionPath,
        isNested,
      );
      validateReturnValue(func, result, "mutation", functionPath);
      return result;
    },

    actionFromPath: async (functionPath: FunctionPath, args: any) => {
      const func = await getFunctionFromPath(functionPath, "action");
      validateValidator(JSON.parse((func as any).exportArgs()), args ?? {});

      const result = await runActionWithHandler(
        getHandler(func),
        args,
        functionPath,
        byTypeWithPath.runQuery,
        byTypeWithPath.runMutation,
        byTypeWithPath.runAction,
      );
      validateReturnValue(func, result, "action", functionPath);
      return result;
    },
    runQuery: async (
      functionReference: FunctionReference<any, any, any, any>,
      args: any,
    ) => {
      const refPath = await getFunctionPathFromReference(functionReference);
      return await byTypeWithPath.queryFromPath(refPath, false, args);
    },
    runMutation: async (
      functionReference: FunctionReference<any, any, any, any>,
      args: any,
    ) => {
      const refPath = await getFunctionPathFromReference(functionReference);
      return await byTypeWithPath.mutationFromPath(refPath, false, args);
    },
    runAction: async (
      functionReference: FunctionReference<any, any, any, any>,
      args: any,
    ) => {
      const refPath = await getFunctionPathFromReference(functionReference);
      return await byTypeWithPath.actionFromPath(refPath, args);
    },
  };

  const byType = {
    query: async (functionReferenceOrHandler: any, args?: any) => {
      if (typeof functionReferenceOrHandler === "function") {
        return await runQueryWithHandler(
          functionReferenceOrHandler,
          args ?? {},
          { componentPath: getCurrentComponentPath(), udfPath: "inline" },
          /* isNested */ false,
        );
      }
      const functionPath = await getFunctionPathFromReference(
        functionReferenceOrHandler,
      );
      return await byTypeWithPath.queryFromPath(
        functionPath,
        /* isNested */ false,
        args,
      );
    },

    mutation: async (
      functionReferenceOrHandler: any,
      args?: any,
    ): Promise<Value> => {
      if (typeof functionReferenceOrHandler === "function") {
        return await runMutationWithHandler(
          functionReferenceOrHandler,
          args ?? {},
          /* extraCtx */ {},
          { componentPath: getCurrentComponentPath(), udfPath: "inline" },
          /* isNested */ false,
        );
      }
      const functionPath = await getFunctionPathFromReference(
        functionReferenceOrHandler,
      );
      return await byTypeWithPath.mutationFromPath(
        functionPath,
        /* isNested */ false,
        args,
      );
    },

    action: async (functionReferenceOrHandler: any, args?: any) => {
      if (typeof functionReferenceOrHandler === "function") {
        return await runActionWithHandler(
          functionReferenceOrHandler,
          args ?? {},
          { componentPath: getCurrentComponentPath(), udfPath: "inline" },
          byTypeWithPath.runQuery,
          byTypeWithPath.runMutation,
          byTypeWithPath.runAction,
        );
      }
      const functionPath = await getFunctionPathFromReference(
        functionReferenceOrHandler,
      );
      return await byTypeWithPath.actionFromPath(functionPath, args);
    },
  };
  const run = async <T>(
    componentPath: string,
    handler: (ctx: any) => T,
  ): Promise<T> => {
    // Grab StorageActionWriter from action ctx
    const a = actionGeneric({
      handler: async (ctx: GenericActionCtx<any>) => {
        return await runMutationWithHandler(
          handler,
          {},
          {
            storage: ctx.storage,
            // ActionCtx support was added in 0.0.40, removed in 0.0.41, should be added back in the future.
            // runAction: ctx.runAction.bind(ctx),
            // vectorSearch: ctx.vectorSearch.bind(ctx),
          },
          // Fake mutation path.
          {
            componentPath,
            udfPath: "custom",
          },
          false,
        );
      },
    });
    // Real backend uses different ID format
    const requestId = "" + Math.random();
    const rawResult = await (
      a as unknown as {
        invokeAction: (requestId: string, args: string) => Promise<string>;
      }
    ).invokeAction(requestId, JSON.stringify(convexToJson([{}])));
    return jsonToConvex(JSON.parse(rawResult)) as T;
  };
  return {
    ...byType,
    ...byTypeWithPath,

    runInComponent: run,
    run: async <T>(handler: (ctx: any) => T) => {
      return await run(getCurrentComponentPath(), handler);
    },

    fun: async (functionPath: FunctionPath, args: any) => {
      const func = await getFunctionFromPath(functionPath, "any");
      if (func.isQuery) {
        return await byTypeWithPath.queryFromPath(
          functionPath,
          /* isNested */ false,
          args,
        );
      }
      if (func.isMutation) {
        return await byTypeWithPath.mutationFromPath(
          functionPath,
          /* isNested */ false,
          args,
        );
      }
      if (func.isAction) {
        return await byTypeWithPath.actionFromPath(functionPath, args);
      }
    },

    fetch: async (path: string, init?: RequestInit) => {
      const router: HttpRouter = (await getModules()("http"))["default"];
      if (!path.startsWith("/")) {
        throw new Error(`Path given to \`t.fetch\` must start with a \`/\``);
      }
      const url = new URL(`https://some.convex.site${path}`);
      const found = router.lookup(url.pathname, init?.method ?? ("GET" as any));
      if (!found) {
        return new Response(`No HttpAction routed for ${url.pathname}`, {
          status: 404,
        });
      }
      const [func] = found;
      const a = httpActionGeneric((ctx: any, a: any) => {
        const testCtx = {
          ...ctx,
          runQuery: byType.query,
          runMutation: byType.mutation,
          runAction: byType.action,
          auth,
        };
        return getHandler(func)(testCtx, a);
      });
      const response = await (
        a as unknown as {
          invokeHttpAction: (request: Request) => Promise<Response>;
        }
      ).invokeHttpAction(new Request(url, init));
      return response;
    },

    // This is needed because when we execute functions
    // we are performing dynamic `import`s, and those
    // are real work that cannot be force-awaited.
    finishInProgressScheduledFunctions: async (): Promise<void> => {
      await waitForInProgressScheduledFunctions();
    },

    finishAllScheduledFunctions: async (
      advanceTimers: () => void,
      maxIterations: number = 100,
    ): Promise<void> => {
      // Wait for all scheduled functions to finish, advancing time in between
      // each function.
      // Stop after a fixed number of iterations to avoid infinite loops.
      for (let i = 0; i < maxIterations; i++) {
        advanceTimers();
        const hadScheduledFunctions =
          await waitForInProgressScheduledFunctions();
        if (!hadScheduledFunctions) {
          return;
        }
      }
      throw new Error(
        "finishAllScheduledFunctions: too many iterations. " +
          "Check for infinitely recursive scheduled functions, " +
          "or increase maxIterations.",
      );
    },
  };
}

function parseArgs(
  args: Record<string, Value> | undefined,
): Record<string, Value> {
  if (args === undefined) {
    return {};
  }
  if (!isSimpleObject(args)) {
    throw new Error(
      `The arguments to a Convex function must be an object. Received: ${
        args as any
      }`,
    );
  }
  return args;
}

function validateReturnValue(
  func: any,
  result: any,
  functionType: "query" | "mutation" | "action",
  functionPath: FunctionPath,
) {
  const returnsValidator = JSON.parse(func.exportReturns());
  if (returnsValidator !== null) {
    try {
      validateValidator(returnsValidator, result);
    } catch (e) {
      throw new Error(
        `Return value validation failed for ${functionType} "${functionPath.udfPath}":\n${(e as Error).message}`,
      );
    }
  }
}

function createFunctionHandle(functionPath: FunctionPath) {
  return `function://${functionPath.componentPath};${functionPath.udfPath}`;
}

function parseFunctionHandle(handle: string) {
  const [componentPath, udfPath] = handle.split("function://")[1].split(";");
  return { componentPath, udfPath };
}

function getFunctionPathFromAddress(
  functionAddress:
    | { name: string; reference: undefined; functionHandle: undefined }
    | { reference: string; name: undefined; functionHandle: undefined }
    | { functionHandle: string; name: undefined; reference: undefined },
): FunctionPath {
  if (functionAddress.functionHandle !== undefined) {
    return parseFunctionHandle(functionAddress.functionHandle);
  }
  if (functionAddress.name !== undefined) {
    return {
      udfPath: functionAddress.name,
      componentPath: getCurrentComponentPath(),
    };
  }
  if (functionAddress.reference !== undefined) {
    // "_reference/childComponent/aggregate/path/to/file/functionName" -> "aggregate/path/to/file"
    const childComponentName = functionAddress.reference.split("/")[2];
    let componentPath = childComponentName;
    if (getCurrentComponentPath().length > 0) {
      componentPath = `${getCurrentComponentPath()}/${componentPath}`;
    }
    // "_reference/childComponent/aggregate/path/to/file/functionName" -> "path/to/file/functionName"
    const functionNameWithSlashes = functionAddress.reference
      .split("/")
      .slice(3)
      .join("/");
    // "path/to/file/functionName" -> "path/to/file:functionName"
    const filepath = functionNameWithSlashes.split("/").slice(0, -1).join("/");
    const functionName = functionNameWithSlashes.split("/").pop();
    return {
      udfPath: `${filepath}:${functionName}`,
      componentPath,
    };
  }
  throw new Error("Function address not supported");
}

async function getFunctionPathFromReference(
  functionReference: FunctionReference<any, any, any, any>,
) {
  // { name: "messages:list" }
  // { reference: "_reference/childComponent/aggregate/path/to/file/functionName" }
  // { functionHandle: "function://<id>/<path>" }
  const functionAddress = getFunctionAddress(functionReference);
  return getFunctionPathFromAddress(functionAddress);
}

type RegisteredFunctions = {
  query: RegisteredQuery<any, any, Value>;
  mutation: RegisteredMutation<any, any, Value>;
  action: RegisteredAction<any, any, any>;
  any: any;
};

type RegisteredFunctionKind = keyof RegisteredFunctions;

function getHandler<Args extends DefaultFunctionArgs, Returns>(
  func: RegisteredAction<any, Args, Returns>,
): (ctx: GenericActionCtx<any>, args: Args) => Promise<Returns>;
function getHandler<Args extends DefaultFunctionArgs, Returns>(
  func: RegisteredMutation<any, Args, Returns>,
): (ctx: GenericActionCtx<any>, args: Args) => Promise<Returns>;
function getHandler<Args extends DefaultFunctionArgs, Returns>(
  func: RegisteredQuery<any, Args, Returns>,
): (ctx: GenericActionCtx<any>, args: Args) => Promise<Returns>;
function getHandler(
  func: PublicHttpAction,
): (ctx: GenericActionCtx<any>, args: Request) => Promise<Response>;
function getHandler(func: any): (ctx: any, args: any) => any {
  return "_handler" in func ? func["_handler"] : func;
}

async function getFunctionFromPath<T extends RegisteredFunctionKind>(
  functionPath: FunctionPath,
  type: T,
): Promise<RegisteredFunctions[T]> {
  // "queries/messages:list" -> ["queries/messages", "list"]
  const [modulePath, maybeExportName] = functionPath.udfPath.split(":");
  const exportName =
    maybeExportName === undefined ? "default" : maybeExportName;

  const module = await getModulesForComponent(functionPath.componentPath)(
    modulePath,
  );

  const func = module[exportName];
  if (func === undefined) {
    throw new Error(
      `Expected a Convex function exported from module "${modulePath}" as \`${exportName}\`, but there is no such export.`,
    );
  }
  if (typeof getHandler(func) !== "function") {
    throw new Error(
      `Expected a Convex function exported from module "${modulePath}" as \`${exportName}\`, but got: ${func}`,
    );
  }
  switch (type) {
    case "query":
      if (!func.isQuery) {
        throw new Error(
          `Expected a query function, but the function exported from module "${modulePath}" as \`${exportName}\` is not a query.`,
        );
      }
      break;
    case "mutation":
      if (!func.isMutation) {
        throw new Error(
          `Expected a mutation function, but the function exported from module "${modulePath}" as \`${exportName}\` is not a mutation.`,
        );
      }
      break;
    case "action":
      if (!func.isAction) {
        throw new Error(
          `Expected an action function, but the function exported from module "${modulePath}" as \`${exportName}\` is not an action.`,
        );
      }
      break;
    case "any":
      break;
    default:
      // eslint-disable-next-line @typescript-eslint/only-throw-error
      throw type satisfies never;
  }
  return func;
}

function simpleHash(string: string) {
  let hash = 0;
  for (let i = 0; i < string.length; i++) {
    const char = string.charCodeAt(i);
    hash = (hash << 5) - hash + char;
    hash = hash & hash; // Convert to 32bit integer
  }
  return hash;
}