From 356f8c8096aa3b73240a1e71a2fd6548dd293788 Mon Sep 17 00:00:00 2001 From: Yashvardhan Thanvi Date: Thu, 15 Oct 2026 00:22:22 +0530 Subject: [PATCH 1/2] fix(sleep): enforce max_tokens_per_night in the nightly cycle The config default and the Budget controller existed, but nothing in the engine read the key or constructed a Budget, so a night ran every step no matter how many tokens it had already used. Build a Budget from max_tokens_per_night at the start of each cycle and check it, against the backend's token counter, before each step that would start new model work: the consolidation, its skill and memory edit steps, the held-out test replay and every per-skill fan-out group. Steps that never started are listed in a report note, in budget skip evidence events and in the cycle end event. A started step always finishes, so a gate decision never rests on partial evidence and nothing is cut off mid-write. A handoff round counts unanswered prompts and placeholder replies, so when the budget is reached while prompts are still pending the round pauses for answers instead of stopping. Otherwise it would stage and save a night built from placeholders. If the budget is already spent before consolidation, the night records no held-out scores instead of zeros. 0 or null keeps the night unbudgeted, matching Budget's existing semantics. Negative, boolean, fractional, infinite or non-numeric values fail before any model call. Only tonight's spend counts, so a backend reused across nights is not stopped early. Fixes #313 --- docs/sleep/README.md | 46 +++ plugins/README.md | 6 +- plugins/claude-code/README.md | 6 +- plugins/openclaw/README.md | 8 +- skillopt_sleep/consolidate.py | 44 ++- skillopt_sleep/cycle.py | 125 ++++++++- skillopt_sleep/dream.py | 5 +- skillopt_sleep/multi_skill.py | 23 +- tests/test_sleep_token_budget.py | 467 +++++++++++++++++++++++++++++++ 9 files changed, 705 insertions(+), 25 deletions(-) create mode 100644 tests/test_sleep_token_budget.py diff --git a/docs/sleep/README.md b/docs/sleep/README.md index ba84923c3..c8cf8f8e9 100644 --- a/docs/sleep/README.md +++ b/docs/sleep/README.md @@ -344,6 +344,52 @@ documents the separate HTTPS-only boundary for Azure managed-identity credential Deterministic proof (no API key): `python -m skillopt_sleep.experiments.run_experiment --persona researcher --assert-improves`. +### Nightly token budget + +`max_tokens_per_night` in `~/.skillopt-sleep/config.json` (default `400000`) +stops a `run` or `dry-run` from starting more model work once tonight's spend +reaches it. The cycle checks the budget before the consolidation, before its +skill and memory edit steps, before the held-out test replay, and before each +per-skill fan-out group. Steps that did not start are listed in a note in +`report.md`, for example +`max_tokens_per_night reached (431870 of 400000 tokens); skipped: memory edits`. + +```json +{"max_tokens_per_night": 200000} +``` + +Set the value to `0` or `null` for no budget. Anything other than a whole, +non-negative number stops the run with an error before any model call. There +is no elapsed-time budget. + +This is not a hard cap. A step that has started always finishes, so a night +can end above the budget by the cost of one step: the first replay of the +validation and training tasks, one edit step (a training replay, a reflect +call and a validation replay of the candidate; with `dream_rollouts` set to K, +K attempts per training task), or the held-out test replay. Each fan-out group +that starts runs its own first replay. + +Spend comes from the backend's token counter, the one the report header prints +as `tokens used`. The header shows the backend's running total and the note +shows tonight's spend; they are equal unless a caller reuses one backend object +across nights. What the counter covers depends on the backend: + +- CLI backends estimate tokens from prompt and reply length. LLM task mining and + contrastive reflection (`dream_rollouts` above 1) call the model directly and + are not counted. +- `azure_openai` adds the usage reported by the API to that estimate, so it + counts mining, and replay and reflect calls count roughly twice. +- `handoff` counts the prompt and answer text of the current round. A budget + check that finds unanswered prompts pauses the round for answers instead of + stopping it, so a round only stops on the budget once every counted call has + a real answer. +- `mock` reports no usage and never reaches the budget. + +A night that stops on the budget still advances the harvest checkpoint, so its +sessions are not mined again, and skipped work is not retried on a later night. +Tonight's tasks stay in the task archive, which later nights only use through +`recall_k`. + ### Opt-in: per-skill fan-out Set `"multi_skill_fanout": true` in `~/.skillopt-sleep/config.json` to add an diff --git a/plugins/README.md b/plugins/README.md index b3aba6bcf..b76d7e0f7 100644 --- a/plugins/README.md +++ b/plugins/README.md @@ -298,7 +298,9 @@ the shipping CLI defaults. - `mock` and `handoff` make no network calls. - `run` stages proposals; `adopt` is the normal live-change boundary. - Adoption backs up existing target files. -- `--max-sessions` and `--max-tasks` bound work, but the main CLI does not yet - enforce a hard token or elapsed-time budget. +- `--max-sessions` and `--max-tasks` bound work. `max_tokens_per_night` in the + config file (default `400000`) stops a night from starting new steps once its + token count reaches the budget, but a started step finishes, so it is not a + hard cap. There is no elapsed-time budget. - Treat real-backend transcript excerpts as data shared with the selected provider. diff --git a/plugins/claude-code/README.md b/plugins/claude-code/README.md index da1c08b15..9d5f6b954 100644 --- a/plugins/claude-code/README.md +++ b/plugins/claude-code/README.md @@ -175,8 +175,10 @@ The scheduled run stages proposals unless `--auto-adopt` is explicitly selected. - **Read-only** harvest of `~/.claude`. `mock` replay has no side effects. - Proposals are **staged**, never auto-applied (unless you opt in with `--auto-adopt`). - Every adopt writes a backup under the staging dir's `backup/`. -- `--max-sessions` and `--max-tasks` bound work, but the main CLI does not enforce - a hard token or wall-clock budget. +- `--max-sessions` and `--max-tasks` bound work. `max_tokens_per_night` in + `~/.skillopt-sleep/config.json` stops a night from starting new steps once its + token count reaches the budget; a started step finishes, so it is not a hard + cap, and there is no wall-clock budget. - Real backends share truncated session/task content with the selected provider; do not assume outbound prompts have been fully redacted. diff --git a/plugins/openclaw/README.md b/plugins/openclaw/README.md index 76121ea85..129ff8652 100644 --- a/plugins/openclaw/README.md +++ b/plugins/openclaw/README.md @@ -56,8 +56,12 @@ staged files before changing a live skill. paths all point to isolated test locations. 3. Review `config.json`. In particular, do not assume that values such as - `max_tokens_per_night` or `replay_mode` are enforced by this custom wrapper - merely because they appear in the example config. + `replay_mode` are enforced by this custom wrapper merely because they appear + in the example config. `max_tokens_per_night` is applied by the shared cycle, + but it is measured with the wrapper backend's rough `tokens_used()` estimate + (a quarter of the attempt prompt length plus 200 per attempt, 600 per rubric + judge call, and a third of the reflect prompt plus 1500 per reflect call), so + it is not a billing limit. 4. Supply credentials through your normal secret-management mechanism. Do not commit a DeepSeek key or place it in a world-readable file. diff --git a/skillopt_sleep/consolidate.py b/skillopt_sleep/consolidate.py index 5b80b4a98..77e1e6a3e 100644 --- a/skillopt_sleep/consolidate.py +++ b/skillopt_sleep/consolidate.py @@ -10,7 +10,7 @@ import math from dataclasses import dataclass, field -from typing import List, Tuple +from typing import Callable, List, Optional, Tuple from skillopt_sleep.backend import Backend @@ -23,6 +23,10 @@ _HAVE_REPO_GATE = True +# gate_action of an epoch that never ran because the token budget was spent. +# Its scores were not measured, so callers must not report them as results. +BUDGET_EXHAUSTED = "budget_exhausted" + def _finite_score(value: float) -> float | None: """Return a JSON-safe score, preserving finite values only.""" @@ -59,6 +63,8 @@ class ConsolidationResult: # fresh final replay, so aggregate changes never hide an individual task. # Keep new optional fields last to preserve positional construction. gate_trials: List[dict] = field(default_factory=list) + # Steps that never started because the caller's token budget was spent. + budget_skipped: List[str] = field(default_factory=list) def _split(tasks: List[TaskRecord]) -> Tuple[List[TaskRecord], List[TaskRecord], bool]: @@ -192,6 +198,7 @@ def consolidate( evolve_skill: bool = True, evolve_memory: bool = True, night: int = 1, + budget_exhausted: Optional[Callable[[], bool]] = None, ) -> ConsolidationResult: """Run one consolidation epoch: reflect -> bounded edit -> gate. @@ -201,12 +208,42 @@ def consolidate( recorded so the report shows whether quality moved. Skill and memory are evolved in sequence (skill first if both enabled). + + ``budget_exhausted`` is checked before the epoch starts and before each + evolve target. Once it returns True, the steps that have not started are + skipped and listed in ``budget_skipped``. A started step always finishes, + so every gate decision still rests on complete evidence; the final replay + only re-scores documents an earlier replay already scored. """ from skillopt_sleep import evidence as evlog ev = evlog.get(backend) train_tasks, val_tasks, holdout_leaked = _split(tasks) gate_off = str(gate_mode).strip().lower() in {"off", "none", "false", "greedy"} holdout_detail: List[dict] = [] + budget_skipped: List[str] = [] + + def _budget_spent(step: str) -> bool: + if budget_exhausted is None or not budget_exhausted(): + return False + budget_skipped.append(step) + if ev is not None: + ev.log("budget", "skip", step=step, night=night) + return True + + if _budget_spent("consolidation"): + return ConsolidationResult( + accepted=False, + gate_action=BUDGET_EXHAUSTED, + baseline_score=0.0, + candidate_score=0.0, + new_skill=skill, + new_memory=memory, + applied_edits=[], + rejected_edits=[], + holdout_baseline=0.0, + holdout_candidate=0.0, + budget_skipped=budget_skipped, + ) # ── baseline on the VAL slice (the gate reference) ──────────────────── # When the gate is OFF the user has opted out of holding out a validation set @@ -307,7 +344,7 @@ def _gate_apply(doc: str, edits: List[EditRecord], which: str) -> str: all_rejected.extend(applied) return doc - if evolve_skill: + if evolve_skill and not _budget_spent("skill edits"): if rollouts_k > 1: # multi-rollout contrastive reflection: run each train task K times # and distill a rule from the good-vs-bad contrast (the imagination signal). @@ -354,7 +391,7 @@ def _gate_apply(doc: str, edits: List[EditRecord], which: str) -> str: ) cand_skill = _gate_apply(cand_skill, edits, "skill") - if evolve_memory: + if evolve_memory and not _budget_spent("memory edits"): # re-evaluate failures under the (possibly improved) skill evlog.set_phase(backend, "train_post_skill") train_pairs2 = replay_batch(backend, train_tasks, cand_skill, cand_memory) @@ -493,4 +530,5 @@ def _gate_apply(doc: str, edits: List[EditRecord], which: str) -> str: reflect_raw=getattr(backend, "last_reflect_raw", "") or "", call_error=getattr(backend, "last_call_error", "") or "", holdout_leaked=holdout_leaked, + budget_skipped=budget_skipped, ) diff --git a/skillopt_sleep/cycle.py b/skillopt_sleep/cycle.py index 4ca5dc135..353325b93 100644 --- a/skillopt_sleep/cycle.py +++ b/skillopt_sleep/cycle.py @@ -15,20 +15,25 @@ import re import shutil import sys +import time import unicodedata from dataclasses import dataclass from typing import List, Optional, Sequence from skillopt_sleep import evidence from skillopt_sleep.backend import Backend, CursorBackendError, build_backend +from skillopt_sleep.budget import Budget from skillopt_sleep.config import DEFAULTS, SleepConfig, load_config +from skillopt_sleep.consolidate import BUDGET_EXHAUSTED from skillopt_sleep.dream import dream_consolidate from skillopt_sleep.evidence import EvidenceLog +from skillopt_sleep.handoff_backend import PendingCalls from skillopt_sleep.harvest_sources import harvest_for_config from skillopt_sleep.memory import ensure_skill_scaffold from skillopt_sleep.mine import group_tasks_by_skill_hint, mine from skillopt_sleep.replay import aggregate_scores, replay_batch from skillopt_sleep.multi_skill import ( + BUDGET_REACHED, SKIPPED, GroupConsolidation, SkillGroup, @@ -89,6 +94,29 @@ def _resolve_split_fractions(cfg: SleepConfig) -> tuple[float, float]: return val, test +def _night_token_budget(cfg: SleepConfig) -> Budget: + """Return ``max_tokens_per_night`` as a Budget; 0 or null means no cap. + + Raises ``ValueError`` before any model call for anything that is not a + whole, non-negative token count: negative numbers, booleans, fractional or + infinite floats (JSON ``Infinity`` included) and non-numeric strings. + """ + raw = cfg.get("max_tokens_per_night") + if raw is None: + return Budget() + try: + if isinstance(raw, bool) or (isinstance(raw, float) and not raw.is_integer()): + raise ValueError(raw) + limit = int(raw) + except (TypeError, ValueError, OverflowError): + limit = -1 + if limit < 0: + raise ValueError( + f"max_tokens_per_night must be 0 (no cap) or a positive token count, got {raw!r}" + ) + return Budget(max_tokens=limit or None) + + # ── Model-swap detection (F16) ─────────────────────────────── def _make_model_key(cfg: SleepConfig) -> str: """Stable string identifying the backend object(s) actually used. @@ -302,6 +330,10 @@ def _render_report_md(report: SleepReport, cfg: SleepConfig) -> str: backend = _markdown_text(cfg.get("backend")) replay = _markdown_text(cfg.get("replay_mode")) gate_action = _markdown_text(report.gate_action) + # A consolidation skipped on budget measured nothing; show no score. + measured = report.gate_action != BUDGET_EXHAUSTED + baseline_score = report.baseline_score if measured else None + candidate_score = report.candidate_score if measured else None lines = [ f"# SkillOpt-Sleep — night {report.night} report", "", @@ -309,8 +341,8 @@ def _render_report_md(report: SleepReport, cfg: SleepConfig) -> str: f"- backend: `{backend}` replay: `{replay}`", f"- sessions harvested: {report.n_sessions}", f"- tasks mined: {report.n_tasks} (replayed: {report.n_replayed})", - f"- held-out score: {_report_score(report.baseline_score)} " - f"-> {_report_score(report.candidate_score)}", + f"- held-out score: {_report_score(baseline_score)} " + f"-> {_report_score(candidate_score)}", f"- gate: **{gate_action}** (accepted={bool(report.accepted)})", f"- no-regression gate: " f"{'enabled' if cfg.get('gate_no_regression', False) else 'disabled'}", @@ -643,6 +675,7 @@ def run_sleep_cycle( it can inspect the backend's pending calls after the run. """ cfg = cfg or load_config() + budget = _night_token_budget(cfg) state = SleepState.load(cfg.state_path) project = _project_paths(cfg) @@ -668,6 +701,28 @@ def run_sleep_cycle( backend.preferences = cfg.get("preferences", "") _progress(cfg, f"night {night}: project={project} backend={backend.name}") + # max_tokens_per_night is measured with the backend's token counter, the + # one the report prints as "tokens used". Only tonight's spend counts, so a + # reused backend starts from its current total. The budget is checked + # before each step that would start new model work; a step in progress is + # never interrupted. + budget.start(time.monotonic, backend.tokens_used()) + budget_skipped: List[str] = [] + + def budget_exhausted() -> bool: + if not budget.exhausted( + tokens_now=backend.tokens_used(), clock_fn=time.monotonic + ): + return False + # A handoff round counts unanswered prompts and their placeholder + # replies. Stopping on that count would stage and save a night built + # from placeholders, so pause the round for answers instead and decide + # once every counted call is real. + pending = getattr(backend, "pending", None) + if pending: + raise PendingCalls(pending) + return True + # ── evidence log (the night's full evidentiary chain) ──────────────── # Pre-create the staging dir so evidence.jsonl accumulates exactly where # the report will land; dry-runs log into the state dir instead. @@ -696,7 +751,8 @@ def run_sleep_cycle( "gate_mixed_weight", "gate_no_regression", "edit_budget", "holdout_fraction", "val_fraction", "test_fraction", "dream_rollouts", "dream_factor", "recall_k", - "max_tasks_per_night", "lookback_hours", "llm_mine", + "max_tasks_per_night", "max_tokens_per_night", + "lookback_hours", "llm_mine", "evolve_skill", "evolve_memory")} cycle_config["opencode_tool_replay"] = ( cfg.get("opencode_tool_replay", False) is True @@ -744,7 +800,6 @@ def run_sleep_cycle( if since is None: lookback_hours = cfg.get("lookback_hours", 72) if lookback_hours is not None and lookback_hours > 0: - import time ref_time = clock if clock is not None else time.time() cutoff = ref_time - lookback_hours * 3600 since = _now_iso(cutoff) @@ -867,10 +922,15 @@ def run_sleep_cycle( evolve_skill=cfg.get("evolve_skill", True), evolve_memory=cfg.get("evolve_memory", True), night=night, + budget_exhausted=budget_exhausted, ) except CursorBackendError: _discard_unstaged_evidence(staging_dir_pre) raise + budget_skipped.extend(getattr(result, "budget_skipped", [])) + # When the budget was spent before consolidation, nothing was replayed or + # scored; keep zero placeholders out of state history and diagnostics. + scores_measured = result.gate_action != BUDGET_EXHAUSTED # archive tonight's real (non-dream) tasks so future nights can recall them state.add_to_archive([t.to_dict() for t in tasks if t.origin != "dream"]) _progress( @@ -880,7 +940,7 @@ def run_sleep_cycle( + (f" unmatched={len(result.unmatched_edits)}" if result.unmatched_edits else ""), ) - report.n_replayed = len(tasks) + report.n_replayed = len(tasks) if scores_measured else 0 report.baseline_score = result.baseline_score report.candidate_score = result.candidate_score report.accepted = result.accepted @@ -900,7 +960,12 @@ def run_sleep_cycle( # test_fraction=0.0 no test task exists and this block never runs, so # legacy nights are bit-for-bit unchanged and cost nothing extra. test_tasks = [t for t in tasks if t.split == "test"] - if test_tasks and ev is not None: + score_test = bool(test_tasks) and ev is not None + if score_test and budget_exhausted(): + score_test = False + budget_skipped.append("held-out test replay") + ev.log("budget", "skip", step="held-out test replay", night=night) + if score_test: final_skill = result.new_skill if result.accepted else skill final_memory = result.new_memory if result.accepted else memory test_pairs = replay_batch(backend, test_tasks, final_skill, final_memory) @@ -960,10 +1025,26 @@ def run_sleep_cycle( gate_mode=cfg.get("gate_mode", "on"), evolve_skill=cfg.get("evolve_skill", True), night=night, + budget_exhausted=budget_exhausted, ) except CursorBackendError: _discard_unstaged_evidence(staging_dir_pre) raise + budget_groups = [ + (name, outcome) + for name, outcome in consolidated_groups.items() + if outcome.reason.startswith(BUDGET_REACHED) + ] + if budget_groups: + budget_skipped.append( + f"work in {len(budget_groups)} per-skill group(s)" + ) + if ev is not None: + # Groups cut short log their own skip from consolidate(). + for name, outcome in budget_groups: + if outcome.status == SKIPPED: + ev.log("budget", "skip", step="per-skill group", + skill=name, night=night) for raw_name in grouped: name = str(raw_name or "").strip() outcome = skipped_groups.get(name) or consolidated_groups.get(name) @@ -981,6 +1062,14 @@ def run_sleep_cycle( report.tokens_used = backend.tokens_used() report.ended_at = _now_iso(clock) + if budget_skipped: + # Never truncate silently: say which steps the budget stopped. + spent = budget.tokens_spent(report.tokens_used) + report.notes.append( + f"max_tokens_per_night reached ({spent} of {budget.max_tokens} " + f"tokens); skipped: {', '.join(budget_skipped)}" + ) + _progress(cfg, f"token budget reached: skipped {', '.join(budget_skipped)}") # ── 5. stage (unless dry-run) ──────────────────────────────────────── staging_dir = "" @@ -1056,8 +1145,12 @@ def run_sleep_cycle( "gate_mode": cfg.get("gate_mode"), "gate_no_regression": cfg.get("gate_no_regression", False), "n_tasks": len(tasks), - "baseline_score": result.baseline_score, - "candidate_score": result.candidate_score, + "baseline_score": ( + result.baseline_score if scores_measured else None + ), + "candidate_score": ( + result.candidate_score if scores_measured else None + ), "accepted": result.accepted, "gate_action": result.gate_action, "holdout_leaked": getattr(result, "holdout_leaked", False), @@ -1080,9 +1173,13 @@ def run_sleep_cycle( except Exception: pass state.set_last_harvest(project, started) + night_scores = ( + {"baseline": result.baseline_score, "candidate": result.candidate_score} + if scores_measured + else {} + ) state.record_night({ - "night": night, "accepted": result.accepted, - "baseline": result.baseline_score, "candidate": result.candidate_score, + "night": night, "accepted": result.accepted, **night_scores, "n_tasks": len(tasks), "staging": staging_dir, }) state.set_last_model_key(_make_backend_key(backend)) # F16: track resolved model @@ -1095,11 +1192,13 @@ def run_sleep_cycle( if ev is not None: ev.log("cycle", "end", night=night, outcome="completed", gate_action=report.gate_action, accepted=report.accepted, - baseline_score=report.baseline_score, - candidate_score=report.candidate_score, + baseline_score=report.baseline_score if scores_measured else None, + candidate_score=report.candidate_score if scores_measured else None, n_applied_edits=len(report.edits), n_rejected_edits=len(report.rejected_edits), n_unmatched_edits=len(report.unmatched_edits), - tokens_used=report.tokens_used, adopted=adopted) + tokens_used=report.tokens_used, adopted=adopted, + max_tokens_per_night=budget.max_tokens, + budget_skipped=list(budget_skipped)) return CycleOutcome(report, staging_dir, adopted, adopted_paths) diff --git a/skillopt_sleep/dream.py b/skillopt_sleep/dream.py index 9906e07d5..cfb8fba97 100644 --- a/skillopt_sleep/dream.py +++ b/skillopt_sleep/dream.py @@ -17,7 +17,7 @@ from __future__ import annotations import re -from typing import List, Optional +from typing import Callable, List, Optional from skillopt_sleep.consolidate import ConsolidationResult, consolidate from skillopt_sleep.types import TaskRecord @@ -134,6 +134,7 @@ def dream_consolidate( evolve_skill: bool = True, evolve_memory: bool = True, night: int = 1, + budget_exhausted: Optional[Callable[[], bool]] = None, ) -> ConsolidationResult: """Recall similar past experience + dream synthetic variants, then run one gated consolidation epoch over the enlarged training pool. @@ -142,6 +143,7 @@ def dream_consolidate( augmentation only enlarge the TRAIN split, so the val slice the gate scores on is never polluted. With ``recall_k=0`` and ``dream_rollouts=1`` (the defaults) this is exactly the previous single-shot ``consolidate``. + ``budget_exhausted`` is passed through to ``consolidate``. """ train = [t for t in tasks if t.split == "train"] enlarged = list(tasks) @@ -165,4 +167,5 @@ def dream_consolidate( gate_no_regression=gate_no_regression, gate_mode=gate_mode, rollouts_k=dream_rollouts, evolve_skill=evolve_skill, evolve_memory=evolve_memory, night=night, + budget_exhausted=budget_exhausted, ) diff --git a/skillopt_sleep/multi_skill.py b/skillopt_sleep/multi_skill.py index 322476398..59aa0728f 100644 --- a/skillopt_sleep/multi_skill.py +++ b/skillopt_sleep/multi_skill.py @@ -22,6 +22,8 @@ SKIPPED = "skipped" FAILED = "failed" +BUDGET_REACHED = "token budget for the night was reached" + @dataclass class SkillGroup: @@ -54,6 +56,7 @@ def consolidate_groups( *, consolidate_fn: Callable[..., ConsolidationResult] = consolidate, group_kwargs_fn: Optional[Callable[[SkillGroup], Dict[str, object]]] = None, + budget_exhausted: Optional[Callable[[], bool]] = None, **consolidate_kwargs: object, ) -> Dict[str, GroupConsolidation]: """Consolidate each group independently, in order, isolating failures. @@ -68,7 +71,9 @@ def consolidate_groups( A group is ``skipped`` when it is unusable (blank name, or no tasks) and ``failed`` when its own consolidation raised; both leave every other - group's decision untouched. + group's decision untouched. When ``budget_exhausted`` returns True before a + usable group starts, that group is ``skipped`` too; a group that already + started finishes, and the predicate is passed on to ``consolidate_fn``. ``memory`` is the shared agent memory and is passed through read-only: group runs evolve skills only, so no group can rewrite another group's memory. @@ -82,6 +87,8 @@ def consolidate_groups( # duplicate keyword (which would otherwise turn the group into a failure). consolidate_kwargs = dict(consolidate_kwargs) consolidate_kwargs["evolve_memory"] = False + if budget_exhausted is not None: + consolidate_kwargs["budget_exhausted"] = budget_exhausted out: Dict[str, GroupConsolidation] = {} for group in groups: name = (group.skill_name or "").strip() @@ -101,6 +108,11 @@ def consolidate_groups( name, SKIPPED, reason="no mined tasks for this skill" ) continue + if budget_exhausted is not None and budget_exhausted(): + out[name] = GroupConsolidation( + name, SKIPPED, reason=BUDGET_REACHED, n_tasks=len(group.tasks) + ) + continue try: group_kwargs = dict(consolidate_kwargs) if group_kwargs_fn is not None: @@ -122,8 +134,15 @@ def consolidate_groups( n_tasks=len(group.tasks), ) continue + skipped_steps = getattr(result, "budget_skipped", None) or [] + reason = ( + f"{BUDGET_REACHED}; skipped {', '.join(skipped_steps)}" + if skipped_steps + else "" + ) out[name] = GroupConsolidation( - name, CONSOLIDATED, result=result, n_tasks=len(group.tasks) + name, CONSOLIDATED, result=result, reason=reason, + n_tasks=len(group.tasks), ) return out diff --git a/tests/test_sleep_token_budget.py b/tests/test_sleep_token_budget.py new file mode 100644 index 000000000..cd703aa15 --- /dev/null +++ b/tests/test_sleep_token_budget.py @@ -0,0 +1,467 @@ +"""max_tokens_per_night stops a night from scheduling more work (issue #313). + +Pure-stdlib (unittest), deterministic, no API key, no network. +Run: python -m pytest tests/test_sleep_token_budget.py +""" +from __future__ import annotations + +import glob +import json +import os +import tempfile +import unittest +from dataclasses import replace +from unittest import mock + +from skillopt_sleep.backend import MockBackend +from skillopt_sleep.config import load_config +from skillopt_sleep.consolidate import consolidate +from skillopt_sleep.cycle import run_sleep_cycle +from skillopt_sleep.evidence import read_events +from skillopt_sleep.experiments.personas import programmer_persona, researcher_persona +from skillopt_sleep.handoff_backend import HandoffBackend, PendingCalls +from skillopt_sleep.mine import assign_splits +from skillopt_sleep.multi_skill import ( + CONSOLIDATED, + SKIPPED, + SkillGroup, + consolidate_groups, +) +from skillopt_sleep.types import TaskRecord + +COST = 100 # tokens billed per uncached model call +_UNSET = object() + + +class MeteredBackend(MockBackend): + """MockBackend that reports token usage like the CLI backends do. + + Every uncached attempt and every reflect call bills ``COST`` tokens. As in + ``CliBackend``, re-running the same task with the same skill and memory is + served from cache and costs nothing, and the mock judge is local. + """ + + name = "metered" + + def __init__(self) -> None: + self.calls: list[str] = [] + self.attempt_skills: list[str] = [] + self._seen: set = set() + self._tokens = 0 + + def _bill(self, kind: str) -> None: + self.calls.append(kind) + self._tokens += COST + + def attempt(self, task, skill, memory, sample_id=0): + key = (task.id, skill, memory, sample_id) + if key not in self._seen: + self._seen.add(key) + self.attempt_skills.append(skill) + self._bill("attempt") + return super().attempt(task, skill, memory, sample_id) + + def reflect(self, failures, successes, skill, memory, **kwargs): + if failures: + self._bill("reflect") + return super().reflect(failures, successes, skill, memory, **kwargs) + + def tokens_used(self) -> int: + return self._tokens + + +class CountedBeforeConsolidationBackend(MeteredBackend): + """Usage already recorded before consolidation starts. + + Mirrors a backend that counts LLM mining usage: the counter is 0 when the + night starts and over the default budget at every later reading. + """ + + def __init__(self) -> None: + super().__init__() + self._reads = 0 + + def tokens_used(self) -> int: + self._reads += 1 + return 0 if self._reads == 1 else 500_000 + + +def _tasks(**kwargs): + kwargs.setdefault("val_fraction", 0.34) + return assign_splits(researcher_persona(), seed=42, **kwargs) + + +def _replay_cost(tasks) -> int: + """Tokens for the baseline val replay plus the train replay.""" + return COST * sum(1 for t in tasks if t.split in ("train", "val")) + + +def _budget_notes(report) -> list: + return [n for n in report.notes if "max_tokens_per_night" in n] + + +def _arithmetic_tasks(): + tasks = [ + TaskRecord( + id=f"t{i}", project="/p", + intent=f"What is {i} + {i}?", + reference_kind="exact", reference=str(i + i), + ) + for i in range(1, 7) + ] + return assign_splits(tasks, holdout_fraction=0.34, seed=42) + + +class TokenBudgetCycleTest(unittest.TestCase): + def setUp(self) -> None: + self._proj = tempfile.TemporaryDirectory() + self._home = tempfile.TemporaryDirectory() + self.addCleanup(self._proj.cleanup) + self.addCleanup(self._home.cleanup) + + def _cfg(self, budget=_UNSET, **overrides): + cfg = load_config( + invoked_project=self._proj.name, + projects="invoked", + backend="mock", + claude_home=os.path.join(self._home.name, ".claude"), + state_dir=os.path.join(self._home.name, "state"), + auto_adopt=False, + **overrides, + ) + if budget is not _UNSET: + # Set directly: load_config drops None overrides, but a config + # file can still carry an explicit null. + cfg.data["max_tokens_per_night"] = budget + return cfg + + def _run(self, tasks, backend=None, budget=_UNSET, dry_run=False, **overrides): + cfg = self._cfg(budget, **overrides) + backend = backend or MeteredBackend() + outcome = run_sleep_cycle( + cfg, seed_tasks=tasks, backend=backend, dry_run=dry_run + ) + return backend, outcome + + def _report_md(self, outcome) -> str: + with open( + os.path.join(outcome.staging_dir, "report.md"), encoding="utf-8" + ) as handle: + return handle.read() + + def _state_history(self) -> list: + with open(self._cfg().state_path, encoding="utf-8") as handle: + return json.load(handle)["history"] + + def test_night_stops_scheduling_work_once_the_budget_is_spent(self): + tasks = _tasks() + limit = _replay_cost(tasks) + + backend, outcome = self._run(tasks, budget=limit) + + self.assertLessEqual(backend.tokens_used(), limit) + self.assertNotIn("reflect", backend.calls) + report = outcome.report + self.assertFalse(report.accepted) + self.assertEqual(report.edits, []) + notes = _budget_notes(report) + self.assertEqual(len(notes), 1) + self.assertIn(f"{limit} of {limit}", notes[0]) + self.assertIn("skill edits", notes[0]) + self.assertIn("memory edits", notes[0]) + self.assertIn("max\\_tokens\\_per\\_night", self._report_md(outcome)) + self.assertFalse( + os.path.exists(os.path.join(outcome.staging_dir, "proposed_SKILL.md")) + ) + events = read_events(os.path.join(outcome.staging_dir, "evidence.jsonl")) + end = [e for e in events if (e["stage"], e["event"]) == ("cycle", "end")] + self.assertEqual(end[-1]["max_tokens_per_night"], limit) + self.assertEqual(end[-1]["budget_skipped"], ["skill edits", "memory edits"]) + + def test_one_token_of_headroom_lets_the_next_step_run(self): + tasks = _tasks() + limit = _replay_cost(tasks) + 1 + + backend, outcome = self._run(tasks, budget=limit) + + # The skill step started under budget and was allowed to finish, so + # its validated edit is kept; the memory step was never started. + self.assertEqual(backend.calls.count("reflect"), 1) + report = outcome.report + self.assertTrue(report.accepted) + self.assertTrue(report.edits) + notes = _budget_notes(report) + self.assertEqual(len(notes), 1) + self.assertNotIn("skill edits", notes[0]) + self.assertIn("memory edits", notes[0]) + self.assertTrue( + os.path.exists(os.path.join(outcome.staging_dir, "proposed_SKILL.md")) + ) + + def test_dry_run_stops_on_the_budget_without_persisting(self): + tasks = _tasks() + limit = _replay_cost(tasks) + + backend, outcome = self._run(tasks, budget=limit, dry_run=True) + + self.assertNotIn("reflect", backend.calls) + self.assertEqual(len(_budget_notes(outcome.report)), 1) + self.assertEqual(outcome.staging_dir, "") + self.assertFalse(os.path.exists(self._cfg().state_path)) + + def test_spent_before_consolidation_reports_no_scores(self): + backend = CountedBeforeConsolidationBackend() + + _, outcome = self._run(_tasks(), backend=backend) + + self.assertEqual(backend.calls, []) + report = outcome.report + self.assertEqual(report.gate_action, "budget_exhausted") + self.assertFalse(report.accepted) + self.assertEqual(report.n_replayed, 0) + notes = _budget_notes(report) + self.assertEqual(len(notes), 1) + self.assertIn("skipped: consolidation", notes[0]) + held_out = [ + line for line in self._report_md(outcome).splitlines() + if line.startswith("- held-out score:") + ] + self.assertEqual(len(held_out), 1) + self.assertNotIn("0.000", held_out[0]) + (night,) = self._state_history() + self.assertNotIn("baseline", night) + self.assertNotIn("candidate", night) + with open( + os.path.join(outcome.staging_dir, "diagnostics.json"), encoding="utf-8" + ) as handle: + diagnostics = json.load(handle) + self.assertIsNone(diagnostics["baseline_score"]) + self.assertIsNone(diagnostics["candidate_score"]) + + def test_default_budget_leaves_an_ordinary_night_unchanged(self): + backend, outcome = self._run(_tasks()) + self.assertGreater(backend.tokens_used(), _replay_cost(_tasks())) + self.assertIn("reflect", backend.calls) + self.assertTrue(outcome.report.accepted) + self.assertFalse(_budget_notes(outcome.report)) + (night,) = self._state_history() + self.assertIn("baseline", night) + self.assertIn("candidate", night) + + def test_zero_and_null_mean_unlimited(self): + reference, _ = self._run(_tasks()) + for value in (None, 0, "0", 0.0): + with self.subTest(max_tokens_per_night=value): + backend, outcome = self._run(_tasks(), budget=value) + self.assertEqual(backend.calls, reference.calls) + self.assertTrue(outcome.report.accepted) + self.assertFalse(_budget_notes(outcome.report)) + + def test_budget_counts_only_tonight_spend_on_a_reused_backend(self): + tasks = _tasks() + backend = MeteredBackend() + backend._tokens = 10 * _replay_cost(tasks) # spent on an earlier night + + _, outcome = self._run( + tasks, backend=backend, budget=_replay_cost(tasks) + 1 + ) + + self.assertIn("reflect", backend.calls) + self.assertTrue(outcome.report.accepted) + + def test_held_out_test_replay_is_skipped_when_the_budget_is_spent(self): + tasks = _tasks(test_fraction=0.2) + n_test = sum(1 for t in tasks if t.split == "test") + self.assertGreater(n_test, 0) + limit = _replay_cost(tasks) + + backend, outcome = self._run(tasks, budget=limit) + + self.assertLessEqual(backend.tokens_used(), limit) + notes = _budget_notes(outcome.report) + self.assertEqual(len(notes), 1) + self.assertIn("held-out test replay", notes[0]) + events = read_events(os.path.join(outcome.staging_dir, "evidence.jsonl")) + skips = [e["step"] for e in events if (e["stage"], e["event"]) == ("budget", "skip")] + self.assertIn("held-out test replay", skips) + self.assertFalse( + [e for e in events if (e["stage"], e["event"]) == ("test", "held_out_score")] + ) + + def test_fanout_groups_do_not_start_once_the_budget_is_spent(self): + claude_home = os.path.join(self._home.name, ".claude") + group_docs = {} + for name in ("research-skill", "programming-skill"): + path = os.path.join(claude_home, "skills", name, "SKILL.md") + os.makedirs(os.path.dirname(path), exist_ok=True) + group_docs[name] = f"# {name}\n" + with open(path, "w", encoding="utf-8", newline="\n") as handle: + handle.write(group_docs[name]) + tasks = [replace(t, skill_hint="research-skill") for t in _tasks()] + tasks += [ + replace(t, id=f"prog-{t.id}", skill_hint="programming-skill") + for t in assign_splits(programmer_persona(), val_fraction=0.34, seed=1) + ] + + backend, outcome = self._run(tasks, budget=1, multi_skill_fanout=True) + + # The main replay had already started, so it finished; no reflect ran + # and no group replayed anything against its own skill document. + self.assertNotIn("reflect", backend.calls) + self.assertTrue(backend.attempt_skills) + for skill in backend.attempt_skills: + for doc in group_docs.values(): + self.assertNotIn(doc.strip(), skill) + rows = {row.skill_name: row for row in outcome.report.skill_groups} + self.assertEqual(set(rows), set(group_docs)) + for row in rows.values(): + self.assertEqual(row.status, SKIPPED) + self.assertIn("token budget", row.reason) + notes = _budget_notes(outcome.report) + self.assertEqual(len(notes), 1) + self.assertIn("2 per-skill group(s)", notes[0]) + events = read_events(os.path.join(outcome.staging_dir, "evidence.jsonl")) + skipped_groups = { + e.get("skill") for e in events + if (e["stage"], e["event"]) == ("budget", "skip") + and e["step"] == "per-skill group" + } + self.assertEqual(skipped_groups, set(group_docs)) + + def test_invalid_budget_is_rejected_before_any_model_call(self): + for value in (-1, "lots", 0.5, True, float("inf")): + with self.subTest(max_tokens_per_night=value): + backend = MeteredBackend() + with self.assertRaisesRegex(ValueError, "max_tokens_per_night"): + self._run(_tasks(), backend=backend, budget=value) + self.assertEqual(backend.calls, []) + + def test_infinity_from_a_json_config_is_rejected(self): + value = json.loads('{"max_tokens_per_night": Infinity}')["max_tokens_per_night"] + with self.assertRaisesRegex(ValueError, "max_tokens_per_night"): + self._run(_tasks(), budget=value) + + +class TokenBudgetHandoffTest(unittest.TestCase): + """A handoff round never stops on a count that includes unanswered prompts.""" + + def test_cli_round_with_unanswered_prompts_pauses_and_persists_nothing(self): + from skillopt_sleep.__main__ import main + from skillopt_sleep.tasks_file import make_tasks_payload, write_tasks_file + + with tempfile.TemporaryDirectory() as proj, \ + tempfile.TemporaryDirectory() as home: + config_path = os.path.join(home, "config.json") + with open(config_path, "w", encoding="utf-8") as handle: + json.dump({"max_tokens_per_night": 1}, handle) + tasks_path = os.path.join(proj, "tasks.json") + payload = make_tasks_payload(_arithmetic_tasks(), project=proj) + payload["reviewed"] = True + write_tasks_file(tasks_path, payload) + claude_home = os.path.join(home, ".claude") + + with mock.patch( + "skillopt_sleep.config._user_config_path", return_value=config_path + ): + rc = main([ + "run", "--backend", "handoff", + "--project", proj, + "--claude-home", claude_home, + "--tasks-file", tasks_path, + ]) + state_path = load_config( + invoked_project=proj, claude_home=claude_home + ).state_path + + self.assertEqual(rc, 3) + hdir = os.path.join(proj, ".skillopt-sleep-handoff") + self.assertTrue(os.path.exists(os.path.join(hdir, "PROMPTS.md"))) + self.assertFalse(os.path.exists(state_path)) + staged = glob.glob( + os.path.join(proj, ".skillopt-sleep", "staging", "*", "report.md") + ) + self.assertEqual(staged, []) + + def test_budget_is_applied_once_the_round_is_answered(self): + with tempfile.TemporaryDirectory() as proj, \ + tempfile.TemporaryDirectory() as home: + hdir = os.path.join(proj, ".skillopt-sleep-handoff") + cfg = load_config( + invoked_project=proj, projects="invoked", backend="handoff", + claude_home=os.path.join(home, ".claude"), + state_dir=os.path.join(home, "state"), auto_adopt=False, + ) + cfg.data["max_tokens_per_night"] = 1 + + first = HandoffBackend(handoff_dir=hdir) + with self.assertRaises(PendingCalls): + run_sleep_cycle(cfg, seed_tasks=_arithmetic_tasks(), backend=first) + self.assertTrue(first.pending) + self.assertFalse(os.path.exists(cfg.state_path)) + for key in first.pending: + with open(first.answer_path(key), "w", encoding="utf-8") as handle: + handle.write("cannot say") + + second = HandoffBackend(handoff_dir=hdir) + outcome = run_sleep_cycle( + cfg, seed_tasks=_arithmetic_tasks(), backend=second + ) + + self.assertEqual(second.pending, {}) + self.assertEqual(outcome.report.night, 1) + notes = _budget_notes(outcome.report) + self.assertEqual(len(notes), 1) + self.assertIn("skill edits", notes[0]) + self.assertTrue(os.path.exists(cfg.state_path)) + + +class TokenBudgetConsolidateTest(unittest.TestCase): + def test_spent_budget_skips_the_whole_consolidation(self): + backend = MeteredBackend() + result = consolidate( + backend, _tasks(), "", "", budget_exhausted=lambda: True + ) + self.assertEqual(backend.calls, []) + self.assertFalse(result.accepted) + self.assertEqual(result.gate_action, "budget_exhausted") + self.assertEqual(result.budget_skipped, ["consolidation"]) + self.assertEqual(result.applied_edits, []) + + def test_no_budget_keeps_the_previous_behavior(self): + plain = MeteredBackend() + budgeted = MeteredBackend() + before = consolidate(plain, _tasks(), "", "") + after = consolidate(budgeted, _tasks(), "", "", budget_exhausted=lambda: False) + self.assertEqual(plain.calls, budgeted.calls) + self.assertEqual(before.gate_action, after.gate_action) + self.assertEqual(before.new_skill, after.new_skill) + self.assertEqual(after.budget_skipped, []) + + +class TokenBudgetGroupsTest(unittest.TestCase): + def test_groups_not_started_are_skipped_once_the_budget_is_spent(self): + backend = MeteredBackend() + tasks = _tasks() + groups = [ + SkillGroup("first-skill", "", list(tasks)), + SkillGroup("second-skill", "", list(tasks)), + ] + + outcomes = consolidate_groups( + backend, + groups, + budget_exhausted=lambda: backend.tokens_used() > 0, + ) + + first = outcomes["first-skill"] + self.assertEqual(first.status, CONSOLIDATED) + self.assertIn("skipped skill edits", first.reason) + second = outcomes["second-skill"] + self.assertEqual(second.status, SKIPPED) + self.assertIn("token budget", second.reason) + self.assertEqual(second.n_tasks, len(tasks)) + self.assertIsNone(second.result) + + +if __name__ == "__main__": + unittest.main() From 0322fbe1e1e4eb2909d6bbc76c5b16015203d9b8 Mon Sep 17 00:00:00 2001 From: Yashvardhan Thanvi Date: Sun, 18 Oct 2026 18:27:44 +0530 Subject: [PATCH 2/2] fix(sleep): skip the final validation replay once the budget is spent The budget check assumed the final validation replay was free because it re-scores documents an earlier replay already scored. That holds for cached text replies, but tool replay is never cached and a failed or timed out call is not kept in the cache, so on those routes the final replay started new provider calls after both edit steps had been skipped. Check the budget before the final validation replay as well. When it is spent, decide from the validation replay that already scored the final skill and memory (the baseline replay, or the trial replay of the last kept edit) instead of starting a new one, and list the skipped replay in the report note and the budget skip events. Nights that never reach the budget still run the fresh final replay as before. The new tests drive ClaudeCliBackend through its real tool and text routes with a synthetic process standing in for the claude CLI and count every process it starts, for answered and timed out calls. --- docs/sleep/README.md | 20 +++- skillopt_sleep/consolidate.py | 24 ++-- tests/test_sleep_token_budget.py | 197 ++++++++++++++++++++++++++++++- 3 files changed, 226 insertions(+), 15 deletions(-) diff --git a/docs/sleep/README.md b/docs/sleep/README.md index c8cf8f8e9..b803d8f97 100644 --- a/docs/sleep/README.md +++ b/docs/sleep/README.md @@ -349,10 +349,18 @@ Deterministic proof (no API key): `max_tokens_per_night` in `~/.skillopt-sleep/config.json` (default `400000`) stops a `run` or `dry-run` from starting more model work once tonight's spend reaches it. The cycle checks the budget before the consolidation, before its -skill and memory edit steps, before the held-out test replay, and before each -per-skill fan-out group. Steps that did not start are listed in a note in -`report.md`, for example -`max_tokens_per_night reached (431870 of 400000 tokens); skipped: memory edits`. +skill and memory edit steps, before its final validation replay, before the +held-out test replay, and before each per-skill fan-out group. Steps that did +not start are listed in a note in `report.md`, for example +`max_tokens_per_night reached (431870 of 400000 tokens); skipped: memory edits, final validation replay`. + +When the final validation replay is skipped, the gate decides from the +validation replay that already scored the same skill and memory: the baseline +replay when no edit was kept, or the trial replay of the last edit that was +kept. That replay is complete, so no decision rests on partial evidence, and +it is not repeated. This matters for replays the response cache does not +serve, such as tool replay or a call that failed or timed out, which would +otherwise start the model again. ```json {"max_tokens_per_night": 200000} @@ -366,8 +374,8 @@ This is not a hard cap. A step that has started always finishes, so a night can end above the budget by the cost of one step: the first replay of the validation and training tasks, one edit step (a training replay, a reflect call and a validation replay of the candidate; with `dream_rollouts` set to K, -K attempts per training task), or the held-out test replay. Each fan-out group -that starts runs its own first replay. +K attempts per training task), the final validation replay, or the held-out +test replay. Each fan-out group that starts runs its own first replay. Spend comes from the backend's token counter, the one the report header prints as `tokens used`. The header shows the backend's running total and the note diff --git a/skillopt_sleep/consolidate.py b/skillopt_sleep/consolidate.py index 77e1e6a3e..3b63520a7 100644 --- a/skillopt_sleep/consolidate.py +++ b/skillopt_sleep/consolidate.py @@ -209,11 +209,13 @@ def consolidate( Skill and memory are evolved in sequence (skill first if both enabled). - ``budget_exhausted`` is checked before the epoch starts and before each - evolve target. Once it returns True, the steps that have not started are - skipped and listed in ``budget_skipped``. A started step always finishes, - so every gate decision still rests on complete evidence; the final replay - only re-scores documents an earlier replay already scored. + ``budget_exhausted`` is checked before the epoch starts, before each + evolve target and before the final validation replay. Once it returns + True, the steps that have not started are skipped and listed in + ``budget_skipped``. A started step always finishes, so every gate decision + still rests on complete evidence. When the final replay is skipped, the + decision uses the validation replay that already scored the final skill + and memory: the baseline replay, or the trial replay of the last kept edit. """ from skillopt_sleep import evidence as evlog ev = evlog.get(backend) @@ -415,8 +417,16 @@ def _gate_apply(doc: str, edits: List[EditRecord], which: str) -> str: base_gate_score = 0.0 else: # scored on the VAL slice (the gate reference) - evlog.set_phase(backend, "final_val") - final_pairs = replay_batch(backend, val_tasks, cand_skill, cand_memory) + if val_tasks and _budget_spent("final validation replay"): + # Replaying is not free on every route: tool replay and failed calls + # are not served from the response cache. current_pairs already + # scores exactly (cand_skill, cand_memory) on val_tasks (the + # baseline replay, or the trial of the last edit kept), so decide + # from that complete replay instead of starting a new one. + final_pairs = current_pairs + else: + evlog.set_phase(backend, "final_val") + final_pairs = replay_batch(backend, val_tasks, cand_skill, cand_memory) final_hard, final_soft = aggregate_scores(final_pairs) final_score = select_gate_score(final_hard, final_soft, gate_metric, gate_mixed_weight) base_gate_score = select_gate_score(base_hard, base_soft, gate_metric, gate_mixed_weight) diff --git a/tests/test_sleep_token_budget.py b/tests/test_sleep_token_budget.py index cd703aa15..fae389381 100644 --- a/tests/test_sleep_token_budget.py +++ b/tests/test_sleep_token_budget.py @@ -8,12 +8,13 @@ import glob import json import os +import subprocess import tempfile import unittest from dataclasses import replace from unittest import mock -from skillopt_sleep.backend import MockBackend +from skillopt_sleep.backend import ClaudeCliBackend, MockBackend from skillopt_sleep.config import load_config from skillopt_sleep.consolidate import consolidate from skillopt_sleep.cycle import run_sleep_cycle @@ -112,6 +113,68 @@ def _arithmetic_tasks(): return assign_splits(tasks, holdout_fraction=0.34, seed=42) +def _tool_tasks(): + """Tasks whose rule judge needs a real tool call (the tool replay route).""" + tasks = [ + TaskRecord( + id=f"look{i}", project="/p", + intent=f"Find note {i} and summarize it.", + reference_kind="rule", + judge={"kind": "rule", "checks": [{"op": "tool_called", "arg": "search"}]}, + ) + for i in range(1, 7) + ] + return assign_splits(tasks, val_fraction=0.34, seed=42) + + +_REAL_RUN = subprocess.run +_SEARCH_RULE = "Always run ./search before answering." + + +class SyntheticClaudeProcess: + """Stands in for the external ``claude`` process; no model is called. + + Records every process the backend spawns as ``(phase, route)``, where the + route is ``tools`` for ``attempt_with_tools`` and ``text`` for ``_call``. + Replies with a fixed answer, or raises ``TimeoutExpired`` like a process + that ran past the backend's timeout. A reflect call (text route during the + train phase) proposes ``_SEARCH_RULE``; a tool replay whose prompt carries + that rule logs a call to the ``search`` shim, so its rule judge passes. + """ + + def __init__(self, backend: ClaudeCliBackend, *, timeout: bool = False) -> None: + self.backend = backend + self.timeout = timeout + self.calls: list[tuple[str, str]] = [] + + def __call__(self, cmd, **kwargs): + if not cmd or cmd[0] != self.backend.claude_path: + return _REAL_RUN(cmd, **kwargs) + route = "tools" if "--allowedTools" in cmd else "text" + phase = getattr(self.backend, "evidence_phase", "") + self.calls.append((phase, route)) + if self.timeout: + raise subprocess.TimeoutExpired(cmd, kwargs.get("timeout")) + prompt = kwargs.get("input") or "" + if route == "text" and phase == "train": + reply = json.dumps([{"op": "add", "content": _SEARCH_RULE}]) + else: + reply = "Here is the summary." + if route == "tools" and _SEARCH_RULE in prompt: + with open(os.path.join(kwargs["cwd"], "_tool_calls.log"), "a") as log: + log.write("search\n") + return subprocess.CompletedProcess(cmd, 0, stdout=reply, stderr="") + + def phases(self) -> list[str]: + return [phase for phase, _route in self.calls] + + +def _synthetic_claude(timeout: bool = False): + backend = ClaudeCliBackend(claude_path="synthetic-claude-for-tests") + process = SyntheticClaudeProcess(backend, timeout=timeout) + return backend, process + + class TokenBudgetCycleTest(unittest.TestCase): def setUp(self) -> None: self._proj = tempfile.TemporaryDirectory() @@ -176,7 +239,10 @@ def test_night_stops_scheduling_work_once_the_budget_is_spent(self): events = read_events(os.path.join(outcome.staging_dir, "evidence.jsonl")) end = [e for e in events if (e["stage"], e["event"]) == ("cycle", "end")] self.assertEqual(end[-1]["max_tokens_per_night"], limit) - self.assertEqual(end[-1]["budget_skipped"], ["skill edits", "memory edits"]) + self.assertEqual( + end[-1]["budget_skipped"], + ["skill edits", "memory edits", "final validation replay"], + ) def test_one_token_of_headroom_lets_the_next_step_run(self): tasks = _tasks() @@ -463,5 +529,132 @@ def test_groups_not_started_are_skipped_once_the_budget_is_spent(self): self.assertIsNone(second.result) +class TokenBudgetProviderCallTest(unittest.TestCase): + """Count the provider processes a budget-stopped night really starts. + + These drive ``ClaudeCliBackend`` through its own replay routes, with + ``subprocess.run`` answered by ``SyntheticClaudeProcess``. Replies the + response cache does not serve (tool replay, a timed out call) are visible + as new processes, which a caching mock cannot show. + """ + + def setUp(self) -> None: + self._proj = tempfile.TemporaryDirectory() + self._home = tempfile.TemporaryDirectory() + self.addCleanup(self._proj.cleanup) + self.addCleanup(self._home.cleanup) + patcher = mock.patch.dict(os.environ, {"SKILLOPT_SLEEP_WORKERS": "1"}) + patcher.start() + self.addCleanup(patcher.stop) + + def _night(self, tasks, *, timeout: bool): + cfg = load_config( + invoked_project=self._proj.name, + projects="invoked", + backend="claude", + claude_home=os.path.join(self._home.name, ".claude"), + state_dir=os.path.join(self._home.name, "state"), + auto_adopt=False, + ) + # Any first call spends this, so only the initial replay may run. + cfg.data["max_tokens_per_night"] = 1 + backend, process = _synthetic_claude(timeout) + with mock.patch("skillopt_sleep.backend.subprocess.run", side_effect=process): + outcome = run_sleep_cycle(cfg, seed_tasks=tasks, backend=backend) + return process, outcome + + def _assert_only_the_initial_replay_ran(self, tasks, process, outcome) -> None: + n_val = sum(1 for t in tasks if t.split == "val") + n_train = sum(1 for t in tasks if t.split == "train") + self.assertGreater(n_val, 0) + phases = process.phases() + self.assertEqual(phases.count("final_val"), 0) + self.assertEqual( + phases, ["baseline_val"] * n_val + ["train"] * n_train + ) + report = outcome.report + self.assertFalse(report.accepted) + notes = _budget_notes(report) + self.assertEqual(len(notes), 1) + self.assertIn( + "skipped: skill edits, memory edits, final validation replay", notes[0] + ) + events = read_events(os.path.join(outcome.staging_dir, "evidence.jsonl")) + self.assertFalse([ + e for e in events + if (e["stage"], e["event"]) == ("replay", "result") + and e.get("phase") == "final_val" + ]) + skips = [e["step"] for e in events if (e["stage"], e["event"]) == ("budget", "skip")] + self.assertIn("final validation replay", skips) + decision = [e for e in events if (e["stage"], e["event"]) == ("gate", "decision")] + self.assertEqual(len(decision), 1) + self.assertEqual( + decision[0]["candidate_score"], decision[0]["baseline_score"] + ) + + def test_tool_route_answers_make_no_final_validation_call(self): + tasks = _tool_tasks() + process, outcome = self._night(tasks, timeout=False) + self.assertEqual({route for _phase, route in process.calls}, {"tools"}) + self._assert_only_the_initial_replay_ran(tasks, process, outcome) + + def test_tool_route_timeouts_make_no_final_validation_call(self): + tasks = _tool_tasks() + process, outcome = self._night(tasks, timeout=True) + self.assertEqual({route for _phase, route in process.calls}, {"tools"}) + self._assert_only_the_initial_replay_ran(tasks, process, outcome) + + def test_text_route_timeouts_make_no_final_validation_call(self): + # An empty reply from a failed call is not kept in the response cache, + # so re-scoring the same documents would start the process again. + tasks = _arithmetic_tasks() + process, outcome = self._night(tasks, timeout=True) + self.assertEqual({route for _phase, route in process.calls}, {"text"}) + self._assert_only_the_initial_replay_ran(tasks, process, outcome) + + def test_text_route_answers_stay_free_cache_hits(self): + # Control: here the final replay would have been served from cache. + tasks = _arithmetic_tasks() + process, outcome = self._night(tasks, timeout=False) + self._assert_only_the_initial_replay_ran(tasks, process, outcome) + + +class TokenBudgetFinalScoringTest(unittest.TestCase): + """With the budget spent, final scoring reuses the replay of the kept edit.""" + + def test_kept_skill_edit_is_scored_from_its_own_validation_replay(self): + tasks = _tool_tasks() + n_val = sum(1 for t in tasks if t.split == "val") + n_train = sum(1 for t in tasks if t.split == "train") + backend, process = _synthetic_claude() + + def spent() -> bool: + # The budget runs out once the skill step's reflect call started. + return ("train", "text") in process.calls + + with mock.patch.dict(os.environ, {"SKILLOPT_SLEEP_WORKERS": "1"}), \ + mock.patch("skillopt_sleep.backend.subprocess.run", side_effect=process): + result = consolidate( + backend, tasks, "# skill\n", "", budget_exhausted=spent + ) + + self.assertEqual( + process.calls, + [("baseline_val", "tools")] * n_val + + [("train", "tools")] * n_train + + [("train", "text")] + + [("gate_trial:skill", "tools")] * n_val, + ) + self.assertEqual(result.budget_skipped, ["memory edits", "final validation replay"]) + self.assertTrue(result.accepted) + self.assertIn(_SEARCH_RULE, result.new_skill) + self.assertEqual(result.candidate_score, 1.0) + trial, final = result.gate_trials + self.assertEqual(trial["target"], "skill") + self.assertEqual(final["target"], "final") + self.assertEqual(final["candidate_score"], trial["candidate_score"]) + + if __name__ == "__main__": unittest.main()