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

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

2 

3This module provides functions for generating Prometheus metrics initialization 

4and collection code for FastMCP servers built with GolfMCP. 

5""" 

6 

7 

8def 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 

25def 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 

41def 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 

62def 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 

73def 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 

194def 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 ]