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