@powersync/service-core 0.0.0-dev-20250507151436 → 0.0.0-dev-20250611110033

package/src/sync/sync.ts

@@ -7,12 +7,12 @@ import * as auth from '../auth/auth-index.js';
 import * as storage from '../storage/storage-index.js';
 import * as util from '../util/util-index.js';
 
-import { logger } from '@powersync/lib-services-framework';
-import { BucketChecksumState } from './BucketChecksumState.js';
+import { Logger, logger as defaultLogger } from '@powersync/lib-services-framework';
+import { BucketChecksumState, CheckpointLine } from './BucketChecksumState.js';
 import { mergeAsyncIterables } from '../streams/streams-index.js';
 import { acquireSemaphoreAbortable, settledPromise, tokenStream, TokenStreamOptions } from './util.js';
 import { SyncContext } from './SyncContext.js';
-import { RequestTracker } from './RequestTracker.js';
+import { OperationsSentStats, RequestTracker, statsForBatch } from './RequestTracker.js';
 
 export interface SyncStreamParameters {
   syncContext: SyncContext;
@@ -21,6 +21,7 @@ export interface SyncStreamParameters {
   params: util.StreamingSyncRequest;
   syncParams: RequestParameters;
   token: auth.JwtPayload;
+  logger?: Logger;
   /**
    * If this signal is aborted, the stream response ends as soon as possible, without error.
    */
@@ -35,6 +36,8 @@ export async function* streamResponse(
 ): AsyncIterable<util.StreamingSyncLine | string | null> {
   const { syncContext, bucketStorage, syncRules, params, syncParams, token, tokenStreamOptions, tracker, signal } =
     options;
+  const logger = options.logger ?? defaultLogger;
+
   // We also need to be able to abort, so we create our own controller.
   const controller = new AbortController();
   if (signal) {
@@ -57,7 +60,8 @@ export async function* streamResponse(
     params,
     syncParams,
     tracker,
-    controller.signal
+    controller.signal,
+    logger
   );
   // Merge the two streams, and abort as soon as one of the streams end.
   const merged = mergeAsyncIterables([stream, ki], controller.signal);
@@ -77,11 +81,6 @@ export async function* streamResponse(
   }
 }
 
-export type BucketSyncState = {
-  description?: BucketDescription; // Undefined if the bucket has not yet been resolved by us.
-  start_op_id: util.InternalOpId;
-};
-
 async function* streamResponseInner(
   syncContext: SyncContext,
   bucketStorage: storage.SyncRulesBucketStorage,
@@ -89,11 +88,12 @@ async function* streamResponseInner(
   params: util.StreamingSyncRequest,
   syncParams: RequestParameters,
   tracker: RequestTracker,
-  signal: AbortSignal
+  signal: AbortSignal,
+  logger: Logger
 ): AsyncGenerator<util.StreamingSyncLine | string | null> {
   const { raw_data, binary_data } = params;
 
-  const checkpointUserId = util.checkpointUserId(syncParams.token_parameters.user_id as string, params.client_id);
+  const checkpointUserId = util.checkpointUserId(syncParams.tokenParameters.user_id as string, params.client_id);
 
   const checksumState = new BucketChecksumState({
     syncContext,
@@ -103,7 +103,8 @@ async function* streamResponseInner(
     initialBucketPositions: params.buckets?.map((bucket) => ({
       name: bucket.name,
       after: BigInt(bucket.after)
-    }))
+    })),
+    logger: logger
   });
   const stream = bucketStorage.watchCheckpointChanges({
     user_id: checkpointUserId,
@@ -111,16 +112,29 @@ async function* streamResponseInner(
   });
   const newCheckpoints = stream[Symbol.asyncIterator]();
 
+  type CheckpointAndLine = {
+    checkpoint: bigint;
+    line: CheckpointLine | null;
+  };
+
+  async function waitForNewCheckpointLine(): Promise<IteratorResult<CheckpointAndLine>> {
+    const next = await newCheckpoints.next();
+    if (next.done) {
+      return { done: true, value: undefined };
+    }
+
+    const line = await checksumState.buildNextCheckpointLine(next.value);
+    return { done: false, value: { checkpoint: next.value.base.checkpoint, line } };
+  }
+
   try {
-    let nextCheckpointPromise:
-      | Promise<PromiseSettledResult<IteratorResult<storage.StorageCheckpointUpdate>>>
-      | undefined;
+    let nextCheckpointPromise: Promise<PromiseSettledResult<IteratorResult<CheckpointAndLine>>> | undefined;
 
     do {
       if (!nextCheckpointPromise) {
         // Wrap in a settledPromise, so that abort errors after the parent stopped iterating
         // does not result in uncaught errors.
-        nextCheckpointPromise = settledPromise(newCheckpoints.next());
+        nextCheckpointPromise = settledPromise(waitForNewCheckpointLine());
       }
       const next = await nextCheckpointPromise;
       nextCheckpointPromise = undefined;
@@ -130,7 +144,7 @@ async function* streamResponseInner(
       if (next.value.done) {
         break;
       }
-      const line = await checksumState.buildNextCheckpointLine(next.value.value);
+      const line = next.value.value.line;
       if (line == null) {
         // No update to sync
         continue;
@@ -138,7 +152,10 @@ async function* streamResponseInner(
 
       const { checkpointLine, bucketsToFetch } = line;
 
+      // Since yielding can block, we update the state just before yielding the line.
+      line.advance();
       yield checkpointLine;
+
       // Start syncing data for buckets up to the checkpoint. As soon as we have completed at least one priority and
       // at least 1000 operations, we also start listening for new checkpoints concurrently. When a new checkpoint comes
       // in while we're still busy syncing data for lower priorities, interrupt the current operation and start syncing
@@ -164,30 +181,32 @@ async function* streamResponseInner(
       function maybeRaceForNewCheckpoint() {
         if (syncedOperations >= 1000 && nextCheckpointPromise === undefined) {
           nextCheckpointPromise = (async () => {
-            const next = await settledPromise(newCheckpoints.next());
-            if (next.status == 'rejected') {
-              abortCheckpointController.abort();
-            } else if (!next.value.done) {
-              // Stop the running bucketDataInBatches() iterations, making the main flow reach the new checkpoint.
-              abortCheckpointController.abort();
+            while (true) {
+              const next = await settledPromise(waitForNewCheckpointLine());
+              if (next.status == 'rejected') {
+                abortCheckpointController.abort();
+              } else if (!next.value.done) {
+                if (next.value.value.line == null) {
+                  // There's a new checkpoint that doesn't affect this sync stream. Keep listening, but don't
+                  // interrupt this batch.
+                  continue;
+                }
+
+                // A new sync line can be emitted. Stop running the bucketDataInBatches() iterations, making the
+                // main flow reach the new checkpoint.
+                abortCheckpointController.abort();
+              }
+
+              return next;
             }
-
-            return next;
           })();
         }
       }
 
-      function markOperationsSent(operations: number) {
-        syncedOperations += operations;
-        tracker.addOperationsSynced(operations);
-        // Disable interrupting checkpoints for now.
-        // There is a bug with interrupting checkpoints where:
-        // 1. User is in the middle of syncing a large batch of data (for example initial sync).
-        // 2. A new checkpoint comes in, which interrupts the batch.
-        // 3. However, the new checkpoint does not contain any new data for this connection, so nothing further is synced.
-        // This then causes the client to wait indefinitely for the remaining data or checkpoint_complete message. That is
-        // only resolved when a new checkpoint comes in that does have data for this connection, or the connection is restarted.
-        // maybeRaceForNewCheckpoint();
+      function markOperationsSent(stats: OperationsSentStats) {
+        syncedOperations += stats.total;
+        tracker.addOperationsSynced(stats);
+        maybeRaceForNewCheckpoint();
       }
 
       // This incrementally updates dataBuckets with each individual bucket position.
@@ -201,18 +220,19 @@ async function* streamResponseInner(
         yield* bucketDataInBatches({
           syncContext: syncContext,
           bucketStorage: bucketStorage,
-          checkpoint: next.value.value.base.checkpoint,
+          checkpoint: next.value.value.checkpoint,
           bucketsToFetch: buckets,
-          checksumState,
+          checkpointLine: line,
           raw_data,
           binary_data,
           onRowsSent: markOperationsSent,
           abort_connection: signal,
           abort_batch: abortCheckpointSignal,
-          user_id: syncParams.user_id,
+          user_id: syncParams.userId,
           // Passing null here will emit a full sync complete message at the end. If we pass a priority, we'll emit a partial
           // sync complete message instead.
-          forPriority: !isLast ? priority : null
+          forPriority: !isLast ? priority : null,
+          logger
         });
       }
 
@@ -229,9 +249,10 @@ interface BucketDataRequest {
   syncContext: SyncContext;
   bucketStorage: storage.SyncRulesBucketStorage;
   checkpoint: util.InternalOpId;
+  /** Contains current bucket state. Modified by the request as data is sent. */
+  checkpointLine: CheckpointLine;
+  /** Subset of checkpointLine.bucketsToFetch, filtered by priority. */
   bucketsToFetch: BucketDescription[];
-  /** Contains current bucket state. Modified by the request as data is sent.  */
-  checksumState: BucketChecksumState;
   raw_data: boolean | undefined;
   binary_data: boolean | undefined;
   /** Signals that the connection was aborted and that streaming should stop ASAP. */
@@ -243,7 +264,8 @@ interface BucketDataRequest {
   abort_batch: AbortSignal;
   user_id?: string;
   forPriority: BucketPriority | null;
-  onRowsSent: (amount: number) => void;
+  onRowsSent: (stats: OperationsSentStats) => void;
+  logger: Logger;
 }
 
 async function* bucketDataInBatches(request: BucketDataRequest) {
@@ -292,12 +314,13 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
     bucketStorage: storage,
     checkpoint,
     bucketsToFetch,
-    checksumState,
+    checkpointLine,
     raw_data,
     binary_data,
     abort_connection,
     abort_batch,
-    onRowsSent
+    onRowsSent,
+    logger
   } = request;
 
   let checkpointInvalidated = false;
@@ -322,7 +345,7 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
     }
     // Optimization: Only fetch buckets for which the checksums have changed since the last checkpoint
     // For the first batch, this will be all buckets.
-    const filteredBuckets = checksumState.getFilteredBucketPositions(bucketsToFetch);
+    const filteredBuckets = checkpointLine.getFilteredBucketPositions(bucketsToFetch);
     const dataBatches = storage.getBucketDataBatch(checkpoint, filteredBuckets);
 
     let has_more = false;
@@ -369,9 +392,9 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
         // iterator memory in case if large data sent.
         yield { data: null, done: false };
       }
-      onRowsSent(r.data.length);
+      onRowsSent(statsForBatch(r));
 
-      checksumState.updateBucketPosition({ bucket: r.bucket, nextAfter: BigInt(r.next_after), hasMore: r.has_more });
+      checkpointLine.updateBucketPosition({ bucket: r.bucket, nextAfter: BigInt(r.next_after), hasMore: r.has_more });
 
       // Check if syncing bucket data is supposed to stop before fetching more data
       // from storage.
@@ -417,7 +440,7 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
   }
 }
 
-function transformLegacyResponse(bucketData: util.SyncBucketData): any {
+function transformLegacyResponse(bucketData: util.SyncBucketData): util.SyncBucketData<util.ProtocolOplogData> {
   return {
     ...bucketData,
     data: bucketData.data.map((entry) => {

package/src/util/config/collectors/config-collector.ts

@@ -1,8 +1,8 @@
-import { configFile } from '@powersync/service-types';
-import * as t from 'ts-codec';
 import * as yaml from 'yaml';
 
 import { schema } from '@powersync/lib-services-framework';
+import { configFile } from '@powersync/service-types';
+
 import { RunnerConfig } from '../types.js';
 import { YamlEnvTag } from './impl/yaml-env.js';
 
@@ -12,11 +12,7 @@ export enum ConfigFileFormat {
 }
 
 // ts-codec itself doesn't give great validation errors, so we use json schema for that
-const configSchemaValidator = schema
-  .parseJSONSchema(
-    t.generateJSONSchema(configFile.powerSyncConfig, { allowAdditional: true, parsers: [configFile.portParser] })
-  )
-  .validator();
+const configSchemaValidator = schema.parseJSONSchema(configFile.PowerSyncConfigJSONSchema).validator();
 
 export abstract class ConfigCollector {
   abstract get name(): string;

package/src/util/protocol-types.ts

@@ -1,5 +1,6 @@
 import * as t from 'ts-codec';
-import { BucketDescription, BucketPriority, SqliteJsonValue } from '@powersync/service-sync-rules';
+import { BucketDescription, BucketPriority, SqliteJsonRow } from '@powersync/service-sync-rules';
+import { JsonContainer } from '@powersync/service-jsonbig';
 
 export const BucketRequest = t.object({
   name: t.string,
@@ -65,7 +66,7 @@ export interface StreamingSyncCheckpointDiff {
 }
 
 export interface StreamingSyncData {
-  data: SyncBucketData;
+  data: SyncBucketData<ProtocolOplogData>;
 }
 
 export interface StreamingSyncCheckpointComplete {
@@ -109,13 +110,9 @@ export interface BucketState {
   op_id: string;
 }
 
-export interface SyncDataBatch {
-  buckets: SyncBucketData[];
-}
-
-export interface SyncBucketData {
+export interface SyncBucketData<Data extends ProtocolOplogData = StoredOplogData> {
   bucket: string;
-  data: OplogEntry[];
+  data: OplogEntry<Data>[];
   /**
    * True if there _could_ be more data for this bucket, and another request must be made.
    */
@@ -130,12 +127,20 @@ export interface SyncBucketData {
   next_after: ProtocolOpId;
 }
 
-export interface OplogEntry {
+export type StoredOplogData = string | null;
+
+// Note: When clients have both raw_data and binary_data disabled (this only affects legacy
+// clients), data is actually a `Record<string, SqliteJsonValue>`. Oplog entries are always
+// stored as a serialized (JSON) string so that they don't have to be parsed in the sync service,
+// this representation only exists on the way out for legacy clients.
+export type ProtocolOplogData = SqliteJsonRow | JsonContainer | StoredOplogData;
+
+export interface OplogEntry<Data extends ProtocolOplogData = StoredOplogData> {
   op_id: ProtocolOpId;
   op: 'PUT' | 'REMOVE' | 'MOVE' | 'CLEAR';
   object_type?: string;
   object_id?: string;
-  data?: Record<string, SqliteJsonValue> | string | null;
+  data?: Data;
   checksum: number | bigint;
   subkey?: string;
 }

package/test/src/auth.test.ts

@@ -75,21 +75,21 @@ describe('JWT Auth', () => {
         defaultAudiences: ['other'],
         maxAge: '6m'
       })
-    ).rejects.toThrow('unexpected "aud" claim value');
+    ).rejects.toThrow('[PSYNC_S2105] Unexpected "aud" claim value: "tests"');
 
     await expect(
       store.verifyJwt(signedJwt, {
         defaultAudiences: [],
         maxAge: '6m'
       })
-    ).rejects.toThrow('unexpected "aud" claim value');
+    ).rejects.toThrow('[PSYNC_S2105] Unexpected "aud" claim value: "tests"');
 
     await expect(
       store.verifyJwt(signedJwt, {
         defaultAudiences: ['tests'],
         maxAge: '1m'
       })
-    ).rejects.toThrow('Token must expire in a maximum of');
+    ).rejects.toThrow('[PSYNC_S2104] Token must expire in a maximum of 60 seconds, got 300s');
 
     const signedJwt2 = await new jose.SignJWT({})
       .setProtectedHeader({ alg: 'HS256', kid: 'k1' })
@@ -104,7 +104,25 @@ describe('JWT Auth', () => {
         defaultAudiences: ['tests'],
         maxAge: '5m'
       })
-    ).rejects.toThrow('missing required "sub" claim');
+    ).rejects.toThrow('[PSYNC_S2101] JWT payload is missing a required claim "sub"');
+
+    // expired token
+    const d = Math.round(Date.now() / 1000);
+    const signedJwt3 = await new jose.SignJWT({})
+      .setProtectedHeader({ alg: 'HS256', kid: 'k1' })
+      .setSubject('f1')
+      .setIssuedAt(d - 500)
+      .setIssuer('tester')
+      .setAudience('tests')
+      .setExpirationTime(d - 400)
+      .sign(signKey);
+
+    await expect(
+      store.verifyJwt(signedJwt3, {
+        defaultAudiences: ['tests'],
+        maxAge: '5m'
+      })
+    ).rejects.toThrow('[PSYNC_S2103] JWT has expired');
   });
 
   test('Algorithm validation', async () => {
@@ -159,7 +177,7 @@ describe('JWT Auth', () => {
         maxAge: '6m'
       })
     ).rejects.toThrow(
-      'Could not find an appropriate key in the keystore. The key is missing or no key matched the token KID'
+      '[PSYNC_S2101] Could not find an appropriate key in the keystore. The key is missing or no key matched the token KID'
     );
 
     // Wrong kid
@@ -178,7 +196,7 @@ describe('JWT Auth', () => {
         maxAge: '6m'
       })
     ).rejects.toThrow(
-      'Could not find an appropriate key in the keystore. The key is missing or no key matched the token KID'
+      '[PSYNC_S2101] Could not find an appropriate key in the keystore. The key is missing or no key matched the token KID'
     );
 
     // No kid, matches sharedKey2
@@ -255,7 +273,7 @@ describe('JWT Auth', () => {
         defaultAudiences: ['tests'],
         maxAge: '6m'
       })
-    ).rejects.toThrow('unexpected "aud" claim value');
+    ).rejects.toThrow('[PSYNC_S2105] Unexpected "aud" claim value: "tests"');
 
     const signedJwt3 = await new jose.SignJWT({})
       .setProtectedHeader({ alg: 'HS256', kid: 'k1' })
@@ -290,7 +308,7 @@ describe('JWT Auth', () => {
         reject_ip_ranges: ['local']
       }
     });
-    await expect(invalid.getKeys()).rejects.toThrow('IPs in this range are not supported');
+    await expect(invalid.getKeys()).rejects.toThrow('[PSYNC_S2204] JWKS request failed');
 
     // IPs throw an error immediately
     expect(
@@ -345,7 +363,7 @@ describe('JWT Auth', () => {
     expect(key.kid).toEqual(publicKeyRSA.kid!);
 
     cached.addTimeForTests(301_000);
-    currentResponse = Promise.reject('refresh failed');
+    currentResponse = Promise.reject(new Error('refresh failed'));
 
     // Uses the promise, refreshes in the background
     let response = await cached.getKeys();
@@ -357,14 +375,14 @@ describe('JWT Auth', () => {
     response = await cached.getKeys();
     // Still have the cached key, but also have the error
     expect(response.keys[0].kid).toEqual(publicKeyRSA.kid!);
-    expect(response.errors[0].message).toMatch('Failed to fetch');
+    expect(response.errors[0].message).toMatch('[PSYNC_S2201] refresh failed');
 
     await cached.addTimeForTests(3601_000);
     response = await cached.getKeys();
 
     // Now the keys have expired, and the request still fails
     expect(response.keys).toEqual([]);
-    expect(response.errors[0].message).toMatch('Failed to fetch');
+    expect(response.errors[0].message).toMatch('[PSYNC_S2201] refresh failed');
 
     currentResponse = Promise.resolve({
       errors: [],

package/test/src/sync/BucketChecksumState.test.ts

@@ -71,6 +71,7 @@ bucket_definitions:
       writeCheckpoint: null,
       update: CHECKPOINT_INVALIDATE_ALL
     }))!;
+    line.advance();
     expect(line.checkpointLine).toEqual({
       checkpoint: {
         buckets: [{ bucket: 'global[]', checksum: 1, count: 1, priority: 3 }],
@@ -85,10 +86,11 @@ bucket_definitions:
       }
     ]);
     // This is the bucket data to be fetched
-    expect(state.getFilteredBucketPositions(line.bucketsToFetch)).toEqual(new Map([['global[]', 0n]]));
+    expect(line.getFilteredBucketPositions()).toEqual(new Map([['global[]', 0n]]));
 
     // This similuates the bucket data being sent
-    state.updateBucketPosition({ bucket: 'global[]', nextAfter: 1n, hasMore: false });
+    line.advance();
+    line.updateBucketPosition({ bucket: 'global[]', nextAfter: 1n, hasMore: false });
 
     // Update bucket storage state
     storage.updateTestChecksum({ bucket: 'global[]', checksum: 2, count: 2 });
@@ -104,6 +106,7 @@ bucket_definitions:
         invalidateParameterBuckets: false
       }
     }))!;
+    line2.advance();
     expect(line2.checkpointLine).toEqual({
       checkpoint_diff: {
         removed_buckets: [],
@@ -112,7 +115,7 @@ bucket_definitions:
         write_checkpoint: undefined
       }
     });
-    expect(state.getFilteredBucketPositions(line2.bucketsToFetch)).toEqual(new Map([['global[]', 1n]]));
+    expect(line2.getFilteredBucketPositions()).toEqual(new Map([['global[]', 1n]]));
   });
 
   test('global bucket with initial state', async () => {
@@ -137,6 +140,7 @@ bucket_definitions:
       writeCheckpoint: null,
       update: CHECKPOINT_INVALIDATE_ALL
     }))!;
+    line.advance();
     expect(line.checkpointLine).toEqual({
       checkpoint: {
         buckets: [{ bucket: 'global[]', checksum: 1, count: 1, priority: 3 }],
@@ -151,7 +155,7 @@ bucket_definitions:
       }
     ]);
     // This is the main difference between this and the previous test
-    expect(state.getFilteredBucketPositions(line.bucketsToFetch)).toEqual(new Map([['global[]', 1n]]));
+    expect(line.getFilteredBucketPositions()).toEqual(new Map([['global[]', 1n]]));
   });
 
   test('multiple static buckets', async () => {
@@ -192,6 +196,7 @@ bucket_definitions:
         priority: 3
       }
     ]);
+    line.advance();
 
     storage.updateTestChecksum({ bucket: 'global[1]', checksum: 2, count: 2 });
     storage.updateTestChecksum({ bucket: 'global[2]', checksum: 2, count: 2 });
@@ -240,6 +245,7 @@ bucket_definitions:
       writeCheckpoint: null,
       update: CHECKPOINT_INVALIDATE_ALL
     }))!;
+    line.advance();
     expect(line.checkpointLine).toEqual({
       checkpoint: {
         buckets: [{ bucket: 'global[]', checksum: 1, count: 1, priority: 3 }],
@@ -253,7 +259,7 @@ bucket_definitions:
         priority: 3
       }
     ]);
-    expect(state.getFilteredBucketPositions(line.bucketsToFetch)).toEqual(new Map([['global[]', 0n]]));
+    expect(line.getFilteredBucketPositions()).toEqual(new Map([['global[]', 0n]]));
   });
 
   test('invalidating individual bucket', async () => {
@@ -274,14 +280,14 @@ bucket_definitions:
     // We specifically do not set this here, so that we have manual control over the events.
     // storage.filter = state.checkpointFilter;
 
-    await state.buildNextCheckpointLine({
+    const line = await state.buildNextCheckpointLine({
       base: { checkpoint: 1n, lsn: '1' },
       writeCheckpoint: null,
       update: CHECKPOINT_INVALIDATE_ALL
     });
-
-    state.updateBucketPosition({ bucket: 'global[1]', nextAfter: 1n, hasMore: false });
-    state.updateBucketPosition({ bucket: 'global[2]', nextAfter: 1n, hasMore: false });
+    line!.advance();
+    line!.updateBucketPosition({ bucket: 'global[1]', nextAfter: 1n, hasMore: false });
+    line!.updateBucketPosition({ bucket: 'global[2]', nextAfter: 1n, hasMore: false });
 
     storage.updateTestChecksum({ bucket: 'global[1]', checksum: 2, count: 2 });
     storage.updateTestChecksum({ bucket: 'global[2]', checksum: 2, count: 2 });
@@ -330,12 +336,14 @@ bucket_definitions:
     storage.updateTestChecksum({ bucket: 'global[1]', checksum: 1, count: 1 });
     storage.updateTestChecksum({ bucket: 'global[2]', checksum: 1, count: 1 });
 
-    await state.buildNextCheckpointLine({
+    const line = await state.buildNextCheckpointLine({
       base: { checkpoint: 1n, lsn: '1' },
       writeCheckpoint: null,
       update: CHECKPOINT_INVALIDATE_ALL
     });
 
+    line!.advance();
+
     storage.updateTestChecksum({ bucket: 'global[1]', checksum: 2, count: 2 });
     storage.updateTestChecksum({ bucket: 'global[2]', checksum: 2, count: 2 });
 
@@ -380,6 +388,7 @@ bucket_definitions:
       writeCheckpoint: null,
       update: CHECKPOINT_INVALIDATE_ALL
     }))!;
+    line.advance();
     expect(line.checkpointLine).toEqual({
       checkpoint: {
         buckets: [
@@ -402,7 +411,7 @@ bucket_definitions:
     ]);
 
     // This is the bucket data to be fetched
-    expect(state.getFilteredBucketPositions(line.bucketsToFetch)).toEqual(
+    expect(line.getFilteredBucketPositions()).toEqual(
       new Map([
         ['global[1]', 0n],
         ['global[2]', 0n]
@@ -411,8 +420,9 @@ bucket_definitions:
 
     // No data changes here.
     // We simulate partial data sent, before a checkpoint is interrupted.
-    state.updateBucketPosition({ bucket: 'global[1]', nextAfter: 3n, hasMore: false });
-    state.updateBucketPosition({ bucket: 'global[2]', nextAfter: 1n, hasMore: true });
+    line.advance();
+    line.updateBucketPosition({ bucket: 'global[1]', nextAfter: 3n, hasMore: false });
+    line.updateBucketPosition({ bucket: 'global[2]', nextAfter: 1n, hasMore: true });
     storage.updateTestChecksum({ bucket: 'global[1]', checksum: 4, count: 4 });
 
     const line2 = (await state.buildNextCheckpointLine({
@@ -424,6 +434,7 @@ bucket_definitions:
         updatedDataBuckets: new Set(['global[1]'])
       }
     }))!;
+    line2.advance();
     expect(line2.checkpointLine).toEqual({
       checkpoint_diff: {
         removed_buckets: [],
@@ -451,7 +462,7 @@ bucket_definitions:
       }
     ]);
 
-    expect(state.getFilteredBucketPositions(line2.bucketsToFetch)).toEqual(
+    expect(line2.getFilteredBucketPositions()).toEqual(
       new Map([
         ['global[1]', 3n],
         ['global[2]', 1n]
@@ -507,16 +518,18 @@ bucket_definitions:
         priority: 3
       }
     ]);
+    line.advance();
     // This is the bucket data to be fetched
-    expect(state.getFilteredBucketPositions(line.bucketsToFetch)).toEqual(
+    expect(line.getFilteredBucketPositions()).toEqual(
       new Map([
         ['by_project[1]', 0n],
         ['by_project[2]', 0n]
       ])
     );
 
-    state.updateBucketPosition({ bucket: 'by_project[1]', nextAfter: 1n, hasMore: false });
-    state.updateBucketPosition({ bucket: 'by_project[2]', nextAfter: 1n, hasMore: false });
+    line.advance();
+    line.updateBucketPosition({ bucket: 'by_project[1]', nextAfter: 1n, hasMore: false });
+    line.updateBucketPosition({ bucket: 'by_project[2]', nextAfter: 1n, hasMore: false });
 
     storage.getParameterSets = async (
       checkpoint: InternalOpId,
@@ -538,6 +551,7 @@ bucket_definitions:
         invalidateParameterBuckets: false
       }
     }))!;
+    line2.advance();
     expect(line2.checkpointLine).toEqual({
       checkpoint_diff: {
         removed_buckets: [],
@@ -546,7 +560,7 @@ bucket_definitions:
         write_checkpoint: undefined
       }
     });
-    expect(state.getFilteredBucketPositions(line2.bucketsToFetch)).toEqual(new Map([['by_project[3]', 0n]]));
+    expect(line2.getFilteredBucketPositions()).toEqual(new Map([['by_project[3]', 0n]]));
   });
 });