statewire 0.3.1__tar.gz → 0.3.3__tar.gz

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "statewire"
-version = "0.3.1"
+version = "0.3.3"
 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.3}/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.3}/src/statewire/assistant_transport.py

@@ -20,9 +20,11 @@ the command whole):
 - ``add-message`` becomes ``run/steer`` with ``{"message": <wire message>}``:
   a fresh ``legacy_`` id is synthesized, image parts become
   ``{"type": "file", "mediaType": "image/*", "url"}`` parts, and command-level
-  anchors are dropped. When the body carries a top-level ``parentId`` and
-  ``get_message_child_id`` resolves it to a child, the command becomes
-  ``run/edit`` with ``{"sourceId": <child>, "message": <wire message>}``.
+  anchors are dropped. A command-level non-null ``sourceId`` makes it
+  ``run/edit`` with that ``sourceId`` directly; otherwise, when the body
+  carries a top-level ``parentId`` and ``get_message_child_id`` resolves it to
+  a child, the command becomes ``run/edit`` with ``{"sourceId": <child>,
+  "message": <wire message>}``.
 - ``add-tool-result`` becomes ``run/input`` with ``{"requestId": <resolved>,
   "response": {"output": <result>, "isError"?, "artifact"?}}``; the requestId
   comes from ``get_input_request_id(toolCallId)``, whose default peeks the
@@ -33,16 +35,19 @@ the command whole):
   (None rejects).
 
 A host without the target ``run/*`` handler rejects the batch via the
-ordinary unknown-command path. The response streams
-state changes as legacy frames until every submitted command settles, then
-EOF; 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.
+ordinary unknown-command path. The response streams state changes as legacy
+frames until the instance is idle — every submitted command settled and no
+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.
 
 ``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 run settles. If no run ever started, the
+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.
@@ -56,7 +61,7 @@ existed at request time.
 ``POST /assistant-transport/api/status`` responds 200 always:
 ``{"isRunning": true, "status": "running"}`` while a run is active,
 ``{"isRunning": false, "status": "completed", "completedAt": <ms epoch>}``
-after it settles, and ``{"isRunning": false, "status": "not_found",
+once idle, and ``{"isRunning": false, "status": "not_found",
 "message": ...}`` when no run ever started on this instance.
 
 Two stream formats, selected by ``assistant_transport_protocol``:
@@ -74,7 +79,7 @@ import contextlib
 import json
 import secrets
 import time
-from typing import Any, AsyncIterator
+from typing import Any, AsyncIterator, Coroutine
 from uuid import uuid4
 
 import httpx
@@ -102,7 +107,7 @@ class _DataStreamFormat:
         return f"aui-state:[{','.join(op_frames)}]\n"
 
     @staticmethod
-    def error(message: str) -> str:
+    def error(message: str, payload: Any = None) -> str:
         return f"3:{json.dumps(message)}\n"
 
 
@@ -119,9 +124,11 @@ class _AssistantTransportFormat:
         )
 
     @staticmethod
-    def error(message: str) -> str:
-        frame = json.dumps({"type": "error", "error": message})
-        return f"data: {frame}\n\n"
+    def error(message: str, payload: Any = None) -> str:
+        chunk: dict[str, Any] = {"type": "error", "error": message}
+        if payload is not None:
+            chunk["payload"] = payload
+        return f"data: {json.dumps(chunk)}\n\n"
 
 
 _FORMATS = {
@@ -183,7 +190,9 @@ _AT_CONSUMED_FIELDS = {"commands", "threadId", "state"}
 
 
 class _RunError(Exception):
-    pass
+    def __init__(self, message: str, *, payload: Any = None) -> None:
+        super().__init__(message)
+        self.payload = payload
 
 
 def _legacy_wire_message(message: Any) -> dict[str, Any]:
@@ -215,6 +224,7 @@ class _Run:
     def __init__(self, replica: Any) -> None:
         self.replica = replica
         self.error: str | None = None
+        self.error_payload: Any = None
         self.done = False
         self.completed_at: int | None = None
         self.attachments: list["asyncio.Queue[_Item]"] = []
@@ -227,7 +237,7 @@ class _Run:
                 self.attachments.remove(aq)
                 while not aq.empty():
                     aq.get_nowait()
-                aq.put_nowait(("error", "state stream overflowed"))
+                aq.put_nowait(("error", ("state stream overflowed", None)))
                 aq.put_nowait(("end", None))
 
 
@@ -236,6 +246,34 @@ 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.
+
+        The default peeks the RunManager root-state mount: a non-null
+        ``state["error"]`` object with a string ``message`` is the recorded
+        run error; its fields beyond ``message`` travel as the frame's
+        structured payload where the format carries one. None ends the
+        stream cleanly."""
+        state = plain(self.state)
+        error = state.get("error") if isinstance(state, dict) else None
+        if isinstance(error, dict) and isinstance(error.get("message"), str):
+            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."""
 
@@ -270,7 +308,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 +337,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
@@ -308,11 +346,12 @@ class AssistantTransport(Statewire):
             parent_id = None
         run = _Run(json.loads(snapshot)["ops"][0]["value"])
         self._at_run = run
-        self.create_task(
+        watcher = self.create_task(
             self._assistant_transport_run(
                 run, q, client_id, lease, commands, parent_id
             )
         )
+        self._at_tasks.discard(watcher)
         return StreamingResponse(
             self._assistant_transport_attach(fmt, run, initial=True),
             media_type="text/event-stream",
@@ -352,13 +391,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")
@@ -415,11 +454,11 @@ class AssistantTransport(Statewire):
         registered = type(self)._statewire_commands
         if kind == "add-message" and "add-message" not in registered:
             message = _legacy_wire_message(command.get("message"))
-            source_id = (
-                await self.get_message_child_id(parent_id)
-                if parent_id is not None
-                else None
-            )
+            source_id = command.get("sourceId")
+            if source_id is not None and not isinstance(source_id, str):
+                raise _RunError("add-message: sourceId must be a string or null")
+            if source_id is None and parent_id is not None:
+                source_id = await self.get_message_child_id(parent_id)
             if source_id is not None:
                 return {
                     "method": "run/edit",
@@ -457,9 +496,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
@@ -486,48 +525,82 @@ class AssistantTransport(Statewire):
             total = len(commands) or (1 if parent_id is not None else 0)
             acked = 0
             open_pending: set[int] = set()
+
+            def idle() -> bool:
+                return (
+                    submit.done()
+                    and acked >= total
+                    and not open_pending
+                    and not self._at_tasks
+                )
+
+            def consume(item: Any) -> bool:
+                nonlocal acked
+                if item is _LAGGARD:
+                    raise _RunError("state stream overflowed")
+                data, finish = item
+                envelope = json.loads(data)
+                op_frames: list[str] = []
+                for op in envelope.get("ops", []):
+                    run.replica, frame = _translate_op(run.replica, op)
+                    if frame is not None:
+                        op_frames.append(frame)
+                if op_frames:
+                    run.push(("state", op_frames))
+                for res in envelope.get("res", []):
+                    if res["type"] == "pending":
+                        open_pending.add(res["seq"])
+                        continue
+                    open_pending.discard(res["seq"])
+                    if res["type"] in ("rejected", "crashed"):
+                        raise _RunError(
+                            res.get("message") or res["type"],
+                            payload=res.get("payload"),
+                        )
+                if "ack" in envelope:
+                    acked = envelope["ack"]
+                if finish and not idle():
+                    fin = envelope.get("fin", {})
+                    raise _RunError(
+                        fin.get("message") or fin.get("reason", "stream ended")
+                    )
+                return finish
+
             try:
-                while not (submit.done() and acked >= total and not open_pending):
-                    waiters = {q_get} if submit.done() else {q_get, submit}
+                ended = False
+                while not ended:
+                    if idle():
+                        self.drain()
+                        self.flush()
+                        if q_get.done():
+                            ended = consume(q_get.result()) or ended
+                            q_get = asyncio.create_task(q.get())
+                            continue
+                        while not q.empty():
+                            ended = consume(q.get_nowait()) or ended
+                        if idle():
+                            break
+                        continue
+                    waiters = {q_get, *self._at_tasks}
+                    if not submit.done():
+                        waiters.add(submit)
                     done, _ = await asyncio.wait(
                         waiters, return_when=asyncio.FIRST_COMPLETED
                     )
                     if submit in done and submit.exception() is not None:
                         raise _RunError(str(submit.exception()))
-                    if q_get not in done:
-                        continue
-                    item = q_get.result()
-                    if item is _LAGGARD:
-                        raise _RunError("state stream overflowed")
-                    data, finish = item
-                    envelope = json.loads(data)
-                    op_frames: list[str] = []
-                    for op in envelope.get("ops", []):
-                        run.replica, frame = _translate_op(run.replica, op)
-                        if frame is not None:
-                            op_frames.append(frame)
-                    if op_frames:
-                        run.push(("state", op_frames))
-                    for res in envelope.get("res", []):
-                        if res["type"] == "pending":
-                            open_pending.add(res["seq"])
-                            continue
-                        open_pending.discard(res["seq"])
-                        if res["type"] in ("rejected", "crashed"):
-                            raise _RunError(res.get("message") or res["type"])
-                    if "ack" in envelope:
-                        acked = envelope["ack"]
-                    if finish:
-                        if submit.done() and acked >= total and not open_pending:
-                            break
-                        fin = envelope.get("fin", {})
-                        raise _RunError(
-                            fin.get("message") or fin.get("reason", "stream ended")
-                        )
-                    q_get = asyncio.create_task(q.get())
+                    if q_get in done:
+                        ended = consume(q_get.result())
+                        if not ended:
+                            q_get = asyncio.create_task(q.get())
+                error = await self.get_run_error()
+                if error is not None:
+                    payload = {k: v for k, v in error.items() if k != "message"}
+                    raise _RunError(error["message"], payload=payload or None)
             except _RunError as exc:
                 run.error = str(exc)
-                run.push(("error", run.error))
+                run.error_payload = exc.payload
+                run.push(("error", (run.error, run.error_payload)))
         finally:
             q_get.cancel()
             submit.cancel()
@@ -559,11 +632,11 @@ class AssistantTransport(Statewire):
                     if kind == "state":
                         yield fmt.state(value)
                     elif kind == "error":
-                        yield fmt.error(value)
+                        yield fmt.error(*value)
                     else:
                         break
             elif initial and run.error is not None:
-                yield fmt.error(run.error)
+                yield fmt.error(run.error, run.error_payload)
             if fmt.end:
                 yield fmt.end
             completed = True

{statewire-0.3.1 → statewire-0.3.3}/tests/test_assistant_transport.py

@@ -2,13 +2,16 @@
 
 The legacy client POSTs ``{"commands", "state", "threadId", ...}`` and reads
 back a data-stream (``<prefix>:<json>`` lines) or assistant-transport SSE
-response carrying ``set``/``append-text`` state ops until the run settles.
+response carrying ``set``/``append-text`` state ops until the instance is
+idle.
 """
 
 import asyncio
 import contextlib
 import json
 
+import pytest
+
 from statewire_helpers import statewire_client
 
 from statewire import AssistantTransport, StatewireReject, command
@@ -110,7 +113,7 @@ async def post_run(client, body, thread="t1"):
     return await client.post(f"/threads/{thread}/assistant-transport/api/chat", json=body)
 
 
-async def test_data_stream_response_streams_snapshot_then_ops_until_settle():
+async def test_data_stream_response_streams_snapshot_then_ops_until_idle():
     async with statewire_client(Harness) as (app, client):
         response = await post_run(
             client,
@@ -347,6 +350,14 @@ class _PostStream:
     async def next_chunk(self, timeout: float = 5) -> bytes:
         return await asyncio.wait_for(self._chunks.get(), timeout)
 
+    async def until_closed(self, timeout: float = 5) -> bytes:
+        assert self._task is not None
+        await asyncio.wait_for(asyncio.shield(self._task), timeout)
+        tail = b""
+        while not self._chunks.empty():
+            tail += self._chunks.get_nowait()
+        return tail
+
     async def __aexit__(self, *exc):
         self._disconnected.set()
         assert self._task is not None
@@ -597,6 +608,138 @@ async def test_status_reports_not_found_running_and_completed():
         assert isinstance(body["completedAt"], int)
 
 
+class BackgroundHarness(AssistantTransport):
+    """Host whose run/steer settles at enqueue; the run is an instance task."""
+
+    def __init__(self, ctx):
+        super().__init__(ctx)
+        self.release = asyncio.Event()
+        self.fail_with = None
+        self.cancelled = 0
+
+    async def lifespan(self):
+        self.state = {"text": "", "error": None}
+        yield
+
+    @command("run/steer")
+    async def run_steer(self, params):
+        self.create_task(self._run())
+
+    async def _run(self):
+        self.state["text"] += "a"
+        await self.release.wait()
+        if self.fail_with is not None:
+            self.state["error"] = self.fail_with
+        else:
+            self.state["text"] += "b"
+
+    async def on_assistant_transport_disconnect(self):
+        self.cancelled += 1
+
+
+def post_background_chat(app):
+    return _PostStream(
+        app,
+        "/threads/t1/assistant-transport/api/chat",
+        {"commands": [{"type": "run/steer"}], "state": {}, "threadId": "t1"},
+    )
+
+
+async def test_chat_stays_open_past_command_settle_and_eofs_at_idle():
+    async with statewire_client(BackgroundHarness) as (app, client):
+        async with post_background_chat(app) as stream:
+            assert stream.status == 200
+            buffer = b""
+            while b'"a"' not in buffer:
+                buffer += await stream.next_chunk()
+            with pytest.raises(TimeoutError):
+                await stream.next_chunk(0.1)
+            instance = (await app.state.pinned_host.directory.get("t1")).instance
+            instance.release.set()
+            buffer += await stream.until_closed()
+        frames = data_stream_frames(buffer.decode())
+        assert all(prefix == "aui-state" for prefix, _ in frames)
+        replica = apply_legacy_ops({}, state_ops(frames))
+        assert replica["text"] == "ab"
+        body = (
+            await client.post(
+                "/threads/t1/assistant-transport/api/status", json={"threadId": "t1"}
+            )
+        ).json()
+        assert body["status"] == "completed"
+
+
+async def test_run_error_recorded_into_state_is_an_error_frame_before_eof():
+    async with statewire_client(BackgroundHarness) as (app, client):
+        async with post_background_chat(app) as stream:
+            buffer = b""
+            while b'"a"' not in buffer:
+                buffer += await stream.next_chunk()
+            instance = (await app.state.pinned_host.directory.get("t1")).instance
+            instance.fail_with = {
+                "message": "payment required",
+                "reason": "payment-required",
+            }
+            instance.release.set()
+            buffer += await stream.until_closed()
+        frames = data_stream_frames(buffer.decode())
+        assert frames[-1] == ("3", "payment required")
+        replica = apply_legacy_ops({}, state_ops(frames))
+        assert replica["error"] == {
+            "message": "payment required",
+            "reason": "payment-required",
+        }
+
+
+async def test_sse_mode_run_error_carries_the_structured_payload():
+    class SSEBackgroundHarness(BackgroundHarness):
+        assistant_transport_protocol = "assistant-transport"
+
+    async with statewire_client(SSEBackgroundHarness) as (app, client):
+        async with _PostStream(
+            app,
+            "/threads/t1/assistant-transport/api/chat",
+            {"commands": [{"type": "run/steer"}], "state": {}, "threadId": "t1"},
+        ) as stream:
+            buffer = b""
+            while b'"a"' not in buffer:
+                buffer += await stream.next_chunk()
+            instance = (await app.state.pinned_host.directory.get("t1")).instance
+            instance.fail_with = {
+                "message": "payment required",
+                "reason": "payment-required",
+            }
+            instance.release.set()
+            buffer += await stream.until_closed()
+        frames = sse_frames(buffer.decode())
+        assert frames[-1] == "[DONE]"
+        assert frames[-2] == {
+            "type": "error",
+            "error": "payment required",
+            "payload": {"reason": "payment-required"},
+        }
+
+
+async def test_disconnect_before_idle_fires_the_hook_and_the_run_continues():
+    async with statewire_client(BackgroundHarness) as (app, client):
+        async with post_background_chat(app) as stream:
+            buffer = b""
+            while b'"a"' not in buffer:
+                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(
+            "/threads/t1/assistant-transport/api/resume",
+            json={"commands": [], "state": {}, "threadId": "t1"},
+        )
+        assert response.headers["x-stream-status"] == "completed"
+        ops = state_ops(data_stream_frames(response.text))
+        assert ops[0]["value"]["text"] == "ab"
+
+
 class RunHarness(AssistantTransport):
     """Host speaking the canonical run/* dialect; no legacy handlers."""
 
@@ -992,6 +1135,58 @@ async def test_add_message_with_a_resolvable_parent_is_normalized_to_run_edit():
         assert params["message"]["parts"] == [{"type": "text", "text": "again"}]
 
 
+async def test_add_message_with_a_source_id_is_normalized_to_run_edit_with_it():
+    async with statewire_client(AnchoredHarness) as (app, client):
+        response = await post_run(
+            client,
+            {
+                "commands": [
+                    {
+                        "type": "add-message",
+                        "message": {
+                            "role": "user",
+                            "parts": [{"type": "text", "text": "again"}],
+                        },
+                        "sourceId": "m9",
+                    }
+                ],
+                "state": {},
+                "parentId": "m1",
+            },
+        )
+        assert response.status_code == 200
+        instance = (await app.state.pinned_host.directory.get("t1")).instance
+        assert instance.steered == []
+        [params] = instance.edited
+        assert params["sourceId"] == "m9"
+        assert params["message"]["parts"] == [{"type": "text", "text": "again"}]
+
+
+async def test_add_message_with_a_null_source_id_and_no_parent_is_run_steer():
+    async with statewire_client(AnchoredHarness) as (app, client):
+        response = await post_run(
+            client,
+            {
+                "commands": [
+                    {
+                        "type": "add-message",
+                        "message": {
+                            "role": "user",
+                            "parts": [{"type": "text", "text": "hi"}],
+                        },
+                        "sourceId": None,
+                    }
+                ],
+                "state": {},
+            },
+        )
+        assert response.status_code == 200
+        instance = (await app.state.pinned_host.directory.get("t1")).instance
+        assert instance.edited == []
+        [(params, _)] = instance.steered
+        assert set(params) == {"message"}
+
+
 async def test_add_message_with_an_unresolvable_parent_falls_back_to_run_steer():
     async with statewire_client(AnchoredHarness) as (app, client):
         response = await post_run(

statewire-0.3.3/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.3}/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.3}/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,
-            }