harness-tools/extensions/pi-workflow/notifications.ts

import { randomUUID } from "node:crypto";
import type {
  ExtensionAPI,
  ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import type { WorkflowRun } from "./domain.ts";
import type { ContextTransitionRuntime } from "./context-transition.ts";
import { deriveRunLabel } from "./run-label.ts";

const NOTIFICATION_TYPE = "pi-workflow-notification";
const DELIVERY_OWNER_TOKEN = `${process.pid}:${randomUUID()}`;
const ACTIVE_DELIVERIES = new Map<
  string,
  { scheduledAt: number; runId: string }
>();
const DELIVERY_RETRY_MS = 30_000;

function receiptKeys(ctx: ExtensionContext) {
  const keys = new Set<string>();
  for (const entry of ctx.sessionManager.getEntries()) {
    if (entry.type !== "message") continue;
    const message = entry.message as unknown as Record<string, unknown>;
    if (message.role !== "custom" || message.customType !== NOTIFICATION_TYPE) {
      continue;
    }
    const details = message.details;
    if (!details || typeof details !== "object") continue;
    const key = (details as Record<string, unknown>).deliveryKey;
    if (typeof key === "string") keys.add(key);
  }
  return keys;
}

function deliveryIsCurrent(run: WorkflowRun, key: string) {
  if (key.startsWith("gate:")) {
    const gate = [...run.gates]
      .reverse()
      .find((candidate) => key.includes(`:${candidate.id}:`));
    const activeGate = [...run.gates]
      .reverse()
      .find((candidate) => candidate.status === "open");
    return Boolean(gate && gate.status === "open" && activeGate?.id === gate.id);
  }
  if (key.startsWith("terminal:")) {
    return key.includes(`:${run.status}:`) &&
      (run.status === "completed" ||
        run.status === "failed" ||
        run.status === "aborted");
  }
  return false;
}

function deliveryText(run: WorkflowRun, key: string) {
  const label = deriveRunLabel(run);
  if (key.startsWith("gate:")) {
    const gate = [...run.gates]
      .reverse()
      .find((candidate) => key.includes(`:${candidate.id}:`));
    return gate
      ? `pi-workflow ${label} is waiting at the ${gate.kind} Gate (revision ${run.revision}). Inspect the run before obtaining and recording the user's Decision.`
      : `pi-workflow ${label} is waiting for a Gate Decision.`;
  }
  const report = run.reportArtifact?.path
    ? ` Report: ${run.reportArtifact.path}`
    : "";
  return `pi-workflow ${label} finished with status ${run.status}.${report}`;
}

function belongsToSession(run: WorkflowRun, sessionId: string) {
  const activeSession =
    run.lease?.activeParentSessionId ?? run.parentSessionIds.at(-1);
  return activeSession === sessionId;
}

export async function reconcileWorkflowDeliveries(
  pi: Pick<ExtensionAPI, "sendMessage">,
  runtime: ContextTransitionRuntime,
  ctx: ExtensionContext,
) {
  const sessionId = ctx.sessionManager.getSessionId();
  if (!sessionId) return;
  let receipts = receiptKeys(ctx);
  for (const run of await runtime.store.listRuns()) {
    if (!belongsToSession(run, sessionId)) continue;
    if (
      (run.status === "completed" ||
        run.status === "failed" ||
        run.status === "aborted") &&
      run.reportArtifact
    ) {
      await runtime.store.ensureResult(run);
    }
    for (const delivery of run.deliveries) {
      if (delivery.state === "sent") continue;
      if (!deliveryIsCurrent(run, delivery.key)) {
        ACTIVE_DELIVERIES.delete(delivery.key);
        await runtime.store.clearDeliveryClaim(run.runId, delivery.key);
        await runtime.store.markDeliverySent(
          run.runId,
          delivery.key,
          Date.now(),
        );
        continue;
      }
      if (receipts.has(delivery.key)) {
        ACTIVE_DELIVERIES.delete(delivery.key);
        await runtime.store.clearDeliveryClaim(run.runId, delivery.key);
        await runtime.store.markDeliverySent(
          run.runId,
          delivery.key,
          Date.now(),
        );
        continue;
      }
      const active = ACTIVE_DELIVERIES.get(delivery.key);
      if (
        active &&
        Date.now() - active.scheduledAt < DELIVERY_RETRY_MS
      ) {
        continue;
      }
      const claimedAt = Date.now();
      const claimed = await runtime.store.claimDelivery(
        run.runId,
        delivery.key,
        DELIVERY_OWNER_TOKEN,
        claimedAt,
      );
      if (!claimed) continue;
      ACTIVE_DELIVERIES.set(delivery.key, {
        scheduledAt: claimedAt,
        runId: run.runId,
      });
      try {
        pi.sendMessage(
          {
            customType: NOTIFICATION_TYPE,
            content: deliveryText(run, delivery.key),
            display: true,
            details: {
              deliveryKey: delivery.key,
              runId: run.runId,
              revision: run.revision,
              status: run.status,
              reportPath: run.reportArtifact?.path,
            },
          },
          run.background
            ? { deliverAs: "followUp", triggerTurn: true }
            : { deliverAs: "steer" },
        );
        receipts = receiptKeys(ctx);
        if (receipts.has(delivery.key)) {
          ACTIVE_DELIVERIES.delete(delivery.key);
          await runtime.store.clearDeliveryClaim(run.runId, delivery.key);
          await runtime.store.markDeliverySent(
            run.runId,
            delivery.key,
            Date.now(),
          );
        }
      } catch (error) {
        ACTIVE_DELIVERIES.delete(delivery.key);
        await runtime.store.releaseDeliveryClaim(
          run.runId,
          delivery.key,
          DELIVERY_OWNER_TOKEN,
        );
        throw error;
      }
    }
  }
}

export function registerWorkflowNotifications(
  pi: ExtensionAPI,
  runtime: ContextTransitionRuntime,
) {
  const reconcile = async (ctx: ExtensionContext) => {
    try {
      await reconcileWorkflowDeliveries(pi, runtime, ctx);
    } catch (error) {
      ctx.ui.notify(
        `pi-workflow delivery recovery failed: ${error instanceof Error ? error.message : String(error)}`,
        "error",
      );
    }
  };
  pi.on("session_start", async (_event, ctx) => reconcile(ctx));
  pi.on("agent_settled", async (_event, ctx) => reconcile(ctx));
  pi.on("tool_result", async (event, ctx) => {
    if (event.toolName === "pi_workflow") await reconcile(ctx);
  });
  pi.on("session_shutdown", async () => {
    const active = [...ACTIVE_DELIVERIES.entries()];
    ACTIVE_DELIVERIES.clear();
    await Promise.all(
      active.map(([key, delivery]) =>
        runtime.store.releaseDeliveryClaim(
          delivery.runId,
          key,
          DELIVERY_OWNER_TOKEN,
        ),
      ),
    );
  });
}

export const workflowNotificationType = NOTIFICATION_TYPE;