From c67e3f20b17e5df4d89f4115dd4d25445c52f330 Mon Sep 17 00:00:00 2001 From: Naseem Alnaji Date: Fri, 7 Aug 2026 16:05:58 +0000 Subject: [PATCH] feat(redaction): add event-level redact_event hook, protect server/client identity fields Add an event-level redact_event hook (AgentCatOptions.redact_event), matching the redactEvent/RedactEvent hooks already available in the TypeScript and Go SDKs. It runs before redact_sensitive_information, receives the full event, and may return a modified event or None to drop it entirely. id, session_id, project_id, event_type, and timestamp are restored from the original event afterward regardless of what the hook returns. Also add server_name, server_version, client_name, and client_version to PROTECTED_FIELDS: these are system-reported MCP metadata, not user-supplied content, and should always reach the dashboard intact rather than passing through redact_sensitive_information. --- README.md | 17 ++++ src/agentcat/__init__.py | 4 +- src/agentcat/modules/event_queue.py | 24 ++++- src/agentcat/modules/redaction.py | 80 ++++++++++++++++- src/agentcat/types.py | 16 ++++ tests/test_event_queue.py | 135 ++++++++++++++++++++++++++++ tests/test_redaction.py | 131 ++++++++++++++++++++++++++- 7 files changed, 402 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 93ffc5c..961c34b 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,23 @@ async def redact(text: str) -> str: agentcat.track(server, "proj_0000000", AgentCatOptions(redact_sensitive_information=redact)) ``` +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 and returns a modified event, or `None` to drop the event entirely. It may be sync or async, runs before `redact_sensitive_information` (so it sees raw values), and can be combined with it. + +```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": + event.response = None + return event + +agentcat.track(server, "proj_0000000", AgentCatOptions(redact_event=redact_event)) +``` + ### Vendor Support 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 97ce075..641ed51 100644 --- a/src/agentcat/__init__.py +++ b/src/agentcat/__init__.py @@ -25,6 +25,7 @@ EventPropertiesFunction, EventTagsFunction, IdentifyFunction, + RedactEventFunction, RedactionFunction, ResolveSessionIdFunction, UnredactedEvent, @@ -449,8 +450,9 @@ def _wire_payload(field: str, value: Any) -> dict[str, Any] | None: # Types for identify functionality "UserIdentity", "IdentifyFunction", - # Type for redaction functionality + # Types for redaction functionality "RedactionFunction", + "RedactEventFunction", # Types for event metadata callbacks "EventTagsFunction", "EventPropertiesFunction", diff --git a/src/agentcat/modules/event_queue.py b/src/agentcat/modules/event_queue.py index f2c4e79..0fdb5e2 100644 --- a/src/agentcat/modules/event_queue.py +++ b/src/agentcat/modules/event_queue.py @@ -26,7 +26,7 @@ from ..utils import generate_prefixed_ksuid, get_agentcat_version 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 @@ -155,6 +155,27 @@ def _worker(self) -> None: def _process_event(self, event: UnredactedEvent) -> None: """Process a single event.""" + if event and event.event_redaction_fn: + # Event-level redaction hook runs first, on raw values, and may + # drop the event entirely. + try: + if not event.id: + event.id = generate_prefixed_ksuid(EVENT_ID_PREFIX) + redacted_event = apply_event_redaction(event, event.event_redaction_fn) + if redacted_event is None: + write_to_log(f"Event {event.id} dropped by redact_event hook") + return + event = redacted_event + event.event_redaction_fn = None # Clear to avoid reprocessing + except (Exception, SystemExit) as error: + # SystemExit included: a customer hook calling sys.exit() must + # not kill the worker thread. + write_to_log( + f"WARNING: Dropping event {event.id or 'unknown'} due to " + f"event redaction failure: {error}" + ) + return # Skip this event if event redaction fails + if event and event.redaction_fn: # Redact sensitive information if a redaction function is provided try: @@ -394,6 +415,7 @@ def publish_event(server: Any, event: UnredactedEvent) -> None: full_event = UnredactedEvent( **stamped, redaction_fn=data.options.redact_sensitive_information, + event_redaction_fn=data.options.redact_event, ) event_queue.add(full_event) diff --git a/src/agentcat/modules/redaction.py b/src/agentcat/modules/redaction.py index 44d362c..9537042 100644 --- a/src/agentcat/modules/redaction.py +++ b/src/agentcat/modules/redaction.py @@ -16,6 +16,10 @@ "id", "project_id", "server", + "server_name", + "server_version", + "client_name", + "client_version", "identify_actor_given_id", "identify_actor_name", "identify_data", @@ -146,8 +150,82 @@ def redact_event(event: "UnredactedEvent", redact_fn: Callable[[str], str]) -> " if not callable(dump): plain: Event = redact_strings_in_object(event, redact, "", False) return plain - dumped = dump(exclude={"redaction_fn"}, warnings=False) + dumped = dump(exclude={"redaction_fn", "event_redaction_fn"}, warnings=False) redacted: Event = event.model_copy( update=redact_strings_in_object(dumped, redact, "", False) ) return redacted + + +def _sync_event_redactor( + redact_event_fn: Callable[["Event"], Any], +) -> Callable[["Event"], Any]: + """A synchronous view of the customer's event-level redaction hook. + + Same rationale as `_sync_redactor`: the publish worker is a thread with no + event loop of its own, so an async hook's result has to be driven to + completion here rather than assigned straight into the pipeline. + """ + + def run(event: "Event") -> Any: + return drive_hook_result(redact_event_fn(event), "redact_event") + + return run + + +# Fields restored from the original event after the event-level redaction +# hook runs, regardless of what the hook returns. These are system-managed — +# not consumer-settable — mirroring RESTORED_FIELDS in the TypeScript SDK and +# the equivalent snapshot/restore around ApplyEventRedaction in the Go SDK. +RESTORED_FIELDS: Set[str] = { + "id", + "session_id", + "project_id", + "event_type", + "timestamp", +} + + +def apply_event_redaction( + event: "UnredactedEvent", + redact_event_fn: Callable[["Event"], Any], +) -> "UnredactedEvent | None": + """ + Applies the customer's event-level redaction hook to an event. + + Runs before `redact_event` (the string-level hook), so the hook receives + raw, unredacted values. The hook may return a modified event, or None to + drop the event entirely. System-managed fields (RESTORED_FIELDS) are + force-restored from the original event afterward, so a hook cannot forge + or erase what AgentCat itself assigned. + + The hook is handed a plain `Event` built from the dump, excluding + `redaction_fn`/`event_redaction_fn` — they are machinery, not event data, + and the customer's hook must not be handed its own function objects. + + Args: + event: The event to run the hook on + redact_event_fn: The customer's event-level redaction hook + + Returns: + The (possibly modified) event to keep processing, or None if the hook + dropped it. + """ + from agentcat.types import Event as EventModel + + run = _sync_event_redactor(redact_event_fn) + dumped = event.model_dump( + exclude={"redaction_fn", "event_redaction_fn"}, warnings=False + ) + hook_input = EventModel(**dumped) + + result = run(hook_input) + + if result is None: + return None + + restored = {field: getattr(event, field) for field in RESTORED_FIELDS} + updated = result.model_dump(warnings=False) + updated.update(restored) + redacted: UnredactedEvent = event.model_copy(update=updated) + return redacted diff --git a/src/agentcat/types.py b/src/agentcat/types.py index fbeaffd..c52b2b0 100644 --- a/src/agentcat/types.py +++ b/src/agentcat/types.py @@ -38,6 +38,13 @@ ] # Type alias for redaction function RedactionFunction = Callable[[str], str | Awaitable[str]] +# Type alias for the event-level redaction hook. Receives the full event +# (raw, unredacted values — it runs before RedactionFunction) and returns a +# modified event, or None to drop the event entirely. Accepts sync or async +# callables (mirrors RedactionFunction). +RedactEventFunction = Callable[ + ["Event"], Optional["Event"] | Awaitable[Optional["Event"]] +] # 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[ @@ -149,6 +156,7 @@ class EventType(str, Enum): class UnredactedEvent(Event): redaction_fn: RedactionFunction | None = None + event_redaction_fn: RedactEventFunction | None = None # Telemetry Exporter Configuration Types @@ -202,6 +210,14 @@ class AgentCatOptions: # publishes anonymously rather than failing the tool call. identify: IdentifyFunction | None = None redact_sensitive_information: RedactionFunction | None = None + # Event-level redaction hook, invoked with the full event (raw, unredacted + # values) before redact_sensitive_information runs. May return a modified + # event, or None to drop the event entirely. May be sync or async. The + # system-managed fields id, session_id, project_id, event_type, and + # timestamp cannot be changed by this hook — they are restored from the + # original event afterward. If the hook raises (or times out), the event + # is dropped rather than published unredacted. + redact_event: RedactEventFunction | None = None exporters: dict[str, ExporterConfig] | None = None # Debug logging to ~/agentcat.log. Tri-state: None (the default) defers to # the AGENTCAT_DEBUG_MODE env var read at import; explicit True/False wins. diff --git a/tests/test_event_queue.py b/tests/test_event_queue.py index fe3b145..ef5a9b6 100644 --- a/tests/test_event_queue.py +++ b/tests/test_event_queue.py @@ -274,6 +274,120 @@ def test_process_event_redaction_failure(self, mock_log, mock_redact): assert "test-id" in log_message mock_send.assert_not_called() + @patch("agentcat.modules.event_queue.apply_event_redaction") + def test_process_event_with_event_redaction(self, mock_apply): + """Test processing event with the event-level redaction hook.""" + eq = EventQueue() + mock_event_redaction_fn = MagicMock() + + redacted_event = UnredactedEvent( + id="test-id", + event_type="mcp:tools/call", + project_id="project-123", + session_id="session-123", + timestamp=datetime.now(timezone.utc), + resource_name="dropped-and-replaced", + event_redaction_fn=None, # This should be cleared after redaction + ) + mock_apply.return_value = redacted_event + + event = UnredactedEvent( + id="test-id", + event_type="mcp:tools/call", + project_id="project-123", + session_id="session-123", + timestamp=datetime.now(timezone.utc), + resource_name="original", + event_redaction_fn=mock_event_redaction_fn, + ) + + with patch.object(eq, "_send_event") as mock_send: + eq._process_event(event) + + mock_apply.assert_called_once_with(event, mock_event_redaction_fn) + called_event = mock_send.call_args[0][0] + assert called_event.resource_name == "dropped-and-replaced" + assert called_event.event_redaction_fn is None + mock_send.assert_called_once() + + def test_process_event_event_redaction_drops_event(self): + """A hook returning None drops the event before it's ever sent.""" + eq = EventQueue() + + def drop_it(event): + return None + + event = UnredactedEvent( + id="test-id", + event_type="mcp:tools/call", + project_id="project-123", + session_id="session-123", + timestamp=datetime.now(timezone.utc), + event_redaction_fn=drop_it, + ) + + with patch.object(eq, "_send_event") as mock_send: + eq._process_event(event) + mock_send.assert_not_called() + + @patch("agentcat.modules.event_queue.apply_event_redaction") + @patch("agentcat.modules.event_queue.write_to_log") + def test_process_event_event_redaction_failure(self, mock_log, mock_apply): + """Test processing event when the event-level hook raises.""" + eq = EventQueue() + mock_event_redaction_fn = MagicMock() + mock_apply.side_effect = Exception("Event redaction error") + + event = UnredactedEvent( + id="test-id", + event_type="mcp:tools/call", + project_id="project-123", + session_id="session-123", + timestamp=datetime.now(timezone.utc), + event_redaction_fn=mock_event_redaction_fn, + ) + + with patch.object(eq, "_send_event") as mock_send: + eq._process_event(event) + + assert mock_log.called + log_message = str(mock_log.call_args_list[0]) + assert "WARNING" in log_message + assert "event redaction failure" in log_message + assert "test-id" in log_message + mock_send.assert_not_called() + + def test_process_event_event_hook_runs_before_string_hook(self): + """Ordering: the event hook must see raw values, before the string + hook ever runs — nothing mocked, both hooks wired for real.""" + eq = EventQueue() + order = [] + + def event_hook(event): + order.append(("event_hook", event.parameters)) + return event + + def string_hook(s): + order.append(("string_hook", s)) + return "[REDACTED]" + + event = UnredactedEvent( + id="test-id", + event_type="mcp:tools/call", + project_id="project-123", + session_id="session-123", + timestamp=datetime.now(timezone.utc), + parameters={"secret": "raw-value"}, + event_redaction_fn=event_hook, + redaction_fn=string_hook, + ) + + with patch.object(eq, "_send_event"): + eq._process_event(event) + + assert order[0] == ("event_hook", {"secret": "raw-value"}) + assert order[1] == ("string_hook", "raw-value") + @patch("agentcat.modules.event_queue.generate_prefixed_ksuid") def test_process_event_without_id(self, mock_ksuid): """Test processing event without ID generates one.""" @@ -825,6 +939,27 @@ def test_publish_event_with_redaction_function(self, mock_eq, mock_tracking): 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.event_queue") + def test_publish_event_with_event_redaction_function(self, mock_eq, mock_tracking): + """Test publishing event includes the event-level redaction hook.""" + mock_server = MagicMock() + mock_event_redaction_fn = MagicMock() + mock_tracking.return_value = _tracking_data( + options=AgentCatOptions(redact_event=mock_event_redaction_fn) + ) + + event = UnredactedEvent( + event_type="mcp:tools/call", + session_id="ses_task_handle", + timestamp=datetime.now(timezone.utc), + ) + + publish_event(mock_server, event) + + added_event = mock_eq.add.call_args[0][0] + assert added_event.event_redaction_fn == mock_event_redaction_fn + def test_module_import_installs_no_process_hooks(): """The SDK must never register signal handlers — the customer's process diff --git a/tests/test_redaction.py b/tests/test_redaction.py index 9ec2a0f..0f20934 100644 --- a/tests/test_redaction.py +++ b/tests/test_redaction.py @@ -5,7 +5,9 @@ from agentcat.modules.redaction import ( redact_strings_in_object, redact_event, + apply_event_redaction, PROTECTED_FIELDS, + RESTORED_FIELDS, ) @@ -418,7 +420,9 @@ def redact_fn(s: str) -> str: assert result.response == { "content": [{"type": "text", "text": "[REDACTED] answer"}] } - assert result.client_name == "[REDACTED] client" + # client_name is now a protected field (see PROTECTED_FIELDS) — it + # must survive untouched, same as the other protected fields below. + assert result.client_name == "SECRET client" # ...and the original is untouched, so a failure downstream cannot # publish a half-redacted object. assert event.parameters == {"arguments": {"text": "SECRET body"}} @@ -433,6 +437,7 @@ def redact_fn(s: str) -> str: assert result.project_id == "proj_keepme" assert result.event_type == "mcp:tools/call" assert result.resource_name == "add_todo" + assert result.client_name == "SECRET client" assert result.identify_actor_given_id == "SECRET actor" assert result.identify_data == {"email": "SECRET@example.com"} assert result.tags == {"env": "SECRET tag"} @@ -463,5 +468,127 @@ async def redact_fn(s: str) -> str: return s.replace("SECRET", "[REDACTED]") result = redact_event(self._event(), redact_fn) - assert result.client_name == "[REDACTED] client" + assert result.user_intent == "find the [REDACTED]" assert "coroutine" not in result.model_dump_json() + + +class TestApplyEventRedaction: + """Test suite for apply_event_redaction — the whole-event redaction hook. + + Mirrors the TypeScript SDK's applyEventRedaction/redactEvent-option tests + and the Go SDK's ApplyEventRedaction tests, since this hook exists to + close a parity gap: Python previously had only the string-level + redact_sensitive_information hook. + """ + + @staticmethod + def _event(**overrides): + from agentcat.types import UnredactedEvent + + fields = { + "session_id": "ses_keepme", + "id": "evt_keepme", + "project_id": "proj_keepme", + "event_type": "mcp:tools/call", + "resource_name": "get_credentials", + "user_intent": "raw intent", + "parameters": {"secret": "raw-value"}, + "response": {"content": [{"type": "text", "text": "raw response"}]}, + } + fields.update(overrides) + return UnredactedEvent(**fields) + + def test_hook_sees_raw_unredacted_values(self): + seen = {} + + def hook(event): + seen["parameters"] = event.parameters + seen["user_intent"] = event.user_intent + return event + + apply_event_redaction(self._event(), hook) + assert seen["parameters"] == {"secret": "raw-value"} + assert seen["user_intent"] == "raw intent" + + def test_hook_can_modify_the_event(self): + def hook(event): + event.response = None + return event + + result = apply_event_redaction(self._event(), hook) + assert result is not None + assert result.response is None + + def test_hook_returning_none_drops_the_event(self): + def drop_get_credentials(event): + if event.resource_name == "get_credentials": + return None + return event + + result = apply_event_redaction(self._event(), drop_get_credentials) + assert result is None + + def test_restored_fields_survive_forgery_attempts(self): + """A hook cannot forge or erase what AgentCat itself assigned.""" + + def forge(event): + event.id = "forged-id" + event.session_id = "forged-session" + event.project_id = "forged-project" + # A different but still-valid enum member, so the assignment + # itself doesn't raise before restoration gets a chance to run. + event.event_type = "agentcat:custom" + event.timestamp = None + return event + + original = self._event() + result = apply_event_redaction(original, forge) + + assert result.id == original.id + assert result.session_id == original.session_id + assert result.project_id == original.project_id + assert result.event_type == original.event_type + assert result.timestamp == original.timestamp + assert RESTORED_FIELDS == { + "id", + "session_id", + "project_id", + "event_type", + "timestamp", + } + + def test_a_raising_hook_propagates_so_the_queue_drops_the_event(self): + def boom(_event): + raise RuntimeError("event redaction exploded") + + with pytest.raises(RuntimeError, match="event redaction exploded"): + apply_event_redaction(self._event(), boom) + + def test_an_async_hook_is_driven_to_completion(self): + async def hook(event): + event.user_intent = "async-modified" + return event + + result = apply_event_redaction(self._event(), hook) + assert result.user_intent == "async-modified" + + def test_hook_never_sees_the_function_fields(self): + seen = {} + + def hook(event): + seen["has_redaction_fn"] = hasattr(event, "redaction_fn") + seen["has_event_redaction_fn"] = hasattr(event, "event_redaction_fn") + return event + + event = self._event(redaction_fn=lambda s: s, event_redaction_fn=hook) + apply_event_redaction(event, hook) + assert seen["has_redaction_fn"] is False + assert seen["has_event_redaction_fn"] is False + + def test_result_preserves_redaction_fn_for_the_string_hook_to_run_next(self): + def string_redact_fn(s): + return "[REDACTED]" + + event = self._event(redaction_fn=string_redact_fn) + result = apply_event_redaction(event, lambda e: e) + assert result.redaction_fn is string_redact_fn