path: root/src/index.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { type BrowserClient, createBrowserClient } from "./client";
import { getBrowserSystemPrompt } from "./prompt";
import { registerSetupCommand } from "./setup";
import { registerProfileCommand } from "./profile/command";
import { readPin } from "./profile/store";
import { type BrowserState, createStatePersister, defaultState, restoreState } from "./state";
import { registerAllTools } from "./registry";
import { cleanupTempDirs } from "./util/truncate";
import { restoreDownloadBehavior } from "./domains/files";
import { createDaemonTransport } from "./daemon/transport";
import { setDebugClicks } from "./util/debug";
import { asBoolean, asString } from "./util/guards";
import { describeCdpTarget, isValidTargetValue, localAliases, redactWsUrl, resolveCdpTarget } from "./cdp/endpoint";
import { clearTarget, readTarget, writeTarget } from "./cdp/target-store";
import { ensureDaemon, stopDaemon } from "./daemon/spawn";
import { queryLiveWsUrl } from "./cdp/discovery";

export default function browserHarnessExtension(pi: ExtensionAPI): void {
  const flagNs = asString(pi.getFlag("browser-namespace"));
  const namespace = flagNs ?? `pi-${Math.random().toString(36).slice(2, 10)}`;

  let state: BrowserState = defaultState(namespace);
  const persister = createStatePersister(pi);
  let client: BrowserClient | null = null;
  let toolsRegistered = false;

  pi.registerFlag("browser-namespace", {
    description: "Browser daemon namespace. Default: auto-generated",
    type: "string",
  });
  pi.registerFlag("browser-debug-clicks", {
    description: "Enable debug click overlay (saves annotated screenshots to /tmp)",
    type: "boolean",
    default: false,
  });

  pi.registerCommand("browser-status", {
    description: "Show browser connection status and current page",
    handler: async (_args, ctx) => {
      if (!client) {
        ctx.ui.notify("Browser client not started. Run /browser-setup first.", "warning");
        return;
      }
      const s = client.status();
      const pin = client.profilePin();

      // FORK: the connection is lazy — nothing dials the browser until the
      // first browser_* tool runs. A plain "Disconnected" reads as a fault
      // when it is really just "not needed yet", so say which it is, and show
      // what we would connect to.
      const target = resolveCdpTarget(process.env, await readTarget());
      // FORK: a BU_CDP_WS pin is not probed: /json/version answering says nothing
      // about whether the browser id pinned in that URL is still current.
      const reachable =
        target.kind === "remote" && target.source !== "ws"
          ? (await queryLiveWsUrl(target.port, 2_000, target.host)) !== null
          : null;

      const browserLine = s.alive
        ? "🟢 Connected"
        : reachable === false
          ? "🔴 Not connected — target unreachable"
          : "🟡 Not connected yet — connects on the first browser tool call";

      const targetLine =
        target.kind === "remote"
          ? `${describeCdpTarget(target)}${reachable === null ? "" : reachable ? " — reachable ✓" : " — unreachable ✗"}` +
            (reachable === false && target.source === "default"
              ? "\n         will fall back to this machine's browser"
              : "")
          : describeCdpTarget(target);

      const lines = [
        `Browser: ${browserLine}`,
        `Target:  ${targetLine}`,
        `Profile: ${pin ? pin.label : "not selected (uses whichever window is focused) — run /browser-profile"}`,
        `Session: ${s.sessionId ?? "none"}`,
      ];
      if (s.remoteBrowserId) lines.push(`Browser ID: ${s.remoteBrowserId}`);
      if (s.alive) {
        const info = await client.pageInfo();
        if (info.success) {
          if ("dialog" in info.data) {
            lines.push(`\n⚠️  Dialog open: ${info.data.dialog.type} — "${info.data.dialog.message}"`);
          } else {
            lines.push(
              `\nCurrent Page:`,
              `  URL: ${info.data.url}`,
              `  Title: ${info.data.title}`,
              `  Viewport: ${info.data.width}x${info.data.height}`,
            );
          }
        }
      }
      ctx.ui.notify(lines.join("\n"), "info");
    },
  });

  // FORK: switch which browser the harness drives, without leaving pi.
  //
  // The target is resolved per daemon process and the daemon holds the browser
  // connection, so a switch has to replace that process — writing the store
  // alone would leave tool calls going to the old browser while /browser-status
  // reported the new one.
  pi.registerCommand("browser-target", {
    description: "Show or change which browser to drive (local, default, or host:port)",
    handler: async (args, ctx) => {
      const wanted = (typeof args === "string" ? args : "").trim();

      if (wanted === "") {
        const current = resolveCdpTarget(process.env, await readTarget());
        ctx.ui.notify(
          `Browser target: ${describeCdpTarget(current)}\n` +
            `Change with: /browser-target local | default | host:port`,
          "info",
        );
        return;
      }

      const clearing = wanted.toLowerCase() === "default";
      if (!clearing && !isValidTargetValue(wanted)) {
        ctx.ui.notify(
          `Not a usable target: "${wanted}"\n` +
            `Expected "host:port", "default", or one of ${localAliases().join(", ")}.`,
          "error",
        );
        return;
      }

      const saved = clearing ? await clearTarget() : await writeTarget(wanted);
      if (!saved.success) {
        ctx.ui.notify(saved.error, "error");
        return;
      }

      const resolved = resolveCdpTarget(process.env, await readTarget());

      // Warn rather than silently disobey: the environment deliberately wins,
      // so a stored choice made now would not take effect this session.
      // FORK: BU_CDP_WS outranks both (upstream reads it first); show only its host:port.
      const wsOverride = (process.env["BU_CDP_WS"] ?? "").trim();
      const httpOverride = (process.env["BU_CDP_HTTP"] ?? "").trim();
      const envOverride =
        wsOverride !== "" ? `BU_CDP_WS (${redactWsUrl(wsOverride)})` : httpOverride !== "" ? `BU_CDP_HTTP=${httpOverride}` : "";
      if (envOverride !== "") {
        ctx.ui.notify(
          `Saved, but ${envOverride} is set in this environment and takes precedence.\n` +
            `Still using: ${describeCdpTarget(resolved)}`,
          "warning",
        );
        return;
      }

      ctx.ui.notify(`Switching to ${describeCdpTarget(resolved)}...`, "info");

      if (client) await client.stop();
      if (!(await stopDaemon())) {
        ctx.ui.notify("Saved, but the browser daemon did not stop — run /browser-setup to pick it up.", "warning");
        return;
      }
      if (!(await ensureDaemon())) {
        ctx.ui.notify("Saved, but the browser daemon did not come back. Run /browser-setup.", "error");
        return;
      }

      if (!client) {
        ctx.ui.notify(`Browser target set to ${describeCdpTarget(resolved)} ✓ — run /browser-setup to connect.`, "info");
        return;
      }
      const started = await client.start();
      ctx.ui.notify(
        started.success
          ? `Browser target set to ${describeCdpTarget(resolved)} ✓`
          : `Target saved (${describeCdpTarget(resolved)}) but connecting failed: ${started.error.message}`,
        started.success ? "info" : "error",
      );
    },
  });

  pi.registerCommand("browser-reload-daemon", {
    description: "Restart the browser client",
    handler: async (_args, ctx) => {
      if (!client) {
        ctx.ui.notify("Browser client not started.", "warning");
        return;
      }
      ctx.ui.notify("Restarting browser client...", "info");
      await client.stop();
      const r = await client.start();
      if (r.success) {
        ctx.ui.notify("Browser client restarted ✓", "info");
      } else {
        ctx.ui.notify(`Restart failed: ${r.error.message}`, "error");
      }
    },
  });

  pi.on("session_start", async (_event, ctx) => {
    state = restoreState(ctx, state.namespace);
    persister.markPersisted(state);
    setDebugClicks(asBoolean(pi.getFlag("browser-debug-clicks")) ?? false);

    // Re-read the on-disk profile pin every session start: another pi session may have changed it since this client was created.
    const profilePin = await readPin();

    if (!client) {
      const transport = createDaemonTransport(state.namespace);

      const initialOwnership: { ownedTargetIds?: ReadonlyArray<string>; harnessWindowTargetId?: string; harnessWindowId?: number } = {};
      if (state.ownedTargetIds !== undefined) initialOwnership.ownedTargetIds = state.ownedTargetIds;
      if (state.harnessWindowTargetId !== undefined) initialOwnership.harnessWindowTargetId = state.harnessWindowTargetId;
      if (state.harnessWindowId !== undefined) initialOwnership.harnessWindowId = state.harnessWindowId;

      client = createBrowserClient({
        namespace: state.namespace,
        transport,
        profilePin,
        ...(Object.keys(initialOwnership).length > 0 ? { initialOwnership } : {}),
        onOwnershipChange: (snap) => {
          state = {
            ...state,
            ownedTargetIds: snap.ownedTargetIds,
            ...(snap.harnessWindowTargetId !== undefined
              ? { harnessWindowTargetId: snap.harnessWindowTargetId }
              : {}),
            ...(snap.harnessWindowId !== undefined
              ? { harnessWindowId: snap.harnessWindowId }
              : {}),
          };
          try { persister.persist(state); } catch {}
        },
      });
    } else {
      client.setProfilePin(profilePin);
    }

    if (!toolsRegistered) {
      registerAllTools(pi, client);
      toolsRegistered = true;
    }
    registerSetupCommand(pi, client);
    registerProfileCommand(pi, client);
  });

  pi.on("session_shutdown", async () => {
    if (client) {
      try {
        // Warn but keep going: detach and tab cleanup still have to run, and a throw here would skip both.
        const restored = await restoreDownloadBehavior(client);
        if (!restored.success) {
          console.warn(`[pi-browser-harness] ${restored.error.message}`);
        }
        await client.detach();
        await client.closeOwnedTabs();
      } catch (e) {
        console.warn("[pi-browser-harness] browser teardown failed during shutdown:", e);
      }
      // Do NOT stop the client or null it out: the surviving daemon connection is what eliminates the per-session "Allow Remote Debugging" prompt.
    }
    persister.persist(state);
    await cleanupTempDirs();
  });

  // The branch already holds the restored snapshot: re-appending it only grows the session file.
  pi.on("session_tree", async (_event, ctx) => {
    state = restoreState(ctx, client?.namespace);
    persister.markPersisted(state);
  });

  pi.on("before_agent_start", async (event) => {
    const target = resolveCdpTarget(process.env, await readTarget());
    return {
      systemPrompt: event.systemPrompt + getBrowserSystemPrompt(client?.status().alive ?? false, target),
    };
  });
}