Coverage for mcpgateway / utils / correlation_id.py: 100%

1# -*- coding: utf-8 -*- 

2"""Location: ./mcpgateway/utils/correlation_id.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: MCP Gateway Contributors 

6 

7Correlation ID (Request ID) Utilities. 

8 

9This module provides async-safe utilities for managing correlation IDs (also known as 

10request IDs) throughout the request lifecycle using Python's contextvars. 

11 

12The correlation ID is a unique identifier that tracks a single request as it flows 

13through all components of the system (HTTP → Middleware → Services → Plugins → Logs). 

14 

15Key concepts: 

16- ContextVar provides per-request isolation in async environments 

17- Correlation IDs can be client-provided (X-Correlation-ID header) or auto-generated 

18- The same ID is used as request_id throughout logs, services, and plugin contexts 

19- Thread-safe and async-safe (no cross-contamination between concurrent requests) 

20""" 

21 

22# Standard 

23from contextvars import ContextVar 

24import logging 

25from typing import Dict, Optional 

26import uuid 

27 

28logger = logging.getLogger(__name__) 

29 

30# Context variable for storing correlation ID (request ID) per-request 

31# This is async-safe and provides automatic isolation between concurrent requests 

32_correlation_id_context: ContextVar[Optional[str]] = ContextVar("correlation_id", default=None) 

33 

34 

35def get_correlation_id() -> Optional[str]: 

36 """Get the current correlation ID (request ID) from context. 

37 

38 Returns the correlation ID for the current async task/request. Each request 

39 has its own isolated context, so concurrent requests won't interfere. 

40 

41 Returns: 

42 Optional[str]: The correlation ID if set, None otherwise 

43 

44 Example: 

45 >>> clear_correlation_id() # Start fresh 

46 >>> get_correlation_id() is None 

47 True 

48 >>> set_correlation_id("test-123") 

49 >>> get_correlation_id() 

50 'test-123' 

51 >>> clear_correlation_id() 

52 """ 

53 return _correlation_id_context.get() 

54 

55 

56def set_correlation_id(correlation_id: str) -> None: 

57 """Set the correlation ID (request ID) for the current context. 

58 

59 Stores the correlation ID in a context variable that's automatically isolated 

60 per async task. This ID will be used as request_id throughout the system. 

61 

62 Args: 

63 correlation_id: The correlation ID to set (typically a UUID or client-provided ID) 

64 

65 Example: 

66 >>> set_correlation_id("my-request-id") 

67 >>> get_correlation_id() 

68 'my-request-id' 

69 >>> clear_correlation_id() 

70 """ 

71 _correlation_id_context.set(correlation_id) 

72 

73 

74def clear_correlation_id() -> None: 

75 """Clear the correlation ID (request ID) from the current context. 

76 

77 Should be called at the end of request processing to clean up context. 

78 In practice, FastAPI middleware automatically handles context cleanup. 

79 

80 Note: This is optional as ContextVar automatically cleans up when the 

81 async task completes. 

82 

83 Example: 

84 >>> set_correlation_id("temp-id") 

85 >>> get_correlation_id() is not None 

86 True 

87 >>> clear_correlation_id() 

88 >>> get_correlation_id() is None 

89 True 

90 """ 

91 _correlation_id_context.set(None) 

92 

93 

94def generate_correlation_id() -> str: 

95 """Generate a new correlation ID (UUID4 hex format). 

96 

97 Creates a new random UUID suitable for use as a correlation ID. 

98 Uses UUID4 which provides 122 bits of randomness. 

99 

100 Returns: 

101 str: A new UUID in hex format (32 characters, no hyphens) 

102 

103 Example: 

104 >>> cid = generate_correlation_id() 

105 >>> len(cid) == 32 

106 True 

107 >>> cid.isalnum() 

108 True 

109 """ 

110 return uuid.uuid4().hex 

111 

112 

113def extract_correlation_id_from_headers(headers: Dict[str, str], header_name: str = "X-Correlation-ID") -> Optional[str]: 

114 """Extract correlation ID from HTTP headers. 

115 

116 Searches for the correlation ID header (case-insensitive) and returns its value. 

117 Validates that the value is non-empty after stripping whitespace. 

118 

119 Args: 

120 headers: Dictionary of HTTP headers 

121 header_name: Name of the correlation ID header (default: X-Correlation-ID) 

122 

123 Returns: 

124 Optional[str]: The correlation ID if found and valid, None otherwise 

125 

126 Example: 

127 >>> headers = {"X-Correlation-ID": "abc-123"} 

128 >>> extract_correlation_id_from_headers(headers) 

129 'abc-123' 

130 

131 >>> headers = {"x-correlation-id": "def-456"} # Case insensitive 

132 >>> extract_correlation_id_from_headers(headers) 

133 'def-456' 

134 """ 

135 # Headers can be accessed case-insensitively in FastAPI/Starlette 

136 for key, value in headers.items(): 

137 if key.lower() == header_name.lower(): 

138 correlation_id = value.strip() 

139 if correlation_id: 

140 return correlation_id 

141 return None 

142 

143 

144def get_or_generate_correlation_id() -> str: 

145 """Get the current correlation ID or generate a new one if not set. 

146 

147 This is a convenience function that ensures you always have a correlation ID. 

148 If the current context doesn't have a correlation ID, it generates and sets 

149 a new one. 

150 

151 Returns: 

152 str: The correlation ID (either existing or newly generated) 

153 

154 Example: 

155 >>> # First call generates new ID 

156 >>> id1 = get_or_generate_correlation_id() 

157 >>> # Second call returns same ID 

158 >>> id2 = get_or_generate_correlation_id() 

159 >>> assert id1 == id2 

160 """ 

161 correlation_id = get_correlation_id() 

162 if not correlation_id: 

163 correlation_id = generate_correlation_id() 

164 set_correlation_id(correlation_id) 

165 return correlation_id 

166 

167 

168def validate_correlation_id(correlation_id: Optional[str], max_length: int = 255) -> bool: 

169 """Validate a correlation ID for safety and length. 

170 

171 Checks that the correlation ID is: 

172 - Non-empty after stripping whitespace 

173 - Within the maximum length limit 

174 - Contains only safe characters (alphanumeric, hyphens, underscores) 

175 

176 Args: 

177 correlation_id: The correlation ID to validate 

178 max_length: Maximum allowed length (default: 255) 

179 

180 Returns: 

181 bool: True if valid, False otherwise 

182 

183 Example: 

184 >>> validate_correlation_id("abc-123") 

185 True 

186 >>> validate_correlation_id("abc 123") # Spaces not allowed 

187 False 

188 >>> validate_correlation_id("a" * 300) # Too long 

189 False 

190 """ 

191 if not correlation_id or not correlation_id.strip(): 

192 return False 

193 

194 correlation_id = correlation_id.strip() 

195 

196 if len(correlation_id) > max_length: 

197 logger.warning(f"Correlation ID too long: {len(correlation_id)} > {max_length}") 

198 return False 

199 

200 # Allow alphanumeric, hyphens, and underscores only 

201 if not all(c.isalnum() or c in ("-", "_") for c in correlation_id): 

202 logger.warning(f"Correlation ID contains invalid characters: {correlation_id}") 

203 return False 

204 

205 return True