From 765d2f8b2a2e739dbf72c373641eb7199ff0a062 Mon Sep 17 00:00:00 2001 From: Pragadeesh122 Date: Fri, 19 Jun 2026 20:16:53 -0500 Subject: [PATCH 1/2] feat(observability): per-session spend guardrail + RAG/chat latency panels JOB-4: make LLM spend and latency visible and bounded on the existing Prometheus/Grafana/OTel stack. - Add per-session token ceiling guardrail (api/session_budget.py): Redis-backed cumulative prompt+completion token counter per session, enforced at the top of both general and project chat turns. Refuses turns once MAX_SESSION_TOKENS (default 2,000,000) is exceeded; fails open on Redis errors. Trips counted in agenticrag_session_budget_blocked_total{chat_type,limit}. - Add agenticrag_retrieval_duration_seconds histogram (cache_status label) for RAG retrieval p50/p95; emitted from pipeline/retriever.py. - Grafana: add "RAG Retrieval Latency (p50/p95)" and "Chat/Stream Completion Latency (p50/p95) by Mode" panels to UX & Latency dashboard, and "Session Spend Guardrail Blocks" panel to Economics dashboard. - Docs: document both guardrail layers (per-turn context cap + per-session ceiling), defaults, and a spend/latency baseline section with read queries. - Tests: tests/test_session_budget.py covers budget check/record/enforce, fail-open, and the new metric emitters. Co-Authored-By: Paperclip --- api/chat.py | 20 ++++ api/project_chat.py | 20 ++++ api/session_budget.py | 106 ++++++++++++++++++ docs/architecture/observability.md | 56 ++++++++- .../dashboards/agenticrag-economics.json | 19 +++- .../dashboards/agenticrag-ux-latency.json | 46 +++++++- observability/metrics.py | 27 +++++ pipeline/retriever.py | 9 ++ tests/test_session_budget.py | 102 +++++++++++++++++ 9 files changed, 401 insertions(+), 4 deletions(-) create mode 100644 api/session_budget.py create mode 100644 tests/test_session_budget.py diff --git a/api/chat.py b/api/chat.py index 9e516f5..70af299 100644 --- a/api/chat.py +++ b/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) diff --git a/api/project_chat.py b/api/project_chat.py index 2699329..cb08f1b 100644 --- a/api/project_chat.py +++ b/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) diff --git a/api/session_budget.py b/api/session_budget.py new file mode 100644 index 0000000..8786f70 --- /dev/null +++ b/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") diff --git a/docs/architecture/observability.md b/docs/architecture/observability.md index bff22c2..8cdc052 100644 --- a/docs/architecture/observability.md +++ b/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). diff --git a/monitoring/grafana/provisioning/dashboards/agenticrag-economics.json b/monitoring/grafana/provisioning/dashboards/agenticrag-economics.json index a9f568c..3a881c8 100644 --- a/monitoring/grafana/provisioning/dashboards/agenticrag-economics.json +++ b/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 } diff --git a/monitoring/grafana/provisioning/dashboards/agenticrag-ux-latency.json b/monitoring/grafana/provisioning/dashboards/agenticrag-ux-latency.json index 6e61400..a468597 100644 --- a/monitoring/grafana/provisioning/dashboards/agenticrag-ux-latency.json +++ b/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 } diff --git a/observability/metrics.py b/observability/metrics.py index 0672b67..79179f8 100644 --- a/observability/metrics.py +++ b/observability/metrics.py @@ -120,6 +120,18 @@ def _env_bool(name: str, default: str = "true") -> bool: ["agent_name"], buckets=(0, 1, 2, 3, 5, 8, 10, 15, 20, 30, 50, 100), ) +RETRIEVAL_DURATION_SECONDS = Histogram( + "agenticrag_retrieval_duration_seconds", + "End-to-end RAG retrieval latency in seconds (cache lookup + vector search " + "+ rerank), labelled by whether the semantic cache served the result.", + ["cache_status"], + buckets=(0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 3, 5, 10), +) +SESSION_BUDGET_BLOCKED_TOTAL = Counter( + "agenticrag_session_budget_blocked_total", + "Chat turns refused because a per-session spend guardrail ceiling was reached.", + ["chat_type", "limit"], +) ORCHESTRATION_STEPS_TOTAL = Counter( "agenticrag_orchestration_steps_total", "General-chat orchestration step planning decisions.", @@ -337,6 +349,21 @@ def observe_retrieval_results(*, agent_name: str, result_count: int) -> None: ) +def observe_retrieval_latency(*, cache_hit: bool, duration_seconds: float) -> None: + if duration_seconds < 0: + return + RETRIEVAL_DURATION_SECONDS.labels( + cache_status="hit" if cache_hit else "miss" + ).observe(duration_seconds) + + +def observe_session_budget_blocked(*, chat_type: str, limit: str) -> None: + SESSION_BUDGET_BLOCKED_TOTAL.labels( + chat_type=chat_type or "unknown", + limit=limit or "unknown", + ).inc() + + def observe_orchestration_step( *, mode: str, diff --git a/pipeline/retriever.py b/pipeline/retriever.py index 44778aa..90caff0 100644 --- a/pipeline/retriever.py +++ b/pipeline/retriever.py @@ -4,11 +4,13 @@ import os import random import threading +import time from pipeline.embedder import embed_query_dense, embed_query_sparse from pipeline.pinecone_helpers import query_vectors from pipeline.query_rewrite import generate_hyde_passage from pipeline.retrieval_cache import get_cached_retrieval, cache_retrieval +from observability.metrics import observe_retrieval_latency from observability.spans import retrieval_span logger = logging.getLogger("pipeline.retriever") @@ -63,6 +65,7 @@ def retrieve( - results: list of {"id", "score", "text", "source", "page", "document_id"} - info: {"cache_hit": bool} — whether results came from the semantic cache """ + started = time.perf_counter() with retrieval_span( span_name="retrieval.pipeline", **{"retrieval.top_k": top_k, "retrieval.alpha": alpha}, @@ -75,6 +78,9 @@ def retrieve( if span is not None: span.set_attribute("cache.hit", True) span.set_attribute("result_count", len(cached)) + observe_retrieval_latency( + cache_hit=True, duration_seconds=time.perf_counter() - started + ) if CACHE_AUDIT_RATE > 0 and random.random() < CACHE_AUDIT_RATE: _schedule_cache_audit( project_id, query, cached, chunk_count, top_k, alpha, use_hyde @@ -95,6 +101,9 @@ def retrieve( if span is not None: span.set_attribute("result_count", len(results)) + observe_retrieval_latency( + cache_hit=False, duration_seconds=time.perf_counter() - started + ) logger.info(f"retrieved {len(results)} results") return results, {"cache_hit": False} diff --git a/tests/test_session_budget.py b/tests/test_session_budget.py new file mode 100644 index 0000000..c695530 --- /dev/null +++ b/tests/test_session_budget.py @@ -0,0 +1,102 @@ +"""Tests for the per-session spend guardrail (api/session_budget.py) and the +retrieval-latency / guardrail metrics added for LLM cost & latency observability.""" + +import unittest +from unittest.mock import MagicMock, patch + +from prometheus_client import REGISTRY + +import api.session_budget as sb +from observability.metrics import ( + SESSION_BUDGET_BLOCKED_TOTAL, + observe_retrieval_latency, + observe_session_budget_blocked, +) + + +def _count(metric: str, **labels) -> float: + return REGISTRY.get_sample_value(metric, labels) or 0.0 + + +class SessionBudgetTests(unittest.TestCase): + def setUp(self): + self.redis = MagicMock() + self._patcher = patch.object(sb, "redis_client", self.redis) + self._patcher.start() + self.addCleanup(self._patcher.stop) + + @patch.object(sb, "MAX_SESSION_TOKENS", 1000) + def test_under_ceiling_is_allowed(self): + self.redis.get.return_value = "500" + status = sb.check_session_budget("sess-1") + self.assertTrue(status.allowed) + self.assertEqual(status.used_tokens, 500) + self.assertEqual(status.ceiling, 1000) + self.assertEqual(status.remaining, 500) + + @patch.object(sb, "MAX_SESSION_TOKENS", 1000) + def test_at_or_over_ceiling_is_blocked(self): + self.redis.get.return_value = "1000" + status = sb.check_session_budget("sess-1") + self.assertFalse(status.allowed) + self.assertEqual(status.remaining, 0) + + @patch.object(sb, "MAX_SESSION_TOKENS", 0) + def test_zero_ceiling_disables_enforcement(self): + status = sb.check_session_budget("sess-1") + self.assertTrue(status.allowed) + self.redis.get.assert_not_called() + + @patch.object(sb, "MAX_SESSION_TOKENS", 1000) + def test_record_increments_and_refreshes_ttl(self): + self.redis.incrby.return_value = 700 + sb.record_session_tokens("sess-1", 400, 300) + self.redis.incrby.assert_called_once_with("session:sess-1:tokens", 700) + self.redis.expire.assert_called_once() + + @patch.object(sb, "MAX_SESSION_TOKENS", 1000) + def test_record_ignores_zero_usage(self): + sb.record_session_tokens("sess-1", 0, 0) + self.redis.incrby.assert_not_called() + + @patch.object(sb, "MAX_SESSION_TOKENS", 1000) + def test_check_fails_open_on_redis_error(self): + self.redis.get.side_effect = RuntimeError("redis down") + status = sb.check_session_budget("sess-1") + self.assertTrue(status.allowed) + + @patch.object(sb, "MAX_SESSION_TOKENS", 1000) + def test_empty_session_id_is_allowed(self): + status = sb.check_session_budget("") + self.assertTrue(status.allowed) + self.redis.get.assert_not_called() + + +class MetricsTests(unittest.TestCase): + def test_retrieval_latency_records_by_cache_status(self): + metric = "agenticrag_retrieval_duration_seconds_count" + before = _count(metric, cache_status="hit") + observe_retrieval_latency(cache_hit=True, duration_seconds=0.012) + after = _count(metric, cache_status="hit") + self.assertEqual(after, before + 1) + + def test_negative_latency_is_ignored(self): + metric = "agenticrag_retrieval_duration_seconds_count" + before = _count(metric, cache_status="miss") + observe_retrieval_latency(cache_hit=False, duration_seconds=-1.0) + after = _count(metric, cache_status="miss") + self.assertEqual(after, before) + + def test_session_budget_blocked_counter_increments(self): + before = SESSION_BUDGET_BLOCKED_TOTAL.labels( + chat_type="general", limit="session_tokens" + )._value.get() + observe_session_budget_blocked(chat_type="general", limit="session_tokens") + after = SESSION_BUDGET_BLOCKED_TOTAL.labels( + chat_type="general", limit="session_tokens" + )._value.get() + self.assertEqual(after, before + 1) + + +if __name__ == "__main__": + unittest.main() From 95eb9d231b4b94fd49b077537377a314779c5e91 Mon Sep 17 00:00:00 2001 From: Pragadeesh122 Date: Thu, 16 Jul 2026 21:12:05 -0500 Subject: [PATCH 2/2] fix(ci): always run test-api and helm-lint so required status checks never skip --- .github/workflows/pr.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 244af24..2dc6ed5 100644 --- a/.github/workflows/pr.yml +++ b/.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: