fix(optimize): judge an expand by the children it attaches

[justi]

Problem

Flash indexing with optimize="full" sometimes fails with:

RuntimeError: node '8.3.2 NTK-Preactivation Cross Correlations' was dropped or changed after it was marked final

expand accepts a split when its lookahead cost (expand_cost) beats a linear scan. The lookahead leaves out what attach_children adds afterwards: when the first proposed heading sits mid-page, the intro node runs onto that page, and same-page siblings are fused. The attached node can then cost as much as the scan. The next round's merge (span <= tree_cost) folds it back after its children were already marked final, and summarize_tree refuses the changed tree.

Example from a 471-page book (arXiv 2106.10165): node 8.3.2 spans pages 224-229 (S=6); the model proposes headings on pages 228 (mid-page) and 229. expand_cost = 5, so the node is expanded; the intro covers 224-228, so tree_cost = 6 and round 2 merges it back.

Fix

Before committing an expand, attach the children to a copy of the node, fuse same-page siblings, and keep the expand only if that copy's tree_cost beats the scan - the same test merge applies next round.

Verification

• New test tests/test_expand_refold.py reproduces the failure on main and passes with the fix; an expand that pays off is still taken.
• I recorded the model's expand replies in 8 real runs on the book (local gemma4 E4B via Ollama) and replayed them: on main 2 of them fail with the error above; with the fix all 8 pass, and the 6 that passed before give identical trees.
• 20 runs with random expand replies: identical trees with and without the fix.

🤖 Generated with Claude Code