feat(retry): ✨ Add model-switch retry for task/page flows (#63)

* feat(retry): add model switch for task/page retry flows

* fix(renderer): 🐛 resolve typecheck errors in retry model flow

* fix(retry): 🐛 harden model parsing and review tests

Author: jorben

package.json

@@ -32,7 +32,7 @@
     "migrate:reset": "prisma migrate reset --schema=./src/core/infrastructure/db/schema.prisma",
     "generate": "prisma generate --schema=./src/core/infrastructure/db/schema.prisma",
     "logo": "electron-icon-builder --input=./src/renderer/assets/logo.png --output=./public",
-    "test": "vitest",
+    "test": "vitest run",
     "test:unit": "vitest run --config vitest.config.ts",
     "test:renderer": "vitest run --config vitest.config.renderer.ts",
     "test:watch": "vitest watch",

src/core/application/workers/ConverterWorker.ts

@@ -370,7 +370,14 @@ export class ConverterWorker extends WorkerBase {
             // Step 2: Check task is not cancelled
             const task = await tx.task.findUnique({
               where: { id: page.task },
-              select: { status: true, pages: true, completed_count: true, failed_count: true },
+              select: {
+                status: true,
+                pages: true,
+                completed_count: true,
+                failed_count: true,
+                provider: true,
+                model: true,
+              },
             });
 
             if (!task) {
@@ -393,6 +400,8 @@ export class ConverterWorker extends WorkerBase {
                 completed_at: new Date(),
                 worker_id: null, // Release worker
                 error: null,
+                provider: task.provider,
+                model: task.model,
               },
             });
 
@@ -475,7 +484,14 @@ export class ConverterWorker extends WorkerBase {
             // Step 2: Check task is not cancelled
             const task = await tx.task.findUnique({
               where: { id: page.task },
-              select: { status: true, pages: true, completed_count: true, failed_count: true },
+              select: {
+                status: true,
+                pages: true,
+                completed_count: true,
+                failed_count: true,
+                provider: true,
+                model: true,
+              },
             });
 
             if (!task) {
@@ -494,6 +510,8 @@ export class ConverterWorker extends WorkerBase {
                 error: errorMessage,
                 completed_at: new Date(),
                 worker_id: null, // Release worker
+                provider: task.provider,
+                model: task.model,
               },
             });
 

src/core/application/workers/__tests__/ConverterWorker.test.ts

@@ -585,21 +585,27 @@ describe('ConverterWorker', () => {
     };
 
     it('should update page status to COMPLETED', async () => {
+      let pageUpdateData: any;
+
       vi.mocked(prisma.$transaction).mockImplementation(async (callback: any) => {
         const tx = {
           taskDetail: {
             findUnique: vi.fn().mockResolvedValue({
               worker_id: worker.getWorkerId(),
               status: PageStatus.PROCESSING,
             }),
-            update: vi.fn(),
+            update: vi.fn().mockImplementation((params: any) => {
+              pageUpdateData = params.data;
+            }),
           },
           task: {
             findUnique: vi.fn().mockResolvedValue({
               status: TaskStatus.PROCESSING,
               pages: 10,
               completed_count: 5,
               failed_count: 0,
+              provider: 1,
+              model: 'gpt-4o',
             }),
             update: vi.fn().mockResolvedValue({
               completed_count: 6,
@@ -618,6 +624,8 @@ describe('ConverterWorker', () => {
       await (worker as any).completePageSuccess(mockPage, mockResult);
 
       expect(prisma.$transaction).toHaveBeenCalled();
+      expect(pageUpdateData.provider).toBe(1);
+      expect(pageUpdateData.model).toBe('gpt-4o');
     });
 
     it('should skip if page already completed (idempotency)', async () => {
@@ -823,6 +831,8 @@ describe('ConverterWorker', () => {
               pages: 10,
               completed_count: 5,
               failed_count: 1,
+              provider: 9,
+              model: 'claude-3-7-sonnet',
             }),
             update: vi.fn().mockResolvedValue({ failed_count: 2 }),
           },
@@ -839,6 +849,8 @@ describe('ConverterWorker', () => {
 
       expect(pageUpdateData.status).toBe(PageStatus.FAILED);
       expect(pageUpdateData.error).toBe('Test error');
+      expect(pageUpdateData.provider).toBe(9);
+      expect(pageUpdateData.model).toBe('claude-3-7-sonnet');
     });
 
     it('should set task to FAILED when all pages failed', async () => {

src/core/infrastructure/services/CloudService.ts

@@ -12,6 +12,7 @@ import type {
   CloudCancelTaskResponse,
   CloudRetryPageResponse,
   CloudApiPagination,
+  CloudModelTier,
   PaymentCheckoutApiResponse,
   PaymentCheckoutStatusApiResponse,
   PaymentHistoryApiItem,
@@ -433,15 +434,23 @@ class CloudService {
   /**
    * Retry an entire task (creates a new task)
    */
-  public async retryTask(id: string): Promise<{
+  public async retryTask(id: string, model?: CloudModelTier): Promise<{
     success: boolean;
     data?: CreateTaskResponse;
     error?: string;
   }> {
     try {
-      const res = await authManager.fetchWithAuth(`${API_BASE_URL}/api/v1/tasks/${encodeURIComponent(id)}/retry`, {
-        method: 'POST',
-      });
+      const hasModelOverride = typeof model === 'string' && model.length > 0;
+      const res = await authManager.fetchWithAuth(
+        `${API_BASE_URL}/api/v1/tasks/${encodeURIComponent(id)}/retry`,
+        hasModelOverride
+          ? {
+              method: 'POST',
+              headers: { 'Content-Type': 'application/json' },
+              body: JSON.stringify({ model }),
+            }
+          : { method: 'POST' }
+      );
 
       if (!res.ok) {
         const errorBody = await res.json().catch(() => null);
@@ -469,15 +478,22 @@ class CloudService {
   /**
    * Retry a single page
    */
-  public async retryPage(taskId: string, pageNumber: number): Promise<{
+  public async retryPage(taskId: string, pageNumber: number, model?: CloudModelTier): Promise<{
     success: boolean;
     data?: CloudRetryPageResponse;
     error?: string;
   }> {
     try {
+      const hasModelOverride = typeof model === 'string' && model.length > 0;
       const res = await authManager.fetchWithAuth(
         `${API_BASE_URL}/api/v1/tasks/${encodeURIComponent(taskId)}/pages/${encodeURIComponent(String(pageNumber))}/retry`,
-        { method: 'POST' },
+        hasModelOverride
+          ? {
+              method: 'POST',
+              headers: { 'Content-Type': 'application/json' },
+              body: JSON.stringify({ model }),
+            }
+          : { method: 'POST' },
       );
 
       if (!res.ok) {

src/core/infrastructure/services/__tests__/CloudService.test.ts

@@ -205,6 +205,37 @@ describe('CloudService', () => {
     })
   })
 
+  it('retryTask/retryPage send model override payload when model is provided', async () => {
+    const cloudService = (await import('../CloudService.js')).default
+
+    mockAuthManager.fetchWithAuth
+      .mockResolvedValueOnce(makeJsonResponse(200, { success: true, data: { task_id: 'task-2', events_url: '/events' } }))
+      .mockResolvedValueOnce(makeJsonResponse(200, { success: true, data: { task_id: 'task-1', page: 3, status: 'queued' } }))
+
+    await cloudService.retryTask('task-1', 'pro')
+    await cloudService.retryPage('task-1', 3, 'ultra')
+
+    const retryTaskCall = mockAuthManager.fetchWithAuth.mock.calls[0]
+    expect(retryTaskCall[0]).toContain('/api/v1/tasks/task-1/retry')
+    expect(retryTaskCall[1]).toEqual(
+      expect.objectContaining({
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ model: 'pro' }),
+      })
+    )
+
+    const retryPageCall = mockAuthManager.fetchWithAuth.mock.calls[1]
+    expect(retryPageCall[0]).toContain('/api/v1/tasks/task-1/pages/3/retry')
+    expect(retryPageCall[1]).toEqual(
+      expect.objectContaining({
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ model: 'ultra' }),
+      })
+    )
+  })
+
   it('cancelTask/retryTask/retryPage/deleteTask return API errors when non-OK', async () => {
     const cloudService = (await import('../CloudService.js')).default
     mockAuthManager.fetchWithAuth

src/main/ipc/__tests__/handlers.test.ts

@@ -139,6 +139,7 @@ vi.mock('../../../shared/ipc/channels.js', () => ({
       GET_ALL: 'task:getAll',
       GET_BY_ID: 'task:getById',
       UPDATE: 'task:update',
+      RETRY: 'task:retry',
       DELETE: 'task:delete',
       HAS_RUNNING: 'task:hasRunningTasks',
     },

src/main/ipc/handlers/__tests__/cloud.handler.test.ts

@@ -170,10 +170,21 @@ describe('cloud.handler', () => {
 
     expect(mockCloudService.getTasks).toHaveBeenCalledWith(2, 20)
     expect(mockCloudService.getTaskById).toHaveBeenCalledWith('t1')
-    expect(mockCloudService.retryTask).toHaveBeenCalledWith('t1')
+    expect(mockCloudService.retryTask).toHaveBeenCalledWith('t1', undefined)
     expect(mockCloudService.deleteTask).toHaveBeenCalledWith('t1')
   })
 
+  it('passes model override for retryTask and retryPage', async () => {
+    mockCloudService.retryTask.mockResolvedValueOnce({ success: true, data: { task_id: 'new' } })
+    mockCloudService.retryPage.mockResolvedValueOnce({ success: true, data: { task_id: 't1', page: 1, status: 1 } })
+
+    await handlers.get('cloud:retryTask')!({}, { id: 't1', model: 'pro' })
+    await handlers.get('cloud:retryPage')!({}, { taskId: 't1', pageNumber: 1, model: 'ultra' })
+
+    expect(mockCloudService.retryTask).toHaveBeenCalledWith('t1', 'pro')
+    expect(mockCloudService.retryPage).toHaveBeenCalledWith('t1', 1, 'ultra')
+  })
+
   describe('cloud:downloadPdf', () => {
     it('returns service error when download fails', async () => {
       mockCloudService.downloadPdf.mockResolvedValueOnce({ success: false, error: 'bad' })