@electric-sql/client/src/shape-stream-state.ts

Postgres everywhere - your data, in sync, wherever you need it.

/*
 * Shape stream state machine.
 *
 * Class hierarchy:
 *
 *   ShapeStreamState (abstract base)
 *   ├── ActiveState (abstract — shared field storage & helpers)
 *   │   ├── FetchingState (abstract — shared Initial/Syncing/StaleRetry behavior)
 *   │   │   ├── InitialState
 *   │   │   ├── SyncingState
 *   │   │   └── StaleRetryState
 *   │   ├── LiveState
 *   │   └── ReplayingState
 *   ├── PausedState   (delegates to previousState)
 *   └── ErrorState    (delegates to previousState)
 *
 * State transitions:
 *
 *   Initial ─response─► Syncing ─up-to-date─► Live
 *                            │                  │
 *                            └──stale──► StaleRetry
 *                                           │
 *                            Syncing ◄──response──┘
 *
 *   Any state ─pause─► Paused ─resume─► (previous state)
 *   Any state ─error─► Error  ─retry──► (previous state)
 *   Any state ─markMustRefetch─► Initial (offset reset)
 */
import { Offset, Schema } from './types'
import {
  OFFSET_QUERY_PARAM,
  SHAPE_HANDLE_QUERY_PARAM,
  LIVE_CACHE_BUSTER_QUERY_PARAM,
  LIVE_QUERY_PARAM,
  CACHE_BUSTER_QUERY_PARAM,
} from './constants'

export type ShapeStreamStateKind =
  | `initial`
  | `syncing`
  | `live`
  | `replaying`
  | `stale-retry`
  | `paused`
  | `error`

/**
 * Shared fields carried by all active (non-paused, non-error) states.
 */
export interface SharedStateFields {
  readonly handle?: string
  readonly offset: Offset
  readonly schema?: Schema
  readonly liveCacheBuster: string
  readonly lastSyncedAt?: number
}

type ResponseBaseInput = {
  status: number
  responseHandle: string | null
  responseOffset: Offset | null
  responseCursor: string | null
  responseSchema?: Schema
  expiredHandle?: string | null
  now: number
}

export type ResponseMetadataInput = ResponseBaseInput & {
  maxStaleCacheRetries: number
  createCacheBuster: () => string
}

export type ResponseMetadataTransition =
  | { action: `accepted`; state: ShapeStreamState }
  | { action: `ignored`; state: ShapeStreamState }
  | {
      action: `stale-retry`
      state: ShapeStreamState
      exceededMaxRetries: boolean
    }

export interface MessageBatchInput {
  hasMessages: boolean
  hasUpToDateMessage: boolean
  isSse: boolean
  upToDateOffset?: Offset
  now: number
  currentCursor: string
}

export interface MessageBatchTransition {
  state: ShapeStreamState
  suppressBatch: boolean
  becameUpToDate: boolean
}

export interface SseCloseInput {
  connectionDuration: number
  wasAborted: boolean
  minConnectionDuration: number
  maxShortConnections: number
}

export interface SseCloseTransition {
  state: ShapeStreamState
  fellBackToLongPolling: boolean
  wasShortConnection: boolean
}

export interface UrlParamsContext {
  isSnapshotRequest: boolean
  canLongPoll: boolean
}

// ---------------------------------------------------------------------------
// Abstract base — shared by ALL states (including Paused/Error)
// ---------------------------------------------------------------------------

/**
 * Abstract base class for all shape stream states.
 *
 * Each concrete state carries only its relevant fields — there is no shared
 * flat context bag. Transitions create new immutable state objects.
 *
 * `isUpToDate` returns true for LiveState and delegating states wrapping LiveState.
 */
export abstract class ShapeStreamState {
  abstract readonly kind: ShapeStreamStateKind

  // --- Shared field getters (all states expose these) ---
  abstract get handle(): string | undefined
  abstract get offset(): Offset
  abstract get schema(): Schema | undefined
  abstract get liveCacheBuster(): string
  abstract get lastSyncedAt(): number | undefined

  // --- Derived booleans ---
  get isUpToDate(): boolean {
    return false
  }

  // --- Per-state field defaults ---
  get staleCacheBuster(): string | undefined {
    return undefined
  }
  get staleCacheRetryCount(): number {
    return 0
  }
  get sseFallbackToLongPolling(): boolean {
    return false
  }
  get consecutiveShortSseConnections(): number {
    return 0
  }
  get replayCursor(): string | undefined {
    return undefined
  }

  // --- Default no-op methods ---

  canEnterReplayMode(): boolean {
    return false
  }

  enterReplayMode(_cursor: string): ShapeStreamState {
    return this
  }

  shouldUseSse(_opts: {
    liveSseEnabled: boolean
    isRefreshing: boolean
    resumingFromPause: boolean
  }): boolean {
    return false
  }

  handleSseConnectionClosed(_input: SseCloseInput): SseCloseTransition {
    return {
      state: this,
      fellBackToLongPolling: false,
      wasShortConnection: false,
    }
  }

  // --- URL param application ---

  /** Adds state-specific query parameters to the fetch URL. */
  applyUrlParams(_url: URL, _context: UrlParamsContext): void {}

  // --- Default response/message handlers (Paused/Error never receive these) ---

  handleResponseMetadata(
    _input: ResponseMetadataInput
  ): ResponseMetadataTransition {
    return { action: `ignored`, state: this }
  }

  handleMessageBatch(_input: MessageBatchInput): MessageBatchTransition {
    return { state: this, suppressBatch: false, becameUpToDate: false }
  }

  // --- Universal transitions ---

  /** Returns a new state identical to this one but with the handle changed. */
  abstract withHandle(handle: string): ShapeStreamState

  pause(): PausedState {
    return new PausedState(this)
  }

  toErrorState(error: Error): ErrorState {
    return new ErrorState(this, error)
  }

  markMustRefetch(handle?: string): InitialState {
    return new InitialState({
      handle,
      offset: `-1`,
      liveCacheBuster: ``,
      lastSyncedAt: this.lastSyncedAt,
      schema: undefined,
    })
  }
}

// ---------------------------------------------------------------------------
// ActiveState — intermediate base for all non-paused, non-error states
// ---------------------------------------------------------------------------

/**
 * Holds shared field storage and provides helpers for response/message
 * handling. All five active states extend this (via FetchingState or directly).
 */
abstract class ActiveState extends ShapeStreamState {
  readonly #shared: SharedStateFields

  constructor(shared: SharedStateFields) {
    super()
    this.#shared = shared
  }

  get handle() {
    return this.#shared.handle
  }
  get offset() {
    return this.#shared.offset
  }
  get schema() {
    return this.#shared.schema
  }
  get liveCacheBuster() {
    return this.#shared.liveCacheBuster
  }
  get lastSyncedAt() {
    return this.#shared.lastSyncedAt
  }

  /** Expose shared fields to subclasses for spreading into new instances. */
  protected get currentFields(): SharedStateFields {
    return this.#shared
  }

  // --- URL param application ---

  applyUrlParams(url: URL, _context: UrlParamsContext): void {
    url.searchParams.set(OFFSET_QUERY_PARAM, this.#shared.offset)
    if (this.#shared.handle) {
      url.searchParams.set(SHAPE_HANDLE_QUERY_PARAM, this.#shared.handle)
    }
  }

  // --- Helpers for subclass handleResponseMetadata implementations ---

  /** Extracts updated SharedStateFields from response headers. */
  protected parseResponseFields(
    input: ResponseMetadataInput
  ): SharedStateFields {
    const responseHandle = input.responseHandle
    const handle =
      responseHandle && responseHandle !== input.expiredHandle
        ? responseHandle
        : this.#shared.handle
    const offset = input.responseOffset ?? this.#shared.offset
    const liveCacheBuster = input.responseCursor ?? this.#shared.liveCacheBuster
    const schema = this.#shared.schema ?? input.responseSchema
    const lastSyncedAt =
      input.status === 204 ? input.now : this.#shared.lastSyncedAt

    return { handle, offset, schema, liveCacheBuster, lastSyncedAt }
  }

  /**
   * Stale detection. Returns a transition if the response is stale,
   * or null if it is not stale and the caller should proceed normally.
   */
  protected checkStaleResponse(
    input: ResponseMetadataInput
  ): ResponseMetadataTransition | null {
    const responseHandle = input.responseHandle
    const expiredHandle = input.expiredHandle

    if (!responseHandle || responseHandle !== expiredHandle) {
      return null // not stale
    }

    // Stale response detected — always enter stale-retry to get a cache buster.
    // Without a cache buster, the CDN will keep serving the same stale response
    // and the client loops infinitely (the URL never changes).
    // currentFields preserves the valid local handle when we already have one.
    const retryCount = this.staleCacheRetryCount + 1
    return {
      action: `stale-retry`,
      state: new StaleRetryState({
        ...this.currentFields,
        staleCacheBuster: input.createCacheBuster(),
        staleCacheRetryCount: retryCount,
      }),
      exceededMaxRetries: retryCount > input.maxStaleCacheRetries,
    }
  }

  // --- handleMessageBatch: template method with onUpToDate override point ---

  handleMessageBatch(input: MessageBatchInput): MessageBatchTransition {
    if (!input.hasMessages || !input.hasUpToDateMessage) {
      return { state: this, suppressBatch: false, becameUpToDate: false }
    }

    // Has up-to-date message — compute shared fields for the transition
    let offset = this.#shared.offset
    if (input.isSse && input.upToDateOffset) {
      offset = input.upToDateOffset
    }

    const shared: SharedStateFields = {
      handle: this.#shared.handle,
      offset,
      schema: this.#shared.schema,
      liveCacheBuster: this.#shared.liveCacheBuster,
      lastSyncedAt: input.now,
    }

    return this.onUpToDate(shared, input)
  }

  /** Override point for up-to-date handling. Default → LiveState. */
  protected onUpToDate(
    shared: SharedStateFields,
    _input: MessageBatchInput
  ): MessageBatchTransition {
    return {
      state: new LiveState(shared),
      suppressBatch: false,
      becameUpToDate: true,
    }
  }
}

// ---------------------------------------------------------------------------
// FetchingState — Common behavior for Initial/Syncing/StaleRetry
// ---------------------------------------------------------------------------

/**
 * Captures shared behavior of InitialState, SyncingState, StaleRetryState:
 * - handleResponseMetadata: stale check → parse fields → new SyncingState (or LiveState for 204)
 * - enterReplayMode(cursor) → new ReplayingState
 * - canEnterReplayMode(): boolean — returns true (StaleRetryState overrides to return false,
 *   because entering replay would lose the stale-retry count; see C1 in SPEC.md)
 */
abstract class FetchingState extends ActiveState {
  handleResponseMetadata(
    input: ResponseMetadataInput
  ): ResponseMetadataTransition {
    const staleResult = this.checkStaleResponse(input)
    if (staleResult) return staleResult

    const shared = this.parseResponseFields(input)

    // NOTE: 204s are deprecated, the Electric server should not send these
    // in latest versions but this is here for backwards compatibility.
    // A 204 means "no content, you're caught up" — transition to live.
    // Skip SSE detection: a 204 gives no indication SSE will work, and
    // the 3-attempt fallback cycle adds unnecessary latency.
    if (input.status === 204) {
      return {
        action: `accepted`,
        state: new LiveState(shared, { sseFallbackToLongPolling: true }),
      }
    }

    return { action: `accepted`, state: new SyncingState(shared) }
  }

  canEnterReplayMode(): boolean {
    return true
  }

  enterReplayMode(cursor: string): ReplayingState {
    return new ReplayingState({
      ...this.currentFields,
      replayCursor: cursor,
    })
  }
}

// ---------------------------------------------------------------------------
// Concrete states
// ---------------------------------------------------------------------------

export class InitialState extends FetchingState {
  readonly kind = `initial` as const

  constructor(shared: SharedStateFields) {
    super(shared)
  }

  withHandle(handle: string): InitialState {
    return new InitialState({ ...this.currentFields, handle })
  }
}

export class SyncingState extends FetchingState {
  readonly kind = `syncing` as const

  constructor(shared: SharedStateFields) {
    super(shared)
  }

  withHandle(handle: string): SyncingState {
    return new SyncingState({ ...this.currentFields, handle })
  }
}

export class StaleRetryState extends FetchingState {
  readonly kind = `stale-retry` as const
  readonly #staleCacheBuster: string
  readonly #staleCacheRetryCount: number

  constructor(
    fields: SharedStateFields & {
      staleCacheBuster: string
      staleCacheRetryCount: number
    }
  ) {
    const { staleCacheBuster, staleCacheRetryCount, ...shared } = fields
    super(shared)
    this.#staleCacheBuster = staleCacheBuster
    this.#staleCacheRetryCount = staleCacheRetryCount
  }

  get staleCacheBuster() {
    return this.#staleCacheBuster
  }
  get staleCacheRetryCount() {
    return this.#staleCacheRetryCount
  }

  // StaleRetryState must not enter replay mode — it would lose the retry count
  canEnterReplayMode(): boolean {
    return false
  }

  withHandle(handle: string): StaleRetryState {
    return new StaleRetryState({
      ...this.currentFields,
      handle,
      staleCacheBuster: this.#staleCacheBuster,
      staleCacheRetryCount: this.#staleCacheRetryCount,
    })
  }

  applyUrlParams(url: URL, context: UrlParamsContext): void {
    super.applyUrlParams(url, context)
    url.searchParams.set(CACHE_BUSTER_QUERY_PARAM, this.#staleCacheBuster)
  }
}

export class LiveState extends ActiveState {
  readonly kind = `live` as const
  readonly #consecutiveShortSseConnections: number
  readonly #sseFallbackToLongPolling: boolean

  constructor(
    shared: SharedStateFields,
    sseState?: {
      consecutiveShortSseConnections?: number
      sseFallbackToLongPolling?: boolean
    }
  ) {
    super(shared)
    this.#consecutiveShortSseConnections =
      sseState?.consecutiveShortSseConnections ?? 0
    this.#sseFallbackToLongPolling = sseState?.sseFallbackToLongPolling ?? false
  }

  get isUpToDate(): boolean {
    return true
  }

  get consecutiveShortSseConnections(): number {
    return this.#consecutiveShortSseConnections
  }

  get sseFallbackToLongPolling(): boolean {
    return this.#sseFallbackToLongPolling
  }

  withHandle(handle: string): LiveState {
    return new LiveState({ ...this.currentFields, handle }, this.sseState)
  }

  applyUrlParams(url: URL, context: UrlParamsContext): void {
    super.applyUrlParams(url, context)
    // Snapshot requests (with subsetParams) should never use live polling
    if (!context.isSnapshotRequest) {
      url.searchParams.set(LIVE_CACHE_BUSTER_QUERY_PARAM, this.liveCacheBuster)
      if (context.canLongPoll) {
        url.searchParams.set(LIVE_QUERY_PARAM, `true`)
      }
    }
  }

  private get sseState() {
    return {
      consecutiveShortSseConnections: this.#consecutiveShortSseConnections,
      sseFallbackToLongPolling: this.#sseFallbackToLongPolling,
    }
  }

  handleResponseMetadata(
    input: ResponseMetadataInput
  ): ResponseMetadataTransition {
    const staleResult = this.checkStaleResponse(input)
    if (staleResult) return staleResult

    const shared = this.parseResponseFields(input)
    return {
      action: `accepted`,
      state: new LiveState(shared, this.sseState),
    }
  }

  protected onUpToDate(
    shared: SharedStateFields,
    _input: MessageBatchInput
  ): MessageBatchTransition {
    return {
      state: new LiveState(shared, this.sseState),
      suppressBatch: false,
      becameUpToDate: true,
    }
  }

  shouldUseSse(opts: {
    liveSseEnabled: boolean
    isRefreshing: boolean
    resumingFromPause: boolean
  }): boolean {
    return (
      opts.liveSseEnabled &&
      !opts.isRefreshing &&
      !opts.resumingFromPause &&
      !this.#sseFallbackToLongPolling
    )
  }

  handleSseConnectionClosed(input: SseCloseInput): SseCloseTransition {
    let nextConsecutiveShort = this.#consecutiveShortSseConnections
    let nextFallback = this.#sseFallbackToLongPolling
    let fellBackToLongPolling = false
    let wasShortConnection = false

    if (
      input.connectionDuration < input.minConnectionDuration &&
      !input.wasAborted
    ) {
      wasShortConnection = true
      nextConsecutiveShort = nextConsecutiveShort + 1

      if (nextConsecutiveShort >= input.maxShortConnections) {
        nextFallback = true
        fellBackToLongPolling = true
      }
    } else if (input.connectionDuration >= input.minConnectionDuration) {
      nextConsecutiveShort = 0
    }

    return {
      state: new LiveState(this.currentFields, {
        consecutiveShortSseConnections: nextConsecutiveShort,
        sseFallbackToLongPolling: nextFallback,
      }),
      fellBackToLongPolling,
      wasShortConnection,
    }
  }
}

export class ReplayingState extends ActiveState {
  readonly kind = `replaying` as const
  readonly #replayCursor: string

  constructor(fields: SharedStateFields & { replayCursor: string }) {
    const { replayCursor, ...shared } = fields
    super(shared)
    this.#replayCursor = replayCursor
  }

  get replayCursor() {
    return this.#replayCursor
  }

  withHandle(handle: string): ReplayingState {
    return new ReplayingState({
      ...this.currentFields,
      handle,
      replayCursor: this.#replayCursor,
    })
  }

  handleResponseMetadata(
    input: ResponseMetadataInput
  ): ResponseMetadataTransition {
    const staleResult = this.checkStaleResponse(input)
    if (staleResult) return staleResult

    const shared = this.parseResponseFields(input)
    return {
      action: `accepted`,
      state: new ReplayingState({
        ...shared,
        replayCursor: this.#replayCursor,
      }),
    }
  }

  protected onUpToDate(
    shared: SharedStateFields,
    input: MessageBatchInput
  ): MessageBatchTransition {
    // Suppress replayed cache data when cursor has not moved since
    // the previous session (non-SSE only).
    const suppressBatch =
      !input.isSse && this.#replayCursor === input.currentCursor
    return {
      state: new LiveState(shared),
      suppressBatch,
      becameUpToDate: true,
    }
  }
}

// ---------------------------------------------------------------------------
// Delegating states (Paused / Error)
// ---------------------------------------------------------------------------

// Union of all non-delegating states — used to type previousState fields so
// that PausedState.previousState is never another PausedState, and
// ErrorState.previousState is never another ErrorState.
export type ShapeStreamActiveState =
  | InitialState
  | SyncingState
  | LiveState
  | ReplayingState
  | StaleRetryState

export class PausedState extends ShapeStreamState {
  readonly kind = `paused` as const
  readonly previousState: ShapeStreamActiveState | ErrorState

  constructor(previousState: ShapeStreamState) {
    super()
    this.previousState = (
      previousState instanceof PausedState
        ? previousState.previousState
        : previousState
    ) as ShapeStreamActiveState | ErrorState
  }

  get handle(): string | undefined {
    return this.previousState.handle
  }
  get offset(): Offset {
    return this.previousState.offset
  }
  get schema(): Schema | undefined {
    return this.previousState.schema
  }
  get liveCacheBuster(): string {
    return this.previousState.liveCacheBuster
  }
  get lastSyncedAt(): number | undefined {
    return this.previousState.lastSyncedAt
  }
  get isUpToDate(): boolean {
    return this.previousState.isUpToDate
  }
  get staleCacheBuster(): string | undefined {
    return this.previousState.staleCacheBuster
  }
  get staleCacheRetryCount(): number {
    return this.previousState.staleCacheRetryCount
  }
  get sseFallbackToLongPolling(): boolean {
    return this.previousState.sseFallbackToLongPolling
  }
  get consecutiveShortSseConnections(): number {
    return this.previousState.consecutiveShortSseConnections
  }
  get replayCursor(): string | undefined {
    return this.previousState.replayCursor
  }

  handleResponseMetadata(
    input: ResponseMetadataInput
  ): ResponseMetadataTransition {
    const transition = this.previousState.handleResponseMetadata(input)
    if (transition.action === `accepted`) {
      return { action: `accepted`, state: new PausedState(transition.state) }
    }
    if (transition.action === `ignored`) {
      return { action: `ignored`, state: this }
    }
    if (transition.action === `stale-retry`) {
      return {
        action: `stale-retry`,
        state: new PausedState(transition.state),
        exceededMaxRetries: transition.exceededMaxRetries,
      }
    }
    const _exhaustive: never = transition
    throw new Error(
      `PausedState.handleResponseMetadata: unhandled transition action "${(_exhaustive as ResponseMetadataTransition).action}"`
    )
  }

  withHandle(handle: string): PausedState {
    return new PausedState(this.previousState.withHandle(handle))
  }

  applyUrlParams(url: URL, context: UrlParamsContext): void {
    this.previousState.applyUrlParams(url, context)
  }

  pause(): PausedState {
    return this
  }

  resume(): ShapeStreamState {
    return this.previousState
  }
}

export class ErrorState extends ShapeStreamState {
  readonly kind = `error` as const
  readonly previousState: ShapeStreamActiveState | PausedState
  readonly error: Error

  constructor(previousState: ShapeStreamState, error: Error) {
    super()
    this.previousState = (
      previousState instanceof ErrorState
        ? previousState.previousState
        : previousState
    ) as ShapeStreamActiveState | PausedState
    this.error = error
  }

  get handle(): string | undefined {
    return this.previousState.handle
  }
  get offset(): Offset {
    return this.previousState.offset
  }
  get schema(): Schema | undefined {
    return this.previousState.schema
  }
  get liveCacheBuster(): string {
    return this.previousState.liveCacheBuster
  }
  get lastSyncedAt(): number | undefined {
    return this.previousState.lastSyncedAt
  }
  get isUpToDate(): boolean {
    return this.previousState.isUpToDate
  }
  get staleCacheBuster(): string | undefined {
    return this.previousState.staleCacheBuster
  }
  get staleCacheRetryCount(): number {
    return this.previousState.staleCacheRetryCount
  }
  get sseFallbackToLongPolling(): boolean {
    return this.previousState.sseFallbackToLongPolling
  }
  get consecutiveShortSseConnections(): number {
    return this.previousState.consecutiveShortSseConnections
  }
  get replayCursor(): string | undefined {
    return this.previousState.replayCursor
  }

  withHandle(handle: string): ErrorState {
    return new ErrorState(this.previousState.withHandle(handle), this.error)
  }

  applyUrlParams(url: URL, context: UrlParamsContext): void {
    this.previousState.applyUrlParams(url, context)
  }

  retry(): ShapeStreamState {
    return this.previousState
  }

  reset(handle?: string): InitialState {
    return this.previousState.markMustRefetch(handle)
  }
}

// ---------------------------------------------------------------------------
// Type alias & factory
// ---------------------------------------------------------------------------

export function createInitialState(opts: {
  offset: Offset
  handle?: string
}): InitialState {
  return new InitialState({
    handle: opts.handle,
    offset: opts.offset,
    liveCacheBuster: ``,
    lastSyncedAt: undefined,
    schema: undefined,
  })
}