diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 4722699b7fa..7b35b182c5d 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2236,7 +2236,12 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse and update.finish_reason is not None ): response.finish_reason = update.finish_reason - response.continuation_token = update.continuation_token + if update.continuation_token is not None: + response.continuation_token = update.continuation_token + elif update.finish_reason is not None: + # A terminal update withdraws the token: its presence on the aggregate + # means the operation is still in progress. + response.continuation_token = None def _merge_function_call_content(message: Message, content: Content) -> None: diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index ff4bfedf014..9becac7515d 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -1889,6 +1889,40 @@ def test_agent_run_response_from_updates_uses_last_non_none_agent_id() -> None: assert response.agent_id == "source-agent" +def test_chat_response_from_updates_preserves_continuation_token_while_in_progress() -> None: + # A token stamped on an earlier update must survive later updates that do not + # repeat it: presence on the aggregate means the operation is still in progress. + response = ChatResponse.from_updates([ + ChatResponseUpdate( + contents=[Content.from_text(text="working")], + response_id="resp_1", + model="gpt-5", + continuation_token={"response_id": "resp_1"}, + ), + ChatResponseUpdate(contents=[Content.from_text(text="...still working")]), + ]) + + assert response.response_id == "resp_1" + assert response.model == "gpt-5" + assert response.continuation_token == {"response_id": "resp_1"} + + +def test_chat_response_from_updates_clears_continuation_token_on_terminal_update() -> None: + # The update that carries a finish_reason is terminal, so it withdraws the token + # from the aggregate: "None means the operation is complete". + response = ChatResponse.from_updates([ + ChatResponseUpdate( + contents=[Content.from_text(text="working")], + response_id="resp_1", + continuation_token={"response_id": "resp_1"}, + ), + ChatResponseUpdate(contents=[Content.from_text(text="done")], finish_reason="stop"), + ]) + + assert response.finish_reason == "stop" + assert response.continuation_token is None + + def test_agent_run_response_str_method(chat_message: Message) -> None: response = AgentResponse(messages=chat_message) assert str(response) == "Hello"