diff --git a/packages/sdk/src/server/lib/pull-requests/__tests__/source-control-pull-requests.test.ts b/packages/sdk/src/server/lib/pull-requests/__tests__/source-control-pull-requests.test.ts index 55fcea6cf..c2570289e 100644 --- a/packages/sdk/src/server/lib/pull-requests/__tests__/source-control-pull-requests.test.ts +++ b/packages/sdk/src/server/lib/pull-requests/__tests__/source-control-pull-requests.test.ts @@ -110,6 +110,20 @@ const { mockTaskPullRequestUpsert, mockTaskRunAssociationUpdate } = vi.hoisted( }), ); +const { + mockFindManyTaskPullRequests, + mockRedisSet, + mockRedisMultiExec, + mockQueueAdd, + redisMultiCalls, +} = vi.hoisted(() => ({ + mockFindManyTaskPullRequests: vi.fn(), + mockRedisSet: vi.fn(), + mockRedisMultiExec: vi.fn(), + mockQueueAdd: vi.fn(), + redisMultiCalls: [] as Array<{ command: string; args: unknown[] }>, +})); + vi.mock('@roomote/db/server', () => ({ getDeploymentGitHubRoomoteMentionEnabled: (...args: unknown[]) => mockGetDeploymentGitHubRoomoteMentionEnabled(...args), @@ -133,6 +147,9 @@ vi.mock('@roomote/db/server', () => ({ tasks: { findFirst: (...args: unknown[]) => mockTasksFindFirst(...args), }, + taskPullRequests: { + findMany: (...args: unknown[]) => mockFindManyTaskPullRequests(...args), + }, }, insert: () => ({ values: (values: unknown) => ({ @@ -177,7 +194,31 @@ vi.mock('@roomote/db/server', () => ({ eq: vi.fn((left: unknown, right: unknown) => ({ type: 'eq', left, right })), })); +vi.mock('@roomote/redis', () => ({ + getRedis: () => ({ + set: (...args: unknown[]) => mockRedisSet(...args), + multi: () => { + const multi: Record = {}; + for (const command of ['rpush', 'expire']) { + multi[command] = (...args: unknown[]) => { + redisMultiCalls.push({ command, args }); + return multi; + }; + } + multi.exec = (...args: unknown[]) => mockRedisMultiExec(...args); + return multi; + }, + }), +})); + +vi.mock('bullmq', () => ({ + Queue: class MockQueue { + add = (...args: unknown[]) => mockQueueAdd(...args); + }, +})); + import { createOrUpdateSourceControlPullRequestForTaskRun } from '../source-control-pull-requests'; +import { enqueuePrReviewNotification } from '../../task-runs/pr-review-notification'; function makeTaskRun(payload: TaskRun['payload']): TaskRun { return { @@ -795,6 +836,12 @@ describe('platform-managed draft state', () => { describe('optional targetBranch', () => { beforeEach(() => { vi.clearAllMocks(); + redisMultiCalls.length = 0; + mockTaskPullRequestUpsert.mockResolvedValue(undefined); + mockFindManyTaskPullRequests.mockResolvedValue([]); + mockRedisSet.mockResolvedValue('OK'); + mockRedisMultiExec.mockResolvedValue([]); + mockQueueAdd.mockResolvedValue(undefined); mockResolveConfiguredGitHubAppSlugIfConfigured.mockResolvedValue(null); mockResolveRunCommitAuthor.mockResolvedValue({ kind: 'roomote', @@ -908,6 +955,129 @@ describe('optional targetBranch', () => { ); }); + it('surfaces an association write failure after the bounded retry budget', async () => { + vi.useFakeTimers(); + makeOctokit({ + list: [], + created: { + number: 13, + node_id: 'node-13', + html_url: 'https://github.com/acme/web/pull/13', + title: '[Feature] X', + draft: true, + base: { ref: 'develop' }, + }, + }); + mockTaskPullRequestUpsert.mockRejectedValue( + new Error('association unavailable'), + ); + + const resultPromise = createOrUpdateSourceControlPullRequestForTaskRun({ + taskRun: makeTaskRun({ repo: 'acme/web' }), + input: { ...baseInput, targetBranch: 'develop' }, + }); + const rejection = expect(resultPromise).rejects.toThrow( + 'association unavailable', + ); + await vi.runAllTimersAsync(); + + await rejection; + expect(mockTaskPullRequestUpsert).toHaveBeenCalledTimes(3); + vi.useRealTimers(); + }); + + it('recovers from a transient association write failure', async () => { + vi.useFakeTimers(); + makeOctokit({ + list: [], + created: { + number: 13, + node_id: 'node-13', + html_url: 'https://github.com/acme/web/pull/13', + title: '[Feature] X', + draft: true, + base: { ref: 'develop' }, + }, + }); + mockTaskPullRequestUpsert + .mockRejectedValueOnce(new Error('association unavailable')) + .mockResolvedValueOnce(undefined); + + const resultPromise = createOrUpdateSourceControlPullRequestForTaskRun({ + taskRun: makeTaskRun({ repo: 'acme/web' }), + input: { ...baseInput, targetBranch: 'develop' }, + }); + await vi.runAllTimersAsync(); + + await expect(resultPromise).resolves.toMatchObject({ + action: 'created', + number: 13, + }); + expect(mockTaskPullRequestUpsert).toHaveBeenCalledTimes(2); + vi.useRealTimers(); + }); + + it('schedules one review notification when the association commits during persistence', async () => { + vi.useFakeTimers(); + let finishAssociation: () => void; + const associationCanCommit = new Promise((resolve) => { + finishAssociation = resolve; + }); + let associationVisible = false; + mockTaskPullRequestUpsert.mockImplementationOnce(async () => { + await associationCanCommit; + associationVisible = true; + }); + mockFindManyTaskPullRequests.mockImplementation(async () => + associationVisible ? [{ taskId: 'task-123' }] : [], + ); + makeOctokit({ + list: [], + created: { + number: 13, + node_id: 'node-13', + html_url: 'https://github.com/acme/web/pull/13', + title: '[Feature] X', + draft: true, + base: { ref: 'develop' }, + }, + }); + + const mutationPromise = createOrUpdateSourceControlPullRequestForTaskRun({ + taskRun: makeTaskRun({ repo: 'acme/web' }), + input: { ...baseInput, targetBranch: 'develop' }, + }); + await vi.waitFor(() => { + expect(mockTaskPullRequestUpsert).toHaveBeenCalledTimes(1); + }); + + const notificationPromise = enqueuePrReviewNotification({ + repository: 'acme/web', + prNumber: 13, + prUrl: 'https://github.com/acme/web/pull/13', + event: { + kind: 'review', + authorLogin: 'reviewer', + reviewState: 'changes_requested', + }, + }); + await vi.waitFor(() => { + expect(mockFindManyTaskPullRequests).toHaveBeenCalledTimes(1); + }); + finishAssociation!(); + await mutationPromise; + await vi.runAllTimersAsync(); + + await expect(notificationPromise).resolves.toEqual({ + notifiedTaskCount: 1, + }); + expect(mockQueueAdd).toHaveBeenCalledTimes(1); + expect( + redisMultiCalls.filter((call) => call.command === 'rpush'), + ).toHaveLength(1); + vi.useRealTimers(); + }); + it('scopes the lookup by base and updates without retargeting when targetBranch is explicit', async () => { // The base-scoped lookup can only return a pull request that already // targets the requested base, so the update never sends base. diff --git a/packages/sdk/src/server/lib/pull-requests/source-control-pull-requests.ts b/packages/sdk/src/server/lib/pull-requests/source-control-pull-requests.ts index 85a4380b9..e7fe47012 100644 --- a/packages/sdk/src/server/lib/pull-requests/source-control-pull-requests.ts +++ b/packages/sdk/src/server/lib/pull-requests/source-control-pull-requests.ts @@ -452,8 +452,9 @@ async function resolveEffectivePrAction(taskRun: TaskRun): Promise { * 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)); + } } } diff --git a/packages/sdk/src/server/lib/task-runs/__tests__/pr-review-notification.test.ts b/packages/sdk/src/server/lib/task-runs/__tests__/pr-review-notification.test.ts index c6b7c5e70..33afe73e4 100644 --- a/packages/sdk/src/server/lib/task-runs/__tests__/pr-review-notification.test.ts +++ b/packages/sdk/src/server/lib/task-runs/__tests__/pr-review-notification.test.ts @@ -98,13 +98,18 @@ describe('enqueuePrReviewNotification', () => { }, }; - it('returns no_linked_tasks when no task links the PR', async () => { + it('returns no_linked_tasks after a bounded lookup when no task links the PR', async () => { + vi.useFakeTimers(); mockFindManyTaskPullRequests.mockResolvedValue([]); - const result = await enqueuePrReviewNotification(baseInput); + const resultPromise = enqueuePrReviewNotification(baseInput); + await vi.runAllTimersAsync(); + const result = await resultPromise; expect(result).toEqual({ notifiedTaskCount: 0, reason: 'no_linked_tasks' }); + expect(mockFindManyTaskPullRequests).toHaveBeenCalledTimes(4); expect(mockQueueAdd).not.toHaveBeenCalled(); + vi.useRealTimers(); }); it('debounces ordinary notifications for web-only tasks without an originating conversation', async () => { diff --git a/packages/sdk/src/server/lib/task-runs/pr-review-notification.ts b/packages/sdk/src/server/lib/task-runs/pr-review-notification.ts index f606b21ab..77dd4b05f 100644 --- a/packages/sdk/src/server/lib/task-runs/pr-review-notification.ts +++ b/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 { 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)));