[Fix] PR review notifications are lost during pull request creation

packages/sdk/src/server/lib/pull-requests/source-control-pull-requests.ts

@@ -452,8 +452,9 @@ async function resolveEffectivePrAction(taskRun: TaskRun): Promise<PrAction> {
  * delivery path used to create this association by parsing `gh pr create`
  * tool output from the transcript; the server-side mutation path knows the
  * pull request authoritatively for every provider, so it persists the
- * association directly. Association failures must not fail the mutation the
- * agent already performed.
+ * association directly. Retry transient write failures at this idempotent
+ * boundary, then surface exhaustion so the caller can retry the whole
+ * create-or-update operation without losing the authoritative association.
  */
 async function persistSourceControlPullRequestAssociation({
   taskRun,
@@ -470,40 +471,45 @@ async function persistSourceControlPullRequestAssociation({
 
   const status = result.draft ? 'draft' : 'open';
 
-  try {
-    await db
-      .insert(taskPullRequests)
-      .values({
-        taskId: taskRun.taskId,
-        sourceControlProvider: repository.sourceControlProvider,
-        host: repository.host,
-        repositoryId: repository.id,
-        prUrl: result.url,
-        prNumber: result.number,
-        prTitle: result.title,
-        repository: result.repositoryFullName,
-        status,
-        createdByRoomote: result.action === 'created',
-        prBaseRef: result.targetBranch,
-      })
-      .onConflictDoUpdate({
-        target: [taskPullRequests.taskId, taskPullRequests.prUrl],
-        set: {
+  const maxAttempts = 3;
+
+  for (let attempt = 1; attempt <= maxAttempts; attempt += 1) {
+    try {
+      await db
+        .insert(taskPullRequests)
+        .values({
+          taskId: taskRun.taskId,
           sourceControlProvider: repository.sourceControlProvider,
           host: repository.host,
           repositoryId: repository.id,
+          prUrl: result.url,
+          prNumber: result.number,
           prTitle: result.title,
+          repository: result.repositoryFullName,
           status,
+          createdByRoomote: result.action === 'created',
           prBaseRef: result.targetBranch,
-          updatedAt: new Date(),
-        },
-      });
-  } catch (error) {
-    console.warn(
-      `[persistSourceControlPullRequestAssociation] Failed to associate ${result.repositoryFullName}#${result.number} with task ${taskRun.taskId}: ${
-        error instanceof Error ? error.message : String(error)
-      }`,
-    );
+        })
+        .onConflictDoUpdate({
+          target: [taskPullRequests.taskId, taskPullRequests.prUrl],
+          set: {
+            sourceControlProvider: repository.sourceControlProvider,
+            host: repository.host,
+            repositoryId: repository.id,
+            prTitle: result.title,
+            status,
+            prBaseRef: result.targetBranch,
+            updatedAt: new Date(),
+          },
+        });
+      return;
+    } catch (error) {
+      if (attempt === maxAttempts) {
+        throw error;
+      }
+
+      await new Promise((resolve) => setTimeout(resolve, attempt * 100));
+    }
   }
 }
 

packages/sdk/src/server/lib/task-runs/pr-review-notification.ts

@@ -50,6 +50,8 @@ export const PR_REVIEW_NOTIFICATION_MAX_DEFERRALS = 288;
 const PENDING_EVENTS_TTL_SECONDS = 24 * 60 * 60;
 const SCHEDULED_MARKER_TTL_BUFFER_SECONDS = 15 * 60;
 const REVIEW_CYCLE_TTL_SECONDS = 30 * 24 * 60 * 60;
+const PR_ASSOCIATION_LOOKUP_MAX_ATTEMPTS = 4;
+const PR_ASSOCIATION_LOOKUP_RETRY_DELAY_MS = 250;
 const SET_REVIEW_CYCLE_IF_NEWER_SCRIPT = `
 local current = redis.call('GET', KEYS[1])
 if current then
@@ -664,17 +666,36 @@ export async function enqueuePrReviewNotification(
 ): Promise<EnqueuePrReviewNotificationResult> {
   const parsedInput = enqueuePrReviewNotificationInputSchema.parse(input);
 
-  const prTaskLinks = await db.query.taskPullRequests.findMany({
-    where: and(
-      eq(
-        taskPullRequests.sourceControlProvider,
-        parsedInput.sourceControlProvider ?? 'github',
+  let prTaskLinks: Array<{ taskId: string }> = [];
+
+  for (
+    let attempt = 1;
+    attempt <= PR_ASSOCIATION_LOOKUP_MAX_ATTEMPTS;
+    attempt += 1
+  ) {
+    prTaskLinks = await db.query.taskPullRequests.findMany({
+      where: and(
+        eq(
+          taskPullRequests.sourceControlProvider,
+          parsedInput.sourceControlProvider ?? 'github',
+        ),
+        eq(taskPullRequests.repository, parsedInput.repository),
+        eq(taskPullRequests.prNumber, parsedInput.prNumber),
       ),
-      eq(taskPullRequests.repository, parsedInput.repository),
-      eq(taskPullRequests.prNumber, parsedInput.prNumber),
-    ),
-    columns: { taskId: true },
-  });
+      columns: { taskId: true },
+    });
+
+    if (
+      prTaskLinks.length > 0 ||
+      attempt === PR_ASSOCIATION_LOOKUP_MAX_ATTEMPTS
+    ) {
+      break;
+    }
+
+    await new Promise((resolve) =>
+      setTimeout(resolve, PR_ASSOCIATION_LOOKUP_RETRY_DELAY_MS),
+    );
+  }
 
   const taskIds = Array.from(new Set(prTaskLinks.map((link) => link.taskId)));