From 5b8376a55b032f91bf0c9e9511c56662916e7e52 Mon Sep 17 00:00:00 2001 From: Naseem AlNaji Date: Wed, 8 Jul 2026 16:42:09 -0400 Subject: [PATCH 1/3] feat: custom events, PostHog exporter, async identify, and hardened event pipeline Features: - publish_custom_event(): publish agentcat:custom events against a tracked server or an MCP session ID string, with support for message/user intent, parameters, response, duration, error state, tags, and properties. Session IDs derived from MCP session IDs are deterministic (SHA-256 -> KSUID), so custom events correlate with SDK-tracked sessions across languages. - PostHog telemetry exporter: /batch capture events, person profile updates, $exception events on errors, and optional AI-observability spans behind enable_ai_tracing. Session IDs mapped to deterministic UUIDv7. - identify hooks may now be async; sync hooks keep working unchanged. - Identity change detection: identify events are published only when the resolved identity actually changes; user_data is merged across calls and identities are cached (LRU, 1000 entries). Fixes: - redact_sensitive_information is now applied to live events. Previously the redaction pass silently skipped the event model, so custom redaction functions never ran on published events. Async redaction functions are now supported and awaited. - The low-level Server integration now captures full structured error data (stack frames, chained exceptions) instead of only the error message. - Event truncation applies per-field limits (user intent, error message, resource and server/client names, stack frames, response text blocks) and raises general limits to match the published event size contract (32KB per string, depth 10, breadth 100, 100KB total). Reliability: - The request path is fully fail-open: failures in identify hooks, tag and property callbacks, error capture, or event publishing are logged and dropped without affecting the customer's tool call - including exceptions whose own __str__ raises. - The queue worker survives poisoned events (cyclic structures, NaN/Inf, bytes, objects with raising __repr__) and continues delivering subsequent events. - publish_custom_event raises only its documented ValueError/TypeError contract on invalid arguments. Tests: 611 passed, 4 skipped, including new failure-injection, custom event, PostHog, identify, and truncation suites. --- src/agentcat/__init__.py | 116 +++- src/agentcat/modules/constants.py | 1 + src/agentcat/modules/event_queue.py | 82 ++- src/agentcat/modules/exceptions.py | 40 +- src/agentcat/modules/exporters/posthog.py | 299 +++++++++ src/agentcat/modules/identify.py | 143 +++- src/agentcat/modules/internal.py | 48 +- src/agentcat/modules/logging.py | 17 + .../overrides/community/monkey_patch.py | 10 +- .../overrides/community_v3/middleware.py | 16 +- src/agentcat/modules/overrides/mcp_server.py | 26 +- .../overrides/official/monkey_patch.py | 8 +- src/agentcat/modules/redaction.py | 62 +- src/agentcat/modules/session.py | 45 ++ src/agentcat/modules/telemetry.py | 4 + src/agentcat/modules/truncation.py | 129 +++- src/agentcat/modules/validation.py | 85 ++- src/agentcat/types.py | 66 +- tests/e2e/official/test_redaction_http.py | 30 +- tests/test_failure_injection.py | 633 ++++++++++++++++++ tests/test_identify.py | 269 ++++++++ tests/test_lowlevel_error_capture.py | 175 +++++ tests/test_posthog_exporter.py | 269 ++++++++ tests/test_publish_custom_event.py | 166 +++++ tests/test_redaction.py | 131 ++++ tests/test_stateless.py | 24 +- tests/test_truncation.py | 189 +++++- 27 files changed, 2897 insertions(+), 186 deletions(-) create mode 100644 src/agentcat/modules/exporters/posthog.py create mode 100644 tests/test_failure_injection.py create mode 100644 tests/test_identify.py create mode 100644 tests/test_lowlevel_error_capture.py create mode 100644 tests/test_posthog_exporter.py create mode 100644 tests/test_publish_custom_event.py diff --git a/src/agentcat/__init__.py b/src/agentcat/__init__.py index 933c5f3..5fe3778 100644 --- a/src/agentcat/__init__.py +++ b/src/agentcat/__init__.py @@ -9,7 +9,11 @@ __version__ = version("agentcat") from agentcat.modules.overrides.mcp_server import override_lowlevel_mcp_server -from agentcat.modules.session import get_session_info, new_session_id +from agentcat.modules.session import ( + derive_session_id_from_mcp_session, + get_session_info, + new_session_id, +) from .modules.compatibility import ( COMPATIBILITY_ERROR_MESSAGE, @@ -18,16 +22,20 @@ is_compatible_server, is_official_fastmcp_server, ) +from .modules.constants import AGENTCAT_CUSTOM_EVENT_TYPE from .modules.diagnostics import init_diagnostics -from .modules.internal import set_server_tracking_data +from .modules.internal import get_server_tracking_data, set_server_tracking_data from .modules.logging import set_debug_mode, write_to_log +from .modules.validation import validate_tags from .types import ( + CustomEventData, EventPropertiesFunction, EventTagsFunction, IdentifyFunction, AgentCatData, AgentCatOptions, RedactionFunction, + UnredactedEvent, UserIdentity, ) @@ -248,9 +256,111 @@ def _apply_server_tracking( override_lowlevel_mcp_server(lowlevel_server, data) +def publish_custom_event( + server_or_session_id: Any, + project_id: str, + event_data: CustomEventData | None = None, +) -> None: + """ + Publish a custom event to AgentCat with flexible session management. + + Args: + server_or_session_id: Either a tracked MCP server instance or an MCP + session ID string. For a session ID string, a deterministic + AgentCat session ID is derived so the same inputs always correlate + to the same session. + project_id: Your AgentCat project ID (required) + event_data: Optional event data to include with the custom event + + Raises: + ValueError: If project_id is missing, or the server is not tracked + TypeError: If the first parameter is neither a server nor a string + + Example: + >>> import agentcat + >>> data = agentcat.CustomEventData( + ... resource_name="custom-action", + ... parameters={"action": "user-feedback", "rating": 5}, + ... message="User provided feedback", + ... ) + >>> agentcat.publish_custom_event(server, "proj_abc123", data) + """ + from agentcat.modules import event_queue as event_queue_module + + if not project_id: + raise ValueError("project_id is required for publish_custom_event") + + tracked_server: Any | None = None + + if isinstance(server_or_session_id, str): + # Custom session ID provided - derive a deterministic session ID + session_id = derive_session_id_from_mcp_session( + server_or_session_id, project_id + ) + elif server_or_session_id is not None and not isinstance( + server_or_session_id, (int, float, bool) + ): + try: + data = get_server_tracking_data(server_or_session_id) + except TypeError: + # Non-weakref-able / unhashable objects can't be tracked servers; + # surface the documented "not tracked" error instead of leaking a + # WeakKeyDictionary internal TypeError. + data = None + if data is None: + # Server is not tracked - treat it as an error + raise ValueError( + "Server is not tracked. Please call agentcat.track() first or " + "provide a session ID string." + ) + # Use the tracked server's session ID and configuration + tracked_server = server_or_session_id + session_id = data.session_id + else: + raise TypeError( + "First parameter must be either an MCP server object or a session ID string" + ) + + event = UnredactedEvent( + # Core fields + session_id=session_id, + project_id=project_id, + # Fixed event type for custom events + event_type=AGENTCAT_CUSTOM_EVENT_TYPE, + timestamp=datetime.now(timezone.utc), + # Event data from parameters + resource_name=event_data.resource_name if event_data else None, + parameters=event_data.parameters if event_data else None, + response=event_data.response if event_data else None, + user_intent=event_data.message if event_data else None, + duration=event_data.duration if event_data else None, + is_error=event_data.is_error if event_data else None, + error=event_data.error if event_data else None, + ) + + # Wire up customer-defined metadata + if event_data and event_data.tags: + event.tags = validate_tags(event_data.tags) + if event_data and event_data.properties: + event.properties = event_data.properties + + # If we have a tracked server, publish through it (merges session info and + # redaction config); otherwise add directly to the event queue. + if tracked_server is not None: + event_queue_module.publish_event(tracked_server, event) + else: + event_queue_module.event_queue.add(event) + + write_to_log( + f"Published custom event for session {session_id} with type " + f"'{AGENTCAT_CUSTOM_EVENT_TYPE}'" + ) + + __all__ = [ # Main API "track", + "publish_custom_event", # Configuration "AgentCatOptions", # Types for identify functionality @@ -261,4 +371,6 @@ def _apply_server_tracking( # Types for event metadata callbacks "EventTagsFunction", "EventPropertiesFunction", + # Type for custom events + "CustomEventData", ] diff --git a/src/agentcat/modules/constants.py b/src/agentcat/modules/constants.py index 8146fb4..e8c0a86 100644 --- a/src/agentcat/modules/constants.py +++ b/src/agentcat/modules/constants.py @@ -4,6 +4,7 @@ EVENT_ID_PREFIX = "evt" AGENTCAT_API_URL = "https://api.agentcat.com" # Default API URL for AgentCat events AGENTCAT_SOURCE = "agentcat" # Source attribution for telemetry exporters +AGENTCAT_CUSTOM_EVENT_TYPE = "agentcat:custom" # Event type for publish_custom_event DEFAULT_CONTEXT_DESCRIPTION = "Explain why you are calling this tool and how it fits into the user's overall goal. This parameter is used for analytics and user intent tracking. YOU MUST provide 15-25 words (count carefully). NEVER use first person ('I', 'we', 'you') - maintain third-person perspective. NEVER include sensitive information such as credentials, passwords, or personal data. Example (20 words): \"Searching across the organization's repositories to find all open issues related to performance complaints and latency issues for team prioritization.\"" # Maximum number of exceptions to capture in a cause chain diff --git a/src/agentcat/modules/event_queue.py b/src/agentcat/modules/event_queue.py index 473ab2e..f1135f4 100644 --- a/src/agentcat/modules/event_queue.py +++ b/src/agentcat/modules/event_queue.py @@ -20,7 +20,7 @@ from ..utils import generate_prefixed_ksuid from .compatibility import get_mcp_compatible_error_message from .internal import get_server_tracking_data -from .logging import write_to_log +from .logging import safe_error_string, write_to_log from .redaction import redact_event from .sanitization import sanitize_event from .truncation import truncate_event @@ -115,8 +115,10 @@ def _process_event(self, event: UnredactedEvent) -> None: event = redacted_event event.redaction_fn = None # Clear the function to avoid reprocessing except Exception as error: + # safe_error_string: the error comes from the customer's + # redaction function and may have a broken __str__. write_to_log( - f"WARNING: Dropping event {event.id or 'unknown'} due to redaction failure: {error}" + f"WARNING: Dropping event {event.id or 'unknown'} due to redaction failure: {safe_error_string(error)}" ) return # Skip this event if redaction fails @@ -287,40 +289,54 @@ def set_event_queue(new_queue: EventQueue) -> None: def publish_event(server: Any, event: UnredactedEvent) -> None: - """Publish an event to the queue.""" - if not event.duration: - if event.timestamp: - event.duration = int( - (datetime.now(timezone.utc).timestamp() - event.timestamp.timestamp()) - * 1000 - ) - else: - event.duration = None + """Publish an event to the queue. - data = get_server_tracking_data(server) - if not data: - write_to_log( - "Warning: Server tracking data not found. Event will not be published." - ) - return + Never raises: this is called from the customer's request path (tool-call + and initialize wrappers). Any failure degrades to a dropped event plus a + log line. + """ + try: + if not event.duration: + if event.timestamp: + event.duration = int( + ( + datetime.now(timezone.utc).timestamp() + - event.timestamp.timestamp() + ) + * 1000 + ) + else: + event.duration = None - session_info = get_session_info(server, data) + data = get_server_tracking_data(server) + if not data: + write_to_log( + "Warning: Server tracking data not found. Event will not be published." + ) + return - # Create full event with all required fields - # Merge event data with session info - event_data = event.model_dump(exclude_none=True) - session_data = session_info.model_dump(exclude_none=True) + session_info = get_session_info(server, data) - # Merge data, ensuring project_id from data takes precedence - merged_data = {**event_data, **session_data} - merged_data["project_id"] = ( - data.project_id - ) # Override with tracking data's project_id + # Create full event with all required fields + # Merge event data with session info + event_data = event.model_dump(exclude_none=True) + session_data = session_info.model_dump(exclude_none=True) - full_event = UnredactedEvent( - **merged_data, - redaction_fn=data.options.redact_sensitive_information, - ) + # Merge data, ensuring project_id from data takes precedence + merged_data = {**event_data, **session_data} + merged_data["project_id"] = ( + data.project_id + ) # Override with tracking data's project_id - set_last_activity(server) - event_queue.add(full_event) + full_event = UnredactedEvent( + **merged_data, + redaction_fn=data.options.redact_sensitive_information, + ) + + set_last_activity(server) + event_queue.add(full_event) + except Exception as e: + write_to_log( + f"WARNING: Dropping event {getattr(event, 'id', None) or 'unknown'} — " + f"publish failed: {safe_error_string(e)}" + ) diff --git a/src/agentcat/modules/exceptions.py b/src/agentcat/modules/exceptions.py index c540d38..56b1727 100644 --- a/src/agentcat/modules/exceptions.py +++ b/src/agentcat/modules/exceptions.py @@ -17,6 +17,17 @@ ) +def _safe_str(value: Any) -> str: + """str(value) that never raises (exceptions may have a broken __str__).""" + try: + return str(value) + except Exception: + try: + return f"" + except Exception: + return "" + + def capture_exception(exc: BaseException | Any) -> ErrorData: """ Captures detailed exception information including stack traces and cause chains. @@ -26,12 +37,33 @@ def capture_exception(exc: BaseException | Any) -> ErrorData: tracebacks into structured frames and detects whether each frame is user code (in_app: True) or library code (in_app: False). + This function must NEVER raise — it runs on the customer's request path + (tool-call wrappers) where an escaping exception would corrupt the + customer's request handling. Any internal failure degrades to a minimal + ErrorData dict. + Args: exc: The error to capture (can be BaseException, string, object, or any value) Returns: ErrorData dict with structured error information including platform="python" """ + try: + return _capture_exception_unsafe(exc) + except Exception: + try: + type_name = type(exc).__name__ if isinstance(exc, BaseException) else None + except Exception: + type_name = None + return { + "message": _safe_str(exc), + "type": type_name, + "platform": "python", + } + + +def _capture_exception_unsafe(exc: BaseException | Any) -> ErrorData: + """Implementation of capture_exception; see its docstring.""" if is_call_tool_result(exc): return capture_call_tool_result_error(exc) @@ -43,7 +75,7 @@ def capture_exception(exc: BaseException | Any) -> ErrorData: } error_data: ErrorData = { - "message": str(exc), + "message": _safe_str(exc), "type": type(exc).__name__, "platform": "python", } @@ -281,7 +313,7 @@ def unwrap_exception_chain(exc: BaseException) -> list[ChainedErrorData]: break chained_data: ChainedErrorData = { - "message": str(next_exc), + "message": _safe_str(next_exc), "type": type(next_exc).__name__, } @@ -315,7 +347,7 @@ def format_exception_string(exc: BaseException) -> str: try: return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)) except Exception: - return f"{type(exc).__name__}: {exc}" + return f"{type(exc).__name__}: {_safe_str(exc)}" def is_call_tool_result(value: Any) -> bool: @@ -403,7 +435,7 @@ def stringify_non_exception(value: Any) -> str: return json.dumps(value) except Exception: - return str(value) + return _safe_str(value) def store_captured_error(exc: BaseException) -> None: diff --git a/src/agentcat/modules/exporters/posthog.py b/src/agentcat/modules/exporters/posthog.py new file mode 100644 index 0000000..c1d7b23 --- /dev/null +++ b/src/agentcat/modules/exporters/posthog.py @@ -0,0 +1,299 @@ +"""PostHog exporter for AgentCat telemetry.""" + +import hashlib +import re +import time +from datetime import datetime, timezone +from typing import Any + +import requests + +from ...modules.constants import AGENTCAT_SOURCE +from ...modules.logging import write_to_log +from ...thirdparty.ksuid import Ksuid +from ...types import Event, EventType, PostHogExporterConfig +from . import Exporter + +DEFAULT_POSTHOG_HOST = "https://us.i.posthog.com" + +_PREFIX_RE = re.compile(r"^[a-z]+_") + +# Map AgentCat event types to PostHog event names +_EVENT_NAME_MAPPING: dict[str, str] = { + EventType.MCP_TOOLS_CALL.value: "mcp_tool_call", + EventType.MCP_TOOLS_LIST.value: "mcp_tools_list", + EventType.MCP_INITIALIZE.value: "mcp_initialize", + EventType.MCP_RESOURCES_READ.value: "mcp_resource_read", + EventType.MCP_RESOURCES_LIST.value: "mcp_resources_list", + EventType.MCP_PROMPTS_GET.value: "mcp_prompt_get", + EventType.MCP_PROMPTS_LIST.value: "mcp_prompts_list", +} + + +def to_uuidv7(prefixed_id: str) -> str: + """Generate a deterministic UUIDv7 from a prefixed KSUID (e.g. ses_xxx). + + Uses the KSUID's embedded timestamp for the UUIDv7 timestamp portion and + a SHA-256 hash of the full ID for the random bits (mirrors the TypeScript + SDK's toUUIDv7 for cross-SDK determinism). + """ + prefixed_id = prefixed_id or "" + # Strip prefix (ses_, evt_, etc.) and parse KSUID + ksuid_str = _PREFIX_RE.sub("", prefixed_id) + try: + ksuid = Ksuid.from_base62(ksuid_str) + timestamp_ms = int(ksuid.timestamp * 1000) + except Exception: + # Fallback: if KSUID parsing fails, use current time + timestamp_ms = int(time.time() * 1000) + + # Hash the full ID for deterministic random bits + digest = hashlib.sha256(prefixed_id.encode("utf-8")).digest() + + buf = bytearray(16) + + # Bytes 0-5: 48-bit Unix timestamp in milliseconds + buf[0:6] = timestamp_ms.to_bytes(6, "big") + + # Byte 6: version 7 (0111) + high 4 bits of rand_a from hash + buf[6] = 0x70 | (digest[0] & 0x0F) + # Byte 7: low 8 bits of rand_a from hash + buf[7] = digest[1] + + # Byte 8: variant 10 + high 6 bits of rand_b from hash + buf[8] = 0x80 | (digest[2] & 0x3F) + # Bytes 9-15: remaining rand_b from hash + buf[9:16] = digest[3:10] + + hex_str = buf.hex() + return "-".join( + [ + hex_str[0:8], + hex_str[8:12], + hex_str[12:16], + hex_str[16:20], + hex_str[20:32], + ] + ) + + +def _get_distinct_id(event: Event) -> str: + return event.identify_actor_given_id or event.session_id or "anonymous" + + +def _get_timestamp(event: Event) -> str: + timestamp = event.timestamp or datetime.now(timezone.utc) + return timestamp.isoformat() + + +class PostHogExporter(Exporter): + """Exports AgentCat events to PostHog via the /batch capture API.""" + + def __init__(self, config: PostHogExporterConfig): + """ + Initialize PostHog exporter. + + Args: + config: PostHog exporter configuration + """ + self.config = config + host = (config.get("host") or DEFAULT_POSTHOG_HOST).rstrip("/") + self.batch_url = f"{host}/batch" + self.api_key = config["api_key"] + self.enable_ai_tracing = config.get("enable_ai_tracing", False) + + # Create session for connection pooling + self.session = requests.Session() + + write_to_log(f"PostHogExporter: Initialized with endpoint {self.batch_url}") + + def export(self, event: Event) -> None: + """ + Export an event to PostHog. + + Args: + event: AgentCat event to export + """ + try: + batch: list[dict[str, Any]] = [] + + # Always send the regular event + batch.append(self.build_capture_event(event)) + + # Send $exception event alongside if this is an error + if event.is_error and event.error: + batch.append(self.build_exception_event(event)) + + # Send $ai_span for tool calls when AI tracing is enabled + if ( + self.enable_ai_tracing + and event.event_type == EventType.MCP_TOOLS_CALL.value + ): + batch.append(self.build_ai_span_event(event)) + + write_to_log( + f"PostHogExporter: Sending {len(batch)} event(s) for {event.id}" + ) + + response = self.session.post( + self.batch_url, + headers={"Content-Type": "application/json"}, + json={"api_key": self.api_key, "batch": batch}, + timeout=10, + ) + + if not response.ok: + write_to_log( + f"PostHog export failed - Status: {response.status_code}, " + f"Body: {response.text}" + ) + else: + write_to_log(f"PostHog export success - Event: {event.id}") + except Exception as error: + write_to_log(f"PostHog export error: {error}") + + def build_capture_event(self, event: Event) -> dict[str, Any]: + """Build the regular PostHog capture event.""" + properties: dict[str, Any] = { + "$session_id": to_uuidv7(event.session_id), + "source": AGENTCAT_SOURCE, + } + + if event.resource_name: + properties["resource_name"] = event.resource_name + if event.event_type == EventType.MCP_TOOLS_CALL.value: + properties["tool_name"] = event.resource_name + if event.duration is not None: + properties["duration_ms"] = event.duration + if event.server_name: + properties["server_name"] = event.server_name + if event.server_version: + properties["server_version"] = event.server_version + if event.client_name: + properties["client_name"] = event.client_name + if event.client_version: + properties["client_version"] = event.client_version + if event.project_id: + properties["project_id"] = event.project_id + if event.user_intent: + properties["user_intent"] = event.user_intent + if event.is_error is not None: + properties["is_error"] = event.is_error + + if event.parameters is not None: + properties["parameters"] = event.parameters + if event.response is not None: + properties["response"] = event.response + + # Set person properties from identity data + person_props: dict[str, Any] = {} + if event.identify_actor_name: + person_props["name"] = event.identify_actor_name + if event.identify_data: + person_props.update(event.identify_data) + if person_props: + properties["$set"] = person_props + + # Spread customer-defined tags directly (can override AgentCat defaults) + if event.tags: + properties.update(event.tags) + + # Spread customer-defined properties directly (can override AgentCat defaults) + if event.properties: + properties.update(event.properties) + + return { + "event": self.map_event_type(event.event_type), + "distinct_id": _get_distinct_id(event), + "properties": properties, + "timestamp": _get_timestamp(event), + "type": "capture", + } + + def build_exception_event(self, event: Event) -> dict[str, Any]: + """Build a PostHog $exception event for error events.""" + properties: dict[str, Any] = { + "$exception_source": "backend", + "$session_id": to_uuidv7(event.session_id), + } + + error = event.error or {} + if isinstance(error, dict): + if error.get("message"): + properties["$exception_message"] = error["message"] + if error.get("type"): + properties["$exception_type"] = error["type"] + if error.get("stack"): + properties["$exception_stacktrace"] = error["stack"] + + # Add tool/resource context + if event.resource_name: + properties["resource_name"] = event.resource_name + if event.event_type == EventType.MCP_TOOLS_CALL.value: + properties["tool_name"] = event.resource_name + if event.server_name: + properties["server_name"] = event.server_name + if event.server_version: + properties["server_version"] = event.server_version + if event.client_name: + properties["client_name"] = event.client_name + if event.client_version: + properties["client_version"] = event.client_version + + return { + "event": "$exception", + "distinct_id": _get_distinct_id(event), + "properties": properties, + "timestamp": _get_timestamp(event), + "type": "capture", + } + + def build_ai_span_event(self, event: Event) -> dict[str, Any]: + """Build a PostHog $ai_span event for AI observability views.""" + properties: dict[str, Any] = { + "$ai_session_id": f"agentcat_{event.session_id}", + "$ai_trace_id": to_uuidv7(event.session_id), + "$ai_span_id": to_uuidv7(event.id), + "$ai_span_name": event.resource_name or "unknown_tool", + "$ai_is_error": event.is_error or False, + "$session_id": to_uuidv7(event.session_id), + "source": AGENTCAT_SOURCE, + } + + if event.duration is not None: + properties["$ai_latency"] = event.duration / 1000 + if event.is_error and event.error: + properties["$ai_error"] = event.error + if event.parameters is not None: + properties["$ai_input_state"] = event.parameters + if event.response is not None: + properties["$ai_output_state"] = event.response + if event.server_name: + properties["server_name"] = event.server_name + if event.client_name: + properties["client_name"] = event.client_name + + # Spread customer tags directly (can override AgentCat defaults) + if event.tags: + properties.update(event.tags) + + # Spread customer properties directly (can override AgentCat defaults) + if event.properties: + properties.update(event.properties) + + return { + "event": "$ai_span", + "distinct_id": _get_distinct_id(event), + "properties": properties, + "timestamp": _get_timestamp(event), + "type": "capture", + } + + def map_event_type(self, event_type: str | None) -> str: + """Map AgentCat event types to PostHog event names.""" + event_type = event_type or "unknown" + mapped = _EVENT_NAME_MAPPING.get(event_type) + if mapped: + return mapped + stripped = event_type.removeprefix("mcp:").replace("/", "_") + return f"mcp_{stripped}" diff --git a/src/agentcat/modules/identify.py b/src/agentcat/modules/identify.py index 9f002fa..9d4e578 100644 --- a/src/agentcat/modules/identify.py +++ b/src/agentcat/modules/identify.py @@ -1,15 +1,104 @@ +import inspect +from collections import OrderedDict from datetime import datetime, timezone +from typing import Optional from agentcat.modules import event_queue from agentcat.modules.internal import get_server_tracking_data -from agentcat.modules.logging import write_to_log +from agentcat.modules.logging import safe_error_string, write_to_log from agentcat.types import EventType, UnredactedEvent, UserIdentity +# Maximum number of session identities retained in the global cache. +IDENTITY_CACHE_MAX_SIZE = 1000 -def identify_session(server, request: any, context: any) -> UserIdentity | None: - """Run the configured identify hook and publish an `agentcat:identify` event. - Returns the resulting UserIdentity, or None if no hook is configured, the +class IdentityCache: + """Simple LRU cache for session identities. + + Prevents memory leaks by capping at max_size entries. This cache persists + across server instance restarts (mirrors the TypeScript SDK's + IdentityCache in modules/internal.ts). + """ + + def __init__(self, max_size: int = IDENTITY_CACHE_MAX_SIZE): + self._cache: OrderedDict[str, UserIdentity] = OrderedDict() + self._max_size = max_size + + def get(self, session_id: str) -> Optional[UserIdentity]: + identity = self._cache.get(session_id) + if identity is not None: + # Move to end (most recently used) + self._cache.move_to_end(session_id) + return identity + + def set(self, session_id: str, identity: UserIdentity) -> None: + if session_id in self._cache: + # Remove so re-adding places it at the end + del self._cache[session_id] + elif len(self._cache) >= self._max_size: + # Evict least recently used + self._cache.popitem(last=False) + self._cache[session_id] = identity + + def __contains__(self, session_id: str) -> bool: + return session_id in self._cache + + def __len__(self) -> int: + return len(self._cache) + + +# Global identity cache shared across all server instances. +# This prevents duplicate identify events when server objects are recreated. +_global_identity_cache = IdentityCache() + + +def reset_identity_cache() -> None: + """Reset the global identity cache (mainly for testing).""" + global _global_identity_cache + _global_identity_cache = IdentityCache() + + +def are_identities_equal(a: UserIdentity, b: UserIdentity) -> bool: + """Deep comparison of two UserIdentity objects.""" + if a.user_id != b.user_id: + return False + if a.user_name != b.user_name: + return False + return (a.user_data or {}) == (b.user_data or {}) + + +def merge_identities( + previous: UserIdentity | None, next_identity: UserIdentity +) -> UserIdentity: + """Merge two UserIdentity objects. + + Overwrites user_id and user_name with the newest values, but merges + user_data keys across calls (newest values win on key collisions). + """ + if previous is None: + return next_identity + + return UserIdentity( + user_id=next_identity.user_id, + user_name=next_identity.user_name, + user_data={ + **(previous.user_data or {}), + **(next_identity.user_data or {}), + }, + ) + + +async def identify_session(server, request: any, context: any) -> UserIdentity | None: + """Run the configured identify hook for a request. + + The identify hook may be a sync or an async callable; coroutine results + are awaited. The resolved identity is merged with the session's previous + identity (user_id/user_name overwritten, user_data keys merged) and an + `agentcat:identify` event is published ONLY when the merged identity + differs from the cached one — repeated identical identities do not + republish. + + Returns the merged UserIdentity, or None if no hook is configured, the hook raises, or it returns a non-UserIdentity value. """ data = get_server_tracking_data(server) @@ -19,6 +108,8 @@ def identify_session(server, request: any, context: any) -> UserIdentity | None: try: identify_result = data.options.identify(request, context) + if inspect.iscoroutine(identify_result): + identify_result = await identify_result if not identify_result or not isinstance(identify_result, UserIdentity): write_to_log( "User identification function did not return a valid UserIdentity " @@ -26,17 +117,41 @@ def identify_session(server, request: any, context: any) -> UserIdentity | None: ) return None - event = UnredactedEvent( - session_id=data.session_id, - timestamp=datetime.now(timezone.utc), - event_type=EventType.AGENTCAT_IDENTIFY.value, - identify_actor_given_id=identify_result.user_id, - identify_actor_name=identify_result.user_name, - identify_data=identify_result.user_data or {}, + session_id = data.session_id + + # Check the global cache (works across server instance restarts) + previous_identity = _global_identity_cache.get(session_id) + + # Merge identities (overwrite user_id/user_name, merge user_data) + merged_identity = merge_identities(previous_identity, identify_result) + + # Only publish if the identity has changed + has_changed = previous_identity is None or not are_identities_equal( + previous_identity, merged_identity ) - event_queue.publish_event(server, event) - return identify_result + _global_identity_cache.set(session_id, merged_identity) + + if has_changed: + write_to_log( + f"Identified session {session_id} (actor: {merged_identity.user_id})" + ) + event = UnredactedEvent( + session_id=session_id, + timestamp=datetime.now(timezone.utc), + event_type=EventType.AGENTCAT_IDENTIFY.value, + identify_actor_given_id=merged_identity.user_id, + identify_actor_name=merged_identity.user_name, + identify_data=merged_identity.user_data or {}, + ) + event_queue.publish_event(server, event) + + return merged_identity except Exception as e: - write_to_log(f"Error occurred during user identification: {e}") + # safe_error_string: the exception comes from customer code and may + # have a broken __str__ — the log line itself must never raise, since + # several call sites run unguarded on the customer's request path. + write_to_log( + f"Error occurred during user identification: {safe_error_string(e)}" + ) return None diff --git a/src/agentcat/modules/internal.py b/src/agentcat/modules/internal.py index 9884ecf..2d2edd3 100644 --- a/src/agentcat/modules/internal.py +++ b/src/agentcat/modules/internal.py @@ -7,7 +7,7 @@ from ..types import EventType, AgentCatData, ToolRegistration, UnredactedEvent from .compatibility import is_official_fastmcp_server -from .logging import write_to_log +from .logging import safe_error_string, write_to_log from .validation import validate_tags # WeakKeyDictionary to store data associated with server instances @@ -51,6 +51,10 @@ def reset_all_tracking_data() -> None: """Reset all server tracking data (mainly for testing).""" _server_data_map.clear() _original_methods.clear() + # Lazy import to avoid a module-level import cycle (identify → internal). + from .identify import reset_identity_cache + + reset_identity_cache() write_to_log("Reset all server tracking data") @@ -136,15 +140,16 @@ async def resolve_event_tags( result = callback(request, extra) if inspect.iscoroutine(result): result = await result + # Truthiness/validation stay inside the guard: the callback may return + # arbitrary objects (broken __bool__, non-dict) and this runs on the + # customer's request path. + if not result: + return None + return validate_tags(result) except Exception as e: - write_to_log(f"event_tags callback error: {e}") - return None - - if not result: + write_to_log(f"event_tags callback error: {safe_error_string(e)}") return None - return validate_tags(result) - async def resolve_event_properties( data: AgentCatData, request: Any, extra: Any @@ -162,12 +167,13 @@ async def resolve_event_properties( result = callback(request, extra) if inspect.iscoroutine(result): result = await result + # Truthiness stays inside the guard — arbitrary return values may have + # a broken __bool__ and this runs on the customer's request path. + return result or None except Exception as e: - write_to_log(f"event_properties callback error: {e}") + write_to_log(f"event_properties callback error: {safe_error_string(e)}") return None - return result or None - async def attach_event_metadata( event: UnredactedEvent, @@ -178,18 +184,22 @@ async def attach_event_metadata( """Attach customer-defined tags and properties to an event before publish. Safe no-op if data is None or callbacks are unset. Failures in either - callback are logged and swallowed — event is still published. + callback are logged and swallowed — event is still published. Must never + raise: call sites run unguarded on the customer's request path. """ - if data is None: - return + try: + if data is None: + return - tags = await resolve_event_tags(data, request, extra) - if tags: - event.tags = tags + tags = await resolve_event_tags(data, request, extra) + if tags: + event.tags = tags - properties = await resolve_event_properties(data, request, extra) - if properties: - event.properties = properties + properties = await resolve_event_properties(data, request, extra) + if properties: + event.properties = properties + except Exception as e: + write_to_log(f"attach_event_metadata error: {safe_error_string(e)}") def get_tool_timeline(server: Any) -> List[Dict[str, Any]]: diff --git a/src/agentcat/modules/logging.py b/src/agentcat/modules/logging.py index 761dac6..11a9f3f 100644 --- a/src/agentcat/modules/logging.py +++ b/src/agentcat/modules/logging.py @@ -33,6 +33,23 @@ def set_diagnostics_sink(fn: Callable[[str], None] | None) -> None: _diagnostics_sink = fn +def safe_error_string(error: object) -> str: + """Best-effort str(error) that never raises. + + Exceptions surfaced from customer code (identify hooks, tools, callbacks) + may have a broken __str__/__format__. Log formatting on the request path + must never raise, or the log line itself would break the customer's + request handling. + """ + try: + return str(error) + except Exception: + try: + return f"" + except Exception: + return "" + + def write_to_log(message: str) -> None: timestamp = datetime.now(timezone.utc).isoformat() log_entry = f"[{timestamp}] {message}" diff --git a/src/agentcat/modules/overrides/community/monkey_patch.py b/src/agentcat/modules/overrides/community/monkey_patch.py index c7f8106..1c87987 100644 --- a/src/agentcat/modules/overrides/community/monkey_patch.py +++ b/src/agentcat/modules/overrides/community/monkey_patch.py @@ -15,7 +15,7 @@ from agentcat.modules.exceptions import capture_exception from agentcat.modules.identify import identify_session from agentcat.modules.internal import attach_event_metadata, get_server_tracking_data -from agentcat.modules.logging import write_to_log +from agentcat.modules.logging import safe_error_string, write_to_log from agentcat.modules.request_extra import params_with_extra from agentcat.modules.session import ( get_client_info_from_request_context, @@ -100,7 +100,9 @@ async def wrapped_call_tool_handler(request: CallToolRequest) -> ServerResult: # Handle session identification try: client_name, client_version = get_client_info_from_request_context(lowlevel_server, request_context) - identity = identify_session(lowlevel_server, request, request_context) + identity = await identify_session( + lowlevel_server, request, request_context + ) except Exception as e: client_name, client_version = None, None identity = None @@ -166,7 +168,7 @@ async def wrapped_call_tool_handler(request: CallToolRequest) -> ServerResult: return result except Exception as e: - write_to_log(f"Error in wrapped_call_tool_handler: {e}") + write_to_log(f"Error in wrapped_call_tool_handler: {safe_error_string(e)}") event.is_error = True # Use full exception capture with stack trace try: @@ -175,7 +177,7 @@ async def wrapped_call_tool_handler(request: CallToolRequest) -> ServerResult: # Fallback to simple error if capture fails write_to_log(f"Error capturing exception: {capture_err}") event.error = { - "message": str(e), + "message": safe_error_string(e), "type": type(e).__name__, "platform": "python", } diff --git a/src/agentcat/modules/overrides/community_v3/middleware.py b/src/agentcat/modules/overrides/community_v3/middleware.py index b597be4..78604b5 100644 --- a/src/agentcat/modules/overrides/community_v3/middleware.py +++ b/src/agentcat/modules/overrides/community_v3/middleware.py @@ -22,7 +22,7 @@ ) from agentcat.modules.identify import identify_session from agentcat.modules.internal import attach_event_metadata, mark_tool_tracked, register_tool -from agentcat.modules.logging import write_to_log +from agentcat.modules.logging import safe_error_string, write_to_log from agentcat.modules.request_extra import params_with_extra from agentcat.modules.session import ( get_client_info_from_request_context, @@ -112,7 +112,9 @@ async def on_initialize( client_name, client_version = get_client_info_from_request_context(self.server, request_context) else: get_client_info_from_request_context(self.server, request_context) - identity = identify_session(self.server, context.message, request_context) + identity = await identify_session( + self.server, context.message, request_context + ) except Exception as e: identity = None write_to_log(f"Non-critical error in session handling: {e}") @@ -170,7 +172,9 @@ async def on_call_tool( request_context = self._get_request_context(context) try: client_name, client_version = get_client_info_from_request_context(self.server, request_context) - identity = identify_session(self.server, context.message, request_context) + identity = await identify_session( + self.server, context.message, request_context + ) except Exception as e: client_name, client_version = None, None identity = None @@ -238,7 +242,7 @@ async def on_call_tool( return result except Exception as e: - write_to_log(f"Error in on_call_tool: {e}") + write_to_log(f"Error in on_call_tool: {safe_error_string(e)}") event.is_error = True store_captured_error(e) event.error = capture_exception(e) @@ -269,7 +273,9 @@ async def on_list_tools( request_context = self._get_request_context(context) try: client_name, client_version = get_client_info_from_request_context(self.server, request_context) - identity = identify_session(self.server, context.message, request_context) + identity = await identify_session( + self.server, context.message, request_context + ) except Exception as e: client_name, client_version = None, None identity = None diff --git a/src/agentcat/modules/overrides/mcp_server.py b/src/agentcat/modules/overrides/mcp_server.py index 12e632c..d516b2f 100644 --- a/src/agentcat/modules/overrides/mcp_server.py +++ b/src/agentcat/modules/overrides/mcp_server.py @@ -8,6 +8,7 @@ from agentcat.modules import event_queue from agentcat.modules.compatibility import is_mcp_error_response +from agentcat.modules.exceptions import capture_exception from agentcat.modules.identify import identify_session from agentcat.modules.internal import attach_event_metadata from agentcat.modules.logging import write_to_log @@ -43,7 +44,7 @@ async def wrapped_initialize_handler(request: InitializeRequest) -> ServerResult """Intercept initialize requests to add AgentCat data to the request context.""" session_id = get_server_session_id(server) request_context = safe_request_context(server) - identity = identify_session(server, request, request_context) + identity = await identify_session(server, request, request_context) # Extract clientInfo from InitializeRequest params (MCP protocol provides it here) client_name, client_version = None, None @@ -82,7 +83,7 @@ async def wrapped_list_tools_handler(request: ListToolsRequest) -> ServerResult: session_id = get_server_session_id(server) request_context = safe_request_context(server) client_name, client_version = get_client_info_from_request_context(server, request_context) - identity = identify_session(server, request, request_context) + identity = await identify_session(server, request, request_context) event = UnredactedEvent( session_id=session_id, @@ -168,7 +169,7 @@ async def wrapped_call_tool_handler(request: CallToolRequest) -> ServerResult: session_id = get_server_session_id(server) request_context = safe_request_context(server) client_name, client_version = get_client_info_from_request_context(server, request_context) - identity = identify_session(server, request, request_context) + identity = await identify_session(server, request, request_context) write_to_log( f"Intercepted call to tool '{tool_name}' with arguments: {arguments} and request context: {request_context}" @@ -218,18 +219,25 @@ async def wrapped_call_tool_handler(request: CallToolRequest) -> ServerResult: try: # Call the original handler result = await original_call_tool_handler(request) - is_error, error_message = is_mcp_error_response(result) + is_error, _error_message = is_mcp_error_response(result) event.is_error = is_error - event.error = {"message": error_message} if is_error else None + # Full structured capture (message/type/platform; the MCP SDK + # already converted the exception, so no frames survive here) + event.error = ( + capture_exception(getattr(result, "root", result)) + if is_error + else None + ) # Record the trace using existing infrastructure event.response = result.model_dump() if result else None event_queue.publish_event(server, event) return result except Exception as e: - # Record the error trace + # Record the error trace with full structured capture + # (frames, chained errors, platform) event.is_error = True - event.error = {"message": str(e)} + event.error = capture_exception(e) event_queue.publish_event(server, event) raise else: @@ -256,7 +264,7 @@ async def wrapped_initialize_handler(request: InitializeRequest) -> ServerResult session_id = get_server_session_id(server) request_context = safe_request_context(server) try: - identity = identify_session(server, request, request_context) + identity = await identify_session(server, request, request_context) except Exception as e: identity = None write_to_log(f"Ran into an error in session identification, no identity could be determined: {e}") @@ -297,7 +305,7 @@ async def wrapped_list_tools_handler(request: ListToolsRequest) -> ServerResult: session_id = get_server_session_id(server) request_context = safe_request_context(server) client_name, client_version = get_client_info_from_request_context(server, request_context) - identity = identify_session(server, request, request_context) + identity = await identify_session(server, request, request_context) event = UnredactedEvent( session_id=session_id, diff --git a/src/agentcat/modules/overrides/official/monkey_patch.py b/src/agentcat/modules/overrides/official/monkey_patch.py index 28f3d6b..2274fd5 100644 --- a/src/agentcat/modules/overrides/official/monkey_patch.py +++ b/src/agentcat/modules/overrides/official/monkey_patch.py @@ -28,7 +28,7 @@ register_tool, store_original_method, ) -from agentcat.modules.logging import write_to_log +from agentcat.modules.logging import safe_error_string, write_to_log from agentcat.modules.request_extra import params_with_extra from agentcat.modules.session import ( get_client_info_from_request_context, @@ -263,7 +263,7 @@ async def patched_call_tool( }, )() - identity = identify_session(server._mcp_server, mock_request, request_context) + identity = await identify_session(server._mcp_server, mock_request, request_context) except Exception as e: client_name, client_version = None, None identity = None @@ -389,7 +389,7 @@ async def patched_call_tool( except Exception as e: # Log the error - write_to_log(f"Error in patched_call_tool: {e}") + write_to_log(f"Error in patched_call_tool: {safe_error_string(e)}") # Try to mark event as error if it exists if event: @@ -402,7 +402,7 @@ async def patched_call_tool( write_to_log(f"Error capturing exception: {capture_err}") try: event.error = { - "message": str(e), + "message": safe_error_string(e), "type": type(e).__name__, "platform": "python", } diff --git a/src/agentcat/modules/redaction.py b/src/agentcat/modules/redaction.py index 58b0dca..9c23f66 100644 --- a/src/agentcat/modules/redaction.py +++ b/src/agentcat/modules/redaction.py @@ -1,7 +1,12 @@ """PII redaction for AgentCat logs.""" +import asyncio +import inspect +import threading from typing import Any, TYPE_CHECKING, Callable, Set +from pydantic import BaseModel + if TYPE_CHECKING: from agentcat.types import Event, UnredactedEvent @@ -25,6 +30,45 @@ } +def _resolve_redacted_value(result: Any) -> Any: + """Resolve a redaction result that may be an awaitable. + + Redaction functions may be sync or async (see `RedactionFunction`). The + event queue processes events on worker threads with no running event loop, + so coroutine results are resolved with `asyncio.run`. If a loop is already + running in this thread (defensive; not the normal publish path), the + coroutine is resolved on a short-lived helper thread instead, since + `asyncio.run` cannot be called from inside a running loop. + """ + if not inspect.isawaitable(result): + return result + + async def _await() -> Any: + return await result + + try: + asyncio.get_running_loop() + except RuntimeError: + # Normal path: event-queue worker threads have no event loop. + return asyncio.run(_await()) + + # A loop is running in this thread — resolve on a fresh thread. + outcome: dict[str, Any] = {} + + def _runner() -> None: + try: + outcome["value"] = asyncio.run(_await()) + except BaseException as exc: # noqa: BLE001 — propagated to caller below + outcome["error"] = exc + + thread = threading.Thread(target=_runner, daemon=True) + thread.start() + thread.join() + if "error" in outcome: + raise outcome["error"] + return outcome["value"] + + def redact_strings_in_object( obj: Any, redact_fn: Callable[[str], str], @@ -38,7 +82,7 @@ def redact_strings_in_object( Args: obj: The object to redact strings from - redact_fn: The redaction function to apply to each string + redact_fn: The redaction function to apply to each string (sync or async) path: The current path in the object tree (used to check protected fields) is_protected: Whether the current object/value is within a protected field @@ -53,7 +97,7 @@ def redact_strings_in_object( # Don't redact if this field or any parent field is protected if is_protected: return obj - return redact_fn(obj) + return _resolve_redacted_value(redact_fn(obj)) # Handle arrays/lists if isinstance(obj, list): @@ -83,7 +127,7 @@ def redact_strings_in_object( return redacted_obj - # For all other types (numbers, booleans, etc.), return as-is + # For all other types (numbers, booleans, datetimes, etc.), return as-is return obj @@ -93,11 +137,21 @@ def redact_event(event: "UnredactedEvent", redact_fn: Callable[[str], str]) -> " This is the main entry point for redacting sensitive information from events before they are sent to the analytics service. + Accepts either a Pydantic event model (the live publish path — the event is + dumped, redacted recursively, and rebuilt) or a plain dict. + Args: event: The event to redact - redact_fn: The customer's redaction function + redact_fn: The customer's redaction function (sync or async) Returns: A new event object with all strings redacted """ + if isinstance(event, BaseModel): + # Dump the model to a plain dict (dropping the redaction_fn callable), + # redact recursively, then rebuild the same model class. + event_dict = event.model_dump(exclude_none=True, exclude={"redaction_fn"}) + redacted_dict = redact_strings_in_object(event_dict, redact_fn, "", False) + return type(event)(**redacted_dict) + return redact_strings_in_object(event, redact_fn, "", False) diff --git a/src/agentcat/modules/session.py b/src/agentcat/modules/session.py index f6d2603..417008e 100644 --- a/src/agentcat/modules/session.py +++ b/src/agentcat/modules/session.py @@ -1,5 +1,6 @@ """Session management for AgentCat.""" +import hashlib import re import sys from datetime import datetime, timedelta, timezone @@ -11,15 +12,59 @@ from agentcat.modules.internal import get_server_tracking_data, set_server_tracking_data from agentcat.modules.logging import write_to_log +from ..thirdparty.ksuid import Ksuid from ..types import AgentCatData, SessionInfo from ..utils import generate_prefixed_ksuid +# Fixed epoch used when deriving deterministic session IDs (2024-01-01T00:00:00Z, +# in milliseconds). Must match the TypeScript SDK's session.ts for cross-SDK +# determinism. +_DERIVED_SESSION_EPOCH_MS = 1704067200000 +# Maximum timestamp offset added to the epoch (1 year, in milliseconds). +_DERIVED_SESSION_MAX_OFFSET_MS = 365 * 24 * 60 * 60 * 1000 + def new_session_id() -> str: """Generate a new session ID.""" return generate_prefixed_ksuid(SESSION_ID_PREFIX) +def derive_session_id_from_mcp_session( + mcp_session_id: str, project_id: str | None = None +) -> str: + """Create a deterministic KSUID session ID from an MCP session ID. + + The same inputs always produce the same session ID, enabling correlation + across server restarts (mirrors the TypeScript SDK's + deriveSessionIdFromMCPSession). + + Args: + mcp_session_id: The session ID from the MCP protocol + project_id: Optional AgentCat project ID to include in the hash + + Returns: + A KSUID with the "ses" prefix derived deterministically from the inputs + """ + input_str = f"{mcp_session_id}:{project_id}" if project_id else mcp_session_id + + # Hash the input with SHA-256 + digest = hashlib.sha256(input_str.encode("utf-8")).digest() + + # Derive a deterministic but valid timestamp from the first 4 bytes: + # fixed 2024-01-01 epoch plus a hash-based offset capped at 1 year. + timestamp_offset = ( + int.from_bytes(digest[:4], "big") % _DERIVED_SESSION_MAX_OFFSET_MS + ) + timestamp_ms = _DERIVED_SESSION_EPOCH_MS + timestamp_offset + + # Use the next 16 bytes of the hash as the KSUID payload + ksuid = Ksuid( + datetime=datetime.fromtimestamp(timestamp_ms / 1000, tz=timezone.utc), + payload=digest[4:20], + ) + return f"{SESSION_ID_PREFIX}_{ksuid}" + + def get_agentcat_version() -> str | None: """Get the current AgentCat SDK version.""" try: diff --git a/src/agentcat/modules/telemetry.py b/src/agentcat/modules/telemetry.py index 767ed75..4eddcb3 100644 --- a/src/agentcat/modules/telemetry.py +++ b/src/agentcat/modules/telemetry.py @@ -63,6 +63,10 @@ def _create_exporter(self, name: str, config: ExporterConfig) -> Optional[Export from .exporters.sentry import SentryExporter return SentryExporter(config) + elif exporter_type == "posthog": + from .exporters.posthog import PostHogExporter + + return PostHogExporter(config) else: write_to_log(f"Unknown exporter type: {exporter_type}") return None diff --git a/src/agentcat/modules/truncation.py b/src/agentcat/modules/truncation.py index da6799c..6e4fd13 100644 --- a/src/agentcat/modules/truncation.py +++ b/src/agentcat/modules/truncation.py @@ -15,11 +15,20 @@ from .logging import write_to_log MAX_EVENT_BYTES = 102_400 # 100 KB total event size -MAX_STRING_BYTES = 10_240 # 10 KB per individual string -MAX_DEPTH = 5 # Max nesting depth for dicts/lists -MAX_BREADTH = 500 # Max items per dict/list +MAX_STRING_BYTES = 32_768 # 32 KB per individual string +MAX_DEPTH = 10 # Max nesting depth for dicts/lists +MAX_BREADTH = 100 # Max items per dict/list MIN_DEPTH = 1 # Never reduce depth below this to avoid type mismatches +# Field-level limits (mirrors the TypeScript SDK's field-cap layer). +# Applied unconditionally to every event, before the size-budget pass. +MAX_USER_INTENT_LENGTH = 2_048 # user_intent +MAX_ERROR_MESSAGE_LENGTH = 2_048 # error.message +MAX_RESOURCE_NAME_LENGTH = 256 # resource_name +MAX_METADATA_LENGTH = 256 # server_name/server_version/client_name/client_version +MAX_STACK_FRAMES = 50 # error.frames — keep first 25 + last 25 when over +MAX_CONTENT_TEXT_LENGTH = 32_768 # response.content[].text per text block + # Only these fields get truncated; all other top-level event fields pass through untouched. TRUNCATABLE_FIELDS = {"parameters", "response", "error", "identify_data", "user_intent", "additional_properties"} @@ -116,10 +125,116 @@ def _truncate_value( _seen.discard(obj_id) +# --- Field-level cap layer (mirrors TypeScript truncateEvent layer 1) --- + + +def _truncate_stack_frames(frames: Any) -> Any: + """Trim a stack frame list to MAX_STACK_FRAMES, keeping first + last halves.""" + if not isinstance(frames, list) or len(frames) <= MAX_STACK_FRAMES: + return frames + half = MAX_STACK_FRAMES // 2 + return frames[:half] + frames[-half:] + + +def _truncate_response_content(response: Any) -> Any: + """Cap each text content block in a response to MAX_CONTENT_TEXT_LENGTH bytes. + + Returns the original object unchanged (same identity) when no block + needed truncation. + """ + if not isinstance(response, dict): + return response + content = response.get("content") + if not isinstance(content, list): + return response + + changed = False + new_content = [] + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "text" + and isinstance(block.get("text"), str) + ): + capped = _truncate_string(block["text"], max_bytes=MAX_CONTENT_TEXT_LENGTH) + if capped != block["text"]: + block = {**block, "text": capped} + changed = True + new_content.append(block) + + if not changed: + return response + return {**response, "content": new_content} + + +def _apply_field_caps(event: "UnredactedEvent") -> "UnredactedEvent": + """Apply per-field limits to an event, returning a copy only if changed. + + Mirrors the TypeScript SDK's field-level truncation layer, which runs + unconditionally on every event before the 100 KB size-budget pass: + - user_intent -> 2048 + - error.message -> 2048; error.frames -> first 25 + last 25 when > 50 + - resource_name -> 256 + - server_name / server_version / client_name / client_version -> 256 + - response.content[].text -> 32 KB per text block + """ + updates: dict[str, Any] = {} + + for field_name, limit in ( + ("user_intent", MAX_USER_INTENT_LENGTH), + ("resource_name", MAX_RESOURCE_NAME_LENGTH), + ("server_name", MAX_METADATA_LENGTH), + ("server_version", MAX_METADATA_LENGTH), + ("client_name", MAX_METADATA_LENGTH), + ("client_version", MAX_METADATA_LENGTH), + ): + value = getattr(event, field_name, None) + if isinstance(value, str): + capped = _truncate_string(value, max_bytes=limit) + if capped != value: + updates[field_name] = capped + + error = getattr(event, "error", None) + if isinstance(error, dict): + new_error = dict(error) + error_changed = False + + message = new_error.get("message") + if isinstance(message, str): + capped = _truncate_string(message, max_bytes=MAX_ERROR_MESSAGE_LENGTH) + if capped != message: + new_error["message"] = capped + error_changed = True + + frames = new_error.get("frames") + trimmed_frames = _truncate_stack_frames(frames) + if trimmed_frames is not frames: + new_error["frames"] = trimmed_frames + error_changed = True + + if error_changed: + updates["error"] = new_error + + response = getattr(event, "response", None) + capped_response = _truncate_response_content(response) + if capped_response is not response: + updates["response"] = capped_response + + if not updates: + return event + return event.model_copy(update=updates) + + def truncate_event(event: "UnredactedEvent | None") -> "UnredactedEvent | None": - """Return a truncated copy of *event* if it exceeds MAX_EVENT_BYTES. + """Apply layered truncation to *event*. + + Layer 1 — field-level caps (unconditional, mirrors the TypeScript SDK): + user_intent, resource_name, server/client metadata, error.message, + error.frames, and response content text blocks are capped regardless + of overall event size. - Uses size-targeted normalization strategy: normalize with the + Layer 2 — size budget: if the event still exceeds MAX_EVENT_BYTES, + uses a size-targeted normalization strategy: normalize with the current limits, check JSON byte size, and if still over the limit tighten limits and re-normalize until it fits. @@ -136,6 +251,10 @@ def truncate_event(event: "UnredactedEvent | None") -> "UnredactedEvent | None": return None try: + # Layer 1: field-level caps, applied to every event unconditionally + event = _apply_field_caps(event) + + # Layer 2: overall size budget serialized_bytes = event.model_dump_json().encode("utf-8") byte_size = len(serialized_bytes) if byte_size <= MAX_EVENT_BYTES: diff --git a/src/agentcat/modules/validation.py b/src/agentcat/modules/validation.py index eb7af60..e61b0e9 100644 --- a/src/agentcat/modules/validation.py +++ b/src/agentcat/modules/validation.py @@ -15,45 +15,29 @@ def validate_tags(tags: dict) -> Optional[dict]: """Validate and filter a tags dict against AgentCat tag constraints. Invalid entries are dropped with a warning logged via write_to_log. - Returns None if no valid entries remain. + Returns None if no valid entries remain. Never raises — callers include + the customer's request path, and keys/values are arbitrary user objects + (possibly with broken __str__/__format__). """ if not tags: return None + if not isinstance(tags, dict): + write_to_log( + f"Dropping tags — expected a dict, got {type(tags).__name__}" + ) + return None + valid: list[tuple[str, str]] = [] for key, value in tags.items(): - if not isinstance(key, str) or not key or not TAG_KEY_REGEX.match(key): - write_to_log( - f'Dropping invalid tag: "{key}" — key contains invalid characters or is empty' - ) - continue - - if len(key) > MAX_TAG_KEY_LENGTH: - write_to_log( - f'Dropping invalid tag: "{key}" — key exceeds max length of {MAX_TAG_KEY_LENGTH}' - ) - continue - - if not isinstance(value, str): - write_to_log( - f'Dropping invalid tag: "{key}" — non-string value (got {type(value).__name__})' - ) + try: + valid_entry = _validate_tag_entry(key, value) + except Exception: + # e.g. a key/value whose __str__/__format__ raises during logging continue - - if len(value) > MAX_TAG_VALUE_LENGTH: - write_to_log( - f'Dropping invalid tag: "{key}" — value exceeds max length of {MAX_TAG_VALUE_LENGTH}' - ) - continue - - if "\n" in value: - write_to_log( - f'Dropping invalid tag: "{key}" — value contains newline character' - ) - continue - - valid.append((key, value)) + if valid_entry: + valid.append((key, value)) if not valid: return None @@ -66,3 +50,42 @@ def validate_tags(tags: dict) -> Optional[dict]: valid = valid[:MAX_TAG_ENTRIES] return dict(valid) + + +def _validate_tag_entry(key, value) -> bool: + """Validate a single tag entry; returns True if it should be kept. + + May raise if key/value have a broken __str__/__format__ (log formatting); + validate_tags catches that per-entry and drops the tag. + """ + if not isinstance(key, str) or not key or not TAG_KEY_REGEX.match(key): + write_to_log( + f'Dropping invalid tag: "{key}" — key contains invalid characters or is empty' + ) + return False + + if len(key) > MAX_TAG_KEY_LENGTH: + write_to_log( + f'Dropping invalid tag: "{key}" — key exceeds max length of {MAX_TAG_KEY_LENGTH}' + ) + return False + + if not isinstance(value, str): + write_to_log( + f'Dropping invalid tag: "{key}" — non-string value (got {type(value).__name__})' + ) + return False + + if len(value) > MAX_TAG_VALUE_LENGTH: + write_to_log( + f'Dropping invalid tag: "{key}" — value exceeds max length of {MAX_TAG_VALUE_LENGTH}' + ) + return False + + if "\n" in value: + write_to_log( + f'Dropping invalid tag: "{key}" — value contains newline character' + ) + return False + + return True diff --git a/src/agentcat/types.py b/src/agentcat/types.py index 41f15a9..f51ed2b 100644 --- a/src/agentcat/types.py +++ b/src/agentcat/types.py @@ -6,12 +6,19 @@ from enum import Enum from typing import Any, Dict, Optional, Set, TypedDict, Literal, Union, NotRequired from agentcat_api import PublishEventRequest -from pydantic import BaseModel - -from agentcat.modules.constants import DEFAULT_CONTEXT_DESCRIPTION - -# Type alias for identify function -IdentifyFunction = Callable[[dict[str, Any], Any], Optional["UserIdentity"]] +from pydantic import BaseModel, field_validator + +from agentcat.modules.constants import ( + AGENTCAT_CUSTOM_EVENT_TYPE, + DEFAULT_CONTEXT_DESCRIPTION, +) + +# Type alias for identify function. +# Accepts sync or async callables (coroutine results are awaited). +IdentifyFunction = Callable[ + [dict[str, Any], Any], + Optional["UserIdentity"] | Awaitable[Optional["UserIdentity"]], +] # Type alias for redaction function RedactionFunction = Callable[[str], str | Awaitable[str]] # Type alias for event_tags callback — returns str:str map attached to every auto-captured event. @@ -58,6 +65,17 @@ class Event(PublishEventRequest): # at publish time, and telemetry-only mode sends events without one. project_id: Optional[str] = None + @field_validator("event_type") + def event_type_validate_enum( + cls, value: Optional[str] # noqa: N805 — pydantic validator + ) -> Optional[str]: + """Relax the generated enum check to admit SDK-defined event types + (e.g. "agentcat:custom") the generated client doesn't know about.""" + if value == AGENTCAT_CUSTOM_EVENT_TYPE: + return value + validated: Optional[str] = PublishEventRequest.event_type_validate_enum(value) + return validated + # Error tracking types @@ -119,6 +137,23 @@ class UnredactedEvent(Event): redaction_fn: RedactionFunction | None = None +@dataclass +class CustomEventData: + """Optional payload for `agentcat.publish_custom_event`.""" + + resource_name: str | None = None + parameters: dict[str, Any] | None = None + response: dict[str, Any] | None = None + # Human-readable explanation of the event; stored as user_intent. + message: str | None = None + duration: int | None = None # milliseconds + is_error: bool | None = None + error: dict[str, Any] | None = None + # Customer-defined metadata (same semantics as event_tags/event_properties) + tags: dict[str, str] | None = None + properties: dict[str, Any] | None = None + + @dataclass class ToolRegistration: """Metadata about a registered tool.""" @@ -162,8 +197,25 @@ class SentryExporterConfig(TypedDict): enable_tracing: Optional[bool] # Optional, defaults to True +class PostHogExporterConfig(TypedDict): + """Configuration for PostHog exporter.""" + + type: Literal["posthog"] + api_key: str # Required - PostHog project API key (e.g. phc_...) + # Optional, defaults to https://us.i.posthog.com (supports self-hosted & EU region) + host: NotRequired[str] + # Emits $ai_span events for tool calls alongside regular capture events, + # integrating with PostHog's AI observability views. Defaults to False. + enable_ai_tracing: NotRequired[bool] + + # Union type for all exporter configurations -ExporterConfig = Union[OTLPExporterConfig, DatadogExporterConfig, SentryExporterConfig] +ExporterConfig = Union[ + OTLPExporterConfig, + DatadogExporterConfig, + SentryExporterConfig, + PostHogExporterConfig, +] @dataclass diff --git a/tests/e2e/official/test_redaction_http.py b/tests/e2e/official/test_redaction_http.py index 30cb60f..752da16 100644 --- a/tests/e2e/official/test_redaction_http.py +++ b/tests/e2e/official/test_redaction_http.py @@ -1,15 +1,9 @@ """Redaction over real-wire payloads. -KNOWN BUG (xfail-tracked): `agentcat.modules.redaction.redact_event` only -recurses into `dict` and `list` types, not Pydantic `UnredactedEvent` objects. -The event_queue worker invokes `redact_event(event, ...)` where `event` is an -`UnredactedEvent`; the call returns the input unchanged, so customer-supplied -redact functions never actually run on the live event-publish path. - -Tests below are marked xfail so they: -1. Codify the intended behavior. -2. Serve as a regression target — when the bug is fixed, they should be - un-xfailed (the strict=False xfail still passes if the test starts working). +Verifies that customer-supplied redact functions run on the live +event-publish path: `redact_event` handles the Pydantic `UnredactedEvent` +model (dump → redact recursively → rebuild), so every unprotected string +field of a published event passes through the redaction function. """ from __future__ import annotations @@ -32,11 +26,6 @@ def _set_redact(server, fn) -> None: data.options.redact_sensitive_information = fn -@pytest.mark.xfail( - reason="redact_event does not recurse into Pydantic UnredactedEvent; " - "redaction never fires on real events. Track as separate fix.", - strict=False, -) @pytest.mark.asyncio async def test_redact_function_runs_on_real_event_payload( official_http_server, capture_queue @@ -67,11 +56,6 @@ def redact(s: str) -> str: _set_redact(server, None) -@pytest.mark.xfail( - reason="redact_event does not recurse into Pydantic UnredactedEvent; " - "redaction never fires on real events. Track as separate fix.", - strict=False, -) @pytest.mark.asyncio async def test_redaction_can_scrub_authorization_header_in_extra( official_http_server, capture_queue @@ -111,12 +95,6 @@ def redact(s: str) -> str: _set_redact(server, None) -@pytest.mark.xfail( - reason="redact_event does not invoke the user's redact fn on Pydantic " - "events, so redact-fn-raise never fires; the 'drop event on raise' path " - "is unreachable until the redact_event recursion bug is fixed.", - strict=False, -) @pytest.mark.asyncio async def test_redaction_failure_drops_event(official_http_server, capture_queue): url, server = official_http_server diff --git a/tests/test_failure_injection.py b/tests/test_failure_injection.py new file mode 100644 index 0000000..a98a6ca --- /dev/null +++ b/tests/test_failure_injection.py @@ -0,0 +1,633 @@ +"""Adversarial failure-injection tests for the fail-open invariant. + +The SDK embeds in customers' MCP servers, so no analytics code path may ever +raise into the customer's request handling, kill the worker pipeline, or +crash the process. Analytics failures must degrade to dropped events plus a +log line. + +These tests inject hostile customer callbacks (identify/tags/properties/ +redaction functions that raise — sync and async — or return garbage) and +hostile payloads (cyclic structures, NaN, bytes, objects with raising +__repr__/__str__, exceptions with broken __str__) into every layer: + +- the request path (low-level server wrappers, FastMCP over a real in-memory + transport) +- capture_exception on the low-level error path +- publish_custom_event (user-invoked; only ValueError/TypeError by design) +- the worker pipeline (redact -> sanitize -> truncate -> send) +- the PostHog exporter and telemetry manager isolation +""" + +import time +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import pytest +from mcp.server import Server as LowLevelServer +from mcp.types import CallToolRequest, CallToolRequestParams, TextContent + +import agentcat +from agentcat import AgentCatOptions, track +from agentcat.modules import event_queue as event_queue_module +from agentcat.modules.event_queue import EventQueue, set_event_queue +from agentcat.modules.exceptions import capture_exception +from agentcat.modules.identify import reset_identity_cache +from agentcat.modules.internal import ( + reset_all_tracking_data, + set_server_tracking_data, +) +from agentcat.modules.overrides.mcp_server import override_lowlevel_mcp_server +from agentcat.types import ( + AgentCatData, + CustomEventData, + SessionInfo, + UnredactedEvent, + UserIdentity, +) + +from .test_utils.client import create_test_client +from .test_utils.todo_server import create_todo_server + + +# --------------------------------------------------------------------------- +# Hostile objects +# --------------------------------------------------------------------------- + + +class UnprintableError(Exception): + """An exception whose __str__ raises (customer code can do this).""" + + def __str__(self): + raise RuntimeError("__str__ exploded") + + +class Unboolable: + """Truthiness raises (e.g. numpy arrays behave this way).""" + + def __bool__(self): + raise ValueError("truth value is ambiguous") + + +class RaisingEq: + """Equality raises (e.g. numpy arrays in bool context).""" + + def __eq__(self, other): + raise ValueError("cannot compare") + + __hash__ = object.__hash__ + + +class RaisingReprStr: + """Both __repr__ and __str__ raise.""" + + def __repr__(self): + raise RuntimeError("__repr__ exploded") + + def __str__(self): + raise RuntimeError("__str__ exploded") + + +def _cyclic_dict(): + d = {"a": 1} + d["self"] = d + return d + + +# --------------------------------------------------------------------------- +# Low-level server harness (exercises overrides/mcp_server.py wrappers, the +# path with NO call-site guards around identify/tags/publish) +# --------------------------------------------------------------------------- + + +def _make_request(name="echo", arguments=None) -> CallToolRequest: + return CallToolRequest( + method="tools/call", + params=CallToolRequestParams(name=name, arguments=arguments or {}), + ) + + +def _setup_lowlevel_server(raw_call_tool_handler=None, **option_kwargs): + """Build a tracked low-level server. If raw_call_tool_handler is given it + replaces the SDK-generated handler BEFORE the AgentCat override, so the + wrapper's exception branch can be exercised directly.""" + server = LowLevelServer("failure-injection-server") + + @server.list_tools() + async def list_tools(): + return [] + + @server.call_tool() + async def call_tool(name, arguments): + return [TextContent(type="text", text="ok")] + + if raw_call_tool_handler is not None: + server.request_handlers[CallToolRequest] = raw_call_tool_handler + + options = AgentCatOptions( + enable_tracing=True, + enable_tool_call_context=False, + enable_report_missing=False, + **option_kwargs, + ) + data = AgentCatData( + project_id="test_project", + session_id="ses_failure_injection", + session_info=SessionInfo(), + last_activity=datetime.now(timezone.utc), + options=options, + ) + set_server_tracking_data(server, data) + override_lowlevel_mcp_server(server, data) + return server + + +@pytest.fixture(autouse=True) +def _reset_tracking(): + reset_all_tracking_data() + reset_identity_cache() + yield + reset_all_tracking_data() + reset_identity_cache() + + +class TestRequestPathHostileCallbacks: + """Hostile identify/tags/properties callbacks must never break a request + on the unguarded low-level wrapper path.""" + + async def _call(self, server): + handler = server.request_handlers[CallToolRequest] + return await handler(_make_request()) + + @patch("agentcat.modules.identify.event_queue") + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_sync_identify_raises(self, mock_eq, mock_id_eq): + def identify(request, extra): + raise ValueError("identify blew up") + + server = _setup_lowlevel_server(identify=identify) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.identify.event_queue") + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_async_identify_raises(self, mock_eq, mock_id_eq): + async def identify(request, extra): + raise ValueError("async identify blew up") + + server = _setup_lowlevel_server(identify=identify) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.identify.event_queue") + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_identify_raises_unprintable_exception(self, mock_eq, mock_id_eq): + """The identify hook raises an exception whose __str__ raises; even + the SDK's own logging of the failure must not raise.""" + + def identify(request, extra): + raise UnprintableError() + + server = _setup_lowlevel_server(identify=identify) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.identify.event_queue") + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_identify_returns_unboolable_garbage(self, mock_eq, mock_id_eq): + def identify(request, extra): + return Unboolable() + + server = _setup_lowlevel_server(identify=identify) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.identify.event_queue") + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_identify_user_data_with_raising_eq(self, mock_eq, mock_id_eq): + """Change-detection deep-compares arbitrary user_data; values whose + __eq__ raises must not break the second request.""" + + def identify(request, extra): + return UserIdentity( + user_id="alice", user_name="Alice", user_data={"blob": RaisingEq()} + ) + + server = _setup_lowlevel_server(identify=identify) + result1 = await self._call(server) + result2 = await self._call(server) # triggers are_identities_equal + assert not getattr(result1.root, "isError", False) + assert not getattr(result2.root, "isError", False) + + @patch("agentcat.modules.identify.event_queue") + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_identify_user_data_cyclic(self, mock_eq, mock_id_eq): + def identify(request, extra): + return UserIdentity( + user_id="alice", user_name="Alice", user_data=_cyclic_dict() + ) + + server = _setup_lowlevel_server(identify=identify) + result1 = await self._call(server) + result2 = await self._call(server) + assert not getattr(result1.root, "isError", False) + assert not getattr(result2.root, "isError", False) + + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_event_tags_callback_raises(self, mock_eq): + server = _setup_lowlevel_server( + event_tags=lambda req, extra: (_ for _ in ()).throw(RuntimeError("tags")) + ) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_event_tags_returns_non_dict(self, mock_eq): + server = _setup_lowlevel_server(event_tags=lambda req, extra: ["not", "dict"]) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_event_tags_returns_unprintable_keys(self, mock_eq): + server = _setup_lowlevel_server( + event_tags=lambda req, extra: {RaisingReprStr(): "value", "ok": "kept"} + ) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_event_properties_returns_unboolable(self, mock_eq): + server = _setup_lowlevel_server(event_properties=lambda req, extra: Unboolable()) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_async_event_tags_raises(self, mock_eq): + async def tags(req, extra): + raise RuntimeError("async tags") + + server = _setup_lowlevel_server(event_tags=tags) + result = await self._call(server) + assert not getattr(result.root, "isError", False) + + async def test_publish_pipeline_failure_does_not_break_request(self): + """If assembling/enqueueing the event blows up inside publish_event, + the customer's successful tool result must still be returned.""" + server = _setup_lowlevel_server() + with patch( + "agentcat.modules.event_queue.get_session_info", + side_effect=RuntimeError("session info exploded"), + ): + handler = server.request_handlers[CallToolRequest] + result = await handler(_make_request()) + assert not getattr(result.root, "isError", False) + + +class TestLowLevelErrorPathHostileExceptions: + """capture_exception runs unguarded on the request path; exotic + exceptions must not replace the customer's error or lose the event.""" + + @patch("agentcat.modules.overrides.mcp_server.event_queue") + async def test_unprintable_tool_exception_reraised_unchanged(self, mock_eq): + async def raw_handler(request): + raise UnprintableError() + + server = _setup_lowlevel_server(raw_call_tool_handler=raw_handler) + handler = server.request_handlers[CallToolRequest] + + with pytest.raises(UnprintableError): + await handler(_make_request()) + + # The event must still be recorded, with a fallback error message + events = [c.args[1] for c in mock_eq.publish_event.call_args_list] + call_events = [e for e in events if e.event_type == "mcp:tools/call"] + assert call_events, "tools/call error event not published" + event = call_events[0] + assert event.is_error is True + assert isinstance(event.error, dict) + assert event.error.get("type") == "UnprintableError" + assert isinstance(event.error.get("message"), str) + + def test_capture_exception_unprintable_str(self): + error_data = capture_exception(UnprintableError()) + assert error_data["type"] == "UnprintableError" + assert isinstance(error_data["message"], str) + assert error_data["platform"] == "python" + + def test_capture_exception_unprintable_cause_chain(self): + e = ValueError("outer") + e.__cause__ = UnprintableError() + e.__suppress_context__ = True + error_data = capture_exception(e) + assert error_data["message"] == "outer" + chained = error_data.get("chained_errors") or [] + assert chained and chained[0]["type"] == "UnprintableError" + assert isinstance(chained[0]["message"], str) + + def test_capture_exception_cyclic_cause_chain(self): + a = ValueError("a") + b = ValueError("b") + a.__cause__ = b + a.__suppress_context__ = True + b.__cause__ = a + b.__suppress_context__ = True + error_data = capture_exception(a) # must terminate, not loop/raise + assert error_data["message"] == "a" + + def test_capture_exception_non_exception_unprintable_object(self): + error_data = capture_exception(RaisingReprStr()) + assert isinstance(error_data["message"], str) + assert error_data["platform"] == "python" + + def test_capture_exception_base_exception_subclass(self): + error_data = capture_exception(KeyboardInterrupt("interrupted")) + assert error_data["type"] == "KeyboardInterrupt" + + +class TestFastMCPTransportHostileCallbacks: + """End-to-end over a real in-memory transport: hostile identify must not + break initialize or tool calls on the FastMCP (official) path.""" + + async def test_unprintable_identify_over_transport(self): + mock_api_client = MagicMock() + test_queue = EventQueue(api_client=mock_api_client) + original_queue = event_queue_module.event_queue + set_event_queue(test_queue) + try: + server = create_todo_server() + + def identify(request, extra): + raise UnprintableError() + + track( + server, + "test_project", + AgentCatOptions(identify=identify, enable_report_missing=False), + ) + + async with create_test_client(server) as client: + result = await client.call_tool( + "add_todo", {"text": "still works", "context": "testing"} + ) + assert not result.isError + finally: + set_event_queue(original_queue) + + +class TestPublishCustomEventContract: + """publish_custom_event may raise ValueError/TypeError on bad args (TS + parity) but nothing else; hostile payloads are deferred to the worker.""" + + def _tracked_mock_queue(self): + return patch.object(event_queue_module, "event_queue", MagicMock()) + + def test_weird_session_id_strings(self): + weird = ["über-session🚀", "a" * 10_000, "line\nbreak\ttab", " ", "0"] + with self._tracked_mock_queue() as mock_queue: + for session in weird: + agentcat.publish_custom_event(session, "proj_test") + assert mock_queue.add.call_count == len(weird) + + def test_non_serializable_parameters_enqueued_not_raised(self): + data = CustomEventData( + resource_name="custom-action", + parameters={ + "obj": object(), + "nan": float("nan"), + "inf": float("inf"), + "bytes": b"\xff\xfe\x00", + "dt": datetime.now(timezone.utc), + "raising_repr": RaisingReprStr(), + }, + ) + with self._tracked_mock_queue() as mock_queue: + agentcat.publish_custom_event("session-x", "proj_test", data) + assert mock_queue.add.call_count == 1 + + def test_cyclic_parameters_enqueued_not_raised(self): + data = CustomEventData(parameters=_cyclic_dict()) + with self._tracked_mock_queue() as mock_queue: + agentcat.publish_custom_event("session-y", "proj_test", data) + assert mock_queue.add.call_count == 1 + + def test_non_dict_tags_dropped_not_raised(self): + data = CustomEventData(tags="not-a-dict") + with self._tracked_mock_queue() as mock_queue: + agentcat.publish_custom_event("session-z", "proj_test", data) + event = mock_queue.add.call_args.args[0] + assert event.tags is None + + def test_bad_args_raise_by_design(self): + with pytest.raises(ValueError): + agentcat.publish_custom_event("session", "") + with pytest.raises(TypeError): + agentcat.publish_custom_event(42, "proj_test") + with pytest.raises(ValueError): + # An untracked server object + agentcat.publish_custom_event(object(), "proj_test") + + def test_custom_event_type_accepted_by_model(self): + event = UnredactedEvent(event_type="agentcat:custom") + assert event.event_type == "agentcat:custom" + + +# --------------------------------------------------------------------------- +# Worker pipeline: one poisonous event must never kill the worker +# --------------------------------------------------------------------------- + + +def _make_worker_queue(): + captured = [] + mock_api = MagicMock() + mock_api.publish_event = MagicMock( + side_effect=lambda publish_event_request: captured.append( + publish_event_request + ) + ) + return EventQueue(api_client=mock_api), captured + + +def _wait_for(condition, timeout=5.0): + deadline = time.time() + timeout + while time.time() < deadline: + if condition(): + return True + time.sleep(0.05) + return False + + +def _event(resource_name, parameters=None, redaction_fn=None): + return UnredactedEvent( + session_id="ses_worker_test", + project_id="proj_worker", + event_type="mcp:tools/call", + resource_name=resource_name, + timestamp=datetime.now(timezone.utc), + parameters=parameters, + redaction_fn=redaction_fn, + ) + + +class TestWorkerPipelineSurvival: + def _run(self, events, expect_delivered, timeout=8.0): + queue, captured = _make_worker_queue() + try: + for event in events: + queue.add(event) + assert _wait_for( + lambda: { + e.resource_name for e in captured + } >= set(expect_delivered), + timeout=timeout, + ), ( + f"expected {expect_delivered} delivered; got " + f"{[e.resource_name for e in captured]}" + ) + return captured + finally: + queue._shutdown = True + queue._shutdown_event.set() + queue.executor.shutdown(wait=False, cancel_futures=True) + + def test_poisonous_payload_then_healthy(self): + poison = _event( + "poison", + parameters={ + "cycle": _cyclic_dict(), + "nan": float("nan"), + "bytes": b"\xff\xfe", + "raising_repr": RaisingReprStr(), + }, + ) + healthy = _event("healthy") + self._run([poison, healthy], expect_delivered={"healthy"}) + + def test_sync_redaction_raises_drops_event_worker_survives(self): + def bad_redact(text): + raise RuntimeError("redaction blew up") + + poisoned = _event("poisoned", {"secret": "x"}, redaction_fn=bad_redact) + healthy = _event("healthy", {"secret": "x"}, redaction_fn=lambda s: "REDACTED") + captured = self._run([poisoned, healthy], expect_delivered={"healthy"}) + assert all(e.resource_name != "poisoned" for e in captured) + + def test_unprintable_redaction_error_drops_event_worker_survives(self): + def bad_redact(text): + raise UnprintableError() + + poisoned = _event("poisoned", {"secret": "x"}, redaction_fn=bad_redact) + healthy = _event("healthy") + captured = self._run([poisoned, healthy], expect_delivered={"healthy"}) + assert all(e.resource_name != "poisoned" for e in captured) + + def test_async_redaction_fn_resolved_on_worker(self): + async def redact(text): + return "ASYNC_REDACTED" + + event = _event("async-redacted", {"secret": "topsecret"}, redaction_fn=redact) + captured = self._run([event], expect_delivered={"async-redacted"}) + delivered = [e for e in captured if e.resource_name == "async-redacted"][0] + assert delivered.parameters == {"secret": "ASYNC_REDACTED"} + + def test_async_redaction_raises_drops_event_worker_survives(self): + async def bad_redact(text): + raise RuntimeError("async redaction blew up") + + poisoned = _event("poisoned", {"secret": "x"}, redaction_fn=bad_redact) + healthy = _event("healthy") + captured = self._run([poisoned, healthy], expect_delivered={"healthy"}) + assert all(e.resource_name != "poisoned" for e in captured) + + def test_redaction_returning_non_string_never_crashes_worker(self): + poisoned = _event("poisoned", {"secret": "x"}, redaction_fn=lambda s: object()) + healthy = _event("healthy") + self._run([poisoned, healthy], expect_delivered={"healthy"}) + + +class TestRedactionLoopFallback: + """redact_event resolves async redaction fns with asyncio.run on worker + threads; when a loop is already running in the calling thread it must + fall back to a helper thread and complete without deadlocking.""" + + async def test_async_redaction_inside_running_loop_no_deadlock(self): + from agentcat.modules.redaction import redact_event + + async def redact(text): + return "R" + + event = _event("loop-fallback", {"secret": "y"}) + # Called from within a running event loop (this async test) — takes + # the helper-thread fallback path and must return promptly. + result = redact_event(event, redact) + assert result.parameters == {"secret": "R"} + + async def test_async_redaction_raises_inside_running_loop(self): + from agentcat.modules.redaction import redact_event + + async def redact(text): + raise RuntimeError("boom") + + event = _event("loop-fallback-err", {"secret": "y"}) + with pytest.raises(RuntimeError): + # Propagating is correct here: the worker's caller catches it and + # drops the event. The requirement is no deadlock/hang. + redact_event(event, redact) + + +# --------------------------------------------------------------------------- +# PostHog exporter + telemetry manager isolation +# --------------------------------------------------------------------------- + + +class TestPostHogExporterRobustness: + def _exporter(self): + from agentcat.modules.exporters.posthog import PostHogExporter + + return PostHogExporter({"type": "posthog", "api_key": "phc_test"}) + + def test_network_error_does_not_raise(self): + exporter = self._exporter() + exporter.session = MagicMock() + exporter.session.post.side_effect = ConnectionError("network down") + exporter.export(_event("net-fail")) # must not raise + + def test_serialization_error_does_not_raise(self): + exporter = self._exporter() + exporter.session = MagicMock() + exporter.session.post.side_effect = TypeError("not JSON serializable") + event = _event("bad-json", parameters={"obj": object()}) + event.id = "evt_test" + exporter.export(event) # must not raise + + def test_to_uuidv7_garbage_inputs(self): + import re + + from agentcat.modules.exporters.posthog import to_uuidv7 + + uuid_re = re.compile( + r"^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$" + ) + for garbage in ["", "no-prefix", "ses_$$$$", None, "ses_" + "x" * 500]: + assert uuid_re.match(to_uuidv7(garbage)), f"bad uuid for {garbage!r}" + + def test_error_event_with_hostile_error_dict(self): + exporter = self._exporter() + exporter.session = MagicMock() + event = _event("err-tool") + event.id = "evt_err" + event.is_error = True + event.error = {"message": "boom", "weird": RaisingReprStr()} + exporter.export(event) # must not raise + + def test_telemetry_manager_isolates_exporter_failures(self): + from agentcat.modules.telemetry import TelemetryManager + + manager = TelemetryManager({}) + failing = MagicMock() + failing.export.side_effect = RuntimeError("exporter down") + working = MagicMock() + manager.exporters = {"bad": failing, "good": working} + + event = _event("isolated") + event.id = "evt_iso" + manager.export(event) # must not raise + assert working.export.call_count == 1 diff --git a/tests/test_identify.py b/tests/test_identify.py new file mode 100644 index 0000000..a1e9ab5 --- /dev/null +++ b/tests/test_identify.py @@ -0,0 +1,269 @@ +"""Tests for identify: async hooks, change detection, and identity merging. + +Mirrors the TypeScript SDK's handleIdentify (src/modules/internal.ts): the +identify hook runs per-request but an `agentcat:identify` event is published +only when the (merged) identity actually changed for the session. +""" + +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +from agentcat.modules.identify import ( + IdentityCache, + are_identities_equal, + identify_session, + merge_identities, + reset_identity_cache, +) +from agentcat.modules.internal import ( + reset_all_tracking_data, + set_server_tracking_data, +) +from agentcat.types import AgentCatData, AgentCatOptions, SessionInfo, UserIdentity + +from .test_utils.todo_server import create_todo_server + + +def _identify_events(mock_event_queue): + """Extract events passed to publish_event from the mocked event_queue.""" + return [ + call.args[1] + for call in mock_event_queue.publish_event.call_args_list + if call.args[1].event_type == "agentcat:identify" + ] + + +class TestIdentifySession: + """Change-detection and merge behavior of identify_session.""" + + def setup_method(self): + reset_all_tracking_data() + reset_identity_cache() + self.server = create_todo_server() + + def teardown_method(self): + reset_all_tracking_data() + + def _setup_data(self, identify, session_id="ses_identify_test"): + options = AgentCatOptions(identify=identify) + data = AgentCatData( + project_id="test_project", + session_id=session_id, + session_info=SessionInfo(), + last_activity=datetime.now(timezone.utc), + options=options, + ) + set_server_tracking_data(self.server, data) + return data + + @patch("agentcat.modules.identify.event_queue") + async def test_identical_identity_publishes_single_event(self, mock_event_queue): + """Repeated identical identities → identify hook runs each time but + only ONE agentcat:identify event is published.""" + identify_fn = MagicMock( + side_effect=lambda req, ctx: UserIdentity( + user_id="alice", user_name="Alice", user_data={"plan": "pro"} + ) + ) + self._setup_data(identify_fn) + + for _ in range(3): + result = await identify_session(self.server, MagicMock(), MagicMock()) + assert result is not None + assert result.user_id == "alice" + + assert identify_fn.call_count == 3 + assert len(_identify_events(mock_event_queue)) == 1 + + @patch("agentcat.modules.identify.event_queue") + async def test_changed_identity_publishes_new_event(self, mock_event_queue): + """A changed identity → a new agentcat:identify event.""" + identities = iter( + [ + UserIdentity(user_id="alice", user_name="Alice", user_data=None), + UserIdentity(user_id="alice", user_name="Alice", user_data=None), + UserIdentity(user_id="bob", user_name="Bob", user_data=None), + ] + ) + self._setup_data(lambda req, ctx: next(identities)) + + await identify_session(self.server, MagicMock(), MagicMock()) + await identify_session(self.server, MagicMock(), MagicMock()) + result = await identify_session(self.server, MagicMock(), MagicMock()) + + events = _identify_events(mock_event_queue) + assert len(events) == 2 + assert events[0].identify_actor_given_id == "alice" + assert events[1].identify_actor_given_id == "bob" + assert result.user_id == "bob" + + @patch("agentcat.modules.identify.event_queue") + async def test_user_data_merges_across_calls(self, mock_event_queue): + """user_id/user_name are overwritten; user_data keys are merged with + the newest values winning on collision.""" + identities = iter( + [ + UserIdentity( + user_id="alice", + user_name="Alice", + user_data={"plan": "free", "region": "us"}, + ), + UserIdentity( + user_id="alice", + user_name="Alice A.", + user_data={"plan": "pro"}, + ), + ] + ) + self._setup_data(lambda req, ctx: next(identities)) + + await identify_session(self.server, MagicMock(), MagicMock()) + result = await identify_session(self.server, MagicMock(), MagicMock()) + + assert result.user_id == "alice" + assert result.user_name == "Alice A." + assert result.user_data == {"plan": "pro", "region": "us"} + + events = _identify_events(mock_event_queue) + assert len(events) == 2 + assert events[1].identify_actor_name == "Alice A." + assert events[1].identify_data == {"plan": "pro", "region": "us"} + + @patch("agentcat.modules.identify.event_queue") + async def test_merged_but_unchanged_identity_does_not_republish( + self, mock_event_queue + ): + """A subset of previously-seen user_data merges into an identical + identity → no new event.""" + identities = iter( + [ + UserIdentity( + user_id="alice", user_name="Alice", user_data={"a": "1", "b": "2"} + ), + UserIdentity(user_id="alice", user_name="Alice", user_data={"a": "1"}), + ] + ) + self._setup_data(lambda req, ctx: next(identities)) + + await identify_session(self.server, MagicMock(), MagicMock()) + result = await identify_session(self.server, MagicMock(), MagicMock()) + + assert result.user_data == {"a": "1", "b": "2"} + assert len(_identify_events(mock_event_queue)) == 1 + + @patch("agentcat.modules.identify.event_queue") + async def test_async_identify_callback(self, mock_event_queue): + """AgentCatOptions.identify may be an async callable (P4 parity with + TypeScript's async identify).""" + calls = [] + + async def identify(request, context): + calls.append(request) + return UserIdentity( + user_id="async-user", user_name="Async User", user_data={"k": "v"} + ) + + self._setup_data(identify) + + result = await identify_session(self.server, MagicMock(), MagicMock()) + + assert calls, "async identify hook never awaited" + assert isinstance(result, UserIdentity) + assert result.user_id == "async-user" + + events = _identify_events(mock_event_queue) + assert len(events) == 1 + assert events[0].identify_actor_given_id == "async-user" + assert events[0].identify_data == {"k": "v"} + + @patch("agentcat.modules.identify.event_queue") + async def test_async_identify_exception_returns_none(self, mock_event_queue): + async def identify(request, context): + raise RuntimeError("async identify exploded") + + self._setup_data(identify) + + result = await identify_session(self.server, MagicMock(), MagicMock()) + + assert result is None + assert not _identify_events(mock_event_queue) + + @patch("agentcat.modules.identify.event_queue") + async def test_new_session_republishes_identity(self, mock_event_queue): + """Identity cache is keyed per session — a new session ID publishes + again even for an identical identity.""" + identify_fn = lambda req, ctx: UserIdentity( # noqa: E731 + user_id="alice", user_name="Alice", user_data=None + ) + data = self._setup_data(identify_fn, session_id="ses_first") + await identify_session(self.server, MagicMock(), MagicMock()) + + data.session_id = "ses_second" + set_server_tracking_data(self.server, data) + await identify_session(self.server, MagicMock(), MagicMock()) + + assert len(_identify_events(mock_event_queue)) == 2 + + +class TestIdentityHelpers: + """Unit tests for merge/equality helpers and the LRU cache.""" + + def test_are_identities_equal(self): + a = UserIdentity(user_id="u", user_name="n", user_data={"k": "v"}) + b = UserIdentity(user_id="u", user_name="n", user_data={"k": "v"}) + assert are_identities_equal(a, b) + + assert not are_identities_equal( + a, UserIdentity(user_id="x", user_name="n", user_data={"k": "v"}) + ) + assert not are_identities_equal( + a, UserIdentity(user_id="u", user_name="x", user_data={"k": "v"}) + ) + assert not are_identities_equal( + a, UserIdentity(user_id="u", user_name="n", user_data={"k": "other"}) + ) + # None and {} user_data are equivalent + assert are_identities_equal( + UserIdentity(user_id="u", user_name="n", user_data=None), + UserIdentity(user_id="u", user_name="n", user_data={}), + ) + + def test_merge_identities_no_previous(self): + nxt = UserIdentity(user_id="u", user_name="n", user_data={"k": "v"}) + assert merge_identities(None, nxt) is nxt + + def test_merge_identities_overwrites_and_merges(self): + prev = UserIdentity( + user_id="old", user_name="Old", user_data={"a": "1", "b": "2"} + ) + nxt = UserIdentity(user_id="new", user_name="New", user_data={"b": "3"}) + merged = merge_identities(prev, nxt) + assert merged.user_id == "new" + assert merged.user_name == "New" + assert merged.user_data == {"a": "1", "b": "3"} + + def test_identity_cache_lru_eviction(self): + cache = IdentityCache(max_size=2) + alice = UserIdentity(user_id="alice", user_name=None, user_data=None) + bob = UserIdentity(user_id="bob", user_name=None, user_data=None) + carol = UserIdentity(user_id="carol", user_name=None, user_data=None) + + cache.set("s1", alice) + cache.set("s2", bob) + # Touch s1 so s2 becomes least recently used + assert cache.get("s1") is alice + cache.set("s3", carol) + + assert len(cache) == 2 + assert cache.get("s2") is None + assert cache.get("s1") is alice + assert cache.get("s3") is carol + + def test_identity_cache_update_existing(self): + cache = IdentityCache(max_size=2) + v1 = UserIdentity(user_id="v1", user_name=None, user_data=None) + v2 = UserIdentity(user_id="v2", user_name=None, user_data=None) + cache.set("s1", v1) + cache.set("s1", v2) + assert len(cache) == 1 + assert cache.get("s1") is v2 diff --git a/tests/test_lowlevel_error_capture.py b/tests/test_lowlevel_error_capture.py new file mode 100644 index 0000000..aae2d11 --- /dev/null +++ b/tests/test_lowlevel_error_capture.py @@ -0,0 +1,175 @@ +"""Structured error capture on the low-level (non-FastMCP) server path. + +overrides/mcp_server.py must record errors via modules.exceptions.capture_exception +(message/type/frames/platform), matching the official and community paths, +instead of a bare {"message": ...} dict. +""" + +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest +from mcp.server import Server as LowLevelServer +from mcp.types import CallToolRequest, CallToolRequestParams + +from agentcat.modules.internal import ( + reset_all_tracking_data, + set_server_tracking_data, +) +from agentcat.modules.overrides.mcp_server import override_lowlevel_mcp_server +from agentcat.types import AgentCatData, AgentCatOptions, SessionInfo + + +def _make_request(name="boom", arguments=None) -> CallToolRequest: + return CallToolRequest( + method="tools/call", + params=CallToolRequestParams(name=name, arguments=arguments or {}), + ) + + +class TestLowLevelErrorCapture: + def setup_method(self): + reset_all_tracking_data() + + def teardown_method(self): + reset_all_tracking_data() + + def _setup_server(self) -> LowLevelServer: + server = LowLevelServer("lowlevel-error-server") + + @server.list_tools() + async def list_tools(): + return [] + + @server.call_tool() + async def call_tool(name, arguments): + raise ValueError("tool blew up") + + data = AgentCatData( + project_id="test_project", + session_id="ses_lowlevel_err", + session_info=SessionInfo(), + last_activity=datetime.now(timezone.utc), + options=AgentCatOptions( + enable_tracing=True, + enable_tool_call_context=False, + enable_report_missing=False, + ), + ) + set_server_tracking_data(server, data) + override_lowlevel_mcp_server(server, data) + return server + + async def test_error_response_captured_structurally(self): + """The MCP SDK converts tool exceptions into isError CallToolResults; + the published event must carry structured ErrorData, not just a bare + message dict.""" + server = self._setup_server() + handler = server.request_handlers[CallToolRequest] + + with patch( + "agentcat.modules.overrides.mcp_server.event_queue" + ) as mock_event_queue: + await handler(_make_request()) + + events = [ + call.args[1] for call in mock_event_queue.publish_event.call_args_list + ] + call_events = [e for e in events if e.event_type == "mcp:tools/call"] + assert call_events, "tools/call event not published" + event = call_events[0] + + assert event.is_error is True + assert isinstance(event.error, dict) + assert "tool blew up" in event.error["message"] + # Structured ErrorData shape from capture_exception + assert event.error["platform"] == "python" + assert "type" in event.error + + async def test_handler_exception_captured_with_frames(self): + """An exception escaping the original handler must be captured with + full structure: type, stack frames, and platform.""" + server = LowLevelServer("lowlevel-raise-server") + + async def exploding_handler(request): + raise RuntimeError("handler exploded") + + server.request_handlers[CallToolRequest] = exploding_handler + + data = AgentCatData( + project_id="test_project", + session_id="ses_lowlevel_raise", + session_info=SessionInfo(), + last_activity=datetime.now(timezone.utc), + options=AgentCatOptions( + enable_tracing=True, + enable_tool_call_context=False, + enable_report_missing=False, + ), + ) + set_server_tracking_data(server, data) + override_lowlevel_mcp_server(server, data) + handler = server.request_handlers[CallToolRequest] + + with patch( + "agentcat.modules.overrides.mcp_server.event_queue" + ) as mock_event_queue: + with pytest.raises(RuntimeError, match="handler exploded"): + await handler(_make_request()) + + events = [ + call.args[1] for call in mock_event_queue.publish_event.call_args_list + ] + call_events = [e for e in events if e.event_type == "mcp:tools/call"] + assert call_events, "tools/call event not published" + event = call_events[0] + + assert event.is_error is True + error = event.error + assert error["message"] == "handler exploded" + assert error["type"] == "RuntimeError" + assert error["platform"] == "python" + assert error["frames"], "expected structured stack frames" + assert any( + frame["function"] == "exploding_handler" for frame in error["frames"] + ) + assert "handler exploded" in error["stack"] + + async def test_success_response_has_no_error(self): + server = LowLevelServer("lowlevel-ok-server") + + @server.list_tools() + async def list_tools(): + return [] + + @server.call_tool() + async def call_tool(name, arguments): + return [] + + data = AgentCatData( + project_id="test_project", + session_id="ses_lowlevel_ok", + session_info=SessionInfo(), + last_activity=datetime.now(timezone.utc), + options=AgentCatOptions( + enable_tracing=True, + enable_tool_call_context=False, + enable_report_missing=False, + ), + ) + set_server_tracking_data(server, data) + override_lowlevel_mcp_server(server, data) + handler = server.request_handlers[CallToolRequest] + + with patch( + "agentcat.modules.overrides.mcp_server.event_queue" + ) as mock_event_queue: + await handler(_make_request(name="ok")) + + events = [ + call.args[1] for call in mock_event_queue.publish_event.call_args_list + ] + call_events = [e for e in events if e.event_type == "mcp:tools/call"] + assert call_events + assert call_events[0].is_error is False + assert call_events[0].error is None diff --git a/tests/test_posthog_exporter.py b/tests/test_posthog_exporter.py new file mode 100644 index 0000000..9d65ec1 --- /dev/null +++ b/tests/test_posthog_exporter.py @@ -0,0 +1,269 @@ +"""Tests for the PostHog telemetry exporter. + +Mirrors the TypeScript SDK's PostHogExporter (src/modules/exporters/posthog.ts): +events POST to /batch, error events add a $exception capture, and tool calls +optionally add an $ai_span capture when enable_ai_tracing is set. +""" + +import re +from datetime import datetime, timezone +from unittest.mock import MagicMock + +from agentcat.modules.exporters.posthog import ( + DEFAULT_POSTHOG_HOST, + PostHogExporter, + to_uuidv7, +) +from agentcat.modules.telemetry import TelemetryManager +from agentcat.types import Event +from agentcat.utils import generate_prefixed_ksuid + +UUID_V7_RE = re.compile( + r"^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$" +) + + +def make_event(**kwargs) -> Event: + defaults = { + "id": "evt_2QjWl3PYbPnPUyYd5b5mFyJXk1a", + "event_type": "mcp:tools/call", + "project_id": "proj_123", + "session_id": "ses_2QjWl3PYbPnPUyYd5b5mFyJXk1a", + "resource_name": "add_todo", + "duration": 250, + "timestamp": datetime(2025, 6, 1, 12, 0, 0, tzinfo=timezone.utc), + "server_name": "todo-server", + "server_version": "1.0.0", + "client_name": "Cursor", + "client_version": "2.6.22", + } + defaults.update(kwargs) + return Event(**defaults) + + +def export_and_get_payload(exporter: PostHogExporter, event: Event) -> dict: + exporter.session = MagicMock() + exporter.export(event) + assert exporter.session.post.called + call = exporter.session.post.call_args + return call.args[0] if call.args else call.kwargs.get("url"), call.kwargs["json"] + + +class TestToUUIDv7: + def test_deterministic(self): + session_id = generate_prefixed_ksuid("ses") + assert to_uuidv7(session_id) == to_uuidv7(session_id) + + def test_different_inputs_differ(self): + assert to_uuidv7(generate_prefixed_ksuid("ses")) != to_uuidv7( + generate_prefixed_ksuid("ses") + ) + + def test_valid_uuidv7_format(self): + result = to_uuidv7(generate_prefixed_ksuid("ses")) + assert UUID_V7_RE.match(result), result + + def test_invalid_ksuid_falls_back_without_error(self): + result = to_uuidv7("ses_not-a-real-ksuid") + assert UUID_V7_RE.match(result), result + + def test_none_input_does_not_crash(self): + result = to_uuidv7(None) + assert UUID_V7_RE.match(result), result + + +class TestPostHogExporterConfig: + def test_default_host(self): + exporter = PostHogExporter({"type": "posthog", "api_key": "phc_test"}) + assert exporter.batch_url == f"{DEFAULT_POSTHOG_HOST}/batch" + assert exporter.enable_ai_tracing is False + + def test_custom_host_trailing_slash_stripped(self): + exporter = PostHogExporter( + {"type": "posthog", "api_key": "phc_test", "host": "https://eu.i.posthog.com/"} + ) + assert exporter.batch_url == "https://eu.i.posthog.com/batch" + + def test_telemetry_manager_registration(self): + manager = TelemetryManager( + {"ph": {"type": "posthog", "api_key": "phc_test"}} + ) + assert manager.get_exporter_count() == 1 + assert isinstance(manager.exporters["ph"], PostHogExporter) + + +class TestPostHogCaptureEvent: + def _exporter(self, **config): + return PostHogExporter({"type": "posthog", "api_key": "phc_test", **config}) + + def test_batch_payload_structure(self): + _, payload = export_and_get_payload(self._exporter(), make_event()) + assert payload["api_key"] == "phc_test" + assert len(payload["batch"]) == 1 + capture = payload["batch"][0] + assert capture["event"] == "mcp_tool_call" + assert capture["type"] == "capture" + assert capture["timestamp"].startswith("2025-06-01T12:00:00") + + def test_distinct_id_prefers_actor(self): + _, payload = export_and_get_payload( + self._exporter(), make_event(identify_actor_given_id="alice") + ) + assert payload["batch"][0]["distinct_id"] == "alice" + + def test_distinct_id_falls_back_to_session_then_anonymous(self): + event = make_event() + _, payload = export_and_get_payload(self._exporter(), event) + assert payload["batch"][0]["distinct_id"] == event.session_id + + _, payload = export_and_get_payload( + self._exporter(), make_event(session_id=None) + ) + assert payload["batch"][0]["distinct_id"] == "anonymous" + + def test_capture_properties(self): + event = make_event(user_intent="add an item", parameters={"a": 1}) + _, payload = export_and_get_payload(self._exporter(), event) + props = payload["batch"][0]["properties"] + + assert props["$session_id"] == to_uuidv7(event.session_id) + assert UUID_V7_RE.match(props["$session_id"]) + assert props["source"] == "agentcat" + assert props["resource_name"] == "add_todo" + assert props["tool_name"] == "add_todo" # tools/call adds tool_name + assert props["duration_ms"] == 250 + assert props["server_name"] == "todo-server" + assert props["server_version"] == "1.0.0" + assert props["client_name"] == "Cursor" + assert props["client_version"] == "2.6.22" + assert props["project_id"] == "proj_123" + assert props["user_intent"] == "add an item" + assert props["parameters"] == {"a": 1} + + def test_event_name_mapping(self): + exporter = self._exporter() + assert exporter.map_event_type("mcp:tools/call") == "mcp_tool_call" + assert exporter.map_event_type("mcp:tools/list") == "mcp_tools_list" + assert exporter.map_event_type("mcp:initialize") == "mcp_initialize" + assert exporter.map_event_type("mcp:resources/read") == "mcp_resource_read" + assert exporter.map_event_type("mcp:resources/list") == "mcp_resources_list" + assert exporter.map_event_type("mcp:prompts/get") == "mcp_prompt_get" + assert exporter.map_event_type("mcp:prompts/list") == "mcp_prompts_list" + # Fallback for unmapped types + assert exporter.map_event_type("mcp:ping") == "mcp_ping" + assert ( + exporter.map_event_type("mcp:resources/templates/list") + == "mcp_resources_templates_list" + ) + + def test_person_properties_from_identity(self): + event = make_event( + identify_actor_given_id="alice", + identify_actor_name="Alice", + identify_data={"plan": "pro"}, + ) + _, payload = export_and_get_payload(self._exporter(), event) + props = payload["batch"][0]["properties"] + assert props["$set"] == {"name": "Alice", "plan": "pro"} + + def test_customer_tags_and_properties_spread(self): + event = make_event( + tags={"env": "prod", "source": "customer-override"}, + properties={"feature_flags": ["dark_mode"]}, + ) + _, payload = export_and_get_payload(self._exporter(), event) + props = payload["batch"][0]["properties"] + assert props["env"] == "prod" + # Customer tags can override AgentCat defaults + assert props["source"] == "customer-override" + assert props["feature_flags"] == ["dark_mode"] + + +class TestPostHogExceptionEvent: + def _exporter(self): + return PostHogExporter({"type": "posthog", "api_key": "phc_test"}) + + def test_error_event_adds_exception_capture(self): + event = make_event( + is_error=True, + error={ + "message": "boom", + "type": "ValueError", + "stack": "Traceback ...", + }, + ) + _, payload = export_and_get_payload(self._exporter(), event) + assert len(payload["batch"]) == 2 + exception = payload["batch"][1] + assert exception["event"] == "$exception" + props = exception["properties"] + assert props["$exception_source"] == "backend" + assert props["$exception_message"] == "boom" + assert props["$exception_type"] == "ValueError" + assert props["$exception_stacktrace"] == "Traceback ..." + assert props["$session_id"] == to_uuidv7(event.session_id) + assert props["tool_name"] == "add_todo" + + def test_non_error_event_has_no_exception(self): + _, payload = export_and_get_payload(self._exporter(), make_event()) + assert [e["event"] for e in payload["batch"]] == ["mcp_tool_call"] + + +class TestPostHogAISpanEvent: + def _exporter(self, enable_ai_tracing=True): + return PostHogExporter( + { + "type": "posthog", + "api_key": "phc_test", + "enable_ai_tracing": enable_ai_tracing, + } + ) + + def test_ai_span_disabled_by_default(self): + exporter = PostHogExporter({"type": "posthog", "api_key": "phc_test"}) + _, payload = export_and_get_payload(exporter, make_event()) + assert [e["event"] for e in payload["batch"]] == ["mcp_tool_call"] + + def test_ai_span_emitted_for_tool_calls(self): + event = make_event(parameters={"in": 1}, response={"out": 2}) + _, payload = export_and_get_payload(self._exporter(), event) + assert [e["event"] for e in payload["batch"]] == [ + "mcp_tool_call", + "$ai_span", + ] + props = payload["batch"][1]["properties"] + assert props["$ai_session_id"] == f"agentcat_{event.session_id}" + assert props["$ai_trace_id"] == to_uuidv7(event.session_id) + assert props["$ai_span_id"] == to_uuidv7(event.id) + assert props["$ai_span_name"] == "add_todo" + assert props["$ai_is_error"] is False + assert props["$ai_latency"] == 0.25 + assert props["$ai_input_state"] == {"in": 1} + assert props["$ai_output_state"] == {"out": 2} + + def test_ai_span_not_emitted_for_non_tool_calls(self): + event = make_event(event_type="mcp:initialize", resource_name=None) + _, payload = export_and_get_payload(self._exporter(), event) + assert [e["event"] for e in payload["batch"]] == ["mcp_initialize"] + + def test_ai_span_error_fields(self): + event = make_event(is_error=True, error={"message": "boom"}) + _, payload = export_and_get_payload(self._exporter(), event) + # capture + $exception + $ai_span + assert [e["event"] for e in payload["batch"]] == [ + "mcp_tool_call", + "$exception", + "$ai_span", + ] + props = payload["batch"][2]["properties"] + assert props["$ai_is_error"] is True + assert props["$ai_error"] == {"message": "boom"} + + +class TestPostHogExportResilience: + def test_network_error_is_swallowed(self): + exporter = PostHogExporter({"type": "posthog", "api_key": "phc_test"}) + exporter.session = MagicMock() + exporter.session.post.side_effect = RuntimeError("network down") + # Must not raise — telemetry never crashes the host + exporter.export(make_event()) diff --git a/tests/test_publish_custom_event.py b/tests/test_publish_custom_event.py new file mode 100644 index 0000000..9d55901 --- /dev/null +++ b/tests/test_publish_custom_event.py @@ -0,0 +1,166 @@ +"""Tests for agentcat.publish_custom_event. + +Mirrors the TypeScript SDK's publishCustomEvent (src/index.ts): custom events +carry the fixed "agentcat:custom" event type and can be published either +through a tracked server (using its session) or with a raw MCP session ID +string (deriving a deterministic AgentCat session ID). +""" + +from unittest.mock import patch + +import pytest + +import agentcat +from agentcat import AgentCatOptions, CustomEventData, publish_custom_event +from agentcat.modules.constants import AGENTCAT_CUSTOM_EVENT_TYPE +from agentcat.modules.internal import ( + get_server_tracking_data, + reset_all_tracking_data, +) +from agentcat.modules.session import derive_session_id_from_mcp_session +from agentcat.utils import parse_prefixed_ksuid + +from .test_utils.todo_server import create_todo_server + + +class TestDeriveSessionId: + """Deterministic session ID derivation from MCP session IDs.""" + + def test_same_inputs_produce_same_session_id(self): + a = derive_session_id_from_mcp_session("mcp-session-1", "proj_123") + b = derive_session_id_from_mcp_session("mcp-session-1", "proj_123") + assert a == b + + def test_different_session_ids_differ(self): + a = derive_session_id_from_mcp_session("mcp-session-1", "proj_123") + b = derive_session_id_from_mcp_session("mcp-session-2", "proj_123") + assert a != b + + def test_different_projects_differ(self): + a = derive_session_id_from_mcp_session("mcp-session-1", "proj_123") + b = derive_session_id_from_mcp_session("mcp-session-1", "proj_456") + assert a != b + + def test_project_id_optional(self): + a = derive_session_id_from_mcp_session("mcp-session-1") + b = derive_session_id_from_mcp_session("mcp-session-1") + assert a == b + assert a != derive_session_id_from_mcp_session("mcp-session-1", "proj_123") + + def test_result_is_valid_prefixed_ksuid(self): + derived = derive_session_id_from_mcp_session("mcp-session-1", "proj_123") + prefix, ksuid = parse_prefixed_ksuid(derived) + assert prefix == "ses" + # Timestamp must fall within [2024-01-01, 2024-01-01 + 1 year) + assert 1704067200 <= ksuid.timestamp < 1704067200 + 365 * 24 * 60 * 60 + + +class TestPublishCustomEventValidation: + def test_missing_project_id_raises(self): + with pytest.raises(ValueError, match="project_id is required"): + publish_custom_event("mcp-session-1", "") + + def test_none_project_id_raises(self): + with pytest.raises(ValueError, match="project_id is required"): + publish_custom_event("mcp-session-1", None) + + def test_untracked_server_raises(self): + reset_all_tracking_data() + server = create_todo_server() + with pytest.raises(ValueError, match="Server is not tracked"): + publish_custom_event(server, "proj_123") + + def test_invalid_first_parameter_raises(self): + with pytest.raises(TypeError, match="MCP server object or a session ID"): + publish_custom_event(42, "proj_123") + + def test_none_first_parameter_raises(self): + with pytest.raises(TypeError, match="MCP server object or a session ID"): + publish_custom_event(None, "proj_123") + + +class TestPublishCustomEventWithSessionId: + """String session ID path: event goes directly to the event queue.""" + + def _publish(self, event_data=None, mcp_session_id="mcp-session-1"): + with patch("agentcat.modules.event_queue.event_queue") as mock_queue: + publish_custom_event(mcp_session_id, "proj_123", event_data) + assert mock_queue.add.call_count == 1 + return mock_queue.add.call_args.args[0] + + def test_event_core_fields(self): + event = self._publish() + assert event.event_type == AGENTCAT_CUSTOM_EVENT_TYPE + assert event.project_id == "proj_123" + assert event.session_id == derive_session_id_from_mcp_session( + "mcp-session-1", "proj_123" + ) + assert event.timestamp is not None + + def test_event_data_fields_mapped(self): + event = self._publish( + CustomEventData( + resource_name="custom-action", + parameters={"action": "user-feedback", "rating": 5}, + response={"ok": True}, + message="User provided feedback", + duration=1234, + is_error=True, + error={"message": "Custom error occurred", "code": "ERR_001"}, + ) + ) + assert event.resource_name == "custom-action" + assert event.parameters == {"action": "user-feedback", "rating": 5} + assert event.response == {"ok": True} + assert event.user_intent == "User provided feedback" + assert event.duration == 1234 + assert event.is_error is True + assert event.error == {"message": "Custom error occurred", "code": "ERR_001"} + + def test_tags_are_validated(self): + event = self._publish( + CustomEventData( + tags={"env": "prod", "bad\nkey!": "x", "numeric": 5}, + ) + ) + assert event.tags == {"env": "prod"} + + def test_properties_attached(self): + event = self._publish( + CustomEventData(properties={"feature_flags": ["dark_mode"]}) + ) + assert event.properties == {"feature_flags": ["dark_mode"]} + + def test_empty_properties_omitted(self): + event = self._publish(CustomEventData(properties={})) + assert event.properties is None + + +class TestPublishCustomEventWithTrackedServer: + def setup_method(self): + reset_all_tracking_data() + + def teardown_method(self): + reset_all_tracking_data() + + def test_tracked_server_uses_its_session(self): + server = create_todo_server() + agentcat.track(server, "proj_tracked", AgentCatOptions()) + data = get_server_tracking_data(server) + assert data is not None + + with patch("agentcat.modules.event_queue.event_queue") as mock_queue: + publish_custom_event( + server, + "proj_tracked", + CustomEventData(resource_name="feature-usage"), + ) + assert mock_queue.add.call_count == 1 + event = mock_queue.add.call_args.args[0] + + assert event.event_type == AGENTCAT_CUSTOM_EVENT_TYPE + assert event.session_id == data.session_id + assert event.project_id == "proj_tracked" + assert event.resource_name == "feature-usage" + # Published via publish_event → session info merged in + assert event.sdk_language is not None diff --git a/tests/test_redaction.py b/tests/test_redaction.py index 533c649..3a769b8 100644 --- a/tests/test_redaction.py +++ b/tests/test_redaction.py @@ -1,5 +1,7 @@ """Unit tests for the redaction module.""" +from datetime import datetime, timezone + import pytest from typing import Any, Dict from agentcat.modules.redaction import ( @@ -7,6 +9,7 @@ redact_event, PROTECTED_FIELDS, ) +from agentcat.types import UnredactedEvent class TestRedactStringsInObject: @@ -367,3 +370,131 @@ def faulty_redact_fn(s: str) -> str: # The function should propagate the error with pytest.raises(ValueError, match="Redaction error"): redact_event(event, faulty_redact_fn) + + +class TestRedactEventPydanticModel: + """Redaction must fire on the live publish path, where events are + Pydantic UnredactedEvent models rather than plain dicts.""" + + def _make_event(self, **overrides) -> UnredactedEvent: + defaults = { + "id": "evt_123", + "session_id": "ses_abc", + "project_id": "proj_456", + "event_type": "mcp:tools/call", + "resource_name": "add_todo", + "timestamp": datetime.now(timezone.utc), + "parameters": { + "arguments": {"text": "secret-value"}, + "extra": {"headers": {"authorization": "Bearer token-xyz"}}, + }, + "user_intent": "do a secret thing", + } + defaults.update(overrides) + return UnredactedEvent(**defaults) + + def test_pydantic_event_strings_are_redacted(self): + def redact_fn(s: str) -> str: + return "[REDACTED]" + + event = self._make_event() + result = redact_event(event, redact_fn) + + # Returns a rebuilt model of the same class + assert isinstance(result, UnredactedEvent) + + # Protected fields preserved + assert result.id == "evt_123" + assert result.session_id == "ses_abc" + assert result.project_id == "proj_456" + assert result.event_type == "mcp:tools/call" + assert result.resource_name == "add_todo" + + # Unprotected string fields redacted, including nested parameters + assert result.user_intent == "[REDACTED]" + assert result.parameters["arguments"]["text"] == "[REDACTED]" + assert ( + result.parameters["extra"]["headers"]["authorization"] == "[REDACTED]" + ) + + # Non-string fields untouched + assert result.timestamp == event.timestamp + + def test_pydantic_event_original_not_mutated(self): + def redact_fn(s: str) -> str: + return "X" + + event = self._make_event() + redact_event(event, redact_fn) + assert event.parameters["arguments"]["text"] == "secret-value" + + def test_pydantic_event_identify_and_tags_protected(self): + def redact_fn(s: str) -> str: + return "[REDACTED]" + + event = self._make_event( + event_type="agentcat:identify", + identify_actor_given_id="user123", + identify_actor_name="John Doe", + identify_data={"email": "john@example.com"}, + tags={"env": "prod"}, + properties={"flag": "on"}, + ) + result = redact_event(event, redact_fn) + + assert result.identify_actor_given_id == "user123" + assert result.identify_actor_name == "John Doe" + assert result.identify_data == {"email": "john@example.com"} + assert result.tags == {"env": "prod"} + assert result.properties == {"flag": "on"} + + def test_pydantic_event_redaction_error_propagates(self): + def faulty_redact_fn(s: str) -> str: + raise RuntimeError("redaction exploded") + + event = self._make_event() + with pytest.raises(RuntimeError, match="redaction exploded"): + redact_event(event, faulty_redact_fn) + + +class TestAsyncRedactionFunction: + """RedactionFunction accepts sync or async callables; the pipeline must + resolve coroutine results (event queue workers run without a loop).""" + + def test_async_redact_fn_on_plain_object(self): + async def redact_fn(s: str) -> str: + return "[ASYNC-REDACTED]" + + result = redact_strings_in_object({"field": "secret"}, redact_fn) + assert result == {"field": "[ASYNC-REDACTED]"} + + def test_async_redact_fn_on_pydantic_event(self): + async def redact_fn(s: str) -> str: + return s.replace("secret", "[REDACTED]") + + event = UnredactedEvent( + id="evt_1", + session_id="ses_1", + event_type="mcp:tools/call", + parameters={"arguments": {"text": "a secret here"}}, + ) + result = redact_event(event, redact_fn) + assert result.parameters["arguments"]["text"] == "a [REDACTED] here" + assert result.session_id == "ses_1" + + def test_async_redact_fn_error_propagates(self): + async def redact_fn(s: str) -> str: + raise ValueError("async redaction error") + + with pytest.raises(ValueError, match="async redaction error"): + redact_strings_in_object({"field": "secret"}, redact_fn) + + async def test_async_redact_fn_inside_running_loop(self): + """Defensive path: resolving an async redact fn while a loop is + already running in this thread must not deadlock or fail.""" + + async def redact_fn(s: str) -> str: + return "[LOOPED]" + + result = redact_strings_in_object({"field": "secret"}, redact_fn) + assert result == {"field": "[LOOPED]"} diff --git a/tests/test_stateless.py b/tests/test_stateless.py index 222e8cf..33ee334 100644 --- a/tests/test_stateless.py +++ b/tests/test_stateless.py @@ -61,31 +61,31 @@ def test_stateless_session_id_is_none(self): assert session_id is None @patch("agentcat.modules.identify.event_queue") - def test_stateless_identify_runs_every_time(self, mock_event_queue): + async def test_stateless_identify_runs_every_time(self, mock_event_queue): """In stateless mode, identify should run on every call (no early-return guard).""" mock_fn = MagicMock(return_value=UserIdentity( user_id="alice", user_name="Alice", user_data=None )) self._setup_data(stateless=True, identify=mock_fn) - identify_session(self.server, MagicMock(), MagicMock()) - identify_session(self.server, MagicMock(), MagicMock()) + await identify_session(self.server, MagicMock(), MagicMock()) + await identify_session(self.server, MagicMock(), MagicMock()) assert mock_fn.call_count == 2 @patch("agentcat.modules.identify.event_queue") - def test_stateless_identify_returns_identity(self, mock_event_queue): + async def test_stateless_identify_returns_identity(self, mock_event_queue): """In stateless mode, identify_session() should return the UserIdentity.""" self._setup_data(stateless=True, identify=_make_identify_fn()) - result = identify_session(self.server, MagicMock(), MagicMock()) + result = await identify_session(self.server, MagicMock(), MagicMock()) assert isinstance(result, UserIdentity) assert result.user_id == "user_123" assert result.user_name == "Test User" @patch("agentcat.modules.identify.event_queue") - def test_stateful_identify_runs_every_time(self, mock_event_queue): + async def test_stateful_identify_runs_every_time(self, mock_event_queue): """Stateful mode runs identify on every request.""" mock_fn = MagicMock(return_value=UserIdentity( user_id="alice", user_name="Alice", user_data=None @@ -97,8 +97,8 @@ def test_stateful_identify_runs_every_time(self, mock_event_queue): assert isinstance(session_id, str) assert session_id.startswith("ses_") - identify_session(self.server, MagicMock(), MagicMock()) - identify_session(self.server, MagicMock(), MagicMock()) + await identify_session(self.server, MagicMock(), MagicMock()) + await identify_session(self.server, MagicMock(), MagicMock()) assert mock_fn.call_count == 2 @@ -131,23 +131,23 @@ def test_track_stateless_none_auto_detects(self): assert data.is_stateless is False @patch("agentcat.modules.identify.event_queue") - def test_stateless_identify_bad_return(self, mock_event_queue): + async def test_stateless_identify_bad_return(self, mock_event_queue): """In stateless mode, identify returning non-UserIdentity should return None.""" bad_fn = MagicMock(return_value="not a UserIdentity") self._setup_data(stateless=True, identify=bad_fn) - result = identify_session(self.server, MagicMock(), MagicMock()) + result = await identify_session(self.server, MagicMock(), MagicMock()) assert result is None assert bad_fn.call_count == 1 @patch("agentcat.modules.identify.event_queue") - def test_stateless_identify_exception(self, mock_event_queue): + async def test_stateless_identify_exception(self, mock_event_queue): """In stateless mode, identify raising should return None, not propagate.""" raising_fn = MagicMock(side_effect=RuntimeError("identify exploded")) self._setup_data(stateless=True, identify=raising_fn) - result = identify_session(self.server, MagicMock(), MagicMock()) + result = await identify_session(self.server, MagicMock(), MagicMock()) assert result is None assert raising_fn.call_count == 1 diff --git a/tests/test_truncation.py b/tests/test_truncation.py index 324b104..02f243c 100644 --- a/tests/test_truncation.py +++ b/tests/test_truncation.py @@ -16,6 +16,12 @@ MAX_EVENT_BYTES, MIN_DEPTH, TRUNCATABLE_FIELDS, + MAX_USER_INTENT_LENGTH, + MAX_ERROR_MESSAGE_LENGTH, + MAX_RESOURCE_NAME_LENGTH, + MAX_METADATA_LENGTH, + MAX_STACK_FRAMES, + MAX_CONTENT_TEXT_LENGTH, ) from agentcat.types import UnredactedEvent @@ -50,8 +56,8 @@ def test_string_over_limit_truncated_with_marker(self): assert len(result.encode("utf-8")) < len(original.encode("utf-8")) def test_utf8_multibyte_truncated_by_bytes_no_broken_codepoints(self): - # Each emoji is 4 bytes. 2560 emojis = 10,240 bytes = exactly at limit - s = "\U0001f600" * 2561 # 10,244 bytes — just over limit + # Each emoji is 4 bytes. 8192 emojis = 32,768 bytes = exactly at limit + s = "\U0001f600" * 8193 # 32,772 bytes — just over limit result = _truncate_value(s) byte_size = len(s.encode("utf-8")) assert f"[string truncated by AgentCat from {byte_size} bytes]" in result @@ -207,8 +213,12 @@ def test_single_large_string_under_limit(self): assert result_bytes <= MAX_EVENT_BYTES def test_many_large_strings_under_limit(self): - """20 strings of 10 KB each = 200 KB of strings before truncation.""" - params = {f"key_{i}": "x" * 20_000 for i in range(20)} + """20 strings of 40 KB each = 800 KB of strings before truncation. + + Each string exceeds MAX_STRING_BYTES so per-string truncation + applies on the first pass. + """ + params = {f"key_{i}": "x" * 40_000 for i in range(20)} event = _make_event(parameters=params) result = truncate_event(event) result_bytes = len(result.model_dump_json().encode("utf-8")) @@ -226,10 +236,10 @@ def test_deeply_nested_wide_structure_under_limit(self): def test_depth_reduces_progressively(self): """Verify depth reduction kicks in when first pass isn't enough.""" - # Build a structure that's over 100 KB even after depth=5 truncation - # 20 keys * 10 KB string = 200 KB at each level, nested 5 deep + # Build a structure that's over 100 KB even after the first pass at + # MAX_DEPTH: 20 keys * 10 KB string = 200 KB leaf, nested MAX_DEPTH deep level = {f"k{i}": "x" * 10_000 for i in range(20)} - for _ in range(5): + for _ in range(MAX_DEPTH): level = {"nested": level, "extra": "x" * 10_000} event = _make_event(parameters=level) result = truncate_event(event) @@ -557,3 +567,168 @@ def test_large_user_intent_truncated_but_metadata_preserved(self): assert "truncated by AgentCat" in result.user_intent result_bytes = len(result.model_dump_json().encode("utf-8")) assert result_bytes <= MAX_EVENT_BYTES + + +class TestFieldLevelCaps: + """Field-level caps run unconditionally on every event, even ones far + under the 100 KB budget (parity with the TypeScript SDK's layer 1).""" + + # --- user_intent -> 2048 --- + + def test_user_intent_over_cap_truncated_with_marker(self): + event = _make_event(user_intent="x" * (MAX_USER_INTENT_LENGTH + 100)) + result = truncate_event(event) + assert result is not event + assert len(result.user_intent.encode("utf-8")) <= MAX_USER_INTENT_LENGTH + assert "truncated by AgentCat" in result.user_intent + + def test_user_intent_under_cap_untouched(self): + intent = "x" * MAX_USER_INTENT_LENGTH + event = _make_event(user_intent=intent) + result = truncate_event(event) + assert result.user_intent == intent + + # --- error.message -> 2048 --- + + def test_error_message_over_cap_truncated_with_marker(self): + event = _make_event( + error={"message": "e" * (MAX_ERROR_MESSAGE_LENGTH + 100), "platform": "python"} + ) + result = truncate_event(event) + assert len(result.error["message"].encode("utf-8")) <= MAX_ERROR_MESSAGE_LENGTH + assert "truncated by AgentCat" in result.error["message"] + # Sibling fields survive + assert result.error["platform"] == "python" + + def test_error_message_under_cap_untouched(self): + message = "e" * MAX_ERROR_MESSAGE_LENGTH + event = _make_event(error={"message": message, "platform": "python"}) + result = truncate_event(event) + assert result.error["message"] == message + + # --- resource_name -> 256 --- + + def test_resource_name_over_cap_truncated_with_marker(self): + event = _make_event(resource_name="r" * (MAX_RESOURCE_NAME_LENGTH + 50)) + result = truncate_event(event) + assert len(result.resource_name.encode("utf-8")) <= MAX_RESOURCE_NAME_LENGTH + assert "truncated by AgentCat" in result.resource_name + + def test_resource_name_under_cap_untouched(self): + name = "r" * MAX_RESOURCE_NAME_LENGTH + event = _make_event(resource_name=name) + result = truncate_event(event) + assert result.resource_name == name + + # --- server/client metadata -> 256 --- + + @pytest.mark.parametrize( + "field_name", + ["server_name", "server_version", "client_name", "client_version"], + ) + def test_metadata_field_over_cap_truncated_with_marker(self, field_name): + event = _make_event(**{field_name: "m" * (MAX_METADATA_LENGTH + 50)}) + result = truncate_event(event) + value = getattr(result, field_name) + assert len(value.encode("utf-8")) <= MAX_METADATA_LENGTH + assert "truncated by AgentCat" in value + + @pytest.mark.parametrize( + "field_name", + ["server_name", "server_version", "client_name", "client_version"], + ) + def test_metadata_field_under_cap_untouched(self, field_name): + value = "m" * MAX_METADATA_LENGTH + event = _make_event(**{field_name: value}) + result = truncate_event(event) + assert getattr(result, field_name) == value + + # --- error.frames -> max 50, keep first 25 + last 25 --- + + def test_frames_over_cap_keep_first_25_and_last_25(self): + frames = [ + {"filename": f"file_{i}.py", "function": f"fn_{i}", "lineno": i} + for i in range(80) + ] + event = _make_event(error={"message": "boom", "frames": frames}) + result = truncate_event(event) + trimmed = result.error["frames"] + assert len(trimmed) == MAX_STACK_FRAMES + assert trimmed[:25] == frames[:25] + assert trimmed[25:] == frames[-25:] + + def test_frames_at_cap_untouched(self): + frames = [ + {"filename": f"file_{i}.py", "function": f"fn_{i}", "lineno": i} + for i in range(MAX_STACK_FRAMES) + ] + event = _make_event(error={"message": "boom", "frames": frames}) + result = truncate_event(event) + assert result.error["frames"] == frames + + # --- response content text blocks -> 32 KB each --- + + def test_content_text_block_over_cap_truncated_with_marker(self): + big_text = "t" * (MAX_CONTENT_TEXT_LENGTH + 1_000) + event = _make_event( + response={"content": [{"type": "text", "text": big_text}]} + ) + result = truncate_event(event) + block = result.response["content"][0] + assert block["type"] == "text" + assert len(block["text"].encode("utf-8")) <= MAX_CONTENT_TEXT_LENGTH + assert "truncated by AgentCat" in block["text"] + + def test_content_text_block_under_cap_untouched(self): + text = "t" * 1_000 + event = _make_event( + response={"content": [{"type": "text", "text": text}]} + ) + result = truncate_event(event) + assert result.response["content"][0]["text"] == text + + def test_non_text_content_blocks_untouched(self): + blocks = [ + {"type": "image", "data": "abc", "mimeType": "image/png"}, + {"type": "text", "text": "t" * (MAX_CONTENT_TEXT_LENGTH + 500)}, + ] + event = _make_event(response={"content": blocks}) + result = truncate_event(event) + assert result.response["content"][0] == blocks[0] + assert "truncated by AgentCat" in result.response["content"][1]["text"] + + # --- layering / general behavior --- + + def test_caps_apply_even_when_event_far_under_budget(self): + """Field caps are unconditional — a tiny event with one over-cap field + is still truncated, unlike the old budget-only behavior.""" + event = _make_event(user_intent="x" * 5_000) # event well under 100 KB + assert len(event.model_dump_json().encode("utf-8")) <= MAX_EVENT_BYTES + result = truncate_event(event) + assert result is not event + assert len(result.user_intent.encode("utf-8")) <= MAX_USER_INTENT_LENGTH + + def test_event_with_no_over_cap_fields_returned_same_object(self): + event = _make_event( + user_intent="short", + server_name="srv", + error={"message": "boom", "frames": [{"filename": "a.py"}]}, + response={"content": [{"type": "text", "text": "ok"}]}, + ) + result = truncate_event(event) + assert result is event + + def test_original_event_not_mutated_by_field_caps(self): + frames = [{"filename": f"f{i}.py"} for i in range(80)] + error = {"message": "m" * (MAX_ERROR_MESSAGE_LENGTH + 100), "frames": frames} + response = {"content": [{"type": "text", "text": "t" * (MAX_CONTENT_TEXT_LENGTH + 100)}]} + event = _make_event( + user_intent="x" * (MAX_USER_INTENT_LENGTH + 100), + error=error, + response=response, + ) + truncate_event(event) + assert len(event.user_intent) == MAX_USER_INTENT_LENGTH + 100 + assert len(event.error["message"]) == MAX_ERROR_MESSAGE_LENGTH + 100 + assert len(event.error["frames"]) == 80 + assert len(event.response["content"][0]["text"]) == MAX_CONTENT_TEXT_LENGTH + 100 From 386fa2a33071cf1a2ce6e0ff7608136b265cda15 Mon Sep 17 00:00:00 2001 From: Naseem AlNaji Date: Thu, 9 Jul 2026 06:59:13 -0400 Subject: [PATCH 2/3] feat: publish agentcat:identify on every identify-hook return The identify hook already runs on every captured request; an agentcat:identify event is now published each time it returns a valid identity, instead of only when the identity changed. Merge semantics are unchanged: user_id/user_name are overwritten and user_data keys accumulate across calls, with the LRU cache retained for merge continuity across server-object recreation. Stateless mode now bypasses the identity cache and publishes the raw identity per request, so user_data from different actors sharing a server instance can no longer merge together. Tests updated to the new semantics, including a stateless no-cross-request-merge case. Full suite: 611 passed, 4 skipped. --- README.md | 2 +- src/agentcat/modules/identify.py | 73 ++++++++++++++------------------ tests/test_failure_injection.py | 4 +- tests/test_identify.py | 71 +++++++++++++------------------ tests/test_stateless.py | 48 ++++++++++++++++++++- 5 files changed, 111 insertions(+), 87 deletions(-) diff --git a/README.md b/README.md index cab7954..b84a053 100644 --- a/README.md +++ b/README.md @@ -72,7 +72,7 @@ agentcat.track(server, "proj_0000000") ### Identifying users -You can identify your user sessions with a simple callback AgentCat exposes, called `identify`. +You can identify your user sessions with a simple callback AgentCat exposes, called `identify`. The hook runs on every captured request and an `agentcat:identify` event is published each time it returns an identity; `user_id`/`user_name` are overwritten and `user_data` keys merge across calls. ```python from agentcat import AgentCatOptions, UserIdentity diff --git a/src/agentcat/modules/identify.py b/src/agentcat/modules/identify.py index 9d4e578..ccbd40d 100644 --- a/src/agentcat/modules/identify.py +++ b/src/agentcat/modules/identify.py @@ -15,9 +15,10 @@ class IdentityCache: """Simple LRU cache for session identities. - Prevents memory leaks by capping at max_size entries. This cache persists - across server instance restarts (mirrors the TypeScript SDK's - IdentityCache in modules/internal.ts). + Provides user_data merge continuity across identify calls for the same + session, and persists across server instance restarts (mirrors the + TypeScript SDK's IdentityCache in modules/internal.ts). Prevents memory + leaks by capping at max_size entries. """ def __init__(self, max_size: int = IDENTITY_CACHE_MAX_SIZE): @@ -48,7 +49,7 @@ def __len__(self) -> int: # Global identity cache shared across all server instances. -# This prevents duplicate identify events when server objects are recreated. +# This preserves user_data merge continuity when server objects are recreated. _global_identity_cache = IdentityCache() @@ -58,15 +59,6 @@ def reset_identity_cache() -> None: _global_identity_cache = IdentityCache() -def are_identities_equal(a: UserIdentity, b: UserIdentity) -> bool: - """Deep comparison of two UserIdentity objects.""" - if a.user_id != b.user_id: - return False - if a.user_name != b.user_name: - return False - return (a.user_data or {}) == (b.user_data or {}) - - def merge_identities( previous: UserIdentity | None, next_identity: UserIdentity ) -> UserIdentity: @@ -94,9 +86,11 @@ async def identify_session(server, request: any, context: any) -> UserIdentity | The identify hook may be a sync or an async callable; coroutine results are awaited. The resolved identity is merged with the session's previous identity (user_id/user_name overwritten, user_data keys merged) and an - `agentcat:identify` event is published ONLY when the merged identity - differs from the cached one — repeated identical identities do not - republish. + `agentcat:identify` event is published every time the hook returns an + identity; the cache provides merge continuity only. In stateless mode the + cache is bypassed entirely (it is keyed on the instance-wide session_id, + so caching would merge different actors' user_data together) and the raw + identity is published as-is. Returns the merged UserIdentity, or None if no hook is configured, the hook raises, or it returns a non-UserIdentity value. @@ -119,32 +113,29 @@ async def identify_session(server, request: any, context: any) -> UserIdentity | session_id = data.session_id - # Check the global cache (works across server instance restarts) - previous_identity = _global_identity_cache.get(session_id) - - # Merge identities (overwrite user_id/user_name, merge user_data) - merged_identity = merge_identities(previous_identity, identify_result) - - # Only publish if the identity has changed - has_changed = previous_identity is None or not are_identities_equal( - previous_identity, merged_identity + if data.is_stateless: + # Stateless mode: the merge cache is keyed on the instance-wide + # session_id, so different actors' user_data would merge together. + # Bypass the cache entirely and publish the raw identity. + merged_identity = identify_result + else: + # Check the global cache (works across server instance restarts) + previous_identity = _global_identity_cache.get(session_id) + + # Merge identities (overwrite user_id/user_name, merge user_data) + merged_identity = merge_identities(previous_identity, identify_result) + + _global_identity_cache.set(session_id, merged_identity) + + event = UnredactedEvent( + session_id=session_id, + timestamp=datetime.now(timezone.utc), + event_type=EventType.AGENTCAT_IDENTIFY.value, + identify_actor_given_id=merged_identity.user_id, + identify_actor_name=merged_identity.user_name, + identify_data=merged_identity.user_data or {}, ) - - _global_identity_cache.set(session_id, merged_identity) - - if has_changed: - write_to_log( - f"Identified session {session_id} (actor: {merged_identity.user_id})" - ) - event = UnredactedEvent( - session_id=session_id, - timestamp=datetime.now(timezone.utc), - event_type=EventType.AGENTCAT_IDENTIFY.value, - identify_actor_given_id=merged_identity.user_id, - identify_actor_name=merged_identity.user_name, - identify_data=merged_identity.user_data or {}, - ) - event_queue.publish_event(server, event) + event_queue.publish_event(server, event) return merged_identity except Exception as e: diff --git a/tests/test_failure_injection.py b/tests/test_failure_injection.py index a98a6ca..4870b07 100644 --- a/tests/test_failure_injection.py +++ b/tests/test_failure_injection.py @@ -204,7 +204,7 @@ def identify(request, extra): @patch("agentcat.modules.identify.event_queue") @patch("agentcat.modules.overrides.mcp_server.event_queue") async def test_identify_user_data_with_raising_eq(self, mock_eq, mock_id_eq): - """Change-detection deep-compares arbitrary user_data; values whose + """Identity merging handles arbitrary user_data; values whose __eq__ raises must not break the second request.""" def identify(request, extra): @@ -214,7 +214,7 @@ def identify(request, extra): server = _setup_lowlevel_server(identify=identify) result1 = await self._call(server) - result2 = await self._call(server) # triggers are_identities_equal + result2 = await self._call(server) # merges with the cached identity assert not getattr(result1.root, "isError", False) assert not getattr(result2.root, "isError", False) diff --git a/tests/test_identify.py b/tests/test_identify.py index a1e9ab5..ab91337 100644 --- a/tests/test_identify.py +++ b/tests/test_identify.py @@ -1,8 +1,9 @@ -"""Tests for identify: async hooks, change detection, and identity merging. +"""Tests for identify: async hooks, per-request publishing, and identity merging. -Mirrors the TypeScript SDK's handleIdentify (src/modules/internal.ts): the -identify hook runs per-request but an `agentcat:identify` event is published -only when the (merged) identity actually changed for the session. +An `agentcat:identify` event is published every time the identify hook returns +a valid identity — no change-detection dedup. The IdentityCache LRU is retained +solely for user_data merge continuity across calls (and across server-object +recreation): user_id/user_name are overwritten, user_data keys accumulate. """ from datetime import datetime, timezone @@ -10,7 +11,6 @@ from agentcat.modules.identify import ( IdentityCache, - are_identities_equal, identify_session, merge_identities, reset_identity_cache, @@ -34,7 +34,7 @@ def _identify_events(mock_event_queue): class TestIdentifySession: - """Change-detection and merge behavior of identify_session.""" + """Publish-every-time and merge behavior of identify_session.""" def setup_method(self): reset_all_tracking_data() @@ -57,9 +57,9 @@ def _setup_data(self, identify, session_id="ses_identify_test"): return data @patch("agentcat.modules.identify.event_queue") - async def test_identical_identity_publishes_single_event(self, mock_event_queue): - """Repeated identical identities → identify hook runs each time but - only ONE agentcat:identify event is published.""" + async def test_identical_identity_publishes_every_time(self, mock_event_queue): + """Repeated identical identities → the identify hook runs each time + and an agentcat:identify event is published each time.""" identify_fn = MagicMock( side_effect=lambda req, ctx: UserIdentity( user_id="alice", user_name="Alice", user_data={"plan": "pro"} @@ -73,11 +73,14 @@ async def test_identical_identity_publishes_single_event(self, mock_event_queue) assert result.user_id == "alice" assert identify_fn.call_count == 3 - assert len(_identify_events(mock_event_queue)) == 1 + events = _identify_events(mock_event_queue) + assert len(events) == 3 + assert all(e.identify_actor_given_id == "alice" for e in events) @patch("agentcat.modules.identify.event_queue") async def test_changed_identity_publishes_new_event(self, mock_event_queue): - """A changed identity → a new agentcat:identify event.""" + """Every hook return publishes; a changed identity carries the new + user_id in its event.""" identities = iter( [ UserIdentity(user_id="alice", user_name="Alice", user_data=None), @@ -92,15 +95,17 @@ async def test_changed_identity_publishes_new_event(self, mock_event_queue): result = await identify_session(self.server, MagicMock(), MagicMock()) events = _identify_events(mock_event_queue) - assert len(events) == 2 + assert len(events) == 3 assert events[0].identify_actor_given_id == "alice" - assert events[1].identify_actor_given_id == "bob" + assert events[1].identify_actor_given_id == "alice" + assert events[2].identify_actor_given_id == "bob" assert result.user_id == "bob" @patch("agentcat.modules.identify.event_queue") async def test_user_data_merges_across_calls(self, mock_event_queue): """user_id/user_name are overwritten; user_data keys are merged with - the newest values winning on collision.""" + the newest values winning on collision. Each call publishes an event + carrying the merged identity.""" identities = iter( [ UserIdentity( @@ -130,11 +135,10 @@ async def test_user_data_merges_across_calls(self, mock_event_queue): assert events[1].identify_data == {"plan": "pro", "region": "us"} @patch("agentcat.modules.identify.event_queue") - async def test_merged_but_unchanged_identity_does_not_republish( - self, mock_event_queue - ): + async def test_merged_unchanged_identity_republishes(self, mock_event_queue): """A subset of previously-seen user_data merges into an identical - identity → no new event.""" + identity → still republished, and the second event carries the FULL + merged identity (merge continuity + publish-every-time).""" identities = iter( [ UserIdentity( @@ -149,7 +153,10 @@ async def test_merged_but_unchanged_identity_does_not_republish( result = await identify_session(self.server, MagicMock(), MagicMock()) assert result.user_data == {"a": "1", "b": "2"} - assert len(_identify_events(mock_event_queue)) == 1 + events = _identify_events(mock_event_queue) + assert len(events) == 2 + assert events[1].identify_actor_given_id == "alice" + assert events[1].identify_data == {"a": "1", "b": "2"} @patch("agentcat.modules.identify.event_queue") async def test_async_identify_callback(self, mock_event_queue): @@ -190,8 +197,8 @@ async def identify(request, context): @patch("agentcat.modules.identify.event_queue") async def test_new_session_republishes_identity(self, mock_event_queue): - """Identity cache is keyed per session — a new session ID publishes - again even for an identical identity.""" + """Merge cache is keyed per session — a new session ID starts with no + previous identity and publishes on its first identify too.""" identify_fn = lambda req, ctx: UserIdentity( # noqa: E731 user_id="alice", user_name="Alice", user_data=None ) @@ -206,27 +213,7 @@ async def test_new_session_republishes_identity(self, mock_event_queue): class TestIdentityHelpers: - """Unit tests for merge/equality helpers and the LRU cache.""" - - def test_are_identities_equal(self): - a = UserIdentity(user_id="u", user_name="n", user_data={"k": "v"}) - b = UserIdentity(user_id="u", user_name="n", user_data={"k": "v"}) - assert are_identities_equal(a, b) - - assert not are_identities_equal( - a, UserIdentity(user_id="x", user_name="n", user_data={"k": "v"}) - ) - assert not are_identities_equal( - a, UserIdentity(user_id="u", user_name="x", user_data={"k": "v"}) - ) - assert not are_identities_equal( - a, UserIdentity(user_id="u", user_name="n", user_data={"k": "other"}) - ) - # None and {} user_data are equivalent - assert are_identities_equal( - UserIdentity(user_id="u", user_name="n", user_data=None), - UserIdentity(user_id="u", user_name="n", user_data={}), - ) + """Unit tests for the merge helper and the LRU cache.""" def test_merge_identities_no_previous(self): nxt = UserIdentity(user_id="u", user_name="n", user_data={"k": "v"}) diff --git a/tests/test_stateless.py b/tests/test_stateless.py index 33ee334..b906e2c 100644 --- a/tests/test_stateless.py +++ b/tests/test_stateless.py @@ -16,6 +16,15 @@ from .test_utils.todo_server import create_todo_server +def _identify_events(mock_event_queue): + """Extract agentcat:identify events passed to the mocked publish_event.""" + return [ + call.args[1] + for call in mock_event_queue.publish_event.call_args_list + if call.args[1].event_type == "agentcat:identify" + ] + + def _make_identify_fn(user_id="user_123", user_name="Test User"): """Return an identify function that always returns a UserIdentity.""" def identify(request, context): @@ -62,7 +71,8 @@ def test_stateless_session_id_is_none(self): @patch("agentcat.modules.identify.event_queue") async def test_stateless_identify_runs_every_time(self, mock_event_queue): - """In stateless mode, identify should run on every call (no early-return guard).""" + """In stateless mode, identify should run on every call (no early-return guard) + and an agentcat:identify event should be published per request.""" mock_fn = MagicMock(return_value=UserIdentity( user_id="alice", user_name="Alice", user_data=None )) @@ -72,6 +82,42 @@ async def test_stateless_identify_runs_every_time(self, mock_event_queue): await identify_session(self.server, MagicMock(), MagicMock()) assert mock_fn.call_count == 2 + events = _identify_events(mock_event_queue) + assert len(events) == 2 + assert all(e.identify_actor_given_id == "alice" for e in events) + + @patch("agentcat.modules.identify.event_queue") + async def test_stateless_no_cross_request_identity_merge(self, mock_event_queue): + """In stateless mode the merge cache is bypassed: different actors on + consecutive requests must each get their own raw identity, with no + user_data bleeding between them.""" + identities = iter( + [ + UserIdentity( + user_id="alice", user_name="Alice", user_data={"plan": "pro"} + ), + UserIdentity( + user_id="bob", user_name="Bob", user_data={"region": "eu"} + ), + ] + ) + self._setup_data(stateless=True, identify=lambda req, ctx: next(identities)) + + result1 = await identify_session(self.server, MagicMock(), MagicMock()) + result2 = await identify_session(self.server, MagicMock(), MagicMock()) + + # Returned identities are the raw per-request identities (no merge) + assert result1.user_id == "alice" + assert result1.user_data == {"plan": "pro"} + assert result2.user_id == "bob" + assert result2.user_data == {"region": "eu"} + + events = _identify_events(mock_event_queue) + assert len(events) == 2 + assert events[0].identify_actor_given_id == "alice" + assert events[0].identify_data == {"plan": "pro"} + assert events[1].identify_actor_given_id == "bob" + assert events[1].identify_data == {"region": "eu"} @patch("agentcat.modules.identify.event_queue") async def test_stateless_identify_returns_identity(self, mock_event_queue): From 50bb1d31ec57eea28075edab925f4a5e9eaa63a4 Mon Sep 17 00:00:00 2001 From: Naseem AlNaji Date: Thu, 9 Jul 2026 08:47:14 -0400 Subject: [PATCH 3/3] refactor: deduplicate helpers and reduce per-event overhead Code-quality pass over the branch; no behavior changes. - exceptions reuses logging.safe_error_string instead of a local copy. - Async redaction functions now cost one event loop per event instead of one per redacted string; sync paths are unchanged. - Redaction rebuilds the event with model_construct, skipping re-validation of data that came from an already-valid model (serialization verified identical). - PostHog exporter computes the session UUID once per export. - publish_custom_event normalizes its optional payload once instead of guarding every field access. - Event-type validator generalizes to the set of SDK-defined event types instead of special-casing one constant. - Removed dead IdentityCache.__contains__ and a redundant second parse of error results on the low-level path. Full suite: 611 passed, 4 skipped. --- src/agentcat/__init__.py | 21 ++-- src/agentcat/modules/constants.py | 4 + src/agentcat/modules/exceptions.py | 14 +-- src/agentcat/modules/exporters/posthog.py | 30 +++-- src/agentcat/modules/identify.py | 3 - src/agentcat/modules/overrides/mcp_server.py | 4 +- src/agentcat/modules/redaction.py | 121 ++++++++++++++++--- src/agentcat/types.py | 4 +- 8 files changed, 146 insertions(+), 55 deletions(-) diff --git a/src/agentcat/__init__.py b/src/agentcat/__init__.py index 5fe3778..c056017 100644 --- a/src/agentcat/__init__.py +++ b/src/agentcat/__init__.py @@ -287,6 +287,9 @@ def publish_custom_event( """ from agentcat.modules import event_queue as event_queue_module + # Normalize once: an omitted payload behaves like an all-defaults payload. + event_data = event_data or CustomEventData() + if not project_id: raise ValueError("project_id is required for publish_custom_event") @@ -329,19 +332,19 @@ def publish_custom_event( event_type=AGENTCAT_CUSTOM_EVENT_TYPE, timestamp=datetime.now(timezone.utc), # Event data from parameters - resource_name=event_data.resource_name if event_data else None, - parameters=event_data.parameters if event_data else None, - response=event_data.response if event_data else None, - user_intent=event_data.message if event_data else None, - duration=event_data.duration if event_data else None, - is_error=event_data.is_error if event_data else None, - error=event_data.error if event_data else None, + resource_name=event_data.resource_name, + parameters=event_data.parameters, + response=event_data.response, + user_intent=event_data.message, + duration=event_data.duration, + is_error=event_data.is_error, + error=event_data.error, ) # Wire up customer-defined metadata - if event_data and event_data.tags: + if event_data.tags: event.tags = validate_tags(event_data.tags) - if event_data and event_data.properties: + if event_data.properties: event.properties = event_data.properties # If we have a tracked server, publish through it (merges session info and diff --git a/src/agentcat/modules/constants.py b/src/agentcat/modules/constants.py index e8c0a86..69e97a2 100644 --- a/src/agentcat/modules/constants.py +++ b/src/agentcat/modules/constants.py @@ -5,6 +5,10 @@ AGENTCAT_API_URL = "https://api.agentcat.com" # Default API URL for AgentCat events AGENTCAT_SOURCE = "agentcat" # Source attribution for telemetry exporters AGENTCAT_CUSTOM_EVENT_TYPE = "agentcat:custom" # Event type for publish_custom_event +# Event types defined by this SDK (all "agentcat:"-prefixed) that the generated +# API client's enum doesn't know about; types.Event lets these bypass the +# generated enum check. Add new SDK-defined event types here. +SDK_EVENT_TYPES = frozenset({AGENTCAT_CUSTOM_EVENT_TYPE}) DEFAULT_CONTEXT_DESCRIPTION = "Explain why you are calling this tool and how it fits into the user's overall goal. This parameter is used for analytics and user intent tracking. YOU MUST provide 15-25 words (count carefully). NEVER use first person ('I', 'we', 'you') - maintain third-person perspective. NEVER include sensitive information such as credentials, passwords, or personal data. Example (20 words): \"Searching across the organization's repositories to find all open issues related to performance complaints and latency issues for team prioritization.\"" # Maximum number of exceptions to capture in a cause chain diff --git a/src/agentcat/modules/exceptions.py b/src/agentcat/modules/exceptions.py index 56b1727..2619dc6 100644 --- a/src/agentcat/modules/exceptions.py +++ b/src/agentcat/modules/exceptions.py @@ -12,22 +12,14 @@ from agentcat.types import ChainedErrorData, ErrorData, StackFrame from agentcat.modules.constants import MAX_EXCEPTION_CHAIN_DEPTH, MAX_STACK_FRAMES +# str(value) that never raises (exceptions may have a broken __str__). +from agentcat.modules.logging import safe_error_string as _safe_str + _captured_error: contextvars.ContextVar[BaseException | None] = contextvars.ContextVar( "_captured_error", default=None ) -def _safe_str(value: Any) -> str: - """str(value) that never raises (exceptions may have a broken __str__).""" - try: - return str(value) - except Exception: - try: - return f"" - except Exception: - return "" - - def capture_exception(exc: BaseException | Any) -> ErrorData: """ Captures detailed exception information including stack traces and cause chains. diff --git a/src/agentcat/modules/exporters/posthog.py b/src/agentcat/modules/exporters/posthog.py index c1d7b23..e12aaad 100644 --- a/src/agentcat/modules/exporters/posthog.py +++ b/src/agentcat/modules/exporters/posthog.py @@ -117,19 +117,25 @@ def export(self, event: Event) -> None: try: batch: list[dict[str, Any]] = [] + # Compute the deterministic UUIDs once per event (KSUID parse + + # SHA-256 each time) and share them across the builders. + session_uuid = to_uuidv7(event.session_id) + # Always send the regular event - batch.append(self.build_capture_event(event)) + batch.append(self.build_capture_event(event, session_uuid)) # Send $exception event alongside if this is an error if event.is_error and event.error: - batch.append(self.build_exception_event(event)) + batch.append(self.build_exception_event(event, session_uuid)) # Send $ai_span for tool calls when AI tracing is enabled if ( self.enable_ai_tracing and event.event_type == EventType.MCP_TOOLS_CALL.value ): - batch.append(self.build_ai_span_event(event)) + batch.append( + self.build_ai_span_event(event, session_uuid, to_uuidv7(event.id)) + ) write_to_log( f"PostHogExporter: Sending {len(batch)} event(s) for {event.id}" @@ -152,10 +158,10 @@ def export(self, event: Event) -> None: except Exception as error: write_to_log(f"PostHog export error: {error}") - def build_capture_event(self, event: Event) -> dict[str, Any]: + def build_capture_event(self, event: Event, session_uuid: str) -> dict[str, Any]: """Build the regular PostHog capture event.""" properties: dict[str, Any] = { - "$session_id": to_uuidv7(event.session_id), + "$session_id": session_uuid, "source": AGENTCAT_SOURCE, } @@ -210,11 +216,11 @@ def build_capture_event(self, event: Event) -> dict[str, Any]: "type": "capture", } - def build_exception_event(self, event: Event) -> dict[str, Any]: + def build_exception_event(self, event: Event, session_uuid: str) -> dict[str, Any]: """Build a PostHog $exception event for error events.""" properties: dict[str, Any] = { "$exception_source": "backend", - "$session_id": to_uuidv7(event.session_id), + "$session_id": session_uuid, } error = event.error or {} @@ -248,15 +254,17 @@ def build_exception_event(self, event: Event) -> dict[str, Any]: "type": "capture", } - def build_ai_span_event(self, event: Event) -> dict[str, Any]: + def build_ai_span_event( + self, event: Event, session_uuid: str, span_uuid: str + ) -> dict[str, Any]: """Build a PostHog $ai_span event for AI observability views.""" properties: dict[str, Any] = { "$ai_session_id": f"agentcat_{event.session_id}", - "$ai_trace_id": to_uuidv7(event.session_id), - "$ai_span_id": to_uuidv7(event.id), + "$ai_trace_id": session_uuid, + "$ai_span_id": span_uuid, "$ai_span_name": event.resource_name or "unknown_tool", "$ai_is_error": event.is_error or False, - "$session_id": to_uuidv7(event.session_id), + "$session_id": session_uuid, "source": AGENTCAT_SOURCE, } diff --git a/src/agentcat/modules/identify.py b/src/agentcat/modules/identify.py index ccbd40d..71bca17 100644 --- a/src/agentcat/modules/identify.py +++ b/src/agentcat/modules/identify.py @@ -41,9 +41,6 @@ def set(self, session_id: str, identity: UserIdentity) -> None: self._cache.popitem(last=False) self._cache[session_id] = identity - def __contains__(self, session_id: str) -> bool: - return session_id in self._cache - def __len__(self) -> int: return len(self._cache) diff --git a/src/agentcat/modules/overrides/mcp_server.py b/src/agentcat/modules/overrides/mcp_server.py index d516b2f..e8c9456 100644 --- a/src/agentcat/modules/overrides/mcp_server.py +++ b/src/agentcat/modules/overrides/mcp_server.py @@ -219,7 +219,9 @@ async def wrapped_call_tool_handler(request: CallToolRequest) -> ServerResult: try: # Call the original handler result = await original_call_tool_handler(request) - is_error, _error_message = is_mcp_error_response(result) + # Only the error flag is needed here; capture_exception below + # extracts the message itself, so the tuple's message is unused. + is_error, _ = is_mcp_error_response(result) event.is_error = is_error # Full structured capture (message/type/platform; the MCP SDK # already converted the exception, so no frames survive here) diff --git a/src/agentcat/modules/redaction.py b/src/agentcat/modules/redaction.py index 9c23f66..ec49aba 100644 --- a/src/agentcat/modules/redaction.py +++ b/src/agentcat/modules/redaction.py @@ -3,7 +3,7 @@ import asyncio import inspect import threading -from typing import Any, TYPE_CHECKING, Callable, Set +from typing import Any, TYPE_CHECKING, Callable, Coroutine, Set from pydantic import BaseModel @@ -30,34 +30,27 @@ } -def _resolve_redacted_value(result: Any) -> Any: - """Resolve a redaction result that may be an awaitable. +def _run_coroutine(coro: Coroutine[Any, Any, Any]) -> Any: + """Run a coroutine to completion from synchronous code. - Redaction functions may be sync or async (see `RedactionFunction`). The - event queue processes events on worker threads with no running event loop, - so coroutine results are resolved with `asyncio.run`. If a loop is already + The event queue processes events on worker threads with no running event + loop, so the coroutine is resolved with `asyncio.run`. If a loop is already running in this thread (defensive; not the normal publish path), the coroutine is resolved on a short-lived helper thread instead, since `asyncio.run` cannot be called from inside a running loop. """ - if not inspect.isawaitable(result): - return result - - async def _await() -> Any: - return await result - try: asyncio.get_running_loop() except RuntimeError: # Normal path: event-queue worker threads have no event loop. - return asyncio.run(_await()) + return asyncio.run(coro) # A loop is running in this thread — resolve on a fresh thread. outcome: dict[str, Any] = {} def _runner() -> None: try: - outcome["value"] = asyncio.run(_await()) + outcome["value"] = asyncio.run(coro) except BaseException as exc: # noqa: BLE001 — propagated to caller below outcome["error"] = exc @@ -69,6 +62,23 @@ def _runner() -> None: return outcome["value"] +def _resolve_redacted_value(result: Any) -> Any: + """Resolve a redaction result that may be an awaitable. + + Handles the edge case of a *sync* redaction function that returns an + awaitable. Coroutine functions are detected upfront in + `redact_strings_in_object` and resolved on a single event loop for the + whole object instead of one loop per string. + """ + if not inspect.isawaitable(result): + return result + + async def _await() -> Any: + return await result + + return _run_coroutine(_await()) + + def redact_strings_in_object( obj: Any, redact_fn: Callable[[str], str], @@ -80,6 +90,9 @@ def redact_strings_in_object( This ensures that sensitive information is removed from all string fields before events are sent to the analytics service. + Async redaction functions are resolved on a single event loop for the whole + object walk (not one loop per string). + Args: obj: The object to redact strings from redact_fn: The redaction function to apply to each string (sync or async) @@ -89,6 +102,20 @@ def redact_strings_in_object( Returns: A new object with all strings redacted """ + if inspect.iscoroutinefunction(redact_fn): + return _run_coroutine( + _redact_strings_in_object_async(obj, redact_fn, path, is_protected) + ) + return _redact_strings_in_object_sync(obj, redact_fn, path, is_protected) + + +def _redact_strings_in_object_sync( + obj: Any, + redact_fn: Callable[[str], str], + path: str, + is_protected: bool, +) -> Any: + """Walk `obj` applying a sync redaction function (see redact_strings_in_object).""" if obj is None: return obj @@ -102,7 +129,64 @@ def redact_strings_in_object( # Handle arrays/lists if isinstance(obj, list): return [ - redact_strings_in_object(item, redact_fn, f"{path}[{index}]", is_protected) + _redact_strings_in_object_sync( + item, redact_fn, f"{path}[{index}]", is_protected + ) + for index, item in enumerate(obj) + ] + + # Handle dictionaries/objects + if isinstance(obj, dict): + redacted_obj = {} + + for key, value in obj.items(): + # Skip None values + if value is None: + continue + + # Build the path for nested fields + field_path = f"{path}.{key}" if path else key + # Check if this field is protected (only check at top level) + is_field_protected = is_protected or ( + path == "" and key in PROTECTED_FIELDS + ) + redacted_obj[key] = _redact_strings_in_object_sync( + value, redact_fn, field_path, is_field_protected + ) + + return redacted_obj + + # For all other types (numbers, booleans, datetimes, etc.), return as-is + return obj + + +async def _redact_strings_in_object_async( + obj: Any, + redact_fn: Callable[[str], Any], + path: str, + is_protected: bool, +) -> Any: + """Async twin of `_redact_strings_in_object_sync` for coroutine redact fns. + + Driven once per redact call by `_run_coroutine`, so every redacted string + shares one event loop. + """ + if obj is None: + return obj + + # Handle strings + if isinstance(obj, str): + # Don't redact if this field or any parent field is protected + if is_protected: + return obj + return await redact_fn(obj) + + # Handle arrays/lists + if isinstance(obj, list): + return [ + await _redact_strings_in_object_async( + item, redact_fn, f"{path}[{index}]", is_protected + ) for index, item in enumerate(obj) ] @@ -121,7 +205,7 @@ def redact_strings_in_object( is_field_protected = is_protected or ( path == "" and key in PROTECTED_FIELDS ) - redacted_obj[key] = redact_strings_in_object( + redacted_obj[key] = await _redact_strings_in_object_async( value, redact_fn, field_path, is_field_protected ) @@ -149,9 +233,10 @@ def redact_event(event: "UnredactedEvent", redact_fn: Callable[[str], str]) -> " """ if isinstance(event, BaseModel): # Dump the model to a plain dict (dropping the redaction_fn callable), - # redact recursively, then rebuild the same model class. + # redact recursively, then rebuild the same model class. The dump came + # from an already-validated model, so skip re-validation on rebuild. event_dict = event.model_dump(exclude_none=True, exclude={"redaction_fn"}) redacted_dict = redact_strings_in_object(event_dict, redact_fn, "", False) - return type(event)(**redacted_dict) + return type(event).model_construct(**redacted_dict) return redact_strings_in_object(event, redact_fn, "", False) diff --git a/src/agentcat/types.py b/src/agentcat/types.py index f51ed2b..585887a 100644 --- a/src/agentcat/types.py +++ b/src/agentcat/types.py @@ -9,8 +9,8 @@ from pydantic import BaseModel, field_validator from agentcat.modules.constants import ( - AGENTCAT_CUSTOM_EVENT_TYPE, DEFAULT_CONTEXT_DESCRIPTION, + SDK_EVENT_TYPES, ) # Type alias for identify function. @@ -71,7 +71,7 @@ def event_type_validate_enum( ) -> Optional[str]: """Relax the generated enum check to admit SDK-defined event types (e.g. "agentcat:custom") the generated client doesn't know about.""" - if value == AGENTCAT_CUSTOM_EVENT_TYPE: + if value in SDK_EVENT_TYPES: return value validated: Optional[str] = PublishEventRequest.event_type_validate_enum(value) return validated