feat(workers): use types for worker messages (#1767)

Author: paul-nechifor

dimos/core/coordination/python_worker.py

@@ -22,6 +22,17 @@
 import traceback
 from typing import TYPE_CHECKING, Any
 
+from dimos.core.coordination.worker_messages import (
+    CallMethodRequest,
+    DeployModuleRequest,
+    GetAttrRequest,
+    SetRefRequest,
+    ShutdownRequest,
+    SuppressConsoleRequest,
+    UndeployModuleRequest,
+    WorkerRequest,
+    WorkerResponse,
+)
 from dimos.core.global_config import GlobalConfig, global_config
 from dimos.core.library_config import apply_library_config
 from dimos.utils.logging_config import setup_logger
@@ -63,7 +74,12 @@ def __getattr__(self, name: str) -> Any:
 
         def _call(*args: Any, **kwargs: Any) -> ActorFuture:
             result = self._actor._send_request_to_worker(
-                {"type": "call_method", "name": name, "args": args, "kwargs": kwargs}
+                CallMethodRequest(
+                    module_id=self._actor._module_id,
+                    name=name,
+                    args=args,
+                    kwargs=kwargs,
+                )
             )
             return ActorFuture(result)
 
@@ -91,34 +107,33 @@ def __reduce__(self) -> tuple[type, tuple[None, type, int, int, None]]:
         """Exclude the connection and lock when pickling."""
         return (Actor, (None, self._cls, self._worker_id, self._module_id, None))
 
-    def _send_request_to_worker(self, request: dict[str, Any]) -> Any:
+    def _send_request_to_worker(self, request: WorkerRequest) -> Any:
         if self._conn is None:
             raise RuntimeError("Actor connection not available - cannot send requests")
-        request["module_id"] = self._module_id
         if self._lock is not None:
             with self._lock:
                 self._conn.send(request)
-                response = self._conn.recv()
+                response: WorkerResponse = self._conn.recv()
         else:
             self._conn.send(request)
             response = self._conn.recv()
-        if response.get("error"):
-            if "AttributeError" in response["error"]:  # TODO: better error handling
-                raise AttributeError(response["error"])
-            raise RuntimeError(f"Worker error: {response['error']}")
-        return response.get("result")
+        if response.error:
+            if "AttributeError" in response.error:  # TODO: better error handling
+                raise AttributeError(response.error)
+            raise RuntimeError(f"Worker error: {response.error}")
+        return response.result
 
     def set_ref(self, ref: Any) -> ActorFuture:
         """Set the actor reference on the remote module."""
-        result = self._send_request_to_worker({"type": "set_ref", "ref": ref})
+        result = self._send_request_to_worker(SetRefRequest(module_id=self._module_id, ref=ref))
         return ActorFuture(result)
 
     def __getattr__(self, name: str) -> Any:
         """Proxy attribute access to the worker process."""
         if name.startswith("_"):
             raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
 
-        return self._send_request_to_worker({"type": "getattr", "name": name})
+        return self._send_request_to_worker(GetAttrRequest(module_id=self._module_id, name=name))
 
 
 # Global forkserver context. Using `forkserver` instead of `fork` because it
@@ -208,20 +223,14 @@ def deploy_module(
         kwargs["g"] = global_config
         module_id = _module_ids.next()
 
-        # Send deploy_module request to the worker process
-        request = {
-            "type": "deploy_module",
-            "module_id": module_id,
-            "module_class": module_class,
-            "kwargs": kwargs,
-        }
+        request = DeployModuleRequest(module_id=module_id, module_class=module_class, kwargs=kwargs)
         try:
             with self._lock:
                 self._conn.send(request)
-                response = self._conn.recv()
+                response: WorkerResponse = self._conn.recv()
 
-            if response.get("error"):
-                raise RuntimeError(f"Failed to deploy module: {response['error']}")
+            if response.error:
+                raise RuntimeError(f"Failed to deploy module: {response.error}")
 
             actor = Actor(self._conn, module_class, self._worker_id, module_id, self._lock)
             actor.set_ref(actor).result()
@@ -243,11 +252,11 @@ def undeploy_module(self, module_id: int) -> None:
             raise RuntimeError("Worker process not started")
 
         with self._lock:
-            self._conn.send({"type": "undeploy_module", "module_id": module_id})
-            response = self._conn.recv()
+            self._conn.send(UndeployModuleRequest(module_id=module_id))
+            response: WorkerResponse = self._conn.recv()
 
-        if response.get("error"):
-            raise RuntimeError(f"Failed to undeploy module: {response['error']}")
+        if response.error:
+            raise RuntimeError(f"Failed to undeploy module: {response.error}")
 
         self._modules.pop(module_id, None)
 
@@ -256,7 +265,7 @@ def suppress_console(self) -> None:
             return
         try:
             with self._lock:
-                self._conn.send({"type": "suppress_console"})
+                self._conn.send(SuppressConsoleRequest())
                 self._conn.recv()
         except (BrokenPipeError, EOFError, ConnectionResetError):
             pass
@@ -265,7 +274,7 @@ def shutdown(self) -> None:
         if self._conn is not None:
             try:
                 with self._lock:
-                    self._conn.send({"type": "shutdown"})
+                    self._conn.send(ShutdownRequest())
                     if self._conn.poll(timeout=5):
                         self._conn.recv()
                     else:
@@ -353,54 +362,48 @@ def _worker_loop(conn: Connection, instances: dict[int, Any], worker_id: int) ->
         except (EOFError, KeyboardInterrupt):
             break
 
-        response: dict[str, Any] = {}
+        response: WorkerResponse
         try:
-            req_type = request.get("type")
-
-            if req_type == "deploy_module":
-                module_class = request["module_class"]
-                kwargs = request["kwargs"]
-                module_id = request["module_id"]
-                instance = module_class(**kwargs)
-                instances[module_id] = instance
-                response["result"] = module_id
-
-            elif req_type == "set_ref":
-                module_id = request["module_id"]
-                instances[module_id].ref = request.get("ref")
-                response["result"] = worker_id
-
-            elif req_type == "getattr":
-                module_id = request["module_id"]
-                response["result"] = getattr(instances[module_id], request["name"])
-
-            elif req_type == "call_method":
-                module_id = request["module_id"]
-                method = getattr(instances[module_id], request["name"])
-                result = method(*request.get("args", ()), **request.get("kwargs", {}))
-                response["result"] = result
-
-            elif req_type == "undeploy_module":
-                module_id = request["module_id"]
-                instance = instances.pop(module_id, None)
-                if instance is not None:
-                    instance.stop()
-                response["result"] = True
-
-            elif req_type == "suppress_console":
-                _suppress_console_output()
-                response["result"] = True
-
-            elif req_type == "shutdown":
-                response["result"] = True
-                conn.send(response)
-                break
-
-            else:
-                response["error"] = f"Unknown request type: {req_type}"
+            match request:
+                case DeployModuleRequest(
+                    module_id=module_id, module_class=module_class, kwargs=kwargs
+                ):
+                    instance = module_class(**kwargs)
+                    instances[module_id] = instance
+                    response = WorkerResponse(result=module_id)
+
+                case SetRefRequest(module_id=module_id, ref=ref):
+                    instances[module_id].ref = ref
+                    response = WorkerResponse(result=worker_id)
+
+                case GetAttrRequest(module_id=module_id, name=name):
+                    response = WorkerResponse(result=getattr(instances[module_id], name))
+
+                case CallMethodRequest(module_id=module_id, name=name, args=args, kwargs=kwargs):
+                    method = getattr(instances[module_id], name)
+                    response = WorkerResponse(result=method(*args, **kwargs))
+
+                case UndeployModuleRequest(module_id=module_id):
+                    instance = instances.pop(module_id, None)
+                    if instance is not None:
+                        instance.stop()
+                    response = WorkerResponse(result=True)
+
+                case SuppressConsoleRequest():
+                    _suppress_console_output()
+                    response = WorkerResponse(result=True)
+
+                case ShutdownRequest():
+                    conn.send(WorkerResponse(result=True))
+                    break
+
+                case _:
+                    response = WorkerResponse(error=f"Unknown request type: {type(request)}")
 
         except Exception as e:
-            response["error"] = f"{e.__class__.__name__}: {e}\n{traceback.format_exc()}"
+            response = WorkerResponse(
+                error=f"{e.__class__.__name__}: {e}\n{traceback.format_exc()}"
+            )
 
         try:
             conn.send(response)

dimos/core/coordination/worker_messages.py

@@ -0,0 +1,79 @@
+# Copyright 2026 Dimensional Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from dimos.core.module import ModuleBase
+
+
+@dataclass(frozen=True)
+class DeployModuleRequest:
+    module_id: int
+    module_class: type[ModuleBase]
+    kwargs: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class SetRefRequest:
+    module_id: int
+    ref: Any
+
+
+@dataclass(frozen=True)
+class GetAttrRequest:
+    module_id: int
+    name: str
+
+
+@dataclass(frozen=True)
+class CallMethodRequest:
+    module_id: int
+    name: str
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class UndeployModuleRequest:
+    module_id: int
+
+
+@dataclass(frozen=True)
+class SuppressConsoleRequest:
+    pass
+
+
+@dataclass(frozen=True)
+class ShutdownRequest:
+    pass
+
+
+WorkerRequest = (
+    DeployModuleRequest
+    | SetRefRequest
+    | GetAttrRequest
+    | CallMethodRequest
+    | UndeployModuleRequest
+    | SuppressConsoleRequest
+    | ShutdownRequest
+)
+
+
+@dataclass(frozen=True)
+class WorkerResponse:
+    result: Any = None
+    error: str | None = None