Implement citation handling improvements in response handlers and orchestration manager

- Add functions to manage streaming citation buffers and clean citations.
- Update orchestration manager to clear citation buffers and clean final text.
- Enhance unit tests for citation cleaning and streaming callbacks.

Author: Ayaz-Microsoft

src/backend/callbacks/response_handlers.py

@@ -17,6 +17,38 @@
 
 logger = logging.getLogger(__name__)
 
+_stream_citation_buffers: dict[tuple[str, str], str] = {}
+
+
+def clear_streaming_citation_buffers(
+    user_id: str,
+    agent_id: str | None = None,
+) -> None:
+    """Discard partial citation markers retained between streaming chunks."""
+    for key in [
+        key
+        for key in _stream_citation_buffers
+        if key[0] == user_id and (agent_id is None or key[1] == agent_id)
+    ]:
+        _stream_citation_buffers.pop(key, None)
+
+
+def _split_trailing_partial_citation(text: str) -> tuple[str, str]:
+    """Hold a trailing citation prefix until a later chunk completes it."""
+    for opener, closer in (("[", "]"), ("【", "】")):
+        open_index = text.rfind(opener)
+        if open_index == -1 or text.find(closer, open_index) != -1:
+            continue
+
+        candidate = text[open_index + 1:]
+        if not candidate.strip() or re.fullmatch(
+            r"\s*\d+(?:\s*:\s*\d*)?(?:\s*[|†]?\s*[a-zA-Z]*)?",
+            candidate,
+        ):
+            return text[:open_index], text[open_index:]
+
+    return text, ""
+
 
 def format_agent_display_name(raw_name: str) -> str:
     """Convert raw agent IDs (e.g. 'HRHelperAgent', 'hr_helper_agent') to
@@ -61,7 +93,12 @@ def clean_citations(text: str) -> str:
     """Remove citation markers from agent responses while preserving formatting."""
     if not text:
         return text
-    text = re.sub(r'\[\d+:\d+\|source\]', '', text)
+    text = re.sub(
+        r'\[\s*\d+\s*:\s*\d+\s*[|†]\s*source\s*\]',
+        '',
+        text,
+        flags=re.IGNORECASE,
+    )
     text = re.sub(r'\[\s*source\s*\]', '', text, flags=re.IGNORECASE)
     text = re.sub(r'\[\d+\]', '', text)
     text = re.sub(r'【[^】]*】', '', text)
@@ -164,7 +201,17 @@ async def streaming_agent_response_callback(
                     collected.append(str(txt))
             chunk_text = "".join(collected) if collected else ""
 
-        cleaned = clean_citations(chunk_text or "")
+        buffer_key = (user_id, agent_id)
+        combined = _stream_citation_buffers.pop(buffer_key, "") + (chunk_text or "")
+
+        if is_final:
+            emittable, _ = _split_trailing_partial_citation(combined)
+        else:
+            emittable, held_back = _split_trailing_partial_citation(combined)
+            if held_back:
+                _stream_citation_buffers[buffer_key] = held_back
+
+        cleaned = clean_citations(emittable)
 
         contents = getattr(update, "contents", []) or []
         tool_calls = _extract_tool_calls_from_contents(contents)

src/backend/orchestration/orchestration_manager.py

@@ -17,6 +17,8 @@
                                             MagenticPlanReviewRequest)
 from agents.agent_factory import AgentFactory
 from callbacks.response_handlers import (agent_response_callback,
+                                         clean_citations,
+                                         clear_streaming_citation_buffers,
                                          format_agent_display_name,
                                          streaming_agent_response_callback)
 from common.config.app_config import config
@@ -395,6 +397,7 @@ async def run_orchestration(self, user_id: str, input_task) -> None:
             final_output_ref: list = [None]
             orchestrator_chunks: list[str] = []
             current_streaming_agent_ref: list = [None]
+            clear_streaming_citation_buffers(user_id)
 
             # Collect participant names for plan conversion
             participant_names = [
@@ -479,6 +482,9 @@ async def run_orchestration(self, user_id: str, input_task) -> None:
             # accumulated orchestrator streaming chunks.
             final_text = final_output_ref[0] or "".join(orchestrator_chunks)
 
+            # The manager's final answer bypasses the participant callbacks.
+            final_text = clean_citations(final_text)
+
             # Repair collapsed markdown tables before rendering (Bug 47810).
             final_text = _normalize_markdown_tables(final_text)
 
@@ -548,6 +554,7 @@ async def run_orchestration(self, user_id: str, input_task) -> None:
             raise
 
         finally:
+            clear_streaming_citation_buffers(user_id)
             # Clean up MCP connections to avoid noisy cross-task
             # RuntimeError from anyio when async generators are GC'd.
             await self._cleanup_workflow_mcp(user_id)
@@ -1131,6 +1138,7 @@ async def _process_event_stream(
                             if isinstance(msg, Message) and msg.text:
                                 final_output_ref[0] = msg.text
                     else:
+                        clear_streaming_citation_buffers(user_id, agent_id)
                         for msg in event.data:
                             if isinstance(msg, Message) and msg.text:
                                 try:

src/tests/backend/callbacks/test_response_handlers.py

@@ -132,6 +132,7 @@ def __init__(self, text="", role="assistant", author_name=""):
 from backend.callbacks.response_handlers import (
     _extract_tool_calls_from_contents, _is_function_call_item,
     agent_response_callback, clean_citations,
+    clear_streaming_citation_buffers,
     format_agent_display_name,
     streaming_agent_response_callback)
 
@@ -166,6 +167,18 @@ def test_clean_citations_numeric_source(self):
         expected = "This is text  with citations."
         assert clean_citations(text) == expected
 
+    def test_clean_citations_foundry_numeric_source(self):
+        """Test cleaning the Azure Foundry [1:2†source] citation format."""
+        text = "This is text [5:0†source] with citations."
+        expected = "This is text  with citations."
+        assert clean_citations(text) == expected
+
+    def test_clean_citations_foundry_numeric_source_with_spacing(self):
+        """Test cleaning Foundry citations with optional spacing and casing."""
+        text = "This is text [ 5 : 0 † SOURCE ] with citations."
+        expected = "This is text  with citations."
+        assert clean_citations(text) == expected
+
     def test_clean_citations_source_only(self):
         """Test cleaning [source] format citations."""
         text = "Text with [source] citation."
@@ -440,7 +453,7 @@ def test_agent_response_callback_with_chat_message(self, mock_time, mock_create_
 
         # Create an instance of our MockChatMessage
         mock_message = MockChatMessage()
-        mock_message.text = "Test message with citations [1:2|source]"
+        mock_message.text = "Test message with citations [5:0†source]"
         mock_message.author_name = "TestAgent"
         mock_message.role = "assistant"
 
@@ -573,7 +586,7 @@ async def test_streaming_callback_no_user_id(self):
     async def test_streaming_callback_with_text(self):
         """Test streaming callback with update that has text."""
         mock_update = Mock()
-        mock_update.text = "Test streaming text [source]"
+        mock_update.text = "Test streaming text [5:0†source]"
         mock_update.contents = []
 
         with patch('backend.callbacks.response_handlers.AgentMessageStreaming') as mock_streaming:
@@ -596,6 +609,24 @@ async def test_streaming_callback_with_text(self):
                 message_type=WebsocketMessageType.AGENT_MESSAGE_STREAMING
             )
 
+    @pytest.mark.asyncio
+    async def test_streaming_callback_cleans_citation_split_across_chunks(self):
+        """A split citation marker must never reach the thinking-process UI."""
+        clear_streaming_citation_buffers("user_456")
+        first_update = Mock(text="Test streaming text [5:0†sou", contents=[])
+        second_update = Mock(text="rce] continues.", contents=[])
+
+        with patch('backend.callbacks.response_handlers.AgentMessageStreaming') as mock_streaming:
+            await streaming_agent_response_callback(
+                "agent_123", first_update, False, user_id="user_456"
+            )
+            await streaming_agent_response_callback(
+                "agent_123", second_update, False, user_id="user_456"
+            )
+
+        assert mock_streaming.call_args_list[0].kwargs["content"] == "Test streaming text "
+        assert mock_streaming.call_args_list[1].kwargs["content"] == " continues."
+
     @pytest.mark.asyncio
     async def test_streaming_callback_no_text_with_contents(self):
         """Test streaming callback when update has no text but has contents with text.

src/tests/backend/orchestration/test_orchestration_manager.py

@@ -217,6 +217,10 @@ def __init__(self):
 
 sys.modules['callbacks.response_handlers'] = Mock(
     agent_response_callback=Mock(),
+    clean_citations=Mock(
+        side_effect=lambda text: text.replace("[5:0†source]", "")
+    ),
+    clear_streaming_citation_buffers=Mock(),
     streaming_agent_response_callback=AsyncMock(),
 )
 
@@ -316,6 +320,7 @@ async def get_agents(self, user_id, team_config_input, memory_store):
 orchestration_config = sys.modules['orchestration.connection_config'].orchestration_config
 agent_response_callback = sys.modules['callbacks.response_handlers'].agent_response_callback
 streaming_agent_response_callback = sys.modules['callbacks.response_handlers'].streaming_agent_response_callback
+clear_streaming_citation_buffers = sys.modules['callbacks.response_handlers'].clear_streaming_citation_buffers
 
 
 # =========================================================================
@@ -606,6 +611,7 @@ def setup_method(self):
         agent_response_callback.reset_mock()
         streaming_agent_response_callback.reset_mock()
         streaming_agent_response_callback.side_effect = None
+        clear_streaming_citation_buffers.reset_mock()
         mock_wait_approval.reset_mock()
         mock_wait_approval.return_value = MockPlanApprovalResponse(approved=True, m_plan_id="test-plan-id")
         mock_convert.reset_mock()
@@ -661,6 +667,31 @@ async def test_given_executor_completed_when_run_then_captures_final_text(self):
         sent_message = call_args[0][0]
         assert sent_message["data"]["content"] == "Final answer text"
 
+    @pytest.mark.asyncio
+    async def test_given_manager_citation_when_run_then_final_text_is_cleaned(self):
+        final_msg = MockMessage(text="Final answer [5:0†source]")
+        events = [
+            _make_event(
+                "executor_completed",
+                data=[final_msg],
+                executor_id="magentic_orchestrator",
+            ),
+        ]
+        mock_workflow = Mock()
+        mock_workflow.run = Mock(return_value=_async_iter(events))
+        mock_workflow._executors = {}
+        mock_workflow.executors = {}
+        mock_workflow.get_executors_list.return_value = []
+        orchestration_config.get_current_orchestration.return_value = mock_workflow
+
+        await OrchestrationManager().run_orchestration(
+            user_id="user-1",
+            input_task="do stuff",
+        )
+
+        sent_message = connection_config.send_status_update_async.call_args_list[-1][0][0]
+        assert sent_message["data"]["content"] == "Final answer "
+
     @pytest.mark.asyncio
     async def test_given_agent_completed_event_when_run_then_calls_agent_callback(self):
         # Arrange
@@ -1071,4 +1102,3 @@ def test_given_empty_or_none_when_normalized_then_returns_input(self):
 
     def test_given_non_table_pipe_line_when_reflowed_then_returns_none(self):
         assert _reflow_collapsed_table_line("a | b | c") is None
-