diff --git a/src/infospace_bench/budget.py b/src/infospace_bench/budget.py index db0cb29..163cb79 100644 --- a/src/infospace_bench/budget.py +++ b/src/infospace_bench/budget.py @@ -375,6 +375,75 @@ def _post_json(url: str, payload: dict[str, Any], timeout: float) -> None: response.read() +def budget_show(root: str | Path) -> dict[str, Any]: + """Return the full per-infospace budget structure for inspection.""" + root_path = Path(root) + plans_payload = _read_plans(root_path / PLANS_FILE) + usage_payload = _read_usage(root_path / USAGE_FILE) + variance = read_run_variance(root_path) + return { + "root": str(root_path), + "plans": { + "schema_version": plans_payload.get("schema_version"), + "pruned_count": plans_payload.get("pruned_count") or 0, + "snapshots": list(plans_payload.get("snapshots") or []), + }, + "usage": { + "schema_version": usage_payload.get("schema_version"), + "runs": list(usage_payload.get("runs") or []), + }, + "summary": variance, + } + + +def budget_list_workspace(workspace: str | Path) -> list[dict[str, Any]]: + """Walk ``/infospaces/*/output/budget/`` and roll up each infospace. + + Missing budget directories are treated as zero, not as errors. Returns a + list of summaries in slug order. + """ + workspace_path = Path(workspace) + infospaces_dir = workspace_path / "infospaces" + if not infospaces_dir.is_dir(): + return [] + rollups: list[dict[str, Any]] = [] + for entry in sorted(infospaces_dir.iterdir()): + if not entry.is_dir(): + continue + rollups.append(_summarise_infospace(entry)) + return rollups + + +def _summarise_infospace(root: Path) -> dict[str, Any]: + plans = read_plan_snapshots(root) + runs = read_usage_runs(root) + total_tokens = 0 + total_cost_known = 0.0 + total_cost_estimated = 0.0 + last_run_at: str | None = None + for run in runs: + rollup = run.get("rollup") or {} + total_tokens += int(rollup.get("total_tokens") or 0) + total_cost_known += _coerce_float(rollup.get("total_cost_usd_known")) or 0.0 + estimated = _coerce_float(rollup.get("total_cost_usd_estimated")) + if estimated is not None: + total_cost_estimated += estimated + completed_at = run.get("completed_at") + if isinstance(completed_at, str) and (last_run_at is None or completed_at > last_run_at): + last_run_at = completed_at + return { + "slug": root.name, + "root": str(root), + "plans_count": len(plans), + "runs_count": len(runs), + "total_tokens": total_tokens, + "total_cost_usd_known": round(total_cost_known, 6), + "total_cost_usd_estimated": round(total_cost_estimated, 6) if total_cost_estimated else None, + "last_run_at": last_run_at, + "latest_snapshot_id": plans[-1].get("snapshot_id") if plans else None, + } + + def read_run_variance(root: str | Path) -> dict[str, Any] | None: path = Path(root) / SUMMARY_FILE if not path.is_file(): diff --git a/src/infospace_bench/cli.py b/src/infospace_bench/cli.py index 8eab20b..d93adf5 100644 --- a/src/infospace_bench/cli.py +++ b/src/infospace_bench/cli.py @@ -228,6 +228,19 @@ def build_parser() -> argparse.ArgumentParser: generate_from_source.add_argument("--max-chunks", type=int, default=0) generate_from_source.add_argument("--apply", action="store_true") + budget = sub.add_parser("budget", help="Inspect per-infospace budget and usage records") + budget_sub = budget.add_subparsers(dest="budget_command", required=True) + budget_list = budget_sub.add_parser( + "list", + help="Workspace rollup: one row per infospace with plan/run counts, tokens, cost", + ) + budget_list.add_argument("workspace") + budget_show = budget_sub.add_parser( + "show", + help="Show the full plans/usage/summary payload for one infospace", + ) + budget_show.add_argument("root") + archive = sub.add_parser( "archive", help="Archive an infospace into artifact-store (durable, content-addressed)", @@ -538,6 +551,14 @@ def main(argv: list[str] | None = None) -> int: _write_json(plan_generation(infospace.root, stage=args.stage)) else: parser.error(f"Unhandled generate command: {args.generate_command}") + elif args.command == "budget": + from .budget import budget_list_workspace, budget_show + if args.budget_command == "list": + _write_json({"workspace": str(Path(args.workspace)), "infospaces": budget_list_workspace(Path(args.workspace))}) + elif args.budget_command == "show": + _write_json(budget_show(Path(args.root))) + else: + parser.error(f"Unhandled budget command: {args.budget_command}") elif args.command == "archive": record = archive_infospace( Path(args.root), diff --git a/tests/test_budget_registry.py b/tests/test_budget_registry.py index d2cb190..dd02690 100644 --- a/tests/test_budget_registry.py +++ b/tests/test_budget_registry.py @@ -623,6 +623,105 @@ def test_run_generation_never_fails_when_hub_is_down(tmp_path: Path, monkeypatch assert status["completed"] is True +def test_budget_show_returns_full_per_infospace_structure(tmp_path: Path) -> None: + from infospace_bench.budget import budget_show + from infospace_bench.generator import run_generation + + root = _build_infospace(tmp_path) + fixture = tmp_path / "responses.yaml" + _write_minimal_fixture(fixture) + plan_generation(root) + run_generation(root, fixture_responses=fixture) + + payload = budget_show(root) + + assert payload["root"] == str(root) + assert payload["plans"]["snapshots"], "plans must round-trip" + assert payload["usage"]["runs"], "usage runs must round-trip" + assert payload["summary"] is not None + assert payload["summary"]["snapshot_resolved"] is True + + +def test_budget_list_workspace_rolls_up_multiple_infospaces(tmp_path: Path) -> None: + from infospace_bench.budget import budget_list_workspace + from infospace_bench.generator import init_generation_infospace, run_generation + + book = tmp_path / "book.epub" + _write_three_chapter_epub(book) + fixture = tmp_path / "responses.yaml" + _write_minimal_fixture(fixture) + + a = init_generation_infospace(tmp_path, book, "alpha", name="Alpha") + b = init_generation_infospace(tmp_path, book, "beta", name="Beta") + plan_generation(a.root, from_chapter=1, to_chapter=2) + plan_generation(b.root) + run_generation(a.root, fixture_responses=fixture) + # leave beta with a plan but no run, to verify partial rollups + + rollups = budget_list_workspace(tmp_path) + by_slug = {item["slug"]: item for item in rollups} + + assert {"alpha", "beta"} <= set(by_slug.keys()) + assert by_slug["alpha"]["plans_count"] >= 1 + assert by_slug["alpha"]["runs_count"] == 1 + assert by_slug["alpha"]["last_run_at"] is not None + assert by_slug["beta"]["plans_count"] >= 1 + assert by_slug["beta"]["runs_count"] == 0 + assert by_slug["beta"]["last_run_at"] is None + + +def test_budget_list_workspace_handles_missing_or_empty_directories(tmp_path: Path) -> None: + from infospace_bench.budget import budget_list_workspace + from infospace_bench.generator import init_generation_infospace + + # No infospaces/ at all → empty list, not an error + assert budget_list_workspace(tmp_path) == [] + + # An infospace with no budget dir → rollup with zeros, not a crash + book = tmp_path / "book.epub" + _write_three_chapter_epub(book) + init_generation_infospace(tmp_path, book, "zero", name="Zero") + + rollups = budget_list_workspace(tmp_path) + assert len(rollups) == 1 + assert rollups[0]["slug"] == "zero" + assert rollups[0]["plans_count"] == 0 + assert rollups[0]["runs_count"] == 0 + assert rollups[0]["total_tokens"] == 0 + assert rollups[0]["total_cost_usd_known"] == 0.0 + + +def test_budget_cli_list_and_show(tmp_path: Path) -> None: + from infospace_bench.generator import run_generation + + root = _build_infospace(tmp_path) + fixture = tmp_path / "responses.yaml" + _write_minimal_fixture(fixture) + plan_generation(root) + run_generation(root, fixture_responses=fixture) + + env = os.environ.copy() + env["PYTHONPATH"] = "src:/home/worsch/markitect-tool/src" + + list_result = subprocess.run( + [sys.executable, "-m", "infospace_bench", "budget", "list", str(tmp_path)], + check=False, env=env, text=True, capture_output=True, + ) + show_result = subprocess.run( + [sys.executable, "-m", "infospace_bench", "budget", "show", str(root)], + check=False, env=env, text=True, capture_output=True, + ) + + assert list_result.returncode == 0, list_result.stderr + assert show_result.returncode == 0, show_result.stderr + list_payload = json.loads(list_result.stdout) + show_payload = json.loads(show_result.stdout) + assert list_payload["workspace"] == str(tmp_path) + assert any(item["slug"] == "budget-test" for item in list_payload["infospaces"]) + assert show_payload["plans"]["snapshots"] + assert show_payload["usage"]["runs"] + + def test_plan_cli_writes_snapshot(tmp_path: Path) -> None: root = _build_infospace(tmp_path) env = os.environ.copy() diff --git a/workplans/IB-WP-0019-budget-and-usage-registry.md b/workplans/IB-WP-0019-budget-and-usage-registry.md index dd19db8..60116e6 100644 --- a/workplans/IB-WP-0019-budget-and-usage-registry.md +++ b/workplans/IB-WP-0019-budget-and-usage-registry.md @@ -177,7 +177,7 @@ state_hub_task_id: "968bca1d-63ff-4818-83bb-ca314b1e633c" ```task id: IB-WP-0019-T06 -status: todo +status: done priority: medium state_hub_task_id: "7cb34bfc-c562-4dda-a6d4-b44158644e19" ```