boocode/apps/coder/src/services/backends/claude-session-store.ts

import type { SessionStore, SessionKey, SessionStoreEntry } from '@anthropic-ai/claude-agent-sdk';
import type { Sql } from '../../db.js';

/**
 * claude-sdk-sessionstore #9 (Part 1) — clean-room PostgresSessionStore.
 *
 * A Postgres-backed implementation of the Claude Agent SDK's `SessionStore`
 * adapter type. The SDK mirrors each transcript line (a JSON-safe POJO with a
 * `type` discriminant) to this store via `append`; on resume it calls `load`
 * to materialize the full transcript back. We treat entries as opaque blobs and
 * preserve append order via a BIGSERIAL `id` — `load` replays `ORDER BY id`.
 *
 * Storage shape: one row per entry in `claude_session_entries`, keyed by the
 * SDK's `SessionKey` (project_key, session_id, subpath). The SDK uses an
 * *undefined* subpath for the main transcript and disallows the empty string;
 * we collapse `undefined → ''` so the main transcript and subagent files share
 * one table, distinguished by the `subpath` column (`'' = main`).
 *
 * Clean-room: written against the SDK's published `SessionStore` type contract
 * and BooCode's existing SQL conventions (porsager tagged templates, `sql.json`
 * for JSONB). No SDK example/reference code was consulted.
 */
export class PostgresSessionStore implements SessionStore {
  constructor(private readonly sql: Sql) {}

  /**
   * Mirror a batch of transcript entries. No-op on an empty batch; otherwise a
   * single multi-row INSERT writes them in array order. Because `id` is a
   * monotonically-increasing BIGSERIAL, the insert order is the replay order
   * `load` reconstructs — entries within one call land in the order given.
   */
  async append(key: SessionKey, entries: SessionStoreEntry[]): Promise<void> {
    if (entries.length === 0) return;
    const subpath = key.subpath ?? '';
    const rows = entries.map((entry) => ({
      project_key: key.projectKey,
      session_id: key.sessionId,
      subpath,
      entry: this.sql.json(entry as never),
    }));
    await this.sql`
      INSERT INTO claude_session_entries ${this.sql(rows, 'project_key', 'session_id', 'subpath', 'entry')}
    `;
  }

  /**
   * Load a full transcript for resume. Returns the entries in append order, or
   * `null` for a (project_key, session_id, subpath) key that was never written.
   */
  async load(key: SessionKey): Promise<SessionStoreEntry[] | null> {
    const subpath = key.subpath ?? '';
    const rows = await this.sql<{ entry: SessionStoreEntry }[]>`
      SELECT entry
      FROM claude_session_entries
      WHERE project_key = ${key.projectKey}
        AND session_id = ${key.sessionId}
        AND subpath = ${subpath}
      ORDER BY id
    `;
    if (rows.length === 0) return null;
    return rows.map((r) => r.entry);
  }

  /**
   * List the main transcripts for a project. `mtime` is the storage write time
   * (latest `created_at` for the session) in Unix epoch milliseconds; the SDK
   * sorts the result by mtime descending.
   */
  async listSessions(projectKey: string): Promise<Array<{ sessionId: string; mtime: number }>> {
    const rows = await this.sql<{ session_id: string; mtime: string }[]>`
      SELECT session_id, extract(epoch FROM max(created_at)) * 1000 AS mtime
      FROM claude_session_entries
      WHERE project_key = ${projectKey}
        AND subpath = ''
      GROUP BY session_id
    `;
    return rows.map((r) => ({ sessionId: r.session_id, mtime: Number(r.mtime) }));
  }

  /**
   * Delete a session. With a `subpath` set, only that subpath's rows are
   * removed; with `subpath` omitted, every row for the session is removed
   * (all subpaths, including the main transcript).
   */
  async delete(key: SessionKey): Promise<void> {
    if (key.subpath !== undefined) {
      await this.sql`
        DELETE FROM claude_session_entries
        WHERE project_key = ${key.projectKey}
          AND session_id = ${key.sessionId}
          AND subpath = ${key.subpath}
      `;
      return;
    }
    await this.sql`
      DELETE FROM claude_session_entries
      WHERE project_key = ${key.projectKey}
        AND session_id = ${key.sessionId}
    `;
  }

  /**
   * List the distinct non-main subpaths under a session (e.g. subagent files).
   * Used during resume to discover and materialize subagent transcripts; the
   * main transcript (`subpath = ''`) is excluded.
   */
  async listSubkeys(key: { projectKey: string; sessionId: string }): Promise<string[]> {
    const rows = await this.sql<{ subpath: string }[]>`
      SELECT DISTINCT subpath
      FROM claude_session_entries
      WHERE project_key = ${key.projectKey}
        AND session_id = ${key.sessionId}
        AND subpath <> ''
    `;
    return rows.map((r) => r.subpath);
  }
}