JOB-4: LLM cost & latency observability + spend guardrails

.github/workflows/pr.yml

@@ -62,7 +62,6 @@ jobs:
   test-api:
     name: API tests
     needs: changes
-    if: needs.changes.outputs.api == 'true'
     runs-on: ubuntu-24.04-arm
     timeout-minutes: 10
     steps:
@@ -132,7 +131,6 @@ jobs:
   helm-lint:
     name: Helm chart lint
     needs: changes
-    if: needs.changes.outputs.helm == 'true'
     runs-on: ubuntu-24.04-arm
     timeout-minutes: 5
     steps:

api/chat.py

@@ -11,6 +11,12 @@
 from utils.tool_planner import plan_tool_calls
 from utils.summarizer import summarize_messages
 from api.session import get_messages, save_messages, get_session_user
+from api.session_budget import (
+    budget_exceeded_message,
+    check_session_budget,
+    note_session_budget_blocked,
+    record_session_tokens,
+)
 from llm.response_utils import usage_tokens
 from observability.context import pop_context, push_context
 from observability.metrics import (
@@ -186,6 +192,15 @@ def chat_stream(session_id: str, user_message: str, attachments: list[dict] | No
             duration_seconds=0.0,
         )
 
+        # Per-session spend guardrail: refuse turns once the session has spent
+        # its cumulative token budget (see api/session_budget.py).
+        budget = check_session_budget(session_id)
+        if not budget.allowed:
+            note_session_budget_blocked(chat_type="general")
+            yield _sse("error", budget_exceeded_message(budget))
+            yield _sse("done", json.dumps({"tools_used": [], "budget_exceeded": True}))
+            return
+
         messages = get_messages(session_id)
         attachments = attachments or []
         # Persist attachment refs (not resolved bytes) on the user message in
@@ -424,6 +439,11 @@ def messages_for_llm() -> list[dict]:
 
         save_messages(session_id, messages)
 
+        # Charge this turn against the session spend guardrail.
+        if usage:
+            turn_prompt, turn_completion = usage_tokens(usage)
+            record_session_tokens(session_id, turn_prompt, turn_completion)
+
         if user_id:
             schedule_memory_persistence(messages, user_id, session_id=session_id)
 

api/project_chat.py

@@ -12,6 +12,12 @@
 from agents.base import Agent
 from agents.router import route as route_agent
 from api.session import get_messages, save_messages, get_session_user
+from api.session_budget import (
+    budget_exceeded_message,
+    check_session_budget,
+    note_session_budget_blocked,
+    record_session_tokens,
+)
 from functions import tool_schemas
 from functions.tool_router import execute_tool_call
 from llm.response_utils import usage_tokens
@@ -99,6 +105,15 @@ def project_chat_stream(
     _span_ctx = chat_turn_span(span_name="project_chat.turn", chat_type="project")
     _span_ctx.__enter__()
     try:
+        # Per-session spend guardrail: refuse turns once the session has spent
+        # its cumulative token budget (see api/session_budget.py).
+        budget = check_session_budget(session_id)
+        if not budget.allowed:
+            note_session_budget_blocked(chat_type="project")
+            yield _sse("error", budget_exceeded_message(budget))
+            yield _sse("done", json.dumps({"budget_exceeded": True}))
+            return
+
         messages = get_messages(session_id)
 
         # 1. Route to agent (pass full conversation for context-aware classification)
@@ -301,6 +316,11 @@ def project_chat_stream(
 
         save_messages(session_id, messages)
 
+        # Charge this turn against the session spend guardrail.
+        if usage:
+            turn_prompt, turn_completion = usage_tokens(usage)
+            record_session_tokens(session_id, turn_prompt, turn_completion)
+
         if user_id:
             schedule_memory_persistence(messages, user_id, session_id=session_id)
 

api/session_budget.py

@@ -0,0 +1,106 @@
+"""Per-session spend guardrail.
+
+A cheap, Redis-backed ceiling on cumulative token usage per chat session. This
+bounds the blast radius of a runaway loop, an abusive client, or a pathological
+conversation: once a session has spent more than the configured token budget it
+is refused further turns until it ages out (the counter shares the session TTL).
+
+This complements the per-turn context cap (``MAX_PROMPT_TOKENS`` in the chat
+handlers, which triggers summarization). That cap bounds the size of a single
+request; this ceiling bounds the cumulative cost of a whole session.
+
+Defaults are intentionally generous so legitimate long research sessions are not
+interrupted, while still capping spend at a few dollars per session. Tune via the
+``MAX_SESSION_TOKENS`` env var (0 disables enforcement).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+
+from memory.redis_client import redis_client
+from observability.metrics import observe_session_budget_blocked
+
+logger = logging.getLogger("session_budget")
+
+# Cumulative prompt+completion tokens allowed per session before turns are
+# refused. 2,000,000 tokens is ~/session at GPT-4o-class blended pricing and
+# far above any honest single conversation, so it only trips on abuse/runaways.
+MAX_SESSION_TOKENS = int(os.getenv("MAX_SESSION_TOKENS", "2000000"))
+
+# Counter key shares the 24h session TTL so it resets when the session ages out.
+_SESSION_TTL = 60 * 60 * 24
+
+
+def _tokens_key(session_id: str) -> str:
+    return f"session:{session_id}:tokens"
+
+
+@dataclass(frozen=True)
+class SessionBudgetStatus:
+    allowed: bool
+    used_tokens: int
+    ceiling: int
+
+    @property
+    def remaining(self) -> int:
+        return max(0, self.ceiling - self.used_tokens)
+
+
+def _enabled() -> bool:
+    return MAX_SESSION_TOKENS > 0
+
+
+def check_session_budget(session_id: str) -> SessionBudgetStatus:
+    """Return whether ``session_id`` is still under its token ceiling.
+
+    Never raises: a Redis failure fails open (allowed) so an observability
+    guardrail can never take down the chat path.
+    """
+    if not _enabled() or not session_id:
+        return SessionBudgetStatus(True, 0, MAX_SESSION_TOKENS)
+    try:
+        raw = redis_client.get(_tokens_key(session_id))
+        used = int(raw) if raw else 0
+    except Exception as e:  # pragma: no cover - defensive
+        logger.warning(f"session budget check failed, allowing turn: {e}")
+        return SessionBudgetStatus(True, 0, MAX_SESSION_TOKENS)
+    return SessionBudgetStatus(used < MAX_SESSION_TOKENS, used, MAX_SESSION_TOKENS)
+
+
+def record_session_tokens(
+    session_id: str, prompt_tokens: int, completion_tokens: int
+) -> None:
+    """Add a turn's token usage to the session's cumulative counter."""
+    if not _enabled() or not session_id:
+        return
+    total = max(0, int(prompt_tokens or 0)) + max(0, int(completion_tokens or 0))
+    if total <= 0:
+        return
+    try:
+        key = _tokens_key(session_id)
+        new_total = redis_client.incrby(key, total)
+        # Keep the counter alive for the session window; refresh on each turn.
+        redis_client.expire(key, _SESSION_TTL)
+        if new_total >= MAX_SESSION_TOKENS:
+            logger.info(
+                f"session {session_id} reached token ceiling: "
+                f"{new_total}/{MAX_SESSION_TOKENS}"
+            )
+    except Exception as e:  # pragma: no cover - defensive
+        logger.warning(f"failed to record session tokens: {e}")
+
+
+def budget_exceeded_message(status: SessionBudgetStatus) -> str:
+    return (
+        "This conversation has reached its usage limit "
+        f"({status.used_tokens:,}/{status.ceiling:,} tokens). "
+        "Start a new chat to continue."
+    )
+
+
+def note_session_budget_blocked(*, chat_type: str) -> None:
+    """Record the guardrail trip for observability."""
+    observe_session_budget_blocked(chat_type=chat_type, limit="session_tokens")

docs/architecture/observability.md

@@ -65,6 +65,8 @@ Gated behind `OBS_ENABLE_HIGH_CARDINALITY_METRICS` (default: on). These use hash
 | `agenticrag_tool_budget_exhausted_total` | Counter | chat_type, budget | Budget exhaustion events (reasoning_steps or total_tool_calls) |
 | `agenticrag_summarization_events_total` | Counter | chat_type, reason | Conversation summarization triggers |
 | `agenticrag_retrieval_results_count` | Histogram | agent_name | Retrieved chunk count per query |
+| `agenticrag_retrieval_duration_seconds` | Histogram | cache_status | End-to-end RAG retrieval latency (cache lookup + vector search + rerank), `hit`/`miss` |
+| `agenticrag_session_budget_blocked_total` | Counter | chat_type, limit | Chat turns refused by the per-session spend guardrail |
 
 ### HTTP Metrics
 
@@ -136,9 +138,9 @@ Three dashboards are auto-provisioned from `monitoring/grafana/provisioning/dash
 
 | Dashboard | Focus |
 |-----------|-------|
-| **RunaxAI - Economics** | LLM spend by provider/model, token usage trends, cost per chat type |
+| **RunaxAI - Economics** | LLM spend by provider/model, token usage trends, cost per chat type, session spend-guardrail blocks |
 | **RunaxAI - Operations** | Agent routing distribution, tool call counts, orchestration step analysis, retrieval chunk counts, duplicate suppression rates |
-| **RunaxAI - UX & Latency** | TTFT distribution, request duration, streaming output speed, HTTP request rates |
+| **RunaxAI - UX & Latency** | TTFT distribution, request duration, streaming output speed, chat/stream p50/p95 by mode, RAG retrieval p50/p95, HTTP request rates |
 
 ### Alert Rules
 
@@ -154,3 +156,53 @@ LLM spend is estimated using LiteLLM's `cost_per_token()` function, which mainta
 4. Broken down by provider, model, and chat type
 
 When token usage isn't reported by the provider (common with some streaming implementations), the client estimates tokens using `litellm.token_counter()` before computing cost.
+
+## Spend Guardrails
+
+Two layers bound LLM spend, each with sane, env-tunable defaults.
+
+### Per-turn context cap (summarization)
+
+Each chat turn measures its prompt token count and, once it exceeds a threshold, collapses older history into a summary before the next turn. This caps the size — and therefore the per-call cost — of any single request.
+
+| Path | Constant | Default | Behavior on breach |
+|------|----------|---------|--------------------|
+| General chat | `MAX_PROMPT_TOKENS` (`api/chat.py`) | 40,000 | Summarize conversation |
+| Project chat | `MAX_PROMPT_TOKENS` (`api/project_chat.py`) | 60,000 | Summarize conversation |
+| Worker prompts | `MAX_PROMPT_TOKENS` (`main.py`) | 10,000 | Summarize conversation |
+| Per uploaded file | `MAX_TOKENS_PER_DOCUMENT` (`pipeline/chat_attachments.py`) | 25,000 | Reject the file |
+| Per session attachments | `MAX_SESSION_ATTACHMENT_TOKENS` (`pipeline/chat_attachments.py`) | 25,000 | Reject the upload |
+
+### Per-session spend ceiling
+
+`api/session_budget.py` enforces a cumulative token ceiling per chat session, backed by a Redis counter that shares the 24h session TTL. Every turn's prompt+completion tokens are added to the counter; once a session crosses the ceiling, further turns are refused with a user-facing error (`event: error`) and a `budget_exceeded` `done` event. This bounds the blast radius of runaway tool loops, abusive clients, or pathological conversations.
+
+| Env var | Default | Meaning |
+|---------|---------|---------|
+| `MAX_SESSION_TOKENS` | 2,000,000 | Cumulative prompt+completion tokens allowed per session. `0` disables enforcement. |
+
+Design notes:
+
+- **Fails open.** Any Redis error during the budget check allows the turn — an observability guardrail must never take down the chat path.
+- **Generous by default.** 2M tokens is far above any honest single conversation (~/session at GPT-4o-class blended pricing), so it only trips on abuse/runaways. Lower it per-environment for tighter cost control.
+- **Observable.** Trips are counted in `agenticrag_session_budget_blocked_total{chat_type}` and surfaced on the **Economics** dashboard ("Session Spend Guardrail Blocks").
+
+## Baseline
+
+Live spend/latency baselines are read from Grafana once production traffic flows; the dashboards above are the source of truth. Useful baseline queries (Explore -> Prometheus):
+
+```promql
+# Blended cost per 1K tokens (range)
+(1000 * sum(increase(agenticrag_llm_spend_usd_total[$__range])))
+  / clamp_min(sum(increase(agenticrag_llm_tokens_total[$__range])), 1)
+
+# Chat/stream completion latency p50 / p95 by mode
+histogram_quantile(0.50, sum by (le, chat_type) (rate(agenticrag_llm_request_duration_seconds_bucket{operation="completion",stream="true"}[5m])))
+histogram_quantile(0.95, sum by (le, chat_type) (rate(agenticrag_llm_request_duration_seconds_bucket{operation="completion",stream="true"}[5m])))
+
+# RAG retrieval latency p50 / p95 by cache status
+histogram_quantile(0.50, sum by (le, cache_status) (rate(agenticrag_retrieval_duration_seconds_bucket[5m])))
+histogram_quantile(0.95, sum by (le, cache_status) (rate(agenticrag_retrieval_duration_seconds_bucket[5m])))
+```
+
+**Cost-model baseline (bounded worst case).** Because per-turn prompt size is capped (see above), the maximum spend per turn is bounded. At GPT-4o-class blended pricing (~/1M input, ~/1M output), a worst-case general-chat turn (40K prompt + ~2K output) costs ~, and a project-chat turn (60K prompt + ~2K output) ~. With the 2M-token session ceiling, a single session is hard-capped at roughly  of spend. These are upper bounds; typical turns are far smaller. Replace with measured values from the queries above once real traffic is captured (target: record p50/p95 latency and /1K-token blended cost after the first week of production traffic).

monitoring/grafana/provisioning/dashboards/agenticrag-economics.json

@@ -278,6 +278,23 @@
       ],
       "title": "Usage Missing / Estimated Rate",
       "type": "timeseries"
+    },
+    {
+      "datasource": {"type": "prometheus", "uid": "prometheus"},
+      "description": "Chat turns refused by the per-session spend guardrail (MAX_SESSION_TOKENS ceiling), by chat mode. A non-zero count means sessions are hitting their token budget.",
+      "fieldConfig": {"defaults": {"unit": "short", "decimals": 0}, "overrides": []},
+      "gridPos": {"h": 8, "w": 12, "x": 0, "y": 29},
+      "id": 11,
+      "options": {"legend": {"displayMode": "table", "placement": "bottom", "calcs": ["sum"]}},
+      "targets": [
+        {
+          "expr": "sum by (chat_type) (increase(agenticrag_session_budget_blocked_total[$__range]))",
+          "legendFormat": "{{chat_type}}",
+          "refId": "A"
+        }
+      ],
+      "title": "Session Spend Guardrail Blocks",
+      "type": "timeseries"
     }
   ],
   "refresh": "30s",
@@ -289,5 +306,5 @@
   "timezone": "browser",
   "title": "RunaxAI - Economics",
   "uid": "agenticrag-econ-v1",
-  "version": 2
+  "version": 3
 }

monitoring/grafana/provisioning/dashboards/agenticrag-ux-latency.json

@@ -191,6 +191,50 @@
       ],
       "title": "HTTP 5xx Rate by Path",
       "type": "timeseries"
+    },
+    {
+      "datasource": {"type": "prometheus", "uid": "prometheus"},
+      "description": "End-to-end RAG retrieval latency (cache lookup + vector search + rerank), split by cache hit/miss.",
+      "fieldConfig": {"defaults": {"unit": "s", "decimals": 3}, "overrides": []},
+      "gridPos": {"h": 8, "w": 12, "x": 0, "y": 32},
+      "id": 9,
+      "options": {"legend": {"displayMode": "table", "placement": "bottom", "calcs": ["mean", "max"]}},
+      "targets": [
+        {
+          "expr": "histogram_quantile(0.50, sum by (le, cache_status) (rate(agenticrag_retrieval_duration_seconds_bucket[5m])))",
+          "legendFormat": "p50 {{cache_status}}",
+          "refId": "A"
+        },
+        {
+          "expr": "histogram_quantile(0.95, sum by (le, cache_status) (rate(agenticrag_retrieval_duration_seconds_bucket[5m])))",
+          "legendFormat": "p95 {{cache_status}}",
+          "refId": "B"
+        }
+      ],
+      "title": "RAG Retrieval Latency (p50/p95)",
+      "type": "timeseries"
+    },
+    {
+      "datasource": {"type": "prometheus", "uid": "prometheus"},
+      "description": "p50 and p95 end-to-end LLM completion latency for streaming chat turns, broken down by chat mode (general vs project).",
+      "fieldConfig": {"defaults": {"unit": "s", "decimals": 2}, "overrides": []},
+      "gridPos": {"h": 8, "w": 12, "x": 12, "y": 32},
+      "id": 10,
+      "options": {"legend": {"displayMode": "table", "placement": "bottom", "calcs": ["mean", "max"]}},
+      "targets": [
+        {
+          "expr": "histogram_quantile(0.50, sum by (le, chat_type) (rate(agenticrag_llm_request_duration_seconds_bucket{operation=\"completion\", stream=\"true\"}[5m])))",
+          "legendFormat": "p50 {{chat_type}}",
+          "refId": "A"
+        },
+        {
+          "expr": "histogram_quantile(0.95, sum by (le, chat_type) (rate(agenticrag_llm_request_duration_seconds_bucket{operation=\"completion\", stream=\"true\"}[5m])))",
+          "legendFormat": "p95 {{chat_type}}",
+          "refId": "B"
+        }
+      ],
+      "title": "Chat/Stream Completion Latency (p50/p95) by Mode",
+      "type": "timeseries"
     }
   ],
   "refresh": "30s",
@@ -202,5 +246,5 @@
   "timezone": "browser",
   "title": "RunaxAI - UX & Latency",
   "uid": "agenticrag-ux-v1",
-  "version": 2
+  "version": 3
 }