Coverage for gco_mcp/audit_middleware.py: 94.12%

1""" 

2Audit-capture middleware for the GCO MCP server. 

3 

4Wires up two pieces: 

5 

61. ``Context.{warning, info, error, elicit}`` are wrapped at the class level 

7 (once, at module import time) so every Context instance — including the 

8 fresh one FastMCP creates for each tool call — appends to the active 

9 capture buffers. 

102. ``AuditCaptureMiddleware`` sets fresh capture buffers in 

11 ``audit_messages_var`` / ``audit_elicitations_var`` at the start of every 

12 ``on_call_tool`` invocation and resets them on the way out. The audit 

13 decorator (``gco_mcp/audit.py::_build_audit_entry``) reads those buffers 

14 when emitting the entry. 

15 

16The Context class only has its methods patched once. Idempotency is 

17enforced by inspecting an attribute we set on the patched function; 

18re-imports (test reloads, hot-reload during dev) detect the marker 

19and skip re-patching. The wrapped methods short-circuit to the 

20originals when no capture buffer is active, so this patch is a no-op 

21for any code that uses Context outside of a tool call (e.g. unit 

22tests that construct a Context directly). 

23""" 

24 

25from __future__ import annotations 

26 

27import time 

28from typing import Any 

29 

30from audit import ( 

31 _sanitize_value, 

32 _truncate_string, 

33 audit_elicitations_var, 

34 audit_messages_var, 

35 audit_resource_read, 

36) 

37from fastmcp.server.context import Context 

38from fastmcp.server.middleware import Middleware 

39 

40# Attribute we tag each spy function with so ``_install_context_patches`` 

41# can detect that the wrappers are already in place. Reading an 

42# attribute on the live class method is more robust than a separate 

43# module-level boolean: if a re-import re-runs this module, the same 

44# marker survives because the patched method on ``Context`` survives. 

45_SPY_MARKER = "_gco_audit_spy" 

46_MAX_CAPTURED_MESSAGES = 100 

47_MAX_CAPTURED_ELICITATIONS = 100 

48 

49 

50def _install_context_patches() -> None: 

51 """Install the class-level Context method wrappers (once).""" 

52 if getattr(Context.warning, _SPY_MARKER, False): 

53 return 

54 

55 _orig_warning = Context.warning 

56 _orig_info = Context.info 

57 _orig_error = Context.error 

58 _orig_elicit = Context.elicit 

59 

60 async def _spy_warning(self: Context, message: str, *args: Any, **kwargs: Any) -> Any: 

61 lst = audit_messages_var.get() 

62 if lst is not None and len(lst) < _MAX_CAPTURED_MESSAGES: 62 ↛ 64line 62 didn't jump to line 64 because the condition on line 62 was always true

63 lst.append({"level": "warning", "message": _truncate_string(str(message))}) 

64 return await _orig_warning(self, message, *args, **kwargs) 

65 

66 async def _spy_info(self: Context, message: str, *args: Any, **kwargs: Any) -> Any: 

67 lst = audit_messages_var.get() 

68 if lst is not None and len(lst) < _MAX_CAPTURED_MESSAGES: 68 ↛ 70line 68 didn't jump to line 70 because the condition on line 68 was always true

69 lst.append({"level": "info", "message": _truncate_string(str(message))}) 

70 return await _orig_info(self, message, *args, **kwargs) 

71 

72 async def _spy_error(self: Context, message: str, *args: Any, **kwargs: Any) -> Any: 

73 lst = audit_messages_var.get() 

74 if lst is not None and len(lst) < _MAX_CAPTURED_MESSAGES: 74 ↛ 76line 74 didn't jump to line 76 because the condition on line 74 was always true

75 lst.append({"level": "error", "message": _truncate_string(str(message))}) 

76 return await _orig_error(self, message, *args, **kwargs) 

77 

78 async def _spy_elicit(self: Context, message: str, *args: Any, **kwargs: Any) -> Any: 

79 result = await _orig_elicit(self, message, *args, **kwargs) 

80 lst = audit_elicitations_var.get() 

81 if lst is not None and len(lst) < _MAX_CAPTURED_ELICITATIONS: 81 ↛ 90line 81 didn't jump to line 90 because the condition on line 81 was always true

82 entry: dict[str, Any] = { 

83 "message": _truncate_string(str(message)), 

84 "action": _sanitize_value(getattr(result, "action", None), depth=0, seen=set()), 

85 } 

86 data = getattr(result, "data", None) 

87 if data is not None: 87 ↛ 89line 87 didn't jump to line 89 because the condition on line 87 was always true

88 entry["data"] = _sanitize_value(data, depth=0, seen=set()) 

89 lst.append(entry) 

90 return result 

91 

92 # Tag each spy with the marker before attaching so a concurrent 

93 # re-entry observes the marker as soon as the assignment lands. 

94 for spy in (_spy_warning, _spy_info, _spy_error, _spy_elicit): 

95 setattr(spy, _SPY_MARKER, True) 

96 

97 Context.warning = _spy_warning # type: ignore[method-assign] 

98 Context.info = _spy_info # type: ignore[method-assign] 

99 Context.error = _spy_error # type: ignore[method-assign] 

100 Context.elicit = _spy_elicit # type: ignore[method-assign] 

101 

102 

103class AuditCaptureMiddleware(Middleware): 

104 """FastMCP middleware that activates per-invocation audit capture buffers. 

105 

106 On every ``on_call_tool`` call, sets fresh empty lists into 

107 ``audit_messages_var`` and ``audit_elicitations_var`` so the patched 

108 Context methods append into them. Resets the ContextVars on the way 

109 out so concurrent calls don't see each other's captures. 

110 """ 

111 

112 def __init__(self) -> None: 

113 _install_context_patches() 

114 

115 async def on_call_tool(self, context: Any, call_next: Any) -> Any: 

116 messages: list[dict[str, str]] = [] 

117 elicitations: list[dict[str, object]] = [] 

118 msg_token = audit_messages_var.set(messages) 

119 elic_token = audit_elicitations_var.set(elicitations) 

120 try: 

121 return await call_next(context) 

122 finally: 

123 audit_messages_var.reset(msg_token) 

124 audit_elicitations_var.reset(elic_token) 

125 

126 async def on_read_resource(self, context: Any, call_next: Any) -> Any: 

127 """Audit every static and templated resource read in one place.""" 

128 started = time.time() 

129 uri = getattr(getattr(context, "message", None), "uri", "") 

130 fastmcp_context = getattr(context, "fastmcp_context", None) 

131 try: 

132 result = await call_next(context) 

133 except Exception as exc: 

134 audit_resource_read( 

135 uri, 

136 status="error", 

137 duration_ms=(time.time() - started) * 1000, 

138 error=str(exc), 

139 ctx=fastmcp_context, 

140 ) 

141 raise 

142 audit_resource_read( 

143 uri, 

144 status="success", 

145 duration_ms=(time.time() - started) * 1000, 

146 ctx=fastmcp_context, 

147 ) 

148 return result 

149 

150 

151# Install patches eagerly at module import so callers that build their own 

152# pipelines (or tests that bypass middleware wiring) still get the capture 

153# behaviour as long as they install fresh ContextVars themselves. 

154_install_context_patches()