mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-05 20:14:47 +00:00
- Per-user SSE cap uses per-connection leases in a sorted set instead of a shared INCR/DECR counter with a TTL. A stream that outlived the TTL could let the counter expire, then decrement another stream's slot or drive the count negative. Leases refresh while a stream sends frames and age out when a stream dies without cleanup. - Shielded cleanup is bounded per step: stream close and on_close in ClosingStreamingResponse, unsubscribe and close in AsyncTopic, lease release, and the replay-budget check, so a dead Redis connection can't hold a request or a graceful shutdown. - Artifact downloads send an ASCII filename plus an RFC 5987 filename* for non-ASCII names, build headers before opening the file, and close disk handles off the event loop. - Type hints and docstrings on the new helpers.
53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
"""Stream/topic key derivations shared by publisher and SSE consumer.
|
|
|
|
Single source of truth for the per-user Redis Streams key and pub/sub
|
|
topic name. Both must agree exactly — a typo here splits the
|
|
publisher's writes from the consumer's reads.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
def stream_key(user_id: str) -> str:
|
|
"""Redis Streams key holding the durable backlog for ``user_id``."""
|
|
return f"user:{user_id}:stream"
|
|
|
|
|
|
def topic_name(user_id: str) -> str:
|
|
"""Redis pub/sub channel used for live fan-out to ``user_id``."""
|
|
return f"user:{user_id}"
|
|
|
|
|
|
def connection_leases_key(user_id: str) -> str:
|
|
"""Redis sorted set of live SSE connection leases for ``user_id``."""
|
|
return f"user:{user_id}:sse_leases"
|
|
|
|
|
|
def replay_budget_key(user_id: str) -> str:
|
|
"""Redis counter tracking snapshot replays for ``user_id`` in the
|
|
rolling rate-limit window."""
|
|
return f"user:{user_id}:replay_count"
|
|
|
|
|
|
def stream_id_compare(a: str, b: str) -> int:
|
|
"""Compare two Redis Streams ids. Returns -1, 0, 1 like ``cmp``.
|
|
|
|
Stream ids are ``ms-seq`` strings; comparing as strings would be wrong
|
|
once ``ms`` straddles digit-count boundaries. We parse and compare
|
|
as ``(int, int)`` tuples.
|
|
|
|
Raises ``ValueError`` on malformed input. Callers must pre-validate
|
|
against ``_STREAM_ID_RE`` (or equivalent) — a lex fallback here let
|
|
a malformed id compare lex-greater than a real one and silently pin
|
|
dedup forever.
|
|
"""
|
|
a_ms, _, a_seq = a.partition("-")
|
|
b_ms, _, b_seq = b.partition("-")
|
|
a_tuple = (int(a_ms), int(a_seq) if a_seq else 0)
|
|
b_tuple = (int(b_ms), int(b_seq) if b_seq else 0)
|
|
if a_tuple < b_tuple:
|
|
return -1
|
|
if a_tuple > b_tuple:
|
|
return 1
|
|
return 0
|