LLM Streaming Session Pinning for WebSocket AI Gateways

Problem

Modern AI applications increasingly rely on real-time streaming responses to deliver tokens progressively to users. This pattern is common in:

• conversational assistants
• copilots
• agent-based systems
• chat applications powered by LLM APIs

These interactions frequently run over long-lived HTTP or WebSocket connections. Traditional load balancing distributes requests across multiple backend nodes. While this works for stateless workloads, it can cause issues for streaming AI inference, where the interaction often maintains temporary state within the inference gateway or middleware.

If traffic from the same conversation is routed to different backend nodes, several problems can occur:

• broken streaming responses
• loss of conversational continuity
• inconsistent token latency
• reconnection errors in WebSocket sessions
• degraded user experience

In AI applications, the critical unit is not just the request — it is the session or conversation. A delivery layer capable of maintaining session affinity for streaming AI workloads is therefore essential.

Solution

This iRule introduces session pinning for AI streaming traffic at the BIG-IP layer. The rule detects streaming or WebSocket upgrade requests and extracts a session or conversation identifier from incoming traffic. Using this identifier, the iRule applies universal persistence so that all requests belonging to the same conversation remain pinned to the same backend node.

The rule performs the following functions:

• Detects WebSocket upgrade requests or streaming endpoints
• Extracts a Session ID or Conversation ID
• Applies universal persistence based on that identifier
• Inserts observability headers for debugging and telemetry
• Logs session-to-node mapping for operational visibility
• Supported session identifiers may include:
• X-Session-ID
• X-Conversation-ID
• Sec-WebSocket-Key
• API keys
• client IP fallback

By implementing persistence at the application delivery layer, BIG-IP ensures that multi-turn AI interactions remain consistent throughout the entire streaming session.

Impact

This solution enhances the reliability and scalability of AI infrastructure by ensuring stable routing for real-time inference workloads.

Key benefits include:

• Improved User Experience
• Streaming responses remain uninterrupted and consistent during long-lived conversations.
• Session Consistency
• Multi-turn interactions stay pinned to the same inference gateway or middleware node.
• Operational Stability
• Prevents backend errors caused by mid-stream node changes.
• AI Infrastructure Optimization
• Enables load-balanced AI clusters while preserving conversational state.
• Observability
• Provides logging and header-based telemetry for troubleshooting session routing.

This approach demonstrates how BIG-IP can function as an AI-aware traffic control layer, managing not only connectivity but also the behavior of real-time AI application flows.

Code

when HTTP_REQUEST {

    # Detect AI streaming or websocket endpoints
    if { [HTTP::path] starts_with "/ws/" or 
         [HTTP::path] starts_with "/chat" or
         [HTTP::path] starts_with "/v1/stream" } {

        # Attempt to retrieve conversation identifier
        set conversation_id [HTTP::header value "X-Conversation-ID"]

        # Fallback to session ID header
        if { $conversation_id eq "" } {
            set conversation_id [HTTP::header value "X-Session-ID"]
        }

        # If WebSocket handshake exists use websocket key
        if { $conversation_id eq "" && [HTTP::header exists "Sec-WebSocket-Key"] } {
            set conversation_id [HTTP::header value "Sec-WebSocket-Key"]
        }

        # Fallback to API key
        if { $conversation_id eq "" && [HTTP::header exists "X-API-Key"] } {
            set conversation_id [HTTP::header value "X-API-Key"]
        }

        # Final fallback: client IP
        if { $conversation_id eq "" } {
            set conversation_id [IP::client_addr]
        }

        # Apply universal persistence for session pinning
        persist uie $conversation_id 1800

        # Observability headers
        HTTP::header insert "X-AI-Session-Pinning" "enabled"
        HTTP::header insert "X-AI-Conversation-ID" $conversation_id

        log local0. "AI_STREAM_PIN session=$conversation_id uri=[HTTP::uri] client=[IP::client_addr]"
    }
}