Preserve tiling layouts across screen lock

README.md

@@ -262,10 +262,10 @@ How to fix it:
 
 O-Tiling has no `session-modes` in `metadata.json`, so it uses GNOME's default `"user"` mode. This means GNOME fully disables and re-enables the extension on every lock/unlock (and on suspend, since suspend locks the screen).
 
-- `disable()`: disconnects all signals, destroys the `Ext` instance, removes the indicator/overlay.
-- `enable()`: rebuilds everything. A new `Ext` instance, signals reattached, tiling tree reconstructed.
+- `disable()`: when locking, saves a plain in-memory layout snapshot, then disconnects all signals, destroys the `Ext` instance, and removes the indicator/overlay.
+- `enable()`: creates a new `Ext` instance and reconnects signals. On unlock, restores saved split directions/proportions, stack order/active tabs, and manual floating overrides using session-stable window IDs. Closed or moved windows are omitted, new windows are tiled normally, and changed monitor work areas fall back to spatial reconstruction.
 
-This is the standard, EGO-compliant approach (no special lock-screen permission needed), unlike Pop Shell or Forge, which skip or bypass this teardown to preserve window layout across a lock. The trade-off is that O-Tiling's full rebuild runs on every lock, not just occasional suspends, so debug suspend/focus issues with:
+This is the standard, EGO-compliant approach (no special lock-screen permission needed), unlike Pop Shell or Forge, which skip or bypass this teardown to preserve window layout across a lock. The teardown still runs on every lock. The snapshot retains no live window/actor references and is cleared after restoration or ordinary extension disable; it does not persist across logout or reboot. Debug suspend/focus issues with:
 
 ```
 journalctl --user -f -o cat | grep -i o-tiling
@@ -332,3 +332,9 @@ Found a bug or have an idea? Open an issue or pull request on [GitHub](https://g
 ---
 
 Happy tiling.
+
+### Layout snapshot regression tests
+
+Run `pnpm test` with Node.js 24 or newer. Tests cover snapshot pruning, interrupted unlocks, repeated rebuilds, split ratios, stack order, floating overrides, and changed work-area fallback. Adapter tests use lightweight Shell doubles; they do not substitute for a real GNOME lock/unlock test.
+
+Manual validation: create nested splits and a stack, resize a split, and float a window. Lock and unlock several times, checking positions, split sizes, tab order, floating state, and subsequent window moves/resizes. Repeat with a closed/new window and with suspend/resume. Changed monitor configurations should use the existing reconstruction path.

package.json

@@ -13,7 +13,8 @@
     "lint:dist": "../GER/ego-lint ./dist",
     "shexli": "make shexli",
     "log": "journalctl -f -o cat /usr/bin/gnome-shell | grep o-tiling",
-    "local-install": "make install"
+    "local-install": "make install",
+    "test": "node tests/layout-snapshot.test.mjs && node --import ./tests/register-shell-loader.mjs tests/lock-layout.test.mjs"
   },
   "devDependencies": {
     "@eslint/js": "^10.0.1",

src/engine/layout_snapshot.ts

@@ -0,0 +1,62 @@
+/** Plain session-local data: never retain Meta.Window, actors, signals or ECS entities. */
+export type LayoutNode =
+    | { kind: 'window'; id: number }
+    | { kind: 'stack'; ids: number[]; active: number }
+    | { kind: 'split'; horizontal: boolean; ratio: number; left: LayoutNode; right: LayoutNode };
+
+export interface WorkArea { x: number; y: number; width: number; height: number }
+export interface WorkspaceLayout {
+    monitor: number;
+    workspace: number;
+    area: WorkArea;
+    tree: LayoutNode;
+}
+export interface LayoutSnapshot {
+    layouts: WorkspaceLayout[];
+    floating: number[];
+    forceTiled: number[];
+}
+
+export function sameArea(a: WorkArea, b: WorkArea): boolean {
+    return a.x === b.x && a.y === b.y && a.width === b.width && a.height === b.height;
+}
+
+/** Remove vanished/moved windows without changing surviving branch or tab order. */
+export function pruneTree(tree: LayoutNode, available: ReadonlySet<number>): LayoutNode | null {
+    switch (tree.kind) {
+        case 'window': return available.has(tree.id) ? { ...tree } : null;
+        case 'stack': {
+            const ids = tree.ids.filter(id => available.has(id));
+            if (!ids.length) return null;
+            if (ids.length === 1) return { kind: 'window', id: ids[0] };
+            return { kind: 'stack', ids, active: ids.includes(tree.active) ? tree.active : ids[0] };
+        }
+        case 'split': {
+            const left = pruneTree(tree.left, available);
+            const right = pruneTree(tree.right, available);
+            return left && right ? { ...tree, left, right } : left ?? right;
+        }
+    }
+}
+
+export function snapshotIds(tree: LayoutNode | null): number[] {
+    if (!tree) return [];
+    if (tree.kind === 'window') return [tree.id];
+    if (tree.kind === 'stack') return [...tree.ids];
+    return [...snapshotIds(tree.left), ...snapshotIds(tree.right)];
+}
+
+/** A second lock can cancel the deferred restore; retain data until it completes. */
+export class PendingLayout {
+    private pending: LayoutSnapshot | null = null;
+
+    peek(): LayoutSnapshot | null { return this.pending; }
+
+    save(snapshot: LayoutSnapshot | null, locking: boolean): void {
+        this.pending = locking ? snapshot ?? this.pending : null;
+    }
+
+    restored(snapshot: LayoutSnapshot | null): void {
+        if (this.pending === snapshot) this.pending = null;
+    }
+}

src/engine/lock_layout.ts

@@ -0,0 +1,128 @@
+import type { Ext } from '../extension.js';
+import type { AutoTiler } from './auto_tiler.js';
+import type { Entity } from '../core/ecs.js';
+import type { ShellWindow } from '../window/window.js';
+import type { LayoutNode, LayoutSnapshot, WorkspaceLayout } from './layout_snapshot.js';
+import { pruneTree, sameArea, snapshotIds } from './layout_snapshot.js';
+import { Node, NodeKind } from './node.js';
+import type { NodeStack } from './node.js';
+import { Stack } from './stack.js';
+import { Orientation } from '../utils/lib.js';
+import * as Tags from '../utils/tags.js';
+
+export function captureLayout(ext: Ext): LayoutSnapshot {
+    const tiler = ext.auto_tiler!;
+    const windowId = (entity: Entity): number | null =>
+        ext.windows.get(entity)?.meta.get_stable_sequence() ?? null;
+    const capture = (branch: Node): LayoutNode | null => {
+        const inner = branch.inner;
+        if (inner.kind === NodeKind.WINDOW) {
+            const id = windowId(inner.entity);
+            return id === null ? null : { kind: 'window', id };
+        }
+        if (inner.kind === NodeKind.STACK) {
+            const ids = inner.entities.map(windowId).filter((id): id is number => id !== null);
+            if (!ids.length) return null;
+            const active = tiler.forest.stacks.get(inner.idx)?.active;
+            return { kind: 'stack', ids, active: active ? windowId(active) ?? ids[0] : ids[0] };
+        }
+        const fork = tiler.forest.forks.get(inner.entity);
+        if (!fork) return null;
+        const left = capture(fork.left);
+        const right = fork.right ? capture(fork.right) : null;
+        if (!left || !right) return left ?? right;
+        return {
+            kind: 'split', horizontal: fork.is_horizontal(),
+            ratio: fork.length_left / fork.length(), left, right,
+        };
+    };
+    const layouts: WorkspaceLayout[] = [];
+    for (const [entity, [monitor, workspace]] of tiler.forest.toplevel.values()) {
+        const tree = capture(Node.fork(entity));
+        if (!tree) continue;
+        const area = ext.monitor_work_area(monitor);
+        layouts.push({ monitor, workspace, area: { x: area.x, y: area.y, width: area.width, height: area.height }, tree });
+    }
+    const floating: number[] = [];
+    const forceTiled: number[] = [];
+    for (const win of ext.windows.values()) {
+        if (ext.contains_tag(win.entity, Tags.Floating)) floating.push(win.meta.get_stable_sequence());
+        if (ext.contains_tag(win.entity, Tags.ForceTile)) forceTiled.push(win.meta.get_stable_sequence());
+    }
+    return { layouts, floating, forceTiled };
+}
+
+/** Restore only groups whose display/workspace still matches; other groups use normal reconstruction. */
+export function restoreWorkspace(
+    tiler: AutoTiler, ext: Ext, saved: WorkspaceLayout, windows: ShellWindow[],
+): number[] {
+    const area = ext.monitor_work_area(saved.monitor);
+    if (!sameArea(saved.area, area)) return [];
+    const byId = new Map(windows.map(win => [win.meta.get_stable_sequence(), win]));
+    const tree = pruneTree(saved.tree, new Set(byId.keys()));
+    if (!tree) return [];
+    const rootArea = area.clone();
+    if (!ext.settings.smart_gaps() || snapshotIds(tree).length > 1) {
+        rootArea.x += ext.gap_outer;
+        rootArea.y += ext.gap_top;
+        rootArea.width -= ext.gap_outer * 2;
+        rootArea.height -= ext.gap_outer + ext.gap_top;
+    }
+    const stacks: NodeStack[] = [];
+    const build = (savedNode: LayoutNode, rect: typeof rootArea): Node => {
+        if (savedNode.kind === 'window') return Node.window(byId.get(savedNode.id)!.entity);
+        if (savedNode.kind === 'stack') {
+            const active = byId.get(savedNode.active)!;
+            const container = new Stack(ext, active.entity, saved.workspace, saved.monitor);
+            container.active_id = savedNode.ids.indexOf(savedNode.active);
+            const idx = tiler.forest.stacks.insert(container);
+            const result = Node.stacked(active.entity, idx);
+            const inner = result.inner as NodeStack;
+            inner.entities = savedNode.ids.map(id => byId.get(id)!.entity);
+            inner.rect = rect.clone();
+            stacks.push(inner);
+            return result;
+        }
+        const leftRect = rect.clone();
+        const rightRect = rect.clone();
+        const gap = ext.gap_inner_half;
+        const split = (savedNode.horizontal ? rect.width : rect.height) * savedNode.ratio;
+        if (savedNode.horizontal) {
+            leftRect.width = split - gap;
+            rightRect.x += split + gap;
+            rightRect.width -= split + gap;
+        } else {
+            leftRect.height = split - gap;
+            rightRect.y += split + gap;
+            rightRect.height -= split + gap;
+        }
+        const left = build(savedNode.left, leftRect);
+        const right = build(savedNode.right, rightRect);
+        const [entity, fork] = tiler.forest.create_fork(left, right, rect, saved.workspace, saved.monitor);
+        fork.orientation = savedNode.horizontal ? Orientation.HORIZONTAL : Orientation.VERTICAL;
+        fork.set_ratio(split);
+        fork.prev_ratio = savedNode.ratio;
+        link(entity, left);
+        link(entity, right);
+        return Node.fork(entity);
+    };
+    const link = (parent: Entity, child: Node) => {
+        if (child.inner.kind === NodeKind.FORK) tiler.forest.parents.insert(child.inner.entity, parent);
+        else if (child.inner.kind === NodeKind.WINDOW) ext.on_tile_attach(parent, child.inner.entity);
+        else for (const entity of child.inner.entities) ext.on_tile_attach(parent, entity);
+    };
+    let root = build(tree, rootArea);
+    if (root.inner.kind !== NodeKind.FORK) {
+        const [entity] = tiler.forest.create_fork(root, null, rootArea, saved.workspace, saved.monitor);
+        link(entity, root);
+        root = Node.fork(entity);
+    }
+    if (root.inner.kind !== NodeKind.FORK) return [];
+    const entity = root.inner.entity;
+    const fork = tiler.forest.forks.get(entity)!;
+    fork.is_toplevel = true;
+    tiler.forest.toplevel.set(`${entity}`, [entity, [saved.monitor, saved.workspace]]);
+    for (const stack of stacks) tiler.update_stack(ext, stack);
+    tiler.update_toplevel(ext, fork, saved.monitor, ext.settings.smart_gaps());
+    return snapshotIds(tree);
+}

src/extension.ts

@@ -15,6 +15,9 @@ import * as Tiling from './engine/tiling.js';
 import * as Window from './window/window.js';
 import * as auto_tiler from './engine/auto_tiler.js';
 import * as reconstruct from './engine/reconstruct.js';
+import { captureLayout, restoreWorkspace } from './engine/lock_layout.js';
+import { PendingLayout } from './engine/layout_snapshot.js';
+import type { LayoutSnapshot } from './engine/layout_snapshot.js';
 import * as node from './engine/node.js';
 import * as utils from './utils/utils.js';
 import * as add_exception from './ui/dialog_add_exception.js';
@@ -64,6 +67,8 @@ import GObject from 'gi://GObject';
 import Mtk from 'gi://Mtk';
 const { GlobalEvent, WindowEvent } = Events;
 
+const lockLayout = new PendingLayout();
+
 export let ext: Ext | null = null;
 export let indicator: Indicator | null = null;
 export let workspace_number_indicator: WorkspaceNumberIndicator | null = null;
@@ -3307,7 +3312,7 @@ export class Ext extends Ecs.System<ExtEvent> {
         }
     }
 
-    auto_tile_on(save_setting: boolean = true, reposition: boolean = true) {
+    auto_tile_on(save_setting: boolean = true, reposition: boolean = true, snapshot: LayoutSnapshot | null = null) {
         // Do not auto-tile if the extension is soft-disabled
         if (this._ext_soft_disabled) return;
 
@@ -3347,6 +3352,17 @@ export class Ext extends Ecs.System<ExtEvent> {
             quick_settings_indicator.updateIcon(this.button_gio_icon_auto_on);
         }
 
+        // Restore session-local floating overrides before selecting tilable windows.
+        if (snapshot) {
+            const floating = new Set(snapshot.floating);
+            const forceTiled = new Set(snapshot.forceTiled);
+            for (const win of this.windows.values()) {
+                const id = win.meta.get_stable_sequence();
+                if (floating.has(id)) this.add_tag(win.entity, Tags.Floating);
+                if (forceTiled.has(id)) this.add_tag(win.entity, Tags.ForceTile);
+            }
+        }
+
         // Group tileable windows by [monitor, workspace] key
         const groups = new Map<string, Window.ShellWindow[]>();
         for (const window of this.windows.values()) {
@@ -3380,7 +3396,16 @@ export class Ext extends Ecs.System<ExtEvent> {
             const [monitor_s, workspace_s] = key.split(',');
             const monitor = parseInt(monitor_s, 10);
             const workspace = parseInt(workspace_s, 10);
-            reconstruct.reconstruct_workspace(tiler, this, monitor, workspace, group);
+            const saved = snapshot?.layouts.find(layout => layout.monitor === monitor && layout.workspace === workspace);
+            const restored = saved ? restoreWorkspace(tiler, this, saved, group) : [];
+            if (restored.length) {
+                const ids = new Set(restored);
+                for (const win of group) {
+                    if (!ids.has(win.meta.get_stable_sequence())) tiler.auto_tile(this, win, true);
+                }
+            } else {
+                reconstruct.reconstruct_workspace(tiler, this, monitor, workspace, group);
+            }
         }
 
         if (reposition) {
@@ -3391,6 +3416,7 @@ export class Ext extends Ecs.System<ExtEvent> {
         }
 
         this.show_border_on_focused();
+        lockLayout.restored(snapshot);
 
         this.register_fn(() => {
             this._batch_moving = false;
@@ -3848,6 +3874,7 @@ export default class OTilingExtension extends Extension {
     enable() {
         globalThis.oTilingExtension = this;
         log.info('enable');
+        const snapshot = lockLayout.peek();
 
         if (!ext) {
             ext = new Ext();
@@ -3915,7 +3942,7 @@ export default class OTilingExtension extends Extension {
                 }
 
                 const id = GLib.timeout_add(GLib.PRIORITY_DEFAULT, 700, () => {
-                    ext?.auto_tile_on(false, false);
+                    ext?.auto_tile_on(false, false, snapshot);
                     if (ext && ext._timeouts['first_startup_tile'] === id) {
                         ext._timeouts['first_startup_tile'] = null;
                     }
@@ -3924,7 +3951,7 @@ export default class OTilingExtension extends Extension {
                 ext._timeouts['first_startup_tile'] = id;
             } else {
                 // Reconstruct tiling tree from current window positions without repositioning.
-                ext.auto_tile_on(false, false);
+                ext.auto_tile_on(false, false, snapshot);
             }
         }
 
@@ -3940,6 +3967,16 @@ export default class OTilingExtension extends Extension {
 
         log.info('disable');
 
+        // Preserve data only across screen lock; ordinary disable starts fresh.
+        lockLayout.save(null, sessionMode.isLocked);
+        if (sessionMode.isLocked && ext?.auto_tiler && ext.settings.tile_by_default()) {
+            try {
+                lockLayout.save(captureLayout(ext), true);
+            } catch (error) {
+                log.warn(`Could not snapshot layout before lock: ${error}`);
+            }
+        }
+
         if (_osk_startup_signal) {
             layoutManager.disconnect(_osk_startup_signal);
             _osk_startup_signal = 0;

tests/layout-snapshot.test.mjs

@@ -0,0 +1,45 @@
+import {test} from 'node:test';
+import assert from 'node:assert/strict';
+import {pruneTree, sameArea, snapshotIds} from '../src/engine/layout_snapshot.ts';
+const win=id=>({kind:'window',id});
+const tree={kind:'split',horizontal:true,ratio:0.37,left:win(11),right:{kind:'split',horizontal:false,ratio:0.62,left:win(22),right:{kind:'stack',ids:[33,44],active:44}}};
+test('retains nested split ratios and stack order without keeping object references',()=>{
+ const restored=pruneTree(tree,new Set([11,22,33,44]));
+ assert.deepEqual(restored,tree);
+ assert.notEqual(restored,tree);
+ assert.notEqual(restored.right.right.ids,tree.right.right.ids);
+});
+test('removes closed windows and collapses empty split branches',()=>{
+ assert.deepEqual(pruneTree(tree,new Set([22,44])),{kind:'split',horizontal:false,ratio:0.62,left:win(22),right:win(44)});
+});
+test('retains stack order and selects a surviving active tab',()=>{
+ assert.deepEqual(pruneTree({kind:'stack',ids:[3,1,2],active:1},new Set([2,3])),{kind:'stack',ids:[3,2],active:3});
+});
+test('all windows closed produces no stale layout',()=>assert.equal(pruneTree(tree,new Set()),null));
+test('does not absorb new windows into the saved tree',()=>assert.deepEqual(snapshotIds(pruneTree(tree,new Set([11,22,33,44,55]))),[11,22,33,44]));
+test('detects monitor work-area changes before restoring',()=>{
+ const a={x:0,y:32,width:2560,height:1048};
+ assert.ok(sameArea(a,{...a}));
+ for(const key of Object.keys(a)) assert.ok(!sameArea(a,{...a,[key]:a[key]+1}));
+});
+test('closed windows on repeated unlocks cannot reappear',()=>{
+ const first=pruneTree(tree,new Set([11,22,44]));
+ assert.deepEqual(pruneTree(first,new Set([44])),win(44));
+});
+test('interrupted unlock retains its snapshot until restoration completes',async()=>{
+ const {PendingLayout}=await import('../src/engine/layout_snapshot.ts');
+ const state=new PendingLayout(); const snapshot={layouts:[],floating:[11],forceTiled:[]};
+ state.save(snapshot,true);
+ const firstUnlock=state.peek();
+ // Lock again before the startup callback can build a tiler.
+ state.save(null,true);
+ assert.equal(state.peek(),snapshot);
+ state.restored(firstUnlock);
+ assert.equal(state.peek(),null);
+});
+test('ordinary disable clears stale state and an old callback cannot clear a newer snapshot',async()=>{
+ const {PendingLayout}=await import('../src/engine/layout_snapshot.ts');
+ const state=new PendingLayout();const a={layouts:[],floating:[1],forceTiled:[]};const b={layouts:[],floating:[2],forceTiled:[]};
+ state.save(a,true);state.save(b,true);state.restored(a);assert.equal(state.peek(),b);
+ state.save(null,false);assert.equal(state.peek(),null);
+});