diff --git a/.gitignore b/.gitignore index 71c6ba3..b902dab 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,4 @@ saved-config.yaml *.db *.db-wal *.db-shm +uv.lock diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ba4a9c4..e328fa7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -23,7 +23,7 @@ repos: hooks: - id: mypy name: mypy - entry: python -m mypy + entry: .venv/bin/python -m mypy --python-version 3.12 language: system pass_filenames: false types: [python] diff --git a/src/tokenops/control/http.py b/src/tokenops/control/http.py index 0e87f01..ecca0ab 100644 --- a/src/tokenops/control/http.py +++ b/src/tokenops/control/http.py @@ -6,14 +6,18 @@ from __future__ import annotations +import csv +import io +import json import os from collections.abc import Awaitable, Callable, Mapping +from dataclasses import asdict from typing import Any import httpx from chronicle.session import reset_session -from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse, Response +from fastapi import FastAPI, Query, Request +from fastapi.responses import JSONResponse, Response, StreamingResponse from tokenops.control.core import Halt from tokenops.control.engine import Throttled @@ -65,6 +69,70 @@ async def register_run(request: Request) -> JSONResponse: ) +_EXPORT_CSV_COLUMNS = [ + "run_id", + "agent", + "status", + "cost_micros", + "steps", + "started_at", + "ended_at", + "duration_s", + "dims", + "halt_reason", + "detector", + "governance_events", +] + + +def _run_to_csv_row(rec): # type: ignore[no-untyped-def] + d = asdict(rec) + d["duration_s"] = round(rec.ended_at - rec.started_at, 2) if rec.ended_at else "" + d["dims"] = json.dumps(rec.dims) if rec.dims else "" + d["governance_events"] = json.dumps(rec.governance_events) if rec.governance_events else "" + return [d.get(col, "") for col in _EXPORT_CSV_COLUMNS] + + +def mount_export(app: FastAPI, store: Store) -> None: + """Mount ``GET /v1/export`` — on-demand run-record export (CSV / JSON).""" + + @app.get("/v1/export") + def export_runs( + from_at: float | None = Query(None, description="Start timestamp (epoch seconds)"), + to_at: float | None = Query(None, description="End timestamp (epoch seconds)"), + agent: str | None = Query(None, description="Filter by agent name"), + status: str | None = Query(None, description="Filter by run status"), + tenant: str | None = Query(None, description="Filter by tenant (from dims)"), + format: str = Query("json", description="Output format: json or csv"), + limit: int = Query(5000, ge=1, le=10_000, description="Max rows to return"), + ) -> Response: + runs = store.export_runs( + from_at=from_at, + to_at=to_at, + agent=agent, + status=status, + tenant=tenant, + limit=limit, + ) + + if format == "csv": + buf = io.StringIO() + writer = csv.writer(buf) + writer.writerow(_EXPORT_CSV_COLUMNS) + for rec in runs: + writer.writerow(_run_to_csv_row(rec)) + buf.seek(0) + return StreamingResponse( + iter([buf.getvalue()]), + media_type="text/csv", + headers={"Content-Disposition": 'attachment; filename="export.csv"'}, + ) + + # JSON (default) + rows = [asdict(rec) for rec in runs] + return JSONResponse(rows, headers={"X-Total-Count": str(len(rows))}) + + def with_governance_errors(handler: Handler) -> Handler: """Wrap a task handler so Halt → 200 halted and Throttled → 429.""" diff --git a/src/tokenops/control/store.py b/src/tokenops/control/store.py index 9fd5b0e..2bf4a19 100644 --- a/src/tokenops/control/store.py +++ b/src/tokenops/control/store.py @@ -30,6 +30,7 @@ import time import uuid from collections.abc import Callable +from datetime import date, datetime from typing import Any, TypeVar from tokenops.control.ledger import LIFETIME, RUN_TOTAL_BUDGET @@ -64,6 +65,26 @@ def _known_policy_templates() -> frozenset[str]: return frozenset({*_TEMPLATES, "trajectory_hint"}) +def _coerce_epoch(value: float | str | date | None) -> float | None: + """Coerce a timestamp value to epoch float. + + Accepts epoch floats (pass-through), ISO date strings (``"2026-09-19"``), + ``datetime.date`` objects, or ``None``. Dates are converted to start-of-day + epoch seconds so they compare correctly against the ``started_at`` REAL column. + """ + if value is None: + return None + if isinstance(value, int | float): + return float(value) + if isinstance(value, date) and not isinstance(value, datetime): + return datetime.combine(value, datetime.min.time()).timestamp() + if isinstance(value, str): + return datetime.fromisoformat(value).timestamp() + if isinstance(value, datetime): + return value.timestamp() + return float(value) + + _SCHEMA = """ CREATE TABLE IF NOT EXISTS segments ( id TEXT PRIMARY KEY, name TEXT NOT NULL, dimension TEXT NOT NULL, @@ -519,6 +540,49 @@ def list_runs(self, *, problematic_only: bool = False, limit: int = 200) -> list sql += " ORDER BY started_at DESC LIMIT ?" return [self._run_with_ledger_cost(r) for r in self._db.execute(sql, (limit,))] + @_locked + def export_runs( + self, + *, + from_at: float | None = None, + to_at: float | None = None, + agent: str | None = None, + status: str | None = None, + tenant: str | None = None, + limit: int = 5000, + ) -> list[RunRecord]: + """Export run-records filtered by time range, agent, status, or tenant. + + Used by the ``GET /v1/export`` route for FinOps / chargeback CSV and JSON. + Returns at most *limit* rows (capped at 10 000) ordered by ``started_at`` DESC. + """ + limit = min(limit, 10_000) + from_at = _coerce_epoch(from_at) + to_at = _coerce_epoch(to_at) + where: list[str] = [] + params: list[object] = [] + if from_at is not None: + where.append("started_at >= ?") + params.append(from_at) + if to_at is not None: + where.append("started_at <= ?") + params.append(to_at) + if agent is not None: + where.append("agent = ?") + params.append(agent) + if status is not None: + where.append("status = ?") + params.append(status) + if tenant is not None: + where.append("json_extract(dims, '$.tenant') = ?") + params.append(tenant) + sql = "SELECT * FROM runs" + if where: + sql += " WHERE " + " AND ".join(where) + sql += " ORDER BY started_at DESC LIMIT ?" + params.append(limit) + return [self._run_with_ledger_cost(r) for r in self._db.execute(sql, tuple(params))] + def _run_with_ledger_cost(self, row: sqlite3.Row) -> RunRecord: """Build a RunRecord; prefer ``__run_total__`` ledger spend when present.""" rec = _run(row) diff --git a/src/tokenops/server/app.py b/src/tokenops/server/app.py index fca90fd..5059666 100644 --- a/src/tokenops/server/app.py +++ b/src/tokenops/server/app.py @@ -7,7 +7,7 @@ from fastapi import FastAPI from tokenops import __version__ -from tokenops.control.http import mount_run_registration +from tokenops.control.http import mount_export, mount_run_registration from tokenops.control.store import Store @@ -32,6 +32,7 @@ def health() -> dict[str, str]: return {"status": "ok", "service": "tokenops-control-plane"} mount_run_registration(app, store) + mount_export(app, store) # Placeholder for future plane APIs (observe, governance admin over HTTP, etc.). # Agents keep using ControlPlaneClient; expand the plane surface here. diff --git a/src/tokenops/ui/views/dashboard.py b/src/tokenops/ui/views/dashboard.py index c90a2f6..115c39b 100644 --- a/src/tokenops/ui/views/dashboard.py +++ b/src/tokenops/ui/views/dashboard.py @@ -2,6 +2,12 @@ from __future__ import annotations +import csv +import io +import json +from dataclasses import asdict +from datetime import datetime, time + import altair as alt import pandas as pd import streamlit as st @@ -169,6 +175,87 @@ def _seg(r) -> str: ) st.dataframe(table, use_container_width=True, hide_index=True) +# ---- export (CSV / JSON download) ---------------------------------------- # +_EXPORT_COLUMNS = [ + "run_id", + "agent", + "status", + "cost_micros", + "steps", + "started_at", + "ended_at", + "duration_s", + "dims", + "halt_reason", + "detector", + "governance_events", +] + + +def _to_csv(runs_list): + buf = io.StringIO() + writer = csv.writer(buf) + writer.writerow(_EXPORT_COLUMNS) + for rec in runs_list: + d = asdict(rec) + d["duration_s"] = round(rec.ended_at - rec.started_at, 2) if rec.ended_at else "" + d["dims"] = json.dumps(rec.dims) if rec.dims else "" + d["governance_events"] = json.dumps(rec.governance_events) if rec.governance_events else "" + writer.writerow([d.get(c, "") for c in _EXPORT_COLUMNS]) + return buf.getvalue() + + +def _date_to_epoch(d, end_of_day: bool = False) -> float: + """Convert a date to epoch seconds. ``end_of_day`` returns 23:59:59.999.""" + t = time.max if end_of_day else time.min + return datetime.combine(d, t).timestamp() + + +with st.expander("Export run data", expanded=False): + agents = sorted({r.agent for r in runs}) + ecol1, ecol2, ecol3, ecol4 = st.columns(4) + with ecol1: + export_from = st.date_input("From", value=None, key="export_from") + with ecol2: + export_to = st.date_input("To", value=None, key="export_to") + with ecol3: + export_agent = st.selectbox("Agent", ["All"] + agents, key="export_agent") + with ecol4: + export_status = st.selectbox( + "Status", + ["All", "completed", "halted", "error", "running"], + key="export_status", + ) + + export_runs = store.export_runs( + from_at=_date_to_epoch(export_from) if export_from else None, + to_at=_date_to_epoch(export_to, end_of_day=True) if export_to else None, + agent=export_agent if export_agent != "All" else None, + status=export_status if export_status != "All" else None, + limit=10_000, + ) + + st.caption(f"{len(export_runs)} runs matched") + st.markdown( + "", + unsafe_allow_html=True, + ) + dl_col1, dl_col2, _ = st.columns([1, 1, 6]) + with dl_col1: + st.download_button( + "Download CSV", + data=_to_csv(export_runs), + file_name="export.csv", + mime="text/csv", + ) + with dl_col2: + st.download_button( + "Download JSON", + data=json.dumps([asdict(r) for r in export_runs], default=str, indent=2), + file_name="export.json", + mime="application/json", + ) + # ---- run detail picker (when not already focused) ------------------------ # if not focus_run: st.subheader("Run detail") diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 0000000..929a562 --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,328 @@ +"""GET /v1/export — on-demand run-record export (CSV / JSON).""" + +from __future__ import annotations + +import csv +import io +import json + +import pytest + +from tokenops.control.models import RunRecord +from tokenops.control.store import Store + + +@pytest.fixture +def store(tmp_path): + return Store(str(tmp_path / "export.db"), auto_seed=False) + + +def _make_run( + store: Store, + *, + run_id: str = "run-1", + agent: str = "researcher", + status: str = "completed", + cost_micros: int = 1000, + steps: int = 5, + started_at: float = 100.0, + ended_at: float | None = 200.0, + dims: dict[str, str] | None = None, + governance_events: list[dict] | None = None, +) -> RunRecord: + rec = RunRecord( + run_id=run_id, + agent=agent, + status=status, + cost_micros=cost_micros, + steps=steps, + started_at=started_at, + ended_at=ended_at, + dims=dims or {}, + governance_events=governance_events or [], + ) + created = store.create_run(rec) + # governance_events is written via update_run (same as real server flow) + if governance_events: + store.update_run(run_id, governance_events=governance_events) + return store.get_run(run_id) or created + + +# ------------------------------------------------------------------ # +# Store.export_runs unit tests # +# ------------------------------------------------------------------ # + + +def test_export_runs_returns_all(store): + _make_run(store, run_id="r1", agent="a", started_at=100) + _make_run(store, run_id="r2", agent="b", started_at=200) + result = store.export_runs() + assert len(result) == 2 + + +def test_export_runs_filters_from_at(store): + _make_run(store, run_id="r1", started_at=100) + _make_run(store, run_id="r2", started_at=200) + result = store.export_runs(from_at=150) + assert len(result) == 1 + assert result[0].run_id == "r2" + + +def test_export_runs_filters_to_at(store): + _make_run(store, run_id="r1", started_at=100) + _make_run(store, run_id="r2", started_at=200) + result = store.export_runs(to_at=150) + assert len(result) == 1 + assert result[0].run_id == "r1" + + +def test_export_runs_filters_agent(store): + _make_run(store, run_id="r1", agent="researcher", started_at=100) + _make_run(store, run_id="r2", agent="writer", started_at=200) + result = store.export_runs(agent="writer") + assert len(result) == 1 + assert result[0].run_id == "r2" + + +def test_export_runs_filters_status(store): + _make_run(store, run_id="r1", status="completed", started_at=100) + _make_run(store, run_id="r2", status="halted", started_at=200) + result = store.export_runs(status="halted") + assert len(result) == 1 + assert result[0].run_id == "r2" + + +def test_export_runs_filters_tenant(store): + _make_run(store, run_id="r1", dims={"tenant": "acme"}, started_at=100) + _make_run(store, run_id="r2", dims={"tenant": "globex"}, started_at=200) + _make_run(store, run_id="r3", dims={}, started_at=300) + result = store.export_runs(tenant="acme") + assert len(result) == 1 + assert result[0].run_id == "r1" + + +def test_export_runs_combined_filters(store): + _make_run(store, run_id="r1", agent="a", status="completed", started_at=100) + _make_run(store, run_id="r2", agent="a", status="halted", started_at=200) + _make_run(store, run_id="r3", agent="b", status="completed", started_at=300) + result = store.export_runs(agent="a", status="completed") + assert len(result) == 1 + assert result[0].run_id == "r1" + + +def test_export_runs_limit(store): + for i in range(20): + _make_run(store, run_id=f"r{i:02d}", started_at=float(i)) + result = store.export_runs(limit=5) + assert len(result) == 5 + + +def test_export_runs_empty(store): + result = store.export_runs() + assert result == [] + + +def test_export_runs_preserves_governance_events(store): + events = [{"kind": "halt", "reason": "budget exceeded", "policy": "cost_budget"}] + _make_run(store, run_id="r1", governance_events=events) + result = store.export_runs() + assert result[0].governance_events == events + + +# ------------------------------------------------------------------ # +# GET /v1/export integration tests # +# ------------------------------------------------------------------ # + + +def test_export_json_default(store, tmp_path): + from fastapi.testclient import TestClient + + from tokenops.server.app import create_app + + _make_run(store, run_id="r1", agent="a", cost_micros=5000, started_at=100) + app = create_app(store=store) + client = TestClient(app) + + resp = client.get("/v1/export") + assert resp.status_code == 200 + data = resp.json() + assert len(data) == 1 + assert data[0]["run_id"] == "r1" + assert data[0]["cost_micros"] == 5000 + assert resp.headers["X-Total-Count"] == "1" + + +def test_export_json_with_filters(store): + from fastapi.testclient import TestClient + + from tokenops.server.app import create_app + + _make_run(store, run_id="r1", agent="a", started_at=100) + _make_run(store, run_id="r2", agent="b", started_at=200) + app = create_app(store=store) + client = TestClient(app) + + resp = client.get("/v1/export?agent=b") + assert resp.status_code == 200 + data = resp.json() + assert len(data) == 1 + assert data[0]["run_id"] == "r2" + + +def test_export_csv(store): + from fastapi.testclient import TestClient + + from tokenops.server.app import create_app + + _make_run( + store, + run_id="r1", + agent="researcher", + status="completed", + cost_micros=12000, + steps=8, + started_at=100.0, + ended_at=250.0, + dims={"tenant": "acme"}, + governance_events=[{"kind": "mutate", "reason": "cap applied"}], + ) + app = create_app(store=store) + client = TestClient(app) + + resp = client.get("/v1/export?format=csv") + assert resp.status_code == 200 + assert resp.headers["content-type"] == "text/csv; charset=utf-8" + assert "export.csv" in resp.headers.get("content-disposition", "") + + reader = csv.reader(io.StringIO(resp.text)) + rows = list(reader) + header = rows[0] + assert "run_id" in header + assert "cost_micros" in header + assert "duration_s" in header + assert "governance_events" in header + body = rows[1] + assert body[header.index("run_id")] == "r1" + assert body[header.index("agent")] == "researcher" + assert body[header.index("cost_micros")] == "12000" + assert body[header.index("duration_s")] == "150.0" + # dims and governance_events are JSON strings in CSV + dims_val = json.loads(body[header.index("dims")]) + assert dims_val == {"tenant": "acme"} + gov_val = json.loads(body[header.index("governance_events")]) + assert len(gov_val) == 1 + assert gov_val[0]["kind"] == "mutate" + + +def test_export_empty(store): + from fastapi.testclient import TestClient + + from tokenops.server.app import create_app + + app = create_app(store=store) + client = TestClient(app) + + resp = client.get("/v1/export") + assert resp.status_code == 200 + assert resp.json() == [] + + +def test_export_csv_empty(store): + from fastapi.testclient import TestClient + + from tokenops.server.app import create_app + + app = create_app(store=store) + client = TestClient(app) + + resp = client.get("/v1/export?format=csv") + assert resp.status_code == 200 + reader = csv.reader(io.StringIO(resp.text)) + rows = list(reader) + assert len(rows) == 1 # header only + + +def test_export_limit_capped(store): + from fastapi.testclient import TestClient + + from tokenops.server.app import create_app + + for i in range(15): + _make_run(store, run_id=f"r{i:02d}", started_at=float(i)) + app = create_app(store=store) + client = TestClient(app) + + resp = client.get("/v1/export?limit=5") + assert resp.status_code == 200 + assert len(resp.json()) == 5 + + +# ------------------------------------------------------------------ # +# _coerce_epoch unit tests # +# ------------------------------------------------------------------ # + +from datetime import date, datetime + +from tokenops.control.store import _coerce_epoch + + +def test_coerce_epoch_none(): + assert _coerce_epoch(None) is None + + +def test_coerce_epoch_float_passthrough(): + assert _coerce_epoch(1700000000.0) == 1700000000.0 + + +def test_coerce_epoch_int_passthrough(): + assert _coerce_epoch(1700000000) == 1700000000.0 + + +def test_coerce_epoch_date_string(): + result = _coerce_epoch("2026-09-19") + expected = datetime(2026, 9, 19).timestamp() + assert result == pytest.approx(expected) + + +def test_coerce_epoch_iso_datetime_string(): + result = _coerce_epoch("2026-09-19T12:00:00") + expected = datetime(2026, 9, 19, 12, 0, 0).timestamp() + assert result == pytest.approx(expected) + + +def test_coerce_epoch_date_object(): + result = _coerce_epoch(date(2026, 9, 19)) + expected = datetime(2026, 9, 19).timestamp() + assert result == pytest.approx(expected) + + +def test_coerce_epoch_datetime_object(): + dt = datetime(2026, 9, 19, 15, 30) + result = _coerce_epoch(dt) + assert result == dt.timestamp() + + +# ------------------------------------------------------------------ # +# String-date filtering in export_runs (regression for type mismatch)# +# ------------------------------------------------------------------ # + + +def test_export_runs_filters_date_string(store): + """Passing a date string instead of epoch float should still filter correctly.""" + # epoch 1600000000 ≈ 2020-09-13, epoch 1900000000 ≈ 2029-11-20 + _make_run(store, run_id="r1", started_at=1600000000.0) + _make_run(store, run_id="r2", started_at=1900000000.0) + # "2023-11-15" ≈ epoch 1699986600 — should include r2 but not r1 + result = store.export_runs(from_at="2023-11-15") + assert len(result) == 1 + assert result[0].run_id == "r2" + + +def test_export_runs_filters_to_date_string(store): + """Passing a to_at date string should filter correctly.""" + _make_run(store, run_id="r1", started_at=1600000000.0) + _make_run(store, run_id="r2", started_at=1900000000.0) + # "2023-11-15" ≈ epoch 1699986600 — should include r1 but not r2 + result = store.export_runs(to_at="2023-11-15") + assert len(result) == 1 + assert result[0].run_id == "r1"