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/__init__.py b/src/agentcat/__init__.py index 933c5f3..c056017 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,114 @@ 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 + + # 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") + + 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, + 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.tags: + event.tags = validate_tags(event_data.tags) + if 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 +374,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..69e97a2 100644 --- a/src/agentcat/modules/constants.py +++ b/src/agentcat/modules/constants.py @@ -4,6 +4,11 @@ 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 +# 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/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..2619dc6 100644 --- a/src/agentcat/modules/exceptions.py +++ b/src/agentcat/modules/exceptions.py @@ -12,6 +12,9 @@ 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 ) @@ -26,12 +29,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 +67,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 +305,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 +339,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 +427,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..e12aaad --- /dev/null +++ b/src/agentcat/modules/exporters/posthog.py @@ -0,0 +1,307 @@ +"""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]] = [] + + # 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, 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, 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, session_uuid, to_uuidv7(event.id)) + ) + + 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, session_uuid: str) -> dict[str, Any]: + """Build the regular PostHog capture event.""" + properties: dict[str, Any] = { + "$session_id": session_uuid, + "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, session_uuid: str) -> dict[str, Any]: + """Build a PostHog $exception event for error events.""" + properties: dict[str, Any] = { + "$exception_source": "backend", + "$session_id": session_uuid, + } + + 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, 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": 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": session_uuid, + "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..71bca17 100644 --- a/src/agentcat/modules/identify.py +++ b/src/agentcat/modules/identify.py @@ -1,15 +1,95 @@ +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. + + 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): + 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 __len__(self) -> int: + return len(self._cache) + + +# Global identity cache shared across all server instances. +# This preserves user_data merge continuity 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 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 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. """ data = get_server_tracking_data(server) @@ -19,6 +99,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 +108,38 @@ def identify_session(server, request: any, context: any) -> UserIdentity | None: ) return None + session_id = data.session_id + + 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=data.session_id, + session_id=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 {}, + 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 identify_result + 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..e8c9456 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,27 @@ 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 - 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 +266,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 +307,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..ec49aba 100644 --- a/src/agentcat/modules/redaction.py +++ b/src/agentcat/modules/redaction.py @@ -1,6 +1,11 @@ """PII redaction for AgentCat logs.""" -from typing import Any, TYPE_CHECKING, Callable, Set +import asyncio +import inspect +import threading +from typing import Any, TYPE_CHECKING, Callable, Coroutine, Set + +from pydantic import BaseModel if TYPE_CHECKING: from agentcat.types import Event, UnredactedEvent @@ -25,6 +30,55 @@ } +def _run_coroutine(coro: Coroutine[Any, Any, Any]) -> Any: + """Run a coroutine to completion from synchronous code. + + 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. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + # Normal path: event-queue worker threads have no event loop. + 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(coro) + 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 _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], @@ -36,15 +90,87 @@ 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 + 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 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 + + # Handle strings + if isinstance(obj, str): + # Don't redact if this field or any parent field is protected + if is_protected: + return obj + return _resolve_redacted_value(redact_fn(obj)) + + # Handle arrays/lists + if isinstance(obj, list): + return [ + _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 @@ -53,12 +179,14 @@ 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 await redact_fn(obj) # Handle arrays/lists if isinstance(obj, list): return [ - redact_strings_in_object(item, redact_fn, f"{path}[{index}]", is_protected) + await _redact_strings_in_object_async( + item, redact_fn, f"{path}[{index}]", is_protected + ) for index, item in enumerate(obj) ] @@ -77,13 +205,13 @@ 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 ) 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 +221,22 @@ 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. 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).model_construct(**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..585887a 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 ( + DEFAULT_CONTEXT_DESCRIPTION, + SDK_EVENT_TYPES, +) + +# 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 in SDK_EVENT_TYPES: + 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..4870b07 --- /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): + """Identity merging handles 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) # merges with the cached identity + 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..ab91337 --- /dev/null +++ b/tests/test_identify.py @@ -0,0 +1,256 @@ +"""Tests for identify: async hooks, per-request publishing, and identity merging. + +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 +from unittest.mock import MagicMock, patch + +from agentcat.modules.identify import ( + IdentityCache, + 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: + """Publish-every-time 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_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"} + ) + ) + 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 + 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): + """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), + 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) == 3 + assert events[0].identify_actor_given_id == "alice" + 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. Each call publishes an event + carrying the merged identity.""" + 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_unchanged_identity_republishes(self, mock_event_queue): + """A subset of previously-seen user_data merges into an identical + identity → still republished, and the second event carries the FULL + merged identity (merge continuity + publish-every-time).""" + 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"} + 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): + """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): + """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 + ) + 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 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"}) + 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..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): @@ -61,31 +70,68 @@ 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): - """In stateless mode, identify should run on every call (no early-return guard).""" + 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) + 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 )) 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 + 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") - 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 +143,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 +177,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