awp-python 0.1.0a1__py3-none-any.whl

awp_sim/loopback.py

@@ -0,0 +1,222 @@
+"""Run agents against a World in-process on a virtual clock, recording wire traces.
+
+The loopback delivers messages instantly and in order; time moves only when told to. Scenarios and
+tests use it to exercise the real client and world deterministically, including failures.
+"""
+
+from __future__ import annotations
+
+import itertools
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from typing import Any
+
+from awp.client import ClientConnection, ErrorResponse, Event, Response
+from awp.jsonrpc import Message
+
+from .world import Close, Output, Send, SendFrame, World
+
+MS = 1_000_000
+
+
+@dataclass
+class TraceLine:
+    sender: str
+    msg: Message
+    delivered: bool = True
+
+    def to_json(self) -> dict[str, Any]:
+        method = self.msg.get("method") or ("error" if "error" in self.msg else "result")
+        line: dict[str, Any] = {"from": self.sender, "step": method, "msg": self.msg}
+        if not self.delivered:
+            line["delivered"] = False
+        return line
+
+
+class Loopback:
+    def __init__(self, world: World, *, step_ms: float = 1.0, start_ns: int = 1_000 * MS) -> None:
+        self.world = world
+        self.now = start_ns
+        self.step_ns = int(step_ms * MS)
+        self.agents: list[LoopbackAgent] = []
+        self._conn_ids = itertools.count(1)
+
+    def agent(
+        self, name: str = "agent", *, clock_offset_ns: int = -5 * MS, **kw: Any
+    ) -> LoopbackAgent:
+        a = LoopbackAgent(self, name, clock_offset_ns, **kw)
+        self.agents.append(a)
+        return a
+
+    def advance(self, ms: float) -> None:
+        end = self.now + int(ms * MS)
+        while self.now < end:
+            self.now = min(end, self.now + self.step_ns)
+            for a in self.agents:
+                a.tick()
+            self.deliver(self.world.advance(self.now))
+            self.settle()
+
+    def run_until(self, predicate: Callable[[], bool], timeout_ms: float) -> bool:
+        end = self.now + int(timeout_ms * MS)
+        while not predicate():
+            if self.now >= end:
+                return False
+            self.advance(self.step_ns / MS)
+        return True
+
+    def settle(self) -> None:
+        """Exchange messages until neither side has anything left to send."""
+        for _ in range(10_000):
+            if not any(a.flush() for a in self.agents):
+                return
+        raise RuntimeError("loopback did not settle")
+
+    def deliver(self, outputs: Iterable[Output]) -> None:
+        for out in outputs:
+            stream = next(
+                (a for a in self.agents if a.stream is not None and a.stream == out.conn), None
+            )
+            if stream is not None:
+                if isinstance(out, SendFrame):
+                    data = self.world.frame_bytes(out, self.now)
+                    stream.events += stream.client.receive_frame(data)
+                elif isinstance(out, Close):
+                    stream.stream = None
+                continue
+            agent = next((a for a in self.agents if a.conn == out.conn), None)
+            if agent is None or isinstance(out, SendFrame):
+                continue
+            if isinstance(out, Close):
+                agent.transport_closed()
+            else:
+                agent.on_world(out)
+
+
+class LoopbackAgent:
+    def __init__(
+        self,
+        net: Loopback,
+        name: str,
+        clock_offset_ns: int,
+        *,
+        agent: dict[str, str] | None = None,
+        modalities: Iterable[str] = ("proprio/json", "text/event+json"),
+        heartbeat_ms: float | None = 500,
+        approver: bool = False,
+    ) -> None:
+        self.approver = approver
+        self.net = net
+        self.name = name
+        self.offset = clock_offset_ns
+        self.client = ClientConnection(
+            agent or {"name": name, "version": "0.1.0", "vendor": "awp-sim"},
+            modalities,
+            clock_ns=lambda: self.net.now + self.offset,
+        )
+        self.heartbeat_ms = heartbeat_ms
+        self.conn: int | None = None
+        self.stream: int | None = None  # a ws stream connection, when attached
+        self.events: list[Event] = []
+        self.trace: list[TraceLine] = []
+        # Fault: the first world message matching this is lost and the connection goes half-open.
+        self.lose: Callable[[Message], bool] | None = None
+        self._last_ping = 0
+
+    # ------------------------------------------------------------ transport
+
+    def connect(self) -> None:
+        if self.conn is not None:
+            self.drop()
+        self.conn = next(self.net._conn_ids)
+        self.net.deliver(self.net.world.connect(self.conn, self.net.now, approver=self.approver))
+
+    def attach_stream(self) -> None:
+        """Open a stream connection with the session token (AWP-TRN-003, AWP-SEC-004)."""
+        assert self.client.session_token is not None
+        self.stream = next(self.net._conn_ids)
+        self.net.deliver(
+            self.net.world.attach_stream(self.stream, self.client.session_token, self.net.now)
+        )
+
+    def drop_stream(self) -> None:
+        if self.stream is not None:
+            stream, self.stream = self.stream, None
+            self.net.deliver(self.net.world.stream_lost(stream, self.net.now))
+
+    def drop(self) -> None:
+        """The agent's side of the connection dies."""
+        if self.conn is None:
+            return
+        conn, self.conn = self.conn, None
+        self.client.connection_lost()
+        self.net.deliver(self.net.world.disconnect(conn, self.net.now))
+
+    def abandon(self) -> None:
+        """The agent loses the connection but the world has not noticed: it is half-open."""
+        self.conn = None
+        self.client.connection_lost()
+
+    def transport_closed(self) -> None:
+        self.conn = None
+        self.client.connection_lost()
+
+    def flush(self) -> bool:
+        out = self.client.outgoing()
+        if self.conn is None:
+            return False
+        for msg in out:
+            self.trace.append(TraceLine("agent", msg))
+            self.net.deliver(self.net.world.receive(self.conn, msg, self.net.now))
+        return bool(out)
+
+    def on_world(self, send: Send) -> None:
+        msg = (
+            self.net.world.frame_sent(send, self.net.now)
+            if send.msg.get("method") == "obs.frame"
+            else send.msg
+        )
+        if self.lose is not None and self.lose(msg):
+            self.lose = None
+            self.trace.append(TraceLine("world", msg, delivered=False))
+            self.abandon()
+            return
+        self.trace.append(TraceLine("world", msg))
+        self.events += self.client.receive(msg)
+
+    def tick(self) -> None:
+        if self.heartbeat_ms is None or self.conn is None or self.client.ready is None:
+            return
+        if self.net.now - self._last_ping >= self.heartbeat_ms * MS:
+            self._last_ping = self.net.now
+            self.client.ping()
+
+    # ------------------------------------------------------------ requests
+
+    def call(self, rid: int) -> dict[str, Any]:
+        """Settle and return the result of request `rid`; raise AwpError on an error response."""
+        self.net.settle()
+        for e in self.events:
+            if isinstance(e, Response) and e.id == rid:
+                return e.result
+            if isinstance(e, ErrorResponse) and e.id == rid:
+                raise e.error
+        raise TimeoutError(f"no response to request {rid}")
+
+    def submit(self, type: str, params: dict[str, Any], **kw: Any) -> str:
+        action_id = self.client.submit(type, params, **kw)
+        self.call(self.client.last_id)
+        return action_id
+
+    def open(self, **kw: Any) -> dict[str, Any]:
+        self.connect()
+        self.call(self.client.initialize())
+        ready = self.call(self.client.open_session(**kw))
+        self.call(self.client.ping())  # a clock sample on the new session clock (AWP-CLK-008)
+        return ready
+
+    def of(self, kind: type[Any]) -> list[Any]:
+        return [e for e in self.events if isinstance(e, kind)]
+
+    def trace_json(self) -> list[dict[str, Any]]:
+        return [line.to_json() for line in self.trace]

awp_sim/recorder.py

@@ -0,0 +1,72 @@
+"""Record each session's wire traffic as a trace in the spec's JSON Lines format."""
+
+from __future__ import annotations
+
+import json
+from collections import OrderedDict
+from collections.abc import Hashable
+from pathlib import Path
+from typing import IO, Any
+
+from awp.jsonrpc import Message
+
+from .audit import redact
+
+
+class TraceRecorder:
+    """Writes `<directory>/<session_id>.jsonl`, with credentials redacted as in the audit log.
+
+    Messages exchanged before a connection has a session (initialize, a failed open, the
+    reconnect before session.resume) are held until it does, then written ahead of the session's
+    own traffic.
+    """
+
+    PENDING_LIMIT = 256
+    OPEN_FILES = 64
+
+    def __init__(self, directory: Path | str) -> None:
+        self.directory = Path(directory)
+        self.directory.mkdir(parents=True, exist_ok=True)
+        self._pending: dict[Hashable, list[dict[str, Any]]] = {}
+        self._files: OrderedDict[str, IO[str]] = OrderedDict()
+
+    def record(
+        self,
+        conn: Hashable,
+        session: str | None,
+        sender: str,
+        msg: Message,
+        *,
+        delivered: bool = True,
+    ) -> None:
+        method = msg.get("method") or ("error" if "error" in msg else "result")
+        line: dict[str, Any] = {"from": sender, "step": method, "msg": redact(msg)}
+        if not delivered:
+            line["delivered"] = False
+        if session is None:
+            held = self._pending.setdefault(conn, [])
+            if len(held) < self.PENDING_LIMIT:
+                held.append(line)
+            return
+        f = self._file(session)
+        for held_line in self._pending.pop(conn, []):
+            f.write(json.dumps(held_line, separators=(",", ":")) + "\n")
+        f.write(json.dumps(line, separators=(",", ":")) + "\n")
+        f.flush()
+
+    def forget(self, conn: Hashable) -> None:
+        self._pending.pop(conn, None)
+
+    def close(self) -> None:
+        for f in self._files.values():
+            f.close()
+        self._files.clear()
+
+    def _file(self, session: str) -> IO[str]:
+        f = self._files.pop(session, None)
+        if f is None:
+            f = (self.directory / f"{session}.jsonl").open("a", encoding="utf-8")
+            while len(self._files) >= self.OPEN_FILES:
+                self._files.popitem(last=False)[1].close()
+        self._files[session] = f  # most recently used last
+        return f

awp_sim/replay.py

@@ -0,0 +1,134 @@
+"""Replay a replay bundle and compare what the world does with what it recorded (AWP-REP-003).
+
+The world is rebuilt from the bundle's configuration and initial state; the session is opened as
+recorded, and the agent's calls that change action state are fed in their recorded order.
+Compared: every frame's payload hash per channel, and the ordered `(state, reason)` sequence of
+every action. Timestamps, sequence numbers, heartbeats, reports, and telemetry are not.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from awp.client import ActionUpdated, FrameReceived
+
+from .loopback import Loopback
+from .world import World, decode_config, decode_state
+
+# What changes action state: the agent's calls, in their recorded order.
+REPLAYED = (
+    "action.submit",
+    "action.cancel",
+    "world.tick",
+    "world.reset",
+    "world.restore",
+    "session.close",
+)
+
+
+@dataclass
+class Outcome:
+    frames: int
+    transitions: int
+    difference: str | None
+
+    @property
+    def reproduced(self) -> bool:
+        return self.difference is None
+
+
+def _recorded(
+    records: list[dict[str, Any]],
+) -> tuple[dict[str, list[str]], dict[str, list[tuple[str, Any]]]]:
+    channel_of: dict[int, str] = {}
+    frames: dict[str, list[str]] = {}
+    states: dict[str, list[tuple[str, Any]]] = {}
+    for r in records:
+        body = r["body"]
+        if r["kind"] == "frame" and body.get("method") == "obs.frame":
+            name = channel_of.get(body["channel_id"], str(body["channel_id"]))
+            frames.setdefault(name, []).append(body["payload_sha256"])
+        elif r["kind"] == "message":
+            result = body.get("result") or {}
+            for g in (
+                (result.get("granted") or {}).get("channels", [])
+                if isinstance(result.get("granted"), dict)
+                else []
+            ):
+                channel_of[g["channel_id"]] = g["channel"]
+            params = body.get("params") or {}
+            status = params if body.get("method") == "action.status" else None
+            if (
+                status is None
+                and "action_id" in result
+                and "state" in result
+                and "status_seq" in result
+            ):
+                status = result
+            if status is not None and r["direction"] == "world":
+                seq = states.setdefault(status["action_id"], [])
+                entry = (status["state"], status.get("reason"))
+                if not seq or seq[-1] != entry:
+                    seq.append(entry)
+    return frames, states
+
+
+def replay(path: Path | str) -> Outcome:
+    records = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
+    header = records[0]
+    if header.get("class") != "replay_bundle":
+        raise ValueError(f"{path} is a {header.get('class')}, not a replay_bundle (AWP-AUD-007)")
+    body = header["body"]
+    world = World(decode_config(body["config"]))
+    world.load_state(decode_state(body["initial_state"]))
+    net = Loopback(world)
+    agent = net.agent("replay", heartbeat_ms=None)
+    agent.connect()
+    agent.call(agent.client.initialize())
+    opened = next(r["body"]["params"] for r in records if r["body"].get("method") == "session.open")
+    agent.call(agent.client.request("session.open", opened))
+    for r in records:
+        msg = r["body"]
+        if r["direction"] == "agent" and msg.get("method") in REPLAYED and "id" in msg:
+            agent.client.request(msg["method"], msg.get("params") or {})
+            net.settle()
+    want_frames, want_states = _recorded(records)
+    got_frames: dict[str, list[str]] = {}
+    for e in agent.of(FrameReceived):
+        got_frames.setdefault(e.channel, []).append(hashlib.sha256(e.frame.payload).hexdigest())
+    got_states: dict[str, list[tuple[str, Any]]] = {}
+    for e in agent.of(ActionUpdated):
+        seq = got_states.setdefault(e.action.action_id, [])
+        entry = (e.status["state"], e.status.get("reason"))
+        if not seq or seq[-1] != entry:
+            seq.append(entry)
+    difference = None
+    for name, hashes in want_frames.items():
+        if got_frames.get(name, []) != hashes:
+            got = got_frames.get(name, [])
+            at = next(
+                (i for i, (a, b) in enumerate(zip(hashes, got, strict=False)) if a != b),
+                min(len(hashes), len(got)),
+            )
+            difference = (
+                f"channel {name}: frame {at} differs ({len(hashes)} recorded, {len(got)} replayed)"
+            )
+            break
+    if difference is None and got_states != want_states:
+        bad = next(
+            k
+            for k in sorted(set(want_states) | set(got_states))
+            if want_states.get(k) != got_states.get(k)
+        )
+        difference = (
+            f"action {bad}: recorded {want_states.get(bad)}, replayed {got_states.get(bad)}"
+        )
+    return Outcome(
+        sum(len(v) for v in want_frames.values()),
+        sum(len(v) for v in want_states.values()),
+        difference,
+    )