Add post hooks for cleanup (#193)

* Add post hooks for cleanup * Switch to signal-based cleanup * Better exit handling
2026-01-28 01:59:15 +00:00
parent 90945a9481
commit 102417f442
9 changed files with 20395 additions and 20370 deletions
@@ -28,7 +28,7 @@ jobs:
        with:
          fetch-depth: 1
      - name: Run agent
-        uses: pullfrog/pullfrog@v0
+        uses: pullfrog/pullfrog@main
        with:
          prompt: ${{ inputs.prompt }}
        env:
@@ -43,4 +43,3 @@ jobs:
          GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
          DEEPSEEK_API_KEY: ${{ secrets.DEEPSEEK_API_KEY }}
          OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
-
@@ -33,14 +33,11 @@ export const agent = <const input extends AgentInput>(input: input): defineAgent
      const search = ctx.payload.search;
      const write = ctx.payload.write;
      log.info(`» running ${input.name} with effort=${ctx.payload.effort}...`);
-      // build log box content: eventInstructions (if any) + user request (if any) + event data
-      const logParts = [
-        ctx.instructions.eventInstructions
-          ? `EVENT-LEVEL INSTRUCTIONS:\n${ctx.instructions.eventInstructions}`
-          : null,
-        ctx.instructions.user ? `USER REQUEST:\n${ctx.instructions.user}` : null,
-        ctx.instructions.event,
-      ].filter(Boolean);
+      // build log box content: user prompt first, then event data with eventInstructions as property
+      const eventWithInstructions = ctx.instructions.eventInstructions
+        ? `additionalInstructions: ${ctx.instructions.eventInstructions}\n${ctx.instructions.event}`
+        : ctx.instructions.event;
+      const logParts = [ctx.instructions.user, eventWithInstructions].filter(Boolean);
      log.box(logParts.join("\n\n---\n\n"), {
        title: "Instructions",
      });
@@ -6,6 +6,7 @@

 import * as core from "@actions/core";
 import { main } from "./main.ts";
+import { runCleanup } from "./utils/exitHandler.ts";

 async function run(): Promise<void> {
  try {
@@ -17,6 +18,8 @@ async function run(): Promise<void> {
  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : "Unknown error occurred";
    core.setFailed(`Action failed: ${errorMessage}`);
+  } finally {
+    await runCleanup();
  }
 }

@@ -1,5 +1,4 @@
 // changes to tool permissions should be reflected in wiki/granular-tools.md
-import { ensureProgressCommentUpdated } from "./mcp/comment.ts";
 import { initToolState, startMcpHttpServer } from "./mcp/server.ts";
 import { computeModes } from "./modes.ts";
 import { resolveAgent } from "./utils/agent.ts";
@@ -7,6 +6,7 @@ import { validateApiKey } from "./utils/apiKeys.ts";
 import { resolveBody } from "./utils/body.ts";
 import { log, writeSummary } from "./utils/cli.ts";
 import { reportErrorToComment } from "./utils/errorReport.ts";
+import { setupExitHandler } from "./utils/exitHandler.ts";
 import { createOctokit } from "./utils/github.ts";
 import { resolveInstructions } from "./utils/instructions.ts";
 import { normalizeEnv } from "./utils/normalizeEnv.ts";
@@ -34,13 +34,15 @@ export async function main(): Promise<MainResult> {
  process.env.ORIGINAL_GITHUB_TOKEN = process.env.GITHUB_TOKEN;

  const timer = new Timer();
-  await using tokenRef = await resolveInstallationToken();
+  const tokenRef = await resolveInstallationToken();
  process.env.GITHUB_TOKEN = tokenRef.token;

  const octokit = createOctokit(tokenRef.token);
  const runInfo = await resolveRun({ octokit });
  const toolState = initToolState({ runInfo });

+  setupExitHandler(toolState);
+
  try {
    const repo = await resolveRepoData({ octokit, token: tokenRef.token });
    timer.checkpoint("repoData");
@@ -120,8 +122,7 @@ export async function main(): Promise<MainResult> {
      await writeSummary(toolState.lastProgressBody);
    }

-    const mainResult = await handleAgentResult(result);
-    return mainResult;
+    return handleAgentResult(result);
  } catch (error) {
    const errorMessage = error instanceof Error ? error.message : "Unknown error occurred";
    log.error(errorMessage);
@@ -134,13 +135,5 @@ export async function main(): Promise<MainResult> {
      success: false,
      error: errorMessage,
    };
-  } finally {
-    // ensure progress comment is updated if it was never updated during execution
-    // do this before revoking the token so we can still make API calls
-    try {
-      await ensureProgressCommentUpdated(toolState);
-    } catch {
-      // error updating comment, but don't let it mask the original error
-    }
  }
 }
@@ -1,9 +1,8 @@
 import { type } from "arktype";
 import type { Agent } from "../agents/index.ts";
 import { buildPullfrogFooter, stripExistingFooter } from "../utils/buildPullfrogFooter.ts";
-import { createOctokit, type OctokitWithPlugins, parseRepoContext } from "../utils/github.ts";
-import { getGitHubInstallationToken } from "../utils/token.ts";
-import type { ToolContext, ToolState } from "./server.ts";
+import { type OctokitWithPlugins, parseRepoContext } from "../utils/github.ts";
+import type { ToolContext } from "./server.ts";
 import { execute, tool } from "./shared.ts";

 /**
@@ -78,12 +77,12 @@ function buildImplementPlanLink(
  return `[Implement plan ➔](${apiUrl}/trigger/${owner}/${repo}/${issueNumber}?action=implement&comment_id=${commentId})`;
 }

-interface AddFooterCtx {
+export interface AddFooterCtx {
  agent?: Agent | undefined;
  octokit?: OctokitWithPlugins | undefined;
 }

-async function addFooter(ctx: AddFooterCtx, body: string): Promise<string> {
+export async function addFooter(ctx: AddFooterCtx, body: string): Promise<string> {
  const bodyWithoutFooter = stripExistingFooter(body);
  const footer = await buildCommentFooter({ agent: ctx.agent, octokit: ctx.octokit });
  return `${bodyWithoutFooter}${footer}`;
@@ -173,7 +172,7 @@ export async function reportProgress(
  // always track the body for job summary
  ctx.toolState.lastProgressBody = body;

-  const existingCommentId = ctx.toolState.progressComment.id;
+  const existingCommentId = ctx.toolState.progressCommentId;
  const issueNumber =
    ctx.toolState.prNumber ?? ctx.toolState.issueNumber ?? ctx.payload.event.issue_number;
  const isPlanMode = ctx.toolState.selectedMode === "Plan";
@@ -200,7 +199,7 @@ export async function reportProgress(
      body: bodyWithFooter,
    });

-    ctx.toolState.progressComment.wasUpdated = true;
+    ctx.toolState.wasUpdated = true;

    return {
      commentId: result.data.id,
@@ -229,10 +228,8 @@ export async function reportProgress(
  });

  // store the comment ID for future updates
-  ctx.toolState.progressComment = {
-    id: result.data.id,
-    wasUpdated: true,
-  };
+  ctx.toolState.progressCommentId = result.data.id;
+  ctx.toolState.wasUpdated = true;

  // if Plan mode, update the comment to add the "Implement plan" link
  if (isPlanMode) {
@@ -301,7 +298,7 @@ export function ReportProgressTool(ctx: ToolContext) {
 * Used after submitting a PR review since the review body contains all necessary info.
 */
 export async function deleteProgressComment(ctx: ToolContext): Promise<boolean> {
-  const existingCommentId = ctx.toolState.progressComment.id;
+  const existingCommentId = ctx.toolState.progressCommentId;
  if (!existingCommentId) {
    return false;
  }
@@ -321,83 +318,13 @@ export async function deleteProgressComment(ctx: ToolContext): Promise<boolean>
    }
  }

-  // reset state but mark as "updated" so ensureProgressCommentUpdated doesn't try to handle it
-  ctx.toolState.progressComment = {
-    id: null,
-    wasUpdated: true,
-  };
+  // reset state and mark as updated so post script doesn't try to handle it
+  ctx.toolState.progressCommentId = null;
+  ctx.toolState.wasUpdated = true;

  return true;
 }

-/**
- * Ensure the progress comment is updated with a generic error message if it was never updated.
- * This should be called after agent execution completes to handle cases where the agent
- * exited without ever calling reportProgress.
- *
- * Works even if MCP context is not initialized (e.g., if error occurs before MCP server starts).
- * Uses comment ID from toolState (set during initToolState from initial fetch).
- */
-export async function ensureProgressCommentUpdated(toolState: ToolState): Promise<void> {
-  // skip if comment was already updated during execution
-  if (toolState.progressComment.wasUpdated) {
-    return;
-  }
-
-  // skip if there's already a progress body recorded (agent called report_progress)
-  if (toolState.lastProgressBody) {
-    return;
-  }
-
-  // get comment ID from toolState (already fetched during initToolState)
-  const existingCommentId = toolState.progressComment.id;
-
-  // if still no comment ID, nothing to update
-  if (!existingCommentId) {
-    return;
-  }
-
-  // check if comment still says "leaping into action" - if it's been updated with an error, don't overwrite it
-  const repoContext = parseRepoContext();
-  const octokit = createOctokit(getGitHubInstallationToken());
-
-  try {
-    const existingComment = await octokit.rest.issues.getComment({
-      owner: repoContext.owner,
-      repo: repoContext.name,
-      comment_id: existingCommentId,
-    });
-
-    const commentBody = existingComment.data.body || "";
-    // if comment doesn't start with the leaping prefix, it's already been updated with an error or progress
-    if (!commentBody.startsWith(LEAPING_INTO_ACTION_PREFIX)) {
-      return;
-    }
-  } catch {
-    // can't fetch comment, skip update
-    return;
-  }
-
-  const runId = process.env.GITHUB_RUN_ID;
-  const workflowRunLink = runId
-    ? `[workflow run logs](https://github.com/${repoContext.owner}/${repoContext.name}/actions/runs/${runId})`
-    : "workflow run logs";
-
-  const errorMessage = `This run croaked 😵
-
-The workflow encountered an error before any progress could be reported. Please check the ${workflowRunLink} for details.`;
-
-  // add footer without agent info (we don't have context here)
-  const body = await addFooter({ octokit }, errorMessage);
-
-  await octokit.rest.issues.updateComment({
-    owner: repoContext.owner,
-    repo: repoContext.name,
-    comment_id: existingCommentId,
-    body,
-  });
-}
-
 export const ReplyToReviewComment = type({
  pull_number: type.number.describe("the pull request number"),
  comment_id: type.number.describe("the ID of the review comment to reply to"),
@@ -423,8 +350,8 @@ export function ReplyToReviewCommentTool(ctx: ToolContext) {
        body: bodyWithFooter,
      });

-      // mark progress as updated so ensureProgressCommentUpdated doesn't think the run failed
-      ctx.toolState.progressComment.wasUpdated = true;
+      // mark progress as updated so post script doesn't think the run failed
+      ctx.toolState.wasUpdated = true;

      return {
        success: true,
@@ -30,11 +30,9 @@ export interface ToolState {
    promise: Promise<PrepResult[]> | undefined;
    results: PrepResult[] | undefined;
  };
-  progressComment: {
-    id: number | null;
-    wasUpdated: boolean;
-  };
+  progressCommentId: number | null;
  lastProgressBody?: string;
+  wasUpdated?: boolean;
 }

 import type { ResolveRunResult } from "../utils/workflow.ts";
@@ -46,12 +44,10 @@ interface InitToolStateParams {
 export function initToolState(ctx: InitToolStateParams): ToolState {
  const progressCommentIdStr = ctx.runInfo.workflowRunInfo.progressCommentId;
  const progressCommentId = progressCommentIdStr ? parseInt(progressCommentIdStr, 10) : null;
+  const resolvedId = Number.isNaN(progressCommentId) ? null : progressCommentId;

  return {
-    progressComment: {
-      id: Number.isNaN(progressCommentId) ? null : progressCommentId,
-      wasUpdated: false,
-    },
+    progressCommentId: resolvedId,
    backgroundProcesses: new Map(),
  };
 }
@@ -12,7 +12,7 @@ interface ReportErrorParams {
 export async function reportErrorToComment(ctx: ReportErrorParams): Promise<void> {
  const formattedError = ctx.title ? `${ctx.title}\n\n${ctx.error}` : ctx.error;

-  const commentId = ctx.toolState.progressComment.id;
+  const commentId = ctx.toolState.progressCommentId;
  if (!commentId) {
    return;
  }
@@ -24,9 +24,7 @@ export async function reportErrorToComment(ctx: ReportErrorParams): Promise<void
  // build footer with workflow run link
  const footer = buildPullfrogFooter({
    triggeredBy: true,
-    workflowRun: runId
-      ? { owner: repoContext.owner, repo: repoContext.name, runId }
-      : undefined,
+    workflowRun: runId ? { owner: repoContext.owner, repo: repoContext.name, runId } : undefined,
  });

  await octokit.rest.issues.updateComment({
@@ -36,6 +34,6 @@ export async function reportErrorToComment(ctx: ReportErrorParams): Promise<void
    body: `${formattedError}${footer}`,
  });

-  // mark as updated so ensureProgressCommentUpdated doesn't try to update again
-  ctx.toolState.progressComment.wasUpdated = true;
+  // mark as updated so exit handler doesn't try to update again
+  ctx.toolState.wasUpdated = true;
 }
@@ -0,0 +1,102 @@
+import { LEAPING_INTO_ACTION_PREFIX } from "../mcp/comment.ts";
+import type { ToolState } from "../mcp/server.ts";
+import { buildPullfrogFooter } from "./buildPullfrogFooter.ts";
+import { log } from "./cli.ts";
+import { createOctokit, parseRepoContext } from "./github.ts";
+import { revokeGitHubInstallationToken } from "./token.ts";
+
+let cleanupFn: ((isCancellation: boolean) => Promise<void>) | undefined;
+
+export function setupExitHandler(toolState: ToolState): void {
+  let hasCleanedUp = false;
+
+  async function cleanup(isCancellation: boolean): Promise<void> {
+    if (hasCleanedUp) {
+      return;
+    }
+    hasCleanedUp = true;
+
+    const token = process.env.GITHUB_TOKEN;
+    const commentId = toolState.progressCommentId;
+    const wasUpdated = toolState.wasUpdated === true;
+
+    // update progress comment if it was never updated (still shows "leaping into action")
+    if (token && commentId && !wasUpdated) {
+      try {
+        const repoContext = parseRepoContext();
+        const octokit = createOctokit(token);
+
+        const existingComment = await octokit.rest.issues.getComment({
+          owner: repoContext.owner,
+          repo: repoContext.name,
+          comment_id: commentId,
+        });
+
+        const commentBody = existingComment.data.body || "";
+
+        // only update if comment still shows the initial "leaping into action" message
+        if (commentBody.startsWith(LEAPING_INTO_ACTION_PREFIX)) {
+          const runId = process.env.GITHUB_RUN_ID;
+          const workflowRunLink = runId
+            ? `[workflow run logs](https://github.com/${repoContext.owner}/${repoContext.name}/actions/runs/${runId})`
+            : "workflow run logs";
+
+          const errorMessage = isCancellation
+            ? `This run was cancelled 🛑\n\nThe workflow was cancelled before completion. Please check the ${workflowRunLink} for details.`
+            : `This run croaked 😵\n\nThe workflow encountered an error before any progress could be reported. Please check the ${workflowRunLink} for details.`;
+
+          const footer = buildPullfrogFooter({
+            triggeredBy: true,
+            workflowRun: runId
+              ? { owner: repoContext.owner, repo: repoContext.name, runId }
+              : undefined,
+          });
+
+          await octokit.rest.issues.updateComment({
+            owner: repoContext.owner,
+            repo: repoContext.name,
+            comment_id: commentId,
+            body: `${errorMessage}${footer}`,
+          });
+
+          log.info("» updated progress comment with error message");
+        }
+      } catch {
+        // ignore errors during cleanup
+      }
+    }
+
+    // revoke token
+    if (token) {
+      try {
+        await revokeGitHubInstallationToken(token);
+        log.debug("» installation token revoked");
+      } catch {
+        // ignore errors during cleanup
+      }
+    }
+  }
+
+  // store cleanup function for runCleanup()
+  cleanupFn = cleanup;
+
+  // handle cancellation signals
+  function handleSignal(): void {
+    log.info("» workflow cancelled, cleaning up...");
+    cleanup(true).finally(() => process.exit(1));
+  }
+
+  process.on("SIGINT", handleSignal);
+  process.on("SIGTERM", handleSignal);
+}
+
+/**
+ * Run cleanup explicitly. Called from entry.ts in finally block.
+ */
+export async function runCleanup(): Promise<void> {
+  try {
+    await cleanupFn?.(false);
+  } catch {
+    // ignore errors during cleanup
+  }
+}