Python: Preserve Foundry Toolbox container file citations

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,
     }

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