Browse Source

fix(smartfetch): add 30s timeout to secondary model prompt

If the secondary model's upstream hangs without closing the connection,
the prompt call blocks indefinitely, freezing smartfetch. Wrap the call
in Promise.race with a 30s timeout. On timeout the error propagates to
runSecondaryModelWithFallback which tries the next model, and if all
fail, the tool returns raw fetched content as before.
dhaern 1 month ago
parent
commit
4303afa596
2 changed files with 61 additions and 18 deletions
  1. 34 0
      src/tools/smartfetch/secondary-model.test.ts
  2. 27 18
      src/tools/smartfetch/secondary-model.ts

+ 34 - 0
src/tools/smartfetch/secondary-model.test.ts

@@ -147,4 +147,38 @@ describe('smartfetch/secondary-model', () => {
       console.warn = originalWarn;
     }
   });
+
+  test('falls back to next model when prompt times out', async () => {
+    const client = {
+      session: {
+        create: mock(async () => ({ id: 'session-timeout' })),
+        prompt: mock(async (opts: any) => {
+          const model = opts.body.model;
+          if (model.modelID === 'small') {
+            throw new Error('Secondary model timed out');
+          }
+          return {
+            data: {
+              parts: [{ type: 'text', text: 'Fallback answer' }],
+            },
+          };
+        }),
+        delete: mock(async () => ({})),
+      },
+      tool: {
+        ids: mock(async () => ({ data: ['read'] })),
+      },
+    } as any;
+
+    const result = await runSecondaryModelWithFallback(
+      client,
+      '/tmp/project',
+      models,
+      'Summarize',
+      'This is enough fetched content to clear the short-content guard.',
+    );
+
+    expect(result.text).toBe('Fallback answer');
+    expect(result.model).toEqual(models[1]);
+  });
 });

+ 27 - 18
src/tools/smartfetch/secondary-model.ts

@@ -158,6 +158,7 @@ function isUsableSecondaryText(text: string) {
 
 const SESSION_DELETE_RETRIES = 3;
 const SESSION_DELETE_RETRY_DELAY_MS = 500;
+const SECONDARY_MODEL_TIMEOUT_MS = 30_000;
 
 /**
  * Delete a temporary secondary-model session with retry.
@@ -239,24 +240,32 @@ async function runSecondaryModel(
       (toolIDs || []).map((id: string) => [id, false]),
     );
 
-    const result = await client.session.prompt({
-      responseStyle: 'data',
-      throwOnError: true,
-      path: { id: sessionId },
-      query: { directory },
-      body: {
-        model,
-        system:
-          'Answer only from the supplied content. Do not use tools or outside knowledge.',
-        tools: disabledTools,
-        parts: [
-          {
-            type: 'text',
-            text: buildPrompt(truncatedContent, effectivePrompt),
-          },
-        ],
-      },
-    });
+    const result = await Promise.race([
+      client.session.prompt({
+        responseStyle: 'data',
+        throwOnError: true,
+        path: { id: sessionId },
+        query: { directory },
+        body: {
+          model,
+          system:
+            'Answer only from the supplied content. Do not use tools or outside knowledge.',
+          tools: disabledTools,
+          parts: [
+            {
+              type: 'text',
+              text: buildPrompt(truncatedContent, effectivePrompt),
+            },
+          ],
+        },
+      }),
+      new Promise<never>((_, reject) =>
+        setTimeout(
+          () => reject(new Error('Secondary model timed out')),
+          SECONDARY_MODEL_TIMEOUT_MS,
+        ),
+      ),
+    ]);
 
     const parts =
       (result as { data?: { parts?: Array<{ type?: string; text?: string }> } })