Repository: PrefectHQ/prefect
File: src/prefect/_internal/websockets.py (Python)

Code before:


def create_ssl_context_for_websocket(uri: str) -> Optional[ssl.SSLContext]:
    """Create SSL context for WebSocket connections based on URI scheme."""
    u = urlparse(uri)

    if u.scheme != "wss":
        return None

    if get_current_settings().api.tls_insecure_skip_verify:
        # Create an unverified context for insecure connections
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
        return ctx
    else:
        # Create a verified context with the certificate file
        cert_file = get_current_settings().api.ssl_cert_file
        if not cert_file:
            cert_file = certifi.where()
        return ssl.create_default_context(cafile=cert_file)


def websocket_connect(uri: str, **kwargs: Any) -> connect:
    """
    Create a WebSocket connection with proxy and SSL support.

    Proxy support is automatic via HTTP_PROXY/HTTPS_PROXY environment variables.
    The websockets library handles proxy detection and connection automatically.
    """
    # Configure SSL context for HTTPS connections
    ssl_context = create_ssl_context_for_websocket(uri)
    if ssl_context:
        kwargs.setdefault("ssl", ssl_context)

    # Add custom headers from settings
    custom_headers = get_current_settings().client.custom_headers
    if custom_headers:
        # Get existing additional_headers or create new dict
        additional_headers = kwargs.get("additional_headers", {})
        if not isinstance(additional_headers, dict):
            additional_headers = {}

        for header_name, header_value in custom_headers.items():
            # Check for protected headers that shouldn't be overridden
            if header_name.lower() in {
                "user-agent",
                "sec-websocket-key",
                "sec-websocket-version",
                "sec-websocket-extensions",
                "sec-websocket-protocol",

Review comment (suggestion):
Can you add a `@wraps` decorator to preserve typing information for the kwargs?
```suggestion
@wraps(connect)
def websocket_connect(uri: str, **kwargs: Any) -> connect:
```

Code after:
from prefect.settings import get_current_settings


def create_ssl_context_for_websocket(uri: str) -> Optional[ssl.SSLContext]:
    """Create SSL context for WebSocket connections based on URI scheme."""
    u = urlparse(uri)

    if u.scheme != "wss":
        return None

    if get_current_settings().api.tls_insecure_skip_verify:
        # Create an unverified context for insecure connections
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
        return ctx
    else:
        # Create a verified context with the certificate file
        cert_file = get_current_settings().api.ssl_cert_file
        if not cert_file:
            cert_file = certifi.where()
        return ssl.create_default_context(cafile=cert_file)


@wraps(connect)
def websocket_connect(uri: str, **kwargs: Any) -> connect:
    """
    Create a WebSocket connection with proxy and SSL support.

    Proxy support is automatic via HTTP_PROXY/HTTPS_PROXY environment variables.
    The websockets library handles proxy detection and connection automatically.
    """
    # Configure SSL context for HTTPS connections
    ssl_context = create_ssl_context_for_websocket(uri)
    if ssl_context:
        kwargs.setdefault("ssl", ssl_context)

    # Add custom headers from settings
    custom_headers = get_current_settings().client.custom_headers
    if custom_headers:
        # Get existing additional_headers or create new dict
        additional_headers = kwargs.get("additional_headers", {})
        if not isinstance(additional_headers, dict):
            additional_headers = {}

        for header_name, header_value in custom_headers.items():
            # Check for protected headers that shouldn't be overridden
            if header_name.lower() in {
                "user-agent",
                "sec-websocket-key",
                "sec-websocket-version",