Don't crash when optimistic updates exist on queries for which no server result exists yet (#38020)

Fix a bug with setting optimistic updates for recently- or not-yet-subscribed-to queries.

Previously the client threw an error and aborted handling a server update with "Cannot read properties of undefined (reading 'result')" in `BaseConvexClient.notifyOnQueryResultChanges` (among other more obfuscated frames) when the client set an optimistic update for a query for which it did not yet have a value from the server.

GitOrigin-RevId: 1fb59b02d5f14750a4d8bdf8c1bb259c4dfdccb5

Author: thomasballinger

src/browser/sync/client.ts

@@ -185,7 +185,7 @@ export type QueryModification =
  * Object describing a transition passed into the `onTransition` handler.
  *
  * These can be from receiving a transition from the server, or from applying an
- * optimistc update locally.
+ * optimistic update locally.
  *
  * @public
  */
@@ -510,20 +510,28 @@ export class BaseConvexClient {
         queryTokenToValue.set(queryToken, query);
       }
     }
+
+    // Query tokens that are new (because of new server results or new local optimistic updates)
+    // or differ from old values (because of changes from local optimistic updates or new results
+    // from the server).
     const changedQueryTokens =
       this.optimisticQueryResults.ingestQueryResultsFromServer(
         queryTokenToValue,
         new Set(completedRequests.keys()),
       );
 
     this.handleTransition({
-      queries: changedQueryTokens.map((token) => ({
-        token,
-        modification: {
-          kind: "Updated",
-          result: queryTokenToValue.get(token)!.result,
-        },
-      })),
+      queries: changedQueryTokens.map((token) => {
+        const optimisticResult =
+          this.optimisticQueryResults.rawQueryResult(token);
+        return {
+          token,
+          modification: {
+            kind: "Updated",
+            result: optimisticResult!.result,
+          },
+        };
+      }),
       reflectedMutations: Array.from(completedRequests).map(
         ([requestId, result]) => ({
           requestId,

src/browser/sync/client_node.test.ts

@@ -14,8 +14,10 @@ import {
 import {
   encodeServerMessage,
   nodeWebSocket,
+  UpdateQueue,
   withInMemoryWebSocket,
 } from "./client_node_test_helpers.js";
+import { FunctionArgs, makeFunctionReference } from "../../server/index.js";
 
 test("BaseConvexClient protocol in node", async () => {
   await withInMemoryWebSocket(async ({ address, receive }) => {
@@ -190,3 +192,175 @@ test("maxObservedTimestamp is updated on mutation and transition", async () => {
     await client.close();
   });
 });
+
+const apiQueriesA = makeFunctionReference<"query", {}, string>("queries:a");
+const apiQueriesB = makeFunctionReference<"query", {}, string>("queries:b");
+
+const _apiMutationsZ = makeFunctionReference<"mutation", {}>("mutations:z");
+
+/**
+ * Regression test for
+ * - subscribing to query a
+ * - running a mutation that sets an optimistic update for queries a and b
+ * - receiving an update for a
+ */
+test("Setting optimistic updates for queries that have not yet been subscribed to", async () => {
+  await withInMemoryWebSocket(async ({ address, receive, send }) => {
+    const q = new UpdateQueue(10);
+
+    const client = new BaseConvexClient(
+      address,
+      (queryTokens) => {
+        q.onTransition(client)(queryTokens);
+      },
+      {
+        webSocketConstructor: nodeWebSocket,
+        unsavedChangesWarning: false,
+        verbose: true,
+      },
+    );
+
+    client.subscribe("queries:a", {});
+
+    expect((await receive()).type).toEqual("Connect");
+    const modify = await receive();
+
+    expect(modify.type).toEqual("ModifyQuerySet");
+    if (modify.type !== "ModifyQuerySet") {
+      return;
+    }
+    expect(modify.modifications.length).toBe(1);
+    expect(modify.modifications).toEqual([
+      {
+        type: "Add",
+        queryId: 0,
+        udfPath: "queries:a",
+        args: [{}],
+      },
+    ]);
+
+    // Now that we're subscribed to queries:a,
+    // run a mutation that, optimistically and on the server,
+    // - modifies q1
+    // - modifies q2
+
+    const mutP = client.mutation(
+      "mutations:z",
+      {},
+      {
+        optimisticUpdate: (
+          localStore,
+          _args: FunctionArgs<typeof _apiMutationsZ>,
+        ) => {
+          const curA = localStore.getQuery(apiQueriesA, {});
+          localStore.setQuery(
+            apiQueriesA,
+            {},
+            curA === undefined ? "a local" : `${curA} with a local applied`,
+          );
+          const curB = localStore.getQuery(apiQueriesB, {});
+          localStore.setQuery(
+            apiQueriesB,
+            {},
+            curB === undefined ? "b local" : `${curB} with b local applied`,
+          );
+        },
+      },
+    );
+
+    // Synchronously, the local store should update and the changes should be broadcast.
+    expect(client.localQueryResult("queries:a", {})).toEqual("a local");
+    // We haven't actually subscribed to this query but it had a value set in an optimistic update.
+    expect(client.localQueryResult("queries:b", {})).toEqual("b local");
+    const update1 = await q.updatePromises[0];
+    expect(q.updates).toHaveLength(1);
+    expect(update1).toEqual({
+      '{"udfPath":"queries:a","args":{}}': "a local",
+      '{"udfPath":"queries:b","args":{}}': "b local",
+    });
+
+    // Now a transition arrives containing only an update to query a.
+    // This previously crashed this execution context.
+    send({
+      type: "Transition",
+      startVersion: {
+        querySet: 0,
+        identity: 0,
+        ts: Long.fromNumber(0),
+      },
+      endVersion: {
+        querySet: 1,
+        identity: 0,
+        ts: Long.fromNumber(100),
+      },
+      modifications: [
+        {
+          type: "QueryUpdated",
+          queryId: 0,
+          value: "a server",
+          logLines: [],
+          journal: null,
+        },
+      ],
+    });
+
+    const update2 = await q.updatePromises[1];
+    expect(update2).toEqual({
+      '{"udfPath":"queries:a","args":{}}': "a server with a local applied",
+      '{"udfPath":"queries:b","args":{}}': "b local",
+    });
+    expect(q.allResults).toEqual({
+      '{"udfPath":"queries:a","args":{}}': "a server with a local applied",
+      '{"udfPath":"queries:b","args":{}}': "b local",
+    });
+    expect(q.updates).toHaveLength(2);
+
+    const mutationRequest = await receive();
+    expect(mutationRequest.type).toEqual("Mutation");
+    expect(mutationRequest).toEqual({
+      type: "Mutation",
+      requestId: 0,
+      udfPath: "mutations:z",
+      args: [{}],
+    });
+
+    // Now the server sends:
+
+    // 1. MutationResponse saying the mutation has run
+    send({
+      type: "MutationResponse",
+      requestId: 0,
+      success: true,
+      result: null,
+      ts: Long.fromNumber(200), // "ZDhuVB3CRxg=", in example
+      logLines: [],
+    });
+
+    // 2. Transition bringing us up to date with the mutation
+    send({
+      type: "Transition",
+      startVersion: { querySet: 1, identity: 0, ts: Long.fromNumber(100) },
+      endVersion: { querySet: 1, identity: 0, ts: Long.fromNumber(200) },
+      modifications: [
+        {
+          type: "QueryUpdated",
+          queryId: 0,
+          value: "a server",
+          logLines: [],
+          journal: null,
+        },
+      ],
+    });
+
+    expect(await q.updatePromises[2]).toEqual({
+      '{"udfPath":"queries:a","args":{}}': "a server",
+      // Now there's no more optimistic value for b!
+      '{"udfPath":"queries:b","args":{}}': undefined,
+    });
+
+    // After all that the mutation should resolve.
+    await mutP;
+
+    await client.close();
+  }, true);
+});

src/browser/sync/client_node_test_helpers.ts

@@ -8,6 +8,8 @@ import WebSocket, { WebSocketServer } from "ws";
 export const nodeWebSocket = WebSocket as unknown as typeof window.WebSocket;
 
 import { ClientMessage, ServerMessage } from "./protocol.js";
+import { QueryToken } from "./udf_path_utils.js";
+import { BaseConvexClient } from "./client.js";
 
 export type InMemoryWebSocketTest = (args: {
   address: string;
@@ -104,3 +106,50 @@ function encodeLong(n: Long) {
   const integerBytes = Uint8Array.from(n.toBytesLE());
   return Base64.fromByteArray(integerBytes);
 }
+
+/**
+ * const q = new UpdateQueue();
+ * const client = new Client("http://...", q.onTransition);
+ *
+ * await q.updatePromises[3];
+ *
+ */
+export class UpdateQueue {
+  updateResolves: ((v: Record<QueryToken, any>) => void)[];
+  updatePromises: Promise<Record<QueryToken, any>>[];
+  updates: Record<QueryToken, any>[];
+  allResults: Record<QueryToken, any>;
+  nextIndex: number;
+
+  constructor(maxLength = 10) {
+    this.updateResolves = [];
+    this.updatePromises = [];
+    this.allResults = {};
+    this.updates = [];
+    this.nextIndex = 0;
+
+    let nextResolve: (v: Record<QueryToken, any>) => void;
+    let nextPromise: Promise<Record<QueryToken, any>>;
+
+    for (let i = 0; i < maxLength; i++) {
+      nextPromise = new Promise((r) => {
+        nextResolve = r;
+      });
+      this.updateResolves.push(nextResolve!);
+      this.updatePromises.push(nextPromise);
+    }
+  }
+
+  onTransition =
+    (client: BaseConvexClient) => (updatedQueryTokens: QueryToken[]) => {
+      const update: Record<QueryToken, any> = {};
+      for (const queryToken of updatedQueryTokens) {
+        const value = client.localQueryResultByToken(queryToken);
+        update[queryToken] = value;
+        this.allResults[queryToken] = value;
+      }
+      this.updateResolves[this.nextIndex](update);
+      this.updates.push(update);
+      this.nextIndex++;
+    };
+}

src/browser/sync/optimistic_updates_impl.ts

@@ -131,7 +131,8 @@ type OptimisticUpdateAndId = {
 type Query = {
   // undefined means the query was set to be loading (undefined) in an optimistic update.
   // Note that we can also have queries not present in the QueryResultMap
-  // at all because they are still loading from the server.
+  // at all because they are still loading from the server and have no optimistic update
+  // setting an optimistic value in advance.
   result: FunctionResult | undefined;
   udfPath: string;
   args: Record<string, Value>;
@@ -152,6 +153,9 @@ export class OptimisticQueryResults {
     this.optimisticUpdates = [];
   }
 
+  /**
+   * Apply all optimistic updates on top of server query results
+   */
   ingestQueryResultsFromServer(
     serverQueryResults: QueryResultsMap,
     optimisticUpdatesToDrop: Set<RequestId>,
@@ -197,6 +201,13 @@ export class OptimisticQueryResults {
     return localStore.modifiedQueries;
   }
 
+  /**
+   * @internal
+   */
+  rawQueryResult(queryToken: QueryToken): Query | undefined {
+    return this.queryResults.get(queryToken);
+  }
+
   queryResult(queryToken: QueryToken): Value | undefined {
     const query = this.queryResults.get(queryToken);
     if (query === undefined) {

src/browser/sync/remote_query_set.ts

@@ -7,6 +7,10 @@ import { FunctionResult } from "./function_result.js";
 /**
  * A represention of the query results we've received on the current WebSocket
  * connection.
+ *
+ * Queries you won't find here include:
+ * - queries which have been requested, but no query transition has been received yet for
+ * - queries which are populated only though active optimistic updates, but are not subscribed to
  */
 export class RemoteQuerySet {
   private version: StateVersion;