From 11be6b80f408a51e10e15768f1a7e82bdb67c56d Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 25 Aug 2026 14:12:31 +0200 Subject: [PATCH] Python: keep agent loop marker from provider SDKs Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 42a18a81-57d5-4912-8dbd-f1a48db57a45 --- .../packages/core/agent_framework/_agents.py | 3 ++ .../core/tests/core/test_harness_loop.py | 51 ++++++++++++++++++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 288d8a29e3..f622442134 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1512,6 +1512,9 @@ async def _prepare_run_context( # _merge_options strips unset (None) options, so e.g. an unset `store` is not forwarded # and the service decides its own default. co = _merge_options(chat_options, run_opts) + # The loop marker must remain on SessionContext.options for after_run provider + # scoping, but it is framework-private metadata and must not reach the client. + co.pop(_LOOP_ITERATION_TOKEN_KEY, None) # Build session_messages from session context: context messages + input messages session_messages: list[Message] = session_context.get_messages(include_input=True) diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index cbf20c64df..5332f65cd2 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -32,9 +32,9 @@ background_tasks_running, background_tasks_running_message, set_agent_mode, - tool, todos_remaining, todos_remaining_message, + tool, ) from agent_framework._harness._loop import ( DEFAULT_JUDGE_MAX_ITERATIONS, @@ -1424,6 +1424,55 @@ async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict self.after_calls += 1 +class _StrictTransportChatClient(RecordingChatClient): + def __init__(self) -> None: + super().__init__() + self.received_provider_options: list[str] = [] + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool = False, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self._strict_transport(**options) + return super()._inner_get_response(messages=messages, stream=stream, options=options, **kwargs) + + def _strict_transport(self, *, tool_choice: Any, _provider_option: str) -> None: + self.received_provider_options.append(_provider_option) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) +async def test_loop_marker_is_not_forwarded_to_provider_transport(stream: bool) -> None: + client = _StrictTransportChatClient() + turn_scoped = _TurnScopedRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=1)], + context_providers=[turn_scoped], + ) + options = {"_provider_option": "kept"} + + if stream: + response = agent.run( # type: ignore[call-overload] # pyrefly: ignore[no-matching-overload] # ty: ignore[no-matching-overload] + "start", + stream=True, + options=options, # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + ) + _ = [update async for update in response] + await response.get_final_response() + else: + await agent.run( # type: ignore[call-overload] # pyrefly: ignore[no-matching-overload] # ty: ignore[no-matching-overload] + "start", + options=options, # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + ) + + assert client.received_provider_options == ["kept"] + assert turn_scoped.after_calls == 1 + + async def test_turn_scoped_after_run_fires_once_per_loop() -> None: client = RecordingChatClient() turn_scoped = _TurnScopedRecordingProvider()