diff --git a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaServerSideAdInsertionMediaSource.java b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaServerSideAdInsertionMediaSource.java
index ed35d9a57a..4d54f64ebb 100644
--- a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaServerSideAdInsertionMediaSource.java
+++ b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaServerSideAdInsertionMediaSource.java
@@ -92,6 +92,7 @@ import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.checkerframework.checker.nullness.qual.EnsuresNonNull;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
 /**
@@ -474,6 +475,7 @@ public final class ImaServerSideAdInsertionMediaSource extends CompositeMediaSou
   }
 
   @MainThread
+  @EnsuresNonNull("contentTimeline")
   private void setContentTimeline(Timeline contentTimeline) {
     if (contentTimeline.equals(this.contentTimeline)) {
       return;
@@ -638,7 +640,9 @@ public final class ImaServerSideAdInsertionMediaSource extends CompositeMediaSou
           // Map adGroupIndex and adIndexInAdGroup to multi-period window.
           Pair<Integer, Integer> adGroupIndexAndAdIndexInAdGroup =
               getAdGroupAndIndexInMultiPeriodWindow(
-                  oldPosition.periodIndex, adPlaybackState, timeline);
+                  oldPosition.periodIndex - window.firstPeriodIndex,
+                  adPlaybackState,
+                  checkNotNull(contentTimeline));
           adGroupIndex = adGroupIndexAndAdIndexInAdGroup.first;
           adIndexInAdGroup = adGroupIndexAndAdIndexInAdGroup.second;
         }
diff --git a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaUtil.java b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaUtil.java
index ad6021c765..a353878831 100644
--- a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaUtil.java
+++ b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaUtil.java
@@ -475,11 +475,11 @@ import java.util.Set;
    *
    * @param adPeriodIndex The period index of the ad period.
    * @param adPlaybackState The ad playback state that holds the ad group and ad information.
-   * @param timeline The timeline that contains the ad period.
+   * @param contentTimeline The timeline that contains the ad period.
    * @return A pair with the ad group index (first) and the ad index in that ad group (second).
    */
   public static Pair<Integer, Integer> getAdGroupAndIndexInMultiPeriodWindow(
-      int adPeriodIndex, AdPlaybackState adPlaybackState, Timeline timeline) {
+      int adPeriodIndex, AdPlaybackState adPlaybackState, Timeline contentTimeline) {
     Timeline.Period period = new Timeline.Period();
     int periodIndex = 0;
     long totalElapsedContentDurationUs = 0;
@@ -488,8 +488,8 @@ import java.util.Set;
       AdPlaybackState.AdGroup adGroup = adPlaybackState.getAdGroup(/* adGroupIndex= */ i);
       long adGroupDurationUs = sum(adGroup.durationsUs);
       long elapsedAdGroupAdDurationUs = 0;
-      for (int j = periodIndex; j < timeline.getPeriodCount(); j++) {
-        timeline.getPeriod(j, period, /* setIds= */ true);
+      for (int j = periodIndex; j < contentTimeline.getPeriodCount(); j++) {
+        contentTimeline.getPeriod(j, period, /* setIds= */ true);
         // TODO(b/192231683) Remove subtracted US from ad group time when we can upgrade the SDK.
         // Subtract one microsecond to work around rounding errors with adGroup.timeUs.
         if (totalElapsedContentDurationUs < adGroup.timeUs - 1) {