Coverage for src / lexigram / ai / relay / gateway / operations / streams.py: 100%

1"""In-flight stream session registry for the relay gateway. 

2 

3``RelayStreamRegistry`` is the single process-local source of truth for 

4active upstream streams. A stream registers once at start and 

5unregisters when its relay loop ends; an operator force-cancel sets the 

6stream's cancel handle, which the relay loop observes as a truncated 

7termination. Registry mutations are synchronous and therefore atomic on 

8the event loop. 

9""" 

10 

11from __future__ import annotations 

12 

13import asyncio 

14from uuid import uuid4 

15 

16from lexigram.contracts.ai.relay import RelayActiveStream 

17from lexigram.primitives import clock 

18 

19__all__ = ["RelayStreamRegistry"] 

20 

21 

22class RelayStreamRegistry: 

23 """Tracks active streams and their cancel handles. 

24 

25 Attributes: 

26 _active: Stream identifier to stream metadata, oldest first by 

27 insertion order. 

28 _handles: Stream identifier to its cancel handle. 

29 """ 

30 

31 def __init__(self) -> None: 

32 """Create an empty registry.""" 

33 self._active: dict[str, RelayActiveStream] = {} 

34 self._handles: dict[str, asyncio.Event] = {} 

35 

36 def register( 

37 self, 

38 *, 

39 channel: str, 

40 model: str, 

41 request_id: str, 

42 ) -> tuple[str, asyncio.Event]: 

43 """Register a new in-flight stream. 

44 

45 Args: 

46 channel: Channel name serving the stream. 

47 model: Outbound model alias of the stream. 

48 request_id: Gateway request identifier. 

49 

50 Returns: 

51 The new stream identifier and its cancel handle; setting 

52 the handle asks the relay loop to terminate truncated. 

53 """ 

54 stream_id = uuid4().hex 

55 self._active[stream_id] = RelayActiveStream( 

56 stream_id=stream_id, 

57 channel=channel, 

58 model=model, 

59 request_id=request_id, 

60 started_at=clock.now(), 

61 ) 

62 handle = asyncio.Event() 

63 self._handles[stream_id] = handle 

64 return stream_id, handle 

65 

66 def unregister(self, stream_id: str) -> None: 

67 """Forget a finished stream and its handle. 

68 

69 Args: 

70 stream_id: Identifier previously returned by ``register``. 

71 """ 

72 self._active.pop(stream_id, None) 

73 self._handles.pop(stream_id, None) 

74 

75 def list(self) -> tuple[RelayActiveStream, ...]: 

76 """Return active streams, oldest first. 

77 

78 Returns: 

79 A tuple of active stream rows; empty when nothing is 

80 in flight. 

81 """ 

82 return tuple(self._active.values()) 

83 

84 def handle(self, stream_id: str) -> asyncio.Event | None: 

85 """Return the cancel handle of *stream_id*, or ``None``. 

86 

87 Args: 

88 stream_id: Stream identifier. 

89 

90 Returns: 

91 The cancel handle when the stream is active, else ``None``. 

92 """ 

93 return self._handles.get(stream_id) 

94 

95 def cancel(self, stream_id: str) -> bool: 

96 """Request cancellation of *stream_id*. 

97 

98 Args: 

99 stream_id: Stream identifier. 

100 

101 Returns: 

102 ``True`` when the stream was active and its handle was set; 

103 ``False`` when the stream is unknown. 

104 """ 

105 handle = self._handles.get(stream_id) 

106 if handle is None: 

107 return False 

108 handle.set() 

109 return True