Refactor to composable KayaMixin architecture

Replace wrapper-based SessionMiddleware/OIDCApp with KayaMixin subclasses
applied via KayaApp(mixins=[...]). Mixins hook into handle_request and
handle_websocket via before/after hooks, so both ASGI and RSGI keep working.
Mixin dependencies are applied automatically and deduplicated.
This commit is contained in:
2026-07-23 22:09:59 +08:00
parent 24a797e3d2
commit 3ebf079533
14 changed files with 480 additions and 359 deletions
@@ -1,4 +1,4 @@
from ._middleware import SessionMiddleware
from ._mixin import SessionMixin
from ._session import Session
from ._store import InMemorySessionStore, SessionStore
@@ -6,6 +6,6 @@ from ._store import InMemorySessionStore, SessionStore
__all__ = [
'InMemorySessionStore',
'Session',
'SessionMiddleware',
'SessionMixin',
'SessionStore',
]
@@ -1,178 +0,0 @@
from typing import Any, Awaitable, Callable, MutableMapping, Optional, Sequence, cast
from kaya.core import HttpMethod, KayaApp
from ._cookie import format_set_cookie, parse_cookie_value
from ._session import Session
from ._store import SessionStore
type HttpHandler = Callable[..., Awaitable[None]]
type WebSocketHandler = Callable[..., Awaitable[None]]
type RouteDecorator = Callable[[HttpHandler], HttpHandler]
type WebSocketDecorator = Callable[[WebSocketHandler], WebSocketHandler]
type ASGIApp = Callable[
[MutableMapping[str, Any], Callable[[], Awaitable[Any]], Callable[[MutableMapping[str, Any]], Awaitable[None]]],
Awaitable[None],
]
class SessionMiddleware:
"""ASGI middleware providing server-side HTTP sessions.
The middleware reads a session cookie from each HTTP request, loads the
session data via the configured store, and makes the session available to
Kaya handlers as ``ctx.session``. After the request it persists the session
and refreshes the cookie.
Routing methods are delegated to the wrapped ``KayaApp`` so the middleware
can be used as a drop-in replacement when registering handlers.
Example::
app = KayaApp()
session_app = SessionMiddleware(app, InMemorySessionStore())
@session_app.GET('/')
async def home(ctx: HttpContext):
ctx.session['visits'] = ctx.session.get('visits', 0) + 1
await ctx.send_str(200, f"visits: {ctx.session['visits']}")
"""
def __init__(
self,
app: KayaApp,
store: SessionStore,
cookie_name: str = 'session_id',
path: str = '/',
max_age: Optional[int] = 14 * 24 * 60 * 60,
httponly: bool = True,
secure: bool = False,
samesite: Optional[str] = 'Lax',
) -> None:
self._app = app
self._store = store
self._cookie_name = cookie_name
self._path = path
self._max_age = max_age
self._httponly = httponly
self._secure = secure
self._samesite = samesite
def route(
self,
paths: str | Sequence[str],
methods: Optional[HttpMethod | Sequence[HttpMethod]] = None,
recursive: bool = False,
) -> RouteDecorator:
return self._app.route(paths, methods, recursive)
def GET(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.GET(path, recursive)
def POST(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.POST(path, recursive)
def PUT(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.PUT(path, recursive)
def DELETE(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.DELETE(path, recursive)
def OPTIONS(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.OPTIONS(path, recursive)
def HEAD(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.HEAD(path, recursive)
def PATCH(self, path: str, recursive: bool = False) -> RouteDecorator:
return self._app.PATCH(path, recursive)
def websocket(self, path: str, recursive: bool = False) -> WebSocketDecorator:
return self._app.websocket(path, recursive)
async def __call__(
self,
scope: MutableMapping[str, Any],
receive: Callable[[], Awaitable[Any]],
send: Callable[[MutableMapping[str, Any]], Awaitable[None]],
) -> None:
if scope['type'] != 'http':
await cast(ASGIApp, self._app)(scope, receive, send)
return
session_id = self._extract_session_id(scope)
session: Session
if session_id is not None:
loaded = await self._store.load(session_id, self._max_age)
session = loaded if loaded is not None else Session()
else:
session = Session()
state = scope.get('state')
if not isinstance(state, dict):
state = {}
scope['state'] = state
state['kaya_session'] = session
async def wrapped_send(message: MutableMapping[str, Any]) -> None:
if message['type'] == 'http.response.start':
message = dict(message)
final_session_id = self._finalize_session_id(session)
if final_session_id is not None:
cookie_value = format_set_cookie(
self._cookie_name,
final_session_id,
path=self._path,
max_age=0 if session.invalidated else self._max_age,
httponly=self._httponly,
secure=self._secure,
samesite=self._samesite,
)
headers = list(message.get('headers', []))
headers.append((b'Set-Cookie', cookie_value.encode()))
message['headers'] = headers
await send(message)
try:
await cast(ASGIApp, self._app)(scope, receive, wrapped_send)
finally:
await self._persist(session)
def _extract_session_id(self, scope: MutableMapping[str, Any]) -> Optional[str]:
headers = scope.get('headers', [])
for key, value in headers:
key_bytes: bytes = key if isinstance(key, bytes) else key.encode()
if key_bytes.lower() == b'cookie':
value_str: str = value.decode() if isinstance(value, bytes) else value
return parse_cookie_value(value_str, self._cookie_name)
return None
def _finalize_session_id(self, session: Session) -> Optional[str]:
if session.invalidated:
return session.id
if session.id is None:
if session.modified or session.regenerate:
session.set_id(self._store.new_session_id())
elif session.regenerate:
session._old_id = session.id
session.set_id(self._store.new_session_id())
session._regenerate = False
return session.id
async def _persist(self, session: Session) -> None:
if session.invalidated:
old_id = session._old_id or session.id
if old_id is not None:
await self._store.delete(old_id)
return
if session._old_id is not None and session._old_id != session.id:
await self._store.delete(session._old_id)
session._old_id = None
if session.id is None and session.modified:
session.set_id(self._store.new_session_id())
if session.id is not None:
await self._store.save(session.id, session, self._max_age)
@@ -0,0 +1,200 @@
from pathlib import Path
from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Optional, Sequence
from kaya.core import HttpContext, KayaApp, KayaMixin
from kaya.core._types import StrOrStrings
from ._cookie import format_set_cookie, parse_cookie_value
from ._session import Session
from ._store import SessionStore
class SessionHttpContext(HttpContext):
"""HttpContext wrapper that exposes ``session`` and injects the session
cookie into response headers.
Works with any concrete ``HttpContext`` (ASGI or RSGI) because it only
relies on the abstract send methods, which all implementations share.
"""
def __init__(
self,
ctx: HttpContext,
session: Session,
cookie_injector: Callable[[], Optional[str]],
) -> None:
self._ctx = ctx
self.session = session
self._cookie_injector = cookie_injector
self.pathsend = ctx.pathsend
self.receive = ctx.receive
self.send = ctx.send
self.scheme = ctx.scheme
self.method = ctx.method
self.path = ctx.path
self.query_string = ctx.query_string
self.headers = ctx.headers
self.client = ctx.client
self.server = ctx.server
self.request_body = ctx.request_body
def _inject_cookie(self, headers: Optional[Mapping[str, StrOrStrings]]) -> Optional[Mapping[str, StrOrStrings]]:
cookie_value = self._cookie_injector()
if cookie_value is None:
return headers
new_headers: dict[str, StrOrStrings] = dict(headers) if headers else {}
existing = new_headers.get('Set-Cookie')
if existing is None:
new_headers['Set-Cookie'] = cookie_value
elif isinstance(existing, str):
new_headers['Set-Cookie'] = (existing, cookie_value)
else:
new_headers['Set-Cookie'] = (*existing, cookie_value)
return new_headers
async def stream_body(self,
status: int,
body_generator: AsyncGenerator[bytes, None],
headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.stream_body(status, body_generator, self._inject_cookie(headers))
async def send_bytes(self, status: int, body: bytes, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_bytes(status, body, self._inject_cookie(headers))
async def send_str(self, status: int, body: str, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_str(status, body, self._inject_cookie(headers))
async def send_file(self, status: int, path: Path, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_file(status, path, self._inject_cookie(headers))
async def send_empty(self, status: int, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
await self._ctx.send_empty(status, self._inject_cookie(headers))
class _CookieInjector:
"""Computes the Set-Cookie value once (on first response) and caches it."""
def __init__(self, mixin: 'SessionMixin', session: Session) -> None:
self._mixin = mixin
self._session = session
self._value: Optional[str] = None
self._computed = False
def __call__(self) -> Optional[str]:
if not self._computed:
self._value = self._mixin._compute_cookie(self._session)
self._computed = True
return self._value
class SessionMixin(KayaMixin):
"""Kaya mixin providing server-side HTTP sessions.
Registers before/after request hooks that load and persist the session and
injects the session cookie into responses via a wrapped ``HttpContext``.
Because the app stays a ``KayaApp``, both ASGI and RSGI keep working.
Example::
session = SessionMixin(InMemorySessionStore())
app = KayaApp(mixins=[session])
@app.GET('/')
async def home(ctx: HttpContext):
ctx.session['visits'] = ctx.session.get('visits', 0) + 1
await ctx.send_str(200, f"visits: {ctx.session['visits']}")
"""
def __init__(
self,
store: SessionStore,
cookie_name: str = 'session_id',
path: str = '/',
max_age: Optional[int] = 14 * 24 * 60 * 60,
httponly: bool = True,
secure: bool = False,
samesite: Optional[str] = 'Lax',
) -> None:
self._store = store
self._cookie_name = cookie_name
self._path = path
self._max_age = max_age
self._httponly = httponly
self._secure = secure
self._samesite = samesite
def apply(self, app: KayaApp) -> None:
app.add_before_request_hook(self._before_request)
app.add_after_request_hook(self._after_request)
async def _before_request(self, ctx: HttpContext) -> Optional[HttpContext]:
session_id = self._extract_session_id(ctx)
session: Session
if session_id is not None:
loaded = await self._store.load(session_id, self._max_age)
session = loaded if loaded is not None else Session()
else:
session = Session()
injector = _CookieInjector(self, session)
return SessionHttpContext(ctx, session, injector)
async def _after_request(self, ctx: HttpContext) -> None:
session = ctx.session
if not isinstance(session, Session):
return
await self._persist(session)
def _extract_session_id(self, ctx: HttpContext) -> Optional[str]:
cookie_header_values = ctx.headers.get('cookie')
if cookie_header_values is None:
return None
if isinstance(cookie_header_values, str):
return parse_cookie_value(cookie_header_values, self._cookie_name)
for value in cookie_header_values:
found = parse_cookie_value(value, self._cookie_name)
if found is not None:
return found
return None
def _compute_cookie(self, session: Session) -> Optional[str]:
final_session_id = self._finalize_session_id(session)
if final_session_id is None:
return None
return format_set_cookie(
self._cookie_name,
final_session_id,
path=self._path,
max_age=0 if session.invalidated else self._max_age,
httponly=self._httponly,
secure=self._secure,
samesite=self._samesite,
)
def _finalize_session_id(self, session: Session) -> Optional[str]:
if session.invalidated:
return session.id
if session.id is None:
if session.modified or session.regenerate:
session.set_id(self._store.new_session_id())
elif session.regenerate:
session._old_id = session.id
session.set_id(self._store.new_session_id())
session._regenerate = False
return session.id
async def _persist(self, session: Session) -> None:
if session.invalidated:
old_id = session._old_id or session.id
if old_id is not None:
await self._store.delete(old_id)
return
if session._old_id is not None and session._old_id != session.id:
await self._store.delete(session._old_id)
session._old_id = None
if session.id is None and session.modified:
session.set_id(self._store.new_session_id())
if session.id is not None:
await self._store.save(session.id, session, self._max_age)