diff --git a/academy/context.py b/academy/context.py index 0536491c..5e97dcd2 100644 --- a/academy/context.py +++ b/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. diff --git a/academy/middleware.py b/academy/middleware.py new file mode 100644 index 00000000..69d245d2 --- /dev/null +++ b/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 diff --git a/academy/runtime.py b/academy/runtime.py index 744b0f28..16422865 100644 --- a/academy/runtime.py +++ b/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. diff --git a/tests/unit/context_test.py b/tests/unit/context_test.py index 08db167e..2ed60e81 100644 --- a/tests/unit/context_test.py +++ b/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() diff --git a/tests/unit/runtime_test.py b/tests/unit/runtime_test.py index 7bb2b102..48f4eb16 100644 --- a/tests/unit/runtime_test.py +++ b/tests/unit/runtime_test.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import pickle import sys import uuid from collections.abc import AsyncGenerator @@ -36,6 +37,8 @@ from academy.message import ShutdownRequest from academy.message import SuccessResponse from academy.message import UserErrorResponse +from academy.middleware import ActionCall +from academy.middleware import ActionHandler from academy.runtime import Runtime from academy.runtime import RuntimeConfig from academy.serialize import allowed_deserializers @@ -725,7 +728,13 @@ async def test_runtime_agent_self_termination( exchange_factory=exchange_client.factory(), registration=registration, ) as runtime: - await runtime.action('end', AgentId.new(), args=(), kwargs={}) + await runtime.action( + 'end', + AgentId.new(), + args=(), + kwargs={}, + tag=uuid.uuid4(), + ) await runtime.wait_shutdown(timeout=TEST_WAIT_TIMEOUT) @@ -756,6 +765,7 @@ async def test_runtime_agent_action_context( exchange_client.client_id, args=(exchange_client.client_id,), kwargs={}, + tag=uuid.uuid4(), ) @@ -932,3 +942,205 @@ async def test_runtime_uses_exception_serialization( assert isinstance(body, UserErrorResponse) assert body.serialization == SerializationStrategy.JSON + + +async def _passthrough_middleware( + call: ActionCall, + next_handler: ActionHandler, +) -> Any: + return await next_handler(call) + + +@pytest.mark.asyncio +async def test_runtime_action_middleware_order_and_result( + exchange_client: UserExchangeClient[LocalExchangeTransport], +) -> None: + registration = await exchange_client.register_agent(CounterAgent) + events: list[str] = [] + + async def outer(call: ActionCall, next_handler: ActionHandler) -> Any: + events.append(f'outer-before-{call.action}') + result = await next_handler(call) + events.append(f'outer-after-{call.action}') + return result + + async def inner(call: ActionCall, next_handler: ActionHandler) -> Any: + events.append(f'inner-before-{call.action}') + result = await next_handler(call) + events.append(f'inner-after-{call.action}') + return -1 if call.action == 'count' else result + + config = RuntimeConfig(action_middleware=(outer, inner)) + async with Runtime( + CounterAgent(), + config=config, + exchange_factory=exchange_client.factory(), + registration=registration, + ) as runtime: + await runtime.action( + 'add', + exchange_client.client_id, + args=(2,), + kwargs={}, + tag=uuid.uuid4(), + ) + result = await runtime.action( + 'count', + exchange_client.client_id, + args=(), + kwargs={}, + tag=uuid.uuid4(), + ) + + assert result == -1 + assert events == [ + 'outer-before-add', + 'inner-before-add', + 'inner-after-add', + 'outer-after-add', + 'outer-before-count', + 'inner-before-count', + 'inner-after-count', + 'outer-after-count', + ] + + +@pytest.mark.asyncio +async def test_runtime_action_middleware_denies_action( + exchange_client: UserExchangeClient[LocalExchangeTransport], +) -> None: + registration = await exchange_client.register_agent(CounterAgent) + # Cancel listener so test can intercept agent responses + await exchange_client._stop_listener_task() + listener = exchange_client._transport.listen(TEST_SLEEP_INTERVAL) + + async def deny(call: ActionCall, next_handler: ActionHandler) -> Any: + raise PermissionError(f'Denied action "{call.action}".') + + config = RuntimeConfig(action_middleware=(deny,)) + async with Runtime( + CounterAgent(), + config=config, + exchange_factory=exchange_client.factory(), + registration=registration, + ) as runtime: + with pytest.raises(PermissionError, match='Denied action'): + await runtime.action( + 'count', + exchange_client.client_id, + args=(), + kwargs={}, + tag=uuid.uuid4(), + ) + + request = Message.create( + src=exchange_client.client_id, + dest=runtime.agent_id, + body=ActionRequest( + action='count', + serialization=SerializationStrategy.PICKLE, + ), + ) + await exchange_client.send(request) + message = await anext(listener) + body = message.get_body() + assert isinstance(body, ErrorResponse) + exception = body.get_exception() + assert isinstance(exception, PermissionError) + assert 'Denied action "count".' in str(exception) + + +@pytest.mark.asyncio +async def test_runtime_action_middleware_receives_call_metadata( + exchange_client: UserExchangeClient[LocalExchangeTransport], +) -> None: + registration = await exchange_client.register_agent(CounterAgent) + # Cancel listener so test can intercept agent responses + await exchange_client._stop_listener_task() + listener = exchange_client._transport.listen(TEST_SLEEP_INTERVAL) + seen: list[ActionCall] = [] + + async def record(call: ActionCall, next_handler: ActionHandler) -> Any: + seen.append(call) + return await next_handler(call) + + config = RuntimeConfig(action_middleware=(record,)) + async with Runtime( + CounterAgent(), + config=config, + exchange_factory=exchange_client.factory(), + registration=registration, + ) as runtime: + request = Message.create( + src=exchange_client.client_id, + dest=runtime.agent_id, + body=ActionRequest( + action='add', + pargs=(3,), + serialization=SerializationStrategy.PICKLE, + ), + ) + await exchange_client.send(request) + message = await anext(listener) + assert isinstance(message.get_body(), ActionResponse) + + assert len(seen) == 1 + call = seen[0] + assert call.action == 'add' + assert call.action_context.source_id == exchange_client.client_id + assert call.action_context.tag == request.tag + assert call.agent_context.agent_id == registration.agent_id + assert call.args == (3,) + assert call.kwargs == {} + + +@pytest.mark.asyncio +async def test_runtime_action_middleware_skips_unknown_action( + exchange_client: UserExchangeClient[LocalExchangeTransport], +) -> None: + registration = await exchange_client.register_agent(EmptyAgent) + + async def sentinel(call: ActionCall, next_handler: ActionHandler) -> Any: + raise AssertionError( # pragma: no cover + 'Middleware must not be invoked for unknown actions.', + ) + + config = RuntimeConfig(action_middleware=(sentinel,)) + async with Runtime( + EmptyAgent(), + config=config, + exchange_factory=exchange_client.factory(), + registration=registration, + ) as runtime: + with pytest.raises(AttributeError, match='does not have an action'): + await runtime.action( + 'missing', + exchange_client.client_id, + args=(), + kwargs={}, + tag=uuid.uuid4(), + ) + + +@pytest.mark.asyncio +async def test_runtime_config_action_middleware_picklable() -> None: + config = RuntimeConfig(action_middleware=(_passthrough_middleware,)) + recreated = pickle.loads(pickle.dumps(config)) + assert recreated.action_middleware == (_passthrough_middleware,) + + async def next_handler(call: ActionCall) -> Any: + return call.action + + call = ActionCall( + action='noop', + agent_context=mock.MagicMock(), + action_context=ActionContext( + AgentId.new(), + uuid.uuid4(), + mock.MagicMock(), + ), + args=(), + kwargs={}, + ) + middleware = recreated.action_middleware[0] + assert await middleware(call, next_handler) == 'noop'