1"""AgentOS API middleware — request/response processing pipeline.

2

3 Provides authentication, CORS, request tracing, request-ID injection,

4 and rate-limiting middleware for the AgentOS API server.

5"""

6

7 from __future__ import annotations

8

9 import time

10 import uuid

11 from dataclasses import dataclass, field

12 from typing import Callable, Optional

13

14

15# ── Request ID ────────────────────────────────────────────────────────────────

16

17

18@dataclass

19 class RequestContext:

20 """请求上下文。"""

21 request_id: str = ""

22 start_time: float = 0.0

23 method: str = ""

24 path: str = ""

25 client_ip: str = ""

26 user_agent: str = ""

27

28 @property

29 def elapsed_ms(self) -> float:

30 return (time.monotonic() - self.start_time) * 1000

31

32

33 class RequestIDMiddleware:

34 """Inject X-Request-ID into every request and propagate it."""

35

36 def __init__(self, header: str = "X-Request-ID"):

37 self.header = header

38

39 def process_request(self, headers: dict) -> RequestContext:

40 rid = headers.get(self.header.lower(), headers.get(self.header, ""))

41 if not rid:

42 rid = str(uuid.uuid4())[:12]

43 return RequestContext(request_id=rid, start_time=time.monotonic())

44

45

46# ── CORS ──────────────────────────────────────────────────────────────────────

47

48

49@dataclass

50 class CORSConfig:

51 """CORS 配置。"""

52 allow_origins: list[str] = field(default_factory=lambda: ["*"])

53 allow_methods: list[str] = field(default_factory=lambda: ["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])

54 allow_headers: list[str] = field(default_factory=lambda: ["Content-Type", "Authorization", "X-Request-ID"])

55 expose_headers: list[str] = field(default_factory=list)

56 max_age: int = 86400

57 allow_credentials: bool = False

58

59

60 class CORSMiddleware:

61 """Add CORS headers to every response."""

62

63 def __init__(self, config: Optional[CORSConfig] = None):

64 self.config = config or CORSConfig()

65

66 def apply(self, response_headers: dict) -> dict:

67 origin = self.config.allow_origins[0] if self.config.allow_origins else ""

68 response_headers["Access-Control-Allow-Origin"] = origin

69 response_headers["Access-Control-Allow-Methods"] = ", ".join(self.config.allow_methods)

70 response_headers["Access-Control-Allow-Headers"] = ", ".join(self.config.allow_headers)

71 if self.config.expose_headers:

72 response_headers["Access-Control-Expose-Headers"] = ", ".join(self.config.expose_headers)

73 response_headers["Access-Control-Max-Age"] = str(self.config.max_age)

74 if self.config.allow_credentials:

75 response_headers["Access-Control-Allow-Credentials"] = "true"

76 return response_headers

77

78

79# ── Auth ──────────────────────────────────────────────────────────────────────

80

81

82@dataclass

83 class AuthConfig:

84 """认证配置。"""

85 api_key_header: str = "X-API-Key"

86 api_key: str = ""

87 enabled: bool = True

88

89

90 class AuthMiddleware:

91 """Simple API-key authentication middleware."""

92

93 def __init__(self, config: Optional[AuthConfig] = None):

94 self.config = config or AuthConfig()

95

96 def authenticate(self, headers: dict) -> tuple[bool, str]:

97 """Return (authorized, message)."""

98 if not self.config.enabled or not self.config.api_key:

99 return True, ""

100 provided = headers.get(self.config.api_key_header.lower(), headers.get(self.config.api_key_header, ""))

101 if provided != self.config.api_key:

102 return False, "Invalid or missing API key"

103 return True, ""

104

105

106# ── Request logger ────────────────────────────────────────────────────────────

107

108

109 class RequestLogMiddleware:

110 """Log every request with method, path, status, and elapsed time."""

111

112 def __init__(self, logger: Optional[Callable[[str], None]] = None):

113 self._log = logger or print

114

115 def log(self, ctx: RequestContext, status: int) -> str:

116 msg = (

117 f"[{ctx.request_id}] {ctx.method} {ctx.path} "

118 f"→ {status} ({ctx.elapsed_ms:.1f}ms)"

119 )

120 self._log(msg)

121 return msg

122

123

124# ── Middleware stack ──────────────────────────────────────────────────────────

125

126

127 class MiddlewareStack:

128 """Ordered middleware pipeline for the AgentOS API."""

129

130 def __init__(

131 self,

132 cors: Optional[CORSMiddleware] = None,

133 auth: Optional[AuthMiddleware] = None,

134 req_log: Optional[RequestLogMiddleware] = None,

135 req_id: Optional[RequestIDMiddleware] = None,

136 ):

137 self.cors = cors or CORSMiddleware()

138 self.auth = auth or AuthMiddleware()

139 self.req_log = req_log or RequestLogMiddleware()

140 self.req_id = req_id or RequestIDMiddleware()