diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index c5780b1ed5f..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. @@ -680,18 +689,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: @@ -702,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. @@ -814,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) @@ -840,7 +851,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,28 +864,49 @@ def _updates_from_task( ) for message in task_messages ] + _set_operation_state(updates[-1], "terminal") + return updates if task.artifacts: - return [] + 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 [ - AgentResponseUpdate( - contents=[], - role="assistant", - response_id=task.id, - 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, + 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", ) ] @@ -888,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, @@ -914,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 [] @@ -944,47 +986,78 @@ 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, - ) - ] - - if not update_event.status.HasField("message") or not update_event.status.message.parts: - return [] + 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 [] + 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 [ + _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", + ) + ] + message = update_event.status.message contents = self._parse_contents_from_a2a(message.parts) if not contents: - return [] + if not background: + return [] + return [ + _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 {} - event_meta = MessageToDict(update_event.metadata) if update_event.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, - 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 533e2e2a667..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 @@ -1142,6 +1143,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 _get_operation_state(updates[1]) == "terminal" + 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 _get_operation_state(updates[1]) == "terminal" + 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) @@ -1866,6 +1962,73 @@ async def test_streaming_artifact_update_event_does_not_duplicate_terminal_task_ 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 [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 _get_operation_state(updates[3]) == "terminal" + + 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: diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 86f50bf60a2..bd027380b38 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,8 @@ logger = logging.getLogger("agent_framework") _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 @@ -2279,7 +2281,60 @@ 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 + operation_state = _get_operation_state(update) + _copy_operation_state(update, response) + if operation_state == "terminal": + response.continuation_token = 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 _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, +) -> None: + """Copy framework-internal resumable-operation state.""" + _set_operation_state(target, _get_operation_state(source)) def _apply_response_tail_to_update( @@ -2294,6 +2349,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) @@ -2689,6 +2745,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: @@ -3018,12 +3075,35 @@ def __init__( self.created_at = created_at self.finish_reason = finish_reason self.continuation_token = continuation_token + 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]: + """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.""" @@ -3143,6 +3223,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 @@ -3358,6 +3439,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 @@ -3470,12 +3552,35 @@ def __init__( self.created_at = created_at self.finish_reason = finish_reason self.continuation_token = continuation_token + 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]: + """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.""" @@ -3494,7 +3599,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, @@ -3506,6 +3611,8 @@ def map_chat_to_agent_update(update: ChatResponseUpdate, agent_name: str | None) 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 6d01c8366a4..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,25 +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, - 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] = [] @@ -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 + _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 03ce5ff54f4..7d4995e14d7 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -44,9 +44,11 @@ _append_instructions, _get_data_bytes, _get_data_bytes_as_str, + _get_operation_state, _parse_content_list, _parse_structured_response_value, _process_update, + _set_operation_state, _validate_uri, add_usage_details, map_chat_to_agent_update, @@ -6445,6 +6447,260 @@ def test_agent_response_update_serialization_includes_finish_reason() -> None: assert data["finish_reason"] == "tool_calls" +def test_chat_response_preserves_continuation_token_for_explicitly_in_progress_updates() -> None: + """Chat response aggregation should preserve an in-progress continuation token.""" + 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([token_update, tokenless_update]) + + 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"})) + terminal_update = ChatResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + + aggregated = ChatResponse.from_updates([token_update, terminal_update]) + + assert aggregated.continuation_token is None + + +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=cast(Any, {"response_id": "resp_1"})), + ChatResponseUpdate(finish_reason="stop"), + ] + + aggregated = ChatResponse.from_updates(events) + + 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.""" + 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([token_update, tokenless_update]) + + 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"})) + terminal_update = AgentResponseUpdate() + _set_operation_state(token_update, "in_progress") + _set_operation_state(terminal_update, "terminal") + + 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.""" + chat_update = ChatResponseUpdate() + _set_operation_state(chat_update, "terminal") + update = map_chat_to_agent_update(chat_update, agent_name=None) + + assert ( + AgentResponse.from_updates([ + AgentResponseUpdate(continuation_token=cast(Any, {"response_id": "resp_1"})), + update, + ]).continuation_token + is None + ) + + +@pytest.mark.parametrize( + "value", + [ + ChatResponse(), + AgentResponse(), + ], +) +def test_response_serialization_omits_private_operation_state( + value: ChatResponse | AgentResponse, +) -> None: + """Internal operation state should not change the public serialized shape.""" + _set_operation_state(value, "terminal") + + 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"})) + 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 0ef70ca2f6e..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 @@ -863,7 +864,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(). @@ -871,18 +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"}), - 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") @@ -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 _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: @@ -905,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") @@ -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 _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 24ca5066975..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, @@ -71,6 +72,7 @@ Role, TextSpanRegion, UsageDetails, + _set_operation_state, # pyright: ignore[reportPrivateUsage] detect_media_type_from_base64, validate_tool_mode, ) @@ -154,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: @@ -801,6 +804,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}", @@ -866,6 +871,22 @@ 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 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 + # error parsing raises and response.created may itself be terminal. + options.pop("continuation_token", None) update = self._parse_chunk_from_openai( chunk, options=validated_options, @@ -874,19 +895,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) @@ -3649,6 +3657,7 @@ def _parse_chunk_from_openai( created_at: str | None = None continuation_token: OpenAIContinuationToken | None = None finish_reason: FinishReason | None = None + operation_state: Literal["in_progress", "terminal"] = "in_progress" model = self.model def output_text_properties(output: Any) -> dict[str, Any] | None: @@ -3735,6 +3744,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: @@ -3882,16 +3898,16 @@ 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 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": + operation_state = "terminal" response_id = event.response.id conversation_id = self._get_conversation_id(event.response, options.get("store")) model = event.response.model @@ -4192,7 +4208,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, @@ -4204,6 +4220,8 @@ def parse_annotations( 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 bd576e61321..99e1d75c979 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, @@ -50,6 +51,7 @@ ResponseComputerToolCall, ResponseComputerToolCallOutputItem, ResponseContentPartDoneEvent, + ResponseErrorEvent, ResponseFunctionShellToolCall, ResponseFunctionShellToolCallOutput, ResponseOutputItemDoneEvent, @@ -9605,6 +9607,81 @@ 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_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") @@ -9625,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") @@ -9751,6 +9876,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 _get_operation_state(update) == "terminal" + + +@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: