From 27f3ea4c1536f320d0583d393598da8d1d834a61 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Wed, 26 Aug 2026 18:11:50 +0800
Subject: [PATCH 01/18] perf: summaries admit the deepest queued node first
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

summarize_tree visits the tree level by level (asyncio's ready queue is
FIFO), so shallow leaves reach the semaphore first and the deepest
leaves last — yet the deepest leaves are the ones still owed a chain of
parent calls, which then run one after another with the pool nearly
idle. On fed-2023 (182 calls) the 64-wide phase lasted ~22 s and the
tail ~33 s, in-flight decaying 36→15→10→6→3→1.

The semaphore becomes a priority gate: waiters are served deepest first
(FIFO within a depth), a ready parent included. Depth counts the calls
left on a node's path to the root, its own included, so this is
longest-remaining-path-first scheduling; a prompt's content does not
depend on admission order, so the summaries are unchanged. A permit
granted to a waiter that is cancelled before it runs is passed on, as
asyncio.Semaphore does; a cancelled waiter still queued is skipped at
the next release. Fewer than one permit is refused at construction, as
asyncio.Semaphore refused a negative count.

Measured on the stored fed-2023 tree, FIFO and priority alternating:
FIFO 62.3 / 57.1 s, priority 54.9 / 48.7 s (-13%); the landed code
re-run 52.2 s. Mean start time by depth flipped from d3 7.8 s < d4
15.1 s < d5 19.5 s to d5 7.1 s < d4 10.4 s.

Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
---
 pageindex/utils.py   | 68 +++++++++++++++++++++++++++++++++++++-------
 tests/test_client.py | 41 ++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 11 deletions(-)

diff --git a/pageindex/utils.py b/pageindex/utils.py
index 83e9d0e79..37b870aa9 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -9,6 +9,8 @@
 import PyPDF2
 import copy
 import asyncio
+import heapq
+from contextlib import asynccontextmanager
 from io import BytesIO
 from dotenv import find_dotenv, load_dotenv
 load_dotenv(find_dotenv(usecwd=True))
@@ -745,6 +747,47 @@ async def generate_summaries_for_structure(structure, model=None):
 SUMMARY_INTRO_MAX_PAGES = 3     # cap on leading pages fed into a parent summary
 
 
+class _PriorityGate:
+    """Semaphore that admits the highest-priority waiter first, FIFO within a priority."""
+
+    def __init__(self, permits):
+        if permits < 1:
+            raise ValueError("permits must be >= 1")
+        self._free = permits
+        self._waiters = []
+        self._seq = 0
+
+    @asynccontextmanager
+    async def slot(self, prio):
+        await self.acquire(prio)
+        try:
+            yield
+        finally:
+            self.release()
+
+    async def acquire(self, prio):
+        if self._free > 0 and not self._waiters:
+            self._free -= 1
+            return
+        fut = asyncio.get_running_loop().create_future()
+        heapq.heappush(self._waiters, (-prio, self._seq, fut))
+        self._seq += 1
+        try:
+            await fut
+        except asyncio.CancelledError:
+            if fut.done() and not fut.cancelled():
+                self.release()  # granted while cancelling: pass the permit on
+            raise
+
+    def release(self):
+        while self._waiters:
+            _, _, fut = heapq.heappop(self._waiters)
+            if not fut.done():
+                fut.set_result(None)  # the permit moves straight to this waiter
+                return
+        self._free += 1
+
+
 def get_intro_text(node, pdf_pages, max_pages=SUMMARY_INTRO_MAX_PAGES):
     """Pages of the node covered by no child: from its start to just before the
     first child starts. Empty when the first child opens on the node's own page."""
@@ -830,20 +873,22 @@ async def summarize_tree(structure, pdf_pages, model=None,
     child summaries plus the pages no child covers. A parent's summary describes
     its whole subtree (end_index union semantics). Nodes that already carry a
     summary are left untouched; leaves under `small_node_tokens` use their raw
-    text as the summary without a model call."""
-    semaphore = asyncio.Semaphore(concurrency or SUMMARY_CONCURRENCY)
+    text as the summary without a model call. Queued model calls are admitted
+    deepest node first: depth counts the calls left on a node's path to the
+    root, its own included."""
+    gate = _PriorityGate(concurrency or SUMMARY_CONCURRENCY)
     asked = answered = False
 
-    async def ask(prompt):
+    async def ask(prompt, prio):
         nonlocal asked, answered
         asked = True
-        async with semaphore:
+        async with gate.slot(prio):
             reply = await llm_acompletion(model, prompt)
         if reply:
             answered = True
         return reply
 
-    async def leaf_summary(node):
+    async def leaf_summary(node, prio):
         text = get_text_of_pdf_pages(pdf_pages, node['start_index'], node['end_index'])
         if count_tokens(text, model="gpt-4o") < small_node_tokens:
             return text.strip()
@@ -874,14 +919,14 @@ async def leaf_summary(node):
 
     Follow strictly the above JSON return format. Do not include any other text!
     """
-        reply = await ask(prompt)
+        reply = await ask(prompt, prio)
         if retitle:
             written = parse_title(reply)
             if written:
                 node['title'] = written
         return parse_summary(reply)
 
-    async def parent_summary(node):
+    async def parent_summary(node, prio):
         children = node['nodes']
         intro = get_intro_text(node, pdf_pages, max_pages=max_intro_pages)
         listing = json.dumps(
@@ -905,12 +950,12 @@ async def parent_summary(node):
 
     Follow strictly the above JSON return format. Do not include any other text!
     """
-        return parse_summary(await ask(prompt))
+        return parse_summary(await ask(prompt, prio))
 
-    async def visit(node):
+    async def visit(node, depth=1):
         children = node.get('nodes') or []
         if children:
-            done = await asyncio.gather(*(visit(child) for child in children),
+            done = await asyncio.gather(*(visit(child, depth + 1) for child in children),
                                         return_exceptions=True)
             for result in done:
                 if isinstance(result, Exception) and _is_unrecoverable(result):
@@ -918,7 +963,8 @@ async def visit(node):
         if node.get('summary'):
             return
         try:
-            node['summary'] = await (parent_summary(node) if children else leaf_summary(node))
+            node['summary'] = await (parent_summary(node, depth) if children
+                                     else leaf_summary(node, depth))
         except Exception as e:
             node['summary'] = ""
             if _is_unrecoverable(e):
diff --git a/tests/test_client.py b/tests/test_client.py
index ec7ed845e..91898c2a0 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1187,6 +1187,47 @@ async def unexpected(model, prompt):
     assert out[0]["summary"] == "tiny"
 
 
+def test_summarize_tree_admits_deepest_leaf_first(monkeypatch):
+    """The tree is visited level by level, so with one permit the shallow
+    leaf D reaches the gate before the deep leaf C. C must be admitted first
+    anyway: the parents still owed above it are what the run ends on."""
+    started = []
+
+    async def fake(model, prompt):
+        started.append(next(w for w in ("alpha", "gamma", "delta", "Section Title")
+                            if w in prompt))
+        for _ in range(8):      # enough loop turns for every leaf to reach the gate
+            await asyncio.sleep(0)
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+    pdf_pages = [("alpha " * 5, 5), ("gamma " * 5, 5), ("delta " * 5, 5)]
+    structure = [{"title": "R", "start_index": 1, "end_index": 3, "nodes": [
+        {"title": "A", "start_index": 1, "end_index": 1},
+        {"title": "B", "start_index": 2, "end_index": 2, "nodes": [
+            {"title": "C", "start_index": 2, "end_index": 2}]},
+        {"title": "D", "start_index": 3, "end_index": 3}]}]
+    asyncio.run(pageindex.utils.summarize_tree(
+        structure, pdf_pages, small_node_tokens=0, concurrency=1))
+    assert started.index("gamma") < started.index("delta")
+    assert started.count("Section Title") == 2
+
+
+def test_priority_gate_passes_the_permit_on_when_its_taker_is_cancelled():
+    """A waiter cancelled after the permit was granted but before it ran
+    must hand the permit on, or the pool shrinks by one for good."""
+    async def scenario():
+        gate = pageindex.utils._PriorityGate(1)
+        await gate.acquire(1)
+        waiter = asyncio.create_task(gate.acquire(1))
+        await asyncio.sleep(0)          # queued
+        gate.release()                  # granted to `waiter`, not yet resumed
+        waiter.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await waiter
+        await asyncio.wait_for(gate.acquire(1), 1)   # hangs on a leaked permit
+    asyncio.run(scenario())
+
+
 def test_summarize_tree_partial_exhaustion_fails_loud(monkeypatch):
     """One lucky call must not vouch for a model that then went away: a
     ladder-exhausted node raises instead of silently blanking."""

From 8a53457aab5419f374fe7726983a97194d79052b Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Wed, 26 Aug 2026 20:29:30 +0800
Subject: [PATCH 02/18] perf: summaries start the deepest node first

The priority gate orders only the calls that wait; the first
SUMMARY_CONCURRENCY calls walk straight through in arrival order, and
arrival order was the recursive visit's level-by-level expansion, shallow
leaves first. So whenever the gate had free permits, the deep leaves whose
parents the run ends on still started last. Now every node's task is
created up front, deepest first: children always exist before their
parent, which awaits their tasks instead of spawning them, so the deep
leaves reach the gate first and a parent's own call no longer queues
behind leaves that merely arrived earlier.

Same 182 calls on fed, prompts unchanged. Summary stage at 64 wide,
mirror-ordered: gate only 51.1/50.0 s, this 46.8/40.0 s; the landed code
re-runs at 45.4 s against 57.4/58.7 s for the gate alone measured the
same afternoon. Error handling is untouched: parents still gather their
children with return_exceptions and re-raise the unrecoverable ones,
roots the same.

Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
---
 pageindex/utils.py   | 26 ++++++++++++++++++--------
 tests/test_client.py | 36 +++++++++++++++++++++++++++++++-----
 2 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/pageindex/utils.py b/pageindex/utils.py
index 37b870aa9..e61af984d 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -873,9 +873,9 @@ async def summarize_tree(structure, pdf_pages, model=None,
     child summaries plus the pages no child covers. A parent's summary describes
     its whole subtree (end_index union semantics). Nodes that already carry a
     summary are left untouched; leaves under `small_node_tokens` use their raw
-    text as the summary without a model call. Queued model calls are admitted
-    deepest node first: depth counts the calls left on a node's path to the
-    root, its own included."""
+    text as the summary without a model call. Model calls run deepest node
+    first, both in starting order and in leaving the queue: depth counts the
+    calls left on a node's path to the root, its own included."""
     gate = _PriorityGate(concurrency or SUMMARY_CONCURRENCY)
     asked = answered = False
 
@@ -952,11 +952,10 @@ async def parent_summary(node, prio):
     """
         return parse_summary(await ask(prompt, prio))
 
-    async def visit(node, depth=1):
+    async def visit(node, depth, child_tasks):
         children = node.get('nodes') or []
         if children:
-            done = await asyncio.gather(*(visit(child, depth + 1) for child in children),
-                                        return_exceptions=True)
+            done = await asyncio.gather(*child_tasks, return_exceptions=True)
             for result in done:
                 if isinstance(result, Exception) and _is_unrecoverable(result):
                     raise result
@@ -970,8 +969,19 @@ async def visit(node, depth=1):
             if _is_unrecoverable(e):
                 raise
 
-    results = await asyncio.gather(*(visit(root) for root in structure),
-                                    return_exceptions=True)
+    order = []
+
+    def collect(nodes, depth):
+        for node in nodes:
+            collect(node.get('nodes') or [], depth + 1)
+            order.append((depth, node))
+    collect(structure, 1)
+    tasks = {}
+    for depth, node in sorted(order, key=lambda pair: -pair[0]):
+        tasks[id(node)] = asyncio.create_task(visit(
+            node, depth, [tasks[id(child)] for child in node.get('nodes') or []]))
+    results = await asyncio.gather(*(tasks[id(root)] for root in structure),
+                                   return_exceptions=True)
     for r in results:
         if isinstance(r, Exception) and _is_unrecoverable(r):
             raise r
diff --git a/tests/test_client.py b/tests/test_client.py
index 91898c2a0..5a3d6eb21 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1187,10 +1187,10 @@ async def unexpected(model, prompt):
     assert out[0]["summary"] == "tiny"
 
 
-def test_summarize_tree_admits_deepest_leaf_first(monkeypatch):
-    """The tree is visited level by level, so with one permit the shallow
-    leaf D reaches the gate before the deep leaf C. C must be admitted first
-    anyway: the parents still owed above it are what the run ends on."""
+def test_summarize_tree_starts_the_deepest_leaf_first(monkeypatch):
+    """Visited level by level, the shallow leaves A and D would take both
+    permits before the deep leaf C even reached the gate. C must start first:
+    the parents still owed above it are what the run ends on."""
     started = []
 
     async def fake(model, prompt):
@@ -1207,11 +1207,37 @@ async def fake(model, prompt):
             {"title": "C", "start_index": 2, "end_index": 2}]},
         {"title": "D", "start_index": 3, "end_index": 3}]}]
     asyncio.run(pageindex.utils.summarize_tree(
-        structure, pdf_pages, small_node_tokens=0, concurrency=1))
+        structure, pdf_pages, small_node_tokens=0, concurrency=2))
     assert started.index("gamma") < started.index("delta")
     assert started.count("Section Title") == 2
 
 
+def test_summarize_tree_admits_a_ready_parent_before_a_queued_shallow_leaf(monkeypatch):
+    """With one permit the leaf D is already queued when C's child finishes
+    and C becomes ready. C has two calls left above it, D one, so C goes
+    first even though D asked earlier."""
+    started = []
+
+    async def fake(model, prompt):
+        started.append(next(w for w in ("Section Title: C", "Section Title",
+                                        "alpha", "epsilon", "delta")
+                            if w in prompt))
+        for _ in range(8):
+            await asyncio.sleep(0)
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+    pdf_pages = [("alpha " * 5, 5), ("epsilon " * 5, 5), ("delta " * 5, 5)]
+    structure = [{"title": "R", "start_index": 1, "end_index": 3, "nodes": [
+        {"title": "B", "start_index": 2, "end_index": 2, "nodes": [
+            {"title": "C", "start_index": 2, "end_index": 2, "nodes": [
+                {"title": "E", "start_index": 2, "end_index": 2}]}]},
+        {"title": "A", "start_index": 1, "end_index": 1},
+        {"title": "D", "start_index": 3, "end_index": 3}]}]
+    asyncio.run(pageindex.utils.summarize_tree(
+        structure, pdf_pages, small_node_tokens=0, concurrency=1))
+    assert started.index("Section Title: C") < started.index("delta")
+
+
 def test_priority_gate_passes_the_permit_on_when_its_taker_is_cancelled():
     """A waiter cancelled after the permit was granted but before it ran
     must hand the permit on, or the pool shrinks by one for good."""

From 79339fc5c651d74c318ced7b35625f4b8573bf93 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Wed, 26 Aug 2026 21:49:30 +0800
Subject: [PATCH 03/18] refactor: summarize_tree becomes a scheduler that takes
 nodes as they settle

summarize_tree built every task up front from a finished tree, so summaries
could only begin once expand had returned it. SummaryScheduler keeps the same
machinery (one task per node, parents awaiting their children, the priority
gate, deepest-first task creation) behind two calls: mark_final(nodes) says
those nodes will not gain, lose or swap children and spawns tasks for their
subtrees; finish() awaits the roots, applies the every-call-failed check and
strips the bookkeeping keys. A node's task first awaits its own mark, so a
subtree can be handed over while the rest of the tree is still being
decided, and a node is never summarized twice however often it is marked:
tasks are memoized per node and a parent composes whatever its children
already wrote.

summarize_tree marks the whole tree and finishes, which reproduces the old
behavior exactly: same task order, same gate, same error semantics. The fed
summary stage re-runs at 44.2 s against 45.4 s before.

Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
---
 pageindex/utils.py   | 177 ++++++++++++++++++++++++++++---------------
 tests/test_client.py |  29 +++++++
 2 files changed, 147 insertions(+), 59 deletions(-)

diff --git a/pageindex/utils.py b/pageindex/utils.py
index e61af984d..f2d375945 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -866,31 +866,79 @@ def strip_internal_keys(structure):
     return structure
 
 
-async def summarize_tree(structure, pdf_pages, model=None,
-                         small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
-                         max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None):
-    """Bottom-up summaries: leaves from their own pages, parents composed from
-    child summaries plus the pages no child covers. A parent's summary describes
-    its whole subtree (end_index union semantics). Nodes that already carry a
-    summary are left untouched; leaves under `small_node_tokens` use their raw
-    text as the summary without a model call. Model calls run deepest node
-    first, both in starting order and in leaving the queue: depth counts the
-    calls left on a node's path to the root, its own included."""
-    gate = _PriorityGate(concurrency or SUMMARY_CONCURRENCY)
-    asked = answered = False
-
-    async def ask(prompt, prio):
-        nonlocal asked, answered
-        asked = True
-        async with gate.slot(prio):
-            reply = await llm_acompletion(model, prompt)
+def _subtree(nodes):
+    for node in nodes:
+        yield node
+        yield from _subtree(node.get('nodes') or [])
+
+
+class SummaryScheduler:
+    """Bottom-up summaries, taking nodes as they are marked final.
+
+    A node's task waits for its mark (its children will not change any more),
+    then for its children's tasks, then makes its own call: leaves from their
+    own pages, parents composed from child summaries plus the pages no child
+    covers. Marked subtrees get their tasks deepest node first and queued
+    calls leave the gate deepest first: depth counts the calls left on a
+    node's path to the root, its own included."""
+
+    def __init__(self, structure, pdf_pages, model=None,
+                 small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
+                 max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None):
+        self.structure = structure
+        self._pdf_pages = pdf_pages
+        self._model = model
+        self._small_node_tokens = small_node_tokens
+        self._max_intro_pages = max_intro_pages
+        self._gate = _PriorityGate(concurrency or SUMMARY_CONCURRENCY)
+        self._asked = self._answered = False
+        self._marks = {}     # id(node) -> future resolved once the node is final
+        self._tasks = {}     # id(node) -> its summary task
+
+    def mark_final(self, nodes):
+        """These nodes will not gain, lose or swap children: their summaries
+        may start. Their subtrees get tasks, deepest node first."""
+        nodes = list(nodes)
+        for node in nodes:
+            mark = self._mark(node)
+            if not mark.done():
+                mark.set_result(None)
+        marked = {id(node) for node in nodes}
+        order = []
+
+        def walk(nodes, depth, inside):
+            for node in nodes:
+                inside_here = inside or id(node) in marked
+                if inside_here:
+                    order.append((depth, node))
+                walk(node.get('nodes') or [], depth + 1, inside_here)
+        walk(self.structure, 1, False)
+        for depth, node in sorted(order, key=lambda pair: -pair[0]):
+            self._task(node, depth)
+
+    def _mark(self, node):
+        mark = self._marks.get(id(node))
+        if mark is None:
+            mark = self._marks[id(node)] = asyncio.get_running_loop().create_future()
+        return mark
+
+    def _task(self, node, depth):
+        task = self._tasks.get(id(node))
+        if task is None:
+            task = self._tasks[id(node)] = asyncio.create_task(self._visit(node, depth))
+        return task
+
+    async def _ask(self, prompt, prio):
+        self._asked = True
+        async with self._gate.slot(prio):
+            reply = await llm_acompletion(self._model, prompt)
         if reply:
-            answered = True
+            self._answered = True
         return reply
 
-    async def leaf_summary(node, prio):
-        text = get_text_of_pdf_pages(pdf_pages, node['start_index'], node['end_index'])
-        if count_tokens(text, model="gpt-4o") < small_node_tokens:
+    async def _leaf_summary(self, node, prio):
+        text = get_text_of_pdf_pages(self._pdf_pages, node['start_index'], node['end_index'])
+        if count_tokens(text, model="gpt-4o") < self._small_node_tokens:
             return text.strip()
 
         # A node merged from same-page siblings carries a title joined from theirs.
@@ -919,16 +967,16 @@ async def leaf_summary(node, prio):
 
     Follow strictly the above JSON return format. Do not include any other text!
     """
-        reply = await ask(prompt, prio)
+        reply = await self._ask(prompt, prio)
         if retitle:
             written = parse_title(reply)
             if written:
                 node['title'] = written
         return parse_summary(reply)
 
-    async def parent_summary(node, prio):
+    async def _parent_summary(self, node, prio):
         children = node['nodes']
-        intro = get_intro_text(node, pdf_pages, max_pages=max_intro_pages)
+        intro = get_intro_text(node, self._pdf_pages, max_pages=self._max_intro_pages)
         listing = json.dumps(
             [{'title': c.get('title', ''), 'summary': c.get('summary', '')} for c in children],
             ensure_ascii=False)
@@ -950,56 +998,67 @@ async def parent_summary(node, prio):
 
     Follow strictly the above JSON return format. Do not include any other text!
     """
-        return parse_summary(await ask(prompt, prio))
+        return parse_summary(await self._ask(prompt, prio))
 
-    async def visit(node, depth, child_tasks):
+    async def _visit(self, node, depth):
+        await self._mark(node)
         children = node.get('nodes') or []
         if children:
-            done = await asyncio.gather(*child_tasks, return_exceptions=True)
+            done = await asyncio.gather(*(self._task(child, depth + 1) for child in children),
+                                        return_exceptions=True)
             for result in done:
                 if isinstance(result, Exception) and _is_unrecoverable(result):
                     raise result
         if node.get('summary'):
             return
         try:
-            node['summary'] = await (parent_summary(node, depth) if children
-                                     else leaf_summary(node, depth))
+            node['summary'] = await (self._parent_summary(node, depth) if children
+                                     else self._leaf_summary(node, depth))
         except Exception as e:
             node['summary'] = ""
             if _is_unrecoverable(e):
                 raise
 
-    order = []
+    async def finish(self):
+        """Wait for every summary; fails loud if the model never answered."""
+        results = await asyncio.gather(*(self._task(root, 1) for root in self.structure),
+                                       return_exceptions=True)
+        for r in results:
+            if isinstance(r, Exception) and _is_unrecoverable(r):
+                raise r
+
+        # Raw-text leaves summarize without the model, so they cannot vouch for
+        # it: a run whose every model call failed still fails loud.
+        def _any_summary(nodes):
+            return any(n.get('summary') or _any_summary(n.get('nodes') or [])
+                       for n in nodes)
+        if (self._asked and not self._answered) or not _any_summary(self.structure):
+            raise RuntimeError(
+                "Summary generation failed for all nodes "
+                "(every summary call failed or returned empty; "
+                "check the model and its context limits)"
+            )
+
+        strip_internal_keys(self.structure)
+        return self.structure
 
-    def collect(nodes, depth):
-        for node in nodes:
-            collect(node.get('nodes') or [], depth + 1)
-            order.append((depth, node))
-    collect(structure, 1)
-    tasks = {}
-    for depth, node in sorted(order, key=lambda pair: -pair[0]):
-        tasks[id(node)] = asyncio.create_task(visit(
-            node, depth, [tasks[id(child)] for child in node.get('nodes') or []]))
-    results = await asyncio.gather(*(tasks[id(root)] for root in structure),
-                                   return_exceptions=True)
-    for r in results:
-        if isinstance(r, Exception) and _is_unrecoverable(r):
-            raise r
-
-    # Raw-text leaves summarize without the model, so they cannot vouch for
-    # it: a run whose every model call failed still fails loud.
-    def _any_summary(nodes):
-        return any(n.get('summary') or _any_summary(n.get('nodes') or [])
-                   for n in nodes)
-    if (asked and not answered) or not _any_summary(structure):
-        raise RuntimeError(
-            "Summary generation failed for all nodes "
-            "(every summary call failed or returned empty; "
-            "check the model and its context limits)"
-        )
 
-    strip_internal_keys(structure)
-    return structure
+async def summarize_tree(structure, pdf_pages, model=None,
+                         small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
+                         max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None):
+    """Bottom-up summaries: leaves from their own pages, parents composed from
+    child summaries plus the pages no child covers. A parent's summary describes
+    its whole subtree (end_index union semantics). Nodes that already carry a
+    summary are left untouched; leaves under `small_node_tokens` use their raw
+    text as the summary without a model call. Model calls run deepest node
+    first, both in starting order and in leaving the queue: depth counts the
+    calls left on a node's path to the root, its own included."""
+    scheduler = SummaryScheduler(structure, pdf_pages, model=model,
+                                 small_node_tokens=small_node_tokens,
+                                 max_intro_pages=max_intro_pages,
+                                 concurrency=concurrency)
+    scheduler.mark_final(list(_subtree(structure)))
+    return await scheduler.finish()
 
 
 def create_clean_structure_for_description(structure):
diff --git a/tests/test_client.py b/tests/test_client.py
index 5a3d6eb21..93efdcc8f 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1238,6 +1238,35 @@ async def fake(model, prompt):
     assert started.index("Section Title: C") < started.index("delta")
 
 
+def test_summary_scheduler_starts_a_marked_node_without_waiting_for_the_rest(monkeypatch):
+    """A node marked final summarizes right away while its siblings are still
+    unmarked; marking it again, or marking its parent later, never repeats
+    its call: the parent composes the summary already written."""
+    calls = []
+
+    async def fake(model, prompt):
+        calls.append(prompt)
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+    pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+    leaf = {"title": "A", "start_index": 1, "end_index": 1}
+    root = {"title": "R", "start_index": 1, "end_index": 2, "nodes": [
+        leaf, {"title": "B", "start_index": 2, "end_index": 2}]}
+
+    async def scenario():
+        scheduler = pageindex.utils.SummaryScheduler([root], pdf_pages, small_node_tokens=0)
+        scheduler.mark_final([leaf])
+        scheduler.mark_final([leaf])
+        for _ in range(3):
+            await asyncio.sleep(0)
+        assert len(calls) == 1 and "alpha" in calls[0]
+        scheduler.mark_final(list(pageindex.utils._subtree([root])))
+        return await scheduler.finish()
+    asyncio.run(scenario())
+    assert len(calls) == 3
+    assert [n["summary"] for n in (leaf, root["nodes"][1], root)] == ["ok"] * 3
+
+
 def test_priority_gate_passes_the_permit_on_when_its_taker_is_cancelled():
     """A waiter cancelled after the permit was granted but before it ran
     must hand the permit on, or the pool shrinks by one for good."""

From 870a7cb8a1d3c4b669876545f71195ae30d63bf5 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Wed, 26 Aug 2026 22:07:32 +0800
Subject: [PATCH 04/18] perf: summaries start the moment expand can no longer
 touch a node

page_index_flash ran optimize and the summaries as two asyncio.run calls, one
after the other: the whole summary wave waited for the expand chain, a handful
of dependent calls that take 18-67 s while the summary channels sat idle. Now
the two share one loop. optimize(on_final=...) reports which nodes will not
change any more, and the SummaryScheduler starts those at once; a parent's
task is already waiting on its children, so it follows the moment its last
child is done.

A node is final when its children cannot change: a collapsed node under the
trigger, a collapsed node expand has judged (frozen), or any node with
children. That rule holds throughout the run, not only at the start. Expand
touches only collapsed nodes over the trigger. The cost merge never fires on a
surviving node after the first round: keeping X expanded means expand_cost <
S(X), and tree_cost(X) right after attaching the children equals expand_cost
term for term, only falling as children expand further; ancestors' tree_cost
can only fall with it while their S stays put, because assign_ends pins the
last child at the parent's old end. Same-page fusion was the one later
mutation, and it is now done where the duplicates arise, right after a merge
collapses a subtree onto a sibling's exact pages and right after expand
attaches children, instead of at the next round's start, so no node waits a
round for it. Reports go out after each round's merges, at each candidate's
decision (kept collapsed, nothing found, or expanded together with what it
grew), and for the whole tree at the end.

Merge-only trees for the nine corpus PDFs are identical before and after; the
only visible change is SpaceX ending after two rounds instead of a third that
did nothing. Same-hour A/B, end to end: fed 97.9 s -> 72.6 s, PRML 174.3 s ->
136.8 s. Peak in flight is now the expand cap plus the summary cap.

Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
---
 pageindex/flash/api.py     |  66 ++++++++++------
 pageindex/tree_optimize.py |  52 ++++++++++---
 tests/test_client.py       | 151 +++++++++++++++++++++++++++++++++++++
 3 files changed, 234 insertions(+), 35 deletions(-)

diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index 6f7361e27..2a133488a 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -73,21 +73,20 @@ async def _summarize(structure, page_list, model, concurrency=None):
     await summarize_tree(structure, page_list, model=model, concurrency=concurrency)
 
 
-def _optimize(structure, page_texts, do_expand, model):
+async def _optimize_async(structure, page_texts, do_expand, model, on_final=None):
     """Merge/expand refinement between extraction and summaries.
 
     Beyond the merge the default path runs anyway, this adds LLM expand and
-    reports before/after search-cost metrics. Summaries run after, so they
-    describe the final tree. Expand reads the same page text the summaries use.
+    reports before/after search-cost metrics. Expand reads the same page text
+    the summaries use.
     """
-    import asyncio
     from ..tree_optimize import optimize
     lines = [[line_text.strip() for line_text in (page_text or "").splitlines()
               if line_text.strip()]
              for page_text in page_texts]
-    outcome = asyncio.run(optimize(structure, page_texts, lines, model=model,
-                                   do_expand=do_expand,
-                                   page_count=len(page_texts)))
+    outcome = await optimize(structure, page_texts, lines, model=model,
+                             do_expand=do_expand, page_count=len(page_texts),
+                             on_final=on_final)
     return {"merges": outcome["merges"], "expands": outcome["expands"],
             "same_page_merges": outcome["same_page_merges"],
             "same_page_dropped": outcome["same_page_dropped"],
@@ -95,6 +94,24 @@ def _optimize(structure, page_texts, do_expand, model):
             "before": outcome["before"], "after": outcome["after"]}
 
 
+def _optimize(structure, page_texts, do_expand, model):
+    import asyncio
+    return asyncio.run(_optimize_async(structure, page_texts, do_expand, model))
+
+
+async def _optimize_and_summarize(structure, page_texts, model, summary_model,
+                                  concurrency):
+    """Expand and summarize on one loop: a node is summarized as soon as
+    expand can no longer change it, a parent once its children are done."""
+    from ..utils import SummaryScheduler
+    scheduler = SummaryScheduler(structure, [(text, 0) for text in page_texts],
+                                 model=summary_model, concurrency=concurrency)
+    report = await _optimize_async(structure, page_texts, True, model,
+                                   on_final=scheduler.mark_final)
+    await scheduler.finish()
+    return report
+
+
 def page_index_flash(pdf, summary=True, summary_model=None,
                      optimize: str | bool | None = None, optimize_expand=None,
                      optimize_model=None, summary_concurrency=None,
@@ -118,28 +135,31 @@ def page_index_flash(pdf, summary=True, summary_model=None,
             f"optimize must be 'full', 'merge', or False, got {optimize!r}")
     result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc)
     structure = result.get("structure", [])
+    if summary and structure and summary_model is None:
+        from ..utils import ConfigLoader
+        cfg = ConfigLoader().load()
+        summary_model = getattr(cfg, 'summary_model', None) or cfg.model
+    # bookmark-only extractions carry no page_texts and scanned ones
+    # only empty strings; expand needs text
+    pages = result.pop("page_texts", None) or []
+    do_expand = optimize == "full" and any(pages)
+    if optimize and structure and summary and do_expand:
+        import asyncio
+        result["optimize"] = asyncio.run(_optimize_and_summarize(
+            structure, pages, optimize_model or summary_model,
+            summary_model, summary_concurrency))
+        return result
     if optimize and structure:
-        # bookmark-only extractions carry no page_texts and scanned ones
-        # only empty strings; expand needs text
-        pages = result.get("page_texts") or []
-        result["optimize"] = _optimize(structure, pages,
-                                       optimize == "full" and any(pages),
+        result["optimize"] = _optimize(structure, pages, do_expand,
                                        optimize_model or summary_model)
     if summary and structure:
         import asyncio
-        from ..utils import ConfigLoader
-        if summary_model is None:
-            cfg = ConfigLoader().load()
-            summary_model = getattr(cfg, 'summary_model', None) or cfg.model
-        page_texts = result.pop("page_texts", [])
-        page_list = [(text, 0) for text in page_texts]
+        page_list = [(text, 0) for text in pages]
         asyncio.run(_summarize(structure, page_list, summary_model,
                                concurrency=summary_concurrency))
-    else:
-        result.pop("page_texts", None)
-        if structure:
-            from ..utils import strip_internal_keys
-            strip_internal_keys(structure)   # summarize_tree does this on its way out
+    elif structure:
+        from ..utils import strip_internal_keys
+        strip_internal_keys(structure)   # summarize_tree does this on its way out
     return result
 
 
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index f1638ff4e..f87e84b4d 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -705,6 +705,7 @@ async def process(node):
             log.append({"op": "expand", "node_id": node.get("node_id"),
                         "decision": "no_children", "S": span, "attempts": attempts})
             frozen.add(node.get("node_id"))
+            args.settled([node])
             return
 
         scored = []
@@ -739,15 +740,19 @@ async def process(node):
                                  for c in best["children"]]})
 
         frozen.add(node.get("node_id"))
-        if keep:
-            changed = True
-            attach_children(node, best["children"], lines)
-            results = await asyncio.gather(*(process(child)
-                                             for child in node["nodes"]),
-                                           return_exceptions=True)
-            for result in results:
-                if isinstance(result, BaseException):
-                    raise result
+        if not keep:
+            args.settled([node])
+            return
+        changed = True
+        attach_children(node, best["children"], lines)
+        merge_same_page([node], log)
+        args.settled([node])
+        results = await asyncio.gather(*(process(child)
+                                         for child in node["nodes"]),
+                                       return_exceptions=True)
+        for result in results:
+            if isinstance(result, BaseException):
+                raise result
 
     results = await asyncio.gather(*(process(node)
                                      for node, _ in flatten(structure)),
@@ -768,11 +773,19 @@ def default_model():
     return getattr(opt, "summary_model", None) or opt.model
 
 
+def final_nodes(nodes, trigger, frozen):
+    """Nodes whose children will not change any more: everything but a
+    collapsed node over the trigger that expand has not judged yet."""
+    return [node for node, _ in flatten(nodes)
+            if not is_frontier(node) or S(node) <= trigger
+            or node.get("node_id") in frozen]
+
+
 async def optimize(structure, pages, lines, model=None, routing=ROUTING_COST,
                    trigger_pages=TRIGGER_PAGES, min_gain_ratio=0.0,
                    do_merge=True, do_expand=True, max_rounds=3, page_count=None,
                    cache=None, kinds=("section", "table"), empty_retries=1,
-                   do_relabel=True, progress=False):
+                   do_relabel=True, progress=False, on_final=None):
     """Run merge and expand over a tree until neither changes anything.
 
     Mutates `structure` in place and returns a summary.
@@ -782,30 +795,45 @@ async def optimize(structure, pages, lines, model=None, routing=ROUTING_COST,
     ancestors' tree_cost. Nodes decided by either operator are frozen for the rest
     of the run, so a node cannot be collapsed and re-expanded in alternating
     rounds.
+
+    `on_final(nodes)` hears, as the run goes, which nodes will not change any
+    more (see final_nodes): after each round's merges, as expand decides each
+    candidate, and for the whole tree at the end.
     """
     if do_expand and pages is None:
         raise ValueError("expand needs the PDF pages; pass pages/lines or do_expand=False")
+    log, frozen = [], set()
+
+    def settled(nodes):
+        if on_final is not None:
+            on_final(final_nodes(nodes, trigger_pages, frozen))
     opts = SimpleNamespace(model=model or default_model(), routing=routing,
                            trigger_pages=trigger_pages,
                            min_gain_ratio=min_gain_ratio, cache=cache,
                            kinds=set(kinds) if kinds else None,
-                           empty_retries=empty_retries, progress=progress)
+                           empty_retries=empty_retries, progress=progress,
+                           settled=settled)
     baseline = set(validate(structure, page_count)) if page_count else set()
     before = complexity(structure, page_count, routing=routing) if page_count else {}
 
-    log, frozen = [], set()
     rounds = 0
     for round_no in range(1, max_rounds + 1):
         rounds = round_no
         note(progress, f"  round {round_no}")
         same_page = merge_same_page(structure, log) if do_merge else False
         merged = merge(structure, routing, log, frozen, progress) if do_merge else False
+        if merged:
+            # a collapsed subtree can land on a sibling's exact pages
+            same_page = merge_same_page(structure, log) or same_page
+        settled(structure)
         expanded = await expand(structure, pages, lines, opts, log, frozen) \
             if do_expand else False
         log.append({"op": "round", "round": round_no, "same_page": same_page,
                     "merged": merged, "expanded": expanded})
         if not (same_page or merged or expanded):
             break
+    if on_final is not None:
+        on_final([node for node, _ in flatten(structure)])
 
     id_map = relabel(structure) if do_relabel else {}
     after = complexity(structure, page_count, routing=routing) if page_count else {}
diff --git a/tests/test_client.py b/tests/test_client.py
index 93efdcc8f..e9db0746a 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1962,6 +1962,157 @@ async def slow_empty(model, prompt):
     assert inflight["peak"] <= 32
 
 
+def _expand_fixture():
+    """Twelve pages; X spans nine of them, so expand looks at it and, when
+    the model offers 'Sub One' (page 4) and 'Sub Two' (page 8), splits it
+    into two leaves under the trigger. Bodies are long enough that every
+    leaf summary needs the model."""
+    body = "body " * 250
+    pages = [body] * 12
+    pages[3] = "Sub One\n" + body
+    pages[7] = "Sub Two\n" + body
+    lines = [[l for l in p.splitlines() if l.strip()] for p in pages]
+    tree = [{"title": "R", "start_index": 1, "end_index": 12, "node_id": "0000", "nodes": [
+        {"title": "A", "start_index": 1, "end_index": 3, "node_id": "0001"},
+        {"title": "X", "start_index": 4, "end_index": 12, "node_id": "0002"}]}]
+    return tree, pages, lines
+
+
+def test_final_nodes_holds_back_only_undecided_expand_candidates():
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, _, _ = _expand_fixture()
+    titles = lambda nodes: sorted(n["title"] for n in nodes)
+    assert titles(tree_optimize.final_nodes(tree, 5, set())) == ["A", "R"]
+    assert titles(tree_optimize.final_nodes(tree, 5, {"0002"})) == ["A", "R", "X"]
+    assert titles(tree_optimize.final_nodes(tree, 9, set())) == ["A", "R", "X"]
+
+
+def test_optimize_reports_final_nodes_as_expand_decides(monkeypatch):
+    """Before the first expand call, everything expand cannot touch is
+    reported final; the candidate and what it grows are reported the moment
+    it is decided; the closing report covers the whole tree."""
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, pages, lines = _expand_fixture()
+    reports, replied_after = [], []
+
+    async def propose(model, prompt):
+        await asyncio.sleep(0.01)
+        replied_after.append(len(reports))
+        return {"subsections": [{"title": "Sub One", "page": 4},
+                                {"title": "Sub Two", "page": 8}]}
+    monkeypatch.setattr(tree_optimize, "ask_model", propose)
+    asyncio.run(tree_optimize.optimize(
+        tree, pages, lines, model="m", do_expand=True,
+        on_final=lambda nodes: reports.append(sorted(n["title"] for n in nodes))))
+    before_reply = [t for report in reports[:replied_after[0]] for t in report]
+    assert "R" in before_reply and "A" in before_reply and "X" not in before_reply
+    decided = next(r for r in reports[replied_after[0]:] if "X" in r)
+    assert decided == ["Sub One", "Sub Two", "X"]
+    assert reports[-1] == ["A", "R", "Sub One", "Sub Two", "X"]
+
+
+def test_optimize_reports_a_candidate_kept_collapsed_once_decided(monkeypatch):
+    """A candidate the model finds nothing in is final the moment its retry
+    ladder ends, not at the end of the run."""
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, pages, lines = _expand_fixture()
+    reports = []
+
+    async def nothing(model, prompt):
+        return {"subsections": []}
+    monkeypatch.setattr(tree_optimize, "ask_model", nothing)
+    asyncio.run(tree_optimize.optimize(
+        tree, pages, lines, model="m", do_expand=True,
+        on_final=lambda nodes: reports.append(sorted(n["title"] for n in nodes))))
+    assert ["X"] in reports[:-1]
+
+
+def test_merge_fuses_the_same_page_frontier_it_creates_at_once():
+    """Collapsing F leaves it on exactly G's page; the two must fuse right
+    then, not one round later (with a single round they never would)."""
+    import pageindex.tree_optimize as tree_optimize
+
+    tree = [{"title": "R", "start_index": 1, "end_index": 6, "node_id": "0000", "nodes": [
+        {"title": "F", "start_index": 1, "end_index": 1, "node_id": "0001", "nodes": [
+            {"title": "f1", "start_index": 1, "end_index": 1, "node_id": "0002"}]},
+        {"title": "G", "start_index": 1, "end_index": 1, "node_id": "0003"},
+        {"title": "H", "start_index": 2, "end_index": 5, "node_id": "0004"}]}]
+    outcome = asyncio.run(tree_optimize.optimize(
+        tree, None, None, do_expand=False, max_rounds=1))
+    assert outcome["merges"] == 1 and outcome["same_page_merges"] == 1
+    kept = tree[0]["nodes"]
+    assert [n["title"] for n in kept][1:] == ["H"] and kept[0].get("_same_page")
+
+
+def test_expand_fuses_same_page_children_before_reporting_them(monkeypatch):
+    """Two proposed headings on one page become one node as soon as they
+    are attached, so the report never carries a duplicate."""
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, pages, lines = _expand_fixture()
+    pages[3] = "Sub One\nSub Two\n" + pages[3]
+    pages[4] = "Sub Three\n" + pages[4]
+    pages[8] = "Sub Four\n" + pages[8]
+    lines = [[l for l in p.splitlines() if l.strip()] for p in pages]
+
+    async def propose(model, prompt):
+        return {"subsections": [{"title": "Sub One", "page": 4}, {"title": "Sub Two", "page": 4},
+                                {"title": "Sub Three", "page": 5}, {"title": "Sub Four", "page": 9}]}
+    monkeypatch.setattr(tree_optimize, "ask_model", propose)
+    reports = []
+    outcome = asyncio.run(tree_optimize.optimize(
+        tree, pages, lines, model="m", do_expand=True, max_rounds=1,
+        on_final=lambda nodes: reports.append([(n["title"], len(n.get("nodes") or []))
+                                              for n in nodes])))
+    assert outcome["expands"] == 1 and outcome["same_page_merges"] == 1
+    assert [n["title"] for n in tree[0]["nodes"][1]["nodes"]][1:] == ["Sub Three", "Sub Four"]
+    assert all(children != 4 for report in reports for _, children in report)
+
+
+def test_flash_summaries_start_while_expand_is_still_deciding(tmp_path, monkeypatch):
+    """With expand on, summaries run on the same loop: a leaf expand cannot
+    touch is already being summarized when the model answers about X. Every
+    node is summarized exactly once and the tree matches the merge+expand
+    pass run on its own."""
+    from conftest import build_pdf
+    import copy
+    import pageindex.flash.api as flash_api
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, pages, _ = _expand_fixture()
+    pdf = tmp_path / "doc.pdf"
+    pdf.write_bytes(build_pdf(["x"]))
+    monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
+        "structure": copy.deepcopy(tree), "page_texts": list(pages)})
+    started, replied_after = [], []
+
+    async def propose(model, prompt):
+        await asyncio.sleep(0.05)
+        replied_after.append(len(started))
+        return {"subsections": [{"title": "Sub One", "page": 4},
+                                {"title": "Sub Two", "page": 8}]}
+    monkeypatch.setattr(tree_optimize, "ask_model", propose)
+
+    async def summarize(model, prompt):
+        started.append(prompt)
+        await asyncio.sleep(0.01)
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", summarize)
+
+    result = flash_api.page_index_flash(str(pdf), summary_model="m")
+    assert replied_after[0] >= 1
+    assert result["optimize"]["expands"] == 1
+    nodes = list(pageindex.utils._subtree(result["structure"]))
+    assert [n["summary"] for n in nodes] == ["ok"] * 5 and len(started) == 5
+    shape = lambda nodes: [(n["title"], n["start_index"], n["end_index"],
+                            shape(n.get("nodes") or [])) for n in nodes]
+    alone = flash_api.page_index_flash(str(pdf), summary=False)
+    assert shape(result["structure"]) == shape(alone["structure"])
+
+
 def test_mode_declaration_top_level(monkeypatch):
     """mode= states where documents live; always optional, always
     checked, and mode="cloud" alone reads the env key."""

From bb6ee89e3828efef3db477b457ab9c8048b810c1 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 02:17:02 +0800
Subject: [PATCH 05/18] fix: review follow-ups - do_merge gate, tokenizer
 prefilter, finish() contract

Four fixes from the PR #432 review round:

- expand's in-flight same-page fusion now obeys do_merge: a caller who
  disabled merging (--no-merge, optimize_tree(do_merge=False)) keeps every
  proposed child and every document title.
- _leaf_summary consults the tokenizer only when the length leaves "small"
  possible: no text averages 8+ chars per token, so long leaves skip the
  synchronous count that ran back-to-back on the loop before either lane's
  first model call (first-call latency measured 2.6 s -> ~30 ms on a
  760-page synthetic, total 7.5 s -> 4.9 s).
- finish() verifies the promise mark_final rests on (a final node stays in
  the tree and keeps its children) and refuses to wait on tasked nodes that
  were never marked: a broken promise now fails loud with a diagnosis
  instead of a silently wrong tree or an indefinite, symptomless hang.
- the deepest-first test now pins the deep leaf against the shallow leaf
  document order would start first; deleting the sort in mark_final turns
  it red (mutation-verified).

Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
---
 pageindex/tree_optimize.py |  5 ++-
 pageindex/utils.py         | 21 +++++++++-
 tests/test_client.py       | 85 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index f87e84b4d..c8484932e 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -745,7 +745,8 @@ async def process(node):
             return
         changed = True
         attach_children(node, best["children"], lines)
-        merge_same_page([node], log)
+        if args.do_merge:
+            merge_same_page([node], log)
         args.settled([node])
         results = await asyncio.gather(*(process(child)
                                          for child in node["nodes"]),
@@ -812,7 +813,7 @@ def settled(nodes):
                            min_gain_ratio=min_gain_ratio, cache=cache,
                            kinds=set(kinds) if kinds else None,
                            empty_retries=empty_retries, progress=progress,
-                           settled=settled)
+                           settled=settled, do_merge=do_merge)
     baseline = set(validate(structure, page_count)) if page_count else set()
     before = complexity(structure, page_count, routing=routing) if page_count else {}
 
diff --git a/pageindex/utils.py b/pageindex/utils.py
index f2d375945..1ba5d5fde 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -894,6 +894,7 @@ def __init__(self, structure, pdf_pages, model=None,
         self._asked = self._answered = False
         self._marks = {}     # id(node) -> future resolved once the node is final
         self._tasks = {}     # id(node) -> its summary task
+        self._finals = []    # (node, ids of its children when it was marked)
 
     def mark_final(self, nodes):
         """These nodes will not gain, lose or swap children: their summaries
@@ -903,6 +904,7 @@ def mark_final(self, nodes):
             mark = self._mark(node)
             if not mark.done():
                 mark.set_result(None)
+                self._finals.append((node, tuple(id(c) for c in node.get('nodes') or [])))
         marked = {id(node) for node in nodes}
         order = []
 
@@ -938,7 +940,10 @@ async def _ask(self, prompt, prio):
 
     async def _leaf_summary(self, node, prio):
         text = get_text_of_pdf_pages(self._pdf_pages, node['start_index'], node['end_index'])
-        if count_tokens(text, model="gpt-4o") < self._small_node_tokens:
+        # no text averages 8+ chars per token, so length alone clears most
+        # leaves without paying for the tokenizer on the loop
+        if (len(text) <= self._small_node_tokens * 8
+                and count_tokens(text, model="gpt-4o") < self._small_node_tokens):
             return text.strip()
 
         # A node merged from same-page siblings carries a title joined from theirs.
@@ -1021,6 +1026,20 @@ async def _visit(self, node, depth):
 
     async def finish(self):
         """Wait for every summary; fails loud if the model never answered."""
+        # mark_final is a promise: the node stays in the tree and keeps its
+        # children. Verify it before awaiting anything - a broken promise
+        # leaves tasks whose marks can never arrive, a hang with no diagnosis.
+        live = {id(node) for node in _subtree(self.structure)}
+        for node, children in self._finals:
+            if id(node) not in live or tuple(id(c) for c in node.get('nodes') or []) != children:
+                name = node.get('title') or node.get('node_id') or '?'
+                raise RuntimeError(f"node {name!r} was dropped or changed "
+                                   "after it was marked final")
+        undecided = sum(1 for nid in self._tasks
+                        if not (mark := self._marks.get(nid)) or not mark.done())
+        if undecided:
+            raise RuntimeError(f"{undecided} node(s) were tasked but never marked final; "
+                               "their summaries would wait forever")
         results = await asyncio.gather(*(self._task(root, 1) for root in self.structure),
                                        return_exceptions=True)
         for r in results:
diff --git a/tests/test_client.py b/tests/test_client.py
index e9db0746a..892497ef2 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1208,6 +1208,7 @@ async def fake(model, prompt):
         {"title": "D", "start_index": 3, "end_index": 3}]}]
     asyncio.run(pageindex.utils.summarize_tree(
         structure, pdf_pages, small_node_tokens=0, concurrency=2))
+    assert started.index("gamma") < started.index("alpha")
     assert started.index("gamma") < started.index("delta")
     assert started.count("Section Title") == 2
 
@@ -1267,6 +1268,66 @@ async def scenario():
     assert [n["summary"] for n in (leaf, root["nodes"][1], root)] == ["ok"] * 3
 
 
+def test_leaf_summary_skips_the_tokenizer_when_length_already_rules_small_out(monkeypatch):
+    """No text averages 8+ characters per token, so a long leaf cannot come
+    in under the raw-text floor: its length alone settles the small-node
+    gate and the tokenizer, a synchronous cost on the loop, never runs."""
+    counted = []
+    real = pageindex.utils.count_tokens
+    monkeypatch.setattr(pageindex.utils, "count_tokens",
+                        lambda text, model=None: counted.append(len(text)) or real(text, model=model))
+
+    async def fake(model, prompt):
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+    long_leaf = [{"title": "A", "start_index": 1, "end_index": 1}]
+    asyncio.run(pageindex.utils.summarize_tree(long_leaf, [("word " * 1000, 1000)]))
+    assert long_leaf[0]["summary"] == "ok" and counted == []
+    short_leaf = [{"title": "B", "start_index": 1, "end_index": 1}]
+    asyncio.run(pageindex.utils.summarize_tree(short_leaf, [("tiny", 1)]))
+    assert short_leaf[0]["summary"] == "tiny" and counted == [4]
+
+
+def test_finish_fails_loud_when_a_final_node_changes_afterwards(monkeypatch):
+    """mark_final is a promise: the node stays in the tree and keeps its
+    children. finish() verifies it, so a run that broke the promise fails
+    with a diagnosis instead of shipping a tree missing a summarized
+    subtree."""
+    async def fake(model, prompt):
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+    pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+    leaf = {"title": "A", "start_index": 1, "end_index": 1}
+    root = {"title": "R", "start_index": 1, "end_index": 2, "nodes": [leaf]}
+
+    async def scenario():
+        scheduler = pageindex.utils.SummaryScheduler([root], pdf_pages, small_node_tokens=0)
+        scheduler.mark_final(list(pageindex.utils._subtree([root])))
+        root["nodes"] = []            # the promise, broken
+        await scheduler.finish()
+    with pytest.raises(RuntimeError, match="marked final"):
+        asyncio.run(asyncio.wait_for(scenario(), 5))
+
+
+def test_finish_fails_loud_instead_of_waiting_on_a_node_never_marked(monkeypatch):
+    """A tasked node whose mark never arrives would park finish() forever,
+    indistinguishable from a slow model. Marking a proper subset must be
+    named as the protocol breach it is."""
+    async def fake(model, prompt):
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
+    pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+    leaf = {"title": "A", "start_index": 1, "end_index": 1}
+    root = {"title": "R", "start_index": 1, "end_index": 2, "nodes": [leaf]}
+
+    async def scenario():
+        scheduler = pageindex.utils.SummaryScheduler([root], pdf_pages, small_node_tokens=0)
+        scheduler.mark_final([root])  # tasks the subtree, marks only the root
+        await scheduler.finish()
+    with pytest.raises(RuntimeError, match="marked final"):
+        asyncio.run(asyncio.wait_for(scenario(), 5))
+
+
 def test_priority_gate_passes_the_permit_on_when_its_taker_is_cancelled():
     """A waiter cancelled after the permit was granted but before it ran
     must hand the permit on, or the pool shrinks by one for good."""
@@ -2072,6 +2133,30 @@ async def propose(model, prompt):
     assert all(children != 4 for report in reports for _, children in report)
 
 
+def test_expand_keeps_same_page_children_when_merging_is_off(monkeypatch):
+    """do_merge=False turns off every merge, the fusion inside expand
+    included: a caller who disabled merging keeps every proposed node and
+    every document title."""
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, pages, lines = _expand_fixture()
+    pages[3] = "Sub One\nSub Two\n" + pages[3]
+    pages[4] = "Sub Three\n" + pages[4]
+    pages[8] = "Sub Four\n" + pages[8]
+    lines = [[l for l in p.splitlines() if l.strip()] for p in pages]
+
+    async def propose(model, prompt):
+        return {"subsections": [{"title": "Sub One", "page": 4}, {"title": "Sub Two", "page": 4},
+                                {"title": "Sub Three", "page": 5}, {"title": "Sub Four", "page": 9}]}
+    monkeypatch.setattr(tree_optimize, "ask_model", propose)
+    outcome = asyncio.run(tree_optimize.optimize(
+        tree, pages, lines, model="m", do_expand=True, do_merge=False, max_rounds=1))
+    children = tree[0]["nodes"][1]["nodes"]
+    assert [n["title"] for n in children] == ["Sub One", "Sub Two", "Sub Three", "Sub Four"]
+    assert outcome["same_page_merges"] == 0
+    assert not any(n.get("_same_page") for n in children)
+
+
 def test_flash_summaries_start_while_expand_is_still_deciding(tmp_path, monkeypatch):
     """With expand on, summaries run on the same loop: a leaf expand cannot
     touch is already being summarized when the model answers about X. Every

From c3d760a41b180910adb585c68b81a97fab4c4577 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 02:38:02 +0800
Subject: [PATCH 06/18] docs: merge_same_page and _optimize_async docstrings
 catch up with the overlap

Both still described the pre-overlap order: merge_same_page "runs first" /
"before merge() ... before expand()" (it now also runs right after merge and
on freshly attached children inside expand), and _optimize_async sitting
"between extraction and summaries" (summaries now run alongside it).

Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
---
 pageindex/flash/api.py     |  2 +-
 pageindex/tree_optimize.py | 10 ++++++----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index 2a133488a..25750a138 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -74,7 +74,7 @@ async def _summarize(structure, page_list, model, concurrency=None):
 
 
 async def _optimize_async(structure, page_texts, do_expand, model, on_final=None):
-    """Merge/expand refinement between extraction and summaries.
+    """Merge/expand refinement after extraction, overlapped with the summaries.
 
     Beyond the merge the default path runs anyway, this adds LLM expand and
     reports before/after search-cost metrics. Expand reads the same page text
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index c8484932e..3d79945c3 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -36,7 +36,8 @@
 `key_items`: the pages stay reachable by scanning the parent, but the titles
 are routing information that would otherwise be lost.
 
-merge_same_page() runs first, as a special case of the same idea. Retrieval is
+merge_same_page() runs wherever same-page duplicates can appear, as a special
+case of the same idea. Retrieval is
 page-granular, so frontier siblings covering identical pages cannot be told apart:
 an agent routed to any of them reads the same text, and because the leaf summary
 prompt sees only that text, their summaries come back near-identical. They collapse
@@ -491,9 +492,10 @@ def union_title(titles, node):
 def merge_same_page(structure, log):
     """Collapse frontier siblings that cover exactly the same pages.
 
-    Deterministic and free. Runs before merge() because a narrower tree changes
-    its ancestors' tree_cost, and before expand() because children an expand pass
-    lands on one page are the same redundancy arriving later.
+    Deterministic and free. Runs at every seam where the redundancy can appear:
+    before merge() because a narrower tree changes its ancestors' tree_cost,
+    again after it because a collapsed subtree can land on a sibling's exact
+    pages, and on a node's children right after expand attaches them.
     """
     changed = False
 

From 551d5bc41a932da20b0e8e845416c12f6f577cda Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 02:43:32 +0800
Subject: [PATCH 07/18] docs: name the load-bearing strict < in expand's keep

The settled promise (a final node is never touched again) rests on this
inequality staying strict: expansion then strictly lowers tree_cost while
spans never change, so merge's span <= cost can never newly hold in later
rounds. At <=, a zero-gain expand would create a node merge folds right
back. finish()'s contract check catches a break at runtime; this line
warns before the edit.

Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
---
 pageindex/tree_optimize.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index 3d79945c3..ded12c69b 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -722,6 +722,8 @@ async def process(node):
         cost = best["expand_cost"]
         gain = span - cost
         ratio = gain / span if span else 0.0
+        # strict: at cost == span the next round's merge (span <= cost) would
+        # fold the node right back, touching children already marked final
         keep = cost < span and ratio >= args.min_gain_ratio
 
         note(args.progress,

From c25aaea6a110f769dbe9e7d90ee5151f1b03bc93 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 03:49:58 +0800
Subject: [PATCH 08/18] fix: the round log entry counts the fusions inside
 expand

A same-page fusion during expansion bumped the run counters while the
round entry still said same_page=False. The flag is now derived from the
round's own log slice, so it counts every fusion and cannot disagree with
the counters. Loop exit is unchanged: an in-expand fusion only happens
after an attach, which already sets expanded.

Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
---
 pageindex/tree_optimize.py | 9 +++++++--
 tests/test_client.py       | 2 ++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index ded12c69b..4c8fd752f 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -825,14 +825,19 @@ def settled(nodes):
     for round_no in range(1, max_rounds + 1):
         rounds = round_no
         note(progress, f"  round {round_no}")
-        same_page = merge_same_page(structure, log) if do_merge else False
+        round_start = len(log)
+        if do_merge:
+            merge_same_page(structure, log)
         merged = merge(structure, routing, log, frozen, progress) if do_merge else False
         if merged:
             # a collapsed subtree can land on a sibling's exact pages
-            same_page = merge_same_page(structure, log) or same_page
+            merge_same_page(structure, log)
         settled(structure)
         expanded = await expand(structure, pages, lines, opts, log, frozen) \
             if do_expand else False
+        # derived from this round's log slice, so the flag counts the fusions
+        # inside expand too and cannot disagree with the run counters
+        same_page = any(e["op"] == "merge_same_page" for e in log[round_start:])
         log.append({"op": "round", "round": round_no, "same_page": same_page,
                     "merged": merged, "expanded": expanded})
         if not (same_page or merged or expanded):
diff --git a/tests/test_client.py b/tests/test_client.py
index 892497ef2..a701fd952 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -2131,6 +2131,8 @@ async def propose(model, prompt):
     assert outcome["expands"] == 1 and outcome["same_page_merges"] == 1
     assert [n["title"] for n in tree[0]["nodes"][1]["nodes"]][1:] == ["Sub Three", "Sub Four"]
     assert all(children != 4 for report in reports for _, children in report)
+    # the round entry agrees with the run counters: the fusion happened this round
+    assert [e["same_page"] for e in outcome["log"] if e["op"] == "round"] == [True]
 
 
 def test_expand_keeps_same_page_children_when_merging_is_off(monkeypatch):

From 3e83148296d61976594fdc5f1722123344a758b0 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 14:34:58 +0800
Subject: [PATCH 09/18] fix: drop the leaf tokenizer prefilter, count with the
 summary model

The length gate bb6ee89 put in front of count_tokens rested on "no text
averages 8+ chars per token", which is false: dot-leader contents pages
run 9-10 chars per token, rule and whitespace runs 60-120. Because the two
tests were and-ed, a leaf over 1600 chars but under the 200-token floor
paid a model call where it was documented to reuse its raw text.

The 2.6 s the prefilter was measured to save was import litellm, paid by
the first count_tokens in a harness whose faked model never imported it.
On the nine example PDFs with litellm already loaded, as a real run's
first model call leaves it, the prefilter moves expand's first call by
10-200 ms per document (PRML 0.36 s -> 0.17 s, about 0.1% of the run).

The tokenizer now takes the summary model like every other count_tokens
site, instead of a hardcoded gpt-4o. litellm falls back to cl100k_base
for names tiktoken does not know and swallows HF tokenizer failures, so
no model name raises; names outside the gpt-4o family count CJK text
about 2x higher, which narrows the raw-text floor for those runs.
---
 pageindex/utils.py   |  5 +----
 tests/test_client.py | 20 --------------------
 2 files changed, 1 insertion(+), 24 deletions(-)

diff --git a/pageindex/utils.py b/pageindex/utils.py
index 1ba5d5fde..f70ac3cc1 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -940,10 +940,7 @@ async def _ask(self, prompt, prio):
 
     async def _leaf_summary(self, node, prio):
         text = get_text_of_pdf_pages(self._pdf_pages, node['start_index'], node['end_index'])
-        # no text averages 8+ chars per token, so length alone clears most
-        # leaves without paying for the tokenizer on the loop
-        if (len(text) <= self._small_node_tokens * 8
-                and count_tokens(text, model="gpt-4o") < self._small_node_tokens):
+        if count_tokens(text, model=self._model) < self._small_node_tokens:
             return text.strip()
 
         # A node merged from same-page siblings carries a title joined from theirs.
diff --git a/tests/test_client.py b/tests/test_client.py
index a701fd952..38b430993 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1268,26 +1268,6 @@ async def scenario():
     assert [n["summary"] for n in (leaf, root["nodes"][1], root)] == ["ok"] * 3
 
 
-def test_leaf_summary_skips_the_tokenizer_when_length_already_rules_small_out(monkeypatch):
-    """No text averages 8+ characters per token, so a long leaf cannot come
-    in under the raw-text floor: its length alone settles the small-node
-    gate and the tokenizer, a synchronous cost on the loop, never runs."""
-    counted = []
-    real = pageindex.utils.count_tokens
-    monkeypatch.setattr(pageindex.utils, "count_tokens",
-                        lambda text, model=None: counted.append(len(text)) or real(text, model=model))
-
-    async def fake(model, prompt):
-        return '{"points": [], "summary": "ok"}'
-    monkeypatch.setattr(pageindex.utils, "llm_acompletion", fake)
-    long_leaf = [{"title": "A", "start_index": 1, "end_index": 1}]
-    asyncio.run(pageindex.utils.summarize_tree(long_leaf, [("word " * 1000, 1000)]))
-    assert long_leaf[0]["summary"] == "ok" and counted == []
-    short_leaf = [{"title": "B", "start_index": 1, "end_index": 1}]
-    asyncio.run(pageindex.utils.summarize_tree(short_leaf, [("tiny", 1)]))
-    assert short_leaf[0]["summary"] == "tiny" and counted == [4]
-
-
 def test_finish_fails_loud_when_a_final_node_changes_afterwards(monkeypatch):
     """mark_final is a promise: the node stays in the tree and keeps its
     children. finish() verifies it, so a run that broke the promise fails

From c35c152d32be4a1e35ff598bf65d20faf8e23fc6 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 14:57:00 +0800
Subject: [PATCH 10/18] fix: finish() guards every live node; flash names its
 two models

- finish()'s tasked-but-unmarked check ran over self._tasks before the
  root tasks existed, so a root that never received mark_final slipped
  past it and finish() parked forever - the symptomless hang the check
  was added to prevent. It now counts every node in the tree.
- _optimize_and_summarize took the expand model as `model`, the name the
  neighbouring _summarize uses for the summary model, with both passed
  positionally as adjacent strings. Renamed to optimize_model and called
  by keyword; the overlap test now runs the two lanes on different names
  and pins which model each receives.
- expand's keep_collapsed settle had no coverage: the existing test exits
  through no_children. A second test forces the cost-rejected branch with
  min_gain_ratio and checks the node is reported final before the closing
  sweep.
- merge_same_page's docstrings claimed every seam; merge_tree() is a bare
  merge(). They now name optimize()'s seams.

Claude-Session: https://claude.ai/code/session_018aNxBG6NTxK7hVwFCJK7Cx
---
 pageindex/flash/api.py     |  8 ++++----
 pageindex/tree_optimize.py |  6 +++---
 pageindex/utils.py         |  4 ++--
 tests/test_client.py       | 32 ++++++++++++++++++++++++++++++--
 4 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index 25750a138..53047c143 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -99,14 +99,14 @@ def _optimize(structure, page_texts, do_expand, model):
     return asyncio.run(_optimize_async(structure, page_texts, do_expand, model))
 
 
-async def _optimize_and_summarize(structure, page_texts, model, summary_model,
+async def _optimize_and_summarize(structure, page_texts, optimize_model, summary_model,
                                   concurrency):
     """Expand and summarize on one loop: a node is summarized as soon as
     expand can no longer change it, a parent once its children are done."""
     from ..utils import SummaryScheduler
     scheduler = SummaryScheduler(structure, [(text, 0) for text in page_texts],
                                  model=summary_model, concurrency=concurrency)
-    report = await _optimize_async(structure, page_texts, True, model,
+    report = await _optimize_async(structure, page_texts, True, optimize_model,
                                    on_final=scheduler.mark_final)
     await scheduler.finish()
     return report
@@ -146,8 +146,8 @@ def page_index_flash(pdf, summary=True, summary_model=None,
     if optimize and structure and summary and do_expand:
         import asyncio
         result["optimize"] = asyncio.run(_optimize_and_summarize(
-            structure, pages, optimize_model or summary_model,
-            summary_model, summary_concurrency))
+            structure, pages, optimize_model=optimize_model or summary_model,
+            summary_model=summary_model, concurrency=summary_concurrency))
         return result
     if optimize and structure:
         result["optimize"] = _optimize(structure, pages, do_expand,
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index 4c8fd752f..2013b6c87 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -36,8 +36,8 @@
 `key_items`: the pages stay reachable by scanning the parent, but the titles
 are routing information that would otherwise be lost.
 
-merge_same_page() runs wherever same-page duplicates can appear, as a special
-case of the same idea. Retrieval is
+merge_same_page() runs at every seam of optimize() where same-page duplicates can
+appear, as a special case of the same idea. Retrieval is
 page-granular, so frontier siblings covering identical pages cannot be told apart:
 an agent routed to any of them reads the same text, and because the leaf summary
 prompt sees only that text, their summaries come back near-identical. They collapse
@@ -492,7 +492,7 @@ def union_title(titles, node):
 def merge_same_page(structure, log):
     """Collapse frontier siblings that cover exactly the same pages.
 
-    Deterministic and free. Runs at every seam where the redundancy can appear:
+    Deterministic and free. Runs at every optimize() seam where the redundancy can appear:
     before merge() because a narrower tree changes its ancestors' tree_cost,
     again after it because a collapsed subtree can land on a sibling's exact
     pages, and on a node's children right after expand attaches them.
diff --git a/pageindex/utils.py b/pageindex/utils.py
index f70ac3cc1..c44040f51 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -1032,10 +1032,10 @@ async def finish(self):
                 name = node.get('title') or node.get('node_id') or '?'
                 raise RuntimeError(f"node {name!r} was dropped or changed "
                                    "after it was marked final")
-        undecided = sum(1 for nid in self._tasks
+        undecided = sum(1 for nid in live
                         if not (mark := self._marks.get(nid)) or not mark.done())
         if undecided:
-            raise RuntimeError(f"{undecided} node(s) were tasked but never marked final; "
+            raise RuntimeError(f"{undecided} node(s) were never marked final; "
                                "their summaries would wait forever")
         results = await asyncio.gather(*(self._task(root, 1) for root in self.structure),
                                        return_exceptions=True)
diff --git a/tests/test_client.py b/tests/test_client.py
index 38b430993..7bdc84192 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1307,6 +1307,14 @@ async def scenario():
     with pytest.raises(RuntimeError, match="marked final"):
         asyncio.run(asyncio.wait_for(scenario(), 5))
 
+    async def root_scenario():
+        other = {"title": "B", "start_index": 2, "end_index": 2}
+        scheduler = pageindex.utils.SummaryScheduler([root, other], pdf_pages, small_node_tokens=0)
+        scheduler.mark_final([root, leaf])  # a root nobody marked: tasked by finish() itself
+        await scheduler.finish()
+    with pytest.raises(RuntimeError, match="marked final"):
+        asyncio.run(asyncio.wait_for(root_scenario(), 5))
+
 
 def test_priority_gate_passes_the_permit_on_when_its_taker_is_cancelled():
     """A waiter cancelled after the permit was granted but before it ran
@@ -2071,6 +2079,23 @@ async def nothing(model, prompt):
     assert ["X"] in reports[:-1]
 
 
+def test_optimize_reports_a_candidate_kept_collapsed_for_too_little_gain_once_decided(monkeypatch):
+    """The other way to stay collapsed: the model proposes a split that does
+    not pay for itself. That node is final right then too."""
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, pages, lines = _expand_fixture()
+    reports = []
+
+    async def propose(model, prompt):
+        return {"subsections": [{"title": "Sub One", "page": 4}, {"title": "Sub Two", "page": 8}]}
+    monkeypatch.setattr(tree_optimize, "ask_model", propose)
+    outcome = asyncio.run(tree_optimize.optimize(
+        tree, pages, lines, model="m", do_expand=True, min_gain_ratio=0.99,
+        on_final=lambda nodes: reports.append(sorted(n["title"] for n in nodes))))
+    assert outcome["kept_collapsed"] == 1 and ["X"] in reports[:-1]
+
+
 def test_merge_fuses_the_same_page_frontier_it_creates_at_once():
     """Collapsing F leaves it on exactly G's page; the two must fuse right
     then, not one round later (with a single round they never would)."""
@@ -2154,9 +2179,10 @@ def test_flash_summaries_start_while_expand_is_still_deciding(tmp_path, monkeypa
     pdf.write_bytes(build_pdf(["x"]))
     monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
         "structure": copy.deepcopy(tree), "page_texts": list(pages)})
-    started, replied_after = [], []
+    started, replied_after, models = [], [], []
 
     async def propose(model, prompt):
+        models.append(("expand", model))
         await asyncio.sleep(0.05)
         replied_after.append(len(started))
         return {"subsections": [{"title": "Sub One", "page": 4},
@@ -2164,16 +2190,18 @@ async def propose(model, prompt):
     monkeypatch.setattr(tree_optimize, "ask_model", propose)
 
     async def summarize(model, prompt):
+        models.append(("summary", model))
         started.append(prompt)
         await asyncio.sleep(0.01)
         return '{"points": [], "summary": "ok"}'
     monkeypatch.setattr(pageindex.utils, "llm_acompletion", summarize)
 
-    result = flash_api.page_index_flash(str(pdf), summary_model="m")
+    result = flash_api.page_index_flash(str(pdf), summary_model="m", optimize_model="x")
     assert replied_after[0] >= 1
     assert result["optimize"]["expands"] == 1
     nodes = list(pageindex.utils._subtree(result["structure"]))
     assert [n["summary"] for n in nodes] == ["ok"] * 5 and len(started) == 5
+    assert set(models) == {("expand", "x"), ("summary", "m")}
     shape = lambda nodes: [(n["title"], n["start_index"], n["end_index"],
                             shape(n.get("nodes") or [])) for n in nodes]
     alone = flash_api.page_index_flash(str(pdf), summary=False)

From d4f05879a702e6b0f142a3651fa91f90d726f886 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 15:44:50 +0800
Subject: [PATCH 11/18] docs: _optimize_async overlaps with the summaries only
 when on_final is passed

---
 pageindex/flash/api.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index 53047c143..d61eec004 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -74,7 +74,8 @@ async def _summarize(structure, page_list, model, concurrency=None):
 
 
 async def _optimize_async(structure, page_texts, do_expand, model, on_final=None):
-    """Merge/expand refinement after extraction, overlapped with the summaries.
+    """Merge/expand refinement after extraction, overlapped with the summaries
+    when `on_final` is passed; without it the caller runs them after.
 
     Beyond the merge the default path runs anyway, this adds LLM expand and
     reports before/after search-cost metrics. Expand reads the same page text

From 06f9c87d92886c387973644a4c92635e3ff54840 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 18:33:55 +0800
Subject: [PATCH 12/18] perf: summary prompts ask for the summary alone, within
 summary_max_words

The two summary prompts asked for a "points" list next to the summary and
parse_summary threw the list away: about 70% of every reply was generated
to be discarded. Dropping the list alone is not enough. The model then
pours the same content into the summary (656 -> ~2000 chars) and parents,
which read their children's summaries, slow down more than the leaves
gained (fed wall +6%). With a word cap the reply shrinks for real.

Measured on gpt-5.6-luna, mirror A/B, summary stage only:
  per-call median   9.7 s  -> 5.3 s   (-45%)
  fed-2023 wall     47.5 s -> 30.7 s  (-35%)
  PRML wall         71.1 s -> 38.1 s  (-46%)
  reply chars       ~3340  -> ~1175   (-65% output tokens)
  summary length    ~670   -> ~1160 chars (the model writes cap x 1.25)
Blinded pairwise judge (claude-sonnet-5, source in view): 150-word
summaries beat the current ones 21-1-0 on fed + attention leaves. The
specifics that used to sit in the discarded list (law names, rates, terms)
now land in the summary. A full prompt rewrite lost to this two-line edit
8-11-3, so the edit stays minimal.

The cap is a parameter, default 150: page_index_flash(summary_max_words=),
PageIndexLocalClient(summary_max_words=) / index={"summary_max_words": n},
and run_pageindex.py --summary-max-words. The classic pipeline's summary
prompt is a different one without a points field and is untouched.
---
 pageindex/client.py    | 18 ++++++++---
 pageindex/flash/api.py | 20 +++++++-----
 pageindex/local_api.py |  7 +++--
 pageindex/types.py     |  1 +
 pageindex/utils.py     | 23 ++++++++------
 run_pageindex.py       |  5 +++
 tests/test_client.py   | 70 ++++++++++++++++++++++++++++++++++++++++++
 7 files changed, 120 insertions(+), 24 deletions(-)

diff --git a/pageindex/client.py b/pageindex/client.py
index 709b92d22..27409166f 100644
--- a/pageindex/client.py
+++ b/pageindex/client.py
@@ -58,7 +58,8 @@ def _agents_sdk_model_name(model: str) -> str:
     return f"litellm/{model}"
 
 
-_LOCAL_INDEX_KEYS = ("model", "summary_model", "backend", "storage_path")
+_LOCAL_INDEX_KEYS = ("model", "summary_model", "backend", "storage_path",
+                     "summary_max_words")
 
 # Near-synonyms of "cloud" that would otherwise parse as model names —
 # a silent wrong mode. They error, pointing at the real word.
@@ -83,7 +84,7 @@ def _env_cloud_key(spelling: str, inline: str = "api_key=...") -> str:
 # there as a PageIndexAPIError — never later, never silently.
 _ARG_TYPES: "dict[str, tuple[type, ...]]" = {
     "model": (str,), "index_model": (str,), "summary_model": (str,),
-    "chat_model": (str,), "retrieve_model": (str,),
+    "chat_model": (str,), "retrieve_model": (str,), "summary_max_words": (int,),
     "storage_path": (str, os.PathLike), "index_backend": (dict,),
     "chat_backend": (dict,)}
 
@@ -170,7 +171,8 @@ def _resolve_index_slot(index) -> "tuple[_CloudKey, dict[str, Any]]":
         mapped = {"index_model": conf.get("model"),
                   "summary_model": conf.get("summary_model"),
                   "index_backend": conf.get("backend"),
-                  "storage_path": conf.get("storage_path")}
+                  "storage_path": conf.get("storage_path"),
+                  "summary_max_words": conf.get("summary_max_words")}
         return None, {name: value for name, value in mapped.items()
                       if value is not None}
     raise PageIndexAPIError("index must be a string or a dict.")
@@ -264,7 +266,8 @@ class PageIndexClient:
             ``"cloud"`` / ``"pageindex-cloud"`` (cloud, key from the
             environment), ``"local"``, a local index model name, or a
             dict: ``{"api_key": ...}`` for cloud, ``{"model",
-            "summary_model", "backend", "storage_path"}`` for local. An
+            "summary_model", "backend", "storage_path",
+            "summary_max_words"}`` for local. An
             optional ``"mode"`` key (``"cloud"`` / ``"local"``) states
             the side and must agree with the other keys; ``{"mode":
             "cloud"}`` alone reads the key from the environment. Not
@@ -306,6 +309,8 @@ class PageIndexClient:
         summary_model (str, optional): Local mode only — legacy: overrides
             the model used for node summaries and document descriptions;
             ``index_model`` covers this.
+        summary_max_words (int, optional): Local mode only — the word cap
+            each node summary is asked to stay within. Defaults to 150.
         retrieve_model (str, optional): Legacy name for ``chat_model`` —
             same meaning everywhere, cloud clients included.
         storage_path (str or os.PathLike, optional): Local mode only —
@@ -346,6 +351,7 @@ def __init__(
         chat_model: Optional[str] = None,
         model: Optional[str] = None,
         summary_model: Optional[str] = None,
+        summary_max_words: Optional[int] = None,
         retrieve_model: Optional[str] = None,
         storage_path: Optional[Union[str, os.PathLike[str]]] = None,
         index_backend: Optional[dict[str, Any]] = None,
@@ -363,6 +369,7 @@ def __init__(
             (("api_key", api_key),
              ("index_model", index_model),
              ("summary_model", summary_model),
+             ("summary_max_words", summary_max_words),
              ("index_backend", index_backend),
              ("storage_path", storage_path), ("model", model))
             if value is not None}
@@ -506,6 +513,7 @@ def __init__(
                 model=self.model,
                 summary_model=self.summary_model,
                 index_backend=index_conf.get("index_backend"),
+                summary_max_words=index_conf.get("summary_max_words"),
             )
             # LiteLLM's multi-second import would otherwise land on the
             # first chat call; failures resurface there with real context.
@@ -1651,6 +1659,7 @@ def __init__(
         chat_model: Optional[str] = None,
         model: Optional[str] = None,
         summary_model: Optional[str] = None,
+        summary_max_words: Optional[int] = None,
         retrieve_model: Optional[str] = None,
         storage_path: Optional[Union[str, os.PathLike[str]]] = None,
         index_backend: Optional[dict[str, Any]] = None,
@@ -1659,5 +1668,6 @@ def __init__(
         super().__init__(None, index=index, chat=chat,
                          index_model=index_model, chat_model=chat_model,
                          model=model, summary_model=summary_model,
+                         summary_max_words=summary_max_words,
                          retrieve_model=retrieve_model, storage_path=storage_path,
                          index_backend=index_backend, chat_backend=chat_backend)
diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index d61eec004..ce82458a5 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -68,9 +68,10 @@ def _validate_pdf(pdf):
     return pdf
 
 
-async def _summarize(structure, page_list, model, concurrency=None):
+async def _summarize(structure, page_list, model, concurrency=None, max_words=None):
     from ..utils import summarize_tree
-    await summarize_tree(structure, page_list, model=model, concurrency=concurrency)
+    await summarize_tree(structure, page_list, model=model, concurrency=concurrency,
+                         max_words=max_words)
 
 
 async def _optimize_async(structure, page_texts, do_expand, model, on_final=None):
@@ -101,12 +102,13 @@ def _optimize(structure, page_texts, do_expand, model):
 
 
 async def _optimize_and_summarize(structure, page_texts, optimize_model, summary_model,
-                                  concurrency):
+                                  concurrency, max_words=None):
     """Expand and summarize on one loop: a node is summarized as soon as
     expand can no longer change it, a parent once its children are done."""
     from ..utils import SummaryScheduler
     scheduler = SummaryScheduler(structure, [(text, 0) for text in page_texts],
-                                 model=summary_model, concurrency=concurrency)
+                                 model=summary_model, concurrency=concurrency,
+                                 max_words=max_words)
     report = await _optimize_async(structure, page_texts, True, optimize_model,
                                    on_final=scheduler.mark_final)
     await scheduler.finish()
@@ -116,8 +118,8 @@ async def _optimize_and_summarize(structure, page_texts, optimize_model, summary
 def page_index_flash(pdf, summary=True, summary_model=None,
                      optimize: str | bool | None = None, optimize_expand=None,
                      optimize_model=None, summary_concurrency=None,
-                     use_embedded_toc=True) -> dict:
-    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+                     use_embedded_toc=True, summary_max_words=None) -> dict:
+    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
     if optimize_expand is not None:
         import warnings
         warnings.warn(
@@ -148,7 +150,8 @@ def page_index_flash(pdf, summary=True, summary_model=None,
         import asyncio
         result["optimize"] = asyncio.run(_optimize_and_summarize(
             structure, pages, optimize_model=optimize_model or summary_model,
-            summary_model=summary_model, concurrency=summary_concurrency))
+            summary_model=summary_model, concurrency=summary_concurrency,
+            max_words=summary_max_words))
         return result
     if optimize and structure:
         result["optimize"] = _optimize(structure, pages, do_expand,
@@ -157,7 +160,8 @@ def page_index_flash(pdf, summary=True, summary_model=None,
         import asyncio
         page_list = [(text, 0) for text in pages]
         asyncio.run(_summarize(structure, page_list, summary_model,
-                               concurrency=summary_concurrency))
+                               concurrency=summary_concurrency,
+                               max_words=summary_max_words))
     elif structure:
         from ..utils import strip_internal_keys
         strip_internal_keys(structure)   # summarize_tree does this on its way out
diff --git a/pageindex/local_api.py b/pageindex/local_api.py
index 40aa3aedc..d92418406 100644
--- a/pageindex/local_api.py
+++ b/pageindex/local_api.py
@@ -37,11 +37,13 @@ class LocalAPI:
     """Backs PageIndexClient's local mode. One instance per client."""
 
     def __init__(self, storage_path: str, model: str, summary_model: str,
-                 index_backend: dict | None = None):
+                 index_backend: dict | None = None,
+                 summary_max_words: int | None = None):
         self._store = DocStore(storage_path)
         self._model = model
         self._summary_model = summary_model
         self._index_backend = index_backend
+        self._summary_max_words = summary_max_words
         from .utils import ConfigLoader
         self._config_loader = ConfigLoader()
 
@@ -232,7 +234,8 @@ def _index_flash(self, file_path: str) -> tuple[list, str | None]:
         result = page_index_flash(file_path, summary=True,
                                   summary_model=self._summary_model,
                                   optimize="full",
-                                  optimize_model=self._summary_model)
+                                  optimize_model=self._summary_model,
+                                  summary_max_words=self._summary_max_words)
         structure = result.get("structure", [])
         if not structure:
             raise PageIndexAPIError(
diff --git a/pageindex/types.py b/pageindex/types.py
index b02962899..4596f15b8 100644
--- a/pageindex/types.py
+++ b/pageindex/types.py
@@ -32,6 +32,7 @@ class LocalIndexConfig(TypedDict, total=False):
     mode: Literal["local"]
     model: str
     summary_model: str
+    summary_max_words: int
     backend: dict
     storage_path: Union[str, os.PathLike[str]]
 
diff --git a/pageindex/utils.py b/pageindex/utils.py
index c44040f51..08e930bd6 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -745,6 +745,7 @@ async def generate_summaries_for_structure(structure, model=None):
 SUMMARY_CONCURRENCY = 64        # simultaneous summary model calls
 SUMMARY_RAW_TEXT_TOKENS = 200   # leaves under this reuse their raw text as the summary
 SUMMARY_INTRO_MAX_PAGES = 3     # cap on leading pages fed into a parent summary
+SUMMARY_MAX_WORDS = 150         # word cap the summary prompts ask for
 
 
 class _PriorityGate:
@@ -884,12 +885,14 @@ class SummaryScheduler:
 
     def __init__(self, structure, pdf_pages, model=None,
                  small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
-                 max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None):
+                 max_intro_pages=SUMMARY_INTRO_MAX_PAGES, max_words=None,
+                 concurrency=None):
         self.structure = structure
         self._pdf_pages = pdf_pages
         self._model = model
         self._small_node_tokens = small_node_tokens
         self._max_intro_pages = max_intro_pages
+        self._max_words = max_words or SUMMARY_MAX_WORDS
         self._gate = _PriorityGate(concurrency or SUMMARY_CONCURRENCY)
         self._asked = self._answered = False
         self._marks = {}     # id(node) -> future resolved once the node is final
@@ -957,13 +960,12 @@ async def _leaf_summary(self, node, prio):
 
         prompt = f"""You are given a text chunk from a document.
     Your task is to generate a concise description of everything that is covered in the text, summarizing all its points without omitting any type of content.
-    Keep the description concise and to the point, avoiding unnecessary details.{ask_title}
+    Keep the description concise and to the point, avoiding unnecessary details, within {self._max_words} words.{ask_title}
 
     Given Text: {text}
 
     Reply strictly in the following JSON format:
     {{{title_field}
-        "points": <a list of points covered in the text>,
         "summary": <a concise description of everything that is covered in the text, summarizing all its points without omitting any type of content>
     }}
 
@@ -984,7 +986,7 @@ async def _parent_summary(self, node, prio):
             ensure_ascii=False)
         prompt = f"""You are given a section of a document: the text that opens the section (possibly empty) and the titles and summaries of its subsections.
     Your task is to generate a concise description of everything that is covered in the whole section, summarizing all its points without omitting any type of content.
-    Keep the description concise and to the point, avoiding unnecessary details.
+    Keep the description concise and to the point, avoiding unnecessary details, within {self._max_words} words.
 
     Section Title: {node.get('title', '')}
 
@@ -994,7 +996,6 @@ async def _parent_summary(self, node, prio):
 
     Reply strictly in the following JSON format:
     {{
-        "points": <a list of points covered in the section>,
         "summary": <a concise description of everything that is covered in the section, summarizing all its points without omitting any type of content>
     }}
 
@@ -1061,18 +1062,20 @@ def _any_summary(nodes):
 
 async def summarize_tree(structure, pdf_pages, model=None,
                          small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
-                         max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None):
+                         max_intro_pages=SUMMARY_INTRO_MAX_PAGES, max_words=None,
+                         concurrency=None):
     """Bottom-up summaries: leaves from their own pages, parents composed from
     child summaries plus the pages no child covers. A parent's summary describes
     its whole subtree (end_index union semantics). Nodes that already carry a
     summary are left untouched; leaves under `small_node_tokens` use their raw
-    text as the summary without a model call. Model calls run deepest node
-    first, both in starting order and in leaving the queue: depth counts the
-    calls left on a node's path to the root, its own included."""
+    text as the summary without a model call; every prompt asks for at most
+    `max_words` words. Model calls run deepest node first, both in starting
+    order and in leaving the queue: depth counts the calls left on a node's
+    path to the root, its own included."""
     scheduler = SummaryScheduler(structure, pdf_pages, model=model,
                                  small_node_tokens=small_node_tokens,
                                  max_intro_pages=max_intro_pages,
-                                 concurrency=concurrency)
+                                 max_words=max_words, concurrency=concurrency)
     scheduler.mark_final(list(_subtree(structure)))
     return await scheduler.finish()
 
diff --git a/run_pageindex.py b/run_pageindex.py
index f2ae111be..d592c147f 100644
--- a/run_pageindex.py
+++ b/run_pageindex.py
@@ -33,6 +33,8 @@
                       help='(legacy) Same as --index-model')
     parser.add_argument('--summary-model', type=str, default=None,
                       help='Model for node summaries (falls back to config.yaml summary_model, then --index-model, then --model)')
+    parser.add_argument('--summary-max-words', type=int, default=None,
+                      help='Word cap for each node summary (flash mode; default 150)')
 
     parser.add_argument('--toc-check-pages', type=int, default=None,
                       help='Number of pages to check for table of contents (PDF only)')
@@ -74,6 +76,8 @@
         raise ValueError("--embedded-toc requires Flash mode with --pdf_path")
     if args.summary is not None and not (args.pdf_path and args.mode == 'flash'):
         raise ValueError("--summary requires Flash mode with --pdf_path")
+    if args.summary_max_words is not None and not (args.pdf_path and args.mode == 'flash'):
+        raise ValueError("--summary-max-words requires Flash mode with --pdf_path")
     if args.pdf_path and args.mode == 'flash':
         for flag, value in (('--toc-check-pages', args.toc_check_pages),
                             ('--max-pages-per-node', args.max_pages_per_node),
@@ -107,6 +111,7 @@
                 summary_model=summary_model,
                 use_embedded_toc=args.embedded_toc if args.embedded_toc is not None else True,
                 summary=will_summarize,
+                summary_max_words=args.summary_max_words,
             )
             if not toc_with_page_number.get('structure'):
                 raise ValueError("PageIndex Flash could not extract a structure from this PDF; "
diff --git a/tests/test_client.py b/tests/test_client.py
index 7bdc84192..fd8d1ced1 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -2309,3 +2309,73 @@ def test_blank_chat_model_carries_no_model_into_agent_config():
     client = PageIndexClient()
     client.chat_model = "   "
     assert "model" not in client.openai_agent_config()
+
+
+def test_summary_prompts_cap_words_and_omit_points(monkeypatch):
+    """Both summary prompts ask for the summary alone, within the word cap.
+    The points list the model wrote first was parsed and thrown away, and
+    with it gone the summary swallows its content unless a cap holds it."""
+    prompts = []
+
+    async def capture(model, prompt):
+        prompts.append(prompt)
+        return '{"summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", capture)
+    pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+
+    def tree():
+        return [{"title": "R", "start_index": 1, "end_index": 2,
+                 "nodes": [{"title": "A", "start_index": 1, "end_index": 1},
+                           {"title": "B", "start_index": 2, "end_index": 2}]}]
+
+    asyncio.run(pageindex.utils.summarize_tree(tree(), pdf_pages, small_node_tokens=0))
+    assert len(prompts) == 3
+    assert all("within 150 words" in p and '"points"' not in p for p in prompts)
+
+    prompts.clear()
+    asyncio.run(pageindex.utils.summarize_tree(tree(), pdf_pages, small_node_tokens=0,
+                                               max_words=80))
+    assert len(prompts) == 3 and all("within 80 words" in p for p in prompts)
+
+
+def test_page_index_flash_plumbs_summary_max_words(tmp_path, monkeypatch):
+    """summary_max_words reaches the prompts on both summary paths: the
+    plain one and the one overlapped with expand."""
+    from conftest import build_pdf
+    import copy
+    import pageindex.flash.api as flash_api
+    import pageindex.tree_optimize as tree_optimize
+
+    tree, pages, _ = _expand_fixture()
+    pdf = tmp_path / "doc.pdf"
+    pdf.write_bytes(build_pdf(["x"]))
+    monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
+        "structure": copy.deepcopy(tree), "page_texts": list(pages)})
+
+    async def propose(model, prompt):
+        return {"subsections": [{"title": "Sub One", "page": 4},
+                                {"title": "Sub Two", "page": 8}]}
+    monkeypatch.setattr(tree_optimize, "ask_model", propose)
+    prompts = []
+
+    async def capture(model, prompt):
+        prompts.append(prompt)
+        return '{"summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", capture)
+
+    for optimize in (False, "full"):
+        prompts.clear()
+        flash_api.page_index_flash(str(pdf), summary_model="m", optimize=optimize,
+                                   summary_max_words=80)
+        assert prompts and all("within 80 words" in p for p in prompts), optimize
+
+
+def test_summary_max_words_reaches_the_local_indexer():
+    """index["summary_max_words"] and the flat argument both land on the
+    local indexer; a non-int refuses in the constructor like every slot."""
+    from pageindex import PageIndexLocalClient
+    assert PageIndexLocalClient(index={"summary_max_words": 80})._api._summary_max_words == 80
+    assert PageIndexLocalClient(summary_max_words=80)._api._summary_max_words == 80
+    assert PageIndexLocalClient()._api._summary_max_words is None
+    with pytest.raises(PageIndexAPIError, match=r'index\["summary_max_words"\] must be a'):
+        PageIndexLocalClient(index={"summary_max_words": "80"})

From 1770a1a42bacce2d0a7ece074e8062d0a5270b2d Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 18:40:07 +0800
Subject: [PATCH 13/18] feat: summary_concurrency, use_embedded_toc and
 optimize on the client

The CLI could set all three; the SDK could not. They join the local index
side the way summary_max_words did: flat arguments on both clients, keys
in the index= slot, LocalAPI hands them to page_index_flash. optimize
takes the CLI's words ("full" / "merge" / "off"; "off" is no optimize
pass at all), so a wrong value refuses in the constructor like every
other slot. The slot check's empty-value rule now lets False through,
since use_embedded_toc is the first bool in the vocabulary.
summary_concurrency also gets a CLI flag.
---
 pageindex/client.py            | 39 ++++++++++++++++++++++++++++++----
 pageindex/local_api.py         | 14 +++++++++---
 pageindex/types.py             |  3 +++
 run_pageindex.py               |  5 +++++
 tests/test_client.py           | 39 ++++++++++++++++++++++++++++++++++
 tests/test_flash_extraction.py |  6 ++++++
 6 files changed, 99 insertions(+), 7 deletions(-)

diff --git a/pageindex/client.py b/pageindex/client.py
index 27409166f..0df387c80 100644
--- a/pageindex/client.py
+++ b/pageindex/client.py
@@ -59,7 +59,8 @@ def _agents_sdk_model_name(model: str) -> str:
 
 
 _LOCAL_INDEX_KEYS = ("model", "summary_model", "backend", "storage_path",
-                     "summary_max_words")
+                     "summary_max_words", "summary_concurrency",
+                     "use_embedded_toc", "optimize")
 
 # Near-synonyms of "cloud" that would otherwise parse as model names —
 # a silent wrong mode. They error, pointing at the real word.
@@ -85,6 +86,7 @@ def _env_cloud_key(spelling: str, inline: str = "api_key=...") -> str:
 _ARG_TYPES: "dict[str, tuple[type, ...]]" = {
     "model": (str,), "index_model": (str,), "summary_model": (str,),
     "chat_model": (str,), "retrieve_model": (str,), "summary_max_words": (int,),
+    "summary_concurrency": (int,), "use_embedded_toc": (bool,), "optimize": (str,),
     "storage_path": (str, os.PathLike), "index_backend": (dict,),
     "chat_backend": (dict,)}
 
@@ -172,7 +174,10 @@ def _resolve_index_slot(index) -> "tuple[_CloudKey, dict[str, Any]]":
                   "summary_model": conf.get("summary_model"),
                   "index_backend": conf.get("backend"),
                   "storage_path": conf.get("storage_path"),
-                  "summary_max_words": conf.get("summary_max_words")}
+                  "summary_max_words": conf.get("summary_max_words"),
+                  "summary_concurrency": conf.get("summary_concurrency"),
+                  "use_embedded_toc": conf.get("use_embedded_toc"),
+                  "optimize": conf.get("optimize")}
         return None, {name: value for name, value in mapped.items()
                       if value is not None}
     raise PageIndexAPIError("index must be a string or a dict.")
@@ -267,7 +272,8 @@ class PageIndexClient:
             environment), ``"local"``, a local index model name, or a
             dict: ``{"api_key": ...}`` for cloud, ``{"model",
             "summary_model", "backend", "storage_path",
-            "summary_max_words"}`` for local. An
+            "summary_max_words", "summary_concurrency", "use_embedded_toc",
+            "optimize"}`` for local. An
             optional ``"mode"`` key (``"cloud"`` / ``"local"``) states
             the side and must agree with the other keys; ``{"mode":
             "cloud"}`` alone reads the key from the environment. Not
@@ -311,6 +317,14 @@ class PageIndexClient:
             ``index_model`` covers this.
         summary_max_words (int, optional): Local mode only — the word cap
             each node summary is asked to stay within. Defaults to 150.
+        summary_concurrency (int, optional): Local mode only — cap on
+            simultaneous indexing model calls. Defaults to 64.
+        use_embedded_toc (bool, optional): Local mode only — whether flash
+            indexing consumes the PDF's embedded bookmarks when they look
+            trustworthy. Defaults to True.
+        optimize (str, optional): Local mode only — the flash tree
+            refinement pass: ``"full"`` (merge + model expand, the
+            default), ``"merge"`` (deterministic merge only) or ``"off"``.
         retrieve_model (str, optional): Legacy name for ``chat_model`` —
             same meaning everywhere, cloud clients included.
         storage_path (str or os.PathLike, optional): Local mode only —
@@ -352,6 +366,9 @@ def __init__(
         model: Optional[str] = None,
         summary_model: Optional[str] = None,
         summary_max_words: Optional[int] = None,
+        summary_concurrency: Optional[int] = None,
+        use_embedded_toc: Optional[bool] = None,
+        optimize: Optional[str] = None,
         retrieve_model: Optional[str] = None,
         storage_path: Optional[Union[str, os.PathLike[str]]] = None,
         index_backend: Optional[dict[str, Any]] = None,
@@ -370,6 +387,9 @@ def __init__(
              ("index_model", index_model),
              ("summary_model", summary_model),
              ("summary_max_words", summary_max_words),
+             ("summary_concurrency", summary_concurrency),
+             ("use_embedded_toc", use_embedded_toc),
+             ("optimize", optimize),
              ("index_backend", index_backend),
              ("storage_path", storage_path), ("model", model))
             if value is not None}
@@ -450,10 +470,13 @@ def __init__(
                         f"got {type(value).__name__}.")
                 if isinstance(value, str):
                     value = conf[name] = value.strip()
-                if not value:
+                if not value and not isinstance(value, bool):
                     raise PageIndexAPIError(
                         f"{shown} is empty — it configures nothing. Pass a "
                         "real value, or drop the argument.")
+                if name == "optimize" and value not in ("full", "merge", "off"):
+                    raise PageIndexAPIError(
+                        f'{shown} must be "full", "merge" or "off", got {value!r}.')
 
         if cloud_key is not None:
             if index_conf:
@@ -514,6 +537,9 @@ def __init__(
                 summary_model=self.summary_model,
                 index_backend=index_conf.get("index_backend"),
                 summary_max_words=index_conf.get("summary_max_words"),
+                summary_concurrency=index_conf.get("summary_concurrency"),
+                use_embedded_toc=index_conf.get("use_embedded_toc", True),
+                optimize=index_conf.get("optimize", "full"),
             )
             # LiteLLM's multi-second import would otherwise land on the
             # first chat call; failures resurface there with real context.
@@ -1660,6 +1686,9 @@ def __init__(
         model: Optional[str] = None,
         summary_model: Optional[str] = None,
         summary_max_words: Optional[int] = None,
+        summary_concurrency: Optional[int] = None,
+        use_embedded_toc: Optional[bool] = None,
+        optimize: Optional[str] = None,
         retrieve_model: Optional[str] = None,
         storage_path: Optional[Union[str, os.PathLike[str]]] = None,
         index_backend: Optional[dict[str, Any]] = None,
@@ -1669,5 +1698,7 @@ def __init__(
                          index_model=index_model, chat_model=chat_model,
                          model=model, summary_model=summary_model,
                          summary_max_words=summary_max_words,
+                         summary_concurrency=summary_concurrency,
+                         use_embedded_toc=use_embedded_toc, optimize=optimize,
                          retrieve_model=retrieve_model, storage_path=storage_path,
                          index_backend=index_backend, chat_backend=chat_backend)
diff --git a/pageindex/local_api.py b/pageindex/local_api.py
index d92418406..5e8616293 100644
--- a/pageindex/local_api.py
+++ b/pageindex/local_api.py
@@ -38,12 +38,18 @@ class LocalAPI:
 
     def __init__(self, storage_path: str, model: str, summary_model: str,
                  index_backend: dict | None = None,
-                 summary_max_words: int | None = None):
+                 summary_max_words: int | None = None,
+                 summary_concurrency: int | None = None,
+                 use_embedded_toc: bool = True,
+                 optimize: str = "full"):
         self._store = DocStore(storage_path)
         self._model = model
         self._summary_model = summary_model
         self._index_backend = index_backend
         self._summary_max_words = summary_max_words
+        self._summary_concurrency = summary_concurrency
+        self._use_embedded_toc = use_embedded_toc
+        self._optimize = optimize
         from .utils import ConfigLoader
         self._config_loader = ConfigLoader()
 
@@ -233,9 +239,11 @@ def _index_flash(self, file_path: str) -> tuple[list, str | None]:
                             generate_doc_description, write_node_id)
         result = page_index_flash(file_path, summary=True,
                                   summary_model=self._summary_model,
-                                  optimize="full",
+                                  optimize=False if self._optimize == "off" else self._optimize,
                                   optimize_model=self._summary_model,
-                                  summary_max_words=self._summary_max_words)
+                                  summary_concurrency=self._summary_concurrency,
+                                  summary_max_words=self._summary_max_words,
+                                  use_embedded_toc=self._use_embedded_toc)
         structure = result.get("structure", [])
         if not structure:
             raise PageIndexAPIError(
diff --git a/pageindex/types.py b/pageindex/types.py
index 4596f15b8..8a7da077f 100644
--- a/pageindex/types.py
+++ b/pageindex/types.py
@@ -33,6 +33,9 @@ class LocalIndexConfig(TypedDict, total=False):
     model: str
     summary_model: str
     summary_max_words: int
+    summary_concurrency: int
+    use_embedded_toc: bool
+    optimize: Literal["full", "merge", "off"]
     backend: dict
     storage_path: Union[str, os.PathLike[str]]
 
diff --git a/run_pageindex.py b/run_pageindex.py
index d592c147f..716ad4e4c 100644
--- a/run_pageindex.py
+++ b/run_pageindex.py
@@ -35,6 +35,8 @@
                       help='Model for node summaries (falls back to config.yaml summary_model, then --index-model, then --model)')
     parser.add_argument('--summary-max-words', type=int, default=None,
                       help='Word cap for each node summary (flash mode; default 150)')
+    parser.add_argument('--summary-concurrency', type=int, default=None,
+                      help='Cap on simultaneous indexing model calls (flash mode; default 64)')
 
     parser.add_argument('--toc-check-pages', type=int, default=None,
                       help='Number of pages to check for table of contents (PDF only)')
@@ -78,6 +80,8 @@
         raise ValueError("--summary requires Flash mode with --pdf_path")
     if args.summary_max_words is not None and not (args.pdf_path and args.mode == 'flash'):
         raise ValueError("--summary-max-words requires Flash mode with --pdf_path")
+    if args.summary_concurrency is not None and not (args.pdf_path and args.mode == 'flash'):
+        raise ValueError("--summary-concurrency requires Flash mode with --pdf_path")
     if args.pdf_path and args.mode == 'flash':
         for flag, value in (('--toc-check-pages', args.toc_check_pages),
                             ('--max-pages-per-node', args.max_pages_per_node),
@@ -112,6 +116,7 @@
                 use_embedded_toc=args.embedded_toc if args.embedded_toc is not None else True,
                 summary=will_summarize,
                 summary_max_words=args.summary_max_words,
+                summary_concurrency=args.summary_concurrency,
             )
             if not toc_with_page_number.get('structure'):
                 raise ValueError("PageIndex Flash could not extract a structure from this PDF; "
diff --git a/tests/test_client.py b/tests/test_client.py
index fd8d1ced1..b5bcb4708 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -2379,3 +2379,42 @@ def test_summary_max_words_reaches_the_local_indexer():
     assert PageIndexLocalClient()._api._summary_max_words is None
     with pytest.raises(PageIndexAPIError, match=r'index\["summary_max_words"\] must be a'):
         PageIndexLocalClient(index={"summary_max_words": "80"})
+
+
+def test_flash_knobs_reach_the_local_indexer(tmp_path, sample_pdf, monkeypatch):
+    """summary_concurrency, use_embedded_toc and optimize are settable on the
+    client, flat or in the index slot, and land on page_index_flash the way
+    the CLI's flags do ("off" is no optimize pass at all)."""
+    from pageindex import PageIndexLocalClient
+    captured = {}
+    monkeypatch.setattr(pageindex.flash, "page_index_flash",
+                        lambda p, **kw: captured.update(kw) or {
+                            "structure": [{"title": "T", "start_index": 1,
+                                           "end_index": 1, "summary": "s", "nodes": []}]})
+    monkeypatch.setattr(pageindex.utils, "llm_completion",
+                        lambda model, prompt, **kw: "d.")
+    stores = iter(range(10))
+
+    def run(**kwargs):
+        captured.clear()
+        store = str(tmp_path / str(next(stores)))
+        if "index" in kwargs:
+            kwargs["index"] = {**kwargs["index"], "storage_path": store}
+        else:
+            kwargs["storage_path"] = store
+        client = PageIndexLocalClient(**kwargs)
+        client.submit_document(sample_pdf)
+        return (captured.get("summary_concurrency"), captured["use_embedded_toc"],
+                captured["optimize"])
+
+    assert run(summary_concurrency=8, use_embedded_toc=False, optimize="merge") == (8, False, "merge")
+    assert run(index={"summary_concurrency": 8, "use_embedded_toc": False,
+                      "optimize": "off"}) == (8, False, False)
+    assert run() == (None, True, "full")
+    for kwargs, msg in (({"index": {"use_embedded_toc": "no"}},
+                         r'index\["use_embedded_toc"\] must be a bool'),
+                        ({"optimize": "sometimes"},
+                         r'optimize must be "full", "merge" or "off"'),
+                        ({"summary_concurrency": "8"}, "summary_concurrency must be a")):
+        with pytest.raises(PageIndexAPIError, match=msg):
+            PageIndexLocalClient(**kwargs)
diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py
index 28c0a963e..5d05e2c2d 100644
--- a/tests/test_flash_extraction.py
+++ b/tests/test_flash_extraction.py
@@ -467,3 +467,9 @@ def test_flash_cli_rejects_empty_structure(monkeypatch, tmp_path):
     with pytest.raises(ValueError, match="try --mode standard"):
         _run_flash_cli(monkeypatch, tmp_path, [], [])
     assert not (tmp_path / "results").exists()
+
+
+def test_flash_cli_summary_concurrency_reaches_the_indexer(monkeypatch, tmp_path):
+    captured = _run_flash_cli(monkeypatch, tmp_path, ["--summary-concurrency", "8"],
+                              [{"title": "A", "start_index": 1, "end_index": 1}])
+    assert captured["summary_concurrency"] == 8

From d957e8c58cd7edf1b5e496a4b3d152d387f90448 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 18:42:01 +0800
Subject: [PATCH 14/18] feat: summary_concurrency caps expand too

Lowering summary_concurrency for a tight quota used to lower only the
summary lane: expand kept its own 32 and still hit the API in bursts.
The one knob now bounds both lanes: expand's gate is min(32, the cap),
so 8 means 8 in each lane (the lanes overlap, so up to cap + min(32, cap)
calls are in flight), while raising it past 32 leaves expand at its
measured plateau. Defaults are unchanged (64 and 32).
---
 pageindex/flash/api.py         | 18 +++++++++++-------
 pageindex/tree_optimize.py     |  9 ++++++---
 tests/test_client.py           | 34 ++++++++++++++++++++++++++++++++++
 tests/test_flash_extraction.py |  6 +++---
 4 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index ce82458a5..f8f70429b 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -74,7 +74,8 @@ async def _summarize(structure, page_list, model, concurrency=None, max_words=No
                          max_words=max_words)
 
 
-async def _optimize_async(structure, page_texts, do_expand, model, on_final=None):
+async def _optimize_async(structure, page_texts, do_expand, model, on_final=None,
+                          concurrency=None):
     """Merge/expand refinement after extraction, overlapped with the summaries
     when `on_final` is passed; without it the caller runs them after.
 
@@ -88,7 +89,7 @@ async def _optimize_async(structure, page_texts, do_expand, model, on_final=None
              for page_text in page_texts]
     outcome = await optimize(structure, page_texts, lines, model=model,
                              do_expand=do_expand, page_count=len(page_texts),
-                             on_final=on_final)
+                             on_final=on_final, concurrency=concurrency)
     return {"merges": outcome["merges"], "expands": outcome["expands"],
             "same_page_merges": outcome["same_page_merges"],
             "same_page_dropped": outcome["same_page_dropped"],
@@ -96,9 +97,10 @@ async def _optimize_async(structure, page_texts, do_expand, model, on_final=None
             "before": outcome["before"], "after": outcome["after"]}
 
 
-def _optimize(structure, page_texts, do_expand, model):
+def _optimize(structure, page_texts, do_expand, model, concurrency=None):
     import asyncio
-    return asyncio.run(_optimize_async(structure, page_texts, do_expand, model))
+    return asyncio.run(_optimize_async(structure, page_texts, do_expand, model,
+                                       concurrency=concurrency))
 
 
 async def _optimize_and_summarize(structure, page_texts, optimize_model, summary_model,
@@ -110,7 +112,8 @@ async def _optimize_and_summarize(structure, page_texts, optimize_model, summary
                                  model=summary_model, concurrency=concurrency,
                                  max_words=max_words)
     report = await _optimize_async(structure, page_texts, True, optimize_model,
-                                   on_final=scheduler.mark_final)
+                                   on_final=scheduler.mark_final,
+                                   concurrency=concurrency)
     await scheduler.finish()
     return report
 
@@ -119,7 +122,7 @@ def page_index_flash(pdf, summary=True, summary_model=None,
                      optimize: str | bool | None = None, optimize_expand=None,
                      optimize_model=None, summary_concurrency=None,
                      use_embedded_toc=True, summary_max_words=None) -> dict:
-    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: cap on simultaneous indexing model calls per lane: the summaries, and expand up to its own ceiling of 32; None uses the library defaults (64 and 32). summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
     if optimize_expand is not None:
         import warnings
         warnings.warn(
@@ -155,7 +158,8 @@ def page_index_flash(pdf, summary=True, summary_model=None,
         return result
     if optimize and structure:
         result["optimize"] = _optimize(structure, pages, do_expand,
-                                       optimize_model or summary_model)
+                                       optimize_model or summary_model,
+                                       concurrency=summary_concurrency)
     if summary and structure:
         import asyncio
         page_list = [(text, 0) for text in pages]
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index 2013b6c87..4c7df992a 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -661,7 +661,7 @@ async def expand(structure, pages, lines, args, log, frozen):
     priced with expand_cost and the cheapest is kept.
     """
     changed = False
-    semaphore = asyncio.Semaphore(EXPAND_CONCURRENCY)
+    semaphore = asyncio.Semaphore(args.concurrency)
 
     async def proposals_for(node):
         """The model half of one node's lookahead: the empty-retry ladder and
@@ -790,7 +790,8 @@ async def optimize(structure, pages, lines, model=None, routing=ROUTING_COST,
                    trigger_pages=TRIGGER_PAGES, min_gain_ratio=0.0,
                    do_merge=True, do_expand=True, max_rounds=3, page_count=None,
                    cache=None, kinds=("section", "table"), empty_retries=1,
-                   do_relabel=True, progress=False, on_final=None):
+                   do_relabel=True, progress=False, on_final=None,
+                   concurrency=None):
     """Run merge and expand over a tree until neither changes anything.
 
     Mutates `structure` in place and returns a summary.
@@ -817,7 +818,9 @@ def settled(nodes):
                            min_gain_ratio=min_gain_ratio, cache=cache,
                            kinds=set(kinds) if kinds else None,
                            empty_retries=empty_retries, progress=progress,
-                           settled=settled, do_merge=do_merge)
+                           settled=settled, do_merge=do_merge,
+                           concurrency=min(EXPAND_CONCURRENCY,
+                                           concurrency or EXPAND_CONCURRENCY))
     baseline = set(validate(structure, page_count)) if page_count else set()
     before = complexity(structure, page_count, routing=routing) if page_count else {}
 
diff --git a/tests/test_client.py b/tests/test_client.py
index b5bcb4708..dbf51f1a1 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -2418,3 +2418,37 @@ def run(**kwargs):
                         ({"summary_concurrency": "8"}, "summary_concurrency must be a")):
         with pytest.raises(PageIndexAPIError, match=msg):
             PageIndexLocalClient(**kwargs)
+
+
+def test_summary_concurrency_caps_expand_too(tmp_path, monkeypatch):
+    """A user who lowers summary_concurrency for a tight quota gets the whole
+    indexing lane lowered: expand's own gate takes the same cap."""
+    from conftest import build_pdf
+    import pageindex.flash.api as flash_api
+    import pageindex.tree_optimize as tree_optimize
+
+    body = "body " * 250
+    pages = [body] * 30
+    roots = [{"title": f"X{i}", "start_index": 1 + 10 * i, "end_index": 10 + 10 * i,
+              "node_id": f"000{i}"} for i in range(3)]
+    pdf = tmp_path / "doc.pdf"
+    pdf.write_bytes(build_pdf(["x"]))
+    monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
+        "structure": roots, "page_texts": list(pages)})
+    in_flight, peak = 0, 0
+
+    async def propose(model, prompt):
+        nonlocal in_flight, peak
+        in_flight += 1
+        peak = max(peak, in_flight)
+        await asyncio.sleep(0.02)
+        in_flight -= 1
+        return {"subsections": []}
+    monkeypatch.setattr(tree_optimize, "ask_model", propose)
+
+    async def summarize(model, prompt):
+        return '{"summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", summarize)
+
+    flash_api.page_index_flash(str(pdf), summary_model="m", summary_concurrency=1)
+    assert peak == 1
diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py
index 5d05e2c2d..0d13d7115 100644
--- a/tests/test_flash_extraction.py
+++ b/tests/test_flash_extraction.py
@@ -237,7 +237,7 @@ def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch):
     from pageindex.flash import api as flash_api
     seen = {}
 
-    def fake_optimize(structure, pages, do_expand, model):
+    def fake_optimize(structure, pages, do_expand, model, concurrency=None):
         seen["do_expand"] = do_expand
         return {"merges": 0}
 
@@ -356,7 +356,7 @@ def test_optimize_full_skips_expand_without_page_texts(tmp_path, monkeypatch):
     monkeypatch.setenv("OPENAI_API_KEY", "k")
     calls = {}
 
-    def fake_optimize(structure, pages, do_expand, model):
+    def fake_optimize(structure, pages, do_expand, model, concurrency=None):
         calls["pages"] = pages
         calls["do_expand"] = do_expand
         return {"merges": 0}
@@ -382,7 +382,7 @@ def test_optimize_full_skips_expand_on_textless_pages(tmp_path, monkeypatch):
     monkeypatch.setenv("OPENAI_API_KEY", "k")
     calls = {}
 
-    def fake_optimize(structure, pages, do_expand, model):
+    def fake_optimize(structure, pages, do_expand, model, concurrency=None):
         calls["do_expand"] = do_expand
         return {"merges": 0}
 

From 35360f54fec2eb1f9b611f57f27d133af1d0913d Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 18:45:10 +0800
Subject: [PATCH 15/18] test: the CLI's --summary-max-words reaches the indexer

Belongs with 06f9c87; the file lost the test on disk between the test
run and that commit.
---
 tests/test_flash_extraction.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py
index 0d13d7115..c1e423e53 100644
--- a/tests/test_flash_extraction.py
+++ b/tests/test_flash_extraction.py
@@ -473,3 +473,9 @@ def test_flash_cli_summary_concurrency_reaches_the_indexer(monkeypatch, tmp_path
     captured = _run_flash_cli(monkeypatch, tmp_path, ["--summary-concurrency", "8"],
                               [{"title": "A", "start_index": 1, "end_index": 1}])
     assert captured["summary_concurrency"] == 8
+
+
+def test_flash_cli_summary_max_words_reaches_the_indexer(monkeypatch, tmp_path):
+    captured = _run_flash_cli(monkeypatch, tmp_path, ["--summary-max-words", "80"],
+                              [{"title": "A", "start_index": 1, "end_index": 1}])
+    assert captured["summary_max_words"] == 80

From d4b8a2eea0ad187c79e16bd512d32961e0af11fd Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 19:22:14 +0800
Subject: [PATCH 16/18] refactor: review follow-ups for the summary knobs

- max_words goes after concurrency in summarize_tree and
  SummaryScheduler: concurrency shipped in 0.2.11 as the sixth
  positional of summarize_tree, so inserting before it would have turned
  a positional concurrency into a word cap.
- _resolve_index_slot maps slot keys with a two-entry rename instead of a
  table that repeated every local key; conf is annotated dict[str, Any]
  so the comprehension types cleanly.
- run_pageindex.py checks the flash-only flags in one loop, like the
  standard-only ones.
- summary_concurrency's docs say what actually runs: per lane, and the
  lanes overlap, so up to cap + min(32, cap) calls at once.
- test_summary_concurrency_caps_expand_too keeps an uncapped control run
  (peak 3) so the capped assertion cannot pass by accident; the knob test
  uses the default store under tmp_path instead of per-run directories.
- page_index_flash's docstring lists summary_max_words in signature order.
---
 pageindex/client.py    | 20 +++++++-------------
 pageindex/flash/api.py |  2 +-
 pageindex/utils.py     | 10 +++++-----
 run_pageindex.py       | 19 ++++++++-----------
 tests/test_client.py   | 21 ++++++++++-----------
 5 files changed, 31 insertions(+), 41 deletions(-)

diff --git a/pageindex/client.py b/pageindex/client.py
index 0df387c80..da326e459 100644
--- a/pageindex/client.py
+++ b/pageindex/client.py
@@ -131,8 +131,8 @@ def _resolve_index_slot(index) -> "tuple[_CloudKey, dict[str, Any]]":
             'or "cloud".')
     if isinstance(index, Mapping):
         # None-valued keys mean "absent", exactly like the flat arguments.
-        conf = {name: value for name, value in index.items()
-                if value is not None}
+        conf: dict[str, Any] = {name: value for name, value in index.items()
+                                if value is not None}
         declared = _declared_mode(conf.pop("mode", None), "index")
         if not conf:
             if declared == "cloud":
@@ -170,16 +170,8 @@ def _resolve_index_slot(index) -> "tuple[_CloudKey, dict[str, Any]]":
                 'index declares mode "cloud" but carries local keys '
                 f"({', '.join(sorted(conf))}) — the cloud pipeline does "
                 'its own indexing; cloud takes "api_key" only.')
-        mapped = {"index_model": conf.get("model"),
-                  "summary_model": conf.get("summary_model"),
-                  "index_backend": conf.get("backend"),
-                  "storage_path": conf.get("storage_path"),
-                  "summary_max_words": conf.get("summary_max_words"),
-                  "summary_concurrency": conf.get("summary_concurrency"),
-                  "use_embedded_toc": conf.get("use_embedded_toc"),
-                  "optimize": conf.get("optimize")}
-        return None, {name: value for name, value in mapped.items()
-                      if value is not None}
+        rename = {"model": "index_model", "backend": "index_backend"}
+        return None, {rename.get(name, name): value for name, value in conf.items()}
     raise PageIndexAPIError("index must be a string or a dict.")
 
 
@@ -318,7 +310,9 @@ class PageIndexClient:
         summary_max_words (int, optional): Local mode only — the word cap
             each node summary is asked to stay within. Defaults to 150.
         summary_concurrency (int, optional): Local mode only — cap on
-            simultaneous indexing model calls. Defaults to 64.
+            simultaneous indexing model calls per lane: the summaries, and
+            expand up to its own ceiling of 32. The lanes overlap, so up to
+            cap + min(32, cap) calls run at once. Defaults to 64.
         use_embedded_toc (bool, optional): Local mode only — whether flash
             indexing consumes the PDF's embedded bookmarks when they look
             trustworthy. Defaults to True.
diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index f8f70429b..b7b5520f0 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -122,7 +122,7 @@ def page_index_flash(pdf, summary=True, summary_model=None,
                      optimize: str | bool | None = None, optimize_expand=None,
                      optimize_model=None, summary_concurrency=None,
                      use_embedded_toc=True, summary_max_words=None) -> dict:
-    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: cap on simultaneous indexing model calls per lane: the summaries, and expand up to its own ceiling of 32; None uses the library defaults (64 and 32). summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: cap on simultaneous indexing model calls per lane: the summaries, and expand up to its own ceiling of 32; None uses the library defaults (64 and 32). use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
     if optimize_expand is not None:
         import warnings
         warnings.warn(
diff --git a/pageindex/utils.py b/pageindex/utils.py
index 08e930bd6..e43e3bed4 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -885,8 +885,8 @@ class SummaryScheduler:
 
     def __init__(self, structure, pdf_pages, model=None,
                  small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
-                 max_intro_pages=SUMMARY_INTRO_MAX_PAGES, max_words=None,
-                 concurrency=None):
+                 max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None,
+                 max_words=None):
         self.structure = structure
         self._pdf_pages = pdf_pages
         self._model = model
@@ -1062,8 +1062,8 @@ def _any_summary(nodes):
 
 async def summarize_tree(structure, pdf_pages, model=None,
                          small_node_tokens=SUMMARY_RAW_TEXT_TOKENS,
-                         max_intro_pages=SUMMARY_INTRO_MAX_PAGES, max_words=None,
-                         concurrency=None):
+                         max_intro_pages=SUMMARY_INTRO_MAX_PAGES, concurrency=None,
+                         max_words=None):
     """Bottom-up summaries: leaves from their own pages, parents composed from
     child summaries plus the pages no child covers. A parent's summary describes
     its whole subtree (end_index union semantics). Nodes that already carry a
@@ -1075,7 +1075,7 @@ async def summarize_tree(structure, pdf_pages, model=None,
     scheduler = SummaryScheduler(structure, pdf_pages, model=model,
                                  small_node_tokens=small_node_tokens,
                                  max_intro_pages=max_intro_pages,
-                                 max_words=max_words, concurrency=concurrency)
+                                 concurrency=concurrency, max_words=max_words)
     scheduler.mark_final(list(_subtree(structure)))
     return await scheduler.finish()
 
diff --git a/run_pageindex.py b/run_pageindex.py
index 716ad4e4c..bf4b386a4 100644
--- a/run_pageindex.py
+++ b/run_pageindex.py
@@ -36,7 +36,7 @@
     parser.add_argument('--summary-max-words', type=int, default=None,
                       help='Word cap for each node summary (flash mode; default 150)')
     parser.add_argument('--summary-concurrency', type=int, default=None,
-                      help='Cap on simultaneous indexing model calls (flash mode; default 64)')
+                      help='Cap on simultaneous indexing model calls per lane (flash mode; default 64, expand tops out at 32)')
 
     parser.add_argument('--toc-check-pages', type=int, default=None,
                       help='Number of pages to check for table of contents (PDF only)')
@@ -70,18 +70,15 @@
         raise ValueError("Either --pdf_path or --md_path must be specified")
     if args.pdf_path and args.md_path:
         raise ValueError("Only one of --pdf_path or --md_path can be specified")
-    if args.optimize is not None and not (args.pdf_path and args.mode == 'flash'):
-        raise ValueError("--optimize requires Flash mode with --pdf_path")
+    for flag, value in (('--optimize', args.optimize),
+                        ('--embedded-toc', args.embedded_toc),
+                        ('--summary', args.summary),
+                        ('--summary-max-words', args.summary_max_words),
+                        ('--summary-concurrency', args.summary_concurrency)):
+        if value is not None and not (args.pdf_path and args.mode == 'flash'):
+            raise ValueError(f"{flag} requires Flash mode with --pdf_path")
     if args.optimize is None:
         args.optimize = 'full' if args.mode == 'flash' else 'off'
-    if args.embedded_toc is not None and not (args.pdf_path and args.mode == 'flash'):
-        raise ValueError("--embedded-toc requires Flash mode with --pdf_path")
-    if args.summary is not None and not (args.pdf_path and args.mode == 'flash'):
-        raise ValueError("--summary requires Flash mode with --pdf_path")
-    if args.summary_max_words is not None and not (args.pdf_path and args.mode == 'flash'):
-        raise ValueError("--summary-max-words requires Flash mode with --pdf_path")
-    if args.summary_concurrency is not None and not (args.pdf_path and args.mode == 'flash'):
-        raise ValueError("--summary-concurrency requires Flash mode with --pdf_path")
     if args.pdf_path and args.mode == 'flash':
         for flag, value in (('--toc-check-pages', args.toc_check_pages),
                             ('--max-pages-per-node', args.max_pages_per_node),
diff --git a/tests/test_client.py b/tests/test_client.py
index dbf51f1a1..bd29f7349 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -2393,17 +2393,11 @@ def test_flash_knobs_reach_the_local_indexer(tmp_path, sample_pdf, monkeypatch):
                                            "end_index": 1, "summary": "s", "nodes": []}]})
     monkeypatch.setattr(pageindex.utils, "llm_completion",
                         lambda model, prompt, **kw: "d.")
-    stores = iter(range(10))
+    monkeypatch.chdir(tmp_path)  # the default .pageindex store lands here
 
     def run(**kwargs):
         captured.clear()
-        store = str(tmp_path / str(next(stores)))
-        if "index" in kwargs:
-            kwargs["index"] = {**kwargs["index"], "storage_path": store}
-        else:
-            kwargs["storage_path"] = store
-        client = PageIndexLocalClient(**kwargs)
-        client.submit_document(sample_pdf)
+        PageIndexLocalClient(**kwargs).submit_document(sample_pdf)
         return (captured.get("summary_concurrency"), captured["use_embedded_toc"],
                 captured["optimize"])
 
@@ -2429,12 +2423,14 @@ def test_summary_concurrency_caps_expand_too(tmp_path, monkeypatch):
 
     body = "body " * 250
     pages = [body] * 30
-    roots = [{"title": f"X{i}", "start_index": 1 + 10 * i, "end_index": 10 + 10 * i,
-              "node_id": f"000{i}"} for i in range(3)]
+
+    def roots():
+        return [{"title": f"X{i}", "start_index": 1 + 10 * i, "end_index": 10 + 10 * i,
+                 "node_id": f"000{i}"} for i in range(3)]
     pdf = tmp_path / "doc.pdf"
     pdf.write_bytes(build_pdf(["x"]))
     monkeypatch.setattr(flash_api, "extract_toc", lambda pdf, **kw: {
-        "structure": roots, "page_texts": list(pages)})
+        "structure": roots(), "page_texts": list(pages)})
     in_flight, peak = 0, 0
 
     async def propose(model, prompt):
@@ -2452,3 +2448,6 @@ async def summarize(model, prompt):
 
     flash_api.page_index_flash(str(pdf), summary_model="m", summary_concurrency=1)
     assert peak == 1
+    peak = 0
+    flash_api.page_index_flash(str(pdf), summary_model="m")  # control: the three overlap
+    assert peak == 3

From a8f0ccf8e220e93c9a7b12adbf5b3edc4a033a88 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Thu, 27 Aug 2026 20:48:56 +0800
Subject: [PATCH 17/18] fix: count_tokens falls back to the default tokenizer;
 expand's settle order is tested

- count_tokens(model=...) picks the model's tokenizer through litellm, and
  the HuggingFace tokenizers it selects for llama-family names reject a
  lone surrogate with TypeError where tiktoken counts it. _leaf_summary's
  raw-text gate runs inside _visit's except, so that raise blanked the
  leaf's summary silently. On any failure count_tokens now re-counts with
  litellm's bundled default encoding (cl100k: offline, surrogate-safe).
- get_page_tokens (both parsers) and _index_standard route through
  count_tokens instead of calling litellm.token_counter themselves, so the
  standard path shares the fallback; a page whose extract_text() is None
  now counts 0 instead of raising.
- expand settles a node only after merge_same_page has fused its new
  children: mark_final snapshots the children and finish() rejects a later
  change. The order was right but nothing held it; the fusion test now
  drives a real SummaryScheduler through on_final and awaits finish(), so
  swapping the two lines fails it.
---
 pageindex/local_api.py     |  6 ++----
 pageindex/tree_optimize.py |  1 +
 pageindex/utils.py         | 10 ++++++----
 tests/test_client.py       | 36 ++++++++++++++++++++++++++++++++----
 4 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/pageindex/local_api.py b/pageindex/local_api.py
index 5e8616293..463909eac 100644
--- a/pageindex/local_api.py
+++ b/pageindex/local_api.py
@@ -12,7 +12,7 @@
 
 from .errors import PageIndexAPIError
 from .local_store import DocStore
-from .utils import run_off_loop
+from .utils import count_tokens, run_off_loop
 
 logger = logging.getLogger(__name__)
 
@@ -214,9 +214,7 @@ def _extract_page_texts(file_path: str) -> list[str]:
 
     def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]:
         from .page_index_classic import page_index_main
-        import litellm
-        page_list = [(text, litellm.token_counter(model=self._model, text=text))
-                     for text in page_texts]
+        page_list = [(text, count_tokens(text, model=self._model)) for text in page_texts]
         opt = self._config_loader.load({
             "model": self._model,
             "summary_model": self._summary_model,
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index 4c7df992a..ddf397c7d 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -751,6 +751,7 @@ async def process(node):
         attach_children(node, best["children"], lines)
         if args.do_merge:
             merge_same_page([node], log)
+        # settle after the fusion: mark_final snapshots the children, finish() rejects a later change
         args.settled([node])
         results = await asyncio.gather(*(process(child)
                                          for child in node["nodes"]),
diff --git a/pageindex/utils.py b/pageindex/utils.py
index e43e3bed4..30799762a 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -71,7 +71,10 @@ def count_tokens(text, model=None):
     if not text:
         return 0
     import litellm
-    return litellm.token_counter(model=model, text=text)
+    try:
+        return litellm.token_counter(model=model, text=text)
+    except Exception:
+        return litellm.token_counter(model=None, text=text)
 
 
 def _strip_prefix(s, prefix):
@@ -515,14 +518,13 @@ def add_preface_if_needed(data):
 
 
 def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"):
-    import litellm
     if pdf_parser == "PyPDF2":
         pdf_reader = PyPDF2.PdfReader(pdf_path)
         page_list = []
         for page_num in range(len(pdf_reader.pages)):
             page = pdf_reader.pages[page_num]
             page_text = page.extract_text()
-            token_length = litellm.token_counter(model=model, text=page_text)
+            token_length = count_tokens(page_text, model=model)
             page_list.append((page_text, token_length))
         return page_list
     elif pdf_parser == "PyMuPDF":
@@ -535,7 +537,7 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"):
         page_list = []
         for page in doc:
             page_text = page.get_text()
-            token_length = litellm.token_counter(model=model, text=page_text)
+            token_length = count_tokens(page_text, model=model)
             page_list.append((page_text, token_length))
         return page_list
     else:
diff --git a/tests/test_client.py b/tests/test_client.py
index bd29f7349..9eacbe220 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -2128,11 +2128,25 @@ async def propose(model, prompt):
         return {"subsections": [{"title": "Sub One", "page": 4}, {"title": "Sub Two", "page": 4},
                                 {"title": "Sub Three", "page": 5}, {"title": "Sub Four", "page": 9}]}
     monkeypatch.setattr(tree_optimize, "ask_model", propose)
+
+    async def summarize(model, prompt):
+        return '{"summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", summarize)
     reports = []
-    outcome = asyncio.run(tree_optimize.optimize(
-        tree, pages, lines, model="m", do_expand=True, max_rounds=1,
-        on_final=lambda nodes: reports.append([(n["title"], len(n.get("nodes") or []))
-                                              for n in nodes])))
+
+    async def run():
+        # the real scheduler: mark_final snapshots each node's children and
+        # finish() rejects a later change, so settling before the fusion fails here
+        scheduler = pageindex.utils.SummaryScheduler(tree, [(p, 0) for p in pages], model="m")
+
+        def on_final(nodes):
+            reports.append([(n["title"], len(n.get("nodes") or [])) for n in nodes])
+            scheduler.mark_final(nodes)
+        outcome = await tree_optimize.optimize(tree, pages, lines, model="m", do_expand=True,
+                                               max_rounds=1, on_final=on_final)
+        await scheduler.finish()
+        return outcome
+    outcome = asyncio.run(run())
     assert outcome["expands"] == 1 and outcome["same_page_merges"] == 1
     assert [n["title"] for n in tree[0]["nodes"][1]["nodes"]][1:] == ["Sub Three", "Sub Four"]
     assert all(children != 4 for report in reports for _, children in report)
@@ -2451,3 +2465,17 @@ async def summarize(model, prompt):
     peak = 0
     flash_api.page_index_flash(str(pdf), summary_model="m")  # control: the three overlap
     assert peak == 3
+
+
+def test_count_tokens_falls_back_to_the_default_tokenizer(monkeypatch):
+    import litellm
+
+    def token_counter(model=None, text=None, **_):
+        if model is not None:
+            raise TypeError("TextInputSequence must be str")  # HF tokenizer on a lone surrogate
+        return 7
+    monkeypatch.setattr(litellm, "token_counter", token_counter)
+    assert pageindex.utils.count_tokens("x", model="groq/llama-3.1-8b-instant") == 7
+
+    monkeypatch.setattr(litellm, "token_counter", lambda model=None, text=None, **_: 3 if model else 7)
+    assert pageindex.utils.count_tokens("x", model="m") == 3  # the model's own count wins when it works

From 32b3b7510e48f2e5e77696945f660f06fe39b1ba Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Fri, 28 Aug 2026 00:21:02 +0800
Subject: [PATCH 18/18] docs: page_index_flash states the overlap peak for
 summary_concurrency

---
 pageindex/flash/api.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py
index b7b5520f0..1399b43b6 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -122,7 +122,7 @@ def page_index_flash(pdf, summary=True, summary_model=None,
                      optimize: str | bool | None = None, optimize_expand=None,
                      optimize_model=None, summary_concurrency=None,
                      use_embedded_toc=True, summary_max_words=None) -> dict:
-    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: cap on simultaneous indexing model calls per lane: the summaries, and expand up to its own ceiling of 32; None uses the library defaults (64 and 32). use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (a model unreachable after the retry ladder — a missing credential included — fails the run loudly from expand itself; a per-prompt rejection leaves just that node collapsed), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. Expand needs readable page text, so a bookmark-only or scanned PDF runs the merge half only (``expands`` reports 0). optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: cap on simultaneous indexing model calls per lane: the summaries, and expand up to its own ceiling of 32 (the lanes overlap, so up to cap + min(32, cap) calls run at once); None uses the library defaults (64 and 32). use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. summary_max_words: word cap each node summary is asked to stay within; None uses the library default (150). Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
     if optimize_expand is not None:
         import warnings
         warnings.warn(