statewire 0.3.2__tar.gz → 0.3.3__tar.gz

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "statewire"
-version = "0.3.2"
+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.2 → 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."""
 
@@ -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",
@@ -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",
@@ -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.2 → 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(