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 )