Coverage for gco_mcp/audit.py: 94.51%
193 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2Audit logging infrastructure for the GCO MCP server.
4Provides:
5- ``_sanitize_arguments`` — redacts sensitive keys, truncates large values.
6- ``audit_logged`` — decorator that emits structured JSON audit entries for
7 every MCP tool invocation (success or failure). Dispatches on
8 ``inspect.iscoroutinefunction`` so async tools work transparently.
9- ``audit_messages_var`` / ``audit_elicitations_var`` — ContextVars populated
10 by ``gco_mcp/audit_middleware.py`` to surface ``ctx.warning``/``info``/``error``
11 /``elicit`` calls in the audit entry.
12- ``audit_resource_read`` — emits the same structured success/error metadata
13 for every MCP resource read via centralized middleware.
14- Startup audit log entry emitted only when the server entry point starts.
15"""
17import contextlib
18import contextvars
19import functools
20import inspect
21import json
22import logging
23import os
24import re
25import time
26from collections.abc import Callable
27from datetime import UTC, datetime
28from typing import Any
30import feature_flags
31from version import get_project_version
33# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
34# Generated at (UTC): 2026-07-18T01:03:40Z
35# Flowchart(s) generated from this file:
36# * ``audit_logged`` -> ``diagrams/code_diagrams/gco_mcp/audit.audit_logged.html``
37# (PNG: ``diagrams/code_diagrams/gco_mcp/audit.audit_logged.png``)
38# Regenerate with ``python diagrams/code_diagrams/generate.py``.
39# <pyflowchart-code-diagram> END
42# =============================================================================
43# AUDIT LOGGING
44# =============================================================================
46_MCP_SERVER_VERSION = get_project_version()
48audit_logger = logging.getLogger("gco.mcp.audit")
50# Patterns for sensitive argument key names (case-insensitive)
51_SENSITIVE_KEY_PATTERNS = [
52 re.compile(r".*token.*", re.IGNORECASE),
53 re.compile(r".*secret.*", re.IGNORECASE),
54 re.compile(r".*password.*", re.IGNORECASE),
55 re.compile(r".*key.*", re.IGNORECASE),
56]
58_MAX_ARG_VALUE_BYTES = 1024 # 1KB per string leaf
59_MAX_TASK_ID_BYTES = 256
60_MAX_AUDIT_DEPTH = 12
61_MAX_CONTAINER_ITEMS = 100
62_CIRCULAR_VALUE = "<circular-reference>"
63_MAX_DEPTH_VALUE = "<max-depth-exceeded>"
65# Per-invocation capture buffers populated by the audit middleware. The
66# middleware sets fresh lists at the start of every tool call; the audit
67# decorator reads them at the end and includes them in the entry when
68# non-empty. Default ``None`` means "no capture in scope" — the patched
69# Context methods short-circuit to the originals without recording.
70audit_messages_var: contextvars.ContextVar[list[dict[str, str]] | None] = contextvars.ContextVar(
71 "gco_audit_messages", default=None
72)
73audit_elicitations_var: contextvars.ContextVar[list[dict[str, object]] | None] = (
74 contextvars.ContextVar("gco_audit_elicitations", default=None)
75)
78def _is_sensitive_key(key: object) -> bool:
79 """Return whether a mapping key names a value that must be redacted."""
80 return isinstance(key, str) and any(pattern.match(key) for pattern in _SENSITIVE_KEY_PATTERNS)
83def _truncate_string(value: str) -> str:
84 """Bound one string leaf by encoded byte length."""
85 if len(value.encode("utf-8", errors="replace")) <= _MAX_ARG_VALUE_BYTES:
86 return value
87 return value[:100] + "[truncated]"
90def _truncate_task_id(value: str) -> str:
91 """Bound an audit task identifier without assuming an ASCII-only value."""
92 encoded = value.encode("utf-8", errors="replace")
93 if len(encoded) <= _MAX_TASK_ID_BYTES:
94 return value
95 marker = b"[truncated]"
96 prefix = encoded[: _MAX_TASK_ID_BYTES - len(marker)]
97 return prefix.decode("utf-8", errors="ignore") + marker.decode()
100def _sanitize_value(value: Any, *, depth: int, seen: set[int]) -> Any:
101 """Recursively redact and bound one audit value without calling user ``repr``.
103 Mapping keys are inspected at every depth before their values are visited,
104 preventing a nested secret from leaking through a parent container's string
105 representation. Containers are depth/item bounded and cycle-aware so audit
106 logging cannot become an unbounded traversal of attacker-controlled input.
107 """
108 if isinstance(value, str):
109 return _truncate_string(value)
110 if value is None or isinstance(value, (bool, int, float)):
111 try:
112 json.dumps(value, allow_nan=False)
113 return value
114 except TypeError, ValueError:
115 return f"<unserializable: {type(value).__name__}>"
116 if depth >= _MAX_AUDIT_DEPTH: 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true
117 return _MAX_DEPTH_VALUE
119 if isinstance(value, dict):
120 identity = id(value)
121 if identity in seen: 121 ↛ 122line 121 didn't jump to line 122 because the condition on line 121 was never true
122 return _CIRCULAR_VALUE
123 seen.add(identity)
124 try:
125 sanitized: dict[Any, Any] = {}
126 for index, (key, nested) in enumerate(value.items()):
127 if index >= _MAX_CONTAINER_ITEMS:
128 break
129 safe_key: Any = key
130 if not isinstance(key, (str, int, float, bool)) and key is not None: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true
131 safe_key = f"<key:{type(key).__name__}>"
132 sanitized[safe_key] = (
133 "[REDACTED]"
134 if _is_sensitive_key(key)
135 else _sanitize_value(nested, depth=depth + 1, seen=seen)
136 )
137 if len(value) > _MAX_CONTAINER_ITEMS:
138 sanitized["<truncated-items>"] = len(value) - _MAX_CONTAINER_ITEMS
139 return sanitized
140 finally:
141 seen.remove(identity)
143 if isinstance(value, (list, tuple)):
144 identity = id(value)
145 if identity in seen:
146 return _CIRCULAR_VALUE
147 seen.add(identity)
148 try:
149 sanitized_items = [
150 _sanitize_value(item, depth=depth + 1, seen=seen)
151 for item in value[:_MAX_CONTAINER_ITEMS]
152 ]
153 if len(value) > _MAX_CONTAINER_ITEMS: 153 ↛ 154line 153 didn't jump to line 154 because the condition on line 153 was never true
154 sanitized_items.append(f"<truncated-items:{len(value) - _MAX_CONTAINER_ITEMS}>")
155 return sanitized_items
156 finally:
157 seen.remove(identity)
159 # Unknown objects are intentionally not coerced through str/repr: those
160 # methods can expose credentials or execute arbitrary user code.
161 return f"<unserializable: {type(value).__name__}>"
164def _sanitize_arguments(kwargs: dict[str, Any]) -> dict[str, Any]:
165 """Recursively sanitize tool arguments for audit logging.
167 Sensitive mapping keys are redacted at every nesting depth. String leaves,
168 container depth, and container item counts are bounded. Unknown objects are
169 represented by type only, keeping audit emission JSON-safe without invoking
170 potentially secret-bearing ``__str__`` implementations.
171 """
172 sanitized: dict[str, Any] = {}
173 seen: set[int] = {id(kwargs)}
174 for key, value in kwargs.items():
175 sanitized[key] = (
176 "[REDACTED]" if _is_sensitive_key(key) else _sanitize_value(value, depth=0, seen=seen)
177 )
178 return sanitized
181def _try_get_fastmcp_context() -> Any | None:
182 """Return the active FastMCP Context if inside a request, else None.
184 Wrapping the import lets ``audit_logged`` work in unit tests that don't
185 go through an MCP request — ``get_context()`` raises ``RuntimeError`` in
186 that case, which we swallow.
187 """
188 try:
189 from fastmcp.server.dependencies import get_context
191 return get_context()
192 except Exception:
193 return None
196def _try_get_task_id(ctx: Any | None) -> str | None:
197 """Extract the active FastMCP task ID through supported APIs first.
199 Docket-backed FastMCP workers expose the protocol task through
200 ``get_task_context()`` rather than request metadata. ``Context.task_id`` is
201 the next supported surface; the metadata walk remains as a compatibility
202 fallback for older FastMCP releases and focused callers. Any identifier is
203 byte-bounded before it can enter an audit record or task-status decision.
204 """
205 candidates: list[object] = []
206 try:
207 from fastmcp.server.dependencies import get_task_context
209 candidates.append(getattr(get_task_context(), "task_id", None))
210 except Exception:
211 pass
213 if ctx is not None: 213 ↛ 221line 213 didn't jump to line 221 because the condition on line 213 was always true
214 with contextlib.suppress(Exception):
215 candidates.append(getattr(ctx, "task_id", None))
216 with contextlib.suppress(Exception):
217 request_context = getattr(ctx, "request_context", None)
218 meta = getattr(request_context, "meta", None) if request_context is not None else None
219 candidates.append(getattr(meta, "task_id", None) if meta is not None else None)
221 for task_id in candidates:
222 if isinstance(task_id, str) and task_id:
223 return _truncate_task_id(task_id)
224 return None
227def _add_request_context_fields(entry: dict[str, Any], ctx: Any | None = None) -> None:
228 """Add common request/client/task identifiers to one audit entry."""
229 ctx = _try_get_fastmcp_context() if ctx is None else ctx
230 if ctx is None:
231 return
232 try:
233 request_context = getattr(ctx, "request_context", None)
234 request_id = getattr(ctx, "request_id", None) if request_context is not None else None
235 client_id = getattr(ctx, "client_id", None)
236 except Exception:
237 request_id = None
238 client_id = None
239 if request_id:
240 entry["request_id"] = request_id
241 if client_id:
242 entry["client_id"] = client_id
243 task_id = _try_get_task_id(ctx)
244 if task_id:
245 entry["task_id"] = task_id
248def _build_audit_entry(
249 func_name: str,
250 sanitized_args: dict[str, Any],
251 status: str,
252 duration_ms: float,
253 error: str | None,
254 result: Any, # noqa: ARG001 -- reserved for future result-shape capture
255) -> dict[str, Any]:
256 """Build the JSON dict for a single tool-invocation audit entry.
258 Optional fields (``error``, ``request_id``, ``client_id``, ``task_id``,
259 ``client_messages``, ``elicitations``) are omitted when their values
260 are missing or empty. Existing sync-tool entries that don't trigger
261 any new field look identical to the pre-refactor shape.
262 """
263 entry: dict[str, Any] = {
264 "event": "mcp.tool.invocation",
265 "tool": func_name,
266 "arguments": sanitized_args,
267 "status": status,
268 "duration_ms": round(duration_ms, 2),
269 "timestamp": datetime.now(UTC).isoformat(),
270 }
271 if error:
272 entry["error"] = error[:200]
274 _add_request_context_fields(entry)
276 msgs = audit_messages_var.get()
277 if msgs:
278 entry["client_messages"] = list(msgs)
279 elics = audit_elicitations_var.get()
280 if elics:
281 entry["elicitations"] = list(elics)
283 return entry
286def audit_resource_read(
287 resource_uri: object,
288 *,
289 status: str,
290 duration_ms: float,
291 error: str | None = None,
292 ctx: Any | None = None,
293) -> None:
294 """Emit one bounded audit record for an MCP resource read."""
295 entry: dict[str, Any] = {
296 "event": "mcp.resource.read",
297 "resource_uri": _truncate_string(str(resource_uri)),
298 "status": status,
299 "duration_ms": round(duration_ms, 2),
300 "timestamp": datetime.now(UTC).isoformat(),
301 }
302 if error:
303 entry["error"] = _truncate_string(error)[:200]
304 _add_request_context_fields(entry, ctx)
305 audit_logger.info(json.dumps(entry))
308def audit_logged(func: Callable[..., Any]) -> Callable[..., Any]:
309 """Decorator that emits structured JSON audit entries for tool invocations.
311 Dispatches on ``inspect.iscoroutinefunction(func)``: async tools get an
312 async wrapper that ``await``s the call, sync tools keep the existing
313 sync path. Both wrappers share ``_build_audit_entry``.
314 """
315 if inspect.iscoroutinefunction(func):
317 @functools.wraps(func)
318 async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
319 start = time.time()
320 sanitized_args = _sanitize_arguments(kwargs)
321 try:
322 result = await func(*args, **kwargs)
323 duration_ms = (time.time() - start) * 1000
324 audit_logger.info(
325 json.dumps(
326 _build_audit_entry(
327 func.__name__,
328 sanitized_args,
329 "success",
330 duration_ms,
331 None,
332 result,
333 )
334 )
335 )
336 return result
337 except Exception as e:
338 duration_ms = (time.time() - start) * 1000
339 audit_logger.info(
340 json.dumps(
341 _build_audit_entry(
342 func.__name__,
343 sanitized_args,
344 "error",
345 duration_ms,
346 str(e),
347 None,
348 )
349 )
350 )
351 raise
353 return async_wrapper
355 @functools.wraps(func)
356 def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
357 start = time.time()
358 sanitized_args = _sanitize_arguments(kwargs)
359 try:
360 result = func(*args, **kwargs)
361 duration_ms = (time.time() - start) * 1000
362 audit_logger.info(
363 json.dumps(
364 _build_audit_entry(
365 func.__name__,
366 sanitized_args,
367 "success",
368 duration_ms,
369 None,
370 result,
371 )
372 )
373 )
374 return result
375 except Exception as e:
376 duration_ms = (time.time() - start) * 1000
377 audit_logger.info(
378 json.dumps(
379 _build_audit_entry(
380 func.__name__,
381 sanitized_args,
382 "error",
383 duration_ms,
384 str(e),
385 None,
386 )
387 )
388 )
389 raise
391 return sync_wrapper
394# =============================================================================
395# STARTUP LOG
396# =============================================================================
398# Recognised values for the ``GCO_MCP_TOOL_SEARCH`` env var. Anything outside
399# this set normalises to ``"bm25"`` — the same fallback rule that
400# ``gco_mcp/server.py`` uses when wiring the catalog-replacement transform.
401_TOOL_SEARCH_VALUES = ("bm25", "regex", "code_mode", "off")
404def _resolve_tool_search() -> str:
405 """Return the effective ``GCO_MCP_TOOL_SEARCH`` value after normalisation.
407 Mirrors the resolution in ``gco_mcp/server.py``: read the env var, strip and
408 lowercase, then fall back to ``"bm25"`` for unset, empty, or unknown
409 values so the audit entry reports what was actually wired.
410 """
411 raw = os.environ.get("GCO_MCP_TOOL_SEARCH", "bm25").strip().lower()
412 return raw if raw in _TOOL_SEARCH_VALUES else "bm25"
415def emit_startup_log() -> None:
416 """Emit the startup audit log entry."""
417 entry: dict[str, Any] = {
418 "event": "mcp.server.startup",
419 "version": _MCP_SERVER_VERSION,
420 "audit_log_level": logging.getLevelName(audit_logger.getEffectiveLevel()),
421 "timestamp": datetime.now(UTC).isoformat(),
422 }
423 if feature_flags.all_tools_enabled():
424 entry["all_tools_enabled"] = True
425 if feature_flags.is_enabled(feature_flags.FLAG_MISSION):
426 entry["mission_enabled"] = True
427 # Every per-tool flag that is effectively enabled (by its own env var or by
428 # the umbrella), so an audit consumer sees the full gated-tool surface for a
429 # run from a single line instead of diffing env vars. Sorted for stable
430 # output and omitted entirely when nothing beyond the default-on set is
431 # enabled. The all_tools_enabled / mission_enabled booleans above are kept
432 # for backward compatibility with existing audit consumers.
433 enabled_flags = sorted(f for f in feature_flags.ALL_FLAGS if feature_flags.is_enabled(f))
434 if enabled_flags:
435 entry["enabled_flags"] = enabled_flags
436 tool_search = _resolve_tool_search()
437 entry["tool_search"] = tool_search
438 if tool_search == "code_mode":
439 entry["code_mode_experimental"] = True
440 audit_logger.info(json.dumps(entry))