boocode/apps/coder/src/services/__tests__/checkpoints.test.ts

import { describe, it, expect, beforeAll, afterAll } from 'vitest';
import { readFileSync } from 'node:fs';
import { rm, mkdir } from 'node:fs/promises';
import { resolve } from 'node:path';
import postgres from 'postgres';
import {
  buildShadowCommitCommand,
  createCheckpoint,
  restoreCheckpoint,
  CheckpointNotFoundError,
} from '../checkpoints.js';
import { ensureSessionWorktree } from '../worktrees.js';
import { hostExec } from '../host-exec.js';

/**
 * write-edit-robustness #4 — worktree checkpoint tests.
 *
 * Pure-helper coverage (no DB / no host) for the shadow-commit command builder,
 * plus a DB+git integration block (DB-opt-in via DATABASE_URL, skips cleanly
 * otherwise; mirrors reconnect_integration.test.ts) that exercises the real
 * create → restore round trip against a worktree on the host fs.
 */

describe('buildShadowCommitCommand (pure)', () => {
  it('parks the commit under refs/boocode/checkpoints/<id> and prints only the SHA', () => {
    const cmd = buildShadowCommitCommand('/tmp/booworktrees/sess-abc', 'cp-id-123');
    // Uses a temp index so the real working tree/index is untouched.
    expect(cmd).toContain('TMP=$(mktemp)');
    expect(cmd).toContain('GIT_INDEX_FILE="$TMP" git read-tree HEAD');
    expect(cmd).toContain('GIT_INDEX_FILE="$TMP" git add -A');
    expect(cmd).toContain('git write-tree');
    expect(cmd).toContain("git commit-tree \"$TREE\" -p HEAD -m \"boocode checkpoint\"");
    // Ref name matches the row id, and stdout is ONLY the SHA (printf, no newline).
    expect(cmd).toContain("update-ref 'refs/boocode/checkpoints/cp-id-123'");
    expect(cmd).toContain("printf '%s' \"$SHA\"");
    expect(cmd).not.toContain('echo "$SHA"');
  });

  it('shell-escapes the worktree path and the id', () => {
    const cmd = buildShadowCommitCommand("/tmp/it's a path", "id'; rm -rf /");
    // Single quotes inside the path/id are escaped via the '\'' wrapping idiom — no
    // bare interpolation that could break out of the quoting.
    expect(cmd).toContain("cd '/tmp/it'\\''s a path'");
    expect(cmd).toContain("refs/boocode/checkpoints/id'\\''; rm -rf /");
  });
});

describe.runIf(!!process.env.DATABASE_URL)('checkpoint create + restore (DB + git)', () => {
  let sql: ReturnType<typeof postgres>;
  const stamp = Date.now();
  const projectDir = `/tmp/boocode-checkpoint-proj-${stamp}`;
  let projectId: string;
  let sessionId: string;
  let chatId: string;
  let worktreePath: string;

  beforeAll(async () => {
    sql = postgres(process.env.DATABASE_URL!, { max: 3 });

    // Server schema first (FK targets), then coder schema (worktrees + checkpoints).
    const serverSchema = resolve(__dirname, '../../../../server/src/schema.sql');
    const coderSchema = resolve(__dirname, '../../schema.sql');
    await sql.unsafe(readFileSync(serverSchema, 'utf8'));
    await sql.unsafe(readFileSync(coderSchema, 'utf8'));

    await mkdir(projectDir, { recursive: true });
    await hostExec(
      `cd ${projectDir} && git init -q && git config user.email t@t && git config user.name t ` +
        `&& echo hello > README.md && git add -A && git commit -qm init`,
      { timeoutMs: 20_000 },
    );

    const [project] = await sql<{ id: string }[]>`
      INSERT INTO projects (name, path, status) VALUES ('checkpoint-test', ${projectDir}, 'open') RETURNING id
    `;
    projectId = project!.id;
    const [session] = await sql<{ id: string }[]>`
      INSERT INTO sessions (project_id, name, model, status)
      VALUES (${projectId}, 'cp', 'm', 'open') RETURNING id
    `;
    sessionId = session!.id;
    const [chat] = await sql<{ id: string }[]>`
      INSERT INTO chats (session_id, name, status) VALUES (${sessionId}, 'tab', 'open') RETURNING id
    `;
    chatId = chat!.id;

    const wt = await ensureSessionWorktree(sql, projectDir, sessionId);
    worktreePath = wt.worktreePath;
  });

  afterAll(async () => {
    if (sql) {
      const rows = await sql<{ path: string }[]>`SELECT path FROM worktrees WHERE session_id = ${sessionId}`.catch(() => []);
      for (const r of rows) {
        await hostExec(`git -C ${projectDir} worktree remove ${r.path} --force`, { timeoutMs: 10_000 }).catch(() => {});
      }
      await sql`DELETE FROM checkpoints WHERE chat_id = ${chatId}`.catch(() => {});
      await sql`DELETE FROM agent_sessions WHERE chat_id = ${chatId}`.catch(() => {});
      await sql`DELETE FROM worktrees WHERE session_id = ${sessionId}`.catch(() => {});
      await sql`DELETE FROM chats WHERE id = ${chatId}`.catch(() => {});
      await sql`DELETE FROM sessions WHERE id = ${sessionId}`.catch(() => {});
      await sql`DELETE FROM projects WHERE id = ${projectId}`.catch(() => {});
      await sql.end({ timeout: 5 });
    }
    await rm(projectDir, { recursive: true, force: true });
  });

  it('createCheckpoint inserts a row + a private ref capturing tracked + untracked', async () => {
    const [wt] = await sql<{ id: string }[]>`SELECT id FROM worktrees WHERE session_id = ${sessionId} AND status = 'active'`;
    const worktreeId = wt!.id;

    // Pre-turn untracked + tracked-edit state the agent will start from.
    await hostExec(`cd ${worktreePath} && echo edited >> README.md && echo new > extra.txt`, { timeoutMs: 10_000 });

    const [assistantMsg] = await sql<{ id: string }[]>`
      INSERT INTO messages (session_id, chat_id, role, content, status)
      VALUES (${sessionId}, ${chatId}, 'assistant', '', 'streaming') RETURNING id
    `;
    const messageId = assistantMsg!.id;

    const cp = await createCheckpoint(sql, {
      chatId,
      sessionId,
      worktreeId,
      worktreePath,
      messageId,
    });
    expect(cp).not.toBeNull();
    expect(cp!.commit_sha).toMatch(/^[0-9a-f]{40}$/);

    const [row] = await sql<{ commit_sha: string; worktree_id: string; message_id: string }[]>`
      SELECT commit_sha, worktree_id, message_id FROM checkpoints WHERE id = ${cp!.id}
    `;
    expect(row!.commit_sha).toBe(cp!.commit_sha);
    expect(row!.worktree_id).toBe(worktreeId);
    expect(row!.message_id).toBe(messageId);

    // The ref exists and the captured tree carries the untracked file (proves the
    // temp-index `git add -A` snapshotted untracked content).
    const refLs = await hostExec(
      `git -C ${worktreePath} ls-tree -r --name-only ${cp!.commit_sha}`,
      { timeoutMs: 10_000 },
    );
    expect(refLs.exitCode).toBe(0);
    expect(refLs.stdout).toContain('extra.txt');

    // The shadow commit did NOT disturb the real working tree: extra.txt is still
    // present + still untracked (status shows it).
    const status = await hostExec(`git -C ${worktreePath} status --porcelain`, { timeoutMs: 10_000 });
    expect(status.stdout).toContain('extra.txt');
  });

  it('restoreCheckpoint resets the worktree, trims the transcript, and drops later checkpoints', async () => {
    // Clean slate for this test: reset the worktree to HEAD, clear prior rows.
    await hostExec(`git -C ${worktreePath} reset --hard HEAD && git -C ${worktreePath} clean -fd`, { timeoutMs: 10_000 });
    await sql`DELETE FROM checkpoints WHERE chat_id = ${chatId}`;
    await sql`DELETE FROM messages WHERE chat_id = ${chatId}`;

    const [wt] = await sql<{ id: string }[]>`SELECT id FROM worktrees WHERE session_id = ${sessionId} AND status = 'active'`;
    const worktreeId = wt!.id;

    // Turn 1: a user msg, then the assistant turn the checkpoint anchors. The
    // worktree is pristine (matches HEAD) when this checkpoint is captured.
    await sql`INSERT INTO messages (session_id, chat_id, role, content, status) VALUES (${sessionId}, ${chatId}, 'user', 'do it', 'complete')`;
    const [a1] = await sql<{ id: string }[]>`
      INSERT INTO messages (session_id, chat_id, role, content, status)
      VALUES (${sessionId}, ${chatId}, 'assistant', 'turn 1', 'complete') RETURNING id
    `;
    const cp1 = await createCheckpoint(sql, { chatId, sessionId, worktreeId, worktreePath, messageId: a1!.id });
    expect(cp1).not.toBeNull();

    // The agent (turn 1) writes a file into the worktree.
    await hostExec(`cd ${worktreePath} && echo agent-wrote > agent.txt`, { timeoutMs: 10_000 });

    // Turn 2: another user msg + assistant turn, AND a second (later) checkpoint.
    await sql`INSERT INTO messages (session_id, chat_id, role, content, status) VALUES (${sessionId}, ${chatId}, 'user', 'more', 'complete')`;
    const [a2] = await sql<{ id: string }[]>`
      INSERT INTO messages (session_id, chat_id, role, content, status)
      VALUES (${sessionId}, ${chatId}, 'assistant', 'turn 2', 'complete') RETURNING id
    `;
    const cp2 = await createCheckpoint(sql, { chatId, sessionId, worktreeId, worktreePath, messageId: a2!.id });
    expect(cp2).not.toBeNull();

    // An agent_sessions row that restore should mark 'crashed'.
    await sql`
      INSERT INTO agent_sessions (chat_id, session_id, worktree_id, agent, backend, agent_session_id, status, last_active_at)
      VALUES (${chatId}, ${sessionId}, ${worktreeId}, 'goose', 'acp_warm', 'sess-1', 'active', clock_timestamp())
      ON CONFLICT (chat_id, agent) DO UPDATE SET status = 'active'
    `;

    const before = await sql<{ id: string }[]>`SELECT id FROM messages WHERE chat_id = ${chatId} ORDER BY created_at`;
    expect(before.length).toBe(4); // user, a1, user, a2

    // Restore to cp1 (before turn 1's assistant message).
    const result = await restoreCheckpoint(sql, cp1!.id, { sessionId });
    expect(result.checkpoint_id).toBe(cp1!.id);
    expect(result.worktree_reset).toBe(true);
    expect(result.backend_reset).toBe(true);
    // a1, user(turn2), a2 deleted (created_at >= a1) → 3 trimmed.
    expect(result.messages_deleted).toBe(3);

    // Transcript trimmed to just the first user message.
    const after = await sql<{ role: string; content: string }[]>`SELECT role, content FROM messages WHERE chat_id = ${chatId} ORDER BY created_at`;
    expect(after.length).toBe(1);
    expect(after[0]!.role).toBe('user');

    // Worktree reset: the agent's file is gone (it was written after cp1).
    const ls = await hostExec(`ls ${worktreePath}/agent.txt`, { timeoutMs: 10_000 });
    expect(ls.exitCode).not.toBe(0);

    // The agent_sessions row was reset to 'crashed'.
    const [as] = await sql<{ status: string }[]>`SELECT status FROM agent_sessions WHERE chat_id = ${chatId} AND agent = 'goose'`;
    expect(as!.status).toBe('crashed');

    // cp1 survives (re-restorable); cp2 (later) was dropped.
    const cps = await sql<{ id: string }[]>`SELECT id FROM checkpoints WHERE chat_id = ${chatId}`;
    expect(cps.map((c) => c.id)).toEqual([cp1!.id]);
  });

  it('restoreCheckpoint throws CheckpointNotFoundError for an unknown id', async () => {
    await expect(
      restoreCheckpoint(sql, '00000000-0000-0000-0000-000000000000', { sessionId }),
    ).rejects.toBeInstanceOf(CheckpointNotFoundError);
  });

  it('restoreCheckpoint throws when the checkpoint is not in the requested session', async () => {
    // A checkpoint whose session_id differs from the route's sessionId.
    const [wt] = await sql<{ id: string }[]>`SELECT id FROM worktrees WHERE session_id = ${sessionId} AND status = 'active'`;
    const cp = await createCheckpoint(sql, { chatId, sessionId, worktreeId: wt!.id, worktreePath, messageId: null });
    expect(cp).not.toBeNull();
    await expect(
      restoreCheckpoint(sql, cp!.id, { sessionId: '11111111-1111-1111-1111-111111111111' }),
    ).rejects.toBeInstanceOf(CheckpointNotFoundError);
    await sql`DELETE FROM checkpoints WHERE id = ${cp!.id}`;
  });

  it('restoreCheckpoint denies a NULL-session_id checkpoint from another session (no fail-open IDOR)', async () => {
    // Regression for the fail-open authorization bug: a checkpoint row whose
    // denormalized session_id is NULL must STILL be scoped via its chat's owning
    // session (chats.session_id), not skipped. The old guard `cp.session_id &&
    // cp.session_id !== sessionId` fell through on NULL → cross-session restore.
    const [row] = await sql<{ id: string }[]>`
      INSERT INTO checkpoints (chat_id, session_id, message_id, commit_sha)
      VALUES (${chatId}, NULL, NULL, 'deadbeef')
      RETURNING id
    `;
    await expect(
      restoreCheckpoint(sql, row!.id, { sessionId: '22222222-2222-2222-2222-222222222222' }),
    ).rejects.toBeInstanceOf(CheckpointNotFoundError);
    await sql`DELETE FROM checkpoints WHERE id = ${row!.id}`;
  });
});