Keep ready study work moving around unresolved model calls

research/experiments/study_context.py

@@ -63,6 +63,7 @@ class _Node:
     children: list = field(default_factory=list)
     group: Any = None
     characters: int = 0
+    canonical_id: int = 0
 
 
 @dataclass
@@ -73,19 +74,31 @@ class _Group:
 
 
 def _source_tree(value):
-    # The complete literal serialization is the equality key. No content hash,
-    # shortened fingerprint or source identity is created or substituted.
-    groups_by_literal = {}
+    # Intern exact structural identities bottom-up. Re-serializing each nested
+    # subtree as an equality key made deeply nested evidence quadratic in input
+    # size before transport. Child IDs preserve exact JSON types and order while
+    # keeping each parent key proportional to its immediate children.
+    canonical_ids = {}
+    groups_by_identity = {}
     groups = []
     active = set()
 
+    def canonical_id(identity):
+        identifier = canonical_ids.get(identity)
+        if identifier is None:
+            identifier = len(canonical_ids) + 1
+            canonical_ids[identity] = identifier
+        return identifier
+
     def visit(current, path):
         if type(current) not in (dict, list):
             if type(current) not in (str, int, float, bool, type(None)):
                 raise TypeError("Context requires ordinary JSON values at " + (path or "/"))
             if isinstance(current, float) and not math.isfinite(current):
                 raise ValueError("Context requires finite JSON numbers at " + (path or "/"))
-            return _Node("scalar", value=current, characters=len(_dumps(current)))
+            literal = _dumps(current)
+            return _Node("scalar", value=current, characters=len(literal),
+                         canonical_id=canonical_id(("scalar", type(current).__name__, literal)))
         if id(current) in active:
             raise ValueError("Cyclic source is not a JSON tree at " + (path or "/"))
         if isinstance(current, dict) and any(type(key) is not str for key in current):
@@ -102,11 +115,15 @@ def visit(current, path):
             node.children = [visit(child, _pointer(path, index)) for index, child in enumerate(current)]
             node.characters = len('["array",[]]') + sum(child.characters for child in node.children) + max(0, len(node.children)-1)
         active.remove(id(current))
-        literal = _dumps(current)
-        group = groups_by_literal.get(literal)
+        if node.kind == "object":
+            identity = ("object", tuple((key, child.canonical_id) for key, child in node.children))
+        else:
+            identity = ("array", tuple(child.canonical_id for child in node.children))
+        node.canonical_id = canonical_id(identity)
+        group = groups_by_identity.get(node.canonical_id)
         if group is None:
             group = _Group(len(groups), node)
-            groups_by_literal[literal] = group
+            groups_by_identity[node.canonical_id] = group
             groups.append(group)
         else:
             # Literal equality includes scalar types and all original order.

research/experiments/study_runtime.py

@@ -545,7 +545,10 @@ def claim(self, campaign_id: str, owner: str, *, work_ids: Sequence[str] | None
         now = time.time()
         with self.connect() as db:
             db.execute("BEGIN IMMEDIATE")
-            query = "SELECT * FROM forest_study_work WHERE campaign_id=? AND state IN ('queued','retryable')"
+            # Candidate queues may contain large durable research histories.
+            # Sort and check dependencies using compact metadata, then hydrate
+            # only the single row that this claim actually wins.
+            query = "SELECT work_id,kind,created_at FROM forest_study_work WHERE campaign_id=? AND state IN ('queued','retryable')"
             params: list[Any] = [campaign_id]
             if work_ids is not None:
                 if not work_ids:
@@ -581,7 +584,11 @@ def claim(self, campaign_id: str, owner: str, *, work_ids: Sequence[str] | None
                     db.execute("INSERT INTO forest_study_events VALUES(?,?,?,?,?,?)",
                         (str(uuid.uuid4()), campaign_id, row["work_id"], "work_claimed",
                          _json({"owner": owner, "kind": row["kind"]}), now))
-                    value = dict(row)
+                    full_row = db.execute(
+                        "SELECT * FROM forest_study_work WHERE campaign_id=? AND work_id=?",
+                        (campaign_id, row["work_id"]),
+                    ).fetchone()
+                    value = dict(full_row)
                     value["state"] = "claimed"
                     value["lease_owner"] = owner
                     value["lease_at"] = now

research/experiments/study_scheduler.py

@@ -432,6 +432,7 @@ def __init__(self, runtime, *, scheduler_id=None, maximum_model_concurrency=2,
         self._closed = False
         self._futures = {}
         self._wave = None
+        self._unsettled_headroom_active = False
         self._observed_samples = []
         self._tested_degrees = set()
         self._pending_decisions = []
@@ -717,6 +718,97 @@ def _unsettled_model_owners(self):
                 (self.campaign_id,))}
         return sorted(owners - active)
 
+    def _external_codex_cli_process_count(self):
+        """Count Codex CLI requests surviving outside this coordinator process tree."""
+        try:
+            import psutil
+            owned = {os.getpid()}
+            try:
+                owned.update(child.pid for child in psutil.Process().children(recursive=True))
+            except (psutil.Error, OSError):
+                pass
+            count = 0
+            for process in psutil.process_iter(['pid', 'cmdline']):
+                try:
+                    if process.info['pid'] in owned:
+                        continue
+                    argv = process.info.get('cmdline') or []
+                    executable = Path(argv[0]).name.lower() if argv else ''
+                    if 'codex' in executable and 'exec' in argv[1:] and '--json' in argv:
+                        count += 1
+                except (psutil.Error, OSError):
+                    continue
+            return count
+        except ImportError:
+            # Without process visibility, retain the original conservative
+            # reservation and do not infer that a provider request has ended.
+            return None
+
+    def _capacity_occupying_unsettled_model_owners(self, owners=None):
+        """Keep ambiguous receipts, but do not reserve slots for dead CLI work.
+
+        An unresolved request remains attached to its original work item and
+        supervisor receipt. For local Codex CLI only, a transport trace that is
+        still merely ``prepared`` past its own request timeout cannot consume
+        a live concurrency slot when no matching external CLI process exists.
+        This never makes that work retryable or replays its decision.
+        """
+        owners = self._unsettled_model_owners() if owners is None else list(owners)
+        if not owners:
+            return []
+        default = self.runtime.manifest.get('model_provider') or {}
+        secondary = self.runtime.manifest.get('secondary_model_provider') or {}
+        profiles = [profile for profile in (default, secondary) if isinstance(profile, Mapping) and profile]
+        if not profiles or any(profile.get('kind') != 'codex_cli' for profile in profiles):
+            return owners
+
+        with self.ledger.connect() as db:
+            rows = [dict(row) for row in db.execute(
+                "SELECT request_id,work_id,started_at FROM forest_study_model_calls "
+                "WHERE campaign_id=? AND status='requested'", (self.campaign_id,))
+                if row['work_id'] in owners]
+        by_owner = {owner: [] for owner in owners}
+        for row in rows:
+            by_owner.setdefault(row['work_id'], []).append(row)
+
+        now = time.time()
+        releasable, retained = [], []
+        for owner in owners:
+            calls = by_owner.get(owner) or []
+            owner_releasable = bool(calls)
+            for call in calls:
+                try:
+                    request_id = str(uuid.UUID(call['request_id']))
+                except (TypeError, ValueError, AttributeError):
+                    owner_releasable = False
+                    break
+                trace_path = self.runtime.campaign / 'runs' / 'transport_traces' / (request_id + '.json')
+                try:
+                    trace = json.loads(trace_path.read_text(encoding='utf-8'))
+                except (OSError, ValueError, TypeError):
+                    owner_releasable = False
+                    break
+                state = trace.get('state')
+                if state in {'exited', 'timed_out', 'not_started'}:
+                    continue
+                prepared = trace.get('prepared_at')
+                timeout = trace.get('timeout_seconds')
+                if state != 'prepared' or not isinstance(prepared, (int, float)) or \
+                        not isinstance(timeout, (int, float)) or now < prepared + timeout:
+                    owner_releasable = False
+                    break
+            (releasable if owner_releasable else retained).append(owner)
+
+        if not releasable:
+            return owners
+        live_cli = self._external_codex_cli_process_count()
+        if live_cli is None:
+            return owners
+        # Unknown and non-Codex providers stay reserved. Of expired local CLI
+        # traces, retain only as many slots as live external CLI processes.
+        held_releasable = releasable[:min(len(releasable), live_cli)]
+        return sorted(set(retained).union(held_releasable))
+
     @staticmethod
     def _provider_failure(failure):
         from research.agents.provider import ProviderError
@@ -931,7 +1023,9 @@ def _wake_reconciled_parents(self):
                     'SELECT work_id,updated_at FROM forest_study_aide_sessions '
                     'INDEXED BY forest_study_aide_session_live_updated WHERE campaign_id=?',
                     (self.campaign_id,))}
-            states = {row['work_id']: row['state'] for row in db.execute('SELECT work_id,state FROM forest_study_work WHERE campaign_id=?', (self.campaign_id,))}
+            states = {row['work_id']: row['state'] for row in db.execute(
+                'SELECT work_id,state FROM forest_study_work INDEXED BY forest_study_work_live_state WHERE campaign_id=?',
+                (self.campaign_id,))}
             edges = db.execute("SELECT source_id,target_id FROM forest_study_edges WHERE campaign_id=? AND relation='derived_from'", (self.campaign_id,)).fetchall()
             awakened = []
             for row in rows:
@@ -985,7 +1079,9 @@ def _ready_ids(self, kinds, compute):
             # classification.  The full payload remains untouched for the
             # eventual execution claim.
             rows = db.execute("SELECT work_id,kind,state,supervisor_work_id FROM forest_study_work INDEXED BY forest_study_work_live_ready_refs WHERE campaign_id=? AND state IN ('queued','retryable') ORDER BY created_at,work_id", (self.campaign_id,)).fetchall()
-            states = {row['work_id']: row['state'] for row in db.execute('SELECT work_id,state FROM forest_study_work WHERE campaign_id=?', (self.campaign_id,))}
+            states = {row['work_id']: row['state'] for row in db.execute(
+                'SELECT work_id,state FROM forest_study_work INDEXED BY forest_study_work_live_state WHERE campaign_id=?',
+                (self.campaign_id,))}
             edges = db.execute('SELECT source_id,target_id,relation FROM forest_study_edges WHERE campaign_id=?', (self.campaign_id,)).fetchall()
             active = {entry['item']['work_id'] for entry in self._futures.values()}
             ready = []
@@ -1279,7 +1375,15 @@ def tick(self):
         occupancy = {lane: sum(entry['lane'] == lane for entry in self._futures.values()) for lane in self._executors}
         # Unknown legacy/provider subprocesses may survive coordinator death.
         # Reattached OS controllers already own a future and are excluded.
-        occupancy['model'] += len(self._unsettled_model_owners())
+        unsettled_models = self._unsettled_model_owners()
+        capacity_owners = self._capacity_occupying_unsettled_model_owners(unsettled_models)
+        capacity_released = sorted(set(unsettled_models) - set(capacity_owners))
+        occupancy['model'] += len(capacity_owners)
+        if capacity_released and getattr(self, '_last_capacity_released_owners', None) != capacity_released:
+            self._decision({'type': 'unsettled_model_capacity_released',
+                'work_ids': capacity_released,
+                'reason': 'Local Codex CLI transport exceeded its persisted request timeout and no matching process remains; original requests and reconciliation states are preserved'})
+            self._last_capacity_released_owners = capacity_released
         ready = {'gpu': self._ready_ids(GPU_KINDS, compute), 'cpu': self._ready_ids(CPU_KINDS, compute), 'model': self._ready_ids(MODEL_KINDS, compute)}
         for lane,ids in ready.items():
             self._observe_ready(lane, ids)
@@ -1289,7 +1393,26 @@ def tick(self):
             ids = ready['gpu']
             if ids:
                 self._dispatch('gpu', ids)
-        for lane,kinds,capacity in [('cpu', CPU_KINDS, self.cpu_degree), ('model', MODEL_KINDS, self.degree)]:
+        measured_headroom = max((int(sample['degree']) for sample in self._observed_samples
+            if isinstance(sample, Mapping) and isinstance(sample.get('degree'), int)
+            and not isinstance(sample.get('degree'), bool)), default=self.degree)
+        model_capacity = self.degree
+        if ready['model'] and capacity_owners and measured_headroom > self.degree:
+            # Keep an unresolved request reserved while using only separately
+            # measured model capacity for independent ready work.
+            model_capacity = min(self.maximum, measured_headroom,
+                max(self.degree, occupancy['model'] + 1))
+        headroom_active = model_capacity > self.degree
+        if headroom_active and not self._unsettled_headroom_active:
+            self._decision({'type': 'unsettled_model_capacity_headroom',
+                'unsettled_model_work_ids': capacity_owners,
+                'measured_model_concurrency': self.degree,
+                'effective_model_capacity': model_capacity,
+                'verified_capacity_ceiling': measured_headroom,
+                'configured_capacity_ceiling': self.maximum,
+                'reason': 'Keep independent ready work moving while an original model request remains unresolved'})
+        self._unsettled_headroom_active = headroom_active
+        for lane,kinds,capacity in [('cpu', CPU_KINDS, self.cpu_degree), ('model', MODEL_KINDS, model_capacity)]:
             available = capacity - occupancy[lane]
             wave_dispatches = []
             for _ in range(available):
@@ -1439,9 +1562,13 @@ def add(kind, start, end):
                 'completed': len(done), 'remaining': remaining,
                 'observed_wall_seconds': elapsed, 'completed_per_second': stage_rate,
                 'eta_remaining_seconds': remaining/stage_rate if stage_rate else None})
+        unsettled_model_owners = self._unsettled_model_owners()
+        capacity_owners = self._capacity_occupying_unsettled_model_owners(unsettled_model_owners)
         return {'campaign_id': self.campaign_id, 'scheduler_id': self.scheduler_id,
             'model_concurrency': self.degree, 'capacity': self._configuration(),
-            'unsettled_model_work_ids': self._unsettled_model_owners(),
+            'unsettled_model_work_ids': unsettled_model_owners,
+            'capacity_occupying_unsettled_model_work_ids': capacity_owners,
+            'capacity_released_unsettled_model_work_ids': sorted(set(unsettled_model_owners) - set(capacity_owners)),
             'categories': measured, 'categories_can_overlap': True,
             'observed_wall_union_seconds': interval_measure([interval for kind,value in categories.items() if kind != 'unreconciled_model_interval' for interval in value])['wall_union_seconds'],
             'work_states': states, 'remaining_ready': len(ready),

tests/test_study_context.py

@@ -79,6 +79,19 @@ def test_key_order_numeric_types_negative_zero_and_unicode_are_preserved():
     assert math.copysign(1, restored[0]["n"]) == -1
 
 
+def test_deep_repeated_subtrees_use_exact_bottom_up_identity_and_roundtrip():
+    nested = {"value": [1, 1.0, True, "x"], "empty": {}}
+    for index in range(180):
+        nested = {"level": index, "child": nested}
+    source = {"left": deepcopy(nested), "right": deepcopy(nested)}
+
+    document, stats = intern_json_with_stats(source)
+
+    assert document["encoding"] == "interned"
+    assert stats["saved_characters"] > 0
+    assert compact(restore_json(document)) == compact(source)
+
+
 def test_stable_reference_ids_follow_first_source_traversal():
     earlier = {"outer": list(range(200)), "name": "earlier"}
     later = {"outer": list(range(400, 600)), "name": "later"}

tests/test_study_runtime.py

@@ -77,6 +77,38 @@ def traced_connect():
     assert saved["retained_evidence"] == large_result["retained_evidence"]
 
 
+def test_claim_hydrates_only_the_selected_candidate_from_a_large_queue(tmp_path):
+    ledger = StudyLedger(tmp_path / "study.sqlite", "ledger-test")
+    ledger.register(manifest([
+        item("selected", kind="analysis"),
+        item("unselected-large", kind="replay"),
+    ]))
+    large_result = {"retained_evidence": "evidence" * 250_000}
+    ledger.set_state("campaign-test", "unselected-large", "retryable", result=large_result)
+
+    statements = []
+    original_connect = ledger.connect
+    def traced_connect():
+        db = original_connect()
+        db.set_trace_callback(statements.append)
+        return db
+    ledger.connect = traced_connect
+
+    claimed = ledger.claim("campaign-test", "single-claim",
+        work_ids=["selected", "unselected-large"])
+    assert claimed["work_id"] == "selected"
+    full_work_reads = [query.lower() for query in statements
+        if query.strip().lower().startswith("select * from forest_study_work")]
+    assert len(full_work_reads) == 1
+    assert "and work_id='selected'" in full_work_reads[0]
+    assert "state in" not in full_work_reads[0]
+    with original_connect() as db:
+        saved = _decode_study_result_json(db.execute(
+            "SELECT result_json FROM forest_study_work WHERE work_id='unselected-large'"
+        ).fetchone()[0])
+    assert saved == large_result
+
+
 def test_completed_action_id_conflict_returns_replan_feedback_without_reexecution(tmp_path, monkeypatch):
     campaign = tmp_path / "campaign"
     campaign.mkdir()