latch-asgi 0.1.3__py3-none-any.whl → 0.3.0__py3-none-any.whl

latch_asgi/asgi_iface.py

@@ -198,3 +198,131 @@ HTTPSendEvent: TypeAlias = (
     | HTTPDisconnectEvent
 )
 HTTPSendCallable: TypeAlias = Callable[[HTTPSendEvent], Awaitable[None]]
+
+
+# >>> Websocket
+@dataclass(frozen=True)
+class WebsocketScope:
+    type: Literal["websocket"]
+    asgi: ASGIVersions
+    http_version: str
+    scheme: str
+    path: str
+    raw_path: bytes
+    query_string: bytes
+    root_path: str
+    headers: Iterable[tuple[bytes, bytes]]
+    client: tuple[str, int] | None
+    server: tuple[str, int | None] | None
+    subprotocols: Iterable[str]
+    extensions: dict[str, object]
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketScope)
+
+
+@dataclass(frozen=True)
+class WebsocketConnectEvent:
+    type: Literal["websocket.connect"]
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketConnectEvent)
+
+
+@dataclass(frozen=True)
+class WebsocketAcceptEvent:
+    type: Literal["websocket.accept"]
+    subprotocol: str | None
+    headers: Iterable[tuple[bytes, bytes]]
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketAcceptEvent)
+
+
+@dataclass(frozen=True)
+class WebsocketReceiveEvent:
+    type: Literal["websocket.receive"]
+    bytes: bytes | None
+    text: str | None
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketReceiveEvent)
+
+
+@dataclass(frozen=True)
+class WebsocketSendEvent:
+    type: Literal["websocket.send"]
+    bytes: bytes | None
+    text: str | None
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketSendEvent)
+
+
+@dataclass(frozen=True)
+class WebsocketResponseStartEvent:
+    type: Literal["websocket.http.response.start"]
+    status: int
+    headers: Iterable[tuple[bytes, bytes]]
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketResponseStartEvent)
+
+
+@dataclass(frozen=True)
+class WebsocketResponseBodyEvent:
+    type: Literal["websocket.http.response.body"]
+    body: bytes
+    more_body: bool
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketResponseBodyEvent)
+
+
+@dataclass(frozen=True)
+class WebsocketDisconnectEvent:
+    type: Literal["websocket.disconnect"]
+    code: int
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketDisconnectEvent)
+
+
+@dataclass(frozen=True)
+class WebsocketCloseEvent:
+    type: Literal["websocket.close"]
+    code: int
+    reason: str | None
+
+    def as_dict(self):
+        return validate(dataclasses.asdict(self), htyping.WebsocketCloseEvent)
+
+
+WebsocketSendEventT: TypeAlias = (
+    WebsocketAcceptEvent
+    | WebsocketSendEvent
+    | WebsocketResponseBodyEvent
+    | WebsocketResponseStartEvent
+    | WebsocketCloseEvent
+)
+WebsocketSendCallable: TypeAlias = Callable[[WebsocketSendEventT], Awaitable[None]]
+
+
+WebsocketReceiveEventT: TypeAlias = (
+    WebsocketConnectEvent | WebsocketReceiveEvent | WebsocketDisconnectEvent
+)
+WebsocketReceiveCallable: TypeAlias = Callable[[], Awaitable[WebsocketReceiveEventT]]
+
+
+WWWScope: TypeAlias = HTTPScope | WebsocketScope
+Scope: TypeAlias = HTTPScope | WebsocketScope | LifespanScope
+
+WWWSendCallable: TypeAlias = HTTPSendCallable | WebsocketSendCallable
+SendCallable: TypeAlias = (
+    HTTPSendCallable | WebsocketSendCallable | LifespanSendCallable
+)
+
+WWWReceiveCallable: TypeAlias = HTTPReceiveCallable | WebsocketReceiveCallable
+ReceiveCallable: TypeAlias = (
+    HTTPReceiveCallable | WebsocketReceiveCallable | LifespanReceiveCallable
+)

latch_asgi/auth.py

@@ -10,7 +10,7 @@ from jwt import PyJWKClient
 from latch_o11y.o11y import app_tracer, trace_app_function
 
 from .config import config
-from .framework import HTTPErrorResponse
+from .framework.http import HTTPErrorResponse
 
 jwk_client = PyJWKClient("https://latchai.us.auth0.com/.well-known/jwks.json")
 debug_authentication_header_regex = re.compile(
@@ -39,9 +39,11 @@ class _HTTPUnauthorized(HTTPErrorResponse):
     def __init__(
         self,
         error_description: str,
-        error: Literal["invalid_request"]
-        | Literal["invalid_token"]
-        | Literal["insufficient_scope"],
+        error: (
+            Literal["invalid_request"]
+            | Literal["invalid_token"]
+            | Literal["insufficient_scope"]
+        ),
     ):
         escaped_description = error_description.replace('"', '\\"')
         super().__init__(
@@ -138,9 +140,7 @@ def get_signer_sub(auth_header: str) -> Authorization:
                 algorithms=["RS256", "HS256"],
                 # fixme(maximsmol): gut this abomination
                 audience=(
-                    config.audience
-                    if jwt_key != config.self_signed_jwk
-                    else None
+                    config.audience if jwt_key != config.self_signed_jwk else None
                 ),
             )
         except jwt.exceptions.InvalidTokenError as e:

latch_asgi/{context.py → context/common.py}

@@ -1,25 +1,22 @@
 from dataclasses import dataclass, field
-from typing import Any, Awaitable, Callable, TypeAlias, TypeVar
+from typing import Generic, TypeVar
 
-from latch_o11y.o11y import (
-    AttributesDict,
-    app_tracer,
-    dict_to_attrs,
-    trace_app_function,
-)
+from latch_o11y.o11y import AttributesDict, app_tracer, trace_app_function
 
-from .asgi_iface import HTTPReceiveCallable, HTTPScope, HTTPSendCallable
-from .auth import Authorization, get_signer_sub
-from .framework import HTTPMethod, current_http_request_span, receive_class_ext
+from ..asgi_iface import WWWReceiveCallable, WWWScope, WWWSendCallable
+from ..auth import Authorization, get_signer_sub
 
-T = TypeVar("T")
+# todo(ayush): this sucks
+Scope = TypeVar("Scope", bound=WWWScope)
+Send = TypeVar("Send", bound=WWWSendCallable)
+Receive = TypeVar("Receive", bound=WWWReceiveCallable)
 
 
 @dataclass
-class Context:
-    scope: HTTPScope
-    receive: HTTPReceiveCallable
-    send: HTTPSendCallable
+class Context(Generic[Scope, Receive, Send]):
+    scope: Scope
+    receive: Receive
+    send: Send
 
     auth: Authorization = field(default_factory=Authorization, init=False)
 
@@ -33,9 +30,6 @@ class Context:
         if auth_header is not None:
             self.auth = get_signer_sub(auth_header)
 
-        if self.auth.oauth_sub is not None:
-            current_http_request_span().set_attribute("enduser.id", self.auth.oauth_sub)
-
     def header(self, x: str | bytes):
         if isinstance(x, str):
             x = x.encode("utf-8")
@@ -58,7 +52,7 @@ class Context:
         return res.decode("latin-1")
 
     def add_request_span_attrs(self, data: AttributesDict, prefix: str):
-        current_http_request_span().set_attributes(dict_to_attrs(data, prefix))
+        raise NotImplementedError()
 
     @trace_app_function
     def add_db_response(self, data: AttributesDict):
@@ -68,23 +62,3 @@ class Context:
         # )
         self.add_request_span_attrs(data, f"db.response.{self._db_response_idx}")
         self._db_response_idx += 1
-
-    @trace_app_function
-    async def receive_request_payload(self, cls: type[T]) -> T:
-        json, res = await receive_class_ext(self.receive, cls)
-
-        # todo(maximsmol): datadog has shit support for events
-        # current_http_request_span().add_event(
-        #     "request payload", dict_to_attrs(json, "data")
-        # )
-        self.add_request_span_attrs(json, "http.request_payload")
-
-        return res
-
-
-HandlerResult: TypeAlias = Any | None
-Handler: TypeAlias = Callable[
-    [Context],
-    Awaitable[HandlerResult],
-]
-Route: TypeAlias = Handler | tuple[list[HTTPMethod], Handler]

latch_asgi/context/http.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable, TypeAlias, TypeVar
+
+from latch_o11y.o11y import AttributesDict, dict_to_attrs, trace_app_function
+
+from ..asgi_iface import HTTPReceiveCallable, HTTPScope, HTTPSendCallable
+from ..framework.http import HTTPMethod, current_http_request_span, receive_class_ext
+from . import common
+
+T = TypeVar("T")
+
+
+@dataclass
+class Context(common.Context[HTTPScope, HTTPReceiveCallable, HTTPSendCallable]):
+    def __post_init__(self):
+        super().__post_init__()
+
+        if self.auth.oauth_sub is not None:
+            current_http_request_span().set_attribute("enduser.id", self.auth.oauth_sub)
+
+    def add_request_span_attrs(self, data: AttributesDict, prefix: str):
+        current_http_request_span().set_attributes(dict_to_attrs(data, prefix))
+
+    @trace_app_function
+    async def receive_request_payload(self, cls: type[T]) -> T:
+        json, res = await receive_class_ext(self.receive, cls)
+
+        # todo(maximsmol): datadog has shit support for events
+        # current_http_request_span().add_event(
+        #     "request payload", dict_to_attrs(json, "data")
+        # )
+        self.add_request_span_attrs(json, "http.request_payload")
+
+        return res
+
+
+HandlerResult: TypeAlias = Any | None
+Handler: TypeAlias = Callable[
+    [Context],
+    Awaitable[HandlerResult],
+]
+Route: TypeAlias = Handler | tuple[list[HTTPMethod], Handler]

latch_asgi/context/websocket.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+from typing import Awaitable, Callable, TypeAlias, TypeVar
+
+from latch_o11y.o11y import AttributesDict, dict_to_attrs
+
+from ..asgi_iface import WebsocketReceiveCallable, WebsocketScope, WebsocketSendCallable
+from ..framework.websocket import WebsocketStatus, current_websocket_request_span
+from . import common
+
+T = TypeVar("T")
+
+
+@dataclass
+class Context(
+    common.Context[WebsocketScope, WebsocketReceiveCallable, WebsocketSendCallable]
+):
+    def __post_init__(self):
+        super().__post_init__()
+
+        if self.auth.oauth_sub is not None:
+            current_websocket_request_span().set_attribute(
+                "enduser.id", self.auth.oauth_sub
+            )
+
+    def add_request_span_attrs(self, data: AttributesDict, prefix: str):
+        current_websocket_request_span().set_attributes(dict_to_attrs(data, prefix))
+
+
+HandlerResult = str | tuple[WebsocketStatus, str]
+Handler: TypeAlias = Callable[
+    [Context],
+    Awaitable[HandlerResult],
+]
+Route: TypeAlias = Handler

latch_asgi/framework/common.py

@@ -0,0 +1,7 @@
+from typing import TypeAlias
+
+from opentelemetry.trace import get_tracer
+
+Headers: TypeAlias = dict[str | bytes, str | bytes]
+
+tracer = get_tracer(__name__)

latch_asgi/{framework.py → framework/http.py}

@@ -2,26 +2,21 @@ from http import HTTPStatus
 from typing import Any, Literal, TypeAlias, TypeVar, cast
 
 import opentelemetry.context as context
-import simdjson
+import orjson
 from latch_data_validation.data_validation import DataValidationError, validate
 from latch_o11y.o11y import trace_function, trace_function_with_span
-from opentelemetry.trace import get_tracer
 from opentelemetry.trace.span import Span
-from orjson import dumps
 
-from .asgi_iface import (
+from ..asgi_iface import (
     HTTPReceiveCallable,
     HTTPResponseBodyEvent,
     HTTPResponseStartEvent,
     HTTPSendCallable,
 )
-
-Headers: TypeAlias = dict[str | bytes, str | bytes]
+from .common import Headers, tracer
 
 T = TypeVar("T")
 
-tracer = get_tracer(__name__)
-
 HTTPMethod: TypeAlias = (
     Literal["GET"]
     | Literal["HEAD"]
@@ -48,7 +43,6 @@ def current_http_request_span() -> Span:
 
 class HTTPErrorResponse(RuntimeError):
     def __init__(self, status: HTTPStatus, data: Any, *, headers: Headers = {}):
-        super().__init__()
         self.status = status
         self.data = data
         self.headers = headers
@@ -69,8 +63,7 @@ class HTTPForbidden(HTTPErrorResponse):
         super().__init__(HTTPStatus.FORBIDDEN, data, headers=headers)
 
 
-class HTTPConnectionClosedError(RuntimeError):
-    ...
+class HTTPConnectionClosedError(RuntimeError): ...
 
 
 # >>> I/O
@@ -96,13 +89,7 @@ async def receive_class(receive: HTTPReceiveCallable, cls: type[T]) -> T:
 
 @trace_function(tracer)
 async def receive_json(receive: HTTPReceiveCallable) -> Any:
-    res = await receive_data(receive)
-
-    p = simdjson.Parser()
-    try:
-        return p.parse(res, True)
-    except ValueError as e:
-        raise HTTPBadRequest("Failed to parse JSON") from e
+    return orjson.loads(await receive_data(receive))
 
 
 async def receive_data(receive: HTTPReceiveCallable):
@@ -127,7 +114,7 @@ async def receive_data(receive: HTTPReceiveCallable):
 
 
 @trace_function_with_span(tracer)
-async def send_data(
+async def send_http_data(
     s: Span,
     send: HTTPSendCallable,
     status: HTTPStatus,
@@ -178,8 +165,8 @@ async def send_json(
     content_type: str = "application/json",
     headers: Headers = {},
 ):
-    return await send_data(
-        send, status, dumps(data), content_type=content_type, headers=headers
+    return await send_http_data(
+        send, status, orjson.dumps(data), content_type=content_type, headers=headers
     )
 
 
@@ -193,6 +180,6 @@ async def send_auto(
     headers: Headers = {},
 ):
     if isinstance(data, str) or isinstance(data, bytes):
-        return await send_data(send, status, data, headers=headers)
+        return await send_http_data(send, status, data, headers=headers)
 
     return await send_json(send, status, data, headers=headers)

latch_asgi/framework/websocket.py

@@ -0,0 +1,266 @@
+from enum import Enum
+from typing import Any, TypeVar, cast
+
+import opentelemetry.context as context
+from latch_data_validation.data_validation import DataValidationError, validate
+from latch_o11y.o11y import trace_function, trace_function_with_span
+from opentelemetry.trace.span import Span
+import orjson
+
+from ..asgi_iface import (
+    WebsocketAcceptEvent,
+    WebsocketCloseEvent,
+    WebsocketReceiveCallable,
+    WebsocketSendCallable,
+    WebsocketSendEvent,
+)
+from .common import Headers, tracer
+
+T = TypeVar("T")
+
+websocket_request_span_key = context.create_key("websocket_request_span")
+
+
+def current_websocket_request_span() -> Span:
+    return cast(Span, context.get_value(websocket_request_span_key))
+
+
+# >>> Error classes
+
+
+class WebsocketStatus(int, Enum):
+    """
+    https://www.rfc-editor.org/rfc/rfc6455.html#section-7.4.1
+    """
+
+    normal = 1000
+    """
+    1000 indicates a normal closure, meaning that the purpose for
+    which the connection was established has been fulfilled.
+    """
+
+    going_away = 1001
+    """
+    1001 indicates that an endpoint is "going away", such as a server
+    going down or a browser having navigated away from a page.
+    """
+
+    protocol_error = 1002
+    """
+    1002 indicates that an endpoint is terminating the connection due
+    to a protocol error.
+    """
+
+    unsupported = 1003
+    """
+    1003 indicates that an endpoint is terminating the connection
+    because it has received a type of data it cannot accept (e.g., an
+    endpoint that understands only text data MAY send this if it
+    receives a binary message).
+    """
+
+    reserved = 1004
+    """
+    Reserved.  The specific meaning might be defined in the future.
+    """
+
+    no_status = 1005
+    """
+    1005 is a reserved value and MUST NOT be set as a status code in a
+    Close control frame by an endpoint.  It is designated for use in
+    applications expecting a status code to indicate that no status
+    code was actually present.
+    """
+
+    abnormal = 1006
+    """
+    1006 is a reserved value and MUST NOT be set as a status code in a
+    Close control frame by an endpoint.  It is designated for use in
+    applications expecting a status code to indicate that the
+    connection was closed abnormally, e.g., without sending or
+    receiving a Close control frame.
+    """
+
+    unsupported_payload = 1007
+    """
+    1007 indicates that an endpoint is terminating the connection
+    because it has received data within a message that was not
+    consistent with the type of the message (e.g., non-UTF-8 [RFC3629]
+    data within a text message).
+    """
+
+    policy_violation = 1008
+    """
+    1008 indicates that an endpoint is terminating the connection
+    because it has received a message that violates its policy.  This
+    is a generic status code that can be returned when there is no
+    other more suitable status code (e.g., 1003 or 1009) or if there
+    is a need to hide specific details about the policy.
+    """
+
+    too_large = 1009
+    """
+    1009 indicates that an endpoint is terminating the connection
+    because it has received a message that is too big for it to
+    process.
+    """
+
+    mandatory_extension = 1010
+    """
+    1010 indicates that an endpoint (client) is terminating the
+    connection because it has expected the server to negotiate one or
+    more extension, but the server didn't return them in the response
+    message of the WebSocket handshake.  The list of extensions that
+    are needed SHOULD appear in the /reason/ part of the Close frame.
+    Note that this status code is not used by the server, because it
+    can fail the WebSocket handshake instead.
+    """
+
+    server_error = 1011
+    """
+    1011 indicates that a server is terminating the connection because
+    it encountered an unexpected condition that prevented it from
+    fulfilling the request.
+    """
+
+    tls_handshake_fail = 1015
+    """
+    1015 is a reserved value and MUST NOT be set as a status code in a
+    Close control frame by an endpoint.  It is designated for use in
+    applications expecting a status code to indicate that the
+    connection was closed due to a failure to perform a TLS handshake
+    (e.g., the server certificate can't be verified).
+    """
+
+
+class WebsocketErrorResponse(RuntimeError):
+    def __init__(self, status: WebsocketStatus, data: Any, *, headers: Headers = {}):
+        super().__init__()
+        self.status = status
+        self.data = data
+        self.headers = headers
+
+
+class WebsocketBadMessage(WebsocketErrorResponse):
+    def __init__(self, data: Any, *, headers: Headers = {}):
+        super().__init__(WebsocketStatus.policy_violation, data, headers=headers)
+
+
+class WebsocketInternalServerError(WebsocketErrorResponse):
+    def __init__(self, data: Any, *, headers: Headers = {}):
+        super().__init__(WebsocketStatus.server_error, data, headers=headers)
+
+
+class WebsocketConnectionClosedError(RuntimeError): ...
+
+
+# >>> I/O
+
+
+async def receive_websocket_data(receive: WebsocketReceiveCallable):
+    with tracer.start_as_current_span("read websocket message") as s:
+        msg = await receive()
+
+        s.set_attribute("type", msg.type)
+
+        if msg.type == "websocket.connect":
+            # todo(ayush): allow upgrades here as well?
+            raise WebsocketBadMessage("connection has already been established")
+
+        if msg.type == "websocket.disconnect":
+            raise WebsocketConnectionClosedError()
+
+        if msg.bytes is not None:
+            res = msg.bytes
+        elif msg.text is not None:
+            res = msg.text.encode("utf-8")
+        else:
+            raise WebsocketBadMessage("empty message")
+
+        s.set_attribute("size", len(res))
+        return res
+
+
+@trace_function(tracer)
+async def receive_websocket_json(receive: WebsocketReceiveCallable) -> Any:
+    orjson.loads(await receive_websocket_data(receive))
+
+
+@trace_function(tracer)
+async def receive_websocket_class_ext(
+    receive: WebsocketReceiveCallable, cls: type[T]
+) -> tuple[Any, T]:
+    data = await receive_websocket_json(receive)
+
+    try:
+        return data, validate(data, cls)
+    except DataValidationError as e:
+        raise WebsocketBadMessage(e.json()) from None
+
+
+@trace_function(tracer)
+async def receive_websocket_class(receive: WebsocketReceiveCallable, cls: type[T]) -> T:
+    return (await receive_websocket_class_ext(receive, cls))[1]
+
+
+@trace_function_with_span(tracer)
+async def send_websocket_data(
+    s: Span,
+    send: WebsocketSendCallable,
+    data: str | bytes,
+):
+    if isinstance(data, str):
+        data = data.encode("utf-8")
+
+    s.set_attribute("body.size", len(data))
+
+    await send(WebsocketSendEvent(type="websocket.send", bytes=data, text=None))
+
+    current_websocket_request_span().set_attribute(
+        "websocket.sent_message_content_length", len(data)
+    )
+
+
+@trace_function(tracer)
+async def accept_websocket_connection(
+    send: WebsocketSendCallable,
+    receive: WebsocketReceiveCallable,
+    /,
+    *,
+    subprotocol: str | None = None,
+    headers: Headers = {},
+):
+    msg = await receive()
+    if msg.type != "websocket.connect":
+        raise WebsocketBadMessage("cannot accept connection without connection request")
+
+    headers_to_send: list[tuple[bytes, bytes]] = []
+
+    for k, v in headers.items():
+        if isinstance(k, str):
+            k = k.encode("latin-1")
+        if isinstance(v, str):
+            v = v.encode("latin-1")
+        headers_to_send.append((k, v))
+
+    await send(
+        WebsocketAcceptEvent(
+            type="websocket.accept", subprotocol=subprotocol, headers=headers_to_send
+        )
+    )
+
+
+@trace_function_with_span(tracer)
+async def close_websocket_connection(
+    s: Span,
+    send: WebsocketSendCallable,
+    /,
+    *,
+    status: WebsocketStatus,
+    data: str,
+):
+    s.set_attribute("reason", data)
+
+    await send(WebsocketCloseEvent("websocket.close", status.value, data))
+
+    current_websocket_request_span().set_attribute("websocket.http.close_reason", data)

latch_asgi/server.py

@@ -2,17 +2,15 @@ import asyncio
 import dataclasses
 import traceback
 from http import HTTPStatus
+from typing import Awaitable, get_args
 
 import opentelemetry.context as context
 from hypercorn.typing import ASGIReceiveCallable, ASGISendCallable, Scope
-from typing import Awaitable
 from latch_data_validation.data_validation import untraced_validate, validate
 from latch_o11y.o11y import log
 from opentelemetry.propagate import set_global_textmap
 from opentelemetry.trace import get_current_span, get_tracer
 
-from .datadog_propagator import DDTraceContextTextMapPropagator
-
 from .asgi_iface import (
     HTTPDisconnectEvent,
     HTTPReceiveCallable,
@@ -30,16 +28,31 @@ from .asgi_iface import (
     LifespanStartupCompleteEvent,
     LifespanStartupEvent,
     LifespanStartupFailedEvent,
+    WebsocketReceiveCallable,
+    WebsocketReceiveEventT,
+    WebsocketScope,
+    WebsocketSendCallable,
+    WebsocketSendEventT,
     type_str,
 )
-from .context import Context, Route
-from .framework import (
+from .context import http, websocket
+from .datadog_propagator import DDTraceContextTextMapPropagator
+from .framework.http import (
     HTTPErrorResponse,
     HTTPInternalServerError,
     current_http_request_span,
     http_request_span_key,
     send_auto,
-    send_data,
+    send_http_data,
+)
+from .framework.websocket import (
+    WebsocketErrorResponse,
+    WebsocketInternalServerError,
+    WebsocketStatus,
+    accept_websocket_connection,
+    close_websocket_connection,
+    current_websocket_request_span,
+    websocket_request_span_key,
 )
 
 tracer = get_tracer(__name__)
@@ -47,18 +60,30 @@ tracer = get_tracer(__name__)
 
 # todo(maximsmol): ASGI instrumentation should trace lifespan
 
+
 class LatchASGIServer:
-    routes: dict[str, Route]
+    http_routes: dict[str, http.Route]
+    websocket_routes: dict[str, websocket.Route]
     startup_tasks: list[Awaitable] = []
     shutdown_tasks: list[Awaitable] = []
 
-    def __init__(self, routes: dict[str, Route], startup_tasks: list[Awaitable] = [], shutdown_tasks: list[Awaitable] = []):
-        self.routes = routes
+    def __init__(
+        self,
+        http_routes: dict[str, http.Route],
+        websocket_routes: dict[str, websocket.Route],
+        startup_tasks: list[Awaitable] = [],
+        shutdown_tasks: list[Awaitable] = [],
+    ):
+        self.http_routes = http_routes
+        self.websocket_routes = websocket_routes
         self.startup_tasks = startup_tasks
         self.shutdown_tasks = shutdown_tasks
 
     async def scope_lifespan(
-        self, scope: LifespanScope, receive: LifespanReceiveCallable, send: LifespanSendCallable
+        self,
+        scope: LifespanScope,
+        receive: LifespanReceiveCallable,
+        send: LifespanSendCallable,
     ):
         await log.info(
             f"Waiting for lifespan events (ASGI v{scope.asgi.version} @ spec"
@@ -77,7 +102,9 @@ class LatchASGIServer:
 
                         with tracer.start_as_current_span("send completion event"):
                             await send(
-                                LifespanStartupCompleteEvent("lifespan.startup.complete")
+                                LifespanStartupCompleteEvent(
+                                    "lifespan.startup.complete"
+                                )
                             )
                     except Exception as e:
                         with tracer.start_as_current_span("send failure event"):
@@ -95,7 +122,9 @@ class LatchASGIServer:
 
                         with tracer.start_as_current_span("send completion event"):
                             await send(
-                                LifespanShutdownCompleteEvent("lifespan.shutdown.complete")
+                                LifespanShutdownCompleteEvent(
+                                    "lifespan.shutdown.complete"
+                                )
                             )
                     except Exception as e:
                         with tracer.start_as_current_span("send failure event"):
@@ -109,6 +138,60 @@ class LatchASGIServer:
 
                     break
 
+    async def scope_websocket(
+        self,
+        scope: WebsocketScope,
+        receive: WebsocketReceiveCallable,
+        send: WebsocketSendCallable,
+    ):
+        ctx_reset_token: object | None = None
+        try:
+            new_ctx = context.set_value(websocket_request_span_key, get_current_span())
+            ctx_reset_token = context.attach(new_ctx)
+
+            current_websocket_request_span().set_attribute("resource.name", scope.path)
+
+            handler = self.websocket_routes.get(scope.path)
+            if handler is None:
+                msg = f"Websocket {scope.path} not found"
+
+                await log.info(msg)
+                await close_websocket_connection(
+                    send, status=WebsocketStatus.policy_violation, data=msg
+                )
+                return
+
+            await log.info(f"Websocket {scope.path}")
+
+            try:
+                try:
+                    ctx = websocket.Context(scope, receive, send)
+
+                    await accept_websocket_connection(ctx.send, ctx.receive)
+                    res = await handler(ctx)
+
+                    if isinstance(res, tuple):
+                        status, data = res
+                    else:
+                        status = WebsocketStatus.normal
+                        data = res
+
+                except WebsocketErrorResponse as e:
+                    raise e
+                except Exception as e:
+                    raise WebsocketInternalServerError(str(e)) from e
+            except WebsocketErrorResponse as e:
+                await close_websocket_connection(
+                    send, status=WebsocketStatus.server_error, data=str(e.data)
+                )
+
+                if e.status == HTTPStatus.INTERNAL_SERVER_ERROR:
+                    traceback.print_exc()
+            else:
+                await close_websocket_connection(send, status=status, data=data)
+        finally:
+            if ctx_reset_token is not None:
+                context.detach(ctx_reset_token)
 
     async def scope_http(
         self, scope: HTTPScope, receive: HTTPReceiveCallable, send: HTTPSendCallable
@@ -122,7 +205,7 @@ class LatchASGIServer:
             await log.info(f"{scope.method} {scope.path}")
 
             with tracer.start_as_current_span("find route handler"):
-                route = self.routes.get(scope.path)
+                route = self.http_routes.get(scope.path)
 
             if not isinstance(route, tuple):
                 methods = ["POST"]
@@ -138,7 +221,7 @@ class LatchASGIServer:
                 else:
                     methods_str = ", and ".join([", ".join(methods[:-1]), methods[-1]])
 
-                await send_data(
+                await send_http_data(
                     send,
                     HTTPStatus.METHOD_NOT_ALLOWED,
                     f"Only {methods_str} requests are supported",
@@ -148,12 +231,12 @@ class LatchASGIServer:
             if handler is None:
                 # todo(maximsmol): better error message
                 await log.info("Not found")
-                await send_data(send, HTTPStatus.NOT_FOUND, "Not found")
+                await send_http_data(send, HTTPStatus.NOT_FOUND, "Not found")
                 return
 
             try:
                 try:
-                    ctx = Context(scope, receive, send)
+                    ctx = http.Context(scope, receive, send)
                     res = await handler(ctx)
 
                     if res is not None:
@@ -175,8 +258,9 @@ class LatchASGIServer:
             if ctx_reset_token is not None:
                 context.detach(ctx_reset_token)
 
-
-    async def raw_app(self, scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
+    async def raw_app(
+        self, scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
+    ):
         try:
             if scope["type"] == "lifespan":
 
@@ -189,7 +273,9 @@ class LatchASGIServer:
                     if x["type"] == type_str(LifespanShutdownEvent):
                         return untraced_validate(x, LifespanShutdownEvent)
 
-                    raise RuntimeError(f"unknown lifespan event type: {repr(x['type'])}")
+                    raise RuntimeError(
+                        f"unknown lifespan event type: {repr(x['type'])}"
+                    )
 
                 async def ls_send(e: LifespanSendEvent):
                     data = dataclasses.asdict(e)
@@ -199,6 +285,29 @@ class LatchASGIServer:
                     untraced_validate(scope, LifespanScope), ls_receive, ls_send
                 )
 
+            if scope["type"] == "websocket":
+
+                async def ws_receive():
+                    x = await receive()
+
+                    for e in get_args(WebsocketReceiveEventT):
+                        if x["type"] != type_str(e):
+                            continue
+
+                        return untraced_validate(x, e)
+
+                    raise RuntimeError(
+                        f"unknown websocket event type: {repr(x['type'])}"
+                    )
+
+                async def ws_send(e: WebsocketSendEventT):
+                    data = dataclasses.asdict(e)
+                    await send(data)
+
+                return await self.scope_websocket(
+                    untraced_validate(scope, WebsocketScope), ws_receive, ws_send
+                )
+
             if scope["type"] == "http":
 
                 async def http_receive():
@@ -216,7 +325,9 @@ class LatchASGIServer:
                     data = dataclasses.asdict(e)
                     await send(data)
 
-                return await self.scope_http(validate(scope, HTTPScope), http_receive, http_send)
+                return await self.scope_http(
+                    validate(scope, HTTPScope), http_receive, http_send
+                )
 
             raise RuntimeError(f"unsupported protocol: {repr(scope['type'])}")
         except Exception as e:

latch_asgi-0.3.0.dist-info/METADATA

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: latch-asgi
+Version: 0.3.0
+Summary: ASGI python server
+Author-Email: Max Smolin <max@latch.bio>
+License: CC0-1.0
+Requires-Python: <4.0,>=3.11
+Requires-Dist: hypercorn[uvloop]<1.0.0,>=0.14.3
+Requires-Dist: latch-data-validation<1.0.0,>=0.1.3
+Requires-Dist: latch-o11y<1.0.0,>=0.1.4
+Requires-Dist: latch-config<1.0.0,>=0.1.6
+Requires-Dist: PyJWT[crypto]<3.0.0,>=2.6.0
+Requires-Dist: orjson<4.0.0,>=3.8.5
+Requires-Dist: opentelemetry-sdk<2.0.0,>=1.15.0
+Requires-Dist: opentelemetry-api<2.0.0,>=1.15.0
+Requires-Dist: opentelemetry-instrumentation-asgi<1.0,>=0.36b0
+Description-Content-Type: text/markdown
+
+# python-asgi

latch_asgi-0.3.0.dist-info/RECORD

@@ -0,0 +1,19 @@
+latch_asgi-0.3.0.dist-info/METADATA,sha256=x1WxOrLQcsLzJzLEaFCl6yCsK0AkZ8EHdhp82vIYYPo,643
+latch_asgi-0.3.0.dist-info/WHEEL,sha256=N2J68yzZqJh3mI_Wg92rwhw0rtJDFpZj9bwQIMJgaVg,90
+latch_asgi-0.3.0.dist-info/licenses/COPYING,sha256=ogEPNDSH0_dhiv_lT3ifVIdgIzHAqNA_SemnxUfPBJk,7048
+latch_asgi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+latch_asgi/asgi_iface.py,sha256=8cltPG2YEf20RWcMpn2FE5g4wtauiUNNlrVnD4UYDqc,8397
+latch_asgi/auth.py,sha256=gXGVIzsvDt_kBlcS_Sx9939iGlNNs1i8KDOs-tpEBiw,4996
+latch_asgi/config.py,sha256=F-78OqEGqrYoDKAdo5OVWwLB8cm6A2qvAC0xJQUufy0,241
+latch_asgi/context/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+latch_asgi/context/common.py,sha256=EQeVLIbdfX3m5tLkLdzLynEvPUKrN2HWuvjhDrYwpoo,2029
+latch_asgi/context/http.py,sha256=-zwxko3ao1mirdfJLcOacQRCOW1U1bXr40HEA_50Z24,1424
+latch_asgi/context/websocket.py,sha256=NA0pJYdJI2kUnSu1AkPnf9heG8_B4M4_op7S8MPlyT4,1024
+latch_asgi/datadog_propagator.py,sha256=PFKaM87B8Ia_5RBNGIIDSgcgTjxhN0uJou2o8359kzs,3306
+latch_asgi/framework/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+latch_asgi/framework/common.py,sha256=My4ePrmLD8aZd7_PC-ZJPZ9CUD90j5UZvHwA0jfTmFw,157
+latch_asgi/framework/http.py,sha256=8AMzA1FjyaFVexVPQbi1NBdbWaH6594R5AAsjoSL67Y,4992
+latch_asgi/framework/websocket.py,sha256=Xyf4EUurKPxdHGV3CXUpCMLYCE41h_zz9GjJaR3NZgw,7993
+latch_asgi/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+latch_asgi/server.py,sha256=kMZ7ivOk7VkZ5HjrIHcveXd2ZMawnnW7QM5-JuHDXZA,12047
+latch_asgi-0.3.0.dist-info/RECORD,,

{latch_asgi-0.1.3.dist-info → latch_asgi-0.3.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.4.0
+Generator: pdm-backend (2.1.8)
 Root-Is-Purelib: true
 Tag: py3-none-any

latch_asgi-0.3.0.dist-info/licenses/COPYING

@@ -0,0 +1,121 @@
+Creative Commons Legal Code
+
+CC0 1.0 Universal
+
+    CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE
+    LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN
+    ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS
+    INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES
+    REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS
+    PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM
+    THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED
+    HEREUNDER.
+
+Statement of Purpose
+
+The laws of most jurisdictions throughout the world automatically confer
+exclusive Copyright and Related Rights (defined below) upon the creator
+and subsequent owner(s) (each and all, an "owner") of an original work of
+authorship and/or a database (each, a "Work").
+
+Certain owners wish to permanently relinquish those rights to a Work for
+the purpose of contributing to a commons of creative, cultural and
+scientific works ("Commons") that the public can reliably and without fear
+of later claims of infringement build upon, modify, incorporate in other
+works, reuse and redistribute as freely as possible in any form whatsoever
+and for any purposes, including without limitation commercial purposes.
+These owners may contribute to the Commons to promote the ideal of a free
+culture and the further production of creative, cultural and scientific
+works, or to gain reputation or greater distribution for their Work in
+part through the use and efforts of others.
+
+For these and/or other purposes and motivations, and without any
+expectation of additional consideration or compensation, the person
+associating CC0 with a Work (the "Affirmer"), to the extent that he or she
+is an owner of Copyright and Related Rights in the Work, voluntarily
+elects to apply CC0 to the Work and publicly distribute the Work under its
+terms, with knowledge of his or her Copyright and Related Rights in the
+Work and the meaning and intended legal effect of CC0 on those rights.
+
+1. Copyright and Related Rights. A Work made available under CC0 may be
+protected by copyright and related or neighboring rights ("Copyright and
+Related Rights"). Copyright and Related Rights include, but are not
+limited to, the following:
+
+  i. the right to reproduce, adapt, distribute, perform, display,
+     communicate, and translate a Work;
+ ii. moral rights retained by the original author(s) and/or performer(s);
+iii. publicity and privacy rights pertaining to a person's image or
+     likeness depicted in a Work;
+ iv. rights protecting against unfair competition in regards to a Work,
+     subject to the limitations in paragraph 4(a), below;
+  v. rights protecting the extraction, dissemination, use and reuse of data
+     in a Work;
+ vi. database rights (such as those arising under Directive 96/9/EC of the
+     European Parliament and of the Council of 11 March 1996 on the legal
+     protection of databases, and under any national implementation
+     thereof, including any amended or successor version of such
+     directive); and
+vii. other similar, equivalent or corresponding rights throughout the
+     world based on applicable law or treaty, and any national
+     implementations thereof.
+
+2. Waiver. To the greatest extent permitted by, but not in contravention
+of, applicable law, Affirmer hereby overtly, fully, permanently,
+irrevocably and unconditionally waives, abandons, and surrenders all of
+Affirmer's Copyright and Related Rights and associated claims and causes
+of action, whether now known or unknown (including existing as well as
+future claims and causes of action), in the Work (i) in all territories
+worldwide, (ii) for the maximum duration provided by applicable law or
+treaty (including future time extensions), (iii) in any current or future
+medium and for any number of copies, and (iv) for any purpose whatsoever,
+including without limitation commercial, advertising or promotional
+purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each
+member of the public at large and to the detriment of Affirmer's heirs and
+successors, fully intending that such Waiver shall not be subject to
+revocation, rescission, cancellation, termination, or any other legal or
+equitable action to disrupt the quiet enjoyment of the Work by the public
+as contemplated by Affirmer's express Statement of Purpose.
+
+3. Public License Fallback. Should any part of the Waiver for any reason
+be judged legally invalid or ineffective under applicable law, then the
+Waiver shall be preserved to the maximum extent permitted taking into
+account Affirmer's express Statement of Purpose. In addition, to the
+extent the Waiver is so judged Affirmer hereby grants to each affected
+person a royalty-free, non transferable, non sublicensable, non exclusive,
+irrevocable and unconditional license to exercise Affirmer's Copyright and
+Related Rights in the Work (i) in all territories worldwide, (ii) for the
+maximum duration provided by applicable law or treaty (including future
+time extensions), (iii) in any current or future medium and for any number
+of copies, and (iv) for any purpose whatsoever, including without
+limitation commercial, advertising or promotional purposes (the
+"License"). The License shall be deemed effective as of the date CC0 was
+applied by Affirmer to the Work. Should any part of the License for any
+reason be judged legally invalid or ineffective under applicable law, such
+partial invalidity or ineffectiveness shall not invalidate the remainder
+of the License, and in such case Affirmer hereby affirms that he or she
+will not (i) exercise any of his or her remaining Copyright and Related
+Rights in the Work or (ii) assert any associated claims and causes of
+action with respect to the Work, in either case contrary to Affirmer's
+express Statement of Purpose.
+
+4. Limitations and Disclaimers.
+
+ a. No trademark or patent rights held by Affirmer are waived, abandoned,
+    surrendered, licensed or otherwise affected by this document.
+ b. Affirmer offers the Work as-is and makes no representations or
+    warranties of any kind concerning the Work, express, implied,
+    statutory or otherwise, including without limitation warranties of
+    title, merchantability, fitness for a particular purpose, non
+    infringement, or the absence of latent or other defects, accuracy, or
+    the present or absence of errors, whether or not discoverable, all to
+    the greatest extent permissible under applicable law.
+ c. Affirmer disclaims responsibility for clearing rights of other persons
+    that may apply to the Work or any use thereof, including without
+    limitation any person's Copyright and Related Rights in the Work.
+    Further, Affirmer disclaims responsibility for obtaining any necessary
+    consents, permissions or other rights required for any use of the
+    Work.
+ d. Affirmer understands and acknowledges that Creative Commons is not a
+    party to this document and has no duty or obligation with respect to
+    this CC0 or use of the Work.

latch_asgi-0.1.3.dist-info/METADATA

@@ -1,25 +0,0 @@
-Metadata-Version: 2.1
-Name: latch-asgi
-Version: 0.1.3
-Summary: ASGI python server
-License: CC0 1.0
-Author: Max Smolin
-Author-email: max@latch.bio
-Requires-Python: >=3.11,<4.0
-Classifier: License :: Other/Proprietary License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: PyJWT[crypto] (>=2.6.0,<3.0.0)
-Requires-Dist: hypercorn[uvloop] (>=0.14.3,<0.15.0)
-Requires-Dist: latch-config (>=0.1.6,<0.2.0)
-Requires-Dist: latch-data-validation (>=0.1.3,<0.2.0)
-Requires-Dist: latch-o11y (>=0.1.4,<0.2.0)
-Requires-Dist: opentelemetry-api (>=1.15.0,<2.0.0)
-Requires-Dist: opentelemetry-instrumentation-asgi (>=0.36b0,<0.37)
-Requires-Dist: opentelemetry-sdk (>=1.15.0,<2.0.0)
-Requires-Dist: orjson (>=3.8.5,<4.0.0)
-Requires-Dist: pysimdjson (>=5.0.2,<6.0.0)
-Description-Content-Type: text/markdown
-
-# python-asgi
-

latch_asgi-0.1.3.dist-info/RECORD

@@ -1,12 +0,0 @@
-latch_asgi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-latch_asgi/asgi_iface.py,sha256=yDSA2lyymjv66oGjRmweFPy8M48uUe0zNdz4-5c0cMU,5002
-latch_asgi/auth.py,sha256=3ROSIoOTaxJvWB8gGEdLWAK8XEvlbF8GlWwXb8StzjA,4999
-latch_asgi/config.py,sha256=F-78OqEGqrYoDKAdo5OVWwLB8cm6A2qvAC0xJQUufy0,241
-latch_asgi/context.py,sha256=3WQG9-EGzKrDFPw1ASQiG70rrthJic85d-x0DdvEJp4,2802
-latch_asgi/datadog_propagator.py,sha256=PFKaM87B8Ia_5RBNGIIDSgcgTjxhN0uJou2o8359kzs,3306
-latch_asgi/framework.py,sha256=kMlTEkGwfboc_8yCcKGKtqH5_8ElN16FRhoIg1jNeE0,5261
-latch_asgi/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-latch_asgi/server.py,sha256=oZHDIaDPPChvEE5jmE2a4u9jM3uudA5-wNQdMDxSzdE,8366
-latch_asgi-0.1.3.dist-info/WHEEL,sha256=vVCvjcmxuUltf8cYhJ0sJMRDLr1XsPuxEId8YDzbyCY,88
-latch_asgi-0.1.3.dist-info/METADATA,sha256=bSPFmNOLkbSgZoS2XF5FKaQLwErPXk29EJTzuCyPLcQ,870
-latch_asgi-0.1.3.dist-info/RECORD,,