zeno-cli 0.3.4__py3-none-any.whl

zeno_sdk/storage.py

@@ -0,0 +1,487 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import uuid
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import aiosqlite
+from alembic import command
+from alembic.config import Config
+
+from .types import AgentRun, BabysittingTaxPoint, LoadProbe, Session, SupervisionEvent
+
+
+def utcnow() -> datetime:
+    return datetime.now(tz=UTC)
+
+
+@dataclass(slots=True)
+class SessionStats:
+    session_id: str
+    avg_n_agents: float
+    intervention_rate: float
+    accept_rate: float
+    idle_ratio: float
+    composite_load: float | None
+    output_quality: float | None
+    # Behavioral anchors for RTLX-S Stage 2a discriminant-validity check
+    # (research 3, 2026-06-10). Each is a raw count or seconds; the API
+    # passes them through to the rtlxs_responses row.
+    interrupt_count: int = 0
+    agent_turns_count: int = 0
+    verification_seconds: int = 0
+
+
+# Per-gap cap on inter-run intervals when estimating verification_seconds.
+# A gap longer than this is presumed AFK (coffee, break, context switch) and
+# clipped to the cap so a single forgotten terminal does not balloon the
+# reviewer-time estimate. 10 minutes is generous for a real verification
+# pause and conservative for an AFK boundary; tune at Stage 2b if needed.
+_VERIFICATION_GAP_CAP_SECONDS: int = 600
+
+
+class ZenoStorage:
+    def __init__(self, db_path: Path) -> None:
+        self.db_path = db_path
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
+    async def migrate(self) -> None:
+        await asyncio.to_thread(self._upgrade_to_head, self._alembic_config())
+
+    def _alembic_config(self) -> Config:
+        # Resolve the alembic scripts in BOTH layouts: the uv-workspace src tree
+        # (scripts two levels up at packages/sdk-python/alembic) and a built wheel
+        # (scripts bundled inside the package at zeno_sdk/_migrations via the
+        # zeno-cli force-include). Package-relative first, workspace as fallback;
+        # otherwise an installed `zeno status` cannot find the migrations.
+        bundled = Path(__file__).resolve().parent / "_migrations"
+        root = bundled if (bundled / "alembic").is_dir() else Path(__file__).resolve().parents[2]
+        config = Config(str(root / "alembic.ini"))
+        config.set_main_option("script_location", str(root / "alembic"))
+        config.set_main_option("sqlalchemy.url", f"sqlite:///{self.db_path}")
+        return config
+
+    def _upgrade_to_head(self, config: Config) -> None:
+        """Apply migrations to head, self-healing a DB whose tables exist but
+        whose alembic_version is unset.
+
+        A migrate() interrupted mid-run leaves the DB wedged: SQLite auto-commits
+        each CREATE TABLE immediately, but alembic writes the version stamp only
+        after a whole revision succeeds. Kill the process in between and the
+        tables exist while alembic_version is empty, so a plain `alembic upgrade`
+        replays already-applied revisions and dies on "table ... already exists" -
+        which broke `zeno status` and every command that opens storage.
+
+        Recover without dropping data: find the latest revision whose tables are
+        ALL already present, stamp that revision, then `upgrade head` applies only
+        the genuinely-missing ones. Covers both a fully-built DB that lost its
+        stamp (stamp head, upgrade is a no-op) and a chain interrupted at any
+        revision (stamp the last complete one, apply the rest). The expensive
+        replay only runs on the rare unstamped-recovery path, never on a normally
+        stamped DB.
+        """
+        import sqlite3
+
+        if self.db_path.exists():
+            conn = sqlite3.connect(self.db_path)
+            try:
+                existing = {
+                    row[0]
+                    for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
+                    if not row[0].startswith("sqlite_")
+                }
+                if "alembic_version" in existing:
+                    stamped = conn.execute("SELECT COUNT(*) FROM alembic_version").fetchone()[0] > 0
+                else:
+                    stamped = False
+            finally:
+                conn.close()
+            data_tables = existing - {"alembic_version"}
+            if data_tables and not stamped:
+                target = self._latest_satisfied_revision(config, existing)
+                if target is not None:
+                    command.stamp(config, target)
+        command.upgrade(config, "head")
+
+    @staticmethod
+    def _latest_satisfied_revision(config: Config, existing_tables: set[str]) -> str | None:
+        """Return the newest revision whose cumulative table set is fully present
+        in ``existing_tables`` (so ``upgrade head`` then applies only the missing
+        revisions). Returns None if not even the base revision's tables are
+        present. Replays the chain on a throwaway DB to learn each revision's
+        cumulative table set; migrations are additive, so the sets are nested and
+        the first unsatisfied revision ends the search.
+        """
+        import sqlite3
+        import tempfile
+
+        from alembic.script import ScriptDirectory
+
+        script = ScriptDirectory.from_config(config)
+        revisions = [rev.revision for rev in reversed(list(script.walk_revisions()))]
+
+        tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
+        tmp.close()
+        probe_path = Path(tmp.name)
+        try:
+            probe = Config(config.config_file_name)
+            probe.set_main_option("script_location", config.get_main_option("script_location"))
+            probe.set_main_option("sqlalchemy.url", f"sqlite:///{probe_path}")
+            target: str | None = None
+            for revision in revisions:
+                command.upgrade(probe, revision)
+                conn = sqlite3.connect(probe_path)
+                try:
+                    rev_tables = {
+                        row[0]
+                        for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
+                        if not row[0].startswith("sqlite_")
+                    } - {"alembic_version"}
+                finally:
+                    conn.close()
+                if rev_tables <= existing_tables:
+                    target = revision
+                else:
+                    break
+            return target
+        finally:
+            probe_path.unlink(missing_ok=True)
+
+    async def insert_session(self, session: Session) -> None:
+        await self._execute(
+            """
+            INSERT OR REPLACE INTO sessions (
+              id, start_at, end_at, agent_count_max, harness, project_id, synced_at
+            ) VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                str(session.id),
+                session.start_at.isoformat(),
+                session.end_at.isoformat() if session.end_at else None,
+                session.agent_count_max,
+                session.harness,
+                str(session.project_id),
+                None,
+            ),
+        )
+
+    async def insert_agent_run(self, run: AgentRun) -> None:
+        await self._execute(
+            """
+            INSERT OR REPLACE INTO agent_runs (
+              id, session_id, harness, model, started_at, ended_at, outcome, synced_at
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                str(run.id),
+                str(run.session_id),
+                run.harness,
+                run.model,
+                run.started_at.isoformat(),
+                run.ended_at.isoformat() if run.ended_at else None,
+                run.outcome,
+                None,
+            ),
+        )
+
+    async def insert_event(self, event: SupervisionEvent) -> None:
+        await self._execute(
+            """
+            INSERT OR REPLACE INTO supervision_events (
+              id, session_id, agent_run_id, type, timestamp,
+              latency_ms_to_decide, metadata_json, synced_at
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                str(event.id),
+                str(event.session_id),
+                str(event.agent_run_id) if event.agent_run_id else None,
+                event.type,
+                event.timestamp.isoformat(),
+                event.latency_ms_to_decide,
+                json.dumps(event.metadata or {}, sort_keys=True),
+                None,
+            ),
+        )
+
+    async def insert_load_probe(self, probe: LoadProbe) -> None:
+        await self._execute(
+            """
+            INSERT OR REPLACE INTO load_probes (
+              id, session_id, prompted_at, responded_at, skipped, subscales_json, synced_at
+            ) VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                str(probe.id),
+                str(probe.session_id),
+                probe.prompted_at.isoformat(),
+                probe.responded_at.isoformat() if probe.responded_at else None,
+                int(probe.skipped),
+                json.dumps(
+                    probe.subscales.model_dump(exclude_none=True) if probe.subscales else None
+                ),
+                None,
+            ),
+        )
+
+    async def upsert_tax_point(self, point: BabysittingTaxPoint) -> None:
+        await self._execute(
+            """
+            INSERT OR REPLACE INTO babysitting_tax_points (
+              id, session_id, n_agents_active, composite_load, output_quality, synced_at
+            ) VALUES (?, ?, ?, ?, ?, ?)
+            """,
+            (
+                str(point.id),
+                str(point.session_id),
+                point.n_agents_active,
+                point.composite_load,
+                point.output_quality,
+                None,
+            ),
+        )
+
+    async def get_events_for_session(self, session_id: str) -> list[dict[str, Any]]:
+        rows = await self._fetchall(
+            """
+            SELECT id, session_id, agent_run_id, type, timestamp,
+                   latency_ms_to_decide, metadata_json
+            FROM supervision_events
+            WHERE session_id = ?
+            ORDER BY timestamp ASC
+            """,
+            (session_id,),
+        )
+        return [
+            {
+                "id": row[0],
+                "session_id": row[1],
+                "agent_run_id": row[2],
+                "type": row[3],
+                "timestamp": row[4],
+                "latency_ms_to_decide": row[5],
+                "metadata": json.loads(row[6] or "{}"),
+            }
+            for row in rows
+        ]
+
+    async def get_tax_curve(self, project_id: str) -> list[dict[str, Any]]:
+        rows = await self._fetchall(
+            """
+            SELECT b.n_agents_active, AVG(b.composite_load), AVG(b.output_quality)
+            FROM babysitting_tax_points b
+            JOIN sessions s ON s.id = b.session_id
+            WHERE s.project_id = ?
+            GROUP BY b.n_agents_active
+            ORDER BY b.n_agents_active ASC
+            """,
+            (project_id,),
+        )
+        return [
+            {
+                "n_agents_active": row[0],
+                "composite_load": float(row[1]),
+                "output_quality": float(row[2]) if row[2] is not None else None,
+            }
+            for row in rows
+        ]
+
+    async def get_session_points(self, project_id: str) -> list[dict[str, Any]]:
+        rows = await self._fetchall(
+            """
+            SELECT b.session_id, b.n_agents_active, b.composite_load, b.output_quality, s.end_at
+            FROM babysitting_tax_points b
+            JOIN sessions s ON s.id = b.session_id
+            WHERE s.project_id = ?
+            ORDER BY s.end_at ASC
+            """,
+            (project_id,),
+        )
+        return [
+            {
+                "session_id": row[0],
+                "n_agents_active": float(row[1]),
+                "composite_load": float(row[2]),
+                "output_quality": float(row[3]) if row[3] is not None else None,
+                "ended_at": row[4],
+            }
+            for row in rows
+        ]
+
+    async def latest_session_id(self, project_id: str) -> str | None:
+        row = await self._fetchone(
+            """
+            SELECT id
+            FROM sessions
+            WHERE project_id = ?
+            ORDER BY end_at DESC
+            LIMIT 1
+            """,
+            (project_id,),
+        )
+        if row is None:
+            return None
+        return str(row[0])
+
+    async def compute_session_stats(self, session_id: str) -> SessionStats:
+        event_rows = await self._fetchall(
+            """
+            SELECT type FROM supervision_events WHERE session_id = ?
+            """,
+            (session_id,),
+        )
+        n_events = len(event_rows)
+        interventions = sum(1 for row in event_rows if row[0] == "intervene")
+        accepts = sum(1 for row in event_rows if row[0] == "accept")
+        idle = sum(1 for row in event_rows if row[0] in {"idle_start", "idle_end"})
+
+        session_row = await self._fetchone(
+            "SELECT agent_count_max FROM sessions WHERE id = ?",
+            (session_id,),
+        )
+        run_count_row = await self._fetchone(
+            "SELECT COUNT(*) FROM agent_runs WHERE session_id = ?",
+            (session_id,),
+        )
+        persisted_agents = float(session_row[0]) if session_row else 0.0
+        observed_agents = float(run_count_row[0]) if run_count_row else 0.0
+        avg_n_agents = max(persisted_agents, observed_agents)
+
+        # Behavioral anchors for RTLX-S Stage 2a (research 3, 2026-06-10).
+        # verification_seconds is the sum of inter-run gaps (between the end
+        # of one agent_run and the start of the next), capped per gap to
+        # avoid counting AFK time. Stage 2 validation tests whether this
+        # proxy correlates positively with supervision_load.
+        agent_turns_count = int(observed_agents)
+        gap_rows = await self._fetchall(
+            """
+            SELECT started_at, ended_at FROM agent_runs
+            WHERE session_id = ? AND ended_at IS NOT NULL
+            ORDER BY started_at ASC
+            """,
+            (session_id,),
+        )
+        verification_seconds = 0
+        last_ended_at: datetime | None = None
+        for started_raw, ended_raw in gap_rows:
+            try:
+                started = datetime.fromisoformat(started_raw)
+                ended = datetime.fromisoformat(ended_raw)
+            except (TypeError, ValueError):
+                continue
+            if last_ended_at is not None:
+                gap = (started - last_ended_at).total_seconds()
+                if gap > 0:
+                    verification_seconds += int(min(gap, _VERIFICATION_GAP_CAP_SECONDS))
+            last_ended_at = ended
+
+        load_row = await self._fetchone(
+            """
+            SELECT subscales_json
+            FROM load_probes
+            WHERE session_id = ?
+            ORDER BY prompted_at DESC
+            LIMIT 1
+            """,
+            (session_id,),
+        )
+        composite_load: float | None = None
+        if load_row and load_row[0]:
+            payload = json.loads(load_row[0])
+            if isinstance(payload, dict):
+                # Average only numeric subscale values; tolerate nulls / mixed
+                # instruments (RTLX-S 5-item vs legacy 10-item) in the column.
+                vals = [v for v in payload.values() if isinstance(v, (int, float))]
+                composite_load = sum(vals) / len(vals) if vals else None
+
+        run_rows = await self._fetchall(
+            "SELECT outcome FROM agent_runs WHERE session_id = ?",
+            (session_id,),
+        )
+        total_runs = len(run_rows)
+        quality_score = None
+        if total_runs:
+            quality_score = (
+                sum(1 for row in run_rows if row[0] in {"accepted", "merged"}) / total_runs
+            )
+
+        return SessionStats(
+            session_id=session_id,
+            avg_n_agents=avg_n_agents,
+            intervention_rate=interventions / n_events if n_events else 0.0,
+            accept_rate=accepts / n_events if n_events else 0.0,
+            idle_ratio=idle / n_events if n_events else 0.0,
+            composite_load=composite_load,
+            output_quality=quality_score,
+            interrupt_count=interventions,
+            agent_turns_count=agent_turns_count,
+            verification_seconds=verification_seconds,
+        )
+
+    async def find_recent_activity_session_id(self, *, window_minutes: int = 120) -> str | None:
+        """Most-recent session_id with at least one agent_run, in window.
+
+        Used by the survey CLI to attach behavioral anchors to the RTLX-S
+        row (research 3, 2026-06-10). The survey command opens a fresh
+        session for its own bookkeeping, so the work-session telemetry has
+        to be looked up explicitly. Defaults to a 2-hour window so a survey
+        immediately after a long work session still resolves cleanly; tune
+        if a stale lookup ever attaches anchors to the wrong session.
+
+        Returns None if no qualifying session is found.
+        """
+        cutoff = datetime.now(tz=UTC).timestamp() - window_minutes * 60
+        row = await self._fetchone(
+            """
+            SELECT s.id
+            FROM sessions s
+            JOIN agent_runs r ON r.session_id = s.id
+            WHERE s.start_at >= ?
+            GROUP BY s.id
+            HAVING COUNT(r.id) > 0
+            ORDER BY MAX(r.started_at) DESC
+            LIMIT 1
+            """,
+            (datetime.fromtimestamp(cutoff, tz=UTC).isoformat(),),
+        )
+        if row is None:
+            return None
+        return str(row[0])
+
+    async def list_tables(self) -> list[str]:
+        rows = await self._fetchall(
+            "SELECT name FROM sqlite_master WHERE type = 'table' ORDER BY name",
+            (),
+        )
+        return [row[0] for row in rows]
+
+    async def _execute(self, query: str, params: tuple[Any, ...]) -> None:
+        async with aiosqlite.connect(self.db_path) as conn:
+            await conn.execute(query, params)
+            await conn.commit()
+
+    async def _fetchall(self, query: str, params: tuple[Any, ...]) -> list[tuple[Any, ...]]:
+        async with aiosqlite.connect(self.db_path) as conn:
+            cursor = await conn.execute(query, params)
+            rows = await cursor.fetchall()
+            return rows
+
+    async def _fetchone(self, query: str, params: tuple[Any, ...]) -> tuple[Any, ...] | None:
+        async with aiosqlite.connect(self.db_path) as conn:
+            cursor = await conn.execute(query, params)
+            return await cursor.fetchone()
+
+
+def make_tax_point(session_id: str, stats: SessionStats) -> BabysittingTaxPoint:
+    composite = stats.composite_load if stats.composite_load is not None else 0.0
+    return BabysittingTaxPoint(
+        id=uuid.uuid4(),
+        session_id=uuid.UUID(session_id),
+        n_agents_active=max(0, int(round(stats.avg_n_agents))),
+        composite_load=composite,
+        output_quality=stats.output_quality,
+    )

zeno_sdk/types/__init__.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Literal
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class AgentRun(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    session_id: UUID
+    harness: str
+    model: str
+    started_at: datetime
+    ended_at: datetime | None = None
+    outcome: Literal["accepted", "rejected", "edited", "discarded", "merged", "unknown"]
+
+
+class BabysittingTaxPoint(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    session_id: UUID
+    n_agents_active: int = Field(ge=0)
+    composite_load: float = Field(ge=0)
+    output_quality: float | None = None
+
+
+class LoadProbeSubscales(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    mental_demand: int | None = Field(default=None, ge=0, le=100)
+    effort: int | None = Field(default=None, ge=0, le=100)
+    frustration: int | None = Field(default=None, ge=0, le=100)
+    supervision_load: int | None = Field(default=None, ge=0, le=100)
+    execution_load: int | None = Field(default=None, ge=0, le=100)
+    physical_demand: int | None = Field(default=None, ge=0, le=100)
+    temporal_demand: int | None = Field(default=None, ge=0, le=100)
+    performance: int | None = Field(default=None, ge=0, le=100)
+    trust_calibration_demand: int | None = Field(default=None, ge=0, le=100)
+    oversight_demand: int | None = Field(default=None, ge=0, le=100)
+    integration_demand: int | None = Field(default=None, ge=0, le=100)
+    intervention_decision_demand: int | None = Field(default=None, ge=0, le=100)
+
+
+class LoadProbe(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    session_id: UUID
+    prompted_at: datetime
+    responded_at: datetime | None = None
+    skipped: bool
+    subscales: LoadProbeSubscales | None = None
+
+
+class OutcomeMetric(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    agent_run_id: UUID
+    tests_passed: bool
+    merged: bool
+    reverted: bool
+    time_to_resolve_ms: int | None = Field(default=None, ge=0)
+
+
+class Project(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    name: str
+    created_at: datetime
+
+
+class Session(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    start_at: datetime
+    end_at: datetime | None = None
+    agent_count_max: int = Field(ge=1)
+    harness: str
+    project_id: UUID
+
+
+class SupervisionEvent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    session_id: UUID
+    agent_run_id: UUID | None = None
+    type: Literal[
+        "accept",
+        "reject",
+        "edit",
+        "intervene",
+        "switch_focus",
+        "idle_start",
+        "idle_end",
+        "agent_spawn",
+        "agent_terminate",
+    ]
+    timestamp: datetime
+    latency_ms_to_decide: int | None = Field(default=None, ge=0)
+    metadata: dict[str, Any] | None = None
+
+
+class User(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: UUID
+    email: str
+    tier: Literal["free", "pro", "team"]
+    created_at: datetime
+
+
+__all__ = [
+    "AgentRun",
+    "BabysittingTaxPoint",
+    "LoadProbe",
+    "OutcomeMetric",
+    "Project",
+    "Session",
+    "SupervisionEvent",
+    "User",
+]

zeno_session_intel/__init__.py

@@ -0,0 +1,19 @@
+"""zeno-session-intel: passive multi-tool coding-agent session intelligence.
+
+A read-only ingester that parses coding-agent session files (Claude Code, Codex)
+into zeno's local SQLite using an agentsview-shaped schema, plus ccusage-class
+analytics (cache-aware token/cost, tool-mix, session archetypes, peak-context).
+
+Stdlib-only, zero runtime dependencies: importable by the SDK/CLI, by tests, and
+by the dashboard's stdlib export scripts alike. Writes via the canonical schema in
+``schema.py`` (mirrored by alembic ``0004_transcript_intelligence`` and guarded by a
+schema-drift test). The capture is ADDITIVE: it sits alongside zeno's per-turn
+cognition signal (the cc-bridge hook), never replacing it.
+
+Session-intelligence schema + analytics design is adapted from agentsview
+(MIT-licensed). See THIRD_PARTY_LICENSES.md.
+"""
+
+__all__ = ["__version__"]
+
+__version__ = "0.0.0"