Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-relay-gateway/src/lexigram/ai/relay/gateway/operations/controls.py: 37%

139 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""Validated runtime control service for the relay gateway. 

2 

3``RelayControlsService`` applies channel drain/enable and typed policy 

4changes behind an explicit permission gate, persists every mutation 

5through a ``RelayPolicyStoreProtocol`` backend, refuses changes that 

6would strand the gateway without an available converter, and emits an 

7``AIAuditEvent`` for every applied mutation. Audit metadata never 

8carries credentials, upstream URLs, prompt content, or media data. 

9""" 

10 

11from __future__ import annotations 

12 

13import asyncio 

14from dataclasses import replace 

15from typing import Any 

16 

17from lexigram.ai.relay.gateway.channels import RelayChannelRegistry 

18from lexigram.ai.relay.gateway.config import RelayGatewayConfig 

19from lexigram.ai.relay.gateway.operations.streams import RelayStreamRegistry 

20from lexigram.contracts.ai.governance import ( 

21 AIAuditEvent, 

22 AIAuditStoreProtocol, 

23 AuditEventType, 

24) 

25from lexigram.contracts.ai.relay import ( 

26 RelayActiveStream, 

27 RelayGatewayError, 

28 RelayPolicyChange, 

29 RelayPolicySnapshot, 

30 RelayPolicyStoreProtocol, 

31) 

32from lexigram.contracts.auth.guard import AuthorizerProtocol 

33from lexigram.logging import get_logger 

34 

35__all__ = [ 

36 "InMemoryRelayPolicyStore", 

37 "RelayControlsService", 

38] 

39 

40logger = get_logger(__name__) 

41 

42PERMISSION_READ = "relay.read" 

43PERMISSION_CHANNEL_CONTROL = "relay.channel_control" 

44PERMISSION_POLICY_CONTROL = "relay.policy_control" 

45PERMISSION_STREAM_CONTROL = "relay.stream_control" 

46PERMISSION_CHANNEL_MANAGE = "relay.manage" 

47 

48 

49class InMemoryRelayPolicyStore(RelayPolicyStoreProtocol): 

50 """Process-local policy store seeded from the gateway configuration. 

51 

52 ``load`` always returns the current snapshot and ``save`` replaces it 

53 wholesale; the store is the single source of truth between control 

54 mutations in this process. 

55 """ 

56 

57 def __init__(self, initial: RelayPolicySnapshot) -> None: 

58 """Bind the store to its initial snapshot. 

59 

60 Args: 

61 initial: Snapshot the store serves until the first save. 

62 """ 

63 self._snapshot = initial 

64 

65 @classmethod 

66 def with_defaults(cls, config: RelayGatewayConfig) -> InMemoryRelayPolicyStore: 

67 """Build a store seeded from a gateway configuration. 

68 

69 Args: 

70 config: Gateway configuration; each channel contributes its 

71 enabled flag and declared models as the allowed options. 

72 

73 Returns: 

74 A store whose snapshot mirrors the configuration. 

75 """ 

76 enabled_channels = { 

77 channel.name: channel.enabled for channel in config.channels 

78 } 

79 allowed_options = { 

80 channel.name: frozenset(channel.models) for channel in config.channels 

81 } 

82 return cls( 

83 RelayPolicySnapshot( 

84 enabled_channels=enabled_channels, 

85 allowed_model_options=allowed_options, 

86 media_allowed_schemes=frozenset({"https"}), 

87 media_allowed_hosts=frozenset(), 

88 max_request_bytes=1024 * 1024, 

89 max_stream_seconds=300.0, 

90 ) 

91 ) 

92 

93 async def load(self) -> RelayPolicySnapshot: 

94 """Return the current snapshot.""" 

95 return self._snapshot 

96 

97 async def save(self, snapshot: RelayPolicySnapshot) -> None: 

98 """Atomically replace the stored snapshot.""" 

99 self._snapshot = snapshot 

100 

101 

102class RelayControlsService: 

103 """Apply permissioned, validated policy mutations for the gateway. 

104 

105 Every mutation is serialized through an in-process lock, validated 

106 against the static channel table, persisted to the policy store, and 

107 audited. A mutation that would leave the gateway without any 

108 enabled channel that serves at least one model option is rejected 

109 before persisting. 

110 """ 

111 

112 def __init__( 

113 self, 

114 registry: RelayChannelRegistry, 

115 store: RelayPolicyStoreProtocol, 

116 authorizer: AuthorizerProtocol | None = None, 

117 audit: AIAuditStoreProtocol | None = None, 

118 streams: RelayStreamRegistry | None = None, 

119 ) -> None: 

120 """Bind the controls service to its dependencies. 

121 

122 Args: 

123 registry: Channel table defining valid channel names and 

124 model options. 

125 store: Persistent backend that owns the current snapshot. 

126 authorizer: Permission gate for ``relay.*`` actions. When 

127 ``None`` no permission check is performed (development). 

128 audit: Audit backend for mutation events. When ``None`` 

129 mutations still apply without audit emission. 

130 streams: Registry of in-flight upstream streams. When 

131 ``None`` a private empty registry is created; share one 

132 instance with the streaming path so force-cancel reaches 

133 live streams. 

134 """ 

135 self._registry = registry 

136 self._store = store 

137 self._authorizer = authorizer 

138 self._audit = audit 

139 self.streams = streams if streams is not None else RelayStreamRegistry() 

140 self._lock = asyncio.Lock() 

141 

142 async def set_channel_state( 

143 self, 

144 channel: str, 

145 enabled: bool, 

146 actor_id: str, 

147 ) -> None: 

148 """Enable or drain *channel* for new requests. 

149 

150 Args: 

151 channel: Channel name; unknown names are rejected. 

152 enabled: ``False`` drains the channel for new requests while 

153 existing streams finish. 

154 actor_id: Operator identity recorded in the audit event. 

155 

156 Raises: 

157 ValueError: The channel is unknown. 

158 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor 

159 lacks ``relay.channel_control``. 

160 """ 

161 await self._require(actor_id, PERMISSION_CHANNEL_CONTROL, f"channel:{channel}") 

162 async with self._lock: 

163 snapshot = await self._store.load() 

164 if channel not in snapshot.enabled_channels: 

165 raise ValueError(f"unknown channel {channel!r}") 

166 changed = replace( 

167 snapshot, 

168 enabled_channels={**snapshot.enabled_channels, channel: enabled}, 

169 ) 

170 self._require_router_survives(changed) 

171 await self._store.save(changed) 

172 self._registry.set_runtime_enabled(channel, enabled) 

173 await self._audit_change( 

174 actor_id=actor_id, 

175 action=PERMISSION_CHANNEL_CONTROL, 

176 resource=channel, 

177 old={"enabled": not enabled}, 

178 new={"enabled": enabled}, 

179 ) 

180 

181 async def update_policy( 

182 self, 

183 change: RelayPolicyChange, 

184 actor_id: str, 

185 ) -> None: 

186 """Apply a typed policy change. 

187 

188 Args: 

189 change: Partial mutation; only the fields explicitly set 

190 change. 

191 actor_id: Operator identity recorded in the audit event. 

192 

193 Raises: 

194 ValueError: The change references an unknown channel or 

195 model option values, or would remove every available 

196 model option. 

197 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor 

198 lacks ``relay.policy_control``. 

199 """ 

200 await self._require(actor_id, PERMISSION_POLICY_CONTROL, "policy") 

201 async with self._lock: 

202 snapshot = await self._store.load() 

203 if change.channel is not None: 

204 if change.channel not in snapshot.enabled_channels: 

205 raise ValueError(f"unknown channel {change.channel!r}") 

206 if change.enabled is not None and change.channel is None: 

207 raise ValueError("channel is required when changing enabled state") 

208 if change.allowed_model_options is not None and change.channel is None: 

209 raise ValueError( 

210 "channel is required when changing allowed model options" 

211 ) 

212 self._validate_options(change) 

213 changed = self._compose(snapshot, change) 

214 self._require_router_survives(changed) 

215 await self._store.save(changed) 

216 if change.enabled is not None and change.channel is not None: 

217 self._registry.set_runtime_enabled(change.channel, change.enabled) 

218 await self._audit_change( 

219 actor_id=actor_id, 

220 action=PERMISSION_POLICY_CONTROL, 

221 resource="policy", 

222 old=self._changed_fields(snapshot, change), 

223 new=self._changed_fields(changed, change), 

224 ) 

225 

226 async def policy_snapshot(self, actor_id: str) -> RelayPolicySnapshot: 

227 """Return the current runtime policy snapshot. 

228 

229 Args: 

230 actor_id: Operator identity; ``relay.read`` permission is 

231 required. 

232 

233 Returns: 

234 The snapshot persisted by the policy store. 

235 

236 Raises: 

237 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor 

238 lacks ``relay.read``. 

239 """ 

240 await self._require(actor_id, PERMISSION_READ, "policy") 

241 return await self._store.load() 

242 

243 def active_streams(self) -> tuple[RelayActiveStream, ...]: 

244 """Return the currently in-flight upstream streams. 

245 

246 Returns: 

247 One row per active stream, oldest first; an empty tuple when 

248 no stream is in flight. 

249 """ 

250 return self.streams.list() 

251 

252 async def force_cancel_stream( 

253 self, 

254 stream_id: str, 

255 actor_id: str, 

256 ) -> None: 

257 """Force-cancel an in-flight upstream stream. 

258 

259 Args: 

260 stream_id: Identifier of the stream to cancel. 

261 actor_id: Operator identity recorded in the audit event; 

262 ``relay.stream_control`` permission is required. 

263 

264 Raises: 

265 ValueError: The stream identifier is unknown. 

266 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor 

267 lacks ``relay.stream_control``. 

268 """ 

269 await self._require(actor_id, PERMISSION_STREAM_CONTROL, f"stream:{stream_id}") 

270 if not self.streams.cancel(stream_id): 

271 raise ValueError(f"unknown stream {stream_id!r}") 

272 await self._audit_change( 

273 actor_id=actor_id, 

274 action=PERMISSION_STREAM_CONTROL, 

275 resource=stream_id, 

276 old={"cancelled": False}, 

277 new={"cancelled": True}, 

278 ) 

279 

280 def _validate_options(self, change: RelayPolicyChange) -> None: 

281 """Reject option names the target channel does not serve.""" 

282 options = change.allowed_model_options 

283 if options is None or change.channel is None: 

284 return 

285 channel = self._channel(change.channel) 

286 if channel is None: 

287 raise ValueError(f"unknown channel {change.channel!r}") 

288 unknown = options - set(channel.models) 

289 if unknown: 

290 name = sorted(unknown)[0] 

291 raise ValueError( 

292 f"unknown model option {name!r} for channel {change.channel!r}" 

293 ) 

294 

295 def _channel(self, name: str): 

296 """Return the configured channel with *name*, or ``None``.""" 

297 for candidate in self._registry.channels: 

298 if candidate.name == name: 

299 return candidate 

300 return None 

301 

302 async def _require(self, actor_id: str, action: str, resource: str) -> None: 

303 """Enforce *action* permission, raising ``PERMISSION_DENIED``.""" 

304 allowed = True 

305 if self._authorizer is not None: 

306 allowed = await self._authorizer.can(actor_id, action, resource) 

307 if not allowed: 

308 raise RelayGatewayError( 

309 code="PERMISSION_DENIED", 

310 message=f"{action} denied for the operator", 

311 status_code=403, 

312 request_id="", 

313 ) 

314 

315 @staticmethod 

316 def _compose( 

317 snapshot: RelayPolicySnapshot, 

318 change: RelayPolicyChange, 

319 ) -> RelayPolicySnapshot: 

320 """Build the snapshot after *change* is applied.""" 

321 enabled_channels = snapshot.enabled_channels 

322 allowed_model_options = snapshot.allowed_model_options 

323 if change.enabled is not None or change.allowed_model_options is not None: 

324 channel = change.channel 

325 if channel is None: 

326 raise ValueError( 

327 "channel is required for channel-scoped policy changes" 

328 ) 

329 if change.enabled is not None: 

330 enabled_channels = {**enabled_channels, channel: change.enabled} 

331 if change.allowed_model_options is not None: 

332 allowed_model_options = { 

333 **allowed_model_options, 

334 channel: change.allowed_model_options, 

335 } 

336 return replace( 

337 snapshot, 

338 enabled_channels=enabled_channels, 

339 allowed_model_options=allowed_model_options, 

340 media_allowed_schemes=( 

341 change.media_allowed_schemes 

342 if change.media_allowed_schemes is not None 

343 else snapshot.media_allowed_schemes 

344 ), 

345 media_allowed_hosts=( 

346 change.media_allowed_hosts 

347 if change.media_allowed_hosts is not None 

348 else snapshot.media_allowed_hosts 

349 ), 

350 max_request_bytes=( 

351 change.max_request_bytes 

352 if change.max_request_bytes is not None 

353 else snapshot.max_request_bytes 

354 ), 

355 max_stream_seconds=( 

356 change.max_stream_seconds 

357 if change.max_stream_seconds is not None 

358 else snapshot.max_stream_seconds 

359 ), 

360 ) 

361 

362 @staticmethod 

363 def _require_router_survives(snapshot: RelayPolicySnapshot) -> None: 

364 """Reject snapshots with no enabled channel serving options. 

365 

366 Raises: 

367 ValueError: When every channel is disabled or every 

368 remaining channel has no allowed model options. 

369 """ 

370 for name, options in snapshot.allowed_model_options.items(): 

371 if snapshot.enabled_channels.get(name, False) and options: 

372 return 

373 raise ValueError("policy change would remove all available converters") 

374 

375 @staticmethod 

376 def _changed_fields( 

377 snapshot: RelayPolicySnapshot, 

378 change: RelayPolicyChange, 

379 ) -> dict[str, Any]: 

380 """Extract the values changed by *change* from *snapshot*.""" 

381 fields: dict[str, Any] = {} 

382 if change.enabled is not None and change.channel is not None: 

383 fields[f"{change.channel}.enabled"] = snapshot.enabled_channels[ 

384 change.channel 

385 ] 

386 if change.allowed_model_options is not None and change.channel is not None: 

387 fields[f"{change.channel}.allowed_model_options"] = ( 

388 snapshot.allowed_model_options.get(change.channel, frozenset()) 

389 ) 

390 if change.media_allowed_schemes is not None: 

391 fields["media_allowed_schemes"] = snapshot.media_allowed_schemes 

392 if change.media_allowed_hosts is not None: 

393 fields["media_allowed_hosts"] = snapshot.media_allowed_hosts 

394 if change.max_request_bytes is not None: 

395 fields["max_request_bytes"] = snapshot.max_request_bytes 

396 if change.max_stream_seconds is not None: 

397 fields["max_stream_seconds"] = snapshot.max_stream_seconds 

398 return fields 

399 

400 async def _audit_change( 

401 self, 

402 actor_id: str, 

403 action: str, 

404 resource: str, 

405 old: dict[str, Any], 

406 new: dict[str, Any], 

407 ) -> None: 

408 """Persist one audit event for an applied mutation.""" 

409 if self._audit is None: 

410 return 

411 await self._audit.record( 

412 AIAuditEvent( 

413 event_type=AuditEventType.CONFIG_RELOADED, 

414 user_id=actor_id, 

415 status="success", 

416 metadata={ 

417 "action": action, 

418 "resource": resource, 

419 "old": old, 

420 "new": new, 

421 }, 

422 ) 

423 ) 

424 logger.debug( 

425 "relay_controls_mutated", 

426 action=action, 

427 resource=resource, 

428 )