From 99347d3f004f9501dd8d4ae3b310bcf224a1e830 Mon Sep 17 00:00:00 2001
From: Joseph Doherty <dohejw01@gmail.com>
Date: Thu, 22 Oct 2026 06:06:54 -0400
Subject: [PATCH] feat(python): typed
 ping/get_session_state/get_worker_info/drain_events session helpers

---
 clients/python/README.md                      |  7 +++
 .../python/src/zb_mom_ww_mxgateway/session.py | 49 ++++++++++++++++
 .../tests/test_typed_command_helpers.py       | 58 +++++++++++++++++++
 3 files changed, 114 insertions(+)

diff --git a/clients/python/README.md b/clients/python/README.md
index 64ddd9b..5f5a2b9 100644
--- a/clients/python/README.md
+++ b/clients/python/README.md
@@ -523,3 +523,10 @@ a `~/.netrc` entry or `PIP_INDEX_URL=https://<user>:<token>@gitea.dohertylan.com
 - [Client Packaging](../../docs/ClientPackaging.md)
 - [Client Proto Generation](../../docs/ClientProtoGeneration.md)
 - [Python Client Detailed Design](./PythonClientDesign.md)
+
+## Typed diagnostics
+
+`Session.ping(message)`, `get_session_state()`, `get_worker_info()`, and
+`drain_events(max_events)` wrap the worker diagnostic commands. `drain_events`
+returns a list that may be shorter than `max_events`; call it until it returns
+an empty list. Replies drained this way carry no `worker_sequence`.
diff --git a/clients/python/src/zb_mom_ww_mxgateway/session.py b/clients/python/src/zb_mom_ww_mxgateway/session.py
index 8a8d3bf..8ff3f9f 100644
--- a/clients/python/src/zb_mom_ww_mxgateway/session.py
+++ b/clients/python/src/zb_mom_ww_mxgateway/session.py
@@ -838,6 +838,55 @@ class Session:
         )
         return reply.activate.status
 
+    async def ping(self, message: str = "", *, correlation_id: str = "") -> pb.MxCommandReply:
+        """Invoke the worker `Ping` diagnostic and return its reply."""
+        return await self.invoke(
+            pb.MxCommand(kind=pb.MX_COMMAND_KIND_PING, ping=pb.PingCommand(message=message)),
+            correlation_id=correlation_id,
+        )
+
+    async def get_session_state(self, *, correlation_id: str = "") -> int:
+        """Return the worker-reported `SessionState` enum value."""
+        reply = await self.invoke(
+            pb.MxCommand(
+                kind=pb.MX_COMMAND_KIND_GET_SESSION_STATE,
+                get_session_state=pb.GetSessionStateCommand(),
+            ),
+            correlation_id=correlation_id,
+        )
+        return reply.session_state.state
+
+    async def get_worker_info(self, *, correlation_id: str = "") -> pb.WorkerInfoReply:
+        """Return worker process id, version, and MXAccess ProgID/CLSID."""
+        reply = await self.invoke(
+            pb.MxCommand(
+                kind=pb.MX_COMMAND_KIND_GET_WORKER_INFO,
+                get_worker_info=pb.GetWorkerInfoCommand(),
+            ),
+            correlation_id=correlation_id,
+        )
+        return reply.worker_info
+
+    async def drain_events(
+        self,
+        max_events: int = 0,
+        *,
+        correlation_id: str = "",
+    ) -> list[pb.MxEvent]:
+        """Drain queued worker events (diagnostic).
+
+        A reply may carry fewer than *max_events* (server count and frame-size
+        caps); call repeatedly until an empty list is returned.
+        """
+        reply = await self.invoke(
+            pb.MxCommand(
+                kind=pb.MX_COMMAND_KIND_DRAIN_EVENTS,
+                drain_events=pb.DrainEventsCommand(max_events=max_events),
+            ),
+            correlation_id=correlation_id,
+        )
+        return list(reply.drain_events.events)
+
     def stream_events(
         self,
         *,
diff --git a/clients/python/tests/test_typed_command_helpers.py b/clients/python/tests/test_typed_command_helpers.py
index 3e25686..77cbabf 100644
--- a/clients/python/tests/test_typed_command_helpers.py
+++ b/clients/python/tests/test_typed_command_helpers.py
@@ -234,3 +234,61 @@ async def test_suspend_and_activate_return_status() -> None:
     )
     status = await session.activate(12, 34)
     assert status.category == pb.MX_STATUS_CATEGORY_OK
+
+
+@pytest.mark.asyncio
+async def test_ping_sends_message_and_returns_reply() -> None:
+    session, stub = await _session_with([_ok(pb.MX_COMMAND_KIND_PING)])
+    reply = await session.ping("hello")
+    command = stub.invoke.requests[0].command
+    assert command.kind == pb.MX_COMMAND_KIND_PING
+    assert command.ping.message == "hello"
+    assert reply.kind == pb.MX_COMMAND_KIND_PING
+
+
+@pytest.mark.asyncio
+async def test_get_session_state_returns_state() -> None:
+    session, stub = await _session_with(
+        [
+            _ok(
+                pb.MX_COMMAND_KIND_GET_SESSION_STATE,
+                session_state=pb.SessionStateReply(state=pb.SESSION_STATE_CLOSED),
+            ),
+        ],
+    )
+    state = await session.get_session_state()
+    assert state == pb.SESSION_STATE_CLOSED
+    assert stub.invoke.requests[0].command.kind == pb.MX_COMMAND_KIND_GET_SESSION_STATE
+
+
+@pytest.mark.asyncio
+async def test_get_worker_info_returns_info() -> None:
+    session, stub = await _session_with(
+        [
+            _ok(
+                pb.MX_COMMAND_KIND_GET_WORKER_INFO,
+                worker_info=pb.WorkerInfoReply(worker_process_id=42, worker_version="1.2"),
+            ),
+        ],
+    )
+    info = await session.get_worker_info()
+    assert info.worker_process_id == 42
+    assert info.worker_version == "1.2"
+    assert stub.invoke.requests[0].command.kind == pb.MX_COMMAND_KIND_GET_WORKER_INFO
+
+
+@pytest.mark.asyncio
+async def test_drain_events_sends_max_and_returns_events() -> None:
+    session, stub = await _session_with(
+        [
+            _ok(
+                pb.MX_COMMAND_KIND_DRAIN_EVENTS,
+                drain_events=pb.DrainEventsReply(events=[pb.MxEvent(), pb.MxEvent()]),
+            ),
+        ],
+    )
+    events = await session.drain_events(max_events=5)
+    assert len(events) == 2
+    command = stub.invoke.requests[0].command
+    assert command.kind == pb.MX_COMMAND_KIND_DRAIN_EVENTS
+    assert command.drain_events.max_events == 5