tokenjam 0.2.0__py3-none-any.whl

tokenjam/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.9"

tokenjam/api/app.py

@@ -0,0 +1,104 @@
+"""FastAPI application factory. Called by `tj serve`."""
+from __future__ import annotations
+
+from html import escape as html_escape
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+
+from tokenjam.api.middleware import IngestAuthMiddleware
+from tokenjam.core.config import TjConfig
+
+if TYPE_CHECKING:
+    from tokenjam.core.db import StorageBackend
+    from tokenjam.core.ingest import IngestPipeline
+
+_UI_DIR = Path(__file__).resolve().parent.parent / "ui"
+
+
+def create_app(
+    config: TjConfig,
+    db: StorageBackend,
+    ingest_pipeline: IngestPipeline,
+) -> FastAPI:
+    """
+    Build and return the FastAPI app.
+
+    db and ingest_pipeline are passed in (not imported globally) so tests
+    can inject mocks easily.
+    """
+    app = FastAPI(
+        title="TokenJam",
+        version="0.1.0",
+        docs_url="/docs",
+        redoc_url=None,
+    )
+
+    # CORS — local only by default
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origin_regex=r"^https?://(localhost|127\.0\.0\.1)(:\d+)?$",
+        allow_methods=["GET", "POST", "PATCH"],
+        allow_headers=["Authorization", "Content-Type"],
+    )
+
+    # Ingest auth middleware
+    app.add_middleware(IngestAuthMiddleware)
+
+    # Shared state for routes
+    app.state.config = config
+    app.state.db = db
+    app.state.pipeline = ingest_pipeline
+
+    # Register routers
+    from tokenjam.api.routes.spans import router as spans_router
+    from tokenjam.api.routes.traces import router as traces_router
+    from tokenjam.api.routes.cost import router as cost_router
+    from tokenjam.api.routes.tools import router as tools_router
+    from tokenjam.api.routes.alerts import router as alerts_router
+    from tokenjam.api.routes.drift import router as drift_router
+    from tokenjam.api.routes.metrics import router as metrics_router
+    from tokenjam.api.routes.status import router as status_router
+    from tokenjam.api.routes.otlp import router as otlp_router
+    from tokenjam.api.routes.budget import router as budget_router
+    from tokenjam.api.routes.agents import router as agents_router
+
+    app.include_router(spans_router, prefix="/api/v1")
+    app.include_router(traces_router, prefix="/api/v1")
+    app.include_router(cost_router, prefix="/api/v1")
+    app.include_router(tools_router, prefix="/api/v1")
+    app.include_router(alerts_router, prefix="/api/v1")
+    app.include_router(drift_router, prefix="/api/v1")
+    app.include_router(status_router, prefix="/api/v1")
+    app.include_router(budget_router, prefix="/api/v1")
+    app.include_router(agents_router, prefix="/api/v1")
+    app.include_router(metrics_router)  # /metrics — no prefix
+    app.include_router(otlp_router)  # /v1/traces, /v1/metrics, /v1/logs — no prefix
+
+    # --- Web UI ---
+    _index_html = ""
+    index_path = _UI_DIR / "index.html"
+    if index_path.exists():
+        _index_html = index_path.read_text()
+
+    def _serve_ui() -> HTMLResponse:
+        html = _index_html
+        if config.api.auth.enabled and config.api.auth.api_key:
+            html = html.replace(
+                "</head>",
+                f'<meta name="tj-api-key" content="{html_escape(config.api.auth.api_key, quote=True)}">\n</head>',
+            )
+        return HTMLResponse(html)
+
+    @app.get("/", include_in_schema=False)
+    async def ui_root():
+        return _serve_ui()
+
+    @app.get("/ui/{path:path}", include_in_schema=False)
+    async def ui_catchall(path: str):
+        return _serve_ui()
+
+    return app

tokenjam/api/deps.py

@@ -0,0 +1,18 @@
+"""Shared FastAPI dependencies."""
+from __future__ import annotations
+
+from fastapi import Depends, HTTPException, Request
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
+security = HTTPBearer(auto_error=False)
+
+
+def require_api_key(
+    request: Request,
+    credentials: HTTPAuthorizationCredentials | None = Depends(security),
+) -> None:
+    config = request.app.state.config
+    if not config.api.auth.enabled:
+        return
+    if credentials is None or credentials.credentials != config.api.auth.api_key:
+        raise HTTPException(status_code=401, detail="Invalid API key")

tokenjam/api/middleware.py

@@ -0,0 +1,28 @@
+"""Ingest auth middleware — validates Bearer token on POST /api/v1/spans."""
+from __future__ import annotations
+
+from fastapi import Request
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+
+
+class IngestAuthMiddleware(BaseHTTPMiddleware):
+    """
+    Validates the ingest secret on POST /api/v1/spans.
+    If security.ingest_secret is empty string, auth is disabled.
+    Returns 401 with JSON error if secret is wrong or missing.
+    """
+
+    PROTECTED_PATHS = {"/api/v1/spans", "/v1/logs", "/v1/traces"}
+
+    async def dispatch(self, request: Request, call_next):  # type: ignore[override]
+        if request.method == "POST" and request.url.path in self.PROTECTED_PATHS:
+            secret = request.app.state.config.security.ingest_secret
+            if secret:
+                auth = request.headers.get("Authorization", "")
+                if not auth.startswith("Bearer ") or auth[7:] != secret:
+                    return JSONResponse(
+                        status_code=401,
+                        content={"detail": "Invalid ingest secret"},
+                    )
+        return await call_next(request)

tokenjam/api/routes/agents.py

@@ -0,0 +1,33 @@
+"""GET /api/v1/agents — agent registry."""
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, Request
+
+from tokenjam.api.deps import require_api_key
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+@router.get("/agents")
+async def list_agents(request: Request) -> dict:
+    db = request.app.state.db
+    if not hasattr(db, "conn") or db.conn is None:
+        return {"agents": []}
+    rows = db.conn.execute(
+        "SELECT a.agent_id, a.first_seen, a.last_seen, "
+        "COALESCE(SUM(s.cost_usd), 0.0) AS lifetime_cost "
+        "FROM agents a LEFT JOIN spans s ON a.agent_id = s.agent_id "
+        "GROUP BY a.agent_id, a.first_seen, a.last_seen "
+        "ORDER BY a.last_seen DESC NULLS LAST"
+    ).fetchall()
+    return {
+        "agents": [
+            {
+                "agent_id": r[0],
+                "first_seen": r[1].isoformat() if r[1] else None,
+                "last_seen": r[2].isoformat() if r[2] else None,
+                "lifetime_cost_usd": float(r[3]),
+            }
+            for r in rows
+        ]
+    }

tokenjam/api/routes/alerts.py

@@ -0,0 +1,77 @@
+"""GET /api/v1/alerts — alert history."""
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, HTTPException, Request
+
+from tokenjam.api.deps import require_api_key
+from tokenjam.core.models import AlertFilters, AlertType, Severity
+from tokenjam.utils.time_parse import parse_since
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+
+@router.get("/alerts")
+async def get_alerts(
+    request: Request,
+    agent_id: str | None = None,
+    since: str | None = None,
+    severity: str | None = None,
+    type: str | None = None,
+    unread: bool = False,
+) -> dict:
+    db = request.app.state.db
+    try:
+        sev = Severity(severity) if severity else None
+    except ValueError:
+        raise HTTPException(status_code=422, detail=f"Invalid severity: {severity!r}")
+    try:
+        typ = AlertType(type) if type else None
+    except ValueError:
+        raise HTTPException(status_code=422, detail=f"Invalid type: {type!r}")
+    filters = AlertFilters(
+        agent_id=agent_id,
+        since=parse_since(since) if since else None,
+        severity=sev,
+        type=typ,
+        unread=unread,
+    )
+    alerts = db.get_alerts(filters)
+    return {
+        "alerts": [
+            {
+                "alert_id": a.alert_id,
+                "fired_at": a.fired_at.isoformat(),
+                "type": a.type.value,
+                "severity": a.severity.value,
+                "title": a.title,
+                "detail": a.detail,
+                "agent_id": a.agent_id,
+                "session_id": a.session_id,
+                "span_id": a.span_id,
+                "acknowledged": a.acknowledged,
+                "suppressed": a.suppressed,
+            }
+            for a in alerts
+        ],
+        "count": len(alerts),
+    }
+
+
+@router.patch("/alerts/{alert_id}/acknowledge")
+async def acknowledge_alert(alert_id: str, request: Request) -> dict:
+    db = request.app.state.db
+    conn = getattr(db, "conn", None)
+    if conn is None:
+        from fastapi import HTTPException
+        raise HTTPException(status_code=503, detail="Write operations unavailable in read-only mode")
+    result = conn.execute(
+        "SELECT alert_id FROM alerts WHERE alert_id = $1", [alert_id]
+    ).fetchone()
+    if result is None:
+        from fastapi import HTTPException
+        raise HTTPException(status_code=404, detail=f"Alert {alert_id} not found")
+    conn.execute(
+        "UPDATE alerts SET acknowledged = true WHERE alert_id = $1", [alert_id]
+    )
+    return {"acknowledged": True, "alert_id": alert_id}

tokenjam/api/routes/budget.py

@@ -0,0 +1,96 @@
+"""GET /api/v1/budget — read budgets. POST /api/v1/budget — update budgets."""
+from __future__ import annotations
+
+from pathlib import Path
+
+from fastapi import APIRouter, Depends, Request
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+
+from tokenjam.api.deps import require_api_key
+from tokenjam.core.config import (
+    AgentConfig,
+    BudgetConfig,
+    find_config_file,
+    resolve_effective_budget,
+    validate_budget_value,
+    write_config,
+)
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+def _budget_payload(config, agent_ids: list[str]) -> dict:
+    def _b(b):
+        return {"daily_usd": b.daily_usd, "session_usd": b.session_usd}
+
+    agents = {}
+    for aid in agent_ids:
+        agent_cfg = config.agents.get(aid)
+        raw = _b(agent_cfg.budget) if agent_cfg else _b(BudgetConfig())
+        eff = _b(resolve_effective_budget(aid, config))
+        agents[aid] = {"configured": raw, "effective": eff}
+
+    return {
+        "defaults": _b(config.defaults.budget),
+        "agents": agents,
+    }
+
+
+@router.get("/budget")
+async def get_budget(request: Request) -> dict:
+    config = request.app.state.config
+    db = request.app.state.db
+
+    # Merge: config agents + DB-observed agents
+    db_agent_ids: set[str] = set()
+    if hasattr(db, "conn"):
+        rows = db.conn.execute(
+            "SELECT DISTINCT agent_id FROM sessions ORDER BY agent_id"
+        ).fetchall()
+        db_agent_ids = {r[0] for r in rows}
+
+    all_agent_ids = sorted(set(config.agents) | db_agent_ids)
+    return _budget_payload(config, all_agent_ids)
+
+
+class BudgetUpdate(BaseModel):
+    scope: str           # "defaults" or an agent_id
+    daily_usd: float | None = None
+    session_usd: float | None = None
+
+
+@router.post("/budget")
+async def post_budget(request: Request, body: BudgetUpdate):
+    config = request.app.state.config
+    config_path_str = find_config_file()
+    if config_path_str is None:
+        return JSONResponse(status_code=400, content={"error": "No config file found"})
+
+    if body.scope == "defaults":
+        budget = config.defaults.budget
+    else:
+        if body.scope not in config.agents:
+            config.agents[body.scope] = AgentConfig()
+        budget = config.agents[body.scope].budget
+
+    try:
+        if body.daily_usd is not None:
+            budget.daily_usd = validate_budget_value(body.daily_usd, "daily_usd")
+        if body.session_usd is not None:
+            budget.session_usd = validate_budget_value(body.session_usd, "session_usd")
+    except ValueError as e:
+        return JSONResponse(status_code=400, content={"error": str(e)})
+
+    write_config(config, Path(config_path_str))
+
+    # Return updated payload with all known agents
+    db = request.app.state.db
+    db_agent_ids: set[str] = set()
+    if hasattr(db, "conn"):
+        rows = db.conn.execute(
+            "SELECT DISTINCT agent_id FROM sessions ORDER BY agent_id"
+        ).fetchall()
+        db_agent_ids = {r[0] for r in rows}
+    all_agent_ids = sorted(set(config.agents) | db_agent_ids)
+    return _budget_payload(config, all_agent_ids)

tokenjam/api/routes/cost.py

@@ -0,0 +1,43 @@
+"""GET /api/v1/cost — aggregated cost data."""
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, Request
+
+from tokenjam.api.deps import require_api_key
+from tokenjam.core.models import CostFilters
+from tokenjam.utils.time_parse import parse_since
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+@router.get("/cost")
+async def get_cost(
+    request: Request,
+    agent_id: str | None = None,
+    since: str | None = None,
+    until: str | None = None,
+    group_by: str = "day",
+) -> dict:
+    db = request.app.state.db
+    filters = CostFilters(
+        agent_id=agent_id,
+        since=parse_since(since) if since else None,
+        until=parse_since(until) if until else None,
+        group_by=group_by,
+    )
+    rows = db.get_cost_summary(filters)
+    total = sum(r.cost_usd for r in rows)
+    return {
+        "rows": [
+            {
+                "group": r.group,
+                "agent_id": r.agent_id,
+                "model": r.model,
+                "input_tokens": r.input_tokens,
+                "output_tokens": r.output_tokens,
+                "cost_usd": r.cost_usd,
+            }
+            for r in rows
+        ],
+        "total_cost_usd": total,
+    }

tokenjam/api/routes/drift.py

@@ -0,0 +1,63 @@
+"""GET /api/v1/drift — drift baseline and latest session comparison."""
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, Depends, Request
+
+from tokenjam.api.deps import require_api_key
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+def _build_agent_drift(db: Any, agent_id: str) -> dict:
+    """Build drift info dict for a single agent."""
+    baseline = db.get_baseline(agent_id)
+    if baseline is None:
+        return {"agent_id": agent_id, "baseline": None, "latest_session": None}
+
+    sessions = db.get_completed_sessions(agent_id, limit=1)
+    latest = None
+    if sessions:
+        s = sessions[0]
+        latest = {
+            "session_id": s.session_id,
+            "input_tokens": s.input_tokens,
+            "output_tokens": s.output_tokens,
+            "tool_call_count": s.tool_call_count,
+            "duration_seconds": s.duration_seconds,
+        }
+
+    return {
+        "agent_id": agent_id,
+        "baseline": {
+            "sessions_sampled": baseline.sessions_sampled,
+            "computed_at": baseline.computed_at.isoformat() if baseline.computed_at else None,
+            "avg_input_tokens": baseline.avg_input_tokens,
+            "stddev_input_tokens": baseline.stddev_input_tokens,
+            "avg_output_tokens": baseline.avg_output_tokens,
+            "stddev_output_tokens": baseline.stddev_output_tokens,
+            "avg_session_duration_s": baseline.avg_session_duration_s,
+            "stddev_session_duration": baseline.stddev_session_duration,
+            "avg_tool_call_count": baseline.avg_tool_call_count,
+            "stddev_tool_call_count": baseline.stddev_tool_call_count,
+        },
+        "latest_session": latest,
+    }
+
+
+@router.get("/drift", response_model=None)
+async def get_drift(request: Request, agent_id: str | None = None):
+    db = request.app.state.db
+
+    if agent_id:
+        return _build_agent_drift(db, agent_id)
+
+    # No agent_id: return drift info for all agents with baselines.
+    if not hasattr(db, "conn"):
+        return {"agents": []}
+    rows = db.conn.execute(
+        "SELECT DISTINCT agent_id FROM drift_baselines ORDER BY agent_id"
+    ).fetchall()
+    agents = [_build_agent_drift(db, row[0]) for row in rows]
+    return {"agents": agents}