From 0b6f4ee0a333a6cc9055b0e129fdddce27cd59cd Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 8 Jul 2026 15:50:36 +0200 Subject: [PATCH 1/3] Python: Add message injection middleware Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/.github/skills/python-samples/SKILL.md | 80 -------- python/AGENTS.md | 1 - python/packages/core/AGENTS.md | 1 + .../packages/core/agent_framework/__init__.py | 6 + .../core/agent_framework/__init__.pyi | 6 + .../core/agent_framework/_middleware.py | 9 + .../core/agent_framework/_sessions.py | 165 +++++++++++++++- .../packages/core/agent_framework/_tools.py | 2 +- .../tests/core/test_middleware_with_chat.py | 177 +++++++++++++++++- python/samples/02-agents/middleware/README.md | 1 + .../message_injection_middleware.py | 87 +++++++++ python/samples/AGENTS.md | 38 +++- 12 files changed, 486 insertions(+), 87 deletions(-) delete mode 100644 python/.github/skills/python-samples/SKILL.md create mode 100644 python/samples/02-agents/middleware/message_injection_middleware.py diff --git a/python/.github/skills/python-samples/SKILL.md b/python/.github/skills/python-samples/SKILL.md deleted file mode 100644 index be992e07719..00000000000 --- a/python/.github/skills/python-samples/SKILL.md +++ /dev/null @@ -1,80 +0,0 @@ ---- -name: python-samples -description: > - Guidelines for creating and modifying sample code in the Agent Framework - Python codebase. Use this when writing new samples or updating existing ones. ---- - -# Python Samples - -## File Structure - -Every sample file follows this order: - -1. PEP 723 inline script metadata (if external dependencies needed) -2. Copyright header: `# Copyright (c) Microsoft. All rights reserved.` -3. Required imports -4. Module docstring: `"""This sample demonstrates..."""` -5. Helper functions -6. Main function(s) demonstrating functionality -7. Entry point: `if __name__ == "__main__": asyncio.run(main())` - -## External Dependencies - -Use [PEP 723](https://peps.python.org/pep-0723/) inline script metadata for -external packages not in the dev environment: - -```python -# /// script -# requires-python = ">=3.10" -# dependencies = [ -# "some-external-package", -# ] -# /// -# Run with: uv run samples/path/to/script.py - -# Copyright (c) Microsoft. All rights reserved. -``` - -Do **not** add sample-only dependencies to the root `pyproject.toml` dev group. - -## Syntax Checking - -```bash -# Format + lint samples -uv run poe syntax -S - -# Check samples for syntax errors and missing imports -uv run poe pyright -S - -# Lint samples only -uv run poe syntax -S -C -``` - -## Documentation - -Samples should be over-documented: - -1. Include a README.md in each set of samples -2. Add a summary docstring under imports explaining the purpose and key components -3. Mark code sections with numbered comments: - ```python - # 1. Create the client instance. - ... - # 2. Create the agent with the client. - ... - ``` -4. Include expected output at the end of the file: - ```python - """ - Sample output: - User:> Why is the sky blue? - Assistant:> The sky is blue due to Rayleigh scattering... - """ - ``` - -## Guidelines - -- **Incremental complexity** — start simple, build up (step1, step2, ...) -- **Getting started naming**: `step_.py` -- When modifying samples, update associated README files diff --git a/python/AGENTS.md b/python/AGENTS.md index ffb62d2b7ff..4bfa1e7e570 100644 --- a/python/AGENTS.md +++ b/python/AGENTS.md @@ -13,7 +13,6 @@ Instructions for AI coding agents working in the Python codebase. - `python-code-quality` — linting, formatting, type checking, prek hooks, CI workflow - `python-feature-lifecycle` — package vs feature lifecycle stages, decorators, enums, and promotion guidance - `python-package-management` — monorepo structure, lazy loading, versioning, new packages -- `python-samples` — sample file structure, PEP 723, documentation guidelines - `pull-requests` — writing PR descriptions (template) and handling/resolving PR review comments ## Maintaining Documentation diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index d8161022b7e..205d9807797 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -69,6 +69,7 @@ agent_framework/ - **`ChatMiddleware`** - Intercepts chat client `get_response()` calls - **`FunctionMiddleware`** - Intercepts function/tool invocations - **`AgentContext`** / **`ChatContext`** / **`FunctionInvocationContext`** - Context objects passed through middleware. A tool can declare a `FunctionInvocationContext` parameter to receive it; `context.tools` is the live, mutable tools list for the run, and `context.add_tools(...)` / `context.remove_tools(...)` enable progressive tool exposure (changes apply on the next function-calling iteration). +- **`MessageInjectionMiddleware`** - Session-scoped chat middleware that lets tools or other code enqueue messages for the next model call in the current `AgentSession`; it drains queued messages into the next call and loops only when no function calls need to be handled by the function invocation layer. ### Sessions (`_sessions.py`) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 873970179d0..60150f3ca12 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -197,8 +197,11 @@ "FileHistoryProvider", "HistoryProvider", "InMemoryHistoryProvider", + "MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY", + "MessageInjectionMiddleware", "ServiceSessionId", "SessionContext", + "enqueue_messages", "register_state_type", ), "._settings": ("SecretString", "load_settings"), @@ -367,6 +370,7 @@ "GROUP_INDEX_KEY", "GROUP_KIND_KEY", "GROUP_TOKEN_COUNT_KEY", + "MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY", "SKIP_PARSING", "SUMMARIZED_BY_SUMMARY_ID_KEY", "SUMMARY_OF_GROUP_IDS_KEY", @@ -492,6 +496,7 @@ "MemoryStore", "MemoryTopicRecord", "Message", + "MessageInjectionMiddleware", "MiddlewareException", "MiddlewareTermination", "MiddlewareType", @@ -593,6 +598,7 @@ "create_edge_runner", "create_harness_agent", "detect_media_type_from_base64", + "enqueue_messages", "evaluate_agent", "evaluate_workflow", "evaluator", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 8894d75f856..ee72c15cc01 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -151,13 +151,16 @@ from ._middleware import ( function_middleware, ) from ._sessions import ( + MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, AgentSession, ContextProvider, FileHistoryProvider, HistoryProvider, InMemoryHistoryProvider, + MessageInjectionMiddleware, ServiceSessionId, SessionContext, + enqueue_messages, register_state_type, ) from ._settings import SecretString, load_settings @@ -334,6 +337,7 @@ __all__ = [ "GROUP_INDEX_KEY", "GROUP_KIND_KEY", "GROUP_TOKEN_COUNT_KEY", + "MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY", "SKIP_PARSING", "SUMMARIZED_BY_SUMMARY_ID_KEY", "SUMMARY_OF_GROUP_IDS_KEY", @@ -459,6 +463,7 @@ __all__ = [ "MemoryStore", "MemoryTopicRecord", "Message", + "MessageInjectionMiddleware", "MiddlewareException", "MiddlewareTermination", "MiddlewareType", @@ -560,6 +565,7 @@ __all__ = [ "create_edge_runner", "create_harness_agent", "detect_media_type_from_base64", + "enqueue_messages", "evaluate_agent", "evaluate_workflow", "evaluator", diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 5315575b6af..c82b1eab69e 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -381,6 +381,7 @@ class ChatContext: messages: The messages being sent to the chat client. options: The options for the chat request as a dict. stream: Whether this is a streaming invocation. + session: The active agent session for this chat invocation, if any. metadata: Metadata dictionary for sharing data between chat middleware. result: Chat execution result. Can be observed after calling ``call_next()`` to see the actual execution result or can be set to override the execution result. @@ -421,6 +422,7 @@ def __init__( messages: Sequence[Message], options: Mapping[str, Any] | None, stream: bool = False, + session: AgentSession | None = None, metadata: Mapping[str, Any] | None = None, result: ChatResponse | ResponseStream[ChatResponseUpdate, ChatResponse] | None = None, kwargs: Mapping[str, Any] | None = None, @@ -439,6 +441,7 @@ def __init__( messages: The messages being sent to the chat client. options: The options for the chat request as a dict. stream: Whether this is a streaming invocation. + session: The active agent session for this chat invocation, if any. metadata: Metadata dictionary for sharing data between chat middleware. result: Chat execution result. kwargs: Additional keyword arguments passed to the chat client. @@ -451,6 +454,7 @@ def __init__( self.messages = messages self.options = options self.stream = stream + self.session = session self.metadata: dict[str, Any] = dict(metadata) if metadata is not None else {} self.result = result self.kwargs: dict[str, Any] = dict(kwargs) if kwargs is not None else {} @@ -1181,6 +1185,10 @@ def get_response( super_get_response = super().get_response # type: ignore[misc] effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {} call_middleware = effective_client_kwargs.pop("middleware", []) + raw_session = effective_client_kwargs.pop("session", None) + from ._sessions import AgentSession as _AgentSession + + session = raw_session if isinstance(raw_session, _AgentSession) else None context_kwargs = dict(effective_client_kwargs) if compaction_strategy is not None: context_kwargs["compaction_strategy"] = compaction_strategy @@ -1203,6 +1211,7 @@ def get_response( messages=list(messages), options=options, stream=stream, + session=session, kwargs=context_kwargs, function_invocation_kwargs=function_invocation_kwargs, ) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index b896cafccbe..b7d52edcbdb 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -22,7 +22,7 @@ import weakref from abc import abstractmethod from base64 import urlsafe_b64encode -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence from contextlib import suppress from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypeGuard, cast @@ -31,12 +31,15 @@ from ._middleware import ChatContext, ChatMiddleware from ._types import ( AgentResponse, + AgentRunInputs, ChatResponse, + ChatResponseUpdate, Message, ResponseStream, _build_agent_response_from_chat_response, # pyright: ignore[reportPrivateUsage] + normalize_messages, ) -from .exceptions import ChatClientInvalidResponseException +from .exceptions import ChatClientInvalidRequestException, ChatClientInvalidResponseException if TYPE_CHECKING: from ._agents import SupportsAgentRun @@ -47,6 +50,8 @@ # Registry of known types for state deserialization _STATE_TYPE_REGISTRY: dict[str, type] = {} +MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY: str = "message_injection.pending_messages" +_MESSAGE_INJECTION_LOCK = threading.Lock() JsonDumps: TypeAlias = Callable[[Any], str | bytes] JsonLoads: TypeAlias = Callable[[str | bytes], Any] @@ -551,6 +556,8 @@ def is_local_history_conversation_id(conversation_id: str | None) -> bool: def _response_contains_follow_up_request(response: ChatResponse) -> bool: """Return whether a response requires another model call in the current run.""" + # TODO(eavanvalkenburg): When informational-only function call content lands, ignore informational-only calls here + # so hosted/provider-executed tool transcript items do not block injected-message processing. return any( item.type in {"function_call", "function_approval_request"} for message in response.messages @@ -574,6 +581,160 @@ def _split_service_call_messages(messages: Sequence[Message]) -> tuple[list[Mess return input_messages, context_messages +def enqueue_messages(session: AgentSession, messages: AgentRunInputs) -> None: + """Enqueue messages for the next model call in the given session. + + Args: + session: The session whose pending message queue should receive the messages. + messages: The messages to enqueue. Accepts the same flexible shapes as ``Agent.run`` input: + a string, ``Content``, ``Message``, or a sequence of those. + """ + pending_messages = normalize_messages(messages) + if not pending_messages: + return + with _MESSAGE_INJECTION_LOCK: + queue = cast( + list[Message], + session.state.setdefault(MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, []), + ) + queue.extend(pending_messages) + + +class MessageInjectionMiddleware(ChatMiddleware): + """Chat middleware that injects queued session messages into the model call loop. + + Messages can be enqueued for an :class:`AgentSession` before a run starts or while a run is in progress, + including from tool code that receives a :class:`FunctionInvocationContext`. Pending messages are stored in + ``session.state`` and drained into the next model call for that session. After a model call completes, the + middleware loops internally only when there are newly queued messages and the response does not contain function + calls that the function invocation layer must handle. + """ + + def __init__(self) -> None: + """Initialize the middleware.""" + + def enqueue_messages(self, session: AgentSession, messages: AgentRunInputs) -> None: + """Enqueue messages for the next model call in the given session. + + Args: + session: The session whose pending message queue should receive the messages. + messages: The messages to enqueue. Accepts the same flexible shapes as ``Agent.run`` input: + a string, ``Content``, ``Message``, or a sequence of those. + """ + enqueue_messages(session, messages) + + def get_pending_messages(self, session: AgentSession) -> list[Message]: + """Return a snapshot of messages queued for the given session. + + Args: + session: The session whose pending messages should be returned. + + Returns: + A point-in-time copy of the queued messages. The returned list is not updated if the queue is later + drained or extended. + """ + with _MESSAGE_INJECTION_LOCK: + return list(cast(list[Message], session.state.get(MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, []))) + + def _drain_pending_messages(self, session: AgentSession, messages: Sequence[Message]) -> list[Message]: + with _MESSAGE_INJECTION_LOCK: + queue = cast( + list[Message], + session.state.setdefault(MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, []), + ) + if not queue: + return list(messages) + next_messages = [*messages, *queue] + queue.clear() + return next_messages + + def _has_pending_messages(self, session: AgentSession) -> bool: + with _MESSAGE_INJECTION_LOCK: + return bool(session.state.get(MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, [])) + + @staticmethod + def _update_context_conversation_id(context: ChatContext, conversation_id: str | None) -> None: + if conversation_id is None: + return + context.kwargs["conversation_id"] = conversation_id + if context.options is None: + context.options = {"conversation_id": conversation_id} + return + context.options = {**context.options, "conversation_id": conversation_id} + + async def _process_non_streaming( + self, + context: ChatContext, + call_next: Callable[[], Awaitable[None]], + session: AgentSession, + ) -> None: + while True: + context.messages = self._drain_pending_messages(session, context.messages) + context.result = None + await call_next() + if context.result is None: + return + if isinstance(context.result, ResponseStream): + raise ValueError("Non-streaming message injection middleware requires a ChatResponse result.") + response = cast(ChatResponse, context.result) + if _response_contains_follow_up_request(response) or not self._has_pending_messages(session): + return + self._update_context_conversation_id(context, response.conversation_id) + empty_messages: list[Message] = [] + context.messages = empty_messages + + async def _stream_injected_messages( + self, + context: ChatContext, + call_next: Callable[[], Awaitable[None]], + session: AgentSession, + ) -> AsyncIterable[ChatResponseUpdate]: + while True: + context.messages = self._drain_pending_messages(session, context.messages) + context.result = None + await call_next() + if context.result is None: + return + if not isinstance(context.result, ResponseStream): + raise ValueError("Streaming message injection middleware requires a ResponseStream result.") + stream = cast(ResponseStream[ChatResponseUpdate, ChatResponse], context.result) + async for update in stream: + yield update + response = await stream.get_final_response() + if _response_contains_follow_up_request(response) or not self._has_pending_messages(session): + return + self._update_context_conversation_id(context, response.conversation_id) + empty_messages: list[Message] = [] + context.messages = empty_messages + + async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + """Inject pending session messages into chat model calls. + + Args: + context: The chat invocation context for the current model call. + call_next: The next middleware or leaf chat client. + + Raises: + ChatClientInvalidRequestException: If the middleware is used without an active ``AgentSession``. + ValueError: If downstream middleware returns a non-streaming result for streaming mode, or vice versa. + """ + session = context.session + if session is None: + raise ChatClientInvalidRequestException( + "MessageInjectionMiddleware requires an AgentSession. Pass session=... when running the agent." + ) + + if not context.stream: + await self._process_non_streaming(context, call_next, session) + return + + response_format = context.options.get("response_format") if context.options is not None else None + context.result = ResponseStream( + self._stream_injected_messages(context, call_next, session), + finalizer=lambda updates: ChatResponse.from_updates(updates, output_format_type=response_format), + ) + + class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware): """Persist local chat history after each service call when history is framework-managed. diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 76fda457339..14ce1d4b34b 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2531,7 +2531,7 @@ def get_response( invocation_session=invocation_session, middleware_pipeline=function_middleware_pipeline, ) - filtered_kwargs = {k: v for k, v in effective_client_kwargs.items() if k != "session"} + filtered_kwargs = dict(effective_client_kwargs) # Make options mutable so we can update conversation_id during function invocation loop mutable_options: dict[str, Any] = dict(options) if options else {} diff --git a/python/packages/core/tests/core/test_middleware_with_chat.py b/python/packages/core/tests/core/test_middleware_with_chat.py index c4aaffc7a88..8b8d9803f42 100644 --- a/python/packages/core/tests/core/test_middleware_with_chat.py +++ b/python/packages/core/tests/core/test_middleware_with_chat.py @@ -1,11 +1,15 @@ # Copyright (c) Microsoft. All rights reserved. -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterable, Awaitable, Callable, Sequence from typing import Any, cast from unittest.mock import patch +import pytest + from agent_framework import ( + MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, Agent, + AgentSession, ChatContext, ChatMiddleware, ChatMiddlewareTypes, @@ -15,10 +19,15 @@ FunctionInvocationContext, FunctionTool, Message, + MessageInjectionMiddleware, + ResponseStream, SupportsChatGetResponse, chat_middleware, + enqueue_messages, function_middleware, + tool, ) +from agent_framework.exceptions import ChatClientInvalidRequestException from .conftest import MockBaseChatClient @@ -373,6 +382,172 @@ async def kwargs_middleware(context: ChatContext, call_next: Callable[[], Awaita assert modified_options["new_param"] == "added_by_middleware" assert modified_options["custom_param"] == "test_value" + async def test_message_injection_middleware_appends_prequeued_messages( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that queued session messages are appended to the next model call.""" + session = AgentSession() + injection = MessageInjectionMiddleware() + enqueue_messages(session, "queued message") + captured_messages: list[list[str | None]] = [] + + async def fake_get_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + captured_messages.append([message.text for message in messages]) + return ChatResponse(messages=Message(role="assistant", contents=["ok"])) + + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=fake_get_response): + agent = Agent(client=chat_client_base, middleware=[injection]) + response = await agent.run("user message", session=session) + + assert response.messages[0].text == "ok" + assert captured_messages == [["user message", "queued message"]] + assert injection.get_pending_messages(session) == [] + + async def test_message_injection_middleware_loops_when_messages_are_queued_after_call( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that queued messages after a non-tool response trigger another model call.""" + session = AgentSession() + injection = MessageInjectionMiddleware() + captured_messages: list[list[str | None]] = [] + captured_conversation_ids: list[str | None] = [] + + async def fake_get_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + captured_messages.append([message.text for message in messages]) + captured_conversation_ids.append(options.get("conversation_id")) + if len(captured_messages) == 1: + enqueue_messages(session, "queued during call") + return ChatResponse( + messages=Message(role="assistant", contents=["first"]), + conversation_id="conversation-1", + ) + return ChatResponse(messages=Message(role="assistant", contents=["second"])) + + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=fake_get_response): + response = await chat_client_base.get_response( + [Message(role="user", contents=["user message"])], + client_kwargs={"middleware": [injection], "session": session}, + ) + + assert response.messages[0].text == "second" + assert captured_messages == [["user message"], ["queued during call"]] + assert captured_conversation_ids == [None, "conversation-1"] + + async def test_message_injection_middleware_tool_enqueued_messages_wait_for_function_results( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that tool-enqueued messages are injected after function results are available.""" + session = AgentSession() + injection = MessageInjectionMiddleware() + captured_messages: list[list[Message]] = [] + responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[Content.from_function_call(call_id="call-1", name="inject_message", arguments={})], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + @tool(approval_mode="never_require") + def inject_message(ctx: FunctionInvocationContext) -> str: + """Inject a message into the active session.""" + active_session = ctx.session + if active_session is None: + raise AssertionError("Expected an active session.") + assert active_session is session + enqueue_messages(active_session, "queued from tool") + return "tool result" + + async def fake_get_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + captured_messages.append(list(messages)) + return responses.pop(0) + + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=fake_get_response): + agent = Agent(client=chat_client_base, middleware=[injection], tools=[inject_message]) + response = await agent.run("user message", session=session) + + second_call_contents = [content for message in captured_messages[1] for content in message.contents] + assert response.messages[-1].text == "done" + assert [message.text for message in captured_messages[0]] == ["user message"] + assert any(content.type == "function_result" for content in second_call_contents) + assert captured_messages[1][-1].text == "queued from tool" + + async def test_message_injection_middleware_loops_for_streaming_pending_messages( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that queued messages after a streaming response trigger another streaming model call.""" + session = AgentSession() + injection = MessageInjectionMiddleware() + captured_messages: list[list[str | None]] = [] + + def fake_streaming_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + captured_messages.append([message.text for message in messages]) + + async def stream() -> AsyncIterable[ChatResponseUpdate]: + if len(captured_messages) == 1: + yield ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant") + enqueue_messages(session, "queued while streaming") + return + yield ChatResponseUpdate(contents=[Content.from_text("second")], role="assistant") + + return ResponseStream( + stream(), + finalizer=lambda updates: ChatResponse.from_updates( + updates, + output_format_type=options.get("response_format"), + ), + ) + + with patch.object(chat_client_base, "_get_streaming_response", side_effect=fake_streaming_response): + stream = chat_client_base.get_response( + [Message(role="user", contents=["user message"])], + stream=True, + client_kwargs={"middleware": [injection], "session": session}, + ) + updates = [update async for update in stream] + + assert [update.text for update in updates] == ["first", "second"] + assert captured_messages == [["user message"], ["queued while streaming"]] + + def test_enqueue_messages_uses_session_state_queue(self) -> None: + """Test that standalone message injection enqueueing stores messages in session state.""" + session = AgentSession() + + enqueue_messages(session, "queued message") + + queued_messages = session.state[MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY] + assert [message.text for message in queued_messages] == ["queued message"] + + async def test_message_injection_middleware_requires_session(self, chat_client_base: "MockBaseChatClient") -> None: + """Test that message injection middleware fails clearly without an active session.""" + with pytest.raises(ChatClientInvalidRequestException, match="requires an AgentSession"): + await chat_client_base.get_response( + [Message(role="user", contents=["user message"])], + client_kwargs={"middleware": [MessageInjectionMiddleware()]}, + ) + def test_chat_middleware_pipeline_cache_reuses_matching_middleware( self, chat_client_base: "MockBaseChatClient", diff --git a/python/samples/02-agents/middleware/README.md b/python/samples/02-agents/middleware/README.md index af255c47f1e..205dd01504a 100644 --- a/python/samples/02-agents/middleware/README.md +++ b/python/samples/02-agents/middleware/README.md @@ -18,6 +18,7 @@ This folder contains focused middleware samples for `Agent`, chat clients, tools | [`exception_handling_with_middleware.py`](./exception_handling_with_middleware.py) | Shows how middleware can handle failures and recover cleanly. | | [`function_based_middleware.py`](./function_based_middleware.py) | Shows function-based agent and function middleware. | | [`middleware_termination.py`](./middleware_termination.py) | Demonstrates stopping a middleware pipeline early. | +| [`message_injection_middleware.py`](./message_injection_middleware.py) | Demonstrates `MessageInjectionMiddleware` with a real Foundry chat client: enqueueing a follow-up message into the active session while a long-running async tool is awaiting. | | [`override_result_with_middleware.py`](./override_result_with_middleware.py) | Shows how middleware can replace regular and streaming results, then post-process the final response. | | [`runtime_context_delegation.py`](./runtime_context_delegation.py) | Demonstrates delegating arguments with runtime context data. | | [`session_behavior_middleware.py`](./session_behavior_middleware.py) | Shows how middleware interacts with session-backed runs. | diff --git a/python/samples/02-agents/middleware/message_injection_middleware.py b/python/samples/02-agents/middleware/message_injection_middleware.py new file mode 100644 index 00000000000..e3e5c913581 --- /dev/null +++ b/python/samples/02-agents/middleware/message_injection_middleware.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from typing import Annotated + +from agent_framework import Agent, AgentSession, MessageInjectionMiddleware, enqueue_messages, tool +from agent_framework.foundry import FoundryChatClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv + +""" +This sample demonstrates MessageInjectionMiddleware with a real FoundryChatClient. + +The sample starts an agent run that is expected to call a long-running async tool. While that tool is waiting on +``asyncio.sleep()``, the application regains control and enqueues a new user message into the same AgentSession. +After the tool completes, MessageInjectionMiddleware drains that queued message into the next model call so the model +can include it in the final answer without starting a separate agent run. +""" + + +load_dotenv() + + +# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; +# see samples/02-agents/tools/function_tool_with_approval.py +# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. +@tool(approval_mode="never_require") +async def slow_inventory_lookup( + item: Annotated[str, "The item to check inventory for."], +) -> str: + """Look up inventory for an item, intentionally taking long enough to inject a follow-up message.""" + print(f"Tool: checking inventory for {item!r}...") + await asyncio.sleep(8) + print("Tool: inventory lookup finished.") + return f"{item} is in stock, with curbside pickup available today." + + +async def main() -> None: + """Run the message injection middleware sample.""" + print("=== Message Injection Middleware Example ===") + + # 1. Create the message injection middleware and the session that owns its pending-message queue. + message_injection = MessageInjectionMiddleware() + session = AgentSession() + + # 2. Create a regular FoundryChatClient-backed agent. + # For authentication, run `az login` or replace AzureCliCredential with your preferred authentication option. + agent = Agent( + client=FoundryChatClient(credential=AzureCliCredential()), + name="InventoryAgent", + instructions=( + "You help with store inventory questions. Always call slow_inventory_lookup before answering inventory " + "questions. If another user message arrives before your final answer, account for it in that final answer." + ), + middleware=[message_injection], + tools=slow_inventory_lookup, + ) + + # 3. Start the run. The model should call slow_inventory_lookup, which awaits asyncio.sleep(). + question = "Can I pick up a red travel mug today? Check inventory before answering." + print(f"User:> {question}") + run_task = asyncio.ensure_future(agent.run(question, session=session)) + + # 4. While the tool is sleeping, enqueue a new message into the same session. + await asyncio.sleep(2) + follow_up = "Please also mention that I can only pick it up after 5 PM." + print(f"User (injected while tool is running):> {follow_up}") + enqueue_messages(session, follow_up) + + # 5. Await the original run. The final model call sees both the tool result and the injected message. + response = await run_task + print(f"Assistant:> {response.text}") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Sample output: +=== Message Injection Middleware Example === +User:> Can I pick up a red travel mug today? Check inventory before answering. +Tool: checking inventory for 'red travel mug'... +User (injected while tool is running):> Please also mention that I can only pick it up after 5 PM. +Tool: inventory lookup finished. +Assistant:> Yes, the red travel mug is in stock and curbside pickup is available today. Since you can only pick it up +after 5 PM, choose an evening pickup window when placing the order. +""" diff --git a/python/samples/AGENTS.md b/python/samples/AGENTS.md index cf162f3e7da..6bc8581ed47 100644 --- a/python/samples/AGENTS.md +++ b/python/samples/AGENTS.md @@ -71,6 +71,7 @@ with an Azure AI Foundry project endpoint: ```python import os +from agent_framework import Agent from agent_framework.foundry import FoundryChatClient from azure.identity import AzureCliCredential @@ -80,7 +81,7 @@ client = FoundryChatClient( model=os.environ["FOUNDRY_MODEL"], credential=credential, ) -agent = client.as_agent(name="...", instructions="...") +agent = Agent(client=client, name="...", instructions="...") ``` Environment variables: @@ -107,10 +108,43 @@ pip install agent-framework `agent-framework` is released, so `--pre` is not required here. `openai` is a core dependency. +## File structure + +Every sample file follows this order: + +1. PEP 723 inline script metadata (if external dependencies are needed) +2. Copyright header: `# Copyright (c) Microsoft. All rights reserved.` +3. Required imports +4. Module docstring explaining the purpose and key components +5. Helper functions +6. Main function(s) demonstrating functionality +7. Entry point: `if __name__ == "__main__": asyncio.run(main())` + +Use PEP 723 inline script metadata for external sample-only dependencies; do not add sample-only dependencies to +the root `pyproject.toml` dev group. + +## Syntax checking + +Run sample checks from the `python/` directory: + +```bash +uv run poe syntax -S +uv run poe pyright -S +``` + +## Documentation + +Samples should be over-documented: + +1. Include a README.md in each set of samples. +2. Mark code sections with numbered comments. +3. Include expected output at the end of the file. + ## Current API notes - `Agent` class renamed from `ChatAgent` (use `from agent_framework import Agent`) - `Message` class renamed from `ChatMessage` (use `from agent_framework import Message`) - `call_next` in middleware takes NO arguments: `await call_next()` (not `await call_next(context)`) -- Prefer `client.as_agent(...)` over `Agent(client=client, ...)` +- Do not use `client.as_agent(...)` in samples; construct agents explicitly with `Agent(client=client, ...)`. - Tool methods on hosted tools are now functions, not classes (e.g. `hosted_mcp_tool(...)` not `HostedMCPTool(...)`) +- When only using a description for the field of a `@tool` parameter, do not use `Field`; use the string directly. From dbbddd57b64ca918cbb604cf5eb5e3afcd0e7ccb Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 9 Jul 2026 09:03:17 +0200 Subject: [PATCH 2/3] Python: Ignore informational tool calls for message injection Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/agent_framework/_sessions.py | 4 +- .../tests/core/test_middleware_with_chat.py | 93 +++++++++++++++++++ 2 files changed, 94 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index a19d2d6180c..7577cd3e5ef 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -557,10 +557,8 @@ def is_local_history_conversation_id(conversation_id: str | None) -> bool: def _response_contains_follow_up_request(response: ChatResponse) -> bool: """Return whether a response requires another model call in the current run.""" - # TODO(eavanvalkenburg): When informational-only function call content lands, ignore informational-only calls here - # so hosted/provider-executed tool transcript items do not block injected-message processing. return any( - item.type in {"function_call", "function_approval_request"} + item.type == "function_approval_request" or (item.type == "function_call" and not item.informational_only) for message in response.messages for item in message.contents ) diff --git a/python/packages/core/tests/core/test_middleware_with_chat.py b/python/packages/core/tests/core/test_middleware_with_chat.py index 8b8d9803f42..7c2cd471895 100644 --- a/python/packages/core/tests/core/test_middleware_with_chat.py +++ b/python/packages/core/tests/core/test_middleware_with_chat.py @@ -443,6 +443,47 @@ async def fake_get_response( assert captured_messages == [["user message"], ["queued during call"]] assert captured_conversation_ids == [None, "conversation-1"] + async def test_message_injection_middleware_ignores_informational_only_function_calls( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that hosted tool transcript calls do not block injected messages.""" + session = AgentSession() + injection = MessageInjectionMiddleware() + captured_messages: list[list[str | None]] = [] + + async def fake_get_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + captured_messages.append([message.text for message in messages]) + if len(captured_messages) == 1: + enqueue_messages(session, "queued after hosted tool") + return ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="hosted-call", + name="hosted_search", + arguments={"query": "docs"}, + informational_only=True, + ) + ], + ) + ) + return ChatResponse(messages=Message(role="assistant", contents=["done"])) + + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=fake_get_response): + response = await chat_client_base.get_response( + [Message(role="user", contents=["user message"])], + client_kwargs={"middleware": [injection], "session": session}, + ) + + assert response.messages[0].text == "done" + assert captured_messages == [["user message"], ["queued after hosted tool"]] + async def test_message_injection_middleware_tool_enqueued_messages_wait_for_function_results( self, chat_client_base: "MockBaseChatClient" ) -> None: @@ -531,6 +572,58 @@ async def stream() -> AsyncIterable[ChatResponseUpdate]: assert [update.text for update in updates] == ["first", "second"] assert captured_messages == [["user message"], ["queued while streaming"]] + async def test_message_injection_middleware_streaming_ignores_informational_only_function_calls( + self, chat_client_base: "MockBaseChatClient" + ) -> None: + """Test that streamed hosted tool transcript calls do not block injected messages.""" + session = AgentSession() + injection = MessageInjectionMiddleware() + captured_messages: list[list[str | None]] = [] + + def fake_streaming_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + captured_messages.append([message.text for message in messages]) + + async def stream() -> AsyncIterable[ChatResponseUpdate]: + if len(captured_messages) == 1: + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + call_id="hosted-call", + name="hosted_search", + arguments={"query": "docs"}, + informational_only=True, + ) + ], + role="assistant", + ) + enqueue_messages(session, "queued while streaming hosted tool") + return + yield ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant") + + return ResponseStream( + stream(), + finalizer=lambda updates: ChatResponse.from_updates( + updates, + output_format_type=options.get("response_format"), + ), + ) + + with patch.object(chat_client_base, "_get_streaming_response", side_effect=fake_streaming_response): + stream = chat_client_base.get_response( + [Message(role="user", contents=["user message"])], + stream=True, + client_kwargs={"middleware": [injection], "session": session}, + ) + updates = [update async for update in stream] + + assert [update.text for update in updates] == ["", "done"] + assert captured_messages == [["user message"], ["queued while streaming hosted tool"]] + def test_enqueue_messages_uses_session_state_queue(self) -> None: """Test that standalone message injection enqueueing stores messages in session state.""" session = AgentSession() From 16d6258eb45bd9b08b4ef637d7120f964ed058e8 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 9 Jul 2026 09:31:08 +0200 Subject: [PATCH 3/3] Python: Preserve per-service history with message injection Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/_agents.py | 23 +++++----- .../packages/core/tests/core/test_agents.py | 42 +++++++++++++++++++ 2 files changed, 55 insertions(+), 10 deletions(-) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index f38e3d75598..fae0d036417 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1410,6 +1410,7 @@ async def _prepare_run_context( effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {} if active_session is not None: effective_client_kwargs["session"] = active_session + per_service_call_history_middleware: PerServiceCallHistoryPersistingMiddleware | None = None if per_service_call_history_providers and active_session is not None: per_service_call_history_middleware = PerServiceCallHistoryPersistingMiddleware( agent=self, @@ -1417,16 +1418,6 @@ async def _prepare_run_context( providers=per_service_call_history_providers, service_stores_history=service_stores_history, ) - existing_middleware = effective_client_kwargs.get("middleware") - if isinstance(existing_middleware, Sequence) and not isinstance(existing_middleware, (str, bytes)): - effective_client_kwargs["middleware"] = [per_service_call_history_middleware, *existing_middleware] - elif existing_middleware is not None: - effective_client_kwargs["middleware"] = [ - per_service_call_history_middleware, - cast(MiddlewareTypes, existing_middleware), - ] - else: - effective_client_kwargs["middleware"] = [per_service_call_history_middleware] provider_middleware = session_context.get_middleware() if provider_middleware: middleware_list = categorize_middleware(provider_middleware) @@ -1449,6 +1440,18 @@ async def _prepare_run_context( else: effective_client_kwargs["middleware"] = provider_function_chat_middleware + if per_service_call_history_middleware is not None: + existing_middleware = effective_client_kwargs.get("middleware") + if isinstance(existing_middleware, Sequence) and not isinstance(existing_middleware, (str, bytes)): + effective_client_kwargs["middleware"] = [*existing_middleware, per_service_call_history_middleware] + elif existing_middleware is not None: + effective_client_kwargs["middleware"] = [ + cast(MiddlewareTypes, existing_middleware), + per_service_call_history_middleware, + ] + else: + effective_client_kwargs["middleware"] = [per_service_call_history_middleware] + return { "session": active_session, "session_context": session_context, diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 347405b0f6c..67cddeccd15 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -29,6 +29,7 @@ HistoryProvider, InMemoryHistoryProvider, Message, + MessageInjectionMiddleware, ResponseStream, ServiceSessionId, SessionContext, @@ -37,6 +38,7 @@ SupportsChatGetResponse, TruncationStrategy, chat_middleware, + enqueue_messages, tool, ) from agent_framework._agents import _get_tool_name, _merge_options, _sanitize_agent_name @@ -511,6 +513,46 @@ def lookup_weather(location: str) -> str: assert session.service_session_id is None +async def test_message_injection_persists_each_injected_service_call( + chat_client_base: SupportsChatGetResponse, +) -> None: + provider = _RecordingHistoryProvider() + session = AgentSession() + session.state[provider.source_id] = {"messages": []} + captured_messages: list[list[str | None]] = [] + + async def fake_get_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + captured_messages.append([message.text for message in messages]) + if len(captured_messages) == 1: + enqueue_messages(session, "queued during first service call") + return ChatResponse(messages=Message(role="assistant", contents=["first"])) + return ChatResponse(messages=Message(role="assistant", contents=["second"])) + + agent = Agent( + client=chat_client_base, + context_providers=[provider], + middleware=[MessageInjectionMiddleware()], + require_per_service_call_history_persistence=True, + ) + + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=fake_get_response): + result = await agent.run("initial message", session=session) + + provider_state = session.state[provider.source_id] + stored_messages = cast(list[Message], provider_state["messages"]) + + assert result.text == "second" + assert captured_messages == [["initial message"], ["initial message", "first", "queued during first service call"]] + assert provider_state["get_call_count"] == 2 + assert provider_state["save_call_count"] == 2 + assert stored_messages[-1].text == "second" + + async def test_per_service_call_history_provider_receives_full_agent_response_metadata( chat_client_base: SupportsChatGetResponse, ) -> None: