path: root/src/domains/files.ts

import { mkdir, access, constants, stat, readFile, writeFile } from "node:fs/promises";
import { basename } from "node:path";
import { Type } from "typebox";
import type { BrowserClient } from "../client";
import { type Result, err, ok } from "../util/result";
import { defineBrowserTool, type ToolErr, type ToolOk } from "../util/tool";
import { safeJs } from "../util/js-template";
import { pdfPath } from "../util/paths";
import { cdpCall, cdpCallBrowser, evalJs } from "./cdp-call";
import { resolveRefToBackendId } from "./ref-resolve";

const UploadArgs = Type.Object({
  ref: Type.Optional(
    Type.String({ description: "Stable ref of the file <input> from browser_snapshot. PREFERRED over selector." }),
  ),
  selector: Type.Optional(Type.String({ description: "CSS selector of the file <input> (fallback when no ref)" })),
  filePath: Type.String({ description: "Absolute path to the file to upload" }),
});

const tryRefUpload = async (
  client: BrowserClient,
  backendId: number,
  filePath: string,
): Promise<Result<void, ToolErr>> => {
  const set = await cdpCall(client, "DOM.setFileInputFiles", { files: [filePath], backendNodeId: backendId });
  if (!set.success) return set;
  return ok(undefined);
};

const verifyReadable = async (filePath: string): Promise<Result<void, ToolErr>> => {
  try {
    await access(filePath, constants.R_OK);
    return ok(undefined);
  } catch (e) {
    return err({
      kind: "io_error",
      message: `Cannot read file: ${filePath} (${e instanceof Error ? e.message : String(e)})`,
    });
  }
};

const tryCdpUpload = async (
  client: BrowserClient,
  selector: string,
  filePath: string,
): Promise<Result<void, ToolErr>> => {
  const doc = await cdpCall(client, "DOM.getDocument", { depth: -1 });
  if (!doc.success) return doc;
  const q = await cdpCall(client, "DOM.querySelector", { nodeId: doc.data.root.nodeId, selector });
  if (!q.success) return q;
  const nodeId = q.data.nodeId;
  if (!nodeId) return err({ kind: "invalid_state", message: `Selector matched 0 file inputs: ${selector}` });
  const set = await cdpCall(client, "DOM.setFileInputFiles", { files: [filePath], nodeId });
  if (!set.success) return set;
  const verify = await evalJs(client, safeJs`document.querySelector(${selector})?.files?.length || 0`);
  if (!verify.success) return verify;
  if (Number(verify.data ?? 0) === 0) {
    return err({ kind: "invalid_state", message: "CDP upload reported success but file count is 0" });
  }
  return ok(undefined);
};

const inferMime = (name: string): string => {
  if (name.endsWith(".png")) return "image/png";
  if (name.endsWith(".jpg") || name.endsWith(".jpeg")) return "image/jpeg";
  if (name.endsWith(".pdf")) return "application/pdf";
  if (name.endsWith(".json")) return "application/json";
  return "text/plain";
};

const jsFallbackUpload = async (
  client: BrowserClient,
  selector: string,
  filePath: string,
): Promise<Result<void, ToolErr>> => {
  const buf = await readFile(filePath);
  const st = await stat(filePath);
  const name = basename(filePath);
  const mime = inferMime(name);
  const expr = safeJs`
    (() => {
      const input = document.querySelector(${selector});
      if (!input || input.type !== 'file') throw new Error('File input not found');
      const bin = atob(${buf.toString("base64")});
      const bytes = new Uint8Array(bin.length);
      for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
      const file = new File([bytes], ${name}, { type: ${mime}, lastModified: ${st.mtimeMs} });
      const dt = new DataTransfer();
      dt.items.add(file);
      input.files = dt.files;
      input.dispatchEvent(new Event('change', { bubbles: true }));
      return input.files.length;
    })()
  `;
  const r = await evalJs(client, expr);
  if (!r.success) return r;
  if (!Number(r.data ?? 0)) return err({ kind: "invalid_state", message: "JS fallback set 0 files" });
  return ok(undefined);
};

export const uploadFileTool = defineBrowserTool({
  name: "browser_upload_file",
  label: "Browser Upload File",
  description: "Set files on a file <input> via CDP, with a JS-DataTransfer fallback for stubborn pages. PREFERRED: pass `ref` from browser_snapshot; fallback: a CSS `selector`.",
  promptSnippet: "Upload a file to a file input by ref (preferred) or selector",
  promptGuidelines: [
    "File inputs are often hidden; browser_snapshot still surfaces them.",
    "File path must be absolute and readable.",
  ],
  parameters: UploadArgs,
  concurrency: "serialized",
  async handler(args, { client }): Promise<Result<ToolOk, ToolErr>> {
    const readable = await verifyReadable(args.filePath);
    if (!readable.success) return readable;

    if (args.ref !== undefined) {
      const backendId = resolveRefToBackendId(client, args.ref);
      if (!backendId.success) return backendId;
      const up = await tryRefUpload(client, backendId.data, args.filePath);
      if (!up.success) return up;
      return ok({ text: `Uploaded ${args.filePath} via CDP (ref ${args.ref})`, details: { mode: "cdp", ref: args.ref, filePath: args.filePath } });
    }

    if (args.selector === undefined) {
      return err({ kind: "invalid_state", message: "Provide either `ref` or `selector`." });
    }
    const cdp = await tryCdpUpload(client, args.selector, args.filePath);
    if (cdp.success) {
      return ok({
        text: `Uploaded ${args.filePath} via CDP`,
        details: { mode: "cdp", filePath: args.filePath },
      });
    }
    const js = await jsFallbackUpload(client, args.selector, args.filePath);
    if (js.success) {
      return ok({
        text: `Uploaded ${args.filePath} via JS DataTransfer`,
        details: { mode: "js", filePath: args.filePath },
      });
    }
    return err({
      kind: "cdp_error",
      message: `Both CDP and JS fallback failed. CDP: ${cdp.error.message}; JS: ${js.error.message}`,
    });
  },
});

// Chrome stores the override in its default browser context, so it outlives this client and
// silently redirects the user's own downloads. Track it, so session teardown can undo it.
// The flag only records what THIS process set: the daemon is shared by every pi process on the
// machine and the CDP primitive is last-writer-wins, so another session may have reset or
// re-pointed it since. A crashed or killed session never reaches teardown, so the override can
// also outlive the process that set it.
let activeDownloadOverride: string | undefined;

// Shutdown must not block on a browser that is already gone: a pending restore against a
// disconnected daemon would otherwise trigger the one-shot reconnect (and Chrome's consent
// prompt) that 0.11.1 confined to explicit tool calls, and wait up to the connect timeout.
const SHUTDOWN_RESTORE_TIMEOUT_MS = 2_000;

export const restoreDownloadBehavior = async (client: BrowserClient): Promise<Result<void, ToolErr>> => {
  const pending = activeDownloadOverride;
  if (pending === undefined) return ok(undefined);
  if (!client.status().alive) {
    return err({ kind: "not_connected", message: `browser not connected; downloads may keep going to ${pending} until Chrome restarts` });
  }
  const r = await cdpCallBrowser(client, "Browser.setDownloadBehavior", { behavior: "default" }, { timeoutMs: SHUTDOWN_RESTORE_TIMEOUT_MS });
  if (!r.success) {
    return err({ ...r.error, message: `${r.error.message} (downloads may keep going to ${pending})` });
  }
  activeDownloadOverride = undefined;
  return ok(undefined);
};

const DownloadArgs = Type.Object({
  downloadPath: Type.Optional(
    Type.String({ description: "Absolute path to a writable directory where downloads should be saved" }),
  ),
  restore: Type.Optional(
    Type.Boolean({
      description:
        "Undo a previous browser_download: reset Chrome's download behavior back to the user's own default.",
    }),
  ),
});

export const downloadTool = defineBrowserTool({
  name: "browser_download",
  label: "Browser Download",
  description:
    "Configure Chrome's download behavior: set the save directory and disable the save-as prompt, or restore the user's default. Restored automatically when this session shuts down cleanly.",
  promptSnippet: "Configure download directory",
  promptGuidelines: [
    "Pass an absolute path to an existing writable directory.",
    "The override applies to the user's own downloads too, not just yours, and only a clean session shutdown undoes it. Call browser_download({restore:true}) as soon as you are done.",
  ],
  parameters: DownloadArgs,
  concurrency: "serialized",
  async handler(args, { client }): Promise<Result<ToolOk, ToolErr>> {
    if (args.restore === true) {
      if (args.downloadPath !== undefined) {
        return err({ kind: "invalid_state", message: "Pass either downloadPath or restore, not both" });
      }
      const r = await cdpCallBrowser(client, "Browser.setDownloadBehavior", { behavior: "default" });
      if (!r.success) return r;
      activeDownloadOverride = undefined;
      return ok({ text: "Downloads restored to Chrome's default location", details: { restored: true } });
    }
    if (args.downloadPath === undefined) {
      return err({ kind: "invalid_state", message: "Provide downloadPath to set a directory, or restore:true to undo" });
    }
    try {
      await mkdir(args.downloadPath, { recursive: true });
    } catch (e) {
      return err({
        kind: "io_error",
        message: `Cannot create download directory: ${args.downloadPath} (${e instanceof Error ? e.message : String(e)})`,
      });
    }
    try {
      const s = await stat(args.downloadPath);
      if (!s.isDirectory()) {
        return err({ kind: "io_error", message: `Not a directory: ${args.downloadPath}` });
      }
      await access(args.downloadPath, constants.W_OK);
    } catch (e) {
      return err({
        kind: "io_error",
        message: `Download path unusable: ${e instanceof Error ? e.message : String(e)}`,
      });
    }
    const r = await cdpCallBrowser(client, "Browser.setDownloadBehavior", {
      behavior: "allow",
      downloadPath: args.downloadPath,
      eventsEnabled: true,
    });
    if (!r.success) return r;
    activeDownloadOverride = args.downloadPath;
    return ok({
      text: `Downloads will save to: ${args.downloadPath}`,
      details: { downloadPath: args.downloadPath },
    });
  },
});

const PrintPdfArgs = Type.Object({
  outputPath: Type.Optional(Type.String({ description: "Where to save the PDF. Default: tmpdir + uuid." })),
});

export const printToPdfTool = defineBrowserTool({
  name: "browser_print_to_pdf",
  label: "Browser Print to PDF",
  description: "Print the current page to a PDF file using Chrome's Page.printToPDF.",
  promptSnippet: "Print the current page to PDF",
  promptGuidelines: ["Default output path is in tmpdir; pass outputPath to control."],
  parameters: PrintPdfArgs,
  concurrency: "serialized",
  async handler(args, { client }): Promise<Result<ToolOk, ToolErr>> {
    const r = await cdpCall(client, "Page.printToPDF", {
      printBackground: true,
      preferCSSPageSize: true,
    });
    if (!r.success) return r;
    const path = args.outputPath ?? pdfPath(client.namespace);
    await writeFile(path, Buffer.from(r.data.data, "base64"));
    return ok({ text: `PDF saved: ${path}`, details: { path } });
  },
});