From dbe5351c3e7ab0f498313c35fc1040f2d0823593 Mon Sep 17 00:00:00 2001 From: Naseem AlNaji Date: Thu, 16 Jul 2026 13:37:35 +0200 Subject: [PATCH 1/2] feat(redaction): add event-level redact_event hook Add a redact_event option to AgentCatOptions that receives the full Event object before publishing, complementing the existing string-level redact_sensitive_information hook. Consumers can inspect event metadata (resource_name, event_type, parameters, response, etc.) to make context-aware redaction decisions, return a modified event, or return None to drop the event entirely. - New EventRedactionFunction type; Event and EventRedactionFunction are now exported from the package root - Hook runs before string-level redaction, so it sees raw values; string redaction, sanitization, and truncation still run on its output - System-managed fields (id, session_id, project_id, event_type, timestamp) are restored if the hook alters them - Hook errors drop the event and log, matching existing redaction semantics - Attached centrally in publish_event, covering every event producer --- README.md | 21 +++ src/agentcat/__init__.py | 6 +- src/agentcat/modules/event_queue.py | 20 ++- src/agentcat/modules/redaction.py | 61 ++++++++- src/agentcat/types.py | 18 +++ tests/e2e/official/test_redact_event_http.py | 124 +++++++++++++++++ tests/test_event_queue.py | 129 ++++++++++++++++++ tests/test_redaction.py | 134 +++++++++++++++++++ 8 files changed, 510 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/official/test_redact_event_http.py diff --git a/README.md b/README.md index cab7954..3b87b30 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,27 @@ def redact_sync(text): agentcat.track(server, "proj_0000000", AgentCatOptions(redact_sensitive_information=redact_sync)) ``` +For redaction decisions that need more context than a single string — such as which tool was called or what type of event is being published — use the event-level `redact_event` hook. It receives the full event object and returns a modified event, or `None` to drop the event entirely. It may be sync or async, and can be combined with `redact_sensitive_information`. + +```python +from agentcat import AgentCatOptions + +def redact_event(event): + # Drop events from tools that handle secrets entirely + if event.resource_name == "get_credentials": + return None + # Strip response payloads from a specific tool + if event.resource_name == "export_report": + modified = event.model_copy() + modified.response = None + return modified + return event + +agentcat.track(server, "proj_0000000", AgentCatOptions(redact_event=redact_event)) +``` + +When both hooks are configured, `redact_event` runs first and sees the raw, unredacted values; `redact_sensitive_information` then runs on its output as a final string-level scrub. The system-managed fields `id`, `session_id`, `project_id`, `event_type`, and `timestamp` cannot be changed by the hook, and if the hook raises, the event is dropped. + ### Forwarding data to existing observability platforms AgentCat seamlessly integrates with your existing observability stack, providing automatic logging and tracing without the tedious setup typically required. Export telemetry data to multiple platforms simultaneously: diff --git a/src/agentcat/__init__.py b/src/agentcat/__init__.py index 933c5f3..20da5f1 100644 --- a/src/agentcat/__init__.py +++ b/src/agentcat/__init__.py @@ -22,7 +22,9 @@ from .modules.internal import set_server_tracking_data from .modules.logging import set_debug_mode, write_to_log from .types import ( + Event, EventPropertiesFunction, + EventRedactionFunction, EventTagsFunction, IdentifyFunction, AgentCatData, @@ -256,8 +258,10 @@ def _apply_server_tracking( # Types for identify functionality "UserIdentity", "IdentifyFunction", - # Type for redaction functionality + # Types for redaction functionality "RedactionFunction", + "EventRedactionFunction", + "Event", # Types for event metadata callbacks "EventTagsFunction", "EventPropertiesFunction", diff --git a/src/agentcat/modules/event_queue.py b/src/agentcat/modules/event_queue.py index 473ab2e..841a099 100644 --- a/src/agentcat/modules/event_queue.py +++ b/src/agentcat/modules/event_queue.py @@ -21,7 +21,7 @@ from .compatibility import get_mcp_compatible_error_message from .internal import get_server_tracking_data from .logging import write_to_log -from .redaction import redact_event +from .redaction import apply_event_redaction, redact_event from .sanitization import sanitize_event from .truncation import truncate_event from .session import get_session_info, set_last_activity @@ -105,6 +105,23 @@ def _worker(self) -> None: def _process_event(self, event: UnredactedEvent) -> None: """Process a single event.""" + if event and event.redact_event_fn: + # Run the customer's event-level redaction hook first, on raw values + event_redact_fn = event.redact_event_fn + event.redact_event_fn = None + try: + if not apply_event_redaction(event, event_redact_fn): + write_to_log( + f"Event {event.id or 'unknown'} dropped by redact_event hook" + ) + return + except Exception as error: + write_to_log( + f"Failed to redact event (event-level hook), dropping event " + f"{event.id or 'unknown'}: {error}" + ) + return + if event and event.redaction_fn: # Redact sensitive information if a redaction function is provided try: @@ -320,6 +337,7 @@ def publish_event(server: Any, event: UnredactedEvent) -> None: full_event = UnredactedEvent( **merged_data, redaction_fn=data.options.redact_sensitive_information, + redact_event_fn=data.options.redact_event, ) set_last_activity(server) diff --git a/src/agentcat/modules/redaction.py b/src/agentcat/modules/redaction.py index 58b0dca..4646581 100644 --- a/src/agentcat/modules/redaction.py +++ b/src/agentcat/modules/redaction.py @@ -1,9 +1,23 @@ """PII redaction for AgentCat logs.""" +import asyncio +import inspect from typing import Any, TYPE_CHECKING, Callable, Set if TYPE_CHECKING: - from agentcat.types import Event, UnredactedEvent + from agentcat.types import Event, EventRedactionFunction, UnredactedEvent + + +# System-managed fields restored after the event-level redaction hook runs. +# These are required for ingestion and session/project attribution, so +# consumer changes to them are ignored. +RESTORED_FIELDS = ( + "id", + "session_id", + "project_id", + "event_type", + "timestamp", +) # Set of field names that should be protected from redaction. @@ -101,3 +115,48 @@ def redact_event(event: "UnredactedEvent", redact_fn: Callable[[str], str]) -> " A new event object with all strings redacted """ return redact_strings_in_object(event, redact_fn, "", False) + + +def apply_event_redaction( + event: "UnredactedEvent", event_redact_fn: "EventRedactionFunction" +) -> bool: + """ + Applies the customer's event-level redaction hook to an event, in place. + The hook receives the full event (without internal function fields) and may + return a modified event, or None to drop the event entirely. + + The event object is rewritten rather than replaced: the queue pipeline + processes the same object across steps, and copying every field ensures + values the hook cleared stay cleared. System-managed fields are restored + from the original, and the string-level redaction_fn is preserved so it + still runs afterwards. + + Args: + event: The event to run the hook on; mutated with the hook's result + event_redact_fn: The customer's event-level redaction hook + + Returns: + True if the event was kept, False if the hook dropped it + """ + from agentcat.types import Event + + # Hide internal carrier fields from the hook + hook_input = Event(**{name: getattr(event, name) for name in Event.model_fields}) + + result = event_redact_fn(hook_input) + if inspect.iscoroutine(result): + # The queue worker runs in a plain thread with no event loop + result = asyncio.run(result) + + if result is None: + return False + + # Copy every field from the result so values the hook cleared stay + # cleared, then restore the system-managed fields from the original + snapshot = {name: getattr(event, name) for name in RESTORED_FIELDS} + for name in Event.model_fields: + setattr(event, name, getattr(result, name, None)) + for name, value in snapshot.items(): + setattr(event, name, value) + + return True diff --git a/src/agentcat/types.py b/src/agentcat/types.py index 41f15a9..3c49c1e 100644 --- a/src/agentcat/types.py +++ b/src/agentcat/types.py @@ -14,6 +14,13 @@ IdentifyFunction = Callable[[dict[str, Any], Any], Optional["UserIdentity"]] # Type alias for redaction function RedactionFunction = Callable[[str], str | Awaitable[str]] +# Type alias for the event-level redaction hook — receives the full Event and +# returns a modified Event, or None to drop the event entirely. +# Accepts sync or async callables (mirrors RedactionFunction). +EventRedactionFunction = Callable[ + ["Event"], + Union["Event", None, Awaitable[Union["Event", None]]], +] # Type alias for event_tags callback — returns str:str map attached to every auto-captured event. # Accepts sync or async callables (mirrors RedactionFunction). EventTagsFunction = Callable[ @@ -117,6 +124,7 @@ class EventType(str, Enum): class UnredactedEvent(Event): redaction_fn: RedactionFunction | None = None + redact_event_fn: EventRedactionFunction | None = None # Whole-event redaction hook @dataclass @@ -176,6 +184,16 @@ class AgentCatOptions: custom_context_description: str = DEFAULT_CONTEXT_DESCRIPTION identify: IdentifyFunction | None = None redact_sensitive_information: RedactionFunction | None = None + # Event-level redaction hook invoked with the full event (inspect + # resource_name, event_type, parameters, response, etc.) before it is + # published. Return a modified event, or None to drop the event entirely. + # May be sync or async. Runs before redact_sensitive_information, so it + # sees raw, unredacted values; the string-level hook, sanitization, and + # truncation still run on its output. The system-managed fields id, + # session_id, project_id, event_type, and timestamp cannot be changed + # (id may be None at hook time). If the hook raises, the event is dropped + # and the error is logged to ~/agentcat.log. + redact_event: EventRedactionFunction | None = None exporters: dict[str, ExporterConfig] | None = None debug_mode: bool = False api_base_url: str | None = None diff --git a/tests/e2e/official/test_redact_event_http.py b/tests/e2e/official/test_redact_event_http.py new file mode 100644 index 0000000..4b9b118 --- /dev/null +++ b/tests/e2e/official/test_redact_event_http.py @@ -0,0 +1,124 @@ +"""Event-level redaction hook (redact_event) over real-wire payloads. + +Unlike the string-level hook (see test_redaction_http.py, xfail-tracked), the +event-level hook receives the Pydantic event object directly, so it works on +the live publish path today. +""" + +from __future__ import annotations + +import time + +import pytest +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + +from agentcat.modules.internal import get_server_tracking_data + + +pytestmark = pytest.mark.e2e + + +def _set_redact_event(server, fn) -> None: + data = get_server_tracking_data(server) + assert data is not None + data.options.redact_event = fn + + +@pytest.mark.asyncio +async def test_hook_inspects_metadata_and_rewrites_event_fields( + official_http_server, capture_queue +): + url, server = official_http_server + + seen_resource_names: list[str | None] = [] + + def hook(event): + seen_resource_names.append(event.resource_name) + if event.resource_name == "add_todo": + modified = event.model_copy() + modified.parameters = {"replaced": True} + return modified + return event + + _set_redact_event(server, hook) + try: + async with streamablehttp_client(url) as (read, write, _): + async with ClientSession(read, write) as client: + await client.initialize() + await client.call_tool( + "add_todo", {"text": "buy milk", "context": "rewrite"} + ) + + time.sleep(0.5) + call_events = [e for e in capture_queue if e.event_type == "mcp:tools/call"] + assert call_events + assert "add_todo" in seen_resource_names + assert call_events[0].parameters == {"replaced": True} + + # System-managed fields are intact + assert call_events[0].session_id + assert call_events[0].event_type == "mcp:tools/call" + finally: + _set_redact_event(server, None) + + +@pytest.mark.asyncio +async def test_hook_returning_none_drops_the_event( + official_http_server, capture_queue +): + url, server = official_http_server + + def hook(event): + if event.event_type == "mcp:tools/call": + return None + return event + + _set_redact_event(server, hook) + try: + async with streamablehttp_client(url) as (read, write, _): + async with ClientSession(read, write) as client: + await client.initialize() + await client.call_tool("add_todo", {"text": "drop", "context": "drop"}) + await client.list_tools() + + time.sleep(1.0) + call_events = [e for e in capture_queue if e.event_type == "mcp:tools/call"] + assert not call_events, ( + f"redact_event returning None must drop the event, got {len(call_events)}" + ) + # Other event types still publish + assert capture_queue + finally: + _set_redact_event(server, None) + + +@pytest.mark.asyncio +async def test_hook_raising_drops_only_the_affected_event( + official_http_server, capture_queue +): + url, server = official_http_server + + def hook(event): + if event.event_type == "mcp:tools/call": + raise RuntimeError("hook exploded") + return event + + _set_redact_event(server, hook) + try: + async with streamablehttp_client(url) as (read, write, _): + async with ClientSession(read, write) as client: + await client.initialize() + await client.call_tool("add_todo", {"text": "boom", "context": "boom"}) + # Server still serves requests after the hook error + tools = await client.list_tools() + assert tools.tools + + time.sleep(1.0) + call_events = [e for e in capture_queue if e.event_type == "mcp:tools/call"] + assert not call_events, ( + f"redact_event raising must drop the event, got {len(call_events)}" + ) + assert capture_queue # other events still arrived + finally: + _set_redact_event(server, None) diff --git a/tests/test_event_queue.py b/tests/test_event_queue.py index 1e5b987..83b644d 100644 --- a/tests/test_event_queue.py +++ b/tests/test_event_queue.py @@ -719,6 +719,135 @@ def test_publish_event_with_redaction_function( added_event = mock_eq.add.call_args[0][0] assert added_event.redaction_fn == mock_redaction_fn + @patch("agentcat.modules.event_queue.get_server_tracking_data") + @patch("agentcat.modules.event_queue.get_session_info") + @patch("agentcat.modules.event_queue.set_last_activity") + @patch("agentcat.modules.event_queue.event_queue") + def test_publish_event_attaches_redact_event_hook_from_options( + self, mock_eq, mock_set_activity, mock_session, mock_tracking + ): + """Test publishing event includes the event-level hook from options.""" + mock_server = MagicMock() + mock_hook = MagicMock() + mock_data = AgentCatData( + project_id="project-123", + session_id="session-123", + session_info=SessionInfo(), + last_activity=datetime.now(timezone.utc), + options=AgentCatOptions(redact_event=mock_hook), + ) + mock_tracking.return_value = mock_data + mock_session.return_value = SessionInfo() + + event = UnredactedEvent( + event_type="mcp:tools/call", + session_id="session-123", + timestamp=datetime.now(timezone.utc), + ) + + publish_event(mock_server, event) + + added_event = mock_eq.add.call_args[0][0] + assert added_event.redact_event_fn == mock_hook + + +class TestEventLevelRedactionHook: + """Test the redact_event hook in the queue pipeline.""" + + @staticmethod + def _event(**overrides) -> UnredactedEvent: + defaults = dict( + id="test-id", + event_type="mcp:tools/call", + project_id="project-123", + session_id="session-123", + timestamp=datetime.now(timezone.utc), + resource_name="add_todo", + parameters={"text": "raw sensitive value"}, + user_intent="raw secret", + ) + defaults.update(overrides) + return UnredactedEvent(**defaults) + + def test_hook_result_is_published(self): + """The hook sees raw metadata and its rewrite is what gets sent.""" + seen = {} + + def hook(event): + seen["resource_name"] = event.resource_name + seen["parameters"] = event.parameters + modified = event.model_copy() + modified.parameters = {"text": "[EVENT-REDACTED]"} + return modified + + eq = EventQueue() + event = self._event(redact_event_fn=hook) + + with patch.object(eq, "_send_event") as mock_send: + eq._process_event(event) + + assert seen["resource_name"] == "add_todo" + assert seen["parameters"] == {"text": "raw sensitive value"} + sent_event = mock_send.call_args[0][0] + assert sent_event.parameters == {"text": "[EVENT-REDACTED]"} + assert sent_event.redact_event_fn is None + + @patch("agentcat.modules.event_queue.write_to_log") + def test_hook_returning_none_drops_the_event(self, mock_log): + eq = EventQueue() + event = self._event(redact_event_fn=lambda event: None) + + with patch.object(eq, "_send_event") as mock_send: + eq._process_event(event) + + mock_send.assert_not_called() + log_messages = [str(c) for c in mock_log.call_args_list] + assert any("dropped by redact_event hook" in m for m in log_messages) + + @patch("agentcat.modules.event_queue.write_to_log") + def test_hook_raising_drops_the_event(self, mock_log): + def hook(event): + raise ValueError("hook exploded") + + eq = EventQueue() + event = self._event(redact_event_fn=hook) + + with patch.object(eq, "_send_event") as mock_send: + eq._process_event(event) + + mock_send.assert_not_called() + log_messages = [str(c) for c in mock_log.call_args_list] + assert any( + "Failed to redact event (event-level hook)" in m + for m in log_messages + ) + + @patch("agentcat.modules.event_queue.redact_event") + def test_hook_runs_before_string_redaction(self, mock_redact): + """The event hook sees raw values; string redaction runs on its output.""" + raw_seen = {} + + def hook(event): + raw_seen["user_intent"] = event.user_intent + modified = event.model_copy() + modified.user_intent = "hook output" + return modified + + string_fn = MagicMock() + mock_redact.side_effect = lambda event, fn: event + + eq = EventQueue() + event = self._event(redact_event_fn=hook, redaction_fn=string_fn) + + with patch.object(eq, "_send_event"): + eq._process_event(event) + + assert raw_seen["user_intent"] == "raw secret" + mock_redact.assert_called_once() + redacted_arg = mock_redact.call_args[0][0] + assert redacted_arg.user_intent == "hook output" + assert mock_redact.call_args[0][1] is string_fn + @patch("agentcat.modules.event_queue.signal.signal") @patch("agentcat.modules.event_queue.atexit.register") diff --git a/tests/test_redaction.py b/tests/test_redaction.py index 533c649..fa5b6c1 100644 --- a/tests/test_redaction.py +++ b/tests/test_redaction.py @@ -1,12 +1,15 @@ """Unit tests for the redaction module.""" import pytest +from datetime import datetime, timezone from typing import Any, Dict from agentcat.modules.redaction import ( redact_strings_in_object, redact_event, + apply_event_redaction, PROTECTED_FIELDS, ) +from agentcat.types import Event, UnredactedEvent class TestRedactStringsInObject: @@ -367,3 +370,134 @@ 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 TestApplyEventRedaction: + """Test suite for the event-level redaction hook helper.""" + + @staticmethod + def _base_event() -> UnredactedEvent: + return UnredactedEvent( + id="evt_original", + session_id="ses_123", + project_id="proj_789", + event_type="mcp:tools/call", + timestamp=datetime(2024, 1, 1, tzinfo=timezone.utc), + resource_name="add_todo", + parameters={"text": "sensitive text"}, + response={"content": "sensitive response"}, + user_intent="sensitive intent", + ) + + def test_applies_sync_hook_in_place(self): + def hook(event: Event) -> Event: + modified = event.model_copy() + modified.parameters = {"text": "[SCRUBBED]"} + return modified + + event = self._base_event() + kept = apply_event_redaction(event, hook) + + assert kept is True + assert event.parameters == {"text": "[SCRUBBED]"} + assert event.response == {"content": "sensitive response"} + + def test_applies_async_hook_in_place(self): + async def hook(event: Event) -> Event: + modified = event.model_copy() + modified.user_intent = "[SCRUBBED]" + return modified + + event = self._base_event() + kept = apply_event_redaction(event, hook) + + assert kept is True + assert event.user_intent == "[SCRUBBED]" + + def test_reports_drop_and_leaves_event_untouched_on_none(self): + def hook(event: Event) -> None: + return None + + event = self._base_event() + kept = apply_event_redaction(event, hook) + + assert kept is False + assert event.parameters == {"text": "sensitive text"} + assert event.user_intent == "sensitive intent" + + def test_honors_cleared_fields_instead_of_resurrecting_them(self): + def hook(event: Event) -> Event: + modified = event.model_copy() + modified.response = None + modified.user_intent = None + return modified + + event = self._base_event() + kept = apply_event_redaction(event, hook) + + assert kept is True + assert event.response is None + assert event.user_intent is None + assert event.parameters == {"text": "sensitive text"} + + def test_restores_system_managed_fields(self): + def hook(event: Event) -> Event: + modified = event.model_copy() + modified.id = "evt_forged" + modified.session_id = "ses_forged" + modified.project_id = "proj_forged" + # event_type is enum-validated by pydantic, so forge with a + # different-but-valid value + modified.event_type = "mcp:ping" + modified.timestamp = None + return modified + + event = self._base_event() + apply_event_redaction(event, hook) + + assert event.id == "evt_original" + assert event.session_id == "ses_123" + assert event.project_id == "proj_789" + assert event.event_type == "mcp:tools/call" + assert event.timestamp == datetime(2024, 1, 1, tzinfo=timezone.utc) + + def test_hides_internal_function_fields_and_preserves_redaction_fn(self): + saw_event: dict[str, Any] = {} + + def hook(event: Event) -> Event: + saw_event["value"] = event + return event + + def string_redact_fn(text: str) -> str: + return text + + event = self._base_event() + event.redaction_fn = string_redact_fn + event.redact_event_fn = hook + + kept = apply_event_redaction(event, hook) + + assert kept is True + assert not hasattr(saw_event["value"], "redaction_fn") + assert not hasattr(saw_event["value"], "redact_event_fn") + # The string-level hook must survive so it still runs afterwards + assert event.redaction_fn is string_redact_fn + + def test_propagates_hook_errors_to_the_caller(self): + def hook(event: Event) -> Event: + raise ValueError("Hook failure") + + with pytest.raises(ValueError, match="Hook failure"): + apply_event_redaction(self._base_event(), hook) + + def test_keeps_the_same_object_identity_for_pipeline_observers(self): + def hook(event: Event) -> Event: + modified = event.model_copy() + modified.parameters = {"text": "[SCRUBBED]"} + return modified + + event = self._base_event() + observer = event # e.g. a capture holding the reference before processing + apply_event_redaction(event, hook) + + assert observer.parameters == {"text": "[SCRUBBED]"} From 90402d5c6da099b608056caf3c2b822d79fc4bf6 Mon Sep 17 00:00:00 2001 From: Naseem AlNaji Date: Thu, 16 Jul 2026 17:45:18 +0200 Subject: [PATCH 2/2] chore(release): bump version to 1.0.3 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 48d8aad..87d9279 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "agentcat" -version = "1.0.2" +version = "1.0.3" description = "Analytics tool for MCP (Model Context Protocol) servers and AI agents - tracks tool usage patterns and provides insights" authors = [ { name = "AgentCat, Inc.", email = "support@agentcat.com" },