feat: implement in-memory token caching for GitHub requests and enhance job completion handling

Author: devarshishimpi

src/server/core/github.ts

@@ -196,9 +196,22 @@ export class GitHubClient {
     private readonly tracker?: { incrementSubrequests(count?: number): void },
   ) {}
 
+  // In-memory token cache scoped to this client instance (i.e. one Worker invocation). Without it,
+  // every GitHub request re-read the token from KV -- a wasted subrequest per call. A finalize or
+  // review invocation makes many GitHub calls, so that repeated KV read pushed the invocation toward
+  // the Workers-Free 50-subrequest cap (finalize could tip over it right before posting the review).
+  private memoToken: InstallationTokenCacheRecord | null = null;
+
   async getInstallationToken(): Promise<string> {
+    // Reuse the in-memory token while it's comfortably unexpired (invocations are < ~120s; tokens
+    // last ~1h, so this holds for the whole invocation) -- no KV read, no network call.
+    if (this.memoToken?.token && new Date(this.memoToken.expiresAt).getTime() > Date.now() + 60_000) {
+      return this.memoToken.token;
+    }
+
     const cached = await readCachedInstallationToken(this.env, this.installationId, this.tracker);
     if (cached?.token) {
+      this.memoToken = cached;
       return cached.token;
     }
 
@@ -229,10 +242,12 @@ export class GitHubClient {
       }
 
       const data = (await response.json()) as { token: string; expires_at: string };
-      await writeCachedInstallationToken(this.env, this.installationId, {
+      const record: InstallationTokenCacheRecord = {
         token: data.token,
         expiresAt: data.expires_at,
-      }, this.tracker);
+      };
+      await writeCachedInstallationToken(this.env, this.installationId, record, this.tracker);
+      this.memoToken = record;
 
       return data.token;
     });

src/server/core/job-recovery.ts

@@ -41,17 +41,34 @@ export async function completeTerminalCheckRuns(env: AppBindings) {
 
     try {
       const github = new GitHubService(env, job.installation_id);
-      const checkRunPresentation = {
-        superseded: { conclusion: 'neutral' as const, title: 'Review superseded', summary: 'Superseded by a newer commit or job.' },
-        cancelled: { conclusion: 'cancelled' as const, title: 'Review stopped', summary: 'Stopped by user.' },
-        failed: { conclusion: 'failure' as const, title: 'Review failed', summary: 'Review failed.' },
-      };
-      const presentation = checkRunPresentation[job.status as keyof typeof checkRunPresentation] ?? checkRunPresentation.failed;
+
+      let conclusion: 'success' | 'neutral' | 'failure' | 'cancelled';
+      let title: string;
+      let summary: string;
+      if (job.status === 'done') {
+        // A completed review whose inline check-run update didn't land (e.g. finalize ran out of
+        // subrequest budget). Reconstruct the same conclusion finalize would have posted.
+        const partial = (job.error_msg ?? '').startsWith('Partial review');
+        conclusion = partial ? 'failure' : (job.verdict === 'approve' ? 'success' : 'neutral');
+        title = partial ? 'Review partially failed' : (job.verdict === 'approve' ? 'LGTM' : 'Comments posted');
+        summary = job.error_msg ?? `${job.comment_count ?? 0} inline comments across ${job.file_count ?? 0} files.`;
+      } else {
+        const checkRunPresentation = {
+          superseded: { conclusion: 'neutral' as const, title: 'Review superseded', summary: 'Superseded by a newer commit or job.' },
+          cancelled: { conclusion: 'cancelled' as const, title: 'Review stopped', summary: 'Stopped by user.' },
+          failed: { conclusion: 'failure' as const, title: 'Review failed', summary: 'Review failed.' },
+        };
+        const presentation = checkRunPresentation[job.status as keyof typeof checkRunPresentation] ?? checkRunPresentation.failed;
+        conclusion = presentation.conclusion;
+        title = presentation.title;
+        summary = job.error_msg ?? presentation.summary;
+      }
+
       await github.updateCheckRun(job.owner, job.repo, job.check_run_id, {
         status: 'completed',
-        conclusion: presentation.conclusion,
-        title: presentation.title,
-        summary: job.error_msg ?? presentation.summary,
+        conclusion,
+        title,
+        summary,
       });
       await markJobCheckRunCompleted(env, job.id);
     } catch (error) {

src/server/core/review.ts

@@ -3,7 +3,7 @@ import { isSupportedGitHubWebhookEvent, type GitHubWebhookEventName, type GitHub
 import { defaultRepoConfig, normalizeModelId, type ParsedReviewComment, type RepoConfig, type ReviewJobMessage } from '@shared/schema';
 import { isTimeoutMessage, matchesAnyTransientSubstring } from '@shared/transient-errors';
 import type { AppBindings } from '@server/env';
-import { bulkInheritFileReviews, getFileReviewsForJobs, recordRetryableFileReviewFailure, upsertFileReview } from '@server/db/file-reviews';
+import { bulkInheritFileReviews, bulkMarkFilesFailed, getFileReviewsForJobs, recordRetryableFileReviewFailure, upsertFileReview } from '@server/db/file-reviews';
 import { getResolvedModelConfig } from '@server/db/model-configs';
 import {
   claimJobLease,
@@ -1227,16 +1227,17 @@ async function runFinalizePhase(
 
     if (missingFiles.length > 0) {
       logger.warn(`Job ${job.id} reached finalize phase with ${missingFiles.length} missing file reviews. Forcing them to failed state.`);
-      for (const file of missingFiles) {
-        await persistFailedFileReview(env, job.id, {
-          filePath: file.path,
-          modelUsed: config.model?.main ?? 'unconfigured',
-          diffLineCount: file.lineCount,
-          durationMs: 0,
-          errorMessage: 'Review could not be completed after repeated infrastructure limits',
-        });
-      }
-      
+      // Batch the backfill into one INSERT. Doing it per-file (a transaction each) scales the
+      // subrequest cost with the number of missing files and, on a large/growing PR, exhausts the
+      // per-invocation budget right before the review is posted (finalize can't safely hibernate --
+      // it posts the GitHub review -- so it must stay within one invocation's budget).
+      await bulkMarkFilesFailed(
+        env,
+        job.id,
+        missingFiles.map((file) => ({ filePath: file.path, diffLineCount: file.lineCount })),
+        { modelUsed: config.model?.main ?? 'unconfigured', errorMessage: 'Review could not be completed after repeated infrastructure limits' },
+      );
+
       // Refresh reviews list after inserting the missing ones
       reviews = await getFileReviewsForJobs(env, [job.id]);
     } else {
@@ -1305,34 +1306,6 @@ async function runFinalizePhase(
     })),
   });
 
-  if (config.review.labels !== false) {
-    const labels = config.review.labels;
-    const labelMap = {
-      comment: { name: labels.p1, color: 'f79009' },
-      approve: { name: labels.p2, color: '027a48' },
-    } as const;
-    const label = labelMap[verdictSummary.verdict];
-
-    await github.removeIssueLabelsIfPresent(
-      job.owner,
-      job.repo,
-      job.prNumber,
-      [labels.p1, labels.p2, labels.p3].filter(possibleLabel => possibleLabel !== label.name),
-    );
-
-    await github.ensureLabel(job.owner, job.repo, label.name, label.color);
-    await github.addIssueLabels(job.owner, job.repo, job.prNumber, [label.name]);
-  }
-
-  if (job.checkRunId) {
-    await github.updateCheckRun(job.owner, job.repo, job.checkRunId, {
-      status: 'completed',
-      conclusion: hasFailures ? 'failure' : (verdictSummary.verdict === 'approve' ? 'success' : 'neutral'),
-      title: hasFailures ? 'Review partially failed' : (verdictSummary.verdict === 'approve' ? 'LGTM' : 'Comments posted'),
-      summary: `${finalComments.length} inline comments across ${files.length} files.${hasFailures ? ` ${failedFileCount} file${failedFileCount === 1 ? '' : 's'} could not be reviewed after repeated provider outages.` : ''}`,
-    });
-  }
-
   const fileInputTokens = reviews.reduce((sum, review) => sum + (review.input_tokens ?? 0), 0);
   const fileOutputTokens = reviews.reduce((sum, review) => sum + (review.output_tokens ?? 0), 0);
 
@@ -1345,6 +1318,10 @@ async function runFinalizePhase(
   const partialErrorMessage = hasFailures
     ? `Partial review: ${failedFileCount} of ${files.length} file${files.length === 1 ? '' : 's'} could not be reviewed after repeated model/provider outages.`
     : null;
+  // Record the posted review and mark the job done IMMEDIATELY after createReview -- before the
+  // best-effort cosmetics below. The review is already on GitHub at this point; if the remaining
+  // label/check-run calls exhaust this invocation's subrequest budget (large PR), we must not leave
+  // the job stranded as 'failed' with review_id null. This is the critical, must-not-lose write.
   await completeJob(env, job.id, {
     verdict: verdictSummary.verdict,
     fileCount: files.length,
@@ -1358,6 +1335,47 @@ async function runFinalizePhase(
   });
   logger.info(`Review job completed: ${job.owner}/${job.repo} PR #${job.prNumber}`);
 
+  // Cosmetics: labels and the check-run conclusion. Best-effort -- the review is posted and the job
+  // is already 'done', so a failure here (e.g. subrequest budget spent, GitHub blip) must not fail
+  // the job. completeTerminalCheckRuns / a re-run can reconcile a check run left un-updated.
+  try {
+    // Check-run conclusion first: it drives the PR's status badge, so it matters more than labels
+    // if the budget only allows one of them.
+    if (job.checkRunId) {
+      await github.updateCheckRun(job.owner, job.repo, job.checkRunId, {
+        status: 'completed',
+        conclusion: hasFailures ? 'failure' : (verdictSummary.verdict === 'approve' ? 'success' : 'neutral'),
+        title: hasFailures ? 'Review partially failed' : (verdictSummary.verdict === 'approve' ? 'LGTM' : 'Comments posted'),
+        summary: `${finalComments.length} inline comments across ${files.length} files.${hasFailures ? ` ${failedFileCount} file${failedFileCount === 1 ? '' : 's'} could not be reviewed after repeated provider outages.` : ''}`,
+      });
+      // Only now is the check run genuinely completed -- record it so the maintenance sweep doesn't
+      // redo it. If the update above threw, this line is skipped and completeTerminalCheckRuns will
+      // finish the check run on a later invocation with a fresh budget.
+      await markJobCheckRunCompleted(env, job.id);
+    }
+
+    if (config.review.labels !== false) {
+      const labels = config.review.labels;
+      const labelMap = {
+        comment: { name: labels.p1, color: 'f79009' },
+        approve: { name: labels.p2, color: '027a48' },
+      } as const;
+      const label = labelMap[verdictSummary.verdict];
+
+      await github.removeIssueLabelsIfPresent(
+        job.owner,
+        job.repo,
+        job.prNumber,
+        [labels.p1, labels.p2, labels.p3].filter(possibleLabel => possibleLabel !== label.name),
+      );
+
+      await github.ensureLabel(job.owner, job.repo, label.name, label.color);
+      await github.addIssueLabels(job.owner, job.repo, job.prNumber, [label.name]);
+    }
+  } catch (error) {
+    logger.warn(`Post-review labels/check-run update failed for job ${job.id}; review is posted and job is completed, so leaving it best-effort`, error instanceof Error ? error : new Error(String(error)));
+  }
+
   await sendReviewTelemetry(env, job, files, reviews, {
     findingsReported: finalComments.length,
     verdict: verdictSummary.verdict,

src/server/db/client.ts

@@ -62,8 +62,29 @@ export function runWithDb<T>(env: DbEnv, fn: () => T): T {
   return dbStorage.run(createDbClient(env), fn);
 }
 
+// Reused clients for the no-runWithDb() fallback path below. Keyed by connection string so a caller
+// that queries outside a runWithDb() scope shares one bounded pool instead of leaking a fresh pool
+// (up to `max` connections) on every single query.
+const fallbackClients = new Map<string, DbClient>();
+
 export function getDb(env: DbEnv) {
-  return dbStorage.getStore() ?? createDbClient(env);
+  const store = dbStorage.getStore();
+  if (store) return store;
+
+  // Outside a runWithDb() scope. In production this never happens -- fetch, queue, scheduled, and
+  // the workflow all wrap their work in runWithDb() -- so this branch is effectively test-only (the
+  // Hono test client calls app.fetch() directly, and tests call db helpers directly). Creating a new
+  // pool per query there leaked connections across the whole test process and exhausted CI's Postgres
+  // (masked locally by Neon's pooler). Reuse one client per connection string instead. Keeping this
+  // strictly to the store-less path preserves the per-request client production relies on to satisfy
+  // Cloudflare's "no I/O across request contexts" rule.
+  const connectionString = env.HYPERDRIVE.connectionString;
+  let client = fallbackClients.get(connectionString);
+  if (!client) {
+    client = createDbClient(env);
+    fallbackClients.set(connectionString, client);
+  }
+  return client;
 }
 
 export async function queryRows<T>(env: DbEnv, sqlText: string, params: unknown[] = []) {

src/server/db/file-reviews.ts

@@ -335,6 +335,33 @@ export async function bulkInheritFileReviews(
   });
 }
 
+/**
+ * Mark many files 'failed' in a single INSERT. Finalize backfills reviews for files that never got
+ * one (e.g. files that appeared mid-review, or unrecoverable ones); doing that one-by-one through
+ * upsertFileReview runs a transaction per file (several Hyperdrive round-trips each), which for a
+ * large/growing PR can blow the per-invocation subrequest budget right before the review is posted.
+ * This collapses the whole backfill into one statement (one subrequest). Skips files that already
+ * have a row (ON CONFLICT DO NOTHING) so it never clobbers a real review.
+ */
+export async function bulkMarkFilesFailed(
+  env: Pick<AppBindings, 'HYPERDRIVE'>,
+  jobId: string,
+  files: Array<{ filePath: string; diffLineCount: number }>,
+  opts: { modelUsed: string; errorMessage: string },
+): Promise<void> {
+  if (files.length === 0) return;
+  await queryRows(
+    env,
+    `
+      INSERT INTO file_reviews (job_id, file_path, file_status, model_used, diff_line_count, diff_input, error_msg, duration_ms)
+      SELECT $1::uuid, u.file_path, 'failed', $2, u.diff_line_count, '', $3, 0
+      FROM UNNEST($4::text[], $5::int[]) AS u(file_path, diff_line_count)
+      ON CONFLICT (job_id, file_path) DO NOTHING
+    `,
+    [jobId, opts.modelUsed, opts.errorMessage, files.map((f) => f.filePath), files.map((f) => f.diffLineCount)],
+  );
+}
+
 export async function getModelUsageStats(env: Pick<AppBindings, 'HYPERDRIVE'>) {
   return queryRows<{
     model_used: string;

src/server/db/jobs.ts

@@ -206,7 +206,7 @@ export async function hasPendingMaintenanceWork(env: Pick<AppBindings, 'HYPERDRI
       SELECT EXISTS (
         SELECT 1 FROM jobs
         WHERE status IN ('queued', 'running')
-           OR (status IN ('failed', 'superseded', 'cancelled') AND check_run_id IS NOT NULL AND check_run_completed_at IS NULL)
+           OR (status IN ('done', 'failed', 'superseded', 'cancelled') AND check_run_id IS NOT NULL AND check_run_completed_at IS NULL)
       ) AS has_work
     `,
   );
@@ -629,7 +629,11 @@ export async function completeJob(
       UPDATE jobs
       SET status = 'done',
           finished_at = now(),
-          check_run_completed_at = now(),
+          -- NOTE: check_run_completed_at is intentionally NOT set here. It's set only once the
+          -- GitHub check run has actually been updated (markJobCheckRunCompleted, called after the
+          -- best-effort updateCheckRun in finalize). That way, if finalize couldn't update the
+          -- check run (e.g. subrequest budget spent on a huge PR), it stays NULL and the maintenance
+          -- sweep (completeTerminalCheckRuns) reconciles it -- the check run always ends 'completed'.
           lease_owner = NULL,
           lease_expires_at = NULL,
           verdict = $2,
@@ -1003,7 +1007,7 @@ export async function getTerminalJobsNeedingCheckRunCompletion(
       SELECT j.*, r.owner, r.repo, r.installation_id
       FROM jobs j
       JOIN repositories r ON j.repository_id = r.id
-      WHERE j.status IN ('failed', 'superseded', 'cancelled')
+      WHERE j.status IN ('done', 'failed', 'superseded', 'cancelled')
         AND j.check_run_id IS NOT NULL
         AND j.check_run_completed_at IS NULL
       ORDER BY COALESCE(j.finished_at, j.started_at, j.created_at) ASC