diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index dd894913609..558e4c39cb6 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -1725,6 +1725,14 @@ def _parse_contents_from_anthropic( def _parse_citations_from_anthropic( self, content_block: BetaContentBlock | BetaRawContentBlockDelta | BetaTextBlock ) -> list[Annotation] | None: + """Parse Anthropic citations into annotations. + + Note: + Document and search-result citations retain their source coordinates in ``annotated_regions``. + Their ``additional_properties["annotated_regions_coordinate_space"]`` is ``"source"`` to distinguish + these locations from generated-response character offsets. The Responses host does not emit these + regions as output-text spans. + """ content_blocks = getattr(content_block, "citations", None) if not content_blocks: return None @@ -1789,6 +1797,8 @@ def _parse_citations_from_anthropic( ) case _: logger.debug(f"Unknown citation type encountered: {citation.type}") + if cit.get("annotated_regions"): + cit["additional_properties"] = {"annotated_regions_coordinate_space": "source"} annotations.append(cit) return annotations or None diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index 7822da91600..d14027f8b2d 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -4304,6 +4304,50 @@ def test_parse_citations_search_result_location( assert len(result) > 0 +@pytest.mark.parametrize( + ("citation_type", "location_fields", "expected_span"), + [ + ("char_location", {"start_char_index": 2, "end_char_index": 12}, (2, 12)), + ("page_location", {"start_page_number": 1, "end_page_number": 3}, (1, 3)), + ("content_block_location", {"start_block_index": 0, "end_block_index": 2}, (0, 2)), + ("search_result_location", {"start_block_index": 3, "end_block_index": 5}, (3, 5)), + ("web_search_result_location", {}, None), + ], +) +def test_parse_citations_marks_source_coordinate_space( + mock_anthropic_client: MagicMock, + citation_type: str, + location_fields: dict[str, int], + expected_span: tuple[int, int] | None, +) -> None: + client = create_test_anthropic_client(mock_anthropic_client) + citation = MagicMock( + type=citation_type, + title="Source", + document_title="Source", + cited_text="Cited text", + file_id=None, + url="https://example.com", + source="https://example.com", + ) + citation.configure_mock(**location_fields) + block = MagicMock(type="text", text="Answer", citations=[citation]) + + result = client._parse_citations_from_anthropic(block) + + assert result is not None + assert len(result) == 1 + annotation = result[0] + if expected_span is None: + assert "annotated_regions" not in annotation + assert "annotated_regions_coordinate_space" not in annotation.get("additional_properties", {}) + else: + assert annotation["annotated_regions"] == [ + {"type": "text_span", "start_index": expected_span[0], "end_index": expected_span[1]} + ] + assert annotation["additional_properties"]["annotated_regions_coordinate_space"] == "source" + + @pytest.mark.flaky @pytest.mark.integration @skip_if_anthropic_integration_tests_disabled diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 3a352eba761..279fc2c3ee6 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -276,6 +276,18 @@ output item ID without changing the provider's `call_id`. Although screenshots are optional in shared `Content`, this Responses host requires them on computer results. The computer-use `Content` constructors and `ComputerSafetyCheck` type are experimental Agent Framework APIs. +## Responses citation annotations + +`ResponsesHostServer` serializes supported framework URL and file citations into `output_text.annotations`. +URL and container-file citation spans refer to character positions in the generated response text. +Annotations with `additional_properties["annotated_regions_coordinate_space"] == "source"` retain +source-document coordinates (for example, Anthropic character, page, or block locations); the host leaves +those annotations unchanged but does not emit their regions as output-text spans. Other supported citation +metadata, such as file citations with an explicit output index, is still converted. +Completed provider responses retain separate text content parts, keeping citation offsets relative to +their original part. Ordinary streaming deltas still merge within their existing text occurrence, +without shifting citation offsets. + ## State store ### Local persistence diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 09692723734..a6f67174461 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -50,6 +50,9 @@ add_usage_details, get_namespaced_value, ) +from agent_framework import ( + Annotation as FrameworkAnnotation, +) from agent_framework._mcp import _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY # pyright: ignore[reportPrivateUsage] from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentFrameworkException @@ -63,12 +66,17 @@ from azure.ai.agentserver.responses._id_generator import IdGenerator from azure.ai.agentserver.responses.aio import ResponseEventStream from azure.ai.agentserver.responses.hosting import ResponsesAgentServerHost +from azure.ai.agentserver.responses.models import ( + Annotation as ResponsesAnnotation, +) from azure.ai.agentserver.responses.models import ( ComputerAction, ComputerCallSafetyCheckParam, ComputerScreenshotImage, ContainerFileCitationBody, CreateResponse, + FileCitationBody, + FilePath, FunctionShellAction, FunctionShellCallOutputContent, FunctionShellCallOutputExitOutcome, @@ -88,6 +96,7 @@ ResponseUsage, ResponseUsageInputTokensDetails, ResponseUsageOutputTokensDetails, + UrlCitationBody, ) from azure.ai.agentserver.responses.streaming._builders import ( OutputItemBuilder, @@ -272,6 +281,98 @@ def _container_file_citations_from_function_result(content: Content) -> list[_Co return citations +def _annotations_to_output_text( + annotations: Sequence[FrameworkAnnotation] | None, +) -> list[ResponsesAnnotation]: + """Convert framework citations into Responses output-text annotations. + + Regions with ``additional_properties["annotated_regions_coordinate_space"] == "source"`` + describe the source document, not output text, and are not emitted as output spans. + """ + if not annotations or not isinstance(annotations, Sequence) or isinstance(annotations, (str, bytes)): + return [] + + output_annotations: list[ResponsesAnnotation] = [] + for annotation in annotations: + if not isinstance(annotation, Mapping): + continue + annotation_mapping = cast("Mapping[str, Any]", annotation) + if annotation_mapping.get("type") != "citation": + continue + + raw_properties = annotation_mapping.get("additional_properties") + properties: Mapping[str, Any] = ( + cast("Mapping[str, Any]", raw_properties) if isinstance(raw_properties, Mapping) else {} + ) + raw_regions = annotation_mapping.get("annotated_regions") + if properties.get("annotated_regions_coordinate_space") == "source": + raw_regions = None + regions: list[Mapping[str, Any]] = [] + if isinstance(raw_regions, Sequence) and not isinstance(raw_regions, (str, bytes)): + for raw_region in cast("Sequence[Any]", raw_regions): + if not isinstance(raw_region, Mapping): + continue + region = cast("Mapping[str, Any]", raw_region) + if region.get("type") == "text_span": + regions.append(region) + file_id = annotation_mapping.get("file_id") + url = annotation_mapping.get("url") + title = annotation_mapping.get("title") + container_id = properties.get("container_id") + index = properties.get("index") + + if isinstance(container_id, str) and container_id and isinstance(file_id, str) and file_id: + for region in regions: + start_index = region.get("start_index") + end_index = region.get("end_index") + if not isinstance(start_index, int) or not isinstance(end_index, int): + continue + output_annotations.append( + ContainerFileCitationBody( + type="container_file_citation", + container_id=container_id, + file_id=file_id, + filename=url or "", + start_index=start_index, + end_index=end_index, + ) + ) + elif isinstance(url, str) and url and not file_id: + for region in regions: + start_index = region.get("start_index") + end_index = region.get("end_index") + if not isinstance(start_index, int) or not isinstance(end_index, int): + continue + output_annotations.append( + UrlCitationBody( + type="url_citation", + url=url, + title=title or "", + start_index=start_index, + end_index=end_index, + ) + ) + elif isinstance(file_id, str) and file_id and isinstance(url, str) and url and isinstance(index, int): + output_annotations.append( + FileCitationBody( + type="file_citation", + file_id=file_id, + filename=url, + index=index, + ) + ) + elif isinstance(file_id, str) and file_id and isinstance(index, int): + output_annotations.append( + FilePath( + type="file_path", + file_id=file_id, + index=index, + ) + ) + + return output_annotations + + def _is_refusal_text_content(content: Content) -> bool: return ( content.type == "text" @@ -1763,7 +1864,9 @@ async def _handle_inner_agent( for update in _agent_response_updates(response, context.response_id): if context.shutdown.is_set() or cancellation_signal.is_set(): break - async for event in tracker.handle_update(update, approval_storage=approval_storage): + async for event in tracker.handle_update( + update, approval_storage=approval_storage, complete_message=True + ): yield event if context.shutdown.is_set() or cancellation_signal.is_set(): continue @@ -2423,7 +2526,8 @@ def __init__( self._outstanding_computer_calls: set[str] = set() self._oauth_consent_requests: set[tuple[str, str]] = set() self._pending_container_file_citations: dict[str, _ContainerFileCitation] = {} - self._message_text_annotations: dict[int, list[ContainerFileCitationBody]] = {} + self._text_annotations: list[ResponsesAnnotation] = [] + self._message_text_annotations: dict[int, list[ResponsesAnnotation]] = {} # Set when an agent update reports the model stopped early (content filter, token # limit); the response then ends as ``incomplete`` instead of ``completed`` so callers # can tell a cut-short turn from a successful one. Mirrored into the stream's @@ -2554,15 +2658,23 @@ async def handle_update( update: AgentResponseUpdate, *, approval_storage: FunctionApprovalStore | None = None, + complete_message: bool = False, ) -> AsyncGenerator[ResponseStreamEvent]: """Process one agent update: note its finish reason, then handle each of its contents. This is the single entry point for both the plain-agent and the workflow loops, so the - finish reason cannot be forgotten on one of them. + finish reason cannot be forgotten on one of them. Set ``complete_message`` when the + update contains a completed message's independent content parts rather than streaming + deltas; text and refusal parts then retain their original boundaries and citation offsets. """ self.record_finish_reason(update.finish_reason) for content in update.contents: - async for event in self.handle(content, message_id=update.message_id, approval_storage=approval_storage): + async for event in self.handle( + content, + message_id=update.message_id, + approval_storage=approval_storage, + start_new_content_part=complete_message, + ): yield event async def handle( @@ -2571,6 +2683,7 @@ async def handle( message_id: str | None = None, *, approval_storage: FunctionApprovalStore | None = None, + start_new_content_part: bool = False, ) -> AsyncGenerator[ResponseStreamEvent]: """Process a content item, yielding its events. @@ -2582,9 +2695,11 @@ async def handle( merged into it. approval_storage: Used for content types that fall back to one-shot emission (anything not recognized as a streaming delta type) to save/load approval requests. + start_new_content_part: Whether text and refusal content start a new independent + content part instead of continuing the current streaming occurrence. """ if _is_refusal_text_content(content) and content.text is not None: - for event in self._ensure_message_content("refusal", message_id): + for event in self._ensure_message_content("refusal", message_id, start_new_part=start_new_content_part): yield event self._active_message_id = message_id delta = self._consume_text_update(content.text, content.merge_mode) @@ -2592,10 +2707,13 @@ async def handle( yield self._refusal_content.emit_delta(delta) elif content.type == "text" and content.text is not None: - for event in self._ensure_message_content("text", message_id): + for event in self._ensure_message_content("text", message_id, start_new_part=start_new_content_part): yield event self._active_message_id = message_id delta = self._consume_text_update(content.text, content.merge_mode) + for annotation in _annotations_to_output_text(content.annotations): + if annotation not in self._text_annotations: + self._text_annotations.append(annotation) if delta and self._text_content is not None: yield self._text_content.emit_delta(delta) @@ -2905,11 +3023,13 @@ def _ensure_message_content( self, content_type: Literal["text", "refusal"], message_id: str | None, + *, + start_new_part: bool = False, ) -> Generator[ResponseStreamEvent]: message_changed = ( message_id is not None and self._active_message_id is not None and message_id != self._active_message_id ) - if self._active_type == content_type and not message_changed: + if self._active_type == content_type and not message_changed and not start_new_part: return if self._message_item is not None and self._active_type in {"text", "refusal"} and not message_changed: yield from self._close_message_content() @@ -2937,6 +3057,7 @@ def _open_message_content( yield self._refusal_content.emit_added() else: self._text_content = self._message_item.add_text_content() + self._text_annotations.clear() yield self._text_content.emit_added() def _open_reasoning(self, content: Content) -> Generator[ResponseStreamEvent]: @@ -3046,7 +3167,10 @@ def _close(self) -> Generator[ResponseStreamEvent]: def _close_message_content(self) -> Generator[ResponseStreamEvent]: accumulated = "".join(self._accumulated) if self._active_type == "text" and self._text_content is not None: - annotations = self._container_file_annotations_for_text(accumulated) + annotations = list(self._text_annotations) + for annotation in self._container_file_annotations_for_text(accumulated): + if annotation not in annotations: + annotations.append(annotation) yield self._text_content.emit_text_done(accumulated) for annotation in annotations: yield self._text_content.emit_annotation_added(annotation) @@ -3058,6 +3182,7 @@ def _close_message_content(self) -> Generator[ResponseStreamEvent]: self._message_text_annotations[self._text_content.content_index] = annotations yield content_done self._text_content = None + self._text_annotations.clear() elif self._active_type == "refusal" and self._refusal_content is not None: yield self._refusal_content.emit_refusal_done(accumulated) yield self._refusal_content.emit_done() diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index dc9b643bd33..6f164edf958 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -34,6 +34,7 @@ AgentResponse, AgentResponseUpdate, AgentSession, + Annotation, BaseChatClient, ChatMiddlewareLayer, ChatResponse, @@ -94,6 +95,7 @@ CONSENT_ERROR_CODE, ConsentError, _agent_response_updates, # pyright: ignore[reportPrivateUsage] + _annotations_to_output_text, # pyright: ignore[reportPrivateUsage] _await_before_signal, # pyright: ignore[reportPrivateUsage] _container_file_citations_from_function_result, # pyright: ignore[reportPrivateUsage] _is_allowed_oauth_consent_link, # pyright: ignore[reportPrivateUsage] @@ -241,6 +243,65 @@ def _expected_container_file_citation( } +def _url_citation( + *, + url: str = "https://www.python.org/downloads/release/python-31315/", + title: str = "Python 3.13.15", + start_index: int = 0, + end_index: int = 6, +) -> Annotation: + return Annotation( + type="citation", + url=url, + title=title, + annotated_regions=[{"type": "text_span", "start_index": start_index, "end_index": end_index}], + ) + + +def _expected_url_citation( + *, + url: str = "https://www.python.org/downloads/release/python-31315/", + title: str = "Python 3.13.15", + start_index: int = 0, + end_index: int = 6, +) -> dict[str, Any]: + return { + "type": "url_citation", + "url": url, + "title": title, + "start_index": start_index, + "end_index": end_index, + } + + +def _file_citation( + *, + file_id: str = "file_123", + filename: str = "brambleglass-kestrel-st731.txt", + index: int = 0, +) -> Annotation: + return Annotation( + type="citation", + file_id=file_id, + url=filename, + additional_properties={"index": index}, + ) + + +def _expected_file_citation( + *, + file_id: str = "file_123", + filename: str = "brambleglass-kestrel-st731.txt", + index: int = 0, +) -> dict[str, Any]: + return { + "type": "file_citation", + "file_id": file_id, + "filename": filename, + "index": index, + } + + @pytest.mark.parametrize( "raw_citations", [ @@ -274,6 +335,160 @@ def test_container_file_citations_ignore_recursive_json() -> None: assert _container_file_citations_from_function_result(function_result) == [] +def test_framework_annotations_convert_all_responses_citation_variants() -> None: + annotations = [ + Annotation( + type="citation", + title="Source", + url="https://example.com/source", + annotated_regions=[ + {"type": "text_span", "start_index": 0, "end_index": 6}, + {"type": "text_span", "start_index": 12, "end_index": 18}, + ], + ), + _file_citation(), + Annotation( + type="citation", + file_id="file_path_123", + additional_properties={"index": 7}, + ), + Annotation( + type="citation", + file_id="container_file_123", + url="report.csv", + annotated_regions=[{"type": "text_span", "start_index": 20, "end_index": 30}], + additional_properties={"container_id": "container_123"}, + ), + ] + + assert _annotations_to_output_text(annotations) == [ + { + "type": "url_citation", + "url": "https://example.com/source", + "title": "Source", + "start_index": 0, + "end_index": 6, + }, + { + "type": "url_citation", + "url": "https://example.com/source", + "title": "Source", + "start_index": 12, + "end_index": 18, + }, + _expected_file_citation(), + { + "type": "file_path", + "file_id": "file_path_123", + "index": 7, + }, + { + "type": "container_file_citation", + "container_id": "container_123", + "file_id": "container_file_123", + "filename": "report.csv", + "start_index": 20, + "end_index": 30, + }, + ] + + +@pytest.mark.parametrize( + ("coordinate_space", "expected"), + [ + pytest.param(None, [_expected_url_citation()], id="unmarked"), + pytest.param("response", [_expected_url_citation()], id="response"), + pytest.param("source", [], id="source"), + ], +) +def test_framework_annotations_respect_coordinate_space_without_raw_representation( + coordinate_space: str | None, expected: list[dict[str, Any]] +) -> None: + annotation = _url_citation() + if coordinate_space is not None: + annotation["additional_properties"] = {"annotated_regions_coordinate_space": coordinate_space} + annotations = json.loads(json.dumps([annotation])) + original_annotations = copy.deepcopy(annotations) + + assert _annotations_to_output_text(annotations) == expected + assert annotations == original_annotations + + +def test_framework_annotations_preserve_file_citations_with_source_regions() -> None: + annotation = _file_citation() + annotation["annotated_regions"] = [{"type": "text_span", "start_index": 1, "end_index": 3}] + annotation["additional_properties"]["annotated_regions_coordinate_space"] = "source" + original_annotation = copy.deepcopy(annotation) + + assert _annotations_to_output_text([annotation]) == [_expected_file_citation()] + assert annotation == original_annotation + + +def test_framework_annotations_ignore_malformed_runtime_shapes() -> None: + annotations = cast( + Any, + [ + None, + "not-an-annotation", + { + "type": "citation", + "url": "https://example.com/mixed", + "title": "Mixed", + "additional_properties": "not-a-mapping", + "annotated_regions": [ + None, + "not-a-region", + {"type": "source_block", "start_index": 0, "end_index": 5}, + {"type": "text_span", "start_index": 2, "end_index": 7}, + ], + }, + { + "type": "citation", + "url": "https://example.com/string-regions", + "annotated_regions": "not-a-sequence-of-regions", + }, + ], + ) + + assert _annotations_to_output_text(annotations) == [ + _expected_url_citation( + url="https://example.com/mixed", + title="Mixed", + start_index=2, + end_index=7, + ) + ] + assert _annotations_to_output_text(cast(Any, "not-an-annotation-sequence")) == [] + + +@pytest.mark.parametrize( + "annotations", + [ + pytest.param(42, id="integer"), + pytest.param(True, id="boolean"), + pytest.param(object(), id="object"), + ], +) +def test_framework_annotations_ignore_non_sequence_inputs(annotations: Any) -> None: + assert _annotations_to_output_text(annotations) == [] + + +def test_framework_annotations_ignore_unsupported_or_incomplete_values() -> None: + annotations = [ + Annotation(type="citation", url="https://example.com/no-span"), + Annotation( + type="citation", + url="https://example.com/invalid-span", + annotated_regions=[{"type": "text_span", "start_index": 0}], + ), + Annotation(type="citation", file_id="file-without-index"), + Annotation(type="citation", file_id="file-without-index", url="document.txt"), + cast(Annotation, {"type": "other", "url": "https://example.com/unsupported"}), + ] + + assert _annotations_to_output_text(annotations) == [] + + # region Helpers @@ -1940,6 +2155,69 @@ def test_provider_background_forwards_finish_reason_without_output(self) -> None assert updates[0].contents == [] assert updates[0].finish_reason == "content_filter" + @pytest.mark.parametrize("annotate_prefix", [False, True]) + async def test_provider_background_preserves_completed_text_part_citation_offsets( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, annotate_prefix: bool + ) -> None: + response = AgentResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_text( + "See ", annotations=[_url_citation(end_index=3)] if annotate_prefix else None + ), + Content.from_text("Python", annotations=[_url_citation()]), + ], + ) + ], + finish_reason="stop", + ) + original_response = copy.deepcopy(response.to_dict()) + agent = Agent(client=_ServiceStorageRecordingClient()) + + async def run(messages: Any = None, *, session: AgentSession, **kwargs: Any) -> AgentResponse: + del messages, kwargs + session.service_session_id = "private-service-conversation" + return response + + monkeypatch.setattr(agent, "run", MagicMock(side_effect=run)) + server = _make_server( + agent, + session_store=SessionStore(), + history_source="service", + background_source="provider", + options=ResponsesServerOptions(resilient_background=True), + response_store=FileResponseStore(storage_dir=tmp_path), + ) + request = CreateResponse(input="hello", store=True, background=True) + context = ResponseContext(response_id="outer-response", mode_flags=MagicMock()) + + with patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])): + events = [ + event + async for event in server._handle_response(request, context, asyncio.Event()) + if isinstance(event, Mapping) + ] + + completed = next(event["response"] for event in events if event["type"] == "response.completed") + message = next(item for item in completed["output"] if item["type"] == "message") + parts = [part for part in message["content"] if part["type"] == "output_text"] + assert len(parts) == len(message["content"]) + assert [part["text"] for part in parts] == ["See ", "Python"] + assert parts[0]["annotations"] == ([_expected_url_citation(end_index=3)] if annotate_prefix else []) + assert parts[1]["annotations"] == [_expected_url_citation()] + annotation_events = [event for event in events if event["type"] == "response.output_text.annotation.added"] + assert [event["content_index"] for event in annotation_events] == ([0, 1] if annotate_prefix else [1]) + assert annotation_events[-1]["annotation"] == _expected_url_citation() + content_done = next( + event for event in events if event["type"] == "response.content_part.done" and event["content_index"] == 1 + ) + assert content_done["part"] == parts[1] + output_done = next(event for event in events if event["type"] == "response.output_item.done") + assert output_done["item"] == message + assert response.to_dict() == original_response + async def test_provider_background_keeps_private_token_under_outer_response( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -3759,6 +4037,33 @@ async def test_function_result_container_file_citation_is_added_to_output_text(s assert preparing["annotations"] == [] assert final["annotations"] == [_expected_container_file_citation()] + async def test_framework_url_and_file_citations_are_added_to_output_text(self) -> None: + source_citation = _url_citation(start_index=0, end_index=1) + source_citation["additional_properties"] = {"annotated_regions_coordinate_space": "source"} + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate( + contents=[ + Content.from_text( + "Python sources", + annotations=[source_citation, _url_citation(), _file_citation()], + ) + ], + role="assistant", + message_id="msg_citations", + ) + ] + ) + + resp = await _post(_make_server(agent), stream=False) + + assert resp.status_code == 200 + message = next(item for item in resp.json()["output"] if item["type"] == "message") + assert message["content"][0]["annotations"] == [ + _expected_url_citation(), + _expected_file_citation(), + ] + async def test_latest_pending_container_file_citation_replaces_same_filename(self) -> None: agent = _make_agent( stream_updates=[ @@ -4582,6 +4887,60 @@ async def test_container_file_citation_is_consistent_across_stream_events(self) ) assert final_message["content"][0]["annotations"] == [expected_annotation] + async def test_framework_url_citation_is_consistent_across_stream_events(self) -> None: + source_citation = _url_citation(start_index=0, end_index=1) + source_citation["additional_properties"] = {"annotated_regions_coordinate_space": "source"} + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate( + contents=[Content.from_text("See ", annotations=[source_citation])], + role="assistant", + message_id="msg_citation", + ), + AgentResponseUpdate( + contents=[ + Content.from_text( + "Python", + annotations=[_url_citation(start_index=4, end_index=10)], + ) + ], + role="assistant", + message_id="msg_citation", + ), + AgentResponseUpdate( + contents=[Content.from_text("", annotations=[_url_citation(start_index=4, end_index=10)])], + role="assistant", + message_id="msg_citation", + ), + AgentResponseUpdate( + contents=[ + Content.from_text( + "See Python", + annotations=[_url_citation(start_index=4, end_index=10)], + merge_mode="replace", + ) + ], + role="assistant", + message_id="msg_citation", + ), + ] + ) + + resp = await _post(_make_server(agent), stream=True) + + assert resp.status_code == 200 + events = _parse_sse_events(resp.text) + expected_annotation = _expected_url_citation(start_index=4, end_index=10) + annotation_events = [event for event in events if event["event"] == "response.output_text.annotation.added"] + assert [event["data"]["annotation"] for event in annotation_events] == [expected_annotation] + + content_done = next(event for event in events if event["event"] == "response.content_part.done") + assert content_done["data"]["part"]["annotations"] == [expected_annotation] + + completed = next(event["data"]["response"] for event in events if event["event"] == "response.completed") + final_message = next(item for item in completed["output"] if item["type"] == "message") + assert final_message["content"][0]["annotations"] == [expected_annotation] + async def test_pending_container_file_citation_survives_tracker_recovery(self) -> None: stream = ResponseEventStream(response_id="resp_container_file") stream.emit_created()