path: root/crates/dinggerd/assets/pi/pi-agent-state.ts

// installed by dingger
// managed by dingger; reinstalling or updating the integration overwrites this file.
// add custom hooks/plugins beside this file instead of editing it.

import fs from "node:fs";
import net from "node:net";
import path from "node:path";

// The dingger hook receiver listens here; override with DINGGER_AGENT_SOCKET.
const socketPath =
  process.env.DINGGER_AGENT_SOCKET ||
  `${process.env.HOME}/.local/state/dingger/agent.sock`;

// tmux sets TMUX_PANE in the environment of every process it starts in a
// pane, so it names exactly the pane this pi runs in. (Asking tmux for "the
// current pane" is wrong: that is the session's active pane, not ours.)
const paneId = process.env.DINGGER_PANE_ID || process.env.TMUX_PANE;
const source = "dingger:pi";

// Filled in by `dingger install pi`: the daemon reloads an agent that
// reports another version (`dingger reload-agents`).
const extensionVersion = "__DINGGER_EXTENSION_VERSION__";

// Where the daemon asks this extension to reload: one socket per extension
// instance (a reload starts a new one beside the old), next to the daemon's.
const controlPath = path.join(
  path.dirname(socketPath),
  `pi-${process.pid}-${Math.random().toString(36).slice(2, 8)}.sock`,
);
let controlListening = false;

function enabled() {
  return !!paneId;
}

function sendRequestAttempt(request: unknown, timeoutMs: number): Promise<boolean> {
  if (!enabled()) {
    return Promise.resolve(true);
  }

  return new Promise((resolve) => {
    let done = false;
    let timeout: ReturnType<typeof setTimeout> | undefined;
    const finish = (delivered: boolean) => {
      if (done) return;
      done = true;
      if (timeout) {
        clearTimeout(timeout);
      }
      socket.destroy();
      resolve(delivered);
    };

    const socket = net.createConnection(socketPath);
    socket.on("error", () => finish(false));
    socket.on("connect", () => socket.write(`${JSON.stringify(request)}\n`));
    socket.on("data", () => finish(true));
    socket.on("end", () => finish(false));
    timeout = setTimeout(() => finish(false), timeoutMs);
    timeout.unref?.();
  });
}

// Whether the daemon got it.
async function sendRequest(request: unknown): Promise<boolean> {
  if (await sendRequestAttempt(request, 500)) {
    return true;
  }
  return sendRequestAttempt(request, 1500);
}

type AgentState = "working" | "blocked" | "idle";

// The session file, so `dingger restore` can resume this session with
// `pi --session <file>`. It changes with /new, /resume, and forks.
let sessionFile: string | undefined;

function updateSessionFile(ctx: any): void {
  try {
    const file = ctx?.sessionManager?.getSessionFile?.();
    sessionFile = typeof file === "string" && file.length > 0 ? file : undefined;
  } catch {
    sessionFile = undefined;
  }
}

// The agent's importance (/importance): it orders agents within a status
// in dingger's list and jump, and a low agent makes no sound. Saved in the
// session (a custom entry, never sent to the model), so it comes back with
// the conversation: on /resume, `pi --session`, and dingger's restore.
// Other extensions can follow it on the event bus: `dingger:importance`,
// `{ level }`, at session start and on every change.
type Importance = "low" | "normal" | "high";
const importanceLevels: Importance[] = ["high", "normal", "low"];
const importanceEntry = "dingger-importance";
let importance: Importance = "normal";

function savedImportance(ctx: any): Importance {
  try {
    const entries = ctx?.sessionManager?.getEntries?.() ?? [];
    for (let i = entries.length - 1; i >= 0; i--) {
      const entry = entries[i];
      if (entry?.type === "custom" && entry.customType === importanceEntry) {
        const level = entry.data?.level;
        return importanceLevels.includes(level) ? level : "normal";
      }
    }
  } catch {
    // An unreadable session: the default.
  }
  return "normal";
}

// The draft in the input box, as last reported: dingger marks a draft you
// left behind. Each session's draft is also kept in a file of its own, so
// it comes back when the session does after a crash, a restart, or
// `pi --session`.
let draft = "";
// How full the context window is, in percent, as pi's footer shows it;
// "" when unknown (right after compacting, before the next reply).
let context = "";
const draftDir = path.join(path.dirname(socketPath), "drafts");
// Enough to find the agent by it; the file keeps all of it.
const reportedDraftChars = 500;

function draftFile(): string | undefined {
  return sessionFile ? path.join(draftDir, `${path.basename(sessionFile)}.draft`) : undefined;
}

function saveDraft(text: string): void {
  const file = draftFile();
  if (!file) {
    return;
  }
  try {
    if (text.trim()) {
      fs.mkdirSync(draftDir, { recursive: true });
      fs.writeFileSync(file, text);
    } else {
      fs.unlinkSync(file);
    }
  } catch {
    // None to remove, or nowhere to keep it.
  }
}

// The session's saved draft, into an empty input box.
function restoreDraft(ctx: any): void {
  const file = draftFile();
  if (!file) {
    return;
  }
  try {
    if ((ctx.ui.getEditorText() ?? "") !== "") {
      return;
    }
    const saved = fs.readFileSync(file, "utf8");
    if (saved.trim()) {
      ctx.ui.setEditorText(saved);
    }
  } catch {
    // No saved draft.
  }
}

function requestId(): string {
  return `${source}:${Date.now()}:${Math.random().toString(36).slice(2)}`;
}

async function report(state: AgentState | undefined, message?: string): Promise<void> {
  await sendRequest({
    id: requestId(),
    method: "pane.report_agent",
    params: {
      pane_id: paneId!,
      source,
      agent: "pi",
      ...(state ? { state } : {}),
      message,
      importance,
      extension: extensionVersion,
      ...(controlListening ? { control: controlPath } : {}),
      draft: draft.slice(0, reportedDraftChars),
      context,
      seq: Date.now() * 1000,
      ...(sessionFile ? { agent_session_path: sessionFile } : {}),
    },
  });
}

let sendInFlight = false;
let queuedState: { state: AgentState; message?: string } | undefined;

function queueState(state: AgentState, message?: string): void {
  queuedState = { state, message };
  if (!sendInFlight) {
    void drainStateQueue();
  }
}

async function drainStateQueue(): Promise<void> {
  if (sendInFlight) {
    return;
  }

  sendInFlight = true;
  try {
    while (queuedState) {
      const next = queuedState;
      queuedState = undefined;
      await report(next.state, next.message);
    }
  } finally {
    sendInFlight = false;
    if (queuedState) {
      void drainStateQueue();
    }
  }
}

export default function (pi) {
  if (!enabled()) {
    return;
  }

  let agentActive = false;
  let blockedCount = 0;
  let blockedMessage: string | undefined;
  let lastState: AgentState | undefined;
  let lastMessage: string | undefined;

  function desiredState() {
    if (blockedCount > 0) {
      return { state: "blocked" as const, message: blockedMessage };
    }
    if (agentActive) {
      return { state: "working" as const, message: undefined };
    }
    return { state: "idle" as const, message: undefined };
  }

  function publishState(force = false) {
    const next = desiredState();
    if (!force && next.state === lastState && next.message === lastMessage) {
      return;
    }
    lastState = next.state;
    lastMessage = next.message;
    queueState(next.state, next.message);
  }

  pi.events.on("blocked", (data) => {
    if (!data?.active) {
      blockedCount = Math.max(0, blockedCount - 1);
      if (blockedCount === 0) {
        blockedMessage = undefined;
      }
      publishState();
      return;
    }

    blockedCount += 1;
    blockedMessage = data.label;
    publishState();
  });

  // No footer status: it would take a line of its own in pi.
  function announceImportance() {
    pi.events.emit("dingger:importance", { level: importance });
  }

  pi.registerCommand("importance", {
    description: "Set this agent's importance in dingger: high, normal, or low",
    // None once a level is typed out, so Enter runs the command at once.
    getArgumentCompletions: (prefix: string) =>
      importanceLevels.includes(prefix.trim() as Importance)
        ? null
        : importanceLevels
            .filter((level) => level.startsWith(prefix.trim()))
            .map((level) => ({ value: level, label: level })),
    handler: async (args: string, ctx: any) => {
      const wanted = args.trim();
      if (!wanted) {
        ctx.ui.notify(`importance: ${importance}`, "info");
        return;
      }
      if (!importanceLevels.includes(wanted as Importance)) {
        ctx.ui.notify("importance is high, normal, or low", "error");
        return;
      }
      if (wanted !== importance) {
        importance = wanted as Importance;
        pi.appendEntry(importanceEntry, { level: importance });
        announceImportance();
        // Without a state: the state stays as last reported.
        if (enabled()) {
          void report(undefined);
        }
      }
      ctx.ui.notify(`importance: ${importance}`, "info");
    },
  });

  // Like "mark as unread" in mail: dingger shows the agent as done again
  // (a pane option; nothing is saved in the session).
  pi.registerCommand("unread", {
    description: "Mark this agent unread in dingger: it shows as done until you come back to it",
    handler: async (_args: string, ctx: any) => {
      if (agentActive || blockedCount > 0) {
        ctx.ui.notify("only an idle agent can be marked unread", "error");
        return;
      }
      const delivered = await sendRequest({
        id: requestId(),
        method: "pane.mark_unread",
        params: { pane_id: paneId! },
      });
      if (delivered) {
        ctx.ui.notify("marked unread: dingger shows it as done until you come back to it", "info");
      } else {
        ctx.ui.notify("dingger is not running", "error");
      }
    },
  });

  // The latest context, for the control socket's checks and the input box.
  let latestCtx: any;
  let control: net.Server | undefined;

  // Read the input box once typing pauses; report and save it when it changed.
  let draftTimer: ReturnType<typeof setTimeout> | undefined;
  let stopWatchingInput: (() => void) | undefined;

  function checkDraft(): void {
    let text: string;
    try {
      text = latestCtx?.ui?.getEditorText?.() ?? "";
    } catch {
      return;
    }
    if (text === draft) {
      return;
    }
    draft = text;
    saveDraft(text);
    void report(undefined);
  }

  /** Report how full the context window is, when that changed. */
  function checkContext(ctx: any, quiet = false): void {
    let next = "";
    try {
      const percent = ctx?.getContextUsage?.()?.percent;
      if (typeof percent === "number" && Number.isFinite(percent)) {
        next = percent.toFixed(1);
      }
    } catch {
      return;
    }
    if (next === context) {
      return;
    }
    context = next;
    if (!quiet) {
      void report(undefined);
    }
  }

  function checkDraftSoon(): void {
    if (draftTimer) {
      clearTimeout(draftTimer);
    }
    draftTimer = setTimeout(checkDraft, 500);
    draftTimer.unref?.();
  }

  // Run only by the control socket below, which checked that pi is idle
  // with an empty input box: extensions can reload pi only from a command.
  pi.registerCommand("dingger-reload", {
    description: "Reload pi (dingger uses it to update its extension; same as /reload)",
    handler: async (_args: string, ctx: any) => {
      await ctx.reload();
    },
  });

  // Why this pi should not reload now, or undefined if it can.
  function reloadRefusal(): string | undefined {
    try {
      if (agentActive || blockedCount > 0 || latestCtx?.isIdle?.() !== true || latestCtx?.hasPendingMessages?.()) {
        return "busy";
      }
      if ((latestCtx?.ui?.getEditorText?.() ?? "").trim() !== "") {
        return "draft";
      }
    } catch {
      return "no context";
    }
    return undefined;
  }

  function startControl() {
    if (control) {
      return;
    }
    control = net.createServer((socket) => {
      let buffer = "";
      socket.on("error", () => socket.destroy());
      socket.on("data", (data) => {
        buffer += data.toString();
        const newline = buffer.indexOf("\n");
        if (newline < 0) {
          return;
        }
        let method: unknown;
        try {
          method = JSON.parse(buffer.slice(0, newline))?.method;
        } catch {
          method = undefined;
        }
        if (method !== "reload") {
          socket.end(`${JSON.stringify({ ok: false, reason: "unknown method" })}\n`);
          return;
        }
        const refusal = reloadRefusal();
        socket.end(`${JSON.stringify(refusal ? { ok: false, reason: refusal } : { ok: true })}\n`);
        if (!refusal) {
          // Checked and started in one go: no keystroke comes in between.
          pi.sendUserMessage("/dingger-reload", { expandPromptTemplates: true });
        }
      });
    });
    control.on("error", () => {
      controlListening = false;
    });
    try {
      fs.unlinkSync(controlPath);
    } catch {
      // None left over.
    }
    control.listen(controlPath, () => {
      controlListening = true;
      void report(undefined);
    });
    // The socket alone does not keep pi running.
    control.unref();
  }

  pi.on("session_shutdown", () => {
    // The input box as it is now: a quit right after typing comes before
    // the next check.
    if (draftTimer) {
      checkDraft();
    }
    stopWatchingInput?.();
    stopWatchingInput = undefined;
    if (draftTimer) {
      clearTimeout(draftTimer);
      draftTimer = undefined;
    }
    if (!control) {
      return;
    }
    control.close();
    control = undefined;
    controlListening = false;
    try {
      fs.unlinkSync(controlPath);
    } catch {
      // Gone already.
    }
  });

  pi.on("session_start", async (event, ctx) => {
    // TUI only: RPC/JSON/print modes are headless.
    if (ctx?.mode !== "tui") {
      return;
    }
    latestCtx = ctx;
    startControl();
    // A reload can replace this extension mid-run without emitting another agent_start.
    agentActive = ctx?.isIdle?.() === false;
    updateSessionFile(ctx);
    restoreDraft(ctx);
    // Before the first report, which carries it.
    draft = "";
    checkDraft();
    stopWatchingInput?.();
    // Every key reaches this before the input box: read the box after it.
    stopWatchingInput = ctx.ui.onTerminalInput?.(() => {
      checkDraftSoon();
      return undefined;
    });
    // Each session has its own: /new starts at normal, /resume brings its own back.
    importance = savedImportance(ctx);
    announceImportance();
    // The report below carries it.
    checkContext(ctx, true);
    publishState(true);
  });

  pi.on("agent_start", (_event, ctx) => {
    latestCtx = ctx;
    agentActive = true;
    updateSessionFile(ctx);
    publishState();
    // The prompt left the input box.
    checkDraftSoon();
  });

  // Each message fills the context (a prompt by pi's estimate, a reply by
  // its usage); compacting empties it; another model or branch changes it.
  for (const event of ["message_end", "session_compact", "model_select", "session_tree"]) {
    pi.on(event, (_event, ctx) => {
      latestCtx = ctx;
      checkContext(ctx);
    });
  }

  pi.on("agent_settled", (_event, ctx) => {
    latestCtx = ctx;
    if (ctx?.isIdle?.() !== true) {
      return;
    }

    agentActive = false;
    // The last reply's usage: message_end comes before pi counts it.
    checkContext(ctx);
    publishState();
    // Escape puts queued messages back into the input box.
    checkDraftSoon();
  });
}