@powersync/service-core 0.0.0-dev-20250317122913 → 0.0.0-dev-20250325131118

package/src/replication/replication-metrics.ts

@@ -0,0 +1,45 @@
+import { MetricsEngine } from '../metrics/metrics-index.js';
+import { ReplicationMetric } from '@powersync/service-types';
+
+/**
+ *  Create and register the core replication metrics.
+ *  @param engine
+ */
+export function createCoreReplicationMetrics(engine: MetricsEngine): void {
+  engine.createCounter({
+    name: ReplicationMetric.DATA_REPLICATED_BYTES,
+    description: 'Uncompressed size of replicated data',
+    unit: 'bytes'
+  });
+
+  engine.createCounter({
+    name: ReplicationMetric.ROWS_REPLICATED,
+    description: 'Total number of replicated rows'
+  });
+
+  engine.createCounter({
+    name: ReplicationMetric.TRANSACTIONS_REPLICATED,
+    description: 'Total number of replicated transactions'
+  });
+
+  engine.createCounter({
+    name: ReplicationMetric.CHUNKS_REPLICATED,
+    description: 'Total number of replication chunks'
+  });
+}
+
+/**
+ *  Initialise the core replication metrics. This should be called after the metrics have been created.
+ *  @param engine
+ */
+export function initializeCoreReplicationMetrics(engine: MetricsEngine): void {
+  const data_replicated_bytes = engine.getCounter(ReplicationMetric.DATA_REPLICATED_BYTES);
+  const rows_replicated_total = engine.getCounter(ReplicationMetric.ROWS_REPLICATED);
+  const transactions_replicated_total = engine.getCounter(ReplicationMetric.TRANSACTIONS_REPLICATED);
+  const chunks_replicated_total = engine.getCounter(ReplicationMetric.CHUNKS_REPLICATED);
+
+  data_replicated_bytes.add(0);
+  rows_replicated_total.add(0);
+  transactions_replicated_total.add(0);
+  chunks_replicated_total.add(0);
+}

package/src/routes/endpoints/probes.ts

@@ -25,11 +25,27 @@ export const startupCheck = routeDefinition({
 export const livenessCheck = routeDefinition({
   path: ProbeRoutes.LIVENESS,
   method: router.HTTPMethod.GET,
-  handler: async () => {
+  handler: async (params) => {
     const state = container.probes.state();
 
+    /**
+     * The HTTP probes currently only function in the API and UNIFIED
+     * modes.
+     *
+     * For the API mode, we don't really touch the state, but any response from
+     * the request indicates the service is alive.
+     *
+     * For the UNIFIED mode we update the touched_at time while the Replicator engine is running.
+     * If the replication engine is present and the timeDifference from the last
+     * touched_at is large, we report that the service is not live.
+     *
+     * This is only an incremental improvement. In future these values should be configurable.
+     */
+
+    const isAPIOnly = !params.context.service_context.replicationEngine;
     const timeDifference = Date.now() - state.touched_at.getTime();
-    const status = timeDifference < 10000 ? 200 : 400;
+
+    const status = isAPIOnly ? 200 : timeDifference < 10000 ? 200 : 400;
 
     return new router.RouterResponse({
       status,

package/src/routes/endpoints/socket-route.ts

@@ -2,18 +2,19 @@ import { ErrorCode, errors, logger, schema } from '@powersync/lib-services-frame
 import { RequestParameters } from '@powersync/service-sync-rules';
 import { serialize } from 'bson';
 
-import { Metrics } from '../../metrics/Metrics.js';
 import * as sync from '../../sync/sync-index.js';
 import * as util from '../../util/util-index.js';
 import { SocketRouteGenerator } from '../router-socket.js';
 import { SyncRoutes } from './sync-stream.js';
 
+import { APIMetric } from '@powersync/service-types';
+
 export const syncStreamReactive: SocketRouteGenerator = (router) =>
   router.reactiveStream<util.StreamingSyncRequest, any>(SyncRoutes.STREAM, {
     validator: schema.createTsCodecValidator(util.StreamingSyncRequest, { allowAdditional: true }),
     handler: async ({ context, params, responder, observer, initialN, signal: upstreamSignal }) => {
       const { service_context } = context;
-      const { routerEngine, syncContext } = service_context;
+      const { routerEngine, metricsEngine, syncContext } = service_context;
 
       // Create our own controller that we can abort directly
       const controller = new AbortController();
@@ -69,8 +70,8 @@ export const syncStreamReactive: SocketRouteGenerator = (router) =>
         controller.abort();
       });
 
-      Metrics.getInstance().concurrent_connections.add(1);
-      const tracker = new sync.RequestTracker();
+      metricsEngine.getUpDownCounter(APIMetric.CONCURRENT_CONNECTIONS).add(1);
+      const tracker = new sync.RequestTracker(metricsEngine);
       try {
         for await (const data of sync.streamResponse({
           syncContext: syncContext,
@@ -147,7 +148,7 @@ export const syncStreamReactive: SocketRouteGenerator = (router) =>
           operations_synced: tracker.operationsSynced,
           data_synced_bytes: tracker.dataSyncedBytes
         });
-        Metrics.getInstance().concurrent_connections.add(-1);
+        metricsEngine.getUpDownCounter(APIMetric.CONCURRENT_CONNECTIONS).add(-1);
       }
     }
   });

package/src/routes/endpoints/sync-stream.ts

@@ -5,10 +5,11 @@ import { Readable } from 'stream';
 import * as sync from '../../sync/sync-index.js';
 import * as util from '../../util/util-index.js';
 
-import { Metrics } from '../../metrics/Metrics.js';
 import { authUser } from '../auth.js';
 import { routeDefinition } from '../router.js';
 
+import { APIMetric } from '@powersync/service-types';
+
 export enum SyncRoutes {
   STREAM = '/sync/stream'
 }
@@ -20,7 +21,7 @@ export const syncStreamed = routeDefinition({
   validator: schema.createTsCodecValidator(util.StreamingSyncRequest, { allowAdditional: true }),
   handler: async (payload) => {
     const { service_context } = payload.context;
-    const { routerEngine, storageEngine, syncContext } = service_context;
+    const { routerEngine, storageEngine, metricsEngine, syncContext } = service_context;
     const headers = payload.request.headers;
     const userAgent = headers['x-user-agent'] ?? headers['user-agent'];
     const clientId = payload.params.client_id;
@@ -49,9 +50,9 @@ export const syncStreamed = routeDefinition({
     const syncRules = bucketStorage.getParsedSyncRules(routerEngine!.getAPI().getParseSyncRulesOptions());
 
     const controller = new AbortController();
-    const tracker = new sync.RequestTracker();
+    const tracker = new sync.RequestTracker(metricsEngine);
     try {
-      Metrics.getInstance().concurrent_connections.add(1);
+      metricsEngine.getUpDownCounter(APIMetric.CONCURRENT_CONNECTIONS).add(1);
       const stream = Readable.from(
         sync.transformToBytesTracked(
           sync.ndjson(
@@ -96,7 +97,7 @@ export const syncStreamed = routeDefinition({
         data: stream,
         afterSend: async () => {
           controller.abort();
-          Metrics.getInstance().concurrent_connections.add(-1);
+          metricsEngine.getUpDownCounter(APIMetric.CONCURRENT_CONNECTIONS).add(-1);
           logger.info(`Sync stream complete`, {
             user_id: syncParams.user_id,
             client_id: clientId,
@@ -108,7 +109,7 @@ export const syncStreamed = routeDefinition({
       });
     } catch (ex) {
       controller.abort();
-      Metrics.getInstance().concurrent_connections.add(-1);
+      metricsEngine.getUpDownCounter(APIMetric.CONCURRENT_CONNECTIONS).add(-1);
     }
   }
 });

package/src/storage/SyncRulesBucketStorage.ts

@@ -73,6 +73,16 @@ export interface SyncRulesBucketStorage
    */
   getParameterSets(checkpoint: util.InternalOpId, lookups: ParameterLookup[]): Promise<SqliteJsonRow[]>;
 
+  /**
+   * Given two checkpoints, return the changes in bucket data and parameters that may have occurred
+   * in that period.
+   *
+   * This is a best-effort optimization:
+   * 1. This may include more changes than what actually occurred.
+   * 2. This may return invalidateDataBuckets or invalidateParameterBuckets instead of of returning
+   *    specific changes.
+   * @param options
+   */
   getCheckpointChanges(options: GetCheckpointChangesOptions): Promise<CheckpointChanges>;
 
   /**
@@ -251,7 +261,7 @@ export interface GetCheckpointChangesOptions {
 }
 
 export interface CheckpointChanges {
-  updatedDataBuckets: string[];
+  updatedDataBuckets: Set<string>;
   invalidateDataBuckets: boolean;
   /** Serialized using JSONBig */
   updatedParameterLookups: Set<string>;
@@ -259,7 +269,7 @@ export interface CheckpointChanges {
 }
 
 export const CHECKPOINT_INVALIDATE_ALL: CheckpointChanges = {
-  updatedDataBuckets: [],
+  updatedDataBuckets: new Set<string>(),
   invalidateDataBuckets: true,
   updatedParameterLookups: new Set<string>(),
   invalidateParameterBuckets: true

package/src/storage/WriteCheckpointAPI.ts

@@ -94,7 +94,6 @@ export interface BaseWriteCheckpointAPI {
  */
 export interface SyncStorageWriteCheckpointAPI extends BaseWriteCheckpointAPI {
   batchCreateCustomWriteCheckpoints(checkpoints: BatchedCustomWriteCheckpointOptions[]): Promise<void>;
-  createCustomWriteCheckpoint(checkpoint: BatchedCustomWriteCheckpointOptions): Promise<bigint>;
   lastWriteCheckpoint(filters: SyncStorageLastWriteCheckpointFilters): Promise<bigint | null>;
 }
 
@@ -104,7 +103,6 @@ export interface SyncStorageWriteCheckpointAPI extends BaseWriteCheckpointAPI {
  */
 export interface WriteCheckpointAPI extends BaseWriteCheckpointAPI {
   batchCreateCustomWriteCheckpoints(checkpoints: CustomWriteCheckpointOptions[]): Promise<void>;
-  createCustomWriteCheckpoint(checkpoint: CustomWriteCheckpointOptions): Promise<bigint>;
   lastWriteCheckpoint(filters: LastWriteCheckpointFilters): Promise<bigint | null>;
 
   watchUserWriteCheckpoint(options: WatchUserWriteCheckpointOptions): AsyncIterable<WriteCheckpointResult>;

package/src/storage/storage-index.ts

@@ -6,6 +6,7 @@ export * from './SourceEntity.js';
 export * from './SourceTable.js';
 export * from './StorageEngine.js';
 export * from './StorageProvider.js';
+export * from './storage-metrics.js';
 export * from './WriteCheckpointAPI.js';
 export * from './BucketStorageFactory.js';
 export * from './BucketStorageBatch.js';

package/src/storage/storage-metrics.ts

@@ -0,0 +1,67 @@
+import { MetricsEngine } from '../metrics/MetricsEngine.js';
+import { logger } from '@powersync/lib-services-framework';
+import { BucketStorageFactory, StorageMetrics } from './BucketStorageFactory.js';
+import { StorageMetric } from '@powersync/service-types';
+
+export function createCoreStorageMetrics(engine: MetricsEngine): void {
+  engine.createObservableGauge({
+    name: StorageMetric.REPLICATION_SIZE_BYTES,
+    description: 'Size of current data stored in PowerSync',
+    unit: 'bytes'
+  });
+
+  engine.createObservableGauge({
+    name: StorageMetric.OPERATION_SIZE_BYTES,
+    description: 'Size of operations stored in PowerSync',
+    unit: 'bytes'
+  });
+
+  engine.createObservableGauge({
+    name: StorageMetric.PARAMETER_SIZE_BYTES,
+    description: 'Size of parameter data stored in PowerSync',
+    unit: 'bytes'
+  });
+}
+
+export function initializeCoreStorageMetrics(engine: MetricsEngine, storage: BucketStorageFactory): void {
+  const replication_storage_size_bytes = engine.getObservableGauge(StorageMetric.REPLICATION_SIZE_BYTES);
+  const operation_storage_size_bytes = engine.getObservableGauge(StorageMetric.OPERATION_SIZE_BYTES);
+  const parameter_storage_size_bytes = engine.getObservableGauge(StorageMetric.PARAMETER_SIZE_BYTES);
+
+  const MINIMUM_INTERVAL = 60_000;
+
+  let cachedRequest: Promise<StorageMetrics | null> | undefined = undefined;
+  let cacheTimestamp = 0;
+
+  const getMetrics = () => {
+    if (cachedRequest == null || Date.now() - cacheTimestamp > MINIMUM_INTERVAL) {
+      cachedRequest = storage.getStorageMetrics().catch((e) => {
+        logger.error(`Failed to get storage metrics`, e);
+        return null;
+      });
+      cacheTimestamp = Date.now();
+    }
+    return cachedRequest;
+  };
+
+  replication_storage_size_bytes.setValueProvider(async () => {
+    const metrics = await getMetrics();
+    if (metrics) {
+      return metrics.replication_size_bytes;
+    }
+  });
+
+  operation_storage_size_bytes.setValueProvider(async () => {
+    const metrics = await getMetrics();
+    if (metrics) {
+      return metrics.operations_size_bytes;
+    }
+  });
+
+  parameter_storage_size_bytes.setValueProvider(async () => {
+    const metrics = await getMetrics();
+    if (metrics) {
+      return metrics.parameters_size_bytes;
+    }
+  });
+}

package/src/sync/BucketChecksumState.ts

@@ -4,11 +4,11 @@ import * as storage from '../storage/storage-index.js';
 import * as util from '../util/util-index.js';
 
 import { ErrorCode, logger, ServiceAssertionError, ServiceError } from '@powersync/lib-services-framework';
+import { JSONBig } from '@powersync/service-jsonbig';
 import { BucketParameterQuerier } from '@powersync/service-sync-rules/src/BucketParameterQuerier.js';
 import { BucketSyncState } from './sync.js';
 import { SyncContext } from './SyncContext.js';
-import { JSONBig } from '@powersync/service-jsonbig';
-import { hasIntersection } from './util.js';
+import { getIntersection, hasIntersection } from './util.js';
 
 export interface BucketChecksumStateOptions {
   syncContext: SyncContext;
@@ -70,10 +70,9 @@ export class BucketChecksumState {
     const storage = this.bucketStorage;
 
     const update = await this.parameterState.getCheckpointUpdate(next);
-    if (update == null) {
+    if (update == null && this.lastWriteCheckpoint == writeCheckpoint) {
       return null;
     }
-
     const { buckets: allBuckets, updatedBuckets } = update;
 
     let dataBucketsNew = new Map<string, BucketSyncState>();
@@ -92,7 +91,7 @@ export class BucketChecksumState {
     }
 
     let checksumMap: util.ChecksumMap;
-    if (updatedBuckets != null) {
+    if (updatedBuckets != INVALIDATE_ALL_BUCKETS) {
       if (this.lastChecksums == null) {
         throw new ServiceAssertionError(`Bucket diff received without existing checksums`);
       }
@@ -115,9 +114,11 @@ export class BucketChecksumState {
         }
       }
 
-      let updatedChecksums = await storage.getChecksums(base.checkpoint, checksumLookups);
-      for (let [bucket, value] of updatedChecksums.entries()) {
-        newChecksums.set(bucket, value);
+      if (checksumLookups.length > 0) {
+        let updatedChecksums = await storage.getChecksums(base.checkpoint, checksumLookups);
+        for (let [bucket, value] of updatedChecksums.entries()) {
+          newChecksums.set(bucket, value);
+        }
       }
       checksumMap = newChecksums;
     } else {
@@ -125,6 +126,7 @@ export class BucketChecksumState {
       const bucketList = [...dataBucketsNew.keys()];
       checksumMap = await storage.getChecksums(base.checkpoint, bucketList);
     }
+
     // Subset of buckets for which there may be new data in this batch.
     let bucketsToFetch: BucketDescription[];
 
@@ -249,6 +251,8 @@ export class BucketChecksumState {
   }
 }
 
+const INVALIDATE_ALL_BUCKETS = Symbol('INVALIDATE_ALL_BUCKETS');
+
 export interface CheckpointUpdate {
   /**
    * All buckets forming part of the checkpoint.
@@ -260,7 +264,7 @@ export interface CheckpointUpdate {
    *
    * If null, assume that any bucket in `buckets` may have been updated.
    */
-  updatedBuckets: Set<string> | null;
+  updatedBuckets: Set<string> | typeof INVALIDATE_ALL_BUCKETS;
 }
 
 export class BucketParameterState {
@@ -291,19 +295,15 @@ export class BucketParameterState {
     this.lookups = new Set<string>(this.querier.parameterQueryLookups.map((l) => JSONBig.stringify(l.values)));
   }
 
-  async getCheckpointUpdate(checkpoint: storage.StorageCheckpointUpdate): Promise<CheckpointUpdate | null> {
+  async getCheckpointUpdate(checkpoint: storage.StorageCheckpointUpdate): Promise<CheckpointUpdate> {
     const querier = this.querier;
-    let update: CheckpointUpdate | null;
+    let update: CheckpointUpdate;
     if (querier.hasDynamicBuckets) {
       update = await this.getCheckpointUpdateDynamic(checkpoint);
     } else {
       update = await this.getCheckpointUpdateStatic(checkpoint);
     }
 
-    if (update == null) {
-      return null;
-    }
-
     if (update.buckets.length > this.context.maxParameterQueryResults) {
       // TODO: Limit number of results even before we get to this point
       // This limit applies _before_ we get the unique set
@@ -325,32 +325,18 @@ export class BucketParameterState {
   /**
    * For static buckets, we can keep track of which buckets have been updated.
    */
-  private async getCheckpointUpdateStatic(
-    checkpoint: storage.StorageCheckpointUpdate
-  ): Promise<CheckpointUpdate | null> {
+  private async getCheckpointUpdateStatic(checkpoint: storage.StorageCheckpointUpdate): Promise<CheckpointUpdate> {
     const querier = this.querier;
     const update = checkpoint.update;
 
     if (update.invalidateDataBuckets) {
       return {
         buckets: querier.staticBuckets,
-        updatedBuckets: null
+        updatedBuckets: INVALIDATE_ALL_BUCKETS
       };
     }
 
-    let updatedBuckets = new Set<string>();
-
-    for (let bucket of update.updatedDataBuckets ?? []) {
-      if (this.staticBuckets.has(bucket)) {
-        updatedBuckets.add(bucket);
-      }
-    }
-
-    if (updatedBuckets.size == 0) {
-      // No change - skip this checkpoint
-      return null;
-    }
-
+    const updatedBuckets = new Set<string>(getIntersection(this.staticBuckets, update.updatedDataBuckets));
     return {
       buckets: querier.staticBuckets,
       updatedBuckets
@@ -360,9 +346,7 @@ export class BucketParameterState {
   /**
    * For dynamic buckets, we need to re-query the list of buckets every time.
    */
-  private async getCheckpointUpdateDynamic(
-    checkpoint: storage.StorageCheckpointUpdate
-  ): Promise<CheckpointUpdate | null> {
+  private async getCheckpointUpdateDynamic(checkpoint: storage.StorageCheckpointUpdate): Promise<CheckpointUpdate> {
     const querier = this.querier;
     const storage = this.bucketStorage;
     const staticBuckets = querier.staticBuckets;
@@ -401,11 +385,11 @@ export class BucketParameterState {
       dynamicBuckets = this.cachedDynamicBuckets;
 
       if (!invalidateDataBuckets) {
-        // TODO: Do set intersection instead
-        for (let bucket of update.updatedDataBuckets ?? []) {
-          if (this.staticBuckets.has(bucket) || this.cachedDynamicBucketSet.has(bucket)) {
-            updatedBuckets.add(bucket);
-          }
+        for (let bucket of getIntersection(this.staticBuckets, update.updatedDataBuckets)) {
+          updatedBuckets.add(bucket);
+        }
+        for (let bucket of getIntersection(this.cachedDynamicBucketSet, update.updatedDataBuckets)) {
+          updatedBuckets.add(bucket);
         }
       }
     }
@@ -415,7 +399,7 @@ export class BucketParameterState {
       return {
         buckets: allBuckets,
         // We cannot track individual bucket updates for dynamic lookups yet
-        updatedBuckets: null
+        updatedBuckets: INVALIDATE_ALL_BUCKETS
       };
     } else {
       return {

package/src/sync/RequestTracker.ts

@@ -1,4 +1,6 @@
-import { Metrics } from '../metrics/Metrics.js';
+import { MetricsEngine } from '../metrics/MetricsEngine.js';
+
+import { APIMetric } from '@powersync/service-types';
 
 /**
  * Record sync stats per request stream.
@@ -7,15 +9,19 @@ export class RequestTracker {
   operationsSynced = 0;
   dataSyncedBytes = 0;
 
+  constructor(private metrics: MetricsEngine) {
+    this.metrics = metrics;
+  }
+
   addOperationsSynced(operations: number) {
     this.operationsSynced += operations;
 
-    Metrics.getInstance().operations_synced_total.add(operations);
+    this.metrics.getCounter(APIMetric.OPERATIONS_SYNCED).add(operations);
   }
 
   addDataSynced(bytes: number) {
     this.dataSyncedBytes += bytes;
 
-    Metrics.getInstance().data_synced_bytes.add(bytes);
+    this.metrics.getCounter(APIMetric.DATA_SYNCED_BYTES).add(bytes);
   }
 }

package/src/sync/util.ts

@@ -1,8 +1,8 @@
 import * as timers from 'timers/promises';
 
+import { SemaphoreInterface } from 'async-mutex';
 import * as util from '../util/util-index.js';
 import { RequestTracker } from './RequestTracker.js';
-import { SemaphoreInterface } from 'async-mutex';
 
 export type TokenStreamOptions = {
   /**
@@ -154,14 +154,35 @@ export function settledPromise<T>(promise: Promise<T>): Promise<PromiseSettledRe
   );
 }
 
-export function hasIntersection<T>(a: Set<T>, b: Set<T>) {
-  if (a.size > b.size) {
-    [a, b] = [b, a];
+export type MapOrSet<T> = Map<T, any> | Set<T>;
+
+/**
+ * Check if two sets have any element(s) in common.
+ */
+export function hasIntersection<T>(a: MapOrSet<T>, b: MapOrSet<T>) {
+  for (let _ of getIntersection(a, b)) {
+    return true;
   }
-  // Now, a is always smaller than b, so iterate over a
-  for (let value of a) {
-    if (b.has(value)) {
-      return true;
+  return false;
+}
+
+/**
+ * Return the intersection of two sets or maps.
+ */
+export function* getIntersection<T>(a: MapOrSet<T>, b: MapOrSet<T>): IterableIterator<T> {
+  // Iterate over the smaller set to reduce the number of lookups
+  if (a.size < b.size) {
+    for (let key of a.keys()) {
+      if (b.has(key)) {
+        yield key;
+      }
+    }
+    return false;
+  } else {
+    for (let key of b.keys()) {
+      if (a.has(key)) {
+        yield key;
+      }
     }
   }
 }

package/src/system/ServiceContext.ts

@@ -1,7 +1,7 @@
 import { LifeCycledSystem, MigrationManager, ServiceIdentifier, container } from '@powersync/lib-services-framework';
 
 import { framework } from '../index.js';
-import * as metrics from '../metrics/Metrics.js';
+import * as metrics from '../metrics/MetricsEngine.js';
 import { PowerSyncMigrationManager } from '../migrations/PowerSyncMigrationManager.js';
 import * as replication from '../replication/replication-index.js';
 import * as routes from '../routes/routes-index.js';
@@ -12,7 +12,7 @@ import { SyncContext } from '../sync/SyncContext.js';
 export interface ServiceContext {
   configuration: utils.ResolvedPowerSyncConfig;
   lifeCycleEngine: LifeCycledSystem;
-  metrics: metrics.Metrics | null;
+  metricsEngine: metrics.MetricsEngine;
   replicationEngine: replication.ReplicationEngine | null;
   routerEngine: routes.RouterEngine | null;
   storageEngine: storage.StorageEngine;
@@ -37,6 +37,11 @@ export class ServiceContextContainer implements ServiceContext {
       configuration
     });
 
+    this.lifeCycleEngine.withLifecycle(this.storageEngine, {
+      start: (storageEngine) => storageEngine.start(),
+      stop: (storageEngine) => storageEngine.shutDown()
+    });
+
     this.syncContext = new SyncContext({
       maxDataFetchConcurrency: configuration.api_parameters.max_data_fetch_concurrency,
       maxBuckets: configuration.api_parameters.max_buckets_per_connection,
@@ -65,8 +70,8 @@ export class ServiceContextContainer implements ServiceContext {
     return container.getOptional(routes.RouterEngine);
   }
 
-  get metrics(): metrics.Metrics | null {
-    return container.getOptional(metrics.Metrics);
+  get metricsEngine(): metrics.MetricsEngine {
+    return container.getImplementation(metrics.MetricsEngine);
   }
 
   get migrations(): PowerSyncMigrationManager {

package/test/src/routes/probes.integration.test.ts

@@ -1,10 +1,10 @@
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import Fastify, { FastifyInstance } from 'fastify';
 import { container } from '@powersync/lib-services-framework';
-import * as auth from '../../../src/routes/auth.js';
-import * as system from '../../../src/system/system-index.js';
+import Fastify, { FastifyInstance } from 'fastify';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { configureFastifyServer } from '../../../src/index.js';
+import * as auth from '../../../src/routes/auth.js';
 import { ProbeRoutes } from '../../../src/routes/endpoints/probes.js';
+import * as system from '../../../src/system/system-index.js';
 
 vi.mock('@powersync/lib-services-framework', async () => {
   const actual = (await vi.importActual('@powersync/lib-services-framework')) as any;
@@ -25,7 +25,7 @@ describe('Probe Routes Integration', () => {
 
   beforeEach(async () => {
     app = Fastify();
-    mockSystem = { routerEngine: {} } as system.ServiceContext;
+    mockSystem = { routerEngine: {}, replicationEngine: {} } as system.ServiceContext;
     await configureFastifyServer(app, { service_context: mockSystem });
     await app.ready();
   });

package/test/src/routes/probes.test.ts

@@ -1,6 +1,6 @@
-import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { container } from '@powersync/lib-services-framework';
-import { startupCheck, livenessCheck, readinessCheck } from '../../../src/routes/endpoints/probes.js';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { livenessCheck, readinessCheck, startupCheck } from '../../../src/routes/endpoints/probes.js';
 
 // Mock the container
 vi.mock('@powersync/lib-services-framework', () => ({
@@ -83,6 +83,7 @@ describe('Probe Routes', () => {
   });
 
   describe('livenessCheck', () => {
+    const mockedContext = { context: { service_context: { replicationEngine: {} } } } as any;
     it('has the correct route definitions', () => {
       expect(livenessCheck.path).toBe('/probes/liveness');
       expect(livenessCheck.method).toBe('GET');
@@ -97,7 +98,7 @@ describe('Probe Routes', () => {
 
       vi.mocked(container.probes.state).mockReturnValue(mockState);
 
-      const response = await livenessCheck.handler();
+      const response = await livenessCheck.handler(mockedContext);
 
       expect(response.status).toBe(200);
       expect(response.data).toEqual(mockState);
@@ -112,7 +113,7 @@ describe('Probe Routes', () => {
 
       vi.mocked(container.probes.state).mockReturnValue(mockState);
 
-      const response = await livenessCheck.handler();
+      const response = await livenessCheck.handler(mockedContext);
 
       expect(response.status).toBe(400);
       expect(response.data).toEqual(mockState);