woggioni-opencode-agent/kaya
1
0
Fork 0
Code Actions Packages Releases Activity

Add WebSocket support

Browse Source
This commit is contained in:
woggioni
2026-07-14 14:07:03 +00:00
parent d4a466ce71
commit d72ab63e09
10 changed files with 430 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kaya/__init__.py
+4 -1
View File
@@ -3,6 +3,7 @@ from ._http_method import HttpMethod
from ._http_context import HttpContext
from ._tree import Tree, PathIterator
from ._path_handler import PathHandler
from ._websocket import WebSocket, WebSocketMessage
__all__ = [
@@ -11,5 +12,7 @@ __all__ = [
'HttpContext',
'Tree',
'PathHandler',
'PathIterator'
'PathIterator',
'WebSocket',
'WebSocketMessage'
]
src/kaya/_app.py
+44 -11
View File
@@ -2,36 +2,44 @@ from abc import ABC, abstractmethod
from asyncio import Queue, AbstractEventLoop
from asyncio import get_running_loop
from logging import getLogger
from typing import Callable, Awaitable, Any, Mapping, Sequence, Optional, Unpack, Tuple, TYPE_CHECKING
from typing import Callable, Awaitable, Any, Mapping, Sequence, Optional, Unpack, Tuple, TYPE_CHECKING, cast
from pathlib import Path, PurePath
from pwo import Maybe, AsyncQueueIterator
from hashlib import md5
from ._http_context import HttpContext
from ._http_method import HttpMethod
from ._path_handler import Context
from ._types import StrOrStrings
from ._websocket import WebSocket
from base64 import b64encode, b64decode
from mimetypes import guess_type
if TYPE_CHECKING:
from _typeshed import StrOrBytesPath
try:
from ._rsgi import RsgiContext
from granian._granian import RSGIHTTPProtocol, RSGIHTTPScope # type: ignore
from ._rsgi import RsgiContext, RsgiWebSocket
from granian._granian import ( # type: ignore
RSGIHTTPProtocol,
RSGIHTTPScope,
RSGIWebsocketProtocol,
RSGIWebsocketScope as RSGIWebSocketScope,
)
except ImportError:
pass
from ._asgi import AsgiContext
from ._asgi import AsgiContext, AsgiWebSocket
from ._tree import Tree
from ._types.asgi import LifespanScope, HTTPScope as ASGIHTTPScope, WebSocketScope
from ._types.asgi import LifespanScope, HTTPScope as ASGIHTTPScope, WebSocketScope as ASGIWebSocketScope
log = getLogger(__name__)
type HttpHandler = Callable[[HttpContext, Unpack[Any]], Awaitable[None]]
type WebSocketHandler = Callable[[WebSocket, Unpack[Any]], Awaitable[None]]
class AbstractBugisApp(ABC):
async def __call__(self,
scope: ASGIHTTPScope | WebSocketScope | LifespanScope,
scope: ASGIHTTPScope | ASGIWebSocketScope | LifespanScope,
receive: Callable[[], Awaitable[Any]],
send: Callable[[Mapping[str, Any]], Awaitable[None]]) -> None:
loop = get_running_loop()
@@ -59,6 +67,9 @@ class AbstractBugisApp(ABC):
elif message['type'] == 'http.disconnect':
request_handling.cancel()
break
elif scope['type'] == 'websocket':
ws = AsgiWebSocket(scope, receive, send)
await self.handle_websocket(ws)
else:
raise NotImplementedError()
@@ -72,15 +83,23 @@ class AbstractBugisApp(ABC):
async def handle_request(self, ctx: HttpContext) -> None:
raise NotImplementedError()
@abstractmethod
async def handle_websocket(self, ws: WebSocket) -> None:
raise NotImplementedError()
def __rsgi_init__(self, loop: AbstractEventLoop) -> None:
self.setup(loop)
def __rsgi_del__(self, loop: AbstractEventLoop) -> None:
self.shutdown(loop)
async def __rsgi__(self, scope: RSGIHTTPScope, protocol: RSGIHTTPProtocol) -> None:
ctx = RsgiContext(scope, protocol)
await self.handle_request(ctx)
async def __rsgi__(self, scope: RSGIHTTPScope | RSGIWebSocketScope, protocol: RSGIHTTPProtocol | RSGIWebsocketProtocol) -> None:
if scope.proto == 'ws':
ws = RsgiWebSocket(scope, protocol) # type: ignore[arg-type]
await self.handle_websocket(ws)
else:
ctx = RsgiContext(scope, protocol) # type: ignore[arg-type]
await self.handle_request(ctx)
class BugisApp(AbstractBugisApp):
@@ -98,6 +117,15 @@ class BugisApp(AbstractBugisApp):
await ctx.send_empty(404)
pass
async def handle_websocket(self, ws: WebSocket) -> None:
result = self._tree.get_handler(ws.path, HttpMethod.WS)
if result is not None:
handler, captured = result
await handler.handle_request(ws, captured)
else:
await ws.close(1000)
pass
def route(self,
paths: StrOrStrings,
methods: Optional[HttpMethod | Sequence[HttpMethod]] = None,
@@ -120,11 +148,17 @@ class BugisApp(AbstractBugisApp):
_paths = tuple(paths)
for method in _methods:
for path in _paths:
self._tree.register(path, method, handler, recursive)
self._tree.register(path, method, cast(Callable[[Context, Unpack[Any]], Awaitable[None]], handler), recursive)
return handler
return wrapped
def websocket(self, path: str, recursive: bool = False) -> Callable[[WebSocketHandler], WebSocketHandler]:
def wrapped(handler: WebSocketHandler) -> WebSocketHandler:
self._tree.register(path, HttpMethod.WS, cast(Callable[[Context, Unpack[Any]], Awaitable[None]], handler), recursive)
return handler
return wrapped
def GET(self, path: str, recursive: bool = False) -> Callable[[HttpHandler], HttpHandler]:
return self.route(path, (HttpMethod.GET,), recursive)
@@ -146,4 +180,3 @@ class BugisApp(AbstractBugisApp):
def PATCH(self, path: str, recursive: bool = False) -> Callable[[HttpHandler], HttpHandler]:
return self.route(path, (HttpMethod.PATCH,), recursive)
src/kaya/_asgi.py
+62 -1
View File
@@ -17,8 +17,9 @@ from pwo import Maybe
from pathlib import Path
from ._http_method import HttpMethod
from ._http_context import HttpContext
from ._websocket import WebSocket, WebSocketMessage
from ._types import StrOrStrings
from ._types.asgi import HTTPScope
from ._types.asgi import HTTPScope, WebSocketScope
def decode_headers(headers: Iterable[Tuple[bytes, bytes]]) -> Dict[str, Sequence[str]]:
@@ -146,3 +147,63 @@ class AsgiContext(HttpContext):
'body': '',
'more_body': False
})
class AsgiWebSocket(WebSocket):
_receive: Callable[[], Awaitable[Any]]
_send: Callable[[Mapping[str, Any]], Awaitable[None]]
scheme: str
path: str
query_string: str
headers: Mapping[str, Sequence[str]]
client: Optional[Tuple[str, int]]
server: Optional[Tuple[str, Optional[int]]]
def __init__(self,
scope: WebSocketScope,
receive: Callable[[], Awaitable[Any]],
send: Callable[[Mapping[str, Any]], Awaitable[None]]):
self._receive = receive
self._send = send
self.path = scope['path']
self.query_string = scope['query_string'].decode()
self.scheme = scope['scheme']
self.client = scope['client']
self.server = scope['server']
self.headers = decode_headers(scope['headers'])
async def accept(self) -> None:
await self._send({'type': 'websocket.accept'})
async def receive(self) -> WebSocketMessage:
message = await self._receive()
message_type: str = message['type']
if message_type == 'websocket.connect':
return await self.receive()
if message_type == 'websocket.receive':
if 'text' in message:
return WebSocketMessage(kind='text', data=message['text'])
elif 'bytes' in message:
return WebSocketMessage(kind='binary', data=message['bytes'])
else:
return WebSocketMessage(kind='close')
elif message_type == 'websocket.disconnect':
code: int = message.get('code', 1000)
return WebSocketMessage(kind='close', data=code)
else:
return WebSocketMessage(kind='close')
async def send_text(self, data: str) -> None:
await self._send({'type': 'websocket.send', 'text': data})
async def send_bytes(self, data: bytes) -> None:
await self._send({'type': 'websocket.send', 'bytes': data})
async def close(self, code: int = 1000) -> None:
await self._send({'type': 'websocket.close', 'code': code})
async def __anext__(self) -> WebSocketMessage:
message = await self.receive()
if message.kind == 'close':
raise StopAsyncIteration
return message
src/kaya/_http_method.py
+1
View File
@@ -9,3 +9,4 @@ class HttpMethod(StrEnum):
PUT = 'PUT'
DELETE = 'DELETE'
PATCH = 'PATCH'
WS = 'WS'
src/kaya/_path_handler.py
+7 -2
View File
@@ -2,10 +2,12 @@ from abc import ABC, abstractmethod
from typing import (
Sequence,
Dict,
Optional
Optional,
Union
)
from dataclasses import dataclass, field
from ._http_context import HttpContext
from ._websocket import WebSocket
@dataclass
@@ -18,10 +20,13 @@ class Matches:
unmatched_paths: Sequence[str] = field(default_factory=list)
type Context = Union[HttpContext, WebSocket]
class PathHandler(ABC):
@abstractmethod
async def handle_request(self, ctx: HttpContext, captured: Matches) -> None:
async def handle_request(self, ctx: Context, captured: Matches) -> None:
pass
@property
src/kaya/_rsgi.py
+73 -1
View File
@@ -14,12 +14,13 @@ from typing import (
cast
)
from granian._granian import RSGIHTTPProtocol, RSGIHTTPScope # type: ignore[attr-defined]
from granian._granian import RSGIHTTPProtocol, RSGIHTTPScope, RSGIWebsocketProtocol, RSGIWebsocketScope, RSGIWebsocketTransport # type: ignore[attr-defined]
from pwo import Maybe
from ._types import StrOrStrings
from ._http_context import HttpContext
from ._http_method import HttpMethod
from ._websocket import WebSocket, WebSocketMessage
class RsgiContext(HttpContext):
@@ -105,3 +106,74 @@ class RsgiContext(HttpContext):
async def send_empty(self, status: int, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
rearranged_headers = Maybe.of_nullable(headers).map(RsgiContext._rearrange_headers).or_else(list())
self.protocol.response_empty(status, rearranged_headers)
class RsgiWebSocket(WebSocket):
_protocol: RSGIWebsocketProtocol
_transport: Optional[RSGIWebsocketTransport]
scheme: str
path: str
query_string: str
headers: Mapping[str, Sequence[str]]
client: Optional[Tuple[str, int]]
server: Optional[Tuple[str, Optional[int]]]
def __init__(self, scope: RSGIWebsocketScope, protocol: RSGIWebsocketProtocol):
if not hasattr(protocol, 'accept') or not hasattr(protocol, 'close'):
raise RuntimeError(
'Granian was not configured for websockets; '
'ensure Granian is started with websocket support enabled'
)
self._protocol = protocol
self._transport = None
self.scheme = scope.scheme
self.path = scope.path
self.query_string = scope.query_string
def acc(d: Dict[str, List[str]], t: Tuple[str, str]) -> Dict[str, List[str]]:
d.setdefault(t[0].lower(), list()).append(t[1])
return d
fun = cast(Callable[[Mapping[str, Sequence[str]], tuple[str, str]], Mapping[str, Sequence[str]]], acc)
self.headers = reduce(fun, scope.headers.items(), {})
self.client = (Maybe.of(scope.client.split(':'))
.map(lambda it: (it[0], int(it[1])))
.or_else_throw(RuntimeError))
self.server = (Maybe.of(scope.server.split(':'))
.map(lambda it: (it[0], int(it[1])))
.or_else_throw(RuntimeError))
async def accept(self) -> None:
self._transport = await self._protocol.accept()
async def receive(self) -> WebSocketMessage:
if self._transport is None:
raise RuntimeError('WebSocket connection has not been accepted yet')
message = await self._transport.receive()
if message.kind == 0:
return WebSocketMessage(kind='close')
elif message.kind == 1:
return WebSocketMessage(kind='binary', data=message.data)
elif message.kind == 2:
return WebSocketMessage(kind='text', data=message.data)
else:
return WebSocketMessage(kind='close')
async def send_text(self, data: str) -> None:
if self._transport is None:
raise RuntimeError('WebSocket connection has not been accepted yet')
await self._transport.send_str(data)
async def send_bytes(self, data: bytes) -> None:
if self._transport is None:
raise RuntimeError('WebSocket connection has not been accepted yet')
await self._transport.send_bytes(data)
async def close(self, code: int = 1000) -> None:
self._protocol.close(code)
async def __anext__(self) -> WebSocketMessage:
message = await self.receive()
if message.kind == 'close':
raise StopAsyncIteration
return message
src/kaya/_tree.py
+3 -3
View File
@@ -18,7 +18,7 @@ from pwo import Maybe, index_of_with_escape
from ._http_context import HttpContext
from ._http_method import HttpMethod
from ._path_handler import PathHandler
from ._path_handler import PathHandler, Context
from ._path_matcher import PathMatcher, IntMatcher, GlobMatcher, StrMatcher, Node
from ._path_handler import Matches
from ._types import NodeType
@@ -135,11 +135,11 @@ class Tree:
def register(self,
path: str,
method: Optional[HttpMethod],
callback: Callable[[HttpContext, Unpack[Any]], Awaitable[None]],
callback: Callable[[Context, Unpack[Any]], Awaitable[None]],
recursive: bool) -> None:
class Handler(PathHandler):
async def handle_request(self, ctx: HttpContext, captured: Matches) -> None:
async def handle_request(self, ctx: Context, captured: Matches) -> None:
args = Maybe.of_nullable(captured.path).map(lambda it: [it]).or_else([])
await callback(ctx, *args, **captured.kwargs)
src/kaya/_types/rsgi.py
+14
View File
@@ -23,4 +23,18 @@ class HTTPScope(TypedDict):
path: str
query_string: str
headers: Mapping[str, str]
authority: Optional[str]
class WebSocketScope(TypedDict):
proto: Literal['ws']
rsgi_version: str
http_version: str
server: str
client: str
scheme: str
method: str
path: str
query_string: str
headers: Mapping[str, str]
authority: Optional[str]
src/kaya/_websocket.py
+47
View File
@@ -0,0 +1,47 @@
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import AsyncIterator, Literal, Mapping, Optional, Sequence, Tuple, Union
type WebSocketData = Optional[Union[str, bytes, int]]
@dataclass
class WebSocketMessage:
kind: Literal['text', 'binary', 'close']
data: WebSocketData = None
class WebSocket(ABC):
path: str
query_string: str
scheme: str
headers: Mapping[str, Sequence[str]]
client: Optional[Tuple[str, int]]
server: Optional[Tuple[str, Optional[int]]]
@abstractmethod
async def accept(self) -> None:
pass
@abstractmethod
async def receive(self) -> WebSocketMessage:
pass
@abstractmethod
async def send_text(self, data: str) -> None:
pass
@abstractmethod
async def send_bytes(self, data: bytes) -> None:
pass
@abstractmethod
async def close(self, code: int = 1000) -> None:
pass
def __aiter__(self) -> AsyncIterator[WebSocketMessage]:
return self
@abstractmethod
async def __anext__(self) -> WebSocketMessage:
pass
tests/test_websocket.py
+175
View File
@@ -0,0 +1,175 @@
import unittest
from typing import Any, Callable, Awaitable, List, Mapping, Optional
from pwo import async_test
from kaya import BugisApp, WebSocket, WebSocketMessage
def websocket_scope(path: str = '/ws') -> Mapping[str, Any]:
return {
'type': 'websocket',
'asgi': {'spec_version': '2.3', 'version': '3.0'},
'http_version': '1.1',
'scheme': 'ws',
'path': path,
'raw_path': path.encode(),
'query_string': b'',
'root_path': '',
'headers': [],
'client': ('127.0.0.1', 12345),
'server': ('127.0.0.1', 80),
'subprotocols': [],
'extensions': None,
}
class WebSocketTest(unittest.TestCase):
app: BugisApp
def setUp(self):
self.app = BugisApp()
@self.app.websocket('/echo')
async def echo(ws: WebSocket) -> None:
await ws.accept()
async for msg in ws:
if msg.kind == 'text':
await ws.send_text(f"echo: {msg.data}")
elif msg.kind == 'binary':
data = msg.data
assert isinstance(data, bytes)
await ws.send_bytes(data)
@self.app.websocket('/room/${room_id}')
async def room(ws: WebSocket, room_id: str) -> None:
await ws.accept()
async for msg in ws:
if msg.kind == 'text':
await ws.send_text(f"[{room_id}] {msg.data}")
@async_test
async def test_echo_text(self):
sent_messages: List[Mapping[str, Any]] = []
received_messages: List[Mapping[str, Any]] = []
async def receive() -> Mapping[str, Any]:
if not received_messages:
received_messages.append({'type': 'websocket.connect'})
return received_messages[-1]
if len(received_messages) == 1:
received_messages.append({'type': 'websocket.receive', 'text': 'hello'})
return received_messages[-1]
received_messages.append({'type': 'websocket.disconnect', 'code': 1000})
return received_messages[-1]
async def send(message: Mapping[str, Any]) -> None:
sent_messages.append(message)
await self.app(websocket_scope('/echo'), receive, send)
self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
self.assertEqual(sent_messages[1]['type'], 'websocket.send')
self.assertEqual(sent_messages[1]['text'], 'echo: hello')
@async_test
async def test_echo_binary(self):
sent_messages: List[Mapping[str, Any]] = []
received_messages: List[Mapping[str, Any]] = []
async def receive() -> Mapping[str, Any]:
if not received_messages:
received_messages.append({'type': 'websocket.connect'})
return received_messages[-1]
if len(received_messages) == 1:
received_messages.append({'type': 'websocket.receive', 'bytes': b'hello'})
return received_messages[-1]
received_messages.append({'type': 'websocket.disconnect', 'code': 1000})
return received_messages[-1]
async def send(message: Mapping[str, Any]) -> None:
sent_messages.append(message)
await self.app(websocket_scope('/echo'), receive, send)
self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
self.assertEqual(sent_messages[1]['type'], 'websocket.send')
self.assertEqual(sent_messages[1]['bytes'], b'hello')
@async_test
async def test_path_parameter(self):
sent_messages: List[Mapping[str, Any]] = []
received_messages: List[Mapping[str, Any]] = []
async def receive() -> Mapping[str, Any]:
if not received_messages:
received_messages.append({'type': 'websocket.connect'})
return received_messages[-1]
if len(received_messages) == 1:
received_messages.append({'type': 'websocket.receive', 'text': 'hi'})
return received_messages[-1]
received_messages.append({'type': 'websocket.disconnect', 'code': 1000})
return received_messages[-1]
async def send(message: Mapping[str, Any]) -> None:
sent_messages.append(message)
await self.app(websocket_scope('/room/general'), receive, send)
self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
self.assertEqual(sent_messages[1]['text'], '[general] hi')
@async_test
async def test_no_handler(self):
sent_messages: List[Mapping[str, Any]] = []
async def receive() -> Mapping[str, Any]:
return {'type': 'websocket.connect'}
async def send(message: Mapping[str, Any]) -> None:
sent_messages.append(message)
await self.app(websocket_scope('/unknown'), receive, send)
self.assertEqual(sent_messages[0]['type'], 'websocket.close')
self.assertEqual(sent_messages[0]['code'], 1000)
@async_test
async def test_close_from_client(self):
sent_messages: List[Mapping[str, Any]] = []
received_messages: List[Mapping[str, Any]] = []
async def receive() -> Mapping[str, Any]:
if not received_messages:
received_messages.append({'type': 'websocket.connect'})
return received_messages[-1]
received_messages.append({'type': 'websocket.disconnect', 'code': 1001})
return received_messages[-1]
async def send(message: Mapping[str, Any]) -> None:
sent_messages.append(message)
await self.app(websocket_scope('/echo'), receive, send)
self.assertEqual(sent_messages[0]['type'], 'websocket.accept')
self.assertEqual(len(sent_messages), 1)
class RsgiWebSocketTest(unittest.TestCase):
def test_misconfigured_granian(self):
from kaya._rsgi import RsgiWebSocket
class FakeScope:
scheme = 'ws'
path = '/ws'
query_string = ''
headers = {}
client = '127.0.0.1:12345'
server = '127.0.0.1:80'
class FakeProtocol:
pass
with self.assertRaises(RuntimeError) as ctx:
RsgiWebSocket(FakeScope(), FakeProtocol()) # type: ignore[arg-type]
self.assertIn('Granian was not configured for websockets', str(ctx.exception))