statewire 0.4.2__tar.gz → 0.4.3__tar.gz

{statewire-0.4.2 → statewire-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: statewire
-Version: 0.4.2
+Version: 0.4.3
 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.4.2 → statewire-0.4.3}/pyproject.toml

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

{statewire-0.4.2 → statewire-0.4.3}/src/statewire/api.py

@@ -88,7 +88,7 @@ import secrets
 import sys
 import time
 from dataclasses import dataclass, field
-from typing import Any, AsyncIterator, Awaitable, Callable
+from typing import Any, AsyncIterator, Awaitable, Callable, Coroutine
 
 from fastapi import HTTPException, Request, WebSocket, WebSocketException
 from pinned import HostableContext, PinnedAPI
@@ -366,6 +366,33 @@ def _command_seq_of(request: Request) -> int:
     return seq
 
 
+def _ws_frame_headers(raw: str) -> dict[str, str] | None:
+    """The optional first client frame: ``{"type": "headers", "headers": {...}}``.
+    Returns the map, ``None`` for any other frame; raises on a malformed one."""
+    try:
+        frame = json.loads(raw)
+    except (json.JSONDecodeError, UnicodeDecodeError):
+        return None
+    if not isinstance(frame, dict) or frame.get("type") != "headers":
+        return None
+    headers = frame.get("headers")
+    if not isinstance(headers, dict) or not all(
+        isinstance(k, str) and isinstance(v, str) for k, v in headers.items()
+    ):
+        raise ValueError("headers frame must carry a string map")
+    return headers
+
+
+def _headers_connection(websocket: WebSocket, extra: dict[str, str]) -> HTTPConnection:
+    """The websocket's request view with the frame headers merged in (frame wins)."""
+    replaced = {name.lower().encode("latin-1") for name in extra}
+    merged = [(k, v) for k, v in websocket.scope["headers"] if k not in replaced]
+    merged += [
+        (k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in extra.items()
+    ]
+    return HTTPConnection({**websocket.scope, "headers": merged})
+
+
 def _ws_frame_command(raw: str) -> tuple[dict[str, Any], int] | None:
     try:
         frame = json.loads(raw)
@@ -394,6 +421,7 @@ def _reserved_paths() -> set[str]:
 class Statewire(PinnedAPI):
     stream_queue_size: int = 256
     heartbeat_interval: float = 15.0
+    ws_init_timeout: float = 5.0
     client_record_limit: int = 1024
     envelope_ts: bool = False
     strip_invalid_types: bool = False
@@ -438,6 +466,7 @@ class Statewire(PinnedAPI):
         self._flush_mode = False
         self._draining = False
         self._drain_armed = False
+        self._tasks: set["asyncio.Task[Any]"] = set()
 
     async def lifespan(self) -> AsyncIterator[None]:
         if self._state is _UNSET:
@@ -585,6 +614,37 @@ class Statewire(PinnedAPI):
         finally:
             self._draining = False
 
+    # ─── Tasks and idleness ─────────────────────────────────
+
+    def create_task(self, coro: Coroutine[Any, Any, Any]) -> "asyncio.Task[Any]":
+        task = super().create_task(coro)
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+        return task
+
+    def unref(self, task: "asyncio.Task[Any]") -> "asyncio.Task[Any]":
+        self._tasks.discard(task)
+        return super().unref(task)
+
+    async def idle(self) -> None:
+        """Wait until the instance is quiescent: every admitted command settled
+        and no tracked task (``create_task`` minus ``unref``) running."""
+        while True:
+            self.drain()
+            self.flush()
+            waiters = {
+                *self._tasks,
+                *(
+                    inv.task
+                    for client in self._clients.values()
+                    for inv in client.inflight.values()
+                    if inv.task is not None
+                ),
+            }
+            if not waiters:
+                return
+            await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED)
+
     # ─── Transient ops and finish ───────────────────────────
 
     def emit(self, value: dict[str, Any]) -> None:
@@ -1005,26 +1065,68 @@ class Statewire(PinnedAPI):
 
     @route.websocket("/ws")
     async def ws(self, websocket: WebSocket) -> None:
-        try:
-            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
         client_id = websocket.query_params.get("client", "")
         invalid = invalid_id_reason(client_id)
         if invalid is not None:
             raise WebSocketException(code=1008, reason=f"invalid client: {invalid}")
+        try:
+            identity = await self.authorize(websocket)
+            authorized = True
+        except HTTPException as exc:
+            identity = _UNSET
+            authorized = False
+            denial = exc.detail if isinstance(exc.detail, str) else "unauthorized"
         offered = websocket.scope.get("subprotocols") or []
         await websocket.accept(
             subprotocol=ops.WS_SUBPROTOCOL if ops.WS_SUBPROTOCOL in offered else None
         )
+        if not authorized:
+            identity = await self._ws_init_auth(websocket, denial)
+            if identity is _UNSET:
+                return
         q, snapshot = self._register(client_id, identity)
         lease = self._clients[client_id].lease
         try:
-            await self._ws_pump(websocket, q, client_id, snapshot, lease)
+            await self._ws_pump(
+                websocket, q, client_id, snapshot, lease, allow_headers=authorized
+            )
         finally:
             self._subscribers.pop(q, None)
 
+    async def _ws_init_auth(self, websocket: WebSocket, denial: str) -> Any:
+        """Unauthenticated attach: the headers frame must arrive within the init
+        window; the snapshot is held until it is authorized. Returns the identity
+        or ``_UNSET`` after closing the socket."""
+        try:
+            raw = await asyncio.wait_for(
+                websocket.receive_text(), self.ws_init_timeout
+            )
+        except TimeoutError:
+            await websocket.close(code=4401, reason=denial)
+            return _UNSET
+        except WebSocketDisconnect:
+            return _UNSET
+        except KeyError:
+            await websocket.close(code=1008)
+            return _UNSET
+        try:
+            headers = _ws_frame_headers(raw)
+            connection = (
+                None if headers is None else _headers_connection(websocket, headers)
+            )
+        except ValueError:
+            await websocket.close(code=1008)
+            return _UNSET
+        if connection is None:
+            await websocket.close(code=4401, reason=denial)
+            return _UNSET
+        try:
+            return await self.authorize(connection)
+        except HTTPException as exc:
+            reason = exc.detail if isinstance(exc.detail, str) else "unauthorized"
+            await websocket.close(code=4401, reason=reason)
+            return _UNSET
+
     async def _ws_pump(
         self,
         websocket: WebSocket,
@@ -1032,6 +1134,7 @@ class Statewire(PinnedAPI):
         client_id: str,
         snapshot: bytes,
         lease: str,
+        allow_headers: bool = False,
     ) -> None:
         await websocket.send_text(snapshot.decode())
         q_task: asyncio.Task[Any] = asyncio.create_task(q.get())
@@ -1062,6 +1165,20 @@ class Statewire(PinnedAPI):
                     except KeyError:
                         await websocket.close(code=1008)
                         return
+                    first = allow_headers
+                    allow_headers = False
+                    headers = None
+                    if first:
+                        try:
+                            headers = _ws_frame_headers(raw)
+                        except ValueError:
+                            await websocket.close(code=1008)
+                            return
+                    if headers is not None:
+                        if not await self._ws_restamp(websocket, client_id, headers):
+                            return
+                        io_task = asyncio.create_task(websocket.receive_text())
+                        continue
                     parsed = _ws_frame_command(raw)
                     if parsed is None:
                         await websocket.close(code=1008)
@@ -1096,6 +1213,25 @@ class Statewire(PinnedAPI):
             q_task.cancel()
             io_task.cancel()
 
+    async def _ws_restamp(
+        self, websocket: WebSocket, client_id: str, headers: dict[str, str]
+    ) -> bool:
+        try:
+            connection = _headers_connection(websocket, headers)
+        except ValueError:
+            await websocket.close(code=1008)
+            return False
+        try:
+            identity = await self.authorize(connection)
+        except HTTPException as exc:
+            reason = exc.detail if isinstance(exc.detail, str) else "unauthorized"
+            await websocket.close(code=4401, reason=reason)
+            return False
+        client = self._clients.get(client_id)
+        if client is not None:
+            client.identity = identity
+        return True
+
     async def _ws_execute(
         self, client_id: str, seq: int, command: dict[str, Any], lease: str
     ) -> dict[str, Any] | None:

{statewire-0.4.2 → statewire-0.4.3}/src/statewire/assistant_transport.py

@@ -42,22 +42,21 @@ live instance task (``create_task`` minus ``unref``) — then EOF. A run error
 recorded into state (``get_run_error``; default ``state["error"]``) becomes
 an error frame before EOF, its fields beyond ``message`` carried as the
 frame's structured payload where the format allows; a rejection or crash
-becomes an error frame followed by EOF. Client disconnect before EOF invokes
-``on_assistant_transport_disconnect`` (a plain overridable hook, not a
-command; default no-op); the run itself keeps executing server-side.
+becomes an error frame followed by EOF. Client disconnect before EOF
+detaches the stream; the run keeps executing server-side.
 
 ``POST /assistant-transport/api/resume`` (same body shape, ``commands`` must
 be absent or empty) reattaches to the run: a full-state ``set`` at the root,
 then the live tail until the instance is idle. If no run ever started, the
 response is 200 with an empty body and ``X-Stream-Status: not_found``. If the
 run already completed, the replay carries ``X-Stream-Status: completed``. A
-resume disconnect never invokes the cancel hook.
+resume disconnect detaches the same way.
 
-``POST /assistant-transport/api/cancel`` invokes the same disconnect hook,
-dispatches the host's registered ``run/stop`` handler (if any; a rejection —
-e.g. wrong-state while idle — is swallowed), and responds
-``{"success": true, "found": <bool>}``; ``found`` is true when a live run
-existed at request time.
+``POST /assistant-transport/api/cancel`` dispatches the host's registered
+``run/stop`` handler (if any; a rejection — e.g. wrong-state while idle — is
+swallowed). A non-empty JSON-object body forwards as the ``run/stop``
+params; an absent or empty body dispatches with none; any other body is a
+400. It responds ``{"success": true}``.
 
 ``POST /assistant-transport/api/status`` responds 200 always:
 ``{"isRunning": true, "status": "running"}`` while a run is active,
@@ -80,7 +79,7 @@ import contextlib
 import json
 import secrets
 import time
-from typing import Any, AsyncIterator, Coroutine
+from typing import Any, AsyncIterator
 
 import httpx
 from fastapi import FastAPI, HTTPException, Request
@@ -225,20 +224,6 @@ class AssistantTransport(Statewire):
 
     _at_run: _Run | None = None
 
-    def __init__(self, ctx: Any) -> None:
-        super().__init__(ctx)
-        self._at_tasks: set["asyncio.Task[Any]"] = set()
-
-    def create_task(self, coro: Coroutine[Any, Any, Any]) -> "asyncio.Task[Any]":
-        task = super().create_task(coro)
-        self._at_tasks.add(task)
-        task.add_done_callback(self._at_tasks.discard)
-        return task
-
-    def unref(self, task: "asyncio.Task[Any]") -> "asyncio.Task[Any]":
-        self._at_tasks.discard(task)
-        return super().unref(task)
-
     async def get_run_error(self) -> dict[str, Any] | None:
         """The run error to surface as the end-of-stream error frame.
 
@@ -253,9 +238,6 @@ class AssistantTransport(Statewire):
             return error
         return None
 
-    async def on_assistant_transport_disconnect(self) -> None:
-        """The legacy client went away or asked to stop; treat as a cancel signal."""
-
     async def get_input_request_id(self, tool_call_id: str) -> str | None:
         """Resolve a legacy toolCallId to the pending run/input requestId.
 
@@ -330,7 +312,7 @@ class AssistantTransport(Statewire):
                 run, q, client_id, lease, commands, parent_id
             )
         )
-        self._at_tasks.discard(watcher)
+        self._tasks.discard(watcher)
         return StreamingResponse(
             self._assistant_transport_attach(fmt, run, initial=True),
             media_type="text/event-stream",
@@ -371,13 +353,20 @@ class AssistantTransport(Statewire):
     @route.post("/assistant-transport/api/cancel")
     async def assistant_transport_cancel(self, request: Request) -> dict[str, Any]:
         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()
+        raw = await request.body()
+        params: list[Any] = []
+        if raw:
+            try:
+                body = json.loads(raw)
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                raise HTTPException(status_code=400, detail="body must be JSON") from None
+            if not isinstance(body, dict):
+                raise HTTPException(status_code=400, detail="body must be a JSON object")
+            if body:
+                params = [body]
         if "run/stop" in type(self)._statewire_commands:
-            await self._assistant_transport_dispatch("run/stop", identity)
-        return {"success": True, "found": found}
+            await self._assistant_transport_dispatch("run/stop", identity, params)
+        return {"success": True}
 
     @route.post("/assistant-transport/api/status")
     async def assistant_transport_status(self, request: Request) -> dict[str, Any]:
@@ -445,14 +434,16 @@ class AssistantTransport(Statewire):
             return member
         return {"method": kind, "params": [command]}
 
-    async def _assistant_transport_dispatch(self, method: str, identity: Any) -> None:
+    async def _assistant_transport_dispatch(
+        self, method: str, identity: Any, params: list[Any]
+    ) -> None:
         client_id = "at-" + secrets.token_urlsafe(9)
         q, _ = self._register(client_id, identity)
         try:
             lease = self._clients[client_id].lease
             assert lease is not None
             await self._admit(
-                client_id, 1, {"commands": [{"method": method, "params": []}]}, lease
+                client_id, 1, {"commands": [{"method": method, "params": params}]}, lease
             )
         finally:
             self._subscribers.pop(q, None)
@@ -480,7 +471,7 @@ class AssistantTransport(Statewire):
                     submit.done()
                     and acked >= total
                     and not open_pending
-                    and not self._at_tasks
+                    and not self._tasks
                 )
 
             def consume(item: Any) -> bool:
@@ -530,7 +521,7 @@ class AssistantTransport(Statewire):
                         if idle():
                             break
                         continue
-                    waiters = {q_get, *self._at_tasks}
+                    waiters = {q_get, *self._tasks}
                     if not submit.done():
                         waiters.add(submit)
                     done, _ = await asyncio.wait(
@@ -566,7 +557,6 @@ class AssistantTransport(Statewire):
         live = run if not run.done else None
         if live is not None:
             live.attachments.append(aq)
-        completed = False
         try:
             yield fmt.state([_legacy_op("set", [], run.replica)])
             if live is not None:
@@ -588,14 +578,10 @@ class AssistantTransport(Statewire):
                 yield fmt.error(run.error, run.error_payload)
             if fmt.end:
                 yield fmt.end
-            completed = True
         finally:
             if live is not None:
                 with contextlib.suppress(ValueError):
                     live.attachments.remove(aq)
-            if not completed and initial:
-                with contextlib.suppress(Exception):
-                    self.create_task(self.on_assistant_transport_disconnect())
 
 
 _HOP_HEADERS = {"host", "content-length", "transfer-encoding", "connection"}

{statewire-0.4.2 → statewire-0.4.3}/tests/test_assistant_transport.py

@@ -21,7 +21,6 @@ class Harness(AssistantTransport):
     def __init__(self, ctx):
         super().__init__(ctx)
         self.seen = []
-        self.cancelled = 0
 
     async def lifespan(self):
         self.state = {"messages": [], "text": "", "meta": {"title": None}}
@@ -55,9 +54,6 @@ class Harness(AssistantTransport):
         await asyncio.sleep(0.01)
         self.state["text"] += "b"
 
-    async def on_assistant_transport_disconnect(self):
-        self.cancelled += 1
-
 
 def data_stream_frames(text: str) -> list[tuple[str, object]]:
     frames = []
@@ -369,25 +365,6 @@ class _PostStream:
                 await self._task
 
 
-async def test_client_disconnect_invokes_the_disconnect_hook():
-    class SlowHarness(Harness):
-        @command("slow")
-        async def slow(self, cmd, *, ctx):
-            ctx.ack()
-            await asyncio.get_running_loop().create_future()
-
-    async with statewire_client(SlowHarness) as (app, client):
-        host = app.state.pinned_host
-        async with _PostStream(
-            app, "/threads/t1/assistant-transport/api/chat", {"commands": [{"type": "slow"}]}
-        ) as stream:
-            assert stream.status == 200
-            assert (await stream.next_chunk()).startswith(b"aui-state:")
-        await asyncio.sleep(0.05)
-        instance = (await host.directory.get("t1")).instance
-        assert instance.cancelled == 1
-
-
 class GatedHarness(Harness):
     def __init__(self, ctx):
         super().__init__(ctx)
@@ -495,7 +472,7 @@ async def test_mid_run_disconnect_then_resume_streams_snapshot_and_live_tail():
             assert replica["text"] == "ab"
 
 
-async def test_resume_disconnect_does_not_cancel_the_run():
+async def test_resume_disconnect_leaves_the_run_executing():
     async with statewire_client(GatedHarness) as (app, client):
         host = app.state.pinned_host
         async with _PostStream(
@@ -505,13 +482,11 @@ async def test_resume_disconnect_does_not_cancel_the_run():
             await stream.next_chunk()
         await asyncio.sleep(0.05)
         instance = (await host.directory.get("t1")).instance
-        cancelled_after_post = instance.cancelled
 
         async with post_resume(app) as resume:
             assert resume.status == 200
             await resume.next_chunk()
         await asyncio.sleep(0.05)
-        assert instance.cancelled == cancelled_after_post
 
         instance.release.set()
         await asyncio.sleep(0.05)
@@ -540,25 +515,22 @@ async def test_sse_mode_resume_is_symmetric():
         assert frames[0]["operations"][0]["path"] == []
 
 
-async def test_cancel_without_a_live_run_reports_not_found():
+async def test_cancel_without_a_live_run_succeeds():
     async with statewire_client(Harness) as (app, client):
         response = await client.post(
             "/threads/t1/assistant-transport/api/cancel", json={"threadId": "t1"}
         )
         assert response.status_code == 200
-        assert response.json() == {"success": True, "found": False}
-        instance = (await app.state.pinned_host.directory.get("t1")).instance
-        assert instance.cancelled == 0
+        assert response.json() == {"success": True}
 
         await post_run(client, {"commands": [], "state": {}})
         response = await client.post(
             "/threads/t1/assistant-transport/api/cancel", json={"threadId": "t1"}
         )
-        assert response.json() == {"success": True, "found": False}
-        assert instance.cancelled == 0
+        assert response.json() == {"success": True}
 
 
-async def test_cancel_with_a_live_run_cancels_and_reports_found():
+async def test_cancel_during_a_live_run_succeeds():
     async with statewire_client(GatedHarness) as (app, client):
         host = app.state.pinned_host
         async with _PostStream(
@@ -571,9 +543,8 @@ async def test_cancel_with_a_live_run_cancels_and_reports_found():
             response = await client.post(
                 "/threads/t1/assistant-transport/api/cancel", json={"threadId": "t1"}
             )
-            assert response.json() == {"success": True, "found": True}
+            assert response.json() == {"success": True}
             instance = (await host.directory.get("t1")).instance
-            assert instance.cancelled == 1
             instance.release.set()
 
 
@@ -615,7 +586,6 @@ class BackgroundHarness(AssistantTransport):
         super().__init__(ctx)
         self.release = asyncio.Event()
         self.fail_with = None
-        self.cancelled = 0
 
     async def lifespan(self):
         self.state = {"text": "", "error": None}
@@ -633,9 +603,6 @@ class BackgroundHarness(AssistantTransport):
         else:
             self.state["text"] += "b"
 
-    async def on_assistant_transport_disconnect(self):
-        self.cancelled += 1
-
 
 def post_background_chat(app):
     return _PostStream(
@@ -720,7 +687,7 @@ async def test_sse_mode_run_error_carries_the_structured_payload():
         }
 
 
-async def test_disconnect_before_idle_fires_the_hook_and_the_run_continues():
+async def test_disconnect_before_idle_leaves_the_run_executing():
     async with statewire_client(BackgroundHarness) as (app, client):
         async with post_background_chat(app) as stream:
             buffer = b""
@@ -728,7 +695,6 @@ async def test_disconnect_before_idle_fires_the_hook_and_the_run_continues():
                 buffer += await stream.next_chunk()
         await asyncio.sleep(0.05)
         instance = (await app.state.pinned_host.directory.get("t1")).instance
-        assert instance.cancelled == 1
         instance.release.set()
         await asyncio.sleep(0.05)
         response = await client.post(
@@ -749,8 +715,7 @@ class RunHarness(AssistantTransport):
         self.edited = []
         self.inputs = []
         self.reloaded = []
-        self.stopped = 0
-        self.cancelled = 0
+        self.stopped = []
 
     async def lifespan(self):
         self.state = {"messages": []}
@@ -775,10 +740,7 @@ class RunHarness(AssistantTransport):
 
     @command("run/stop")
     async def run_stop(self, params=None, *, ctx):
-        self.stopped += 1
-
-    async def on_assistant_transport_disconnect(self):
-        self.cancelled += 1
+        self.stopped.append(params)
 
 
 class AnchoredHarness(RunHarness):
@@ -938,10 +900,33 @@ async def test_cancel_dispatches_run_stop_when_registered():
             "/threads/t1/assistant-transport/api/cancel", json={"threadId": "t1"}
         )
         assert response.status_code == 200
-        assert response.json() == {"success": True, "found": False}
+        assert response.json() == {"success": True}
+        instance = (await app.state.pinned_host.directory.get("t1")).instance
+        assert instance.stopped == [{"threadId": "t1"}]
+
+
+async def test_cancel_forwards_the_body_as_run_stop_params():
+    async with statewire_client(RunHarness) as (app, client):
+        response = await client.post(
+            "/threads/t1/assistant-transport/api/cancel", json={"reason": "send_now"}
+        )
+        assert response.status_code == 200
+        instance = (await app.state.pinned_host.directory.get("t1")).instance
+        assert instance.stopped == [{"reason": "send_now"}]
+
+        response = await client.post("/threads/t1/assistant-transport/api/cancel")
+        assert response.status_code == 200
+        assert instance.stopped == [{"reason": "send_now"}, None]
+
+
+async def test_cancel_rejects_a_non_object_body():
+    async with statewire_client(RunHarness) as (app, client):
+        response = await client.post(
+            "/threads/t1/assistant-transport/api/cancel", json=[1]
+        )
+        assert response.status_code == 400
         instance = (await app.state.pinned_host.directory.get("t1")).instance
-        assert instance.stopped == 1
-        assert instance.cancelled == 0
+        assert instance.stopped == []
 
 
 async def test_cancel_swallows_a_run_stop_rejection():
@@ -955,7 +940,7 @@ async def test_cancel_swallows_a_run_stop_rejection():
             "/threads/t1/assistant-transport/api/cancel", json={"threadId": "t1"}
         )
         assert response.status_code == 200
-        assert response.json() == {"success": True, "found": False}
+        assert response.json() == {"success": True}
 
 
 async def test_add_tool_result_is_normalized_to_run_input():

{statewire-0.4.2 → statewire-0.4.3}/tests/test_authorize.py

@@ -201,6 +201,78 @@ async def test_identity_is_not_settable_via_context_updates():
             assert res["payload"] == {"identity": "alice"}
 
 
+async def test_ws_headers_frame_authorizes_and_releases_the_snapshot():
+    async with statewire_client(Guarded) as (app, client):
+        async with ws_of(app) as ws:
+            assert ws.accepted
+            await ws.send_frame({"type": "headers", "headers": {"x-token": "secret"}})
+            frame = await ws.next_frame()
+            assert "ops" in frame and "syn" in frame
+            assert (await ws.command({"method": "inc", "params": []}))["ack"] == 1
+
+
+async def test_ws_unauthorized_headers_frame_closes_4401():
+    async with statewire_client(Guarded) as (app, client):
+        async with ws_of(app) as ws:
+            assert ws.accepted
+            await ws.send_frame({"type": "headers", "headers": {"x-token": "wrong"}})
+            assert await ws.next_frame() is None
+            assert ws.close_code == 4401
+            assert ws.close_reason == "unauthorized"
+        assert instances[-1]._subscribers == {}
+
+
+async def test_ws_init_timeout_closes_unauthenticated_sockets_4401():
+    class Quick(Guarded):
+        ws_init_timeout = 0.05
+
+    async with statewire_client(Quick) as (app, client):
+        async with ws_of(app) as ws:
+            assert ws.accepted
+            assert await ws.next_frame() is None
+            assert ws.close_code == 4401
+
+
+async def test_ws_non_headers_first_frame_on_a_guarded_host_closes_4401():
+    async with statewire_client(Guarded) as (app, client):
+        async with ws_of(app) as ws:
+            assert ws.accepted
+            await ws.send_frame({"method": "inc", "params": [], "seq": 1})
+            assert await ws.next_frame() is None
+            assert ws.close_code == 4401
+
+
+async def test_ws_headers_frame_identity_reaches_ctx_caller():
+    async with statewire_client(Identified) as (app, client):
+        async with ws_of(app) as ws:
+            await ws.send_frame({"type": "headers", "headers": {"x-user": "wanda"}})
+            await ws.next_frame()
+            env = await ws.command({"method": "whoami", "params": []})
+            assert env["res"] == [
+                {"seq": 1, "type": "accepted", "payload": {"identity": "wanda"}}
+            ]
+
+
+async def test_ws_malformed_headers_frame_closes_1008():
+    async with statewire_client(Identified) as (app, client):
+        async with ws_of(app) as ws:
+            await ws.next_frame()
+            await ws.send_frame({"type": "headers", "headers": {"a": 1}})
+            assert await ws.next_frame() is None
+            assert ws.close_code == 1008
+
+
+async def test_ws_without_headers_frame_attaches_as_before_when_unguarded():
+    async with statewire_client(Identified) as (app, client):
+        async with ws_of(app) as ws:
+            assert ws.accepted
+            assert "syn" in (await ws.next_frame())
+            env = await ws.command({"method": "whoami", "params": []})
+            assert env["res"] == [
+                {"seq": 1, "type": "accepted", "payload": {"identity": None}}
+            ]
+
+
 class LegacyIdentified(AssistantTransport):
     seen = None
     stopped = None
@@ -217,7 +289,7 @@ class LegacyIdentified(AssistantTransport):
         self.seen = ctx.caller.identity
 
     @command("run/stop")
-    async def stop(self, *, ctx):
+    async def stop(self, params=None, *, ctx):
         self.stopped = ctx.caller.identity
 
 

statewire-0.4.3/tests/test_idle.py

@@ -0,0 +1,91 @@
+import asyncio
+
+from statewire_helpers import attach, post_command, statewire_client
+
+from statewire import Statewire, command
+
+
+def make_host():
+    class Host(Statewire):
+        def __init__(self, ctx):
+            super().__init__(ctx)
+            self.gates = [asyncio.Event(), asyncio.Event(), asyncio.Event()]
+
+        async def lifespan(self):
+            self.state = {"log": []}
+            yield
+
+        @command
+        async def hold(self, *, ctx):
+            ctx.ack()
+            await self.gates[0].wait()
+            self.state["log"].append("held")
+
+        @command
+        async def chain(self, *, ctx):
+            ctx.ack()
+            await self.gates[0].wait()
+            self.create_task(self._stage_a())
+
+        async def _stage_a(self):
+            await self.gates[1].wait()
+            self.create_task(self._stage_b())
+
+        async def _stage_b(self):
+            await self.gates[2].wait()
+            self.state["log"].append("done")
+
+    return Host
+
+
+async def _instance_of(app):
+    return (await app.state.pinned_host.directory.get("t1")).instance
+
+
+async def _settles(idle: "asyncio.Task[None]") -> bool:
+    await asyncio.sleep(0.05)
+    return idle.done()
+
+
+async def test_idle_returns_promptly_on_a_fresh_instance():
+    async with statewire_client(make_host()) as (app, client):
+        await attach(app)
+        instance = await _instance_of(app)
+        await asyncio.wait_for(instance.idle(), timeout=1)
+
+
+async def test_idle_blocks_on_a_pending_command_until_settle():
+    async with statewire_client(make_host()) as (app, client):
+        await attach(app)
+        assert (await post_command(client, {"method": "hold", "params": []}, seq=1)).status_code == 200
+        instance = await _instance_of(app)
+        idle = asyncio.create_task(instance.idle())
+        assert not await _settles(idle)
+        instance.gates[0].set()
+        await asyncio.wait_for(idle, timeout=1)
+        assert list(instance.state["log"]) == ["held"]
+
+
+async def test_idle_blocks_on_a_tracked_task_until_it_finishes():
+    async with statewire_client(make_host()) as (app, client):
+        await attach(app)
+        instance = await _instance_of(app)
+        instance.create_task(instance._stage_b())
+        idle = asyncio.create_task(instance.idle())
+        assert not await _settles(idle)
+        instance.gates[2].set()
+        await asyncio.wait_for(idle, timeout=1)
+        assert list(instance.state["log"]) == ["done"]
+
+
+async def test_one_idle_call_spans_a_command_task_task_chain():
+    async with statewire_client(make_host()) as (app, client):
+        await attach(app)
+        assert (await post_command(client, {"method": "chain", "params": []}, seq=1)).status_code == 200
+        instance = await _instance_of(app)
+        idle = asyncio.create_task(instance.idle())
+        for gate in instance.gates:
+            assert not await _settles(idle)
+            gate.set()
+        await asyncio.wait_for(idle, timeout=1)
+        assert list(instance.state["log"]) == ["done"]

{statewire-0.4.2 → statewire-0.4.3}/tests/test_ws.py

@@ -396,16 +396,19 @@ async def test_gone_finish_reaches_ws_then_close():
             assert await ws.next_frame() is None
 
 
-async def test_authorize_rejection_closes_1008_and_registers_nothing():
+async def test_authorize_rejection_closes_4401_and_registers_nothing():
     class Guarded(make_counter()):
+        ws_init_timeout = 0.05
+
         async def authorize(self, request: Request) -> None:
             if request.headers.get("x-token") != "secret":
                 raise HTTPException(status_code=401, detail="unauthorized")
 
     async with statewire_client(Guarded) as (app, client):
         async with ws_of(app) as ws:
-            assert ws.accepted is False
-            assert ws.close_code == 1008
+            assert ws.accepted
+            assert await ws.next_frame() is None
+            assert ws.close_code == 4401
             assert ws.close_reason == "unauthorized"
         host = app.state.pinned_host
         record = await host.directory.get("t1")