From 1bee606f23ee225423d8708e2a46fdf53ee5eab8 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Sun, 26 Jul 2026 11:20:46 -0500 Subject: [PATCH 1/2] fix(python): handle callable class middleware safely in _determine_middleware_type (#6697) --- .../core/agent_framework/_middleware.py | 10 ++-- .../tests/core/test_middleware_with_agent.py | 47 +++++++++++++++++++ 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index c82b1eab69e..c72e1a7277f 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -1445,7 +1445,7 @@ def _middleware_handler( def _determine_middleware_type(middleware: Any) -> MiddlewareType: - """Determine middleware type using decorator and/or parameter type annotation. + """Determine the middleware type from function annotations or decorators. Args: middleware: The middleware function to analyze. @@ -1456,6 +1456,8 @@ def _determine_middleware_type(middleware: Any) -> MiddlewareType: Raises: MiddlewareException: When middleware type cannot be determined or there's a mismatch. """ + middleware_name = getattr(middleware, "__name__", type(middleware).__name__) + # Check for decorator marker decorator_type: MiddlewareType | None = getattr(middleware, "_middleware_type", None) @@ -1480,7 +1482,7 @@ def _determine_middleware_type(middleware: Any) -> MiddlewareType: # Not enough parameters - can't be valid middleware raise MiddlewareException( f"Middleware function must have at least 2 parameters (context, call_next), " - f"but {middleware.__name__} has {len(params)}" + f"but {middleware_name} has {len(params)}" ) except Exception as e: if isinstance(e, MiddlewareException): @@ -1493,7 +1495,7 @@ def _determine_middleware_type(middleware: Any) -> MiddlewareType: if decorator_type != param_type: raise MiddlewareException( f"MiddlewareTypes type mismatch: decorator indicates '{decorator_type.value}' " - f"but parameter type indicates '{param_type.value}' for function {middleware.__name__}" + f"but parameter type indicates '{param_type.value}' for function {middleware_name}" ) return decorator_type @@ -1507,7 +1509,7 @@ def _determine_middleware_type(middleware: Any) -> MiddlewareType: # Neither decorator nor parameter type specified - throw exception raise MiddlewareException( - f"Cannot determine middleware type for function {middleware.__name__}. " + f"Cannot determine middleware type for function {middleware_name}. " f"Please either use @agent_middleware/@function_middleware/@chat_middleware decorators " f"or specify parameter types (AgentContext, FunctionInvocationContext, or ChatContext)." ) diff --git a/python/packages/core/tests/core/test_middleware_with_agent.py b/python/packages/core/tests/core/test_middleware_with_agent.py index fc9bbe5488b..c88193345e1 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -2302,3 +2302,50 @@ async def kwargs_middleware(context: AgentContext, call_next: Callable[[], Await # response = await agent.run("test message") # assert response is not None # assert execution_order == ["before", "after"] + + +class TestCallableClassMiddlewareErrorHandling: + """Tests for exception handling when using callable class instances as middleware.""" + + def test_callable_class_middleware_insufficient_params_raises_middleware_exception(self) -> None: + """Test that callable class instance with insufficient params raises MiddlewareException.""" + + class InsufficientParamsMiddleware: + async def __call__(self, ctx: Any) -> None: + pass + + client = MockBaseChatClient() + with pytest.raises(MiddlewareException) as exc_info: + Agent(client=client, middleware=[InsufficientParamsMiddleware()]) + + assert "InsufficientParamsMiddleware" in str(exc_info.value) + assert "must have at least 2 parameters" in str(exc_info.value) + + def test_callable_class_middleware_type_mismatch_raises_middleware_exception(self) -> None: + """Test that callable class instance with decorator/annotation mismatch raises MiddlewareException.""" + + @agent_middleware + class MismatchedCallableMiddleware: + async def __call__(self, context: FunctionInvocationContext, call_next: Any) -> None: + await call_next() + + client = MockBaseChatClient() + with pytest.raises(MiddlewareException) as exc_info: + Agent(client=client, middleware=[MismatchedCallableMiddleware()]) + + assert "MismatchedCallableMiddleware" in str(exc_info.value) + assert "MiddlewareTypes type mismatch" in str(exc_info.value) + + def test_callable_class_middleware_undetermined_type_raises_middleware_exception(self) -> None: + """Test that a callable class instance without annotations or decorator raises MiddlewareException.""" + + class UndeterminedCallableMiddleware: + async def __call__(self, arg1: Any, arg2: Any) -> None: + pass + + client = MockBaseChatClient() + with pytest.raises(MiddlewareException) as exc_info: + Agent(client=client, middleware=[UndeterminedCallableMiddleware()]) + + assert "UndeterminedCallableMiddleware" in str(exc_info.value) + assert "Cannot determine middleware type" in str(exc_info.value) From 2d716dce8a31ab96b1473b7cdb4e8dfa71487640 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Mon, 27 Jul 2026 08:27:04 -0500 Subject: [PATCH 2/2] test(python): type-annotate test middleware lists to pass test-typing checks --- .../core/tests/core/test_middleware_with_agent.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/packages/core/tests/core/test_middleware_with_agent.py b/python/packages/core/tests/core/test_middleware_with_agent.py index c88193345e1..723ef0c81b8 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -2315,8 +2315,9 @@ async def __call__(self, ctx: Any) -> None: pass client = MockBaseChatClient() + insufficient_middleware: list[Any] = [InsufficientParamsMiddleware()] with pytest.raises(MiddlewareException) as exc_info: - Agent(client=client, middleware=[InsufficientParamsMiddleware()]) + Agent(client=client, middleware=insufficient_middleware) assert "InsufficientParamsMiddleware" in str(exc_info.value) assert "must have at least 2 parameters" in str(exc_info.value) @@ -2324,14 +2325,16 @@ async def __call__(self, ctx: Any) -> None: def test_callable_class_middleware_type_mismatch_raises_middleware_exception(self) -> None: """Test that callable class instance with decorator/annotation mismatch raises MiddlewareException.""" - @agent_middleware class MismatchedCallableMiddleware: + _middleware_type = MiddlewareType.AGENT + async def __call__(self, context: FunctionInvocationContext, call_next: Any) -> None: await call_next() client = MockBaseChatClient() + mismatched_middleware: list[Any] = [MismatchedCallableMiddleware()] with pytest.raises(MiddlewareException) as exc_info: - Agent(client=client, middleware=[MismatchedCallableMiddleware()]) + Agent(client=client, middleware=mismatched_middleware) assert "MismatchedCallableMiddleware" in str(exc_info.value) assert "MiddlewareTypes type mismatch" in str(exc_info.value) @@ -2344,8 +2347,9 @@ async def __call__(self, arg1: Any, arg2: Any) -> None: pass client = MockBaseChatClient() + undetermined_middleware: list[Any] = [UndeterminedCallableMiddleware()] with pytest.raises(MiddlewareException) as exc_info: - Agent(client=client, middleware=[UndeterminedCallableMiddleware()]) + Agent(client=client, middleware=undetermined_middleware) assert "UndeterminedCallableMiddleware" in str(exc_info.value) assert "Cannot determine middleware type" in str(exc_info.value)