Add timeout

Signed-off-by: Leslie Lazzarino <leslie.lazzarino@gmail.com>

Signed-off-by: Leslie Lazzarino <leslie.lazzarino@gmail.com>

Author: leslielazzarino

manifest.json

@@ -11,7 +11,7 @@
       "update_url": "https://raw.githubusercontent.com/TNG/tb-llm-composer/refs/heads/main/updates.json"
     }
   },
-  "permissions": ["menus", "compose", "storage", "notifications", "messagesRead", "accountsRead"],
+  "permissions": ["menus", "compose", "storage", "notifications", "messagesRead", "accountsRead", "alarms"],
   "icons": {
     "64": "icons/icon-64px.png",
     "32": "icons/icon-32px.png",

public/options.html

@@ -126,6 +126,15 @@ <h2>LLM API Settings</h2>
             <label for="api_token">Api token:</label>
             <input type="text" id="api_token" name="api token" />
           </div>
+          <div class="form-group">
+            <label for="timeout">
+              Request timeout (seconds):
+              <sup title="Timeout for API requests in seconds. Leave empty or set to 0 for no timeout.">
+                &#8505;&#65039;
+              </sup>
+            </label>
+            <input type="number" id="timeout" name="timeout" min="0" step="1" placeholder="No timeout" />
+          </div>
         </div>
 
         <div class="settings-section">

src/__tests__/llmConnection.test.ts

@@ -1,6 +1,6 @@
 import { afterAll, describe, expect, test } from "vitest";
 import { type LlmApiRequestMessage, LlmRoles, sendContentToLlm } from "../llmConnection";
-import { getExpectedRequestContent, getMockResponseBody, mockBrowserAndFetch } from "./testUtils";
+import { getMockResponseBody, mockBrowserAndFetch } from "./testUtils";
 
 const originalBrowser = global.browser;
 const originalFetch = global.fetch;
@@ -39,7 +39,12 @@ describe("Testing sentContentToLlm", () => {
     expect(global.fetch).toHaveBeenCalledTimes(1);
     expect(global.fetch).toHaveBeenCalledWith(
       MOCK_MODEL_URL,
-      getExpectedRequestContent([MOCK_CONTEXT, MOCK_PROMPT], abortSignal),
+      expect.objectContaining({
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({ messages: [MOCK_CONTEXT, MOCK_PROMPT] }),
+        signal: expect.any(AbortSignal),
+      }),
     );
     expect(result).toEqual(mockResponseBody);
   });
@@ -57,7 +62,12 @@ describe("Testing sentContentToLlm", () => {
     expect(global.fetch).toHaveBeenCalledTimes(1);
     expect(global.fetch).toHaveBeenCalledWith(
       MOCK_MODEL_URL,
-      getExpectedRequestContent([MOCK_CONTEXT, MOCK_PROMPT], abortSignal, mockToken),
+      expect.objectContaining({
+        method: "POST",
+        headers: { "Content-Type": "application/json", Authorization: `Bearer ${mockToken}` },
+        body: JSON.stringify({ messages: [MOCK_CONTEXT, MOCK_PROMPT] }),
+        signal: expect.any(AbortSignal),
+      }),
     );
     expect(result).toEqual(mockResponseBody);
   });

src/__tests__/testUtils.ts

@@ -1,4 +1,5 @@
 import { vi } from "vitest";
+import type { LlmPluginAction } from "../llmButtonClickHandling";
 import {
   type LlmApiRequestMessage,
   type LlmChoice,
@@ -11,8 +12,6 @@ import { DEFAULT_OPTIONS, type LlmParameters, type Options } from "../optionsPar
 import ComposeDetails = browser.compose.ComposeDetails;
 import _CreateCreateProperties = browser.menus._CreateCreateProperties;
 
-import type { LlmPluginAction } from "../llmButtonClickHandling";
-
 const MOCK_IDENTITY_ID = "MOCK_IDENTITY_ID";
 export const MOCK_TAB_DETAILS: browser.compose.ComposeDetails = { identityId: MOCK_IDENTITY_ID };
 export const MOCK_USER_NAME = "MOCK_USER_NAME";
@@ -135,6 +134,16 @@ export function mockBrowser(args: mockBrowserArgs) {
     commands: {
       getAll: async () => allShortcuts,
     },
+    // @ts-ignore
+    alarms: {
+      create: vi.fn(),
+      clear: vi.fn(),
+      onAlarm: {
+        addListener: vi.fn(),
+        removeListener: vi.fn(),
+        hasListener: vi.fn(),
+      },
+    },
   };
 }
 
@@ -203,13 +212,17 @@ export function getExpectedRequestContent(
     ...params,
   };
 
+  const headers: { [key: string]: string } = {
+    "Content-Type": "application/json",
+  };
+  if (api_token) {
+    headers.Authorization = `Bearer ${api_token}`;
+  }
+
   return {
     body: JSON.stringify(expectedRequestBody),
     signal,
-    headers: {
-      "Content-Type": "application/json",
-      Authorization: api_token ? `Bearer ${api_token}` : undefined,
-    },
+    headers,
     method: "POST",
   };
 }

src/background.ts

@@ -1,3 +1,4 @@
+import { handleKeepAliveAlarm } from "./keepAlive";
 import { executeLlmAction, type LlmPluginAction } from "./llmButtonClickHandling";
 import { addLlmActionsToMenu, enableSummarizeMenuEntryIfReply, handleMenuClickListener } from "./menu";
 import { deleteFromOriginalTabCache, storeOriginalReplyText } from "./originalTabConversation";
@@ -8,6 +9,9 @@ import Tab = browser.tabs.Tab;
 // Otherwise, the shortcuts may not work if the background script is not running (which is after 90s of idling or so)
 browser.commands.onCommand.addListener((command, tab) => executeLlmAction(command as LlmPluginAction, tab));
 
+// Keep the background page alive during long-running LLM requests
+browser.alarms.onAlarm.addListener(handleKeepAliveAlarm);
+
 browser.tabs.onCreated.addListener(async (tab: Tab) => {
   await storeOriginalReplyText(tab);
   await enableSummarizeMenuEntryIfReply(tab);

src/keepAlive.ts

@@ -0,0 +1,40 @@
+const KEEP_ALIVE_ALARM_NAME = "llm-keep-alive";
+const KEEP_ALIVE_INTERVAL_MINUTES = 0.4; // ~24 seconds, well under the ~90s idle threshold
+
+let activeRequests = 0;
+
+/**
+ * Starts a periodic alarm to keep the background page alive during long-running LLM requests.
+ * Thunderbird MV3 event pages can be suspended after ~90s of inactivity.
+ * While an active fetch() should prevent suspension, this alarm acts as a safety net
+ * for very long local LLM inference times (minutes to hours).
+ */
+export async function startKeepAlive(): Promise<void> {
+  activeRequests++;
+  if (activeRequests === 1) {
+    console.log("KEEP-ALIVE: Starting keep-alive alarm for long-running LLM request");
+    browser.alarms.create(KEEP_ALIVE_ALARM_NAME, {
+      periodInMinutes: KEEP_ALIVE_INTERVAL_MINUTES,
+    });
+  }
+}
+
+/**
+ * Stops the keep-alive alarm when no more LLM requests are active.
+ */
+export async function stopKeepAlive(): Promise<void> {
+  activeRequests = Math.max(0, activeRequests - 1);
+  if (activeRequests === 0) {
+    console.log("KEEP-ALIVE: Stopping keep-alive alarm, no active LLM requests");
+    await browser.alarms.clear(KEEP_ALIVE_ALARM_NAME);
+  }
+}
+
+/**
+ * Handler for the keep-alive alarm. Simply logs to keep the background page active.
+ */
+export function handleKeepAliveAlarm(alarm: browser.alarms.Alarm): void {
+  if (alarm.name === KEEP_ALIVE_ALARM_NAME) {
+    console.log("KEEP-ALIVE: Heartbeat - background page still active");
+  }
+}

src/llmConnection.ts

@@ -1,3 +1,4 @@
+import { startKeepAlive, stopKeepAlive } from "./keepAlive";
 import { getPluginOptions, type LlmParameters } from "./optionsParams";
 
 export enum LlmRoles {
@@ -81,14 +82,15 @@ export async function sendContentToLlm(
     ...options.params,
   };
 
-  return callLlmApi(options.model, requestBody, abortSignal, options.api_token);
+  return callLlmApi(options.model, requestBody, abortSignal, options.api_token, options.timeout);
 }
 
 async function callLlmApi(
   url: string,
   requestBody: LlmApiRequestBody,
   signal: AbortSignal,
   token?: string,
+  timeout?: number,
 ): Promise<LlmTextCompletionResponse | TgiErrorResponse> {
   const headers: { [key: string]: string } = {
     "Content-Type": "application/json",
@@ -97,21 +99,51 @@ async function callLlmApi(
     headers.Authorization = `Bearer ${token}`;
   }
 
-  console.log(`LLM-CONNECTION: Sending request to LLM: POST ${url} with body:\n`, JSON.stringify(requestBody));
-  const response = await fetch(url, {
-    signal: signal,
-    method: "POST",
-    headers: headers,
-    body: JSON.stringify(requestBody),
-  });
-  if (!response.ok) {
-    const errorResponseBody = await response.text();
-    throw Error(`LLM-CONNECTION: Error response from ${url}: ${errorResponseBody}`);
+  // Create a combined abort controller for both user cancellation and timeout
+  const combinedAbortController = new AbortController();
+  let timeoutId: ReturnType<typeof setTimeout> | undefined;
+
+  // Listen to the user's abort signal
+  const abortHandler = () => combinedAbortController.abort(signal.reason);
+  signal.addEventListener("abort", abortHandler);
+
+  // Set up timeout if specified
+  if (timeout !== undefined && timeout > 0) {
+    timeoutId = setTimeout(() => {
+      combinedAbortController.abort(
+        new DOMException(
+          `LLM request timed out after ${timeout / 1000} seconds. You can increase or disable the timeout in the extension settings.`,
+          "TimeoutError",
+        ),
+      );
+    }, timeout);
   }
-  const responseBody = (await response.json()) as LlmTextCompletionResponse | TgiErrorResponse;
-  console.log("LLM-CONNECTION: LLM responded with:", response.status, responseBody);
 
-  return responseBody;
+  await startKeepAlive();
+  try {
+    console.log(`LLM-CONNECTION: Sending request to LLM: POST ${url} with body:\n`, JSON.stringify(requestBody));
+    const response = await fetch(url, {
+      signal: combinedAbortController.signal,
+      method: "POST",
+      headers: headers,
+      body: JSON.stringify(requestBody),
+    });
+    if (!response.ok) {
+      const errorResponseBody = await response.text();
+      throw Error(`LLM-CONNECTION: Error response from ${url}: ${errorResponseBody}`);
+    }
+    const responseBody = (await response.json()) as LlmTextCompletionResponse | TgiErrorResponse;
+    console.log("LLM-CONNECTION: LLM responded with:", response.status, responseBody);
+
+    return responseBody;
+  } finally {
+    await stopKeepAlive();
+    // Clean up
+    signal.removeEventListener("abort", abortHandler);
+    if (timeoutId !== undefined) {
+      clearTimeout(timeoutId);
+    }
+  }
 }
 
 export function isLlmTextCompletionResponse(response: LlmTextCompletionResponse | TgiErrorResponse) {

src/notifications.ts

@@ -14,6 +14,14 @@ export function notifyOnError<T>(callback: () => Promise<T>) {
       console.debug("User cancelled request, do not notify", e);
       return;
     }
+    if (e.name === "TimeoutError") {
+      timedNotification(
+        "LLM Request Timeout",
+        e.message || "The LLM request timed out. You can increase the timeout in the extension settings.",
+        15000,
+      );
+      return;
+    }
     const message = e?.message || e.toString();
     timedNotification("Thunderbird LLM Extension Error", message);
   });

src/options.ts

@@ -5,6 +5,7 @@ import { getInputElement } from "./utils";
 document.addEventListener("DOMContentLoaded", restoreOptions);
 document.querySelector("#url")?.addEventListener("change", updateUrl);
 document.querySelector("#api_token")?.addEventListener("change", updateApiToken);
+document.querySelector("#timeout")?.addEventListener("change", updateTimeout);
 document.querySelector("#llm_context")?.addEventListener("change", updateLlmContext);
 document.querySelector("#use_last_mails")?.addEventListener("change", updateUseLastMails);
 document.querySelector("#context_window")?.addEventListener("change", updateContextWindow);
@@ -29,6 +30,15 @@ async function updateApiToken(event: Event) {
   await browser.storage.sync.set({ options });
 }
 
+async function updateTimeout(event: Event) {
+  const timeoutInput = event.target as HTMLInputElement;
+  const options = await getPluginOptions();
+  const timeoutSeconds = timeoutInput.valueAsNumber;
+  // Convert seconds to milliseconds, or set to undefined if 0 or empty
+  options.timeout = timeoutSeconds > 0 ? timeoutSeconds * 1000 : undefined;
+  await browser.storage.sync.set({ options });
+}
+
 async function updateLlmContext(event: Event) {
   const llmContextInput = event.target as HTMLTextAreaElement;
   const options = await getPluginOptions();
@@ -64,6 +74,7 @@ export async function restoreOptions(): Promise<void> {
 
   getInputElement("#url").value = options.model;
   getInputElement("#api_token").value = options.api_token || "";
+  getInputElement("#timeout").value = options.timeout ? `${options.timeout / 1000}` : "";
   getInputElement("#context_window").value = `${options.context_window}`;
   getInputElement("#use_last_mails").checked = options.include_recent_mails;
   getInputElement("#other_options").value = JSON.stringify(options.params, null, 2);

src/optionsParams.ts

@@ -23,6 +23,7 @@ export interface Options {
   include_recent_mails: boolean;
   params: LlmParameters;
   llmContext: string;
+  timeout?: number; // Timeout in milliseconds, undefined means no timeout
 }
 
 export const DEFAULT_PARAMS: LlmParameters = {};