diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ef0ff8a61..f23c51fdb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -156,7 +156,9 @@ jobs: # the lockfile's mcp 1.x (also exercised incidentally by the `tests` # matrix — this leg exists as an explicit, named signal); the v2 leg # (spec 2026-07-28) swaps in mcp>=2 and standalone FastMCP 4, which uses - # the v2 registry. The v2-bare leg runs mcp>=2 without FastMCP: importing + # the v2 registry. The v1-fastmcp2 leg keeps mcp 1.x and swaps in + # standalone FastMCP 2, which dispatches tools without a public + # call_tool. The v2-bare leg runs mcp>=2 without FastMCP: importing # FastMCP 4 patches camelCase properties onto mcp_types models for the # whole process, which hides reads that plain mcp 2.x rejects. # posthog/test/mcp/conftest.py splits collection by major. @@ -166,8 +168,10 @@ jobs: strategy: matrix: python-version: ['3.10', '3.14'] - mcp-major: ['v1', 'v2', 'v2-bare'] + mcp-major: ['v1', 'v1-fastmcp2', 'v2', 'v2-bare'] include: + - mcp-major: 'v1-fastmcp2' + mcp-packages: "'mcp>=1.28.1,<2' 'fastmcp>=2,<3'" - mcp-major: 'v2' mcp-packages: "'mcp>=2,<3' 'fastmcp>=4,<5'" - mcp-major: 'v2-bare' @@ -194,7 +198,7 @@ jobs: run: | UV_PROJECT_ENVIRONMENT=$pythonLocation uv sync --extra test - - name: Swap in MCP SDK v2 + - name: Swap in MCP packages if: matrix.mcp-major != 'v1' shell: bash run: | diff --git a/.sampo/changesets/fastmcp-tool-error-exception.md b/.sampo/changesets/fastmcp-tool-error-exception.md new file mode 100644 index 000000000..ab1520bd5 --- /dev/null +++ b/.sampo/changesets/fastmcp-tool-error-exception.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +`posthog.mcp`: a failed tool call on a standalone FastMCP server now reports the exception the tool raised, with its cause chain and stack frames, and `$mcp_error_type` names its class. Before, FastMCP converted the exception into an error result first, so the failure was a generic `Error` with no frames. diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 98a73bede..9f9db0211 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -290,6 +290,7 @@ def instrument( from ._instrument_fastmcp import instrument_fastmcp from ._instrument_lowlevel import instrument_fastmcp_v2, instrument_low_level from ._instrument_v2 import instrument_lowlevel_v2, instrument_mcpserver_v2 + from ._raised_tool_error import wrap_tool_dispatch client = _resolve_client(posthog_client) if client is None: @@ -306,8 +307,10 @@ def instrument( server_build=server_build, ) - if is_fastmcp_v2(server) and uses_v2_handler_registry(key): - data.standalone_fastmcp = weakref.ref(server) + if is_fastmcp_v2(server): + wrap_tool_dispatch(server) + if uses_v2_handler_registry(key): + data.standalone_fastmcp = weakref.ref(server) # A standalone FastMCP wrapper and its low-level server share one tracking # key, so instrumenting the second of the pair must still attach what only diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index a8ab58671..8d6cf4393 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -54,6 +54,7 @@ from ._model_parameters import request_meta_from_context from ._model_parameters import can_inject_model_parameter, is_capture_model_enabled from ._output_instructions import mirror_instructions_into_structured_content +from ._raised_tool_error import raised_tool_error from .logger import log, warn from .tools import get_more_tools_result_text @@ -282,7 +283,8 @@ async def handler(req: Any) -> Any: start = time.monotonic() try: - result = await original(req) + with raised_tool_error(high_level) as raised: + result = await original(req) except Exception as error: if await lifecycle.record_if_unknown_tool( await registered_tool_exists(high_level, name) @@ -349,6 +351,7 @@ async def handler(req: Any) -> Any: call_result, duration_ms, conversation_id_delivered=delivered, + error=raised.error, ) return result diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index b7d597789..da85d2c06 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -75,6 +75,7 @@ request_meta_from_context, ) from ._output_instructions import mirror_instructions_into_structured_content +from ._raised_tool_error import raised_tool_error from ._tool_input import get_registered_tool_input_schema from .asgi import get_mcp_session from .logger import log @@ -673,7 +674,8 @@ async def handler(ctx: Any, params: Any) -> Any: start = time.monotonic() try: - result = await original(ctx, params) + with raised_tool_error(standalone) as raised: + result = await original(ctx, params) except Exception as error: if await lifecycle.record_if_unknown_tool( await registered_tool_exists(standalone, name) @@ -702,7 +704,10 @@ async def handler(ctx: Any, params: Any) -> Any: return result await lifecycle.record_result( - result, duration_ms, conversation_id_delivered=delivered + result, + duration_ms, + conversation_id_delivered=delivered, + error=raised.error, ) return result diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 9b1db44c7..023497620 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -665,7 +665,12 @@ async def record_error(self, error: Any, duration_ms: float) -> None: ) async def record_result( - self, result: Any, duration_ms: float, *, conversation_id_delivered: bool + self, + result: Any, + duration_ms: float, + *, + conversation_id_delivered: bool, + error: Optional[BaseException] = None, ) -> None: conversation_id = self._anchored_conversation_id(conversation_id_delivered) session_id = await self.prepare_session(conversation_id) @@ -678,6 +683,7 @@ async def record_result( allow_self_reported_model=self.allow_self_reported_model, analytics_owned_parameters=self.analytics_owned_parameters, result=result, + error=error, duration_ms=duration_ms, client_name=self.client_name, client_version=self.client_version, @@ -813,14 +819,13 @@ async def record_tool_call( event["llm_model"] = model event["llm_model_source"] = source + if result is not None: + event["response"] = _wrap_response(result) + if error is None and is_tool_result_error(result): + error = result if error is not None: event["is_error"] = True event["error"] = capture_exception(error) - elif result is not None: - event["response"] = _wrap_response(result) - if is_tool_result_error(result): - event["is_error"] = True - event["error"] = capture_exception(result) props = await resolve_event_properties(data, request, extra) input_aliases = None diff --git a/posthog/mcp/_posthog_events.py b/posthog/mcp/_posthog_events.py index f7b58f0e7..e20971aed 100644 --- a/posthog/mcp/_posthog_events.py +++ b/posthog/mcp/_posthog_events.py @@ -178,6 +178,7 @@ def _add_common_properties(event: Event, properties: Dict[str, Any]) -> None: # its text, so reporting it loses nothing. _DISPATCH_WRAPPERS = ( (("ToolError", "UnexpectedToolError"), "Error executing tool"), + (("ToolError",), "Error calling tool '"), (("MCPError", "UnexpectedResourceError"), "Error reading resource"), ) diff --git a/posthog/mcp/_raised_tool_error.py b/posthog/mcp/_raised_tool_error.py new file mode 100644 index 000000000..e694cf6fe --- /dev/null +++ b/posthog/mcp/_raised_tool_error.py @@ -0,0 +1,83 @@ +"""Recover the exception behind a standalone FastMCP ``isError`` tool result. + +FastMCP's ``tools/call`` handler, the one we wrap, turns whatever its tool +dispatch raised into ``CallToolResult(isError=True)``, keeping only the +message text. Wrapping that dispatch on the server instance sees the exception +first, and hands it to the enclosing handler through a context variable. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Any, Iterator, Optional + +_WRAPPED_FLAG = "__posthog_mcp_wrapped__" + +# fastmcp 3+ dispatches tools/call through the public `call_tool`; 2.x has no +# `call_tool` and dispatches through `_call_tool_middleware`. +_DISPATCH_METHODS = ("call_tool", "_call_tool_middleware") + + +@dataclass +class RaisedToolError: + error: Optional[BaseException] = None + + +@dataclass +class _Slot: + server: Any + error: Optional[BaseException] = None + open: bool = True + + +_current: ContextVar[Optional[_Slot]] = ContextVar( + "posthog_mcp_raised_tool_error", default=None +) + + +def wrap_tool_dispatch(server: Any) -> None: + method = next((m for m in _DISPATCH_METHODS if hasattr(server, m)), None) + if method is None: + return + original = getattr(server, method) + if getattr(original, _WRAPPED_FLAG, False): + return + + async def dispatch(*args: Any, **kwargs: Any) -> Any: + # A tool calling another tool nests these calls. The outermost call + # settles last, so the slot ends holding the outcome the client sees. + try: + result = await original(*args, **kwargs) + except Exception as error: + _settle(server, error) + raise + _settle(server, None) + return result + + setattr(dispatch, _WRAPPED_FLAG, True) + setattr(server, method, dispatch) + + +def _settle(server: Any, error: Optional[BaseException]) -> None: + # Another server's dispatch, reached from inside a tool, must not write + # its outcome onto this call, nor may a background task that outlives it. + slot = _current.get() + if slot is not None and slot.open and slot.server is server: + slot.error = error + + +@contextmanager +def raised_tool_error(server: Any) -> Iterator[RaisedToolError]: + slot = _Slot(server) + raised = RaisedToolError() + token = _current.set(slot) + try: + yield raised + finally: + _current.reset(token) + # A background task started by the tool keeps a copy of this context, + # so the slot outlives the call; it must not keep the traceback alive. + raised.error, slot.error = slot.error, None + slot.open = False diff --git a/posthog/mcp/_sanitization.py b/posthog/mcp/_sanitization.py index 426d2d29f..12b442245 100644 --- a/posthog/mcp/_sanitization.py +++ b/posthog/mcp/_sanitization.py @@ -10,6 +10,7 @@ from __future__ import annotations +import json import re from typing import Any, Dict, List, Optional, Set, Tuple from urllib.parse import SplitResult, parse_qsl, urlencode, urlsplit, urlunsplit @@ -32,6 +33,10 @@ re.IGNORECASE, ) +_JSON_STRING_STOP = re.compile(r'["\\\n]') +_JSON_NUMBER = re.compile(r"-?\d[\d.eE+-]*") +_NON_SPACE = re.compile(r"\S") + # No leading `\b`: a word boundary needs a non-word character before the scheme, # so `resource_https://user:pw@host` (an `_` before the scheme) would match # nothing and stay unredacted. Without it the leftmost match wins — `foo.https://x` @@ -442,9 +447,155 @@ def _redact_credentials(value: str) -> str: entropy detector scans. Running the URL pass last loses nothing — it only redacts or percent-encodes, so it never exposes a credential the detectors could have matched.""" + value = _redact_json_members(value) return _redact_secret_tokens(_POSTHOG_TOKEN_PATTERN.sub(_REDACTED_VALUE, value)) +def _redact_json_members(value: str) -> str: + if '"' not in value: + return value + return _JsonMemberScan(value).redacted() + + +class _JsonMemberScan: + """Redacts the members posthog-go's jsonMemberPattern and posthog-rs's + JSON_MEMBER match: `"key": value`, where the value is a string, number, + boolean, or flat array of those (with every item after the first needing its + comma). A tool that returns structuredContent also returns it as JSON text, + where a credential is a member inside a string rather than a dict key. + + Like those patterns, any quote may open a key, and the scan resumes after + each member. They run on linear-time engines; Python's `re` backtracks and + would rescan the rest of the text from every quote, so this caches where a + string body ends and what follows each closing quote. + + A JSON string never holds a raw newline, so one cut off by truncation ends + at the line or the text instead of swallowing prose up to the next quote. + Values that nest objects, and JSON inside a JSON string, are out of scope. + """ + + def __init__(self, text: str) -> None: + self.text = text + self.body_ends: Dict[int, int] = {} + self.members: Dict[int, Optional[Tuple[int, int]]] = {} + + def redacted(self) -> str: + text = self.text + parts: List[str] = [] + last = 0 + quote = text.find('"') + while quote != -1: + key_end = self.body_end(quote + 1) + member = ( + self.member_after(key_end + 1) + if key_end < len(text) and text[key_end] == '"' + else None + ) + if member is None: + quote = text.find('"', quote + 1) + continue + value_start, value_end = member + if _should_redact_key(_decoded_json_key(text[quote + 1 : key_end])): + parts.append(text[last:value_start]) + parts.append(f'"{_REDACTED_VALUE}"') + last = value_end + quote = text.find('"', value_end) + parts.append(text[last:]) + return "".join(parts) + + def body_end(self, start: int) -> int: + """Where a string body starting at `start` stops: its closing quote, a + newline, a backslash with nothing left to escape, or the end.""" + text = self.text + escapes = [] + pos = start + while True: + stop = _JSON_STRING_STOP.search(text, pos) + end = stop.start() if stop else len(text) + if end in self.body_ends: + end = self.body_ends[end] + break + if text[end : end + 1] != "\\" or text[end + 1 : end + 2] in ("", "\n"): + break + escapes.append(end) + pos = end + 2 + for escape in escapes: + self.body_ends[escape] = end + return end + + def member_after(self, pos: int) -> Optional[Tuple[int, int]]: + """Where the value after a key that closed just before `pos` starts and + ends, or None when no `: value` follows. Every quote inside a run of + escaped quotes closes on the same key, so this is computed once.""" + if pos not in self.members: + colon = self.skip_space(pos) + member = None + if self.text[colon : colon + 1] == ":": + value_start = self.skip_space(colon + 1) + value_end = self.value_end(value_start) + if value_end is not None: + member = (value_start, value_end) + self.members[pos] = member + return self.members[pos] + + def value_end(self, pos: int) -> Optional[int]: + if self.text[pos : pos + 1] == "[": + return self.array_end(pos + 1) + return self.scalar_end(pos) + + def scalar_end(self, pos: int) -> Optional[int]: + text = self.text + if text[pos : pos + 1] == '"': + end = self.body_end(pos + 1) + return end + 1 if text[end : end + 1] in ('"', "\\") else end + number = _JSON_NUMBER.match(text, pos) + if number: + return number.end() + for literal in ("true", "false"): + if text.startswith(literal, pos): + return pos + len(literal) + return None + + def item_end(self, pos: int) -> Optional[int]: + end = self.scalar_end(pos) + if end is None and self.text.startswith("null", pos): + end = pos + 4 + return None if end is None else self.skip_space(end) + + def array_end(self, pos: int) -> Optional[int]: + text = self.text + pos = self.skip_space(pos) + item = self.item_end(pos) + if item is not None: + pos = item + while text[pos : pos + 1] == ",": + after_comma = self.skip_space(pos + 1) + item = self.item_end(after_comma) + if item is None: + pos = after_comma + break + pos = item + if text[pos : pos + 1] == "]": + return pos + 1 + if pos == len(text) or (pos == len(text) - 1 and text[pos] == "\n"): + return pos + return None + + def skip_space(self, pos: int) -> int: + non_space = _NON_SPACE.search(self.text, pos) + return non_space.start() if non_space else len(self.text) + + +def _decoded_json_key(raw: str) -> str: + if "\\" not in raw: + return raw + try: + decoded = json.loads(f'"{raw}"') + except ValueError: + return raw + return decoded if isinstance(decoded, str) else raw + + def sanitize_free_text(value: Any) -> Any: """Sanitize agent-narrated free text (the intent, the send_feedback fields): the binary gate, the credential passes, structured PII, then URLs. @@ -680,14 +831,39 @@ def _sanitize_exception_values(error: Any) -> Any: return { **error, "$exception_list": [ - {**exception, "value": sanitize_captured_value(exception.get("value"))} - if isinstance(exception, dict) - else exception + _sanitize_exception(exception) if isinstance(exception, dict) else exception for exception in exception_list ], } +_FRAME_SOURCE_KEYS = ("pre_context", "context_line", "post_context") + + +def _sanitize_exception(exception: Dict[str, Any]) -> Dict[str, Any]: + sanitized = {**exception, "value": sanitize_captured_value(exception.get("value"))} + stacktrace = exception.get("stacktrace") + if isinstance(stacktrace, dict) and isinstance(stacktrace.get("frames"), list): + sanitized["stacktrace"] = { + **stacktrace, + "frames": [_sanitize_frame(frame) for frame in stacktrace["frames"]], + } + return sanitized + + +def _sanitize_frame(frame: Any) -> Any: + if not isinstance(frame, dict): + return frame + return { + **frame, + **{ + key: sanitize_captured_value(frame[key]) + for key in _FRAME_SOURCE_KEYS + if key in frame + }, + } + + def _sanitize_response(response: Any) -> Any: if response is None or not isinstance(response, (dict, list, str)): return sanitize_captured_value(response) diff --git a/posthog/test/mcp/test_fastmcp_tool_errors.py b/posthog/test/mcp/test_fastmcp_tool_errors.py new file mode 100644 index 000000000..c1f1100cd --- /dev/null +++ b/posthog/test/mcp/test_fastmcp_tool_errors.py @@ -0,0 +1,209 @@ +"""Standalone FastMCP turns a raised tool exception into an ``isError`` result +before the handler we wrap returns. The failure must still carry what was raised +(posthog-python#1223).""" + +import asyncio +import gc +import logging +import weakref + +import pytest + +pytest.importorskip("fastmcp") + +from fastmcp import Client, FastMCP # noqa: E402 +from fastmcp.exceptions import ToolError # noqa: E402 + +from posthog.mcp import MCPAnalyticsOptions, instrument # noqa: E402 +from posthog.test.mcp._helpers import FakeClient, events_named, flush_background # noqa: E402 + + +def _upstream(): + raise ConnectionError("upstream down") + + +def _chained(): + try: + _upstream() + except ConnectionError as exc: + raise RuntimeError("lookup failed") from exc + + +def _tool_error(): + raise ToolError("bad input") + + +def _tool_error_worded_like_fastmcp(): + try: + _upstream() + except ConnectionError as exc: + raise ToolError("Error calling tool: lookup failed") from exc + + +@pytest.mark.parametrize( + "raise_in_tool, mask_error_details, error_type, message, chain", + [ + ( + _chained, + False, + "RuntimeError", + "lookup failed", + ["RuntimeError", "ConnectionError"], + ), + ( + _chained, + True, + "RuntimeError", + "lookup failed", + ["RuntimeError", "ConnectionError"], + ), + (_tool_error, False, "ToolError", "bad input", ["ToolError"]), + ( + _tool_error_worded_like_fastmcp, + False, + "ToolError", + "Error calling tool: lookup failed", + ["ToolError", "ConnectionError"], + ), + ], +) +async def test_failed_call_carries_the_raised_exception( + raise_in_tool, mask_error_details, error_type, message, chain +): + server = FastMCP("example-server", mask_error_details=mask_error_details) + + @server.tool + def fail() -> str: + raise_in_tool() + return "unreachable" + + sink = FakeClient() + instrument(server, sink, MCPAnalyticsOptions()) + + async with Client(server) as client: + result = await client.call_tool("fail", {}, raise_on_error=False) + await flush_background() + + assert result.is_error + [call] = events_named(sink, "$mcp_tool_call") + assert call["properties"]["$mcp_is_error"] is True + assert call["properties"]["$mcp_error_type"] == error_type + assert call["properties"]["$mcp_error_message"] == message + [exception] = events_named(sink, "$exception") + exception_list = exception["properties"]["$exception_list"] + raised = [e for e in exception_list if e["type"] in chain] + assert [e["type"] for e in raised] == chain + assert all(e["stacktrace"]["frames"] for e in raised) + + +@pytest.mark.parametrize( + "failing_server, instrumented", + [ + ("same", "wrapper"), + ("other", "wrapper"), + ("other", "low_level_only"), + ], +) +async def test_exception_caught_by_the_tool_is_not_reported( + failing_server, instrumented +): + if not hasattr(FastMCP, "call_tool"): + pytest.skip("fastmcp 2 has no public call_tool for a tool to call") + server = FastMCP("example-server") + other = FastMCP("other-server") + instrument(other, FakeClient(), MCPAnalyticsOptions()) + backend = server if failing_server == "same" else other + + @backend.tool + def inner() -> str: + raise RuntimeError("inner failure") + + @server.tool + async def outer() -> str: + try: + await backend.call_tool("inner", {}) + except ToolError: + pass + return "recovered" + + sink = FakeClient() + target = server if instrumented == "wrapper" else server._mcp_server + instrument(target, sink, MCPAnalyticsOptions()) + + async with Client(server) as client: + result = await client.call_tool("outer", {}, raise_on_error=False) + await flush_background() + + assert not result.is_error + [call] = events_named(sink, "$mcp_tool_call") + assert call["properties"]["$mcp_is_error"] is False + assert events_named(sink, "$exception") == [] + + +async def test_secrets_in_frame_source_lines_are_redacted(): + server = FastMCP("example-server") + + @server.tool + def fail() -> str: + dsn = "postgres://admin:hunter2@db.example.com/app" + credentials = {"password": "swordfish", "account": "reader"} + raise RuntimeError(f"connect failed for {credentials['account']}@{dsn[-6:]}") + + sink = FakeClient() + instrument(server, sink, MCPAnalyticsOptions()) + + async with Client(server) as client: + await client.call_tool("fail", {}, raise_on_error=False) + await flush_background() + + [exception] = events_named(sink, "$exception") + frames = [ + frame + for entry in exception["properties"]["$exception_list"] + for frame in entry["stacktrace"]["frames"] + ] + assert "db.example.com/app" in repr(frames) + assert '"account": "reader"' in repr(frames) + assert "hunter2" not in repr(frames) + assert "swordfish" not in repr(frames) + + +@pytest.fixture +def logging_off(): + # pytest's log capture keeps fastmcp's "Error calling tool" records, and + # their exc_info would keep the traceback alive on its own. + logging.disable(logging.CRITICAL) + yield + logging.disable(logging.NOTSET) + + +async def test_a_background_task_does_not_keep_the_failure_alive(logging_off): + server = FastMCP("example-server") + release = asyncio.Event() + lingering = [] + payload_refs = [] + + class Payload: + pass + + @server.tool + async def fail() -> str: + payload = Payload() + payload_refs.append(weakref.ref(payload)) + lingering.append(asyncio.get_running_loop().create_task(release.wait())) + raise RuntimeError("failed") + + sink = FakeClient() + instrument(server, sink, MCPAnalyticsOptions()) + + async with Client(server) as client: + await client.call_tool("fail", {}, raise_on_error=False) + await flush_background() + gc.collect() + + try: + assert not lingering[0].done() + assert payload_refs[0]() is None + finally: + release.set() + await lingering[0] diff --git a/posthog/test/mcp/test_pipeline.py b/posthog/test/mcp/test_pipeline.py index ea7457023..1e8728920 100644 --- a/posthog/test/mcp/test_pipeline.py +++ b/posthog/test/mcp/test_pipeline.py @@ -1,5 +1,8 @@ """Unit tests for the MCP analytics core pipeline (Milestone 1, no server).""" +import json +import re +import random from datetime import datetime, timezone import pytest @@ -9,6 +12,7 @@ PostHogMCPAnalyticsEvent, PostHogMCPAnalyticsProperty, ) +import posthog.mcp._sanitization as sanitization from posthog.mcp._event_types import MCPAnalyticsEventType from posthog.mcp._exceptions import capture_exception from posthog.mcp._ids import deterministic_prefixed_id, new_prefixed_id @@ -73,6 +77,138 @@ def test_sanitize_redacts_sensitive_keys(): assert out["safe"] == "keep" +@pytest.mark.parametrize( + "value, expected", + [ + pytest.param( + '{"user":"ada","password":"hunter2"}', + '{"user":"ada","password":"[redacted]"}', + id="string", + ), + pytest.param( + '{\n "Password" : "hun\\"ter2",\n "n": 1\n}', + '{\n "Password" : "[redacted]",\n "n": 1\n}', + id="escaped quote and spacing", + ), + pytest.param('{"api_key":12345}', '{"api_key":"[redacted]"}', id="number"), + pytest.param( + '{"access-token":true}', '{"access-token":"[redacted]"}', id="boolean" + ), + pytest.param( + '{"refresh_token":"abcdefgh', + '{"refresh_token":"[redacted]"', + id="cut off at the end of the text", + ), + pytest.param( + '{"private_key":"abc\\', + '{"private_key":"[redacted]"', + id="cut off after an escape", + ), + pytest.param( + '{"Authorization":["Bearer x"]}', + '{"Authorization":"[redacted]"}', + id="array", + ), + pytest.param( + '{"cookie": [\n "a=1",\n "b=2"\n]}', + '{"cookie": "[redacted]"}', + id="indented array", + ), + pytest.param( + '{"set-cookie":["a=1","b=2', + '{"set-cookie":"[redacted]"', + id="array cut off", + ), + pytest.param( + 'retry failed: "token": "abc"', + 'retry failed: "token": "[redacted]"', + id="inside prose", + ), + pytest.param( + '"token": "abc\nnext "line"', + '"token": "[redacted]"\nnext "line"', + id="unclosed in prose stops at the line", + ), + pytest.param( + '{"pass\\u0077ord":"hunter2"}', + '{"pass\\u0077ord":"[redacted]"}', + id="escaped key", + ), + pytest.param( + ' credentials = {"password": "hunter2", "account": "reader"}', + ' credentials = {"password": "[redacted]", "account": "reader"}', + id="source line", + ), + pytest.param( + '{"password_hint":"pet"}', + '{"password_hint":"pet"}', + id="key only contains a sensitive word", + ), + pytest.param('{"token":null}', '{"token":null}', id="null"), + pytest.param( + '{"secret":{"v":"x"}}', + '{"secret":{"v":"x"}}', + id="out of scope: nested object", + ), + pytest.param( + '{"body":"{\\"password\\":\\"x\\"}"}', + '{"body":"{\\"password\\":\\"x\\"}"}', + id="out of scope: JSON inside a JSON string", + ), + ], +) +def test_sanitize_redacts_sensitive_json_members_inside_strings(value, expected): + assert sanitize_captured_value(value) == expected + + +@pytest.mark.timeout(5) +@pytest.mark.parametrize( + "value", + [ + pytest.param('{"note":[' + "1" * 40 + "!", id="digits"), + pytest.param('{"note":[' + "1 " * 40 + "!", id="spaced digits"), + pytest.param('{"note":[' + "1," * 20_000 + "!", id="long array"), + pytest.param('\\"' * 32_000, id="escaped quotes"), + pytest.param(json.dumps('"' * 32_000), id="JSON-encoded quotes"), + ], +) +def test_json_member_redaction_runs_in_linear_time_on_adversarial_text(value): + assert sanitize_captured_value(value) == value + + +# The pattern posthog-go and posthog-rs match JSON members with (with every +# array item after the first needing its comma), run on Python's re. The SDK's +# scanner must find exactly its members, only without its quadratic retries. +_JSON_STRING_ORACLE = r'"(?:[^"\\\n]|\\.)*(?:"|\\?(?m:$))' +_JSON_SCALAR_ORACLE = _JSON_STRING_ORACLE + r"|-?\d[\d.eE+-]*|true|false" +_JSON_ITEM_ORACLE = r"(?:" + _JSON_SCALAR_ORACLE + r"|null)\s*" +_JSON_MEMBER_ORACLE = re.compile( + r'("((?:[^"\\\n]|\\.)*)"\s*:\s*)(?:' + + _JSON_SCALAR_ORACLE + + r"|\[\s*(?:" + + _JSON_ITEM_ORACLE + + r"(?:,\s*" + + _JSON_ITEM_ORACLE + + r")*(?:,\s*)?)?(?:\]|$))" +) + + +def test_json_member_scan_finds_what_the_reference_pattern_finds(monkeypatch): + monkeypatch.setattr(sanitization, "_should_redact_key", lambda key: True) + rng = random.Random(1223) + alphabet = ['"', '"', "\\", ":", ",", "[", "]", " ", "\n", "1", "-", "e", "x"] + words = ["true", "false", "null", '"k"', '"k": ', "[1, 2]", '\\"'] + texts = [ + "".join(rng.choice(alphabet + words) for _ in range(rng.randrange(16))) + for _ in range(30_000) + ] + for text in texts: + expected = _JSON_MEMBER_ORACLE.sub( + lambda member: member.group(1) + '"[redacted]"', text + ) + assert sanitization._redact_json_members(text) == expected, repr(text) + + def test_sanitize_redacts_large_base64(): blob = "A" * 11000 assert sanitize_captured_value(blob).startswith("[binary data redacted")