Support multiple windows in layers (#728)

apps/studio/src/lib/editor/engine/ast/index.ts

@@ -1,41 +1,47 @@
-import { makeAutoObservable } from 'mobx';
-import { TemplateNodeMap } from './map';
 import { EditorAttributes, MainChannels } from '@onlook/models/constants';
+import type { LayerNode, TemplateNode } from '@onlook/models/element';
+import { makeAutoObservable } from 'mobx';
+import { AstRelationshipManager } from './map';
 import { getUniqueSelector, isOnlookInDoc } from '/common/helpers';
 import { getTemplateNode } from '/common/helpers/template';
-import type { LayerNode } from '@onlook/models/element';
-import type { TemplateNode } from '@onlook/models/element';
 
 export class AstManager {
-    private doc: Document | undefined;
-    private displayLayers: LayerNode[] = [];
-    templateNodeMap: TemplateNodeMap = new TemplateNodeMap();
+    private relationshipMap: AstRelationshipManager = new AstRelationshipManager();
 
     constructor() {
         makeAutoObservable(this);
     }
 
     get layers() {
-        return this.displayLayers;
+        return this.relationshipMap.getRootLayers();
     }
 
-    set layers(layers: LayerNode[]) {
-        this.displayLayers = layers;
+    setLayers(webviewId: string, layer: LayerNode) {
+        this.relationshipMap.setRootLayer(webviewId, layer);
     }
 
-    replaceElement(selector: string, newNode: LayerNode) {
-        const element = this.doc?.querySelector(selector);
+    replaceElement(webviewId: string, selector: string, newNode: LayerNode) {
+        const doc = this.relationshipMap.getDocument(webviewId);
+        const element = doc?.querySelector(selector);
         if (!element) {
             console.warn('Failed to replaceElement: Element not found');
             return;
         }
+
         const parent = element.parentElement;
         if (!parent) {
             console.warn('Failed to replaceElement: Parent not found');
             return;
         }
+
+        const rootNode = this.relationshipMap.getRootLayer(webviewId);
+        if (!rootNode) {
+            console.warn('Failed to replaceElement: Root node not found');
+            return;
+        }
+
         const parentSelector = getUniqueSelector(parent, parent.ownerDocument.body);
-        const parentNode = this.findInLayersTree(parentSelector, this.displayLayers[0]);
+        const parentNode = this.findInLayersTree(parentSelector, rootNode);
         if (!parentNode || !parentNode.children) {
             console.warn('Failed to replaceElement: Parent node not found');
             return;
@@ -48,7 +54,7 @@ export class AstManager {
             parentNode.children = parentNode.children?.filter((child) => child.id !== selector);
         }
 
-        this.processNode(parent as HTMLElement);
+        this.processNode(webviewId, parent as HTMLElement);
     }
 
     findInLayersTree(selector: string, node: LayerNode | undefined): LayerNode | undefined {
@@ -74,30 +80,34 @@ export class AstManager {
     }
 
     getInstance(selector: string): TemplateNode | undefined {
-        return this.templateNodeMap.getInstance(selector);
+        return this.relationshipMap.getTemplateInstance(selector);
     }
 
     getRoot(selector: string): TemplateNode | undefined {
-        return this.templateNodeMap.getRoot(selector);
+        return this.relationshipMap.getTemplateRoot(selector);
+    }
+
+    getWebviewId(selector: string): string | undefined {
+        return this.relationshipMap.getWebviewId(selector);
     }
 
-    setDoc(doc: Document) {
-        this.doc = doc;
+    setDoc(webviewId: string, doc: Document) {
+        this.relationshipMap.setDocument(webviewId, doc);
     }
 
-    setMapRoot(rootElement: Element) {
-        this.setDoc(rootElement.ownerDocument);
+    setMapRoot(webviewId: string, rootElement: Element) {
+        this.setDoc(webviewId, rootElement.ownerDocument);
 
         if (isOnlookInDoc(rootElement.ownerDocument)) {
-            this.processNode(rootElement as HTMLElement);
+            this.processNode(webviewId, rootElement as HTMLElement);
         } else {
             console.warn('Page is not Onlook enabled');
         }
     }
 
-    processNode(node: HTMLElement) {
+    processNode(webviewId: string, node: HTMLElement) {
         this.dfs(node as HTMLElement, (node) => {
-            this.processNodeForMap(node as HTMLElement);
+            this.processNodeForMap(webviewId, node as HTMLElement);
         });
     }
 
@@ -115,8 +125,9 @@ export class AstManager {
         }
     }
 
-    private processNodeForMap(node: HTMLElement) {
-        const selector = getUniqueSelector(node, this.doc?.body);
+    private processNodeForMap(webviewId: string, node: HTMLElement) {
+        const doc = this.relationshipMap.getDocument(webviewId);
+        const selector = getUniqueSelector(node, doc?.body);
         if (!selector) {
             return;
         }
@@ -125,12 +136,13 @@ export class AstManager {
             return;
         }
 
-        this.templateNodeMap.setRoot(selector, templateNode);
+        this.relationshipMap.setTemplateRoot(webviewId, selector, templateNode);
         const dataOnlookId = node.getAttribute(EditorAttributes.DATA_ONLOOK_ID) as string;
-        this.findNodeInstance(node, node, templateNode, selector, dataOnlookId);
+        this.findNodeInstance(webviewId, node, node, templateNode, selector, dataOnlookId);
     }
 
     private async findNodeInstance(
+        webviewId: string,
         originalNode: HTMLElement,
         node: HTMLElement,
         templateNode: TemplateNode,
@@ -157,9 +169,10 @@ export class AstManager {
                 { parent: parentTemplateNode, child: templateNode, index },
             );
             if (instance) {
-                this.templateNodeMap.setInstance(selector, instance);
+                this.relationshipMap.setTemplateInstance(webviewId, selector, instance);
             } else {
                 await this.findNodeInstance(
+                    webviewId,
                     originalNode,
                     parent,
                     templateNode,
@@ -171,7 +184,6 @@ export class AstManager {
     }
 
     clear() {
-        this.templateNodeMap = new TemplateNodeMap();
-        this.displayLayers = [];
+        this.relationshipMap = new AstRelationshipManager();
     }
 }

apps/studio/src/lib/editor/engine/ast/map.ts

@@ -1,46 +1,81 @@
-import type { TemplateNode } from '@onlook/models/element';
+import type { LayerNode, TemplateNode } from '@onlook/models/element';
+import { makeAutoObservable } from 'mobx';
 
-export class TemplateNodeMap {
+export class AstRelationshipManager {
     templateToSelectors: Map<TemplateNode, string[]> = new Map();
+
     selectorToInstance: Map<string, TemplateNode> = new Map();
     selectorToRoot: Map<string, TemplateNode> = new Map();
+    selectorToWebviewId: Map<string, string> = new Map();
+
+    webviewIdToDocument: Map<string, Document> = new Map();
+    webviewIdToRootNode: Map<string, LayerNode> = new Map();
 
+    constructor() {
+        makeAutoObservable(this);
+    }
     isProcessed(selector: string): boolean {
         return this.selectorToInstance.has(selector) || this.selectorToRoot.has(selector);
     }
 
     remove(selector: string) {
         this.selectorToInstance.delete(selector);
         this.selectorToRoot.delete(selector);
+        this.selectorToWebviewId.delete(selector);
     }
 
     getSelectors(templateNode: TemplateNode): string[] {
         return this.templateToSelectors.get(templateNode) || [];
     }
 
-    getInstance(selector: string): TemplateNode | undefined {
+    getTemplateInstance(selector: string): TemplateNode | undefined {
         return this.selectorToInstance.get(selector);
     }
 
-    getRoot(selector: string): TemplateNode | undefined {
+    getTemplateRoot(selector: string): TemplateNode | undefined {
         return this.selectorToRoot.get(selector);
     }
 
-    setSelector(templateNode: TemplateNode, selector: string) {
+    getWebviewId(selector: string): string | undefined {
+        return this.selectorToWebviewId.get(selector);
+    }
+
+    setSelector(webviewId: string, templateNode: TemplateNode, selector: string) {
         const existing = this.templateToSelectors.get(templateNode) || [];
         if (!existing.includes(selector)) {
             existing.push(selector);
             this.templateToSelectors.set(templateNode, existing);
         }
+        this.selectorToWebviewId.set(selector, webviewId);
     }
 
-    setRoot(selector: string, templateNode: TemplateNode) {
+    setTemplateRoot(webviewId: string, selector: string, templateNode: TemplateNode) {
         this.selectorToRoot.set(selector, templateNode);
-        this.setSelector(templateNode, selector);
+        this.setSelector(webviewId, templateNode, selector);
     }
 
-    setInstance(selector: string, templateNode: TemplateNode) {
+    setTemplateInstance(webviewId: string, selector: string, templateNode: TemplateNode) {
         this.selectorToInstance.set(selector, templateNode);
-        this.setSelector(templateNode, selector);
+        this.setSelector(webviewId, templateNode, selector);
+    }
+
+    getRootLayers(): LayerNode[] {
+        return Array.from(this.webviewIdToRootNode.values());
+    }
+
+    getRootLayer(webviewId: string): LayerNode | undefined {
+        return this.webviewIdToRootNode.get(webviewId);
+    }
+
+    getDocument(webviewId: string): Document | undefined {
+        return this.webviewIdToDocument.get(webviewId);
+    }
+
+    setDocument(webviewId: string, doc: Document) {
+        this.webviewIdToDocument.set(webviewId, doc);
+    }
+
+    setRootLayer(webviewId: string, layer: LayerNode) {
+        this.webviewIdToRootNode.set(webviewId, layer);
     }
 }

apps/studio/src/lib/editor/engine/dom/index.ts

@@ -18,7 +18,7 @@ export class DomManager {
     }
 
     setDom(webviewId: string, root: Element) {
-        this.editorEngine.ast.setMapRoot(root);
+        this.editorEngine.ast.setMapRoot(webviewId, root);
         this.webviewToRootElement.set(webviewId, root);
         this.webviewToRootElement = new Map(this.webviewToRootElement);
     }
@@ -30,7 +30,7 @@ export class DomManager {
 
     async refreshAstDoc(webview: WebviewTag) {
         const root = await this.getBodyFromWebview(webview);
-        this.editorEngine.ast.setDoc(root.ownerDocument);
+        this.editorEngine.ast.setDoc(webview.id, root.ownerDocument);
     }
 
     getElementBySelector(selector: string, webviewId: string) {

apps/studio/src/lib/editor/eventHandler.ts

@@ -1,9 +1,8 @@
+import { WebviewChannels } from '@onlook/models/constants';
+import type { DomElement, LayerNode } from '@onlook/models/element';
 import { debounce } from 'lodash';
 import { EditorMode } from '../models';
 import type { EditorEngine } from './engine';
-import { WebviewChannels } from '@onlook/models/constants';
-import type { DomElement } from '@onlook/models/element';
-import type { LayerNode } from '@onlook/models/element';
 
 export class WebviewEventHandler {
     eventCallbacks: Record<string, (e: any) => void>;
@@ -37,7 +36,7 @@ export class WebviewEventHandler {
             const body = await this.editorEngine.dom.getBodyFromWebview(webview);
             this.editorEngine.dom.setDom(webview.id, body);
             const layerTree = e.args[0] as LayerNode;
-            this.editorEngine.ast.layers = [layerTree as LayerNode];
+            this.editorEngine.ast.setLayers(webview.id, layerTree);
         };
     }
 
@@ -62,7 +61,7 @@ export class WebviewEventHandler {
                 };
                 await this.editorEngine.dom.refreshAstDoc(webview);
                 [...added, ...removed].forEach((layerNode: LayerNode) => {
-                    this.editorEngine.ast.replaceElement(layerNode.id, layerNode);
+                    this.editorEngine.ast.replaceElement(webview.id, layerNode.id, layerNode);
                 });
             },
             1000,
@@ -171,7 +170,7 @@ export class WebviewEventHandler {
     ) {
         this.editorEngine.mode = EditorMode.DESIGN;
         await this.editorEngine.dom.refreshAstDoc(webview);
-        this.editorEngine.ast.replaceElement(layerNode.id, layerNode);
+        this.editorEngine.ast.replaceElement(webview.id, layerNode.id, layerNode);
         this.editorEngine.elements.click([domEl], webview);
     }
 

apps/studio/src/routes/editor/EditPanel/StylesTab/single/TailwindInput.tsx

@@ -1,12 +1,12 @@
 import { useEditorEngine } from '@/components/Context';
-import { Textarea } from '@onlook/ui/textarea';
 import { sendAnalytics } from '@/lib/utils';
-import { observer } from 'mobx-react-lite';
-import { useEffect, useRef, useState } from 'react';
-import { MainChannels } from '@onlook/models/constants';
 import type { CodeDiffRequest } from '@onlook/models/code';
+import { MainChannels } from '@onlook/models/constants';
 import type { TemplateNode } from '@onlook/models/element';
 import { Icons } from '@onlook/ui/icons';
+import { Textarea } from '@onlook/ui/textarea';
+import { observer } from 'mobx-react-lite';
+import { useEffect, useRef, useState } from 'react';
 
 const TailwindInput = observer(() => {
     const editorEngine = useEditorEngine();
@@ -47,7 +47,7 @@ const TailwindInput = observer(() => {
         if (root) {
             const rootClasses: string[] = await window.api.invoke(
                 MainChannels.GET_TEMPLATE_NODE_CLASS,
-                root,
+                JSON.parse(JSON.stringify(root)),
             );
             setRootClasses(rootClasses.join(' '));
         }

apps/studio/src/routes/editor/LayersPanel/LayersTab.tsx

@@ -1,4 +1,5 @@
 import { useEditorEngine } from '@/components/Context';
+import type { LayerNode } from '@onlook/models/element';
 import { observer } from 'mobx-react-lite';
 import { useEffect, useRef, useState } from 'react';
 import { type NodeApi, Tree, type TreeApi } from 'react-arborist';
@@ -7,14 +8,12 @@ import RightClickMenu from '../RightClickMenu';
 import TreeNode from './Tree/TreeNode';
 import TreeRow from './Tree/TreeRow';
 import { escapeSelector } from '/common/helpers';
-import type { LayerNode } from '@onlook/models/element';
 
 const LayersTab = observer(() => {
     const treeRef = useRef<TreeApi<LayerNode>>();
     const editorEngine = useEditorEngine();
     const [treeHovered, setTreeHovered] = useState(false);
     const { ref, width, height } = useResizeObserver();
-    const domTree = editorEngine.ast.layers;
 
     useEffect(handleSelectChange, [editorEngine.elements.selected]);
 
@@ -121,7 +120,7 @@ const LayersTab = observer(() => {
             <RightClickMenu>
                 <Tree
                     ref={treeRef}
-                    data={domTree}
+                    data={editorEngine.ast.layers}
                     openByDefault={true}
                     overscanCount={1}
                     indent={8}