Coverage for /Users/antonigmitruk/golf/src/golf/core/builder_metrics.py: 0%

1"""Metrics integration for the GolfMCP build process.

2

3 This module provides functions for generating Prometheus metrics initialization

4 and collection code for FastMCP servers built with GolfMCP.

5"""

6

7

8 def generate_metrics_imports() -> list[str]:

9 """Generate import statements for metrics collection.

10

11 Returns:

12 List of import statements for metrics

13 """

14 return [

15 "# Prometheus metrics imports",

16 "from golf.metrics import init_metrics, get_metrics_collector",

17 "from prometheus_client import generate_latest, CONTENT_TYPE_LATEST",

18 "from starlette.responses import Response",

19 "from starlette.middleware.base import BaseHTTPMiddleware",

20 "from starlette.requests import Request",

21 "import time",

22 ]

23

24

25 def generate_metrics_initialization(server_name: str) -> list[str]:

26 """Generate metrics initialization code.

27

28 Args:

29 server_name: Name of the MCP server

30

31 Returns:

32 List of code lines for metrics initialization

33 """

34 return [

35 "# Initialize metrics collection",

36 "init_metrics(enabled=True)",

37 "",

38 ]

39

40

41 def generate_metrics_route(metrics_path: str) -> list[str]:

42 """Generate the metrics endpoint route code.

43

44 Args:

45 metrics_path: Path for the metrics endpoint (e.g., "/metrics")

46

47 Returns:

48 List of code lines for the metrics route

49 """

50 return [

51 "# Add metrics endpoint",

52 f'@mcp.custom_route("{metrics_path}", methods=["GET"])',

53 "async def metrics_endpoint(request):",

54 ' """Prometheus metrics endpoint for monitoring."""',

55 " # Update uptime before returning metrics",

56 " update_uptime()",

57 " return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)",

58 "",

59 ]

60

61

62 def get_metrics_dependencies() -> list[str]:

63 """Get list of metrics dependencies to add to pyproject.toml.

64

65 Returns:

66 List of package requirements strings

67 """

68 return [

69 "prometheus-client>=0.19.0",

70 ]

71

72

73 def generate_metrics_instrumentation() -> list[str]:

74 """Generate metrics instrumentation wrapper functions.

75

76 Returns:

77 List of code lines for metrics instrumentation

78 """

79 return [

80 "# Metrics instrumentation wrapper functions",

81 "import time",

82 "import functools",

83 "from typing import Any, Callable",

84 "",

85 "def instrument_tool(func: Callable, tool_name: str) -> Callable:",

86 ' """Wrap a tool function with metrics collection."""',

87 " @functools.wraps(func)",

88 " async def wrapper(*args, **kwargs) -> Any:",

89 " collector = get_metrics_collector()",

90 " start_time = time.time()",

91 " status = 'success'",

92 " try:",

93 " result = await func(*args, **kwargs)",

94 " return result",

95 " except Exception as e:",

96 " status = 'error'",

97 " collector.increment_error('tool', type(e).__name__)",

98 " raise",

99 " finally:",

100 " duration = time.time() - start_time",

101 " collector.increment_tool_execution(tool_name, status)",

102 " collector.record_tool_duration(tool_name, duration)",

103 " return wrapper",

104 "",

105 "def instrument_resource(func: Callable, resource_name: str) -> Callable:",

106 ' """Wrap a resource function with metrics collection."""',

107 " @functools.wraps(func)",

108 " async def wrapper(*args, **kwargs) -> Any:",

109 " collector = get_metrics_collector()",

110 " try:",

111 " result = await func(*args, **kwargs)",

112 " # Extract URI from args if available for resource_reads metric",

113 " if args and len(args) > 0:",

114 " uri = str(args[0]) if args[0] else resource_name",

115 " else:",

116 " uri = resource_name",

117 " collector.increment_resource_read(uri)",

118 " return result",

119 " except Exception as e:",

120 " collector.increment_error('resource', type(e).__name__)",

121 " raise",

122 " return wrapper",

123 "",

124 "def instrument_prompt(func: Callable, prompt_name: str) -> Callable:",

125 ' """Wrap a prompt function with metrics collection."""',

126 " @functools.wraps(func)",

127 " async def wrapper(*args, **kwargs) -> Any:",

128 " collector = get_metrics_collector()",

129 " try:",

130 " result = await func(*args, **kwargs)",

131 " collector.increment_prompt_generation(prompt_name)",

132 " return result",

133 " except Exception as e:",

134 " collector.increment_error('prompt', type(e).__name__)",

135 " raise",

136 " return wrapper",

137 "",

138 "# HTTP Request Metrics Middleware",

139 "class MetricsMiddleware(BaseHTTPMiddleware):",

140 ' """Middleware to collect HTTP request metrics."""',

141 "",

142 " async def dispatch(self, request: Request, call_next):",

143 " collector = get_metrics_collector()",

144 " start_time = time.time()",

145 " ",

146 " # Extract path and method",

147 " method = request.method",

148 " path = request.url.path",

149 " ",

150 " try:",

151 " response = await call_next(request)",

152 " status_code = response.status_code",

153 " except Exception as e:",

154 " status_code = 500",

155 " collector.increment_error('http', type(e).__name__)",

156 " raise",

157 " finally:",

158 " duration = time.time() - start_time",

159 " collector.increment_http_request(method, status_code, path)",

160 " collector.record_http_duration(method, path, duration)",

161 " ",

162 " return response",

163 "",

164 "# Session tracking helpers",

165 "import atexit",

166 "from contextlib import asynccontextmanager",

167 "",

168 "# Global server start time for uptime tracking",

169 "_server_start_time = time.time()",

170 "",

171 "def track_session_start():",

172 ' """Track when a new session starts."""',

173 " collector = get_metrics_collector()",

174 " collector.increment_session()",

175 "",

176 "def track_session_end(start_time: float):",

177 ' """Track when a session ends."""',

178 " collector = get_metrics_collector()",

179 " duration = time.time() - start_time",

180 " collector.record_session_duration(duration)",

181 "",

182 "def update_uptime():",

183 ' """Update the uptime metric."""',

184 " collector = get_metrics_collector()",

185 " uptime = time.time() - _server_start_time",

186 " collector.set_uptime(uptime)",

187 "",

188 "# Initialize uptime tracking",

189 "update_uptime()",

190 "",

191 ]

192

193

194 def generate_session_tracking() -> list[str]:

195 """Generate session tracking integration code.

196

197 Returns:

198 List of code lines for session tracking

199 """

200 return [

201 "# Session tracking integration",

202 "import asyncio",

203 "from typing import Dict",

204 "",

205 "# Track active sessions",

206 "_active_sessions: Dict[str, float] = {}",

207 "",

208 "# Hook into FastMCP's session lifecycle if available",

209 "try:",

210 " from fastmcp.server import SessionManager",

211 " ",

212 " # Monkey patch session creation if possible",

213 " _original_create_session = getattr(mcp, '_create_session', None)",

214 " if _original_create_session:",

215 " async def _patched_create_session(*args, **kwargs):",

216 " session_id = str(id(args)) if args else 'unknown'",

217 " _active_sessions[session_id] = time.time()",

218 " track_session_start()",

219 " try:",

220 " return await _original_create_session(*args, **kwargs)",

221 " except Exception:",

222 " # If session creation fails, clean up",

223 " if session_id in _active_sessions:",

224 " del _active_sessions[session_id]",

225 " raise",

226 " ",

227 " mcp._create_session = _patched_create_session",

228 "except (ImportError, AttributeError):",

229 " # Fallback: track sessions via request patterns",

230 " pass",

231 "",

232 ]