statewire 0.8.2__tar.gz → 0.8.4__tar.gz

{statewire-0.8.2 → statewire-0.8.4}/PKG-INFO

@@ -1,15 +1,17 @@
 Metadata-Version: 2.5
 Name: statewire
-Version: 0.8.2
+Version: 0.8.4
 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
 Requires-Python: <4.0,>=3.12
 Requires-Dist: fastapi>=0.115
 Requires-Dist: httpx>=0.27
-Requires-Dist: pinned<0.5,>=0.4.0
+Requires-Dist: pinned<0.8,>=0.7.0
 Provides-Extra: langgraph
 Requires-Dist: langchain-core>=0.3; extra == 'langgraph'
+Provides-Extra: orjson
+Requires-Dist: orjson>=3.9; extra == 'orjson'
 Description-Content-Type: text/markdown
 
 # statewire

{statewire-0.8.2 → statewire-0.8.4}/pyproject.toml

@@ -1,18 +1,19 @@
 [project]
 name = "statewire"
-version = "0.8.2"
+version = "0.8.4"
 description = "Replicate one JSON object over an SSE op stream + a command endpoint"
 readme = "README.md"
 license = "MIT"
 requires-python = ">=3.12,<4.0"
 dependencies = [
-    "pinned>=0.4.0,<0.5",
+    "pinned>=0.7.0,<0.8",
     "fastapi>=0.115",
     "httpx>=0.27",
 ]
 
 [project.optional-dependencies]
 langgraph = ["langchain-core>=0.3"]
+orjson = ["orjson>=3.9"]
 
 [project.urls]
 Repository = "https://github.com/assistant-ui/harness-sdk"

{statewire-0.8.2 → statewire-0.8.4}/src/statewire/api.py

@@ -130,7 +130,15 @@ from starlette.responses import JSONResponse, Response, StreamingResponse
 from starlette.websockets import WebSocketDisconnect
 
 from . import metrics, ops
-from .state import StateProxy, _decode_meta, _encoded, _frozen, plain
+from .state import (
+    StateProxy,
+    _date_meta,
+    _decode_meta,
+    _frozen,
+    _frozen_encoded,
+    _track_dates,
+    plain,
+)
 
 if sys.version_info < (3, 12):
     raise RuntimeError("statewire requires Python 3.12+ (eager task execution)")
@@ -258,7 +266,7 @@ def _parse_frame(frame: Any, *, first: bool) -> _Frame:
 
 
 def _payload_fields(value: Any) -> dict[str, Any]:
-    payload, meta = _encoded(_frozen(value))
+    payload, meta = _frozen_encoded(value)
     return {"payload": payload, **({"meta": meta} if meta else {})}
 
 
@@ -266,10 +274,11 @@ def _ws_auth_rejected(exc: AuthRejected) -> AuthRejected:
     fin: dict[str, Any] = {"reason": "unauthorized", "message": exc.message}
     if exc.payload is not None:
         fin.update(_payload_fields(exc.payload))
+    # Pinned serializes this payload itself: hand it wire values, not datetimes.
     return AuthRejected(
         exc.message,
         status_code=exc.status_code,
-        payload={"fin": fin},
+        payload={"fin": _load_json(ops.encode_envelope(fin))},
     )
 
 
@@ -670,6 +679,9 @@ class Statewire(PinnedAPI):
     def __init__(self, ctx: HostableContext) -> None:
         super().__init__(ctx)
         self._state: Any = _UNSET
+        self._state_version = 0
+        self._date_paths: set[tuple[str | int, ...]] = set()
+        self._snapshot_cache: tuple[int, bytes] | None = None
         self._pending_ops: list[ops.Op] = []
         self._dirty_acks: set[str] = set()
         self._pending_res: dict[str, list[dict[str, Any]]] = {}
@@ -747,14 +759,15 @@ class Statewire(PinnedAPI):
 
     @state.setter
     def state(self, value: Any) -> None:
-        value = _frozen(plain(value), lenient=self.strip_invalid_types)
-        if not isinstance(value, dict):
+        stored, meta = _frozen_encoded(plain(value), lenient=self.strip_invalid_types)
+        if not isinstance(stored, dict):
             raise TypeError("state must be a dict")
-        self._state = value
-        wire, meta = _encoded(value)
-        self._statewire_record(ops.replace_op([], wire, meta))
+        self._state = stored
+        self._statewire_record(ops.replace_op([], stored, meta))
 
     def _statewire_record(self, op: ops.Op) -> None:
+        self._state_version += 1
+        _track_dates(self._date_paths, self._state, op)
         self._pending_ops.append(op)
         self._schedule_flush()
 
@@ -837,8 +850,6 @@ class Statewire(PinnedAPI):
                     for seq in (cmd_applied.get(key, []) if key is not None else [])
                 ]
                 envelope: dict[str, Any] = {}
-                if entries:
-                    envelope["ops"] = entries
                 if key is not None and key in res:
                     envelope["res"] = res[key]  # legacy: envelope res dual-emit
                 if key is not None and key in acks:
@@ -854,10 +865,12 @@ class Statewire(PinnedAPI):
                         envelope["cmd"] = applied_entries
                     if idle_edge:
                         envelope["idle"] = level
-                elif applied_entries:
-                    msgs.append(ops.encode_envelope({"cmd": applied_entries}))
-                if envelope:
-                    msgs.append(ops.encode_envelope(envelope))
+                    msgs.append(ops.ops_envelope(entries, envelope))
+                else:
+                    if applied_entries:
+                        msgs.append(ops.encode_envelope({"cmd": applied_entries}))
+                    if envelope:
+                        msgs.append(ops.encode_envelope(envelope))
                 verdicts = cmd_verdicts.get(key, []) if key is not None else []
                 if verdicts:
                     msgs.append(ops.encode_envelope({"cmd": verdicts}))
@@ -1019,7 +1032,6 @@ class Statewire(PinnedAPI):
         q: asyncio.Queue[Any] = asyncio.Queue(maxsize=self.stream_queue_size)
         self._subscribers[q] = client_id
         self._gc_clients()
-        wire_state, state_meta = _encoded(self._state)
         syn: dict[str, Any] = {
             "seq": client.admitted_through(),
             "ctxSeq": client.ctx_seq,
@@ -1029,8 +1041,17 @@ class Statewire(PinnedAPI):
         }
         if self._idle_wire:
             syn["idle"] = True
-        snapshot = ops.encode_envelope(
-            {"ops": ops.snapshot_event(wire_state, state_meta), "syn": syn}
+        if (
+            self._snapshot_cache is None
+            or self._snapshot_cache[0] != self._state_version
+        ):
+            op = ops.replace_op(
+                [], self._state, _date_meta(self._state, self._date_paths)
+            )
+            self._snapshot_cache = (self._state_version, ops.ops_envelope([op])[:-1])
+        # splice the cached b'{"ops":[...]' fragment with the per-attach syn
+        snapshot = (
+            self._snapshot_cache[1] + b"," + ops.encode_envelope({"syn": syn})[1:]
         )
         return q, snapshot
 

{statewire-0.8.2 → statewire-0.8.4}/src/statewire/client.py

@@ -25,7 +25,7 @@ from urllib.parse import quote, urlencode
 
 import httpx
 
-from .state import _iso
+from .ops import _iso
 
 logger = logging.getLogger("statewire.client")
 
@@ -50,14 +50,20 @@ def _describe(failure: dict[str, Any]) -> str:
     match fate:
         case "rejected":
             code = failure.get("code")
-            detail = (f" ({code})" if code else "") + (f": {message}" if message else "")
+            detail = (f" ({code})" if code else "") + (
+                f": {message}" if message else ""
+            )
             return f"statewire: command rejected{detail}"
         case "aborted":
             return f"statewire: command aborted ({cause})"
         case "lost":
             return f"statewire: command lost ({cause})"
         case "failed":
-            return f"statewire: command failed: {message}" if message else "statewire: command failed"
+            return (
+                f"statewire: command failed: {message}"
+                if message
+                else "statewire: command failed"
+            )
         case "crashed":
             return "statewire: command crashed"
         case "expired":
@@ -103,7 +109,9 @@ def _revive_value(value: Any, meta: Any) -> Any:
             parent, key = node, seg
             if isinstance(seg, str) and isinstance(node, dict) and seg in node:
                 node = node[seg]
-            elif isinstance(seg, int) and isinstance(node, list) and 0 <= seg < len(node):
+            elif (
+                isinstance(seg, int) and isinstance(node, list) and 0 <= seg < len(node)
+            ):
                 node = node[seg]
             else:
                 raise StatewireProtocolError(f"meta path {path} does not resolve")
@@ -326,7 +334,9 @@ class SseTransport:
         except httpx.HTTPError as exc:
             raise TransportError(str(exc) or "network") from exc
 
-    async def _envelopes(self, response: httpx.Response) -> AsyncIterator[dict[str, Any]]:
+    async def _envelopes(
+        self, response: httpx.Response
+    ) -> AsyncIterator[dict[str, Any]]:
         buffer = ""
         try:
             async for chunk in response.aiter_text():
@@ -444,7 +454,9 @@ class WsTransport:
     async def _frames(self, sock: Any) -> AsyncIterator[dict[str, Any]]:
         while True:
             try:
-                raw = await asyncio.wait_for(sock.recv(), self._timing.heartbeat_timeout)
+                raw = await asyncio.wait_for(
+                    sock.recv(), self._timing.heartbeat_timeout
+                )
             except TimeoutError as exc:
                 raise TransportError("timeout") from exc
             except asyncio.CancelledError:
@@ -1105,11 +1117,15 @@ class StatewireClient:
             if outcome == "idle":
                 self._attempt = 0
                 self._set_connection({"status": "idle"})
-                while self._running and not self._has_work() and not self._reconnect_now:
+                while (
+                    self._running and not self._has_work() and not self._reconnect_now
+                ):
                     await self._wait_wake()
                 continue
             self._attempt += 1
-            delay = 0.0 if self._reconnect_now else _backoff(self._timing, self._attempt)
+            delay = (
+                0.0 if self._reconnect_now else _backoff(self._timing, self._attempt)
+            )
             if self._session_known or self._reconnect_now:
                 reconnecting: dict[str, Any] = {
                     "status": "reconnecting",

statewire-0.8.4/src/statewire/langgraph.py

@@ -0,0 +1,343 @@
+import weakref
+from typing import Any, NamedTuple
+
+from langchain_core.messages.ai import AIMessageChunk, add_ai_message_chunks
+
+from .state import plain
+
+
+def _message_matches(existing_message: Any, message_dict: dict[str, Any]) -> bool:
+    existing_message = plain(existing_message)
+    if not isinstance(existing_message, dict):
+        return False
+
+    message_id = message_dict.get("id")
+    if message_id is not None and existing_message.get("id") == message_id:
+        return True
+
+    tool_call_id = message_dict.get("tool_call_id")
+    return (
+        tool_call_id is not None
+        and existing_message.get("tool_call_id") == tool_call_id
+    )
+
+
+def _find_existing_message_index(
+    messages: Any, message_dict: dict[str, Any]
+) -> int | None:
+    for i, existing_message in enumerate(messages):
+        if _message_matches(existing_message, message_dict):
+            return i
+    return None
+
+
+def _is_provisional_id(message_id: Any) -> bool:
+    return message_id is None or (
+        isinstance(message_id, str) and message_id.startswith("lc_")
+    )
+
+
+def _open_chunk_tail_index(messages: Any, message_dict: dict[str, Any]) -> int | None:
+    """Providers churn chunk ids mid-generation (lc_run-* vs provider ids);
+    an open chunk accumulation at the tail is the same generation."""
+    if not len(messages):
+        return None
+    index = len(messages) - 1
+    tail = plain(messages[index])
+    if not isinstance(tail, dict) or tail.get("type") != "ai":
+        return None
+    if "chunk_position" not in tail or tail["chunk_position"] == "last":
+        return None
+    if _is_provisional_id(tail.get("id")) or _is_provisional_id(message_dict.get("id")):
+        return index
+    return None
+
+
+def _can_patch_value(current_value: Any, next_value: Any) -> bool:
+    # Keep these guards in sync with _patch_child so the later mutation cannot
+    # partially apply before discovering an unsupported deletion/truncation.
+    current_value = plain(current_value)
+
+    if current_value == next_value:
+        return True
+
+    if isinstance(current_value, dict) and isinstance(next_value, dict):
+        current_keys = set(current_value.keys())
+        next_keys = set(next_value.keys())
+        return current_keys.issubset(next_keys) and all(
+            _can_patch_value(current_value[key], next_value[key])
+            for key in current_keys
+        )
+
+    if isinstance(current_value, list) and isinstance(next_value, list):
+        return len(current_value) <= len(next_value) and all(
+            _can_patch_value(item, next_value[index])
+            for index, item in enumerate(current_value)
+        )
+
+    return True
+
+
+def _patch_child(parent: Any, key: str | int, next_value: Any) -> None:
+    """Patch through StateProxy containers so writes emit granular ops."""
+    current_value = plain(parent[key])
+
+    if current_value == next_value:
+        return
+
+    if isinstance(current_value, dict) and isinstance(next_value, dict):
+        current_keys = set(current_value.keys())
+        next_keys = set(next_value.keys())
+        if not current_keys.issubset(next_keys):
+            raise ValueError("Cannot represent deleted dictionary keys as ops")
+
+        target = parent[key]
+        for child_key in next_keys:
+            if child_key not in current_value:
+                target[child_key] = next_value[child_key]
+            else:
+                _patch_child(target, child_key, next_value[child_key])
+        return
+
+    if isinstance(current_value, list) and isinstance(next_value, list):
+        if len(next_value) < len(current_value):
+            raise ValueError("Cannot represent list truncation as ops")
+
+        target = parent[key]
+        for index, item in enumerate(next_value):
+            if index >= len(current_value):
+                target.append(item)
+            else:
+                _patch_child(target, index, item)
+        return
+
+    parent[key] = next_value
+
+
+_CHUNK_FIELDS = frozenset(AIMessageChunk(content="").model_dump())
+
+
+class _OpenChunk(NamedTuple):
+    index: int
+    chunk: AIMessageChunk
+    stored: dict[str, Any]
+    churn_id: str | None
+
+
+_open_chunks: weakref.WeakKeyDictionary[Any, _OpenChunk] = weakref.WeakKeyDictionary()
+
+
+def _fast_chunk_merge(
+    state: Any, entry: _OpenChunk, message: Any, message_dict: dict[str, Any]
+) -> AIMessageChunk | None:
+    """Merge a streamed chunk into the retained open accumulation in O(chunk);
+    None on any mismatch, which falls back to the full scan-and-patch path."""
+    messages = plain(state["messages"])
+    if entry.index >= len(messages) or messages[entry.index] is not entry.stored:
+        return None
+    if message_dict.get("chunk_position") == "last":
+        return None
+    if not _CHUNK_FIELDS.issuperset(message_dict):
+        return None
+    incoming_id = message_dict.get("id")
+    if incoming_id != entry.stored.get("id"):
+        # Only an id the transcript scan is known to miss may churn into the tail.
+        if incoming_id is not None and incoming_id != entry.churn_id:
+            return None
+        if entry.index != len(messages) - 1:
+            return None
+        if not (
+            _is_provisional_id(entry.stored.get("id"))
+            or _is_provisional_id(incoming_id)
+        ):
+            return None
+
+    merged = add_ai_message_chunks(entry.chunk, message)
+
+    changed: dict[str, Any] = {}
+    if message_dict["content"] != "":
+        changed["content"] = merged.content
+    if message_dict["additional_kwargs"]:
+        changed["additional_kwargs"] = merged.additional_kwargs
+    if message_dict["response_metadata"]:
+        changed["response_metadata"] = merged.response_metadata
+    if message_dict["tool_call_chunks"]:
+        changed["tool_call_chunks"] = merged.tool_call_chunks
+        changed["tool_calls"] = merged.tool_calls
+        changed["invalid_tool_calls"] = merged.invalid_tool_calls
+    if message_dict["usage_metadata"] is not None:
+        changed["usage_metadata"] = merged.usage_metadata
+    for key in ("id", "name", "chunk_position"):
+        value = getattr(merged, key)
+        if entry.stored.get(key) != value:
+            changed[key] = value
+
+    if not all(
+        _can_patch_value(entry.stored[key], value) for key, value in changed.items()
+    ):
+        return None
+
+    target = state["messages"][entry.index]
+    for key, value in changed.items():
+        _patch_child(target, key, value)
+    return merged
+
+
+def _retain_open_chunk(
+    state: Any,
+    host: Any,
+    message_dict: dict[str, Any],
+    index: int,
+    merged: AIMessageChunk | None,
+    churn_id: str | None,
+) -> None:
+    messages = plain(state["messages"])
+    if index >= len(messages):
+        return
+    stored = messages[index]
+    chunk = (
+        merged
+        if merged is not None
+        else AIMessageChunk(**{**message_dict, "type": "AIMessageChunk"})
+    )
+    if (
+        not isinstance(stored, dict)
+        or stored.get("id") != chunk.id
+        or stored.get("chunk_position", "last") == "last"
+    ):
+        return
+    _open_chunks[host] = _OpenChunk(index, chunk, stored, churn_id)
+
+
+def _patch_message(messages: Any, index: int, next_message: dict[str, Any]) -> None:
+    current_message = plain(messages[index])
+    if not isinstance(current_message, dict) or not _can_patch_value(
+        current_message, next_message
+    ):
+        messages[index] = next_message
+        return
+
+    _patch_child(messages, index, next_message)
+
+
+def _dump_message(message: Any) -> dict[str, Any]:
+    message_dict = message if isinstance(message, dict) else message.model_dump()
+    if message_dict.get("type") == "AIMessageChunk":
+        message_dict = {**message_dict, "type": "ai"}
+    return message_dict
+
+
+def _patch_message_fields(
+    target: Any, current: dict[str, Any], next_message: dict[str, Any]
+) -> None:
+    for key in [key for key in current if key not in next_message]:
+        del target[key]
+    for key, value in next_message.items():
+        if key not in current or current[key] != value:
+            target[key] = value
+
+
+def replace_messages(state: Any, messages: list[Any]) -> None:
+    """Rebaseline ``state["messages"]`` to ``messages``, emitting granular ops."""
+    dumped = [_dump_message(message) for message in messages]
+
+    host = getattr(state, "_host", None)
+    if host is not None:
+        _open_chunks.pop(host, None)
+
+    if "messages" not in state:
+        state["messages"] = dumped
+        return
+
+    target = state["messages"]
+    stored = plain(target)
+    stored_len = len(stored)
+
+    matched = 0
+    while (
+        matched < stored_len
+        and matched < len(dumped)
+        and isinstance(stored[matched], dict)
+        and stored[matched].get("id") == dumped[matched].get("id")
+    ):
+        _patch_message_fields(target[matched], stored[matched], dumped[matched])
+        matched += 1
+
+    for index in range(matched, min(stored_len, len(dumped))):
+        target[index] = dumped[index]
+    for index in range(stored_len, len(dumped)):
+        target.append(dumped[index])
+    for index in range(stored_len - 1, len(dumped) - 1, -1):
+        del target[index]
+
+
+def append_langgraph_event(
+    state: Any, _namespace: str, type: str, payload: Any
+) -> None:
+    """Fold a LangGraph stream event ('messages' or 'updates') into the state."""
+    if type == "messages":
+        if "messages" not in state:
+            state["messages"] = []
+
+        message = payload[0]
+        message_dict = message.model_dump()
+
+        is_ai_message_chunk = message_dict.get("type") == "AIMessageChunk"
+        if is_ai_message_chunk:
+            message_dict["type"] = "ai"
+
+        host = getattr(state, "_host", None)
+        entry = _open_chunks.get(host) if host is not None else None
+        if entry is not None and is_ai_message_chunk:
+            merged = _fast_chunk_merge(state, entry, message, message_dict)
+            if merged is not None:
+                _open_chunks[host] = entry._replace(chunk=merged)
+                return
+        if host is not None:
+            _open_chunks.pop(host, None)
+
+        merged = None
+        scan_index = _find_existing_message_index(state["messages"], message_dict)
+        existing_message_index = scan_index
+        if existing_message_index is None and is_ai_message_chunk:
+            existing_message_index = _open_chunk_tail_index(
+                state["messages"], message_dict
+            )
+
+        if existing_message_index is not None:
+            if is_ai_message_chunk:
+                existing_message = plain(state["messages"][existing_message_index])
+                merged = add_ai_message_chunks(
+                    AIMessageChunk(**{**existing_message, "type": "AIMessageChunk"}),
+                    AIMessageChunk(**{**message_dict, "type": "AIMessageChunk"}),
+                )
+                new_message_dict = merged.model_dump()
+                new_message_dict["type"] = "ai"
+                _patch_message(
+                    state["messages"],
+                    existing_message_index,
+                    new_message_dict,
+                )
+            else:
+                state["messages"][existing_message_index] = message_dict
+        else:
+            state["messages"].append(message_dict)
+
+        if host is not None and is_ai_message_chunk:
+            index = (
+                existing_message_index
+                if existing_message_index is not None
+                else len(plain(state["messages"])) - 1
+            )
+            churn_id = message_dict.get("id") if scan_index is None else None
+            _retain_open_chunk(state, host, message_dict, index, merged, churn_id)
+
+    elif type == "updates":
+        for _node_name, channels in payload.items():
+            if not isinstance(channels, dict):
+                continue
+            for channel_name, channel_value in channels.items():
+                if channel_name == "messages":
+                    continue
+
+                state[channel_name] = channel_value

{statewire-0.8.2 → statewire-0.8.4}/src/statewire/ops.py

@@ -125,29 +125,44 @@ silence; 30s of receive silence means the line is dead.
 """
 
 import json
-from typing import Any
+from datetime import datetime, timezone
+from typing import Any, NamedTuple
+
+try:
+    import orjson
+except ImportError:
+    orjson = None
 
-Op = dict[str, Any]
 Path = list[str | int]
 Meta = list[dict[str, Any]] | None
 
 
-def replace_op(path: Path, value: Any, meta: Meta = None) -> Op:
-    op = {"op": "replace", "path": path, "value": value}
+class Op(NamedTuple):
+    """A recorded state change: kind/path/meta as data for the host's date
+    registry, plus its wire bytes, encoded at record time."""
+
+    kind: str
+    path: Path
+    meta: Meta
+    data: bytes
+
+
+def _op(kind: str, path: Path, body: dict[str, Any], meta: Meta) -> Op:
     if meta:
-        op["meta"] = meta
-    return op
+        body["meta"] = meta
+    return Op(kind, path, meta or None, encode_envelope(body))
+
+
+def replace_op(path: Path, value: Any, meta: Meta = None) -> Op:
+    return _op("replace", path, {"op": "replace", "path": path, "value": value}, meta)
 
 
 def add_op(path: Path, value: Any, meta: Meta = None) -> Op:
-    op = {"op": "add", "path": path, "value": value}
-    if meta:
-        op["meta"] = meta
-    return op
+    return _op("add", path, {"op": "add", "path": path, "value": value}, meta)
 
 
 def remove_op(path: Path) -> Op:
-    return {"op": "remove", "path": path}
+    return _op("remove", path, {"op": "remove", "path": path}, None)
 
 
 PROTOCOL_VERSION = 1
@@ -155,12 +170,33 @@ WS_SUBPROTOCOL = "statewire.v1"
 VERSION_HEADER = "Statewire-Version"
 
 
-def snapshot_event(state: dict[str, Any], meta: Meta = None) -> list[Op]:
-    return [replace_op([], state, meta)]
+def _iso(value: datetime) -> str:
+    return value.astimezone(timezone.utc).isoformat(timespec="milliseconds")[:-6] + "Z"
+
+
+def _iso_default(value: Any) -> str:
+    if not isinstance(value, datetime):
+        raise TypeError(f"{type(value).__name__} is not wire-encodable")
+    return _iso(value)
 
 
 def encode_envelope(envelope: dict[str, Any]) -> bytes:
-    return json.dumps(envelope, separators=(",", ":"), allow_nan=False).encode()
+    if orjson is not None:
+        # orjson nullifies non-finite floats instead of raising; _frozen
+        # rejects them before any envelope value is built.
+        return orjson.dumps(
+            envelope, default=_iso_default, option=orjson.OPT_PASSTHROUGH_DATETIME
+        )
+    return json.dumps(
+        envelope, separators=(",", ":"), allow_nan=False, default=_iso_default
+    ).encode()
+
+
+def ops_envelope(entries: list[Op], rest: dict[str, Any] | None = None) -> bytes:
+    head = b'{"ops":[' + b",".join(op.data for op in entries) + b"]"
+    if not rest:
+        return head + b"}"
+    return head + b"," + encode_envelope(rest)[1:]
 
 
 def sse_frame(data: bytes) -> bytes: