Update schema version when modifying HEAD with a CLEAR_ALL.

PiperOrigin-RevId: 251490149
Change-Id: Ifc3daaa24e498c8a48bdfbbca805d743d9021149
diff --git a/src/main/java/com/google/android/libraries/feed/api/internal/common/Model.java b/src/main/java/com/google/android/libraries/feed/api/internal/common/Model.java
index f6e3d14..d8c49a8 100644
--- a/src/main/java/com/google/android/libraries/feed/api/internal/common/Model.java
+++ b/src/main/java/com/google/android/libraries/feed/api/internal/common/Model.java
@@ -34,7 +34,11 @@
   }
 
   public static Model of(List<StreamDataOperation> streamDataOperations) {
-    return new Model(streamDataOperations, CURRENT_SCHEMA_VERSION);
+    return of(streamDataOperations, CURRENT_SCHEMA_VERSION);
+  }
+
+  public static Model of(List<StreamDataOperation> streamDataOperations, int schemaVersion) {
+    return new Model(streamDataOperations, schemaVersion);
   }
 
   public static Model empty() {
diff --git a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImpl.java b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImpl.java
index bee56f3..d602ab9 100644
--- a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImpl.java
+++ b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImpl.java
@@ -77,10 +77,15 @@
   public void updateSession(
       boolean clearHead,
       List<StreamStructure> streamStructures,
+      int schemaVersion,
       /*@Nullable*/ MutationContext mutationContext) {
     ElapsedTimeTracker timeTracker = timingUtils.getElapsedTimeTracker(TAG);
     updateCount++;
 
+    if (clearHead) {
+      this.schemaVersion = schemaVersion;
+    }
+
     StreamToken mutationSourceToken =
         mutationContext != null ? mutationContext.getContinuationToken() : null;
     if (mutationSourceToken != null) {
diff --git a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/Session.java b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/Session.java
index 836abeb..2cfd000 100644
--- a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/Session.java
+++ b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/Session.java
@@ -27,6 +27,7 @@
   void updateSession(
       boolean clearHead,
       List<StreamStructure> streamStructures,
+      int schemaVersion,
       /*@Nullable*/ MutationContext mutationContext);
 
   boolean invalidateOnResetHead();
diff --git a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCache.java b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCache.java
index 10e2157..7f6002f 100644
--- a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCache.java
+++ b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCache.java
@@ -340,15 +340,21 @@
     }
   }
 
-  /** Updates the last time content was added to HEAD and writes data to disk. */
-  void updateHeadLastAddedTimeMillis(long lastAddedTimeMillis) {
+  /**
+   * Updates the last time content was added to HEAD and its schema version and writes data to disk.
+   */
+  void updateHeadMetadata(long lastAddedTimeMillis, int schemaVersion) {
     threadUtils.checkNotMainThread();
     synchronized (lock) {
       SessionMetadata metadata = sessionsMetadata.get(head.getSessionId());
       SessionMetadata.Builder builder =
           metadata == null ? SessionMetadata.newBuilder() : metadata.toBuilder();
       sessionsMetadata.put(
-          head.getSessionId(), builder.setLastAddedTimeMillis(lastAddedTimeMillis).build());
+          head.getSessionId(),
+          builder
+              .setLastAddedTimeMillis(lastAddedTimeMillis)
+              .setSchemaVersion(schemaVersion)
+              .build());
     }
     updatePersistedSessionsMetadata();
   }
diff --git a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImpl.java b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImpl.java
index 23e9530..293093e 100644
--- a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImpl.java
+++ b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImpl.java
@@ -141,6 +141,7 @@
   public void updateSession(
       boolean clearHead,
       List<StreamStructure> streamStructures,
+      int schemaVersion,
       /*@Nullable*/ MutationContext mutationContext) {
     String localSessionId = Validators.checkNotNull(sessionId);
     if (clearHead) {
diff --git a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutation.java b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutation.java
index cc61d23..905c696 100644
--- a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutation.java
+++ b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutation.java
@@ -225,7 +225,7 @@
     private final List<StreamStructure> streamStructures = new ArrayList<>();
 
     @VisibleForTesting boolean clearedHead = false;
-    private List<StreamDataOperation> dataOperations;
+    private Model model;
 
     private MutationCommitter(
         String task,
@@ -264,8 +264,8 @@
         }
         return;
       }
-      dataOperations = updateResults.getValue().streamDataOperations;
-      for (StreamDataOperation operation : dataOperations) {
+      model = updateResults.getValue();
+      for (StreamDataOperation operation : model.streamDataOperations) {
         if (operation.getStreamStructure().getOperation() == Operation.CLEAR_ALL) {
           clearedHead = true;
           break;
@@ -302,7 +302,7 @@
       ContentMutation contentMutation = store.editContent();
       SemanticPropertiesMutation semanticPropertiesMutation = store.editSemanticProperties();
       ActionPropertiesMutation actionPropertiesMutation = store.editActionProperties();
-      for (StreamDataOperation dataOperation : dataOperations) {
+      for (StreamDataOperation dataOperation : model.streamDataOperations) {
         Operation operation = dataOperation.getStreamStructure().getOperation();
         if (operation == Operation.CLEAR_ALL) {
           streamStructures.add(dataOperation.getStreamStructure());
@@ -368,7 +368,7 @@
             }
             contentCache.finishMutation();
           });
-      timeTracker.stop("", "contentUpdate", "items", dataOperations.size());
+      timeTracker.stop("", "contentUpdate", "items", model.streamDataOperations.size());
     }
 
     private void commitSessionUpdates() {
@@ -399,8 +399,9 @@
         if (session == head) {
           long updateTime = clock.currentTimeMillis();
           if (clearedHead) {
-            session.updateSession(clearedHead, streamStructures, mutationContext);
-            sessionCache.updateHeadLastAddedTimeMillis(updateTime);
+            session.updateSession(
+                clearedHead, streamStructures, model.schemaVersion, mutationContext);
+            sessionCache.updateHeadMetadata(updateTime, model.schemaVersion);
             schedulerApi.onReceiveNewContent(updateTime);
             if (knownContentListener != null) {
               knownContentListener.onNewContentReceived(/* isNewRefresh */ true, updateTime);
@@ -415,7 +416,7 @@
           }
         }
         Logger.i(TAG, "Update Session %s", session.getSessionId());
-        session.updateSession(clearedHead, streamStructures, mutationContext);
+        session.updateSession(clearedHead, streamStructures, model.schemaVersion, mutationContext);
       }
       timeTracker.stop("", "sessionUpdate", "sessions", updates.size());
     }
diff --git a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImpl.java b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImpl.java
index b6bcdb7..9cb0c62 100644
--- a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImpl.java
+++ b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImpl.java
@@ -60,6 +60,7 @@
   public void updateSession(
       boolean clearHead,
       List<StreamStructure> streamStructures,
+      int schemaVersion,
       /*@Nullable*/ MutationContext mutationContext) {
     String localSessionId = Validators.checkNotNull(sessionId);
     Logger.i(
diff --git a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/testing/AbstractSessionImplTest.java b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/testing/AbstractSessionImplTest.java
index 543e449..ba23743 100644
--- a/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/testing/AbstractSessionImplTest.java
+++ b/src/main/java/com/google/android/libraries/feed/feedsessionmanager/internal/testing/AbstractSessionImplTest.java
@@ -39,6 +39,7 @@
 /** Abstract class implementing many of the core SessionImpl tests. */
 public abstract class AbstractSessionImplTest {
   protected static final String TEST_SESSION_ID = "TEST$1";
+  protected static final int SCHEMA_VERSION = 1;
 
   protected final ContentIdGenerators contentIdGenerators = new ContentIdGenerators();
 
@@ -80,7 +81,7 @@
   @Test
   public void testUpdateSession_notFullyInitialized() {
     SessionImpl session = getSessionImpl();
-    assertThatRunnable(() -> session.updateSession(false, new ArrayList<>(), null))
+    assertThatRunnable(() -> session.updateSession(false, new ArrayList<>(), SCHEMA_VERSION, null))
         .throwsAnExceptionOfType(NullPointerException.class);
   }
 
@@ -96,7 +97,7 @@
 
     // 1 clear, 3 features
     assertThat(streamStructures).hasSize(4);
-    session.updateSession(false, streamStructures, null);
+    session.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     assertThat(fakeSessionMutation.streamStructures).hasSize(featureCnt);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.addedChildren).hasSize(featureCnt);
@@ -126,11 +127,11 @@
 
     List<StreamStructure> tokenStructures =
         new InternalProtocolBuilder().addToken(token.getContentId()).buildAsStreamStructure();
-    session.updateSession(false, tokenStructures, null);
+    session.updateSession(false, tokenStructures, SCHEMA_VERSION, null);
 
     session.bindModelProvider(modelProvider, null);
 
-    session.updateSession(false, streamStructures, context);
+    session.updateSession(false, streamStructures, SCHEMA_VERSION, context);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.mutationContext).isEqualTo(context);
   }
@@ -150,7 +151,7 @@
     StreamToken token =
         StreamToken.newBuilder().setContentId(contentIdGenerators.createTokenContentId(2)).build();
     MutationContext context = new MutationContext.Builder().setContinuationToken(token).build();
-    session.updateSession(false, streamStructures, context);
+    session.updateSession(false, streamStructures, SCHEMA_VERSION, context);
     assertThat(session.getContentInSession()).isEmpty();
   }
 
@@ -165,7 +166,7 @@
     List<StreamStructure> streamStructures = protocolBuilder.buildAsStreamStructure();
 
     sessionMutationResults = false;
-    session.updateSession(false, streamStructures, null);
+    session.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.isCommitted()).isTrue(); // Optimistic write will still call this
   }
@@ -181,7 +182,7 @@
     protocolBuilder.removeFeature(
         contentIdGenerators.createFeatureContentId(1), contentIdGenerators.createRootContentId(0));
     List<StreamStructure> streamStructures = protocolBuilder.buildAsStreamStructure();
-    session.updateSession(false, streamStructures, null);
+    session.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.removedChildren).hasSize(1);
     assertThat(session.getContentInSession()).hasSize(1);
diff --git a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/BUILD b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/BUILD
index 00472e9..b66b251 100644
--- a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/BUILD
+++ b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/BUILD
@@ -104,6 +104,7 @@
         "//src/test/proto/search/now/wire/feed:feed_test_java_proto_lite",
         "//third_party:robolectric",
         "@com_google_protobuf_javalite//:protobuf_java_lite",
+        "@maven//:com_google_guava_guava",
         "@maven//:com_google_truth_truth",
         "@maven//:org_mockito_mockito_core",
         "@robolectric//bazel:android-all",
@@ -119,6 +120,7 @@
     manifest_values = DEFAULT_ANDROID_LOCAL_TEST_MANIFEST,
     deps = [
         "//src/main/java/com/google/android/libraries/feed/api/host/config",
+        "//src/main/java/com/google/android/libraries/feed/api/internal/common",
         "//src/main/java/com/google/android/libraries/feed/api/internal/common/testing",
         "//src/main/java/com/google/android/libraries/feed/api/internal/modelprovider",
         "//src/main/java/com/google/android/libraries/feed/api/internal/store",
diff --git a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImplTest.java b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImplTest.java
index dbbef79..89b60c6 100644
--- a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImplTest.java
+++ b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/HeadSessionImplTest.java
@@ -37,6 +37,8 @@
 /** Tests of the {@link HeadSessionImpl} class. */
 @RunWith(RobolectricTestRunner.class)
 public class HeadSessionImplTest {
+  private static final int SCHEMA_VERSION = 1;
+
   private final ContentIdGenerators contentIdGenerators = new ContentIdGenerators();
   private final FakeClock fakeClock = new FakeClock();
   private final FakeThreadUtils fakeThreadUtils = FakeThreadUtils.withoutThreadChecks();
@@ -64,7 +66,7 @@
 
     // 1 clear, 3 features
     assertThat(streamStructures).hasSize(featureCnt + 1);
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
 
     // expect: 3 features
     assertThat(headSession.getContentInSession()).hasSize(featureCnt);
@@ -80,7 +82,7 @@
     addFeatures(protocolBuilder, featureCnt, 1);
     List<StreamStructure> streamStructures = protocolBuilder.buildAsStreamStructure();
     assertThat(streamStructures).hasSize(featureCnt + 1);
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     assertThat(headSession.getContentInSession()).hasSize(featureCnt);
 
     headSession.reset();
@@ -97,7 +99,7 @@
 
     // 1 clear, 3 features, token
     assertThat(streamStructures).hasSize(5);
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
 
     // expect: 3 features, 1 token
     assertThat(headSession.getContentInSession()).hasSize(featureCnt + 1);
@@ -119,10 +121,10 @@
     // The token needs to be in the session so update its content IDs with the token.
     List<StreamStructure> tokenStructures =
         new InternalProtocolBuilder().addToken(token.getContentId()).buildAsStreamStructure();
-    headSession.updateSession(false, tokenStructures, null);
+    headSession.updateSession(false, tokenStructures, SCHEMA_VERSION, null);
 
     MutationContext context = new MutationContext.Builder().setContinuationToken(token).build();
-    headSession.updateSession(false, streamStructures, context);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, context);
     // features 3, plus the token added above
     assertThat(headSession.getContentInSession()).hasSize(featureCnt + 1);
   }
@@ -139,7 +141,7 @@
 
     // The token needs to be in the session, if not we ignore the update
     MutationContext context = new MutationContext.Builder().setContinuationToken(token).build();
-    headSession.updateSession(false, streamStructures, context);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, context);
     assertThat(headSession.getContentInSession()).isEmpty();
   }
 
@@ -154,7 +156,7 @@
 
     // 1 clear, 3 features, 1 remove
     assertThat(streamStructures).hasSize(5);
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
 
     // expect: 2 features (3 added, then 1 removed)
     assertThat(headSession.getContentInSession()).hasSize(featureCnt - 1);
@@ -174,7 +176,7 @@
     assertThat(streamStructures).hasSize(4);
 
     // 1 clear, 3 features
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     assertThat(headSession.getContentInSession()).hasSize(featureCnt);
     assertThat(getContentInSession()).hasSize(featureCnt);
 
@@ -185,7 +187,7 @@
     assertThat(streamStructures).hasSize(1);
 
     // 0 features
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     assertThat(headSession.getContentInSession()).hasSize(featureCnt);
     assertThat(getContentInSession()).hasSize(featureCnt);
   }
@@ -199,7 +201,7 @@
     assertThat(streamStructures).hasSize(4);
 
     // 1 clear, 3 features
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     assertThat(headSession.getContentInSession()).hasSize(featureCnt);
     assertThat(getContentInSession()).hasSize(featureCnt);
 
@@ -211,13 +213,15 @@
     assertThat(streamStructures).hasSize(additionalFeatureCnt);
 
     // 0 features
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION, null);
     assertThat(headSession.getContentInSession()).hasSize(featureCnt + additionalFeatureCnt);
     assertThat(getContentInSession()).hasSize(featureCnt + additionalFeatureCnt);
   }
 
   @Test
-  public void testUpdateSession_clearHead() {
+  public void testUpdateSession_storeClearHead() {
+    headSession.initializeSession(ImmutableList.of(), SCHEMA_VERSION);
+
     int featureCnt = 3;
     InternalProtocolBuilder protocolBuilder = new InternalProtocolBuilder().addClearOperation();
     addFeatures(protocolBuilder, featureCnt, 1);
@@ -225,9 +229,10 @@
     assertThat(streamStructures).hasSize(4);
 
     // 1 clear, 3 features
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION + 1, null);
     assertThat(headSession.getContentInSession()).hasSize(featureCnt);
     assertThat(getContentInSession()).hasSize(featureCnt);
+    assertThat(headSession.getSchemaVersion()).isEqualTo(SCHEMA_VERSION);
 
     // Clear head and add 2 features, make sure we have new content ids
     int newFeatureCnt = 2;
@@ -240,9 +245,29 @@
 
     // 0 features
     fakeStore.clearHead();
-    headSession.updateSession(false, streamStructures, null);
+    headSession.updateSession(false, streamStructures, SCHEMA_VERSION + 1, null);
     assertThat(headSession.getContentInSession()).hasSize(newFeatureCnt);
     assertThat(getContentInSession()).hasSize(newFeatureCnt);
+    assertThat(headSession.getSchemaVersion()).isEqualTo(SCHEMA_VERSION);
+  }
+
+  @Test
+  public void testUpdateSession_clearHeadUpdatesSchemaVersion() {
+    headSession.initializeSession(ImmutableList.of(), SCHEMA_VERSION);
+    headSession.updateSession(
+        /* clearHead= */ true, ImmutableList.of(), SCHEMA_VERSION + 1, /* mutationContext= */ null);
+    assertThat(headSession.getSchemaVersion()).isEqualTo(SCHEMA_VERSION + 1);
+  }
+
+  @Test
+  public void testUpdateSession_schemaVersionUnchanged() {
+    headSession.initializeSession(ImmutableList.of(), SCHEMA_VERSION);
+    headSession.updateSession(
+        /* clearHead= */ false,
+        ImmutableList.of(),
+        SCHEMA_VERSION + 1,
+        /* mutationContext= */ null);
+    assertThat(headSession.getSchemaVersion()).isEqualTo(SCHEMA_VERSION);
   }
 
   @Test
@@ -254,6 +279,7 @@
     headSession.updateSession(
         /* clearHead= */ false,
         protocolBuilder.buildAsStreamStructure(),
+        SCHEMA_VERSION,
         /* mutationContext= */ null);
     assertThat(headSession.getContentInSession()).hasSize(1);
     assertThat(getContentInSession()).hasSize(1);
diff --git a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCacheTest.java b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCacheTest.java
index 4ef662c..bc90038 100644
--- a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCacheTest.java
+++ b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionCacheTest.java
@@ -20,6 +20,7 @@
 import static org.mockito.MockitoAnnotations.initMocks;
 
 import com.google.android.libraries.feed.api.host.config.Configuration;
+import com.google.android.libraries.feed.api.internal.common.PayloadWithId;
 import com.google.android.libraries.feed.api.internal.common.testing.ContentIdGenerators;
 import com.google.android.libraries.feed.api.internal.modelprovider.ModelProvider.ViewDepthProvider;
 import com.google.android.libraries.feed.api.internal.store.Store;
@@ -343,6 +344,39 @@
   }
 
   @Test
+  public void testUpdateHeadMetadata() {
+    long currentTime = 2L;
+    int schemaVersion = 3;
+    StreamSessions streamSessions =
+        StreamSessions.newBuilder()
+            .addStreamSession(
+                StreamSession.newBuilder()
+                    .setSessionId(Store.HEAD_SESSION_ID)
+                    .setLegacyTimeMillis(1))
+            .build();
+
+    mockStreamSessions(streamSessions);
+    sessionCache.initialize();
+    sessionCache.updateHeadMetadata(currentTime, schemaVersion);
+
+    List<Object> content = fakeStore.getContentById(SessionCache.STREAM_SESSION_CONTENT_ID);
+    assertThat(content).hasSize(1);
+    assertThat(((PayloadWithId) content.get(0)).payload)
+        .isEqualTo(
+            StreamPayload.newBuilder()
+                .setStreamSessions(
+                    StreamSessions.newBuilder()
+                        .addStreamSession(
+                            StreamSession.newBuilder()
+                                .setSessionId(Store.HEAD_SESSION_ID)
+                                .setSessionMetadata(
+                                    SessionMetadata.newBuilder()
+                                        .setLastAddedTimeMillis(currentTime)
+                                        .setSchemaVersion(schemaVersion))))
+                .build());
+  }
+
+  @Test
   public void testUpdatePersistedSessions() {
     sessionCache.initialize();
 
diff --git a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImplTest.java b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImplTest.java
index c720b11..a3ece68 100644
--- a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImplTest.java
+++ b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionImplTest.java
@@ -68,7 +68,7 @@
 
     // clear head will be ignored
     assertThat(streamStructures).hasSize(4);
-    session.updateSession(true, streamStructures, null);
+    session.updateSession(true, streamStructures, SCHEMA_VERSION, null);
     assertThat(fakeSessionMutation.streamStructures).isEmpty();
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.addedChildren).isEmpty();
@@ -88,6 +88,7 @@
     session.updateSession(
         true,
         streamStructures,
+        SCHEMA_VERSION,
         new MutationContext.Builder()
             .setContinuationToken(StreamToken.newBuilder().setContentId("token").build())
             .setRequestingSessionId(TEST_SESSION_ID)
@@ -117,6 +118,7 @@
     session.updateSession(
         /* clearHead= */ false,
         protocolBuilder.buildAsStreamStructure(),
+        SCHEMA_VERSION,
         /* mutationContext= */ null);
 
     assertThat(fakeSessionMutation.streamStructures).hasSize(1);
diff --git a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutationTest.java b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutationTest.java
index f3cc34b..8c4af39 100644
--- a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutationTest.java
+++ b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/SessionManagerMutationTest.java
@@ -54,6 +54,7 @@
 import com.google.android.libraries.feed.hostimpl.storage.InMemoryContentStorage;
 import com.google.android.libraries.feed.hostimpl.storage.InMemoryJournalStorage;
 import com.google.android.libraries.feed.testing.host.logging.FakeBasicLoggingApi;
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import com.google.search.now.feed.client.StreamDataProto.StreamDataOperation;
 import com.google.search.now.feed.client.StreamDataProto.StreamFeature;
@@ -169,6 +170,24 @@
   }
 
   @Test
+  public void testUpdateHeadMetadata() {
+    int schemaVersion = 3;
+    long currentTime = 8L;
+    fakeClock.set(currentTime);
+    MutationCommitter mutationCommitter = getMutationCommitter(MutationContext.EMPTY_CONTEXT);
+    mutationCommitter.accept(
+        Result.success(
+            Model.of(
+                ImmutableList.of(
+                    getStreamDataOperation(
+                        StreamStructure.newBuilder().setOperation(Operation.CLEAR_ALL).build(),
+                        /* payload= */ null)),
+                schemaVersion)));
+    assertThat(sessionCache.getHead().getSchemaVersion()).isEqualTo(schemaVersion);
+    assertThat(sessionCache.getHeadLastAddedTimeMillis()).isEqualTo(currentTime);
+  }
+
+  @Test
   public void testUpdateContent() {
     fakeClock.set(8L);
     List<String> contentIds = getContentIds(3);
diff --git a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImplTest.java b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImplTest.java
index 929d2a6..2060f41 100644
--- a/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImplTest.java
+++ b/src/test/java/com/google/android/libraries/feed/feedsessionmanager/internal/TimeoutSessionImplTest.java
@@ -77,7 +77,7 @@
 
     // 1 clear, 3 features
     assertThat(streamStructures).hasSize(4);
-    session.updateSession(true, streamStructures, null);
+    session.updateSession(true, streamStructures, SCHEMA_VERSION, null);
     assertThat(fakeSessionMutation.streamStructures).hasSize(featureCnt);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.addedChildren).hasSize(featureCnt);
@@ -102,6 +102,7 @@
     session.updateSession(
         true,
         streamStructures,
+        SCHEMA_VERSION,
         new MutationContext.Builder()
             .setRequestingSessionId(TEST_SESSION_ID)
             .setContinuationToken(StreamToken.newBuilder().setContentId("token").build())
@@ -172,7 +173,8 @@
         session, contentIdGenerators.createFeatureContentId(6), modelChildren);
 
     // Items below the lowest child should be removed.
-    session.updateSession(/* clearHead= */ true, new ArrayList<>(), /* mutationContext= */ null);
+    session.updateSession(
+        /* clearHead= */ true, new ArrayList<>(), SCHEMA_VERSION, /* mutationContext= */ null);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.removedChildren).hasSize(3);
     assertThat(fakeModelMutation.removedChildren.get(0).getContentId())
@@ -192,7 +194,8 @@
     setupForViewDepthProvider(session, tokenContentId, modelChildren);
 
     // Because the lowest child is a token it should also be removed.
-    session.updateSession(/* clearHead= */ true, new ArrayList<>(), /* mutationContext= */ null);
+    session.updateSession(
+        /* clearHead= */ true, new ArrayList<>(), SCHEMA_VERSION, /* mutationContext= */ null);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.removedChildren).hasSize(1);
     assertThat(fakeModelMutation.removedChildren.get(0).getContentId()).isEqualTo(tokenContentId);
@@ -209,7 +212,8 @@
     setupForViewDepthProvider(session, /* lowestChild= */ null, modelChildren);
 
     // The token should be removed even when there is a null lowest child.
-    session.updateSession(/* clearHead= */ true, new ArrayList<>(), /* mutationContext= */ null);
+    session.updateSession(
+        /* clearHead= */ true, new ArrayList<>(), SCHEMA_VERSION, /* mutationContext= */ null);
     FakeModelMutation fakeModelMutation = fakeModelProvider.getLatestModelMutation();
     assertThat(fakeModelMutation.removedChildren).hasSize(1);
     assertThat(fakeModelMutation.removedChildren.get(0).getContentId())