From d7540693e67dfdd4bb09ffa625ac872ea555aa0f Mon Sep 17 00:00:00 2001 From: MAF Dashboard Bot Date: Thu, 23 Apr 2026 19:05:37 +0000 Subject: [PATCH 1/3] Fix continuation_token and background leaking across tool-loop iterations (#5394) Strip continuation_token and background from mutable_options after a background response completes (continuation_token becomes None). Without this fix, subsequent tool-loop iterations in FunctionInvocationLayer would repeatedly retrieve the same completed response via GET instead of POSTing tool results, causing an infinite loop until max_iterations. Applied to both non-streaming and streaming tool loops. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/_tools.py | 12 +++ .../core/test_function_invocation_logic.py | 94 +++++++++++++++++++ 2 files changed, 106 insertions(+) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 3f15472a5a2..33551cbdc31 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2349,6 +2349,12 @@ async def _get_response() -> ChatResponse[Any]: options=mutable_options, ) + # Once a background response completes, strip polling/background + # options so subsequent tool-loop iterations POST results normally. + if response.continuation_token is None: + mutable_options.pop("continuation_token", None) + mutable_options.pop("background", None) + if response.conversation_id is not None: prepped_messages = [] @@ -2496,6 +2502,12 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: options=mutable_options, ) + # Once a background response completes, strip polling/background + # options so subsequent tool-loop iterations POST results normally. + if response.continuation_token is None: + mutable_options.pop("continuation_token", None) + mutable_options.pop("background", None) + if not any( item.type in ("function_call", "function_approval_request") for msg in response.messages diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index fe9a8145724..74d10c762cc 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3807,3 +3807,97 @@ def empty_request(task: str) -> str: ] assert len(function_results) >= 1 assert any("user input" in (fr.result or "").lower() for fr in function_results) + + +async def test_continuation_token_stripped_after_completed_response(chat_client_base: SupportsChatGetResponse): + """continuation_token from options must not persist across tool-loop iterations. + + When a background response completes (continuation_token becomes None), + subsequent tool-result submissions should POST normally instead of + retrieving the same completed response. + """ + recorded_options: list[dict[str, Any]] = [] + original_get = chat_client_base._get_non_streaming_response + + async def _tracking_get(*, messages, options, **kwargs): + recorded_options.append(dict(options)) + return await original_get(messages=messages, options=options, **kwargs) + + chat_client_base._get_non_streaming_response = _tracking_get + + @tool(name="lookup", approval_mode="never_require") + def lookup(query: str) -> str: + return f"found: {query}" + + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="lookup", arguments='{"query": "test"}'), + ], + ), + ), + ChatResponse(messages=Message(role="assistant", contents=["answer"])), + ] + + await chat_client_base.get_response( + [Message(role="user", contents=["find test"])], + options={ + "tools": [lookup], + "tool_choice": "auto", + "continuation_token": {"response_id": "resp_abc123"}, + }, + ) + + assert chat_client_base.call_count >= 2 + assert "continuation_token" in recorded_options[0], "First call should include continuation_token" + assert "continuation_token" not in recorded_options[1], ( + "continuation_token must be stripped after the background response completes" + ) + + +async def test_background_option_stripped_after_completed_response(chat_client_base: SupportsChatGetResponse): + """background=True from options must not persist across tool-loop iterations. + + Tool-result submissions should not start new background jobs. + """ + recorded_options: list[dict[str, Any]] = [] + original_get = chat_client_base._get_non_streaming_response + + async def _tracking_get(*, messages, options, **kwargs): + recorded_options.append(dict(options)) + return await original_get(messages=messages, options=options, **kwargs) + + chat_client_base._get_non_streaming_response = _tracking_get + + @tool(name="lookup", approval_mode="never_require") + def lookup(query: str) -> str: + return f"found: {query}" + + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="lookup", arguments='{"query": "test"}'), + ], + ), + ), + ChatResponse(messages=Message(role="assistant", contents=["answer"])), + ] + + await chat_client_base.get_response( + [Message(role="user", contents=["find test"])], + options={ + "tools": [lookup], + "tool_choice": "auto", + "background": True, + }, + ) + + assert chat_client_base.call_count >= 2 + assert recorded_options[0].get("background") is True, "First call should include background=True" + assert not recorded_options[1].get("background", False), ( + "background must be stripped after the background response completes" + ) From a9db85d2f1717000e5d6855ba9257db370922125 Mon Sep 17 00:00:00 2001 From: MAF Dashboard Bot Date: Thu, 23 Apr 2026 19:43:04 +0000 Subject: [PATCH 2/3] fix(#5394): tighten polling-option guard and add streaming/negative tests Address PR review feedback: - Tighten the guard condition in both non-streaming and streaming paths to only fire when continuation_token or background keys are actually present in mutable_options, making the intent explicit and avoiding unnecessary pop calls on every iteration. - Replace the two separate non-streaming tests with a single parameterized test to reduce duplication. - Add test for both continuation_token and background present simultaneously to verify both are stripped together. - Add negative test verifying options are preserved when the response has a non-None continuation_token (background job still in progress). - Add streaming-mode parameterized test covering the identical logic in the _stream() generator path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/_tools.py | 8 +- .../core/test_function_invocation_logic.py | 149 ++++++++++++++++-- 2 files changed, 142 insertions(+), 15 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 33551cbdc31..f2d105e2bf9 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2351,7 +2351,9 @@ async def _get_response() -> ChatResponse[Any]: # Once a background response completes, strip polling/background # options so subsequent tool-loop iterations POST results normally. - if response.continuation_token is None: + if response.continuation_token is None and ( + "continuation_token" in mutable_options or "background" in mutable_options + ): mutable_options.pop("continuation_token", None) mutable_options.pop("background", None) @@ -2504,7 +2506,9 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: # Once a background response completes, strip polling/background # options so subsequent tool-loop iterations POST results normally. - if response.continuation_token is None: + if response.continuation_token is None and ( + "continuation_token" in mutable_options or "background" in mutable_options + ): mutable_options.pop("continuation_token", None) mutable_options.pop("background", None) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 74d10c762cc..559f64db3f6 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3808,13 +3808,22 @@ def empty_request(task: str) -> str: assert len(function_results) >= 1 assert any("user input" in (fr.result or "").lower() for fr in function_results) - -async def test_continuation_token_stripped_after_completed_response(chat_client_base: SupportsChatGetResponse): - """continuation_token from options must not persist across tool-loop iterations. +@pytest.mark.parametrize( + "option_key,option_value", + [ + ("continuation_token", {"response_id": "resp_abc123"}), + ("background", True), + ], + ids=["continuation_token", "background"], +) +async def test_polling_option_stripped_after_completed_response( + chat_client_base: SupportsChatGetResponse, option_key: str, option_value: Any +): + """Polling/background options must not persist across tool-loop iterations. When a background response completes (continuation_token becomes None), subsequent tool-result submissions should POST normally instead of - retrieving the same completed response. + retrieving the same completed response or starting new background jobs. """ recorded_options: list[dict[str, Any]] = [] original_get = chat_client_base._get_non_streaming_response @@ -3841,27 +3850,72 @@ def lookup(query: str) -> str: ChatResponse(messages=Message(role="assistant", contents=["answer"])), ] + await chat_client_base.get_response( + [Message(role="user", contents=["find test"])], + options={ + "tools": [lookup], + "tool_choice": "auto", + option_key: option_value, + }, + ) + + assert chat_client_base.call_count >= 2 + assert option_key in recorded_options[0], f"First call should include {option_key}" + assert option_key not in recorded_options[1], ( + f"{option_key} must be stripped after the background response completes" + ) + + +async def test_both_polling_options_stripped_after_completed_response(chat_client_base: SupportsChatGetResponse): + """Both continuation_token and background are stripped together on completion.""" + recorded_options: list[dict[str, Any]] = [] + original_get = chat_client_base._get_non_streaming_response + + async def _tracking_get(*, messages, options, **kwargs): + recorded_options.append(dict(options)) + return await original_get(messages=messages, options=options, **kwargs) + + chat_client_base._get_non_streaming_response = _tracking_get + + @tool(name="lookup", approval_mode="never_require") + def lookup(query: str) -> str: + return f"found: {query}" + + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="lookup", arguments='{"query": "test"}'), + ], + ), + ), + ChatResponse(messages=Message(role="assistant", contents=["answer"])), + ] + await chat_client_base.get_response( [Message(role="user", contents=["find test"])], options={ "tools": [lookup], "tool_choice": "auto", "continuation_token": {"response_id": "resp_abc123"}, + "background": True, }, ) assert chat_client_base.call_count >= 2 - assert "continuation_token" in recorded_options[0], "First call should include continuation_token" + assert "continuation_token" in recorded_options[0] + assert recorded_options[0].get("background") is True assert "continuation_token" not in recorded_options[1], ( "continuation_token must be stripped after the background response completes" ) + assert not recorded_options[1].get("background", False), ( + "background must be stripped after the background response completes" + ) -async def test_background_option_stripped_after_completed_response(chat_client_base: SupportsChatGetResponse): - """background=True from options must not persist across tool-loop iterations. - - Tool-result submissions should not start new background jobs. - """ +async def test_polling_options_preserved_while_background_in_progress(chat_client_base: SupportsChatGetResponse): + """Options are preserved when continuation_token is non-None (background job still in progress).""" recorded_options: list[dict[str, Any]] = [] original_get = chat_client_base._get_non_streaming_response @@ -3876,6 +3930,7 @@ def lookup(query: str) -> str: return f"found: {query}" chat_client_base.run_responses = [ + # First response: still in progress (non-None continuation_token) with a function call ChatResponse( messages=Message( role="assistant", @@ -3883,7 +3938,9 @@ def lookup(query: str) -> str: Content.from_function_call(call_id="c1", name="lookup", arguments='{"query": "test"}'), ], ), + continuation_token={"response_id": "resp_in_progress"}, ), + # Second response: completed (continuation_token defaults to None) ChatResponse(messages=Message(role="assistant", contents=["answer"])), ] @@ -3892,12 +3949,78 @@ def lookup(query: str) -> str: options={ "tools": [lookup], "tool_choice": "auto", + "continuation_token": {"response_id": "resp_abc123"}, "background": True, }, ) assert chat_client_base.call_count >= 2 - assert recorded_options[0].get("background") is True, "First call should include background=True" - assert not recorded_options[1].get("background", False), ( - "background must be stripped after the background response completes" + # After the first response with non-None continuation_token, options should be preserved + assert "continuation_token" in recorded_options[1], ( + "continuation_token should be preserved while background job is in progress" + ) + assert recorded_options[1].get("background") is True, ( + "background should be preserved while background job is in progress" + ) + + +@pytest.mark.parametrize( + "option_key,option_value", + [ + ("continuation_token", {"response_id": "resp_abc123"}), + ("background", True), + ], + ids=["continuation_token", "background"], +) +async def test_polling_option_stripped_after_completed_streaming_response( + chat_client_base: SupportsChatGetResponse, option_key: str, option_value: Any +): + """Streaming path: polling/background options must not persist across tool-loop iterations.""" + recorded_options: list[dict[str, Any]] = [] + original_stream = chat_client_base._get_streaming_response + + def _tracking_stream(*, messages, options, **kwargs): + recorded_options.append(dict(options)) + return original_stream(messages=messages, options=options, **kwargs) + + chat_client_base._get_streaming_response = _tracking_stream + + @tool(name="lookup", approval_mode="never_require") + def lookup(query: str) -> str: + return f"found: {query}" + + chat_client_base.streaming_responses = [ + # First streaming round: function call + [ + ChatResponseUpdate( + contents=[Content.from_function_call(call_id="c1", name="lookup", arguments='{"query": "test"}')], + role="assistant", + finish_reason="function_call", + ), + ], + # Second streaming round: final answer + [ + ChatResponseUpdate( + contents=[Content.from_text("answer")], + role="assistant", + finish_reason="stop", + ), + ], + ] + + async for _ in chat_client_base.get_response( + [Message(role="user", contents=["find test"])], + options={ + "tools": [lookup], + "tool_choice": "auto", + option_key: option_value, + }, + stream=True, + ): + pass + + assert chat_client_base.call_count >= 2 + assert option_key in recorded_options[0], f"First call should include {option_key}" + assert option_key not in recorded_options[1], ( + f"{option_key} must be stripped after the background streaming response completes" ) From c33ec18ae94b050922f73d66a130583346e57dc5 Mon Sep 17 00:00:00 2001 From: MAF Dashboard Bot Date: Thu, 23 Apr 2026 19:57:20 +0000 Subject: [PATCH 3/3] Address review feedback for #5394: Python: [Bug]: `background=True` causes infinite tool-call loop, tool-result submissions inherit background mode --- .../packages/core/tests/core/test_function_invocation_logic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 559f64db3f6..e08646951b0 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3808,6 +3808,7 @@ def empty_request(task: str) -> str: assert len(function_results) >= 1 assert any("user input" in (fr.result or "").lower() for fr in function_results) + @pytest.mark.parametrize( "option_key,option_value", [