Fix extension losing workflow state when interacting with webpage (#86)

Author: SamuelHassine

src/background/handlers/misc-handlers.ts

@@ -12,6 +12,7 @@ import { successResponse, errorResponse } from '../../shared/types/common';
 import type { MessageHandler } from './types';
 import { loggers } from '../../shared/utils/logger';
 import { generateNativePDF, isNativePDFAvailable } from '../../shared/extraction/native-pdf';
+import { savePanelWorkflowState, loadPanelWorkflowState, clearPanelWorkflowState, type PanelWorkflowState } from '../../shared/utils/storage';
 
 const log = loggers.background;
 
@@ -182,12 +183,47 @@ export const handleInjectAllTabs: MessageHandler = async (_payload, sendResponse
 };
 
 /**
- * Get panel state
+ * Get panel state - returns any saved workflow state so the panel can restore it
  */
 export const handleGetPanelState: MessageHandler = async (_payload, sendResponse) => {
-  // Panel state is managed by the content script
-  // This just confirms the panel is ready
-  sendResponse(successResponse({ ready: true }));
+  try {
+    const workflowState = await loadPanelWorkflowState();
+    sendResponse(successResponse({ ready: true, workflowState }));
+  } catch (error) {
+    log.warn('Failed to load panel workflow state:', error);
+    sendResponse(successResponse({ ready: true, workflowState: null }));
+  }
+};
+
+/**
+ * Save panel workflow state for persistence across close/reopen
+ */
+export const handleSavePanelState: MessageHandler = async (payload, sendResponse) => {
+  const state = payload as PanelWorkflowState | undefined;
+  if (!state || typeof state !== 'object' || !('panelMode' in state) || !state.panelMode) {
+    sendResponse(errorResponse('Invalid panel state'));
+    return;
+  }
+  try {
+    await savePanelWorkflowState({ ...state, timestamp: Date.now() });
+    sendResponse(successResponse({ saved: true }));
+  } catch (error) {
+    log.warn('Failed to save panel workflow state:', error);
+    sendResponse(errorResponse('Failed to save state'));
+  }
+};
+
+/**
+ * Clear panel workflow state (e.g. after completing a workflow)
+ */
+export const handleClearPanelState: MessageHandler = async (_payload, sendResponse) => {
+  try {
+    await clearPanelWorkflowState();
+    sendResponse(successResponse({ cleared: true }));
+  } catch (error) {
+    log.warn('Failed to clear panel workflow state:', error);
+    sendResponse(errorResponse('Failed to clear state'));
+  }
 };
 
 /**
@@ -293,6 +329,8 @@ export const miscHandlers: Record<string, MessageHandler> = {
   INJECT_CONTENT_SCRIPT: handleInjectContentScript,
   INJECT_ALL_TABS: handleInjectAllTabs,
   GET_PANEL_STATE: handleGetPanelState,
+  SAVE_PANEL_STATE: handleSavePanelState,
+  CLEAR_PANEL_STATE: handleClearPanelState,
   GENERATE_NATIVE_PDF: handleGenerateNativePDF,
   FETCH_IMAGE_AS_DATA_URL: handleFetchImageAsDataURL,
 };

src/content/panel.ts

@@ -14,7 +14,6 @@
  */
 
 import { loggers } from '../shared/utils/logger';
-import { PANEL_WIDTH_PX } from '../shared/constants';
 import type { DetectedObservable } from '../shared/types/observables';
 import type { DetectedOCTIEntity } from '../shared/types/opencti';
 import { extractArticleContent, extractFirstParagraph } from './extraction';
@@ -29,7 +28,6 @@ let panelFrame: HTMLIFrameElement | null = null;
 let panelOverlay: HTMLDivElement | null = null;
 let isPanelReady = false;
 const panelMessageQueue: Array<{ type: string; payload?: unknown }> = [];
-let documentClickHandlerInstalled = false;
 let highlightClickInProgress = false;
 // Split screen mode - uses browser's native side panel instead of floating iframe
 let splitScreenMode = false;
@@ -431,11 +429,6 @@ export function ensurePanelElements(): void {
     }, 100);
   }
 
-  // Install document click handler
-  if (!documentClickHandlerInstalled) {
-    document.addEventListener('click', handleDocumentClickForPanel, true);
-    documentClickHandlerInstalled = true;
-  }
 }
 
 /**
@@ -482,44 +475,6 @@ export function isPanelHidden(): boolean {
   return panelFrame?.classList.contains('hidden') ?? true;
 }
 
-// ============================================================================
-// Click Handling
-// ============================================================================
-
-/**
- * Document click handler for closing panel when clicking outside
- */
-function handleDocumentClickForPanel(e: MouseEvent): void {
-  if (panelFrame?.classList.contains('hidden')) {
-    return;
-  }
-  
-  if (highlightClickInProgress) {
-    return;
-  }
-  
-  const target = e.target as HTMLElement;
-  
-  if (panelFrame && (target === panelFrame || panelFrame.contains(target))) {
-    return;
-  }
-  
-  if (target.closest('.xtm-highlight')) {
-    return;
-  }
-  
-  if (target.closest('[class*="xtm-"]')) {
-    return;
-  }
-  
-  const panelAreaStart = window.innerWidth - PANEL_WIDTH_PX;
-  if (e.clientX >= panelAreaStart) {
-    return;
-  }
-  
-  hidePanel();
-}
-
 // ============================================================================
 // Theme Helper
 // ============================================================================

src/panel/App.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useState, useMemo, useCallback } from 'react';
+import React, { useEffect, useRef, useState, useMemo, useCallback } from 'react';
 import {
   ThemeProvider,
   createTheme,
@@ -357,6 +357,75 @@ const App: React.FC = () => {
     showToast,
   });
 
+  // Modes that represent an active workflow worth persisting
+  const PERSISTABLE_MODES = new Set([
+    'scan-results', 'preview', 'platform-select', 'container-type',
+    'container-form', 'existing-containers', 'entity', 'not-found',
+    'investigation', 'scenario-overview', 'scenario-form',
+    'atomic-testing', 'unified-search', 'add', 'add-selection',
+    'add-to-scan-results',
+  ]);
+
+  // Track whether we've restored state to avoid save/restore loops
+  const hasRestoredState = useRef(false);
+  // Platform ID that needs URL resolution once platforms finish loading
+  const pendingPlatformIdRef = useRef<string | null>(null);
+
+  // Debounced persistence of panel workflow state.
+  // Reads scan data from refs to always capture the latest values without
+  // requiring them in the dependency array (which would trigger too-frequent saves).
+  useEffect(() => {
+    if (!hasRestoredState.current) return;
+    if (!PERSISTABLE_MODES.has(panelMode)) return;
+    if (typeof chrome === 'undefined' || !chrome.runtime?.sendMessage) return;
+
+    const timer = setTimeout(() => {
+      const currentScanEntities = scanResultsEntitiesRef.current;
+      const stateToSave = {
+        panelMode,
+        containerType,
+        containerForm,
+        containerWorkflowOrigin,
+        selectedPlatformId,
+        currentPageUrl,
+        currentPageTitle,
+        currentPdfFileName,
+        isPdfSource,
+        scanResultsEntities: currentScanEntities.length > 0 ? currentScanEntities : undefined,
+        selectedScanItems: selectedScanItems.size > 0 ? Array.from(selectedScanItems) : undefined,
+        scanPageContent: scanPageContent || undefined,
+        timestamp: Date.now(),
+      };
+      chrome.runtime.sendMessage({ type: 'SAVE_PANEL_STATE', payload: stateToSave });
+    }, 500);
+
+    return () => clearTimeout(timer);
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [panelMode, containerType, containerForm, containerWorkflowOrigin,
+      selectedPlatformId, currentPageUrl, currentPageTitle,
+      selectedScanItems, scanPageContent]);
+
+  // Resolve deferred platform URL once platforms finish loading
+  useEffect(() => {
+    const pid = pendingPlatformIdRef.current;
+    if (pid && availablePlatforms.length > 0) {
+      const platform = availablePlatforms.find(p => p.id === pid);
+      if (platform) {
+        setPlatformUrl(platform.url);
+        pendingPlatformIdRef.current = null;
+      }
+    }
+  }, [availablePlatforms, setPlatformUrl]);
+
+  // Clear persisted state when a workflow completes or user returns to empty
+  useEffect(() => {
+    if (panelMode === 'empty' || panelMode === 'import-results') {
+      if (typeof chrome !== 'undefined' && chrome.runtime?.sendMessage) {
+        chrome.runtime.sendMessage({ type: 'CLEAR_PANEL_STATE' });
+      }
+    }
+  }, [panelMode]);
+
   // Sync entitiesToAdd with selectedScanItems when in preview mode
   React.useEffect(() => {
     if (panelMode === 'preview' && scanResultsEntitiesRef.current.length > 0) {
@@ -645,7 +714,15 @@ const App: React.FC = () => {
     setTimeout(() => {
       chrome.runtime.sendMessage({ type: 'GET_PANEL_STATE' }, (response) => {
         if (chrome.runtime.lastError) return;
-        if (response?.success && response.data) handlePanelState(response.data);
+        if (response?.success && response.data) {
+          handlePanelState(response.data);
+          // Restore persisted workflow state if available
+          const ws = response.data.workflowState;
+          if (ws && ws.panelMode && ws.panelMode !== 'empty' && ws.panelMode !== 'loading') {
+            restoreWorkflowState(ws);
+          }
+        }
+        hasRestoredState.current = true;
       });
     }, 50);
 
@@ -845,6 +922,49 @@ const App: React.FC = () => {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [panelMode, selectedPlatformId]);
 
+  // Restore workflow state from persisted data (called on panel mount)
+  const restoreWorkflowState = (ws: {
+    panelMode?: string;
+    containerType?: string;
+    containerForm?: { name: string; description: string; content?: string };
+    containerWorkflowOrigin?: string;
+    selectedPlatformId?: string;
+    currentPageUrl?: string;
+    currentPageTitle?: string;
+    currentPdfFileName?: string;
+    isPdfSource?: boolean;
+    scanResultsEntities?: unknown[];
+    selectedScanItems?: string[];
+    scanPageContent?: string;
+  }) => {
+    log.debug('[PANEL] Restoring workflow state, mode:', ws.panelMode);
+    if (ws.currentPageUrl) setCurrentPageUrl(ws.currentPageUrl);
+    if (ws.currentPageTitle) setCurrentPageTitle(ws.currentPageTitle);
+    if (ws.currentPdfFileName) setCurrentPdfFileName(ws.currentPdfFileName);
+    if (ws.isPdfSource !== undefined) setIsPdfSource(ws.isPdfSource);
+    if (ws.selectedPlatformId) {
+      setSelectedPlatformId(ws.selectedPlatformId);
+      const platform = availablePlatformsRef.current.find(p => p.id === ws.selectedPlatformId);
+      if (platform) {
+        setPlatformUrl(platform.url);
+      } else {
+        pendingPlatformIdRef.current = ws.selectedPlatformId;
+      }
+    }
+    if (ws.scanResultsEntities && Array.isArray(ws.scanResultsEntities) && ws.scanResultsEntities.length > 0) {
+      setScanResultsEntities(ws.scanResultsEntities as typeof scanResultsEntities);
+      scanResultsEntitiesRef.current = ws.scanResultsEntities as typeof scanResultsEntities;
+    }
+    if (ws.selectedScanItems && Array.isArray(ws.selectedScanItems)) {
+      setSelectedScanItems(new Set(ws.selectedScanItems));
+    }
+    if (ws.scanPageContent) setScanPageContent(ws.scanPageContent);
+    if (ws.containerType) setOCTIContainerType(ws.containerType as Parameters<typeof setOCTIContainerType>[0]);
+    if (ws.containerForm) setContainerForm(ws.containerForm as Parameters<typeof setContainerForm>[0]);
+    if (ws.containerWorkflowOrigin) setContainerWorkflowOrigin(ws.containerWorkflowOrigin as 'preview' | 'direct' | 'import');
+    if (ws.panelMode) setPanelMode(ws.panelMode as PanelMode);
+  };
+
   const handleMessage = (event: MessageEvent) => {
     const { type, payload } = event.data;
     handlePanelState({ type, payload });

src/panel/views/CommonScanResultsView.tsx

@@ -915,6 +915,10 @@ export const CommonScanResultsView: React.FC<ExtendedScanResultsViewProps> = ({
               // Reset filters
               setScanResultsTypeFilter('all');
               setScanResultsFoundFilter('all');
+              // Clear persisted workflow state so stale results don't restore on reopen
+              if (typeof chrome !== 'undefined' && chrome.runtime?.sendMessage) {
+                chrome.runtime.sendMessage({ type: 'CLEAR_PANEL_STATE' });
+              }
             }}
             sx={{
               textTransform: 'none',

src/shared/types/messages.ts

@@ -36,6 +36,8 @@ export type MessageType =
   | 'GET_CACHE_STATS'
   | 'SELECTION_CHANGED'
   | 'GET_PANEL_STATE'
+  | 'SAVE_PANEL_STATE'
+  | 'CLEAR_PANEL_STATE'
   | 'OPEN_SIDE_PANEL'
   | 'OPEN_SIDE_PANEL_IMMEDIATE'
   | 'FORWARD_TO_PANEL'

src/shared/utils/storage.ts

@@ -82,6 +82,59 @@ export const sessionStorage = {
   },
 };
 
+// ============================================================================
+// Panel Workflow State Persistence
+// ============================================================================
+
+const PANEL_WORKFLOW_STATE_KEY = 'panelWorkflowState';
+const PANEL_WORKFLOW_STATE_TTL = 30 * 60 * 1000; // 30 minutes
+
+export interface PanelWorkflowState {
+  panelMode: string;
+  containerType?: string;
+  containerForm?: { name: string; description: string; content?: string };
+  containerWorkflowOrigin?: string;
+  selectedPlatformId?: string;
+  currentPageUrl?: string;
+  currentPageTitle?: string;
+  currentPdfFileName?: string;
+  isPdfSource?: boolean;
+  scanResultsEntities?: unknown[];
+  selectedScanItems?: string[];
+  scanPageContent?: string;
+  timestamp: number;
+}
+
+/**
+ * Save panel workflow state to session storage.
+ * This preserves the user's place in a workflow (e.g. container creation)
+ * across panel close/reopen cycles.
+ */
+export async function savePanelWorkflowState(state: PanelWorkflowState): Promise<void> {
+  await sessionStorage.set(PANEL_WORKFLOW_STATE_KEY, state);
+}
+
+/**
+ * Load panel workflow state from session storage.
+ * Returns null if no state exists or the state has expired.
+ */
+export async function loadPanelWorkflowState(): Promise<PanelWorkflowState | null> {
+  const state = await sessionStorage.get<PanelWorkflowState>(PANEL_WORKFLOW_STATE_KEY);
+  if (!state) return null;
+  if (Date.now() - state.timestamp > PANEL_WORKFLOW_STATE_TTL) {
+    await sessionStorage.remove(PANEL_WORKFLOW_STATE_KEY);
+    return null;
+  }
+  return state;
+}
+
+/**
+ * Clear panel workflow state (e.g. after completing a workflow).
+ */
+export async function clearPanelWorkflowState(): Promise<void> {
+  await sessionStorage.remove(PANEL_WORKFLOW_STATE_KEY);
+}
+
 // ============================================================================
 // Settings Storage
 // ============================================================================

tests/unit/messages.test.ts

@@ -454,6 +454,8 @@ describe('Handler Registries', () => {
       'INJECT_CONTENT_SCRIPT',
       'INJECT_ALL_TABS',
       'GET_PANEL_STATE',
+      'SAVE_PANEL_STATE',
+      'CLEAR_PANEL_STATE',
       'GENERATE_NATIVE_PDF',
       'FETCH_IMAGE_AS_DATA_URL',
     ];

tests/unit/misc-handlers.test.ts

@@ -299,7 +299,7 @@ describe('handleGetPanelState', () => {
     expect(mockSendResponse).toHaveBeenCalledWith(
       expect.objectContaining({
         success: true,
-        data: { ready: true },
+        data: { ready: true, workflowState: null },
       })
     );
   });