@hydranium/data-server 1.0.0-next.117 → 1.0.0-next.123

package/src/data-server.ts

@@ -16,6 +16,7 @@ import {
    isSyntheticSource,
    messageError,
    ReferenceSource,
+   SessionClosedError,
    type CloseModelArgs,
    type HydraniumResponseError,
    type Disposable,
@@ -46,13 +47,17 @@ import {
    type StartProfilingArgs,
    type StopProfilingArgs,
    type WriteServerHeapSnapshotArgs,
+   type CloseSessionArgs,
+   type CreateModelDocumentArgs,
+   type CreateSessionArgs,
    type GetModelDocumentArgs,
    type GetProjectForUriArgs,
    type TransferSaveDocumentArgs,
    type WatchModelDocumentArgs,
    type TransferDocumentSavedEvent,
    type TransferDocumentUpdatedEvent,
-   type TransferUpdateDocumentArgs
+   type TransferUpdateDocumentArgs,
+   type TransferUpdateDocumentsArgs
 } from '@hydranium/protocol/data';
 import { REVERT_ON_CLOSE_CLIENT_ID, UNKNOWN_CLIENT_ID } from '@hydranium/core';
 import { defaultDataServerDiagnostics } from './default-diagnostics.js';
@@ -99,6 +104,8 @@ export const referenceSettleTimeoutError = (elapsedMs: number): HydraniumRespons
 
 import type { DataServerDiagnosticsProvider, DataServerProfileCapture } from './diagnostics-provider.js';
 import type {
+   ClientSession,
+   ClientSessionWriteArgs,
    ClientTextDocumentChangeEvent,
    HydraniumLanguageServices,
    LogNameOptions,
@@ -106,6 +113,7 @@ import type {
    ModelService,
    ProjectChangeEvent,
    ServerSharedServices,
+   SessionEndCause,
    TransferEncoder
 } from '@hydranium/core';
 import type { TextDocument } from 'vscode-languageserver-textdocument';
@@ -150,6 +158,18 @@ function fingerprintHash(parts: readonly string[]): string {
    return high + low;
 }
 
+/**
+ * A session a data connection registered with a resume token, and how to end
+ * it from another connection. See {@link DataServer.resumableSessions}.
+ */
+export interface ResumableSession {
+   readonly token: string;
+   /** End the session as its connection closing would. */
+   readonly end: () => void;
+}
+
+const resumableByModelService = new WeakMap<object, Map<string, ResumableSession>>();
+
 /**
  * Which observable state {@link DataServer.computeDocumentFingerprint} hashes to
  * de-dup `onDocumentUpdated` emissions.
@@ -410,17 +430,15 @@ export class DataServer<
    /** Subscription bookkeeping: URI → set of clientIds that subscribed for that URI. */
    protected readonly subscriptions = new Map<string, Set<string>>();
    /**
-    * Open-document bookkeeping: URI → set of clientIds that opened it over THIS
-    * connection.
+    * The client sessions this connection registered, by client id. A document
+    * request carrying one of these ids acts as that session; one carrying any
+    * other id fails.
     *
-    * Kept because the document store releases a client's hold only from an
-    * explicit close, and a client that dies without closing would otherwise keep
-    * the document open forever — resident, and with its last-close revert
-    * suppressed. The store has no notion of which connection a clientId reached
-    * it over, so the connection-scoped set has to live here; teaching it a client
-    * identity is the more robust and much larger alternative.
+    * Entries stay after {@link dispose} has ended their sessions, so a request
+    * still running from before the connection closed reaches the ended handle
+    * and fails with the session's own error.
     */
-   protected readonly openedDocuments = new Map<string, Set<string>>();
+   protected readonly clientSessions = new Map<string, ClientSession<AstNode, AstDiagnostic, TTransfer>>();
    /** Typed client proxy — sends `data-server/on*` notifications back over the same wire. */
    protected readonly clientProxy: DataClientProtocol<TTransfer, TDiagnostic, TProject>;
    protected readonly disposables = new DisposableCollection();
@@ -472,12 +490,12 @@ export class DataServer<
     */
    protected readonly lastEmittedFingerprint = new Map<string, string>();
    /**
-    * URIs whose LAST client just closed. The update handler rebuilds such a
-    * document from its disk content (discarding unsaved in-session edits),
-    * but {@link dispatchPhaseEvent} gates on the subscription map — and the
-    * last close typically also removed the last watcher, so consumers that
-    * only ever fetch via `getModelDocument` would keep showing the
-    * discarded state forever. Marked on the last-close transition (see
+    * URIs the text store just released after their last close. The store
+    * rebuilds such a document from its disk content (discarding unsaved
+    * in-session edits), but {@link dispatchPhaseEvent} gates on the
+    * subscription map — and the last close typically also removed the last
+    * watcher, so consumers that only ever fetch via `getModelDocument` would
+    * keep showing the discarded state forever. Marked on the release (see
     * {@link subscribeToTextDocumentCloses}) and consumed by
     * {@link dispatchPhaseEvent}, which broadcasts the following rebuild's
     * phase event even without a subscription — de-duplicated against
@@ -556,8 +574,10 @@ export class DataServer<
       this.subscribeToProjectManager();
       // Self-register teardown so an adopter that keeps no reference to the
       // server still releases per-connection state, with no lifecycle hook of
-      // its own. See `dispose`.
-      this.disposables.push(connection.onClose(() => this.dispose()));
+      // its own. The client did not close anything itself, so its sessions end
+      // as lost and each document it was the last to have open waits out the
+      // store's revert grace. See `dispose`.
+      this.disposables.push(connection.onClose(() => this.dispose('lost')));
    }
 
    /**
@@ -565,16 +585,19 @@ export class DataServer<
     * handlers and every listener — and clear the subscription map and the
     * per-URI emission-fingerprint cache, so a long-lived shared services
     * bundle does not retain per-connection memory after the connection
-    * closes. Also closes every document still open over this connection (see
-    * {@link closeOpenDocuments}), which is the SHARED store's state rather than
-    * this server's and so outlives the connection unless released here.
+    * closes. Also ends every session this connection registered, which closes
+    * every document it has open: the SHARED store's state rather than this
+    * server's, which outlives the connection unless released here.
     *
-    * Idempotent: subsequent calls are no-ops. Self-fires on
-    * `connection.onClose` so adopters who don't hold a reference still
-    * get per-connection cleanup; adopters that DO hold a reference may
-    * call `dispose()` directly for early teardown.
+    * Idempotent: subsequent calls are no-ops. Runs by itself when the
+    * connection closes, so adopters who don't hold a reference still get
+    * per-connection cleanup; adopters that DO hold a reference may call
+    * `dispose()` directly for early teardown. The sessions end with `cause`:
+    * `'closed'` reverts every document whose last open they close at once,
+    * and `'lost'`, the connection's own close, lets each wait out the revert
+    * grace.
     */
-   dispose(): void {
+   dispose(cause: SessionEndCause = 'closed'): void {
       this.disposed = true;
       // Release an in-flight interactive capture so the process-wide inspector
       // singleton is not left active after the connection closes.
@@ -583,129 +606,184 @@ export class DataServer<
          this.activeProfile = undefined;
          void capture.stop({}).catch(() => undefined);
       }
-      // AFTER `disposables.dispose()`, deliberately: closing a document fires
-      // `onDidClose`, and this server's own close listener would otherwise mark a
-      // revert broadcast for a connection that is already gone.
+      // AFTER `disposables.dispose()`, deliberately: ending a session releases
+      // each document it was the last to hold, and this server's own
+      // `onDidCloseLastOpen` listener would otherwise mark a revert broadcast
+      // for a connection that is already gone.
       this.disposables.dispose();
-      this.closeOpenDocuments();
+      for (const [clientId, session] of this.clientSessions) {
+         this.endSession(clientId, session, cause);
+      }
       this.subscriptions.clear();
       this.lastEmittedFingerprint.clear();
       this.pendingRevertBroadcasts.clear();
    }
 
    // ============================================================
-   // DataServerProtocol implementation — delegates to ModelService.
+   // DataServerProtocol implementation
    // ============================================================
    //
-   // The lifecycle methods (`get` / `update` / `save` / `ready`) forward
-   // straight to `ModelService` and encode the returned `AstDocument` on
-   // the wire boundary via `encoder.astDocumentToTransferDocument`.
-   // Adopters customising lifecycle behaviour (normalisation, supersession,
-   // settled-phase choice, etc.) override on the ModelService subclass —
-   // no override on DataServer is needed.
-
-   async openModelDocument(args: OpenModelArgs): Promise<TransferDocument<TTransfer, TDiagnostic>> {
-      // Register the editor session (idempotent — an already-open document is
-      // attached without rebuilding), then return a fresh snapshot at the
-      // configured target phase. Subsequent build-phase events arrive via
-      // `watchModelDocument`.
-      await this.modelService.open(args);
-      // Record the hold so `dispose` can release it for a client that never
-      // closes. Keyed by the URI as given, because that is what `close` takes.
-      let holders = this.openedDocuments.get(args.uri);
-      if (!holders) {
-         holders = new Set<string>();
-         this.openedDocuments.set(args.uri, holders);
+   // Reads go to `ModelService`; a request that opens or writes a document
+   // goes to the session this connection registered under its `clientId`.
+   // Each answer is encoded on the wire boundary via
+   // `encoder.astDocumentToTransferDocument`. To change how documents are
+   // read, override on the ModelService subclass; to change how they are
+   // written, on the session class the bound `ClientSessionFactory` builds.
+
+   async createSession(args: CreateSessionArgs): Promise<void> {
+      // After teardown nothing would ever end a session registered now, and its
+      // id would stay taken for the life of the process.
+      if (this.disposed) {
+         throw new SessionClosedError(args.clientId);
       }
-      const heldBefore = holders.has(args.clientId);
-      holders.add(args.clientId);
-      let document: TransferDocument<TTransfer, TDiagnostic>;
-      try {
-         document = await this.getModelDocument({ uri: args.uri });
-      } catch (error: unknown) {
-         // The snapshot is the last step, and failing it fails the RPC — so the
-         // caller sees no open and has no reason to close. Release what this
-         // call took rather than leaving a hold nothing on either side is
-         // tracking. A repeat open by a client that already held the URI keeps
-         // it: the hold is one per `(uri, clientId)`, so releasing here would
-         // revoke an earlier open that is still legitimately in use.
-         if (!heldBefore) {
-            try {
-               // Untracked only once the close has actually LANDED. Dropping the
-               // record first gives a failed close the same effect as a
-               // successful one: the hold survives on the document store while
-               // the connection-close drain — the only thing left that would
-               // release it — has no record of it.
-               await this.modelService.close({ uri: args.uri, clientId: args.clientId });
-               this.forgetOpenDocument(args.uri, args.clientId);
-            } catch {
-               // Swallowed so the caller receives the reason the OPEN failed
-               // rather than a secondary one about the cleanup. The record stays
-               // behind deliberately, so teardown retries the release.
-            }
-         }
-         throw error;
+      const resumable = this.resumableSessions();
+      const previous = resumable.get(args.clientId);
+      if (previous && args.resumeToken !== undefined && previous.token === args.resumeToken) {
+         previous.end();
+      }
+      const session = this.modelService.createSession(args.label, args.clientId);
+      this.clientSessions.set(args.clientId, session);
+      if (args.resumeToken !== undefined) {
+         // Taken over by a client registering again after its connection
+         // dropped, so the old session ends as lost: the documents it was the
+         // last to have open wait out the grace, and the new session's opens
+         // find their unsaved text.
+         resumable.set(args.clientId, { token: args.resumeToken, end: () => this.endSession(args.clientId, session, 'lost') });
       }
-      // Project the authoritative client-facing version onto the open snapshot.
-      // `getModelDocument` encodes the freshly-built AST snapshot, whose version
-      // is the `LangiumDocument`'s own `textDocument.version`. That lags the
-      // multi-client synced version whenever the open seeded the synced document
-      // with a caller-supplied version id — and the synced version is what
-      // `ModelService.update`'s optimistic-concurrency gate compares
-      // `basedOn` against, so reporting the snapshot's would make the
-      // caller's first tagged write self-conflict. A genuine concurrent edit
-      // still trips the gate and is reconciled by the caller's replay rather
-      // than predicted here.
-      return { ...document, version: asSnapshotVersion(this.services.workspace.TextDocuments.version(args.uri)) };
    }
 
-   async closeModelDocument(args: CloseModelArgs): Promise<void> {
-      // Closing a session also releases its watch for (uri, clientId) — a
-      // forgotten unwatch would otherwise leak phase-event dispatch until the
-      // connection closes. Idempotent: a close without a prior watch is a no-op.
-      await this.unwatchModelDocument({ uri: args.uri, clientId: args.clientId });
-      this.forgetOpenDocument(args.uri, args.clientId);
-      await this.modelService.close(args);
+   /**
+    * End a session this connection registered. The ended handle stays in
+    * {@link clientSessions}, as it does after {@link dispose}, so a request for
+    * the id still arriving fails with the ended session's own error.
+    */
+   async closeSession(args: CloseSessionArgs): Promise<void> {
+      const session = this.clientSessions.get(args.clientId);
+      if (session) {
+         this.endSession(args.clientId, session);
+      }
+   }
+
+   /** End `session`, registered here under `clientId`, with `cause`, and drop its watches on this connection. */
+   protected endSession(
+      clientId: string,
+      session: ClientSession<AstNode, AstDiagnostic, TTransfer>,
+      cause: SessionEndCause = 'closed'
+   ): void {
+      for (const [uri, watchers] of this.subscriptions) {
+         if (watchers.delete(clientId) && watchers.size === 0) {
+            this.subscriptions.delete(uri);
+            this.lastEmittedFingerprint.delete(uri);
+         }
+      }
+      session.dispose(cause);
    }
 
    /**
-    * Drop the recorded hold for `(uri, clientId)` so {@link dispose} does not
-    * close it a second time. Idempotent.
+    * The resumable sessions of this services tree, by client id, across all of
+    * its data connections: a client resumes on another connection than the one
+    * that registered it. An entry leaves when its session ends, by any path, so
+    * it keeps no ended connection alive.
+    *
+    * A takeover by resume token ends the old session even when its connection
+    * is still alive, so two clients that shared an id and its token would end
+    * each other's sessions. And the token is a guard against colliding with a
+    * session the server has not yet seen end, not a secret: the wire carries no
+    * authentication, and a peer that learns the token can end the session.
     */
-   protected forgetOpenDocument(uri: string, clientId: string): void {
-      const holders = this.openedDocuments.get(uri);
-      if (!holders) {
-         return;
+   protected resumableSessions(): Map<string, ResumableSession> {
+      let sessions = resumableByModelService.get(this.modelService);
+      if (!sessions) {
+         const created = new Map<string, ResumableSession>();
+         // Subscribed once per services tree, and never disposed: it lives as
+         // long as the text store it listens to.
+         this.services.workspace.TextDocuments.onDidCloseSession(event => created.delete(event.clientId));
+         resumableByModelService.set(this.modelService, created);
+         sessions = created;
       }
-      holders.delete(clientId);
-      if (holders.size === 0) {
-         this.openedDocuments.delete(uri);
+      return sessions;
+   }
+
+   async createModelDocument(args: CreateModelDocumentArgs): Promise<TransferDocument<TTransfer, TDiagnostic>> {
+      const session = this.requireSession(args.clientId);
+      await session.create(args.uri, args.text);
+      return this.openedSnapshot(args.uri, () => session.close(args.uri));
+   }
+
+   async updateModelDocuments(args: TransferUpdateDocumentsArgs<TTransfer>): Promise<TransferDocument<TTransfer, TDiagnostic>[]> {
+      const astDocuments = await this.requireSession(args.clientId).updateAll({
+         updates: args.updates.map(update => this.toSessionWrite(update))
+      });
+      return astDocuments.map(
+         astDocument =>
+            this.encoder.astDocumentToTransferDocument(astDocument as never) as unknown as TransferDocument<TTransfer, TDiagnostic>
+      );
+   }
+
+   /**
+    * The session write a wire write request carries. Only the fields a session
+    * takes: a field an adopter adds to a request reaches the session only
+    * through an override of this.
+    */
+   protected toSessionWrite(request: Omit<TransferUpdateDocumentArgs<TTransfer>, 'clientId'>): ClientSessionWriteArgs<TTransfer> {
+      return { uri: request.uri, model: request.model, basedOn: request.basedOn };
+   }
+
+   /**
+    * The session this connection registered under `clientId`, which every
+    * document request acts as. An id that is not one fails with the
+    * closed-session code: opening or writing under it would leave an open that
+    * no session end ever closes.
+    */
+   protected requireSession(clientId: string): ClientSession<AstNode, AstDiagnostic, TTransfer> {
+      const session = this.clientSessions.get(clientId);
+      if (!session) {
+         throw new SessionClosedError(clientId, `Client session ${clientId} was never registered on this connection`);
       }
+      return session;
+   }
+
+   async openModelDocument(args: Pick<OpenModelArgs, 'uri' | 'clientId' | 'options'>): Promise<TransferDocument<TTransfer, TDiagnostic>> {
+      const session = this.requireSession(args.clientId);
+      const wasOpen = this.services.workspace.TextDocuments.isOpenInClient(args.uri, args.clientId);
+      await session.open(args.uri, args.options);
+      // A repeat open keeps the earlier one, which is still in use.
+      return this.openedSnapshot(args.uri, wasOpen ? undefined : () => session.close(args.uri));
    }
 
    /**
-    * Close every document still open over this connection, for the clients that
-    * opened it here.
+    * The snapshot a session's open or create answers with. When the read
+    * fails, `rollback` undoes the open this call made before the failure is
+    * rethrown: the caller sees no open, so nothing on its side would ever close
+    * it.
     *
-    * The document store releases a per-URI hold only from an explicit close, so
-    * without this a client that dies mid-session keeps its documents open for the
-    * lifetime of the process: `isOpenInAnyClient` stays true, the document stays
-    * resident, and the last-close revert never runs. A long-lived multi-client
-    * head is the configuration where a dead client is normal rather than
-    * exceptional, so the leak accumulates there.
-    *
-    * Runs from {@link dispose}, which is synchronous, so each close is fired and
-    * its failure swallowed — a teardown must not reject, and a URI whose close
-    * fails is no worse off than it was before this drain existed.
+    * Versioned with the store's shared version rather than the snapshot's own:
+    * the built document's `textDocument.version` lags the store's whenever the
+    * store assigned the version, and the store's is what a write's `basedOn` is
+    * compared against, so reporting the snapshot's would make the caller's
+    * first based-on write conflict with itself.
     */
-   protected closeOpenDocuments(): void {
-      const held = [...this.openedDocuments];
-      this.openedDocuments.clear();
-      for (const [uri, clientIds] of held) {
-         for (const clientId of clientIds) {
-            void Promise.resolve(this.modelService.close({ uri, clientId })).catch(() => undefined);
-         }
+   protected async openedSnapshot(uri: string, rollback?: () => Promise<void>): Promise<TransferDocument<TTransfer, TDiagnostic>> {
+      let document: TransferDocument<TTransfer, TDiagnostic>;
+      try {
+         document = await this.getModelDocument({ uri });
+      } catch (error: unknown) {
+         // The open's failure is the one the caller needs, not the cleanup's.
+         await rollback?.().catch(() => undefined);
+         throw error;
       }
+      return { ...document, version: asSnapshotVersion(this.services.workspace.TextDocuments.version(uri)) };
+   }
+
+   async closeModelDocument(args: CloseModelArgs): Promise<void> {
+      // Through the handle, so a close under an ended session fails as that
+      // session's calls do, and a session class's own close is honoured.
+      const session = this.requireSession(args.clientId);
+      // Closing also releases the watch for (uri, clientId) — a forgotten
+      // unwatch would otherwise leak phase-event dispatch until the connection
+      // closes. Idempotent: a close without a prior watch is a no-op.
+      await this.unwatchModelDocument({ uri: args.uri, clientId: args.clientId });
+      await session.close(args.uri);
    }
 
    async getModelDocument(args: GetModelDocumentArgs): Promise<TransferDocument<TTransfer, TDiagnostic>> {
@@ -725,12 +803,12 @@ export class DataServer<
    }
 
    async updateModelDocument(args: TransferUpdateDocumentArgs<TTransfer>): Promise<TransferDocument<TTransfer, TDiagnostic>> {
-      const astDocument = await this.modelService.update(args);
+      const astDocument = await this.requireSession(args.clientId).update(this.toSessionWrite(args));
       return this.encoder.astDocumentToTransferDocument(astDocument as never) as unknown as TransferDocument<TTransfer, TDiagnostic>;
    }
 
    async saveModelDocument(args: TransferSaveDocumentArgs<TTransfer>): Promise<TransferDocument<TTransfer, TDiagnostic>> {
-      const astDocument = await this.modelService.save(args);
+      const astDocument = await this.requireSession(args.clientId).save(this.toSessionWrite(args));
       return this.encoder.astDocumentToTransferDocument(astDocument as never) as unknown as TransferDocument<TTransfer, TDiagnostic>;
    }
 
@@ -1196,18 +1274,17 @@ export class DataServer<
    }
 
    /**
-    * Mark the last-close transition per URI (see
-    * {@link pendingRevertBroadcasts}). The listener consults
-    * `isOpenInAnyClient` AFTER the store decremented the closing client's
-    * hold, so a `false` answer means this close was the last one.
+    * Mark each document the store releases after its last close (see
+    * {@link pendingRevertBroadcasts}). A release, not the close itself: a
+    * document whose last client lost its connection is released only once
+    * the revert grace runs out, or when another client opens it meanwhile, and
+    * not at all when a client lost from it opens it again within its own
+    * grace.
     */
    protected subscribeToTextDocumentCloses(): void {
-      const textDocuments = this.services.workspace.TextDocuments;
       this.disposables.push(
-         textDocuments.onDidClose(event => {
-            if (!textDocuments.isOpenInAnyClient(event.document.uri)) {
-               this.pendingRevertBroadcasts.add(this.canonicalKey(event.document.uri));
-            }
+         this.services.workspace.TextDocuments.onDidCloseLastOpen(event => {
+            this.pendingRevertBroadcasts.add(this.canonicalKey(event.uri));
          })
       );
    }