@dxos/echo-pipeline 0.8.3 → 0.8.4-main.05e74ebcff

package/src/db-host/stub.ts

@@ -0,0 +1,33 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import type {
+  DeleteFromQueueRequest,
+  InsertIntoQueueRequest,
+  QueryQueueRequest,
+  QueueQueryResult,
+  QueueService,
+  SyncQueueRequest,
+} from '@dxos/protocols/proto/dxos/client/services';
+
+/**
+ * Stub implementation for when Edge is not available.
+ */
+export class QueueServiceStub implements QueueService {
+  queryQueue(request: QueryQueueRequest): Promise<QueueQueryResult> {
+    throw new Error('Not available.');
+  }
+
+  insertIntoQueue(request: InsertIntoQueueRequest): Promise<void> {
+    throw new Error('Not available.');
+  }
+
+  deleteFromQueue(request: DeleteFromQueueRequest): Promise<void> {
+    throw new Error('Not available.');
+  }
+
+  syncQueue(request: SyncQueueRequest): Promise<void> {
+    throw new Error('Not available.');
+  }
+}

package/src/edge/echo-edge-replicator.test.ts

@@ -7,7 +7,8 @@ import { getRandomPort } from 'get-port-please';
 import { describe, expect, onTestFinished, test } from 'vitest';
 
 import { Event } from '@dxos/async';
-import { createEphemeralEdgeIdentity, EdgeClient, MessageSchema } from '@dxos/edge-client';
+import { Context } from '@dxos/context';
+import { EdgeClient, type EdgeHttpClient, MessageSchema, createEphemeralEdgeIdentity } from '@dxos/edge-client';
 import { createTestEdgeWsServer } from '@dxos/edge-client/testing';
 import { PublicKey, SpaceId } from '@dxos/keys';
 import { EdgeService } from '@dxos/protocols';
@@ -16,8 +17,8 @@ import { createBuf } from '@dxos/protocols/buf';
 import type { Peer } from '@dxos/protocols/proto/dxos/edge/messenger';
 import { openAndClose } from '@dxos/test-utils';
 
+import type { AutomergeReplicatorConnection, AutomergeReplicatorContext } from '../automerge';
 import { EchoEdgeReplicator } from './echo-edge-replicator';
-import type { EchoReplicatorContext, ReplicatorConnection } from '../automerge';
 
 describe('EchoEdgeReplicator', () => {
   test('reconnects', async () => {
@@ -27,7 +28,7 @@ describe('EchoEdgeReplicator', () => {
 
     const { context, connectionOpen } = createMockContext();
     const replicator = await connectReplicator(client, context);
-    await replicator.connectToSpace(spaceId);
+    await replicator.connectToSpace(Context.default(), spaceId);
 
     client.setIdentity(await createEphemeralEdgeIdentity());
     await connectionOpen.waitForCount(1);
@@ -54,7 +55,7 @@ describe('EchoEdgeReplicator', () => {
         documentSpaceId: { [documentId]: spaceId },
       });
       const replicator = await connectReplicator(client, context);
-      await replicator.connectToSpace(spaceId);
+      await replicator.connectToSpace(Context.default(), spaceId);
 
       await expect.poll(() => openConnections.length === 1).toBeTruthy();
       expect(openConnections[0].shouldAdvertise({ documentId })).toBeTruthy();
@@ -68,7 +69,7 @@ describe('EchoEdgeReplicator', () => {
       const remoteCollections: { [peerId: string]: { [documentId: string]: boolean } } = {};
       const { context, openConnections } = createMockContext({ remoteCollections });
       const replicator = await connectReplicator(client, context);
-      await replicator.connectToSpace(spaceId);
+      await replicator.connectToSpace(Context.default(), spaceId);
 
       await expect.poll(() => openConnections.length === 1).toBeTruthy();
       const connection = openConnections[0];
@@ -78,9 +79,10 @@ describe('EchoEdgeReplicator', () => {
     });
   });
 
-  const connectReplicator = async (client: EdgeClient, context: EchoReplicatorContext) => {
-    const replicator = new EchoEdgeReplicator({ edgeConnection: client });
-    await replicator.connect(context);
+  const connectReplicator = async (client: EdgeClient, context: AutomergeReplicatorContext) => {
+    // EdgeHttpClient functionality is not tested here.
+    const replicator = new EchoEdgeReplicator({ edgeConnection: client, edgeHttpClient: {} as EdgeHttpClient });
+    await replicator.connect(Context.default(), context);
     onTestFinished(() => replicator.disconnect());
     return replicator;
   };
@@ -99,7 +101,7 @@ const createMockContext = (args?: {
   documentSpaceId?: { [documentId: string]: SpaceId };
 }) => {
   const connectionOpen = new Event();
-  const openConnections: ReplicatorConnection[] = [];
+  const openConnections: AutomergeReplicatorConnection[] = [];
   return {
     context: {
       getContainingSpaceIdForDocument: async (documentId) => args?.documentSpaceId?.[documentId] ?? null,
@@ -119,7 +121,7 @@ const createMockContext = (args?: {
       },
 
       peerId: PublicKey.random().toHex(),
-    } satisfies EchoReplicatorContext,
+    } satisfies AutomergeReplicatorContext,
 
     openConnections,
     connectionOpen,

package/src/edge/echo-edge-replicator.ts

@@ -2,33 +2,43 @@
 // Copyright 2024 DXOS.org
 //
 
-import { cbor } from '@automerge/automerge-repo';
+import * as A from '@automerge/automerge';
+import { type DocumentId, type Heads, cbor } from '@automerge/automerge-repo';
 
-import { Mutex, scheduleTask, scheduleMicroTask } from '@dxos/async';
+import { Mutex, scheduleMicroTask, scheduleTask } from '@dxos/async';
 import { Context, Resource } from '@dxos/context';
 import { randomUUID } from '@dxos/crypto';
 import type { CollectionId } from '@dxos/echo-protocol';
-import { type EdgeConnection } from '@dxos/edge-client';
+import { type EdgeConnection, type EdgeHttpClient } from '@dxos/edge-client';
 import { invariant } from '@dxos/invariant';
 import type { SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
-import { EdgeService, type AutomergeProtocolMessage, type PeerId } from '@dxos/protocols';
+import {
+  type AutomergeProtocolMessage,
+  DocumentCodec,
+  EdgeService,
+  type ErrorProtocolMessage,
+  type ExportBundleRequest,
+  type ImportBundleRequest,
+  type PeerId,
+} from '@dxos/protocols';
 import { buf } from '@dxos/protocols/buf';
 import {
   type Message as RouterMessage,
   MessageSchema as RouterMessageSchema,
 } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
-import { bufferToArray } from '@dxos/util';
+import { trace } from '@dxos/tracing';
+import { bufferToArray, compositeKey, setDeep } from '@dxos/util';
 
-import { InflightRequestLimiter } from './inflight-request-limiter';
 import {
+  type AutomergeReplicatorConnection,
+  type AutomergeReplicatorContext,
+  type EdgeAutomergeReplicator,
+  type ShouldAdvertiseProps,
+  type ShouldSyncCollectionProps,
   getSpaceIdFromCollectionId,
-  type EchoReplicator,
-  type EchoReplicatorContext,
-  type ReplicatorConnection,
-  type ShouldAdvertiseParams,
-  type ShouldSyncCollectionParams,
 } from '../automerge';
+import { InflightRequestLimiter } from './inflight-request-limiter';
 
 /**
  * Delay before restarting the connection after receiving a forbidden error.
@@ -37,30 +47,34 @@ const INITIAL_RESTART_DELAY = 500;
 const RESTART_DELAY_JITTER = 250;
 const MAX_RESTART_DELAY = 5000;
 
-export type EchoEdgeReplicatorParams = {
+export type EchoEdgeReplicatorProps = {
   edgeConnection: EdgeConnection;
+  edgeHttpClient: EdgeHttpClient;
   disableSharePolicy?: boolean;
 };
 
-export class EchoEdgeReplicator implements EchoReplicator {
+@trace.resource()
+export class EchoEdgeReplicator implements EdgeAutomergeReplicator {
   private readonly _edgeConnection: EdgeConnection;
+  private readonly _edgeHttpClient: EdgeHttpClient;
   private readonly _mutex = new Mutex();
 
   private _ctx?: Context = undefined;
-  private _context: EchoReplicatorContext | null = null;
+  private _context: AutomergeReplicatorContext | null = null;
   private _connectedSpaces = new Set<SpaceId>();
   private _connections = new Map<SpaceId, EdgeReplicatorConnection>();
   private _sharePolicyEnabled = true;
 
-  constructor({ edgeConnection, disableSharePolicy }: EchoEdgeReplicatorParams) {
+  constructor({ edgeConnection, edgeHttpClient, disableSharePolicy }: EchoEdgeReplicatorProps) {
     this._edgeConnection = edgeConnection;
+    this._edgeHttpClient = edgeHttpClient;
     this._sharePolicyEnabled = !disableSharePolicy;
   }
 
-  async connect(context: EchoReplicatorContext): Promise<void> {
+  async connect(ctx: Context, context: AutomergeReplicatorContext): Promise<void> {
     log('connecting...', { peerId: context.peerId, connectedSpaces: this._connectedSpaces.size });
     this._context = context;
-    this._ctx = Context.default();
+    this._ctx = ctx.derive();
     this._ctx.onDispose(
       this._edgeConnection.onReconnected(() => {
         this._ctx && scheduleMicroTask(this._ctx, () => this._handleReconnect());
@@ -79,7 +93,7 @@ export class EchoEdgeReplicator implements EchoReplicator {
 
     if (this._context !== null) {
       for (const spaceId of spaces) {
-        await this._openConnection(spaceId);
+        await this._openConnection(this._ctx!, spaceId);
       }
     }
   }
@@ -94,7 +108,9 @@ export class EchoEdgeReplicator implements EchoReplicator {
     this._connections.clear();
   }
 
-  async connectToSpace(spaceId: SpaceId): Promise<void> {
+  @trace.span({ showInBrowserTimeline: true, op: 'replication' })
+  async connectToSpace(ctx: Context, spaceId: SpaceId): Promise<void> {
+    log('connectToSpace', { spaceId });
     using _guard = await this._mutex.acquire();
 
     if (this._connectedSpaces.has(spaceId)) {
@@ -102,9 +118,8 @@ export class EchoEdgeReplicator implements EchoReplicator {
     }
     this._connectedSpaces.add(spaceId);
 
-    // Check if AM-repo requested that we connect to remote peers.
     if (this._context !== null) {
-      await this._openConnection(spaceId);
+      await this._openConnection(ctx, spaceId);
     }
   }
 
@@ -120,21 +135,25 @@ export class EchoEdgeReplicator implements EchoReplicator {
     }
   }
 
-  private async _openConnection(spaceId: SpaceId, reconnects: number = 0): Promise<void> {
+  @trace.span({ showInBrowserTimeline: true, op: 'replication' })
+  private async _openConnection(ctx: Context, spaceId: SpaceId, reconnects: number = 0): Promise<void> {
     invariant(this._context);
     invariant(!this._connections.has(spaceId));
 
     let restartScheduled = false;
 
     const connection = new EdgeReplicatorConnection({
+      edgeHttpClient: this._edgeHttpClient,
       edgeConnection: this._edgeConnection,
       spaceId,
       context: this._context,
       sharedPolicyEnabled: this._sharePolicyEnabled,
       onRemoteConnected: async () => {
+        log.trace('dxos.echo.edge.replicator.onRemoteConnected', { spaceId });
         this._context?.onConnectionOpen(connection);
       },
       onRemoteDisconnected: async () => {
+        log.trace('dxos.echo.edge.replicator.onRemoteDisconnected', { spaceId });
         this._context?.onConnectionClosed(connection);
       },
       onRestartRequested: async () => {
@@ -149,20 +168,21 @@ export class EchoEdgeReplicator implements EchoReplicator {
 
         restartScheduled = true;
         scheduleTask(
-          this._ctx,
+          this._ctx!,
           async () => {
             using _guard = await this._mutex.acquire();
             if (this._connections.get(spaceId) !== connection) {
               return;
             }
 
-            const ctx = this._ctx;
-            await connection.close(); // Will call onRemoteDisconnected
+            const replicatorCtx = this._ctx;
+            await connection.close();
             this._connections.delete(spaceId);
-            if (ctx?.disposed) {
+            if (replicatorCtx?.disposed) {
               return;
             }
-            await this._openConnection(spaceId, reconnects + 1);
+            log.trace('dxos.echo.edge.replicator.restart', { spaceId, reconnects, restartDelay });
+            await this._openConnection(replicatorCtx!, spaceId, reconnects + 1);
           },
           restartDelay,
         );
@@ -170,14 +190,15 @@ export class EchoEdgeReplicator implements EchoReplicator {
     });
     this._connections.set(spaceId, connection);
 
-    await connection.open();
+    await connection.open(ctx);
   }
 }
 
-type EdgeReplicatorConnectionsParams = {
+type EdgeReplicatorConnectionsProps = {
   edgeConnection: EdgeConnection;
+  edgeHttpClient: EdgeHttpClient;
   spaceId: SpaceId;
-  context: EchoReplicatorContext;
+  context: AutomergeReplicatorContext;
   sharedPolicyEnabled: boolean;
   onRemoteConnected: () => Promise<void>;
   onRemoteDisconnected: () => Promise<void>;
@@ -187,16 +208,19 @@ type EdgeReplicatorConnectionsParams = {
 const MAX_INFLIGHT_REQUESTS = 5;
 const MAX_RATE_LIMIT_WAIT_TIME_MS = 3000;
 
-class EdgeReplicatorConnection extends Resource implements ReplicatorConnection {
+class EdgeReplicatorConnection extends Resource implements AutomergeReplicatorConnection {
+  private readonly _connectionId = randomUUID();
   private readonly _edgeConnection: EdgeConnection;
+  private readonly _edgeHttpClient: EdgeHttpClient;
   private readonly _remotePeerId: string | null = null;
   private readonly _targetServiceId: string;
   private readonly _spaceId: SpaceId;
-  private readonly _context: EchoReplicatorContext;
+  private readonly _context: AutomergeReplicatorContext;
   private readonly _sharedPolicyEnabled: boolean;
   private readonly _onRemoteConnected: () => Promise<void>;
   private readonly _onRemoteDisconnected: () => Promise<void>;
   private readonly _onRestartRequested: () => void;
+  private _sequence = 0;
 
   private _requestLimiter = new InflightRequestLimiter({
     maxInflightRequests: MAX_INFLIGHT_REQUESTS,
@@ -210,23 +234,26 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
 
   constructor({
     edgeConnection,
+    edgeHttpClient,
     spaceId,
     context,
     sharedPolicyEnabled,
     onRemoteConnected,
     onRemoteDisconnected,
     onRestartRequested,
-  }: EdgeReplicatorConnectionsParams) {
+  }: EdgeReplicatorConnectionsProps) {
     super();
     this._edgeConnection = edgeConnection;
+    this._edgeHttpClient = edgeHttpClient;
     this._spaceId = spaceId;
     this._context = context;
     // Generate a unique peer id for every connection.
     // This way automerge-repo will have separate sync states for every connection.
     // This is important because the previous connection might have had some messages that failed to deliver
     // abd if we don't clear the sync-state, automerge will not attempt to deliver them again.
-    this._remotePeerId = `${EdgeService.AUTOMERGE_REPLICATOR}:${spaceId}-${randomUUID()}`;
-    this._targetServiceId = `${EdgeService.AUTOMERGE_REPLICATOR}:${spaceId}`;
+    this._targetServiceId = compositeKey(EdgeService.AUTOMERGE_REPLICATOR, spaceId);
+    // TODO(wittjosiah): Use compositeKey.
+    this._remotePeerId = `${this._targetServiceId}-${this._connectionId}`;
     this._sharedPolicyEnabled = sharedPolicyEnabled;
     this._onRemoteConnected = onRemoteConnected;
     this._onRemoteDisconnected = onRemoteDisconnected;
@@ -242,7 +269,7 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
       write: async (message: AutomergeProtocolMessage, controller) => {
         await this._requestLimiter.rateLimit(message);
 
-        await this._sendMessage(message);
+        await this._sendMessage(this._ctx, message);
       },
     });
   }
@@ -252,17 +279,30 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
 
     await this._requestLimiter.open();
 
-    // TODO: handle reconnects
     this._ctx.onDispose(
       this._edgeConnection.onMessage((msg: RouterMessage) => {
         this._onMessage(msg);
       }),
     );
 
+    let firstReconnect = true;
+    this._ctx.onDispose(
+      // NOTE: This will fire immediately if the connection is already open.
+      this._edgeConnection.onReconnected(async () => {
+        if (firstReconnect) {
+          log.verbose('first reconnect skipped');
+          firstReconnect = false;
+          return;
+        }
+
+        this._onRestartRequested();
+      }),
+    );
+
     await this._onRemoteConnected();
   }
 
-  protected override async _close(): Promise<void> {
+  protected override async _close(ctx: Context): Promise<void> {
     log('closing...');
     this._readableStreamController.close();
 
@@ -276,7 +316,7 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
     return this._remotePeerId;
   }
 
-  async shouldAdvertise(params: ShouldAdvertiseParams): Promise<boolean> {
+  async shouldAdvertise(params: ShouldAdvertiseProps): Promise<boolean> {
     if (!this._sharedPolicyEnabled) {
       return true;
     }
@@ -301,7 +341,7 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
     return spaceId === this._spaceId;
   }
 
-  shouldSyncCollection(params: ShouldSyncCollectionParams): boolean {
+  shouldSyncCollection(params: ShouldSyncCollectionProps): boolean {
     if (!this._sharedPolicyEnabled) {
       return true;
     }
@@ -317,8 +357,7 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
 
     const payload = cbor.decode(message.payload!.value) as AutomergeProtocolMessage;
     log.verbose('received', {
-      type: payload.type,
-      documentId: payload.type === 'sync' && payload.documentId,
+      ...getMessageInfo(payload),
       remoteId: this._remotePeerId,
     });
 
@@ -327,11 +366,33 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
     this._processMessage(payload);
   }
 
+  get bundleSyncEnabled(): boolean {
+    return true;
+  }
+
+  async pushBundle(ctx: Context, bundle: { documentId: DocumentId; data: Uint8Array; heads: Heads }[]) {
+    const request: ImportBundleRequest = {
+      bundle: bundle.map(({ documentId, data, heads }) => ({
+        documentId,
+        mutation: DocumentCodec.encode(data),
+        heads,
+      })),
+    };
+    await this._edgeHttpClient.importBundle(ctx, this._spaceId, request);
+  }
+
+  async pullBundle(ctx: Context, docHeads: Record<DocumentId, Heads>): Promise<Record<DocumentId, Uint8Array>> {
+    const request: ExportBundleRequest = { docHeads };
+    const response = await this._edgeHttpClient.exportBundle(ctx, this._spaceId, request);
+    return Object.fromEntries(response.bundle.map((doc) => [doc.documentId, DocumentCodec.decode(doc.mutation)]));
+  }
+
   private _processMessage(message: AutomergeProtocolMessage): void {
     // There's a race between the credentials being replicated that are needed for access control and the data replication.
     // AutomergeReplicator might return a Forbidden error if the credentials are not yet replicated.
     // We restart the connection with some delay to account for that.
-    if (isForbiddenErrorMessage(message)) {
+    if (isErrorMessage(message)) {
+      log.verbose('stream error', { error: (message as ErrorProtocolMessage).message });
       this._onRestartRequested();
       return;
     }
@@ -341,33 +402,58 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
     this._readableStreamController.enqueue(message);
   }
 
-  private async _sendMessage(message: AutomergeProtocolMessage): Promise<void> {
+  private async _sendMessage(ctx: Context, message: AutomergeProtocolMessage): Promise<void> {
     // Fix the peer id.
     (message as any).targetId = this._targetServiceId as PeerId;
 
+    // Note: This is used on EDGE to detect out-of-order messages per connection.
+    setDeep(message, ['metadata', 'dxos_sequence'], this._getSequence());
+    setDeep(message, ['metadata', 'dxos_connectionId'], this._connectionId);
+
     log.verbose('sending...', {
-      type: message.type,
-      documentId: message.type === 'sync' && message.documentId,
+      ...getMessageInfo(message),
       remoteId: this._remotePeerId,
     });
 
     const encoded = cbor.encode(message);
 
-    await this._edgeConnection.send(
-      buf.create(RouterMessageSchema, {
-        serviceId: this._targetServiceId,
-        source: {
-          identityKey: this._edgeConnection.identityKey,
-          peerKey: this._edgeConnection.peerKey,
-        },
-        payload: { value: bufferToArray(encoded) },
-      }),
-    );
+    try {
+      await this._edgeConnection.send(
+        ctx,
+        buf.create(RouterMessageSchema, {
+          serviceId: this._targetServiceId,
+          source: {
+            identityKey: this._edgeConnection.identityKey,
+            peerKey: this._edgeConnection.peerKey,
+          },
+          payload: { value: bufferToArray(encoded) },
+        }),
+      );
+    } catch (err) {
+      log.error('failed to send message', { err });
+    }
+  }
+
+  private _getSequence(): number {
+    return this._sequence++;
   }
 }
 
 /**
  * This message is sent by EDGE AutomergeReplicator when the authorization is denied.
  */
-const isForbiddenErrorMessage = (message: AutomergeProtocolMessage) =>
-  message.type === 'error' && message.message === 'Forbidden';
+const isErrorMessage = (message: AutomergeProtocolMessage) => message.type === 'error';
+
+const getMessageInfo = (msg: AutomergeProtocolMessage) => {
+  const { have, heads, need, changes } = msg.type === 'sync' ? A.decodeSyncMessage(msg.data) : {};
+  return {
+    type: msg.type,
+    documentId: 'documentId' in msg ? msg.documentId : undefined,
+    collectionId: 'collectionId' in msg ? msg.collectionId : undefined,
+    have,
+    heads,
+    need,
+    changes: changes?.length,
+    sequence: msg.metadata?.dxos_sequence,
+  };
+};