statewire 0.3.0__tar.gz → 0.3.2__tar.gz

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

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

{statewire-0.3.0 → statewire-0.3.2}/pyproject.toml

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

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

@@ -36,6 +36,16 @@ body over HTTP, as a ``res`` on the stream over WS.
 Handlers that declare a keyword-only ``ctx`` parameter receive a
 ``StatewireCommandContext``: ``ctx.caller`` is the submitting client's
 handle (holdable across a long run), ``ctx.ack()`` the invocation ack.
+``authorize`` gates every ingress and returns the caller's identity
+(``None`` = anonymous); the server stamps it on the client record per
+attach and per admitted command submission — never client-supplied,
+never on a rejected ingress — and drops it with the record. Handlers
+read it as ``ctx.caller.identity``: the current session owner, live —
+a newer authorized attach re-stamps it, so per-command trust decisions
+read it before the first await. A client id is a within-thread session
+id, not an authorization boundary: ``authorize`` must gate the thread
+itself, since any authorized principal knowing a client id can attach,
+receive the snapshot, and supersede that session.
 
 The host also owns a scheduler: ``self.schedule(fn)`` runs ``fn`` at the
 next drain (duplicate schedules coalesce), ``self.drain()`` drains now.
@@ -161,6 +171,7 @@ class _Client:
         self.results: dict[int, tuple[float, dict[str, Any]]] = {}
         self.inflight: dict[int, _Invocation] = {}
         self.context: dict[str, Any] | None = None
+        self.identity: Any = None
 
     def settle(self, seq: int, response: dict[str, Any] | None = None) -> None:
         full = {"seq": seq, **(response if response is not None else {"type": "accepted"})}
@@ -207,8 +218,9 @@ class _Client:
 class StatewireClientHandle:
     """Read-only view of the submitting client, reachable as ``ctx.caller``
     in a ``@command`` handler. Holdable across a long run: ``context`` reads
-    the client's current stored context live, ``is_connected`` reports
-    whether the client has an active stream attachment."""
+    the client's current stored context live, ``identity`` the current
+    session owner's identity, and ``is_connected`` reports whether the
+    client has an active stream attachment."""
 
     def __init__(self, host: "Statewire", client: _Client) -> None:
         self._host = host
@@ -222,6 +234,13 @@ class StatewireClientHandle:
     def context(self) -> dict[str, Any] | None:
         return self._client.context
 
+    @property
+    def identity(self) -> Any:
+        """Current session owner's identity — live, re-stamped by a newer
+        authorized attach; per-command trust decisions read it before the
+        handler's first await."""
+        return self._client.identity
+
     @property
     def is_connected(self) -> bool:
         client = self._client
@@ -391,11 +410,14 @@ class Statewire(PinnedAPI):
                         f"{cls.__name__}.{fn.__name__} declares @route on reserved "
                         f"path {spec.path!r}; it is a Statewire protocol endpoint"
                     )
-        merged = dict(cls._statewire_commands)
-        for name, fn in vars(cls).items():
-            method = getattr(fn, "_statewire_command_method", None)
-            if method is not None:
-                merged[method] = name
+        # Full-MRO scan so plain mixins (no Statewire base) contribute their
+        # @command handlers; the most derived registration of a method wins.
+        merged: dict[str, str] = {}
+        for base in reversed(cls.__mro__):
+            for name, fn in vars(base).items():
+                method = getattr(fn, "_statewire_command_method", None)
+                if method is not None:
+                    merged[method] = name
         cls._statewire_commands = merged
         lifespan = cls.__dict__.get("lifespan")
         if lifespan is not None and not getattr(lifespan, "_statewire_checked", False):
@@ -589,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(
@@ -650,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",
@@ -707,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)
@@ -716,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
@@ -785,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:
@@ -800,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)
@@ -811,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)
@@ -967,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
@@ -979,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.0 → statewire-0.3.2}/src/statewire/assistant_transport.py

@@ -12,7 +12,28 @@ context, readable via ``ctx.caller`` in handlers that declare a keyword-only
 ``ctx`` parameter. Unknown fields below the top level are ignored. All
 commands ``{"type": <name>, ...}`` are submitted as one batch through the
 normal Statewire admission path to the ``@command`` handlers registered under
-those names, each with the whole command object as the single param. The response streams
+those names, each with the whole command object as the single param — except
+the legacy commands, which are normalized to the canonical ``run/*`` dialect
+(a host-registered handler of the same legacy name always wins and receives
+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>}``.
+- ``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
+  unanswered tool-call request in ``state["inputRequests"]`` (None rejects),
+  and a ``modelContent`` field rejects (unsupported in ``run/input``).
+- A body with a ``parentId`` and no commands is the legacy reload: it becomes
+  ``run/reload`` with ``{"sourceId": <child>}`` via ``get_message_child_id``
+  (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
@@ -26,9 +47,11 @@ 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.
 
-``POST /assistant-transport/api/cancel`` invokes the same disconnect hook and
-responds ``{"success": true, "found": <bool>}``; ``found`` is true when a live
-run existed at request time.
+``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/status`` responds 200 always:
 ``{"isRunning": true, "status": "running"}`` while a run is active,
@@ -52,6 +75,7 @@ import json
 import secrets
 import time
 from typing import Any, AsyncIterator
+from uuid import uuid4
 
 import httpx
 from fastapi import FastAPI, HTTPException, Request
@@ -60,7 +84,7 @@ from starlette.background import BackgroundTask
 from starlette.responses import Response, StreamingResponse
 
 from .api import _LAGGARD, Statewire
-from .state import _frozen
+from .state import _frozen, plain
 
 _STREAM_HEADERS = {
     "Cache-Control": "no-cache, no-store, no-transform",
@@ -162,6 +186,28 @@ class _RunError(Exception):
     pass
 
 
+def _legacy_wire_message(message: Any) -> dict[str, Any]:
+    if not isinstance(message, dict):
+        raise _RunError("add-message: message must be an object")
+    parts = message.get("parts")
+    if not isinstance(parts, list):
+        raise _RunError("add-message: message.parts must be an array")
+    converted = []
+    for part in parts:
+        if isinstance(part, dict) and part.get("type") == "image":
+            image = part.get("image")
+            if not isinstance(image, str):
+                raise _RunError("add-message: image part requires a string image field")
+            converted.append({"type": "file", "mediaType": "image/*", "url": image})
+        else:
+            converted.append(part)
+    return {
+        "id": f"legacy_{uuid4().hex}",
+        "role": message.get("role"),
+        "parts": converted,
+    }
+
+
 _Item = tuple[str, Any]
 
 
@@ -193,9 +239,38 @@ class AssistantTransport(Statewire):
     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.
+
+        The default peeks the RunManager root-state mount:
+        ``state["inputRequests"]`` holds ``{..., "response": null}`` entries;
+        the unanswered tool-call request matching ``tool_call_id`` wins. None
+        rejects the legacy ``add-tool-result`` command."""
+        state = plain(self.state)
+        requests = state.get("inputRequests") if isinstance(state, dict) else None
+        if not isinstance(requests, list):
+            return None
+        for request in requests:
+            if (
+                isinstance(request, dict)
+                and request.get("type") == "tool-call"
+                and request.get("toolCallId") == tool_call_id
+                and request.get("response") is None
+            ):
+                return request.get("id")
+        return None
+
+    async def get_message_child_id(self, parent_id: str) -> str | None:
+        """Resolve a message id to its child message id.
+
+        Anchors the legacy reload (``run/reload``) and edit (``run/edit``)
+        translations. None (the default) rejects a legacy reload and sends
+        ``add-message`` as ``run/steer``."""
+        return None
+
     @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:
@@ -224,14 +299,19 @@ 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
+        parent_id = body.get("parentId")
+        if not isinstance(parent_id, str):
+            parent_id = None
         run = _Run(json.loads(snapshot)["ops"][0]["value"])
         self._at_run = run
         self.create_task(
-            self._assistant_transport_run(run, q, client_id, lease, commands)
+            self._assistant_transport_run(
+                run, q, client_id, lease, commands, parent_id
+            )
         )
         return StreamingResponse(
             self._assistant_transport_attach(fmt, run, initial=True),
@@ -272,11 +352,13 @@ class AssistantTransport(Statewire):
 
     @route.post("/assistant-transport/api/cancel")
     async def assistant_transport_cancel(self, request: Request) -> dict[str, Any]:
-        await self.authorize(request)
+        identity = await self.authorize(request)
         run = self._at_run
         found = run is not None and not run.done
         if found:
             await self.on_assistant_transport_disconnect()
+        if "run/stop" in type(self)._statewire_commands:
+            await self._assistant_transport_dispatch("run/stop", identity)
         return {"success": True, "found": found}
 
     @route.post("/assistant-transport/api/status")
@@ -299,16 +381,22 @@ class AssistantTransport(Statewire):
         }
 
     async def _assistant_transport_submit(
-        self, client_id: str, lease: str, commands: list[dict[str, Any]]
+        self,
+        client_id: str,
+        lease: str,
+        commands: list[dict[str, Any]],
+        parent_id: str | None,
     ) -> None:
         if not commands:
-            return
-        body = {
-            "commands": [
-                {"method": command["type"], "params": [command]}
+            if parent_id is None:
+                return
+            members = [await self._legacy_reload(parent_id)]
+        else:
+            members = [
+                await self._assistant_transport_member(command, parent_id)
                 for command in commands
             ]
-        }
+        body = {"commands": members}
         outcome, detail = await self._admit(client_id, 1, body, lease)
         if outcome == "rejected":
             raise _RunError(
@@ -320,6 +408,67 @@ class AssistantTransport(Statewire):
         if outcome != "executed":
             raise _RunError(f"command submission failed: {outcome}")
 
+    async def _assistant_transport_member(
+        self, command: dict[str, Any], parent_id: str | None
+    ) -> dict[str, Any]:
+        kind = command["type"]
+        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
+            )
+            if source_id is not None:
+                return {
+                    "method": "run/edit",
+                    "params": [{"sourceId": source_id, "message": message}],
+                }
+            return {"method": "run/steer", "params": [{"message": message}]}
+        if kind == "add-tool-result" and "add-tool-result" not in registered:
+            return {
+                "method": "run/input",
+                "params": [await self._legacy_input(command)],
+            }
+        return {"method": kind, "params": [command]}
+
+    async def _legacy_input(self, command: dict[str, Any]) -> dict[str, Any]:
+        if command.get("modelContent") is not None:
+            raise _RunError("add-tool-result: modelContent is not supported")
+        tool_call_id = command.get("toolCallId")
+        if not isinstance(tool_call_id, str):
+            raise _RunError("add-tool-result: toolCallId must be a string")
+        request_id = await self.get_input_request_id(tool_call_id)
+        if request_id is None:
+            raise _RunError(
+                f"add-tool-result: no input request for toolCallId {tool_call_id!r}"
+            )
+        response: dict[str, Any] = {"output": command.get("result")}
+        if command.get("isError"):
+            response["isError"] = True
+        if "artifact" in command:
+            response["artifact"] = command["artifact"]
+        return {"requestId": request_id, "response": response}
+
+    async def _legacy_reload(self, parent_id: str) -> dict[str, Any]:
+        source_id = await self.get_message_child_id(parent_id)
+        if source_id is None:
+            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, identity: 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
+            )
+        finally:
+            self._subscribers.pop(q, None)
+
     async def _assistant_transport_run(
         self,
         run: _Run,
@@ -327,13 +476,14 @@ class AssistantTransport(Statewire):
         client_id: str,
         lease: str,
         commands: list[dict[str, Any]],
+        parent_id: str | None,
     ) -> None:
         submit = asyncio.create_task(
-            self._assistant_transport_submit(client_id, lease, commands)
+            self._assistant_transport_submit(client_id, lease, commands, parent_id)
         )
         q_get: asyncio.Task[Any] = asyncio.create_task(q.get())
         try:
-            total = len(commands)
+            total = len(commands) or (1 if parent_id is not None else 0)
             acked = 0
             open_pending: set[int] = set()
             try: