zeroae / zeroae/zae-limiter

✨ Add WebSocket rate limiting support for streaming scenarios

Open
#257 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

area/limiter
Dominant language
Python
Stars
0
Forks
0
Avg merge
6h 51m
Merged PRs (30d)
104

Description

Problem or Use Case

LLM chat applications commonly use WebSocket connections for streaming responses. The current acquire() pattern works well for HTTP requests where:

  1. A single request has a bounded lifecycle
  2. Token consumption is reconciled once at the end

WebSocket scenarios have different characteristics:

  • Long-lived connections - A single connection may span many message exchanges
  • Bidirectional streaming - Both client messages (prompts) and server messages (tokens) need tracking
  • Per-message rate limiting - Need to limit messages/tokens across the connection lifetime
  • Connection-level budgets - May want to cap total tokens per session/connection
  • Graceful degradation - When rate limited, may want to slow down rather than disconnect

Current workaround requires wrapping each message in a separate acquire() context, which:

  • Creates overhead for high-frequency messages
  • Doesn't provide connection-level budget tracking
  • Makes it awkward to implement backpressure

Proposed Solution

Add WebSocket-aware rate limiting primitives:

# Connection-level budget with per-message tracking
async with limiter.websocket_session(
    entity_id="user-123",
    resource="gpt-4",
    limits=[
        Limit.per_minute("messages", 60),     # Message rate
        Limit.per_minute("tpm", 10_000),      # Token rate
        Limit.per_connection("total", 50_000), # Session budget
    ],
) as session:
    async for message in websocket:
        # Lightweight per-message check (no DynamoDB call if cached)
        await session.check_rate("messages", 1)
        
        # Process and stream response
        tokens_used = 0
        async for chunk in llm_stream(message):
            tokens_used += chunk.token_count
            await websocket.send(chunk)
        
        # Reconcile tokens at end of response
        await session.consume("tpm", tokens_used)
        await session.consume("total", tokens_used)
Alternative: Message-level helper

For simpler cases, a decorator/context manager for individual messages:

@limiter.websocket_message(
    entity_id_header="X-API-Key",
    limits=[Limit.per_minute("messages", 60)],
)
async def on_message(websocket, message, lease):
    # lease automatically committed on success
    response = await process(message)
    await lease.adjust(tokens=response.token_count)
    return response

Alternatives Considered

  1. Status quo - Wrap each message in acquire(). Works but inefficient and lacks connection-level features.

  2. External connection tracking - Track connections separately and combine with per-message rate limiting. More complex integration.

  3. Framework-specific only - Only support FastAPI WebSocket. Limits reusability.

Acceptance Criteria

  • websocket_session() context manager exists in RateLimiter class
  • Session supports check_rate() for lightweight in-memory checks between DynamoDB syncs
  • Session supports consume() for recording usage with periodic DynamoDB flush
  • Session supports Limit.per_connection() for session-scoped budgets
  • Graceful handling when session budget exceeded (configurable: disconnect vs backpressure)
  • Unit tests cover session lifecycle (create, consume, close, timeout)
  • Integration tests verify DynamoDB state after session completion
  • Documentation in docs/guide/ covers WebSocket patterns with code examples
  • FastAPI WebSocket example added to examples/fastapi-demo/

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the RateLimiter class and its existing acquire() path, then review the proposed session lifecycle and DynamoDB synchronization requirements. Use the unit and integration tests described in the acceptance criteria, and inspect docs/guide/ and examples/fastapi-demo/ for the required documentation and example coverage. Done means all listed session, budget, graceful-handling, test, documentation, and example criteria are met.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, fastapi, python
Domain
api, backend, cloud, databases
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.