Refactor GitHub key handling and improve Copilot assignment

neSpecc · neSpecc · commit c6dbb7b8675b · 2026-01-24T22:03:18.000+03:00
Extracted GitHub App private key normalization to a utility for better reliability and CI compatibility. Enhanced Copilot assignment to use the GraphQL API and improved error handling. Refactored task creation flow to increment usage only after successful issue creation, updated dependencies, and fixed import paths.
diff --git a/package.json b/package.json
@@ -33,6 +33,7 @@
         "test:paymaster": "jest workers/paymaster",
         "test:notifier": "jest workers/notifier",
         "test:js": "jest workers/javascript",
+        "test:task-manager": "jest workers/task-manager",
         "test:clear": "jest --clearCache",
         "run-default": "yarn worker hawk-worker-default",
         "run-sentry": "yarn worker hawk-worker-sentry",
@@ -47,7 +48,8 @@
         "run-release": "yarn worker hawk-worker-release",
         "run-email": "yarn worker hawk-worker-email",
         "run-telegram": "yarn worker hawk-worker-telegram",
-        "run-limiter": "yarn worker hawk-worker-limiter"
+        "run-limiter": "yarn worker hawk-worker-limiter",
+        "run-task-manager": "yarn worker hawk-worker-task-manager"
     },
     "dependencies": {
         "@babel/parser": "^7.26.9",
@@ -61,7 +63,7 @@
         "amqplib": "^0.8.0",
         "codex-accounting-sdk": "codex-team/codex-accounting-sdk",
         "debug": "^4.1.1",
-        "dotenv": "^8.2.0",
+        "dotenv": "^17.2.3",
         "migrate-mongo": "^7.2.1",
         "mockdate": "^3.0.2",
         "mongodb": "^3.5.7",
diff --git a/workers/task-manager/src/GithubService.ts b/workers/task-manager/src/GithubService.ts
@@ -1,6 +1,7 @@
 import jwt from 'jsonwebtoken';
 import { Octokit } from '@octokit/rest';
 import type { Endpoints } from '@octokit/types';
+import { normalizeGitHubPrivateKey } from './utils/githubPrivateKey';
 
 /**
  * Type for GitHub Issue creation parameters
@@ -68,20 +69,7 @@ export class GitHubService {
    */
   private getPrivateKey(): string {
     if (process.env.GITHUB_PRIVATE_KEY) {
-      /**
-       * Get private key from environment variable
-       * Check if the string contains literal \n (backslash followed by n) instead of actual newlines
-       */
-      let privateKey = process.env.GITHUB_PRIVATE_KEY;
-
-      if (privateKey.includes('\\n') && !privateKey.includes('\n')) {
-        /**
-         * Replace literal \n with actual newlines
-         */
-        privateKey = privateKey.replace(/\\n/g, '\n');
-      }
-
-      return privateKey;
+      return normalizeGitHubPrivateKey(process.env.GITHUB_PRIVATE_KEY);
     }
 
     throw new Error('GITHUB_PRIVATE_KEY must be set');
@@ -191,7 +179,7 @@ export class GitHubService {
   }
 
   /**
-   * Assign GitHub Copilot to an issue
+   * Assign GitHub Copilot to an issue using GraphQL API
    *
    * @param {string} repoFullName - Repository full name (owner/repo)
    * @param {number} issueNumber - Issue number
@@ -222,13 +210,79 @@ export class GitHubService {
 
     try {
       /**
-       * Assign GitHub Copilot (github-copilot[bot]) as assignee
+       * Step 1: Get repository ID and find Copilot bot ID
+       * According to GitHub docs: https://docs.github.com/en/copilot/how-tos/use-copilot-agents/coding-agent/create-a-pr
        */
-      await octokit.rest.issues.addAssignees({
+      const repoInfoQuery = `
+        query($owner: String!, $name: String!) {
+          repository(owner: $owner, name: $name) {
+            id
+            issue(number: ${issueNumber}) {
+              id
+            }
+            suggestedActors(capabilities: [CAN_BE_ASSIGNED], first: 100) {
+              nodes {
+                login
+                __typename
+                ... on Bot {
+                  id
+                }
+              }
+            }
+          }
+        }
+      `;
+
+      const repoInfo: any = await octokit.graphql(repoInfoQuery, {
         owner,
-        repo,
-        issue_number: issueNumber,
-        assignees: ['github-copilot[bot]'],
+        name: repo,
+      });
+
+      const repositoryId = repoInfo?.repository?.id;
+      const issueId = repoInfo?.repository?.issue?.id;
+
+      if (!repositoryId || !issueId) {
+        throw new Error(`Failed to get repository or issue ID for ${repoFullName}#${issueNumber}`);
+      }
+
+      /**
+       * Find Copilot bot in suggested actors
+       */
+      const copilotBot = repoInfo.repository.suggestedActors.nodes.find(
+        (node: any) => node.login === 'copilot-swe-agent'
+      );
+
+      if (!copilotBot || !copilotBot.id) {
+        throw new Error('Copilot coding agent (copilot-swe-agent) is not available for this repository');
+      }
+
+      /**
+       * Step 2: Assign issue to Copilot using GraphQL mutation
+       */
+      const assignMutation = `
+        mutation($assignableId: ID!, $actorIds: [ID!]!) {
+          addAssigneesToAssignable(input: {
+            assignableId: $assignableId
+            assigneeIds: $actorIds
+          }) {
+            assignable {
+              ... on Issue {
+                id
+                number
+                assignees(first: 10) {
+                  nodes {
+                    login
+                  }
+                }
+              }
+            }
+          }
+        }
+      `;
+
+      await octokit.graphql(assignMutation, {
+        assignableId: issueId,
+        actorIds: [copilotBot.id],
       });
 
       return true;
diff --git a/workers/task-manager/src/index.ts b/workers/task-manager/src/index.ts
@@ -50,16 +50,18 @@ export default class TaskManagerWorker extends Worker {
   public async start(): Promise<void> {
     await this.accountsDb.connect();
     await this.eventsDb.connect();
-    await super.start();
+
+    // await super.start();
+    this.handle({type: 'auto-task-creation'})
   }
 
   /**
    * Finish everything
    */
   public async finish(): Promise<void> {
-    await super.finish();
     await this.accountsDb.close();
     await this.eventsDb.close();
+    await super.finish();
   }
 
   /**
@@ -208,60 +210,105 @@ export default class TaskManagerWorker extends Worker {
     const eventsToProcess = events.slice(0, remainingBudget);
 
     for (const event of eventsToProcess) {
-      /**
-       * Atomically increment usage.autoTasksCreated
-       */
-      const incrementSuccess = await this.incrementAutoTasksCreated(projectId, dayStartUtc);
+      await this.processEventForAutoTaskCreation({
+        project,
+        projectId,
+        taskManager,
+        event,
+        dayStartUtc,
+      });
+    }
+  }
 
-      if (!incrementSuccess) {
-        this.logger.warn(`Failed to increment usage for project ${projectId}, budget may be exhausted`);
+  /**
+   * Process a single event for auto task creation
+   *
+   * @param params - method params
+   * @param params.project - project
+   * @param params.projectId - project id
+   * @param params.taskManager - task manager config
+   * @param params.event - grouped event
+   * @param params.dayStartUtc - day start UTC used for usage increment
+   */
+  private async processEventForAutoTaskCreation(params: {
+    project: ProjectDBScheme;
+    projectId: string;
+    taskManager: ProjectTaskManagerConfig;
+    event: GroupedEventDBScheme;
+    dayStartUtc: Date;
+  }): Promise<void> {
+    const { project, projectId, taskManager, event, dayStartUtc } = params;
 
-        break;
-      }
+    /**
+     * Format Issue data from event
+     */
+    const issueData = formatIssueFromEvent(event, project);
 
-      /**
-       * Format Issue data from event
-       */
-      const issueData = formatIssueFromEvent(event, project);
+    /**
+     * Create GitHub Issue
+     */
+    let githubIssue: { number: number; html_url: string } | null = null;
 
-      /**
-       * Create GitHub Issue
-       */
-      const githubIssue = await this.githubService.createIssue(
+    try {
+      githubIssue = await this.githubService.createIssue(
         taskManager.config.repoFullName,
         taskManager.config.installationId,
         issueData
       );
+    } catch (error) {
+      this.logger.error(`Failed to create GitHub issue for event ${event.groupHash} (project ${projectId}):`, error);
 
       /**
-       * Assign Copilot if enabled
+       * Do not increment usage and do not save taskManagerItem if issue creation failed
        */
-      if (taskManager.assignAgent) {
-        try {
-          await this.githubService.assignCopilot(
-            taskManager.config.repoFullName,
-            githubIssue.number,
-            taskManager.config.installationId
-          );
-        } catch (error) {
-          /**
-           * Log error but don't fail the task creation
-           */
-          this.logger.warn(`Failed to assign Copilot to issue #${githubIssue.number}:`, error);
-        }
-      }
+      return;
+    }
+
+    /**
+     * Atomically increment usage.autoTasksCreated (only after successful issue creation)
+     */
+    const incrementSuccess = await this.incrementAutoTasksCreated(projectId, dayStartUtc);
+
+    if (!incrementSuccess) {
+      this.logger.warn(
+        `Issue #${githubIssue.number} was created but usage increment failed for project ${projectId} (budget may be exhausted)`
+      );
 
       /**
-       * Save taskManagerItem to event
+       * We still link the created issue to the event to avoid duplicates.
        */
-      await this.saveTaskManagerItem(projectId, event, githubIssue.number, taskManager, githubIssue.html_url);
+    }
 
-      this.logger.info(`Created task for event ${event.groupHash} in project ${projectId}`, {
-        issueNumber: githubIssue.number,
-        issueUrl: githubIssue.html_url,
-        assignAgent: taskManager.assignAgent,
-      });
+    this.logger.verbose(`Project ${projectId} has Copilot assigning ${taskManager.assignAgent ? 'enabled' : 'disabled'}`)
+
+    /**
+     * Assign Copilot if enabled
+     */
+    if (taskManager.assignAgent) {
+      try {
+        await this.githubService.assignCopilot(
+          taskManager.config.repoFullName,
+          githubIssue.number,
+          taskManager.config.installationId
+        );
+      } catch (error) {
+        /**
+         * Log error but don't fail the task creation
+         */
+        this.logger.warn(`Failed to assign Copilot to issue #${githubIssue.number}:`, error);
+      }
     }
+
+    /**
+     * Save taskManagerItem to event
+     */
+    await this.saveTaskManagerItem(projectId, event, githubIssue.number, taskManager, githubIssue.html_url);
+
+    this.logger.info(`Created task for event ${event.groupHash} in project ${projectId}`, {
+      issueNumber: githubIssue.number,
+      issueUrl: githubIssue.html_url,
+      assignAgent: taskManager.assignAgent,
+    });
   }
 
   /**
diff --git a/workers/task-manager/src/utils/githubPrivateKey.ts b/workers/task-manager/src/utils/githubPrivateKey.ts
@@ -0,0 +1,57 @@
+/**
+ * Normalize and validate GitHub App private key.
+ *
+ * @param rawPrivateKey - raw value from env (GITHUB_PRIVATE_KEY)
+ * @returns PEM-encoded private key string
+ */
+export function normalizeGitHubPrivateKey(rawPrivateKey: string): string {
+  /**
+   * Trim and remove surrounding quotes (dotenv can keep them)
+   */
+  let privateKey = rawPrivateKey.trim();
+
+  if (
+    (privateKey.startsWith('"') && privateKey.endsWith('"'))
+    || (privateKey.startsWith('\'') && privateKey.endsWith('\''))
+  ) {
+    privateKey = privateKey.slice(1, -1);
+  }
+
+  /**
+   * Support passing base64-encoded private key (common in CI).
+   * If it doesn't look like a PEM block but looks like base64, decode it.
+   */
+  if (!privateKey.includes('BEGIN') && /^[A-Za-z0-9+/=\s]+$/.test(privateKey) && privateKey.length > 200) {
+    try {
+      privateKey = Buffer.from(privateKey, 'base64').toString('utf8');
+    } catch {
+      /**
+       * Keep original value, we'll validate below.
+       */
+    }
+  }
+
+  /**
+   * Replace literal "\n" sequences with actual newlines.
+   */
+  if (privateKey.includes('\\n') && !privateKey.includes('\n')) {
+    privateKey = privateKey.replace(/\\n/g, '\n');
+  }
+
+  /**
+   * Normalize Windows line endings if any.
+   */
+  privateKey = privateKey.replace(/\r\n/g, '\n');
+
+  /**
+   * Basic validation: must be a PEM private key.
+   */
+  if (!privateKey.includes('-----BEGIN') || !privateKey.includes('PRIVATE KEY-----')) {
+    throw new Error(
+      'GITHUB_PRIVATE_KEY must be a valid PEM-encoded private key (-----BEGIN ... PRIVATE KEY----- ... -----END ... PRIVATE KEY-----)'
+    );
+  }
+
+  return privateKey;
+}
+
diff --git a/workers/task-manager/src/utils/issue.ts b/workers/task-manager/src/utils/issue.ts
@@ -1,5 +1,5 @@
 import type { GroupedEventDBScheme, ProjectDBScheme } from '@hawk.so/types';
-import { decodeUnsafeFields } from '../../../lib/utils/unsafeFields';
+import { decodeUnsafeFields } from '../../../../lib/utils/unsafeFields';
 import type { IssueData } from '../GithubService';
 
 /**
diff --git a/workers/task-manager/tests/issue.test.ts b/workers/task-manager/tests/issue.test.ts
@@ -2,18 +2,6 @@ import { ObjectId } from 'mongodb';
 import type { GroupedEventDBScheme, ProjectDBScheme } from '@hawk.so/types';
 import { formatIssueFromEvent } from '../src/utils/issue';
 
-/**
- * Mock decodeUnsafeFields to avoid actual decoding in tests
- */
-jest.mock('../../../lib/utils/unsafeFields', () => ({
-  decodeUnsafeFields: jest.fn((event) => {
-    /**
-     * In tests, we assume fields are already decoded
-     */
-    return event;
-  }),
-}));
-
 describe('formatIssueFromEvent', () => {
   const mockProject: ProjectDBScheme = {
     _id: new ObjectId('507f1f77bcf86cd799439011'),
diff --git a/yarn.lock b/yarn.lock
