Fixed and improved token tracking

migrations/versions/v9q0r1s2t3u4_token_event_provenance.py

@@ -0,0 +1,128 @@
+"""add token event provenance fields
+
+Revision ID: v9q0r1s2t3u4
+Revises: u8p9q0r1s2t3
+Create Date: 2026-05-23
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+revision = "v9q0r1s2t3u4"
+down_revision = "u8p9q0r1s2t3"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "token_events",
+        sa.Column("measurement_kind", sa.Text(), nullable=False, server_default="estimated"),
+    )
+    op.add_column(
+        "token_events",
+        sa.Column("source_provider", sa.Text(), nullable=False, server_default="manual"),
+    )
+    op.add_column("token_events", sa.Column("source_id", sa.Text(), nullable=True))
+    op.add_column("token_events", sa.Column("source_path", sa.Text(), nullable=True))
+    op.add_column(
+        "token_events",
+        sa.Column("source_created_at", sa.TIMESTAMP(timezone=True), nullable=True),
+    )
+    op.add_column(
+        "token_events",
+        sa.Column("ingested_at", sa.TIMESTAMP(timezone=True), nullable=False, server_default=sa.text("now()")),
+    )
+    op.add_column("token_events", sa.Column("parser_version", sa.Text(), nullable=True))
+    op.add_column(
+        "token_events",
+        sa.Column("confidence", sa.Float(), nullable=False, server_default="0.35"),
+    )
+    op.add_column(
+        "token_events",
+        sa.Column("cached_input_tokens", sa.Integer(), nullable=False, server_default="0"),
+    )
+    op.add_column(
+        "token_events",
+        sa.Column("reasoning_output_tokens", sa.Integer(), nullable=False, server_default="0"),
+    )
+    op.add_column("token_events", sa.Column("raw_total_tokens", sa.Integer(), nullable=True))
+    op.add_column("token_events", sa.Column("cost_estimated_usd", sa.Float(), nullable=True))
+    op.add_column(
+        "token_events",
+        sa.Column(
+            "raw_metadata",
+            postgresql.JSONB(astext_type=sa.Text()),
+            nullable=False,
+            server_default=sa.text("'{}'::jsonb"),
+        ),
+    )
+
+    op.execute(
+        """
+        UPDATE token_events
+        SET
+          measurement_kind = CASE
+            WHEN note = 'heuristic_superseded_by_codex_backfill' THEN 'superseded'
+            WHEN note = 'workplan' THEN 'allocated'
+            WHEN note = 'heuristic' THEN 'estimated'
+            WHEN note = 'measured' OR note LIKE 'backfill:codex-session%' THEN 'measured'
+            ELSE measurement_kind
+          END,
+          source_provider = CASE
+            WHEN note = 'heuristic' THEN 'task_fallback'
+            WHEN note LIKE 'backfill:codex-session%' OR ref_id LIKE 'codex:%' THEN 'codex_session'
+            WHEN note = 'measured' AND agent ILIKE '%claude%' THEN 'claude_transcript'
+            ELSE source_provider
+          END,
+          source_id = CASE
+            WHEN source_id IS NULL AND (note LIKE 'backfill:codex-session%' OR ref_id LIKE 'codex:%')
+              THEN ref_id
+            ELSE source_id
+          END,
+          raw_total_tokens = CASE
+            WHEN raw_total_tokens IS NULL THEN tokens_in + tokens_out
+            ELSE raw_total_tokens
+          END,
+          confidence = CASE
+            WHEN note = 'heuristic_superseded_by_codex_backfill' THEN 0.0
+            WHEN note = 'heuristic' THEN 0.35
+            WHEN note = 'workplan' THEN 0.70
+            WHEN note = 'measured' OR note LIKE 'backfill:codex-session%' THEN 1.0
+            ELSE confidence
+          END
+        """
+    )
+
+    op.create_index("ix_token_events_measurement_kind", "token_events", ["measurement_kind"])
+    op.create_index("ix_token_events_source_provider", "token_events", ["source_provider"])
+    op.create_index("ix_token_events_source_id", "token_events", ["source_id"])
+    op.create_index("ix_token_events_source_created_at", "token_events", ["source_created_at"])
+    op.create_index("ix_token_events_ingested_at", "token_events", ["ingested_at"])
+    op.create_unique_constraint(
+        "uq_token_events_source_identity",
+        "token_events",
+        ["measurement_kind", "source_provider", "source_id"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_constraint("uq_token_events_source_identity", "token_events", type_="unique")
+    op.drop_index("ix_token_events_ingested_at", table_name="token_events")
+    op.drop_index("ix_token_events_source_created_at", table_name="token_events")
+    op.drop_index("ix_token_events_source_id", table_name="token_events")
+    op.drop_index("ix_token_events_source_provider", table_name="token_events")
+    op.drop_index("ix_token_events_measurement_kind", table_name="token_events")
+    op.drop_column("token_events", "raw_metadata")
+    op.drop_column("token_events", "cost_estimated_usd")
+    op.drop_column("token_events", "raw_total_tokens")
+    op.drop_column("token_events", "reasoning_output_tokens")
+    op.drop_column("token_events", "cached_input_tokens")
+    op.drop_column("token_events", "confidence")
+    op.drop_column("token_events", "parser_version")
+    op.drop_column("token_events", "ingested_at")
+    op.drop_column("token_events", "source_created_at")
+    op.drop_column("token_events", "source_path")
+    op.drop_column("token_events", "source_id")
+    op.drop_column("token_events", "source_provider")
+    op.drop_column("token_events", "measurement_kind")

migrations/versions/w0r1s2t3u4v5_token_event_legacy_source_ids.py

@@ -0,0 +1,33 @@
+"""assign legacy source ids to measured token events
+
+Revision ID: w0r1s2t3u4v5
+Revises: v9q0r1s2t3u4
+Create Date: 2026-05-23
+"""
+from alembic import op
+
+revision = "w0r1s2t3u4v5"
+down_revision = "v9q0r1s2t3u4"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.execute(
+        """
+        UPDATE token_events
+        SET source_id = source_provider || ':legacy:' || id::text
+        WHERE measurement_kind = 'measured'
+          AND source_id IS NULL
+        """
+    )
+
+
+def downgrade() -> None:
+    op.execute(
+        """
+        UPDATE token_events
+        SET source_id = NULL
+        WHERE source_id = source_provider || ':legacy:' || id::text
+        """
+    )