tokenjam 0.2.0__py3-none-any.whl

tokenjam/api/routes/spans.py

@@ -0,0 +1,202 @@
+"""POST /api/v1/spans — OTLP JSON span ingest endpoint."""
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
+
+from tokenjam.core.ingest import SpanRejectedError
+from tokenjam.core.models import NormalizedSpan, SpanKind, SpanStatus
+from tokenjam.otel.semconv import GenAIAttributes, TjAttributes
+from tokenjam.utils.ids import new_span_id
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.post("/spans")
+async def ingest_spans(request: Request) -> JSONResponse:
+    """
+    Accept a batch of spans in OTLP JSON format.
+    Auth is enforced by IngestAuthMiddleware.
+    Returns 200 even on partial rejection; 400 only if body is entirely malformed.
+    """
+    try:
+        body = await request.json()
+    except Exception:
+        return JSONResponse(
+            status_code=400,
+            content={"error": "Invalid JSON body"},
+        )
+
+    if not isinstance(body, dict) or "resourceSpans" not in body:
+        return JSONResponse(
+            status_code=400,
+            content={"error": "Expected OTLP JSON with 'resourceSpans' key"},
+        )
+
+    pipeline = request.app.state.pipeline
+    ingested = 0
+    rejections: list[dict[str, str]] = []
+
+    for resource_span in body.get("resourceSpans", []):
+        resource_attrs = _extract_resource_attrs(resource_span)
+        for scope_span in resource_span.get("scopeSpans", []):
+            for raw_span in scope_span.get("spans", []):
+                span_id = raw_span.get("spanId", new_span_id())
+                try:
+                    span = _parse_span(raw_span, resource_attrs)
+                    pipeline.process(span)
+                    ingested += 1
+                except SpanRejectedError as exc:
+                    rejections.append({"span_id": span_id, "reason": str(exc)})
+                except Exception as exc:
+                    logger.warning("Failed to process span %s: %s", span_id, exc)
+                    rejections.append({"span_id": span_id, "reason": str(exc)})
+
+    return JSONResponse(
+        status_code=200,
+        content={
+            "ingested": ingested,
+            "rejected": len(rejections),
+            "rejections": rejections,
+        },
+    )
+
+
+def _extract_resource_attrs(resource_span: dict) -> dict[str, Any]:
+    """Pull flat attributes from the resource section."""
+    resource = resource_span.get("resource", {})
+    attrs: dict[str, Any] = {}
+    for attr in resource.get("attributes", []):
+        key = attr.get("key", "")
+        value = _otlp_value(attr.get("value", {}))
+        if key and value is not None:
+            attrs[key] = value
+    return attrs
+
+
+def _parse_span(raw: dict, resource_attrs: dict[str, Any]) -> NormalizedSpan:
+    """Convert an OTLP JSON span dict into a NormalizedSpan."""
+    # Merge resource + span attributes
+    attrs: dict[str, Any] = dict(resource_attrs)
+    for attr in raw.get("attributes", []):
+        key = attr.get("key", "")
+        value = _otlp_value(attr.get("value", {}))
+        if key and value is not None:
+            attrs[key] = value
+
+    # Parse timestamps (OTLP uses nanoseconds as strings)
+    start_ns = int(raw.get("startTimeUnixNano", 0))
+    end_ns = int(raw.get("endTimeUnixNano", 0))
+    start_time = datetime.fromtimestamp(start_ns / 1e9, tz=timezone.utc) if start_ns else datetime.now(tz=timezone.utc)
+    end_time = datetime.fromtimestamp(end_ns / 1e9, tz=timezone.utc) if end_ns else None
+
+    duration_ms = None
+    if start_ns and end_ns:
+        duration_ms = (end_ns - start_ns) / 1e6
+
+    # Parse status
+    status_raw = raw.get("status", {})
+    status_code_int = status_raw.get("code", 0)
+    status_map = {0: SpanStatus.UNSET, 1: SpanStatus.OK, 2: SpanStatus.ERROR}
+    status_code = status_map.get(status_code_int, SpanStatus.UNSET)
+
+    # Parse kind
+    kind_int = raw.get("kind", 1)
+    kind_map = {
+        1: SpanKind.INTERNAL, 2: SpanKind.SERVER, 3: SpanKind.CLIENT,
+        4: SpanKind.PRODUCER, 5: SpanKind.CONSUMER,
+    }
+    kind = kind_map.get(kind_int, SpanKind.INTERNAL)
+
+    # --- OpenClaw / generic OTLP attribute enrichment ---
+    span_name = raw.get("name", "")
+    agent_id = attrs.get(GenAIAttributes.AGENT_ID)
+    tool_name = attrs.get(GenAIAttributes.TOOL_NAME)
+    provider = attrs.get(GenAIAttributes.PROVIDER_NAME)
+
+    # Fall back to service.name for agent_id (OpenClaw sets service.name)
+    if not agent_id:
+        agent_id = attrs.get("service.name") or None
+
+    # Fall back to gen_ai.system for provider (OpenClaw uses this)
+    if not provider:
+        provider = attrs.get("gen_ai.system") or None
+
+    # Extract tool_name from span names like "tool.Read", "tool.exec"
+    if not tool_name and span_name.startswith("tool."):
+        tool_name = span_name[5:]  # strip "tool." prefix
+
+    return NormalizedSpan(
+        span_id=raw.get("spanId", new_span_id()),
+        trace_id=raw.get("traceId", ""),
+        name=span_name,
+        kind=kind,
+        status_code=status_code,
+        status_message=status_raw.get("message"),
+        start_time=start_time,
+        end_time=end_time,
+        duration_ms=duration_ms,
+        parent_span_id=raw.get("parentSpanId"),
+        attributes=attrs,
+        events=[
+            {"name": e.get("name", ""), "time": e.get("timeUnixNano"),
+             "attributes": {a.get("key", ""): _otlp_value(a.get("value", {})) for a in e.get("attributes", [])}}
+            for e in raw.get("events", [])
+        ],
+        # Extract indexed fields from merged attributes
+        agent_id=agent_id,
+        provider=provider,
+        model=attrs.get(GenAIAttributes.REQUEST_MODEL),
+        tool_name=tool_name,
+        input_tokens=_safe_int(attrs.get(GenAIAttributes.INPUT_TOKENS)),
+        output_tokens=_safe_int(attrs.get(GenAIAttributes.OUTPUT_TOKENS)),
+        cache_tokens=_safe_int(attrs.get(GenAIAttributes.CACHE_READ_TOKENS)),
+        cost_usd=_safe_float(attrs.get(TjAttributes.COST_USD)),
+        request_type=attrs.get(GenAIAttributes.REQUEST_TYPE),
+        conversation_id=attrs.get(GenAIAttributes.CONVERSATION_ID),
+        session_id=attrs.get("session.id"),
+    )
+
+
+def _otlp_value(v: dict) -> Any:
+    """Extract a value from an OTLP AttributeValue wrapper."""
+    if "stringValue" in v:
+        return v["stringValue"]
+    if "intValue" in v:
+        return int(v["intValue"])
+    if "doubleValue" in v:
+        return float(v["doubleValue"])
+    if "boolValue" in v:
+        return v["boolValue"]
+    if "arrayValue" in v:
+        return [_otlp_value(item) for item in v["arrayValue"].get("values", [])]
+    if "kvlistValue" in v:
+        return {
+            kv["key"]: _otlp_value(kv["value"])
+            for kv in v["kvlistValue"].get("values", [])
+        }
+    return None
+
+
+def _safe_int(v: Any) -> int | None:
+    if v is None:
+        return None
+    try:
+        return int(v)
+    except (TypeError, ValueError):
+        return None
+
+
+def _safe_float(v: Any) -> float | None:
+    if v is None:
+        return None
+    try:
+        return float(v)
+    except (TypeError, ValueError):
+        return None

tokenjam/api/routes/status.py

@@ -0,0 +1,84 @@
+"""GET /api/v1/status — agent status overview."""
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, Request
+
+from tokenjam.api.deps import require_api_key
+from tokenjam.core.db import _row_to_session
+from tokenjam.core.models import AlertFilters
+from tokenjam.utils.time_parse import utcnow
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+@router.get("/status")
+async def get_status(
+    request: Request,
+    agent_id: str | None = None,
+) -> dict:
+    db = request.app.state.db
+
+    # Discover agent IDs
+    if agent_id:
+        agent_ids = [agent_id]
+    elif hasattr(db, "conn"):
+        rows = db.conn.execute(
+            "SELECT DISTINCT agent_id FROM sessions WHERE agent_id IS NOT NULL "
+            "UNION "
+            "SELECT DISTINCT agent_id FROM spans WHERE agent_id IS NOT NULL "
+            "ORDER BY agent_id"
+        ).fetchall()
+        agent_ids = [r[0] for r in rows]
+    else:
+        agent_ids = []
+
+    has_active_alerts = False
+    agents_data = []
+
+    for aid in agent_ids:
+        session = None
+
+        # Check for active session first, then fall back to latest completed
+        if hasattr(db, "conn"):
+            active_rows = db.conn.execute(
+                "SELECT * FROM sessions WHERE agent_id = $1 AND status = 'active' "
+                "ORDER BY started_at DESC LIMIT 1",
+                [aid],
+            ).fetchall()
+            if active_rows:
+                cols = [d[0] for d in db.conn.description]
+                session = _row_to_session(active_rows[0], cols)
+
+        if session is None:
+            sessions = db.get_completed_sessions(aid, limit=1)
+            if sessions:
+                session = sessions[0]
+
+        today_cost = db.get_daily_cost(aid, utcnow().date())
+
+        # Active (unacknowledged, unsuppressed) alerts
+        alerts = db.get_alerts(AlertFilters(agent_id=aid, unread=True, limit=50))
+        active_alerts = [a for a in alerts if not a.acknowledged and not a.suppressed]
+        if active_alerts:
+            has_active_alerts = True
+
+        agent_data = {
+            "agent_id": aid,
+            "status": session.status if session else "idle",
+            "session_id": session.session_id if session else None,
+            "cost_today": today_cost,
+            "input_tokens": session.input_tokens if session else 0,
+            "output_tokens": session.output_tokens if session else 0,
+            "tool_call_count": session.tool_call_count if session else 0,
+            "error_count": session.error_count if session else 0,
+            "active_alerts": len(active_alerts),
+            "duration_seconds": session.duration_seconds if session else None,
+            "started_at": session.started_at.isoformat() if session and session.started_at else None,
+            "total_cost_usd": float(session.total_cost_usd) if session and session.total_cost_usd is not None else 0.0,
+        }
+        agents_data.append(agent_data)
+
+    return {
+        "agents": agents_data,
+        "has_active_alerts": has_active_alerts,
+    }

tokenjam/api/routes/tools.py

@@ -0,0 +1,22 @@
+"""GET /api/v1/tools — tool call records with aggregated stats."""
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, Request
+
+from tokenjam.api.deps import require_api_key
+from tokenjam.utils.time_parse import parse_since
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+@router.get("/tools")
+async def get_tools(
+    request: Request,
+    agent_id: str | None = None,
+    since: str | None = None,
+    tool_name: str | None = None,
+) -> dict:
+    db = request.app.state.db
+    since_dt = parse_since(since) if since else None
+    rows = db.get_tool_calls(agent_id, since_dt, tool_name)
+    return {"tools": rows, "count": len(rows)}

tokenjam/api/routes/traces.py

@@ -0,0 +1,92 @@
+"""GET /api/v1/traces — trace listing and detail."""
+from __future__ import annotations
+
+
+from fastapi import APIRouter, Depends, Request
+
+from tokenjam.api.deps import require_api_key
+from tokenjam.core.models import TraceFilters
+from tokenjam.utils.time_parse import parse_since
+
+router = APIRouter(dependencies=[Depends(require_api_key)])
+
+
+@router.get("/traces")
+async def list_traces(
+    request: Request,
+    agent_id: str | None = None,
+    since: str | None = None,
+    until: str | None = None,
+    limit: int = 50,
+    offset: int = 0,
+    status: str | None = None,
+    span_name: str | None = None,
+) -> dict:
+    db = request.app.state.db
+    filters = TraceFilters(
+        agent_id=agent_id,
+        since=parse_since(since) if since else None,
+        until=parse_since(until) if until else None,
+        limit=limit,
+        offset=offset,
+        status=status,
+        span_name=span_name,
+    )
+    traces = db.get_traces(filters)
+    return {
+        "traces": [
+            {
+                "trace_id": t.trace_id,
+                "agent_id": t.agent_id,
+                "name": t.name,
+                "start_time": t.start_time.isoformat() if t.start_time else None,
+                "duration_ms": t.duration_ms,
+                "cost_usd": t.cost_usd,
+                "status_code": t.status_code,
+                "span_count": t.span_count,
+            }
+            for t in traces
+        ],
+        "count": len(traces),
+    }
+
+
+@router.get("/traces/{trace_id}")
+async def get_trace(request: Request, trace_id: str) -> dict:
+    db = request.app.state.db
+    spans = db.get_trace_spans(trace_id)
+    return {
+        "trace_id": trace_id,
+        "spans": [_span_to_dict(s) for s in spans],
+        "span_count": len(spans),
+    }
+
+
+def _span_to_dict(span: object) -> dict:
+    """Serialise a NormalizedSpan to a JSON-safe dict."""
+    from tokenjam.core.models import NormalizedSpan
+    assert isinstance(span, NormalizedSpan)
+    return {
+        "span_id": span.span_id,
+        "trace_id": span.trace_id,
+        "parent_span_id": span.parent_span_id,
+        "name": span.name,
+        "kind": span.kind.value,
+        "status_code": span.status_code.value,
+        "status_message": span.status_message,
+        "start_time": span.start_time.isoformat() if span.start_time else None,
+        "end_time": span.end_time.isoformat() if span.end_time else None,
+        "duration_ms": span.duration_ms,
+        "agent_id": span.agent_id,
+        "session_id": span.session_id,
+        "provider": span.provider,
+        "model": span.model,
+        "tool_name": span.tool_name,
+        "input_tokens": span.input_tokens,
+        "output_tokens": span.output_tokens,
+        "cache_tokens": span.cache_tokens,
+        "cost_usd": span.cost_usd,
+        "request_type": span.request_type,
+        "conversation_id": span.conversation_id,
+        "attributes": span.attributes,
+    }

tokenjam/cli/cmd_alerts.py

@@ -0,0 +1,94 @@
+import json
+
+import click
+
+from tokenjam.core.models import AlertFilters, AlertType, Severity
+from tokenjam.utils.formatting import console, make_table, severity_colour
+from tokenjam.utils.time_parse import parse_since
+
+
+@click.command("alerts")
+@click.option("--agent", default=None, help="Filter to specific agent_id")
+@click.option("--since", default="24h", help="Time window (e.g. 1h, 24h, 7d)")
+@click.option(
+    "--severity",
+    type=click.Choice(["critical", "warning", "info"]),
+    default=None,
+    help="Filter by minimum severity",
+)
+@click.option("--type", "alert_type", default=None, help="Filter by alert type")
+@click.option("--unread", is_flag=True, help="Show only unacknowledged alerts")
+@click.option("--json", "output_json", is_flag=True, help="Output as JSON")
+@click.pass_context
+def cmd_alerts(
+    ctx: click.Context,
+    agent: str | None,
+    since: str,
+    severity: str | None,
+    alert_type: str | None,
+    unread: bool,
+    output_json: bool,
+) -> None:
+    """Show alert history."""
+    db = ctx.obj["db"]
+    try:
+        since_dt = parse_since(since)
+    except ValueError as exc:
+        raise click.BadParameter(str(exc), param_hint="'--since'") from exc
+    filters = AlertFilters(
+        agent_id=agent,
+        since=since_dt,
+        severity=Severity(severity) if severity else None,
+        type=AlertType(alert_type) if alert_type else None,
+        unread=unread,
+    )
+    alerts = db.get_alerts(filters)
+
+    if output_json:
+        click.echo(json.dumps(
+            [
+                {
+                    "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
+            ],
+            default=str,
+        ))
+        return
+
+    if not alerts:
+        console.print("[dim]No alerts found for the given filters.[/dim]")
+        return
+
+    critical_count = sum(1 for a in alerts if a.severity == Severity.CRITICAL)
+    warning_count = sum(1 for a in alerts if a.severity == Severity.WARNING)
+    console.print(
+        f"[bold]Alerts \u2014 last {since}[/bold]   "
+        f"({len(alerts)} total: {critical_count} critical, {warning_count} warning)"
+    )
+
+    table = make_table("TIME", "SEVERITY", "TYPE", "AGENT", "DETAIL")
+    for a in alerts:
+        time_str = a.fired_at.strftime("%H:%M:%S")
+        colour = severity_colour(a.severity.value)
+        detail_msg = a.detail.get("message", a.title)
+        if len(detail_msg) > 60:
+            detail_msg = detail_msg[:57] + "..."
+        table.add_row(
+            time_str,
+            f"[{colour}]{a.severity.value.upper()}[/]",
+            a.type.value,
+            a.agent_id or "-",
+            detail_msg,
+        )
+    console.print(table)

tokenjam/cli/cmd_budget.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import click
+
+from tokenjam.core.config import (
+    AgentConfig,
+    find_config_file,
+    resolve_effective_budget,
+    validate_budget_value,
+    write_config,
+)
+from tokenjam.utils.formatting import console
+
+
+@click.command("budget")
+@click.option("--agent", default=None, help="Agent ID to target (omit for global defaults)")
+@click.option("--daily", "daily_usd", type=float, default=None,
+              help="Daily budget in USD (0 = remove limit)")
+@click.option("--session", "session_usd", type=float, default=None,
+              help="Per-session budget in USD (0 = remove limit)")
+@click.pass_context
+def cmd_budget(
+    ctx: click.Context,
+    agent: str | None,
+    daily_usd: float | None,
+    session_usd: float | None,
+) -> None:
+    """View or set cost budgets for agents."""
+    config = ctx.obj["config"]
+    writing = daily_usd is not None or session_usd is not None
+
+    if not writing:
+        _show_budgets(config, ctx.obj.get("db"))
+        return
+
+    # Write mode — find config file on disk
+    config_path_str = find_config_file()
+    if config_path_str is None:
+        raise click.ClickException(
+            "No config file found. Run 'tj onboard' to create one."
+        )
+
+    if agent:
+        if agent not in config.agents:
+            config.agents[agent] = AgentConfig()
+        budget = config.agents[agent].budget
+        scope = f"agent '{agent}'"
+    else:
+        budget = config.defaults.budget
+        scope = "global defaults"
+
+    try:
+        if daily_usd is not None:
+            budget.daily_usd = validate_budget_value(daily_usd, "daily_usd")
+        if session_usd is not None:
+            budget.session_usd = validate_budget_value(session_usd, "session_usd")
+    except ValueError as e:
+        raise click.ClickException(str(e))
+
+    write_config(config, Path(config_path_str))
+
+    console.print(f"[green]\u2713[/green] Budget updated for {scope}")
+    if daily_usd is not None:
+        val = f"${daily_usd:.2f}" if daily_usd > 0 else "no limit"
+        console.print(f"  Daily:   {val}")
+    if session_usd is not None:
+        val = f"${session_usd:.2f}" if session_usd > 0 else "no limit"
+        console.print(f"  Session: {val}")
+
+
+def _show_budgets(config, db) -> None:
+    from rich.table import Table
+
+    table = Table(show_header=True, header_style="bold", box=None, padding=(0, 2))
+    table.add_column("Scope")
+    table.add_column("Daily", justify="right")
+    table.add_column("Session", justify="right")
+
+    def _fmt(val: float | None) -> str:
+        return f"${val:.2f}" if val is not None else "[dim]no limit[/dim]"
+
+    def _fmt_effective(raw: float | None, effective: float | None) -> str:
+        if raw is not None:
+            return f"${raw:.2f}"
+        if effective is not None:
+            return f"[dim]${effective:.2f}[/dim] [dim](default)[/dim]"
+        return "[dim]no limit[/dim]"
+
+    table.add_row(
+        "[bold]defaults[/bold]",
+        _fmt(config.defaults.budget.daily_usd),
+        _fmt(config.defaults.budget.session_usd),
+    )
+
+    # Merge agent IDs from config + DB-observed agents
+    agent_ids = set(config.agents)
+    if db is not None and hasattr(db, "conn"):
+        rows = db.conn.execute(
+            "SELECT DISTINCT agent_id FROM sessions ORDER BY agent_id"
+        ).fetchall()
+        agent_ids |= {r[0] for r in rows}
+
+    for agent_id in sorted(agent_ids):
+        agent_cfg = config.agents.get(agent_id)
+        eff = resolve_effective_budget(agent_id, config)
+        raw_daily = agent_cfg.budget.daily_usd if agent_cfg else None
+        raw_session = agent_cfg.budget.session_usd if agent_cfg else None
+        table.add_row(
+            agent_id,
+            _fmt_effective(raw_daily, eff.daily_usd),
+            _fmt_effective(raw_session, eff.session_usd),
+        )
+
+    if not agent_ids:
+        table.add_row("[dim]no agents configured[/dim]", "", "")
+
+    console.print(table)