Add websocket session support to kaya-session

- kaya-core: WebSocket ABC gains session attribute and accept(headers=...)
- kaya-core: AsgiWebSocket injects headers into websocket.accept message
- kaya-rsgi: RsgiWebSocket accepts headers param (ignored — Granian's
  accept() takes no args)
- kaya-session: SessionWebSocket wrapper exposes ws.session and injects
  Set-Cookie on accept()
- kaya-session: SessionMixin registers before/after websocket hooks;
  session loaded at connect, persisted on close if modified
- 10 new WV session tests covering read, persist, handshake cookie,
  regenerate, invalidate, isolation
- Example and README updated

packages/kaya-core/src/kaya/core/_asgi.py

@@ -173,8 +173,13 @@ class AsgiWebSocket(WebSocket):
         self.server = scope['server']
         self.headers = decode_headers(scope['headers'])
 
-    async def accept(self) -> None:
-        await self._send({'type': 'websocket.accept'})
+    async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
+        message: Dict[str, Any] = {'type': 'websocket.accept'}
+        if headers is not None:
+            # Emit a list rather than a tuple: wsproto's AcceptConnection (used
+            # by httpx_ws and others) requires list concatenation.
+            message['headers'] = list(encode_headers(headers))
+        await self._send(message)
 
     async def receive(self) -> WebSocketMessage:
         message = await self._receive()

packages/kaya-core/src/kaya/core/_websocket.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import AsyncIterator, Literal, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, AsyncIterator, Literal, Mapping, Optional, Sequence, Tuple, Union
+
+from ._types.base import StrOrStrings
 
 type WebSocketData = Optional[Union[str, bytes, int]]
 
@@ -18,9 +20,10 @@ class WebSocket(ABC):
     headers: Mapping[str, Sequence[str]]
     client: Optional[Tuple[str, int]]
     server: Optional[Tuple[str, Optional[int]]]
+    session: Optional[Any] = None
 
     @abstractmethod
-    async def accept(self) -> None:
+    async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
         pass
 
     @abstractmethod

packages/kaya-core/tests/test_websocket.py

@@ -4,6 +4,7 @@ from pwo import async_test
 from httpx_ws import aconnect_ws, WebSocketDisconnect
 from httpx_ws.transport import ASGIWebSocketTransport
 from kaya.core import KayaApp, WebSocket
+from kaya.core._asgi import AsgiWebSocket
 
 
 class WebSocketTest(unittest.TestCase):
@@ -78,3 +79,55 @@ class WebSocketTest(unittest.TestCase):
         async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
             async with aconnect_ws("/echo", client) as ws:
                 pass
+
+    @async_test
+    async def test_accept_with_headers(self):
+        sent_messages = []
+
+        async def send(message):
+            sent_messages.append(message)
+
+        async def receive():
+            return {'type': 'websocket.connect'}
+
+        scope = {
+            'type': 'websocket',
+            'path': '/echo',
+            'query_string': b'',
+            'scheme': 'ws',
+            'client': ('127.0.0.1', 12345),
+            'server': ('127.0.0.1', 80),
+            'headers': [],
+        }
+        ws = AsgiWebSocket(scope, receive, send)
+        await ws.accept(headers={'Set-Cookie': 'sid=abc; Path=/', 'X-Custom': ('a', 'b')})
+        self.assertEqual(1, len(sent_messages))
+        message = sent_messages[0]
+        self.assertEqual('websocket.accept', message['type'])
+        self.assertIn((b'Set-Cookie', b'sid=abc; Path=/'), message['headers'])
+        self.assertIn((b'X-Custom', b'a'), message['headers'])
+        self.assertIn((b'X-Custom', b'b'), message['headers'])
+
+    @async_test
+    async def test_accept_without_headers(self):
+        sent_messages = []
+
+        async def send(message):
+            sent_messages.append(message)
+
+        async def receive():
+            return {'type': 'websocket.connect'}
+
+        scope = {
+            'type': 'websocket',
+            'path': '/echo',
+            'query_string': b'',
+            'scheme': 'ws',
+            'client': ('127.0.0.1', 12345),
+            'server': ('127.0.0.1', 80),
+            'headers': [],
+        }
+        ws = AsgiWebSocket(scope, receive, send)
+        await ws.accept()
+        self.assertEqual(1, len(sent_messages))
+        self.assertEqual({'type': 'websocket.accept'}, sent_messages[0])