archify-vscode-ext/src/sync/liveSpace.ts

import * as path from 'node:path';
import * as vscode from 'vscode';
import * as Y from 'yjs';
import { Awareness } from 'y-protocols/awareness';
import { syncSocketUrl } from './api';
import { SyncAuth } from './auth';
import { ConnectionState, LiveConnection } from './connection';
import { FileBridge } from './fileBridge';
import { OfflineStore } from './offlineStore';
import { colorFor } from './cursors';
import type { LiveCursor, LiveSelection } from './liveBridge';
import { Link } from './protocol';

/** How a space starts: as found in the workspace, just shared from here, or just joined. */
export type StartMode = 'open' | 'share' | 'join';

/** Edits loaded from the offline copy; not stored again. */
const FROM_COPY = Symbol('offline copy');
const COMPACT_AFTER = 200;

export interface Presence {
  user: { id: string; name: string };
  /** Model-relative path of the file the person has in front of them. */
  file?: string;
  /** Where the person's cursor is: Yjs relative positions, so it stays put while others type. */
  cursor?: { file: string; anchor: unknown; head: unknown };
}

/**
 * One shared model in this window: the Yjs document, its offline copy, the
 * connection to the server, the bridge to the files, and presence.
 */
export class LiveSpace implements vscode.Disposable {
  readonly doc = new Y.Doc();
  readonly awareness = new Awareness(this.doc);
  private readonly store: OfflineStore;
  private readonly local = { origin: 'local' };
  private readonly bridge: FileBridge;
  private connection?: LiveConnection;
  private readonly disposables: vscode.Disposable[] = [];
  private hadCopy = false;
  private firstSync = true;

  constructor(
    readonly root: vscode.Uri,
    readonly link: Link,
    private readonly mode: StartMode,
    private readonly services: {
      context: vscode.ExtensionContext;
      auth: SyncAuth;
      log: vscode.LogOutputChannel;
      onChange: () => void;
      /** Someone's cursor moved. */
      onCursors?: () => void;
    },
  ) {
    const dir = path.join(services.context.globalStorageUri.fsPath, 'sync');
    this.store = new OfflineStore(dir, link.space);
    this.bridge = new FileBridge(root, this.doc, this.local, services.log);
  }

  get state(): ConnectionState {
    return this.connection?.state ?? 'connecting';
  }

  get name(): string {
    return this.link.name ?? path.basename(this.root.fsPath);
  }

  /** Distinct people in the space, this one included. */
  people(): Presence[] {
    const seen = new Map<string, Presence>();
    for (const state of this.awareness.getStates().values()) {
      const p = state as Partial<Presence>;
      if (p.user?.id) seen.set(p.user.id, { user: p.user, file: p.file });
    }
    return [...seen.values()];
  }

  /** Puts this person's cursor in a shared file; false when the file is not part of this model. */
  setCursor(uri: vscode.Uri, selection: LiveSelection | undefined): boolean {
    const rel = this.bridge.relative(uri);
    if (!rel) return false;
    const current = this.awareness.getLocalState() as Presence | null;
    if (!current) return true;
    if (!selection) {
      // Only the file the cursor is in can take it away.
      if (current.cursor?.file === rel) this.awareness.setLocalStateField('cursor', undefined);
      return true;
    }
    const text = this.doc.getMap<Y.Text>('files').get(rel);
    if (!(text instanceof Y.Text)) return true;
    const at = (offset: number) => Y.relativePositionToJSON(Y.createRelativePositionFromTypeIndex(text, Math.max(0, Math.min(offset, text.length))));
    this.awareness.setLocalStateField('cursor', { file: rel, anchor: at(selection.anchor), head: at(selection.head) });
    return true;
  }

  /** Other people's cursors in a file, as offsets into its text here. */
  cursors(uri: vscode.Uri): LiveCursor[] {
    const rel = this.bridge.relative(uri);
    if (!rel) return [];
    const found: LiveCursor[] = [];
    for (const [client, state] of this.awareness.getStates()) {
      if (client === this.doc.clientID) continue;
      const p = state as Partial<Presence>;
      if (!p.user?.id || p.cursor?.file !== rel) continue;
      const index = (json: unknown) => {
        try {
          return Y.createAbsolutePositionFromRelativePosition(Y.createRelativePositionFromJSON(json), this.doc)?.index;
        } catch {
          return undefined;
        }
      };
      const anchor = index(p.cursor.anchor);
      const head = index(p.cursor.head);
      if (anchor === undefined || head === undefined) continue;
      found.push({ name: p.user.name, color: colorFor(p.user.id), anchor, head });
    }
    return found;
  }

  async start(): Promise<void> {
    const { log } = this.services;
    this.hadCopy = this.store.exists();
    for (const update of this.store.load()) Y.applyUpdate(this.doc, update, FROM_COPY);
    // Registered before the connection's listener, so every edit is on disk before it is sent.
    this.doc.on('update', (update: Uint8Array, origin: unknown) => {
      if (origin === FROM_COPY) return;
      this.store.append(update);
      if (this.store.pending >= COMPACT_AFTER) this.compact();
    });

    if (this.hadCopy) {
      // Edits made while the extension was not running (git pull, another tool).
      await this.bridge.pushAll(true);
    } else if (this.mode === 'share') {
      await this.bridge.pushAll(false);
    }
    // Without a copy, local edits wait for the server's files, so a file
    // does not start twice (once here, once there) and lose one side.
    this.bridge.start(this.hadCopy || this.mode === 'share');

    const session = await this.services.auth.session(this.link.server);
    if (session) this.awareness.setLocalState({ user: { id: session.account.id, name: session.account.display_name } } satisfies Presence);
    this.disposables.push(
      vscode.window.onDidChangeActiveTextEditor(() => this.updateFile()),
      vscode.window.tabGroups.onDidChangeTabs(() => this.updateFile()),
      { dispose: () => this.awareness.destroy() },
    );
    this.updateFile();
    let people = '';
    this.awareness.on('change', () => {
      // Cursors move all the time; the status bar and pages only care who is where.
      const now = JSON.stringify(this.people());
      if (now !== people) {
        people = now;
        this.services.onChange();
      }
      this.services.onCursors?.();
    });

    this.connection = new LiveConnection({
      url: syncSocketUrl(this.link.server, this.link.space),
      token: (force) => this.services.auth.accessToken(this.link.server, force),
      doc: this.doc,
      awareness: this.awareness,
      onState: (state) => {
        log.info(`sync: ${this.name} is ${state}`);
        this.services.onChange();
      },
      onSynced: () => void this.synced(),
      log: (message) => log.info(message),
    });
  }

  private async synced(): Promise<void> {
    if (!this.firstSync) return;
    this.firstSync = false;
    if (!this.hadCopy && this.mode !== 'share') {
      // First time here: the server's files come to disk, then local-only files go up.
      await this.bridge.pullAll();
      this.bridge.start(true);
      this.hadCopy = true;
    }
    this.compact();
  }

  private compact(): void {
    try {
      this.store.compact(Y.encodeStateAsUpdate(this.doc));
    } catch (error) {
      this.services.log.warn(`sync: could not compact the offline copy (${String(error)})`);
    }
  }

  private updateFile(): void {
    const current = this.awareness.getLocalState() as Presence | null;
    if (!current) return;
    // The text editor, or the tab of an Archify page (the model page, a diagram, a registry).
    const input = vscode.window.tabGroups.activeTabGroup.activeTab?.input;
    const tab = input instanceof vscode.TabInputCustom || input instanceof vscode.TabInputText ? input.uri : undefined;
    const uri = vscode.window.activeTextEditor?.document.uri ?? tab;
    const file = uri ? this.bridge.relative(uri) : undefined;
    if (current.file !== file) this.awareness.setLocalState({ ...current, file });
  }

  /** After signing in again. */
  reconnect(): void {
    void this.services.auth.session(this.link.server).then((session) => {
      if (session && !this.awareness.getLocalState()) {
        this.awareness.setLocalState({ user: { id: session.account.id, name: session.account.display_name } } satisfies Presence);
      }
      this.connection?.reconnectNow();
    });
  }

  /** Stops syncing this model here and forgets the offline copy. */
  forget(): void {
    this.dispose();
    this.store.remove();
  }

  dispose(): void {
    this.connection?.dispose();
    this.connection = undefined;
    this.bridge.dispose();
    for (const d of this.disposables.splice(0)) d.dispose();
    if (this.store.pending) this.compact();
    this.doc.destroy();
  }
}