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()