fix: Use correct limit when retrying a limit query stream with a cursor

dev/src/reference/query-util.ts

@@ -181,6 +181,7 @@ export class QueryUtil<
     const startTime = Date.now();
     const isExplain = explainOptions !== undefined;
 
+    let numDocumentsReceived = 0;
     let lastReceivedDocument: QueryDocumentSnapshot<
       AppModelType,
       DbModelType
@@ -239,6 +240,7 @@ export class QueryUtil<
           );
         }
 
+        ++numDocumentsReceived;
         callback(undefined, output);
 
         if (proto.done) {
@@ -317,6 +319,12 @@ export class QueryUtil<
                   stream.destroy(err);
                   streamActive.resolve(/* active= */ false);
                 } else if (lastReceivedDocument && retryWithCursor) {
+                  if (query instanceof VectorQuery) {
+                    throw new Error(
+                      'Unimplemented: Vector query does not support cursors yet.'
+                    );
+                  }
+
                   logger(
                     'Query._stream',
                     tag,
@@ -330,12 +338,30 @@ export class QueryUtil<
                   // the query cursor. Note that we do not use backoff here. The
                   // call to `requestStream()` will backoff should the restart
                   // fail before delivering any results.
+                  let newQuery: Query<AppModelType, DbModelType>;
+                  if (!this._queryOptions.limit) {
+                    newQuery = query;
+                  } else {
+                    const newLimit =
+                      this._queryOptions.limit - numDocumentsReceived;
+                    if (
+                      this._queryOptions.limitType === undefined ||
+                      this._queryOptions.limitType === LimitType.First
+                    ) {
+                      newQuery = query.limit(newLimit);
+                    } else {
+                      newQuery = query.limitToLast(newLimit);
+                    }
+                  }
+
                   if (this._queryOptions.requireConsistency) {
-                    request = query
+                    request = newQuery
                       .startAfter(lastReceivedDocument)
                       .toProto(lastReceivedDocument.readTime);
                   } else {
-                    request = query.startAfter(lastReceivedDocument).toProto();
+                    request = newQuery
+                      .startAfter(lastReceivedDocument)
+                      .toProto();
                   }
 
                   // Set lastReceivedDocument to null before each retry attempt to ensure the retry makes progress