fix(alerts): evaluate active session limits during ingest

Author: mikemikimike

tests/integration/test_full_pipeline.py

@@ -13,11 +13,10 @@
 from __future__ import annotations
 
 import json
-import threading
+from datetime import timedelta
 from typing import Sequence
 
 import pytest
-from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider, ReadableSpan
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult
 
@@ -39,6 +38,7 @@
 from tokenjam.sdk.agent import watch, AgentSession, record_llm_call, record_tool_call
 from tokenjam.utils.time_parse import utcnow
 import tokenjam.sdk.agent as agent_mod
+from tests.factories import make_llm_span
 
 
 
@@ -156,6 +156,7 @@ class _Stack:
 
     stack = _Stack()
     stack.db = db
+    stack.pipeline = pipeline
 
     yield stack
 
@@ -268,6 +269,67 @@ def my_agent():
     assert len(llm_spans) == 3
 
 
+def test_llm_only_session_enforces_session_budget_while_still_active(full_stack):
+    """A stream without invoke_agent still reaches session-scoped alerts."""
+    full_stack.pipeline.process(
+        make_llm_span(
+            agent_id="test-agent",
+            session_id="llm-only-budget",
+            input_tokens=30_000_000,
+            output_tokens=0,
+        )
+    )
+    full_stack.pipeline.process(
+        make_llm_span(
+            agent_id="test-agent",
+            session_id="llm-only-budget",
+            input_tokens=30_000_000,
+            output_tokens=0,
+        )
+    )
+    full_stack.pipeline.process(
+        make_llm_span(
+            agent_id="test-agent",
+            session_id="another-llm-only-budget",
+            input_tokens=30_000_000,
+            output_tokens=0,
+        )
+    )
+
+    rows = full_stack.db.conn.execute(
+        "SELECT type, suppressed FROM alerts ORDER BY fired_at"
+    ).fetchall()
+    assert [(row[0], row[1]) for row in rows] == [
+        ("cost_budget_session", False),
+        ("cost_budget_session", False),
+    ]
+
+
+def test_llm_only_session_enforces_duration_limit_while_still_active(full_stack):
+    """Duration is checked from the running session's observed span bounds."""
+    start = utcnow() - timedelta(seconds=4001)
+    full_stack.pipeline.process(
+        make_llm_span(
+            agent_id="test-agent",
+            session_id="llm-only-duration",
+            start_time=start,
+            cost_usd=0.01,
+        )
+    )
+    full_stack.pipeline.process(
+        make_llm_span(
+            agent_id="test-agent",
+            session_id="llm-only-duration",
+            cost_usd=0.01,
+        )
+    )
+
+    rows = full_stack.db.conn.execute(
+        "SELECT type FROM alerts WHERE session_id = ?", ["llm-only-duration"]
+    ).fetchall()
+    assert [row[0] for row in rows] == ["session_duration"]
+
+
 def test_tool_call_flows_to_db(full_stack):
     """record_tool_call() should produce a tool span in the DB."""
 

tests/synthetic/test_alert_rules.py

@@ -1,15 +1,13 @@
 """Synthetic tests for alert rules — uses span factories + mock StorageBackend."""
 from __future__ import annotations
 
-from datetime import date
+from datetime import timedelta
 from unittest.mock import MagicMock
 
 from tokenjam.core.alerts import (
-    AlertDispatcher,
     AlertEngine,
     FileChannel,
     StdoutChannel,
-    WebhookChannel,
     _alert_to_dict,
     _strip_sensitive,
     SENSITIVE_DETAIL_KEYS,
@@ -220,6 +218,36 @@ def test_cost_budget_session_does_not_fire_under_budget():
     db.insert_alert.assert_not_called()
 
 
+def test_cost_budget_session_fires_during_session_progress():
+    config = _make_config(agents={
+        "test-agent": AgentConfig(budget=BudgetConfig(session_usd=1.00)),
+    })
+    engine, db = _make_engine(config)
+    session = make_session(
+        agent_id="test-agent", total_cost_usd=1.50, status="active", ended_at=None
+    )
+    engine.evaluate_session_progress(session)
+    db.insert_alert.assert_called_once()
+    assert db.insert_alert.call_args[0][0].type == AlertType.COST_BUDGET_SESSION
+
+
+def test_session_duration_fires_during_session_progress():
+    config = _make_config(agents={
+        "test-agent": AgentConfig(budget=BudgetConfig(session_usd=5.00)),
+    })
+    engine, db = _make_engine(config)
+    now = utcnow()
+    session = make_session(
+        agent_id="test-agent",
+        started_at=now - timedelta(seconds=4000),
+        ended_at=now,
+        status="active",
+    )
+    engine.evaluate_session_progress(session)
+    db.insert_alert.assert_called_once()
+    assert db.insert_alert.call_args[0][0].type == AlertType.SESSION_DURATION
+
+
 def test_cost_budget_daily_fires_when_exceeded():
     config = _make_config(agents={
         "test-agent": AgentConfig(budget=BudgetConfig(daily_usd=10.00)),

tokenjam/core/alerts.py

@@ -51,6 +51,10 @@
 _FAILURE_RATE_THRESHOLD = 0.20
 _FAILURE_RATE_CHECK_INTERVAL = 5
 _SESSION_DURATION_DEFAULT = 3600  # seconds
+_SESSION_SCOPED_ALERT_TYPES = frozenset({
+    AlertType.COST_BUDGET_SESSION,
+    AlertType.SESSION_DURATION,
+})
 
 # Row cap for the startup queries that rebuild CooldownTracker/_failure_rate_fired
 # from the alerts table (#592). Generous rather than exact: missing a stale row
@@ -198,17 +202,23 @@ class CooldownTracker:
 
     def __init__(self, cooldown_seconds: int = 60) -> None:
         self.cooldown_seconds = cooldown_seconds
-        self._last_fired: dict[tuple[str, str], datetime] = {}
+        self._last_fired: dict[tuple[str, str, str], datetime] = {}
 
-    def is_suppressed(self, agent_id: str | None, alert_type: AlertType) -> bool:
-        key = (agent_id or "", alert_type.value)
+    def is_suppressed(
+        self, agent_id: str | None, alert_type: AlertType, session_id: str | None = None,
+    ) -> bool:
+        scope = session_id if alert_type in _SESSION_SCOPED_ALERT_TYPES else ""
+        key = (agent_id or "", alert_type.value, scope or "")
         last = self._last_fired.get(key)
         if last is None:
             return False
         return (utcnow() - last).total_seconds() < self.cooldown_seconds
 
-    def record(self, agent_id: str | None, alert_type: AlertType) -> None:
-        key = (agent_id or "", alert_type.value)
+    def record(
+        self, agent_id: str | None, alert_type: AlertType, session_id: str | None = None,
+    ) -> None:
+        scope = session_id if alert_type in _SESSION_SCOPED_ALERT_TYPES else ""
+        key = (agent_id or "", alert_type.value, scope or "")
         self._last_fired[key] = utcnow()
 
     def hydrate(self, alerts: list[Alert]) -> None:
@@ -221,7 +231,8 @@ def hydrate(self, alerts: list[Alert]) -> None:
         for alert in alerts:
             if alert.suppressed:
                 continue
-            key = (alert.agent_id or "", alert.type.value)
+            scope = alert.session_id if alert.type in _SESSION_SCOPED_ALERT_TYPES else ""
+            key = (alert.agent_id or "", alert.type.value, scope or "")
             existing = self._last_fired.get(key)
             if existing is None or alert.fired_at > existing:
                 self._last_fired[key] = alert.fired_at
@@ -245,6 +256,7 @@ def __init__(self, db: StorageBackend, config: TjConfig) -> None:
         # threshold re-fired (5/20, 6/20, 7/20 …), turning a few incidents into a
         # cascade of near-identical alerts.
         self._failure_rate_fired: set[str] = set()
+        self._session_limit_fired: set[tuple[str, str]] = set()
         self._hydrate_from_db()
 
     def _hydrate_from_db(self) -> None:
@@ -275,6 +287,9 @@ def _hydrate_from_db(self) -> None:
             since = utcnow() - timedelta(seconds=self.cooldown.cooldown_seconds)
             recent = self.db.get_alerts(AlertFilters(since=since, limit=_HYDRATION_LIMIT))
             self.cooldown.hydrate(recent)
+            for alert in recent:
+                if alert.type in _SESSION_SCOPED_ALERT_TYPES and alert.session_id:
+                    self._session_limit_fired.add((alert.session_id, alert.type.value))
 
             fired = self.db.get_alerts(
                 AlertFilters(type=AlertType.FAILURE_RATE, limit=_HYDRATION_LIMIT)
@@ -305,6 +320,18 @@ def evaluate_session_end(self, session: SessionRecord) -> None:
         self._check_cost_budgets(session)
         self._check_session_duration(session)
 
+    def evaluate_session_progress(self, session: SessionRecord) -> None:
+        """Evaluate limits for sessions that do not emit an end marker.
+
+        OTLP and transcript-based SDK integrations may emit only individual
+        ``gen_ai.llm.call`` spans. Those sessions remain active, so waiting
+        for ``evaluate_session_end`` would make session cost and duration
+        alerts unreachable. Daily budgets stay end-triggered because they
+        are not session-scoped.
+        """
+        self._check_session_budget(session)
+        self._check_session_duration(session)
+
     def fire(
         self,
         alert_type: AlertType,
@@ -517,7 +544,17 @@ def _check_cost_budgets(self, session: SessionRecord) -> None:
         """
         budget = resolve_effective_budget(session.agent_id, self.config)
 
-        # Session budget (per-agent, both kinds; backward-compat only).
+        self._check_session_budget(session)
+
+        kind = classify_agent_kind(session.agent_id)
+        if kind.is_coding and kind.group:
+            self._check_coding_group_daily_budget(session, kind.group)
+        else:
+            self._check_agent_daily_budget(session, budget)
+
+    def _check_session_budget(self, session: SessionRecord) -> None:
+        """Check the per-session cost threshold for an active or ended session."""
+        budget = resolve_effective_budget(session.agent_id, self.config)
         if budget.session_usd is not None and session.total_cost_usd is not None:
             if session.total_cost_usd > budget.session_usd:
                 alert = Alert(
@@ -536,12 +573,6 @@ def _check_cost_budgets(self, session: SessionRecord) -> None:
                 )
                 self._fire(alert)
 
-        kind = classify_agent_kind(session.agent_id)
-        if kind.is_coding and kind.group:
-            self._check_coding_group_daily_budget(session, kind.group)
-        else:
-            self._check_agent_daily_budget(session, budget)
-
     def _check_coding_group_daily_budget(self, session: SessionRecord, group_id: str) -> None:
         """Daily cap for a coding-tool GROUP: compares the SUM of today's
         spend across every agent_id that classifies into `group_id` (not
@@ -634,12 +665,18 @@ def _check_session_duration(self, session: SessionRecord) -> None:
 
     def _fire(self, alert: Alert) -> None:
         """Persist alert to DB and dispatch. Suppressed alerts are persisted but not dispatched."""
-        if self.cooldown.is_suppressed(alert.agent_id, alert.type):
+        session_key = (alert.session_id, alert.type.value)
+        if alert.type in _SESSION_SCOPED_ALERT_TYPES and alert.session_id:
+            if session_key in self._session_limit_fired:
+                return
+        if self.cooldown.is_suppressed(alert.agent_id, alert.type, alert.session_id):
             alert.suppressed = True
             self.db.insert_alert(alert)
             return
         self.db.insert_alert(alert)
-        self.cooldown.record(alert.agent_id, alert.type)
+        if alert.type in _SESSION_SCOPED_ALERT_TYPES and alert.session_id:
+            self._session_limit_fired.add(session_key)
+        self.cooldown.record(alert.agent_id, alert.type, alert.session_id)
         self.dispatcher.dispatch(alert)
 
 

tokenjam/core/ingest.py

@@ -333,6 +333,20 @@ def process(self, span: NormalizedSpan) -> None:
         # 6. Post-ingest hooks (never let hook errors kill the pipeline)
         self._run_hooks(span)
 
+        # CostEngine runs in the post-ingest hooks and may replace the span's
+        # incoming cost, so evaluate active-session limits against the
+        # persisted record after those hooks have completed.
+        if not self._is_session_end(span) and self.alert_engine:
+            try:
+                evaluate_progress = getattr(
+                    self.alert_engine, "evaluate_session_progress", None
+                )
+                if evaluate_progress is not None:
+                    current_session = self.db.get_session(span.session_id) or session
+                    evaluate_progress(current_session)
+            except Exception as exc:
+                logger.warning("AlertEngine progress hook failed: %s", exc)
+
     def _is_duplicate_observation(self, span: NormalizedSpan) -> bool:
         """True when another ingest source already recorded this exact call.