Session recovery

// The whole point of this example, with no DOM in it.
//
// Everything a Cap'n Web client has to do about disconnection lives here:
// noticing one, throwing away the capabilities it invalidated, establishing a
// fresh session, and picking the event stream back up without a gap.

import { newWebSocketRpcSession, RpcTarget } from './vendor/capnweb.js';

/**
 * The object the server calls back into.
 *
 * Passing this over RPC gives the server a stub for it, and calling a method
 * on that stub is an RPC in the other direction. This is all "bidirectional
 * calling" is: there is no separate subscription mechanism.
 */
class EventSink extends RpcTarget {
  #onEvent;
  #onGap;

  constructor({ onEvent, onGap }) {
    super();
    this.#onEvent = onEvent;
    this.#onGap = onGap;
  }

  onEvent(event) {
    this.#onEvent(event);
  }

  onGap(sinceId) {
    this.#onGap(sinceId);
  }
}

/**
 * A client that reconnects.
 *
 * `report` is called with a log line for the UI; `onEvent` with each event as
 * it arrives. Everything else is internal.
 */
export class RecoveringClient {
  #url;
  #token;
  #report;
  #onEvent;
  #onStateChange;

  /** Set while connected. All four are invalidated together by a disconnect. */
  #socket = null;
  #api = null;
  #authed = null;
  #subscription = null;

  /**
   * The last authenticated stub we held, kept after teardown purely so the
   * demo can call a method on it and show what a dead stub does.
   */
  #staleAuthed = null;

  /**
   * The resume token: the id of the last event we actually processed.
   *
   * This is the only thing that survives a reconnect, and it survives because
   * it lives out here in our own state rather than in anything the session
   * owns. A stub cannot survive; a number can.
   */
  #cursor = null;

  /** Set when the caller asked to stop, to tell a deliberate close from a drop. */
  #closing = false;

  #state = 'offline';

  constructor({ url, token, report, onEvent, onStateChange }) {
    this.#url = url;
    this.#token = token;
    this.#report = report;
    this.#onEvent = onEvent;
    this.#onStateChange = onStateChange ?? (() => {});
  }

  get state() {
    return this.#state;
  }

  get cursor() {
    return this.#cursor;
  }

  #setState(state) {
    this.#state = state;
    this.#onStateChange(state);
  }

  /**
   * Connect, authenticate, and subscribe -- in one round trip.
   *
   * `authenticate()` returns a promise for the authenticated API, and we call
   * `subscribe()` on that promise without awaiting it first. That is promise
   * pipelining: the second call is sent immediately, carrying a reference to
   * the not-yet-existing result of the first.
   *
   * @param {{ resume?: boolean }} options
   *   `resume: false` deliberately throws the cursor away, so you can watch
   *   the gap appear that a resume token exists to prevent.
   */
  async connect({ resume = true } = {}) {
    if (this.#state !== 'offline') return;
    this.#closing = false;
    this.#setState('connecting');

    // We construct the socket ourselves rather than passing a URL string, so
    // that we hold it and can close it on demand. `newWebSocketRpcSession`
    // accepts either.
    const socket = new WebSocket(this.#url);
    this.#socket = socket;

    const api = newWebSocketRpcSession(socket, undefined);
    this.#api = api;

    // Fires for any end of session: a clean close, a dropped connection, or a
    // protocol error. There is no separate "disconnected" event to listen for.
    api.onRpcBroken((error) => this.#onBroken(error));

    const sink = new EventSink({
      onEvent: (event) => {
        this.#cursor = event.id;
        this.#onEvent(event);
      },
      onGap: (sinceId) => {
        this.#report(
          `server dropped history before #${sinceId}: too far behind to replay`,
          'warn',
        );
      },
    });

    const sinceId = resume ? this.#cursor : null;

    try {
      const authed = api.authenticate(this.#token);
      const subscription = authed.subscribe(sinceId, sink);

      // One await, so everything above cost a single round trip.
      const user = await authed.whoami();

      this.#authed = authed;
      this.#subscription = subscription;
      this.#setState('online');

      this.#report(
        sinceId === null
          ? `connected as ${user.name}; streaming from now (no resume)`
          : `connected as ${user.name}; resuming after #${sinceId}`,
        'good',
      );
    } catch (error) {
      this.#report(`connect failed: ${error.message}`, 'bad');
      this.#teardown();
      this.#setState('offline');
    }
  }

  /**
   * Prove that the capability really is gone after a drop.
   *
   * Calling a method on a stub from a dead session does not hang or silently
   * no-op; it rejects. This is the check the demo runs to make the point.
   */
  async probeStaleStub() {
    const stub = this.#authed ?? this.#staleAuthed;
    if (!stub) return 'nothing to probe -- connect first';
    try {
      const user = await stub.whoami();
      return `stub still works: ${user.name}`;
    } catch (error) {
      return `stub is broken: ${error.message}`;
    }
  }

  /** Simulate losing the network. The socket dies without a clean handshake. */
  sever() {
    if (!this.#socket) return;
    this.#report('severing the connection', 'warn');
    this.#socket.close(4000, 'simulated network loss');
  }

  /** A deliberate shutdown, so `onRpcBroken` is not treated as a failure. */
  disconnect() {
    if (!this.#socket) return;
    this.#closing = true;
    this.#report('disconnecting', 'plain');

    // Disposing the main stub closes the session, and with it the connection.
    this.#api[Symbol.dispose]();
    this.#teardown();
    this.#setState('offline');
  }

  #onBroken(error) {
    if (this.#state === 'offline') return;

    this.#teardown();
    this.#setState('offline');

    if (this.#closing) return;

    this.#report(`session broken: ${error.message}`, 'bad');
    this.#report(
      this.#cursor === null
        ? 'every stub from that session is now dead'
        : `every stub from that session is now dead; cursor held at #${this.#cursor}`,
      'plain',
    );
  }

  /**
   * Drop our references to the session.
   *
   * Deliberately does *not* touch `#cursor`. Everything the session owned is
   * gone; the resume token is ours.
   */
  #teardown() {
    this.#staleAuthed = this.#authed ?? this.#staleAuthed;
    this.#socket = null;
    this.#api = null;
    this.#authed = null;
    this.#subscription = null;
  }
}

// The RPC API for the session-recovery example, shared by the Cloudflare
// Worker (`worker.js`) and the in-page playground in the docs.
//
// `capnweb` is a bare specifier here rather than a relative path into `dist/`.
// Under Node it resolves through the repo's own workspace self-link; under
// Workers it is mapped to the workerd build by the `alias` block in
// `wrangler.jsonc`. Either way there is exactly one copy of the library, which
// matters because `RpcTarget` identity is checked at the session boundary.

import { RpcTarget } from 'capnweb';

/** The only credential this demo knows about. */
const TOKENS = new Map([
  ['demo-token', { id: 'u_1', name: 'Ada Lovelace' }],
  ['other-token', { id: 'u_2', name: 'Alan Turing' }],
]);

const HEADLINES = [
  'Order filled',
  'Deployment finished',
  'Invoice paid',
  'Container recycled',
  'Cache purged',
  'Alert cleared',
  'Backup completed',
  'Certificate renewed',
];

/**
 * How much history a client may ask for in one go. A resume token from a
 * client that has been gone for a week should not turn into an unbounded
 * replay: past this, the client is told it fell too far behind and should
 * resynchronize from scratch.
 */
export const MAX_REPLAY = 40;

/**
 * The event log.
 *
 * Deliberately created *outside* any session and passed in, because that is
 * the whole point of the example: an RPC session is per-connection memory that
 * dies with the socket, and anything that must outlive a disconnect has to
 * live somewhere else.
 *
 * Events are derived from the clock rather than stored, so this needs no
 * timer, no storage, and behaves identically whether it is running in a Worker
 * isolate or inside the docs page. Event `n` is defined to have happened at
 * `epoch + n * intervalMs`.
 */
export function createEventLog({ intervalMs = 1200, epoch = Date.now() } = {}) {
  const at = (id) => ({
    id,
    at: epoch + id * intervalMs,
    text: `${HEADLINES[id % HEADLINES.length]} #${1000 + id}`,
  });

  return {
    intervalMs,

    /** Sequence number of the most recent event that has already happened. */
    latestId() {
      return Math.max(0, Math.floor((Date.now() - epoch) / intervalMs));
    },

    /**
     * Everything after `sinceId`. Returns `{ events, truncated }` so the caller
     * can tell "nothing happened" apart from "you missed more than we keep".
     */
    since(sinceId) {
      const latest = this.latestId();
      const from = Math.max(sinceId, latest - MAX_REPLAY);
      const events = [];
      for (let id = from + 1; id <= latest; id++) events.push(at(id));
      return { events, truncated: from > sinceId };
    },

    /** Milliseconds until event `id` happens. Negative if it already has. */
    msUntil(id) {
      return epoch + id * intervalMs - Date.now();
    },
  };
}

/**
 * A live subscription.
 *
 * Returned by `AuthedApi.subscribe()` rather than being a fire-and-forget
 * call, so the client holds a capability it can dispose. Disposal happens
 * either explicitly or when the session drops -- see `[Symbol.dispose]`.
 */
class Subscription extends RpcTarget {
  #log;
  #sink;
  #lastId;
  #timer = null;
  #stopped = false;

  constructor(log, sink, sinceId) {
    super();
    this.#log = log;
    this.#sink = sink;
    this.#lastId = sinceId;
    this.#pump();
  }

  /** The highest event id delivered so far. The client's resume token. */
  get cursor() {
    return this.#lastId;
  }

  #pump() {
    if (this.#stopped) return;

    const { events, truncated } = this.#log.since(this.#lastId);
    if (truncated) {
      // Fire-and-forget, but still settled -- see the note in the loop below.
      this.#sink.onGap(this.#lastId).catch(() => {});
    }

    for (const event of events) {
      this.#lastId = event.id;

      // The client's sink is a stub, so this is an RPC back to the browser.
      // We do not need the result, but we do settle the promise: an RPC
      // promise that is never awaited and never disposed keeps an entry in the
      // session's tables alive for as long as the session lasts.
      this.#sink.onEvent(event).catch(() => {});
    }

    const wait = Math.max(20, this.#log.msUntil(this.#lastId + 1));
    this.#timer = setTimeout(() => this.#pump(), wait);
  }

  /**
   * Runs when the client disposes this stub, and also when the session dies,
   * which is what stops the timer on an abrupt disconnect.
   */
  [Symbol.dispose]() {
    this.#stopped = true;
    if (this.#timer !== null) clearTimeout(this.#timer);
    this.#sink[Symbol.dispose]();
  }
}

/**
 * The authenticated API.
 *
 * The client can only obtain one of these by calling `authenticate()` with a
 * valid token. Holding the stub *is* the authorization: there is no session
 * cookie, no bearer header on subsequent calls, and no way to reach these
 * methods without the capability. It also means the credential crosses the
 * wire exactly once per connection.
 */
class AuthedApi extends RpcTarget {
  #log;
  #user;

  constructor(log, user) {
    super();
    this.#log = log;
    this.#user = user;
  }

  whoami() {
    return { ...this.#user };
  }

  /**
   * Start streaming events after `sinceId`.
   *
   * Pass `sinceId: null` to start from the present and accept a gap; pass the
   * last id you actually processed to have the gap replayed.
   */
  subscribe(sinceId, sink) {
    const from = sinceId === null || sinceId === undefined ? this.#log.latestId() : sinceId;

    // Stubs received as parameters are disposed when the call returns, so a
    // callback that will be used later has to be duplicated first.
    return new Subscription(this.#log, sink.dup(), from);
  }
}

/** The interface a fresh connection starts with. */
export class PublicApi extends RpcTarget {
  #log;

  constructor(log) {
    super();
    this.#log = log;
  }

  /** Exchange a token for the authenticated API. */
  authenticate(token) {
    const user = TOKENS.get(token);
    if (!user) throw new Error(`Unknown API token: ${token}`);
    return new AuthedApi(this.#log, user);
  }

  /** Available without authenticating, so the page has something to show. */
  serverInfo() {
    return { intervalMs: this.#log.intervalMs, maxReplay: MAX_REPLAY };
  }
}

// Cloudflare Worker serving the session-recovery demo.
//
// Static assets are served ahead of this Worker by the `assets` config, so
// `fetch` only ever sees `/ws` and unknown paths.

import { newWorkersRpcResponse } from 'capnweb';
import { createEventLog, PublicApi } from './api.mjs';

/**
 * The event log outlives any one connection.
 *
 * Module scope means it lives as long as the isolate, which is enough for a
 * demo and is exactly the wrong answer for production: isolates come and go,
 * and two clients can easily land on two different ones. Anything that must
 * genuinely survive a disconnect belongs in storage that is addressable --
 * a Durable Object, a database, a queue. The point being made here is only
 * that it has to live *somewhere that is not the session*.
 */
const log = createEventLog();

/**
 * The main interface handed to each new connection.
 *
 * Also imported directly by the docs playground, which runs both ends of the
 * session inside one page and so never goes through `fetch` at all.
 */
export function createMain() {
  return new PublicApi(log);
}

export default {
  async fetch(request) {
    const url = new URL(request.url);

    if (url.pathname !== '/ws') {
      return new Response('Not found', { status: 404 });
    }

    if (request.headers.get('Upgrade')?.toLowerCase() !== 'websocket') {
      return new Response('This endpoint speaks WebSocket only.', { status: 426 });
    }

    return await newWorkersRpcResponse(request, createMain());
  },
};

// DOM wiring for the session-recovery demo. The RPC lives in `session.js`;
// this file only moves its output onto the page.

import { RecoveringClient } from './session.js';

const $ = (id) => document.getElementById(id);

const els = {
  state: $('state'),
  cursor: $('cursor'),
  feed: $('feed'),
  log: $('log'),
  connect: $('connect'),
  sever: $('sever'),
  probe: $('probe'),
  disconnect: $('disconnect'),
  resume: $('resume'),
};

/** Events the page has seen, so a gap in the sequence is visible. */
let seen = [];

function renderFeed() {
  els.feed.replaceChildren(
    ...seen.slice(-14).map((entry) => {
      const li = document.createElement('li');
      li.className = entry.gap ? 'event event--gap' : 'event';
      li.innerHTML = entry.gap
        ? `<span class="event__id">gap</span><span class="event__text">${entry.missed} event${
            entry.missed === 1 ? '' : 's'
          } never arrived</span>`
        : `<span class="event__id">#${entry.id}</span><span class="event__text"></span>`;
      if (!entry.gap) li.querySelector('.event__text').textContent = entry.text;
      return li;
    }),
  );
  els.feed.scrollTop = els.feed.scrollHeight;
}

function report(message, tone = 'plain') {
  const li = document.createElement('li');
  li.className = `line line--${tone}`;
  const time = new Date().toLocaleTimeString([], { hour12: false });
  li.innerHTML = `<span class="line__time">${time}</span><span class="line__text"></span>`;
  li.querySelector('.line__text').textContent = message;
  els.log.append(li);
  while (els.log.children.length > 40) els.log.firstChild.remove();
  els.log.scrollTop = els.log.scrollHeight;
}

/**
 * Record an event, inserting a marker when the sequence jumps.
 *
 * This is what makes the demo worth looking at: without a resume token the
 * numbers skip, and the marker says how many were lost.
 */
function pushEvent(event) {
  const previous = seen.filter((entry) => !entry.gap).at(-1);
  if (previous && event.id > previous.id + 1) {
    seen.push({ gap: true, missed: event.id - previous.id - 1 });
  }
  seen.push(event);
  renderFeed();
}

const client = new RecoveringClient({
  // Same-origin, so this works under `wrangler dev` and in the docs playground
  // without either of them knowing anything about the other.
  url: new URL('/ws', location.href).href.replace(/^http/, 'ws'),
  token: 'demo-token',
  report,
  onEvent: pushEvent,
  onStateChange: (state) => {
    els.state.textContent = state;
    els.state.dataset.state = state;
    els.connect.disabled = state !== 'offline';
    els.sever.disabled = state !== 'online';
    els.disconnect.disabled = state !== 'online';
    els.cursor.textContent = client.cursor === null ? '--' : `#${client.cursor}`;
  },
});

els.connect.addEventListener('click', () => client.connect({ resume: els.resume.checked }));
els.sever.addEventListener('click', () => client.sever());
els.disconnect.addEventListener('click', () => client.disconnect());
els.probe.addEventListener('click', async () => report(await client.probeStaleStub(), 'plain'));

// Keep the cursor readout live while events stream in.
setInterval(() => {
  els.cursor.textContent = client.cursor === null ? '--' : `#${client.cursor}`;
}, 250);

report('idle -- press Connect', 'plain');