diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index dc4795bb347..367bdcdc7f5 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -934,8 +934,12 @@ def test_parse_tool_result_content_modes_for_complementary_and_duplicate_payload async def test_generated_mcp_tool_preserves_complete_host_payload_once() -> None: """The generated FunctionTool path retains one complete, persistent Host payload.""" + file_meta = { + "container_id": "cntr_123", + "container_file_citations": '[{"file_id":"cfile_123","filename":"result.txt"}]', + } mcp_result = types.CallToolResult( - content=[types.TextContent(type="text", text="Summary")], + content=[types.TextContent(type="text", text="Summary", _meta=file_meta)], structuredContent={"image_url": "https://example.test/widget.png"}, isError=False, _meta={"widget": "image"}, @@ -948,7 +952,7 @@ async def test_generated_mcp_tool_preserves_complete_host_payload_once() -> None assert function_result.items is not None expected_host_payload = { "_meta": {"widget": "image"}, - "content": [{"type": "text", "text": "Summary"}], + "content": [{"type": "text", "text": "Summary", "_meta": file_meta}], "structuredContent": {"image_url": "https://example.test/widget.png"}, "isError": False, } 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 fc926f20dc7..93bc39c7fdd 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -49,6 +49,7 @@ WorkflowAgent, add_usage_details, ) +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 from azure.ai.agentserver.core import get_request_context @@ -64,6 +65,7 @@ ComputerAction, ComputerCallSafetyCheckParam, ComputerScreenshotImage, + ContainerFileCitationBody, CreateResponse, FunctionShellAction, FunctionShellCallOutputContent, @@ -137,11 +139,93 @@ _HOSTED_CONVERSATION_COMMITTED_KEY = "_foundry_conversation_committed" _HOSTED_PROVIDER_OUTPUT_COUNT_KEY = "_foundry_provider_output_count" _HOSTED_PROVIDER_USAGE_KEY = "_foundry_provider_usage" +_PENDING_CONTAINER_FILE_CITATIONS_KEY = "_foundry_pending_container_file_citations" +_CONTAINER_FILE_CITATIONS_KEY = "container_file_citations" _HistorySource = Literal["agent_server", "agent", "service"] _AGENT_SOURCE_UNSET = object() +@dataclass(frozen=True) +class _ContainerFileCitation: + container_id: str + file_id: str + filename: str + + +def _has_container_filename_boundaries(text: str, start_index: int, end_index: int) -> bool: + """Return whether a filename match is not embedded in another filename-like token.""" + before = text[start_index - 1] if start_index > 0 else None + after = text[end_index] if end_index < len(text) else None + return (before is None or not (before.isalnum() or before in "._-")) and ( + after is None or not (after.isalnum() or after in "._-") + ) + + +def _container_file_citations_from_function_result(content: Content) -> list[_ContainerFileCitation]: + """Extract validated container file citations from a core-preserved MCP Host payload.""" + raw_payload = content.additional_properties.get(_MCP_TOOL_RESULT_HOST_PAYLOAD_KEY) + if not isinstance(raw_payload, Mapping): + return [] + payload = cast(Mapping[str, Any], raw_payload) + + metadata_blocks: list[Mapping[str, Any]] = [] + raw_result_meta = payload.get("_meta") + if isinstance(raw_result_meta, Mapping): + metadata_blocks.append(cast(Mapping[str, Any], raw_result_meta)) + + raw_content = payload.get("content") + if isinstance(raw_content, Sequence) and not isinstance(raw_content, (str, bytes, bytearray)): + for raw_item in cast(Sequence[Any], raw_content): + if not isinstance(raw_item, Mapping): + continue + item = cast(Mapping[str, Any], raw_item) + raw_item_meta = item.get("_meta") + if isinstance(raw_item_meta, Mapping): + metadata_blocks.append(cast(Mapping[str, Any], raw_item_meta)) + + citations: list[_ContainerFileCitation] = [] + for metadata in metadata_blocks: + raw_citations = metadata.get(_CONTAINER_FILE_CITATIONS_KEY) + if isinstance(raw_citations, str): + try: + raw_citations = json.loads(raw_citations) + except (json.JSONDecodeError, RecursionError): + continue + + citation_values: list[Any] + if isinstance(raw_citations, Mapping): + citation_values = [cast(Mapping[str, Any], raw_citations)] + elif isinstance(raw_citations, Sequence) and not isinstance(raw_citations, (str, bytes, bytearray)): + citation_values = list(cast(Sequence[Any], raw_citations)) + else: + continue + + fallback_container_id = metadata.get("container_id") + for raw_citation in citation_values: + if not isinstance(raw_citation, Mapping): + continue + citation = cast(Mapping[str, Any], raw_citation) + container_id = citation.get("container_id") or fallback_container_id + file_id = citation.get("file_id") + filename = citation.get("filename") + if not isinstance(container_id, str) or not container_id: + continue + if not isinstance(file_id, str) or not file_id: + continue + if not isinstance(filename, str) or not filename: + continue + citations.append( + _ContainerFileCitation( + container_id=container_id, + file_id=file_id, + filename=filename, + ) + ) + + return citations + + def _is_refusal_text_content(content: Content) -> bool: return content.type == "text" and content.additional_properties.get(_MODEL_OUTPUT_KIND_KEY) == _MODEL_OUTPUT_REFUSAL @@ -1189,6 +1273,7 @@ async def _handle_prepared_response( # selecting the terminal event. Same ``client_cancelled`` gate as above. return + tracker.discard_pending_container_file_citations() incomplete_reason = tracker.incomplete_reason if tracker.oauth_consent_requested or incomplete_reason is not None: yield response_event_stream.emit_incomplete(reason=incomplete_reason, usage=tracker.usage) @@ -2122,6 +2207,7 @@ def _emit_failure( yield from tracker.close() except Exception: logger.exception("Error while closing streaming tracker after failure") + tracker.discard_pending_container_file_citations() message = str(ex) or type(ex).__name__ yield response_event_stream.emit_failed(message=message, usage=tracker.usage if tracker is not None else None) @@ -2177,6 +2263,8 @@ def __init__( self._outstanding_function_calls: dict[str, str | None] = {} 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]] = {} # 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 @@ -2187,6 +2275,33 @@ def __init__( if isinstance(persisted_reason, str): with suppress(ValueError): self._incomplete_reason = ResponseIncompleteReason(persisted_reason) + persisted_citations = stream.internal_metadata.get(_PENDING_CONTAINER_FILE_CITATIONS_KEY) + if persisted_citations is not None: + if not isinstance(persisted_citations, Sequence) or isinstance( + persisted_citations, (str, bytes, bytearray) + ): + raise RuntimeError("The persisted container file citations are invalid.") + for raw_citation in cast(Sequence[Any], persisted_citations): + if not isinstance(raw_citation, Mapping): + raise RuntimeError("The persisted container file citations are invalid.") + citation = cast(Mapping[str, Any], raw_citation) + container_id = citation.get("container_id") + file_id = citation.get("file_id") + filename = citation.get("filename") + if ( + not isinstance(container_id, str) + or not container_id + or not isinstance(file_id, str) + or not file_id + or not isinstance(filename, str) + or not filename + ): + raise RuntimeError("The persisted container file citations are invalid.") + self._pending_container_file_citations[filename] = _ContainerFileCitation( + container_id=container_id, + file_id=file_id, + filename=filename, + ) for item in stream.response.get("output", []): if not isinstance(item, Mapping): continue @@ -2344,6 +2459,11 @@ async def handle( elif content.type == "function_result": for event in self._close(): yield event + citations = _container_file_citations_from_function_result(content) + if citations: + for citation in citations: + self._pending_container_file_citations[citation.filename] = citation + self._persist_pending_container_file_citations() async for event in self._stream.output_item_function_call_output( content.call_id, # type: ignore[arg-type] _json_safe_to_str(content.result), @@ -2590,9 +2710,14 @@ async def handle( logger.warning(f"Content type '{content.type}' is not supported yet. This is usually safe to ignore.") def close(self) -> Generator[ResponseStreamEvent]: - """Close any remaining active builder.""" + """Flush any remaining active builder without discarding recoverable state.""" yield from self._close() + def discard_pending_container_file_citations(self) -> None: + """Discard unmatched citations immediately before a terminal response event.""" + self._pending_container_file_citations.clear() + self._persist_pending_container_file_citations() + # -- Private open/close helpers -- def _ensure_message_content( @@ -2613,6 +2738,7 @@ def _ensure_message_content( yield from self._open_message(content_type) def _open_message(self, content_type: Literal["text", "refusal"]) -> Generator[ResponseStreamEvent]: + self._message_text_annotations.clear() self._message_item = self._stream.add_output_item_message() yield self._message_item.emit_added() yield from self._open_message_content(content_type) @@ -2680,8 +2806,27 @@ def _close(self) -> Generator[ResponseStreamEvent]: if self._active_type in {"text", "refusal"}: yield from self._close_message_content() if self._message_item is not None: - yield self._message_item.emit_done() + message_done = self._message_item.emit_done() + if self._message_text_annotations: + message_done_dict = cast(dict[str, Any], message_done) + item = cast(dict[str, Any], message_done_dict["item"]) + content_parts = cast(list[dict[str, Any]], item["content"]) + for content_index, annotations in self._message_text_annotations.items(): + if content_index >= len(content_parts): + raise RuntimeError("Container file citation content index is out of range.") + content_parts[content_index]["annotations"] = annotations + + response_output = self._stream.response.get("output") + output_index = self._message_item.output_index + if not isinstance(response_output, list): + raise RuntimeError("Container file citation output index is out of range.") + response_output_items = cast(list[Any], response_output) + if output_index >= len(response_output_items): + raise RuntimeError("Container file citation output index is out of range.") + response_output_items[output_index] = item + yield message_done self._message_item = None + self._message_text_annotations.clear() elif self._active_type == "text_reasoning" and self._summary_part and self._reasoning_item: accumulated = "".join(self._accumulated) @@ -2720,8 +2865,17 @@ 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) yield self._text_content.emit_text_done(accumulated) - yield self._text_content.emit_done() + for annotation in annotations: + yield self._text_content.emit_annotation_added(annotation) + content_done = self._text_content.emit_done() + if annotations: + content_done_dict = cast(dict[str, Any], content_done) + part = cast(dict[str, Any], content_done_dict["part"]) + part["annotations"] = annotations + self._message_text_annotations[self._text_content.content_index] = annotations + yield content_done self._text_content = None elif self._active_type == "refusal" and self._refusal_content is not None: yield self._refusal_content.emit_refusal_done(accumulated) @@ -2731,6 +2885,64 @@ def _close_message_content(self) -> Generator[ResponseStreamEvent]: self._active_id = None self._accumulated.clear() + def _container_file_annotations_for_text(self, text: str) -> list[ContainerFileCitationBody]: + candidates: list[tuple[int, int, str, _ContainerFileCitation]] = [] + for filename, citation in self._pending_container_file_citations.items(): + search_index = 0 + while True: + start_index = text.find(filename, search_index) + if start_index < 0: + break + end_index = start_index + len(filename) + if _has_container_filename_boundaries(text, start_index, end_index): + candidates.append((start_index, end_index, filename, citation)) + search_index = start_index + 1 + + candidates.sort(key=lambda match: (-(match[1] - match[0]), match[0], match[2])) + selected: list[tuple[int, int, str, _ContainerFileCitation]] = [] + selected_filenames: set[str] = set() + for candidate in candidates: + start_index, end_index, filename, _ = candidate + if filename in selected_filenames: + continue + if any( + start_index < selected_end and selected_start < end_index + for selected_start, selected_end, _, _ in selected + ): + continue + selected.append(candidate) + selected_filenames.add(filename) + + annotations: list[ContainerFileCitationBody] = [] + for start_index, end_index, filename, citation in sorted(selected, key=lambda match: match[0]): + annotations.append( + ContainerFileCitationBody( + type="container_file_citation", + container_id=citation.container_id, + file_id=citation.file_id, + filename=citation.filename, + start_index=start_index, + end_index=end_index, + ) + ) + del self._pending_container_file_citations[filename] + if annotations: + self._persist_pending_container_file_citations() + return annotations + + def _persist_pending_container_file_citations(self) -> None: + if not self._pending_container_file_citations: + self._stream.internal_metadata.pop(_PENDING_CONTAINER_FILE_CITATIONS_KEY, None) + return + self._stream.internal_metadata[_PENDING_CONTAINER_FILE_CITATIONS_KEY] = [ + { + "container_id": citation.container_id, + "file_id": citation.file_id, + "filename": citation.filename, + } + for citation in self._pending_container_file_citations.values() + ] + # endregion diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 6fb67cdc7a3..d3c518dfe98 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -58,6 +58,7 @@ executor, tool, ) +from agent_framework._mcp import _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY # pyright: ignore[reportPrivateUsage] from agent_framework.ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore from agent_framework.openai import OpenAIChatClient, OpenAIChatOptions, OpenAIContinuationToken from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context @@ -94,6 +95,7 @@ ConsentError, _agent_response_updates, # 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] _item_to_message, # pyright: ignore[reportPrivateUsage] _json_safe_to_str, # pyright: ignore[reportPrivateUsage] @@ -151,6 +153,121 @@ def _make_function_approval_request_content( return Content.from_function_approval_request(request_id, function_call) +def _mcp_function_result(call_id: str, *, item_meta: Mapping[str, Any]) -> Content: + return Content.from_function_result( + call_id, + result="Created the requested file.", + additional_properties={ + _MCP_TOOL_RESULT_HOST_PAYLOAD_KEY: { + "content": [ + { + "type": "text", + "text": "Created the requested file.", + "_meta": dict(item_meta), + } + ], + "isError": False, + } + }, + ) + + +def _container_file_result( + call_id: str, + *, + container_id: str = "cntr_123", + file_id: str = "cfile_123", + filename: str = "result.txt", +) -> Content: + return _mcp_function_result( + call_id, + item_meta={ + "container_id": container_id, + "container_file_citations": json.dumps([ + { + "file_id": file_id, + "filename": filename, + } + ]), + }, + ) + + +def _container_file_citation_updates() -> list[AgentResponseUpdate]: + return [ + AgentResponseUpdate( + contents=[Content.from_function_call("call_1", "code_interpreter", arguments={})], + role="assistant", + message_id="msg_call", + ), + AgentResponseUpdate( + contents=[_container_file_result("call_1")], + role="tool", + message_id="msg_result", + ), + AgentResponseUpdate( + contents=[Content.from_text("Preparing the download.")], + role="assistant", + message_id="msg_preparing", + ), + AgentResponseUpdate( + contents=[Content.from_text("Download result.txt")], + role="assistant", + message_id="msg_final", + ), + ] + + +def _expected_container_file_citation( + *, + container_id: str = "cntr_123", + file_id: str = "cfile_123", + filename: str = "result.txt", + start_index: int = 9, +) -> dict[str, Any]: + return { + "type": "container_file_citation", + "container_id": container_id, + "file_id": file_id, + "filename": filename, + "start_index": start_index, + "end_index": start_index + len(filename), + } + + +@pytest.mark.parametrize( + "raw_citations", + [ + "not-json", + [{"file_id": "cfile_123"}], + 123, + ], +) +def test_container_file_citations_ignore_malformed_metadata(raw_citations: Any) -> None: + function_result = _mcp_function_result( + "call_1", + item_meta={ + "container_id": "cntr_123", + "container_file_citations": raw_citations, + }, + ) + + assert _container_file_citations_from_function_result(function_result) == [] + + +def test_container_file_citations_ignore_recursive_json() -> None: + function_result = _mcp_function_result( + "call_1", + item_meta={ + "container_id": "cntr_123", + "container_file_citations": "[]", + }, + ) + + with patch("agent_framework_foundry_hosting._responses.json.loads", side_effect=RecursionError): + assert _container_file_citations_from_function_result(function_result) == [] + + # region Helpers @@ -3588,6 +3705,153 @@ async def test_function_call_and_result(self) -> None: assert "function_call_output" in types assert "message" in types + async def test_function_result_container_file_citation_is_added_to_output_text(self) -> None: + resp = await _post(_make_server(_make_agent(stream_updates=_container_file_citation_updates())), stream=False) + + assert resp.status_code == 200 + messages = [item for item in resp.json()["output"] if item["type"] == "message"] + preparing = next( + part for message in messages for part in message["content"] if part.get("text") == "Preparing the download." + ) + final = next( + part for message in messages for part in message["content"] if part.get("text") == "Download result.txt" + ) + assert preparing["annotations"] == [] + assert final["annotations"] == [_expected_container_file_citation()] + + async def test_latest_pending_container_file_citation_replaces_same_filename(self) -> None: + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate( + contents=[Content.from_function_call("call_1", "code_interpreter", arguments={})], + role="assistant", + ), + AgentResponseUpdate( + contents=[_container_file_result("call_1", container_id="cntr_old", file_id="cfile_old")], + role="tool", + ), + AgentResponseUpdate( + contents=[Content.from_text("Still working.")], + role="assistant", + message_id="msg_working", + ), + AgentResponseUpdate( + contents=[Content.from_function_call("call_2", "code_interpreter", arguments={})], + role="assistant", + ), + AgentResponseUpdate( + contents=[_container_file_result("call_2", container_id="cntr_new", file_id="cfile_new")], + role="tool", + ), + AgentResponseUpdate( + contents=[Content.from_text("Download result.txt")], + role="assistant", + message_id="msg_final", + ), + ] + ) + + resp = await _post(_make_server(agent), stream=False) + + final_message = next( + item + for item in resp.json()["output"] + if item["type"] == "message" and item["content"][0].get("text") == "Download result.txt" + ) + assert final_message["content"][0]["annotations"] == [ + _expected_container_file_citation(container_id="cntr_new", file_id="cfile_new") + ] + + async def test_container_file_citations_require_complete_non_overlapping_filename_matches(self) -> None: + updates: list[AgentResponseUpdate] = [] + for suffix, filename in [ + ("a", "a"), + ("data", "data.csv"), + ("metadata", "metadata.csv"), + ("report", "report.txt"), + ("my_report", "my report.txt"), + ]: + updates.extend([ + AgentResponseUpdate( + contents=[Content.from_function_call(f"call_{suffix}", "code_interpreter", arguments={})], + role="assistant", + ), + AgentResponseUpdate( + contents=[ + _container_file_result( + f"call_{suffix}", + container_id=f"cntr_{suffix}", + file_id=f"cfile_{suffix}", + filename=filename, + ) + ], + role="tool", + ), + ]) + for suffix, text in [ + ("preparing", "Preparing the download."), + ("metadata", "Download metadata.csv"), + ("data", "Download data.csv"), + ("a", "Download a"), + ("my_report", "Download my report.txt"), + ("report", "Download report.txt"), + ]: + updates.append( + AgentResponseUpdate( + contents=[Content.from_text(text)], + role="assistant", + message_id=f"msg_{suffix}", + ) + ) + + agent = _make_agent(stream_updates=updates) + + resp = await _post(_make_server(agent), stream=False) + + parts_by_text = { + part["text"]: part + for item in resp.json()["output"] + if item["type"] == "message" + for part in item["content"] + if part["type"] == "output_text" + } + assert parts_by_text["Preparing the download."]["annotations"] == [] + assert parts_by_text["Download metadata.csv"]["annotations"] == [ + _expected_container_file_citation( + container_id="cntr_metadata", + file_id="cfile_metadata", + filename="metadata.csv", + ) + ] + assert parts_by_text["Download data.csv"]["annotations"] == [ + _expected_container_file_citation( + container_id="cntr_data", + file_id="cfile_data", + filename="data.csv", + ) + ] + assert parts_by_text["Download a"]["annotations"] == [ + _expected_container_file_citation( + container_id="cntr_a", + file_id="cfile_a", + filename="a", + ) + ] + assert parts_by_text["Download my report.txt"]["annotations"] == [ + _expected_container_file_citation( + container_id="cntr_my_report", + file_id="cfile_my_report", + filename="my report.txt", + ) + ] + assert parts_by_text["Download report.txt"]["annotations"] == [ + _expected_container_file_citation( + container_id="cntr_report", + file_id="cfile_report", + filename="report.txt", + ) + ] + async def test_native_computer_call_and_result(self) -> None: item_id = IdGenerator.new_computer_call_item_id() actions: list[dict[str, Any]] = [ @@ -4183,6 +4447,83 @@ async def test_basic_text_streaming(self) -> None: assert len(done_events) == 1 assert done_events[0]["data"]["text"] == "Hello world!" + async def test_container_file_citation_is_consistent_across_stream_events(self) -> None: + resp = await _post( + _make_server(_make_agent(stream_updates=_container_file_citation_updates())), + stream=True, + ) + + assert resp.status_code == 200 + events = _parse_sse_events(resp.text) + expected_annotation = _expected_container_file_citation() + + annotation_event = next(event for event in events if event["event"] == "response.output_text.annotation.added") + assert annotation_event["data"]["annotation"] == expected_annotation + text_done_index = next( + index + for index, event in enumerate(events) + if event["event"] == "response.output_text.done" and event["data"]["text"] == "Download result.txt" + ) + annotation_index = events.index(annotation_event) + content_done_index = next( + index + for index, event in enumerate(events) + if event["event"] == "response.content_part.done" + and event["data"]["part"].get("text") == "Download result.txt" + ) + assert text_done_index < annotation_index < content_done_index + + content_done = events[content_done_index] + assert content_done["data"]["part"]["annotations"] == [expected_annotation] + + output_done = next( + event + for event in events + if event["event"] == "response.output_item.done" + and event["data"]["item"].get("type") == "message" + and event["data"]["item"]["content"][0].get("text") == "Download result.txt" + ) + assert output_done["data"]["item"]["content"][0]["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" and item["content"][0].get("text") == "Download result.txt" + ) + 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() + stream.emit_in_progress() + tracker = _OutputItemTracker(stream) + + function_call = Content.from_function_call("call_1", "code_interpreter", arguments={}) + _ = [event async for event in tracker.handle(function_call)] + _ = [event async for event in tracker.handle(_container_file_result("call_1"))] + _ = list(tracker.close()) + _ = stream.checkpoint() + + recovered = _OutputItemTracker(stream) + events = [ + event + async for event in recovered.handle( + Content.from_text("Download result.txt"), + message_id="msg_final", + ) + ] + events.extend(recovered.close()) + + annotation_event = next(event for event in events if event["type"] == "response.output_text.annotation.added") + assert annotation_event["annotation"] == _expected_container_file_citation() + final_message = next( + item + for item in stream.response["output"] + if item["type"] == "message" and item["content"][0].get("text") == "Download result.txt" + ) + assert final_message["content"][0]["annotations"] == [_expected_container_file_citation()] + async def test_computer_call_streaming_emits_complete_items_once(self) -> None: item_id = "cu_" + "a" * 32 call = Content.from_computer_tool_call(