Python: Add message injection middleware

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

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`)
 

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",

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",

python/packages/core/agent_framework/_agents.py

@@ -1410,23 +1410,14 @@ 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,
                 session=active_session,
                 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,

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,
         )