diff --git a/datasource-catalog.json b/datasource-catalog.json index 40081905e0..27c398f351 100644 --- a/datasource-catalog.json +++ b/datasource-catalog.json @@ -222,7 +222,7 @@ }, { "id": "grafana/loki", - "name": "Grafana Loki", + "name": "Loki", "description": "Search and query Grafana Loki logs", "categories": ["logs"], "iconUrl": "https://raw.githubusercontent.com/gilbarbara/logos/de2c1f96ff6e74ea7ea979b43202e8d4b863c655/logos/grafana.svg", @@ -294,7 +294,7 @@ }, { "id": "grafana/tempo", - "name": "Grafana Tempo", + "name": "Tempo", "description": "Search and query Grafana Tempo traces", "categories": ["traces"], "iconUrl": "https://raw.githubusercontent.com/gilbarbara/logos/de2c1f96ff6e74ea7ea979b43202e8d4b863c655/logos/grafana.svg", diff --git a/experimental/ag-ui/server-agui.py b/experimental/ag-ui/server-agui.py index 269b04415f..dd5c6a88a4 100644 --- a/experimental/ag-ui/server-agui.py +++ b/experimental/ag-ui/server-agui.py @@ -24,6 +24,11 @@ from starlette.responses import PlainTextResponse from holmes.utils.stream import StreamMessage, StreamEvents +from holmes.core.usage_recorder import ( + UsageRecorderState, + resolve_provider, + stream_with_usage_recording, +) from holmes.common.env_vars import ( HOLMES_HOST, HOLMES_PORT, @@ -139,9 +144,47 @@ async def event_generator(message_history): run_id=input_data.run_id, ) ) - hgpt_chat_stream_response: StreamMessage = ai.call_stream( - msgs=message_history, - enable_tool_approval=chat_request.enable_tool_approval or False, + # Build the AI usage recorder state. AG-UI clients pass FE flags + # via input_data.context if they want to populate them; otherwise + # request_source / source_ref / user_id stay NULL. + ctx = input_data.context or {} + agui_user_id = None + try: + agui_user_id = getattr(input_data, "user_id", None) or ( + ctx.get("user_id") if isinstance(ctx, dict) else None + ) + except (AttributeError, TypeError) as e: + # Telemetry-only fallback: if input_data shape changed and + # neither attribute nor mapping access works, log at debug + # so we notice during dev rather than silently dropping + # user attribution. Don't break the response path. + logging.debug( + "agui_chat: failed to extract user_id (input_data=%r ctx=%r): %s", + type(input_data).__name__, + type(ctx).__name__, + e, + ) + ai_model = getattr(ai.llm, "model", None) or chat_request.model or "unknown" + ai_provider = resolve_provider(ai_model) + recorder_state = UsageRecorderState( + dal=dal, + request_type="agui_chat", + request_source=ctx.get("request_source") if isinstance(ctx, dict) else None, + source_ref=ctx.get("source_ref") if isinstance(ctx, dict) else None, + conversation_id=getattr(input_data, "thread_id", None), + conversation_source=None, # AG-UI doesn't write Conversations or ChatHistory + user_id=agui_user_id, + is_streaming=True, + model=ai_model, + provider=ai_provider, + is_robusta_model=getattr(ai.llm, "is_robusta_model", False), + ) + hgpt_chat_stream_response: StreamMessage = stream_with_usage_recording( + ai.call_stream( + msgs=message_history, + enable_tool_approval=chat_request.enable_tool_approval or False, + ), + recorder_state, ) for chunk in hgpt_chat_stream_response: if hasattr(chunk, "event"): diff --git a/holmes/checks/checks.py b/holmes/checks/checks.py index 8bbcc7bf6c..99263ca6f2 100644 --- a/holmes/checks/checks.py +++ b/holmes/checks/checks.py @@ -24,6 +24,11 @@ from holmes.config import Config from holmes.core.issue import Issue, IssueStatus from holmes.core.tool_calling_llm import LLMResult, ToolCallingLLM +from holmes.core.usage_recorder import ( + UsageRecorderState, + record_error, + record_from_llm_result, +) from holmes.plugins.destinations.pagerduty.plugin import PagerDutyDestination from holmes.plugins.destinations.slack.plugin import SlackDestination @@ -89,6 +94,7 @@ def execute_check( ai: ToolCallingLLM, verbose: bool = False, console: Optional[Console] = None, + recorder_state: Optional[UsageRecorderState] = None, ) -> CheckResult: """ Execute a single health check. @@ -101,6 +107,10 @@ def execute_check( ai: The LLM instance to use for evaluation verbose: Whether to print verbose output console: Optional console for output (only used if verbose=True) + recorder_state: Optional UsageRecorderState. When supplied (e.g. by the + /api/checks/execute endpoint) a usage event is recorded for this + LLM call. The CLI runner doesn't pass one and is therefore not + tracked, by design. Returns: CheckResult with status, message, and metadata @@ -114,6 +124,10 @@ def execute_check( start_time = time.time() try: response = _execute_ai_check(check, ai) + if recorder_state is not None: + # Fire the usage recorder. response IS-A LLMResult (RequestStats + # subclass) so cost/token fields come straight off it. + record_from_llm_result(recorder_state, response) check_response = _parse_check_response(response) if verbose and console: status_str = "PASS" if check_response.passed else "FAIL" @@ -142,6 +156,8 @@ def execute_check( ) except Exception as e: + if recorder_state is not None: + record_error(recorder_state, e) duration = time.time() - start_time result = CheckResult( check_name=check.name, diff --git a/holmes/checks/checks_api.py b/holmes/checks/checks_api.py index fead5d69a1..fa978ce5ba 100644 --- a/holmes/checks/checks_api.py +++ b/holmes/checks/checks_api.py @@ -13,6 +13,7 @@ from holmes.core.issue import Issue, IssueStatus from holmes.core.tool_calling_llm import LLMResult, ToolCallingLLM from holmes.core.tools import PrerequisiteCacheMode, ToolsetTag +from holmes.core.usage_recorder import UsageRecorderState, resolve_provider from holmes.plugins.destinations.slack.plugin import SlackDestination checks_app = FastAPI() @@ -110,12 +111,32 @@ def execute_health_check( destinations=destination_names, ) + # Build the recorder state so the operator-driven check shows up in + # HolmesUsageEvents with request_type='health_check' and source_ref + # set to the check name (per-check cost reporting key). + ai_model = getattr(ai.llm, "model", None) or request.model or "unknown" + recorder_state = UsageRecorderState( + dal=_CONFIG.dal, + request_type="health_check", + request_source="operator", + source_ref=request.name or "api-check", + conversation_id=None, + conversation_source=None, + user_id=None, + is_streaming=False, + model=ai_model, + provider=resolve_provider(ai_model), + is_robusta_model=getattr(ai.llm, "is_robusta_model", False), + meta={"check_mode": request.mode.value, "timeout": request.timeout}, + ) + # Execute the check using the shared function result: CheckResult = execute_check( check=check, ai=ai, verbose=False, console=None, + recorder_state=recorder_state, ) # Track notification statuses diff --git a/holmes/core/conversations_worker/models.py b/holmes/core/conversations_worker/models.py index 6f59315ab2..84df46c478 100644 --- a/holmes/core/conversations_worker/models.py +++ b/holmes/core/conversations_worker/models.py @@ -28,6 +28,12 @@ class ConversationTask(BaseModel): request_sequence: int metadata: Dict[str, Any] = Field(default_factory=dict) title: Optional[str] = None + # The Conversations row's user_id column (the human who started the + # chat). Used as a fallback for HolmesUsageEvents.user_id when the FE + # didn't include user_id in the user_message event's data — common + # because the runner-side Conversations row already has the value, so + # the FE has no reason to duplicate it into every per-turn event. + user_id: Optional[str] = None # Hydrated post-construction from events; not part of the validated row schema. _user_message_data: Dict[str, Any] = PrivateAttr(default_factory=dict) diff --git a/holmes/core/conversations_worker/worker.py b/holmes/core/conversations_worker/worker.py index 2f1f1f009d..b978fd4dfa 100644 --- a/holmes/core/conversations_worker/worker.py +++ b/holmes/core/conversations_worker/worker.py @@ -41,6 +41,10 @@ inject_frontend_tools, ) from holmes.core.tracing import TracingFactory +from holmes.core.usage_recorder import ( + build_chat_recorder_state, + stream_with_usage_recording, +) from holmes.utils.holmes_status import update_holmes_status_in_db from holmes.utils.stream import StreamEvents @@ -483,6 +487,11 @@ def _build_task_from_conversation_row( request_sequence=int(conv.get("request_sequence", 1)), metadata=conv.get("metadata") or {}, title=conv.get("title"), + # Conversations.user_id (set by the FE when it created the row) + # — surfaced on the task so per-turn ChatRequest construction + # can use it as a fallback when the user_message event's data + # doesn't carry user_id explicitly. + user_id=conv.get("user_id"), ) except Exception: logging.exception( @@ -617,6 +626,20 @@ def _process_conversation(self, task: ConversationTask) -> None: if data.get("tool_decisions"): enable_tool_approval = True + # AI usage tracking (HolmesUsageEvents) — resolve user_id and + # request_source with row-level fallbacks. The FE writes both onto + # the Conversations row when it creates the chat (user_id as a + # column, request_source under metadata) but doesn't necessarily + # repeat them in every user_message event's data. Without this + # fallback, follow-up turns produce HolmesUsageEvents rows with + # NULL user_id / request_source even though the values are known. + # Per-event data still wins so the FE can override per-turn (e.g. + # an alert-investigation chat that pivots to a freeform question). + resolved_user_id = data.get("user_id") or task.user_id + resolved_request_source = data.get("request_source") or ( + task.metadata.get("request_source") if task.metadata else None + ) + chat_request = ChatRequest( ask=ask, images=data.get("images"), @@ -630,7 +653,27 @@ def _process_conversation(self, task: ConversationTask) -> None: frontend_tool_results=data.get("frontend_tool_results"), # type: ignore[arg-type] response_format=data.get("response_format"), behavior_controls=data.get("behavior_controls"), - user_id=data.get("user_id"), + # source_ref / meta / is_internal still come from the per-event + # blob only — they're per-turn signals (which alert this + # follow-up question was about, etc.), not Conversation-level + # state. user_id / request_source fall back to the Conversations + # row when the FE didn't repeat them in the event. + user_id=resolved_user_id, + # request_type: pass through whatever the FE sent (None if absent) + # rather than hard-coding 'user_chat' here. The recorder helper + # (build_chat_recorder_state) handles the default and runs Slack + # auto-detection — hard-coding 'user_chat' would defeat the + # auto-detection because the helper bails out if request_type is + # already truthy. Today only /api/chat hits the Slack-prefix + # path, but the runner could route Slack through Conversations + # at any time without a code change here. + request_type=data.get("request_type"), + request_source=resolved_request_source, + source_ref=data.get("source_ref"), + conversation_id=task.conversation_id, + conversation_source="conversations", + meta=data.get("meta"), + is_internal=data.get("is_internal"), ) self._run_chat_and_publish( @@ -787,7 +830,19 @@ def _run_chat_and_publish( request_context = {"user_id": chat_request.user_id} try: - stream = request_ai.call_stream( + # Wrap the raw stream with the usage recorder BEFORE the + # publisher consumes it, so the recorder sees Holmes' native + # StreamMessage events (TOOL_RESULT / ANSWER_END / etc.) and + # can fire one HolmesUsageEvents row per worker-driven turn. + # Mirrors the wiring in server.py::chat() for the streaming + # path; without this the worker bypasses the recorder entirely. + recorder_state = build_chat_recorder_state( + chat_request, + request_ai, + dal=self.dal, + is_streaming=True, + ) + raw_stream = request_ai.call_stream( msgs=messages, enable_tool_approval=chat_request.enable_tool_approval or False, tool_decisions=chat_request.tool_decisions, @@ -796,6 +851,7 @@ def _run_chat_and_publish( request_context=request_context, trace_span=trace_span, ) + stream = stream_with_usage_recording(raw_stream, recorder_state) terminal = publisher.consume(stream) if terminal is None: diff --git a/holmes/core/models.py b/holmes/core/models.py index d0c9c5332d..bb1cbd055e 100644 --- a/holmes/core/models.py +++ b/holmes/core/models.py @@ -214,6 +214,69 @@ class ChatRequestBaseModel(BaseModel): ) user_id: Optional[str] = None # User ID from relay session token validation + # ── AI usage tracking fields (HolmesUsageEvents). All optional / additive; + # old clients that don't supply them keep working unchanged. ── + request_type: Optional[str] = Field( + default=None, + description=( + "Backend-set classification: 'user_chat' (default for /api/chat), " + "'scheduled_prompt' (set by ScheduledPromptsExecutor), 'agui_chat' " + "(set by AG-UI handler), 'health_check' (set by /api/checks/execute)." + ), + ) + request_source: Optional[str] = Field( + default=None, + description=( + "FE-supplied UI flow label, free-form. Examples: 'freeform', " + "'followup_logs', 'alert_investigation', 'resource_chat'." + ), + ) + source_ref: Optional[str] = Field( + default=None, + description=( + "FE-supplied opaque pointer to the entity the chat is about " + "(e.g. an issue id when request_source='alert_investigation'). " + "Meaning is implied by request_source." + ), + ) + conversation_id: Optional[str] = Field( + default=None, + description=( + "Stable id grouping multi-turn chats. Soft reference (NOT a FK): " + "matches Conversations.conversation_id when worker handles the chat, " + "or the FE-owned ChatHistory id for direct /api/chat traffic. NULL for " + "single-turn / non-UI flows." + ), + ) + conversation_source: Optional[str] = Field( + default=None, + description=( + "Discriminator telling dashboards which table conversation_id targets: " + "'conversations' (worker path) or 'chat_history' (direct /api/chat). " + "Worker sets it explicitly; chat() defaults to 'chat_history' when " + "conversation_id is non-NULL and not already set." + ), + ) + meta: Optional[Dict[str, Any]] = Field( + default=None, + description=( + "Forward-compatibility metadata bag. FE-supplied opaque dict; the " + "server shallow-merges with backend-derived keys (backend wins on " + "collision). Keep small; promote stable keys to real columns over time. " + "Do NOT put PII / large strings (prompts, completions, tool outputs) here." + ), + ) + is_internal: Optional[bool] = Field( + default=None, + description=( + "Marks server-internal calls (title generation, classification, " + "summarization, etc.) so dashboards can filter them out of user-facing " + "metrics. FE sets True for those. When unset, the server defaults it " + "to True if request_source starts with 'internal_' (backwards compat " + "with the prefix convention) — otherwise False." + ), + ) + # In our setup with litellm, the first message in conversation_history # should follow the structure [{"role": "system", "content": ...}], # where the "role" field is expected to be "system". diff --git a/holmes/core/scheduled_prompts/executor.py b/holmes/core/scheduled_prompts/executor.py index afb0d035c2..8594bed61b 100644 --- a/holmes/core/scheduled_prompts/executor.py +++ b/holmes/core/scheduled_prompts/executor.py @@ -204,6 +204,10 @@ def _execute_prompt( additional_system_prompt=additional_system_prompt, trace_span=heartbeat_span, behavior_controls=behavior_controls, + # AI usage tracking — these runs are server-driven, not user-driven. + request_type="scheduled_prompt", + request_source="scheduler", + source_ref=sp.id, ) empty_request = Request(scope={"type": "http", "headers": []}) diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index 2132cd065b..d54619fdab 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -7,7 +7,7 @@ import threading from datetime import datetime, timedelta from enum import Enum -from typing import Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple from uuid import uuid4 import sentry_sdk @@ -49,6 +49,12 @@ from holmes.utils.global_instructions import Instructions from holmes.utils.krr_utils import calculate_krr_savings +if TYPE_CHECKING: + # Forward reference only — `usage_recorder` already TYPE_CHECKING-imports + # this module, so importing the other direction at runtime would close + # the cycle. We just need the name for the parameter annotation. + from holmes.core.usage_recorder import UsageRecorderState + SUPABASE_TIMEOUT_SECONDS = int(os.getenv("SUPABASE_TIMEOUT_SECONDS", 60)) # Maximum total rows to fetch from KRR scans, regardless of number of clusters @@ -69,6 +75,7 @@ CONVERSATIONS_TABLE = "Conversations" CONVERSATION_EVENTS_TABLE = "ConversationEvents" OAUTH_TOKENS_TABLE = "OAuthTokens" +HOLMES_USAGE_EVENTS_TABLE = "HolmesUsageEvents" ENRICHMENT_BLACKLIST = ["text_file", "graph", "ai_analysis", "holmes"] ENRICHMENT_BLACKLIST_SET = set(ENRICHMENT_BLACKLIST) @@ -820,6 +827,70 @@ def sync_toolsets(self, toolsets: list[dict], cluster_name: str) -> None: f"An error occurred during toolset synchronization: {e}", exc_info=True ) + def record_usage_event(self, state: "UsageRecorderState") -> None: + """Record one HolmesUsageEvents row. Best-effort: swallows DB errors. + + Called from UsageRecorderState._fire on a daemon thread, so + errors here only affect the telemetry row, never the request + response. Takes a ``UsageRecorderState`` and reads only the fields + that map to columns — this is the single place that knows the + column shape, so adding a new field is "add it on the state, read + it here, write the migration." The DAL doesn't import the state + class at runtime (TYPE_CHECKING-only); attribute access is duck- + typed, so any object with the right shape works (handy for tests). + """ + if not self.enabled: + return + try: + stats = state.stats # may be None on aborted/error rows + self.client.table(HOLMES_USAGE_EVENTS_TABLE).insert({ + "account_id": self.account_id, + "cluster_id": state.cluster_id or self.cluster, + "user_id": state.user_id, + "conversation_id": state.conversation_id, + "conversation_source": state.conversation_source, + "request_id": state.request_id, + "request_type": state.request_type, + "request_source": state.request_source, + "source_ref": state.source_ref, + "status": state.status, + "model": state.model, + "provider": state.provider, + "is_robusta_model": state.is_robusta_model, + # Stats may be None when the request never reached a terminal + # event with cost data (aborted / pre-LLM error). The getattr + # default keeps the row writable in those cases. + "prompt_tokens": getattr(stats, "prompt_tokens", 0) or 0, + "completion_tokens": getattr(stats, "completion_tokens", 0) or 0, + "cached_tokens": getattr(stats, "cached_tokens", None), + "reasoning_tokens": getattr(stats, "reasoning_tokens", 0) or 0, + "total_tokens": getattr(stats, "total_tokens", 0) or 0, + "total_cost": float(getattr(stats, "total_cost", 0.0) or 0.0), + "num_compactions": getattr(stats, "num_compactions", 0) or 0, + "iterations": state.iterations, + "max_prompt_tokens_per_call": getattr( + stats, "max_prompt_tokens_per_call", 0 + ) or 0, + "max_completion_tokens_per_call": getattr( + stats, "max_completion_tokens_per_call", 0 + ) or 0, + "tool_call_count": state.tool_call_count, + "duration_ms": state.duration_ms, + "is_streaming": state.is_streaming, + "is_internal": state.is_internal, + "finish_reason": state.finish_reason, + "meta": state.meta or {}, + }).execute() + except Exception: + logging.exception("Failed to record usage event") + + # NOTE: feedback writes (thumbs up/down + category + comment) do NOT go + # through Holmes. The frontend calls the public.record_feedback() Postgres + # function directly via supabase.rpc('record_feedback', ...). The function + # runs `security invoker` and scopes by `auth.uid()`, which is a stricter + # user-scoping than any FE-supplied user_id we could pass through here. + # See plan section G and the migration script for the function body. + def has_scheduled_prompt_definitions(self) -> bool: """ Check if the account has any scheduled prompt definitions. diff --git a/holmes/core/tool_calling_llm.py b/holmes/core/tool_calling_llm.py index 6062f5247f..8d8d109069 100644 --- a/holmes/core/tool_calling_llm.py +++ b/holmes/core/tool_calling_llm.py @@ -178,6 +178,7 @@ class LLMResult(RequestStats): instructions: List[str] = Field(default_factory=list) messages: Optional[List[dict]] = None metadata: Optional[Dict[Any, Any]] = None + finish_reason: Optional[str] = None # Last LLM iteration's finish_reason (stop / length / tool_calls / content_filter) class ToolCallWithDecision(BaseModel): @@ -597,6 +598,7 @@ def call( # type: ignore num_llm_calls=total_num_llm_calls, messages=terminal_data.get("messages"), metadata=terminal_data.get("metadata"), + finish_reason=(terminal_data.get("metadata") or {}).get("finish_reason"), **accumulated_stats.model_dump(), ) @@ -618,6 +620,7 @@ def call( # type: ignore num_llm_calls=total_num_llm_calls, messages=terminal_data["messages"], metadata=terminal_data.get("metadata"), + finish_reason=(terminal_data.get("metadata") or {}).get("finish_reason"), **accumulated_stats.model_dump(), ) @@ -1174,6 +1177,18 @@ def call_stream( tools_to_call = getattr(response_message, "tool_calls", None) if not tools_to_call: + # Capture the final iteration's finish_reason for usage tracking + # (HolmesUsageEvents.finish_reason). Earlier iterations always end + # with 'tool_calls'; this last one tells us why the loop terminated + # (stop / length / content_filter / etc.). Skip if the value isn't + # a real string (e.g. MagicMock in tests), so pydantic validation + # of LLMResult below doesn't blow up. + try: + fr = full_response.choices[0].finish_reason # type: ignore + if isinstance(fr, str): + metadata["finish_reason"] = fr + except (AttributeError, IndexError, TypeError): + pass yield StreamMessage( event=StreamEvents.ANSWER_END, data={ diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py new file mode 100644 index 0000000000..bcde3e238d --- /dev/null +++ b/holmes/core/usage_recorder.py @@ -0,0 +1,580 @@ +"""Shared helper for recording AI usage events to HolmesUsageEvents. + +Used by every LLM-consuming entry point (server.py /api/chat, the +ConversationWorker, scheduled prompts, the AG-UI server, and +holmes/checks/checks_api.py) so usage tracking is consistent and there's +exactly one place to update if the recording shape changes. + +The recorder is fire-and-forget: each call spawns a daemon thread to do +the DB write. Telemetry must never block or break the response path. +""" +from __future__ import annotations + +import logging +import re +import time +import uuid +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Any, Dict, Generator, Optional + +from holmes.core.llm_usage import RequestStats +from holmes.utils.stream import StreamEvents, StreamMessage + +if TYPE_CHECKING: + from holmes.core.models import ChatRequest + from holmes.core.supabase_dal import SupabaseDal + from holmes.core.tool_calling_llm import LLMResult + + +# Bounded thread pool for recorder DB writes. Caps concurrent supabase-py +# connection acquisitions from telemetry so a slow Supabase / connection +# leak in the recorder can't starve foreground writes (HolmesStatus, +# ToolStatus). max_workers=4 is enough for Holmes' single-pod-per-customer +# load — one Supabase write is ~50–200ms, so 4 workers handle ~80 events/sec +# sustained, well above current request rate. Process-exit semantics: the +# stdlib's atexit handler drains live executors, which is *better* than the +# previous daemon-thread fire-and-forget — we lose fewer rows on graceful +# shutdown. Threads are spawned lazily on first submit; importing this +# module doesn't start any. +_RECORDER_EXECUTOR = ThreadPoolExecutor( + max_workers=4, + thread_name_prefix="usage-recorder", +) + + +# Slack auto-detection: the Robusta runner's Slack handler currently prepends +# a fixed prefix to the user's message before POSTing /api/chat. Example: +# "**@user_U0AKMP2CZ97** • 2026-05-04T05:10:04Z\n\nhigh cpu in pod alert" +# Extracted into a shared regex so both the direct /api/chat path (server.py) +# and the worker path (conversations_worker/worker.py) can run the same +# detection. Heuristic — fragile if the runner format changes. +_SLACK_ASK_PREFIX_RE = re.compile( + r"^\*\*@user_(?PU[A-Z0-9]+)\*\*\s*•\s*" + r"(?P\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z)", +) + + +def detect_slack_origin(ask: Optional[str]) -> Optional[Dict[str, Any]]: + """If `ask` matches the runner's Slack-prefix shape, return parsed + metadata (slack_user_id + slack_triggered_at). Otherwise None. + """ + if not ask: + return None + m = _SLACK_ASK_PREFIX_RE.match(ask) + if not m: + return None + return { + "slack_user_id": m.group("slack_user_id"), + "slack_triggered_at": m.group("slack_triggered_at"), + } + + +def resolve_provider(model: Optional[str]) -> str: + """Best-effort: return the canonical litellm provider for `model`. + + Falls back to splitting on the litellm prefix (`openai/...`, + `anthropic/...`) if the helper raises (e.g. for unrecognized models). + Importing litellm is deferred so this module stays import-light for + consumers that only want the dataclass. + """ + if not model: + return "unknown" + try: + import litellm # local import — keep usage_recorder cheap to import + return litellm.get_llm_provider(model)[1] or "unknown" + except Exception: + return model.split("/")[0] if "/" in model else "unknown" + + +def build_chat_recorder_state( + chat_request: "ChatRequest", + request_ai: Any, + *, + dal: Any, + is_streaming: bool, +) -> "UsageRecorderState": + """Construct a UsageRecorderState from a ChatRequest. + + Used by every code path that consumes a ChatRequest and wraps a stream: + server.py:chat() (direct /api/chat) and ConversationWorker._run_chat_and_publish + (worker path). Centralizes the request_type / request_source / is_internal + / Slack auto-detection logic so all entry points get identical behavior. + """ + # Default conversation_source to 'chat_history' when conversation_id is set + # but the caller didn't override (i.e. direct /api/chat). The worker passes + # 'conversations' explicitly. + conversation_source = chat_request.conversation_source + if conversation_source is None and chat_request.conversation_id: + conversation_source = "chat_history" + + model_name = ( + getattr(request_ai.llm, "model", None) + or chat_request.model + or "unknown" + ) + + # Internal calls (title generation, classification, summarization, etc.) + # get filtered out of user-facing dashboards. FE sets is_internal=True + # explicitly for those. Backwards compat: if FE didn't set it, fall back + # to detecting the legacy 'internal_' prefix on request_source. + if chat_request.is_internal is None: + is_internal = bool( + chat_request.request_source + and chat_request.request_source.startswith("internal_") + ) + else: + is_internal = bool(chat_request.is_internal) + + # Slack auto-detection: tag both request_type='slack_chat' and + # request_source='slack' as defaults that the caller can still override. + slack_info = detect_slack_origin(chat_request.ask) + if chat_request.request_type: + request_type = chat_request.request_type + elif slack_info is not None: + request_type = "slack_chat" + else: + request_type = "user_chat" + + request_source = chat_request.request_source + if request_source is None and slack_info is not None: + request_source = "slack" + + # Merge meta: FE-supplied keys, then backend-derived keys (backend wins + # on collision). Slack info goes under a 'slack' sub-key so it doesn't + # clutter the top level. + merged_meta: Dict[str, Any] = dict(chat_request.meta or {}) + if slack_info is not None: + merged_meta["slack"] = slack_info + + return UsageRecorderState( + dal=dal, + request_type=request_type, + request_source=request_source, + source_ref=chat_request.source_ref, + conversation_id=chat_request.conversation_id, + conversation_source=conversation_source, + user_id=chat_request.user_id, + is_streaming=is_streaming, + is_internal=is_internal, + model=model_name, + provider=resolve_provider(model_name), + is_robusta_model=getattr(request_ai.llm, "is_robusta_model", False), + meta=merged_meta, + ) + + +class RequestStatus(str, Enum): + """Final outcome of a request, written to HolmesUsageEvents.status. + + Subclassing both ``str`` and ``Enum`` keeps the values transparent to + JSON serializers (supabase-py sends them as plain strings) and + backwards-compatible with anything that compares against the old + string literals. Add a new variant here whenever the recorder needs + to surface a new outcome — the DB column is plain ``text`` so no + migration is required to widen the set. + """ + + SUCCESS = "success" # terminal ANSWER_END event seen + APPROVAL_REQUIRED = "approval_required" # terminal APPROVAL_REQUIRED seen + ERROR = "error" # terminal ERROR event or unhandled exception + RATE_LIMITED = "rate_limited" # provider rate-limit detected by record_error + ABORTED = "aborted" # stream ended without any terminal event + + +@dataclass +class UsageRecorderState: + """All the data needed to write one HolmesUsageEvents row. + + Fields fall into three groups: + + 1. **Required identity / classification** — set at construction time + by the entry point (server.chat, the worker, AG-UI, scheduled + prompts, health checks). The recorder cannot run without these. + + 2. **Optional identity / classification** — also set at construction. + NULL/default values are written through to the DB column as-is. + + 3. **Mutable runtime fields** — left at their defaults at construction + and filled in by ``stream_with_usage_recording`` (streaming path) + or ``record_from_llm_result`` / ``record_error`` (non-streaming + path) just before the row is fired. Don't set these on the entry + point side; they get clobbered. + + For each field below: the comment lists where the value comes from, + what it means, and what gets written if the field is left at its + default. + """ + + # ── Group 1: required identity / classification ───────────────────── + + # SupabaseDal handle. The recorder calls dal.record_usage_event(state) + # on a daemon thread when the request finishes — passing this entire + # state object positionally; the DAL reads the fields it needs. + # Typed Any to avoid importing SupabaseDal at runtime (circular import + # via ChatRequest). + dal: Any + + # Backend taxonomy of which call surface initiated this chat. Stable + # values dashboards group by: + # 'user_chat' — direct POST /api/chat (default) + # 'slack_chat' — auto-detected from the runner's Slack prefix + # in chat_request.ask, OR set explicitly by a + # future runner-side change + # 'agui_chat' — set by the AG-UI handler + # 'scheduled_prompt' — set by ScheduledPromptsExecutor + # 'health_check' — set by /api/checks/execute + request_type: str + + # LLM model string after Holmes' model routing. Whatever litellm + # accepts: 'anthropic/claude-sonnet-4-5', 'openai/gpt-4o', etc. + # Sourced from request_ai.llm.model in build_chat_recorder_state. + model: str + + # Canonical litellm provider derived from `model`. Use the helper + # `resolve_provider(model)`. Examples: 'anthropic', 'openai', 'azure', + # 'bedrock'. Falls back to 'unknown' if litellm can't classify. + provider: str + + # True when the LLM call hit a Robusta-managed model (Robusta paid + # the token bill). Read from request_ai.llm.is_robusta_model. SaaS + # dashboards filter on this column to hide cost / show only input + # tokens for managed-model rows. + is_robusta_model: bool + + # ── Group 2: optional identity / classification ───────────────────── + + # FE/runner-supplied finer UI flow label. Free-form text — adding new + # values doesn't need a Holmes migration. Examples: + # 'freeform' — user typed a question in the FE + # 'alert_investigation' — chat opened from an alert + # 'followup_logs' — user clicked a follow-up action button + # 'resource_chat' — chat opened from a Kubernetes resource + # 'slack' / 'teams' — auto-detected for messaging-platform + # chats; runner can override with finer + # values like 'slack_alert_investigation' + # 'scheduler' / 'operator' — set by scheduled-prompt / health-check + request_source: Optional[str] = None + + # Opaque pointer to the entity the chat is *about*. Meaning is + # implied by `request_source`. Examples: + # request_source='alert_investigation' → source_ref= + # request_source='resource_chat' → source_ref= + # request_source='operator' (checks) → source_ref= + # Free-form text; not a foreign key — the table this points to varies. + source_ref: Optional[str] = None + + # Stable id grouping multi-turn chats so dashboards can show per- + # conversation cost / token totals. Soft reference (NOT a FK). + # Either matches Conversations.conversation_id (worker path) or the + # FE-owned ChatHistory id (legacy /api/chat path); the discriminator + # below tells dashboards which table to LEFT JOIN. NULL for single- + # turn / non-UI flows (CLI, scheduled prompts, health checks). + conversation_id: Optional[str] = None + + # Discriminator telling dashboards which table `conversation_id` + # targets when it's non-NULL: + # 'conversations' — worker path (set explicitly by the worker) + # 'chat_history' — direct /api/chat path (defaulted by chat() when + # conversation_id is set and no override given) + # None — no conversation context + conversation_source: Optional[str] = None + + # UUID of the human who started the chat. Sourced from the auth- + # validated session token (server.chat) or the Conversations row + # (worker fallback). NULL for system / scheduled flows where there + # is no human user. Used for per-user analytics and (today) by the + # feedback RPC's auth.uid() check. + user_id: Optional[str] = None + + # Holmes' cluster id (env-var string, e.g. 'production-east'). Falls + # back to dal.cluster inside record_usage_event when left at None + # here. NULL for CLI / ad-hoc flows that have no cluster context. + cluster_id: Optional[str] = None + + # Per-request UUID. Auto-generated; you don't set this manually. The + # stream wrapper injects it into the terminal event's `metadata` so + # the FE can read it from `ai_answer_end` and pass it to the + # `record_feedback` Supabase RPC later when the user clicks 👍/👎. + request_id: str = field(default_factory=lambda: str(uuid.uuid4())) + + # True for streaming responses (SSE), False for non-streaming. The + # worker is always streaming; /api/chat respects chat_request.stream; + # health checks are always non-streaming. Set by the entry point. + is_streaming: bool = False + + # Marks server-internal calls (title generation, classifier prompts, + # summarization, etc.) so user-facing dashboards can filter them out + # of activity / cost metrics. Defaults False; set True at the entry + # point for those flows. build_chat_recorder_state also auto-sets + # True when request_source starts with 'internal_' (legacy convention). + is_internal: bool = False + + # Forward-compatibility metadata bag. JSONB-serializable dict that + # gets shallow-merged with backend-derived keys (backend wins on + # collision) and stored as-is in the row's `meta` column. + # Conventions: + # - Slack auto-detect populates meta['slack'] with slack_user_id / + # slack_triggered_at when the regex matches. + # - FE can add provisional fields (experiment_id, etc.) that + # haven't earned a real column yet. + # - Don't put PII or large strings (prompts, completions, tool + # output) here — that's out of scope for v1. + meta: Dict[str, Any] = field(default_factory=dict) + + # ── Group 3: mutable runtime fields — DO NOT set at construction ──── + + # Wall-clock start of the request (time.monotonic seconds). Used to + # compute `duration_ms` when the row is written. Defaults to + # construction time, which is when the entry point starts the work. + t_start: float = field(default_factory=time.monotonic) + + # Aggregate token / cost counters. Filled by `stream_with_usage_recording` + # from the ANSWER_END event's `metadata.costs`, or by + # `record_from_llm_result` from the LLMResult's RequestStats fields. + # None means the request didn't reach a terminal event with stats — + # the recorder still writes a zero-stats row with status='aborted'/'error'. + stats: Optional[RequestStats] = None + + # Number of LLM round-trips in this request (1 for a simple chat, N + # for an agentic loop with N-1 tool turns). Filled by the wrapper + # from the terminal event's `num_llm_calls`. 0 means no LLM call + # completed (aborted / errored before the first response). + iterations: int = 0 + + # Number of TOOL_RESULT events the wrapper saw flow past during the + # stream. Filled by the wrapper. The non-streaming path + # (record_from_llm_result) instead reads len(llm_result.tool_calls). + tool_call_count: int = 0 + + # Last LLM iteration's finish reason: 'stop', 'length', 'tool_calls', + # 'content_filter', etc. Filled by the wrapper from the terminal + # event's `metadata.finish_reason`. Earlier iterations always end in + # 'tool_calls'; only the final iteration carries a meaningful value. + finish_reason: Optional[str] = None + + # Final outcome of the request. Filled by the wrapper or the + # record_error helper. See ``RequestStatus`` above for the full + # set of values and their meanings. The default ``SUCCESS`` is + # overwritten by the wrapper's finally-block to ``ABORTED`` if no + # terminal event was ever observed. + status: RequestStatus = RequestStatus.SUCCESS + + @property + def duration_ms(self) -> int: + """Wall-clock milliseconds since ``t_start``. + + Computed on read so the value reflects "now minus when the request + started" at the moment the row is written. Used by + ``SupabaseDal.record_usage_event`` to populate the duration column. + """ + return int((time.monotonic() - self.t_start) * 1000) + + # ── private helpers — kept on the class because they only operate on + # this state. The public entry points (stream_with_usage_recording, + # record_from_llm_result, record_error) stay as module-level functions + # — the stream wrapper in particular has the stream as its primary + # input, so a "method on state" shape would invert its natural reading. + + def _capture_costs(self, data: Dict[str, Any]) -> None: + """Replace ``self.stats`` from an event's ``metadata.costs``. + + Called for both terminal events (ANSWER_END / APPROVAL_REQUIRED / + ERROR) and mid-stream TOKEN_COUNT events. Each event carries the + cumulative cost up to that point, so the latest one always wins — + which is what we want: ANSWER_END's costs == final TOKEN_COUNT's + costs in the success case, and the last seen TOKEN_COUNT gives + partial cost in the mid-loop-exception case. + """ + metadata = data.get("metadata") or {} + costs = metadata.get("costs") or {} + if not costs: + return + try: + self.stats = RequestStats(**costs) + except Exception: + logging.debug( + "Failed to materialize RequestStats from event costs", + exc_info=True, + ) + + def _capture_terminal(self, data: Dict[str, Any]) -> None: + """Pull cost/iterations/finish_reason from a terminal event's data. + + Terminal events (ANSWER_END / APPROVAL_REQUIRED / ERROR) carry the + full picture: costs, iteration count, finish reason. Mid-stream + TOKEN_COUNT events carry only costs — those go through + ``_capture_costs`` directly. + """ + self._capture_costs(data) + # Explicit None-check rather than `or` so a legitimate 0 (unlikely + # but not impossible) is preserved instead of falling back to + # self.iterations. + raw_iterations = data.get("num_llm_calls") + if raw_iterations is not None: + self.iterations = raw_iterations + metadata = data.get("metadata") or {} + self.finish_reason = ( + metadata.get("finish_reason") or self.finish_reason + ) + + def _fire(self) -> None: + """Submit the dal write to the shared recorder thread pool. + + Fire-and-forget — the response path never waits on this. The + executor caps concurrent writes (see ``_RECORDER_EXECUTOR`` at + module top); under burst load, additional submissions queue + inside the executor rather than spawning unbounded fresh threads. + + ``executor.submit`` raises ``RuntimeError`` if the executor has + already been shut down (process exiting). Treat that as accepted + loss — same fate as in-flight rows on the previous daemon-thread + fire-and-forget shape. + """ + if self.dal is None or not getattr(self.dal, "enabled", False): + return + try: + _RECORDER_EXECUTOR.submit(self.dal.record_usage_event, self) + except RuntimeError: + # Executor was shut down — accept the loss. + logging.debug( + "Usage recorder executor is shut down; dropping row", + exc_info=True, + ) + except Exception: + # Defense in depth — record_usage_event has its own try/except too. + logging.exception("Failed to submit usage recorder write") + + +def stream_with_usage_recording( + stream: Generator[StreamMessage, None, None], + state: UsageRecorderState, +) -> Generator[StreamMessage, None, None]: + """Forward stream events; capture state; record on stream end. + + Used by chat() and AG-UI. Watches for terminal events (ANSWER_END, + APPROVAL_REQUIRED, ERROR) to extract final stats / counts / reason, + counts TOOL_RESULT events along the way, and fires the recorder in + a `finally` block so the row is written even on exceptions or + client disconnects. + + Also injects ``state.request_id`` into the terminal event's + ``metadata`` dict so the SSE formatter ships it back to the FE. The + FE saves it from ``ai_answer_end`` and passes it to the + ``public.record_feedback()`` Supabase RPC when the user clicks + thumbs up/down. + """ + saw_terminal = False + try: + for msg in stream: + if msg.event == StreamEvents.TOOL_RESULT: + state.tool_call_count += 1 + elif msg.event == StreamEvents.TOKEN_COUNT: + # Cumulative cost broadcast after each successful LLM iteration + # (and after compaction). Capturing it here is the only way to + # record partial cost when the agentic loop raises mid-loop — + # call_stream's local `stats` accumulator gets GC'd along with + # the function frame on exception, so we'd otherwise write a + # zero-stats row even though earlier iterations burned real + # tokens. Each TOKEN_COUNT carries the running total, so the + # last one we see before the error tells us exactly how much + # the failed turn cost up to that point. + state._capture_costs(msg.data) + elif msg.event == StreamEvents.ANSWER_END: + state._capture_terminal(msg.data) + _inject_request_id(msg.data, state.request_id) + state.status = RequestStatus.SUCCESS + saw_terminal = True + elif msg.event == StreamEvents.APPROVAL_REQUIRED: + state._capture_terminal(msg.data) + _inject_request_id(msg.data, state.request_id) + state.status = RequestStatus.APPROVAL_REQUIRED + saw_terminal = True + elif msg.event == StreamEvents.ERROR: + state._capture_terminal(msg.data) + _inject_request_id(msg.data, state.request_id) + state.status = RequestStatus.ERROR + saw_terminal = True + yield msg + except Exception: + if not saw_terminal: + state.status = RequestStatus.ERROR + raise + finally: + # If the inner stream ended without yielding any terminal event + # (client disconnected mid-stream, generator exhausted abnormally), + # `state.status` would still be the constructor default SUCCESS. + # That's wrong — mark such cases as ABORTED so dashboards can + # filter incomplete runs out of "successful chat" metrics. + if not saw_terminal and state.status == RequestStatus.SUCCESS: + state.status = RequestStatus.ABORTED + state._fire() + + +def _inject_request_id(data: Dict[str, Any], request_id: str) -> None: + """Drop request_id into data['metadata'] so the SSE formatter ships it + to the FE. Creates the metadata dict if missing or non-dict-shaped. + + Stays a module-level function (not a method on UsageRecorderState) + because it operates on the stream event's data dict, not on state — + the only state field it reads is request_id, which it takes as an arg. + """ + md = data.get("metadata") + if not isinstance(md, dict): + md = {} + data["metadata"] = md + md["request_id"] = request_id + + +def record_from_llm_result( + state: UsageRecorderState, + llm_result: "LLMResult", +) -> None: + """Record a usage event from a non-streaming `ai.call(...)` result. + + Used by `holmes/checks/checks.py:execute_check` and any other caller + that gets back an LLMResult directly. LLMResult IS-A RequestStats + (it inherits the cost / token fields), so we copy them out via + model_dump. + """ + try: + # LLMResult inherits from RequestStats and adds extra fields + # (tool_calls, messages, finish_reason, ...). Ask Pydantic to filter + # the dump down to RequestStats's own model_fields so the extras get + # dropped without us hardcoding the stats field set here — when + # RequestStats grows a new column, this stays correct automatically. + state.stats = RequestStats( + **llm_result.model_dump(include=set(RequestStats.model_fields)) + ) + except Exception: + logging.debug("Failed to extract stats from LLMResult", exc_info=True) + state.stats = RequestStats() + + state.iterations = getattr(llm_result, "num_llm_calls", None) or 1 + state.tool_call_count = len(getattr(llm_result, "tool_calls", None) or []) + state.finish_reason = getattr(llm_result, "finish_reason", None) + state.status = RequestStatus.SUCCESS + state._fire() + + +def record_error(state: UsageRecorderState, exc: Exception) -> None: + """Record a failed call where an exception bubbled before getting a result.""" + msg = str(exc).lower() + if "rate" in msg and "limit" in msg: + state.status = RequestStatus.RATE_LIMITED + else: + state.status = RequestStatus.ERROR + state._fire() + + +__all__ = [ + "RequestStatus", + "UsageRecorderState", + "build_chat_recorder_state", + "detect_slack_origin", + "record_error", + "record_from_llm_result", + "resolve_provider", + "stream_with_usage_recording", +] diff --git a/server.py b/server.py index a393f90c36..adf3db67fd 100644 --- a/server.py +++ b/server.py @@ -27,7 +27,6 @@ from fastapi import FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, StreamingResponse from litellm.exceptions import AuthenticationError - from holmes import get_version, is_official_release from holmes.common.env_vars import ( DEVELOPMENT_MODE, @@ -70,6 +69,12 @@ inject_frontend_tools, ) from holmes.core.tracing import TracingFactory +from holmes.core.usage_recorder import ( + build_chat_recorder_state, + record_error, + record_from_llm_result, + stream_with_usage_recording, +) from holmes.utils.stream import stream_chat_formatter @@ -517,7 +522,12 @@ def chat(chat_request: ChatRequest, http_request: Request): inv_attrs = {"gen_ai_request_model": chat_request.model or config.model or "unknown"} otel_metrics.investigation_count.add(1, inv_attrs) - stream = stream_chat_formatter( + # Build the usage recorder state and wrap the raw stream BEFORE the + # SSE formatter so the wrapper sees Holmes' native StreamMessage events. + recorder_state = build_chat_recorder_state( + chat_request, request_ai, dal=dal, is_streaming=True + ) + recorded_stream = stream_with_usage_recording( request_ai.call_stream( msgs=messages, enable_tool_approval=chat_request.enable_tool_approval or False, @@ -527,6 +537,10 @@ def chat(chat_request: ChatRequest, http_request: Request): request_context=request_context, trace_span=trace_span, ), + recorder_state, + ) + stream = stream_chat_formatter( + recorded_stream, [f.model_dump() for f in follow_up_actions], model=chat_request.model or config.model, ) @@ -535,6 +549,9 @@ def chat(chat_request: ChatRequest, http_request: Request): media_type="text/event-stream", ) else: + recorder_state = build_chat_recorder_state( + chat_request, request_ai, dal=dal, is_streaming=False + ) try: # Use provided trace_span or create a root investigation span trace_span = chat_request.trace_span @@ -554,6 +571,9 @@ def chat(chat_request: ChatRequest, http_request: Request): request_context=request_context, ) + # Record usage event for non-streaming path (fire-and-forget). + record_from_llm_result(recorder_state, llm_call) + # Record investigation metrics otel_metrics = TracingFactory.get_metrics() if otel_metrics: @@ -572,14 +592,25 @@ def chat(chat_request: ChatRequest, http_request: Request): ) else: logging.info(f"Completed {req_info}") + # Surface request_id in the response metadata so the FE has a + # handle for the public.record_feedback() RPC later. Streaming + # path does the same via _inject_request_id in the stream wrapper. + response_metadata = dict(llm_call.metadata or {}) + response_metadata["request_id"] = recorder_state.request_id response = ChatResponse( analysis=llm_call.result, tool_calls=llm_call.tool_calls, conversation_history=llm_call.messages, follow_up_actions=follow_up_actions, - metadata=llm_call.metadata, + metadata=response_metadata, ) return response + except Exception as e: + # Non-streaming path: record the failed event so it shows up in dashboards + # with status='error' / 'rate_limited'. Streaming path records via the + # wrapper's `finally` automatically. + record_error(recorder_state, e) + raise finally: if trace_span is not None: trace_span.end() diff --git a/tests/core/conversations_worker/test_worker_lifecycle.py b/tests/core/conversations_worker/test_worker_lifecycle.py index c369345283..425f6bc191 100644 --- a/tests/core/conversations_worker/test_worker_lifecycle.py +++ b/tests/core/conversations_worker/test_worker_lifecycle.py @@ -45,6 +45,7 @@ def test_build_task_from_conversation_row_parses_required_fields(): "request_sequence": 3, "metadata": {"foo": "bar"}, "title": "hello", + "user_id": "u-42", } task = w._build_task_from_conversation_row(row) assert task is not None @@ -52,6 +53,10 @@ def test_build_task_from_conversation_row_parses_required_fields(): assert task.request_sequence == 3 assert task.metadata == {"foo": "bar"} assert task.title == "hello" + # user_id from the Conversations row is surfaced on the task so the + # ChatRequest construction can fall back to it when the per-event + # data doesn't carry user_id. + assert task.user_id == "u-42" def test_build_task_from_conversation_row_tolerates_missing_fields(): @@ -61,6 +66,9 @@ def test_build_task_from_conversation_row_tolerates_missing_fields(): assert task is not None assert task.request_sequence == 1 assert task.origin == "chat" + # user_id is optional on the Conversations row (e.g. older rows that + # predate the column); the task should still build cleanly. + assert task.user_id is None def test_build_task_from_conversation_row_returns_none_on_bad_input(): diff --git a/tests/core/conversations_worker/test_worker_usage_recorder.py b/tests/core/conversations_worker/test_worker_usage_recorder.py new file mode 100644 index 0000000000..2af3bf9ff7 --- /dev/null +++ b/tests/core/conversations_worker/test_worker_usage_recorder.py @@ -0,0 +1,386 @@ +"""Verify the ConversationWorker wires the usage recorder around the LLM +stream so worker-driven chats produce HolmesUsageEvents rows. + +Context: the worker takes a code path that bypasses server.py::chat() and +calls request_ai.call_stream(...) directly. Before this wiring, that path +never invoked the recorder, so worker-driven conversations (the new +/api/conversations flow) were silently missing from HolmesUsageEvents +while only server.py-direct calls were tracked. + +These tests assert the integration without re-testing the recorder itself +(which is covered in tests/core/test_usage_recorder.py): +1. ``stream_with_usage_recording`` is invoked with the raw stream. +2. The wrapped stream — not the raw stream — is what ``publisher.consume`` + receives, so the recorder gets a chance to observe terminal events. +3. The recorder state passed in carries the worker's classification + (``conversation_source='conversations'``, ``request_type='user_chat'``, + etc.) so dashboards can attribute these rows correctly. +""" +import threading +from collections import deque +from unittest.mock import MagicMock, patch + +from holmes.core.conversations_worker.models import ConversationTask +from holmes.core.conversations_worker.worker import ConversationWorker +from holmes.core.models import ChatRequest + + +def _bare_worker(): + w = ConversationWorker.__new__(ConversationWorker) + w.dal = MagicMock() + w.dal.enabled = True + w.dal.update_conversation_status = MagicMock(return_value=True) + w.dal.get_global_instructions_for_account = MagicMock(return_value=None) + w.config = MagicMock() + # create_toolcalling_llm returns the AI; we configure its llm attrs so + # build_chat_recorder_state can read model / is_robusta_model. + ai = MagicMock() + ai.llm = MagicMock() + ai.llm.model = "anthropic/claude-sonnet-4-5" + ai.llm.is_robusta_model = False + w.config.create_toolcalling_llm = MagicMock(return_value=ai) + w.config.get_skill_catalog = MagicMock(return_value=[]) + w.chat_function = MagicMock() + w.holmes_id = "h-test" + w._running = True + w._claim_thread = None + w._notify_event = threading.Event() + w._executor = MagicMock() + w._active_conversation_ids = set() + w._active_lock = threading.Lock() + w._queued_tasks = deque() + w._queued_lock = threading.Lock() + w._dispatch_lock = threading.Lock() + w._realtime_manager = None + return w, ai + + +def _task(): + return ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + ) + + +def _chat_request(): + return ChatRequest( + ask="why is my pod failing?", + stream=True, + request_type="user_chat", + conversation_id="c1", + conversation_source="conversations", # worker sets this explicitly + user_id="u-1", + ) + + +def _run(worker, ai): + """Drive _run_chat_and_publish with all heavy collaborators mocked. + + Returns the captured (raw_stream, recorder_state, wrapped_stream) so + individual tests can assert on each. + """ + raw_stream = iter(["raw-event-1", "raw-event-2"]) + wrapped_stream_sentinel = object() + + ai.call_stream = MagicMock(return_value=raw_stream) + # _inject_frontend_tools just returns the AI back when there are no + # frontend tools; bypass it so we don't need to mock the helper module. + worker._inject_frontend_tools = MagicMock(return_value=ai) + + publisher = MagicMock() + # consume returns ANSWER_END so the worker doesn't take the failed-conversation + # branch and try to call _fail_conversation. + from holmes.utils.stream import StreamEvents + publisher.consume = MagicMock(return_value=StreamEvents.ANSWER_END) + + captured = {} + with patch( + "holmes.core.conversations_worker.worker.stream_with_usage_recording" + ) as mock_wrap, patch( + "holmes.core.conversations_worker.worker.build_chat_recorder_state" + ) as mock_build_state, patch( + "holmes.core.conversations_worker.worker.build_chat_messages" + ) as mock_build_messages, patch( + "holmes.core.conversations_worker.worker.tool_result_storage" + ) as mock_storage, patch( + "holmes.core.conversations_worker.worker.TracingFactory" + ) as mock_tracing: + # build_chat_messages is heavy (Jinja, prompts) — return a fake list. + mock_build_messages.return_value = [{"role": "user", "content": "fake"}] + # tool_result_storage is a context manager. + mock_storage.return_value.__enter__ = MagicMock(return_value="/tmp/x") + mock_storage.return_value.__exit__ = MagicMock(return_value=False) + # Tracing returns a tracer that returns a span with .log/.end. + tracer = MagicMock() + span = MagicMock() + tracer.start_trace.return_value = span + mock_tracing.create_tracer.return_value = tracer + + recorder_state_sentinel = MagicMock(name="recorder_state") + mock_build_state.return_value = recorder_state_sentinel + mock_wrap.return_value = wrapped_stream_sentinel + + worker._run_chat_and_publish( + task=_task(), + chat_request=_chat_request(), + publisher=publisher, + ) + + captured["raw_stream"] = raw_stream + captured["wrap_call"] = mock_wrap.call_args + captured["build_state_call"] = mock_build_state.call_args + captured["wrapped_stream"] = wrapped_stream_sentinel + captured["recorder_state"] = recorder_state_sentinel + captured["publisher"] = publisher + + return captured + + +def test_stream_is_wrapped_with_usage_recorder(): + """The recorder wrapper must see the raw stream so it can observe + TOOL_RESULT / ANSWER_END events as they flow past.""" + worker, ai = _bare_worker() + captured = _run(worker, ai) + + wrap_call = captured["wrap_call"] + assert wrap_call is not None, ( + "stream_with_usage_recording was never called — the worker is " + "still bypassing the recorder." + ) + # First positional arg is the raw stream. + assert wrap_call.args[0] is captured["raw_stream"] + # Second positional arg is the recorder state. + assert wrap_call.args[1] is captured["recorder_state"] + + +def test_publisher_consumes_wrapped_stream_not_raw(): + """If the publisher consumed the raw stream directly, the recorder's + finally-block would never see the terminal event and would mark the row + 'aborted'. The wrapped stream must be the one passed to the publisher.""" + worker, ai = _bare_worker() + captured = _run(worker, ai) + + consume_args, _ = captured["publisher"].consume.call_args + assert consume_args[0] is captured["wrapped_stream"], ( + "publisher.consume() must receive the wrapped stream, not the raw one. " + f"Got {consume_args[0]!r}" + ) + + +def test_recorder_state_uses_workers_dal_and_streaming_flag(): + """build_chat_recorder_state must be called with the worker's dal and + is_streaming=True (worker is always streaming). Without this, telemetry + would either fall on the floor (no dal) or be misclassified as + non-streaming.""" + worker, ai = _bare_worker() + captured = _run(worker, ai) + + build_call = captured["build_state_call"] + assert build_call.kwargs.get("dal") is worker.dal + assert build_call.kwargs.get("is_streaming") is True + # Positional args are (chat_request, request_ai). + assert build_call.args[1] is ai + + +# -------------------------------------------------------------------------- +# user_id / request_source fallback into Conversations row + metadata. +# +# The FE writes user_id (column) and request_source (under metadata) onto +# the Conversations row when it creates a chat. It does NOT necessarily +# repeat them in every per-turn user_message event's data. Without these +# fallbacks, follow-up turns produce HolmesUsageEvents rows with NULL +# user_id and request_source even though the values are sitting on the +# Conversations row the worker already loaded. These tests pin the +# fallback behavior so a future refactor can't silently re-introduce the +# NULL-row bug the user reported. +# -------------------------------------------------------------------------- + +def _capture_chat_request_from_process(task, user_message_data): + """Drive _process_conversation just far enough to capture the + ChatRequest it constructs. Patches _run_chat_and_publish so the + LLM pipeline never runs.""" + worker, _ = _bare_worker() + worker.dal.get_conversation_events = MagicMock( + return_value=[{"event": "user_message", "data": user_message_data, "ts": "1"}] + ) + + captured = {} + + def capture(self, t, chat_request, publisher, resume_only=False): + captured["chat_request"] = chat_request + + with patch.object(ConversationWorker, "_run_chat_and_publish", capture): + worker._process_conversation(task) + + return captured.get("chat_request") + + +def test_user_id_falls_back_to_conversations_row_when_event_omits_it(): + # FE wrote user_id onto the Conversations row but didn't repeat it in + # the per-turn user_message data — exactly the symptom the user + # reported (HolmesUsageEvents.user_id NULL despite the value being + # known to the worker). + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + user_id="u-conversations-row", + ) + cr = _capture_chat_request_from_process(task, {"ask": "follow-up?"}) + assert cr is not None + assert cr.user_id == "u-conversations-row" + + +def test_event_user_id_wins_over_conversations_row(): + # If the FE DOES repeat user_id in the event data, that value wins — + # so a future flow that lets users hand-off a chat could still record + # the per-turn rater. The Conversations-row value is a fallback, not + # an override. + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + user_id="u-conversations-row", + ) + cr = _capture_chat_request_from_process( + task, {"ask": "q", "user_id": "u-from-event"} + ) + assert cr is not None + assert cr.user_id == "u-from-event" + + +def test_request_source_falls_back_to_conversations_metadata(): + # FE puts request_source under Conversations.metadata when it creates + # the row. Per-turn events typically don't repeat it. Worker should + # pull from metadata so dashboards can slice by request_source even + # for follow-up turns. + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + metadata={"request_source": "alert_investigation", "other": "x"}, + ) + cr = _capture_chat_request_from_process(task, {"ask": "follow-up?"}) + assert cr is not None + assert cr.request_source == "alert_investigation" + + +def test_event_request_source_wins_over_conversations_metadata(): + # Same caller-wins semantic as user_id. The conversation may have been + # created from one surface ('alert_investigation') but a follow-up + # turn could legitimately re-classify itself ('freeform'). + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + metadata={"request_source": "alert_investigation"}, + ) + cr = _capture_chat_request_from_process( + task, {"ask": "q", "request_source": "freeform"} + ) + assert cr is not None + assert cr.request_source == "freeform" + + +def test_no_fallback_values_anywhere_yields_null(): + # Defense: if neither the row nor the event carries either field, the + # ChatRequest must still build (just with NULLs that the recorder + # writes through to the row). + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + ) + cr = _capture_chat_request_from_process(task, {"ask": "q"}) + assert cr is not None + assert cr.user_id is None + assert cr.request_source is None + + +# -------------------------------------------------------------------------- +# request_type passthrough. +# +# The worker used to hard-code request_type='user_chat' on the ChatRequest +# it constructed. That defeated build_chat_recorder_state's auto-detection +# logic: the helper only auto-classifies (Slack-prefix etc.) when +# chat_request.request_type is falsy, so a hard-coded value short-circuited +# every detection path. These tests pin the passthrough behavior so future +# refactors can't silently re-introduce that bug. +# -------------------------------------------------------------------------- + + +def test_request_type_passes_through_from_event_data(): + # FE-supplied request_type wins. Today only /api/chat sees this, but + # the runner could write request_type into the user_message blob at + # any time without code changes here. + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + ) + cr = _capture_chat_request_from_process( + task, {"ask": "q", "request_type": "scheduled_prompt"} + ) + assert cr is not None + assert cr.request_type == "scheduled_prompt" + + +def test_request_type_unset_when_event_omits_it(): + # Critical: when the event doesn't supply request_type, the worker + # MUST leave it None on the ChatRequest so build_chat_recorder_state's + # auto-detection (Slack prefix, fallback default) gets to run. If the + # worker hard-codes 'user_chat' here, Slack rows get mis-tagged. + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + ) + cr = _capture_chat_request_from_process(task, {"ask": "q"}) + assert cr is not None + assert cr.request_type is None, ( + "Worker must leave request_type=None so build_chat_recorder_state " + "can auto-detect (e.g. Slack prefix → 'slack_chat'). Hard-coding " + "'user_chat' defeats the helper's detection path." + ) + + +def test_slack_prefix_in_event_ask_routes_to_slack_chat_via_helper(): + # End-to-end check: an ask carrying the runner's Slack prefix must + # arrive at build_chat_recorder_state with request_type=None so the + # helper tags it as 'slack_chat'. We don't call the helper here — + # that's covered in tests/test_chat_recorder_state.py — we only + # assert the worker passes the right inputs into it. + task = ConversationTask( + conversation_id="c1", + account_id="a1", + cluster_id="cl1", + origin="chat", + request_sequence=1, + ) + slack_ask = ( + "**@user_U0AKMP2CZ97** • 2026-05-04T05:10:04Z\n\n" + "high cpu in pod alert" + ) + cr = _capture_chat_request_from_process(task, {"ask": slack_ask}) + assert cr is not None + # ChatRequest.ask carries the original prefix (helper inspects this). + assert cr.ask.startswith("**@user_U0AKMP2CZ97**") + # And request_type is None so the helper's auto-detect runs. + assert cr.request_type is None diff --git a/tests/core/test_supabase_dal_usage.py b/tests/core/test_supabase_dal_usage.py new file mode 100644 index 0000000000..8551c9c670 --- /dev/null +++ b/tests/core/test_supabase_dal_usage.py @@ -0,0 +1,208 @@ +"""Unit tests for SupabaseDal.record_usage_event. + +Feedback writes are not handled by Holmes (the FE calls the +public.record_feedback() Postgres RPC directly via supabase.rpc), so there +are no Holmes-side feedback unit tests here — coverage lives in the +migration's integration verification. + +record_usage_event is best-effort: it swallows Supabase errors so the +response path can never be broken by a telemetry write. The tests verify +both the happy path (correct payload sent to .insert) and the failure path +(exceptions are absorbed). + +The DAL takes a ``UsageRecorderState`` positional arg (Moshe's review on +PR #1969 — single-object signature avoids drift between the recorder and +DAL when fields are added). Tests construct a real UsageRecorderState via +the ``_make_state`` helper below; mock_dal is duck-typed so any object +with the same attributes would also work. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from holmes.core.llm_usage import RequestStats +from holmes.core.supabase_dal import ( + HOLMES_USAGE_EVENTS_TABLE, + SupabaseDal, +) +from holmes.core.usage_recorder import UsageRecorderState + + +@pytest.fixture +def mock_dal(): + """A SupabaseDal with mocked Supabase client and account_id.""" + with patch("holmes.core.supabase_dal.create_client"): + dal = SupabaseDal(cluster="test-cluster") + dal.enabled = True + dal.account_id = "00000000-0000-0000-0000-000000000001" + dal.cluster = "test-cluster" + dal.client = MagicMock() + return dal + + +def _stats() -> RequestStats: + return RequestStats( + total_cost=0.0123, + total_tokens=1234, + prompt_tokens=1000, + completion_tokens=234, + cached_tokens=50, + reasoning_tokens=12, + max_completion_tokens_per_call=234, + max_prompt_tokens_per_call=1000, + num_compactions=1, + ) + + +def _make_state(**overrides) -> UsageRecorderState: + """Build a minimal UsageRecorderState. record_usage_event ignores + ``dal`` (it's the DAL itself being called), so we leave it None. + Override any field via kwargs.""" + defaults = dict( + dal=None, + request_type="user_chat", + model="m", + provider="p", + is_robusta_model=False, + stats=_stats(), + iterations=1, + ) + defaults.update(overrides) + return UsageRecorderState(**defaults) # type: ignore[arg-type] + + +# ────────────────────────────────────────────────────────────────── +# record_usage_event +# ────────────────────────────────────────────────────────────────── + + +class TestRecordUsageEvent: + def test_no_op_when_dal_disabled(self, mock_dal): + mock_dal.enabled = False + mock_dal.record_usage_event(_make_state()) + # No client interaction at all when disabled. + mock_dal.client.table.assert_not_called() + + def test_inserts_row_with_correct_payload(self, mock_dal): + mock_dal.record_usage_event(_make_state( + request_type="user_chat", + request_source="freeform", + source_ref="issue-42", + conversation_id="conv-abc", + conversation_source="chat_history", + model="anthropic/claude-sonnet-4-5", + provider="anthropic", + is_robusta_model=False, + iterations=3, + tool_call_count=5, + is_streaming=True, + finish_reason="stop", + user_id="user-xyz", + request_id="req-uuid-123", + meta={"experiment_id": "abc"}, + )) + + # client.table().insert().execute() + mock_dal.client.table.assert_called_once_with(HOLMES_USAGE_EVENTS_TABLE) + insert_call = mock_dal.client.table.return_value.insert + insert_call.assert_called_once() + + payload = insert_call.call_args.args[0] + + # Identity + assert payload["account_id"] == mock_dal.account_id + assert payload["cluster_id"] == "test-cluster" + assert payload["user_id"] == "user-xyz" + assert payload["conversation_id"] == "conv-abc" + assert payload["conversation_source"] == "chat_history" + assert payload["request_id"] == "req-uuid-123" + + # Classification + assert payload["request_type"] == "user_chat" + assert payload["request_source"] == "freeform" + assert payload["source_ref"] == "issue-42" + assert payload["status"] == "success" # default on the state + assert payload["model"] == "anthropic/claude-sonnet-4-5" + assert payload["provider"] == "anthropic" + assert payload["is_robusta_model"] is False + + # Stats + assert payload["prompt_tokens"] == 1000 + assert payload["completion_tokens"] == 234 + assert payload["cached_tokens"] == 50 + assert payload["reasoning_tokens"] == 12 + assert payload["total_tokens"] == 1234 + assert payload["total_cost"] == pytest.approx(0.0123) + assert payload["num_compactions"] == 1 + assert payload["iterations"] == 3 + assert payload["max_prompt_tokens_per_call"] == 1000 + assert payload["max_completion_tokens_per_call"] == 234 + + # Outcome + assert payload["tool_call_count"] == 5 + # duration_ms is computed from state.t_start at write time, so we + # don't pin a value — just assert the type contract. + assert isinstance(payload["duration_ms"], int) + assert payload["duration_ms"] >= 0 + assert payload["is_streaming"] is True + assert payload["finish_reason"] == "stop" + assert payload["meta"] == {"experiment_id": "abc"} + + def test_falls_back_to_dal_cluster_when_cluster_id_not_supplied(self, mock_dal): + mock_dal.record_usage_event(_make_state()) # cluster_id left None + payload = mock_dal.client.table.return_value.insert.call_args.args[0] + assert payload["cluster_id"] == "test-cluster" + + def test_explicit_cluster_id_overrides_dal_default(self, mock_dal): + mock_dal.record_usage_event(_make_state(cluster_id="other-cluster")) + payload = mock_dal.client.table.return_value.insert.call_args.args[0] + assert payload["cluster_id"] == "other-cluster" + + def test_meta_defaults_to_empty_dict_when_none(self, mock_dal): + # The state defaults meta to {}, so the column should be {} too. + mock_dal.record_usage_event(_make_state()) + payload = mock_dal.client.table.return_value.insert.call_args.args[0] + assert payload["meta"] == {} + + def test_swallows_supabase_errors(self, mock_dal): + # Supabase client raises — record_usage_event must not bubble up. + mock_dal.client.table.return_value.insert.return_value.execute.side_effect = ( + RuntimeError("supabase down") + ) + # Should not raise. + mock_dal.record_usage_event(_make_state()) + + def test_handles_stats_with_none_cached_tokens(self, mock_dal): + # Some providers don't report cached_tokens — should land as NULL. + stats = RequestStats( + total_cost=0.001, + total_tokens=100, + prompt_tokens=80, + completion_tokens=20, + cached_tokens=None, + reasoning_tokens=0, + ) + mock_dal.record_usage_event(_make_state(stats=stats)) + payload = mock_dal.client.table.return_value.insert.call_args.args[0] + assert payload["cached_tokens"] is None + + def test_handles_state_with_no_stats(self, mock_dal): + # state.stats is None when the request never reached a terminal + # event with cost data (aborted / pre-LLM error). The DAL must + # still write a row with zero/NULL token columns. + mock_dal.record_usage_event(_make_state(stats=None, status="aborted")) + payload = mock_dal.client.table.return_value.insert.call_args.args[0] + assert payload["status"] == "aborted" + assert payload["prompt_tokens"] == 0 + assert payload["completion_tokens"] == 0 + assert payload["total_tokens"] == 0 + assert payload["total_cost"] == 0.0 + assert payload["cached_tokens"] is None + + +# Feedback writes are no longer Holmes' responsibility — the FE calls the +# public.record_feedback() Postgres function directly via supabase.rpc(...). +# That function lives in the migration script and is verified against a real +# Postgres in the integration suite. There is no Holmes-side code path to +# unit-test here. diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py new file mode 100644 index 0000000000..de01df81fe --- /dev/null +++ b/tests/core/test_usage_recorder.py @@ -0,0 +1,578 @@ +"""Unit tests for holmes.core.usage_recorder. + +The recorder is fire-and-forget — it spawns a daemon thread to write the row. +For deterministic tests we patch threading.Thread so the target runs inline, +which lets us assert against the exact UsageRecorderState passed to +dal.record_usage_event. + +Per Moshe's review on PR #1969, ``record_usage_event`` now takes the entire +``UsageRecorderState`` positionally instead of ~20 individual kwargs (drops +the old ``to_kwargs()`` indirection — the DAL is the single place that +knows the column shape). Tests assert on +``state.dal.record_usage_event.call_args.args[0]``, which is the live state +object the recorder passed in. +""" + +from typing import List +from unittest.mock import MagicMock + +import pytest + +from holmes.core.usage_recorder import ( + UsageRecorderState, + record_error, + record_from_llm_result, + stream_with_usage_recording, +) +from holmes.utils.stream import StreamEvents, StreamMessage + + +# ────────────────────────────────────────────────────────────────── +# Helpers +# ────────────────────────────────────────────────────────────────── + + +def _make_state(**overrides) -> UsageRecorderState: + """Build a UsageRecorderState with sensible defaults for tests.""" + base = dict( + dal=MagicMock(enabled=True), + request_type="user_chat", + model="openai/gpt-4", + provider="openai", + is_robusta_model=False, + request_source="freeform", + source_ref=None, + conversation_id="conv-123", + conversation_source="chat_history", + user_id="user-abc", + is_streaming=True, + ) + base.update(overrides) + return UsageRecorderState(**base) + + +def _stream(*events: StreamMessage): + for e in events: + yield e + + +def _terminal_data(costs: dict, num_llm_calls: int = 1, finish_reason: str = "stop") -> dict: + return { + "content": "ok", + "messages": [], + "metadata": {"costs": costs, "finish_reason": finish_reason}, + "num_llm_calls": num_llm_calls, + "costs": costs, + } + + +def _patch_inline_thread(monkeypatch): + """Replace the recorder's ThreadPoolExecutor with an inline stub so + target(state) runs synchronously in the test, letting us assert + against state.dal.record_usage_event right after _fire returns. + + The stub mimics ``ThreadPoolExecutor.submit`` semantics: if the + callable raises, the exception is captured (not propagated) — the + real executor would put it on the returned Future, which production + code never awaits. Without this, the inline path would bubble + target exceptions that the production path never would, breaking + "fire-and-forget never propagates" tests. + + Kept under the historical name (_patch_inline_thread) so existing + test bodies don't churn — the underlying mechanism is now the + bounded executor at module-level rather than a fresh Thread per call. + """ + import holmes.core.usage_recorder as mod + + class _InlineExecutor: + def submit(self, fn, *args, **kwargs): + try: + fn(*args, **kwargs) + except Exception: + # Real executor stashes the exception on the Future. We + # don't return one, so just swallow — the test doesn't + # await results. + pass + + monkeypatch.setattr(mod, "_RECORDER_EXECUTOR", _InlineExecutor()) + + +def _state_arg(state: UsageRecorderState) -> UsageRecorderState: + """Pull the state object out of the recorded dal.record_usage_event call. + + The recorder fires it positionally — args[0]. Centralized so test bodies + don't repeat the indexing ceremony. + """ + return state.dal.record_usage_event.call_args.args[0] + + +# ────────────────────────────────────────────────────────────────── +# UsageRecorderState basics — direct attribute access + duration_ms property +# (Replaces the old TestToKwargs class; to_kwargs() no longer exists.) +# ────────────────────────────────────────────────────────────────── + + +class TestStateBasics: + def test_default_values_match_spec(self): + state = _make_state() + # Identity + assert state.request_type == "user_chat" + assert state.request_source == "freeform" + assert state.conversation_id == "conv-123" + assert state.conversation_source == "chat_history" + assert state.user_id == "user-abc" + assert state.request_id # auto-generated UUID + + # Classification + assert state.model == "openai/gpt-4" + assert state.provider == "openai" + assert state.is_robusta_model is False + assert state.is_streaming is True + + # Mutable defaults — these get filled by the wrapper at runtime + assert state.status == "success" # RequestStatus.SUCCESS == "success" + assert state.iterations == 0 + assert state.tool_call_count == 0 + assert state.finish_reason is None + assert state.meta == {} + assert state.stats is None # not pre-populated + + def test_duration_ms_property_grows_with_time(self): + state = _make_state() + # Force t_start to be in the past so duration_ms > 0 + state.t_start -= 1.0 + assert state.duration_ms >= 1000 + + def test_duration_ms_is_an_int(self): + # The DB column is `int`; the property must always return an int. + state = _make_state() + assert isinstance(state.duration_ms, int) + + def test_is_internal_defaults_to_false(self): + state = _make_state() + assert state.is_internal is False + + def test_is_internal_round_trips(self): + state = _make_state(is_internal=True) + assert state.is_internal is True + + +# ────────────────────────────────────────────────────────────────── +# stream_with_usage_recording +# ────────────────────────────────────────────────────────────────── + + +class TestStreamWithUsageRecording: + def test_success_path_records_with_status_success(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + + costs = {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150, "total_cost": 0.001} + events = [ + StreamMessage(event=StreamEvents.START_TOOL, data={}), + StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "kubectl"}), + StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "prom"}), + StreamMessage(event=StreamEvents.ANSWER_END, data=_terminal_data(costs, num_llm_calls=3)), + ] + + # Drain the wrapped stream — the recorder fires in the finally block. + consumed: List[StreamMessage] = list( + stream_with_usage_recording(_stream(*events), state) + ) + + assert len(consumed) == 4 + state.dal.record_usage_event.assert_called_once() + s = _state_arg(state) + assert s.status == "success" + assert s.tool_call_count == 2 + assert s.iterations == 3 + assert s.finish_reason == "stop" + assert s.stats.prompt_tokens == 100 + assert s.stats.total_tokens == 150 + + def test_error_event_marks_status_error(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + + events = [ + StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "k"}), + StreamMessage(event=StreamEvents.ERROR, data={"metadata": {}}), + ] + list(stream_with_usage_recording(_stream(*events), state)) + + s = _state_arg(state) + assert s.status == "error" + assert s.tool_call_count == 1 + + def test_approval_required_marks_status_approval_required(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + list(stream_with_usage_recording( + _stream(StreamMessage(event=StreamEvents.APPROVAL_REQUIRED, data={"metadata": {}})), + state, + )) + assert _state_arg(state).status == "approval_required" + + def test_exception_in_inner_stream_still_fires_recorder_with_error_status(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + + def failing_stream(): + yield StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "x"}) + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + list(stream_with_usage_recording(failing_stream(), state)) + + # Recorder still fires from the finally + state.dal.record_usage_event.assert_called_once() + s = _state_arg(state) + assert s.status == "error" + # And the tool we saw before the exception was counted + assert s.tool_call_count == 1 + + def test_token_count_event_captures_cumulative_costs(self, monkeypatch): + """TOKEN_COUNT events broadcast the running cost after each + successful LLM iteration. The wrapper must capture them so partial + spend is recorded even when the agentic loop raises mid-loop + (call_stream's local stats accumulator gets GC'd on exception).""" + _patch_inline_thread(monkeypatch) + state = _make_state() + + # Two successful iterations broadcast their cumulative cost... + iter1_costs = {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, "total_cost": 0.001} + iter2_costs = {"prompt_tokens": 250, "completion_tokens": 50, "total_tokens": 300, "total_cost": 0.003} + + events = [ + StreamMessage(event=StreamEvents.TOKEN_COUNT, data={"metadata": {"costs": iter1_costs}}), + StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "k"}), + StreamMessage(event=StreamEvents.TOKEN_COUNT, data={"metadata": {"costs": iter2_costs}}), + StreamMessage(event=StreamEvents.ANSWER_END, data=_terminal_data(iter2_costs, num_llm_calls=2)), + ] + list(stream_with_usage_recording(_stream(*events), state)) + + s = _state_arg(state) + # Last cumulative wins — ANSWER_END's costs override the earlier + # TOKEN_COUNT (in this case they're identical, mirroring real flow). + assert s.stats.total_tokens == 300 + assert s.stats.prompt_tokens == 250 + assert s.stats.total_cost == pytest.approx(0.003) + + def test_partial_costs_captured_when_loop_raises_mid_iteration(self, monkeypatch): + """The whole point of listening to TOKEN_COUNT: when the agentic + loop hits an exception in iteration 3 after iterations 1 and 2 + succeeded, the recorder should still record the cost of iters 1+2. + Without TOKEN_COUNT capture this row would have status=error and + zero tokens, even though real spend already happened.""" + _patch_inline_thread(monkeypatch) + state = _make_state() + + # Cumulative costs after iterations 1 and 2 succeeded. + partial_costs = { + "prompt_tokens": 1500, + "completion_tokens": 200, + "total_tokens": 1700, + "total_cost": 0.012, + } + + def failing_stream(): + # Iteration 1 succeeds. + yield StreamMessage(event=StreamEvents.TOKEN_COUNT, data={"metadata": {"costs": {"prompt_tokens": 800, "total_tokens": 850, "total_cost": 0.005}}}) + yield StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "kubectl"}) + # Iteration 2 succeeds. + yield StreamMessage(event=StreamEvents.TOKEN_COUNT, data={"metadata": {"costs": partial_costs}}) + yield StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "prom"}) + # Iteration 3's LLM call raises — TOKEN_COUNT for iter 3 is + # never emitted; the exception unwinds call_stream's frame. + raise RuntimeError("rate limit hit") + + with pytest.raises(RuntimeError, match="rate limit"): + list(stream_with_usage_recording(failing_stream(), state)) + + s = _state_arg(state) + # Partial spend from iterations 1+2 was captured — NOT zero. + assert s.stats is not None + assert s.stats.total_tokens == 1700 + assert s.stats.prompt_tokens == 1500 + assert s.stats.total_cost == pytest.approx(0.012) + # Status correctly marked as error. + assert s.status == "error" + # Tool calls before the exception were counted. + assert s.tool_call_count == 2 + + def test_partial_costs_captured_on_client_disconnect(self, monkeypatch): + """Client-disconnect case: stream ends cleanly without a terminal + event after some TOKEN_COUNT events. Recorder marks 'aborted' but + still captures the partial spend so the row reflects real cost.""" + _patch_inline_thread(monkeypatch) + state = _make_state() + + partial_costs = { + "prompt_tokens": 600, + "completion_tokens": 100, + "total_tokens": 700, + "total_cost": 0.004, + } + events = [ + StreamMessage(event=StreamEvents.TOKEN_COUNT, data={"metadata": {"costs": partial_costs}}), + StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "k"}), + # Stream just ends — no terminal event (client closed the tab). + ] + list(stream_with_usage_recording(_stream(*events), state)) + + s = _state_arg(state) + # Aborted, but with the real partial cost. + assert s.status == "aborted" + assert s.stats is not None + assert s.stats.total_tokens == 700 + assert s.stats.total_cost == pytest.approx(0.004) + + def test_stream_without_terminal_event_still_records_as_aborted(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + + # No terminal — e.g. client disconnected. + list(stream_with_usage_recording( + _stream(StreamMessage(event=StreamEvents.TOOL_RESULT, data={"tool_name": "y"})), + state, + )) + + # finally block still fired, but status downgraded from default + # "success" to "aborted" because no terminal event was seen. + state.dal.record_usage_event.assert_called_once() + assert _state_arg(state).status == "aborted" + + def test_terminal_event_keeps_its_explicit_status(self, monkeypatch): + """Sanity: the abort downgrade only applies when no terminal was seen.""" + _patch_inline_thread(monkeypatch) + state = _make_state() + + list(stream_with_usage_recording( + _stream(StreamMessage( + event=StreamEvents.ANSWER_END, + data={"metadata": {}, "num_llm_calls": 1}, + )), + state, + )) + + assert _state_arg(state).status == "success" + + def test_request_id_injected_into_answer_end_metadata(self, monkeypatch): + """The FE needs request_id from ai_answer_end so it can post feedback later.""" + _patch_inline_thread(monkeypatch) + state = _make_state() + state.request_id = "deadbeef-1234" + + answer_end = StreamMessage( + event=StreamEvents.ANSWER_END, + data={ + "content": "ok", + "messages": [], + "metadata": {"costs": {"total_tokens": 100}}, + "num_llm_calls": 1, + }, + ) + consumed = list(stream_with_usage_recording(_stream(answer_end), state)) + + # The same StreamMessage flows through; its metadata now has request_id. + out = consumed[0] + assert out.event == StreamEvents.ANSWER_END + assert out.data["metadata"]["request_id"] == "deadbeef-1234" + # Existing metadata content (costs) is preserved alongside. + assert out.data["metadata"]["costs"] == {"total_tokens": 100} + + def test_request_id_injected_when_metadata_missing(self, monkeypatch): + """If the upstream event has no metadata key, _inject_request_id creates it.""" + _patch_inline_thread(monkeypatch) + state = _make_state() + state.request_id = "uuid-xyz" + + answer_end = StreamMessage( + event=StreamEvents.ANSWER_END, + data={"content": "ok", "messages": [], "num_llm_calls": 1}, # no metadata key + ) + consumed = list(stream_with_usage_recording(_stream(answer_end), state)) + assert consumed[0].data["metadata"] == {"request_id": "uuid-xyz"} + + def test_request_id_injected_into_approval_required(self, monkeypatch): + """Feedback should be possible on paused turns too — request_id must be there.""" + _patch_inline_thread(monkeypatch) + state = _make_state() + state.request_id = "rid-paused" + + approval = StreamMessage( + event=StreamEvents.APPROVAL_REQUIRED, + data={"metadata": {}, "pending_approvals": []}, + ) + consumed = list(stream_with_usage_recording(_stream(approval), state)) + assert consumed[0].data["metadata"]["request_id"] == "rid-paused" + + def test_request_id_injected_into_error_event(self, monkeypatch): + """Surface request_id even on ERROR so the FE can report 'this request failed'.""" + _patch_inline_thread(monkeypatch) + state = _make_state() + state.request_id = "rid-err" + + err = StreamMessage(event=StreamEvents.ERROR, data={"metadata": {}}) + consumed = list(stream_with_usage_recording(_stream(err), state)) + assert consumed[0].data["metadata"]["request_id"] == "rid-err" + + +# ────────────────────────────────────────────────────────────────── +# record_from_llm_result (non-streaming) +# ────────────────────────────────────────────────────────────────── + + +class TestRecordFromLlmResult: + def test_extracts_stats_iterations_finish_reason_and_tool_count(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state(is_streaming=False) + + # Build a fake LLMResult-shaped object. RequestStats fields are inherited, + # so we can dump a flat dict via model_dump() in the helper. + fake_result = MagicMock() + fake_result.model_dump.return_value = { + "total_cost": 0.005, + "total_tokens": 250, + "prompt_tokens": 200, + "completion_tokens": 50, + "cached_tokens": None, + "reasoning_tokens": 0, + "max_completion_tokens_per_call": 50, + "max_prompt_tokens_per_call": 200, + "num_compactions": 0, + } + fake_result.num_llm_calls = 4 + fake_result.tool_calls = [object(), object(), object()] + fake_result.finish_reason = "stop" + + record_from_llm_result(state, fake_result) + + s = _state_arg(state) + assert s.status == "success" + assert s.iterations == 4 + assert s.tool_call_count == 3 + assert s.finish_reason == "stop" + assert s.stats.total_tokens == 250 + assert s.stats.total_cost == 0.005 + + def test_handles_missing_attrs_gracefully(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + + # An LLMResult-like object with only the bare minimum + bare = MagicMock() + bare.model_dump.return_value = {} + bare.num_llm_calls = None + bare.tool_calls = None + bare.finish_reason = None + + record_from_llm_result(state, bare) + + s = _state_arg(state) + # iterations falls back to 1 when num_llm_calls is None + assert s.iterations == 1 + assert s.tool_call_count == 0 + + +# ────────────────────────────────────────────────────────────────── +# record_error +# ────────────────────────────────────────────────────────────────── + + +class TestRecordError: + def test_marks_rate_limited_when_rate_limit_in_message(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + record_error(state, RuntimeError("rate limit exceeded for model")) + assert _state_arg(state).status == "rate_limited" + + def test_marks_error_for_other_exceptions(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state() + record_error(state, ValueError("invalid model")) + assert _state_arg(state).status == "error" + + +# ────────────────────────────────────────────────────────────────── +# Disabled-DAL no-op behavior +# ────────────────────────────────────────────────────────────────── + + +class TestDisabledDalNoop: + def test_no_thread_spawned_when_dal_disabled(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state(dal=MagicMock(enabled=False)) + record_error(state, RuntimeError("anything")) + # The disabled DAL has the method mocked but should never be called. + state.dal.record_usage_event.assert_not_called() + + def test_no_thread_spawned_when_dal_is_none(self, monkeypatch): + _patch_inline_thread(monkeypatch) + state = _make_state(dal=None) + # No exception; nothing happens. + record_error(state, RuntimeError("x")) + + +# ────────────────────────────────────────────────────────────────── +# Real executor mode (no inline patching) — verifies fire-and-forget +# ────────────────────────────────────────────────────────────────── + + +class TestFireAndForgetThreadMode: + def test_record_runs_dal_in_background_via_executor(self): + """The shared ThreadPoolExecutor runs record_usage_event off the + caller's thread. The caller returns before the dal write completes.""" + import threading + import time + + called = threading.Event() + + # Real executor → state arrives positionally. Match the + # production call signature. + def slow_record(state): + time.sleep(0.05) + called.set() + + dal = MagicMock(enabled=True) + dal.record_usage_event = slow_record + state = _make_state(dal=dal) + record_error(state, RuntimeError("x")) + # Caller returns immediately; the executor worker is still running. + # Wait briefly for it to finish. + assert called.wait(timeout=2.0), ( + "executor worker did not run record_usage_event in the background" + ) + + def test_dal_exception_does_not_propagate(self, monkeypatch): + """If record_usage_event raises, the caller of _fire must not see it. + In production the real ThreadPoolExecutor parks the exception on the + returned Future (which we never await), so nothing bubbles. The + inline stub mimics that swallow semantic for deterministic testing.""" + _patch_inline_thread(monkeypatch) + + dal = MagicMock(enabled=True) + dal.record_usage_event.side_effect = RuntimeError("supabase down") + state = _make_state(dal=dal) + + # Should NOT raise — fire-and-forget contract. + record_error(state, RuntimeError("x")) + + def test_executor_shutdown_is_silently_handled(self, monkeypatch): + """If the executor was shut down (process exiting), submit raises + RuntimeError. _fire catches that and accepts the loss — same fate + as in-flight rows on the previous daemon-thread shape during + process shutdown.""" + import holmes.core.usage_recorder as mod + + class _ShutdownExecutor: + def submit(self, fn, *args, **kwargs): + raise RuntimeError("cannot schedule new futures after shutdown") + + monkeypatch.setattr(mod, "_RECORDER_EXECUTOR", _ShutdownExecutor()) + + dal = MagicMock(enabled=True) + state = _make_state(dal=dal) + + # Should NOT raise — _fire's `except RuntimeError` accepts the loss. + record_error(state, RuntimeError("x")) diff --git a/tests/test_chat_recorder_state.py b/tests/test_chat_recorder_state.py new file mode 100644 index 0000000000..760076e2bc --- /dev/null +++ b/tests/test_chat_recorder_state.py @@ -0,0 +1,233 @@ +"""Unit tests for usage_recorder.build_chat_recorder_state. + +The helper is shared by every entry point that consumes a ChatRequest and +wraps a stream — direct /api/chat (server.py) and the worker path +(conversations_worker/worker.py). These tests cover the derivation logic +without touching either entry point so they remain stable while the +calling code evolves. + +Coverage: +- is_internal derivation (explicit True/False vs internal_-prefix fallback) +- Slack-prefix auto-detection on `ask` (request_type / request_source / + meta.slack) +- conversation_source default ('chat_history' when conversation_id set + but unset by caller) +- Smoke: model / provider / user_id / source_ref / meta passthrough +""" + +from unittest.mock import MagicMock + +from holmes.core.usage_recorder import build_chat_recorder_state + + +def _make_request_ai(model="openai/gpt-4", is_robusta=False): + ai = MagicMock() + ai.llm = MagicMock() + ai.llm.model = model + ai.llm.is_robusta_model = is_robusta + return ai + + +def _chat_request(**overrides): + """Build a minimal ChatRequest. Imports lazily so server-import side + effects don't hit collection-time.""" + from holmes.core.models import ChatRequest + + base = dict(ask="test question", stream=False) + base.update(overrides) + return ChatRequest(**base) + + +def _dal(): + """A stand-in dal — build_chat_recorder_state only stores it on the + state, doesn't call it. We never fire the recorder in these tests.""" + return MagicMock() + + +class TestIsInternalDerivation: + def test_explicit_true_wins(self): + req = _chat_request(is_internal=True, request_source="freeform") + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.is_internal is True + + def test_explicit_false_wins_even_with_internal_prefix(self): + # FE may have a "freeform" request labeled with an internal_-prefixed + # request_source for some reason; the explicit False should still win. + req = _chat_request( + is_internal=False, request_source="internal_legacy_user_chat" + ) + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.is_internal is False + + def test_unset_falls_back_to_internal_prefix(self): + # Backwards-compat: existing FE clients use the internal_ prefix + # convention without setting is_internal explicitly. + req = _chat_request(request_source="internal_title_generation") + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.is_internal is True + + def test_unset_with_no_prefix_yields_false(self): + req = _chat_request(request_source="freeform") + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.is_internal is False + + def test_unset_with_no_request_source_yields_false(self): + # No FE labeling at all → not internal. + req = _chat_request() + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.is_internal is False + + +class TestRecorderStateSmoke: + """Catch obvious wiring regressions in build_chat_recorder_state.""" + + def test_carries_through_basic_fields(self): + req = _chat_request( + user_id="u-abc", + conversation_id="conv-123", + request_source="alert_investigation", + source_ref="issue-42", + meta={"experiment_id": "x"}, + ) + state = build_chat_recorder_state( + req, + _make_request_ai(model="anthropic/claude-sonnet-4-5"), + dal=_dal(), + is_streaming=True, + ) + + assert state.request_type == "user_chat" # default + assert state.request_source == "alert_investigation" + assert state.source_ref == "issue-42" + assert state.conversation_id == "conv-123" + # Default for direct /api/chat: chat_history when conversation_id is set. + assert state.conversation_source == "chat_history" + assert state.user_id == "u-abc" + assert state.is_streaming is True + assert state.is_internal is False + assert state.model == "anthropic/claude-sonnet-4-5" + assert state.meta == {"experiment_id": "x"} + + def test_explicit_conversation_source_wins(self): + # Worker passes conversation_source='conversations' explicitly when + # constructing the ChatRequest; the helper must not clobber it with + # its 'chat_history' default. + req = _chat_request( + conversation_id="conv-123", conversation_source="conversations" + ) + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=True + ) + assert state.conversation_source == "conversations" + + def test_no_conversation_id_means_no_conversation_source(self): + # CLI / scheduled-prompt-style requests have no conversation; the + # discriminator must stay None so dashboards know not to join either + # table. + req = _chat_request() + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.conversation_id is None + assert state.conversation_source is None + + +# Sample of the prefix the Robusta runner's Slack handler prepends to `ask`. +SLACK_ASK = ( + "**@user_U0AKMP2CZ97** • 2026-05-04T05:10:04Z\n\nhigh cpu in pod alert" +) + + +class TestSlackAutoDetect: + def test_slack_prefix_sets_request_type_to_slack_chat(self): + req = _chat_request(ask=SLACK_ASK) + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=True + ) + assert state.request_type == "slack_chat" + + def test_slack_prefix_sets_request_source_to_slack(self): + req = _chat_request(ask=SLACK_ASK) + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=True + ) + assert state.request_source == "slack" + + def test_slack_prefix_captures_user_id_and_ts_in_meta(self): + req = _chat_request(ask=SLACK_ASK) + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=True + ) + assert state.meta.get("slack") == { + "slack_user_id": "U0AKMP2CZ97", + "slack_triggered_at": "2026-05-04T05:10:04Z", + } + + def test_explicit_request_type_wins_over_slack_detection(self): + # Even with the Slack-shaped prefix, an explicit request_type must win + # (e.g. a future caller that overrides for some reason). + req = _chat_request(ask=SLACK_ASK, request_type="user_chat") + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=True + ) + assert state.request_type == "user_chat" + # Slack metadata is still extracted — we don't drop the signal just + # because the type was overridden. + assert state.meta.get("slack", {}).get("slack_user_id") == "U0AKMP2CZ97" + + def test_explicit_request_source_wins_over_slack_default(self): + # Same caller-wins semantic for request_source: if the runner ever + # ships finer values like 'slack_mention' / 'slack_alert_investigation', + # those should not be clobbered by the auto-detected default. + req = _chat_request(ask=SLACK_ASK, request_source="slack_mention") + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=True + ) + assert state.request_source == "slack_mention" + # request_type still auto-set since it wasn't explicitly provided. + assert state.request_type == "slack_chat" + + def test_no_slack_prefix_uses_default_request_type_and_no_source(self): + req = _chat_request(ask="why is my-service crashing?") + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.request_type == "user_chat" + assert state.request_source is None + assert "slack" not in state.meta + + def test_slack_meta_merges_with_fe_meta(self): + req = _chat_request(ask=SLACK_ASK, meta={"experiment_id": "abc"}) + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=True + ) + # Both keys preserved; backend doesn't clobber FE meta. + assert state.meta == { + "experiment_id": "abc", + "slack": { + "slack_user_id": "U0AKMP2CZ97", + "slack_triggered_at": "2026-05-04T05:10:04Z", + }, + } + + def test_partial_slack_prefix_does_not_match(self): + # Just a markdown bold, no • or timestamp — must not falsely match. + req = _chat_request( + ask="**@user_U0AKMP2CZ97** asked: why is my pod down?" + ) + state = build_chat_recorder_state( + req, _make_request_ai(), dal=_dal(), is_streaming=False + ) + assert state.request_type == "user_chat" + assert state.request_source is None + assert "slack" not in state.meta