@blackbelt-technology/pi-agent-dashboard 0.2.0

package/packages/server/src/__tests__/git-operations.test.ts

@@ -0,0 +1,251 @@
+import { describe, it, expect, beforeEach, afterEach } from "vitest";
+import { execSync } from "node:child_process";
+import { mkdtempSync, rmSync, writeFileSync, mkdirSync } from "node:fs";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import {
+  isGitRepo,
+  getDirtyFiles,
+  listBranches,
+  checkoutBranch,
+  gitInit,
+  stashPop,
+} from "../git-operations.js";
+
+function git(cmd: string, cwd: string) {
+  execSync(`git ${cmd}`, { cwd, stdio: "pipe" });
+}
+
+function makeRepo(): string {
+  const dir = mkdtempSync(join(tmpdir(), "git-ops-test-"));
+  git("init", dir);
+  git("config user.email test@test.com", dir);
+  git("config user.name Test", dir);
+  // Initial commit so we have a branch
+  writeFileSync(join(dir, "README.md"), "init");
+  git("add .", dir);
+  git("commit -m init", dir);
+  return dir;
+}
+
+describe("git-operations", () => {
+  let repo: string;
+
+  beforeEach(() => {
+    repo = makeRepo();
+  });
+
+  afterEach(() => {
+    rmSync(repo, { recursive: true, force: true });
+  });
+
+  describe("isGitRepo", () => {
+    it("returns true for a git repo", () => {
+      expect(isGitRepo(repo)).toBe(true);
+    });
+
+    it("returns false for a non-git directory", () => {
+      const plain = mkdtempSync(join(tmpdir(), "no-git-"));
+      try {
+        expect(isGitRepo(plain)).toBe(false);
+      } finally {
+        rmSync(plain, { recursive: true, force: true });
+      }
+    });
+  });
+
+  describe("getDirtyFiles", () => {
+    it("returns empty for a clean repo", () => {
+      expect(getDirtyFiles(repo)).toEqual([]);
+    });
+
+    it("returns modified files", () => {
+      writeFileSync(join(repo, "README.md"), "changed");
+      const files = getDirtyFiles(repo);
+      expect(files).toContain("README.md");
+    });
+
+    it("returns untracked files", () => {
+      writeFileSync(join(repo, "new.txt"), "hello");
+      const files = getDirtyFiles(repo);
+      expect(files).toContain("new.txt");
+    });
+  });
+
+  describe("listBranches", () => {
+    it("lists the current branch", () => {
+      const info = listBranches(repo);
+      expect(info.detached).toBe(false);
+      expect(info.branches.length).toBeGreaterThanOrEqual(1);
+      const current = info.branches.find((b) => b.isCurrent);
+      expect(current).toBeDefined();
+      expect(current!.name).toBe(info.current);
+    });
+
+    it("lists multiple local branches", () => {
+      git("checkout -b feature-a", repo);
+      git("checkout -b feature-b", repo);
+      const info = listBranches(repo);
+      const names = info.branches.map((b) => b.name);
+      expect(names).toContain("feature-a");
+      expect(names).toContain("feature-b");
+    });
+
+    it("handles empty repo (no commits)", () => {
+      const emptyRepo = mkdtempSync(join(tmpdir(), "git-empty-"));
+      try {
+        git("init", emptyRepo);
+        const info = listBranches(emptyRepo);
+        expect(info.detached).toBe(false);
+        expect(info.branches).toEqual([]);
+        expect(info.current).toBeTruthy(); // default branch name
+      } finally {
+        rmSync(emptyRepo, { recursive: true, force: true });
+      }
+    });
+
+    it("detects detached HEAD", () => {
+      const sha = execSync("git rev-parse HEAD", { cwd: repo, encoding: "utf-8" }).trim();
+      git(`checkout ${sha}`, repo);
+      const info = listBranches(repo);
+      expect(info.detached).toBe(true);
+      expect(info.current).toMatch(/^[0-9a-f]+$/);
+    });
+
+    it("lists remote branches", () => {
+      // Create a "remote" by cloning
+      const clone = mkdtempSync(join(tmpdir(), "git-ops-clone-"));
+      try {
+        execSync(`git clone ${repo} ${clone}`, { stdio: "pipe" });
+        // Create a branch in origin that doesn't exist locally
+        git("checkout -b remote-only", repo);
+        writeFileSync(join(repo, "remote.txt"), "data");
+        git("add .", repo);
+        git("commit -m remote-only", repo);
+        git("checkout master", repo);
+
+        // Fetch in clone
+        git("fetch origin", clone);
+        const info = listBranches(clone);
+        const remotes = info.branches.filter((b) => b.isRemote);
+        const remoteNames = remotes.map((b) => b.name);
+        expect(remoteNames.some((n) => n.includes("remote-only"))).toBe(true);
+      } finally {
+        rmSync(clone, { recursive: true, force: true });
+      }
+    });
+
+    it("excludes origin/HEAD from remote branches", () => {
+      const clone = mkdtempSync(join(tmpdir(), "git-ops-clone-"));
+      try {
+        execSync(`git clone ${repo} ${clone}`, { stdio: "pipe" });
+        const info = listBranches(clone);
+        const remoteNames = info.branches.filter((b) => b.isRemote).map((b) => b.name);
+        expect(remoteNames.every((n) => !n.endsWith("/HEAD"))).toBe(true);
+      } finally {
+        rmSync(clone, { recursive: true, force: true });
+      }
+    });
+  });
+
+  describe("checkoutBranch", () => {
+    it("checks out a local branch on clean repo", () => {
+      git("checkout -b feature-x", repo);
+      git("checkout master", repo);
+      const result = checkoutBranch(repo, "feature-x", false);
+      expect(result.success).toBe(true);
+      const head = execSync("git rev-parse --abbrev-ref HEAD", { cwd: repo, encoding: "utf-8" }).trim();
+      expect(head).toBe("feature-x");
+    });
+
+    it("returns dirty when working tree is dirty and stash=false", () => {
+      git("checkout -b feature-y", repo);
+      git("checkout master", repo);
+      writeFileSync(join(repo, "README.md"), "dirty");
+      const result = checkoutBranch(repo, "feature-y", false);
+      expect(result.success).toBe(false);
+      if (!result.success) {
+        expect(result.dirty).toBe(true);
+        expect(result.files.length).toBeGreaterThan(0);
+      }
+    });
+
+    it("stashes and checks out when stash=true", () => {
+      git("checkout -b feature-z", repo);
+      git("checkout master", repo);
+      writeFileSync(join(repo, "README.md"), "dirty");
+      const result = checkoutBranch(repo, "feature-z", true);
+      expect(result.success).toBe(true);
+      if (result.success) {
+        expect(result.stashed).toBe(true);
+      }
+      const head = execSync("git rev-parse --abbrev-ref HEAD", { cwd: repo, encoding: "utf-8" }).trim();
+      expect(head).toBe("feature-z");
+    });
+
+    it("returns success when already on target branch", () => {
+      const result = checkoutBranch(repo, "master", false);
+      expect(result.success).toBe(true);
+    });
+
+    it("creates local tracking branch for remote branch", () => {
+      // Create remote-only branch
+      git("checkout -b only-remote", repo);
+      writeFileSync(join(repo, "r.txt"), "data");
+      git("add .", repo);
+      git("commit -m r", repo);
+      git("checkout master", repo);
+
+      const clone = mkdtempSync(join(tmpdir(), "git-ops-clone-"));
+      try {
+        execSync(`git clone ${repo} ${clone}`, { stdio: "pipe" });
+        const result = checkoutBranch(clone, "origin/only-remote", false);
+        expect(result.success).toBe(true);
+        const head = execSync("git rev-parse --abbrev-ref HEAD", { cwd: clone, encoding: "utf-8" }).trim();
+        expect(head).toBe("only-remote");
+      } finally {
+        rmSync(clone, { recursive: true, force: true });
+      }
+    });
+  });
+
+  describe("gitInit", () => {
+    it("initializes a new git repo", () => {
+      const dir = mkdtempSync(join(tmpdir(), "git-init-"));
+      try {
+        gitInit(dir);
+        expect(isGitRepo(dir)).toBe(true);
+      } finally {
+        rmSync(dir, { recursive: true, force: true });
+      }
+    });
+
+    it("throws if already a git repo", () => {
+      expect(() => gitInit(repo)).toThrow("already a git repository");
+    });
+  });
+
+  describe("stashPop", () => {
+    it("pops stash cleanly", () => {
+      writeFileSync(join(repo, "README.md"), "stashed-content");
+      git("stash push -u", repo);
+      const result = stashPop(repo);
+      expect(result.conflicts).toBe(false);
+    });
+
+    it("throws when no stash entries", () => {
+      expect(() => stashPop(repo)).toThrow("no stash entries");
+    });
+
+    it("detects conflicts on stash pop", () => {
+      // Create a stash, then modify same file on current branch
+      writeFileSync(join(repo, "README.md"), "stash-version");
+      git("stash push -u", repo);
+      writeFileSync(join(repo, "README.md"), "branch-version");
+      git("add .", repo);
+      git("commit -m conflict", repo);
+      const result = stashPop(repo);
+      expect(result.conflicts).toBe(true);
+    });
+  });
+});

package/packages/server/src/__tests__/headless-pid-registry.test.ts

@@ -0,0 +1,233 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { createHeadlessPidRegistry } from "../headless-pid-registry.js";
+import { EventEmitter } from "node:events";
+import { join } from "node:path";
+import { mkdtempSync, readFileSync, writeFileSync, mkdirSync } from "node:fs";
+import { tmpdir } from "node:os";
+import type { ChildProcess } from "node:child_process";
+
+function mockProcess(): ChildProcess {
+  return new EventEmitter() as any;
+}
+
+function makeTempDir() {
+  return mkdtempSync(join(tmpdir(), "pid-reg-test-"));
+}
+
+describe("HeadlessPidRegistry", () => {
+  it("should register and track a process", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+    expect(registry.size()).toBe(1);
+  });
+
+  it("should remove entry on process exit", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+    expect(registry.size()).toBe(1);
+    proc.emit("exit");
+    expect(registry.size()).toBe(0);
+  });
+
+  it("should link session ID by cwd", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+    const linked = registry.linkSession("session-1", "/projects/app");
+    expect(linked).toBe(true);
+    expect(registry.getPid("session-1")).toBe(100);
+  });
+
+  it("should return false when linking unknown cwd", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const linked = registry.linkSession("session-1", "/unknown");
+    expect(linked).toBe(false);
+  });
+
+  it("should use FIFO matching for same cwd", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc1 = mockProcess();
+    const proc2 = mockProcess();
+    registry.register(100, "/projects/app", proc1);
+    registry.register(200, "/projects/app", proc2);
+
+    registry.linkSession("session-1", "/projects/app");
+    expect(registry.getPid("session-1")).toBe(100);
+
+    registry.linkSession("session-2", "/projects/app");
+    expect(registry.getPid("session-2")).toBe(200);
+  });
+
+  it("should not link to already-linked entries", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+    registry.linkSession("session-1", "/projects/app");
+
+    const linked = registry.linkSession("session-2", "/projects/app");
+    expect(linked).toBe(false);
+  });
+
+  it("should return undefined for unknown session ID", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    expect(registry.getPid("unknown")).toBeUndefined();
+  });
+
+  it("should kill process by session ID", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc = mockProcess();
+    registry.register(process.pid, "/projects/app", proc);
+    const killSpy = vi.spyOn(process, "kill").mockImplementation(() => true);
+
+    registry.linkSession("session-1", "/projects/app");
+    const killed = registry.killBySessionId("session-1");
+    expect(killed).toBe(true);
+    expect(killSpy).toHaveBeenCalledWith(-process.pid, "SIGTERM");
+    expect(registry.size()).toBe(0);
+
+    killSpy.mockRestore();
+  });
+
+  it("should return false when killing unknown session", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const killed = registry.killBySessionId("unknown");
+    expect(killed).toBe(false);
+  });
+
+  it("should handle kill failure gracefully", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc = mockProcess();
+    registry.register(999999, "/projects/app", proc);
+    registry.linkSession("session-1", "/projects/app");
+
+    const killSpy = vi.spyOn(process, "kill").mockImplementation(() => {
+      throw new Error("ESRCH");
+    });
+
+    const killed = registry.killBySessionId("session-1");
+    expect(killed).toBe(false);
+    expect(registry.size()).toBe(0);
+
+    killSpy.mockRestore();
+  });
+
+  it("should remove by PID", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+    registry.remove(100);
+    expect(registry.size()).toBe(0);
+  });
+
+  it("should kill all tracked processes", () => {
+    const registry = createHeadlessPidRegistry({ pidFilePath: join(makeTempDir(), "pids.json") });
+    const proc1 = mockProcess();
+    const proc2 = mockProcess();
+    registry.register(100, "/a", proc1);
+    registry.register(200, "/b", proc2);
+
+    const killSpy = vi.spyOn(process, "kill").mockImplementation(() => true);
+    registry.killAll();
+    expect(killSpy).toHaveBeenCalledTimes(2);
+    expect(registry.size()).toBe(0);
+    killSpy.mockRestore();
+  });
+});
+
+describe("HeadlessPidRegistry persistence", () => {
+  it("should persist entries to disk on register", () => {
+    const dir = makeTempDir();
+    const pidFile = join(dir, "pids.json");
+    const registry = createHeadlessPidRegistry({ pidFilePath: pidFile });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+
+    const data = JSON.parse(readFileSync(pidFile, "utf-8"));
+    expect(data.entries).toHaveLength(1);
+    expect(data.entries[0].pid).toBe(100);
+    expect(data.entries[0].cwd).toBe("/projects/app");
+    expect(data.entries[0].spawnedAt).toBeDefined();
+  });
+
+  it("should remove entry from disk on process exit", () => {
+    const dir = makeTempDir();
+    const pidFile = join(dir, "pids.json");
+    const registry = createHeadlessPidRegistry({ pidFilePath: pidFile });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+    proc.emit("exit");
+
+    const data = JSON.parse(readFileSync(pidFile, "utf-8"));
+    expect(data.entries).toHaveLength(0);
+  });
+
+  it("should remove entry from disk on remove()", () => {
+    const dir = makeTempDir();
+    const pidFile = join(dir, "pids.json");
+    const registry = createHeadlessPidRegistry({ pidFilePath: pidFile });
+    const proc = mockProcess();
+    registry.register(100, "/projects/app", proc);
+    registry.remove(100);
+
+    const data = JSON.parse(readFileSync(pidFile, "utf-8"));
+    expect(data.entries).toHaveLength(0);
+  });
+});
+
+describe("HeadlessPidRegistry orphan cleanup", () => {
+  it("should reclaim alive processes from disk", () => {
+    const dir = makeTempDir();
+    const pidFile = join(dir, "pids.json");
+
+    // Pre-populate the PID file with current process PID (guaranteed alive)
+    writeFileSync(pidFile, JSON.stringify({
+      entries: [{ pid: process.pid, cwd: "/projects/app", spawnedAt: new Date().toISOString() }],
+    }));
+
+    const registry = createHeadlessPidRegistry({ pidFilePath: pidFile });
+    registry.cleanupOrphans();
+
+    expect(registry.size()).toBe(1);
+    expect(registry.getPid("any")).toBeUndefined(); // not linked yet
+  });
+
+  it("should remove dead processes from disk", () => {
+    const dir = makeTempDir();
+    const pidFile = join(dir, "pids.json");
+
+    // Use a PID that's almost certainly dead
+    writeFileSync(pidFile, JSON.stringify({
+      entries: [{ pid: 999999, cwd: "/projects/app", spawnedAt: new Date().toISOString() }],
+    }));
+
+    const registry = createHeadlessPidRegistry({ pidFilePath: pidFile });
+    registry.cleanupOrphans();
+
+    expect(registry.size()).toBe(0);
+    const data = JSON.parse(readFileSync(pidFile, "utf-8"));
+    expect(data.entries).toHaveLength(0);
+  });
+
+  it("should kill very old alive orphans (>7 days)", () => {
+    const dir = makeTempDir();
+    const pidFile = join(dir, "pids.json");
+
+    const oldDate = new Date(Date.now() - 8 * 24 * 60 * 60 * 1000).toISOString(); // 8 days ago
+    writeFileSync(pidFile, JSON.stringify({
+      entries: [{ pid: process.pid, cwd: "/projects/app", spawnedAt: oldDate }],
+    }));
+
+    const killSpy = vi.spyOn(process, "kill").mockImplementation(() => true);
+    const registry = createHeadlessPidRegistry({ pidFilePath: pidFile });
+    registry.cleanupOrphans();
+
+    // Should have tried to kill the process group
+    expect(killSpy).toHaveBeenCalledWith(-process.pid, "SIGTERM");
+    // Should NOT be reclaimed
+    expect(registry.size()).toBe(0);
+
+    killSpy.mockRestore();
+  });
+});

package/packages/server/src/__tests__/headless-shutdown-fallback.test.ts

@@ -0,0 +1,109 @@
+/**
+ * Integration test: shutdown fallback kills headless process when bridge is disconnected.
+ */
+import { describe, it, expect, afterAll } from "vitest";
+import { createServer, type DashboardServer } from "../server.js";
+import { WebSocket } from "ws";
+import { spawn } from "node:child_process";
+
+function waitForOpen(ws: WebSocket): Promise<void> {
+  return new Promise((resolve, reject) => {
+    if (ws.readyState === WebSocket.OPEN) return resolve();
+    ws.on("open", resolve);
+    ws.on("error", reject);
+    setTimeout(() => reject(new Error("open timeout")), 3000);
+  });
+}
+
+function collectMsgs(ws: WebSocket, ms: number): Promise<any[]> {
+  return new Promise((resolve) => {
+    const arr: any[] = [];
+    const h = (raw: any) => arr.push(JSON.parse(raw.toString()));
+    ws.on("message", h);
+    setTimeout(() => { ws.off("message", h); resolve(arr); }, ms);
+  });
+}
+
+const delay = (ms: number) => new Promise((r) => setTimeout(r, ms));
+const httpPort = 19190;
+const piPort = 19191;
+let server: DashboardServer;
+
+describe("Headless shutdown fallback", () => {
+  afterAll(async () => {
+    if (server) await server.stop();
+  });
+
+  it("should kill headless process via SIGTERM when bridge is disconnected", async () => {
+    server = await createServer({
+      port: httpPort, piPort, dev: true,
+      autoShutdown: false, shutdownIdleSeconds: 999, tunnel: false,
+    editor: { idleTimeoutMinutes: 10, maxInstances: 3 },
+    });
+    await server.start();
+
+    // Spawn a real dummy process (sleep) to act as the headless pi session
+    const dummy = spawn("sleep", ["60"], { detached: true, stdio: "ignore" });
+    dummy.unref();
+    const pid = dummy.pid!;
+
+    // Register it in the headless registry with a known cwd
+    const registry = server.browserGateway.headlessPidRegistry;
+    registry.register(pid, "/test/cwd", dummy);
+
+    // Simulate bridge connecting and registering with that cwd
+    const bridge = new WebSocket(`ws://localhost:${piPort}`);
+    await waitForOpen(bridge);
+    bridge.send(JSON.stringify({
+      type: "session_register", sessionId: "headless-1", cwd: "/test/cwd", source: "tui",
+    }));
+    // Wait for session_register to be processed (may need longer under load)
+    for (let i = 0; i < 20; i++) {
+      if (registry.getPid("headless-1") !== undefined) break;
+      await delay(50);
+    }
+
+    // Verify the session got linked
+    expect(registry.getPid("headless-1")).toBe(pid);
+
+    // Now disconnect the bridge (simulating bridge gone)
+    bridge.close();
+    await delay(200);
+
+    // Browser sends shutdown — bridge is disconnected, should fallback to kill
+    const browser = new WebSocket(`ws://localhost:${httpPort}/ws`);
+    await waitForOpen(browser);
+    await delay(100);
+
+    browser.send(JSON.stringify({ type: "shutdown", sessionId: "headless-1" }));
+    await delay(300);
+
+    // Verify process was killed
+    let alive = true;
+    try {
+      process.kill(pid, 0); // signal 0 = check if alive
+    } catch {
+      alive = false;
+    }
+    expect(alive).toBe(false);
+
+    browser.close();
+    await delay(50);
+  }, 15000);
+
+  it("should not crash when no PID is linked for shutdown", async () => {
+    const browser = new WebSocket(`ws://localhost:${httpPort}/ws`);
+    await waitForOpen(browser);
+    await delay(100);
+
+    // Send shutdown for an unknown session — should not crash
+    browser.send(JSON.stringify({ type: "shutdown", sessionId: "nonexistent" }));
+    await delay(200);
+
+    // If we get here without crashing, it's a pass
+    expect(browser.readyState).toBe(WebSocket.OPEN);
+
+    browser.close();
+    await delay(50);
+  }, 10000);
+});

package/packages/server/src/__tests__/health-endpoint.test.ts

@@ -0,0 +1,35 @@
+/**
+ * Tests for GET /api/health endpoint.
+ */
+import { describe, it, expect, afterEach } from "vitest";
+import { createServer, type DashboardServer } from "../server.js";
+
+const httpPort = 19090;
+const piPort = 19091;
+let server: DashboardServer;
+
+describe("GET /api/health", () => {
+  afterEach(async () => {
+    if (server) {
+      try { await server.stop(); } catch { /* already stopped */ }
+    }
+  });
+
+  it("should return ok, pid, and uptime", async () => {
+    server = await createServer({
+      port: httpPort, piPort, dev: true,
+      autoShutdown: false, shutdownIdleSeconds: 999, tunnel: false,
+    editor: { idleTimeoutMinutes: 10, maxInstances: 3 },
+    });
+    await server.start();
+
+    const res = await fetch(`http://localhost:${httpPort}/api/health`);
+    expect(res.status).toBe(200);
+
+    const body = await res.json();
+    expect(body.ok).toBe(true);
+    expect(body.pid).toBe(process.pid);
+    expect(typeof body.uptime).toBe("number");
+    expect(body.uptime).toBeGreaterThanOrEqual(0);
+  });
+});

package/packages/server/src/__tests__/heartbeat-ack.test.ts

@@ -0,0 +1,63 @@
+/**
+ * Tests for heartbeat_ack response in pi-gateway.
+ */
+import { describe, it, expect, afterEach } from "vitest";
+import { createPiGateway } from "../pi-gateway.js";
+import { createMemorySessionManager } from "../memory-session-manager.js";
+import { WebSocket } from "ws";
+
+const delay = (ms: number) => new Promise((r) => setTimeout(r, ms));
+
+function waitForOpen(ws: WebSocket): Promise<void> {
+  return new Promise((resolve, reject) => {
+    if (ws.readyState === WebSocket.OPEN) return resolve();
+    ws.on("open", resolve);
+    ws.on("error", reject);
+    setTimeout(() => reject(new Error("open timeout")), 3000);
+  });
+}
+
+let portCounter = 19550;
+
+describe("heartbeat_ack", () => {
+  let gateway: ReturnType<typeof createPiGateway>;
+
+  afterEach(() => {
+    gateway?.stop();
+  });
+
+  it("should respond with heartbeat_ack when receiving session_heartbeat", async () => {
+    const sessionManager = createMemorySessionManager();
+    gateway = createPiGateway(sessionManager, { heartbeatTimeout: 5000 });
+    const port = portCounter++;
+    gateway.start(port);
+
+    const ws = new WebSocket(`ws://localhost:${port}`);
+    await waitForOpen(ws);
+
+    // Register session first
+    ws.send(JSON.stringify({
+      type: "session_register", sessionId: "ack-test", cwd: "/tmp", source: "tui",
+    }));
+    await delay(100);
+
+    // Collect messages
+    const messages: any[] = [];
+    ws.on("message", (raw) => {
+      messages.push(JSON.parse(raw.toString()));
+    });
+
+    // Send heartbeat
+    ws.send(JSON.stringify({
+      type: "session_heartbeat", sessionId: "ack-test",
+    }));
+    await delay(100);
+
+    // Should have received heartbeat_ack
+    const ack = messages.find((m) => m.type === "heartbeat_ack");
+    expect(ack).toBeDefined();
+    expect(ack.type).toBe("heartbeat_ack");
+
+    ws.close();
+  }, 10000);
+});