# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import asyncio import hashlib import hmac import json import secrets import time from collections import deque from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field from types import MappingProxyType class InvalidResumeTokenError(RuntimeError): pass class DuplexJournalGapError(RuntimeError): pass class DuplexJournalOverflowError(RuntimeError): pass @dataclass(frozen=True) class ResumeToken: plaintext: str = field(repr=False) @classmethod def generate(cls) -> ResumeToken: return cls(secrets.token_urlsafe(32)) @dataclass class DuplexResumeCredential: token_digest: bytes @classmethod def from_token(cls, token: ResumeToken) -> DuplexResumeCredential: return cls(token_digest=cls._digest(token.plaintext)) @staticmethod def _digest(plaintext: str) -> bytes: return hashlib.sha256(plaintext.encode("utf-8")).digest() def verify(self, plaintext: str) -> bool: return hmac.compare_digest(self.token_digest, self._digest(plaintext)) def rotate(self) -> ResumeToken: token = ResumeToken.generate() self.token_digest = self._digest(token.plaintext) return token @dataclass(frozen=True) class DuplexTransportAttachment: generation: int send: Callable[[dict[str, object]], Awaitable[None]] = field(repr=False) close: Callable[[str], Awaitable[None]] = field(repr=False) @dataclass(frozen=True) class JournalEntry: sequence: int created_monotonic: float encoded_bytes: int payload: Mapping[str, object] = field(repr=False) def __post_init__(self) -> None: object.__setattr__(self, "payload", MappingProxyType(dict(self.payload))) class DuplexEventJournal: def __init__( self, *, max_bytes: int, ttl_s: float, clock: Callable[[], float] | None = None, ) -> None: if max_bytes <= 0: raise ValueError("journal max_bytes must be positive") if ttl_s <= 0: raise ValueError("journal ttl_s must be positive") self._max_bytes = max_bytes self._ttl_s = ttl_s self._clock = clock or time.monotonic self._entries: deque[JournalEntry] = deque() self._next_sequence = 1 self._dropped_through = 0 self._retained_bytes = 0 self._overflowed = False @property def retained_bytes(self) -> int: return self._retained_bytes @property def overflowed(self) -> bool: return self._overflowed @property def last_sequence(self) -> int: return self._next_sequence - 1 def record(self, payload: Mapping[str, object]) -> JournalEntry: if self._overflowed: raise DuplexJournalOverflowError("duplex event journal already exceeded its byte limit") self.prune() sequence = self._next_sequence sequenced_payload = dict(payload) sequenced_payload["server_event_seq"] = sequence encoded_bytes = len( json.dumps( sequenced_payload, separators=(",", ":"), ensure_ascii=False, ).encode("utf-8") ) if self._retained_bytes + encoded_bytes > self._max_bytes: self._overflowed = True raise DuplexJournalOverflowError( f"duplex event journal byte limit exceeded: {self._retained_bytes + encoded_bytes} > {self._max_bytes}" ) entry = JournalEntry( sequence=sequence, created_monotonic=self._clock(), encoded_bytes=encoded_bytes, payload=sequenced_payload, ) self._entries.append(entry) self._retained_bytes += encoded_bytes self._next_sequence += 1 return entry def acknowledge(self, sequence: int) -> int: if sequence < 0: raise ValueError("acknowledged sequence must not be negative") if sequence > self.last_sequence: raise ValueError(f"acknowledged sequence {sequence} is newer than journal head {self.last_sequence}") removed = 0 while self._entries and self._entries[0].sequence <= sequence: entry = self._entries.popleft() self._retained_bytes -= entry.encoded_bytes removed += 1 self._dropped_through = max(self._dropped_through, sequence) return removed def prune(self, now: float | None = None) -> int: effective_now = self._clock() if now is None else now cutoff = effective_now - self._ttl_s removed = 0 while self._entries and self._entries[0].created_monotonic <= cutoff: entry = self._entries.popleft() self._retained_bytes -= entry.encoded_bytes self._dropped_through = max(self._dropped_through, entry.sequence) removed += 1 return removed def replay_after(self, sequence: int) -> tuple[JournalEntry, ...]: if sequence < 0: raise ValueError("replay sequence must not be negative") self.prune() if self._overflowed: raise DuplexJournalGapError("duplex event journal overflowed; replay is incomplete") if sequence < self._dropped_through: raise DuplexJournalGapError( f"requested sequence {sequence} is older than retained journal boundary {self._dropped_through}" ) if sequence > self.last_sequence: raise ValueError(f"requested sequence {sequence} is newer than journal head {self.last_sequence}") return tuple(entry for entry in self._entries if entry.sequence > sequence) @dataclass(frozen=True) class DuplexSessionAttachmentCreated: session_id: str attachment_generation: int resume_token: ResumeToken = field(repr=False) @dataclass(frozen=True) class DuplexSessionResumeResult: session_id: str attachment_generation: int resume_token: ResumeToken = field(repr=False) replay_entries: tuple[JournalEntry, ...] = () replaced_attachment: DuplexTransportAttachment | None = None @dataclass class _DuplexSessionAttachmentState: session_id: str credential: DuplexResumeCredential journal: DuplexEventJournal attachment: DuplexTransportAttachment | None attachment_generation: int outbound_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) recovery_token_digest: bytes | None = field(default=None, repr=False) class DuplexSessionAttachmentRegistry: def __init__( self, *, replay_ttl_s: float, replay_max_bytes_per_session: int, clock: Callable[[], float] | None = None, ) -> None: if replay_ttl_s <= 0: raise ValueError("replay_ttl_s must be positive") if replay_max_bytes_per_session <= 0: raise ValueError("replay_max_bytes_per_session must be positive") self._replay_ttl_s = replay_ttl_s self._replay_max_bytes_per_session = replay_max_bytes_per_session self._clock = clock or time.monotonic self._sessions: dict[str, _DuplexSessionAttachmentState] = {} self._lock = asyncio.Lock() def __repr__(self) -> str: return f"{type(self).__name__}(session_ids={sorted(self._sessions)})" async def create( self, session_id: str, *, send: Callable[[dict[str, object]], Awaitable[None]], close: Callable[[str], Awaitable[None]], ) -> DuplexSessionAttachmentCreated: async with self._lock: if session_id in self._sessions: raise ValueError(f"duplex attachment session already exists: {session_id}") token = ResumeToken.generate() generation = 1 self._sessions[session_id] = _DuplexSessionAttachmentState( session_id=session_id, credential=DuplexResumeCredential.from_token(token), journal=DuplexEventJournal( max_bytes=self._replay_max_bytes_per_session, ttl_s=self._replay_ttl_s, clock=self._clock, ), attachment=DuplexTransportAttachment( generation=generation, send=send, close=close, ), attachment_generation=generation, ) return DuplexSessionAttachmentCreated( session_id=session_id, attachment_generation=generation, resume_token=token, ) async def send_event( self, session_id: str, payload: Mapping[str, object], *, journal: bool = True, on_accepted: Callable[[], None] | None = None, ) -> JournalEntry | None: """Sequence and dispatch one event to the current attachment. The per-session lock keeps wire order equal to journal order without serializing unrelated sessions. A detached session still records replayable events, but has no transport side effect. ``on_accepted`` runs synchronously once the event is journaled, or after a successful transport send when journaling is disabled. A later send failure cannot undo acceptance into the journal. Detached, non-journaled events do not invoke it. The callback must not raise. """ async with self._lock: state = self._require(session_id) async with state.outbound_lock: async with self._lock: if self._sessions.get(session_id) is not state: raise KeyError(f"unknown duplex attachment session: {session_id}") entry = state.journal.record(payload) if journal else None attachment = state.attachment wire_payload = dict(entry.payload) if entry is not None else dict(payload) if entry is not None and on_accepted is not None: on_accepted() if attachment is not None: await attachment.send(wire_payload) if entry is None and on_accepted is not None: on_accepted() return entry async def acknowledge(self, session_id: str, sequence: int) -> int: async with self._lock: return self._require(session_id).journal.acknowledge(sequence) async def detach(self, session_id: str, *, attachment_generation: int | None = None) -> bool: """Drop the current transport; the engine lease owns the disconnect grace. ``attachment_generation`` names the connection asking to detach, so a socket that already lost a takeover cannot detach the winner. ``None`` means "whichever connection is attached right now" and is for callers that only know the session (the outbound pump, whose send just failed). Returns whether this call is the one that detached: an already-detached session answers ``False`` so a second disconnect signal for the same socket cannot restart the engine's disconnect grace window. """ async with self._lock: state = self._sessions.get(session_id) if state is None or state.attachment is None: return False if attachment_generation is not None and state.attachment_generation != attachment_generation: return False state.attachment = None return True async def has_attachment(self, session_id: str) -> bool: """Whether some connection is attached right now, whoever it is. A resume that failed to activate has no generation of its own, so it cannot ask ``is_current_attachment``. What it needs to know before rolling the engine lease back into its disconnect grace is only whether it would be rolling back somebody else's live attachment. """ async with self._lock: state = self._sessions.get(session_id) return state is not None and state.attachment is not None async def is_current_attachment(self, session_id: str, attachment_generation: int) -> bool: async with self._lock: state = self._sessions.get(session_id) return ( state is not None and state.attachment is not None and state.attachment_generation == attachment_generation ) async def authenticate_resume( self, session_id: str, *, resume_token: str, last_received_server_event_seq: int, ) -> None: """Validate transport credentials before any engine resume control.""" async with self._lock: state = self._require(session_id) self._validate_resume_identity(state, resume_token=resume_token) state.journal.replay_after(last_received_server_event_seq) async def resume( self, session_id: str, *, resume_token: str, last_received_server_event_seq: int, send: Callable[[dict[str, object]], Awaitable[None]], close: Callable[[str], Awaitable[None]], activation_payload_factory: Callable[[ResumeToken, int], Mapping[str, object]] | None = None, ) -> DuplexSessionResumeResult: async with self._lock: state = self._require(session_id) async with state.outbound_lock: async with self._lock: if self._sessions.get(session_id) is not state: raise KeyError(f"unknown duplex attachment session: {session_id}") used_recovery = self._validate_resume_identity(state, resume_token=resume_token) replay_entries = state.journal.replay_after(last_received_server_event_seq) accepted_token_digest = ( state.recovery_token_digest if used_recovery else bytes(state.credential.token_digest) ) state.recovery_token_digest = None rotated_token = state.credential.rotate() replaced = state.attachment state.attachment_generation += 1 attachment_generation = state.attachment_generation state.attachment = DuplexTransportAttachment( generation=attachment_generation, send=send, close=close, ) if activation_payload_factory is not None: try: await send(dict(activation_payload_factory(rotated_token, attachment_generation))) for entry in replay_entries: await send(dict(entry.payload)) except (Exception, asyncio.CancelledError): # CancelledError derives from BaseException, so an # ``except Exception`` here would let a cancellation during # activation or replay skip the rollback: the dead # attachment would stay current while the token it rotated # past is gone, leaving the session attached to an unusable # transport with no way back. Both failures take the same # path deliberately, so they cannot drift apart. await self._rollback_resume( session_id, state=state, attachment_generation=attachment_generation, accepted_token_digest=accepted_token_digest, ) raise return DuplexSessionResumeResult( session_id=session_id, attachment_generation=attachment_generation, resume_token=rotated_token, replay_entries=replay_entries, replaced_attachment=replaced, ) async def _rollback_resume( self, session_id: str, *, state: _DuplexSessionAttachmentState, attachment_generation: int, accepted_token_digest: bytes, ) -> None: """Undo a resume whose activation never reached the client. Generation-checked: a later resume that already took over must keep its attachment. Restoring ``recovery_token_digest`` is what lets the caller retry with the token it presented, which ``resume`` had rotated past. """ async with self._lock: if self._sessions.get(session_id) is state and state.attachment_generation == attachment_generation: state.attachment = None state.recovery_token_digest = accepted_token_digest async def close(self, session_id: str) -> DuplexTransportAttachment | None: async with self._lock: state = self._sessions.pop(session_id, None) return state.attachment if state is not None else None def _require(self, session_id: str) -> _DuplexSessionAttachmentState: state = self._sessions.get(session_id) if state is None: raise KeyError(f"unknown duplex attachment session: {session_id}") return state @staticmethod def _validate_resume_identity(state: _DuplexSessionAttachmentState, *, resume_token: str) -> bool: if state.credential.verify(resume_token): return False recovery_digest = state.recovery_token_digest if ( state.attachment is None and recovery_digest is not None and hmac.compare_digest(recovery_digest, DuplexResumeCredential._digest(resume_token)) ): return True raise InvalidResumeTokenError(f"invalid resume token for duplex session {state.session_id}") __all__ = [ "DuplexEventJournal", "DuplexJournalGapError", "DuplexJournalOverflowError", "DuplexResumeCredential", "DuplexSessionAttachmentCreated", "DuplexSessionAttachmentRegistry", "DuplexSessionResumeResult", "DuplexTransportAttachment", "InvalidResumeTokenError", "JournalEntry", "ResumeToken", ]