Add post hooks for cleanup (#193)

* Add post hooks for cleanup

* Switch to signal-based cleanup

* Better exit handling

.github/workflows/pullfrog.yml

@@ -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 }}
-

agents/shared.ts

@@ -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
+      // build log box content: user prompt first, then event data with eventInstructions as property
-      const logParts = [
+      const eventWithInstructions = ctx.instructions.eventInstructions
-        ctx.instructions.eventInstructions
+        ? `additionalInstructions: ${ctx.instructions.eventInstructions}\n${ctx.instructions.event}`
-          ? `EVENT-LEVEL INSTRUCTIONS:\n${ctx.instructions.eventInstructions}`
+        : ctx.instructions.event;
-          : null,
+      const logParts = [ctx.instructions.user, eventWithInstructions].filter(Boolean);
-        ctx.instructions.user ? `USER REQUEST:\n${ctx.instructions.user}` : null,
-        ctx.instructions.event,
-      ].filter(Boolean);
       log.box(logParts.join("\n\n---\n\n"), {
         title: "Instructions",
       });

entry.ts

@@ -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();
   }
 }
 

main.ts

@@ -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 handleAgentResult(result);
-    return mainResult;
   } 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
-    }
   }
 }

mcp/comment.ts

@@ -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 { type OctokitWithPlugins, parseRepoContext } from "../utils/github.ts";
-import { getGitHubInstallationToken } from "../utils/token.ts";
+import type { ToolContext } from "./server.ts";
-import type { ToolContext, ToolState } 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 = {
+  ctx.toolState.progressCommentId = result.data.id;
-    id: result.data.id,
+  ctx.toolState.wasUpdated = true;
-    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
+  // reset state and mark as updated so post script doesn't try to handle it
-  ctx.toolState.progressComment = {
+  ctx.toolState.progressCommentId = null;
-    id: null,
+  ctx.toolState.wasUpdated = true;
-    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
+      // mark progress as updated so post script doesn't think the run failed
-      ctx.toolState.progressComment.wasUpdated = true;
+      ctx.toolState.wasUpdated = true;
 
       return {
         success: true,

mcp/server.ts

@@ -30,11 +30,9 @@ export interface ToolState {
     promise: Promise<PrepResult[]> | undefined;
     results: PrepResult[] | undefined;
   };
-  progressComment: {
+  progressCommentId: number | null;
-    id: number | null;
-    wasUpdated: boolean;
-  };
   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: {
+    progressCommentId: resolvedId,
-      id: Number.isNaN(progressCommentId) ? null : progressCommentId,
-      wasUpdated: false,
-    },
     backgroundProcesses: new Map(),
   };
 }

utils/errorReport.ts

@@ -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
+    workflowRun: runId ? { owner: repoContext.owner, repo: repoContext.name, runId } : undefined,
-      ? { 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
+  // mark as updated so exit handler doesn't try to update again
-  ctx.toolState.progressComment.wasUpdated = true;
+  ctx.toolState.wasUpdated = true;
 }

utils/exitHandler.ts

@@ -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
+  }
+}