diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 3f15472a5a2..f2d105e2bf9 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2349,6 +2349,14 @@ 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 and ( + "continuation_token" in mutable_options or "background" in mutable_options + ): + mutable_options.pop("continuation_token", None) + mutable_options.pop("background", None) + if response.conversation_id is not None: prepped_messages = [] @@ -2496,6 +2504,14 @@ 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 and ( + "continuation_token" in mutable_options or "background" in mutable_options + ): + 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..e08646951b0 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,221 @@ 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", + [ + ("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 or starting 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", + 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] + 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_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 + + 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 = [ + # First response: still in progress (non-None continuation_token) with a function call + ChatResponse( + messages=Message( + role="assistant", + contents=[ + 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"])), + ] + + 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 + # 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" + )