diff --git a/src/langbot/pkg/agent/runner/config_schema.py b/src/langbot/pkg/agent/runner/config_schema.py index aa7d13f7e..1fad3193f 100644 --- a/src/langbot/pkg/agent/runner/config_schema.py +++ b/src/langbot/pkg/agent/runner/config_schema.py @@ -16,6 +16,65 @@ KB_SELECTOR_TYPES = {'knowledge-base-multi-selector'} PROMPT_EDITOR_TYPES = {'prompt-editor'} NONE_SENTINELS = {'', '__none__', '__none'} +# Selector field types that pick a Host resource. The Runner config form is +# plugin-declared and open-ended (DynamicFormItemSchema allows extra types), so +# this is a registry, not a fixed list: a new selector type is one line here and +# is then reported by telemetry without touching the extraction call sites. +RESOURCE_SELECTOR_TYPES: dict[str, str] = { + 'knowledge-base-multi-selector': 'knowledge_bases', + 'tool-multi-selector': 'tools', + 'mcp-server-multi-selector': 'mcp_servers', + 'skill-multi-selector': 'skills', +} + + +def declared_resource_selectors(descriptor: RunnerDescriptor | None) -> dict[str, list[str]]: + """Map each selected resource kind to the schema field names that select it. + + Only the descriptor is read, so this reports what a Runner *can* select even + when the current configuration leaves the field empty. + """ + declared: dict[str, list[str]] = {} + if descriptor is None: + return declared + for item in descriptor.config_schema or []: + if not isinstance(item, dict): + continue + kind = RESOURCE_SELECTOR_TYPES.get(normalize_schema_item_type(item.get('type'))) + field_name = item.get('name') + if not kind or not isinstance(field_name, str) or not field_name: + continue + declared.setdefault(kind, []) + if field_name not in declared[kind]: + declared[kind].append(field_name) + return declared + + +def extract_selected_resources( + descriptor: RunnerDescriptor | None, + runner_config: dict[str, typing.Any] | None, +) -> dict[str, list[str]]: + """Extract the resource ids selected through schema-declared selector fields. + + Field names are never hard-coded: they come from the Runner's own config + schema, so a form that renames or adds a selector keeps working. + """ + selected: dict[str, list[str]] = {} + if not isinstance(runner_config, dict): + return selected + for kind, field_names in declared_resource_selectors(descriptor).items(): + values: list[str] = [] + for field_name in field_names: + value = runner_config.get(field_name) + if isinstance(value, str): + if value not in NONE_SENTINELS: + values.append(value) + elif isinstance(value, (list, tuple)): + values.extend(item for item in value if isinstance(item, str) and item not in NONE_SENTINELS) + if values: + selected[kind] = list(dict.fromkeys(values)) + return selected + def normalize_schema_item_type(item_type: typing.Any) -> typing.Any: """Normalize legacy/frontend DynamicForm aliases to protocol field types.""" diff --git a/src/langbot/pkg/agent/runner/orchestrator.py b/src/langbot/pkg/agent/runner/orchestrator.py index c7c9443eb..b04ec8394 100644 --- a/src/langbot/pkg/agent/runner/orchestrator.py +++ b/src/langbot/pkg/agent/runner/orchestrator.py @@ -37,6 +37,7 @@ from .run_journal import AgentRunJournal from .session_registry import AgentRunSessionRegistry, get_session_registry from .state_scope import build_state_context from ...provider.tools.loaders import skill as skill_loader +from ...telemetry import resources as telemetry_resources from ...telemetry import trace as trace_mod from ...telemetry.execution import bind_trace as bind_execution_trace from ...telemetry.execution import close_trace as close_execution_trace @@ -142,6 +143,22 @@ class AgentRunOrchestrator: binding=binding, descriptor=descriptor, ) + # Telemetry: report what this run was granted (never what it used). The + # materialized resource list is authoritative for tools/skills/models; + # the Runner's own config form is authoritative for its selectors. + try: + from . import config_schema as runner_config_schema + + capabilities = getattr(descriptor, 'capabilities', None) + telemetry_resources.note_agent( + runner_id=descriptor.id, + resources=resources, + selectors=runner_config_schema.extract_selected_resources(descriptor, binding.runner_config), + capabilities=capabilities.model_dump() if hasattr(capabilities, 'model_dump') else None, + tools_all=getattr(binding.resource_policy, 'allow_all_tools', None), + ) + except Exception: + pass context = await self.context_builder.build_context_from_event( event=execution_event, diff --git a/src/langbot/pkg/api/http/service/agent.py b/src/langbot/pkg/api/http/service/agent.py index 260b1f2df..5cfcf116d 100644 --- a/src/langbot/pkg/api/http/service/agent.py +++ b/src/langbot/pkg/api/http/service/agent.py @@ -348,25 +348,42 @@ class AgentService: execution_events.append(visible_result) final_text = '' - async for output in self.ap.agent_run_orchestrator.run( - event, - binding, - adapter_context={ - '_query': execution_query, - '_execution_context': execution_context, - '_result_observer': observe_result, - }, + # A debug run has no platform event: the synthetic event envelope is the + # execution's inbound boundary, so the chain is opened here with the real + # event id and marked as a debug WebUI/API execution. The Runner inside + # then joins this chain instead of becoming a trace origin of its own. + from ....telemetry.execution import ingress as execution_ingress + + principal_type = str(getattr(getattr(context, 'principal', None), 'principal_type', '') or '') + origin = 'api' if principal_type == 'api_key' else 'webui' + with execution_ingress( + self.ap, + 'event_done', + execution_context, + event_id, + debug=True, + origin=origin, + synthetic_event=event_type, ): - output_text = self._provider_output_to_text(output) - if output_text: - final_text = output_text - output_items.append( - { - 'kind': output.__class__.__name__, - 'role': str(getattr(output, 'role', '') or ''), - 'text': output_text, - } - ) + async for output in self.ap.agent_run_orchestrator.run( + event, + binding, + adapter_context={ + '_query': execution_query, + '_execution_context': execution_context, + '_result_observer': observe_result, + }, + ): + output_text = self._provider_output_to_text(output) + if output_text: + final_text = output_text + output_items.append( + { + 'kind': output.__class__.__name__, + 'role': str(getattr(output, 'role', '') or ''), + 'text': output_text, + } + ) return { 'event_id': event_id, diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index 4f51b15aa..a18ebd914 100644 --- a/src/langbot/pkg/box/service.py +++ b/src/langbot/pkg/box/service.py @@ -20,7 +20,6 @@ from langbot_plugin.box.security import BOX_SHARED_WORKSPACE_PROBE_PREFIX from .admission import SandboxAdmissionController, require_cloud_admission_policy from .connector import BoxRuntimeConnector, _get_box_config from . import secure_fs -from ..telemetry import features as telemetry_features from ..utils import httpclient from ..api.http.context import ExecutionContext from ..api.http.service.tenant import TenantContext, require_workspace_uuid @@ -636,7 +635,6 @@ class BoxService: f'query_id={query.query_id} ' f'summary={json.dumps(self._summarize_result(result), ensure_ascii=False)}' ) - telemetry_features.increment(query, 'sandbox', 'execs') return self._serialize_result(result) def resolve_box_session_id(self, query: pipeline_query.Query) -> str: @@ -2011,7 +2009,6 @@ class BoxService: # ── Observability ───────────────────────────────────────────────── def _record_error(self, exc: Exception, query: pipeline_query.Query): - telemetry_features.increment(query, 'sandbox', 'errors') self._recent_errors.append( { 'timestamp': _dt.datetime.now(_UTC).isoformat(), diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index 7da4a0a15..cf3ecd13a 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -309,6 +309,14 @@ class Application: if runtime_ops_reporter is not None: runtime_ops_stats = runtime_ops_reporter.get_stats() + telemetry_stats = {} + telemetry_health = getattr(self.telemetry, 'health', None) + if callable(telemetry_health): + try: + telemetry_stats = telemetry_health() + except Exception: + telemetry_stats = {} + return { 'asyncio_tasks': asyncio_tasks, 'event_loop': self.event_loop_monitor.snapshot(), @@ -320,7 +328,7 @@ class Application: 'models': model_stats, 'runtimes': runtime_stats, 'runtime_ops': runtime_ops_stats, - 'telemetry_tasks': len(getattr(self.telemetry, 'send_tasks', ())), + 'telemetry': telemetry_stats, } async def initialize(self): diff --git a/src/langbot/pkg/pipeline/pipelinemgr.py b/src/langbot/pkg/pipeline/pipelinemgr.py index c98b94c8a..ad45b4c50 100644 --- a/src/langbot/pkg/pipeline/pipelinemgr.py +++ b/src/langbot/pkg/pipeline/pipelinemgr.py @@ -194,6 +194,19 @@ class RuntimePipeline: query.variables['_pipeline_bound_mcp_servers'] = self.bound_mcp_servers query.variables['_pipeline_mcp_resource_attachments'] = self.mcp_resource_attachments query.variables['_pipeline_mcp_resource_agent_read_enabled'] = self.mcp_resource_agent_read_enabled + # Telemetry: the bindings this lane runs with. None means "all enabled", + # which the *_all flags make explicit in the chain record. + try: + from ..telemetry import resources as telemetry_resources + + telemetry_resources.note_pipeline( + plugins=self.bound_plugins, + plugins_all=self.enable_all_plugins, + mcp_servers=self.bound_mcp_servers, + mcp_all=self.enable_all_mcp_servers, + ) + except Exception: + pass # Record query start for monitoring try: @@ -370,7 +383,38 @@ class RuntimePipeline: i += 1 + @staticmethod + def _synthetic_origin(query: pipeline_query.Query) -> str: + """Origin of a lane that starts without an inbound platform event.""" + try: + variables = getattr(query, 'variables', None) + declared = str(variables.get('_telemetry_origin') or '').strip() if isinstance(variables, dict) else '' + except Exception: + declared = '' + if declared in ('webui', 'api'): + return declared + adapter = getattr(query, 'adapter', None) + if adapter is not None and adapter.__class__.__name__ == 'WebSocketAdapter': + # WebChat and the CLI diagnostic surface both speak through the + # WebSocket proxy adapter; the CLI marks itself in its variables. + try: + variables = getattr(query, 'variables', None) + if isinstance(variables, dict) and '_cli_run_status' in variables: + return 'api' + except Exception: + pass + return 'webui' + return 'api' + + @staticmethod + def _synthetic_event_type(query: pipeline_query.Query) -> str: + """Event type of the virtual inbound event a synthetic lane stands for.""" + event = getattr(query, 'message_event', None) + event_type = str(getattr(event, 'type', None) or '').strip() + return event_type[:160] or 'message.received' + async def process_query(self, query: pipeline_query.Query): + from ..telemetry import trace as trace_mod from ..telemetry.execution import ingress from ..telemetry.execution import record as record_execution from ..telemetry.platform import processing_mode @@ -380,10 +424,21 @@ class RuntimePipeline: # The Workspace-scoped opaque query uuid is the execution identity Space # shows for this lane; it is a stable UUID for every pooled query. execution_id = str(getattr(query, 'query_uuid', '') or '').strip() or str(query.query_id) + # A lane with no inbound platform event in flight (WebChat, WebUI debug, + # HTTP pipeline run) synthesizes its virtual inbound event here, so every + # execution starts at an event boundary and the Runner is never the trace + # origin. A lane running under a platform route nests inside that ingress. + synthesizing = trace_mod.current() is None try: - # Callers without a platform event (Webchat, HTTP API) still get one - # trace for the whole Pipeline lane; nested calls reuse the trace. - with ingress(self.ap, 'pipeline_done', getattr(self, 'execution_context', None), execution_id=execution_id): + with ingress( + self.ap, + 'pipeline_done', + getattr(self, 'execution_context', None), + execution_id=execution_id, + debug=synthesizing, + origin=self._synthetic_origin(query) if synthesizing else 'platform', + synthetic_event=self._synthetic_event_type(query) if synthesizing else '', + ): # The lane is its own workflow step: acknowledgements and replies it # sends directly, and the runner it drives, hang under this node. # When the lane runs under a platform route the node nests there. @@ -576,14 +631,8 @@ class RuntimePipeline: state = current_trace() # The runner orchestrator already lands its own failed ``runner/execute`` # node when the break happens inside a runner; in that case only carry the - # reason over instead of repeating the node on the trace. - already_recorded = bool( - state is not None - and any( - stage.get('family') == 'runner' and stage.get('outcome') not in ('success', 'skipped') - for stage in state.stages - ) - ) + # reason over instead of repeating the node on the chain. + already_recorded = bool(state is not None and state.runner_failed) if not already_recorded: # The lane failure is its own workflow step, hanging under whatever # step routed into it (or a root when nothing did). diff --git a/src/langbot/pkg/pipeline/preproc/preproc.py b/src/langbot/pkg/pipeline/preproc/preproc.py index 67db3b1a4..2a0878170 100644 --- a/src/langbot/pkg/pipeline/preproc/preproc.py +++ b/src/langbot/pkg/pipeline/preproc/preproc.py @@ -108,6 +108,12 @@ class PreProcessor(stage.PipelineStage): ) catalog = ResourcePolicyProjector.filter_tools(catalog, policy) ToolManager.bind_query_tool_sources(query, catalog) + try: + from ...telemetry import resources as telemetry_resources + + telemetry_resources.note_pipeline(tools=catalog, tools_all=policy.allow_all_tools) + except Exception: + pass return ToolManager.tools_from_catalog(catalog) def _runner_accepts_multimodal_input(self, descriptor: RunnerDescriptor | None) -> bool: @@ -393,6 +399,18 @@ class PreProcessor(stage.PipelineStage): descriptor, runner_config, ) + # Telemetry: the knowledge bases this lane may retrieve from, plus the + # schema-declared selector fields that selected them. + try: + from ...telemetry import resources as telemetry_resources + + telemetry_resources.note_pipeline( + knowledge_bases=query.variables['_knowledge_base_uuids'], + runner_id=descriptor.id if descriptor is not None else '', + selectors=config_schema.extract_selected_resources(descriptor, runner_config), + ) + except Exception: + pass # Emit PromptPreProcessing before the runner receives the query. @@ -427,4 +445,10 @@ class PreProcessor(stage.PipelineStage): bound_skills = extensions_prefs['skills'] query.variables['_pipeline_bound_skills'] = bound_skills + try: + from ...telemetry import resources as telemetry_resources + + telemetry_resources.note_pipeline(skills=bound_skills) + except Exception: + pass return entities.StageProcessResult(result_type=entities.ResultType.CONTINUE, new_query=query) diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index 48501292b..7b5a83afd 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -1464,9 +1464,6 @@ class RuntimeMCPSession: if query is None: return try: - from langbot.pkg.telemetry import features as telemetry_features - - telemetry_features.increment(query, 'mcp_resource_reads', envelope.get('source') or 'unknown') query.variables.setdefault(MCP_RESOURCE_TRACE_QUERY_KEY, []).append( { 'server_name': envelope.get('server_name'), diff --git a/src/langbot/pkg/provider/tools/toolmgr.py b/src/langbot/pkg/provider/tools/toolmgr.py index b09dff9d8..26329dad0 100644 --- a/src/langbot/pkg/provider/tools/toolmgr.py +++ b/src/langbot/pkg/provider/tools/toolmgr.py @@ -480,8 +480,6 @@ class ToolManager: query: pipeline_query.Query, source_ref: ToolSourceRef | None = None, ) -> typing.Any: - from langbot.pkg.telemetry import features as telemetry_features - source_ref = source_ref or self.get_query_tool_source(query, name) if source_ref is not None: execution_context = get_query_execution_context(query) @@ -534,7 +532,6 @@ class ToolManager: return await loader.invoke_tool(name, parameters, query, source_id=source_id) return await loader.invoke_tool(name, parameters, query) - telemetry_features.increment(query, 'tool_calls', telemetry_source) return await self._invoke_tool_with_monitoring( source=telemetry_source, name=name, @@ -547,7 +544,6 @@ class ToolManager: await self._bind_plugin_workspace(execution_context) sandbox_available = await self._workspace_sandbox_available(execution_context) if sandbox_available and await self.native_tool_loader.has_tool(name): - telemetry_features.increment(query, 'tool_calls', 'native') return await self._invoke_tool_with_monitoring( source='native', name=name, @@ -556,7 +552,6 @@ class ToolManager: invoke=lambda: self.native_tool_loader.invoke_tool(name, parameters, query), ) if await self.plugin_tool_loader.has_tool(name): - telemetry_features.increment(query, 'tool_calls', 'plugin') return await self._invoke_tool_with_monitoring( source='plugin', name=name, @@ -565,7 +560,6 @@ class ToolManager: invoke=lambda: self.plugin_tool_loader.invoke_tool(name, parameters, query), ) if await self.mcp_tool_loader.has_tool(execution_context, name): - telemetry_features.increment(query, 'tool_calls', 'mcp') return await self._invoke_tool_with_monitoring( source='mcp', name=name, @@ -574,7 +568,6 @@ class ToolManager: invoke=lambda: self.mcp_tool_loader.invoke_tool(name, parameters, query), ) if sandbox_available and await self.skill_tool_loader.has_tool(name): - telemetry_features.increment(query, 'tool_calls', 'skill') return await self._invoke_tool_with_monitoring( source='skill', name=name, diff --git a/src/langbot/pkg/telemetry/execution.py b/src/langbot/pkg/telemetry/execution.py index da0111f47..1e5a16a1c 100644 --- a/src/langbot/pkg/telemetry/execution.py +++ b/src/langbot/pkg/telemetry/execution.py @@ -1,18 +1,16 @@ -"""Content-free execution records for the existing telemetry sender. +"""Point-wise execution records for the existing telemetry sender. -This module reports observations only. Coverage catalogs and acceptance rules -belong to Space. +One execution is reported as a chain of independent records that all carry the +same ``event_id``: one ``execution_node`` per workflow node, emitted the moment +that node completes, plus exactly one ``execution_chain`` emitted when the +execution closes. The sender never assembles a trace locally, so there is no +per-trace buffer, no TTL and no per-process trace cap; memory is bounded by the +outbound record buffer, the flush batch size and the small cross-task registry. -One complete execution is one record: the execution identity (Pipeline query, -inbound platform event or Agent run), its bounded ordered stages and its -terminal outcome. Nothing is aggregated or counted across executions — every -record carries the id of the execution it belongs to, so Space can fetch the -whole chain by primary identity. Records are buffered whole and flushed as -``{"records": [...]}`` on the configured cadence and on shutdown; an in-flight -trace is never flushed half-written. - -Memory is bounded in stages per trace, buffered traces, buffered records, the -cross-task execution registry and flush batch size. +Sampling is decided exactly once per event, at ingress, from a stable hash of +the execution identity, so every record of one execution carries the same +decision. A chain that broke is always emitted even when sampling excluded it, +which is why a chain may arrive with zero or partial nodes. """ from __future__ import annotations @@ -21,29 +19,48 @@ import asyncio import contextlib import contextvars import json +import re import time +import zlib from datetime import datetime, timezone from . import trace as trace_mod from .identity import workspace_identity +from .resources import snapshot as resources_snapshot +from .resources import begin as resources_begin +from .resources import end as resources_end from .trace import TraceState MODES = frozenset({'pipeline', 'agent', 'event_processor', 'none'}) OUTCOMES = frozenset({'success', 'failed', 'cancelled', 'timeout', 'skipped', 'unknown'}) +FAMILIES = frozenset({'platform_event', 'event_route', 'pipeline', 'runner', 'platform_api'}) +ORIGINS = frozenset({'platform', 'webui', 'api'}) +CLOSED_BY = frozenset({'event_done', 'pipeline_done', 'runner_done', 'timeout'}) +CHAIN_FAILURES = frozenset({'failed', 'cancelled', 'timeout'}) -# Trace bounds: stages per trace, traces buffered per process, trace lifetime. -MAX_TRACES = 64 -TRACE_TTL_SECONDS = 120 +SCHEMA = 2 TRACE_MODES = frozenset({'off', 'failures', 'sampled', 'all'}) -DEFAULT_TRACE_MODE = 'sampled' +DEFAULT_TRACE_MODE = 'all' DEFAULT_TRACE_SAMPLE = 20 -# Record bounds: buffered records, one flush batch, flush cadence, registry. +# Sampling bounds: outbound record buffer, one flush batch, flush cadence. MAX_BUFFERED_RECORDS = 512 MAX_RECORDS_PER_FLUSH = 128 MAX_FLUSH_BYTES = 200 * 1024 DEFAULT_FLUSH_SECONDS = 180 +MAX_FLUSH_ATTEMPTS = 3 +FLUSH_RETRY_BACKOFF_SECONDS = 0.2 + +# Cross-task registry: execution id -> owning execution, bounded and time-limited. MAX_REGISTRY_KEYS = 256 +REGISTRY_TTL_SECONDS = 300 + +# Failure details are operator-facing: keep the failure class and a short +# machine-ish prefix, never the free text a user or a model produced. +ERROR_LIMIT = 200 +ERROR_PREFIX_CHARS = 120 +ERROR_PREFIX_TOKENS = 12 +_ERROR_CLASS = re.compile(r'[A-Za-z_][\w.]*') # Set by cross-task call sites (plugin/RPC actions) that run without the # ingress context but know which execution they belong to. @@ -68,16 +85,36 @@ def current_execution_id() -> str: return _execution_id.get() -def _bounded_error(value, limit: int = 400) -> str: - """Bound and de-fang one failure detail before it leaves the process. +def _short_prefix(text: str) -> str: + """Keep at most a short machine-ish prefix of one failure message.""" + if len(text) > ERROR_PREFIX_CHARS: + text = text[:ERROR_PREFIX_CHARS] + parts = text.split(' ') + if len(parts) > ERROR_PREFIX_TOKENS: + text = ' '.join(parts[:ERROR_PREFIX_TOKENS]) + return text - Error text is operator-facing: keep it single-line, printable and short so a - noisy failure cannot smuggle control characters or bloat the payload. + +def _bounded_error(value, limit: int = ERROR_LIMIT) -> str: + """Bound and redact one failure detail before it leaves the process. + + The failure class and a short machine-ish prefix survive; control characters + are stripped, whitespace is collapsed and the result is capped, so a noisy + failure cannot smuggle user or model content into the payload. """ if value is None: return '' text = ''.join(ch if ch.isprintable() else ' ' for ch in str(value)) text = ' '.join(text.split()) + if not text: + return '' + head, separator, rest = text.partition(':') + head = head.strip() + if separator and len(head) <= 80 and _ERROR_CLASS.fullmatch(head): + prefix = _short_prefix(rest.strip()) + text = f'{head}: {prefix}' if prefix else head + else: + text = _short_prefix(text) return text[:limit] @@ -91,18 +128,21 @@ def _duration_ms(started_at: str, ended_at: str) -> int: class ExecutionCounters: + """Point-wise execution reporting: one record per node, one per closure.""" + def __init__(self, manager): self.manager = manager self.records: list[dict] = [] self.task: asyncio.Task | None = None + # Exposed health: records dropped or accepted by the sender. self.dropped = 0 - self.traces: dict[str, TraceState] = {} - self.trace_deadlines: dict[str, float] = {} - # Execution identity -> in-flight trace, for stages recorded by tasks + self.sent = 0 + # Execution identity -> owning execution, for nodes recorded by tasks # that no longer share the ingress context (plugin/RPC platform calls). + # It resolves identity and the sampling decision only; node payloads are + # never accumulated here. self.registry: dict[str, TraceState] = {} self.registry_deadlines: dict[str, float] = {} - self.dropped_traces = 0 # ------------------------------------------------------------------ config @@ -130,6 +170,63 @@ class ExecutionCounters: seconds = DEFAULT_FLUSH_SECONDS return seconds if 1 <= seconds <= 86400 else DEFAULT_FLUSH_SECONDS + def health(self) -> dict: + """Real telemetry health for the instance liveness snapshot.""" + return { + 'records_buffered': len(self.records), + 'records_sent': self.sent, + 'records_dropped': self.dropped, + 'executions_in_flight': len({id(state) for state in self.registry.values()}), + } + + # ----------------------------------------------------------------- sampling + + def _sampling_decision(self, event_id: str) -> tuple[str, int, bool]: + """Decide once, deterministically, whether this event is reported. + + The identity is hashed instead of being parsed, so every id shape + (``platform:...``, uuids, debug ids) samples the same way. + """ + mode = self.trace_mode() + if mode == 'sampled': + denominator = self.trace_sample() + try: + sampled_in = zlib.crc32(str(event_id).encode('utf-8')) % denominator == 0 + except Exception: + sampled_in = False + return mode, denominator, sampled_in + # 'off' reports nothing, 'failures' keeps nodes out but still reports a + # broken chain, 'all' reports everything; denominator is 1 for all three. + return mode, 1, mode == 'all' + + def configure( + self, + state: TraceState, + context=None, + *, + debug: bool = False, + origin: str = 'platform', + identity=None, + ) -> None: + """Bind the execution identity, the sampling decision and the flags.""" + if state.configured: + return + state.configured = True + state.debug = bool(debug) + state.origin = origin if origin in ORIGINS else 'platform' + mode, denominator, sampled_in = self._sampling_decision(state.event_id) + state.mode = mode + state.denominator = denominator + # Debug chains are always emitted (unless reporting is off entirely). + state.emit_nodes = mode != 'off' and (state.debug or sampled_in) + if identity is not None: + state.identity = dict(identity) + elif context is not None: + try: + state.identity = workspace_identity(context) + except Exception: + pass + # --------------------------------------------------------------- recording def record( @@ -148,102 +245,96 @@ class ExecutionCounters: node: str = '', ): try: - cfg = self.manager.telemetry_config - if not cfg or cfg.get('disable_telemetry', False) or not cfg.get('url'): + if not self._enabled(): return - if family not in {'platform_event', 'event_route', 'pipeline', 'runner', 'platform_api'}: - return - if mode not in MODES or outcome not in OUTCOMES: + if family not in FAMILIES or mode not in MODES or outcome not in OUTCOMES: return # Only code-defined identifiers are accepted; never pass user values. if any(not isinstance(v, str) or len(v) > 160 for v in (operation, adapter, runner, node)): return state = trace_mod.current() if state is None: - # Cross-task work has no inherited trace; resolve the owning - # execution through the registry instead. + # Cross-task work has no inherited context; resolve the owning + # execution through the registry. It has no open step here, so + # its node becomes a root of that execution instead of being + # mirrored from another task's stack. lookup = str(execution_id or '').strip() or current_execution_id() if lookup: - state = self.registry.get(lookup) - if state is None: + state = self._resolve(lookup) + if state is None or state.closed: # Every observation belongs to exactly one execution chain. return + self.configure(state) try: identity = workspace_identity(context) except Exception: # A cross-task caller may only carry the Workspace id; reuse the - # identity already bound to the owning trace. + # identity already bound to the owning execution. identity = state.identity or {} if not identity.get('instance_id') or not identity.get('workspace_uuid'): return - self._record_trace_stage( - state, - identity=identity, - family=family, - operation=operation, - mode=mode, - adapter=adapter, - runner=runner, - outcome=outcome, - synthetic=synthetic, - error=_bounded_error(error), - node=node, + state.identity = identity + + ended_at = datetime.now(timezone.utc).isoformat() + node_id, seq, started_at, parent = trace_mod.resolve_node(state, str(node or ''), ended_at) + root = parent == '' and not state.root_seen + if root: + state.root_seen = True + bounded_error = _bounded_error(error) + state.note_node(outcome=outcome, adapter=adapter, runner=runner, family=family, error=bounded_error) + if not state.emit_nodes: + # Sampling excluded this execution: the node is never buffered, + # only counted so the terminal chain can explain the gap. + state.dropped_nodes += 1 + return + self._enqueue_record( + { + 'event_type': 'execution_node', + 'schema': SCHEMA, + **self._common_fields(state, ended_at), + 'node_id': node_id, + 'parent_node_id': parent, + 'root': root, + 'seq': seq, + 'family': family, + 'operation': operation, + 'mode': mode, + 'adapter': adapter, + 'runner': runner, + 'outcome': outcome, + 'error': bounded_error, + 'route_ref': trace_mod.current_route(), + 'run_id': trace_mod.current_run(), + 'synthetic': bool(synthetic), + 'started_at': started_at, + 'ended_at': ended_at, + } ) - self._ensure_loop() except Exception: # Observability must never change execution behavior. return - def _record_trace_stage( - self, - state: TraceState, - *, - identity, - family, - operation, - mode, - adapter, - runner, - outcome, - synthetic, - error: str = '', - node: str = '', - ) -> None: - if state.abandoned: - return - if trace_mod.current() is state: - # Recorded by the task that owns the step: its step stack is exact. - # A step's own record sits under its enclosing step, never itself. - parent = trace_mod.current_parent() - if node and parent == node: - parent = trace_mod.enclosing_parent() - else: - # A cross-task observation (plugin/RPC) cannot see the step - # contextvar; attach it to the innermost step the owning task has - # open on the trace. - parent = state.open_nodes[-1] if state.open_nodes else '' - if state.trace_id not in self.traces: - # Never buffer traces this configuration would discard anyway. - if self.trace_mode() == 'off' or len(self.traces) >= MAX_TRACES: - self.dropped += 1 - state.abandoned = True - self._unregister_state(state) - return - self.traces[state.trace_id] = state - self.trace_deadlines[state.trace_id] = time.monotonic() + TRACE_TTL_SECONDS - state.identity = identity - state.append( - family=family, - operation=operation, - mode=mode, - adapter=adapter, - runner=runner, - outcome=outcome, - synthetic=bool(synthetic), - error=error, - node=node, - parent=parent, - ) + def _common_fields(self, state: TraceState, timestamp: str) -> dict: + identity = state.identity + from ..utils import constants + + return { + 'event_id': state.event_id, + 'instance_id': identity['instance_id'], + 'workspace_uuid': identity['workspace_uuid'], + 'version': constants.semantic_version, + 'edition': constants.edition, + 'timestamp': timestamp, + 'debug': bool(state.debug), + 'sample': {'mode': state.mode, 'denominator': state.denominator}, + } + + def _enabled(self) -> bool: + try: + cfg = self.manager.telemetry_config + except Exception: + return False + return bool(cfg) and not cfg.get('disable_telemetry', False) and bool(cfg.get('url')) def _ensure_loop(self) -> None: if self.task is None or self.task.done(): @@ -252,20 +343,27 @@ class ExecutionCounters: # --------------------------------------------------------------- registry def register_trace(self, state: TraceState, execution_id) -> None: - """Alias one more execution identity onto an in-flight trace.""" + """Alias one more execution identity onto an in-flight execution.""" try: exec_id = str(execution_id or '').strip() - if not exec_id or state.abandoned: + if not exec_id: return + self.configure(state) + now = time.monotonic() + self._sweep_registry(now) if exec_id not in self.registry and len(self.registry) >= MAX_REGISTRY_KEYS: oldest = next(iter(self.registry)) self.registry.pop(oldest, None) self.registry_deadlines.pop(oldest, None) self.registry[exec_id] = state - self.registry_deadlines[exec_id] = time.monotonic() + TRACE_TTL_SECONDS + self.registry_deadlines[exec_id] = now + REGISTRY_TTL_SECONDS except Exception: return + def _resolve(self, execution_id: str) -> TraceState | None: + self._sweep_registry(time.monotonic()) + return self.registry.get(execution_id) + def _unregister_state(self, state: TraceState) -> None: for exec_id in [key for key, value in self.registry.items() if value is state]: self.registry.pop(exec_id, None) @@ -277,26 +375,68 @@ class ExecutionCounters: self.registry.pop(exec_id, None) self.registry_deadlines.pop(exec_id, None) - # ------------------------------------------------------------------ traces + # ------------------------------------------------------------------ chains def close_trace(self, state: TraceState, reason: str = 'event_done', error: str = '') -> None: - """Buffer one complete execution record when it matches the sample.""" + """Emit the one terminal ``execution_chain`` record of this execution.""" try: if state.closed: return state.closed = True - registered = self.traces.pop(state.trace_id, None) is not None - self.trace_deadlines.pop(state.trace_id, None) self._unregister_state(state) - if not registered and not state.failure_reason: + if not self._enabled(): + # Nothing can be uploaded; never grow the buffer for a record + # the sender would refuse to send. return - if (not state.stages and not state.failure_reason) or not self._trace_emitted(state): + self.configure(state) + if not self._chain_emitted(state): return - record = self._build_record(state, reason, error) - if record is not None: - self._enqueue_record(record) + self._enqueue_record(self._build_chain_record(state, reason, error)) except Exception: return + finally: + # The snapshot was read while building the chain record; stop + # collecting so a late sibling task cannot leak into the next chain. + resources_end() + + def _chain_emitted(self, state: TraceState) -> bool: + if state.mode == 'off': + return False + if state.emit_nodes or state.debug: + return True + # A chain that broke is always uploaded, even when sampling excluded it. + return bool(state.failure_reason) or state.outcome() in CHAIN_FAILURES + + def _build_chain_record(self, state: TraceState, reason: str, error: str = '') -> dict: + ended_at = datetime.now(timezone.utc).isoformat() + record = { + 'event_type': 'execution_chain', + 'schema': SCHEMA, + **self._common_fields(state, ended_at), + 'closed_by': reason if reason in CLOSED_BY else 'event_done', + 'outcome': state.outcome(), + # First failure detail in this chain, so Space can show why it broke. + 'error': _bounded_error(error) + or _bounded_error(state.failure_reason) + or _bounded_error(state.first_error) + or '', + 'duration_ms': _duration_ms(state.started_at, ended_at), + 'stage_count': state.stage_count, + 'dropped_nodes': state.dropped_nodes, + 'origin': state.origin, + 'model_name': state.model_name or '', + 'adapter': state.adapter or state.node_adapter, + 'runner': state.runner or state.node_runner, + 'runner_category': state.runner_category or '', + 'pipeline_plugins': state.pipeline_plugins, + } + # Bound resources of this execution, when any lane reported them. The + # key is omitted entirely when nothing was noted, so the payload stays + # byte-identical for executions with no bindings. + resources = resources_snapshot() + if resources: + record['resources'] = resources + return record def _enqueue_record(self, record: dict) -> None: if len(self.records) >= MAX_BUFFERED_RECORDS: @@ -305,137 +445,52 @@ class ExecutionCounters: self.records.append(record) self._ensure_loop() - def _trace_emitted(self, state: TraceState) -> bool: - mode = self.trace_mode() - if mode == 'off': - return False - # A chain that failed must always be uploaded, even when no stage was - # recorded before it broke: the reason is the whole point of the trace. - if state.failure_reason: - return True - if mode == 'all': - return True - # A failed, cancelled or timed-out stage is always worth keeping. - if any(stage['outcome'] not in ('success', 'skipped') for stage in state.stages): - return True - # WebUI debug runs are deliberately traced; they stay flagged synthetic. - if state.synthetic: - return True - if mode == 'failures': - return False - try: - return int(state.trace_id[:8], 16) % self.trace_sample() == 0 - except ValueError: - return False - - def _build_record(self, state: TraceState, reason: str, error: str = '') -> dict | None: - from ..utils import constants - - identity = state.identity - if not identity.get('instance_id') or not identity.get('workspace_uuid'): - return None - ended_at = datetime.now(timezone.utc).isoformat() - stages = state.stages - observations = [] - for stage in stages: - observation = { - 'seq': stage['seq'], - 'family': stage['family'], - 'operation': stage['operation'], - 'mode': stage['mode'], - 'adapter': stage['adapter'], - 'runner': stage['runner'], - 'outcome': stage['outcome'], - 'synthetic': stage['synthetic'], - 'error': stage['error'], - 'route_ref': stage['route_ref'], - 'run_id': stage['run_id'], - 'first_seen': stage['first_seen'], - 'last_seen': stage['last_seen'], - 'trace_id': state.trace_id, - } - # Workflow position is optional: records outside any step omit both - # keys entirely, so their payloads stay byte-identical. - if stage.get('node'): - observation['node'] = stage['node'] - if stage.get('parent'): - observation['parent'] = stage['parent'] - observations.append(observation) - return { - 'event_type': 'feature_execution', - # The execution identity: one record == one chain == one Space row. - 'query_id': state.trace_id, - 'instance_id': identity['instance_id'], - 'workspace_uuid': identity['workspace_uuid'], - 'version': constants.semantic_version, - 'edition': constants.edition, - 'timestamp': ended_at, - 'trusted': True, - # First failure detail in this chain, so Space can show why it broke. - 'error': _bounded_error(error) - or state.failure_reason - or next((stage['error'] for stage in stages if stage['error']), ''), - # Legacy-compatible columns Space already maps; best effort. - 'duration_ms': _duration_ms(state.started_at, ended_at), - 'model_name': state.model_name or '', - 'adapter': state.adapter or next((stage['adapter'] for stage in stages if stage['adapter']), ''), - 'runner': state.runner or next((stage['runner'] for stage in stages if stage['runner']), ''), - 'runner_category': state.runner_category or '', - 'pipeline_plugins': state.pipeline_plugins, - 'features': { - 'schema': 1, - 'observations': observations, - 'trace': { - 'closed_by': reason, - 'started_at': state.started_at, - 'ended_at': ended_at, - 'route_ref': next((stage['route_ref'] for stage in stages if stage['route_ref']), ''), - 'run_id': next((stage['run_id'] for stage in stages if stage['run_id']), ''), - 'outcome': state.outcome(), - 'dropped_stages': state.dropped_stages, - }, - }, - } - - def _sweep_traces(self, now: float) -> None: - for trace_id in [key for key, deadline in self.trace_deadlines.items() if deadline <= now]: - state = self.traces.get(trace_id) - if state is not None: - self.close_trace(state, 'ttl') - # ----------------------------------------------------------------- flushing async def _loop(self): interval = self.flush_seconds() - last_flush = time.monotonic() - while self.records or self.traces or self.registry: - # Traces expire on a short TTL; records keep their own cadence. - await asyncio.sleep(1 if (self.traces or self.registry) else interval) - now = time.monotonic() - self._sweep_traces(now) - self._sweep_registry(now) - if self.records and now - last_flush >= interval: - last_flush = now + try: + # The outbound loop exists only to flush the record buffer. + while self.records: + await asyncio.sleep(interval) await self.flush() - self.task = None + finally: + self.task = None - async def flush(self): + async def flush(self) -> bool: + """Upload the buffered prefix; records stay buffered until accepted.""" if not self.records: - return + return True batch: list[dict] = [] size = 0 - while self.records and len(batch) < MAX_RECORDS_PER_FLUSH: - encoded = len(json.dumps(self.records[0], ensure_ascii=False, default=str)) + for record in self.records: + if len(batch) >= MAX_RECORDS_PER_FLUSH: + break + encoded = len(json.dumps(record, ensure_ascii=False, default=str)) if batch and size + encoded > MAX_FLUSH_BYTES: break - batch.append(self.records.pop(0)) + batch.append(record) size += encoded if not batch: - return + return True payload = {'records': batch} - if not await self.manager.send(payload): - await asyncio.sleep(1) - await self.manager.send(payload) + for attempt in range(MAX_FLUSH_ATTEMPTS): + try: + if await self.manager.send(payload): + # Peek-then-send: nothing is removed until the sender + # accepted the exact batch it was given. + del self.records[: len(batch)] + self.sent += len(batch) + return True + except asyncio.CancelledError: + raise + except Exception: + pass + if attempt + 1 < MAX_FLUSH_ATTEMPTS: + await asyncio.sleep(FLUSH_RETRY_BACKOFF_SECONDS * (2**attempt)) + # The batch stays buffered for the next cycle; a permanently unaccepted + # record is dropped by the buffer bound above and counted there. + return False async def shutdown(self): if self.task is not None: @@ -451,12 +506,8 @@ class ExecutionCounters: break except (Exception, asyncio.CancelledError): pass + self.dropped += len(self.records) self.records.clear() - # Executions still in flight cannot be completed during shutdown. - self.dropped += len(self.traces) - self.dropped_traces += len(self.traces) - self.traces.clear() - self.trace_deadlines.clear() self.registry.clear() self.registry_deadlines.clear() @@ -477,7 +528,7 @@ def record(ap, context, **observation): def close_trace(ap, state, reason: str = 'event_done', error: str = '') -> None: - """Best-effort trace close usable with optional telemetry and test doubles.""" + """Best-effort chain close usable with optional telemetry and test doubles.""" try: counters = _execution_counters(ap) if counters is not None: @@ -487,8 +538,10 @@ def close_trace(ap, state, reason: str = 'event_done', error: str = '') -> None: def bind_trace(ap, execution_id: str | None = None) -> trace_mod.TraceBinding: - """Bind an execution trace and alias its identity for cross-task stages.""" + """Bind an execution and alias its identity for cross-task nodes.""" binding = trace_mod.bind(execution_id) + if binding.created: + resources_begin() if execution_id: try: counters = _execution_counters(ap) @@ -500,36 +553,61 @@ def bind_trace(ap, execution_id: str | None = None) -> trace_mod.TraceBinding: @contextlib.contextmanager -def ingress(ap, reason: str = 'event_done', context=None, execution_id: str | None = None): - """Trace one execution boundary; only the owner of the trace closes it. +def ingress( + ap, + reason: str = 'event_done', + context=None, + execution_id: str | None = None, + *, + debug: bool = False, + origin: str = 'platform', + synthetic_event: str = '', +): + """Open one execution boundary; only the owner of the chain closes it. Nested boundaries (an event routed into a Pipeline, a Runner invoked from a - dispatch) reuse the in-flight trace instead of starting a second one, and - only alias their own execution id onto it. + dispatch) reuse the in-flight chain instead of starting a second one, and + only alias their own execution id onto it. ``debug``/``origin`` apply when + this boundary starts the chain: a WebUI/API run without a real platform + event synthesizes its virtual inbound event here, so the chain has a real + identity and a node for the event it stands for. """ binding = trace_mod.bind(execution_id) - if binding.created and context is not None: - try: - binding.state.identity = workspace_identity(context) - except Exception: - pass - if execution_id and not binding.state.abandoned: - try: - counters = _execution_counters(ap) - if counters is not None: - counters.register_trace(binding.state, execution_id) - except Exception: - pass + state = binding.state + if binding.created: + resources_begin() + counters = _execution_counters(ap) + if counters is not None: + if binding.created: + counters.configure(state, context, debug=debug, origin=origin) + if execution_id: + counters.register_trace(state, execution_id) + scope = None + if binding.created and synthetic_event: + scope = trace_mod.stage_scope() + node = scope.__enter__() + record( + ap, + context, + family='platform_event', + operation=str(synthetic_event)[:160], + outcome='success', + synthetic=True, + node=node, + ) failure = '' try: yield binding except BaseException as exc: - # The boundary itself failed: keep the reason on the trace so the uploaded - # payload explains the failure instead of reporting a bare outcome. + # The boundary itself failed: keep the reason on the chain so the + # uploaded payload explains the failure instead of reporting a bare + # outcome, even when sampling had excluded this execution. if binding.created: failure = f'{type(exc).__name__}: {exc}' - binding.state.mark_failure(failure) + state.mark_failure(failure) raise finally: + if scope is not None: + scope.__exit__(None, None, None) if trace_mod.unbind_root(binding): - close_trace(ap, binding.state, reason, failure) + close_trace(ap, state, reason, failure) diff --git a/src/langbot/pkg/telemetry/features.py b/src/langbot/pkg/telemetry/features.py deleted file mode 100644 index ab66f9884..000000000 --- a/src/langbot/pkg/telemetry/features.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Per-query telemetry feature counters. - -Collects anonymous, content-free usage signals (tool call counts, knowledge -base usage, sandbox executions, ...) into ``query.variables`` during query -processing. The chat handler reads the accumulated dict when building the -telemetry payload and ships it as the ``features`` JSON object. - -Every helper here is defensive: telemetry must NEVER break the pipeline, so -all mutations are wrapped and failures are silently ignored. -""" - -from __future__ import annotations - -import typing - -if typing.TYPE_CHECKING: - import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query - - -FEATURES_KEY = '_telemetry_features' - - -def get_features(query: pipeline_query.Query) -> dict: - """Return the mutable features dict for this query, creating it if needed.""" - try: - return query.variables.setdefault(FEATURES_KEY, {}) - except Exception: - return {} - - -def increment(query: pipeline_query.Query, group: str, key: str | None = None, amount: int = 1) -> None: - """Increment a counter. - - ``increment(q, 'sandbox', 'execs')`` -> features['sandbox']['execs'] += 1 - ``increment(q, 'tool_call_rounds')`` -> features['tool_call_rounds'] += 1 - """ - try: - features = get_features(query) - if key is None: - features[group] = int(features.get(group, 0)) + amount - else: - nested = features.setdefault(group, {}) - if isinstance(nested, dict): - nested[key] = int(nested.get(key, 0)) + amount - except Exception: - pass - - -def set_value(query: pipeline_query.Query, group: str, value: typing.Any) -> None: - """Set a feature value (overwrites).""" - try: - get_features(query)[group] = value - except Exception: - pass - - -def collect_features(query: pipeline_query.Query) -> dict: - """Build the final ``features`` object for the telemetry payload. - - Combines the counters accumulated during processing with end-of-query - snapshots (activated skills, bound MCP servers). Returns a plain dict - that must be JSON-serializable; non-serializable values are dropped. - """ - features: dict = {} - try: - accumulated = query.variables.get(FEATURES_KEY) - if isinstance(accumulated, dict): - features.update(accumulated) - except Exception: - pass - - # Activated skills (names only, registered by the activate tool) - try: - activated = query.variables.get('_activated_skills', {}) - if isinstance(activated, dict) and activated: - features['activated_skills'] = sorted(activated.keys()) - except Exception: - pass - - # MCP servers bound to the pipeline (names only; None means "all enabled") - try: - bound_mcp = query.variables.get('_pipeline_bound_mcp_servers', None) - if bound_mcp is not None: - features['mcp_servers'] = list(bound_mcp) - except Exception: - pass - - # Drop anything that is not JSON-serializable - import json - - try: - json.dumps(features) - return features - except Exception: - safe: dict = {} - for k, v in features.items(): - try: - json.dumps({k: v}) - safe[k] = v - except Exception: - continue - return safe diff --git a/src/langbot/pkg/telemetry/heartbeat.py b/src/langbot/pkg/telemetry/heartbeat.py index 2499c7e06..7deb20400 100644 --- a/src/langbot/pkg/telemetry/heartbeat.py +++ b/src/langbot/pkg/telemetry/heartbeat.py @@ -13,6 +13,7 @@ credentials. from __future__ import annotations import asyncio +import time import typing from datetime import datetime, timezone @@ -36,10 +37,66 @@ class WorkspaceResourceSnapshot(typing.TypedDict): mcp_server_count: int extension_count: int skill_count: int + agent_count: int + event_processor_count: int + plugin_processor_count: int + event_binding_count: int adapters: list[str] execution_generation: int +# Agent rows are configuration, not runtime registries, so a Workspace count +# needs one scoped query. The heartbeat runs every five minutes on Cloud; a +# one-hour cache keeps the cost at one bounded query per Workspace per hour +# instead of one per cycle. Unavailable counts stay -1, never a fake zero. +_AGENT_COUNT_CACHE_TTL_SECONDS = 3600 +_agent_count_cache: dict[str, tuple[float, dict[str, int]]] = {} + + +async def _workspace_agent_kind_counts(ap: core_app.Application, workspace_uuid: str) -> dict[str, int]: + """Count one Workspace's Agents per ``kind``; empty when unavailable.""" + if not workspace_uuid: + return {} + now = time.monotonic() + cached = _agent_count_cache.get(workspace_uuid) + if cached is not None and cached[0] > now: + return dict(cached[1]) + counts: dict[str, int] = {} + try: + from ..entity.persistence import agent as persistence_agent + + statement = sqlalchemy.select(persistence_agent.Agent.kind, sqlalchemy.func.count()).group_by( + persistence_agent.Agent.kind + ) + tenant_uow = getattr(ap.persistence_mgr, 'tenant_uow', None) + if callable(tenant_uow): + async with tenant_uow(workspace_uuid): + result = await ap.persistence_mgr.execute_async(statement) + else: + result = await ap.persistence_mgr.execute_async(statement) + for kind, count in result.all(): + counts[str(kind or 'agent')] = int(count or 0) + except Exception: + counts = {} + _agent_count_cache[workspace_uuid] = (now + _AGENT_COUNT_CACHE_TTL_SECONDS, dict(counts)) + return counts + + +def _bot_processor_counts(bots) -> tuple[int, int]: + """Count plugin-processor subscriptions and event bindings across bots.""" + processors = 0 + bindings = 0 + for bot in bots: + entity = getattr(bot, 'bot_entity', bot) + declared_processors = getattr(entity, 'plugin_processors', None) + if isinstance(declared_processors, list): + processors += len(declared_processors) + declared_bindings = getattr(entity, 'event_bindings', None) + if isinstance(declared_bindings, list): + bindings += len(declared_bindings) + return processors, bindings + + async def _count( ap: core_app.Application, table, @@ -81,6 +138,10 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application, bindings) - 'mcp_server_count': 0, 'extension_count': 0, 'skill_count': 0, + 'agent_count': -1, + 'event_processor_count': -1, + 'plugin_processor_count': 0, + 'event_binding_count': 0, 'adapters': [], 'execution_generation': binding.placement_generation, } @@ -91,6 +152,9 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application, bindings) - for key, bot in getattr(ap.platform_mgr, '_bots_by_key', {}).items(): if len(key) >= 2 and key[1] in counts: counts[key[1]]['bot_count'] += 1 + processors, bindings = _bot_processor_counts([bot]) + counts[key[1]]['plugin_processor_count'] += processors + counts[key[1]]['event_binding_count'] += bindings adapter = getattr(bot, 'adapter', None) if adapter is not None and getattr(bot, 'enable', False): adapter_sets[key[1]].add(adapter.__class__.__name__) @@ -113,6 +177,10 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application, bindings) - for workspace_uuid, resource in counts.items(): resource['extension_count'] = resource['plugin_count'] + resource['mcp_server_count'] resource['adapters'] = sorted(adapter_sets[workspace_uuid]) + agent_counts = await _workspace_agent_kind_counts(ap, workspace_uuid) + if agent_counts: + resource['agent_count'] = int(agent_counts.get('agent', 0)) + resource['event_processor_count'] = int(agent_counts.get('event_processor', 0)) return list(counts.values()) @@ -163,6 +231,9 @@ async def build_heartbeat_payload( features['bot_count'] = len(platform_mgr.bots) adapters = sorted({bot.adapter.__class__.__name__ for bot in enabled_bots if getattr(bot, 'adapter', None)}) features['adapters'] = adapters + processors, bindings = _bot_processor_counts(platform_mgr.bots) + features['plugin_processor_count'] = processors + features['event_binding_count'] = bindings except Exception: pass @@ -206,6 +277,13 @@ async def build_heartbeat_payload( except Exception: pass + # Agent processors: configuration rows, so one scoped count per Workspace. + # The Cloud snapshot supplies these already and overrides below. + if workspace_resource is None: + agent_counts = await _workspace_agent_kind_counts(ap, workspace_uuid) + features['agent_count'] = int(agent_counts.get('agent', 0)) if agent_counts else -1 + features['event_processor_count'] = int(agent_counts.get('event_processor', 0)) if agent_counts else -1 + if workspace_resource is not None: features.update({key: value for key, value in workspace_resource.items() if key != 'workspace_uuid'}) diff --git a/src/langbot/pkg/telemetry/resources.py b/src/langbot/pkg/telemetry/resources.py new file mode 100644 index 000000000..5164e090e --- /dev/null +++ b/src/langbot/pkg/telemetry/resources.py @@ -0,0 +1,244 @@ +"""Bound-resource reporting for one execution. + +Answers "what was this execution given", never "what did it use": the ids and +names of the Plugins, MCP servers, knowledge bases, tools, skills and models +that were bound or granted to a Pipeline lane or an Agent run. Nothing here is +user content - only code-defined identifiers (uuids, plugin names, tool names) +and booleans. + +Two sources feed one execution-scoped snapshot: + +* the Host policy actually granted to the run (``ResourcePolicy`` / + materialized ``AgentResources``), which is authoritative for tools, skills and + platform APIs; +* the Runner's own configuration form (``DynamicFormItemSchema``), which is + where knowledge bases and models are selected. Selector field types are read + through :mod:`langbot.pkg.agent.runner.config_schema`, so a Runner that + declares a new selector type is picked up without a change here. + +The snapshot is attached to the terminal ``execution_chain`` record. It is +bounded twice - per list and in total - and it is best effort: telemetry must +never change execution behaviour, so every entry point swallows failures. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import json +import typing + +MAX_ITEMS_PER_LIST = 64 +MAX_NAME_CHARS = 160 +MAX_RESOURCES_BYTES = 8 * 1024 + +_resources: contextvars.ContextVar[dict[str, typing.Any] | None] = contextvars.ContextVar( + 'telemetry_resources', default=None +) + + +def begin() -> None: + """Start collecting for one execution (called when a chain is opened).""" + _resources.set({}) + + +def end() -> None: + """Stop collecting; the snapshot has already been read.""" + _resources.set(None) + + +def _bucket(section: str) -> dict[str, typing.Any] | None: + root = _resources.get() + if root is None: + return None + bucket = root.get(section) + if not isinstance(bucket, dict): + bucket = {} + root[section] = bucket + return bucket + + +def _names(values: typing.Any) -> list[str]: + """Bound one list of code-defined identifiers, preserving source order.""" + if not isinstance(values, (list, tuple)): + return [] + names: list[str] = [] + seen: set[str] = set() + for value in values: + name = '' + if isinstance(value, str): + name = value + elif isinstance(value, dict): + name = str(value.get('name') or value.get('kb_id') or value.get('tool_name') or '') + if not name: + continue + name = name[:MAX_NAME_CHARS] + if name in seen: + continue + seen.add(name) + names.append(name) + if len(names) >= MAX_ITEMS_PER_LIST: + break + return names + + +def _tool_refs(values: typing.Any) -> list[dict[str, str]]: + """Bound one list of tool grants, keeping their Host source.""" + if not isinstance(values, (list, tuple)): + return [] + refs: list[dict[str, str]] = [] + seen: set[str] = set() + for value in values: + if isinstance(value, str): + name, source, source_id = value, '', '' + elif isinstance(value, dict): + name = str(value.get('tool_name') or value.get('name') or '') + source = str(value.get('tool_type') or value.get('source') or '') + raw_id = value.get('source_id') + source_id = raw_id if isinstance(raw_id, str) else '' + else: + continue + name = name[:MAX_NAME_CHARS] + if not name or name in seen: + continue + seen.add(name) + ref = {'name': name} + if source: + ref['source'] = source[:64] + if source_id: + ref['source_id'] = source_id[:MAX_NAME_CHARS] + refs.append(ref) + if len(refs) >= MAX_ITEMS_PER_LIST: + break + return refs + + +def _set(bucket: dict[str, typing.Any] | None, key: str, value: typing.Any) -> None: + if bucket is None or value in (None, [], {}, ''): + return + bucket[key] = value + + +def note_pipeline( + *, + plugins: typing.Any = None, + plugins_all: typing.Any = None, + mcp_servers: typing.Any = None, + mcp_all: typing.Any = None, + knowledge_bases: typing.Any = None, + skills: typing.Any = None, + tools: typing.Any = None, + tools_all: typing.Any = None, + runner_id: str = '', + selectors: typing.Any = None, +) -> None: + """Record the resources one Pipeline lane is bound to. + + ``plugins``/``mcp_servers`` accept ``None`` to mean "all enabled", matching + the binding model; the ``*_all`` flags make that explicit in the payload. + """ + try: + bucket = _bucket('pipeline') + if bucket is None: + return + _set(bucket, 'plugins', _names(plugins)) + _set(bucket, 'mcp_servers', _names(mcp_servers)) + _set(bucket, 'knowledge_bases', _names(knowledge_bases)) + _set(bucket, 'skills', _names(skills)) + _set(bucket, 'tools', _tool_refs(tools)) + _set(bucket, 'runner_id', str(runner_id or '')[:MAX_NAME_CHARS]) + if isinstance(selectors, dict): + cleaned = { + str(kind)[:64]: _names(fields) + for kind, fields in selectors.items() + if isinstance(fields, (list, tuple)) and fields + } + _set(bucket, 'selectors', cleaned) + if plugins_all is not None: + bucket['plugins_all'] = bool(plugins_all) + if mcp_all is not None: + bucket['mcp_all'] = bool(mcp_all) + if tools_all is not None: + bucket['tools_all'] = bool(tools_all) + except Exception: + return + + +def note_agent( + *, + runner_id: str = '', + resources: typing.Any = None, + selectors: typing.Any = None, + capabilities: typing.Any = None, + tools_all: typing.Any = None, +) -> None: + """Record the resources one Agent run was granted. + + ``resources`` is the materialized ``AgentResources`` mapping the Host + actually handed to the Runner (tools carry their source), ``selectors`` + names the Runner config-form fields that selected a resource, and + ``capabilities`` records which resource categories the Runner declared. + """ + try: + bucket = _bucket('agent') + if bucket is None: + return + _set(bucket, 'runner_id', str(runner_id or '')[:MAX_NAME_CHARS]) + if isinstance(resources, dict): + _set(bucket, 'knowledge_bases', _names(resources.get('knowledge_bases'))) + _set(bucket, 'skills', _names(resources.get('skills'))) + _set(bucket, 'models', _names(resources.get('models'))) + _set(bucket, 'tools', _tool_refs(resources.get('tools'))) + _set(bucket, 'platform_apis', _names(resources.get('platform_tools'))) + if isinstance(selectors, dict): + cleaned = { + str(kind)[:64]: _names(fields) + for kind, fields in selectors.items() + if isinstance(fields, (list, tuple)) and fields + } + _set(bucket, 'selectors', cleaned) + if isinstance(capabilities, dict): + declared = {key: bool(value) for key, value in capabilities.items() if value} + _set(bucket, 'capabilities', declared) + if tools_all is not None: + bucket['tools_all'] = bool(tools_all) + except Exception: + return + + +def snapshot() -> dict[str, typing.Any] | None: + """Return the bounded snapshot of this execution; None when nothing was noted.""" + try: + root = _resources.get() + if not isinstance(root, dict) or not root: + return None + + def encoded_size(value: typing.Any) -> int: + return len(json.dumps(value, ensure_ascii=False, default=str).encode('utf-8')) + + if encoded_size(root) <= MAX_RESOURCES_BYTES: + return root + # Over budget: shed the widest detail (per-tool grants) first, then fall + # back to a marker, so an oversized binding can never block the chain + # record or bloat the upload. + reduced = json.loads(json.dumps(root, ensure_ascii=False, default=str)) + for section in ('pipeline', 'agent'): + bucket = reduced.get(section) + if isinstance(bucket, dict): + bucket.pop('tools', None) + reduced['truncated'] = True + if encoded_size(reduced) <= MAX_RESOURCES_BYTES: + return reduced + return {'truncated': True} + except Exception: + return None + + +@contextlib.contextmanager +def collected() -> typing.Iterator[None]: + """Bind a collection scope, for callers outside an execution boundary.""" + token = _resources.set({}) + try: + yield + finally: + _resources.reset(token) diff --git a/src/langbot/pkg/telemetry/telemetry.py b/src/langbot/pkg/telemetry/telemetry.py index c36b35e3c..9188a5c54 100644 --- a/src/langbot/pkg/telemetry/telemetry.py +++ b/src/langbot/pkg/telemetry/telemetry.py @@ -12,9 +12,6 @@ from ..utils import httpclient from .execution import ExecutionCounters -_MAX_INFLIGHT_TELEMETRY_TASKS = 8 - - class TelemetryManager: """TelemetryManager handles sending telemetry for a given application instance. @@ -27,36 +24,18 @@ class TelemetryManager: self.ap = ap self.telemetry_config: dict[str, typing.Any] = {} - self.send_tasks: list[asyncio.Task] = [] self._client: httpx.AsyncClient | None = None self.execution = ExecutionCounters(self) async def initialize(self): self.telemetry_config = self.ap.instance_config.data.get('space', {}) - async def start_send_task(self, payload: dict): - self.send_tasks = [task for task in self.send_tasks if not task.done()] - if len(self.send_tasks) >= _MAX_INFLIGHT_TELEMETRY_TASKS: - self.ap.logger.debug('Telemetry queue is full; dropping best-effort event') - return - task = asyncio.create_task(self.send(payload)) - self.send_tasks.append(task) - task.add_done_callback(self._send_task_done) - - def _send_task_done(self, task: asyncio.Task) -> None: - try: - self.send_tasks.remove(task) - except ValueError: - pass + def health(self) -> dict: + """Telemetry health for the instance liveness snapshot.""" + return self.execution.health() async def shutdown(self) -> None: await self.execution.shutdown() - tasks = list(self.send_tasks) - for task in tasks: - task.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - self.send_tasks.clear() if self._client is not None: await self._client.aclose() self._client = None @@ -95,45 +74,53 @@ class TelemetryManager: url = server.rstrip('/') + '/api/v1/telemetry' try: - # Sanitize payload so string fields are strings and not nulls + # Sanitize payload so string fields are strings and not nulls. + # A record batch is an envelope of whole execution records; it + # carries no top-level record fields, so only a single-record + # payload gets the legacy defaults. sanitized = dict(payload) - if 'query_id' in sanitized: - try: - sanitized['query_id'] = '' if sanitized['query_id'] is None else str(sanitized['query_id']) - except Exception: - sanitized['query_id'] = str(sanitized.get('query_id', '')) + if isinstance(sanitized.get('records'), list): + sanitized['records'] = [item for item in sanitized['records'] if isinstance(item, dict)] + else: + if 'query_id' in sanitized: + try: + sanitized['query_id'] = ( + '' if sanitized['query_id'] is None else str(sanitized['query_id']) + ) + except Exception: + sanitized['query_id'] = str(sanitized.get('query_id', '')) - for sfield in ( - 'adapter', - 'runner', - 'runner_category', - 'model_name', - 'version', - 'edition', - 'error', - 'timestamp', - 'event_type', - ): - if sfield not in sanitized: - continue - v = sanitized.get(sfield) - sanitized[sfield] = '' if v is None else str(v) + for sfield in ( + 'adapter', + 'runner', + 'runner_category', + 'model_name', + 'version', + 'edition', + 'error', + 'timestamp', + 'event_type', + ): + if sfield not in sanitized: + continue + v = sanitized.get(sfield) + sanitized[sfield] = '' if v is None else str(v) - # event_type defaults to 'query' for backward compatibility - if not sanitized.get('event_type'): - sanitized['event_type'] = 'query' + # event_type defaults to 'query' for backward compatibility + if not sanitized.get('event_type'): + sanitized['event_type'] = 'query' - # features must be a JSON object - if 'features' in sanitized and not isinstance(sanitized['features'], dict): - sanitized['features'] = {} + # features must be a JSON object + if 'features' in sanitized and not isinstance(sanitized['features'], dict): + sanitized['features'] = {} - if 'duration_ms' in sanitized: - try: - sanitized['duration_ms'] = ( - int(sanitized['duration_ms']) if sanitized['duration_ms'] is not None else 0 - ) - except Exception: - sanitized['duration_ms'] = 0 + if 'duration_ms' in sanitized: + try: + sanitized['duration_ms'] = ( + int(sanitized['duration_ms']) if sanitized['duration_ms'] is not None else 0 + ) + except Exception: + sanitized['duration_ms'] = 0 async with self._client_context() as client: try: diff --git a/src/langbot/pkg/telemetry/trace.py b/src/langbot/pkg/telemetry/trace.py index 445bb470e..fb4d18031 100644 --- a/src/langbot/pkg/telemetry/trace.py +++ b/src/langbot/pkg/telemetry/trace.py @@ -1,22 +1,22 @@ -"""Content-free per-execution traces. +"""Per-execution telemetry context for the v2 execution chain. -One trace identity is one execution identity: the id a Pipeline query, inbound -platform event or Agent run is known by, so a stored chain can be joined with -the Space view of that same execution. Stage records appended under it describe -how that one execution was routed and processed, using only code-defined -identifiers. Nothing here is aggregated: a trace is a bounded, ordered sequence -for one execution. +One execution identity (a Pipeline query, inbound platform event or Agent run) +is one chain. The chain is not assembled here: every workflow node is emitted +by ``telemetry.execution`` the moment it completes, carrying its own identity +and the identity of the execution it belongs to. This module only holds the +small execution-scoped context the emitter needs: -Three bounds keep memory independent of traffic volume: +* the execution state (identity, the once-per-event sampling decision, the + execution-scoped display fields); +* the routing/run identity pinned around the block that produced a node; +* the open-step stack, which lives in a ContextVar holding *one task's own + immutable tuple*. Child tasks inherit a snapshot at creation; no shared + mutable stack and no cross-task mirror exists, so sibling tasks can never + corrupt each other's parent links. -* ``MAX_STAGES`` stages per trace (further stages are counted, not kept); -* ``MAX_TRACES`` traces in flight, enforced by the sender that owns the buffer; -* a wall-clock TTL, enforced by the sender's sweep. - -The identity itself lives in a ContextVar so that asynchronous work spawned -while handling one execution inherits it, mirroring ``telemetry.platform``. -Cross-task stages (plugin/RPC work with no inherited context) are resolved by -the owner's execution-id registry instead. +Cross-task callers (plugin/RPC work with no inherited context) resolve their +execution through the owner's execution-id registry instead, which only +carries the execution state - never node payloads. """ from __future__ import annotations @@ -27,64 +27,79 @@ import typing from datetime import datetime, timezone from uuid import uuid4 -MAX_STAGES = 32 - _current: contextvars.ContextVar['TraceState | None'] = contextvars.ContextVar('telemetry_trace', default=None) _route: contextvars.ContextVar[str] = contextvars.ContextVar('telemetry_trace_route', default='') _run: contextvars.ContextVar[str] = contextvars.ContextVar('telemetry_trace_run', default='') -# Open workflow step nodes, innermost last. Records emitted inside a step carry -# the innermost node as their ``parent``; the step's own record carries the one -# below it (its enclosing step), never itself. -_node: contextvars.ContextVar[tuple[str, ...]] = contextvars.ContextVar('telemetry_trace_node', default=()) +# Open workflow step frames of the current task, innermost last. The tuple is +# never mutated in place: opening a step publishes a new tuple on this task's +# context only, and closing it restores the previous one. +_node: contextvars.ContextVar[tuple['NodeFrame', ...]] = contextvars.ContextVar('telemetry_trace_node', default=()) def _now() -> str: return datetime.now(timezone.utc).isoformat() +class NodeFrame(typing.NamedTuple): + """One workflow step open in the task that opened it.""" + + node_id: str + seq: int + started_at: str + + class TraceState: - """Bounded stage buffer for exactly one execution.""" + """The small per-execution context; node payloads never accumulate here.""" __slots__ = ( - 'trace_id', - 'execution_id', - 'started_at', - 'stages', - 'dropped_stages', - 'synthetic', - 'sequence', - 'node_sequence', - 'open_nodes', + 'event_id', 'identity', - 'abandoned', + 'configured', + 'mode', + 'denominator', + 'emit_nodes', + 'debug', + 'origin', + 'started_at', + 'node_seq', + 'root_seen', + 'stage_count', + 'dropped_nodes', + 'last_outcome', + 'first_error', 'failure_reason', + 'runner_failed', 'closed', 'adapter', 'runner', 'runner_category', 'model_name', 'pipeline_plugins', + 'node_adapter', + 'node_runner', ) def __init__(self, execution_id: str | None = None) -> None: - # The execution identity is the trace identity so a stored chain can be - # joined to the Space view of that same Pipeline query / event / run. - self.execution_id = str(execution_id).strip() if execution_id else '' - self.trace_id = self.execution_id or str(uuid4()) - self.started_at = _now() - self.stages: list[dict[str, typing.Any]] = [] - self.dropped_stages = 0 - self.synthetic = False - self.sequence = 0 - # Node ids number the workflow steps of this one execution ("n1", "n2"...). - self.node_sequence = 0 - # Mirror of the open step stack, readable from any task: cross-task - # observations (plugin/RPC) resolve their parent through this. - self.open_nodes: tuple[str, ...] = () + # The execution identity every record of this chain carries. + self.event_id = str(execution_id).strip() if execution_id else str(uuid4()) self.identity: dict[str, str] = {} - # Set when the sender refused to buffer this trace: stop appending. - self.abandoned = False + self.configured = False + # Sampling, decided exactly once per event (at ingress or first touch). + self.mode = 'all' + self.denominator = 1 + self.emit_nodes = False + self.debug = False + self.origin = 'platform' + self.started_at = _now() + self.node_seq = 0 + self.root_seen = False + self.stage_count = 0 + self.dropped_nodes = 0 + self.last_outcome = '' + self.first_error = '' self.failure_reason = '' + # Set when a runner node already reported a non-success outcome. + self.runner_failed = False self.closed = False # Execution-scoped record fields attached by the owning lane when known. self.adapter = '' @@ -92,54 +107,40 @@ class TraceState: self.runner_category = '' self.model_name = '' self.pipeline_plugins: typing.Any = None + # First adapter/runner seen on a node, used for the chain record. + self.node_adapter = '' + self.node_runner = '' - def append( + def allocate_seq(self) -> int: + """Mint the next chain-local node order.""" + seq = self.node_seq + self.node_seq += 1 + return seq + + def note_node( self, *, - family: str, - operation: str, - mode: str, - adapter: str, - runner: str, outcome: str, - synthetic: bool, + adapter: str = '', + runner: str = '', + family: str = '', error: str = '', - node: str = '', - parent: str = '', ) -> None: - if len(self.stages) >= MAX_STAGES: - self.dropped_stages += 1 - return - seen = _now() - entry = { - 'family': family, - 'operation': operation, - 'mode': mode, - 'adapter': adapter, - 'runner': runner, - 'outcome': outcome, - 'synthetic': bool(synthetic), - 'seq': self.sequence, - 'route_ref': _route.get(), - 'run_id': _run.get(), - 'first_seen': seen, - 'last_seen': seen, - 'error': error, - } - # Only workflow steps carry these; legacy records stay byte-identical. - if node: - entry['node'] = node - if parent: - entry['parent'] = parent - self.stages.append(entry) - self.sequence += 1 - if synthetic: - self.synthetic = True - - def allocate_node(self) -> str: - """Mint the next step-node id for this execution ("n1", "n2", ...).""" - self.node_sequence += 1 - return f'n{self.node_sequence}' + """Remember only what the terminal chain record needs from a node.""" + self.stage_count += 1 + self.last_outcome = outcome or 'unknown' + if adapter and not self.node_adapter: + self.node_adapter = adapter + if runner and not self.node_runner: + self.node_runner = runner + if family == 'runner' and outcome not in ('success', 'skipped'): + # The lane records its own failed runner step only when the runner + # orchestrator has not already landed one. + self.runner_failed = True + if outcome in ('failed', 'cancelled', 'timeout') and error and not self.first_error: + # The chain must explain itself even when the break was reported as + # a failed node rather than as an exception at the boundary. + self.first_error = error def mark_failure(self, reason: str) -> None: """Remember why this chain broke so the payload can explain itself.""" @@ -147,18 +148,10 @@ class TraceState: self.failure_reason = reason def outcome(self) -> str: - """Terminal outcome of this chain, for space-side display. - - A chain that carries a failure reason never reports ``success``: the - reason is the point of the trace. - """ - if self.stages: - last = self.stages[-1]['outcome'] - else: - last = 'unknown' - if self.failure_reason and last in ('success', 'skipped', 'unknown'): + """Terminal outcome of this chain; a broken chain never reports success.""" + if self.failure_reason and self.last_outcome in ('', 'success', 'skipped', 'unknown'): return 'failed' - return last + return self.last_outcome or 'unknown' class TraceBinding(typing.NamedTuple): @@ -168,9 +161,9 @@ class TraceBinding(typing.NamedTuple): def bind(execution_id: str | None = None) -> TraceBinding: - """Start a trace unless one is already in flight in this context. + """Start a chain unless one is already in flight in this context. - The execution id becomes the trace id only when this call starts the trace; + The execution id becomes the chain identity only when this call starts it; a nested boundary keeps the enclosing execution's identity. """ existing = _current.get() @@ -181,7 +174,7 @@ def bind(execution_id: str | None = None) -> TraceBinding: def unbind(binding: TraceBinding) -> bool: - """Detach this binding. Returns True when this caller owns the trace.""" + """Detach this binding. Returns True when this caller owns the chain.""" if binding.token is not None: try: _current.reset(binding.token) @@ -192,7 +185,7 @@ def unbind(binding: TraceBinding) -> bool: def unbind_root(binding: TraceBinding) -> bool: - """Detach only when this caller started the trace, else leave it in place.""" + """Detach only when this caller started the chain, else leave it in place.""" if not binding.created: return False return unbind(binding) @@ -202,39 +195,64 @@ def current() -> TraceState | None: return _current.get() +def node_stack() -> tuple[NodeFrame, ...]: + """This task's own open-step frames, innermost last.""" + return _node.get() + + def current_parent() -> str: """The step node that records emitted right now attach to ('' when none).""" stack = _node.get() - return stack[-1] if stack else '' + return stack[-1].node_id if stack else '' def enclosing_parent() -> str: """The step enclosing the innermost open one; '' when it is a root step.""" stack = _node.get() - return stack[-2] if len(stack) > 1 else '' + return stack[-2].node_id if len(stack) > 1 else '' + + +def resolve_node( + state: TraceState, + node: str, + fallback_started_at: str, +) -> tuple[str, int, str, str]: + """Resolve one record to ``(node_id, seq, started_at, parent_node_id)``. + + A record stamped with a ``node`` the calling task has open reuses that + step's identity, order and start time; anything else (a plain observation, + or a node opened by another task) is a fresh node under this task's + innermost open step. + """ + stack = _node.get() + if node: + for index in range(len(stack) - 1, -1, -1): + frame = stack[index] + if frame.node_id == node: + parent = stack[index - 1].node_id if index > 0 else '' + return frame.node_id, frame.seq, frame.started_at, parent + parent = stack[-1].node_id if stack else '' + return node or uuid4().hex, state.allocate_seq(), fallback_started_at, parent @contextlib.contextmanager def stage_scope() -> typing.Iterator[str]: """Open one workflow step node for the duration of the block. - Allocates a node id unique within the execution, publishes it as the current - parent so that nested observations recorded inside the block attach to it, - and yields that id so the step's own records can be stamped with ``node=``. - Outside any trace nothing is allocated and the empty id is yielded. + Allocates a globally unique node id and the chain-local order, publishes the + frame on this task's own stack so nested observations attach to it as their + parent, and yields the node id so the step's own record can be stamped with + ``node=``. Outside any execution nothing is allocated and '' is yielded. """ state = _current.get() if state is None: yield '' return - node = state.allocate_node() - previous_open = state.open_nodes - state.open_nodes = previous_open + (node,) - token = _node.set(_node.get() + (node,)) + frame = NodeFrame(uuid4().hex, state.allocate_seq(), _now()) + token = _node.set(_node.get() + (frame,)) try: - yield node + yield frame.node_id finally: - state.open_nodes = previous_open try: _node.reset(token) except (ValueError, RuntimeError): @@ -258,7 +276,7 @@ def reset_run(token: typing.Any) -> None: @contextlib.contextmanager def scope(*, route_ref: str | None = None) -> typing.Iterator[None]: - """Pin routing identity for stages recorded inside this block.""" + """Pin routing identity for nodes recorded inside this block.""" if not route_ref: yield return @@ -270,3 +288,11 @@ def scope(*, route_ref: str | None = None) -> typing.Iterator[None]: _route.reset(token) except (ValueError, RuntimeError): pass + + +def current_route() -> str: + return _route.get() + + +def current_run() -> str: + return _run.get() diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index 9e5a0f7ac..f8752314f 100644 --- a/src/langbot/templates/config.yaml +++ b/src/langbot/templates/config.yaml @@ -444,11 +444,13 @@ space: # heartbeat and execution reporting; restart after changing this setting. disable_telemetry: false # Per-event execution traces: off | failures | sampled | all. - # 'sampled' keeps every failed run plus every N-th successful one - # (N is execution_trace_sample) and always keeps WebUI debug runs. - execution_trace: sampled + # Every execution is reported by default: one record per workflow node plus + # one terminal chain record. 'failures' keeps only broken chains, 'sampled' + # keeps broken chains plus every N-th successful one (N is + # execution_trace_sample), and 'off' reports no executions at all. + execution_trace: all execution_trace_sample: 20 # How often buffered execution records are uploaded, in seconds. - # Records of complete executions are batched as {"records": [...]}; a - # still-running execution is never uploaded half-written. + # Complete execution records are batched as {"records": [...]}; records are + # removed from the buffer only once the sender accepted them. telemetry_flush_seconds: 180 diff --git a/tests/unit_tests/platform/test_botmgr_execution_trace.py b/tests/unit_tests/platform/test_botmgr_execution_trace.py index 75605aa04..8fbad1463 100644 --- a/tests/unit_tests/platform/test_botmgr_execution_trace.py +++ b/tests/unit_tests/platform/test_botmgr_execution_trace.py @@ -1,4 +1,4 @@ -"""RuntimeBot ingress tracing: one inbound event owns one execution trace.""" +"""RuntimeBot ingress tracing: one inbound event owns one v2 execution chain.""" from __future__ import annotations @@ -28,9 +28,6 @@ class FakeTelemetryManager: self.sent.append(payload) return True - def start_send_task(self, payload: dict) -> None: - self.sent.append(payload) - def make_bot(event_bindings: list[dict], config=None, agent=None): from langbot.pkg.platform.botmgr import RuntimeBot @@ -76,7 +73,7 @@ class FakeAdapter: async def flushed_records(counters, manager) -> list[dict]: - """Flush buffered telemetry and return everything the pipeline emitted. + """Flush buffered telemetry and return everything the event emitted. Records are batched and uploaded by ``ExecutionCounters.flush()``; there is no per-trace direct send any more. @@ -87,34 +84,48 @@ async def flushed_records(counters, manager) -> list[dict]: return [record for payload in manager.sent for record in payload.get('records', [])] +def execution_nodes(records: list[dict]) -> list[dict]: + return [record for record in records if record.get('event_type') == 'execution_node'] + + +def execution_chains(records: list[dict]) -> list[dict]: + return [record for record in records if record.get('event_type') == 'execution_chain'] + + @pytest.mark.asyncio -async def test_route_miss_emits_one_closed_trace(): +async def test_route_miss_emits_one_closed_chain(): bot, manager, counters = make_bot([]) await bot._handle_platform_event(message_received_event(), FakeAdapter()) records = await flushed_records(counters, manager) assert len(manager.sent) == 1 - trace_records = [item for item in records if item.get('event_type') == 'feature_execution'] - assert len(trace_records) == 1 - payload = trace_records[0] - assert payload['event_type'] == 'feature_execution' + chains = execution_chains(records) + assert len(chains) == 1 + chain = chains[0] + assert chain['schema'] == 2 # The execution identity is deterministic: one inbound event, one chain. - assert payload['query_id'] == 'platform:bot-1:message-1' - assert payload['instance_id'] == 'instance-test' - assert payload['workspace_uuid'] == 'workspace-test' - features = payload['features'] - assert features['schema'] == 1 - assert [(row['family'], row['operation'], row['outcome']) for row in features['observations']] == [ + assert chain['event_id'] == 'platform:bot-1:message-1' + assert chain['instance_id'] == 'instance-test' + assert chain['workspace_uuid'] == 'workspace-test' + assert chain['closed_by'] == 'event_done' + assert chain['stage_count'] == 2 + + rows = execution_nodes(records) + assert [(row['family'], row['operation'], row['outcome']) for row in rows] == [ ('event_route', 'message.received', 'skipped'), ('platform_event', 'message.received', 'success'), ] - assert [row['seq'] for row in features['observations']] == [0, 1] - assert all(row['trace_id'] == payload['query_id'] for row in features['observations']) - assert all(row['adapter'] == 'FakeAdapter' for row in features['observations']) - assert features['trace']['closed_by'] == 'event_done' - # The trace is closed and deregistered: nothing dangles after the event. - assert counters.traces == {} + assert {row['seq'] for row in rows} == {0, 1} + assert all(row['event_id'] == chain['event_id'] for row in rows) + assert all(row['adapter'] == 'FakeAdapter' for row in rows) + # The inbound event is the root step; the routing decision nests under it. + root = next(row for row in rows if row['root']) + route = next(row for row in rows if row['family'] == 'event_route') + assert (root['family'], root['parent_node_id']) == ('platform_event', '') + assert route['parent_node_id'] == root['node_id'] + # The chain is closed and deregistered: nothing dangles after the event. + assert counters.registry == {} @pytest.mark.asyncio @@ -133,13 +144,13 @@ async def test_unavailable_route_target_carries_route_identity(): await bot._handle_platform_event(message_received_event(), FakeAdapter()) - stages = (await flushed_records(counters, manager))[0]['features']['observations'] - assert [(stage['family'], stage['outcome'], stage['mode']) for stage in stages] == [ + nodes = execution_nodes(await flushed_records(counters, manager)) + assert [(node['family'], node['outcome'], node['mode']) for node in nodes] == [ ('event_route', 'failed', 'agent'), ('platform_event', 'success', 'none'), ] - assert stages[0]['route_ref'] == 'agent:agent-1' - assert stages[1]['route_ref'] == '' + assert nodes[0]['route_ref'] == 'agent:agent-1' + assert nodes[1]['route_ref'] == '' @pytest.mark.asyncio @@ -157,10 +168,10 @@ async def test_discarded_route_is_visible_without_user_values(): await bot._handle_platform_event(SimpleNamespace(type='platform.member.joined'), FakeAdapter()) - stages = (await flushed_records(counters, manager))[0]['features']['observations'] - assert [stage['outcome'] for stage in stages] == ['skipped', 'success'] - assert stages[0]['operation'] == 'platform.member.joined' - assert stages[0]['mode'] == 'none' + nodes = execution_nodes(await flushed_records(counters, manager)) + assert [node['outcome'] for node in nodes] == ['skipped', 'success'] + assert nodes[0]['operation'] == 'platform.member.joined' + assert nodes[0]['mode'] == 'none' @pytest.mark.asyncio @@ -171,7 +182,7 @@ async def test_telemetry_opt_out_emits_nothing(): assert manager.sent == [] assert counters.records == [] - assert counters.traces == {} + assert counters.registry == {} @pytest.mark.asyncio @@ -185,13 +196,13 @@ async def test_route_trace_scope_is_restored_after_dispatch(): assert trace_mod.current() is None buffered = list(counters.records) - # A later record outside the ingress must not join the closed trace, and an + # A later record outside the ingress must not join the closed chain, and an # observation without an owning chain is not buffered at all. from langbot.pkg.telemetry.execution import record record(bot.ap, TEST_CONTEXT, family='platform_api', operation='send_message', outcome='success') assert counters.records == buffered - assert counters.traces == {} + assert counters.registry == {} @pytest.mark.asyncio diff --git a/tests/unit_tests/provider/test_mcp_resources.py b/tests/unit_tests/provider/test_mcp_resources.py index 0ba164e9a..97b136d56 100644 --- a/tests/unit_tests/provider/test_mcp_resources.py +++ b/tests/unit_tests/provider/test_mcp_resources.py @@ -24,7 +24,6 @@ from langbot.pkg.provider.tools.loaders.mcp import ( MCPToolCallTimeoutError, RuntimeMCPSession, ) -from langbot.pkg.telemetry import features as telemetry_features from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError @@ -298,10 +297,6 @@ async def test_read_resource_envelope_truncates_caches_and_records_trace(): traces = query.variables[MCP_RESOURCE_TRACE_QUERY_KEY] assert [trace['source'] for trace in traces] == ['ui_preview', 'agent_tool'] assert traces[1]['cache_hit'] is True - assert query.variables[telemetry_features.FEATURES_KEY]['mcp_resource_reads'] == { - 'ui_preview': 1, - 'agent_tool': 1, - } @pytest.mark.asyncio diff --git a/tests/unit_tests/telemetry/test_cloud_sampling.py b/tests/unit_tests/telemetry/test_cloud_sampling.py index b1a604489..ef4de5e2e 100644 --- a/tests/unit_tests/telemetry/test_cloud_sampling.py +++ b/tests/unit_tests/telemetry/test_cloud_sampling.py @@ -1,6 +1,7 @@ from types import SimpleNamespace from langbot.pkg.telemetry.execution import ExecutionCounters +from langbot.pkg.telemetry.trace import TraceState from langbot.pkg.utils import constants @@ -8,8 +9,10 @@ def test_cloud_counts_are_not_sampled(monkeypatch): monkeypatch.setattr(constants, 'edition', 'cloud') counters = ExecutionCounters(SimpleNamespace(telemetry_config={'execution_trace': 'sampled'})) assert counters.trace_mode() == 'all' - state = SimpleNamespace(failure_reason='', stages=[], synthetic=False, trace_id='00000001') - assert counters._trace_emitted(state) + state = TraceState('00000001') + counters.configure(state) + assert state.emit_nodes is True + assert counters._chain_emitted(state) is True def test_community_retains_sampling(monkeypatch): diff --git a/tests/unit_tests/telemetry/test_features.py b/tests/unit_tests/telemetry/test_features.py deleted file mode 100644 index 6e19d0828..000000000 --- a/tests/unit_tests/telemetry/test_features.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Unit tests for telemetry feature counters (pkg/telemetry/features.py).""" - -from __future__ import annotations - -from importlib import import_module - - -def get_features_module(): - return import_module('langbot.pkg.telemetry.features') - - -class FakeQuery: - def __init__(self): - self.variables = {} - - -class TestIncrement: - def test_increment_nested_counter(self): - features = get_features_module() - q = FakeQuery() - features.increment(q, 'tool_calls', 'native') - features.increment(q, 'tool_calls', 'native') - features.increment(q, 'tool_calls', 'mcp') - assert q.variables[features.FEATURES_KEY]['tool_calls'] == {'native': 2, 'mcp': 1} - - def test_increment_flat_counter(self): - features = get_features_module() - q = FakeQuery() - features.increment(q, 'something') - features.increment(q, 'something', amount=2) - assert q.variables[features.FEATURES_KEY]['something'] == 3 - - def test_increment_never_raises_on_broken_query(self): - features = get_features_module() - - class Broken: - @property - def variables(self): - raise RuntimeError('boom') - - # Must not raise - features.increment(Broken(), 'tool_calls', 'native') - - def test_set_value(self): - features = get_features_module() - q = FakeQuery() - features.set_value(q, 'tool_call_rounds', 5) - assert q.variables[features.FEATURES_KEY]['tool_call_rounds'] == 5 - - -class TestCollectFeatures: - def test_collect_empty(self): - features = get_features_module() - q = FakeQuery() - assert features.collect_features(q) == {} - - def test_collect_combines_counters_and_snapshots(self): - features = get_features_module() - q = FakeQuery() - features.increment(q, 'sandbox', 'execs') - features.set_value(q, 'kb', {'kb_count': 2, 'engine_plugins': ['builtin'], 'retrieved_entries': 7}) - q.variables['_activated_skills'] = {'pdf-tools': {}, 'a-skill': {}} - q.variables['_pipeline_bound_mcp_servers'] = ['srv1', 'srv2'] - - result = features.collect_features(q) - assert result['sandbox'] == {'execs': 1} - assert result['kb']['kb_count'] == 2 - assert result['activated_skills'] == ['a-skill', 'pdf-tools'] # sorted - assert result['mcp_servers'] == ['srv1', 'srv2'] - - def test_collect_omits_mcp_when_all_enabled(self): - """None means 'all enabled' and is not reported.""" - features = get_features_module() - q = FakeQuery() - q.variables['_pipeline_bound_mcp_servers'] = None - assert 'mcp_servers' not in features.collect_features(q) - - def test_collect_drops_non_json_serializable(self): - features = get_features_module() - q = FakeQuery() - features.set_value(q, 'good', 1) - features.set_value(q, 'bad', object()) - result = features.collect_features(q) - assert result == {'good': 1} - - def test_collect_is_json_serializable(self): - import json - - features = get_features_module() - q = FakeQuery() - features.increment(q, 'tool_calls', 'skill') - json.dumps(features.collect_features(q)) diff --git a/tests/unit_tests/telemetry/test_pipeline_lane.py b/tests/unit_tests/telemetry/test_pipeline_lane.py index f091f2bb8..4b3e8625e 100644 --- a/tests/unit_tests/telemetry/test_pipeline_lane.py +++ b/tests/unit_tests/telemetry/test_pipeline_lane.py @@ -1,4 +1,4 @@ -"""The Pipeline lane must land its failure on the owning execution trace.""" +"""The Pipeline lane must land its nodes on one v2 execution chain.""" from __future__ import annotations @@ -43,7 +43,7 @@ def build_pipeline(pipelinemgr, ap): return pipeline -def test_lane_failure_marks_the_owned_trace_without_any_prior_stage(): +def test_lane_failure_marks_the_owned_chain_without_any_prior_stage(): pipelinemgr, execution = get_modules() manager = FakeManager() counters = execution.ExecutionCounters(manager) @@ -54,15 +54,16 @@ def test_lane_failure_marks_the_owned_trace_without_any_prior_stage(): with execution.ingress(ap, 'pipeline_done', CONTEXT, execution_id='query-7'): pipeline._record_lane_failure(query, 'sandbox exited 127') - assert len(counters.records) == 1 - record = counters.records[0] - assert record['query_id'] == 'query-7' - assert record['error'] == 'sandbox exited 127' - assert record['features']['trace']['outcome'] == 'failed' - observations = record['features']['observations'] - assert [(row['family'], row['operation'], row['outcome']) for row in observations] == [ - ('runner', 'execute', 'failed') - ] + rows = [record for record in counters.records if record['event_type'] == 'execution_node'] + assert len(rows) == 1 + node = rows[0] + assert node['event_id'] == 'query-7' + assert node['error'] == 'sandbox exited 127' + assert (node['family'], node['operation'], node['outcome']) == ('runner', 'execute', 'failed') + chain = [record for record in counters.records if record['event_type'] == 'execution_chain'][0] + assert chain['event_id'] == 'query-7' + assert chain['outcome'] == 'failed' + assert chain['error'] == 'sandbox exited 127' def test_lane_failure_does_not_duplicate_the_runner_node_the_orchestrator_recorded(): @@ -87,13 +88,13 @@ def test_lane_failure_does_not_duplicate_the_runner_node_the_orchestrator_record ) pipeline._record_lane_failure(query, 'No authorized model for local-agent') - observations = [ - (row['family'], row['operation'], row['outcome']) - for record in counters.records - for row in record['features']['observations'] + rows = [record for record in counters.records if record['event_type'] == 'execution_node'] + assert [(record['family'], record['operation'], record['outcome']) for record in rows] == [ + ('runner', 'execute', 'failed') ] - assert observations == [('runner', 'execute', 'failed')] - assert counters.records[-1]['features']['trace']['outcome'] == 'failed' + chain = [record for record in counters.records if record['event_type'] == 'execution_chain'][-1] + assert chain['outcome'] == 'failed' + assert chain['error'] == 'No authorized model for local-agent' def build_counters(): @@ -118,11 +119,14 @@ def build_query(**overrides): return query +def lane_nodes(counters): + return [record for record in counters.records if record['event_type'] == 'execution_node'] + + def lane_shapes(counters): return [ - (row['family'], row['operation'], row.get('node'), row.get('parent')) - for record in counters.records - for row in record['features']['observations'] + (record['family'], record['operation'], record['node_id'], record['parent_node_id']) + for record in lane_nodes(counters) ] @@ -138,7 +142,7 @@ async def stop_counters(counters): counters.task = None -def test_lane_step_owns_the_apis_and_runner_it_calls(): +async def test_lane_step_owns_the_apis_and_runner_it_calls(): pipelinemgr, execution = get_modules() trace = get_trace() counters, ap = build_counters() @@ -168,21 +172,28 @@ def test_lane_step_owns_the_apis_and_runner_it_calls(): await lane.process_query(query) await stop_counters(counters) - asyncio.run(scenario()) + await scenario() shapes = lane_shapes(counters) - lane_node = shapes[-1][2] - assert lane_node == 'n1' - assert shapes == [ - ('platform_api', 'send_message', None, lane_node), - ('runner', 'execute', seen['run'], lane_node), - ('platform_api', 'reply_message', None, lane_node), - # The lane step is recorded last; a standalone lane is a root. - ('pipeline', 'run', lane_node, None), + # The synthetic inbound event roots the chain; the lane step hangs under it. + assert [shape[0:2] for shape in shapes] == [ + ('platform_event', 'message.received'), + ('platform_api', 'send_message'), + ('runner', 'execute'), + ('platform_api', 'reply_message'), + ('pipeline', 'run'), ] + event_node = shapes[0][2] + lane_node = shapes[-1][2] + assert shapes[0][3] == '' + assert shapes[1][3] == lane_node + assert shapes[2] == ('runner', 'execute', seen['run'], lane_node) + assert shapes[3][3] == lane_node + # The lane step is recorded last and nests under the synthetic event node. + assert shapes[4] == ('pipeline', 'run', lane_node, event_node) -def test_lane_node_nests_under_the_platform_route_that_called_it(): +async def test_lane_node_nests_under_the_platform_route_that_called_it(): pipelinemgr, execution = get_modules() trace = get_trace() counters, ap = build_counters() @@ -197,7 +208,7 @@ def test_lane_node_nests_under_the_platform_route_that_called_it(): lane._process_query = fake_process async def scenario(): - # A platform event owns the trace; the route step wraps the Pipeline lane. + # A platform event owns the chain; the route step wraps the Pipeline lane. with execution.ingress(ap, 'event_done', CONTEXT, execution_id='exec-routed'): with trace.stage_scope() as event_node: execution.record( @@ -213,19 +224,21 @@ def test_lane_node_nests_under_the_platform_route_that_called_it(): await stop_counters(counters) return event_node, route_node - event_node, route_node = asyncio.run(scenario()) + event_node, route_node = await scenario() - assert (event_node, route_node) == ('n1', 'n2') - assert lane_shapes(counters) == [ - ('platform_event', 'message.received', event_node, None), - ('event_route', 'message.received', route_node, event_node), - # The API is called by the lane, so it hangs on the lane node. - ('platform_api', 'send_message', None, 'n3'), - ('pipeline', 'run', 'n3', route_node), - ] + assert event_node != route_node + shapes = lane_shapes(counters) + assert shapes[0] == ('platform_event', 'message.received', event_node, '') + assert shapes[1] == ('event_route', 'message.received', route_node, event_node) + lane_node = shapes[-1][2] + # The API is called by the lane, so it hangs on the lane node. + assert shapes[2][0:2] == ('platform_api', 'send_message') + assert shapes[2][3] == lane_node and shapes[2][2] != lane_node + # The lane nests under the route step that called it. + assert shapes[3] == ('pipeline', 'run', lane_node, route_node) -def test_broken_lane_records_a_failed_step_with_its_reason(): +async def test_broken_lane_records_a_failed_step_with_its_reason(): pipelinemgr, execution = get_modules() counters, ap = build_counters() lane = build_lane(pipelinemgr, ap) @@ -242,16 +255,19 @@ def test_broken_lane_records_a_failed_step_with_its_reason(): await stop_counters(counters) with pytest.raises(RuntimeError, match='lane exploded'): - asyncio.run(scenario()) + await scenario() - family, operation, node, parent = lane_shapes(counters)[-1] - row = counters.records[0]['features']['observations'][-1] - assert (family, operation, node, parent) == ('pipeline', 'run', 'n1', None) + shapes = lane_shapes(counters) + family, operation, node, parent = shapes[-1] + row = lane_nodes(counters)[-1] + # The lane step nests under the synthetic inbound event that roots the chain. + assert (family, operation) == ('pipeline', 'run') + assert shapes[0] == ('platform_event', 'message.received', parent, '') assert row['outcome'] == 'failed' assert row['error'] == 'lane exploded' -def test_lane_reported_failure_marks_the_step_failed(): +async def test_lane_reported_failure_marks_the_step_failed(): pipelinemgr, execution = get_modules() counters, ap = build_counters() lane = build_lane(pipelinemgr, ap) @@ -266,7 +282,7 @@ def test_lane_reported_failure_marks_the_step_failed(): await lane.process_query(query) await stop_counters(counters) - asyncio.run(scenario()) + await scenario() - row = counters.records[0]['features']['observations'][-1] + row = lane_nodes(counters)[-1] assert (row['family'], row['operation'], row['outcome']) == ('pipeline', 'run', 'failed') diff --git a/tests/unit_tests/telemetry/test_telemetry.py b/tests/unit_tests/telemetry/test_telemetry.py index 004009a68..87699a18d 100644 --- a/tests/unit_tests/telemetry/test_telemetry.py +++ b/tests/unit_tests/telemetry/test_telemetry.py @@ -6,7 +6,6 @@ Tests cover: - Early return conditions (disabled, empty config, no server) - URL construction (with actual URL verification) - HTTP request success/failure scenarios -- Source code bug: send_tasks should be instance variable """ from __future__ import annotations @@ -629,53 +628,3 @@ class TestAuthenticatedWorkspaceReporter: 'Authorization': 'Bearer refreshed-workspace-owner-token' } - -class TestStartSendTask: - """Tests for start_send_task() method.""" - - @pytest.mark.asyncio - async def test_start_send_task_creates_task(self): - """Test that start_send_task creates an asyncio task.""" - telemetry = get_telemetry_module() - mock_app = Mock() - mock_app.logger = Mock() - mock_app.instance_config = Mock() - mock_app.instance_config.data = {} - - manager = telemetry.TelemetryManager(mock_app) - manager.telemetry_config = {} - - await manager.start_send_task({'query_id': 'test'}) - - # Task should be added to send_tasks list - assert len(manager.send_tasks) >= 1 - - # Clean up the task - for task in manager.send_tasks: - if not task.done(): - task.cancel() - manager.send_tasks.clear() - - @pytest.mark.asyncio - async def test_start_send_task_multiple_tasks(self): - """Test that multiple tasks are tracked.""" - telemetry = get_telemetry_module() - mock_app = Mock() - mock_app.logger = Mock() - mock_app.instance_config = Mock() - mock_app.instance_config.data = {} - - manager = telemetry.TelemetryManager(mock_app) - manager.telemetry_config = {} - - await manager.start_send_task({'query_id': 'test1'}) - await manager.start_send_task({'query_id': 'test2'}) - await manager.start_send_task({'query_id': 'test3'}) - - assert len(manager.send_tasks) >= 3 - - # Clean up - for task in manager.send_tasks: - if not task.done(): - task.cancel() - manager.send_tasks.clear() diff --git a/tests/unit_tests/telemetry/test_trace.py b/tests/unit_tests/telemetry/test_trace.py index 283cdac65..95b9f06f1 100644 --- a/tests/unit_tests/telemetry/test_trace.py +++ b/tests/unit_tests/telemetry/test_trace.py @@ -1,4 +1,10 @@ -"""Unit tests for per-execution execution traces (pkg/telemetry/trace.py + execution.py).""" +"""Unit tests for per-execution execution chains (pkg/telemetry/trace.py + execution.py). + +One execution is one chain of independent records: one ``execution_node`` per +workflow node, emitted the moment the node completes, plus exactly one +``execution_chain`` when the execution closes. Every record of a chain carries +the same ``event_id``. +""" from __future__ import annotations @@ -31,8 +37,21 @@ class FakeManager: self.sent.append(payload) return True - def start_send_task(self, payload: dict) -> None: + +class RejectingManager(FakeManager): + """Stand-in whose sender refuses every batch until told to accept them.""" + + def __init__(self, config=None): + super().__init__(config) + self.rejecting = True + self.attempts = 0 + + async def send(self, payload: dict) -> bool: + self.attempts += 1 + if self.rejecting: + return False self.sent.append(payload) + return True def make_counters(config=None): @@ -56,6 +75,14 @@ def records_sent(manager): return [record for payload in manager.sent for record in payload.get('records', [])] +def nodes(counters): + return [record for record in counters.records if record.get('event_type') == 'execution_node'] + + +def chains(counters): + return [record for record in counters.records if record.get('event_type') == 'execution_chain'] + + CONTEXT = types.SimpleNamespace(instance_uuid='instance-1', workspace_uuid='workspace-1') STAGE = { @@ -66,6 +93,35 @@ STAGE = { 'outcome': 'success', } +NODE_KEYS = { + 'event_type', + 'schema', + 'event_id', + 'instance_id', + 'workspace_uuid', + 'version', + 'edition', + 'timestamp', + 'debug', + 'sample', + 'node_id', + 'parent_node_id', + 'root', + 'seq', + 'family', + 'operation', + 'mode', + 'adapter', + 'runner', + 'outcome', + 'error', + 'route_ref', + 'run_id', + 'synthetic', + 'started_at', + 'ended_at', +} + class TestTraceIdentity: def test_nested_bind_reuses_the_in_flight_trace(self): @@ -93,13 +149,12 @@ class TestTraceIdentity: # A repeated unbind of the same binding must never raise. assert trace.unbind_root(outer) is True - def test_trace_id_is_the_execution_id(self): + def test_event_id_is_the_execution_id(self): trace, _ = get_modules() binding = trace.bind('exec-1234') try: assert binding.created is True - assert binding.state.trace_id == 'exec-1234' - assert binding.state.execution_id == 'exec-1234' + assert binding.state.event_id == 'exec-1234' finally: trace.unbind_root(binding) @@ -109,19 +164,20 @@ class TestTraceIdentity: inner = trace.bind('query-9') try: assert inner.created is False - assert inner.state.trace_id == 'event-1' + assert inner.state.event_id == 'event-1' finally: trace.unbind_root(inner) trace.unbind_root(outer) def test_route_and_run_are_only_pinned_inside_the_block(self): trace, _ = get_modules() + _, counters = make_counters(trace_config()) binding = trace.bind() token = trace.set_run('run-1') try: with trace.scope(route_ref='agent:agent-1'): - inside = trace.current() - inside.append( + counters.record( + CONTEXT, family='runner', operation='execute', mode='agent', @@ -130,8 +186,8 @@ class TestTraceIdentity: outcome='success', synthetic=False, ) - outside = trace.current() - outside.append( + counters.record( + CONTEXT, family='platform_api', operation='send_message', mode='agent', @@ -143,7 +199,7 @@ class TestTraceIdentity: finally: trace.reset_run(token) trace.unbind_root(binding) - assert [(s['route_ref'], s['run_id']) for s in binding.state.stages] == [ + assert [(row['route_ref'], row['run_id']) for row in nodes(counters)] == [ ('agent:agent-1', 'run-1'), ('', 'run-1'), ] @@ -152,65 +208,42 @@ class TestTraceIdentity: class TestTraceRecording: def test_stage_carries_sequence_identity_and_timestamps(self): trace, _ = get_modules() - _, counters = make_counters() + _, counters = make_counters(trace_config()) binding = trace.bind() try: counters.record(CONTEXT, **STAGE) counters.record(CONTEXT, **{**STAGE, 'family': 'runner', 'operation': 'execute', 'mode': 'pipeline'}) + counters.close_trace(binding.state, 'event_done') finally: trace.unbind_root(binding) - stages = binding.state.stages - assert [stage['seq'] for stage in stages] == [0, 1] - assert stages[0]['first_seen'] == stages[0]['last_seen'] - assert binding.state.identity == {'instance_id': 'instance-1', 'workspace_uuid': 'workspace-1'} - assert binding.state.outcome() == 'success' - - def test_stage_buffer_is_bounded_and_counts_overflow(self): - trace, _ = get_modules() - _, counters = make_counters() - binding = trace.bind() - try: - for _ in range(trace.MAX_STAGES + 5): - counters.record(CONTEXT, **STAGE) - finally: - trace.unbind_root(binding) - assert len(binding.state.stages) == trace.MAX_STAGES - assert binding.state.dropped_stages == 5 + rows = nodes(counters) + assert [row['seq'] for row in rows] == [0, 1] + # A plain observation starts and ends at the same instant. + assert rows[0]['started_at'] == rows[0]['ended_at'] + assert {(row['instance_id'], row['workspace_uuid']) for row in rows} == {('instance-1', 'workspace-1')} + chain = chains(counters)[0] + assert chain['outcome'] == 'success' + assert chain['stage_count'] == 2 def test_oversized_identifier_never_becomes_a_stage(self): trace, _ = get_modules() - _, counters = make_counters() + _, counters = make_counters(trace_config()) binding = trace.bind() try: counters.record(CONTEXT, **{**STAGE, 'operation': 'x' * 200}) counters.record(CONTEXT, **{**STAGE, 'family': 'not-a-family'}) finally: trace.unbind_root(binding) - assert binding.state.stages == [] - - def test_trace_limit_abandons_further_traces(self): - trace, execution = get_modules() - _, counters = make_counters() - bindings = [] - for _ in range(execution.MAX_TRACES + 1): - binding = trace.bind() - bindings.append(binding) - counters.record(CONTEXT, **STAGE) - trace.unbind_root(binding) - assert trace.current() is None - assert len(counters.traces) == execution.MAX_TRACES - assert counters.dropped == 1 - assert bindings[-1].state.abandoned is True - assert bindings[-1].state.stages == [] + assert nodes(counters) == [] def test_observation_without_an_owning_execution_is_dropped(self): - trace, _ = get_modules() + _, _ = get_modules() _, counters = make_counters() - # Clustering is gone: an observation that belongs to no execution is - # discarded instead of being aggregated into a synthetic bucket. + # An observation that belongs to no execution is discarded instead of + # being aggregated into a synthetic bucket. counters.record(CONTEXT, **STAGE) assert counters.records == [] - assert counters.traces == {} + assert counters.registry == {} class TestTraceNodeScope: @@ -238,11 +271,13 @@ class TestTraceNodeScope: finally: trace.unbind_root(binding) - step_row, api_row = binding.state.stages - assert step_row['node'] == step - assert 'parent' not in step_row # a root step carries no parent - assert api_row['parent'] == step - assert 'node' not in api_row # observations are not step nodes + step_row, api_row = nodes(counters) + assert step_row['node_id'] == step + assert step_row['parent_node_id'] == '' # a root step carries no parent + assert step_row['root'] is True + assert api_row['parent_node_id'] == step + assert api_row['root'] is False + assert api_row['node_id'] != step # observations are not step nodes def test_nested_step_records_its_enclosing_step_as_parent(self): trace, _ = get_modules() @@ -257,11 +292,11 @@ class TestTraceNodeScope: finally: trace.unbind_root(binding) - outer_row, inner_row, api_row = binding.state.stages - assert (outer_row['node'], outer_row.get('parent')) == (outer, None) + outer_row, inner_row, api_row = nodes(counters) + assert (outer_row['node_id'], outer_row['parent_node_id']) == (outer, '') # The step's own record sits under its enclosing step, never under itself. - assert (inner_row['node'], inner_row.get('parent')) == (inner, outer) - assert api_row.get('parent') == inner and 'node' not in api_row + assert (inner_row['node_id'], inner_row['parent_node_id']) == (inner, outer) + assert api_row['parent_node_id'] == inner and api_row['node_id'] != inner def test_sequential_steps_get_distinct_ids_and_roots_have_no_parent(self): trace, _ = get_modules() @@ -275,13 +310,16 @@ class TestTraceNodeScope: finally: trace.unbind_root(binding) - assert (first, second) == ('n1', 'n2') - assert [row['node'] for row in binding.state.stages] == [first, second] - assert all('parent' not in row for row in binding.state.stages) + assert first and second and first != second + rows = nodes(counters) + assert [row['node_id'] for row in rows] == [first, second] + assert [row['seq'] for row in rows] == [0, 1] + assert rows[0]['root'] is True + assert all(row['parent_node_id'] == '' for row in rows) - def test_cross_task_observation_attaches_to_the_open_step(self): + def test_cross_task_observation_joins_the_owning_execution(self): trace, execution = get_modules() - _, counters = make_counters() + _, counters = make_counters(trace_config()) binding = execution.bind_trace(make_ap(counters), 'exec-cross') try: with trace.stage_scope() as lane: @@ -290,7 +328,7 @@ class TestTraceNodeScope: ) # A plugin/RPC observation arrives on a task that never had the # trace context; it is resolved through the execution registry and - # must still hang off the step the owning task has open. + # still lands on the owning chain. foreign = contextvars.Context() foreign.run(execution.set_execution_id, 'exec-cross') foreign.run( @@ -304,13 +342,16 @@ class TestTraceNodeScope: finally: trace.unbind_root(binding) - rows = binding.state.stages - assert rows[0]['node'] == lane and 'parent' not in rows[0] - assert rows[1].get('parent') == lane and 'node' not in rows[1] + lane_row, foreign_row = nodes(counters) + assert lane_row['node_id'] == lane and lane_row['parent_node_id'] == '' and lane_row['root'] is True + # The cross-task node joins the same execution, but never mirrors another + # task's open-step stack: it is a root of the chain. + assert foreign_row['event_id'] == 'exec-cross' + assert foreign_row['parent_node_id'] == '' def test_cross_task_observation_without_an_open_step_has_no_parent(self): trace, execution = get_modules() - _, counters = make_counters() + _, counters = make_counters(trace_config()) binding = execution.bind_trace(make_ap(counters), 'exec-bare') try: foreign = contextvars.Context() @@ -326,8 +367,8 @@ class TestTraceNodeScope: finally: trace.unbind_root(binding) - row = binding.state.stages[0] - assert 'node' not in row and 'parent' not in row + row = nodes(counters)[0] + assert row['parent_node_id'] == '' and row['root'] is True def test_scope_outside_any_trace_is_inert(self): trace, _ = get_modules() @@ -335,9 +376,9 @@ class TestTraceNodeScope: assert node == '' assert trace.current_parent() == '' - def test_records_outside_any_scope_serialize_without_node_keys(self): - trace, execution = get_modules() - manager, counters = make_counters(trace_config()) + def test_records_outside_any_scope_are_root_nodes(self): + trace, _ = get_modules() + _, counters = make_counters(trace_config()) binding = trace.bind('exec-legacy') try: counters.record(CONTEXT, **STAGE) @@ -345,29 +386,16 @@ class TestTraceNodeScope: finally: trace.unbind_root(binding) - stage = binding.state.stages[0] - assert set(stage) == { - 'family', - 'operation', - 'mode', - 'adapter', - 'runner', - 'outcome', - 'synthetic', - 'seq', - 'route_ref', - 'run_id', - 'first_seen', - 'last_seen', - 'error', - } - observation = counters.records[0]['features']['observations'][0] - assert 'node' not in json.dumps(observation, sort_keys=True) - assert 'parent' not in json.dumps(observation, sort_keys=True) + node = nodes(counters)[0] + assert set(node) == NODE_KEYS + assert node['node_id'] and node['parent_node_id'] == '' and node['root'] is True + assert node['seq'] == 0 + # Node placement lives on node records only; the chain carries none. + assert 'node_id' not in json.dumps(chains(counters)[0], sort_keys=True) class TestTraceDelivery: - def test_close_buffers_one_record_per_trace(self): + def test_close_emits_one_terminal_chain_record_per_execution(self): trace, _ = get_modules() manager, counters = make_counters(trace_config()) binding = trace.bind('exec-1') @@ -383,49 +411,55 @@ class TestTraceDelivery: # Closing buffers; nothing is sent before the flush cadence. assert manager.sent == [] - assert len(counters.records) == 1 - assert counters.traces == {} - - record = counters.records[0] - assert record['event_type'] == 'feature_execution' - assert record['query_id'] == 'exec-1' + assert len(nodes(counters)) == 2 + chain_rows = chains(counters) + assert len(chain_rows) == 1 + record = chain_rows[0] + assert record['schema'] == 2 + assert record['event_id'] == 'exec-1' assert record['instance_id'] == 'instance-1' assert record['workspace_uuid'] == 'workspace-1' - assert record['trusted'] is True - features = record['features'] - assert features['schema'] == 1 - assert 'count' not in json.dumps(record) - assert [row['seq'] for row in features['observations']] == [0, 1] - assert [row['trace_id'] for row in features['observations']] == ['exec-1'] * 2 - assert features['trace'] == { - 'closed_by': 'event_done', - 'started_at': binding.state.started_at, - 'ended_at': features['trace']['ended_at'], - 'route_ref': '', - 'run_id': '', - 'outcome': 'success', - 'dropped_stages': 0, - } + assert record['closed_by'] == 'event_done' + assert record['outcome'] == 'success' + assert record['stage_count'] == 2 + assert isinstance(record['duration_ms'], int) and record['duration_ms'] >= 0 + assert 'query_id' not in record and 'features' not in record def test_sampling_decision_is_deterministic_and_keeps_failures(self): trace, _ = get_modules() _, counters = make_counters(trace_config(execution_trace='sampled', execution_trace_sample=2)) - binding = trace.bind() - try: - binding.state.trace_id = '00000001-0000-4000-8000-000000000000' - counters.record(CONTEXT, **STAGE) - assert counters._trace_emitted(binding.state) is False - binding.state.trace_id = '00000000-0000-4000-8000-000000000000' - assert counters._trace_emitted(binding.state) is True - binding.state.trace_id = '00000001-0000-4000-8000-000000000001' - counters.record(CONTEXT, **{**STAGE, 'outcome': 'timeout'}) - assert counters._trace_emitted(binding.state) is True - finally: - trace.unbind_root(binding) + in_id = next(i for i in (f'sample-{n}' for n in range(64)) if counters._sampling_decision(i)[2]) + out_id = next(i for i in (f'sample-{n}' for n in range(64)) if not counters._sampling_decision(i)[2]) + # Same id, same answer: the decision is a pure function of the identity. + assert counters._sampling_decision(in_id) == counters._sampling_decision(in_id) - def test_failures_mode_keeps_failed_and_debug_traces_only(self): + excluded = trace.bind(out_id) + counters.record(CONTEXT, **STAGE) + assert excluded.state.emit_nodes is False + assert counters._chain_emitted(excluded.state) is False + counters.close_trace(excluded.state, 'event_done') + trace.unbind_root(excluded) + assert counters.records == [] + + included = trace.bind(in_id) + counters.record(CONTEXT, **STAGE) + assert included.state.emit_nodes is True + counters.close_trace(included.state, 'event_done') + trace.unbind_root(included) + assert [record['event_type'] for record in counters.records] == ['execution_node', 'execution_chain'] + + counters.records.clear() + broken = trace.bind(out_id) + counters.record(CONTEXT, **{**STAGE, 'outcome': 'timeout'}) + assert counters._chain_emitted(broken.state) is True + counters.close_trace(broken.state, 'event_done') + trace.unbind_root(broken) + assert [record['event_type'] for record in counters.records] == ['execution_chain'] + assert counters.records[0]['outcome'] == 'timeout' + + def test_failures_mode_keeps_failed_and_debug_chains_only(self): trace, _ = get_modules() - _, counters = make_counters(trace_config(execution_trace='failures', execution_trace_sample=1000000)) + _, counters = make_counters(trace_config(execution_trace='failures')) failed = trace.bind('exec-failed') counters.record(CONTEXT, **{**STAGE, 'outcome': 'failed'}) @@ -433,7 +467,8 @@ class TestTraceDelivery: trace.unbind_root(failed) debug = trace.bind('exec-debug') - counters.record(CONTEXT, **{**STAGE, 'synthetic': True}) + counters.configure(debug.state, debug=True) + counters.record(CONTEXT, **STAGE) counters.close_trace(debug.state, 'event_done') trace.unbind_root(debug) @@ -442,9 +477,12 @@ class TestTraceDelivery: counters.close_trace(sampled_out.state, 'event_done') trace.unbind_root(sampled_out) - assert [record['query_id'] for record in counters.records] == ['exec-failed', 'exec-debug'] - assert counters.records[0]['features']['observations'][0]['outcome'] == 'failed' - assert counters.records[1]['features']['observations'][0]['synthetic'] is True + assert [record['event_id'] for record in chains(counters)] == ['exec-failed', 'exec-debug'] + assert chains(counters)[0]['outcome'] == 'failed' + # Its node was sampled out, but the broken chain still explains itself. + assert chains(counters)[0]['dropped_nodes'] == 1 + assert chains(counters)[1]['debug'] is True + assert [record['event_id'] for record in nodes(counters)] == ['exec-debug'] def test_off_mode_and_opt_out_buffer_nothing(self): trace, _ = get_modules() @@ -469,24 +507,9 @@ class TestTraceDelivery: trace.unbind_root(binding) assert disabled_manager.sent == [] assert disabled.records == [] - assert disabled.traces == {} + assert disabled.registry == {} - def test_ttl_sweep_closes_stale_trace(self): - trace, _ = get_modules() - manager, counters = make_counters(trace_config()) - binding = trace.bind('exec-ttl') - try: - counters.record(CONTEXT, **{**STAGE, 'outcome': 'failed'}) - counters.trace_deadlines[binding.state.trace_id] = 0 - counters._sweep_traces(time.monotonic()) - finally: - trace.unbind_root(binding) - assert counters.records[0]['features']['trace']['closed_by'] == 'ttl' - assert counters.traces == {} - assert counters.trace_deadlines == {} - assert manager.sent == [] - - async def test_shutdown_drops_open_traces(self): + async def test_shutdown_does_not_claim_an_unfinished_execution(self): trace, _ = get_modules() manager, counters = make_counters(trace_config()) binding = trace.bind('exec-open') @@ -495,26 +518,12 @@ class TestTraceDelivery: await counters.shutdown() finally: trace.unbind_root(binding) - # Shutdown must not claim a trace that never completed. - assert records_sent(manager) == [] + # The already-emitted node is flushed, but shutdown never fabricates a + # terminal chain for an execution that never completed. + assert [record['event_type'] for record in records_sent(manager)] == ['execution_node'] assert counters.records == [] - assert counters.traces == {} + assert chains(counters) == [] assert counters.registry == {} - assert counters.dropped_traces == 1 - - async def test_loop_sweeps_and_flushes_expired_trace(self): - trace, _ = get_modules() - manager, counters = make_counters(trace_config(telemetry_flush_seconds=1)) - binding = trace.bind('exec-ttl-loop') - try: - counters.record(CONTEXT, **{**STAGE, 'outcome': 'failed'}) - counters.trace_deadlines[binding.state.trace_id] = 0 - await asyncio.sleep(1.3) - assert [record['features']['trace']['closed_by'] for record in records_sent(manager)] == ['ttl'] - assert counters.traces == {} - finally: - await counters.shutdown() - trace.unbind_root(binding) async def test_flush_sends_one_batch_of_complete_records(self): _, execution = get_modules() @@ -523,14 +532,20 @@ class TestTraceDelivery: for index in range(2): with execution.ingress(ap, 'event_done', CONTEXT, execution_id=f'exec-{index}'): counters.record(CONTEXT, **STAGE) - assert len(counters.records) == 2 + assert len(counters.records) == 4 await counters.flush() assert len(manager.sent) == 1 payload = manager.sent[0] assert set(payload) == {'records'} - assert [record['query_id'] for record in payload['records']] == ['exec-0', 'exec-1'] - assert 'count' not in json.dumps(payload) + assert [record['event_type'] for record in payload['records']] == [ + 'execution_node', + 'execution_chain', + 'execution_node', + 'execution_chain', + ] + assert [record['event_id'] for record in payload['records']] == ['exec-0', 'exec-0', 'exec-1', 'exec-1'] + await counters.shutdown() async def test_flush_interval_is_honoured(self): _, execution = get_modules() @@ -541,7 +556,9 @@ class TestTraceDelivery: assert counters.flush_seconds() == 1 assert manager.sent == [] await asyncio.sleep(1.3) - assert [record['query_id'] for record in records_sent(manager)] == ['exec-timed'] + sent = records_sent(manager) + assert [record['event_type'] for record in sent] == ['execution_node', 'execution_chain'] + assert {record['event_id'] for record in sent} == {'exec-timed'} await counters.shutdown() def test_flush_interval_config_is_bounded(self): @@ -563,11 +580,10 @@ class TestTraceDelivery: for index in range(5): binding = trace.bind(f'exec-{index}') counters.record(CONTEXT, **STAGE) - counters.close_trace(binding.state, 'event_done') trace.unbind_root(binding) finally: execution.MAX_BUFFERED_RECORDS = original - assert [record['query_id'] for record in counters.records] == ['exec-2', 'exec-3', 'exec-4'] + assert [record['event_id'] for record in nodes(counters)] == ['exec-2', 'exec-3', 'exec-4'] assert counters.dropped == 2 @@ -579,9 +595,9 @@ class TestIngress: with execution.ingress(ap, 'event_done', execution_id='exec-outer'): with execution.ingress(ap, 'pipeline_done', execution_id='exec-inner'): counters.record(CONTEXT, **{**STAGE, 'outcome': 'failed'}) - assert len(counters.records) == 1 - assert counters.records[0]['query_id'] == 'exec-outer' - assert counters.records[0]['features']['trace']['closed_by'] == 'event_done' + assert [record['event_type'] for record in counters.records] == ['execution_node', 'execution_chain'] + assert all(record['event_id'] == 'exec-outer' for record in counters.records) + assert chains(counters)[0]['closed_by'] == 'event_done' assert manager.sent == [] def test_ingress_without_telemetry_is_a_no_op(self): @@ -602,7 +618,7 @@ class TestIngress: class TestCrossTaskAttribution: - def test_registry_joins_stages_from_a_task_without_a_trace(self): + def test_registry_joins_nodes_from_a_task_without_a_trace(self): trace, execution = get_modules() _, counters = make_counters(trace_config()) ap = make_ap(counters) @@ -615,23 +631,24 @@ class TestCrossTaskAttribution: counters.record( CONTEXT, **{**STAGE, 'family': 'platform_api', 'operation': 'send_message'}, execution_id='run-1' ) - assert len(binding.state.stages) == 1 - assert binding.state.stages[0]['operation'] == 'send_message' + rows = nodes(counters) + assert len(rows) == 1 + assert rows[0]['operation'] == 'send_message' and rows[0]['event_id'] == 'run-1' counters.close_trace(binding.state, 'event_done') - assert counters.records[0]['query_id'] == 'run-1' + assert chains(counters)[0]['event_id'] == 'run-1' def test_registry_aliases_a_nested_execution(self): _, execution = get_modules() _, counters = make_counters(trace_config()) ap = make_ap(counters) - with execution.ingress(ap, 'event_done', CONTEXT, execution_id='event-1') as binding: + with execution.ingress(ap, 'event_done', CONTEXT, execution_id='event-1'): with execution.ingress(ap, 'pipeline_done', CONTEXT, execution_id='query-1'): pass assert set(counters.registry) == {'event-1', 'query-1'} counters.record(CONTEXT, **STAGE, execution_id='query-1') - assert counters.records[0]['query_id'] == 'event-1' - assert len(binding.state.stages) == 1 + rows = nodes(counters) + assert len(rows) == 1 and rows[0]['event_id'] == 'event-1' def test_contextvar_execution_id_is_used_when_no_kwarg(self): trace, execution = get_modules() @@ -644,7 +661,8 @@ class TestCrossTaskAttribution: counters.record(CONTEXT, **STAGE) finally: execution.reset_execution_id(token) - assert len(binding.state.stages) == 1 + rows = nodes(counters) + assert len(rows) == 1 and rows[0]['event_id'] == 'run-2' def test_registry_alias_is_swept_when_never_closed(self): trace, execution = get_modules() @@ -667,8 +685,10 @@ class TestTraceDispatch: with execution.ingress(ap, 'event_done', CONTEXT, execution_id='exec-async'): counters.record(CONTEXT, **STAGE) await counters.flush() - assert [record['query_id'] for record in records_sent(manager)] == ['exec-async'] - assert manager.sent[0]['records'][0]['event_type'] == 'feature_execution' + sent = records_sent(manager) + assert [record['event_type'] for record in sent] == ['execution_node', 'execution_chain'] + assert all(record['event_id'] == 'exec-async' for record in sent) + await counters.shutdown() class AsyncSendManager(FakeManager): @@ -689,9 +709,9 @@ class TestTraceFailureDetail: counters.close_trace(binding.state, 'event_done') finally: trace.unbind_root(binding) - record = counters.records[0] - assert record['error'] == 'boom: sandbox exited 127' - assert record['features']['observations'][0]['error'] == 'boom: sandbox exited 127' + assert nodes(counters)[0]['error'] == 'boom: sandbox exited 127' + # The detail lives on the node; the chain carries the terminal outcome. + assert chains(counters)[0]['outcome'] == 'failed' def test_successful_stage_reports_no_error(self): trace, _ = get_modules() @@ -702,9 +722,8 @@ class TestTraceFailureDetail: counters.close_trace(binding.state, 'event_done') finally: trace.unbind_root(binding) - record = counters.records[0] - assert record['error'] == '' - assert record['features']['observations'][0]['error'] == '' + assert nodes(counters)[0]['error'] == '' + assert chains(counters)[0]['error'] == '' def test_error_detail_is_bounded_and_single_line(self): trace, _ = get_modules() @@ -715,8 +734,8 @@ class TestTraceFailureDetail: counters.close_trace(binding.state, 'event_done') finally: trace.unbind_root(binding) - detail = counters.records[0]['error'] - assert len(detail) <= 400 + detail = chains(counters)[0]['error'] + assert len(detail) <= 200 assert '\n' not in detail and '\t' not in detail @@ -733,10 +752,10 @@ class TestIngressFailureDetail: except RuntimeError: raised = True assert raised is True - record = counters.records[0] - assert record['error'].startswith('RuntimeError: dispatch exploded') - assert record['features']['trace']['outcome'] == 'failed' - assert record['features']['trace']['closed_by'] == 'event_done' + chain = chains(counters)[0] + assert chain['error'].startswith('RuntimeError: dispatch exploded') + assert chain['outcome'] == 'failed' + assert chain['closed_by'] == 'event_done' def test_successful_ingress_still_reports_success(self): _, execution = get_modules() @@ -744,27 +763,29 @@ class TestIngressFailureDetail: ap = make_ap(counters) with execution.ingress(ap, 'event_done', execution_id='exec-fine'): counters.record(CONTEXT, **STAGE) - record = counters.records[0] - assert record['error'] == '' - assert record['features']['trace']['outcome'] == 'success' + chain = chains(counters)[0] + assert chain['error'] == '' + assert chain['outcome'] == 'success' class TestFailureOnlyTrace: def test_failure_without_stages_still_uploads(self): _, execution = get_modules() - # sampled-out config: only failures are guaranteed to be kept - _, counters = make_counters({'url': 'https://space.example.test', 'execution_trace_sample': 1}) + # failures mode: nodes are not reported, but a broken chain always is. + _, counters = make_counters({'url': 'https://space.example.test', 'execution_trace': 'failures'}) ap = make_ap(counters) try: with execution.ingress(ap, 'pipeline_done', CONTEXT, execution_id='exec-empty'): raise RuntimeError('pipeline exploded before any stage') except RuntimeError: pass - assert len(counters.records) == 1 - record = counters.records[0] + rows = chains(counters) + assert len(rows) == 1 + record = rows[0] assert record['error'].startswith('RuntimeError: pipeline exploded') - assert record['features']['trace']['outcome'] == 'failed' - assert record['features']['observations'] == [] + assert record['outcome'] == 'failed' + assert record['stage_count'] == 0 + assert nodes(counters) == [] def test_cancelled_and_timeout_chains_report_their_outcome(self): trace, _ = get_modules() @@ -776,15 +797,14 @@ class TestFailureOnlyTrace: counters.close_trace(binding.state, 'runner_done') finally: trace.unbind_root(binding) - assert len(counters.records) == 1 - record = counters.records.pop(0) - assert record['query_id'] == f'exec-{outcome}' - assert record['features']['trace']['outcome'] == outcome - assert record['error'] == reason + record = chains(counters)[-1] + assert record['event_id'] == f'exec-{outcome}' + assert record['outcome'] == outcome + assert nodes(counters)[-1]['error'] == reason class TestLegacyRecordFields: - def test_record_carries_legacy_top_level_columns(self): + def test_chain_carries_execution_scoped_columns(self): trace, _ = get_modules() _, counters = make_counters(trace_config()) binding = trace.bind('exec-legacy') @@ -800,14 +820,13 @@ class TestLegacyRecordFields: counters.close_trace(binding.state, 'event_done') finally: trace.unbind_root(binding) - record = counters.records[0] + record = chains(counters)[0] assert record['model_name'] == 'gpt-4o' assert record['pipeline_plugins'] == ['plugin:a'] assert record['runner_category'] == 'cloud' assert record['adapter'] == 'AiocqhttpAdapter' assert record['runner'] == 'runner-1' assert record['duration_ms'] >= 1000 - assert record['timestamp'] == record['features']['trace']['ended_at'] def test_legacy_string_fields_are_never_null(self): trace, _ = get_modules() @@ -818,9 +837,135 @@ class TestLegacyRecordFields: counters.close_trace(binding.state, 'event_done') finally: trace.unbind_root(binding) - record = counters.records[0] + record = chains(counters)[0] assert record['model_name'] == '' assert record['runner'] == '' assert record['runner_category'] == '' assert isinstance(record['duration_ms'], int) and record['duration_ms'] >= 0 assert record['pipeline_plugins'] is None + + +class TestNoTraceTtl: + def test_nodes_more_than_two_minutes_apart_still_emit(self): + trace, _ = get_modules() + _, counters = make_counters(trace_config()) + binding = trace.bind('exec-long') + try: + counters.record(CONTEXT, **STAGE) + # A gap far beyond the old 120s TTL must not sweep the chain: there is + # no local trace deadline any more, so both nodes and the terminal + # chain are still emitted. + binding.state.started_at = (datetime.now(timezone.utc) - timedelta(seconds=200)).isoformat() + counters.record(CONTEXT, **{**STAGE, 'operation': 'still.alive'}) + counters.close_trace(binding.state, 'event_done') + finally: + trace.unbind_root(binding) + rows = nodes(counters) + assert [record['seq'] for record in rows] == [0, 1] + assert [record['operation'] for record in rows] == ['message.received', 'still.alive'] + chain = chains(counters)[0] + assert chain['stage_count'] == 2 + assert chain['duration_ms'] >= 120_000 + + +class TestFlushDelivery: + async def test_rejected_batch_stays_buffered_and_is_retried(self): + trace, execution = get_modules() + manager = RejectingManager(trace_config()) + counters = execution.ExecutionCounters(manager) + binding = trace.bind('exec-retry') + try: + counters.record(CONTEXT, **STAGE) + assert await counters.flush() is False + assert manager.attempts == execution.MAX_FLUSH_ATTEMPTS + # The record was never lost: it is still buffered and unsent. + assert len(counters.records) == 1 and manager.sent == [] + manager.rejecting = False + assert await counters.flush() is True + finally: + trace.unbind_root(binding) + assert counters.records == [] + assert [record['event_id'] for record in records_sent(manager)] == ['exec-retry'] + await counters.shutdown() + + async def test_health_reports_buffered_sent_and_dropped(self): + trace, execution = get_modules() + _, counters = make_counters(trace_config()) + ap = make_ap(counters) + binding = execution.bind_trace(ap, 'exec-health') + try: + counters.record(CONTEXT, **STAGE) + health = counters.health() + assert health['records_buffered'] == 1 + assert health['records_sent'] == 0 + assert health['records_dropped'] == 0 + assert health['executions_in_flight'] == 1 + await counters.flush() + health = counters.health() + assert health['records_buffered'] == 0 + assert health['records_sent'] == 1 + finally: + trace.unbind_root(binding) + + # Overflowing the outbound buffer is surfaced as dropped records. + original = execution.MAX_BUFFERED_RECORDS + execution.MAX_BUFFERED_RECORDS = 1 + try: + for index in range(3): + inner = trace.bind(f'exec-drop-{index}') + counters.record(CONTEXT, **STAGE) + trace.unbind_root(inner) + finally: + execution.MAX_BUFFERED_RECORDS = original + health = counters.health() + assert health['records_dropped'] == 2 + assert health['records_buffered'] == 1 + await counters.shutdown() + + +class TestSiblingTaskIsolation: + async def test_sibling_stage_scopes_do_not_corrupt_parent_links(self): + trace, execution = get_modules() + _, counters = make_counters(trace_config()) + ap = make_ap(counters) + opened = {} + + async def work(name): + with trace.stage_scope() as node: + opened[name] = node + counters.record(CONTEXT, **{**STAGE, 'operation': name, 'node': node}) + + async def scenario(): + with execution.ingress(ap, 'event_done', CONTEXT, execution_id='exec-siblings'): + await asyncio.gather(work('alpha'), work('beta')) + snapshot = list(counters.records) + await counters.shutdown() + return snapshot + + records = await scenario() + rows = {record['operation']: record for record in records if record['event_type'] == 'execution_node'} + assert opened['alpha'] != opened['beta'] + # Each sibling's step is its own root: neither becomes the other's parent. + assert rows['alpha']['node_id'] == opened['alpha'] + assert rows['beta']['node_id'] == opened['beta'] + assert rows['alpha']['parent_node_id'] == '' + assert rows['beta']['parent_node_id'] == '' + assert [record['event_type'] for record in records].count('execution_chain') == 1 + + +class TestSamplingShape: + def test_platform_prefixed_event_id_is_emitted_at_default_config(self): + trace, _ = get_modules() + _, counters = make_counters() # no execution_trace override -> default mode + assert counters.trace_mode() == 'all' + # The identity is hashed, not parsed: a platform-prefixed id samples the + # same way any other id shape does. + assert counters._sampling_decision('platform:bot-1:message-1') == ('all', 1, True) + binding = trace.bind('platform:bot-1:message-1') + try: + counters.record(CONTEXT, **STAGE) + counters.close_trace(binding.state, 'event_done') + finally: + trace.unbind_root(binding) + assert [record['event_type'] for record in counters.records] == ['execution_node', 'execution_chain'] + assert all(record['event_id'] == 'platform:bot-1:message-1' for record in counters.records) diff --git a/tests/unit_tests/test_telemetry.py b/tests/unit_tests/test_telemetry.py deleted file mode 100644 index ed3994923..000000000 --- a/tests/unit_tests/test_telemetry.py +++ /dev/null @@ -1,29 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace -import asyncio - -import pytest - -from langbot.pkg.telemetry.telemetry import TelemetryManager - - -@pytest.mark.asyncio -async def test_send_tasks_are_scoped_to_manager_instance(monkeypatch): - async def fake_send(self, payload): - return payload - - monkeypatch.setattr(TelemetryManager, 'send', fake_send) - - first = TelemetryManager(SimpleNamespace()) - second = TelemetryManager(SimpleNamespace()) - - assert first.send_tasks is not second.send_tasks - - await first.start_send_task({'event': 'first'}) - task = first.send_tasks[0] - await task - await asyncio.sleep(0) - - assert first.send_tasks == [] - assert second.send_tasks == []