Repository: proxystore/aeris

Issue:
Add ZMQ-based exchange server and interface

Hints from issue comments:

Solution patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0de13c3..7867c97 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -8,7 +8,6 @@ repos:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-added-large-files
-      - id: check-docstring-first
       - id: check-json
       - id: check-yaml
       - id: check-merge-conflict
diff --git a/aeris/agent.py b/aeris/agent.py
index ef7acba..b0ffa22 100644
--- a/aeris/agent.py
+++ b/aeris/agent.py
@@ -146,7 +146,7 @@ def _message_handler(self, message: Message) -> Message | None:
                     message.kwargs,
                 )
             except Exception as e:
-                return message.response(exception=e)
+                return message.error(exception=e)
             else:
                 return message.response(result=result)
         elif isinstance(message, PingRequest):
@@ -157,7 +157,7 @@ def _message_handler(self, message: Message) -> Message | None:
             return None
         else:
             raise BadMessageTypeError(
-                f'Agent cannot handle message type: {message}',
+                f'Agent cannot handle message type: {message!r}',
             )
 
     def _message_listener(self) -> None:
diff --git a/aeris/exception.py b/aeris/exception.py
index ad1f3c1..e28e88d 100644
--- a/aeris/exception.py
+++ b/aeris/exception.py
@@ -13,9 +13,17 @@ class BadMessageTypeError(Exception):
     pass
 
 
+class ExchangeRegistrationError(Exception):
+    """Registration with exchange failed."""
+
+    pass
+
+
 class HandleClosedError(Exception):
     """Agent handle has been closed."""
 
+    pass
+
 
 class MailboxClosedError(Exception):
     """Mailbox is closed and cannot send or receive messages."""
diff --git a/aeris/exchange/message.py b/aeris/exchange/message.py
new file mode 100644
index 0000000..5a6da96
--- /dev/null
+++ b/aeris/exchange/message.py
@@ -0,0 +1,169 @@
+"""Exchange message types.
+
+Warning:
+    This module is not considered part of the public API. These message
+    types are helpers for exchange implementations.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+from typing import Optional
+from typing import Union
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import TypeAdapter
+
+from aeris.identifier import Identifier
+from aeris.message import Message
+
+
+class BaseExchangeMessage(BaseModel):
+    """Base exchange message."""
+
+    model_config = ConfigDict(
+        extra='forbid',
+        use_enum_values=True,
+        validate_default=True,
+    )
+
+    @classmethod
+    def model_deserialize(cls, raw: bytes) -> ExchangeMessage:
+        """Reconstruct a specific message from a serialized message.
+
+        Example:
+            ```python
+            from aeris.exchange.message import BaseExchangeMessage
+            from aeris.exchange.message import RegisterMessage
+
+            message = RegisterMessage(...)
+            raw = RegisterMessage.model_serialize()
+            assert BaseExchangeMessage.model_deserialize(raw) == message
+            ```
+        """
+        dump = raw.decode()
+        return TypeAdapter(ExchangeMessage).validate_json(dump)
+
+    def model_serialize(self) -> bytes:
+        """Serialize a message to a bytestring.
+
+        Messages are serialized by dumping the model to a JSON-compatible
+        string and then encoding the string into bytes.
+        """
+        dump = self.model_dump_json()
+        return dump.encode()
+
+
+class RegisterMessage(BaseExchangeMessage):
+    """Exchange registration request message.
+
+    Args:
+        src: Identifier of the entity registering with the exchange.
+    """
+
+    src: Identifier
+    kind: Literal['register'] = Field('register', repr=False)
+
+    def response(self, error: str | None = None) -> ExchangeResponseMessage:
+        """Construct an exchange response message.
+
+        Args:
+            error: Error message if the exchange could not complete the
+                request.
+        """
+        return ExchangeResponseMessage(src=self.src, request=self, error=error)
+
+
+class UnregisterMessage(BaseExchangeMessage):
+    """Exchange unregistration request message.
+
+    Args:
+        src: Identifier of the entity unregistering with the exchange.
+    """
+
+    src: Identifier
+    kind: Literal['unregister'] = Field('unregister', repr=False)
+
+    def response(self, error: str | None = None) -> ExchangeResponseMessage:
+        """Construct an exchange response message.
+
+        Args:
+            error: Error message if the exchange could not complete the
+                request.
+        """
+        return ExchangeResponseMessage(src=self.src, request=self, error=error)
+
+
+class ForwardMessage(BaseExchangeMessage):
+    """Message for exchange to forward to another mailbox.
+
+    Args:
+        src: Source of the message.
+        dest: Destination of the message.
+        message: The message to send forward to `dest` via the exchange.
+    """
+
+    src: Identifier
+    dest: Identifier
+    message: Message
+    kind: Literal['forward'] = Field('forward', repr=False)
+
+    def response(self, error: str | None = None) -> ExchangeResponseMessage:
+        """Construct an exchange response message.
+
+        Args:
+            error: Error message if the exchange could not complete the
+                request.
+        """
+        return ExchangeResponseMessage(src=self.src, request=self, error=error)
+
+
+ExchangeRequestMessage = Union[
+    RegisterMessage,
+    UnregisterMessage,
+    ForwardMessage,
+]
+
+
+class ExchangeResponseMessage(BaseExchangeMessage):
+    """Response message from the exchange to clients.
+
+    Args:
+        src: Identifier of the client that made the request.
+        op: Type of the request the exchange received.
+        error: Error message from the exchange.
+    """
+
+    src: Identifier
+    request: ExchangeRequestMessage
+    error: Optional[str] = None  # noqa: UP007
+    kind: Literal['response'] = Field('response', repr=False)
+
+    @property
+    def success(self) -> bool:
+        """Check if the exchange completed the request.
+
+        If `False`, `error` will be a string containing an error message
+        from the exchange.
+        """
+        return self.error is None
+
+
+ExchangeMessage = Union[ExchangeRequestMessage, ExchangeResponseMessage]
+"""Exchange message union type for type annotations.
+
+Tip:
+    This is a parameterized generic type meaning that this type cannot
+    be used for [`isinstance`][`builtins.isinstance`] checks:
+    ```python
+    isinstance(message, ExchangeMessage)  # Fails
+    ```
+    Instead, use [`typing.get_args()`][typing.get_args]:
+    ```
+    from typing import get_args
+
+    isinstance(message, get_args(ExchangeMessage))  # Works
+    ```
+"""
diff --git a/aeris/exchange/simple.py b/aeris/exchange/simple.py
new file mode 100644
index 0000000..57f832e
--- /dev/null
+++ b/aeris/exchange/simple.py
@@ -0,0 +1,659 @@
+"""Simple message exchange client and server.
+
+To start the exchange:
+```bash
+python -m aeris.exchange.simple --host localhost --port 1234
+```
+
+Connect to the exchange through the client.
+```python
+from aeris.exchange.simple import SimpleExchange
+
+with SimpleExchange('localhost', 1234) as exchange:
+    aid = exchange.register_agent()
+    mailbox = exchange.get_mailbox(aid)
+    ...
+    mailbox.close()
+```
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import contextlib
+import dataclasses
+import io
+import logging
+import queue
+import signal
+import socket
+import sys
+import threading
+from collections.abc import AsyncGenerator
+from collections.abc import Sequence
+from concurrent.futures import Future
+from types import TracebackType
+from typing import cast
+from typing import Generic
+from typing import get_args
+from typing import TypeVar
+
+if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
+    from typing import Self
+else:  # pragma: <3.11 cover
+    from typing_extensions import Self
+
+from aeris.exception import BadIdentifierError
+from aeris.exception import ExchangeRegistrationError
+from aeris.exception import MailboxClosedError
+from aeris.exchange.message import BaseExchangeMessage
+from aeris.exchange.message import ExchangeMessage
+from aeris.exchange.message import ExchangeResponseMessage
+from aeris.exchange.message import ForwardMessage
+from aeris.exchange.message import RegisterMessage
+from aeris.exchange.message import UnregisterMessage
+from aeris.handle import Handle
+from aeris.identifier import AgentIdentifier
+from aeris.identifier import ClientIdentifier
+from aeris.identifier import Identifier
+from aeris.message import Message
+from aeris.message import RequestMessage
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar('T')
+
+DEFAULT_PRIORITY = 0
+DEFAULT_SERVER_TIMEOUT = 30
+CLOSE_PRIORITY = DEFAULT_PRIORITY + 1
+CLOSE_SENTINAL = object()
+
+
+@dataclasses.dataclass(order=True)
+class _QueueItem(Generic[T]):
+    priority: int
+    message: T | object = dataclasses.field(compare=False)
+
+
+class SimpleMailbox:
+    """Thread-safe queue-based mailbox."""
+
+    def __init__(self, uid: Identifier, exchange: SimpleExchange) -> None:
+        self.uid = uid
+        self._exchange = exchange
+        self._queue: queue.PriorityQueue[_QueueItem[Message]] = (
+            queue.PriorityQueue()
+        )
+        self._closed = False
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        exc_traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def _push(self, message: Message) -> None:
+        if self._closed:
+            raise MailboxClosedError
+        self._queue.put(_QueueItem(DEFAULT_PRIORITY, message))
+
+    def send(self, message: Message) -> None:
+        """Send a message to this mailbox.
+
+        Raises:
+            MailboxClosedError: if [`close()`][aeris.exchange.Mailbox.close]
+                has been called.
+        """
+        if self._closed:
+            raise MailboxClosedError
+        wrapped = ForwardMessage(
+            src=message.src,
+            dest=message.dest,
+            message=message,
+        )
+        self._exchange._send_server_message(wrapped)
+
+    def recv(self) -> Message:
+        """Get the next message from this mailbox.
+
+        Raises:
+            MailboxClosedError: if [`close()`][aeris.exchange.Mailbox.close]
+                has been called.
+        """
+        if self._closed:
+            raise MailboxClosedError
+        item = self._queue.get(block=True)
+        message = item.message
+        if message is CLOSE_SENTINAL:  # pragma: no cover
+            raise MailboxClosedError
+        assert isinstance(message, get_args(Message))
+        return message
+
+    def close(self) -> None:
+        """Close the mailbox.
+
+        This unregisters the entity from the exchange.
+        """
+        if not self._closed:
+            self._closed = True
+            self._queue.put(_QueueItem(CLOSE_PRIORITY, CLOSE_SENTINAL))
+            self._exchange._unregister_entity(self.uid)
+
+
+class SimpleExchange:
+    """Simple exchange client.
+
+    Args:
+        host: Host of the exchange server.
+        port: Port of the exchange server.
+        timeout: Timeout when waiting for server responses.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        timeout: float = DEFAULT_SERVER_TIMEOUT,
+    ) -> None:
+        self.host = host
+        self.port = port
+        self.timeout = timeout
+
+        self._socket = socket.create_connection(
+            (self.host, self.port),
+            timeout=self.timeout,
+        )
+        self._socket.setblocking(False)
+        self._handler_thread = threading.Thread(
+            target=self._listen_server_messages,
+        )
+        self._handler_thread.start()
+        logging.debug('%r started server message handler thread', self)
+        self._mailboxes: dict[Identifier, SimpleMailbox] = {}
+
+        self._pending_registration: dict[Identifier, Future[None]] = {}
+        self._pending_unregistration: dict[Identifier, Future[None]] = {}
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        exc_traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def __reduce__(
+        self,
+    ) -> tuple[type[Self], tuple[str, int]]:
+        return (type(self), (self.host, self.port))
+
+    def __repr__(self) -> str:
+        return f'{type(self).__name__}("{self.host}:{self.port}")'
+
+    def __str__(self) -> str:
+        return f'{type(self).__name__}<{self.host}:{self.port}>'
+
+    def _handle_response_message(
+        self,
+        message: ExchangeResponseMessage,
+    ) -> None:
+        if isinstance(message.request, (RegisterMessage, UnregisterMessage)):
+            if isinstance(message.request, RegisterMessage):
+                future = self._pending_registration[message.src]
+            elif isinstance(message.request, UnregisterMessage):
+                future = self._pending_unregistration[message.src]
+            else:
+                raise AssertionError('Unreachable.')
+            if message.error is None:
+                future.set_result(None)
+            else:
+                future.set_exception(ExchangeRegistrationError(message.error))
+        elif (
+            isinstance(message.request, ForwardMessage) and not message.success
+        ):
+            logger.warning(
+                '%s forward message failed: %s',
+                self,
+                message.error,
+            )
+            if isinstance(message.request.message, get_args(RequestMessage)):
+                response = message.request.message.error(
+                    BadIdentifierError(message.error),
+                )
+                self._mailboxes[message.src]._push(response)
+            else:  # pragma: >=3.10 cover
+                # This is covered, by coverage in Python 3.9 doesn't not
+                # detect the empty else pass.
+                # If the failed forward was for a response message, then the
+                # dest entity is likely offline and we cannot recover.
+                pass
+        elif isinstance(message.request, ForwardMessage):
+            # Nothing needs to be done in this case.
+            pass
+        else:
+            raise AssertionError('Unreachable.')
+
+    def _handle_server_message(self, message: ExchangeMessage) -> None:
+        logger.debug('%s received message from server: %r', self, message)
+        if isinstance(message, ForwardMessage):
+            self._mailboxes[message.dest]._push(message.message)
+        elif isinstance(message, ExchangeResponseMessage):
+            self._handle_response_message(message)
+        else:
+            logger.warning(
+                'Unhandled message type from exchange: %r',
+                message,
+            )
+
+    def _listen_server_messages(self) -> None:
+        buffer = io.BytesIO()
+        while True:
+            try:
+                raw = self._socket.recv(1024)
+            except BlockingIOError:
+                continue
+            except OSError:
+                break
+
+            if len(raw) == 0:  # pragma: no cover
+                break
+
+            buffer.write(raw)
+            buffer.seek(0)
+            start_index = 0
+            for line in buffer:
+                start_index += len(line)
+
+                raw = line.strip()
+                if len(raw) == 0:
+                    continue
+
+                message = BaseExchangeMessage.model_deserialize(raw)
+                self._handle_server_message(message)
+
+            if start_index > 0:
+                buffer.seek(start_index)
+                remaining = buffer.read()
+                buffer.truncate(0)
+                buffer.seek(0)
+                buffer.write(remaining)
+            else:  # pragma: no cover
+                buffer.seek(0, 2)
+        buffer.close()
+
+    def _send_server_message(self, message: ExchangeMessage) -> None:
+        self._socket.send(message.model_serialize() + b'\n')
+        logger.debug('%s sent message to server: %r', self, message)
+
+    def close(self) -> None:
+        """Close the connection to the exchange."""
+        logger.debug('%s closing socket connection to server', self)
+        self._socket.close()
+        self._handler_thread.join(timeout=1)
+
+    def _register_entity(self, uid: Identifier) -> None:
+        self._mailboxes[uid] = SimpleMailbox(uid, self)
+        future: Future[None] = Future()
+        self._pending_registration[uid] = future
+        message = RegisterMessage(src=uid)
+        self._send_server_message(message)
+        future.result(timeout=self.timeout)
+        del self._pending_registration[uid]
+        logger.info('%s registered %r', self, uid)
+
+    def register_agent(self, name: str | None = None) -> AgentIdentifier:
+        """Create a mailbox for a new agent in the system.
+
+        Args:
+            name: Optional human-readable name for the agent.
+        """
+        aid = AgentIdentifier.new(name=name)
+        self._register_entity(aid)
+        return aid
+
+    def register_client(self, name: str | None = None) -> ClientIdentifier:
+        """Create a mailbox for a new client in the system.
+
+        Args:
+            name: Optional human-readable name for the client.
+        """
+        cid = ClientIdentifier.new(name=name)
+        self._register_entity(cid)
+        return cid
+
+    def _unregister_entity(self, uid: Identifier) -> None:
+        future: Future[None] = Future()
+        self._pending_unregistration[uid] = future
+        message = UnregisterMessage(src=uid)
+        self._send_server_message(message)
+        future.result(timeout=self.timeout)
+        logger.info('%s unregistered %r', self, uid)
+
+    def unregister(self, uid: Identifier) -> None:
+        """Unregister the entity (either agent or client).
+
+        Args:
+            uid: Identifier of the entity to unregister.
+        """
+        mailbox = self._mailboxes.pop(uid, None)
+        if mailbox is not None:
+            mailbox.close()
+
+    def create_handle(self, aid: AgentIdentifier) -> Handle:
+        """Create a handle to an agent in the system.
+
+        A handle enables a client to invoke actions on the agent.
+
+        Note:
+            It is not possible to create a handle to a client since a handle
+            is essentially a new client of a specific agent.
+
+        Args:
+            aid: Identifier of agent in the system to create a handle to.
+
+        Returns:
+            Handle to the agent.
+
+        Raises:
+            TypeError: if `aid` is not an instance of
+                [`AgentIdentifier`][aeris.identifier.AgentIdentifier].
+        """
+        if not isinstance(aid, AgentIdentifier):
+            raise TypeError(
+                f'Handle must be created from an {AgentIdentifier.__name__} '
+                f'but got identifier with type {type(aid).__name__}.',
+            )
+        return Handle(aid, self)
+
+    def get_mailbox(self, uid: Identifier) -> SimpleMailbox:
+        """Get the mailbox for an entity in the system.
+
+        Args:
+            uid: Identifier of entity in the system.
+
+        Returns:
+            Mailbox for the entity.
+
+        Raises:
+            BadIdentifierError: if an entity with `uid` is not
+                registered with the exchange.
+        """
+        try:
+            return self._mailboxes[uid]
+        except KeyError as e:
+            raise BadIdentifierError(
+                f'{uid} is not registered with this exchange.',
+            ) from e
+
+
+class _AsyncQueue(Generic[T]):
+    def __init__(self) -> None:
+        self._queue: asyncio.PriorityQueue[_QueueItem[T]] = (
+            asyncio.PriorityQueue()
+        )
+        self._closed = False
+
+    async def close(self, immediate: bool = False) -> None:
+        if not self.closed():
+            self._closed = True
+            priority = CLOSE_PRIORITY if immediate else DEFAULT_PRIORITY
+            await self._queue.put(_QueueItem(priority, CLOSE_SENTINAL))
+
+    def closed(self) -> bool:
+        return self._closed
+
+    async def get(self) -> T:
+        item = await self._queue.get()
+        if item.message is CLOSE_SENTINAL:
+            raise MailboxClosedError
+        return cast(T, item.message)
+
+    async def put(self, message: T) -> None:
+        if self.closed():
+            raise MailboxClosedError
+        await self._queue.put(_QueueItem(DEFAULT_PRIORITY, message))
+
+    async def subscribe(self) -> AsyncGenerator[T]:
+        while True:
+            try:
+                yield await self.get()
+            except MailboxClosedError:
+                return
+
+
+class _MailboxManager:
+    def __init__(self) -> None:
+        self._mailboxes: dict[Identifier, _AsyncQueue[ForwardMessage]] = {}
+
+    def register(self, uid: Identifier) -> None:
+        if uid not in self._mailboxes or self._mailboxes[uid].closed():
+            # If the old mailbox was closed, it gets thrown away.
+            self._mailboxes[uid] = _AsyncQueue()
+        else:
+            raise ExchangeRegistrationError(f'{uid!r} is already registered.')
+
+    async def unregister(self, uid: Identifier) -> None:
+        mailbox = self._mailboxes.pop(uid, None)
+        if mailbox is not None:
+            await mailbox.close(immediate=True)
+        else:
+            raise ExchangeRegistrationError(f'{uid!r} is not registered.')
+
+    async def send(self, message: ForwardMessage) -> None:
+        try:
+            await self._mailboxes[message.dest].put(message)
+        except KeyError as e:
+            raise BadIdentifierError(
+                f'No mailbox associated with {message.dest!r}',
+            ) from e
+
+    async def subscribe(
+        self,
+        uid: Identifier,
+    ) -> AsyncGenerator[ForwardMessage]:
+        try:
+            return self._mailboxes[uid].subscribe()
+        except KeyError as e:
+            raise BadIdentifierError(
+                f'No mailbox associated with {uid}',
+            ) from e
+
+
+class SimpleServer:
+    """Simple asyncio mailbox exchange server.
+
+    Args:
+        host: Host interface to bind to.
+        port: Port to bind to.
+    """
+
+    def __init__(self, host: str, port: int) -> None:
+        self.host = host
+        self.port = port
+        self.manager = _MailboxManager()
+        self._subscriber_tasks: dict[Identifier, asyncio.Task[None]] = {}
+
+    def __repr__(self) -> str:
+        return f'{type(self).__name__}("{self.host}:{self.port}")'
+
+    def __str__(self) -> str:
+        return f'{type(self).__name__}<{self.host}:{self.port}>'
+
+    async def _subscribe(
+        self,
+        uid: Identifier,
+        writer: asyncio.StreamWriter,
+    ) -> None:
+        messages = await self.manager.subscribe(uid)
+        logger.info('%s started subscriber task for %r', self, uid)
+
+        while not writer.is_closing():  # pragma: no branch
+            try:
+                message = await asyncio.wait_for(
+                    messages.__anext__(),
+                    timeout=1,
+                )
+            except asyncio.TimeoutError:  # pragma: no cover
+                continue
+            except StopAsyncIteration:
+                break
+
+            encoded = message.model_serialize()
+            writer.write(encoded)
+            writer.write(b'\n')
+            try:
+                await writer.drain()
+                logger.debug('%s sent message to %r: %r', self, uid, message)
+            except OSError:  # pragma: no cover
+                logger.warning(
+                    '%s failed to send message to %r: %r',
+                    self,
+                    uid,
+                    message,
+                )
+
+        # Note: we don't close the writer here. The happy path is that
+        # the server replies to the client's unregister request and the
+        # client closes the socket once the response is received.
+        logger.info('%s exited subscriber task for %r', self, uid)
+
+    async def _handle_client_message(
+        self,
+        message: ExchangeMessage,
+        writer: asyncio.StreamWriter,
+    ) -> ExchangeResponseMessage:
+        if isinstance(message, ForwardMessage):
+            try:
+                await self.manager.send(message)
+            except BadIdentifierError as e:
+                response = message.response(error=str(e))
+            else:
+                response = message.response()
+        elif isinstance(message, RegisterMessage):
+            try:
+                self.manager.register(message.src)
+            except ExchangeRegistrationError as e:
+                response = message.response(error=str(e))
+            else:
+                task = asyncio.create_task(
+                    self._subscribe(message.src, writer),
+                    name=f'{message.src.uid}-subscriber',
+                )
+                self._subscriber_tasks[message.src] = task
+                logger.info('%s registered client %r', self, message.src)
+                response = message.response()
+        elif isinstance(message, UnregisterMessage):
+            try:
+                await self.manager.unregister(message.src)
+            except ExchangeRegistrationError as e:
+                response = message.response(error=str(e))
+            else:
+                task = self._subscriber_tasks.pop(message.src)
+                await task
+                logger.info('%s unregistered client %r', self, message.src)
+                response = message.response()
+        else:
+            raise AssertionError('Unreachable.')
+
+        return response
+
+    async def _handle_client(
+        self,
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+    ) -> None:
+        logger.debug('%s started new client handle', self)
+        while not reader.at_eof():
+            raw = await reader.readline()
+            if raw == b'':
+                reader.feed_eof()
+                continue
+
+            message = BaseExchangeMessage.model_deserialize(raw)
+            logger.debug('%s received: %r', self, message)
+            response = await self._handle_client_message(message, writer)
+            encoded = response.model_serialize()
+            writer.write(encoded)
+            writer.write(b'\n')
+            await writer.drain()
+
+        writer.close()
+        await writer.wait_closed()
+        logger.info('%s exited client handle', self)
+
+    async def serve_forever(self, stop: asyncio.Future[None]) -> None:
+        """Accept and handles connections forever."""
+        server = await asyncio.start_server(
+            self._handle_client,
+            host=self.host,
+            port=self.port,
+        )
+
+        async with server:
+            await server.start_serving()
+            logger.info(
+                'Server listening on %r:%r (ctrl-C to exit)',
+                self.host,
+                self.port,
+            )
+            await stop
+            logger.info('Closing server...')
+            for task in self._subscriber_tasks.values():
+                task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
+
+        if sys.version_info >= (3, 13):  # pragma: >=3.13 cover
+            server.close_clients()
+
+
+async def _serve_forever(
+    server: SimpleServer,
+    stop: asyncio.Future[None] | None = None,
+) -> None:
+    loop = asyncio.get_running_loop()
+    stop = loop.create_future() if stop is None else stop
+    # Set the stop condition when receiving SIGINT (ctrl-C) and SIGTERM.
+    loop.add_signal_handler(signal.SIGINT, stop.set_result, None)
+    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
+    logger.debug('Registered signal handlers for SIGINT and SIGTERM')
+
+    await server.serve_forever(stop)
+
+    loop.remove_signal_handler(signal.SIGINT)
+    loop.remove_signal_handler(signal.SIGTERM)
+
+
+def _main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--host', default='0.0.0.0')
+    parser.add_argument('--port', type=int, required=True)
+    parser.add_argument('--log-level', default='INFO')
+
+    argv = sys.argv[1:] if argv is None else argv
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(
+        format='[%(asctime)s] %(levelname)-5s (%(name)s) :: %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S',
+        level=args.log_level,
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+
+    server = SimpleServer(host=args.host, port=args.port)
+    asyncio.run(_serve_forever(server))
+
+    return 0
+
+
+if __name__ == '__main__':
+    raise SystemExit(_main())
diff --git a/aeris/exchange/thread.py b/aeris/exchange/thread.py
index 6662ac4..da32748 100644
--- a/aeris/exchange/thread.py
+++ b/aeris/exchange/thread.py
@@ -7,6 +7,7 @@
 import sys
 from collections import defaultdict
 from types import TracebackType
+from typing import get_args
 
 if sys.version_info >= (3, 10):  # pragma: >=3.10 cover
     from typing import TypeAlias
@@ -85,7 +86,7 @@ def recv(self) -> Message:
         message = item.message
         if message is CLOSE_SENTINAL:
             raise MailboxClosedError
-        assert isinstance(message, Message)
+        assert isinstance(message, get_args(Message))
         return message
 
     def close(self) -> None:
diff --git a/aeris/handle.py b/aeris/handle.py
index 9455819..3aacedd 100644
--- a/aeris/handle.py
+++ b/aeris/handle.py
@@ -41,14 +41,14 @@
 logger = logging.getLogger(__name__)
 
 P = ParamSpec('P')
-R_co = TypeVar('R_co', covariant=True)
+R = TypeVar('R')
 
 
 def _validate_state(
-    method: Callable[Concatenate[Handle, P], R_co],
-) -> Callable[Concatenate[Handle, P], R_co]:
+    method: Callable[Concatenate[Handle, P], R],
+) -> Callable[Concatenate[Handle, P], R]:
     @functools.wraps(method)
-    def _wrapper(self: Handle, *args: P.args, **kwargs: P.kwargs) -> R_co:
+    def _wrapper(self: Handle, *args: P.args, **kwargs: P.kwargs) -> R:
         if self._closed:
             raise HandleClosedError()
         return method(self, *args, **kwargs)
@@ -189,7 +189,7 @@ def action(
         /,
         *args: Any,
         **kwargs: Any,
-    ) -> Future[R_co]:
+    ) -> Future[R]:
         """Invoke an action on the agent.
 
         Args:
@@ -207,7 +207,7 @@ def action(
             args=args,
             kwargs=kwargs,
         )
-        future: Future[R_co] = Future()
+        future: Future[R] = Future()
         self._futures[request.mid] = future
         self._agent_mailbox.send(request)
         logger.debug(f'{self} sent {request}')
diff --git a/aeris/identifier.py b/aeris/identifier.py
index cf9b32e..717f04b 100644
--- a/aeris/identifier.py
+++ b/aeris/identifier.py
@@ -1,37 +1,40 @@
 from __future__ import annotations
 
-import abc
-import enum
 import sys
 import uuid
+from typing import Literal
+from typing import Optional
+from typing import Union
 
 if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
     from typing import Self
 else:  # pragma: <3.11 cover
     from typing_extensions import Self
 
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
 
-class Role(enum.Enum):
-    """Roles of entities in a multi-agent system."""
 
-    AGENT = 'agent'
-    CLIENT = 'client'
-
-
-class Identifier(abc.ABC):
+class BaseIdentifier(BaseModel):
     """Unique identifier of an entity in a multi-agent system.
 
     Internally an entity is represented by a UUID. An identify can also
     have an optional name which is only used for human-readability in logging.
 
     Args:
-        uuid: Unique identifier for the entity.
+        uid: Unique identifier for the entity.
         name: Optional human-readable name for the entity.
     """
 
-    def __init__(self, uid: uuid.UUID, *, name: str | None = None) -> None:
-        self._uid = uid
-        self._name = name
+    model_config = ConfigDict(
+        extra='forbid',
+        frozen=True,
+        validate_default=True,
+    )
+
+    uid: uuid.UUID = Field()
+    name: Optional[str] = Field(None)  # noqa: UP007
 
     @classmethod
     def new(cls, name: str | None = None) -> Self:
@@ -40,58 +43,26 @@ def new(cls, name: str | None = None) -> Self:
         Args:
             name: Optional human-readable name for the entity.
         """
-        return cls(uuid.uuid4(), name=name)
-
-    @property
-    def name(self) -> str | None:
-        """Human-readable name of entity."""
-        return self._name
-
-    @property
-    @abc.abstractmethod
-    def role(self) -> Role:
-        """Get the role of this entity."""
-        ...
+        return cls(uid=uuid.uuid4(), name=name)
 
     def __eq__(self, other: object, /) -> bool:
-        if isinstance(other, Identifier):
-            return self._uid == other._uid and self.role == other.role
-        else:
-            return False
+        return isinstance(other, type(self)) and self.uid == other.uid
 
     def __hash__(self) -> int:
-        return hash(self.role) + hash(self._uid)
-
-    def __repr__(self) -> str:
-        if self.name is not None:
-            return (
-                f'{type(self).__name__}(uid={self._uid!s}, name={self.name})'
-            )
-        else:
-            return f'{type(self).__name__}(uid={self._uid!s})'
-
-    def __str__(self) -> str:
-        kind = f'{self.role.name.title()}ID'
-        if self.name is not None:
-            uid_bits = str(self._uid)[:8]
-            return f'{kind}<{uid_bits}.., name={self.name}>'
-        else:
-            return f'{kind}<{self._uid}>'
-
-
-class AgentIdentifier(Identifier):
+        return hash(type(self).__name__) + hash(self.uid)
+
+
+class AgentIdentifier(BaseIdentifier):
     """Unique identifier of an agent in a multi-agent system."""
 
-    @property
-    def role(self) -> Role:
-        """Get the role of this entity."""
-        return Role.AGENT
+    role: Literal['agent'] = Field('agent', repr=False)
 
 
-class ClientIdentifier(Identifier):
+class ClientIdentifier(BaseIdentifier):
     """Unique identifier of a client in a multi-agent system."""
 
-    @property
-    def role(self) -> Role:
-        """Get the role of this entity."""
-        return Role.CLIENT
+    role: Literal['client'] = Field('client', repr=False)
+
+
+Identifier = Union[AgentIdentifier, ClientIdentifier]
+"""Identifier union type for type annotations."""
diff --git a/aeris/message.py b/aeris/message.py
index fd82d94..7e3e04d 100644
--- a/aeris/message.py
+++ b/aeris/message.py
@@ -1,29 +1,34 @@
 from __future__ import annotations
 
-import enum
+import base64
+import pickle
 import uuid
 from typing import Any
+from typing import Literal
 from typing import Optional
+from typing import Union
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import field_serializer
+from pydantic import field_validator
+from pydantic import TypeAdapter
 
+from aeris.identifier import AgentIdentifier
 from aeris.identifier import Identifier
 
 NO_RESULT = object()
 
 
-class MessageKind(enum.Enum):
-    """Message variants."""
+class BaseMessage(BaseModel):
+    """Base message type for messages between entities (agents or clients).
 
-    ACTION = 'action'
-    PING = 'ping'
-    SHUTDOWN = 'shutdown'
-
-
-class Message(BaseModel):
-    """Base exchange message type."""
+    Args:
+        mid: Unique message ID.
+        src: Source entity.
+        dest: Destination entity.
+    """
 
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
@@ -33,80 +38,287 @@ class Message(BaseModel):
         validate_default=True,
     )
 
-    kind: MessageKind
     mid: uuid.UUID = Field(default_factory=uuid.uuid4)
     src: Identifier
     dest: Identifier
 
-    def __str__(self) -> str:
-        name = type(self).__name__
-        return f'{name}<from {self.src}; to {self.dest}; {self.mid}>'
+    @classmethod
+    def model_from_json(cls, data: str) -> Message:
+        """Reconstruct a specific message from a JSON dump.
+
+        Example:
+            ```python
+            from aeris.message import BaseMessage, ActionRequest
+
+            message = ActionRequest(...)
+            dump = message.model_dump_json()
+            assert BaseMessage.model_from_json(dump) == message
+            ```
+        """
+        return TypeAdapter(Message).validate_json(data)
+
 
+class ActionRequest(BaseMessage):
+    """Agent action request message.
 
-class ActionRequest(Message):
-    """Action request message."""
+    When this message is dumped to a JSON string, the `args` and `kwargs`
+    are pickled and then base64 encoded to a string. This can have non-trivial
+    time and space overheads for large arguments.
+
+    Args:
+        action: Name of the action requested.
+        args: Positional arguments for the action method.
+        kwargs: Keyword arguments for the action method.
+    """
 
-    kind: MessageKind = Field(MessageKind.ACTION)
     action: str
     args: tuple[Any, ...] = Field(default_factory=tuple)
     kwargs: dict[str, Any] = Field(default_factory=dict)
+    kind: Literal['action-request'] = Field('action-request', repr=False)
 
-    def response(
-        self,
-        *,
-        exception: Exception | None = None,
-        result: Any = NO_RESULT,
-    ) -> ActionResponse:
-        """Construct response to action request.
+    @field_serializer('args', 'kwargs', when_used='json')
+    def _pickle_and_encode_obj(self, obj: Any) -> str:
+        raw = pickle.dumps(obj)
+        return base64.b64encode(raw).decode('utf-8')
 
-        Args:
-            exception: Exception raised by the action.
-            result: Result of the action.
+    @field_validator('args', 'kwargs', mode='before')
+    @classmethod
+    def _decode_pickled_obj(cls, obj: Any) -> Any:
+        if not isinstance(obj, str):
+            return obj
+        return pickle.loads(base64.b64decode(obj))
 
-        Raises:
-            TypeError: if neither or both of `exception` and `result` are set.
-        """
-        if exception is None and result is NO_RESULT:
-            raise TypeError('One of exception or result must be provided.')
-        elif exception is not None and result is not NO_RESULT:
-            raise TypeError('Cannot provide exception and result.')
+    def error(self, exception: Exception) -> ActionResponse:
+        """Construct an error response to action request.
 
+        Args:
+            exception: Error of the action.
+        """
         return ActionResponse(
             mid=self.mid,
             src=self.dest,
             dest=self.src,
             action=self.action,
+            result=None,
             exception=exception,
-            result=None if result is NO_RESULT else result,
+        )
+
+    def response(self, result: Any) -> ActionResponse:
+        """Construct a success response to action request.
+
+        Args:
+            result: Result of the action.
+        """
+        return ActionResponse(
+            mid=self.mid,
+            src=self.dest,
+            dest=self.src,
+            action=self.action,
+            result=result,
+            exception=None,
         )
 
 
-class ActionResponse(Message):
-    """Action response message."""
+class ActionResponse(BaseMessage):
+    """Agent action response message.
+
+    Args:
+        action: Name of the action requested.
+        result: Result of the action if successful.
+        exception: Exception of the action if unsuccessful.
+    """
 
-    kind: MessageKind = Field(MessageKind.ACTION)
     action: str
-    exception: Optional[Exception] = None  # noqa: UP007
     result: Any = None
+    exception: Optional[Exception] = None  # noqa: UP007
+    kind: Literal['action-response'] = Field('action-response', repr=False)
+
+    @field_serializer('exception', 'result', when_used='json')
+    def _pickle_and_encode_obj(self, obj: Any) -> Optional[tuple[str, str]]:  # noqa: UP007
+        if obj is None:
+            return None
+        raw = pickle.dumps(obj)
+        # This sential value at the start of the tuple is so we can
+        # disambiguate a result that is a str versus the string of a
+        # serialized result.
+        return ('__pickled__', base64.b64encode(raw).decode('utf-8'))
+
+    @field_validator('exception', 'result', mode='before')
+    @classmethod
+    def _decode_pickled_obj(cls, obj: Any) -> Any:
+        if (
+            isinstance(obj, (tuple, list))
+            and len(obj) == 2  # noqa: PLR2004
+            and obj[0] == '__pickled__'
+        ):
+            return pickle.loads(base64.b64decode(obj[1]))
+        return obj
+
+    def __eq__(self, other: object, /) -> bool:
+        if not isinstance(other, ActionResponse):
+            return False
+        return (
+            self.mid == other.mid
+            and self.src == other.src
+            and self.dest == other.dest
+            and self.action == other.action
+            and self.result == other.result
+            # Custom __eq__ is required because exception instances need
+            # to be compared by type. I.e., Exception() == Exception() is
+            # always False.
+            and type(self.exception) is type(other.exception)
+        )
 
 
-class PingRequest(Message):
+class PingRequest(BaseMessage):
     """Ping request message."""
 
-    kind: MessageKind = Field(MessageKind.PING)
+    kind: Literal['ping-request'] = Field('ping-request', repr=False)
 
     def response(self) -> PingResponse:
         """Construct a ping response message."""
-        return PingResponse(mid=self.mid, src=self.dest, dest=self.src)
+        return PingResponse(
+            mid=self.mid,
+            src=self.dest,
+            dest=self.src,
+            exception=None,
+        )
+
+    def error(self, exception: Exception) -> PingResponse:
+        """Construct an error response to ping request.
+
+        Args:
+            exception: Error of the action.
+        """
+        return PingResponse(
+            mid=self.mid,
+            src=self.src,
+            dest=self.dest,
+            exception=exception,
+        )
 
 
-class PingResponse(Message):
+class PingResponse(BaseMessage):
     """Ping response message."""
 
-    kind: MessageKind = Field(MessageKind.PING)
+    exception: Optional[Exception] = None  # noqa: UP007
+    kind: Literal['ping-response'] = Field('ping-response', repr=False)
+
+    @field_serializer('exception', when_used='json')
+    def _pickle_and_encode_obj(self, obj: Any) -> Optional[str]:  # noqa: UP007
+        if obj is None:
+            return None
+        raw = pickle.dumps(obj)
+        return base64.b64encode(raw).decode('utf-8')
+
+    @field_validator('exception', mode='before')
+    @classmethod
+    def _decode_pickled_obj(cls, obj: Any) -> Any:
+        if not isinstance(obj, str):
+            return obj
+        return pickle.loads(base64.b64decode(obj))
+
+    def __eq__(self, other: object, /) -> bool:
+        if not isinstance(other, PingResponse):
+            return False
+        return (
+            self.mid == other.mid
+            and self.src == other.src
+            and self.dest == other.dest
+            # Custom __eq__ is required because exception instances need
+            # to be compared by type. I.e., Exception() == Exception() is
+            # always False.
+            and type(self.exception) is type(other.exception)
+        )
+
+
+class ShutdownRequest(BaseMessage):
+    """Agent shutdown request message."""
+
+    kind: Literal['shutdown-request'] = Field('shutdown-request', repr=False)
+
+    @field_validator('dest', mode='after')
+    @classmethod
+    def _validate_agent(cls, dest: Identifier) -> Identifier:
+        if not isinstance(dest, AgentIdentifier):
+            raise ValueError(
+                'Shutdown requests can only be send to an agent. '
+                f'Destination identifier has the {dest.role} role.',
+            )
+        return dest
+
+    def response(self) -> ShutdownResponse:
+        """Construct a shutdown response message."""
+        return ShutdownResponse(
+            mid=self.mid,
+            src=self.dest,
+            dest=self.src,
+            exception=None,
+        )
 
+    def error(self, exception: Exception) -> ShutdownResponse:
+        """Construct an error response to shutdown request.
+
+        Args:
+            exception: Error of the action.
+        """
+        return ShutdownResponse(
+            mid=self.mid,
+            src=self.src,
+            dest=self.dest,
+            exception=exception,
+        )
+
+
+class ShutdownResponse(BaseMessage):
+    """Agent shutdown response message."""
+
+    exception: Optional[Exception] = None  # noqa: UP007
+    kind: Literal['shutdown-response'] = Field('shutdown-response', repr=False)
+
+    @field_serializer('exception', when_used='json')
+    def _pickle_and_encode_obj(self, obj: Any) -> Optional[str]:  # noqa: UP007
+        if obj is None:
+            return None
+        raw = pickle.dumps(obj)
+        return base64.b64encode(raw).decode('utf-8')
+
+    @field_validator('exception', mode='before')
+    @classmethod
+    def _decode_pickled_obj(cls, obj: Any) -> Any:
+        if not isinstance(obj, str):
+            return obj
+        return pickle.loads(base64.b64decode(obj))
+
+    def __eq__(self, other: object, /) -> bool:
+        if not isinstance(other, ShutdownResponse):
+            return False
+        return (
+            self.mid == other.mid
+            and self.src == other.src
+            and self.dest == other.dest
+            # Custom __eq__ is required because exception instances need
+            # to be compared by type. I.e., Exception() == Exception() is
+            # always False.
+            and type(self.exception) is type(other.exception)
+        )
 
-class ShutdownRequest(Message):
-    """Shutdown request message."""
 
-    kind: MessageKind = Field(MessageKind.SHUTDOWN)
+RequestMessage = Union[ActionRequest, PingRequest, ShutdownRequest]
+ResponseMessage = Union[ActionResponse, PingResponse, ShutdownResponse]
+Message = Union[RequestMessage, ResponseMessage]
+"""Message union type for type annotations.
+
+Tip:
+    This is a parameterized generic type meaning that this type cannot
+    be used for [`isinstance`][`builtins.isinstance`] checks:
+    ```python
+    isinstance(message, Message)  # Fails
+    ```
+    Instead, use [`typing.get_args()`][typing.get_args]:
+    ```
+    from typing import get_args
+
+    isinstance(message, get_args(Message))  # Works
+    ```
+"""
diff --git a/pyproject.toml b/pyproject.toml
index eeca76c..af7b53d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,12 +33,14 @@ Documentation = "https://aeris.proxystore.dev"
 Source = "https://github.com/proxystore/aeris"
 
 [project.optional-dependencies]
+all = ["aeris"]
 dev = [
     "covdefaults>=2.2",
     "coverage",
     "mypy",
     "pre-commit",
     "pytest",
+    "pytest-asyncio",
     "pytest-cov",
     "ruff>=0.2.0",
     "tox",
@@ -87,6 +89,9 @@ module = "tests.*"
 allow_incomplete_defs = true
 allow_untyped_defs = true
 
+[tool.pytest.ini_options]
+asyncio_default_fixture_loop_scope = "function"
+
 [tool.ruff]
 line-length = 79
 target-version = "py39"
@@ -150,7 +155,7 @@ parametrize-values-type = "tuple"
 
 [tool.ruff.lint.flake8-quotes]
 inline-quotes = "single"
-multiline-quotes = "single"
+multiline-quotes = "double"
 
 [tool.ruff.lint.isort]
 force-single-line = true
diff --git a/testing/constant.py b/testing/constant.py
index 4deb8fb..72838a8 100644
--- a/testing/constant.py
+++ b/testing/constant.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
 
-TEST_SLEEP = 0.001
+TEST_CONNECTION_TIMEOUT = 5
 TEST_LOOP_SLEEP = 0.001
+TEST_SLEEP = 0.001
diff --git a/testing/sys.py b/testing/sys.py
new file mode 100644
index 0000000..ac1c6ec
--- /dev/null
+++ b/testing/sys.py
@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import socket
+
+_used_ports: set[int] = set()
+
+
+def open_port() -> int:
+    """Return open port.
+
+    Source: https://stackoverflow.com/questions/2838244
+    """
+    while True:
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        s.bind(('', 0))
+        s.listen(1)
+        port = s.getsockname()[1]
+        s.close()
+        if port not in _used_ports:  # pragma: no branch
+            _used_ports.add(port)
+            return port


Test patch:
diff --git a/tests/exchange/message_test.py b/tests/exchange/message_test.py
new file mode 100644
index 0000000..6dac5f3
--- /dev/null
+++ b/tests/exchange/message_test.py
@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import get_args
+
+import pytest
+
+from aeris.exchange.message import BaseExchangeMessage
+from aeris.exchange.message import ExchangeMessage
+from aeris.exchange.message import ExchangeRequestMessage
+from aeris.exchange.message import ExchangeResponseMessage
+from aeris.exchange.message import ForwardMessage
+from aeris.exchange.message import RegisterMessage
+from aeris.exchange.message import UnregisterMessage
+from aeris.identifier import AgentIdentifier
+from aeris.message import PingRequest
+
+
+@pytest.mark.parametrize(
+    'message',
+    (
+        RegisterMessage(src=AgentIdentifier.new()),
+        UnregisterMessage(src=AgentIdentifier.new()),
+        ForwardMessage(
+            src=AgentIdentifier.new(),
+            dest=AgentIdentifier.new(),
+            message=PingRequest(
+                src=AgentIdentifier.new(),
+                dest=AgentIdentifier.new(),
+            ),
+        ),
+        ExchangeResponseMessage(
+            src=AgentIdentifier.new(),
+            request=RegisterMessage(src=AgentIdentifier.new()),
+        ),
+    ),
+)
+def test_exchange_message_serialize(message: ExchangeMessage) -> None:
+    raw = message.model_serialize()
+    recreated = BaseExchangeMessage.model_deserialize(raw)
+    assert isinstance(recreated, get_args(ExchangeMessage))
+    assert message == recreated
+
+
+@pytest.mark.parametrize(
+    'message',
+    (
+        RegisterMessage(src=AgentIdentifier.new()),
+        UnregisterMessage(src=AgentIdentifier.new()),
+        ForwardMessage(
+            src=AgentIdentifier.new(),
+            dest=AgentIdentifier.new(),
+            message=PingRequest(
+                src=AgentIdentifier.new(),
+                dest=AgentIdentifier.new(),
+            ),
+        ),
+    ),
+)
+def test_create_response(message: ExchangeRequestMessage) -> None:
+    response = message.response(error='error')
+    assert isinstance(response, ExchangeResponseMessage)
+    assert isinstance(response.request, type(message))
+    assert not response.success
diff --git a/tests/exchange/simple_test.py b/tests/exchange/simple_test.py
new file mode 100644
index 0000000..0a2e371
--- /dev/null
+++ b/tests/exchange/simple_test.py
@@ -0,0 +1,347 @@
+from __future__ import annotations
+
+import asyncio
+import pickle
+import socket
+import threading
+import time
+from collections.abc import Generator
+from typing import get_args
+from unittest import mock
+
+import pytest
+
+from aeris.exception import BadIdentifierError
+from aeris.exception import ExchangeRegistrationError
+from aeris.exception import MailboxClosedError
+from aeris.exchange import Exchange
+from aeris.exchange import Mailbox
+from aeris.exchange.message import ExchangeResponseMessage
+from aeris.exchange.message import ForwardMessage
+from aeris.exchange.message import RegisterMessage
+from aeris.exchange.simple import _AsyncQueue
+from aeris.exchange.simple import _MailboxManager
+from aeris.exchange.simple import _main
+from aeris.exchange.simple import _serve_forever
+from aeris.exchange.simple import SimpleExchange
+from aeris.exchange.simple import SimpleServer
+from aeris.identifier import AgentIdentifier
+from aeris.identifier import ClientIdentifier
+from aeris.message import PingRequest
+from aeris.message import PingResponse
+from aeris.message import ResponseMessage
+from testing.constant import TEST_CONNECTION_TIMEOUT
+from testing.constant import TEST_LOOP_SLEEP
+from testing.constant import TEST_SLEEP
+from testing.sys import open_port
+
+
+@pytest.fixture
+def server_thread() -> Generator[tuple[str, int]]:
+    host, port = 'localhost', open_port()
+    server = SimpleServer(host, port)
+    loop = asyncio.new_event_loop()
+    stop = loop.create_future()
+
+    def _target() -> None:
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(server.serve_forever(stop))
+        loop.close()
+
+    handle = threading.Thread(target=_target)
+    handle.start()
+
+    # Wait for server to be listening
+    waited = 0.0
+    while True:
+        try:
+            start = time.perf_counter()
+            with socket.create_connection(
+                (host, port),
+                timeout=TEST_LOOP_SLEEP,
+            ):
+                break
+        except OSError as e:
+            if waited > TEST_CONNECTION_TIMEOUT:  # pragma: no cover
+                raise TimeoutError from e
+            end = time.perf_counter()
+            sleep = max(0, TEST_LOOP_SLEEP - (end - start))
+            time.sleep(sleep)
+            waited += sleep
+
+    yield host, port
+
+    loop.call_soon_threadsafe(stop.set_result, None)
+    handle.join(timeout=TEST_CONNECTION_TIMEOUT)
+    if handle.is_alive():  # pragma: no cover
+        raise TimeoutError(
+            'Server thread did not gracefully exit within '
+            f'{TEST_CONNECTION_TIMEOUT} seconds.',
+        )
+
+
+@pytest.mark.asyncio
+async def test_async_queue() -> None:
+    queue: _AsyncQueue[str] = _AsyncQueue()
+
+    message = 'foo'
+    await queue.put(message)
+    received = await queue.get()
+    assert message == received
+
+    await queue.close()
+    await queue.close()  # Idempotent check
+
+    assert queue.closed()
+    with pytest.raises(MailboxClosedError):
+        await queue.put(message)
+    with pytest.raises(MailboxClosedError):
+        await queue.get()
+
+
+@pytest.mark.asyncio
+async def test_async_queue_subscribe() -> None:
+    queue: _AsyncQueue[int] = _AsyncQueue()
+
+    await queue.put(1)
+    await queue.put(2)
+    await queue.put(3)
+    await queue.close(immediate=False)
+
+    messages = [m async for m in queue.subscribe()]
+    assert set(messages) == {1, 2, 3}
+
+
+@pytest.mark.asyncio
+async def test_mailbox_manager_registration() -> None:
+    manager = _MailboxManager()
+    uid = AgentIdentifier.new()
+    manager.register(uid)
+    await manager.unregister(uid)
+
+
+@pytest.mark.asyncio
+async def test_mailbox_manager_registration_failure() -> None:
+    manager = _MailboxManager()
+    uid = AgentIdentifier.new()
+    manager.register(uid)
+    with pytest.raises(ExchangeRegistrationError):
+        manager.register(uid)
+    await manager.unregister(uid)
+    with pytest.raises(ExchangeRegistrationError):
+        await manager.unregister(uid)
+
+
+@pytest.mark.asyncio
+async def test_mailbox_manager_subscribe() -> None:
+    manager = _MailboxManager()
+    uid = AgentIdentifier.new()
+    request = PingRequest(src=uid, dest=uid)
+    messages = [
+        ForwardMessage(src=uid, dest=uid, message=request),
+        ForwardMessage(src=uid, dest=uid, message=request),
+        ForwardMessage(src=uid, dest=uid, message=request),
+    ]
+
+    manager.register(uid)
+    for message in messages:
+        await manager.send(message)
+
+    received: list[ForwardMessage] = []
+    async for message in await manager.subscribe(uid):  # pragma: no branch
+        received.append(message)
+        if len(received) == len(messages):
+            break
+
+    assert received == messages
+    await manager.unregister(uid)
+
+
+@pytest.mark.asyncio
+async def test_mailbox_manager_bad_identifier() -> None:
+    manager = _MailboxManager()
+    uid = AgentIdentifier.new()
+    request = PingRequest(src=uid, dest=uid)
+    message = ForwardMessage(src=uid, dest=uid, message=request)
+
+    with pytest.raises(BadIdentifierError):
+        await manager.send(message)
+
+    with pytest.raises(BadIdentifierError):
+        await manager.subscribe(uid)
+
+
+def test_mailbox_serve() -> None:
+    with mock.patch('aeris.exchange.simple._serve_forever'):
+        assert _main(['--port', '0']) == 0
+
+
+@pytest.mark.asyncio
+async def test_mailbox_server_serve_forever() -> None:
+    server = SimpleServer('localhost', open_port())
+    assert isinstance(repr(server), str)
+    assert isinstance(str(server), str)
+
+    stop = asyncio.get_running_loop().create_future()
+    task = asyncio.create_task(_serve_forever(server, stop))
+    await asyncio.sleep(TEST_SLEEP)
+    stop.set_result(None)
+    await task
+    task.result()
+
+
+@pytest.mark.asyncio
+async def test_exchange_register_entity(
+    server_thread: tuple[str, int],
+) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        aid = exchange.register_agent()
+        cid = exchange.register_client()
+        exchange.unregister(aid)
+        exchange.unregister(aid)  # Idempotent check
+        exchange.unregister(cid)
+
+
+@pytest.mark.asyncio
+async def test_exchange_register_failure(
+    server_thread: tuple[str, int],
+) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        aid = exchange.register_agent()
+        with pytest.raises(ExchangeRegistrationError):
+            exchange._register_entity(aid)
+        exchange.unregister(aid)
+
+
+@pytest.mark.asyncio
+async def test_exchange_unregister_failure(
+    server_thread: tuple[str, int],
+) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        with pytest.raises(ExchangeRegistrationError):
+            # A client typically won't reach this error because the exchange
+            # unregisters an entity by closing it's mailbox if one exists.
+            exchange._unregister_entity(AgentIdentifier.new())
+
+
+@pytest.mark.asyncio
+async def test_exchange_serialize(server_thread: tuple[str, int]) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange1:
+        pickled = pickle.dumps(exchange1)
+        with pickle.loads(pickled) as exchange2:
+            assert repr(exchange1) == repr(exchange2)
+            assert str(exchange1) == str(exchange2)
+
+
+@pytest.mark.asyncio
+async def test_exchange_bad_server_message(
+    server_thread: tuple[str, int],
+) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        message = RegisterMessage(src=AgentIdentifier.new())
+        # The exchange client should never receive this type of message
+        # from the server, but if it does it should just ignore it.
+        exchange._handle_server_message(message)
+
+
+@pytest.mark.asyncio
+async def test_mailbox_manager_send_response_failure(
+    server_thread: tuple[str, int],
+) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        uid = AgentIdentifier.new()
+        message = ForwardMessage(
+            src=uid,
+            dest=uid,
+            message=PingResponse(src=uid, dest=uid),
+        )
+        response = ExchangeResponseMessage(
+            src=uid,
+            request=message,
+            error='bad request',
+        )
+        # This should just drop the response.
+        exchange._handle_server_message(response)
+
+
+@pytest.mark.asyncio
+async def test_exchange_mailbox_errors(server_thread: tuple[str, int]) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        aid = exchange.register_agent()
+
+        with pytest.raises(BadIdentifierError):
+            exchange.get_mailbox(AgentIdentifier.new())
+
+        mailbox = exchange.get_mailbox(aid)
+        mailbox.close()
+
+        message = PingRequest(src=aid, dest=AgentIdentifier.new())
+
+        with pytest.raises(MailboxClosedError):
+            mailbox._push(message)
+
+        with pytest.raises(MailboxClosedError):
+            mailbox.send(message)
+
+        with pytest.raises(MailboxClosedError):
+            mailbox.recv()
+
+        mailbox.close()
+
+
+@pytest.mark.asyncio
+async def test_exchange_mailbox_send_messages(
+    server_thread: tuple[str, int],
+) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        assert isinstance(exchange, Exchange)
+        aid1 = exchange.register_agent()
+        aid2 = exchange.register_agent()
+
+        with (
+            exchange.get_mailbox(aid1) as mailbox1,
+            exchange.get_mailbox(aid2) as mailbox2,
+        ):
+            assert isinstance(mailbox1, Mailbox)
+            message = PingRequest(src=aid1, dest=aid2)
+            mailbox1.send(message)
+            assert mailbox2.recv() == message
+
+
+@pytest.mark.asyncio
+async def test_exchange_mailbox_send_message_failure(
+    server_thread: tuple[str, int],
+) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        assert isinstance(exchange, Exchange)
+        aid1 = exchange.register_agent()
+        aid2 = AgentIdentifier.new()
+
+        with exchange.get_mailbox(aid1) as mailbox:
+            message = PingRequest(src=aid1, dest=aid2)
+            mailbox.send(message)
+            response = mailbox.recv()
+            assert isinstance(response, get_args(ResponseMessage))
+            assert isinstance(response.exception, BadIdentifierError)
+
+
+@pytest.mark.asyncio
+async def test_exchange_create_handle(server_thread: tuple[str, int]) -> None:
+    host, port = server_thread
+    with SimpleExchange(host, port) as exchange:
+        cid = ClientIdentifier.new()
+        with pytest.raises(TypeError):
+            exchange.create_handle(cid)  # type: ignore[arg-type]
+
+        aid = exchange.register_agent()
+        handle = exchange.create_handle(aid)
+        handle.close()
diff --git a/tests/exchange/thread_test.py b/tests/exchange/thread_test.py
index 3b9ddec..f9a979c 100644
--- a/tests/exchange/thread_test.py
+++ b/tests/exchange/thread_test.py
@@ -9,8 +9,8 @@
 from aeris.exchange import Mailbox
 from aeris.exchange import MailboxClosedError
 from aeris.exchange.thread import ThreadExchange
+from aeris.identifier import AgentIdentifier
 from aeris.identifier import ClientIdentifier
-from aeris.identifier import Role
 from aeris.message import PingRequest
 
 
@@ -22,8 +22,8 @@ def test_protocol() -> None:
         agent_id = exchange.register_agent()
         client_id = exchange.register_client()
 
-        assert agent_id.role == Role.AGENT
-        assert client_id.role == Role.CLIENT
+        assert isinstance(agent_id, AgentIdentifier)
+        assert isinstance(client_id, ClientIdentifier)
 
         mailbox = exchange.get_mailbox(agent_id)
         assert isinstance(mailbox, Mailbox)
diff --git a/tests/identifier_test.py b/tests/identifier_test.py
index caec796..e719173 100644
--- a/tests/identifier_test.py
+++ b/tests/identifier_test.py
@@ -4,39 +4,30 @@
 
 from aeris.identifier import AgentIdentifier
 from aeris.identifier import ClientIdentifier
-from aeris.identifier import Role
 
 
 def tests_identifier_equality() -> None:
     aid = AgentIdentifier.new()
     cid = ClientIdentifier.new()
 
-    assert aid.role == Role.AGENT
-    assert cid.role == Role.CLIENT
+    assert isinstance(aid, AgentIdentifier)
+    assert isinstance(cid, ClientIdentifier)
 
-    assert aid != aid._uid
     assert aid != cid
     assert hash(aid) != hash(cid)
 
-    uuid_ = uuid.uuid4()
-    aid1 = AgentIdentifier(uuid_)
-    aid2 = AgentIdentifier(uuid_)
-    cid1 = ClientIdentifier(uuid_)
+    uid = uuid.uuid4()
+    aid1 = AgentIdentifier(uid=uid)
+    aid2 = AgentIdentifier(uid=uid)
+    cid1 = ClientIdentifier(uid=uid)
 
     assert aid1 == aid2
-    assert repr(aid1) == repr(aid2)
-    assert str(aid1) == str(aid2)
-    assert hash(aid1) == hash(aid2)
     assert aid1 != cid1
-    assert hash(aid1) != hash(cid1)
 
 
 def tests_identifier_equality_ignore_name() -> None:
-    uuid_ = uuid.uuid4()
-    aid1 = AgentIdentifier(uuid_, name='aid1')
-    aid2 = AgentIdentifier(uuid_, name='aid2')
+    uid = uuid.uuid4()
+    aid1 = AgentIdentifier(uid=uid, name='aid1')
+    aid2 = AgentIdentifier(uid=uid, name='aid2')
 
     assert aid1 == aid2
-    assert repr(aid1) != repr(aid2)
-    assert str(aid1) != str(aid2)
-    assert hash(aid1) == hash(aid2)
diff --git a/tests/message_test.py b/tests/message_test.py
index 255323b..a1677c0 100644
--- a/tests/message_test.py
+++ b/tests/message_test.py
@@ -1,48 +1,101 @@
 from __future__ import annotations
 
+from typing import get_args
+
 import pytest
 
 from aeris.identifier import AgentIdentifier
 from aeris.identifier import ClientIdentifier
 from aeris.message import ActionRequest
 from aeris.message import ActionResponse
+from aeris.message import BaseMessage
+from aeris.message import Message
 from aeris.message import PingRequest
 from aeris.message import PingResponse
+from aeris.message import RequestMessage
+from aeris.message import ResponseMessage
+from aeris.message import ShutdownRequest
+from aeris.message import ShutdownResponse
+
 
+def test_shutdown_dest_type() -> None:
+    with pytest.raises(
+        ValueError,
+        match='Destination identifier has the client role.',
+    ):
+        ShutdownRequest(src=AgentIdentifier.new(), dest=ClientIdentifier.new())
 
-def test_message_repr() -> None:
-    request = ActionRequest(
-        src=ClientIdentifier.new(),
-        dest=AgentIdentifier.new(),
-        action='foo',
-    )
-    assert isinstance(repr(request), str)
-    assert isinstance(str(request), str)
 
+_src = AgentIdentifier.new()
+_dest = AgentIdentifier.new()
 
-def test_construct_action_response() -> None:
-    request = ActionRequest(
-        src=ClientIdentifier.new(),
-        dest=AgentIdentifier.new(),
-        action='foo',
-    )
 
-    assert isinstance(request.response(result=42), ActionResponse)
+@pytest.mark.parametrize(
+    'message',
+    (
+        ActionRequest(src=_src, dest=_dest, action='foo', args=(b'bar',)),
+        ActionResponse(src=_src, dest=_dest, action='foo', result=b'bar'),
+        ActionResponse(
+            src=_src,
+            dest=_dest,
+            action='foo',
+            exception=Exception(),
+        ),
+        PingRequest(src=_src, dest=_dest),
+        PingResponse(src=_src, dest=_dest),
+        PingResponse(src=_src, dest=_dest, exception=Exception()),
+        ShutdownRequest(src=_src, dest=_dest),
+        ShutdownResponse(src=_src, dest=_dest),
+        ShutdownResponse(src=_src, dest=_dest, exception=Exception()),
+    ),
+)
+def test_message_representations(message: Message) -> None:
+    assert isinstance(str(message), str)
+    assert isinstance(repr(message), str)
+    jsoned = message.model_dump_json()
+    recreated = BaseMessage.model_from_json(jsoned)
+    assert message == recreated
 
-    with pytest.raises(
-        TypeError,
-        match='One of exception or result must be provided',
-    ):
-        request.response()
 
-    with pytest.raises(TypeError, match='Cannot provide exception and result'):
-        request.response(exception=Exception(), result=42)
+@pytest.mark.parametrize(
+    'request_',
+    (
+        ActionRequest(src=_src, dest=_dest, action='foo', args=(b'bar',)),
+        PingRequest(src=_src, dest=_dest),
+        ShutdownRequest(src=_src, dest=_dest),
+    ),
+)
+def test_create_response_message(request_: RequestMessage) -> None:
+    if isinstance(request_, ActionRequest):
+        assert isinstance(
+            request_.response(result=42),
+            get_args(ResponseMessage),
+        )
+    else:
+        assert isinstance(request_.response(), get_args(ResponseMessage))
 
+    exception = Exception('foo')
+    response = request_.error(exception)
+    assert isinstance(response, get_args(ResponseMessage))
+    assert response.exception == exception
 
-def test_construct_ping_response() -> None:
-    request = PingRequest(
-        src=ClientIdentifier.new(),
-        dest=AgentIdentifier.new(),
-    )
 
-    assert isinstance(request.response(), PingResponse)
+@pytest.mark.parametrize(
+    'response',
+    (
+        ActionResponse(
+            src=_src,
+            dest=_dest,
+            action='foo',
+            exception=Exception(),
+        ),
+        PingResponse(src=_src, dest=_dest, exception=Exception()),
+        ShutdownResponse(src=_src, dest=_dest, exception=Exception()),
+    ),
+)
+def test_action_response_exception_equality(response: ResponseMessage) -> None:
+    dump = response.model_dump()
+    dump['exception'] = Exception()
+    other = type(response).model_validate(dump)
+    assert response == other
+    assert response != BaseMessage(src=response.src, dest=response.dest)