From 0cbf22caa6b17a18079418fc80b9f9a7c2851392 Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 29 Apr 2026 15:21:08 +0300 Subject: [PATCH 01/21] Add AI usage tracking via HolmesUsageEvents (account/cluster/user/feature/cost) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires every LLM-consuming entry point through a shared usage recorder so each request lands as one row in the new HolmesUsageEvents Supabase table. Enables per-account / per-cluster / per-user / per-feature cost reporting, plus thumbs-up/down feedback on individual chat turns. Why: per requirements doc (Hypatos / Adobe Firefly / Playtika), customers need to see "cost per alert investigation", "questions asked this week", "who used Holmes from Teams vs UI", "thumbs-up rate by feature" — none of which are answerable today because Holmes persists no AI metrics. What's covered: POST /api/chat (direct + worker path), scheduled prompts, POST /api/agui/chat, POST /api/checks/execute. CLI flows are intentionally out of scope (no FastAPI server -> no recorder hook). How it works: - New holmes/core/usage_recorder.py with UsageRecorderState dataclass + stream_with_usage_recording (streaming) + record_from_llm_result (non-streaming) + record_error. Every entry point builds one UsageRecorderState and either wraps its stream or calls record_from_llm_result. The recorder fires fire-and-forget on a daemon thread so telemetry never blocks the response. - SupabaseDal.record_usage_event / record_feedback write through the existing self.client (same email+password user-JWT pattern as upsert_holmes_status). No new env vars, no service_role key. RLS mirrors HolmesStatus's policy. - ChatRequest gains 6 optional fields (request_type, request_source, source_ref, conversation_id, conversation_source, meta). All additive / backwards-compatible — old clients keep working. - LLMResult gains a finish_reason field, captured from the last LLM iteration's response.choices[0].finish_reason. - POST /api/feedback endpoint accepts request_id + sentiment + optional category/comment, UPDATEs the matching event row. Schema (separately applied to Supabase): HolmesUsageEvents covers account_id (uuid), cluster_id, user_id (uuid), conversation_id + conversation_source discriminator, request_type, request_source, source_ref, model, provider, is_robusta_model, full RequestStats fields, iterations, tool_call_count, duration_ms, is_streaming, finish_reason, feedback_* columns, and a meta JSONB for forward-compat. Optional HolmesUsageEventCalls per-LLM-call child table is gated by HOLMES_RECORD_LLM_CALLS=true (not enabled in v1). Tests: - tests/core/test_usage_recorder.py — 15 tests covering streaming / non-streaming / error / approval paths, disabled-DAL no-op, and fire-and-forget thread mode. - tests/core/test_supabase_dal_usage.py — 11 tests covering payload shape, account-scoped UPDATE, NULL handling, error swallowing. - tests/test_tool_calling_llm.py — all 41 existing tests still pass (one regression I caught and fixed: finish_reason capture now guards with isinstance(str) so MagicMock'd test responses don't break pydantic validation). Plan: ~/.claude/plans/does-holmes-save-any-sunny-globe.md Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- experimental/ag-ui/server-agui.py | 41 ++- holmes/checks/checks.py | 16 + holmes/checks/checks_api.py | 33 ++ holmes/core/conversations_worker/worker.py | 10 + holmes/core/models.py | 53 ++++ holmes/core/scheduled_prompts/executor.py | 4 + holmes/core/supabase_dal.py | 111 +++++++ holmes/core/tool_calling_llm.py | 15 + holmes/core/usage_recorder.py | 217 +++++++++++++ server.py | 111 ++++++- tests/core/test_supabase_dal_usage.py | 333 ++++++++++++++++++++ tests/core/test_usage_recorder.py | 342 +++++++++++++++++++++ 12 files changed, 1282 insertions(+), 4 deletions(-) create mode 100644 holmes/core/usage_recorder.py create mode 100644 tests/core/test_supabase_dal_usage.py create mode 100644 tests/core/test_usage_recorder.py diff --git a/experimental/ag-ui/server-agui.py b/experimental/ag-ui/server-agui.py index 269b04415f..5db0cd7494 100644 --- a/experimental/ag-ui/server-agui.py +++ b/experimental/ag-ui/server-agui.py @@ -24,6 +24,10 @@ from starlette.responses import PlainTextResponse from holmes.utils.stream import StreamMessage, StreamEvents +from holmes.core.usage_recorder import ( + UsageRecorderState, + stream_with_usage_recording, +) from holmes.common.env_vars import ( HOLMES_HOST, HOLMES_PORT, @@ -139,9 +143,40 @@ 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") + except Exception: + pass + ai_model = getattr(ai.llm, "model", None) or chat_request.model or "unknown" + try: + import litellm as _litellm # local import: this file is opt-in + ai_provider = _litellm.get_llm_provider(ai_model)[1] or "unknown" + except Exception: + ai_provider = ai_model.split("/")[0] if "/" in ai_model else "unknown" + 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..274c26a578 100644 --- a/holmes/checks/checks_api.py +++ b/holmes/checks/checks_api.py @@ -13,8 +13,21 @@ 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 from holmes.plugins.destinations.slack.plugin import SlackDestination + +def _resolve_provider(model: Optional[str]) -> str: + """Best-effort: return the canonical litellm provider for `model`.""" + if not model: + return "unknown" + try: + import litellm + + return litellm.get_llm_provider(model)[1] or "unknown" + except Exception: + return model.split("/")[0] if "/" in model else "unknown" + checks_app = FastAPI() _CONFIG: Config @@ -110,12 +123,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/worker.py b/holmes/core/conversations_worker/worker.py index 574af8cb66..d3429f2147 100644 --- a/holmes/core/conversations_worker/worker.py +++ b/holmes/core/conversations_worker/worker.py @@ -468,6 +468,16 @@ 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"), + # AI usage tracking (HolmesUsageEvents). user_id / request_source / + # source_ref / meta come from the FE-supplied user_message blob; + # conversation_id matches the Conversations row the worker is processing. + user_id=data.get("user_id"), + request_type="user_chat", + request_source=data.get("request_source"), + source_ref=data.get("source_ref"), + conversation_id=task.conversation_id, + conversation_source="conversations", + meta=data.get("meta"), ) self._run_chat_and_publish( diff --git a/holmes/core/models.py b/holmes/core/models.py index 17feb21f33..682e751179 100644 --- a/holmes/core/models.py +++ b/holmes/core/models.py @@ -213,6 +213,59 @@ 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." + ), + ) + # 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 b7f02b2ff0..bad28c4e1e 100644 --- a/holmes/core/scheduled_prompts/executor.py +++ b/holmes/core/scheduled_prompts/executor.py @@ -198,6 +198,10 @@ def _execute_prompt( stream=False, additional_system_prompt=additional_system_prompt, trace_span=heartbeat_span, + # 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 c2cfc8bccb..a73cac2074 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -64,6 +64,8 @@ CONVERSATIONS_TABLE = "Conversations" CONVERSATION_EVENTS_TABLE = "ConversationEvents" OAUTH_TOKENS_TABLE = "OAuthTokens" +HOLMES_USAGE_EVENTS_TABLE = "HolmesUsageEvents" +HOLMES_USAGE_EVENT_CALLS_TABLE = "HolmesUsageEventCalls" ENRICHMENT_BLACKLIST = ["text_file", "graph", "ai_analysis", "holmes"] ENRICHMENT_BLACKLIST_SET = set(ENRICHMENT_BLACKLIST) @@ -817,6 +819,115 @@ 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, + *, + request_type: str, + request_source: Optional[str], + source_ref: Optional[str], + conversation_id: Optional[str], + conversation_source: Optional[str], + status: str, + model: str, + provider: str, + is_robusta_model: bool, + stats, # holmes.core.llm_usage.RequestStats; not typed to avoid circular import + iterations: int, + duration_ms: Optional[int], + tool_call_count: int, + is_streaming: bool, + finish_reason: Optional[str], + user_id: Optional[str], + cluster_id: Optional[str] = None, + request_id: Optional[str] = None, + meta: Optional[Dict] = None, + ) -> None: + """Record one HolmesUsageEvents row. Best-effort: swallows DB errors. + + Called from holmes.core.usage_recorder._fire on a daemon thread, so + errors here only affect the telemetry row, never the request response. + """ + if not self.enabled: + return + try: + self.client.table(HOLMES_USAGE_EVENTS_TABLE).insert({ + "account_id": self.account_id, + "cluster_id": cluster_id or self.cluster, + "user_id": user_id, + "conversation_id": conversation_id, + "conversation_source": conversation_source, + "request_id": request_id, + "request_type": request_type, + "request_source": request_source, + "source_ref": source_ref, + "status": status, + "model": model, + "provider": provider, + "is_robusta_model": is_robusta_model, + "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": 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": tool_call_count, + "duration_ms": duration_ms, + "is_streaming": is_streaming, + "finish_reason": finish_reason, + "meta": meta or {}, + }).execute() + except Exception: + logging.exception("Failed to record usage event") + + def record_feedback( + self, + *, + request_id: str, + sentiment: str, + category: Optional[str], + comment: Optional[str], + user_id: Optional[str], + ) -> None: + """Update the feedback_* columns on a HolmesUsageEvents row. + + Best-effort. The UPDATE is account-scoped via WHERE so RLS and an + explicit predicate both prevent cross-account writes if request_id + ever collides. + + ``user_id`` is accepted for symmetry / future use (e.g. a + feedback_user_id column when chats become shareable) but is not + currently written, since v1's auth model assumes the rater is the + asker (events.user_id == feedback giver). + """ + if not self.enabled: + return + if sentiment not in ("thumbs_up", "thumbs_down"): + logging.warning( + "record_feedback: ignoring invalid sentiment %r", sentiment + ) + return + try: + self.client.table(HOLMES_USAGE_EVENTS_TABLE).update( + { + "feedback_sentiment": sentiment, + "feedback_category": category, + "feedback_comment": comment, + "feedback_at": datetime.now().isoformat(), + } + ).eq("account_id", self.account_id).eq( + "request_id", request_id + ).execute() + except Exception: + logging.exception("Failed to record feedback") + 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 ae7a98399b..94255041d3 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): @@ -577,6 +578,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(), ) @@ -598,6 +600,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(), ) @@ -1154,6 +1157,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..2a5bcf195f --- /dev/null +++ b/holmes/core/usage_recorder.py @@ -0,0 +1,217 @@ +"""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 threading +import time +import uuid +from dataclasses import dataclass, field +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.supabase_dal import SupabaseDal + from holmes.core.tool_calling_llm import LLMResult + + +@dataclass +class UsageRecorderState: + """All the data needed to write one HolmesUsageEvents row. + + Identity / classification fields are set by the entry point at + construction time. Mutable runtime fields (`stats`, `iterations`, + `tool_call_count`, `finish_reason`, `status`) are filled by the + stream wrapper or by `record_from_llm_result` before firing. + """ + + # required identity / classification — set by the entry point + dal: Any # SupabaseDal; typed Any to avoid circular import at runtime + request_type: str + model: str + provider: str + is_robusta_model: bool + + # optional identity / classification + request_source: Optional[str] = None + source_ref: Optional[str] = None + conversation_id: Optional[str] = None + conversation_source: Optional[str] = None + user_id: Optional[str] = None + cluster_id: Optional[str] = None + request_id: str = field(default_factory=lambda: str(uuid.uuid4())) + is_streaming: bool = False + meta: Dict[str, Any] = field(default_factory=dict) + + # mutable — filled during the call by the wrapper or recorder + t_start: float = field(default_factory=time.monotonic) + stats: Optional[RequestStats] = None + iterations: int = 0 + tool_call_count: int = 0 + finish_reason: Optional[str] = None + status: str = "success" + + def to_kwargs(self) -> Dict[str, Any]: + """Pack the state into the kwargs `SupabaseDal.record_usage_event` expects.""" + return { + "request_type": self.request_type, + "request_source": self.request_source, + "source_ref": self.source_ref, + "conversation_id": self.conversation_id, + "conversation_source": self.conversation_source, + "status": self.status, + "model": self.model, + "provider": self.provider, + "is_robusta_model": self.is_robusta_model, + "stats": self.stats or RequestStats(), + "iterations": self.iterations, + "duration_ms": int((time.monotonic() - self.t_start) * 1000), + "tool_call_count": self.tool_call_count, + "is_streaming": self.is_streaming, + "finish_reason": self.finish_reason, + "user_id": self.user_id, + "cluster_id": self.cluster_id, + "request_id": self.request_id, + "meta": self.meta, + } + + +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. + """ + saw_terminal = False + try: + for msg in stream: + if msg.event == StreamEvents.TOOL_RESULT: + state.tool_call_count += 1 + elif msg.event == StreamEvents.ANSWER_END: + _capture_terminal(state, msg.data) + state.status = "success" + saw_terminal = True + elif msg.event == StreamEvents.APPROVAL_REQUIRED: + _capture_terminal(state, msg.data) + state.status = "approval_required" + saw_terminal = True + elif msg.event == StreamEvents.ERROR: + _capture_terminal(state, msg.data) + state.status = "error" + saw_terminal = True + yield msg + except Exception: + if not saw_terminal: + state.status = "error" + raise + finally: + _fire(state) + + +def _capture_terminal(state: UsageRecorderState, data: Dict[str, Any]) -> None: + """Pull cost/iterations/finish_reason from a terminal event's data.""" + metadata = data.get("metadata") or {} + costs = metadata.get("costs") or {} + if costs: + try: + state.stats = RequestStats(**costs) + except Exception: + logging.debug( + "Failed to materialize RequestStats from terminal event costs", + exc_info=True, + ) + state.iterations = data.get("num_llm_calls", state.iterations) or state.iterations + state.finish_reason = ( + metadata.get("finish_reason") or state.finish_reason + ) + + +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, so dump the stats fields out. + stats_fields = { + k: v + for k, v in llm_result.model_dump().items() + if k + in { + "total_cost", + "total_tokens", + "prompt_tokens", + "completion_tokens", + "cached_tokens", + "reasoning_tokens", + "max_completion_tokens_per_call", + "max_prompt_tokens_per_call", + "num_compactions", + } + } + state.stats = RequestStats(**stats_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 = "success" + _fire(state) + + +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 = "rate_limited" + else: + state.status = "error" + _fire(state) + + +def _fire(state: UsageRecorderState) -> None: + """Background-thread the dal write so the response path never blocks.""" + if state.dal is None or not getattr(state.dal, "enabled", False): + return + try: + threading.Thread( + target=state.dal.record_usage_event, + kwargs=state.to_kwargs(), + daemon=True, + name="usage-recorder", + ).start() + except Exception: + # Defense in depth — record_usage_event has its own try/except too. + logging.exception("Failed to spawn usage recorder thread") + + +__all__ = [ + "UsageRecorderState", + "stream_with_usage_recording", + "record_from_llm_result", + "record_error", +] diff --git a/server.py b/server.py index 75b4536b22..be70fee39e 100644 --- a/server.py +++ b/server.py @@ -27,6 +27,7 @@ from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse from litellm.exceptions import AuthenticationError +from pydantic import BaseModel from holmes import get_version, is_official_release from holmes.common.env_vars import ( @@ -67,9 +68,64 @@ from holmes.core.models import FrontendToolMode from holmes.core.tools_utils.frontend_tools import build_frontend_noop_tool, build_frontend_pause_tool from holmes.core.tracing import TracingFactory +from holmes.core.usage_recorder import ( + UsageRecorderState, + record_error, + record_from_llm_result, + stream_with_usage_recording, +) from holmes.utils.stream import stream_chat_formatter +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). + """ + if not model: + return "unknown" + try: + # litellm.get_llm_provider returns (model, provider, api_key, api_base). + 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, + is_streaming: bool, +) -> UsageRecorderState: + """Construct a UsageRecorderState for /api/chat. All entry points (direct, + worker, scheduled, agui) flow through chat() — request_type, request_source, + source_ref, conversation_id, conversation_source, user_id are read off + chat_request (set by whoever called chat()). + """ + # Default conversation_source to 'chat_history' when conversation_id is set + # but the caller didn't override (i.e. direct /api/chat). + 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" + + return UsageRecorderState( + dal=dal, + request_type=chat_request.request_type or "user_chat", + request_source=chat_request.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, + model=model_name, + provider=_resolve_provider(model_name), + is_robusta_model=getattr(request_ai.llm, "is_robusta_model", False), + meta=dict(chat_request.meta or {}), + ) + + def init_logging(): # Filter out periodical healniss and readiness probe. uvicorn_logger = logging.getLogger("uvicorn.access") @@ -504,7 +560,10 @@ 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, 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, @@ -514,6 +573,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, ) @@ -522,6 +585,7 @@ def chat(chat_request: ChatRequest, http_request: Request): media_type="text/event-stream", ) else: + recorder_state = _build_chat_recorder_state(chat_request, request_ai, is_streaming=False) try: # Use provided trace_span or create a root investigation span trace_span = chat_request.trace_span @@ -541,6 +605,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: @@ -567,6 +634,12 @@ def chat(chat_request: ChatRequest, http_request: Request): metadata=llm_call.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() @@ -580,6 +653,42 @@ def chat(chat_request: ChatRequest, http_request: Request): raise HTTPException(status_code=500, detail=str(e)) +class FeedbackRequest(BaseModel): + """Body for POST /api/feedback. Defined inline to avoid bloating models.py + with an endpoint-specific schema.""" + request_id: str + sentiment: str # 'thumbs_up' | 'thumbs_down' + category: Optional[str] = None + comment: Optional[str] = None + + +@app.post("/api/feedback") +def feedback(req: FeedbackRequest, http_request: Request) -> dict: + """Record user feedback (thumbs up/down + optional category/comment) on + a previously-completed chat call. Looks up the row by request_id and + UPDATEs feedback_* columns. Best-effort: returns 200 even if the row + isn't found yet (rare network reorder).""" + if req.sentiment not in ("thumbs_up", "thumbs_down"): + raise HTTPException(status_code=400, detail="invalid sentiment") + user_id: Optional[str] = None + # Mirror /api/chat's user_id resolution path: prefer header passthrough + # (Robusta relay token) but accept a body field if the caller set one. + try: + body_user_id = http_request.query_params.get("user_id") + if body_user_id: + user_id = body_user_id + except Exception: + pass + dal.record_feedback( + request_id=req.request_id, + sentiment=req.sentiment, + category=req.category, + comment=req.comment, + user_id=user_id, + ) + return {"ok": True} + + scheduled_prompts_executor = ScheduledPromptsExecutor( dal=dal, config=config, chat_function=chat ) diff --git a/tests/core/test_supabase_dal_usage.py b/tests/core/test_supabase_dal_usage.py new file mode 100644 index 0000000000..bb6b76cd7d --- /dev/null +++ b/tests/core/test_supabase_dal_usage.py @@ -0,0 +1,333 @@ +"""Unit tests for SupabaseDal.record_usage_event and SupabaseDal.record_feedback. + +These methods are best-effort: they swallow 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 / .update) and the failure path +(exceptions are absorbed). +""" + +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, +) + + +@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, + ) + + +# ────────────────────────────────────────────────────────────────── +# record_usage_event +# ────────────────────────────────────────────────────────────────── + + +class TestRecordUsageEvent: + def test_no_op_when_dal_disabled(self, mock_dal): + mock_dal.enabled = False + mock_dal.record_usage_event( + request_type="user_chat", + request_source=None, + source_ref=None, + conversation_id=None, + conversation_source=None, + status="success", + model="openai/gpt-4", + provider="openai", + is_robusta_model=False, + stats=_stats(), + iterations=1, + duration_ms=42, + tool_call_count=0, + is_streaming=False, + finish_reason=None, + user_id=None, + ) + # 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( + request_type="user_chat", + request_source="freeform", + source_ref="issue-42", + conversation_id="conv-abc", + conversation_source="chat_history", + status="success", + model="anthropic/claude-sonnet-4-5", + provider="anthropic", + is_robusta_model=False, + stats=_stats(), + iterations=3, + duration_ms=1500, + 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" + 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 + assert payload["duration_ms"] == 1500 + 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( + request_type="user_chat", + request_source=None, + source_ref=None, + conversation_id=None, + conversation_source=None, + status="success", + model="m", + provider="p", + is_robusta_model=False, + stats=_stats(), + iterations=1, + duration_ms=10, + tool_call_count=0, + is_streaming=False, + finish_reason=None, + user_id=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( + request_type="user_chat", + request_source=None, + source_ref=None, + conversation_id=None, + conversation_source=None, + status="success", + model="m", + provider="p", + is_robusta_model=False, + stats=_stats(), + iterations=1, + duration_ms=10, + tool_call_count=0, + is_streaming=False, + finish_reason=None, + user_id=None, + 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): + mock_dal.record_usage_event( + request_type="user_chat", + request_source=None, + source_ref=None, + conversation_id=None, + conversation_source=None, + status="success", + model="m", + provider="p", + is_robusta_model=False, + stats=_stats(), + iterations=1, + duration_ms=10, + tool_call_count=0, + is_streaming=False, + finish_reason=None, + user_id=None, + meta=None, + ) + 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( + request_type="user_chat", + request_source=None, + source_ref=None, + conversation_id=None, + conversation_source=None, + status="success", + model="m", + provider="p", + is_robusta_model=False, + stats=_stats(), + iterations=1, + duration_ms=10, + tool_call_count=0, + is_streaming=False, + finish_reason=None, + user_id=None, + ) + + 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( + request_type="user_chat", + request_source=None, + source_ref=None, + conversation_id=None, + conversation_source=None, + status="success", + model="m", + provider="p", + is_robusta_model=False, + stats=stats, + iterations=1, + duration_ms=10, + tool_call_count=0, + is_streaming=False, + finish_reason=None, + user_id=None, + ) + payload = mock_dal.client.table.return_value.insert.call_args.args[0] + assert payload["cached_tokens"] is None + + +# ────────────────────────────────────────────────────────────────── +# record_feedback +# ────────────────────────────────────────────────────────────────── + + +class TestRecordFeedback: + def test_no_op_when_dal_disabled(self, mock_dal): + mock_dal.enabled = False + mock_dal.record_feedback( + request_id="r-1", + sentiment="thumbs_up", + category=None, + comment=None, + user_id=None, + ) + mock_dal.client.table.assert_not_called() + + def test_invalid_sentiment_logs_and_returns(self, mock_dal): + # Should not raise, should not call .update. + mock_dal.record_feedback( + request_id="r-1", + sentiment="meh", # not allowed + category=None, + comment=None, + user_id=None, + ) + mock_dal.client.table.assert_not_called() + + def test_updates_correct_columns_with_account_and_request_id_filter(self, mock_dal): + mock_dal.record_feedback( + request_id="req-uuid-123", + sentiment="thumbs_down", + category="wrong_answer", + comment="missed the OOM", + user_id="user-xyz", + ) + + # Chain: client.table(...).update({...}).eq("account_id", ...).eq("request_id", ...).execute() + mock_dal.client.table.assert_called_once_with(HOLMES_USAGE_EVENTS_TABLE) + update_mock = mock_dal.client.table.return_value.update + update_mock.assert_called_once() + update_payload = update_mock.call_args.args[0] + assert update_payload["feedback_sentiment"] == "thumbs_down" + assert update_payload["feedback_category"] == "wrong_answer" + assert update_payload["feedback_comment"] == "missed the OOM" + assert "feedback_at" in update_payload # an ISO timestamp string + + # Account-scoped + request_id WHERE clauses + eq_calls = update_mock.return_value.eq.call_args_list + # First .eq is account_id, second is request_id (chained on the return value) + assert eq_calls[0].args == ("account_id", mock_dal.account_id) + # The second .eq is on the result of the first, but with our MagicMock chain + # it's also captured here. In any case, both account_id and request_id + # must be in the predicate chain. + first_eq_return = update_mock.return_value.eq.return_value + first_eq_return.eq.assert_called_with("request_id", "req-uuid-123") + + def test_swallows_supabase_errors(self, mock_dal): + mock_dal.client.table.return_value.update.return_value.eq.return_value.eq.return_value.execute.side_effect = ( + RuntimeError("update failed") + ) + # Should not raise. + mock_dal.record_feedback( + request_id="r-1", + sentiment="thumbs_up", + category=None, + comment=None, + user_id=None, + ) diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py new file mode 100644 index 0000000000..b64fc6c5b0 --- /dev/null +++ b/tests/core/test_usage_recorder.py @@ -0,0 +1,342 @@ +"""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 kwargs passed to dal.record_usage_event. +""" + +from typing import List +from unittest.mock import MagicMock, patch + +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 threading.Thread inside usage_recorder so target() runs inline.""" + import holmes.core.usage_recorder as mod + + class _InlineThread: + def __init__(self, target=None, kwargs=None, daemon=None, name=None): + self._target = target + self._kwargs = kwargs or {} + + def start(self): + self._target(**self._kwargs) + + monkeypatch.setattr(mod.threading, "Thread", _InlineThread) + + +# ────────────────────────────────────────────────────────────────── +# UsageRecorderState.to_kwargs +# ────────────────────────────────────────────────────────────────── + + +class TestToKwargs: + def test_packs_all_required_fields(self): + state = _make_state() + kwargs = state.to_kwargs() + + # Identity + assert kwargs["request_type"] == "user_chat" + assert kwargs["request_source"] == "freeform" + assert kwargs["conversation_id"] == "conv-123" + assert kwargs["conversation_source"] == "chat_history" + assert kwargs["user_id"] == "user-abc" + assert "request_id" in kwargs and kwargs["request_id"] + + # Classification + assert kwargs["model"] == "openai/gpt-4" + assert kwargs["provider"] == "openai" + assert kwargs["is_robusta_model"] is False + assert kwargs["is_streaming"] is True + + # Mutable defaults + assert kwargs["status"] == "success" + assert kwargs["iterations"] == 0 + assert kwargs["tool_call_count"] == 0 + assert kwargs["finish_reason"] is None + assert kwargs["meta"] == {} + + # Stats default to an empty RequestStats, not None + assert kwargs["stats"] is not None + assert kwargs["stats"].total_tokens == 0 + assert kwargs["stats"].total_cost == 0.0 + + def test_duration_ms_is_computed_from_t_start(self): + state = _make_state() + # Force t_start to be in the past so duration_ms > 0 + state.t_start -= 1.0 + kwargs = state.to_kwargs() + assert kwargs["duration_ms"] >= 1000 + + +# ────────────────────────────────────────────────────────────────── +# 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() + kw = state.dal.record_usage_event.call_args.kwargs + assert kw["status"] == "success" + assert kw["tool_call_count"] == 2 + assert kw["iterations"] == 3 + assert kw["finish_reason"] == "stop" + assert kw["stats"].prompt_tokens == 100 + assert kw["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)) + + kw = state.dal.record_usage_event.call_args.kwargs + assert kw["status"] == "error" + assert kw["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.dal.record_usage_event.call_args.kwargs["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() + assert state.dal.record_usage_event.call_args.kwargs["status"] == "error" + # And the tool we saw before the exception was counted + assert state.dal.record_usage_event.call_args.kwargs["tool_call_count"] == 1 + + def test_stream_without_terminal_event_still_records(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 + state.dal.record_usage_event.assert_called_once() + + +# ────────────────────────────────────────────────────────────────── +# 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) + + kw = state.dal.record_usage_event.call_args.kwargs + assert kw["status"] == "success" + assert kw["iterations"] == 4 + assert kw["tool_call_count"] == 3 + assert kw["finish_reason"] == "stop" + assert kw["stats"].total_tokens == 250 + assert kw["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) + + kw = state.dal.record_usage_event.call_args.kwargs + # iterations falls back to 1 when num_llm_calls is None + assert kw["iterations"] == 1 + assert kw["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")) + kw = state.dal.record_usage_event.call_args.kwargs + assert kw["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")) + kw = state.dal.record_usage_event.call_args.kwargs + assert kw["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 Thread mode (no inline patching) — verifies fire-and-forget +# ────────────────────────────────────────────────────────────────── + + +class TestFireAndForgetThreadMode: + def test_record_calls_dal_in_background_thread(self): + import threading + import time + + called = threading.Event() + + def slow_record(**kwargs): + 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 thread is still running. + # Wait briefly for it to finish. + assert called.wait(timeout=2.0), "background thread did not run record_usage_event" + + def test_dal_exception_does_not_propagate(self, monkeypatch): + # _fire wraps Thread.start() in a try/except so even with inline-thread + # patching (where target runs synchronously in start()) downstream + # exceptions don't bubble out to the caller. Logged via + # logging.exception(), but the caller is unaffected. This is the + # defense-in-depth contract — telemetry must never break the response. + _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. The inner try/except in _fire swallows it. + record_error(state, RuntimeError("x")) From c9341381d4999914c9bf5f7de610691880134c37 Mon Sep 17 00:00:00 2001 From: alonelish Date: Thu, 30 Apr 2026 05:51:31 +0300 Subject: [PATCH 02/21] Surface request_id in /api/chat response metadata for feedback The FE needs the chat-call's request_id to send POST /api/feedback {request_id, sentiment} when the user clicks thumbs up/down. Previously the recorder generated request_id internally but never sent it back. Streaming path: stream_with_usage_recording now injects state.request_id into the terminal event's metadata dict before yielding it. The SSE formatter then ships it as part of metadata in ai_answer_end (and approval_required / error). Same treatment for all three terminals so feedback works on paused / errored chats too. Non-streaming path: chat() copies llm_call.metadata, drops in recorder_state.request_id, and uses the merged dict as ChatResponse.metadata. Same shape on the wire as the streaming case. Tests: 4 new cases in test_usage_recorder.py covering injection into ANSWER_END / APPROVAL_REQUIRED / ERROR plus the missing-metadata case where _inject_request_id has to create the dict. All 71 tests (usage_recorder + supabase_dal_usage + tool_calling_llm) pass. FE contract: read response.metadata.request_id from ai_answer_end (or the JSON body for non-stream); save it; POST it back to /api/feedback. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/usage_recorder.py | 19 ++++++++++ server.py | 7 +++- tests/core/test_usage_recorder.py | 60 +++++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 1 deletion(-) diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 2a5bcf195f..639354805b 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -97,6 +97,11 @@ def stream_with_usage_recording( 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 posts it back to + ``POST /api/feedback`` when the user clicks thumbs up/down. """ saw_terminal = False try: @@ -105,14 +110,17 @@ def stream_with_usage_recording( state.tool_call_count += 1 elif msg.event == StreamEvents.ANSWER_END: _capture_terminal(state, msg.data) + _inject_request_id(msg.data, state.request_id) state.status = "success" saw_terminal = True elif msg.event == StreamEvents.APPROVAL_REQUIRED: _capture_terminal(state, msg.data) + _inject_request_id(msg.data, state.request_id) state.status = "approval_required" saw_terminal = True elif msg.event == StreamEvents.ERROR: _capture_terminal(state, msg.data) + _inject_request_id(msg.data, state.request_id) state.status = "error" saw_terminal = True yield msg @@ -124,6 +132,17 @@ def stream_with_usage_recording( _fire(state) +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. + """ + md = data.get("metadata") + if not isinstance(md, dict): + md = {} + data["metadata"] = md + md["request_id"] = request_id + + def _capture_terminal(state: UsageRecorderState, data: Dict[str, Any]) -> None: """Pull cost/iterations/finish_reason from a terminal event's data.""" metadata = data.get("metadata") or {} diff --git a/server.py b/server.py index be70fee39e..7176aca14e 100644 --- a/server.py +++ b/server.py @@ -626,12 +626,17 @@ 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 POST /api/feedback 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: diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py index b64fc6c5b0..f7c4f9f76b 100644 --- a/tests/core/test_usage_recorder.py +++ b/tests/core/test_usage_recorder.py @@ -203,6 +203,66 @@ def test_stream_without_terminal_event_still_records(self, monkeypatch): # finally block still fired state.dal.record_usage_event.assert_called_once() + 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) From 6c80253355ce7cb272b20726ea4da198dc35c5c8 Mon Sep 17 00:00:00 2001 From: alonelish Date: Thu, 30 Apr 2026 07:02:30 +0300 Subject: [PATCH 03/21] Add is_internal flag to HolmesUsageEvents to filter internal calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Server-internal calls (title generation, classification, summarization, follow-up suggestion ranking, etc.) currently use an 'internal_' prefix on request_source as a convention. That works but forces every dashboard to do `request_source NOT LIKE 'internal\_%'` and is easy to forget. Adds a dedicated boolean is_internal column. Dashboards default-filter with `WHERE NOT is_internal`. Index-friendly (low-cardinality boolean), self-documenting in schema, removes string-pattern coupling. Backwards-compatible: when chat_request.is_internal is unset, the server falls back to detecting the legacy 'internal_' prefix on request_source. Existing FE clients keep working unchanged. New FE code can set is_internal=true explicitly and stop relying on naming. Wiring: - ChatRequest gains optional is_internal: bool field - UsageRecorderState carries it; to_kwargs includes it - SupabaseDal.record_usage_event accepts and writes it - server._build_chat_recorder_state derives it (explicit > prefix > false) - ConversationWorker forwards is_internal from task.user_message_data - Scheduled prompts and health checks NOT marked internal — their output is user-facing (alert routing, scheduled-prompt results), so request_type='scheduled_prompt'/'health_check' is the right discriminator for those, not is_internal. Tests: - 2 new in test_usage_recorder.py for default-false / round-trip - New test_chat_recorder_state.py with 6 tests covering all 4 cases of the prefix-fallback derivation (explicit true / explicit false with prefix / unset with prefix / unset without prefix / unset without request_source) plus a smoke test on the rest of the wiring DB-side: ALTER TABLE HolmesUsageEvents ADD COLUMN is_internal boolean NOT NULL DEFAULT false (run separately by Alon). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/conversations_worker/worker.py | 6 +- holmes/core/models.py | 10 ++ holmes/core/supabase_dal.py | 2 + holmes/core/usage_recorder.py | 2 + server.py | 13 +++ tests/core/test_usage_recorder.py | 9 ++ tests/test_chat_recorder_state.py | 107 +++++++++++++++++++++ 7 files changed, 147 insertions(+), 2 deletions(-) create mode 100644 tests/test_chat_recorder_state.py diff --git a/holmes/core/conversations_worker/worker.py b/holmes/core/conversations_worker/worker.py index d3429f2147..564476f001 100644 --- a/holmes/core/conversations_worker/worker.py +++ b/holmes/core/conversations_worker/worker.py @@ -469,8 +469,9 @@ def _process_conversation(self, task: ConversationTask) -> None: response_format=data.get("response_format"), behavior_controls=data.get("behavior_controls"), # AI usage tracking (HolmesUsageEvents). user_id / request_source / - # source_ref / meta come from the FE-supplied user_message blob; - # conversation_id matches the Conversations row the worker is processing. + # source_ref / meta / is_internal come from the FE-supplied + # user_message blob; conversation_id matches the Conversations row + # the worker is processing. user_id=data.get("user_id"), request_type="user_chat", request_source=data.get("request_source"), @@ -478,6 +479,7 @@ def _process_conversation(self, task: ConversationTask) -> None: conversation_id=task.conversation_id, conversation_source="conversations", meta=data.get("meta"), + is_internal=data.get("is_internal"), ) self._run_chat_and_publish( diff --git a/holmes/core/models.py b/holmes/core/models.py index 682e751179..47e393048c 100644 --- a/holmes/core/models.py +++ b/holmes/core/models.py @@ -265,6 +265,16 @@ class ChatRequestBaseModel(BaseModel): "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": ...}], diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index a73cac2074..f3e321d66f 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -841,6 +841,7 @@ def record_usage_event( cluster_id: Optional[str] = None, request_id: Optional[str] = None, meta: Optional[Dict] = None, + is_internal: bool = False, ) -> None: """Record one HolmesUsageEvents row. Best-effort: swallows DB errors. @@ -881,6 +882,7 @@ def record_usage_event( "tool_call_count": tool_call_count, "duration_ms": duration_ms, "is_streaming": is_streaming, + "is_internal": is_internal, "finish_reason": finish_reason, "meta": meta or {}, }).execute() diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 639354805b..bf43acd343 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -51,6 +51,7 @@ class UsageRecorderState: cluster_id: Optional[str] = None request_id: str = field(default_factory=lambda: str(uuid.uuid4())) is_streaming: bool = False + is_internal: bool = False meta: Dict[str, Any] = field(default_factory=dict) # mutable — filled during the call by the wrapper or recorder @@ -78,6 +79,7 @@ def to_kwargs(self) -> Dict[str, Any]: "duration_ms": int((time.monotonic() - self.t_start) * 1000), "tool_call_count": self.tool_call_count, "is_streaming": self.is_streaming, + "is_internal": self.is_internal, "finish_reason": self.finish_reason, "user_id": self.user_id, "cluster_id": self.cluster_id, diff --git a/server.py b/server.py index 7176aca14e..91de62d704 100644 --- a/server.py +++ b/server.py @@ -110,6 +110,18 @@ def _build_chat_recorder_state( 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) + return UsageRecorderState( dal=dal, request_type=chat_request.request_type or "user_chat", @@ -119,6 +131,7 @@ def _build_chat_recorder_state( 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), diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py index f7c4f9f76b..1ab19d9579 100644 --- a/tests/core/test_usage_recorder.py +++ b/tests/core/test_usage_recorder.py @@ -116,6 +116,15 @@ def test_duration_ms_is_computed_from_t_start(self): kwargs = state.to_kwargs() assert kwargs["duration_ms"] >= 1000 + def test_is_internal_defaults_to_false(self): + state = _make_state() + assert state.is_internal is False + assert state.to_kwargs()["is_internal"] is False + + def test_is_internal_true_round_trips(self): + state = _make_state(is_internal=True) + assert state.to_kwargs()["is_internal"] is True + # ────────────────────────────────────────────────────────────────── # stream_with_usage_recording diff --git a/tests/test_chat_recorder_state.py b/tests/test_chat_recorder_state.py new file mode 100644 index 0000000000..8644c06276 --- /dev/null +++ b/tests/test_chat_recorder_state.py @@ -0,0 +1,107 @@ +"""Unit tests for server._build_chat_recorder_state. + +Specifically covers the is_internal derivation logic: +- explicit chat_request.is_internal=True wins +- explicit chat_request.is_internal=False wins (even with internal_ prefix) +- is_internal=None falls back to detecting the legacy 'internal_' prefix + on request_source +- is_internal=None with no internal_ prefix yields False + +Plus a smoke check that the rest of the recorder state is built correctly. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +# Importing server is heavy (initializes dal/config), so do it lazily inside a +# fixture. The test still exercises the real function. + + +@pytest.fixture(scope="module") +def build_chat_recorder_state(): + # Patch the module-level dal in server.py before importing so init doesn't + # try to authenticate against a real Supabase URL. + with patch("holmes.core.supabase_dal.create_client"): + from server import _build_chat_recorder_state + return _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 happen only once via the fixture.""" + from holmes.core.models import ChatRequest + + base = dict(ask="test question", stream=False) + base.update(overrides) + return ChatRequest(**base) + + +class TestIsInternalDerivation: + def test_explicit_true_wins(self, build_chat_recorder_state): + req = _chat_request(is_internal=True, request_source="freeform") + state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=False) + assert state.is_internal is True + + def test_explicit_false_wins_even_with_internal_prefix( + self, build_chat_recorder_state + ): + # 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(), is_streaming=False) + assert state.is_internal is False + + def test_unset_falls_back_to_internal_prefix(self, build_chat_recorder_state): + # 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(), is_streaming=False) + assert state.is_internal is True + + def test_unset_with_no_prefix_yields_false(self, build_chat_recorder_state): + req = _chat_request(request_source="freeform") + state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=False) + assert state.is_internal is False + + def test_unset_with_no_request_source_yields_false(self, build_chat_recorder_state): + # No FE labeling at all → not internal. + req = _chat_request() + state = build_chat_recorder_state(req, _make_request_ai(), 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, build_chat_recorder_state): + 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"), 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"} From 456765690978f480a51287921a10ff77a3194b11 Mon Sep 17 00:00:00 2001 From: alonelish Date: Thu, 30 Apr 2026 08:05:23 +0300 Subject: [PATCH 04/21] Address CodeRabbit feedback on AI usage tracking PR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 5 fixes from CodeRabbit review (3 reviews on commits 0cbf22ca, c9341381, 6c802533): - holmes/checks/checks_api.py: move `import litellm` to module scope per CLAUDE.md import rule (was inline inside _resolve_provider). - holmes/core/supabase_dal.py: scope record_feedback UPDATE by user_id when provided. Defense in depth — even though v1's auth model assumes rater == asker, adding `.eq("user_id", user_id)` when present prevents cross-user overwrites if a request_id ever leaks. Skip the filter when user_id is None to keep the path open for future system / scheduled flows. - server.py /api/feedback: drop the unnecessary try/except around http_request.query_params.get("user_id") (.get() doesn't raise), rename body_user_id → user_id directly, simplify to a single expression. - experimental/ag-ui/server-agui.py: narrow the user_id extraction except to AttributeError/TypeError and log at debug instead of silent pass. Adds an isinstance(ctx, dict) guard since ctx is opaque per the AG-UI protocol. - holmes/core/usage_recorder.py _capture_terminal: replace `data.get("num_llm_calls", state.iterations) or state.iterations` with an explicit None check so a legitimate 0 isn't silently overwritten. Edge case (the agentic loop always makes >=1 call) but cheap to fix. Not addressed (intentional): - __all__ in usage_recorder.py is not alphabetically sorted (RUF022). Current order is intentional (public class first, then helpers grouped by usage frequency: streaming most common, then non-streaming, then error path). CodeRabbit acknowledged this is "low value" / cosmetic. All 79 tests still pass. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- experimental/ag-ui/server-agui.py | 17 ++++++++++++++--- holmes/checks/checks_api.py | 3 +-- holmes/core/supabase_dal.py | 16 ++++++++++++---- holmes/core/usage_recorder.py | 6 +++++- server.py | 13 ++++--------- 5 files changed, 36 insertions(+), 19 deletions(-) diff --git a/experimental/ag-ui/server-agui.py b/experimental/ag-ui/server-agui.py index 5db0cd7494..5da8ca10ad 100644 --- a/experimental/ag-ui/server-agui.py +++ b/experimental/ag-ui/server-agui.py @@ -149,9 +149,20 @@ async def event_generator(message_history): ctx = input_data.context or {} agui_user_id = None try: - agui_user_id = getattr(input_data, "user_id", None) or ctx.get("user_id") - except Exception: - pass + 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" try: import litellm as _litellm # local import: this file is opt-in diff --git a/holmes/checks/checks_api.py b/holmes/checks/checks_api.py index 274c26a578..48867f4af0 100644 --- a/holmes/checks/checks_api.py +++ b/holmes/checks/checks_api.py @@ -3,6 +3,7 @@ import time from typing import Optional +import litellm from fastapi import FastAPI, HTTPException from litellm.exceptions import AuthenticationError from pydantic import BaseModel, Field @@ -22,8 +23,6 @@ def _resolve_provider(model: Optional[str]) -> str: if not model: return "unknown" try: - import litellm - return litellm.get_llm_provider(model)[1] or "unknown" except Exception: return model.split("/")[0] if "/" in model else "unknown" diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index f3e321d66f..3c25c376fe 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -917,16 +917,24 @@ def record_feedback( ) return try: - self.client.table(HOLMES_USAGE_EVENTS_TABLE).update( + # Defense in depth: when user_id is provided, scope the UPDATE so + # one user can't overwrite another user's feedback within the same + # account if a request_id ever leaks. v1's auth model assumes + # rater == asker, so the user_id sent on POST /api/feedback + # matches HolmesUsageEvents.user_id for normal flows. Skip the + # filter when user_id is None (e.g. system / scheduled flows that + # might emit feedback in the future) to avoid breaking those. + query = self.client.table(HOLMES_USAGE_EVENTS_TABLE).update( { "feedback_sentiment": sentiment, "feedback_category": category, "feedback_comment": comment, "feedback_at": datetime.now().isoformat(), } - ).eq("account_id", self.account_id).eq( - "request_id", request_id - ).execute() + ).eq("account_id", self.account_id).eq("request_id", request_id) + if user_id: + query = query.eq("user_id", user_id) + query.execute() except Exception: logging.exception("Failed to record feedback") diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index bf43acd343..3eb3ba3e08 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -157,7 +157,11 @@ def _capture_terminal(state: UsageRecorderState, data: Dict[str, Any]) -> None: "Failed to materialize RequestStats from terminal event costs", exc_info=True, ) - state.iterations = data.get("num_llm_calls", state.iterations) or state.iterations + # Explicit None-check rather than `or` so a legitimate 0 (unlikely but + # not impossible) is preserved instead of falling back to state.iterations. + raw_iterations = data.get("num_llm_calls") + if raw_iterations is not None: + state.iterations = raw_iterations state.finish_reason = ( metadata.get("finish_reason") or state.finish_reason ) diff --git a/server.py b/server.py index 91de62d704..b92990f514 100644 --- a/server.py +++ b/server.py @@ -688,15 +688,10 @@ def feedback(req: FeedbackRequest, http_request: Request) -> dict: isn't found yet (rare network reorder).""" if req.sentiment not in ("thumbs_up", "thumbs_down"): raise HTTPException(status_code=400, detail="invalid sentiment") - user_id: Optional[str] = None - # Mirror /api/chat's user_id resolution path: prefer header passthrough - # (Robusta relay token) but accept a body field if the caller set one. - try: - body_user_id = http_request.query_params.get("user_id") - if body_user_id: - user_id = body_user_id - except Exception: - pass + # Mirror /api/chat's user_id resolution: pull from the request's query + # params (e.g. when posted by the Robusta relay). `.get()` returns None + # for missing keys without raising, so no try/except needed. + user_id: Optional[str] = http_request.query_params.get("user_id") dal.record_feedback( request_id=req.request_id, sentiment=req.sentiment, From 0f8bd384d4fa19bacd8751e1a487e4f0a76e231f Mon Sep 17 00:00:00 2001 From: alonelish Date: Sun, 3 May 2026 16:52:18 +0300 Subject: [PATCH 05/21] Address second-round CodeRabbit feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 3 fixes: - usage_recorder: stream_with_usage_recording now downgrades status from default 'success' to 'aborted' in the finally block when no terminal event was seen (e.g. client disconnect mid-stream). Without this the row was recorded as success even though the chat never completed. CodeRabbit catch. - supabase_dal: record_feedback docstring was stale ("user_id is accepted but not currently written") — updated to describe the defense-in-depth user_id scoping that was added in 45676569. - server.py /api/feedback: require user_id (return 400 if missing). The DAL keeps `if user_id:` guarding for future system/scheduled feedback flows, but the user-facing endpoint must always have it to preserve the cross-user defense. Tests: - test_stream_without_terminal_event_still_records renamed to ..._as_aborted and now asserts the downgraded status - New test_terminal_event_keeps_its_explicit_status sanity-checks that the downgrade doesn't fire when a terminal event was seen Not addressed (intentional): - usage_recorder thread-pool refactor (CodeRabbit "🏗️ Heavy lift"). Per-event daemon thread is fine at Holmes' expected volume. Reasoning posted in PR reply: at peak ~100 chats/min and a 60s Supabase outage, peak in-flight threads is ~100 (~800KB total). Daemon-thread overhead is negligible. A bounded ThreadPoolExecutor becomes worth it if sustained traffic ever exceeds ~1000/min OR Supabase outages become routine — neither is the case for v1. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/supabase_dal.py | 12 ++++++++---- holmes/core/usage_recorder.py | 7 +++++++ server.py | 7 ++++++- tests/core/test_usage_recorder.py | 21 +++++++++++++++++++-- 4 files changed, 40 insertions(+), 7 deletions(-) diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index 3c25c376fe..5c69931cd7 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -904,10 +904,14 @@ def record_feedback( explicit predicate both prevent cross-account writes if request_id ever collides. - ``user_id`` is accepted for symmetry / future use (e.g. a - feedback_user_id column when chats become shareable) but is not - currently written, since v1's auth model assumes the rater is the - asker (events.user_id == feedback giver). + ``user_id`` is optional. When provided, the UPDATE is additionally + scoped by ``.eq("user_id", user_id)`` as defense in depth — even + though v1's auth model assumes rater == asker (so events.user_id + already matches the user posting the thumb), the extra predicate + prevents cross-user overwrites if a request_id ever leaks. When + omitted, no user filter is applied, leaving room for future + system / scheduled flows that might emit feedback without a + specific human user. """ if not self.enabled: return diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 3eb3ba3e08..fdd5761589 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -131,6 +131,13 @@ def stream_with_usage_recording( state.status = "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 == "success": + state.status = "aborted" _fire(state) diff --git a/server.py b/server.py index b72c020ee4..b2e6e002d6 100644 --- a/server.py +++ b/server.py @@ -706,8 +706,13 @@ def feedback(req: FeedbackRequest, http_request: Request) -> dict: raise HTTPException(status_code=400, detail="invalid sentiment") # Mirror /api/chat's user_id resolution: pull from the request's query # params (e.g. when posted by the Robusta relay). `.get()` returns None - # for missing keys without raising, so no try/except needed. + # for missing keys without raising. Require it so the user-scoping in + # record_feedback's WHERE clause is always applied — without user_id + # the UPDATE would fall back to (account_id, request_id) only, which + # weakens the cross-user defense. user_id: Optional[str] = http_request.query_params.get("user_id") + if not user_id: + raise HTTPException(status_code=400, detail="missing user_id") dal.record_feedback( request_id=req.request_id, sentiment=req.sentiment, diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py index 1ab19d9579..b6caf3a640 100644 --- a/tests/core/test_usage_recorder.py +++ b/tests/core/test_usage_recorder.py @@ -199,7 +199,7 @@ def failing_stream(): # And the tool we saw before the exception was counted assert state.dal.record_usage_event.call_args.kwargs["tool_call_count"] == 1 - def test_stream_without_terminal_event_still_records(self, monkeypatch): + def test_stream_without_terminal_event_still_records_as_aborted(self, monkeypatch): _patch_inline_thread(monkeypatch) state = _make_state() @@ -209,8 +209,25 @@ def test_stream_without_terminal_event_still_records(self, monkeypatch): state, )) - # finally block still fired + # 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.dal.record_usage_event.call_args.kwargs["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.dal.record_usage_event.call_args.kwargs["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.""" From f14450be9a0a79f5bb724657ba0bc1c8d83481e8 Mon Sep 17 00:00:00 2001 From: alonelish Date: Mon, 4 May 2026 09:26:22 +0300 Subject: [PATCH 06/21] Auto-detect Slack-driven /api/chat calls and tag as request_type='slack_chat' MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Robusta runner's Slack handler currently doesn't populate request_type / user_id / conversation_id structurally on /api/chat — instead it prepends a fixed marker to the `ask` text before sending: "**@user_U0AKMP2CZ97** • 2026-05-04T05:10:04Z\n\nhigh cpu in pod alert" Without runner changes, those rows land in HolmesUsageEvents with NULL request_type / user_id and look identical to a generic web client. That prevents dashboards from grouping by client surface. This change adds a heuristic in _build_chat_recorder_state that: 1. Detects the prefix via regex (matches '**@user_U** • '). 2. When matched and the caller didn't explicitly set request_type, tags the row as request_type='slack_chat' (matches the existing _chat suffix convention used by user_chat / agui_chat). 3. Captures the parsed slack_user_id and slack_triggered_at into meta.slack so dashboards can drill in (per-Slack-user cost, etc.) without waiting for the runner to send them as structured fields. Caller-supplied request_type still wins, so scheduled prompts / agui / checks remain unaffected. Slack metadata extraction also runs even when request_type was overridden, so the meta.slack signal is preserved regardless. Caveats documented inline: - Heuristic / fragile if the runner format ever changes. - Doesn't recover conversation_id (Slack thread_ts isn't in the ask text), so multi-turn Slack threads still won't group until the runner sends it. - Doesn't fix the SaaS-side timeout that's causing 5006 errors mid-stream. Long-term proper fix is the runner sending request_type / user_id / conversation_id explicitly. This heuristic is a tactical workaround until that ships. Tests: - 6 new cases in tests/test_chat_recorder_state.py covering: prefix sets request_type; user_id/ts captured in meta; explicit request_type wins over detection; no prefix uses default; meta merges with FE-supplied meta; partial / non-matching prefix doesn't false-positive. 86 tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- server.py | 62 +++++++++++++++++++++++++++++-- tests/test_chat_recorder_state.py | 58 +++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 3 deletions(-) diff --git a/server.py b/server.py index b2e6e002d6..78ae687add 100644 --- a/server.py +++ b/server.py @@ -11,11 +11,12 @@ # IMPORTING ABOVE MIGHT INITIALIZE AN HTTPS CLIENT THAT DOESN'T TRUST THE CUSTOM CERTIFICATE import json import logging +import re import threading import time from datetime import datetime from pathlib import Path -from typing import List, Optional +from typing import Any, Dict, List, Optional import colorlog import litellm @@ -95,6 +96,41 @@ def _resolve_provider(model: Optional[str]) -> str: return model.split("/")[0] if "/" in model else "unknown" +# 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" +# This regex matches that prefix so we can tag Slack-driven rows in +# HolmesUsageEvents without requiring the runner to populate request_type +# explicitly. Heuristic — fragile if the runner format changes; the proper +# long-term fix is the runner sending request_type/user_id/conversation_id +# itself. Captured groups also let us stash the slack user id and timestamp +# in `meta.slack` for dashboards, since they aren't (yet) sent as structured +# fields. +_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. + + Used as a heuristic in _build_chat_recorder_state to auto-classify + Slack-driven calls until the runner is updated to send request_type + and user_id explicitly. + """ + 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 _build_chat_recorder_state( chat_request: ChatRequest, request_ai, @@ -125,9 +161,29 @@ def _build_chat_recorder_state( else: is_internal = bool(chat_request.is_internal) + # Slack auto-detection: the Robusta runner's Slack handler doesn't yet + # populate request_type / user_id / conversation_id structurally — instead + # it prepends a '**@user_X** • \n\n' marker to `ask`. Detect that and + # tag the row as request_type='slack_chat' so dashboards group Slack + # traffic separately. Caller's explicit request_type still wins. + 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" + + # 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=chat_request.request_type or "user_chat", + request_type=request_type, request_source=chat_request.request_source, source_ref=chat_request.source_ref, conversation_id=chat_request.conversation_id, @@ -138,7 +194,7 @@ def _build_chat_recorder_state( model=model_name, provider=_resolve_provider(model_name), is_robusta_model=getattr(request_ai.llm, "is_robusta_model", False), - meta=dict(chat_request.meta or {}), + meta=merged_meta, ) diff --git a/tests/test_chat_recorder_state.py b/tests/test_chat_recorder_state.py index 8644c06276..c4c3ed7ceb 100644 --- a/tests/test_chat_recorder_state.py +++ b/tests/test_chat_recorder_state.py @@ -105,3 +105,61 @@ def test_carries_through_basic_fields(self, build_chat_recorder_state): assert state.is_internal is False assert state.model == "anthropic/claude-sonnet-4-5" assert state.meta == {"experiment_id": "x"} + + +# 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, build_chat_recorder_state): + req = _chat_request(ask=SLACK_ASK) + state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=True) + assert state.request_type == "slack_chat" + + def test_slack_prefix_captures_user_id_and_ts_in_meta(self, build_chat_recorder_state): + req = _chat_request(ask=SLACK_ASK) + state = build_chat_recorder_state(req, _make_request_ai(), 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, build_chat_recorder_state + ): + # 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(), 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_no_slack_prefix_uses_default_request_type(self, build_chat_recorder_state): + req = _chat_request(ask="why is my-service crashing?") + state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=False) + assert state.request_type == "user_chat" + assert "slack" not in state.meta + + def test_slack_meta_merges_with_fe_meta(self, build_chat_recorder_state): + req = _chat_request(ask=SLACK_ASK, meta={"experiment_id": "abc"}) + state = build_chat_recorder_state(req, _make_request_ai(), 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, build_chat_recorder_state): + # 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(), is_streaming=False) + assert state.request_type == "user_chat" + assert "slack" not in state.meta From 6c7474ce555ced1a0eb4342878e8f0f4634e62e3 Mon Sep 17 00:00:00 2001 From: alonelish Date: Mon, 4 May 2026 10:29:01 +0300 Subject: [PATCH 07/21] Also auto-set request_source='slack' when Slack prefix is detected MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to f14450be: the previous commit only set request_type='slack_chat' when the Slack-runner prefix was detected, leaving request_source NULL. That worked for the FE's request_type-based client grouping but left an asymmetry — other request_source-driven dashboards saw NULL for every Slack row. Now request_source defaults to 'slack' when auto-detection fires AND the caller didn't explicitly send a value. Caller-supplied request_source still wins, so when the runner eventually ships finer values like 'slack_mention' / 'slack_alert_investigation' those override the default with no code change here. Tests: - New test_slack_prefix_sets_request_source_to_slack - New test_explicit_request_source_wins_over_slack_default (caller-wins) - Updated test_no_slack_prefix_... and test_partial_slack_prefix_... to also assert request_source is None when the prefix doesn't match 47 tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- server.py | 18 +++++++++++++----- tests/test_chat_recorder_state.py | 23 ++++++++++++++++++++++- 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/server.py b/server.py index 78ae687add..9f7bf25381 100644 --- a/server.py +++ b/server.py @@ -162,10 +162,14 @@ def _build_chat_recorder_state( is_internal = bool(chat_request.is_internal) # Slack auto-detection: the Robusta runner's Slack handler doesn't yet - # populate request_type / user_id / conversation_id structurally — instead - # it prepends a '**@user_X** • \n\n' marker to `ask`. Detect that and - # tag the row as request_type='slack_chat' so dashboards group Slack - # traffic separately. Caller's explicit request_type still wins. + # populate request_type / request_source / user_id / conversation_id + # structurally — instead it prepends a '**@user_X** • \n\n' marker to + # `ask`. Detect that and tag both request_type='slack_chat' (matches the + # _chat suffix of user_chat / agui_chat) and request_source='slack' + # (default UI flow label) so dashboards group Slack traffic on either + # axis. Caller-supplied values for either field still win — the runner + # can later override with finer values like 'slack_mention' / + # 'slack_alert_investigation' without code changes here. slack_info = _detect_slack_origin(chat_request.ask) if chat_request.request_type: request_type = chat_request.request_type @@ -174,6 +178,10 @@ def _build_chat_recorder_state( 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. @@ -184,7 +192,7 @@ def _build_chat_recorder_state( return UsageRecorderState( dal=dal, request_type=request_type, - request_source=chat_request.request_source, + request_source=request_source, source_ref=chat_request.source_ref, conversation_id=chat_request.conversation_id, conversation_source=conversation_source, diff --git a/tests/test_chat_recorder_state.py b/tests/test_chat_recorder_state.py index c4c3ed7ceb..d8d6bc4c2d 100644 --- a/tests/test_chat_recorder_state.py +++ b/tests/test_chat_recorder_state.py @@ -119,6 +119,11 @@ def test_slack_prefix_sets_request_type_to_slack_chat(self, build_chat_recorder_ state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=True) assert state.request_type == "slack_chat" + def test_slack_prefix_sets_request_source_to_slack(self, build_chat_recorder_state): + req = _chat_request(ask=SLACK_ASK) + state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=True) + assert state.request_source == "slack" + def test_slack_prefix_captures_user_id_and_ts_in_meta(self, build_chat_recorder_state): req = _chat_request(ask=SLACK_ASK) state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=True) @@ -139,10 +144,25 @@ def test_explicit_request_type_wins_over_slack_detection( # because the type was overridden. assert state.meta.get("slack", {}).get("slack_user_id") == "U0AKMP2CZ97" - def test_no_slack_prefix_uses_default_request_type(self, build_chat_recorder_state): + def test_explicit_request_source_wins_over_slack_default( + self, build_chat_recorder_state + ): + # 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(), 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, build_chat_recorder_state + ): req = _chat_request(ask="why is my-service crashing?") state = build_chat_recorder_state(req, _make_request_ai(), 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, build_chat_recorder_state): @@ -162,4 +182,5 @@ def test_partial_slack_prefix_does_not_match(self, build_chat_recorder_state): req = _chat_request(ask="**@user_U0AKMP2CZ97** asked: why is my pod down?") state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=False) assert state.request_type == "user_chat" + assert state.request_source is None assert "slack" not in state.meta From b7969f979cbe71b8f94fc8224e69fdda15275d89 Mon Sep 17 00:00:00 2001 From: alonelish Date: Mon, 4 May 2026 16:03:58 +0300 Subject: [PATCH 08/21] Make user_id optional on POST /api/feedback Reverts the 400-on-missing-user_id check from 0f8bd384. The check was added in response to a CodeRabbit suggestion about cross-user tampering defense, but the simpler model is fine for v1: - account_id-scoped RLS already prevents cross-account writes. - (account_id, request_id) in the UPDATE WHERE clause prevents another account's request_id from matching. - The remaining "another user in the same account spoofs your thumb" attack vector is low-impact (worst case: skews the attacker's own account's analytics). When user_id IS supplied as a query param, the DAL still adds .eq("user_id", user_id) as defense-in-depth (record_feedback unchanged). Just no longer required. Endpoint now succeeds with or without ?user_id=... Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- server.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/server.py b/server.py index 9f7bf25381..05f0b7f933 100644 --- a/server.py +++ b/server.py @@ -768,15 +768,11 @@ def feedback(req: FeedbackRequest, http_request: Request) -> dict: isn't found yet (rare network reorder).""" if req.sentiment not in ("thumbs_up", "thumbs_down"): raise HTTPException(status_code=400, detail="invalid sentiment") - # Mirror /api/chat's user_id resolution: pull from the request's query - # params (e.g. when posted by the Robusta relay). `.get()` returns None - # for missing keys without raising. Require it so the user-scoping in - # record_feedback's WHERE clause is always applied — without user_id - # the UPDATE would fall back to (account_id, request_id) only, which - # weakens the cross-user defense. + # user_id is optional — when supplied via query param (mirrors /api/chat's + # resolution path), the DAL adds it as defense-in-depth on the UPDATE. + # When absent, account_id RLS + (account_id, request_id) WHERE is enough + # to prevent cross-account writes. user_id: Optional[str] = http_request.query_params.get("user_id") - if not user_id: - raise HTTPException(status_code=400, detail="missing user_id") dal.record_feedback( request_id=req.request_id, sentiment=req.sentiment, From df209d5f0643fef333be0cb61cdcf47932b30f8f Mon Sep 17 00:00:00 2001 From: alonelish Date: Tue, 5 May 2026 13:21:25 +0300 Subject: [PATCH 09/21] Wire usage recorder into ConversationWorker so worker chats are tracked MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The worker's _run_chat_and_publish path bypassed server.py::chat() and called request_ai.call_stream() directly, so the recorder wrapper never fired. Result: every worker-driven (new conversation API) turn was silently missing from HolmesUsageEvents — only title-generation / classifier sub-calls that go through server.py::chat() were tracked. Fix: - Promote build_chat_recorder_state, detect_slack_origin, and resolve_provider from server.py-private helpers to public usage_recorder.py exports, parameterized on dal so any caller (server, worker, scheduled prompts, agui, checks) can use them. - Wrap the worker's call_stream() output with stream_with_usage_recording before passing it to publisher.consume, mirroring the wiring in server.py::chat() for the streaming path. - Worker passes self.dal and is_streaming=True into the helper; the ChatRequest already carries conversation_source='conversations' and request_type='user_chat' from the existing construction at worker.py:624-652. Tests: - tests/core/conversations_worker/test_worker_usage_recorder.py (new): asserts the wrapper is called with the raw stream, the publisher consumes the wrapped stream, and the helper receives the worker's dal + is_streaming=True. - tests/test_chat_recorder_state.py: re-pointed at the public helper, added coverage for explicit conversation_source='conversations' (worker-set) winning over the chat_history default. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/conversations_worker/worker.py | 19 +- holmes/core/usage_recorder.py | 130 +++++++++++- server.py | 138 +------------ .../test_worker_usage_recorder.py | 185 ++++++++++++++++++ tests/test_chat_recorder_state.py | 175 +++++++++++------ 5 files changed, 450 insertions(+), 197 deletions(-) create mode 100644 tests/core/conversations_worker/test_worker_usage_recorder.py diff --git a/holmes/core/conversations_worker/worker.py b/holmes/core/conversations_worker/worker.py index 2c4986dcab..552bd4a539 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 @@ -801,7 +805,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, @@ -810,6 +826,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/usage_recorder.py b/holmes/core/usage_recorder.py index fdd5761589..39d17d4a3e 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -11,6 +11,7 @@ from __future__ import annotations import logging +import re import threading import time import uuid @@ -21,10 +22,132 @@ 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 +# 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, + ) + + @dataclass class UsageRecorderState: """All the data needed to write one HolmesUsageEvents row. @@ -243,7 +366,10 @@ def _fire(state: UsageRecorderState) -> None: __all__ = [ "UsageRecorderState", - "stream_with_usage_recording", - "record_from_llm_result", + "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 05f0b7f933..9755c3b8ef 100644 --- a/server.py +++ b/server.py @@ -11,12 +11,11 @@ # IMPORTING ABOVE MIGHT INITIALIZE AN HTTPS CLIENT THAT DOESN'T TRUST THE CUSTOM CERTIFICATE import json import logging -import re import threading import time from datetime import datetime from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import List, Optional import colorlog import litellm @@ -73,7 +72,7 @@ ) from holmes.core.tracing import TracingFactory from holmes.core.usage_recorder import ( - UsageRecorderState, + build_chat_recorder_state, record_error, record_from_llm_result, stream_with_usage_recording, @@ -81,131 +80,6 @@ from holmes.utils.stream import stream_chat_formatter -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). - """ - if not model: - return "unknown" - try: - # litellm.get_llm_provider returns (model, provider, api_key, api_base). - return litellm.get_llm_provider(model)[1] or "unknown" - except Exception: - return model.split("/")[0] if "/" in model else "unknown" - - -# 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" -# This regex matches that prefix so we can tag Slack-driven rows in -# HolmesUsageEvents without requiring the runner to populate request_type -# explicitly. Heuristic — fragile if the runner format changes; the proper -# long-term fix is the runner sending request_type/user_id/conversation_id -# itself. Captured groups also let us stash the slack user id and timestamp -# in `meta.slack` for dashboards, since they aren't (yet) sent as structured -# fields. -_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. - - Used as a heuristic in _build_chat_recorder_state to auto-classify - Slack-driven calls until the runner is updated to send request_type - and user_id explicitly. - """ - 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 _build_chat_recorder_state( - chat_request: ChatRequest, - request_ai, - is_streaming: bool, -) -> UsageRecorderState: - """Construct a UsageRecorderState for /api/chat. All entry points (direct, - worker, scheduled, agui) flow through chat() — request_type, request_source, - source_ref, conversation_id, conversation_source, user_id are read off - chat_request (set by whoever called chat()). - """ - # Default conversation_source to 'chat_history' when conversation_id is set - # but the caller didn't override (i.e. direct /api/chat). - 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: the Robusta runner's Slack handler doesn't yet - # populate request_type / request_source / user_id / conversation_id - # structurally — instead it prepends a '**@user_X** • \n\n' marker to - # `ask`. Detect that and tag both request_type='slack_chat' (matches the - # _chat suffix of user_chat / agui_chat) and request_source='slack' - # (default UI flow label) so dashboards group Slack traffic on either - # axis. Caller-supplied values for either field still win — the runner - # can later override with finer values like 'slack_mention' / - # 'slack_alert_investigation' without code changes here. - 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, - ) - - def init_logging(): # Filter out periodical healniss and readiness probe. uvicorn_logger = logging.getLogger("uvicorn.access") @@ -652,7 +526,9 @@ def chat(chat_request: ChatRequest, http_request: Request): # 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, is_streaming=True) + 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, @@ -675,7 +551,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, is_streaming=False) + 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 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..d5cc3e2fd3 --- /dev/null +++ b/tests/core/conversations_worker/test_worker_usage_recorder.py @@ -0,0 +1,185 @@ +"""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 diff --git a/tests/test_chat_recorder_state.py b/tests/test_chat_recorder_state.py index d8d6bc4c2d..760076e2bc 100644 --- a/tests/test_chat_recorder_state.py +++ b/tests/test_chat_recorder_state.py @@ -1,30 +1,23 @@ -"""Unit tests for server._build_chat_recorder_state. - -Specifically covers the is_internal derivation logic: -- explicit chat_request.is_internal=True wins -- explicit chat_request.is_internal=False wins (even with internal_ prefix) -- is_internal=None falls back to detecting the legacy 'internal_' prefix - on request_source -- is_internal=None with no internal_ prefix yields False - -Plus a smoke check that the rest of the recorder state is built correctly. +"""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, patch - -import pytest +from unittest.mock import MagicMock -# Importing server is heavy (initializes dal/config), so do it lazily inside a -# fixture. The test still exercises the real function. - - -@pytest.fixture(scope="module") -def build_chat_recorder_state(): - # Patch the module-level dal in server.py before importing so init doesn't - # try to authenticate against a real Supabase URL. - with patch("holmes.core.supabase_dal.create_client"): - from server import _build_chat_recorder_state - return _build_chat_recorder_state +from holmes.core.usage_recorder import build_chat_recorder_state def _make_request_ai(model="openai/gpt-4", is_robusta=False): @@ -37,7 +30,7 @@ def _make_request_ai(model="openai/gpt-4", is_robusta=False): def _chat_request(**overrides): """Build a minimal ChatRequest. Imports lazily so server-import side - effects happen only once via the fixture.""" + effects don't hit collection-time.""" from holmes.core.models import ChatRequest base = dict(ask="test question", stream=False) @@ -45,44 +38,60 @@ def _chat_request(**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, build_chat_recorder_state): + def test_explicit_true_wins(self): req = _chat_request(is_internal=True, request_source="freeform") - state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=False) + 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, build_chat_recorder_state - ): + 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(), is_streaming=False) + 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, build_chat_recorder_state): + 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(), is_streaming=False) + 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, build_chat_recorder_state): + def test_unset_with_no_prefix_yields_false(self): req = _chat_request(request_source="freeform") - state = build_chat_recorder_state(req, _make_request_ai(), is_streaming=False) + 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, build_chat_recorder_state): + 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(), is_streaming=False) + 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.""" + """Catch obvious wiring regressions in build_chat_recorder_state.""" - def test_carries_through_basic_fields(self, build_chat_recorder_state): + def test_carries_through_basic_fields(self): req = _chat_request( user_id="u-abc", conversation_id="conv-123", @@ -91,7 +100,10 @@ def test_carries_through_basic_fields(self, build_chat_recorder_state): meta={"experiment_id": "x"}, ) state = build_chat_recorder_state( - req, _make_request_ai(model="anthropic/claude-sonnet-4-5"), is_streaming=True + req, + _make_request_ai(model="anthropic/claude-sonnet-4-5"), + dal=_dal(), + is_streaming=True, ) assert state.request_type == "user_chat" # default @@ -106,6 +118,29 @@ def test_carries_through_basic_fields(self, build_chat_recorder_state): 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 = ( @@ -114,60 +149,68 @@ def test_carries_through_basic_fields(self, build_chat_recorder_state): class TestSlackAutoDetect: - def test_slack_prefix_sets_request_type_to_slack_chat(self, build_chat_recorder_state): + 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(), is_streaming=True) + 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, build_chat_recorder_state): + 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(), is_streaming=True) + 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, build_chat_recorder_state): + 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(), is_streaming=True) + 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, build_chat_recorder_state - ): + 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(), is_streaming=True) + 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, build_chat_recorder_state - ): + 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(), is_streaming=True) + 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, build_chat_recorder_state - ): + 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(), is_streaming=False) + 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, build_chat_recorder_state): + 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(), is_streaming=True) + 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", @@ -177,10 +220,14 @@ def test_slack_meta_merges_with_fe_meta(self, build_chat_recorder_state): }, } - def test_partial_slack_prefix_does_not_match(self, build_chat_recorder_state): + 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(), is_streaming=False) + 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 From ec2236e4f61a64f47b7aa163fec6b602ea8b34e5 Mon Sep 17 00:00:00 2001 From: alonelish Date: Tue, 5 May 2026 13:33:38 +0300 Subject: [PATCH 10/21] Fall back to Conversations row for user_id / request_source on worker turns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After the previous fix, worker-driven conversations were producing HolmesUsageEvents rows but with NULL user_id and NULL request_source. Root cause: the FE writes both onto the Conversations row when it creates a chat (user_id as a column, request_source under metadata) but doesn't repeat them in every per-turn user_message event's data. The worker was reading only from the per-event blob, so follow-up turns came in with both fields NULL. Fix: - ConversationTask gains an optional ``user_id`` field; the Conversations row's ``user_id`` column is surfaced onto the task in ``_build_task_from_conversation_row``. - ``_process_conversation`` now resolves both fields with a per-event- wins fallback: ``data.get("user_id") or task.user_id`` and ``data.get("request_source") or task.metadata.get("request_source")``. The Conversations row is the fallback, not an override, so any future flow that legitimately wants to override per-turn (e.g. a chat pivoting from alert_investigation to freeform) still can. - source_ref / meta / is_internal stay per-event only — they're per-turn signals, not Conversation-level state. Tests: - test_worker_lifecycle.py: extend existing row-parsing tests to assert user_id round-trips off the row, and stays None when the row omits it. - test_worker_usage_recorder.py: five new tests pinning the four fallback combinations (row-only, event-only, both, neither) for both user_id and request_source. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/conversations_worker/models.py | 6 + holmes/core/conversations_worker/worker.py | 35 +++-- .../test_worker_lifecycle.py | 8 ++ .../test_worker_usage_recorder.py | 125 ++++++++++++++++++ 4 files changed, 165 insertions(+), 9 deletions(-) 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 552bd4a539..e6b21c8e52 100644 --- a/holmes/core/conversations_worker/worker.py +++ b/holmes/core/conversations_worker/worker.py @@ -487,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( @@ -621,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"), @@ -634,16 +653,14 @@ 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"), - # AI usage tracking (HolmesUsageEvents). user_id / request_source / - # source_ref / meta / is_internal come from the FE-supplied - # user_message blob; conversation_id matches the Conversations row - # the worker is processing. Note: upstream/master added - # `user_id=data.get("user_id")` independently for its own purposes - # (frontend tools / approval flow); our block already includes it, - # so the merge is a clean superset. - 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="user_chat", - request_source=data.get("request_source"), + request_source=resolved_request_source, source_ref=data.get("source_ref"), conversation_id=task.conversation_id, conversation_source="conversations", 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 index d5cc3e2fd3..8347102ebf 100644 --- a/tests/core/conversations_worker/test_worker_usage_recorder.py +++ b/tests/core/conversations_worker/test_worker_usage_recorder.py @@ -183,3 +183,128 @@ def test_recorder_state_uses_workers_dal_and_streaming_flag(): 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 From e776bd0bedd15ac0a85b294a4038134ed5b7cfbd Mon Sep 17 00:00:00 2001 From: alonelish Date: Tue, 5 May 2026 14:34:58 +0300 Subject: [PATCH 11/21] Use UTC for feedback_at and let worker request_type auto-detect MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bugs surfaced during PR review: 1. record_feedback wrote feedback_at via datetime.now().isoformat() — a NAIVE local-time ISO string. Postgres timestamptz then interprets the value relative to the *server* timezone, which differs from the Holmes pod's TZ. Both interpretations produce shifted timestamps. Fix: datetime.now(timezone.utc).isoformat() so the column always stores the actual feedback moment regardless of pod / DB timezone. 2. ConversationWorker hard-coded 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 → 'slack_chat', etc.) when chat_request.request_type is falsy, so a hard-coded value short-circuited every detection path. Today only /api/chat sees the Slack prefix, but the runner could route Slack through Conversations at any time without a code change here, and then those rows would be mis-tagged. Fix: pass data.get("request_type") through (None when absent) so the helper handles the default + Slack detection consistently for both /api/chat and the worker path. Tests: - test_supabase_dal_usage.py: existing record_feedback test now asserts feedback_at carries a UTC offset (+00:00 or Z), not a naive string. - test_worker_usage_recorder.py: three new tests pinning the worker's request_type passthrough behavior — FE-set value wins, unset stays None for the helper, Slack-shaped ask propagates with request_type still None so the helper can detect it. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/conversations_worker/worker.py | 10 ++- holmes/core/supabase_dal.py | 11 ++- .../test_worker_usage_recorder.py | 76 +++++++++++++++++++ tests/core/test_supabase_dal_usage.py | 12 ++- 4 files changed, 105 insertions(+), 4 deletions(-) diff --git a/holmes/core/conversations_worker/worker.py b/holmes/core/conversations_worker/worker.py index e6b21c8e52..b978fd4dfa 100644 --- a/holmes/core/conversations_worker/worker.py +++ b/holmes/core/conversations_worker/worker.py @@ -659,7 +659,15 @@ def _process_conversation(self, task: ConversationTask) -> None: # 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="user_chat", + # 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, diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index a2e873ce47..4aafc4318c 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -5,7 +5,7 @@ import logging import os import threading -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from enum import Enum from typing import Dict, List, Optional, Tuple from uuid import uuid4 @@ -936,7 +936,14 @@ def record_feedback( "feedback_sentiment": sentiment, "feedback_category": category, "feedback_comment": comment, - "feedback_at": datetime.now().isoformat(), + # Use UTC explicitly. datetime.now() (no tz) returns the + # Holmes pod's local time and produces a naive ISO string; + # Postgres timestamptz then interprets it relative to the + # *server*'s timezone, which differs from the pod's. Both + # cases produce shifted timestamps. Always pass an aware + # UTC datetime so the column stores the actual feedback + # moment regardless of pod / DB timezone settings. + "feedback_at": datetime.now(timezone.utc).isoformat(), } ).eq("account_id", self.account_id).eq("request_id", request_id) if user_id: diff --git a/tests/core/conversations_worker/test_worker_usage_recorder.py b/tests/core/conversations_worker/test_worker_usage_recorder.py index 8347102ebf..2af3bf9ff7 100644 --- a/tests/core/conversations_worker/test_worker_usage_recorder.py +++ b/tests/core/conversations_worker/test_worker_usage_recorder.py @@ -308,3 +308,79 @@ def test_no_fallback_values_anywhere_yields_null(): 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 index bb6b76cd7d..1ee43ca147 100644 --- a/tests/core/test_supabase_dal_usage.py +++ b/tests/core/test_supabase_dal_usage.py @@ -307,7 +307,17 @@ def test_updates_correct_columns_with_account_and_request_id_filter(self, mock_d assert update_payload["feedback_sentiment"] == "thumbs_down" assert update_payload["feedback_category"] == "wrong_answer" assert update_payload["feedback_comment"] == "missed the OOM" - assert "feedback_at" in update_payload # an ISO timestamp string + # feedback_at must be a UTC ISO timestamp — without an explicit + # timezone, Postgres timestamptz interprets the value relative to + # the *server*'s local time which differs from the Holmes pod's, + # producing shifted timestamps. Assert a UTC offset is present. + feedback_at = update_payload["feedback_at"] + assert isinstance(feedback_at, str) + # Python's datetime.isoformat() with timezone.utc produces + # '2026-01-02T03:04:05.678+00:00' — accept either '+00:00' or 'Z'. + assert feedback_at.endswith("+00:00") or feedback_at.endswith("Z"), ( + f"feedback_at must include a UTC offset, got {feedback_at!r}" + ) # Account-scoped + request_id WHERE clauses eq_calls = update_mock.return_value.eq.call_args_list From 0101f1f50ea1bd774c17c0d8774d94e1b3771752 Mon Sep 17 00:00:00 2001 From: alonelish Date: Tue, 5 May 2026 14:42:03 +0300 Subject: [PATCH 12/21] Dedupe _resolve_provider via shared usage_recorder.resolve_provider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three byte-equivalent copies of the same litellm-provider lookup had accumulated: - holmes/core/usage_recorder.py:resolve_provider (the canonical, public) - holmes/checks/checks_api.py:_resolve_provider (private duplicate) - experimental/ag-ui/server-agui.py inline try/except (third copy) Replaced both copies with imports of the canonical helper. Drops the now-unused `import litellm` from checks_api.py and the local `import litellm as _litellm` from server-agui.py. Behavior unchanged — the bodies were already identical; the only practical effect is that future tweaks to provider resolution (e.g. handling a litellm API shift) need to land in exactly one place. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- experimental/ag-ui/server-agui.py | 7 ++----- holmes/checks/checks_api.py | 15 ++------------- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/experimental/ag-ui/server-agui.py b/experimental/ag-ui/server-agui.py index 5da8ca10ad..dd5c6a88a4 100644 --- a/experimental/ag-ui/server-agui.py +++ b/experimental/ag-ui/server-agui.py @@ -26,6 +26,7 @@ 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 ( @@ -164,11 +165,7 @@ async def event_generator(message_history): e, ) ai_model = getattr(ai.llm, "model", None) or chat_request.model or "unknown" - try: - import litellm as _litellm # local import: this file is opt-in - ai_provider = _litellm.get_llm_provider(ai_model)[1] or "unknown" - except Exception: - ai_provider = ai_model.split("/")[0] if "/" in ai_model else "unknown" + ai_provider = resolve_provider(ai_model) recorder_state = UsageRecorderState( dal=dal, request_type="agui_chat", diff --git a/holmes/checks/checks_api.py b/holmes/checks/checks_api.py index 48867f4af0..fa978ce5ba 100644 --- a/holmes/checks/checks_api.py +++ b/holmes/checks/checks_api.py @@ -3,7 +3,6 @@ import time from typing import Optional -import litellm from fastapi import FastAPI, HTTPException from litellm.exceptions import AuthenticationError from pydantic import BaseModel, Field @@ -14,19 +13,9 @@ 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 +from holmes.core.usage_recorder import UsageRecorderState, resolve_provider from holmes.plugins.destinations.slack.plugin import SlackDestination - -def _resolve_provider(model: Optional[str]) -> str: - """Best-effort: return the canonical litellm provider for `model`.""" - if not model: - return "unknown" - try: - return litellm.get_llm_provider(model)[1] or "unknown" - except Exception: - return model.split("/")[0] if "/" in model else "unknown" - checks_app = FastAPI() _CONFIG: Config @@ -136,7 +125,7 @@ def execute_health_check( user_id=None, is_streaming=False, model=ai_model, - provider=_resolve_provider(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}, ) From 321f9739aeeb058fa6c4a0e6b58b50ef1e06e9e1 Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 6 May 2026 13:09:00 +0300 Subject: [PATCH 13/21] Drop redundant 'Grafana' prefix from Loki / Tempo catalog names Both entries already use the Grafana icon and live alongside other Grafana-family integrations; the prefix in the display name was just duplicating that signal. Brings them in line with how the rest of the catalog labels integrations (e.g. just 'Loki' / 'Tempo' rather than 'Grafana Loki' / 'Grafana Tempo'). IDs unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- datasource-catalog.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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", From a17fc1b8e501d22e39b1c22e2cba350f54b57b36 Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 6 May 2026 15:57:39 +0300 Subject: [PATCH 14/21] =?UTF-8?q?Drop=20/api/feedback=20endpoint=20?= =?UTF-8?q?=E2=80=94=20FE=20will=20call=20record=5Ffeedback()=20RPC=20dire?= =?UTF-8?q?ctly?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback (Moshe Morad on PR #1969): the /api/feedback endpoint was a thin pass-through with no LLM/tool involvement, and using a Postgres RPC instead lets the function scope by `auth.uid()` (JWT-derived) rather than the FE-supplied query-param `user_id` Holmes was trusting blindly. Strictly stronger user-scoping. Removed from Holmes: - server.py: `class FeedbackRequest`, `def feedback`, the BaseModel import that became unused. - holmes/core/supabase_dal.py: `def record_feedback`, the `timezone` import (only used here). Replaced with a 6-line comment pointer to the RPC for future readers. - tests/core/test_supabase_dal_usage.py: the four record_feedback tests (no Holmes-side code path remains to unit-test). - Stale doc strings referencing `POST /api/feedback` updated to point at the RPC (server.py response-metadata comment, usage_recorder.py stream-wrapper docstring, test module docstring). The RPC body lives in the migration script (see plan file). FE migration from any prior endpoint to `supabase.rpc('record_feedback', ...)` is on the FE team — coordinating with this PR's deploy. Tests: 144 pass (was 148; the 4 dropped were record_feedback unit tests). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/supabase_dal.py | 67 ++---------------- holmes/core/usage_recorder.py | 5 +- server.py | 38 +---------- tests/core/test_supabase_dal_usage.py | 97 ++++----------------------- 4 files changed, 26 insertions(+), 181 deletions(-) diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index 4aafc4318c..3f5580b009 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -5,7 +5,7 @@ import logging import os import threading -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from enum import Enum from typing import Dict, List, Optional, Tuple from uuid import uuid4 @@ -892,65 +892,12 @@ def record_usage_event( except Exception: logging.exception("Failed to record usage event") - def record_feedback( - self, - *, - request_id: str, - sentiment: str, - category: Optional[str], - comment: Optional[str], - user_id: Optional[str], - ) -> None: - """Update the feedback_* columns on a HolmesUsageEvents row. - - Best-effort. The UPDATE is account-scoped via WHERE so RLS and an - explicit predicate both prevent cross-account writes if request_id - ever collides. - - ``user_id`` is optional. When provided, the UPDATE is additionally - scoped by ``.eq("user_id", user_id)`` as defense in depth — even - though v1's auth model assumes rater == asker (so events.user_id - already matches the user posting the thumb), the extra predicate - prevents cross-user overwrites if a request_id ever leaks. When - omitted, no user filter is applied, leaving room for future - system / scheduled flows that might emit feedback without a - specific human user. - """ - if not self.enabled: - return - if sentiment not in ("thumbs_up", "thumbs_down"): - logging.warning( - "record_feedback: ignoring invalid sentiment %r", sentiment - ) - return - try: - # Defense in depth: when user_id is provided, scope the UPDATE so - # one user can't overwrite another user's feedback within the same - # account if a request_id ever leaks. v1's auth model assumes - # rater == asker, so the user_id sent on POST /api/feedback - # matches HolmesUsageEvents.user_id for normal flows. Skip the - # filter when user_id is None (e.g. system / scheduled flows that - # might emit feedback in the future) to avoid breaking those. - query = self.client.table(HOLMES_USAGE_EVENTS_TABLE).update( - { - "feedback_sentiment": sentiment, - "feedback_category": category, - "feedback_comment": comment, - # Use UTC explicitly. datetime.now() (no tz) returns the - # Holmes pod's local time and produces a naive ISO string; - # Postgres timestamptz then interprets it relative to the - # *server*'s timezone, which differs from the pod's. Both - # cases produce shifted timestamps. Always pass an aware - # UTC datetime so the column stores the actual feedback - # moment regardless of pod / DB timezone settings. - "feedback_at": datetime.now(timezone.utc).isoformat(), - } - ).eq("account_id", self.account_id).eq("request_id", request_id) - if user_id: - query = query.eq("user_id", user_id) - query.execute() - except Exception: - logging.exception("Failed to record feedback") + # 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: """ diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 39d17d4a3e..b77de82595 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -225,8 +225,9 @@ def stream_with_usage_recording( 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 posts it back to - ``POST /api/feedback`` when the user clicks thumbs up/down. + 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: diff --git a/server.py b/server.py index 9755c3b8ef..adf3db67fd 100644 --- a/server.py +++ b/server.py @@ -27,8 +27,6 @@ from fastapi import FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, StreamingResponse from litellm.exceptions import AuthenticationError -from pydantic import BaseModel - from holmes import get_version, is_official_release from holmes.common.env_vars import ( DEVELOPMENT_MODE, @@ -595,8 +593,8 @@ 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 POST /api/feedback later. Streaming path does the - # same via _inject_request_id in the stream wrapper. + # 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( @@ -629,38 +627,6 @@ def chat(chat_request: ChatRequest, http_request: Request): raise HTTPException(status_code=500, detail=str(e)) -class FeedbackRequest(BaseModel): - """Body for POST /api/feedback. Defined inline to avoid bloating models.py - with an endpoint-specific schema.""" - request_id: str - sentiment: str # 'thumbs_up' | 'thumbs_down' - category: Optional[str] = None - comment: Optional[str] = None - - -@app.post("/api/feedback") -def feedback(req: FeedbackRequest, http_request: Request) -> dict: - """Record user feedback (thumbs up/down + optional category/comment) on - a previously-completed chat call. Looks up the row by request_id and - UPDATEs feedback_* columns. Best-effort: returns 200 even if the row - isn't found yet (rare network reorder).""" - if req.sentiment not in ("thumbs_up", "thumbs_down"): - raise HTTPException(status_code=400, detail="invalid sentiment") - # user_id is optional — when supplied via query param (mirrors /api/chat's - # resolution path), the DAL adds it as defense-in-depth on the UPDATE. - # When absent, account_id RLS + (account_id, request_id) WHERE is enough - # to prevent cross-account writes. - user_id: Optional[str] = http_request.query_params.get("user_id") - dal.record_feedback( - request_id=req.request_id, - sentiment=req.sentiment, - category=req.category, - comment=req.comment, - user_id=user_id, - ) - return {"ok": True} - - scheduled_prompts_executor = ScheduledPromptsExecutor( dal=dal, config=config, chat_function=chat ) diff --git a/tests/core/test_supabase_dal_usage.py b/tests/core/test_supabase_dal_usage.py index 1ee43ca147..5f00316371 100644 --- a/tests/core/test_supabase_dal_usage.py +++ b/tests/core/test_supabase_dal_usage.py @@ -1,8 +1,13 @@ -"""Unit tests for SupabaseDal.record_usage_event and SupabaseDal.record_feedback. +"""Unit tests for SupabaseDal.record_usage_event. -These methods are best-effort: they swallow 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 / .update) and the failure path +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). """ @@ -262,82 +267,8 @@ def test_handles_stats_with_none_cached_tokens(self, mock_dal): assert payload["cached_tokens"] is None -# ────────────────────────────────────────────────────────────────── -# record_feedback -# ────────────────────────────────────────────────────────────────── - - -class TestRecordFeedback: - def test_no_op_when_dal_disabled(self, mock_dal): - mock_dal.enabled = False - mock_dal.record_feedback( - request_id="r-1", - sentiment="thumbs_up", - category=None, - comment=None, - user_id=None, - ) - mock_dal.client.table.assert_not_called() - - def test_invalid_sentiment_logs_and_returns(self, mock_dal): - # Should not raise, should not call .update. - mock_dal.record_feedback( - request_id="r-1", - sentiment="meh", # not allowed - category=None, - comment=None, - user_id=None, - ) - mock_dal.client.table.assert_not_called() - - def test_updates_correct_columns_with_account_and_request_id_filter(self, mock_dal): - mock_dal.record_feedback( - request_id="req-uuid-123", - sentiment="thumbs_down", - category="wrong_answer", - comment="missed the OOM", - user_id="user-xyz", - ) - - # Chain: client.table(...).update({...}).eq("account_id", ...).eq("request_id", ...).execute() - mock_dal.client.table.assert_called_once_with(HOLMES_USAGE_EVENTS_TABLE) - update_mock = mock_dal.client.table.return_value.update - update_mock.assert_called_once() - update_payload = update_mock.call_args.args[0] - assert update_payload["feedback_sentiment"] == "thumbs_down" - assert update_payload["feedback_category"] == "wrong_answer" - assert update_payload["feedback_comment"] == "missed the OOM" - # feedback_at must be a UTC ISO timestamp — without an explicit - # timezone, Postgres timestamptz interprets the value relative to - # the *server*'s local time which differs from the Holmes pod's, - # producing shifted timestamps. Assert a UTC offset is present. - feedback_at = update_payload["feedback_at"] - assert isinstance(feedback_at, str) - # Python's datetime.isoformat() with timezone.utc produces - # '2026-01-02T03:04:05.678+00:00' — accept either '+00:00' or 'Z'. - assert feedback_at.endswith("+00:00") or feedback_at.endswith("Z"), ( - f"feedback_at must include a UTC offset, got {feedback_at!r}" - ) - - # Account-scoped + request_id WHERE clauses - eq_calls = update_mock.return_value.eq.call_args_list - # First .eq is account_id, second is request_id (chained on the return value) - assert eq_calls[0].args == ("account_id", mock_dal.account_id) - # The second .eq is on the result of the first, but with our MagicMock chain - # it's also captured here. In any case, both account_id and request_id - # must be in the predicate chain. - first_eq_return = update_mock.return_value.eq.return_value - first_eq_return.eq.assert_called_with("request_id", "req-uuid-123") - - def test_swallows_supabase_errors(self, mock_dal): - mock_dal.client.table.return_value.update.return_value.eq.return_value.eq.return_value.execute.side_effect = ( - RuntimeError("update failed") - ) - # Should not raise. - mock_dal.record_feedback( - request_id="r-1", - sentiment="thumbs_up", - category=None, - comment=None, - user_id=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. From ab4958b231ed26914d1bcd15301b05620c78dacf Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 6 May 2026 17:12:59 +0300 Subject: [PATCH 15/21] Document each UsageRecorderState field inline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback (Moshe Morad on PR #1969): the dataclass had 18 fields with only group-level comments, making it hard for callers to know which to set, what valid values are, and which are filled in later by the wrapper. Each field now carries a 1–8 line comment covering: - what the value represents - where it comes from (entry point vs wrapper-filled) - the canonical / accepted values for enum-ish fields (request_type, request_source, conversation_source, status, finish_reason, ...) - what gets written to the DB if the field is left at its default - cross-references to the helpers that fill mutable fields The class-level docstring also gained a clearer three-group breakdown (required identity, optional identity, mutable runtime) so a reader scanning the file knows up front which fields belong to which phase. No behavior change — pure documentation. All 144 existing tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/usage_recorder.py | 160 ++++++++++++++++++++++++++++++++-- 1 file changed, 152 insertions(+), 8 deletions(-) diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index b77de82595..5155a53db2 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -152,37 +152,181 @@ def build_chat_recorder_state( class UsageRecorderState: """All the data needed to write one HolmesUsageEvents row. - Identity / classification fields are set by the entry point at - construction time. Mutable runtime fields (`stats`, `iterations`, - `tool_call_count`, `finish_reason`, `status`) are filled by the - stream wrapper or by `record_from_llm_result` before firing. + 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. """ - # required identity / classification — set by the entry point - dal: Any # SupabaseDal; typed Any to avoid circular import at runtime + # ── Group 1: required identity / classification ───────────────────── + + # SupabaseDal handle. The recorder calls dal.record_usage_event(...) + # on a daemon thread when the request finishes. 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 - # optional identity / classification + # ── 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) - # mutable — filled during the call by the wrapper or recorder + # ── 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: + # 'success' — terminal ANSWER_END event seen + # 'approval_required' — terminal APPROVAL_REQUIRED event seen + # 'error' — terminal ERROR event or exception + # 'rate_limited' — record_error detected a rate-limit message + # 'aborted' — stream ended without any terminal event + # (client disconnect / generator exhaustion) + # Default 'success' is overwritten by the wrapper's finally-block + # to 'aborted' if no terminal event was ever observed. status: str = "success" def to_kwargs(self) -> Dict[str, Any]: From f942de1db4780ec42190df4e6750a81dfd19e0ec Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 6 May 2026 17:20:08 +0300 Subject: [PATCH 16/21] Replace string-literal status values with RequestStatus enum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback (Moshe Morad on PR #1969): the recorder's `status` column had five string-literal values ('success' / 'approval_required' / 'error' / 'rate_limited' / 'aborted') sprinkled across the field default, three branches of stream_with_usage_recording, and both record_from_llm_result / record_error. Easy to typo, hard to discover the full set, no compiler help. Introduced RequestStatus(str, Enum) — same pattern as ConversationStatus in conversations_worker/models.py. Subclassing str keeps the values JSON-serializable as plain strings (verified: json.dumps produces "success", not "RequestStatus.SUCCESS") and equality with old string literals still works, so existing tests that assert `kwargs["status"] == "success"` keep passing without changes. Also dropped the unused HOLMES_USAGE_EVENT_CALLS_TABLE constant from supabase_dal.py — leftover from the child-table approach we explicitly dropped earlier in this PR. No code referenced it. All 144 existing tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/supabase_dal.py | 1 - holmes/core/usage_recorder.py | 57 ++++++++++++++++++++++------------- 2 files changed, 36 insertions(+), 22 deletions(-) diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index 3f5580b009..9cd77534bc 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -70,7 +70,6 @@ CONVERSATION_EVENTS_TABLE = "ConversationEvents" OAUTH_TOKENS_TABLE = "OAuthTokens" HOLMES_USAGE_EVENTS_TABLE = "HolmesUsageEvents" -HOLMES_USAGE_EVENT_CALLS_TABLE = "HolmesUsageEventCalls" ENRICHMENT_BLACKLIST = ["text_file", "graph", "ai_analysis", "holmes"] ENRICHMENT_BLACKLIST_SET = set(ENRICHMENT_BLACKLIST) diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 5155a53db2..c1213de032 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -16,6 +16,7 @@ import time import uuid 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 @@ -148,6 +149,24 @@ def build_chat_recorder_state( ) +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. @@ -318,16 +337,11 @@ class UsageRecorderState: finish_reason: Optional[str] = None # Final outcome of the request. Filled by the wrapper or the - # record_error helper: - # 'success' — terminal ANSWER_END event seen - # 'approval_required' — terminal APPROVAL_REQUIRED event seen - # 'error' — terminal ERROR event or exception - # 'rate_limited' — record_error detected a rate-limit message - # 'aborted' — stream ended without any terminal event - # (client disconnect / generator exhaustion) - # Default 'success' is overwritten by the wrapper's finally-block - # to 'aborted' if no terminal event was ever observed. - status: str = "success" + # 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 def to_kwargs(self) -> Dict[str, Any]: """Pack the state into the kwargs `SupabaseDal.record_usage_event` expects.""" @@ -381,31 +395,31 @@ def stream_with_usage_recording( elif msg.event == StreamEvents.ANSWER_END: _capture_terminal(state, msg.data) _inject_request_id(msg.data, state.request_id) - state.status = "success" + state.status = RequestStatus.SUCCESS saw_terminal = True elif msg.event == StreamEvents.APPROVAL_REQUIRED: _capture_terminal(state, msg.data) _inject_request_id(msg.data, state.request_id) - state.status = "approval_required" + state.status = RequestStatus.APPROVAL_REQUIRED saw_terminal = True elif msg.event == StreamEvents.ERROR: _capture_terminal(state, msg.data) _inject_request_id(msg.data, state.request_id) - state.status = "error" + state.status = RequestStatus.ERROR saw_terminal = True yield msg except Exception: if not saw_terminal: - state.status = "error" + 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 + # `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 == "success": - state.status = "aborted" + if not saw_terminal and state.status == RequestStatus.SUCCESS: + state.status = RequestStatus.ABORTED _fire(state) @@ -479,7 +493,7 @@ def record_from_llm_result( 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 = "success" + state.status = RequestStatus.SUCCESS _fire(state) @@ -487,9 +501,9 @@ 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 = "rate_limited" + state.status = RequestStatus.RATE_LIMITED else: - state.status = "error" + state.status = RequestStatus.ERROR _fire(state) @@ -510,6 +524,7 @@ def _fire(state: UsageRecorderState) -> None: __all__ = [ + "RequestStatus", "UsageRecorderState", "build_chat_recorder_state", "detect_slack_origin", From 32fc36d973548c34851a5a2f973ec9e3f1149a4e Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 6 May 2026 17:27:42 +0300 Subject: [PATCH 17/21] record_usage_event takes UsageRecorderState directly; drop to_kwargs() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback (Moshe Morad on PR #1969): the DAL had a 20-kwarg signature that mirrored every field on UsageRecorderState plus a to_kwargs() method on the state class to pack the dict. Net cost: adding a new field meant editing the dataclass, the to_kwargs method, the DAL signature, the DAL body, and any test that constructed kwargs. Five places, one logical change. Refactor: record_usage_event(state) takes the entire state object positionally and reads fields off it. The DAL is now the single place that knows the column shape — adding a new field is "add it on the state, read it here, write the migration." Net diff: - Drops UsageRecorderState.to_kwargs() (~24 lines). - Replaces the 20-kwarg DAL signature with `state: UsageRecorderState` (~50 lines lighter, more readable). - Adds duration_ms as a @property on the state — was previously computed inside to_kwargs(), now the DAL just reads state.duration_ms. - TYPE_CHECKING import in supabase_dal.py for UsageRecorderState (no runtime circular-import risk since usage_recorder also TYPE_CHECKING- imports SupabaseDal). - Tests refactored: kwargs-based call sites become `mock_dal.record_usage_event(_make_state(...))`, and the recorder tests assert on `state.dal.record_usage_event.call_args.args[0]` (the live state object) instead of the old kwargs dict. Added one new test for state.stats=None handling and one for the duration_ms property type contract. 146 pass (was 144). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/supabase_dal.py | 83 ++++++------- holmes/core/usage_recorder.py | 43 +++---- tests/core/test_supabase_dal_usage.py | 168 ++++++++------------------ tests/core/test_usage_recorder.py | 160 +++++++++++++----------- 4 files changed, 195 insertions(+), 259 deletions(-) diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index 9cd77534bc..3a8db9c3f0 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 @@ -821,52 +827,39 @@ 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, - *, - request_type: str, - request_source: Optional[str], - source_ref: Optional[str], - conversation_id: Optional[str], - conversation_source: Optional[str], - status: str, - model: str, - provider: str, - is_robusta_model: bool, - stats, # holmes.core.llm_usage.RequestStats; not typed to avoid circular import - iterations: int, - duration_ms: Optional[int], - tool_call_count: int, - is_streaming: bool, - finish_reason: Optional[str], - user_id: Optional[str], - cluster_id: Optional[str] = None, - request_id: Optional[str] = None, - meta: Optional[Dict] = None, - is_internal: bool = False, - ) -> None: + def record_usage_event(self, state: "UsageRecorderState") -> None: """Record one HolmesUsageEvents row. Best-effort: swallows DB errors. Called from holmes.core.usage_recorder._fire on a daemon thread, so - errors here only affect the telemetry row, never the request response. + 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": cluster_id or self.cluster, - "user_id": user_id, - "conversation_id": conversation_id, - "conversation_source": conversation_source, - "request_id": request_id, - "request_type": request_type, - "request_source": request_source, - "source_ref": source_ref, - "status": status, - "model": model, - "provider": provider, - "is_robusta_model": is_robusta_model, + "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), @@ -874,19 +867,19 @@ def record_usage_event( "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": iterations, + "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": tool_call_count, - "duration_ms": duration_ms, - "is_streaming": is_streaming, - "is_internal": is_internal, - "finish_reason": finish_reason, - "meta": meta or {}, + "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") diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index c1213de032..2a676614af 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -193,9 +193,11 @@ class UsageRecorderState: # ── Group 1: required identity / classification ───────────────────── - # SupabaseDal handle. The recorder calls dal.record_usage_event(...) - # on a daemon thread when the request finishes. Typed Any to avoid - # importing SupabaseDal at runtime (circular import via ChatRequest). + # 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 @@ -343,30 +345,15 @@ class UsageRecorderState: # terminal event was ever observed. status: RequestStatus = RequestStatus.SUCCESS - def to_kwargs(self) -> Dict[str, Any]: - """Pack the state into the kwargs `SupabaseDal.record_usage_event` expects.""" - return { - "request_type": self.request_type, - "request_source": self.request_source, - "source_ref": self.source_ref, - "conversation_id": self.conversation_id, - "conversation_source": self.conversation_source, - "status": self.status, - "model": self.model, - "provider": self.provider, - "is_robusta_model": self.is_robusta_model, - "stats": self.stats or RequestStats(), - "iterations": self.iterations, - "duration_ms": int((time.monotonic() - self.t_start) * 1000), - "tool_call_count": self.tool_call_count, - "is_streaming": self.is_streaming, - "is_internal": self.is_internal, - "finish_reason": self.finish_reason, - "user_id": self.user_id, - "cluster_id": self.cluster_id, - "request_id": self.request_id, - "meta": self.meta, - } + @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) def stream_with_usage_recording( @@ -514,7 +501,7 @@ def _fire(state: UsageRecorderState) -> None: try: threading.Thread( target=state.dal.record_usage_event, - kwargs=state.to_kwargs(), + args=(state,), daemon=True, name="usage-recorder", ).start() diff --git a/tests/core/test_supabase_dal_usage.py b/tests/core/test_supabase_dal_usage.py index 5f00316371..8551c9c670 100644 --- a/tests/core/test_supabase_dal_usage.py +++ b/tests/core/test_supabase_dal_usage.py @@ -9,6 +9,12 @@ 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 @@ -20,6 +26,7 @@ HOLMES_USAGE_EVENTS_TABLE, SupabaseDal, ) +from holmes.core.usage_recorder import UsageRecorderState @pytest.fixture @@ -48,6 +55,23 @@ def _stats() -> RequestStats: ) +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 # ────────────────────────────────────────────────────────────────── @@ -56,48 +80,28 @@ def _stats() -> RequestStats: class TestRecordUsageEvent: def test_no_op_when_dal_disabled(self, mock_dal): mock_dal.enabled = False - mock_dal.record_usage_event( - request_type="user_chat", - request_source=None, - source_ref=None, - conversation_id=None, - conversation_source=None, - status="success", - model="openai/gpt-4", - provider="openai", - is_robusta_model=False, - stats=_stats(), - iterations=1, - duration_ms=42, - tool_call_count=0, - is_streaming=False, - finish_reason=None, - user_id=None, - ) + 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( + 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", - status="success", model="anthropic/claude-sonnet-4-5", provider="anthropic", is_robusta_model=False, - stats=_stats(), iterations=3, - duration_ms=1500, 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) @@ -118,7 +122,7 @@ def test_inserts_row_with_correct_payload(self, mock_dal): assert payload["request_type"] == "user_chat" assert payload["request_source"] == "freeform" assert payload["source_ref"] == "issue-42" - assert payload["status"] == "success" + 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 @@ -137,76 +141,27 @@ def test_inserts_row_with_correct_payload(self, mock_dal): # Outcome assert payload["tool_call_count"] == 5 - assert payload["duration_ms"] == 1500 + # 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( - request_type="user_chat", - request_source=None, - source_ref=None, - conversation_id=None, - conversation_source=None, - status="success", - model="m", - provider="p", - is_robusta_model=False, - stats=_stats(), - iterations=1, - duration_ms=10, - tool_call_count=0, - is_streaming=False, - finish_reason=None, - user_id=None, - ) + 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( - request_type="user_chat", - request_source=None, - source_ref=None, - conversation_id=None, - conversation_source=None, - status="success", - model="m", - provider="p", - is_robusta_model=False, - stats=_stats(), - iterations=1, - duration_ms=10, - tool_call_count=0, - is_streaming=False, - finish_reason=None, - user_id=None, - cluster_id="other-cluster", - ) + 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): - mock_dal.record_usage_event( - request_type="user_chat", - request_source=None, - source_ref=None, - conversation_id=None, - conversation_source=None, - status="success", - model="m", - provider="p", - is_robusta_model=False, - stats=_stats(), - iterations=1, - duration_ms=10, - tool_call_count=0, - is_streaming=False, - finish_reason=None, - user_id=None, - meta=None, - ) + # 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"] == {} @@ -216,24 +171,7 @@ def test_swallows_supabase_errors(self, mock_dal): RuntimeError("supabase down") ) # Should not raise. - mock_dal.record_usage_event( - request_type="user_chat", - request_source=None, - source_ref=None, - conversation_id=None, - conversation_source=None, - status="success", - model="m", - provider="p", - is_robusta_model=False, - stats=_stats(), - iterations=1, - duration_ms=10, - tool_call_count=0, - is_streaming=False, - finish_reason=None, - user_id=None, - ) + 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. @@ -245,25 +183,21 @@ def test_handles_stats_with_none_cached_tokens(self, mock_dal): cached_tokens=None, reasoning_tokens=0, ) - mock_dal.record_usage_event( - request_type="user_chat", - request_source=None, - source_ref=None, - conversation_id=None, - conversation_source=None, - status="success", - model="m", - provider="p", - is_robusta_model=False, - stats=stats, - iterations=1, - duration_ms=10, - tool_call_count=0, - is_streaming=False, - finish_reason=None, - user_id=None, - ) + 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 diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py index b6caf3a640..98e3c1565f 100644 --- a/tests/core/test_usage_recorder.py +++ b/tests/core/test_usage_recorder.py @@ -2,11 +2,19 @@ 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 kwargs passed to dal.record_usage_event. +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, patch +from unittest.mock import MagicMock import pytest @@ -59,71 +67,84 @@ def _terminal_data(costs: dict, num_llm_calls: int = 1, finish_reason: str = "st def _patch_inline_thread(monkeypatch): - """Replace threading.Thread inside usage_recorder so target() runs inline.""" + """Replace threading.Thread inside usage_recorder so target() runs inline. + + _fire spawns the recorder thread with ``args=(state,)`` (positional + state arg, since the DAL takes a single state object now). Mirror that + exactly — pass through both args and kwargs to be tolerant of either. + """ import holmes.core.usage_recorder as mod class _InlineThread: - def __init__(self, target=None, kwargs=None, daemon=None, name=None): + def __init__(self, target=None, args=None, kwargs=None, daemon=None, name=None): self._target = target + self._args = args or () self._kwargs = kwargs or {} def start(self): - self._target(**self._kwargs) + self._target(*self._args, **self._kwargs) monkeypatch.setattr(mod.threading, "Thread", _InlineThread) +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.to_kwargs +# UsageRecorderState basics — direct attribute access + duration_ms property +# (Replaces the old TestToKwargs class; to_kwargs() no longer exists.) # ────────────────────────────────────────────────────────────────── -class TestToKwargs: - def test_packs_all_required_fields(self): +class TestStateBasics: + def test_default_values_match_spec(self): state = _make_state() - kwargs = state.to_kwargs() - # Identity - assert kwargs["request_type"] == "user_chat" - assert kwargs["request_source"] == "freeform" - assert kwargs["conversation_id"] == "conv-123" - assert kwargs["conversation_source"] == "chat_history" - assert kwargs["user_id"] == "user-abc" - assert "request_id" in kwargs and kwargs["request_id"] + 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 kwargs["model"] == "openai/gpt-4" - assert kwargs["provider"] == "openai" - assert kwargs["is_robusta_model"] is False - assert kwargs["is_streaming"] is True - - # Mutable defaults - assert kwargs["status"] == "success" - assert kwargs["iterations"] == 0 - assert kwargs["tool_call_count"] == 0 - assert kwargs["finish_reason"] is None - assert kwargs["meta"] == {} - - # Stats default to an empty RequestStats, not None - assert kwargs["stats"] is not None - assert kwargs["stats"].total_tokens == 0 - assert kwargs["stats"].total_cost == 0.0 - - def test_duration_ms_is_computed_from_t_start(self): + 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 - kwargs = state.to_kwargs() - assert kwargs["duration_ms"] >= 1000 + 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 - assert state.to_kwargs()["is_internal"] is False - def test_is_internal_true_round_trips(self): + def test_is_internal_round_trips(self): state = _make_state(is_internal=True) - assert state.to_kwargs()["is_internal"] is True + assert state.is_internal is True # ────────────────────────────────────────────────────────────────── @@ -151,13 +172,13 @@ def test_success_path_records_with_status_success(self, monkeypatch): assert len(consumed) == 4 state.dal.record_usage_event.assert_called_once() - kw = state.dal.record_usage_event.call_args.kwargs - assert kw["status"] == "success" - assert kw["tool_call_count"] == 2 - assert kw["iterations"] == 3 - assert kw["finish_reason"] == "stop" - assert kw["stats"].prompt_tokens == 100 - assert kw["stats"].total_tokens == 150 + 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) @@ -169,9 +190,9 @@ def test_error_event_marks_status_error(self, monkeypatch): ] list(stream_with_usage_recording(_stream(*events), state)) - kw = state.dal.record_usage_event.call_args.kwargs - assert kw["status"] == "error" - assert kw["tool_call_count"] == 1 + 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) @@ -180,7 +201,7 @@ def test_approval_required_marks_status_approval_required(self, monkeypatch): _stream(StreamMessage(event=StreamEvents.APPROVAL_REQUIRED, data={"metadata": {}})), state, )) - assert state.dal.record_usage_event.call_args.kwargs["status"] == "approval_required" + 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) @@ -195,9 +216,10 @@ def failing_stream(): # Recorder still fires from the finally state.dal.record_usage_event.assert_called_once() - assert state.dal.record_usage_event.call_args.kwargs["status"] == "error" + s = _state_arg(state) + assert s.status == "error" # And the tool we saw before the exception was counted - assert state.dal.record_usage_event.call_args.kwargs["tool_call_count"] == 1 + assert s.tool_call_count == 1 def test_stream_without_terminal_event_still_records_as_aborted(self, monkeypatch): _patch_inline_thread(monkeypatch) @@ -212,7 +234,7 @@ def test_stream_without_terminal_event_still_records_as_aborted(self, monkeypatc # 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.dal.record_usage_event.call_args.kwargs["status"] == "aborted" + 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.""" @@ -227,7 +249,7 @@ def test_terminal_event_keeps_its_explicit_status(self, monkeypatch): state, )) - assert state.dal.record_usage_event.call_args.kwargs["status"] == "success" + 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.""" @@ -320,13 +342,13 @@ def test_extracts_stats_iterations_finish_reason_and_tool_count(self, monkeypatc record_from_llm_result(state, fake_result) - kw = state.dal.record_usage_event.call_args.kwargs - assert kw["status"] == "success" - assert kw["iterations"] == 4 - assert kw["tool_call_count"] == 3 - assert kw["finish_reason"] == "stop" - assert kw["stats"].total_tokens == 250 - assert kw["stats"].total_cost == 0.005 + 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) @@ -341,10 +363,10 @@ def test_handles_missing_attrs_gracefully(self, monkeypatch): record_from_llm_result(state, bare) - kw = state.dal.record_usage_event.call_args.kwargs + s = _state_arg(state) # iterations falls back to 1 when num_llm_calls is None - assert kw["iterations"] == 1 - assert kw["tool_call_count"] == 0 + assert s.iterations == 1 + assert s.tool_call_count == 0 # ────────────────────────────────────────────────────────────────── @@ -357,15 +379,13 @@ 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")) - kw = state.dal.record_usage_event.call_args.kwargs - assert kw["status"] == "rate_limited" + 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")) - kw = state.dal.record_usage_event.call_args.kwargs - assert kw["status"] == "error" + assert _state_arg(state).status == "error" # ────────────────────────────────────────────────────────────────── @@ -400,7 +420,9 @@ def test_record_calls_dal_in_background_thread(self): called = threading.Event() - def slow_record(**kwargs): + # Real thread → state arrives positionally (one arg). Match the + # production call signature. + def slow_record(state): time.sleep(0.05) called.set() From cbdfc54d664aebebc4ca926a4bd85df922dd40ce Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 6 May 2026 17:40:03 +0300 Subject: [PATCH 18/21] record_from_llm_result: filter stats via Pydantic, not a hardcoded set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback (Moshe Morad on PR #1969): the helper was building RequestStats from llm_result.model_dump() by manually filtering the dict against a hardcoded set of 9 field names. Cleaner and self-maintaining: ask Pydantic to do the include-filter using RequestStats.model_fields. state.stats = RequestStats( **llm_result.model_dump(include=set(RequestStats.model_fields)) ) LLMResult IS-A RequestStats but adds extra fields (tool_calls, messages, finish_reason, ...). The `include=` parameter tells model_dump() to emit only the keys that match RequestStats's own model_fields, so the extras get dropped without us hardcoding which keys belong to stats. When RequestStats grows a new column, this code stays correct automatically. No behavior change — same fields end up on state.stats. Net −10 lines. All 146 tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/usage_recorder.py | 26 ++++++++------------------ 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 2a676614af..01e9ca9eab 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -455,24 +455,14 @@ def record_from_llm_result( model_dump. """ try: - # LLMResult inherits from RequestStats, so dump the stats fields out. - stats_fields = { - k: v - for k, v in llm_result.model_dump().items() - if k - in { - "total_cost", - "total_tokens", - "prompt_tokens", - "completion_tokens", - "cached_tokens", - "reasoning_tokens", - "max_completion_tokens_per_call", - "max_prompt_tokens_per_call", - "num_compactions", - } - } - state.stats = RequestStats(**stats_fields) + # 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() From e22fa33d4c7e7b0312e0eb76a1f5df7be5e06847 Mon Sep 17 00:00:00 2001 From: alonelish Date: Wed, 6 May 2026 18:05:39 +0300 Subject: [PATCH 19/21] Move _fire and _capture_terminal into UsageRecorderState MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback (Moshe Morad on PR #1969): the two private helpers that operate solely on the state are a natural fit as methods. Before After ────────────────────────── ────────────────────────── _fire(state) state._fire() _capture_terminal(state, data) state._capture_terminal(data) Public functions stay at module level — they're either operating on external inputs (stream_with_usage_recording's primary arg is a stream, not state) or are simple recording entry points where the function signature reads naturally: record_from_llm_result(state, llm_result) # stays a function record_error(state, exc) # stays a function stream_with_usage_recording(stream, state) # stays a function This addresses Moshe's "lots of functions take state" concern for the helpers that genuinely belong on the class without converting the whole module to OO style — the data/behavior split (passive UsageRecorderState + recording functions over it) is the same pattern Holmes uses elsewhere (build_chat_messages, format_tool_result_data, etc.). _inject_request_id stays a function too — it operates on the stream event's data dict, not on state, and only takes request_id as an arg. Updated supabase_dal.py docstring to reference UsageRecorderState._fire instead of the old module path. All 146 tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/supabase_dal.py | 2 +- holmes/core/usage_recorder.py | 97 +++++++++++++++++++---------------- 2 files changed, 54 insertions(+), 45 deletions(-) diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index 3a8db9c3f0..d54619fdab 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -830,7 +830,7 @@ def sync_toolsets(self, toolsets: list[dict], cluster_name: str) -> None: def record_usage_event(self, state: "UsageRecorderState") -> None: """Record one HolmesUsageEvents row. Best-effort: swallows DB errors. - Called from holmes.core.usage_recorder._fire on a daemon thread, so + 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 diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 01e9ca9eab..2ce24a8ef0 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -355,6 +355,49 @@ def duration_ms(self) -> int: """ 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_terminal(self, data: Dict[str, Any]) -> None: + """Pull cost/iterations/finish_reason from a terminal event's data.""" + metadata = data.get("metadata") or {} + costs = metadata.get("costs") or {} + if costs: + try: + self.stats = RequestStats(**costs) + except Exception: + logging.debug( + "Failed to materialize RequestStats from terminal event costs", + exc_info=True, + ) + # 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 + self.finish_reason = ( + metadata.get("finish_reason") or self.finish_reason + ) + + def _fire(self) -> None: + """Background-thread the dal write so the response path never blocks.""" + if self.dal is None or not getattr(self.dal, "enabled", False): + return + try: + threading.Thread( + target=self.dal.record_usage_event, + args=(self,), + daemon=True, + name="usage-recorder", + ).start() + except Exception: + # Defense in depth — record_usage_event has its own try/except too. + logging.exception("Failed to spawn usage recorder thread") + def stream_with_usage_recording( stream: Generator[StreamMessage, None, None], @@ -380,17 +423,17 @@ def stream_with_usage_recording( if msg.event == StreamEvents.TOOL_RESULT: state.tool_call_count += 1 elif msg.event == StreamEvents.ANSWER_END: - _capture_terminal(state, msg.data) + 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: - _capture_terminal(state, msg.data) + 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: - _capture_terminal(state, msg.data) + state._capture_terminal(msg.data) _inject_request_id(msg.data, state.request_id) state.status = RequestStatus.ERROR saw_terminal = True @@ -407,12 +450,16 @@ def stream_with_usage_recording( # filter incomplete runs out of "successful chat" metrics. if not saw_terminal and state.status == RequestStatus.SUCCESS: state.status = RequestStatus.ABORTED - _fire(state) + 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): @@ -421,28 +468,6 @@ def _inject_request_id(data: Dict[str, Any], request_id: str) -> None: md["request_id"] = request_id -def _capture_terminal(state: UsageRecorderState, data: Dict[str, Any]) -> None: - """Pull cost/iterations/finish_reason from a terminal event's data.""" - metadata = data.get("metadata") or {} - costs = metadata.get("costs") or {} - if costs: - try: - state.stats = RequestStats(**costs) - except Exception: - logging.debug( - "Failed to materialize RequestStats from terminal event costs", - exc_info=True, - ) - # Explicit None-check rather than `or` so a legitimate 0 (unlikely but - # not impossible) is preserved instead of falling back to state.iterations. - raw_iterations = data.get("num_llm_calls") - if raw_iterations is not None: - state.iterations = raw_iterations - state.finish_reason = ( - metadata.get("finish_reason") or state.finish_reason - ) - - def record_from_llm_result( state: UsageRecorderState, llm_result: "LLMResult", @@ -471,7 +496,7 @@ def record_from_llm_result( 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 - _fire(state) + state._fire() def record_error(state: UsageRecorderState, exc: Exception) -> None: @@ -481,23 +506,7 @@ def record_error(state: UsageRecorderState, exc: Exception) -> None: state.status = RequestStatus.RATE_LIMITED else: state.status = RequestStatus.ERROR - _fire(state) - - -def _fire(state: UsageRecorderState) -> None: - """Background-thread the dal write so the response path never blocks.""" - if state.dal is None or not getattr(state.dal, "enabled", False): - return - try: - threading.Thread( - target=state.dal.record_usage_event, - args=(state,), - daemon=True, - name="usage-recorder", - ).start() - except Exception: - # Defense in depth — record_usage_event has its own try/except too. - logging.exception("Failed to spawn usage recorder thread") + state._fire() __all__ = [ From 1921789c8906875af54a6f7823dcbbaa8e54ca8b Mon Sep 17 00:00:00 2001 From: alonelish Date: Thu, 7 May 2026 10:42:30 +0300 Subject: [PATCH 20/21] Capture partial token costs from TOKEN_COUNT events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Before this change: when the agentic loop raised an exception after some iterations had already succeeded (e.g. rate-limit on iteration 3 after 1 and 2 burned real tokens), the recorder fired a row with stats=None and status=error — the tokens those successful iterations consumed were silently lost from analytics. Root cause: call_stream accumulates `stats` in a local variable inside the generator function. The accumulated total only escapes via the terminal event's metadata.costs. When an exception unwinds the frame, the local `stats` is GC'd; the recorder has nothing to record. But the data is already on the wire — call_stream emits a TOKEN_COUNT event after every successful iteration with metadata.costs set to the running cumulative total. The wrapper just wasn't listening. Fix: one new branch in stream_with_usage_recording for TOKEN_COUNT that calls the freshly extracted state._capture_costs(). Each TOKEN_COUNT overwrites state.stats with the latest cumulative figure, so when the loop later raises, state.stats holds iterations-1..N-1's cost. Also benefits the client-disconnect case (was: status=aborted with $0; now: status=aborted with real partial cost). Refactor: split _capture_terminal into _capture_costs (pure cost extraction, called from both TOKEN_COUNT and terminal branches) + the iteration / finish_reason extraction (terminal-only). Tests: - test_token_count_event_captures_cumulative_costs — verifies TOKEN_COUNT updates state.stats and ANSWER_END's terminal capture works alongside it (last cumulative wins; equal in success case). - test_partial_costs_captured_when_loop_raises_mid_iteration — the case the user asked about: 2 iterations succeed (TOKEN_COUNT each), iteration 3 raises before its TOKEN_COUNT, recorder fires with partial cost preserved and status=error. - test_partial_costs_captured_on_client_disconnect — same partial capture for the abort path; status=aborted with real cost (was $0 before this change). 149 pass (was 146). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/usage_recorder.py | 51 ++++++++++++---- tests/core/test_usage_recorder.py | 96 +++++++++++++++++++++++++++++++ 2 files changed, 137 insertions(+), 10 deletions(-) diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 2ce24a8ef0..337efd56d8 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -361,24 +361,44 @@ def duration_ms(self) -> int: # — 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_terminal(self, data: Dict[str, Any]) -> None: - """Pull cost/iterations/finish_reason from a terminal event's data.""" + 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 costs: - try: - self.stats = RequestStats(**costs) - except Exception: - logging.debug( - "Failed to materialize RequestStats from terminal event costs", - exc_info=True, - ) + 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 ) @@ -422,6 +442,17 @@ def stream_with_usage_recording( 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) diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py index 98e3c1565f..817134e4ba 100644 --- a/tests/core/test_usage_recorder.py +++ b/tests/core/test_usage_recorder.py @@ -221,6 +221,102 @@ def failing_stream(): # 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() From e629689d0866833b9b29cf0b878f375f427eab8e Mon Sep 17 00:00:00 2001 From: alonelish Date: Sun, 10 May 2026 17:11:48 +0300 Subject: [PATCH 21/21] Replace per-request Thread with bounded ThreadPoolExecutor for recorder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review feedback (Moshe Morad on PR #1969): _fire previously spawned a fresh OS thread per HolmesUsageEvents write. Switch to a module-level ThreadPoolExecutor(max_workers=4, thread_name_prefix="usage-recorder") that all states share. Net behavior changes: - Caps concurrent recorder writes at 4 (caps blast radius if Supabase is slow or supabase-py's connection pool is contended). - Removes per-request thread-spawn overhead (~50–200μs). - Cleaner shutdown semantics — Python's atexit handler drains live executors, so we lose fewer rows on graceful shutdown than the previous daemon-thread fire-and-forget shape (which was dropped abruptly on process exit). - Under burst, additional submissions queue inside the executor rather than spawning unbounded threads. Threads are spawned lazily on first submit, so importing the module doesn't start any. max_workers=4 is sized 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 rates and below supabase-py's default pool size, so we won't starve foreground writes. submit() raises RuntimeError if the executor is shut down (process exiting); _fire catches that explicitly and accepts the loss (same fate as in-flight rows on the previous daemon-thread shape). Tests: _patch_inline_thread now stubs the executor with an inline submit() that mimics ThreadPoolExecutor.submit's swallow-target-exception semantic (real executor parks them on the Future). Added one new test covering the executor-shutdown branch. 156 pass (was 155). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: alonelish --- holmes/core/usage_recorder.py | 46 +++++++++++++---- tests/core/test_usage_recorder.py | 82 +++++++++++++++++++++---------- 2 files changed, 94 insertions(+), 34 deletions(-) diff --git a/holmes/core/usage_recorder.py b/holmes/core/usage_recorder.py index 337efd56d8..bcde3e238d 100644 --- a/holmes/core/usage_recorder.py +++ b/holmes/core/usage_recorder.py @@ -12,9 +12,9 @@ import logging import re -import threading 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 @@ -28,6 +28,22 @@ 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" @@ -404,19 +420,31 @@ def _capture_terminal(self, data: Dict[str, Any]) -> None: ) def _fire(self) -> None: - """Background-thread the dal write so the response path never blocks.""" + """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: - threading.Thread( - target=self.dal.record_usage_event, - args=(self,), - daemon=True, - name="usage-recorder", - ).start() + _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 spawn usage recorder thread") + logging.exception("Failed to submit usage recorder write") def stream_with_usage_recording( diff --git a/tests/core/test_usage_recorder.py b/tests/core/test_usage_recorder.py index 817134e4ba..de01df81fe 100644 --- a/tests/core/test_usage_recorder.py +++ b/tests/core/test_usage_recorder.py @@ -67,24 +67,34 @@ def _terminal_data(costs: dict, num_llm_calls: int = 1, finish_reason: str = "st def _patch_inline_thread(monkeypatch): - """Replace threading.Thread inside usage_recorder so target() runs inline. - - _fire spawns the recorder thread with ``args=(state,)`` (positional - state arg, since the DAL takes a single state object now). Mirror that - exactly — pass through both args and kwargs to be tolerant of either. + """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 _InlineThread: - def __init__(self, target=None, args=None, kwargs=None, daemon=None, name=None): - self._target = target - self._args = args or () - self._kwargs = kwargs or {} - - def start(self): - self._target(*self._args, **self._kwargs) + 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.threading, "Thread", _InlineThread) + monkeypatch.setattr(mod, "_RECORDER_EXECUTOR", _InlineExecutor()) def _state_arg(state: UsageRecorderState) -> UsageRecorderState: @@ -505,18 +515,20 @@ def test_no_thread_spawned_when_dal_is_none(self, monkeypatch): # ────────────────────────────────────────────────────────────────── -# Real Thread mode (no inline patching) — verifies fire-and-forget +# Real executor mode (no inline patching) — verifies fire-and-forget # ────────────────────────────────────────────────────────────────── class TestFireAndForgetThreadMode: - def test_record_calls_dal_in_background_thread(self): + 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 thread → state arrives positionally (one arg). Match the + # Real executor → state arrives positionally. Match the # production call signature. def slow_record(state): time.sleep(0.05) @@ -526,21 +538,41 @@ def slow_record(state): dal.record_usage_event = slow_record state = _make_state(dal=dal) record_error(state, RuntimeError("x")) - # Caller returns immediately; the thread is still running. + # Caller returns immediately; the executor worker is still running. # Wait briefly for it to finish. - assert called.wait(timeout=2.0), "background thread did not run record_usage_event" + 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): - # _fire wraps Thread.start() in a try/except so even with inline-thread - # patching (where target runs synchronously in start()) downstream - # exceptions don't bubble out to the caller. Logged via - # logging.exception(), but the caller is unaffected. This is the - # defense-in-depth contract — telemetry must never break the response. + """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. The inner try/except in _fire swallows it. + # 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"))