statewire 0.3.1__tar.gz → 0.3.2__tar.gz

{statewire-0.3.1 → statewire-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: statewire
-Version: 0.3.1
+Version: 0.3.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.3.1 → statewire-0.3.2}/pyproject.toml

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

{statewire-0.3.1 → statewire-0.3.2}/src/statewire/api.py

@@ -36,6 +36,16 @@ body over HTTP, as a ``res`` on the stream over WS.
 Handlers that declare a keyword-only ``ctx`` parameter receive a
 ``StatewireCommandContext``: ``ctx.caller`` is the submitting client's
 handle (holdable across a long run), ``ctx.ack()`` the invocation ack.
+``authorize`` gates every ingress and returns the caller's identity
+(``None`` = anonymous); the server stamps it on the client record per
+attach and per admitted command submission — never client-supplied,
+never on a rejected ingress — and drops it with the record. Handlers
+read it as ``ctx.caller.identity``: the current session owner, live —
+a newer authorized attach re-stamps it, so per-command trust decisions
+read it before the first await. A client id is a within-thread session
+id, not an authorization boundary: ``authorize`` must gate the thread
+itself, since any authorized principal knowing a client id can attach,
+receive the snapshot, and supersede that session.
 
 The host also owns a scheduler: ``self.schedule(fn)`` runs ``fn`` at the
 next drain (duplicate schedules coalesce), ``self.drain()`` drains now.
@@ -161,6 +171,7 @@ class _Client:
         self.results: dict[int, tuple[float, dict[str, Any]]] = {}
         self.inflight: dict[int, _Invocation] = {}
         self.context: dict[str, Any] | None = None
+        self.identity: Any = None
 
     def settle(self, seq: int, response: dict[str, Any] | None = None) -> None:
         full = {"seq": seq, **(response if response is not None else {"type": "accepted"})}
@@ -207,8 +218,9 @@ class _Client:
 class StatewireClientHandle:
     """Read-only view of the submitting client, reachable as ``ctx.caller``
     in a ``@command`` handler. Holdable across a long run: ``context`` reads
-    the client's current stored context live, ``is_connected`` reports
-    whether the client has an active stream attachment."""
+    the client's current stored context live, ``identity`` the current
+    session owner's identity, and ``is_connected`` reports whether the
+    client has an active stream attachment."""
 
     def __init__(self, host: "Statewire", client: _Client) -> None:
         self._host = host
@@ -222,6 +234,13 @@ class StatewireClientHandle:
     def context(self) -> dict[str, Any] | None:
         return self._client.context
 
+    @property
+    def identity(self) -> Any:
+        """Current session owner's identity — live, re-stamped by a newer
+        authorized attach; per-command trust decisions read it before the
+        handler's first await."""
+        return self._client.identity
+
     @property
     def is_connected(self) -> bool:
         client = self._client
@@ -592,14 +611,22 @@ class Statewire(PinnedAPI):
 
     # ─── GET /stream ────────────────────────────────────────
 
-    async def authorize(self, request: HTTPConnection) -> None:
-        """Auth seam: override to gate /stream, /commands and /ws; raise HTTPException to reject."""
+    async def authorize(self, request: HTTPConnection) -> Any:
+        """Auth seam: override to gate /stream, /commands and /ws; raise
+        HTTPException to reject. The return value is the caller's identity —
+        any host-chosen object, ``None`` for anonymous — stamped on the client
+        record at each attach and each admitted command submission (fresh
+        identity wins) and read by handlers via ``ctx.caller.identity``."""
+        return None
 
-    def _register(self, client_id: str) -> tuple[asyncio.Queue[Any], bytes]:
+    def _register(
+        self, client_id: str, identity: Any
+    ) -> tuple[asyncio.Queue[Any], bytes]:
         client = self._clients.get(client_id)
         last = client.last_seq if client is not None else -1
         if client is None:
             client = self._clients[client_id] = _Client(self, client_id)
+        client.identity = identity
         client.touched = time.monotonic()
         self._flush()
         self._evict_streams(
@@ -653,9 +680,9 @@ class Statewire(PinnedAPI):
 
     @route.get("/stream")
     async def stream(self, request: Request) -> Response:
-        await self.authorize(request)
+        identity = await self.authorize(request)
         client_id = _client_id_of(request)
-        q, snapshot = self._register(client_id)
+        q, snapshot = self._register(client_id, identity)
         return StreamingResponse(
             self._stream_events(q, snapshot),
             media_type="text/event-stream",
@@ -710,7 +737,7 @@ class Statewire(PinnedAPI):
 
     @route.post("/commands")
     async def commands(self, request: Request) -> Any:
-        await self.authorize(request)
+        identity = await self.authorize(request)
         client_id = _client_id_of(request)
         seq = _command_seq_of(request)
         lease = request.headers.get(LEASE_HEADER)
@@ -719,7 +746,7 @@ class Statewire(PinnedAPI):
             body = json.loads(raw)
         except (json.JSONDecodeError, UnicodeDecodeError):
             raise HTTPException(status_code=400, detail="body must be JSON") from None
-        return await self._submit(client_id, seq, body, lease)
+        return await self._submit(client_id, seq, body, lease, identity)
 
     def _validate_submission(
         self, body: Any
@@ -788,7 +815,12 @@ class Statewire(PinnedAPI):
         return _Prepared(method, handler, params, wants_ctx), None
 
     async def _admit(
-        self, client_id: str, seq: int, body: Any, lease: str | None
+        self,
+        client_id: str,
+        seq: int,
+        body: Any,
+        lease: str | None,
+        identity: Any = _UNSET,
     ) -> tuple[str, Any]:
         client = self._clients.get(client_id)
         if client is None:
@@ -803,6 +835,10 @@ class Statewire(PinnedAPI):
         context, members, rejections = self._validate_submission(body)
         if rejections:
             return "rejected", rejections
+        # Only an admitted submission may re-stamp; _UNSET keeps the
+        # attach-time identity (WS frames, facade dispatch).
+        if identity is not _UNSET:
+            client.identity = identity
         if context is not _UNSET:
             client.context = context
         invocations = self._start_batch(client, seq, members)
@@ -814,9 +850,9 @@ class Statewire(PinnedAPI):
         return "executed", None
 
     async def _submit(
-        self, client_id: str, seq: int, body: Any, lease: str | None
+        self, client_id: str, seq: int, body: Any, lease: str | None, identity: Any
     ) -> Any:
-        outcome, detail = await self._admit(client_id, seq, body, lease)
+        outcome, detail = await self._admit(client_id, seq, body, lease, identity)
         match outcome:
             case "unknown-client":
                 return JSONResponse({"error": "unknown-client"}, status_code=412)
@@ -970,7 +1006,7 @@ class Statewire(PinnedAPI):
     @route.websocket("/ws")
     async def ws(self, websocket: WebSocket) -> None:
         try:
-            await self.authorize(websocket)
+            identity = await self.authorize(websocket)
         except HTTPException as exc:
             reason = exc.detail if isinstance(exc.detail, str) else "unauthorized"
             raise WebSocketException(code=1008, reason=reason) from None
@@ -982,7 +1018,7 @@ class Statewire(PinnedAPI):
         await websocket.accept(
             subprotocol=ops.WS_SUBPROTOCOL if ops.WS_SUBPROTOCOL in offered else None
         )
-        q, snapshot = self._register(client_id)
+        q, snapshot = self._register(client_id, identity)
         lease = self._clients[client_id].lease
         try:
             await self._ws_pump(websocket, q, client_id, snapshot, lease)

{statewire-0.3.1 → statewire-0.3.2}/src/statewire/assistant_transport.py

@@ -270,7 +270,7 @@ class AssistantTransport(Statewire):
 
     @route.post("/assistant-transport/api/chat")
     async def assistant_transport_chat(self, request: Request) -> Response:
-        await self.authorize(request)
+        identity = await self.authorize(request)
         fmt = _FORMATS[self.assistant_transport_protocol]
         raw = await request.body()
         try:
@@ -299,7 +299,7 @@ class AssistantTransport(Statewire):
                 status_code=400, detail=f"invalid context: {exc}"
             ) from None
         client_id = "at-" + secrets.token_urlsafe(9)
-        q, snapshot = self._register(client_id)
+        q, snapshot = self._register(client_id, identity)
         self._clients[client_id].context = context
         lease = self._clients[client_id].lease
         assert lease is not None
@@ -352,13 +352,13 @@ class AssistantTransport(Statewire):
 
     @route.post("/assistant-transport/api/cancel")
     async def assistant_transport_cancel(self, request: Request) -> dict[str, Any]:
-        await self.authorize(request)
+        identity = await self.authorize(request)
         run = self._at_run
         found = run is not None and not run.done
         if found:
             await self.on_assistant_transport_disconnect()
         if "run/stop" in type(self)._statewire_commands:
-            await self._assistant_transport_dispatch("run/stop")
+            await self._assistant_transport_dispatch("run/stop", identity)
         return {"success": True, "found": found}
 
     @route.post("/assistant-transport/api/status")
@@ -457,9 +457,9 @@ class AssistantTransport(Statewire):
             raise _RunError(f"reload: no child message for parentId {parent_id!r}")
         return {"method": "run/reload", "params": [{"sourceId": source_id}]}
 
-    async def _assistant_transport_dispatch(self, method: str) -> None:
+    async def _assistant_transport_dispatch(self, method: str, identity: Any) -> None:
         client_id = "at-" + secrets.token_urlsafe(9)
-        q, _ = self._register(client_id)
+        q, _ = self._register(client_id, identity)
         try:
             lease = self._clients[client_id].lease
             assert lease is not None

statewire-0.3.2/tests/test_authorize.py

@@ -0,0 +1,256 @@
+from statewire_helpers import post_command, statewire_client, stream_of, ws_of
+from fastapi import HTTPException, Request
+
+from statewire import AssistantTransport, Statewire, command
+
+
+instances = []
+runs = []
+
+
+class Guarded(Statewire):
+    async def lifespan(self):
+        instances.append(self)
+        self.state = {"count": 0}
+        yield
+
+    async def authorize(self, request: Request) -> None:
+        if request.headers.get("x-token") != "secret":
+            raise HTTPException(status_code=401, detail="unauthorized")
+
+    @command
+    async def inc(self):
+        runs.append("inc")
+        self.state["count"] += 1
+
+
+async def test_stream_attach_without_token_is_401_and_registers_nothing():
+    async with statewire_client(Guarded) as (app, client):
+        async with stream_of(app) as stream:
+            assert stream.status == 401
+        assert instances[-1]._subscribers == {}
+
+
+async def test_command_without_token_is_401_and_handler_never_runs():
+    async with statewire_client(Guarded) as (app, client):
+        async with stream_of(app, headers={"x-token": "secret"}) as stream:
+            await stream.next_event()
+        before = len(runs)
+        response = await post_command(client, {"method": "inc", "params": []}, seq=1)
+        assert response.status_code == 401
+        assert len(runs) == before
+
+
+async def test_authorized_requests_succeed():
+    async with statewire_client(Guarded) as (app, client):
+        async with stream_of(app, headers={"x-token": "secret"}) as stream:
+            assert stream.status == 200
+            await stream.next_event()
+            response = await post_command(
+                client, {"method": "inc", "params": []}, seq=1, headers={"x-token": "secret"}
+            )
+            assert response.status_code == 200
+            assert await stream.next_event() == {
+                "ops": [{"op": "replace", "path": ["count"], "value": 1}],
+                "ack": 1,
+            }
+
+
+class Identified(Statewire):
+    caller = None
+
+    async def lifespan(self):
+        self.state = {}
+        yield
+
+    async def authorize(self, request):
+        return request.headers.get("x-user")
+
+    @command
+    async def whoami(self, *, ctx):
+        self.caller = ctx.caller
+        return {"identity": ctx.caller.identity}
+
+
+async def instance_of(app, thread="t1"):
+    return (await app.state.pinned_host.directory.get(thread)).instance
+
+
+async def whoami_res(client, *, seq, headers=None):
+    body = {"method": "whoami", "params": []}
+    first = await post_command(client, body, seq=seq, headers=headers)
+    assert first.status_code == 200
+    dup = await post_command(client, body, seq=seq, headers=headers)
+    return dup.json()["res"]
+
+
+async def test_identity_returned_by_authorize_is_visible_via_ctx_caller():
+    async with statewire_client(Identified) as (app, client):
+        async with stream_of(app, headers={"x-user": "alice"}) as stream:
+            await stream.next_event()
+            res = await whoami_res(client, seq=1, headers={"x-user": "alice"})
+            assert res["payload"] == {"identity": "alice"}
+
+
+async def test_identity_is_visible_over_ws():
+    async with statewire_client(Identified) as (app, client):
+        async with ws_of(app, headers={"x-user": "wanda"}) as ws:
+            await ws.next_frame()
+            env = await ws.command({"method": "whoami", "params": []})
+            assert env["res"] == [
+                {"seq": 1, "type": "accepted", "payload": {"identity": "wanda"}}
+            ]
+
+
+async def test_reattach_restamps_the_identity_on_a_held_handle():
+    async with statewire_client(Identified) as (app, client):
+        async with stream_of(app, headers={"x-user": "alice"}) as stream:
+            await stream.next_event()
+            res = await whoami_res(client, seq=1, headers={"x-user": "alice"})
+            assert res["payload"] == {"identity": "alice"}
+        instance = await instance_of(app)
+        assert instance.caller.identity == "alice"
+        async with stream_of(app, headers={"x-user": "bob"}) as stream:
+            await stream.next_event()
+            assert instance.caller.identity == "bob"
+
+
+async def test_identity_defaults_to_none_when_authorize_is_not_overridden():
+    class Anonymous(Statewire):
+        async def lifespan(self):
+            self.state = {}
+            yield
+
+        @command
+        async def whoami(self, *, ctx):
+            return {"identity": ctx.caller.identity}
+
+    async with statewire_client(Anonymous) as (app, client):
+        async with stream_of(app) as stream:
+            await stream.next_event()
+            body = {"method": "whoami", "params": []}
+            assert (await post_command(client, body, seq=1)).status_code == 200
+            dup = await post_command(client, body, seq=1)
+            assert dup.json()["res"]["payload"] == {"identity": None}
+
+
+async def test_forged_client_id_with_bad_lease_cannot_overwrite_identity():
+    async with statewire_client(Identified) as (app, client):
+        async with stream_of(app, headers={"x-user": "alice"}) as stream:
+            await stream.next_event()
+            res = await whoami_res(client, seq=1, headers={"x-user": "alice"})
+            assert res["payload"] == {"identity": "alice"}
+            instance = await instance_of(app)
+            for lease in (None, "stolen"):
+                response = await post_command(
+                    client,
+                    {"method": "whoami", "params": []},
+                    seq=2,
+                    headers={"x-user": "mallory"},
+                    lease=lease,
+                )
+                assert response.status_code == 423
+                assert instance.caller.identity == "alice"
+
+
+async def test_unadmitted_lease_valid_submissions_do_not_restamp_identity():
+    async with statewire_client(Identified) as (app, client):
+        async with stream_of(app, headers={"x-user": "alice"}) as stream:
+            await stream.next_event()
+            res = await whoami_res(client, seq=1, headers={"x-user": "alice"})
+            assert res["payload"] == {"identity": "alice"}
+            instance = await instance_of(app)
+            duplicate = await post_command(
+                client,
+                {"method": "whoami", "params": []},
+                seq=1,
+                headers={"x-user": "mallory"},
+            )
+            assert duplicate.status_code == 200
+            assert instance.caller.identity == "alice"
+            gap = await post_command(
+                client,
+                {"method": "whoami", "params": []},
+                seq=9,
+                headers={"x-user": "mallory"},
+            )
+            assert gap.status_code == 409
+            assert instance.caller.identity == "alice"
+            rejected = await post_command(
+                client,
+                {"method": "no-such-method", "params": []},
+                seq=2,
+                headers={"x-user": "mallory"},
+            )
+            assert rejected.status_code == 422
+            assert instance.caller.identity == "alice"
+
+
+async def test_identity_is_not_settable_via_context_updates():
+    async with statewire_client(Identified) as (app, client):
+        async with stream_of(app, headers={"x-user": "alice"}) as stream:
+            await stream.next_event()
+            response = await post_command(
+                client,
+                {"context": [{"op": "replace", "path": [], "value": {"identity": "hacker"}}]},
+                seq=1,
+                headers={"x-user": "alice"},
+            )
+            assert response.status_code == 200
+            res = await whoami_res(client, seq=2, headers={"x-user": "alice"})
+            assert res["payload"] == {"identity": "alice"}
+
+
+class LegacyIdentified(AssistantTransport):
+    seen = None
+    stopped = None
+
+    async def lifespan(self):
+        self.state = {}
+        yield
+
+    async def authorize(self, request):
+        return request.headers.get("x-user")
+
+    @command("run/steer")
+    async def steer(self, params, *, ctx):
+        self.seen = ctx.caller.identity
+
+    @command("run/stop")
+    async def stop(self, *, ctx):
+        self.stopped = ctx.caller.identity
+
+
+async def test_shim_run_client_carries_the_identity_into_translated_commands():
+    async with statewire_client(LegacyIdentified) as (app, client):
+        response = await client.post(
+            "/threads/t1/assistant-transport/api/chat",
+            headers={"x-user": "legacy"},
+            json={
+                "threadId": "t1",
+                "commands": [
+                    {
+                        "type": "add-message",
+                        "message": {
+                            "role": "user",
+                            "parts": [{"type": "text", "text": "hi"}],
+                        },
+                    }
+                ],
+            },
+        )
+        assert response.status_code == 200
+        instance = await instance_of(app)
+        assert instance.seen == "legacy"
+
+
+async def test_shim_cancel_dispatch_carries_the_identity():
+    async with statewire_client(LegacyIdentified) as (app, client):
+        response = await client.post(
+            "/threads/t1/assistant-transport/api/cancel",
+            headers={"x-user": "canceller"},
+            json={"threadId": "t1"},
+        )
+        assert response.status_code == 200
+        instance = await instance_of(app)
+        assert instance.stopped == "canceller"

{statewire-0.3.1 → statewire-0.3.2}/tests/test_state_proxy.py

@@ -258,7 +258,7 @@ async def test_get_value_matches_assistant_streams_unwrap_pattern():
 
 async def test_flush_delivers_the_pending_envelope_synchronously():
     inst = make({"a": 1})
-    q, _ = inst._register("c1")
+    q, _ = inst._register("c1", None)
     inst.state["a"] = 2
     assert q.empty()
     inst.flush()

{statewire-0.3.1 → statewire-0.3.2}/tests/test_writer_lease.py

@@ -87,7 +87,7 @@ async def test_ws_command_on_a_superseded_socket_closes_1008_without_executing()
             await ws.next_frame()
             record = await app.state.pinned_host.directory.get("t1")
             await ws.send_frame({"method": "inc", "params": [], "seq": 1})
-            record.instance._register("c1")
+            record.instance._register("c1", None)
             while await ws.next_frame() is not None:
                 pass
             assert ws.close_code == 1008

statewire-0.3.1/tests/test_authorize.py

@@ -1,56 +0,0 @@
-from statewire_helpers import post_command, statewire_client, stream_of
-from fastapi import HTTPException, Request
-
-from statewire import Statewire, command
-
-
-instances = []
-runs = []
-
-
-class Guarded(Statewire):
-    async def lifespan(self):
-        instances.append(self)
-        self.state = {"count": 0}
-        yield
-
-    async def authorize(self, request: Request) -> None:
-        if request.headers.get("x-token") != "secret":
-            raise HTTPException(status_code=401, detail="unauthorized")
-
-    @command
-    async def inc(self):
-        runs.append("inc")
-        self.state["count"] += 1
-
-
-async def test_stream_attach_without_token_is_401_and_registers_nothing():
-    async with statewire_client(Guarded) as (app, client):
-        async with stream_of(app) as stream:
-            assert stream.status == 401
-        assert instances[-1]._subscribers == {}
-
-
-async def test_command_without_token_is_401_and_handler_never_runs():
-    async with statewire_client(Guarded) as (app, client):
-        async with stream_of(app, headers={"x-token": "secret"}) as stream:
-            await stream.next_event()
-        before = len(runs)
-        response = await post_command(client, {"method": "inc", "params": []}, seq=1)
-        assert response.status_code == 401
-        assert len(runs) == before
-
-
-async def test_authorized_requests_succeed():
-    async with statewire_client(Guarded) as (app, client):
-        async with stream_of(app, headers={"x-token": "secret"}) as stream:
-            assert stream.status == 200
-            await stream.next_event()
-            response = await post_command(
-                client, {"method": "inc", "params": []}, seq=1, headers={"x-token": "secret"}
-            )
-            assert response.status_code == 200
-            assert await stream.next_event() == {
-                "ops": [{"op": "replace", "path": ["count"], "value": 1}],
-                "ack": 1,
-            }