From bf7f2e7138145601576a7f74a726b7026f3b5835 Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Fri, 2 Oct 2026 18:00:41 -0700 Subject: [PATCH 1/7] Python: add terminal state check for continuation token --- .../a2a/agent_framework_a2a/_agent.py | 55 +++++++++-- python/packages/a2a/tests/test_a2a_agent.py | 97 +++++++++++++++++++ .../packages/core/agent_framework/_types.py | 28 +++++- .../core/agent_framework/_workflows/_agent.py | 2 + python/packages/core/tests/core/test_types.py | 84 ++++++++++++++++ .../tests/workflow/test_workflow_agent.py | 5 +- .../agent_framework_openai/_chat_client.py | 3 + .../tests/openai/test_openai_chat_client.py | 41 ++++++++ 8 files changed, 304 insertions(+), 11 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index c5780b1ed5f..dde63c99c04 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -680,18 +680,21 @@ async def _map_a2a_stream( all_updates.append(update) yield update elif is_terminal or is_input_required: - if updates: + if updates and any(update.contents for update in updates): # Terminal/input-required event with content — discard accumulated intermediates pending_updates_by_task.pop(status_event.task_id, None) for update in updates: all_updates.append(update) yield update elif is_terminal: - # Terminal event with NO content — flush accumulated updates + # Terminal event with no content — flush accumulated updates before its metadata marker. pending = pending_updates_by_task.pop(status_event.task_id, []) for update in pending: all_updates.append(update) yield update + for update in updates: + all_updates.append(update) + yield update else: # Non-streaming intermediate: accumulate for later if updates: @@ -840,7 +843,7 @@ def _updates_from_task( if getattr(message.raw_representation, "artifact_id", None) not in streamed_artifact_ids ] if task_messages: - return [ + updates = [ AgentResponseUpdate( contents=message.contents, role=message.role, @@ -853,13 +856,25 @@ def _updates_from_task( ) for message in task_messages ] + updates[-1].is_operation_terminal = True + return updates if task.artifacts: - return [] + return [ + AgentResponseUpdate( + contents=[], + role="assistant", + response_id=task.id, + is_operation_terminal=True, + additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, + raw_representation=task, + ) + ] return [ AgentResponseUpdate( contents=[], role="assistant", response_id=task.id, + is_operation_terminal=True, additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, raw_representation=task, ) @@ -962,18 +977,39 @@ def _updates_from_task_update_event( ) ] - if not update_event.status.HasField("message") or not update_event.status.message.parts: - return [] - if state not in TERMINAL_TASK_STATES: return [] + event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {} + if not update_event.status.HasField("message") or not update_event.status.message.parts: + return [ + AgentResponseUpdate( + contents=[], + role="assistant", + response_id=update_event.task_id, + is_operation_terminal=True, + continuation_token=continuation_token, + additional_properties={"a2a_metadata": event_meta} if event_meta else None, + raw_representation=update_event, + ) + ] + message = update_event.status.message contents = self._parse_contents_from_a2a(message.parts) if not contents: - return [] + return [ + AgentResponseUpdate( + contents=[], + role="assistant" if message.role == A2ARole.ROLE_AGENT else "user", + response_id=update_event.task_id, + message_id=message.message_id, + is_operation_terminal=True, + continuation_token=continuation_token, + additional_properties={"a2a_metadata": event_meta} if event_meta else None, + raw_representation=update_event, + ) + ] msg_meta = MessageToDict(message.metadata) if message.metadata else {} - event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {} merged_metadata = {**msg_meta, **event_meta} or None return [ @@ -982,6 +1018,7 @@ def _updates_from_task_update_event( role="assistant" if message.role == A2ARole.ROLE_AGENT else "user", response_id=update_event.task_id, message_id=message.message_id, + is_operation_terminal=True, continuation_token=continuation_token, additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, raw_representation=update_event, diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 533e2e2a667..c130ec82801 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -1142,6 +1142,101 @@ async def test_resume_streaming_via_continuation_token(a2a_agent: A2AAgent, mock assert updates[1].contents[0].text == "Stream resumed" +@mark.parametrize( + "terminal_state", + [ + TaskState.TASK_STATE_COMPLETED, + TaskState.TASK_STATE_FAILED, + TaskState.TASK_STATE_CANCELED, + TaskState.TASK_STATE_REJECTED, + ], +) +async def test_resume_streaming_terminal_task_clears_final_continuation_token( + a2a_agent: A2AAgent, + mock_a2a_client: MockA2AClient, + terminal_state: TaskState, +) -> None: + """A terminal A2A task should clear the aggregated continuation token.""" + working_task = Task( + id="task-rs-terminal", + context_id="ctx-rs-terminal", + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + completed_task = Task( + id="task-rs-terminal", + context_id="ctx-rs-terminal", + status=TaskStatus(state=terminal_state), + artifacts=[ + Artifact( + artifact_id="art-rs-terminal", + name="result", + parts=[Part(text="Stream completed")], + ) + ], + ) + mock_a2a_client.subscribe_responses.extend([ + StreamResponse(task=working_task), + StreamResponse(task=completed_task), + ]) + + stream = a2a_agent.run( + stream=True, + continuation_token=A2AContinuationToken( + task_id="task-rs-terminal", + context_id="ctx-rs-terminal", + ), + background=True, + ) + updates = [update async for update in stream] + response = await stream.get_final_response() + + assert len(updates) == 2 + assert updates[0].continuation_token is not None + assert updates[1].continuation_token is None + assert updates[1].is_operation_terminal is True + assert response.continuation_token is None + assert response.finish_reason is None + + +async def test_streaming_terminal_status_without_message_clears_final_continuation_token( + a2a_agent: A2AAgent, + mock_a2a_client: MockA2AClient, +) -> None: + """A message-less terminal status update should clear the aggregated continuation token.""" + working_task = Task( + id="task-status-terminal", + context_id="ctx-status-terminal", + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + completed_event = TaskStatusUpdateEvent( + task_id="task-status-terminal", + context_id="ctx-status-terminal", + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED), + ) + mock_a2a_client.subscribe_responses.extend([ + StreamResponse(task=working_task), + StreamResponse(status_update=completed_event), + ]) + + stream = a2a_agent.run( + stream=True, + continuation_token=A2AContinuationToken( + task_id="task-status-terminal", + context_id="ctx-status-terminal", + ), + background=True, + ) + updates = [update async for update in stream] + response = await stream.get_final_response() + + assert len(updates) == 2 + assert updates[0].continuation_token is not None + assert updates[1].contents == [] + assert updates[1].is_operation_terminal is True + assert response.continuation_token is None + assert response.finish_reason is None + + async def test_poll_task_in_progress(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None: """Test poll_task returns continuation token when task is still in progress.""" status = TaskStatus(state=TaskState.TASK_STATE_WORKING, message=None) @@ -1862,6 +1957,8 @@ async def test_streaming_artifact_update_event_does_not_duplicate_terminal_task_ response = await stream.get_final_response() assert [update.text for update in updates] == ["Hello ", "world"] + assert updates[-1].contents == [] + assert updates[-1].is_operation_terminal is True assert response.text == "Hello world" assert len(response.messages) == 1 diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 4e24e79ea0b..3fe4627e9bc 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2279,7 +2279,10 @@ 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.is_operation_terminal: + response.continuation_token = None + elif update.continuation_token is not None: + response.continuation_token = update.continuation_token def _apply_response_tail_to_update( @@ -2927,6 +2930,7 @@ class ChatResponseUpdate(SerializationMixin): model: The model associated with this response update. created_at: A timestamp for the chat response update. finish_reason: The finish reason for the operation. + is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Any additional properties associated with the chat response update. raw_representation: The raw representation of the chat response update from an underlying implementation. @@ -2973,6 +2977,7 @@ def __init__( created_at: CreatedAtT | None = None, finish_reason: FinishReasonLiteral | FinishReason | None = None, continuation_token: ContinuationToken | None = None, + is_operation_terminal: bool = False, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, ) -> None: @@ -2990,6 +2995,7 @@ def __init__( finish_reason: Optional finish reason for the operation. continuation_token: Optional token for resuming a long-running background operation. When present, indicates the operation is still in progress. + is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Optional additional properties associated with the chat response update. raw_representation: Optional raw representation of the chat response update from an underlying implementation. @@ -3018,12 +3024,20 @@ def __init__( self.created_at = created_at self.finish_reason = finish_reason self.continuation_token = continuation_token + self.is_operation_terminal = is_operation_terminal self.additional_properties = _restore_compaction_annotation_in_additional_properties( additional_properties, allow_none=True, ) self.raw_representation = raw_representation + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Convert the update to a dictionary.""" + result = super().to_dict(exclude=exclude, exclude_none=exclude_none) + if not self.is_operation_terminal: + result.pop("is_operation_terminal", None) + return result + @property def text(self) -> str: """Returns the concatenated text of all contents in the update.""" @@ -3378,6 +3392,7 @@ class AgentResponseUpdate(SerializationMixin): response_id: The ID of the response of which this update is a part. message_id: The ID of the message of which this update is a part. created_at: A timestamp for the response update. + is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Any additional properties associated with the update. raw_representation: The raw representation from an underlying implementation. @@ -3425,6 +3440,7 @@ def __init__( created_at: CreatedAtT | None = None, finish_reason: FinishReasonLiteral | FinishReason | None = None, continuation_token: ContinuationToken | None = None, + is_operation_terminal: bool = False, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, ) -> None: @@ -3444,6 +3460,7 @@ def __init__( ``"tool_calls"`` (the model invoked a tool). continuation_token: Optional token for resuming a long-running background operation. When present, indicates the operation is still in progress. + is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Optional additional properties associated with the chat response update. raw_representation: Optional raw representation of the chat response update. @@ -3470,12 +3487,20 @@ def __init__( self.created_at = created_at self.finish_reason = finish_reason self.continuation_token = continuation_token + self.is_operation_terminal = is_operation_terminal self.additional_properties = _restore_compaction_annotation_in_additional_properties( additional_properties, allow_none=True, ) self.raw_representation: Any | list[Any] | None = raw_representation + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Convert the update to a dictionary.""" + result = super().to_dict(exclude=exclude, exclude_none=exclude_none) + if not self.is_operation_terminal: + result.pop("is_operation_terminal", None) + return result + @property def text(self) -> str: """Get the concatenated text of all TextContent objects in contents.""" @@ -3503,6 +3528,7 @@ def map_chat_to_agent_update(update: ChatResponseUpdate, agent_name: str | None) created_at=update.created_at, finish_reason=update.finish_reason, # type: ignore[arg-type] continuation_token=update.continuation_token, + is_operation_terminal=update.is_operation_terminal, additional_properties=update.additional_properties, raw_representation=update, ) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 6d01c8366a4..f2870ffa8f5 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -669,6 +669,7 @@ def _convert_workflow_event_to_agent_response_updates( # connectors may set any string); forward the value unchanged. finish_reason=cast(FinishReasonLiteral | FinishReason | None, data.finish_reason), continuation_token=data.continuation_token, + is_operation_terminal=data.is_operation_terminal, additional_properties=dict(data.additional_properties) if data.additional_properties is not None else None, @@ -695,6 +696,7 @@ def _convert_workflow_event_to_agent_response_updates( updates[-1].agent_id = data.agent_id updates[-1].finish_reason = data.finish_reason updates[-1].continuation_token = data.continuation_token + updates[-1].is_operation_terminal = data.continuation_token is None updates[-1].additional_properties = dict(data.additional_properties) return updates if isinstance(data, Message): diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 1f1bd2fc313..9fffcb86903 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -6206,6 +6206,90 @@ def test_agent_response_update_serialization_includes_finish_reason() -> None: assert data["finish_reason"] == "tool_calls" +async def test_chat_client_agent_streaming_continuation_token() -> None: + """Chat response aggregation should preserve an in-progress continuation token.""" + events = [ + 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")]), + ] + + aggregated = ChatResponse.from_updates(events) + + assert aggregated.continuation_token == {"response_id": "resp_1"} + + +def test_chat_response_finished_update_clears_continuation_token() -> None: + """A finished chat response should clear its continuation token.""" + events = [ + ChatResponseUpdate(continuation_token={"response_id": "resp_1"}), + ChatResponseUpdate(is_operation_terminal=True), + ] + + aggregated = ChatResponse.from_updates(events) + + assert aggregated.continuation_token is None + + +def test_chat_response_finish_reason_does_not_clear_continuation_token() -> None: + """A finish reason alone should not be treated as the terminal streaming signal.""" + events = [ + ChatResponseUpdate(continuation_token={"response_id": "resp_1"}), + ChatResponseUpdate(finish_reason="stop"), + ] + + aggregated = ChatResponse.from_updates(events) + + assert aggregated.continuation_token == {"response_id": "resp_1"} + assert aggregated.finish_reason == "stop" + + +def test_agent_response_update_preserves_continuation_token() -> None: + """Agent response aggregation should preserve an in-progress continuation token.""" + events = [ + AgentResponseUpdate(continuation_token={"response_id": "resp_1"}), + AgentResponseUpdate(contents=[Content.from_text(text="...still working")]), + ] + + aggregated = AgentResponse.from_updates(events) + + assert aggregated.continuation_token == {"response_id": "resp_1"} + + +def test_agent_response_finished_update_clears_continuation_token() -> None: + """A finished agent response should clear its continuation token.""" + events = [ + AgentResponseUpdate(continuation_token={"response_id": "resp_1"}), + AgentResponseUpdate(is_operation_terminal=True), + ] + + aggregated = AgentResponse.from_updates(events) + + assert aggregated.continuation_token is None + + +def test_map_chat_to_agent_update_preserves_terminal_signal() -> None: + """Chat-to-agent update conversion should preserve terminality.""" + update = map_chat_to_agent_update(ChatResponseUpdate(is_operation_terminal=True), agent_name=None) + + assert update.is_operation_terminal is True + + +@pytest.mark.parametrize("update_type", [ChatResponseUpdate, AgentResponseUpdate]) +def test_response_update_serialization_omits_false_terminal_signal( + update_type: type[ChatResponseUpdate] | type[AgentResponseUpdate], +) -> None: + """Nonterminal updates should retain their existing serialized shape.""" + assert "is_operation_terminal" not in update_type().to_dict() + serialized = update_type(is_operation_terminal=True).to_dict() + assert serialized["is_operation_terminal"] is True + assert update_type.from_dict(serialized).is_operation_terminal is True + + # endregion diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 0ef70ca2f6e..81b5509d682 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -863,7 +863,7 @@ async def yielding_executor(messages: list[Message], ctx: WorkflowContext[Never, assert "second output" in texts async def test_workflow_as_agent_stream_preserves_response_update_metadata(self) -> None: - """Test that streaming forwards finish_reason, continuation_token and additional_properties. + """Test that streaming forwards finish_reason, continuation_token, terminality, and additional properties. This validates the fix for issue #7952: AgentResponseUpdate metadata should be forwarded as-is when the workflow is wrapped via .as_agent(). @@ -880,6 +880,7 @@ async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, message_id="source-message", finish_reason="stop", continuation_token=cast(Any, {"token": "resume-token"}), + is_operation_terminal=False, additional_properties={"provider_marker": "preserve-me"}, ) ) @@ -898,6 +899,7 @@ async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, assert update.agent_id == "source-agent" assert update.finish_reason == "stop" assert update.continuation_token == {"token": "resume-token"} + assert update.is_operation_terminal is False assert update.additional_properties == {"provider_marker": "preserve-me"} async def test_workflow_as_agent_stream_preserves_response_metadata(self) -> None: @@ -934,6 +936,7 @@ async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, assert updates[-1].agent_id == "source-agent" assert updates[-1].finish_reason == "length" assert updates[-1].continuation_token == {"token": "response-resume-token"} + assert updates[-1].is_operation_terminal is False assert updates[-1].additional_properties == {"provider_marker": "preserve-response"} assert final_response.agent_id == "source-agent" assert final_response.finish_reason == "length" diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 24ca5066975..eba70784660 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -3649,6 +3649,7 @@ def _parse_chunk_from_openai( created_at: str | None = None continuation_token: OpenAIContinuationToken | None = None finish_reason: FinishReason | None = None + is_operation_terminal = False model = self.model def output_text_properties(output: Any) -> dict[str, Any] | None: @@ -3892,6 +3893,7 @@ def parse_annotations( conversation_id = self._get_conversation_id(event.response, options.get("store")) continuation_token = OpenAIContinuationToken(response_id=event.response.id) case "response.completed" | "response.incomplete" | "response.failed": + is_operation_terminal = True response_id = event.response.id conversation_id = self._get_conversation_id(event.response, options.get("store")) model = event.response.model @@ -4201,6 +4203,7 @@ def parse_annotations( created_at=created_at, continuation_token=continuation_token, finish_reason=finish_reason, + is_operation_terminal=is_operation_terminal, additional_properties=metadata, raw_representation=event, ) diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index bd576e61321..f5aa728d707 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -9751,6 +9751,47 @@ def test_streaming_terminal_response_sets_finish_reason( update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids={}) assert update.finish_reason == expected_finish_reason + assert update.is_operation_terminal is True + + +@pytest.mark.parametrize( + ("event_type", "status", "incomplete_reason"), + [ + ("response.failed", "failed", None), + ("response.incomplete", "incomplete", "other"), + ], +) +def test_streaming_terminal_response_clears_aggregated_continuation_token( + event_type: str, + status: str, + incomplete_reason: str | None, +) -> None: + """Terminal Responses API events should clear an earlier continuation token.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + in_progress_event = MagicMock() + in_progress_event.type = "response.in_progress" + in_progress_event.response.id = "resp_terminal" + in_progress_event.response.conversation = None + + terminal_event = MagicMock() + terminal_event.type = event_type + terminal_event.response.id = "resp_terminal" + terminal_event.response.conversation = None + terminal_event.response.model = "test-model" + terminal_event.response.created_at = 1000000000 + terminal_event.response.usage = None + terminal_event.response.status = status + terminal_event.response.incomplete_details = ( + MagicMock(reason=incomplete_reason) if incomplete_reason is not None else None + ) + terminal_event.response.output = [] + + response = ChatResponse.from_updates([ + client._parse_chunk_from_openai(in_progress_event, options={}, function_call_ids={}), + client._parse_chunk_from_openai(terminal_event, options={}, function_call_ids={}), + ]) + + assert response.continuation_token is None def test_map_chat_to_agent_update_preserves_continuation_token() -> None: From 9dd12d440d16b695f145191dcd880c3bd82d0760 Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Fri, 2 Oct 2026 18:48:27 -0700 Subject: [PATCH 2/7] Python: update continuation token in background --- .../a2a/agent_framework_a2a/_agent.py | 6 ++ python/packages/a2a/tests/test_a2a_agent.py | 69 ++++++++++++++++++- 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index dde63c99c04..bb4b004b3f4 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -859,6 +859,8 @@ def _updates_from_task( updates[-1].is_operation_terminal = True return updates if task.artifacts: + if not background: + return [] return [ AgentResponseUpdate( contents=[], @@ -982,6 +984,8 @@ def _updates_from_task_update_event( event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {} if not update_event.status.HasField("message") or not update_event.status.message.parts: + if not background: + return [] return [ AgentResponseUpdate( contents=[], @@ -997,6 +1001,8 @@ def _updates_from_task_update_event( message = update_event.status.message contents = self._parse_contents_from_a2a(message.parts) if not contents: + if not background: + return [] return [ AgentResponseUpdate( contents=[], diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index c130ec82801..3c350ef3ed0 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -1957,12 +1957,77 @@ async def test_streaming_artifact_update_event_does_not_duplicate_terminal_task_ response = await stream.get_final_response() assert [update.text for update in updates] == ["Hello ", "world"] - assert updates[-1].contents == [] - assert updates[-1].is_operation_terminal is True assert response.text == "Hello world" assert len(response.messages) == 1 +async def test_streaming_background_artifacts_emit_terminal_marker_without_duplication( + a2a_agent: A2AAgent, + mock_a2a_client: MockA2AClient, +) -> None: + """A background artifact stream should emit a terminal marker without duplicating content.""" + working_task = Task( + id="task-art-background", + context_id="ctx-art-background", + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + first_chunk = TaskArtifactUpdateEvent( + task_id="task-art-background", + context_id="ctx-art-background", + artifact=Artifact( + artifact_id="artifact-background", + parts=[Part(text="Hello ")], + ), + append=False, + ) + second_chunk = TaskArtifactUpdateEvent( + task_id="task-art-background", + context_id="ctx-art-background", + artifact=Artifact( + artifact_id="artifact-background", + parts=[Part(text="world")], + ), + append=True, + ) + terminal_task = Task( + id="task-art-background", + context_id="ctx-art-background", + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED), + artifacts=[ + Artifact( + artifact_id="artifact-background", + parts=[Part(text="Hello world")], + ) + ], + ) + + mock_a2a_client.responses.extend([ + StreamResponse(task=working_task), + StreamResponse(artifact_update=first_chunk), + StreamResponse(artifact_update=second_chunk), + StreamResponse(task=terminal_task), + ]) + + stream = a2a_agent.run("Hello", stream=True, background=True) + updates = [update async for update in stream] + response = await stream.get_final_response() + + assert len(updates) == 4 + + assert updates[0].contents == [] + assert updates[0].continuation_token is not None + assert updates[0].is_operation_terminal is False + + assert [update.text for update in updates[1:3]] == ["Hello ", "world"] + + assert updates[3].contents == [] + assert updates[3].continuation_token is None + assert updates[3].is_operation_terminal is True + + assert response.text == "Hello world" + assert response.continuation_token is None + + async def test_streaming_terminal_task_artifacts_are_emitted_when_terminal_event_has_no_content( a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient ) -> None: From df1a3e8ee5b680a9b6e81a5da426e0475162b219 Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Wed, 7 Oct 2026 17:12:56 -0700 Subject: [PATCH 3/7] replace is_operation_terminal with _OperationState --- .../a2a/agent_framework_a2a/_agent.py | 212 ++++++++++-------- python/packages/a2a/tests/test_a2a_agent.py | 9 +- .../packages/core/agent_framework/_types.py | 68 +++--- .../core/agent_framework/_workflows/_agent.py | 42 ++-- python/packages/core/tests/core/test_types.py | 195 ++++++++++++---- .../tests/workflow/test_workflow_agent.py | 52 ++--- .../agent_framework_openai/_chat_client.py | 10 +- .../tests/openai/test_openai_chat_client.py | 25 ++- 8 files changed, 395 insertions(+), 218 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index bb4b004b3f4..4afdd0cd3bc 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -44,7 +44,7 @@ prepend_agent_framework_to_user_agent, ) from agent_framework._telemetry import mark_feature_used -from agent_framework._types import AgentRunInputs +from agent_framework._types import AgentRunInputs, _set_operation_state # pyright: ignore[reportPrivateUsage] from agent_framework.exceptions import AgentInvalidRequestException from agent_framework.observability import AgentTelemetryLayer from google.protobuf.json_format import MessageToDict @@ -66,6 +66,15 @@ class A2AServiceSessionId(TypedDict): task_state: TaskState | None +def _with_operation_state( + update: AgentResponseUpdate, + state: Literal["in_progress", "terminal"], +) -> AgentResponseUpdate: + """Attach framework-internal resumable-operation state.""" + _set_operation_state(update, state) + return update + + class A2AAgentSession(AgentSession): """Session for A2A-based agents. @@ -705,7 +714,7 @@ async def _map_a2a_stream( last_task_id = artifact_event.task_id if artifact_event.context_id: last_context_id = artifact_event.context_id - updates = self._updates_from_task_update_event(artifact_event) + updates = self._updates_from_task_update_event(artifact_event, background=background) # Always yield artifact updates — they carry actual response # content (files, data). Track IDs so that a subsequent # terminal Task doesn't duplicate the same artifacts. @@ -817,22 +826,21 @@ def _updates_from_task( if status.state == TaskState.TASK_STATE_INPUT_REQUIRED: message = status.message if status.HasField("message") and status.message.parts else None occurrence_id = input_request_occurrence_id or task.id - return [ - AgentResponseUpdate( - contents=[ - self._input_required_request( - task.id, - message, - fallback_occurrence_id=occurrence_id, - ) - ], - role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user", - response_id=task.id, - continuation_token=self._build_continuation_token(task) if background else None, - additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, - raw_representation=task, - ) - ] + update = AgentResponseUpdate( + contents=[ + self._input_required_request( + task.id, + message, + fallback_occurrence_id=occurrence_id, + ) + ], + role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user", + response_id=task.id, + continuation_token=self._build_continuation_token(task) if background else None, + additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, + raw_representation=task, + ) + return [_with_operation_state(update, "in_progress")] if background else [update] if status.state in TERMINAL_TASK_STATES: task_messages = self._parse_messages_from_task(task) @@ -856,42 +864,49 @@ def _updates_from_task( ) for message in task_messages ] - updates[-1].is_operation_terminal = True + _set_operation_state(updates[-1], "terminal") return updates if task.artifacts: if not background: return [] return [ + _with_operation_state( + AgentResponseUpdate( + contents=[], + role="assistant", + response_id=task.id, + additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, + raw_representation=task, + ), + "terminal", + ) + ] + return [ + _with_operation_state( AgentResponseUpdate( contents=[], role="assistant", response_id=task.id, - is_operation_terminal=True, additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, raw_representation=task, - ) - ] - return [ - AgentResponseUpdate( - contents=[], - role="assistant", - response_id=task.id, - is_operation_terminal=True, - additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, - raw_representation=task, + ), + "terminal", ) ] if background and status.state in IN_PROGRESS_TASK_STATES: token = self._build_continuation_token(task) return [ - AgentResponseUpdate( - contents=[], - role="assistant", - response_id=task.id, - continuation_token=token, - additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, - raw_representation=task, + _with_operation_state( + AgentResponseUpdate( + contents=[], + role="assistant", + response_id=task.id, + continuation_token=token, + additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, + raw_representation=task, + ), + "in_progress", ) ] @@ -905,7 +920,18 @@ def _updates_from_task( contents = self._parse_contents_from_a2a(status.message.parts) if contents: return [ - AgentResponseUpdate( + _with_operation_state( + AgentResponseUpdate( + contents=contents, + role="assistant" if status.message.role == A2ARole.ROLE_AGENT else "user", + response_id=task.id, + additional_properties={"a2a_metadata": task_metadata} if task_metadata else None, + raw_representation=task, + ), + "in_progress", + ) + if background + else AgentResponseUpdate( contents=contents, role="assistant" if status.message.role == A2ARole.ROLE_AGENT else "user", response_id=task.id, @@ -931,16 +957,15 @@ def _updates_from_task_update_event( artifact_meta = MessageToDict(update_event.artifact.metadata) if update_event.artifact.metadata else {} event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {} merged_metadata = {**artifact_meta, **event_meta} or None - return [ - AgentResponseUpdate( - contents=contents, - role="assistant", - response_id=update_event.task_id, - message_id=update_event.artifact.artifact_id, - additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, - raw_representation=update_event, - ) - ] + update = AgentResponseUpdate( + contents=contents, + role="assistant", + response_id=update_event.task_id, + message_id=update_event.artifact.artifact_id, + additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, + raw_representation=update_event, + ) + return [_with_operation_state(update, "in_progress")] if background else [update] if not isinstance(update_event, TaskStatusUpdateEvent): return [] @@ -961,23 +986,22 @@ def _updates_from_task_update_event( event_meta = MessageToDict(update_event.metadata) if update_event.metadata else {} merged_metadata = {**message_meta, **event_meta} or None occurrence_id = input_request_occurrence_id or update_event.task_id - return [ - AgentResponseUpdate( - contents=[ - self._input_required_request( - update_event.task_id, - message, - fallback_occurrence_id=occurrence_id, - ) - ], - role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user", - response_id=update_event.task_id, - message_id=message.message_id if message is not None else None, - continuation_token=continuation_token, - additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, - raw_representation=update_event, - ) - ] + update = AgentResponseUpdate( + contents=[ + self._input_required_request( + update_event.task_id, + message, + fallback_occurrence_id=occurrence_id, + ) + ], + role="assistant" if message is None or message.role == A2ARole.ROLE_AGENT else "user", + response_id=update_event.task_id, + message_id=message.message_id if message is not None else None, + continuation_token=continuation_token, + additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, + raw_representation=update_event, + ) + return [_with_operation_state(update, "in_progress")] if background else [update] if state not in TERMINAL_TASK_STATES: return [] @@ -987,14 +1011,16 @@ def _updates_from_task_update_event( if not background: return [] return [ - AgentResponseUpdate( - contents=[], - role="assistant", - response_id=update_event.task_id, - is_operation_terminal=True, - continuation_token=continuation_token, - additional_properties={"a2a_metadata": event_meta} if event_meta else None, - raw_representation=update_event, + _with_operation_state( + AgentResponseUpdate( + contents=[], + role="assistant", + response_id=update_event.task_id, + continuation_token=continuation_token, + additional_properties={"a2a_metadata": event_meta} if event_meta else None, + raw_representation=update_event, + ), + "terminal", ) ] @@ -1004,30 +1030,34 @@ def _updates_from_task_update_event( if not background: return [] return [ - AgentResponseUpdate( - contents=[], - role="assistant" if message.role == A2ARole.ROLE_AGENT else "user", - response_id=update_event.task_id, - message_id=message.message_id, - is_operation_terminal=True, - continuation_token=continuation_token, - additional_properties={"a2a_metadata": event_meta} if event_meta else None, - raw_representation=update_event, + _with_operation_state( + AgentResponseUpdate( + contents=[], + role="assistant" if message.role == A2ARole.ROLE_AGENT else "user", + response_id=update_event.task_id, + message_id=message.message_id, + continuation_token=continuation_token, + additional_properties={"a2a_metadata": event_meta} if event_meta else None, + raw_representation=update_event, + ), + "terminal", ) ] msg_meta = MessageToDict(message.metadata) if message.metadata else {} merged_metadata = {**msg_meta, **event_meta} or None return [ - AgentResponseUpdate( - contents=contents, - role="assistant" if message.role == A2ARole.ROLE_AGENT else "user", - response_id=update_event.task_id, - message_id=message.message_id, - is_operation_terminal=True, - continuation_token=continuation_token, - additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, - raw_representation=update_event, + _with_operation_state( + AgentResponseUpdate( + contents=contents, + role="assistant" if message.role == A2ARole.ROLE_AGENT else "user", + response_id=update_event.task_id, + message_id=message.message_id, + continuation_token=continuation_token, + additional_properties={"a2a_metadata": merged_metadata} if merged_metadata else None, + raw_representation=update_event, + ), + "terminal", ) ] diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 3c350ef3ed0..d578f1af039 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -29,6 +29,7 @@ Message, SessionContext, ) +from agent_framework._types import _get_operation_state from agent_framework.a2a import A2AAgent from agent_framework.exceptions import AgentInvalidRequestException from google.protobuf.json_format import MessageToDict @@ -1193,7 +1194,7 @@ async def test_resume_streaming_terminal_task_clears_final_continuation_token( assert len(updates) == 2 assert updates[0].continuation_token is not None assert updates[1].continuation_token is None - assert updates[1].is_operation_terminal is True + assert _get_operation_state(updates[1]) == "terminal" assert response.continuation_token is None assert response.finish_reason is None @@ -1232,7 +1233,7 @@ async def test_streaming_terminal_status_without_message_clears_final_continuati assert len(updates) == 2 assert updates[0].continuation_token is not None assert updates[1].contents == [] - assert updates[1].is_operation_terminal is True + assert _get_operation_state(updates[1]) == "terminal" assert response.continuation_token is None assert response.finish_reason is None @@ -2016,13 +2017,13 @@ async def test_streaming_background_artifacts_emit_terminal_marker_without_dupli assert updates[0].contents == [] assert updates[0].continuation_token is not None - assert updates[0].is_operation_terminal is False assert [update.text for update in updates[1:3]] == ["Hello ", "world"] + assert AgentResponse.from_updates(updates[:3]).continuation_token is not None assert updates[3].contents == [] assert updates[3].continuation_token is None - assert updates[3].is_operation_terminal is True + assert _get_operation_state(updates[3]) == "terminal" assert response.text == "Hello world" assert response.continuation_token is None diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 3fe4627e9bc..28a466c0f28 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -24,7 +24,7 @@ from copy import deepcopy from datetime import datetime from inspect import isawaitable -from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NamedTuple, NewType, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NamedTuple, NewType, TypeAlias, cast, overload from typing_extensions import Required, TypedDict @@ -40,6 +40,7 @@ logger = logging.getLogger("agent_framework") _SERIALIZED_EXCEPTION_MARKER: Final[str] = "FunctionInvocationError" +_OperationState: TypeAlias = Literal["in_progress", "terminal"] if TYPE_CHECKING: from pydantic import BaseModel @@ -2279,12 +2280,40 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse and update.finish_reason is not None ): response.finish_reason = update.finish_reason - if update.is_operation_terminal: + operation_state = _get_operation_state(update) + _copy_operation_state(update, response) + if operation_state == "terminal": response.continuation_token = None - elif update.continuation_token is not None: + elif operation_state == "in_progress": + if update.continuation_token is not None: + response.continuation_token = update.continuation_token + else: response.continuation_token = update.continuation_token +def _get_operation_state( + value: ChatResponse[Any] | ChatResponseUpdate | AgentResponse[Any] | AgentResponseUpdate, +) -> _OperationState | None: + """Get framework-internal resumable-operation state.""" + return value._operation_state # pyright: ignore[reportPrivateUsage] + + +def _set_operation_state( + value: ChatResponse[Any] | ChatResponseUpdate | AgentResponse[Any] | AgentResponseUpdate, + state: _OperationState | None, +) -> None: + """Set framework-internal resumable-operation state.""" + value._operation_state = state # pyright: ignore[reportPrivateUsage] + + +def _copy_operation_state( + source: ChatResponse[Any] | ChatResponseUpdate | AgentResponse[Any] | AgentResponseUpdate, + target: ChatResponse[Any] | ChatResponseUpdate | AgentResponse[Any] | AgentResponseUpdate, +) -> None: + """Copy framework-internal resumable-operation state.""" + _set_operation_state(target, _get_operation_state(source)) + + def _apply_response_tail_to_update( response: ChatResponse[Any] | AgentResponse[Any], update: ChatResponseUpdate | AgentResponseUpdate, @@ -2297,6 +2326,7 @@ def _apply_response_tail_to_update( """ update.finish_reason = response.finish_reason update.continuation_token = response.continuation_token + _copy_operation_state(response, update) if response.additional_properties: merged = dict(update.additional_properties) if update.additional_properties else {} merged.update(response.additional_properties) @@ -2692,6 +2722,7 @@ def __init__( _restore_compaction_annotation_in_additional_properties(additional_properties) or {} ) self.continuation_token = continuation_token + self._operation_state: _OperationState | None = None self.raw_representation: Any | list[Any] | None = raw_representation def mark_internal_conversation_id(self) -> None: @@ -2930,7 +2961,6 @@ class ChatResponseUpdate(SerializationMixin): model: The model associated with this response update. created_at: A timestamp for the chat response update. finish_reason: The finish reason for the operation. - is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Any additional properties associated with the chat response update. raw_representation: The raw representation of the chat response update from an underlying implementation. @@ -2977,7 +3007,6 @@ def __init__( created_at: CreatedAtT | None = None, finish_reason: FinishReasonLiteral | FinishReason | None = None, continuation_token: ContinuationToken | None = None, - is_operation_terminal: bool = False, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, ) -> None: @@ -2995,7 +3024,6 @@ def __init__( finish_reason: Optional finish reason for the operation. continuation_token: Optional token for resuming a long-running background operation. When present, indicates the operation is still in progress. - is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Optional additional properties associated with the chat response update. raw_representation: Optional raw representation of the chat response update from an underlying implementation. @@ -3024,20 +3052,13 @@ def __init__( self.created_at = created_at self.finish_reason = finish_reason self.continuation_token = continuation_token - self.is_operation_terminal = is_operation_terminal + self._operation_state: _OperationState | None = None self.additional_properties = _restore_compaction_annotation_in_additional_properties( additional_properties, allow_none=True, ) self.raw_representation = raw_representation - def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: - """Convert the update to a dictionary.""" - result = super().to_dict(exclude=exclude, exclude_none=exclude_none) - if not self.is_operation_terminal: - result.pop("is_operation_terminal", None) - return result - @property def text(self) -> str: """Returns the concatenated text of all contents in the update.""" @@ -3157,6 +3178,7 @@ def __init__( _restore_compaction_annotation_in_additional_properties(additional_properties) or {} ) self.continuation_token = continuation_token + self._operation_state: _OperationState | None = None self.raw_representation = raw_representation @property @@ -3372,6 +3394,7 @@ def _build_agent_response_from_chat_response( # pyright: ignore[reportUnusedFun if response._value_parsed: # pyright: ignore[reportPrivateUsage] agent_response._value = response._value # pyright: ignore[reportPrivateUsage] agent_response._value_parsed = True # pyright: ignore[reportPrivateUsage] + _copy_operation_state(response, agent_response) return agent_response @@ -3392,7 +3415,6 @@ class AgentResponseUpdate(SerializationMixin): response_id: The ID of the response of which this update is a part. message_id: The ID of the message of which this update is a part. created_at: A timestamp for the response update. - is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Any additional properties associated with the update. raw_representation: The raw representation from an underlying implementation. @@ -3440,7 +3462,6 @@ def __init__( created_at: CreatedAtT | None = None, finish_reason: FinishReasonLiteral | FinishReason | None = None, continuation_token: ContinuationToken | None = None, - is_operation_terminal: bool = False, additional_properties: dict[str, Any] | None = None, raw_representation: Any | None = None, ) -> None: @@ -3460,7 +3481,6 @@ def __init__( ``"tool_calls"`` (the model invoked a tool). continuation_token: Optional token for resuming a long-running background operation. When present, indicates the operation is still in progress. - is_operation_terminal: Whether this update marks the resumable operation as complete. additional_properties: Optional additional properties associated with the chat response update. raw_representation: Optional raw representation of the chat response update. @@ -3487,20 +3507,13 @@ def __init__( self.created_at = created_at self.finish_reason = finish_reason self.continuation_token = continuation_token - self.is_operation_terminal = is_operation_terminal + self._operation_state: _OperationState | None = None self.additional_properties = _restore_compaction_annotation_in_additional_properties( additional_properties, allow_none=True, ) self.raw_representation: Any | list[Any] | None = raw_representation - def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: - """Convert the update to a dictionary.""" - result = super().to_dict(exclude=exclude, exclude_none=exclude_none) - if not self.is_operation_terminal: - result.pop("is_operation_terminal", None) - return result - @property def text(self) -> str: """Get the concatenated text of all TextContent objects in contents.""" @@ -3519,7 +3532,7 @@ def __str__(self) -> str: def map_chat_to_agent_update(update: ChatResponseUpdate, agent_name: str | None) -> AgentResponseUpdate: - return AgentResponseUpdate( + agent_update = AgentResponseUpdate( contents=update.contents, role=update.role, author_name=update.author_name or agent_name, @@ -3528,10 +3541,11 @@ def map_chat_to_agent_update(update: ChatResponseUpdate, agent_name: str | None) created_at=update.created_at, finish_reason=update.finish_reason, # type: ignore[arg-type] continuation_token=update.continuation_token, - is_operation_terminal=update.is_operation_terminal, additional_properties=update.additional_properties, raw_representation=update, ) + _copy_operation_state(update, agent_update) + return agent_update # Type variables for ResponseStream diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index f2870ffa8f5..197b7298aec 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -29,6 +29,7 @@ Message, ResponseStream, UsageDetails, + _copy_operation_state, # pyright: ignore[reportPrivateUsage] add_usage_details, ) from ..exceptions import AgentException, AgentInvalidRequestException, AgentInvalidResponseException @@ -656,26 +657,25 @@ def _convert_workflow_event_to_agent_response_updates( if isinstance(data, AgentResponseUpdate): # Construct a fresh AgentResponseUpdate so we don't mutate a payload # that AgentExecutor still holds a reference to in its `updates` list. - return [ - AgentResponseUpdate( - contents=list(data.contents), - role=data.role, - author_name=data.author_name or executor_id, - agent_id=data.agent_id, - response_id=data.response_id, - message_id=data.message_id, - created_at=data.created_at, - # The attribute is typed wider than the constructor accepts (custom - # connectors may set any string); forward the value unchanged. - finish_reason=cast(FinishReasonLiteral | FinishReason | None, data.finish_reason), - continuation_token=data.continuation_token, - is_operation_terminal=data.is_operation_terminal, - additional_properties=dict(data.additional_properties) - if data.additional_properties is not None - else None, - raw_representation=data.raw_representation, - ) - ] + update = AgentResponseUpdate( + contents=list(data.contents), + role=data.role, + author_name=data.author_name or executor_id, + agent_id=data.agent_id, + response_id=data.response_id, + message_id=data.message_id, + created_at=data.created_at, + # The attribute is typed wider than the constructor accepts (custom + # connectors may set any string); forward the value unchanged. + finish_reason=cast(FinishReasonLiteral | FinishReason | None, data.finish_reason), + continuation_token=data.continuation_token, + additional_properties=dict(data.additional_properties) + if data.additional_properties is not None + else None, + raw_representation=data.raw_representation, + ) + _copy_operation_state(data, update) + return [update] if isinstance(data, AgentResponse): # Convert each message in AgentResponse to an AgentResponseUpdate updates: list[AgentResponseUpdate] = [] @@ -696,7 +696,7 @@ def _convert_workflow_event_to_agent_response_updates( updates[-1].agent_id = data.agent_id updates[-1].finish_reason = data.finish_reason updates[-1].continuation_token = data.continuation_token - updates[-1].is_operation_terminal = data.continuation_token is None + _copy_operation_state(cast(AgentResponse[Any], data), updates[-1]) updates[-1].additional_properties = dict(data.additional_properties) return updates if isinstance(data, Message): diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 9fffcb86903..7734dda334e 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -47,6 +47,7 @@ _parse_content_list, _parse_structured_response_value, _process_update, + _set_operation_state, _validate_uri, add_usage_details, map_chat_to_agent_update, @@ -6206,88 +6207,194 @@ def test_agent_response_update_serialization_includes_finish_reason() -> None: assert data["finish_reason"] == "tool_calls" -async def test_chat_client_agent_streaming_continuation_token() -> None: +def test_chat_response_preserves_continuation_token_for_explicitly_in_progress_updates() -> None: """Chat response aggregation should preserve an in-progress continuation token.""" - events = [ - 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")]), - ] + token_update = ChatResponseUpdate( + contents=[Content.from_text(text="working")], + response_id="resp_1", + model="gpt-5", + continuation_token=cast(Any, {"response_id": "resp_1"}), + ) + tokenless_update = ChatResponseUpdate(contents=[Content.from_text(text="...still working")]) + _set_operation_state(token_update, "in_progress") + _set_operation_state(tokenless_update, "in_progress") - aggregated = ChatResponse.from_updates(events) + aggregated = ChatResponse.from_updates([token_update, tokenless_update]) assert aggregated.continuation_token == {"response_id": "resp_1"} -def test_chat_response_finished_update_clears_continuation_token() -> None: - """A finished chat response should clear its continuation token.""" - events = [ - ChatResponseUpdate(continuation_token={"response_id": "resp_1"}), - ChatResponseUpdate(is_operation_terminal=True), - ] +def test_chat_response_terminal_update_clears_continuation_token() -> None: + """An explicitly terminal chat response should clear its continuation token.""" + token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = ChatResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") - aggregated = ChatResponse.from_updates(events) + aggregated = ChatResponse.from_updates([token_update, terminal_update]) assert aggregated.continuation_token is None -def test_chat_response_finish_reason_does_not_clear_continuation_token() -> None: - """A finish reason alone should not be treated as the terminal streaming signal.""" +def test_chat_response_legacy_tokenless_update_clears_continuation_token() -> None: + """Providers without explicit operation state should retain legacy token overwrite behavior.""" events = [ - ChatResponseUpdate(continuation_token={"response_id": "resp_1"}), + ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})), ChatResponseUpdate(finish_reason="stop"), ] aggregated = ChatResponse.from_updates(events) - assert aggregated.continuation_token == {"response_id": "resp_1"} + assert aggregated.continuation_token is None + assert aggregated.finish_reason == "stop" + + +def test_agent_response_legacy_tokenless_update_clears_continuation_token() -> None: + """Legacy agent providers should retain token overwrite behavior.""" + aggregated = AgentResponse.from_updates([ + AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})), + AgentResponseUpdate(finish_reason="stop"), + ]) + + assert aggregated.continuation_token is None assert aggregated.finish_reason == "stop" def test_agent_response_update_preserves_continuation_token() -> None: """Agent response aggregation should preserve an in-progress continuation token.""" - events = [ - AgentResponseUpdate(continuation_token={"response_id": "resp_1"}), - AgentResponseUpdate(contents=[Content.from_text(text="...still working")]), - ] + token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + tokenless_update = AgentResponseUpdate(contents=[Content.from_text(text="...still working")]) + _set_operation_state(token_update, "in_progress") + _set_operation_state(tokenless_update, "in_progress") - aggregated = AgentResponse.from_updates(events) + aggregated = AgentResponse.from_updates([token_update, tokenless_update]) assert aggregated.continuation_token == {"response_id": "resp_1"} -def test_agent_response_finished_update_clears_continuation_token() -> None: - """A finished agent response should clear its continuation token.""" - events = [ - AgentResponseUpdate(continuation_token={"response_id": "resp_1"}), - AgentResponseUpdate(is_operation_terminal=True), - ] +def test_agent_response_terminal_update_clears_continuation_token() -> None: + """An explicitly terminal agent response should clear its continuation token.""" + token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = AgentResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") - aggregated = AgentResponse.from_updates(events) + aggregated = AgentResponse.from_updates([token_update, terminal_update]) assert aggregated.continuation_token is None def test_map_chat_to_agent_update_preserves_terminal_signal() -> None: """Chat-to-agent update conversion should preserve terminality.""" - update = map_chat_to_agent_update(ChatResponseUpdate(is_operation_terminal=True), agent_name=None) + chat_update = ChatResponseUpdate() + _set_operation_state(chat_update, "terminal") + update = map_chat_to_agent_update(chat_update, agent_name=None) - assert update.is_operation_terminal is True + assert ( + AgentResponse.from_updates([ + AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})), + update, + ]).continuation_token + is None + ) -@pytest.mark.parametrize("update_type", [ChatResponseUpdate, AgentResponseUpdate]) -def test_response_update_serialization_omits_false_terminal_signal( - update_type: type[ChatResponseUpdate] | type[AgentResponseUpdate], +@pytest.mark.parametrize( + "value", + [ + ChatResponse(), + ChatResponseUpdate(), + AgentResponse(), + AgentResponseUpdate(), + ], +) +def test_response_serialization_omits_private_operation_state( + value: ChatResponse | ChatResponseUpdate | AgentResponse | AgentResponseUpdate, ) -> None: - """Nonterminal updates should retain their existing serialized shape.""" - assert "is_operation_terminal" not in update_type().to_dict() - serialized = update_type(is_operation_terminal=True).to_dict() - assert serialized["is_operation_terminal"] is True - assert update_type.from_dict(serialized).is_operation_terminal is True + """Internal operation state should not change the public serialized shape.""" + _set_operation_state(value, "terminal") + + assert "_operation_state" not in value.to_dict() + + +def test_chat_response_round_trip_preserves_terminal_operation_state() -> None: + """Chat response buffering should preserve terminality.""" + token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = ChatResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + response = ChatResponse.from_updates([token_update, terminal_update]) + + round_tripped_updates = response.to_updates() + legacy_token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "stale"})) + round_tripped = ChatResponse.from_updates([legacy_token_update, *round_tripped_updates]) + + assert round_tripped.continuation_token is None + + +def test_agent_response_round_trip_preserves_terminal_operation_state() -> None: + """Agent response buffering should preserve terminality.""" + token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = AgentResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + response = AgentResponse.from_updates([token_update, terminal_update]) + + round_tripped_updates = response.to_updates() + legacy_token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "stale"})) + round_tripped = AgentResponse.from_updates([legacy_token_update, *round_tripped_updates]) + + assert round_tripped.continuation_token is None + + +async def test_buffered_chat_response_replacement_preserves_terminal_operation_state() -> None: + """Buffered chat conversion should retain terminality on rebuilt updates.""" + token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = ChatResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield token_update + yield terminal_update + + stream = ResponseStream( + updates(), + finalizer=ChatResponse.from_updates, + result_transforms=[lambda response: response], + stream_updates=False, + result_to_updates=ChatResponse.to_updates, + ) + + released = [update async for update in stream] + stale_token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "stale"})) + + assert ChatResponse.from_updates([stale_token_update, *released]).continuation_token is None + + +async def test_buffered_agent_response_replacement_preserves_terminal_operation_state() -> None: + """Buffered agent conversion should retain terminality on rebuilt updates.""" + token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = AgentResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + + async def updates() -> AsyncIterable[AgentResponseUpdate]: + yield token_update + yield terminal_update + + stream = ResponseStream( + updates(), + finalizer=AgentResponse.from_updates, + result_transforms=[lambda response: response], + stream_updates=False, + result_to_updates=AgentResponse.to_updates, + ) + + released = [update async for update in stream] + stale_token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "stale"})) + + assert AgentResponse.from_updates([stale_token_update, *released]).continuation_token is None # endregion diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 81b5509d682..92f7bcfa545 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -33,6 +33,7 @@ handler, response_handler, ) +from agent_framework._types import _get_operation_state, _set_operation_state from agent_framework._workflows._typing_utils import deserialize_type @@ -871,19 +872,18 @@ async def test_workflow_as_agent_stream_preserves_response_update_metadata(self) @executor async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponseUpdate]) -> None: # type: ignore[valid-type] - await ctx.yield_output( - AgentResponseUpdate( - contents=[Content.from_text(text="payload")], - role="assistant", - agent_id="source-agent", - response_id="source-response", - message_id="source-message", - finish_reason="stop", - continuation_token=cast(Any, {"token": "resume-token"}), - is_operation_terminal=False, - additional_properties={"provider_marker": "preserve-me"}, - ) + response_update = AgentResponseUpdate( + contents=[Content.from_text(text="payload")], + role="assistant", + agent_id="source-agent", + response_id="source-response", + message_id="source-message", + finish_reason="stop", + continuation_token=cast(Any, {"token": "resume-token"}), + additional_properties={"provider_marker": "preserve-me"}, ) + _set_operation_state(response_update, "in_progress") + await ctx.yield_output(response_update) workflow = WorkflowBuilder(start_executor=metadata_executor).build() agent = workflow.as_agent("metadata-test-agent") @@ -899,7 +899,7 @@ async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, assert update.agent_id == "source-agent" assert update.finish_reason == "stop" assert update.continuation_token == {"token": "resume-token"} - assert update.is_operation_terminal is False + assert _get_operation_state(update) == "in_progress" assert update.additional_properties == {"provider_marker": "preserve-me"} async def test_workflow_as_agent_stream_preserves_response_metadata(self) -> None: @@ -907,19 +907,19 @@ async def test_workflow_as_agent_stream_preserves_response_metadata(self) -> Non @executor async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] - await ctx.yield_output( - AgentResponse( - messages=[ - Message(role="assistant", contents=["first"]), - Message(role="assistant", contents=["second"]), - ], - agent_id="source-agent", - response_id="source-response", - finish_reason="length", - continuation_token=cast(Any, {"token": "response-resume-token"}), - additional_properties={"provider_marker": "preserve-response"}, - ) + response = AgentResponse( + messages=[ + Message(role="assistant", contents=["first"]), + Message(role="assistant", contents=["second"]), + ], + agent_id="source-agent", + response_id="source-response", + finish_reason="length", + continuation_token=cast(Any, {"token": "response-resume-token"}), + additional_properties={"provider_marker": "preserve-response"}, ) + _set_operation_state(response, "in_progress") + await ctx.yield_output(response) workflow = WorkflowBuilder(start_executor=metadata_executor).build() agent = workflow.as_agent("response-metadata-test-agent") @@ -936,7 +936,7 @@ async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, assert updates[-1].agent_id == "source-agent" assert updates[-1].finish_reason == "length" assert updates[-1].continuation_token == {"token": "response-resume-token"} - assert updates[-1].is_operation_terminal is False + assert _get_operation_state(updates[-1]) == "in_progress" assert updates[-1].additional_properties == {"provider_marker": "preserve-response"} assert final_response.agent_id == "source-agent" assert final_response.finish_reason == "length" diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index eba70784660..5b324139b19 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -71,6 +71,7 @@ Role, TextSpanRegion, UsageDetails, + _set_operation_state, # pyright: ignore[reportPrivateUsage] detect_media_type_from_base64, validate_tool_mode, ) @@ -3649,7 +3650,7 @@ def _parse_chunk_from_openai( created_at: str | None = None continuation_token: OpenAIContinuationToken | None = None finish_reason: FinishReason | None = None - is_operation_terminal = False + operation_state: Literal["in_progress", "terminal"] = "in_progress" model = self.model def output_text_properties(output: Any) -> dict[str, Any] | None: @@ -3893,7 +3894,7 @@ def parse_annotations( conversation_id = self._get_conversation_id(event.response, options.get("store")) continuation_token = OpenAIContinuationToken(response_id=event.response.id) case "response.completed" | "response.incomplete" | "response.failed": - is_operation_terminal = True + operation_state = "terminal" response_id = event.response.id conversation_id = self._get_conversation_id(event.response, options.get("store")) model = event.response.model @@ -4194,7 +4195,7 @@ def parse_annotations( if not isinstance(event.type, str) or not event.type.startswith(_AZURE_AI_SEARCH_OUTPUT_EVENT_PREFIX): logger.debug("Unparsed event of type: %s: %s", event.type, event) - return ChatResponseUpdate( + update = ChatResponseUpdate( contents=contents, conversation_id=conversation_id, response_id=response_id, @@ -4203,10 +4204,11 @@ def parse_annotations( created_at=created_at, continuation_token=continuation_token, finish_reason=finish_reason, - is_operation_terminal=is_operation_terminal, additional_properties=metadata, raw_representation=event, ) + _set_operation_state(update, operation_state) + return update def _parse_usage_from_openai(self, usage: ResponseUsage) -> UsageDetails | None: details = UsageDetails( diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index f5aa728d707..1995964a9c8 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -39,6 +39,7 @@ SessionContext, _filter_approval_control_messages, ) +from agent_framework._types import _get_operation_state from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value from agent_framework.exceptions import ( ChatClientException, @@ -9605,6 +9606,28 @@ def test_streaming_response_in_progress_sets_continuation_token() -> None: assert _response_id_from_token(update.continuation_token) == "resp_stream_123" +def test_streaming_tokenless_intermediate_update_preserves_continuation_token() -> None: + """Tokenless Responses API events should not erase an active background token.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + in_progress_event = MagicMock() + in_progress_event.type = "response.in_progress" + in_progress_event.response.id = "resp_stream_123" + in_progress_event.response.conversation = None + + text_delta_event = MagicMock() + text_delta_event.type = "response.output_text.delta" + text_delta_event.delta = "working" + text_delta_event.logprobs = None + + response = ChatResponse.from_updates([ + client._parse_chunk_from_openai(in_progress_event, options={}, function_call_ids={}), + client._parse_chunk_from_openai(text_delta_event, options={}, function_call_ids={}), + ]) + + assert response.continuation_token is not None + assert _response_id_from_token(response.continuation_token) == "resp_stream_123" + + def test_streaming_response_created_with_in_progress_status_sets_continuation_token() -> None: """Test that response.created with in_progress status sets continuation_token.""" client = OpenAIChatClient(model="test-model", api_key="test-key") @@ -9751,7 +9774,7 @@ def test_streaming_terminal_response_sets_finish_reason( update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids={}) assert update.finish_reason == expected_finish_reason - assert update.is_operation_terminal is True + assert _get_operation_state(update) == "terminal" @pytest.mark.parametrize( From c4312cade7bc2aabc10007ee74f6e857bc7a5787 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Thu, 8 Oct 2026 11:16:08 +0200 Subject: [PATCH 4/7] Potential fix for pull request finding 'Unused import' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- python/packages/core/agent_framework/_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 28a466c0f28..b7446b68114 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -24,7 +24,7 @@ from copy import deepcopy from datetime import datetime from inspect import isawaitable -from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NamedTuple, NewType, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, NewType, TypeAlias, cast, overload from typing_extensions import Required, TypedDict From aea1a64d9f38eea397f2c6c1e56eed2fb8e8b90c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:24:06 +0000 Subject: [PATCH 5/7] Restore Generic import for response classes Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com> --- python/packages/core/agent_framework/_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 59ea842e715..9a209d0a39b 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -24,7 +24,7 @@ from copy import deepcopy from datetime import datetime from inspect import isawaitable -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, NewType, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NamedTuple, NewType, TypeAlias, cast, overload from typing_extensions import Required, TypedDict From d545fcc4df7a7fde12c7e3615e4c388f545c4b28 Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Thu, 8 Oct 2026 15:48:06 -0700 Subject: [PATCH 6/7] error handling and serialization fix --- .../packages/core/agent_framework/_types.py | 67 +++++++++++++++++ python/packages/core/tests/core/test_types.py | 71 ++++++++++++++++++- .../agent_framework_openai/_chat_client.py | 31 ++++---- .../tests/openai/test_openai_chat_client.py | 54 ++++++++++++++ 4 files changed, 207 insertions(+), 16 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 9a209d0a39b..bd027380b38 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -41,6 +41,7 @@ _SERIALIZED_EXCEPTION_MARKER: Final[str] = "FunctionInvocationError" _OperationState: TypeAlias = Literal["in_progress", "terminal"] +_OPERATION_STATE_SERIALIZATION_KEY: Final[str] = "_operation_state" if TYPE_CHECKING: from pydantic import BaseModel @@ -2306,6 +2307,28 @@ def _set_operation_state( value._operation_state = state # pyright: ignore[reportPrivateUsage] +def _serialize_operation_state( + value: ChatResponseUpdate | AgentResponseUpdate, + data: dict[str, Any], + exclude: set[str] | None, +) -> dict[str, Any]: + """Add explicitly set operation state to an update's serialized representation.""" + if (not exclude or _OPERATION_STATE_SERIALIZATION_KEY not in exclude) and ( + operation_state := _get_operation_state(value) + ): + data[_OPERATION_STATE_SERIALIZATION_KEY] = operation_state + return data + + +def _deserialize_operation_state(value: MutableMapping[str, Any]) -> tuple[dict[str, Any], _OperationState | None]: + """Remove and validate operation state from an update's serialized representation.""" + data = dict(value) + operation_state = data.pop(_OPERATION_STATE_SERIALIZATION_KEY, None) + if operation_state not in (None, "in_progress", "terminal"): + raise ValueError(f"Invalid operation state: {operation_state!r}") + return data, operation_state + + def _copy_operation_state( source: ChatResponse[Any] | ChatResponseUpdate | AgentResponse[Any] | AgentResponseUpdate, target: ChatResponse[Any] | ChatResponseUpdate | AgentResponse[Any] | AgentResponseUpdate, @@ -3059,6 +3082,28 @@ def __init__( ) self.raw_representation = raw_representation + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Serialize the update, including framework-internal operation state when set.""" + return _serialize_operation_state( + self, + super().to_dict(exclude=exclude, exclude_none=exclude_none), + exclude, + ) + + @classmethod + def from_dict( + cls, + value: MutableMapping[str, Any], + /, + *, + dependencies: MutableMapping[str, Any] | None = None, + ) -> ChatResponseUpdate: + """Deserialize an update and restore framework-internal operation state.""" + data, operation_state = _deserialize_operation_state(value) + update = super().from_dict(data, dependencies=dependencies) + _set_operation_state(update, operation_state) + return update + @property def text(self) -> str: """Returns the concatenated text of all contents in the update.""" @@ -3514,6 +3559,28 @@ def __init__( ) self.raw_representation: Any | list[Any] | None = raw_representation + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: + """Serialize the update, including framework-internal operation state when set.""" + return _serialize_operation_state( + self, + super().to_dict(exclude=exclude, exclude_none=exclude_none), + exclude, + ) + + @classmethod + def from_dict( + cls, + value: MutableMapping[str, Any], + /, + *, + dependencies: MutableMapping[str, Any] | None = None, + ) -> AgentResponseUpdate: + """Deserialize an update and restore framework-internal operation state.""" + data, operation_state = _deserialize_operation_state(value) + update = super().from_dict(data, dependencies=dependencies) + _set_operation_state(update, operation_state) + return update + @property def text(self) -> str: """Get the concatenated text of all TextContent objects in contents.""" diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 7ecfa992c70..7d4995e14d7 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -44,6 +44,7 @@ _append_instructions, _get_data_bytes, _get_data_bytes_as_str, + _get_operation_state, _parse_content_list, _parse_structured_response_value, _process_update, @@ -6463,6 +6464,34 @@ def test_chat_response_preserves_continuation_token_for_explicitly_in_progress_u assert aggregated.continuation_token == {"response_id": "resp_1"} +def test_chat_response_update_serialization_preserves_in_progress_continuation_token() -> None: + """Serialized in-progress chat updates should retain an active continuation token.""" + token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + tokenless_update = ChatResponseUpdate(contents=[Content.from_text(text="...still working")]) + _set_operation_state(token_update, "in_progress") + _set_operation_state(tokenless_update, "in_progress") + + serialized_updates = json.dumps([token_update.to_dict(), tokenless_update.to_dict()]) + restored_updates = [ChatResponseUpdate.from_dict(update_data) for update_data in json.loads(serialized_updates)] + + assert ChatResponse.from_updates(restored_updates).continuation_token == {"response_id": "resp_1"} + + +def test_chat_response_update_serialization_preserves_terminal_state() -> None: + """Serialized terminal chat updates should retain terminal operation state.""" + token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = ChatResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + + serialized_updates = json.dumps([token_update.to_dict(), terminal_update.to_dict()]) + restored_updates = [ChatResponseUpdate.from_dict(update_data) for update_data in json.loads(serialized_updates)] + + assert _get_operation_state(restored_updates[0]) == "in_progress" + assert _get_operation_state(restored_updates[1]) == "terminal" + assert ChatResponse.from_updates(restored_updates).continuation_token is None + + def test_chat_response_terminal_update_clears_continuation_token() -> None: """An explicitly terminal chat response should clear its continuation token.""" token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) @@ -6511,6 +6540,34 @@ def test_agent_response_update_preserves_continuation_token() -> None: assert aggregated.continuation_token == {"response_id": "resp_1"} +def test_agent_response_update_serialization_preserves_in_progress_continuation_token() -> None: + """Serialized in-progress agent updates should retain an active continuation token.""" + token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + tokenless_update = AgentResponseUpdate(contents=[Content.from_text(text="...still working")]) + _set_operation_state(token_update, "in_progress") + _set_operation_state(tokenless_update, "in_progress") + + serialized_updates = json.dumps([token_update.to_dict(), tokenless_update.to_dict()]) + restored_updates = [AgentResponseUpdate.from_dict(update_data) for update_data in json.loads(serialized_updates)] + + assert AgentResponse.from_updates(restored_updates).continuation_token == {"response_id": "resp_1"} + + +def test_agent_response_update_serialization_preserves_terminal_state() -> None: + """Serialized terminal agent updates should retain terminal operation state.""" + token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) + terminal_update = AgentResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + + serialized_updates = json.dumps([token_update.to_dict(), terminal_update.to_dict()]) + restored_updates = [AgentResponseUpdate.from_dict(update_data) for update_data in json.loads(serialized_updates)] + + assert _get_operation_state(restored_updates[0]) == "in_progress" + assert _get_operation_state(restored_updates[1]) == "terminal" + assert AgentResponse.from_updates(restored_updates).continuation_token is None + + def test_agent_response_terminal_update_clears_continuation_token() -> None: """An explicitly terminal agent response should clear its continuation token.""" token_update = AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) @@ -6542,13 +6599,11 @@ def test_map_chat_to_agent_update_preserves_terminal_signal() -> None: "value", [ ChatResponse(), - ChatResponseUpdate(), AgentResponse(), - AgentResponseUpdate(), ], ) def test_response_serialization_omits_private_operation_state( - value: ChatResponse | ChatResponseUpdate | AgentResponse | AgentResponseUpdate, + value: ChatResponse | AgentResponse, ) -> None: """Internal operation state should not change the public serialized shape.""" _set_operation_state(value, "terminal") @@ -6556,6 +6611,16 @@ def test_response_serialization_omits_private_operation_state( assert "_operation_state" not in value.to_dict() +@pytest.mark.parametrize("value", [ChatResponseUpdate(), AgentResponseUpdate()]) +def test_update_serialization_includes_explicit_operation_state( + value: ChatResponseUpdate | AgentResponseUpdate, +) -> None: + """Update serialization should preserve explicitly set internal operation state.""" + _set_operation_state(value, "terminal") + + assert value.to_dict()["_operation_state"] == "terminal" + + def test_chat_response_round_trip_preserves_terminal_operation_state() -> None: """Chat response buffering should preserve terminality.""" token_update = ChatResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 5b324139b19..e3ebcea7a90 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -802,6 +802,8 @@ async def _prepare_request( def _handle_request_error(self, ex: Exception) -> NoReturn: """Convert exceptions to appropriate service exceptions. Always raises.""" + if isinstance(ex, ChatClientException): + raise ex if isinstance(ex, BadRequestError) and ex.code == "content_filter": raise OpenAIContentFilterException( f"{type(self)} service encountered a content error: {ex}", @@ -867,6 +869,15 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: served_model = self._extract_served_model(getattr(raw_stream_response, "headers", None)) async with _open_event_stream(raw_stream_response) as stream_response: async for chunk in stream_response: + if chunk.type in ( + "error", + "response.completed", + "response.incomplete", + "response.failed", + ) and isinstance(options, dict): + # Clear the caller-owned resume token before parsing because + # ResponseErrorEvent parsing raises instead of yielding an update. + options.pop("continuation_token", None) update = self._parse_chunk_from_openai( chunk, options=validated_options, @@ -875,19 +886,6 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: ) if served_model is not None: update.model = served_model - if chunk.type in ( - "response.completed", - "response.incomplete", - "response.failed", - ) and isinstance(options, dict): - # Same as the non-streaming path (issue #5394): once the resumed - # background response has finished, drop the continuation_token - # from the caller's options dict. FunctionInvocationLayer reuses - # that dict, so a leftover token makes the next tool-loop iteration - # retrieve this response again instead of POSTing the tool results, - # and the tools run again each time. Do it before yielding, so a - # consumer that stops at the terminal update doesn't keep it. - options.pop("continuation_token", None) yield update except Exception as ex: self._handle_request_error(ex) @@ -3737,6 +3735,13 @@ def parse_annotations( # ResponseQueuedEvent, # ResponseCustomToolCallInputDeltaEvent, # ResponseCustomToolCallInputDoneEvent, + case "error": + error_details = event.message + if event.code: + error_details = f"{event.code}: {error_details}" + if event.param: + error_details = f"{error_details} (parameter: {event.param})" + self._handle_request_error(RuntimeError(f"OpenAI streaming error: {error_details}")) case "response.content_part.added": event_part = event.part match event_part.type: diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 1995964a9c8..a55eeb54d5e 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -51,6 +51,7 @@ ResponseComputerToolCall, ResponseComputerToolCallOutputItem, ResponseContentPartDoneEvent, + ResponseErrorEvent, ResponseFunctionShellToolCall, ResponseFunctionShellToolCallOutput, ResponseOutputItemDoneEvent, @@ -9628,6 +9629,59 @@ def test_streaming_tokenless_intermediate_update_preserves_continuation_token() assert _response_id_from_token(response.continuation_token) == "resp_stream_123" +def test_streaming_error_event_raises_chat_client_exception() -> None: + """Responses API error events should surface their provider details.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + error_event = ResponseErrorEvent( + code="server_error", + message="The response stream failed.", + param="input", + sequence_number=2, + type="error", + ) + + with pytest.raises(ChatClientException) as exc_info: + client._parse_chunk_from_openai(error_event, options={}, function_call_ids={}) + + assert "OpenAI streaming error: server_error: The response stream failed. (parameter: input)" in str(exc_info.value) + + +async def test_resumed_stream_error_clears_continuation_token() -> None: + """A failed resumed stream should clear its caller-owned continuation token.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + in_progress_event = MagicMock() + in_progress_event.type = "response.in_progress" + in_progress_event.response.id = "resp_stream_error" + in_progress_event.response.conversation = None + error_event = ResponseErrorEvent( + code="server_error", + message="The response stream failed.", + param=None, + sequence_number=2, + type="error", + ) + retrieve = AsyncMock(return_value=_FakeAsyncEventStream([in_progress_event, error_event])) + options: OpenAIChatOptions[None] = {"continuation_token": {"response_id": "resp_stream_error"}} + updates: list[ChatResponseUpdate] = [] + + with ( + patch.object(client.client.responses.with_raw_response, "retrieve", new=retrieve), + pytest.raises(ChatClientException, match="OpenAI streaming error: server_error"), + ): + stream = client._inner_get_response( + messages=[Message(role="user", contents=["resume"])], + stream=True, + options=options, + ) + assert isinstance(stream, ResponseStream) + async for update in stream: + updates.append(update) + + assert len(updates) == 1 + assert updates[0].continuation_token == {"response_id": "resp_stream_error"} + assert "continuation_token" not in options + + def test_streaming_response_created_with_in_progress_status_sets_continuation_token() -> None: """Test that response.created with in_progress status sets continuation_token.""" client = OpenAIChatClient(model="test-model", api_key="test-key") From 46d13cdf4f6858e9bc7f126da40a20c26a07b2d9 Mon Sep 17 00:00:00 2001 From: sophia-ramsey Date: Thu, 8 Oct 2026 16:41:59 -0700 Subject: [PATCH 7/7] fix default event status --- .../agent_framework_openai/_chat_client.py | 30 +++++++----- .../tests/openai/test_openai_chat_client.py | 48 +++++++++++++++++++ 2 files changed, 67 insertions(+), 11 deletions(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index e3ebcea7a90..a42424457d8 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -27,6 +27,7 @@ TYPE_CHECKING, Any, ClassVar, + Final, Generic, Literal, NoReturn, @@ -155,6 +156,7 @@ class _PromptCacheOptions(TypedDict, total=False): _MODEL_OUTPUT_KIND_KEY = "model_output_kind" _MODEL_OUTPUT_REFUSAL = "refusal" +_TERMINAL_RESPONSE_STATUSES: Final[frozenset[str]] = frozenset({"completed", "failed", "incomplete", "cancelled"}) def _is_refusal_text_content(content: Content) -> bool: @@ -869,14 +871,21 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: served_model = self._extract_served_model(getattr(raw_stream_response, "headers", None)) async with _open_event_stream(raw_stream_response) as stream_response: async for chunk in stream_response: - if chunk.type in ( - "error", - "response.completed", - "response.incomplete", - "response.failed", - ) and isinstance(options, dict): + if isinstance(options, dict) and ( + chunk.type + in ( + "error", + "response.completed", + "response.incomplete", + "response.failed", + ) + or ( + chunk.type == "response.created" + and chunk.response.status in _TERMINAL_RESPONSE_STATUSES + ) + ): # Clear the caller-owned resume token before parsing because - # ResponseErrorEvent parsing raises instead of yielding an update. + # error parsing raises and response.created may itself be terminal. options.pop("continuation_token", None) update = self._parse_chunk_from_openai( chunk, @@ -3889,10 +3898,9 @@ def parse_annotations( case "response.created": response_id = event.response.id conversation_id = self._get_conversation_id(event.response, options.get("store")) - if event.response.status and event.response.status in ( - "in_progress", - "queued", - ): + if event.response.status in _TERMINAL_RESPONSE_STATUSES: + operation_state = "terminal" + elif event.response.status in ("in_progress", "queued"): continuation_token = OpenAIContinuationToken(response_id=event.response.id) case "response.in_progress": response_id = event.response.id diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index a55eeb54d5e..99e1d75c979 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -9702,6 +9702,54 @@ def test_streaming_response_created_with_in_progress_status_sets_continuation_to assert _response_id_from_token(update.continuation_token) == "resp_created_123" +@pytest.mark.parametrize("status", ["completed", "failed", "incomplete", "cancelled"]) +def test_streaming_response_created_with_terminal_status_clears_continuation_token(status: str) -> None: + """A terminal status embedded in response.created should terminate the operation.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + in_progress_event = MagicMock() + in_progress_event.type = "response.in_progress" + in_progress_event.response.id = "resp_created_terminal" + in_progress_event.response.conversation = None + + created_event = MagicMock() + created_event.type = "response.created" + created_event.response.id = "resp_created_terminal" + created_event.response.conversation = None + created_event.response.status = status + + updates = [ + client._parse_chunk_from_openai(in_progress_event, options={}, function_call_ids={}), + client._parse_chunk_from_openai(created_event, options={}, function_call_ids={}), + ] + + assert _get_operation_state(updates[-1]) == "terminal" + assert ChatResponse.from_updates(updates).continuation_token is None + + +async def test_resumed_stream_terminal_created_event_clears_continuation_token() -> None: + """A terminal response.created event should clear the resumed stream token.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + created_event = MagicMock() + created_event.type = "response.created" + created_event.response.id = "resp_created_cancelled" + created_event.response.conversation = None + created_event.response.status = "cancelled" + retrieve = AsyncMock(return_value=_FakeAsyncEventStream([created_event])) + options: OpenAIChatOptions[None] = {"continuation_token": {"response_id": "resp_created_cancelled"}} + + with patch.object(client.client.responses.with_raw_response, "retrieve", new=retrieve): + stream = client._inner_get_response( + messages=[Message(role="user", contents=["resume"])], + stream=True, + options=options, + ) + assert isinstance(stream, ResponseStream) + response = await stream.get_final_response() + + assert response.continuation_token is None + assert "continuation_token" not in options + + def test_streaming_response_completed_no_continuation_token() -> None: """Test that response.completed does NOT set continuation_token.""" client = OpenAIChatClient(model="test-model", api_key="test-key")