diff --git a/src/lib/api.test.ts b/src/lib/api.test.ts index 5ef1179..824e330 100644 --- a/src/lib/api.test.ts +++ b/src/lib/api.test.ts @@ -104,28 +104,21 @@ describe('ApiClient', () => { expect(callBody.repositoryName).toBeUndefined(); }); - it('decodes base64 generalComment in response', async () => { - const encoded = Buffer.from('Good code').toString('base64'); - mockOkResponse({ generalComment: encoded }); + it('sets async true in the request body', async () => { + mockOkResponse({ reviewId: 'apirev_x' }); + await client.review({ patch: 'x' }); - const result = await client.review({ patch: 'x' }); - expect(result.generalComment).toBe('Good code'); + const callBody = JSON.parse(fetchMock.mock.calls[0][1].body); + expect(callBody.async).toBe(true); }); - it('decodes base64 fileComments array in response', async () => { - const c1 = Buffer.from('comment 1').toString('base64'); - const c2 = Buffer.from('comment 2').toString('base64'); - mockOkResponse({ fileComments: [c1, c2] }); - - const result = await client.review({ patch: 'x' }); - expect(result.fileComments).toEqual(['comment 1', 'comment 2']); - }); + it('returns the reviewId and reviewCount from the 202 receipt', async () => { + const reviewCount = { current: 5, limit: 100, remaining: 95 }; + mockOkResponse({ reviewId: 'apirev_abc', reviewCount }); - it('handles response with no generalComment or fileComments', async () => { - mockOkResponse({}); const result = await client.review({ patch: 'x' }); - expect(result.generalComment).toBeUndefined(); - expect(result.fileComments).toBeUndefined(); + expect(result.reviewId).toBe('apirev_abc'); + expect(result.reviewCount).toEqual(reviewCount); }); it('passes through reviewCount unchanged', async () => { @@ -296,21 +289,51 @@ describe('ApiClient', () => { }); }); - describe('review with reviewSessionId', () => { - it('includes reviewSessionId in body when provided', async () => { - mockOkResponse({}); - await client.review({ patch: 'x', reviewSessionId: 'sess-123' }); + describe('getReviewResult', () => { + it('GETs /api/review/result/:reviewId with auth', async () => { + mockOkResponse({ status: 'pending' }); + await client.getReviewResult('apirev_abc'); - const callBody = JSON.parse(fetchMock.mock.calls[0][1].body); - expect(callBody.reviewSessionId).toBe('sess-123'); + expect(fetchMock).toHaveBeenCalledWith( + 'http://test-api.local/api/review/result/apirev_abc', + expect.objectContaining({ + method: 'GET', + headers: expect.objectContaining({ 'Authorization': 'Bearer test-api-key' }), + }) + ); }); - it('omits reviewSessionId when not provided', async () => { - mockOkResponse({}); - await client.review({ patch: 'x' }); + it('returns pending while the review runs', async () => { + mockOkResponse({ status: 'pending' }); + expect(await client.getReviewResult('apirev_abc')).toEqual({ status: 'pending' }); + }); - const callBody = JSON.parse(fetchMock.mock.calls[0][1].body); - expect(callBody.reviewSessionId).toBeUndefined(); + it('decodes base64 comments on a done result', async () => { + const gc = Buffer.from('All good').toString('base64'); + const fc = Buffer.from('line 1').toString('base64'); + mockOkResponse({ status: 'done', generalComment: gc, fileComments: [fc], isOptibotInstalled: true }); + + expect(await client.getReviewResult('apirev_abc')).toEqual({ + status: 'done', + generalComment: 'All good', + fileComments: ['line 1'], + isOptibotInstalled: true, + }); + }); + + it('returns failed with the error message', async () => { + mockOkResponse({ status: 'failed', error: 'boom' }); + expect(await client.getReviewResult('apirev_abc')).toEqual({ status: 'failed', error: 'boom' }); + }); + + it('returns not_found on a 404', async () => { + mockErrorResponse(404, { status: 'not_found' }); + expect(await client.getReviewResult('apirev_abc')).toEqual({ status: 'not_found' }); + }); + + it('throws on a non-404 error', async () => { + mockErrorResponse(401, { message: 'Unauthorized' }); + await expect(client.getReviewResult('apirev_abc')).rejects.toThrow('Unauthorized'); }); }); diff --git a/src/lib/api.ts b/src/lib/api.ts index a71e993..192fe09 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -1,5 +1,6 @@ import { - ReviewResponse, + AsyncReviewResponse, + ReviewResultResponse, ApiKeyCreateResponse, ApiKeyListItem, ApiKeyListResponse, @@ -44,24 +45,25 @@ export class ApiClient { throw error; } + /** + * Request a review asynchronously. Returns a receipt ({ reviewId }) rather + * than the review itself: the backend enqueues the work and the result is + * fetched separately via {@link getReviewResult}. Requires a backend with + * async review support deployed. + */ async review(params: { patch: string; repositoryName?: string; files?: Record; - reviewSessionId?: string; - }): Promise { + }): Promise { const patchBase64 = Buffer.from(params.patch, 'utf-8').toString('base64'); - const body: Record = { patch: patchBase64 }; + const body: Record = { patch: patchBase64, async: true }; if (params.repositoryName) { body.repositoryName = params.repositoryName; } - if (params.reviewSessionId) { - body.reviewSessionId = params.reviewSessionId; - } - if (params.files && Object.keys(params.files).length > 0) { // Null-prototype map: filenames from a repo are untrusted input // (a file literally named `__proto__` would otherwise pollute). @@ -86,20 +88,49 @@ export class ApiClient { await this.throwApiError(response); } - const result = await response.json() as ReviewResponse; + return await response.json() as AsyncReviewResponse; + } + + /** + * Fetch the result of an async review by its reviewId. Returns `pending` + * while the review runs, `done` with the decoded payload when it finishes, + * `failed` with an error, or `not_found` (HTTP 404) when the reviewId is + * unknown or expired. + */ + async getReviewResult(reviewId: string): Promise { + const response = await fetch( + `${API_BASE_URL}/api/review/result/${encodeURIComponent(reviewId)}`, + { + method: 'GET', + headers: { + ...CLIENT_HEADERS, + 'Authorization': `Bearer ${this.apiKey}`, + }, + } + ); + + if (response.status === 404) { + return { status: 'not_found' }; + } - // Decode base64-encoded fields - if (result.generalComment && typeof result.generalComment === 'string') { - result.generalComment = Buffer.from(result.generalComment, 'base64').toString('utf-8'); + if (!response.ok) { + await this.throwApiError(response); } - if (result.fileComments && Array.isArray(result.fileComments)) { - result.fileComments = result.fileComments.map((comment: string) => { - if (typeof comment === 'string') { - return Buffer.from(comment, 'base64').toString('utf-8'); - } - return comment; - }); + const result = await response.json() as ReviewResultResponse; + + // Decode base64-encoded comments on a completed review. + if (result.status === 'done') { + if (typeof result.generalComment === 'string') { + result.generalComment = Buffer.from(result.generalComment, 'base64').toString('utf-8'); + } + if (Array.isArray(result.fileComments)) { + result.fileComments = result.fileComments.map((comment) => + typeof comment === 'string' + ? Buffer.from(comment, 'base64').toString('utf-8') + : comment + ); + } } return result; diff --git a/src/lib/reviewPolling.test.ts b/src/lib/reviewPolling.test.ts new file mode 100644 index 0000000..3715950 --- /dev/null +++ b/src/lib/reviewPolling.test.ts @@ -0,0 +1,62 @@ +import { describe, it, expect, vi } from 'vitest'; +import { waitForReviewResult, ReviewResultFetcher } from './reviewPolling.js'; +import { ReviewResultResponse } from '../types.js'; + +const fetcher = (results: ReviewResultResponse[]): ReviewResultFetcher => { + const queue = [...results]; + return { + getReviewResult: vi.fn(async () => queue.shift() ?? { status: 'pending' }), + }; +}; + +// No real waiting in tests. +const noSleep = () => Promise.resolve(); + +describe('waitForReviewResult', () => { + it('returns immediately when the first poll is done', async () => { + const done: ReviewResultResponse = { status: 'done', generalComment: 'ok' }; + const client = fetcher([done]); + + const result = await waitForReviewResult(client, 'apirev_x', { sleep: noSleep }); + + expect(result).toEqual(done); + expect(client.getReviewResult).toHaveBeenCalledTimes(1); + }); + + it('returns the failed result', async () => { + const client = fetcher([{ status: 'failed', error: 'boom' }]); + const result = await waitForReviewResult(client, 'apirev_x', { sleep: noSleep }); + expect(result).toEqual({ status: 'failed', error: 'boom' }); + }); + + it('keeps polling through pending and not_found until done', async () => { + const client = fetcher([ + { status: 'not_found' }, + { status: 'pending' }, + { status: 'pending' }, + { status: 'done', generalComment: 'ok' }, + ]); + + const result = await waitForReviewResult(client, 'apirev_x', { sleep: noSleep }); + + expect(result).toEqual({ status: 'done', generalComment: 'ok' }); + expect(client.getReviewResult).toHaveBeenCalledTimes(4); + }); + + it('throws when it exceeds the timeout', async () => { + // Always pending; a clock that jumps past the timeout after the first poll. + const client: ReviewResultFetcher = { + getReviewResult: vi.fn(async () => ({ status: 'pending' })), + }; + let t = 0; + const now = () => { + const value = t; + t += 1000; // advance 1s per read + return value; + }; + + await expect( + waitForReviewResult(client, 'apirev_x', { timeoutMs: 1500, sleep: noSleep, now }), + ).rejects.toThrow('Timed out'); + }); +}); diff --git a/src/lib/reviewPolling.ts b/src/lib/reviewPolling.ts new file mode 100644 index 0000000..aa9c10a --- /dev/null +++ b/src/lib/reviewPolling.ts @@ -0,0 +1,56 @@ +import { ReviewResultResponse } from '../types.js'; + +/** Terminal states of a review. */ +export type TerminalReviewResult = Extract< + ReviewResultResponse, + { status: 'done' } | { status: 'failed' } +>; + +/** Minimal surface needed to poll for a review result (the ApiClient satisfies it). */ +export interface ReviewResultFetcher { + getReviewResult(reviewId: string): Promise; +} + +export interface WaitForReviewResultOptions { + /** Give up after this long. Default 20 minutes. */ + timeoutMs?: number; + /** Delay between polls. Default 3 seconds. */ + intervalMs?: number; + /** Injectable clock, for tests. */ + now?: () => number; + /** Injectable delay, for tests. */ + sleep?: (ms: number) => Promise; +} + +/** + * Poll for an async review's result until it reaches a terminal state. + * + * `pending` and the brief post-enqueue `not_found` race both mean "keep + * waiting": the job was accepted (the POST returned a reviewId), so the result + * is on its way. Throws on timeout. + */ +export async function waitForReviewResult( + client: ReviewResultFetcher, + reviewId: string, + options: WaitForReviewResultOptions = {}, +): Promise { + const timeoutMs = options.timeoutMs ?? 20 * 60 * 1000; + const intervalMs = options.intervalMs ?? 3000; + const now = options.now ?? (() => Date.now()); + const sleep = options.sleep ?? ((ms) => new Promise((resolve) => setTimeout(resolve, ms))); + + const start = now(); + + for (;;) { + const result = await client.getReviewResult(reviewId); + if (result.status === 'done' || result.status === 'failed') { + return result; + } + + if (now() - start >= timeoutMs) { + throw new Error('Timed out waiting for the review result.'); + } + + await sleep(intervalMs); + } +} diff --git a/src/lib/reviewProgress.ts b/src/lib/reviewProgress.ts index 640ddbe..79ea1d2 100644 --- a/src/lib/reviewProgress.ts +++ b/src/lib/reviewProgress.ts @@ -27,10 +27,16 @@ export class ReviewProgressService { private progressCallback: ProgressCallback | null = null; private connectionTimeout: NodeJS.Timeout | null = null; - async startSession(onProgress: ProgressCallback): Promise { + /** + * Open the progress socket and join a review's room. Pass `sessionId` to join + * a known room — for async reviews this is the server's reviewId, so progress + * events reach this client (handshake "Option A"). When omitted, a random id + * is generated (legacy behavior). + */ + async startSession(onProgress: ProgressCallback, sessionId?: string): Promise { this.cleanup(); - this.sessionId = randomUUID(); + this.sessionId = sessionId ?? randomUUID(); this.progressCallback = onProgress; return new Promise((resolve) => { diff --git a/src/tools/review.test.ts b/src/tools/review.test.ts index 8d72e37..cf33836 100644 --- a/src/tools/review.test.ts +++ b/src/tools/review.test.ts @@ -15,6 +15,7 @@ const mockCheckMergeConflicts = vi.fn(); const mockFormatReview = vi.fn(); const mockFormatError = vi.fn(); const mockApiReview = vi.fn(); +const mockGetReviewResult = vi.fn(); const mockStartSession = vi.fn(); const mockEndSession = vi.fn(); @@ -37,6 +38,7 @@ vi.mock('../lib/git.js', () => ({ vi.mock('../lib/api.js', () => ({ ApiClient: class { review(...args: any[]) { return mockApiReview(...args); } + getReviewResult(...args: any[]) { return mockGetReviewResult(...args); } }, })); @@ -82,6 +84,8 @@ describe('review tools', () => { beforeEach(() => { mockStartSession.mockResolvedValue('session-id-123'); + mockApiReview.mockResolvedValue({ reviewId: 'apirev_test', reviewCount: { current: 1, limit: 5, remaining: 4 } }); + mockGetReviewResult.mockResolvedValue({ status: 'done', generalComment: 'Good' }); }); describe('review_local_changes', () => { @@ -127,7 +131,7 @@ describe('review tools', () => { expect(result.content[0].text).toBe('Auth error'); }); - it('ends progress session even when review fails', async () => { + it('returns an error when the submit fails', async () => { mockReadConfig.mockResolvedValue({ apiKey: 'key' }); mockGetRepoRoot.mockResolvedValue('/repo'); mockGetRepoName.mockResolvedValue('my-repo'); @@ -140,26 +144,47 @@ describe('review tools', () => { const handler = registeredTools.get('review_local_changes')!; const result = await handler(mockExtra); + expect(result.isError).toBe(true); + }); + + it('ends the progress session when the review result fails', async () => { + mockReadConfig.mockResolvedValue({ apiKey: 'key' }); + mockGetRepoRoot.mockResolvedValue('/repo'); + mockGetRepoName.mockResolvedValue('my-repo'); + mockGetDiffHead.mockResolvedValue('diff'); + mockGetChangedFiles.mockResolvedValue([]); + mockGetFileContents.mockResolvedValue({}); + mockApiReview.mockResolvedValue({ reviewId: 'apirev_test' }); + mockGetReviewResult.mockResolvedValue({ status: 'failed', error: 'boom' }); + mockFormatError.mockReturnValue('Review error'); + + const handler = registeredTools.get('review_local_changes')!; + const result = await handler(mockExtra); + expect(result.isError).toBe(true); expect(mockEndSession).toHaveBeenCalled(); }); - it('passes reviewSessionId to API client', async () => { + it('joins the progress room by the reviewId and fetches by it (Option A)', async () => { mockReadConfig.mockResolvedValue({ apiKey: 'key' }); mockGetRepoRoot.mockResolvedValue('/repo'); mockGetRepoName.mockResolvedValue('my-repo'); mockGetDiffHead.mockResolvedValue('diff'); mockGetChangedFiles.mockResolvedValue([]); mockGetFileContents.mockResolvedValue({}); - mockApiReview.mockResolvedValue({}); + mockApiReview.mockResolvedValue({ reviewId: 'apirev_test' }); + mockGetReviewResult.mockResolvedValue({ status: 'done', generalComment: 'Good' }); mockFormatReview.mockReturnValue('Review'); const handler = registeredTools.get('review_local_changes')!; await handler(mockExtra); + // review() no longer carries a session id; the socket joins by reviewId. expect(mockApiReview).toHaveBeenCalledWith( - expect.objectContaining({ reviewSessionId: 'session-id-123' }) + expect.not.objectContaining({ reviewSessionId: expect.anything() }) ); + expect(mockStartSession).toHaveBeenCalledWith(expect.any(Function), 'apirev_test'); + expect(mockGetReviewResult).toHaveBeenCalledWith('apirev_test'); }); }); diff --git a/src/tools/review.ts b/src/tools/review.ts index 5778d89..a10d75c 100644 --- a/src/tools/review.ts +++ b/src/tools/review.ts @@ -5,6 +5,7 @@ import * as git from '../lib/git.js'; import { ApiClient } from '../lib/api.js'; import { formatReview, formatError } from '../lib/output.js'; import { ReviewProgressService, ReviewProgressEvent } from '../lib/reviewProgress.js'; +import { waitForReviewResult } from '../lib/reviewPolling.js'; import { safeSendLog } from '../lib/notify.js'; const ReviewBranchSchema = { @@ -54,15 +55,25 @@ export function registerReviewTools(server: McpServer): void { const changedFiles = await git.getChangedFiles(repoRoot); const files = await git.getFileContents(changedFiles, repoRoot); - // Start progress session + // Submit async, then watch progress by reviewId and poll for the result. + const client = new ApiClient(config.apiKey); + const receipt = await client.review({ patch, repositoryName: repoName, files }); + const progressService = new ReviewProgressService(); - const reviewSessionId = await progressService.startSession((event: ReviewProgressEvent) => { + await progressService.startSession((event: ReviewProgressEvent) => { safeSendLog(extra, 'optibot', formatProgressStep(event)); - }); + }, receipt.reviewId); try { - const client = new ApiClient(config.apiKey); - const response = await client.review({ patch, repositoryName: repoName, files, reviewSessionId }); + const result = await waitForReviewResult(client, receipt.reviewId); + if (result.status === 'failed') { + throw new Error(result.error || 'The review failed. Please try again.'); + } + const response = { + generalComment: result.generalComment, + fileComments: result.fileComments, + reviewCount: receipt.reviewCount, + }; return { content: [{ type: 'text' as const, text: formatReview(response) }] }; } finally { progressService.endSession(); @@ -103,24 +114,34 @@ export function registerReviewTools(server: McpServer): void { const changedFiles = await git.getChangedFiles(repoRoot, targetBranch); const files = await git.getFileContents(changedFiles, repoRoot); - // Start progress session + // Submit async, then watch progress by reviewId and poll for the result. + const client = new ApiClient(config.apiKey); + const receipt = await client.review({ patch, repositoryName: repoName, files }); + const progressService = new ReviewProgressService(); - const reviewSessionId = await progressService.startSession((event: ReviewProgressEvent) => { + await progressService.startSession((event: ReviewProgressEvent) => { safeSendLog(extra, 'optibot', formatProgressStep(event)); - }); + }, receipt.reviewId); try { - const client = new ApiClient(config.apiKey); - const response = await client.review({ patch, repositoryName: repoName, files, reviewSessionId }); + const reviewResult = await waitForReviewResult(client, receipt.reviewId); + if (reviewResult.status === 'failed') { + throw new Error(reviewResult.error || 'The review failed. Please try again.'); + } + const response = { + generalComment: reviewResult.generalComment, + fileComments: reviewResult.fileComments, + reviewCount: receipt.reviewCount, + }; - let result = ''; + let output = ''; if (hasConflicts) { - result += `**Warning:** Merge conflicts detected with ${targetBranch}. Review results may not reflect the final merged state.\n\n`; + output += `**Warning:** Merge conflicts detected with ${targetBranch}. Review results may not reflect the final merged state.\n\n`; } - result += `Comparing against: ${targetBranch}\n\n`; - result += formatReview(response); + output += `Comparing against: ${targetBranch}\n\n`; + output += formatReview(response); - return { content: [{ type: 'text' as const, text: result }] }; + return { content: [{ type: 'text' as const, text: output }] }; } finally { progressService.endSession(); } @@ -145,15 +166,25 @@ export function registerReviewTools(server: McpServer): void { return { content: [{ type: 'text' as const, text: 'The diff file is empty. Nothing to review.' }] }; } - // Start progress session + // Submit async, then watch progress by reviewId and poll for the result. + const client = new ApiClient(config.apiKey); + const receipt = await client.review({ patch }); + const progressService = new ReviewProgressService(); - const reviewSessionId = await progressService.startSession((event: ReviewProgressEvent) => { + await progressService.startSession((event: ReviewProgressEvent) => { safeSendLog(extra, 'optibot', formatProgressStep(event)); - }); + }, receipt.reviewId); try { - const client = new ApiClient(config.apiKey); - const response = await client.review({ patch, reviewSessionId }); + const result = await waitForReviewResult(client, receipt.reviewId); + if (result.status === 'failed') { + throw new Error(result.error || 'The review failed. Please try again.'); + } + const response = { + generalComment: result.generalComment, + fileComments: result.fileComments, + reviewCount: receipt.reviewCount, + }; return { content: [{ type: 'text' as const, text: formatReview(response) }] }; } finally { progressService.endSession(); diff --git a/src/types.ts b/src/types.ts index ec89f90..b2bf51b 100644 --- a/src/types.ts +++ b/src/types.ts @@ -20,6 +20,24 @@ export interface ReviewResponse { }; } +/** Receipt returned by an async `POST /api/review` ({ async: true }). */ +export interface AsyncReviewResponse { + reviewId: string; + reviewCount?: ReviewResponse['reviewCount']; +} + +/** Result of polling `GET /api/review/result/:reviewId`. */ +export type ReviewResultResponse = + | { status: 'pending' } + | { status: 'not_found' } + | { status: 'failed'; error?: string } + | { + status: 'done'; + generalComment?: string; + fileComments?: string[]; + isOptibotInstalled?: boolean; + }; + export interface ParsedFileComment { filePath: string; startLine: number;