vllm-omni/vllm_omni/entrypoints/duplex/session_attachment.py

# 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",
]