From 4a9c3d8e4b40cdb6377ef7f7cdedc5e1fd15901a Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Tue, 17 Feb 2026 14:49:29 -0800 Subject: [PATCH 1/5] Fix streaming branch in weather override middleware sample The streaming branch of weather_override_middleware only prefixed the original weather data via a transform hook instead of replacing the content with the 'perfect weather' override like the non-streaming branch does. Replace with a new ResponseStream that yields the override content as ChatResponseUpdate chunks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../override_result_with_middleware.py | 28 +++++++++---------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/python/samples/02-agents/middleware/override_result_with_middleware.py b/python/samples/02-agents/middleware/override_result_with_middleware.py index efaae28a9b..a72aeb7fc9 100644 --- a/python/samples/02-agents/middleware/override_result_with_middleware.py +++ b/python/samples/02-agents/middleware/override_result_with_middleware.py @@ -2,7 +2,7 @@ import asyncio import re -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterable, Awaitable, Callable from random import randint from typing import Annotated @@ -13,9 +13,9 @@ ChatContext, ChatResponse, ChatResponseUpdate, + Content, Message, ResponseStream, - Role, tool, ) from agent_framework.openai import OpenAIResponsesClient @@ -66,22 +66,20 @@ async def weather_override_middleware(context: ChatContext, call_next: Callable[ ] if context.stream and isinstance(context.result, ResponseStream): - index = {"value": 0} - def _update_hook(update: ChatResponseUpdate) -> ChatResponseUpdate: - for content in update.contents or []: - if not content.text: - continue - content.text = f"Weather Advisory: [{index['value']}] {content.text}" - index["value"] += 1 - return update + async def _override_stream() -> AsyncIterable[ChatResponseUpdate]: + for i, chunk_text in enumerate(chunks): + yield ChatResponseUpdate( + contents=[Content.from_text(text=f"Weather Advisory: [{i}] {chunk_text}")], + role="assistant", + ) - context.result.with_transform_hook(_update_hook) + context.result = ResponseStream(_override_stream()) else: # For non-streaming: just replace with a new message current_text = context.result.text if isinstance(context.result, ChatResponse) else "" custom_message = f"Weather Advisory: [0] {''.join(chunks)} Original message was: {current_text}" - context.result = ChatResponse(messages=[Message(role=Role.ASSISTANT, text=custom_message)]) + context.result = ChatResponse(messages=[Message(role="assistant", text=custom_message)]) async def validate_weather_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: @@ -96,12 +94,12 @@ async def validate_weather_middleware(context: ChatContext, call_next: Callable[ if context.stream and isinstance(context.result, ResponseStream): def _append_validation_note(response: ChatResponse) -> ChatResponse: - response.messages.append(Message(role=Role.ASSISTANT, text=validation_note)) + response.messages.append(Message(role="assistant", text=validation_note)) return response context.result.with_finalizer(_append_validation_note) elif isinstance(context.result, ChatResponse): - context.result.messages.append(Message(role=Role.ASSISTANT, text=validation_note)) + context.result.messages.append(Message(role="assistant", text=validation_note)) async def agent_cleanup_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: @@ -154,7 +152,7 @@ def _sanitize(response: AgentResponse) -> AgentResponse: if not found_validation: raise RuntimeError("Expected validation note not found in agent response.") - cleaned_messages.append(Message(role=Role.ASSISTANT, text=" Agent: OK")) + cleaned_messages.append(Message(role="assistant", text=" Agent: OK")) response.messages = cleaned_messages return response From 9916a92d1e2ffe4b1870e5eda4e5275e05a6f8d7 Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:10:58 -0800 Subject: [PATCH 2/5] Fixed exception handling middleware sample --- .../middleware/exception_handling_with_middleware.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/samples/02-agents/middleware/exception_handling_with_middleware.py b/python/samples/02-agents/middleware/exception_handling_with_middleware.py index d8626a095e..83f604c25c 100644 --- a/python/samples/02-agents/middleware/exception_handling_with_middleware.py +++ b/python/samples/02-agents/middleware/exception_handling_with_middleware.py @@ -47,8 +47,8 @@ async def exception_handling_middleware( print(f"[ExceptionHandlingMiddleware] Caught TimeoutError: {e}") # Override function result to provide custom message in response. context.result = ( - "Request Timeout: The data service is taking longer than expected to respond.", - "Respond with message - 'Sorry for the inconvenience, please try again later.'", + "Request Timeout: The data service is taking longer than expected to respond." + "Respond with message - 'Sorry for the inconvenience, please try again later.'" ) From 7255fa218c098c6d78697f0373a510b04ba57d41 Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:24:08 -0800 Subject: [PATCH 3/5] Fixed runtime context delegation middleware example --- .../packages/core/agent_framework/_agents.py | 7 +- .../packages/core/agent_framework/_types.py | 7 +- .../tests/core/test_middleware_with_agent.py | 173 ++++++++++++++++++ 3 files changed, 180 insertions(+), 7 deletions(-) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 4756bd74e8..890168c6d8 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1024,6 +1024,11 @@ async def _prepare_run_context( await self._async_exit_stack.enter_async_context(mcp_server) final_tools.extend(mcp_server.functions) + # Merge runtime kwargs into additional_function_arguments so they're available + # in function middleware context and tool invocation. + existing_additional_args = opts.pop("additional_function_arguments", None) or {} + additional_function_arguments = {**kwargs, **existing_additional_args} + # Build options dict from run() options merged with provided options run_opts: dict[str, Any] = { "model_id": opts.pop("model_id", None), @@ -1031,7 +1036,7 @@ async def _prepare_run_context( if active_session else opts.pop("conversation_id", None), "allow_multiple_tool_calls": opts.pop("allow_multiple_tool_calls", None), - "additional_function_arguments": opts.pop("additional_function_arguments", None), + "additional_function_arguments": additional_function_arguments or None, "frequency_penalty": opts.pop("frequency_penalty", None), "logit_bias": opts.pop("logit_bias", None), "max_tokens": opts.pop("max_tokens", None), diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index fb43b4e9e4..8446e3ec53 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2871,12 +2871,7 @@ class _ChatOptionsBase(TypedDict, total=False): presence_penalty: float # Tool configuration (forward reference to avoid circular import) - tools: ( - ToolTypes - | Callable[..., Any] - | Sequence[ToolTypes | Callable[..., Any]] - | None - ) + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None tool_choice: ToolMode | Literal["auto", "required", "none"] allow_multiple_tool_calls: bool 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 7a7011412b..bfe5ec1293 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -769,6 +769,179 @@ async def kwargs_middleware( assert modified_kwargs["new_param"] == "added_by_middleware" assert modified_kwargs["custom_param"] == "test_value" + async def test_run_kwargs_available_in_function_middleware(self, chat_client_base: "MockBaseChatClient") -> None: + """Test that kwargs passed directly to agent.run() appear in FunctionInvocationContext.kwargs, + including complex nested values like dicts.""" + captured_kwargs: dict[str, Any] = {} + + @function_middleware + async def capture_middleware( + context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]] + ) -> None: + captured_kwargs.update(context.kwargs) + await call_next() + + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}' + ) + ], + ) + ] + ), + ChatResponse(messages=[Message(role="assistant", text="Done!")]), + ] + + agent = Agent(client=chat_client_base, middleware=[capture_middleware], tools=[sample_tool_function]) + + session_metadata = {"tenant": "acme-corp", "region": "us-west"} + await agent.run( + [Message(role="user", text="Get weather")], + user_id="user-456", + session_metadata=session_metadata, + ) + + assert "user_id" in captured_kwargs, f"Expected 'user_id' in kwargs: {captured_kwargs}" + assert captured_kwargs["user_id"] == "user-456" + assert captured_kwargs["session_metadata"] == {"tenant": "acme-corp", "region": "us-west"} + + async def test_run_kwargs_merged_with_additional_function_arguments( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that explicit additional_function_arguments in options take precedence over run kwargs.""" + captured_kwargs: dict[str, Any] = {} + + @function_middleware + async def capture_middleware( + context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]] + ) -> None: + captured_kwargs.update(context.kwargs) + await call_next() + + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}' + ) + ], + ) + ] + ), + ChatResponse(messages=[Message(role="assistant", text="Done!")]), + ] + + agent = Agent(client=chat_client_base, middleware=[capture_middleware], tools=[sample_tool_function]) + + await agent.run( + [Message(role="user", text="Get weather")], + # This kwarg should be overridden by additional_function_arguments + user_id="from-kwargs", + tenant_id="from-kwargs", + options={ + "additional_function_arguments": { + "user_id": "from-options", + "extra_key": "only-in-options", + } + }, + ) + + # additional_function_arguments takes precedence for overlapping keys + assert captured_kwargs["user_id"] == "from-options" + # Non-overlapping kwargs from run() still come through + assert captured_kwargs["tenant_id"] == "from-kwargs" + # Keys only in additional_function_arguments are present + assert captured_kwargs["extra_key"] == "only-in-options" + + async def test_run_kwargs_consistent_across_multiple_tool_calls( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that kwargs are consistent across multiple tool invocations in a single run.""" + invocation_kwargs: list[dict[str, Any]] = [] + + @function_middleware + async def capture_middleware( + context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]] + ) -> None: + invocation_kwargs.append(dict(context.kwargs)) + await call_next() + + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}' + ), + Content.from_function_call( + call_id="call_2", name="sample_tool_function", arguments='{"location": "Portland"}' + ), + ], + ) + ] + ), + ChatResponse(messages=[Message(role="assistant", text="Done!")]), + ] + + agent = Agent(client=chat_client_base, middleware=[capture_middleware], tools=[sample_tool_function]) + + await agent.run( + [Message(role="user", text="Get weather for both cities")], + user_id="user-456", + request_id="req-001", + ) + + assert len(invocation_kwargs) == 2 + for kw in invocation_kwargs: + assert kw["user_id"] == "user-456" + assert kw["request_id"] == "req-001" + + async def test_run_without_kwargs_produces_empty_context_kwargs( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that when no kwargs are passed to run(), FunctionInvocationContext.kwargs is empty.""" + captured_kwargs: dict[str, Any] = {} + + @function_middleware + async def capture_middleware( + context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]] + ) -> None: + captured_kwargs.update(context.kwargs) + await call_next() + + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}' + ) + ], + ) + ] + ), + ChatResponse(messages=[Message(role="assistant", text="Done!")]), + ] + + agent = Agent(client=chat_client_base, middleware=[capture_middleware], tools=[sample_tool_function]) + + await agent.run([Message(role="user", text="Get weather")]) + + # No runtime kwargs should be present + assert "user_id" not in captured_kwargs + class TestMiddlewareDynamicRebuild: """Test cases for dynamic middleware pipeline rebuilding with Agent.""" From 44c79e7bf200bd502f9770b919152bf02ec362b8 Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:33:39 -0800 Subject: [PATCH 4/5] Fixed multimodal input examples --- .../02-agents/multimodal_input/azure_chat_multimodal.py | 4 ++-- .../multimodal_input/azure_responses_multimodal.py | 4 ++-- .../02-agents/multimodal_input/openai_chat_multimodal.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/samples/02-agents/multimodal_input/azure_chat_multimodal.py b/python/samples/02-agents/multimodal_input/azure_chat_multimodal.py index 369221ac36..969f3b4640 100644 --- a/python/samples/02-agents/multimodal_input/azure_chat_multimodal.py +++ b/python/samples/02-agents/multimodal_input/azure_chat_multimodal.py @@ -9,7 +9,7 @@ def create_sample_image() -> str: """Create a simple 1x1 pixel PNG image for testing.""" - # This is a tiny red pixel in PNG format + # This is a tiny yellow pixel in PNG format png_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" return f"data:image/png;base64,{png_data}" @@ -32,7 +32,7 @@ async def test_image() -> None: ], ) - response = await client.get_response(message) + response = await client.get_response([message]) print(f"Image Response: {response}") diff --git a/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py b/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py index 7a71553f1e..d9a6c18cef 100644 --- a/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py +++ b/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py @@ -41,7 +41,7 @@ async def test_image() -> None: ], ) - response = await client.get_response(message) + response = await client.get_response([message]) print(f"Image Response: {response}") @@ -62,7 +62,7 @@ async def test_pdf() -> None: ], ) - response = await client.get_response(message) + response = await client.get_response([message]) print(f"PDF Response: {response}") diff --git a/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py b/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py index f752d8a52c..d2d2d0eea2 100644 --- a/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py +++ b/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py @@ -53,7 +53,7 @@ async def test_image() -> None: ], ) - response = await client.get_response(message) + response = await client.get_response([message]) print(f"Image Response: {response}") @@ -70,7 +70,7 @@ async def test_audio() -> None: ], ) - response = await client.get_response(message) + response = await client.get_response([message]) print(f"Audio Response: {response}") @@ -89,7 +89,7 @@ async def test_pdf() -> None: ], ) - response = await client.get_response(message) + response = await client.get_response([message]) print(f"PDF Response: {response}") From 07ed686014ad1c9d5d59427d015d26ed0dcc0029 Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:34:08 -0800 Subject: [PATCH 5/5] Small update --- .../02-agents/multimodal_input/azure_responses_multimodal.py | 2 +- .../02-agents/multimodal_input/openai_chat_multimodal.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py b/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py index d9a6c18cef..c020add6f6 100644 --- a/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py +++ b/python/samples/02-agents/multimodal_input/azure_responses_multimodal.py @@ -18,7 +18,7 @@ def load_sample_pdf() -> bytes: def create_sample_image() -> str: """Create a simple 1x1 pixel PNG image for testing.""" - # This is a tiny red pixel in PNG format + # This is a tiny yellow pixel in PNG format png_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" return f"data:image/png;base64,{png_data}" diff --git a/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py b/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py index d2d2d0eea2..22029c8c33 100644 --- a/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py +++ b/python/samples/02-agents/multimodal_input/openai_chat_multimodal.py @@ -19,7 +19,7 @@ def load_sample_pdf() -> bytes: def create_sample_image() -> str: """Create a simple 1x1 pixel PNG image for testing.""" - # This is a tiny red pixel in PNG format + # This is a tiny yellow pixel in PNG format png_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" return f"data:image/png;base64,{png_data}"