unique_orchestrator 2026.32.0.dev0__tar.gz → 2026.32.0.dev2__tar.gz

{unique_orchestrator-2026.32.0.dev0 → unique_orchestrator-2026.32.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-orchestrator
-Version: 2026.32.0.dev0
+Version: 2026.32.0.dev2
 Summary: 
 Author: Andreas Hauri
 Author-email: Andreas Hauri <andreas.hauri@unique.ai>
@@ -9,15 +9,15 @@ Requires-Dist: pydantic>=2.8.2,<3
 Requires-Dist: pydantic-settings>=2.10.1,<3
 Requires-Dist: typing-extensions>=4.9.0,<5
 Requires-Dist: jinja2>=3.1.0,<4
-Requires-Dist: unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0
-Requires-Dist: unique-stock-ticker>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-toolkit>=2026.32.0.dev6,<2026.32.0rc0
+Requires-Dist: unique-stock-ticker>=2026.32.0.dev1,<2026.32.0rc0
 Requires-Dist: unique-follow-up-questions>=2026.32.0.dev0,<2026.32.0rc0
-Requires-Dist: unique-internal-search>=2026.32.0.dev0,<2026.32.0rc0
-Requires-Dist: unique-deep-research>=2026.32.0.dev0,<2026.32.0rc0
-Requires-Dist: unique-web-search>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-internal-search>=2026.32.0.dev1,<2026.32.0rc0
+Requires-Dist: unique-deep-research>=2026.32.0.dev1,<2026.32.0rc0
+Requires-Dist: unique-web-search>=2026.32.0.dev2,<2026.32.0rc0
 Requires-Dist: unique-swot>=2026.32.0.dev0,<2026.32.0rc0
 Requires-Dist: unique-skill-tool>=2026.32.0.dev0,<2026.32.0rc0
-Requires-Dist: unique-user-memory>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-user-memory>=2026.32.0.dev2,<2026.32.0rc0
 Requires-Dist: openai>=1.109.1,<3
 Requires-Python: >=3.12, <4
 Description-Content-Type: text/markdown

{unique_orchestrator-2026.32.0.dev0 → unique_orchestrator-2026.32.0.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_orchestrator"
-version = "2026.32.0.dev0"
+version = "2026.32.0.dev2"
 description = ""
 readme = "README.md"
 license = { text = "Proprietary" }
@@ -13,15 +13,15 @@ dependencies = [
     "pydantic-settings>=2.10.1,<3",
     "typing-extensions>=4.9.0,<5",
     "jinja2>=3.1.0,<4",
-    "unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0",
-    "unique-stock-ticker>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-toolkit>=2026.32.0.dev6,<2026.32.0rc0",
+    "unique-stock-ticker>=2026.32.0.dev1,<2026.32.0rc0",
     "unique-follow-up-questions>=2026.32.0.dev0,<2026.32.0rc0",
-    "unique-internal-search>=2026.32.0.dev0,<2026.32.0rc0",
-    "unique-deep-research>=2026.32.0.dev0,<2026.32.0rc0",
-    "unique-web-search>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-internal-search>=2026.32.0.dev1,<2026.32.0rc0",
+    "unique-deep-research>=2026.32.0.dev1,<2026.32.0rc0",
+    "unique-web-search>=2026.32.0.dev2,<2026.32.0rc0",
     "unique-swot>=2026.32.0.dev0,<2026.32.0rc0",
     "unique-skill-tool>=2026.32.0.dev0,<2026.32.0rc0",
-    "unique-user-memory>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-user-memory>=2026.32.0.dev2,<2026.32.0rc0",
     "openai>=1.109.1,<3",
 ]
 

{unique_orchestrator-2026.32.0.dev0 → unique_orchestrator-2026.32.0.dev2}/unique_orchestrator/tests/test_unique_ai_token_usage.py

@@ -1,9 +1,14 @@
+from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from unique_toolkit.agentic.debug_info_manager.debug_info_manager import (
+    DebugInfoManager,
+)
 from unique_toolkit.agentic.loop_runner import SupportsInvocationStats
 from unique_toolkit.agentic.tools.schemas import ToolCallResponse
 from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
+from unique_toolkit.language_model.model_costs import MODEL_COSTS_FILE_ENV
 from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
 
 from unique_orchestrator.unique_ai import UniqueAI
@@ -73,6 +78,10 @@ def _build_run_ua(loop_responses: list[MagicMock], max_iterations: int):
     mock_postprocessor_manager.run_postprocessors = AsyncMock(return_value=None)
     mock_postprocessor_manager.get_execution_times.return_value = {}
     mock_postprocessor_manager.get_invocation_stats.return_value = []
+    # A fresh list per call: `.return_value = []` would hand back the *same*
+    # list object every time, so mutations from one `ua.run()` (e.g. the
+    # main-loop usage append) would leak into the next run's list.
+    mock_postprocessor_manager.take_pending_invocation_stats.side_effect = lambda: []
 
     mock_evaluation_manager = MagicMock()
     mock_evaluation_manager.run_evaluations = AsyncMock(return_value=[])
@@ -204,9 +213,25 @@ class TestAccumulateUsage:
 class TestRunTokenUsageIntegration:
     @pytest.mark.ai
     @pytest.mark.asyncio
-    async def test_run__debug_info_manager__receives_llm_invocations_key(self):
+    async def test_run__debug_info_manager__receives_llm_invocations_key(
+        self,
+        tmp_path: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
         """A single-iteration run with real usage on the loop response must
         surface it under debug_info's 'llm_invocations' key."""
+        cost_file = tmp_path / "costs.yaml"
+        cost_file.write_text(
+            f"""
+costSchemaVersion: 1
+models:
+  {MODEL_NAME}:
+    input: 2
+    completion: 8
+""",
+            encoding="utf-8",
+        )
+        monkeypatch.setenv(MODEL_COSTS_FILE_ENV, str(cost_file))
         loop_response = _make_loop_response(
             LanguageModelTokenUsage(
                 completion_tokens=12, prompt_tokens=34, total_tokens=46
@@ -224,6 +249,7 @@ class TestRunTokenUsageIntegration:
         assert len(payload) == 1
         assert payload[0]["modelName"] == MODEL_NAME
         assert payload[0]["source"] == "main_loop[1]"
+        assert payload[0]["costUsd"] == pytest.approx(0.000164)
         assert payload[0]["tokenUsage"] == {
             "completionTokens": 12,
             "promptTokens": 34,
@@ -251,6 +277,203 @@ class TestRunTokenUsageIntegration:
         assert "llm_invocations" in calls_by_key
         assert calls_by_key["llm_invocations"] == []
 
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__cancelled_before_postprocessing__keeps_invocations_incomplete(
+        self,
+    ):
+        """
+        Purpose: Verify cancellation after the model call leaves invocation usage partial.
+        Why this matters: Postprocessor and evaluation usage has not been aggregated yet.
+        Setup summary: Cancel after planning, then assert the completion flag remains false.
+        """
+        # Arrange
+        ua = _build_run_ua([_make_loop_response(None)], max_iterations=1)
+        ua._chat_service.cancellation.check_cancellation_async = AsyncMock(
+            side_effect=[False, True]
+        )
+
+        # Act
+        await ua.run()
+
+        # Assert
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        assert calls_by_key["llm_invocations_complete"] is False
+        ua._postprocessor_manager.run_postprocessors.assert_not_awaited()
+        ua._evaluation_manager.run_evaluations.assert_not_awaited()
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__cancelled_before_postprocessing__still_reports_pending_stats(
+        self,
+    ):
+        """
+        Purpose: Verify pre-run usage (e.g. a user-memory load-time condense
+            call) is still reported when a turn exits before postprocessors run.
+        Why this matters: That usage previously lived inside a postprocessor
+            and was only surfaced when its `run()` executed; a turn that
+            cancels first silently dropped it from analytics.
+        Setup summary: Seed pending postprocessor stats, cancel before
+            postprocessing, and assert the stats still reach `llm_invocations`.
+        """
+        # Arrange
+        pending_stats = LanguageModelInvocationStats.from_usage(
+            MODEL_NAME,
+            LanguageModelTokenUsage(total_tokens=7),
+            source="user_memory_load_condense",
+        )
+        ua = _build_run_ua([_make_loop_response(None)], max_iterations=1)
+        ua._postprocessor_manager.take_pending_invocation_stats.side_effect = None
+        ua._postprocessor_manager.take_pending_invocation_stats.return_value = [
+            pending_stats
+        ]
+        ua._chat_service.cancellation.check_cancellation_async = AsyncMock(
+            side_effect=[False, True]
+        )
+
+        # Act
+        await ua.run()
+
+        # Assert
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        assert calls_by_key["llm_invocations"] == [
+            pending_stats.model_dump(by_alias=True)
+        ]
+        assert calls_by_key["llm_invocations_complete"] is False
+        ua._postprocessor_manager.run_postprocessors.assert_not_awaited()
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__tool_takes_control__keeps_invocations_incomplete(self):
+        """
+        Purpose: Verify a control-taking tool leaves invocation usage partial.
+        Why this matters: The agent exits before postprocessors and evaluations can add usage.
+        Setup summary: Execute a control-taking tool and assert completion remains false.
+        """
+        # Arrange
+        tool_call = MagicMock()
+        tool_call.name = "SubAgent"
+        loop_response = _make_loop_response(None)
+        loop_response.tool_calls = [tool_call]
+        ua = _build_run_ua([loop_response], max_iterations=1)
+        ua._tool_manager.does_a_tool_take_control.return_value = True
+        ua._tool_manager.execute_selected_tools = AsyncMock(
+            return_value=[
+                ToolCallResponse(id="call_1", name="SubAgent", invocation_stats=[])
+            ]
+        )
+
+        # Act
+        await ua.run()
+
+        # Assert
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        assert calls_by_key["llm_invocations_complete"] is False
+        ua._postprocessor_manager.run_postprocessors.assert_not_awaited()
+        ua._evaluation_manager.run_evaluations.assert_not_awaited()
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__deep_research_takes_control__marks_invocations_complete(self):
+        """
+        Purpose: Verify a Deep Research execution run signals complete usage despite
+            exiting before postprocessors and evaluations run.
+        Why this matters: Deep Research always takes control, so it never reaches
+            `_handle_no_tool_calls` and `_invocation_stats_finalized` never becomes
+            true; the completion signal for it must come from `message_execution_id`
+            alone, not be gated behind that flag.
+        Setup summary: Call Deep Research as a tool with an execution id set and
+            assert completion is still reported true.
+        """
+        # Arrange
+        tool_call = MagicMock()
+        tool_call.name = "DeepResearch"
+        loop_response = _make_loop_response(None)
+        loop_response.tool_calls = [tool_call]
+        ua = _build_run_ua([loop_response], max_iterations=1)
+        ua._debug_info_manager = DebugInfoManager()
+        ua._event.payload.message_execution_id = "execution-1"
+        ua._tool_manager.does_a_tool_take_control.return_value = True
+        ua._tool_manager.execute_selected_tools = AsyncMock(
+            return_value=[
+                ToolCallResponse(id="call_1", name="DeepResearch", invocation_stats=[])
+            ]
+        )
+
+        # Act
+        await ua.run()
+
+        # Assert
+        debug_info = ua._chat_service.update_debug_info_async.call_args.kwargs[
+            "debug_info"
+        ]
+        assert debug_info["llm_invocations_complete"] is True
+        ua._postprocessor_manager.run_postprocessors.assert_not_awaited()
+        ua._evaluation_manager.run_evaluations.assert_not_awaited()
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__max_iterations_with_tool_calls__keeps_invocations_incomplete(
+        self,
+    ):
+        """
+        Purpose: Verify tool-only exhaustion leaves invocation usage partial.
+        Why this matters: Max-iteration exit skips postprocessor and evaluation aggregation.
+        Setup summary: Exhaust one tool-call iteration and assert completion remains false.
+        """
+        # Arrange
+        tool_call = MagicMock()
+        tool_call.name = "WebSearch"
+        loop_response = _make_loop_response(None)
+        loop_response.tool_calls = [tool_call]
+        ua = _build_run_ua([loop_response], max_iterations=1)
+        ua._tool_manager.execute_selected_tools = AsyncMock(
+            return_value=[
+                ToolCallResponse(id="call_1", name="WebSearch", invocation_stats=[])
+            ]
+        )
+
+        # Act
+        await ua.run()
+
+        # Assert
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        assert calls_by_key["llm_invocations_complete"] is False
+        ua._postprocessor_manager.run_postprocessors.assert_not_awaited()
+        ua._evaluation_manager.run_evaluations.assert_not_awaited()
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__empty_response__keeps_invocations_incomplete(self):
+        """
+        Purpose: Verify an empty model response leaves invocation usage partial.
+        Why this matters: Empty responses exit before downstream LLM usage is aggregated.
+        Setup summary: Return an empty response and assert completion remains false.
+        """
+        # Arrange
+        loop_response = _make_loop_response(None)
+        loop_response.is_empty.return_value = True
+        ua = _build_run_ua([loop_response], max_iterations=1)
+
+        # Act
+        await ua.run()
+
+        # Assert
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        assert calls_by_key["llm_invocations_complete"] is False
+        ua._postprocessor_manager.run_postprocessors.assert_not_awaited()
+        ua._evaluation_manager.run_evaluations.assert_not_awaited()
+
     @pytest.mark.ai
     @pytest.mark.asyncio
     async def test_run__postprocessor_usage__added_to_loop_usage(self):
@@ -461,6 +684,136 @@ class TestRunTokenUsageIntegration:
         assert payload[0]["tokenUsage"]["totalTokens"] == 3
         assert payload[1]["tokenUsage"]["totalTokens"] == 9
 
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__deep_research__merges_previous_turn_invocations(self):
+        """
+        Purpose: Verify Deep Research preserves invocation stats from its earlier turns.
+        Why this matters: Clarification and message-execution runs update the same debugInfo.
+        Setup summary: Seed one persisted invocation, run another turn, and assert both and their analytics remain.
+        """
+        # Arrange
+        current_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=10,
+                prompt_tokens=20,
+                total_tokens=30,
+            )
+        )
+        ua = _build_run_ua([current_response], max_iterations=1)
+        ua._debug_info_manager = DebugInfoManager()
+        ua._debug_info_manager.debug_info["tools"] = [
+            {"name": "DeepResearch", "info": {}}
+        ]
+        ua._event.payload.message_execution_id = "execution-1"
+        previous_invocation = LanguageModelInvocationStats.from_usage(
+            MODEL_NAME,
+            LanguageModelTokenUsage(
+                completion_tokens=2,
+                prompt_tokens=3,
+                total_tokens=5,
+            ),
+            source="deep_research.clarification",
+        )
+        ua._chat_service.get_debug_info_async.return_value = {
+            "llm_invocations": [previous_invocation.model_dump(by_alias=True)],
+            "preserved": True,
+        }
+
+        # Act
+        await ua.run()
+
+        # Assert
+        debug_info = ua._chat_service.update_debug_info_async.call_args.kwargs[
+            "debug_info"
+        ]
+        assert debug_info["preserved"] is True
+        assert [
+            invocation["source"] for invocation in debug_info["llm_invocations"]
+        ] == ["deep_research.clarification", "main_loop[1]"]
+        assert debug_info["analytics"]["consumption_by_llm"][0]["total_tokens"] == 35
+        assert debug_info["analytics"]["consumption"]["total_tokens"] == 35
+        assert debug_info["llm_invocations_complete"] is True
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__deep_research_partial_persist__does_not_duplicate_previous_invocations(
+        self,
+    ):
+        """
+        Purpose: Verify failed Deep Research persistence does not merge prior usage twice.
+        Why this matters: Duplicate invocations overstate token analytics for aborted runs.
+        Setup summary: Fail the full persist after merging stored usage, then inspect the partial retry.
+        """
+        # Arrange
+        current_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=10,
+                prompt_tokens=20,
+                total_tokens=30,
+            )
+        )
+        ua = _build_run_ua([current_response], max_iterations=1)
+        ua._debug_info_manager = DebugInfoManager()
+        ua._debug_info_manager.debug_info["tools"] = [
+            {"name": "DeepResearch", "info": {}}
+        ]
+        ua._event.payload.message_execution_id = "execution-1"
+        previous_invocation = LanguageModelInvocationStats.from_usage(
+            MODEL_NAME,
+            LanguageModelTokenUsage(
+                completion_tokens=2,
+                prompt_tokens=3,
+                total_tokens=5,
+            ),
+            source="deep_research.clarification",
+        )
+        ua._chat_service.get_debug_info_async.return_value = {
+            "llm_invocations": [previous_invocation.model_dump(by_alias=True)],
+        }
+        ua._chat_service.update_debug_info_async.side_effect = [
+            RuntimeError("full persist failed"),
+            None,
+        ]
+
+        # Act
+        with pytest.raises(RuntimeError, match="full persist failed"):
+            await ua.run()
+
+        # Assert
+        partial_debug_info = ua._chat_service.update_debug_info_async.call_args_list[
+            1
+        ].kwargs["debug_info"]
+        assert [
+            invocation["source"] for invocation in partial_debug_info["llm_invocations"]
+        ] == ["deep_research.clarification", "main_loop[1]"]
+        assert partial_debug_info["llm_invocations_complete"] is False
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__deep_research_clarification__keeps_invocations_incomplete(self):
+        """
+        Purpose: Verify Deep Research does not signal complete usage before research execution.
+        Why this matters: Readers otherwise stop waiting after seeing only clarification calls.
+        Setup summary: Run a Deep Research turn without an execution ID and assert the flag is false.
+        """
+        # Arrange
+        ua = _build_run_ua([_make_loop_response(None)], max_iterations=1)
+        ua._debug_info_manager = DebugInfoManager()
+        ua._debug_info_manager.debug_info["tools"] = [
+            {"name": "DeepResearch", "info": {}}
+        ]
+        ua._event.payload.message_execution_id = None
+
+        # Act
+        await ua.run()
+
+        # Assert
+        debug_info = ua._chat_service.update_debug_info_async.call_args.kwargs[
+            "debug_info"
+        ]
+        assert debug_info["llm_invocations_complete"] is False
+
     @pytest.mark.ai
     @pytest.mark.asyncio
     async def test_run__reset_between_runs__does_not_leak_across_calls(self):

{unique_orchestrator-2026.32.0.dev0 → unique_orchestrator-2026.32.0.dev2}/unique_orchestrator/unique_ai.py

@@ -73,6 +73,29 @@ from unique_orchestrator.settings import env_settings
 from unique_orchestrator.utils import resolve_other_options
 
 
+def _load_invocation_stats_from_debug_info(
+    debug_info: dict[str, Any],
+    logger: Logger,
+) -> list[LanguageModelInvocationStats]:
+    """Deserialize previously persisted invocation stats, ignoring malformed entries."""
+    raw_invocations = debug_info.get("llm_invocations")
+    if not isinstance(raw_invocations, list):
+        return []
+
+    invocations: list[LanguageModelInvocationStats] = []
+    for raw_invocation in raw_invocations:
+        try:
+            invocations.append(
+                LanguageModelInvocationStats.model_validate(raw_invocation)
+            )
+        except Exception:
+            logger.warning(
+                "Ignoring malformed persisted LLM invocation stats",
+                exc_info=True,
+            )
+    return invocations
+
+
 class UniqueAI:
     start_text = ""
     current_iteration_index = 0
@@ -213,6 +236,7 @@ class UniqueAI:
         # True/False when it ran and did/didn't update the stored memory profile.
         self._context_memory_updated: bool | None = None
         self._invocation_stats: list[LanguageModelInvocationStats] = []
+        self._invocation_stats_finalized = False
 
     async def _on_cancellation(self, _event: CancellationEvent) -> None:
         """Subscriber called by the cancellation event bus."""
@@ -243,7 +267,17 @@ class UniqueAI:
         # cancellation) must not report the previous run's artifacts.
         self._generated_files_info = None
         self._context_memory_updated = None
-        self._invocation_stats = []
+        # Pending pre-run usage (e.g. the user-memory load-time condense call)
+        # must be captured unconditionally, before it's known whether this
+        # turn will even reach postprocessors -- otherwise a turn that exits
+        # early (cancellation, empty response, control-taking tool) drops it.
+        self._invocation_stats = (
+            self._postprocessor_manager.take_pending_invocation_stats()
+        )
+        self._invocation_stats_finalized = False
+        self._debug_info_manager.add("llm_invocations_complete", False)
+        invocations_persisted = False
+        persisted_invocations_merged = False
         run_start = time.perf_counter()
 
         await preload_invoked_skills(
@@ -361,6 +395,20 @@ class UniqueAI:
             self._debug_info_manager.add("loop_params", self._loop_debug_params)
             skills_debug_info = self._get_activated_skills_debug_info()
             self._debug_info_manager.add("skills", skills_debug_info)
+            tool_names = {
+                tool["name"] for tool in self._debug_info_manager.get()["tools"]
+            }
+            existing_debug_info: dict[str, Any] = {}
+            if "DeepResearch" in tool_names:
+                existing_debug_info = await self._chat_service.get_debug_info_async()
+                self._invocation_stats = [
+                    *_load_invocation_stats_from_debug_info(
+                        existing_debug_info,
+                        self._logger,
+                    ),
+                    *self._invocation_stats,
+                ]
+                persisted_invocations_merged = True
             self._debug_info_manager.add(
                 "llm_invocations",
                 [
@@ -380,9 +428,6 @@ class UniqueAI:
                 str(tool.name): tool.display_name() or str(tool.name)
                 for tool in self._tool_manager.available_tools
             }
-            tool_names = [
-                tool["name"] for tool in self._debug_info_manager.get()["tools"]
-            ]
 
             total_time_to_answer_ms: int | None = None
             if not self._chat_service.cancellation.is_cancelled:
@@ -415,16 +460,92 @@ class UniqueAI:
                 total_time_to_answer_ms=total_time_to_answer_ms,
                 artifacts=self._generated_files_info,
                 context_memory_updated=self._context_memory_updated,
+                invocations=self._invocation_stats,
             )
 
-            # Get current debug info from chat service and add debug info from run. Do not update if DeepResearch is in the tool names.
-            if "DeepResearch" not in tool_names:
+            # DeepResearch always takes control (see `takes_control()`), so a run
+            # that invokes it exits before `_handle_no_tool_calls` and never sets
+            # `_invocation_stats_finalized` -- gating on that flag would make this
+            # branch unreachable. Its completion signal is instead the message
+            # execution callback that runs the actual research (`message_execution_id`
+            # set), which is when its merged invocation stats are truly final.
+            llm_invocations_complete = (
+                self._event.payload.message_execution_id is not None
+                if "DeepResearch" in tool_names
+                else self._invocation_stats_finalized
+            )
+            self._debug_info_manager.add(
+                "llm_invocations_complete",
+                llm_invocations_complete,
+            )
+            run_debug_info = self._debug_info_manager.get()
+            if "DeepResearch" in tool_names:
                 debug_info = {
-                    **await self._chat_service.get_debug_info_async(),
-                    **self._debug_info_manager.get(),
+                    **existing_debug_info,
+                    "llm_invocations": run_debug_info["llm_invocations"],
+                    "llm_invocations_complete": llm_invocations_complete,
+                    "analytics": run_debug_info["analytics"],
                 }
-                await self._chat_service.update_debug_info_async(debug_info=debug_info)
+            else:
+                existing_debug_info = await self._chat_service.get_debug_info_async()
+                debug_info = {**existing_debug_info, **run_debug_info}
+            await self._chat_service.update_debug_info_async(debug_info=debug_info)
+            invocations_persisted = True
         finally:
+            if not invocations_persisted:
+                if isinstance(self._loop_iteration_runner, SupportsInvocationStats):
+                    self._invocation_stats.extend(
+                        self._loop_iteration_runner.get_invocation_stats()
+                    )
+                self._debug_info_manager.add(
+                    "llm_invocations",
+                    [
+                        invocation.model_dump(by_alias=True)
+                        for invocation in self._invocation_stats
+                    ],
+                )
+                self._debug_info_manager.add("llm_invocations_complete", False)
+                try:
+                    existing_debug_info = (
+                        await self._chat_service.get_debug_info_async()
+                    )
+                    partial_debug_info = self._debug_info_manager.get()
+                    is_deep_research = any(
+                        tool.get("name") == "DeepResearch"
+                        for tool in partial_debug_info.get("tools", [])
+                    )
+                    if is_deep_research:
+                        if persisted_invocations_merged:
+                            merged_invocations = partial_debug_info["llm_invocations"]
+                        else:
+                            previous_invocations = (
+                                _load_invocation_stats_from_debug_info(
+                                    existing_debug_info,
+                                    self._logger,
+                                )
+                            )
+                            merged_invocations = [
+                                *[
+                                    invocation.model_dump(by_alias=True)
+                                    for invocation in previous_invocations
+                                ],
+                                *partial_debug_info["llm_invocations"],
+                            ]
+                        debug_info = {
+                            **existing_debug_info,
+                            "llm_invocations": merged_invocations,
+                            "llm_invocations_complete": False,
+                        }
+                    else:
+                        debug_info = {**existing_debug_info, **partial_debug_info}
+                    await self._chat_service.update_debug_info_async(
+                        debug_info=debug_info
+                    )
+                except Exception:
+                    self._logger.warning(
+                        "Failed to persist partial LLM invocation usage",
+                        exc_info=True,
+                    )
             sub.cancel()
 
     @staticmethod
@@ -801,6 +922,7 @@ class UniqueAI:
                 name_str, 0
             )
         self._invocation_stats.extend(self._evaluation_manager.get_invocation_stats())
+        self._invocation_stats_finalized = True
 
         if evaluation_results.success and not all(
             result.is_positive for result in evaluation_results.unpack()