Add action middleware hook to RuntimeConfig

academy/context.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import dataclasses
+import uuid
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 from typing import Generic
@@ -25,9 +26,11 @@ class ActionContext:
     def __init__(
         self,
         source_id: EntityId,
+        tag: uuid.UUID,
         exchange_client: ae.AgentExchangeClient[Any, Any],
     ) -> None:
         self._source_id = source_id
+        self._tag = tag
         self._exchange_client = exchange_client
         self._source_handle: Handle[Any] | None = None
 
@@ -36,6 +39,11 @@ def source_id(self) -> EntityId:
         """ID of the source that requested the action."""
         return self._source_id
 
+    @property
+    def tag(self) -> uuid.UUID:
+        """Unique tag of the request message invoking the action."""
+        return self._tag
+
     @property
     def source_handle(self) -> Handle[Any]:
         """Get a handle to the source agent of the request.

academy/middleware.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Awaitable
+from collections.abc import Callable
+from typing import Any
+
+from academy.context import ActionContext
+from academy.context import AgentContext
+
+
+@dataclasses.dataclass(frozen=True)
+class ActionCall:
+    """Description of one action invocation dispatched by a runtime.
+
+    Instances are passed to each
+    [`ActionMiddleware`][academy.middleware.ActionMiddleware] configured on
+    the [`RuntimeConfig`][academy.runtime.RuntimeConfig].
+
+    Attributes:
+        action: Name of the action being invoked.
+        source_id: ID of the entity that requested the action.
+        agent_id: ID of the agent executing the action.
+        args: Tuple of positional arguments.
+        kwargs: Dictionary of keyword arguments.
+        tag: Unique tag of the request message that initiated this
+            invocation, if the action was initiated by a message.
+    """
+
+    action: str
+    agent_context: AgentContext[Any]
+    action_context: ActionContext
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+
+
+ActionHandler = Callable[[ActionCall], Awaitable[Any]]
+"""Signature of the next handler passed to an action middleware."""
+
+ActionMiddleware = Callable[[ActionCall, ActionHandler], Awaitable[Any]]
+"""Signature of an action middleware.
+
+A middleware is an async callable invoked with the
+[`ActionCall`][academy.middleware.ActionCall] and the next handler in the
+chain. A middleware may inspect or validate the call before awaiting the
+next handler, transform or record the result after, short-circuit by
+returning without awaiting the next handler, or deny the call by raising
+an exception. Exceptions raised by middleware propagate exactly like
+exceptions raised by the action itself (i.e., returned to a remote caller
+as a [`UserErrorResponse`][academy.message.UserErrorResponse]).
+"""
+
+
+def _chain_middleware(
+    middleware: ActionMiddleware,
+    next_handler: ActionHandler,
+) -> ActionHandler:
+    async def handler(call: ActionCall) -> Any:
+        return await middleware(call, next_handler)
+
+    return handler

academy/runtime.py

@@ -48,6 +48,10 @@
 from academy.message import ShutdownRequest
 from academy.message import SuccessResponse
 from academy.message import UserErrorResponse
+from academy.middleware import _chain_middleware
+from academy.middleware import ActionCall
+from academy.middleware import ActionHandler
+from academy.middleware import ActionMiddleware
 from academy.serialize import allowed_deserializers
 from academy.serialize import default_serializer
 from academy.serialize import NoPickleMixin
@@ -76,6 +80,12 @@ class RuntimeConfig(BaseModel):
     """Agent runtime configuration.
 
     Attributes:
+        action_middleware: Ordered middleware applied around every action
+            dispatch. The first middleware is outermost: it is entered first
+            and returns last. Middleware must be picklable (such as
+            module-level functions or instances of importable classes
+            defining `async def __call__`) when agents are launched on
+            remote executors.
         cancel_actions_on_shutdown: Cancel running actions when the agent
             is shutdown, otherwise wait for the actions to finish.
         max_sync_concurrency: Maximum number of concurrent sync tasks allowed
@@ -98,6 +108,7 @@ class RuntimeConfig(BaseModel):
 
     model_config = ConfigDict(extra='forbid')
 
+    action_middleware: tuple[ActionMiddleware, ...] = ()
     cancel_actions_on_shutdown: bool = True
     max_sync_concurrency: int | None = None
     raise_loop_errors_on_shutdown: bool = True
@@ -162,6 +173,11 @@ def __init__(
         self._actions = agent._agent_actions()
         self._loops = agent._agent_loops()
 
+        handler: ActionHandler = self._invoke_action
+        for middleware in reversed(self.config.action_middleware):
+            handler = _chain_middleware(middleware, handler)
+        self._action_handler = handler
+
         self._started_event = asyncio.Event()
         self._shutdown_event = asyncio.Event()
         self._shutdown_options = _ShutdownState()
@@ -269,6 +285,7 @@ async def _execute_action(self, request: Message[ActionRequest]) -> None:
                 request.src,
                 args=body.get_args(),
                 kwargs=body.get_kwargs(),
+                tag=invocation_id,
             )
 
             # Keep response in try/except so serialization errors are caught
@@ -438,21 +455,33 @@ async def action(
         *,
         args: Any,
         kwargs: Any,
+        tag: uuid.UUID,
     ) -> Any:
         """Invoke an action of the agent's agent.
 
+        The invocation passes through any
+        [`ActionMiddleware`][academy.middleware.ActionMiddleware] configured
+        on the [`RuntimeConfig`][academy.runtime.RuntimeConfig] before the
+        action method is called. This call requires the `runtime` to be
+        started. Otherwise, it will raise a `AgentNotInitializedError`.
+
         Args:
             action: Name of action to invoke.
             source_id: ID of the source that requested the action.
             args: Tuple of positional arguments.
             kwargs: Dictionary of keyword arguments.
+            tag: Unique tag of the request message that initiated this
+                invocation.
 
         Returns:
             Result of the action.
 
         Raises:
             AttributeError: If an action with this name is not implemented by
-                the agent's agent.
+                the agent's agent. Middleware is not invoked in this case.
+            AgentNotInitializedError: If this was called when the runtime
+                has not been started.
+            Exception: Any exception raised by a configured middleware.
         """
         logger.debug(
             'Invoking "%s" action on %s',
@@ -467,13 +496,33 @@ async def action(
             raise AttributeError(
                 f'{self.agent} does not have an action named "{action}".',
             )
-        action_method = self._actions[action]
+
+        assert self._exchange_client is not None
+        context = ActionContext(
+            source_id,
+            tag,
+            self._exchange_client,
+        )
+
+        call = ActionCall(
+            action=action,
+            agent_context=self.agent.agent_context,
+            action_context=context,
+            args=tuple(args),
+            kwargs=kwargs,
+        )
+        return await self._action_handler(call)
+
+    async def _invoke_action(self, call: ActionCall) -> Any:
+        action_method = self._actions[call.action]
         if action_method._action_method_context:
-            assert self._exchange_client is not None
-            context = ActionContext(source_id, self._exchange_client)
-            return await action_method(*args, context=context, **kwargs)
+            return await action_method(
+                *call.args,
+                context=call.action_context,
+                **call.kwargs,
+            )
         else:
-            return await action_method(*args, **kwargs)
+            return await action_method(*call.args, **call.kwargs)
 
     async def run_until_complete(self) -> None:
         """Run the agent until shutdown.

tests/unit/context_test.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import uuid
 from typing import Any
 
 import pytest
@@ -18,6 +19,7 @@ async def test_action_context_agent_source(
 ) -> None:
     factory = exchange_client.factory()
     source_id: AgentId[EmptyAgent] = AgentId.new()
+    tag = uuid.uuid4()
     registration = await exchange_client.register_agent(EmptyAgent)
 
     async def _request_handler(_: Any) -> None:  # pragma: no cover
@@ -27,13 +29,15 @@ async def _request_handler(_: Any) -> None:  # pragma: no cover
         registration,
         _request_handler,
     ) as agent_client:
-        context = ActionContext(source_id, agent_client)
+        context = ActionContext(source_id, tag, agent_client)
         assert context.is_agent_source()
         assert not context.is_user_source()
         handle = context.source_handle
         assert handle.agent_id == source_id
         # Check that handle.source_handle is cached
         assert context.source_handle is handle
+        # Check tag is set
+        assert context.tag == tag
 
 
 @pytest.mark.asyncio
@@ -42,6 +46,7 @@ async def test_action_context_user_source(
 ) -> None:
     factory = exchange_client.factory()
     source_id = UserId.new()
+    tag = uuid.uuid4()
     registration = await exchange_client.register_agent(EmptyAgent)
 
     async def _request_handler(_: Any) -> None:  # pragma: no cover
@@ -51,7 +56,7 @@ async def _request_handler(_: Any) -> None:  # pragma: no cover
         registration,
         _request_handler,
     ) as agent_client:
-        context = ActionContext(source_id, agent_client)
+        context = ActionContext(source_id, tag, agent_client)
         assert not context.is_agent_source()
         assert context.is_user_source()