From 21a398c013aaaabb9b1527191536f3b7ce530346 Mon Sep 17 00:00:00 2001
From: milaGGL <107142260+milaGGL@users.noreply.github.com>
Date: Tue, 25 Jun 2024 15:15:08 -0400
Subject: [PATCH] sync stored metadata state

---
 .../firestore/src/core/sync_engine_impl.ts    |  8 +++---
 .../src/local/shared_client_state.ts          | 25 ++++++++++++++-----
 .../unit/specs/listen_source_spec.test.ts     | 16 ++++++------
 3 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/packages/firestore/src/core/sync_engine_impl.ts b/packages/firestore/src/core/sync_engine_impl.ts
index bba07f4f4bc..fa40f302bd8 100644
--- a/packages/firestore/src/core/sync_engine_impl.ts
+++ b/packages/firestore/src/core/sync_engine_impl.ts
@@ -353,9 +353,11 @@ async function allocateTargetAndMaybeListen(
   // not registering it in shared client state, and directly calculate initial snapshots and
   // subsequent updates from cache. Otherwise, register the target ID with local Firestore client
   // as active watch target.
-  const status: QueryTargetState = shouldListenToRemote
-    ? syncEngineImpl.sharedClientState.addLocalQueryTarget(targetId)
-    : 'not-current';
+  const status: QueryTargetState =
+    syncEngineImpl.sharedClientState.addLocalQueryTarget(
+      targetId,
+      shouldListenToRemote
+    );
 
   let viewSnapshot;
   if (shouldInitializeView) {
diff --git a/packages/firestore/src/local/shared_client_state.ts b/packages/firestore/src/local/shared_client_state.ts
index 0d0e430c12f..dd7f2b465ef 100644
--- a/packages/firestore/src/local/shared_client_state.ts
+++ b/packages/firestore/src/local/shared_client_state.ts
@@ -104,7 +104,10 @@ export interface SharedClientState {
    * If the target id is already associated with local client, the method simply
    * returns its `QueryTargetState`.
    */
-  addLocalQueryTarget(targetId: TargetId): QueryTargetState;
+  addLocalQueryTarget(
+    targetId: TargetId,
+    addActiveTargetId?: boolean
+  ): QueryTargetState;
 
   /** Removes the Query Target ID association from the local client. */
   removeLocalQueryTarget(targetId: TargetId): void;
@@ -655,7 +658,10 @@ export class WebStorageSharedClientState implements SharedClientState {
     this.removeMutationState(batchId);
   }
 
-  addLocalQueryTarget(targetId: TargetId): QueryTargetState {
+  addLocalQueryTarget(
+    targetId: TargetId,
+    addActiveTargetId: boolean = true
+  ): QueryTargetState {
     let queryState: QueryTargetState = 'not-current';
 
     // Lookup an existing query state if the target ID was already registered
@@ -676,8 +682,10 @@ export class WebStorageSharedClientState implements SharedClientState {
       }
     }
 
-    this.localClientState.addQueryTarget(targetId);
-    this.persistClientState();
+    if (addActiveTargetId) {
+      this.localClientState.addQueryTarget(targetId);
+      this.persistClientState();
+    }
 
     return queryState;
   }
@@ -1110,8 +1118,13 @@ export class MemorySharedClientState implements SharedClientState {
     // No op.
   }
 
-  addLocalQueryTarget(targetId: TargetId): QueryTargetState {
-    this.localState.addQueryTarget(targetId);
+  addLocalQueryTarget(
+    targetId: TargetId,
+    addActiveTargetId: boolean = true
+  ): QueryTargetState {
+    if (addActiveTargetId) {
+      this.localState.addQueryTarget(targetId);
+    }
     return this.queryState[targetId] || 'not-current';
   }
 
diff --git a/packages/firestore/test/unit/specs/listen_source_spec.test.ts b/packages/firestore/test/unit/specs/listen_source_spec.test.ts
index ee2a9cf5944..cb5a89bfa75 100644
--- a/packages/firestore/test/unit/specs/listen_source_spec.test.ts
+++ b/packages/firestore/test/unit/specs/listen_source_spec.test.ts
@@ -491,10 +491,10 @@ describeSpec('Listens source options:', [], () => {
           // Listen to cache in secondary clients
           .client(1)
           .userListensToCache(query1)
-          .expectEvents(query1, { added: [docA], fromCache: true })
+          .expectEvents(query1, { added: [docA] })
           .client(2)
           .userListensToCache(query1)
-          .expectEvents(query1, { added: [docA], fromCache: true })
+          .expectEvents(query1, { added: [docA] })
           // Updates in the primary client notifies listeners sourcing from cache
           // in secondary clients.
           .client(0)
@@ -748,7 +748,7 @@ describeSpec('Listens source options:', [], () => {
           //Listen to cache in secondary client
           .client(1)
           .userListensToCache(limitToLast)
-          .expectEvents(limitToLast, { added: [docB, docA], fromCache: true })
+          .expectEvents(limitToLast, { added: [docB, docA] })
           // Watch sends document changes
           .client(0)
           .watchSends({ affects: [limit] }, docC)
@@ -794,7 +794,7 @@ describeSpec('Listens source options:', [], () => {
           // Listen to cache in primary client
           .client(0)
           .userListensToCache(limitToLast)
-          .expectEvents(limitToLast, { added: [docB, docA], fromCache: true })
+          .expectEvents(limitToLast, { added: [docB, docA] })
           // Watch sends document changes
           .watchSends({ affects: [limit] }, docC)
           .watchSnapshots(2000)
@@ -825,16 +825,18 @@ describeSpec('Listens source options:', [], () => {
           // Listen to cache in secondary client
           .client(1)
           .userListensToCache(query1)
-          .expectEvents(query1, { added: [docA], fromCache: true })
+          .expectEvents(query1, { added: [docA] })
+
           .client(0)
           .userUnlistens(query1)
+          .watchRemoves(query1)
+
           .userSets('collection/b', { key: 'b' })
           // The other listener should work as expected
           .client(1)
           .expectEvents(query1, {
             hasPendingWrites: true,
-            added: [docB],
-            fromCache: true
+            added: [docB]
           })
           .userUnlistensToCache(query1)
       );