diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 4e82bf4..233fe55 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -10,7 +10,7 @@ "name": "nexum", "source": "./", "description": "Context-token and model-cost optimization for Claude Code.", - "version": "0.3.0", + "version": "0.4.0", "author": { "name": "Rahul Tyagi", "email": "rahul.1992.tyagi@gmail.com" diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json index d29d3e6..a911187 100644 --- a/.claude-plugin/plugin.json +++ b/.claude-plugin/plugin.json @@ -5,5 +5,5 @@ }, "description": "Context-token and model-cost optimization for Claude Code.", "name": "nexum", - "version": "0.3.0" + "version": "0.4.0" } diff --git a/CHANGELOG.md b/CHANGELOG.md index 53a9906..15e3559 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,19 @@ All notable changes to nexum are documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.4.0] - 2026-06-18 +### Added +- **Estimated per-tier usage recording in `/nx-build`** so the cost report breakdown is populated with token estimates per execution tier, giving visibility into where the session's cost is incurred. +- **Per-repo routing calibration** that biases `/nx-plan` from historical step outcomes, routing subsequent steps based on learned success rates and execution patterns rather than static tier assignments. +- **Throttled SessionStart audit nudge** (`scripts/audit_nudge.py`, SessionStart hook) that checks ignore-config rot on session start without spamming the user, surfacing a one-line reminder when the ignore config has actionable findings (missing ignore file, unignored noise dirs, or large/binary files). Throttled to once per repo per `audit_nudge_throttle_hours`. Config: `audit_nudge_enabled` (default true), `audit_nudge_throttle_hours` (default 24). + +### Changed +- **`max_steps_per_dispatch` default lowered 6 → 4.** A 6-step grouped Sonnet dispatch stalled the stream watchdog mid-batch (600s no-progress) and lost the whole batch; smaller batches bound both the per-dispatch context and a single failure's blast radius. +- **Size-aware dispatch batching.** `/nx-build` now partitions a route tier with `plan_preview.py --plan … --indices …`, which estimates each step's context load (per-step base + declared-file bytes ÷ 4) and bounds each sub-batch by **both** `max_dispatch_context_tokens` (default 50000) and the `max_steps_per_dispatch` count cap — so a few large-file steps split off while a single over-budget step still dispatches alone (steps are never split). New `store.partition_steps_by_size`; config `max_dispatch_context_tokens`, `dispatch_step_base_tokens`. The count-only `store.py plan-batches` helper remains for when no plan file is available. + +### Fixed +- **predup freshness guard.** A recorded `tool_calls` row only proves an output was injected once, not that it is still in the live context — subagents share the parent's DB, and compaction/resume evicts output while the row persists, so predup could deny a legitimate read whose content was no longer (or never) in context. predup now lets a call through when its prior row is older than `predup_max_age_seconds` (default 900; 0 restores the prior ever-recorded behaviour). + ## [0.3.0] - 2026-06-18 ### Added - **Pre-emptive dedup hook** (`scripts/predup.py`, PreToolUse). Denies an identical repeated `Read`, `Grep`, or `Glob` call that already ran in the same session, with an mtime guard for `Read` to allow re-reads of changed files. Unlike the PostToolUse dedup (currently inert), a PreToolUse `deny` is actually honored by Claude Code, so the avoided re-injection is a real saving — the `saved` figure in the status line now moves on deduped calls. Config: `predup_enabled` (default true), `predup_decision` (`deny` | `ask`, default `deny`), `predup_bash_readonly` (extend coverage to read-only Bash; default false). diff --git a/commands/nx-build.md b/commands/nx-build.md index 41d0774..a99514c 100644 --- a/commands/nx-build.md +++ b/commands/nx-build.md @@ -64,13 +64,15 @@ Grouping keeps each model's prompt prefix stable, maximising cache hits. - **`group` (default):** Send the route group to **one** executor dispatch. The executor reads the shared spec and source files once and reuses them across every step in the group — one warm context, one cached prefix, instead of N cold starts that each re-derive the same context. The executor returns a per-step result list. - **`step`:** Send one dispatch per step. More isolation (a mid-batch failure can't affect siblings), but pays a cold start and re-reads context for every step. Use only when steps in a tier are large or risky enough that isolation outweighs the re-derivation cost. -**Cap the batch size — compute the split, don't eyeball it.** Under `group`, a single oversized dispatch can overflow the executor's own context (forcing a mid-batch compaction that re-derives everything and wipes the token savings) and widens the blast radius if one step fails. Do not judge "too big" by eye — ask the helper for the deterministic, order-preserving partition of the tier's step indices (in execution order): +**Cap the batch size — compute the split, don't eyeball it.** Under `group`, a single oversized dispatch can overflow the executor's own context (forcing a mid-batch compaction that re-derives everything and wipes the token savings, or — as observed — stalling the stream watchdog and losing the whole batch) and widens the blast radius if one step fails. Do not judge "too big" by eye — ask the helper for the deterministic, order-preserving partition of the tier's step indices (in execution order): ``` -python3 ${CLAUDE_PLUGIN_ROOT}/scripts/store.py plan-batches --indices "" +python3 ${CLAUDE_PLUGIN_ROOT}/scripts/plan_preview.py --plan --indices "" --root ``` -It returns a JSON array of sub-batches (e.g. `[[0,1,2,3,4,5],[6,7]]`), each at most `max_steps_per_dispatch` (config, default 6; `0` disables the cap). Dispatch the sub-batches **sequentially**; each reuses the same shared-context prefix, so the cache stays warm while each dispatch is bounded. Order is preserved, so a dependent step never runs before its prerequisite. +This is **size-aware**: it estimates each step's context load (a per-step base + the byte size of its declared `files` ÷ 4) and packs the tier into sub-batches bounded by **both** `max_dispatch_context_tokens` (config, default 50000) and the `max_steps_per_dispatch` count cap (config, default 4). So a tier of a few large-file steps splits more aggressively than a tier of small ones, while a single step over the token budget still dispatches alone (a step is never split). It returns a JSON array of sub-batches (e.g. `[[1,2,3,4],[5,7]]`). Dispatch the sub-batches **sequentially**; each reuses the same shared-context prefix, so the cache stays warm while each dispatch is bounded. Order is preserved, so a dependent step never runs before its prerequisite. + +(The older count-only helper `store.py plan-batches --indices "…"` still exists and caps by `max_steps_per_dispatch` alone — use it only if a plan file isn't available to size against.) ## 4. Skip the spawn when the tier already matches the session model @@ -136,6 +138,31 @@ When a previously-failed step later passes, overwrite it with `--status done` as If `orchestrator_resume_enabled` is false, skip all ledger writes. +**Immediately after recording the ledger for each step**, also record estimated usage and (when enabled) calibration counters: + +**Usage recording (always; read heuristics from `store.py config`):** + +``` +python3 ${CLAUDE_PLUGIN_ROOT}/scripts/store.py record-usage \ + --session \ + --model \ + --input-tok \ + --output-tok +``` + +Read `plan_preview_input_tok_per_step` and `plan_preview_output_tok_per_step` from `store.py config`. These are per-step heuristic estimates; multiply by the number of steps in the dispatch to get the total. An escalated step counts against the higher tier (the one that actually ran), not the original dispatched tier. + +**Calibration recording (when `route_calib_enabled` from `store.py config` is true):** + +``` +python3 ${CLAUDE_PLUGIN_ROOT}/scripts/store.py calib-record \ + --repo \ + --route \ + --dispatched 1 +``` + +Append `--passed-first-try 1` when the step passed on the first attempt with no escalation. Append `--escalated 1` when the step required escalation to a higher tier. Usage recording is always performed; calibration is skipped only when `route_calib_enabled` is false. + ## 7. Retry and escalation ladder On a FAIL: diff --git a/commands/nx-plan.md b/commands/nx-plan.md index a08d90b..40dd421 100644 --- a/commands/nx-plan.md +++ b/commands/nx-plan.md @@ -30,6 +30,14 @@ Assign every step exactly one route. Use the cheapest tier that is sufficient: When in doubt, prefer **standard** over **mechanical** — a false mechanical that fails wastes more time than a conservative standard. +**Calibration nudge (per-repo history).** Before finalising routes, consult this repo's historical outcomes when `route_calib_enabled` (from `python3 ${CLAUDE_PLUGIN_ROOT}/scripts/store.py config`) is true. Run: + +``` +python3 ${CLAUDE_PLUGIN_ROOT}/scripts/store.py calib-list --repo +``` + +For any route whose `dispatched >= route_calib_min_samples` (config, default 5) **and** `passed_first_try / dispatched < route_calib_min_success_ratio` (config, default 0.6), nudge the steps you would have assigned to that route **up one tier** (mechanical→standard, standard→needs-strong) and note the calibration reason in that step's `objective` (e.g. "routed up from mechanical: this repo's mechanical steps pass first-try only 40% of the time"). This is **advisory and one-directional**: never downgrade a route, and never move a `needs-strong` step. When calibration is disabled, the repo has no rows, or a route's sample count is below `route_calib_min_samples`, route by the static rubric above unchanged. + **Dependency-vs-tier rule.** The implementer executes tiers in the order mechanical → standard → needs-strong, so a step must never be routed to a tier that runs *before* a step it depends on. If step B consumes step A's output, B's tier must be the same as or costlier than A's. In particular: a test step that exercises code written in a `standard` step is itself `standard` (not `mechanical`); and a final full-suite / verification step takes the highest tier of any step it validates (or, if you keep it `mechanical`, state explicitly in its `objective` that it runs last). Ordering steps so prerequisites come first is not enough — the tier assignment must also respect the dependency, or the cheaper tier will run first and fail. ## 3. Step schema diff --git a/hooks/hooks.json b/hooks/hooks.json index c217e38..399314d 100644 --- a/hooks/hooks.json +++ b/hooks/hooks.json @@ -43,6 +43,15 @@ "type": "command" } ] + }, + { + "hooks": [ + { + "command": "python3 ${CLAUDE_PLUGIN_ROOT}/scripts/audit_nudge.py", + "timeout": 10, + "type": "command" + } + ] } ], "UserPromptSubmit": [ diff --git a/scripts/audit_nudge.py b/scripts/audit_nudge.py new file mode 100644 index 0000000..44085b0 --- /dev/null +++ b/scripts/audit_nudge.py @@ -0,0 +1,140 @@ +""" +audit_nudge.py — Nexum SessionStart ignore-config nudge hook. + +When the project's ignore config has actionable findings (missing ignore file, +unignored noise dirs, or large/binary files), surfaces a one-line +"/nx-audit" hint — throttled to once per repo per audit_nudge_throttle_hours. + +Hook contract: + stdin → single JSON object (Claude Code SessionStart payload) + stdout → single JSON object (hint or {}) + exit 0 always (fail-open) +""" + +import json +import os +import sys +import time + +# --------------------------------------------------------------------------- +# sys.path: ensure scripts/ dir is importable as "import store" / "import audit" +# --------------------------------------------------------------------------- +_SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__)) +if _SCRIPTS_DIR not in sys.path: + sys.path.insert(0, _SCRIPTS_DIR) + +import store # noqa: E402 +import audit # noqa: E402 + + +def main() -> None: + """SessionStart hook entry point.""" + try: + # ---------------------------------------------------------------- + # 1. Parse stdin JSON + # ---------------------------------------------------------------- + try: + raw = sys.stdin.read() + data = json.loads(raw) + except Exception: + print("{}") + return + + if not isinstance(data, dict): + print("{}") + return + + # ---------------------------------------------------------------- + # 2. Check config gate + # ---------------------------------------------------------------- + cfg = store.get_config() + if not cfg.get("audit_nudge_enabled", True): + print("{}") + return + + # ---------------------------------------------------------------- + # 3. Skip continued sessions (already resumed or compacted) + # ---------------------------------------------------------------- + source = data.get("source") + if source in {"resume", "compact"}: + print("{}") + return + + # ---------------------------------------------------------------- + # 4. Throttle: once per repo per audit_nudge_throttle_hours + # ---------------------------------------------------------------- + cwd = data.get("cwd") or os.getcwd() + repo_key = store.current_repo(cwd) + throttle_hours = float(cfg.get("audit_nudge_throttle_hours", 24)) + + last_str = store.get_flag("_audit_nudge", repo_key) + now = time.time() + if last_str is not None: + try: + last_ts = float(last_str) + if now - last_ts < throttle_hours * 3600: + print("{}") + return + except ValueError: + pass + + # ---------------------------------------------------------------- + # 5. Run audit + # ---------------------------------------------------------------- + try: + findings = audit.run_audit(cwd) + except Exception: + print("{}") + return + + # ---------------------------------------------------------------- + # 6. Check for actionable findings + # ---------------------------------------------------------------- + missing = findings.get("missing_ignore", False) + noise_dirs = findings.get("unignored_noise_dirs") or [] + large_bin = findings.get("large_or_binary") or [] + + if not missing and not noise_dirs and not large_bin: + print("{}") + return + + # ---------------------------------------------------------------- + # 7. Build hint + # ---------------------------------------------------------------- + count = sum([bool(missing), bool(noise_dirs), bool(large_bin)]) + + if noise_dirs: + example = noise_dirs[0] + "/" + elif missing: + example = "no ignore file" + else: + # large_or_binary entries are tuples of (rel_path, size, is_binary) + first = large_bin[0] + example = first[0] if isinstance(first, (tuple, list)) else str(first) + + hint = ( + f"[nexum] Ignore-config has {count} issue(s) " + f"(e.g. {example}); run /nx-audit." + ) + + out = { + "hookSpecificOutput": { + "hookEventName": "SessionStart", + "additionalContext": hint, + }, + "systemMessage": hint, + } + + # ---------------------------------------------------------------- + # 8. Record throttle timestamp and emit + # ---------------------------------------------------------------- + store.set_flag("_audit_nudge", repo_key, str(now)) + print(json.dumps(out, sort_keys=True)) + + except Exception: + # Fail-open: never crash the Claude Code session + print("{}") + + +if __name__ == "__main__": + main() diff --git a/scripts/cost_report.py b/scripts/cost_report.py index 2004e89..45c01d4 100644 --- a/scripts/cost_report.py +++ b/scripts/cost_report.py @@ -122,13 +122,16 @@ def build_report(rows: list) -> str: saved = total_baseline - total_actual lines = [] - lines.append("[nexum] Cost report") - lines.append("=" * 48) + lines.append( + "[nexum] Cost report " + "(per-tier breakdown is ESTIMATED — heuristic attribution per dispatch)" + ) + lines.append("=" * 72) lines.append(f" Actual cost: ${total_actual:>10.4f}") lines.append(f" All-opus baseline: ${total_baseline:>10.4f}") lines.append(f" Saved vs opus: ${saved:>10.4f}") lines.append("") - lines.append("Per-model breakdown:") + lines.append("Per-model breakdown (ESTIMATED — heuristic attribution recorded by /nx-build):") lines.append(f" {'Model':<20} {'Input tok':>12} {'Output tok':>12} " f"{'Cache-R tok':>12} {'Actual $':>10} {'Baseline $':>10}") lines.append(" " + "-" * 80) @@ -147,6 +150,8 @@ def build_report(rows: list) -> str: else: lines.append("[nexum] Note: token yield needs shipped-token tagging " "(no shipped-token field in v1 usage rows).") + lines.append("[nexum] The authoritative, cache-accurate total is the metered " + "section below.") return "\n".join(lines) diff --git a/scripts/plan_preview.py b/scripts/plan_preview.py index 0173f19..aaa25f1 100644 --- a/scripts/plan_preview.py +++ b/scripts/plan_preview.py @@ -12,6 +12,7 @@ """ import argparse +import json import os import re import sys @@ -41,13 +42,25 @@ # Plan parser # --------------------------------------------------------------------------- +def _parse_files_field(value: str) -> List[str]: + """Split a step's ``files:`` value into a list of paths. + + Returns [] for ``none`` / empty. Splits on commas and trims whitespace. + """ + value = (value or "").strip() + if not value or value.lower() == "none": + return [] + return [p.strip() for p in value.split(",") if p.strip()] + + def parse_plan_steps(text: str) -> List[dict]: - """Parse a nexum plan markdown and return steps with index, title, and route. + """Parse a nexum plan markdown and return steps with index, title, route, files. Scans for ``### Step : `` headers and reads the first - ``- route: <value>`` line that follows. Returns a list of - ``{"index": int, "title": str, "route": str}`` in file order. - Only real ``### Step`` headers are parsed; route-rubric examples + ``- route: <value>`` and ``- files: <value>`` lines that follow within the + step block. Returns a list of + ``{"index": int, "title": str, "route": str, "files": [str, ...]}`` in file + order. Only real ``### Step`` headers are parsed; route-rubric examples inside ``|``-quoted table cells are ignored. """ steps = [] @@ -60,26 +73,49 @@ def parse_plan_steps(text: str) -> List[dict]: if m: index = int(m.group(1)) title = m.group(2).strip() - # Scan forward for the route line (first occurrence within the block) + # Scan forward for the route + files lines within this step block. route = "standard" + files: List[str] = [] for j in range(i + 1, min(i + 30, len(lines))): next_line = lines[j] # Stop at the next ### header (new step) if re.match(r'^###\s+Step\s+\d+:', next_line): break - # Only match a plain list item (not inside a table cell) + # Only match plain list items (not inside a table cell) rm = re.match(r'^- route:\s+(\S+)', next_line) if rm: raw_route = rm.group(1).rstrip(',;|') - # Map to known routes; default to "standard" if raw_route in ROUTE_TIER: route = raw_route - break - steps.append({"index": index, "title": title, "route": route}) + continue + fm = re.match(r'^- files:\s+(.+)', next_line) + if fm: + files = _parse_files_field(fm.group(1)) + continue + steps.append({"index": index, "title": title, "route": route, "files": files}) i += 1 return steps +def estimate_step_tokens(files: List[str], root: str, base: int) -> int: + """Estimate a step's context-token load: *base* + (file bytes ÷ 4) summed + over the step's declared files that exist on disk. + + Missing files (to be created by the step) contribute only the base. Paths + may be absolute or relative to *root*. Fail-open per file (an unstattable + path is skipped). + """ + total = int(base) + for f in files: + p = f if os.path.isabs(f) else os.path.join(root, f) + try: + if os.path.isfile(p): + total += os.path.getsize(p) // 4 + except OSError: + pass + return total + + # --------------------------------------------------------------------------- # Preview builder # --------------------------------------------------------------------------- @@ -176,20 +212,75 @@ def main() -> None: ) parser.add_argument("--plan", required=True, help="Path to the nexum plan file.") parser.add_argument("--session", default=None, help="Session ID (unused, accepted for parity).") + parser.add_argument( + "--indices", default=None, + help="Comma-separated step indices (in execution order) to partition " + "into size-aware sub-batches. When given, prints a JSON array of " + "sub-batches instead of the cost preview.", + ) + parser.add_argument( + "--root", default=None, + help="Repo root for resolving step file paths (default: cwd).", + ) args = parser.parse_args() try: with open(args.plan, "r", encoding="utf-8", errors="replace") as fh: text = fh.read() except Exception: - print(f"[nexum] No plan file at {args.plan}.") + if args.indices is not None: + # Batch mode must still emit valid JSON for the caller. + print("[]") + else: + print(f"[nexum] No plan file at {args.plan}.") return cfg = store.get_config() - print(build_preview(parse_plan_steps(text), cfg)) + steps = parse_plan_steps(text) + + if args.indices is not None: + print(_build_batches(steps, args.indices, args.root, cfg)) + return + + print(build_preview(steps, cfg)) except Exception: # Fail-open - pass + if "--indices" in sys.argv: + print("[]") + + +def _build_batches(steps: list, indices_arg: str, root: str, cfg: dict) -> str: + """Return a JSON array of size-aware sub-batches for the requested indices. + + Estimates each requested step's context tokens (base + file bytes ÷ 4) and + packs them with store.partition_steps_by_size, bounded by + max_dispatch_context_tokens and max_steps_per_dispatch. Order is preserved. + Unknown indices fall back to base-only size so they still dispatch. + """ + root = root or os.getcwd() + by_index = {s["index"]: s for s in steps} + base = int(cfg.get("dispatch_step_base_tokens", 1500)) + max_size = int(cfg.get("max_dispatch_context_tokens", 50000)) + max_per = int(cfg.get("max_steps_per_dispatch", 4)) + + items: List[int] = [] + sizes: List[int] = [] + for tok in (indices_arg or "").split(","): + tok = tok.strip() + if not tok: + continue + try: + idx = int(tok) + except ValueError: + continue + items.append(idx) + step = by_index.get(idx) + files = step.get("files", []) if step else [] + sizes.append(estimate_step_tokens(files, root, base)) + + return json.dumps( + store.partition_steps_by_size(items, sizes, max_size, max_per) + ) if __name__ == "__main__": diff --git a/scripts/predup.py b/scripts/predup.py index da7510f..eedd84d 100644 --- a/scripts/predup.py +++ b/scripts/predup.py @@ -16,6 +16,7 @@ import os import shlex import sys +import time # --------------------------------------------------------------------------- # sys.path: ensure scripts/ dir is importable as "import store" @@ -113,6 +114,23 @@ def main() -> None: print("{}") return + # ---------------------------------------------------------------- + # 5b. Freshness guard — a recorded row only proves the output was + # injected once, not that it is still in the live context. Subagents + # share this DB and compaction/resume evicts output while the row + # persists, so beyond predup_max_age_seconds we let the call through + # rather than deny a read whose content is no longer in context. + # ---------------------------------------------------------------- + max_age = float(cfg.get("predup_max_age_seconds", 900) or 0) + if max_age > 0: + prior_ts = prior.get("ts") + try: + if prior_ts is not None and (time.time() - float(prior_ts)) > max_age: + print("{}") + return + except (TypeError, ValueError): + pass + # ---------------------------------------------------------------- # 6. Read-state guard (only for Read tool) # ---------------------------------------------------------------- diff --git a/scripts/store.py b/scripts/store.py index 5f7e09b..c8da7fe 100644 --- a/scripts/store.py +++ b/scripts/store.py @@ -101,7 +101,22 @@ # dispatch can't overflow the executor's context (which would force a # mid-batch compaction and re-derivation) or widen the blast radius of a # single failure. 0 disables the cap (send the entire tier at once). - "max_steps_per_dispatch": 6, + # Default lowered from 6 to 4 after a 6-step grouped Sonnet dispatch stalled + # the stream watchdog mid-batch (600s no-progress) and lost the whole batch; + # smaller batches bound the blast radius and the per-dispatch context. + "max_steps_per_dispatch": 4, + # Size-aware dispatch cap (used by plan_preview.py --indices). The step-count + # cap above is blunt: 4 trivial steps and 4 huge-file steps are very + # different context loads. This bounds a single grouped dispatch by ESTIMATED + # context tokens (per-step base + the byte size of each declared file ÷ 4), + # so a few large-file steps get split off even under the count cap. A single + # step over budget still dispatches alone (a step is never split). 0 disables + # the size bound (revert to the count-only cap). + "max_dispatch_context_tokens": 50000, + # Per-step base overhead (tokens) added to each step's file-size estimate + # when computing the size-aware partition — covers the shared spec/prompt and + # the step's own fields, independent of how big its files are. + "dispatch_step_base_tokens": 1500, # Resume: /nx-build persists each step's verdict to the step_ledger # table and, on a re-run for the same plan, skips already-`done` steps and # patch-retries `failed` ones from their saved diff — so a session that died @@ -117,6 +132,14 @@ # When True, also predup a conservative allowlist of read-only Bash commands # (cat, head, tail, ls, wc, grep family, find, git log/diff/show/status/branch). "predup_bash_readonly": False, + # Max age (seconds) of a recorded tool_calls row before predup STOPS treating + # it as "still in context." A row only proves the output was injected once — + # not that it survived. Subagents share the parent's tool_calls DB, and + # compaction/resume silently evicts output while the row persists, so an old + # row can deny a legitimate read whose content is no longer (or never was) in + # the live context. Beyond this window predup lets the call through. 0 + # disables the age check (revert to the original "ever-recorded" behaviour). + "predup_max_age_seconds": 900, # SessionStart resume nudge: when a recent handoff for the current branch # exists, surface a one-line hint without auto-loading anything. "resume_nudge_enabled": True, @@ -129,6 +152,18 @@ "plan_preview_input_tok_per_step": 8000, # Per-step token heuristic for the plan cost preview (output side). "plan_preview_output_tok_per_step": 2000, + # Route calibration: track per-route dispatch outcomes (pass/escalate) and + # feed the data back into routing decisions over time. + "route_calib_enabled": True, + # Minimum ratio of passed_first_try / dispatched to consider a route reliable. + "route_calib_min_success_ratio": 0.6, + # Minimum number of dispatches before calibration data is trusted. + "route_calib_min_samples": 5, + # SessionStart audit nudge: surface /nx-audit hint when ignore config has + # findings (throttled to once per repo per audit_nudge_throttle_hours). + "audit_nudge_enabled": True, + # Hours between successive audit nudges for the same repo. + "audit_nudge_throttle_hours": 24, } # --------------------------------------------------------------------------- @@ -231,6 +266,20 @@ PRIMARY KEY(session_id, plan_hash, step_index) ) """, + # Route calibration — durable per-(repo, route) dispatch outcome counters. + # Accumulated incrementally by /nx-build so routing decisions can be tuned + # over time. updated_ts is epoch seconds of the last upsert. + """ + CREATE TABLE IF NOT EXISTS route_calibration( + repo TEXT, + route TEXT, + dispatched INTEGER, + passed_first_try INTEGER, + escalated INTEGER, + updated_ts REAL, + PRIMARY KEY(repo, route) + ) + """, ] # Column migrations for databases created by an earlier schema version. @@ -714,6 +763,94 @@ def usage_rows(session_id: Optional[str] = None) -> List[Dict[str, Any]]: return [] +# --------------------------------------------------------------------------- +# current_repo helper +# --------------------------------------------------------------------------- + +def current_repo(cwd: Optional[str] = None) -> str: + """Return the basename of the git toplevel for *cwd* (or os.getcwd()). + + Falls back to the string ``"default"`` when not inside a git repo or on + any error (fail-open: never raises). + """ + try: + import subprocess as _subprocess + effective_cwd = cwd or os.getcwd() + result = _subprocess.run( + ["git", "-C", effective_cwd, "rev-parse", "--show-toplevel"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + return os.path.basename(result.stdout.strip()) + except Exception: + pass + return "default" + + +# --------------------------------------------------------------------------- +# Route calibration helpers +# --------------------------------------------------------------------------- + +def record_calibration( + repo: str, + route: str, + *, + dispatched: int = 0, + passed_first_try: int = 0, + escalated: int = 0, +) -> None: + """Upsert per-(repo, route) calibration counters, accumulating increments. + + Reads any existing row and writes back the summed counters via INSERT OR + REPLACE so concurrent writes naturally accumulate. Fail-open. + """ + try: + conn = db() + existing = conn.execute( + "SELECT dispatched, passed_first_try, escalated " + "FROM route_calibration WHERE repo=? AND route=?", + (repo, route), + ).fetchone() + if existing: + new_dispatched = existing["dispatched"] + dispatched + new_passed = existing["passed_first_try"] + passed_first_try + new_escalated = existing["escalated"] + escalated + else: + new_dispatched = dispatched + new_passed = passed_first_try + new_escalated = escalated + with conn: + conn.execute( + "INSERT OR REPLACE INTO route_calibration" + "(repo, route, dispatched, passed_first_try, escalated, updated_ts) " + "VALUES (?, ?, ?, ?, ?, ?)", + (repo, route, new_dispatched, new_passed, new_escalated, time.time()), + ) + conn.close() + except Exception: + pass + + +def calibration_rows(repo: str) -> List[Dict[str, Any]]: + """Return all route_calibration rows for *repo*, ordered by route. + + Fail-open: returns ``[]`` on any error. + """ + try: + conn = db() + rows = conn.execute( + "SELECT repo, route, dispatched, passed_first_try, escalated, updated_ts " + "FROM route_calibration WHERE repo=? ORDER BY route", + (repo,), + ).fetchall() + conn.close() + return [dict(r) for r in rows] + except Exception: + return [] + + def record_saving( session_id: str, source: str, @@ -965,6 +1102,50 @@ def partition_steps(items: List[Any], max_per: int) -> List[List[Any]]: return [items[i:i + max_per] for i in range(0, len(items), max_per)] +def partition_steps_by_size( + items: List[Any], + sizes: List[int], + max_size: int, + max_per: int = 0, +) -> List[List[Any]]: + """Pack ordered step indices into ordered sub-batches bounded by *both* a + per-batch size budget (*max_size*, e.g. estimated context tokens) and an + optional item-count cap (*max_per*), preserving order. + + Greedy and order-preserving (like ``partition_steps``): a dependent never + lands in a batch that runs before its prerequisite's. A batch is closed + before adding an item when, with the batch already non-empty, adding that + item would exceed *max_size*, or the batch already holds *max_per* items. + A single item larger than *max_size* gets its own batch (a step is never + split). ``max_size <= 0`` disables the size bound (falls back to the + count-only cap). If *sizes* doesn't line up with *items*, falls back to + ``partition_steps`` (count-only) so the caller never crashes. + """ + items = list(items) + sizes = list(sizes) + if len(sizes) != len(items): + return partition_steps(items, max_per) + if not items: + return [] + + batches: List[List[Any]] = [] + cur: List[Any] = [] + cur_size = 0 + for it, sz in zip(items, sizes): + sz = max(0, int(sz)) + over_size = max_size and max_size > 0 and cur and (cur_size + sz) > max_size + over_count = max_per and max_per > 0 and cur and len(cur) >= max_per + if over_size or over_count: + batches.append(cur) + cur = [] + cur_size = 0 + cur.append(it) + cur_size += sz + if cur: + batches.append(cur) + return batches + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- @@ -1054,6 +1235,35 @@ def _cmd_plan_batches(args) -> None: print(json.dumps(partition_steps(indices, max_per))) +def _cmd_record_usage(args) -> None: + """Append a usage row (estimated per-tier attribution recorded by /nx-build).""" + add_usage( + session_id=args.session, + model=args.model, + input_tok=args.input_tok, + output_tok=args.output_tok, + cache_read_tok=args.cache_read_tok, + ) + print(json.dumps({"ok": True}, sort_keys=True)) + + +def _cmd_calib_record(args) -> None: + """Incrementally upsert per-repo, per-route calibration counters.""" + record_calibration( + args.repo, + args.route, + dispatched=args.dispatched, + passed_first_try=args.passed_first_try, + escalated=args.escalated, + ) + print(json.dumps({"ok": True}, sort_keys=True)) + + +def _cmd_calib_list(args) -> None: + """Print all calibration rows for a repo as a JSON array.""" + print(json.dumps(calibration_rows(args.repo), sort_keys=True)) + + def main() -> None: parser = argparse.ArgumentParser( prog="store.py", @@ -1100,6 +1310,23 @@ def main() -> None: p_batches.add_argument("--max", type=int, default=None, help="Cap per batch; defaults to max_steps_per_dispatch.") + p_usage = sub.add_parser("record-usage", help="Append an estimated usage row.") + p_usage.add_argument("--session", required=True) + p_usage.add_argument("--model", required=True) + p_usage.add_argument("--input-tok", type=int, required=True) + p_usage.add_argument("--output-tok", type=int, required=True) + p_usage.add_argument("--cache-read-tok", type=int, default=0) + + p_calib = sub.add_parser("calib-record", help="Increment per-repo route calibration counters.") + p_calib.add_argument("--repo", required=True) + p_calib.add_argument("--route", required=True) + p_calib.add_argument("--dispatched", type=int, default=0) + p_calib.add_argument("--passed-first-try", type=int, default=0) + p_calib.add_argument("--escalated", type=int, default=0) + + p_clist = sub.add_parser("calib-list", help="Print calibration rows for a repo as JSON.") + p_clist.add_argument("--repo", required=True) + args = parser.parse_args() if args.command == "init": @@ -1118,6 +1345,12 @@ def main() -> None: _cmd_step_clear(args) elif args.command == "plan-batches": _cmd_plan_batches(args) + elif args.command == "record-usage": + _cmd_record_usage(args) + elif args.command == "calib-record": + _cmd_calib_record(args) + elif args.command == "calib-list": + _cmd_calib_list(args) else: parser.print_help() sys.exit(1) diff --git a/tests/test_audit_nudge.py b/tests/test_audit_nudge.py new file mode 100644 index 0000000..8fd48c3 --- /dev/null +++ b/tests/test_audit_nudge.py @@ -0,0 +1,126 @@ +""" +test_audit_nudge.py — stdlib unittest tests for scripts/audit_nudge.py + +Covers: + (a) actionable findings (unignored noise dir, no ignore file) → output has /nx-audit + (b) second run within the throttle window → {} + (c) audit_nudge_enabled=false → {} + (d) audit.run_audit raising → {} exit 0 (fail-open) +""" + +import io +import json +import os +import subprocess +import sys +import tempfile +import unittest +from contextlib import redirect_stdout +from unittest import mock + +_SCRIPTS_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "scripts" +) +if _SCRIPTS_DIR not in sys.path: + sys.path.insert(0, _SCRIPTS_DIR) + +_AUDIT_NUDGE = os.path.join(_SCRIPTS_DIR, "audit_nudge.py") + + +def _make_dirty_cwd() -> str: + """Create a temp dir with an unignored noise dir and no ignore file.""" + cwd = tempfile.mkdtemp() + os.makedirs(os.path.join(cwd, "node_modules"), exist_ok=True) + return cwd + + +def _run_nudge(payload: dict, data_dir: str, cwd: str, extra_config: dict = None): + """Run audit_nudge.py and return (parsed_output, exit_code).""" + env = os.environ.copy() + env["CLAUDE_PLUGIN_DATA"] = data_dir + env["PYTHONPATH"] = _SCRIPTS_DIR + os.pathsep + env.get("PYTHONPATH", "") + + if extra_config is not None: + with open(os.path.join(data_dir, "config.json"), "w") as fh: + json.dump(extra_config, fh) + + result = subprocess.run( + [sys.executable, _AUDIT_NUDGE], + input=json.dumps(payload).encode(), + capture_output=True, + env=env, + cwd=cwd, + timeout=15, + ) + return json.loads(result.stdout.decode()), result.returncode + + +class TestAuditNudgeFires(unittest.TestCase): + """Actionable findings → output contains /nx-audit.""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + self._cwd = _make_dirty_cwd() + + def test_fires_with_findings(self): + out, rc = _run_nudge({"cwd": self._cwd}, self._tmp, self._cwd) + self.assertEqual(rc, 0) + self.assertNotEqual(out, {}, f"Expected a hint, got: {out}") + self.assertIn("/nx-audit", json.dumps(out)) + + +class TestAuditNudgeThrottle(unittest.TestCase): + """Second run within the throttle window → {}.""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + self._cwd = _make_dirty_cwd() + + def test_throttled_second_run(self): + first, rc1 = _run_nudge({"cwd": self._cwd}, self._tmp, self._cwd) + self.assertEqual(rc1, 0) + self.assertNotEqual(first, {}, "First run should nudge") + + second, rc2 = _run_nudge({"cwd": self._cwd}, self._tmp, self._cwd) + self.assertEqual(rc2, 0) + self.assertEqual(second, {}, f"Second run should be throttled, got: {second}") + + +class TestAuditNudgeDisabled(unittest.TestCase): + """audit_nudge_enabled=false → {}.""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + self._cwd = _make_dirty_cwd() + + def test_disabled_suppresses(self): + out, rc = _run_nudge({"cwd": self._cwd}, self._tmp, self._cwd, + extra_config={"audit_nudge_enabled": False}) + self.assertEqual(rc, 0) + self.assertEqual(out, {}, f"Expected {{}} when disabled, got: {out}") + + +class TestAuditNudgeFailOpen(unittest.TestCase): + """audit.run_audit raising → {} exit 0 (fail-open).""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + self._cwd = _make_dirty_cwd() + os.environ["CLAUDE_PLUGIN_DATA"] = self._tmp + + def tearDown(self): + os.environ.pop("CLAUDE_PLUGIN_DATA", None) + + def test_run_audit_raises_is_fail_open(self): + import audit_nudge + payload = json.dumps({"cwd": self._cwd}) + buf = io.StringIO() + with mock.patch("audit.run_audit", side_effect=RuntimeError("boom")), \ + mock.patch("sys.stdin", io.StringIO(payload)), \ + redirect_stdout(buf): + audit_nudge.main() + self.assertEqual(buf.getvalue().strip(), "{}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_cost_report.py b/tests/test_cost_report.py index d570f45..a03fcc5 100644 --- a/tests/test_cost_report.py +++ b/tests/test_cost_report.py @@ -202,6 +202,8 @@ def test_report_header_present(self): out, rc = _run_cost_report(self._tmp) self.assertEqual(rc, 0) self.assertIn("[nexum] Cost report", out) + self.assertIn("ESTIMATED", out) + self.assertIn("authoritative", out) def test_token_yield_note_present(self): """v1 must include the token yield note.""" diff --git a/tests/test_plan_preview.py b/tests/test_plan_preview.py index be3123a..f6d2acf 100644 --- a/tests/test_plan_preview.py +++ b/tests/test_plan_preview.py @@ -167,5 +167,78 @@ def test_real_plan_exit_0(self): self.assertIn("Projected:", output) +class TestFilesParsingAndSizing(unittest.TestCase): + """parse_plan_steps captures files; estimate_step_tokens sizes from disk.""" + + def setUp(self): + self._root = tempfile.mkdtemp() + + def _write(self, rel, nbytes): + path = os.path.join(self._root, rel) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as fh: + fh.write("x" * nbytes) + return path + + def test_parse_captures_files_and_none(self): + plan = ( + "### Step 1: a\n- route: standard\n- files: scripts/a.py, tests/b.py\n" + "- acceptance: true\n\n" + "### Step 2: b\n- route: mechanical\n- files: none\n- acceptance: true\n" + ) + steps = plan_preview.parse_plan_steps(plan) + self.assertEqual(steps[0]["files"], ["scripts/a.py", "tests/b.py"]) + self.assertEqual(steps[1]["files"], []) + + def test_estimate_counts_existing_bytes(self): + self._write("big.py", 4000) # ~1000 tokens + base = 500 + est = plan_preview.estimate_step_tokens( + ["big.py", "missing.py"], self._root, base) + # base + 4000//4 (missing file contributes nothing) + self.assertEqual(est, base + 1000) + + def test_build_batches_isolates_large_step(self): + # Two small files + one large; size budget forces the large one alone. + self._write("small1.py", 400) + self._write("small2.py", 400) + self._write("huge.py", 400000) # ~100k tokens + plan = ( + "### Step 1: s1\n- route: standard\n- files: small1.py\n- acceptance: true\n\n" + "### Step 2: big\n- route: standard\n- files: huge.py\n- acceptance: true\n\n" + "### Step 3: s2\n- route: standard\n- files: small2.py\n- acceptance: true\n" + ) + plan_path = os.path.join(self._root, "plan.md") + with open(plan_path, "w") as fh: + fh.write(plan) + + env = os.environ.copy() + env["PYTHONPATH"] = _SCRIPTS_DIR + os.pathsep + env.get("PYTHONPATH", "") + env["CLAUDE_PLUGIN_DATA"] = tempfile.mkdtemp() + result = subprocess.run( + [sys.executable, os.path.join(_SCRIPTS_DIR, "plan_preview.py"), + "--plan", plan_path, "--indices", "1,2,3", "--root", self._root], + capture_output=True, env=env, timeout=15, + ) + self.assertEqual(result.returncode, 0) + batches = json.loads(result.stdout.decode()) + # The huge step is over the 50k default budget → its own batch. + self.assertIn([2], batches) + # Order preserved across the flattened result. + flat = [x for b in batches for x in b] + self.assertEqual(flat, [1, 2, 3]) + + def test_build_batches_missing_plan_emits_empty_json(self): + env = os.environ.copy() + env["PYTHONPATH"] = _SCRIPTS_DIR + os.pathsep + env.get("PYTHONPATH", "") + result = subprocess.run( + [sys.executable, os.path.join(_SCRIPTS_DIR, "plan_preview.py"), + "--plan", os.path.join(self._root, "nope.md"), "--indices", "1,2"], + capture_output=True, env=env, timeout=15, + ) + self.assertEqual(result.returncode, 0) + self.assertEqual(json.loads(result.stdout.decode()), []) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_predup.py b/tests/test_predup.py index 79501e5..73aa4a4 100644 --- a/tests/test_predup.py +++ b/tests/test_predup.py @@ -318,5 +318,79 @@ def test_bash_readonly_enabled_cat_denies(self): self.assertEqual(decision, "deny", f"Expected deny for cat with readonly=True, got: {out}") +def _age_tool_call(data_dir, session_id, sig, age_seconds): + """Backdate a seeded tool_call row's ts by age_seconds.""" + import time as _time + old_val = os.environ.get("CLAUDE_PLUGIN_DATA") + os.environ["CLAUDE_PLUGIN_DATA"] = data_dir + try: + import importlib + import store as _store + importlib.reload(_store) + conn = _store.db() + with conn: + conn.execute( + "UPDATE tool_calls SET ts=? WHERE session_id=? AND input_sig=?", + (_time.time() - age_seconds, session_id, sig), + ) + conn.close() + finally: + if old_val is None: + os.environ.pop("CLAUDE_PLUGIN_DATA", None) + else: + os.environ["CLAUDE_PLUGIN_DATA"] = old_val + + +class TestPredupFreshnessGuard(unittest.TestCase): + """A stale prior row (older than predup_max_age_seconds) → allow {}.""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + + def test_stale_row_allows(self): + payload = { + "session_id": "s_age", + "tool_name": "Read", + "tool_input": {"file_path": "/some/file.py"}, + } + sig = _seed_tool_call(self._tmp, "s_age", "Read", + {"file_path": "/some/file.py"}) + # Backdate well beyond the default 900s window. + _age_tool_call(self._tmp, "s_age", sig, age_seconds=5000) + out, rc = _run_predup(payload, self._tmp) + self.assertEqual(rc, 0) + self.assertEqual(out, {}, f"Expected {{}} for stale row, got: {out}") + + def test_fresh_row_still_denies(self): + """Control: a recent row within the window is still deduped.""" + payload = { + "session_id": "s_fresh", + "tool_name": "Read", + "tool_input": {"file_path": "/some/file.py"}, + } + _seed_tool_call(self._tmp, "s_fresh", "Read", + {"file_path": "/some/file.py"}) + out, rc = _run_predup(payload, self._tmp) + self.assertEqual(rc, 0) + decision = out.get("hookSpecificOutput", {}).get("permissionDecision") + self.assertEqual(decision, "deny", f"Expected deny for fresh row, got: {out}") + + def test_age_check_disabled_denies_stale(self): + """predup_max_age_seconds=0 reverts to ever-recorded behaviour.""" + payload = { + "session_id": "s_off", + "tool_name": "Read", + "tool_input": {"file_path": "/some/file.py"}, + } + sig = _seed_tool_call(self._tmp, "s_off", "Read", + {"file_path": "/some/file.py"}) + _age_tool_call(self._tmp, "s_off", sig, age_seconds=5000) + out, rc = _run_predup(payload, self._tmp, + extra_config={"predup_max_age_seconds": 0}) + self.assertEqual(rc, 0) + decision = out.get("hookSpecificOutput", {}).get("permissionDecision") + self.assertEqual(decision, "deny", f"Expected deny when age check off, got: {out}") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_store.py b/tests/test_store.py index 10f7f54..b7b083a 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -12,6 +12,7 @@ import multiprocessing import os import sqlite3 +import subprocess import sys import tempfile import time @@ -785,5 +786,162 @@ def test_env_override_honored(self): os.environ.pop("CLAUDE_PLUGIN_DATA", None) +class TestRecordUsageCLI(unittest.TestCase): + """record-usage CLI populates a usage row the cost report can render.""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + os.environ["CLAUDE_PLUGIN_DATA"] = self._tmp + + def tearDown(self): + os.environ.pop("CLAUDE_PLUGIN_DATA", None) + + def test_record_usage_cli_then_report_non_empty(self): + import store + import cost_report + store_py = os.path.join(_SCRIPTS_DIR, "store.py") + env = os.environ.copy() + env["CLAUDE_PLUGIN_DATA"] = self._tmp + env["PYTHONPATH"] = _SCRIPTS_DIR + os.pathsep + env.get("PYTHONPATH", "") + result = subprocess.run( + [sys.executable, store_py, "record-usage", "--session", "s1", + "--model", "haiku", "--input-tok", "1000", "--output-tok", "200"], + capture_output=True, env=env, timeout=15, + ) + self.assertEqual(result.returncode, 0) + self.assertEqual(json.loads(result.stdout.decode()), {"ok": True}) + + rows = store.usage_rows("s1") + self.assertEqual(len(rows), 1) + report = cost_report.build_report(rows) + self.assertNotIn("No usage rows", report) + + +class TestCalibrationHelpers(unittest.TestCase): + """record_calibration upserts incrementally; calibration_rows reads back.""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + os.environ["CLAUDE_PLUGIN_DATA"] = self._tmp + + def tearDown(self): + os.environ.pop("CLAUDE_PLUGIN_DATA", None) + + def test_record_calibration_accumulates(self): + import store + store.record_calibration("r", "standard", dispatched=1, passed_first_try=1) + store.record_calibration("r", "standard", dispatched=1, passed_first_try=1) + rows = store.calibration_rows("r") + self.assertEqual(len(rows), 1) + row = rows[0] + self.assertEqual(row["route"], "standard") + self.assertEqual(row["dispatched"], 2) + self.assertEqual(row["passed_first_try"], 2) + self.assertEqual(row["escalated"], 0) + + def test_record_calibration_escalated_independent_route(self): + import store + store.record_calibration("r", "mechanical", dispatched=1, escalated=1) + store.record_calibration("r", "standard", dispatched=1, passed_first_try=1) + rows = {row["route"]: row for row in store.calibration_rows("r")} + self.assertEqual(rows["mechanical"]["escalated"], 1) + self.assertEqual(rows["mechanical"]["passed_first_try"], 0) + self.assertEqual(rows["standard"]["passed_first_try"], 1) + + def test_calibration_rows_empty_for_unknown_repo(self): + import store + self.assertEqual(store.calibration_rows("does-not-exist"), []) + + +class TestCalibrationCLI(unittest.TestCase): + """calib-record then calib-list round-trip through the CLI.""" + + def setUp(self): + self._tmp = tempfile.mkdtemp() + os.environ["CLAUDE_PLUGIN_DATA"] = self._tmp + + def tearDown(self): + os.environ.pop("CLAUDE_PLUGIN_DATA", None) + + def _run(self, *args): + store_py = os.path.join(_SCRIPTS_DIR, "store.py") + env = os.environ.copy() + env["CLAUDE_PLUGIN_DATA"] = self._tmp + env["PYTHONPATH"] = _SCRIPTS_DIR + os.pathsep + env.get("PYTHONPATH", "") + return subprocess.run( + [sys.executable, store_py, *args], + capture_output=True, env=env, timeout=15, + ) + + def test_calib_record_list_roundtrip(self): + rec = self._run("calib-record", "--repo", "proj", "--route", "standard", + "--dispatched", "1", "--passed-first-try", "1") + self.assertEqual(rec.returncode, 0) + self.assertEqual(json.loads(rec.stdout.decode()), {"ok": True}) + + lst = self._run("calib-list", "--repo", "proj") + self.assertEqual(lst.returncode, 0) + rows = json.loads(lst.stdout.decode()) + self.assertTrue(any(r["route"] == "standard" and r["dispatched"] == 1 + for r in rows)) + + +class TestPartitionStepsBySize(unittest.TestCase): + """store.partition_steps_by_size — size + count bounded, order-preserving.""" + + def test_packs_within_size_budget(self): + import store + self.assertEqual( + store.partition_steps_by_size([1, 2, 3, 4], [10, 10, 30, 5], + max_size=25, max_per=0), + [[1, 2], [3], [4]], + ) + + def test_oversized_item_dispatches_alone(self): + import store + self.assertEqual( + store.partition_steps_by_size([1, 2, 3], [100, 5, 5], + max_size=25, max_per=0), + [[1], [2, 3]], + ) + + def test_count_cap_applies_when_size_disabled(self): + import store + self.assertEqual( + store.partition_steps_by_size([1, 2, 3, 4, 5], [1, 1, 1, 1, 1], + max_size=0, max_per=2), + [[1, 2], [3, 4], [5]], + ) + + def test_count_cap_closes_batch_before_size(self): + import store + # Size budget would allow all five, but max_per=2 closes earlier. + self.assertEqual( + store.partition_steps_by_size([1, 2, 3, 4, 5], [1, 1, 1, 1, 1], + max_size=1000, max_per=2), + [[1, 2], [3, 4], [5]], + ) + + def test_size_mismatch_falls_back_to_count(self): + import store + self.assertEqual( + store.partition_steps_by_size([1, 2, 3], [1, 2], + max_size=10, max_per=2), + [[1, 2], [3]], + ) + + def test_empty(self): + import store + self.assertEqual( + store.partition_steps_by_size([], [], max_size=10, max_per=2), []) + + def test_order_preserved(self): + import store + out = store.partition_steps_by_size( + [5, 1, 9, 3], [10, 10, 10, 10], max_size=20, max_per=0) + flat = [x for batch in out for x in batch] + self.assertEqual(flat, [5, 1, 9, 3]) + + if __name__ == "__main__": unittest.main()