fix: Regenerate installation token before reinvoking run (#656)

* fix: Regenerate installation token before reinvoking run

* cr

apps/open-swe/src/graphs/manager/nodes/create-new-session.ts

@@ -6,6 +6,9 @@ import {
 } from "@open-swe/shared/open-swe/manager/types";
 import { createIssueFieldsFromMessages } from "../utils/generate-issue-fields.js";
 import {
+  GITHUB_INSTALLATION_ID,
+  GITHUB_INSTALLATION_TOKEN_COOKIE,
+  LOCAL_MODE_HEADER,
   MANAGER_GRAPH_ID,
   OPEN_SWE_STREAM_MODE,
 } from "@open-swe/shared/constants";
@@ -24,6 +27,11 @@ import { getBranchName } from "../../../utils/github/git.js";
 import { getDefaultHeaders } from "../../../utils/default-headers.js";
 import { getCustomConfigurableFields } from "../../../utils/config.js";
 import { StreamMode } from "@langchain/langgraph-sdk";
+import { isLocalMode } from "@open-swe/shared/open-swe/local-mode";
+import { regenerateInstallationToken } from "../../../utils/github/regenerate-token.js";
+import { createLogger, LogLevel } from "../../../utils/logger.js";
+
+const logger = createLogger(LogLevel.INFO, "CreateNewSession");
 
 /**
  * Create new manager session.
@@ -73,8 +81,20 @@ ${ISSUE_CONTENT_CLOSE_TAG}`,
     }),
   ];
 
+  const isLocal = isLocalMode(config);
+  const defaultHeaders = isLocal
+    ? { [LOCAL_MODE_HEADER]: "true" }
+    : getDefaultHeaders(config);
+
+  if (!isLocal) {
+    logger.info("Regenerating installation token before starting new session.");
+    defaultHeaders[GITHUB_INSTALLATION_TOKEN_COOKIE] =
+      await regenerateInstallationToken(defaultHeaders[GITHUB_INSTALLATION_ID]);
+    logger.info("Regenerated installation token before starting new session.");
+  }
+
   const langGraphClient = createLangGraphClient({
-    defaultHeaders: getDefaultHeaders(config),
+    defaultHeaders,
   });
 
   const newManagerThreadId = uuidv4();

apps/open-swe/src/graphs/manager/nodes/start-planner.ts

@@ -10,6 +10,8 @@ import {
   OPEN_SWE_STREAM_MODE,
   PLANNER_GRAPH_ID,
   LOCAL_MODE_HEADER,
+  GITHUB_INSTALLATION_ID,
+  GITHUB_INSTALLATION_TOKEN_COOKIE,
 } from "@open-swe/shared/constants";
 import { createLogger, LogLevel } from "../../../utils/logger.js";
 import { getBranchName } from "../../../utils/github/git.js";
@@ -18,6 +20,7 @@ import { getDefaultHeaders } from "../../../utils/default-headers.js";
 import { getCustomConfigurableFields } from "../../../utils/config.js";
 import { getRecentUserRequest } from "../../../utils/user-request.js";
 import { StreamMode } from "@langchain/langgraph-sdk";
+import { regenerateInstallationToken } from "../../../utils/github/regenerate-token.js";
 
 const logger = createLogger(LogLevel.INFO, "StartPlanner");
 
@@ -37,14 +40,20 @@ export async function startPlanner(
   });
 
   const localMode = isLocalMode(config);
+  const defaultHeaders = localMode
+    ? { [LOCAL_MODE_HEADER]: "true" }
+    : getDefaultHeaders(config);
+
+  if (!localMode) {
+    logger.info("Regenerating installation token before starting planner run.");
+    defaultHeaders[GITHUB_INSTALLATION_TOKEN_COOKIE] =
+      await regenerateInstallationToken(defaultHeaders[GITHUB_INSTALLATION_ID]);
+    logger.info("Regenerated installation token before starting planner run.");
+  }
 
   try {
     const langGraphClient = createLangGraphClient({
-      defaultHeaders: localMode
+      defaultHeaders,
-        ? {
-            [LOCAL_MODE_HEADER]: "true",
-          }
-        : getDefaultHeaders(config),
     });
 
     const runInput: PlannerGraphUpdate = {

apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts

@@ -26,6 +26,8 @@ import {
   PROGRAMMER_GRAPH_ID,
   OPEN_SWE_STREAM_MODE,
   LOCAL_MODE_HEADER,
+  GITHUB_INSTALLATION_ID,
+  GITHUB_INSTALLATION_TOKEN_COOKIE,
 } from "@open-swe/shared/constants";
 import { PlannerGraphState } from "@open-swe/shared/open-swe/planner/types";
 import { createLangGraphClient } from "../../../utils/langgraph-client.js";
@@ -45,6 +47,7 @@ import {
   cleanTaskItems,
 } from "../../../utils/github/plan.js";
 import { isLocalMode } from "@open-swe/shared/open-swe/local-mode";
+import { regenerateInstallationToken } from "../../../utils/github/regenerate-token.js";
 
 const logger = createLogger(LogLevel.INFO, "ProposedPlan");
 
@@ -87,10 +90,24 @@ async function startProgrammerRun(input: {
   newMessages?: BaseMessage[];
 }) {
   const { runInput, state, config, newMessages } = input;
+  const isLocal = isLocalMode(config);
+  const defaultHeaders = isLocal
+    ? { [LOCAL_MODE_HEADER]: "true" }
+    : getDefaultHeaders(config);
+
+  if (!isLocal) {
+    logger.info(
+      "Regenerating installation token before starting programmer run.",
+    );
+    defaultHeaders[GITHUB_INSTALLATION_TOKEN_COOKIE] =
+      await regenerateInstallationToken(defaultHeaders[GITHUB_INSTALLATION_ID]);
+    logger.info(
+      "Regenerated installation token before starting programmer run.",
+    );
+  }
+
   const langGraphClient = createLangGraphClient({
-    defaultHeaders: isLocalMode(config)
+    defaultHeaders,
-      ? { [LOCAL_MODE_HEADER]: "true" }
-      : getDefaultHeaders(config),
   });
 
   const programmerThreadId = uuidv4();

apps/open-swe/src/utils/github/regenerate-token.ts

@@ -0,0 +1,28 @@
+import { encryptSecret } from "@open-swe/shared/crypto";
+import { getInstallationToken } from "@open-swe/shared/github/auth";
+
+export async function regenerateInstallationToken(
+  installationId: string | undefined,
+): Promise<string> {
+  if (!installationId) {
+    throw new Error(
+      "Missing installation ID for regenerating installation token.",
+    );
+  }
+
+  const appId = process.env.GITHUB_APP_ID;
+  const privateKey = process.env.GITHUB_APP_PRIVATE_KEY;
+  const secretsEncryptionKey = process.env.SECRETS_ENCRYPTION_KEY;
+  if (!appId || !privateKey || !secretsEncryptionKey) {
+    throw new Error(
+      "Missing environment variables for regenerating installation token.",
+    );
+  }
+
+  const newInstallationToken = await getInstallationToken(
+    installationId,
+    appId,
+    privateKey,
+  );
+  return encryptSecret(newInstallationToken, secretsEncryptionKey);
+}