convex/src/browser/sync/client.ts

Client for the Convex Cloud

import { v4 as uuidv4 } from "uuid";
import { GenericAPI } from "../../api/index.js";
import { version } from "../../index.js";
import { convexToJson, JSONValue, Value } from "../../values/index.js";
import { logFatalError } from "../logging.js";
import { ActionManager } from "./action_manager.js";
import { LocalSyncState } from "./local_state.js";
import { RequestManager } from "./request_manager.js";
import {
  OptimisticLocalStore,
  OptimisticUpdate,
} from "./optimistic_updates.js";
import {
  OptimisticQueryResults,
  QueryResultsMap,
} from "./optimistic_updates_impl.js";
import {
  ActionRequest,
  AuthError,
  MutationRequest,
  QueryId,
  QueryJournal,
  RequestId,
  ServerMessage,
} from "./protocol.js";
import { QueryResult, RemoteQuerySet } from "./remote_query_set.js";
import { QueryToken, serializePathAndArgs } from "./udf_path_utils.js";
import { ReconnectMetadata, WebSocketManager } from "./web_socket_manager.js";

/**
 * Options for {@link InternalConvexClient}.
 *
 * @public
 */
export interface ClientOptions {
  /**
   * Whether to prompt the user if they have unsaved changes pending
   * when navigating away or closing a web page.
   *
   * This is only possible when the `window` object exists, i.e. in a browser.
   *
   * The default value is `true`.
   */
  unsavedChangesWarning?: boolean;
  /**
   * Specifies an alternate
   * [WebSocket](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket)
   * constructor to use for client communication with the Convex cloud.
   * The default behavior is to use `WebSocket` from the global environment.
   */
  webSocketConstructor?: typeof WebSocket;
}

/**
 * State describing the client's connection with the Convex backend.
 *
 * @public
 */
export type ConnectionState = {
  hasInflightRequests: boolean;
  isWebSocketConnected: boolean;
};

/**
 * An async function returning the JWT-encoded OpenID Connect Identity Token
 * if available.
 * See {@link ConvexReactClient.setAuth}.
 *
 * @public
 */
export type AuthTokenFetcher = () => Promise<string | null | undefined>;

const DEFAULT_OPTIONS = {
  unsavedChangesWarning: true,
};

/**
 * Low-level client for directly integrating state management libraries
 * with Convex.
 *
 * Most developers should use higher level clients, like
 * the {@link ConvexHttpClient} or the React hook based {@link react.ConvexReactClient}.
 *
 * @public
 */
export class InternalConvexClient {
  private readonly state: LocalSyncState;
  private readonly requestManager: RequestManager;
  private readonly actionManager: ActionManager;
  private readonly webSocketManager: WebSocketManager;
  private remoteQuerySet: RemoteQuerySet;
  private readonly optimisticQueryResults: OptimisticQueryResults;
  private readonly onTransition: (updatedQueries: QueryToken[]) => void;
  private nextRequestId: RequestId;
  private readonly sessionId: string;
  private fetchToken: null | (() => Promise<string | null | undefined>) = null;

  /**
   * @param address - The url of your Convex deployment, often provided
   * by an environment variable. E.g. `https://small-mouse-123.convex.cloud`.
   * @param onTransition - A callback receiving an array of query tokens
   * corresponding to query results that have changed.
   * @param options - See {@link ClientOptions} for a full description.
   */
  constructor(
    address: string,
    onTransition: (updatedQueries: QueryToken[]) => void,
    options?: ClientOptions
  ) {
    if (typeof address === "object") {
      throw new Error(
        "Passing a ClientConfig object is no longer supported. Pass the URL of the Convex deployment as a string directly."
      );
    }
    options = { ...DEFAULT_OPTIONS, ...options };
    let webSocketConstructor = options.webSocketConstructor;
    if (!webSocketConstructor && typeof WebSocket === "undefined") {
      throw new Error(
        "No WebSocket global variable defined! To use Convex in an environment without WebSocket try the HTTP client: https://docs.convex.dev/api/classes/browser.ConvexHttpClient"
      );
    }
    webSocketConstructor = webSocketConstructor || WebSocket;

    // Substitute http(s) with ws(s)
    const i = address.search("://");
    if (i === -1) {
      throw new Error("Provided address was not an absolute URL.");
    }
    const origin = address.substring(i + 3); // move past the double slash
    const protocol = address.substring(0, i);
    let wsProtocol;
    if (protocol === "http") {
      wsProtocol = "ws";
    } else if (protocol === "https") {
      wsProtocol = "wss";
    } else {
      throw new Error(`Unknown parent protocol ${protocol}`);
    }
    const wsUri = `${wsProtocol}://${origin}/api/${version}/sync`;

    this.state = new LocalSyncState();
    this.remoteQuerySet = new RemoteQuerySet(queryId =>
      this.state.queryPath(queryId)
    );
    this.requestManager = new RequestManager();
    this.actionManager = new ActionManager();
    this.optimisticQueryResults = new OptimisticQueryResults();
    this.onTransition = onTransition;
    this.nextRequestId = 0;
    this.sessionId = uuidv4();

    const { unsavedChangesWarning } = options;
    if (typeof window === "undefined" && unsavedChangesWarning) {
      throw new Error(
        "unsavedChangesWarning enabled, but no window object found! Navigating away from the page could cause in-flight mutations to be dropped. Pass {unsavedChangesWarning: false} in Convex client options to disable this feature."
      );
    }

    // Listen for tab close events and notify the user on unsaved changes.
    unsavedChangesWarning &&
      window.addEventListener("beforeunload", e => {
        if (
          this.requestManager.hasIncompleteRequests() ||
          this.actionManager.hasInflightActions()
        ) {
          // There are 3 different ways to trigger this pop up so just try all of
          // them.

          e.preventDefault();
          // This confirmation message doesn't actually appear in most modern
          // browsers but we tried.
          const confirmationMessage =
            "Are you sure you want to leave? Your changes may not be saved.";
          (e || window.event).returnValue = confirmationMessage;
          return confirmationMessage;
        }
      });

    this.webSocketManager = new WebSocketManager(
      wsUri,
      (reconnectMetadata: ReconnectMetadata) => {
        // We have a new WebSocket!

        this.webSocketManager.sendMessage({
          ...reconnectMetadata,
          type: "Connect",
          sessionId: this.sessionId,
        });

        // Throw out our remote query, reissue queries
        // and outstanding mutations, and reauthenticate.
        this.remoteQuerySet = new RemoteQuerySet(queryId =>
          this.state.queryPath(queryId)
        );
        const [querySetModification, authModification] = this.state.restart();
        if (authModification) {
          this.webSocketManager.sendMessage(authModification);
        }
        this.webSocketManager.sendMessage(querySetModification);
        this.actionManager.restart();
        for (const message of this.requestManager.restart()) {
          this.webSocketManager.sendMessage(message);
        }
      },
      (serverMessage: ServerMessage) => {
        switch (serverMessage.type) {
          case "Transition": {
            this.remoteQuerySet.transition(serverMessage);
            this.state.saveQueryJournals(serverMessage);
            // No optimistic updates for workflows.
            const completedMutations = this.requestManager.removeCompleted(
              this.remoteQuerySet.timestamp()
            );
            this.notifyOnQueryResultChanges(completedMutations);
            break;
          }
          case "MutationResponse": {
            const completedMutationId =
              this.requestManager.onResponse(serverMessage);
            if (completedMutationId) {
              this.notifyOnQueryResultChanges(new Set([completedMutationId]));
            }
            break;
          }
          case "ActionResponse": {
            this.actionManager.onResponse(serverMessage);
            break;
          }
          case "AuthError": {
            this.tryToReauthenticate(serverMessage)
              .then()
              .catch(error => {
                logFatalError(error);
                // TODO(CX-3070): This ignores a failed Promise
                void this.webSocketManager.stop();
              });
            break;
          }
          case "FatalError": {
            const error = logFatalError(serverMessage.error);
            void this.webSocketManager.stop();
            throw error;
          }
          case "Ping":
            break; // do nothing
          default: {
            const _typeCheck: never = serverMessage;
          }
        }
      },
      webSocketConstructor
    );
  }

  /**
   * Compute the current query results based on the remoteQuerySet and the
   * current optimistic updates and call `onTransition` for all the changed
   * queries.
   *
   * @param completedMutations - A set of mutation IDs whose optimistic updates
   * are no longer needed.
   */
  private notifyOnQueryResultChanges(completedMutations: Set<RequestId>) {
    const remoteQueryResults: Map<QueryId, QueryResult> =
      this.remoteQuerySet.remoteQueryResults();
    const queryTokenToValue: QueryResultsMap = new Map();
    for (const [queryId, result] of remoteQueryResults) {
      const queryToken = this.state.queryToken(queryId);
      // It's possible that we've already unsubscribed to this query but
      // the server hasn't learned about that yet. If so, ignore this one.

      if (queryToken !== null) {
        const query = {
          result,
          udfPath: this.state.queryPath(queryId)!,
          args: this.state.queryArgs(queryId)!,
        };
        queryTokenToValue.set(queryToken, query);
      }
    }

    this.onTransition(
      this.optimisticQueryResults.ingestQueryResultsFromServer(
        queryTokenToValue,
        completedMutations
      )
    );
  }

  async setAuth(fetchToken: AuthTokenFetcher) {
    this.fetchToken = fetchToken;
    const token = await fetchToken();
    if (token) {
      this.authenticate(token);
    } else if (this.state.hasAuth()) {
      this.clearAuth();
    }
  }

  // Stop the webSocket so that we don't retry with bad auth
  private async tryToReauthenticate(serverMessage: AuthError) {
    if (!this.fetchToken) {
      throw new Error(serverMessage.error);
    }
    console.log("Attempting to reauthenticate");
    await this.webSocketManager.pause();
    const token = await this.fetchToken();
    if (token && this.state.isNewAuth(token)) {
      this.state.setAuth(token);
    } else if (this.state.hasAuth()) {
      console.log("Reauthentication failed, clearing auth");
      this.state.clearAuth();
    }
    await this.webSocketManager.resume();
  }

  private authenticate(token: string) {
    const message = this.state.setAuth(token);
    this.webSocketManager.sendMessage(message);
  }

  /** @internal */
  setAdminAuth(value: string) {
    const message = this.state.setAdminAuth(value);
    this.webSocketManager.sendMessage(message);
  }

  clearAuth() {
    const message = this.state.clearAuth();
    this.webSocketManager.sendMessage(message);
  }

  /**
   * Subscribe to a query function.
   *
   * Whenever this query's result changes, the `onTransition` callback
   * passed into the constructor will be called.
   *
   * @param name - The name of the query.
   * @param args - An array of the arguments to the query.
   * @param journal - An (optional) journal produced from a previous
   * execution of this query function. Note that if this query function with
   * these arguments has already been requested the journal will have no effect.
   * @returns An object containing a {@link QueryToken} corresponding to this
   * query and an `unsubscribe` callback.
   */
  subscribe(
    name: string,
    args: any[],
    journal?: QueryJournal
  ): { queryToken: QueryToken; unsubscribe: () => void } {
    // `subscribe` used to collect the arguments with a rest operator
    // (like `...args`). Double check that it's an array to make sure developers
    // have updated their code.
    if (!Array.isArray(args)) {
      throw new Error(
        `Query arguments to \`InternalConvexClient.subcribe\` must be an array. Received ${args}.`
      );
    }

    const { modification, queryToken, unsubscribe } = this.state.subscribe(
      name,
      args,
      journal
    );
    if (modification !== null) {
      this.webSocketManager.sendMessage(modification);
    }
    // TODO: Use FinalizationRegistry?
    return {
      queryToken,
      unsubscribe: () => {
        const modification = unsubscribe();
        if (modification) {
          this.webSocketManager.sendMessage(modification);
        }
      },
    };
  }

  /**
   * A query result based only on the current, local state.
   *
   * The only way this will return a value is if we're already subscribed to the
   * query or its value has been set optimistically.
   */
  localQueryResult(udfPath: string, args: any[]): Value | undefined {
    const queryToken = serializePathAndArgs(udfPath, args);
    return this.optimisticQueryResults.queryResult(queryToken);
  }

  /**
   * Retrieve the current {@link QueryJournal} for this query function.
   *
   * If we have not yet received a result for this query, this will be `undefined`.
   *
   * @param name - The name of the query.
   * @param args - An array of arguments to this query.
   * @returns The query's {@link QueryJournal} or `undefined`.
   */
  queryJournal(name: string, args: any[]): QueryJournal | undefined {
    const queryToken = serializePathAndArgs(name, args);
    return this.state.queryJournal(queryToken);
  }

  /**
   * Get the current {@link ConnectionState} between the client and the Convex
   * backend.
   *
   * @returns The {@link ConnectionState} with the Convex backend.
   */
  connectionState(): ConnectionState {
    return {
      hasInflightRequests:
        this.requestManager.hasInflightRequests() ||
        this.actionManager.hasInflightActions(),
      isWebSocketConnected: this.webSocketManager.socketState() === "ready",
    };
  }

  async mutate<Args extends any[]>(
    udfPath: string,
    args: Args,
    optimisticUpdate: OptimisticUpdate<GenericAPI, Args> | null = null
  ): Promise<any> {
    const requestId = this.nextRequestId;
    this.nextRequestId++;

    if (optimisticUpdate !== null) {
      const wrappedUpdate = (localQueryStore: OptimisticLocalStore) => {
        optimisticUpdate(localQueryStore, ...args);
      };
      const changedQueries = this.optimisticQueryResults.applyOptimisticUpdate(
        wrappedUpdate,
        requestId
      );
      this.onTransition(changedQueries);
    }

    const message: MutationRequest = {
      type: "Mutation",
      requestId,
      udfPath,
      args: convexToJson(args) as JSONValue[],
    };
    const result = this.requestManager.request(message);
    this.webSocketManager.sendMessage(message);
    return result;
  }

  async action<Args extends any[]>(udfPath: string, args: Args): Promise<any> {
    const requestId = this.nextRequestId;
    this.nextRequestId++;

    const message: ActionRequest = {
      type: "Action",
      requestId,
      udfPath,
      args: convexToJson(args) as JSONValue[],
    };

    const result = this.actionManager.request(message);
    this.webSocketManager.sendMessage(message);
    return result;
  }

  async close(): Promise<void> {
    return this.webSocketManager.stop();
  }
}