drydock-cli 3.1.30__tar.gz → 3.1.31__tar.gz

{drydock_cli-3.1.30/drydock_cli.egg-info → drydock_cli-3.1.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drydock-cli
-Version: 3.1.30
+Version: 3.1.31
 Summary: Drydock — a local, provider-agnostic terminal coding agent for local LLMs
 Author: Frank Bobe III
 License-Expression: Apache-2.0

drydock_cli-3.1.31/drydock/asr.py

@@ -0,0 +1,227 @@
+"""Adaptive Swarm Runtime (ASR) — the search-decision core.
+
+This is the resource-agnostic brain of the adaptive swarm (PRD "Resource-Aware Adaptive
+Swarms"): it turns a compute/time BUDGET into decisions about *where the next unit of inference
+should go* and *when to stop*. It deliberately holds no model, GPU, or scheduler state — those
+live in capacity.py (concurrency), swarm.py (blackboard/candidates/worktrees) and the mission
+store (durable state). Keeping the decisions pure makes them unit-testable and reusable.
+
+Three pieces, in dependency order:
+  * Budget      — a finite termination envelope parsed from a user string (§15).
+  * Convergence — should the swarm keep going? Stops on success / budget / no-improvement /
+                  when marginal value Δquality/Δcompute falls below a floor (§23).
+  * AdaptiveAllocator — split the next compute pool across live branches by expected value,
+                  pruning the weak ones (§17). Bad ideas must not get the same compute as good.
+"""
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+# ── Budget (§15) ────────────────────────────────────────────────────────────
+_DUR = re.compile(r"^\s*([\d.]+)\s*([hms])\s*$", re.I)
+_TOK = re.compile(r"^\s*([\d.]+)\s*([kmg]?)\s*[-_]?tokens?\s*$", re.I)
+_MULT = {"": 1, "k": 1_000, "m": 1_000_000, "g": 1_000_000_000}
+
+
+@dataclass
+class Budget:
+    """A finite stopping envelope (§10/§15). Every field is optional, but a Budget with no
+    finite limit is a bug — `is_finite` guards against "run forever"."""
+    wall_secs: float | None = None
+    tokens: int | None = None
+    generations: int | None = None
+
+    @property
+    def is_finite(self) -> bool:
+        return any(v is not None for v in (self.wall_secs, self.tokens, self.generations))
+
+    @classmethod
+    def parse(cls, spec: str) -> "Budget":
+        """Parse one user token: '2h' / '30m' / '90s' (wall) or '100M-tokens' / '50k-tokens'
+        (compute). Unknown strings raise ValueError so a typo never silently means 'unbounded'."""
+        s = (spec or "").strip()
+        md = _DUR.match(s)
+        if md:
+            n, unit = float(md.group(1)), md.group(2).lower()
+            return cls(wall_secs=n * {"h": 3600, "m": 60, "s": 1}[unit])
+        mt = _TOK.match(s)
+        if mt:
+            return cls(tokens=int(float(mt.group(1)) * _MULT[mt.group(2).lower()]))
+        raise ValueError(f"unparseable budget: {spec!r} (try '2h', '30m', '100M-tokens')")
+
+
+# ── Convergence (§23) ─────────────────────────────────────────────────────────
+@dataclass
+class GenerationRecord:
+    gen: int
+    best_score: float          # best branch quality so far (higher = better)
+    cumulative_tokens: int     # total inference tokens spent through this generation
+    elapsed_s: float           # wall-clock since start
+
+
+@dataclass
+class Convergence:
+    """Decides whether the swarm should run another generation (§23). "More agents are not
+    automatically better": we stop the moment any finite condition trips, INCLUDING the marginal
+    one — when Δquality per unit compute drops below `min_marginal`, extra inference is waste."""
+    budget: Budget
+    success_score: float = 100.0     # stop when best branch reaches this (e.g. all tests pass)
+    no_improve_gens: int = 3         # stop after this many generations with no best-score gain
+    min_marginal: float = 0.0        # floor on Δscore per 1k tokens; <=0 disables the check
+    eps: float = 1e-9
+    history: list[GenerationRecord] = field(default_factory=list)
+
+    def record(self, best_score: float, cumulative_tokens: int, elapsed_s: float) -> None:
+        self.history.append(GenerationRecord(len(self.history) + 1, best_score,
+                                             cumulative_tokens, elapsed_s))
+
+    def _gens_since_improvement(self) -> int:
+        """Generations since the best score was FIRST reached (a plateau at the best counts)."""
+        if not self.history:
+            return 0
+        best = max(r.best_score for r in self.history)
+        first_best = next(i for i, r in enumerate(self.history) if r.best_score >= best - self.eps)
+        return (len(self.history) - 1) - first_best
+
+    def marginal_value(self) -> float | None:
+        """Δscore per 1,000 tokens over the last generation, or None with <2 records."""
+        if len(self.history) < 2:
+            return None
+        a, b = self.history[-2], self.history[-1]
+        dtok = b.cumulative_tokens - a.cumulative_tokens
+        if dtok <= 0:
+            return None
+        return (b.best_score - a.best_score) / (dtok / 1000.0)
+
+    def should_continue(self) -> tuple[bool, str]:
+        """(keep_going?, reason). Checked in priority order; the reason names the stop trigger."""
+        last = self.history[-1] if self.history else None
+        if last is not None and last.best_score >= self.success_score - self.eps:
+            return False, "success score reached"
+        b = self.budget
+        if b.wall_secs is not None and last is not None and last.elapsed_s >= b.wall_secs:
+            return False, "wall-clock budget exhausted"
+        if b.tokens is not None and last is not None and last.cumulative_tokens >= b.tokens:
+            return False, "token budget exhausted"
+        if b.generations is not None and len(self.history) >= b.generations:
+            return False, "generation budget reached"
+        if self.no_improve_gens > 0 and self._gens_since_improvement() >= self.no_improve_gens:
+            return False, f"no improvement in {self.no_improve_gens} generations"
+        mv = self.marginal_value()
+        if self.min_marginal > 0 and mv is not None and mv < self.min_marginal:
+            return False, f"marginal value {mv:.3f} < floor {self.min_marginal:.3f} (Δquality/Δcompute)"
+        return True, "continue"
+
+
+# ── Adaptive compute allocation (§17) ──────────────────────────────────────────
+@dataclass
+class Allocation:
+    grants: dict[str, int]     # branch id -> tokens granted this round
+    terminated: list[str]      # branch ids pruned (0 tokens, killed)
+
+
+def allocate(branches: list[tuple[str, float]], pool_tokens: int, *,
+             prune_below: float = 0.0, keep_top: int | None = None,
+             bias: float = 1.5, min_grant: int = 0) -> Allocation:
+    """Split `pool_tokens` across branches by expected value (§17).
+
+    branches: (id, score) with higher score = more promising. A branch is TERMINATED (gets 0)
+    if its score is below `prune_below` or it falls outside the top `keep_top`. Survivors share
+    the pool in proportion to score**bias, so leaders get disproportionately more (bias>1) — the
+    system must not spend equally on good and bad ideas. `min_grant` guarantees a floor so a
+    kept-but-trailing branch still gets a probe. Deterministic; ties broken by input order."""
+    if pool_tokens <= 0 or not branches:
+        return Allocation({}, [b for b, _ in branches])
+    ranked = sorted(branches, key=lambda x: x[1], reverse=True)
+    survivors = [(b, s) for b, s in ranked if s >= prune_below]
+    if keep_top is not None:
+        survivors = survivors[:max(0, keep_top)]
+    kept_ids = {b for b, _ in survivors}
+    terminated = [b for b, _ in branches if b not in kept_ids]
+    if not survivors:
+        return Allocation({}, terminated)
+    weights = [(b, max(s, 0.0) ** bias) for b, s in survivors]
+    total_w = sum(w for _, w in weights) or float(len(weights))
+    grants: dict[str, int] = {}
+    for b, w in weights:
+        share = (w / total_w) if total_w else 1.0 / len(weights)
+        grants[b] = max(min_grant, int(round(pool_tokens * share)))
+    # correct rounding drift so grants sum to the pool (adjust the top branch)
+    drift = pool_tokens - sum(grants.values())
+    if drift and weights:
+        top = weights[0][0]
+        grants[top] = max(min_grant, grants[top] + drift)
+    return Allocation(grants, terminated)
+
+
+# ── Planner: budget + intensity + hardware → population plan (§15/§16, §5/§9) ──
+# Diverse roles for solution-space coverage (§9/§10). Correlated agents waste compute.
+ROLES = [
+    "first-principles solver", "conventional solver", "skeptic", "debugger",
+    "security reviewer", "researcher", "minimalist", "adversarial critic",
+    "performance specialist", "alternative-architecture agent",
+]
+
+# Per-intensity search shape (§16): initial population, critic count, generation cap, and how
+# many branches survive each prune. Population is a SEARCH choice — independent of hardware.
+INTENSITY = {
+    "light":    {"population": 3,  "critics": 1, "generations": 2,  "keep": 2},
+    "standard": {"population": 8,  "critics": 2, "generations": 4,  "keep": 4},
+    "deep":     {"population": 16, "critics": 3, "generations": 6,  "keep": 6},
+    "extreme":  {"population": 24, "critics": 5, "generations": 10, "keep": 8},
+}
+
+
+@dataclass
+class SwarmPlan:
+    intensity: str
+    initial_population: int      # logical agents in the explore stage (§9)
+    max_concurrency: int         # simultaneous inference — HARDWARE bound (§5), != population
+    critics: int                 # judge/critic agents (§11)
+    generations: int             # generation cap (a convergence bound, §23)
+    keep: int                    # survivors promoted past each prune (§12)
+    roles: list[str]             # diversified assignments (§10)
+    estimated_logical: int       # rough total agents incl. spawned descendants (§13) — info only
+
+    @property
+    def queued_at_start(self) -> int:
+        """Logical agents that must wait because concurrency < population (§18)."""
+        return max(0, self.initial_population - self.max_concurrency)
+
+
+def choose_intensity(budget: Budget) -> str:
+    """Map a budget to a search intensity for `--swarm auto` (§16). Bigger budget → deeper
+    search. Falls back to 'standard' when the budget carries no size signal."""
+    if budget.wall_secs is not None:
+        s = budget.wall_secs
+        return "light" if s < 900 else "standard" if s < 3600 else "deep" if s < 14400 else "extreme"
+    if budget.tokens is not None:
+        t = budget.tokens
+        return ("light" if t < 1_000_000 else "standard" if t < 10_000_000
+                else "deep" if t < 50_000_000 else "extreme")
+    if budget.generations is not None:
+        g = budget.generations
+        return "light" if g <= 2 else "standard" if g <= 4 else "deep" if g <= 6 else "extreme"
+    return "standard"
+
+
+def plan(budget: Budget, concurrency: int, *, intensity: str = "auto") -> SwarmPlan:
+    """Turn a budget + measured hardware concurrency into a population plan (§15/§16).
+
+    KEY: logical population comes from the intensity (a search decision); `concurrency` only
+    caps how many run at once (§5). On a 1-GPU box a 16-agent 'deep' plan still has 16 logical
+    agents — the rest queue (§18); the penalty is wall-clock, not lost search capability. This
+    is the opposite of capacity.swarm_size(), which caps logical count AT concurrency."""
+    intensity = choose_intensity(budget) if intensity == "auto" else intensity
+    if intensity not in INTENSITY:
+        raise ValueError(f"unknown intensity {intensity!r} (light|standard|deep|extreme|auto)")
+    p = INTENSITY[intensity]
+    pop = p["population"]
+    roles = [ROLES[i % len(ROLES)] for i in range(pop)]
+    # estimate total logical agents: initial + ~one descendant per survivor per later generation
+    estimated = pop + p["keep"] * max(0, p["generations"] - 1)
+    return SwarmPlan(intensity=intensity, initial_population=pop,
+                     max_concurrency=max(1, concurrency), critics=p["critics"],
+                     generations=p["generations"], keep=p["keep"], roles=roles,
+                     estimated_logical=estimated)

{drydock_cli-3.1.30 → drydock_cli-3.1.31}/drydock/config.py

@@ -27,6 +27,12 @@ DEFAULTS: dict[str, object] = {
     # default invisibly, which left base_url absent from the file. For a non-vllm
     # provider, override base_url (or use --base-url / the first-run prompt).
     "base_url": "http://localhost:8000/v1",
+    # API key for the MAIN model endpoint, sent as `Authorization: Bearer <key>`.
+    # Leave "" for a local server that needs no auth (llama.cpp/Ollama/LM Studio);
+    # set it for a hosted OpenAI-compatible endpoint (e.g. a vLLM behind a gateway,
+    # OpenRouter, Together). Mirrors `advisor_api_key`. Can also come from the
+    # provider's api_key_env (e.g. OPENAI_API_KEY); the config value wins.
+    "api_key": "",
     "max_tokens": 8192,  # 4096 truncated large file writes mid-JSON (→ _raw fail)
     "temperature": 0.2,
     # The model server's context window (llama.cpp -c / vLLM --max-model-len).

{drydock_cli-3.1.30 → drydock_cli-3.1.31}/drydock/mission_run.py

@@ -270,14 +270,21 @@ class TaskOutcome:
 
 def run_task(store: M.MissionStore, task: dict, *, mission: dict, cwd: str, repo: str,
              worker: WorkerFn, evaluator: EvaluatorFn, base_config: dict | None = None,
-             worker_id: str = "worker-1") -> TaskOutcome | None:
-    """Execute one leased task with checkpoint + KEEP/REVERT. None if the lease was lost."""
+             worker_id: str = "worker-1",
+             checkpoint_fn: "Callable[[str], str]" = checkpoint,
+             restore_fn: "Callable[[str, str], bool]" = restore) -> TaskOutcome | None:
+    """Execute one leased task with checkpoint + KEEP/REVERT. None if the lease was lost.
+
+    checkpoint_fn/restore_fn default to git (`repo` = a checkout dir), but are injectable so a
+    non-git target can plug in its own snapshot mechanism — e.g. `docker commit` for a task that
+    lives in a container (the tbench/ddt harness). `repo` is then just an opaque handle passed
+    to those callables; an empty `repo` disables checkpointing."""
     mid = mission["id"]
     tid = task["id"]
     if not store.claim_task(tid, worker_id):
         return None
     before = mission.get("current_metric") or 0.0
-    cp = checkpoint(repo) if repo else ""
+    cp = checkpoint_fn(repo) if repo else ""
     store.event(mid, "task_started", task=tid, objective=task.get("objective", ""))
     # Context reconstruction (§13): surface known failed approaches so the worker doesn't
     # repeat them (§16). Key the lookup on the SPECIFIC target (task reason, e.g. the failing
@@ -295,7 +302,7 @@ def run_task(store: M.MissionStore, task: dict, *, mission: dict, cwd: str, repo
     store.add_usage(mid, tokens=wr.in_tokens + wr.out_tokens, wall_s=time.monotonic() - t0)
     if not wr.ok:
         if repo and cp:
-            restore(repo, cp)
+            restore_fn(repo, cp)
         store.complete_task(tid, M.T_FAILED, {"error": wr.error})
         return TaskOutcome(tid, accept=False, metric_after=before, reverted=True, error=wr.error)
 
@@ -309,7 +316,7 @@ def run_task(store: M.MissionStore, task: dict, *, mission: dict, cwd: str, repo
     tamper = tampered_paths(changed, protected)
     if tamper:
         if repo and cp:
-            restore(repo, cp)
+            restore_fn(repo, cp)
         store.add_usage(mid, experiments=1)
         store.complete_task(tid, M.T_FAILED,
                             {"summary": wr.summary, "decision": "REVERT", "reverted": True,
@@ -327,7 +334,7 @@ def run_task(store: M.MissionStore, task: dict, *, mission: dict, cwd: str, repo
     ev = evaluator(task, mission, cwd, before)
     store.add_usage(mid, experiments=1)
     if ev.accept:
-        end = checkpoint(repo) if repo else ""
+        end = checkpoint_fn(repo) if repo else ""
         store.set_metric(mid, ev.metric_after)
         store.complete_task(tid, M.T_COMPLETED,
                             {"summary": wr.summary, "metric": ev.metric_after, "decision": "KEEP",
@@ -340,7 +347,7 @@ def run_task(store: M.MissionStore, task: dict, *, mission: dict, cwd: str, repo
                             f"{(wr.summary or '')[:220]}", confidence=0.7, sources=[tid])
     else:
         if repo and cp:
-            restore(repo, cp)      # auto-revert the regression (§19)
+            restore_fn(repo, cp)      # auto-revert the regression (§19)
         store.complete_task(tid, M.T_COMPLETED,
                             {"summary": wr.summary, "decision": "REVERT", "reverted": True})
         # negative knowledge stays in history even though the code is reverted (§16/§19)
@@ -505,6 +512,8 @@ def run_mission(store: M.MissionStore, mission_id: str, *, cwd: str, repo: str =
                 review_every_tasks: int = 10, review_every_secs: float = 3600.0,
                 reviewer: Callable[[M.MissionStore, dict, dict], dict] | None = None,
                 critic: Callable[[M.MissionStore, dict, dict], str] | None = None,
+                checkpoint_fn: "Callable[[str], str]" = checkpoint,
+                restore_fn: "Callable[[str, str], bool]" = restore,
                 on_event: Callable[[str, dict], None] | None = None) -> RunSummary:
     """Drive a mission to a stopping condition with a SINGLE worker (§41/§49). Stops on
     success (§30), budget (§30), stagnation (§22), or an empty queue (planner is a later
@@ -557,7 +566,8 @@ def run_mission(store: M.MissionStore, mission_id: str, *, cwd: str, repo: str =
         task = ready[0]
         out = run_task(store, task, mission=m, cwd=cwd, repo=repo, worker=worker,
                        evaluator=(evaluator or (lambda *_: Evaluation(accept=True))),
-                       base_config=base_config, worker_id=worker_id)
+                       base_config=base_config, worker_id=worker_id,
+                       checkpoint_fn=checkpoint_fn, restore_fn=restore_fn)
         progressed = bool(out and out.accept)
         _ev("task_done", task=task["id"], accepted=progressed,
             metric=out.metric_after if out else None)

{drydock_cli-3.1.30 → drydock_cli-3.1.31/drydock_cli.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drydock-cli
-Version: 3.1.30
+Version: 3.1.31
 Summary: Drydock — a local, provider-agnostic terminal coding agent for local LLMs
 Author: Frank Bobe III
 License-Expression: Apache-2.0

{drydock_cli-3.1.30 → drydock_cli-3.1.31}/drydock_cli.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ drydock/__init__.py
 drydock/__main__.py
 drydock/advisor.py
 drydock/agent.py
+drydock/asr.py
 drydock/bash_safety.py
 drydock/bottleneck.py
 drydock/budget.py
@@ -95,6 +96,7 @@ drydock_cli.egg-info/requires.txt
 drydock_cli.egg-info/top_level.txt
 tests/test_advisor.py
 tests/test_approval.py
+tests/test_asr.py
 tests/test_back_command.py
 tests/test_bash_background.py
 tests/test_bash_binary_output.py
@@ -152,6 +154,7 @@ tests/test_jobs.py
 tests/test_knowledge_tool_available.py
 tests/test_leaked_tool_call.py
 tests/test_loop_detect.py
+tests/test_main_api_key.py
 tests/test_mcp.py
 tests/test_mission.py
 tests/test_mission_cli.py

{drydock_cli-3.1.30 → drydock_cli-3.1.31}/pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 # PyPI distribution name is drydock-cli (the established install name, continued
 # from the retired v2 fork); the import package + CLI command stay `drydock`.
 name = "drydock-cli"
-version = "3.1.30"
+version = "3.1.31"
 description = "Drydock — a local, provider-agnostic terminal coding agent for local LLMs"
 readme = "README.md"
 requires-python = ">=3.11"

drydock_cli-3.1.31/tests/test_asr.py

@@ -0,0 +1,140 @@
+"""Tests for the Adaptive Swarm Runtime search core (drydock/asr.py) — pure, no model/GPU."""
+from __future__ import annotations
+
+import pytest
+
+from drydock import asr
+
+
+# ── Budget (§15) ────────────────────────────────────────────────────────────
+def test_budget_parse_time_and_tokens():
+    assert asr.Budget.parse("2h").wall_secs == 7200
+    assert asr.Budget.parse("30m").wall_secs == 1800
+    assert asr.Budget.parse("90s").wall_secs == 90
+    assert asr.Budget.parse("100M-tokens").tokens == 100_000_000
+    assert asr.Budget.parse("50k-tokens").tokens == 50_000
+    assert asr.Budget.parse("2h").is_finite
+
+
+def test_budget_rejects_garbage_never_silently_unbounded():
+    with pytest.raises(ValueError):
+        asr.Budget.parse("whenever")
+    assert asr.Budget().is_finite is False        # empty budget is explicitly not finite
+
+
+# ── Convergence (§23) ─────────────────────────────────────────────────────────
+def _conv(**kw):
+    return asr.Convergence(budget=asr.Budget(generations=100), **kw)
+
+
+def test_convergence_stops_on_success():
+    c = _conv(success_score=100.0)
+    c.record(best_score=100.0, cumulative_tokens=5000, elapsed_s=60)
+    go, why = c.should_continue()
+    assert not go and "success" in why
+
+
+def test_convergence_stops_on_budget():
+    c = asr.Convergence(budget=asr.Budget(tokens=10_000))
+    c.record(best_score=40.0, cumulative_tokens=12_000, elapsed_s=30)
+    go, why = c.should_continue()
+    assert not go and "token budget" in why
+
+
+def test_convergence_stops_on_no_improvement():
+    c = _conv(no_improve_gens=2, success_score=100.0)
+    c.record(50.0, 1000, 10)
+    c.record(50.0, 2000, 20)     # no gain
+    c.record(50.0, 3000, 30)     # still no gain -> 2 stagnant gens
+    go, why = c.should_continue()
+    assert not go and "no improvement" in why
+
+
+def test_convergence_stops_on_low_marginal_value():
+    # big jump, then a tiny gain for a lot of tokens -> marginal value below floor
+    c = _conv(min_marginal=1.0, no_improve_gens=0, success_score=100.0)
+    c.record(40.0, 1000, 10)
+    c.record(80.0, 3000, 20)                 # +40 over 2k tok = 20/1k -> continues
+    assert c.should_continue()[0] is True
+    c.record(80.5, 13000, 40)                # +0.5 over 10k tok = 0.05/1k -> below 1.0 floor
+    go, why = c.should_continue()
+    assert not go and "marginal value" in why
+
+
+def test_convergence_continues_while_improving():
+    c = _conv(no_improve_gens=3, min_marginal=0.5, success_score=100.0)
+    c.record(40.0, 1000, 10)
+    c.record(60.0, 2000, 20)                 # +20/1k, improving, under budget
+    assert c.should_continue() == (True, "continue")
+
+
+def test_marginal_value_needs_two_records():
+    c = _conv()
+    assert c.marginal_value() is None
+    c.record(10.0, 500, 5)
+    assert c.marginal_value() is None
+
+
+# ── Adaptive allocation (§17) ──────────────────────────────────────────────────
+def test_allocate_favors_leaders_and_prunes_weak():
+    branches = [("A", 0.91), ("B", 0.86), ("C", 0.43), ("D", 0.21)]
+    alloc = asr.allocate(branches, 8000, prune_below=0.3)
+    assert "D" in alloc.terminated and alloc.grants.get("D", 0) == 0   # weakest pruned
+    assert set(alloc.grants) == {"A", "B", "C"}
+    assert sum(alloc.grants.values()) == 8000                          # whole pool spent
+    assert alloc.grants["A"] > alloc.grants["B"] > alloc.grants["C"]   # EV ordering
+    assert alloc.grants["C"] < alloc.grants["A"] / 2                   # bias toward the leader
+
+
+def test_allocate_keep_top_terminates_the_rest():
+    branches = [("A", 0.9), ("B", 0.8), ("C", 0.7), ("D", 0.6)]
+    alloc = asr.allocate(branches, 6000, keep_top=2)
+    assert set(alloc.grants) == {"A", "B"}
+    assert sorted(alloc.terminated) == ["C", "D"]
+    assert sum(alloc.grants.values()) == 6000
+
+
+def test_allocate_empty_pool_terminates_all():
+    alloc = asr.allocate([("A", 0.9), ("B", 0.5)], 0)
+    assert alloc.grants == {} and sorted(alloc.terminated) == ["A", "B"]
+
+
+def test_allocate_min_grant_floor_for_kept_branches():
+    branches = [("A", 0.99), ("B", 0.01)]
+    alloc = asr.allocate(branches, 10_000, min_grant=200)
+    assert alloc.grants["B"] >= 200                    # trailing-but-kept branch still probed
+    assert alloc.grants["A"] > alloc.grants["B"]
+
+
+# ── Planner (§15/§16, §5/§18) ──────────────────────────────────────────────────
+def test_choose_intensity_scales_with_budget():
+    assert asr.choose_intensity(asr.Budget.parse("10m")) == "light"
+    assert asr.choose_intensity(asr.Budget.parse("45m")) == "standard"
+    assert asr.choose_intensity(asr.Budget.parse("2h")) == "deep"
+    assert asr.choose_intensity(asr.Budget.parse("8h")) == "extreme"
+    assert asr.choose_intensity(asr.Budget.parse("5M-tokens")) == "standard"
+
+
+def test_plan_decouples_logical_population_from_concurrency():
+    # THE core ASR property (§5/§18): a deep plan on a 1-slot box keeps all 16 logical agents;
+    # they queue, they aren't discarded down to concurrency (unlike capacity.swarm_size).
+    weak = asr.plan(asr.Budget.parse("2h"), concurrency=1, intensity="deep")
+    assert weak.initial_population == 16 and weak.max_concurrency == 1
+    assert weak.queued_at_start == 15
+    strong = asr.plan(asr.Budget.parse("2h"), concurrency=8, intensity="deep")
+    assert strong.initial_population == 16 and strong.max_concurrency == 8
+    assert strong.queued_at_start == 8          # same search, more parallelism → less queue
+
+
+def test_plan_auto_and_intensity_ladder():
+    auto = asr.plan(asr.Budget.parse("30m"), concurrency=4)   # auto -> standard
+    assert auto.intensity == "standard"
+    pops = [asr.plan(asr.Budget.parse("2h"), 4, intensity=i).initial_population
+            for i in ("light", "standard", "deep", "extreme")]
+    assert pops == sorted(pops) and pops[0] < pops[-1]        # deeper => bigger population
+    assert len(auto.roles) == auto.initial_population and len(set(auto.roles)) > 1  # diverse
+
+
+def test_plan_rejects_unknown_intensity():
+    with pytest.raises(ValueError):
+        asr.plan(asr.Budget.parse("1h"), 4, intensity="ludicrous")

drydock_cli-3.1.31/tests/test_main_api_key.py

@@ -0,0 +1,76 @@
+"""Main-model API key in config.toml (mirrors advisor_api_key)."""
+from __future__ import annotations
+
+from drydock import config as C
+
+
+def test_api_key_is_a_declared_default():
+    # the fix: api_key is a first-class default (so it appears in generated config.toml
+    # and survives resolve), not just something providers.py happened to read.
+    assert "api_key" in C.DEFAULTS and C.DEFAULTS["api_key"] == ""
+
+
+def test_config_file_api_key_round_trips(tmp_path):
+    p = tmp_path / "config.toml"
+    C.save_file({**C.DEFAULTS, "api_key": "sk-main-123"}, p)
+    assert "api_key" in p.read_text()                 # written to the file (discoverable)
+    cfg = C.resolve({}, p)
+    assert cfg["api_key"] == "sk-main-123"            # and resolved back
+
+
+def test_generated_default_config_includes_api_key(tmp_path):
+    p = tmp_path / "config.toml"
+    C.resolve({}, p)                                  # non-existent -> writes DEFAULTS
+    assert "api_key" in p.read_text()
+
+
+def test_stream_sends_config_api_key_to_client(monkeypatch):
+    import openai
+
+    from drydock import providers as P
+    captured: dict = {}
+
+    class _Stop(Exception):
+        pass
+
+    class _FakeClient:
+        def __init__(self, **kw):
+            captured.update(kw)
+
+        def __getattr__(self, _name):     # any use (client.chat...) aborts after construction
+            raise _Stop()
+
+    monkeypatch.setattr(openai, "OpenAI", _FakeClient)
+    gen = P.stream("m", "sys", [{"role": "user", "content": "hi"}], [],
+                   {"provider": "vllm", "base_url": "http://x/v1", "api_key": "sk-MAIN"})
+    try:
+        next(gen)
+    except Exception:
+        pass
+    assert captured.get("api_key") == "sk-MAIN"       # main-model key reaches the client
+
+
+def test_stream_empty_api_key_falls_back_not_crashes(monkeypatch):
+    import openai
+
+    from drydock import providers as P
+    captured: dict = {}
+
+    class _Stop(Exception):
+        pass
+
+    class _FakeClient:
+        def __init__(self, **kw):
+            captured.update(kw)
+
+        def __getattr__(self, _name):
+            raise _Stop()
+
+    monkeypatch.setattr(openai, "OpenAI", _FakeClient)
+    gen = P.stream("m", "sys", [{"role": "user", "content": "hi"}], [],
+                   {"provider": "vllm", "base_url": "http://x/v1", "api_key": ""})
+    try:
+        next(gen)
+    except Exception:
+        pass
+    assert captured.get("api_key") == "dummy"         # empty -> provider default (unchanged behavior)

{drydock_cli-3.1.30 → drydock_cli-3.1.31}/tests/test_mission_run.py

@@ -347,6 +347,33 @@ def test_restore_preserves_mission_workspace(tmp_path):
     assert (ws / "state.db").read_text() == "STATE-V2"  # mission state SURVIVES the revert
 
 
+def test_injectable_checkpoint_restore(tmp_path):
+    # a non-git target (e.g. a container via `docker commit`) plugs in its own snapshot fns
+    snaps = {"made": [], "restored": []}
+
+    def cp_fn(handle):
+        ref = f"snap-{len(snaps['made'])}"
+        snaps["made"].append(ref)
+        return ref
+
+    def rs_fn(handle, ref):
+        snaps["restored"].append(ref)
+        return True
+
+    s = M.MissionStore(tmp_path / "s.db")
+    mid = s.create_mission("obj")
+    s.set_metric(mid, 50.0)
+    tid = s.add_task(mid, "risky")
+    out = R.run_task(s, s.get_task(tid), mission=s.get_mission(mid), cwd=".", repo="ctr:task",
+                     worker=lambda *a: R.WorkerResult(ok=True, summary="tried"),
+                     evaluator=lambda *a: R.Evaluation(accept=False, metric_before=50, metric_after=40,
+                                                       reason="worse"),
+                     checkpoint_fn=cp_fn, restore_fn=rs_fn)
+    assert out and out.reverted
+    assert snaps["made"] == ["snap-0"]           # checkpointed via the injected fn (not git)
+    assert snaps["restored"] == ["snap-0"]        # reverted via the injected fn
+
+
 def test_verify_passes_reflects_the_tree(tmp_path):
     repo = _repo(tmp_path)
     (Path(repo) / "answer.txt").write_text("PASS")