From 569f673dd30dc7639f2f019482f99a0ddfa4a1de Mon Sep 17 00:00:00 2001 From: Ian Date: Thu, 1 Oct 2026 10:09:32 -0500 Subject: [PATCH 1/3] Add action middleware hook to RuntimeConfig --- academy/runtime.py | 101 +++++++++++++++++++- tests/unit/runtime_test.py | 186 +++++++++++++++++++++++++++++++++++++ 2 files changed, 282 insertions(+), 5 deletions(-) diff --git a/academy/runtime.py b/academy/runtime.py index 744b0f28..df087560 100644 --- a/academy/runtime.py +++ b/academy/runtime.py @@ -34,6 +34,7 @@ from academy.exchange.client import exchange_context from academy.exchange.transport import AgentRegistrationT from academy.exchange.transport import ExchangeTransportT +from academy.identifier import AgentId from academy.identifier import EntityId from academy.message import AcademyErrorResponse from academy.message import ActionRequest @@ -72,10 +73,69 @@ class _ShutdownState: terminate_override: bool | None = None +@dataclasses.dataclass(frozen=True) +class ActionCall: + """Description of one action invocation dispatched by a runtime. + + Instances are passed to each + [`ActionMiddleware`][academy.runtime.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 + source_id: EntityId + agent_id: AgentId[Any] + args: tuple[Any, ...] + kwargs: dict[str, Any] + tag: uuid.UUID | None = None + + +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.runtime.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 + + 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 +158,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 +223,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 +335,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 +505,30 @@ async def action( *, args: Any, kwargs: Any, + tag: uuid.UUID | None = None, ) -> Any: """Invoke an action of the agent's agent. + The invocation passes through any + [`ActionMiddleware`][academy.runtime.ActionMiddleware] configured + on the [`RuntimeConfig`][academy.runtime.RuntimeConfig] before the + action method is called. + 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, if the action was initiated by a message. 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. + Exception: Any exception raised by a configured middleware. """ logger.debug( 'Invoking "%s" action on %s', @@ -467,13 +543,28 @@ async def action( raise AttributeError( f'{self.agent} does not have an action named "{action}".', ) - action_method = self._actions[action] + call = ActionCall( + action=action, + source_id=source_id, + agent_id=self.agent_id, + args=tuple(args), + kwargs=kwargs, + tag=tag, + ) + 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) + context = ActionContext(call.source_id, self._exchange_client) + return await action_method( + *call.args, + context=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/runtime_test.py b/tests/unit/runtime_test.py index 7bb2b102..91c150c5 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.runtime import ActionCall +from academy.runtime import ActionHandler from academy.runtime import Runtime from academy.runtime import RuntimeConfig from academy.serialize import allowed_deserializers @@ -932,3 +935,186 @@ 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={}, + ) + result = await runtime.action( + 'count', + exchange_client.client_id, + args=(), + kwargs={}, + ) + + 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={}, + ) + + 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.source_id == exchange_client.client_id + assert call.agent_id == registration.agent_id + assert call.args == (3,) + assert call.kwargs == {} + assert call.tag == request.tag + + +@pytest.mark.asyncio +async def test_runtime_action_middleware_skips_unknown_action( + exchange_client: UserExchangeClient[LocalExchangeTransport], +) -> None: + registration = await exchange_client.register_agent(EmptyAgent) + invoked = False + + async def record(call: ActionCall, next_handler: ActionHandler) -> Any: + nonlocal invoked + invoked = True + return await next_handler(call) + + config = RuntimeConfig(action_middleware=(record,)) + 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={}, + ) + + assert not invoked + + +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,) From fa4be6761c3bee340059650f112e4f119d328259 Mon Sep 17 00:00:00 2001 From: Ian Date: Thu, 1 Oct 2026 10:41:24 -0500 Subject: [PATCH 2/3] Close coverage gaps in middleware tests --- tests/unit/runtime_test.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/tests/unit/runtime_test.py b/tests/unit/runtime_test.py index 91c150c5..24b399cb 100644 --- a/tests/unit/runtime_test.py +++ b/tests/unit/runtime_test.py @@ -1089,14 +1089,13 @@ async def test_runtime_action_middleware_skips_unknown_action( exchange_client: UserExchangeClient[LocalExchangeTransport], ) -> None: registration = await exchange_client.register_agent(EmptyAgent) - invoked = False - async def record(call: ActionCall, next_handler: ActionHandler) -> Any: - nonlocal invoked - invoked = True - return await next_handler(call) + 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=(record,)) + config = RuntimeConfig(action_middleware=(sentinel,)) async with Runtime( EmptyAgent(), config=config, @@ -1111,10 +1110,22 @@ async def record(call: ActionCall, next_handler: ActionHandler) -> Any: kwargs={}, ) - assert not invoked - -def test_runtime_config_action_middleware_picklable() -> None: +@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', + source_id=AgentId.new(), + agent_id=AgentId.new(), + args=(), + kwargs={}, + ) + middleware = recreated.action_middleware[0] + assert await middleware(call, next_handler) == 'noop' From 823ad11aaa66889fe469814b28439cdedaaaf019 Mon Sep 17 00:00:00 2001 From: Alok Kamatar Date: Fri, 9 Oct 2026 11:51:56 -0500 Subject: [PATCH 3/3] Refactoring middleware into a new file, and to use context strucutres --- academy/context.py | 8 ++++ academy/middleware.py | 61 +++++++++++++++++++++++++++ academy/runtime.py | 86 ++++++++++---------------------------- tests/unit/context_test.py | 9 +++- tests/unit/runtime_test.py | 31 ++++++++++---- 5 files changed, 121 insertions(+), 74 deletions(-) create mode 100644 academy/middleware.py 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 df087560..16422865 100644 --- a/academy/runtime.py +++ b/academy/runtime.py @@ -34,7 +34,6 @@ from academy.exchange.client import exchange_context from academy.exchange.transport import AgentRegistrationT from academy.exchange.transport import ExchangeTransportT -from academy.identifier import AgentId from academy.identifier import EntityId from academy.message import AcademyErrorResponse from academy.message import ActionRequest @@ -49,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 @@ -73,59 +76,6 @@ class _ShutdownState: terminate_override: bool | None = None -@dataclasses.dataclass(frozen=True) -class ActionCall: - """Description of one action invocation dispatched by a runtime. - - Instances are passed to each - [`ActionMiddleware`][academy.runtime.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 - source_id: EntityId - agent_id: AgentId[Any] - args: tuple[Any, ...] - kwargs: dict[str, Any] - tag: uuid.UUID | None = None - - -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.runtime.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 - - class RuntimeConfig(BaseModel): """Agent runtime configuration. @@ -505,14 +455,15 @@ async def action( *, args: Any, kwargs: Any, - tag: uuid.UUID | None = None, + tag: uuid.UUID, ) -> Any: """Invoke an action of the agent's agent. The invocation passes through any - [`ActionMiddleware`][academy.runtime.ActionMiddleware] configured + [`ActionMiddleware`][academy.middleware.ActionMiddleware] configured on the [`RuntimeConfig`][academy.runtime.RuntimeConfig] before the - action method is called. + 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. @@ -520,7 +471,7 @@ async def 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. + invocation. Returns: Result of the action. @@ -528,6 +479,8 @@ async def action( Raises: AttributeError: If an action with this name is not implemented by 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( @@ -543,24 +496,29 @@ async def action( raise AttributeError( f'{self.agent} does not have an action named "{action}".', ) + + assert self._exchange_client is not None + context = ActionContext( + source_id, + tag, + self._exchange_client, + ) + call = ActionCall( action=action, - source_id=source_id, - agent_id=self.agent_id, + agent_context=self.agent.agent_context, + action_context=context, args=tuple(args), kwargs=kwargs, - tag=tag, ) 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(call.source_id, self._exchange_client) return await action_method( *call.args, - context=context, + context=call.action_context, **call.kwargs, ) else: 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 24b399cb..48f4eb16 100644 --- a/tests/unit/runtime_test.py +++ b/tests/unit/runtime_test.py @@ -37,8 +37,8 @@ from academy.message import ShutdownRequest from academy.message import SuccessResponse from academy.message import UserErrorResponse -from academy.runtime import ActionCall -from academy.runtime import ActionHandler +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 @@ -728,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) @@ -759,6 +765,7 @@ async def test_runtime_agent_action_context( exchange_client.client_id, args=(exchange_client.client_id,), kwargs={}, + tag=uuid.uuid4(), ) @@ -975,12 +982,14 @@ async def inner(call: ActionCall, next_handler: ActionHandler) -> Any: 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 @@ -1021,6 +1030,7 @@ async def deny(call: ActionCall, next_handler: ActionHandler) -> Any: exchange_client.client_id, args=(), kwargs={}, + tag=uuid.uuid4(), ) request = Message.create( @@ -1077,11 +1087,11 @@ async def record(call: ActionCall, next_handler: ActionHandler) -> Any: assert len(seen) == 1 call = seen[0] assert call.action == 'add' - assert call.source_id == exchange_client.client_id - assert call.agent_id == registration.agent_id + 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 == {} - assert call.tag == request.tag @pytest.mark.asyncio @@ -1108,6 +1118,7 @@ async def sentinel(call: ActionCall, next_handler: ActionHandler) -> Any: exchange_client.client_id, args=(), kwargs={}, + tag=uuid.uuid4(), ) @@ -1122,8 +1133,12 @@ async def next_handler(call: ActionCall) -> Any: call = ActionCall( action='noop', - source_id=AgentId.new(), - agent_id=AgentId.new(), + agent_context=mock.MagicMock(), + action_context=ActionContext( + AgentId.new(), + uuid.uuid4(), + mock.MagicMock(), + ), args=(), kwargs={}, )