statewire 0.6.0__tar.gz → 0.6.2__tar.gz

{statewire-0.6.0 → statewire-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: statewire
-Version: 0.6.0
+Version: 0.6.2
 Summary: Replicate one JSON object over an SSE op stream + a command endpoint
 Project-URL: Repository, https://github.com/assistant-ui/harness-sdk
 License-Expression: MIT

{statewire-0.6.0 → statewire-0.6.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "statewire"
-version = "0.6.0"
+version = "0.6.2"
 description = "Replicate one JSON object over an SSE op stream + a command endpoint"
 readme = "README.md"
 license = "MIT"

{statewire-0.6.0 → statewire-0.6.2}/src/statewire/api.py

@@ -61,6 +61,12 @@ stream of that client id with ``fin {"reason": "superseded"}``. HTTP
 commands must carry the current token in ``Statewire-Lease`` (else 423);
 over WS, holding the connection is the lease.
 
+Subclass ``@route`` handlers are gated behind ``authorize`` automatically:
+the request is authorized before the handler body runs (an ``HTTPException``
+propagates as its status; a WebSocket handshake is denied with close code
+1008), and a handler that declares an ``identity`` parameter receives the
+authorized identity in it.
+
 Expected usage::
 
     class Thread(Statewire):
@@ -81,12 +87,14 @@ one atomic envelope, flushed when control returns to the event loop.
 
 import asyncio
 import contextlib
+import functools
 import inspect
 import json
 import logging
 import secrets
 import sys
 import time
+import typing
 from dataclasses import dataclass, field
 from typing import Any, AsyncIterator, Awaitable, Callable, Coroutine
 
@@ -444,6 +452,54 @@ def _reserved_paths() -> set[str]:
     return {spec.path for fn in vars(Statewire).values() for spec in routes_of(fn)}
 
 
+def _authorized_route(fn: Any) -> Any:
+    """Gate a subclass ``@route`` handler behind ``authorize``: the identity is
+    injected into an ``identity`` parameter when the handler declares one."""
+    is_ws = any(spec.method == "WEBSOCKET" for spec in routes_of(fn))
+    conn_name = "websocket" if is_ws else "request"
+    sig = inspect.signature(fn)
+    injects_identity = "identity" in sig.parameters
+    has_conn = conn_name in sig.parameters
+
+    @functools.wraps(fn)
+    async def wrapped(self: "Statewire", **kwargs: Any) -> Any:
+        connection = kwargs[conn_name] if has_conn else kwargs.pop(conn_name)
+        if is_ws:
+            try:
+                identity = await self.authorize(connection)
+            except HTTPException as exc:
+                detail = exc.detail if isinstance(exc.detail, str) else "unauthorized"
+                raise WebSocketException(code=1008, reason=detail) from exc
+        else:
+            identity = await self.authorize(connection)
+        if injects_identity:
+            kwargs["identity"] = identity
+        return await fn(self, **kwargs)
+
+    # Resolve hints against fn's own module so pinned's signature-driven mount
+    # sees real annotations, and expose a FastAPI-injectable connection param.
+    conn_type = WebSocket if is_ws else Request
+    hints = typing.get_type_hints(fn)
+    hints.setdefault(conn_name, conn_type)
+    params = [
+        p.replace(annotation=hints.get(p.name, p.annotation))
+        for p in sig.parameters.values()
+        if p.name != "identity"
+    ]
+    if not has_conn:
+        params.insert(
+            1,
+            inspect.Parameter(
+                conn_name,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=conn_type,
+            ),
+        )
+    wrapped.__signature__ = sig.replace(parameters=params)  # type: ignore[attr-defined]
+    wrapped.__annotations__ = {k: v for k, v in hints.items() if k != "identity"}
+    return wrapped
+
+
 class Statewire(PinnedAPI):
     stream_queue_size: int = 256
     heartbeat_interval: float = 15.0
@@ -458,13 +514,15 @@ class Statewire(PinnedAPI):
     def __init_subclass__(cls, **kwargs: Any) -> None:
         super().__init_subclass__(**kwargs)
         reserved = _reserved_paths()
-        for fn in vars(cls).values():
+        for name, fn in list(vars(cls).items()):
             for spec in routes_of(fn):
                 if spec.path in reserved:
                     raise TypeError(
                         f"{cls.__name__}.{fn.__name__} declares @route on reserved "
                         f"path {spec.path!r}; it is a Statewire protocol endpoint"
                     )
+            if routes_of(fn):
+                setattr(cls, name, _authorized_route(fn))
         # Full-MRO scan so plain mixins (no Statewire base) contribute their
         # @command handlers; the most derived registration of a method wins.
         merged: dict[str, str] = {}

statewire-0.6.2/tests/test_route_auth.py

@@ -0,0 +1,106 @@
+"""Contract: subclass ``@route`` handlers are gated behind ``authorize``.
+
+- an unauthorized request answers with the authorize HTTPException's status
+  and the handler never runs,
+- a handler declaring ``identity`` receives the authorized identity,
+- a handler without ``identity`` is still gated,
+- FastAPI body/query parameters on the handler keep working through the wrap,
+- a WebSocket subclass route denies the handshake when authorize rejects.
+"""
+
+from statewire_helpers import WSStream, statewire_client
+from fastapi import HTTPException, Request
+from pinned.routing import route
+
+from statewire import Statewire
+
+handled: list[str] = []
+
+
+class Domain(Statewire):
+    async def lifespan(self):
+        self.state = {"count": 0}
+        yield
+
+    async def authorize(self, request) -> str:
+        if request.headers.get("x-token") != "secret":
+            raise HTTPException(status_code=401, detail="unauthorized")
+        return request.headers.get("x-user")
+
+    @route.get("/whoami")
+    async def whoami(self, identity):
+        return {"identity": identity}
+
+    @route.post("/bump")
+    async def bump(self, payload: dict, by: int = 1):
+        handled.append("bump")
+        return {"count": self.state["count"] + by, "echo": payload}
+
+    @route.get("/peek")
+    async def peek(self, request: Request, identity):
+        return {"identity": identity, "path": request.url.path}
+
+    @route.websocket("/tap")
+    async def tap(self, websocket, identity):
+        await websocket.accept()
+        await websocket.send_json({"tapped": identity})
+        await websocket.close()
+
+
+async def test_unauthorized_route_answers_the_authorize_status():
+    async with statewire_client(Domain) as (app, client):
+        before = len(handled)
+        response = await client.post("/threads/t1/bump", json={"k": "v"})
+        assert response.status_code == 401
+        assert len(handled) == before
+
+
+async def test_identity_parameter_receives_the_authorized_identity():
+    async with statewire_client(Domain) as (app, client):
+        response = await client.get(
+            "/threads/t1/whoami", headers={"x-token": "secret", "x-user": "alice"}
+        )
+        assert response.status_code == 200
+        assert response.json() == {"identity": "alice"}
+
+
+async def test_gated_handler_without_identity_still_runs_authorized():
+    async with statewire_client(Domain) as (app, client):
+        response = await client.post(
+            "/threads/t1/bump", json={"k": "v"}, headers={"x-token": "secret"}
+        )
+        assert response.status_code == 200
+        assert response.json() == {"count": 1, "echo": {"k": "v"}}
+
+
+async def test_body_and_query_params_survive_the_wrap():
+    async with statewire_client(Domain) as (app, client):
+        response = await client.post(
+            "/threads/t1/bump",
+            params={"by": 5},
+            json={"k": "v"},
+            headers={"x-token": "secret"},
+        )
+        assert response.status_code == 200
+        assert response.json() == {"count": 5, "echo": {"k": "v"}}
+
+
+async def test_declared_request_param_reaches_the_handler():
+    async with statewire_client(Domain) as (app, client):
+        response = await client.get(
+            "/threads/t1/peek", headers={"x-token": "secret", "x-user": "bob"}
+        )
+        assert response.status_code == 200
+        assert response.json() == {"identity": "bob", "path": "/threads/t1/peek"}
+
+
+async def test_websocket_route_denies_the_handshake_when_unauthorized():
+    async with statewire_client(Domain) as (app, client):
+        async with WSStream(app, "/threads/t1/tap") as ws:
+            assert ws.accepted is False
+            assert ws.close_code == 1008
+        async with WSStream(
+            app, "/threads/t1/tap", headers={"x-token": "secret", "x-user": "carol"}
+        ) as ws:
+            assert ws.accepted is True
+            assert (await ws.next_frame()) == {"tapped": "carol"}