From f69db2a75c3602aeaaa3ebe0e7d0c57ad98ab54e Mon Sep 17 00:00:00 2001
From: Ayaz-Microsoft <v-ayazkhan@microsoft.com>
Date: Fri, 14 Aug 2026 12:19:15 +0530
Subject: [PATCH 1/4] 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.
---
 src/backend/callbacks/response_handlers.py    | 51 ++++++++++++++++++-
 .../orchestration/orchestration_manager.py    |  8 +++
 .../callbacks/test_response_handlers.py       | 35 ++++++++++++-
 .../test_orchestration_manager.py             | 32 +++++++++++-
 4 files changed, 121 insertions(+), 5 deletions(-)

diff --git a/src/backend/callbacks/response_handlers.py b/src/backend/callbacks/response_handlers.py
index c2bf654ff..78e4d921d 100644
--- a/src/backend/callbacks/response_handlers.py
+++ b/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)
diff --git a/src/backend/orchestration/orchestration_manager.py b/src/backend/orchestration/orchestration_manager.py
index 77846e39f..66154c7b8 100644
--- a/src/backend/orchestration/orchestration_manager.py
+++ b/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:
diff --git a/src/tests/backend/callbacks/test_response_handlers.py b/src/tests/backend/callbacks/test_response_handlers.py
index 23a0c0ac5..4849e79db 100644
--- a/src/tests/backend/callbacks/test_response_handlers.py
+++ b/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.
diff --git a/src/tests/backend/orchestration/test_orchestration_manager.py b/src/tests/backend/orchestration/test_orchestration_manager.py
index f0a0b6166..93f1d3c23 100644
--- a/src/tests/backend/orchestration/test_orchestration_manager.py
+++ b/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
-

From 11cc79eb037a65f9e06c13c6dd17c7c4367cbf5b Mon Sep 17 00:00:00 2001
From: Ayaz-Microsoft <v-ayazkhan@microsoft.com>
Date: Fri, 14 Aug 2026 12:47:49 +0530
Subject: [PATCH 2/4] Enhance citation parsing in
 _split_trailing_partial_citation function and add unit tests for citation
 handling

---
 src/backend/callbacks/response_handlers.py       | 16 ++++++++++++----
 .../backend/callbacks/test_response_handlers.py  | 16 ++++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/backend/callbacks/response_handlers.py b/src/backend/callbacks/response_handlers.py
index 78e4d921d..d989fe12c 100644
--- a/src/backend/callbacks/response_handlers.py
+++ b/src/backend/callbacks/response_handlers.py
@@ -40,11 +40,19 @@ def _split_trailing_partial_citation(text: str) -> tuple[str, str]:
         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 = text[open_index + 1:].strip()
+        source_prefix = r"(?:s|so|sou|sour|sourc|source)"
+        is_numeric_citation = re.fullmatch(
+            rf"\d+\s*:\s*\d*(?:\s*[|†]\s*{source_prefix}?)?",
             candidate,
-        ):
+            flags=re.IGNORECASE,
+        )
+        is_source_citation = re.fullmatch(
+            rf"{source_prefix}(?::[^\]]*)?",
+            candidate,
+            flags=re.IGNORECASE,
+        )
+        if is_numeric_citation or is_source_citation:
             return text[:open_index], text[open_index:]
 
     return text, ""
diff --git a/src/tests/backend/callbacks/test_response_handlers.py b/src/tests/backend/callbacks/test_response_handlers.py
index 4849e79db..4d4813aeb 100644
--- a/src/tests/backend/callbacks/test_response_handlers.py
+++ b/src/tests/backend/callbacks/test_response_handlers.py
@@ -131,6 +131,7 @@ def __init__(self, text="", role="assistant", author_name=""):
 # Now import our module under test
 from backend.callbacks.response_handlers import (
     _extract_tool_calls_from_contents, _is_function_call_item,
+    _split_trailing_partial_citation,
     agent_response_callback, clean_citations,
     clear_streaming_citation_buffers,
     format_agent_display_name,
@@ -228,6 +229,21 @@ def test_clean_citations_preserves_formatting(self):
         assert clean_citations(text) == expected
 
 
+class TestSplitTrailingPartialCitation:
+    """Tests for distinguishing partial citations from normal bracketed text."""
+
+    def test_holds_partial_foundry_citation(self):
+        text = "Answer [5:0†sou"
+        assert _split_trailing_partial_citation(text) == ("Answer ", "[5:0†sou")
+
+    @pytest.mark.parametrize("text", [
+        "Consider [5 reasons",
+        "See the [2024 report",
+    ])
+    def test_preserves_normal_unclosed_bracketed_text(self, text):
+        assert _split_trailing_partial_citation(text) == (text, "")
+
+
 class TestFormatAgentDisplayName:
     """Tests for the format_agent_display_name function."""
 

From dd67f30cc05aaf5fb8b5387f49ebb7747b28eee5 Mon Sep 17 00:00:00 2001
From: Ayaz-Microsoft <v-ayazkhan@microsoft.com>
Date: Fri, 14 Aug 2026 15:53:53 +0530
Subject: [PATCH 3/4] revert: remove citation handling changes

Reverts commits 11cc79eb and f69db2a7.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 156d17a7-8650-4c22-b4e3-b83fa00ef7be
---
 src/backend/callbacks/response_handlers.py    | 59 +------------------
 .../orchestration/orchestration_manager.py    |  8 ---
 .../callbacks/test_response_handlers.py       | 51 +---------------
 .../test_orchestration_manager.py             | 32 +---------
 4 files changed, 5 insertions(+), 145 deletions(-)

diff --git a/src/backend/callbacks/response_handlers.py b/src/backend/callbacks/response_handlers.py
index d989fe12c..c2bf654ff 100644
--- a/src/backend/callbacks/response_handlers.py
+++ b/src/backend/callbacks/response_handlers.py
@@ -17,46 +17,6 @@
 
 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:].strip()
-        source_prefix = r"(?:s|so|sou|sour|sourc|source)"
-        is_numeric_citation = re.fullmatch(
-            rf"\d+\s*:\s*\d*(?:\s*[|†]\s*{source_prefix}?)?",
-            candidate,
-            flags=re.IGNORECASE,
-        )
-        is_source_citation = re.fullmatch(
-            rf"{source_prefix}(?::[^\]]*)?",
-            candidate,
-            flags=re.IGNORECASE,
-        )
-        if is_numeric_citation or is_source_citation:
-            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
@@ -101,12 +61,7 @@ def clean_citations(text: str) -> str:
     """Remove citation markers from agent responses while preserving formatting."""
     if not text:
         return text
-    text = re.sub(
-        r'\[\s*\d+\s*:\s*\d+\s*[|†]\s*source\s*\]',
-        '',
-        text,
-        flags=re.IGNORECASE,
-    )
+    text = re.sub(r'\[\d+:\d+\|source\]', '', text)
     text = re.sub(r'\[\s*source\s*\]', '', text, flags=re.IGNORECASE)
     text = re.sub(r'\[\d+\]', '', text)
     text = re.sub(r'【[^】]*】', '', text)
@@ -209,17 +164,7 @@ async def streaming_agent_response_callback(
                     collected.append(str(txt))
             chunk_text = "".join(collected) if collected else ""
 
-        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)
+        cleaned = clean_citations(chunk_text or "")
 
         contents = getattr(update, "contents", []) or []
         tool_calls = _extract_tool_calls_from_contents(contents)
diff --git a/src/backend/orchestration/orchestration_manager.py b/src/backend/orchestration/orchestration_manager.py
index 66154c7b8..77846e39f 100644
--- a/src/backend/orchestration/orchestration_manager.py
+++ b/src/backend/orchestration/orchestration_manager.py
@@ -17,8 +17,6 @@
                                             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
@@ -397,7 +395,6 @@ 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 = [
@@ -482,9 +479,6 @@ 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)
 
@@ -554,7 +548,6 @@ 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)
@@ -1138,7 +1131,6 @@ 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:
diff --git a/src/tests/backend/callbacks/test_response_handlers.py b/src/tests/backend/callbacks/test_response_handlers.py
index 4d4813aeb..23a0c0ac5 100644
--- a/src/tests/backend/callbacks/test_response_handlers.py
+++ b/src/tests/backend/callbacks/test_response_handlers.py
@@ -131,9 +131,7 @@ def __init__(self, text="", role="assistant", author_name=""):
 # Now import our module under test
 from backend.callbacks.response_handlers import (
     _extract_tool_calls_from_contents, _is_function_call_item,
-    _split_trailing_partial_citation,
     agent_response_callback, clean_citations,
-    clear_streaming_citation_buffers,
     format_agent_display_name,
     streaming_agent_response_callback)
 
@@ -168,18 +166,6 @@ 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."
@@ -229,21 +215,6 @@ def test_clean_citations_preserves_formatting(self):
         assert clean_citations(text) == expected
 
 
-class TestSplitTrailingPartialCitation:
-    """Tests for distinguishing partial citations from normal bracketed text."""
-
-    def test_holds_partial_foundry_citation(self):
-        text = "Answer [5:0†sou"
-        assert _split_trailing_partial_citation(text) == ("Answer ", "[5:0†sou")
-
-    @pytest.mark.parametrize("text", [
-        "Consider [5 reasons",
-        "See the [2024 report",
-    ])
-    def test_preserves_normal_unclosed_bracketed_text(self, text):
-        assert _split_trailing_partial_citation(text) == (text, "")
-
-
 class TestFormatAgentDisplayName:
     """Tests for the format_agent_display_name function."""
 
@@ -469,7 +440,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 [5:0†source]"
+        mock_message.text = "Test message with citations [1:2|source]"
         mock_message.author_name = "TestAgent"
         mock_message.role = "assistant"
 
@@ -602,7 +573,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 [5:0†source]"
+        mock_update.text = "Test streaming text [source]"
         mock_update.contents = []
 
         with patch('backend.callbacks.response_handlers.AgentMessageStreaming') as mock_streaming:
@@ -625,24 +596,6 @@ 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.
diff --git a/src/tests/backend/orchestration/test_orchestration_manager.py b/src/tests/backend/orchestration/test_orchestration_manager.py
index 93f1d3c23..f0a0b6166 100644
--- a/src/tests/backend/orchestration/test_orchestration_manager.py
+++ b/src/tests/backend/orchestration/test_orchestration_manager.py
@@ -217,10 +217,6 @@ 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(),
 )
 
@@ -320,7 +316,6 @@ 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
 
 
 # =========================================================================
@@ -611,7 +606,6 @@ 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()
@@ -667,31 +661,6 @@ 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
@@ -1102,3 +1071,4 @@ 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
+

From eb10518243a1219317e8319bb01825fb36b5ba76 Mon Sep 17 00:00:00 2001
From: Ayaz-Microsoft <v-ayazkhan@microsoft.com>
Date: Fri, 14 Aug 2026 20:09:54 +0530
Subject: [PATCH 4/4] Implement citation suppression in agent instructions and
 tests

---
 src/backend/agents/agent_factory.py           | 10 +++++++
 .../orchestration/plan_review_helpers.py      |  2 ++
 .../backend/agents/test_agent_factory.py      | 30 +++++++++++++++++++
 .../orchestration/test_plan_review_helpers.py |  6 ++++
 4 files changed, 48 insertions(+)

diff --git a/src/backend/agents/agent_factory.py b/src/backend/agents/agent_factory.py
index c27bc2106..cd7de8a90 100644
--- a/src/backend/agents/agent_factory.py
+++ b/src/backend/agents/agent_factory.py
@@ -55,6 +55,13 @@ class UnsupportedModelError(Exception):
 6. Do NOT re-ask anything already answered in the conversation history.
 """
 
+_KNOWLEDGE_BASE_NO_CITATIONS_PROMPT = """
+
+RESPONSE CITATION POLICY (CRITICAL):
+- Do not include any citation markers, source-reference tokens, attribution
+  markers, or footnotes in your response.
+"""
+
 
 class AgentFactory:
     """Create and manage teams of agents from JSON configuration.
@@ -158,6 +165,9 @@ async def create_agent_from_config(
         # Build agent instructions from system_message + optional interaction rules
         instructions = getattr(agent_obj, "system_message", "")
 
+        if kb_config:
+            instructions += _KNOWLEDGE_BASE_NO_CITATIONS_PROMPT
+
         # Universal user-interaction rules for agents that have
         # user_responses=true — tells them to call request_user_clarification.
         if user_responses:
diff --git a/src/backend/orchestration/plan_review_helpers.py b/src/backend/orchestration/plan_review_helpers.py
index efb2d536e..2effb598d 100644
--- a/src/backend/orchestration/plan_review_helpers.py
+++ b/src/backend/orchestration/plan_review_helpers.py
@@ -180,6 +180,8 @@ def get_magentic_prompt_kwargs(
   recommend, or guess any specific team, do NOT claim any action was performed, and
   do NOT attempt to answer the out-of-scope request itself.
 - Compile ONLY from messages agents actually produced. Quote verbatim where appropriate.
+- Do not include any citation markers, source-reference tokens, attribution
+  markers, or footnotes in your response.
 - Do NOT fabricate URLs, results, or content that no agent produced.
 - If a required agent step did not run, state it plainly — do not pretend it did.
 - If an agent produced an image (a markdown image ![alt](url) or an image URL such as one
diff --git a/src/tests/backend/agents/test_agent_factory.py b/src/tests/backend/agents/test_agent_factory.py
index 77f03d62b..e483e5054 100644
--- a/src/tests/backend/agents/test_agent_factory.py
+++ b/src/tests/backend/agents/test_agent_factory.py
@@ -46,6 +46,7 @@
 # --- agents sub-modules (short absolute imports in factory code)
 mock_agent_template_cls = Mock()
 mock_mcp_config_cls = Mock()
+mock_knowledge_base_config_cls = Mock()
 
 sys.modules.setdefault("agents", Mock())  # parent package stub
 _mock_agent_template_mod = Mock()
@@ -57,6 +58,7 @@
 _mock_mcp_config_mod = Mock()
 _mock_mcp_config_mod.MCPConfig = mock_mcp_config_cls
 _mock_mcp_config_mod.VectorStoreConfig = mock_vector_store_config_cls
+_mock_mcp_config_mod.KnowledgeBaseConfig = mock_knowledge_base_config_cls
 sys.modules["config.mcp_config"] = _mock_mcp_config_mod
 
 # Now import the module under test (full backend.* path as per project convention)
@@ -76,6 +78,8 @@ def _agent_obj(**overrides) -> SimpleNamespace:
         coding_tools=False,
         use_toolbox=False,
         use_file_search=False,
+        use_knowledge_base=False,
+        knowledge_base_name=None,
         user_responses=False,
         vector_store_name=None,
     )
@@ -113,6 +117,7 @@ def setup_method(self):
         self.memory_store = Mock()
         mock_agent_template_cls.reset_mock()
         mock_mcp_config_cls.reset_mock()
+        mock_knowledge_base_config_cls.reset_mock()
         mock_vector_store_config_cls.reset_mock()
 
     @pytest.mark.asyncio
@@ -164,6 +169,31 @@ async def test_user_responses_false_no_mcp_config(self):
 
         mock_mcp_config_cls.from_env.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_knowledge_base_agent_appends_no_citations_prompt(self):
+        """KB-backed agents receive citation cleanup instructions."""
+        kb_instance = Mock()
+        mock_knowledge_base_config_cls.from_env.return_value = kb_instance
+        agent_instance = Mock()
+        agent_instance.open = AsyncMock()
+        mock_agent_template_cls.return_value = agent_instance
+
+        await self.factory.create_agent_from_config(
+            "user123",
+            _agent_obj(
+                use_knowledge_base=True,
+                knowledge_base_name="test-kb",
+                system_message="Use retrieved facts.",
+            ),
+            self.team_config,
+            self.memory_store,
+        )
+
+        mock_knowledge_base_config_cls.from_env.assert_called_once_with("test-kb")
+        instructions = mock_agent_template_cls.call_args[1]["agent_instructions"]
+        assert "RESPONSE CITATION POLICY" in instructions
+        assert "Do not include any citation markers" in instructions
+
     @pytest.mark.asyncio
     async def test_use_toolbox_takes_priority_over_user_responses(self):
         """use_toolbox=True takes priority; MCPConfig uses the toolbox_filter, not 'user_responses'."""
diff --git a/src/tests/backend/orchestration/test_plan_review_helpers.py b/src/tests/backend/orchestration/test_plan_review_helpers.py
index 9bc9a595d..2ebf03047 100644
--- a/src/tests/backend/orchestration/test_plan_review_helpers.py
+++ b/src/tests/backend/orchestration/test_plan_review_helpers.py
@@ -237,6 +237,12 @@ def test_given_no_user_responses_when_called_then_final_has_answer_rules(self):
         # Assert
         assert "FINAL ANSWER RULES" in result["final_answer_prompt"]
 
+    def test_given_any_team_when_called_then_final_suppresses_citations(self):
+        result = get_magentic_prompt_kwargs(has_user_responses=False)
+
+        final_prompt = result["final_answer_prompt"]
+        assert "Do not include any citation markers" in final_prompt
+
     def test_given_default_when_called_then_user_responses_is_false(self):
         # Act
         result = get_magentic_prompt_kwargs()