Python: Add message injection middleware#6998
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
giles17
added
documentation
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||||||||||||
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Automated Code Review
Reviewers: 4 | Confidence: 83%
✓ Correctness
The implementation is well-structured and correct. The MessageInjectionMiddleware properly coordinates with the function invocation layer by yielding control when function calls are present, and only loping internally when messages are queued without pending function calls. Thread safety is appropriate for the brief list operations. The session propagation through the middleware and tool layers is clean—session is extracted at the middleware layer and placed on ChatContext without leaking to the leaf chat client. No correctness issues found.
✓ Security Reliability
The MessageInjectionMiddleware implementation is well-structured with proper thread-safe locking, correct session flow through the middleware/tool layers, and clear error handling for missing sessions. The main reliability concern is the unbounded
while Trueloops in the non-streaming and streaming paths — unlike the function invocation layer which has aDEFAULT_MAX_ITERATIONS = 40guard, the message injection loops have no iteration limit. A misbehaving downstream component that continuously enqueues messages could cause runaway API calls.
✓ Test Coverage
The test suite for MessageInjectionMiddleware is well-structured with focused tests for the key non-streaming scenarios (pre-queued messages, loop-on-queue, tool-enqueued deferral), the streaming loop path, session-required validation, and state-key storage. The primary test coverage gap is the absence of a streaming-mode test where the response contains function calls—a distinct code path in
_stream_injected_messages(line 704) that should NOT loop, validated only indirectly by the non-streaming tool test. A minor gap is the lack of a test for accumulating multipleenqueue_messagescalls before drain.
✓ Design Approach
I found one blocking design issue in the new streaming test coverage. The PR currently codifies a queued-message loop that only works when the chat service stores history server-side; for the framework’s default stateless client model, the second streamed model call would lose the prior transcript. Existing function-loop code in the repo already handles this distinction, so the new test is locking in the wrong behavior for one important client class.
Suggestions
- Consider adding a configurable max-iteration guard to the
while Trueloops in_process_non_streamingand_stream_injected_messages, similar toDEFAULT_MAX_ITERATIONS = 40in the function invocation layer (_tools.py:90). This provides defense-in-depth against runaway API calls if a caller bug continuously enqueues messages.
Automated review by eavanvalkenburg's agents
There was a problem hiding this comment.
Pull request overview
Adds a Python-native “message injection” mechanism that lets host code/tools enqueue additional messages into an active AgentSession and have them automatically included in the next model call within the same run (including during streaming), without introducing a chat-client wrapper pattern.
Changes:
- Introduces
MessageInjectionMiddleware,enqueue_messages(session, ...), andMESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, enabling session-scoped message queueing and draining between model calls. - Plumbs
AgentSessionthroughChatContext(middleware-visible only) so chat middleware can coordinate with session state without forwardingsessionto leaf chat clients. - Adds focused core tests and a Foundry-backed middleware sample; consolidates sample guidelines into
python/samples/AGENTS.mdand removes the duplicatepython-samplesskill.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| python/samples/AGENTS.md | Consolidates sample guidance; updates Foundry snippet to construct Agent(...) explicitly. |
| python/samples/02-agents/middleware/README.md | Adds the new message-injection sample to the middleware index. |
| python/samples/02-agents/middleware/message_injection_middleware.py | New sample demonstrating injecting a follow-up user message while a long-running tool is awaiting. |
| python/packages/core/tests/core/test_middleware_with_chat.py | Adds tests covering prequeued messages, post-call queueing, tool-queued messages, streaming, and missing-session behavior. |
| python/packages/core/AGENTS.md | Documents MessageInjectionMiddleware as a session-scoped middleware concept. |
| python/packages/core/agent_framework/_tools.py | Ensures session remains available to middleware layers by no longer stripping it early. |
| python/packages/core/agent_framework/_sessions.py | Implements session-state queueing and the MessageInjectionMiddleware loop behavior. |
| python/packages/core/agent_framework/_middleware.py | Adds session to ChatContext and extracts it from client_kwargs for middleware use. |
| python/packages/core/agent_framework/init.pyi | Exposes new public API symbols in typing stubs. |
| python/packages/core/agent_framework/init.py | Exposes new public API symbols at runtime via lazy exports. |
| python/AGENTS.md | Removes the python-samples skill reference from Python agent instructions. |
| python/.github/skills/python-samples/SKILL.md | Deletes the duplicated samples skill (now covered by python/samples/AGENTS.md). |
| 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 |
| ## File structure | ||
|
|
||
| Every sample file follows this order: | ||
|
|
||
| 1. PEP 723 inline script metadata (if external dependencies are needed) |
Motivation & Context
Python needs a native way to inject messages into an active agent run while the function-calling loop is in progress, matching the scenario covered by the .NET message-injecting chat client without introducing a chat-client wrapper pattern into Python.
This enables tools or host code to enqueue additional session messages while a long-running tool is executing. The queued messages are drained into the next model call for the same
AgentSession, so the final answer can account for the new input without starting a separate run.This PR is draft because it depends on #6997. Until informational-only function call content is available, message injection treats all function-call content as actionable; a TODO marks the follow-up point where informational-only hosted/provider-executed calls should be ignored by the follow-up detection.
Description & Review Guide
MessageInjectionMiddleware,enqueue_messages(session, messages), andMESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEYto the Python core API.AgentSessionthroughChatContextso chat middleware can coordinate with session-scoped state without forwarding the session to leaf chat clients.python/samples/AGENTS.mdand removes the duplicatepython-samplesskill.sessionkwarg at the leaf layer._sessions.pywith other session-scoped chat middleware.enqueue_messagesandMESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY) are the right API shape.Related Issue
Fixes #6996.
Depends on #6997.
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.