Coverage for src / lexigram / ai / relay / gateway / operations / controls.py: 93%
138 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-08 23:08 +0800
1"""Validated runtime control service for the relay gateway.
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"""
11from __future__ import annotations
13import asyncio
14from dataclasses import replace
15from typing import Any
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
35__all__ = [
36 "InMemoryRelayPolicyStore",
37 "RelayControlsService",
38]
40logger = get_logger(__name__)
42PERMISSION_READ = "relay.read"
43PERMISSION_CHANNEL_CONTROL = "relay.channel_control"
44PERMISSION_POLICY_CONTROL = "relay.policy_control"
45PERMISSION_STREAM_CONTROL = "relay.stream_control"
48class InMemoryRelayPolicyStore(RelayPolicyStoreProtocol):
49 """Process-local policy store seeded from the gateway configuration.
51 ``load`` always returns the current snapshot and ``save`` replaces it
52 wholesale; the store is the single source of truth between control
53 mutations in this process.
54 """
56 def __init__(self, initial: RelayPolicySnapshot) -> None:
57 """Bind the store to its initial snapshot.
59 Args:
60 initial: Snapshot the store serves until the first save.
61 """
62 self._snapshot = initial
64 @classmethod
65 def with_defaults(cls, config: RelayGatewayConfig) -> InMemoryRelayPolicyStore:
66 """Build a store seeded from a gateway configuration.
68 Args:
69 config: Gateway configuration; each channel contributes its
70 enabled flag and declared models as the allowed options.
72 Returns:
73 A store whose snapshot mirrors the configuration.
74 """
75 enabled_channels = {
76 channel.name: channel.enabled for channel in config.channels
77 }
78 allowed_options = {
79 channel.name: frozenset(channel.models) for channel in config.channels
80 }
81 return cls(
82 RelayPolicySnapshot(
83 enabled_channels=enabled_channels,
84 allowed_model_options=allowed_options,
85 media_allowed_schemes=frozenset({"https"}),
86 media_allowed_hosts=frozenset(),
87 max_request_bytes=1024 * 1024,
88 max_stream_seconds=300.0,
89 )
90 )
92 async def load(self) -> RelayPolicySnapshot:
93 """Return the current snapshot."""
94 return self._snapshot
96 async def save(self, snapshot: RelayPolicySnapshot) -> None:
97 """Atomically replace the stored snapshot."""
98 self._snapshot = snapshot
101class RelayControlsService:
102 """Apply permissioned, validated policy mutations for the gateway.
104 Every mutation is serialized through an in-process lock, validated
105 against the static channel table, persisted to the policy store, and
106 audited. A mutation that would leave the gateway without any
107 enabled channel that serves at least one model option is rejected
108 before persisting.
109 """
111 def __init__(
112 self,
113 registry: RelayChannelRegistry,
114 store: RelayPolicyStoreProtocol,
115 authorizer: AuthorizerProtocol | None = None,
116 audit: AIAuditStoreProtocol | None = None,
117 streams: RelayStreamRegistry | None = None,
118 ) -> None:
119 """Bind the controls service to its dependencies.
121 Args:
122 registry: Channel table defining valid channel names and
123 model options.
124 store: Persistent backend that owns the current snapshot.
125 authorizer: Permission gate for ``relay.*`` actions. When
126 ``None`` no permission check is performed (development).
127 audit: Audit backend for mutation events. When ``None``
128 mutations still apply without audit emission.
129 streams: Registry of in-flight upstream streams. When
130 ``None`` a private empty registry is created; share one
131 instance with the streaming path so force-cancel reaches
132 live streams.
133 """
134 self._registry = registry
135 self._store = store
136 self._authorizer = authorizer
137 self._audit = audit
138 self.streams = streams if streams is not None else RelayStreamRegistry()
139 self._lock = asyncio.Lock()
141 async def set_channel_state(
142 self,
143 channel: str,
144 enabled: bool,
145 actor_id: str,
146 ) -> None:
147 """Enable or drain *channel* for new requests.
149 Args:
150 channel: Channel name; unknown names are rejected.
151 enabled: ``False`` drains the channel for new requests while
152 existing streams finish.
153 actor_id: Operator identity recorded in the audit event.
155 Raises:
156 ValueError: The channel is unknown.
157 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor
158 lacks ``relay.channel_control``.
159 """
160 await self._require(actor_id, PERMISSION_CHANNEL_CONTROL, f"channel:{channel}")
161 async with self._lock:
162 snapshot = await self._store.load()
163 if channel not in snapshot.enabled_channels:
164 raise ValueError(f"unknown channel {channel!r}")
165 changed = replace(
166 snapshot,
167 enabled_channels={**snapshot.enabled_channels, channel: enabled},
168 )
169 self._require_router_survives(changed)
170 await self._store.save(changed)
171 self._registry.set_runtime_enabled(channel, enabled)
172 await self._audit_change(
173 actor_id=actor_id,
174 action=PERMISSION_CHANNEL_CONTROL,
175 resource=channel,
176 old={"enabled": not enabled},
177 new={"enabled": enabled},
178 )
180 async def update_policy(
181 self,
182 change: RelayPolicyChange,
183 actor_id: str,
184 ) -> None:
185 """Apply a typed policy change.
187 Args:
188 change: Partial mutation; only the fields explicitly set
189 change.
190 actor_id: Operator identity recorded in the audit event.
192 Raises:
193 ValueError: The change references an unknown channel or
194 model option values, or would remove every available
195 model option.
196 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor
197 lacks ``relay.policy_control``.
198 """
199 await self._require(actor_id, PERMISSION_POLICY_CONTROL, "policy")
200 async with self._lock:
201 snapshot = await self._store.load()
202 if change.channel is not None:
203 if change.channel not in snapshot.enabled_channels:
204 raise ValueError(f"unknown channel {change.channel!r}")
205 if change.enabled is not None and change.channel is None:
206 raise ValueError("channel is required when changing enabled state")
207 if change.allowed_model_options is not None and change.channel is None:
208 raise ValueError(
209 "channel is required when changing allowed model options"
210 )
211 self._validate_options(change)
212 changed = self._compose(snapshot, change)
213 self._require_router_survives(changed)
214 await self._store.save(changed)
215 if change.enabled is not None and change.channel is not None:
216 self._registry.set_runtime_enabled(change.channel, change.enabled)
217 await self._audit_change(
218 actor_id=actor_id,
219 action=PERMISSION_POLICY_CONTROL,
220 resource="policy",
221 old=self._changed_fields(snapshot, change),
222 new=self._changed_fields(changed, change),
223 )
225 async def policy_snapshot(self, actor_id: str) -> RelayPolicySnapshot:
226 """Return the current runtime policy snapshot.
228 Args:
229 actor_id: Operator identity; ``relay.read`` permission is
230 required.
232 Returns:
233 The snapshot persisted by the policy store.
235 Raises:
236 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor
237 lacks ``relay.read``.
238 """
239 await self._require(actor_id, PERMISSION_READ, "policy")
240 return await self._store.load()
242 def active_streams(self) -> tuple[RelayActiveStream, ...]:
243 """Return the currently in-flight upstream streams.
245 Returns:
246 One row per active stream, oldest first; an empty tuple when
247 no stream is in flight.
248 """
249 return self.streams.list()
251 async def force_cancel_stream(
252 self,
253 stream_id: str,
254 actor_id: str,
255 ) -> None:
256 """Force-cancel an in-flight upstream stream.
258 Args:
259 stream_id: Identifier of the stream to cancel.
260 actor_id: Operator identity recorded in the audit event;
261 ``relay.stream_control`` permission is required.
263 Raises:
264 ValueError: The stream identifier is unknown.
265 RelayGatewayError: With ``PERMISSION_DENIED`` when the actor
266 lacks ``relay.stream_control``.
267 """
268 await self._require(actor_id, PERMISSION_STREAM_CONTROL, f"stream:{stream_id}")
269 if not self.streams.cancel(stream_id):
270 raise ValueError(f"unknown stream {stream_id!r}")
271 await self._audit_change(
272 actor_id=actor_id,
273 action=PERMISSION_STREAM_CONTROL,
274 resource=stream_id,
275 old={"cancelled": False},
276 new={"cancelled": True},
277 )
279 def _validate_options(self, change: RelayPolicyChange) -> None:
280 """Reject option names the target channel does not serve."""
281 options = change.allowed_model_options
282 if options is None or change.channel is None:
283 return
284 channel = self._channel(change.channel)
285 if channel is None:
286 raise ValueError(f"unknown channel {change.channel!r}")
287 unknown = options - set(channel.models)
288 if unknown:
289 name = sorted(unknown)[0]
290 raise ValueError(
291 f"unknown model option {name!r} for channel {change.channel!r}"
292 )
294 def _channel(self, name: str):
295 """Return the configured channel with *name*, or ``None``."""
296 for candidate in self._registry.channels:
297 if candidate.name == name:
298 return candidate
299 return None
301 async def _require(self, actor_id: str, action: str, resource: str) -> None:
302 """Enforce *action* permission, raising ``PERMISSION_DENIED``."""
303 allowed = True
304 if self._authorizer is not None:
305 allowed = await self._authorizer.can(actor_id, action, resource)
306 if not allowed:
307 raise RelayGatewayError(
308 code="PERMISSION_DENIED",
309 message=f"{action} denied for the operator",
310 status_code=403,
311 request_id="",
312 )
314 @staticmethod
315 def _compose(
316 snapshot: RelayPolicySnapshot,
317 change: RelayPolicyChange,
318 ) -> RelayPolicySnapshot:
319 """Build the snapshot after *change* is applied."""
320 enabled_channels = snapshot.enabled_channels
321 allowed_model_options = snapshot.allowed_model_options
322 if change.enabled is not None or change.allowed_model_options is not None:
323 channel = change.channel
324 if channel is None:
325 raise ValueError(
326 "channel is required for channel-scoped policy changes"
327 )
328 if change.enabled is not None:
329 enabled_channels = {**enabled_channels, channel: change.enabled}
330 if change.allowed_model_options is not None:
331 allowed_model_options = {
332 **allowed_model_options,
333 channel: change.allowed_model_options,
334 }
335 return replace(
336 snapshot,
337 enabled_channels=enabled_channels,
338 allowed_model_options=allowed_model_options,
339 media_allowed_schemes=(
340 change.media_allowed_schemes
341 if change.media_allowed_schemes is not None
342 else snapshot.media_allowed_schemes
343 ),
344 media_allowed_hosts=(
345 change.media_allowed_hosts
346 if change.media_allowed_hosts is not None
347 else snapshot.media_allowed_hosts
348 ),
349 max_request_bytes=(
350 change.max_request_bytes
351 if change.max_request_bytes is not None
352 else snapshot.max_request_bytes
353 ),
354 max_stream_seconds=(
355 change.max_stream_seconds
356 if change.max_stream_seconds is not None
357 else snapshot.max_stream_seconds
358 ),
359 )
361 @staticmethod
362 def _require_router_survives(snapshot: RelayPolicySnapshot) -> None:
363 """Reject snapshots with no enabled channel serving options.
365 Raises:
366 ValueError: When every channel is disabled or every
367 remaining channel has no allowed model options.
368 """
369 for name, options in snapshot.allowed_model_options.items():
370 if snapshot.enabled_channels.get(name, False) and options:
371 return
372 raise ValueError("policy change would remove all available converters")
374 @staticmethod
375 def _changed_fields(
376 snapshot: RelayPolicySnapshot,
377 change: RelayPolicyChange,
378 ) -> dict[str, Any]:
379 """Extract the values changed by *change* from *snapshot*."""
380 fields: dict[str, Any] = {}
381 if change.enabled is not None and change.channel is not None:
382 fields[f"{change.channel}.enabled"] = snapshot.enabled_channels[
383 change.channel
384 ]
385 if change.allowed_model_options is not None and change.channel is not None:
386 fields[f"{change.channel}.allowed_model_options"] = (
387 snapshot.allowed_model_options.get(change.channel, frozenset())
388 )
389 if change.media_allowed_schemes is not None:
390 fields["media_allowed_schemes"] = snapshot.media_allowed_schemes
391 if change.media_allowed_hosts is not None:
392 fields["media_allowed_hosts"] = snapshot.media_allowed_hosts
393 if change.max_request_bytes is not None:
394 fields["max_request_bytes"] = snapshot.max_request_bytes
395 if change.max_stream_seconds is not None:
396 fields["max_stream_seconds"] = snapshot.max_stream_seconds
397 return fields
399 async def _audit_change(
400 self,
401 actor_id: str,
402 action: str,
403 resource: str,
404 old: dict[str, Any],
405 new: dict[str, Any],
406 ) -> None:
407 """Persist one audit event for an applied mutation."""
408 if self._audit is None:
409 return
410 await self._audit.record(
411 AIAuditEvent(
412 event_type=AuditEventType.CONFIG_RELOADED,
413 user_id=actor_id,
414 status="success",
415 metadata={
416 "action": action,
417 "resource": resource,
418 "old": old,
419 "new": new,
420 },
421 )
422 )
423 logger.debug(
424 "relay_controls_mutated",
425 action=action,
426 resource=resource,
427 )