Fix bug with offscreen API

Author: f1shy-dev

.github/workflows/release-server.yml

@@ -78,10 +78,7 @@ jobs:
 
       - name: Publish to NPM
         working-directory: packages/server
-        env:
-          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
-        run: |
-          npm publish --access public
+        run: npm publish --access public
 
       - name: Commit version change
         if: github.event_name == 'workflow_dispatch'

packages/extension/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@distracted/extension",
-  "version": "1.3.0",
+  "version": "1.3.1",
   "private": true,
   "description": "blocks distracting websites! do mini tasks to get back on them...",
   "keywords": [

packages/extension/src/entrypoints/background/index.ts

@@ -36,30 +36,73 @@ type GuardEntry = {
   settings: unknown;
 };
 
-const offscreenApi = (globalThis as { chrome?: { offscreen?: any } }).chrome?.offscreen;
+// MV3-only (Chrome): WXT provides typings for `browser.offscreen`.
+// Avoid `globalThis` shims/`any` here so type checking stays strict.
+const offscreenApi = isMV3 ? browser.offscreen : undefined;
+
+function isNoOffscreenDocumentError(error: unknown): boolean {
+  const message =
+    error instanceof Error
+      ? error.message
+      : typeof error === "string"
+        ? error
+        : JSON.stringify(error);
+  return message.includes("No current offscreen document");
+}
+
+function isReceivingEndDoesNotExistError(error: unknown): boolean {
+  const message =
+    error instanceof Error
+      ? error.message
+      : typeof error === "string"
+        ? error
+        : JSON.stringify(error);
+  return (
+    message.includes("Receiving end does not exist") ||
+    message.includes("Could not establish connection") ||
+    message.includes("No current offscreen document")
+  );
+}
 
 async function ensureOffscreenDocument(): Promise<void> {
   if (!isMV3 || !offscreenApi) return;
-  const hasDocument = (await offscreenApi.hasDocument?.()) as boolean | undefined;
-  if (hasDocument) return;
-
-  await offscreenApi.createDocument?.({
-    url: browser.runtime.getURL(
-      "/offscreen.html" as unknown as Parameters<typeof browser.runtime.getURL>[0],
-    ),
-    reasons: ["WORKERS"],
-    justification: "Maintain unlock guards that require real-time server state via WebSocket.",
-  });
+  try {
+    const hasDocument = (await offscreenApi.hasDocument()) as boolean | undefined;
+    if (hasDocument) return;
+  } catch (error) {
+    // Treat errors as "no document" and attempt to create one.
+    console.warn("[distracted] offscreen.hasDocument failed:", error);
+  }
+
+  try {
+    await offscreenApi.createDocument({
+      url: browser.runtime.getURL(
+        "/offscreen.html" as unknown as Parameters<typeof browser.runtime.getURL>[0],
+      ),
+      reasons: ["WORKERS"],
+      justification: "Maintain unlock guards that require real-time server state via WebSocket.",
+    });
+  } catch (error) {
+    // Chrome can be flaky here; if the document was created concurrently, this is safe to ignore.
+    console.warn("[distracted] offscreen.createDocument failed:", error);
+  }
 }
 
 async function maybeCloseOffscreenDocument(): Promise<void> {
   if (!isMV3 || !offscreenApi) return;
   const session = await browser.storage.session.get();
   const hasGuards = Object.keys(session).some((key) => key.startsWith(GUARD_PREFIX));
   if (!hasGuards) {
-    const hasDocument = (await offscreenApi.hasDocument?.()) as boolean | undefined;
-    if (hasDocument) {
-      await offscreenApi.closeDocument?.();
+    try {
+      const hasDocument = (await offscreenApi.hasDocument()) as boolean | undefined;
+      if (hasDocument) {
+        await offscreenApi.closeDocument();
+      }
+    } catch (error) {
+      // If the document is already gone, treat close as a no-op.
+      if (!isNoOffscreenDocumentError(error)) {
+        console.warn("[distracted] offscreen.closeDocument failed:", error);
+      }
     }
   }
 }
@@ -135,6 +178,24 @@ async function startGuardForSite(site: BlockedSite): Promise<void> {
         pollIntervalMs: guard.pollIntervalMs,
       });
     } catch (error) {
+      // MV3 service worker/offscreen can race; retry once after ensuring document exists.
+      if (isReceivingEndDoesNotExistError(error)) {
+        try {
+          await ensureOffscreenDocument();
+          await browser.runtime.sendMessage({
+            type: "GUARD_START",
+            siteId: entry.siteId,
+            method: entry.method,
+            settings: entry.settings,
+            heartbeatMs: GUARD_HEARTBEAT_MS,
+            pollIntervalMs: guard.pollIntervalMs,
+          });
+          return;
+        } catch (retryError) {
+          console.warn("[distracted] Failed to start guard in offscreen (retry):", retryError);
+          return;
+        }
+      }
       console.warn("[distracted] Failed to start guard in offscreen:", error);
     }
   } else {
@@ -152,7 +213,10 @@ async function stopGuardForSite(siteId: string): Promise<void> {
         siteId,
       });
     } catch (error) {
-      console.warn("[distracted] Failed to stop guard in offscreen:", error);
+      // If offscreen isn't around anymore, stopping is effectively complete.
+      if (!isReceivingEndDoesNotExistError(error)) {
+        console.warn("[distracted] Failed to stop guard in offscreen:", error);
+      }
     }
     await maybeCloseOffscreenDocument();
   } else {

packages/extension/src/entrypoints/blocked/App.tsx

@@ -112,9 +112,11 @@ export default function BlockedPage() {
     if (!blockedSite || !originalUrl) return;
 
     setUnlocking(true);
+    let navigationStarted = false;
 
     try {
       if (alreadyUnlocked) {
+        navigationStarted = true;
         window.location.href = originalUrl;
         return;
       }
@@ -146,6 +148,7 @@ export default function BlockedPage() {
           });
         }
 
+        navigationStarted = true;
         window.location.href = originalUrl;
       } else {
         setError(result.error || "Failed to unlock");
@@ -155,6 +158,13 @@ export default function BlockedPage() {
       console.error("[distracted] Error unlocking:", err);
       setError("Failed to unlock site");
       setUnlocking(false);
+    } finally {
+      // If navigation doesn't happen (e.g. browser blocks it), don't leave the UI stuck forever.
+      if (!navigationStarted) {
+        setUnlocking(false);
+      } else {
+        setTimeout(() => setUnlocking(false), 5000);
+      }
     }
   }, [blockedSite, originalUrl, statsEnabled, alreadyUnlocked]);
 

packages/extension/src/lib/claude-blocker.ts

@@ -74,7 +74,10 @@ export async function getClaudeBlockerStatus(serverUrl: string): Promise<ClaudeB
   }
 
   try {
-    const response = await fetch(statusUrl, { cache: "no-store" });
+    const controller = new AbortController();
+    const timeout = setTimeout(() => controller.abort(), 2000);
+    const response = await fetch(statusUrl, { cache: "no-store", signal: controller.signal });
+    clearTimeout(timeout);
     if (!response.ok) {
       return {
         active: false,

packages/extension/src/lib/guard-watcher.ts

@@ -27,6 +27,9 @@ export function startGuardWatcher({
   let socket: WebSocket | null = null;
   let pollTimer: ReturnType<typeof setInterval> | null = null;
   let heartbeatTimer: ReturnType<typeof setInterval> | null = null;
+  let wsPingTimer: ReturnType<typeof setInterval> | null = null;
+  let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+  let reconnectAttempts = 0;
   let lastState: UnlockGuardState | null = null;
 
   const emitState = (state: UnlockGuardState) => {
@@ -40,6 +43,8 @@ export function startGuardWatcher({
     stopped = true;
     if (pollTimer) clearInterval(pollTimer);
     if (heartbeatTimer) clearInterval(heartbeatTimer);
+    if (wsPingTimer) clearInterval(wsPingTimer);
+    if (reconnectTimer) clearTimeout(reconnectTimer);
     if (socket) {
       socket.close();
       socket = null;
@@ -80,25 +85,67 @@ export function startGuardWatcher({
 
   const wsUrl = guard.getWebSocketUrl?.(settings as never);
   if (wsUrl) {
-    socket = new WebSocket(wsUrl);
-    void runCheck();
-    socket.addEventListener("message", (event) => {
+    // WebSocket mode: keep a polling fallback and auto-reconnect the socket.
+    // This avoids instantly revoking access on transient WS drops (common in MV3/offscreen).
+    startPolling();
+
+    const scheduleReconnect = () => {
+      if (stopped) return;
+      if (reconnectTimer) return;
+      reconnectAttempts += 1;
+      const backoffMs = Math.min(10_000, 500 * 2 ** Math.min(6, reconnectAttempts)); // 0.5s .. 10s
+      reconnectTimer = setTimeout(() => {
+        reconnectTimer = null;
+        connect();
+      }, backoffMs);
+    };
+
+    const connect = () => {
+      if (stopped) return;
       try {
-        const payload = JSON.parse(event.data as string);
-        const state = guard.parseWebSocketMessage?.(payload, settings as never);
-        if (state) {
-          handleGuardState(state);
+        socket?.close();
+      } catch {
+        // ignore
+      }
+      socket = new WebSocket(wsUrl);
+
+      socket.addEventListener("open", () => {
+        reconnectAttempts = 0;
+        void runCheck();
+      });
+
+      socket.addEventListener("message", (event) => {
+        try {
+          const payload = JSON.parse(event.data as string);
+          const state = guard.parseWebSocketMessage?.(payload, settings as never);
+          if (state) {
+            handleGuardState(state);
+          }
+        } catch {
+          // Ignore malformed messages
         }
+      });
+
+      socket.addEventListener("close", () => {
+        scheduleReconnect();
+      });
+      socket.addEventListener("error", () => {
+        scheduleReconnect();
+      });
+    };
+
+    connect();
+
+    // App-level ping to keep the server (and any intermediaries) from timing out idle connections.
+    wsPingTimer = setInterval(() => {
+      if (stopped) return;
+      if (!socket || socket.readyState !== WebSocket.OPEN) return;
+      try {
+        socket.send(JSON.stringify({ type: "ping" }));
       } catch {
-        // Ignore malformed messages
+        // ignore; reconnect loop will handle close/error
       }
-    });
-    socket.addEventListener("close", () => {
-      handleGuardState({ active: false, reason: "offline" });
-    });
-    socket.addEventListener("error", () => {
-      handleGuardState({ active: false, reason: "offline" });
-    });
+    }, 10_000);
 
     startHeartbeat();
   } else {

packages/server/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@distracted/server",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "description": "Local server + Claude Code hook installer for distracted's Claude Blocker integration.",
   "license": "MIT",
   "repository": {