Aipex Performance Optimization: Making AI Smarter at Understanding Web Pages

// Multiple separate calls - inefficient and exposes too much data
async function getPageDataTraditional() {
  // Call 1: Get all page content
  const pageContent = await getPageContent();
  // Returns: Full HTML, all styles, all attributes, selectors

  // Call 2: Get interactive elements
  const interactiveElements = await getInteractiveElements();
  // Returns: Complex objects with selectors, styles, positions

  // Call 3: Get page links
  const pageLinks = await getPageLinks();
  // Returns: All links and their attributes

  return {
    content: pageContent,      // ~50KB of data
    elements: interactiveElements, // ~30KB of data
    links: pageLinks          // ~15KB of data
  };
  // Total: ~95KB of data with sensitive information exposed
}

// Real CDP implementation from Aipex's actual code
/**
 * Get REAL accessibility tree using Chrome DevTools Protocol
 * This is the ACTUAL browser's native accessibility tree - exactly like Puppeteer's page.accessibility.snapshot()
 */
async function getRealAccessibilityTree(tabId: number): Promise<AccessibilityTree | null> {
  return new Promise(async (resolve, reject) => {
    console.log('🔍 [DEBUG] Connecting to tab via Chrome DevTools Protocol:', tabId);

    // Safely attach debugger to the tab
    const attached = await safeAttachDebugger(tabId);
    if (!attached) {
      reject(new Error('Failed to attach debugger'));
      return;
    }

    // STEP 1: Enable accessibility domain - REQUIRED for consistent AXNodeIds
    chrome.debugger.sendCommand({ tabId }, "Accessibility.enable", {}, () => {
      if (chrome.runtime.lastError) {
        console.error('❌ [DEBUG] Failed to enable Accessibility domain:', chrome.runtime.lastError.message);
        safeDetachDebugger(tabId);
        reject(new Error(`Failed to enable Accessibility domain: ${chrome.runtime.lastError.message}`));
        return;
      }

      console.log('✅ [DEBUG] Accessibility domain enabled');

      // STEP 2: Get the full accessibility tree
      // This is the same as Puppeteer's page.accessibility.snapshot()
      chrome.debugger.sendCommand({ tabId }, "Accessibility.getFullAXTree", {
        // depth: undefined - get full tree (not just top level)
        // frameId: undefined - get main frame
      }, (result: any) => {
        if (chrome.runtime.lastError) {
          console.error('❌ [DEBUG] Failed to get accessibility tree:', chrome.runtime.lastError.message);
          // Disable accessibility before detaching
          chrome.debugger.sendCommand({ tabId }, "Accessibility.disable", {}, () => {
            safeDetachDebugger(tabId);
          });
          reject(new Error(`Failed to get accessibility tree: ${chrome.runtime.lastError.message}`));
          return;
        }

        console.log('✅ [DEBUG] Got accessibility tree with', result.nodes?.length || 0, 'nodes');

        // STEP 3: Disable accessibility and detach debugger
        chrome.debugger.sendCommand({ tabId }, "Accessibility.disable", {}, () => {
          // Add a small delay to ensure accessibility is properly disabled
          setTimeout(() => {
            safeDetachDebugger(tabId);
          }, 100);
        });

        resolve(result);
      });
    });
  });
}

/**
 * Convert CDP accessibility tree to Puppeteer-like SerializedAXNode tree
 * This uses Puppeteer's TWO-PASS approach: collect interesting nodes, then serialize
 */
function convertAccessibilityTreeToSnapshot(
  snapshotResult: any,
  snapshotId: string
): TextSnapshotNode | null {
  const nodes = snapshotResult.nodes;
  if (!nodes || nodes.length === 0) {
    return null;
  }

  console.log('🔍 [DEBUG] Processing', nodes.length, 'raw CDP nodes');

  // Build nodeId -> AXNode map
  const nodeMap = new Map<string, AXNode>();
  for (const node of nodes) {
    nodeMap.set(node.nodeId, node);
  }

  // Find root (no parentId)
  const rootNode = nodes.find((n: AXNode) => !n.parentId);
  if (!rootNode) {
    return null;
  }

  // PASS 1: Collect interesting nodes (Puppeteer's approach)
  const interestingNodes = new Set<string>(); // Store nodeIds

  function collectInterestingNodes(nodeId: string) {
    const node = nodeMap.get(nodeId);
    if (!node) return;

    // Check if this node is interesting (same logic as Puppeteer)
    if (isInterestingNode(node)) {
      interestingNodes.add(nodeId);
    }

    // Recursively check children
    if (node.childIds) {
      for (const childId of node.childIds) {
        collectInterestingNodes(childId);
      }
    }
  }

  // Start from root
  collectInterestingNodes(rootNode.nodeId);

  // PASS 2: Build tree structure with only interesting nodes
  const idToNode = new Map<string, TextSnapshotNode>();

  function buildSnapshotNode(nodeId: string): TextSnapshotNode | null {
    const node = nodeMap.get(nodeId);
    if (!node || !interestingNodes.has(nodeId)) {
      return null;
    }

    const snapshotNode: TextSnapshotNode = {
      id: nodeId,
      role: node.role?.value || 'unknown',
      name: node.name?.value,
      value: node.value?.value,
      description: node.description?.value,
      children: [],
      backendDOMNodeId: node.backendDOMNodeId
    };

    // Add children
    if (node.childIds) {
      for (const childId of node.childIds) {
        const child = buildSnapshotNode(childId);
        if (child) {
          snapshotNode.children.push(child);
        }
      }
    }

    idToNode.set(nodeId, snapshotNode);
    return snapshotNode;
  }

  const root = buildSnapshotNode(rootNode.nodeId);
  console.log(`✅ [DEBUG] Built accessibility tree with ${idToNode.size} interesting nodes`);

  return root;
}

/**
 * Check if a node is "interesting" - optimized for DevTools MCP-like output
 * This matches Puppeteer's interestingOnly: true logic exactly
 */
function isInterestingNode(node: AXNode): boolean {
  // Skip ignored nodes
  if (node.ignored) {
    return false;
  }

  const role = node.role?.value;
  if (!role) return false;

  // Include interactive elements
  const interactiveRoles = [
    'button', 'link', 'textbox', 'checkbox', 'radio', 'combobox',
    'listbox', 'menu', 'menuitem', 'tab', 'tabpanel', 'slider',
    'progressbar', 'spinbutton', 'switch', 'tree', 'treeitem'
  ];

  if (interactiveRoles.includes(role)) {
    return true;
  }

  // Include structural elements with meaningful content
  const structuralRoles = [
    'heading', 'main', 'navigation', 'banner', 'contentinfo',
    'complementary', 'region', 'article', 'section', 'aside'
  ];

  if (structuralRoles.includes(role) && (node.name?.value || node.description?.value)) {
    return true;
  }

  // Include elements with names or descriptions
  if (node.name?.value || node.description?.value) {
    return true;
  }

  return false;
}

// Traditional approach - fragile and unreliable
async function clickElementTraditional(selector: string) {
  // Problem 1: Selectors break when page changes
  const element = await page.$(selector); // "#login-button" might not exist

  // Problem 2: No validation of element state
  if (!element) {
    throw new Error(`Element not found: ${selector}`);
  }

  // Problem 3: No retry mechanism for dynamic content
  await element.click();

  // Problem 4: No way to verify if click was successful
  return { success: true }; // Assumes success
}

async function fillElementTraditional(selector: string, value: string) {
  // Same problems as clicking
  const element = await page.$(selector);
  if (!element) {
    throw new Error(`Element not found: ${selector}`);
  }

  await element.fill(value);
  return { success: true };
}

// Usage - fragile and error-prone
await clickElementTraditional("#login-button"); // Breaks if ID changes
await fillElementTraditional("input[name='email']", "user@example.com"); // Breaks if name changes

// Aipex's snapshot-based approach - reliable and stable
export async function takeSnapshot(): Promise<{
  success: boolean;
  snapshotId: string;
  snapshot: string;
  title: string;
  url: string;
  message?: string;
}> {
  const [tab] = await chrome.tabs.query({ active: true, currentWindow: true })
  if (!tab || typeof tab.id !== "number") return {
    success: false,
    snapshotId: '',
    snapshot: '',
    title: '',
    url: '',
    message: 'No active tab found'
  }

  // Get accessibility tree from browser
  const accessibilityTree = await getRealAccessibilityTree(tab.id);

  if (!accessibilityTree || !accessibilityTree.nodes) {
    return {
      success: false,
      snapshotId: '',
      snapshot: '',
      title: tab.title || '',
      url: tab.url || '',
      message: "Failed to get accessibility tree"
    }
  }

  // Generate unique snapshot ID
  const snapshotId = `snapshot_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;

  // Convert to snapshot format with UID mapping
  const root = convertAccessibilityTreeToSnapshot(accessibilityTree, snapshotId);

  if (!root) {
    return {
      success: false,
      snapshotId: '',
      snapshot: '',
      title: tab.title || '',
      url: tab.url || '',
      message: "Failed to convert accessibility tree"
    }
  }

  // Store snapshot globally - creates UID-to-element mapping
  currentSnapshot = {
    root,
    idToNode,
    snapshotId
  };

  // Format as text for AI consumption
  const snapshotText = formatSnapshotAsText(root);

  return {
    success: true,
    snapshotId,
    snapshot: snapshotText,
    title: tab.title || '',
    url: tab.url || '',
    message: "Snapshot taken successfully"
  }
}

// Click element using UID from snapshot
export async function clickElementByUid(uid: string, dblClick = false): Promise<{
  success: boolean;
  message: string;
  title: string;
  url: string;
}> {
  const tab = await getCurrentTab();

  if (!tab || typeof tab.id !== "number") {
    return {
      success: false,
      message: "No accessible tab found",
      title: "",
      url: ""
    };
  }

  let handle: ElementHandle | null = null;

  try {
    console.log('🔍 [DEBUG] Starting clickElementByUid using snapshot UID mapping for uid:', uid);

    // Step 1: Get element handle using snapshot UID mapping
    handle = await getElementByUid(uid);
    if (!handle) {
      return {
        success: false,
        message: "Element not found in current snapshot. Call take_snapshot first to get fresh element UIDs.",
        title: tab.title || "",
        url: tab.url || ""
      };
    }

    console.log('✅ [DEBUG] Found element handle via snapshot UID mapping');

    // Step 2: Use Locator system to click the element
    await waitForEventsAfterAction(async () => {
      await handle!.asLocator().click({ count: dblClick ? 2 : 1 });
    });

    return {
      success: true,
      message: `Element ${dblClick ? 'double ' : ''}clicked successfully using snapshot UID mapping`,
      title: tab.title || "",
      url: tab.url || ""
    };

  } catch (error) {
    console.error('❌ [DEBUG] Error in clickElementByUid:', error);
    return {
      success: false,
      message: `Error clicking element: ${error instanceof Error ? error.message : 'Unknown error'}`,
      title: tab?.title || "",
      url: tab?.url || ""
    };
  } finally {
    // Clean up resources
    if (handle) {
      handle.dispose();
    }
  }
}

// Fill element using UID from snapshot
export async function fillElementByUid(uid: string, value: string): Promise<{
  success: boolean;
  message: string;
  title: string;
  url: string;
}> {
  const tab = await getCurrentTab();

  if (!tab || typeof tab.id !== "number") {
    return {
      success: false,
      message: "No accessible tab found",
      title: "",
      url: ""
    };
  }

  let handle: ElementHandle | null = null;

  try {
    console.log('🔍 [DEBUG] Starting fillElementByUid using snapshot UID mapping for uid:', uid);

    // Step 1: Get element handle using snapshot UID mapping
    handle = await getElementByUid(uid);
    if (!handle) {
      return {
        success: false,
        message: "Element not found in current snapshot. Call take_snapshot first to get fresh element UIDs.",
        title: tab.title || "",
        url: tab.url || ""
      };
    }

    console.log('✅ [DEBUG] Found element handle via snapshot UID mapping');

    // Step 2: Use Locator system to fill the element
    await waitForEventsAfterAction(async () => {
      await handle!.asLocator().fill(value);
    });

    return {
      success: true,
      message: "Element filled successfully using snapshot UID mapping",
      title: tab.title || "",
      url: tab.url || ""
    };

  } catch (error) {
    console.error('❌ [DEBUG] Error in fillElementByUid:', error);
    return {
      success: false,
      message: `Error filling element: ${error instanceof Error ? error.message : 'Unknown error'}`,
      title: tab?.title || "",
      url: tab?.url || ""
    };
  } finally {
    // Clean up resources
    if (handle) {
      handle.dispose();
    }
  }
}

// Get element by UID - core of the snapshot system
export async function getElementByUid(uid: string): Promise<ElementHandle | null> {
  if (!currentSnapshot) {
    throw new Error('No snapshot available. Call take_snapshot first.');
  }

  // Validate snapshot ID
  const [snapshotId] = uid.split('_');
  if (currentSnapshot.snapshotId !== snapshotId) {
    throw new Error('This uid is coming from a stale snapshot. Call take_snapshot to get a fresh snapshot.');
  }

  // Get node from snapshot
  const node = currentSnapshot.idToNode.get(uid);
  if (!node) {
    throw new Error('No such element found in the snapshot');
  }

  console.log('🔍 [DEBUG] Found node in snapshot for uid:', uid, {
    role: node.role,
    name: node.name,
    description: node.description,
    backendDOMNodeId: node.backendDOMNodeId,
    value: node.value
  });

  // Get current tab
  const tab = await getCurrentTab();
  if (!tab || typeof tab.id !== "number") {
    throw new Error('No accessible tab found');
  }

  // Return ElementHandle if we have backendDOMNodeId
  if (node.backendDOMNodeId) {
    console.log('✅ [DEBUG] Creating SmartElementHandle with backendDOMNodeId:', node.backendDOMNodeId);
    return new SmartElementHandle(tab.id, node, node.backendDOMNodeId);
  }

  return null;
}

// Usage - reliable and stable
const snapshot = await takeSnapshot();
// AI gets: "Login button (uid: snapshot_123_abc_0)"
await clickElementByUid("snapshot_123_abc_0"); // Always works, even if page changes
await fillElementByUid("snapshot_123_abc_1", "user@example.com"); // Stable reference

// Inefficient approach - all snapshots sent to AI
async function runChatWithTools(userMessages: any[], messageId?: string) {
  let messages = [systemPrompt, ...userMessages]

  // AI calls take_snapshot multiple times during conversation
  // Each snapshot gets added to conversation history
  while (hasToolCalls) {
    for (const toolCall of toolCalls) {
      if (toolCall.name === 'take_snapshot') {
        const result = await executeToolCall(toolCall.name, toolCall.args)

        // PROBLEM: Every snapshot is added to conversation
        messages.push({
          role: 'tool',
          name: 'take_snapshot',
          content: JSON.stringify(result) // Full snapshot data
        })
      }
    }
  }

  // AI receives ALL snapshots - wastes tokens and causes confusion
  return messages; // Contains multiple snapshots of same page
}

// Optimized approach from Aipex's actual implementation
async function runChatWithTools(userMessages: any[], messageId?: string) {
  let messages = [systemPrompt, ...userMessages]

  while (hasToolCalls) {
    for (const toolCall of toolCalls) {
      if (toolCall.name === 'take_snapshot') {
        const result = await executeToolCall(toolCall.name, toolCall.args)

        // Add current snapshot to conversation
        messages.push({
          role: 'tool',
          name: 'take_snapshot',
          content: JSON.stringify(result)
        })

        // CRITICAL: Implement smart deduplication
        if (toolCall.name === 'take_snapshot') {
          const currentTabUrl = result.data?.url || result.url || '';
          const currentSnapshotId = result.data?.snapshotId || result.snapshotId || '';

          // Replace ALL previous take_snapshot results with fake result
          // This ensures only the latest snapshot is real, all previous ones are treated as duplicate calls
          let replacedCount = 0;

          // Go through messages in reverse order to find all previous real snapshots
          for (let i = messages.length - 1; i >= 0; i--) {
            const msg = messages[i];
            if (msg.role === 'tool' && msg.name === 'take_snapshot') {
              try {
                const content = JSON.parse(msg.content);
                const existingUrl = content.data?.url || content.url || '';
                const existingSnapshotId = content.data?.snapshotId || content.snapshotId || '';

                if (!content.skipped) {
                  // Replace this real snapshot with fake result
                  replacedCount++;
                  messages[i] = {
                    ...msg,
                    content: JSON.stringify({
                      skipped: true,
                      reason: "replaced_by_later_snapshot",
                      url: existingUrl,
                      originalSnapshotId: existingSnapshotId,
                      message: "This snapshot was replaced by a later snapshot (duplicate call)"
                    })
                  };
                }
              } catch {
                // Keep if parsing fails
              }
            }
          }

          if (replacedCount > 0) {
            console.log(`🔄 [SNAPSHOT DEDUP] Replaced ${replacedCount} previous snapshot(s) with fake result`);
            console.log(`🔄 [SNAPSHOT DEDUP] Keeping latest snapshot - URL: ${currentTabUrl}, ID: ${currentSnapshotId}`);
          }
        }
      }
    }
  }

  // AI only receives the latest snapshot - saves tokens and prevents confusion
  return messages; // Contains only the most recent snapshot
}

// Global snapshot storage - only one current snapshot
let currentSnapshot: TextSnapshot | null = null;

export async function takeSnapshot(): Promise<{
  success: boolean;
  snapshotId: string;
  snapshot: string;
  title: string;
  url: string;
  message?: string;
}> {
  const [tab] = await chrome.tabs.query({ active: true, currentWindow: true })
  if (!tab || typeof tab.id !== "number") return {
    success: false,
    snapshotId: '',
    snapshot: '',
    title: '',
    url: '',
    message: 'No active tab found'
  }

  // Get accessibility tree from browser
  const accessibilityTree = await getRealAccessibilityTree(tab.id);

  if (!accessibilityTree || !accessibilityTree.nodes) {
    return {
      success: false,
      snapshotId: '',
      snapshot: '',
      title: tab.title || '',
      url: tab.url || '',
      message: "Failed to get accessibility tree"
    }
  }

  // Generate unique snapshot ID
  const snapshotId = `snapshot_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;

  // Convert to snapshot format
  const root = convertAccessibilityTreeToSnapshot(accessibilityTree, snapshotId);

  if (!root) {
    return {
      success: false,
      snapshotId: '',
      snapshot: '',
      title: tab.title || '',
      url: tab.url || '',
      message: "Failed to convert accessibility tree"
    }
  }

  // Store snapshot globally - REPLACES previous snapshot
  currentSnapshot = {
    root,
    idToNode,
    snapshotId
  };

  // Format as text for AI consumption
  const snapshotText = formatSnapshotAsText(root);

  return {
    success: true,
    snapshotId,
    snapshot: snapshotText,
    title: tab.title || '',
    url: tab.url || '',
    message: "Snapshot taken successfully"
  }
}